""" 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())