From 8d8a8713432a88737c4400610eef11c5c8457b86 Mon Sep 17 00:00:00 2001 From: Keryn Knight Date: Tue, 20 Jul 2021 13:04:51 +0100 Subject: Refs #28455 -- Implemented private API methods for preventing QuerySet cloning. MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Multiple calls are idempotent assuming they're balanced. Also, multiple calls to disable cloning followed by a single call to re-enable cloning will subsequently cause clones to occur - it is not a stack, just a toggle. @contextlib.contextmanager is intentionally not used for performance reasons: - decorator takes 1.1µs to execute, or 2µs if used correctly in a `with ...:` statement - custom class takes 300ns to execute, or 900ns if used correctly in a `with ...:` statement Based on work originally done by Anssi Kääriäinen and Tim Graham. --- django/db/models/query.py | 63 ++++++++++++++++++++++++++++++++++++++++++++++- tests/queries/tests.py | 46 ++++++++++++++++++++++++++++++++++ 2 files changed, 108 insertions(+), 1 deletion(-) diff --git a/django/db/models/query.py b/django/db/models/query.py index cbe77caea9..d4775308b8 100644 --- a/django/db/models/query.py +++ b/django/db/models/query.py @@ -300,6 +300,28 @@ class FlatValuesListIterable(BaseIterable): yield row[0] +class PreventQuerySetCloning: + """ + Temporarily prevent the given QuerySet from creating new QuerySet instances + on each mutating operation (e.g: filter(), exclude() etc), instead + modifying the QuerySet in-place. + + @contextlib.contextmanager is intentionally not used for performance + reasons. + """ + + __slots__ = ("queryset",) + + def __init__(self, queryset): + self.queryset = queryset + + def __enter__(self): + return self.queryset._disable_cloning() + + def __exit__(self, exc_type, exc_value, traceback): + self.queryset._enable_cloning() + + class QuerySet(AltersData): """Represent a lazy database lookup for a set of objects.""" @@ -319,6 +341,7 @@ class QuerySet(AltersData): self._fields = None self._defer_next_filter = False self._deferred_filter = None + self._cloning_enabled = True @property def query(self): @@ -2134,12 +2157,50 @@ class QuerySet(AltersData): ) return inserted_rows + def _disable_cloning(self): + """ + Prevent calls to _chain() from creating a new QuerySet via _clone(). + All subsequent QuerySet mutations will occur on this instance until + _enable_cloning() is used. + """ + self._cloning_enabled = False + return self + + def _enable_cloning(self): + """ + Allow calls to _chain() to create a new QuerySet via _clone(). Restores + the default behavior where any QuerySet mutation will return a new + QuerySet instance. Necessary only when there has been a + _disable_cloning() call previously. + """ + self._cloning_enabled = True + return self + + def _avoid_cloning(self): + """ + Temporarily prevent QuerySet _clone() operations, restoring the default + behavior on exit. For the duration of the context managed statement, + all operations (e.g. filter(), exclude(), etc.) will mutate the same + QuerySet instance. + + @contextlib.contextmanager is intentionally not used for performance + reasons. + """ + return PreventQuerySetCloning(self) + def _chain(self): """ Return a copy of the current QuerySet that's ready for another operation. + + If the QuerySet has opted in to in-place mutations via + _disable_cloning() temporarily, the copy doesn't occur and instead the + same QuerySet instance will be modified. """ - obj = self._clone() + if not self._cloning_enabled: + obj = self + else: + obj = self._clone() if obj._sticky_filter: obj.query.filter_is_sticky = True obj._sticky_filter = False diff --git a/tests/queries/tests.py b/tests/queries/tests.py index 74929e4944..af657b2580 100644 --- a/tests/queries/tests.py +++ b/tests/queries/tests.py @@ -4629,3 +4629,49 @@ class Ticket23622Tests(TestCase): set(Ticket23605A.objects.filter(qy).values_list("pk", flat=True)), ) self.assertSequenceEqual(Ticket23605A.objects.filter(qx), [a2]) + + +class QuerySetCloningTests(TestCase): + @classmethod + def setUpTestData(cls): + SimpleCategory.objects.bulk_create( + [ + SimpleCategory(name="first"), + SimpleCategory(name="second"), + SimpleCategory(name="third"), + SimpleCategory(name="fourth"), + ] + ) + + def test_context_manager(self): + """ + _avoid_cloning() makes modifications apply to the original QuerySet. + """ + qs = SimpleCategory.objects.all() + with qs._avoid_cloning(): + qs2 = qs.filter(name__in={"first", "second"}).exclude(name="second") + self.assertIs(qs2, qs) + qs3 = qs2.exclude(name__in={"third", "fourth"}) + # qs3 is not a mutation of qs2 (which is actually also qs) but a new + # instance entirely. + self.assertIsNot(qs3, qs) + self.assertIsNot(qs3, qs2) + + def test_explicit_toggling(self): + qs = SimpleCategory.objects.filter(name__in={"first", "second"}) + qs2 = qs._disable_cloning() + # The _disable_cloning() method doesn't return a new QuerySet, but + # toggles the value on the current instance. qs2 can be ignored. + self.assertIs(qs2, qs) + qs3 = qs.filter(name__in={"first", "second"}) + qs3 = qs3.exclude(name="second") + qs3._enable_cloning() + # These are still both references to the same QuerySet, despite + # re-binding as if they were normal chained operations providing new + # QuerySet instances. + self.assertIs(qs3, qs) + qs3 = qs3.filter(name="second") + # Cloning has been re-enabled so subsequent operations yield a new + # QuerySet. qs3 is now all of the filters applied to qs + an additional + # filter. + self.assertIsNot(qs3, qs) -- cgit v1.3