72 lines
1.9 KiB
Python
72 lines
1.9 KiB
Python
from typing import Optional, cast
|
|
|
|
from flask import Flask
|
|
from sqlalchemy import create_engine
|
|
from sqlalchemy.engine import Engine
|
|
from sqlalchemy.orm import Session, scoped_session, sessionmaker
|
|
|
|
from tofu_api.common.config import Config
|
|
|
|
__all__ = [
|
|
'SQLAlchemy',
|
|
]
|
|
|
|
|
|
class SQLAlchemy:
|
|
"""
|
|
Wrapper class for integrating SQLAlchemy into the application as a Flask extension.
|
|
"""
|
|
_engine: Optional[Engine] = None
|
|
_scoped_session: Optional[scoped_session] = None
|
|
|
|
def init_database(self, app: Flask):
|
|
"""
|
|
Initializes the SQLAlchemy engine.
|
|
"""
|
|
self._engine = self._create_engine(app.config)
|
|
self._scoped_session = self._create_scoped_session()
|
|
|
|
@app.teardown_appcontext
|
|
def shutdown_session(_exception=None):
|
|
self._scoped_session.remove()
|
|
|
|
@staticmethod
|
|
def _create_engine(config: Config) -> Engine:
|
|
"""
|
|
Create the database engine using the app configuration.
|
|
"""
|
|
return create_engine(
|
|
config.sqlalchemy_database_uri,
|
|
echo=config.sqlalchemy_echo,
|
|
)
|
|
|
|
def _create_scoped_session(self) -> scoped_session:
|
|
"""
|
|
Create a scoped session.
|
|
"""
|
|
return scoped_session(
|
|
sessionmaker(
|
|
autocommit=False,
|
|
autoflush=False,
|
|
bind=self._engine,
|
|
)
|
|
)
|
|
|
|
@property
|
|
def engine(self) -> Engine:
|
|
"""
|
|
Database engine.
|
|
"""
|
|
assert self._engine, 'Engine not ready yet.'
|
|
return self._engine
|
|
|
|
@property
|
|
def session(self) -> Session:
|
|
"""
|
|
Scoped database session.
|
|
"""
|
|
assert self._scoped_session, 'Session not ready yet.'
|
|
# For all further purposes, the scoped session should be treated like a regular Session object.
|
|
# Use cast() so we can use Session as the type annotation.
|
|
return cast(Session, self._scoped_session)
|