Skip to content

Commit

Permalink
fix: missing __init__ on module sqllab commands (#23107)
Browse files Browse the repository at this point in the history
  • Loading branch information
dpgaspar authored Feb 17, 2023
1 parent 7160dae commit 4b03d25
Show file tree
Hide file tree
Showing 6 changed files with 37 additions and 43 deletions.
6 changes: 2 additions & 4 deletions superset/sqllab/api.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down Expand Up @@ -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,
)

Expand Down
16 changes: 16 additions & 0 deletions superset/sqllab/commands/__init__.py
Original file line number Diff line number Diff line change
@@ -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.
8 changes: 3 additions & 5 deletions superset/sqllab/commands/export.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -74,15 +72,15 @@ def validate(self) -> None:

try:
self._query.raise_for_access()
except SupersetSecurityException:
except SupersetSecurityException as ex:
raise SupersetErrorException(
SupersetError(
message=__("Cannot access the query"),
error_type=SupersetErrorType.QUERY_SECURITY_ACCESS_ERROR,
level=ErrorLevel.ERROR,
),
status=403,
)
) from ex

def run(
self,
Expand Down
3 changes: 1 addition & 2 deletions superset/sqllab/commands/results.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion superset/views/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
)

Expand Down
45 changes: 14 additions & 31 deletions tests/integration_tests/sql_lab/commands_tests.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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)
Expand All @@ -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)
Expand All @@ -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):
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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

0 comments on commit 4b03d25

Please sign in to comment.