Database adapter to work with both flask-sqlachemy and sqlalchemy

flask_database_adapter.py

import abc
from sys import stderr

from flask import Flask
from flask_sqlalchemy import SQLAlchemy
from sqlalchemy import create_engine
from sqlalchemy.orm import scoped_session, sessionmaker


class DatabaseAdapter(abc.ABC):
    @abc.abstractmethod
    def connect(self, *args):
        raise NotImplementedError


class SqlAlchemyDatabaseAdapter(DatabaseAdapter):
    def __init__(self, session_factory=None):
        self.session_factory = session_factory if session_factory else sessionmaker()
        self.session: scoped_session = scoped_session(self.session_factory)

    def connect(self, app: Flask):
        database_url = app.config.get("SQLALCHEMY_DATABASE_URI")
        self.session.remove()
        self.session.configure(
            bind=create_engine(database_url),
            autoflush=False,
            autocommit=False,
            expire_on_commit=False,
        )
        self._setup_hooks(app)

    def _setup_hooks(self, app):
        @app.teardown_request
        def close_and_rollback(error=None):
            if error:
                print(error, file=stderr)
            self.session.rollback()
            self.session.close()

        @app.teardown_appcontext
        def rollback_and_remove(error=None):
            if error:
                print(error, file=stderr)
            self.session.rollback()
            self.session.remove()

    def teardown(self):
        self.session.remove()


class FlaskSqlAlchemyDatabaseAdapter(DatabaseAdapter):
    def __init__(self):
        self.db = SQLAlchemy()

    def connect(self, app: Flask):
        self.db.init_app(app)

    @property
    def session(self):
        return self.db.session
      
"""
Usage example.

in src/__init__.py:

db = FlaskSqlAlchemyDatabaseAdapter()  # or SqlAlchemyDatabaseAdapter

def create_app(testing=False):
  app = Flask(__name__)
  setup_config(testing, app)
  
  db.connect(app)
  
  return app
"""