Further Improvements

Author: ravenac95

dagster_sqlmesh/controller/base.py

@@ -80,89 +80,41 @@ def parse_fqn(self):
         return parse_fqn(self.fqn)
 
 
-class SQLMeshController:
-    """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
-    `setup` class method"""
-
+class SQLMeshInstance:
     config: SQLMeshContextConfig
     console: EventConsole
     logger: logging.Logger
+    context: Context
+    environment: str
 
-    @classmethod
-    def setup(
-        cls,
+    def __init__(
+        self,
+        environment: str,
+        console: EventConsole,
         config: SQLMeshContextConfig,
-        debug_console: t.Optional[Console] = None,
-        log_override: t.Optional[logging.Logger] = None,
+        context: Context,
+        logger: logging.Logger,
     ):
-        console = EventConsole(log_override=log_override)
-        if debug_console:
-            console = DebugEventConsole(debug_console)
-        controller = cls(
-            console=console,
-            config=config,
-        )
-        return controller
-
-    def __init__(self, config: SQLMeshContextConfig, console: EventConsole):
-        self.config = config
+        self.environment = environment
         self.console = console
+        self.config = config
+        self.context = context
         self.logger = logger
 
-    def set_logger(self, logger: logging.Logger):
-        self.logger = logger
-
-    def add_event_handler(self, handler: ConsoleEventHandler):
-        return self.console.add_handler(handler)
-
-    def remove_event_handler(self, handler_id: str):
-        return self.console.remove_handler(handler_id)
-
-    def models_dag(self):
-        with self.context() as context:
-            return context.dag
-
-    def models(self):
-        with self.context() as context:
-            return context.models
-
-    def _create_context(self):
-        options: t.Dict[str, t.Any] = dict(
-            paths=self.config.path,
-            gateway=self.config.gateway,
-            console=self.console,
-        )
-        if self.config.sqlmesh_config:
-            options["config"] = self.config.sqlmesh_config
-        return Context(**options)
-
     @contextmanager
-    def context(self):
-        context = self._create_context()
-        yield context
-        context.close()
+    def console_context(self, handler: ConsoleEventHandler):
+        id = self.console.add_handler(handler)
+        yield
+        self.console.remove_handler(id)
 
     def plan(
         self,
-        environment: str,
         categorizer: t.Optional[SnapshotCategorizer] = None,
         default_catalog: t.Optional[str] = None,
         **plan_options: t.Unpack[PlanOptions],
     ):
-        with self.context() as context:
-            return self._plan(
-                context, environment, categorizer, default_catalog, plan_options
-            )
+        context = self.context
 
-    def _plan(
-        self,
-        context: Context,
-        environment: str,
-        categorizer: t.Optional[SnapshotCategorizer],
-        default_catalog: t.Optional[str],
-        plan_options: PlanOptions,
-    ):
         # Runs things in thread
         def run_sqlmesh_thread(
             logger: logging.Logger,
@@ -190,46 +142,30 @@ def run_sqlmesh_thread(
                 controller.console.exception(e)
 
         generator = ConsoleGenerator(self.logger)
-        event_id = self.add_event_handler(generator)
-
-        thread = threading.Thread(
-            target=run_sqlmesh_thread,
-            args=(
-                self.logger,
-                context,
-                self,
-                environment,
-                plan_options,
-                default_catalog,
-            ),
-        )
-        thread.start()
-
-        for event in generator.events(thread):
-            match event:
-                case ConsoleException(e):
-                    raise e
-                case _:
-                    yield (context, event)
-
-        thread.join()
+        with self.console_context(generator):
+            thread = threading.Thread(
+                target=run_sqlmesh_thread,
+                args=(
+                    self.logger,
+                    context,
+                    self,
+                    self.environment,
+                    plan_options,
+                    default_catalog,
+                ),
+            )
+            thread.start()
 
-        self.remove_event_handler(event_id)
+            for event in generator.events(thread):
+                match event:
+                    case ConsoleException(e):
+                        raise e
+                    case _:
+                        yield event
 
-    def run(
-        self,
-        environment: str,
-        **run_options: t.Unpack[RunOptions],
-    ):
-        with self.context() as context:
-            return self._run(context, environment, run_options=run_options)
+            thread.join()
 
-    def _run(
-        self,
-        context: Context,
-        environment: str,
-        run_options: RunOptions,
-    ):
+    def run(self, **run_options: t.Unpack[RunOptions]):
         # Runs things in thread
         def run_sqlmesh_thread(
             logger: logging.Logger,
@@ -245,29 +181,130 @@ def run_sqlmesh_thread(
                 controller.console.exception(e)
 
         generator = ConsoleGenerator(self.logger)
-        event_id = self.add_event_handler(generator)
-
-        thread = threading.Thread(
-            target=run_sqlmesh_thread,
-            args=(
-                self.logger,
-                context,
-                self,
-                environment,
-                run_options or {},
-            ),
+        with self.console_context(generator):
+            thread = threading.Thread(
+                target=run_sqlmesh_thread,
+                args=(
+                    self.logger,
+                    self.context,
+                    self,
+                    self.environment,
+                    run_options or {},
+                ),
+            )
+            thread.start()
+
+            for event in generator.events(thread):
+                match event:
+                    case ConsoleException(e):
+                        raise e
+                    case _:
+                        yield event
+
+            thread.join()
+
+    def plan_and_run(
+        self,
+        categorizer: t.Optional[SnapshotCategorizer] = None,
+        default_catalog: t.Optional[str] = None,
+        plan_options: t.Optional[PlanOptions] = None,
+        run_options: t.Optional[RunOptions] = None,
+    ):
+        run_options = run_options or {}
+        plan_options = plan_options or {}
+
+        yield from self.plan(categorizer, default_catalog, **plan_options)
+        yield from self.run(**run_options)
+
+    def models(self):
+        return self.context.models
+
+    def models_dag(self):
+        return self.context.dag
+
+
+class SQLMeshController:
+    """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
+    `setup` class method"""
+
+    config: SQLMeshContextConfig
+    console: EventConsole
+    logger: logging.Logger
+
+    @classmethod
+    def setup(
+        cls,
+        config: SQLMeshContextConfig,
+        debug_console: t.Optional[Console] = None,
+        log_override: t.Optional[logging.Logger] = None,
+    ):
+        console = EventConsole(log_override=log_override)
+        if debug_console:
+            console = DebugEventConsole(debug_console)
+        controller = cls(
+            console=console,
+            config=config,
         )
-        thread.start()
+        return controller
 
-        for event in generator.events(thread):
-            match event:
-                case ConsoleException(e):
-                    raise e
-                case _:
-                    yield (context, event)
+    def __init__(self, config: SQLMeshContextConfig, console: EventConsole):
+        self.config = config
+        self.console = console
+        self.logger = logger
+        self._context_open = False
 
-        thread.join()
-        self.remove_event_handler(event_id)
+    def set_logger(self, logger: logging.Logger):
+        self.logger = logger
+
+    def add_event_handler(self, handler: ConsoleEventHandler):
+        return self.console.add_handler(handler)
+
+    def remove_event_handler(self, handler_id: str):
+        return self.console.remove_handler(handler_id)
+
+    def _create_context(self):
+        options: t.Dict[str, t.Any] = dict(
+            paths=self.config.path,
+            gateway=self.config.gateway,
+            console=self.console,
+        )
+        if self.config.sqlmesh_config:
+            options["config"] = self.config.sqlmesh_config
+        return Context(**options)
+
+    @contextmanager
+    def instance(self, environment: str):
+        if self._context_open:
+            raise Exception("Only one sqlmesh instance at a time")
+
+        context = self._create_context()
+        self._context_open = True
+        try:
+            yield SQLMeshInstance(
+                environment, self.console, self.config, context, self.logger
+            )
+        finally:
+            self._context_open = False
+            context.close()
+
+    def run(
+        self,
+        environment: str,
+        **run_options: t.Unpack[RunOptions],
+    ):
+        with self.instance(environment) as mesh:
+            yield from mesh.run(**run_options)
+
+    def plan(
+        self,
+        environment: str,
+        categorizer: t.Optional[SnapshotCategorizer],
+        default_catalog: t.Optional[str],
+        plan_options: PlanOptions,
+    ):
+        with self.instance(environment) as mesh:
+            yield from mesh.plan(categorizer, default_catalog, **plan_options)
 
     def plan_and_run(
         self,
@@ -277,17 +314,10 @@ def plan_and_run(
         plan_options: t.Optional[PlanOptions] = None,
         run_options: t.Optional[RunOptions] = None,
     ):
-        with self.context() as context:
-            yield from self._plan(
-                context,
-                environment=environment,
+        with self.instance(environment) as mesh:
+            yield from mesh.plan_and_run(
                 categorizer=categorizer,
                 default_catalog=default_catalog,
-                plan_options=plan_options or {},
-            )
-
-            yield from self._run(
-                context,
-                environment=environment,
-                run_options=run_options or {},
+                plan_options=plan_options,
+                run_options=run_options,
             )