allow setting internal freshness policy in asset decorator

Author: anuthebananu

python_modules/dagster/dagster/_core/definitions/decorators/asset_decorator.py

@@ -34,6 +34,7 @@
     CoercibleToAssetKey,
     CoercibleToAssetKeyPrefix,
 )
+from dagster._core.definitions.freshness import INTERNAL_FRESHNESS_POLICY_METADATA_KEY
 from dagster._core.definitions.freshness_policy import FreshnessPolicy
 from dagster._core.definitions.input import GraphIn
 from dagster._core.definitions.metadata import ArbitraryMetadataMapping, RawMetadataMapping
@@ -50,6 +51,7 @@
 from dagster._core.errors import DagsterInvalidDefinitionError
 from dagster._core.storage.tags import KIND_PREFIX
 from dagster._core.types.dagster_type import DagsterType
+from dagster._serdes import serialize_value
 from dagster._utils.tags import normalize_tags
 from dagster._utils.warnings import disable_dagster_warnings
 
@@ -299,6 +301,13 @@ def downstream_asset(conditional_asset):
         **{f"{KIND_PREFIX}{kind}": "" for kind in kinds or []},
     }
 
+    internal_freshness_policy = kwargs.get("internal_freshness_policy")
+    if internal_freshness_policy:
+        metadata = {
+            **(metadata or {}),
+            INTERNAL_FRESHNESS_POLICY_METADATA_KEY: serialize_value(internal_freshness_policy),
+        }
+
     only_allow_hidden_params_in_kwargs(asset, kwargs)
 
     args = AssetDecoratorArgs(

python_modules/dagster/dagster_tests/asset_defs_tests/test_internal_freshness.py

@@ -2,6 +2,8 @@
 
 from dagster._core.definitions.asset_key import AssetKey
 from dagster._core.definitions.asset_spec import AssetSpec, attach_internal_freshness_policy
+from dagster._core.definitions.assets import AssetsDefinition
+from dagster._core.definitions.decorators.asset_decorator import asset
 from dagster._core.definitions.definitions_class import Definitions
 from dagster._core.definitions.freshness import (
     INTERNAL_FRESHNESS_POLICY_METADATA_KEY,
@@ -16,6 +18,26 @@
 )
 
 
+def test_asset_decorator_with_internal_freshness_policy() -> None:
+    """Can we define an asset from decorator with an internal freshness policy?"""
+
+    @asset(
+        internal_freshness_policy=TimeWindowFreshnessPolicy.from_timedeltas(
+            fail_window=timedelta(minutes=10), warn_window=timedelta(minutes=5)
+        )
+    )
+    def asset_with_internal_freshness_policy():
+        pass
+
+    spec = asset_with_internal_freshness_policy.get_asset_spec()
+    policy = spec.metadata.get(INTERNAL_FRESHNESS_POLICY_METADATA_KEY)
+    assert policy is not None
+    deserialized = deserialize_value(policy)
+    assert isinstance(deserialized, TimeWindowFreshnessPolicy)
+    assert deserialized.fail_window == SerializableTimeDelta.from_timedelta(timedelta(minutes=10))
+    assert deserialized.warn_window == SerializableTimeDelta.from_timedelta(timedelta(minutes=5))
+
+
 def test_asset_spec_with_internal_freshness_policy() -> None:
     """Can we define an asset spec with an internal freshness policy?"""
 
@@ -98,21 +120,32 @@ def assert_freshness_policy(spec, expected_fail_window, expected_warn_window=Non
 
 
 def test_map_asset_specs_attach_internal_freshness_policy() -> None:
-    """Can we map attach_internal_freshness_policy over a selection of asset specs?"""
-    asset_specs = [AssetSpec(key="foo"), AssetSpec(key="bar"), AssetSpec(key="baz")]
+    """Can we map attach_internal_freshness_policy over a selection of assets and asset specs?"""
+
+    @asset
+    def foo_asset():
+        pass
+
+    asset_specs = [foo_asset, AssetSpec(key="bar"), AssetSpec(key="baz")]
+    defs: Definitions = Definitions(assets=asset_specs)
+
     freshness_policy = TimeWindowFreshnessPolicy.from_timedeltas(
         fail_window=timedelta(minutes=10), warn_window=timedelta(minutes=5)
     )
-    defs: Definitions = Definitions(assets=asset_specs)
     mapped_defs = defs.map_asset_specs(
         func=lambda spec: attach_internal_freshness_policy(spec, freshness_policy)
     )
 
-    assets = mapped_defs.assets
-    assert len(assets) == 3
-    for asset in assets:
-        assert INTERNAL_FRESHNESS_POLICY_METADATA_KEY in asset.metadata
-        policy = deserialize_value(asset.metadata[INTERNAL_FRESHNESS_POLICY_METADATA_KEY])
+    assets_and_specs = mapped_defs.assets
+    assert len(assets_and_specs) == 3
+    for asset_or_spec in assets_and_specs:
+        spec = (
+            asset_or_spec.get_asset_spec()
+            if isinstance(asset_or_spec, AssetsDefinition)
+            else asset_or_spec
+        )
+        assert INTERNAL_FRESHNESS_POLICY_METADATA_KEY in spec.metadata
+        policy = deserialize_value(spec.metadata[INTERNAL_FRESHNESS_POLICY_METADATA_KEY])
         assert isinstance(policy, TimeWindowFreshnessPolicy)
         assert policy.fail_window == SerializableTimeDelta.from_timedelta(timedelta(minutes=10))
         assert policy.warn_window == SerializableTimeDelta.from_timedelta(timedelta(minutes=5))