Coverage for polar/postgres.py: 96%

1from collections.abc import AsyncGenerator 

2from typing import Literal, TypeAlias 

3 

4from fastapi import Request 

5from starlette.types import ASGIApp, Receive, Scope, Send 

6 

7from polar.config import settings 

8from polar.kit.db.postgres import ( 

9 AsyncEngine, 

10 AsyncReadSession, 

11 AsyncReadSessionMaker, 

12 AsyncSession, 

13 AsyncSessionMaker, 

14 Engine, 

15 sql, 

16) 

17from polar.kit.db.postgres import create_async_engine as _create_async_engine 

18from polar.kit.db.postgres import create_sync_engine as _create_sync_engine 

19 

20ProcessName: TypeAlias = Literal["app", "worker", "scheduler", "script"] 

21 

22 

23def create_async_engine(process_name: ProcessName) -> AsyncEngine: 

24 return _create_async_engine( 

25 dsn=str(settings.get_postgres_dsn("asyncpg")), 

26 application_name=f"{settings.ENV.value}.{process_name}", 

27 debug=settings.SQLALCHEMY_DEBUG, 

28 pool_size=settings.DATABASE_POOL_SIZE, 

29 pool_recycle=settings.DATABASE_POOL_RECYCLE_SECONDS, 

30 command_timeout=settings.DATABASE_COMMAND_TIMEOUT_SECONDS, 

31 ) 

32 

33 

34def create_async_read_engine(process_name: ProcessName) -> AsyncEngine: 

35 return _create_async_engine( 

36 dsn=str(settings.get_postgres_read_dsn("asyncpg")), 

37 application_name=f"{settings.ENV.value}.{process_name}", 

38 debug=settings.SQLALCHEMY_DEBUG, 

39 pool_size=settings.DATABASE_POOL_SIZE, 

40 pool_recycle=settings.DATABASE_POOL_RECYCLE_SECONDS, 

41 command_timeout=settings.DATABASE_COMMAND_TIMEOUT_SECONDS, 

42 ) 

43 

44 

45def create_sync_engine(process_name: ProcessName) -> Engine: 

46 return _create_sync_engine( 

47 dsn=str(settings.get_postgres_dsn("psycopg2")), 

48 application_name=f"{settings.ENV.value}.{process_name}", 

49 debug=settings.SQLALCHEMY_DEBUG, 

50 pool_size=settings.DATABASE_SYNC_POOL_SIZE, 

51 pool_recycle=settings.DATABASE_POOL_RECYCLE_SECONDS, 

52 command_timeout=settings.DATABASE_COMMAND_TIMEOUT_SECONDS, 

53 ) 

54 

55 

56class AsyncSessionMiddleware: 

57 def __init__(self, app: ASGIApp) -> None: 

58 self.app = app 

59 

60 async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: 

61 if scope["type"] not in ("http", "websocket"): 

62 return await self.app(scope, receive, send) 

63 

64 sessionmaker: AsyncSessionMaker = scope["state"]["async_sessionmaker"] 

65 async with sessionmaker() as session: 

66 scope["state"]["async_session"] = session 

67 await self.app(scope, receive, send) 

68 

69 

70async def get_db_sessionmaker(request: Request) -> AsyncSessionMaker: 

71 return request.state.async_sessionmaker 

72 

73 

74async def get_db_session(request: Request) -> AsyncGenerator[AsyncSession]: 

75 try: 

76 session = request.state.async_session 

77 except AttributeError as e: 

78 raise RuntimeError( 

79 "Session is not present in the request state. " 

80 "Did you forget to add AsyncSessionMiddleware?" 

81 ) from e 

82 

83 try: 

84 yield session 

85 except: 

86 await session.rollback() 

87 raise 

88 else: 

89 await session.commit() 

90 

91 

92async def get_db_read_session(request: Request) -> AsyncGenerator[AsyncReadSession]: 

93 sessionmaker: AsyncReadSessionMaker = request.state.async_read_sessionmaker 

94 async with sessionmaker() as session: 

95 yield session 

96 

97 

98__all__ = [ 

99 "AsyncEngine", 

100 "AsyncSession", 

101 "AsyncReadSession", 

102 "sql", 

103 "create_async_engine", 

104 "create_async_read_engine", 

105 "create_sync_engine", 

106 "get_db_session", 

107 "get_db_read_session", 

108 "get_db_sessionmaker", 

109]