|
9 | 9 | from sqlmesh.core.console import set_console |
10 | 10 | from sqlmesh.core.context import Context |
11 | 11 | from sqlmesh.core.model import Model |
12 | | -from sqlmesh.core.plan import PlanBuilder |
| 12 | +from sqlmesh.core.plan import Plan as SQLMeshPlan, PlanBuilder |
13 | 13 | from sqlmesh.utils.dag import DAG |
14 | 14 | from sqlmesh.utils.date import TimeLike |
15 | 15 |
|
|
19 | 19 | ConsoleEventHandler, |
20 | 20 | ConsoleException, |
21 | 21 | EventConsole, |
22 | | - Plan, |
23 | 22 | SnapshotCategorizer, |
24 | 23 | ) |
25 | 24 | from dagster_sqlmesh.events import ConsoleGenerator |
@@ -179,16 +178,7 @@ def run_sqlmesh_thread( |
179 | 178 | ) -> None: |
180 | 179 | logger.debug("dagster-sqlmesh: thread started") |
181 | 180 |
|
182 | | - def auto_execute_plan(event: ConsoleEvent): |
183 | | - if isinstance(event, Plan): |
184 | | - try: |
185 | | - event.plan_builder.apply() |
186 | | - except Exception as e: |
187 | | - controller.console.exception(e) |
188 | | - return None |
189 | | - |
190 | 181 | try: |
191 | | - controller.console.add_handler(auto_execute_plan) |
192 | 182 | builder = t.cast( |
193 | 183 | PlanBuilder, |
194 | 184 | context.plan_builder( |
@@ -381,6 +371,11 @@ def non_external_models_dag(self) -> t.Iterable[tuple[Model, set[str]]]: |
381 | 371 | continue |
382 | 372 | yield (model, deps) |
383 | 373 |
|
| 374 | + |
| 375 | +class ContextApplyFunction(t.Protocol): |
| 376 | + def __call__(self, plan: SQLMeshPlan, *args: t.Any, **kwargs: t.Any) -> None: |
| 377 | + ... |
| 378 | + |
384 | 379 | class SQLMeshController(t.Generic[ContextCls]): |
385 | 380 | """Allows control of sqlmesh via a python interface. It is not suggested to |
386 | 381 | use the constructor of this class directly, but instead use the provided |
@@ -480,7 +475,21 @@ def _create_context(self) -> ContextCls: |
480 | 475 | if self.config.sqlmesh_config: |
481 | 476 | options["config"] = self.config.sqlmesh_config |
482 | 477 | set_console(self.console) |
483 | | - return self._context_factory(**options) |
| 478 | + context = self._context_factory(**options) |
| 479 | + |
| 480 | + # As part of the context, it specifies a method "apply" that we would |
| 481 | + # like to introspect. To do so we replace "apply" with a wrapped |
| 482 | + # function that generates a special console event unique to |
| 483 | + # dagster-sqlmesh |
| 484 | + def wrap_apply_event(f: ContextApplyFunction) -> ContextApplyFunction: |
| 485 | + def wrapped_apply(plan: SQLMeshPlan, *args: t.Any, **kwargs: t.Any): |
| 486 | + self.logger.debug("capturing plan as event") |
| 487 | + self.console.capture_built_plan(plan) |
| 488 | + result = f(plan, *args, **kwargs) |
| 489 | + return result |
| 490 | + return wrapped_apply |
| 491 | + context.apply = wrap_apply_event(context.apply) |
| 492 | + return context |
484 | 493 |
|
485 | 494 | @contextmanager |
486 | 495 | def instance( |
|
0 commit comments