marriott-box-image-video-tagging/db.py

"""
Postgres logging for the Marriott Box tagger.

One row per file Gemini was called on (success or error). The DB is auxiliary —
all functions swallow exceptions and print to stderr so a Postgres outage cannot
stop the tagging pass. Box remains the source of truth.
"""

import json
import os
import sys
from pathlib import Path

import psycopg
from psycopg.types.json import Jsonb

SCHEMA_PATH = Path(__file__).parent / "schema.sql"

INSERT_SQL = """
INSERT INTO tagging_events (
    run_id, file_id, file_name, folder_path, media_type, gemini_model,
    prompt, raw_response, description, scenes, validated_metadata,
    metadata_write_success, description_write_success, scene_comment_write_success,
    status, error_message, duration_ms
) VALUES (
    %(run_id)s, %(file_id)s, %(file_name)s, %(folder_path)s, %(media_type)s, %(gemini_model)s,
    %(prompt)s, %(raw_response)s, %(description)s, %(scenes)s, %(validated_metadata)s,
    %(metadata_write_success)s, %(description_write_success)s, %(scene_comment_write_success)s,
    %(status)s, %(error_message)s, %(duration_ms)s
)
"""


def _dsn():
    dsn = os.getenv("DATABASE_URL")
    if not dsn:
        raise RuntimeError("DATABASE_URL not set")
    return dsn


def get_conn():
    """Open a Postgres connection. Caller owns close()."""
    return psycopg.connect(_dsn(), autocommit=True)


def ensure_schema(conn):
    """Apply schema.sql idempotently."""
    sql = SCHEMA_PATH.read_text()
    with conn.cursor() as cur:
        cur.execute(sql)


def _jsonable(value):
    if value is None:
        return None
    return Jsonb(value)


def log_event(
    conn,
    *,
    run_id,
    file_id,
    file_name,
    folder_path,
    media_type,
    gemini_model,
    status,
    prompt=None,
    raw_response=None,
    description=None,
    scenes=None,
    validated_metadata=None,
    metadata_write_success=None,
    description_write_success=None,
    scene_comment_write_success=None,
    error_message=None,
    duration_ms=None,
):
    """
    Insert one tagging_events row. Never raises — DB problems are reported to stderr
    and the tagger continues.
    """
    if conn is None:
        return
    params = {
        "run_id": str(run_id),
        "file_id": str(file_id) if file_id is not None else None,
        "file_name": file_name,
        "folder_path": folder_path,
        "media_type": media_type,
        "gemini_model": gemini_model,
        "prompt": prompt,
        "raw_response": _jsonable(raw_response),
        "description": description,
        "scenes": _jsonable(scenes),
        "validated_metadata": _jsonable(validated_metadata),
        "metadata_write_success": metadata_write_success,
        "description_write_success": description_write_success,
        "scene_comment_write_success": scene_comment_write_success,
        "status": status,
        "error_message": error_message,
        "duration_ms": duration_ms,
    }
    try:
        with conn.cursor() as cur:
            cur.execute(INSERT_SQL, params)
    except Exception as e:
        print(
            f"  WARN: DB log_event failed ({type(e).__name__}: {e}) — continuing",
            file=sys.stderr,
        )