sadpandajoe commented on code in PR #42598:
URL: https://github.com/apache/superset/pull/42598#discussion_r3707708326
##########
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: `extract_tables_from_statement` returns a set, so two aliases of
the same virtual table collapse to one and the fallback cap is still applied
before the self-join. Could this count table occurrences rather than distinct
table identities and add a late-match self-join regression?
##########
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: this assertion is unchanged if the length check is removed
because the oversized SQL still parses to one table. Could the fixture
reference two tables (or assert the length-check call) so the test fails when
the guard is bypassed?
--
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]
