fix!: Runs should trigger materialization result

dagster_sqlmesh/console.py

@@ -10,7 +10,7 @@
 from sqlmesh.core.console import Console
 from sqlmesh.core.context_diff import ContextDiff
 from sqlmesh.core.environment import EnvironmentNamingInfo
-from sqlmesh.core.plan import EvaluatablePlan, PlanBuilder
+from sqlmesh.core.plan import EvaluatablePlan, Plan as SQLMeshPlan, PlanBuilder
 from sqlmesh.core.snapshot import Snapshot, SnapshotChangeCategory, SnapshotInfoLike
 from sqlmesh.core.table_diff import RowDiff, SchemaDiff, TableDiff
 from sqlmesh.utils.concurrency import NodeExecutionFailedError
@@ -219,6 +219,10 @@ class PrintEnvironments(BaseConsoleEvent):
 class ShowTableDiffSummary(BaseConsoleEvent):
     table_diff: TableDiff
 
+@dataclass(kw_only=True)
+class PlanBuilt(BaseConsoleEvent):
+    plan: SQLMeshPlan 
+
 ConsoleEvent = (
     StartPlanEvaluation
     | StopPlanEvaluation
@@ -263,6 +267,7 @@ class ShowTableDiffSummary(BaseConsoleEvent):
     | ConsoleException
     | PrintEnvironments
     | ShowTableDiffSummary
+    | PlanBuilt
 )
 
 ConsoleEventHandler = t.Callable[[ConsoleEvent], None]
@@ -424,7 +429,22 @@ def add_handler(self, handler: ConsoleEventHandler) -> str:
 
     def remove_handler(self, handler_id: str) -> None:
         del self._handlers[handler_id]
-
+
+    def plan(self, plan_builder: PlanBuilder, auto_apply: bool, default_catalog: str | None, no_diff: bool = False, no_prompts: bool = False) -> None:
+        """Plan is not a console event. This triggers building of a plan and
+        applying said plan
+
+        This method is called by SQLMesh to start the plan process (when you
+        call Context#plan)
+
+        This overriden method ignores the options passed in at this time
+        """
+
+        plan_builder.apply()
+
+    def capture_built_plan(self, plan: SQLMeshPlan) -> None:
+        """Capture the built plan and publish a PlanBuilt event."""
+        self.publish(PlanBuilt(plan=plan))
 
 class EventConsole(IntrospectingConsole):
     """

dagster_sqlmesh/controller/base.py

@@ -9,7 +9,7 @@
 from sqlmesh.core.console import set_console
 from sqlmesh.core.context import Context
 from sqlmesh.core.model import Model
-from sqlmesh.core.plan import PlanBuilder
+from sqlmesh.core.plan import Plan as SQLMeshPlan, PlanBuilder
 from sqlmesh.utils.dag import DAG
 from sqlmesh.utils.date import TimeLike
 
@@ -19,7 +19,6 @@
     ConsoleEventHandler,
     ConsoleException,
     EventConsole,
-    Plan,
     SnapshotCategorizer,
 )
 from dagster_sqlmesh.events import ConsoleGenerator
@@ -179,16 +178,7 @@ def run_sqlmesh_thread(
         ) -> None:
             logger.debug("dagster-sqlmesh: thread started")
 
-            def auto_execute_plan(event: ConsoleEvent):
-                if isinstance(event, Plan):
-                    try:
-                        event.plan_builder.apply()
-                    except Exception as e:
-                        controller.console.exception(e)
-                return None
-
             try:
-                controller.console.add_handler(auto_execute_plan)
                 builder = t.cast(
                     PlanBuilder,
                     context.plan_builder(
@@ -381,6 +371,11 @@ def non_external_models_dag(self) -> t.Iterable[tuple[Model, set[str]]]:
                 continue
             yield (model, deps)
 
+
+class ContextApplyFunction(t.Protocol):
+    def __call__(self, plan: SQLMeshPlan, *args: t.Any, **kwargs: t.Any) -> None:
+        ...
+
 class SQLMeshController(t.Generic[ContextCls]):
     """Allows control of sqlmesh via a python interface. It is not suggested to
     use the constructor of this class directly, but instead use the provided
@@ -480,7 +475,21 @@ def _create_context(self) -> ContextCls:
         if self.config.sqlmesh_config:
             options["config"] = self.config.sqlmesh_config
         set_console(self.console)
-        return self._context_factory(**options)
+        context = self._context_factory(**options)
+
+        # As part of the context, it specifies a method "apply" that we would
+        # like to introspect. To do so we replace "apply" with a wrapped
+        # function that generates a special console event unique to
+        # dagster-sqlmesh
+        def wrap_apply_event(f: ContextApplyFunction) -> ContextApplyFunction:
+            def wrapped_apply(plan: SQLMeshPlan, *args: t.Any, **kwargs: t.Any):
+                self.logger.debug("capturing plan as event")
+                self.console.capture_built_plan(plan)
+                result = f(plan, *args, **kwargs)
+                return result
+            return wrapped_apply
+        context.apply = wrap_apply_event(context.apply)
+        return context
 
     @contextmanager
     def instance(