fix: remove translator from controller and use helper for internal ref

Author: lucargir

README.md

@@ -28,9 +28,9 @@ from dagster_sqlmesh import sqlmesh_assets, SQLMeshContextConfig, SQLMeshResourc
 
 sqlmesh_config = SQLMeshContextConfig(path="/home/foo/sqlmesh_project", gateway="name-of-your-gateway")
 
-@sqlmesh_assets(environment="dev", config=sqlmesh_config)
+@sqlmesh_assets(environment="dev", config=sqlmesh_config, translator=SQLMeshDagsterTranslator())
 def sqlmesh_project(context: AssetExecutionContext, sqlmesh: SQLMeshResource):
-    yield from sqlmesh.run(context, translator=SQLMeshDagsterTranslator())
+    yield from sqlmesh.run(context)
 
 defs = Definitions(
     assets=[sqlmesh_project],

dagster_sqlmesh/controller/base.py

@@ -23,7 +23,6 @@
     SnapshotCategorizer,
 )
 from ..events import ConsoleGenerator
-from ..translator import SQLMeshDagsterTranslator
 
 logger = logging.getLogger(__name__)
 
@@ -417,7 +416,6 @@ class SQLMeshController(t.Generic[ContextCls]):
     config: SQLMeshContextConfig
     console: EventConsole
     logger: logging.Logger
-    translator: SQLMeshDagsterTranslator
 
     @classmethod
     def setup(
@@ -427,12 +425,10 @@ def setup(
         context_factory: ContextFactory[ContextCls],
         gateway: str = "local",
         log_override: logging.Logger | None = None,
-        translator: SQLMeshDagsterTranslator = SQLMeshDagsterTranslator(),
     ) -> t.Self:
         return cls.setup_with_config(
             config=SQLMeshContextConfig(path=path, gateway=gateway),
             log_override=log_override,
-            translator=translator,
             context_factory=context_factory,
         )
 
@@ -443,15 +439,13 @@ def setup_with_config(
         config: SQLMeshContextConfig,
         context_factory: ContextFactory[ContextCls] = DEFAULT_CONTEXT_FACTORY,
         log_override: logging.Logger | None = None,
-        translator: SQLMeshDagsterTranslator = SQLMeshDagsterTranslator(),
     ) -> t.Self:
         console = EventConsole(log_override=log_override) # type: ignore
         controller = cls(
             console=console,
             config=config,
             log_override=log_override,
-            context_factory=context_factory,
-            translator=translator
+            context_factory=context_factory
         )
         return controller
 
@@ -461,13 +455,11 @@ def __init__(
         console: EventConsole,
         context_factory: ContextFactory[ContextCls],
         log_override: logging.Logger | None = None,
-        translator: SQLMeshDagsterTranslator = SQLMeshDagsterTranslator(),
     ) -> None:
         self.config = config
         self.console = console
         self.logger = log_override or logger
         self._context_factory = context_factory
-        self.translator = translator
         self._context_open = False
 
     def set_logger(self, logger: logging.Logger) -> None:

dagster_sqlmesh/controller/dagster.py

@@ -5,6 +5,8 @@
 from dagster import AssetDep, AssetKey, AssetOut
 from dagster._core.definitions.asset_dep import CoercibleToAssetDep
 
+from dagster_sqlmesh.utils import get_asset_key_str
+
 from ..translator import SQLMeshDagsterTranslator
 from ..types import SQLMeshModelDep, SQLMeshMultiAssetOptions
 from .base import ContextCls, SQLMeshController
@@ -16,10 +18,9 @@ class DagsterSQLMeshController(SQLMeshController[ContextCls]):
     """An extension of the sqlmesh controller specifically for dagster use"""
 
     def to_asset_outs(
-        self, environment: str, translator: SQLMeshDagsterTranslator | None = None,
+        self, environment: str, translator: SQLMeshDagsterTranslator = SQLMeshDagsterTranslator(),
     ) -> SQLMeshMultiAssetOptions:
         with self.instance(environment, "to_asset_outs") as instance:
-            translator = translator or SQLMeshDagsterTranslator()
             context = instance.context
             output = SQLMeshMultiAssetOptions()
             depsMap: dict[str, CoercibleToAssetDep] = {}
@@ -39,12 +40,12 @@ def to_asset_outs(
                             translator.get_asset_key(context, dep.model.fqn)
                         )
                     else:
-                        table = translator.get_asset_key_str(dep.fqn)
+                        table = get_asset_key_str(dep.fqn)
                         key = translator.get_asset_key(context, dep.fqn)
                         internal_asset_deps.add(key)
                         # create an external dep
                         depsMap[table] = AssetDep(key)
-                model_key = translator.get_asset_key_str(model.fqn)
+                model_key = get_asset_key_str(model.fqn)
                 # If current Dagster supports "kinds", add labels for Dagster UI
                 if "kinds" in signature(AssetOut).parameters:
                     output.outs[model_key] = AssetOut(

dagster_sqlmesh/resource.py

@@ -20,6 +20,7 @@
     ContextCls,
     ContextFactory,
 )
+from dagster_sqlmesh.utils import get_asset_key_str
 
 from . import console
 from .config import SQLMeshContextConfig
@@ -177,7 +178,7 @@ def notify_success(
             # If the model is not in models_map, we can skip any notification
             if model:
                 # Passing model.fqn to translator
-                output_key = self.translator.get_asset_key_str(model.fqn)
+                output_key = get_asset_key_str(model.fqn)
                 if not self._is_testing:
                     # Stupidly dagster when testing cannot use the following
                     # method so we must specifically skip this when testing
@@ -231,7 +232,7 @@ def report_event(self, event: console.ConsoleEvent) -> None:
                 log_context.info(
                     "Snapshot progress update",
                     {
-                        "asset_key": self.translator.get_asset_key_str(snapshot.model.name),
+                        "asset_key": get_asset_key_str(snapshot.model.name),
                         "progress": f"{done}/{expected}",
                         "duration_ms": duration_ms,
                     },
@@ -332,7 +333,11 @@ def run(
 
         logger = context.log
 
-        controller = self.get_controller(context_factory, logger, translator)
+        controller = self.get_controller(
+            context_factory=context_factory,
+            log_override=logger,
+            translator=translator
+        )
 
         with controller.instance(environment) as mesh:
             dag = mesh.models_dag()
@@ -404,7 +409,7 @@ def _get_selected_models_from_context(
         select_models: list[str] = []
         models_map = {}
         for key, model in models.items():
-            if translator.get_asset_key_str(model.fqn) in selected_output_names:
+            if get_asset_key_str(model.fqn) in selected_output_names:
                 models_map[key] = model
                 select_models.append(model.name)
         return (
@@ -416,12 +421,11 @@ def _get_selected_models_from_context(
     def get_controller(
         self,
         context_factory: ContextFactory[ContextCls],
+        translator: SQLMeshDagsterTranslator,
         log_override: logging.Logger | None = None,
-        translator: SQLMeshDagsterTranslator = SQLMeshDagsterTranslator(),
     ) -> DagsterSQLMeshController[ContextCls]:
         return DagsterSQLMeshController.setup_with_config(
             config=self.config,
             context_factory=context_factory,
-            log_override=log_override,
-            translator=translator
+            log_override=log_override
         )

dagster_sqlmesh/translator.py

@@ -1,7 +1,7 @@
-import re
 from collections.abc import Sequence
 
 from dagster import AssetKey
+from sqlglot import exp
 from sqlmesh.core.context import Context
 from sqlmesh.core.model import Model
 
@@ -15,16 +15,10 @@ def get_asset_key(self, context: Context, fqn: str) -> AssetKey:
         return AssetKey(path)
 
     def get_asset_key_name(self, fqn: str) -> Sequence[str]:
-        asset_path = re.findall(r"[A-Za-z0-9_\-]+", fqn)
-        return asset_path
+        table = exp.to_table(fqn)
+        asset_key_name = [table.catalog, table.db, table.name]
 
-    def get_asset_key_str(self, fqn: str) -> str:
-        # This is an internal identifier used to map outputs and dependencies
-        # it will not affect the existing AssetKeys
-        # Only alphanumeric characters and underscores
-        path = self.get_asset_key_name(fqn)
-        
-        return "__dot__".join(path).replace("-", "__dash__")
+        return asset_key_name
 
     def get_group_name(self, context: Context, model: Model) -> str:
         path = self.get_asset_key_name(model.fqn)

dagster_sqlmesh/utils.py

@@ -1,6 +1,16 @@
+from sqlglot import exp
 from sqlmesh.core.snapshot import SnapshotId
 
 
+def get_asset_key_str(fqn: str) -> str:
+    # This is an internal identifier used to map outputs and dependencies
+    # it will not affect the existing AssetKeys
+    # Only alphanumeric characters and underscores
+    table = exp.to_table(fqn)
+    asset_key_name = [table.catalog, table.db, table.name]
+    
+    return "sqlmesh__" + "_".join(asset_key_name)
+
 def snapshot_id_to_model_name(snapshot_id: SnapshotId) -> str:
     """Convert a SnapshotId to its model name.
 

sample/dagster_project/definitions.py

@@ -21,7 +21,7 @@
 sqlmesh_config = SQLMeshContextConfig(path=SQLMESH_PROJECT_PATH, gateway="local")
 
 
-@asset
+@asset(key=["db", "sources", "reset_asset"])
 def reset_asset() -> MaterializeResult:
     """An asset used for testing this entire workflow. If the duckdb database is
     found, this will delete it. This allows us to continously test this dag if
@@ -34,7 +34,7 @@ def reset_asset() -> MaterializeResult:
     return MaterializeResult(metadata={"deleted": deleted})
 
 
-@asset(deps=[reset_asset])
+@asset(deps=[reset_asset], key=["db", "sources", "test_source"])
 def test_source() -> pl.DataFrame:
     """Sets up the `test_source` table in duckdb that one of the sample sqlmesh
     models depends on"""