add features to handle unnamed queries, extracting properly hive tables from queries with subscripts, some additional issues that were already fixed were documented by tests, some cleanup and refactor to decrease unreachable paths

Author: collerek

sql_metadata/column_extractor.py

@@ -37,6 +37,7 @@ class ExtractionResult:
     alias_map: Dict[str, Union[str, list]]
     cte_names: UniqueList
     subquery_names: UniqueList
+    output_columns: list
 
 
 # ---------------------------------------------------------------------------
@@ -149,6 +150,7 @@ class _Collector:
         "cte_names",
         "cte_alias_names",
         "subquery_items",
+        "output_columns",
     )
 
     def __init__(self, table_aliases: Dict[str, str]):
@@ -161,6 +163,7 @@ def __init__(self, table_aliases: Dict[str, str]):
         self.cte_names = UniqueList()
         self.cte_alias_names: set = set()
         self.subquery_items: list = []
+        self.output_columns: list = []
 
     def add_column(self, name: str, clause: str) -> None:
         """Record a column name, filing it into the appropriate section."""
@@ -227,12 +230,6 @@ def extract(self) -> ExtractionResult:
 
         self._seed_cte_names()
 
-        # Handle CREATE TABLE with column defs (no SELECT)
-        if isinstance(self._ast, exp.Create) and not self._ast.find(exp.Select):
-            for col_def in self._ast.find_all(exp.ColumnDef):
-                c.add_column(col_def.name, "")
-            return self._build_result()
-
         # Reset cte_names — walk will re-collect them in order
         c.cte_names = UniqueList()
         self._walk(self._ast)
@@ -251,6 +248,7 @@ def extract(self) -> ExtractionResult:
             alias_map=c.alias_map,
             cte_names=final_cte,
             subquery_names=self._build_subquery_names(),
+            output_columns=c.output_columns,
         )
 
     # -------------------------------------------------------------------
@@ -279,19 +277,6 @@ def _build_subquery_names(self) -> UniqueList:
             names.append(name)
         return names
 
-    def _build_result(self) -> ExtractionResult:
-        """Build result from collector (used for early-return CREATE TABLE path)."""
-        c = self._collector
-        alias_dict = c.alias_dict if c.alias_dict else None
-        return ExtractionResult(
-            columns=c.columns,
-            columns_dict=c.columns_dict,
-            alias_names=c.alias_names,
-            alias_dict=alias_dict,
-            alias_map=c.alias_map,
-            cte_names=c.cte_names,
-            subquery_names=self._build_subquery_names(),
-        )
 
     # -------------------------------------------------------------------
     # Column name helpers
@@ -366,15 +351,11 @@ def _dispatch_leaf(self, node: exp.Expression, clause: str, depth: int) -> bool:
 
         Returns ``True`` if handled (stop recursion), ``False`` to continue.
         """
-        if isinstance(node, (exp.Values, exp.Star, exp.ColumnDef, exp.Identifier)):
-            # TODO: revisit if Stars appear outside Select.expressions
-            if isinstance(node, exp.Star):  # pragma: no cover
-                self._handle_star(node, clause)
-            # TODO: revisit if CREATE TABLE walk stops returning early
-            elif isinstance(node, exp.ColumnDef):  # pragma: no cover
+        if isinstance(node, exp.Values) and not node.find(exp.Select):
+            return True
+        if isinstance(node, (exp.Star, exp.ColumnDef, exp.Identifier)):
+            if isinstance(node, exp.ColumnDef):
                 self._collector.add_column(node.name, clause)
-            elif isinstance(node, exp.Identifier):
-                self._handle_identifier(node, clause)
             return True
         if isinstance(node, exp.CTE):
             self._handle_cte(node, depth)
@@ -417,23 +398,6 @@ def _recurse_child(self, child: Any, clause: str, depth: int) -> None:
     # Node handlers
     # -------------------------------------------------------------------
 
-    # TODO: revisit if Stars reach _dispatch_leaf
-    def _handle_star(self, node: exp.Star, clause: str) -> None:  # pragma: no cover
-        """Handle a standalone Star node (not inside a Column or function)."""
-        not_in_col = not isinstance(node.parent, exp.Column)
-        if not_in_col and not self._is_star_inside_function(node):
-            self._collector.add_column("*", clause)
-
-    def _handle_identifier(self, node: exp.Identifier, clause: str) -> None:
-        """Handle an Identifier in a USING clause (not inside a Column)."""
-        if not isinstance(
-            node.parent,
-            (exp.Column, exp.Table, exp.TableAlias, exp.CTE),
-        ):
-            # TODO: revisit if JOIN produces bare Identifiers
-            if clause == "join":  # pragma: no cover
-                self._collector.add_column(node.name, clause)
-
     def _handle_insert_schema(self, node: exp.Insert) -> None:
         """Extract column names from the Schema of an INSERT statement."""
         schema = node.find(exp.Schema)
@@ -482,18 +446,23 @@ def _handle_column(self, col: exp.Column, clause: str) -> None:
     def _handle_select_exprs(self, exprs: list, clause: str, depth: int) -> None:
         """Handle the expressions list of a SELECT clause."""
         assert isinstance(exprs, list)
+        out = self._collector.output_columns
 
         for expr in exprs:
             if isinstance(expr, exp.Alias):
                 self._handle_alias(expr, clause, depth)
+                out.append(expr.alias)
             elif isinstance(expr, exp.Star):
                 self._collector.add_column("*", clause)
+                out.append("*")
             elif isinstance(expr, exp.Column):
                 self._handle_column(expr, clause)
+                out.append(self._column_full_name(expr))
             else:
                 cols = self._flat_columns(expr)
                 for col in cols:
                     self._collector.add_column(col, clause)
+                out.append(cols[0] if len(cols) == 1 else str(expr))
 
     def _handle_alias(self, alias_node: exp.Alias, clause: str, depth: int) -> None:
         """Handle an Alias node inside a SELECT expression list."""

sql_metadata/dialect_parser.py

@@ -103,12 +103,12 @@ def _detect_dialects(sql: str) -> list:
             return [HashVarDialect, None, "mysql"]
         if "`" in sql:
             return ["mysql", None]
+        if "LATERAL VIEW" in upper:
+            return ["spark", None, "mysql"]
         if "[" in sql or " TOP " in upper:
             return [BracketedTableDialect, None, "mysql"]
         if " UNIQUE " in upper:
             return [None, "mysql", "oracle"]
-        if "LATERAL VIEW" in upper:
-            return ["spark", None, "mysql"]
         return [None, "mysql"]
 
     # -- parsing ------------------------------------------------------------

sql_metadata/nested_resolver.py

@@ -9,7 +9,7 @@
 from __future__ import annotations
 
 import copy
-from typing import TYPE_CHECKING, Dict, List, Optional, Set, Union
+from typing import TYPE_CHECKING, Dict, List, Optional, Set, Tuple, Union
 
 if TYPE_CHECKING:
     from sql_metadata.parser import Parser
@@ -167,24 +167,49 @@ def extract_cte_names(
         return names
 
     @staticmethod
-    def extract_subquery_names(ast: Optional[exp.Expression]) -> List[str]:
-        """Extract aliased subquery names from the AST in post-order.
+    def extract_subqueries(
+        ast: Optional[exp.Expression],
+    ) -> Tuple[List[str], Dict[str, str]]:
+        """Extract subquery names and bodies in a single post-order walk.
+
+        Aliased subqueries keep their alias as the name.  Unaliased
+        subqueries (e.g. ``WHERE id IN (SELECT …)``) get auto-generated
+        names ``subquery_1``, ``subquery_2``, etc.
 
-        Called by :attr:`Parser.subqueries_names`.
+        :returns: ``(names, bodies)`` where *names* is ordered innermost-first.
         """
         if ast is None:  # pragma: no cover — Parser ensures AST exists
-            return []
-        names = UniqueList()
-        NestedResolver._collect_subquery_names_postorder(ast, names)
-        return names
+            return [], {}
+        names: list = UniqueList()
+        bodies: Dict[str, str] = {}
+        NestedResolver._walk_subqueries(ast, names, bodies, 0)
+        return names, bodies
 
     @staticmethod
-    def _collect_subquery_names_postorder(node: exp.Expression, out: list) -> None:
-        """Recursively collect subquery aliases in post-order."""
+    def _walk_subqueries(
+        node: exp.Expression,
+        names: list,
+        bodies: Dict[str, str],
+        counter: int,
+    ) -> int:
+        """Post-order walk collecting subquery names and bodies.
+
+        Returns the updated *counter* so unnamed subqueries are numbered
+        sequentially.
+        """
         for child in node.iter_expressions():
-            NestedResolver._collect_subquery_names_postorder(child, out)
-        if isinstance(node, exp.Subquery) and node.alias:
-            out.append(node.alias)
+            counter = NestedResolver._walk_subqueries(
+                child, names, bodies, counter
+            )
+        if isinstance(node, exp.Subquery):
+            if node.alias:
+                name = node.alias
+            else:
+                counter += 1
+                name = f"subquery_{counter}"
+            names.append(name)
+            bodies[name] = NestedResolver._body_sql(node.this)
+        return counter
 
     # -------------------------------------------------------------------
     # Body extraction
@@ -226,39 +251,6 @@ def extract_cte_bodies(
 
         return results
 
-    def extract_subquery_bodies(
-        self,
-        subquery_names: List[str],
-    ) -> Dict[str, str]:
-        """Extract subquery body SQL for each name in *subquery_names*.
-
-        Uses a post-order AST walk so that inner subqueries appear before
-        outer ones.
-
-        :param subquery_names: List of subquery alias names to extract.
-        :returns: Mapping of ``{subquery_name: body_sql}``.
-        """
-        if not self._ast or not subquery_names:
-            return {}
-
-        names_upper = {n.upper(): n for n in subquery_names}
-        results: Dict[str, str] = {}
-        self._collect_subqueries_postorder(self._ast, names_upper, results)
-        return results
-
-    @staticmethod
-    def _collect_subqueries_postorder(
-        node: exp.Expression, names_upper: Dict[str, str], out: Dict[str, str]
-    ) -> None:
-        """Recursively collect subquery bodies in post-order."""
-        for child in node.iter_expressions():
-            NestedResolver._collect_subqueries_postorder(child, names_upper, out)
-        if isinstance(node, exp.Subquery) and node.alias:
-            alias_upper = node.alias.upper()
-            if alias_upper in names_upper:
-                original_name = names_upper[alias_upper]
-                out[original_name] = NestedResolver._body_sql(node.this)
-
     # -------------------------------------------------------------------
     # Column resolution (from parser.py)
     # -------------------------------------------------------------------