feat(GAQ): Add Redis Sentinel Support for Global Async Queries (#29912)

Co-authored-by: Sivarajan Narayanan <narayanan_sivarajan@apple.com>

tests/integration_tests/tasks/async_queries_tests.py

@@ -20,8 +20,14 @@ from unittest import mock
 from uuid import uuid4
 
 import pytest
+import redis
 from celery.exceptions import SoftTimeLimitExceeded
+from parameterized import parameterized
 
+from superset.async_events.cache_backend import (
+    RedisCacheBackend,
+    RedisSentinelCacheBackend,
+)
 from superset.commands.chart.data.get_data_command import ChartDataCommand
 from superset.commands.chart.exceptions import ChartDataQueryFailedError
 from superset.exceptions import SupersetException
@@ -38,17 +44,29 @@ from tests.integration_tests.fixtures.tags import (
 from tests.integration_tests.test_app import app
 
 
+@pytest.mark.usefixtures(
+    "load_birth_names_data", "load_birth_names_dashboard_with_slices"
+)
 class TestAsyncQueries(SupersetTestCase):
-    @pytest.mark.usefixtures(
-        "load_birth_names_data", "load_birth_names_dashboard_with_slices"
+    @parameterized.expand(
+        [
+            ("RedisCacheBackend", mock.Mock(spec=RedisCacheBackend)),
+            ("RedisSentinelCacheBackend", mock.Mock(spec=RedisSentinelCacheBackend)),
+            ("redis.Redis", mock.Mock(spec=redis.Redis)),
+        ]
     )
-    @mock.patch.object(async_query_manager, "update_job")
     @mock.patch("superset.tasks.async_queries.set_form_data")
-    def test_load_chart_data_into_cache(self, mock_set_form_data, mock_update_job):
+    @mock.patch.object(async_query_manager, "update_job")
+    def test_load_chart_data_into_cache(
+        self, cache_type, cache_backend, mock_update_job, mock_set_form_data
+    ):
         from superset.tasks.async_queries import load_chart_data_into_cache
 
         app._got_first_request = False
+
+        async_query_manager.get_cache_backend = mock.Mock(return_value=cache_backend)
         async_query_manager.init_app(app)
+
         query_context = get_query_context("birth_names")
         user = security_manager.find_user("gamma")
         job_metadata = {
@@ -60,20 +78,33 @@ class TestAsyncQueries(SupersetTestCase):
         }
 
         load_chart_data_into_cache(job_metadata, query_context)
+
         mock_set_form_data.assert_called_once_with(query_context)
         mock_update_job.assert_called_once_with(
             job_metadata, "done", result_url=mock.ANY
         )
 
+    @parameterized.expand(
+        [
+            ("RedisCacheBackend", mock.Mock(spec=RedisCacheBackend)),
+            ("RedisSentinelCacheBackend", mock.Mock(spec=RedisSentinelCacheBackend)),
+            ("redis.Redis", mock.Mock(spec=redis.Redis)),
+        ]
+    )
     @mock.patch.object(
         ChartDataCommand, "run", side_effect=ChartDataQueryFailedError("Error: foo")
     )
     @mock.patch.object(async_query_manager, "update_job")
-    def test_load_chart_data_into_cache_error(self, mock_update_job, mock_run_command):
+    def test_load_chart_data_into_cache_error(
+        self, cache_type, cache_backend, mock_update_job, mock_run_command
+    ):
         from superset.tasks.async_queries import load_chart_data_into_cache
 
         app._got_first_request = False
+
+        async_query_manager.get_cache_backend = mock.Mock(return_value=cache_backend)
         async_query_manager.init_app(app)
+
         query_context = get_query_context("birth_names")
         user = security_manager.find_user("gamma")
         job_metadata = {
@@ -90,15 +121,25 @@ class TestAsyncQueries(SupersetTestCase):
         errors = [{"message": "Error: foo"}]
         mock_update_job.assert_called_once_with(job_metadata, "error", errors=errors)
 
+    @parameterized.expand(
+        [
+            ("RedisCacheBackend", mock.Mock(spec=RedisCacheBackend)),
+            ("RedisSentinelCacheBackend", mock.Mock(spec=RedisSentinelCacheBackend)),
+            ("redis.Redis", mock.Mock(spec=redis.Redis)),
+        ]
+    )
     @mock.patch.object(ChartDataCommand, "run")
     @mock.patch.object(async_query_manager, "update_job")
     def test_soft_timeout_load_chart_data_into_cache(
-        self, mock_update_job, mock_run_command
+        self, cache_type, cache_backend, mock_update_job, mock_run_command
     ):
         from superset.tasks.async_queries import load_chart_data_into_cache
 
         app._got_first_request = False
+
+        async_query_manager.get_cache_backend = mock.Mock(return_value=cache_backend)
         async_query_manager.init_app(app)
+
         user = security_manager.find_user("gamma")
         form_data = {}
         job_metadata = {
@@ -118,13 +159,25 @@ class TestAsyncQueries(SupersetTestCase):
                 load_chart_data_into_cache(job_metadata, form_data)
             set_form_data.assert_called_once_with(form_data, "error", errors=errors)
 
+    @parameterized.expand(
+        [
+            ("RedisCacheBackend", mock.Mock(spec=RedisCacheBackend)),
+            ("RedisSentinelCacheBackend", mock.Mock(spec=RedisSentinelCacheBackend)),
+            ("redis.Redis", mock.Mock(spec=redis.Redis)),
+        ]
+    )
     @pytest.mark.usefixtures("load_birth_names_dashboard_with_slices")
     @mock.patch.object(async_query_manager, "update_job")
-    def test_load_explore_json_into_cache(self, mock_update_job):
+    def test_load_explore_json_into_cache(
+        self, cache_type, cache_backend, mock_update_job
+    ):
         from superset.tasks.async_queries import load_explore_json_into_cache
 
         app._got_first_request = False
+
+        async_query_manager.get_cache_backend = mock.Mock(return_value=cache_backend)
         async_query_manager.init_app(app)
+
         table = self.get_table(name="birth_names")
         user = security_manager.find_user("gamma")
         form_data = {
@@ -146,19 +199,30 @@ class TestAsyncQueries(SupersetTestCase):
         }
 
         load_explore_json_into_cache(job_metadata, form_data)
+
         mock_update_job.assert_called_once_with(
             job_metadata, "done", result_url=mock.ANY
         )
 
+    @parameterized.expand(
+        [
+            ("RedisCacheBackend", mock.Mock(spec=RedisCacheBackend)),
+            ("RedisSentinelCacheBackend", mock.Mock(spec=RedisSentinelCacheBackend)),
+            ("redis.Redis", mock.Mock(spec=redis.Redis)),
+        ]
+    )
     @mock.patch.object(async_query_manager, "update_job")
     @mock.patch("superset.tasks.async_queries.set_form_data")
     def test_load_explore_json_into_cache_error(
-        self, mock_set_form_data, mock_update_job
+        self, cache_type, cache_backend, mock_set_form_data, mock_update_job
     ):
         from superset.tasks.async_queries import load_explore_json_into_cache
 
         app._got_first_request = False
+
+        async_query_manager.get_cache_backend = mock.Mock(return_value=cache_backend)
         async_query_manager.init_app(app)
+
         user = security_manager.find_user("gamma")
         form_data = {}
         job_metadata = {
@@ -176,15 +240,25 @@ class TestAsyncQueries(SupersetTestCase):
         errors = ["The dataset associated with this chart no longer exists"]
         mock_update_job.assert_called_once_with(job_metadata, "error", errors=errors)
 
+    @parameterized.expand(
+        [
+            ("RedisCacheBackend", mock.Mock(spec=RedisCacheBackend)),
+            ("RedisSentinelCacheBackend", mock.Mock(spec=RedisSentinelCacheBackend)),
+            ("redis.Redis", mock.Mock(spec=redis.Redis)),
+        ]
+    )
     @mock.patch.object(ChartDataCommand, "run")
     @mock.patch.object(async_query_manager, "update_job")
     def test_soft_timeout_load_explore_json_into_cache(
-        self, mock_update_job, mock_run_command
+        self, cache_type, cache_backend, mock_update_job, mock_run_command
     ):
         from superset.tasks.async_queries import load_explore_json_into_cache
 
         app._got_first_request = False
+
+        async_query_manager.get_cache_backend = mock.Mock(return_value=cache_backend)
         async_query_manager.init_app(app)
+
         user = security_manager.find_user("gamma")
         form_data = {}
         job_metadata = {
@@ -194,7 +268,7 @@ class TestAsyncQueries(SupersetTestCase):
             "status": "pending",
             "errors": [],
         }
-        errors = ["A timeout occurred while loading explore json, error"]
+        errors = ["A timeout occurred while loading explore JSON data"]
 
         with pytest.raises(SoftTimeLimitExceeded):
             with mock.patch(