Files
superset2/tests/unit_tests/sql_lab_test.py
T
6c2fef29cb feat(snowflake): Add support for OAuth 2.0 authentication (#36856)
Co-authored-by: Evan Rusackas <evan@preset.io>
Co-authored-by: Evan Rusackas <evan@rusackas.com>
Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: Joe Li <joe@preset.io>
2026-08-24 21:22:14 -07:00

730 lines
26 KiB
Python

# 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.
# pylint: disable=import-outside-toplevel, invalid-name, unused-argument, too-many-locals
import json # noqa: TID251
from typing import Any
from unittest.mock import MagicMock, patch
from urllib.parse import parse_qs, urlparse
from uuid import UUID
import pytest
from freezegun import freeze_time
from pytest_mock import MockerFixture
from sqlalchemy import text
from sqlalchemy.orm import Session
from superset.app import SupersetApp
from superset.common.db_query_status import QueryStatus
from superset.db_engine_specs.postgres import PostgresEngineSpec
from superset.errors import ErrorLevel, SupersetErrorType
from superset.exceptions import OAuth2Error, SupersetErrorException
from superset.models.core import Database
from superset.sql.parse import SQLStatement, Table
from superset.sql_lab import (
execute_query,
execute_sql_statements,
get_query,
get_sql_results,
SqlLabException,
)
from superset.utils.rls import apply_rls, get_predicates_for_table
from tests.conftest import with_config
from tests.unit_tests.models.core_test import oauth2_client_info
def test_execute_query(mocker: MockerFixture, app: None) -> None:
"""
Simple test for `execute_sql_statement`.
"""
query = mocker.MagicMock()
query.executed_sql = "SELECT 42 AS answer"
query.limit = 1
database = query.database
database.allow_dml = False
db_engine_spec = database.db_engine_spec
db_engine_spec.fetch_data.return_value = [(42,)]
cursor = mocker.MagicMock()
SupersetResultSet = mocker.patch("superset.sql_lab.SupersetResultSet") # noqa: N806
# Mock db.session.refresh to avoid AttributeError during session refresh
mocker.patch("superset.sql_lab.db.session.refresh", return_value=None)
execute_query(query, cursor=cursor, log_params={})
db_engine_spec.execute_with_cursor.assert_called_with(
cursor,
"SELECT 42 AS answer",
query,
)
SupersetResultSet.assert_called_with([(42,)], cursor.description, db_engine_spec)
def test_get_query_rolls_back_session_before_retrying(
mocker: MockerFixture, app: SupersetApp
) -> None:
"""
A broken transaction (e.g. `PendingRollbackError` following a failed flush)
leaves the session unusable until `session.rollback()` is called, so without
it every `backoff` retry would reuse the same poisoned session and fail
identically. `get_query` must roll back on failure so each retry gets a
clean session and has a real chance to succeed.
"""
# avoid actually sleeping through the `backoff` decorator's retry interval
mocker.patch("backoff._sync.time.sleep")
expected_query = mocker.MagicMock()
mock_one = mocker.patch("superset.sql_lab.db.session.query")
mock_one.return_value.filter_by.return_value.one.side_effect = [
Exception("session is broken"),
expected_query,
]
mock_rollback = mocker.patch("superset.sql_lab.db.session.rollback")
result = get_query(query_id=1)
assert result is expected_query
assert mock_one.return_value.filter_by.return_value.one.call_count == 2
mock_rollback.assert_called_once()
def test_get_query_swallows_rollback_failure(
mocker: MockerFixture, app: SupersetApp
) -> None:
"""
If the session/connection is too broken for `rollback()` itself to succeed,
that failure must not replace the original lookup error: `get_query` still
needs to raise `SqlLabException` so the `backoff` decorator's retry contract
(which only matches on `SqlLabException`) isn't bypassed.
"""
mocker.patch("backoff._sync.time.sleep")
mock_one = mocker.patch("superset.sql_lab.db.session.query")
mock_one.return_value.filter_by.return_value.one.side_effect = Exception(
"session is broken"
)
mocker.patch(
"superset.sql_lab.db.session.rollback",
side_effect=Exception("connection already closed"),
)
with pytest.raises(SqlLabException):
get_query(query_id=1)
@with_config(
{
"SQLLAB_PAYLOAD_MAX_MB": 50,
"DISALLOWED_SQL_FUNCTIONS": {},
"SQLLAB_CTAS_NO_LIMIT": False,
"SQL_MAX_ROW": 100000,
"QUERY_LOGGER": None,
"TROUBLESHOOTING_LINK": None,
"STATS_LOGGER": MagicMock(),
}
)
def test_execute_sql_statement_exceeds_payload_limit(
mocker: MockerFixture, app
) -> None:
"""
Test for `execute_sql_statements` when the result payload size exceeds the limit.
"""
# Mock the query object and database
query = mocker.MagicMock()
query.limit = 1
query.database = mocker.MagicMock()
query.database.cache_timeout = 100
query.status = "RUNNING"
query.select_as_cta = False
query.database.allow_run_async = True
# Mock get_query to return our mocked query object
mocker.patch("superset.sql_lab.get_query", return_value=query)
# Mock sys.getsizeof to simulate a large payload size
mocker.patch("sys.getsizeof", return_value=100000000) # 100 MB
# Mock _serialize_payload
def mock_serialize_payload(payload, use_msgpack):
return "serialized_payload"
mocker.patch(
"superset.sql_lab._serialize_payload", side_effect=mock_serialize_payload
)
# Mock db.session.refresh to avoid AttributeError during session refresh
mocker.patch("superset.sql_lab.db.session.refresh", return_value=None)
# Mock the results backend to avoid "Results backend is not configured" error
mocker.patch("superset.sql_lab.results_backend", return_value=True)
# Test that the exception is raised when the payload exceeds the limit
with pytest.raises(SupersetErrorException):
execute_sql_statements(
query_id=1,
rendered_query="SELECT 42 AS answer",
return_results=True, # Simulate that results are being returned
store_results=True, # Not storing results but returning them
start_time=None,
expand_data=False,
log_params={},
)
@with_config(
{
"SQLLAB_PAYLOAD_MAX_MB": 50,
"DISALLOWED_SQL_FUNCTIONS": {},
"SQLLAB_CTAS_NO_LIMIT": False,
"SQL_MAX_ROW": 100000,
"QUERY_LOGGER": None,
"TROUBLESHOOTING_LINK": None,
"STATS_LOGGER": MagicMock(),
}
)
def test_execute_sql_statement_within_payload_limit(mocker: MockerFixture, app) -> None:
"""
Test for `execute_sql_statements` when the result payload size is within the limit,
and check if the flow executes smoothly without raising any exceptions.
"""
# Mock the query object and database
query = mocker.MagicMock()
query.limit = 1
query.database = mocker.MagicMock()
query.database.cache_timeout = 100
query.status = "RUNNING"
query.select_as_cta = False
query.database.allow_run_async = True
# Mock get_query to return our mocked query object
mocker.patch("superset.sql_lab.get_query", return_value=query)
# Mock sys.getsizeof to simulate a payload size that is within the limit
mocker.patch("sys.getsizeof", return_value=10000000) # 10 MB (within limit)
# Mock _serialize_payload
def mock_serialize_payload(payload, use_msgpack):
return "serialized_payload"
mocker.patch(
"superset.sql_lab._serialize_payload", side_effect=mock_serialize_payload
)
# Mock db.session.refresh to avoid AttributeError during session refresh
mocker.patch("superset.sql_lab.db.session.refresh", return_value=None)
# Mock the results backend to avoid "Results backend is not configured" error
mocker.patch("superset.sql_lab.results_backend", return_value=True)
# Test that no exception is raised and the function executes smoothly
try:
execute_sql_statements(
query_id=1,
rendered_query="SELECT 42 AS answer",
return_results=True, # Simulate that results are being returned
store_results=True, # Not storing results but returning them
start_time=None,
expand_data=False,
log_params={},
)
except SupersetErrorException:
pytest.fail(
"SupersetErrorException should not have been raised for payload within the limit" # noqa: E501
)
def test_execute_sql_statements_mutates_before_split_by_default(
mocker: MockerFixture, app: SupersetApp
) -> None:
"""
With the default `MUTATE_AFTER_SPLIT=False`, `execute_sql_statements` should
mutate the whole, un-split query once before splitting it into individual
statement blocks, for engines that execute statements individually rather
than as one. Regression guard for issue #30169.
"""
query = mocker.MagicMock()
query.limit = 1
query.database = mocker.MagicMock()
query.database.cache_timeout = 100
query.status = "RUNNING"
query.select_as_cta = False
query.database.allow_run_async = True
query.database.db_engine_spec.engine = "sqlite"
query.database.db_engine_spec.run_multiple_statements_as_one = False
query.database.db_engine_spec.allows_sql_comments = True
mutate_mock = mocker.patch.object(
query.database,
"mutate_sql_based_on_config",
side_effect=lambda sql, **kw: sql,
)
mocker.patch("superset.sql_lab.get_query", return_value=query)
mocker.patch("sys.getsizeof", return_value=10000000)
mocker.patch(
"superset.sql_lab._serialize_payload",
side_effect=lambda payload, use_msgpack: "serialized_payload",
)
mocker.patch("superset.sql_lab.db.session.refresh", return_value=None)
mocker.patch("superset.sql_lab.results_backend", return_value=True)
execute_sql_statements(
query_id=1,
rendered_query="SELECT 1; SELECT 2;",
return_results=True,
store_results=True,
start_time=None,
expand_data=False,
log_params={},
)
is_split_values = [
call.kwargs.get("is_split") for call in mutate_mock.call_args_list
]
# The mutator is called once on the whole, un-split query before splitting...
assert is_split_values[0] is False
first_call_sql = mutate_mock.call_args_list[0].args[0]
assert "1" in first_call_sql
assert "2" in first_call_sql
# Both statements are present in a single, un-split call.
assert first_call_sql.count("SELECT") == 2
# ...and once again per already-split statement (a no-op when
# `MUTATE_AFTER_SPLIT=False`, since `is_split=True` won't match the config).
assert all(value is True for value in is_split_values[1:])
assert len(is_split_values) == 3
def test_execute_sql_statements_mutates_per_statement_when_run_as_one(
mocker: MockerFixture, app: SupersetApp
) -> None:
"""
Engines that always run statements as a single block (e.g. BigQuery, Kusto)
never see `is_split=True` in the per-block mutation call further down, so with
`MUTATE_AFTER_SPLIT=True` the mutator must instead be applied to each
statement up front, before they're joined into that single block.
"""
mocker.patch.dict(app.config, {"MUTATE_AFTER_SPLIT": True})
query = mocker.MagicMock()
query.limit = 1
query.database = mocker.MagicMock()
query.database.cache_timeout = 100
query.status = "RUNNING"
query.select_as_cta = False
query.database.allow_run_async = True
query.database.db_engine_spec.engine = "bigquery"
query.database.db_engine_spec.run_multiple_statements_as_one = True
query.database.db_engine_spec.allows_sql_comments = True
mutate_mock = mocker.patch.object(
query.database,
"mutate_sql_based_on_config",
side_effect=lambda sql, **kw: sql,
)
mocker.patch("superset.sql_lab.get_query", return_value=query)
mocker.patch("sys.getsizeof", return_value=10000000)
mocker.patch(
"superset.sql_lab._serialize_payload",
side_effect=lambda payload, use_msgpack: "serialized_payload",
)
mocker.patch("superset.sql_lab.db.session.refresh", return_value=None)
mocker.patch("superset.sql_lab.results_backend", return_value=True)
execute_sql_statements(
query_id=1,
rendered_query="SELECT 1; SELECT 2;",
return_results=True,
store_results=True,
start_time=None,
expand_data=False,
log_params={},
)
is_split_values = [
call.kwargs.get("is_split") for call in mutate_mock.call_args_list
]
# Mutated once per statement before joining into the single block...
assert is_split_values[0] is True
assert is_split_values[1] is True
first_call_sql = mutate_mock.call_args_list[0].args[0]
second_call_sql = mutate_mock.call_args_list[1].args[0]
assert "1" in first_call_sql
assert "2" in second_call_sql
# ...and the later per-block call is a no-op (`is_split=False` never matches
# `MUTATE_AFTER_SPLIT=True`), so the mutator isn't applied a second time.
assert is_split_values[2] is False
assert len(is_split_values) == 3
def test_execute_sql_statements_raises_when_mutator_strips_all_statements(
mocker: MockerFixture, app: SupersetApp
) -> None:
"""
A `SQL_QUERY_MUTATOR` that strips a query down to nothing (e.g. only
comments/whitespace) must raise a clean error instead of silently
producing an empty block list.
"""
query = mocker.MagicMock()
query.limit = 1
query.database = mocker.MagicMock()
query.database.cache_timeout = 100
query.status = "RUNNING"
query.select_as_cta = False
query.database.allow_run_async = True
query.database.db_engine_spec.engine = "sqlite"
query.database.db_engine_spec.run_multiple_statements_as_one = False
query.database.db_engine_spec.allows_sql_comments = True
mocker.patch.object(
query.database,
"mutate_sql_based_on_config",
side_effect=lambda sql, **kw: "-- just a comment",
)
mocker.patch("superset.sql_lab.get_query", return_value=query)
mocker.patch("superset.sql_lab.db.session.refresh", return_value=None)
mocker.patch("superset.sql_lab.results_backend", return_value=True)
with pytest.raises(SupersetErrorException):
execute_sql_statements(
query_id=1,
rendered_query="SELECT 1;",
return_results=True,
store_results=True,
start_time=None,
expand_data=False,
log_params={},
)
def test_execute_sql_statements_raises_when_mutator_strips_single_block(
mocker: MockerFixture, app: SupersetApp
) -> None:
"""
The empty-statement guard must also cover engines that run all statements
as one block: with `MUTATE_AFTER_SPLIT=True` the per-statement mutator
outputs are joined into a single block, and a comment-only/empty result
must raise a clean error instead of reaching execution as an empty block.
"""
mocker.patch.dict(app.config, {"MUTATE_AFTER_SPLIT": True})
query = mocker.MagicMock()
query.limit = 1
query.database = mocker.MagicMock()
query.database.cache_timeout = 100
query.status = "RUNNING"
query.select_as_cta = False
query.database.allow_run_async = True
query.database.db_engine_spec.engine = "bigquery"
query.database.db_engine_spec.run_multiple_statements_as_one = True
query.database.db_engine_spec.allows_sql_comments = True
mocker.patch.object(
query.database,
"mutate_sql_based_on_config",
side_effect=lambda sql, **kw: "-- just a comment",
)
mocker.patch("superset.sql_lab.get_query", return_value=query)
mocker.patch("superset.sql_lab.db.session.refresh", return_value=None)
mocker.patch("superset.sql_lab.results_backend", return_value=True)
with pytest.raises(SupersetErrorException):
execute_sql_statements(
query_id=1,
rendered_query="SELECT 1; SELECT 2;",
return_results=True,
store_results=True,
start_time=None,
expand_data=False,
log_params={},
)
@freeze_time("2021-04-01T00:00:00Z")
def test_get_sql_results_oauth2(mocker: MockerFixture, app) -> None:
"""
Test that `get_sql_results` works with OAuth2.
"""
# Pushed/popped manually (rather than via a ``with`` block) so the
# ``finally`` below still pops it if an assertion fails, preventing the
# request context from leaking into later tests in the same session.
app_context = app.test_request_context()
app_context.push()
try:
mocker.patch(
"superset.db_engine_specs.base.uuid4",
return_value=UUID("fb11f528-6eba-4a8a-837e-6b0d39ee9187"),
)
mocker.patch(
"superset.db_engine_specs.base.generate_code_verifier",
return_value="xkBPVZoFChVcy3VZ2l5u7d0FZPTU-olO7HtsAOok2IUGigyoZ62tG_oldy2xg9_HdqPKrWUmKZLmU-CUqz_SQ",
)
mocker.patch("superset.daos.key_value.KeyValueDAO.delete_expired_entries")
mocker.patch("superset.daos.key_value.KeyValueDAO.create_entry")
mocker.patch("superset.db_engine_specs.base.db.session.commit")
g = mocker.patch("superset.db_engine_specs.base.g")
g.user = mocker.MagicMock()
g.user.id = 42
database = Database(
id=1,
database_name="my_db",
sqlalchemy_uri="sqlite://",
encrypted_extra=json.dumps(oauth2_client_info),
)
database.db_engine_spec.oauth2_exception = OAuth2Error
get_sqla_engine = mocker.patch.object(database, "get_sqla_engine")
get_sqla_engine().__enter__().raw_connection.side_effect = OAuth2Error(
"OAuth2 required"
)
# `limit` and `select_as_cta_used` must match the real `Query` model's
# defaults (nullable Integer -> None, Boolean default=False) so that
# `apply_limit` -- called unconditionally before the mocked OAuth2 error
# is ever reached -- doesn't try to compare an unconfigured MagicMock
# against an int.
query = mocker.MagicMock(
select_as_cta=False,
select_as_cta_used=False,
limit=None,
database=database,
)
mocker.patch("superset.sql_lab.get_query", return_value=query)
payload = get_sql_results(query_id=1, rendered_query="SELECT 1")
assert payload["status"] == QueryStatus.FAILED
assert payload["error"] == "You don't have permission to access the data."
assert len(payload["errors"]) == 1
error = payload["errors"][0]
assert error["message"] == "You don't have permission to access the data."
assert error["error_type"] == SupersetErrorType.OAUTH2_REDIRECT
assert error["level"] == ErrorLevel.WARNING
assert error["extra"]["tab_id"] == "fb11f528-6eba-4a8a-837e-6b0d39ee9187"
assert (
error["extra"]["redirect_uri"]
== "http://example.com/api/v1/database/oauth2/"
)
# Parse the OAuth2 authorization URL and verify components individually,
# since the JWT state and PKCE code_challenge are computed deterministically
# from mocked inputs but their exact encoding depends on library internals.
url = urlparse(error["extra"]["url"])
assert url.scheme == "https"
assert url.netloc == "abcd1234.snowflakecomputing.com"
assert url.path == "/oauth/authorize"
params = parse_qs(url.query)
assert params["scope"] == ["refresh_token session:role:USERADMIN"]
assert params["response_type"] == ["code"]
assert params["redirect_uri"] == ["http://example.com/api/v1/database/oauth2/"]
assert params["client_id"] == ["my_client_id"]
assert params["code_challenge_method"] == ["S256"]
# Verify PKCE code_challenge matches the mocked code_verifier
from superset.utils.oauth2 import generate_code_challenge
expected_code_challenge = generate_code_challenge(
"xkBPVZoFChVcy3VZ2l5u7d0FZPTU-olO7HtsAOok2IUGigyoZ62tG_oldy2xg9_HdqPKrWUmKZLmU-CUqz_SQ"
)
assert params["code_challenge"] == [expected_code_challenge]
finally:
app_context.pop()
def test_apply_rls(mocker: MockerFixture) -> None:
"""
Test the ``apply_rls`` helper function.
"""
database = mocker.MagicMock()
database.get_default_schema_for_query.return_value = "public"
database.get_default_catalog.return_value = "examples"
database.db_engine_spec = PostgresEngineSpec
get_predicates_for_table = mocker.patch(
"superset.utils.rls.get_predicates_for_table",
side_effect=[["c1 = 1"], ["c2 = 2"]],
)
parsed_statement = SQLStatement("SELECT * FROM t1, t2", "postgresql")
parsed_statement.tables = sorted(parsed_statement.tables, key=lambda x: x.table) # type: ignore
apply_rls(database, "examples", "public", parsed_statement)
get_predicates_for_table.assert_has_calls(
[
mocker.call(
Table("t1", "public", "examples"),
database,
"examples",
exclude_dataset_id=None,
),
mocker.call(
Table("t2", "public", "examples"),
database,
"examples",
exclude_dataset_id=None,
),
]
)
assert (
parsed_statement.format()
== """
SELECT
*
FROM (
SELECT
*
FROM t1
WHERE
c1 = 1
) AS "t1", (
SELECT
*
FROM t2
WHERE
c2 = 2
) AS "t2"
""".strip()
)
def test_get_predicates_for_table(mocker: MockerFixture) -> None:
"""
Test the ``get_predicates_for_table`` helper function.
"""
database = mocker.MagicMock()
dataset = mocker.MagicMock()
predicate = mocker.MagicMock()
predicate.compile.return_value = "c1 = 1"
dataset.get_sqla_row_level_filters.return_value = [predicate]
db = mocker.patch("superset.utils.rls.db")
db.session.query().filter().one_or_none.return_value = dataset
table = Table("t1", "public", "examples")
assert get_predicates_for_table(table, database, "examples") == ["c1 = 1"]
dataset.get_sqla_row_level_filters.assert_called_once_with(
include_global_guest_rls=False
)
def test_get_predicates_for_table_null_schema_dataset(session: Session) -> None:
"""
A dataset stored with a NULL schema is scoped to the database's default
schema, mirroring the existing null-catalog fallback.
A query resolving to that default schema must find the dataset, so its RLS
predicates are applied instead of being silently dropped. A query against a
different schema must not, since the null-schema dataset doesn't describe it.
"""
from superset.connectors.sqla.models import SqlaTable
SqlaTable.metadata.create_all(session.get_bind())
database = Database(database_name="rls_db", sqlalchemy_uri="sqlite://")
# registered without an explicit schema, e.g. via the dataset API
dataset = SqlaTable(table_name="t1", schema=None, catalog=None, database=database)
session.add_all([database, dataset])
session.flush()
with (
patch.object(
SqlaTable, "get_sqla_row_level_filters", return_value=[text("c1 = 1")]
),
patch.object(Database, "get_default_schema", return_value="public"),
):
assert get_predicates_for_table(
Table("t1", "public", None), database, None
) == ["c1 = 1"]
assert (
get_predicates_for_table(Table("t1", "sales", None), database, None) == []
)
def test_get_predicates_for_table_prefers_exact_schema_match(session: Session) -> None:
"""
A dataset stored without a schema and one stored with the default schema can
coexist for the same table. The exact match must win, and the lookup must stay
unambiguous rather than treating both rows as candidates for a single dataset.
"""
from superset.connectors.sqla.models import SqlaTable
SqlaTable.metadata.create_all(session.get_bind())
database = Database(database_name="rls_db_exact", sqlalchemy_uri="sqlite://")
session.add_all(
[
database,
SqlaTable(table_name="t1", schema=None, catalog=None, database=database),
SqlaTable(
table_name="t1", schema="public", catalog=None, database=database
),
]
)
session.flush()
def row_level_filters(
self: Any, include_global_guest_rls: bool = True
) -> list[Any]:
return [text(f"c1 = '{self.schema}'")]
with (
patch.object(
SqlaTable,
"get_sqla_row_level_filters",
autospec=True,
side_effect=row_level_filters,
),
patch.object(Database, "get_default_schema", return_value="public"),
):
assert get_predicates_for_table(
Table("t1", "public", None), database, None
) == ["c1 = 'public'"]
def test_get_predicates_for_table_excludes_self(mocker: MockerFixture) -> None:
"""
When ``exclude_dataset_id`` is supplied, the lookup query must add an
``id != exclude_dataset_id`` filter so a virtual dataset whose
``table_name`` matches a table referenced inside its own SQL doesn't get
its own RLS injected into the inner SQL (would double-apply on top of the
outer WHERE). Regression test for the physical→virtual conversion bug.
"""
database = mocker.MagicMock()
db = mocker.patch("superset.utils.rls.db")
db.session.query().filter().one_or_none.return_value = None
table = Table("orders", "public", "examples")
assert (
get_predicates_for_table(table, database, "examples", exclude_dataset_id=42)
== []
)
# The filter call should have received four base filters plus the exclusion
# filter, i.e. five total positional args inside and_().
filter_call = db.session.query().filter.call_args
and_clause = filter_call.args[0]
assert len(and_clause.clauses) == 5