Spaces:
Running on Zero
Running on Zero
| """ | |
| 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] | |
| def rowcount(self) -> int: | |
| """Rows produced by the statement (works with `RETURNING`).""" | |
| if self._rows is None: | |
| return 0 | |
| return len(self._rows) | |
| 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) | |