diff --git a/superset/sqllab/api.py b/superset/sqllab/api.py index f73ef749d4936..afe68323fb6b0 100644 --- a/superset/sqllab/api.py +++ b/superset/sqllab/api.py @@ -121,9 +121,7 @@ def export_csv(self, client_id: str) -> CsvResponse: """ result = SqlResultExportCommand(client_id=client_id).run() - query = result.get("query") - data = result.get("data") - row_count = result.get("count") + query, data, row_count = result["query"], result["data"], result["count"] quoted_csv_name = parse.quote(query.name) response = CsvResponse( @@ -292,7 +290,7 @@ def _create_sql_json_command( SqlQueryRenderImpl(get_template_processor), sql_json_executor, execution_context_convertor, - config.get("SQLLAB_CTAS_NO_LIMIT"), + config["SQLLAB_CTAS_NO_LIMIT"], log_params, ) diff --git a/superset/sqllab/commands/__init__.py b/superset/sqllab/commands/__init__.py new file mode 100644 index 0000000000000..13a83393a9124 --- /dev/null +++ b/superset/sqllab/commands/__init__.py @@ -0,0 +1,16 @@ +# 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. diff --git a/superset/sqllab/commands/export.py b/superset/sqllab/commands/export.py index feca664225c7e..e9559be3b97f4 100644 --- a/superset/sqllab/commands/export.py +++ b/superset/sqllab/commands/export.py @@ -14,14 +14,13 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -# pylint: disable=too-few-public-methods, too-many-arguments from __future__ import annotations import logging from typing import Any, cast, List, TypedDict import pandas as pd -from flask_babel import gettext as __, lazy_gettext as _ +from flask_babel import gettext as __ from superset import app, db, results_backend, results_backend_use_msgpack from superset.commands.base import BaseCommand @@ -31,7 +30,6 @@ from superset.sql_parse import ParsedQuery from superset.sqllab.limiting_factor import LimitingFactor from superset.utils import core as utils, csv -from superset.utils.dates import now_as_float from superset.views.utils import _deserialize_results_payload config = app.config @@ -74,7 +72,7 @@ def validate(self) -> None: try: self._query.raise_for_access() - except SupersetSecurityException: + except SupersetSecurityException as ex: raise SupersetErrorException( SupersetError( message=__("Cannot access the query"), @@ -82,7 +80,7 @@ def validate(self) -> None: level=ErrorLevel.ERROR, ), status=403, - ) + ) from ex def run( self, diff --git a/superset/sqllab/commands/results.py b/superset/sqllab/commands/results.py index 9aef5ab4620d7..d6c415a09fef6 100644 --- a/superset/sqllab/commands/results.py +++ b/superset/sqllab/commands/results.py @@ -14,13 +14,12 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -# pylint: disable=too-few-public-methods, too-many-arguments from __future__ import annotations import logging from typing import Any, cast, Dict, Optional -from flask_babel import gettext as __, lazy_gettext as _ +from flask_babel import gettext as __ from superset import app, db, results_backend, results_backend_use_msgpack from superset.commands.base import BaseCommand diff --git a/superset/views/core.py b/superset/views/core.py index cad8483bfda3a..d3dfdb017cb48 100755 --- a/superset/views/core.py +++ b/superset/views/core.py @@ -2363,7 +2363,7 @@ def _create_sql_json_command( SqlQueryRenderImpl(get_template_processor), sql_json_executor, execution_context_convertor, - config.get("SQLLAB_CTAS_NO_LIMIT"), + config["SQLLAB_CTAS_NO_LIMIT"], log_params, ) diff --git a/tests/integration_tests/sql_lab/commands_tests.py b/tests/integration_tests/sql_lab/commands_tests.py index edb71552370b7..cf0aebf001e01 100644 --- a/tests/integration_tests/sql_lab/commands_tests.py +++ b/tests/integration_tests/sql_lab/commands_tests.py @@ -104,13 +104,9 @@ def test_run_no_results_backend_select_sql(self, get_df_mock: Mock) -> None: get_df_mock.return_value = pd.DataFrame({"foo": [1, 2, 3]}) result = command.run() - data = result.get("data") - count = result.get("count") - query = result.get("query") - - assert data == "foo\n1\n2\n3\n" - assert count == 3 - assert query.client_id == "test" + assert result["data"] == "foo\n1\n2\n3\n" + assert result["count"] == 3 + assert result["query"].client_id == "test" @pytest.mark.usefixtures("create_database_and_query") @patch("superset.models.sql_lab.Query.raise_for_access", lambda _: None) @@ -126,13 +122,9 @@ def test_run_no_results_backend_executed_sql(self, get_df_mock: Mock) -> None: get_df_mock.return_value = pd.DataFrame({"foo": [1, 2, 3]}) result = command.run() - data = result.get("data") - count = result.get("count") - query = result.get("query") - - assert data == "foo\n1\n2\n" - assert count == 2 - assert query.client_id == "test" + assert result["data"] == "foo\n1\n2\n" + assert result["count"] == 2 + assert result["query"].client_id == "test" @pytest.mark.usefixtures("create_database_and_query") @patch("superset.models.sql_lab.Query.raise_for_access", lambda _: None) @@ -152,13 +144,9 @@ def test_run_no_results_backend_executed_sql_limiting_factor( result = command.run() - data = result.get("data") - count = result.get("count") - query = result.get("query") - - assert data == "foo\n1\n" - assert count == 1 - assert query.client_id == "test" + assert result["data"] == "foo\n1\n" + assert result["count"] == 1 + assert result["query"].client_id == "test" @pytest.mark.usefixtures("create_database_and_query") @patch("superset.models.sql_lab.Query.raise_for_access", lambda _: None) @@ -179,13 +167,9 @@ def test_run_with_results_backend(self) -> None: result = command.run() - data = result.get("data") - count = result.get("count") - query = result.get("query") - - assert data == "foo\n0\n1\n2\n3\n4\n" - assert count == 5 - assert query.client_id == "test" + assert result["data"] == "foo\n0\n1\n2\n3\n4\n" + assert result["count"] == 5 + assert result["query"].client_id == "test" class TestSqlExecutionResultsCommand(SupersetTestCase): @@ -217,9 +201,8 @@ def create_database_and_query(self): db.session.commit() @patch("superset.sqllab.commands.results.results_backend_use_msgpack", False) + @patch("superset.sqllab.commands.results.results_backend", None) def test_validation_no_results_backend(self) -> None: - results.results_backend = None - command = results.SqlExecutionResultsCommand("test", 1000) with pytest.raises(SupersetErrorException) as ex_info: @@ -306,5 +289,5 @@ def test_run_succeeds(self) -> None: result = command.run() assert result.get("status") == "success" - assert result.get("query").get("rows") == 104 + assert result["query"].get("rows") == 104 assert result.get("data") == data