From: Mike Bayer Date: Sat, 11 Jun 2022 15:33:46 +0000 (-0400) Subject: add auto_recurse option to selectinload, immediateload X-Git-Tag: rel_2_0_0b1~245^2 X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=b3a1162553879d1369154e920f3f4129020bb88e;p=thirdparty%2Fsqlalchemy%2Fsqlalchemy.git add auto_recurse option to selectinload, immediateload Added very experimental feature to the :func:`_orm.selectinload` and :func:`_orm.immediateload` loader options called :paramref:`_orm.selectinload.auto_recurse` / :paramref:`_orm.immediateload.auto_recurse` , which when set to True will cause a self-referential relationship load to continue loading with arbitrary depth until no further objects are found. This may be useful for self-referential structures that must be loaded fully eagerly, such as when using asyncio. Fixes: #8126 Change-Id: I5bbd00bd0ca43f4649b44680fea1e84680f0a5db --- diff --git a/doc/build/changelog/unreleased_20/8126.rst b/doc/build/changelog/unreleased_20/8126.rst new file mode 100644 index 0000000000..7afd351746 --- /dev/null +++ b/doc/build/changelog/unreleased_20/8126.rst @@ -0,0 +1,12 @@ +.. change:: + :tags: feature, orm + :tickets: 8126 + + Added very experimental feature to the :func:`_orm.selectinload` and + :func:`_orm.immediateload` loader options called + :paramref:`_orm.selectinload.auto_recurse` / + :paramref:`_orm.immediateload.auto_recurse` , which when set to True will + cause a self-referential relationship load to continue loading with + arbitrary depth until no further objects are found. This may be useful for + self-referential structures that must be loaded fully eagerly, such as when + using asyncio. diff --git a/lib/sqlalchemy/orm/strategies.py b/lib/sqlalchemy/orm/strategies.py index c4c0fb180f..d4297f84c5 100644 --- a/lib/sqlalchemy/orm/strategies.py +++ b/lib/sqlalchemy/orm/strategies.py @@ -43,6 +43,7 @@ from .interfaces import LoaderStrategy from .interfaces import StrategizedProperty from .session import _state_session from .state import InstanceState +from .strategy_options import Load from .util import _none_set from .util import AliasedClass from .. import event @@ -1259,8 +1260,24 @@ class ImmediateLoader(PostLoader): populators, ): def load_immediate(state, dict_, row): + if auto_recurse: + new_opt = Load(loadopt.path.entity) + new_opt.context = (loadopt._recurse(),) + state.load_options += (new_opt,) + state.get_impl(self.key).get(state, dict_, flags) + if loadopt and loadopt.local_opts.get("auto_recurse", False): + if not self.parent_property._is_self_referential: + raise sa_exc.InvalidRequestError( + f"auto_recurse option on relationship " + f"{self.parent_property} not valid for " + "non-self-referential relationship" + ) + auto_recurse = True + else: + auto_recurse = False + if self._check_recursive_postload(context, path): # this will not emit SQL and will only emit for a many-to-one # "use get" load. the "_RELATED" part means it may return @@ -3003,6 +3020,21 @@ class SelectInLoader(PostLoader, util.MemoizedSlots): ), ) + if loadopt and loadopt.local_opts.get("auto_recurse", False): + if not self.parent_property._is_self_referential: + # how could we do more recursion? auto_recurse is put on + # the last option in a chain, then _recurse() will walk + # backwards up the path to see where the pattern repeats + # and splice it based on that + raise sa_exc.InvalidRequestError( + f"auto_recurse option on relationship " + f"{self.parent_property} not valid for " + "non-self-referential relationship" + ) + new_opt = Load(effective_entity) + new_opt.context = (loadopt._recurse(),) + new_options += (new_opt,) + q = q.options(*new_options)._update_compile_options( {"_current_path": effective_path} ) diff --git a/lib/sqlalchemy/orm/strategy_options.py b/lib/sqlalchemy/orm/strategy_options.py index 593e2abd24..05089291f2 100644 --- a/lib/sqlalchemy/orm/strategy_options.py +++ b/lib/sqlalchemy/orm/strategy_options.py @@ -343,7 +343,7 @@ class _AbstractLoad(traversals.GenerativeOnTraversal, LoaderOption): return self._set_relationship_strategy(attr, {"lazy": "subquery"}) def selectinload( - self: Self_AbstractLoad, attr: _AttrType + self: Self_AbstractLoad, attr: _AttrType, auto_recurse: bool = False ) -> Self_AbstractLoad: """Indicate that the given attribute should be loaded using SELECT IN eager loading. @@ -365,7 +365,21 @@ class _AbstractLoad(traversals.GenerativeOnTraversal, LoaderOption): query(Order).options( lazyload(Order.items).selectinload(Item.keywords)) - .. versionadded:: 1.2 + :param auto_recurse: boolean; may be set to True when setting the + option on a self-referential relationship, which indicates that + "selectin" loading will continue an arbitrary number of levels deep + until no more items are found. + + .. note:: The :paramref:`_orm.selectinload.auto_recurse` option + currently supports only self-referential relationships. There + is not yet an option to automatically traverse recursive structures + with more than one relationship involved. + + .. warning:: This parameter is new and experimental and should be + treated as "alpha" status + + .. versionadded:: 2.0 added :paramref:`_orm.selectinload.auto_recurse` + .. seealso:: @@ -374,7 +388,9 @@ class _AbstractLoad(traversals.GenerativeOnTraversal, LoaderOption): :ref:`selectin_eager_loading` """ - return self._set_relationship_strategy(attr, {"lazy": "selectin"}) + return self._set_relationship_strategy( + attr, {"lazy": "selectin"}, opts={"auto_recurse": auto_recurse} + ) def lazyload( self: Self_AbstractLoad, attr: _AttrType @@ -395,7 +411,7 @@ class _AbstractLoad(traversals.GenerativeOnTraversal, LoaderOption): return self._set_relationship_strategy(attr, {"lazy": "select"}) def immediateload( - self: Self_AbstractLoad, attr: _AttrType + self: Self_AbstractLoad, attr: _AttrType, auto_recurse: bool = False ) -> Self_AbstractLoad: """Indicate that the given attribute should be loaded using an immediate load with a per-attribute SELECT statement. @@ -410,6 +426,23 @@ class _AbstractLoad(traversals.GenerativeOnTraversal, LoaderOption): This function is part of the :class:`_orm.Load` interface and supports both method-chained and standalone operation. + :param auto_recurse: boolean; may be set to True when setting the + option on a self-referential relationship, which indicates that + "immediate" loading will continue an arbitrary number of levels deep + until no more items are found. + + .. note:: The :paramref:`_orm.immediateload.auto_recurse` option + currently supports only self-referential relationships. There + is not yet an option to automatically traverse recursive structures + with more than one relationship involved. + + .. warning:: This parameter is new and experimental and should be + treated as "alpha" status + + .. versionadded:: 2.0 added + :paramref:`_orm.immediateload.auto_recurse` + + .. seealso:: :ref:`loading_toplevel` @@ -417,7 +450,9 @@ class _AbstractLoad(traversals.GenerativeOnTraversal, LoaderOption): :ref:`selectin_eager_loading` """ - loader = self._set_relationship_strategy(attr, {"lazy": "immediate"}) + loader = self._set_relationship_strategy( + attr, {"lazy": "immediate"}, opts={"auto_recurse": auto_recurse} + ) return loader def noload(self: Self_AbstractLoad, attr: _AttrType) -> Self_AbstractLoad: @@ -1690,7 +1725,15 @@ class _LoadElement( def __init__(self) -> None: raise NotImplementedError() - def _prepend_path_from(self, parent): + def _recurse(self) -> _LoadElement: + cloned = self._clone() + cloned.path = PathRegistry.coerce(self.path[:] + self.path[-2:]) + + return cloned + + def _prepend_path_from( + self, parent: Union[Load, _LoadElement] + ) -> _LoadElement: """adjust the path of this :class:`._LoadElement` to be a subpath of that of the given parent :class:`_orm.Load` object's path. @@ -2337,8 +2380,12 @@ def subqueryload(*keys: _AttrType) -> _AbstractLoad: @loader_unbound_fn -def selectinload(*keys: _AttrType) -> _AbstractLoad: - return _generate_from_keys(Load.selectinload, keys, False, {}) +def selectinload( + *keys: _AttrType, auto_recurse: bool = False +) -> _AbstractLoad: + return _generate_from_keys( + Load.selectinload, keys, False, {"auto_recurse": auto_recurse} + ) @loader_unbound_fn @@ -2347,8 +2394,12 @@ def lazyload(*keys: _AttrType) -> _AbstractLoad: @loader_unbound_fn -def immediateload(*keys: _AttrType) -> _AbstractLoad: - return _generate_from_keys(Load.immediateload, keys, False, {}) +def immediateload( + *keys: _AttrType, auto_recurse: bool = False +) -> _AbstractLoad: + return _generate_from_keys( + Load.immediateload, keys, False, {"auto_recurse": auto_recurse} + ) @loader_unbound_fn diff --git a/lib/sqlalchemy/testing/assertions.py b/lib/sqlalchemy/testing/assertions.py index e33d61f6db..9888d7c18d 100644 --- a/lib/sqlalchemy/testing/assertions.py +++ b/lib/sqlalchemy/testing/assertions.py @@ -801,7 +801,7 @@ class AssertsExecutionResults: return self.assert_sql_execution(db, callable_, *newrules) def assert_sql_count(self, db, callable_, count): - self.assert_sql_execution( + return self.assert_sql_execution( db, callable_, assertsql.CountStatements(count) ) diff --git a/test/orm/test_immediate_load.py b/test/orm/test_immediate_load.py index 376157bf7f..399b8f5e04 100644 --- a/test/orm/test_immediate_load.py +++ b/test/orm/test_immediate_load.py @@ -1,10 +1,19 @@ """basic tests of lazy loaded attributes""" +import sqlalchemy as sa +from sqlalchemy import ForeignKey +from sqlalchemy import Integer +from sqlalchemy import select +from sqlalchemy import String from sqlalchemy import testing from sqlalchemy.orm import immediateload from sqlalchemy.orm import relationship from sqlalchemy.testing import eq_ +from sqlalchemy.testing import expect_raises_message +from sqlalchemy.testing import fixtures from sqlalchemy.testing.fixtures import fixture_session +from sqlalchemy.testing.schema import Column +from sqlalchemy.testing.schema import Table from test.orm import _fixtures @@ -54,6 +63,35 @@ class ImmediateTest(_fixtures.FixtureTest): result, ) + def test_no_auto_recurse_non_self_referential(self): + users, Address, addresses, User = ( + self.tables.users, + self.classes.Address, + self.tables.addresses, + self.classes.User, + ) + + self.mapper_registry.map_imperatively( + User, + users, + properties={ + "addresses": relationship( + self.mapper_registry.map_imperatively(Address, addresses), + order_by=Address.id, + ) + }, + ) + sess = fixture_session() + + stmt = select(User).options( + immediateload(User.addresses, auto_recurse=True) + ) + with expect_raises_message( + sa.exc.InvalidRequestError, + "auto_recurse option on relationship User.addresses not valid", + ): + sess.execute(stmt).all() + @testing.combinations( ("raise",), ("raise_on_sql",), @@ -201,3 +239,107 @@ class ImmediateTest(_fixtures.FixtureTest): # aren't fired off. This is because the "lazyload" strategy # does not invoke eager loaders. assert "user" not in u1.addresses[0].__dict__ + + +class SelfReferentialTest(fixtures.MappedTest): + @classmethod + def define_tables(cls, metadata): + Table( + "nodes", + metadata, + Column( + "id", Integer, primary_key=True, test_needs_autoincrement=True + ), + Column("parent_id", Integer, ForeignKey("nodes.id")), + Column("data", String(30)), + ) + + @classmethod + def setup_classes(cls): + class Node(cls.Comparable): + def append(self, node): + self.children.append(node) + + @testing.fixture + def data_fixture(self): + def go(sess): + Node = self.classes.Node + n1 = Node(data="n1") + n1.append(Node(data="n11")) + n1.append(Node(data="n12")) + n1.append(Node(data="n13")) + + n1.children[0].children = [Node(data="n111"), Node(data="n112")] + + n1.children[1].append(Node(data="n121")) + n1.children[1].append(Node(data="n122")) + n1.children[1].append(Node(data="n123")) + n2 = Node(data="n2") + n2.append(Node(data="n21")) + n2.children[0].append(Node(data="n211")) + n2.children[0].append(Node(data="n212")) + sess.add(n1) + sess.add(n2) + sess.flush() + sess.expunge_all() + return n1, n2 + + return go + + def _full_structure(self): + Node = self.classes.Node + return [ + Node( + data="n1", + children=[ + Node(data="n11"), + Node( + data="n12", + children=[ + Node(data="n121"), + Node(data="n122"), + Node(data="n123"), + ], + ), + Node(data="n13"), + ], + ), + Node( + data="n2", + children=[ + Node( + data="n21", + children=[ + Node(data="n211"), + Node(data="n212"), + ], + ) + ], + ), + ] + + def test_auto_recurse_opt(self, data_fixture): + nodes = self.tables.nodes + Node = self.classes.Node + + self.mapper_registry.map_imperatively( + Node, + nodes, + properties={"children": relationship(Node)}, + ) + sess = fixture_session() + n1, n2 = data_fixture(sess) + + def go(): + return ( + sess.query(Node) + .filter(Node.data.in_(["n1", "n2"])) + .options(immediateload(Node.children, auto_recurse=True)) + .order_by(Node.data) + .all() + ) + + result = self.assert_sql_count(testing.db, go, 14) + sess.close() + + eq_(result, self._full_structure()) diff --git a/test/orm/test_selectin_relations.py b/test/orm/test_selectin_relations.py index 0458d616ed..5ba0bf8d67 100644 --- a/test/orm/test_selectin_relations.py +++ b/test/orm/test_selectin_relations.py @@ -21,6 +21,7 @@ from sqlalchemy.orm import with_polymorphic from sqlalchemy.testing import assert_raises_message from sqlalchemy.testing import assert_warns from sqlalchemy.testing import eq_ +from sqlalchemy.testing import expect_raises_message from sqlalchemy.testing import fixtures from sqlalchemy.testing import is_ from sqlalchemy.testing import is_not @@ -89,6 +90,35 @@ class EagerTest(_fixtures.FixtureTest, testing.AssertsCompiledSQL): self.assert_sql_count(testing.db, go, 2) + def test_no_auto_recurse_non_self_referential(self): + users, Address, addresses, User = ( + self.tables.users, + self.classes.Address, + self.tables.addresses, + self.classes.User, + ) + + self.mapper_registry.map_imperatively( + User, + users, + properties={ + "addresses": relationship( + self.mapper_registry.map_imperatively(Address, addresses), + order_by=Address.id, + ) + }, + ) + sess = fixture_session() + + stmt = select(User).options( + selectinload(User.addresses, auto_recurse=True) + ) + with expect_raises_message( + sa.exc.InvalidRequestError, + "auto_recurse option on relationship User.addresses not valid", + ): + sess.execute(stmt).all() + @testing.combinations(True, False) def test_from_statement(self, legacy): users, Address, addresses, User = ( @@ -2663,13 +2693,75 @@ class SelfReferentialTest(fixtures.MappedTest): Column("data", String(30)), ) - def test_basic(self): - nodes = self.tables.nodes - - class Node(fixtures.ComparableEntity): + @classmethod + def setup_classes(cls): + class Node(cls.Comparable): def append(self, node): self.children.append(node) + @testing.fixture + def data_fixture(self): + def go(sess): + Node = self.classes.Node + n1 = Node(data="n1") + n1.append(Node(data="n11")) + n1.append(Node(data="n12")) + n1.append(Node(data="n13")) + + n1.children[0].children = [Node(data="n111"), Node(data="n112")] + + n1.children[1].append(Node(data="n121")) + n1.children[1].append(Node(data="n122")) + n1.children[1].append(Node(data="n123")) + n2 = Node(data="n2") + n2.append(Node(data="n21")) + n2.children[0].append(Node(data="n211")) + n2.children[0].append(Node(data="n212")) + sess.add(n1) + sess.add(n2) + sess.flush() + sess.expunge_all() + return n1, n2 + + return go + + def _full_structure(self): + Node = self.classes.Node + return [ + Node( + data="n1", + children=[ + Node(data="n11"), + Node( + data="n12", + children=[ + Node(data="n121"), + Node(data="n122"), + Node(data="n123"), + ], + ), + Node(data="n13"), + ], + ), + Node( + data="n2", + children=[ + Node( + data="n21", + children=[ + Node(data="n211"), + Node(data="n212"), + ], + ) + ], + ), + ] + + def test_basic(self, data_fixture): + nodes = self.tables.nodes + + Node = self.classes.Node + self.mapper_registry.map_imperatively( Node, nodes, @@ -2680,22 +2772,7 @@ class SelfReferentialTest(fixtures.MappedTest): }, ) sess = fixture_session() - n1 = Node(data="n1") - n1.append(Node(data="n11")) - n1.append(Node(data="n12")) - n1.append(Node(data="n13")) - n1.children[1].append(Node(data="n121")) - n1.children[1].append(Node(data="n122")) - n1.children[1].append(Node(data="n123")) - n2 = Node(data="n2") - n2.append(Node(data="n21")) - n2.children[0].append(Node(data="n211")) - n2.children[0].append(Node(data="n212")) - - sess.add(n1) - sess.add(n2) - sess.flush() - sess.expunge_all() + n1, n2 = data_fixture(sess) def go(): d = ( @@ -2705,46 +2782,15 @@ class SelfReferentialTest(fixtures.MappedTest): .all() ) eq_( - [ - Node( - data="n1", - children=[ - Node(data="n11"), - Node( - data="n12", - children=[ - Node(data="n121"), - Node(data="n122"), - Node(data="n123"), - ], - ), - Node(data="n13"), - ], - ), - Node( - data="n2", - children=[ - Node( - data="n21", - children=[ - Node(data="n211"), - Node(data="n212"), - ], - ) - ], - ), - ], + self._full_structure(), d, ) self.assert_sql_count(testing.db, go, 4) - def test_lazy_fallback_doesnt_affect_eager(self): + def test_lazy_fallback_doesnt_affect_eager(self, data_fixture): nodes = self.tables.nodes - - class Node(fixtures.ComparableEntity): - def append(self, node): - self.children.append(node) + Node = self.classes.Node self.mapper_registry.map_imperatively( Node, @@ -2756,18 +2802,7 @@ class SelfReferentialTest(fixtures.MappedTest): }, ) sess = fixture_session() - n1 = Node(data="n1") - n1.append(Node(data="n11")) - n1.append(Node(data="n12")) - n1.append(Node(data="n13")) - n1.children[0].append(Node(data="n111")) - n1.children[0].append(Node(data="n112")) - n1.children[1].append(Node(data="n121")) - n1.children[1].append(Node(data="n122")) - n1.children[1].append(Node(data="n123")) - sess.add(n1) - sess.flush() - sess.expunge_all() + n1, n2 = data_fixture(sess) def go(): allnodes = sess.query(Node).order_by(Node.data).all() @@ -2785,12 +2820,9 @@ class SelfReferentialTest(fixtures.MappedTest): self.assert_sql_count(testing.db, go, 2) - def test_with_deferred(self): + def test_with_deferred(self, data_fixture): nodes = self.tables.nodes - - class Node(fixtures.ComparableEntity): - def append(self, node): - self.children.append(node) + Node = self.classes.Node self.mapper_registry.map_imperatively( Node, @@ -2803,39 +2835,55 @@ class SelfReferentialTest(fixtures.MappedTest): }, ) sess = fixture_session() - n1 = Node(data="n1") - n1.append(Node(data="n11")) - n1.append(Node(data="n12")) - sess.add(n1) - sess.flush() - sess.expunge_all() + n1, n2 = data_fixture(sess) def go(): eq_( - Node(data="n1", children=[Node(data="n11"), Node(data="n12")]), + Node( + data="n1", + children=[ + Node(data="n11"), + Node(data="n12"), + Node(data="n13"), + ], + ), sess.query(Node).order_by(Node.id).first(), ) - self.assert_sql_count(testing.db, go, 6) + self.assert_sql_count(testing.db, go, 8) sess.expunge_all() def go(): eq_( - Node(data="n1", children=[Node(data="n11"), Node(data="n12")]), + Node( + data="n1", + children=[ + Node(data="n11"), + Node(data="n12"), + Node(data="n13"), + ], + ), sess.query(Node) .options(undefer(Node.data)) .order_by(Node.id) .first(), ) - self.assert_sql_count(testing.db, go, 5) + self.assert_sql_count(testing.db, go, 7) sess.expunge_all() def go(): eq_( - Node(data="n1", children=[Node(data="n11"), Node(data="n12")]), + Node( + data="n1", + children=[ + Node(data="n11"), + Node(data="n12"), + Node(data="n13"), + ], + ), sess.query(Node) .options( undefer(Node.data), @@ -2844,14 +2892,11 @@ class SelfReferentialTest(fixtures.MappedTest): .first(), ) - self.assert_sql_count(testing.db, go, 3) + self.assert_sql_count(testing.db, go, 4) - def test_options(self): + def test_options(self, data_fixture): nodes = self.tables.nodes - - class Node(fixtures.ComparableEntity): - def append(self, node): - self.children.append(node) + Node = self.classes.Node self.mapper_registry.map_imperatively( Node, @@ -2859,16 +2904,7 @@ class SelfReferentialTest(fixtures.MappedTest): properties={"children": relationship(Node, order_by=nodes.c.id)}, ) sess = fixture_session() - n1 = Node(data="n1") - n1.append(Node(data="n11")) - n1.append(Node(data="n12")) - n1.append(Node(data="n13")) - n1.children[1].append(Node(data="n121")) - n1.children[1].append(Node(data="n122")) - n1.children[1].append(Node(data="n123")) - sess.add(n1) - sess.flush() - sess.expunge_all() + n1, n2 = data_fixture(sess) def go(): d = ( @@ -2901,14 +2937,11 @@ class SelfReferentialTest(fixtures.MappedTest): self.assert_sql_count(testing.db, go, 3) - def test_no_depth(self): + def test_no_depth(self, data_fixture): """no join depth is set, so no eager loading occurs.""" nodes = self.tables.nodes - - class Node(fixtures.ComparableEntity): - def append(self, node): - self.children.append(node) + Node = self.classes.Node self.mapper_registry.map_imperatively( Node, @@ -2916,19 +2949,7 @@ class SelfReferentialTest(fixtures.MappedTest): properties={"children": relationship(Node, lazy="selectin")}, ) sess = fixture_session() - n1 = Node(data="n1") - n1.append(Node(data="n11")) - n1.append(Node(data="n12")) - n1.append(Node(data="n13")) - n1.children[1].append(Node(data="n121")) - n1.children[1].append(Node(data="n122")) - n1.children[1].append(Node(data="n123")) - n2 = Node(data="n2") - n2.append(Node(data="n21")) - sess.add(n1) - sess.add(n2) - sess.flush() - sess.expunge_all() + n1, n2 = data_fixture(sess) def go(): d = ( @@ -2961,6 +2982,32 @@ class SelfReferentialTest(fixtures.MappedTest): self.assert_sql_count(testing.db, go, 4) + def test_auto_recurse_opt(self, data_fixture): + nodes = self.tables.nodes + Node = self.classes.Node + + self.mapper_registry.map_imperatively( + Node, + nodes, + properties={"children": relationship(Node)}, + ) + sess = fixture_session() + n1, n2 = data_fixture(sess) + + def go(): + return ( + sess.query(Node) + .filter(Node.data.in_(["n1", "n2"])) + .options(selectinload(Node.children, auto_recurse=True)) + .order_by(Node.data) + .all() + ) + + result = self.assert_sql_count(testing.db, go, 4) + sess.close() + + eq_(result, self._full_structure()) + class SelfRefInheritanceAliasedTest( fixtures.DeclarativeMappedTest, testing.AssertsCompiledSQL