fix get_bind with polymorphic table inheritance (#1157)

Author: davidism

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

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

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]

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"