AIP-38 Invalidate DryRun query cache on submit

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:

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]
 

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",
       },
     });

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
          */

airflow/ui/src/components/MarkAs/TaskInstance/MarkTaskInstanceAsButton.tsx

@@ -64,7 +64,9 @@ const MarkTaskInstanceAsButton = ({ taskInstance, withText = true }: Props) => {
               }}
               value={menuState}
             >
-              <StateBadge state={menuState}>{menuState}</StateBadge>
+              <StateBadge my={1} state={menuState}>
+                {menuState}
+              </StateBadge>
             </Menu.Item>
           ))}
         </Menu.Content>

airflow/ui/src/queries/useClearDagRunDryRun.ts

@@ -28,6 +28,8 @@ type Props<TData, TError> = {
   requestBody: DAGRunClearBody;
 };
 
+export const useClearDagRunDryRunKey = "clearRunDryRun";
+
 export const useClearDagRunDryRun = <TData = TaskInstanceCollectionResponse, TError = unknown>({
   dagId,
   dagRunId,
@@ -45,5 +47,5 @@ export const useClearDagRunDryRun = <TData = TaskInstanceCollectionResponse, TEr
           ...requestBody,
         },
       }) as TData,
-    queryKey: ["clearDagRun", dagId, requestBody.only_failed],
+    queryKey: [useClearDagRunDryRunKey, dagId, { only_failed: requestBody.only_failed }],
   });

airflow/ui/src/queries/useClearRun.ts

@@ -27,6 +27,8 @@ import {
 } from "openapi/queries";
 import { toaster } from "src/components/ui";
 
+import { useClearDagRunDryRunKey } from "./useClearDagRunDryRun";
+
 const onError = () => {
   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 })));

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 })));

airflow/ui/src/queries/useClearTaskInstancesDryRun.ts

@@ -27,6 +27,8 @@ type Props<TData, TError> = {
   requestBody: ClearTaskInstancesBody;
 };
 
+export const useClearTaskInstancesDryRunKey = "clearTaskInstanceDryRun";
+
 export const useClearTaskInstancesDryRun = <TData = PostClearTaskInstancesResponse, TError = unknown>({
   dagId,
   options,
@@ -42,15 +44,5 @@ export const useClearTaskInstancesDryRun = <TData = PostClearTaskInstancesRespon
           ...requestBody,
         },
       }) as TData,
-    queryKey: [
-      "clearTaskInstance",
-      dagId,
-      requestBody.dag_run_id,
-      requestBody.only_failed,
-      requestBody.task_ids,
-      requestBody.include_downstream,
-      requestBody.include_future,
-      requestBody.include_past,
-      requestBody.include_upstream,
-    ],
+    queryKey: [useClearTaskInstancesDryRunKey, dagId, requestBody],
   });

airflow/ui/src/queries/usePatchTaskInstance.ts

@@ -26,6 +26,9 @@ import {
 } from "openapi/queries";
 import { toaster } from "src/components/ui";
 
+import { useClearTaskInstancesDryRunKey } from "./useClearTaskInstancesDryRun";
+import { usePatchTaskInstanceDryRunKey } from "./usePatchTaskInstanceDryRun";
+
 const onError = () => {
   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 })));

airflow/ui/src/queries/usePatchTaskInstanceDryRun.ts

@@ -30,6 +30,8 @@ type Props<TData, TError> = {
   taskId: string;
 };
 
+export const usePatchTaskInstanceDryRunKey = "patchTaskInstanceDryRun";
+
 export const usePatchTaskInstanceDryRun = <TData = PatchTaskInstanceDryRunResponse, TError = unknown>({
   dagId,
   dagRunId,
@@ -49,15 +51,17 @@ export const usePatchTaskInstanceDryRun = <TData = PatchTaskInstanceDryRunRespon
         taskId,
       }) as TData,
     queryKey: [
-      "patchTaskInstanceDryRun",
+      usePatchTaskInstanceDryRunKey,
       dagId,
       dagRunId,
-      taskId,
-      mapIndex,
-      requestBody.new_state,
-      requestBody.include_downstream,
-      requestBody.include_future,
-      requestBody.include_past,
-      requestBody.include_upstream,
+      {
+        include_downstream: requestBody.include_downstream,
+        include_future: requestBody.include_future,
+        include_past: requestBody.include_past,
+        include_upstream: requestBody.include_upstream,
+        mapIndex,
+        new_state: requestBody.new_state,
+        taskId,
+      },
     ],
   });

tests/api_fastapi/core_api/routes/public/test_task_instances.py

@@ -3492,7 +3492,7 @@ def test_update_mask_should_call_mocked_api(
         assert mock_set_ti_state.call_count == set_ti_state_call_count
 
     @mock.patch("airflow.models.dag.DAG.set_task_instance_state")
-    def test_should_raise_409_for_updating_same_task_instance_state(
+    def test_should_return_empty_list_for_updating_same_task_instance_state(
         self, mock_set_ti_state, test_client, session
     ):
         self.create_task_instances(session)
@@ -3505,5 +3505,5 @@ def test_should_raise_409_for_updating_same_task_instance_state(
                 "new_state": "success",
             },
         )
-        assert response.status_code == 409
-        assert "Task id print_the_context is already in success state" in response.text
+        assert response.status_code == 200
+        assert response.json() == {"task_instances": [], "total_entries": 0}