fix: optional translator, define otherwise in config and use sqlglot in translator

Author: lucargir

dagster_sqlmesh/controller/base.py

@@ -427,7 +427,7 @@ def setup(
         context_factory: ContextFactory[ContextCls],
         gateway: str = "local",
         log_override: logging.Logger | None = None,
-        translator: SQLMeshDagsterTranslator = SQLMeshDagsterTranslator(),
+        translator: SQLMeshDagsterTranslator | None = None,
     ) -> t.Self:
         return cls.setup_with_config(
             config=SQLMeshContextConfig(path=path, gateway=gateway),
@@ -443,9 +443,11 @@ def setup_with_config(
         config: SQLMeshContextConfig,
         context_factory: ContextFactory[ContextCls] = DEFAULT_CONTEXT_FACTORY,
         log_override: logging.Logger | None = None,
-        translator: SQLMeshDagsterTranslator = SQLMeshDagsterTranslator(),
+        translator: SQLMeshDagsterTranslator | None = None,
     ) -> t.Self:
         console = EventConsole(log_override=log_override) # type: ignore
+        if not translator:
+            translator = SQLMeshDagsterTranslator()
         controller = cls(
             console=console,
             config=config,
@@ -460,8 +462,8 @@ def __init__(
         config: SQLMeshContextConfig,
         console: EventConsole,
         context_factory: ContextFactory[ContextCls],
+        translator: SQLMeshDagsterTranslator,
         log_override: logging.Logger | None = None,
-        translator: SQLMeshDagsterTranslator = SQLMeshDagsterTranslator(),
     ) -> None:
         self.config = config
         self.console = console

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,8 +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]
+
+        return asset_key_name
 
     def get_asset_key_str(self, fqn: str) -> str:
         # This is an internal identifier used to map outputs and dependencies

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"""