From 4a93a3c916ecea809a05897cd5e6425281fd8fea Mon Sep 17 00:00:00 2001 From: "Hugh A. Miles II" Date: Fri, 13 May 2022 16:51:50 +0000 Subject: [PATCH 1/9] add save_metadata function to QueryDAO --- superset/queries/dao.py | 20 ++++++++++++++++++++ superset/sqllab/command.py | 15 ++++++++++++--- 2 files changed, 32 insertions(+), 3 deletions(-) diff --git a/superset/queries/dao.py b/superset/queries/dao.py index 2f438bdb369ff..0015abd77ad2c 100644 --- a/superset/queries/dao.py +++ b/superset/queries/dao.py @@ -14,8 +14,10 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. +import json import logging from datetime import datetime +from typing import Any, Dict, List from superset.dao.base import BaseDAO from superset.extensions import db @@ -48,3 +50,21 @@ def update_saved_query_exec_info(query_id: int) -> None: saved_query.rows = query.rows saved_query.last_run = datetime.now() db.session.commit() + + @staticmethod + def save_metadata(query: Query, query_payload: str) -> None: + # parse payload + try: + payload: Dict[str, Any] = json.loads(query_payload) + except json.JSONDecodeError: + logger.warning("Error parsing query_payload: %s", query_payload) + return None + + # pull relevant data from payload and store in json_metadata + columns = payload.get("columns", {}) + metadata: Dict[str, Any] = {"columns": columns} + + # save payload into query object + query.json_metadata = json.dumps(metadata) + db.session.commit() + return None diff --git a/superset/sqllab/command.py b/superset/sqllab/command.py index ff50e18eda9e5..69420f657e0f4 100644 --- a/superset/sqllab/command.py +++ b/superset/sqllab/command.py @@ -94,11 +94,20 @@ def run( # pylint: disable=too-many-statements,useless-suppression status = SqlJsonExecutionStatus.QUERY_ALREADY_CREATED else: status = self._run_sql_json_exec_from_scratch() + + payload = self._execution_context_convertor.to_payload( + self._execution_context, status + ) + + # save columns into metadata_json + if query: + self._query_dao.save_metadata(query, payload) + else: + self._query_dao.save_metadata(self._execution_context.query, payload) + return { "status": status, - "payload": self._execution_context_convertor.to_payload( - self._execution_context, status - ), + "payload": payload, } except (SqlLabException, SupersetErrorsException) as ex: raise ex From 30a613bc9e5f3ba66588b66fae8f0a4f6fa75ca1 Mon Sep 17 00:00:00 2001 From: "Hugh A. Miles II" Date: Mon, 16 May 2022 15:29:19 +0000 Subject: [PATCH 2/9] use set_extra_json_key --- superset/queries/dao.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/superset/queries/dao.py b/superset/queries/dao.py index 0015abd77ad2c..dfd3759aa097e 100644 --- a/superset/queries/dao.py +++ b/superset/queries/dao.py @@ -62,9 +62,8 @@ def save_metadata(query: Query, query_payload: str) -> None: # pull relevant data from payload and store in json_metadata columns = payload.get("columns", {}) - metadata: Dict[str, Any] = {"columns": columns} # save payload into query object - query.json_metadata = json.dumps(metadata) + query.set_extra_json_key("columns", columns) db.session.commit() return None From d34a054d9203eb751b08c8a237685fea2cfdff46 Mon Sep 17 00:00:00 2001 From: "Hugh A. Miles II" Date: Mon, 16 May 2022 16:29:25 +0000 Subject: [PATCH 3/9] added test --- tests/unit_tests/dao/queries_test.py | 39 ++++++++++++++++++++++++++++ 1 file changed, 39 insertions(+) create mode 100644 tests/unit_tests/dao/queries_test.py diff --git a/tests/unit_tests/dao/queries_test.py b/tests/unit_tests/dao/queries_test.py new file mode 100644 index 0000000000000..60e99c229c56d --- /dev/null +++ b/tests/unit_tests/dao/queries_test.py @@ -0,0 +1,39 @@ +import json +from typing import Iterator + +import pytest +from sqlalchemy.orm.session import Session + + +def test_query_dao_save_metadata(app_context: None, session: Session) -> None: + from superset.models.core import Database + from superset.models.sql_lab import Query + + engine = session.get_bind() + Query.metadata.create_all(engine) # pylint: disable=no-member + + db = Database(database_name="my_database", sqlalchemy_uri="sqlite://") + + query_obj = Query( + client_id="foo", + database=db, + tab_name="test_tab", + sql_editor_id="test_editor_id", + sql="select * from bar", + select_sql="select * from bar", + executed_sql="select * from bar", + limit=100, + select_as_cta=False, + rows=100, + error_message="none", + results_key="abc", + ) + + session.add(db) + session.add(query_obj) + + from superset.queries.dao import QueryDAO + + query = session.query(Query).one() + QueryDAO.save_metadata(query=query, query_payload=json.dumps({"columns": []})) + assert query.extra.get("columns", None) == [] From ee932e0c960971843eac4d46a72bd959222e2020 Mon Sep 17 00:00:00 2001 From: "Hugh A. Miles II" Date: Mon, 16 May 2022 12:35:31 -0400 Subject: [PATCH 4/9] Update queries_test.py --- tests/unit_tests/dao/queries_test.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/tests/unit_tests/dao/queries_test.py b/tests/unit_tests/dao/queries_test.py index 60e99c229c56d..98caad6dfe59e 100644 --- a/tests/unit_tests/dao/queries_test.py +++ b/tests/unit_tests/dao/queries_test.py @@ -1,3 +1,19 @@ +# 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. import json from typing import Iterator From 8ff1349fd9024048ca7dc47160ff44244246adb3 Mon Sep 17 00:00:00 2001 From: "Hugh A. Miles II" Date: Mon, 16 May 2022 19:28:42 +0000 Subject: [PATCH 5/9] fix pylint --- superset/queries/dao.py | 2 +- superset/sql_lab.py | 2 +- superset/views/core.py | 4 ++-- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/superset/queries/dao.py b/superset/queries/dao.py index dfd3759aa097e..a0398a7eea834 100644 --- a/superset/queries/dao.py +++ b/superset/queries/dao.py @@ -17,7 +17,7 @@ import json import logging from datetime import datetime -from typing import Any, Dict, List +from typing import Any, Dict from superset.dao.base import BaseDAO from superset.extensions import db diff --git a/superset/sql_lab.py b/superset/sql_lab.py index cd3684b1d8ea8..2eeb2976b4126 100644 --- a/superset/sql_lab.py +++ b/superset/sql_lab.py @@ -187,7 +187,7 @@ def get_sql_results( # pylint: disable=too-many-arguments return handle_query_error(ex, query, session) -def execute_sql_statement( # pylint: disable=too-many-arguments,too-many-locals,too-many-statements +def execute_sql_statement( # pylint: disable=too-many-arguments,too-many-statements sql_statement: str, query: Query, session: Session, diff --git a/superset/views/core.py b/superset/views/core.py index fa3437e47b9c0..a42596323c65d 100755 --- a/superset/views/core.py +++ b/superset/views/core.py @@ -1331,7 +1331,7 @@ def add_slices( # pylint: disable=no-self-use @has_access_api @event_logger.log_this @expose("/testconn", methods=["POST", "GET"]) - def testconn(self) -> FlaskResponse: # pylint: disable=no-self-use + def testconn(self) -> FlaskResponse: """Tests a sqla connection""" logger.warning( "%s.testconn " @@ -2306,7 +2306,7 @@ def stop_query(self) -> FlaskResponse: @event_logger.log_this @expose("/validate_sql_json/", methods=["POST", "GET"]) def validate_sql_json( - # pylint: disable=too-many-locals,no-self-use + # pylint: disable=too-many-locals self, ) -> FlaskResponse: """Validates that arbitrary sql is acceptable for the given database. From 3618c8afc4b995ca1176572fbe2cd8f31e0f8ffa Mon Sep 17 00:00:00 2001 From: "Hugh A. Miles II" Date: Mon, 16 May 2022 20:29:09 +0000 Subject: [PATCH 6/9] add to session --- superset/queries/dao.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/superset/queries/dao.py b/superset/queries/dao.py index a0398a7eea834..9a072dc0a7c72 100644 --- a/superset/queries/dao.py +++ b/superset/queries/dao.py @@ -64,6 +64,6 @@ def save_metadata(query: Query, query_payload: str) -> None: columns = payload.get("columns", {}) # save payload into query object + db.session.add(query) query.set_extra_json_key("columns", columns) - db.session.commit() return None From c80634706be23b2471dca2b281ca0260e82e2b15 Mon Sep 17 00:00:00 2001 From: "Hugh A. Miles II" Date: Mon, 16 May 2022 20:29:09 +0000 Subject: [PATCH 7/9] add to session --- superset/queries/dao.py | 2 +- superset/sqllab/command.py | 5 +---- 2 files changed, 2 insertions(+), 5 deletions(-) diff --git a/superset/queries/dao.py b/superset/queries/dao.py index a0398a7eea834..9a072dc0a7c72 100644 --- a/superset/queries/dao.py +++ b/superset/queries/dao.py @@ -64,6 +64,6 @@ def save_metadata(query: Query, query_payload: str) -> None: columns = payload.get("columns", {}) # save payload into query object + db.session.add(query) query.set_extra_json_key("columns", columns) - db.session.commit() return None diff --git a/superset/sqllab/command.py b/superset/sqllab/command.py index 69420f657e0f4..f8aa612c0bfe9 100644 --- a/superset/sqllab/command.py +++ b/superset/sqllab/command.py @@ -100,10 +100,7 @@ def run( # pylint: disable=too-many-statements,useless-suppression ) # save columns into metadata_json - if query: - self._query_dao.save_metadata(query, payload) - else: - self._query_dao.save_metadata(self._execution_context.query, payload) + self._query_dao.save_metadata(self._execution_context.query, payload) return { "status": status, From 4c7549dd020daa4265600616c6cb9bf7c4c706c1 Mon Sep 17 00:00:00 2001 From: "Hugh A. Miles II" Date: Tue, 17 May 2022 20:24:58 +0000 Subject: [PATCH 8/9] refactor --- superset/queries/dao.py | 10 +--- superset/sqllab/command.py | 19 +++----- .../sqllab/execution_context_convertor.py | 47 ++++++++++--------- superset/views/core.py | 4 +- 4 files changed, 35 insertions(+), 45 deletions(-) diff --git a/superset/queries/dao.py b/superset/queries/dao.py index 9a072dc0a7c72..97fab40053dbe 100644 --- a/superset/queries/dao.py +++ b/superset/queries/dao.py @@ -52,13 +52,7 @@ def update_saved_query_exec_info(query_id: int) -> None: db.session.commit() @staticmethod - def save_metadata(query: Query, query_payload: str) -> None: - # parse payload - try: - payload: Dict[str, Any] = json.loads(query_payload) - except json.JSONDecodeError: - logger.warning("Error parsing query_payload: %s", query_payload) - return None + def save_metadata(query: Query, payload: Dict[str, Any]) -> None: # pull relevant data from payload and store in json_metadata columns = payload.get("columns", {}) @@ -66,4 +60,4 @@ def save_metadata(query: Query, query_payload: str) -> None: # save payload into query object db.session.add(query) query.set_extra_json_key("columns", columns) - return None + return diff --git a/superset/sqllab/command.py b/superset/sqllab/command.py index f8aa612c0bfe9..9525fa2de6345 100644 --- a/superset/sqllab/command.py +++ b/superset/sqllab/command.py @@ -42,6 +42,8 @@ from superset.sqllab.sql_json_executer import SqlJsonExecutor from superset.sqllab.sqllab_execution_context import SqlJsonExecutionContext +from superset.sqllab.execution_context_convertor import ExecutionContextConvertor + logger = logging.getLogger(__name__) CommandResult = Dict[str, Any] @@ -95,16 +97,18 @@ def run( # pylint: disable=too-many-statements,useless-suppression else: status = self._run_sql_json_exec_from_scratch() - payload = self._execution_context_convertor.to_payload( + self._execution_context_convertor.set_payload( self._execution_context, status ) # save columns into metadata_json - self._query_dao.save_metadata(self._execution_context.query, payload) + self._query_dao.save_metadata( + self._execution_context.query, self._execution_context_convertor.payload + ) return { "status": status, - "payload": payload, + "payload": self._execution_context_convertor.serialize_payload(), } except (SqlLabException, SupersetErrorsException) as ex: raise ex @@ -215,12 +219,3 @@ def validate(self, query: Query) -> None: class SqlQueryRender: def render(self, execution_context: SqlJsonExecutionContext) -> str: raise NotImplementedError() - - -class ExecutionContextConvertor: - def to_payload( - self, - execution_context: SqlJsonExecutionContext, - execution_status: SqlJsonExecutionStatus, - ) -> str: - raise NotImplementedError() diff --git a/superset/sqllab/execution_context_convertor.py b/superset/sqllab/execution_context_convertor.py index 20eabfa7b3f4d..c9be8d3123e9b 100644 --- a/superset/sqllab/execution_context_convertor.py +++ b/superset/sqllab/execution_context_convertor.py @@ -16,52 +16,53 @@ # under the License. from __future__ import annotations -from typing import TYPE_CHECKING +import logging +from typing import Any, Dict, TYPE_CHECKING import simplejson as json import superset.utils.core as utils -from superset.sqllab.command import ExecutionContextConvertor from superset.sqllab.command_status import SqlJsonExecutionStatus from superset.sqllab.utils import apply_display_max_row_configuration_if_require +logger = logging.getLogger(__name__) + if TYPE_CHECKING: from superset.models.sql_lab import Query from superset.sqllab.sql_json_executer import SqlResults from superset.sqllab.sqllab_execution_context import SqlJsonExecutionContext -class ExecutionContextConvertorImpl(ExecutionContextConvertor): +class ExecutionContextConvertor: _max_row_in_display_configuration: int # pylint: disable=invalid-name + _exc_status: SqlJsonExecutionStatus + payload: Dict[str, Any] def set_max_row_in_display(self, value: int) -> None: self._max_row_in_display_configuration = value # pylint: disable=invalid-name - def to_payload( + def set_payload( self, execution_context: SqlJsonExecutionContext, execution_status: SqlJsonExecutionStatus, - ) -> str: - + ) -> None: + self._exc_status = execution_status if execution_status == SqlJsonExecutionStatus.HAS_RESULTS: - return self._to_payload_results_based( - execution_context.get_execution_result() or {} - ) - return self._to_payload_query_based(execution_context.query) + self.payload = execution_context.get_execution_result() or {} + else: + self.payload = execution_context.query.to_dict() - def _to_payload_results_based(self, execution_result: SqlResults) -> str: - return json.dumps( - apply_display_max_row_configuration_if_require( - execution_result, self._max_row_in_display_configuration - ), - default=utils.pessimistic_json_iso_dttm_ser, - ignore_nan=True, - encoding=None, - ) + def serialize_payload(self) -> str: + if self._exc_status == SqlJsonExecutionStatus.HAS_RESULTS: + json.dumps( + apply_display_max_row_configuration_if_require( + self.payload, self._max_row_in_display_configuration + ), + default=utils.pessimistic_json_iso_dttm_ser, + ignore_nan=True, + encoding=None, + ) - def _to_payload_query_based( # pylint: disable=no-self-use - self, query: Query - ) -> str: return json.dumps( - {"query": query.to_dict()}, default=utils.json_int_dttm_ser, ignore_nan=True + {"query": self.payload}, default=utils.json_int_dttm_ser, ignore_nan=True ) diff --git a/superset/views/core.py b/superset/views/core.py index a42596323c65d..15ff3b1620e92 100755 --- a/superset/views/core.py +++ b/superset/views/core.py @@ -113,7 +113,7 @@ QueryIsForbiddenToAccessException, SqlLabException, ) -from superset.sqllab.execution_context_convertor import ExecutionContextConvertorImpl +from superset.sqllab.execution_context_convertor import ExecutionContextConvertor from superset.sqllab.limiting_factor import LimitingFactor from superset.sqllab.query_render import SqlQueryRenderImpl from superset.sqllab.sql_json_executer import ( @@ -2406,7 +2406,7 @@ def _create_sql_json_command( sql_json_executor = Superset._create_sql_json_executor( execution_context, query_dao ) - execution_context_convertor = ExecutionContextConvertorImpl() + execution_context_convertor = ExecutionContextConvertor() execution_context_convertor.set_max_row_in_display( int(config.get("DISPLAY_MAX_ROW")) # type: ignore ) From 6af71ee29a95111c309457d8f60a86b58bfb2fba Mon Sep 17 00:00:00 2001 From: "Hugh A. Miles II" Date: Tue, 17 May 2022 20:41:46 +0000 Subject: [PATCH 9/9] forgot the return --- superset/queries/dao.py | 7 +------ superset/sqllab/command.py | 2 +- superset/sqllab/execution_context_convertor.py | 2 +- tests/unit_tests/dao/queries_test.py | 2 +- 4 files changed, 4 insertions(+), 9 deletions(-) diff --git a/superset/queries/dao.py b/superset/queries/dao.py index 97fab40053dbe..c7d59343e8587 100644 --- a/superset/queries/dao.py +++ b/superset/queries/dao.py @@ -14,7 +14,6 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -import json import logging from datetime import datetime from typing import Any, Dict @@ -53,11 +52,7 @@ def update_saved_query_exec_info(query_id: int) -> None: @staticmethod def save_metadata(query: Query, payload: Dict[str, Any]) -> None: - - # pull relevant data from payload and store in json_metadata + # pull relevant data from payload and store in extra_json columns = payload.get("columns", {}) - - # save payload into query object db.session.add(query) query.set_extra_json_key("columns", columns) - return diff --git a/superset/sqllab/command.py b/superset/sqllab/command.py index 9525fa2de6345..690a8c4c9ebe5 100644 --- a/superset/sqllab/command.py +++ b/superset/sqllab/command.py @@ -34,6 +34,7 @@ QueryIsForbiddenToAccessException, SqlLabException, ) +from superset.sqllab.execution_context_convertor import ExecutionContextConvertor from superset.sqllab.limiting_factor import LimitingFactor if TYPE_CHECKING: @@ -42,7 +43,6 @@ from superset.sqllab.sql_json_executer import SqlJsonExecutor from superset.sqllab.sqllab_execution_context import SqlJsonExecutionContext -from superset.sqllab.execution_context_convertor import ExecutionContextConvertor logger = logging.getLogger(__name__) diff --git a/superset/sqllab/execution_context_convertor.py b/superset/sqllab/execution_context_convertor.py index c9be8d3123e9b..f49fbd9a31db5 100644 --- a/superset/sqllab/execution_context_convertor.py +++ b/superset/sqllab/execution_context_convertor.py @@ -54,7 +54,7 @@ def set_payload( def serialize_payload(self) -> str: if self._exc_status == SqlJsonExecutionStatus.HAS_RESULTS: - json.dumps( + return json.dumps( apply_display_max_row_configuration_if_require( self.payload, self._max_row_in_display_configuration ), diff --git a/tests/unit_tests/dao/queries_test.py b/tests/unit_tests/dao/queries_test.py index 98caad6dfe59e..8df6d2066aaca 100644 --- a/tests/unit_tests/dao/queries_test.py +++ b/tests/unit_tests/dao/queries_test.py @@ -51,5 +51,5 @@ def test_query_dao_save_metadata(app_context: None, session: Session) -> None: from superset.queries.dao import QueryDAO query = session.query(Query).one() - QueryDAO.save_metadata(query=query, query_payload=json.dumps({"columns": []})) + QueryDAO.save_metadata(query=query, payload={"columns": []}) assert query.extra.get("columns", None) == []