handle redshift append clause with custom dialect, clean up table extractor and add more descriptive docstrings

Author: collerek

sql_metadata/ast_parser.py

@@ -10,6 +10,7 @@
 from typing import Optional
 
 from sqlglot import exp
+from sqlglot.dialects.dialect import DialectType
 
 from sql_metadata.dialect_parser import DialectParser
 from sql_metadata.sql_cleaner import SqlCleaner
@@ -30,7 +31,7 @@ class ASTParser:
     def __init__(self, sql: str) -> None:
         self._raw_sql = sql
         self._ast: Optional[exp.Expression] = None
-        self._dialect: object = None
+        self._dialect: DialectType = None
         self._parsed = False
         self._is_replace = False
         self._cte_name_map: dict[str, str] = {}
@@ -50,7 +51,7 @@ def ast(self) -> Optional[exp.Expression]:
         return self._ast
 
     @property
-    def dialect(self) -> object:
+    def dialect(self) -> DialectType:
         """The sqlglot dialect that produced the current AST.
 
         Set as a side-effect of :attr:`ast` access.  May be ``None``

sql_metadata/dialect_parser.py

@@ -9,9 +9,12 @@ class so that callers only need to call :meth:`DialectParser.parse`.
 from typing import Any, Optional
 
 import sqlglot
-from sqlglot import Dialect, exp
+from sqlglot import exp
+from sqlglot.dialects.dialect import Dialect, DialectType
+from sqlglot.dialects.redshift import Redshift
 from sqlglot.dialects.tsql import TSQL
 from sqlglot.errors import ParseError, TokenError
+from sqlglot.parsers.redshift import RedshiftParser
 from sqlglot.tokens import Tokenizer as BaseTokenizer
 
 from sql_metadata.comments import _has_hash_variables
@@ -47,6 +50,31 @@ class Tokenizer(BaseTokenizer):
         VAR_SINGLE_TOKENS = {*BaseTokenizer.VAR_SINGLE_TOKENS, "#"}
 
 
+class _RedshiftAppendParser(RedshiftParser):
+    """Redshift parser extended with ``ALTER TABLE ... APPEND FROM``."""
+
+    def _parse_alter_table_append(self) -> "exp.Expr | None":
+        self._match_text_seq("FROM")
+        return self._parse_table()
+
+    ALTER_PARSERS = {
+        **RedshiftParser.ALTER_PARSERS,
+        "APPEND": lambda self: self._parse_alter_table_append(),
+    }
+
+
+class RedshiftAppendDialect(Redshift):
+    """Redshift dialect extended with ``ALTER TABLE ... APPEND FROM`` support.
+
+    Redshift's ``APPEND FROM`` syntax is not natively supported by sqlglot,
+    which causes the statement to degrade to ``exp.Command``.  This dialect
+    adds an ``APPEND`` entry to ``ALTER_PARSERS`` so the statement is parsed
+    as a proper ``exp.Alter`` with ``exp.Table`` nodes.
+    """
+
+    Parser = _RedshiftAppendParser
+
+
 class BracketedTableDialect(TSQL):
     """TSQL dialect for queries containing ``[bracketed]`` identifiers.
 
@@ -65,7 +93,7 @@ class BracketedTableDialect(TSQL):
 class DialectParser:
     """Detect the appropriate sqlglot dialect and parse SQL into an AST."""
 
-    def parse(self, clean_sql: str) -> tuple[exp.Expression, object]:
+    def parse(self, clean_sql: str) -> tuple[exp.Expression, DialectType]:
         """Parse *clean_sql*, returning ``(ast, dialect)``.
 
         Detects candidate dialects via heuristics, tries each in order,
@@ -110,13 +138,15 @@ def _detect_dialects(sql: str) -> list:
             return [BracketedTableDialect, None, "mysql"]
         if " UNIQUE " in upper:
             return [None, "mysql", "oracle"]
+        if "APPEND FROM" in upper:
+            return [RedshiftAppendDialect, None, "mysql"]
         return [None, "mysql"]
 
     # -- parsing ------------------------------------------------------------
 
     def _try_dialects(
         self, clean_sql: str, dialects: list
-    ) -> tuple[exp.Expression, object]:
+    ) -> tuple[exp.Expression, DialectType]:
         """Try parsing *clean_sql* with each dialect, returning the best.
 
         :returns: 2-tuple of ``(ast_node, winning_dialect)``.

sql_metadata/parser.py

@@ -287,11 +287,13 @@ def tables(self) -> List[str]:
         if self._tables is not None:
             return self._tables
         _ = self.query_type
+        ast = self._ast_parser.ast
+        assert ast is not None  # guaranteed by query_type raising on None
         cte_names = set(self.with_names)
         for placeholder in self._ast_parser.cte_name_map:
             cte_names.add(placeholder)
         extractor = TableExtractor(
-            self._ast_parser.ast,
+            ast,
             self._raw_query,
             cte_names,
             dialect=self._ast_parser.dialect,
@@ -304,7 +306,9 @@ def tables_aliases(self) -> Dict[str, str]:
         """Return the table alias mapping for this query."""
         if self._table_aliases is not None:
             return self._table_aliases
-        extractor = TableExtractor(self._ast_parser.ast)
+        ast = self._ast_parser.ast
+        assert ast is not None  # guaranteed by prior tables/query_type access
+        extractor = TableExtractor(ast)
         self._table_aliases = extractor.extract_aliases(self.tables)
         return self._table_aliases
 

sql_metadata/query_type_extractor.py

@@ -89,10 +89,15 @@ def _unwrap_parens(ast: exp.Expression) -> exp.Expression:
 
     @staticmethod
     def _resolve_command_type(root: exp.Expression) -> Optional[QueryType]:
-        """Determine query type for an opaque Command node."""
+        """Determine query type for an opaque ``exp.Command`` node.
+
+        Hive ``CREATE FUNCTION ... USING JAR ... WITH SERDEPROPERTIES``
+        is not supported by any sqlglot dialect and degrades to
+        ``exp.Command(this='CREATE', ...)``.  This fallback extracts
+        the query type from the command text so callers still get
+        ``QueryType.CREATE``.
+        """
         expression_text = str(root.this).upper() if root.this else ""
-        if expression_text == "ALTER":
-            return QueryType.ALTER
         if expression_text == "CREATE":
             return QueryType.CREATE
         return None