dtoro
574d71371f
Drops face recognition, OCR, object detection, and semantic embeddings. The sole remaining vision task is a CLIP-based binary classifier (photography vs other); photos in "other" get needs_review=true so screenshots, documents, memes and scans can be triaged from a new filter pill in the UI. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
205 lines
8.2 KiB
Python
205 lines
8.2 KiB
Python
"""
|
|
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,
|
|
pool_timeout=10, # fail fast if pool exhausted (default 30)
|
|
# Kill connections idle in a transaction for >60s. Prevents leaked
|
|
# sessions from thumbnail requests that disconnect mid-flight.
|
|
connect_args={"server_settings": {"idle_in_transaction_session_timeout": "60000"}},
|
|
)
|
|
|
|
# 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.
|
|
|
|
Rolls back any uncommitted transaction before closing so a client
|
|
disconnect doesn't leave idle-in-transaction connections in the pool.
|
|
"""
|
|
async with AsyncSessionLocal() as session:
|
|
try:
|
|
yield session
|
|
except Exception:
|
|
await session.rollback()
|
|
raise
|
|
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 # noqa: F401
|
|
|
|
# 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")
|