Coverage for src/lumen/db/pool.py: 63%

1""" 

2This file manages the PostgreSQL connection pool using asyncpg. 

3It handles pool initialization, lifecycle management, and registers JSONB serialization helpers. 

4""" 

5 

6from __future__ import annotations 

7 

8import asyncio 

9import json 

10 

11import asyncpg 

12 

13from lumen.config import DatabaseSettings 

14 

15_pool: asyncpg.Pool[asyncpg.Record] | None = None 

16 

17 

18async def _init_connection(conn: asyncpg.Connection[asyncpg.Record]) -> None: 

19 """ 

20 handle jsonb decoder/encoder for asyncpg connection 

21 """ 

22 await conn.set_type_codec( 

23 "jsonb", 

24 encoder=json.dumps, 

25 decoder=json.loads, 

26 schema="pg_catalog", 

27 ) 

28 

29 

30async def init_pool(settings: DatabaseSettings | None = None) -> asyncpg.Pool[asyncpg.Record]: 

31 global _pool 

32 try: 

33 current_loop = asyncio.get_running_loop() 

34 except RuntimeError: 

35 current_loop = None 

36 

37 if _pool is not None: 

38 try: 

39 pool_loop = getattr(_pool, "_loop", None) 

40 if not getattr(_pool, "_closed", True) and pool_loop == current_loop: 

41 return _pool 

42 await _pool.close() 

43 except Exception: 

44 pass 

45 _pool = None 

46 

47 if settings is None: 

48 from lumen.config import get_settings 

49 

50 settings = get_settings().db 

51 

52 _pool = await asyncpg.create_pool( 

53 dsn=settings.url, 

54 min_size=settings.min_connections, 

55 max_size=settings.max_connections, 

56 init=_init_connection, 

57 ) 

58 return _pool 

59 

60 

61async def close_pool() -> None: 

62 global _pool 

63 if _pool is not None: 

64 await _pool.close() 

65 _pool = None 

66 

67 

68def get_pool() -> asyncpg.Pool[asyncpg.Record]: 

69 if _pool is None: 

70 raise RuntimeError("Database pool has not been initialised") 

71 return _pool