From 144b279aa26a3b19b0ef02b0bd7437e511fa910f Mon Sep 17 00:00:00 2001 From: Daniel Vaz Gaspar Date: Mon, 26 Oct 2020 18:20:07 +0000 Subject: [PATCH] feat: saved queries with execution info (#11391) * feat: add rows and last_run info to saved queries * feat: add rows to saved query * refactor and tests * lint * fix tests --- ...0b1a7b18_add_exec_info_to_saved_queries.py | 40 +++++++++++++++++++ superset/models/sql_lab.py | 16 ++++++++ superset/queries/dao.py | 24 ++++++++++- superset/queries/saved_queries/api.py | 4 ++ superset/views/core.py | 10 +++++ tests/queries/saved_queries/api_tests.py | 6 ++- tests/sqllab_tests.py | 32 ++++++++++++++- 7 files changed, 128 insertions(+), 4 deletions(-) create mode 100644 superset/migrations/versions/585b0b1a7b18_add_exec_info_to_saved_queries.py diff --git a/superset/migrations/versions/585b0b1a7b18_add_exec_info_to_saved_queries.py b/superset/migrations/versions/585b0b1a7b18_add_exec_info_to_saved_queries.py new file mode 100644 index 000000000..54755c429 --- /dev/null +++ b/superset/migrations/versions/585b0b1a7b18_add_exec_info_to_saved_queries.py @@ -0,0 +1,40 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""add exec info to saved queries + +Revision ID: 585b0b1a7b18 +Revises: af30ca79208f +Create Date: 2020-10-20 17:28:22.857694 + +""" + +# revision identifiers, used by Alembic. +revision = "585b0b1a7b18" +down_revision = "af30ca79208f" + +import sqlalchemy as sa +from alembic import op + + +def upgrade(): + op.add_column("saved_query", sa.Column("last_run", sa.DateTime(), nullable=True)) + op.add_column("saved_query", sa.Column("rows", sa.Integer(), nullable=True)) + + +def downgrade(): + op.drop_column("saved_query", "rows") + op.drop_column("saved_query", "last_run") diff --git a/superset/models/sql_lab.py b/superset/models/sql_lab.py index f7d484887..f7f8f86db 100644 --- a/superset/models/sql_lab.py +++ b/superset/models/sql_lab.py @@ -23,6 +23,8 @@ import simplejson as json import sqlalchemy as sqla from flask import Markup from flask_appbuilder import Model +from flask_appbuilder.models.decorators import renders +from humanize import naturaltime from sqlalchemy import ( Boolean, Column, @@ -181,6 +183,8 @@ class SavedQuery(Model, AuditMixinNullable, ExtraJSONMixin): foreign_keys=[db_id], backref=backref("saved_queries", cascade="all, delete-orphan"), ) + rows = Column(Integer, nullable=True) + last_run = Column(DateTime, nullable=True) def __repr__(self) -> str: return str(self.label) @@ -210,6 +214,18 @@ class SavedQuery(Model, AuditMixinNullable, ExtraJSONMixin): def sql_tables(self) -> List[Table]: return list(ParsedQuery(self.sql).tables) + @property + def last_run_humanized(self) -> str: + return naturaltime(datetime.now() - self.changed_on) + + @property + def _last_run_delta_humanized(self) -> str: + return naturaltime(datetime.now() - self.changed_on) + + @renders("changed_on") + def last_run_delta_humanized(self) -> str: + return self._last_run_delta_humanized + class TabState(Model, AuditMixinNullable, ExtraJSONMixin): diff --git a/superset/queries/dao.py b/superset/queries/dao.py index a3317d2e8..2f438bdb3 100644 --- a/superset/queries/dao.py +++ b/superset/queries/dao.py @@ -15,9 +15,11 @@ # specific language governing permissions and limitations # under the License. import logging +from datetime import datetime from superset.dao.base import BaseDAO -from superset.models.sql_lab import Query +from superset.extensions import db +from superset.models.sql_lab import Query, SavedQuery from superset.queries.filters import QueryFilter logger = logging.getLogger(__name__) @@ -26,3 +28,23 @@ logger = logging.getLogger(__name__) class QueryDAO(BaseDAO): model_cls = Query base_filter = QueryFilter + + @staticmethod + def update_saved_query_exec_info(query_id: int) -> None: + """ + Propagates query execution info back to saved query if applicable + + :param query_id: The query id + :return: + """ + query = db.session.query(Query).get(query_id) + related_saved_queries = ( + db.session.query(SavedQuery) + .filter(SavedQuery.database == query.database) + .filter(SavedQuery.sql == query.sql) + ).all() + if related_saved_queries: + for saved_query in related_saved_queries: + saved_query.rows = query.rows + saved_query.last_run = datetime.now() + db.session.commit() diff --git a/superset/queries/saved_queries/api.py b/superset/queries/saved_queries/api.py index 37e30841f..794ab2572 100644 --- a/superset/queries/saved_queries/api.py +++ b/superset/queries/saved_queries/api.py @@ -88,6 +88,8 @@ class SavedQueryRestApi(BaseSupersetModelRestApi): "schema", "sql", "sql_tables", + "rows", + "last_run_delta_humanized", ] add_columns = ["db_id", "description", "label", "schema", "sql"] edit_columns = add_columns @@ -96,10 +98,12 @@ class SavedQueryRestApi(BaseSupersetModelRestApi): "label", "description", "sql", + "rows", "created_by.first_name", "database.database_name", "created_on", "changed_on_delta_humanized", + "last_run_delta_humanized", ] search_columns = ["id", "database", "label", "schema"] diff --git a/superset/views/core.py b/superset/views/core.py index 8380823b2..da6c0b1df 100755 --- a/superset/views/core.py +++ b/superset/views/core.py @@ -82,6 +82,7 @@ from superset.models.datasource_access_request import DatasourceAccessRequest from superset.models.slice import Slice from superset.models.sql_lab import Query, TabState from superset.models.user_attributes import UserAttribute +from superset.queries.dao import QueryDAO from superset.security.analytics_db_safety import ( check_sqlalchemy_uri, DBSecurityException, @@ -2144,6 +2145,7 @@ class Superset(BaseSupersetView): # pylint: disable=too-many-public-methods """ logger.info("Query %i: Running query on a Celery worker", query.id) # Ignore the celery future object and the request may time out. + query_id = query.id try: task = sql_lab.get_sql_results.delay( query.id, @@ -2170,6 +2172,10 @@ class Superset(BaseSupersetView): # pylint: disable=too-many-public-methods query.error_message = msg session.commit() return json_error_response("{}".format(msg)) + + # Update saved query with execution info from the query execution + QueryDAO.update_saved_query_exec_info(query_id) + resp = json_success( json.dumps( {"query": query.to_dict()}, @@ -2204,6 +2210,7 @@ class Superset(BaseSupersetView): # pylint: disable=too-many-public-methods is_feature_enabled("SQLLAB_BACKEND_PERSISTENCE") and not query.select_as_cta ) + query_id = query.id with utils.timeout(seconds=timeout, error_message=timeout_msg): # pylint: disable=no-value-for-parameter data = sql_lab.get_sql_results( @@ -2216,6 +2223,9 @@ class Superset(BaseSupersetView): # pylint: disable=too-many-public-methods log_params=log_params, ) + # Update saved query if needed + QueryDAO.update_saved_query_exec_info(query_id) + payload = json.dumps( apply_display_max_row_limit(data), default=utils.pessimistic_json_iso_dttm_ser, diff --git a/tests/queries/saved_queries/api_tests.py b/tests/queries/saved_queries/api_tests.py index ace164baa..0b7df8254 100644 --- a/tests/queries/saved_queries/api_tests.py +++ b/tests/queries/saved_queries/api_tests.py @@ -168,8 +168,10 @@ class TestSavedQueryApi(SupersetTestCase): """ admin = self.get_user("admin") saved_queries = ( - db.session.query(SavedQuery).filter(SavedQuery.created_by == admin).all() - ) + db.session.query(SavedQuery) + .filter(SavedQuery.created_by == admin) + .order_by(SavedQuery.schema.asc()) + ).all() self.login(username="admin") query_string = {"order_column": "schema", "order_direction": "asc"} uri = f"api/v1/saved_query/?q={prison.dumps(query_string)}" diff --git a/tests/sqllab_tests.py b/tests/sqllab_tests.py index 97e4f3fe0..52fd32c48 100644 --- a/tests/sqllab_tests.py +++ b/tests/sqllab_tests.py @@ -28,7 +28,7 @@ import tests.test_app from superset import db, security_manager from superset.connectors.sqla.models import SqlaTable from superset.db_engine_specs import BaseEngineSpec -from superset.models.sql_lab import Query +from superset.models.sql_lab import Query, SavedQuery from superset.result_set import SupersetResultSet from superset.sql_parse import CtasMethod from superset.utils.core import ( @@ -71,6 +71,36 @@ class TestSqlLab(SupersetTestCase): data = self.run_sql("SELECT * FROM unexistant_table", "2") self.assertLess(0, len(data["error"])) + def test_sql_json_to_saved_query_info(self): + """ + SQLLab: Test SQLLab query execution info propagation to saved queries + """ + from freezegun import freeze_time + + self.login("admin") + + sql_statement = "SELECT * FROM birth_names LIMIT 10" + examples_db_id = get_example_database().id + saved_query = SavedQuery(db_id=examples_db_id, sql=sql_statement) + db.session.add(saved_query) + db.session.commit() + + with freeze_time("2020-01-01T00:00:00Z"): + self.run_sql(sql_statement, "1") + saved_query_ = ( + db.session.query(SavedQuery) + .filter( + SavedQuery.db_id == examples_db_id, SavedQuery.sql == sql_statement + ) + .one_or_none() + ) + assert saved_query_.rows is not None + assert saved_query_.last_run == datetime.now() + + # Rollback changes + db.session.delete(saved_query_) + db.session.commit() + @parameterized.expand([CtasMethod.TABLE, CtasMethod.VIEW]) def test_sql_json_cta_dynamic_db(self, ctas_method): examples_db = get_example_database()