spb/crypto-trade-bot

1"""asyncpg pool + tiny migrator (plain SQL files in migrations/, tracked in schema_migrations)."""23from __future__ import annotations45import json6import re7from pathlib import Path8from typing import Any910import asyncpg1112from ktrade.logging import get_logger1314log = get_logger("db")15MIGRATIONS_DIR = Path(__file__).resolve().parents[2] / "migrations"161718def dsn(url: str) -> str:19    return re.sub(r"^postgresql\+\w+://", "postgresql://", url)202122def _json_default(o: Any) -> Any:23    import datetime as _dt24    import decimal25    import uuid2627    import numpy as np2829    if isinstance(o, np.bool_):30        return bool(o)31    if isinstance(o, np.integer):32        return int(o)33    if isinstance(o, np.floating):34        v = float(o)35        return v if v == v and v not in (float("inf"), float("-inf")) else None36    if isinstance(o, np.ndarray):37        return o.tolist()38    if isinstance(o, decimal.Decimal | uuid.UUID):39        return str(o)40    if isinstance(o, _dt.datetime | _dt.date):41        return o.isoformat()42    raise TypeError(f"not JSON serializable: {type(o).__name__}")434445def json_dumps(o: Any) -> str:46    return json.dumps(o, default=_json_default, allow_nan=False)474849async def _init_conn(conn: asyncpg.Connection) -> None:50    await conn.set_type_codec("jsonb", encoder=json_dumps, decoder=json.loads, schema="pg_catalog")51    await conn.set_type_codec("json", encoder=json_dumps, decoder=json.loads, schema="pg_catalog")525354async def create_pool(url: str, min_size: int = 1, max_size: int = 10) -> asyncpg.Pool:55    return await asyncpg.create_pool(dsn(url), min_size=min_size, max_size=max_size, init=_init_conn, command_timeout=60)565758async def migrate(url: str) -> list[str]:59    applied: list[str] = []60    conn = await asyncpg.connect(dsn(url))61    try:62        await conn.execute("CREATE TABLE IF NOT EXISTS schema_migrations (version TEXT PRIMARY KEY, applied_at TIMESTAMPTZ NOT NULL DEFAULT now())")63        done = {r["version"] for r in await conn.fetch("SELECT version FROM schema_migrations")}64        for f in sorted(MIGRATIONS_DIR.glob("*.sql")):65            if f.stem in done:66                continue67            async with conn.transaction():68                await conn.execute(f.read_text())69                await conn.execute("INSERT INTO schema_migrations (version) VALUES ($1)", f.stem)70            applied.append(f.stem)71            log.info("db.migrated", version=f.stem)72    finally:73        await conn.close()74    return applied757677async def has_timescale(pool: asyncpg.Pool) -> bool:78    return bool(await pool.fetchval("SELECT EXISTS (SELECT 1 FROM pg_extension WHERE extname='timescaledb')"))798081def rows_to_dicts(rows: list[asyncpg.Record]) -> list[dict[str, Any]]:82    return [dict(r) for r in rows]83