Hugging Face's logo

Spaces:
Leon4gr45
/
powerbi
Paused

App Files Community
1
powerbi / tests /unit_tests /sql /execution /test_executor.py
Leon4gr45's picture
Leon4gr45
Upload folder using huggingface_hub (part 17)
5b7ed07 verified
Raw History Blame Contribute Delete
87.4 kB
# 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.queries.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 (
_passthrough_mutate_sql_based_on_config,
create_mock_cursor,
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
def test_execute_disallowed_tables(
mocker: MockerFixture, database: Database, app_context: None
) -> None:
"""Test that disallowed SQL tables are blocked."""
mocker.patch.dict(
current_app.config,
{
"SQL_QUERY_MUTATOR": None,
"SQLLAB_TIMEOUT": 30,
"DISALLOWED_SQL_FUNCTIONS": {},
"DISALLOWED_SQL_TABLES": {"sqlite": {"pg_stat_activity", "pg_roles"}},
},
)
result = database.execute("SELECT * FROM pg_stat_activity")
assert result.status == QueryStatus.FAILED
assert result.error_message is not None
assert "Disallowed SQL tables: pg_stat_activity" == result.error_message
def test_execute_allowed_tables(
mocker: MockerFixture, database: Database, app_context: None
) -> None:
"""Test that allowed SQL tables work normally."""
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,
"DISALLOWED_SQL_FUNCTIONS": {},
"DISALLOWED_SQL_TABLES": {"sqlite": {"pg_stat_activity", "pg_roles"}},
"QUERY_LOGGER": None,
},
)
result = database.execute("SELECT * FROM users")
assert result.status == QueryStatus.SUCCESS
def test_check_disallowed_tables_no_config(
mocker: MockerFixture, database: Database, app_context: None
) -> None:
"""Test disallowed tables check when no config exists."""
from superset.sql.execution.executor import SQLExecutor
mocker.patch.dict(current_app.config, {"DISALLOWED_SQL_TABLES": {}})
executor = SQLExecutor(database)
script = MagicMock()
result = executor._check_disallowed_tables(script)
assert result is None
# =============================================================================
# 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
def test_execute_oauth2_redirect_error_propagates(
mocker: MockerFixture, database: Database, app_context: None
) -> None:
"""Test that OAuth2RedirectError propagates instead of being swallowed."""
from superset.exceptions import OAuth2RedirectError
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=OAuth2RedirectError(
url="https://oauth.example.com/authorize",
tab_id="test-tab",
redirect_uri="https://superset.example.com/callback",
),
)
mocker.patch.dict(
current_app.config,
{
"SQL_QUERY_MUTATOR": None,
"SQLLAB_TIMEOUT": 30,
"SQL_MAX_ROW": None,
"QUERY_LOGGER": None,
},
)
with pytest.raises(OAuth2RedirectError):
database.execute("SELECT 1")
def test_execute_oauth2_error_propagates(
mocker: MockerFixture, database: Database, app_context: None
) -> None:
"""Test that OAuth2Error propagates instead of being swallowed."""
from superset.exceptions import OAuth2Error
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=OAuth2Error("No configuration found for OAuth2"),
)
mocker.patch.dict(
current_app.config,
{
"SQL_QUERY_MUTATOR": None,
"SQLLAB_TIMEOUT": 30,
"SQL_MAX_ROW": None,
"QUERY_LOGGER": None,
},
)
with pytest.raises(OAuth2Error):
database.execute("SELECT 1")
# =============================================================================
# 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()
@pytest.mark.parametrize("is_split", [True, False])
def test_execute_sql_with_cursor_forwards_is_split(
mocker: MockerFixture,
database: Database,
app_context: None,
mock_db_session: MagicMock,
mock_query: MagicMock,
is_split: bool,
) -> None:
"""
`execute_sql_with_cursor` must forward `is_split` to the SQL mutator.
`Database.mutate_sql_based_on_config` only fires `SQL_QUERY_MUTATOR` when
`is_split == MUTATE_AFTER_SPLIT`, so passing the wrong value silently skips
mutation (the SQL Lab bug behind issue #30169). This guards the contract.
"""
from superset.sql.execution.executor import execute_sql_with_cursor
mutate_mock: MagicMock = mocker.patch.object(
database,
"mutate_sql_based_on_config",
side_effect=_passthrough_mutate_sql_based_on_config,
)
mocker.patch.object(database.db_engine_spec, "execute")
mocker.patch.object(database.db_engine_spec, "fetch_data", return_value=[(1,)])
mocker.patch("superset.result_set.SupersetResultSet", return_value=MagicMock())
cursor: MagicMock = create_mock_cursor(["id"], data=[(1,)])
execute_sql_with_cursor(
database=database,
cursor=cursor,
statements=["SELECT id FROM t"],
query=mock_query,
is_split=is_split,
)
mutate_mock.assert_called_once()
assert mutate_mock.call_args.kwargs["is_split"] is is_split
# =============================================================================
# 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
@pytest.mark.parametrize("mutated_output", [" \n", "-- governance comment\n"])
def test_execute_sql_with_cursor_empty_after_mutation(
app_context: None,
mutated_output: str,
) -> None:
"""A mutator that strips a statement to nothing executable (whitespace or
comments only) raises a clean error."""
from superset.errors import SupersetErrorType
from superset.exceptions import SupersetErrorException
from superset.sql.execution.executor import execute_sql_with_cursor
mock_database = MagicMock()
mock_database.db_engine_spec.engine = "postgresql"
mock_database.mutate_sql_based_on_config = lambda sql, **kw: mutated_output
mock_cursor = MagicMock()
mock_query = MagicMock()
with pytest.raises(SupersetErrorException) as excinfo:
execute_sql_with_cursor(
database=mock_database,
cursor=mock_cursor,
statements=["SELECT 1"],
query=mock_query,
)
assert excinfo.value.error.error_type == SupersetErrorType.INVALID_SQL_ERROR
mock_database.db_engine_spec.execute.assert_not_called()
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"
)
# Schema is resolved through the query-aware ``get_default_schema_for_query``
# (so per-query engine gates run), not the static ``get_default_schema``.
get_default_schema_mock = mocker.patch.object(
database, "get_default_schema_for_query", 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()
def test_resolve_query_schema_uses_query_aware_resolution(
mocker: MockerFixture, database: Database, app_context: None
) -> None:
"""``_resolve_query_schema`` resolves through the query-aware
``get_default_schema_for_query`` (which runs per-query engine security gates),
handing it a transient probe Query that carries the request's SQL, schema,
catalog and template params."""
from superset.models.sql_lab import Query
from superset.sql.execution.executor import SQLExecutor
resolve_mock = mocker.patch.object(
database, "get_default_schema_for_query", return_value="resolved_schema"
)
executor = SQLExecutor(database)
options = QueryOptions(schema="explicit", template_params={"p": 1})
result = executor._resolve_query_schema("SELECT 1", options, "cat")
assert result == "resolved_schema"
probe, template_params = resolve_mock.call_args.args
assert isinstance(probe, Query)
assert probe.sql == "SELECT 1"
assert probe.schema == "explicit"
assert probe.catalog == "cat"
assert template_params == {"p": 1}
def test_resolve_query_schema_omits_blank_schema(
mocker: MockerFixture, database: Database, app_context: None
) -> None:
"""An unset request schema reaches the probe as ``None`` so the engine spec
resolves the runtime default instead of matching on an empty string."""
resolve_mock = mocker.patch.object(
database, "get_default_schema_for_query", return_value="public"
)
from superset.sql.execution.executor import SQLExecutor
executor = SQLExecutor(database)
result = executor._resolve_query_schema("SELECT 1", QueryOptions(), None)
assert result == "public"
probe = resolve_mock.call_args.args[0]
assert probe.schema is None
assert probe.catalog is None
def test_prepare_sql_runs_schema_gate_with_explicit_schema(
mocker: MockerFixture, database: Database, app_context: None
) -> None:
"""The per-query schema gate must run even when an explicit schema is
supplied, so an explicit-schema request cannot smuggle a ``SET search_path``
past the gate the resolver enforces (parity with the estimate path, which
resolves unconditionally)."""
from superset.errors import ErrorLevel, SupersetError, SupersetErrorType
from superset.exceptions import SupersetSecurityException
from superset.sql.execution.executor import SQLExecutor
gate = mocker.patch.object(
database,
"get_default_schema_for_query",
side_effect=SupersetSecurityException(
SupersetError(
message="blocked",
error_type=SupersetErrorType.QUERY_SECURITY_ACCESS_ERROR,
level=ErrorLevel.ERROR,
)
),
)
executor = SQLExecutor(database)
with pytest.raises(SupersetSecurityException):
executor._prepare_sql(
"SET search_path = secret; SELECT 1",
QueryOptions(schema="explicit"),
)
gate.assert_called_once()
# =============================================================================
# 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 query_result.error_message is not None
assert "not found" in query_result.error_message.lower()
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 query_result.error_message is not None
assert "Error loading results" in query_result.error_message
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 query_result.error_message is not None
assert "Results not available" in query_result.error_message
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.queries.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.queries.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.queries.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 query_result.error_message is not None
assert "Results not available" in query_result.error_message
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 query_result.error_message is not None
assert "Results not available" in query_result.error_message
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.queries.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
def test_build_statement_blocks_skips_validation_for_unparseable_mutated_sql(
mocker: MockerFixture, mock_database: MagicMock, app_context: None
) -> None:
"""
A mutator may emit engine-specific SQL the parser can't handle. The
empty-statement validation on the joined block is skipped in that case
instead of blocking execution, leaving the database as the authority
on validity.
"""
from superset.exceptions import SupersetParseError
from superset.sql.execution.executor import build_statement_blocks
from superset.sql.parse import SQLScript
mocker.patch.dict(current_app.config, {"MUTATE_AFTER_SPLIT": True})
mock_database.db_engine_spec.run_multiple_statements_as_one = True
mocker.patch.object(
mock_database,
"mutate_sql_based_on_config",
side_effect=lambda sql, **kw: f"ENGINE SPECIFIC {sql}",
)
parsed_script = SQLScript("SELECT 1; SELECT 2;", engine="bigquery")
mocker.patch(
"superset.sql.execution.executor.SQLScript",
side_effect=SupersetParseError("ENGINE SPECIFIC SQL"),
)
_, blocks = build_statement_blocks(
parsed_script, mock_database.db_engine_spec, mock_database
)
assert len(blocks) == 1
assert blocks[0].count("ENGINE SPECIFIC") == 2
# =============================================================================
# Cache Key Identity Tests
# =============================================================================
def test_generate_cache_key_ignores_user_without_impersonation(
mocker: MockerFixture, database: Database, app_context: None
) -> None:
"""
When the connection does not carry per-user identity, the cache key
must not depend on the calling user -- unrelated users legitimately
share cache entries for identical queries.
"""
from superset.sql.execution.executor import SQLExecutor
executor = SQLExecutor(database)
options = QueryOptions()
mocker.patch("superset.sql.execution.executor.utils.get_user_id", return_value=1)
key_user_1 = executor._generate_cache_key("SELECT * FROM salaries", options)
mocker.patch("superset.sql.execution.executor.utils.get_user_id", return_value=2)
key_user_2 = executor._generate_cache_key("SELECT * FROM salaries", options)
assert key_user_1 is not None
assert key_user_1 == key_user_2
def test_generate_cache_key_scopes_by_user_when_impersonation_enabled(
mocker: MockerFixture, database: Database, app_context: None
) -> None:
"""
Regression: when the database impersonates the connecting user (or
uses per-user OAuth2), the effective identity the database sees
differs per user even for byte-identical SQL text. The cache key must
incorporate that identity so one user's cached rows can never be
served to another user the database itself would have denied.
"""
from superset.sql.execution.executor import SQLExecutor
database.impersonate_user = True
executor = SQLExecutor(database)
options = QueryOptions()
mocker.patch("superset.sql.execution.executor.utils.get_user_id", return_value=1)
key_analyst_a = executor._generate_cache_key("SELECT * FROM salaries", options)
mocker.patch("superset.sql.execution.executor.utils.get_user_id", return_value=2)
key_analyst_b = executor._generate_cache_key("SELECT * FROM salaries", options)
assert key_analyst_a is not None
assert key_analyst_b is not None
assert key_analyst_a != key_analyst_b
def test_generate_cache_key_includes_impersonation_key_when_present(
mocker: MockerFixture, database: Database, app_context: None
) -> None:
"""
The executor's result cache mirrors the chart-data cache-key path:
when CACHE_IMPERSONATION / CACHE_QUERY_BY_USER / per_user_caching
resolves an impersonation key for the connection, that key must be
folded into the generated cache key so it can't collide with a key
generated for a different impersonated identity.
"""
from superset.sql.execution.executor import SQLExecutor
executor = SQLExecutor(database)
options = QueryOptions()
mocker.patch(
"superset.utils.cache_keys.add_impersonation_cache_key_if_needed",
side_effect=lambda db, cache_dict: cache_dict.__setitem__(
"impersonation_key", "engineer_a"
),
)
key_engineer_a = executor._generate_cache_key("SELECT * FROM salaries", options)
mocker.patch(
"superset.utils.cache_keys.add_impersonation_cache_key_if_needed",
side_effect=lambda db, cache_dict: cache_dict.__setitem__(
"impersonation_key", "engineer_b"
),
)
key_engineer_b = executor._generate_cache_key("SELECT * FROM salaries", options)
mocker.patch(
"superset.utils.cache_keys.add_impersonation_cache_key_if_needed",
side_effect=lambda db, cache_dict: None,
)
key_no_impersonation = executor._generate_cache_key(
"SELECT * FROM salaries", options
)
assert key_engineer_a is not None
assert key_engineer_b is not None
assert key_no_impersonation is not None
assert len({key_engineer_a, key_engineer_b, key_no_impersonation}) == 3
def test_generate_cache_key_none_when_identity_unknown(
mocker: MockerFixture, database: Database, app_context: None
) -> None:
"""
If the database carries per-user identity but the effective user
cannot be determined (e.g. no request context), the query must be
treated as uncacheable rather than risk a shared cache entry.
"""
from superset.sql.execution.executor import SQLExecutor
database.impersonate_user = True
executor = SQLExecutor(database)
options = QueryOptions()
mocker.patch("superset.sql.execution.executor.utils.get_user_id", return_value=None)
assert executor._generate_cache_key("SELECT * FROM salaries", options) is None
def test_get_from_cache_skips_when_identity_unknown(
mocker: MockerFixture, database: Database, app_context: None
) -> None:
"""
``_get_from_cache`` must not read from the cache backend at all when
``_generate_cache_key`` reports the query as uncacheable.
"""
from superset.extensions import cache_manager
from superset.sql.execution.executor import SQLExecutor
database.impersonate_user = True
executor = SQLExecutor(database)
mocker.patch("superset.sql.execution.executor.utils.get_user_id", return_value=None)
mock_cache_get = mocker.patch.object(cache_manager.data_cache, "get")
result = executor._get_from_cache("SELECT * FROM salaries", QueryOptions())
assert result is None
mock_cache_get.assert_not_called()
def test_store_in_cache_skips_when_identity_unknown(
mocker: MockerFixture, database: Database, app_context: None
) -> None:
"""
``_store_in_cache`` must not write to the cache backend at all when
``_generate_cache_key`` reports the query as uncacheable.
"""
from superset_core.queries.types import (
QueryResult as QueryResultType,
StatementResult,
)
from superset.extensions import cache_manager
from superset.sql.execution.executor import SQLExecutor
database.impersonate_user = True
executor = SQLExecutor(database)
mocker.patch("superset.sql.execution.executor.utils.get_user_id", return_value=None)
mock_cache_set = mocker.patch.object(cache_manager.data_cache, "set")
result = QueryResultType(
status=QueryStatus.SUCCESS,
statements=[
StatementResult(
original_sql="SELECT * FROM salaries",
executed_sql="SELECT * FROM salaries",
data=pd.DataFrame({"salary": [1]}),
row_count=1,
)
],
)
executor._store_in_cache(result, "SELECT * FROM salaries", QueryOptions())
mock_cache_set.assert_not_called()