From d3d1dde9faa75fce287c3dd317e4ae77b904ac98 Mon Sep 17 00:00:00 2001 From: Pierre Jeambrun Date: Wed, 29 Jan 2025 21:30:41 +0100 Subject: [PATCH] AIP-38 Invalidate DryRun query cache on submit (#46238) --- .../core_api/openapi/v1-generated.yaml | 12 ------- .../core_api/routes/public/task_instances.py | 33 +++++++++---------- .../ui/openapi-gen/requests/services.gen.ts | 2 -- airflow/ui/openapi-gen/requests/types.gen.ts | 8 ----- .../TaskInstance/MarkTaskInstanceAsButton.tsx | 4 ++- .../ui/src/queries/useClearDagRunDryRun.ts | 4 ++- airflow/ui/src/queries/useClearRun.ts | 3 ++ .../ui/src/queries/useClearTaskInstances.ts | 5 +++ .../queries/useClearTaskInstancesDryRun.ts | 14 ++------ .../ui/src/queries/usePatchTaskInstance.ts | 5 +++ .../src/queries/usePatchTaskInstanceDryRun.ts | 20 ++++++----- .../routes/public/test_task_instances.py | 6 ++-- 12 files changed, 53 insertions(+), 63 deletions(-) diff --git a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml index 62df27b568ecd..a85947c6451c2 100644 --- a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml +++ b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml @@ -5754,12 +5754,6 @@ paths: schema: $ref: '#/components/schemas/HTTPExceptionResponse' description: Not Found - '409': - content: - application/json: - schema: - $ref: '#/components/schemas/HTTPExceptionResponse' - description: Conflict '422': description: Validation Error content: @@ -5846,12 +5840,6 @@ paths: schema: $ref: '#/components/schemas/HTTPExceptionResponse' description: Not Found - '409': - content: - application/json: - schema: - $ref: '#/components/schemas/HTTPExceptionResponse' - description: Conflict '422': description: Validation Error content: diff --git a/airflow/api_fastapi/core_api/routes/public/task_instances.py b/airflow/api_fastapi/core_api/routes/public/task_instances.py index c97e190fb2cd8..9c74112097cac 100644 --- a/airflow/api_fastapi/core_api/routes/public/task_instances.py +++ b/airflow/api_fastapi/core_api/routes/public/task_instances.py @@ -717,13 +717,13 @@ def _patch_ti_validate_request( @task_instances_router.patch( task_instances_prefix + "/{task_id}/dry_run", responses=create_openapi_http_exception_doc( - [status.HTTP_404_NOT_FOUND, status.HTTP_400_BAD_REQUEST, status.HTTP_409_CONFLICT], + [status.HTTP_404_NOT_FOUND, status.HTTP_400_BAD_REQUEST], ), ) @task_instances_router.patch( task_instances_prefix + "/{task_id}/{map_index}/dry_run", responses=create_openapi_http_exception_doc( - [status.HTTP_404_NOT_FOUND, status.HTTP_400_BAD_REQUEST, status.HTTP_409_CONFLICT], + [status.HTTP_404_NOT_FOUND, status.HTTP_400_BAD_REQUEST], ), ) def patch_task_instance_dry_run( @@ -744,23 +744,22 @@ def patch_task_instance_dry_run( tis: list[TI] = [] if data.get("new_state"): - tis = dag.set_task_instance_state( - task_id=task_id, - run_id=dag_run_id, - map_indexes=[map_index], - state=data["new_state"], - upstream=body.include_upstream, - downstream=body.include_downstream, - future=body.include_future, - past=body.include_past, - commit=False, - session=session, + tis = ( + dag.set_task_instance_state( + task_id=task_id, + run_id=dag_run_id, + map_indexes=[map_index], + state=data["new_state"], + upstream=body.include_upstream, + downstream=body.include_downstream, + future=body.include_future, + past=body.include_past, + commit=False, + session=session, + ) + or [] ) - if not tis: - raise HTTPException( - status.HTTP_409_CONFLICT, f"Task id {task_id} is already in {data['new_state']} state" - ) elif "note" in data: tis = [ti] diff --git a/airflow/ui/openapi-gen/requests/services.gen.ts b/airflow/ui/openapi-gen/requests/services.gen.ts index bf6e528da11ed..1d64d9d03999f 100644 --- a/airflow/ui/openapi-gen/requests/services.gen.ts +++ b/airflow/ui/openapi-gen/requests/services.gen.ts @@ -2525,7 +2525,6 @@ export class TaskInstanceService { 401: "Unauthorized", 403: "Forbidden", 404: "Not Found", - 409: "Conflict", 422: "Validation Error", }, }); @@ -2566,7 +2565,6 @@ export class TaskInstanceService { 401: "Unauthorized", 403: "Forbidden", 404: "Not Found", - 409: "Conflict", 422: "Validation Error", }, }); diff --git a/airflow/ui/openapi-gen/requests/types.gen.ts b/airflow/ui/openapi-gen/requests/types.gen.ts index 6ff1a083ec6e8..9300f92898a95 100644 --- a/airflow/ui/openapi-gen/requests/types.gen.ts +++ b/airflow/ui/openapi-gen/requests/types.gen.ts @@ -4306,10 +4306,6 @@ export type $OpenApiTs = { * Not Found */ 404: HTTPExceptionResponse; - /** - * Conflict - */ - 409: HTTPExceptionResponse; /** * Validation Error */ @@ -4341,10 +4337,6 @@ export type $OpenApiTs = { * Not Found */ 404: HTTPExceptionResponse; - /** - * Conflict - */ - 409: HTTPExceptionResponse; /** * Validation Error */ diff --git a/airflow/ui/src/components/MarkAs/TaskInstance/MarkTaskInstanceAsButton.tsx b/airflow/ui/src/components/MarkAs/TaskInstance/MarkTaskInstanceAsButton.tsx index fb888e424fcc5..15e4f1fa408e3 100644 --- a/airflow/ui/src/components/MarkAs/TaskInstance/MarkTaskInstanceAsButton.tsx +++ b/airflow/ui/src/components/MarkAs/TaskInstance/MarkTaskInstanceAsButton.tsx @@ -64,7 +64,9 @@ const MarkTaskInstanceAsButton = ({ taskInstance, withText = true }: Props) => { }} value={menuState} > - {menuState} + + {menuState} + ))} diff --git a/airflow/ui/src/queries/useClearDagRunDryRun.ts b/airflow/ui/src/queries/useClearDagRunDryRun.ts index a15099e69ecc5..e082635c7a15c 100644 --- a/airflow/ui/src/queries/useClearDagRunDryRun.ts +++ b/airflow/ui/src/queries/useClearDagRunDryRun.ts @@ -28,6 +28,8 @@ type Props = { requestBody: DAGRunClearBody; }; +export const useClearDagRunDryRunKey = "clearRunDryRun"; + export const useClearDagRunDryRun = ({ dagId, dagRunId, @@ -45,5 +47,5 @@ export const useClearDagRunDryRun = { toaster.create({ description: "Clear Dag Run request failed", @@ -52,6 +54,7 @@ export const useClearDagRun = ({ UseDagServiceGetDagDetailsKeyFn({ dagId }), UseDagRunServiceGetDagRunKeyFn({ dagId, dagRunId }), [useDagRunServiceGetDagRunsKey], + [useClearDagRunDryRunKey, dagId], ]; await Promise.all(queryKeys.map((key) => queryClient.invalidateQueries({ queryKey: key }))); diff --git a/airflow/ui/src/queries/useClearTaskInstances.ts b/airflow/ui/src/queries/useClearTaskInstances.ts index 4c0e7be662b9c..8082483586f88 100644 --- a/airflow/ui/src/queries/useClearTaskInstances.ts +++ b/airflow/ui/src/queries/useClearTaskInstances.ts @@ -27,6 +27,9 @@ import { import type { ClearTaskInstancesBody, TaskInstanceCollectionResponse } from "openapi/requests/types.gen"; import { toaster } from "src/components/ui"; +import { useClearTaskInstancesDryRunKey } from "./useClearTaskInstancesDryRun"; +import { usePatchTaskInstanceDryRunKey } from "./usePatchTaskInstanceDryRun"; + const onError = () => { toaster.create({ description: "Clear Task Instance request failed", @@ -76,6 +79,8 @@ export const useClearTaskInstances = ({ ...taskInstanceKeys, UseDagRunServiceGetDagRunKeyFn({ dagId, dagRunId }), [useDagRunServiceGetDagRunsKey], + [useClearTaskInstancesDryRunKey, dagId], + [usePatchTaskInstanceDryRunKey, dagId, dagRunId], ]; await Promise.all(queryKeys.map((key) => queryClient.invalidateQueries({ queryKey: key }))); diff --git a/airflow/ui/src/queries/useClearTaskInstancesDryRun.ts b/airflow/ui/src/queries/useClearTaskInstancesDryRun.ts index e85fae59a57b8..af565fdb3488f 100644 --- a/airflow/ui/src/queries/useClearTaskInstancesDryRun.ts +++ b/airflow/ui/src/queries/useClearTaskInstancesDryRun.ts @@ -27,6 +27,8 @@ type Props = { requestBody: ClearTaskInstancesBody; }; +export const useClearTaskInstancesDryRunKey = "clearTaskInstanceDryRun"; + export const useClearTaskInstancesDryRun = ({ dagId, options, @@ -42,15 +44,5 @@ export const useClearTaskInstancesDryRun = { toaster.create({ description: "Patch Task Instance request failed", @@ -54,6 +57,8 @@ export const usePatchTaskInstance = ({ UseTaskInstanceServiceGetTaskInstanceKeyFn({ dagId, dagRunId, taskId }), UseTaskInstanceServiceGetMappedTaskInstanceKeyFn({ dagId, dagRunId, mapIndex, taskId }), [useTaskInstanceServiceGetTaskInstancesKey], + [usePatchTaskInstanceDryRunKey, dagId, dagRunId, { mapIndex, taskId }], + [useClearTaskInstancesDryRunKey, dagId], ]; await Promise.all(queryKeys.map((key) => queryClient.invalidateQueries({ queryKey: key }))); diff --git a/airflow/ui/src/queries/usePatchTaskInstanceDryRun.ts b/airflow/ui/src/queries/usePatchTaskInstanceDryRun.ts index 70da064b63b36..cbae4449ae51d 100644 --- a/airflow/ui/src/queries/usePatchTaskInstanceDryRun.ts +++ b/airflow/ui/src/queries/usePatchTaskInstanceDryRun.ts @@ -30,6 +30,8 @@ type Props = { taskId: string; }; +export const usePatchTaskInstanceDryRunKey = "patchTaskInstanceDryRun"; + export const usePatchTaskInstanceDryRun = ({ dagId, dagRunId, @@ -49,15 +51,17 @@ export const usePatchTaskInstanceDryRun =