additional changes and optimizations after initial review

Author: collerek

sql_metadata/ast_parser.py

@@ -48,6 +48,10 @@ def ast(self) -> exp.Expression | None:
         self._ast = self._parse(self._raw_sql)
         return self._ast
 
+    def _ensure_parsed(self) -> None:
+        """Trigger lazy parsing so side-effect fields are populated."""
+        _ = self.ast
+
     @property
     def dialect(self) -> DialectType:
         """The sqlglot dialect that produced the current AST.
@@ -58,7 +62,7 @@ def dialect(self) -> DialectType:
 
         :rtype: DialectType
         """
-        _ = self.ast
+        self._ensure_parsed()
         return self._dialect
 
     @property
@@ -72,7 +76,7 @@ def is_replace(self) -> bool:
 
         :rtype: bool
         """
-        _ = self.ast
+        self._ensure_parsed()
         return self._is_replace
 
     @property
@@ -84,7 +88,7 @@ def cte_name_map(self) -> dict[str, str]:
 
         :rtype: dict[str, str]
         """
-        _ = self.ast
+        self._ensure_parsed()
         return self._cte_name_map
 
     def _parse(self, sql: str) -> exp.Expression | None:

sql_metadata/column_extractor.py

@@ -251,7 +251,7 @@ def __init__(
         self._table_aliases = table_aliases
         self._cte_name_map = cte_name_map or {}
         self._collector = _Collector()
-        self._reverse_cte_map = self._cte_name_map
+        self._cte_restore_map = self._cte_name_map
 
     # -------------------------------------------------------------------
     # Public API
@@ -290,7 +290,7 @@ def extract(self) -> ExtractionResult:
         # Restore qualified CTE names (reverse placeholder mapping)
         final_cte = UniqueList()
         for name in c.cte_names:
-            final_cte.append(self._reverse_cte_map.get(name, name))
+            final_cte.append(self._cte_restore_map.get(name, name))
 
         alias_dict = c.alias_dict
         return ExtractionResult(
@@ -327,7 +327,7 @@ def _seed_cte_names(self) -> None:
             alias = cte.alias
             if alias:
                 self._collector.cte_names.append(
-                    self._reverse_cte_map.get(alias, alias)
+                    self._cte_restore_map.get(alias, alias)
                 )
 
     def _build_subquery_names(self) -> UniqueList:

sql_metadata/parser.py

@@ -21,6 +21,7 @@
 from sql_metadata.column_extractor import ColumnExtractor
 from sql_metadata.comments import extract_comments, strip_comments
 from sql_metadata.generalizator import Generalizator
+from sql_metadata.keywords_lists import QueryType
 from sql_metadata.nested_resolver import NestedResolver
 from sql_metadata.query_type_extractor import QueryTypeExtractor
 from sql_metadata.sql_cleaner import SqlCleaner
@@ -47,7 +48,7 @@ def __init__(self, sql: str = "", disable_logging: bool = False) -> None:
         self._logger.disabled = disable_logging
 
         self._raw_query = sql
-        self._query_type: str | None = None
+        self._query_type: QueryType | None = None
 
         self._ast_parser = ASTParser(sql)
         self._resolver: NestedResolver | None = None
@@ -104,7 +105,7 @@ def query(self) -> str:
         return SqlCleaner.preprocess_query(self._raw_query)
 
     @property
-    def query_type(self) -> str:
+    def query_type(self) -> QueryType | None:
         """Return the type of the SQL query.
 
         Delegates to :class:`QueryTypeExtractor` which maps the top-level
@@ -113,7 +114,7 @@ def query_type(self) -> str:
         (unparseable SQL), the extractor falls back to keyword matching on
         the raw string.
 
-        :rtype: str
+        :rtype: QueryType | None
         """
         if self._query_type:
             return self._query_type
@@ -320,7 +321,7 @@ def tables_aliases(self) -> dict[str, str]:
         if self._table_aliases is not None:
             return self._table_aliases
         ast = self._ast_parser.ast
-        assert ast is not None  # guaranteed by prior tables/query_type access
+        assert ast is not None  # mypy
         extractor = TableExtractor(ast)
         self._table_aliases = extractor.extract_aliases(self.tables)
         return self._table_aliases

test/test_edge_cases.py

@@ -1,4 +1,10 @@
-"""Edge-case tests for internals not covered by feature-specific test files."""
+"""Edge-case tests for internal utilities.
+
+These tests exercise code paths (depth guards, degenerate inputs) that
+are difficult or impossible to trigger through the public Parser API.
+They test internal symbols directly and may need updating if those
+internals are refactored.
+"""
 
 from sql_metadata import Parser
 from sql_metadata.sql_cleaner import SqlCleaner, _strip_outer_parens

test/test_getting_columns.py

@@ -772,4 +772,5 @@ def test_cte_with_table_star_in_body():
         "WITH cte(a) AS (SELECT t.* FROM t) "
         "SELECT a FROM cte"
     )
-    assert "t.*" in p.columns or "a" in p.columns_aliases_names
+    assert p.columns == ["t.*"]
+    assert p.columns_aliases_names == ["a"]

test/test_getting_tables.py

@@ -943,11 +943,9 @@ def test_on_keyword_not_table_alias():
 def test_unmatched_parentheses_graceful():
     # solved: https://github.com/macbre/sql-metadata/issues/532
     # Should not raise IndexError; graceful handling of malformed SQL
-    try:
-        parser = Parser("SELECT arrayJoin(tags.key)) FROM foo")
-        _ = parser.tables
-    except (ValueError, Exception):
-        pass
+    parser = Parser("SELECT arrayJoin(tags.key)) FROM foo")
+    tables = parser.tables
+    assert isinstance(tables, list)
 
 
 def test_degraded_parse_falls_through_to_last_dialect():

test/test_query_type.py

@@ -6,25 +6,25 @@
 def test_insert_query():
     queries = [
         "INSERT IGNORE /* foo */ INTO bar VALUES (1, '123', '2017-01-01');",
-        "/* foo */ INSERT IGNORE INTO bar VALUES (1, '123', '2017-01-01');"
-        "-- foo\nINSERT IGNORE INTO bar VALUES (1, '123', '2017-01-01');"
+        "/* foo */ INSERT IGNORE INTO bar VALUES (1, '123', '2017-01-01');",
+        "-- foo\nINSERT IGNORE INTO bar VALUES (1, '123', '2017-01-01');",
         "# foo\nINSERT IGNORE INTO bar VALUES (1, '123', '2017-01-01');",
     ]
 
     for query in queries:
-        assert "INSERT" == Parser(query).query_type
+        assert QueryType.INSERT == Parser(query).query_type
 
 
 def test_select_query():
     queries = [
         "SELECT /* foo */ foo FROM bar",
-        "/* foo */ SELECT foo FROM bar"
-        "-- foo\nSELECT foo FROM bar"
+        "/* foo */ SELECT foo FROM bar",
+        "-- foo\nSELECT foo FROM bar",
         "# foo\nSELECT foo FROM bar",
     ]
 
     for query in queries:
-        assert "SELECT" == Parser(query).query_type
+        assert QueryType.SELECT == Parser(query).query_type
 
 
 def test_delete_query():
@@ -35,7 +35,7 @@ def test_delete_query():
 
     for query in queries:
         for comment in ["", "/* foo */", "\n--foo\n", "\n# foo\n"]:
-            assert "DELETE" == Parser(query.format(comment)).query_type
+            assert QueryType.DELETE == Parser(query.format(comment)).query_type
 
 
 def test_drop_table_query():
@@ -45,7 +45,7 @@ def test_drop_table_query():
 
     for query in queries:
         for comment in ["", "/* foo */", "\n--foo\n", "\n# foo\n"]:
-            assert "DROP TABLE" == Parser(query.format(comment)).query_type
+            assert QueryType.DROP == Parser(query.format(comment)).query_type
 
 
 def test_unsupported_query(caplog):
@@ -171,7 +171,7 @@ def test_unrecognized_command_type():
 def test_deeply_parenthesized_query():
     """Triple-parenthesized SELECT parses correctly."""
     p = Parser("(((SELECT col FROM t)))")
-    assert p.query_type == "SELECT"
+    assert p.query_type == QueryType.SELECT
     assert p.tables == ["t"]
     assert p.columns == ["col"]