| """Async Postgres access layer. |
| |
| Provides a lazily-created ``asyncpg`` connection pool plus the small set of |
| queries needed by the auth routes (``/auth/signup``, ``/auth/login``), the |
| ``langgraph_sdk.Auth`` authenticate handler, and per-user monthly token |
| usage tracking. |
| |
| Requires the ``DATABASE_URL`` environment variable, e.g.:: |
| |
| postgresql://user:password@localhost:5432/personal_assistant |
| |
| The schema is in ``db/schema.sql``. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import os |
| from datetime import date |
| from typing import Any, Optional |
|
|
| import asyncpg |
|
|
| _pool: Optional[asyncpg.Pool] = None |
|
|
|
|
| async def get_pool() -> asyncpg.Pool: |
| """Return the shared connection pool, creating it on first use.""" |
| global _pool |
| if _pool is None: |
| database_url = os.environ.get("DATABASE_URL") |
| if not database_url: |
| raise RuntimeError( |
| "DATABASE_URL environment variable is required for auth. " |
| "Example: postgresql://user:password@localhost:5432/personal_assistant" |
| ) |
| _pool = await asyncpg.create_pool(dsn=database_url, min_size=1, max_size=10) |
| return _pool |
|
|
|
|
| async def close_pool() -> None: |
| """Close the connection pool (called on app shutdown).""" |
| global _pool |
| if _pool is not None: |
| await _pool.close() |
| _pool = None |
|
|
|
|
| async def create_user( |
| username: str, |
| password_hash: str, |
| display_name: Optional[str] = None, |
| email: Optional[str] = None, |
| ) -> dict[str, Any]: |
| """Insert a new user and return the created row as a dict. |
| |
| Raises ``asyncpg.UniqueViolationError`` if the username (case-insensitive) |
| or email is already taken. |
| """ |
| pool = await get_pool() |
| async with pool.acquire() as conn: |
| row = await conn.fetchrow( |
| """ |
| INSERT INTO users (username, password_hash, display_name, email) |
| VALUES ($1, $2, $3, $4) |
| RETURNING id, username, display_name, email, role, created_at |
| """, |
| username, |
| password_hash, |
| display_name or username, |
| email, |
| ) |
| return dict(row) |
|
|
|
|
| async def get_user_by_username(username: str) -> Optional[dict[str, Any]]: |
| """Fetch a user (including password hash) by case-insensitive username.""" |
| pool = await get_pool() |
| async with pool.acquire() as conn: |
| row = await conn.fetchrow( |
| """ |
| SELECT id, username, display_name, email, role, |
| password_hash, is_active |
| FROM users |
| WHERE lower(username) = lower($1) |
| """, |
| username, |
| ) |
| return dict(row) if row else None |
|
|
|
|
| async def get_user_by_id(user_id: str) -> Optional[dict[str, Any]]: |
| """Fetch a user (without password hash) by primary key.""" |
| pool = await get_pool() |
| async with pool.acquire() as conn: |
| row = await conn.fetchrow( |
| """ |
| SELECT id, username, display_name, email, role, is_active |
| FROM users |
| WHERE id = $1 |
| """, |
| user_id, |
| ) |
| return dict(row) if row else None |
|
|
|
|
| |
| |
| |
|
|
| def _current_period_start() -> date: |
| """First day of the current calendar month (UTC date).""" |
| today = date.today() |
| return today.replace(day=1) |
|
|
|
|
| async def get_token_usage(user_id: str) -> int: |
| """Tokens used by the user in the current month (0 if no row yet). |
| |
| Because usage is keyed by ``period_start`` (first of the month), the |
| allowance automatically refreshes on the first day of every month — a |
| new month simply has no row yet and reads as zero. |
| """ |
| pool = await get_pool() |
| async with pool.acquire() as conn: |
| row = await conn.fetchrow( |
| """ |
| SELECT tokens_used |
| FROM token_usage |
| WHERE user_id = $1 AND period_start = $2 |
| """, |
| user_id, |
| _current_period_start(), |
| ) |
| return int(row["tokens_used"]) if row else 0 |
|
|
|
|
| async def add_token_usage(user_id: str, tokens: int) -> int: |
| """Add ``tokens`` to the user's current-month usage; return the new total.""" |
| if tokens <= 0: |
| return await get_token_usage(user_id) |
| pool = await get_pool() |
| async with pool.acquire() as conn: |
| row = await conn.fetchrow( |
| """ |
| INSERT INTO token_usage (user_id, period_start, tokens_used) |
| VALUES ($1, $2, $3) |
| ON CONFLICT (user_id, period_start) |
| DO UPDATE SET tokens_used = token_usage.tokens_used + EXCLUDED.tokens_used |
| RETURNING tokens_used |
| """, |
| user_id, |
| _current_period_start(), |
| tokens, |
| ) |
| return int(row["tokens_used"]) |
|
|
