Coverage for src / idx_api / database.py: 73%

1"""Database connection and session management for PostgreSQL with pgvector.""" 

2 

3from collections.abc import Generator 

4from contextlib import contextmanager 

5 

6from sqlalchemy import create_engine, event, text 

7from sqlalchemy.orm import Session, sessionmaker 

8 

9from idx_api.config import settings 

10 

11# Create PostgreSQL engine with connection pooling 

12engine = create_engine( 

13 settings.database_url, 

14 echo=settings.environment == "dev", 

15 pool_size=10, 

16 max_overflow=20, 

17 pool_pre_ping=True, # Verify connections before use 

18) 

19 

20 

21@event.listens_for(engine, "connect") 

22def enable_pgvector(dbapi_connection, connection_record): 

23 """Enable pgvector extension on first connection.""" 

24 with dbapi_connection.cursor() as cursor: 

25 cursor.execute("CREATE EXTENSION IF NOT EXISTS vector") 

26 

27 

28SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine) 

29 

30 

31def get_db() -> Generator[Session, None, None]: 

32 """Dependency for FastAPI routes.""" 

33 db = SessionLocal() 

34 try: 

35 yield db 

36 finally: 

37 db.close() 

38 

39 

40@contextmanager 

41def get_db_session() -> Generator[Session, None, None]: 

42 """Context manager for non-FastAPI usage.""" 

43 db = SessionLocal() 

44 try: 

45 yield db 

46 finally: 

47 db.close()