hf_database.py

"""
Database initialization for Hugging Face Spaces environment.
This creates an in-memory SQLite database for demo purposes.
"""
import logging
import os
import sqlite3
from sqlalchemy import create_engine, event
from sqlalchemy.orm import sessionmaker
from sqlalchemy.pool import StaticPool
from src.models.base import Base
from src.models.user import User
try:
    # Try to import from src.api.security first (for local development)
    from src.api.security import get_password_hash
except ImportError:
    # Fall back to simplified security module for HF (copied during deployment)
    from security_hf import get_password_hash

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Database URL for SQLite in-memory
DATABASE_URL = "sqlite:///:memory:"

# Create engine with special configuration for in-memory SQLite
engine = create_engine(
    DATABASE_URL,
    connect_args={"check_same_thread": False},
    poolclass=StaticPool,
    echo=False
)

# Add pragma for foreign key support
@event.listens_for(engine, "connect")
def set_sqlite_pragma(dbapi_connection, connection_record):
    cursor = dbapi_connection.cursor()
    cursor.execute("PRAGMA foreign_keys=ON")
    cursor.close()

# Create all tables
Base.metadata.create_all(engine)

# Create session factory
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)

def init_demo_data():
    """Initialize demo data for the in-memory database."""
    session = SessionLocal()
    try:
        # Check if we already have users
        user_count = session.query(User).count()
        if user_count == 0:
            # Create admin user
            admin_user = User(
                username="admin",
                email="[email protected]",
                full_name="Admin User",
                hashed_password=get_password_hash("adminpassword"),
                is_active=True,
                is_superuser=True
            )
            session.add(admin_user)
            
            # Create regular user
            regular_user = User(
                username="user",
                email="[email protected]",
                full_name="Regular User",
                hashed_password=get_password_hash("userpassword"),
                is_active=True,
                is_superuser=False
            )
            session.add(regular_user)
            
            # Create API user
            api_user = User(
                username="api_user",
                email="[email protected]",
                full_name="API User",
                hashed_password=get_password_hash("apipassword"),
                is_active=True,
                is_superuser=False
            )
            session.add(api_user)
            
            # Commit the session
            session.commit()
            logger.info("Demo users created successfully")
        else:
            logger.info("Demo data already exists")
            
        # Here you would add other demo data like threats, indicators, etc.
            
    except Exception as e:
        session.rollback()
        logger.error(f"Error initializing demo data: {e}")
    finally:
        session.close()

# Initialize demo data
init_demo_data()

logger.info("Hugging Face database initialized with demo data")