Skip to content

Commit

Permalink
feat: add hook for dataset health check
Browse files Browse the repository at this point in the history
  • Loading branch information
Grace committed Dec 15, 2020
1 parent cc44a2c commit 88a96b3
Show file tree
Hide file tree
Showing 6 changed files with 51 additions and 2 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Expand Down Expand Up @@ -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' }}
Expand All @@ -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']}
Expand Down
5 changes: 5 additions & 0 deletions superset/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
20 changes: 20 additions & 0 deletions superset/connectors/sqla/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)
Expand Down
2 changes: 2 additions & 0 deletions superset/datasets/dao.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
4 changes: 4 additions & 0 deletions superset/views/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
2 changes: 2 additions & 0 deletions superset/views/datasource.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()

Expand Down

0 comments on commit 88a96b3

Please sign in to comment.