Coverage for src / quber / playground / db.py: 88%

1"""Postgres/pgvector connection helper for the ADE RAG playground. 

2 

3Reuses `quber.settings` for credentials (it loads the repo `.env` 

4automatically) and registers the pgvector type adapters so `VECTOR` 

5columns round-trip as Python lists / numpy arrays. 

6""" 

7 

8from __future__ import annotations 

9 

10from contextlib import asynccontextmanager, contextmanager 

11from pathlib import Path 

12from typing import AsyncGenerator, Generator, LiteralString, cast 

13 

14import psycopg 

15from pgvector.psycopg import register_vector, register_vector_async 

16 

17from quber.settings import get_settings 

18 

19SCHEMA_FILE = Path(__file__).with_name("schema.sql") 

20 

21 

22def dsn() -> str: 

23 db = get_settings().db 

24 return f"host={db.host} port={db.port} dbname={db.db} user={db.user} password={db.password}" 

25 

26 

27@contextmanager 

28def connect() -> Generator[psycopg.Connection]: 

29 """Yield a connection with pgvector adapters registered and autocommit on. 

30 

31 For code off the event loop: ingest, the CLI, scripts. A request handler 

32 or anything else running under asyncio uses `connect_async`, because a 

33 blocking query there holds the loop and stalls every other request's 

34 in-flight work. 

35 """ 

36 with psycopg.connect(dsn(), autocommit=True) as conn: 

37 register_vector(conn) 

38 yield conn 

39 

40 

41@asynccontextmanager 

42async def connect_async() -> AsyncGenerator[psycopg.AsyncConnection]: 

43 """The same connection for code on the event loop: queries are awaited, 

44 so the loop keeps serving other requests while Postgres works.""" 

45 async with await psycopg.AsyncConnection.connect(dsn(), autocommit=True) as conn: 

46 await register_vector_async(conn) 

47 yield conn 

48 

49 

50def apply_schema() -> None: 

51 """Create the `ade_playground` schema from schema.sql (drops existing tables).""" 

52 # psycopg types `execute` to LiteralString to keep injection out of query 

53 # strings; a schema file we author and read off disk is the same trust level. 

54 sql = cast(LiteralString, SCHEMA_FILE.read_text()) 

55 with connect() as conn: 

56 conn.execute(sql)