Allow deleting queued partitioned asset events via REST API (#73119)

* Allow deleting queued partitioned asset events through the REST API

The DELETE queuedEvents endpoints only cleared queued events of
non-partitioned assets, so a partitioned event that was wrong or no
longer relevant could not be removed and would still trigger the
downstream Dag run. They now accept an optional partition_key to delete
the queued events of that partition whose Dag run has not been created
yet.

* Match partition_key against the asset event partition key

The partition_key on the DELETE queued events endpoints now refers to
the partition key of the queued asset event, the same key used by
POST /assets/{asset_id}/materialize, instead of the partition key of
the Dag run the event would create. This lets a single wrong upstream
partition be cleared when a partition mapper rolls several of them up
into one Dag run.

* Test that deleting queued partitioned events respects Dag read access

* Mention the partition key when no queued partitioned event is found

* Keep events of deleted Dag runs when deleting queued partitioned events

A partitioned asset event row outlives its partition record once the Dag
run it created is deleted, since only the partition record cascades.
Those rows were matched as queued events, so the endpoint deleted them
and returned 204 although nothing was queued. Only events of a partition
still waiting for its Dag run are deleted now.

* Test that Dag scoped queued event deletes leave other Dags' partitions alone

* Scope the pending partition cleanup to the partitions the call clears

Deleting queued partitioned events swept every empty pending partition in the
table, so clearing one asset's key could drop a partition of another Dag that
the request never matched. The partitions to clear are now collected before
anything is deleted, and both deletes are restricted to them.

The asset route now deletes what the caller may edit rather than what it may
read, so its test has to tell the two apart.

* Lock the partitions a queued event delete is about to clear

Two calls clearing different assets of the same partition could each still see
the other's uncommitted events under READ COMMITTED, so neither would notice
the partition had been emptied and it would stay pending with nothing left to
wait for.

* Name the queued event 404 detail helper with an action verb

* Say which partition key a queued event delete matches

A partition mapper can map an asset event's partition key to a different
Dag run partition key, so a caller reading only the parameter name cannot
tell which of the two the filter applies to.

* Build the partitioned queued event filter with the shared helper

The non-partitioned and partitioned queue tables carry the same asset,
Dag and timestamp columns, so two copies of the same filter drift apart
as soon as one of those columns changes.

* Regenerate the API spec and UI client

Keeps the generated clients carrying the same parameter documentation the
route declares.

* Say that one asset event can feed several partitioned Dag runs

A fan-out mapper turns one upstream partition key into several downstream
keys, and a rollup Dag run is fed by more than one event, so describing
the filter against a single Dag run's key mis-states what it matches.

airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml

@@ -722,6 +722,24 @@ paths:
           - type: string
           - type: 'null'
           title: Before
+      - name: partition_key
+        in: query
+        required: false
+        schema:
+          anyOf:
+          - type: string
+          - type: 'null'
+          description: Delete queued events of partitioned assets emitted for this
+            partition key whose Dag run has not been created yet, instead of queued
+            events of non-partitioned assets. The key is matched against the partition
+            key the asset event was emitted for, not against the partition keys of
+            the Dag runs the event feeds; a partition mapper can make them differ.
+          title: Partition Key
+        description: Delete queued events of partitioned assets emitted for this partition
+          key whose Dag run has not been created yet, instead of queued events of
+          non-partitioned assets. The key is matched against the partition key the
+          asset event was emitted for, not against the partition keys of the Dag runs
+          the event feeds; a partition mapper can make them differ.
       responses:
         '204':
           description: Successful Response
@@ -870,6 +888,24 @@ paths:
           - type: string
           - type: 'null'
           title: Before
+      - name: partition_key
+        in: query
+        required: false
+        schema:
+          anyOf:
+          - type: string
+          - type: 'null'
+          description: Delete queued events of partitioned assets emitted for this
+            partition key whose Dag run has not been created yet, instead of queued
+            events of non-partitioned assets. The key is matched against the partition
+            key the asset event was emitted for, not against the partition keys of
+            the Dag runs the event feeds; a partition mapper can make them differ.
+          title: Partition Key
+        description: Delete queued events of partitioned assets emitted for this partition
+          key whose Dag run has not been created yet, instead of queued events of
+          non-partitioned assets. The key is matched against the partition key the
+          asset event was emitted for, not against the partition keys of the Dag runs
+          the event feeds; a partition mapper can make them differ.
       responses:
         '204':
           description: Successful Response
@@ -995,6 +1031,24 @@ paths:
           - type: string
           - type: 'null'
           title: Before
+      - name: partition_key
+        in: query
+        required: false
+        schema:
+          anyOf:
+          - type: string
+          - type: 'null'
+          description: Delete queued events of partitioned assets emitted for this
+            partition key whose Dag run has not been created yet, instead of queued
+            events of non-partitioned assets. The key is matched against the partition
+            key the asset event was emitted for, not against the partition keys of
+            the Dag runs the event feeds; a partition mapper can make them differ.
+          title: Partition Key
+        description: Delete queued events of partitioned assets emitted for this partition
+          key whose Dag run has not been created yet, instead of queued events of
+          non-partitioned assets. The key is matched against the partition key the
+          asset event was emitted for, not against the partition keys of the Dag runs
+          the event feeds; a partition mapper can make them differ.
       responses:
         '204':
           description: Successful Response

airflow-core/src/airflow/api_fastapi/core_api/routes/public/assets.py

@@ -21,8 +21,8 @@ from datetime import datetime
 from typing import TYPE_CHECKING, Annotated, cast
 from uuid import UUID
 
-from fastapi import Depends, HTTPException, status
-from sqlalchemy import and_, delete, func, select
+from fastapi import Depends, HTTPException, Query, status
+from sqlalchemy import and_, delete, exists, func, select
 from sqlalchemy.engine import CursorResult
 from sqlalchemy.orm import joinedload, subqueryload
 
@@ -90,41 +90,143 @@ from airflow.models.asset import (
     AssetDagRunQueue,
     AssetEvent,
     AssetModel,
+    AssetPartitionDagRun,
     AssetWatcherModel,
+    PartitionedAssetKeyLog,
     TaskOutletAssetReference,
 )
 from airflow.models.dag import DagModel
 from airflow.models.dag_version import DagVersion
+from airflow.utils.sqlalchemy import with_row_locks
 from airflow.utils.state import DagRunState
 from airflow.utils.types import DagRunTriggeredByType, DagRunType
 
 if TYPE_CHECKING:
     from sqlalchemy.engine import Result
-    from sqlalchemy.sql import Select
+    from sqlalchemy.orm import Session
+    from sqlalchemy.sql import ColumnElement, Select
 
 assets_router = AirflowRouter(tags=["Asset"])
 
 
 def _generate_queued_event_where_clause(
     *,
+    model: type[AssetDagRunQueue] | type[PartitionedAssetKeyLog] = AssetDagRunQueue,
     asset_id: int | None = None,
     dag_id: str | None = None,
     before: datetime | str | None = None,
     permitted_dag_ids: set[str] | None = None,
-) -> list:
-    """Get AssetDagRunQueue where clause."""
-    where_clause = []
+) -> list[ColumnElement[bool]]:
+    """Get the asset, Dag and time filters shared by the queued event tables."""
+    where_clause: list[ColumnElement[bool]] = []
     if dag_id is not None:
-        where_clause.append(AssetDagRunQueue.target_dag_id == dag_id)
+        where_clause.append(model.target_dag_id == dag_id)
     if asset_id is not None:
-        where_clause.append(AssetDagRunQueue.asset_id == asset_id)
+        where_clause.append(model.asset_id == asset_id)
     if before is not None:
-        where_clause.append(AssetDagRunQueue.created_at < before)
+        where_clause.append(model.created_at < before)
     if permitted_dag_ids is not None:
-        where_clause.append(AssetDagRunQueue.target_dag_id.in_(permitted_dag_ids))
+        where_clause.append(model.target_dag_id.in_(permitted_dag_ids))
     return where_clause
 
 
+QueuedEventPartitionKeyQuery = Annotated[
+    str | None,
+    Query(
+        description=(
+            "Delete queued events of partitioned assets emitted for this partition key whose "
+            "Dag run has not been created yet, instead of queued events of non-partitioned assets. "
+            "The key is matched against the partition key the asset event was emitted for, not "
+            "against the partition keys of the Dag runs the event feeds; a partition mapper can "
+            "make them differ."
+        ),
+    ),
+]
+
+
+def _delete_pending_partitioned_queued_events(
+    *,
+    partition_key: str,
+    session: Session,
+    asset_id: int | None = None,
+    dag_id: str | None = None,
+    before: datetime | str | None = None,
+    permitted_dag_ids: set[str] | None = None,
+) -> int:
+    """
+    Delete queued partitioned asset events whose Dag run has not been created yet.
+
+    ``partition_key`` is matched against the partition key of the asset event, as in
+    ``POST /assets/{asset_id}/materialize``, not against the partition keys of the Dag
+    runs the event feeds; a partition mapper can make them differ.
+
+    Only events of an existing ``AssetPartitionDagRun`` whose Dag run has not been created
+    yet are deleted; ``PartitionedAssetKeyLog`` rows left behind once their Dag run was
+    deleted are history, not queued events. Pending ``AssetPartitionDagRun`` rows left
+    without any contributing events are deleted afterwards, so clearing some of the events
+    of a partition keeps the partition waiting for them rather than dropping the progress
+    made by the others.
+
+    Both deletes are restricted to the partitions this call actually clears events from. A
+    pending partition that is already empty for an unrelated reason belongs to whoever left
+    it behind, and a call that never matched any of its events must not remove it.
+
+    :return: The number of deleted queued events.
+    """
+    pakl_where_clause = _generate_queued_event_where_clause(
+        model=PartitionedAssetKeyLog,
+        asset_id=asset_id,
+        dag_id=dag_id,
+        before=before,
+        permitted_dag_ids=permitted_dag_ids,
+    )
+    pakl_where_clause.append(PartitionedAssetKeyLog.source_partition_key == partition_key)
+    pakl_of_apdr = PartitionedAssetKeyLog.asset_partition_dag_run_id == AssetPartitionDagRun.id
+
+    # Locking the partitions serialises calls that clear different assets of the same one.
+    # Under READ COMMITTED each would otherwise still see the events the other has deleted
+    # but not committed, conclude the partition is not empty, and leave it pending forever.
+    apdr_ids = session.scalars(
+        with_row_locks(
+            select(AssetPartitionDagRun.id).where(
+                AssetPartitionDagRun.created_dag_run_id.is_(None),
+                exists().where(pakl_of_apdr, *pakl_where_clause),
+            ),
+            session=session,
+            key_share=False,
+        )
+    ).all()
+    if not apdr_ids:
+        return 0
+
+    # Nothing is loaded into the session beforehand, so skip synchronizing it; the subquery
+    # criteria would otherwise make SQLAlchemy read back the deleted keys.
+    pakl_result = cast(
+        "CursorResult",
+        session.execute(
+            delete(PartitionedAssetKeyLog)
+            .where(
+                *pakl_where_clause,
+                PartitionedAssetKeyLog.asset_partition_dag_run_id.in_(apdr_ids),
+            )
+            .execution_options(synchronize_session=False)
+        ),
+    )
+
+    session.execute(
+        delete(AssetPartitionDagRun)
+        .where(AssetPartitionDagRun.id.in_(apdr_ids), ~exists().where(pakl_of_apdr))
+        .execution_options(synchronize_session=False)
+    )
+    return pakl_result.rowcount
+
+
+def _build_queued_event_not_found_detail(subject: str, partition_key: str | None) -> str:
+    if partition_key is not None:
+        subject = f"{subject} and partition_key: `{partition_key}`"
+    return f"{subject} was not found"
+
+
 class OnlyActiveFilter(BaseParam[bool]):
     """Filter on asset activeness."""
 
@@ -705,19 +807,31 @@ def delete_asset_queued_events(
     editable_dags_filter: EditableDagsFilterDep,
     session: SessionDep,
     before: OptionalDateTimeQuery = None,
+    partition_key: QueuedEventPartitionKeyQuery = None,
 ):
     """Delete queued asset events for an asset."""
     # This route is not tied to a single Dag, so the Dag axis is enforced per row: only the
     # queued events of Dags the user may edit are deleted, and the others are left in place.
-    where_clause = _generate_queued_event_where_clause(
-        asset_id=asset_id, before=before, permitted_dag_ids=editable_dags_filter.value
-    )
-    delete_stmt = delete(AssetDagRunQueue).where(*where_clause)
-    result = cast("CursorResult", session.execute(delete_stmt))
-    if result.rowcount == 0:
+    if partition_key is not None:
+        deleted = _delete_pending_partitioned_queued_events(
+            partition_key=partition_key,
+            asset_id=asset_id,
+            before=before,
+            permitted_dag_ids=editable_dags_filter.value,
+            session=session,
+        )
+    else:
+        where_clause = _generate_queued_event_where_clause(
+            asset_id=asset_id, before=before, permitted_dag_ids=editable_dags_filter.value
+        )
+        delete_stmt = delete(AssetDagRunQueue).where(*where_clause)
+        deleted = cast("CursorResult", session.execute(delete_stmt)).rowcount
+    if deleted == 0:
         raise HTTPException(
             status.HTTP_404_NOT_FOUND,
-            detail=f"Queue event with asset_id: `{asset_id}` was not found",
+            detail=_build_queued_event_not_found_detail(
+                f"Queue event with asset_id: `{asset_id}`", partition_key
+            ),
         )
 
 
@@ -741,16 +855,28 @@ def delete_dag_asset_queued_events(
     readable_dags_filter: ReadableDagsFilterDep,
     session: SessionDep,
     before: OptionalDateTimeQuery = None,
+    partition_key: QueuedEventPartitionKeyQuery = None,
 ):
-    where_clause = _generate_queued_event_where_clause(
-        dag_id=dag_id, before=before, permitted_dag_ids=readable_dags_filter.value
-    )
+    if partition_key is not None:
+        deleted = _delete_pending_partitioned_queued_events(
+            partition_key=partition_key,
+            dag_id=dag_id,
+            before=before,
+            permitted_dag_ids=readable_dags_filter.value,
+            session=session,
+        )
+    else:
+        where_clause = _generate_queued_event_where_clause(
+            dag_id=dag_id, before=before, permitted_dag_ids=readable_dags_filter.value
+        )
+        delete_statement = delete(AssetDagRunQueue).where(*where_clause)
+        deleted = cast("CursorResult", session.execute(delete_statement)).rowcount
 
-    delete_statement = delete(AssetDagRunQueue).where(*where_clause)
-    result = cast("CursorResult", session.execute(delete_statement))
-
-    if result.rowcount == 0:
-        raise HTTPException(status.HTTP_404_NOT_FOUND, f"Queue event with dag_id: `{dag_id}` was not found")
+    if deleted == 0:
+        raise HTTPException(
+            status.HTTP_404_NOT_FOUND,
+            _build_queued_event_not_found_detail(f"Queue event with dag_id: `{dag_id}`", partition_key),
+        )
 
 
 @assets_router.delete(
@@ -774,15 +900,28 @@ def delete_dag_asset_queued_event(
     readable_dags_filter: ReadableDagsFilterDep,
     session: SessionDep,
     before: OptionalDateTimeQuery = None,
+    partition_key: QueuedEventPartitionKeyQuery = None,
 ):
     """Delete a queued asset event for a Dag."""
-    where_clause = _generate_queued_event_where_clause(
-        dag_id=dag_id, before=before, asset_id=asset_id, permitted_dag_ids=readable_dags_filter.value
-    )
-    delete_statement = delete(AssetDagRunQueue).where(*where_clause)
-    result = cast("CursorResult", session.execute(delete_statement))
-    if result.rowcount == 0:
+    if partition_key is not None:
+        deleted = _delete_pending_partitioned_queued_events(
+            partition_key=partition_key,
+            dag_id=dag_id,
+            asset_id=asset_id,
+            before=before,
+            permitted_dag_ids=readable_dags_filter.value,
+            session=session,
+        )
+    else:
+        where_clause = _generate_queued_event_where_clause(
+            dag_id=dag_id, before=before, asset_id=asset_id, permitted_dag_ids=readable_dags_filter.value
+        )
+        delete_statement = delete(AssetDagRunQueue).where(*where_clause)
+        deleted = cast("CursorResult", session.execute(delete_statement)).rowcount
+    if deleted == 0:
         raise HTTPException(
             status.HTTP_404_NOT_FOUND,
-            detail=f"Queued event with dag_id: `{dag_id}` and asset_id: `{asset_id}` was not found",
+            detail=_build_queued_event_not_found_detail(
+                f"Queued event with dag_id: `{dag_id}` and asset_id: `{asset_id}`", partition_key
+            ),
         )

airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts

@@ -3524,31 +3524,37 @@ export const useVariableServiceBulkVariables = <TData = Common.VariableServiceBu
 * @param data The data for the request.
 * @param data.assetId
 * @param data.before
+* @param data.partitionKey Delete queued events of partitioned assets emitted for this partition key whose Dag run has not been created yet, instead of queued events of non-partitioned assets. The key is matched against the partition key the asset event was emitted for, not against the partition keys of the Dag runs the event feeds; a partition mapper can make them differ.
 * @returns void Successful Response
 * @throws ApiError
 */
 export const useAssetServiceDeleteAssetQueuedEvents = <TData = Common.AssetServiceDeleteAssetQueuedEventsMutationResult, TError = unknown, TContext = unknown>(options?: Omit<UseMutationOptions<TData, TError, {
   assetId: number;
   before?: string;
+  partitionKey?: string;
 }, TContext>, "mutationFn">) => useMutation<TData, TError, {
   assetId: number;
   before?: string;
-}, TContext>({ mutationFn: ({ assetId, before }) => AssetService.deleteAssetQueuedEvents({ assetId, before }) as unknown as Promise<TData>, ...options });
+  partitionKey?: string;
+}, TContext>({ mutationFn: ({ assetId, before, partitionKey }) => AssetService.deleteAssetQueuedEvents({ assetId, before, partitionKey }) as unknown as Promise<TData>, ...options });
 /**
 * Delete Dag Asset Queued Events
 * @param data The data for the request.
 * @param data.dagId
 * @param data.before
+* @param data.partitionKey Delete queued events of partitioned assets emitted for this partition key whose Dag run has not been created yet, instead of queued events of non-partitioned assets. The key is matched against the partition key the asset event was emitted for, not against the partition keys of the Dag runs the event feeds; a partition mapper can make them differ.
 * @returns void Successful Response
 * @throws ApiError
 */
 export const useAssetServiceDeleteDagAssetQueuedEvents = <TData = Common.AssetServiceDeleteDagAssetQueuedEventsMutationResult, TError = unknown, TContext = unknown>(options?: Omit<UseMutationOptions<TData, TError, {
   before?: string;
   dagId: string;
+  partitionKey?: string;
 }, TContext>, "mutationFn">) => useMutation<TData, TError, {
   before?: string;
   dagId: string;
-}, TContext>({ mutationFn: ({ before, dagId }) => AssetService.deleteDagAssetQueuedEvents({ before, dagId }) as unknown as Promise<TData>, ...options });
+  partitionKey?: string;
+}, TContext>({ mutationFn: ({ before, dagId, partitionKey }) => AssetService.deleteDagAssetQueuedEvents({ before, dagId, partitionKey }) as unknown as Promise<TData>, ...options });
 /**
 * Delete Dag Asset Queued Event
 * Delete a queued asset event for a Dag.
@@ -3556,6 +3562,7 @@ export const useAssetServiceDeleteDagAssetQueuedEvents = <TData = Common.AssetSe
 * @param data.dagId
 * @param data.assetId
 * @param data.before
+* @param data.partitionKey Delete queued events of partitioned assets emitted for this partition key whose Dag run has not been created yet, instead of queued events of non-partitioned assets. The key is matched against the partition key the asset event was emitted for, not against the partition keys of the Dag runs the event feeds; a partition mapper can make them differ.
 * @returns void Successful Response
 * @throws ApiError
 */
@@ -3563,11 +3570,13 @@ export const useAssetServiceDeleteDagAssetQueuedEvent = <TData = Common.AssetSer
   assetId: number;
   before?: string;
   dagId: string;
+  partitionKey?: string;
 }, TContext>, "mutationFn">) => useMutation<TData, TError, {
   assetId: number;
   before?: string;
   dagId: string;
-}, TContext>({ mutationFn: ({ assetId, before, dagId }) => AssetService.deleteDagAssetQueuedEvent({ assetId, before, dagId }) as unknown as Promise<TData>, ...options });
+  partitionKey?: string;
+}, TContext>({ mutationFn: ({ assetId, before, dagId, partitionKey }) => AssetService.deleteDagAssetQueuedEvent({ assetId, before, dagId, partitionKey }) as unknown as Promise<TData>, ...options });
 /**
 * Delete Connection
 * Delete a connection entry.

airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts

@@ -247,6 +247,7 @@ export class AssetService {
      * @param data The data for the request.
      * @param data.assetId
      * @param data.before
+     * @param data.partitionKey Delete queued events of partitioned assets emitted for this partition key whose Dag run has not been created yet, instead of queued events of non-partitioned assets. The key is matched against the partition key the asset event was emitted for, not against the partition keys of the Dag runs the event feeds; a partition mapper can make them differ.
      * @returns void Successful Response
      * @throws ApiError
      */
@@ -258,7 +259,8 @@ export class AssetService {
                 asset_id: data.assetId
             },
             query: {
-                before: data.before
+                before: data.before,
+                partition_key: data.partitionKey
             },
             errors: {
                 401: 'Unauthorized',
@@ -325,6 +327,7 @@ export class AssetService {
      * @param data The data for the request.
      * @param data.dagId
      * @param data.before
+     * @param data.partitionKey Delete queued events of partitioned assets emitted for this partition key whose Dag run has not been created yet, instead of queued events of non-partitioned assets. The key is matched against the partition key the asset event was emitted for, not against the partition keys of the Dag runs the event feeds; a partition mapper can make them differ.
      * @returns void Successful Response
      * @throws ApiError
      */
@@ -336,7 +339,8 @@ export class AssetService {
                 dag_id: data.dagId
             },
             query: {
-                before: data.before
+                before: data.before,
+                partition_key: data.partitionKey
             },
             errors: {
                 400: 'Bad Request',
@@ -385,6 +389,7 @@ export class AssetService {
      * @param data.dagId
      * @param data.assetId
      * @param data.before
+     * @param data.partitionKey Delete queued events of partitioned assets emitted for this partition key whose Dag run has not been created yet, instead of queued events of non-partitioned assets. The key is matched against the partition key the asset event was emitted for, not against the partition keys of the Dag runs the event feeds; a partition mapper can make them differ.
      * @returns void Successful Response
      * @throws ApiError
      */
@@ -397,7 +402,8 @@ export class AssetService {
                 asset_id: data.assetId
             },
             query: {
-                before: data.before
+                before: data.before,
+                partition_key: data.partitionKey
             },
             errors: {
                 400: 'Bad Request',

airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts

@@ -3808,6 +3808,10 @@ export type GetAssetQueuedEventsResponse = QueuedEventCollectionResponse;
 export type DeleteAssetQueuedEventsData = {
     assetId: number;
     before?: string | null;
+    /**
+     * Delete queued events of partitioned assets emitted for this partition key whose Dag run has not been created yet, instead of queued events of non-partitioned assets. The key is matched against the partition key the asset event was emitted for, not against the partition keys of the Dag runs the event feeds; a partition mapper can make them differ.
+     */
+    partitionKey?: string | null;
 };
 
 export type DeleteAssetQueuedEventsResponse = void;
@@ -3828,6 +3832,10 @@ export type GetDagAssetQueuedEventsResponse = QueuedEventCollectionResponse;
 export type DeleteDagAssetQueuedEventsData = {
     before?: string | null;
     dagId: string;
+    /**
+     * Delete queued events of partitioned assets emitted for this partition key whose Dag run has not been created yet, instead of queued events of non-partitioned assets. The key is matched against the partition key the asset event was emitted for, not against the partition keys of the Dag runs the event feeds; a partition mapper can make them differ.
+     */
+    partitionKey?: string | null;
 };
 
 export type DeleteDagAssetQueuedEventsResponse = void;
@@ -3844,6 +3852,10 @@ export type DeleteDagAssetQueuedEventData = {
     assetId: number;
     before?: string | null;
     dagId: string;
+    /**
+     * Delete queued events of partitioned assets emitted for this partition key whose Dag run has not been created yet, instead of queued events of non-partitioned assets. The key is matched against the partition key the asset event was emitted for, not against the partition keys of the Dag runs the event feeds; a partition mapper can make them differ.
+     */
+    partitionKey?: string | null;
 };
 
 export type DeleteDagAssetQueuedEventResponse = void;

airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_assets.py

@@ -40,8 +40,10 @@ from airflow.models.asset import (
     AssetDagRunQueue,
     AssetEvent,
     AssetModel,
+    AssetPartitionDagRun,
     AssetWatcherModel,
     DagScheduleAssetReference,
+    PartitionedAssetKeyLog,
     TaskInletAssetReference,
     TaskOutletAssetReference,
 )
@@ -2015,6 +2017,13 @@ class TestQueuedEventEndpoint(TestAssets):
         session.commit()
         return adrq
 
+    @staticmethod
+    def _authorized_dag_ids_by_method(*, readable: set[str], editable: set[str]):
+        def get_authorized_dag_ids(*, user, method="GET", **kwargs):
+            return set(editable) if method == "PUT" else set(readable)
+
+        return get_authorized_dag_ids
+
 
 class TestGetDagAssetQueuedEvents(TestQueuedEventEndpoint):
     @pytest.mark.usefixtures("time_freezer")
@@ -2200,6 +2209,446 @@ class TestDeleteDagDatasetQueuedEvents(TestQueuedEventEndpoint):
         assert response.json()["detail"] == "Queue event with dag_id: `dag` was not found"
 
 
+class TestDeletePartitionedQueuedEvents(TestQueuedEventEndpoint):
+    @pytest.fixture(autouse=True)
+    def clear_partitioned_queue(self, session):
+        session.execute(delete(PartitionedAssetKeyLog))
+        session.execute(delete(AssetPartitionDagRun))
+        session.commit()
+        yield
+        session.execute(delete(PartitionedAssetKeyLog))
+        session.execute(delete(AssetPartitionDagRun))
+        session.commit()
+
+    @staticmethod
+    def _queue_partition(
+        session,
+        *,
+        dag_id: str,
+        partition_key: str,
+        source_keys_by_asset_id: dict[int, list[str]],
+        created_dag_run_id: int | None = None,
+    ) -> AssetPartitionDagRun:
+        apdr = AssetPartitionDagRun(
+            target_dag_id=dag_id, partition_key=partition_key, created_dag_run_id=created_dag_run_id
+        )
+        session.add(apdr)
+        session.flush()
+        for asset_id, source_keys in source_keys_by_asset_id.items():
+            for source_key in source_keys:
+                event = AssetEvent(asset_id=asset_id, partition_key=source_key, timestamp=timezone.utcnow())
+                session.add(event)
+                session.flush()
+                session.add(
+                    PartitionedAssetKeyLog(
+                        asset_id=asset_id,
+                        asset_event_id=event.id,
+                        asset_partition_dag_run_id=apdr.id,
+                        source_partition_key=source_key,
+                        target_dag_id=dag_id,
+                        target_partition_key=partition_key,
+                    )
+                )
+        session.commit()
+        return apdr
+
+    @staticmethod
+    def _remaining(session) -> tuple[set[tuple[str, int | None]], set[tuple[int, int, str]]]:
+        session.expire_all()
+        apdrs = {
+            (apdr.partition_key, apdr.created_dag_run_id)
+            for apdr in session.scalars(select(AssetPartitionDagRun))
+        }
+        pakls = {
+            (pakl.asset_partition_dag_run_id, pakl.asset_id, pakl.source_partition_key)
+            for pakl in session.scalars(select(PartitionedAssetKeyLog))
+        }
+        return apdrs, pakls
+
+    def test_dag_endpoint_deletes_only_pending_partition_with_matching_key(
+        self, test_client, session, create_dummy_dag
+    ):
+        dag, _ = create_dummy_dag()
+        dag_run_id = session.scalar(select(DagRun.id).where(DagRun.dag_id == dag.dag_id))
+        asset_1, asset_2 = self.create_assets(session=session, num=2)
+        self._create_asset_dag_run_queues(dag.dag_id, asset_1.id, session)
+        self._queue_partition(
+            session,
+            dag_id=dag.dag_id,
+            partition_key="2026-09-02",
+            source_keys_by_asset_id={asset_1.id: ["2026-09-02"], asset_2.id: ["2026-09-02"]},
+        )
+        other = self._queue_partition(
+            session,
+            dag_id=dag.dag_id,
+            partition_key="2026-09-03",
+            source_keys_by_asset_id={asset_1.id: ["2026-09-03"]},
+        )
+        self._queue_partition(
+            session,
+            dag_id=dag.dag_id,
+            partition_key="2026-09-02",
+            source_keys_by_asset_id={},
+            created_dag_run_id=dag_run_id,
+        )
+
+        response = test_client.delete(
+            f"/dags/{dag.dag_id}/assets/queuedEvents", params={"partition_key": "2026-09-02"}
+        )
+
+        assert response.status_code == 204
+        apdrs, pakls = self._remaining(session)
+        assert apdrs == {("2026-09-03", None), ("2026-09-02", dag_run_id)}
+        assert pakls == {(other.id, asset_1.id, "2026-09-03")}
+        # Queued events of non-partitioned assets are left alone when a partition key is given.
+        assert len(session.scalars(select(AssetDagRunQueue)).all()) == 1
+        check_last_log(session, dag_id=dag.dag_id, event="delete_dag_asset_queued_events", logical_date=None)
+
+    @pytest.mark.parametrize(
+        "path_template",
+        [
+            pytest.param("/assets/{asset_id}/queuedEvents", id="asset"),
+            pytest.param("/dags/{dag_id}/assets/{asset_id}/queuedEvents", id="dag-asset"),
+        ],
+    )
+    def test_asset_endpoints_keep_partition_pending_until_no_asset_remains(
+        self, test_client, session, create_dummy_dag, path_template
+    ):
+        dag, _ = create_dummy_dag()
+        asset_1, asset_2 = self.create_assets(session=session, num=2)
+        apdr = self._queue_partition(
+            session,
+            dag_id=dag.dag_id,
+            partition_key="2026-09-02",
+            source_keys_by_asset_id={asset_1.id: ["2026-09-02"], asset_2.id: ["2026-09-02"]},
+        )
+
+        response = test_client.delete(
+            path_template.format(dag_id=dag.dag_id, asset_id=asset_1.id),
+            params={"partition_key": "2026-09-02"},
+        )
+
+        assert response.status_code == 204
+        apdrs, pakls = self._remaining(session)
+        assert apdrs == {("2026-09-02", None)}
+        assert pakls == {(apdr.id, asset_2.id, "2026-09-02")}
+
+        response = test_client.delete(
+            path_template.format(dag_id=dag.dag_id, asset_id=asset_2.id),
+            params={"partition_key": "2026-09-02"},
+        )
+
+        assert response.status_code == 204
+        assert self._remaining(session) == (set(), set())
+
+    @pytest.mark.parametrize(
+        "path_template",
+        [
+            pytest.param("/assets/{asset_id}/queuedEvents", id="asset"),
+            pytest.param("/dags/{dag_id}/assets/queuedEvents", id="dag"),
+            pytest.param("/dags/{dag_id}/assets/{asset_id}/queuedEvents", id="dag-asset"),
+        ],
+    )
+    def test_partition_key_matches_asset_event_partition_key_not_dag_run_partition_key(
+        self, test_client, session, create_dummy_dag, path_template
+    ):
+        dag, _ = create_dummy_dag()
+        (asset,) = self.create_assets(session=session, num=1)
+        apdr = self._queue_partition(
+            session,
+            dag_id=dag.dag_id,
+            partition_key="2026-09-02",
+            source_keys_by_asset_id={asset.id: ["2026-09-02T00", "2026-09-02T01"]},
+        )
+        path = path_template.format(dag_id=dag.dag_id, asset_id=asset.id)
+
+        response = test_client.delete(path, params={"partition_key": "2026-09-02"})
+        assert response.status_code == 404
+        assert len(self._remaining(session)[1]) == 2
+
+        response = test_client.delete(path, params={"partition_key": "2026-09-02T00"})
+        assert response.status_code == 204
+        assert self._remaining(session) == (
+            {("2026-09-02", None)},
+            {(apdr.id, asset.id, "2026-09-02T01")},
+        )
+
+        response = test_client.delete(path, params={"partition_key": "2026-09-02T01"})
+        assert response.status_code == 204
+        assert self._remaining(session) == (set(), set())
+
+    @pytest.mark.parametrize(
+        "path_template",
+        [
+            pytest.param("/dags/{dag_id}/assets/queuedEvents", id="dag"),
+            pytest.param("/dags/{dag_id}/assets/{asset_id}/queuedEvents", id="dag-asset"),
+        ],
+    )
+    def test_dag_endpoints_only_delete_events_queued_for_that_dag(
+        self, test_client, session, create_dummy_dag, path_template
+    ):
+        create_dummy_dag()
+        (asset,) = self.create_assets(session=session, num=1)
+        self._queue_partition(
+            session,
+            dag_id="dag",
+            partition_key="2026-09-02",
+            source_keys_by_asset_id={asset.id: ["2026-09-02"]},
+        )
+        other = self._queue_partition(
+            session,
+            dag_id="other_dag",
+            partition_key="2026-09-02",
+            source_keys_by_asset_id={asset.id: ["2026-09-02"]},
+        )
+        self._queue_partition(
+            session, dag_id="other_dag", partition_key="2026-09-03", source_keys_by_asset_id={}
+        )
+
+        with mock.patch.object(
+            BaseAuthManager, "get_authorized_dag_ids", autospec=True, return_value={"dag", "other_dag"}
+        ):
+            response = test_client.delete(
+                path_template.format(dag_id="dag", asset_id=asset.id),
+                params={"partition_key": "2026-09-02"},
+            )
+
+        assert response.status_code == 204
+        session.expire_all()
+        assert {
+            (apdr.target_dag_id, apdr.partition_key) for apdr in session.scalars(select(AssetPartitionDagRun))
+        } == {("other_dag", "2026-09-02"), ("other_dag", "2026-09-03")}
+        assert set(session.scalars(select(PartitionedAssetKeyLog.asset_partition_dag_run_id))) == {other.id}
+
+    def test_keeps_empty_pending_partitions_the_call_did_not_clear(
+        self, test_client, session, create_dummy_dag
+    ):
+        create_dummy_dag()
+        (asset,) = self.create_assets(session=session, num=1)
+        self._queue_partition(
+            session,
+            dag_id="dag",
+            partition_key="2026-09-02",
+            source_keys_by_asset_id={asset.id: ["2026-09-02"]},
+        )
+        # Pending and already empty before the call, for another Dag and another key.
+        self._queue_partition(
+            session, dag_id="other_dag", partition_key="2026-09-03", source_keys_by_asset_id={}
+        )
+
+        with mock.patch.object(
+            BaseAuthManager, "get_authorized_dag_ids", autospec=True, return_value={"dag", "other_dag"}
+        ):
+            response = test_client.delete(
+                f"/assets/{asset.id}/queuedEvents", params={"partition_key": "2026-09-02"}
+            )
+
+        assert response.status_code == 204
+        session.expire_all()
+        assert {
+            (apdr.target_dag_id, apdr.partition_key) for apdr in session.scalars(select(AssetPartitionDagRun))
+        } == {("other_dag", "2026-09-03")}
+
+    @pytest.mark.parametrize(
+        ("editable_dags", "expected_status", "expected_remaining_dag_ids"),
+        [
+            pytest.param({"dag", "hidden_dag"}, 204, set(), id="both-dags-editable"),
+            pytest.param({"dag"}, 204, {"hidden_dag"}, id="one-dag-editable"),
+            pytest.param(set(), 404, {"dag", "hidden_dag"}, id="no-dag-editable"),
+        ],
+    )
+    def test_only_deletes_events_queued_for_editable_dags(
+        self,
+        test_client,
+        session,
+        create_dummy_dag,
+        editable_dags,
+        expected_status,
+        expected_remaining_dag_ids,
+    ):
+        """Being able to read a Dag is not enough: this route deletes only what the user may edit."""
+        create_dummy_dag()
+        (asset,) = self.create_assets(session=session, num=1)
+        for dag_id in ("dag", "hidden_dag"):
+            self._queue_partition(
+                session,
+                dag_id=dag_id,
+                partition_key="2026-09-02",
+                source_keys_by_asset_id={asset.id: ["2026-09-02"]},
+            )
+
+        with mock.patch.object(
+            BaseAuthManager,
+            "get_authorized_dag_ids",
+            side_effect=self._authorized_dag_ids_by_method(
+                readable={"dag", "hidden_dag"}, editable=editable_dags
+            ),
+        ):
+            response = test_client.delete(
+                f"/assets/{asset.id}/queuedEvents", params={"partition_key": "2026-09-02"}
+            )
+
+        assert response.status_code == expected_status
+        session.expire_all()
+        assert set(session.scalars(select(AssetPartitionDagRun.target_dag_id))) == expected_remaining_dag_ids
+        assert (
+            set(session.scalars(select(PartitionedAssetKeyLog.target_dag_id))) == expected_remaining_dag_ids
+        )
+
+    def test_before_only_deletes_older_events(self, test_client, session, create_dummy_dag):
+        dag, _ = create_dummy_dag()
+        (asset,) = self.create_assets(session=session, num=1)
+        with time_machine.travel(DEFAULT_DATE, tick=False):
+            apdr = self._queue_partition(
+                session,
+                dag_id=dag.dag_id,
+                partition_key="us",
+                source_keys_by_asset_id={asset.id: ["us"]},
+            )
+        with time_machine.travel(DEFAULT_DATE + timedelta(hours=2), tick=False):
+            newer = AssetEvent(asset_id=asset.id, partition_key="us", timestamp=timezone.utcnow())
+            session.add(newer)
+            session.flush()
+            session.add(
+                PartitionedAssetKeyLog(
+                    asset_id=asset.id,
+                    asset_event_id=newer.id,
+                    asset_partition_dag_run_id=apdr.id,
+                    source_partition_key="us",
+                    target_dag_id=dag.dag_id,
+                    target_partition_key="us",
+                )
+            )
+            session.commit()
+
+        response = test_client.delete(
+            f"/assets/{asset.id}/queuedEvents",
+            params={
+                "partition_key": "us",
+                "before": from_datetime_to_zulu_without_ms(DEFAULT_DATE + timedelta(hours=1)),
+            },
+        )
+
+        assert response.status_code == 204
+        apdrs, _ = self._remaining(session)
+        assert apdrs == {("us", None)}
+        assert session.scalars(select(PartitionedAssetKeyLog.asset_event_id)).all() == [newer.id]
+
+    @pytest.mark.parametrize(
+        ("path_template", "expected_detail"),
+        [
+            pytest.param(
+                "/assets/{asset_id}/queuedEvents",
+                "Queue event with asset_id: `{asset_id}` and partition_key: `2026-09-02` was not found",
+                id="asset",
+            ),
+            pytest.param(
+                "/dags/{dag_id}/assets/queuedEvents",
+                "Queue event with dag_id: `{dag_id}` and partition_key: `2026-09-02` was not found",
+                id="dag",
+            ),
+            pytest.param(
+                "/dags/{dag_id}/assets/{asset_id}/queuedEvents",
+                "Queued event with dag_id: `{dag_id}` and asset_id: `{asset_id}` and partition_key: `2026-09-02` was not found",
+                id="dag-asset",
+            ),
+        ],
+    )
+    def test_should_respond_404_when_partition_already_created_dag_run(
+        self, test_client, session, create_dummy_dag, path_template, expected_detail
+    ):
+        dag, _ = create_dummy_dag()
+        dag_run_id = session.scalar(select(DagRun.id).where(DagRun.dag_id == dag.dag_id))
+        (asset,) = self.create_assets(session=session, num=1)
+        self._create_asset_dag_run_queues(dag.dag_id, asset.id, session)
+        self._queue_partition(
+            session,
+            dag_id=dag.dag_id,
+            partition_key="2026-09-02",
+            source_keys_by_asset_id={asset.id: ["2026-09-02"]},
+            created_dag_run_id=dag_run_id,
+        )
+
+        response = test_client.delete(
+            path_template.format(dag_id=dag.dag_id, asset_id=asset.id),
+            params={"partition_key": "2026-09-02"},
+        )
+
+        assert response.status_code == 404
+        assert response.json()["detail"] == expected_detail.format(dag_id=dag.dag_id, asset_id=asset.id)
+        assert len(self._remaining(session)[1]) == 1
+        assert len(session.scalars(select(AssetDagRunQueue)).all()) == 1
+
+    def test_should_respond_404_for_events_of_deleted_dag_run(self, test_client, session, create_dummy_dag):
+        dag, _ = create_dummy_dag()
+        dag_run = session.scalar(select(DagRun).where(DagRun.dag_id == dag.dag_id))
+        (asset,) = self.create_assets(session=session, num=1)
+        self._queue_partition(
+            session,
+            dag_id=dag.dag_id,
+            partition_key="2026-09-02",
+            source_keys_by_asset_id={asset.id: ["2026-09-02"]},
+            created_dag_run_id=dag_run.id,
+        )
+        session.delete(dag_run)
+        session.commit()
+        apdrs, pakls = self._remaining(session)
+        assert apdrs == set()
+        assert len(pakls) == 1
+
+        response = test_client.delete(
+            f"/assets/{asset.id}/queuedEvents", params={"partition_key": "2026-09-02"}
+        )
+
+        assert response.status_code == 404
+        assert len(self._remaining(session)[1]) == 1
+
+    def test_delete_does_not_read_back_deleted_row_keys(self, test_client, session, create_dummy_dag):
+        from sqlalchemy import event
+
+        import airflow.settings
+
+        dag, _ = create_dummy_dag()
+        (asset,) = self.create_assets(session=session, num=1)
+        self._queue_partition(
+            session,
+            dag_id=dag.dag_id,
+            partition_key="2026-09-02",
+            source_keys_by_asset_id={asset.id: ["2026-09-02"]},
+        )
+
+        executed_statements: list[str] = []
+
+        def capture(_conn, _cursor, statement, _parameters, _context, _executemany):
+            executed_statements.append(" ".join(statement.split()).upper())
+
+        event.listen(airflow.settings.engine, "before_cursor_execute", capture)
+        try:
+            response = test_client.delete(
+                f"/dags/{dag.dag_id}/assets/{asset.id}/queuedEvents",
+                params={"partition_key": "2026-09-02"},
+            )
+        finally:
+            event.remove(airflow.settings.engine, "before_cursor_execute", capture)
+
+        assert response.status_code == 204
+        deletes = [s for s in executed_statements if s.startswith("DELETE")]
+        assert len(deletes) == 2
+        assert [s for s in deletes if "RETURNING" in s] == [], "DELETE must not read back deleted keys"
+        partition_selects = [
+            s for s in executed_statements if s.startswith("SELECT") and "PARTITIONED_ASSET_KEY_LOG" in s
+        ]
+        # Collecting the partitions to clear is the one read of the queued events, and it
+        # projects their ``AssetPartitionDagRun`` id rather than the keys being deleted.
+        assert len(partition_selects) == 1
+        assert partition_selects[0].startswith("SELECT ASSET_PARTITION_DAG_RUN.ID")
+        assert executed_statements.index(partition_selects[0]) < executed_statements.index(deletes[0])
+        after_first_delete = executed_statements[executed_statements.index(deletes[0]) :]
+        assert [s for s in after_first_delete if s.startswith("SELECT")] == [], (
+            "No SELECT may follow a DELETE to read back the keys it removed"
+        )
+
+
 class TestPostAssetEvents(TestAssets):
     @pytest.mark.usefixtures("time_freezer")
     def test_should_respond_200(self, test_client, session):
@@ -2975,13 +3424,6 @@ class TestDeleteAssetQueuedEvents(TestQueuedEventEndpoint):
         session.commit()
         return asset
 
-    @staticmethod
-    def _authorized_dag_ids_by_method(*, readable: set[str], editable: set[str]):
-        def get_authorized_dag_ids(*, user, method="GET", **kwargs):
-            return set(editable) if method == "PUT" else set(readable)
-
-        return get_authorized_dag_ids
-
     @pytest.mark.parametrize("query_string", ["", "?dag_id=dag_a"], ids=["no-dag-id", "editable-dag-id"])
     def test_deletes_only_queued_events_of_editable_dags(self, test_client, session, query_string):
         """A user who may edit dag_a but only read dag_b must leave dag_b's queued events in place."""