app/core/database.py

"""
Database configuration and session management
Per-session databases - each user session gets its own SQLite file
"""
from sqlalchemy import create_engine, text
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker, Session
from typing import Optional
import os
import uuid

# Ensure data directory exists
os.makedirs("data/sessions", exist_ok=True)

# Base class for models
Base = declarative_base()

# Cache for session engines
_session_engines = {}
_session_makers = {}


def get_session_engine(session_id: str):
    """Get or create engine for a specific session"""
    if session_id not in _session_engines:
        db_path = f"data/sessions/{session_id}.db"
        engine = create_engine(
            f"sqlite:///./{db_path}",
            connect_args={"check_same_thread": False}
        )
        _session_engines[session_id] = engine
        _session_makers[session_id] = sessionmaker(autocommit=False, autoflush=False, bind=engine)
        
        # Initialize tables for this session
        Base.metadata.create_all(bind=engine)
        _run_migrations(engine)
        
    return _session_engines[session_id]


def get_session_db(session_id: str):
    """Get database session for a specific user session"""
    get_session_engine(session_id)  # Ensure engine exists
    SessionLocal = _session_makers[session_id]
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()


def get_db_for_session(session_id: str) -> Session:
    """Direct session getter (non-generator) for routes"""
    get_session_engine(session_id)
    SessionLocal = _session_makers[session_id]
    return SessionLocal()


# Legacy - default database for backwards compatibility
from app.config import settings
engine = create_engine(
    settings.database_url,
    connect_args={"check_same_thread": False}
)
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)


def get_default_session() -> Session:
    """Create a new session for the default database."""
    return SessionLocal()


def get_db():
    """Legacy: Default database session"""
    db = get_default_session()
    try:
        yield db
    finally:
        db.close()


def _run_migrations(eng):
    """Run migrations on an engine"""
    with eng.connect() as conn:
        try:
            conn.execute(text("ALTER TABLE entities ADD COLUMN event_date DATETIME"))
            conn.commit()
        except Exception:
            pass
        try:
            conn.execute(text("ALTER TABLE relationships ADD COLUMN event_date DATETIME"))
            conn.commit()
        except Exception:
            pass
        try:
            conn.execute(text("ALTER TABLE entities ADD COLUMN project_id VARCHAR(36)"))
            conn.commit()
        except Exception:
            pass
        try:
            conn.execute(text("ALTER TABLE relationships ADD COLUMN project_id VARCHAR(36)"))
            conn.commit()
        except Exception:
            pass


def init_db():
    """Initialize default database tables"""
    Base.metadata.create_all(bind=engine)
    _run_migrations(engine)


def create_new_session_id() -> str:
    """Generate a new session ID"""
    return str(uuid.uuid4())