""" Database configuration and session management. Schema management strategy -------------------------- Postgres (default): Alembic owns schema deltas. `alembic upgrade head` is run before the app starts (in the container CMD). `init_db()` calls `create_all` afterward as the source of truth for fresh installs — it is idempotent for existing tables and creates any tables defined on `Base.metadata` that don't yet exist. Future Alembic migrations should be written defensively (`IF NOT EXISTS` etc.) so they remain safe to run on a fresh DB where `create_all` has already laid down the same objects. SQLite (escape hatch via docker-compose.sqlite.yml): no Alembic. The historical inline ALTER TABLE block stays in place so existing dev installs keep upgrading. """ import os from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine, async_sessionmaker from sqlalchemy.orm import declarative_base from sqlalchemy.pool import NullPool from sqlalchemy import text import logging from pathlib import Path from app.config import settings logger = logging.getLogger(__name__) _is_sqlite = settings.database_url.startswith("sqlite") _is_postgres = settings.database_url.startswith("postgresql") # When running inside a Celery worker we use NullPool rather than the # default connection pool. The reasons stack up: # # 1. Celery's prefork model forks the master *after* imports, so every # child inherits the same asyncpg Connection objects — they share # a socket, and two children using one concurrently raises # "another operation is in progress". # # 2. Task bodies run under `asyncio.run()`, which spins up a fresh # event loop per invocation. A pooled asyncpg Connection created # on loop A, returned to the pool, and checked out on loop B # raises "Future attached to a different loop". # # NullPool dodges both: every session checkout opens a brand-new # connection on the *current* loop and the connection is closed at # session end. Connection setup is cheap compared to task cost, so this # is the right default for the worker. The FastAPI backend keeps the # normal pool because it serves many short requests on a single long- # lived event loop, where pooling is a clear win. _is_celery_worker = os.environ.get("MULITA_CELERY_WORKER") == "1" if _is_sqlite: db_path = Path(settings.database_url.replace("sqlite+aiosqlite:///", "")) db_path.parent.mkdir(parents=True, exist_ok=True) engine = create_async_engine( settings.database_url, echo=False, connect_args={ "check_same_thread": False, "timeout": 30, }, ) elif _is_celery_worker: engine = create_async_engine( settings.database_url, echo=False, poolclass=NullPool, ) else: engine = create_async_engine( settings.database_url, echo=False, pool_size=settings.performance.db_pool_size, max_overflow=settings.performance.db_pool_max_overflow, pool_recycle=settings.performance.db_pool_recycle, pool_pre_ping=True, ) # Create async session factory AsyncSessionLocal = async_sessionmaker( engine, class_=AsyncSession, expire_on_commit=False ) # Base class for models Base = declarative_base() async def get_db() -> AsyncSession: """Dependency to get database session""" async with AsyncSessionLocal() as session: try: yield session finally: await session.close() async def init_db(): """Initialize database, create tables if they don't exist""" async with engine.begin() as conn: # Import all models to register them with Base from app.models import User, Photo, Folder, SourceRoot, Tag, PhotoTag, Heap, HeapPhoto, Embedding # Postgres: ensure pgvector is available before create_all touches # any Vector columns (added in later PRs but the extension is cheap # and idempotent to create now). if _is_postgres: await conn.execute(text("CREATE EXTENSION IF NOT EXISTS vector")) # Create all tables. Note: create_all only creates *missing* tables — # it does NOT add new columns to existing tables when the model gains # them. On Postgres, Alembic handles deltas; on SQLite, the inline # ALTER block below is the legacy fallback. await conn.run_sync(Base.metadata.create_all) gps_columns_added = False if _is_sqlite: # WAL mode for better concurrency. await conn.execute(text("PRAGMA journal_mode=WAL")) await conn.execute(text("PRAGMA synchronous=NORMAL")) await conn.execute(text("PRAGMA cache_size=10000")) await conn.execute(text("PRAGMA temp_store=MEMORY")) # ── Idempotent column adds (SQLite only) ───────────────────── # SQLite supports ADD COLUMN but not "IF NOT EXISTS" for # columns, so introspect via PRAGMA first. Each entry is # (column_name, ALTER statement). Add new columns at the # bottom. On Postgres these live in Alembic migrations. existing_cols = { row[1] for row in ( await conn.execute(text("PRAGMA table_info(photos)")) ).fetchall() } pending_alters: list[tuple[str, str]] = [ ("phash", "ALTER TABLE photos ADD COLUMN phash VARCHAR(16)"), ( "duplicate_group_id", "ALTER TABLE photos ADD COLUMN duplicate_group_id VARCHAR", ), ("latitude", "ALTER TABLE photos ADD COLUMN latitude REAL"), ("longitude", "ALTER TABLE photos ADD COLUMN longitude REAL"), ] for col_name, alter_sql in pending_alters: if col_name not in existing_cols: logger.info(f"Adding photos.{col_name} column") await conn.execute(text(alter_sql)) if col_name in ("latitude", "longitude"): gps_columns_added = True await conn.execute( text("CREATE INDEX IF NOT EXISTS ix_photos_phash ON photos(phash)") ) await conn.execute( text( "CREATE INDEX IF NOT EXISTS ix_photos_duplicate_group_id " "ON photos(duplicate_group_id)" ) ) await conn.execute( text( "CREATE INDEX IF NOT EXISTS ix_photos_lat_lon " "ON photos(latitude, longitude)" ) ) logger.info("Database initialized successfully") # If we just introduced the GPS columns on an existing SQLite # install, kick off a one-shot backfill so the Map view is # populated without a manual full re-scan. Postgres installs are # always fresh (no SQLite→PG migration path), so this code path # is SQLite-only. if _is_sqlite and gps_columns_added: try: from app.tasks.scan import backfill_gps backfill_gps.delay() logger.info("Queued one-shot backfill_gps task after column add") except Exception as e: logger.warning(f"Could not queue backfill_gps task: {e}") async def create_fts_table(): """Create Full-Text Search table for SQLite. On Postgres this is replaced by a tsvector column on the photos table (added in PR5).""" if _is_sqlite: async with engine.begin() as conn: # Create FTS5 virtual table for full-text search await conn.execute(text(""" CREATE VIRTUAL TABLE IF NOT EXISTS photos_fts USING fts5( photo_id UNINDEXED, filename, user_title, user_notes, exif_text, tokenize='unicode61' ) """)) logger.info("FTS5 table created successfully")