Add support for disabling cachalot (#158)

* This context manager allows for disabling cachalot using a context manager
* Allow for disabling all queries within context manager

cachalot/api.py

@@ -1,3 +1,5 @@
+from contextlib import contextmanager
+
 from django.apps import apps
 from django.conf import settings
 from django.db import connections
@@ -9,7 +11,15 @@ from .transaction import AtomicCache
 from .utils import _invalidate_tables
 
 
-__all__ = ('invalidate', 'get_last_invalidation')
+try:
+    from asgiref.local import Local
+    LOCAL_STORAGE = Local()
+except ImportError:
+    import threading
+    LOCAL_STORAGE = threading.local()
+
+
+__all__ = ('invalidate', 'get_last_invalidation', 'cachalot_disabled')
 
 
 def _cache_db_tables_iterator(tables, cache_alias, db_alias):
@@ -121,3 +131,36 @@ def get_last_invalidation(*tables_or_models, **kwargs):
             if current_last_invalidation > last_invalidation:
                 last_invalidation = current_last_invalidation
     return last_invalidation
+
+
+@contextmanager
+def cachalot_disabled(all_queries=False):
+    """
+    Context manager for temporarily disabling cachalot.
+    If you evaluate the same queryset a second time,
+    like normally for Django querysets, this will access
+    the variable that saved it in-memory.
+
+    For example:
+    with cachalot_disabled():
+        qs = Test.objects.filter(blah=blah)
+        # Does a single query to the db
+        list(qs)  # Evaluates queryset
+        # Because the qs was evaluated, it's
+        # saved in memory:
+        list(qs)  # this does 0 queries.
+        # This does 1 query to the db
+        list(Test.objects.filter(blah=blah))
+
+    If you evaluate the queryset outside the context manager, any duplicate
+    query will use the cached result unless an object creation happens in between
+    the original and duplicate query.
+
+    :arg all_queries: Any query, including already evaluated queries, are re-evaluated.
+    :type all_queries: bool
+    """
+    was_enabled = getattr(LOCAL_STORAGE, "cachalot_enabled", cachalot_settings.CACHALOT_ENABLED)
+    LOCAL_STORAGE.enabled = False
+    LOCAL_STORAGE.disable_on_all = all_queries
+    yield
+    LOCAL_STORAGE.enabled = was_enabled

cachalot/monkey_patch.py

@@ -10,7 +10,7 @@ from django.db.models.sql.compiler import (
 )
 from django.db.transaction import Atomic, get_connection
 
-from .api import invalidate
+from .api import invalidate, LOCAL_STORAGE
 from .cache import cachalot_caches
 from .settings import cachalot_settings, ITERABLES
 from .utils import (
@@ -66,6 +66,10 @@ def _patch_compiler(original):
     @_unset_raw_connection
     def inner(compiler, *args, **kwargs):
         execute_query_func = lambda: original(compiler, *args, **kwargs)
+        # Checks if utils/cachalot_disabled
+        if not getattr(LOCAL_STORAGE, "cachalot_enabled", True):
+            return execute_query_func()
+
         db_alias = compiler.using
         if db_alias not in cachalot_settings.CACHALOT_DATABASES \
                 or isinstance(compiler, WRITE_COMPILERS):

cachalot/tests/read.py

@@ -18,6 +18,7 @@ from django.test import (
 from pytz import UTC
 
 from cachalot.cache import cachalot_caches
+from ..api import cachalot_disabled
 from ..settings import cachalot_settings
 from ..utils import UncachableQuery
 from .models import Test, TestChild, TestParent, UnmanagedModel
@@ -862,6 +863,47 @@ class ReadTestCase(TestUtilsMixin, TransactionTestCase):
         self.assert_tables(qs, UnmanagedModel)
         self.assert_query_cached(qs)
 
+    def test_cachalot_disabled_multiple_queries_ignoring_in_mem_cache(self):
+        """
+        Test that when queries are given the `cachalot_disabled` context manager,
+        the queries will not be cached.
+        """
+        with cachalot_disabled(True):
+            qs = Test.objects.all()
+            with self.assertNumQueries(1):
+                data1 = list(qs.all())
+            Test.objects.create(
+                name='test3', owner=self.user,
+                date='1789-07-14', datetime='1789-07-14T16:43:27',
+                permission=self.t1__permission)
+            with self.assertNumQueries(1):
+                data2 = list(qs.all())
+            self.assertNotEqual(data1, data2)
+
+    def test_query_cachalot_disabled_even_if_already_cached(self):
+        """
+        Test that when a query is given the `cachalot_disabled` context manager,
+        the query outside of the context manager will be cached. Any duplicated
+        query will use the original query's cached result.
+        """
+        qs = Test.objects.all()
+        self.assert_query_cached(qs)
+        with cachalot_disabled() and self.assertNumQueries(0):
+            list(qs.all())
+
+    def test_duplicate_query_execute_anyways(self):
+        """After an object is created, a duplicate query should execute
+        rather than use the cached result.
+        """
+        qs = Test.objects.all()
+        self.assert_query_cached(qs)
+        Test.objects.create(
+            name='test3', owner=self.user,
+            date='1789-07-14', datetime='1789-07-14T16:43:27',
+            permission=self.t1__permission)
+        with cachalot_disabled() and self.assertNumQueries(1):
+            list(qs.all())
+
 
 class ParameterTypeTestCase(TestUtilsMixin, TransactionTestCase):
     def test_tuple(self):