Repository: adamchainz/django-mysql

Issue:
MySQLCache - prefix operations
Could extend cache further with e.g. `get_with_prefix` -> `SELECT ... WHERE cache_key LIKE 'x_%'`



Discussion:

Resolving patch:
diff --git a/django_mysql/cache.py b/django_mysql/cache.py
--- a/django_mysql/cache.py
+++ b/django_mysql/cache.py
@@ -1,16 +1,20 @@
 # -*- coding:utf-8 -*-
 from __future__ import unicode_literals
 
+import re
 import zlib
 from random import random
 from textwrap import dedent
 from time import time
 
 import django
-from django.core.cache.backends.base import DEFAULT_TIMEOUT, BaseCache
+from django.core.cache.backends.base import (
+    DEFAULT_TIMEOUT, BaseCache, default_key_func
+)
 from django.db import connections, router
 from django.utils import six
 from django.utils.encoding import force_bytes
+from django.utils.module_loading import import_string
 
 from django_mysql.utils import collapse_spaces
 
@@ -59,6 +63,36 @@ class CacheEntry(object):
         self.cache_model_class = CacheEntry
 
 
+reverse_key_re = re.compile(r'^([^:]*):(\d+):(.*)')
+
+
+def default_reverse_key_func(full_key):
+    """
+    Reverse of Django's default_key_func, i.e. undoing:
+
+        def default_key_func(key, key_prefix, version):
+            return '%s:%s:%s' % (key_prefix, version, key)
+    """
+    match = reverse_key_re.match(full_key)
+    return match.group(3), match.group(1), int(match.group(2))
+
+
+def get_reverse_key_func(reverse_key_func):
+    """
+    Function to decide which reverse key function to use
+
+    Defaults to ``None``, as any other value might not apply to the given
+    KEY_FUNCTION. Also the user may not use any of the operations that require
+    reversing the key_func.
+    """
+    if reverse_key_func is not None:
+        if callable(reverse_key_func):
+            return reverse_key_func
+        else:
+            return import_string(reverse_key_func)
+    return None
+
+
 class MySQLCache(BaseDatabaseCache):
 
     # Got an error with the add() query using BIGINT_UNSIGNED_MAX, so use a
@@ -89,6 +123,19 @@ def __init__(self, table, params):
         self._compress_level = options.get('COMPRESS_LEVEL', 6)
         self._cull_probability = options.get('CULL_PROBABILITY', 0.01)
 
+        # Figure out our *reverse* key function
+        if self.key_func is default_key_func:
+            self.reverse_key_func = default_reverse_key_func
+            if ':' in self.key_prefix:
+                raise ValueError(
+                    "Cannot use the default KEY_FUNCTION and "
+                    "REVERSE_KEY_FUNCTION if you have a colon in your "
+                    "KEY_PREFIX."
+                )
+        else:
+            reverse_key_func = params.get('REVERSE_KEY_FUNCTION', None)
+            self.reverse_key_func = get_reverse_key_func(reverse_key_func)
+
     # Django API + helpers
 
     def get(self, key, default=None, version=None):
@@ -408,6 +455,88 @@ def get_backend_timeout(self, timeout=DEFAULT_TIMEOUT):
 
     # Our API extensions
 
+    def keys_with_prefix(self, prefix, version=None):
+        if self.reverse_key_func is None:
+            raise ValueError(
+                "To use the _with_prefix commands with a custom KEY_FUNCTION, "
+                "you need to specify a custom REVERSE_KEY_FUNCTION too."
+            )
+
+        if version is None:
+            version = self.version
+
+        db = router.db_for_read(self.cache_model_class)
+        table = connections[db].ops.quote_name(self._table)
+
+        prefix = self.make_key(prefix + '%', version=version)
+
+        with connections[db].cursor() as cursor:
+            cursor.execute(
+                """SELECT cache_key FROM {table}
+                   WHERE cache_key LIKE %s AND
+                         expires >= %s""".format(table=table),
+                (prefix, self._now())
+            )
+            rows = cursor.fetchall()
+            full_keys = {row[0] for row in rows}
+
+            keys = {}
+            for full_key in full_keys:
+                key, key_prefix, key_version = self.reverse_key_func(full_key)
+
+                if key_version == version:
+                    keys[key] = key_version
+            return set(six.iterkeys(keys))
+
+    def get_with_prefix(self, prefix, version=None):
+        if self.reverse_key_func is None:
+            raise ValueError(
+                "To use the _with_prefix commands with a custom KEY_FUNCTION, "
+                "you need to specify a custom REVERSE_KEY_FUNCTION too."
+            )
+
+        if version is None:
+            version = self.version
+
+        db = router.db_for_read(self.cache_model_class)
+        table = connections[db].ops.quote_name(self._table)
+
+        prefix = self.make_key(prefix + '%', version=version)
+        version = six.text_type(version)
+
+        with connections[db].cursor() as cursor:
+            cursor.execute(
+                """SELECT cache_key, value, value_type
+                   FROM {table}
+                   WHERE cache_key LIKE %s AND
+                         expires >= %s""".format(table=table),
+                (prefix, self._now())
+            )
+            rows = cursor.fetchall()
+
+            data = {}
+            for made_key, value, value_type in rows:
+                key, key_prefix, key_version = self.reverse_key_func(made_key)
+                data[key] = self.decode(value, value_type)
+
+            return data
+
+    def delete_with_prefix(self, prefix, version=None):
+        if version is None:
+            version = self.version
+
+        db = router.db_for_write(self.cache_model_class)
+        table = connections[db].ops.quote_name(self._table)
+
+        prefix = self.make_key(prefix + '%', version=version)
+
+        with connections[db].cursor() as cursor:
+            return cursor.execute(
+                """DELETE FROM {table}
+                   WHERE cache_key LIKE %s""".format(table=table),
+                (prefix,)
+            )
+
     def cull(self):
         db = router.db_for_write(self.cache_model_class)
         table = connections[db].ops.quote_name(self._table)


Test patch:
diff --git a/tests/django_mysql_tests/test_cache.py b/tests/django_mysql_tests/test_cache.py
--- a/tests/django_mysql_tests/test_cache.py
+++ b/tests/django_mysql_tests/test_cache.py
@@ -7,6 +7,7 @@
 import warnings
 from decimal import Decimal
 
+import ddt
 import pytest
 from django.core.cache import CacheKeyWarning, cache, caches
 from django.core.management import CommandError, call_command
@@ -57,13 +58,33 @@ def custom_key_func(key, key_prefix, version):
     return 'CUSTOM-' + '-'.join([key_prefix, str(version), key])
 
 
+def reverse_custom_key_func(full_key):
+    "The reverse of custom_key_func"
+    # Remove CUSTOM-
+    full_key = full_key[len('CUSTOM-'):]
+
+    first_dash = full_key.find('-')
+    key_prefix = full_key[:first_dash]
+
+    second_dash = full_key.find('-', first_dash + 1)
+    version = int(full_key[first_dash + 1:second_dash])
+
+    key = full_key[second_dash + 1:]
+
+    return key, key_prefix, version
+
+
 _caches_setting_base = {
     'default': {},
     'prefix': {'KEY_PREFIX': 'cacheprefix{}'.format(os.getpid())},
     'v2': {'VERSION': 2},
-    'custom_key': {'KEY_FUNCTION': custom_key_func},
+    'custom_key': {'KEY_FUNCTION': custom_key_func,
+                   'REVERSE_KEY_FUNCTION': reverse_custom_key_func},
     'custom_key2': {
-        'KEY_FUNCTION': 'django_mysql_tests.test_cache.custom_key_func'
+        'KEY_FUNCTION':
+            'django_mysql_tests.test_cache.custom_key_func',
+        'REVERSE_KEY_FUNCTION':
+            'django_mysql_tests.test_cache.reverse_custom_key_func',
     },
     'cull': {'OPTIONS': {'CULL_PROBABILITY': 1,
                          'MAX_ENTRIES': 30}},
@@ -105,6 +126,7 @@ def override_cache_settings(BACKEND='django_mysql.cache.MySQLCache',
 
 
 @override_cache_settings()
+@ddt.ddt
 class MySQLCacheTests(TransactionTestCase):
 
     def setUp(self):
@@ -1018,6 +1040,156 @@ def test_value_type_case_sensitivity(self):
             n = cursor.fetchone()[0]
             assert n == 0
 
+    def test_bad_key_prefix_for_reverse_function(self):
+        override = override_cache_settings(KEY_PREFIX='a:bad:prefix')
+        with override, pytest.raises(ValueError) as excinfo:
+            caches['default']
+        assert str(excinfo.value).startswith(
+            "Cannot use the default KEY_FUNCTION")
+
+    @ddt.data('default', 'prefix', 'custom_key', 'custom_key2')
+    def test_keys_with_prefix(self, cache_name):
+        cache = caches[cache_name]
+        assert cache.keys_with_prefix('') == set()
+        assert cache.keys_with_prefix('K') == set()
+
+        cache.set('A2', True)
+        cache.set('K1', True)
+        cache.set('K23', True, 1000)
+        cache.set('K99', True, 0.1)
+        time.sleep(0.2)
+        assert cache.keys_with_prefix('') == {'A2', 'K1', 'K23'}
+        assert cache.keys_with_prefix('K') == {'K1', 'K23'}
+
+        cache.delete('K1')
+        assert cache.keys_with_prefix('K') == {'K23'}
+
+        cache.clear()
+        assert cache.keys_with_prefix('') == set()
+        assert cache.keys_with_prefix('K') == set()
+
+    @ddt.data('default', 'prefix', 'custom_key', 'custom_key2')
+    def test_keys_with_prefix_version(self, cache_name):
+        cache = caches[cache_name]
+
+        cache.set('V12', True, version=1)
+        cache.set('V12', True, version=2)
+        cache.set('V2', True, version=2)
+        cache.set('V3', True, version=3)
+        assert cache.keys_with_prefix('V', version=1) == {'V12'}
+        assert cache.keys_with_prefix('V', version=2) == {'V12', 'V2'}
+        assert cache.keys_with_prefix('V', version=3) == {'V3'}
+
+    @override_cache_settings(KEY_FUNCTION=custom_key_func)
+    def test_keys_with_prefix_with_bad_cache(self):
+        with pytest.raises(ValueError) as excinfo:
+            cache.keys_with_prefix('')
+        assert str(excinfo.value).startswith(
+            "To use the _with_prefix commands")
+
+    @ddt.data('default', 'prefix', 'custom_key', 'custom_key2')
+    def test_get_with_prefix(self, cache_name):
+        cache = caches[cache_name]
+        assert cache.get_with_prefix('') == {}
+        assert cache.get_with_prefix('K') == {}
+
+        cache.set('A2', [True])
+        cache.set('K1', "Value1")
+        cache.set('K23', 2, 1000)
+        cache.set('K99', ["Value", 99], 0.1)
+        time.sleep(0.2)
+        assert (
+            cache.get_with_prefix('') ==
+            {'A2': [True], 'K1': "Value1", 'K23': 2}
+        )
+        assert (
+            cache.get_with_prefix('K') ==
+            {'K1': "Value1", 'K23': 2}
+        )
+
+        cache.delete('K1')
+        assert cache.get_with_prefix('K') == {'K23': 2}
+
+        cache.clear()
+        assert cache.get_with_prefix('') == {}
+        assert cache.get_with_prefix('K') == {}
+
+    @ddt.data('default', 'prefix', 'custom_key', 'custom_key2')
+    def test_get_with_prefix_version(self, cache_name):
+        cache = caches[cache_name]
+
+        cache.set('V12', ('version1',), version=1)
+        cache.set('V12', "str", version=2)
+        cache.set('V2', 2, version=2)
+        cache.set('V3', object, version=3)
+        assert cache.get_with_prefix('V', version=1) == {'V12': ('version1',)}
+        assert cache.get_with_prefix('V', version=2) == {'V12': "str", 'V2': 2}
+        assert cache.get_with_prefix('V', version=3) == {'V3': object}
+
+    @override_cache_settings(KEY_FUNCTION=custom_key_func)
+    def test_get_with_prefix_with_bad_cache(self):
+        with pytest.raises(ValueError) as excinfo:
+            cache.get_with_prefix('')
+        assert str(excinfo.value).startswith(
+            "To use the _with_prefix commands")
+
+    @ddt.data('default', 'prefix', 'custom_key', 'custom_key2')
+    def test_delete_with_prefix(self, cache_name):
+        cache = caches[cache_name]
+
+        # Check it runs on an empty cache
+        assert cache.delete_with_prefix('') == 0
+        assert cache.delete_with_prefix('K') == 0
+
+        cache.set('A1', True)
+        cache.set('A2', True)
+        cache.set('K2', True)
+        cache.set('K44', True)
+
+        assert cache.keys_with_prefix('') == {'A1', 'A2', 'K2', 'K44'}
+        assert cache.delete_with_prefix('A') == 2
+        assert cache.keys_with_prefix('') == {'K2', 'K44'}
+        assert cache.delete_with_prefix('A') == 0
+        assert cache.keys_with_prefix('') == {'K2', 'K44'}
+        assert cache.delete_with_prefix('K') == 2
+        assert cache.keys_with_prefix('K') == set()
+        assert cache.keys_with_prefix('') == set()
+
+    @ddt.data('default', 'prefix', 'custom_key', 'custom_key2')
+    def test_delete_with_prefix_version(self, cache_name):
+        cache = caches[cache_name]
+
+        cache.set('V12', True, version=1)
+        cache.set('V12', True, version=2)
+        cache.set('V2', True, version=2)
+        cache.set('V3', True, version=3)
+
+        has_key = cache.has_key  # avoid lint error
+
+        assert cache.delete_with_prefix('V', version=1) == 1
+        assert not has_key('V12', version=1)
+        assert has_key('V12', version=2)
+        assert has_key('V2', version=2)
+        assert has_key('V3', version=3)
+
+        assert cache.delete_with_prefix('V', version=2) == 2
+        assert not has_key('V12', version=1)
+        assert not has_key('V12', version=2)
+        assert not has_key('V2', version=2)
+        assert has_key('V3', version=3)
+
+        assert cache.delete_with_prefix('V', version=3) == 1
+        assert not has_key('V12', version=1)
+        assert not has_key('V12', version=2)
+        assert not has_key('V2', version=2)
+        assert not has_key('V3', version=3)
+
+    @override_cache_settings(KEY_FUNCTION=custom_key_func)
+    def test_delete_with_prefix_with_no_reverse_works(self):
+        cache.set_many({'K1': 'value', 'K2': 'value2', 'B2': 'Anothervalue'})
+        assert cache.delete_with_prefix('K') == 2
+        assert cache.get_many(['K1', 'K2', 'B2']) == {'B2': 'Anothervalue'}
+
     # mysql_cache_migration tests
 
     def test_mysql_cache_migration(self):