AlankritVerma01 commented on code in PR #45313:
URL: https://github.com/apache/superset/pull/45313#discussion_r4420530261
##########
tests/unit_tests/db_engine_specs/test_bigquery.py:
##########
@@ -1598,3 +1621,290 @@ def test_expand_nested_columns_keeps_user_alias() ->
None:
"calendar.day",
]
assert result[2] is columns[1]
+
+
[email protected](
+ "value,accuracy,row_security,expected",
+ [
+ ("0", "PRECISE", False, 0),
+ ("100", "PRECISE", False, 100),
+ (str(2**60 + 1), "PRECISE", False, 2**60 + 1),
+ ("101", "LOWER_BOUND", False, 101),
+ ("100", "LOWER_BOUND", False, 100),
+ ("101", "UPPER_BOUND", False, 101),
+ ("100", "UPPER_BOUND", False, 100),
+ ("0", "UNKNOWN", False, 0),
+ ("101", None, False, 101),
+ ("101", "PRECISE", True, None),
+ (None, "PRECISE", False, None),
+ ("masked", "PRECISE", False, None),
+ ("1.5", "PRECISE", False, None),
+ (True, "PRECISE", False, None),
+ (-1, "PRECISE", False, None),
+ pytest.param("9" * 5000, "PRECISE", False, None,
id="exceeds-int-digit-limit"),
+ ],
+)
+def test_warning_dry_run_uses_execution_client_and_raw_statistics(
+ mocker: MockerFixture,
+ value: object,
+ accuracy: str | None,
+ row_security: bool,
+ expected: int | None,
+) -> None:
+ """Exercise the SDK and driver's real connection shapes without a
warehouse."""
+ from google.auth.credentials import AnonymousCredentials
+ from google.cloud.bigquery import Client, QueryJob
+ from sqlalchemy import create_engine
+
+ from superset.db_engine_specs.bigquery import BigQueryEngineSpec
+ from superset.utils.core import QuerySource
+
+ assert BigQueryEngineSpec.get_query_cost_warning_metric() ==
"bytes_processed"
+
+ # Patch only credential discovery; let the installed driver construct its
client.
+ mocker.patch(
+ "google.auth.default", return_value=(AnonymousCredentials(), "project")
+ )
+ engine = create_engine(
+
"bigquery://project/dataset?location=EU&priority=BATCH&maximum_bytes_billed=10000"
+ )
+ database = mocker.MagicMock()
+ database.mutate_sql_based_on_config.side_effect = lambda sql, **kwargs: sql
+ connection = engine.raw_connection()
+ database.get_raw_connection.return_value.__enter__.return_value =
connection
+ client = connection.driver_connection._client
+ assert isinstance(client, Client)
+ job = QueryJob.from_api_repr(
+ {
+ "jobReference": {
+ "jobId": "dry-run",
+ "projectId": "project",
+ "location": "EU",
+ },
+ "configuration": {"query": {"query": "SELECT 1"}, "dryRun": True},
+ "statistics": {
+ "query": {
+ "totalBytesProcessed": value,
+ "totalBytesProcessedAccuracy": accuracy,
+ },
+ "rowLevelSecurityStatistics": {"rowLevelSecurityApplied":
row_security},
+ },
+ },
+ client,
+ )
+ api_call = mocker.patch.object(client, "_call_api",
return_value=job._properties)
+ query = mocker.spy(client, "query")
+ try:
+ estimate = BigQueryEngineSpec.estimate_query_cost_for_warning(
+ database, "project", "dataset", "SELECT 1"
+ )
+ query.assert_called_once()
+ assert estimate.accuracy == (None if row_security else accuracy)
+ assert bool(estimate.unavailable_reason) is row_security
+ assert estimate.value == expected
+ if row_security:
+ assert estimate.unavailable_reason is not None
+ assert "row-level security" in estimate.unavailable_reason
+ if value == str(2**60 + 1):
+ assert estimate.value == 2**60 + 1
+ config = query.call_args.kwargs["job_config"]
+ assert config.dry_run is True
+ submitted = api_call.call_args.kwargs["data"]
+ assert submitted["configuration"]["dryRun"] is True
+ assert submitted["configuration"]["query"]["defaultDataset"] == {
+ "projectId": "project",
+ "datasetId": "dataset",
+ }
+ assert submitted["configuration"]["query"]["priority"] == "BATCH"
+ assert submitted["configuration"]["query"]["maximumBytesBilled"] ==
"10000"
+ assert submitted["jobReference"]["location"] == "EU"
+ assert query.call_args.kwargs["job_retry"] is None
+ assert client.project == "project"
+ assert client.location == "EU"
+ defaults = client.default_query_job_config
+ assert defaults.default_dataset.dataset_id == "dataset"
+ assert defaults.priority == "BATCH"
+ assert defaults.maximum_bytes_billed == 10000
+ database.get_raw_connection.assert_called_once_with(
+ catalog="project", schema="dataset", source=QuerySource.SQL_LAB
+ )
+ finally:
+ connection.close()
+ engine.dispose()
+
+
[email protected](
+ "sql,expected",
+ [
+ (
+ "SELECT PARSE_DATE('%%Y-%%m-%%d', '2026-01-01')",
+ "SELECT PARSE_DATE('%Y-%m-%d', '2026-01-01')",
+ ),
+ ("SELECT '50%'", "SELECT '50%'"),
+ ("SELECT '%%%%'", "SELECT '%%'"),
+ ("SELECT 1", "SELECT 1"),
+ ],
+)
+def test_warning_sql_matches_dbapi_execution_formatting(
+ mocker: MockerFixture, sql: str, expected: str
+) -> None:
+ """Direct dry runs match the real driver's single percent-unescaping
pass."""
+ from google.auth.credentials import AnonymousCredentials
+ from google.cloud.bigquery import Client, QueryJob
+ from google.cloud.bigquery.dbapi import Connection
+
+ from superset.db_engine_specs.base import BaseEngineSpec
+ from superset.db_engine_specs.bigquery import BigQueryEngineSpec
+
+ client = Client(project="project", credentials=AnonymousCredentials())
+ connection = Connection(client=client, prefer_bqstorage_client=False)
+ database = mocker.MagicMock()
+ database.mutate_sql_based_on_config.side_effect = lambda sql, **kwargs: sql
+ database.get_raw_connection.return_value.__enter__.return_value =
SimpleNamespace(
+ driver_connection=connection
+ )
+ job = QueryJob("dry-run", sql, client)
+ api_call = mocker.patch.object(client, "_call_api",
return_value=job._properties)
+ process = mocker.spy(BigQueryEngineSpec, "process_statement")
+ try:
+ BigQueryEngineSpec.estimate_query_cost_for_warning(
+ database, "project", "dataset", sql
+ )
+ process.assert_called_once()
+ api_call.assert_called_once()
+ warning_request = api_call.call_args.kwargs["data"]
+ assert warning_request["configuration"]["dryRun"] is True
+
+ # Supply completed REST responses for either jobs.query or jobs.insert.
+ api_call.reset_mock()
+ api_call.return_value = {
+ "jobReference": {"projectId": "project", "jobId": "execution"},
+ "configuration": {"query": {"query": expected}},
+ "status": {"state": "DONE"},
+ "jobComplete": True,
+ "schema": {"fields": []},
+ "totalRows": "0",
+ "rows": [],
+ }
+ cursor = connection.cursor()
+ execute = mocker.spy(cursor, "execute")
+ BaseEngineSpec.execute(cursor, process.spy_return, database)
+ execute.assert_called_once_with(process.spy_return)
+ execution_request = api_call.call_args_list[0].kwargs["data"]
+ execution_sql = (
+ execution_request["configuration"]["query"]["query"]
+ if "configuration" in execution_request
+ else execution_request["query"]
+ )
+ assert warning_request["configuration"]["query"]["query"] ==
execution_sql
+ assert " ".join(execution_sql.split()) == expected
+ finally:
+ connection.close()
+ client.close()
+
+
+def test_warning_dry_run_unverified_bigquery_subclass_is_unsupported(
+ mocker: MockerFixture,
+) -> None:
+ """A subclass must explicitly verify warning support before opting in."""
+ from superset.db_engine_specs.bigquery import BigQueryEngineSpec
+
+ class UnverifiedBigQuery(BigQueryEngineSpec):
+ """A connector with unverified dry-run semantics."""
+
+ database = mocker.MagicMock()
+ assert UnverifiedBigQuery.get_query_cost_warning_metric() is None
+ assert UnverifiedBigQuery.get_query_warning_check() is None
+ assert (
+ UnverifiedBigQuery.estimate_query_cost_for_warning(
+ database, None, "", "SELECT 1"
+ ).value
+ is None
+ )
+ database.get_raw_connection.assert_not_called()
+
+
[email protected](
+ "statistics",
+ [
+ None,
+ [],
+ {},
+ {"query": []},
+ {"query": {"totalBytesProcessed": "1"}},
+ {"rowLevelSecurityStatistics": "masked"},
+ {"rowLevelSecurityStatistics": {"rowLevelSecurityApplied": "false"}},
Review Comment:
Thanks, fixed. The case now includes a valid byte count and PRECISE
accuracy, so it specifically tests the row-security guard. I also fixed a few
other test issues: mock leakage between tests, cancellation assertions that ran
too early, and malformed-response cases that weren't testing the intended guard.
##########
tests/unit_tests/commands/sql_lab/test_query_cost_warning.py:
##########
@@ -0,0 +1,899 @@
+# 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.
+"""Warning checks must describe execution SQL without hiding errors or
bounds."""
+
+import os
+import subprocess
+import sys
+from pathlib import Path
+from textwrap import dedent
+from typing import Any, cast
+from unittest.mock import MagicMock
+
+import pytest
+from flask import Flask
+from pytest_mock import MockerFixture
+
+from superset import config
+from superset.commands.sql_lab.estimate import (
+ EstimateQueryCostType,
+ QueryEstimationCommand,
+)
+from superset.commands.sql_lab.execute import ExecuteSqlCommand
+from superset.db_engine_specs.base import BaseEngineSpec
+from superset.db_engine_specs.query_warning import (
+ QueryCostEstimate,
+ QueryCostThresholdCheck,
+)
+from superset.errors import ErrorLevel, SupersetError, SupersetErrorType
+from superset.exceptions import (
+ OAuth2RedirectError,
+ SupersetDisallowedSQLFunctionException,
+ SupersetDisallowedSQLTableException,
+ SupersetDMLNotAllowedException,
+ SupersetErrorException,
+ SupersetParseError,
+ SupersetSecurityException,
+ SupersetTimeoutException,
+)
+from superset.models.core import Database
+from superset.models.sql_lab import Query
+from superset.sql.parse import SQLScript
+from superset.sql_lab import apply_limit
+from superset.sqllab.sqllab_execution_context import SqlJsonExecutionContext
+from superset.utils import json
+
+
+class WarningTestEngineSpec(BaseEngineSpec):
+ """Exercise a verified metric independently of its production engine
adapters."""
+
+ engine = "sqlite"
+
+ @classmethod
+ def get_allow_cost_estimate(cls, extra: dict[str, Any]) -> bool:
+ """Allow estimates when the connection has also opted in."""
+ return True
+
+ @classmethod
+ def get_query_warning_check(cls) -> QueryCostThresholdCheck | None:
+ """Opt this test engine into the shared numeric check."""
+ return QueryCostThresholdCheck("total_cost")
+
+ @classmethod
+ def estimate_query_cost_for_warning(
+ cls, database: Database, catalog: str | None, schema: str, sql: str
+ ) -> QueryCostEstimate:
+ """Return an unrounded decimal estimate without contacting a
database."""
+ return QueryCostEstimate.from_value(10.25)
+
+
+def _make_command(
+ mocker: MockerFixture, engine: str = "postgresql", **params: Any
+) -> QueryEstimationCommand:
+ """Use real engine specs, SQL parsing, rule validation and execution
limits."""
+ metric = "bytes_processed" if engine == "bigquery" else "total_cost"
+ database = Database(
+ id=1,
+ database_name="warning_test",
+ sqlalchemy_uri=f"{engine}://localhost/test",
+ extra=json.dumps(
+ {
+ "cost_estimate_enabled": True,
+ "query_cost_warning": {
+ "enabled": True,
+ "metric": metric,
+ "threshold": 10,
+ },
+ }
+ ),
+ allow_dml=False,
+ )
+ parameters: dict[str, Any] = {
+ "database_id": 1,
+ "sql": "SELECT * FROM sales",
+ "catalog": "selected_catalog",
+ "schema": "selected_schema",
+ "template_params": {},
+ "query_limit": 100,
+ "check_warning": True,
+ **params,
+ }
+ command = QueryEstimationCommand(cast(EstimateQueryCostType, parameters))
+ mocker.patch(
+ "superset.commands.sql_lab.estimate.DatabaseDAO.find_by_id",
+ return_value=database,
+ )
+ mocker.patch(
+ "superset.commands.sql_lab.estimate.security_manager",
new_callable=MagicMock
+ )
+ mocker.patch(
+ "superset.commands.sql_lab.estimate.is_feature_enabled",
+ side_effect=lambda name: name
+ in {"SQLLAB_QUERY_COST_WARNINGS", "ESTIMATE_QUERY_COST"},
+ )
+ mocker.patch(
+ "superset.db_engine_specs.query_warning.is_feature_enabled",
return_value=True
+ )
+ mocker.patch.object(database, "get_default_catalog",
return_value="default_catalog")
+ mocker.patch.object(
+ database.db_engine_spec,
+ "estimate_query_cost_for_warning",
+ return_value=QueryCostEstimate.from_value(5),
+ )
+ return command
+
+
+def _set_extra(command: QueryEstimationCommand, extra: dict[str, Any]) -> None:
+ """Replace raw saved settings so invalid data cannot use the display
fallback."""
+ command.validate()
+ command._database.extra = json.dumps(extra)
+
+
+def _estimator(command: QueryEstimationCommand) -> MagicMock:
+ """Type the engine hook replaced by the test setup's synchronous mock."""
+ return cast(
+ MagicMock,
command._database.db_engine_spec.estimate_query_cost_for_warning
+ )
+
+
+def _run_warning(command: QueryEstimationCommand) -> dict[str, Any]:
+ """Narrow the endpoint's manual-or-warning result to the requested
warning."""
+ result = command.run()
+ assert isinstance(result, dict)
+ return dict(result)
+
+
[email protected]("engine", ["bigquery", "postgresql"])
[email protected](
+ "value,status",
+ [(0, "within_threshold"), (10, "within_threshold"), (11, "exceeded")],
+)
+def test_warning_comparison(
+ mocker: MockerFixture, engine: str, value: int, status: str
+) -> None:
+ """Compare raw native values with strict greater-than, including valid
zero."""
+ command = _make_command(mocker, engine)
+ command.validate()
+ _estimator(command).return_value = QueryCostEstimate.from_value(value)
+ result = _run_warning(command)
+ if status == "within_threshold":
+ assert result == {"action": "proceed"}
+ else:
+ assert result["action"] == "confirm"
+ assert result["reason"] == "warning"
+ assert result["message"]
+ unit = "bytes" if engine == "bigquery" else "planner cost units"
+ assert result["details"][0]["value"] == f"{value:,} {unit}"
+ assert result["details"][1] == {
+ "label": "Warning threshold",
+ "value": f"10 {unit}",
+ }
+
+
+def test_warning_compares_postgres_decimals_before_display(
+ mocker: MockerFixture,
+) -> None:
+ """A planner estimate just above a decimal threshold must warn before
rounding."""
+ command = _make_command(mocker)
+ _set_extra(
+ command,
+ {
+ "cost_estimate_enabled": True,
+ "query_cost_warning": {
+ "enabled": True,
+ "metric": "total_cost",
+ "threshold": 10.125,
+ },
+ },
+ )
+ _estimator(command).return_value = QueryCostEstimate.from_value(10.12501)
+ result = _run_warning(command)
+ assert result["reason"] == "warning"
+ assert result["details"][0]["value"] == "10.12501 planner cost units"
+ assert result["details"][1]["value"] == "10.125 planner cost units"
+
+
[email protected]("metric", ["total_cost", "unknown", ""])
+def test_warning_command_uses_metric_contract_for_test_engine(
+ mocker: MockerFixture, metric: str
+) -> None:
+ """Known metrics work without engine-name branches; unknown ones never
estimate."""
+ command = _make_command(mocker, "sqlite")
+ mocker.patch.object(
+ Database,
+ "db_engine_spec",
+ new_callable=mocker.PropertyMock,
+ return_value=WarningTestEngineSpec,
+ )
+ if metric != "total_cost":
+ mocker.patch.object(
+ WarningTestEngineSpec,
+ "get_query_warning_check",
+ return_value=QueryCostThresholdCheck(metric),
+ )
+ estimator = mocker.spy(WarningTestEngineSpec,
"estimate_query_cost_for_warning")
+ if metric != "total_cost":
+ with pytest.raises(SupersetErrorException) as error:
+ _run_warning(command)
+ assert error.value.status == 400
+ estimator.assert_not_called()
+ else:
+ result = _run_warning(command)
+ assert result["action"] == "confirm"
+ assert result["reason"] == "warning"
+ assert result["details"] == [
+ {"label": "Estimated planner cost", "value": "10.25 planner cost
units"},
+ {"label": "Warning threshold", "value": "10 planner cost units"},
+ ]
+ estimator.assert_called_once()
+
+
[email protected]("enabled", [True, False])
+def test_warning_command_delegates_to_engine_check(
+ mocker: MockerFixture, client: Any, full_api_access: None, enabled: bool
+) -> None:
+ """The protected API uses a message-only check with manual estimation
off."""
+ from tests.fixtures.query_warning import MessageOnlyQueryWarningCheck
+
+ command = _make_command(mocker)
+ rule = {"enabled": enabled}
+ _set_extra(command, {"query_cost_warning": rule, "cost_estimate_enabled":
False})
+ check = MessageOnlyQueryWarningCheck()
+ validate = mocker.spy(check, "validate")
+ assess = mocker.spy(check, "assess")
+ mocker.patch(
+ "superset.commands.sql_lab.estimate.is_feature_enabled",
+ side_effect=lambda name: name == "SQLLAB_QUERY_COST_WARNINGS",
+ )
+ mocker.patch(
+ "superset.db_engine_specs.query_warning.is_feature_enabled",
return_value=False
+ )
+ mocker.patch.object(
+ command._database.db_engine_spec, "get_query_warning_check",
return_value=check
+ )
+ response = client.post(
+ "/api/v1/sqllab/estimate/",
+ json={
+ "database_id": 1,
+ "sql": "SELECT * FROM sales",
+ "catalog": "selected_catalog",
+ "schema": "selected_schema",
+ "query_limit": 100,
+ "check_warning": True,
+ },
+ )
+ assert response.status_code == 200
+ assert response.json == {
+ "result": (
+ {"action": "confirm", "reason": "warning", "message": "Review this
query."}
+ if enabled
+ else {"action": "proceed"}
+ )
+ }
+ validate.assert_called_once_with(rule)
+ assert command._database.allows_cost_estimate is False
+ if enabled:
+ assess.assert_called_once()
+ database, catalog, schema, sql, assessed_rule = assess.call_args.args
+ assert database is command._database
+ assert catalog == "selected_catalog"
+ assert schema == "selected_schema"
+ assert (
+ SQLScript(sql,
engine="postgresql").statements[0].get_limit_value() == 101
+ )
+ assert assessed_rule == rule
+ else:
+ assess.assert_not_called()
+ _estimator(command).assert_not_called()
+
+
[email protected](
+ "value,accuracy,status",
+ [
+ (10, "LOWER_BOUND", "unavailable"),
+ (9, "LOWER_BOUND", "unavailable"),
+ (11, "LOWER_BOUND", "exceeded"),
+ (10, "UPPER_BOUND", "within_threshold"),
+ (11, "UPPER_BOUND", "exceeded"),
+ (0, "UNKNOWN", "unavailable"),
+ (11, None, "unavailable"),
+ (2**63 - 1, "PRECISE", "exceeded"),
+ ],
+)
+def test_warning_preserves_bigquery_bounds(
+ mocker: MockerFixture, value: int, accuracy: str | None, status: str
+) -> None:
+ """Unknown bounds never imply safety; usable bounds keep their meaning."""
+ command = _make_command(mocker, "bigquery")
+ command.validate()
+ _estimator(command).return_value = QueryCostEstimate.from_value(value,
accuracy)
+ result = _run_warning(command)
+ if status == "within_threshold":
+ assert result == {"action": "proceed"}
+ else:
+ assert result["action"] == "confirm"
+ assert result["reason"] == (
+ "warning" if status == "exceeded" else "unavailable"
+ )
+ assert result["message"]
+ if accuracy in {"PRECISE", "LOWER_BOUND", "UPPER_BOUND"}:
+ prefix = {"LOWER_BOUND": "At least ", "UPPER_BOUND": "Up to "}.get(
+ accuracy, ""
+ )
+ assert result["details"][0]["value"] == f"{prefix}{value:,} bytes"
+ else:
+ assert result["details"] == [
+ {"label": "Warning threshold", "value": "10 bytes"}
+ ]
+
+
+def test_warning_compares_large_bytes_exactly(mocker: MockerFixture) -> None:
+ """A one-byte difference near JavaScript's integer boundary is compared
here."""
+ command = _make_command(mocker, "bigquery")
+ _set_extra(
+ command,
+ {
+ "cost_estimate_enabled": True,
+ "query_cost_warning": {
+ "enabled": True,
+ "metric": "bytes_processed",
+ "threshold": 2**53 - 1,
+ },
+ },
+ )
+ _estimator(command).return_value = QueryCostEstimate.from_value(2**53)
+ result = _run_warning(command)
+ assert result["reason"] == "warning"
+ assert result["details"][0]["value"] == "9,007,199,254,740,992 bytes"
+ assert result["details"][1]["value"] == "9,007,199,254,740,991 bytes"
+
+
[email protected](
+ "unavailable_reason", [None, "The planner supplied no estimate."]
+)
+def test_warning_unusable_estimate(
+ mocker: MockerFixture, unavailable_reason: str | None
+) -> None:
+ """An unusable engine estimate is unavailable, not within threshold."""
+ command = _make_command(mocker)
+ command.validate()
+ _estimator(command).return_value = QueryCostEstimate(
+ unavailable_reason=unavailable_reason
+ )
+ result = _run_warning(command)
+ assert result["action"] == "confirm"
+ assert result["reason"] == "unavailable"
+ assert result["message"]
+ if unavailable_reason:
+ assert result["message"] == unavailable_reason
+ assert result["details"] == [
+ {"label": "Warning threshold", "value": "10 planner cost units"}
+ ]
+
+
[email protected]("flag", ["SQLLAB_QUERY_COST_WARNINGS",
"ESTIMATE_QUERY_COST"])
+def test_warning_flag_off_still_checks_access(mocker: MockerFixture, flag:
str) -> None:
+ """Even disabled checks first run the existing database/table
authorization."""
+ command = _make_command(mocker)
+ flags = mocker.patch(
+ "superset.commands.sql_lab.estimate.is_feature_enabled",
+ side_effect=lambda name: name != flag,
+ )
+ mocker.patch(
+ "superset.db_engine_specs.query_warning.is_feature_enabled",
+ side_effect=lambda name: name != flag,
+ )
+ assert _run_warning(command) == {"action": "proceed"}
+ _estimator(command).assert_not_called()
+ manager = mocker.patch(
+ "superset.commands.sql_lab.estimate.security_manager",
new_callable=MagicMock
+ )
+ access = manager.raise_for_access
+ access.side_effect = SupersetSecurityException(
+ SupersetError(
+ message="Denied",
+ error_type=SupersetErrorType.TABLE_SECURITY_ACCESS_ERROR,
+ level=ErrorLevel.ERROR,
+ )
+ )
+ flags.reset_mock()
+ with pytest.raises(SupersetSecurityException):
+ _run_warning(command)
+ access.assert_called_once()
+ flags.assert_not_called()
+
+
[email protected](
+ "extra",
+ [
+ {
+ "cost_estimate_enabled": False,
+ "query_cost_warning": {
+ "enabled": True,
+ "metric": "total_cost",
+ "threshold": 10,
+ },
+ },
+ {
+ "cost_estimate_enabled": True,
+ "query_cost_warning": {
+ "enabled": False,
+ "metric": "total_cost",
+ "threshold": 10,
+ },
+ },
+ {"cost_estimate_enabled": True},
+ ],
+)
+def test_warning_connection_opt_out(
+ mocker: MockerFixture, extra: dict[str, Any]
+) -> None:
+ """Connection estimation and explicit warning opt-in are independent
switches."""
+ command = _make_command(mocker)
+ _set_extra(command, extra)
+ assert _run_warning(command) == {"action": "proceed"}
+ _estimator(command).assert_not_called()
+
+
[email protected]("enabled", [False, True])
+def test_warning_unsupported_engine(mocker: MockerFixture, enabled: bool) ->
None:
+ """Unsupported engines differ from supported engines with unestimatable
SQL."""
+ command = _make_command(mocker, "sqlite")
+ _set_extra(command, {"query_cost_warning": {"enabled": enabled}})
+ if enabled:
+ with pytest.raises(SupersetErrorException) as error:
+ _run_warning(command)
+ assert error.value.status == 400
+ else:
+ assert _run_warning(command) == {"action": "proceed"}
+ _estimator(command).assert_not_called()
+
+
[email protected](
+ "rule",
+ [
+ None,
+ {"enabled": "true"},
+ {"enabled": True},
+ {"enabled": True, "metric": "bytes_processed", "threshold": 10},
+ {"enabled": True, "metric": "total_cost", "threshold": "10"},
+ {"enabled": True, "metric": "unknown", "threshold": 10},
+ {"enabled": False, "threshold": -1},
+ ],
+)
+def test_warning_invalid_rule_is_an_error(mocker: MockerFixture, rule: object)
-> None:
+ """Invalid saved rules remain errors instead of allowing execution."""
+ command = _make_command(mocker)
+ _set_extra(
+ command,
+ {
+ "cost_estimate_enabled": True,
+ "query_cost_warning": rule,
+ },
+ )
+ with pytest.raises(SupersetErrorException) as error:
+ _run_warning(command)
+ assert error.value.status == 400
+ _estimator(command).assert_not_called()
+
+
[email protected](
+ "sql,params,reason",
+ [
+ ("SELECT {{ amount }}", {}, "templates"),
+ ("SELECT 1", {"amount": 1}, "templates"),
+ ("SELECT 1 {# comment #}", {}, "templates"),
+ ("SELECT 1; SELECT 2", {}, "single read-only SELECT"),
+ ("SHOW TABLES", {}, "single read-only SELECT"),
+ ("CREATE TABLE t AS SELECT 1", {}, "single read-only SELECT"),
+ ("SELECT * INTO t FROM sales", {}, "single read-only SELECT"),
+ (
+ "WITH deleted AS (DELETE FROM sales RETURNING *) SELECT * FROM
deleted",
+ {},
+ "single read-only SELECT",
+ ),
+ ],
+)
+def test_warning_unestimatable_sql(
+ mocker: MockerFixture, sql: str, params: dict[str, Any], reason: str
+) -> None:
+ """Only a single read-only SELECT is eligible for the initial warning
check."""
+ command = _make_command(mocker, sql=sql, template_params=params)
+ command.validate()
+ command._database.allow_dml = True
+ result = _run_warning(command)
+ assert result["action"] == "confirm"
+ assert result["reason"] == "unavailable"
+ assert reason in result["message"]
+ _estimator(command).assert_not_called()
+
+
[email protected]("custom", ["SQL_QUERY_MUTATOR",
"CUSTOM_TEMPLATE_PROCESSORS"])
+def test_warning_rejects_custom_transformations(
+ mocker: MockerFixture, app: Flask, custom: str
+) -> None:
+ """Do not run arbitrary transformations to predict a future execution
render."""
+ command = _make_command(mocker)
+ transform = mocker.MagicMock()
+ mocker.patch.dict(
+ app.config,
+ {
+ custom: transform
+ if custom == "SQL_QUERY_MUTATOR"
+ else {"postgresql": transform}
+ },
+ )
+ mocker.patch(
+ "superset.commands.sql_lab.estimate.is_feature_enabled",
+ side_effect=lambda name: name != "RLS_IN_SQLLAB",
+ )
+ result = _run_warning(command)
+ assert result["action"] == "confirm"
+ assert result["reason"] == "unavailable"
+ transform.assert_not_called()
+
+
+def test_warning_rejects_configured_sql_mutator(tmp_path: Path) -> None:
+ """Detect overrides loaded before the command imports the config module."""
+ config_file = tmp_path / "superset_config.py"
+ config_file.write_text(
+ "from unittest.mock import Mock\nSQL_QUERY_MUTATOR = Mock()\n",
+ encoding="utf-8",
+ )
+ result = subprocess.run( # noqa: S603
+ [
+ sys.executable,
+ "-c",
+ dedent(
+ """
+ from unittest.mock import MagicMock, patch
+
+ from superset import config
+ from superset.app import create_app
+
+ app =
create_app("tests.integration_tests.superset_test_config")
+
+ from superset.commands.sql_lab.estimate import
QueryEstimationCommand
+ from superset.db_engine_specs.base import QueryCostEstimate
+ from superset.db_engine_specs.postgres import
PostgresEngineSpec
+ from superset.models.core import Database
+ from superset.utils import json
+
+ database = Database(
+ sqlalchemy_uri="postgresql://localhost/test",
+ extra=json.dumps({
+ "cost_estimate_enabled": True,
+ "query_cost_warning": {
+ "enabled": True, "metric": "total_cost",
"threshold": 10,
+ },
+ }),
+ )
+ command = QueryEstimationCommand({
+ "database_id": 1, "sql": "SELECT 1", "catalog": "test",
+ "schema": "public", "template_params": {},
"check_warning": True,
+ })
+ with (
+ app.app_context(),
+
patch("superset.commands.sql_lab.estimate.DatabaseDAO.find_by_id",
+ return_value=database),
+
patch("superset.commands.sql_lab.estimate.security_manager",
+ new_callable=MagicMock),
+
patch("superset.commands.sql_lab.estimate.is_feature_enabled",
+ side_effect=lambda name: name in {
+ "SQLLAB_QUERY_COST_WARNINGS",
"ESTIMATE_QUERY_COST",
+ }),
+
patch("superset.db_engine_specs.query_warning.is_feature_enabled",
+ return_value=True),
+ patch.object(
+ PostgresEngineSpec, "estimate_query_cost_for_warning",
+ return_value=QueryCostEstimate.from_value(1),
+ ) as estimator,
+ ):
+ response = command.run()
+ assert response["action"] == "confirm", response
+ assert response["reason"] == "unavailable", response
+ assert "custom SQL transformations" in response["message"]
+ config.SQL_QUERY_MUTATOR.assert_not_called()
+ estimator.assert_not_called()
+ """
+ ),
+ ],
+ env={**os.environ, "SUPERSET_CONFIG_PATH": str(config_file)},
+ capture_output=True,
+ text=True,
+ check=False,
+ timeout=60,
+ )
+ assert result.returncode == 0, result.stdout + result.stderr
+
+
[email protected]("engine", ["postgresql", "bigquery"])
[email protected](
+ "sql_limit,dropdown,maximum,expected",
+ [
+ (None, 100, 50000, 101),
+ (5, 100, 50000, 6),
+ (500, 100, 50000, 101),
+ (100, 100, 50000, 101),
+ (None, 100, 20, 21),
+ (None, None, 20, 21),
+ (None, 0, 20, 21),
+ (0, 100, 50000, 50001),
+ (None, -1, 20, 21),
+ ],
+)
+def test_warning_matches_execution_sql_limits(
+ mocker: MockerFixture,
+ app: Flask,
+ engine: str,
+ sql_limit: int | None,
+ dropdown: int | None,
+ maximum: int,
+ expected: int,
+) -> None:
+ """Exercise the actual execution context, command and final extra-row
rewrite."""
+ sql = "SELECT * FROM sales" + ( # noqa: S608
+ f" LIMIT {sql_limit}" if sql_limit is not None else ""
+ )
+ command = _make_command(mocker, engine, sql=sql, query_limit=dropdown)
+ mocker.patch.dict(
+ app.config,
+ {"SQL_MAX_ROW": maximum, "SQL_QUERY_MUTATOR":
config.SQL_QUERY_MUTATOR},
+ )
+ _run_warning(command)
+ database = command._database
+ estimator = _estimator(command)
+ prepared = estimator.call_args.args[3]
+
+ context = SqlJsonExecutionContext(
+ {"database_id": 1, "sql": sql, "queryLimit": dropdown}
+ )
+ context.set_database(database)
+ context.set_query(Query(database=database))
+ execute = ExecuteSqlCommand(
+ context,
+ MagicMock(),
+ MagicMock(),
+ MagicMock(),
+ MagicMock(),
+ MagicMock(),
+ MagicMock(),
+ False,
+ )
+ execute._set_query_limit_if_required(sql)
Review Comment:
Thanks, added unit coverage for all three cases, asserting both the
effective limit and limiting_factor.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
