SQLALCHEMY_BINDS how to query several db with the same model

[0ry0n]

I read the documentation but I did not find this use case.

If I have several databases to query, and they all have the same data model, how do I query by specifying the bind (bind_key) to which one to bind to?

Example:

settings/conf.py

class Config:
    SQLALCHEMY_DATABASE_URI = "sqlite:///app.db"

    SQLALCHEMY_BINDS = {
        "db1": "sqlite:////db1.db",
        "db2": "sqlite:////db2.db"
    }
    SQLALCHEMY_TRACK_MODIFICATIONS = False
    SECRET_KEY = "xxx"

models/init.py

from sqlalchemy import MetaData
from sqlalchemy.orm import DeclarativeBase
from flask_sqlalchemy import SQLAlchemy


class Base(DeclarativeBase):
    naming_convention = {
        "ix": "ix_%(column_0_label)s",
        "uq": "uq_%(table_name)s_%(column_0_name)s",
        "ck": "ck_%(table_name)s_%(constraint_name)s",
        "fk": "fk_%(table_name)s_%(column_0_name)s_%(referred_table_name)s",
        "pk": "pk_%(table_name)s",
    }

    metadata = MetaData(naming_convention=naming_convention)


db = SQLAlchemy(model_class=Base)

models/user.py (remode model)

from sqlalchemy.orm import Mapped, mapped_column

from models import db


class User(db.Model):
    """
    Table user in database
    """
    __tablename__ = "user"
    id: Mapped[int] = mapped_column(db.Integer, primary_key=True)
    username: Mapped[str] = mapped_column(db.String(50))
    _password: Mapped[str] = mapped_column(db.String(250), name="password")

main.py

import os

from flask import Flask
from models import db

from settings import config

conf = config.Config

app = Flask(__name__)
app.config.from_object(conf)

db.init_app(app)

if __name__ == "__main__":
    with app.app_context():
        from models.user import User

        users_from_dblocal = db.session.execute(select(User)).all()

        db.session.configure(bind=db.engines['db1'])
        users_from_db1 = db.session.execute(select(User)).all()

        db.session.configure(bind=db.engines['db2'])
        users_from_db2 = db.session.execute(select(User)).all()
    app.run(debug=True) 

I tried this db.session.configure(bind=db.engines['db1']) but it doesn’t work.
The real project is much more complicated than that.
Am I doing something wrong? Can anyone help me?

Thanks.

[kingkong-cmd]

I am wondering the same thing. I experimented with the approach from the gist below. But I keep getting this error:

sqlalchemy.exc.UnboundExecutionError: Bind key 'None' is not in 'SQLALCHEMY_BINDS' config.

Edit:

I see now, get_engine() is depreceated. Anyone know how i should restructure to achieve same with flask-sqlalchemy 3+?

Referenced gist:
https://gist.github.com/miikka/28a7bd77574a00fcec8d

Posting full stacktrace in case any smart people come looking:

Basically, get_engine is never called when the query is executed.

Query is current_user from flask_login: User.query.get(int(id))

class MultiTenantSQLAlchemy(SQLAlchemy):
    def choose_tenant(self, bind_key):
        if hasattr(g, 'tenant'):
            raise RuntimeError('Switching tenant in the middle of the request.')
        g.tenant = bind_key

    def get_engine(self, app=None, bind=None):
        if bind is None:
            if not hasattr(g, 'tenant'):
                raise RuntimeError('No tenant chosen.')
            bind = g.tenant
        return super().get_engine(app=app, bind=bind)

File "/media/Data/Projects/app/models.py", line 247, in load_user
  return User.query.get(int(id))
File "<string>", line 2, in get
File "/media/Data/Projects/venv_linux/lib/python3.10/site-packages/sqlalchemy/util/deprecations.py", line 386, in warned
  return fn(*args, **kwargs)  # type: ignore[no-any-return]
File "/media/Data/Projects/venv_linux/lib/python3.10/site-packages/sqlalchemy/orm/query.py", line 1136, in get
  return self._get_impl(ident, loading.load_on_pk_identity)
File "/media/Data/Projects/venv_linux/lib/python3.10/site-packages/sqlalchemy/orm/query.py", line 1145, in _get_impl
  return self.session._get_impl(
File "/media/Data/Projects/venv_linux/lib/python3.10/site-packages/sqlalchemy/orm/session.py", line 3833, in _get_impl
  return db_load_fn(
File "/media/Data/Projects/venv_linux/lib/python3.10/site-packages/sqlalchemy/orm/loading.py", line 690, in load_on_pk_identity
  session.execute(
File "/media/Data/Projects/venv_linux/lib/python3.10/site-packages/sqlalchemy/orm/session.py", line 2308, in execute
  return self._execute_internal(
File "/media/Data/Projects/venv_linux/lib/python3.10/site-packages/sqlalchemy/orm/session.py", line 2178, in _execute_internal
  bind = self.get_bind(**bind_arguments)
File "/media/Data/Projects/venv_linux/lib/python3.10/site-packages/flask_sqlalchemy/session.py", line 64, in get_bind
  engine = _clause_to_engine(mapper.local_table, engines)
File "/media/Data/Projects/venv_linux/lib/python3.10/site-packages/flask_sqlalchemy/session.py", line 100, in _clause_to_engine
  raise sa_exc.UnboundExecutionError(
sqlalchemy.exc.UnboundExecutionError: Bind key 'None' is not in 'SQLALCHEMY_BINDS' config.

[kingkong-cmd]

All right, I figured out one way to to it. If anyone sees this i would love it if someone picked it apart and told me why this is a stupid way of doing it.

The way i implemented it, you select a bind_key before each request (based on your own logic) and save it to g. Then you override some of flask-sqlalchemy, so that the bind_key is matched with the bind engine every time a session is created (before every query basically).

in config variables I did not set a SQLALCHEMY_DATABASE_URI value. But i set this:

SQLALCHEMY_BINDS = {
    'a': {'url' :'postgresql+psycopg2://postgres:[email protected]:5432/a'},
    'b': {'url': 'postgresql+psycopg2://postgres:[email protected]:5432/b'}
    }

Then in init.py you have to override both SQLalchemy (from flask-sqlalchemy) and the Session (also from flask-sqlalchemy). I did it like this:

from flask_sqlalchemy import SQLAlchemy
from flask_sqlalchemy.session import Session

class MultiTenantSession(Session):

    def get_bind(self, mapper=None, clause=None):
        if not hasattr(g, 'tenant'):
            raise RuntimeError('No tenant chosen.')
        bind_key = g.tenant
        return db.engines[bind_key]

class MultiTenantSQLAlchemy(SQLAlchemy):

    def choose_tenant(self, bind_key):
        if hasattr(g, 'tenant'):
            raise RuntimeError('Switching tenant in the middle of the request.')
        g.tenant = bind_key

    def _make_session_factory(self, options):
        options.setdefault("class_", MultiTenantSession)
        options.setdefault("query_cls", self.Query)

        return sa.orm.sessionmaker(db=self, **options)

Then, also in init.py file, instead of

db = SQLAlchemy()

use

db = MultiTenantSQLAlchemy()

Then, set up a way to select which bind_key to use. I used app.before_request, which will run before any request.

@app.before_request
def before_request():
    # Some logic to set the bind_key
    bind_key = 'a'
    db.choose_tenant(bind_key)

[davidism]

Using the standard SQLAlchemy pattern of creating a Session subclass to do what you want is correct, although I can't speak to the "completeness" of this specific solution. There's nothing specific to Flask-SQLAlchemy besides subclassing its Session and passing that to the extension.

Note that you don't need to subclass the extension itself, you can pass session options during init. SQLAlchemy(session_options={"class_": TenantSession}) https://flask-sqlalchemy.palletsprojects.com/en/3.1.x/customizing/#session-class

[kingkong-cmd]

Many thanks for having a look @davidism. I didn't know I could bypass subclassing the extension itself that way.

I must have read that part of the docs while trying to figure this out, but reading it again now, I assume i must have skipped it because the example showing subclassing with ShardedSession is not easy to understand.

[spmsh]

I am using that same solution @kingkong-cmd to the only difference that the SQLALCHEMY_BINDS need to be dynamic.

if it can be of any use, this is how I am doing in the MultiTenantSQLAlchemy class (feedback/improvements welcomed ofc).
The MultiTenantSession remains the same with the enhancement mentioned by @davidism (passing the options during init).

from flask import current_app, g

class MultiTenantSQLAlchemy(SQLAlchemy):
    def bind_tenant(self, tenant):
        if tenant not in self.engines:
            # creating the engine
            engine_options = self._engine_options.copy()
            engine_options["url"] = f"postgresql://user:password@localhost:5432/{tenant}",
            self.engines[tenant] = self._make_engine(tenant, engine_options, current_app)

    def switch_tenant(self, tenant):
        self.bind_tenant(tenant)
        g.tenant = tenant