add chunking to selectin_polymorphic

Improved the implementation of the :func:`_orm.selectin_polymorphic`
inheritance loader strategy to properly render the IN expressions using
chunks of 500 records each, in the same manner as that of the
:func:`_orm.selectinload` relationship loader strategy.  Previously, the IN
expression would be arbitrarily large, leading to failures on databases
that have limits on the size of IN expressions including Oracle Database.

Fixes: #12790
Change-Id: I8df7f34d4fdf73996780772d983b72ea0fc8309e

diff --git a/doc/build/changelog/unreleased_20/12790.rst b/doc/build/changelog/unreleased_20/12790.rst
new file mode 100644 (file)
index 0000000..0f28b76
--- /dev/null
+++ b/doc/build/changelog/unreleased_20/12790.rst
@@ -0,0 +1,10 @@
+.. change::
+    :tags: bug, orm
+    :tickets: 12790
+
+    Improved the implementation of the :func:`_orm.selectin_polymorphic`
+    inheritance loader strategy to properly render the IN expressions using
+    chunks of 500 records each, in the same manner as that of the
+    :func:`_orm.selectinload` relationship loader strategy.  Previously, the IN
+    expression would be arbitrarily large, leading to failures on databases
+    that have limits on the size of IN expressions including Oracle Database.

diff --git a/lib/sqlalchemy/orm/loading.py b/lib/sqlalchemy/orm/loading.py
index ff28e2e2049f46c9d7ee6b619ecf8a7b89dfaa37..0e28a7a4682d0d9e06cd23a5120ede23a98dc136 100644 (file)
--- a/lib/sqlalchemy/orm/loading.py
+++ b/lib/sqlalchemy/orm/loading.py
@@ -39,6 +39,7 @@ from .base import PassiveFlag
 from .context import _ORMCompileState
 from .context import FromStatement
 from .context import QueryContext
+from .strategies import _SelectInLoader
 from .util import _none_set
 from .util import state_str
 from .. import exc as sa_exc
@@ -1309,15 +1310,18 @@ def _load_subclass_via_in(
         if context.populate_existing:
             q2 = q2.execution_options(populate_existing=True)
 
-        context.session.execute(
-            q2,
-            dict(
-                primary_keys=[
-                    state.key[1][0] if zero_idx else state.key[1]
-                    for state, load_attrs in states
-                ]
-            ),
-        ).unique().scalars().all()
+        while states:
+            chunk = states[0 : _SelectInLoader._chunksize]
+            states = states[_SelectInLoader._chunksize :]
+            context.session.execute(
+                q2,
+                dict(
+                    primary_keys=[
+                        state.key[1][0] if zero_idx else state.key[1]
+                        for state, load_attrs in chunk
+                    ]
+                ),
+            ).unique().scalars().all()
 
     return do_load
 

diff --git a/test/orm/inheritance/test_poly_loading.py b/test/orm/inheritance/test_poly_loading.py
index 58cf7b54271f302040558f5f8919ee918132476a..cf7d3146652b4ffe0029dd5729f3c51bfe95d0b3 100644 (file)
--- a/test/orm/inheritance/test_poly_loading.py
+++ b/test/orm/inheritance/test_poly_loading.py
@@ -199,6 +199,77 @@ class LoadBaseAndSubWEagerRelMapped(
         self._assert_all_selectin(q)
 
 
+class ChunkingTest(
+    fixtures.DeclarativeMappedTest, testing.AssertsExecutionResults
+):
+    @classmethod
+    def setup_classes(cls):
+        Base = cls.DeclarativeBasic
+
+        class A(Base):
+            __tablename__ = "a"
+            id = Column(Integer, primary_key=True)
+            adata = Column(String(50))
+            type = Column(String(50))
+
+            __mapper_args__ = {
+                "polymorphic_on": type,
+                "polymorphic_identity": "a",
+            }
+
+        class ASub(A):
+            __tablename__ = "asub"
+            id = Column(ForeignKey("a.id"), primary_key=True)
+            asubdata = Column(String(50))
+
+            __mapper_args__ = {
+                "polymorphic_load": "selectin",
+                "polymorphic_identity": "asub",
+            }
+
+    @classmethod
+    def insert_data(cls, connection):
+        ASub = cls.classes.ASub
+        s = Session(connection)
+        s.add_all(
+            [
+                ASub(id=i, adata=f"adata {i}", asubdata=f"asubdata {i}")
+                for i in range(1, 1255)
+            ]
+        )
+
+        s.commit()
+
+    def test_chunking(self):
+        A = self.classes.A
+        s = fixture_session()
+
+        with self.sql_execution_asserter(testing.db) as asserter:
+            asubs = s.scalars(select(A).order_by(A.id))
+            eq_(len(asubs.all()), 1254)
+
+        poly_load_sql = (
+            "SELECT asub.id AS asub_id, a.id AS a_id, a.type AS a_type, "
+            "asub.asubdata AS asub_asubdata FROM a JOIN asub "
+            "ON a.id = asub.id WHERE a.id "
+            "IN (__[POSTCOMPILE_primary_keys]) ORDER BY a.id"
+        )
+        asserter.assert_(
+            CompiledSQL(
+                "SELECT a.id, a.adata, a.type FROM a ORDER BY a.id", []
+            ),
+            CompiledSQL(
+                poly_load_sql, [{"primary_keys": list(range(1, 501))}]
+            ),
+            CompiledSQL(
+                poly_load_sql, [{"primary_keys": list(range(501, 1001))}]
+            ),
+            CompiledSQL(
+                poly_load_sql, [{"primary_keys": list(range(1001, 1255))}]
+            ),
+        )
+
+
 class FixtureLoadTest(_Polymorphic, testing.AssertsExecutionResults):
     def test_person_selectin_subclasses(self):
         s = fixture_session()