add to asset spec

dagster-io · May 24, 2024 · 7ade35a · 7ade35a
1 parent 4517e11
commit 7ade35a
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 1 deletion.
diff --git a/python_modules/dagster/dagster/_core/definitions/asset_spec.py b/python_modules/dagster/dagster/_core/definitions/asset_spec.py
@@ -11,6 +11,7 @@
 
 import dagster._check as check
 from dagster._annotations import PublicAttr, experimental_param
+from dagster._core.definitions.tags import StorageKindTagSet
 from dagster._serdes.serdes import whitelist_for_serdes
 
 from .auto_materialize_policy import AutoMaterializePolicy
@@ -125,12 +126,18 @@ def __new__(
         auto_materialize_policy: Optional[AutoMaterializePolicy] = None,
         owners: Optional[Sequence[str]] = None,
         tags: Optional[Mapping[str, str]] = None,
+        storage_kind: Optional[str] = None,
     ):
         from dagster._core.definitions.asset_dep import coerce_to_deps_and_check_duplicates
 
         key = AssetKey.from_coercible(key)
         asset_deps = coerce_to_deps_and_check_duplicates(deps, key)
 
+        all_tags = {
+            **(tags or {}),
+            **(StorageKindTagSet(storage_kind=storage_kind) if storage_kind else {}),
+        } or None
+
         return super().__new__(
             cls,
             key=key,
@@ -151,5 +158,5 @@ def __new__(
                 AutoMaterializePolicy,
             ),
             owners=check.opt_sequence_param(owners, "owners", of_type=str),
-            tags=validate_tags_strict(tags),
+            tags=validate_tags_strict(all_tags),
         )
diff --git a/python_modules/dagster/dagster_tests/asset_defs_tests/test_assets.py b/python_modules/dagster/dagster_tests/asset_defs_tests/test_assets.py
@@ -2243,6 +2243,20 @@ def assets2(): ...
 
 
 def test_multi_assets_asset_spec_with_storage_kind_override() -> None:
+    @multi_asset(
+        specs=[
+            AssetSpec("asset1"),
+            AssetSpec("asset2", storage_kind="bigquery"),
+        ],
+        storage_kind="snowflake",
+    )
+    def assets2(): ...
+
+    assert assets2.tags_by_key[AssetKey("asset1")] == {"dagster/storage_kind": "snowflake"}
+    assert assets2.tags_by_key[AssetKey("asset2")] == {"dagster/storage_kind": "bigquery"}
+
+
+def test_multi_assets_asset_spec_with_storage_kind_tag_override() -> None:
     @multi_asset(
         specs=[
             AssetSpec("asset1"),