Allow deleting queued partitioned asset events via REST API

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

@@ -722,6 +722,20 @@ 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.
+          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.
       responses:
         '204':
           description: Successful Response
@@ -870,6 +884,20 @@ 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.
+          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.
       responses:
         '204':
           description: Successful Response
@@ -995,6 +1023,20 @@ 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.
+          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.
       responses:
         '204':
           description: Successful Response

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

@@ -21,8 +21,8 @@
 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,17 +90,21 @@
     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"])
 
@@ -125,6 +129,103 @@ def _generate_queued_event_where_clause(
     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."
+        ),
+    ),
+]
+
+
+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 key of the Dag run
+    the event would create; a partition mapper can make the two 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: list[ColumnElement[bool]] = [
+        PartitionedAssetKeyLog.source_partition_key == partition_key,
+    ]
+    if asset_id is not None:
+        pakl_where_clause.append(PartitionedAssetKeyLog.asset_id == asset_id)
+    if dag_id is not None:
+        pakl_where_clause.append(PartitionedAssetKeyLog.target_dag_id == dag_id)
+    if before is not None:
+        pakl_where_clause.append(PartitionedAssetKeyLog.created_at < before)
+    if permitted_dag_ids is not None:
+        pakl_where_clause.append(PartitionedAssetKeyLog.target_dag_id.in_(permitted_dag_ids))
+    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."""
 
@@ -703,19 +804,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
+            ),
         )
 
 
@@ -739,16 +852,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
-    )
-
-    delete_statement = delete(AssetDagRunQueue).where(*where_clause)
-    result = cast("CursorResult", session.execute(delete_statement))
+    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
 
-    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(
@@ -772,15 +897,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

@@ -3495,50 +3495,59 @@ 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.
 * @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.
 * @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.
 * @param data The data for the request.
 * @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.
 * @returns void Successful Response
 * @throws ApiError
 */
 export const useAssetServiceDeleteDagAssetQueuedEvent = <TData = Common.AssetServiceDeleteDagAssetQueuedEventMutationResult, TError = unknown, TContext = unknown>(options?: Omit<UseMutationOptions<TData, TError, {
   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.