Adds the ability to replace/extend caching backend (#7856)

* Add ability to override cache backend with custom module

* Tests for setup_cache

* Add custom CACHE_CONFIG documentation

* Fix linter errors and documentation link

* Fix black formatting errors

docs/installation.rst

@@ -521,6 +521,24 @@ into your global default defined in ``CACHE_CONFIG``.
         'CACHE_REDIS_URL': 'redis://localhost:6379/0',
     }
 
+It is also possible to pass a custom cache initialization function in the
+config to handle additional caching use cases. The function must return an
+object that is compatible with the `Flask-Cache <https://pythonhosted.org/Flask-Cache/>`_ API.
+
+.. code-block:: python
+
+    from custom_caching import CustomCache
+
+    def init_cache(app):
+        """Takes an app instance and returns a custom cache backend"""
+        config = {
+            'CACHE_DEFAULT_TIMEOUT': 60 * 60 * 24, # 1 day default (in secs)
+            'CACHE_KEY_PREFIX': 'superset_results',
+        }
+        return CustomCache(app, config)
+
+    CACHE_CONFIG = init_cache
+
 Superset has a Celery task that will periodically warm up the cache based on
 different strategies. To use it, add the following to the `CELERYBEAT_SCHEDULE`
 section in `config.py`:

superset/utils/core.py

@@ -775,8 +775,14 @@ def choicify(values):
 
 def setup_cache(app: Flask, cache_config) -> Optional[Cache]:
     """Setup the flask-cache on a flask app"""
-    if cache_config and cache_config.get("CACHE_TYPE") != "null":
-        return Cache(app, config=cache_config)
+    if cache_config:
+        if isinstance(cache_config, dict):
+            if cache_config.get("CACHE_TYPE") != "null":
+                return Cache(app, config=cache_config)
+        else:
+            # Accepts a custom cache initialization function,
+            # returning an object compatible with Flask-Caching API
+            return cache_config(app)
 
     return None
 

tests/utils_tests.py

@@ -20,6 +20,8 @@ import unittest
 from unittest.mock import patch
 import uuid
 
+from flask import Flask
+from flask_caching import Cache
 import numpy
 
 from superset import app
@@ -39,6 +41,7 @@ from superset.utils.core import (
     parse_human_timedelta,
     parse_js_uri_path_item,
     parse_past_timedelta,
+    setup_cache,
     validate_json,
     zlib_compress,
     zlib_decompress_to_string,
@@ -766,6 +769,35 @@ class UtilsTestCase(unittest.TestCase):
         self.assertIsNone(parse_js_uri_path_item(None))
         self.assertIsNotNone(parse_js_uri_path_item("item"))
 
+    def test_setup_cache_no_config(self):
+        app = Flask(__name__)
+        cache_config = None
+        self.assertIsNone(setup_cache(app, cache_config))
+
+    def test_setup_cache_null_config(self):
+        app = Flask(__name__)
+        cache_config = {"CACHE_TYPE": "null"}
+        self.assertIsNone(setup_cache(app, cache_config))
+
+    def test_setup_cache_standard_config(self):
+        app = Flask(__name__)
+        cache_config = {
+            "CACHE_TYPE": "redis",
+            "CACHE_DEFAULT_TIMEOUT": 60,
+            "CACHE_KEY_PREFIX": "superset_results",
+            "CACHE_REDIS_URL": "redis://localhost:6379/0",
+        }
+        assert isinstance(setup_cache(app, cache_config), Cache) is True
+
+    def test_setup_cache_custom_function(self):
+        app = Flask(__name__)
+        CustomCache = type("CustomCache", (object,), {"__init__": lambda *args: None})
+
+        def init_cache(app):
+            return CustomCache(app, {})
+
+        assert isinstance(setup_cache(app, init_cache), CustomCache) is True
+
     def test_get_stacktrace(self):
         with app.app_context():
             app.config["SHOW_STACKTRACE"] = True