aetherdialect 0.1.0__py3-none-any.whl

text2sql/dialect.py

@@ -0,0 +1,1134 @@
+"""Database dialect abstraction for SQL validation and introspection.
+
+Provides AST-based and EXPLAIN-based SQL validation, join-pair extraction, CTE body parsing, and enum reflection for PostgreSQL and Databricks.
+
+The module conditionally imports pglast for PostgreSQL and exposes a Dialect base class with dialect-specific implementations.
+
+A factory function returns the active dialect based on EngineConfig.
+"""
+
+from __future__ import annotations
+
+import re
+from typing import Any
+
+from .config import DATABRICKS_TABLE_QUALIFY_SKIP_IDENTIFIERS, EngineConfig
+from .core_utils import canonicalize_sql, debug
+
+
+def _extract_join_pairs_sqlglot(sql: str) -> set[tuple[str, str]]:
+    """Extract equality join pairs from Spark SQL using sqlglot AST."""
+    import sqlglot
+    from sqlglot import exp
+
+    pairs: set[tuple[str, str]] = set()
+    try:
+        parsed = sqlglot.parse_one(sql, dialect="spark")
+    except Exception:
+        return pairs
+
+    alias_to_table: dict[str, str] = {}
+    for table in parsed.find_all(exp.Table):
+        alias = table.alias_or_name
+        name = table.name
+        alias_to_table[alias] = name
+        alias_to_table[name] = name
+
+    def _resolve(table_or_alias: str) -> str:
+        return alias_to_table.get(table_or_alias, table_or_alias)
+
+    def _add_pair(left: exp.Column, right: exp.Column) -> None:
+        lt = left.table or ""
+        lc = left.name or ""
+        rt = right.table or ""
+        rc = right.name or ""
+        if lt and lc and rt and rc:
+            la = _resolve(lt)
+            ra = _resolve(rt)
+            pair_a = f"{la}.{lc}"
+            pair_b = f"{ra}.{rc}"
+            pairs.add(tuple(sorted([pair_a, pair_b])))
+
+    for join in parsed.find_all(exp.Join):
+        on_expr = join.args.get("on")
+        if on_expr is None:
+            continue
+        for eq in on_expr.find_all(exp.EQ):
+            left = eq.this
+            right = eq.expression
+            if isinstance(left, exp.Column) and isinstance(right, exp.Column):
+                _add_pair(left, right)
+
+    return pairs
+
+
+def _extract_cte_bodies_sqlglot(sql: str) -> dict[str, str]:
+    """Extract CTE name to body SQL mapping from Spark SQL using sqlglot."""
+    import sqlglot
+    from sqlglot import exp
+
+    cte_bodies: dict[str, str] = {}
+    try:
+        parsed = sqlglot.parse_one(sql, dialect="spark")
+    except Exception:
+        return cte_bodies
+
+    for cte in parsed.find_all(exp.CTE):
+        alias = cte.alias
+        if alias and cte.this:
+            cte_bodies[alias.lower()] = cte.this.sql(dialect="spark")
+
+    return cte_bodies
+
+if EngineConfig.TYPE == "postgresql":
+
+    def _strip_schema(ident: str) -> str:
+        """Strip schema prefix from identifier and return lowercase table name."""
+        s = (ident or "").strip().lower()
+        if "." in s:
+            s = s.split(".")[-1]
+        return s
+
+    def _colref_to_parts(n: Any) -> tuple[str, str] | None:
+        """Parse AST ColumnRef node into a (table, column) tuple.
+
+        Args:
+
+            n: An AST node, expected to be a ColumnRef with one to three String fields.
+
+        Returns:
+
+            Tuple of (table_alias_or_name, column_name), or None if the node is not a valid ColumnRef or contains unsupported field types.
+        """
+        tag = getattr(n, "__class__", type("x", (), {})).__name__
+        if tag != "ColumnRef":
+            return None
+        fields = getattr(n, "fields", None)
+        if not isinstance(fields, list | tuple) or not fields:
+            return None
+        parts: list[str] = []
+        for f in fields:
+            ftag = getattr(f, "__class__", type("x", (), {})).__name__
+            if ftag != "String":
+                return None
+            sval = getattr(f, "sval", None)
+            if isinstance(sval, str) and sval:
+                parts.append(sval)
+        if len(parts) == 1:
+            return ("", _strip_schema(parts[0]))
+        if len(parts) == 2:
+            return (_strip_schema(parts[0]), _strip_schema(parts[1]))
+        if len(parts) == 3:
+            return (_strip_schema(parts[1]), _strip_schema(parts[2]))
+        return None
+
+    def _extract_eq_pairs(expr: Any, alias_to_table: dict[str, str], out: set[tuple[str, str]]) -> bool:
+        """Recursively extract equality join pairs from an AST expression.
+
+        Args:
+
+            expr: AST expression node (A_Expr or BoolExpr).
+            alias_to_table: Mapping from table alias to canonical table name.
+            out: Set to accumulate discovered (tableA.col, tableB.col) join pairs.
+
+        Returns:
+
+            True if the expression was fully traversed without unsupported constructs, False if a non-equality or disallowed expression was found.
+        """
+        tag = getattr(expr, "__class__", type("x", (), {})).__name__
+        if tag == "BoolExpr":
+            if getattr(expr, "boolop", None) != 0:
+                return False
+            args = getattr(expr, "args", None)
+            if not isinstance(args, list):
+                return False
+            for a in args:
+                if not _extract_eq_pairs(a, alias_to_table, out):
+                    return False
+            return True
+        if tag != "A_Expr":
+            return True
+        k = getattr(expr, "kind", None)
+        if isinstance(k, int):
+            if k != 0:
+                return False
+        else:
+            if getattr(k, "name", "") != "AEXPR_OP":
+                return False
+        name = getattr(expr, "name", None)
+        if not isinstance(name, list | tuple) or len(name) != 1:
+            return False
+        op = name[0]
+        if getattr(op, "__class__", type("x", (), {})).__name__ != "String":
+            return False
+        op_str = getattr(op, "sval", None)
+        if op_str != "=":
+            return True
+        lhs = _colref_to_parts(getattr(expr, "lexpr", None))
+        r = _colref_to_parts(getattr(expr, "rexpr", None))
+        if not lhs or not r:
+            return True
+        lt, lc = lhs
+        rt, rc = r
+        lt = alias_to_table.get(lt, lt) if lt else ""
+        rt = alias_to_table.get(rt, rt) if rt else ""
+        if not lt or not rt:
+            return True
+        if not lc or not rc:
+            return True
+        a = f"{_strip_schema(lt)}.{lc}"
+        b = f"{_strip_schema(rt)}.{rc}"
+        out.add(tuple(sorted((a, b))))
+        return True
+
+    def _collect_from_items(
+        fr: Any,
+    ) -> tuple[bool, dict[str, str], bool, bool, bool, bool, bool]:
+        """Collect table aliases and detect forbidden structures from a FROM clause.
+
+        Args:
+
+            fr: FROM clause node or list of FROM clause items from the AST.
+
+        Returns:
+
+            Tuple of (ok, alias_to_table, has_subquery, has_using, has_cross_join, has_self_join, has_items) where ok is False if an unsupported node type was encountered.
+        """
+        alias_to_table: dict[str, str] = {}
+        has_subquery = False
+        has_using = False
+        has_cross_join = False
+        has_self_join = False
+        seen_tables: set[str] = set()
+        ok = True
+
+        def add_alias(relname: str, alias: Any):
+            nonlocal alias_to_table, has_self_join, seen_tables
+            t = _strip_schema(relname)
+            if t in seen_tables:
+                has_self_join = True
+            seen_tables.add(t)
+            if alias is None:
+                alias_to_table[t] = t
+                return
+            an = getattr(alias, "aliasname", None)
+            if isinstance(an, str) and an:
+                alias_to_table[_strip_schema(an)] = t
+            alias_to_table[t] = t
+
+        def walk(item: Any):
+            nonlocal has_subquery, has_using, has_cross_join, ok
+            if item is None:
+                ok = False
+                return False
+            tag = getattr(item, "__class__", type("x", (), {})).__name__
+            if tag == "RangeVar":
+                add_alias(getattr(item, "relname", "") or "", getattr(item, "alias", None))
+                return True
+            if tag == "JoinExpr":
+                if getattr(item, "usingClause", None) is not None or getattr(item, "isNatural", False):
+                    has_using = True
+                join_type = getattr(item, "jointype", None)
+                if join_type is not None and str(join_type) == "JoinType.JOIN_INNER":
+                    quals = getattr(item, "quals", None)
+                    if quals is None:
+                        has_cross_join = True
+                if not walk(getattr(item, "larg", None)):
+                    return False
+                if not walk(getattr(item, "rarg", None)):
+                    return False
+                return True
+            if tag in {
+                "RangeSubselect",
+                "RangeFunction",
+                "RangeTableFunc",
+                "RangeTableSample",
+            }:
+                has_subquery = True
+                ok = False
+                return False
+            ok = False
+            return False
+
+        if fr is None:
+            return False, {}, False, False, False, False, False
+        for it in fr if isinstance(fr, list | tuple) else [fr]:
+            if not walk(it):
+                ok = False
+                break
+            return (
+                ok,
+                alias_to_table,
+                has_subquery,
+                has_using,
+                has_cross_join,
+                has_self_join,
+                True,
+            )
+        return (
+            ok,
+            alias_to_table,
+            has_subquery,
+            has_using,
+            has_cross_join,
+            has_self_join,
+            True,
+        )
+
+    def _normalized_join_pairs_from_sql_ast(
+        sql: str,
+    ) -> tuple[bool, set[tuple[str, str]]]:
+        """Extract normalized join pairs from SQL using AST parsing and handle CTEs.
+
+        Args:
+
+            sql: SQL query string to parse.
+
+        Returns:
+
+            Tuple of (ok, pairs) where ok is False if parsing failed or a forbidden construct was found, and pairs is the set of (tableA.col, tableB.col) tuples sorted so the lexicographically smaller element comes first.
+        """
+        s = canonicalize_sql(sql)
+        if not s:
+            debug("[dialect.normalized_join_pairs_from_sql_ast] empty_sql")
+            return False, set()
+        try:
+            import pglast
+        except ImportError:
+            debug("[dialect.normalized_join_pairs_from_sql_ast] pglast_unavailable")
+            return False, set()
+        parse_sql = getattr(pglast, "parse_sql", None)
+        if parse_sql is None:
+            debug("[dialect.normalized_join_pairs_from_sql_ast] pglast_unavailable")
+            return False, set()
+        try:
+            stmts = parse_sql(s)
+        except Exception as e:
+            debug(f"[dialect.normalized_join_pairs_from_sql_ast] parse_failed: {e}")
+            return False, set()
+        if not stmts:
+            debug("[dialect.normalized_join_pairs_from_sql_ast] no_statements")
+            return False, set()
+        first = stmts[0]
+        root = getattr(first, "stmt", None)
+        if root is None:
+            debug("[dialect.normalized_join_pairs_from_sql_ast] no_stmt_attr")
+            return False, set()
+        if getattr(root, "__class__", type("x", (), {})).__name__ != "SelectStmt":
+            debug("[dialect.normalized_join_pairs_from_sql_ast] not_select_stmt")
+            return False, set()
+
+        with_clause = getattr(root, "withClause", None)
+        cte_names = set()
+        all_pairs: set[tuple[str, str]] = set()
+
+        if with_clause is not None:
+            ctes = getattr(with_clause, "ctes", [])
+            for cte in ctes:
+                cte_name = getattr(getattr(cte, "ctename", None), "sval", None) or getattr(cte, "ctename", "")
+                if cte_name:
+                    cte_names.add(cte_name.lower())
+
+                cte_query = getattr(cte, "ctequery", None)
+                if cte_query is None:
+                    continue
+
+                cte_from = getattr(cte_query, "fromClause", None)
+                ok_cte_from, cte_alias_map, has_subq, has_using, _, _, _ = _collect_from_items(cte_from)
+                if not ok_cte_from or has_subq or has_using:
+                    return False, set()
+
+                def walk_cte_joins(item: Any, alias_map: dict[str, str], pairs: set[tuple[str, str]]) -> bool:
+                    if getattr(item, "__class__", type("x", (), {})).__name__ == "JoinExpr":
+                        if getattr(item, "quals", None) is None:
+                            return False
+                        if not _extract_eq_pairs(getattr(item, "quals", None), alias_map, pairs):
+                            return False
+                        if not walk_cte_joins(getattr(item, "larg", None), alias_map, pairs):
+                            return False
+                        if not walk_cte_joins(getattr(item, "rarg", None), alias_map, pairs):
+                            return False
+                    return True
+
+                if cte_from:
+                    for it in cte_from if isinstance(cte_from, list | tuple) else [cte_from]:
+                        if not walk_cte_joins(it, cte_alias_map, all_pairs):
+                            return False, set()
+
+                cte_where = getattr(cte_query, "whereClause", None)
+                if cte_where is not None:
+                    tmp: set[tuple[str, str]] = set()
+                    if not _extract_eq_pairs(cte_where, cte_alias_map, tmp):
+                        return False, set()
+                    all_pairs |= tmp
+
+        fr = getattr(root, "fromClause", None)
+        ok_from, alias_to_table, has_subq, has_using, _, _, _ = _collect_from_items(fr)
+        if not ok_from:
+            debug("[dialect.normalized_join_pairs_from_sql_ast] collect_from_failed")
+            return False, set()
+        if has_subq:
+            debug("[dialect.normalized_join_pairs_from_sql_ast] has_subquery")
+            return False, set()
+        if has_using:
+            debug("[dialect.normalized_join_pairs_from_sql_ast] has_using")
+            return False, set()
+
+        def walk_join_ons(item: Any) -> bool:
+            if getattr(item, "__class__", type("x", (), {})).__name__ == "JoinExpr":
+                quals = getattr(item, "quals", None)
+                if quals is None:
+                    debug("[dialect.normalized_join_pairs_from_sql_ast] missing_quals")
+                    return False
+                if not _extract_eq_pairs(quals, alias_to_table, all_pairs):
+                    debug("[dialect.normalized_join_pairs_from_sql_ast] extract_pairs_failed")
+                    return False
+                if not walk_join_ons(getattr(item, "larg", None)):
+                    return False
+                if not walk_join_ons(getattr(item, "rarg", None)):
+                    return False
+            return True
+
+        for it in fr if isinstance(fr, list | tuple) else []:
+            if not walk_join_ons(it):
+                return False, set()
+
+        where = getattr(root, "whereClause", None)
+        if where is not None:
+            tmp: set[tuple[str, str]] = set()
+            if not _extract_eq_pairs(where, alias_to_table, tmp):
+                debug("[dialect.normalized_join_pairs_from_sql_ast] where_extract_failed")
+                return False, set()
+            all_pairs |= tmp
+
+        if getattr(root, "groupClause", None) is not None:
+            for g in getattr(root, "groupClause", []) or []:
+                if getattr(g, "__class__", type("x", (), {})).__name__ == "SelectStmt":
+                    return False, set()
+        if getattr(root, "sortClause", None) is not None:
+            for so in getattr(root, "sortClause", []) or []:
+                if getattr(so, "__class__", type("x", (), {})).__name__ == "SelectStmt":
+                    return False, set()
+        return True, all_pairs
+
+    def _validate_cte_bodies(with_clause: Any) -> tuple[bool, str]:
+        """Validate CTE bodies against structural restrictions.
+
+        Forbids recursive CTEs, subqueries, window functions, set operations, and CASE expressions inside any CTE body.
+
+        Args:
+
+            with_clause: AST WithClause node or None.
+
+        Returns:
+
+            Tuple of (ok, error_code). ok is True when all CTEs pass validation. error_code is an empty string on success or a short string token describing the first violation found.
+        """
+        if with_clause is None:
+            return True, ""
+
+        if getattr(with_clause, "recursive", False):
+            return False, "cte_recursive"
+
+        ctes = getattr(with_clause, "ctes", [])
+        if not ctes:
+            return True, ""
+
+        for cte in ctes:
+            cte_query = getattr(cte, "ctequery", None)
+            if cte_query is None:
+                return False, "cte_malformed"
+
+            def walk_cte(n):
+                tag = getattr(n, "__class__", type("x", (), {})).__name__
+                if tag in {"RangeSubselect", "SubLink"}:
+                    if tag == "SubLink":
+                        sublink_type = getattr(n, "subLinkType", None)
+                        if sublink_type is not None and sublink_type == 0:
+                            return "cte_contains_exists"
+                    return "cte_contains_subquery"
+                if tag == "SetOperationStmt":
+                    return "cte_contains_set_op"
+                if tag == "WindowDef":
+                    return "cte_contains_window"
+                if tag == "CaseExpr":
+                    return "cte_contains_case"
+                try:
+                    attrs = vars(n)
+                except TypeError:
+                    return None
+                for attr in attrs.values():
+                    if isinstance(attr, list):
+                        for x in attr:
+                            if hasattr(x, "__class__"):
+                                err = walk_cte(x)
+                                if err:
+                                    return err
+                    elif hasattr(attr, "__class__"):
+                        err = walk_cte(attr)
+                        if err:
+                            return err
+                return None
+
+            err = walk_cte(cte_query)
+            if err:
+                return False, err
+
+        return True, ""
+
+    def _ast_structural_valid(sql: str) -> tuple[bool, str]:
+        """Validate SQL structure using the pglast AST.
+
+        Checks that the SQL is a single SELECT statement free of subqueries, window functions, CASE expressions, CROSS JOINs, self-joins, USING clauses, EXISTS, LATERAL, and set operations. Also validates any CTE bodies.
+
+        Args:
+
+            sql: SQL query string to validate.
+
+        Returns:
+
+            Tuple of (ok, error_code). ok is True when the SQL passes all structural checks. error_code is an empty string on success or a short token identifying the first violation.
+        """
+        try:
+            import pglast
+        except ImportError:
+            return False, "ast_unavailable"
+        parse_sql = getattr(pglast, "parse_sql", None)
+        if parse_sql is None:
+            return False, "ast_unavailable"
+
+        try:
+            stmts = parse_sql(canonicalize_sql(sql))
+        except Exception:
+            return False, "ast_parse_failed"
+
+        if not stmts or len(stmts) != 1:
+            return False, "multiple_statements"
+
+        root = getattr(stmts[0], "stmt", None)
+        if root is None:
+            return False, "no_root"
+
+        if getattr(root, "__class__", type("x", (), {})).__name__ != "SelectStmt":
+            return False, "not_select"
+
+        with_clause = getattr(root, "withClause", None)
+        has_cte = with_clause is not None
+
+        if has_cte:
+            ok, err = _validate_cte_bodies(with_clause)
+            if not ok:
+                return False, err
+
+        fr = getattr(root, "fromClause", None)
+        if fr is not None:
+            _, _, has_subq, has_using, has_cross, has_self, _ = _collect_from_items(fr)
+            if has_subq:
+                return False, "subquery_not_allowed"
+            if has_using:
+                return False, "using_not_allowed"
+            if has_cross:
+                return False, "cross_join_not_allowed"
+            if has_self:
+                return False, "self_join_not_allowed"
+
+        has_window = False
+        has_case = False
+        has_exists = False
+        has_lateral = False
+
+        def walk(n, check_window=True):
+            nonlocal has_window, has_case, has_exists, has_lateral
+            tag = getattr(n, "__class__", type("x", (), {})).__name__
+            if tag in {
+                "RangeSubselect",
+                "SubLink",
+                "SetOperationStmt",
+            }:
+                if tag == "SubLink":
+                    sublink_type = getattr(n, "subLinkType", None)
+                    if sublink_type is not None and sublink_type == 0:
+                        has_exists = True
+                return False
+            if tag == "CaseExpr":
+                has_case = True
+                return False
+            if tag == "RangeFunction":
+                is_lateral = getattr(n, "lateral", False)
+                if is_lateral:
+                    has_lateral = True
+                    return False
+            if check_window and tag == "WindowDef":
+                has_window = True
+                return False
+            if check_window and tag == "FuncCall":
+                over = getattr(n, "over", None)
+                if over is not None:
+                    has_window = True
+                    return False
+
+            try:
+                attrs = vars(n)
+            except TypeError:
+                return True
+
+            for attr in attrs.values():
+                if isinstance(attr, list):
+                    for x in attr:
+                        if hasattr(x, "__class__") and not walk(x, check_window):
+                            return False
+                elif hasattr(attr, "__class__"):
+                    if not walk(attr, check_window):
+                        return False
+            return True
+
+        if not walk(root):
+            if has_window:
+                return False, "window_not_allowed"
+            if has_case:
+                return False, "case_when_not_allowed"
+            if has_exists:
+                return False, "exists_not_allowed"
+            if has_lateral:
+                return False, "lateral_not_allowed"
+            return False, "forbidden_structure"
+
+        return True, ""
+
+
+class Dialect:
+    """Base dialect interface for database-specific operations."""
+
+    name: str = "base"
+
+    def __init__(self, config):
+        """Initialize the dialect with a runtime configuration object.
+
+        Args:
+
+            config: Runtime config instance, such as PostgresRuntimeConfig or DatabricksRuntimeConfig, providing connection details.
+        """
+        self.config = config
+
+    def ast_validate(self, sql: str) -> tuple[bool, str]:
+        """Validate SQL structure via AST parsing without a database connection.
+
+        Args:
+
+            sql: SQL query string to validate.
+
+        Returns:
+
+            Tuple of (ok, error_code). ok is True when the SQL passes structural checks.
+        """
+        raise NotImplementedError
+
+    def extract_join_pairs(self, sql: str) -> tuple[bool, set[tuple[str, str]]]:
+        """Extract normalized join column pairs from the SQL AST.
+
+        Args:
+
+            sql: SQL query string to parse.
+
+        Returns:
+
+            Tuple of (ok, pairs) where pairs is a set of sorted (tableA.col, tableB.col) tuples found in ON conditions.
+        """
+        raise NotImplementedError
+
+    def extract_cte_bodies(self, sql: str) -> dict[str, str]:
+        """Extract the SQL body for each named CTE from a WITH clause.
+
+        Args:
+
+            sql: SQL query string that may contain a WITH clause.
+
+        Returns:
+
+            Dictionary mapping lowercase CTE name to its inner SELECT SQL string.
+        """
+        raise NotImplementedError
+
+    def explain_sql(self, engine: Any, sql: str, params: dict[str, Any] | None = None) -> tuple[bool, str]:
+        """Test SQL executability by running EXPLAIN against the database.
+
+        Args:
+
+            engine: SQLAlchemy Engine connected to the target database.
+            sql: SQL query string to explain.
+            params: Optional bind-parameter values keyed by placeholder name.
+
+        Returns:
+
+            Tuple of (ok, error_message). ok is True when EXPLAIN succeeds.
+        """
+        raise NotImplementedError
+
+    def reflect_enums(self, engine: Any, schema_name: str) -> dict[str, list[str]]:
+        """Discover enum types and their allowed values from the database.
+
+        Args:
+
+            engine: SQLAlchemy Engine connected to the target database.
+            schema_name: Database schema name to introspect.
+
+        Returns:
+
+            Dictionary mapping enum type name to an ordered list of its values.
+        """
+        raise NotImplementedError
+
+
+class PostgresDialect(Dialect):
+    """PostgreSQL-specific dialect implementation."""
+
+    name: str = "postgresql"
+
+    def __init__(self, config):
+        """Initialize the PostgreSQL dialect and create a SQLAlchemy engine.
+
+        Args:
+
+            config: PostgresRuntimeConfig instance providing connection credentials.
+        """
+        super().__init__(config)
+        from sqlalchemy import create_engine
+
+        self.engine = create_engine(config.db_url())
+
+    def ast_validate(self, sql: str) -> tuple[bool, str]:
+        """Validate SQL structure via AST parsing."""
+        return _ast_structural_valid(sql)
+
+    def extract_join_pairs(self, sql: str) -> tuple[bool, set[tuple[str, str]]]:
+        """Extract normalized join pairs from SQL AST."""
+        return _normalized_join_pairs_from_sql_ast(sql)
+
+    def extract_cte_bodies(self, sql: str) -> dict[str, str]:
+        """Extract SQL body for each CTE from a WITH clause using the pglast AST."""
+        try:
+            import pglast
+        except ImportError:
+            return self._extract_cte_bodies_depth_tracking(sql)
+        parse_sql = getattr(pglast, "parse_sql", None)
+        if parse_sql is None:
+            return self._extract_cte_bodies_depth_tracking(sql)
+        cte_bodies: dict[str, str] = {}
+        try:
+            stmts = parse_sql(canonicalize_sql(sql))
+            if not stmts:
+                return self._extract_cte_bodies_depth_tracking(sql)
+            root = getattr(stmts[0], "stmt", None)
+            if root is None:
+                return self._extract_cte_bodies_depth_tracking(sql)
+            with_clause = getattr(root, "withClause", None)
+            if with_clause is None:
+                return cte_bodies
+            ctes = getattr(with_clause, "ctes", [])
+            for cte in ctes:
+                cte_name = getattr(getattr(cte, "ctename", None), "sval", None) or getattr(cte, "ctename", "")
+                cte_query = getattr(cte, "ctequery", None)
+                if cte_name and cte_query:
+                    try:
+                        from pglast import prettify
+
+                        cte_sql = prettify(cte_query)
+                        cte_bodies[cte_name.lower()] = cte_sql
+                    except Exception as exc:
+                        debug(f"[dialect.extract_cte_bodies] prettify failed for CTE '{cte_name}': {exc}")
+            if cte_bodies:
+                debug(
+                    f"[dialect.extract_cte_bodies] pglast extracted {len(cte_bodies)} CTEs: {list(cte_bodies.keys())}"
+                )
+                return cte_bodies
+        except Exception as e:
+            debug(f"[dialect.extract_cte_bodies] pglast failed: {e}, falling back to depth-tracking")
+        return self._extract_cte_bodies_depth_tracking(sql)
+
+    def _extract_cte_bodies_depth_tracking(self, sql: str) -> dict[str, str]:
+        """Extract CTE bodies using depth-tracking character parser."""
+        cte_bodies: dict[str, str] = {}
+        sql_upper = sql.upper()
+        if not sql_upper.strip().startswith("WITH"):
+            return cte_bodies
+        with_match = re.match(r"\s*WITH\s+", sql, re.IGNORECASE)
+        if not with_match:
+            return cte_bodies
+        pos = with_match.end()
+        while pos < len(sql):
+            while pos < len(sql) and sql[pos] in " \t\n\r,":
+                pos += 1
+            if pos >= len(sql):
+                break
+            name_match = re.match(r"(\w+)\s+AS\s*\(", sql[pos:], re.IGNORECASE)
+            if not name_match:
+                break
+            cte_name = name_match.group(1).lower()
+            pos += name_match.end()
+            start_pos = pos
+            depth = 1
+            in_string = False
+            string_char = None
+            while pos < len(sql) and depth > 0:
+                c = sql[pos]
+                if in_string:
+                    if c == string_char and (pos + 1 >= len(sql) or sql[pos + 1] != string_char):
+                        in_string = False
+                    elif c == string_char and pos + 1 < len(sql) and sql[pos + 1] == string_char:
+                        pos += 1
+                else:
+                    if c in ("'", '"'):
+                        in_string = True
+                        string_char = c
+                    elif c == "(":
+                        depth += 1
+                    elif c == ")":
+                        depth -= 1
+                pos += 1
+            if depth == 0:
+                cte_sql = sql[start_pos : pos - 1].strip()
+                cte_bodies[cte_name] = cte_sql
+            rest = sql[pos:].strip().upper()
+            if rest.startswith("SELECT") or rest.startswith("(SELECT"):
+                break
+        debug(
+            f"[dialect.extract_cte_bodies] depth-tracking extracted {len(cte_bodies)} CTEs: {list(cte_bodies.keys())}"
+        )
+        return cte_bodies
+
+    def explain_sql(self, engine: Any, sql: str, params: dict[str, Any] | None = None) -> tuple[bool, str]:
+        """Test SQL executability via PostgreSQL EXPLAIN."""
+        from sqlalchemy import text
+
+        try:
+            with engine.connect() as conn:
+                conn.execute(text(f"EXPLAIN {sql}"), params or {})
+            return True, ""
+        except Exception as e:
+            return False, str(e)
+
+    def reflect_enums(self, engine: Any, schema_name: str) -> dict[str, list[str]]:
+        """Discover PostgreSQL enum types and their values."""
+        if EngineConfig.TYPE != "postgresql":
+            return {}
+
+        from sqlalchemy import text
+
+        enum_values: dict[str, list[str]] = {}
+        try:
+            with engine.connect() as conn:
+                enum_query = text("""
+                    SELECT t.typname, e.enumlabel
+                    FROM pg_type t
+                    JOIN pg_enum e ON t.oid = e.enumtypid
+                    JOIN pg_namespace n ON t.typnamespace = n.oid
+                    WHERE n.nspname = :schema
+                    ORDER BY t.typname, e.enumsortorder
+                """)
+                for row in conn.execute(enum_query, {"schema": schema_name}):
+                    enum_name = row[0]
+                    if enum_name not in enum_values:
+                        enum_values[enum_name] = []
+                    enum_values[enum_name].append(row[1])
+        except Exception as exc:
+            debug(f"[dialect.reflect_enums] failed to reflect enums: {exc}")
+        return enum_values
+
+
+class DatabricksDialect(Dialect):
+    """Databricks Spark dialect with Spark EXPLAIN validation."""
+
+    name: str = "databricks"
+
+    def __init__(self, config):
+        """Initialize the Databricks dialect with optional native connection.
+
+        When ``config.has_native_connection()`` is True, creates a connection via
+        ``databricks-sql-connector`` and skips SparkSession. Otherwise falls
+        back to PySpark.
+
+        Args:
+
+            config: DatabricksRuntimeConfig instance providing catalog and schema.
+
+        Raises:
+
+            RuntimeError: If neither a native connection nor a SparkSession can be obtained.
+        """
+        super().__init__(config)
+
+        self.connection = None
+        self.spark = None
+
+        if config.has_native_connection():
+            try:
+                from databricks import sql as dbsql
+
+                self.connection = dbsql.connect(
+                    server_hostname=config.SERVER_HOSTNAME,
+                    http_path=config.HTTP_PATH,
+                    access_token=config.ACCESS_TOKEN,
+                )
+            except Exception as e:
+                raise RuntimeError(f"Failed to connect via databricks-sql-connector: {e}") from e
+        else:
+            try:
+                from pyspark.sql import SparkSession
+
+                self.spark = SparkSession.builder.getOrCreate()
+            except Exception as e:
+                raise RuntimeError(f"Failed to get SparkSession: {e}") from e
+
+    def ast_validate(self, sql: str) -> tuple[bool, str]:
+        """Validate SQL syntax via Spark EXPLAIN on native connection or PySpark."""
+        qualified = self._qualify_table_references(sql)
+
+        if self.connection is not None:
+            try:
+                cursor = self.connection.cursor()
+                cursor.execute(f"EXPLAIN {qualified}")
+                cursor.fetchall()
+                cursor.close()
+            except Exception as e:
+                err_str = str(e).lower()
+                if "syntax" in err_str or "parse" in err_str:
+                    return False, f"Databricks syntax error: {e}"
+                debug(f"[dialect.ast_validate] non-syntax Databricks error (treating as valid): {e}")
+            return True, ""
+
+        try:
+            explain_df = self.spark.sql(f"EXPLAIN {qualified}")
+            explain_df.collect()
+        except Exception as e:
+            err_str = str(e).lower()
+            if "syntax" in err_str or "parse" in err_str:
+                return False, f"Spark syntax error: {e}"
+            debug(f"[dialect.ast_validate] non-syntax Spark error (treating as valid): {e}")
+
+        return True, ""
+
+    def extract_join_pairs(self, sql: str) -> tuple[bool, set[tuple[str, str]]]:
+        """Extract join pairs from Spark SQL using sqlglot."""
+        try:
+            pairs = _extract_join_pairs_sqlglot(sql)
+            return True, pairs
+        except Exception as e:
+            debug(f"[dialect.extract_join_pairs] failed: {e}")
+            return False, set()
+
+    def extract_cte_bodies(self, sql: str) -> dict[str, str]:
+        """Extract SQL body for each CTE from a WITH clause using sqlglot."""
+        cte_bodies = _extract_cte_bodies_sqlglot(sql)
+        if cte_bodies:
+            debug(
+                f"[dialect.extract_cte_bodies] sqlglot extracted {len(cte_bodies)} CTEs: {list(cte_bodies.keys())}"
+            )
+            return cte_bodies
+        return self._extract_cte_bodies_depth_tracking(sql)
+
+    def _extract_cte_bodies_depth_tracking(self, sql: str) -> dict[str, str]:
+        """Extract CTE bodies using depth-tracking character parser."""
+        cte_bodies: dict[str, str] = {}
+        sql_upper = sql.upper()
+        if not sql_upper.strip().startswith("WITH"):
+            return cte_bodies
+        with_match = re.match(r"\s*WITH\s+", sql, re.IGNORECASE)
+        if not with_match:
+            return cte_bodies
+        pos = with_match.end()
+        while pos < len(sql):
+            while pos < len(sql) and sql[pos] in " \t\n\r,":
+                pos += 1
+            if pos >= len(sql):
+                break
+            name_match = re.match(r"(\w+)\s+AS\s*\(", sql[pos:], re.IGNORECASE)
+            if not name_match:
+                break
+            cte_name = name_match.group(1).lower()
+            pos += name_match.end()
+            start_pos = pos
+            depth = 1
+            in_string = False
+            string_char = None
+            while pos < len(sql) and depth > 0:
+                c = sql[pos]
+                if in_string:
+                    if c == string_char and (pos + 1 >= len(sql) or sql[pos + 1] != string_char):
+                        in_string = False
+                    elif c == string_char and pos + 1 < len(sql) and sql[pos + 1] == string_char:
+                        pos += 1
+                else:
+                    if c in ("'", '"'):
+                        in_string = True
+                        string_char = c
+                    elif c == "(":
+                        depth += 1
+                    elif c == ")":
+                        depth -= 1
+                pos += 1
+            if depth == 0:
+                cte_sql = sql[start_pos : pos - 1].strip()
+                cte_bodies[cte_name] = cte_sql
+            rest = sql[pos:].strip().upper()
+            if rest.startswith("SELECT") or rest.startswith("(SELECT"):
+                break
+        debug(
+            f"[dialect.extract_cte_bodies] depth-tracking extracted {len(cte_bodies)} CTEs: {list(cte_bodies.keys())}"
+        )
+        return cte_bodies
+
+    def explain_sql(self, engine: Any, sql: str, params: dict[str, Any] | None = None) -> tuple[bool, str]:
+        """Test SQL executability via Spark EXPLAIN."""
+        try:
+            qualified = self._qualify_table_references(sql)
+            explain_df = self.spark.sql(f"EXPLAIN {qualified}")
+            explain_df.collect()
+            return True, ""
+        except Exception as e:
+            return False, str(e)
+
+    def reflect_enums(self, engine: Any, schema_name: str) -> dict[str, list[str]]:
+        """Databricks does not support native enums."""
+        return {}
+
+    def prepare_for_execution(self, sql: str) -> str:
+        """Qualify table references with catalog.schema for Spark execution."""
+        return self._qualify_table_references(sql)
+
+    def _qualify_table_references(self, sql: str) -> str:
+        """Add catalog.schema prefix to physical table references only."""
+        catalog = self.config.CATALOG
+        schema_name = self.config.SCHEMA
+
+        pattern = r'\b(FROM|JOIN)\s+(["`]?)(\w+)\2\b'
+
+        def replace_table(match: re.Match[str]) -> str:
+            keyword = match.group(1)
+            quote = match.group(2) or "`"
+            table = match.group(3)
+            if table.lower() in DATABRICKS_TABLE_QUALIFY_SKIP_IDENTIFIERS:
+                return match.group(0)
+            return (
+                f"{keyword} {quote}{catalog}{quote}.{quote}{schema_name}{quote}."
+                f"{quote}{table}{quote}"
+            )
+
+        return re.sub(pattern, replace_table, sql, flags=re.IGNORECASE)
+
+    def execute_sql_spark(self, sql: str, sql_already_spark: bool = False) -> list[tuple]:
+        """Execute SQL on Databricks via native connector or Spark."""
+        spark_sql = self._qualify_table_references(sql)
+
+        if self.config.DEBUG:
+            print(f"[DEBUG] Executing on Databricks:\n{spark_sql}")
+
+        if self.connection is not None:
+            cursor = self.connection.cursor()
+            cursor.execute(spark_sql)
+            rows = cursor.fetchall()
+            cursor.close()
+            return [tuple(row) for row in rows]
+
+        df = self.spark.sql(spark_sql)
+        rows = df.collect()
+
+        return [tuple(row) for row in rows]
+
+
+def render_date_diff_expr(
+    dialect_type: str, left_expr: str, op: str, unit: str, amount: int
+) -> str:
+    """Render a date-difference filter expression for dialect.
+
+    Args:
+
+        dialect_type: One of 'postgresql' or 'databricks'.
+        left_expr: SQL expression for date subtraction (e.g. col1 - col2).
+        op: Comparison operator.
+        unit: 'day', 'week', 'month', or 'year'.
+        amount: Numeric amount for the interval.
+
+    Returns:
+
+        SQL predicate string.
+    """
+    if dialect_type == "postgresql":
+        interval_suffix = "s" if amount != 1 else ""
+        return f"({left_expr}) {op} INTERVAL '{amount} {unit}{interval_suffix}'"
+    if unit == "day":
+        return f"({left_expr}) {op} {amount}"
+    if unit == "week":
+        return f"({left_expr}) {op} {amount * 7}"
+    if unit == "month":
+        return f"({left_expr}) {op} {amount * 30}"
+    if unit == "year":
+        return f"({left_expr}) {op} {amount * 365}"
+    return f"({left_expr}) {op} {amount}"
+
+
+def render_date_window_expr(dialect_type: str, column: str, op: str, unit: str, offset: int) -> str:
+    """Render a dialect-specific date window filter expression.
+
+    Args:
+
+        dialect_type: One of 'postgresql' or 'databricks'.
+        column: Column reference string to compare against the date boundary.
+        op: Comparison operator, for example '>=' or '<'.
+        unit: Date unit: 'day', 'week', 'month', or 'year'.
+        offset: Number of units in the past. 0 means the start of the current unit.
+
+    Returns:
+
+        SQL expression string suitable for use in a WHERE clause.
+    """
+    if offset == 0:
+        if dialect_type == "postgresql":
+            return f"{column} {op} DATE_TRUNC('{unit}', CURRENT_DATE)"
+        return f"{column} {op} date_trunc('{unit}', current_date())"
+    if dialect_type == "postgresql":
+        suffix = "s" if offset != 1 else ""
+        return f"{column} {op} CURRENT_DATE - INTERVAL '{offset} {unit}{suffix}'"
+    if unit == "day":
+        return f"{column} {op} date_sub(current_date(), {offset})"
+    if unit == "week":
+        return f"{column} {op} date_sub(current_date(), {offset * 7})"
+    if unit == "month":
+        return f"{column} {op} add_months(current_date(), -{offset})"
+    if unit == "year":
+        return f"{column} {op} add_months(current_date(), -{offset * 12})"
+    if unit == "hour":
+        return f"{column} {op} (current_timestamp() - INTERVAL '{offset} HOURS')"
+    if unit == "minute":
+        return f"{column} {op} (current_timestamp() - INTERVAL '{offset} MINUTES')"
+    if unit == "second":
+        return f"{column} {op} (current_timestamp() - INTERVAL '{offset} SECONDS')"
+    return f"{column} {op} date_sub(current_date(), {offset})"
+
+
+def get_dialect(engine_type: str = None, config=None) -> Dialect:
+    """Return a Dialect instance for the specified engine type.
+
+    Args:
+
+        engine_type: One of 'postgresql' or 'databricks'. Defaults to EngineConfig.TYPE.
+        config: Runtime config object. Defaults to EngineConfig.RUNTIME.
+
+    Returns:
+
+        Configured Dialect subclass instance ready for use.
+
+    Raises:
+
+        ValueError: If engine_type is not a supported dialect name.
+    """
+    if engine_type is None:
+        engine_type = EngineConfig.TYPE
+    if config is None:
+        config = EngineConfig.RUNTIME
+
+    if engine_type == "postgresql":
+        return PostgresDialect(config)
+    elif engine_type == "databricks":
+        return DatabricksDialect(config)
+    else:
+        raise ValueError(f"Unsupported dialect: {engine_type}")