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."""