mirror of
https://github.com/apache/superset.git
synced 2024-09-12 08:39:45 -04:00
367 lines
14 KiB
Python
367 lines
14 KiB
Python
# 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.
|
|
from unittest import mock
|
|
from unittest.mock import Mock, patch
|
|
|
|
import pandas as pd
|
|
import pytest
|
|
from flask_babel import gettext as __
|
|
|
|
from superset import app, db, sql_lab
|
|
from superset.common.db_query_status import QueryStatus
|
|
from superset.errors import ErrorLevel, SupersetError, SupersetErrorType
|
|
from superset.exceptions import (
|
|
SerializationError,
|
|
SupersetErrorException,
|
|
SupersetSecurityException,
|
|
SupersetTimeoutException,
|
|
)
|
|
from superset.models.core import Database
|
|
from superset.models.sql_lab import Query
|
|
from superset.sqllab.commands import estimate, export, results
|
|
from superset.sqllab.limiting_factor import LimitingFactor
|
|
from superset.sqllab.schemas import EstimateQueryCostSchema
|
|
from superset.utils import core as utils
|
|
from superset.utils.database import get_example_database
|
|
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"}
|
|
schema = EstimateQueryCostSchema()
|
|
data: EstimateQueryCostSchema = schema.dump(params)
|
|
command = estimate.QueryEstimationCommand(data)
|
|
|
|
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}}
|
|
schema = EstimateQueryCostSchema()
|
|
data: EstimateQueryCostSchema = schema.dump(params)
|
|
command = estimate.QueryEstimationCommand(data)
|
|
|
|
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"}
|
|
schema = EstimateQueryCostSchema()
|
|
data: EstimateQueryCostSchema = schema.dump(params)
|
|
command = estimate.QueryEstimationCommand(data)
|
|
|
|
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 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()
|
|
|
|
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)
|
|
@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()
|
|
|
|
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)
|
|
@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()
|
|
|
|
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)
|
|
@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()
|
|
|
|
assert result["data"] == "foo\n0\n1\n2\n3\n4\n"
|
|
assert result["count"] == 5
|
|
assert result["query"].client_id == "test"
|
|
|
|
|
|
class TestSqlExecutionResultsCommand(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()
|
|
|
|
@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:
|
|
command = results.SqlExecutionResultsCommand("test", 1000)
|
|
|
|
with pytest.raises(SupersetErrorException) as ex_info:
|
|
command.run()
|
|
assert (
|
|
ex_info.value.error.error_type
|
|
== SupersetErrorType.RESULTS_BACKEND_NOT_CONFIGURED_ERROR
|
|
)
|
|
|
|
@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
|
|
|
|
command = results.SqlExecutionResultsCommand("test", 1000)
|
|
|
|
with pytest.raises(SupersetErrorException) as ex_info:
|
|
command.run()
|
|
assert ex_info.value.error.error_type == SupersetErrorType.RESULTS_BACKEND_ERROR
|
|
|
|
@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,
|
|
"query": {"rows": 100},
|
|
"data": data,
|
|
}
|
|
serialized_payload = sql_lab._serialize_payload(payload, False)
|
|
compressed = utils.zlib_compress(serialized_payload)
|
|
|
|
results.results_backend = mock.Mock()
|
|
results.results_backend.get.return_value = compressed
|
|
|
|
command = results.SqlExecutionResultsCommand("test", 1000)
|
|
|
|
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")
|
|
@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,
|
|
"query": {"rows": 104},
|
|
"data": data,
|
|
}
|
|
serialized_payload = sql_lab._serialize_payload(payload, False)
|
|
compressed = utils.zlib_compress(serialized_payload)
|
|
|
|
results.results_backend = mock.Mock()
|
|
results.results_backend.get.return_value = compressed
|
|
|
|
with mock.patch(
|
|
"superset.views.utils._deserialize_results_payload",
|
|
side_effect=SerializationError(),
|
|
):
|
|
with pytest.raises(SupersetErrorException) as ex_info:
|
|
command = results.SqlExecutionResultsCommand("test_other", 1000)
|
|
command.run()
|
|
assert (
|
|
ex_info.value.error.error_type
|
|
== SupersetErrorType.RESULTS_BACKEND_ERROR
|
|
)
|
|
|
|
@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 = {
|
|
"status": QueryStatus.SUCCESS,
|
|
"query": {"rows": 104},
|
|
"data": data,
|
|
}
|
|
serialized_payload = sql_lab._serialize_payload(payload, False)
|
|
compressed = utils.zlib_compress(serialized_payload)
|
|
|
|
results.results_backend = mock.Mock()
|
|
results.results_backend.get.return_value = compressed
|
|
|
|
command = results.SqlExecutionResultsCommand("abc_query", 1000)
|
|
result = command.run()
|
|
|
|
assert result.get("status") == "success"
|
|
assert result["query"].get("rows") == 104
|
|
assert result.get("data") == data
|