fix: more fixes

Author: ravenac95

dagster_sqlmesh/resource.py

@@ -29,10 +29,9 @@
 
 logger = logging.getLogger(__name__)
 
+
 def _START_OF_UNIX_TIME():
-    dt = datetime.strptime(
-        "1970-01-01T00:00:00Z", "%Y-%m-%dT%H:%M:%SZ"
-    )
+    dt = datetime.strptime("1970-01-01T00:00:00Z", "%Y-%m-%dT%H:%M:%SZ")
     return dt.astimezone(UTC)
 
 
@@ -72,7 +71,9 @@ def as_dagster_metadata(
             last_promoted = dg.MetadataValue.timestamp(
                 max(previous.last_promoted, self.last_promoted)
             )
-            last_backfill = dg.MetadataValue.timestamp(max(previous.last_backfill, self.last_backfill))
+            last_backfill = dg.MetadataValue.timestamp(
+                max(previous.last_backfill, self.last_backfill)
+            )
             created_at = dg.MetadataValue.timestamp(previous.created_at)
         else:
             # If there is no previous materialization status all dates can use
@@ -98,21 +99,25 @@ def from_dagster_metadata(
         # convert metadata values
         converted: dict[str, dg.MetadataValue] = {}
         for key, value in metadata.items():
-            assert isinstance(value, dg.MetadataValue), f"Expected MetadataValue for {key}, got {type(value)}"
+            assert isinstance(
+                value, dg.MetadataValue
+            ), f"Expected MetadataValue for {key}, got {type(value)}"
             converted[key] = value
 
-        return cls.model_validate(dict(
-            model_fqn=converted["model_fqn"].value,
-            snapshot_id=converted["snapshot_id"].value,
-            created_at=converted["created_at"].value,
-            last_updated_or_restated=converted["last_updated_or_restated"].value,
-            last_promoted=converted["last_promoted"].value,
-            last_backfill=converted["last_backfill"].value,
-        ))
-    
+        return cls.model_validate(
+            dict(
+                model_fqn=converted["model_fqn"].value,
+                snapshot_id=converted["snapshot_id"].value,
+                created_at=converted["created_at"].value,
+                last_updated_or_restated=converted["last_updated_or_restated"].value,
+                last_promoted=converted["last_promoted"].value,
+                last_backfill=converted["last_backfill"].value,
+            )
+        )
+
     def as_glot_table(self) -> exp.Table:
         return sqlglot.to_table(self.model_fqn)
-    
+
     def is_match(self, input: str, ignore_catalog: bool = False) -> bool:
         """Tests if the passed in string matches this model's table
 
@@ -133,7 +138,7 @@ def is_match(self, input: str, ignore_catalog: bool = False) -> bool:
             return False
         if input_as_table.db != table.db:
             return False
-        
+
         if not ignore_catalog:
             if input_as_table.catalog != table.catalog:
                 return False
@@ -251,12 +256,17 @@ def notify_queue_next(self) -> tuple[str, ModelMaterializationStatus] | None:
 
             if model_name_for_notification in self._non_model_names:
                 self._current_index += 1
-                self.logger.debug(f"skipping non-model snapshot {model_name_for_notification}")
+                self.logger.debug(
+                    f"skipping non-model snapshot {model_name_for_notification}"
+                )
                 continue
 
             if model_name_for_notification in self._model_metadata:
                 self._current_index += 1
-                return (model_name_for_notification, self._model_metadata[model_name_for_notification])
+                return (
+                    model_name_for_notification,
+                    self._model_metadata[model_name_for_notification],
+                )
             return None
 
 
@@ -320,7 +330,23 @@ def __init__(
         dag: DAG[t.Any],
         prefix: str,
         is_testing: bool = False,
+        materializations_enabled: bool = True,
     ) -> None:
+        """Dagster event handler for SQLMesh models.
+
+        The handler is responsible for reporting events from sqlmesh to dagster.
+
+        Args:
+            context: The Dagster asset execution context.
+            models_map: A mapping of model names to their SQLMesh model instances.
+            dag: The directed acyclic graph representing the SQLMesh models.
+            prefix: A prefix to use for all asset keys generated by this handler.
+            is_testing: Whether the handler is being used in a testing context.
+            materializations_enabled: Whether the handler is to generate
+                materializations, this should be disabled if you with to run a
+                sqlmesh plan or run in an environment different from the normal
+                target environment.
+        """
         self._models_map = models_map
         self._prefix = prefix
         self._context = context
@@ -331,6 +357,7 @@ def __init__(
         self._stage = "plan"
         self._errors: list[Exception] = []
         self._is_testing = is_testing
+        self._materializations_enabled = materializations_enabled
 
     def process_events(self, event: console.ConsoleEvent) -> None:
         self.report_event(event)
@@ -363,7 +390,14 @@ def notify_success(
                     )
                 else:
                     asset_key = self._context.asset_key_for_output(output_key)
-                yield self.create_materialize_result(self._context, asset_key, materialization_status)
+                if self._materializations_enabled:
+                    yield self.create_materialize_result(
+                        self._context, asset_key, materialization_status
+                    )
+                else:
+                    self._logger.debug(
+                        f"Materializations disabled. Would have materialized for {asset_key.to_user_string()}"
+                    )
             notify = self._tracker.notify_queue_next()
         else:
             self._logger.debug("No more materializations to process")
@@ -384,12 +418,14 @@ def create_materialize_result(
             )
             last_materialization_status = None
         else:
-            assert last_materialization.asset_materialization is not None, (
-                "Expected asset materialization to be present."
-            )
+            assert (
+                last_materialization.asset_materialization is not None
+            ), "Expected asset materialization to be present."
             try:
-                last_materialization_status = ModelMaterializationStatus.from_dagster_metadata(
-                    dict(last_materialization.asset_materialization.metadata)
+                last_materialization_status = (
+                    ModelMaterializationStatus.from_dagster_metadata(
+                        dict(last_materialization.asset_materialization.metadata)
+                    )
                 )
             except Exception as e:
                 self._logger.warning(
@@ -399,7 +435,9 @@ def create_materialize_result(
 
         return dg.MaterializeResult(
             asset_key=asset_key,
-            metadata=current_materialization_status.as_dagster_metadata(last_materialization_status),
+            metadata=current_materialization_status.as_dagster_metadata(
+                last_materialization_status
+            ),
         )
 
     def report_event(self, event: console.ConsoleEvent) -> None:
@@ -554,6 +592,7 @@ def run(
         skip_run: bool = False,
         plan_options: PlanOptions | None = None,
         run_options: RunOptions | None = None,
+        materializations_enabled: bool = True,
     ) -> t.Iterable[dg.MaterializeResult]:
         """Execute SQLMesh based on the configuration given"""
         plan_options = plan_options or {}
@@ -592,6 +631,7 @@ def run(
                 dag=dag,
                 prefix="sqlmesh: ",
                 is_testing=self.is_testing,
+                materializations_enabled=materializations_enabled,
             )
 
             def raise_for_sqlmesh_errors(
@@ -645,13 +685,15 @@ def create_event_handler(
         models_map: dict[str, Model],
         prefix: str,
         is_testing: bool,
+        materializations_enabled: bool,
     ) -> DagsterSQLMeshEventHandler:
         return DagsterSQLMeshEventHandler(
             context=context,
             dag=dag,
             models_map=models_map,
             prefix=prefix,
             is_testing=is_testing,
+            materializations_enabled=materializations_enabled,
         )
 
     def _get_selected_models_from_context(