feat: add hook for dataset health check (#11970)

* feat: add hook for dataset health check

* add event log

* optimize datasource json data like certified data

* add unit test

* fix review comments

* extra code review comments

superset-frontend/spec/javascripts/explore/components/DatasourceControl_spec.jsx

@@ -24,6 +24,8 @@ import { Menu } from 'src/common/components';
 import DatasourceModal from 'src/datasource/DatasourceModal';
 import ChangeDatasourceModal from 'src/datasource/ChangeDatasourceModal';
 import DatasourceControl from 'src/explore/components/controls/DatasourceControl';
+import Icon from 'src/components/Icon';
+import { Tooltip } from 'src/common/components/Tooltip';
 
 const defaultProps = {
   name: 'datasource',
@@ -40,6 +42,7 @@ const defaultProps = {
       backend: 'mysql',
       name: 'main',
     },
+    health_check_message: 'Warning message!',
   },
   actions: {
     setDatasource: sinon.spy(),
@@ -91,4 +94,14 @@ describe('DatasourceControl', () => {
     );
     expect(menuWrapper.find(Menu.Item)).toHaveLength(2);
   });
+
+  it('should render health check message', () => {
+    const wrapper = setup();
+    const alert = wrapper.find(Icon).first();
+    expect(alert.prop('name')).toBe('alert-solid');
+    const tooltip = wrapper.find(Tooltip).at(1);
+    expect(tooltip.prop('title')).toBe(
+      defaultProps.datasource.health_check_message,
+    );
+  });
 });

superset-frontend/src/explore/components/controls/DatasourceControl.jsx

@@ -19,7 +19,7 @@
 import React from 'react';
 import PropTypes from 'prop-types';
 import { Col, Collapse, Row, Well } from 'react-bootstrap';
-import { t, styled } from '@superset-ui/core';
+import { t, styled, supersetTheme } from '@superset-ui/core';
 import { ColumnOption, MetricOption } from '@superset-ui/chart-controls';
 
 import { Dropdown, Menu } from 'src/common/components';
@@ -73,6 +73,10 @@ const Styles = styled.div`
     vertical-align: middle;
     cursor: pointer;
   }
+
+  .datasource-controls {
+    display: flex;
+  }
 `;
 
 /**
@@ -213,10 +217,13 @@ class DatasourceControl extends React.PureComponent {
       </Menu>
     );
 
+    // eslint-disable-next-line camelcase
+    const { health_check_message: healthCheckMessage } = datasource;
+
     return (
       <Styles className="DatasourceControl">
         <ControlHeader {...this.props} />
-        <div>
+        <div className="datasource-controls">
           <Tooltip title={t('Expand/collapse dataset configuration')}>
             <Label
               style={{ textTransform: 'none' }}
@@ -230,6 +237,14 @@ class DatasourceControl extends React.PureComponent {
               />
             </Label>
           </Tooltip>
+          {healthCheckMessage && (
+            <Tooltip title={healthCheckMessage}>
+              <Icon
+                name="alert-solid"
+                color={supersetTheme.colors.warning.base}
+              />
+            </Tooltip>
+          )}
           <Dropdown
             overlay={datasourceMenu}
             trigger={['click']}

superset/config.py

@@ -1012,3 +1012,8 @@ class CeleryConfig:  # pylint: disable=too-few-public-methods
     except Exception:
         logger.exception("Found but failed to import local superset_config")
         raise
+
+
+# It's possible to add a dataset health check logic which is specific to your system.
+# It will get executed each time when user open a chart's explore view.
+DATASET_HEALTH_CHECK = None

superset/connectors/sqla/models.py

@@ -57,6 +57,7 @@
 from superset.db_engine_specs.base import TimestampExpression
 from superset.errors import ErrorLevel, SupersetError, SupersetErrorType
 from superset.exceptions import QueryObjectValidationError, SupersetSecurityException
+from superset.extensions import event_logger
 from superset.jinja_context import (
     BaseTemplateProcessor,
     ExtraCache,
@@ -686,6 +687,10 @@ def select_star(self) -> Optional[str]:
             self.table_name, schema=self.schema, show_cols=False, latest_partition=False
         )
 
+    @property
+    def health_check_message(self) -> Optional[str]:
+        return self.extra_dict.get("health_check", {}).get("message")
+
     @property
     def data(self) -> Dict[str, Any]:
         data_ = super().data
@@ -699,6 +704,7 @@ def data(self) -> Dict[str, Any]:
             data_["fetch_values_predicate"] = self.fetch_values_predicate
             data_["template_params"] = self.template_params
             data_["is_sqllab_view"] = self.is_sqllab_view
+            data_["health_check_message"] = self.health_check_message
         return data_
 
     @property
@@ -1468,6 +1474,26 @@ class and any keys added via `ExtraCache`.
             extra_cache_keys += sqla_query.extra_cache_keys
         return extra_cache_keys
 
+    def health_check(self, commit: bool = False, force: bool = False) -> None:
+        check = config.get("DATASET_HEALTH_CHECK")
+        if check is None:
+            return
+
+        extra = self.extra_dict
+        # force re-run health check, or health check is updated
+        if force or extra.get("health_check", {}).get("version") != check.version:
+            with event_logger.log_context(action="dataset_health_check"):
+                message = check(self)
+                extra["health_check"] = {
+                    "version": check.version,
+                    "message": message,
+                }
+                self.extra = json.dumps(extra)
+
+                db.session.merge(self)
+                if commit:
+                    db.session.commit()
+
 
 sa.event.listen(SqlaTable, "after_insert", security_manager.set_perm)
 sa.event.listen(SqlaTable, "after_update", security_manager.set_perm)

superset/datasets/dao.py

@@ -186,6 +186,8 @@ def update(  # pylint: disable=W:279
             super().update(model, properties, commit=commit)
             properties["columns"] = original_properties
 
+        super().update(model, properties, commit=False)
+        model.health_check(force=True, commit=False)
         return super().update(model, properties, commit=commit)
 
     @classmethod

superset/views/core.py

@@ -717,6 +717,10 @@ def explore(  # pylint: disable=too-many-locals,too-many-return-statements
                 f"datasource_id={datasource_id}&"
             )
 
+        # if feature enabled, run some health check rules for sqla datasource
+        if hasattr(datasource, "health_check"):
+            datasource.health_check()
+
         viz_type = form_data.get("viz_type")
         if not viz_type and datasource.default_endpoint:
             return redirect(datasource.default_endpoint)

superset/views/datasource.py

@@ -70,6 +70,8 @@ def save(self) -> FlaskResponse:
                 f"Duplicate column name(s): {','.join(duplicates)}", status=409
             )
         orm_datasource.update_from_object(datasource_dict)
+        if hasattr(orm_datasource, "health_check"):
+            orm_datasource.health_check(force=True, commit=False)
         data = orm_datasource.data
         db.session.commit()
 

tests/datasource_tests.py

@@ -18,7 +18,7 @@
 import json
 from copy import deepcopy
 
-from superset import db
+from superset import app, db
 from superset.connectors.sqla.models import SqlaTable
 from superset.utils.core import get_example_database
 
@@ -190,6 +190,22 @@ def test_get_datasource(self):
             },
         )
 
+    def test_get_datasource_with_health_check(self):
+        def my_check(datasource):
+            return "Warning message!"
+
+        app.config["DATASET_HEALTH_CHECK"] = my_check
+        my_check.version = 0.1
+
+        self.login(username="admin")
+        tbl = self.get_table_by_name("birth_names")
+        url = f"/datasource/get/{tbl.type}/{tbl.id}/"
+        tbl.health_check(commit=True, force=True)
+        resp = self.get_json_resp(url)
+        self.assertEqual(resp["health_check_message"], "Warning message!")
+
+        del app.config["DATASET_HEALTH_CHECK"]
+
     def test_get_datasource_failed(self):
         self.login(username="admin")
         url = f"/datasource/get/druid/500000/"