feat: Stop pending queries when user close dashboard

superset-frontend/src/dashboard/components/Dashboard.jsx

@@ -34,12 +34,14 @@ import {
   LOG_ACTIONS_MOUNT_DASHBOARD,
   Logger,
 } from '../../logger/LogUtils';
+import { isFeatureEnabled, FeatureFlag } from '../../featureFlags';
 import OmniContainer from '../../components/OmniContainer';
 import { areObjectsEqual } from '../../reduxUtils';
 
 import '../stylesheets/index.less';
 import getLocationHash from '../util/getLocationHash';
 import isDashboardEmpty from '../util/isDashboardEmpty';
+import isDashboardLoading from '../util/isDashboardLoading';
 
 const propTypes = {
   actions: PropTypes.shape({
@@ -68,16 +70,7 @@ const defaultProps = {
 };
 
 class Dashboard extends React.PureComponent {
-  // eslint-disable-next-line react/sort-comp
-  static onBeforeUnload(hasChanged) {
-    if (hasChanged) {
-      window.addEventListener('beforeunload', Dashboard.unload);
-    } else {
-      window.removeEventListener('beforeunload', Dashboard.unload);
-    }
-  }
-
-  static unload() {
+  static showUnsavedMessage() {
     const message = t('You have unsaved changes.');
     window.event.returnValue = message; // Gecko + IE
     return message; // Gecko + Webkit, Safari, Chrome etc.
@@ -86,7 +79,11 @@ class Dashboard extends React.PureComponent {
   constructor(props) {
     super(props);
     this.appliedFilters = props.activeFilters || {};
+    this.canStopPendingQueries = isFeatureEnabled(
+      FeatureFlag.STOP_DASHBOARD_PENDING_QUERIES,
+    );
 
+    this.stopPendingQueries = this.stopPendingQueries.bind(this);
     this.onVisibilityChange = this.onVisibilityChange.bind(this);
   }
 
@@ -143,7 +140,8 @@ class Dashboard extends React.PureComponent {
   }
 
   componentDidUpdate() {
-    const { hasUnsavedChanges, editMode } = this.props.dashboardState;
+    const { charts, dashboardState } = this.props;
+    const { hasUnsavedChanges, editMode } = dashboardState;
 
     const { appliedFilters } = this;
     const { activeFilters } = this.props;
@@ -153,9 +151,15 @@ class Dashboard extends React.PureComponent {
     }
 
     if (hasUnsavedChanges) {
-      Dashboard.onBeforeUnload(true);
+      window.addEventListener('beforeunload', Dashboard.showUnsavedMessage);
+    } else {
+      window.removeEventListener('beforeunload', Dashboard.showUnsavedMessage);
+    }
+
+    if (this.canStopPendingQueries && isDashboardLoading(charts)) {
+      window.addEventListener('beforeunload', this.stopPendingQueries);
     } else {
-      Dashboard.onBeforeUnload(false);
+      window.removeEventListener('beforeunload', this.stopPendingQueries);
     }
   }
 
@@ -185,6 +189,14 @@ class Dashboard extends React.PureComponent {
     return Object.values(this.props.charts);
   }
 
+  stopPendingQueries() {
+    if (navigator && navigator.sendBeacon) {
+      navigator.sendBeacon(
+        `/superset/dashboard/${this.props.dashboardInfo.id}/stop/`,
+      );
+    }
+  }
+
   applyFilters() {
     const { appliedFilters } = this;
     const { activeFilters } = this.props;

superset-frontend/src/featureFlags.ts

@@ -27,6 +27,7 @@ export enum FeatureFlag {
   SHARE_QUERIES_VIA_KV_STORE = 'SHARE_QUERIES_VIA_KV_STORE',
   SQLLAB_BACKEND_PERSISTENCE = 'SQLLAB_BACKEND_PERSISTENCE',
   THUMBNAILS = 'THUMBNAILS',
+  STOP_DASHBOARD_PENDING_QUERIES = 'STOP_DASHBOARD_PENDING_QUERIES',
   SIP_34_SAVED_QUERIES_UI = 'SIP_34_SAVED_QUERIES_UI',
 }
 

superset/config.py

@@ -172,7 +172,11 @@ def _try_json_readsha(  # pylint: disable=unused-argument
 WTF_CSRF_ENABLED = True
 
 # Add endpoints that need to be exempt from CSRF protection
-WTF_CSRF_EXEMPT_LIST = ["superset.views.core.log", "superset.charts.api.data"]
+WTF_CSRF_EXEMPT_LIST = [
+    "superset.views.core.log",
+    "superset.charts.api.data",
+    "superset.views.core.stop_dashboard_queries",
+]
 
 # Whether to run the web server in debug mode or not
 DEBUG = os.environ.get("FLASK_ENV") == "development"
@@ -309,6 +313,8 @@ def _try_json_readsha(  # pylint: disable=unused-argument
     "SIP_38_VIZ_REARCHITECTURE": False,
     "TAGGING_SYSTEM": False,
     "SQLLAB_BACKEND_PERSISTENCE": False,
+    # stop pending queries when user close/reload dashboard in browser
+    "STOP_DASHBOARD_PENDING_QUERIES": False,
 }
 
 # This is merely a default.
@@ -627,11 +633,11 @@ class CeleryConfig:  # pylint: disable=too-few-public-methods
 # db configuration and a result of this function.
 
 # mypy doesn't catch that if case ensures list content being always str
-ALLOWED_USER_CSV_SCHEMA_FUNC: Callable[
-    ["Database", "models.User"], List[str]
-] = lambda database, user: [
-    UPLOADED_CSV_HIVE_NAMESPACE
-] if UPLOADED_CSV_HIVE_NAMESPACE else []
+ALLOWED_USER_CSV_SCHEMA_FUNC: Callable[["Database", "models.User"], List[str]] = (
+    lambda database, user: [UPLOADED_CSV_HIVE_NAMESPACE]
+    if UPLOADED_CSV_HIVE_NAMESPACE
+    else []
+)
 
 # Values that should be treated as nulls for the csv uploads.
 CSV_DEFAULT_NA_NAMES = list(STR_NA_VALUES)

superset/db_engine_specs/base.py

@@ -1025,3 +1025,13 @@ def get_extra_params(database: "Database") -> Dict[str, Any]:
                 logger.error(ex)
                 raise ex
         return extra
+
+    @classmethod
+    def stop_queries(cls, username: str, dashboard_id: int) -> None:
+        """
+        An empty function. The actual stop implementation depends on the engine
+
+        :param: username: user sends out queries
+        :param dashboard_id: dashboard has charts that waiting for queries
+        """
+        return None

superset/views/core.py

@@ -118,6 +118,7 @@
     check_slice_perms,
     get_cta_schema_name,
     get_dashboard_extra_filters,
+    get_database_ids,
     get_datasource_info,
     get_form_data,
     get_viz,
@@ -1421,6 +1422,25 @@ def fave_slices(  # pylint: disable=no-self-use
             payload.append(dash)
         return json_success(json.dumps(payload, default=utils.json_int_dttm_ser))
 
+    @api
+    @has_access_api
+    @event_logger.log_this
+    @expose("/dashboard/<int:dashboard_id>/stop/", methods=["POST"])
+    def stop_dashboard_queries(  # pylint: disable=no-self-use
+        self, dashboard_id: int
+    ) -> FlaskResponse:
+        if is_feature_enabled("STOP_DASHBOARD_PENDING_QUERIES"):
+            username = g.user.username
+            database_ids = get_database_ids(dashboard_id)
+
+            # stop pending query is only available for certain database(s)
+            for dbid in database_ids:
+                mydb = db.session.query(models.Database).get(dbid)
+                if mydb:
+                    mydb.db_engine_spec.stop_queries(username, int(dashboard_id))
+
+        return Response(status=200)
+
     @event_logger.log_this
     @api
     @has_access_api
@@ -1778,7 +1798,7 @@ def sync_druid_source(self) -> FlaskResponse:  # pylint: disable=no-self-use
     @expose("/get_or_create_table/", methods=["POST"])
     @event_logger.log_this
     def sqllab_table_viz(self) -> FlaskResponse:  # pylint: disable=no-self-use
-        """ Gets or creates a table object with attributes passed to the API.
+        """Gets or creates a table object with attributes passed to the API.
 
         It expects the json with params:
         * datasourceName - e.g. table name, required

superset/views/utils.py

@@ -213,6 +213,33 @@ def get_datasource_info(
     return datasource_id, datasource_type
 
 
+def get_database_ids(dashboard_id: int) -> List[int]:
+    """
+    Find all database ids used by a given dashboard
+
+    :param dashboard_id: The dashboard id
+    :returns: A list of database ids used by the given dashboard
+    """
+    dashboard = db.session.query(Dashboard).filter_by(id=dashboard_id).one()
+    slices = dashboard.slices
+    datasource_ids: Set[int] = set()
+    database_ids: Set[int] = set()
+
+    for slc in slices:
+        datasource = slc.datasource
+        if (
+            datasource
+            and datasource.type == "table"
+            and datasource.id not in datasource_ids
+        ):
+            datasource_ids.add(datasource.id)
+            database = datasource.database
+            if database:
+                database_ids.add(database.id)
+
+    return list(database_ids)
+
+
 def apply_display_max_row_limit(
     sql_results: Dict[str, Any], rows: Optional[int] = None
 ) -> Dict[str, Any]:

tests/core_tests.py

@@ -1237,6 +1237,22 @@ def test_get_column_names_from_metric(self):
             "my_col"
         ]
 
+    @mock.patch.dict(
+        "superset.extensions.feature_flag_manager._feature_flags",
+        {"STOP_DASHBOARD_PENDING_QUERIES": True},
+        clear=True,
+    )
+    def test_stop_dashboard_queries(self):
+        username = "admin"
+        self.login(username)
+        dashboard = self.get_dash_by_slug("births")
+        with mock.patch.object(BaseEngineSpec, "stop_queries") as mock_stop_queries:
+            resp = self.client.post(f"/superset/dashboard/{dashboard.id}/stop/")
+
+            self.assertTrue(is_feature_enabled("STOP_DASHBOARD_PENDING_QUERIES"))
+            self.assertEqual(resp.status_code, 200)
+            mock_stop_queries.assert_called_once_with(username, dashboard.id)
+
 
 if __name__ == "__main__":
     unittest.main()

tests/utils_tests.py

@@ -23,6 +23,7 @@
 import json
 import os
 import re
+from typing import List
 from unittest.mock import Mock, patch
 
 import numpy
@@ -32,6 +33,7 @@
 
 import tests.test_app
 from superset import app, db, security_manager
+from superset.connectors.base.models import BaseDatasource
 from superset.exceptions import CertificateException, SupersetException
 from superset.models.core import Database, Log
 from superset.models.dashboard import Dashboard
@@ -45,6 +47,7 @@
     get_form_data_token,
     get_iterable,
     get_email_address_list,
+    get_example_database,
     get_or_create_db,
     get_since_until,
     get_stacktrace,
@@ -67,6 +70,7 @@
 from superset.utils import schema
 from superset.views.utils import (
     build_extra_filters,
+    get_database_ids,
     get_form_data,
     get_time_range_endpoints,
 )
@@ -1134,3 +1138,80 @@ def test_get_form_data_token(self):
         assert get_form_data_token({"token": "token_abcdefg1"}) == "token_abcdefg1"
         generated_token = get_form_data_token({})
         assert re.match(r"^token_[a-z0-9]{8}$", generated_token) is not None
+
+    def test_get_database_ids(self) -> None:
+        world_health = db.session.query(Dashboard).filter_by(slug="world_health").one()
+        dash_id = world_health.id
+        database_ids = get_database_ids(dash_id)
+        assert len(database_ids) == 1
+        assert database_ids == [get_example_database().id]
+
+    def test_get_database_ids_empty_dash(self) -> None:
+        # test dash with no slice
+        dashboard = Dashboard(dashboard_title="no slices", id=101, slices=[])
+        with patch("superset.db.session.query") as mock_query:
+            mock_query.return_value.filter_by.return_value.one.return_value = dashboard
+            database_ids = get_database_ids(dashboard.id)
+            assert database_ids == []
+
+    def test_get_database_ids_multiple_databases(self) -> None:
+        # test dash with 2 databases
+        datasource_1 = Mock()
+        datasource_1.type = "table"
+        datasource_1.datasource_name = "table_datasource_1"
+        datasource_1.database = Mock()
+
+        datasource_2 = Mock()
+        datasource_2.type = "table"
+        datasource_2.datasource_name = "table_datasource_2"
+        datasource_2.database = Mock()
+
+        slices = [
+            Slice(
+                datasource_id=datasource_1.id,
+                datasource_type=datasource_1.type,
+                datasource_name=datasource_1.datasource_name,
+                slice_name="slice_name_1",
+            ),
+            Slice(
+                datasource_id=datasource_2.id,
+                datasource_type=datasource_2.type,
+                datasource_name=datasource_2.datasource_name,
+                slice_name="slice_name_2",
+            ),
+        ]
+        dashboard = Dashboard(dashboard_title="with 2 slices", id=102, slices=slices)
+        with patch("superset.db.session.query") as mock_query:
+            mock_query.return_value.filter_by.return_value.one.return_value = dashboard
+            mock_query.return_value.filter_by.return_value.first.side_effect = [
+                datasource_1,
+                datasource_2,
+            ]
+            database_ids = get_database_ids(dashboard.id)
+            self.assertCountEqual(
+                database_ids, [datasource_1.database.id, datasource_2.database.id]
+            )
+
+    def test_get_database_ids_druid(self) -> None:
+        druid_datasource = Mock()
+        druid_datasource.type = "druid"
+        druid_datasource.datasource_name = "druid_datasource_1"
+        druid_datasource.cluster = Mock()
+
+        slices = [
+            Slice(
+                datasource_id=druid_datasource.id,
+                datasource_type=druid_datasource.type,
+                datasource_name=druid_datasource.datasource_name,
+                slice_name="slice_name_1",
+            ),
+        ]
+        dashboard = Dashboard(dashboard_title="druid dash", id=103, slices=slices)
+        with patch("superset.db.session.query") as mock_query:
+            mock_query.return_value.filter_by.return_value.one.return_value = dashboard
+            mock_query.return_value.filter_by.return_value.first.return_value = (
+                druid_datasource
+            )
+            database_ids = get_database_ids(dashboard.id)
+            # druid slice has no database id
+            assert database_ids == []