codeant-ai-for-open-source[bot] commented on code in PR #42598:
URL: https://github.com/apache/superset/pull/42598#discussion_r3707710760
##########
tests/unit_tests/sql/parse_tests.py:
##########
@@ -234,6 +235,45 @@ def test_extract_tables_from_sql() -> None:
) == {Table("forbidden_table")}
+def test_count_referenced_tables() -> None:
+ """
+ Test that ``count_referenced_tables`` counts distinct table references,
+ 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
+
+
+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.
+ """
+ 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" WHERE a IN (' + padding + "1)"
+ assert len(statement.encode("utf-8")) > 100
+ assert count_referenced_tables(statement, Dialects.SQLITE) == 1
Review Comment:
Confirmed—the review comment is valid. The current fixture cannot
distinguish between enforcing and bypassing `_check_script_length`, because
both paths return `1` for a single-table statement.
A minimal fix is to make the oversized statement reference two distinct
tables:
```python
padding = "1, " * 50
statement = (
'SELECT * FROM "db.table1" AS t1 '
'JOIN "db.table2" AS t2 ON t1.id = t2.id '
f"WHERE t1.a IN ({padding}1)"
)
assert len(statement.encode("utf-8")) > 100
assert count_referenced_tables(statement, Dialects.SQLITE) == 1
```
With the length guard active, parsing is skipped and the conservative
fallback returns `1`. If the guard is bypassed, `sqlglot` parses both table
references and returns `2`, causing the test to fail. This directly verifies
that `count_referenced_tables` respects `SQL_MAX_PARSE_LENGTH`; no additional
production change is required.
##########
superset/extensions/metadb.py:
##########
@@ -119,6 +148,53 @@ def create_connect_args(self, url: URL) ->
tuple[tuple[()], dict[str, Any]]:
},
)
+ def do_execute(
+ self,
+ cursor: Any,
+ statement: str,
+ parameters: Any,
+ context: Any = None,
+ ) -> None:
+ with self._flag_multi_table_query(statement):
+ super().do_execute(cursor, statement, parameters, context)
+
+ def do_execute_no_params(
+ self,
+ cursor: Any,
+ statement: str,
+ context: Any = None,
+ ) -> None:
+ with self._flag_multi_table_query(statement):
+ super().do_execute_no_params(cursor, statement, context)
+
+ def do_executemany(
+ self,
+ cursor: Any,
+ statement: str,
+ parameters: Any,
+ context: Any = None,
+ ) -> None:
+ with self._flag_multi_table_query(statement):
+ super().do_executemany(cursor, statement, parameters, context)
+
+ @staticmethod
+ @contextmanager
+ def _flag_multi_table_query(statement: str) -> Iterator[None]:
+ """
+ Record, for the duration of executing ``statement``, whether it
references
+ more than one `superset://` virtual table.
+
+ `SupersetShillelaghAdapter.get_data` reads this to decide whether it's
safe
+ to apply `SUPERSET_META_DB_LIMIT` to the table it's fetching (see
`get_data`).
+ """
+ token = _executing_multi_table_query.set(
+ _count_referenced_tables(statement) > 1
+ )
Review Comment:
Confirmed—the issue is valid. `extract_tables_from_statement()` returns a
`set`, so a self-join such as:
```sql
SELECT l.a, r.a
FROM "database1.table1_large" AS l
JOIN "database1.table1_large" AS r
ON l.a = r.a
```
is counted as one table. `SUPERSET_META_DB_LIMIT` can therefore still
truncate the shared virtual table before SQLite evaluates both aliases.
The fix should make `count_referenced_tables()` count table **occurrences**,
while retaining the existing parser/CTE handling rather than counting distinct
`Table` values. For example, factor the source traversal into a helper that
returns a list, and use its length for this function:
```python
def count_referenced_tables(statement: str, dialect: Dialects | str | None)
-> int:
try:
_check_script_length(statement, str(dialect) if dialect else None)
parsed = sqlglot.parse_one(statement, dialect=dialect)
return len(_extract_table_references(parsed, dialect))
except Exception: # pylint: disable=broad-except
return 1
```
Add a regression with a late-match self-join and `SUPERSET_META_DB_LIMIT=2`,
asserting that the expected row is returned. This ensures aliases of the same
physical table disable the per-table cap just like joins across different
tables.
This change is necessary because the bug affects query structure, not
physical-table identity; counting occurrences closes the self-join gap while
preserving the existing behavior for single-table statements.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
