fix get_bind with polymorphic table inheritance (#1157)

Author: davidism

Diff for: CHANGES.rst

@@ -5,6 +5,8 @@ Unreleased
 
 -   Show helpful errors when mistakenly using multiple ``SQLAlchemy`` instances for the
     same app, or without calling ``init_app``. :pr:`1151`
+-   Fix issue with getting the engine associated with a model that uses polymorphic
+    table inheritance. :issue:`1155`
 
 
 Version 3.0.2

Diff for: setup.cfg

@@ -1,21 +1,24 @@
 # Config goes in pyproject.toml unless a tool doesn't support that.
 
 [flake8]
-# B = bugbear
-# E = pycodestyle errors
-# F = flake8 pyflakes
-# W = pycodestyle warnings
-# B9 = bugbear opinions
-# ISC = implicit-str-concat
-select = B, E, F, W, B9, ISC
-ignore =
+extend-select =
+    # bugbear
+    B
+    # bugbear opinions
+    B9
+    # implicit str concat
+    ISC
+extend-ignore =
     # slice notation whitespace, invalid
     E203
     # line length, handled by bugbear B950
     E501
     # bare except, handled by bugbear B001
     E722
-    # bin op line break, invalid
-    W503
+    # zip with strict=, requires python >= 3.10
+    B905
+    # string formatting opinion, B028 renamed to B907
+    B028
+    B907
 # up to 88 allowed by bugbear B950
 max-line-length = 80

Diff for: src/flask_sqlalchemy/session.py

@@ -38,6 +38,9 @@ def get_bind(  # type: ignore[override]
         """Select an engine based on the ``bind_key`` of the metadata associated with
         the model or table being queried. If no bind key is set, uses the default bind.
 
+        .. versionchanged:: 3.0.3
+            Fix finding the bind for a joined inheritance model.
+
         .. versionchanged:: 3.0
             The implementation more closely matches the base SQLAlchemy implementation.
 
@@ -47,6 +50,8 @@ def get_bind(  # type: ignore[override]
         if bind is not None:
             return bind
 
+        engines = self._db.engines
+
         if mapper is not None:
             try:
                 mapper = sa.inspect(mapper)
@@ -56,26 +61,42 @@ def get_bind(  # type: ignore[override]
 
                 raise
 
-            clause = mapper.persist_selectable
+            engine = _clause_to_engine(mapper.local_table, engines)
 
-        engines = self._db.engines
-
-        if isinstance(clause, sa.Table) and "bind_key" in clause.metadata.info:
-            key = clause.metadata.info["bind_key"]
+            if engine is not None:
+                return engine
 
-            if key not in engines:
-                raise sa.exc.UnboundExecutionError(
-                    f"Bind key '{key}' is not in 'SQLALCHEMY_BINDS' config."
-                )
+        if clause is not None:
+            engine = _clause_to_engine(clause, engines)
 
-            return engines[key]
+            if engine is not None:
+                return engine
 
         if None in engines:
             return engines[None]
 
         return super().get_bind(mapper=mapper, clause=clause, bind=bind, **kwargs)
 
 
+def _clause_to_engine(
+    clause: t.Any | 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"]
+
+        if key not in engines:
+            raise sa.exc.UnboundExecutionError(
+                f"Bind key '{key}' is not in 'SQLALCHEMY_BINDS' config."
+            )
+
+        return engines[key]
+
+    return None
+
+
 def _app_ctx_id() -> int:
     """Get the id of the current Flask application context for the session scope."""
     return id(app_ctx._get_current_object())  # type: ignore[attr-defined]

Diff for: tests/test_session.py

@@ -64,3 +64,31 @@ class Post(db.Model):
 
     assert db.session.get_bind(mapper=User) is db.engine
     assert db.session.get_bind(mapper=Post) is db.engines["a"]
+
+
+@pytest.mark.usefixtures("app_ctx")
+def test_get_bind_inheritance(app: Flask) -> None:
+    app.config["SQLALCHEMY_BINDS"] = {"a": "sqlite://"}
+    db = SQLAlchemy(app)
+
+    class User(db.Model):
+        __bind_key__ = "a"
+        id = sa.Column(sa.Integer, primary_key=True)
+        type = sa.Column(sa.String, nullable=False)
+
+        __mapper_args__ = {"polymorphic_on": type, "polymorphic_identity": "user"}
+
+    class Admin(User):
+        id = sa.Column(  # type: ignore[assignment]
+            sa.ForeignKey(User.id), primary_key=True
+        )
+        org = sa.Column(sa.String, nullable=False)
+
+        __mapper_args__ = {"polymorphic_identity": "admin"}
+
+    db.create_all()
+    db.session.add(Admin(org="pallets"))
+    db.session.commit()
+    admin = db.session.execute(db.select(Admin)).scalar_one()
+    db.session.expire(admin)
+    assert admin.org == "pallets"