add optimizer debug script

Author: treysp

sqlmesh/core/renderer.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import logging
+import os
 import typing as t
 from contextlib import contextmanager
 from functools import partial
@@ -45,6 +46,29 @@
 logger = logging.getLogger(__name__)
 
 
+def _optimizer_trace_enabled(model_fqn: t.Optional[str]) -> bool:
+    """Enable verbose optimizer tracing for selected model(s).
+
+    Controlled via env vars:
+    - SQLMESH_OPTIMIZER_TRACE: truthy enables tracing
+    - SQLMESH_OPTIMIZER_TRACE_MODEL: optional comma-separated substrings to match model fqn(s)
+    """
+    from sqlmesh.utils import str_to_bool
+
+    if not str_to_bool(os.environ.get("SQLMESH_OPTIMIZER_TRACE")):
+        return False
+
+    selector = (os.environ.get("SQLMESH_OPTIMIZER_TRACE_MODEL") or "").strip()
+    if not selector:
+        return True
+
+    if not model_fqn:
+        return False
+
+    patterns = [p.strip() for p in selector.split(",") if p.strip()]
+    return any(p in model_fqn for p in patterns)
+
+
 class BaseExpressionRenderer:
     def __init__(
         self,
@@ -652,39 +676,96 @@ def _optimize_query(self, query: exp.Query, all_deps: t.Set[str]) -> exp.Query:
 
         # We don't want to normalize names in the schema because that's handled by the optimizer
         original = query
+        trace_enabled = _optimizer_trace_enabled(self._model_fqn)
         missing_deps = set()
         all_deps = all_deps - {self._model_fqn}
         should_optimize = not self.schema.empty or not all_deps
+        stage = "precheck"
 
         for dep in all_deps:
             if not self.schema.find(exp.to_table(dep)):
                 should_optimize = False
                 missing_deps.add(dep)
 
+        if trace_enabled and not should_optimize:
+            logger.info(
+                "OPTIMIZER TRACE skip model=%s missing_deps=%s",
+                self._model_fqn,
+                sorted(missing_deps),
+            )
+
         if self._model_fqn and not should_optimize and any(s.is_star for s in query.selects):
             deps = ", ".join(f"'{dep}'" for dep in sorted(missing_deps))
             self._violated_rules[InvalidSelectStarExpansion] = deps
 
         try:
             if should_optimize:
+                if trace_enabled:
+                    logger.info(
+                        "OPTIMIZER TRACE model=%s dialect=%s default_catalog=%s deps=%s",
+                        self._model_fqn,
+                        self._dialect,
+                        self._default_catalog,
+                        sorted(all_deps),
+                    )
+                    for dep in sorted(all_deps):
+                        logger.info(
+                            "OPTIMIZER TRACE dep_schema dep=%s schema=%s",
+                            dep,
+                            self.schema.find(exp.to_table(dep)),
+                        )
+                    logger.info(
+                        "OPTIMIZER TRACE stage=input sql=%s",
+                        original.sql(dialect=self._dialect, identify=True, comments=False),
+                    )
+
                 query = query.copy()
-                simplify(
-                    annotate_types(
-                        qualify(
-                            query,
-                            dialect=self._dialect,
-                            schema=self.schema,
-                            infer_schema=False,
-                            catalog=self._default_catalog,
-                            quote_identifiers=self._quote_identifiers,
-                        ),
-                        schema=self.schema,
-                        dialect=self._dialect,
-                    ),
+                stage = "qualify"
+                qualified = qualify(
+                    query,
                     dialect=self._dialect,
+                    schema=self.schema,
+                    infer_schema=False,
+                    catalog=self._default_catalog,
+                    quote_identifiers=self._quote_identifiers,
                 )
+                if trace_enabled:
+                    logger.info(
+                        "OPTIMIZER TRACE stage=qualified sql=%s",
+                        qualified.sql(dialect=self._dialect, identify=True, comments=False),
+                    )
+
+                stage = "annotate_types"
+                typed = annotate_types(qualified, schema=self.schema, dialect=self._dialect)
+                if trace_enabled:
+                    logger.info(
+                        "OPTIMIZER TRACE stage=typed sql=%s",
+                        typed.sql(dialect=self._dialect, identify=True, comments=False),
+                    )
+
+                stage = "simplify"
+                simplified = simplify(typed, dialect=self._dialect)
+                if trace_enabled:
+                    logger.info(
+                        "OPTIMIZER TRACE stage=simplified sql=%s",
+                        simplified.sql(dialect=self._dialect, identify=True, comments=False),
+                    )
+                query = t.cast(exp.Query, simplified)
         except SqlglotError as ex:
             self._violated_rules[AmbiguousOrInvalidColumn] = ex
+            if trace_enabled:
+                logger.info(
+                    "OPTIMIZER TRACE failed model=%s stage=%s error=%s",
+                    self._model_fqn,
+                    stage,
+                    ex,
+                    exc_info=True,
+                )
+                logger.info(
+                    "OPTIMIZER TRACE failed_sql model=%s sql=%s",
+                    self._model_fqn,
+                    original.sql(dialect=self._dialect, identify=True, comments=False),
+                )
 
             query = original