diff --git a/docs/static/resources/openapi.json b/docs/static/resources/openapi.json index cc92f091e2f9e..18ea7a47f8f19 100644 --- a/docs/static/resources/openapi.json +++ b/docs/static/resources/openapi.json @@ -345,7 +345,7 @@ "AnnotationLayerRestApi.get_list": { "properties": { "changed_by": { - "$ref": "#/components/schemas/AnnotationLayerRestApi.get_list.User1" + "$ref": "#/components/schemas/AnnotationLayerRestApi.get_list.User" }, "changed_on": { "format": "date-time", @@ -356,7 +356,7 @@ "readOnly": true }, "created_by": { - "$ref": "#/components/schemas/AnnotationLayerRestApi.get_list.User" + "$ref": "#/components/schemas/AnnotationLayerRestApi.get_list.User1" }, "created_on": { "format": "date-time", @@ -502,13 +502,13 @@ "AnnotationRestApi.get_list": { "properties": { "changed_by": { - "$ref": "#/components/schemas/AnnotationRestApi.get_list.User1" + "$ref": "#/components/schemas/AnnotationRestApi.get_list.User" }, "changed_on_delta_humanized": { "readOnly": true }, "created_by": { - "$ref": "#/components/schemas/AnnotationRestApi.get_list.User" + "$ref": "#/components/schemas/AnnotationRestApi.get_list.User1" }, "end_dttm": { "format": "date-time", @@ -1223,6 +1223,7 @@ "example": false }, "periods": { + "description": "Time periods (in units of `time_grain`) to predict into the future", "example": 7, "format": "int32", "type": "integer" @@ -1578,6 +1579,7 @@ "type": "string" }, "from_dttm": { + "description": "Start timestamp of time range", "format": "int32", "nullable": true, "type": "integer" @@ -1603,6 +1605,7 @@ "type": "integer" }, "stacktrace": { + "description": "Stacktrace if there was an error", "nullable": true, "type": "string" }, @@ -1620,6 +1623,7 @@ "type": "string" }, "to_dttm": { + "description": "End timestamp of time range", "format": "int32", "nullable": true, "type": "integer" @@ -2232,6 +2236,7 @@ "type": "string" }, "rolling_type_options": { + "description": "Optional options to pass to rolling method. Needed for e.g. quantile operation.", "example": {}, "type": "object" }, @@ -3027,13 +3032,13 @@ "CssTemplateRestApi.get_list": { "properties": { "changed_by": { - "$ref": "#/components/schemas/CssTemplateRestApi.get_list.User1" + "$ref": "#/components/schemas/CssTemplateRestApi.get_list.User" }, "changed_on_delta_humanized": { "readOnly": true }, "created_by": { - "$ref": "#/components/schemas/CssTemplateRestApi.get_list.User" + "$ref": "#/components/schemas/CssTemplateRestApi.get_list.User1" }, "created_on": { "format": "date-time", @@ -5230,7 +5235,7 @@ "DatasetRestApi.get_list": { "properties": { "changed_by": { - "$ref": "#/components/schemas/DatasetRestApi.get_list.User1" + "$ref": "#/components/schemas/DatasetRestApi.get_list.User" }, "changed_by_name": { "readOnly": true @@ -5273,7 +5278,7 @@ "readOnly": true }, "owners": { - "$ref": "#/components/schemas/DatasetRestApi.get_list.User" + "$ref": "#/components/schemas/DatasetRestApi.get_list.User1" }, "schema": { "maxLength": 255, @@ -5317,14 +5322,6 @@ "maxLength": 64, "type": "string" }, - "id": { - "format": "int32", - "type": "integer" - }, - "last_name": { - "maxLength": 64, - "type": "string" - }, "username": { "maxLength": 64, "type": "string" @@ -5332,7 +5329,6 @@ }, "required": [ "first_name", - "last_name", "username" ], "type": "object" @@ -5343,6 +5339,14 @@ "maxLength": 64, "type": "string" }, + "id": { + "format": "int32", + "type": "integer" + }, + "last_name": { + "maxLength": 64, + "type": "string" + }, "username": { "maxLength": 64, "type": "string" @@ -5350,6 +5354,7 @@ }, "required": [ "first_name", + "last_name", "username" ], "type": "object" @@ -19997,6 +20002,57 @@ ] } }, + "/api/v1/sqllab/export/{client_id}/": { + "get": { + "parameters": [ + { + "description": "The SQL query result identifier", + "in": "path", + "name": "client_id", + "required": true, + "schema": { + "type": "integer" + } + } + ], + "responses": { + "200": { + "content": { + "text/csv": { + "schema": { + "type": "string" + } + } + }, + "description": "SQL query results" + }, + "400": { + "$ref": "#/components/responses/400" + }, + "401": { + "$ref": "#/components/responses/401" + }, + "403": { + "$ref": "#/components/responses/403" + }, + "404": { + "$ref": "#/components/responses/404" + }, + "500": { + "$ref": "#/components/responses/500" + } + }, + "security": [ + { + "jwt": [] + } + ], + "summary": "Exports the SQL query results to a CSV", + "tags": [ + "SQL Lab" + ] + } + }, "/api/v1/sqllab/results/": { "get": { "parameters": [ diff --git a/superset-frontend/src/SqlLab/components/ResultSet/index.tsx b/superset-frontend/src/SqlLab/components/ResultSet/index.tsx index 62912d66a274d..fad6c98bc94b8 100644 --- a/superset-frontend/src/SqlLab/components/ResultSet/index.tsx +++ b/superset-frontend/src/SqlLab/components/ResultSet/index.tsx @@ -219,6 +219,9 @@ const ResultSet = ({ } }; + const getExportCsvUrl = (clientId: string) => + `/api/v1/sqllab/export/${clientId}/`; + const renderControls = () => { if (search || visualize || csv) { let { data } = query.results; @@ -257,7 +260,7 @@ const ResultSet = ({ /> )} {csv && ( - )} diff --git a/superset/sqllab/api.py b/superset/sqllab/api.py index 283c3ab638707..f73ef749d4936 100644 --- a/superset/sqllab/api.py +++ b/superset/sqllab/api.py @@ -16,6 +16,7 @@ # under the License. import logging from typing import Any, cast, Dict, Optional +from urllib import parse import simplejson as json from flask import request @@ -32,6 +33,7 @@ from superset.sql_lab import get_sql_results from superset.sqllab.command_status import SqlJsonExecutionStatus from superset.sqllab.commands.execute import CommandResult, ExecuteSqlCommand +from superset.sqllab.commands.export import SqlResultExportCommand from superset.sqllab.commands.results import SqlExecutionResultsCommand from superset.sqllab.exceptions import ( QueryIsForbiddenToAccessException, @@ -53,7 +55,7 @@ from superset.sqllab.validators import CanAccessQueryValidatorImpl from superset.superset_typing import FlaskResponse from superset.utils import core as utils -from superset.views.base import json_success +from superset.views.base import CsvResponse, generate_download_headers, json_success from superset.views.base_api import BaseSupersetApi, requires_json, statsd_metrics config = app.config @@ -79,6 +81,69 @@ class SqlLabRestApi(BaseSupersetApi): QueryExecutionResponseSchema, ) + @expose("/export//") + @protect() + @statsd_metrics + @event_logger.log_this_with_context( + action=lambda self, *args, **kwargs: f"{self.__class__.__name__}" + f".export_csv", + log_to_statsd=False, + ) + def export_csv(self, client_id: str) -> CsvResponse: + """Exports the SQL query results to a CSV + --- + get: + summary: >- + Exports the SQL query results to a CSV + parameters: + - in: path + schema: + type: integer + name: client_id + description: The SQL query result identifier + responses: + 200: + description: SQL query results + content: + text/csv: + schema: + type: string + 400: + $ref: '#/components/responses/400' + 401: + $ref: '#/components/responses/401' + 403: + $ref: '#/components/responses/403' + 404: + $ref: '#/components/responses/404' + 500: + $ref: '#/components/responses/500' + """ + result = SqlResultExportCommand(client_id=client_id).run() + + query = result.get("query") + data = result.get("data") + row_count = result.get("count") + + quoted_csv_name = parse.quote(query.name) + response = CsvResponse( + data, headers=generate_download_headers("csv", quoted_csv_name) + ) + event_info = { + "event_type": "data_export", + "client_id": client_id, + "row_count": row_count, + "database": query.database.name, + "schema": query.schema, + "sql": query.sql, + "exported_format": "csv", + } + event_rep = repr(event_info) + logger.debug( + "CSV exported: %s", event_rep, extra={"superset_event": event_info} + ) + return response + @expose("/results/") @protect() @statsd_metrics diff --git a/superset/sqllab/commands/export.py b/superset/sqllab/commands/export.py new file mode 100644 index 0000000000000..feca664225c7e --- /dev/null +++ b/superset/sqllab/commands/export.py @@ -0,0 +1,136 @@ +# 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. +# 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 superset import app, db, results_backend, results_backend_use_msgpack +from superset.commands.base import BaseCommand +from superset.errors import ErrorLevel, SupersetError, SupersetErrorType +from superset.exceptions import SupersetErrorException, SupersetSecurityException +from superset.models.sql_lab import Query +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 + +logger = logging.getLogger(__name__) + + +class SqlExportResult(TypedDict): + query: Query + count: int + data: List[Any] + + +class SqlResultExportCommand(BaseCommand): + _client_id: str + _query: Query + + def __init__( + self, + client_id: str, + ) -> None: + self._client_id = client_id + + def validate(self) -> None: + self._query = ( + db.session.query(Query).filter_by(client_id=self._client_id).one_or_none() + ) + if self._query is None: + raise SupersetErrorException( + SupersetError( + message=__( + "The query associated with these results could not be found. " + "You need to re-run the original query." + ), + error_type=SupersetErrorType.RESULTS_BACKEND_ERROR, + level=ErrorLevel.ERROR, + ), + status=404, + ) + + try: + self._query.raise_for_access() + except SupersetSecurityException: + raise SupersetErrorException( + SupersetError( + message=__("Cannot access the query"), + error_type=SupersetErrorType.QUERY_SECURITY_ACCESS_ERROR, + level=ErrorLevel.ERROR, + ), + status=403, + ) + + def run( + self, + ) -> SqlExportResult: + self.validate() + blob = None + if results_backend and self._query.results_key: + logger.info( + "Fetching CSV from results backend [%s]", self._query.results_key + ) + blob = results_backend.get(self._query.results_key) + if blob: + logger.info("Decompressing") + payload = utils.zlib_decompress( + blob, decode=not results_backend_use_msgpack + ) + obj = _deserialize_results_payload( + payload, self._query, cast(bool, results_backend_use_msgpack) + ) + + df = pd.DataFrame( + data=obj["data"], + dtype=object, + columns=[c["name"] for c in obj["columns"]], + ) + + logger.info("Using pandas to convert to CSV") + else: + logger.info("Running a query to turn into CSV") + if self._query.select_sql: + sql = self._query.select_sql + limit = None + else: + sql = self._query.executed_sql + limit = ParsedQuery(sql).limit + if limit is not None and self._query.limiting_factor in { + LimitingFactor.QUERY, + LimitingFactor.DROPDOWN, + LimitingFactor.QUERY_AND_DROPDOWN, + }: + # remove extra row from `increased_limit` + limit -= 1 + df = self._query.database.get_df(sql, self._query.schema)[:limit] + + csv_data = csv.df_to_escaped_csv(df, index=False, **config["CSV_EXPORT"]) + + return { + "query": self._query, + "count": len(df.index), + "data": csv_data, + } diff --git a/superset/views/core.py b/superset/views/core.py index 94cabfa96c181..07ea47a1bc00d 100755 --- a/superset/views/core.py +++ b/superset/views/core.py @@ -2393,6 +2393,7 @@ def _create_response_from_execution_context( # pylint: disable=invalid-name, no @has_access @event_logger.log_this @expose("/csv/") + @deprecated() def csv(self, client_id: str) -> FlaskResponse: # pylint: disable=no-self-use """Download the query results as csv.""" logger.info("Exporting CSV file [%s]", client_id) diff --git a/tests/integration_tests/sql_lab/api_tests.py b/tests/integration_tests/sql_lab/api_tests.py index 4c2080ad4cc2f..93beb380f0db6 100644 --- a/tests/integration_tests/sql_lab/api_tests.py +++ b/tests/integration_tests/sql_lab/api_tests.py @@ -19,6 +19,9 @@ import datetime import json import random +import csv +import pandas as pd +import io import pytest import prison @@ -26,7 +29,7 @@ from unittest import mock from tests.integration_tests.test_app import app -from superset import sql_lab +from superset import db, sql_lab from superset.common.db_query_status import QueryStatus from superset.models.core import Database from superset.utils.database import get_example_database, get_main_database @@ -176,3 +179,37 @@ def test_get_results_with_display_limit(self): self.assertEqual(result_limited, expected_limited) app.config["RESULTS_BACKEND_USE_MSGPACK"] = use_msgpack + + @mock.patch("superset.models.sql_lab.Query.raise_for_access", lambda _: None) + @mock.patch("superset.models.core.Database.get_df") + def test_export_results(self, get_df_mock: mock.Mock) -> None: + self.login() + + database = get_example_database() + query_obj = Query( + client_id="test", + database=database, + tab_name="test_tab", + sql_editor_id="test_editor_id", + sql="select * from bar", + select_sql=None, + executed_sql="select * from bar limit 2", + limit=100, + select_as_cta=False, + rows=104, + error_message="none", + results_key="test_abc", + ) + + db.session.add(query_obj) + db.session.commit() + + get_df_mock.return_value = pd.DataFrame({"foo": [1, 2, 3]}) + + resp = self.get_resp("/api/v1/sqllab/export/test/") + data = csv.reader(io.StringIO(resp)) + expected_data = csv.reader(io.StringIO("foo\n1\n2")) + + self.assertEqual(list(expected_data), list(data)) + db.session.delete(query_obj) + db.session.commit() diff --git a/tests/integration_tests/sql_lab/commands_tests.py b/tests/integration_tests/sql_lab/commands_tests.py index 74c1fe7082103..edb71552370b7 100644 --- a/tests/integration_tests/sql_lab/commands_tests.py +++ b/tests/integration_tests/sql_lab/commands_tests.py @@ -15,23 +15,208 @@ # specific language governing permissions and limitations # under the License. from unittest import mock, skip -from unittest.mock import patch +from unittest.mock import Mock, patch +import pandas as pd import pytest from superset import db, sql_lab from superset.common.db_query_status import QueryStatus -from superset.errors import SupersetErrorType -from superset.exceptions import SerializationError, SupersetErrorException +from superset.errors import ErrorLevel, SupersetErrorType +from superset.exceptions import ( + SerializationError, + SupersetError, + SupersetErrorException, + SupersetSecurityException, +) from superset.models.core import Database from superset.models.sql_lab import Query -from superset.sqllab.commands import results +from superset.sqllab.commands import export, results +from superset.sqllab.limiting_factor import LimitingFactor from superset.utils import core as utils +from superset.utils.database import get_example_database from tests.integration_tests.base_tests import SupersetTestCase +class TestSqlResultExportCommand(SupersetTestCase): + @pytest.fixture() + def create_database_and_query(self): + with self.create_app().app_context(): + database = get_example_database() + query_obj = Query( + client_id="test", + database=database, + 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=104, + error_message="none", + results_key="abc_query", + ) + + db.session.add(query_obj) + db.session.commit() + + yield + + db.session.delete(query_obj) + db.session.commit() + + @pytest.mark.usefixtures("create_database_and_query") + def test_validation_query_not_found(self) -> None: + command = export.SqlResultExportCommand("asdf") + + with pytest.raises(SupersetErrorException) as ex_info: + command.run() + assert ex_info.value.error.error_type == SupersetErrorType.RESULTS_BACKEND_ERROR + + @pytest.mark.usefixtures("create_database_and_query") + def test_validation_invalid_access(self) -> None: + command = export.SqlResultExportCommand("test") + + with mock.patch( + "superset.security_manager.raise_for_access", + side_effect=SupersetSecurityException( + SupersetError( + "dummy", + SupersetErrorType.DATASOURCE_SECURITY_ACCESS_ERROR, + ErrorLevel.ERROR, + ) + ), + ): + with pytest.raises(SupersetErrorException) as ex_info: + command.run() + assert ( + ex_info.value.error.error_type + == SupersetErrorType.QUERY_SECURITY_ACCESS_ERROR + ) + + @pytest.mark.usefixtures("create_database_and_query") + @patch("superset.models.sql_lab.Query.raise_for_access", lambda _: None) + @patch("superset.models.core.Database.get_df") + def test_run_no_results_backend_select_sql(self, get_df_mock: Mock) -> None: + command = export.SqlResultExportCommand("test") + + 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" + + @pytest.mark.usefixtures("create_database_and_query") + @patch("superset.models.sql_lab.Query.raise_for_access", lambda _: None) + @patch("superset.models.core.Database.get_df") + def test_run_no_results_backend_executed_sql(self, get_df_mock: Mock) -> None: + query_obj = db.session.query(Query).filter_by(client_id="test").one() + query_obj.executed_sql = "select * from bar limit 2" + query_obj.select_sql = None + db.session.commit() + + command = export.SqlResultExportCommand("test") + + 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" + + @pytest.mark.usefixtures("create_database_and_query") + @patch("superset.models.sql_lab.Query.raise_for_access", lambda _: None) + @patch("superset.models.core.Database.get_df") + def test_run_no_results_backend_executed_sql_limiting_factor( + self, get_df_mock: Mock + ) -> None: + query_obj = db.session.query(Query).filter_by(results_key="abc_query").one() + query_obj.executed_sql = "select * from bar limit 2" + query_obj.select_sql = None + query_obj.limiting_factor = LimitingFactor.DROPDOWN + db.session.commit() + + command = export.SqlResultExportCommand("test") + + 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\n" + assert count == 1 + assert query.client_id == "test" + + @pytest.mark.usefixtures("create_database_and_query") + @patch("superset.models.sql_lab.Query.raise_for_access", lambda _: None) + @patch("superset.sqllab.commands.export.results_backend_use_msgpack", False) + def test_run_with_results_backend(self) -> None: + command = export.SqlResultExportCommand("test") + + data = [{"foo": i} for i in range(5)] + payload = { + "columns": [{"name": "foo"}], + "data": data, + } + serialized_payload = sql_lab._serialize_payload(payload, False) + compressed = utils.zlib_compress(serialized_payload) + + export.results_backend = mock.Mock() + export.results_backend.get.return_value = compressed + + 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" + + class TestSqlExecutionResultsCommand(SupersetTestCase): - @mock.patch("superset.sqllab.commands.results.results_backend_use_msgpack", False) + @pytest.fixture() + def create_database_and_query(self): + with self.create_app().app_context(): + database = get_example_database() + query_obj = Query( + client_id="test", + database=database, + 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=104, + error_message="none", + results_key="abc_query", + ) + + db.session.add(query_obj) + db.session.commit() + + yield + + db.session.delete(query_obj) + db.session.commit() + + @patch("superset.sqllab.commands.results.results_backend_use_msgpack", False) def test_validation_no_results_backend(self) -> None: results.results_backend = None @@ -44,7 +229,7 @@ def test_validation_no_results_backend(self) -> None: == SupersetErrorType.RESULTS_BACKEND_NOT_CONFIGURED_ERROR ) - @mock.patch("superset.sqllab.commands.results.results_backend_use_msgpack", False) + @patch("superset.sqllab.commands.results.results_backend_use_msgpack", False) def test_validation_data_cannot_be_retrieved(self) -> None: results.results_backend = mock.Mock() results.results_backend.get.return_value = None @@ -55,8 +240,8 @@ def test_validation_data_cannot_be_retrieved(self) -> None: command.run() assert ex_info.value.error.error_type == SupersetErrorType.RESULTS_BACKEND_ERROR - @mock.patch("superset.sqllab.commands.results.results_backend_use_msgpack", False) - def test_validation_query_not_found(self) -> None: + @patch("superset.sqllab.commands.results.results_backend_use_msgpack", False) + def test_validation_data_not_found(self) -> None: data = [{"col_0": i} for i in range(100)] payload = { "status": QueryStatus.SUCCESS, @@ -75,8 +260,9 @@ def test_validation_query_not_found(self) -> None: command.run() assert ex_info.value.error.error_type == SupersetErrorType.RESULTS_BACKEND_ERROR - @mock.patch("superset.sqllab.commands.results.results_backend_use_msgpack", False) - def test_validation_query_not_found2(self) -> None: + @pytest.mark.usefixtures("create_database_and_query") + @patch("superset.sqllab.commands.results.results_backend_use_msgpack", False) + def test_validation_query_not_found(self) -> None: data = [{"col_0": i} for i in range(104)] payload = { "status": QueryStatus.SUCCESS, @@ -89,38 +275,20 @@ def test_validation_query_not_found2(self) -> None: results.results_backend = mock.Mock() results.results_backend.get.return_value = compressed - database = Database(database_name="my_database", sqlalchemy_uri="sqlite://") - query_obj = Query( - client_id="foo", - database=database, - 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=104, - error_message="none", - results_key="test_abc", - ) - - db.session.add(database) - db.session.add(query_obj) - with mock.patch( "superset.views.utils._deserialize_results_payload", side_effect=SerializationError(), ): with pytest.raises(SupersetErrorException) as ex_info: - command = results.SqlExecutionResultsCommand("test", 1000) + command = results.SqlExecutionResultsCommand("test_other", 1000) command.run() assert ( ex_info.value.error.error_type == SupersetErrorType.RESULTS_BACKEND_ERROR ) - @mock.patch("superset.sqllab.commands.results.results_backend_use_msgpack", False) + @pytest.mark.usefixtures("create_database_and_query") + @patch("superset.sqllab.commands.results.results_backend_use_msgpack", False) def test_run_succeeds(self) -> None: data = [{"col_0": i} for i in range(104)] payload = { @@ -134,26 +302,7 @@ def test_run_succeeds(self) -> None: results.results_backend = mock.Mock() results.results_backend.get.return_value = compressed - database = Database(database_name="my_database", sqlalchemy_uri="sqlite://") - query_obj = Query( - client_id="foo", - database=database, - 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=104, - error_message="none", - results_key="test_abc", - ) - - db.session.add(database) - db.session.add(query_obj) - - command = results.SqlExecutionResultsCommand("test_abc", 1000) + command = results.SqlExecutionResultsCommand("abc_query", 1000) result = command.run() assert result.get("status") == "success"