mirror of
https://github.com/apache/superset.git
synced 2026-07-05 06:15:31 +00:00
Compare commits
1 Commits
chore/ci/s
...
adopt/sqll
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
6a2cd30ac4 |
@@ -2005,6 +2005,9 @@ def SQL_QUERY_MUTATOR( # pylint: disable=invalid-name,unused-argument # noqa:
|
|||||||
# An example use case is if data has role based access controls, and you want to apply
|
# An example use case is if data has role based access controls, and you want to apply
|
||||||
# a SET ROLE statement alongside every user query. Changing this variable maintains
|
# a SET ROLE statement alongside every user query. Changing this variable maintains
|
||||||
# functionality for both the SQL_Lab and Charts.
|
# functionality for both the SQL_Lab and Charts.
|
||||||
|
# This applies consistently in SQL Lab: with MUTATE_AFTER_SPLIT = True the mutator runs
|
||||||
|
# on each individual statement, and with MUTATE_AFTER_SPLIT = False it runs once on the
|
||||||
|
# un-split query block.
|
||||||
MUTATE_AFTER_SPLIT = False
|
MUTATE_AFTER_SPLIT = False
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -445,6 +445,9 @@ def _execute_sql_statements(
|
|||||||
log_query_fn=_make_log_query_fn(database),
|
log_query_fn=_make_log_query_fn(database),
|
||||||
check_stopped_fn=_make_check_stopped_fn(query),
|
check_stopped_fn=_make_check_stopped_fn(query),
|
||||||
execute_fn=_make_execute_fn(query, db_engine_spec),
|
execute_fn=_make_execute_fn(query, db_engine_spec),
|
||||||
|
# `blocks` is a single un-split block when the engine runs multiple
|
||||||
|
# statements as one; otherwise each block is an individual statement.
|
||||||
|
is_split=not db_engine_spec.run_multiple_statements_as_one,
|
||||||
)
|
)
|
||||||
except SoftTimeLimitExceeded as ex:
|
except SoftTimeLimitExceeded as ex:
|
||||||
query.status = QueryStatus.TIMED_OUT
|
query.status = QueryStatus.TIMED_OUT
|
||||||
|
|||||||
@@ -101,6 +101,7 @@ def execute_sql_with_cursor(
|
|||||||
log_query_fn: Any | None = None,
|
log_query_fn: Any | None = None,
|
||||||
check_stopped_fn: Any | None = None,
|
check_stopped_fn: Any | None = None,
|
||||||
execute_fn: Any | None = None,
|
execute_fn: Any | None = None,
|
||||||
|
is_split: bool = True,
|
||||||
) -> list[tuple[str, SupersetResultSet | None, float, int]]:
|
) -> list[tuple[str, SupersetResultSet | None, float, int]]:
|
||||||
"""
|
"""
|
||||||
Execute SQL statements with a cursor and return all result sets.
|
Execute SQL statements with a cursor and return all result sets.
|
||||||
@@ -119,6 +120,10 @@ def execute_sql_with_cursor(
|
|||||||
:param execute_fn: Optional custom execute function. If not provided, uses
|
:param execute_fn: Optional custom execute function. If not provided, uses
|
||||||
database.db_engine_spec.execute(cursor, sql, database). Custom function
|
database.db_engine_spec.execute(cursor, sql, database). Custom function
|
||||||
should accept (cursor, sql) and handle execution.
|
should accept (cursor, sql) and handle execution.
|
||||||
|
:param is_split: Whether `statements` are individual split-out statements (True)
|
||||||
|
or a single un-split block (False, e.g. when the engine spec runs multiple
|
||||||
|
statements as one). Passed to the SQL mutator so `MUTATE_AFTER_SPLIT` can
|
||||||
|
decide whether to fire.
|
||||||
:returns: List of (statement_sql, result_set, execution_time_ms, rowcount) tuples
|
:returns: List of (statement_sql, result_set, execution_time_ms, rowcount) tuples
|
||||||
Returns empty list if stopped. Raises exception on error (fail-fast).
|
Returns empty list if stopped. Raises exception on error (fail-fast).
|
||||||
"""
|
"""
|
||||||
@@ -140,7 +145,7 @@ def execute_sql_with_cursor(
|
|||||||
# Apply SQL mutation
|
# Apply SQL mutation
|
||||||
stmt_sql = database.mutate_sql_based_on_config(
|
stmt_sql = database.mutate_sql_based_on_config(
|
||||||
statement,
|
statement,
|
||||||
is_split=True,
|
is_split=is_split,
|
||||||
)
|
)
|
||||||
|
|
||||||
# Log query
|
# Log query
|
||||||
|
|||||||
@@ -474,7 +474,7 @@ def execute_sql_statements( # noqa: C901
|
|||||||
for statement in parsed_script.statements:
|
for statement in parsed_script.statements:
|
||||||
apply_limit(query, statement)
|
apply_limit(query, statement)
|
||||||
|
|
||||||
# some databases (like BigQuery and Kusto) do not persist state across mmultiple
|
# 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
|
# statements if they're run separately (especially when using `NullPool`), so we run
|
||||||
# the query as a single block.
|
# the query as a single block.
|
||||||
if db_engine_spec.run_multiple_statements_as_one:
|
if db_engine_spec.run_multiple_statements_as_one:
|
||||||
@@ -517,8 +517,15 @@ def execute_sql_statements( # noqa: C901
|
|||||||
query.set_extra_json_key("progress", msg)
|
query.set_extra_json_key("progress", msg)
|
||||||
db.session.commit()
|
db.session.commit()
|
||||||
|
|
||||||
# Hook to allow environment-specific mutation (usually comments) to the SQL
|
# Hook to allow environment-specific mutation (usually comments) to the SQL.
|
||||||
query.executed_sql = database.mutate_sql_based_on_config(block)
|
# `is_split` reflects whether this block is an individual statement: when
|
||||||
|
# the engine runs everything as one block the SQL is not split, otherwise
|
||||||
|
# each block is a single split-out statement. This lets `MUTATE_AFTER_SPLIT`
|
||||||
|
# decide correctly whether the mutator fires here.
|
||||||
|
query.executed_sql = database.mutate_sql_based_on_config(
|
||||||
|
block,
|
||||||
|
is_split=not db_engine_spec.run_multiple_statements_as_one,
|
||||||
|
)
|
||||||
|
|
||||||
try:
|
try:
|
||||||
result_set = execute_query(query, cursor, log_params)
|
result_set = execute_query(query, cursor, log_params)
|
||||||
|
|||||||
@@ -262,6 +262,44 @@ def test_table_column_database() -> None:
|
|||||||
assert TableColumn(database=database).database is database
|
assert TableColumn(database=database).database is database
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.parametrize(
|
||||||
|
"is_split,mutate_after_split,expect_mutated",
|
||||||
|
[
|
||||||
|
# A split-out statement is mutated only when the mutator is meant to run
|
||||||
|
# after the split, and an un-split block only when it runs before.
|
||||||
|
(True, True, True),
|
||||||
|
(True, False, False),
|
||||||
|
(False, False, True),
|
||||||
|
(False, True, False),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
def test_mutate_sql_based_on_config_respects_is_split(
|
||||||
|
app_context: None,
|
||||||
|
mocker: MockerFixture,
|
||||||
|
is_split: bool,
|
||||||
|
mutate_after_split: bool,
|
||||||
|
expect_mutated: bool,
|
||||||
|
) -> None:
|
||||||
|
"""
|
||||||
|
`mutate_sql_based_on_config` fires `SQL_QUERY_MUTATOR` only when the call
|
||||||
|
site's `is_split` matches the `MUTATE_AFTER_SPLIT` config. Regression guard
|
||||||
|
for issue #30169, where SQL Lab always passed the default `is_split=False`
|
||||||
|
and so never mutated when `MUTATE_AFTER_SPLIT=True`.
|
||||||
|
"""
|
||||||
|
database = Database(database_name="db", sqlalchemy_uri="sqlite://")
|
||||||
|
mocker.patch.dict(
|
||||||
|
current_app.config,
|
||||||
|
{
|
||||||
|
"SQL_QUERY_MUTATOR": lambda sql, **kwargs: f"-- mutated\n{sql}",
|
||||||
|
"MUTATE_AFTER_SPLIT": mutate_after_split,
|
||||||
|
},
|
||||||
|
)
|
||||||
|
|
||||||
|
result = database.mutate_sql_based_on_config("SELECT 1", is_split=is_split)
|
||||||
|
|
||||||
|
assert result == ("-- mutated\nSELECT 1" if expect_mutated else "SELECT 1")
|
||||||
|
|
||||||
|
|
||||||
def test_catalog_cache() -> None:
|
def test_catalog_cache() -> None:
|
||||||
"""
|
"""
|
||||||
Test the catalog cache.
|
Test the catalog cache.
|
||||||
|
|||||||
@@ -47,7 +47,7 @@ from superset.models.core import Database
|
|||||||
|
|
||||||
# Note: database, database_with_dml, mock_db_session fixtures and
|
# Note: database, database_with_dml, mock_db_session fixtures and
|
||||||
# mock_query_execution helper are imported from conftest.py
|
# mock_query_execution helper are imported from conftest.py
|
||||||
from .conftest import mock_query_execution
|
from .conftest import create_mock_cursor, mock_query_execution
|
||||||
|
|
||||||
# =============================================================================
|
# =============================================================================
|
||||||
# Basic Execution Tests
|
# Basic Execution Tests
|
||||||
@@ -871,6 +871,45 @@ def test_execute_applies_sql_mutator(
|
|||||||
mutate_mock.assert_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 = 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,)])
|
||||||
|
mocker.patch("superset.result_set.SupersetResultSet", return_value=MagicMock())
|
||||||
|
|
||||||
|
cursor = 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
|
# Progress Tracking Tests
|
||||||
# =============================================================================
|
# =============================================================================
|
||||||
|
|||||||
Reference in New Issue
Block a user