Add asset partition sensor

airflow-core/newsfragments/67941.feature.rst

@@ -0,0 +1 @@
+Add an asset partition sensor for waiting on a specific asset event partition.

airflow-core/src/airflow/example_dags/example_asset_partition.py

@@ -19,6 +19,7 @@
 
 from typing import TYPE_CHECKING
 
+from airflow.providers.standard.sensors.asset import AssetPartitionSensor
 from airflow.sdk import (
     DAG,
     AllowedKeyMapper,
@@ -112,6 +113,39 @@ def combine_player_stats(dag_run=None):
     combine_player_stats()
 
 
+with DAG(
+    dag_id="wait_for_combined_player_stats_partition",
+    schedule="@hourly",
+    catchup=False,
+    tags=["example", "player-stats", "sensor"],
+):
+    """
+    Wait for a specific ``combined_player_stats`` partition from a time-scheduled Dag.
+
+    ``AssetPartitionSensor`` bridges time-based scheduling and asset partitioning: this hourly
+    Dag blocks until the combined-stats partition for its own data interval has an asset event.
+    ``after`` bounds the lookup to the current interval so a stale event with the same partition
+    key from an earlier run does not satisfy the wait.
+    """
+
+    wait_for_combined_partition = AssetPartitionSensor(
+        task_id="wait_for_combined_partition",
+        asset=combined_player_stats,
+        partition_key="{{ data_interval_start.strftime('%Y-%m-%dT%H') }}",
+        after="{{ data_interval_start }}",
+        deferrable=True,
+    )
+
+    @task
+    def report_partition_ready(dag_run=None):
+        """Run once the awaited combined-stats partition is available."""
+        if TYPE_CHECKING:
+            assert dag_run
+        print(f"combined_player_stats partition ready for {dag_run.logical_date}")
+
+    wait_for_combined_partition >> report_partition_ready()
+
+
 @asset(
     uri="file://analytics/player-stats/computed-player-odds.csv",
     # Fallback to IdentityMapper if no partition_mapper is specified.

airflow-core/src/airflow/jobs/triggerer_job_runner.py

@@ -73,6 +73,8 @@
     DeleteXCom,
     DRCount,
     ErrorResponse,
+    GetAssetEventByAsset,
+    GetAssetEventByAssetAlias,
     GetAssetStateStoreByName,
     GetAssetStateStoreByUri,
     GetConnection,
@@ -108,6 +110,8 @@
     handle_delete_asset_state_store_by_uri,
     handle_delete_variable,
     handle_delete_xcom,
+    handle_get_asset_event_by_asset,
+    handle_get_asset_event_by_asset_alias,
     handle_get_asset_state_store_by_name,
     handle_get_asset_state_store_by_uri,
     handle_get_connection,
@@ -641,6 +645,10 @@ def _handle_request(self, msg: ToTriggerSupervisor, log: FilteringBoundLogger, r
             resp, dump_opts = handle_get_dr_count(self.client, msg)
         elif isinstance(msg, GetDagRunState):
             resp, dump_opts = handle_get_dag_run_state(self.client, msg)
+        elif isinstance(msg, GetAssetEventByAsset):
+            resp, dump_opts = handle_get_asset_event_by_asset(self.client, msg)
+        elif isinstance(msg, GetAssetEventByAssetAlias):
+            resp, dump_opts = handle_get_asset_event_by_asset_alias(self.client, msg)
 
         elif isinstance(msg, GetTICount):
             resp, dump_opts = handle_get_ti_count(self.client, msg)

providers/standard/provider.yaml

@@ -106,6 +106,7 @@ sensors:
       - airflow.providers.standard.sensors.python
       - airflow.providers.standard.sensors.filesystem
       - airflow.providers.standard.sensors.external_task
+      - airflow.providers.standard.sensors.asset
 hooks:
   - integration-name: Standard
     python-modules:
@@ -120,6 +121,7 @@ triggers:
       - airflow.providers.standard.triggers.file
       - airflow.providers.standard.triggers.temporal
       - airflow.providers.standard.triggers.hitl
+      - airflow.providers.standard.triggers.asset
 
 extra-links:
   - airflow.providers.standard.operators.trigger_dagrun.TriggerDagRunLink

providers/standard/src/airflow/providers/standard/exceptions.py

@@ -67,3 +67,7 @@ class HITLTimeoutError(HITLTriggerEventError):
 
 class HITLRejectException(AirflowException):
     """Raised when an ApprovalOperator receives a "Reject" response when fail_on_reject is set to True."""
+
+
+class AssetPartitionTriggerEventError(AirflowException):
+    """Raised when AssetPartitionTrigger reports an error event."""

providers/standard/src/airflow/providers/standard/get_provider_info.py

@@ -75,6 +75,7 @@ def get_provider_info():
                     "airflow.providers.standard.sensors.python",
                     "airflow.providers.standard.sensors.filesystem",
                     "airflow.providers.standard.sensors.external_task",
+                    "airflow.providers.standard.sensors.asset",
                 ],
             }
         ],
@@ -96,6 +97,7 @@ def get_provider_info():
                     "airflow.providers.standard.triggers.file",
                     "airflow.providers.standard.triggers.temporal",
                     "airflow.providers.standard.triggers.hitl",
+                    "airflow.providers.standard.triggers.asset",
                 ],
             }
         ],

providers/standard/src/airflow/providers/standard/sensors/asset.py

@@ -0,0 +1,123 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import datetime
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Any
+
+from airflow.providers.common.compat.sdk import (
+    AirflowOptionalProviderFeatureException,
+    BaseSensorOperator,
+    conf,
+    timezone,
+)
+from airflow.providers.standard.version_compat import AIRFLOW_V_3_4_PLUS
+
+if not AIRFLOW_V_3_4_PLUS:
+    raise AirflowOptionalProviderFeatureException("Asset partition sensor needs Airflow 3.4+.")
+
+from airflow.providers.standard.exceptions import AssetPartitionTriggerEventError
+from airflow.providers.standard.triggers.asset import AssetPartitionTrigger
+
+if TYPE_CHECKING:
+    from airflow.providers.common.compat.sdk import Asset, Context
+
+
+class AssetPartitionSensor(BaseSensorOperator):
+    """
+    Wait for an asset event with the given partition key.
+
+    :param asset: asset to wait for.
+    :param partition_key: partition key for the asset event to wait for.
+    :param after: only match events whose timestamp is at or after this (timezone-aware)
+        datetime. When unset, any historical event with the partition key satisfies the
+        wait. Bound it (e.g. ``after="{{ data_interval_start }}"``) when the partition key
+        can be reused across events, so a stale event from an earlier run is not matched.
+    :param deferrable: If waiting for completion, whether to defer the task until done.
+    """
+
+    template_fields: Sequence[str] = ("partition_key", "after")
+    ui_color = "#e6f1f2"
+
+    def __init__(
+        self,
+        *,
+        asset: Asset,
+        partition_key: str,
+        after: datetime.datetime | str | None = None,
+        deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.asset = asset
+        self.partition_key = partition_key
+        self.after = after
+        self.deferrable = deferrable
+
+    def poke(self, context: Context) -> bool:
+        from airflow.sdk.exceptions import AirflowRuntimeError
+        from airflow.sdk.execution_time.comms import AssetEventsResult, ErrorResponse, GetAssetEventByAsset
+        from airflow.sdk.execution_time.task_runner import SUPERVISOR_COMMS
+
+        self.log.info("Poking for asset event: asset=%s, partition_key=%s", self.asset, self.partition_key)
+        after = timezone.parse(self.after) if isinstance(self.after, str) else self.after
+        response = SUPERVISOR_COMMS.send(
+            GetAssetEventByAsset(
+                name=self.asset.name,
+                uri=self.asset.uri,
+                partition_key=self.partition_key,
+                after=after,
+                ascending=False,
+                limit=1,
+            )
+        )
+        if isinstance(response, ErrorResponse):
+            raise AirflowRuntimeError(response)
+        if TYPE_CHECKING:
+            assert isinstance(response, AssetEventsResult)
+        return bool(response and response.asset_events)
+
+    def execute(self, context: Context) -> None:
+        if not self.deferrable:
+            super().execute(context=context)
+            return
+
+        if not self.poke(context=context):
+            self.defer(
+                timeout=datetime.timedelta(seconds=self.timeout),
+                trigger=AssetPartitionTrigger(
+                    asset_name=self.asset.name,
+                    asset_uri=self.asset.uri,
+                    partition_key=self.partition_key,
+                    after=self.after,
+                    poke_interval=self.poke_interval,
+                ),
+                method_name="execute_complete",
+            )
+
+    def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> None:
+        if event and event.get("status") == "success":
+            self.log.info(
+                "Asset partition event found: asset=%s, partition_key=%s",
+                self.asset,
+                self.partition_key,
+            )
+            return
+        message = event.get("message") if event else "Trigger completed without an event"
+        raise AssetPartitionTriggerEventError(message)

providers/standard/src/airflow/providers/standard/triggers/asset.py

@@ -0,0 +1,107 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import asyncio
+from collections.abc import AsyncIterator
+from typing import TYPE_CHECKING, Any
+
+from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException, timezone
+from airflow.providers.standard.version_compat import AIRFLOW_V_3_4_PLUS
+
+if not AIRFLOW_V_3_4_PLUS:
+    raise AirflowOptionalProviderFeatureException("Asset partition sensor needs Airflow 3.4+.")
+
+from airflow.triggers.base import BaseTrigger, TriggerEvent
+
+if TYPE_CHECKING:
+    import datetime
+
+
+class AssetPartitionTrigger(BaseTrigger):
+    """
+    Trigger when an asset event exists for the given partition key.
+
+    :param asset_name: name of the asset to wait for.
+    :param asset_uri: URI of the asset to wait for.
+    :param partition_key: partition key for the asset event to wait for.
+    :param after: only match events whose timestamp is at or after this (timezone-aware)
+        datetime. Leave unset to match any event with the partition key.
+    :param poke_interval: polling interval in seconds.
+    """
+
+    def __init__(
+        self,
+        *,
+        asset_name: str | None,
+        asset_uri: str | None,
+        partition_key: str,
+        after: datetime.datetime | str | None = None,
+        poke_interval: float = 5.0,
+    ) -> None:
+        super().__init__()
+        self.asset_name = asset_name
+        self.asset_uri = asset_uri
+        self.partition_key = partition_key
+        self.after = after
+        self.poke_interval = poke_interval
+
+    def serialize(self) -> tuple[str, dict[str, Any]]:
+        """Serialize AssetPartitionTrigger arguments and classpath."""
+        return (
+            "airflow.providers.standard.triggers.asset.AssetPartitionTrigger",
+            {
+                "asset_name": self.asset_name,
+                "asset_uri": self.asset_uri,
+                "partition_key": self.partition_key,
+                "after": self.after,
+                "poke_interval": self.poke_interval,
+            },
+        )
+
+    async def run(self) -> AsyncIterator[TriggerEvent]:
+        """Poll until the requested asset partition event exists."""
+        from airflow.sdk.execution_time.comms import AssetEventsResult, ErrorResponse, GetAssetEventByAsset
+        from airflow.sdk.execution_time.task_runner import SUPERVISOR_COMMS
+
+        after = timezone.parse(self.after) if isinstance(self.after, str) else self.after
+        while True:
+            response = await SUPERVISOR_COMMS.asend(
+                GetAssetEventByAsset(
+                    name=self.asset_name,
+                    uri=self.asset_uri,
+                    partition_key=self.partition_key,
+                    after=after,
+                    ascending=False,
+                    limit=1,
+                )
+            )
+            if isinstance(response, ErrorResponse):
+                yield TriggerEvent(
+                    {
+                        "status": "error",
+                        "message": f"{response.error.value}: {response.detail}",
+                    }
+                )
+                return
+            if TYPE_CHECKING:
+                assert isinstance(response, AssetEventsResult)
+            if response and response.asset_events:
+                yield TriggerEvent({"status": "success"})
+                return
+            await asyncio.sleep(self.poke_interval)

providers/standard/src/airflow/providers/standard/version_compat.py

@@ -37,6 +37,7 @@ def get_base_airflow_version_tuple() -> tuple[int, int, int]:
 AIRFLOW_V_3_1_3_PLUS: bool = get_base_airflow_version_tuple() >= (3, 1, 3)
 AIRFLOW_V_3_2_PLUS: bool = get_base_airflow_version_tuple() >= (3, 2, 0)
 AIRFLOW_V_3_3_PLUS: bool = get_base_airflow_version_tuple() >= (3, 3, 0)
+AIRFLOW_V_3_4_PLUS: bool = get_base_airflow_version_tuple() >= (3, 4, 0)
 
 # BaseOperator: Use 3.1+ due to xcom_push method missing in SDK BaseOperator 3.0.x
 # This is needed for DecoratedOperator compatibility
@@ -60,6 +61,7 @@ def is_arg_set(value):  # type: ignore[misc,no-redef]
     "AIRFLOW_V_3_1_PLUS",
     "AIRFLOW_V_3_2_PLUS",
     "AIRFLOW_V_3_3_PLUS",
+    "AIRFLOW_V_3_4_PLUS",
     "ArgNotSet",
     "BaseOperator",
     "is_arg_set",