""" database/connection.py Async PostgreSQL connection manager using asyncpg. This module intentionally mimics the small subset of the aiosqlite API the rest of the codebase relies on, so services did not have to be rewritten: db = await get_db() await db.execute("INSERT ...", (a, b)) async with db.execute("SELECT ...", (a,)) as cursor: row = await cursor.fetchone() Placeholders use the SQLite-style "?" which is translated to Postgres "$1, $2" on the fly. Only "?" characters that are not inside single-quoted string literals are translated, so SQL containing a literal '?' still works. """ import logging from datetime import date, datetime from typing import Any, Iterable, Sequence import asyncpg from config.settings import DATABASE_URL logger = logging.getLogger(__name__) _pool: asyncpg.Pool | None = None # Minimum / maximum pool size. A Telegram bot is I/O bound with a handful of # concurrent updates, so a small pool keeps Postgres connection count low. POOL_MIN_SIZE = 1 POOL_MAX_SIZE = 10 # ── Value coercion ──────────────────────────────────────────────────────────── def as_date(value: Any) -> date | None: """Coerce a value into a `date` so asyncpg can bind it to a DATE column. asyncpg rejects ISO strings for DATE/TIMESTAMPTZ parameters, so every date passed as a query argument must go through here. """ if value is None: return None if isinstance(value, datetime): return value.date() if isinstance(value, date): return value text = str(value).strip() if not text: return None try: return date.fromisoformat(text[:10]) except ValueError: return datetime.fromisoformat(text).date() def as_datetime(value: Any) -> datetime | None: """Coerce a value into a `datetime` for a TIMESTAMPTZ column.""" if value is None: return None if isinstance(value, datetime): return value text = str(value).strip() if not text: return None try: return datetime.fromisoformat(text) except ValueError: pass for fmt in ("%Y-%m-%d %H:%M:%S", "%Y-%m-%d %H:%M", "%Y-%m-%d"): try: return datetime.strptime(text[:19], fmt) except ValueError: continue return None # ── Placeholder translation ────────────────────────────────────────────────── def translate_placeholders(sql: str) -> str: """Rewrite SQLite-style "?" placeholders to Postgres "$n" placeholders. Skips "?" occurrences inside single-quoted literals so values such as 'apakah?' are left untouched. """ out: list[str] = [] index = 0 n = 0 length = len(sql) while index < length: char = sql[index] if char == "'": # Consume the whole string literal (handles '' escapes). out.append(char) index += 1 while index < length: out.append(sql[index]) if sql[index] == "'": if index + 1 < length and sql[index + 1] == "'": out.append(sql[index + 1]) index += 2 continue index += 1 break index += 1 continue if char == "?": n += 1 out.append(f"${n}") index += 1 continue out.append(char) index += 1 return "".join(out) class QueryResult: """Cursor-like result object. Usable both as an awaitable (`await db.execute(...)`) and as an async context manager (`async with db.execute(...) as cursor`), which is exactly how the services were written against aiosqlite. """ def __init__(self, pool: asyncpg.Pool, sql: str, args: Sequence[Any]): self._pool = pool self._stmt = translate_placeholders(sql) self._args = tuple(args) self._rows: list[asyncpg.Record] | None = None async def _run(self) -> "QueryResult": if self._rows is None: async with self._pool.acquire() as conn: self._rows = await conn.fetch(self._stmt, *self._args) return self def __await__(self): return self._run().__await__() async def __aenter__(self) -> "QueryResult": await self._run() return self async def __aexit__(self, *exc_info) -> bool: return False # ── Cursor API ──────────────────────────────────────────────────────────── async def fetchone(self) -> asyncpg.Record | None: await self._run() return self._rows[0] if self._rows else None async def fetchall(self) -> list[asyncpg.Record]: await self._run() return list(self._rows) async def fetchval(self) -> Any: """Return the first column of the first row, else None.""" await self._run() if not self._rows: return None return self._rows[0][0] @property def rowcount(self) -> int: """Rows produced by the statement (works with `RETURNING`).""" if self._rows is None: return 0 return len(self._rows) @property def lastrowid(self) -> Any: """Compat shim: first column of the first row. Callers that need a generated id should use `INSERT ... RETURNING id` together with `fetchval()`, which is explicit and race-free. """ if self._rows: return self._rows[0][0] return None class Database: """Thin async wrapper around an asyncpg connection pool.""" def __init__(self, pool: asyncpg.Pool): self._pool = pool def execute(self, sql: str, args: Sequence[Any] | Iterable[Any] = ()) -> QueryResult: """Execute a statement. Returns a cursor-like object (await it).""" return QueryResult(self._pool, sql, tuple(args)) async def commit(self) -> None: """No-op kept for API compatibility. asyncpg runs each statement in autocommit mode, so writes are already durable once `execute()` resolves. Call sites keep their `commit()` calls so they read naturally (and so adding explicit transactions later is a one-line change). """ async def rollback(self) -> None: """No-op kept for API compatibility (see `commit`).""" # ── Convenience helpers used by the newer agentic code paths ───────────── async def fetchone(self, sql: str, args: Sequence[Any] | Iterable[Any] = ()) -> asyncpg.Record | None: return await self.execute(sql, args).fetchone() async def fetchall(self, sql: str, args: Sequence[Any] | Iterable[Any] = ()) -> list[asyncpg.Record]: return await self.execute(sql, args).fetchall() async def fetchval(self, sql: str, args: Sequence[Any] | Iterable[Any] = ()) -> Any: return await self.execute(sql, args).fetchval() def transaction(self) -> "_Transaction": """Run a block of statements atomically. Usage: async with db.transaction() as tx: await tx.execute("DELETE ...") await tx.execute("INSERT ...") """ return _Transaction(self._pool) class _Transaction: def __init__(self, pool: asyncpg.Pool): self._pool = pool self._conn: asyncpg.Connection | None = None self._tx: Any = None async def __aenter__(self) -> "_Transaction": self._conn = await self._pool.acquire() self._tx = self._conn.transaction() await self._tx.start() return self def execute(self, sql: str, args: Sequence[Any] | Iterable[Any] = ()) -> QueryResult: assert self._conn is not None, "transaction not started" return _ConnectionQuery(self._conn, sql, tuple(args)) async def __aexit__(self, exc_type, exc, tb) -> bool: assert self._tx is not None if exc_type is None: await self._tx.commit() else: await self._tx.rollback() await self._pool.release(self._conn) self._conn = None return False class _ConnectionQuery(QueryResult): """QueryResult bound to a specific connection (inside a transaction).""" def __init__(self, conn: asyncpg.Connection, sql: str, args: Sequence[Any]): # Bypass Database.__init__/pool handling by calling the helpers directly. self._pool = None # type: ignore[assignment] self._stmt = translate_placeholders(sql) self._args = tuple(args) self._rows = None self._conn = conn async def _run(self) -> "QueryResult": if self._rows is None: self._rows = await self._conn.fetch(self._stmt, *self._args) return self # ── Module-level singleton ─────────────────────────────────────────────────── async def init_db() -> Database: """Create the connection pool. Safe to call repeatedly.""" global _pool if _pool is not None: return Database(_pool) if not DATABASE_URL: raise RuntimeError( "DATABASE_URL is not set. Create a Postgres database (Neon, Supabase, " "or local) and export DATABASE_URL before starting the bot." ) logger.info("Connecting to PostgreSQL...") _pool = await asyncpg.create_pool( dsn=DATABASE_URL, min_size=POOL_MIN_SIZE, max_size=POOL_MAX_SIZE, command_timeout=30, ) logger.info("PostgreSQL connection pool ready (min=%d max=%d).", POOL_MIN_SIZE, POOL_MAX_SIZE) return Database(_pool) async def close_db() -> None: global _pool if _pool is not None: await _pool.close() _pool = None logger.info("PostgreSQL connection pool closed.") async def get_db() -> Database: """Return the shared database handle.""" if _pool is None: raise RuntimeError("Database not initialized. Call init_db() first.") return Database(_pool)