From 88a96b3f6c774a3259662763e334ee28903baba2 Mon Sep 17 00:00:00 2001 From: Grace Date: Mon, 14 Dec 2020 17:31:29 -0800 Subject: [PATCH 1/6] 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 { ); + let healthCheckMessage = ''; + const { extra: rawExtra } = datasource; + if (rawExtra) { + const extra = JSON.parse(rawExtra) || {}; + // eslint-disable-next-line camelcase + healthCheckMessage = extra?.health_check?.message; + } + return ( -
+
+ {healthCheckMessage && ( + + + + )} 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() From 420adccfc2ebe3dd21954326d9f79bdcef44146d Mon Sep 17 00:00:00 2001 From: Grace Date: Mon, 14 Dec 2020 20:11:58 -0800 Subject: [PATCH 2/6] add event log --- superset/connectors/sqla/models.py | 25 +++++++++++++++---------- 1 file changed, 15 insertions(+), 10 deletions(-) diff --git a/superset/connectors/sqla/models.py b/superset/connectors/sqla/models.py index dcec4b32665b0..517cc41eade3f 100644 --- a/superset/connectors/sqla/models.py +++ b/superset/connectors/sqla/models.py @@ -58,6 +58,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, @@ -1475,17 +1476,21 @@ def health_check(self, commit: bool = False, force: bool = False) -> None: 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) + if force or extra.get("health_check", {}).get("version") != check.version: + with event_logger.log_context(action="dataset_health_check"): + 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() + db.session.merge(self) + if commit: + db.session.commit() sa.event.listen(SqlaTable, "after_insert", security_manager.set_perm) From c17f86ab264a0a5230c2e625c7b18ef9e0f42b69 Mon Sep 17 00:00:00 2001 From: Grace Date: Tue, 15 Dec 2020 12:23:00 -0800 Subject: [PATCH 3/6] optimize datasource json data like certified data --- .../explore/components/DatasourceControl_spec.jsx | 8 ++++++++ .../explore/components/controls/DatasourceControl.jsx | 9 ++------- superset/connectors/sqla/models.py | 6 +++++- 3 files changed, 15 insertions(+), 8 deletions(-) diff --git a/superset-frontend/spec/javascripts/explore/components/DatasourceControl_spec.jsx b/superset-frontend/spec/javascripts/explore/components/DatasourceControl_spec.jsx index 8f6b8e9492c6a..2484f8bd27ba9 100644 --- a/superset-frontend/spec/javascripts/explore/components/DatasourceControl_spec.jsx +++ b/superset-frontend/spec/javascripts/explore/components/DatasourceControl_spec.jsx @@ -24,6 +24,7 @@ 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'; const defaultProps = { name: 'datasource', @@ -40,6 +41,7 @@ const defaultProps = { backend: 'mysql', name: 'main', }, + health_check_message: 'Warning message!', }, actions: { setDatasource: sinon.spy(), @@ -91,4 +93,10 @@ describe('DatasourceControl', () => { ); expect(menuWrapper.find(Menu.Item)).toHaveLength(2); }); + + it('should render health check message', () => { + const wrapper = setup(); + const icons = wrapper.find(Icon); + expect(icons.first().prop('name')).toBe('alert-solid'); + }); }); diff --git a/superset-frontend/src/explore/components/controls/DatasourceControl.jsx b/superset-frontend/src/explore/components/controls/DatasourceControl.jsx index cda1fe9bdcff1..b7622d4c6470e 100644 --- a/superset-frontend/src/explore/components/controls/DatasourceControl.jsx +++ b/superset-frontend/src/explore/components/controls/DatasourceControl.jsx @@ -213,13 +213,8 @@ class DatasourceControl extends React.PureComponent { ); - let healthCheckMessage = ''; - const { extra: rawExtra } = datasource; - if (rawExtra) { - const extra = JSON.parse(rawExtra) || {}; - // eslint-disable-next-line camelcase - healthCheckMessage = extra?.health_check?.message; - } + // eslint-disable-next-line camelcase + const { health_check_message: healthCheckMessage } = datasource; return ( diff --git a/superset/connectors/sqla/models.py b/superset/connectors/sqla/models.py index 517cc41eade3f..6214a984e700c 100644 --- a/superset/connectors/sqla/models.py +++ b/superset/connectors/sqla/models.py @@ -686,6 +686,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,7 +703,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 + data_["health_check_message"] = self.health_check_message return data_ @property From 5146afa83217c5468efc6b525cd702e94dc353ff Mon Sep 17 00:00:00 2001 From: Grace Date: Tue, 15 Dec 2020 15:30:45 -0800 Subject: [PATCH 4/6] add unit test --- superset/connectors/sqla/models.py | 2 +- tests/datasource_tests.py | 18 +++++++++++++++++- 2 files changed, 18 insertions(+), 2 deletions(-) diff --git a/superset/connectors/sqla/models.py b/superset/connectors/sqla/models.py index 6214a984e700c..d8b912b239cc9 100644 --- a/superset/connectors/sqla/models.py +++ b/superset/connectors/sqla/models.py @@ -1474,7 +1474,7 @@ class and any keys added via `ExtraCache`. return extra_cache_keys def health_check(self, commit: bool = False, force: bool = False) -> None: - check = config["DATASET_HEALTH_CHECK"] + check = config.get("DATASET_HEALTH_CHECK") if check is None: return diff --git a/tests/datasource_tests.py b/tests/datasource_tests.py index 31a4c96334188..890b4a63e42c6 100644 --- a/tests/datasource_tests.py +++ b/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/" From 975d275422b87d92b7aa6a22db61c72a42af62ca Mon Sep 17 00:00:00 2001 From: Grace Date: Tue, 15 Dec 2020 15:45:22 -0800 Subject: [PATCH 5/6] fix review comments --- superset/connectors/sqla/models.py | 11 ++++------- superset/views/core.py | 2 +- 2 files changed, 5 insertions(+), 8 deletions(-) diff --git a/superset/connectors/sqla/models.py b/superset/connectors/sqla/models.py index d8b912b239cc9..ab6b88ea317c3 100644 --- a/superset/connectors/sqla/models.py +++ b/superset/connectors/sqla/models.py @@ -1483,13 +1483,10 @@ def health_check(self, commit: bool = False, force: bool = False) -> None: if force or extra.get("health_check", {}).get("version") != check.version: with event_logger.log_context(action="dataset_health_check"): message = check(self) - if message: - extra["health_check"] = { - "version": check.version, - "message": message, - } - else: - extra.pop("health_check", None) + extra["health_check"] = { + "version": check.version, + "message": message, + } self.extra = json.dumps(extra) db.session.merge(self) diff --git a/superset/views/core.py b/superset/views/core.py index cd29b3dd18fe4..5c5881d5a9997 100755 --- a/superset/views/core.py +++ b/superset/views/core.py @@ -651,7 +651,7 @@ def explore( # pylint: disable=too-many-locals,too-many-return-statements ) # if feature enabled, run some health check rules for sqla datasource - if hasattr(datasource, "health_check") and conf["DATASET_HEALTH_CHECK"]: + if hasattr(datasource, "health_check"): datasource.health_check() viz_type = form_data.get("viz_type") From 98006068fe0410ed4567cfb23499066d40731782 Mon Sep 17 00:00:00 2001 From: Grace Date: Tue, 15 Dec 2020 16:49:07 -0800 Subject: [PATCH 6/6] extra code review comments --- .../explore/components/DatasourceControl_spec.jsx | 9 +++++++-- .../explore/components/controls/DatasourceControl.jsx | 6 +++++- 2 files changed, 12 insertions(+), 3 deletions(-) diff --git a/superset-frontend/spec/javascripts/explore/components/DatasourceControl_spec.jsx b/superset-frontend/spec/javascripts/explore/components/DatasourceControl_spec.jsx index 2484f8bd27ba9..1c91c1cd25256 100644 --- a/superset-frontend/spec/javascripts/explore/components/DatasourceControl_spec.jsx +++ b/superset-frontend/spec/javascripts/explore/components/DatasourceControl_spec.jsx @@ -25,6 +25,7 @@ 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', @@ -96,7 +97,11 @@ describe('DatasourceControl', () => { it('should render health check message', () => { const wrapper = setup(); - const icons = wrapper.find(Icon); - expect(icons.first().prop('name')).toBe('alert-solid'); + 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, + ); }); }); diff --git a/superset-frontend/src/explore/components/controls/DatasourceControl.jsx b/superset-frontend/src/explore/components/controls/DatasourceControl.jsx index b7622d4c6470e..f8f79d36499ff 100644 --- a/superset-frontend/src/explore/components/controls/DatasourceControl.jsx +++ b/superset-frontend/src/explore/components/controls/DatasourceControl.jsx @@ -73,6 +73,10 @@ const Styles = styled.div` vertical-align: middle; cursor: pointer; } + + .datasource-controls { + display: flex; + } `; /** @@ -219,7 +223,7 @@ class DatasourceControl extends React.PureComponent { return ( -
+