fix: address more issues with moving config around

Author: ravenac95

dagster_sqlmesh/config.py

@@ -7,8 +7,7 @@
 from sqlmesh.core.config import Config as MeshConfig
 from sqlmesh.core.config.loader import load_configs
 
-if t.TYPE_CHECKING:
-    from dagster_sqlmesh.translator import SQLMeshDagsterTranslator
+from dagster_sqlmesh.translator import SQLMeshDagsterTranslator
 
 
 @dataclass
@@ -73,7 +72,7 @@ class and return a different translator. However, due to the way that
     gateway: str
     config_override: dict[str, t.Any] | None = Field(default_factory=lambda: None)
 
-    def get_translator(self) -> "SQLMeshDagsterTranslator":
+    def get_translator(self) -> SQLMeshDagsterTranslator:
         """Get a translator instance. Override this method to provide a custom translator.
         
         Returns:

dagster_sqlmesh/resource.py

@@ -606,7 +606,9 @@ def run(
         logger = context.log
 
         controller = self.get_controller(
-            context_factory=context_factory, log_override=logger
+            config=config,
+            context_factory=context_factory, 
+            log_override=logger
         )
 
         with controller.instance(environment) as mesh:
@@ -738,11 +740,12 @@ def _get_selected_models_from_context(
 
     def get_controller(
         self,
+        config: SQLMeshContextConfig,
         context_factory: ContextFactory[ContextCls],
         log_override: logging.Logger | None = None,
     ) -> DagsterSQLMeshController[ContextCls]:
         return DagsterSQLMeshController.setup_with_config(
-            config=self.config,
+            config=config,
             context_factory=context_factory,
             log_override=log_override,
         )

dagster_sqlmesh/test_resource.py

@@ -59,10 +59,11 @@ def test_sqlmesh_resource_should_report_no_errors(
 ):
     resource = sample_sqlmesh_resource_initialization.resource
     dg_context = sample_sqlmesh_resource_initialization.dagster_context
+    context_config = sample_sqlmesh_resource_initialization.test_context.context_config
 
     success = True
     try:
-        for result in resource.run(dg_context):
+        for result in resource.run(context=dg_context, config=context_config):
             pass
     except PlanOrRunFailedError as e:
         success = False
@@ -83,10 +84,11 @@ def test_sqlmesh_resource_properly_reports_errors(
     )
     resource = sqlmesh_resource_initialization.resource
     dg_context = sqlmesh_resource_initialization.dagster_context
+    context_config = sqlmesh_resource_initialization.test_context.context_config
 
     caught_failure = False
     try:
-        for result in resource.run(dg_context):
+        for result in resource.run(context=dg_context, config=context_config):
             pass
     except PlanOrRunFailedError as e:
         caught_failure = True
@@ -106,6 +108,7 @@ def test_sqlmesh_resource_properly_reports_errors_not_thrown(
 ):
     dg_context = sample_sqlmesh_resource_initialization.dagster_context
     resource = sample_sqlmesh_resource_initialization.resource
+    context_config = sample_sqlmesh_resource_initialization.test_context.context_config
 
     def event_handler_factory(
         *args: t.Any, **kwargs: t.Any
@@ -120,7 +123,7 @@ def event_handler_factory(
 
     caught_failure = False
     try:
-        for result in resource.run(dg_context):
+        for result in resource.run(context=dg_context, config=context_config):
             pass
     except PlanOrRunFailedError as e:
         caught_failure = True
@@ -151,17 +154,19 @@ def test_sqlmesh_resource_should_properly_materialize_results_when_no_plan_is_ru
     resource = sample_sqlmesh_resource_initialization.resource
     dg_context = sample_sqlmesh_resource_initialization.dagster_context
     dg_instance = sample_sqlmesh_resource_initialization.dagster_instance
+    context_config = sample_sqlmesh_resource_initialization.test_context.context_config
 
     # First run should materialize all models
     initial_results: list[dg.MaterializeResult] = []
-    for result in resource.run(dg_context):
+    for result in resource.run(context=dg_context, config=context_config):
         initial_results.append(result)
         assert result.asset_key is not None, "Expected asset key to be present."
-        dg_instance.report_runless_asset_event(dg.AssetMaterialization(
-            asset_key=result.asset_key,
-            metadata=result.metadata,
-        ))
-
+        dg_instance.report_runless_asset_event(
+            dg.AssetMaterialization(
+                asset_key=result.asset_key,
+                metadata=result.metadata,
+            )
+        )
 
     # All metadata times should be set to the same time
     initial_times: set[float] = set()
@@ -180,7 +185,7 @@ def test_sqlmesh_resource_should_properly_materialize_results_when_no_plan_is_ru
 
     # Second run should also materialize all models
     second_results: list[dg.MaterializeResult] = []
-    for result in resource.run(dg_context):
+    for result in resource.run(context=dg_context, config=context_config):
         second_results.append(result)
 
     assert len(second_results) > 0, "Expected second results to be non-empty."
@@ -204,7 +209,9 @@ def test_sqlmesh_resource_should_properly_materialize_results_when_no_plan_is_ru
     # Third run will restate the full model
     third_results: list[dg.MaterializeResult] = []
     for result in resource.run(
-        dg_context, restate_models=["sqlmesh_example.full_model"]
+        context=dg_context,
+        config=context_config,
+        restate_models=["sqlmesh_example.full_model"],
     ):
         third_results.append(result)
 

dagster_sqlmesh/testing/context.py

@@ -56,8 +56,8 @@ class TestSQLMeshResource(SQLMeshResource):
     It allows for easy setup and teardown of the SQLMesh context.
     """
 
-    def __init__(self, config: SQLMeshContextConfig, is_testing: bool = False):
-        super().__init__(config=config, is_testing=is_testing)
+    def __init__(self, is_testing: bool = False):
+        super().__init__(is_testing=is_testing)
         def default_event_handler_factory(*args: t.Any, **kwargs: t.Any) -> DagsterSQLMeshEventHandler:
             """Default event handler factory for the SQLMesh resource."""
             return DagsterSQLMeshEventHandler(*args, **kwargs)
@@ -82,7 +82,9 @@ def create_event_handler(self, *args: t.Any, **kwargs: t.Any) -> DagsterSQLMeshE
             DagsterSQLMeshEventHandler: The created event handler.
         """
         # Ensure translator is passed to the event handler factory
-        kwargs['translator'] = self.config.get_translator()
+        # FIXME: this is a hack to deal with an older signature that didn't expected the config
+        config = t.cast(SQLMeshContextConfig, kwargs.pop("config"))
+        kwargs["translator"] = config.get_translator()
         return self._event_handler_factory(*args, **kwargs)
 
 
@@ -99,9 +101,7 @@ def create_controller(self) -> DagsterSQLMeshController[Context]:
         )
 
     def create_resource(self) -> TestSQLMeshResource:
-        return TestSQLMeshResource(
-            config=self.context_config, is_testing=True,
-        )
+        return TestSQLMeshResource(is_testing=True)
 
     def query(self, *args: t.Any, **kwargs: t.Any) -> list[t.Any]:
         conn = duckdb.connect(self.db_path)