Coverage for src/common/database_engine.py: 78%

1"""Management of SQLAlchemy database engine""" 

2 

3from sqlalchemy import create_engine 

4from sqlalchemy.orm import scoped_session, sessionmaker 

5 

6from .app_config import DB_URI, SQLALCHEMY_POOL_RECYCLE, SQLALCHEMY_POOL_TIMEOUT 

7 

8# See https://docs.sqlalchemy.org/en/20/core/pooling.html#dealing-with-disconnects 

9# on how the pre_pool_ping prevents errors. 

10 

11# TODO: Flask SQLALchemy configuration function. pylint: disable=fixme 

12 

13 

14class DBEngine: # pylint: disable=too-few-public-methods 

15 """Wrapper class for use in Dramatiq worker processes, not tied to 

16 Flask and does not instantiate connections by importing this module; must 

17 create a class instance first, which is good because this code will 

18 be imported by both the Flask app container code and the Dramatiq worker 

19 container code.""" 

20 

21 def __init__(self): 

22 self.engine = create_engine( 

23 DB_URI, 

24 pool_pre_ping=True, 

25 pool_recycle=SQLALCHEMY_POOL_RECYCLE, 

26 pool_timeout=SQLALCHEMY_POOL_TIMEOUT, 

27 ) 

28 

29 def get_session(self) -> scoped_session: 

30 """Get a SQLALchemy session from the DB engine""" 

31 session_factory = sessionmaker(self.engine) 

32 return scoped_session(session_factory)