fix(embedded): Retry when executing alert queries to avoid sending transient errors to users as alert failure notifications (#20419)

Co-authored-by: Rui Zhao <zhaorui@dropbox.com>

tests/integration_tests/reports/alert_tests.py

@@ -0,0 +1,119 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=invalid-name, unused-argument, import-outside-toplevel
+
+import pandas as pd
+from pytest_mock import MockFixture
+
+
+def test_execute_query_succeeded_no_retry(
+    mocker: MockFixture, app_context: None
+) -> None:
+
+    from superset.reports.commands.alert import AlertCommand
+
+    execute_query_mock = mocker.patch(
+        "superset.reports.commands.alert.AlertCommand._execute_query",
+        side_effect=lambda: pd.DataFrame([{"sample_col": 0}]),
+    )
+
+    command = AlertCommand(report_schedule=mocker.Mock())
+
+    command.validate()
+
+    assert execute_query_mock.call_count == 1
+
+
+def test_execute_query_succeeded_with_retries(
+    mocker: MockFixture, app_context: None
+) -> None:
+    from superset.reports.commands.alert import AlertCommand, AlertQueryError
+
+    execute_query_mock = mocker.patch(
+        "superset.reports.commands.alert.AlertCommand._execute_query"
+    )
+
+    query_executed_count = 0
+    # Should match the value defined in superset_test_config.py
+    expected_max_retries = 3
+
+    def _mocked_execute_query() -> pd.DataFrame:
+        nonlocal query_executed_count
+        query_executed_count += 1
+
+        if query_executed_count < expected_max_retries:
+            raise AlertQueryError()
+        else:
+            return pd.DataFrame([{"sample_col": 0}])
+
+    execute_query_mock.side_effect = _mocked_execute_query
+    execute_query_mock.__name__ = "mocked_execute_query"
+
+    command = AlertCommand(report_schedule=mocker.Mock())
+
+    command.validate()
+
+    assert execute_query_mock.call_count == expected_max_retries
+
+
+def test_execute_query_failed_no_retry(mocker: MockFixture, app_context: None) -> None:
+    from superset.reports.commands.alert import AlertCommand, AlertQueryTimeout
+
+    execute_query_mock = mocker.patch(
+        "superset.reports.commands.alert.AlertCommand._execute_query"
+    )
+
+    def _mocked_execute_query() -> None:
+        raise AlertQueryTimeout
+
+    execute_query_mock.side_effect = _mocked_execute_query
+    execute_query_mock.__name__ = "mocked_execute_query"
+
+    command = AlertCommand(report_schedule=mocker.Mock())
+
+    try:
+        command.validate()
+    except AlertQueryTimeout:
+        pass
+
+    assert execute_query_mock.call_count == 1
+
+
+def test_execute_query_failed_max_retries(
+    mocker: MockFixture, app_context: None
+) -> None:
+    from superset.reports.commands.alert import AlertCommand, AlertQueryError
+
+    execute_query_mock = mocker.patch(
+        "superset.reports.commands.alert.AlertCommand._execute_query"
+    )
+
+    def _mocked_execute_query() -> None:
+        raise AlertQueryError
+
+    execute_query_mock.side_effect = _mocked_execute_query
+    execute_query_mock.__name__ = "mocked_execute_query"
+
+    command = AlertCommand(report_schedule=mocker.Mock())
+
+    try:
+        command.validate()
+    except AlertQueryError:
+        pass
+
+    # Should match the value defined in superset_test_config.py
+    assert execute_query_mock.call_count == 3

tests/integration_tests/superset_test_config.py

@@ -115,6 +115,8 @@ GLOBAL_ASYNC_QUERIES_JWT_SECRET = "test-secret-change-me-test-secret-change-me"
 
 ALERT_REPORTS_WORKING_TIME_OUT_KILL = True
 
+ALERT_REPORTS_QUERY_EXECUTION_MAX_TRIES = 3
+
 
 class CeleryConfig(object):
     BROKER_URL = f"redis://{REDIS_HOST}:{REDIS_PORT}/{REDIS_CELERY_DB}"