From f40a56dfe15c652ae9d5d5ab6293b5208220e3d1 Mon Sep 17 00:00:00 2001 From: Andrew Chen Wang <60190294+Andrew-Chen-Wang@users.noreply.github.com> Date: Wed, 29 Jul 2020 14:00:04 -0400 Subject: [PATCH] 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 | 45 +++++++++++++++++++++++++++++++++++++++- cachalot/monkey_patch.py | 6 +++++- cachalot/tests/read.py | 42 +++++++++++++++++++++++++++++++++++++ 3 files changed, 91 insertions(+), 2 deletions(-) diff --git a/cachalot/api.py b/cachalot/api.py index 18110d1..c64511c 100644 --- a/cachalot/api.py +++ b/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 diff --git a/cachalot/monkey_patch.py b/cachalot/monkey_patch.py index 8f14d3a..46e8c8f 100644 --- a/cachalot/monkey_patch.py +++ b/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): diff --git a/cachalot/tests/read.py b/cachalot/tests/read.py index f9fc03b..1b2a47c 100644 --- a/cachalot/tests/read.py +++ b/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):