villebro commented on code in PR #43942:
URL: https://github.com/apache/superset/pull/43942#discussion_r3959313256
##########
superset/commands/sql_lab/execute.py:
##########
@@ -164,13 +172,144 @@ def _run_sql_json_exec_from_scratch(self) ->
SqlJsonExecutionStatus:
self._query_dao.update(
query, {"limit": self._execution_context.query.limit}
)
- return self._sql_json_executor.execute(
- self._execution_context, rendered_query, self._log_params
- )
+ return self._execute(rendered_query)
except Exception:
self._query_dao.update(query, {"status": QueryStatus.FAILED})
raise
+ def _execute(self, rendered_query: str) -> SqlJsonExecutionStatus:
+ """Dispatch to synchronous (inline) or asynchronous (GTF task)
execution.
+
+ Sync runs the query in-process via the unified SQL-Lab executor entry.
+ Async only *prepares* here (the ``Query`` row is committed by this
+ command's transaction); the actual GTF task is scheduled by
+ ``submit_async``, which the endpoint calls after the transaction
commits
+ because scheduling a task cannot run inside an outer ``@transaction``.
+ """
+ if self._execution_context.is_run_asynchronous():
+ return self._prepare_async(rendered_query)
+ return self._execute_sync(rendered_query)
+
+ def _execute_sync(self, rendered_query: str) -> SqlJsonExecutionStatus:
+ """Run the query inline via ``execute_sql_lab_query`` and set the
result."""
+ from superset.sql.execution.sqllab_executor import
execute_sql_lab_query
+ from superset.sql_lab import get_query, handle_query_error
+
+ context = self._execution_context
+ query = context.query
+ timeout = app.config["SQLLAB_TIMEOUT"]
+ store_results = (
+ is_feature_enabled("SQLLAB_BACKEND_PERSISTENCE")
+ and not context.select_as_cta
+ )
+ try:
+ with utils.timeout(
+ seconds=timeout,
+ error_message=f"The query exceeded the {timeout} seconds
timeout.",
+ ):
+ try:
+ data = execute_sql_lab_query(
+ query,
+ rendered_query,
+ return_results=True,
+ store_results=store_results,
+ expand_data=context.expand_data,
+ log_params=self._log_params,
+ )
+ except Exception as ex: # pylint: disable=broad-except
+ # Re-fetch (the session may be poisoned) and build the
error
+ # payload, mirroring the classic synchronous path.
+ data = handle_query_error(ex, get_query(query_id=query.id))
+ except SupersetTimeoutException:
+ raise
+ except Exception as ex:
+ logger.exception("Query %i failed unexpectedly", query.id)
+ raise SupersetGenericDBErrorException(
+ utils.error_msg_from_exception(ex)
+ ) from ex
+
+ context.set_execution_result(data)
+ if data and data.get("status") == QueryStatus.FAILED:
+ if data.get("errors"):
+ errors = [SupersetError(**params) for params in data["errors"]]
+ status = (
+ 500
+ if any(error.level == ErrorLevel.ERROR for error in errors)
+ else 400
+ )
+ raise SupersetErrorsException(errors, status=status)
+ raise SupersetGenericDBErrorException(data["error"])
+ return SqlJsonExecutionStatus.HAS_RESULTS
Review Comment:
This is intended parity with the classic synchronous path
(`SynchronousSqlJsonExecutor.execute` likewise returned `HAS_RESULTS` carrying
a STOPPED payload). The client keys off `query.state == "stopped"` in the
returned payload, not the HTTP/command status, and there is no `QUERY_STOPPED`
command-status to return. A STOPPED payload on the sync path only arises if a
concurrent `/query/stop` lands mid multi-statement query. Leaving as-is to
preserve behavior.
##########
superset/commands/sql_lab/execute.py:
##########
@@ -164,13 +172,144 @@ def _run_sql_json_exec_from_scratch(self) ->
SqlJsonExecutionStatus:
self._query_dao.update(
query, {"limit": self._execution_context.query.limit}
)
- return self._sql_json_executor.execute(
- self._execution_context, rendered_query, self._log_params
- )
+ return self._execute(rendered_query)
except Exception:
self._query_dao.update(query, {"status": QueryStatus.FAILED})
raise
+ def _execute(self, rendered_query: str) -> SqlJsonExecutionStatus:
+ """Dispatch to synchronous (inline) or asynchronous (GTF task)
execution.
+
+ Sync runs the query in-process via the unified SQL-Lab executor entry.
+ Async only *prepares* here (the ``Query`` row is committed by this
+ command's transaction); the actual GTF task is scheduled by
+ ``submit_async``, which the endpoint calls after the transaction
commits
+ because scheduling a task cannot run inside an outer ``@transaction``.
+ """
+ if self._execution_context.is_run_asynchronous():
+ return self._prepare_async(rendered_query)
+ return self._execute_sync(rendered_query)
+
+ def _execute_sync(self, rendered_query: str) -> SqlJsonExecutionStatus:
+ """Run the query inline via ``execute_sql_lab_query`` and set the
result."""
+ from superset.sql.execution.sqllab_executor import
execute_sql_lab_query
+ from superset.sql_lab import get_query, handle_query_error
+
+ context = self._execution_context
+ query = context.query
+ timeout = app.config["SQLLAB_TIMEOUT"]
+ store_results = (
+ is_feature_enabled("SQLLAB_BACKEND_PERSISTENCE")
+ and not context.select_as_cta
+ )
+ try:
+ with utils.timeout(
+ seconds=timeout,
+ error_message=f"The query exceeded the {timeout} seconds
timeout.",
+ ):
+ try:
+ data = execute_sql_lab_query(
+ query,
+ rendered_query,
+ return_results=True,
+ store_results=store_results,
+ expand_data=context.expand_data,
+ log_params=self._log_params,
+ )
+ except Exception as ex: # pylint: disable=broad-except
+ # Re-fetch (the session may be poisoned) and build the
error
+ # payload, mirroring the classic synchronous path.
+ data = handle_query_error(ex, get_query(query_id=query.id))
+ except SupersetTimeoutException:
+ raise
+ except Exception as ex:
+ logger.exception("Query %i failed unexpectedly", query.id)
+ raise SupersetGenericDBErrorException(
+ utils.error_msg_from_exception(ex)
+ ) from ex
+
+ context.set_execution_result(data)
+ if data and data.get("status") == QueryStatus.FAILED:
+ if data.get("errors"):
+ errors = [SupersetError(**params) for params in data["errors"]]
+ status = (
+ 500
+ if any(error.level == ErrorLevel.ERROR for error in errors)
+ else 400
+ )
+ raise SupersetErrorsException(errors, status=status)
+ raise SupersetGenericDBErrorException(data["error"])
+ return SqlJsonExecutionStatus.HAS_RESULTS
+
+ def _prepare_async(self, rendered_query: str) -> SqlJsonExecutionStatus:
+ """Validate async prerequisites and defer scheduling to
``submit_async``.
+
+ Async SQL Lab execution runs as a GTF task, so it requires the
+ ``GLOBAL_TASK_FRAMEWORK`` feature flag; fail fast with a clear error
when
+ it is disabled rather than surfacing a raw framework error at schedule
time.
+ """
+ if not is_feature_enabled("GLOBAL_TASK_FRAMEWORK"):
+ error = SupersetError(
+ message=__(
+ "Asynchronous SQL Lab execution requires the "
+ "GLOBAL_TASK_FRAMEWORK feature flag to be enabled."
+ ),
+ error_type=SupersetErrorType.ASYNC_WORKERS_ERROR,
+ level=ErrorLevel.ERROR,
+ )
+ self._fail_query(error)
+ raise SupersetErrorException(error)
+ self._rendered_query = rendered_query
+ self._pending_async = True
+ return SqlJsonExecutionStatus.QUERY_IS_RUNNING
+
+ def submit_async(self) -> None:
+ """Schedule the async GTF SQL task, outside this command's transaction.
+
+ Called by the endpoint after ``run`` commits. Scheduling a GTF task
+ acquires its own lock/transaction and refuses to run inside an outer
+ ``@transaction`` (see ``SubmitTaskCommand``), so it must happen here.
+ """
+ if not self._pending_async:
+ return
+ from superset_core.tasks.types import TaskOptions
+
+ from superset.tasks.sql_queries import run_sql_lab_query
+
+ context = self._execution_context
+ query = context.query
+ try:
+ run_sql_lab_query.schedule(
+ query.id,
+ self._rendered_query,
+ store_results=not context.select_as_cta,
+ expand_data=context.expand_data,
+ username=get_username(),
+ start_time=now_as_float(),
+ log_params=self._log_params,
+ # PRIVATE dedup on the browser-generated client_id subsumes the
+ # classic ``is_query_handled`` idempotency guard.
+ options=TaskOptions(task_key=query.client_id),
+ )
Review Comment:
This is by design and matches the classic async path
(`ASynchronousSqlJsonExecutor` also used `store_results = not select_as_cta`).
Async results are retrieved later via `results_key`, so the async path always
persists and already *requires* a results backend — `execute_sql_lab_query`
raises `SupersetResultsBackendNotConfigureException` when `allow_run_async` is
set without one. `SQLLAB_BACKEND_PERSISTENCE` only gates persistence on the
*synchronous* path, which returns results inline. Making async honor it would
break async result retrieval.
--
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]
