fixup! feat: redesign PartitionedAssetTimetable

Author: Lee-W

airflow-core/src/airflow/models/asset.py

@@ -383,9 +383,6 @@ def to_serialized(self) -> SerializedAsset:
             group=self.group,
             extra=self.extra,
             watchers=[],
-            # AIP-76: we allow user to specify per asset partition_mapper, but this exists in timetable
-            # instead of asset model itself. Thus, from asset model to SerializedAsset will always set
-            # partition_mapper to None
         )
 
     def add_trigger(self, trigger: Trigger, watcher_name: str):

airflow-core/src/airflow/serialization/decoders.py

@@ -105,7 +105,6 @@ def _decode_asset(var: dict[str, Any]):
             )
             for watcher in watchers
         ],
-        partition_mapper=var.get("partition_mapper", None),
     )
 
 
@@ -128,10 +127,8 @@ def decode_asset_like(var: dict[str, Any]) -> SerializedAssetBase:
         case DAT.ASSET_ANY:
             return SerializedAssetAny([decode_asset_like(x) for x in var["objects"]])
         case DAT.ASSET_ALIAS:
-            # TODO: (AIP-76) partition_mapper
             return SerializedAssetAlias(name=var["name"], group=var["group"])
         case DAT.ASSET_REF:
-            # TODO: (AIP-76) partition_mapper
             if "name" in var:
                 return SerializedAssetNameRef(**var)
             return SerializedAssetUriRef(**var)

airflow-core/src/airflow/serialization/definitions/assets.py

@@ -120,7 +120,6 @@ class SerializedAsset(SerializedAssetBase):
     group: str
     extra: dict[str, Any]
     watchers: MutableSequence[SerializedAssetWatcher]
-    partition_mapper: dict | None = None
 
     def as_expression(self) -> Any:
         """
@@ -164,7 +163,6 @@ def asprofile(self) -> AssetProfile:
         return AssetProfile(name=self.name or None, uri=self.uri or None, type="Asset")
 
 
-# TODO: (AIP-76) add partition_mapper: dict | None = None
 class SerializedAssetRef(SerializedAssetBase, AttrsInstance):
     """Serialized representation of an asset reference."""
 
@@ -193,7 +191,6 @@ def iter_dag_dependencies(self, *, source: str = "", target: str = "") -> Iterat
         )
 
 
-# TODO: (AIP-76) add partition_mapper: dict | None = None
 @attrs.define(hash=True)
 class SerializedAssetNameRef(SerializedAssetRef):
     """Serialized representation of an asset reference by name."""
@@ -203,7 +200,6 @@ class SerializedAssetNameRef(SerializedAssetRef):
     _dependency_type = "asset-name-ref"
 
 
-# TODO: (AIP-76) add partition_mapper: dict | None = None
 @attrs.define(hash=True)
 class SerializedAssetUriRef(SerializedAssetRef):
     """Serialized representation of an asset reference by URI."""
@@ -213,7 +209,6 @@ class SerializedAssetUriRef(SerializedAssetRef):
     _dependency_type = "asset-uri-ref"
 
 
-# TODO: (AIP-76) add partition_mapper: dict | None = None
 @attrs.define
 class SerializedAssetAlias(SerializedAssetBase):
     """Serialized representation of an asset alias."""

airflow-core/src/airflow/serialization/encoders.py

@@ -168,42 +168,17 @@ def encode_asset_like(a: BaseAsset | SerializedAssetBase) -> dict[str, Any]:
     d: dict[str, Any]
     match a:
         case Asset() | SerializedAsset():
-            d = {
-                "__type": DAT.ASSET,
-                "name": a.name,
-                "uri": a.uri,
-                "group": a.group,
-                "extra": a.extra,
-            }
+            d = {"__type": DAT.ASSET, "name": a.name, "uri": a.uri, "group": a.group, "extra": a.extra}
             if a.watchers:
                 d["watchers"] = [{"name": w.name, "trigger": encode_trigger(w.trigger)} for w in a.watchers]
-
-            if a.partition_mapper:
-                if isinstance(a.partition_mapper, dict):
-                    d["partition_mapper"] = a.partition_mapper
-                else:
-                    d["partition_mapper"] = encode_partition_mapper(a.partition_mapper)
             return d
         case AssetAlias() | SerializedAssetAlias():
-            return {
-                "__type": DAT.ASSET_ALIAS,
-                "name": a.name,
-                "group": a.group,
-                # TODO: (AIP_76) add partition_mapper
-                # if a.partition_mapper:
-                #     if isinstance(a.partition_mapper, dict):
-                #         d["partition_mapper"] = a.partition_mapper
-                #     else:
-                #         d["partition_mapper"] = encode_partition_mapper(a.partition_mapper)
-            }
+            return {"__type": DAT.ASSET_ALIAS, "name": a.name, "group": a.group}
         case AssetAll() | SerializedAssetAll():
-            # TODO: (AIP_76) add partition_mapper
             return {"__type": DAT.ASSET_ALL, "objects": [encode_asset_like(x) for x in a.objects]}
         case AssetAny() | SerializedAssetAny():
-            # TODO: (AIP_76) add partition_mapper
             return {"__type": DAT.ASSET_ANY, "objects": [encode_asset_like(x) for x in a.objects]}
         case AssetRef() | SerializedAssetRef():
-            # TODO: (AIP_76) add partition_mapper
             return {"__type": DAT.ASSET_REF, **attrs.asdict(a)}
     raise ValueError(f"serialization not implemented for {type(a).__name__!r}")
 

airflow-core/src/airflow/timetables/simple.py

@@ -289,3 +289,7 @@ def deserialize(cls, data: dict[str, Any]) -> Timetable:
             else None,
         )
         return timetable
+
+    def get_partition_mapper(self, *, name: str, uri: str) -> PartitionMapper:
+        # TODO: implement
+        ...

airflow-core/tests/unit/jobs/test_scheduler_job.py

@@ -8702,11 +8702,9 @@ def test_partitioned_dag_run_with_customized_mapper(
         with dag_maker(
             dag_id="asset-event-consumer",
             schedule=PartitionedAssetTimetable(
+                assets=Asset(name="asset-1"),
                 # TODO: (AIP-76) fix typing
-                assets=Asset(
-                    name="asset-1",
-                    partition_mapper=Key1Mapper(),  # type: ignore[call-overload]
-                ),
+                default_partition_mapper=Key1Mapper(),  # type: ignore[call-overload]
             ),
             session=session,
         ):
@@ -8858,6 +8856,10 @@ def test_consumer_dag_listen_to_two_partitioned_asset_with_key_1_mapper(
                     Asset(name="asset-1", partition_mapper=Key1Mapper())  # type: ignore[call-overload]
                     & Asset(name="asset-2", partition_mapper=Key1Mapper())  # type: ignore[call-overload]
                 ),
+                partition_mapper_mapping={
+                    Asset(name="asset-1"): Key1Mapper(),
+                    Asset(name="asset-2"): Key1Mapper(),
+                },
             ),
             session=session,
         ):

airflow-core/tests/unit/models/test_taskinstance.py

@@ -3052,10 +3052,7 @@ def test_when_dag_run_has_partition_and_downstreams_listening_then_tables_popula
     with dag_maker(
         dag_id="asset_event_listener",
         schedule=PartitionedAssetTimetable(
-            assets=Asset(
-                name="hello",
-                partition_mapper=IdentityMapper(),
-            )
+            assets=Asset(name="hello"), default_partition_mapper=IdentityMapper()
         ),
         session=session,
     ):

task-sdk/src/airflow/sdk/definitions/asset/__init__.py

@@ -37,7 +37,6 @@
     from pydantic.types import JsonValue
     from typing_extensions import Self
 
-    from airflow.sdk import PartitionMapper
     from airflow.sdk.api.datamodels._generated import AssetProfile
     from airflow.sdk.io.path import ObjectStoragePath
     from airflow.triggers.base import BaseEventTrigger
@@ -230,8 +229,6 @@ class BaseAsset:
     :meta private:
     """
 
-    partition_mapper: PartitionMapper | None = None
-
     def __or__(self, other: BaseAsset) -> BaseAsset:
         if not isinstance(other, BaseAsset):
             return NotImplemented
@@ -281,7 +278,6 @@ class Asset(os.PathLike, BaseAsset):
     watchers: list[AssetWatcher] = attrs.field(
         factory=list,
     )
-    partition_mapper: PartitionMapper | None = None
 
     asset_type: ClassVar[str] = "asset"
     __version__: ClassVar[int] = 1
@@ -295,7 +291,6 @@ def __init__(
         group: str = ...,
         extra: dict[str, JsonValue] | None = None,
         watchers: list[AssetWatcher] = ...,
-        partition_mapper: PartitionMapper | None = None,
     ) -> None:
         """Canonical; both name and uri are provided."""
 
@@ -307,7 +302,6 @@ def __init__(
         group: str = ...,
         extra: dict[str, JsonValue] | None = None,
         watchers: list[AssetWatcher] = ...,
-        partition_mapper: PartitionMapper | None = None,
     ) -> None:
         """It's possible to only provide the name, either by keyword or as the only positional argument."""
 
@@ -319,7 +313,6 @@ def __init__(
         group: str = ...,
         extra: dict[str, JsonValue] | None = None,
         watchers: list[AssetWatcher] = ...,
-        partition_mapper: PartitionMapper | None = None,
     ) -> None:
         """It's possible to only provide the URI as a keyword argument."""
 
@@ -331,7 +324,6 @@ def __init__(
         group: str | None = None,
         extra: dict[str, JsonValue] | None = None,
         watchers: list[AssetWatcher] | None = None,
-        partition_mapper: PartitionMapper | None = None,
     ) -> None:
         if name is None and uri is None:
             raise TypeError("Asset() requires either 'name' or 'uri'")
@@ -353,12 +345,9 @@ def __init__(
             kwargs["extra"] = extra
         if watchers is not None:
             kwargs["watchers"] = watchers
-        if partition_mapper is not None:
-            kwargs["partition_mapper"] = partition_mapper
 
         self.__attrs_init__(name=name, uri=uri, **kwargs)
 
-    # TODO: AIP-76: support something like Asset.ref(name=..., partition_mapper=...)?
     @overload
     @staticmethod
     def ref(*, name: str) -> AssetNameRef: ...