diff --git a/docs/static/resources/openapi.json b/docs/static/resources/openapi.json index cc92f091e2f9e..ba5e46a4bf14d 100644 --- a/docs/static/resources/openapi.json +++ b/docs/static/resources/openapi.json @@ -1768,7 +1768,7 @@ "type": "string" }, "changed_by": { - "$ref": "#/components/schemas/ChartDataRestApi.get_list.User" + "$ref": "#/components/schemas/ChartDataRestApi.get_list.User3" }, "changed_by_name": { "readOnly": true @@ -1783,7 +1783,7 @@ "readOnly": true }, "created_by": { - "$ref": "#/components/schemas/ChartDataRestApi.get_list.User2" + "$ref": "#/components/schemas/ChartDataRestApi.get_list.User1" }, "created_on_delta_humanized": { "readOnly": true @@ -1830,10 +1830,10 @@ "type": "string" }, "last_saved_by": { - "$ref": "#/components/schemas/ChartDataRestApi.get_list.User3" + "$ref": "#/components/schemas/ChartDataRestApi.get_list.User2" }, "owners": { - "$ref": "#/components/schemas/ChartDataRestApi.get_list.User1" + "$ref": "#/components/schemas/ChartDataRestApi.get_list.User" }, "params": { "nullable": true, @@ -1900,11 +1900,16 @@ "last_name": { "maxLength": 64, "type": "string" + }, + "username": { + "maxLength": 64, + "type": "string" } }, "required": [ "first_name", - "last_name" + "last_name", + "username" ], "type": "object" }, @@ -1921,16 +1926,11 @@ "last_name": { "maxLength": 64, "type": "string" - }, - "username": { - "maxLength": 64, - "type": "string" } }, "required": [ "first_name", - "last_name", - "username" + "last_name" ], "type": "object" }, @@ -1961,10 +1961,6 @@ "maxLength": 64, "type": "string" }, - "id": { - "format": "int32", - "type": "integer" - }, "last_name": { "maxLength": 64, "type": "string" @@ -2560,7 +2556,7 @@ "type": "string" }, "changed_by": { - "$ref": "#/components/schemas/ChartRestApi.get_list.User" + "$ref": "#/components/schemas/ChartRestApi.get_list.User3" }, "changed_by_name": { "readOnly": true @@ -2575,7 +2571,7 @@ "readOnly": true }, "created_by": { - "$ref": "#/components/schemas/ChartRestApi.get_list.User2" + "$ref": "#/components/schemas/ChartRestApi.get_list.User1" }, "created_on_delta_humanized": { "readOnly": true @@ -2622,10 +2618,10 @@ "type": "string" }, "last_saved_by": { - "$ref": "#/components/schemas/ChartRestApi.get_list.User3" + "$ref": "#/components/schemas/ChartRestApi.get_list.User2" }, "owners": { - "$ref": "#/components/schemas/ChartRestApi.get_list.User1" + "$ref": "#/components/schemas/ChartRestApi.get_list.User" }, "params": { "nullable": true, @@ -2692,11 +2688,16 @@ "last_name": { "maxLength": 64, "type": "string" + }, + "username": { + "maxLength": 64, + "type": "string" } }, "required": [ "first_name", - "last_name" + "last_name", + "username" ], "type": "object" }, @@ -2713,16 +2714,11 @@ "last_name": { "maxLength": 64, "type": "string" - }, - "username": { - "maxLength": 64, - "type": "string" } }, "required": [ "first_name", - "last_name", - "username" + "last_name" ], "type": "object" }, @@ -2753,10 +2749,6 @@ "maxLength": 64, "type": "string" }, - "id": { - "format": "int32", - "type": "integer" - }, "last_name": { "maxLength": 64, "type": "string" @@ -3400,7 +3392,7 @@ "type": "string" }, "changed_by": { - "$ref": "#/components/schemas/DashboardRestApi.get_list.User" + "$ref": "#/components/schemas/DashboardRestApi.get_list.User2" }, "changed_by_name": { "readOnly": true @@ -3415,7 +3407,7 @@ "readOnly": true }, "created_by": { - "$ref": "#/components/schemas/DashboardRestApi.get_list.User2" + "$ref": "#/components/schemas/DashboardRestApi.get_list.User1" }, "created_on_delta_humanized": { "readOnly": true @@ -3441,7 +3433,7 @@ "type": "string" }, "owners": { - "$ref": "#/components/schemas/DashboardRestApi.get_list.User1" + "$ref": "#/components/schemas/DashboardRestApi.get_list.User" }, "position_json": { "nullable": true, @@ -3489,6 +3481,10 @@ }, "DashboardRestApi.get_list.User": { "properties": { + "email": { + "maxLength": 64, + "type": "string" + }, "first_name": { "maxLength": 64, "type": "string" @@ -3507,6 +3503,7 @@ } }, "required": [ + "email", "first_name", "last_name", "username" @@ -3515,10 +3512,6 @@ }, "DashboardRestApi.get_list.User1": { "properties": { - "email": { - "maxLength": 64, - "type": "string" - }, "first_name": { "maxLength": 64, "type": "string" @@ -3530,17 +3523,11 @@ "last_name": { "maxLength": 64, "type": "string" - }, - "username": { - "maxLength": 64, - "type": "string" } }, "required": [ - "email", "first_name", - "last_name", - "username" + "last_name" ], "type": "object" }, @@ -3557,11 +3544,16 @@ "last_name": { "maxLength": 64, "type": "string" + }, + "username": { + "maxLength": 64, + "type": "string" } }, "required": [ "first_name", - "last_name" + "last_name", + "username" ], "type": "object" }, @@ -4898,7 +4890,7 @@ "type": "integer" }, "changed_by": { - "$ref": "#/components/schemas/DatasetRestApi.get.User" + "$ref": "#/components/schemas/DatasetRestApi.get.User1" }, "changed_on": { "format": "date-time", @@ -4976,7 +4968,7 @@ "type": "integer" }, "owners": { - "$ref": "#/components/schemas/DatasetRestApi.get.User1" + "$ref": "#/components/schemas/DatasetRestApi.get.User" }, "schema": { "maxLength": 255, @@ -5173,14 +5165,23 @@ "maxLength": 64, "type": "string" }, + "id": { + "format": "int32", + "type": "integer" + }, "last_name": { "maxLength": 64, "type": "string" + }, + "username": { + "maxLength": 64, + "type": "string" } }, "required": [ "first_name", - "last_name" + "last_name", + "username" ], "type": "object" }, @@ -5190,23 +5191,14 @@ "maxLength": 64, "type": "string" }, - "id": { - "format": "int32", - "type": "integer" - }, "last_name": { "maxLength": 64, "type": "string" - }, - "username": { - "maxLength": 64, - "type": "string" } }, "required": [ "first_name", - "last_name", - "username" + "last_name" ], "type": "object" }, @@ -5623,6 +5615,29 @@ }, "type": "object" }, + "EstimateQueryCostSchema": { + "properties": { + "database_id": { + "format": "int32", + "type": "integer" + }, + "schema": { + "nullable": true, + "type": "string" + }, + "sql": { + "type": "string" + }, + "template_params": { + "type": "object" + } + }, + "required": [ + "database_id", + "sql" + ], + "type": "object" + }, "ExecutePayloadSchema": { "properties": { "client_id": { @@ -6950,7 +6965,7 @@ "type": "boolean" }, "changed_by": { - "$ref": "#/components/schemas/ReportScheduleRestApi.get_list.User" + "$ref": "#/components/schemas/ReportScheduleRestApi.get_list.User2" }, "changed_on": { "format": "date-time", @@ -6966,7 +6981,7 @@ "type": "integer" }, "created_by": { - "$ref": "#/components/schemas/ReportScheduleRestApi.get_list.User2" + "$ref": "#/components/schemas/ReportScheduleRestApi.get_list.User1" }, "created_on": { "format": "date-time", @@ -7016,7 +7031,7 @@ "type": "string" }, "owners": { - "$ref": "#/components/schemas/ReportScheduleRestApi.get_list.User1" + "$ref": "#/components/schemas/ReportScheduleRestApi.get_list.User" }, "recipients": { "$ref": "#/components/schemas/ReportScheduleRestApi.get_list.ReportRecipients" @@ -7060,6 +7075,10 @@ "maxLength": 64, "type": "string" }, + "id": { + "format": "int32", + "type": "integer" + }, "last_name": { "maxLength": 64, "type": "string" @@ -7077,10 +7096,6 @@ "maxLength": 64, "type": "string" }, - "id": { - "format": "int32", - "type": "integer" - }, "last_name": { "maxLength": 64, "type": "string" @@ -19936,6 +19951,62 @@ ] } }, + "/api/v1/sqllab/estimate/": { + "post": { + "requestBody": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/EstimateQueryCostSchema" + } + } + }, + "description": "SQL query and params", + "required": true + }, + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "properties": { + "result": { + "type": "object" + } + }, + "type": "object" + } + } + }, + "description": "Query estimation result" + }, + "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": "Estimates the SQL query execution cost", + "tags": [ + "SQL Lab" + ] + } + }, "/api/v1/sqllab/execute/": { "post": { "description": "Starts the execution of a SQL query", diff --git a/superset-frontend/src/SqlLab/actions/sqlLab.js b/superset-frontend/src/SqlLab/actions/sqlLab.js index 40aea66301214..6a5d08358edb7 100644 --- a/superset-frontend/src/SqlLab/actions/sqlLab.js +++ b/superset-frontend/src/SqlLab/actions/sqlLab.js @@ -184,18 +184,20 @@ export function estimateQueryCost(queryEditor) { const { dbId, schema, sql, selectedText, templateParams } = getUpToDateQuery(getState(), queryEditor); const requestSql = selectedText || sql; - const endpoint = - schema === null - ? `/superset/estimate_query_cost/${dbId}/` - : `/superset/estimate_query_cost/${dbId}/${schema}/`; + + const postPayload = { + database_id: dbId, + schema, + sql: requestSql, + template_params: JSON.parse(templateParams || '{}'), + }; + return Promise.all([ dispatch({ type: COST_ESTIMATE_STARTED, query: queryEditor }), SupersetClient.post({ - endpoint, - postPayload: { - sql: requestSql, - templateParams: JSON.parse(templateParams || '{}'), - }, + endpoint: '/api/v1/sqllab/estimate/', + body: JSON.stringify(postPayload), + headers: { 'Content-Type': 'application/json' }, }) .then(({ json }) => dispatch({ type: COST_ESTIMATE_RETURNED, query: queryEditor, json }), diff --git a/superset-frontend/src/SqlLab/reducers/sqlLab.js b/superset-frontend/src/SqlLab/reducers/sqlLab.js index e3bb196fbcef0..a110914b814cf 100644 --- a/superset-frontend/src/SqlLab/reducers/sqlLab.js +++ b/superset-frontend/src/SqlLab/reducers/sqlLab.js @@ -335,7 +335,7 @@ export default function sqlLabReducer(state = {}, action) { ...state.queryCostEstimates, [action.query.id]: { completed: true, - cost: action.json, + cost: action.json.result, error: null, }, }, diff --git a/superset/sqllab/api.py b/superset/sqllab/api.py index 283c3ab638707..d85574eba2569 100644 --- a/superset/sqllab/api.py +++ b/superset/sqllab/api.py @@ -18,7 +18,7 @@ from typing import Any, cast, Dict, Optional import simplejson as json -from flask import request +from flask import request, Response from flask_appbuilder.api import expose, protect, rison from flask_appbuilder.models.sqla.interface import SQLAInterface from marshmallow import ValidationError @@ -31,6 +31,7 @@ from superset.queries.dao import QueryDAO from superset.sql_lab import get_sql_results from superset.sqllab.command_status import SqlJsonExecutionStatus +from superset.sqllab.commands.estimate import QueryEstimationCommand from superset.sqllab.commands.execute import CommandResult, ExecuteSqlCommand from superset.sqllab.commands.results import SqlExecutionResultsCommand from superset.sqllab.exceptions import ( @@ -40,6 +41,7 @@ from superset.sqllab.execution_context_convertor import ExecutionContextConvertor from superset.sqllab.query_render import SqlQueryRenderImpl from superset.sqllab.schemas import ( + EstimateQueryCostSchema, ExecutePayloadSchema, QueryExecutionResponseSchema, sql_lab_get_results_schema, @@ -68,6 +70,7 @@ class SqlLabRestApi(BaseSupersetApi): class_permission_name = "Query" + estimate_model_schema = EstimateQueryCostSchema() execute_model_schema = ExecutePayloadSchema() apispec_parameter_schemas = { @@ -75,10 +78,63 @@ class SqlLabRestApi(BaseSupersetApi): } openapi_spec_tag = "SQL Lab" openapi_spec_component_schemas = ( + EstimateQueryCostSchema, ExecutePayloadSchema, QueryExecutionResponseSchema, ) + @expose("/estimate/", methods=["POST"]) + @protect() + @statsd_metrics + @requires_json + @event_logger.log_this_with_context( + action=lambda self, *args, **kwargs: f"{self.__class__.__name__}" + f".estimate_query_cost", + log_to_statsd=False, + ) + def estimate_query_cost(self) -> Response: + """Estimates the SQL query execution cost + --- + post: + summary: >- + Estimates the SQL query execution cost + requestBody: + description: SQL query and params + required: true + content: + application/json: + schema: + $ref: '#/components/schemas/EstimateQueryCostSchema' + responses: + 200: + description: Query estimation result + content: + application/json: + schema: + type: object + properties: + result: + type: object + 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' + """ + try: + model = self.estimate_model_schema.load(request.json) + except ValidationError as error: + return self.response_400(message=error.messages) + + command = QueryEstimationCommand(model) + result = command.run() + return self.response(200, result=result) + @expose("/results/") @protect() @statsd_metrics diff --git a/superset/sqllab/commands/estimate.py b/superset/sqllab/commands/estimate.py new file mode 100644 index 0000000000000..ee0a084ac60fd --- /dev/null +++ b/superset/sqllab/commands/estimate.py @@ -0,0 +1,106 @@ +# 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, Dict, List + +from flask_babel import gettext as __, lazy_gettext as _ + +from superset import app, db +from superset.commands.base import BaseCommand +from superset.errors import ErrorLevel, SupersetError, SupersetErrorType +from superset.exceptions import SupersetErrorException, SupersetTimeoutException +from superset.jinja_context import get_template_processor +from superset.models.core import Database +from superset.sqllab.schemas import EstimateQueryCostSchema +from superset.utils import core as utils + +config = app.config +SQLLAB_QUERY_COST_ESTIMATE_TIMEOUT = config["SQLLAB_QUERY_COST_ESTIMATE_TIMEOUT"] +stats_logger = config["STATS_LOGGER"] + +logger = logging.getLogger(__name__) + + +class QueryEstimationCommand(BaseCommand): + _database_id: int + _sql: str + _template_params: Dict[str, Any] + _schema: str + _database: Database + + def __init__(self, params: EstimateQueryCostSchema) -> None: + self._database_id = params.get("database_id") + self._sql = params.get("sql", "") + self._template_params = params.get("template_params", {}) + self._schema = params.get("schema", "") + + def validate(self) -> None: + self._database = db.session.query(Database).get(self._database_id) + if not self._database: + raise SupersetErrorException( + SupersetError( + message=__("The database could not be found"), + error_type=SupersetErrorType.RESULTS_BACKEND_ERROR, + level=ErrorLevel.ERROR, + ), + status=404, + ) + + def run( + self, + ) -> List[Dict[str, Any]]: + self.validate() + + sql = self._sql + if self._template_params: + template_processor = get_template_processor(self._database) + sql = template_processor.process_template(sql, **self._template_params) + + timeout = SQLLAB_QUERY_COST_ESTIMATE_TIMEOUT + timeout_msg = f"The estimation exceeded the {timeout} seconds timeout." + try: + with utils.timeout(seconds=timeout, error_message=timeout_msg): + cost = self._database.db_engine_spec.estimate_query_cost( + self._database, self._schema, sql, utils.QuerySource.SQL_LAB + ) + except SupersetTimeoutException as ex: + logger.exception(ex) + raise SupersetErrorException( + SupersetError( + message=__( + "The query estimation was killed after %(sqllab_timeout)s seconds. It might " + "be too complex, or the database might be under heavy load.", + sqllab_timeout=SQLLAB_QUERY_COST_ESTIMATE_TIMEOUT, + ), + error_type=SupersetErrorType.SQLLAB_TIMEOUT_ERROR, + level=ErrorLevel.ERROR, + ), + status=500, + ) + + spec = self._database.db_engine_spec + query_cost_formatters: Dict[str, Any] = app.config[ + "QUERY_COST_FORMATTERS_BY_ENGINE" + ] + query_cost_formatter = query_cost_formatters.get( + spec.engine, spec.query_cost_formatter + ) + cost = query_cost_formatter(cost) + return cost diff --git a/superset/sqllab/schemas.py b/superset/sqllab/schemas.py index f238fda5c918f..134b9ea7bb799 100644 --- a/superset/sqllab/schemas.py +++ b/superset/sqllab/schemas.py @@ -25,6 +25,15 @@ } +class EstimateQueryCostSchema(Schema): + database_id = fields.Integer(required=True, description="The database id") + sql = fields.String(required=True, description="The SQL query to estimate") + template_params = fields.Dict( + keys=fields.String(), description="The SQL query template params" + ) + schema = fields.String(allow_none=True, description="The database schema") + + class ExecutePayloadSchema(Schema): database_id = fields.Integer(required=True) sql = fields.String(required=True) diff --git a/superset/views/core.py b/superset/views/core.py index 94cabfa96c181..f408e9828f7a7 100755 --- a/superset/views/core.py +++ b/superset/views/core.py @@ -2049,6 +2049,7 @@ def extra_table_metadata( # pylint: disable=no-self-use @expose("/estimate_query_cost//", methods=["POST"]) @expose("/estimate_query_cost///", methods=["POST"]) @event_logger.log_this + @deprecated() def estimate_query_cost( # pylint: disable=no-self-use self, database_id: int, schema: Optional[str] = None ) -> FlaskResponse: diff --git a/tests/integration_tests/sql_lab/api_tests.py b/tests/integration_tests/sql_lab/api_tests.py index 4c2080ad4cc2f..8c34e0a79b5af 100644 --- a/tests/integration_tests/sql_lab/api_tests.py +++ b/tests/integration_tests/sql_lab/api_tests.py @@ -39,6 +39,73 @@ class TestSqlLabApi(SupersetTestCase): + def test_estimate_required_params(self): + self.login() + + rv = self.client.post( + "/api/v1/sqllab/estimate/", + json={}, + ) + failed_resp = { + "message": { + "sql": ["Missing data for required field."], + "database_id": ["Missing data for required field."], + } + } + resp_data = json.loads(rv.data.decode("utf-8")) + self.assertDictEqual(resp_data, failed_resp) + self.assertEqual(rv.status_code, 400) + + data = {"sql": "SELECT 1"} + rv = self.client.post( + "/api/v1/sqllab/estimate/", + json=data, + ) + failed_resp = {"message": {"database_id": ["Missing data for required field."]}} + resp_data = json.loads(rv.data.decode("utf-8")) + self.assertDictEqual(resp_data, failed_resp) + self.assertEqual(rv.status_code, 400) + + data = {"database_id": 1} + rv = self.client.post( + "/api/v1/sqllab/estimate/", + json=data, + ) + failed_resp = {"message": {"sql": ["Missing data for required field."]}} + resp_data = json.loads(rv.data.decode("utf-8")) + self.assertDictEqual(resp_data, failed_resp) + self.assertEqual(rv.status_code, 400) + + def test_estimate_valid_request(self): + self.login() + + formatter_response = [ + { + "value": 100, + } + ] + + db_mock = mock.Mock() + db_mock.db_engine_spec = mock.Mock() + db_mock.db_engine_spec.estimate_query_cost = mock.Mock(return_value=100) + db_mock.db_engine_spec.query_cost_formatter = mock.Mock( + return_value=formatter_response + ) + + with mock.patch("superset.sqllab.commands.estimate.db") as mock_superset_db: + mock_superset_db.session.query().get.return_value = db_mock + + data = {"database_id": 1, "sql": "SELECT 1"} + rv = self.client.post( + "/api/v1/sqllab/estimate/", + json=data, + ) + + success_resp = {"result": formatter_response} + resp_data = json.loads(rv.data.decode("utf-8")) + self.assertDictEqual(resp_data, success_resp) + self.assertEqual(rv.status_code, 200) + @mock.patch("superset.sqllab.commands.results.results_backend_use_msgpack", False) def test_execute_required_params(self): self.login() diff --git a/tests/integration_tests/sql_lab/commands_tests.py b/tests/integration_tests/sql_lab/commands_tests.py index 74c1fe7082103..84e227294740d 100644 --- a/tests/integration_tests/sql_lab/commands_tests.py +++ b/tests/integration_tests/sql_lab/commands_tests.py @@ -18,18 +18,88 @@ from unittest.mock import patch import pytest +from flask_babel import gettext as __ -from superset import db, sql_lab +from superset import app, 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, + SupersetErrorException, + SupersetTimeoutException, +) from superset.models.core import Database from superset.models.sql_lab import Query -from superset.sqllab.commands import results +from superset.sqllab.commands import estimate, results from superset.utils import core as utils from tests.integration_tests.base_tests import SupersetTestCase +class TestQueryEstimationCommand(SupersetTestCase): + def test_validation_no_database(self) -> None: + params = {"database_id": 1, "sql": "SELECT 1"} + command = estimate.QueryEstimationCommand(params) + + with mock.patch("superset.sqllab.commands.estimate.db") as mock_superset_db: + mock_superset_db.session.query().get.return_value = None + with pytest.raises(SupersetErrorException) as ex_info: + command.validate() + assert ( + ex_info.value.error.error_type + == SupersetErrorType.RESULTS_BACKEND_ERROR + ) + + @patch("superset.tasks.scheduler.is_feature_enabled") + def test_run_timeout(self, is_feature_enabled) -> None: + params = {"database_id": 1, "sql": "SELECT 1", "template_params": {"temp": 123}} + command = estimate.QueryEstimationCommand(params) + + db_mock = mock.Mock() + db_mock.db_engine_spec = mock.Mock() + db_mock.db_engine_spec.estimate_query_cost = mock.Mock( + side_effect=SupersetTimeoutException( + error_type=SupersetErrorType.CONNECTION_DATABASE_TIMEOUT, + message=( + "Please check your connection details and database settings, " + "and ensure that your database is accepting connections, " + "then try connecting again." + ), + level=ErrorLevel.ERROR, + ) + ) + db_mock.db_engine_spec.query_cost_formatter = mock.Mock(return_value=None) + is_feature_enabled.return_value = False + + with mock.patch("superset.sqllab.commands.estimate.db") as mock_superset_db: + mock_superset_db.session.query().get.return_value = db_mock + with pytest.raises(SupersetErrorException) as ex_info: + command.run() + assert ( + ex_info.value.error.error_type == SupersetErrorType.SQLLAB_TIMEOUT_ERROR + ) + assert ex_info.value.error.message == __( + "The query estimation was killed after %(sqllab_timeout)s seconds. It might " + "be too complex, or the database might be under heavy load.", + sqllab_timeout=app.config["SQLLAB_QUERY_COST_ESTIMATE_TIMEOUT"], + ) + + def test_run_success(self) -> None: + params = {"database_id": 1, "sql": "SELECT 1"} + command = estimate.QueryEstimationCommand(params) + + payload = {"value": 100} + + db_mock = mock.Mock() + db_mock.db_engine_spec = mock.Mock() + db_mock.db_engine_spec.estimate_query_cost = mock.Mock(return_value=100) + db_mock.db_engine_spec.query_cost_formatter = mock.Mock(return_value=payload) + + with mock.patch("superset.sqllab.commands.estimate.db") as mock_superset_db: + mock_superset_db.session.query().get.return_value = db_mock + result = command.run() + assert result == payload + + class TestSqlExecutionResultsCommand(SupersetTestCase): @mock.patch("superset.sqllab.commands.results.results_backend_use_msgpack", False) def test_validation_no_results_backend(self) -> None: