further fixes and duplication cleanup

Author: collerek

sql_metadata/column_extractor.py

@@ -16,7 +16,7 @@
 
 from sqlglot import exp
 
-from sql_metadata.utils import UniqueList, _make_reverse_cte_map
+from sql_metadata.utils import UniqueList, _make_reverse_cte_map, last_segment
 
 # ---------------------------------------------------------------------------
 # Result dataclass
@@ -516,10 +516,10 @@ def _handle_alias(self, alias_node: exp.Alias, clause: str, depth: int) -> None:
             for col in inner_cols:
                 c.add_column(col, clause)
 
-            unique_inner = list(dict.fromkeys(inner_cols))
+            unique_inner = UniqueList(inner_cols)
             is_self_alias = len(unique_inner) == 1 and (
                 unique_inner[0] == alias_name
-                or unique_inner[0].split(".")[-1] == alias_name
+                or last_segment(unique_inner[0]) == alias_name
             )
             is_direct = isinstance(inner, exp.Column)
 

sql_metadata/nested_resolver.py

@@ -17,7 +17,13 @@
 from sqlglot import exp
 from sqlglot.generator import Generator
 
-from sql_metadata.utils import UniqueList, _make_reverse_cte_map, flatten_list
+from sql_metadata.utils import (
+    DOT_PLACEHOLDER,
+    UniqueList,
+    _make_reverse_cte_map,
+    flatten_list,
+    last_segment,
+)
 
 # ---------------------------------------------------------------------------
 # Custom SQL generator — preserves function signatures
@@ -119,6 +125,7 @@ def __init__(
         self._subqueries_parsers: Dict = {}
         self._with_parsers: Dict = {}
         self._columns_aliases: Dict = {}
+        self._cached_cte_nodes: Optional[list] = None
 
         # Set by resolve() caller
         self._subqueries_names: List[str] = []
@@ -130,21 +137,29 @@ def __init__(
     # Name extraction (CTE and subquery names from the AST)
     # -------------------------------------------------------------------
 
-    @staticmethod
+    def _cte_nodes(self) -> list:
+        """Return all ``exp.CTE`` nodes from the AST (cached)."""
+        if self._cached_cte_nodes is None:
+            if self._ast is None:
+                self._cached_cte_nodes = []
+            else:
+                self._cached_cte_nodes = list(self._ast.find_all(exp.CTE))
+        return self._cached_cte_nodes
+
     def extract_cte_names(
-        ast: Optional[exp.Expression],
+        self,
         cte_name_map: Optional[Dict] = None,
     ) -> List[str]:
         """Extract CTE names from the AST.
 
         Called by :attr:`Parser.with_names`.
         """
-        if ast is None:
+        if self._ast is None:
             return []
         cte_name_map = cte_name_map or {}
         reverse_map = _make_reverse_cte_map(cte_name_map)
         names = UniqueList()
-        for cte in ast.find_all(exp.CTE):
+        for cte in self._cte_nodes():
             alias = cte.alias
             if alias:
                 names.append(reverse_map.get(alias, alias))
@@ -196,13 +211,13 @@ def extract_cte_bodies(
 
         alias_to_name: Dict[str, str] = {}
         for name in cte_names:
-            placeholder = name.replace(".", "__DOT__")
+            placeholder = name.replace(".", DOT_PLACEHOLDER)
             alias_to_name[placeholder.upper()] = name
             alias_to_name[name.upper()] = name
-            alias_to_name[name.split(".")[-1].upper()] = name
+            alias_to_name[last_segment(name).upper()] = name
 
         results: Dict[str, str] = {}
-        for cte in self._ast.find_all(exp.CTE):
+        for cte in self._cte_nodes():
             alias = cte.alias
             if alias.upper() in alias_to_name:
                 original_name = alias_to_name[alias.upper()]
@@ -312,42 +327,34 @@ def _resolve_and_filter(
             final.append(col)
         return final
 
+    def _nested_sources(self) -> list:
+        """Return the (names, defs, cache) tuples for subqueries then CTEs."""
+        return [
+            (self._subqueries_names, self._subqueries, self._subqueries_parsers),
+            (self._with_names, self._with_queries, self._with_parsers),
+        ]
+
     def _resolve_sub_queries(self, column: str) -> Union[str, List[str]]:
         """Resolve a ``subquery.column`` reference to actual column(s)."""
-        result = self._resolve_nested_query(
-            subquery_alias=column,
-            nested_queries_names=self._subqueries_names,
-            nested_queries=self._subqueries,
-            already_parsed=self._subqueries_parsers,
-        )
-        if isinstance(result, str):
-            result = self._resolve_nested_query(
-                subquery_alias=result,
-                nested_queries_names=self._with_names,
-                nested_queries=self._with_queries,
-                already_parsed=self._with_parsers,
-            )
+        result: Union[str, List[str]] = column
+        for names, defs, cache in self._nested_sources():
+            if isinstance(result, str):
+                result = self._resolve_nested_query(
+                    subquery_alias=result,
+                    nested_queries_names=names,
+                    nested_queries=defs,
+                    already_parsed=cache,
+                )
         return result if isinstance(result, list) else [result]
 
     def _resolve_bare_through_nested(self, col_name: str) -> Union[str, List[str]]:
         """Resolve a bare column name through subquery/CTE alias definitions."""
-        result = self._lookup_alias_in_nested(
-            col_name,
-            self._subqueries_names,
-            self._subqueries,
-            self._subqueries_parsers,
-            check_columns=True,
-        )
-        if result is not None:
-            return result
-        result = self._lookup_alias_in_nested(
-            col_name,
-            self._with_names,
-            self._with_queries,
-            self._with_parsers,
-        )
-        if result is not None:
-            return result
+        for i, (names, defs, cache) in enumerate(self._nested_sources()):
+            result = self._lookup_alias_in_nested(
+                col_name, names, defs, cache, check_columns=(i == 0)
+            )
+            if result is not None:
+                return result
         return col_name
 
     def _lookup_alias_in_nested(
@@ -447,7 +454,7 @@ def _find_column_fallback(
     ) -> Union[str, List[str]]:
         """Find a column by name in the subparser with wildcard fallbacks."""
         try:
-            idx = [x.split(".")[-1] for x in subparser.columns].index(column_name)
+            idx = [last_segment(x) for x in subparser.columns].index(column_name)
         except ValueError:
             if "*" in subparser.columns:
                 return column_name

sql_metadata/parser.py

@@ -298,8 +298,9 @@ def with_names(self) -> List[str]:
         """Return the CTE (Common Table Expression) names from the query."""
         if self._with_names is not None:
             return self._with_names
-        self._with_names = NestedResolver.extract_cte_names(
-            self._ast_parser.ast, self._ast_parser.cte_name_map
+        resolver = self._get_resolver()
+        self._with_names = resolver.extract_cte_names(
+            self._ast_parser.cte_name_map
         )
         return self._with_names
 

sql_metadata/sql_cleaner.py

@@ -11,6 +11,7 @@
 from typing import NamedTuple, Optional
 
 from sql_metadata.comments import strip_comments_for_parsing as _strip_comments
+from sql_metadata.utils import DOT_PLACEHOLDER
 
 
 class CleanResult(NamedTuple):
@@ -70,7 +71,7 @@ def replacer(match: re.Match[str]) -> str:
         prefix = match.group(1)
         qualified_name = match.group(2)
         suffix = match.group(3)
-        placeholder = qualified_name.replace(".", "__DOT__")
+        placeholder = qualified_name.replace(".", DOT_PLACEHOLDER)
         name_map[placeholder] = qualified_name
         return f"{prefix}{placeholder}{suffix}"
 

sql_metadata/table_extractor.py

@@ -13,7 +13,7 @@
 
 from sqlglot import exp
 
-from sql_metadata.utils import UniqueList
+from sql_metadata.utils import UniqueList, last_segment
 
 # ---------------------------------------------------------------------------
 # Pure static helpers (no instance state needed)
@@ -127,6 +127,7 @@ def __init__(
         self._bracket_mode = isinstance(dialect, type) and issubclass(
             dialect, BracketedTableDialect
         )
+        self._cached_table_nodes: Optional[List[exp.Table]] = None
 
     # -------------------------------------------------------------------
     # Public API
@@ -152,6 +153,13 @@ def extract(self) -> List[str]:
         collected_sorted = sorted(collected, key=lambda t: self._first_position(t))
         return self._place_tables_in_order(create_target, collected_sorted)
 
+    def _table_nodes(self) -> List[exp.Table]:
+        """Return all ``exp.Table`` nodes from the AST (cached)."""
+        if self._cached_table_nodes is None:
+            assert self._ast is not None
+            self._cached_table_nodes = list(self._ast.find_all(exp.Table))
+        return self._cached_table_nodes
+
     def extract_aliases(self, tables: List[str]) -> Dict[str, str]:
         """Extract table alias mappings from the AST.
 
@@ -162,7 +170,7 @@ def extract_aliases(self, tables: List[str]) -> Dict[str, str]:
             return {}
 
         aliases = {}
-        for table in self._ast.find_all(exp.Table):
+        for table in self._table_nodes():
             alias = table.alias
             if not alias:
                 continue
@@ -203,19 +211,25 @@ def _first_position(self, name: str) -> int:
         if pos >= 0:
             return pos
 
-        last_part = name_upper.split(".")[-1]
+        last_part = last_segment(name_upper)
         pos = self._find_word_in_table_context(last_part)
         if pos >= 0:
             return pos
 
         pos = self._find_word(name_upper)
         return pos if pos >= 0 else len(self._raw_sql)
 
+    _pattern_cache: Dict[str, re.Pattern[str]] = {}
+
     @staticmethod
     def _word_pattern(name_upper: str) -> re.Pattern[str]:
-        """Build a regex matching *name_upper* as a whole word."""
-        escaped = re.escape(name_upper)
-        return re.compile(r"(?<![A-Za-z0-9_])" + escaped + r"(?![A-Za-z0-9_])")
+        """Build a regex matching *name_upper* as a whole word (cached)."""
+        pat = TableExtractor._pattern_cache.get(name_upper)
+        if pat is None:
+            escaped = re.escape(name_upper)
+            pat = re.compile(r"(?<![A-Za-z0-9_])" + escaped + r"(?![A-Za-z0-9_])")
+            TableExtractor._pattern_cache[name_upper] = pat
+        return pat
 
     def _find_word(self, name_upper: str, start: int = 0) -> int:
         """Find *name_upper* as a whole word in the upper-cased SQL."""
@@ -271,7 +285,7 @@ def _collect_all(self) -> UniqueList:
         """Collect table names from Table and Lateral AST nodes."""
         assert self._ast is not None
         collected = UniqueList()
-        for table in self._ast.find_all(exp.Table):
+        for table in self._table_nodes():
             full_name = self._table_full_name(table)
             if full_name and full_name not in self._cte_names:
                 collected.append(full_name)

sql_metadata/utils.py

@@ -7,6 +7,10 @@
 
 from typing import Any, Dict, Iterable, List
 
+#: Placeholder used to encode dots in qualified CTE names so that sqlglot
+#: does not misinterpret ``db.cte_name`` as a table reference.
+DOT_PLACEHOLDER = "__DOT__"
+
 
 class UniqueList(list):
     """A list subclass that silently rejects duplicate items.
@@ -17,9 +21,14 @@ class UniqueList(list):
     an internal ``set`` for O(1) membership checks.
     """
 
-    def __init__(self, *args: Any, **kwargs: Any) -> None:
-        super().__init__(*args, **kwargs)
-        self._seen: set = set(self)
+    def __init__(self, iterable: Any = None, **kwargs: Any) -> None:
+        self._seen: set = set()
+        if iterable is not None:
+            super().__init__(**kwargs)
+            self.extend(iterable)
+        else:
+            super().__init__(**kwargs)
+            self._seen = set(self)
 
     def append(self, item: Any) -> None:
         """Append *item* only if it is not already present (O(1) check)."""
@@ -32,6 +41,10 @@ def extend(self, items: Iterable[Any]) -> None:  # type: ignore[override]
         for item in items:
             self.append(item)
 
+    def __contains__(self, item: Any) -> bool:
+        """O(1) membership check using the internal set."""
+        return item in self._seen
+
     def __sub__(self, other: Any) -> List:
         """Return a plain list of elements in *self* that are not in *other*."""
         other_set = set(other)
@@ -40,11 +53,16 @@ def __sub__(self, other: Any) -> List:
 
 def _make_reverse_cte_map(cte_name_map: Dict) -> Dict[str, str]:
     """Build reverse mapping from placeholder CTE names to originals."""
-    reverse = {v.replace(".", "__DOT__"): v for v in cte_name_map.values()}
+    reverse = {v.replace(".", DOT_PLACEHOLDER): v for v in cte_name_map.values()}
     reverse.update(cte_name_map)
     return reverse
 
 
+def last_segment(name: str) -> str:
+    """Return the last dot-separated segment of a qualified name."""
+    return name.rsplit(".", 1)[-1]
+
+
 def flatten_list(input_list: List) -> List[str]:
     """Recursively flatten a list that may contain nested lists.