AIP-76: Add PartitionAtRuntime authoring API to Task SDK

airflow-core/docs/authoring-and-scheduling/assets.rst

@@ -188,6 +188,8 @@ Declaring an ``@asset`` automatically creates:
 * A ``DAG`` with *dag_id* set to the function name.
 * A task inside the ``DAG`` with *task_id* set to the function name, and *outlet* to the created ``Asset``.
 
+The parameter names ``self``, ``context``, and ``outlet_events`` are **reserved** in an ``@asset`` function: they are populated by Airflow at runtime (with the asset itself, the execution context, and the outlet event accessor respectively) and are never treated as inlet asset references.
+
 Attaching extra information to an emitting asset event
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 

airflow-core/src/airflow/api_fastapi/execution_api/datamodels/asset_event.py

@@ -19,6 +19,7 @@
 
 from datetime import datetime
 
+from pydantic import Field
 from pydantic.types import JsonValue
 
 from airflow.api_fastapi.core_api.base import BaseModel, StrictBaseModel
@@ -54,6 +55,7 @@ class AssetEventResponse(BaseModel):
     source_run_id: str | None = None
     source_map_index: int | None = None
     partition_key: str | None = None
+    partition_keys: list[str] = Field(default_factory=list)
 
 
 class AssetEventsResponse(BaseModel):

airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py

@@ -32,6 +32,7 @@
     AddDagEndpoint,
     AddDagRunDetailEndpoint,
     AddNoteField,
+    AddOutletPartitionKeysField,
     AddPartitionKeyField,
     AddRunAfterField,
     AddTaskInstanceStartDateField,
@@ -54,6 +55,7 @@
     Version(
         "2026-04-06",
         AddPartitionKeyField,
+        AddOutletPartitionKeysField,
         MovePreviousRunEndpoint,
         AddDagRunDetailEndpoint,
         MakeDagRunStartDateNullable,

airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_04_06.py

@@ -61,6 +61,23 @@ def remove_partition_key_from_asset_events(response: ResponseInfo) -> None:  # t
             elem.pop("partition_key", None)
 
 
+class AddOutletPartitionKeysField(VersionChange):
+    """Add the `partition_keys` field to AssetEventResponse for runtime-partitioned outlet events."""
+
+    description = __doc__
+
+    instructions_to_migrate_to_previous_version = (
+        schema(AssetEventResponse).field("partition_keys").didnt_exist,
+    )
+
+    @convert_response_to_previous_version_for(AssetEventsResponse)  # type: ignore[arg-type]
+    def remove_partition_keys_from_asset_events(response: ResponseInfo) -> None:  # type: ignore[misc]
+        """Remove the `partition_keys` field from each asset event when converting to the previous version."""
+        events = response.body["asset_events"]
+        for elem in events:
+            elem.pop("partition_keys", None)
+
+
 class MovePreviousRunEndpoint(VersionChange):
     """Add new previous-run endpoint and migrate old endpoint."""
 

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

@@ -55,6 +55,7 @@
 from airflow.sdk.definitions.partition_mappers.temporal import StartOfHourMapper
 from airflow.sdk.definitions.timetables.assets import (
     AssetTriggeredTimetable,
+    PartitionAtRuntime,
     PartitionedAssetTimetable,
 )
 from airflow.sdk.definitions.timetables.simple import ContinuousTimetable, NullTimetable, OnceTimetable
@@ -293,6 +294,7 @@ class _Serializer:
         MultipleCronTriggerTimetable: "airflow.timetables.trigger.MultipleCronTriggerTimetable",
         NullTimetable: "airflow.timetables.simple.NullTimetable",
         OnceTimetable: "airflow.timetables.simple.OnceTimetable",
+        PartitionAtRuntime: "airflow.timetables.simple.PartitionAtRuntime",
         PartitionedAssetTimetable: "airflow.timetables.simple.PartitionedAssetTimetable",
     }
 
@@ -318,7 +320,10 @@ def serialize_timetable(self, timetable: BaseTimetable | CoreTimetable) -> dict[
     @serialize_timetable.register(ContinuousTimetable)
     @serialize_timetable.register(NullTimetable)
     @serialize_timetable.register(OnceTimetable)
-    def _(self, timetable: ContinuousTimetable | NullTimetable | OnceTimetable) -> dict[str, Any]:
+    @serialize_timetable.register(PartitionAtRuntime)
+    def _(
+        self, timetable: ContinuousTimetable | NullTimetable | OnceTimetable | PartitionAtRuntime
+    ) -> dict[str, Any]:
         return {}
 
     @serialize_timetable.register

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

@@ -218,6 +218,13 @@ class Timetable(Protocol):
     instead of the traditional logic based on logical dates and data intervals.
     """
 
+    partitioned_at_runtime: bool = False
+    """Whether this timetable defers partition selection to task runtime.
+
+    *True* for :class:`~airflow.timetables.simple.PartitionAtRuntime`;
+    downstream code can branch on this flag instead of using ``isinstance``.
+    """
+
     @classmethod
     def deserialize(cls, data: dict[str, Any]) -> Timetable:
         """

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

@@ -93,6 +93,7 @@ class NullTimetable(_TrivialTimetable):
     """
 
     can_be_scheduled = False  # TODO (GH-52141): Find a way to keep this and one in Core in sync.
+    partitioned_at_runtime = False
     description: str = "Never, external triggers only"
 
     @property
@@ -183,6 +184,21 @@ def next_dagrun_info(
         return DagRunInfo.interval(start, end)
 
 
+class PartitionAtRuntime(NullTimetable):
+    """
+    Timetable that never schedules anything; partition keys are set at runtime.
+
+    This corresponds to ``schedule=PartitionAtRuntime()``.
+    """
+
+    description: str = "Never, partition key(s) set at runtime"
+    partitioned_at_runtime = True
+
+    @property
+    def summary(self) -> str:
+        return "PartitionAtRuntime"
+
+
 class AssetTriggeredTimetable(_TrivialTimetable):
     """
     Timetable that never schedules anything.

airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_asset_events.py

@@ -124,6 +124,7 @@ def test_get_by_asset(self, uri, name, client):
                     },
                     "timestamp": "2021-01-01T00:00:00Z",
                     "partition_key": None,
+                    "partition_keys": [],
                 },
                 {
                     "id": 2,
@@ -141,6 +142,7 @@ def test_get_by_asset(self, uri, name, client):
                     "created_dagruns": [],
                     "timestamp": "2021-01-02T00:00:00Z",
                     "partition_key": None,
+                    "partition_keys": [],
                 },
                 {
                     "id": 3,
@@ -158,6 +160,7 @@ def test_get_by_asset(self, uri, name, client):
                     "created_dagruns": [],
                     "timestamp": "2021-01-03T00:00:00Z",
                     "partition_key": None,
+                    "partition_keys": [],
                 },
             ]
         }
@@ -195,6 +198,7 @@ def test_get_by_asset_with_after_filter(self, uri, name, client):
                     "created_dagruns": [],
                     "timestamp": "2021-01-02T00:00:00Z",
                     "partition_key": None,
+                    "partition_keys": [],
                 },
                 {
                     "id": 3,
@@ -212,6 +216,7 @@ def test_get_by_asset_with_after_filter(self, uri, name, client):
                     "created_dagruns": [],
                     "timestamp": "2021-01-03T00:00:00Z",
                     "partition_key": None,
+                    "partition_keys": [],
                 },
             ]
         }
@@ -249,6 +254,7 @@ def test_get_by_asset_with_before_filter(self, uri, name, client):
                     "created_dagruns": [],
                     "timestamp": "2021-01-01T00:00:00Z",
                     "partition_key": None,
+                    "partition_keys": [],
                 },
                 {
                     "id": 2,
@@ -266,6 +272,7 @@ def test_get_by_asset_with_before_filter(self, uri, name, client):
                     "created_dagruns": [],
                     "timestamp": "2021-01-02T00:00:00Z",
                     "partition_key": None,
+                    "partition_keys": [],
                 },
             ]
         }
@@ -308,6 +315,7 @@ def test_get_by_asset_with_before_and_after_filters(self, uri, name, client):
                     "created_dagruns": [],
                     "timestamp": "2021-01-02T00:00:00Z",
                     "partition_key": None,
+                    "partition_keys": [],
                 },
             ]
         }
@@ -345,6 +353,7 @@ def test_get_by_asset_with_descending_order(self, uri, name, client):
                     "created_dagruns": [],
                     "timestamp": "2021-01-03T00:00:00Z",
                     "partition_key": None,
+                    "partition_keys": [],
                 },
                 {
                     "id": 2,
@@ -362,6 +371,7 @@ def test_get_by_asset_with_descending_order(self, uri, name, client):
                     "created_dagruns": [],
                     "timestamp": "2021-01-02T00:00:00Z",
                     "partition_key": None,
+                    "partition_keys": [],
                 },
                 {
                     "id": 1,
@@ -379,6 +389,7 @@ def test_get_by_asset_with_descending_order(self, uri, name, client):
                     "created_dagruns": [],
                     "timestamp": "2021-01-01T00:00:00Z",
                     "partition_key": None,
+                    "partition_keys": [],
                 },
             ]
         }
@@ -416,6 +427,7 @@ def test_get_by_asset_get_first(self, uri, name, client):
                     "created_dagruns": [],
                     "timestamp": "2021-01-01T00:00:00Z",
                     "partition_key": None,
+                    "partition_keys": [],
                 },
             ]
         }
@@ -453,6 +465,7 @@ def test_get_by_asset_get_last(self, uri, name, client):
                     "created_dagruns": [],
                     "timestamp": "2021-01-03T00:00:00Z",
                     "partition_key": None,
+                    "partition_keys": [],
                 },
             ]
         }
@@ -484,6 +497,7 @@ def test_get_by_asset(self, client):
                     "created_dagruns": [],
                     "timestamp": "2021-01-01T00:00:00Z",
                     "partition_key": None,
+                    "partition_keys": [],
                 },
                 {
                     "id": 2,
@@ -501,6 +515,7 @@ def test_get_by_asset(self, client):
                     "created_dagruns": [],
                     "timestamp": "2021-01-02T00:00:00Z",
                     "partition_key": None,
+                    "partition_keys": [],
                 },
                 {
                     "id": 3,
@@ -518,6 +533,7 @@ def test_get_by_asset(self, client):
                     "created_dagruns": [],
                     "timestamp": "2021-01-03T00:00:00Z",
                     "partition_key": None,
+                    "partition_keys": [],
                 },
             ]
         }

task-sdk/docs/api.rst

@@ -209,6 +209,8 @@ Timetables
 
 .. autoapiclass:: airflow.sdk.MultipleCronTriggerTimetable
 
+.. autoapiclass:: airflow.sdk.PartitionAtRuntime
+
 .. autoapiclass:: airflow.sdk.PartitionedAssetTimetable
 
 

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

@@ -59,6 +59,7 @@
     "ObjectStoragePath",
     "Param",
     "ParamsDict",
+    "PartitionAtRuntime",
     "PartitionedAssetTimetable",
     "PartitionMapper",
     "PokeReturnValue",
@@ -118,7 +119,13 @@
     from airflow.sdk.bases.skipmixin import SkipMixin
     from airflow.sdk.bases.xcom import BaseXCom
     from airflow.sdk.configuration import AirflowSDKConfigParser
-    from airflow.sdk.definitions.asset import Asset, AssetAlias, AssetAll, AssetAny, AssetWatcher
+    from airflow.sdk.definitions.asset import (
+        Asset,
+        AssetAlias,
+        AssetAll,
+        AssetAny,
+        AssetWatcher,
+    )
     from airflow.sdk.definitions.asset.decorators import asset
     from airflow.sdk.definitions.asset.metadata import Metadata
     from airflow.sdk.definitions.callback import AsyncCallback, SyncCallback
@@ -154,6 +161,7 @@
     from airflow.sdk.definitions.template import literal
     from airflow.sdk.definitions.timetables.assets import (
         AssetOrTimeSchedule,
+        PartitionAtRuntime,
         PartitionedAssetTimetable,
     )
     from airflow.sdk.definitions.timetables.events import EventsTimetable
@@ -215,6 +223,7 @@
     "ObjectStoragePath": ".io.path",
     "Param": ".definitions.param",
     "ParamsDict": ".definitions.param",
+    "PartitionAtRuntime": ".definitions.timetables.assets",
     "PartitionedAssetTimetable": ".definitions.timetables.assets",
     "PartitionMapper": ".definitions.partition_mappers.base",
     "PokeReturnValue": ".bases.sensor",

task-sdk/src/airflow/sdk/__init__.pyi

@@ -86,6 +86,7 @@ from airflow.sdk.definitions.taskgroup import TaskGroup as TaskGroup
 from airflow.sdk.definitions.template import literal as literal
 from airflow.sdk.definitions.timetables.assets import (
     AssetOrTimeSchedule,
+    PartitionAtRuntime,
     PartitionedAssetTimetable,
 )
 from airflow.sdk.definitions.timetables.events import EventsTimetable
@@ -145,6 +146,7 @@ __all__ = [
     "ObjectStoragePath",
     "Param",
     "PokeReturnValue",
+    "PartitionAtRuntime",
     "PartitionedAssetTimetable",
     "PartitionMapper",
     "ProductMapper",

task-sdk/src/airflow/sdk/api/datamodels/_generated.py

@@ -660,6 +660,7 @@ class AssetEventResponse(BaseModel):
     source_run_id: Annotated[str | None, Field(title="Source Run Id")] = None
     source_map_index: Annotated[int | None, Field(title="Source Map Index")] = None
     partition_key: Annotated[str | None, Field(title="Partition Key")] = None
+    partition_keys: Annotated[list[str] | None, Field(title="Partition Keys")] = None
 
 
 class AssetEventsResponse(BaseModel):

task-sdk/src/airflow/sdk/bases/timetable.py

@@ -47,6 +47,14 @@ class BaseTimetable:
 
     asset_condition: BaseAsset | None = None
 
+    partitioned_at_runtime: bool = False
+    """
+    Whether this timetable defers partition selection to task runtime.
+
+    *True* for :class:`~airflow.sdk.PartitionAtRuntime`; downstream code can
+    branch on this flag instead of using ``isinstance``.
+    """
+
     def validate(self) -> None:
         """
         Validate the timetable is correctly specified.