Add router to read and write in the cache

Author: WaVEV

django_mongodb_backend/cache.py

@@ -5,6 +5,7 @@
 from django.core.cache.backends.db import Options
 from django.db import connections, router
 from django.utils.functional import cached_property
+from pymongo import IndexModel
 from pymongo.errors import DuplicateKeyError
 
 
@@ -41,20 +42,29 @@ class CacheEntry:
         self.cache_model_class = CacheEntry
 
     def create_indexes(self):
-        self.collection.create_index("expires_at", expireAfterSeconds=0)
-        self.collection.create_index("key", unique=True)
+        expires_index = IndexModel("expires_at", expireAfterSeconds=0)
+        key_index = IndexModel("key", unique=True)
+        self.collection_to_write.create_indexes([expires_index, key_index])
 
     @cached_property
     def serializer(self):
         return MongoSerializer(self.pickle_protocol)
 
     @property
-    def _db(self):
+    def _db_to_read(self):
         return connections[router.db_for_read(self.cache_model_class)]
 
     @property
-    def collection(self):
-        return self._db.get_collection(self._collection_name)
+    def collection_to_read(self):
+        return self._db_to_read.get_collection(self._collection_name)
+
+    @property
+    def _db_to_write(self):
+        return connections[router.db_for_write(self.cache_model_class)]
+
+    @property
+    def collection_to_write(self):
+        return self._db_to_read.get_collection(self._collection_name)
 
     def get(self, key, default=None, version=None):
         result = self.get_many([key], version)
@@ -71,18 +81,18 @@ def get_many(self, keys, version=None):
         if not keys:
             return {}
         keys_map = {self.make_and_validate_key(key, version=version): key for key in keys}
-        with self.collection.find(
+        with self.collection_to_read.find(
             {"key": {"$in": tuple(keys_map)}, **self._filter_expired(expired=False)}
         ) as cursor:
             return {keys_map[row["key"]]: self.serializer.loads(row["value"]) for row in cursor}
 
     def set(self, key, value, timeout=DEFAULT_TIMEOUT, version=None):
         key = self.make_and_validate_key(key, version=version)
         serialized_data = self.serializer.dumps(value)
-        num = self.collection.count_documents({})
+        num = self.collection_to_write.count_documents({})
         if num >= self._max_entries:
             self._cull(num)
-        return self.collection.update_one(
+        return self.collection_to_write.update_one(
             {"key": key},
             {
                 "$set": {
@@ -97,11 +107,11 @@ def set(self, key, value, timeout=DEFAULT_TIMEOUT, version=None):
     def add(self, key, value, timeout=DEFAULT_TIMEOUT, version=None):
         key = self.make_and_validate_key(key, version=version)
         serialized_data = self.serializer.dumps(value)
-        num = self.collection.count_documents({})
+        num = self.collection_to_write.count_documents({})
         if num >= self._max_entries:
             self._cull(num)
         try:
-            self.collection.update_one(
+            self.collection_to_write.update_one(
                 {"key": key, **self._filter_expired(expired=True)},
                 {
                     "$set": {
@@ -124,7 +134,7 @@ def _cull(self, num):
             try:
                 # Delete the first expiration date.
                 deleted_from = next(
-                    self.collection.aggregate(
+                    self.collection_to_write.aggregate(
                         [
                             {"$sort": {"expires_at": -1, "key": 1}},
                             {"$skip": keep_num},
@@ -136,7 +146,7 @@ def _cull(self, num):
             except StopIteration:
                 pass
             else:
-                self.collection.delete_many(
+                self.collection_to_write.delete_many(
                     {
                         "$or": [
                             {"expires_at": {"$lt": deleted_from["expires_at"]}},
@@ -152,7 +162,7 @@ def _cull(self, num):
 
     def touch(self, key, timeout=DEFAULT_TIMEOUT, version=None):
         key = self.make_and_validate_key(key, version=version)
-        res = self.collection.update_one(
+        res = self.collection_to_write.update_one(
             {"key": key}, {"$set": {"expires_at": self._get_expiration_time(timeout)}}
         )
         return res.matched_count > 0
@@ -173,13 +183,16 @@ def _delete_many(self, keys, version=None):
         if not keys:
             return False
         keys = tuple(self.make_and_validate_key(key, version=version) for key in keys)
-        return bool(self.collection.delete_many({"key": {"$in": keys}}).deleted_count)
+        return bool(self.collection_to_write.delete_many({"key": {"$in": keys}}).deleted_count)
 
     def has_key(self, key, version=None):
         key = self.make_and_validate_key(key, version=version)
         return (
-            self.collection.count_documents({"key": key, **self._filter_expired(expired=False)}) > 0
+            self.collection_to_read.count_documents(
+                {"key": key, **self._filter_expired(expired=False)}
+            )
+            > 0
         )
 
     def clear(self):
-        self.collection.delete_many({})
+        self.collection_to_write.delete_many({})

django_mongodb_backend/creation.py

@@ -25,9 +25,9 @@ def create_test_db(self, *args, **kwargs):
         # Create cache collections
         for cache_alias in settings.CACHES:
             cache = caches[cache_alias]
-            connection = cache._db
-            if cache._collection_name in connection.introspection.table_names():
-                continue
-            cache = MongoDBCache(cache._collection_name, {})
-            cache.create_indexes()
+            if isinstance(cache, MongoDBCache):
+                connection = cache._db_to_write
+                if cache._collection_name in connection.introspection.table_names():
+                    continue
+                cache.create_indexes()
         return test_database_name

django_mongodb_backend/management/commands/createcachecollection.py

@@ -43,15 +43,14 @@ def handle(self, *collection_names, **options):
             for cache_alias in settings.CACHES:
                 cache = caches[cache_alias]
                 if isinstance(cache, MongoDBCache):
-                    self.check_collection(db, cache._collection_name)
+                    self.check_collection(db, cache)
 
-    def check_collection(self, database, collection_name):
-        cache = MongoDBCache(collection_name, {})
+    def check_collection(self, database, cache):
         if not router.allow_migrate_model(database, cache.cache_model_class):
             return
         connection = connections[database]
-        if collection_name in connection.introspection.table_names():
+        if cache._collection_name in connection.introspection.table_names():
             if self.verbosity > 0:
-                self.stdout.write("Cache collection '%s' already exists." % collection_name)
+                self.stdout.write("Cache collection '%s' already exists." % cache._collection_name)
             return
         cache.create_indexes()

tests/cache_/tests.py

@@ -929,7 +929,7 @@ def test_get_or_set_racing(self):
             self.assertEqual(cache.get_or_set("key", "default"), "default")
 
     def test_collection_has_indexes(self):
-        indexes = list(cache.collection.list_indexes())
+        indexes = list(cache.collection_to_read.list_indexes())
         self.assertTrue(
             any(
                 index["key"] == SON([("expires_at", 1)]) and index.get("expireAfterSeconds") == 0
@@ -962,7 +962,7 @@ def setUp(self):
         self.addCleanup(self.drop_collection)
 
     def drop_collection(self):
-        cache.collection.drop()
+        cache.collection_to_write.drop()
 
     def create_cache_collection(self):
         management.call_command("createcachecollection", verbosity=0)
@@ -971,3 +971,49 @@ def create_cache_collection(self):
 @override_settings(USE_TZ=True)
 class DBCacheWithTimeZoneTests(DBCacheTests):
     pass
+
+
+class DBCacheRouter:
+    """A router that puts the cache table on the 'other' database."""
+
+    def db_for_read(self, model, **hints):
+        if model._meta.app_label == "django_cache":
+            return "other"
+        return None
+
+    def db_for_write(self, model, **hints):
+        if model._meta.app_label == "django_cache":
+            return "other"
+        return None
+
+    def allow_migrate(self, db, app_label, **hints):
+        if app_label == "django_cache":
+            return db == "other"
+        return None
+
+
+@override_settings(
+    CACHES={
+        "default": {
+            "BACKEND": "django_mongodb_backend.cache.MongoDBCache",
+            "LOCATION": "my_cache_table",
+        },
+    },
+)
+@modify_settings(
+    INSTALLED_APPS={"prepend": "django_mongodb_backend"},
+)
+class CreateCacheTableForDBCacheTests(TestCase):
+    databases = {"default", "other"}
+
+    @override_settings(DATABASE_ROUTERS=[DBCacheRouter()])
+    def test_createcachetable_observes_database_router(self):
+        # cache table should not be created on 'default'
+        with self.assertNumQueries(0, using="default"):
+            management.call_command("createcachecollection", database="default", verbosity=0)
+        # cache table should be created on 'other'
+        # Queries:
+        #   1: Create indexes
+        num = 1
+        with self.assertNumQueries(num, using="other"):
+            management.call_command("createcachecollection", database="other", verbosity=0)