feat: move translator to config, add support for kinds label and specify group through translator

The SQLMeshDagsterTranslator has been simplified to return an AssetKey, the name of
the AssetKey or a string equivalent, which deprecates utils.sqlmesh_model_name_to_key
and utils.key_to_sqlmesh_model_name.

The kinds label is now also added, to show technology labels on the UI, only when
Dagster's version allows for it.

Groups can now be specified through the translator, instead of having a fixed method.

Author: lucargir

README.md

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

dagster_sqlmesh/asset.py

@@ -25,10 +25,11 @@ def sqlmesh_assets(
     # For now we don't set this by default
     enabled_subsetting: bool = False,
 ) -> t.Callable[[t.Callable[..., t.Any]], AssetsDefinition]:
-    controller = DagsterSQLMeshController.setup_with_config(config)
     if not dagster_sqlmesh_translator:
-        dagster_sqlmesh_translator = SQLMeshDagsterTranslator()
-    conversion = controller.to_asset_outs(environment, dagster_sqlmesh_translator)
+        dagster_sqlmesh_translator = dagster_sqlmesh_translator
+    
+    controller = DagsterSQLMeshController.setup_with_config(config)
+    conversion = controller.to_asset_outs(environment, translator=dagster_sqlmesh_translator)
 
     return multi_asset(
         name=name,

dagster_sqlmesh/controller/base.py

@@ -24,6 +24,7 @@
     SnapshotCategorizer,
 )
 from ..events import ConsoleGenerator
+from ..translator import SQLMeshDagsterTranslator
 
 logger = logging.getLogger(__name__)
 
@@ -400,30 +401,34 @@ class SQLMeshController:
     config: SQLMeshContextConfig
     console: EventConsole
     logger: logging.Logger
+    translator: SQLMeshDagsterTranslator
 
     @classmethod
     def setup(
         cls,
         path: str,
         gateway: str = "local",
         log_override: logging.Logger | None = None,
+        translator_override: SQLMeshDagsterTranslator | None = None,
     ) -> "SQLMeshController":
         return cls.setup_with_config(
             config=SQLMeshContextConfig(path=path, gateway=gateway),
-            log_override=log_override,
+            log_override=log_override, translator_override=translator_override
         )
 
     @classmethod
     def setup_with_config(
         cls: type[T],
         config: SQLMeshContextConfig,
         log_override: logging.Logger | None = None,
+        translator_override: SQLMeshDagsterTranslator | None = None,
     ) -> T:
         console = EventConsole(log_override=log_override) # type: ignore
         controller = cls(
             console=console,
             config=config,
             log_override=log_override,
+            translator_override=translator_override
         )
         return controller
 
@@ -432,10 +437,12 @@ def __init__(
         config: SQLMeshContextConfig,
         console: EventConsole,
         log_override: logging.Logger | None = None,
+        translator_override: SQLMeshDagsterTranslator | None = None,
     ) -> None:
         self.config = config
         self.console = console
         self.logger = log_override or logger
+        self.translator = translator_override or SQLMeshDagsterTranslator()
         self._context_open = False
 
     def set_logger(self, logger: logging.Logger) -> None:
@@ -457,7 +464,7 @@ def _create_context(self) -> Context:
             options["config"] = self.config.sqlmesh_config
         set_console(self.console)
         return Context(**options)
-
+    
     @contextmanager
     def instance(
         self, environment: str, component: str = "unknown"

dagster_sqlmesh/controller/dagster.py

@@ -1,12 +1,12 @@
 # pyright: reportPrivateImportUsage=false
 import logging
+from inspect import signature
 
 from dagster import AssetDep, AssetKey, AssetOut
 from dagster._core.definitions.asset_dep import CoercibleToAssetDep
 
 from ..translator import SQLMeshDagsterTranslator
 from ..types import SQLMeshModelDep, SQLMeshMultiAssetOptions
-from ..utils import sqlmesh_model_name_to_key
 from .base import SQLMeshController
 
 logger = logging.getLogger(__name__)
@@ -16,18 +16,16 @@ class DagsterSQLMeshController(SQLMeshController):
     """An extension of the sqlmesh controller specifically for dagster use"""
 
     def to_asset_outs(
-        self, environment: str, translator: SQLMeshDagsterTranslator
+        self, environment: str, translator: SQLMeshDagsterTranslator | None = None,
     ) -> SQLMeshMultiAssetOptions:
         with self.instance(environment, "to_asset_outs") as instance:
+            translator = translator or SQLMeshDagsterTranslator()
             context = instance.context
             output = SQLMeshMultiAssetOptions()
             depsMap: dict[str, CoercibleToAssetDep] = {}
 
             for model, deps in instance.non_external_models_dag():
-                asset_key = translator.get_asset_key_from_model(
-                    context,
-                    model,
-                )
+                asset_key = translator.get_asset_key(context=context, fqn=model.fqn)
                 model_deps = [
                     SQLMeshModelDep(fqn=dep, model=context.get_model(dep))
                     for dep in deps
@@ -38,18 +36,27 @@ def to_asset_outs(
                 for dep in model_deps:
                     if dep.model:
                         internal_asset_deps.add(
-                            translator.get_asset_key_from_model(context, dep.model)
+                            translator.get_asset_key(context, dep.model.fqn)
                         )
                     else:
-                        table = translator.get_fqn_to_table(context, dep.fqn)
-                        key = translator.get_asset_key_fqn(context, dep.fqn)
+                        table = translator.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.name] = AssetDep(key)
-                model_key = sqlmesh_model_name_to_key(model.name)
-                output.outs[model_key] = AssetOut(
-                    key=asset_key, tags=asset_tags, is_required=False
-                )
+                        depsMap[table] = AssetDep(key)
+                model_key = translator.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(
+                        key=asset_key, tags=asset_tags, is_required=False,
+                        group_name=translator.get_group_name(context, model),
+                        kinds={"sqlmesh", translator._get_context_dialect(context).lower()}
+                    )
+                else:
+                    output.outs[model_key] = AssetOut(
+                        key=asset_key, tags=asset_tags, is_required=False,
+                        group_name=translator.get_group_name(context, model)
+                    )
                 output.internal_asset_deps[model_key] = internal_asset_deps
 
             output.deps = list(depsMap.values())

dagster_sqlmesh/resource.py

@@ -16,16 +16,17 @@
 from .config import SQLMeshContextConfig
 from .controller import PlanOptions, RunOptions
 from .controller.dagster import DagsterSQLMeshController
-from .utils import sqlmesh_model_name_to_key
+from .translator import SQLMeshDagsterTranslator
 
 
 class MaterializationTracker:
     """Tracks sqlmesh materializations and notifies dagster in the correct
     order. This is necessary because sqlmesh may skip some materializations that
     have no changes and those will be reported as completed out of order."""
 
-    def __init__(self, sorted_dag: list[str], logger: logging.Logger) -> None:
+    def __init__(self, sorted_dag: list[str], logger: logging.Logger, translator: SQLMeshDagsterTranslator) -> None:
         self.logger = logger
+        self.translator = translator
         self._batches: dict[Snapshot, int] = {}
         self._count: dict[Snapshot, int] = {}
         self._complete_update_status: dict[str, bool] = {}
@@ -114,12 +115,14 @@ def __init__(
         models_map: dict[str, Model],
         dag: DAG[t.Any],
         prefix: str,
+        translator: SQLMeshDagsterTranslator
     ) -> None:
         self._models_map = models_map
         self._prefix = prefix
         self._context = context
         self._logger = context.log
-        self._tracker = MaterializationTracker(dag.sorted[:], self._logger)
+        self.translator = translator
+        self._tracker = MaterializationTracker(sorted_dag=dag.sorted[:], logger=self._logger, translator=self.translator)
         self._stage = "plan"
 
     def process_events(self, event: console.ConsoleEvent) -> None:
@@ -143,7 +146,8 @@ def notify_success(
             # We allow selecting models. That value is mapped to models_map.
             # If the model is not in models_map, we can skip any notification
             if model:
-                output_key = sqlmesh_model_name_to_key(model.name)
+                # Passing model.fqn to translator
+                output_key = self.translator.get_asset_key_str(model.fqn)
                 asset_key = self._context.asset_key_for_output(output_key)
                 yield MaterializeResult(
                     asset_key=asset_key,
@@ -192,7 +196,7 @@ def report_event(self, event: console.ConsoleEvent) -> None:
                 log_context.info(
                     "Snapshot progress update",
                     {
-                        "asset_key": sqlmesh_model_name_to_key(snapshot.model.name),
+                        "asset_key": self.translator.get_asset_key_str(snapshot.model.name),
                         "progress": f"{done}/{expected}",
                         "duration_ms": duration_ms,
                     },
@@ -263,6 +267,7 @@ def run(
         self,
         context: AssetExecutionContext,
         *,
+        translator: SQLMeshDagsterTranslator | None = None,
         environment: str = "dev",
         start: TimeLike | None = None,
         end: TimeLike | None = None,
@@ -277,9 +282,12 @@ def run(
         plan_options = plan_options or {}
         run_options = run_options or {}
 
+        if translator is None:
+            translator = SQLMeshDagsterTranslator()
+
         logger = context.log
 
-        controller = self.get_controller(logger)
+        controller = self.get_controller(logger, translator)
 
         with controller.instance(environment) as mesh:
             dag = mesh.models_dag()
@@ -295,7 +303,7 @@ def run(
                 models_map = {}
                 for key, model in models.items():
                     if (
-                        sqlmesh_model_name_to_key(model.name)
+                        translator.get_asset_key_str(model.fqn)
                         in context.selected_output_names
                     ):
                         models_map[key] = model
@@ -312,7 +320,8 @@ def run(
                 logger.info(f"selected models: {select_models}")
 
             event_handler = DagsterSQLMeshEventHandler(
-                context, models_map, dag, "sqlmesh: "
+                context=context, models_map=models_map, dag=dag,
+                prefix="sqlmesh: ", translator=translator
             )
 
             for event in mesh.plan_and_run(
@@ -330,8 +339,8 @@ def run(
             yield from event_handler.notify_success(mesh.context)
 
     def get_controller(
-        self, log_override: logging.Logger | None = None
+        self, log_override: logging.Logger | None = None, translator: SQLMeshDagsterTranslator | None = None
     ) -> DagsterSQLMeshController:
         return DagsterSQLMeshController.setup_with_config(
-            self.config, log_override=log_override
+            self.config, log_override=log_override, translator_override=translator
         )

dagster_sqlmesh/test_asset.py

@@ -1,10 +1,8 @@
-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)
+    outs = controller.to_asset_outs("dev")
     assert len(list(outs.deps)) == 1
     assert len(outs.outs) == 9

dagster_sqlmesh/translator.py

@@ -1,27 +1,34 @@
+import re
+from collections.abc import Sequence
 
-import sqlglot
 from dagster import AssetKey
-from sqlglot import exp
 from sqlmesh.core.context import Context
 from sqlmesh.core.model import Model
 
 
 class SQLMeshDagsterTranslator:
     """Translates sqlmesh objects for dagster"""
 
-    def get_asset_key_from_model(self, context: Context, model: Model) -> AssetKey:
+    def get_asset_key(self, context: Context, fqn: str) -> AssetKey:
         """Given the sqlmesh context and a model return the asset key"""
-        return AssetKey(model.view_name)
-
-    def get_asset_key_fqn(self, context: Context, fqn: str) -> AssetKey:
-        """Given the sqlmesh context and a fqn of a model return an asset key"""
-        table = self.get_fqn_to_table(context, fqn)
-        return AssetKey(table.name)
-
-    def get_fqn_to_table(self, context: Context, fqn: str) -> exp.Table:
-        """Given the sqlmesh context and a fqn return the table"""
-        dialect = self._get_context_dialect(context)
-        return sqlglot.to_table(fqn, dialect=dialect)
+        path = self.get_asset_key_name(fqn)
+        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
+
+    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__")
+    
+    def get_group_name(self, context: Context, model: Model) -> str:
+        path = self.get_asset_key_name(model.fqn)
+        return path[-2]
 
     def _get_context_dialect(self, context: Context) -> str:
         return context.engine_adapter.dialect

dagster_sqlmesh/utils.py

@@ -1,14 +1,6 @@
 from sqlmesh.core.snapshot import SnapshotId
 
 
-def sqlmesh_model_name_to_key(name: str) -> str:
-    return name.replace(".", "_dot__")
-
-
-def key_to_sqlmesh_model_name(key: str) -> str:
-    return key.replace("_dot__", ".")
-
-
 def snapshot_id_to_model_name(snapshot_id: SnapshotId) -> str:
     """Convert a SnapshotId to its model name.