diff --git a/airflow/api_connexion/endpoints/task_instance_endpoint.py b/airflow/api_connexion/endpoints/task_instance_endpoint.py index 5f528a0b33773..b599f24ab311c 100644 --- a/airflow/api_connexion/endpoints/task_instance_endpoint.py +++ b/airflow/api_connexion/endpoints/task_instance_endpoint.py @@ -814,6 +814,7 @@ def _query(orm_object): ) +@mark_fastapi_migration_done @security.requires_access_dag("GET", DagAccessEntity.TASK_INSTANCE) @provide_session def get_mapped_task_instance_tries( diff --git a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml index a4b1188a6b8be..2610dbcf6123c 100644 --- a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml +++ b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml @@ -4146,6 +4146,75 @@ paths: schema: type: string title: Task Id + - name: map_index + in: query + required: false + schema: + type: integer + default: -1 + title: Map Index + responses: + '200': + description: Successful Response + content: + application/json: + schema: + $ref: '#/components/schemas/TaskInstanceHistoryCollectionResponse' + '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/dags/{dag_id}/dagRuns/{dag_run_id}/taskInstances/{task_id}/{map_index}/tries: + get: + tags: + - Task Instance + summary: Get Mapped Task Instance Tries + operationId: get_mapped_task_instance_tries + 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 + - name: task_id + in: path + required: true + schema: + type: string + title: Task Id + - name: map_index + in: path + required: true + schema: + type: integer + title: Map Index responses: '200': description: Successful Response 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 a866a7d09c9f8..304634cefc68f 100644 --- a/airflow/api_fastapi/core_api/routes/public/task_instances.py +++ b/airflow/api_fastapi/core_api/routes/public/task_instances.py @@ -247,9 +247,9 @@ def get_task_instance_tries( dag_run_id: str, task_id: str, session: Annotated[Session, Depends(get_session)], + map_index: int = -1, ) -> TaskInstanceHistoryCollectionResponse: """Get list of task instances history.""" - map_index = -1 def _query(orm_object: Base) -> Select: query = select(orm_object).where( @@ -278,6 +278,26 @@ def _query(orm_object: Base) -> Select: ) +@task_instances_router.get( + task_instances_prefix + "/{task_id}/{map_index}/tries", + responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]), +) +def get_mapped_task_instance_tries( + dag_id: str, + dag_run_id: str, + task_id: str, + session: Annotated[Session, Depends(get_session)], + map_index: int, +) -> TaskInstanceHistoryCollectionResponse: + return get_task_instance_tries( + dag_id=dag_id, + dag_run_id=dag_run_id, + task_id=task_id, + map_index=map_index, + session=session, + ) + + @task_instances_router.get( task_instances_prefix + "/{task_id}/{map_index}", responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]), diff --git a/airflow/ui/openapi-gen/queries/common.ts b/airflow/ui/openapi-gen/queries/common.ts index 8a186e22980d3..4696713d203c5 100644 --- a/airflow/ui/openapi-gen/queries/common.ts +++ b/airflow/ui/openapi-gen/queries/common.ts @@ -979,16 +979,43 @@ export const UseTaskInstanceServiceGetTaskInstanceTriesKeyFn = ( { dagId, dagRunId, + mapIndex, taskId, }: { dagId: string; dagRunId: string; + mapIndex?: number; taskId: string; }, queryKey?: Array, ) => [ useTaskInstanceServiceGetTaskInstanceTriesKey, - ...(queryKey ?? [{ dagId, dagRunId, taskId }]), + ...(queryKey ?? [{ dagId, dagRunId, mapIndex, taskId }]), +]; +export type TaskInstanceServiceGetMappedTaskInstanceTriesDefaultResponse = + Awaited>; +export type TaskInstanceServiceGetMappedTaskInstanceTriesQueryResult< + TData = TaskInstanceServiceGetMappedTaskInstanceTriesDefaultResponse, + TError = unknown, +> = UseQueryResult; +export const useTaskInstanceServiceGetMappedTaskInstanceTriesKey = + "TaskInstanceServiceGetMappedTaskInstanceTries"; +export const UseTaskInstanceServiceGetMappedTaskInstanceTriesKeyFn = ( + { + dagId, + dagRunId, + mapIndex, + taskId, + }: { + dagId: string; + dagRunId: string; + mapIndex: number; + taskId: string; + }, + queryKey?: Array, +) => [ + useTaskInstanceServiceGetMappedTaskInstanceTriesKey, + ...(queryKey ?? [{ dagId, dagRunId, mapIndex, taskId }]), ]; export type TaskInstanceServiceGetMappedTaskInstanceDefaultResponse = Awaited< ReturnType diff --git a/airflow/ui/openapi-gen/queries/prefetch.ts b/airflow/ui/openapi-gen/queries/prefetch.ts index 4ec4242505f7d..ff83018fc89d7 100644 --- a/airflow/ui/openapi-gen/queries/prefetch.ts +++ b/airflow/ui/openapi-gen/queries/prefetch.ts @@ -1292,6 +1292,7 @@ export const prefetchUseTaskInstanceServiceGetTaskInstanceDependencies1 = ( * @param data.dagId * @param data.dagRunId * @param data.taskId + * @param data.mapIndex * @returns TaskInstanceHistoryCollectionResponse Successful Response * @throws ApiError */ @@ -1300,10 +1301,12 @@ export const prefetchUseTaskInstanceServiceGetTaskInstanceTries = ( { dagId, dagRunId, + mapIndex, taskId, }: { dagId: string; dagRunId: string; + mapIndex?: number; taskId: string; }, ) => @@ -1311,10 +1314,55 @@ export const prefetchUseTaskInstanceServiceGetTaskInstanceTries = ( queryKey: Common.UseTaskInstanceServiceGetTaskInstanceTriesKeyFn({ dagId, dagRunId, + mapIndex, + taskId, + }), + queryFn: () => + TaskInstanceService.getTaskInstanceTries({ + dagId, + dagRunId, + mapIndex, + taskId, + }), + }); +/** + * Get Mapped Task Instance Tries + * @param data The data for the request. + * @param data.dagId + * @param data.dagRunId + * @param data.taskId + * @param data.mapIndex + * @returns TaskInstanceHistoryCollectionResponse Successful Response + * @throws ApiError + */ +export const prefetchUseTaskInstanceServiceGetMappedTaskInstanceTries = ( + queryClient: QueryClient, + { + dagId, + dagRunId, + mapIndex, + taskId, + }: { + dagId: string; + dagRunId: string; + mapIndex: number; + taskId: string; + }, +) => + queryClient.prefetchQuery({ + queryKey: Common.UseTaskInstanceServiceGetMappedTaskInstanceTriesKeyFn({ + dagId, + dagRunId, + mapIndex, taskId, }), queryFn: () => - TaskInstanceService.getTaskInstanceTries({ dagId, dagRunId, taskId }), + TaskInstanceService.getMappedTaskInstanceTries({ + dagId, + dagRunId, + mapIndex, + taskId, + }), }); /** * Get Mapped Task Instance diff --git a/airflow/ui/openapi-gen/queries/queries.ts b/airflow/ui/openapi-gen/queries/queries.ts index 94e2d80c40c9b..ccc8bc1a12dd4 100644 --- a/airflow/ui/openapi-gen/queries/queries.ts +++ b/airflow/ui/openapi-gen/queries/queries.ts @@ -1553,6 +1553,7 @@ export const useTaskInstanceServiceGetTaskInstanceDependencies1 = < * @param data.dagId * @param data.dagRunId * @param data.taskId + * @param data.mapIndex * @returns TaskInstanceHistoryCollectionResponse Successful Response * @throws ApiError */ @@ -1564,10 +1565,12 @@ export const useTaskInstanceServiceGetTaskInstanceTries = < { dagId, dagRunId, + mapIndex, taskId, }: { dagId: string; dagRunId: string; + mapIndex?: number; taskId: string; }, queryKey?: TQueryKey, @@ -1575,13 +1578,57 @@ export const useTaskInstanceServiceGetTaskInstanceTries = < ) => useQuery({ queryKey: Common.UseTaskInstanceServiceGetTaskInstanceTriesKeyFn( - { dagId, dagRunId, taskId }, + { dagId, dagRunId, mapIndex, taskId }, queryKey, ), queryFn: () => TaskInstanceService.getTaskInstanceTries({ dagId, dagRunId, + mapIndex, + taskId, + }) as TData, + ...options, + }); +/** + * Get Mapped Task Instance Tries + * @param data The data for the request. + * @param data.dagId + * @param data.dagRunId + * @param data.taskId + * @param data.mapIndex + * @returns TaskInstanceHistoryCollectionResponse Successful Response + * @throws ApiError + */ +export const useTaskInstanceServiceGetMappedTaskInstanceTries = < + TData = Common.TaskInstanceServiceGetMappedTaskInstanceTriesDefaultResponse, + TError = unknown, + TQueryKey extends Array = unknown[], +>( + { + dagId, + dagRunId, + mapIndex, + taskId, + }: { + dagId: string; + dagRunId: string; + mapIndex: number; + taskId: string; + }, + queryKey?: TQueryKey, + options?: Omit, "queryKey" | "queryFn">, +) => + useQuery({ + queryKey: Common.UseTaskInstanceServiceGetMappedTaskInstanceTriesKeyFn( + { dagId, dagRunId, mapIndex, taskId }, + queryKey, + ), + queryFn: () => + TaskInstanceService.getMappedTaskInstanceTries({ + dagId, + dagRunId, + mapIndex, taskId, }) as TData, ...options, diff --git a/airflow/ui/openapi-gen/queries/suspense.ts b/airflow/ui/openapi-gen/queries/suspense.ts index 5a577d6ba33c1..49125fef08de4 100644 --- a/airflow/ui/openapi-gen/queries/suspense.ts +++ b/airflow/ui/openapi-gen/queries/suspense.ts @@ -1531,6 +1531,7 @@ export const useTaskInstanceServiceGetTaskInstanceDependencies1Suspense = < * @param data.dagId * @param data.dagRunId * @param data.taskId + * @param data.mapIndex * @returns TaskInstanceHistoryCollectionResponse Successful Response * @throws ApiError */ @@ -1542,10 +1543,12 @@ export const useTaskInstanceServiceGetTaskInstanceTriesSuspense = < { dagId, dagRunId, + mapIndex, taskId, }: { dagId: string; dagRunId: string; + mapIndex?: number; taskId: string; }, queryKey?: TQueryKey, @@ -1553,13 +1556,57 @@ export const useTaskInstanceServiceGetTaskInstanceTriesSuspense = < ) => useSuspenseQuery({ queryKey: Common.UseTaskInstanceServiceGetTaskInstanceTriesKeyFn( - { dagId, dagRunId, taskId }, + { dagId, dagRunId, mapIndex, taskId }, queryKey, ), queryFn: () => TaskInstanceService.getTaskInstanceTries({ dagId, dagRunId, + mapIndex, + taskId, + }) as TData, + ...options, + }); +/** + * Get Mapped Task Instance Tries + * @param data The data for the request. + * @param data.dagId + * @param data.dagRunId + * @param data.taskId + * @param data.mapIndex + * @returns TaskInstanceHistoryCollectionResponse Successful Response + * @throws ApiError + */ +export const useTaskInstanceServiceGetMappedTaskInstanceTriesSuspense = < + TData = Common.TaskInstanceServiceGetMappedTaskInstanceTriesDefaultResponse, + TError = unknown, + TQueryKey extends Array = unknown[], +>( + { + dagId, + dagRunId, + mapIndex, + taskId, + }: { + dagId: string; + dagRunId: string; + mapIndex: number; + taskId: string; + }, + queryKey?: TQueryKey, + options?: Omit, "queryKey" | "queryFn">, +) => + useSuspenseQuery({ + queryKey: Common.UseTaskInstanceServiceGetMappedTaskInstanceTriesKeyFn( + { dagId, dagRunId, mapIndex, taskId }, + queryKey, + ), + queryFn: () => + TaskInstanceService.getMappedTaskInstanceTries({ + dagId, + dagRunId, + mapIndex, taskId, }) as TData, ...options, diff --git a/airflow/ui/openapi-gen/requests/services.gen.ts b/airflow/ui/openapi-gen/requests/services.gen.ts index 7c819ae0b7b4c..1f14639496d0c 100644 --- a/airflow/ui/openapi-gen/requests/services.gen.ts +++ b/airflow/ui/openapi-gen/requests/services.gen.ts @@ -110,6 +110,8 @@ import type { GetTaskInstanceDependencies1Response, GetTaskInstanceTriesData, GetTaskInstanceTriesResponse, + GetMappedTaskInstanceTriesData, + GetMappedTaskInstanceTriesResponse, GetMappedTaskInstanceData, GetMappedTaskInstanceResponse, PatchTaskInstance1Data, @@ -1906,6 +1908,7 @@ export class TaskInstanceService { * @param data.dagId * @param data.dagRunId * @param data.taskId + * @param data.mapIndex * @returns TaskInstanceHistoryCollectionResponse Successful Response * @throws ApiError */ @@ -1920,6 +1923,40 @@ export class TaskInstanceService { dag_run_id: data.dagRunId, task_id: data.taskId, }, + query: { + map_index: data.mapIndex, + }, + errors: { + 401: "Unauthorized", + 403: "Forbidden", + 404: "Not Found", + 422: "Validation Error", + }, + }); + } + + /** + * Get Mapped Task Instance Tries + * @param data The data for the request. + * @param data.dagId + * @param data.dagRunId + * @param data.taskId + * @param data.mapIndex + * @returns TaskInstanceHistoryCollectionResponse Successful Response + * @throws ApiError + */ + public static getMappedTaskInstanceTries( + data: GetMappedTaskInstanceTriesData, + ): CancelablePromise { + return __request(OpenAPI, { + method: "GET", + url: "/public/dags/{dag_id}/dagRuns/{dag_run_id}/taskInstances/{task_id}/{map_index}/tries", + path: { + dag_id: data.dagId, + dag_run_id: data.dagRunId, + task_id: data.taskId, + map_index: data.mapIndex, + }, errors: { 401: "Unauthorized", 403: "Forbidden", diff --git a/airflow/ui/openapi-gen/requests/types.gen.ts b/airflow/ui/openapi-gen/requests/types.gen.ts index 754a22cdf7690..f130c187c711a 100644 --- a/airflow/ui/openapi-gen/requests/types.gen.ts +++ b/airflow/ui/openapi-gen/requests/types.gen.ts @@ -1670,12 +1670,23 @@ export type GetTaskInstanceDependencies1Response = export type GetTaskInstanceTriesData = { dagId: string; dagRunId: string; + mapIndex?: number; taskId: string; }; export type GetTaskInstanceTriesResponse = TaskInstanceHistoryCollectionResponse; +export type GetMappedTaskInstanceTriesData = { + dagId: string; + dagRunId: string; + mapIndex: number; + taskId: string; +}; + +export type GetMappedTaskInstanceTriesResponse = + TaskInstanceHistoryCollectionResponse; + export type GetMappedTaskInstanceData = { dagId: string; dagRunId: string; @@ -3363,6 +3374,33 @@ export type $OpenApiTs = { }; }; }; + "/public/dags/{dag_id}/dagRuns/{dag_run_id}/taskInstances/{task_id}/{map_index}/tries": { + get: { + req: GetMappedTaskInstanceTriesData; + res: { + /** + * Successful Response + */ + 200: TaskInstanceHistoryCollectionResponse; + /** + * Unauthorized + */ + 401: HTTPExceptionResponse; + /** + * Forbidden + */ + 403: HTTPExceptionResponse; + /** + * Not Found + */ + 404: HTTPExceptionResponse; + /** + * Validation Error + */ + 422: HTTPValidationError; + }; + }; + }; "/public/dags/{dag_id}/dagRuns/{dag_run_id}/taskInstances/{task_id}/{map_index}": { get: { req: GetMappedTaskInstanceData; diff --git a/tests/api_fastapi/core_api/routes/public/test_task_instances.py b/tests/api_fastapi/core_api/routes/public/test_task_instances.py index 0ae2f4a9a9f15..899af4640ea9d 100644 --- a/tests/api_fastapi/core_api/routes/public/test_task_instances.py +++ b/tests/api_fastapi/core_api/routes/public/test_task_instances.py @@ -2417,6 +2417,96 @@ def test_ti_in_retry_state_not_returned(self, test_client, session): "total_entries": 1, } + def test_mapped_task_should_respond_200(self, test_client, session): + tis = self.create_task_instances(session, task_instances=[{"state": State.FAILED}]) + old_ti = tis[0] + for idx in (1, 2): + ti = TaskInstance(task=old_ti.task, run_id=old_ti.run_id, map_index=idx) + for attr in ["duration", "end_date", "pid", "start_date", "state", "queue"]: + setattr(ti, attr, getattr(old_ti, attr)) + ti.try_number = 1 + session.add(ti) + session.commit() + tis = session.query(TaskInstance).all() + + # Record the task instance history + from airflow.models.taskinstance import clear_task_instances + + clear_task_instances(tis, session) + # Simulate the try_number increasing to new values in TI + for ti in tis: + if ti.map_index > 0: + ti.try_number += 1 + ti.queue = "default_queue" + session.merge(ti) + session.commit() + + # in each loop, we should get the right mapped TI back + for map_index in (1, 2): + # Get the info from TIHistory: try_number 1, try_number 2 is TI table(latest) + response = test_client.get( + "/public/dags/example_python_operator/dagRuns/TEST_DAG_RUN_ID/taskInstances" + f"/print_the_context/{map_index}/tries", + ) + assert response.status_code == 200 + assert ( + response.json()["total_entries"] == 2 + ) # the mapped task was cleared. So both the task instance and its history + assert len(response.json()["task_instances"]) == 2 + assert response.json() == { + "task_instances": [ + { + "dag_id": "example_python_operator", + "duration": 10000.0, + "end_date": "2020-01-03T00:00:00Z", + "executor": None, + "executor_config": "{}", + "hostname": "", + "map_index": map_index, + "max_tries": 0, + "operator": "PythonOperator", + "pid": 100, + "pool": "default_pool", + "pool_slots": 1, + "priority_weight": 9, + "queue": "default_queue", + "queued_when": None, + "start_date": "2020-01-02T00:00:00Z", + "state": "failed", + "task_id": "print_the_context", + "task_display_name": "print_the_context", + "try_number": 1, + "unixname": getuser(), + "dag_run_id": "TEST_DAG_RUN_ID", + }, + { + "dag_id": "example_python_operator", + "duration": 10000.0, + "end_date": "2020-01-03T00:00:00Z", + "executor": None, + "executor_config": "{}", + "hostname": "", + "map_index": map_index, + "max_tries": 1, + "operator": "PythonOperator", + "pid": 100, + "pool": "default_pool", + "pool_slots": 1, + "priority_weight": 9, + "queue": "default_queue", + "queued_when": None, + "start_date": "2020-01-02T00:00:00Z", + "state": None, + "task_id": "print_the_context", + "task_display_name": "print_the_context", + "try_number": 2, + "unixname": getuser(), + "dag_run_id": "TEST_DAG_RUN_ID", + }, + ], + "total_entries": 2, + } + def test_raises_404_for_nonexistent_task_instance(self, test_client, session): self.create_task_instances(session) response = test_client.get(