From 8d8a8713432a88737c4400610eef11c5c8457b86 Mon Sep 17 00:00:00 2001
From: Keryn Knight <keryn@kerynknight.com>
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 ++++++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 62 insertions(+), 1 deletion(-)

(limited to 'django/db/models/query.py')

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
-- 
cgit v1.3