""" SQLAlchemy database setup with SQLite. Uses synchronous SQLAlchemy for simplicity and compatibility. """ import os import time import logging from contextlib import contextmanager from sqlalchemy import create_engine, event, inspect, text from sqlalchemy.orm import sessionmaker, declarative_base from app.config import get_settings settings = get_settings() logger = logging.getLogger(__name__) # ── Engine & Session ───────────────────────────────── is_sqlite = settings.DATABASE_URL.startswith("sqlite") if is_sqlite: # ── Ensure data directory exists ───────────────────── db_path = settings.DATABASE_URL.replace("sqlite:///", "") if db_path and os.path.dirname(db_path): os.makedirs(os.path.dirname(db_path), exist_ok=True) engine = create_engine( settings.DATABASE_URL, connect_args={"check_same_thread": False}, # Required for SQLite echo=settings.DEBUG, ) @event.listens_for(engine, "connect") def set_sqlite_pragma(dbapi_connection, connection_record): cursor = dbapi_connection.cursor() cursor.execute("PRAGMA journal_mode=WAL") cursor.execute("PRAGMA synchronous=NORMAL") cursor.execute("PRAGMA busy_timeout=5000") cursor.close() else: engine = create_engine( settings.DATABASE_URL, echo=settings.DEBUG, pool_size=settings.DATABASE_POOL_SIZE, max_overflow=settings.DATABASE_MAX_OVERFLOW, pool_pre_ping=settings.DATABASE_POOL_PRE_PING, pool_recycle=settings.DATABASE_POOL_RECYCLE, # Recycle stale connections ) SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine) Base = declarative_base() def get_db(): """FastAPI dependency — yields a DB session per request.""" db = SessionLocal() try: yield db finally: db.close() @contextmanager def get_db_session(): """Context manager for background tasks, streaming, and other uses outside FastAPI DI. Creates a new session, commits on success, rolls back SQLAlchemy errors (converting them to typed AppException), re-raises non-DB exceptions, and always closes the session. """ from sqlalchemy.exc import SQLAlchemyError from app.exceptions import AppException session = SessionLocal() try: yield session session.commit() except SQLAlchemyError as e: session.rollback() raise AppException( "DATABASE_ERROR", "A database error occurred while processing your request.", 500, {"error": str(e)[:200]}, ) from e except Exception: session.rollback() raise finally: session.close() _SLOW_QUERY_THRESHOLD = settings.DATABASE_SLOW_QUERY_THRESHOLD @event.listens_for(engine, "before_cursor_execute") def _before_cursor_execute(conn, cursor, statement, parameters, context, executemany): conn.info.setdefault("query_start_time", []).append(time.perf_counter()) @event.listens_for(engine, "after_cursor_execute") def _after_cursor_execute(conn, cursor, statement, parameters, context, executemany): start_times = conn.info.get("query_start_time") if not start_times: return total_time = time.perf_counter() - start_times.pop(-1) if total_time > _SLOW_QUERY_THRESHOLD: logger.warning( "Slow query detected (%.2fs): %s", total_time, statement[:500], ) # ── Session Lifecycle Logging (DEBUG only) ─────────── if settings.DEBUG: @event.listens_for(SessionLocal, "after_begin") def _receive_after_begin(session, transaction, connection): logger.debug("Session %s began transaction", id(session)) @event.listens_for(SessionLocal, "after_commit") def _receive_after_commit(session): logger.debug("Session %s committed", id(session)) @event.listens_for(SessionLocal, "after_rollback") def _receive_after_rollback(session): logger.debug("Session %s rolled back", id(session)) @event.listens_for(SessionLocal, "after_close") def _receive_after_close(session): logger.debug("Session %s closed", id(session)) def _migrate_schema(): """Apply schema migrations for existing databases (SQLite-compatible). SQLAlchemy's ``create_all`` only creates new tables and does **not** add missing columns to existing tables. This helper fills that gap for non-destructive changes such as new nullable columns. """ inspector = inspect(engine) # Migrate users existing_users_columns = {c["name"] for c in inspector.get_columns("users")} users_migrations = [ ("users", "hf_token", "ALTER TABLE users ADD COLUMN hf_token VARCHAR(255)"), ("users", "google_refresh_token", "ALTER TABLE users ADD COLUMN google_refresh_token TEXT"), ("users", "role", "ALTER TABLE users ADD COLUMN role VARCHAR(20) DEFAULT 'user'"), ("users", "last_login", "ALTER TABLE users ADD COLUMN last_login TIMESTAMP"), ("users", "is_verified", "ALTER TABLE users ADD COLUMN is_verified BOOLEAN DEFAULT TRUE NOT NULL"), ("users", "verification_token_hash", "ALTER TABLE users ADD COLUMN verification_token_hash VARCHAR(64)"), ( "users", "verification_token_created_at", "ALTER TABLE users ADD COLUMN verification_token_created_at TIMESTAMP", ), ("users", "display_name", "ALTER TABLE users ADD COLUMN display_name VARCHAR(120)"), ("users", "avatar_url", "ALTER TABLE users ADD COLUMN avatar_url VARCHAR(500)"), ] for table, column, ddl in users_migrations: if column not in existing_users_columns: try: with engine.begin() as conn: conn.execute(text(ddl)) logger.info("Migration: added column %s.%s", table, column) except Exception: logger.warning( "Migration skipped (may already exist): %s.%s", table, column ) # Migrate api_keys try: existing_keys_columns = {c["name"] for c in inspector.get_columns("api_keys")} except Exception: existing_keys_columns = set() keys_migrations = [ ("api_keys", "name", "ALTER TABLE api_keys ADD COLUMN name VARCHAR(100) DEFAULT 'default'"), ("api_keys", "is_active", "ALTER TABLE api_keys ADD COLUMN is_active BOOLEAN DEFAULT 1 NOT NULL"), ("api_keys", "last_used_at", "ALTER TABLE api_keys ADD COLUMN last_used_at TIMESTAMP"), ] for table, column, ddl in keys_migrations: if column not in existing_keys_columns: try: with engine.begin() as conn: conn.execute(text(ddl)) logger.info("Migration: added column %s.%s", table, column) except Exception: logger.warning( "Migration skipped (may already exist): %s.%s", table, column ) # Migrate documents existing_docs_columns = {c["name"] for c in inspector.get_columns("documents")} docs_migrations = [ ("documents", "last_accessed_at", "ALTER TABLE documents ADD COLUMN last_accessed_at TIMESTAMP"), ("documents", "is_deleted", "ALTER TABLE documents ADD COLUMN is_deleted BOOLEAN DEFAULT FALSE NOT NULL"), ("documents", "deleted_at", "ALTER TABLE documents ADD COLUMN deleted_at TIMESTAMP"), ("documents", "summary", "ALTER TABLE documents ADD COLUMN summary TEXT"), ("documents", "chunk_size", "ALTER TABLE documents ADD COLUMN chunk_size INTEGER"), ("documents", "chunk_overlap", "ALTER TABLE documents ADD COLUMN chunk_overlap INTEGER"), ("documents", "drive_file_id", "ALTER TABLE documents ADD COLUMN drive_file_id VARCHAR(255)"), ("documents", "drive_folder_id", "ALTER TABLE documents ADD COLUMN drive_folder_id VARCHAR(255)"), ("documents", "drive_synced_at", "ALTER TABLE documents ADD COLUMN drive_synced_at TIMESTAMP"), ("documents", "processing_progress", "ALTER TABLE documents ADD COLUMN processing_progress INTEGER DEFAULT 0"), ("documents", "processing_stage", "ALTER TABLE documents ADD COLUMN processing_stage VARCHAR(20) DEFAULT 'queued'"), ("documents", "retry_count", "ALTER TABLE documents ADD COLUMN retry_count INTEGER DEFAULT 0"), ("documents", "last_error_traceback", "ALTER TABLE documents ADD COLUMN last_error_traceback TEXT"), ("documents", "processing_started_at", "ALTER TABLE documents ADD COLUMN processing_started_at TIMESTAMP"), ("documents", "completed_at", "ALTER TABLE documents ADD COLUMN completed_at TIMESTAMP"), ("documents", "extracted_urls", "ALTER TABLE documents ADD COLUMN extracted_urls TEXT"), ("documents", "keywords", "ALTER TABLE documents ADD COLUMN keywords TEXT"), ] for table, column, ddl in docs_migrations: if column not in existing_docs_columns: try: with engine.begin() as conn: conn.execute(text(ddl)) logger.info("Migration: added column %s.%s", table, column) except Exception: logger.warning( "Migration skipped (may already exist): %s.%s", table, column ) # Migrate chat_messages try: existing_chat_columns = {c["name"] for c in inspector.get_columns("chat_messages")} except Exception: existing_chat_columns = set() chat_migrations = [ ("chat_messages", "feedback", "ALTER TABLE chat_messages ADD COLUMN feedback VARCHAR(10)"), ("chat_messages", "session_id", "ALTER TABLE chat_messages ADD COLUMN session_id CHAR(36)"), ] for table, column, ddl in chat_migrations: if column not in existing_chat_columns: try: with engine.begin() as conn: conn.execute(text(ddl)) logger.info("Migration: added column %s.%s", table, column) except Exception: logger.warning( "Migration skipped (may already exist): %s.%s", table, column ) # Ensure index exists on session_id try: with engine.begin() as conn: conn.execute(text("CREATE INDEX IF NOT EXISTS ix_chat_messages_session_id ON chat_messages (session_id)")) except Exception: pass # Migrate documents — embedding cache tracking try: existing_docs_columns = {c["name"] for c in inspector.get_columns("documents")} except Exception: existing_docs_columns = set() embedding_cache_migrations = [ ( "documents", "extracted_urls", "ALTER TABLE documents ADD COLUMN extracted_urls TEXT", ), ] for table, column, ddl in embedding_cache_migrations: if column not in existing_docs_columns: try: with engine.begin() as conn: conn.execute(text(ddl)) logger.info("Migration: added column %s.%s", table, column) except Exception: logger.warning( "Migration skipped (may already exist): %s.%s", table, column ) # ── Workspace tables ────────────────────────────────────────────────── existing_tables = set(inspector.get_table_names()) if "workspaces" not in existing_tables: try: with engine.begin() as conn: conn.execute(text(""" CREATE TABLE workspaces ( id CHAR(36) PRIMARY KEY, name VARCHAR(255) NOT NULL, created_by CHAR(36) NOT NULL REFERENCES users(id), created_at TIMESTAMP ) """)) conn.execute(text( "CREATE INDEX IF NOT EXISTS ix_workspaces_created_by " "ON workspaces (created_by)" )) logger.info("Migration: created table workspaces") except Exception: logger.warning("Migration skipped (may already exist): workspaces") if "workspace_members" not in existing_tables: try: with engine.begin() as conn: conn.execute(text(""" CREATE TABLE workspace_members ( id CHAR(36) PRIMARY KEY, workspace_id CHAR(36) NOT NULL REFERENCES workspaces(id), user_id CHAR(36) NOT NULL REFERENCES users(id), role VARCHAR(20) NOT NULL DEFAULT 'viewer', joined_at TIMESTAMP, CONSTRAINT uq_workspace_member UNIQUE (workspace_id, user_id) ) """)) conn.execute(text( "CREATE INDEX IF NOT EXISTS ix_workspace_members_workspace_id " "ON workspace_members (workspace_id)" )) conn.execute(text( "CREATE INDEX IF NOT EXISTS ix_workspace_members_user_id " "ON workspace_members (user_id)" )) logger.info("Migration: created table workspace_members") except Exception: logger.warning("Migration skipped (may already exist): workspace_members") if "workspace_invitations" not in existing_tables: try: with engine.begin() as conn: conn.execute(text(""" CREATE TABLE workspace_invitations ( id CHAR(36) PRIMARY KEY, email VARCHAR(120) NOT NULL, token_hash VARCHAR(255) NOT NULL UNIQUE, inviter_id CHAR(36) NOT NULL REFERENCES users(id), workspace_name VARCHAR(255) NOT NULL, created_at TIMESTAMP, expires_at TIMESTAMP NOT NULL, accepted_at TIMESTAMP ) """)) conn.execute(text( "CREATE INDEX IF NOT EXISTS ix_workspace_invitations_email " "ON workspace_invitations (email)" )) conn.execute(text( "CREATE INDEX IF NOT EXISTS ix_workspace_invitations_token_hash " "ON workspace_invitations (token_hash)" )) conn.execute(text( "CREATE INDEX IF NOT EXISTS ix_workspace_invitations_inviter_id " "ON workspace_invitations (inviter_id)" )) logger.info("Migration: created table workspace_invitations") except Exception: logger.warning("Migration skipped (may already exist): workspace_invitations") def advisory_lock(lock_id: int): """Context manager that acquires a PostgreSQL advisory lock (xact scope). On SQLite the lock is a no-op because SQLite serializes all writes anyway. On PostgreSQL the lock is released automatically at transaction commit. Usage:: with advisory_lock(hash("cleanup_inactive") & 0x7FFFFFFF): ... """ if is_sqlite: # SQLite serializes writes; no explicit lock needed. return _noop_contextmanager() from contextlib import contextmanager @contextmanager def _pg_lock(): with engine.begin() as conn: conn.execute(text("SELECT pg_advisory_xact_lock(:id)"), {"id": lock_id}) yield return _pg_lock() def _noop_contextmanager(): from contextlib import contextmanager as _cm @_cm def _noop(): yield return _noop() def init_db(): """Create all tables on startup and apply schema migrations.""" from app import models # noqa: F401 — import to register models Base.metadata.create_all(bind=engine) _migrate_schema()