Fixes Dagster Resource (#6)

* Fixes Dagster Resource

* dev publish

* Fixed broken usage of instance

* fix

* Better CI

Author: ravenac95

.github/workflows/ci-default.yml

@@ -69,5 +69,8 @@ jobs:
       - name: Test (Python)
         run: |
           poetry run pytest
-      
-     
+
+      - name: Run the sample dagster project
+        run: |
+          poetry run dagster job execute -f sample/dagster_project/definitions.py -j all_assets_job
+      

dagster_sqlmesh/controller/base.py

@@ -350,13 +350,19 @@ def setup_with_config(
         controller = cls(
             console=console,
             config=config,
+            log_override=log_override,
         )
         return controller
 
-    def __init__(self, config: SQLMeshContextConfig, console: EventConsole):
+    def __init__(
+        self,
+        config: SQLMeshContextConfig,
+        console: EventConsole,
+        log_override: t.Optional[logging.Logger] = None,
+    ):
         self.config = config
         self.console = console
-        self.logger = logger
+        self.logger = log_override or logger
         self._context_open = False
 
     def set_logger(self, logger: logging.Logger):
@@ -379,7 +385,10 @@ def _create_context(self):
         return Context(**options)
 
     @contextmanager
-    def instance(self, environment: str):
+    def instance(self, environment: str, component: str = "unknown"):
+        self.logger.info(
+            f"Opening sqlmesh instance for env={environment} component={component}"
+        )
         if self._context_open:
             raise Exception("Only one sqlmesh instance at a time")
 
@@ -390,6 +399,9 @@ def instance(self, environment: str):
                 environment, self.console, self.config, context, self.logger
             )
         finally:
+            self.logger.info(
+                f"Closing sqlmesh instance for env={environment} component={component}"
+            )
             self._context_open = False
             context.close()
 
@@ -398,7 +410,7 @@ def run(
         environment: str,
         **run_options: t.Unpack[RunOptions],
     ):
-        with self.instance(environment) as mesh:
+        with self.instance(environment, "run") as mesh:
             yield from mesh.run(**run_options)
 
     def plan(
@@ -408,7 +420,7 @@ def plan(
         default_catalog: t.Optional[str],
         plan_options: PlanOptions,
     ):
-        with self.instance(environment) as mesh:
+        with self.instance(environment, "plan") as mesh:
             yield from mesh.plan(categorizer, default_catalog, **plan_options)
 
     def plan_and_run(
@@ -419,7 +431,7 @@ def plan_and_run(
         plan_options: t.Optional[PlanOptions] = None,
         run_options: t.Optional[RunOptions] = None,
     ):
-        with self.instance(environment) as mesh:
+        with self.instance(environment, "plan_and_run") as mesh:
             yield from mesh.plan_and_run(
                 categorizer=categorizer,
                 default_catalog=default_catalog,

dagster_sqlmesh/controller/dagster.py

@@ -22,7 +22,7 @@ class DagsterSQLMeshController(SQLMeshController):
     def to_asset_outs(
         self, environment: str, translator: SQLMeshDagsterTranslator
     ) -> SQLMeshMultiAssetOptions:
-        with self.instance(environment) as instance:
+        with self.instance(environment, "to_asset_outs") as instance:
             context = instance.context
             dag = context.dag
             output = SQLMeshMultiAssetOptions()

dagster_sqlmesh/resource.py

@@ -248,12 +248,13 @@ def run(
                 context, models_map, dag, "sqlmesh: "
             )
 
-            for event in controller.plan_and_run(
-                environment=environment,
+            for event in mesh.plan_and_run(
                 plan_options=plan_options,
                 run_options=run_options,
             ):
                 yield from event_handler.process_events(mesh.context, event)
 
     def get_controller(self, log_override: t.Optional[logging.Logger] = None):
-        return SQLMeshController.setup(self.config, log_override=log_override)
+        return SQLMeshController.setup_with_config(
+            self.config, log_override=log_override
+        )

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "dagster-sqlmesh"
-version = "0.3.0"
+version = "0.3.1"
 description = ""
 authors = ["Reuven Gonzales <[email protected]>"]
 license = "Apache 2.0"

sample/dagster_project/definitions.py

@@ -6,6 +6,7 @@
     asset,
     AssetExecutionContext,
     Definitions,
+    define_asset_job,
 )
 from dagster_duckdb_polars import DuckDBPolarsIOManager
 import polars as pl
@@ -55,6 +56,8 @@ def sqlmesh_project(context: AssetExecutionContext, sqlmesh: SQLMeshResource):
     yield from sqlmesh.run(context)
 
 
+all_assets_job = define_asset_job(name="all_assets_job")
+
 defs = Definitions(
     assets=[sqlmesh_project, test_source, reset_asset],
     resources={
@@ -64,4 +67,5 @@ def sqlmesh_project(context: AssetExecutionContext, sqlmesh: SQLMeshResource):
             schema="sources",
         ),
     },
+    jobs=[all_assets_job],
 )