diff --git a/superset-core/src/superset_core/api/models.py b/superset-core/src/superset_core/api/models.py index 2b78665a655..59cb07dc382 100644 --- a/superset-core/src/superset_core/api/models.py +++ b/superset-core/src/superset_core/api/models.py @@ -30,13 +30,22 @@ Usage: session = get_session() """ +from __future__ import annotations + from datetime import datetime -from typing import Any +from typing import Any, TYPE_CHECKING from uuid import UUID from flask_appbuilder import Model from sqlalchemy.orm import scoped_session +if TYPE_CHECKING: + from superset_core.api.types import ( + AsyncQueryHandle, + QueryOptions, + QueryResult, + ) + class CoreModel(Model): """ @@ -75,6 +84,83 @@ class Database(CoreModel): def data(self) -> dict[str, Any]: raise NotImplementedError + def execute( + self, + sql: str, + options: QueryOptions | None = None, + ) -> QueryResult: + """ + Execute SQL synchronously. + + :param sql: SQL query to execute + :param options: Query execution options (see `QueryOptions`). + If not provided, defaults are used. + :returns: QueryResult with status, data (DataFrame), and metadata + + Example: + from superset_core.api.daos import DatabaseDAO + from superset_core.api.types import QueryOptions, QueryStatus + + db = DatabaseDAO.find_one_or_none(id=1) + result = db.execute( + "SELECT * FROM users WHERE active = true", + options=QueryOptions(schema="public", limit=100) + ) + if result.status == QueryStatus.SUCCESS: + df = result.data + print(f"Found {sum(s.row_count for s in result.statements)} rows") + + Example with templates: + result = db.execute( + "SELECT * FROM {{ table }} WHERE date > '{{ start_date }}'", + options=QueryOptions( + schema="analytics", + template_params={"table": "events", "start_date": "2024-01-01"} + ) + ) + + Example with dry_run: + result = db.execute( + "SELECT * FROM users", + options=QueryOptions(schema="public", limit=100, dry_run=True) + ) + print(f"Would execute: {result.statements[0].statement}") + """ + raise NotImplementedError("Method will be replaced during initialization") + + def execute_async( + self, + sql: str, + options: QueryOptions | None = None, + ) -> AsyncQueryHandle: + """ + Execute SQL asynchronously. + + Returns immediately with a handle for tracking progress and retrieving + results from the background worker. + + :param sql: SQL query to execute + :param options: Query execution options (see `QueryOptions`). + If not provided, defaults are used. + :returns: AsyncQueryHandle for tracking the query + + Example: + handle = db.execute_async( + "SELECT * FROM large_table", + options=QueryOptions(schema="analytics") + ) + + # Check status and get results + status = handle.get_status() + if status == QueryStatus.SUCCESS: + query_result = handle.get_result() + df = query_result.statements[0].data + + # Cancel if needed + handle.cancel() + """ + raise NotImplementedError("Method will be replaced during initialization") + class Dataset(CoreModel): """ diff --git a/superset-core/src/superset_core/api/types.py b/superset-core/src/superset_core/api/types.py new file mode 100644 index 00000000000..1e4d5972657 --- /dev/null +++ b/superset-core/src/superset_core/api/types.py @@ -0,0 +1,177 @@ +# 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. + +""" +Query execution types for superset-core. + +Provides type definitions for query execution that are partially aligned +with frontend types in superset-ui-core/src/query/types/. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from datetime import datetime +from enum import Enum +from typing import Any, TYPE_CHECKING + +if TYPE_CHECKING: + import pandas as pd + + +class QueryStatus(Enum): + """ + Status of query execution. + """ + + PENDING = "pending" + RUNNING = "running" + SUCCESS = "success" + FAILED = "failed" + TIMED_OUT = "timed_out" + STOPPED = "stopped" + + +@dataclass +class CacheOptions: + """ + Options for query result caching. + """ + + timeout: int | None = None # Override default cache timeout (seconds) + force_refresh: bool = False # Bypass cache and re-execute query + + +@dataclass +class QueryOptions: + """ + Options for query execution via Database.execute() and execute_async(). + + Supports customization of: + - Basic: catalog, schema, limit, timeout + - Templates: Jinja2 template parameters + - Caching: Cache timeout and refresh control + - Dry run: Return transformed SQL without execution + """ + + # Basic options + catalog: str | None = None + schema: str | None = None + limit: int | None = None + timeout_seconds: int | None = None + + # Template options + template_params: dict[str, Any] | None = None # For Jinja2 rendering + + # Caching options + cache: CacheOptions | None = None + + # Dry run option + dry_run: bool = False # Return transformed SQL without executing + + +@dataclass +class StatementResult: + """ + Result of a single SQL statement execution. + + For SELECT queries: data contains DataFrame, row_count is len(data) + For DML queries: data is None, row_count contains affected rows + """ + + original_sql: str # The SQL statement as submitted by the user + executed_sql: ( + str # The SQL statement after transformations (RLS, mutations, limits) + ) + data: pd.DataFrame | None = None + row_count: int = 0 + execution_time_ms: float | None = None + + +@dataclass +class QueryResult: + """ + Result of a multi-statement query execution. + + On success: statements contains all executed statements + On failure: statements contains successful statements before failure + + Fields: + status: Overall query status (SUCCESS or FAILED) + statements: Results from each executed statement + query_id: Query model ID for entire execution (None if dry_run=True) + total_execution_time_ms: Total execution time across all statements + is_cached: Whether result came from cache + error_message: Query-level error (e.g., "Statement 2 of 3: error") + """ + + status: QueryStatus + statements: list[StatementResult] = field(default_factory=list) + query_id: int | None = None + total_execution_time_ms: float | None = None + is_cached: bool = False + error_message: str | None = None + + +@dataclass +class AsyncQueryHandle: + """ + Handle for tracking an asynchronous query. + + Provides methods to check status, retrieve results, and cancel the query. + The methods are bound to concrete implementations at runtime. + + This is the return type of Database.execute_async(). + """ + + query_id: int | None # None for cached results + status: QueryStatus = field(default=QueryStatus.PENDING) + started_at: datetime | None = None + + def get_status(self) -> QueryStatus: + """ + Get the current status of the async query. + + :returns: Current QueryStatus + """ + raise NotImplementedError("Method will be replaced during initialization") + + def get_result(self) -> QueryResult: + """ + Get the result of the async query. + + :returns: QueryResult with data if successful + """ + raise NotImplementedError("Method will be replaced during initialization") + + def cancel(self) -> bool: + """ + Cancel the async query. + + :returns: True if cancellation was successful + """ + raise NotImplementedError("Method will be replaced during initialization") + + +__all__ = [ + "QueryStatus", + "QueryOptions", + "QueryResult", + "StatementResult", + "AsyncQueryHandle", + "CacheOptions", +] diff --git a/superset/commands/security/create.py b/superset/commands/security/create.py index 65bbf28900f..f75bcdf9476 100644 --- a/superset/commands/security/create.py +++ b/superset/commands/security/create.py @@ -44,7 +44,9 @@ class CreateRLSRuleCommand(BaseCommand): def validate(self) -> None: roles = populate_roles(self._roles) tables = ( - db.session.query(SqlaTable).filter(SqlaTable.id.in_(self._tables)).all() # type: ignore[attr-defined] + db.session.query(SqlaTable) + .filter(SqlaTable.id.in_(self._tables)) # type: ignore[attr-defined] + .all() ) if len(tables) != len(self._tables): raise DatasourceNotFoundValidationError() diff --git a/superset/commands/security/update.py b/superset/commands/security/update.py index 54b1dce8d07..8a2e06c44c1 100644 --- a/superset/commands/security/update.py +++ b/superset/commands/security/update.py @@ -51,7 +51,9 @@ class UpdateRLSRuleCommand(BaseCommand): raise RLSRuleNotFoundError() roles = populate_roles(self._roles) tables = ( - db.session.query(SqlaTable).filter(SqlaTable.id.in_(self._tables)).all() # type: ignore[attr-defined] + db.session.query(SqlaTable) + .filter(SqlaTable.id.in_(self._tables)) # type: ignore[attr-defined] + .all() ) if len(tables) != len(self._tables): raise DatasourceNotFoundValidationError() diff --git a/superset/daos/base.py b/superset/daos/base.py index 557bae54e10..4897f796cf8 100644 --- a/superset/daos/base.py +++ b/superset/daos/base.py @@ -173,7 +173,7 @@ class BaseDAO(CoreBaseDAO[T], Generic[T]): def __init_subclass__(cls) -> None: cls.model_cls = get_args( - cls.__orig_bases__[0] # type: ignore # pylint: disable=no-member + cls.__orig_bases__[0] # type: ignore[attr-defined] # pylint: disable=no-member )[0] @classmethod diff --git a/superset/models/core.py b/superset/models/core.py index 6f96a7f433b..8fe49119968 100755 --- a/superset/models/core.py +++ b/superset/models/core.py @@ -95,6 +95,8 @@ metadata = Model.metadata # pylint: disable=no-member logger = logging.getLogger(__name__) if TYPE_CHECKING: + from superset_core.api.types import AsyncQueryHandle, QueryOptions, QueryResult + from superset.models.sql_lab import Query @@ -1276,6 +1278,38 @@ class Database(CoreDatabase, AuditMixinNullable, ImportExportMixin): # pylint: DatabaseUserOAuth2Tokens.id == self.id ).delete() + def execute( + self, + sql: str, + options: QueryOptions | None = None, + ) -> QueryResult: + """ + Execute SQL synchronously. + + :param sql: SQL query to execute + :param options: QueryOptions with execution settings + :returns: QueryResult with status, data, and metadata + """ + from superset.sql.execution import SQLExecutor + + return SQLExecutor(self).execute(sql, options) + + def execute_async( + self, + sql: str, + options: QueryOptions | None = None, + ) -> AsyncQueryHandle: + """ + Execute SQL asynchronously via Celery. + + :param sql: SQL query to execute + :param options: QueryOptions with execution settings + :returns: AsyncQueryHandle for tracking the query + """ + from superset.sql.execution import SQLExecutor + + return SQLExecutor(self).execute_async(sql, options) + sqla.event.listen(Database, "after_insert", security_manager.database_after_insert) sqla.event.listen(Database, "after_update", security_manager.database_after_update) diff --git a/superset/sql/execution/__init__.py b/superset/sql/execution/__init__.py new file mode 100644 index 00000000000..c376911abf9 --- /dev/null +++ b/superset/sql/execution/__init__.py @@ -0,0 +1,20 @@ +# 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. + +from .executor import SQLExecutor + +__all__ = ["SQLExecutor"] diff --git a/superset/sql/execution/celery_task.py b/superset/sql/execution/celery_task.py new file mode 100644 index 00000000000..c5dd1e9f1c9 --- /dev/null +++ b/superset/sql/execution/celery_task.py @@ -0,0 +1,486 @@ +# 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. +""" +Celery task for async SQL execution. + +This module provides the Celery task for executing SQL queries asynchronously. +It is used by SQLExecutor.execute_async() to run queries in the background. +""" + +from __future__ import annotations + +import dataclasses +import logging +import uuid +from typing import Any, TYPE_CHECKING + +import msgpack +from celery.exceptions import SoftTimeLimitExceeded +from flask import current_app as app, has_app_context +from flask_babel import gettext as __ + +from superset import ( + db, + results_backend, + security_manager, +) +from superset.common.db_query_status import QueryStatus +from superset.constants import QUERY_CANCEL_KEY +from superset.errors import ErrorLevel, SupersetError, SupersetErrorType +from superset.exceptions import ( + SupersetErrorException, + SupersetErrorsException, +) +from superset.extensions import celery_app +from superset.models.sql_lab import Query +from superset.result_set import SupersetResultSet +from superset.sql.execution.executor import execute_sql_with_cursor +from superset.sql.parse import SQLScript +from superset.sqllab.utils import write_ipc_buffer +from superset.utils import json +from superset.utils.core import override_user, zlib_compress +from superset.utils.dates import now_as_float +from superset.utils.decorators import stats_timing + +if TYPE_CHECKING: + pass + +logger = logging.getLogger(__name__) + +BYTES_IN_MB = 1024 * 1024 + + +def _get_query(query_id: int) -> Query: + """Get the query by ID.""" + return db.session.query(Query).filter_by(id=query_id).one() + + +def _handle_query_error( + ex: Exception, + query: Query, + payload: dict[str, Any] | None = None, + prefix_message: str = "", +) -> dict[str, Any]: + """Handle error while processing the SQL query.""" + payload = payload or {} + msg = f"{prefix_message} {str(ex)}".strip() + query.error_message = msg + query.tmp_table_name = None + # Preserve TIMED_OUT status if already set (from SoftTimeLimitExceeded handler) + if query.status != QueryStatus.TIMED_OUT: + query.status = QueryStatus.FAILED + + if not query.end_time: + query.end_time = now_as_float() + + # Extract DB-specific errors + if isinstance(ex, SupersetErrorException): + errors = [ex.error] + elif isinstance(ex, SupersetErrorsException): + errors = ex.errors + else: + errors = query.database.db_engine_spec.extract_errors( + str(ex), database_name=query.database.unique_name + ) + + errors_payload = [dataclasses.asdict(error) for error in errors] + if errors: + query.set_extra_json_key("errors", errors_payload) + + db.session.commit() # pylint: disable=consider-using-transaction + payload.update( + {"status": query.status.value, "error": msg, "errors": errors_payload} + ) + if troubleshooting_link := app.config.get("TROUBLESHOOTING_LINK"): + payload["link"] = troubleshooting_link + return payload + + +def _serialize_payload(payload: dict[Any, Any]) -> bytes: + """Serialize payload for storage based on RESULTS_BACKEND_USE_MSGPACK config.""" + from superset import results_backend_use_msgpack + + if results_backend_use_msgpack: + return msgpack.dumps(payload, default=json.json_iso_dttm_ser, use_bin_type=True) + return json.dumps(payload, default=json.json_iso_dttm_ser, ignore_nan=True).encode( + "utf-8" + ) + + +def _prepare_statement_blocks( + rendered_query: str, + db_engine_spec: Any, +) -> tuple[SQLScript, list[str]]: + """ + Parse SQL and build statement blocks for execution. + + Some databases (like BigQuery and Kusto) do not persist state across multiple + statements if they're run separately (especially when using `NullPool`), so we run + the query as a single block when the database engine spec requires it. + """ + parsed_script = SQLScript(rendered_query, engine=db_engine_spec.engine) + + # Build statement blocks for execution + if db_engine_spec.run_multiple_statements_as_one: + blocks = [parsed_script.format(comments=db_engine_spec.allows_sql_comments)] + else: + blocks = [ + statement.format(comments=db_engine_spec.allows_sql_comments) + for statement in parsed_script.statements + ] + + return parsed_script, blocks + + +def _finalize_successful_query( + query: Query, + original_script: SQLScript, + execution_results: list[tuple[str, SupersetResultSet | None, float, int]], + payload: dict[str, Any], + total_execution_time_ms: float, +) -> None: + """Update query metadata and payload after successful execution.""" + # Calculate total rows across all statements + total_rows = 0 + statements_data: list[dict[str, Any]] = [] + + # Get original statement strings + original_sqls = [stmt.format() for stmt in original_script.statements] + + for orig_sql, (exec_sql, result_set, exec_time, rowcount) in zip( + original_sqls, execution_results, strict=True + ): + if result_set is not None: + # SELECT statement + total_rows += result_set.size + data, columns = _serialize_result_set(result_set) + statements_data.append( + { + "original_sql": orig_sql, + "executed_sql": exec_sql, + "data": data, + "columns": columns, + "row_count": result_set.size, + "execution_time_ms": exec_time, + } + ) + else: + # DML statement - no data, just row count + statements_data.append( + { + "original_sql": orig_sql, + "executed_sql": exec_sql, + "data": None, + "columns": [], + "row_count": rowcount, + "execution_time_ms": exec_time, + } + ) + + query.rows = total_rows + query.progress = 100 + query.set_extra_json_key("progress", None) + # Store columns from last statement (for compatibility) + if execution_results and execution_results[-1][1] is not None: + query.set_extra_json_key("columns", execution_results[-1][1].columns) + query.end_time = now_as_float() + + payload.update( + { + "status": QueryStatus.SUCCESS.value, + "statements": statements_data, + "total_execution_time_ms": total_execution_time_ms, + "query": query.to_dict(), + } + ) + payload["query"]["state"] = QueryStatus.SUCCESS.value + + +def _store_results_in_backend( + query: Query, + payload: dict[str, Any], + database: Any, +) -> None: + """Store query results in the results backend.""" + key = str(uuid.uuid4()) + payload["query"]["resultsKey"] = key + logger.info( + "Query %s: Storing results in results backend, key: %s", + str(query.id), + key, + ) + stats_logger = app.config["STATS_LOGGER"] + with stats_timing("sqllab.query.results_backend_write", stats_logger): + with stats_timing( + "sqllab.query.results_backend_write_serialization", stats_logger + ): + serialized_payload = _serialize_payload(payload) + + # Check payload size limit + if sql_lab_payload_max_mb := app.config.get("SQLLAB_PAYLOAD_MAX_MB"): + serialized_payload_size = len(serialized_payload) + max_bytes = sql_lab_payload_max_mb * BYTES_IN_MB + + if serialized_payload_size > max_bytes: + logger.info("Result size exceeds the allowed limit.") + raise SupersetErrorException( + SupersetError( + message=( + f"Result size " + f"({serialized_payload_size / BYTES_IN_MB:.2f} MB) " + f"exceeds the allowed limit of " + f"{sql_lab_payload_max_mb} MB." + ), + error_type=SupersetErrorType.RESULT_TOO_LARGE_ERROR, + level=ErrorLevel.ERROR, + ) + ) + + cache_timeout = database.cache_timeout + if cache_timeout is None: + cache_timeout = app.config["CACHE_DEFAULT_TIMEOUT"] + + compressed = zlib_compress(serialized_payload) + logger.debug("*** serialized payload size: %i", len(serialized_payload)) + logger.debug("*** compressed payload size: %i", len(compressed)) + + write_success = results_backend.set(key, compressed, cache_timeout) + if not write_success: + logger.error( + "Query %s: Failed to store results in backend, key: %s", + str(query.id), + key, + ) + stats_logger.incr("sqllab.results_backend.write_failure") + query.results_key = None + query.status = QueryStatus.FAILED + query.error_message = ( + "Failed to store query results in the results backend. " + "Please try again or contact your administrator." + ) + db.session.commit() # pylint: disable=consider-using-transaction + raise SupersetErrorException( + SupersetError( + message=__("Failed to store query results. Please try again."), + error_type=SupersetErrorType.RESULTS_BACKEND_ERROR, + level=ErrorLevel.ERROR, + ) + ) + else: + query.results_key = key + logger.info( + "Query %s: Successfully stored results in backend, key: %s", + str(query.id), + key, + ) + + +def _serialize_result_set( + result_set: SupersetResultSet, +) -> tuple[bytes | list[Any], list[Any]]: + """ + Serialize result set based on RESULTS_BACKEND_USE_MSGPACK config. + + When msgpack is enabled, uses Apache Arrow IPC format for efficiency. + Otherwise, falls back to JSON-serializable records. + + :param result_set: Query result set to serialize + :returns: Tuple of (serialized_data, columns) + """ + from superset import results_backend_use_msgpack + from superset.dataframe import df_to_records + + if results_backend_use_msgpack: + if has_app_context(): + stats_logger = app.config["STATS_LOGGER"] + with stats_timing( + "sqllab.query.results_backend_pa_serialization", stats_logger + ): + data: bytes | list[Any] = write_ipc_buffer( + result_set.pa_table + ).to_pybytes() + else: + data = write_ipc_buffer(result_set.pa_table).to_pybytes() + else: + df = result_set.to_pandas_df() + data = df_to_records(df) or [] + + return (data, result_set.columns) + + +@celery_app.task(name="query_execution.execute_sql") +def execute_sql_task( + query_id: int, + rendered_query: str, + username: str | None = None, + start_time: float | None = None, +) -> dict[str, Any] | None: + """ + Execute SQL query asynchronously via Celery. + + This task is used by SQLExecutor.execute_async() to run queries + in background workers with full feature support. + + :param query_id: ID of the Query model + :param rendered_query: Pre-rendered SQL query to execute + :param username: Username for context override + :param start_time: Query start time for timing metrics + :returns: Query result payload or None + """ + with app.test_request_context(): + with override_user(security_manager.find_user(username)): + try: + return _execute_sql_statements( + query_id, + rendered_query, + start_time=start_time, + ) + except Exception as ex: + logger.debug("Query %d: %s", query_id, ex) + stats_logger = app.config["STATS_LOGGER"] + stats_logger.incr("error_sqllab_unhandled") + query = _get_query(query_id=query_id) + return _handle_query_error(ex, query) + + +def _make_check_stopped_fn(query: Query) -> Any: + """Create a function to check if query was stopped.""" + + def check_stopped() -> bool: + db.session.refresh(query) + return query.status == QueryStatus.STOPPED + + return check_stopped + + +def _make_execute_fn(query: Query, db_engine_spec: Any) -> Any: + """Create an execute function with stats timing.""" + + def execute_with_stats(cursor: Any, sql: str) -> None: + query.executed_sql = sql + stats_logger = app.config["STATS_LOGGER"] + with stats_timing("sqllab.query.time_executing_query", stats_logger): + db_engine_spec.execute_with_cursor(cursor, sql, query) + + return execute_with_stats + + +def _make_log_query_fn(database: Any) -> Any: + """Create a query logging function.""" + + def log_query(sql: str, schema: str | None) -> None: + if log_query_fn := app.config.get("QUERY_LOGGER"): + log_query_fn( + database.sqlalchemy_uri, + sql, + schema, + __name__, + security_manager, + None, + ) + + return log_query + + +def _execute_sql_statements( + query_id: int, + rendered_query: str, + start_time: float | None, +) -> dict[str, Any] | None: + """Execute SQL statements and store results.""" + if start_time: + stats_logger = app.config["STATS_LOGGER"] + stats_logger.timing("sqllab.query.time_pending", now_as_float() - start_time) + + query = _get_query(query_id=query_id) + payload: dict[str, Any] = {"query_id": query_id} + database = query.database + db_engine_spec = database.db_engine_spec + db_engine_spec.patch() + + logger.info("Query %s: Set query to 'running'", str(query_id)) + query.status = QueryStatus.RUNNING + query.start_running_time = now_as_float() + execution_start_time = now_as_float() + db.session.commit() # pylint: disable=consider-using-transaction + + # Parse original SQL (from user) to preserve before transformations + original_script = SQLScript(query.sql, engine=db_engine_spec.engine) + + # Parse transformed SQL (with RLS, limits, etc.) + parsed_script, blocks = _prepare_statement_blocks(rendered_query, db_engine_spec) + + with database.get_raw_connection( + catalog=query.catalog, + schema=query.schema, + ) as conn: + cursor = conn.cursor() + + cancel_query_id = db_engine_spec.get_cancel_query_id(cursor, query) + if cancel_query_id is not None: + query.set_extra_json_key(QUERY_CANCEL_KEY, cancel_query_id) + db.session.commit() # pylint: disable=consider-using-transaction + + try: + execution_results = execute_sql_with_cursor( + database=database, + cursor=cursor, + statements=blocks, + query=query, + log_query_fn=_make_log_query_fn(database), + check_stopped_fn=_make_check_stopped_fn(query), + execute_fn=_make_execute_fn(query, db_engine_spec), + ) + except SoftTimeLimitExceeded as ex: + query.status = QueryStatus.TIMED_OUT + logger.warning("Query %d: Time limit exceeded", query.id) + timeout_sec = app.config["SQLLAB_ASYNC_TIME_LIMIT_SEC"] + raise SupersetErrorException( + SupersetError( + message=__( + "The query was killed after %(sqllab_timeout)s seconds. " + "It might be too complex, or the database might be " + "under heavy load.", + sqllab_timeout=timeout_sec, + ), + error_type=SupersetErrorType.SQLLAB_TIMEOUT_ERROR, + level=ErrorLevel.ERROR, + ) + ) from ex + + # Check if stopped + if not execution_results: + payload.update({"status": QueryStatus.STOPPED.value}) + return payload + + # Commit for mutations + if parsed_script.has_mutation() or query.select_as_cta: + conn.commit() # pylint: disable=consider-using-transaction + + total_execution_time_ms = (now_as_float() - execution_start_time) * 1000 + _finalize_successful_query( + query, original_script, execution_results, payload, total_execution_time_ms + ) + + if results_backend: + _store_results_in_backend(query, payload, database) + + if query.status != QueryStatus.FAILED: + query.status = QueryStatus.SUCCESS + db.session.commit() # pylint: disable=consider-using-transaction + + return payload diff --git a/superset/sql/execution/executor.py b/superset/sql/execution/executor.py new file mode 100644 index 00000000000..b1f5e794049 --- /dev/null +++ b/superset/sql/execution/executor.py @@ -0,0 +1,1108 @@ +# 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. + +""" +SQL Executor implementation for Database.execute() and execute_async(). + +This module provides the SQLExecutor class that implements the query execution +methods defined in superset_core.api.models.Database. + +Implementation Features +----------------------- + +Query Preparation (applies to both sync and async): +- Jinja2 template rendering (via template_params in QueryOptions) +- SQL mutation via SQL_QUERY_MUTATOR config hook +- DML permission checking (requires database.allow_dml=True for DML) +- Disallowed functions checking via DISALLOWED_SQL_FUNCTIONS config +- Row-level security (RLS) via AST transformation (always applied) +- Result limit application via SQL_MAX_ROW config +- Catalog/schema resolution and validation + +Synchronous Execution (execute): +- Multi-statement SQL parsing and execution +- Progress tracking via Query model +- Result caching via cache_manager.data_cache +- Query logging via QUERY_LOGGER config hook +- Timeout protection via SQLLAB_TIMEOUT config +- Dry run mode (returns transformed SQL without execution) + +Asynchronous Execution (execute_async): +- Celery task submission for background execution +- Security validation before submission +- Query model creation with PENDING status +- Result caching check (returns cached if available) +- Background execution with timeout via SQLLAB_ASYNC_TIME_LIMIT_SEC +- Results stored in results backend for retrieval +- Handle-based progress tracking and cancellation + +See Database.execute() and Database.execute_async() docstrings in +superset_core.api.models for the public API contract. +""" + +from __future__ import annotations + +import contextlib +import hashlib +import logging +import time +import uuid +from datetime import datetime +from typing import Any, TYPE_CHECKING + +from flask import current_app as app, g, has_app_context + +from superset import db +from superset.errors import ErrorLevel, SupersetError, SupersetErrorType +from superset.exceptions import ( + SupersetSecurityException, + SupersetTimeoutException, +) +from superset.extensions import cache_manager +from superset.sql.parse import SQLScript +from superset.utils import core as utils + +if TYPE_CHECKING: + from superset_core.api.types import ( + AsyncQueryHandle, + QueryOptions, + QueryResult, + QueryStatus, + StatementResult, + ) + + from superset.models.core import Database + from superset.result_set import SupersetResultSet + +logger = logging.getLogger(__name__) + + +def execute_sql_with_cursor( + database: Database, + cursor: Any, + statements: list[str], + query: Any, + log_query_fn: Any | None = None, + check_stopped_fn: Any | None = None, + execute_fn: Any | None = None, +) -> list[tuple[str, SupersetResultSet | None, float, int]]: + """ + Execute SQL statements with a cursor and return all result sets. + + This is the shared execution logic used by both sync (SQLExecutor) and + async (celery_task) execution paths. It handles multi-statement execution + with progress tracking via the Query model. + + :param database: Database model to execute against + :param cursor: Database cursor to use for execution + :param statements: List of SQL statements to execute + :param query: Query model for progress tracking + :param log_query_fn: Optional function to log queries, called as fn(sql, schema) + :param check_stopped_fn: Optional function to check if query was stopped. + Should return True if stopped. Used by async execution for cancellation. + :param execute_fn: Optional custom execute function. If not provided, uses + database.db_engine_spec.execute(cursor, sql, database). Custom function + should accept (cursor, sql) and handle execution. + :returns: List of (statement_sql, result_set, execution_time_ms, rowcount) tuples + Returns empty list if stopped. Raises exception on error (fail-fast). + """ + from superset.result_set import SupersetResultSet + + total = len(statements) + if total == 0: + return [] + + results: list[tuple[str, SupersetResultSet | None, float, int]] = [] + + for i, statement in enumerate(statements): + # Check if query was stopped (async cancellation) + if check_stopped_fn and check_stopped_fn(): + return results + + stmt_start_time = time.time() + + # Apply SQL mutation + stmt_sql = database.mutate_sql_based_on_config( + statement, + is_split=True, + ) + + # Log query + if log_query_fn: + log_query_fn(stmt_sql, query.schema) + + # Execute - use custom function or default + if execute_fn: + execute_fn(cursor, stmt_sql) + else: + database.db_engine_spec.execute(cursor, stmt_sql, database) + + stmt_execution_time = (time.time() - stmt_start_time) * 1000 + + # Fetch results from ALL statements + description = cursor.description + if description: + rows = database.db_engine_spec.fetch_data(cursor) + result_set = SupersetResultSet( + rows, + description, + database.db_engine_spec, + ) + else: + # DML statement - no result set + result_set = None + + # Get row count for DML statements + rowcount = cursor.rowcount if hasattr(cursor, "rowcount") else 0 + + results.append((stmt_sql, result_set, stmt_execution_time, rowcount)) + + # Update progress on Query model + progress_pct = int(((i + 1) / total) * 100) + query.progress = progress_pct + query.set_extra_json_key( + "progress", + f"Running statement {i + 1} of {total}", + ) + db.session.commit() # pylint: disable=consider-using-transaction + + return results + + +class SQLExecutor: + """ + SQL query executor implementation. + + Implements Database.execute() and execute_async() methods. + See superset_core.api.models.Database for the full public API documentation. + """ + + def __init__(self, database: Database) -> None: + """ + Initialize the executor with a database. + + :param database: Database model instance to execute queries against + """ + self.database = database + + def execute( + self, + sql: str, + options: QueryOptions | None = None, + ) -> QueryResult: + """ + Execute SQL synchronously. + + If options.dry_run=True, returns the transformed SQL without execution. + All transformations (RLS, templates, limits) are still applied. + + See superset_core.api.models.Database.execute() for full documentation. + """ + from superset_core.api.types import ( + QueryOptions as QueryOptionsType, + QueryResult as QueryResultType, + QueryStatus, + StatementResult, + ) + + opts: QueryOptionsType = options or QueryOptionsType() + start_time = time.time() + + try: + # 1. Prepare SQL (assembly only, no security checks) + original_script, transformed_script, catalog, schema = self._prepare_sql( + sql, opts + ) + + # 2. Security checks on transformed script + self._check_security(transformed_script) + + # 3. Get mutation status and format SQL + has_mutation = transformed_script.has_mutation() + final_sql = transformed_script.format() + + # DRY RUN: Return transformed SQL without execution + if opts.dry_run: + total_execution_time_ms = (time.time() - start_time) * 1000 + # Create a StatementResult for each statement in dry-run mode + original_sqls = [stmt.format() for stmt in original_script.statements] + transformed_sqls = [ + stmt.format() for stmt in transformed_script.statements + ] + dry_run_statements = [ + StatementResult( + original_sql=orig_sql, + executed_sql=trans_sql, + data=None, + row_count=0, + execution_time_ms=0, + ) + for orig_sql, trans_sql in zip( + original_sqls, transformed_sqls, strict=True + ) + ] + return QueryResultType( + status=QueryStatus.SUCCESS, + statements=dry_run_statements, + query_id=None, + total_execution_time_ms=total_execution_time_ms, + is_cached=False, + ) + + # 4. Check cache + cached_result = self._try_get_cached_result(has_mutation, final_sql, opts) + if cached_result: + return cached_result + + # 5. Create Query model for audit + query = self._create_query_record( + final_sql, opts, catalog, schema, status=QueryStatus.RUNNING + ) + + # 6. Execute with timeout + timeout = opts.timeout_seconds or app.config.get("SQLLAB_TIMEOUT", 30) + timeout_msg = f"Query exceeded the {timeout} seconds timeout." + + with utils.timeout(seconds=timeout, error_message=timeout_msg): + statement_results = self._execute_statements( + original_script, + transformed_script, + catalog, + schema, + query, + ) + + total_execution_time_ms = (time.time() - start_time) * 1000 + + # Calculate total row count for Query model + total_rows = sum(stmt.row_count for stmt in statement_results) + + # Update query record + query.status = "success" + query.rows = total_rows + query.progress = 100 + db.session.commit() # pylint: disable=consider-using-transaction + + result = QueryResultType( + status=QueryStatus.SUCCESS, + statements=statement_results, + query_id=query.id, + total_execution_time_ms=total_execution_time_ms, + ) + + # Store in cache (if SELECT and caching enabled) + if not has_mutation: + self._store_in_cache(result, final_sql, opts) + + return result + + except SupersetTimeoutException: + return self._create_error_result( + QueryStatus.TIMED_OUT, + "Query exceeded the timeout limit", + start_time, + ) + except SupersetSecurityException as ex: + return self._create_error_result(QueryStatus.FAILED, str(ex), start_time) + except Exception as ex: + error_msg = self.database.db_engine_spec.extract_error_message(ex) + return self._create_error_result(QueryStatus.FAILED, error_msg, start_time) + + def execute_async( + self, + sql: str, + options: QueryOptions | None = None, + ) -> AsyncQueryHandle: + """ + Execute SQL asynchronously via Celery. + + If options.dry_run=True, returns the transformed SQL as a completed + AsyncQueryHandle without submitting to Celery. + + See superset_core.api.models.Database.execute_async() for full documentation. + """ + from superset_core.api.types import ( + QueryOptions as QueryOptionsType, + QueryResult as QueryResultType, + QueryStatus, + ) + + opts: QueryOptionsType = options or QueryOptionsType() + + # 1. Prepare SQL (assembly only, no security checks) + original_script, transformed_script, catalog, schema = self._prepare_sql( + sql, opts + ) + + # 2. Security checks on transformed script + self._check_security(transformed_script) + + # 3. Get mutation status and format SQL + has_mutation = transformed_script.has_mutation() + final_sql = transformed_script.format() + + # DRY RUN: Return transformed SQL as completed async handle + if opts.dry_run: + from superset_core.api.types import StatementResult + + original_sqls = [stmt.format() for stmt in original_script.statements] + transformed_sqls = [stmt.format() for stmt in transformed_script.statements] + dry_run_statements = [ + StatementResult( + original_sql=orig_sql, + executed_sql=trans_sql, + data=None, + row_count=0, + execution_time_ms=0, + ) + for orig_sql, trans_sql in zip( + original_sqls, transformed_sqls, strict=True + ) + ] + dry_run_result = QueryResultType( + status=QueryStatus.SUCCESS, + statements=dry_run_statements, + query_id=None, + total_execution_time_ms=0, + is_cached=False, + ) + return self._create_cached_handle(dry_run_result) + + # 4. Check cache + if cached_result := self._try_get_cached_result(has_mutation, final_sql, opts): + return self._create_cached_handle(cached_result) + + # 5. Create Query model for audit + query = self._create_query_record( + final_sql, opts, catalog, schema, status=QueryStatus.PENDING + ) + + # 6. Submit to Celery + self._submit_query_to_celery(query, final_sql) + + # 7. Create and return handle with bound methods + return self._create_async_handle(query.id) + + def _prepare_sql( + self, + sql: str, + opts: QueryOptions, + ) -> tuple[SQLScript, SQLScript, str | None, str | None]: + """ + Prepare SQL for execution (no side effects, no security checks). + + This method performs SQL preprocessing: + 1. Template rendering + 2. SQL parsing + 3. Catalog/schema resolution + 4. RLS application + 5. Limit application (if not mutation) + + Security checks (disallowed functions, DML permission) are performed + by the caller after receiving the prepared scripts. + + :param sql: Original SQL query + :param opts: Query options + :returns: Tuple of (original_script, transformed_script, catalog, schema) + """ + # 1. Render Jinja2 templates + rendered_sql = self._render_sql_template(sql, opts.template_params) + + # 2. Parse SQL with SQLScript - this is the ORIGINAL script + original_script = SQLScript(rendered_sql, self.database.db_engine_spec.engine) + + # 3. Create a copy for transformation + transformed_script = SQLScript( + rendered_sql, self.database.db_engine_spec.engine + ) + + # 4. Get catalog and schema + catalog = opts.catalog or self.database.get_default_catalog() + schema = opts.schema or self.database.get_default_schema(catalog) + + # 5. Apply RLS to transformed script only + self._apply_rls_to_script(transformed_script, catalog, schema) + + # 6. Apply limit only if not a mutation + if not transformed_script.has_mutation(): + self._apply_limit_to_script(transformed_script, opts) + + return original_script, transformed_script, catalog, schema + + def _check_security(self, script: SQLScript) -> None: + """ + Perform security checks on prepared SQL script. + + :param script: Prepared SQLScript + :raises SupersetSecurityException: If security checks fail + """ + # Check disallowed functions + if disallowed := self._check_disallowed_functions(script): + raise SupersetSecurityException( + SupersetError( + message=f"Disallowed SQL functions: {', '.join(disallowed)}", + error_type=SupersetErrorType.INVALID_SQL_ERROR, + level=ErrorLevel.ERROR, + ) + ) + + # Check DML permission + if script.has_mutation() and not self.database.allow_dml: + raise SupersetSecurityException( + SupersetError( + message="DML queries are not allowed on this database", + error_type=SupersetErrorType.DML_NOT_ALLOWED_ERROR, + level=ErrorLevel.ERROR, + ) + ) + + def _execute_statements( + self, + original_script: SQLScript, + transformed_script: SQLScript, + catalog: str | None, + schema: str | None, + query: Any, + ) -> list[StatementResult]: + """ + Execute SQL statements and return per-statement results. + + Progress is tracked via Query.progress field. + Uses the same execution path for both single and multi-statement queries. + + :param original_script: SQLScript with original SQL (before transformations) + :param transformed_script: SQLScript with transformed SQL + (after RLS, limits, etc.) + :param catalog: Catalog name + :param schema: Schema name + :param query: Query model for progress tracking + :returns: List of StatementResult objects + """ + from superset_core.api.types import StatementResult + + # Get original statement strings + original_sqls = [stmt.format() for stmt in original_script.statements] + + # Handle empty script + if not transformed_script.statements: + return [] + + results_list = [] + + # Use consistent execution path for all queries + with self.database.get_raw_connection(catalog=catalog, schema=schema) as conn: + with contextlib.closing(conn.cursor()) as cursor: + execution_results = execute_sql_with_cursor( + database=self.database, + cursor=cursor, + statements=[ + stmt.format() for stmt in transformed_script.statements + ], + query=query, + log_query_fn=self._log_query, + ) + + # If execution was stopped or returned no results, return early + if not execution_results: + return [] + + # Build StatementResult for each executed statement + # with both original and executed SQL + for orig_sql, (exec_sql, result_set, exec_time, rowcount) in zip( + original_sqls, execution_results, strict=True + ): + if result_set is not None: + # SELECT statement + df = result_set.to_pandas_df() + stmt_result = StatementResult( + original_sql=orig_sql, + executed_sql=exec_sql, + data=df, + row_count=len(df), + execution_time_ms=exec_time, + ) + else: + # DML statement - no data, just row count + stmt_result = StatementResult( + original_sql=orig_sql, + executed_sql=exec_sql, + data=None, + row_count=rowcount, + execution_time_ms=exec_time, + ) + + results_list.append(stmt_result) + + return results_list + + def _log_query( + self, + sql: str, + schema: str | None, + ) -> None: + """ + Log query using QUERY_LOGGER config. + + :param sql: SQL to log + :param schema: Schema name + """ + from superset import security_manager + + if log_query := app.config.get("QUERY_LOGGER"): + log_query( + self.database, + sql, + schema, + __name__, + security_manager, + {}, + ) + + def _create_error_result( + self, + status: Any, + error_message: str, + start_time: float, + partial_results: list[Any] | None = None, + ) -> QueryResult: + """ + Create a QueryResult for error cases. + + :param status: QueryStatus enum value + :param error_message: Error message to include + :param start_time: Start time for calculating execution duration + :param partial_results: Optional list of StatementResult from successful + statements before the failure + :returns: QueryResult with error status + """ + from superset_core.api.types import QueryResult as QueryResultType + + return QueryResultType( + status=status, + statements=partial_results or [], + error_message=error_message, + total_execution_time_ms=(time.time() - start_time) * 1000, + ) + + def _render_sql_template( + self, sql: str, template_params: dict[str, Any] | None + ) -> str: + """ + Render Jinja2 template with params. + + :param sql: SQL string potentially containing Jinja2 templates + :param template_params: Parameters to pass to the template + :returns: Rendered SQL string + """ + if template_params is None: + return sql + + from superset.jinja_context import get_template_processor + + tp = get_template_processor(database=self.database) + return tp.process_template(sql, **template_params) + + def _apply_limit_to_script(self, script: SQLScript, opts: QueryOptions) -> None: + """ + Apply limit to the last statement in the script in place. + + :param script: SQLScript object to modify + :param opts: Query options + """ + # Skip if no limit requested + if opts.limit is None: + return + + sql_max_row = app.config.get("SQL_MAX_ROW") + effective_limit = opts.limit + if sql_max_row and opts.limit > sql_max_row: + effective_limit = sql_max_row + + # Apply limit to last statement only + if script.statements: + script.statements[-1].set_limit_value( + effective_limit, + self.database.db_engine_spec.limit_method, + ) + + def _try_get_cached_result( + self, + has_mutation: bool, + sql: str, + opts: QueryOptions, + ) -> QueryResult | None: + """ + Try to get a cached result if conditions allow. + + :param has_mutation: Whether the query contains mutations (DML) + :param sql: SQL query + :param opts: Query options + :returns: Cached QueryResult or None + """ + if has_mutation or (opts.cache and opts.cache.force_refresh): + return None + + return self._get_from_cache(sql, opts) + + def _check_disallowed_functions(self, script: SQLScript) -> set[str] | None: + """ + Check for disallowed SQL functions. + + :param script: Parsed SQL script + :returns: Set of disallowed functions found, or None if none found + """ + disallowed_config = app.config.get("DISALLOWED_SQL_FUNCTIONS", {}) + engine_name = self.database.db_engine_spec.engine + + # Get disallowed functions for this engine + engine_disallowed = disallowed_config.get(engine_name, set()) + if not engine_disallowed: + return None + + # Check each statement for disallowed functions + found = set() + for statement in script.statements: + # Use the statement's AST to check for function calls + statement_str = str(statement).upper() + for func in engine_disallowed: + if func.upper() in statement_str: + found.add(func) + + return found if found else None + + def _apply_rls_to_script( + self, script: SQLScript, catalog: str | None, schema: str | None + ) -> None: + """ + Apply Row-Level Security to SQLScript statements in place. + + :param script: SQLScript object to modify + :param catalog: Catalog name + :param schema: Schema name + """ + from superset.utils.rls import apply_rls + + # Apply RLS to each statement in the script + for statement in script.statements: + apply_rls(self.database, catalog, schema or "", statement) + + def _create_query_record( + self, + sql: str, + opts: QueryOptions, + catalog: str | None, + schema: str | None, + status: QueryStatus, + ) -> Any: + """ + Create Query model for audit/tracking. + + :param sql: SQL to execute + :param opts: Query options + :param catalog: Catalog name + :param schema: Schema name + :param status: Initial QueryStatus (RUNNING for sync, PENDING for async) + :returns: Query model instance + """ + from superset.models.sql_lab import Query as QueryModel + + user_id = None + if has_app_context() and hasattr(g, "user") and g.user: + user_id = g.user.get_id() + + # Generate client_id for Query model + client_id = uuid.uuid4().hex[:11] + + query = QueryModel( + client_id=client_id, + database_id=self.database.id, + sql=sql, + catalog=catalog, + schema=schema, + user_id=user_id, + status=status.value, + limit=opts.limit, + ) + db.session.add(query) + db.session.commit() # pylint: disable=consider-using-transaction + + return query + + def _get_from_cache(self, sql: str, opts: QueryOptions) -> QueryResult | None: + """ + Check results cache for existing result. + + :param sql: SQL query + :param opts: Query options + :returns: Cached QueryResult if found, None otherwise + """ + from superset_core.api.types import ( + QueryResult as QueryResultType, + QueryStatus, + StatementResult, + ) + + cache_key = self._generate_cache_key(sql, opts) + + if (cached := cache_manager.data_cache.get(cache_key)) is not None: + # Reconstruct statement results from cached data + statements = [ + StatementResult( + original_sql=stmt_data["original_sql"], + executed_sql=stmt_data["executed_sql"], + data=stmt_data["data"], + row_count=stmt_data["row_count"], + execution_time_ms=stmt_data["execution_time_ms"], + ) + for stmt_data in cached.get("statements", []) + ] + return QueryResultType( + status=QueryStatus.SUCCESS, + statements=statements, + query_id=None, + is_cached=True, + total_execution_time_ms=cached.get("total_execution_time_ms", 0), + ) + + return None + + def _store_in_cache( + self, result: QueryResult, sql: str, opts: QueryOptions + ) -> None: + """ + Store result in cache. + + :param result: Query result to cache + :param sql: SQL query (for cache key) + :param opts: Query options + """ + from superset_core.api.types import QueryStatus + + if result.status != QueryStatus.SUCCESS: + return + + cache_key = self._generate_cache_key(sql, opts) + timeout = ( + (opts.cache.timeout if opts.cache else None) + or self.database.cache_timeout + or app.config.get("CACHE_DEFAULT_TIMEOUT", 300) + ) + + # Serialize statement results for caching + cached_data = { + "statements": [ + { + "original_sql": stmt.original_sql, + "executed_sql": stmt.executed_sql, + "data": stmt.data, + "row_count": stmt.row_count, + "execution_time_ms": stmt.execution_time_ms, + } + for stmt in result.statements + ], + "total_execution_time_ms": result.total_execution_time_ms, + } + + cache_manager.data_cache.set( + cache_key, + cached_data, + timeout=timeout, + ) + + def _generate_cache_key(self, sql: str, opts: QueryOptions) -> str: + """ + Generate cache key for query result. + + :param sql: SQL query + :param opts: Query options + :returns: Cache key string + """ + # Include relevant options in the cache key + key_parts = [ + str(self.database.id), + sql, + opts.catalog or "", + opts.schema or "", + str(opts.limit) if opts.limit is not None else "", + ] + key_string = "|".join(key_parts) + return hashlib.sha256(key_string.encode()).hexdigest() + + def _submit_query_to_celery( + self, + query: Any, + rendered_sql: str, + ) -> None: + """ + Submit query to Celery for async execution. + + :param query: Query model instance + :param rendered_sql: Rendered SQL to execute + :raises: Re-raises any exception after marking query as failed + """ + from superset.sql.execution.celery_task import execute_sql_task + from superset.utils.core import get_username + from superset.utils.dates import now_as_float + + try: + task = execute_sql_task.delay( + query.id, + rendered_sql, + username=get_username(), + start_time=now_as_float(), + ) + task.forget() # Don't track task result in Celery backend + except Exception as ex: + query.status = "failed" + query.error_message = str(ex) + db.session.commit() # pylint: disable=consider-using-transaction + raise + + def _create_async_handle(self, query_id: int) -> AsyncQueryHandle: + """ + Create AsyncQueryHandle with bound methods for tracking the query. + + :param query_id: ID of the Query model + :returns: AsyncQueryHandle with configured methods + """ + from superset_core.api.types import ( + AsyncQueryHandle as AsyncQueryHandleType, + QueryResult as QueryResultType, + QueryStatus, + ) + + handle = AsyncQueryHandleType( + query_id=query_id, + status=QueryStatus.PENDING, + started_at=datetime.now(), + ) + + # Create bound closures for handle methods + def get_status_impl() -> QueryStatus: + return SQLExecutor._get_async_query_status(query_id) + + def get_result_impl() -> QueryResultType: + return SQLExecutor._get_async_query_result(query_id) + + def cancel_impl() -> bool: + return SQLExecutor._cancel_async_query(query_id, self.database) + + # Use object.__setattr__ to bypass dataclass frozen-like behavior + object.__setattr__(handle, "get_status", get_status_impl) + object.__setattr__(handle, "get_result", get_result_impl) + object.__setattr__(handle, "cancel", cancel_impl) + + return handle + + def _create_cached_handle(self, cached_result: QueryResult) -> AsyncQueryHandle: + """ + Create AsyncQueryHandle for a cached result. + + When cache hits occur for async queries, we return an AsyncQueryHandle + that immediately provides the cached data without submitting to Celery. + + :param cached_result: The cached QueryResult + :returns: AsyncQueryHandle that returns the cached data + """ + from superset_core.api.types import ( + AsyncQueryHandle as AsyncQueryHandleType, + QueryResult as QueryResultType, + QueryStatus, + ) + + handle = AsyncQueryHandleType( + query_id=None, + status=QueryStatus.SUCCESS, + started_at=datetime.now(), + ) + + # Create closures that return the cached result + def get_status_impl() -> QueryStatus: + return QueryStatus.SUCCESS + + def get_result_impl() -> QueryResultType: + return cached_result + + def cancel_impl() -> bool: + return False # Nothing to cancel for cached results + + object.__setattr__(handle, "get_status", get_status_impl) + object.__setattr__(handle, "get_result", get_result_impl) + object.__setattr__(handle, "cancel", cancel_impl) + + return handle + + @staticmethod + def _get_async_query_status(query_id: int) -> Any: + """Get the current status of an async query.""" + from superset_core.api.types import QueryStatus as QueryStatusType + + from superset.models.sql_lab import Query as QueryModel + + query = db.session.query(QueryModel).filter_by(id=query_id).one_or_none() + if not query: + return QueryStatusType.FAILED + + status_map = { + "pending": QueryStatusType.PENDING, + "running": QueryStatusType.RUNNING, + "success": QueryStatusType.SUCCESS, + "failed": QueryStatusType.FAILED, + "timed_out": QueryStatusType.TIMED_OUT, + "stopped": QueryStatusType.STOPPED, + } + return status_map.get(query.status, QueryStatusType.FAILED) + + @staticmethod + def _get_async_query_result(query_id: int) -> Any: + """Get the result of an async query.""" + import pandas as pd + from superset_core.api.types import ( + QueryResult as QueryResultType, + QueryStatus as QueryStatusType, + StatementResult, + ) + + from superset.models.sql_lab import Query as QueryModel + + query = db.session.query(QueryModel).filter_by(id=query_id).one_or_none() + if not query: + return QueryResultType( + status=QueryStatusType.FAILED, + error_message="Query not found", + ) + + status = SQLExecutor._get_async_query_status(query_id) + if status != QueryStatusType.SUCCESS: + return QueryResultType( + status=status, + error_message=query.error_message, + query_id=query_id, + ) + + # Fetch results from results backend + if query.results_key: + import msgpack + + from superset import results_backend_manager + + results_backend = results_backend_manager.results_backend + if results_backend is not None: + blob = results_backend.get(query.results_key) + if blob: + try: + from superset.utils.core import zlib_decompress + + payload = msgpack.loads(zlib_decompress(blob)) + + statements = [ + StatementResult( + original_sql=stmt_data.get("original_sql", ""), + executed_sql=stmt_data.get("executed_sql", ""), + data=( + pd.DataFrame( + stmt_data.get("data", []), + columns=[ + c.get("column_name", c.get("name", "")) + for c in stmt_data.get("columns", []) + ], + ) + if stmt_data.get("data") + else None + ), + row_count=stmt_data.get("row_count", 0), + execution_time_ms=stmt_data.get("execution_time_ms"), + ) + for stmt_data in payload.get("statements", []) + ] + return QueryResultType( + status=QueryStatusType.SUCCESS, + statements=statements, + query_id=query_id, + total_execution_time_ms=payload.get( + "total_execution_time_ms" + ), + is_cached=True, + ) + except Exception as ex: + logger.exception("Error loading async query results") + return QueryResultType( + status=QueryStatusType.FAILED, + error_message=f"Error loading results: {str(ex)}", + query_id=query_id, + ) + + return QueryResultType( + status=QueryStatusType.FAILED, + error_message="Results not available", + query_id=query_id, + ) + + @staticmethod + def _cancel_async_query(query_id: int, database: Database) -> bool: + """Cancel an async query.""" + from superset.models.sql_lab import Query as QueryModel + + query = db.session.query(QueryModel).filter_by(id=query_id).one_or_none() + if not query: + return False + + return SQLExecutor._cancel_query(database, query) + + @staticmethod + def _cancel_query(database: Database, query: Any) -> bool: + """ + Cancel a running query. + + This method handles query cancellation for different database types. + Some databases have implicit cancellation, others require explicit + cursor-based cancellation. + + :param database: Database model instance + :param query: Query model instance to cancel + :returns: True if cancelled successfully, False otherwise + """ + from superset.constants import QUERY_CANCEL_KEY, QUERY_EARLY_CANCEL_KEY + from superset.utils.core import QuerySource + + # Some engines implicitly handle cancellation + if database.db_engine_spec.has_implicit_cancel(): + return True + + # Some databases may need to make preparations for query cancellation + database.db_engine_spec.prepare_cancel_query(query) + + # Check early cancellation flag + if query.extra.get(QUERY_EARLY_CANCEL_KEY): + return True + + # Get cancel ID + cancel_query_id = query.extra.get(QUERY_CANCEL_KEY) + if cancel_query_id is None: + return False + + # Execute cancellation + with database.get_sqla_engine( + catalog=query.catalog, + schema=query.schema, + source=QuerySource.SQL_LAB, + ) as engine: + with contextlib.closing(engine.raw_connection()) as conn: + with contextlib.closing(conn.cursor()) as cursor: + return database.db_engine_spec.cancel_query( + cursor, query, cancel_query_id + ) diff --git a/tests/unit_tests/models/core_test.py b/tests/unit_tests/models/core_test.py index 26e4cd2736f..2667d359d2e 100644 --- a/tests/unit_tests/models/core_test.py +++ b/tests/unit_tests/models/core_test.py @@ -1060,3 +1060,87 @@ def test_apply_limit_to_sql( limited = db.apply_limit_to_sql(sql, limit, force) assert limited == expected + + +def test_database_execute_delegates_to_sql_executor(mocker: MockerFixture) -> None: + """Test that Database.execute() delegates to SQLExecutor.execute().""" + from unittest.mock import MagicMock + + mock_executor_class = mocker.patch("superset.sql.execution.SQLExecutor") + mock_executor = MagicMock() + mock_executor_class.return_value = mock_executor + + mock_result = MagicMock() + mock_executor.execute.return_value = mock_result + + database = Database(database_name="test_db", sqlalchemy_uri="sqlite://") + mock_options = MagicMock() + + result = database.execute("SELECT 1", mock_options) + + mock_executor_class.assert_called_once_with(database) + mock_executor.execute.assert_called_once_with("SELECT 1", mock_options) + assert result == mock_result + + +def test_database_execute_without_options(mocker: MockerFixture) -> None: + """Test that Database.execute() works without options.""" + from unittest.mock import MagicMock + + mock_executor_class = mocker.patch("superset.sql.execution.SQLExecutor") + mock_executor = MagicMock() + mock_executor_class.return_value = mock_executor + + mock_result = MagicMock() + mock_executor.execute.return_value = mock_result + + database = Database(database_name="test_db", sqlalchemy_uri="sqlite://") + + result = database.execute("SELECT 1") + + mock_executor_class.assert_called_once_with(database) + mock_executor.execute.assert_called_once_with("SELECT 1", None) + assert result == mock_result + + +def test_database_execute_async_delegates_to_sql_executor( + mocker: MockerFixture, +) -> None: + """Test that Database.execute_async() delegates to SQLExecutor.execute_async().""" + from unittest.mock import MagicMock + + mock_executor_class = mocker.patch("superset.sql.execution.SQLExecutor") + mock_executor = MagicMock() + mock_executor_class.return_value = mock_executor + + mock_handle = MagicMock() + mock_executor.execute_async.return_value = mock_handle + + database = Database(database_name="test_db", sqlalchemy_uri="sqlite://") + mock_options = MagicMock() + + result = database.execute_async("SELECT 1", mock_options) + + mock_executor_class.assert_called_once_with(database) + mock_executor.execute_async.assert_called_once_with("SELECT 1", mock_options) + assert result == mock_handle + + +def test_database_execute_async_without_options(mocker: MockerFixture) -> None: + """Test that Database.execute_async() works without options.""" + from unittest.mock import MagicMock + + mock_executor_class = mocker.patch("superset.sql.execution.SQLExecutor") + mock_executor = MagicMock() + mock_executor_class.return_value = mock_executor + + mock_handle = MagicMock() + mock_executor.execute_async.return_value = mock_handle + + database = Database(database_name="test_db", sqlalchemy_uri="sqlite://") + + result = database.execute_async("SELECT 1") + + mock_executor_class.assert_called_once_with(database) + mock_executor.execute_async.assert_called_once_with("SELECT 1", None) + assert result == mock_handle diff --git a/tests/unit_tests/sql/execution/__init__.py b/tests/unit_tests/sql/execution/__init__.py new file mode 100644 index 00000000000..13a83393a91 --- /dev/null +++ b/tests/unit_tests/sql/execution/__init__.py @@ -0,0 +1,16 @@ +# 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. diff --git a/tests/unit_tests/sql/execution/conftest.py b/tests/unit_tests/sql/execution/conftest.py new file mode 100644 index 00000000000..630f6884529 --- /dev/null +++ b/tests/unit_tests/sql/execution/conftest.py @@ -0,0 +1,324 @@ +# 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. + +""" +Shared fixtures and helpers for SQL execution tests. + +This module provides common mocks, fixtures, and helper functions used across +test_celery_task.py and test_executor.py to reduce code duplication. +""" + +from contextlib import contextmanager +from typing import Any +from unittest.mock import MagicMock + +import pandas as pd +import pytest +from flask import current_app +from pytest_mock import MockerFixture + +from superset.common.db_query_status import QueryStatus as QueryStatusEnum +from superset.models.core import Database + +# ============================================================================= +# Core Fixtures +# ============================================================================= + + +@pytest.fixture(autouse=True) +def mock_db_session(mocker: MockerFixture) -> MagicMock: + """Mock database session for all tests to avoid foreign key constraints.""" + mock_session = MagicMock() + mocker.patch("superset.sql.execution.executor.db.session", mock_session) + mocker.patch("superset.sql.execution.celery_task.db.session", mock_session) + return mock_session + + +@pytest.fixture +def mock_query() -> MagicMock: + """Create a mock Query model.""" + query = MagicMock() + query.id = 123 + query.database_id = 1 + query.sql = "SELECT * FROM users" + query.status = QueryStatusEnum.PENDING + query.error_message = None + query.progress = 0 + query.end_time = None + query.start_running_time = None + query.executed_sql = None + query.tmp_table_name = None + query.catalog = None + query.schema = "public" + query.extra = {} + query.set_extra_json_key = MagicMock() + query.results_key = None + query.select_as_cta = False + query.rows = 0 + query.to_dict = MagicMock(return_value={"id": 123}) + query.database = MagicMock() + query.database.db_engine_spec.extract_errors.return_value = [] + query.database.unique_name = "test_db" + query.database.cache_timeout = 300 + return query + + +@pytest.fixture +def mock_database() -> MagicMock: + """Create a mock Database.""" + database = MagicMock() + database.id = 1 + database.unique_name = "test_db" + database.cache_timeout = 300 + database.sqlalchemy_uri = "postgresql://localhost/test" + database.db_engine_spec = MagicMock() + database.db_engine_spec.engine = "postgresql" + database.db_engine_spec.run_multiple_statements_as_one = False + database.db_engine_spec.allows_sql_comments = True + database.db_engine_spec.extract_errors = MagicMock(return_value=[]) + database.db_engine_spec.execute_with_cursor = MagicMock() + database.db_engine_spec.get_cancel_query_id = MagicMock(return_value=None) + database.db_engine_spec.patch = MagicMock() + database.db_engine_spec.fetch_data = MagicMock(return_value=[]) + return database + + +@pytest.fixture +def mock_result_set() -> MagicMock: + """Create a mock SupersetResultSet.""" + result_set = MagicMock() + result_set.size = 2 + result_set.columns = [{"name": "id"}, {"name": "name"}] + result_set.pa_table = MagicMock() + result_set.to_pandas_df = MagicMock( + return_value=pd.DataFrame({"id": [1, 2], "name": ["Alice", "Bob"]}) + ) + return result_set + + +@pytest.fixture +def database() -> Database: + """Create a real test database instance.""" + return Database( + id=1, + database_name="test_db", + sqlalchemy_uri="sqlite://", + allow_dml=False, + ) + + +@pytest.fixture +def database_with_dml() -> Database: + """Create a real test database instance with DML allowed.""" + return Database( + id=2, + database_name="test_db_dml", + sqlalchemy_uri="sqlite://", + allow_dml=True, + ) + + +# ============================================================================= +# Helper Functions for Mock Creation +# ============================================================================= + + +def create_mock_cursor( + column_names: list[str], + data: list[tuple[Any, ...]] | None = None, +) -> MagicMock: + """ + Create a mock database cursor with column description. + + Args: + column_names: List of column names + data: Optional data to return from fetchall() + + Returns: + Configured MagicMock cursor + """ + mock_cursor = MagicMock() + mock_cursor.description = [ + (name, None, None, None, None, None, None) for name in column_names + ] + if data is not None: + mock_cursor.fetchall.return_value = data + return mock_cursor + + +def create_mock_connection(mock_cursor: MagicMock | None = None) -> MagicMock: + """ + Create a mock database connection. + + Args: + mock_cursor: Optional cursor to return from cursor() + + Returns: + Configured MagicMock connection with context manager support + """ + if mock_cursor is None: + mock_cursor = create_mock_cursor([]) + + mock_conn = MagicMock() + mock_conn.cursor.return_value = mock_cursor + mock_conn.close = MagicMock() + mock_conn.__enter__ = MagicMock(return_value=mock_conn) + mock_conn.__exit__ = MagicMock(return_value=False) + return mock_conn + + +def setup_mock_raw_connection( + mock_database: MagicMock, + mock_connection: MagicMock | None = None, +) -> MagicMock: + """ + Setup get_raw_connection as a context manager on a mock database. + + Args: + mock_database: The database mock to configure + mock_connection: Optional connection to yield + + Returns: + The configured mock connection + """ + if mock_connection is None: + mock_connection = create_mock_connection() + + @contextmanager + def _raw_connection( + catalog: str | None = None, + schema: str | None = None, + nullpool: bool = True, + source: Any | None = None, + ): + yield mock_connection + + mock_database.get_raw_connection = _raw_connection + return mock_connection + + +def setup_db_session_query_mock( + mock_db_session: MagicMock, + return_value: Any = None, +) -> None: + """ + Setup database session query chain for query lookup. + + Args: + mock_db_session: The database session mock + return_value: Value to return from one_or_none() + """ + filter_mock = mock_db_session.query.return_value.filter_by.return_value + filter_mock.one_or_none.return_value = return_value + + +def mock_query_execution( + mocker: MockerFixture, + database: Database, + return_data: list[tuple[Any, ...]], + column_names: list[str], +) -> MagicMock: + """ + Mock the raw connection execution path for testing. + + This helper sets up all necessary mocks for executing a query through + the database engine spec and returning results. + + Args: + mocker: pytest-mock fixture + database: Database instance to mock + return_data: Data to return from fetch_data, e.g. [(1, "Alice"), (2, "Bob")] + column_names: Column names for the result, e.g. ["id", "name"] + + Returns: + The mock for get_raw_connection, so tests can make assertions on it + """ + from superset.result_set import SupersetResultSet + + # Mock cursor and connection + mock_cursor = create_mock_cursor(column_names, return_data) + mock_conn = create_mock_connection(mock_cursor) + + get_raw_conn_mock = mocker.patch.object( + database, "get_raw_connection", return_value=mock_conn + ) + mocker.patch.object( + database, "mutate_sql_based_on_config", side_effect=lambda sql, **kw: sql + ) + mocker.patch.object(database.db_engine_spec, "execute") + mocker.patch.object(database.db_engine_spec, "fetch_data", return_value=return_data) + + # Create a real SupersetResultSet that converts to DataFrame properly + mock_result_set = MagicMock(spec=SupersetResultSet) + mock_result_set.to_pandas_df.return_value = pd.DataFrame( + return_data, columns=column_names + ) + mocker.patch("superset.result_set.SupersetResultSet", return_value=mock_result_set) + + return get_raw_conn_mock + + +# ============================================================================= +# Composite Fixtures for Common Test Patterns +# ============================================================================= + + +@pytest.fixture +def default_sql_config(mocker: MockerFixture) -> None: + """Patch app config with default SQL execution settings.""" + mocker.patch.dict( + current_app.config, + { + "SQL_QUERY_MUTATOR": None, + "SQLLAB_TIMEOUT": 30, + "SQL_MAX_ROW": None, + "QUERY_LOGGER": None, + }, + ) + + +@pytest.fixture +def mock_celery_task(mocker: MockerFixture) -> MagicMock: + """Mock the Celery task for SQL execution.""" + return mocker.patch("superset.sql.execution.celery_task.execute_sql_task") + + +def setup_cache_mocks( + mocker: MockerFixture, + get_result: Any = None, + store_side_effect: Any = None, +) -> tuple[MagicMock, MagicMock]: + """ + Setup cache get/store mocks for executor tests. + + Args: + mocker: pytest-mock fixture + get_result: Value to return from _get_from_cache + store_side_effect: Optional side effect for _store_in_cache + + Returns: + Tuple of (mock_get, mock_store) + """ + from superset.sql.execution.executor import SQLExecutor + + mock_get = mocker.patch.object( + SQLExecutor, "_get_from_cache", return_value=get_result + ) + mock_store = mocker.patch.object( + SQLExecutor, "_store_in_cache", side_effect=store_side_effect + ) + return mock_get, mock_store diff --git a/tests/unit_tests/sql/execution/test_celery_task.py b/tests/unit_tests/sql/execution/test_celery_task.py new file mode 100644 index 00000000000..9b2737cc0ea --- /dev/null +++ b/tests/unit_tests/sql/execution/test_celery_task.py @@ -0,0 +1,1077 @@ +# 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. + +"""Tests for celery_task.py - async SQL execution via Celery.""" + +from typing import Any +from unittest.mock import MagicMock + +import msgpack +import pytest +from celery.exceptions import SoftTimeLimitExceeded +from flask import current_app +from pytest_mock import MockerFixture + +from superset.common.db_query_status import QueryStatus as QueryStatusEnum +from superset.errors import ErrorLevel, SupersetError, SupersetErrorType +from superset.exceptions import SupersetErrorException, SupersetErrorsException + +# Note: mock_query, mock_database, mock_result_set, and mock_db_session +# fixtures are imported from conftest.py + + +# ============================================================================= +# Query Retrieval Tests +# ============================================================================= + + +def test_get_query_success( + mocker: MockerFixture, app_context: None, mock_query: MagicMock +) -> None: + """Test successful query retrieval.""" + from superset.sql.execution.celery_task import _get_query + + mock_session = mocker.patch("superset.sql.execution.celery_task.db.session") + mock_session.query.return_value.filter_by.return_value.one.return_value = mock_query + + result = _get_query(123) + + assert result == mock_query + + +# ============================================================================= +# Error Handling Tests +# ============================================================================= + + +def test_handle_query_error_basic( + mocker: MockerFixture, app_context: None, mock_query: MagicMock +) -> None: + """Test basic error handling.""" + from superset.sql.execution.celery_task import _handle_query_error + + mocker.patch("superset.sql.execution.celery_task.db.session") + + ex = Exception("Something went wrong") + payload = _handle_query_error(ex, mock_query) + + assert payload["status"] == QueryStatusEnum.FAILED + assert "Something went wrong" in payload["error"] + + +def test_handle_query_error_with_end_time_set( + mocker: MockerFixture, app_context: None, mock_query: MagicMock +) -> None: + """Test error handling when end_time is already set (line 116->120).""" + from superset.sql.execution.celery_task import _handle_query_error + + mocker.patch("superset.sql.execution.celery_task.db.session") + + # Set end_time to trigger the branch skip + mock_query.end_time = 12345.0 + + ex = Exception("Error with end_time set") + payload = _handle_query_error(ex, mock_query) + + assert payload["status"] == QueryStatusEnum.FAILED + # end_time should not be modified + assert mock_query.end_time == 12345.0 + + +def test_handle_query_error_sets_end_time( + mocker: MockerFixture, app_context: None, mock_query: MagicMock +) -> None: + """Test error handling sets end_time when not set.""" + from superset.sql.execution.celery_task import _handle_query_error + + mocker.patch("superset.sql.execution.celery_task.db.session") + mocker.patch( + "superset.sql.execution.celery_task.now_as_float", return_value=99999.0 + ) + + # end_time is None + mock_query.end_time = None + + ex = Exception("Error") + _handle_query_error(ex, mock_query) + + # Should set end_time + assert mock_query.end_time == 99999.0 + + +def test_handle_query_error_superset_error_exception( + mocker: MockerFixture, app_context: None, mock_query: MagicMock +) -> None: + """Test error handling with SupersetErrorException.""" + from superset.sql.execution.celery_task import _handle_query_error + + mocker.patch("superset.sql.execution.celery_task.db.session") + + error = SupersetError( + message="Test error", + error_type=SupersetErrorType.GENERIC_DB_ENGINE_ERROR, + level=ErrorLevel.ERROR, + ) + ex = SupersetErrorException(error) + + payload = _handle_query_error(ex, mock_query) + + assert len(payload["errors"]) == 1 + assert payload["errors"][0]["message"] == "Test error" + + +def test_handle_query_error_superset_errors_exception( + mocker: MockerFixture, app_context: None, mock_query: MagicMock +) -> None: + """Test error handling with SupersetErrorsException.""" + from superset.sql.execution.celery_task import _handle_query_error + + mocker.patch("superset.sql.execution.celery_task.db.session") + + errors = [ + SupersetError( + message="Error 1", + error_type=SupersetErrorType.GENERIC_DB_ENGINE_ERROR, + level=ErrorLevel.ERROR, + ), + SupersetError( + message="Error 2", + error_type=SupersetErrorType.GENERIC_DB_ENGINE_ERROR, + level=ErrorLevel.ERROR, + ), + ] + ex = SupersetErrorsException(errors) + + payload = _handle_query_error(ex, mock_query) + + assert len(payload["errors"]) == 2 + + +def test_handle_query_error_with_troubleshooting_link( + mocker: MockerFixture, app_context: None, mock_query: MagicMock +) -> None: + """Test error handling includes troubleshooting link.""" + from superset.sql.execution.celery_task import _handle_query_error + + mocker.patch("superset.sql.execution.celery_task.db.session") + mocker.patch.dict( + current_app.config, {"TROUBLESHOOTING_LINK": "https://help.example.com"} + ) + + ex = Exception("Error") + payload = _handle_query_error(ex, mock_query) + + assert payload["link"] == "https://help.example.com" + + +# ============================================================================= +# Serialization Tests +# ============================================================================= + + +def test_serialize_payload_json(mocker: MockerFixture, app_context: None) -> None: + """Test JSON serialization when msgpack config is False.""" + from superset.sql.execution.celery_task import _serialize_payload + + mocker.patch("superset.results_backend_use_msgpack", False) + payload = {"status": "success", "data": [1, 2, 3]} + + result = _serialize_payload(payload) + + # Now always returns bytes (encoded UTF-8) + assert isinstance(result, bytes) + assert b"success" in result + + +def test_serialize_payload_msgpack(mocker: MockerFixture, app_context: None) -> None: + """Test msgpack serialization when msgpack config is True.""" + from superset.sql.execution.celery_task import _serialize_payload + + mocker.patch("superset.results_backend_use_msgpack", True) + payload = {"status": "success", "data": [1, 2, 3]} + + result = _serialize_payload(payload) + + assert isinstance(result, bytes) + unpacked = msgpack.loads(result) + assert unpacked["status"] == "success" + + +# ============================================================================= +# Statement Preparation Tests +# ============================================================================= + + +def test_prepare_statement_blocks_single_statement( + app_context: None, mock_database: MagicMock +) -> None: + """Test statement block preparation for single statement.""" + from superset.sql.execution.celery_task import _prepare_statement_blocks + + sql = "SELECT * FROM users" + + script, blocks = _prepare_statement_blocks(sql, mock_database.db_engine_spec) + + assert len(blocks) == 1 + + +def test_prepare_statement_blocks_multiple_statements( + app_context: None, mock_database: MagicMock +) -> None: + """Test statement block preparation for multiple statements.""" + from superset.sql.execution.celery_task import _prepare_statement_blocks + + sql = "SELECT * FROM users; SELECT * FROM orders;" + + script, blocks = _prepare_statement_blocks(sql, mock_database.db_engine_spec) + + assert len(blocks) == 2 + + +def test_prepare_statement_blocks_run_as_one( + app_context: None, mock_database: MagicMock +) -> None: + """Test statement block preparation when engine runs multiple as one.""" + from superset.sql.execution.celery_task import _prepare_statement_blocks + + mock_database.db_engine_spec.run_multiple_statements_as_one = True + sql = "SELECT * FROM users; SELECT * FROM orders;" + + script, blocks = _prepare_statement_blocks(sql, mock_database.db_engine_spec) + + assert len(blocks) == 1 + + +# ============================================================================= +# Result Finalization Tests +# ============================================================================= + + +def test_finalize_successful_query( + mocker: MockerFixture, + app_context: None, + mock_query: MagicMock, + mock_result_set: MagicMock, + mock_database: MagicMock, +) -> None: + """Test successful query finalization.""" + from superset.sql.execution.celery_task import _finalize_successful_query + from superset.sql.parse import SQLScript + + mocker.patch("superset.results_backend_use_msgpack", False) + mocker.patch("superset.dataframe.df_to_records", return_value=[{"id": 1}]) + payload: dict[str, Any] = {} + + # Create original script + original_script = SQLScript( + "SELECT * FROM users", mock_database.db_engine_spec.engine + ) + + # New signature: (query, original_script, execution_results, payload, + # total_execution_time_ms). execution_results is a list of tuples: + # (executed_sql, result_set, exec_time, rowcount) + execution_results = [ + ("SELECT * FROM users WHERE rls_filter", mock_result_set, 10.5, 2) + ] + _finalize_successful_query( + mock_query, + original_script, + execution_results, # type: ignore[arg-type] + payload, + 10.5, + ) + + assert mock_query.rows == 2 + assert mock_query.progress == 100 + assert payload["status"] == QueryStatusEnum.SUCCESS + assert "statements" in payload + # SQL is formatted by SQLScript, so we can't compare exact whitespace + assert "SELECT" in payload["statements"][0]["original_sql"] + assert "FROM users" in payload["statements"][0]["original_sql"] + assert ( + payload["statements"][0]["executed_sql"] + == "SELECT * FROM users WHERE rls_filter" + ) + + +def test_finalize_successful_query_with_msgpack( + mocker: MockerFixture, + app_context: None, + mock_query: MagicMock, + mock_result_set: MagicMock, + mock_database: MagicMock, +) -> None: + """Test successful query finalization with Arrow/msgpack.""" + from superset.sql.execution.celery_task import _finalize_successful_query + from superset.sql.parse import SQLScript + + mocker.patch("superset.results_backend_use_msgpack", True) + mock_buffer = MagicMock() + mock_buffer.to_pybytes.return_value = b"arrow_data" + mocker.patch( + "superset.sql.execution.celery_task.write_ipc_buffer", return_value=mock_buffer + ) + + # Mock stats_logger to cover the stats timing branch + mock_stats = MagicMock() + mocker.patch.dict(current_app.config, {"STATS_LOGGER": mock_stats}) + + payload: dict[str, Any] = {} + + # Create original script + original_script = SQLScript( + "SELECT * FROM users", mock_database.db_engine_spec.engine + ) + + execution_results = [ + ("SELECT * FROM users WHERE rls_filter", mock_result_set, 10.5, 2) + ] + _finalize_successful_query( + mock_query, + original_script, + execution_results, # type: ignore[arg-type] + payload, + 10.5, + ) + + assert payload["statements"][0]["data"] == b"arrow_data" + + +def test_finalize_successful_query_msgpack_no_stats( + mocker: MockerFixture, + app_context: None, + mock_query: MagicMock, + mock_result_set: MagicMock, + mock_database: MagicMock, +) -> None: + """Test finalization with msgpack when has_app_context() is False.""" + from superset.sql.execution.celery_task import _finalize_successful_query + from superset.sql.parse import SQLScript + + mocker.patch("superset.results_backend_use_msgpack", True) + mocker.patch( + "superset.sql.execution.celery_task.has_app_context", return_value=False + ) + mock_buffer = MagicMock() + mock_buffer.to_pybytes.return_value = b"arrow_data" + mocker.patch( + "superset.sql.execution.celery_task.write_ipc_buffer", return_value=mock_buffer + ) + + payload: dict[str, Any] = {} + + # Create original script + original_script = SQLScript( + "SELECT * FROM users", mock_database.db_engine_spec.engine + ) + + execution_results = [ + ("SELECT * FROM users WHERE rls_filter", mock_result_set, 10.5, 2) + ] + _finalize_successful_query( + mock_query, + original_script, + execution_results, # type: ignore[arg-type] + payload, + 10.5, + ) + + assert payload["statements"][0]["data"] == b"arrow_data" + + +def test_finalize_successful_query_with_dml( + mocker: MockerFixture, + app_context: None, + mock_query: MagicMock, + mock_database: MagicMock, +) -> None: + """Test successful query finalization with DML statement (no result_set).""" + from superset.sql.execution.celery_task import _finalize_successful_query + from superset.sql.parse import SQLScript + + payload: dict[str, Any] = {} + + # Create original script + original_script = SQLScript( + "INSERT INTO users VALUES (1)", mock_database.db_engine_spec.engine + ) + + # DML statement: result_set is None, rowcount indicates affected rows + execution_results: list[tuple[str, None, float, int]] = [ + ("INSERT INTO users VALUES (1)", None, 5.0, 1) + ] + _finalize_successful_query( + mock_query, + original_script, + execution_results, # type: ignore[arg-type] + payload, + 5.0, + ) + + assert mock_query.rows == 0 # No rows returned for DML + assert mock_query.progress == 100 + assert payload["status"] == QueryStatusEnum.SUCCESS + assert payload["statements"][0]["data"] is None + assert payload["statements"][0]["row_count"] == 1 + assert payload["statements"][0]["columns"] == [] + + +# ============================================================================= +# Results Storage Tests +# ============================================================================= + + +def test_store_results_in_backend_success( + mocker: MockerFixture, + app_context: None, + mock_query: MagicMock, + mock_database: MagicMock, +) -> None: + """Test successful results storage.""" + from superset.sql.execution.celery_task import _store_results_in_backend + + mock_results_backend = MagicMock() + mock_results_backend.set.return_value = True + mocker.patch( + "superset.sql.execution.celery_task.results_backend", mock_results_backend + ) + mocker.patch("superset.results_backend_use_msgpack", False) + mocker.patch( + "superset.sql.execution.celery_task.zlib_compress", return_value=b"compressed" + ) + mocker.patch("superset.sql.execution.celery_task.db.session") + + payload = {"status": "success", "data": [], "query": {}} + _store_results_in_backend(mock_query, payload, mock_database) + + assert mock_query.results_key is not None + mock_results_backend.set.assert_called_once() + + +def test_store_results_in_backend_with_size_check( + mocker: MockerFixture, + app_context: None, + mock_query: MagicMock, + mock_database: MagicMock, +) -> None: + """Test results storage with payload size check (covers lines 232-247).""" + from superset.sql.execution.celery_task import _store_results_in_backend + + mock_results_backend = MagicMock() + mock_results_backend.set.return_value = True + mocker.patch( + "superset.sql.execution.celery_task.results_backend", mock_results_backend + ) + mocker.patch("superset.results_backend_use_msgpack", False) + mocker.patch( + "superset.sql.execution.celery_task.zlib_compress", return_value=b"compressed" + ) + mocker.patch("superset.sql.execution.celery_task.db.session") + + # Set a high payload max to pass the size check + mocker.patch.dict(current_app.config, {"SQLLAB_PAYLOAD_MAX_MB": 100}) + + payload = {"status": "success", "data": [], "query": {}} + _store_results_in_backend(mock_query, payload, mock_database) + + mock_results_backend.set.assert_called_once() + + +def test_store_results_in_backend_payload_too_large( + mocker: MockerFixture, + app_context: None, + mock_query: MagicMock, + mock_database: MagicMock, +) -> None: + """Test results storage with payload exceeding size limit.""" + from superset.sql.execution.celery_task import _store_results_in_backend + + mocker.patch("superset.results_backend_use_msgpack", False) + # Set very low limit + mocker.patch.dict(current_app.config, {"SQLLAB_PAYLOAD_MAX_MB": 0.000001}) + + large_payload = {"data": "x" * 1000, "query": {}} + + with pytest.raises(SupersetErrorException) as exc_info: + _store_results_in_backend(mock_query, large_payload, mock_database) + + assert exc_info.value.error.error_type == SupersetErrorType.RESULT_TOO_LARGE_ERROR + + +def test_store_results_in_backend_default_cache_timeout( + mocker: MockerFixture, + app_context: None, + mock_query: MagicMock, + mock_database: MagicMock, +) -> None: + """Test storage uses default cache timeout when database timeout is None.""" + from superset.sql.execution.celery_task import _store_results_in_backend + + mock_results_backend = MagicMock() + mock_results_backend.set.return_value = True + mocker.patch( + "superset.sql.execution.celery_task.results_backend", mock_results_backend + ) + mocker.patch("superset.results_backend_use_msgpack", False) + mocker.patch( + "superset.sql.execution.celery_task.zlib_compress", return_value=b"compressed" + ) + mocker.patch("superset.sql.execution.celery_task.db.session") + + # Set database cache_timeout to None + mock_database.cache_timeout = None + + payload = {"status": "success", "data": [], "query": {}} + _store_results_in_backend(mock_query, payload, mock_database) + + mock_results_backend.set.assert_called_once() + + +def test_store_results_in_backend_write_failure( + mocker: MockerFixture, + app_context: None, + mock_query: MagicMock, + mock_database: MagicMock, +) -> None: + """Test results storage write failure.""" + from superset.sql.execution.celery_task import _store_results_in_backend + + mock_results_backend = MagicMock() + mock_results_backend.set.return_value = False + mocker.patch( + "superset.sql.execution.celery_task.results_backend", mock_results_backend + ) + mocker.patch("superset.results_backend_use_msgpack", False) + mocker.patch( + "superset.sql.execution.celery_task.zlib_compress", return_value=b"compressed" + ) + mocker.patch("superset.sql.execution.celery_task.db.session") + + payload = {"status": "success", "data": [], "query": {}} + + with pytest.raises(SupersetErrorException) as exc_info: + _store_results_in_backend(mock_query, payload, mock_database) + + assert exc_info.value.error.error_type == SupersetErrorType.RESULTS_BACKEND_ERROR + + +# ============================================================================= +# Data Serialization Tests +# ============================================================================= + + +def test_serialize_result_set_msgpack( + mocker: MockerFixture, app_context: None, mock_result_set: MagicMock +) -> None: + """Test result set serialization with msgpack/Arrow when config is True.""" + from superset.sql.execution.celery_task import _serialize_result_set + + mocker.patch("superset.results_backend_use_msgpack", True) + mock_buffer = MagicMock() + mock_buffer.to_pybytes.return_value = b"arrow_data" + mocker.patch( + "superset.sql.execution.celery_task.write_ipc_buffer", return_value=mock_buffer + ) + mocker.patch.dict(current_app.config, {"STATS_LOGGER": MagicMock()}) + + data, columns = _serialize_result_set(mock_result_set) + + assert data == b"arrow_data" + assert columns == mock_result_set.columns + + +def test_serialize_result_set_json( + mocker: MockerFixture, app_context: None, mock_result_set: MagicMock +) -> None: + """Test result set serialization with JSON when msgpack config is False.""" + from superset.sql.execution.celery_task import _serialize_result_set + + mocker.patch("superset.results_backend_use_msgpack", False) + mocker.patch( + "superset.dataframe.df_to_records", + return_value=[{"id": 1, "name": "Alice"}], + ) + + data, columns = _serialize_result_set(mock_result_set) + + assert data == [{"id": 1, "name": "Alice"}] + assert columns == mock_result_set.columns + + +# ============================================================================= +# Helper Function Tests +# ============================================================================= + + +@pytest.mark.parametrize( + "query_status,expected_result", + [ + (QueryStatusEnum.STOPPED, True), + (QueryStatusEnum.RUNNING, False), + ], +) +def test_make_check_stopped_fn( + mocker: MockerFixture, + app_context: None, + mock_query: MagicMock, + query_status: QueryStatusEnum, + expected_result: bool, +) -> None: + """Test check_stopped function returns correct value based on query status.""" + from superset.sql.execution.celery_task import _make_check_stopped_fn + + mocker.patch("superset.sql.execution.celery_task.db.session") + mock_query.status = query_status + + check_stopped = _make_check_stopped_fn(mock_query) + result = check_stopped() + + assert result is expected_result + + +def test_make_execute_fn( + mocker: MockerFixture, + app_context: None, + mock_query: MagicMock, + mock_database: MagicMock, +) -> None: + """Test execute function creation.""" + from superset.sql.execution.celery_task import _make_execute_fn + + mock_cursor = MagicMock() + mocker.patch.dict(current_app.config, {"STATS_LOGGER": MagicMock()}) + + execute_fn = _make_execute_fn(mock_query, mock_database.db_engine_spec) + execute_fn(mock_cursor, "SELECT * FROM users") + + assert mock_query.executed_sql == "SELECT * FROM users" + + +@pytest.mark.parametrize( + "logger_configured,should_be_called", + [ + (True, True), + (False, False), + ], + ids=["with_logger", "no_logger"], +) +def test_make_log_query_fn( + mocker: MockerFixture, + app_context: None, + mock_database: MagicMock, + logger_configured: bool, + should_be_called: bool, +) -> None: + """Test log query function with and without logger configured.""" + from superset.sql.execution.celery_task import _make_log_query_fn + + mock_logger = MagicMock() if logger_configured else None + mocker.patch.dict(current_app.config, {"QUERY_LOGGER": mock_logger}) + mocker.patch("superset.sql.execution.celery_task.security_manager", MagicMock()) + + log_fn = _make_log_query_fn(mock_database) + log_fn("SELECT * FROM users", "public") + + if should_be_called: + assert mock_logger is not None + mock_logger.assert_called_once() + # If no logger, the function should complete without error + + +# ============================================================================= +# Main Task Execution Tests +# ============================================================================= + + +def test_execute_sql_task_success( + mocker: MockerFixture, + app_context: None, + mock_query: MagicMock, + mock_database: MagicMock, + mock_result_set: MagicMock, +) -> None: + """Test successful execute_sql_task (covers lines 339-352, 400-473).""" + from superset.sql.execution.celery_task import execute_sql_task + + from .conftest import setup_mock_raw_connection + + mock_query.database = mock_database + mock_query.status = QueryStatusEnum.PENDING + + mocker.patch( + "superset.sql.execution.celery_task._get_query", return_value=mock_query + ) + # execute_sql_with_cursor returns (exec_sql, result_set, time, rowcount) + mocker.patch( + "superset.sql.execution.celery_task.execute_sql_with_cursor", + return_value=[("SELECT * FROM users", mock_result_set, 10.5, 2)], + ) + mocker.patch("superset.sql.execution.celery_task.results_backend", None) + mocker.patch("superset.results_backend_use_msgpack", False) + mocker.patch("superset.sql.execution.celery_task.db.session") + mocker.patch("superset.dataframe.df_to_records", return_value=[]) + mocker.patch("superset.sql.execution.celery_task.security_manager") + mocker.patch.dict(current_app.config, {"STATS_LOGGER": MagicMock()}) + + setup_mock_raw_connection(mock_database) + + result = execute_sql_task(123, "SELECT * FROM users", username="admin") + + assert result is not None # Success returns payload + assert result["status"] == QueryStatusEnum.SUCCESS + assert mock_query.status == QueryStatusEnum.SUCCESS + + +def test_execute_sql_task_with_start_time( + mocker: MockerFixture, + app_context: None, + mock_query: MagicMock, + mock_database: MagicMock, + mock_result_set: MagicMock, +) -> None: + """Test execute_sql_task accepts start_time parameter (covers line 400-402).""" + import time + + from superset.sql.execution.celery_task import execute_sql_task + + mock_query.database = mock_database + + mocker.patch( + "superset.sql.execution.celery_task._get_query", return_value=mock_query + ) + # execute_sql_with_cursor returns (exec_sql, result_set, time, rowcount) + mocker.patch( + "superset.sql.execution.celery_task.execute_sql_with_cursor", + return_value=[("SELECT * FROM users", mock_result_set, 10.5, 2)], + ) + mocker.patch("superset.sql.execution.celery_task.results_backend", None) + mocker.patch("superset.results_backend_use_msgpack", False) + mocker.patch("superset.sql.execution.celery_task.db.session") + mocker.patch("superset.dataframe.df_to_records", return_value=[]) + mocker.patch("superset.sql.execution.celery_task.security_manager") + mocker.patch.dict(current_app.config, {"STATS_LOGGER": MagicMock()}) + + from .conftest import setup_mock_raw_connection + + setup_mock_raw_connection(mock_database) + + start_time = time.time() - 1.0 + result = execute_sql_task(123, "SELECT * FROM users", start_time=start_time) + + # Verify task completes successfully with start_time + assert result is not None + assert result["status"] == QueryStatusEnum.SUCCESS + + +def test_execute_sql_task_with_cancel_query_id( + mocker: MockerFixture, + app_context: None, + mock_query: MagicMock, + mock_database: MagicMock, + mock_result_set: MagicMock, +) -> None: + """Test execute_sql_task sets cancel_query_id when available.""" + from superset.sql.execution.celery_task import execute_sql_task + + mock_query.database = mock_database + mock_database.db_engine_spec.get_cancel_query_id.return_value = "cancel_123" + + mocker.patch( + "superset.sql.execution.celery_task._get_query", return_value=mock_query + ) + # execute_sql_with_cursor returns (exec_sql, result_set, time, rowcount) + mocker.patch( + "superset.sql.execution.celery_task.execute_sql_with_cursor", + return_value=[("SELECT * FROM users", mock_result_set, 10.5, 2)], + ) + mocker.patch("superset.sql.execution.celery_task.results_backend", None) + mocker.patch("superset.results_backend_use_msgpack", False) + mocker.patch("superset.sql.execution.celery_task.db.session") + mocker.patch("superset.dataframe.df_to_records", return_value=[]) + mocker.patch("superset.sql.execution.celery_task.security_manager") + mocker.patch.dict(current_app.config, {"STATS_LOGGER": MagicMock()}) + + from .conftest import setup_mock_raw_connection + + setup_mock_raw_connection(mock_database) + + execute_sql_task(123, "SELECT * FROM users") + + # Verify cancel_query_id was set via db_engine_spec + mock_database.db_engine_spec.get_cancel_query_id.assert_called_once() + mock_query.set_extra_json_key.assert_any_call("cancel_query", "cancel_123") + + +def test_execute_sql_task_stopped( + mocker: MockerFixture, + app_context: None, + mock_query: MagicMock, + mock_database: MagicMock, +) -> None: + """Test execute_sql_task when query is stopped (covers lines 456-458).""" + from superset.sql.execution.celery_task import execute_sql_task + + mock_query.database = mock_database + + mocker.patch( + "superset.sql.execution.celery_task._get_query", return_value=mock_query + ) + # Empty list indicates stopped (check_stopped_fn returned True mid-execution) + mocker.patch( + "superset.sql.execution.celery_task.execute_sql_with_cursor", + return_value=[], + ) + mocker.patch("superset.sql.execution.celery_task.db.session") + mocker.patch("superset.sql.execution.celery_task.security_manager") + mocker.patch.dict(current_app.config, {"STATS_LOGGER": MagicMock()}) + + from .conftest import setup_mock_raw_connection + + setup_mock_raw_connection(mock_database) + + result = execute_sql_task(123, "SELECT * FROM users") + + assert result["status"] == QueryStatusEnum.STOPPED + + +def test_execute_sql_task_with_mutation( + mocker: MockerFixture, + app_context: None, + mock_query: MagicMock, + mock_database: MagicMock, + mock_result_set: MagicMock, +) -> None: + """Test execute_sql_task commits for mutations (covers lines 461-462).""" + from superset.sql.execution.celery_task import execute_sql_task + + mock_query.database = mock_database + mock_query.select_as_cta = True # Trigger mutation commit + + mocker.patch( + "superset.sql.execution.celery_task._get_query", return_value=mock_query + ) + # execute_sql_with_cursor returns (exec_sql, result_set, time, rowcount) + mocker.patch( + "superset.sql.execution.celery_task.execute_sql_with_cursor", + return_value=[("INSERT INTO users VALUES (1)", mock_result_set, 5.0, 1)], + ) + mocker.patch("superset.sql.execution.celery_task.results_backend", None) + mocker.patch("superset.sql.execution.celery_task.db.session") + mocker.patch("superset.dataframe.df_to_records", return_value=[]) + mocker.patch("superset.sql.execution.celery_task.security_manager") + mocker.patch.dict(current_app.config, {"STATS_LOGGER": MagicMock()}) + + from .conftest import setup_mock_raw_connection + + mock_conn = setup_mock_raw_connection(mock_database) + + execute_sql_task(123, "INSERT INTO users VALUES (1)") + + mock_conn.commit.assert_called() + + +def test_execute_sql_task_with_results_backend( + mocker: MockerFixture, + app_context: None, + mock_query: MagicMock, + mock_database: MagicMock, + mock_result_set: MagicMock, +) -> None: + """Test execute_sql_task stores results in backend (covers lines 466-467).""" + from superset.sql.execution.celery_task import execute_sql_task + + mock_query.database = mock_database + + mocker.patch( + "superset.sql.execution.celery_task._get_query", return_value=mock_query + ) + # execute_sql_with_cursor returns (exec_sql, result_set, time, rowcount) + mocker.patch( + "superset.sql.execution.celery_task.execute_sql_with_cursor", + return_value=[("SELECT * FROM users", mock_result_set, 10.5, 2)], + ) + + mock_results_backend = MagicMock() + mock_results_backend.set.return_value = True + mocker.patch( + "superset.sql.execution.celery_task.results_backend", mock_results_backend + ) + mocker.patch("superset.results_backend_use_msgpack", False) + mocker.patch( + "superset.sql.execution.celery_task.zlib_compress", return_value=b"data" + ) + mocker.patch("superset.sql.execution.celery_task.db.session") + mocker.patch("superset.dataframe.df_to_records", return_value=[]) + mocker.patch("superset.sql.execution.celery_task.security_manager") + mocker.patch.dict(current_app.config, {"STATS_LOGGER": MagicMock()}) + + from .conftest import setup_mock_raw_connection + + setup_mock_raw_connection(mock_database) + + result = execute_sql_task(123, "SELECT * FROM users") + + mock_results_backend.set.assert_called_once() + assert result is not None + assert result["status"] == QueryStatusEnum.SUCCESS + + +def test_execute_sql_task_timeout( + mocker: MockerFixture, + app_context: None, + mock_query: MagicMock, + mock_database: MagicMock, +) -> None: + """Test execute_sql_task handles timeout (covers lines 438-453).""" + from superset.sql.execution.celery_task import execute_sql_task + + mock_query.database = mock_database + + mocker.patch( + "superset.sql.execution.celery_task._get_query", return_value=mock_query + ) + mocker.patch( + "superset.sql.execution.celery_task.execute_sql_with_cursor", + side_effect=SoftTimeLimitExceeded(), + ) + mocker.patch("superset.sql.execution.celery_task.db.session") + mocker.patch("superset.sql.execution.celery_task.security_manager") + mocker.patch.dict( + current_app.config, + {"STATS_LOGGER": MagicMock(), "SQLLAB_ASYNC_TIME_LIMIT_SEC": 300}, + ) + + from .conftest import setup_mock_raw_connection + + setup_mock_raw_connection(mock_database) + + result = execute_sql_task(123, "SELECT * FROM users") + + # TIMED_OUT status is preserved (not overwritten to FAILED) + assert result["status"] == QueryStatusEnum.TIMED_OUT.value + + +def test_execute_sql_task_unhandled_exception( + mocker: MockerFixture, + app_context: None, + mock_query: MagicMock, +) -> None: + """Test execute_sql_task handles unhandled exceptions (covers lines 347-352).""" + from superset.sql.execution.celery_task import execute_sql_task + + # Mock _get_query to succeed first time (for override_user), then return mock_query + mocker.patch( + "superset.sql.execution.celery_task._get_query", return_value=mock_query + ) + mocker.patch( + "superset.sql.execution.celery_task._execute_sql_statements", + side_effect=Exception("Unexpected error"), + ) + mocker.patch("superset.sql.execution.celery_task.db.session") + mocker.patch("superset.sql.execution.celery_task.security_manager") + mocker.patch.dict(current_app.config, {"STATS_LOGGER": MagicMock()}) + + result = execute_sql_task(123, "SELECT * FROM users") + + assert result["status"] == QueryStatusEnum.FAILED + + +def test_execute_sql_task_success_final_commit( + mocker: MockerFixture, + app_context: None, + mock_query: MagicMock, + mock_database: MagicMock, + mock_result_set: MagicMock, +) -> None: + """Test execute_sql_task final success path (covers lines 469-473).""" + from superset.sql.execution.celery_task import execute_sql_task + + mock_query.database = mock_database + mock_query.status = QueryStatusEnum.RUNNING # Will be changed to SUCCESS + + mocker.patch( + "superset.sql.execution.celery_task._get_query", return_value=mock_query + ) + # execute_sql_with_cursor returns (exec_sql, result_set, time, rowcount) + mocker.patch( + "superset.sql.execution.celery_task.execute_sql_with_cursor", + return_value=[("SELECT * FROM users", mock_result_set, 10.5, 2)], + ) + mocker.patch("superset.sql.execution.celery_task.results_backend", None) + mocker.patch("superset.results_backend_use_msgpack", False) + mock_session = mocker.patch("superset.sql.execution.celery_task.db.session") + mocker.patch("superset.dataframe.df_to_records", return_value=[]) + mocker.patch("superset.sql.execution.celery_task.security_manager") + mocker.patch.dict(current_app.config, {"STATS_LOGGER": MagicMock()}) + + from .conftest import setup_mock_raw_connection + + setup_mock_raw_connection(mock_database) + + result = execute_sql_task(123, "SELECT * FROM users") + + assert result is not None + assert result["status"] == QueryStatusEnum.SUCCESS + assert mock_query.status == QueryStatusEnum.SUCCESS + mock_session.commit.assert_called() + + +def test_execute_sql_task_with_failed_status_before_final_commit( + mocker: MockerFixture, + app_context: None, + mock_query: MagicMock, + mock_database: MagicMock, + mock_result_set: MagicMock, +) -> None: + """Test execute_sql_task final commit when query.status is already FAILED.""" + from superset.sql.execution.celery_task import execute_sql_task + + mock_query.database = mock_database + + mocker.patch( + "superset.sql.execution.celery_task._get_query", return_value=mock_query + ) + # execute_sql_with_cursor returns (exec_sql, result_set, time, rowcount) + mocker.patch( + "superset.sql.execution.celery_task.execute_sql_with_cursor", + return_value=[("SELECT * FROM users", mock_result_set, 10.5, 2)], + ) + + # Mock _store_results_in_backend to set status to FAILED without raising + def mock_store_results(query, payload, database): + query.status = QueryStatusEnum.FAILED + + mocker.patch("superset.sql.execution.celery_task.results_backend", MagicMock()) + mocker.patch("superset.results_backend_use_msgpack", False) + mocker.patch( + "superset.sql.execution.celery_task._store_results_in_backend", + side_effect=mock_store_results, + ) + mocker.patch("superset.sql.execution.celery_task.db.session") + mocker.patch("superset.dataframe.df_to_records", return_value=[]) + mocker.patch("superset.sql.execution.celery_task.security_manager") + mocker.patch.dict(current_app.config, {"STATS_LOGGER": MagicMock()}) + + from .conftest import setup_mock_raw_connection + + setup_mock_raw_connection(mock_database) + + result = execute_sql_task(123, "SELECT * FROM users") + + # Verify query status remains FAILED and is not changed to SUCCESS + assert result is not None + assert mock_query.status == QueryStatusEnum.FAILED diff --git a/tests/unit_tests/sql/execution/test_executor.py b/tests/unit_tests/sql/execution/test_executor.py new file mode 100644 index 00000000000..5168959e9ea --- /dev/null +++ b/tests/unit_tests/sql/execution/test_executor.py @@ -0,0 +1,2114 @@ +# 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. + +""" +Tests for SQLExecutor. + +These tests cover the SQL execution API including: +- Basic execution +- DML handling +- Jinja2 template rendering +- CTAS/CVAS support +- Security features (RLS, disallowed functions) +- Result caching +- Query model persistence +- Async execution +""" + +from typing import Any +from unittest.mock import MagicMock + +import msgpack +import pandas as pd +import pytest +from flask import current_app +from pytest_mock import MockerFixture +from superset_core.api.types import ( + CacheOptions, + QueryOptions, + QueryStatus, +) + +from superset.models.core import Database + +# Note: database, database_with_dml, mock_db_session fixtures and +# mock_query_execution helper are imported from conftest.py +from .conftest import mock_query_execution + +# ============================================================================= +# Basic Execution Tests +# ============================================================================= + + +def test_execute_select_success( + mocker: MockerFixture, database: Database, app_context: None +) -> None: + """Test successful SELECT query execution.""" + mock_query_execution( + mocker, + database, + return_data=[(1, "Alice"), (2, "Bob")], + column_names=["id", "name"], + ) + mocker.patch.dict( + current_app.config, + { + "SQL_QUERY_MUTATOR": None, + "SQLLAB_TIMEOUT": 30, + "SQL_MAX_ROW": None, + "QUERY_LOGGER": None, + }, + ) + + result = database.execute("SELECT id, name FROM users") + + assert result.status == QueryStatus.SUCCESS + assert len(result.statements) == 1 + assert result.statements[0].data is not None + assert result.statements[0].row_count == 2 + assert result.error_message is None + + +def test_execute_with_options( + mocker: MockerFixture, database: Database, app_context: None +) -> None: + """Test query execution with custom options.""" + get_raw_conn_mock = mock_query_execution( + mocker, database, return_data=[(100,)], column_names=["count"] + ) + mocker.patch.dict( + current_app.config, + { + "SQL_QUERY_MUTATOR": None, + "SQLLAB_TIMEOUT": 30, + "SQL_MAX_ROW": None, + "QUERY_LOGGER": None, + }, + ) + + options = QueryOptions(catalog="main", schema="public", limit=50) + result = database.execute("SELECT COUNT(*) FROM users", options=options) + + assert result.status == QueryStatus.SUCCESS + get_raw_conn_mock.assert_called_once() + call_kwargs = get_raw_conn_mock.call_args[1] + assert call_kwargs["catalog"] == "main" + assert call_kwargs["schema"] == "public" + + +def test_execute_records_execution_time( + mocker: MockerFixture, database: Database, app_context: None +) -> None: + """Test that execution time is recorded.""" + mock_query_execution(mocker, database, return_data=[(1,)], column_names=["id"]) + mocker.patch.dict( + current_app.config, + { + "SQL_QUERY_MUTATOR": None, + "SQLLAB_TIMEOUT": 30, + "SQL_MAX_ROW": None, + "QUERY_LOGGER": None, + }, + ) + + result = database.execute("SELECT id FROM users") + + assert result.status == QueryStatus.SUCCESS + assert result.total_execution_time_ms is not None + assert result.total_execution_time_ms >= 0 + + +def test_execute_creates_query_record( + mocker: MockerFixture, + database: Database, + app_context: None, + mock_db_session: MagicMock, +) -> None: + """Test that execute creates a Query record for audit.""" + from superset.sql.execution.executor import SQLExecutor + + mock_query_execution(mocker, database, return_data=[(1,)], column_names=["id"]) + mocker.patch.dict( + current_app.config, + { + "SQL_QUERY_MUTATOR": None, + "SQLLAB_TIMEOUT": 30, + "SQL_MAX_ROW": None, + "QUERY_LOGGER": None, + }, + ) + + # Mock _create_query_record to return a mock query with ID + mock_query = MagicMock() + mock_query.id = 123 + mock_create_query = mocker.patch.object( + SQLExecutor, "_create_query_record", return_value=mock_query + ) + + result = database.execute("SELECT id FROM users") + + assert result.status == QueryStatus.SUCCESS + assert result.query_id == 123 + mock_create_query.assert_called_once() + + +# ============================================================================= +# DML Handling Tests +# ============================================================================= + + +def test_execute_dml_without_permission( + mocker: MockerFixture, database: Database, app_context: None +) -> None: + """Test that DML queries fail when database.allow_dml is False.""" + mocker.patch.dict( + current_app.config, + {"SQL_QUERY_MUTATOR": None, "SQLLAB_TIMEOUT": 30}, + ) + + result = database.execute("INSERT INTO users (name) VALUES ('test')") + + assert result.status == QueryStatus.FAILED + assert result.error_message is not None + assert "DML queries are not allowed" in result.error_message + + +def test_execute_dml_with_permission( + mocker: MockerFixture, database_with_dml: Database, app_context: None +) -> None: + """Test that DML queries succeed when database.allow_dml is True.""" + mock_query_execution(mocker, database_with_dml, return_data=[], column_names=[]) + mocker.patch.dict( + current_app.config, + { + "SQL_QUERY_MUTATOR": None, + "SQLLAB_TIMEOUT": 30, + "SQL_MAX_ROW": None, + "QUERY_LOGGER": None, + }, + ) + + result = database_with_dml.execute("INSERT INTO users (name) VALUES ('test')") + + assert result.status == QueryStatus.SUCCESS + + +def test_execute_update_without_permission( + mocker: MockerFixture, database: Database, app_context: None +) -> None: + """Test that UPDATE queries fail when database.allow_dml is False.""" + mocker.patch.dict( + current_app.config, + {"SQL_QUERY_MUTATOR": None, "SQLLAB_TIMEOUT": 30}, + ) + + result = database.execute("UPDATE users SET name = 'test' WHERE id = 1") + + assert result.status == QueryStatus.FAILED + assert result.error_message is not None + assert "DML queries are not allowed" in result.error_message + + +def test_execute_delete_without_permission( + mocker: MockerFixture, database: Database, app_context: None +) -> None: + """Test that DELETE queries fail when database.allow_dml is False.""" + mocker.patch.dict( + current_app.config, + {"SQL_QUERY_MUTATOR": None, "SQLLAB_TIMEOUT": 30}, + ) + + result = database.execute("DELETE FROM users WHERE id = 1") + + assert result.status == QueryStatus.FAILED + assert result.error_message is not None + assert "DML queries are not allowed" in result.error_message + + +# ============================================================================= +# Jinja2 Template Rendering Tests +# ============================================================================= + + +def test_execute_with_template_params( + mocker: MockerFixture, database: Database, app_context: None +) -> None: + """Test query execution with Jinja2 template parameters.""" + mock_query_execution(mocker, database, return_data=[(1,)], column_names=["id"]) + mocker.patch.dict( + current_app.config, + { + "SQL_QUERY_MUTATOR": None, + "SQLLAB_TIMEOUT": 30, + "SQL_MAX_ROW": None, + "QUERY_LOGGER": None, + }, + ) + + # Mock the template processor + mock_tp = MagicMock() + mock_tp.process_template.return_value = ( + "SELECT * FROM events WHERE date > '2024-01-01'" + ) + mocker.patch( + "superset.jinja_context.get_template_processor", + return_value=mock_tp, + ) + + options = QueryOptions( + template_params={"table": "events", "start_date": "2024-01-01"} + ) + result = database.execute( + "SELECT * FROM {{ table }} WHERE date > '{{ start_date }}'", + options=options, + ) + + assert result.status == QueryStatus.SUCCESS + mock_tp.process_template.assert_called_once() + + +def test_execute_without_template_params_no_rendering( + mocker: MockerFixture, database: Database, app_context: None +) -> None: + """Test that template rendering is skipped when no params provided.""" + mock_query_execution(mocker, database, return_data=[(1,)], column_names=["id"]) + mocker.patch.dict( + current_app.config, + { + "SQL_QUERY_MUTATOR": None, + "SQLLAB_TIMEOUT": 30, + "SQL_MAX_ROW": None, + "QUERY_LOGGER": None, + }, + ) + + mock_get_tp = mocker.patch("superset.jinja_context.get_template_processor") + + result = database.execute("SELECT * FROM users") + + assert result.status == QueryStatus.SUCCESS + mock_get_tp.assert_not_called() + + +# ============================================================================= +# Disallowed Functions Tests +# ============================================================================= + + +def test_execute_disallowed_functions( + mocker: MockerFixture, database: Database, app_context: None +) -> None: + """Test that disallowed SQL functions are blocked.""" + mocker.patch.dict( + current_app.config, + { + "SQL_QUERY_MUTATOR": None, + "SQLLAB_TIMEOUT": 30, + "DISALLOWED_SQL_FUNCTIONS": {"sqlite": {"LOAD_EXTENSION"}}, + }, + ) + + result = database.execute("SELECT LOAD_EXTENSION('malicious.so')") + + assert result.status == QueryStatus.FAILED + assert result.error_message is not None + assert "Disallowed SQL functions" in result.error_message + + +def test_execute_allowed_functions( + mocker: MockerFixture, database: Database, app_context: None +) -> None: + """Test that allowed SQL functions work normally.""" + mock_query_execution(mocker, database, return_data=[(5,)], column_names=["count"]) + mocker.patch.dict( + current_app.config, + { + "SQL_QUERY_MUTATOR": None, + "SQLLAB_TIMEOUT": 30, + "SQL_MAX_ROW": None, + "DISALLOWED_SQL_FUNCTIONS": {"sqlite": {"LOAD_EXTENSION"}}, + "QUERY_LOGGER": None, + }, + ) + + result = database.execute("SELECT COUNT(*) FROM users") + + assert result.status == QueryStatus.SUCCESS + + +# ============================================================================= +# Row-Level Security Tests +# ============================================================================= + + +def test_execute_rls_applied( + mocker: MockerFixture, database: Database, app_context: None +) -> None: + """Test that RLS is always applied.""" + from superset.sql.execution.executor import SQLExecutor + + mock_query_execution(mocker, database, return_data=[(1,)], column_names=["id"]) + mocker.patch.dict( + current_app.config, + { + "SQL_QUERY_MUTATOR": None, + "SQLLAB_TIMEOUT": 30, + "SQL_MAX_ROW": None, + "QUERY_LOGGER": None, + }, + ) + + # Mock _apply_rls_to_script to verify it's always called + mock_apply_rls = mocker.patch.object(SQLExecutor, "_apply_rls_to_script") + + result = database.execute("SELECT * FROM users") + + assert result.status == QueryStatus.SUCCESS + mock_apply_rls.assert_called() + + +# ============================================================================= +# Result Caching Tests +# ============================================================================= + + +def test_execute_returns_cached_result( + mocker: MockerFixture, database: Database, app_context: None +) -> None: + """Test that cached results are returned when available.""" + from superset.sql.execution.executor import SQLExecutor + + cached_df = pd.DataFrame({"id": [1, 2]}) + + mocker.patch.dict( + current_app.config, + { + "SQL_QUERY_MUTATOR": None, + "SQLLAB_TIMEOUT": 30, + "SQL_MAX_ROW": None, + "QUERY_LOGGER": None, + }, + ) + + # Mock _get_from_cache to simulate a cache hit + cached_result = MagicMock() + cached_result.status = QueryStatus.SUCCESS + cached_result.data = cached_df + cached_result.is_cached = True + mocker.patch.object(SQLExecutor, "_get_from_cache", return_value=cached_result) + + # get_raw_connection should NOT be called if cache hit + get_conn_mock = mocker.patch.object(database, "get_raw_connection") + + result = database.execute("SELECT * FROM users") + + assert result.status == QueryStatus.SUCCESS + assert result.is_cached is True + get_conn_mock.assert_not_called() + + +def test_execute_force_cache_refresh( + mocker: MockerFixture, database: Database, app_context: None +) -> None: + """Test that force_cache_refresh bypasses the cache.""" + from superset.sql.execution.executor import SQLExecutor + + mock_query_execution(mocker, database, return_data=[(1,)], column_names=["id"]) + mocker.patch.dict( + current_app.config, + { + "SQL_QUERY_MUTATOR": None, + "SQLLAB_TIMEOUT": 30, + "SQL_MAX_ROW": None, + "QUERY_LOGGER": None, + }, + ) + + # Mock _get_from_cache - should NOT be called when force_refresh=True + mock_get_cache = mocker.patch.object(SQLExecutor, "_get_from_cache") + + options = QueryOptions(cache=CacheOptions(force_refresh=True)) + result = database.execute("SELECT * FROM users", options=options) + + assert result.status == QueryStatus.SUCCESS + assert result.is_cached is False + assert sum(s.row_count for s in result.statements) == 1 # Fresh result + mock_get_cache.assert_not_called() + + +def test_execute_stores_in_cache( + mocker: MockerFixture, database: Database, app_context: None +) -> None: + """Test that results are stored in cache.""" + from superset.sql.execution.executor import SQLExecutor + + mock_query_execution(mocker, database, return_data=[(1,)], column_names=["id"]) + mocker.patch.dict( + current_app.config, + { + "SQL_QUERY_MUTATOR": None, + "SQLLAB_TIMEOUT": 30, + "SQL_MAX_ROW": None, + "CACHE_DEFAULT_TIMEOUT": 300, + "QUERY_LOGGER": None, + }, + ) + + # Mock _get_from_cache to return None (cache miss) + mocker.patch.object(SQLExecutor, "_get_from_cache", return_value=None) + # Mock _store_in_cache to verify it gets called + mock_store_cache = mocker.patch.object(SQLExecutor, "_store_in_cache") + + result = database.execute("SELECT * FROM users") + + assert result.status == QueryStatus.SUCCESS + mock_store_cache.assert_called_once() + + +# ============================================================================= +# Timeout Tests +# ============================================================================= + + +def test_execute_timeout( + mocker: MockerFixture, database: Database, app_context: None +) -> None: + """Test query timeout handling.""" + from superset.errors import ErrorLevel, SupersetErrorType + from superset.exceptions import SupersetTimeoutException + + # Mock get_raw_connection to raise timeout + mock_conn = MagicMock() + mock_conn.__enter__ = MagicMock( + side_effect=SupersetTimeoutException( + error_type=SupersetErrorType.GENERIC_BACKEND_ERROR, + message="Query timed out", + level=ErrorLevel.ERROR, + ) + ) + mock_conn.__exit__ = MagicMock(return_value=False) + mocker.patch.object(database, "get_raw_connection", return_value=mock_conn) + mocker.patch.dict( + current_app.config, + { + "SQL_QUERY_MUTATOR": None, + "SQLLAB_TIMEOUT": 1, + "SQL_MAX_ROW": None, + "QUERY_LOGGER": None, + }, + ) + + result = database.execute("SELECT * FROM large_table") + + assert result.status == QueryStatus.TIMED_OUT + assert result.error_message is not None + + +def test_execute_custom_timeout( + mocker: MockerFixture, database: Database, app_context: None +) -> None: + """Test query with custom timeout option.""" + mock_query_execution(mocker, database, return_data=[(1,)], column_names=["id"]) + mocker.patch.dict( + current_app.config, + { + "SQL_QUERY_MUTATOR": None, + "SQLLAB_TIMEOUT": 30, + "SQL_MAX_ROW": None, + "QUERY_LOGGER": None, + }, + ) + + mock_timeout = mocker.patch("superset.sql.execution.executor.utils.timeout") + mock_timeout.return_value.__enter__ = MagicMock() + mock_timeout.return_value.__exit__ = MagicMock(return_value=False) + + options = QueryOptions(timeout_seconds=60) + result = database.execute("SELECT * FROM users", options=options) + + assert result.status == QueryStatus.SUCCESS + mock_timeout.assert_called_with( + seconds=60, + error_message="Query exceeded the 60 seconds timeout.", + ) + + +# ============================================================================= +# Error Handling Tests +# ============================================================================= + + +def test_execute_error( + mocker: MockerFixture, database: Database, app_context: None +) -> None: + """Test general error handling.""" + # Mock get_raw_connection to raise an error + mock_conn = MagicMock() + mock_cursor = MagicMock() + mock_conn.cursor.return_value = mock_cursor + mock_conn.__enter__ = MagicMock(return_value=mock_conn) + mock_conn.__exit__ = MagicMock(return_value=False) + mocker.patch.object(database, "get_raw_connection", return_value=mock_conn) + mocker.patch.object( + database, "mutate_sql_based_on_config", side_effect=lambda sql, **kw: sql + ) + mocker.patch.object( + database.db_engine_spec, "execute", side_effect=Exception("Database error") + ) + mocker.patch.dict( + current_app.config, + { + "SQL_QUERY_MUTATOR": None, + "SQLLAB_TIMEOUT": 30, + "SQL_MAX_ROW": None, + "QUERY_LOGGER": None, + }, + ) + + result = database.execute("SELECT * FROM nonexistent") + + assert result.status == QueryStatus.FAILED + assert result.error_message is not None + assert "Database error" in result.error_message + + +# ============================================================================= +# Async Execution Tests +# ============================================================================= + + +def test_execute_async_creates_query( + mocker: MockerFixture, + database: Database, + app_context: None, + mock_db_session: MagicMock, +) -> None: + """Test that execute_async creates a Query record and submits to Celery.""" + mocker.patch.dict( + current_app.config, {"SQL_QUERY_MUTATOR": None, "SQLLAB_TIMEOUT": 30} + ) + + # Mock db.session.add to set query.id (simulating database auto-increment) + def set_query_id(query: Any) -> None: + if not hasattr(query, "id") or query.id is None: + query.id = 123 + + mock_db_session.add.side_effect = set_query_id + + mock_celery_task = mocker.patch( + "superset.sql.execution.celery_task.execute_sql_task" + ) + + result = database.execute_async("SELECT * FROM users") + + assert result.status == QueryStatus.PENDING + assert result.query_id is not None + assert result.query_id == 123 + mock_db_session.add.assert_called() + mock_celery_task.delay.assert_called() + + +def test_execute_async_with_options( + mocker: MockerFixture, database: Database, app_context: None +) -> None: + """Test async execution with custom options.""" + mocker.patch.dict( + current_app.config, {"SQL_QUERY_MUTATOR": None, "SQLLAB_TIMEOUT": 30} + ) + + mock_celery_task = mocker.patch( + "superset.sql.execution.celery_task.execute_sql_task" + ) + + options = QueryOptions(catalog="analytics", schema="reports") + result = database.execute_async("SELECT * FROM sales", options=options) + + assert result.status == QueryStatus.PENDING + mock_celery_task.delay.assert_called_once() + + +def test_execute_async_dml_without_permission_raises( + mocker: MockerFixture, database: Database, app_context: None +) -> None: + """Test that async DML queries raise exception when not allowed.""" + from superset.exceptions import SupersetSecurityException + + mocker.patch.dict( + current_app.config, {"SQL_QUERY_MUTATOR": None, "SQLLAB_TIMEOUT": 30} + ) + + with pytest.raises(SupersetSecurityException, match="DML queries are not allowed"): + database.execute_async("INSERT INTO users (name) VALUES ('test')") + + +def test_async_handle_get_status( + mocker: MockerFixture, + database: Database, + app_context: None, + mock_db_session: MagicMock, +) -> None: + """Test that async handle can retrieve query status.""" + from superset.models.sql_lab import Query + + mock_query = MagicMock(spec=Query) + mock_query.status = "success" + filter_mock = mock_db_session.query.return_value.filter_by.return_value + filter_mock.one_or_none.return_value = mock_query + + mocker.patch.dict( + current_app.config, {"SQL_QUERY_MUTATOR": None, "SQLLAB_TIMEOUT": 30} + ) + mocker.patch("superset.sql.execution.celery_task.execute_sql_task") + + result = database.execute_async("SELECT * FROM users") + + status = result.get_status() + assert status == QueryStatus.SUCCESS + + +def test_async_handle_cancel( + mocker: MockerFixture, + database: Database, + app_context: None, + mock_db_session: MagicMock, +) -> None: + """Test that async handle can cancel a query.""" + from superset.models.sql_lab import Query + from superset.sql.execution.executor import SQLExecutor + + mock_query = MagicMock(spec=Query) + mock_query.status = "running" + mock_query.database = database + filter_mock = mock_db_session.query.return_value.filter_by.return_value + filter_mock.one_or_none.return_value = mock_query + + mocker.patch.dict( + current_app.config, {"SQL_QUERY_MUTATOR": None, "SQLLAB_TIMEOUT": 30} + ) + mocker.patch("superset.sql.execution.celery_task.execute_sql_task") + mock_cancel = mocker.patch.object(SQLExecutor, "_cancel_query", return_value=True) + + result = database.execute_async("SELECT * FROM users") + + cancelled = result.cancel() + assert cancelled is True + mock_cancel.assert_called_once() + + +# ============================================================================= +# SQL Mutation Tests +# ============================================================================= + + +def test_execute_applies_sql_mutator( + mocker: MockerFixture, database: Database, app_context: None +) -> None: + """Test that SQL mutator is applied internally.""" + mock_query_execution(mocker, database, return_data=[(1,)], column_names=["id"]) + + # Track what SQL gets mutated + mutate_mock = mocker.patch.object( + database, "mutate_sql_based_on_config", side_effect=lambda sql, **kw: sql + ) + mocker.patch.dict( + current_app.config, + { + "SQL_QUERY_MUTATOR": lambda sql, **kwargs: f"/* mutated */ {sql}", + "SQLLAB_TIMEOUT": 30, + "SQL_MAX_ROW": None, + "QUERY_LOGGER": None, + }, + ) + + result = database.execute("SELECT id FROM users") + + assert result.status == QueryStatus.SUCCESS + # Verify mutate_sql_based_on_config was called + mutate_mock.assert_called() + + +# ============================================================================= +# Progress Tracking Tests +# ============================================================================= + + +def test_execute_multi_statement_updates_query_progress( + mocker: MockerFixture, + database: Database, + app_context: None, + mock_db_session: MagicMock, +) -> None: + """Test that multi-statement execution updates Query.progress.""" + from superset.models.sql_lab import Query as QueryModel + from superset.result_set import SupersetResultSet + + # Mock raw connection for multi-statement execution + mock_cursor = MagicMock() + mock_cursor.description = [("id",), ("name",)] + mock_conn = MagicMock() + mock_conn.cursor.return_value = mock_cursor + mock_conn.__enter__ = MagicMock(return_value=mock_conn) + mock_conn.__exit__ = MagicMock(return_value=False) + + mocker.patch.object(database, "get_raw_connection", return_value=mock_conn) + mocker.patch.object( + database, + "mutate_sql_based_on_config", + side_effect=lambda sql, **kw: sql, + ) + mocker.patch.object(database.db_engine_spec, "execute") + mocker.patch.object( + database.db_engine_spec, "fetch_data", return_value=[("1", "Alice")] + ) + + mock_result_set = MagicMock(spec=SupersetResultSet) + mock_result_set.to_pandas_df.return_value = pd.DataFrame( + {"id": ["1"], "name": ["Alice"]} + ) + mocker.patch("superset.result_set.SupersetResultSet", return_value=mock_result_set) + + mocker.patch.dict( + current_app.config, + { + "SQL_QUERY_MUTATOR": None, + "SQLLAB_TIMEOUT": 30, + "SQL_MAX_ROW": None, + "QUERY_LOGGER": None, + }, + ) + + # Track progress updates on the Query model + mock_query = MagicMock(spec=QueryModel) + mock_query.id = 123 + mocker.patch("superset.models.sql_lab.Query", return_value=mock_query) + + # Execute multiple statements + result = database.execute("SELECT 1; SELECT 2;") + + assert result.status == QueryStatus.SUCCESS + # Query.progress should have been updated + assert mock_query.progress == 100 # Final progress after completion + # set_extra_json_key should have been called for progress messages + assert mock_query.set_extra_json_key.call_count >= 2 + + +# ============================================================================= +# Query Logging Tests +# ============================================================================= + + +def test_execute_calls_query_logger( + mocker: MockerFixture, database: Database, app_context: None +) -> None: + """Test that QUERY_LOGGER is called when configured.""" + mock_query_execution(mocker, database, return_data=[(1,)], column_names=["id"]) + log_calls: list[tuple[str, str, str | None]] = [] + + def mock_logger(db, sql, schema, module, security_manager, context) -> None: # noqa: ARG001 + log_calls.append((db.database_name, sql, schema)) + + mocker.patch.dict( + current_app.config, + { + "SQL_QUERY_MUTATOR": None, + "SQLLAB_TIMEOUT": 30, + "SQL_MAX_ROW": None, + "QUERY_LOGGER": mock_logger, + }, + ) + + result = database.execute( + "SELECT * FROM users", options=QueryOptions(schema="public") + ) + + assert result.status == QueryStatus.SUCCESS + assert len(log_calls) == 1 + assert log_calls[0][0] == "test_db" + # SQL may be formatted/prettified, so check for key parts + logged_sql = log_calls[0][1] + assert "SELECT" in logged_sql + assert "FROM users" in logged_sql + assert log_calls[0][2] == "public" + + +def test_execute_no_query_logger_configured( + mocker: MockerFixture, database: Database, app_context: None +) -> None: + """Test that execution works when QUERY_LOGGER is not configured.""" + mock_query_execution(mocker, database, return_data=[(1,)], column_names=["id"]) + mocker.patch.dict( + current_app.config, + { + "SQL_QUERY_MUTATOR": None, + "SQLLAB_TIMEOUT": 30, + "SQL_MAX_ROW": None, + "QUERY_LOGGER": None, + }, + ) + + # Should not raise any errors + result = database.execute("SELECT * FROM users") + + assert result.status == QueryStatus.SUCCESS + + +# ============================================================================= +# Dry Run Tests +# ============================================================================= + + +def test_execute_dry_run_returns_transformed_sql( + mocker: MockerFixture, database: Database, app_context: None +) -> None: + """Test dry run returns transformed SQL without execution.""" + from superset.sql.execution.executor import SQLExecutor + + mocker.patch.dict( + current_app.config, + {"SQL_QUERY_MUTATOR": None, "SQLLAB_TIMEOUT": 30, "SQL_MAX_ROW": 100}, + ) + + # Mock _apply_rls_to_script to verify it's called even in dry run + mock_apply_rls = mocker.patch.object(SQLExecutor, "_apply_rls_to_script") + + # get_raw_connection should NOT be called in dry run + get_conn_mock = mocker.patch.object(database, "get_raw_connection") + + options = QueryOptions(dry_run=True, limit=50) + result = database.execute("SELECT * FROM users", options=options) + + assert result.status == QueryStatus.SUCCESS + assert len(result.statements) > 0 # Transformed SQL returned in statements + assert result.statements[0].original_sql is not None # Has original SQL + assert result.statements[0].executed_sql is not None # Has transformed SQL + assert result.statements[0].data is None # No data in dry run + assert result.query_id is None # No Query model created + mock_apply_rls.assert_called() # RLS still applied + get_conn_mock.assert_not_called() # No execution + + +def test_execute_async_dry_run_returns_transformed_sql( + mocker: MockerFixture, database: Database, app_context: None +) -> None: + """Test async dry run returns transformed SQL without Celery submission.""" + mocker.patch.dict( + current_app.config, + {"SQL_QUERY_MUTATOR": None, "SQLLAB_TIMEOUT": 30}, + ) + + # Mock Celery task - should NOT be called + mock_celery_task = mocker.patch( + "superset.sql.execution.celery_task.execute_sql_task" + ) + + options = QueryOptions(dry_run=True) + result = database.execute_async("SELECT * FROM users", options=options) + + assert result.status == QueryStatus.SUCCESS + assert result.query_id is None # None for cached/dry-run results + mock_celery_task.delay.assert_not_called() + + # Handle should return the dry run result + status = result.get_status() + assert status == QueryStatus.SUCCESS + + +def test_execute_async_cached_result_returns_immediately( + mocker: MockerFixture, database: Database, app_context: None +) -> None: + """Test async execution with cached result returns immediately.""" + from superset.sql.execution.executor import SQLExecutor + + cached_df = pd.DataFrame({"id": [1, 2]}) + cached_result = MagicMock() + cached_result.status = QueryStatus.SUCCESS + cached_result.data = cached_df + cached_result.is_cached = True + + mocker.patch.dict( + current_app.config, + {"SQL_QUERY_MUTATOR": None, "SQLLAB_TIMEOUT": 30}, + ) + mocker.patch.object(SQLExecutor, "_get_from_cache", return_value=cached_result) + + # Mock Celery task - should NOT be called + mock_celery_task = mocker.patch( + "superset.sql.execution.celery_task.execute_sql_task" + ) + + result = database.execute_async("SELECT * FROM users") + + assert result.status == QueryStatus.SUCCESS + assert result.query_id is None + mock_celery_task.delay.assert_not_called() + + +# ============================================================================= +# Empty Statement Tests +# ============================================================================= + + +def test_execute_empty_sql( + mocker: MockerFixture, database: Database, app_context: None +) -> None: + """Test execution with empty SQL.""" + from superset.sql.execution.executor import SQLExecutor + + mock_query_execution(mocker, database, return_data=[], column_names=[]) + mocker.patch.dict( + current_app.config, + { + "SQL_QUERY_MUTATOR": None, + "SQLLAB_TIMEOUT": 30, + "SQL_MAX_ROW": None, + "QUERY_LOGGER": None, + }, + ) + + # Mock _create_query_record + mock_query = MagicMock() + mock_query.id = 123 + mocker.patch.object(SQLExecutor, "_create_query_record", return_value=mock_query) + + result = database.execute("") + + assert result.status == QueryStatus.SUCCESS + assert sum(s.row_count for s in result.statements) == 0 + + +def test_execute_sql_with_cursor_empty_statements(app_context: None) -> None: + """Test execute_sql_with_cursor with empty statement list.""" + from superset.sql.execution.executor import execute_sql_with_cursor + + mock_database = MagicMock() + mock_cursor = MagicMock() + mock_query = MagicMock() + + result = execute_sql_with_cursor( + database=mock_database, + cursor=mock_cursor, + statements=[], + query=mock_query, + ) + + assert result == [] # Returns empty list for empty statements + + +def test_execute_sql_with_cursor_stopped_mid_execution( + mocker: MockerFixture, app_context: None +) -> None: + """Test execute_sql_with_cursor when query is stopped mid-execution.""" + from superset.sql.execution.executor import execute_sql_with_cursor + + mock_database = MagicMock() + mock_database.mutate_sql_based_on_config = lambda sql, **kw: sql + mock_database.db_engine_spec.execute = MagicMock() + mock_database.db_engine_spec.fetch_data = MagicMock(return_value=[]) + + mock_cursor = MagicMock() + mock_cursor.description = [("id",)] + mock_cursor.fetchall = MagicMock() + + mock_query = MagicMock() + mock_query.schema = "public" + mock_query.progress = 0 + mock_query.set_extra_json_key = MagicMock() + + mocker.patch("superset.sql.execution.executor.db.session") + + # Check stopped function returns True after first statement + call_count = {"count": 0} + + def check_stopped(): + call_count["count"] += 1 + return call_count["count"] > 1 + + result = execute_sql_with_cursor( + database=mock_database, + cursor=mock_cursor, + statements=["SELECT 1", "SELECT 2", "SELECT 3"], + query=mock_query, + check_stopped_fn=check_stopped, + ) + + # Returns results collected before stopped (first statement completed) + assert len(result) == 1 # Only first statement completed before stop + + +def test_execute_sql_with_cursor_custom_execute_fn( + mocker: MockerFixture, app_context: None +) -> None: + """Test execute_sql_with_cursor with custom execute function.""" + from superset.sql.execution.executor import execute_sql_with_cursor + + mock_database = MagicMock() + mock_database.mutate_sql_based_on_config = lambda sql, **kw: sql + mock_database.db_engine_spec = MagicMock() + + mock_cursor = MagicMock() + mock_cursor.description = [("count",)] + mock_cursor.fetchall = MagicMock() + + mock_query = MagicMock() + mock_query.schema = "public" + mock_query.progress = 0 + mock_query.set_extra_json_key = MagicMock() + + mocker.patch("superset.sql.execution.executor.db.session") + mock_database.db_engine_spec.fetch_data = MagicMock(return_value=[(100,)]) + + custom_execute_calls = [] + + def custom_execute(cursor, sql): + custom_execute_calls.append(sql) + + result = execute_sql_with_cursor( + database=mock_database, + cursor=mock_cursor, + statements=["SELECT COUNT(*) FROM users"], + query=mock_query, + execute_fn=custom_execute, + ) + + assert len(custom_execute_calls) == 1 + assert result is not None + + +# ============================================================================= +# Limit Application Tests +# ============================================================================= + + +def test_execute_applies_limit( + mocker: MockerFixture, database: Database, app_context: None +) -> None: + """Test that limit is applied to SELECT queries.""" + from superset.sql.execution.executor import SQLExecutor + + mock_query_execution(mocker, database, return_data=[(1,)], column_names=["id"]) + mocker.patch.dict( + current_app.config, + { + "SQL_QUERY_MUTATOR": None, + "SQLLAB_TIMEOUT": 30, + "SQL_MAX_ROW": None, + "QUERY_LOGGER": None, + }, + ) + + # Mock _apply_limit_to_script + apply_limit_mock = mocker.patch.object(SQLExecutor, "_apply_limit_to_script") + + options = QueryOptions(limit=50) + result = database.execute("SELECT * FROM users", options=options) + + assert result.status == QueryStatus.SUCCESS + # Verify limit was applied + assert apply_limit_mock.called + + +def test_execute_respects_sql_max_row( + mocker: MockerFixture, database: Database, app_context: None +) -> None: + """Test that SQL_MAX_ROW config limits the effective limit.""" + from superset.sql.execution.executor import SQLExecutor + + mock_query_execution(mocker, database, return_data=[(1,)], column_names=["id"]) + mocker.patch.dict( + current_app.config, + { + "SQL_QUERY_MUTATOR": None, + "SQLLAB_TIMEOUT": 30, + "SQL_MAX_ROW": 100, + "QUERY_LOGGER": None, + }, + ) + + apply_limit_mock = mocker.patch.object(SQLExecutor, "_apply_limit_to_script") + + # Request 1000 but should be capped at 100 + options = QueryOptions(limit=1000) + result = database.execute("SELECT * FROM users", options=options) + + assert result.status == QueryStatus.SUCCESS + # Verify limit was applied + assert apply_limit_mock.called + + +def test_execute_no_limit_for_dml( + mocker: MockerFixture, database_with_dml: Database, app_context: None +) -> None: + """Test that limit is not applied to DML queries.""" + from superset.sql.execution.executor import SQLExecutor + + mock_query_execution(mocker, database_with_dml, return_data=[], column_names=[]) + mocker.patch.dict( + current_app.config, + { + "SQL_QUERY_MUTATOR": None, + "SQLLAB_TIMEOUT": 30, + "SQL_MAX_ROW": 100, + "QUERY_LOGGER": None, + }, + ) + + apply_limit_mock = mocker.patch.object(SQLExecutor, "_apply_limit_to_script") + + options = QueryOptions(limit=50) + database_with_dml.execute("INSERT INTO users VALUES (1)", options=options) + + # Should not apply limit to DML + apply_limit_mock.assert_not_called() + + +def test_apply_limit_to_script_respects_sql_max_row( + mocker: MockerFixture, database: Database, app_context: None +) -> None: + """Test that _apply_limit_to_script caps limit at SQL_MAX_ROW.""" + from superset.sql.execution.executor import SQLExecutor + from superset.sql.parse import SQLScript + + mocker.patch.dict( + current_app.config, + { + "SQL_MAX_ROW": 100, + }, + ) + + executor = SQLExecutor(database) + + # Create a script with a statement + script = SQLScript("SELECT * FROM users", database.db_engine_spec.engine) + + # Mock set_limit_value on the first statement + set_limit_mock = mocker.patch.object(script.statements[0], "set_limit_value") + + options = QueryOptions(limit=1000) # Request 1000 but should be capped at 100 + executor._apply_limit_to_script(script, options) + + # Verify limit was capped at SQL_MAX_ROW (100) + set_limit_mock.assert_called_once_with(100, database.db_engine_spec.limit_method) + + +def test_apply_limit_to_script_with_empty_statements( + mocker: MockerFixture, database: Database, app_context: None +) -> None: + """Test that _apply_limit_to_script handles empty script.statements.""" + from superset.sql.execution.executor import SQLExecutor + from superset.sql.parse import SQLScript + + mocker.patch.dict( + current_app.config, + { + "SQL_MAX_ROW": 100, + }, + ) + + executor = SQLExecutor(database) + + # Create a script with no statements (empty string) + script = SQLScript("", database.db_engine_spec.engine) + + options = QueryOptions(limit=50) + # Should not raise any errors when statements is empty + executor._apply_limit_to_script(script, options) + + # Verify no error occurred and statements is still empty + assert script.statements == [] + + +# ============================================================================= +# Catalog/Schema Resolution Tests +# ============================================================================= + + +def test_execute_uses_default_catalog_and_schema( + mocker: MockerFixture, database: Database, app_context: None +) -> None: + """Test that default catalog and schema are used when not specified.""" + mock_query_execution(mocker, database, return_data=[(1,)], column_names=["id"]) + get_default_catalog_mock = mocker.patch.object( + database, "get_default_catalog", return_value="main" + ) + get_default_schema_mock = mocker.patch.object( + database, "get_default_schema", return_value="public" + ) + mocker.patch.dict( + current_app.config, + { + "SQL_QUERY_MUTATOR": None, + "SQLLAB_TIMEOUT": 30, + "SQL_MAX_ROW": None, + "QUERY_LOGGER": None, + }, + ) + + result = database.execute("SELECT * FROM users") + + assert result.status == QueryStatus.SUCCESS + # Verify default catalog/schema were fetched + get_default_catalog_mock.assert_called() + get_default_schema_mock.assert_called() + + +# ============================================================================= +# Async Query Status and Result Tests +# ============================================================================= + + +def test_async_handle_get_result_query_not_found( + mocker: MockerFixture, + database: Database, + app_context: None, + mock_db_session: MagicMock, +) -> None: + """Test getting result for non-existent query.""" + + # Query not found + filter_mock = mock_db_session.query.return_value.filter_by.return_value + filter_mock.one_or_none.return_value = None + + mocker.patch.dict( + current_app.config, {"SQL_QUERY_MUTATOR": None, "SQLLAB_TIMEOUT": 30} + ) + mocker.patch("superset.sql.execution.celery_task.execute_sql_task") + + result = database.execute_async("SELECT * FROM users") + + # Try to get result + query_result = result.get_result() + + assert query_result.status == QueryStatus.FAILED + assert "not found" in query_result.error_message.lower() # type: ignore[union-attr] + + +def test_async_handle_get_result_pending( + mocker: MockerFixture, + database: Database, + app_context: None, + mock_db_session: MagicMock, +) -> None: + """Test getting result for pending query.""" + from superset.models.sql_lab import Query + + mock_query = MagicMock(spec=Query) + mock_query.status = "pending" + mock_query.error_message = None + mock_query.executed_sql = "SELECT * FROM users" + mock_query.results_key = None + + filter_mock = mock_db_session.query.return_value.filter_by.return_value + filter_mock.one_or_none.return_value = mock_query + + mocker.patch.dict( + current_app.config, {"SQL_QUERY_MUTATOR": None, "SQLLAB_TIMEOUT": 30} + ) + mocker.patch("superset.sql.execution.celery_task.execute_sql_task") + + result = database.execute_async("SELECT * FROM users") + + query_result = result.get_result() + + assert query_result.status == QueryStatus.PENDING + + +def test_async_handle_get_result_with_results_backend( + mocker: MockerFixture, + database: Database, + app_context: None, + mock_db_session: MagicMock, +) -> None: + """Test getting result from results backend.""" + from superset.models.sql_lab import Query + + mock_query = MagicMock(spec=Query) + mock_query.status = "success" + mock_query.error_message = None + mock_query.executed_sql = "SELECT * FROM users" + mock_query.results_key = "result_key_123" + + filter_mock = mock_db_session.query.return_value.filter_by.return_value + filter_mock.one_or_none.return_value = mock_query + + # Mock results backend with new multi-statement format + payload = msgpack.dumps( + { + "statements": [ + { + "original_sql": "SELECT * FROM users", + "executed_sql": "SELECT * FROM users", + "data": [{"id": 1, "name": "Alice"}, {"id": 2, "name": "Bob"}], + "columns": [ + {"column_name": "id", "name": "id"}, + {"column_name": "name", "name": "name"}, + ], + "row_count": 2, + "execution_time_ms": 10.0, + } + ], + "total_execution_time_ms": 10.0, + } + ) + compressed_payload = b"compressed_data" + + mock_results_backend = MagicMock() + mock_results_backend.get.return_value = compressed_payload + + mock_results_backend_manager = MagicMock() + mock_results_backend_manager.results_backend = mock_results_backend + + mocker.patch( + "superset.results_backend_manager", + mock_results_backend_manager, + ) + mocker.patch("superset.utils.core.zlib_decompress", return_value=payload) + mocker.patch.dict( + current_app.config, {"SQL_QUERY_MUTATOR": None, "SQLLAB_TIMEOUT": 30} + ) + mocker.patch("superset.sql.execution.celery_task.execute_sql_task") + + result = database.execute_async("SELECT * FROM users") + + query_result = result.get_result() + + assert query_result.status == QueryStatus.SUCCESS + assert len(query_result.statements) > 0 + assert query_result.statements[0].data is not None + assert sum(s.row_count for s in query_result.statements) == 2 + + +def test_async_handle_get_result_backend_load_error( + mocker: MockerFixture, + database: Database, + app_context: None, + mock_db_session: MagicMock, +) -> None: + """Test error handling when loading results from backend.""" + from superset.models.sql_lab import Query + + mock_query = MagicMock(spec=Query) + mock_query.status = "success" + mock_query.error_message = None + mock_query.executed_sql = "SELECT * FROM users" + mock_query.results_key = "result_key_123" + + filter_mock = mock_db_session.query.return_value.filter_by.return_value + filter_mock.one_or_none.return_value = mock_query + + # Mock results backend with error + mock_results_backend = MagicMock() + mock_results_backend.get.return_value = b"invalid_data" + + mock_results_backend_manager = MagicMock() + mock_results_backend_manager.results_backend = mock_results_backend + + mocker.patch( + "superset.results_backend_manager", + mock_results_backend_manager, + ) + mocker.patch( + "superset.utils.core.zlib_decompress", + side_effect=Exception("Decompression failed"), + ) + mocker.patch.dict( + current_app.config, {"SQL_QUERY_MUTATOR": None, "SQLLAB_TIMEOUT": 30} + ) + mocker.patch("superset.sql.execution.celery_task.execute_sql_task") + + result = database.execute_async("SELECT * FROM users") + + query_result = result.get_result() + + assert query_result.status == QueryStatus.FAILED + assert "Error loading results" in query_result.error_message # type: ignore[operator] + + +def test_async_handle_get_result_no_results_key( + mocker: MockerFixture, + database: Database, + app_context: None, + mock_db_session: MagicMock, +) -> None: + """Test getting result when results_key is missing.""" + from superset.models.sql_lab import Query + + mock_query = MagicMock(spec=Query) + mock_query.status = "success" + mock_query.error_message = None + mock_query.executed_sql = "SELECT * FROM users" + mock_query.results_key = None # No results key + + filter_mock = mock_db_session.query.return_value.filter_by.return_value + filter_mock.one_or_none.return_value = mock_query + + mocker.patch.dict( + current_app.config, {"SQL_QUERY_MUTATOR": None, "SQLLAB_TIMEOUT": 30} + ) + mocker.patch("superset.sql.execution.celery_task.execute_sql_task") + + result = database.execute_async("SELECT * FROM users") + + query_result = result.get_result() + + assert query_result.status == QueryStatus.FAILED + assert "Results not available" in query_result.error_message # type: ignore[operator] + + +def test_async_handle_get_status_query_not_found( + mocker: MockerFixture, + database: Database, + app_context: None, + mock_db_session: MagicMock, +) -> None: + """Test getting status for non-existent query.""" + # Query not found + filter_mock = mock_db_session.query.return_value.filter_by.return_value + filter_mock.one_or_none.return_value = None + + mocker.patch.dict( + current_app.config, {"SQL_QUERY_MUTATOR": None, "SQLLAB_TIMEOUT": 30} + ) + mocker.patch("superset.sql.execution.celery_task.execute_sql_task") + + result = database.execute_async("SELECT * FROM users") + + # Override the query_id to simulate looking up a non-existent query + object.__setattr__(result, "query_id", 999999) + + status = result.get_status() + + assert status == QueryStatus.FAILED + + +# ============================================================================= +# Query Cancellation Tests +# ============================================================================= + + +def test_cancel_query_implicit_cancel( + mocker: MockerFixture, database: Database, app_context: None +) -> None: + """Test cancel_query with implicit cancellation.""" + from superset.sql.execution.executor import SQLExecutor + + mock_query = MagicMock() + mock_query.extra = {} + + mocker.patch.object( + database.db_engine_spec, "has_implicit_cancel", return_value=True + ) + + result = SQLExecutor._cancel_query(database, mock_query) + + assert result is True + + +def test_cancel_query_early_cancel_flag( + mocker: MockerFixture, database: Database, app_context: None +) -> None: + """Test cancel_query with early cancel flag set.""" + from superset.constants import QUERY_EARLY_CANCEL_KEY + from superset.sql.execution.executor import SQLExecutor + + mock_query = MagicMock() + mock_query.extra = {QUERY_EARLY_CANCEL_KEY: True} + + mocker.patch.object( + database.db_engine_spec, "has_implicit_cancel", return_value=False + ) + prepare_cancel_mock = mocker.patch.object( + database.db_engine_spec, "prepare_cancel_query" + ) + + result = SQLExecutor._cancel_query(database, mock_query) + + assert result is True + prepare_cancel_mock.assert_called_with(mock_query) + + +def test_cancel_query_no_cancel_id( + mocker: MockerFixture, database: Database, app_context: None +) -> None: + """Test cancel_query when no cancel ID is available.""" + from superset.sql.execution.executor import SQLExecutor + + mock_query = MagicMock() + mock_query.extra = {} + + mocker.patch.object( + database.db_engine_spec, "has_implicit_cancel", return_value=False + ) + mocker.patch.object(database.db_engine_spec, "prepare_cancel_query") + + result = SQLExecutor._cancel_query(database, mock_query) + + assert result is False + + +def test_cancel_query_with_cancel_id( + mocker: MockerFixture, database: Database, app_context: None +) -> None: + """Test cancel_query executes cancellation with cancel ID.""" + from superset.constants import QUERY_CANCEL_KEY + from superset.sql.execution.executor import SQLExecutor + from superset.utils.core import QuerySource + + mock_query = MagicMock() + mock_query.extra = {QUERY_CANCEL_KEY: "cancel_123"} + mock_query.catalog = "main" + mock_query.schema = "public" + + mocker.patch.object( + database.db_engine_spec, "has_implicit_cancel", return_value=False + ) + mocker.patch.object(database.db_engine_spec, "prepare_cancel_query") + cancel_query_mock = mocker.patch.object( + database.db_engine_spec, "cancel_query", return_value=True + ) + + # Mock engine and connection + mock_cursor = MagicMock() + mock_conn = MagicMock() + mock_conn.cursor.return_value.__enter__ = MagicMock(return_value=mock_cursor) + mock_conn.cursor.return_value.__exit__ = MagicMock(return_value=False) + mock_conn.__enter__ = MagicMock(return_value=mock_conn) + mock_conn.__exit__ = MagicMock(return_value=False) + + mock_engine = MagicMock() + mock_engine.raw_connection.return_value.__enter__ = MagicMock( + return_value=mock_conn + ) + mock_engine.raw_connection.return_value.__exit__ = MagicMock(return_value=False) + mock_engine.__enter__ = MagicMock(return_value=mock_engine) + mock_engine.__exit__ = MagicMock(return_value=False) + + get_sqla_engine_mock = mocker.patch.object( + database, "get_sqla_engine", return_value=mock_engine + ) + + result = SQLExecutor._cancel_query(database, mock_query) + + assert result is True + get_sqla_engine_mock.assert_called_with( + catalog="main", schema="public", source=QuerySource.SQL_LAB + ) + cancel_query_mock.assert_called_once() + + +def test_async_handle_cancel_query_not_found( + mocker: MockerFixture, + database: Database, + app_context: None, + mock_db_session: MagicMock, +) -> None: + """Test cancelling non-existent query.""" + + # Query not found + filter_mock = mock_db_session.query.return_value.filter_by.return_value + filter_mock.one_or_none.return_value = None + + mocker.patch.dict( + current_app.config, {"SQL_QUERY_MUTATOR": None, "SQLLAB_TIMEOUT": 30} + ) + mocker.patch("superset.sql.execution.celery_task.execute_sql_task") + + result = database.execute_async("SELECT * FROM users") + + # Override to simulate non-existent query + object.__setattr__(result, "query_id", 999999) + + cancelled = result.cancel() + + assert cancelled is False + + +# ============================================================================= +# Cache Timeout Tests +# ============================================================================= + + +def test_execute_uses_database_cache_timeout( + mocker: MockerFixture, database: Database, app_context: None +) -> None: + """Test that database cache timeout is used when available.""" + from superset.sql.execution.executor import SQLExecutor + + mock_query_execution(mocker, database, return_data=[(1,)], column_names=["id"]) + mocker.patch.dict( + current_app.config, + { + "SQL_QUERY_MUTATOR": None, + "SQLLAB_TIMEOUT": 30, + "SQL_MAX_ROW": None, + "QUERY_LOGGER": None, + "CACHE_DEFAULT_TIMEOUT": 300, + }, + ) + + database.cache_timeout = 600 # Custom timeout + + # Mock cache operations + mocker.patch.object(SQLExecutor, "_get_from_cache", return_value=None) + mock_cache_set = mocker.patch("superset.extensions.cache_manager.data_cache.set") + + result = database.execute("SELECT * FROM users") + + assert result.status == QueryStatus.SUCCESS + # Verify cache timeout used + if mock_cache_set.called: + call_kwargs = mock_cache_set.call_args[1] + assert call_kwargs.get("timeout") == 600 + + +def test_execute_uses_custom_cache_timeout_option( + mocker: MockerFixture, database: Database, app_context: None +) -> None: + """Test that custom cache timeout from options is used.""" + from superset.sql.execution.executor import SQLExecutor + + mock_query_execution(mocker, database, return_data=[(1,)], column_names=["id"]) + mocker.patch.dict( + current_app.config, + { + "SQL_QUERY_MUTATOR": None, + "SQLLAB_TIMEOUT": 30, + "SQL_MAX_ROW": None, + "QUERY_LOGGER": None, + "CACHE_DEFAULT_TIMEOUT": 300, + }, + ) + + # Mock cache operations + mocker.patch.object(SQLExecutor, "_get_from_cache", return_value=None) + mock_cache_set = mocker.patch("superset.extensions.cache_manager.data_cache.set") + + options = QueryOptions(cache=CacheOptions(timeout=1200)) + result = database.execute("SELECT * FROM users", options=options) + + assert result.status == QueryStatus.SUCCESS + # Verify custom timeout used + if mock_cache_set.called: + call_kwargs = mock_cache_set.call_args[1] + assert call_kwargs.get("timeout") == 1200 + + +# ============================================================================= +# Celery Submission Error Tests +# ============================================================================= + + +def test_execute_async_celery_submission_error( + mocker: MockerFixture, + database: Database, + app_context: None, + mock_db_session: MagicMock, +) -> None: + """Test error handling when Celery submission fails.""" + mocker.patch.dict( + current_app.config, {"SQL_QUERY_MUTATOR": None, "SQLLAB_TIMEOUT": 30} + ) + + # Mock Celery task to raise exception + mock_celery_task = mocker.patch( + "superset.sql.execution.celery_task.execute_sql_task" + ) + mock_celery_task.delay.side_effect = Exception("Celery connection failed") + + with pytest.raises(Exception, match="Celery connection failed"): + database.execute_async("SELECT * FROM users") + + +# ============================================================================= +# Additional Edge Case Tests for 100% Coverage +# ============================================================================= + + +def test_execute_sql_with_cursor_no_rows_or_description( + mocker: MockerFixture, app_context: None +) -> None: + """Test execute_sql_with_cursor when cursor returns no rows and description.""" + from superset.sql.execution.executor import execute_sql_with_cursor + + mock_database = MagicMock() + mock_database.mutate_sql_based_on_config = lambda sql, **kw: sql + mock_database.db_engine_spec.execute = MagicMock() + mock_database.db_engine_spec.fetch_data = MagicMock(return_value=None) + + mock_cursor = MagicMock() + mock_cursor.description = None # No description + mock_cursor.fetchall = MagicMock() + + mock_query = MagicMock() + mock_query.schema = "public" + mock_query.progress = 0 + mock_query.set_extra_json_key = MagicMock() + + mocker.patch("superset.sql.execution.executor.db.session") + + result = execute_sql_with_cursor( + database=mock_database, + cursor=mock_cursor, + statements=["INSERT INTO users VALUES (1)"], + query=mock_query, + ) + + # DML statement returns result with result_set=None + assert len(result) == 1 + assert result[0][1] is None # result_set is None for DML + + +def test_execute_with_exception_on_execute( + mocker: MockerFixture, database: Database, app_context: None +) -> None: + """Test that _execute_statements returns empty DataFrame on None result.""" + from superset.sql.execution.executor import SQLExecutor + + # Mock to simulate None result_set + mocker.patch( + "superset.sql.execution.executor.execute_sql_with_cursor", + return_value=[], # Empty list for no results + ) + + mock_query = MagicMock() + mock_query.id = 123 + mocker.patch.object(SQLExecutor, "_create_query_record", return_value=mock_query) + + mock_conn = MagicMock() + mock_conn.cursor.return_value = MagicMock() + mock_conn.__enter__ = MagicMock(return_value=mock_conn) + mock_conn.__exit__ = MagicMock(return_value=False) + mocker.patch.object(database, "get_raw_connection", return_value=mock_conn) + + mocker.patch.dict( + current_app.config, + { + "SQL_QUERY_MUTATOR": None, + "SQLLAB_TIMEOUT": 30, + "SQL_MAX_ROW": None, + "QUERY_LOGGER": None, + }, + ) + + result = database.execute("SELECT * FROM nonexistent") + + # Should handle None result_set gracefully + assert result.status == QueryStatus.SUCCESS + assert sum(s.row_count for s in result.statements) == 0 + + +def test_check_disallowed_functions_no_config( + mocker: MockerFixture, database: Database, app_context: None +) -> None: + """Test disallowed functions check when no config exists.""" + from superset.sql.execution.executor import SQLExecutor + + mocker.patch.dict(current_app.config, {"DISALLOWED_SQL_FUNCTIONS": {}}) + + executor = SQLExecutor(database) + script = MagicMock() + result = executor._check_disallowed_functions(script) + + assert result is None + + +def test_try_get_cached_result_with_mutation( + mocker: MockerFixture, database: Database, app_context: None +) -> None: + """Test that cache is skipped for mutation queries.""" + from superset.sql.execution.executor import SQLExecutor + + executor = SQLExecutor(database) + + script = MagicMock() + script.has_mutation.return_value = True + + result = executor._try_get_cached_result(script, "INSERT INTO foo", MagicMock()) + + # Should skip cache for mutations + assert result is None + + +def test_store_in_cache_with_failed_status( + mocker: MockerFixture, database: Database, app_context: None +) -> None: + """Test that failed queries are not cached.""" + from superset_core.api.types import QueryResult as QueryResultType + + from superset.sql.execution.executor import SQLExecutor + + executor = SQLExecutor(database) + + failed_result = QueryResultType( + status=QueryStatus.FAILED, + error_message="Test error", + ) + + mock_cache_set = mocker.patch("superset.extensions.cache_manager.data_cache.set") + + executor._store_in_cache(failed_result, "SELECT 1", QueryOptions()) + + # Should not cache failed queries + mock_cache_set.assert_not_called() + + +def test_store_in_cache_with_no_data( + mocker: MockerFixture, database: Database, app_context: None +) -> None: + """Test that DML queries (with no data) are cached.""" + from superset_core.api.types import QueryResult as QueryResultType, StatementResult + + from superset.sql.execution.executor import SQLExecutor + + executor = SQLExecutor(database) + + result_no_data = QueryResultType( + status=QueryStatus.SUCCESS, + statements=[ + StatementResult( + original_sql="INSERT INTO t VALUES (1)", + executed_sql="INSERT INTO t VALUES (1)", + data=None, + row_count=1, + ) + ], + ) + + mock_cache_set = mocker.patch("superset.extensions.cache_manager.data_cache.set") + + executor._store_in_cache(result_no_data, "INSERT INTO t VALUES (1)", QueryOptions()) + + # DML queries are cached too + mock_cache_set.assert_called_once() + + +def test_create_cached_async_result_cancel( + mocker: MockerFixture, database: Database, app_context: None +) -> None: + """Test that cached async result cancel returns False.""" + from superset_core.api.types import QueryResult as QueryResultType, StatementResult + + from superset.sql.execution.executor import SQLExecutor + + mocker.patch.dict( + current_app.config, {"SQL_QUERY_MUTATOR": None, "SQLLAB_TIMEOUT": 30} + ) + + executor = SQLExecutor(database) + + cached_result = QueryResultType( + status=QueryStatus.SUCCESS, + statements=[ + StatementResult( + original_sql="SELECT 1", + executed_sql="SELECT 1", + data=pd.DataFrame({"id": [1]}), + row_count=1, + ) + ], + ) + + async_result = executor._create_cached_handle(cached_result) + + # Try to cancel a cached result + cancelled = async_result.cancel() + + # Should return False (nothing to cancel) + assert cancelled is False + + +def test_async_handle_get_result_with_empty_blob( + mocker: MockerFixture, + database: Database, + app_context: None, + mock_db_session: MagicMock, +) -> None: + """Test getting result when backend returns None for blob.""" + from superset.models.sql_lab import Query + + mock_query = MagicMock(spec=Query) + mock_query.status = "success" + mock_query.error_message = None + mock_query.executed_sql = "SELECT * FROM users" + mock_query.results_key = "result_key_123" + + filter_mock = mock_db_session.query.return_value.filter_by.return_value + filter_mock.one_or_none.return_value = mock_query + + # Mock results backend returning None for blob + mock_results_backend = MagicMock() + mock_results_backend.get.return_value = None # No blob found + + mock_results_backend_manager = MagicMock() + mock_results_backend_manager.results_backend = mock_results_backend + + mocker.patch( + "superset.results_backend_manager", + mock_results_backend_manager, + ) + mocker.patch.dict( + current_app.config, {"SQL_QUERY_MUTATOR": None, "SQLLAB_TIMEOUT": 30} + ) + mocker.patch("superset.sql.execution.celery_task.execute_sql_task") + + result = database.execute_async("SELECT * FROM users") + + query_result = result.get_result() + + # Should return failure when blob not found + assert query_result.status == QueryStatus.FAILED + assert "Results not available" in query_result.error_message # type: ignore[operator] + + +def test_async_handle_get_result_no_results_backend( + mocker: MockerFixture, + database: Database, + app_context: None, + mock_db_session: MagicMock, +) -> None: + """Test getting result when results_backend is None.""" + from superset.models.sql_lab import Query + + mock_query = MagicMock(spec=Query) + mock_query.status = "success" + mock_query.error_message = None + mock_query.executed_sql = "SELECT * FROM users" + mock_query.results_key = "result_key_123" + + filter_mock = mock_db_session.query.return_value.filter_by.return_value + filter_mock.one_or_none.return_value = mock_query + + # Mock results_backend_manager with None backend + mock_results_backend_manager = MagicMock() + mock_results_backend_manager.results_backend = None # No backend configured + + mocker.patch( + "superset.results_backend_manager", + mock_results_backend_manager, + ) + mocker.patch.dict( + current_app.config, {"SQL_QUERY_MUTATOR": None, "SQLLAB_TIMEOUT": 30} + ) + mocker.patch("superset.sql.execution.celery_task.execute_sql_task") + + result = database.execute_async("SELECT * FROM users") + + query_result = result.get_result() + + # Should return failure when no results backend + assert query_result.status == QueryStatus.FAILED + assert "Results not available" in query_result.error_message # type: ignore[operator] + + +def test_create_query_record_with_user( + mocker: MockerFixture, database: Database, app_context: None +) -> None: + """Test that _create_query_record captures user_id when user exists.""" + from flask import g + + mock_query_execution(mocker, database, return_data=[(1,)], column_names=["id"]) + + # Mock a user with get_id + mock_user = MagicMock() + mock_user.get_id.return_value = 42 + + mocker.patch("superset.sql.execution.executor.has_app_context", return_value=True) + mocker.patch.object(g, "user", mock_user, create=True) + + mocker.patch.dict( + current_app.config, + { + "SQL_QUERY_MUTATOR": None, + "SQLLAB_TIMEOUT": 30, + "SQL_MAX_ROW": None, + "QUERY_LOGGER": None, + }, + ) + + result = database.execute("SELECT * FROM users") + + assert result.status == QueryStatus.SUCCESS + mock_user.get_id.assert_called_once() + + +def test_get_from_cache_returns_cached_result( + mocker: MockerFixture, database: Database, app_context: None +) -> None: + """Test that _get_from_cache returns cached result when available.""" + from superset.extensions import cache_manager + from superset.sql.execution.executor import SQLExecutor + + executor = SQLExecutor(database) + + cached_data = { + "statements": [ + { + "original_sql": "SELECT * FROM users", + "executed_sql": "SELECT * FROM users", + "data": pd.DataFrame({"id": [1, 2]}), + "row_count": 2, + "execution_time_ms": 10.0, + } + ], + "total_execution_time_ms": 10.0, + } + + mocker.patch.object(cache_manager.data_cache, "get", return_value=cached_data) + + options = QueryOptions() + result = executor._get_from_cache("SELECT * FROM users", options) + + assert result is not None + assert result.status == QueryStatus.SUCCESS + assert result.is_cached is True + assert sum(s.row_count for s in result.statements) == 2 + + +def test_cached_async_result_get_result_returns_cached( + mocker: MockerFixture, database: Database, app_context: None +) -> None: + """Test that cached async result returns the original cached result.""" + from superset_core.api.types import QueryResult as QueryResultType, StatementResult + + from superset.sql.execution.executor import SQLExecutor + + mocker.patch.dict( + current_app.config, {"SQL_QUERY_MUTATOR": None, "SQLLAB_TIMEOUT": 30} + ) + + executor = SQLExecutor(database) + + cached_result = QueryResultType( + status=QueryStatus.SUCCESS, + statements=[ + StatementResult( + original_sql="SELECT 1", + executed_sql="SELECT 1", + data=pd.DataFrame({"id": [1, 2, 3]}), + row_count=3, + ) + ], + ) + + async_result = executor._create_cached_handle(cached_result) + + # Get the result back + retrieved_result = async_result.get_result() + + # Should return the same cached result + assert retrieved_result.status == QueryStatus.SUCCESS + assert sum(s.row_count for s in retrieved_result.statements) == 3 + assert retrieved_result is cached_result