fix: move get_asset_key_str to translator and define translator from config

Author: lucargir

README.md

@@ -24,11 +24,11 @@ from dagster import (
     AssetExecutionContext,
     Definitions,
 )
-from dagster_sqlmesh import sqlmesh_assets, SQLMeshContextConfig, SQLMeshResource, SQLMeshDagsterTranslator
+from dagster_sqlmesh import sqlmesh_assets, SQLMeshContextConfig, SQLMeshResource
 
 sqlmesh_config = SQLMeshContextConfig(path="/home/foo/sqlmesh_project", gateway="name-of-your-gateway")
 
-@sqlmesh_assets(environment="dev", config=sqlmesh_config, translator=SQLMeshDagsterTranslator())
+@sqlmesh_assets(environment="dev", config=sqlmesh_config)
 def sqlmesh_project(context: AssetExecutionContext, sqlmesh: SQLMeshResource):
     yield from sqlmesh.run(context)
 
@@ -40,6 +40,34 @@ defs = Definitions(
 )
 ```
 
+## Advanced Usage
+
+### Custom Translator
+
+The translator is centrally configured and ensures consistency across all components. You can customize the translator by specifying a custom class in the config:
+
+```python
+from dagster_sqlmesh import SQLMeshDagsterTranslator
+
+class CustomSQLMeshTranslator(SQLMeshDagsterTranslator):
+    def get_asset_key_str(self, fqn: str) -> str:
+        # Custom asset key generation logic
+        return f"custom_prefix__{super().get_asset_key_str(fqn)}"
+
+# Configure with custom translator
+sqlmesh_config = SQLMeshContextConfig(
+    path="/home/foo/sqlmesh_project", 
+    gateway="name-of-your-gateway",
+    translator_class_name="your_module.CustomSQLMeshTranslator"
+)
+
+@sqlmesh_assets(environment="dev", config=sqlmesh_config)
+def sqlmesh_project(context: AssetExecutionContext, sqlmesh: SQLMeshResource):
+    yield from sqlmesh.run(context)
+```
+
+This approach ensures that both the `SQLMeshResource` and the `@sqlmesh_assets` decorator use the same translator instance, preventing inconsistencies. The translator is created using `config.get_translator()` and passed to all components that need it, including the `DagsterSQLMeshEventHandler`.
+
 
 ## Contributing
 

dagster_sqlmesh/asset.py

@@ -10,7 +10,6 @@
     ContextFactory,
     DagsterSQLMeshController,
 )
-from dagster_sqlmesh.translator import SQLMeshDagsterTranslator
 from dagster_sqlmesh.types import SQLMeshMultiAssetOptions
 
 logger = logging.getLogger(__name__)
@@ -20,20 +19,18 @@ def sqlmesh_to_multi_asset_options(
     environment: str,
     config: SQLMeshContextConfig,
     context_factory: ContextFactory[ContextCls] = lambda **kwargs: Context(**kwargs),
-    dagster_sqlmesh_translator: SQLMeshDagsterTranslator | None = None,
 ) -> SQLMeshMultiAssetOptions:
     """Converts sqlmesh project into a SQLMeshMultiAssetOptions object which is
     an intermediate representation of the SQLMesh project that can be used to
     create a dagster multi_asset definition."""
     controller = DagsterSQLMeshController.setup_with_config(
         config=config, context_factory=context_factory
     )
-    if not dagster_sqlmesh_translator:
-        dagster_sqlmesh_translator = SQLMeshDagsterTranslator()
+    dagster_sqlmesh_translator = config.get_translator()
 
     conversion = controller.to_asset_outs(
         environment,
-        translator=dagster_sqlmesh_translator,
+        dagster_sqlmesh_translator=dagster_sqlmesh_translator,
     )
     return conversion
 
@@ -74,7 +71,6 @@ def sqlmesh_assets(
     config: SQLMeshContextConfig,
     context_factory: ContextFactory[ContextCls] = lambda **kwargs: Context(**kwargs),
     name: str | None = None,
-    dagster_sqlmesh_translator: SQLMeshDagsterTranslator | None = None,
     compute_kind: str = "sqlmesh",
     op_tags: t.Mapping[str, t.Any] | None = None,
     required_resource_keys: set[str] | None = None,
@@ -86,7 +82,6 @@ def sqlmesh_assets(
         environment=environment,
         config=config,
         context_factory=context_factory,
-        dagster_sqlmesh_translator=dagster_sqlmesh_translator,
     )
 
     return sqlmesh_asset_from_multi_asset_options(

dagster_sqlmesh/config.py

@@ -1,12 +1,16 @@
+import inspect
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 from dagster import Config
 from pydantic import Field
 from sqlmesh.core.config import Config as MeshConfig
 from sqlmesh.core.config.loader import load_configs
 
+if TYPE_CHECKING:
+    from dagster_sqlmesh.translator import SQLMeshDagsterTranslator
+
 
 @dataclass
 class ConfigOverride:
@@ -22,11 +26,56 @@ class SQLMeshContextConfig(Config):
     sqlmesh project define all the configuration in it's own directory which
     also ensures that configuration is consistent if running sqlmesh locally vs
     running via dagster.
+    
+    The config also manages the translator class used for converting SQLMesh
+    models to Dagster assets. You can specify a custom translator by setting
+    the translator_class_name field to the fully qualified class name.
     """
 
     path: str
     gateway: str
     config_override: dict[str, Any] | None = Field(default_factory=lambda: None)
+    translator_class_name: str = Field(
+        default="dagster_sqlmesh.translator.SQLMeshDagsterTranslator",
+        description="Fully qualified class name of the SQLMesh Dagster translator to use"
+    )
+    
+    def get_translator(self) -> "SQLMeshDagsterTranslator":
+        """Get a translator instance using the configured class name.
+        
+        Imports and validates the translator class, then creates a new instance.
+        The class must inherit from SQLMeshDagsterTranslator.
+        
+        Returns:
+            SQLMeshDagsterTranslator: A new instance of the configured translator class
+            
+        Raises:
+            ValueError: If the imported object is not a class or does not inherit 
+                       from SQLMeshDagsterTranslator
+        """
+        from importlib import import_module
+        
+        from dagster_sqlmesh.translator import SQLMeshDagsterTranslator
+        
+        module_name, class_name = self.translator_class_name.rsplit(".", 1)
+        module = import_module(module_name)
+        translator_class = getattr(module, class_name)
+        
+        # Validate that the imported class inherits from SQLMeshDagsterTranslator
+        if not inspect.isclass(translator_class):
+            raise ValueError(
+                f"'{self.translator_class_name}' is not a class. "
+                f"Expected a class that inherits from SQLMeshDagsterTranslator."
+            )
+        
+        if not issubclass(translator_class, SQLMeshDagsterTranslator):
+            raise ValueError(
+                f"Translator class '{self.translator_class_name}' must inherit from "
+                f"SQLMeshDagsterTranslator. Found class that inherits from: "
+                f"{[base.__name__ for base in translator_class.__bases__]}"
+            )
+        
+        return translator_class()
 
     @property
     def sqlmesh_config(self) -> MeshConfig:

dagster_sqlmesh/controller/dagster.py

@@ -12,7 +12,6 @@
     SQLMeshModelDep,
     SQLMeshMultiAssetOptions,
 )
-from dagster_sqlmesh.utils import get_asset_key_str
 
 logger = logging.getLogger(__name__)
 
@@ -23,7 +22,7 @@ class DagsterSQLMeshController(SQLMeshController[ContextCls]):
     def to_asset_outs(
         self,
         environment: str,
-        translator: SQLMeshDagsterTranslator,
+        dagster_sqlmesh_translator: SQLMeshDagsterTranslator,
     ) -> SQLMeshMultiAssetOptions:
         """Loads all the asset outs of the current sqlmesh environment. If a
         cache is provided, it will be tried first to load the asset outs."""
@@ -36,42 +35,42 @@ def to_asset_outs(
             context = instance.context
 
             for model, deps in instance.non_external_models_dag():
-                asset_key = translator.get_asset_key(context=context, fqn=model.fqn)
+                asset_key = dagster_sqlmesh_translator.get_asset_key(context=context, fqn=model.fqn)
                 asset_key_str = asset_key.to_user_string()
                 model_deps = [
                     SQLMeshModelDep(fqn=dep, model=context.get_model(dep))
                     for dep in deps
                 ]
                 internal_asset_deps: set[str] = set()
-                asset_tags = translator.get_tags(context, model)
+                asset_tags = dagster_sqlmesh_translator.get_tags(context, model)
 
                 for dep in model_deps:
                     if dep.model:
-                        dep_asset_key_str = translator.get_asset_key(
+                        dep_asset_key_str = dagster_sqlmesh_translator.get_asset_key(
                             context, dep.model.fqn
                         ).to_user_string()
 
                         internal_asset_deps.add(dep_asset_key_str)
                     else:
-                        table = get_asset_key_str(dep.fqn)
-                        key = translator.get_asset_key(
+                        table = dagster_sqlmesh_translator.get_asset_key_str(dep.fqn)
+                        key = dagster_sqlmesh_translator.get_asset_key(
                             context, dep.fqn
                         ).to_user_string()
                         internal_asset_deps.add(key)
 
                         # create an external dep
-                        deps_map[table] = translator.create_asset_dep(key=key)
+                        deps_map[table] = dagster_sqlmesh_translator.create_asset_dep(key=key)
 
-                model_key = get_asset_key_str(model.fqn)
-                asset_outs[model_key] = translator.create_asset_out(
+                model_key = dagster_sqlmesh_translator.get_asset_key_str(model.fqn)
+                asset_outs[model_key] = dagster_sqlmesh_translator.create_asset_out(
                     model_key=model_key,
                     asset_key=asset_key_str,
                     tags=asset_tags,
                     is_required=False,
-                    group_name=translator.get_group_name(context, model),
+                    group_name=dagster_sqlmesh_translator.get_group_name(context, model),
                     kinds={
                         "sqlmesh",
-                        translator.get_context_dialect(context).lower(),
+                        dagster_sqlmesh_translator.get_context_dialect(context).lower(),
                     },
                 )
                 internal_asset_deps_map[model_key] = internal_asset_deps

dagster_sqlmesh/resource.py

@@ -25,7 +25,9 @@
     ContextFactory,
 )
 from dagster_sqlmesh.controller.dagster import DagsterSQLMeshController
-from dagster_sqlmesh.utils import get_asset_key_str
+
+if t.TYPE_CHECKING:
+    from dagster_sqlmesh.translator import SQLMeshDagsterTranslator
 
 logger = logging.getLogger(__name__)
 
@@ -329,6 +331,7 @@ def __init__(
         models_map: dict[str, Model],
         dag: DAG[t.Any],
         prefix: str,
+        dagster_sqlmesh_translator: "SQLMeshDagsterTranslator",
         is_testing: bool = False,
         materializations_enabled: bool = True,
     ) -> None:
@@ -341,6 +344,7 @@ def __init__(
             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.
+            dagster_sqlmesh_translator: The SQLMesh Dagster translator instance.
             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
@@ -351,6 +355,7 @@ def __init__(
         self._prefix = prefix
         self._context = context
         self._logger = context.log
+        self._translator = dagster_sqlmesh_translator
         self._tracker = MaterializationTracker(
             sorted_dag=dag.sorted[:], logger=self._logger
         )
@@ -382,7 +387,7 @@ def notify_success(
             # If the model is not in models_map, we can skip any notification
             if model:
                 # Passing model.fqn to get internal unique asset key
-                output_key = get_asset_key_str(model.fqn)
+                output_key = self._translator.get_asset_key_str(model.fqn)
                 if self._is_testing:
                     asset_key = dg.AssetKey(["testing", output_key])
                     self._logger.warning(
@@ -491,15 +496,15 @@ def report_event(self, event: console.ConsoleEvent) -> None:
                     log_context.info(
                         "Snapshot progress complete",
                         {
-                            "asset_key": get_asset_key_str(snapshot.model.name),
+                            "asset_key": self._translator.get_asset_key_str(snapshot.model.name),
                         },
                     )
                     self._tracker.update_run(snapshot)
                 else:
                     log_context.info(
                         "Snapshot progress update",
                         {
-                            "asset_key": get_asset_key_str(snapshot.model.name),
+                            "asset_key": self._translator.get_asset_key_str(snapshot.model.name),
                             "progress": f"{done}/{expected}",
                             "duration_ms": duration_ms,
                         },
@@ -687,11 +692,13 @@ def create_event_handler(
         is_testing: bool,
         materializations_enabled: bool,
     ) -> DagsterSQLMeshEventHandler:
+        translator = self.config.get_translator()
         return DagsterSQLMeshEventHandler(
             context=context,
             dag=dag,
             models_map=models_map,
             prefix=prefix,
+            dagster_sqlmesh_translator=translator,
             is_testing=is_testing,
             materializations_enabled=materializations_enabled,
         )
@@ -701,7 +708,7 @@ def _get_selected_models_from_context(
     ) -> tuple[set[str], dict[str, Model], list[str] | None]:
         models_map = models.copy()
         try:
-            selected_output_names = set(context.selected_output_names)
+            selected_output_names = set(context.op_execution_context.selected_output_names)
         except (DagsterInvalidPropertyError, AttributeError) as e:
             # Special case for direct execution context when testing. This is related to:
             # https://github.com/dagster-io/dagster/issues/23633
@@ -711,10 +718,11 @@ def _get_selected_models_from_context(
             else:
                 raise e
 
+        translator = self.config.get_translator()
         select_models: list[str] = []
         models_map = {}
         for key, model in models.items():
-            if get_asset_key_str(model.fqn) in selected_output_names:
+            if translator.get_asset_key_str(model.fqn) in selected_output_names:
                 models_map[key] = model
                 select_models.append(model.name)
         return (

dagster_sqlmesh/test_asset.py

@@ -1,10 +1,9 @@
-from dagster_sqlmesh.asset import SQLMeshDagsterTranslator
 from dagster_sqlmesh.conftest import SQLMeshTestContext
 
 
 def test_sqlmesh_context_to_asset_outs(sample_sqlmesh_test_context: SQLMeshTestContext):
     controller = sample_sqlmesh_test_context.create_controller()
-    translator = SQLMeshDagsterTranslator()
-    outs = controller.to_asset_outs("dev", translator)
+    translator = sample_sqlmesh_test_context.context_config.get_translator()
+    outs = controller.to_asset_outs("dev", dagster_sqlmesh_translator=translator)
     assert len(list(outs.deps)) == 1
     assert len(outs.outs) == 10