rusackas commented on code in PR #42598:
URL: https://github.com/apache/superset/pull/42598#discussion_r3832501445
##########
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:
Good catch, real gap. `scope.sources` collapses a self-joined CTE to a
single dict entry per alias, so the definition scope's own tables only got
counted once no matter how many times it's consumed. Weighted it through
`selected_sources` so each reference counts on its own (recurses through nested
CTEs too), pushed with a 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]
