Fixed clause_to_engine bug not finding the right db bind key (#1211)

Co-authored-by: Pamela Fox <[email protected]>

Author: ivan-gerov

CHANGES.rst

@@ -8,6 +8,7 @@ Unreleased
 -   Bump minimum version of SQLAlchemy to 2.0.16.
 -   Remove previously deprecated code.
 -   Pass extra keyword arguments from ``get_or_404`` to ``session.get``. :issue:`1149`
+-   Fix bug with finding right bind key for clause statements. :issue:`1211`
 
 
 Version 3.0.5

src/flask_sqlalchemy/session.py

@@ -79,13 +79,22 @@ def get_bind(
 
 
 def _clause_to_engine(
-    clause: t.Any | None, engines: t.Mapping[str | None, sa.engine.Engine]
+    clause: sa.ClauseElement | None,
+    engines: t.Mapping[str | None, sa.engine.Engine],
 ) -> sa.engine.Engine | None:
     """If the clause is a table, return the engine associated with the table's
     metadata's bind key.
     """
-    if isinstance(clause, sa.Table) and "bind_key" in clause.metadata.info:
-        key = clause.metadata.info["bind_key"]
+    table = None
+
+    if clause is not None:
+        if isinstance(clause, sa.Table):
+            table = clause
+        elif isinstance(clause, sa.UpdateBase) and isinstance(clause.table, sa.Table):
+            table = clause.table
+
+    if table is not None and "bind_key" in table.metadata.info:
+        key = table.metadata.info["bind_key"]
 
         if key not in engines:
             raise sa_exc.UnboundExecutionError(

tests/test_session.py

@@ -140,3 +140,68 @@ class Admin(User):  # type: ignore[no-redef]
     admin = db.session.execute(db.select(Admin)).scalar_one()
     db.session.expire(admin)
     assert admin.org == "pallets"
+
+
+@pytest.mark.usefixtures("app_ctx")
+def test_session_multiple_dbs(app: Flask, model_class: t.Any) -> None:
+    app.config["SQLALCHEMY_BINDS"] = {"db1": "sqlite:///"}
+    db = SQLAlchemy(app, model_class=model_class)
+
+    if issubclass(db.Model, (sa_orm.MappedAsDataclass)):
+
+        class User(db.Model):
+            id: sa_orm.Mapped[int] = sa_orm.mapped_column(
+                sa.Integer, primary_key=True, init=False
+            )
+            name: sa_orm.Mapped[str] = sa_orm.mapped_column(
+                sa.String(50), nullable=False, init=False
+            )
+
+        class Product(db.Model):
+            __bind_key__ = "db1"
+            id: sa_orm.Mapped[int] = sa_orm.mapped_column(
+                sa.Integer, primary_key=True, init=False
+            )
+            name: sa_orm.Mapped[str] = sa_orm.mapped_column(
+                sa.String(50), nullable=False, init=False
+            )
+
+    elif issubclass(db.Model, (sa_orm.DeclarativeBase, sa_orm.DeclarativeBaseNoMeta)):
+
+        class User(db.Model):  # type: ignore[no-redef]
+            id: sa_orm.Mapped[int] = sa_orm.mapped_column(sa.Integer, primary_key=True)
+            name: sa_orm.Mapped[str] = sa_orm.mapped_column(
+                sa.String(50), nullable=False
+            )
+
+        class Product(db.Model):  # type: ignore[no-redef]
+            __bind_key__ = "db1"
+            id: sa_orm.Mapped[int] = sa_orm.mapped_column(sa.Integer, primary_key=True)
+            name: sa_orm.Mapped[str] = sa_orm.mapped_column(
+                sa.String(50), nullable=False
+            )
+
+    else:
+
+        class User(db.Model):  # type: ignore[no-redef]
+            id = sa.Column(sa.Integer, primary_key=True)
+            name = sa.Column(sa.String(50), nullable=False)
+
+        class Product(db.Model):  # type: ignore[no-redef]
+            __bind_key__ = "db1"
+            id = sa.Column(sa.Integer, primary_key=True)
+            name = sa.Column(sa.String(50), nullable=False)
+
+    db.create_all()
+
+    db.session.execute(User.__table__.insert(), [{"name": "User1"}, {"name": "User2"}])
+    db.session.commit()
+    users = db.session.execute(db.select(User)).scalars().all()
+    assert len(users) == 2
+
+    db.session.execute(
+        Product.__table__.insert(), [{"name": "Product1"}, {"name": "Product2"}]
+    )
+    db.session.commit()
+    products = db.session.execute(db.select(Product)).scalars().all()
+    assert len(products) == 2