db.py

"""Postgres connection and transaction helpers.

Postgres is the shared medium between tasks. Every task run gets a fresh
instance with a fresh filesystem, and task arguments and results are JSON moving
through the Render API, so nothing large or stateful travels between tasks — it
travels through here.
"""

from __future__ import annotations

import json
import logging
import os
import sys
from contextlib import contextmanager
from typing import Iterator

import psycopg
from psycopg import sql
from psycopg.rows import dict_row

from config import sync_config

logging.basicConfig(
    level=os.getenv("LOG_LEVEL", "INFO").upper(),
    format="%(asctime)s %(levelname)s %(name)s %(message)s",
    stream=sys.stdout,
)
logger = logging.getLogger("api_sync")


def database_url() -> str:
    url = os.getenv("DATABASE_URL")
    if not url:
        raise RuntimeError(
            "DATABASE_URL is not set (wire it from sync-db — see render.yaml)"
        )
    return url


@contextmanager
def connect() -> Iterator[psycopg.Connection]:
    """Open one connection with autocommit off.

    Callers wrap each unit of work in `conn.transaction()`. `fetch_slice` holds a
    single connection across its whole window and commits once per page, which is
    what makes the checkpoint and the staged rows for that page atomic.
    """
    conn = psycopg.connect(database_url(), row_factory=dict_row, autocommit=False)
    try:
        with conn.cursor() as cur:
            # Applied to every session. A runaway query should fail with its own
            # error before the enclosing Workflow task timeout kills the process.
            cur.execute(
                sql.SQL("set statement_timeout = {}").format(
                    sql.Literal(sync_config().statement_timeout_ms)
                )
            )
        conn.commit()
        yield conn
    finally:
        conn.close()


@contextmanager
def transaction() -> Iterator[psycopg.Cursor]:
    """One connection, one transaction, one commit — for single-shot work."""
    with connect() as conn:
        with conn.transaction():
            with conn.cursor() as cur:
                yield cur


def identifiers(names: tuple[str, ...] | list[str]) -> sql.Composed:
    return sql.SQL(", ").join(sql.Identifier(name) for name in names)


def log_event(level: int, event: str, **fields) -> None:
    """One JSON object per line, so the Dashboard's log search is usable."""
    logger.log(
        level, json.dumps({"event": event, **fields}, default=str, sort_keys=True)
    )