bito-code-review[bot] commented on code in PR #42785:
URL: https://github.com/apache/superset/pull/42785#discussion_r4073044688
##########
tests/unit_tests/commands/sql_lab/test_estimate.py:
##########
@@ -30,8 +32,11 @@
OAuth2RedirectError,
SupersetErrorException,
SupersetGenericDBErrorException,
+ SupersetParseError,
SupersetSecurityException,
)
+from superset.models.core import Database
+from tests.unit_tests.conftest import with_feature_flags # noqa: E402
Review Comment:
<!-- Bito Reply -->
The suggestion to remove the dead `# noqa: E402` suppression is appropriate.
Since the import is located within the standard top-of-module import block, the
E402 rule does not trigger, and removing the unnecessary comment improves code
cleanliness and consistency with other test files.
**tests/unit_tests/commands/sql_lab/test_estimate.py**
```
from superset.models.core import Database
from tests.unit_tests.conftest import with_feature_flags
```
##########
tests/unit_tests/commands/sql_lab/test_estimate.py:
##########
@@ -577,3 +617,447 @@ def
test_run_reraises_oauth2_redirect_error_from_cost_estimation(
command.run()
assert exc_info.value.status == 403
+
+
+# ---------------------------------------------------------------------------
+# Templates are rendered before estimating, as on the execution path
+# ---------------------------------------------------------------------------
+
+
+@patch("superset.commands.sql_lab.estimate.app")
+@patch("superset.commands.sql_lab.estimate.get_template_processor")
+@patch("superset.commands.sql_lab.estimate.security_manager",
new_callable=MagicMock)
+@patch("superset.commands.sql_lab.estimate.DatabaseDAO")
+def test_run_renders_a_template_without_template_params(
+ mock_dao: MagicMock,
+ mock_security_manager: MagicMock,
+ mock_get_template_processor: MagicMock,
+ mock_app: MagicMock,
+) -> None:
+ """A query needs no declared parameter to need rendering:
``get_time_filter()``
+ and friends take none, and SQL Lab posts an empty ``template_params`` for
an
+ estimate, so gating the render on it left the template in place for the
+ parser to choke on."""
+ mock_app.config = {
+ "DISALLOWED_SQL_FUNCTIONS": {},
+ "DISALLOWED_SQL_TABLES": {},
+ "SQLLAB_QUERY_COST_ESTIMATE_TIMEOUT": 10,
+ "QUERY_COST_FORMATTERS_BY_ENGINE": {},
+ }
+ mock_database = MagicMock()
+ mock_database.db_engine_spec.engine = "postgresql"
+ mock_database.allow_dml = False
+ mock_database.db_engine_spec.query_cost_formatter.return_value = [{"Cost":
"1"}]
+ mock_dao.find_by_id.return_value = mock_database
+ mock_security_manager.raise_for_access.return_value = None
+ processor = mock_get_template_processor.return_value
+ processor.process_template.return_value = "SELECT 1"
+ processor.get_undefined_parameters.return_value = set()
+
+ sql = "{% set tf = get_time_filter('ds') %}SELECT 1 {% if tf %}{% endif %}"
+ command = QueryEstimationCommand(_make_params(sql=sql))
+
+ assert command.run() == [{"Cost": "1"}]
+
mock_get_template_processor.return_value.process_template.assert_called_once_with(
+ sql
+ )
+ # What reaches the engine is the rendered SQL.
+ assert (
+ mock_database.db_engine_spec.estimate_query_cost.call_args.args[3] ==
"SELECT 1"
+ )
+
+
+@patch("superset.commands.sql_lab.estimate.app")
+@patch("superset.commands.sql_lab.estimate.get_template_processor")
+@patch("superset.commands.sql_lab.estimate.security_manager",
new_callable=MagicMock)
+@patch("superset.commands.sql_lab.estimate.DatabaseDAO")
+def test_run_estimates_a_template_its_parameters_fully_bind(
+ mock_dao: MagicMock,
+ mock_security_manager: MagicMock,
+ mock_get_template_processor: MagicMock,
+ mock_app: MagicMock,
+) -> None:
+ """A template whose parameters are all supplied renders to the same SQL the
+ query would run, so it is estimated rather than refused."""
+ mock_app.config = {
+ "DISALLOWED_SQL_FUNCTIONS": {},
+ "DISALLOWED_SQL_TABLES": {},
+ "SQLLAB_QUERY_COST_ESTIMATE_TIMEOUT": 10,
+ "QUERY_COST_FORMATTERS_BY_ENGINE": {},
+ }
+ mock_database = MagicMock()
+ mock_database.db_engine_spec.engine = "postgresql"
+ mock_database.allow_dml = False
+ mock_database.db_engine_spec.query_cost_formatter.return_value = [{"Cost":
"2"}]
+ mock_dao.find_by_id.return_value = mock_database
+ mock_security_manager.raise_for_access.return_value = None
+ processor = mock_get_template_processor.return_value
+ processor.process_template.return_value = "SELECT '2026-08-20'"
+ processor.get_undefined_parameters.return_value = set()
+
+ command = QueryEstimationCommand(
+ _make_params(sql="SELECT '{{ ds }}'", template_params={"ds":
"2026-08-20"})
+ )
+
+ assert command.run() == [{"Cost": "2"}]
+
mock_get_template_processor.return_value.process_template.assert_called_once_with(
+ "SELECT '{{ ds }}'", ds="2026-08-20"
+ )
+ # What reaches the engine is the rendered SQL, not the template.
+ assert (
+ mock_database.db_engine_spec.estimate_query_cost.call_args.args[3]
+ == "SELECT '2026-08-20'"
+ )
+
+
+@patch("superset.commands.sql_lab.estimate.app")
+@patch("superset.commands.sql_lab.estimate.get_template_processor")
+@patch("superset.commands.sql_lab.estimate.security_manager",
new_callable=MagicMock)
+@patch("superset.commands.sql_lab.estimate.DatabaseDAO")
+def test_run_reports_an_unprovided_parameter_as_missing(
+ mock_dao: MagicMock,
+ mock_security_manager: MagicMock,
+ mock_get_template_processor: MagicMock,
+ mock_app: MagicMock,
+) -> None:
+ """``DebugUndefined`` leaves an unprovided parameter in place instead of
+ raising, and in a position like a string literal the leftover still parses.
+ Estimating it would describe a query the user cannot run, so it gets the
+ same typed response the execution path gives it."""
+ mock_app.config = {
+ "DISALLOWED_SQL_FUNCTIONS": {},
+ "DISALLOWED_SQL_TABLES": {},
+ "SQLLAB_QUERY_COST_ESTIMATE_TIMEOUT": 10,
+ "QUERY_COST_FORMATTERS_BY_ENGINE": {},
+ }
+ mock_database = MagicMock()
+ mock_database.db_engine_spec.engine = "postgresql"
+ mock_database.allow_dml = False
+ mock_dao.find_by_id.return_value = mock_database
+ mock_security_manager.raise_for_access.return_value = None
+ processor = mock_get_template_processor.return_value
+ processor.process_template.return_value = "SELECT '{{ ds }}' AS d"
+ processor.get_undefined_parameters.return_value = {"ds"}
+
+ command = QueryEstimationCommand(_make_params(sql="SELECT '{{ ds }}' AS
d"))
+ with pytest.raises(SupersetErrorException) as exc_info:
+ command.run()
+
+ error = exc_info.value.error
+ assert exc_info.value.status == 400
+ assert error.error_type == SupersetErrorType.MISSING_TEMPLATE_PARAMS_ERROR
+ assert error.message.startswith('The parameter "ds" in your query is
undefined.')
+ # The execution path's suggestion travels with it.
+ assert "Set Parameters" in error.message
+ assert error.extra["undefined_parameters"] == ["ds"]
+ assert error.extra["issue_codes"][0]["code"] == 1006
+ # Nothing was estimated.
+ mock_database.db_engine_spec.estimate_query_cost.assert_not_called()
+
+
+@patch("superset.commands.sql_lab.estimate.app")
+@patch("superset.commands.sql_lab.estimate.get_template_processor")
+@patch("superset.commands.sql_lab.estimate.security_manager",
new_callable=MagicMock)
+@patch("superset.commands.sql_lab.estimate.DatabaseDAO")
+def test_run_leaves_a_genuine_syntax_error_alone(
+ mock_dao: MagicMock,
+ mock_security_manager: MagicMock,
+ mock_get_template_processor: MagicMock,
+ mock_app: MagicMock,
+) -> None:
+ """SQL that fails to parse with nothing undefined in it keeps the parser's
+ own error -- the query really is malformed."""
+ mock_app.config = {
+ "DISALLOWED_SQL_FUNCTIONS": {},
+ "DISALLOWED_SQL_TABLES": {},
+ "SQLLAB_QUERY_COST_ESTIMATE_TIMEOUT": 10,
+ "QUERY_COST_FORMATTERS_BY_ENGINE": {},
+ }
+ mock_database = MagicMock()
+ mock_database.db_engine_spec.engine = "postgresql"
+ mock_database.allow_dml = False
+ mock_dao.find_by_id.return_value = mock_database
+ mock_security_manager.raise_for_access.return_value = None
+ processor = mock_get_template_processor.return_value
+ processor.process_template.return_value = "SELECT FROM FROM"
+ processor.get_undefined_parameters.return_value = set()
+
+ command = QueryEstimationCommand(_make_params(sql="SELECT FROM FROM"))
+ with pytest.raises(SupersetParseError) as exc_info:
+ command.run()
+
+ assert exc_info.value.error.error_type ==
SupersetErrorType.INVALID_SQL_ERROR
+
+
+# ---------------------------------------------------------------------------
+# What is authorized is what is estimated
+# ---------------------------------------------------------------------------
+
+
+@patch("superset.commands.sql_lab.estimate.app")
+@patch("superset.commands.sql_lab.estimate.get_template_processor")
+@patch("superset.commands.sql_lab.estimate.security_manager",
new_callable=MagicMock)
+@patch("superset.commands.sql_lab.estimate.DatabaseDAO")
+def test_run_reauthorizes_the_rendered_sql(
+ mock_dao: MagicMock,
+ mock_security_manager: MagicMock,
+ mock_get_template_processor: MagicMock,
+ mock_app: MagicMock,
+) -> None:
+ """``validate()`` authorizes a render of its own, and a template need not
+ render the same way twice. The SQL that will be estimated is authorized as
+ a literal, as ``_validate_rendered_access`` does on the execution path."""
+ mock_app.config = {
+ "DISALLOWED_SQL_FUNCTIONS": {},
+ "DISALLOWED_SQL_TABLES": {},
+ "SQLLAB_QUERY_COST_ESTIMATE_TIMEOUT": 10,
+ "QUERY_COST_FORMATTERS_BY_ENGINE": {},
+ }
+ mock_database = MagicMock()
+ mock_database.db_engine_spec.engine = "postgresql"
+ mock_database.allow_dml = False
+ mock_database.db_engine_spec.query_cost_formatter.return_value = [{"Cost":
"1"}]
+ mock_dao.find_by_id.return_value = mock_database
+ mock_security_manager.raise_for_access.return_value = None
+ processor = mock_get_template_processor.return_value
+ processor.process_template.return_value = "SELECT * FROM allowed_ds"
+ processor.get_undefined_parameters.return_value = set()
+
+ sql = "SELECT * FROM {{ ['allowed_ds', 'secret_tbl'] | random }}"
Review Comment:
<!-- Bito Reply -->
The suggestion to add explicit type annotations to the local variables
`processor` and `sql` is appropriate. Applying these annotations improves
consistency with the existing typing conventions in the test file and enhances
code readability. You should apply this change to the new tests.
**tests/unit_tests/commands/sql_lab/test_estimate.py**
```
processor: MagicMock = mock_get_template_processor.return_value
processor.process_template.return_value = "SELECT 1"
processor.get_undefined_parameters.return_value = set()
sql: str = "{% set tf = get_time_filter('ds') %}SELECT 1 {% if tf %}{%
endif %}"
```
##########
tests/unit_tests/commands/sql_lab/test_estimate.py:
##########
@@ -577,3 +617,447 @@ def
test_run_reraises_oauth2_redirect_error_from_cost_estimation(
command.run()
assert exc_info.value.status == 403
+
+
+# ---------------------------------------------------------------------------
+# Templates are rendered before estimating, as on the execution path
+# ---------------------------------------------------------------------------
+
+
+@patch("superset.commands.sql_lab.estimate.app")
+@patch("superset.commands.sql_lab.estimate.get_template_processor")
+@patch("superset.commands.sql_lab.estimate.security_manager",
new_callable=MagicMock)
+@patch("superset.commands.sql_lab.estimate.DatabaseDAO")
+def test_run_renders_a_template_without_template_params(
+ mock_dao: MagicMock,
+ mock_security_manager: MagicMock,
+ mock_get_template_processor: MagicMock,
+ mock_app: MagicMock,
+) -> None:
+ """A query needs no declared parameter to need rendering:
``get_time_filter()``
+ and friends take none, and SQL Lab posts an empty ``template_params`` for
an
+ estimate, so gating the render on it left the template in place for the
+ parser to choke on."""
+ mock_app.config = {
+ "DISALLOWED_SQL_FUNCTIONS": {},
+ "DISALLOWED_SQL_TABLES": {},
+ "SQLLAB_QUERY_COST_ESTIMATE_TIMEOUT": 10,
+ "QUERY_COST_FORMATTERS_BY_ENGINE": {},
+ }
+ mock_database = MagicMock()
+ mock_database.db_engine_spec.engine = "postgresql"
+ mock_database.allow_dml = False
+ mock_database.db_engine_spec.query_cost_formatter.return_value = [{"Cost":
"1"}]
+ mock_dao.find_by_id.return_value = mock_database
+ mock_security_manager.raise_for_access.return_value = None
+ processor = mock_get_template_processor.return_value
+ processor.process_template.return_value = "SELECT 1"
+ processor.get_undefined_parameters.return_value = set()
+
+ sql = "{% set tf = get_time_filter('ds') %}SELECT 1 {% if tf %}{% endif %}"
+ command = QueryEstimationCommand(_make_params(sql=sql))
+
+ assert command.run() == [{"Cost": "1"}]
+
mock_get_template_processor.return_value.process_template.assert_called_once_with(
+ sql
+ )
+ # What reaches the engine is the rendered SQL.
+ assert (
+ mock_database.db_engine_spec.estimate_query_cost.call_args.args[3] ==
"SELECT 1"
+ )
+
+
+@patch("superset.commands.sql_lab.estimate.app")
+@patch("superset.commands.sql_lab.estimate.get_template_processor")
+@patch("superset.commands.sql_lab.estimate.security_manager",
new_callable=MagicMock)
+@patch("superset.commands.sql_lab.estimate.DatabaseDAO")
+def test_run_estimates_a_template_its_parameters_fully_bind(
+ mock_dao: MagicMock,
+ mock_security_manager: MagicMock,
+ mock_get_template_processor: MagicMock,
+ mock_app: MagicMock,
+) -> None:
+ """A template whose parameters are all supplied renders to the same SQL the
+ query would run, so it is estimated rather than refused."""
+ mock_app.config = {
+ "DISALLOWED_SQL_FUNCTIONS": {},
+ "DISALLOWED_SQL_TABLES": {},
+ "SQLLAB_QUERY_COST_ESTIMATE_TIMEOUT": 10,
+ "QUERY_COST_FORMATTERS_BY_ENGINE": {},
+ }
+ mock_database = MagicMock()
+ mock_database.db_engine_spec.engine = "postgresql"
+ mock_database.allow_dml = False
+ mock_database.db_engine_spec.query_cost_formatter.return_value = [{"Cost":
"2"}]
+ mock_dao.find_by_id.return_value = mock_database
+ mock_security_manager.raise_for_access.return_value = None
+ processor = mock_get_template_processor.return_value
+ processor.process_template.return_value = "SELECT '2026-08-20'"
+ processor.get_undefined_parameters.return_value = set()
+
+ command = QueryEstimationCommand(
+ _make_params(sql="SELECT '{{ ds }}'", template_params={"ds":
"2026-08-20"})
+ )
+
+ assert command.run() == [{"Cost": "2"}]
+
mock_get_template_processor.return_value.process_template.assert_called_once_with(
+ "SELECT '{{ ds }}'", ds="2026-08-20"
+ )
+ # What reaches the engine is the rendered SQL, not the template.
+ assert (
+ mock_database.db_engine_spec.estimate_query_cost.call_args.args[3]
+ == "SELECT '2026-08-20'"
+ )
+
+
+@patch("superset.commands.sql_lab.estimate.app")
+@patch("superset.commands.sql_lab.estimate.get_template_processor")
+@patch("superset.commands.sql_lab.estimate.security_manager",
new_callable=MagicMock)
+@patch("superset.commands.sql_lab.estimate.DatabaseDAO")
+def test_run_reports_an_unprovided_parameter_as_missing(
+ mock_dao: MagicMock,
+ mock_security_manager: MagicMock,
+ mock_get_template_processor: MagicMock,
+ mock_app: MagicMock,
+) -> None:
+ """``DebugUndefined`` leaves an unprovided parameter in place instead of
+ raising, and in a position like a string literal the leftover still parses.
+ Estimating it would describe a query the user cannot run, so it gets the
+ same typed response the execution path gives it."""
+ mock_app.config = {
+ "DISALLOWED_SQL_FUNCTIONS": {},
+ "DISALLOWED_SQL_TABLES": {},
+ "SQLLAB_QUERY_COST_ESTIMATE_TIMEOUT": 10,
+ "QUERY_COST_FORMATTERS_BY_ENGINE": {},
+ }
+ mock_database = MagicMock()
+ mock_database.db_engine_spec.engine = "postgresql"
+ mock_database.allow_dml = False
+ mock_dao.find_by_id.return_value = mock_database
+ mock_security_manager.raise_for_access.return_value = None
+ processor = mock_get_template_processor.return_value
+ processor.process_template.return_value = "SELECT '{{ ds }}' AS d"
+ processor.get_undefined_parameters.return_value = {"ds"}
+
+ command = QueryEstimationCommand(_make_params(sql="SELECT '{{ ds }}' AS
d"))
+ with pytest.raises(SupersetErrorException) as exc_info:
+ command.run()
+
+ error = exc_info.value.error
+ assert exc_info.value.status == 400
+ assert error.error_type == SupersetErrorType.MISSING_TEMPLATE_PARAMS_ERROR
+ assert error.message.startswith('The parameter "ds" in your query is
undefined.')
+ # The execution path's suggestion travels with it.
+ assert "Set Parameters" in error.message
+ assert error.extra["undefined_parameters"] == ["ds"]
+ assert error.extra["issue_codes"][0]["code"] == 1006
+ # Nothing was estimated.
+ mock_database.db_engine_spec.estimate_query_cost.assert_not_called()
+
+
+@patch("superset.commands.sql_lab.estimate.app")
+@patch("superset.commands.sql_lab.estimate.get_template_processor")
+@patch("superset.commands.sql_lab.estimate.security_manager",
new_callable=MagicMock)
+@patch("superset.commands.sql_lab.estimate.DatabaseDAO")
+def test_run_leaves_a_genuine_syntax_error_alone(
+ mock_dao: MagicMock,
+ mock_security_manager: MagicMock,
+ mock_get_template_processor: MagicMock,
+ mock_app: MagicMock,
+) -> None:
+ """SQL that fails to parse with nothing undefined in it keeps the parser's
+ own error -- the query really is malformed."""
+ mock_app.config = {
+ "DISALLOWED_SQL_FUNCTIONS": {},
+ "DISALLOWED_SQL_TABLES": {},
+ "SQLLAB_QUERY_COST_ESTIMATE_TIMEOUT": 10,
+ "QUERY_COST_FORMATTERS_BY_ENGINE": {},
+ }
+ mock_database = MagicMock()
+ mock_database.db_engine_spec.engine = "postgresql"
+ mock_database.allow_dml = False
+ mock_dao.find_by_id.return_value = mock_database
+ mock_security_manager.raise_for_access.return_value = None
+ processor = mock_get_template_processor.return_value
+ processor.process_template.return_value = "SELECT FROM FROM"
+ processor.get_undefined_parameters.return_value = set()
+
+ command = QueryEstimationCommand(_make_params(sql="SELECT FROM FROM"))
+ with pytest.raises(SupersetParseError) as exc_info:
+ command.run()
+
+ assert exc_info.value.error.error_type ==
SupersetErrorType.INVALID_SQL_ERROR
+
+
+# ---------------------------------------------------------------------------
+# What is authorized is what is estimated
+# ---------------------------------------------------------------------------
+
+
+@patch("superset.commands.sql_lab.estimate.app")
+@patch("superset.commands.sql_lab.estimate.get_template_processor")
+@patch("superset.commands.sql_lab.estimate.security_manager",
new_callable=MagicMock)
+@patch("superset.commands.sql_lab.estimate.DatabaseDAO")
+def test_run_reauthorizes_the_rendered_sql(
+ mock_dao: MagicMock,
+ mock_security_manager: MagicMock,
+ mock_get_template_processor: MagicMock,
+ mock_app: MagicMock,
+) -> None:
+ """``validate()`` authorizes a render of its own, and a template need not
+ render the same way twice. The SQL that will be estimated is authorized as
+ a literal, as ``_validate_rendered_access`` does on the execution path."""
+ mock_app.config = {
+ "DISALLOWED_SQL_FUNCTIONS": {},
+ "DISALLOWED_SQL_TABLES": {},
+ "SQLLAB_QUERY_COST_ESTIMATE_TIMEOUT": 10,
+ "QUERY_COST_FORMATTERS_BY_ENGINE": {},
+ }
+ mock_database = MagicMock()
+ mock_database.db_engine_spec.engine = "postgresql"
+ mock_database.allow_dml = False
+ mock_database.db_engine_spec.query_cost_formatter.return_value = [{"Cost":
"1"}]
+ mock_dao.find_by_id.return_value = mock_database
+ mock_security_manager.raise_for_access.return_value = None
+ processor = mock_get_template_processor.return_value
+ processor.process_template.return_value = "SELECT * FROM allowed_ds"
+ processor.get_undefined_parameters.return_value = set()
+
+ sql = "SELECT * FROM {{ ['allowed_ds', 'secret_tbl'] | random }}"
+ command = QueryEstimationCommand(_make_params(sql=sql, schema="public"))
+
+ assert command.run() == [{"Cost": "1"}]
+
+ first, second = mock_security_manager.raise_for_access.call_args_list
Review Comment:
<!-- Bito Reply -->
The suggestion to explicitly annotate `first` and `second` as `MagicMock` is
appropriate. It improves type safety and consistency with the existing mock
parameters in the test file, aligning with the project's coding standards for
test files.
**tests/unit_tests/commands/sql_lab/test_estimate.py**
```
first: MagicMock
second: MagicMock
first, second = mock_security_manager.raise_for_access.call_args_list
```
--
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]
