Coverage for src/qdrant_loader/core/state/session.py: 93%

1from __future__ import annotations 

2 

3from sqlalchemy.ext.asyncio import AsyncEngine, async_sessionmaker, create_async_engine 

4from sqlalchemy.pool import StaticPool 

5 

6from qdrant_loader.config.state import StateManagementConfig 

7from qdrant_loader.core.state.models import Base 

8from qdrant_loader.core.state.utils import generate_database_url 

9 

10 

11def initialize_engine_and_session( 

12 config: StateManagementConfig, 

13) -> tuple[AsyncEngine, async_sessionmaker]: 

14 """Create the async engine and session factory for the state DB. 

15 

16 SQLite uses a StaticPool with check_same_thread disabled (single in-process 

17 file/memory DB). PostgreSQL uses a real connection pool sized from 

18 config.connection_pool, with pool_pre_ping so RDS idle-connection drops are 

19 detected and recycled rather than surfacing as errors. 

20 """ 

21 database_url = generate_database_url(config) 

22 

23 if database_url.startswith("sqlite"): 

24 is_in_memory = ":memory:" in database_url or "mode=memory" in database_url 

25 connect_args = {"check_same_thread": False, "timeout": 30} 

26 

27 if is_in_memory: 

28 engine = create_async_engine( 

29 database_url, 

30 poolclass=StaticPool, 

31 connect_args=connect_args, 

32 echo=False, 

33 ) 

34 else: 

35 engine = create_async_engine( 

36 database_url, 

37 connect_args=connect_args, 

38 echo=False, 

39 ) 

40 else: 

41 # asyncpg is an optional dependency (the `postgres` extra), imported 

42 # lazily only on the Postgres path so SQLite installs don't need it. 

43 try: 

44 import asyncpg # noqa: F401 

45 except ModuleNotFoundError as exc: 

46 raise RuntimeError( 

47 "PostgreSQL state backend requires the 'postgres' extra: " 

48 "install with `pip install qdrant-loader[postgres]`." 

49 ) from exc 

50 

51 pool_cfg = config.connection_pool or {} 

52 engine = create_async_engine( 

53 database_url, 

54 pool_size=pool_cfg.get("size", 5), 

55 pool_timeout=pool_cfg.get("timeout", 30), 

56 pool_pre_ping=True, 

57 pool_recycle=1800, # recycle connections every 30 min (RDS-friendly) 

58 echo=False, 

59 ) 

60 

61 session_factory = async_sessionmaker(engine, expire_on_commit=False) 

62 return engine, session_factory 

63 

64 

65async def create_tables(engine: AsyncEngine) -> None: 

66 """Create database tables if they do not exist.""" 

67 async with engine.begin() as conn: 

68 await conn.run_sync(Base.metadata.create_all) 

69 

70 

71async def dispose_engine(engine: AsyncEngine) -> None: 

72 """Dispose the async engine and free resources.""" 

73 await engine.dispose()