From 6aae6784031fa014c0eb942e40d30c404d0c22a2 Mon Sep 17 00:00:00 2001 From: Henry Chen Date: Wed, 4 Nov 2026 01:05:13 +0800 Subject: [PATCH] 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. --- .../openapi/v2-rest-api-generated.yaml | 54 +++ .../core_api/routes/public/assets.py | 203 ++++++-- .../airflow/ui/openapi-gen/queries/queries.ts | 15 +- .../ui/openapi-gen/requests/services.gen.ts | 12 +- .../ui/openapi-gen/requests/types.gen.ts | 12 + .../core_api/routes/public/test_assets.py | 456 +++++++++++++++++- 6 files changed, 707 insertions(+), 45 deletions(-) diff --git a/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml b/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml index e135844ebec..8d77a7b77f7 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml +++ b/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 diff --git a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/assets.py b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/assets.py index 003e880b81d..3c58dfbf812 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/assets.py +++ b/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 + ), ) diff --git a/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts b/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts index 6cf9354c512..734261d0663 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts @@ -3524,31 +3524,37 @@ export const useVariableServiceBulkVariables = (options?: Omit, "mutationFn">) => useMutation({ mutationFn: ({ assetId, before }) => AssetService.deleteAssetQueuedEvents({ assetId, before }) as unknown as Promise, ...options }); + partitionKey?: string; +}, TContext>({ mutationFn: ({ assetId, before, partitionKey }) => AssetService.deleteAssetQueuedEvents({ assetId, before, partitionKey }) as unknown as Promise, ...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 = (options?: Omit, "mutationFn">) => useMutation({ mutationFn: ({ before, dagId }) => AssetService.deleteDagAssetQueuedEvents({ before, dagId }) as unknown as Promise, ...options }); + partitionKey?: string; +}, TContext>({ mutationFn: ({ before, dagId, partitionKey }) => AssetService.deleteDagAssetQueuedEvents({ before, dagId, partitionKey }) as unknown as Promise, ...options }); /** * Delete Dag Asset Queued Event * Delete a queued asset event for a Dag. @@ -3556,6 +3562,7 @@ export const useAssetServiceDeleteDagAssetQueuedEvents = , "mutationFn">) => useMutation({ mutationFn: ({ assetId, before, dagId }) => AssetService.deleteDagAssetQueuedEvent({ assetId, before, dagId }) as unknown as Promise, ...options }); + partitionKey?: string; +}, TContext>({ mutationFn: ({ assetId, before, dagId, partitionKey }) => AssetService.deleteDagAssetQueuedEvent({ assetId, before, dagId, partitionKey }) as unknown as Promise, ...options }); /** * Delete Connection * Delete a connection entry. diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts b/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts index 3296ee3690a..c103a857259 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts +++ b/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', diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts b/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts index 601b906b99f..f1979800647 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts +++ b/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; diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_assets.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_assets.py index 655a30e7d96..aace92dc03e 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_assets.py +++ b/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."""