1"""PostgreSQL access (psycopg 3, connection pool) and SQL migrations."""2from __future__ import annotations34import hashlib5from collections.abc import Iterator6from contextlib import contextmanager7from importlib import resources8from typing import Any910import psycopg11from psycopg.rows import dict_row12from psycopg_pool import ConnectionPool1314from findata_core.config import get_settings1516_pool: ConnectionPool | None = None171819def get_pool() -> ConnectionPool:20 global _pool21 if _pool is None:22 s = get_settings()23 _pool = ConnectionPool(s.database_url.get_secret_value(), min_size=1, max_size=8, open=True,24 kwargs={"row_factory": dict_row, "autocommit": False})25 return _pool262728def close_pool() -> None:29 global _pool30 if _pool is not None:31 _pool.close()32 _pool = None333435@contextmanager36def connection() -> Iterator[psycopg.Connection[Any]]:37 pool = get_pool()38 with pool.connection() as conn:39 yield conn404142@contextmanager43def transaction() -> Iterator[psycopg.Connection[Any]]:44 with connection() as conn, conn.transaction():45 yield conn464748def fetch_all(sql: str, params: Any = None) -> list[dict[str, Any]]:49 with connection() as conn:50 return conn.execute(sql, params).fetchall()515253def fetch_one(sql: str, params: Any = None) -> dict[str, Any] | None:54 with connection() as conn:55 return conn.execute(sql, params).fetchone()565758def execute(sql: str, params: Any = None) -> None:59 with transaction() as conn:60 conn.execute(sql, params)616263def migrate(verbose: bool = True) -> list[str]:64 """Apply every `migrations/NNNN_*.sql` not yet recorded in `schema_migrations`. Idempotent."""65 applied: list[str] = []66 files = sorted(p for p in resources.files("findata_core.migrations").iterdir() if p.name.endswith(".sql"))67 with connection() as conn:68 conn.execute("CREATE TABLE IF NOT EXISTS schema_migrations (name text primary key, sha256 text not null, "69 "applied_at timestamptz not null default now())")70 conn.commit()71 done = {r["name"]: r["sha256"] for r in conn.execute("SELECT name, sha256 FROM schema_migrations").fetchall()}72 for f in files:73 sql = f.read_text(encoding="utf-8")74 digest = hashlib.sha256(sql.encode()).hexdigest()75 if f.name in done:76 if done[f.name] != digest:77 raise RuntimeError(f"migration {f.name} was modified after being applied")78 continue79 with conn.transaction():80 conn.execute(sql)81 conn.execute("INSERT INTO schema_migrations(name, sha256) VALUES (%s, %s)", (f.name, digest))82 applied.append(f.name)83 if verbose:84 print(f"applied {f.name}")85 return applied86