diff --git a/airflow/api_connexion/endpoints/dag_run_endpoint.py b/airflow/api_connexion/endpoints/dag_run_endpoint.py index 6a38eb27ff45..a97f7d0b20ef 100644 --- a/airflow/api_connexion/endpoints/dag_run_endpoint.py +++ b/airflow/api_connexion/endpoints/dag_run_endpoint.py @@ -115,6 +115,7 @@ def get_dag_run( raise BadRequest("DAGRunSchema error", detail=str(e)) +@mark_fastapi_migration_done @security.requires_access_dag("GET", DagAccessEntity.RUN) @security.requires_access_asset("GET") @provide_session diff --git a/airflow/api_fastapi/core_api/datamodels/assets.py b/airflow/api_fastapi/core_api/datamodels/assets.py index 85e41ff7b569..340daad3c16a 100644 --- a/airflow/api_fastapi/core_api/datamodels/assets.py +++ b/airflow/api_fastapi/core_api/datamodels/assets.py @@ -73,7 +73,7 @@ class DagRunAssetReference(BaseModel): dag_id: str execution_date: datetime = Field(alias="logical_date") start_date: datetime - end_date: datetime + end_date: datetime | None state: str data_interval_start: datetime data_interval_end: datetime diff --git a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml index b99b389de51f..1b6cd24f1110 100644 --- a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml +++ b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml @@ -1164,6 +1164,58 @@ paths: application/json: schema: $ref: '#/components/schemas/HTTPValidationError' + /public/dags/{dag_id}/dagRuns/{dag_run_id}/upstreamAssetEvents: + get: + tags: + - DagRun + summary: Get Upstream Asset Events + description: If dag run is asset-triggered, return the asset events that triggered + it. + operationId: get_upstream_asset_events + parameters: + - name: dag_id + in: path + required: true + schema: + type: string + title: Dag Id + - name: dag_run_id + in: path + required: true + schema: + type: string + title: Dag Run Id + responses: + '200': + description: Successful Response + content: + application/json: + schema: + $ref: '#/components/schemas/AssetEventCollectionResponse' + '401': + content: + application/json: + schema: + $ref: '#/components/schemas/HTTPExceptionResponse' + description: Unauthorized + '403': + content: + application/json: + schema: + $ref: '#/components/schemas/HTTPExceptionResponse' + description: Forbidden + '404': + content: + application/json: + schema: + $ref: '#/components/schemas/HTTPExceptionResponse' + description: Not Found + '422': + description: Validation Error + content: + application/json: + schema: + $ref: '#/components/schemas/HTTPValidationError' /public/dagSources/{file_token}: get: tags: @@ -4817,8 +4869,10 @@ components: format: date-time title: Start Date end_date: - type: string - format: date-time + anyOf: + - type: string + format: date-time + - type: 'null' title: End Date state: type: string diff --git a/airflow/api_fastapi/core_api/routes/public/dag_run.py b/airflow/api_fastapi/core_api/routes/public/dag_run.py index 810896806eea..698e322a12ec 100644 --- a/airflow/api_fastapi/core_api/routes/public/dag_run.py +++ b/airflow/api_fastapi/core_api/routes/public/dag_run.py @@ -30,6 +30,7 @@ ) from airflow.api_fastapi.common.db.common import get_session from airflow.api_fastapi.common.router import AirflowRouter +from airflow.api_fastapi.core_api.datamodels.assets import AssetEventCollectionResponse, AssetEventResponse from airflow.api_fastapi.core_api.datamodels.dag_run import ( DAGRunPatchBody, DAGRunPatchStates, @@ -142,3 +143,35 @@ def patch_dag_run( dag_run = session.get(DagRun, dag_run.id) return DAGRunResponse.model_validate(dag_run, from_attributes=True) + + +@dag_run_router.get( + "/{dag_run_id}/upstreamAssetEvents", + responses=create_openapi_http_exception_doc( + [ + status.HTTP_404_NOT_FOUND, + ] + ), +) +def get_upstream_asset_events( + dag_id: str, dag_run_id: str, session: Annotated[Session, Depends(get_session)] +) -> AssetEventCollectionResponse: + """If dag run is asset-triggered, return the asset events that triggered it.""" + dag_run: DagRun | None = session.scalar( + select(DagRun).where( + DagRun.dag_id == dag_id, + DagRun.run_id == dag_run_id, + ) + ) + if dag_run is None: + raise HTTPException( + status.HTTP_404_NOT_FOUND, + f"The DagRun with dag_id: `{dag_id}` and run_id: `{dag_run_id}` was not found", + ) + events = dag_run.consumed_asset_events + return AssetEventCollectionResponse( + asset_events=[ + AssetEventResponse.model_validate(asset_event, from_attributes=True) for asset_event in events + ], + total_entries=len(events), + ) diff --git a/airflow/ui/openapi-gen/queries/common.ts b/airflow/ui/openapi-gen/queries/common.ts index bc3cabf37929..8a5d75cc132e 100644 --- a/airflow/ui/openapi-gen/queries/common.ts +++ b/airflow/ui/openapi-gen/queries/common.ts @@ -305,6 +305,28 @@ export const UseDagRunServiceGetDagRunKeyFn = ( }, queryKey?: Array, ) => [useDagRunServiceGetDagRunKey, ...(queryKey ?? [{ dagId, dagRunId }])]; +export type DagRunServiceGetUpstreamAssetEventsDefaultResponse = Awaited< + ReturnType +>; +export type DagRunServiceGetUpstreamAssetEventsQueryResult< + TData = DagRunServiceGetUpstreamAssetEventsDefaultResponse, + TError = unknown, +> = UseQueryResult; +export const useDagRunServiceGetUpstreamAssetEventsKey = + "DagRunServiceGetUpstreamAssetEvents"; +export const UseDagRunServiceGetUpstreamAssetEventsKeyFn = ( + { + dagId, + dagRunId, + }: { + dagId: string; + dagRunId: string; + }, + queryKey?: Array, +) => [ + useDagRunServiceGetUpstreamAssetEventsKey, + ...(queryKey ?? [{ dagId, dagRunId }]), +]; export type DagSourceServiceGetDagSourceDefaultResponse = Awaited< ReturnType >; diff --git a/airflow/ui/openapi-gen/queries/prefetch.ts b/airflow/ui/openapi-gen/queries/prefetch.ts index f5f120e55554..83e4b2a55cad 100644 --- a/airflow/ui/openapi-gen/queries/prefetch.ts +++ b/airflow/ui/openapi-gen/queries/prefetch.ts @@ -386,6 +386,32 @@ export const prefetchUseDagRunServiceGetDagRun = ( queryKey: Common.UseDagRunServiceGetDagRunKeyFn({ dagId, dagRunId }), queryFn: () => DagRunService.getDagRun({ dagId, dagRunId }), }); +/** + * Get Upstream Asset Events + * If dag run is asset-triggered, return the asset events that triggered it. + * @param data The data for the request. + * @param data.dagId + * @param data.dagRunId + * @returns AssetEventCollectionResponse Successful Response + * @throws ApiError + */ +export const prefetchUseDagRunServiceGetUpstreamAssetEvents = ( + queryClient: QueryClient, + { + dagId, + dagRunId, + }: { + dagId: string; + dagRunId: string; + }, +) => + queryClient.prefetchQuery({ + queryKey: Common.UseDagRunServiceGetUpstreamAssetEventsKeyFn({ + dagId, + dagRunId, + }), + queryFn: () => DagRunService.getUpstreamAssetEvents({ dagId, dagRunId }), + }); /** * Get Dag Source * Get source code using file token. diff --git a/airflow/ui/openapi-gen/queries/queries.ts b/airflow/ui/openapi-gen/queries/queries.ts index f16ebde095fb..b4e0c6f7628b 100644 --- a/airflow/ui/openapi-gen/queries/queries.ts +++ b/airflow/ui/openapi-gen/queries/queries.ts @@ -484,6 +484,39 @@ export const useDagRunServiceGetDagRun = < queryFn: () => DagRunService.getDagRun({ dagId, dagRunId }) as TData, ...options, }); +/** + * Get Upstream Asset Events + * If dag run is asset-triggered, return the asset events that triggered it. + * @param data The data for the request. + * @param data.dagId + * @param data.dagRunId + * @returns AssetEventCollectionResponse Successful Response + * @throws ApiError + */ +export const useDagRunServiceGetUpstreamAssetEvents = < + TData = Common.DagRunServiceGetUpstreamAssetEventsDefaultResponse, + TError = unknown, + TQueryKey extends Array = unknown[], +>( + { + dagId, + dagRunId, + }: { + dagId: string; + dagRunId: string; + }, + queryKey?: TQueryKey, + options?: Omit, "queryKey" | "queryFn">, +) => + useQuery({ + queryKey: Common.UseDagRunServiceGetUpstreamAssetEventsKeyFn( + { dagId, dagRunId }, + queryKey, + ), + queryFn: () => + DagRunService.getUpstreamAssetEvents({ dagId, dagRunId }) as TData, + ...options, + }); /** * Get Dag Source * Get source code using file token. diff --git a/airflow/ui/openapi-gen/queries/suspense.ts b/airflow/ui/openapi-gen/queries/suspense.ts index e1d8d3f9d4ee..8a75084926bc 100644 --- a/airflow/ui/openapi-gen/queries/suspense.ts +++ b/airflow/ui/openapi-gen/queries/suspense.ts @@ -469,6 +469,39 @@ export const useDagRunServiceGetDagRunSuspense = < queryFn: () => DagRunService.getDagRun({ dagId, dagRunId }) as TData, ...options, }); +/** + * Get Upstream Asset Events + * If dag run is asset-triggered, return the asset events that triggered it. + * @param data The data for the request. + * @param data.dagId + * @param data.dagRunId + * @returns AssetEventCollectionResponse Successful Response + * @throws ApiError + */ +export const useDagRunServiceGetUpstreamAssetEventsSuspense = < + TData = Common.DagRunServiceGetUpstreamAssetEventsDefaultResponse, + TError = unknown, + TQueryKey extends Array = unknown[], +>( + { + dagId, + dagRunId, + }: { + dagId: string; + dagRunId: string; + }, + queryKey?: TQueryKey, + options?: Omit, "queryKey" | "queryFn">, +) => + useSuspenseQuery({ + queryKey: Common.UseDagRunServiceGetUpstreamAssetEventsKeyFn( + { dagId, dagRunId }, + queryKey, + ), + queryFn: () => + DagRunService.getUpstreamAssetEvents({ dagId, dagRunId }) as TData, + ...options, + }); /** * Get Dag Source * Get source code using file token. diff --git a/airflow/ui/openapi-gen/requests/schemas.gen.ts b/airflow/ui/openapi-gen/requests/schemas.gen.ts index 7bf8f4b02966..a79e11c02157 100644 --- a/airflow/ui/openapi-gen/requests/schemas.gen.ts +++ b/airflow/ui/openapi-gen/requests/schemas.gen.ts @@ -1924,8 +1924,15 @@ export const $DagRunAssetReference = { title: "Start Date", }, end_date: { - type: "string", - format: "date-time", + anyOf: [ + { + type: "string", + format: "date-time", + }, + { + type: "null", + }, + ], title: "End Date", }, state: { diff --git a/airflow/ui/openapi-gen/requests/services.gen.ts b/airflow/ui/openapi-gen/requests/services.gen.ts index ee4ed1e4c41c..ed657317ac9b 100644 --- a/airflow/ui/openapi-gen/requests/services.gen.ts +++ b/airflow/ui/openapi-gen/requests/services.gen.ts @@ -45,6 +45,8 @@ import type { DeleteDagRunResponse, PatchDagRunData, PatchDagRunResponse, + GetUpstreamAssetEventsData, + GetUpstreamAssetEventsResponse, GetDagSourceData, GetDagSourceResponse, GetDagStatsData, @@ -735,6 +737,34 @@ export class DagRunService { }, }); } + + /** + * Get Upstream Asset Events + * If dag run is asset-triggered, return the asset events that triggered it. + * @param data The data for the request. + * @param data.dagId + * @param data.dagRunId + * @returns AssetEventCollectionResponse Successful Response + * @throws ApiError + */ + public static getUpstreamAssetEvents( + data: GetUpstreamAssetEventsData, + ): CancelablePromise { + return __request(OpenAPI, { + method: "GET", + url: "/public/dags/{dag_id}/dagRuns/{dag_run_id}/upstreamAssetEvents", + path: { + dag_id: data.dagId, + dag_run_id: data.dagRunId, + }, + errors: { + 401: "Unauthorized", + 403: "Forbidden", + 404: "Not Found", + 422: "Validation Error", + }, + }); + } } export class DagSourceService { diff --git a/airflow/ui/openapi-gen/requests/types.gen.ts b/airflow/ui/openapi-gen/requests/types.gen.ts index c4eee6b97c21..086ca98b71bd 100644 --- a/airflow/ui/openapi-gen/requests/types.gen.ts +++ b/airflow/ui/openapi-gen/requests/types.gen.ts @@ -419,7 +419,7 @@ export type DagRunAssetReference = { dag_id: string; logical_date: string; start_date: string; - end_date: string; + end_date: string | null; state: string; data_interval_start: string; data_interval_end: string; @@ -1114,6 +1114,13 @@ export type PatchDagRunData = { export type PatchDagRunResponse = DAGRunResponse; +export type GetUpstreamAssetEventsData = { + dagId: string; + dagRunId: string; +}; + +export type GetUpstreamAssetEventsResponse = AssetEventCollectionResponse; + export type GetDagSourceData = { accept?: string; fileToken: string; @@ -1978,6 +1985,33 @@ export type $OpenApiTs = { }; }; }; + "/public/dags/{dag_id}/dagRuns/{dag_run_id}/upstreamAssetEvents": { + get: { + req: GetUpstreamAssetEventsData; + res: { + /** + * Successful Response + */ + 200: AssetEventCollectionResponse; + /** + * Unauthorized + */ + 401: HTTPExceptionResponse; + /** + * Forbidden + */ + 403: HTTPExceptionResponse; + /** + * Not Found + */ + 404: HTTPExceptionResponse; + /** + * Validation Error + */ + 422: HTTPValidationError; + }; + }; + }; "/public/dagSources/{file_token}": { get: { req: GetDagSourceData; diff --git a/tests/api_fastapi/core_api/routes/public/test_dag_run.py b/tests/api_fastapi/core_api/routes/public/test_dag_run.py index 64c3512e88b7..ed11e0ad933f 100644 --- a/tests/api_fastapi/core_api/routes/public/test_dag_run.py +++ b/tests/api_fastapi/core_api/routes/public/test_dag_run.py @@ -21,6 +21,8 @@ import pytest +from airflow import Asset +from airflow.models.asset import AssetEvent, AssetModel from airflow.operators.empty import EmptyOperator from airflow.utils.session import provide_session from airflow.utils.state import DagRunState @@ -49,6 +51,7 @@ DAG2_RUN1_TRIGGERED_BY = DagRunTriggeredByType.CLI DAG2_RUN2_TRIGGERED_BY = DagRunTriggeredByType.REST_API START_DATE = datetime(2024, 6, 15, 0, 0, tzinfo=timezone.utc) +END_DATE = datetime(2024, 6, 15, 0, 0, tzinfo=timezone.utc) EXECUTION_DATE = datetime(2024, 6, 16, 0, 0, tzinfo=timezone.utc) DAG1_RUN1_NOTE = "test_note" @@ -254,3 +257,75 @@ def test_delete_dag_run_not_found(self, test_client): assert response.status_code == 404 body = response.json() assert body["detail"] == "The DagRun with dag_id: `test_dag1` and run_id: `invalid` was not found" + + +class TestGetDagRunAssetTriggerEvents: + def test_should_respond_200(self, test_client, dag_maker, session): + asset1 = Asset(uri="ds1") + + with dag_maker(dag_id="source_dag", start_date=START_DATE, session=session): + EmptyOperator(task_id="task", outlets=[asset1]) + dr = dag_maker.create_dagrun() + ti = dr.task_instances[0] + + asset1_id = session.query(AssetModel.id).filter_by(uri=asset1.uri).scalar() + event = AssetEvent( + asset_id=asset1_id, + source_task_id=ti.task_id, + source_dag_id=ti.dag_id, + source_run_id=ti.run_id, + source_map_index=ti.map_index, + ) + session.add(event) + + with dag_maker(dag_id="TEST_DAG_ID", start_date=START_DATE, session=session): + pass + dr = dag_maker.create_dagrun(run_id="TEST_DAG_RUN_ID", run_type=DagRunType.ASSET_TRIGGERED) + dr.consumed_asset_events.append(event) + + session.commit() + assert event.timestamp + + response = test_client.get( + "/public/dags/TEST_DAG_ID/dagRuns/TEST_DAG_RUN_ID/upstreamAssetEvents", + ) + assert response.status_code == 200 + expected_response = { + "asset_events": [ + { + "timestamp": event.timestamp.isoformat().replace("+00:00", "Z"), + "asset_id": asset1_id, + "uri": asset1.uri, + "extra": {}, + "id": event.id, + "source_dag_id": ti.dag_id, + "source_map_index": ti.map_index, + "source_run_id": ti.run_id, + "source_task_id": ti.task_id, + "created_dagruns": [ + { + "dag_id": "TEST_DAG_ID", + "run_id": "TEST_DAG_RUN_ID", + "data_interval_end": dr.data_interval_end.isoformat().replace("+00:00", "Z"), + "data_interval_start": dr.data_interval_start.isoformat().replace("+00:00", "Z"), + "end_date": None, + "logical_date": dr.logical_date.isoformat().replace("+00:00", "Z"), + "start_date": dr.start_date.isoformat().replace("+00:00", "Z"), + "state": "running", + } + ], + } + ], + "total_entries": 1, + } + assert response.json() == expected_response + + def test_should_respond_404(self, test_client): + response = test_client.get( + "public/dags/invalid-id/dagRuns/invalid-id/upstreamAssetEvents", + ) + assert response.status_code == 404 + assert ( + "The DagRun with dag_id: `invalid-id` and run_id: `invalid-id` was not found" + == response.json()["detail"] + )