From 88a96b3f6c774a3259662763e334ee28903baba2 Mon Sep 17 00:00:00 2001
From: Grace <grace.guo@airbnb.com>
Date: Mon, 14 Dec 2020 17:31:29 -0800
Subject: [PATCH] feat: add hook for dataset health check

---
 .../components/controls/DatasourceControl.jsx | 20 +++++++++++++++++--
 superset/config.py                            |  5 +++++
 superset/connectors/sqla/models.py            | 20 +++++++++++++++++++
 superset/datasets/dao.py                      |  2 ++
 superset/views/core.py                        |  4 ++++
 superset/views/datasource.py                  |  2 ++
 6 files changed, 51 insertions(+), 2 deletions(-)

diff --git a/superset-frontend/src/explore/components/controls/DatasourceControl.jsx b/superset-frontend/src/explore/components/controls/DatasourceControl.jsx
index d7c2a5546a574..cda1fe9bdcff1 100644
--- a/superset-frontend/src/explore/components/controls/DatasourceControl.jsx
+++ b/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';
@@ -213,10 +213,18 @@ class DatasourceControl extends React.PureComponent {
       </Menu>
     );
 
+    let healthCheckMessage = '';
+    const { extra: rawExtra } = datasource;
+    if (rawExtra) {
+      const extra = JSON.parse(rawExtra) || {};
+      // eslint-disable-next-line camelcase
+      healthCheckMessage = extra?.health_check?.message;
+    }
+
     return (
       <Styles className="DatasourceControl">
         <ControlHeader {...this.props} />
-        <div>
+        <div style={{ display: 'flex' }}>
           <Tooltip title={t('Expand/collapse dataset configuration')}>
             <Label
               style={{ textTransform: 'none' }}
@@ -230,6 +238,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']}
diff --git a/superset/config.py b/superset/config.py
index 22d3cc5262de6..61f4647b2ed46 100644
--- a/superset/config.py
+++ b/superset/config.py
@@ -991,3 +991,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
diff --git a/superset/connectors/sqla/models.py b/superset/connectors/sqla/models.py
index 6b2f57f6bb7df..dcec4b32665b0 100644
--- a/superset/connectors/sqla/models.py
+++ b/superset/connectors/sqla/models.py
@@ -698,6 +698,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_["extra"] = self.extra
         return data_
 
     @property
@@ -1467,6 +1468,25 @@ 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["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 not extra.get("health_check", {}).get("version") == check.version:
+            message = check(self)
+            if message:
+                extra["health_check"] = {"version": check.version, "message": message}
+            else:
+                extra.pop("health_check", None)
+            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)
diff --git a/superset/datasets/dao.py b/superset/datasets/dao.py
index 284a43508068a..aaede30f26b4a 100644
--- a/superset/datasets/dao.py
+++ b/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
diff --git a/superset/views/core.py b/superset/views/core.py
index 66cb05141e6be..cd29b3dd18fe4 100755
--- a/superset/views/core.py
+++ b/superset/views/core.py
@@ -650,6 +650,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") and conf["DATASET_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)
diff --git a/superset/views/datasource.py b/superset/views/datasource.py
index 92cf5c1540b35..7724d1e26c6b3 100644
--- a/superset/views/datasource.py
+++ b/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()