1"""PostgreSQL access (psycopg 3, connection pool) and SQL migrations."""2from __future__ import annotations34import hashlib5from contextlib import contextmanager6from importlib import resources7from typing import Any, Iterator89import psycopg10from psycopg.rows import dict_row11from psycopg_pool import ConnectionPool1213from findata_core.config import get_settings1415_pool: ConnectionPool | None = None161718def get_pool() -> ConnectionPool:19 global _pool20 if _pool is None:21 s = get_settings()22 _pool = ConnectionPool(s.database_url.get_secret_value(), min_size=1, max_size=8, open=True,23 kwargs={"row_factory": dict_row, "autocommit": False})24 return _pool252627def close_pool() -> None:28 global _pool29 if _pool is not None:30 _pool.close()31 _pool = None323334@contextmanager35def connection() -> Iterator[psycopg.Connection[Any]]:36 pool = get_pool()37 with pool.connection() as conn:38 yield conn394041@contextmanager42def transaction() -> Iterator[psycopg.Connection[Any]]:43 with connection() as conn:44 with 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