diff --git a/superset/extensions/metadb.py b/superset/extensions/metadb.py index a1a26b488fd..ebd256f6b63 100644 --- a/superset/extensions/metadb.py +++ b/superset/extensions/metadb.py @@ -37,6 +37,7 @@ joins and unions are done in memory, using the SQLite engine. from __future__ import annotations +import contextvars import datetime import decimal import operator @@ -68,7 +69,34 @@ from sqlalchemy.exc import NoSuchTableError from sqlalchemy.sql import Select, select from superset import db, feature_flag_manager, security_manager -from superset.sql.parse import Table +from superset.sql.parse import count_referenced_tables, Table + + +def _count_referenced_tables(statement: str) -> int: + """ + Count the distinct `superset://` virtual tables a statement references, + so ``get_data`` can tell whether it's being asked for a standalone table + or for one side of a multi-table statement (see ``get_data`` for why this + matters). Shillelagh calls `SupersetShillelaghAdapter.get_data` once per + underlying table, independently of any other table referenced by the + same statement, so it has no way on its own to tell the two cases apart. + + Uses the real SQL parser rather than pattern-matching on quoted + identifiers, since a naive `"db.table"`-shaped regex also matches + dotted, double-quoted column aliases (e.g. `AS "metric.value"`) that + have nothing to do with table references, and would misclassify a + single-table statement as multi-table. + """ + return count_referenced_tables(statement, "sqlite") + + +# `SupersetAPSWDialect.on_connect` populates `_executing_multi_table_query` for +# the duration of a statement so that `get_data` can tell whether it's being +# asked for a standalone table or for one side of a multi-table query (see +# `get_data` for why this matters). +_executing_multi_table_query: contextvars.ContextVar[bool] = contextvars.ContextVar( + "_executing_multi_table_query", default=False +) # pylint: disable=abstract-method @@ -119,6 +147,49 @@ class SupersetAPSWDialect(APSWDialect): }, ) + def on_connect(self) -> Callable[[Any], None]: + """ + Wrap cursor creation on every new DBAPI connection so ``execute`` tracks + whether the statement it's about to run references more than one + `superset://` virtual table, no matter how that statement reaches the + cursor. + + SQLAlchemy's `do_execute*` hooks only fire for statements executed + through a SQLAlchemy `Connection` (the ORM/Core path). SQL Lab, the + primary way users query these tables, instead pulls a raw DBAPI cursor + via `engine.raw_connection()` and calls `cursor.execute()` on it + directly, bypassing those hooks entirely -- which would leave + `_executing_multi_table_query` permanently `False` for that path, and + `get_data` back to silently truncating one side of a join (see + `get_data` and #36304). Patching the cursor factory here, at the point + a new physical connection is established, catches every path, since + each one ultimately calls `execute()` on a cursor obtained from this + same connection. + """ + + def setup(dbapi_connection: Any) -> None: + original_cursor = dbapi_connection.cursor + + def cursor(*args: Any, **kwargs: Any) -> Any: + raw_cursor = original_cursor(*args, **kwargs) + original_execute = raw_cursor.execute + + def execute(operation: str, parameters: Any = None) -> Any: + token = _executing_multi_table_query.set( + _count_referenced_tables(operation) > 1 + ) + try: + return original_execute(operation, parameters) + finally: + _executing_multi_table_query.reset(token) + + raw_cursor.execute = execute + return raw_cursor + + dbapi_connection.cursor = cursor + + return setup + F = TypeVar("F", bound=Callable[..., Any]) @@ -409,7 +480,18 @@ class SupersetShillelaghAdapter(Adapter): """ app_limit: int | None = current_app.config["SUPERSET_META_DB_LIMIT"] if limit is None: - limit = app_limit + # Shillelagh calls `get_data` once per table, independently of any + # other table referenced by the same statement, so a value of `None` + # here doesn't necessarily mean this table is the whole query -- it + # can equally mean this table is one side of a join (or other + # multi-table statement). Applying the app-wide default in that case + # would silently truncate this table before the in-memory join runs, + # dropping rows that have a genuine match on the other side with no + # error (see #36304). Only fall back to the default for statements + # that reference a single table, where truncating it can't hide + # otherwise-valid matches. + if app_limit is not None and not _executing_multi_table_query.get(): + limit = app_limit elif app_limit is not None: limit = min(limit, app_limit) diff --git a/superset/sql/parse.py b/superset/sql/parse.py index b1b93892d47..b90cc55c471 100644 --- a/superset/sql/parse.py +++ b/superset/sql/parse.py @@ -22,7 +22,6 @@ import enum import logging import re import urllib.parse -from collections.abc import Iterable from dataclasses import dataclass from typing import Any, Generic, Optional, TYPE_CHECKING, TypeVar @@ -2152,12 +2151,47 @@ class SQLScript: return len(self.statements) == 1 and self.statements[0].is_select() -def extract_tables_from_statement( +def _find_show_statement_tables(statement: exp.Show) -> set[Table]: + """ + Build the table references for a ``SHOW`` statement. + + Structured metadata statements (`SHOW CREATE TABLE foo.bar`, + `SHOW COLUMNS FROM foo`, ...) reference their target via dedicated + args rather than query sources, so build the table references + explicitly. Statements with no extractable target (e.g. + `SHOW TABLES FROM some_schema`) yield an empty set and are treated + as unparseable for authorization purposes (see + `SQLScript.has_unparseable_statement`). + + ``SHOW`` statements reference a single metadata target, never a join, so + (unlike ``_find_table_sources``) there is no distinct occurrence-counting + variant of this helper: the deduplicated set is always the right count. + """ + show_tables = { + Table( + source.name, + source.db if source.db != "" else None, + source.catalog if source.catalog != "" else None, + ) + for source in statement.find_all(exp.Table) + } + if target := statement.args.get("target"): + db = statement.args.get("db") + show_tables.add( + Table( + target.name if isinstance(target, exp.Expression) else str(target), + db.name if isinstance(db, exp.Expression) else db, + ) + ) + return show_tables + + +def _find_table_sources( statement: exp.Expression, dialect: Dialects | None, -) -> set[Table]: +) -> list[exp.Table]: """ - Extract all table references in a single statement. + Find every table reference (occurrence, not deduplicated) in a statement. Please note that this is not trivial; consider the following queries: @@ -2165,60 +2199,45 @@ def extract_tables_from_statement( SHOW PARTITIONS FROM some_table; WITH masked_name AS (SELECT * FROM some_table) SELECT * FROM masked_name; - See the unit tests for other tricky cases. + See the unit tests for other tricky cases. Note that `exp.Show` statements + are not handled here: see `_find_show_statement_tables`. """ - sources: Iterable[exp.Table] - if isinstance(statement, exp.Describe): # A `DESCRIBE` query has no sources in sqlglot, so we need to explicitly # query for all tables. - sources = statement.find_all(exp.Table) - elif isinstance(statement, exp.Command): + return list(statement.find_all(exp.Table)) + if isinstance(statement, exp.Command): # Commands, like `SHOW COLUMNS FROM foo`, have to be converted into a # `SELECT` statetement in order to extract tables. literal = statement.find(exp.Literal) if not literal: - return set() + return [] pseudo_sql = f"SELECT {literal.this}" try: _check_script_length(pseudo_sql, None) pseudo_query = sqlglot.parse_one(pseudo_sql, dialect=dialect) except (ParseError, SupersetParseError): - return set() - sources = pseudo_query.find_all(exp.Table) - elif isinstance(statement, exp.Show): - # Structured metadata statements (`SHOW CREATE TABLE foo.bar`, - # `SHOW COLUMNS FROM foo`, ...) reference their target via dedicated - # args rather than query sources, so build the table references - # explicitly. Statements with no extractable target (e.g. - # `SHOW TABLES FROM some_schema`) yield an empty set and are treated - # as unparseable for authorization purposes (see - # `SQLScript.has_unparseable_statement`). - show_tables = { - Table( - source.name, - source.db if source.db != "" else None, - source.catalog if source.catalog != "" else None, - ) - for source in statement.find_all(exp.Table) - } - if target := statement.args.get("target"): - db = statement.args.get("db") - show_tables.add( - Table( - target.name if isinstance(target, exp.Expression) else str(target), - db.name if isinstance(db, exp.Expression) else db, - ) - ) - return show_tables - else: - sources = [ - source - for scope in traverse_scope(statement) - for source in scope.sources.values() - if isinstance(source, exp.Table) and not is_cte(source, scope) - ] + return [] + return list(pseudo_query.find_all(exp.Table)) + + return [ + source + for scope in traverse_scope(statement) + for source in scope.sources.values() + if isinstance(source, exp.Table) and not is_cte(source, scope) + ] + + +def extract_tables_from_statement( + statement: exp.Expression, + dialect: Dialects | None, +) -> set[Table]: + """ + Extract all distinct table references in a single statement. + """ + if isinstance(statement, exp.Show): + return _find_show_statement_tables(statement) return { Table( @@ -2226,10 +2245,79 @@ def extract_tables_from_statement( source.db if source.db != "" else None, source.catalog if source.catalog != "" else None, ) - for source in sources + for source in _find_table_sources(statement, dialect) } +def count_referenced_tables(statement: str, dialect: Dialects | str | None) -> int: + """ + Count the table references in a raw SQL string. + + This counts occurrences, not distinct tables, so a self-join referencing + the same physical table twice (via two aliases) is still counted as 2 - + callers use this count to decide whether a statement is a join, and a + self-join needs the same treatment as a join across different tables. + A CTE that's referenced more than once (e.g. self-joined) is weighted the + same way: each reference to it counts its own underlying tables again, + since a CTE is inlined at every place it's used (see + ``_count_weighted_table_references``). + + Falls back to a conservative count of 1 (i.e. "not multi-table") if the + statement can't be parsed, since callers gating multi-table-only behavior + on this count should default to treating an unparseable statement as a + single table. + """ + try: + _check_script_length(statement, str(dialect) if dialect else None) + parsed = sqlglot.parse_one(statement, dialect=dialect) + if isinstance(parsed, exp.Show): + return len(_find_show_statement_tables(parsed)) + if isinstance(parsed, (exp.Describe, exp.Command)): + # Neither has join semantics for a per-table row cap to interact + # with, so the plain (unweighted) extraction already used for + # permissioning is fine here too. + return len(_find_table_sources(parsed, dialect)) + return _count_weighted_table_references(parsed) + except Exception: # pylint: disable=broad-except + return 1 + + +def _count_weighted_table_references(statement: exp.Expression) -> int: + """ + Count table references the way callers gating multi-table-only behavior + need: weighting each CTE by how many times it's actually referenced, + not by how many distinct tables its own definition reads. + + ``_find_table_sources`` (used for permissioning) intentionally counts a + CTE's underlying tables exactly once regardless of how many times the + CTE is referenced downstream, since permission checks only care about + the *set* of tables read. But a CTE that wraps a single virtual table + and is then self-joined N ways is inlined at each of those N places, so + it triggers N separate reads of that table -- one per join side -- and + must count as N here too. Otherwise a per-table row cap (see + ``SUPERSET_META_DB_LIMIT`` and #36304) looks safe to apply and silently + truncates one side of the self-join away before the join runs. + """ + + def resolve(scope: Scope, seen: frozenset[int]) -> list[exp.Table]: + if id(scope) in seen: + return [] # guards a WITH RECURSIVE self-reference from looping forever + seen = seen | {id(scope)} + tables: list[exp.Table] = [] + for _, source in scope.selected_sources.values(): + if isinstance(source, exp.Table) and not is_cte(source, scope): + tables.append(source) + elif isinstance(source, Scope) and source.scope_type == ScopeType.CTE: + tables.extend(resolve(source, seen)) + return tables + + return sum( + len(resolve(scope, frozenset())) + for scope in traverse_scope(statement) + if scope.scope_type != ScopeType.CTE + ) + + def is_cte(source: exp.Table, scope: Scope) -> bool: """ Does this reference resolve to a CTE rather than to a real table? diff --git a/tests/unit_tests/extensions/test_sqlalchemy.py b/tests/unit_tests/extensions/test_sqlalchemy.py index 01ce1ced5fb..998534a015b 100644 --- a/tests/unit_tests/extensions/test_sqlalchemy.py +++ b/tests/unit_tests/extensions/test_sqlalchemy.py @@ -22,6 +22,7 @@ from collections.abc import Iterator from typing import TYPE_CHECKING import pytest +from flask import current_app from pytest_mock import MockerFixture from sqlalchemy import text from sqlalchemy.engine import create_engine @@ -239,6 +240,410 @@ def test_superset_joins( assert list(results) == [(10, "ten"), (20, "twenty")] +@pytest.mark.parametrize( + ("statement", "expected"), + [ + # A single table reference is not a multi-table statement... + ('SELECT * FROM "database1.table1"', 1), + # ...even when it has a dotted, double-quoted column alias, which a + # naive `"[^"]*\.[^"]*"`-shaped regex would also match, misidentifying + # a single-table statement as multi-table and silently skipping + # SUPERSET_META_DB_LIMIT for it. + ( + 'SELECT COUNT(id) AS "metric.value" FROM "database1.table1"', + 1, + ), + ( + 'SELECT t1.b, t2.b FROM "database1.table1" AS t1 ' + 'JOIN "database2.table2" AS t2 ON t1.a = t2.a', + 2, + ), + ( + 'SELECT * FROM "database1.table1", "database2.table2" WHERE t1.a = t2.a', + 2, + ), + # Statements the parser can't handle fall back to the safe default + # (treat as single-table, so the app-wide limit still applies). + ("this is not valid sql (((", 1), + ], +) +def test_count_referenced_tables(statement: str, expected: int) -> None: + """ + Regression for a review comment on #42598/#36304: the multi-table + detection used to gate SUPERSET_META_DB_LIMIT must count actual table + references via the real SQL parser, not pattern-match dotted quoted + identifiers, which also matches dotted column aliases. + """ + from superset.extensions.metadb import _count_referenced_tables + + assert _count_referenced_tables(statement) == expected + + +@pytest.fixture +def table1_large(session: Session, database1: "Database") -> Iterator[None]: + with database1.get_sqla_engine() as engine: + with engine.begin() as conn: + conn.execute( + text( + "CREATE TABLE table1_large (a INTEGER NOT NULL PRIMARY KEY, " + "b INTEGER)" + ) + ) + conn.execute( + text("INSERT INTO table1_large (a, b) VALUES (1, 10), (2, 20), (3, 30)") + ) + db.session.commit() + + yield + + with engine.begin() as conn: + conn.execute(text("DROP TABLE table1_large")) + db.session.commit() + + +@pytest.fixture +def table2_late_match(session: Session, database2: "Database") -> Iterator[None]: + with database2.get_sqla_engine() as engine: + with engine.begin() as conn: + conn.execute( + text( + "CREATE TABLE table2_late_match (a INTEGER NOT NULL PRIMARY KEY, " + "b TEXT)" + ) + ) + conn.execute( + text("INSERT INTO table2_late_match (a, b) VALUES (3, 'thirty')") + ) + db.session.commit() + + yield + + with engine.begin() as conn: + conn.execute(text("DROP TABLE table2_late_match")) + db.session.commit() + + +@pytest.fixture +def table2_multi_late_match(session: Session, database2: "Database") -> Iterator[None]: + with database2.get_sqla_engine() as engine: + with engine.begin() as conn: + conn.execute( + text( + "CREATE TABLE table2_multi_late_match " + "(a INTEGER NOT NULL PRIMARY KEY, b TEXT)" + ) + ) + conn.execute( + text( + "INSERT INTO table2_multi_late_match (a, b) " + "VALUES (2, 'twenty'), (3, 'thirty')" + ) + ) + db.session.commit() + + yield + + with engine.begin() as conn: + conn.execute(text("DROP TABLE table2_multi_late_match")) + db.session.commit() + + +@pytest.fixture +def table2_fanout_match(session: Session, database2: "Database") -> Iterator[None]: + with database2.get_sqla_engine() as engine: + with engine.begin() as conn: + conn.execute( + text( + "CREATE TABLE table2_fanout_match " + "(id INTEGER NOT NULL PRIMARY KEY, a INTEGER, b TEXT)" + ) + ) + # `a` is deliberately not unique (unlike table2_late_match, where + # it's the primary key): a single outer row matching on `a=3` + # fans out into two inner rows here, so reading the match + # requires pulling more than one row through the cursor per + # outer probe, instead of a single unique-index lookup. + conn.execute( + text( + "INSERT INTO table2_fanout_match (a, b) " + "VALUES (3, 'thirty-x'), (3, 'thirty-y')" + ) + ) + db.session.commit() + + yield + + with engine.begin() as conn: + conn.execute(text("DROP TABLE table2_fanout_match")) + db.session.commit() + + +@with_feature_flags(ENABLE_SUPERSET_META_DB=True) +def test_superset_joins_with_limit_drops_fanout_matches( + mocker: MockerFixture, + monkeypatch: pytest.MonkeyPatch, + app_context: None, + table1_large: None, + table2_fanout_match: None, +) -> None: + """ + Coverage note from review of #42598: the other join regression tests + match on a primary key on both sides, so each inner lookup returns at + most one row. Here `table1_large`'s single match (a=3) fans out into two + rows in `table2_fanout_match`, so satisfying it means pulling more than + one row through the cursor for a single outer probe, rather than a single + unique-index lookup. + """ + monkeypatch.setitem(current_app.config, "DB_SQLA_URI_VALIDATOR", None) + monkeypatch.setitem(current_app.config, "SUPERSET_META_DB_LIMIT", 2) + monkeypatch.setitem(current_app.config, "DATABASE_OAUTH2_CLIENTS", {}) + monkeypatch.setitem(current_app.config, "SQLALCHEMY_CUSTOM_PASSWORD_STORE", None) + + mocker.patch( + "superset.extensions.metadb.security_manager.raise_for_access", + return_value=None, + ) + + from flask import g + + g.user = mocker.MagicMock() + g.user.is_anonymous = False + + try: + engine = create_engine("superset://", future=True) + except Exception as e: + pytest.skip(f"Superset dialect not available: {e}") + + with engine.connect() as conn: + results = conn.execute( + text(""" + SELECT t1.b, t2.b + FROM "database1.table1_large" AS t1 + JOIN "database2.table2_fanout_match" AS t2 + ON t1.a = t2.a + ORDER BY t2.b + """) + ) + assert list(results) == [(30, "thirty-x"), (30, "thirty-y")] + + +@with_feature_flags(ENABLE_SUPERSET_META_DB=True) +def test_superset_joins_with_limit_multiple_late_matches( + mocker: MockerFixture, + monkeypatch: pytest.MonkeyPatch, + app_context: None, + table1_large: None, + table2_multi_late_match: None, +) -> None: + """ + Diagnostic probe raised in review of #42598: does the per-table + ``SUPERSET_META_DB_LIMIT`` skip (keyed on the + ``_executing_multi_table_query`` ContextVar) hold for every row of a + multi-row join result, or only the first? ``table1_large`` has two + genuine matches in ``table2_multi_late_match`` (a=2 and a=3), both of + which fall past SUPERSET_META_DB_LIMIT=2 in table1_large's own row + order for a naive per-table truncation. + """ + monkeypatch.setitem(current_app.config, "DB_SQLA_URI_VALIDATOR", None) + monkeypatch.setitem(current_app.config, "SUPERSET_META_DB_LIMIT", 2) + monkeypatch.setitem(current_app.config, "DATABASE_OAUTH2_CLIENTS", {}) + monkeypatch.setitem(current_app.config, "SQLALCHEMY_CUSTOM_PASSWORD_STORE", None) + + mocker.patch( + "superset.extensions.metadb.security_manager.raise_for_access", + return_value=None, + ) + + from flask import g + + g.user = mocker.MagicMock() + g.user.is_anonymous = False + + try: + engine = create_engine("superset://", future=True) + except Exception as e: + pytest.skip(f"Superset dialect not available: {e}") + + with engine.connect() as conn: + results = conn.execute( + text(""" + SELECT t1.b, t2.b + FROM "database1.table1_large" AS t1 + JOIN "database2.table2_multi_late_match" AS t2 + ON t1.a = t2.a + ORDER BY t1.b + """) + ) + assert list(results) == [(20, "twenty"), (30, "thirty")] + + +@with_feature_flags(ENABLE_SUPERSET_META_DB=True) +def test_superset_joins_with_limit_drops_matches( + mocker: MockerFixture, + monkeypatch: pytest.MonkeyPatch, + app_context: None, + table1_large: None, + table2_late_match: None, +) -> None: + """ + Regression for #36304: SUPERSET_META_DB_LIMIT is applied to each + underlying table independently, before the in-memory join runs. A row + that has a genuine match on the other side of the join but falls past + the per-table limit is silently dropped from the join result, with no + error or truncation warning. + """ + # Use monkeypatch (rather than the `@with_config` decorator) so the + # config overrides are guaranteed to be undone even though this test is + # expected to fail its assertion until the underlying bug is fixed. + # `@with_config` only restores the original values after the wrapped + # test function returns normally, so an assertion failure here would + # otherwise leak SUPERSET_META_DB_LIMIT=2 into later tests. + monkeypatch.setitem(current_app.config, "DB_SQLA_URI_VALIDATOR", None) + monkeypatch.setitem(current_app.config, "SUPERSET_META_DB_LIMIT", 2) + monkeypatch.setitem(current_app.config, "DATABASE_OAUTH2_CLIENTS", {}) + monkeypatch.setitem(current_app.config, "SQLALCHEMY_CUSTOM_PASSWORD_STORE", None) + + mocker.patch( + "superset.extensions.metadb.security_manager.raise_for_access", + return_value=None, + ) + + from flask import g + + g.user = mocker.MagicMock() + g.user.is_anonymous = False + + try: + engine = create_engine("superset://", future=True) + except Exception as e: + # Skip test if superset:// dialect can't be loaded (common in Docker) + pytest.skip(f"Superset dialect not available: {e}") + + with engine.connect() as conn: + results = conn.execute( + text(""" + SELECT t1.b, t2.b + FROM "database1.table1_large" AS t1 + JOIN "database2.table2_late_match" AS t2 + ON t1.a = t2.a + """) + ) + # table2_late_match's only row (a=3) has a genuine match in + # table1_large (a=3, b=30), but SUPERSET_META_DB_LIMIT=2 truncates + # table1_large to its first two rows (a=1, a=2) before the join + # runs, so the join comes back empty instead of finding the match. + assert list(results) == [(30, "thirty")] + + +@with_feature_flags(ENABLE_SUPERSET_META_DB=True) +def test_superset_comma_join_with_limit_drops_matches( + mocker: MockerFixture, + monkeypatch: pytest.MonkeyPatch, + app_context: None, + table1_large: None, + table2_late_match: None, +) -> None: + """ + Regression for #36304: an implicit comma join (``FROM a, b WHERE ...``) + references two tables just like an explicit ``JOIN``, but doesn't contain + the literal `JOIN` keyword. Multi-table detection has to catch this shape + too, or the per-table limit still gets applied and silently drops matches. + """ + # See test_superset_joins_with_limit_drops_matches for why monkeypatch is + # used here instead of `@with_config`. + monkeypatch.setitem(current_app.config, "DB_SQLA_URI_VALIDATOR", None) + monkeypatch.setitem(current_app.config, "SUPERSET_META_DB_LIMIT", 2) + monkeypatch.setitem(current_app.config, "DATABASE_OAUTH2_CLIENTS", {}) + monkeypatch.setitem(current_app.config, "SQLALCHEMY_CUSTOM_PASSWORD_STORE", None) + + mocker.patch( + "superset.extensions.metadb.security_manager.raise_for_access", + return_value=None, + ) + + from flask import g + + g.user = mocker.MagicMock() + g.user.is_anonymous = False + + try: + engine = create_engine("superset://", future=True) + except Exception as e: + # Skip test if superset:// dialect can't be loaded (common in Docker) + pytest.skip(f"Superset dialect not available: {e}") + + with engine.connect() as conn: + results = conn.execute( + text(""" + SELECT t1.b, t2.b + FROM "database1.table1_large" AS t1, "database2.table2_late_match" AS t2 + WHERE t1.a = t2.a + """) + ) + # Same scenario as test_superset_joins_with_limit_drops_matches, but + # using a comma join instead of the `JOIN` keyword. + assert list(results) == [(30, "thirty")] + + +@with_feature_flags(ENABLE_SUPERSET_META_DB=True) +def test_superset_joins_via_raw_cursor_drops_matches( + mocker: MockerFixture, + monkeypatch: pytest.MonkeyPatch, + app_context: None, + table1_large: None, + table2_late_match: None, +) -> None: + """ + Regression for #36304: SQL Lab executes statements through a raw DBAPI + cursor (``engine.raw_connection().cursor()``), not through SQLAlchemy's + ``Connection.execute()``. That path never reaches + ``SupersetAPSWDialect.do_execute*``, so a fix keyed only on those hooks + leaves the per-table ``SUPERSET_META_DB_LIMIT`` skip blind to exactly the + statements SQL Lab runs, and the same join match SQL Lab users see would + still be silently dropped even though + ``test_superset_joins_with_limit_drops_matches`` (which goes through + ``Connection.execute()``) passes. + """ + monkeypatch.setitem(current_app.config, "DB_SQLA_URI_VALIDATOR", None) + monkeypatch.setitem(current_app.config, "SUPERSET_META_DB_LIMIT", 2) + monkeypatch.setitem(current_app.config, "DATABASE_OAUTH2_CLIENTS", {}) + monkeypatch.setitem(current_app.config, "SQLALCHEMY_CUSTOM_PASSWORD_STORE", None) + + mocker.patch( + "superset.extensions.metadb.security_manager.raise_for_access", + return_value=None, + ) + + from flask import g + + g.user = mocker.MagicMock() + g.user.is_anonymous = False + + try: + engine = create_engine("superset://", future=True) + except Exception as e: + # Skip test if superset:// dialect can't be loaded (common in Docker) + pytest.skip(f"Superset dialect not available: {e}") + + raw_connection = engine.raw_connection() + try: + cursor = raw_connection.cursor() + cursor.execute( + """ + SELECT t1.b, t2.b + FROM "database1.table1_large" AS t1 + JOIN "database2.table2_late_match" AS t2 + ON t1.a = t2.a + """ + ) + # Same scenario as test_superset_joins_with_limit_drops_matches, but + # executed the way SQL Lab actually runs queries: a raw DBAPI cursor + # obtained from `engine.raw_connection()`, bypassing `do_execute*`. + assert list(cursor) == [(30, "thirty")] + finally: + raw_connection.close() + + @with_feature_flags(ENABLE_SUPERSET_META_DB=True) def test_dml( mocker: MockerFixture, diff --git a/tests/unit_tests/sql/parse_tests.py b/tests/unit_tests/sql/parse_tests.py index 3a8370ace7b..9749ef12c9b 100644 --- a/tests/unit_tests/sql/parse_tests.py +++ b/tests/unit_tests/sql/parse_tests.py @@ -28,7 +28,9 @@ from superset.exceptions import QueryClauseValidationException, SupersetParseErr from superset.jinja_context import JinjaTemplateProcessor from superset.sql.parse import ( _check_script_length, + _count_weighted_table_references, BaseSQLStatement, + count_referenced_tables, CTASMethod, extract_tables_from_statement, has_aggregate, @@ -234,6 +236,170 @@ def test_extract_tables_from_sql() -> None: ) == {Table("forbidden_table")} +def test_count_referenced_tables() -> None: + """ + Test that ``count_referenced_tables`` counts table reference occurrences + (not distinct tables), ignoring dotted quoted aliases, and falls back to + 1 for unparseable SQL. + """ + assert count_referenced_tables('SELECT * FROM "db.table1"', Dialects.SQLITE) == 1 + assert ( + count_referenced_tables( + 'SELECT COUNT(id) AS "metric.value" FROM "db.table1"', Dialects.SQLITE + ) + == 1 + ) + assert ( + count_referenced_tables( + 'SELECT t1.b, t2.b FROM "db.table1" AS t1 ' + 'JOIN "db.table2" AS t2 ON t1.a = t2.a', + Dialects.SQLITE, + ) + == 2 + ) + assert count_referenced_tables("this is not valid sql (((", Dialects.SQLITE) == 1 + assert count_referenced_tables("SHOW CREATE TABLE s1.t1", "mysql") == 1 + + +def test_count_referenced_tables_self_join() -> None: + """ + A self-join references the same physical table twice via two aliases; + it must still count as 2 (a join), not 1 (deduplicated to a single + table), or the caller's multi-table detection would incorrectly treat + it as single-table. + """ + assert ( + count_referenced_tables( + 'SELECT l.a, r.a FROM "db.table1" AS l JOIN "db.table1" AS r ON l.a = r.a', + Dialects.SQLITE, + ) + == 2 + ) + + +def test_count_referenced_tables_cte_self_join() -> None: + """ + A CTE that reads a single virtual table and is then self-joined must + count as 2, matching the direct self-join case, since the CTE is + inlined at each of its two consumption sites and triggers a read of + that table for both sides of the join. + """ + assert ( + count_referenced_tables( + 'WITH cte AS (SELECT a FROM "db.table1") ' + "SELECT l.a, r.a FROM cte AS l JOIN cte AS r ON l.a = r.a", + Dialects.SQLITE, + ) + == 2 + ) + # A CTE used exactly once, with no join, still counts as a single table. + assert ( + count_referenced_tables( + 'WITH cte AS (SELECT a FROM "db.table1") SELECT a FROM cte', + Dialects.SQLITE, + ) + == 1 + ) + # A CTE joined against a distinct real table also counts as 2. + assert ( + count_referenced_tables( + 'WITH cte AS (SELECT a FROM "db.table1") ' + 'SELECT l.a, r.a FROM cte AS l JOIN "db.table2" AS r ON l.a = r.a', + Dialects.SQLITE, + ) + == 2 + ) + # Nested CTEs: a CTE built on top of another CTE, then self-joined, + # still weights the base CTE's own table by the self-join count. + assert ( + count_referenced_tables( + 'WITH base AS (SELECT a FROM "db.table1"), derived AS (SELECT a FROM base) ' + "SELECT l.a, r.a FROM derived AS l JOIN derived AS r ON l.a = r.a", + Dialects.SQLITE, + ) + == 2 + ) + + +def test_count_referenced_tables_describe() -> None: + """ + ``DESCRIBE`` (and other ``exp.Describe``/``exp.Command`` statements) has + no join semantics for a per-table row cap to interact with, so it takes + the plain unweighted table-extraction path rather than + ``_count_weighted_table_references``. + """ + assert count_referenced_tables("DESCRIBE table1", Dialects.SQLITE) == 1 + + +def test_count_referenced_tables_derived_subqueries() -> None: + """ + Two distinct derived (non-CTE) subqueries joined together must each + resolve their own tables directly, without recursing as if they were + CTE sources -- covering the branch in ``_count_weighted_table_references`` + where a selected source is a ``Scope`` but not a CTE. + """ + assert ( + count_referenced_tables( + 'SELECT l.a, r.a FROM (SELECT a FROM "db.table1") AS l ' + 'JOIN (SELECT a FROM "db.table1") AS r ON l.a = r.a', + Dialects.SQLITE, + ) + == 2 + ) + + +def test_count_weighted_table_references_self_referential_scope_guard( + mocker: MockerFixture, +) -> None: + """ + ``_count_weighted_table_references`` must not recurse forever on a + self-referential ``Scope`` graph, the shape a ``WITH RECURSIVE`` CTE + could in principle produce if sqlglot ever resolved its own + self-reference to the same ``Scope`` object instead of a bare + ``exp.Table``. The ``seen`` guard must catch the repeat visit and treat + it as contributing no further table reads. + """ + from sqlglot.optimizer.scope import Scope, ScopeType # noqa: PLC0415 + + cte_scope = Scope.__new__(Scope) + cte_scope.scope_type = ScopeType.CTE + # The CTE's own body references itself. + cte_scope._selected_sources = {"t": (None, cte_scope)} # noqa: SLF001 + + root_scope = Scope.__new__(Scope) + root_scope.scope_type = ScopeType.ROOT + root_scope._selected_sources = {"t": (None, cte_scope)} # noqa: SLF001 + + mocker.patch( + "superset.sql.parse.traverse_scope", + return_value=[cte_scope, root_scope], + ) + + assert _count_weighted_table_references(mocker.MagicMock()) == 0 + + +def test_count_referenced_tables_respects_parse_length_cap( + mocker: MockerFixture, +) -> None: + """ + ``count_referenced_tables`` must not bypass ``SQL_MAX_PARSE_LENGTH``: an + oversized statement should fail the length check before reaching + sqlglot, and fall back to the conservative single-table count. The + statement references two tables so that bypassing the guard (and + reaching sqlglot) would produce a different, detectable result. + """ + mocker.patch("superset.config.SQL_MAX_PARSE_LENGTH", 100) + mocker.patch("superset.sql.parse.has_app_context", return_value=False) + padding = "1, " * 50 + statement = ( + 'SELECT * FROM "db.table1" AS t1 ' # noqa: S608 + 'JOIN "db.table2" AS t2 ON t1.a = t2.a ' + f"WHERE t1.a IN ({padding}1)" + ) + assert len(statement.encode("utf-8")) > 100 + assert count_referenced_tables(statement, Dialects.SQLITE) == 1 + + def test_extract_tables_subselect() -> None: """ Test that tables inside subselects are parsed correctly.