AlankritVerma01 commented on code in PR #45313:
URL: https://github.com/apache/superset/pull/45313#discussion_r4418296224
##########
superset/commands/database/update.py:
##########
@@ -281,8 +283,27 @@ def validate(self) -> None:
raise
DatabaseInvalidError(exceptions=[DatabaseExistsValidationError()])
if self._model:
+ self._validate_query_cost_warning()
self._check_no_unsafe_secret_rebind()
+ def _validate_query_cost_warning(self) -> None:
+ """Validate replacement or retained settings without adding omitted
extra."""
+ assert self._model
+ serialized_extra = self._properties.get("extra", self._model.extra) or
"{}"
+ if "query_cost_warning" not in serialized_extra:
Review Comment:
Hey, thanks for catching this. Fixed by parsing the JSON before checking for
the warning setting. Added regression coverage for escaped keys in both
submitted and retained settings, including invalid thresholds and unsupported
metrics.
##########
superset/commands/sql_lab/estimate.py:
##########
@@ -164,11 +178,107 @@ def _apply_sql_security(self, sql: str) -> str:
return sql
+ def _get_warning_rule(self, check: QueryWarningCheck | None) -> dict[str,
Any]:
+ """Validate saved settings directly so invalid rules remain errors."""
+ try:
+ rule = json.loads(self._database.extra or "{}").get(
+ "query_cost_warning", {"enabled": False}
+ )
+ rule = validate_query_warning_rule(rule, check)
+ except ValidationError as ex:
+ raise SupersetErrorException(
+ SupersetError(
+ message=__("The saved query cost warning rule is
invalid."),
+ error_type=SupersetErrorType.GENERIC_COMMAND_ERROR,
+ level=ErrorLevel.ERROR,
+ extra={"messages": ex.messages},
+ ),
+ status=400,
+ ) from ex
+ return rule
+
+ def _prepare_warning_sql(self) -> tuple[str | None, str | None]:
+ """Prepare execution-equivalent SQL, or explain why it cannot be
estimated."""
+ spec = self._database.db_engine_spec
+ if self._template_params or any(
+ kind != "data" for _, kind, _ in
Environment(autoescape=True).lex(self._sql)
+ ):
+ return None, __("Warning checks are unavailable for SQL
templates.")
+ mutator = app.config["SQL_QUERY_MUTATOR"]
+ if (mutator and mutator is not config._default_sql_query_mutator) or (
+ is_feature_enabled("ENABLE_TEMPLATE_PROCESSING")
+ and spec.engine in app.config["CUSTOM_TEMPLATE_PROCESSORS"]
+ ):
+ return None, __(
+ "Warning checks are unavailable for custom SQL
transformations."
+ )
+
+ if self._catalog is None:
+ self._catalog = self._database.get_default_catalog()
+ sql = self._apply_sql_security(self._sql)
+ script = SQLScript(sql, engine=spec.engine)
+ if (
+ len(script.statements) != 1
+ or not script.statements[0].is_select()
+ or script.statements[0].is_mutating()
+ ):
+ return None, __("Warning checks require a single read-only
SELECT.")
+
+ # Use execution's dropdown/SQL limits and its extra truncation row.
+ # This Query is never persisted or dispatched for execution.
+ query = Query(database=self._database, select_as_cta_used=False)
+ if query in db.session:
+ db.session.expunge(query)
+ set_query_limit(query, sql, get_query_limit(self._query_limit), spec)
+ apply_limit(query, script.statements[0])
+ return script.format(comments=spec.allows_sql_comments), None
+
+ def _run_warning_check(self) -> QueryWarningResult:
+ """Prepare the intended query and delegate assessment to its engine
check."""
+ if not is_feature_enabled("SQLLAB_QUERY_COST_WARNINGS"):
+ return {"action": "proceed"}
+ check = self._database.db_engine_spec.get_query_warning_check()
+ rule = self._get_warning_rule(check)
+ if check is None or not rule["enabled"] or not
check.is_enabled(self._database):
+ return {"action": "proceed"}
+ sql, reason = self._prepare_warning_sql()
+ if sql is None:
+ return {
+ "action": "confirm",
+ "reason": "unavailable",
+ "message": reason
+ or __(
+ "The estimate is unavailable or cannot establish whether "
+ "usage is within the warning threshold."
+ ),
+ }
+ timeout = app.config["SQLLAB_QUERY_COST_ESTIMATE_TIMEOUT"]
+ try:
+ with utils.timeout(
Review Comment:
Thanks, agreed - the previous test only covered handling the exception.
Updated it to check that the timeout context is entered with the configured
duration, and that a timeout returns confirm/unavailable without calling the
estimator.
--
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]
