mule-image/backend/app/database.py

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