Coverage for app/core/db.py: 93%

1import sqlite3 

2from datetime import date, datetime 

3from typing import Optional 

4 

5from sqlalchemy import Connection, Engine, create_engine, event 

6 

7from app.core.config import settings 

8from app.core.logger import logger 

9 

10 

11class DBConnectionManager: 

12 def __init__(self, db_uri: str) -> None: 

13 self.base_uri: str = db_uri 

14 self.read_engine: Engine = self._create_engine(read_only=True) 

15 self.write_engine: Engine = self._create_engine(read_only=False) 

16 self._register_event_listeners() 

17 

18 def _create_engine(self, read_only: bool) -> Engine: 

19 uri = self.base_uri 

20 if read_only: 

21 uri += "&mode=ro" 

22 return create_engine(uri, echo=False, connect_args={"check_same_thread": False}) 

23 

24 def _register_event_listeners(self) -> None: 

25 for engine in [self.read_engine, self.write_engine]: 

26 

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

28 def apply_conn_settings( 

29 dbapi_connection: Connection, connection_record: Optional[object] 

30 ) -> None: 

31 sqlite3.register_adapter(date, date.isoformat) 

32 sqlite3.register_adapter(datetime, datetime.isoformat) 

33 cursor = dbapi_connection.cursor() 

34 cursor.execute("PRAGMA foreign_keys = ON") 

35 cursor.execute("PRAGMA busy_timeout = 5000") 

36 cursor.execute("PRAGMA synchronous = normal;") 

37 cursor.execute("PRAGMA case_sensitive_like = ON") 

38 cursor.close() 

39 

40 def get_engine(self, read_only: bool = True) -> Engine: 

41 engine = self.read_engine if read_only else self.write_engine 

42 return engine 

43 

44 

45try: 

46 logger.info("Creating database manager instance") 

47 db_manager = DBConnectionManager(settings.SQLALCHEMY_DATABASE_URI) 

48 read_engine = db_manager.read_engine 

49 write_engine = db_manager.write_engine 

50 logger.info("Database engines initialized successfully") 

51except Exception as e: 

52 logger.critical(f"Failed to initialize database: {str(e)}") 

53 raise