rusackas commented on code in PR #42598:
URL: https://github.com/apache/superset/pull/42598#discussion_r3709888878
##########
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:
Good catch, fixed. `count_referenced_tables` now counts table reference
occurrences instead of distinct tables (factored the source traversal into
`_find_table_sources`, dedup only happens in `extract_tables_from_statement`),
so a self-join no longer collapses to a single-table count. Added
`test_count_referenced_tables_self_join` covering it directly.
##########
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:
Good catch, fixed. The statement now references two tables, so bypassing
`_check_script_length` would parse to `2` instead of `1` - the test actually
exercises the guard now instead of passing regardless.
--
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]
