sadpandajoe commented on code in PR #42598:
URL: https://github.com/apache/superset/pull/42598#discussion_r3830710073
##########
superset/sql/parse.py:
##########
@@ -2129,84 +2128,128 @@ def is_valid_cvas(self) -> bool:
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:
DESCRIBE some_table;
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)
Review Comment:
This intentionally ignores CTE references, but the count now decides whether
the source read may be capped. A CTE that reads one virtual table and is joined
to itself still counts as one here, so its source is limited before the CTE
self-join and a late match disappears. Could multi-source classification
account for repeated CTE consumers and add that regression?
##########
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):
Review Comment:
SQL Lab executes this engine through `raw_connection()` and a DBAPI
`cursor.execute()`, which bypasses the SQLAlchemy dialect's `do_execute*`
hooks. That leaves this flag false for the user-facing path, so a join with a
match beyond `SUPERSET_META_DB_LIMIT` is still truncated before SQLite joins
it. Could the execution state be established for raw-cursor execution too, with
an SQL Lab-path regression test?
--
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]
