88 lines
2.4 KiB
Python
88 lines
2.4 KiB
Python
from typing import Optional, cast
|
|
|
|
from flask import Flask
|
|
from sqlalchemy import MetaData, create_engine
|
|
from sqlalchemy.engine import Engine
|
|
from sqlalchemy.orm import Session, scoped_session, sessionmaker
|
|
|
|
from .metadata import metadata_obj
|
|
|
|
__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()
|
|
self._scoped_session = self._create_scoped_session()
|
|
|
|
@app.teardown_appcontext
|
|
def shutdown_session(_exception=None):
|
|
self._scoped_session.remove()
|
|
|
|
def _create_engine(self) -> Engine:
|
|
"""
|
|
Create the database engine using the app configuration.
|
|
"""
|
|
# TODO: Use config
|
|
return create_engine('sqlite:////tmp/test.db')
|
|
|
|
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)
|
|
|
|
@property
|
|
def metadata(self) -> MetaData:
|
|
"""
|
|
Database metadata object.
|
|
"""
|
|
return metadata_obj
|
|
|
|
def create_all_tables(self) -> None:
|
|
"""
|
|
Create tables in the database for all models defined in the metadata.
|
|
"""
|
|
self.metadata.create_all(self.engine)
|
|
|
|
def drop_all_tables(self) -> None:
|
|
"""
|
|
Delete tables in the database for all models defined in the metadata.
|
|
"""
|
|
self.metadata.drop_all(self.engine)
|