rusackas commented on code in PR #41127:
URL: https://github.com/apache/superset/pull/41127#discussion_r3616409617
##########
superset/sql/execution/executor.py:
##########
@@ -87,12 +90,102 @@
StatementResult,
)
+ from superset.db_engine_specs.base import BaseEngineSpec
from superset.models.core import Database
from superset.result_set import SupersetResultSet
logger = logging.getLogger(__name__)
+def _raise_all_statements_stripped() -> NoReturn:
+ """Raise a clean error for a mutator that stripped a query down to
nothing."""
+ raise SupersetErrorException(
+ SupersetError(
+ message=__(
+ "The SQL query mutator removed all executable "
+ "statements from this query."
+ ),
+ error_type=SupersetErrorType.INVALID_SQL_ERROR,
+ level=ErrorLevel.ERROR,
+ )
+ )
+
+
+def _has_executable_statements(sql: str, engine: str) -> bool:
+ """Best-effort check that mutated SQL still contains executable
statements."""
+ try:
+ return bool(SQLScript(sql, engine=engine).statements)
+ except SupersetParseError:
+ # A mutator may emit engine-specific SQL our parser can't handle; the
+ # database itself is the authority on validity in that case.
+ return True
+
+
+def build_statement_blocks(
+ parsed_script: SQLScript,
+ db_engine_spec: type[BaseEngineSpec],
+ database: Database,
+) -> tuple[SQLScript, list[str]]:
+ """
+ Build the SQL blocks to execute from a parsed script, applying
+ ``SQL_QUERY_MUTATOR`` according to ``MUTATE_AFTER_SPLIT``.
+
+ Some databases (like BigQuery and Kusto) do not persist state across
multiple
+ statements if they're run separately (especially when using `NullPool`),
so the
+ query runs as a single joined block when the engine spec requires it;
otherwise
+ each statement becomes its own block. Shared by the sync (``sql_lab``) and
+ async (``celery_task``) SQL Lab paths so the
+ ``run_multiple_statements_as_one`` × ``MUTATE_AFTER_SPLIT`` matrix behaves
+ identically in both.
+
+ Returns the (possibly re-parsed) script and the blocks to execute.
+
+ :raises SupersetErrorException: if the mutator strips the query down to
+ nothing executable (e.g. only comments/whitespace)
+ """
+ blocks: list[str]
+ if db_engine_spec.run_multiple_statements_as_one:
+ if app.config["MUTATE_AFTER_SPLIT"]:
+ # These engines never actually execute statements individually, so
+ # the per-block mutation call at execution time (whose `is_split`
is
+ # always `False` here) would never fire. Mutate each statement
here,
+ # before joining them into the single block this engine requires,
so
+ # `MUTATE_AFTER_SPLIT=True` still applies the mutator per
statement.
+ joined_block = ";\n".join(
+ database.mutate_sql_based_on_config(
+
statement.format(comments=db_engine_spec.allows_sql_comments),
+ is_split=True,
+ )
+ for statement in parsed_script.statements
+ )
+ if not _has_executable_statements(joined_block,
db_engine_spec.engine):
+ _raise_all_statements_stripped()
+ blocks = [joined_block]
+ else:
+ blocks =
[parsed_script.format(comments=db_engine_spec.allows_sql_comments)]
+ else:
+ if not app.config["MUTATE_AFTER_SPLIT"]:
+ # `MUTATE_AFTER_SPLIT=False` means the mutator should see the
whole,
+ # un-split query, but this engine executes statements individually.
+ # Mutate the whole block up front and re-parse it, so the
+ # per-statement split below (and the per-block mutation call at
+ # execution time, which is a no-op here since its `is_split=True`
no
+ # longer matches the config) operate on the already-mutated SQL.
+ mutated_sql: str = database.mutate_sql_based_on_config(
+
parsed_script.format(comments=db_engine_spec.allows_sql_comments),
+ is_split=False,
+ )
+ parsed_script = SQLScript(mutated_sql,
engine=db_engine_spec.engine)
+ if not parsed_script.statements:
+ _raise_all_statements_stripped()
Review Comment:
As noted on the earlier thread, the re-parse is unavoidable on this path...
the engine executes statements individually, so the mutated block has to be
split, and splitting means parsing. Where validation is optional
(`_has_executable_statements`) parse failures already fall back to letting the
database judge, and where it isn't, `SupersetParseError` surfaces as a clean
error rather than a crash.
##########
superset/sql/execution/executor.py:
##########
@@ -140,8 +238,23 @@ def execute_sql_with_cursor(
# Apply SQL mutation
stmt_sql = database.mutate_sql_based_on_config(
statement,
- is_split=True,
+ is_split=is_split,
)
+ if not stmt_sql.strip():
+ # A `SQL_QUERY_MUTATOR` that strips a statement down to nothing
+ # would otherwise be sent to the database engine as an empty
+ # query, surfacing a confusing engine-specific error instead of
+ # a clean one.
+ raise SupersetErrorException(
+ SupersetError(
+ message=__(
+ "The SQL query mutator removed all executable "
+ "statements from this query."
+ ),
+ error_type=SupersetErrorType.INVALID_SQL_ERROR,
+ level=ErrorLevel.ERROR,
+ )
+ )
Review Comment:
Good catch, fixed! The per-statement guard uses the same parser-based
`_has_executable_statements` check as the joined-block path, so comments-only
mutator output raises the clean `INVALID_SQL_ERROR` too. Covered both cases in
the test.
--
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]
