From 29e8e92cf832dec781b3aabd2e7fd78c7b631844 Mon Sep 17 00:00:00 2001 From: kalyanr Date: Sat, 8 Mar 2025 16:10:37 +0530 Subject: [PATCH 01/27] add dag_parsing --- .../api_fastapi/core_api/openapi/v1-generated.yaml | 10 ++++++++++ .../core_api/routes/public/dag_parsing.py | 3 ++- airflow/ui/openapi-gen/queries/queries.ts | 7 +++++-- airflow/ui/openapi-gen/requests/services.gen.ts | 4 ++++ airflow/ui/openapi-gen/requests/types.gen.ts | 1 + .../core_api/routes/public/test_dag_parsing.py | 12 ++++++++++++ 6 files changed, 34 insertions(+), 3 deletions(-) diff --git a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml index ac9079c251d50..d4329b9c0e113 100644 --- a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml +++ b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml @@ -6814,6 +6814,8 @@ paths: summary: Reparse Dag File description: Request re-parsing a DAG file. operationId: reparse_dag_file + security: + - OAuth2PasswordBearer: [] parameters: - name: file_token in: path @@ -6821,6 +6823,14 @@ paths: schema: type: string title: File Token + - name: dag_id + in: query + required: false + schema: + anyOf: + - type: string + - type: 'null' + title: Dag Id responses: '201': description: Successful Response diff --git a/airflow/api_fastapi/core_api/routes/public/dag_parsing.py b/airflow/api_fastapi/core_api/routes/public/dag_parsing.py index f5b7f2c359335..a4deb3b18de02 100644 --- a/airflow/api_fastapi/core_api/routes/public/dag_parsing.py +++ b/airflow/api_fastapi/core_api/routes/public/dag_parsing.py @@ -27,6 +27,7 @@ from airflow.api_fastapi.common.db.common import SessionDep from airflow.api_fastapi.common.router import AirflowRouter from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc +from airflow.api_fastapi.core_api.security import requires_access_dag from airflow.api_fastapi.logging.decorators import action_logging from airflow.models.dag import DagModel from airflow.models.dagbag import DagPriorityParsingRequest @@ -41,7 +42,7 @@ "", responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]), status_code=status.HTTP_201_CREATED, - dependencies=[Depends(action_logging())], + dependencies=[Depends(requires_access_dag(method="PUT")), Depends(action_logging())], ) def reparse_dag_file( file_token: str, diff --git a/airflow/ui/openapi-gen/queries/queries.ts b/airflow/ui/openapi-gen/queries/queries.ts index 955bb41075246..842c271814c71 100644 --- a/airflow/ui/openapi-gen/queries/queries.ts +++ b/airflow/ui/openapi-gen/queries/queries.ts @@ -3599,6 +3599,7 @@ export const useBackfillServiceCancelBackfill = < * Request re-parsing a DAG file. * @param data The data for the request. * @param data.fileToken + * @param data.dagId * @returns null Successful Response * @throws ApiError */ @@ -3612,6 +3613,7 @@ export const useDagParsingServiceReparseDagFile = < TData, TError, { + dagId?: string; fileToken: string; }, TContext @@ -3623,12 +3625,13 @@ export const useDagParsingServiceReparseDagFile = < TData, TError, { + dagId?: string; fileToken: string; }, TContext >({ - mutationFn: ({ fileToken }) => - DagParsingService.reparseDagFile({ fileToken }) as unknown as Promise, + mutationFn: ({ dagId, fileToken }) => + DagParsingService.reparseDagFile({ dagId, fileToken }) as unknown as Promise, ...options, }); /** diff --git a/airflow/ui/openapi-gen/requests/services.gen.ts b/airflow/ui/openapi-gen/requests/services.gen.ts index 56b080ef68172..d7f07322284fb 100644 --- a/airflow/ui/openapi-gen/requests/services.gen.ts +++ b/airflow/ui/openapi-gen/requests/services.gen.ts @@ -3394,6 +3394,7 @@ export class DagParsingService { * Request re-parsing a DAG file. * @param data The data for the request. * @param data.fileToken + * @param data.dagId * @returns null Successful Response * @throws ApiError */ @@ -3404,6 +3405,9 @@ export class DagParsingService { path: { file_token: data.fileToken, }, + query: { + dag_id: data.dagId, + }, 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 724b6d5841090..d2cd8bf45b6fc 100644 --- a/airflow/ui/openapi-gen/requests/types.gen.ts +++ b/airflow/ui/openapi-gen/requests/types.gen.ts @@ -2494,6 +2494,7 @@ export type BulkVariablesData = { export type BulkVariablesResponse = BulkResponse; export type ReparseDagFileData = { + dagId?: string | null; fileToken: string; }; diff --git a/tests/api_fastapi/core_api/routes/public/test_dag_parsing.py b/tests/api_fastapi/core_api/routes/public/test_dag_parsing.py index 0943684e32192..00ab7dbbaf324 100644 --- a/tests/api_fastapi/core_api/routes/public/test_dag_parsing.py +++ b/tests/api_fastapi/core_api/routes/public/test_dag_parsing.py @@ -65,6 +65,18 @@ def test_201_and_400_requests(self, url_safe_serializer, session, test_client): assert parsing_requests[0].fileloc == test_dag.fileloc _check_last_log(session, dag_id=None, event="reparse_dag_file", logical_date=None) + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.put( + "/public/parseDagFile/token", headers={"Accept": "application/json"} + ) + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.put( + "/public/parseDagFile/token", headers={"Accept": "application/json"} + ) + assert response.status_code == 403 + def test_bad_file_request(self, url_safe_serializer, session, test_client): url = f"/public/parseDagFile/{url_safe_serializer.dumps('/some/random/file.py')}" response = test_client.put(url, headers={"Accept": "application/json"}) From 3cb1a1137d9d8c888b8c5482a5f6daf9c013fffe Mon Sep 17 00:00:00 2001 From: kalyanr Date: Sat, 1 Mar 2025 18:09:03 +0530 Subject: [PATCH 02/27] add dag_run --- .../core_api/openapi/v1-generated.yaml | 46 ++++++++++--- .../core_api/routes/public/dag_run.py | 41 ++++++++--- airflow/ui/openapi-gen/queries/queries.ts | 8 +-- airflow/ui/openapi-gen/requests/types.gen.ts | 16 ++--- .../core_api/routes/public/test_dag_run.py | 68 +++++++++++++++++++ 5 files changed, 151 insertions(+), 28 deletions(-) diff --git a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml index d4329b9c0e113..7c1a165fab058 100644 --- a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml +++ b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml @@ -2085,12 +2085,16 @@ paths: - DagRun summary: Get Dag Run operationId: get_dag_run + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id - name: dag_run_id in: path @@ -2135,12 +2139,16 @@ paths: summary: Delete Dag Run description: Delete a DAG Run entry. operationId: delete_dag_run + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id - name: dag_run_id in: path @@ -2187,12 +2195,16 @@ paths: summary: Patch Dag Run description: Modify a DAG Run. operationId: patch_dag_run + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id - name: dag_run_id in: path @@ -2268,7 +2280,9 @@ paths: in: path required: true schema: - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id - name: dag_run_id in: path @@ -2313,12 +2327,16 @@ paths: - DagRun summary: Clear Dag Run operationId: clear_dag_run + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id - name: dag_run_id in: path @@ -2377,12 +2395,16 @@ paths: This endpoint allows specifying `~` as the dag_id to retrieve Dag Runs for all DAGs.' operationId: get_dag_runs + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id - name: limit in: query @@ -2542,11 +2564,16 @@ paths: summary: Trigger Dag Run description: Trigger a DAG. operationId: trigger_dag_run + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: + anyOf: + - type: string + - type: 'null' title: Dag Id requestBody: required: true @@ -2604,13 +2631,16 @@ paths: summary: Get List Dag Runs Batch description: Get a list of DAG Runs. operationId: get_list_dag_runs_batch + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - const: '~' - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id requestBody: required: true 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 8703fe42e423a..eff8d729b18c4 100644 --- a/airflow/api_fastapi/core_api/routes/public/dag_run.py +++ b/airflow/api_fastapi/core_api/routes/public/dag_run.py @@ -29,6 +29,7 @@ set_dag_run_state_to_queued, set_dag_run_state_to_success, ) +from airflow.api_fastapi.auth.managers.models.resource_details import DagAccessEntity from airflow.api_fastapi.common.db.common import SessionDep, paginated_select from airflow.api_fastapi.common.parameters import ( FilterOptionEnum, @@ -59,7 +60,7 @@ TaskInstanceResponse, ) from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc -from airflow.api_fastapi.core_api.security import requires_access_asset +from airflow.api_fastapi.core_api.security import requires_access_asset, requires_access_dag from airflow.api_fastapi.logging.decorators import action_logging from airflow.exceptions import ParamValidationError from airflow.listeners.listener import get_listener_manager @@ -78,6 +79,7 @@ status.HTTP_404_NOT_FOUND, ] ), + dependencies=[Depends(requires_access_dag(method="GET", access_entity=DagAccessEntity.RUN))], ) def get_dag_run(dag_id: str, dag_run_id: str, session: SessionDep) -> DAGRunResponse: dag_run = session.scalar(select(DagRun).filter_by(dag_id=dag_id, run_id=dag_run_id)) @@ -99,7 +101,10 @@ def get_dag_run(dag_id: str, dag_run_id: str, session: SessionDep) -> DAGRunResp status.HTTP_404_NOT_FOUND, ], ), - dependencies=[Depends(action_logging())], + dependencies=[ + Depends(requires_access_dag(method="DELETE", access_entity=DagAccessEntity.RUN)), + Depends(action_logging()), + ], ) def delete_dag_run(dag_id: str, dag_run_id: str, session: SessionDep): """Delete a DAG Run entry.""" @@ -121,7 +126,10 @@ def delete_dag_run(dag_id: str, dag_run_id: str, session: SessionDep): status.HTTP_404_NOT_FOUND, ], ), - dependencies=[Depends(action_logging())], + dependencies=[ + Depends(requires_access_dag(method="PUT", access_entity=DagAccessEntity.RUN)), + Depends(action_logging()), + ], ) def patch_dag_run( dag_id: str, @@ -190,7 +198,10 @@ def patch_dag_run( status.HTTP_404_NOT_FOUND, ] ), - dependencies=[Depends(requires_access_asset(method="GET"))], + dependencies=[ + Depends(requires_access_asset(method="GET")), + Depends(requires_access_dag(method="GET", access_entity=DagAccessEntity.RUN)), + ], ) def get_upstream_asset_events( dag_id: str, dag_run_id: str, session: SessionDep @@ -217,7 +228,10 @@ def get_upstream_asset_events( @dag_run_router.post( "/{dag_run_id}/clear", responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]), - dependencies=[Depends(action_logging())], + dependencies=[ + Depends(requires_access_dag(method="POST", access_entity=DagAccessEntity.RUN)), + Depends(action_logging()), + ], ) def clear_dag_run( dag_id: str, @@ -263,7 +277,11 @@ def clear_dag_run( return dag_run_cleared -@dag_run_router.get("", responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND])) +@dag_run_router.get( + "", + responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]), + dependencies=[Depends(requires_access_dag(method="GET", access_entity=DagAccessEntity.RUN))], +) def get_dag_runs( dag_id: str, limit: QueryLimit, @@ -337,7 +355,10 @@ def get_dag_runs( status.HTTP_409_CONFLICT, ] ), - dependencies=[Depends(action_logging())], + dependencies=[ + Depends(requires_access_dag(method="POST", access_entity=DagAccessEntity.RUN)), + Depends(action_logging()), + ], ) def trigger_dag_run( dag_id, @@ -383,7 +404,11 @@ def trigger_dag_run( raise HTTPException(status.HTTP_400_BAD_REQUEST, str(e)) -@dag_run_router.post("/list", responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND])) +@dag_run_router.post( + "/list", + responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]), + dependencies=[Depends(requires_access_dag(method="POST", access_entity=DagAccessEntity.RUN))], +) def get_list_dag_runs_batch( dag_id: Literal["~"], body: DAGRunsBatchBody, session: SessionDep ) -> DAGRunCollectionResponse: diff --git a/airflow/ui/openapi-gen/queries/queries.ts b/airflow/ui/openapi-gen/queries/queries.ts index 842c271814c71..be0e58f3afcb0 100644 --- a/airflow/ui/openapi-gen/queries/queries.ts +++ b/airflow/ui/openapi-gen/queries/queries.ts @@ -3219,7 +3219,7 @@ export const useDagRunServiceTriggerDagRun = < TData, TError, { - dagId: unknown; + dagId: string; requestBody: TriggerDAGRunPostBody; }, TContext @@ -3231,7 +3231,7 @@ export const useDagRunServiceTriggerDagRun = < TData, TError, { - dagId: unknown; + dagId: string; requestBody: TriggerDAGRunPostBody; }, TContext @@ -3259,7 +3259,7 @@ export const useDagRunServiceGetListDagRunsBatch = < TData, TError, { - dagId: "~"; + dagId: string; requestBody: DAGRunsBatchBody; }, TContext @@ -3271,7 +3271,7 @@ export const useDagRunServiceGetListDagRunsBatch = < TData, TError, { - dagId: "~"; + dagId: string; requestBody: DAGRunsBatchBody; }, TContext diff --git a/airflow/ui/openapi-gen/requests/types.gen.ts b/airflow/ui/openapi-gen/requests/types.gen.ts index d2cd8bf45b6fc..dfac68c9320cd 100644 --- a/airflow/ui/openapi-gen/requests/types.gen.ts +++ b/airflow/ui/openapi-gen/requests/types.gen.ts @@ -1899,21 +1899,21 @@ export type TestConnectionResponse = ConnectionTestResponse; export type CreateDefaultConnectionsResponse = void; export type GetDagRunData = { - dagId: string; + dagId: string | null; dagRunId: string; }; export type GetDagRunResponse = DAGRunResponse; export type DeleteDagRunData = { - dagId: string; + dagId: string | null; dagRunId: string; }; export type DeleteDagRunResponse = void; export type PatchDagRunData = { - dagId: string; + dagId: string | null; dagRunId: string; requestBody: DAGRunPatchBody; updateMask?: Array | null; @@ -1922,14 +1922,14 @@ export type PatchDagRunData = { export type PatchDagRunResponse = DAGRunResponse; export type GetUpstreamAssetEventsData = { - dagId: string; + dagId: string | null; dagRunId: string; }; export type GetUpstreamAssetEventsResponse = AssetEventCollectionResponse; export type ClearDagRunData = { - dagId: string; + dagId: string | null; dagRunId: string; requestBody: DAGRunClearBody; }; @@ -1937,7 +1937,7 @@ export type ClearDagRunData = { export type ClearDagRunResponse = TaskInstanceCollectionResponse | DAGRunResponse; export type GetDagRunsData = { - dagId: string; + dagId: string | null; endDateGte?: string | null; endDateLte?: string | null; limit?: number; @@ -1957,14 +1957,14 @@ export type GetDagRunsData = { export type GetDagRunsResponse = DAGRunCollectionResponse; export type TriggerDagRunData = { - dagId: unknown; + dagId: string | null; requestBody: TriggerDAGRunPostBody; }; export type TriggerDagRunResponse = DAGRunResponse; export type GetListDagRunsBatchData = { - dagId: "~"; + dagId: string | null; requestBody: DAGRunsBatchBody; }; 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 f62bff0665ec6..61199fc1f37a5 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 @@ -244,6 +244,14 @@ def test_get_dag_run_not_found(self, test_client): body = response.json() assert body["detail"] == "The DagRun with dag_id: `test_dag1` and run_id: `invalid` was not found" + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.get(f"/public/dags/{DAG1_ID}/dagRuns/invalid") + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.get(f"/public/dags/{DAG1_ID}/dagRuns/invalid") + assert response.status_code == 403 + class TestGetDagRuns: @pytest.mark.parametrize("dag_id, total_entries", [(DAG1_ID, 2), (DAG2_ID, 2), ("~", 4)]) @@ -277,6 +285,14 @@ def test_invalid_order_by_raises_400(self, test_client): == "Ordering with 'invalid' is disallowed or the attribute does not exist on the model" ) + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.get("/public/dags/test_dag1/dagRuns?order_by=invalid") + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.get("/public/dags/test_dag1/dagRuns?order_by=invalid") + assert response.status_code == 403 + @pytest.mark.parametrize( "order_by,expected_order", [ @@ -550,6 +566,14 @@ def test_list_dag_runs_return_200(self, test_client, session): expected = get_dag_run_dict(run) assert each == expected + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.post("/public/dags/~/dagRuns/list", json={}) + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.post("/public/dags/~/dagRuns/list", json={}) + assert response.status_code == 403 + def test_list_dag_runs_with_invalid_dag_id(self, test_client): response = test_client.post("/public/dags/invalid/dagRuns/list", json={}) assert response.status_code == 422 @@ -909,6 +933,14 @@ def test_patch_dag_run(self, test_client, dag_id, run_id, patch_body, response_b assert body.get("note") == response_body.get("note") _check_last_log(session, dag_id=dag_id, event="patch_dag_run", logical_date=None) + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.patch("/public/dags/dag_1/dagRuns/run_1", json={}) + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.patch("/public/dags/dag_1/dagRuns/run_1", json={}) + assert response.status_code == 403 + @pytest.mark.parametrize( "query_params, patch_body, response_body, expected_status_code", [ @@ -1008,6 +1040,14 @@ def test_delete_dag_run_not_found(self, test_client): body = response.json() assert body["detail"] == "The DagRun with dag_id: `test_dag1` and run_id: `invalid` was not found" + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.delete(f"/public/dags/{DAG1_ID}/dagRuns/invalid") + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.delete(f"/public/dags/{DAG1_ID}/dagRuns/invalid") + assert response.status_code == 403 + class TestGetDagRunAssetTriggerEvents: @pytest.mark.usefixtures("configure_git_connection_for_dag_bundle") @@ -1115,6 +1155,20 @@ def test_clear_dag_run(self, test_client, session): logical_date=None, ) + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.post( + f"/public/dags/{DAG1_ID}/dagRuns/{DAG1_RUN1_ID}/clear", + json={"dry_run": False}, + ) + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.post( + f"/public/dags/{DAG1_ID}/dagRuns/{DAG1_RUN1_ID}/clear", + json={"dry_run": False}, + ) + assert response.status_code == 403 + @pytest.mark.parametrize( "body, dag_run_id, expected_state", [ @@ -1246,6 +1300,20 @@ def test_should_respond_200( assert response.json() == expected_response_json _check_last_log(session, dag_id=DAG1_ID, event="trigger_dag_run", logical_date=None) + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.post( + f"/public/dags/{DAG1_ID}/dagRuns", + json={}, + ) + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.post( + f"/public/dags/{DAG1_ID}/dagRuns", + json={}, + ) + assert response.status_code == 403 + @pytest.mark.parametrize( "post_body, expected_detail", [ From 5a404d7af9574c4bd698fc7c04aae769e4284e69 Mon Sep 17 00:00:00 2001 From: kalyanr Date: Sat, 1 Mar 2025 18:40:05 +0530 Subject: [PATCH 03/27] add dag_source --- .../api_fastapi/core_api/openapi/v1-generated.yaml | 6 +++++- .../core_api/routes/public/dag_sources.py | 5 ++++- airflow/ui/openapi-gen/requests/types.gen.ts | 2 +- .../core_api/routes/public/test_dag_sources.py | 12 ++++++++++++ 4 files changed, 22 insertions(+), 3 deletions(-) diff --git a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml index 7c1a165fab058..225dc6097ab41 100644 --- a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml +++ b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml @@ -2686,12 +2686,16 @@ paths: summary: Get Dag Source description: Get source code using file token. operationId: get_dag_source + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id - name: version_number in: query diff --git a/airflow/api_fastapi/core_api/routes/public/dag_sources.py b/airflow/api_fastapi/core_api/routes/public/dag_sources.py index 4337c92107322..fb00a4dd553b5 100644 --- a/airflow/api_fastapi/core_api/routes/public/dag_sources.py +++ b/airflow/api_fastapi/core_api/routes/public/dag_sources.py @@ -16,14 +16,16 @@ # under the License. from __future__ import annotations -from fastapi import HTTPException, Response, status +from fastapi import Depends, HTTPException, Response, status +from airflow.api_fastapi.auth.managers.models.resource_details import DagAccessEntity from airflow.api_fastapi.common.db.common import SessionDep from airflow.api_fastapi.common.headers import HeaderAcceptJsonOrText from airflow.api_fastapi.common.router import AirflowRouter from airflow.api_fastapi.common.types import Mimetype from airflow.api_fastapi.core_api.datamodels.dag_sources import DAGSourceResponse from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc +from airflow.api_fastapi.core_api.security import requires_access_dag from airflow.models.dag_version import DagVersion dag_sources_router = AirflowRouter(tags=["DagSource"], prefix="/dagSources") @@ -47,6 +49,7 @@ }, }, response_model=DAGSourceResponse, + dependencies=[Depends(requires_access_dag(method="GET", access_entity=DagAccessEntity.CODE))], ) def get_dag_source( accept: HeaderAcceptJsonOrText, diff --git a/airflow/ui/openapi-gen/requests/types.gen.ts b/airflow/ui/openapi-gen/requests/types.gen.ts index dfac68c9320cd..bf9a05f25fa34 100644 --- a/airflow/ui/openapi-gen/requests/types.gen.ts +++ b/airflow/ui/openapi-gen/requests/types.gen.ts @@ -1972,7 +1972,7 @@ export type GetListDagRunsBatchResponse = DAGRunCollectionResponse; export type GetDagSourceData = { accept?: "application/json" | "text/plain" | "*/*"; - dagId: string; + dagId: string | null; versionNumber?: number | null; }; diff --git a/tests/api_fastapi/core_api/routes/public/test_dag_sources.py b/tests/api_fastapi/core_api/routes/public/test_dag_sources.py index 6da8902b232a0..8bd53ea2f97b9 100644 --- a/tests/api_fastapi/core_api/routes/public/test_dag_sources.py +++ b/tests/api_fastapi/core_api/routes/public/test_dag_sources.py @@ -77,6 +77,18 @@ def test_should_respond_200_text(self, test_client, test_dag): json.loads(response.content.decode()) assert response.headers["Content-Type"].startswith("text/plain") + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.get( + f"{API_PREFIX}/{TEST_DAG_ID}", headers={"Accept": "text/plain"} + ) + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.get( + f"{API_PREFIX}/{TEST_DAG_ID}", headers={"Accept": "text/plain"} + ) + assert response.status_code == 403 + @pytest.mark.parametrize( "headers", [{"Accept": "application/json"}, {"Accept": "application/json; charset=utf-8"}, {}] ) From ff0210be345f6b501ba9c2cb58c1d71ed30f5e10 Mon Sep 17 00:00:00 2001 From: kalyanr Date: Sat, 1 Mar 2025 19:01:08 +0530 Subject: [PATCH 04/27] add dag_stats --- .../core_api/openapi/v1-generated.yaml | 10 ++++++++++ .../core_api/routes/public/dag_stats.py | 3 +++ airflow/ui/openapi-gen/queries/common.ts | 4 +++- airflow/ui/openapi-gen/queries/prefetch.ts | 7 +++++-- airflow/ui/openapi-gen/queries/queries.ts | 7 +++++-- airflow/ui/openapi-gen/queries/suspense.ts | 7 +++++-- .../ui/openapi-gen/requests/services.gen.ts | 2 ++ airflow/ui/openapi-gen/requests/types.gen.ts | 1 + .../core_api/routes/public/test_dag_stats.py | 20 +++++++++++++------ 9 files changed, 48 insertions(+), 13 deletions(-) diff --git a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml index 225dc6097ab41..6735f60a0e38c 100644 --- a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml +++ b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml @@ -2770,7 +2770,17 @@ paths: summary: Get Dag Stats description: Get Dag statistics. operationId: get_dag_stats + security: + - OAuth2PasswordBearer: [] parameters: + - name: dag_id + in: query + required: false + schema: + anyOf: + - type: string + - type: 'null' + title: Dag Id - name: dag_ids in: query required: false diff --git a/airflow/api_fastapi/core_api/routes/public/dag_stats.py b/airflow/api_fastapi/core_api/routes/public/dag_stats.py index a6aa6063c263b..d70ef3fe6143f 100644 --- a/airflow/api_fastapi/core_api/routes/public/dag_stats.py +++ b/airflow/api_fastapi/core_api/routes/public/dag_stats.py @@ -21,6 +21,7 @@ from fastapi import Depends, status +from airflow.api_fastapi.auth.managers.models.resource_details import DagAccessEntity from airflow.api_fastapi.common.db.common import ( SessionDep, paginated_select, @@ -38,6 +39,7 @@ DagStatsStateResponse, ) from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc +from airflow.api_fastapi.core_api.security import requires_access_dag from airflow.models.dagrun import DagRun from airflow.utils.state import DagRunState @@ -52,6 +54,7 @@ status.HTTP_404_NOT_FOUND, ] ), + dependencies=[Depends(requires_access_dag(method="GET", access_entity=DagAccessEntity.RUN))], ) def get_dag_stats( session: SessionDep, diff --git a/airflow/ui/openapi-gen/queries/common.ts b/airflow/ui/openapi-gen/queries/common.ts index 8e22ea4d4a8cc..99cfdeac4d3ae 100644 --- a/airflow/ui/openapi-gen/queries/common.ts +++ b/airflow/ui/openapi-gen/queries/common.ts @@ -683,12 +683,14 @@ export type DagStatsServiceGetDagStatsQueryResult< export const useDagStatsServiceGetDagStatsKey = "DagStatsServiceGetDagStats"; export const UseDagStatsServiceGetDagStatsKeyFn = ( { + dagId, dagIds, }: { + dagId?: string; dagIds?: string[]; } = {}, queryKey?: Array, -) => [useDagStatsServiceGetDagStatsKey, ...(queryKey ?? [{ dagIds }])]; +) => [useDagStatsServiceGetDagStatsKey, ...(queryKey ?? [{ dagId, dagIds }])]; export type DagReportServiceGetDagReportsDefaultResponse = Awaited< ReturnType >; diff --git a/airflow/ui/openapi-gen/queries/prefetch.ts b/airflow/ui/openapi-gen/queries/prefetch.ts index 7987c721b2006..63cd1639e7f72 100644 --- a/airflow/ui/openapi-gen/queries/prefetch.ts +++ b/airflow/ui/openapi-gen/queries/prefetch.ts @@ -923,6 +923,7 @@ export const prefetchUseDagSourceServiceGetDagSource = ( * Get Dag Stats * Get Dag statistics. * @param data The data for the request. + * @param data.dagId * @param data.dagIds * @returns DagStatsCollectionResponse Successful Response * @throws ApiError @@ -930,14 +931,16 @@ export const prefetchUseDagSourceServiceGetDagSource = ( export const prefetchUseDagStatsServiceGetDagStats = ( queryClient: QueryClient, { + dagId, dagIds, }: { + dagId?: string; dagIds?: string[]; } = {}, ) => queryClient.prefetchQuery({ - queryKey: Common.UseDagStatsServiceGetDagStatsKeyFn({ dagIds }), - queryFn: () => DagStatsService.getDagStats({ dagIds }), + queryKey: Common.UseDagStatsServiceGetDagStatsKeyFn({ dagId, dagIds }), + queryFn: () => DagStatsService.getDagStats({ dagId, dagIds }), }); /** * Get Dag Reports diff --git a/airflow/ui/openapi-gen/queries/queries.ts b/airflow/ui/openapi-gen/queries/queries.ts index be0e58f3afcb0..44df779c7e744 100644 --- a/airflow/ui/openapi-gen/queries/queries.ts +++ b/airflow/ui/openapi-gen/queries/queries.ts @@ -1115,6 +1115,7 @@ export const useDagSourceServiceGetDagSource = < * Get Dag Stats * Get Dag statistics. * @param data The data for the request. + * @param data.dagId * @param data.dagIds * @returns DagStatsCollectionResponse Successful Response * @throws ApiError @@ -1125,16 +1126,18 @@ export const useDagStatsServiceGetDagStats = < TQueryKey extends Array = unknown[], >( { + dagId, dagIds, }: { + dagId?: string; dagIds?: string[]; } = {}, queryKey?: TQueryKey, options?: Omit, "queryKey" | "queryFn">, ) => useQuery({ - queryKey: Common.UseDagStatsServiceGetDagStatsKeyFn({ dagIds }, queryKey), - queryFn: () => DagStatsService.getDagStats({ dagIds }) as TData, + queryKey: Common.UseDagStatsServiceGetDagStatsKeyFn({ dagId, dagIds }, queryKey), + queryFn: () => DagStatsService.getDagStats({ dagId, dagIds }) as TData, ...options, }); /** diff --git a/airflow/ui/openapi-gen/queries/suspense.ts b/airflow/ui/openapi-gen/queries/suspense.ts index fc25561641a11..699d7fd817c5a 100644 --- a/airflow/ui/openapi-gen/queries/suspense.ts +++ b/airflow/ui/openapi-gen/queries/suspense.ts @@ -1092,6 +1092,7 @@ export const useDagSourceServiceGetDagSourceSuspense = < * Get Dag Stats * Get Dag statistics. * @param data The data for the request. + * @param data.dagId * @param data.dagIds * @returns DagStatsCollectionResponse Successful Response * @throws ApiError @@ -1102,16 +1103,18 @@ export const useDagStatsServiceGetDagStatsSuspense = < TQueryKey extends Array = unknown[], >( { + dagId, dagIds, }: { + dagId?: string; dagIds?: string[]; } = {}, queryKey?: TQueryKey, options?: Omit, "queryKey" | "queryFn">, ) => useSuspenseQuery({ - queryKey: Common.UseDagStatsServiceGetDagStatsKeyFn({ dagIds }, queryKey), - queryFn: () => DagStatsService.getDagStats({ dagIds }) as TData, + queryKey: Common.UseDagStatsServiceGetDagStatsKeyFn({ dagId, dagIds }, queryKey), + queryFn: () => DagStatsService.getDagStats({ dagId, dagIds }) as TData, ...options, }); /** diff --git a/airflow/ui/openapi-gen/requests/services.gen.ts b/airflow/ui/openapi-gen/requests/services.gen.ts index d7f07322284fb..b1714e823174a 100644 --- a/airflow/ui/openapi-gen/requests/services.gen.ts +++ b/airflow/ui/openapi-gen/requests/services.gen.ts @@ -1581,6 +1581,7 @@ export class DagStatsService { * Get Dag Stats * Get Dag statistics. * @param data The data for the request. + * @param data.dagId * @param data.dagIds * @returns DagStatsCollectionResponse Successful Response * @throws ApiError @@ -1590,6 +1591,7 @@ export class DagStatsService { method: "GET", url: "/public/dagStats", query: { + dag_id: data.dagId, dag_ids: data.dagIds, }, errors: { diff --git a/airflow/ui/openapi-gen/requests/types.gen.ts b/airflow/ui/openapi-gen/requests/types.gen.ts index bf9a05f25fa34..8d121e7ade0f5 100644 --- a/airflow/ui/openapi-gen/requests/types.gen.ts +++ b/airflow/ui/openapi-gen/requests/types.gen.ts @@ -1979,6 +1979,7 @@ export type GetDagSourceData = { export type GetDagSourceResponse = DAGSourceResponse; export type GetDagStatsData = { + dagId?: string | null; dagIds?: Array; }; diff --git a/tests/api_fastapi/core_api/routes/public/test_dag_stats.py b/tests/api_fastapi/core_api/routes/public/test_dag_stats.py index 8a0dae3604c1d..e7264bd48d145 100644 --- a/tests/api_fastapi/core_api/routes/public/test_dag_stats.py +++ b/tests/api_fastapi/core_api/routes/public/test_dag_stats.py @@ -129,7 +129,7 @@ def teardown_method(self) -> None: class TestGetDagStats(TestDagStatsEndpoint): """Unit tests for Get DAG Stats.""" - def test_should_respond_200(self, client, session): + def test_should_respond_200(self, test_client, session): self._create_dag_and_runs(session) exp_payload = { "dags": [ @@ -179,13 +179,21 @@ def test_should_respond_200(self, client, session): "total_entries": 2, } - response = client().get(f"{API_PREFIX}?dag_ids={DAG1_ID}&dag_ids={DAG2_ID}") + response = test_client.get(f"{API_PREFIX}?dag_ids={DAG1_ID}&dag_ids={DAG2_ID}") assert response.status_code == 200 res_json = response.json() assert res_json["total_entries"] == len(res_json["dags"]) assert res_json == exp_payload - def test_all_dags_should_respond_200(self, client, session): + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.get(f"{API_PREFIX}?dag_ids={DAG1_ID}&dag_ids={DAG2_ID}") + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.get(f"{API_PREFIX}?dag_ids={DAG1_ID}&dag_ids={DAG2_ID}") + assert response.status_code == 403 + + def test_all_dags_should_respond_200(self, test_client, session): self._create_dag_and_runs(session) exp_payload = { "dags": [ @@ -256,7 +264,7 @@ def test_all_dags_should_respond_200(self, client, session): "total_entries": 3, } - response = client().get(API_PREFIX) + response = test_client.get(API_PREFIX) assert response.status_code == 200 res_json = response.json() assert res_json["total_entries"] == len(res_json["dags"]) @@ -403,9 +411,9 @@ def test_all_dags_should_respond_200(self, client, session): ), ], ) - def test_single_dag_in_dag_ids(self, client, session, url, params, exp_payload): + def test_single_dag_in_dag_ids(self, test_client, session, url, params, exp_payload): self._create_dag_and_runs(session) - response = client().get(url, params=params) + response = test_client.get(url, params=params) assert response.status_code == 200 res_json = response.json() assert res_json["total_entries"] == len(res_json["dags"]) From 92b519b86fd37ec815aa861dd3c7ff44b5644613 Mon Sep 17 00:00:00 2001 From: kalyanr Date: Sat, 1 Mar 2025 23:18:08 +0530 Subject: [PATCH 05/27] add dag_warning --- airflow/api_fastapi/core_api/openapi/v1-generated.yaml | 2 ++ airflow/api_fastapi/core_api/routes/public/dag_warning.py | 3 +++ .../core_api/routes/public/test_dag_warning.py | 8 ++++++++ 3 files changed, 13 insertions(+) diff --git a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml index 6735f60a0e38c..f5fcdc1d5f8bb 100644 --- a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml +++ b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml @@ -3038,6 +3038,8 @@ paths: summary: List Dag Warnings description: Get a list of DAG warnings. operationId: list_dag_warnings + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: query diff --git a/airflow/api_fastapi/core_api/routes/public/dag_warning.py b/airflow/api_fastapi/core_api/routes/public/dag_warning.py index 2c964efce7e45..69b3354a26783 100644 --- a/airflow/api_fastapi/core_api/routes/public/dag_warning.py +++ b/airflow/api_fastapi/core_api/routes/public/dag_warning.py @@ -22,6 +22,7 @@ from fastapi import Depends from sqlalchemy import select +from airflow.api_fastapi.auth.managers.models.resource_details import DagAccessEntity from airflow.api_fastapi.common.db.common import ( SessionDep, paginated_select, @@ -37,6 +38,7 @@ from airflow.api_fastapi.core_api.datamodels.dag_warning import ( DAGWarningCollectionResponse, ) +from airflow.api_fastapi.core_api.security import requires_access_dag from airflow.models.dagwarning import DagWarning, DagWarningType dag_warning_router = AirflowRouter(tags=["DagWarning"]) @@ -44,6 +46,7 @@ @dag_warning_router.get( "/dagWarnings", + dependencies=[Depends(requires_access_dag(method="GET", access_entity=DagAccessEntity.WARNING))], ) def list_dag_warnings( dag_id: Annotated[FilterParam[str | None], Depends(filter_param_factory(DagWarning.dag_id, str | None))], diff --git a/tests/api_fastapi/core_api/routes/public/test_dag_warning.py b/tests/api_fastapi/core_api/routes/public/test_dag_warning.py index 61237bd10299a..c3dfc304cd8a3 100644 --- a/tests/api_fastapi/core_api/routes/public/test_dag_warning.py +++ b/tests/api_fastapi/core_api/routes/public/test_dag_warning.py @@ -78,6 +78,14 @@ def test_get_dag_warnings(self, test_client, query_params, expected_total_entrie assert len(response_json["dag_warnings"]) == len(expected_messages) assert [dag_warning["message"] for dag_warning in response_json["dag_warnings"]] == expected_messages + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.get("/public/dagWarnings", params={}) + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.get("/public/dagWarnings", params={}) + assert response.status_code == 403 + def test_get_dag_warnings_bad_request(self, test_client): response = test_client.get("/public/dagWarnings", params={"warning_type": "invalid"}) response_json = response.json() From a067840d8b31445b3afefaa1fb5f745663939dd8 Mon Sep 17 00:00:00 2001 From: kalyanr Date: Sun, 2 Mar 2025 00:07:55 +0530 Subject: [PATCH 06/27] add ti --- .../core_api/openapi/v1-generated.yaml | 97 ++++++++++-- .../core_api/routes/public/task_instances.py | 18 +++ airflow/ui/openapi-gen/queries/queries.ts | 4 +- airflow/ui/openapi-gen/requests/types.gen.ts | 32 ++-- .../routes/public/test_task_instances.py | 146 ++++++++++++++++++ 5 files changed, 262 insertions(+), 35 deletions(-) diff --git a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml index f5fcdc1d5f8bb..9255839be0c80 100644 --- a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml +++ b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml @@ -4895,12 +4895,16 @@ paths: summary: Get Task Instance description: Get task instance. operationId: get_task_instance + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id - name: dag_run_id in: path @@ -4951,12 +4955,16 @@ paths: summary: Patch Task Instance description: Update a task instance. operationId: patch_task_instance + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id - name: dag_run_id in: path @@ -5043,12 +5051,16 @@ paths: summary: Get Mapped Task Instances description: Get list of mapped task instances. operationId: get_mapped_task_instances + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id - name: dag_run_id in: path @@ -5269,12 +5281,16 @@ paths: summary: Get Task Instance Dependencies description: Get dependencies blocking task from getting scheduled. operationId: get_task_instance_dependencies + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id - name: dag_run_id in: path @@ -5332,12 +5348,16 @@ paths: summary: Get Task Instance Dependencies description: Get dependencies blocking task from getting scheduled. operationId: get_task_instance_dependencies + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id - name: dag_run_id in: path @@ -5396,12 +5416,16 @@ paths: summary: Get Task Instance Tries description: Get list of task instances history. operationId: get_task_instance_tries + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id - name: dag_run_id in: path @@ -5459,12 +5483,16 @@ paths: - Task Instance summary: Get Mapped Task Instance Tries operationId: get_mapped_task_instance_tries + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id - name: dag_run_id in: path @@ -5522,12 +5550,16 @@ paths: summary: Get Mapped Task Instance description: Get task instance. operationId: get_mapped_task_instance + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id - name: dag_run_id in: path @@ -5584,12 +5616,16 @@ paths: summary: Patch Task Instance description: Update a task instance. operationId: patch_task_instance + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id - name: dag_run_id in: path @@ -5681,12 +5717,16 @@ paths: and DAG runs.' operationId: get_task_instances + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id - name: dag_run_id in: path @@ -5917,13 +5957,16 @@ paths: summary: Get Task Instances Batch description: Get list of task instances. operationId: get_task_instances_batch + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - const: '~' - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id - name: dag_run_id in: path @@ -5976,12 +6019,16 @@ paths: summary: Get Task Instance Try Details description: Get task instance details by try number. operationId: get_task_instance_try_details + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id - name: dag_run_id in: path @@ -6045,12 +6092,16 @@ paths: - Task Instance summary: Get Mapped Task Instance Try Details operationId: get_mapped_task_instance_try_details + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id - name: dag_run_id in: path @@ -6114,12 +6165,16 @@ paths: summary: Post Clear Task Instances description: Clear task instances. operationId: post_clear_task_instances + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id requestBody: required: true @@ -6165,12 +6220,16 @@ paths: summary: Patch Task Instance Dry Run description: Update a task instance dry_run mode. operationId: patch_task_instance_dry_run + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id - name: dag_run_id in: path @@ -6250,12 +6309,16 @@ paths: summary: Patch Task Instance Dry Run description: Update a task instance dry_run mode. operationId: patch_task_instance_dry_run + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id - name: dag_run_id in: path 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 a43c3c7078083..1642a605142a1 100644 --- a/airflow/api_fastapi/core_api/routes/public/task_instances.py +++ b/airflow/api_fastapi/core_api/routes/public/task_instances.py @@ -27,6 +27,7 @@ from sqlalchemy.orm import joinedload from sqlalchemy.sql.selectable import Select +from airflow.api_fastapi.auth.managers.models.resource_details import DagAccessEntity from airflow.api_fastapi.common.db.common import SessionDep, paginated_select from airflow.api_fastapi.common.parameters import ( FilterOptionEnum, @@ -60,6 +61,7 @@ TaskInstancesBatchBody, ) from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc +from airflow.api_fastapi.core_api.security import requires_access_dag from airflow.exceptions import TaskNotFound from airflow.models import Base, DagRun from airflow.models.dag import DAG @@ -77,6 +79,7 @@ @task_instances_router.get( task_instances_prefix + "/{task_id}", responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]), + dependencies=[Depends(requires_access_dag(method="GET", access_entity=DagAccessEntity.TASK_INSTANCE))], ) def get_task_instance( dag_id: str, dag_run_id: str, task_id: str, session: SessionDep @@ -107,6 +110,7 @@ def get_task_instance( @task_instances_router.get( task_instances_prefix + "/{task_id}/listMapped", responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]), + dependencies=[Depends(requires_access_dag(method="GET", access_entity=DagAccessEntity.TASK_INSTANCE))], ) def get_mapped_task_instances( dag_id: str, @@ -210,10 +214,12 @@ def get_mapped_task_instances( @task_instances_router.get( task_instances_prefix + "/{task_id}/dependencies", responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]), + dependencies=[Depends(requires_access_dag(method="GET", access_entity=DagAccessEntity.TASK_INSTANCE))], ) @task_instances_router.get( task_instances_prefix + "/{task_id}/{map_index}/dependencies", responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]), + dependencies=[Depends(requires_access_dag(method="GET", access_entity=DagAccessEntity.TASK_INSTANCE))], ) def get_task_instance_dependencies( dag_id: str, @@ -264,6 +270,7 @@ def get_task_instance_dependencies( @task_instances_router.get( task_instances_prefix + "/{task_id}/tries", responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]), + dependencies=[Depends(requires_access_dag(method="GET", access_entity=DagAccessEntity.TASK_INSTANCE))], ) def get_task_instance_tries( dag_id: str, @@ -307,6 +314,7 @@ 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]), + dependencies=[Depends(requires_access_dag(method="GET", access_entity=DagAccessEntity.TASK_INSTANCE))], ) def get_mapped_task_instance_tries( dag_id: str, @@ -327,6 +335,7 @@ def get_mapped_task_instance_tries( @task_instances_router.get( task_instances_prefix + "/{task_id}/{map_index}", responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]), + dependencies=[Depends(requires_access_dag(method="GET", access_entity=DagAccessEntity.TASK_INSTANCE))], ) def get_mapped_task_instance( dag_id: str, @@ -357,6 +366,7 @@ def get_mapped_task_instance( @task_instances_router.get( task_instances_prefix, responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]), + dependencies=[Depends(requires_access_dag(method="GET", access_entity=DagAccessEntity.TASK_INSTANCE))], ) def get_task_instances( dag_id: str, @@ -463,6 +473,7 @@ def get_task_instances( @task_instances_router.post( task_instances_prefix + "/list", responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]), + dependencies=[Depends(requires_access_dag(method="GET", access_entity=DagAccessEntity.TASK_INSTANCE))], ) def get_task_instances_batch( dag_id: Literal["~"], @@ -544,6 +555,7 @@ def get_task_instances_batch( @task_instances_router.get( task_instances_prefix + "/{task_id}/tries/{task_try_number}", responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]), + dependencies=[Depends(requires_access_dag(method="GET", access_entity=DagAccessEntity.TASK_INSTANCE))], ) def get_task_instance_try_details( dag_id: str, @@ -579,6 +591,7 @@ def _query(orm_object: Base) -> TI | TIH | None: @task_instances_router.get( task_instances_prefix + "/{task_id}/{map_index}/tries/{task_try_number}", responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]), + dependencies=[Depends(requires_access_dag(method="GET", access_entity=DagAccessEntity.TASK_INSTANCE))], ) def get_mapped_task_instance_try_details( dag_id: str, @@ -601,6 +614,7 @@ def get_mapped_task_instance_try_details( @task_instances_router.post( "/clearTaskInstances", responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]), + dependencies=[Depends(requires_access_dag(method="PUT", access_entity=DagAccessEntity.TASK_INSTANCE))], ) def post_clear_task_instances( dag_id: str, @@ -745,12 +759,14 @@ def _patch_ti_validate_request( responses=create_openapi_http_exception_doc( [status.HTTP_404_NOT_FOUND, status.HTTP_400_BAD_REQUEST], ), + dependencies=[Depends(requires_access_dag(method="PUT", access_entity=DagAccessEntity.TASK_INSTANCE))], ) @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], ), + dependencies=[Depends(requires_access_dag(method="PUT", access_entity=DagAccessEntity.TASK_INSTANCE))], ) def patch_task_instance_dry_run( dag_id: str, @@ -805,12 +821,14 @@ def patch_task_instance_dry_run( responses=create_openapi_http_exception_doc( [status.HTTP_404_NOT_FOUND, status.HTTP_400_BAD_REQUEST, status.HTTP_409_CONFLICT], ), + dependencies=[Depends(requires_access_dag(method="PUT", access_entity=DagAccessEntity.TASK_INSTANCE))], ) @task_instances_router.patch( task_instances_prefix + "/{task_id}/{map_index}", responses=create_openapi_http_exception_doc( [status.HTTP_404_NOT_FOUND, status.HTTP_400_BAD_REQUEST, status.HTTP_409_CONFLICT], ), + dependencies=[Depends(requires_access_dag(method="PUT", access_entity=DagAccessEntity.TASK_INSTANCE))], ) def patch_task_instance( dag_id: str, diff --git a/airflow/ui/openapi-gen/queries/queries.ts b/airflow/ui/openapi-gen/queries/queries.ts index 44df779c7e744..f2a087b183823 100644 --- a/airflow/ui/openapi-gen/queries/queries.ts +++ b/airflow/ui/openapi-gen/queries/queries.ts @@ -3303,7 +3303,7 @@ export const useTaskInstanceServiceGetTaskInstancesBatch = < TData, TError, { - dagId: "~"; + dagId: string; dagRunId: "~"; requestBody: TaskInstancesBatchBody; }, @@ -3316,7 +3316,7 @@ export const useTaskInstanceServiceGetTaskInstancesBatch = < TData, TError, { - dagId: "~"; + dagId: string; dagRunId: "~"; requestBody: TaskInstancesBatchBody; }, diff --git a/airflow/ui/openapi-gen/requests/types.gen.ts b/airflow/ui/openapi-gen/requests/types.gen.ts index 8d121e7ade0f5..ba812d953d65e 100644 --- a/airflow/ui/openapi-gen/requests/types.gen.ts +++ b/airflow/ui/openapi-gen/requests/types.gen.ts @@ -2109,7 +2109,7 @@ export type GetExtraLinksData = { export type GetExtraLinksResponse = ExtraLinksResponse; export type GetTaskInstanceData = { - dagId: string; + dagId: string | null; dagRunId: string; taskId: string; }; @@ -2117,7 +2117,7 @@ export type GetTaskInstanceData = { export type GetTaskInstanceResponse = TaskInstanceResponse; export type PatchTaskInstanceData = { - dagId: string; + dagId: string | null; dagRunId: string; mapIndex?: number; requestBody: PatchTaskInstanceBody; @@ -2128,7 +2128,7 @@ export type PatchTaskInstanceData = { export type PatchTaskInstanceResponse = TaskInstanceResponse; export type GetMappedTaskInstancesData = { - dagId: string; + dagId: string | null; dagRunId: string; durationGte?: number | null; durationLte?: number | null; @@ -2156,7 +2156,7 @@ export type GetMappedTaskInstancesData = { export type GetMappedTaskInstancesResponse = TaskInstanceCollectionResponse; export type GetTaskInstanceDependenciesData = { - dagId: string; + dagId: string | null; dagRunId: string; mapIndex: number; taskId: string; @@ -2165,7 +2165,7 @@ export type GetTaskInstanceDependenciesData = { export type GetTaskInstanceDependenciesResponse = TaskDependencyCollectionResponse; export type GetTaskInstanceDependencies1Data = { - dagId: string; + dagId: string | null; dagRunId: string; mapIndex?: number; taskId: string; @@ -2174,7 +2174,7 @@ export type GetTaskInstanceDependencies1Data = { export type GetTaskInstanceDependencies1Response = TaskDependencyCollectionResponse; export type GetTaskInstanceTriesData = { - dagId: string; + dagId: string | null; dagRunId: string; mapIndex?: number; taskId: string; @@ -2183,7 +2183,7 @@ export type GetTaskInstanceTriesData = { export type GetTaskInstanceTriesResponse = TaskInstanceHistoryCollectionResponse; export type GetMappedTaskInstanceTriesData = { - dagId: string; + dagId: string | null; dagRunId: string; mapIndex: number; taskId: string; @@ -2192,7 +2192,7 @@ export type GetMappedTaskInstanceTriesData = { export type GetMappedTaskInstanceTriesResponse = TaskInstanceHistoryCollectionResponse; export type GetMappedTaskInstanceData = { - dagId: string; + dagId: string | null; dagRunId: string; mapIndex: number; taskId: string; @@ -2201,7 +2201,7 @@ export type GetMappedTaskInstanceData = { export type GetMappedTaskInstanceResponse = TaskInstanceResponse; export type PatchTaskInstance1Data = { - dagId: string; + dagId: string | null; dagRunId: string; mapIndex: number; requestBody: PatchTaskInstanceBody; @@ -2212,7 +2212,7 @@ export type PatchTaskInstance1Data = { export type PatchTaskInstance1Response = TaskInstanceResponse; export type GetTaskInstancesData = { - dagId: string; + dagId: string | null; dagRunId: string; durationGte?: number | null; durationLte?: number | null; @@ -2241,7 +2241,7 @@ export type GetTaskInstancesData = { export type GetTaskInstancesResponse = TaskInstanceCollectionResponse; export type GetTaskInstancesBatchData = { - dagId: "~"; + dagId: string | null; dagRunId: "~"; requestBody: TaskInstancesBatchBody; }; @@ -2249,7 +2249,7 @@ export type GetTaskInstancesBatchData = { export type GetTaskInstancesBatchResponse = TaskInstanceCollectionResponse; export type GetTaskInstanceTryDetailsData = { - dagId: string; + dagId: string | null; dagRunId: string; mapIndex?: number; taskId: string; @@ -2259,7 +2259,7 @@ export type GetTaskInstanceTryDetailsData = { export type GetTaskInstanceTryDetailsResponse = TaskInstanceHistoryResponse; export type GetMappedTaskInstanceTryDetailsData = { - dagId: string; + dagId: string | null; dagRunId: string; mapIndex: number; taskId: string; @@ -2269,14 +2269,14 @@ export type GetMappedTaskInstanceTryDetailsData = { export type GetMappedTaskInstanceTryDetailsResponse = TaskInstanceHistoryResponse; export type PostClearTaskInstancesData = { - dagId: string; + dagId: string | null; requestBody: ClearTaskInstancesBody; }; export type PostClearTaskInstancesResponse = TaskInstanceCollectionResponse; export type PatchTaskInstanceDryRunData = { - dagId: string; + dagId: string | null; dagRunId: string; mapIndex: number; requestBody: PatchTaskInstanceBody; @@ -2287,7 +2287,7 @@ export type PatchTaskInstanceDryRunData = { export type PatchTaskInstanceDryRunResponse = TaskInstanceCollectionResponse; export type PatchTaskInstanceDryRun1Data = { - dagId: string; + dagId: string | null; dagRunId: string; mapIndex?: number; requestBody: PatchTaskInstanceBody; 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 176709c321f90..465c26278cdec 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 @@ -206,6 +206,18 @@ def test_should_respond_200(self, test_client, session): "triggerer_job": None, } + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.get( + "/public/dags/example_python_operator/dagRuns/TEST_DAG_RUN_ID/taskInstances/print_the_context" + ) + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.get( + "/public/dags/example_python_operator/dagRuns/TEST_DAG_RUN_ID/taskInstances/print_the_context" + ) + assert response.status_code == 403 + @pytest.mark.parametrize( "run_id, expected_version_number", [ @@ -524,6 +536,18 @@ def test_should_respond_200_mapped_task_instance_with_rtif(self, test_client, se "triggerer_job": None, } + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.get( + "/public/dags/example_python_operator/dagRuns/TEST_DAG_RUN_ID/taskInstances/print_the_context/1", + ) + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.get( + "/public/dags/example_python_operator/dagRuns/TEST_DAG_RUN_ID/taskInstances/print_the_context/1", + ) + assert response.status_code == 403 + def test_should_respond_404_wrong_map_index(self, test_client, session): self.create_task_instances(session) @@ -664,6 +688,18 @@ def one_task_with_zero_mapped_tis(self, dag_maker, session): }, ) + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.get( + "/public/dags/mapped_tis/dagRuns/run_mapped_tis/taskInstances/task_2/listMapped", + ) + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.get( + "/public/dags/mapped_tis/dagRuns/run_mapped_tis/taskInstances/task_2/listMapped", + ) + assert response.status_code == 403 + def test_should_respond_404(self, test_client): response = test_client.get( "/public/dags/mapped_tis/dagRuns/run_mapped_tis/taskInstances/task_2/listMapped", @@ -1067,6 +1103,18 @@ def test_should_respond_200( assert response.json()["total_entries"] == expected_ti assert len(response.json()["task_instances"]) == expected_ti + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.get( + "/public/dags/example_python_operator/dagRuns/~/taskInstances", + ) + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.get( + "/public/dags/example_python_operator/dagRuns/~/taskInstances", + ) + assert response.status_code == 403 + def test_not_found(self, test_client): response = test_client.get("/public/dags/invalid/dagRuns/~/taskInstances") assert response.status_code == 404 @@ -1293,6 +1341,20 @@ def test_should_respond_dependencies_mapped(self, test_client, session): ) assert response.status_code == 200, response.text + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.get( + "/public/dags/example_python_operator/dagRuns/TEST_DAG_RUN_ID/taskInstances/" + "print_the_context/0/dependencies", + ) + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.get( + "/public/dags/example_python_operator/dagRuns/TEST_DAG_RUN_ID/taskInstances/" + "print_the_context/0/dependencies", + ) + assert response.status_code == 403 + class TestGetTaskInstancesBatch(TestTaskInstanceEndpoint): @pytest.mark.parametrize( @@ -1506,6 +1568,20 @@ def test_should_raise_400_for_no_json(self, test_client): }, ] + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.post( + "/public/dags/~/dagRuns/~/taskInstances/list", + json={}, + ) + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.post( + "/public/dags/~/dagRuns/~/taskInstances/list", + json={}, + ) + assert response.status_code == 403 + def test_should_respond_422_for_non_wildcard_path_parameters(self, test_client): response = test_client.post( "/public/dags/non_wildcard/dagRuns/~/taskInstances/list", @@ -1818,6 +1894,18 @@ def test_should_respond_200_with_task_state_in_removed(self, test_client, sessio "dag_version": None, } + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.get( + "/public/dags/example_python_operator/dagRuns/TEST_DAG_RUN_ID/taskInstances/print_the_context/tries/1", + ) + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.get( + "/public/dags/example_python_operator/dagRuns/TEST_DAG_RUN_ID/taskInstances/print_the_context/tries/1", + ) + assert response.status_code == 403 + def test_raises_404_for_nonexistent_task_instance(self, test_client, session): self.create_task_instances(session) response = test_client.get( @@ -2087,6 +2175,20 @@ def test_dag_run_with_future_or_past_flag_returns_400(self, test_client, session in response.json()["detail"] ) + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.post( + "/public/dags/dag_id/clearTaskInstances", + json={}, + ) + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.post( + "/public/dags/dag_id/clearTaskInstances", + json={}, + ) + assert response.status_code == 403 + @pytest.mark.parametrize( "main_dag, task_instances, request_dag, payload, expected_ti", [ @@ -2711,6 +2813,18 @@ def test_should_respond_200(self, test_client, session): "total_entries": 2, } + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.get( + "/public/dags/example_python_operator/dagRuns/TEST_DAG_RUN_ID/taskInstances/print_the_context/tries" + ) + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.get( + "/public/dags/example_python_operator/dagRuns/TEST_DAG_RUN_ID/taskInstances/print_the_context/tries" + ) + assert response.status_code == 403 + def test_ti_in_retry_state_not_returned(self, test_client, session): self.create_task_instances( session=session, task_instances=[{"state": State.SUCCESS}], with_ti_history=True @@ -3025,6 +3139,24 @@ def test_should_update_mapped_task_instance_state(self, test_client, session): assert response2.status_code == 200 assert response2.json()["state"] == self.NEW_STATE + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.patch( + self.ENDPOINT_URL, + json={ + "new_state": self.NEW_STATE, + }, + ) + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.patch( + self.ENDPOINT_URL, + json={ + "new_state": self.NEW_STATE, + }, + ) + assert response.status_code == 403 + @pytest.mark.parametrize( "error, code, payload", [ @@ -3525,6 +3657,20 @@ def test_should_not_update(self, test_client, session, payload): assert task_before == task_after + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.patch( + f"{self.ENDPOINT_URL}/dry_run", + json={}, + ) + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.patch( + f"{self.ENDPOINT_URL}/dry_run", + json={}, + ) + assert response.status_code == 403 + def test_should_not_update_mapped_task_instance(self, test_client, session): map_index = 1 tis = self.create_task_instances(session) From 62c17c0e60ce5132e177abf3a32269250a053ccd Mon Sep 17 00:00:00 2001 From: kalyanr Date: Sun, 2 Mar 2025 00:38:02 +0530 Subject: [PATCH 07/27] add xcom --- .../core_api/openapi/v1-generated.yaml | 24 +++++++-- .../core_api/routes/public/xcom.py | 8 ++- airflow/ui/openapi-gen/requests/types.gen.ts | 8 +-- .../core_api/routes/public/test_xcom.py | 54 +++++++++++++++++++ 4 files changed, 85 insertions(+), 9 deletions(-) diff --git a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml index 9255839be0c80..7122cf8120e43 100644 --- a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml +++ b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml @@ -4559,12 +4559,16 @@ paths: summary: Get Xcom Entry description: Get an XCom entry. operationId: get_xcom_entry + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id - name: task_id in: path @@ -4652,12 +4656,16 @@ paths: summary: Update Xcom Entry description: Update an existing XCom entry. operationId: update_xcom_entry + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id - name: task_id in: path @@ -4731,12 +4739,16 @@ paths: This endpoint allows specifying `~` as the dag_id, dag_run_id, task_id to retrieve XCom entries for all DAGs.' operationId: get_xcom_entries + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id - name: dag_run_id in: path @@ -4826,12 +4838,16 @@ paths: summary: Create Xcom Entry description: Create an XCom entry. operationId: create_xcom_entry + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id - name: task_id in: path diff --git a/airflow/api_fastapi/core_api/routes/public/xcom.py b/airflow/api_fastapi/core_api/routes/public/xcom.py index 3da163f3e4033..572e7482540e2 100644 --- a/airflow/api_fastapi/core_api/routes/public/xcom.py +++ b/airflow/api_fastapi/core_api/routes/public/xcom.py @@ -19,9 +19,10 @@ import copy from typing import Annotated -from fastapi import HTTPException, Query, Request, status +from fastapi import Depends, HTTPException, Query, Request, status from sqlalchemy import and_, select +from airflow.api_fastapi.auth.managers.models.resource_details import DagAccessEntity from airflow.api_fastapi.common.db.common import SessionDep, paginated_select from airflow.api_fastapi.common.parameters import QueryLimit, QueryOffset from airflow.api_fastapi.common.router import AirflowRouter @@ -33,6 +34,7 @@ XComUpdateBody, ) from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc +from airflow.api_fastapi.core_api.security import requires_access_dag from airflow.exceptions import TaskNotFound from airflow.models import DAG, DagRun as DR, XCom from airflow.settings import conf @@ -50,6 +52,7 @@ status.HTTP_404_NOT_FOUND, ] ), + dependencies=[Depends(requires_access_dag(method="GET", access_entity=DagAccessEntity.XCOM))], ) def get_xcom_entry( dag_id: str, @@ -105,6 +108,7 @@ def get_xcom_entry( status.HTTP_404_NOT_FOUND, ] ), + dependencies=[Depends(requires_access_dag(method="GET", access_entity=DagAccessEntity.XCOM))], ) def get_xcom_entries( dag_id: str, @@ -155,6 +159,7 @@ def get_xcom_entries( status.HTTP_404_NOT_FOUND, ] ), + dependencies=[Depends(requires_access_dag(method="POST", access_entity=DagAccessEntity.XCOM))], ) def create_xcom_entry( dag_id: str, @@ -234,6 +239,7 @@ def create_xcom_entry( status.HTTP_404_NOT_FOUND, ] ), + dependencies=[Depends(requires_access_dag(method="PUT", access_entity=DagAccessEntity.XCOM))], ) def update_xcom_entry( dag_id: str, diff --git a/airflow/ui/openapi-gen/requests/types.gen.ts b/airflow/ui/openapi-gen/requests/types.gen.ts index ba812d953d65e..8a7d35f239cb2 100644 --- a/airflow/ui/openapi-gen/requests/types.gen.ts +++ b/airflow/ui/openapi-gen/requests/types.gen.ts @@ -2397,7 +2397,7 @@ export type GetProvidersData = { export type GetProvidersResponse = ProviderCollectionResponse; export type GetXcomEntryData = { - dagId: string; + dagId: string | null; dagRunId: string; deserialize?: boolean; mapIndex?: number; @@ -2409,7 +2409,7 @@ export type GetXcomEntryData = { export type GetXcomEntryResponse = XComResponseNative | XComResponseString; export type UpdateXcomEntryData = { - dagId: string; + dagId: string | null; dagRunId: string; requestBody: XComUpdateBody; taskId: string; @@ -2419,7 +2419,7 @@ export type UpdateXcomEntryData = { export type UpdateXcomEntryResponse = XComResponseNative; export type GetXcomEntriesData = { - dagId: string; + dagId: string | null; dagRunId: string; limit?: number; mapIndex?: number | null; @@ -2431,7 +2431,7 @@ export type GetXcomEntriesData = { export type GetXcomEntriesResponse = XComCollectionResponse; export type CreateXcomEntryData = { - dagId: string; + dagId: string | null; dagRunId: string; requestBody: XComCreateBody; taskId: string; diff --git a/tests/api_fastapi/core_api/routes/public/test_xcom.py b/tests/api_fastapi/core_api/routes/public/test_xcom.py index dd5a073c1baae..c665a5299ced3 100644 --- a/tests/api_fastapi/core_api/routes/public/test_xcom.py +++ b/tests/api_fastapi/core_api/routes/public/test_xcom.py @@ -149,6 +149,18 @@ def test_should_respond_200_native(self, test_client): "value": TEST_XCOM_VALUE, } + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.get( + f"/public/dags/{TEST_DAG_ID}/dagRuns/{run_id}/taskInstances/{TEST_TASK_ID}/xcomEntries/{TEST_XCOM_KEY}" + ) + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.get( + f"/public/dags/{TEST_DAG_ID}/dagRuns/{run_id}/taskInstances/{TEST_TASK_ID}/xcomEntries/{TEST_XCOM_KEY}" + ) + assert response.status_code == 403 + def test_should_raise_404_for_non_existent_xcom(self, test_client): response = test_client.get( f"/public/dags/{TEST_DAG_ID}/dagRuns/{run_id}/taskInstances/{TEST_TASK_ID}/xcomEntries/{TEST_XCOM_KEY_2}" @@ -437,6 +449,20 @@ def _create_xcom_entries(self, dag_id, run_id, logical_date, task_id, mapped_ti= map_index=map_index, ) + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.get( + "/public/dags/~/dagRuns/~/taskInstances/~/xcomEntries", + params={}, + ) + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.get( + "/public/dags/~/dagRuns/~/taskInstances/~/xcomEntries", + params={}, + ) + assert response.status_code == 403 + class TestPaginationGetXComEntries(TestXComEndpoint): @pytest.mark.parametrize( @@ -583,6 +609,20 @@ def test_create_xcom_entry( assert current_data["run_id"] == dag_run_id assert current_data["map_index"] == request_body.map_index + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.post( + "/public/dags/dag_id/dagRuns/dag_run_id/taskInstances/task_id/xcomEntries", + json={}, + ) + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.post( + "/public/dags/dag_id/dagRuns/dag_run_id/taskInstances/task_id/xcomEntries", + json={}, + ) + assert response.status_code == 403 + class TestPatchXComEntry(TestXComEndpoint): @pytest.mark.parametrize( @@ -623,3 +663,17 @@ def test_patch_xcom_entry(self, key, patch_body, expected_status, expected_detai assert response.json()["value"] == XCom.serialize_value(new_value) else: assert response.json()["detail"] == expected_detail + + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.patch( + f"/public/dags/{TEST_DAG_ID}/dagRuns/run_id/taskInstances/TEST_TASK_ID/xcomEntries/key", + json={}, + ) + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.patch( + f"/public/dags/{TEST_DAG_ID}/dagRuns/run_id/taskInstances/TEST_TASK_ID/xcomEntries/key", + json={}, + ) + assert response.status_code == 403 From 6ed55d30bd06e996312d913b2b416bbb9f709009 Mon Sep 17 00:00:00 2001 From: kalyanr Date: Sat, 8 Mar 2025 16:10:37 +0530 Subject: [PATCH 08/27] add dag_parsing --- .../api_fastapi/core_api/openapi/v1-generated.yaml | 10 ++++++++++ .../core_api/routes/public/dag_parsing.py | 3 ++- airflow/ui/openapi-gen/queries/queries.ts | 7 +++++-- airflow/ui/openapi-gen/requests/services.gen.ts | 4 ++++ airflow/ui/openapi-gen/requests/types.gen.ts | 1 + .../core_api/routes/public/test_dag_parsing.py | 12 ++++++++++++ 6 files changed, 34 insertions(+), 3 deletions(-) diff --git a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml index d10c6969ba7da..df6ec337f40f8 100644 --- a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml +++ b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml @@ -6816,6 +6816,8 @@ paths: summary: Reparse Dag File description: Request re-parsing a DAG file. operationId: reparse_dag_file + security: + - OAuth2PasswordBearer: [] parameters: - name: file_token in: path @@ -6823,6 +6825,14 @@ paths: schema: type: string title: File Token + - name: dag_id + in: query + required: false + schema: + anyOf: + - type: string + - type: 'null' + title: Dag Id responses: '201': description: Successful Response diff --git a/airflow/api_fastapi/core_api/routes/public/dag_parsing.py b/airflow/api_fastapi/core_api/routes/public/dag_parsing.py index f5b7f2c359335..a4deb3b18de02 100644 --- a/airflow/api_fastapi/core_api/routes/public/dag_parsing.py +++ b/airflow/api_fastapi/core_api/routes/public/dag_parsing.py @@ -27,6 +27,7 @@ from airflow.api_fastapi.common.db.common import SessionDep from airflow.api_fastapi.common.router import AirflowRouter from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc +from airflow.api_fastapi.core_api.security import requires_access_dag from airflow.api_fastapi.logging.decorators import action_logging from airflow.models.dag import DagModel from airflow.models.dagbag import DagPriorityParsingRequest @@ -41,7 +42,7 @@ "", responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]), status_code=status.HTTP_201_CREATED, - dependencies=[Depends(action_logging())], + dependencies=[Depends(requires_access_dag(method="PUT")), Depends(action_logging())], ) def reparse_dag_file( file_token: str, diff --git a/airflow/ui/openapi-gen/queries/queries.ts b/airflow/ui/openapi-gen/queries/queries.ts index 02b647cdfa104..d460a56a66389 100644 --- a/airflow/ui/openapi-gen/queries/queries.ts +++ b/airflow/ui/openapi-gen/queries/queries.ts @@ -3596,6 +3596,7 @@ export const useBackfillServiceCancelBackfill = < * Request re-parsing a DAG file. * @param data The data for the request. * @param data.fileToken + * @param data.dagId * @returns null Successful Response * @throws ApiError */ @@ -3609,6 +3610,7 @@ export const useDagParsingServiceReparseDagFile = < TData, TError, { + dagId?: string; fileToken: string; }, TContext @@ -3620,12 +3622,13 @@ export const useDagParsingServiceReparseDagFile = < TData, TError, { + dagId?: string; fileToken: string; }, TContext >({ - mutationFn: ({ fileToken }) => - DagParsingService.reparseDagFile({ fileToken }) as unknown as Promise, + mutationFn: ({ dagId, fileToken }) => + DagParsingService.reparseDagFile({ dagId, fileToken }) as unknown as Promise, ...options, }); /** diff --git a/airflow/ui/openapi-gen/requests/services.gen.ts b/airflow/ui/openapi-gen/requests/services.gen.ts index 11b54997040d6..3acae741015da 100644 --- a/airflow/ui/openapi-gen/requests/services.gen.ts +++ b/airflow/ui/openapi-gen/requests/services.gen.ts @@ -3388,6 +3388,7 @@ export class DagParsingService { * Request re-parsing a DAG file. * @param data The data for the request. * @param data.fileToken + * @param data.dagId * @returns null Successful Response * @throws ApiError */ @@ -3398,6 +3399,9 @@ export class DagParsingService { path: { file_token: data.fileToken, }, + query: { + dag_id: data.dagId, + }, 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 3b7a01c260817..7d39b5115eacc 100644 --- a/airflow/ui/openapi-gen/requests/types.gen.ts +++ b/airflow/ui/openapi-gen/requests/types.gen.ts @@ -2492,6 +2492,7 @@ export type BulkVariablesData = { export type BulkVariablesResponse = BulkResponse; export type ReparseDagFileData = { + dagId?: string | null; fileToken: string; }; diff --git a/tests/api_fastapi/core_api/routes/public/test_dag_parsing.py b/tests/api_fastapi/core_api/routes/public/test_dag_parsing.py index 0943684e32192..00ab7dbbaf324 100644 --- a/tests/api_fastapi/core_api/routes/public/test_dag_parsing.py +++ b/tests/api_fastapi/core_api/routes/public/test_dag_parsing.py @@ -65,6 +65,18 @@ def test_201_and_400_requests(self, url_safe_serializer, session, test_client): assert parsing_requests[0].fileloc == test_dag.fileloc _check_last_log(session, dag_id=None, event="reparse_dag_file", logical_date=None) + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.put( + "/public/parseDagFile/token", headers={"Accept": "application/json"} + ) + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.put( + "/public/parseDagFile/token", headers={"Accept": "application/json"} + ) + assert response.status_code == 403 + def test_bad_file_request(self, url_safe_serializer, session, test_client): url = f"/public/parseDagFile/{url_safe_serializer.dumps('/some/random/file.py')}" response = test_client.put(url, headers={"Accept": "application/json"}) From 4b6b015cfe5a980c53250f0fc357ebd40027a0be Mon Sep 17 00:00:00 2001 From: kalyanr Date: Sat, 1 Mar 2025 18:09:03 +0530 Subject: [PATCH 09/27] add dag_run --- .../core_api/openapi/v1-generated.yaml | 46 ++++++++++--- .../core_api/routes/public/dag_run.py | 41 ++++++++--- airflow/ui/openapi-gen/queries/queries.ts | 8 +-- airflow/ui/openapi-gen/requests/types.gen.ts | 16 ++--- .../core_api/routes/public/test_dag_run.py | 68 +++++++++++++++++++ 5 files changed, 151 insertions(+), 28 deletions(-) diff --git a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml index df6ec337f40f8..e44e74d64b955 100644 --- a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml +++ b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml @@ -2077,12 +2077,16 @@ paths: - DagRun summary: Get Dag Run operationId: get_dag_run + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id - name: dag_run_id in: path @@ -2127,12 +2131,16 @@ paths: summary: Delete Dag Run description: Delete a DAG Run entry. operationId: delete_dag_run + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id - name: dag_run_id in: path @@ -2179,12 +2187,16 @@ paths: summary: Patch Dag Run description: Modify a DAG Run. operationId: patch_dag_run + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id - name: dag_run_id in: path @@ -2260,7 +2272,9 @@ paths: in: path required: true schema: - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id - name: dag_run_id in: path @@ -2305,12 +2319,16 @@ paths: - DagRun summary: Clear Dag Run operationId: clear_dag_run + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id - name: dag_run_id in: path @@ -2369,12 +2387,16 @@ paths: This endpoint allows specifying `~` as the dag_id to retrieve Dag Runs for all DAGs.' operationId: get_dag_runs + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id - name: limit in: query @@ -2534,11 +2556,16 @@ paths: summary: Trigger Dag Run description: Trigger a DAG. operationId: trigger_dag_run + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: + anyOf: + - type: string + - type: 'null' title: Dag Id requestBody: required: true @@ -2596,13 +2623,16 @@ paths: summary: Get List Dag Runs Batch description: Get a list of DAG Runs. operationId: get_list_dag_runs_batch + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - const: '~' - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id requestBody: required: true 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 8703fe42e423a..eff8d729b18c4 100644 --- a/airflow/api_fastapi/core_api/routes/public/dag_run.py +++ b/airflow/api_fastapi/core_api/routes/public/dag_run.py @@ -29,6 +29,7 @@ set_dag_run_state_to_queued, set_dag_run_state_to_success, ) +from airflow.api_fastapi.auth.managers.models.resource_details import DagAccessEntity from airflow.api_fastapi.common.db.common import SessionDep, paginated_select from airflow.api_fastapi.common.parameters import ( FilterOptionEnum, @@ -59,7 +60,7 @@ TaskInstanceResponse, ) from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc -from airflow.api_fastapi.core_api.security import requires_access_asset +from airflow.api_fastapi.core_api.security import requires_access_asset, requires_access_dag from airflow.api_fastapi.logging.decorators import action_logging from airflow.exceptions import ParamValidationError from airflow.listeners.listener import get_listener_manager @@ -78,6 +79,7 @@ status.HTTP_404_NOT_FOUND, ] ), + dependencies=[Depends(requires_access_dag(method="GET", access_entity=DagAccessEntity.RUN))], ) def get_dag_run(dag_id: str, dag_run_id: str, session: SessionDep) -> DAGRunResponse: dag_run = session.scalar(select(DagRun).filter_by(dag_id=dag_id, run_id=dag_run_id)) @@ -99,7 +101,10 @@ def get_dag_run(dag_id: str, dag_run_id: str, session: SessionDep) -> DAGRunResp status.HTTP_404_NOT_FOUND, ], ), - dependencies=[Depends(action_logging())], + dependencies=[ + Depends(requires_access_dag(method="DELETE", access_entity=DagAccessEntity.RUN)), + Depends(action_logging()), + ], ) def delete_dag_run(dag_id: str, dag_run_id: str, session: SessionDep): """Delete a DAG Run entry.""" @@ -121,7 +126,10 @@ def delete_dag_run(dag_id: str, dag_run_id: str, session: SessionDep): status.HTTP_404_NOT_FOUND, ], ), - dependencies=[Depends(action_logging())], + dependencies=[ + Depends(requires_access_dag(method="PUT", access_entity=DagAccessEntity.RUN)), + Depends(action_logging()), + ], ) def patch_dag_run( dag_id: str, @@ -190,7 +198,10 @@ def patch_dag_run( status.HTTP_404_NOT_FOUND, ] ), - dependencies=[Depends(requires_access_asset(method="GET"))], + dependencies=[ + Depends(requires_access_asset(method="GET")), + Depends(requires_access_dag(method="GET", access_entity=DagAccessEntity.RUN)), + ], ) def get_upstream_asset_events( dag_id: str, dag_run_id: str, session: SessionDep @@ -217,7 +228,10 @@ def get_upstream_asset_events( @dag_run_router.post( "/{dag_run_id}/clear", responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]), - dependencies=[Depends(action_logging())], + dependencies=[ + Depends(requires_access_dag(method="POST", access_entity=DagAccessEntity.RUN)), + Depends(action_logging()), + ], ) def clear_dag_run( dag_id: str, @@ -263,7 +277,11 @@ def clear_dag_run( return dag_run_cleared -@dag_run_router.get("", responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND])) +@dag_run_router.get( + "", + responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]), + dependencies=[Depends(requires_access_dag(method="GET", access_entity=DagAccessEntity.RUN))], +) def get_dag_runs( dag_id: str, limit: QueryLimit, @@ -337,7 +355,10 @@ def get_dag_runs( status.HTTP_409_CONFLICT, ] ), - dependencies=[Depends(action_logging())], + dependencies=[ + Depends(requires_access_dag(method="POST", access_entity=DagAccessEntity.RUN)), + Depends(action_logging()), + ], ) def trigger_dag_run( dag_id, @@ -383,7 +404,11 @@ def trigger_dag_run( raise HTTPException(status.HTTP_400_BAD_REQUEST, str(e)) -@dag_run_router.post("/list", responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND])) +@dag_run_router.post( + "/list", + responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]), + dependencies=[Depends(requires_access_dag(method="POST", access_entity=DagAccessEntity.RUN))], +) def get_list_dag_runs_batch( dag_id: Literal["~"], body: DAGRunsBatchBody, session: SessionDep ) -> DAGRunCollectionResponse: diff --git a/airflow/ui/openapi-gen/queries/queries.ts b/airflow/ui/openapi-gen/queries/queries.ts index d460a56a66389..f6cebbf884ebe 100644 --- a/airflow/ui/openapi-gen/queries/queries.ts +++ b/airflow/ui/openapi-gen/queries/queries.ts @@ -3216,7 +3216,7 @@ export const useDagRunServiceTriggerDagRun = < TData, TError, { - dagId: unknown; + dagId: string; requestBody: TriggerDAGRunPostBody; }, TContext @@ -3228,7 +3228,7 @@ export const useDagRunServiceTriggerDagRun = < TData, TError, { - dagId: unknown; + dagId: string; requestBody: TriggerDAGRunPostBody; }, TContext @@ -3256,7 +3256,7 @@ export const useDagRunServiceGetListDagRunsBatch = < TData, TError, { - dagId: "~"; + dagId: string; requestBody: DAGRunsBatchBody; }, TContext @@ -3268,7 +3268,7 @@ export const useDagRunServiceGetListDagRunsBatch = < TData, TError, { - dagId: "~"; + dagId: string; requestBody: DAGRunsBatchBody; }, TContext diff --git a/airflow/ui/openapi-gen/requests/types.gen.ts b/airflow/ui/openapi-gen/requests/types.gen.ts index 7d39b5115eacc..561d0ab6fb115 100644 --- a/airflow/ui/openapi-gen/requests/types.gen.ts +++ b/airflow/ui/openapi-gen/requests/types.gen.ts @@ -1898,21 +1898,21 @@ export type TestConnectionResponse = ConnectionTestResponse; export type CreateDefaultConnectionsResponse = void; export type GetDagRunData = { - dagId: string; + dagId: string | null; dagRunId: string; }; export type GetDagRunResponse = DAGRunResponse; export type DeleteDagRunData = { - dagId: string; + dagId: string | null; dagRunId: string; }; export type DeleteDagRunResponse = void; export type PatchDagRunData = { - dagId: string; + dagId: string | null; dagRunId: string; requestBody: DAGRunPatchBody; updateMask?: Array | null; @@ -1921,14 +1921,14 @@ export type PatchDagRunData = { export type PatchDagRunResponse = DAGRunResponse; export type GetUpstreamAssetEventsData = { - dagId: string; + dagId: string | null; dagRunId: string; }; export type GetUpstreamAssetEventsResponse = AssetEventCollectionResponse; export type ClearDagRunData = { - dagId: string; + dagId: string | null; dagRunId: string; requestBody: DAGRunClearBody; }; @@ -1936,7 +1936,7 @@ export type ClearDagRunData = { export type ClearDagRunResponse = TaskInstanceCollectionResponse | DAGRunResponse; export type GetDagRunsData = { - dagId: string; + dagId: string | null; endDateGte?: string | null; endDateLte?: string | null; limit?: number; @@ -1956,14 +1956,14 @@ export type GetDagRunsData = { export type GetDagRunsResponse = DAGRunCollectionResponse; export type TriggerDagRunData = { - dagId: unknown; + dagId: string | null; requestBody: TriggerDAGRunPostBody; }; export type TriggerDagRunResponse = DAGRunResponse; export type GetListDagRunsBatchData = { - dagId: "~"; + dagId: string | null; requestBody: DAGRunsBatchBody; }; 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 f62bff0665ec6..61199fc1f37a5 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 @@ -244,6 +244,14 @@ def test_get_dag_run_not_found(self, test_client): body = response.json() assert body["detail"] == "The DagRun with dag_id: `test_dag1` and run_id: `invalid` was not found" + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.get(f"/public/dags/{DAG1_ID}/dagRuns/invalid") + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.get(f"/public/dags/{DAG1_ID}/dagRuns/invalid") + assert response.status_code == 403 + class TestGetDagRuns: @pytest.mark.parametrize("dag_id, total_entries", [(DAG1_ID, 2), (DAG2_ID, 2), ("~", 4)]) @@ -277,6 +285,14 @@ def test_invalid_order_by_raises_400(self, test_client): == "Ordering with 'invalid' is disallowed or the attribute does not exist on the model" ) + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.get("/public/dags/test_dag1/dagRuns?order_by=invalid") + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.get("/public/dags/test_dag1/dagRuns?order_by=invalid") + assert response.status_code == 403 + @pytest.mark.parametrize( "order_by,expected_order", [ @@ -550,6 +566,14 @@ def test_list_dag_runs_return_200(self, test_client, session): expected = get_dag_run_dict(run) assert each == expected + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.post("/public/dags/~/dagRuns/list", json={}) + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.post("/public/dags/~/dagRuns/list", json={}) + assert response.status_code == 403 + def test_list_dag_runs_with_invalid_dag_id(self, test_client): response = test_client.post("/public/dags/invalid/dagRuns/list", json={}) assert response.status_code == 422 @@ -909,6 +933,14 @@ def test_patch_dag_run(self, test_client, dag_id, run_id, patch_body, response_b assert body.get("note") == response_body.get("note") _check_last_log(session, dag_id=dag_id, event="patch_dag_run", logical_date=None) + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.patch("/public/dags/dag_1/dagRuns/run_1", json={}) + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.patch("/public/dags/dag_1/dagRuns/run_1", json={}) + assert response.status_code == 403 + @pytest.mark.parametrize( "query_params, patch_body, response_body, expected_status_code", [ @@ -1008,6 +1040,14 @@ def test_delete_dag_run_not_found(self, test_client): body = response.json() assert body["detail"] == "The DagRun with dag_id: `test_dag1` and run_id: `invalid` was not found" + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.delete(f"/public/dags/{DAG1_ID}/dagRuns/invalid") + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.delete(f"/public/dags/{DAG1_ID}/dagRuns/invalid") + assert response.status_code == 403 + class TestGetDagRunAssetTriggerEvents: @pytest.mark.usefixtures("configure_git_connection_for_dag_bundle") @@ -1115,6 +1155,20 @@ def test_clear_dag_run(self, test_client, session): logical_date=None, ) + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.post( + f"/public/dags/{DAG1_ID}/dagRuns/{DAG1_RUN1_ID}/clear", + json={"dry_run": False}, + ) + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.post( + f"/public/dags/{DAG1_ID}/dagRuns/{DAG1_RUN1_ID}/clear", + json={"dry_run": False}, + ) + assert response.status_code == 403 + @pytest.mark.parametrize( "body, dag_run_id, expected_state", [ @@ -1246,6 +1300,20 @@ def test_should_respond_200( assert response.json() == expected_response_json _check_last_log(session, dag_id=DAG1_ID, event="trigger_dag_run", logical_date=None) + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.post( + f"/public/dags/{DAG1_ID}/dagRuns", + json={}, + ) + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.post( + f"/public/dags/{DAG1_ID}/dagRuns", + json={}, + ) + assert response.status_code == 403 + @pytest.mark.parametrize( "post_body, expected_detail", [ From 94d772330bd001194c6a2c2d922abd33b73f2228 Mon Sep 17 00:00:00 2001 From: kalyanr Date: Sat, 1 Mar 2025 18:40:05 +0530 Subject: [PATCH 10/27] add dag_source --- .../api_fastapi/core_api/openapi/v1-generated.yaml | 6 +++++- .../core_api/routes/public/dag_sources.py | 5 ++++- airflow/ui/openapi-gen/requests/types.gen.ts | 2 +- .../core_api/routes/public/test_dag_sources.py | 12 ++++++++++++ 4 files changed, 22 insertions(+), 3 deletions(-) diff --git a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml index e44e74d64b955..8a0b2989f3015 100644 --- a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml +++ b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml @@ -2678,12 +2678,16 @@ paths: summary: Get Dag Source description: Get source code using file token. operationId: get_dag_source + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id - name: version_number in: query diff --git a/airflow/api_fastapi/core_api/routes/public/dag_sources.py b/airflow/api_fastapi/core_api/routes/public/dag_sources.py index 4337c92107322..fb00a4dd553b5 100644 --- a/airflow/api_fastapi/core_api/routes/public/dag_sources.py +++ b/airflow/api_fastapi/core_api/routes/public/dag_sources.py @@ -16,14 +16,16 @@ # under the License. from __future__ import annotations -from fastapi import HTTPException, Response, status +from fastapi import Depends, HTTPException, Response, status +from airflow.api_fastapi.auth.managers.models.resource_details import DagAccessEntity from airflow.api_fastapi.common.db.common import SessionDep from airflow.api_fastapi.common.headers import HeaderAcceptJsonOrText from airflow.api_fastapi.common.router import AirflowRouter from airflow.api_fastapi.common.types import Mimetype from airflow.api_fastapi.core_api.datamodels.dag_sources import DAGSourceResponse from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc +from airflow.api_fastapi.core_api.security import requires_access_dag from airflow.models.dag_version import DagVersion dag_sources_router = AirflowRouter(tags=["DagSource"], prefix="/dagSources") @@ -47,6 +49,7 @@ }, }, response_model=DAGSourceResponse, + dependencies=[Depends(requires_access_dag(method="GET", access_entity=DagAccessEntity.CODE))], ) def get_dag_source( accept: HeaderAcceptJsonOrText, diff --git a/airflow/ui/openapi-gen/requests/types.gen.ts b/airflow/ui/openapi-gen/requests/types.gen.ts index 561d0ab6fb115..ffbcb2cbddfe3 100644 --- a/airflow/ui/openapi-gen/requests/types.gen.ts +++ b/airflow/ui/openapi-gen/requests/types.gen.ts @@ -1971,7 +1971,7 @@ export type GetListDagRunsBatchResponse = DAGRunCollectionResponse; export type GetDagSourceData = { accept?: "application/json" | "text/plain" | "*/*"; - dagId: string; + dagId: string | null; versionNumber?: number | null; }; diff --git a/tests/api_fastapi/core_api/routes/public/test_dag_sources.py b/tests/api_fastapi/core_api/routes/public/test_dag_sources.py index 6da8902b232a0..8bd53ea2f97b9 100644 --- a/tests/api_fastapi/core_api/routes/public/test_dag_sources.py +++ b/tests/api_fastapi/core_api/routes/public/test_dag_sources.py @@ -77,6 +77,18 @@ def test_should_respond_200_text(self, test_client, test_dag): json.loads(response.content.decode()) assert response.headers["Content-Type"].startswith("text/plain") + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.get( + f"{API_PREFIX}/{TEST_DAG_ID}", headers={"Accept": "text/plain"} + ) + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.get( + f"{API_PREFIX}/{TEST_DAG_ID}", headers={"Accept": "text/plain"} + ) + assert response.status_code == 403 + @pytest.mark.parametrize( "headers", [{"Accept": "application/json"}, {"Accept": "application/json; charset=utf-8"}, {}] ) From d7a932e0d0c59e2aa4b6355dd11728f0999d8cae Mon Sep 17 00:00:00 2001 From: kalyanr Date: Sat, 1 Mar 2025 19:01:08 +0530 Subject: [PATCH 11/27] add dag_stats --- .../core_api/openapi/v1-generated.yaml | 10 ++++++++++ .../core_api/routes/public/dag_stats.py | 3 +++ airflow/ui/openapi-gen/queries/common.ts | 4 +++- airflow/ui/openapi-gen/queries/prefetch.ts | 7 +++++-- airflow/ui/openapi-gen/queries/queries.ts | 7 +++++-- airflow/ui/openapi-gen/queries/suspense.ts | 7 +++++-- .../ui/openapi-gen/requests/services.gen.ts | 2 ++ airflow/ui/openapi-gen/requests/types.gen.ts | 1 + .../core_api/routes/public/test_dag_stats.py | 20 +++++++++++++------ 9 files changed, 48 insertions(+), 13 deletions(-) diff --git a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml index 8a0b2989f3015..cb3e709c8100b 100644 --- a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml +++ b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml @@ -2762,7 +2762,17 @@ paths: summary: Get Dag Stats description: Get Dag statistics. operationId: get_dag_stats + security: + - OAuth2PasswordBearer: [] parameters: + - name: dag_id + in: query + required: false + schema: + anyOf: + - type: string + - type: 'null' + title: Dag Id - name: dag_ids in: query required: false diff --git a/airflow/api_fastapi/core_api/routes/public/dag_stats.py b/airflow/api_fastapi/core_api/routes/public/dag_stats.py index a6aa6063c263b..d70ef3fe6143f 100644 --- a/airflow/api_fastapi/core_api/routes/public/dag_stats.py +++ b/airflow/api_fastapi/core_api/routes/public/dag_stats.py @@ -21,6 +21,7 @@ from fastapi import Depends, status +from airflow.api_fastapi.auth.managers.models.resource_details import DagAccessEntity from airflow.api_fastapi.common.db.common import ( SessionDep, paginated_select, @@ -38,6 +39,7 @@ DagStatsStateResponse, ) from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc +from airflow.api_fastapi.core_api.security import requires_access_dag from airflow.models.dagrun import DagRun from airflow.utils.state import DagRunState @@ -52,6 +54,7 @@ status.HTTP_404_NOT_FOUND, ] ), + dependencies=[Depends(requires_access_dag(method="GET", access_entity=DagAccessEntity.RUN))], ) def get_dag_stats( session: SessionDep, diff --git a/airflow/ui/openapi-gen/queries/common.ts b/airflow/ui/openapi-gen/queries/common.ts index e16aef391704b..727dc9b6befdd 100644 --- a/airflow/ui/openapi-gen/queries/common.ts +++ b/airflow/ui/openapi-gen/queries/common.ts @@ -683,12 +683,14 @@ export type DagStatsServiceGetDagStatsQueryResult< export const useDagStatsServiceGetDagStatsKey = "DagStatsServiceGetDagStats"; export const UseDagStatsServiceGetDagStatsKeyFn = ( { + dagId, dagIds, }: { + dagId?: string; dagIds?: string[]; } = {}, queryKey?: Array, -) => [useDagStatsServiceGetDagStatsKey, ...(queryKey ?? [{ dagIds }])]; +) => [useDagStatsServiceGetDagStatsKey, ...(queryKey ?? [{ dagId, dagIds }])]; export type DagReportServiceGetDagReportsDefaultResponse = Awaited< ReturnType >; diff --git a/airflow/ui/openapi-gen/queries/prefetch.ts b/airflow/ui/openapi-gen/queries/prefetch.ts index 8f50495dbdabd..f495c57194eaa 100644 --- a/airflow/ui/openapi-gen/queries/prefetch.ts +++ b/airflow/ui/openapi-gen/queries/prefetch.ts @@ -923,6 +923,7 @@ export const prefetchUseDagSourceServiceGetDagSource = ( * Get Dag Stats * Get Dag statistics. * @param data The data for the request. + * @param data.dagId * @param data.dagIds * @returns DagStatsCollectionResponse Successful Response * @throws ApiError @@ -930,14 +931,16 @@ export const prefetchUseDagSourceServiceGetDagSource = ( export const prefetchUseDagStatsServiceGetDagStats = ( queryClient: QueryClient, { + dagId, dagIds, }: { + dagId?: string; dagIds?: string[]; } = {}, ) => queryClient.prefetchQuery({ - queryKey: Common.UseDagStatsServiceGetDagStatsKeyFn({ dagIds }), - queryFn: () => DagStatsService.getDagStats({ dagIds }), + queryKey: Common.UseDagStatsServiceGetDagStatsKeyFn({ dagId, dagIds }), + queryFn: () => DagStatsService.getDagStats({ dagId, dagIds }), }); /** * Get Dag Reports diff --git a/airflow/ui/openapi-gen/queries/queries.ts b/airflow/ui/openapi-gen/queries/queries.ts index f6cebbf884ebe..853a53f26ed38 100644 --- a/airflow/ui/openapi-gen/queries/queries.ts +++ b/airflow/ui/openapi-gen/queries/queries.ts @@ -1115,6 +1115,7 @@ export const useDagSourceServiceGetDagSource = < * Get Dag Stats * Get Dag statistics. * @param data The data for the request. + * @param data.dagId * @param data.dagIds * @returns DagStatsCollectionResponse Successful Response * @throws ApiError @@ -1125,16 +1126,18 @@ export const useDagStatsServiceGetDagStats = < TQueryKey extends Array = unknown[], >( { + dagId, dagIds, }: { + dagId?: string; dagIds?: string[]; } = {}, queryKey?: TQueryKey, options?: Omit, "queryKey" | "queryFn">, ) => useQuery({ - queryKey: Common.UseDagStatsServiceGetDagStatsKeyFn({ dagIds }, queryKey), - queryFn: () => DagStatsService.getDagStats({ dagIds }) as TData, + queryKey: Common.UseDagStatsServiceGetDagStatsKeyFn({ dagId, dagIds }, queryKey), + queryFn: () => DagStatsService.getDagStats({ dagId, dagIds }) as TData, ...options, }); /** diff --git a/airflow/ui/openapi-gen/queries/suspense.ts b/airflow/ui/openapi-gen/queries/suspense.ts index dd2c5fe4fcca9..b8c9689b4123e 100644 --- a/airflow/ui/openapi-gen/queries/suspense.ts +++ b/airflow/ui/openapi-gen/queries/suspense.ts @@ -1092,6 +1092,7 @@ export const useDagSourceServiceGetDagSourceSuspense = < * Get Dag Stats * Get Dag statistics. * @param data The data for the request. + * @param data.dagId * @param data.dagIds * @returns DagStatsCollectionResponse Successful Response * @throws ApiError @@ -1102,16 +1103,18 @@ export const useDagStatsServiceGetDagStatsSuspense = < TQueryKey extends Array = unknown[], >( { + dagId, dagIds, }: { + dagId?: string; dagIds?: string[]; } = {}, queryKey?: TQueryKey, options?: Omit, "queryKey" | "queryFn">, ) => useSuspenseQuery({ - queryKey: Common.UseDagStatsServiceGetDagStatsKeyFn({ dagIds }, queryKey), - queryFn: () => DagStatsService.getDagStats({ dagIds }) as TData, + queryKey: Common.UseDagStatsServiceGetDagStatsKeyFn({ dagId, dagIds }, queryKey), + queryFn: () => DagStatsService.getDagStats({ dagId, dagIds }) as TData, ...options, }); /** diff --git a/airflow/ui/openapi-gen/requests/services.gen.ts b/airflow/ui/openapi-gen/requests/services.gen.ts index 3acae741015da..d0cd4ede930c5 100644 --- a/airflow/ui/openapi-gen/requests/services.gen.ts +++ b/airflow/ui/openapi-gen/requests/services.gen.ts @@ -1579,6 +1579,7 @@ export class DagStatsService { * Get Dag Stats * Get Dag statistics. * @param data The data for the request. + * @param data.dagId * @param data.dagIds * @returns DagStatsCollectionResponse Successful Response * @throws ApiError @@ -1588,6 +1589,7 @@ export class DagStatsService { method: "GET", url: "/public/dagStats", query: { + dag_id: data.dagId, dag_ids: data.dagIds, }, errors: { diff --git a/airflow/ui/openapi-gen/requests/types.gen.ts b/airflow/ui/openapi-gen/requests/types.gen.ts index ffbcb2cbddfe3..30fc01d6ff28e 100644 --- a/airflow/ui/openapi-gen/requests/types.gen.ts +++ b/airflow/ui/openapi-gen/requests/types.gen.ts @@ -1978,6 +1978,7 @@ export type GetDagSourceData = { export type GetDagSourceResponse = DAGSourceResponse; export type GetDagStatsData = { + dagId?: string | null; dagIds?: Array; }; diff --git a/tests/api_fastapi/core_api/routes/public/test_dag_stats.py b/tests/api_fastapi/core_api/routes/public/test_dag_stats.py index 8a0dae3604c1d..e7264bd48d145 100644 --- a/tests/api_fastapi/core_api/routes/public/test_dag_stats.py +++ b/tests/api_fastapi/core_api/routes/public/test_dag_stats.py @@ -129,7 +129,7 @@ def teardown_method(self) -> None: class TestGetDagStats(TestDagStatsEndpoint): """Unit tests for Get DAG Stats.""" - def test_should_respond_200(self, client, session): + def test_should_respond_200(self, test_client, session): self._create_dag_and_runs(session) exp_payload = { "dags": [ @@ -179,13 +179,21 @@ def test_should_respond_200(self, client, session): "total_entries": 2, } - response = client().get(f"{API_PREFIX}?dag_ids={DAG1_ID}&dag_ids={DAG2_ID}") + response = test_client.get(f"{API_PREFIX}?dag_ids={DAG1_ID}&dag_ids={DAG2_ID}") assert response.status_code == 200 res_json = response.json() assert res_json["total_entries"] == len(res_json["dags"]) assert res_json == exp_payload - def test_all_dags_should_respond_200(self, client, session): + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.get(f"{API_PREFIX}?dag_ids={DAG1_ID}&dag_ids={DAG2_ID}") + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.get(f"{API_PREFIX}?dag_ids={DAG1_ID}&dag_ids={DAG2_ID}") + assert response.status_code == 403 + + def test_all_dags_should_respond_200(self, test_client, session): self._create_dag_and_runs(session) exp_payload = { "dags": [ @@ -256,7 +264,7 @@ def test_all_dags_should_respond_200(self, client, session): "total_entries": 3, } - response = client().get(API_PREFIX) + response = test_client.get(API_PREFIX) assert response.status_code == 200 res_json = response.json() assert res_json["total_entries"] == len(res_json["dags"]) @@ -403,9 +411,9 @@ def test_all_dags_should_respond_200(self, client, session): ), ], ) - def test_single_dag_in_dag_ids(self, client, session, url, params, exp_payload): + def test_single_dag_in_dag_ids(self, test_client, session, url, params, exp_payload): self._create_dag_and_runs(session) - response = client().get(url, params=params) + response = test_client.get(url, params=params) assert response.status_code == 200 res_json = response.json() assert res_json["total_entries"] == len(res_json["dags"]) From 250d5db25ce58b76d236245d4153fc9d7d020b6d Mon Sep 17 00:00:00 2001 From: kalyanr Date: Sat, 1 Mar 2025 23:18:08 +0530 Subject: [PATCH 12/27] add dag_warning --- airflow/api_fastapi/core_api/openapi/v1-generated.yaml | 2 ++ airflow/api_fastapi/core_api/routes/public/dag_warning.py | 3 +++ .../core_api/routes/public/test_dag_warning.py | 8 ++++++++ 3 files changed, 13 insertions(+) diff --git a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml index cb3e709c8100b..454677c8b8d87 100644 --- a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml +++ b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml @@ -3030,6 +3030,8 @@ paths: summary: List Dag Warnings description: Get a list of DAG warnings. operationId: list_dag_warnings + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: query diff --git a/airflow/api_fastapi/core_api/routes/public/dag_warning.py b/airflow/api_fastapi/core_api/routes/public/dag_warning.py index 2c964efce7e45..69b3354a26783 100644 --- a/airflow/api_fastapi/core_api/routes/public/dag_warning.py +++ b/airflow/api_fastapi/core_api/routes/public/dag_warning.py @@ -22,6 +22,7 @@ from fastapi import Depends from sqlalchemy import select +from airflow.api_fastapi.auth.managers.models.resource_details import DagAccessEntity from airflow.api_fastapi.common.db.common import ( SessionDep, paginated_select, @@ -37,6 +38,7 @@ from airflow.api_fastapi.core_api.datamodels.dag_warning import ( DAGWarningCollectionResponse, ) +from airflow.api_fastapi.core_api.security import requires_access_dag from airflow.models.dagwarning import DagWarning, DagWarningType dag_warning_router = AirflowRouter(tags=["DagWarning"]) @@ -44,6 +46,7 @@ @dag_warning_router.get( "/dagWarnings", + dependencies=[Depends(requires_access_dag(method="GET", access_entity=DagAccessEntity.WARNING))], ) def list_dag_warnings( dag_id: Annotated[FilterParam[str | None], Depends(filter_param_factory(DagWarning.dag_id, str | None))], diff --git a/tests/api_fastapi/core_api/routes/public/test_dag_warning.py b/tests/api_fastapi/core_api/routes/public/test_dag_warning.py index 61237bd10299a..c3dfc304cd8a3 100644 --- a/tests/api_fastapi/core_api/routes/public/test_dag_warning.py +++ b/tests/api_fastapi/core_api/routes/public/test_dag_warning.py @@ -78,6 +78,14 @@ def test_get_dag_warnings(self, test_client, query_params, expected_total_entrie assert len(response_json["dag_warnings"]) == len(expected_messages) assert [dag_warning["message"] for dag_warning in response_json["dag_warnings"]] == expected_messages + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.get("/public/dagWarnings", params={}) + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.get("/public/dagWarnings", params={}) + assert response.status_code == 403 + def test_get_dag_warnings_bad_request(self, test_client): response = test_client.get("/public/dagWarnings", params={"warning_type": "invalid"}) response_json = response.json() From 6836c5907b18706a83096b1fbfda18ae000a2957 Mon Sep 17 00:00:00 2001 From: kalyanr Date: Sun, 2 Mar 2025 00:07:55 +0530 Subject: [PATCH 13/27] add ti --- .../core_api/openapi/v1-generated.yaml | 97 ++++++++++-- .../core_api/routes/public/task_instances.py | 34 +++- airflow/ui/openapi-gen/queries/queries.ts | 4 +- airflow/ui/openapi-gen/requests/types.gen.ts | 32 ++-- .../routes/public/test_task_instances.py | 146 ++++++++++++++++++ 5 files changed, 274 insertions(+), 39 deletions(-) diff --git a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml index 454677c8b8d87..f3e59c148be7f 100644 --- a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml +++ b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml @@ -4895,12 +4895,16 @@ paths: summary: Get Task Instance description: Get task instance. operationId: get_task_instance + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id - name: dag_run_id in: path @@ -4951,12 +4955,16 @@ paths: summary: Patch Task Instance description: Update a task instance. operationId: patch_task_instance + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id - name: dag_run_id in: path @@ -5043,12 +5051,16 @@ paths: summary: Get Mapped Task Instances description: Get list of mapped task instances. operationId: get_mapped_task_instances + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id - name: dag_run_id in: path @@ -5269,12 +5281,16 @@ paths: summary: Get Task Instance Dependencies description: Get dependencies blocking task from getting scheduled. operationId: get_task_instance_dependencies + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id - name: dag_run_id in: path @@ -5332,12 +5348,16 @@ paths: summary: Get Task Instance Dependencies description: Get dependencies blocking task from getting scheduled. operationId: get_task_instance_dependencies + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id - name: dag_run_id in: path @@ -5396,12 +5416,16 @@ paths: summary: Get Task Instance Tries description: Get list of task instances history. operationId: get_task_instance_tries + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id - name: dag_run_id in: path @@ -5459,12 +5483,16 @@ paths: - Task Instance summary: Get Mapped Task Instance Tries operationId: get_mapped_task_instance_tries + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id - name: dag_run_id in: path @@ -5522,12 +5550,16 @@ paths: summary: Get Mapped Task Instance description: Get task instance. operationId: get_mapped_task_instance + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id - name: dag_run_id in: path @@ -5584,12 +5616,16 @@ paths: summary: Patch Task Instance description: Update a task instance. operationId: patch_task_instance + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id - name: dag_run_id in: path @@ -5681,12 +5717,16 @@ paths: and DAG runs.' operationId: get_task_instances + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id - name: dag_run_id in: path @@ -5917,13 +5957,16 @@ paths: summary: Get Task Instances Batch description: Get list of task instances. operationId: get_task_instances_batch + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - const: '~' - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id - name: dag_run_id in: path @@ -5976,12 +6019,16 @@ paths: summary: Get Task Instance Try Details description: Get task instance details by try number. operationId: get_task_instance_try_details + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id - name: dag_run_id in: path @@ -6045,12 +6092,16 @@ paths: - Task Instance summary: Get Mapped Task Instance Try Details operationId: get_mapped_task_instance_try_details + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id - name: dag_run_id in: path @@ -6114,12 +6165,16 @@ paths: summary: Post Clear Task Instances description: Clear task instances. operationId: post_clear_task_instances + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id requestBody: required: true @@ -6165,12 +6220,16 @@ paths: summary: Patch Task Instance Dry Run description: Update a task instance dry_run mode. operationId: patch_task_instance_dry_run + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id - name: dag_run_id in: path @@ -6250,12 +6309,16 @@ paths: summary: Patch Task Instance Dry Run description: Update a task instance dry_run mode. operationId: patch_task_instance_dry_run + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id - name: dag_run_id in: path 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 7d130bc1e37e9..89bba0aafbc74 100644 --- a/airflow/api_fastapi/core_api/routes/public/task_instances.py +++ b/airflow/api_fastapi/core_api/routes/public/task_instances.py @@ -27,6 +27,7 @@ from sqlalchemy.orm import joinedload from sqlalchemy.sql.selectable import Select +from airflow.api_fastapi.auth.managers.models.resource_details import DagAccessEntity from airflow.api_fastapi.common.db.common import SessionDep, paginated_select from airflow.api_fastapi.common.parameters import ( FilterOptionEnum, @@ -60,6 +61,7 @@ TaskInstancesBatchBody, ) from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc +from airflow.api_fastapi.core_api.security import requires_access_dag from airflow.api_fastapi.logging.decorators import action_logging from airflow.exceptions import TaskNotFound from airflow.models import Base, DagRun @@ -78,6 +80,7 @@ @task_instances_router.get( task_instances_prefix + "/{task_id}", responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]), + dependencies=[Depends(requires_access_dag(method="GET", access_entity=DagAccessEntity.TASK_INSTANCE))], ) def get_task_instance( dag_id: str, dag_run_id: str, task_id: str, session: SessionDep @@ -108,6 +111,7 @@ def get_task_instance( @task_instances_router.get( task_instances_prefix + "/{task_id}/listMapped", responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]), + dependencies=[Depends(requires_access_dag(method="GET", access_entity=DagAccessEntity.TASK_INSTANCE))], ) def get_mapped_task_instances( dag_id: str, @@ -211,10 +215,12 @@ def get_mapped_task_instances( @task_instances_router.get( task_instances_prefix + "/{task_id}/dependencies", responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]), + dependencies=[Depends(requires_access_dag(method="GET", access_entity=DagAccessEntity.TASK_INSTANCE))], ) @task_instances_router.get( task_instances_prefix + "/{task_id}/{map_index}/dependencies", responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]), + dependencies=[Depends(requires_access_dag(method="GET", access_entity=DagAccessEntity.TASK_INSTANCE))], ) def get_task_instance_dependencies( dag_id: str, @@ -265,6 +271,7 @@ def get_task_instance_dependencies( @task_instances_router.get( task_instances_prefix + "/{task_id}/tries", responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]), + dependencies=[Depends(requires_access_dag(method="GET", access_entity=DagAccessEntity.TASK_INSTANCE))], ) def get_task_instance_tries( dag_id: str, @@ -308,6 +315,7 @@ 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]), + dependencies=[Depends(requires_access_dag(method="GET", access_entity=DagAccessEntity.TASK_INSTANCE))], ) def get_mapped_task_instance_tries( dag_id: str, @@ -328,6 +336,7 @@ def get_mapped_task_instance_tries( @task_instances_router.get( task_instances_prefix + "/{task_id}/{map_index}", responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]), + dependencies=[Depends(requires_access_dag(method="GET", access_entity=DagAccessEntity.TASK_INSTANCE))], ) def get_mapped_task_instance( dag_id: str, @@ -358,6 +367,7 @@ def get_mapped_task_instance( @task_instances_router.get( task_instances_prefix, responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]), + dependencies=[Depends(requires_access_dag(method="GET", access_entity=DagAccessEntity.TASK_INSTANCE))], ) def get_task_instances( dag_id: str, @@ -464,7 +474,10 @@ def get_task_instances( @task_instances_router.post( task_instances_prefix + "/list", responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]), - dependencies=[Depends(action_logging())], + dependencies=[ + Depends(action_logging()), + Depends(requires_access_dag(method="GET", access_entity=DagAccessEntity.TASK_INSTANCE)), + ], ) def get_task_instances_batch( dag_id: Literal["~"], @@ -546,6 +559,7 @@ def get_task_instances_batch( @task_instances_router.get( task_instances_prefix + "/{task_id}/tries/{task_try_number}", responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]), + dependencies=[Depends(requires_access_dag(method="GET", access_entity=DagAccessEntity.TASK_INSTANCE))], ) def get_task_instance_try_details( dag_id: str, @@ -581,6 +595,7 @@ def _query(orm_object: Base) -> TI | TIH | None: @task_instances_router.get( task_instances_prefix + "/{task_id}/{map_index}/tries/{task_try_number}", responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]), + dependencies=[Depends(requires_access_dag(method="GET", access_entity=DagAccessEntity.TASK_INSTANCE))], ) def get_mapped_task_instance_try_details( dag_id: str, @@ -603,7 +618,10 @@ def get_mapped_task_instance_try_details( @task_instances_router.post( "/clearTaskInstances", responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]), - dependencies=[Depends(action_logging())], + dependencies=[ + Depends(action_logging()), + Depends(requires_access_dag(method="PUT", access_entity=DagAccessEntity.TASK_INSTANCE)), + ], ) def post_clear_task_instances( dag_id: str, @@ -748,12 +766,14 @@ def _patch_ti_validate_request( responses=create_openapi_http_exception_doc( [status.HTTP_404_NOT_FOUND, status.HTTP_400_BAD_REQUEST], ), + dependencies=[Depends(requires_access_dag(method="PUT", access_entity=DagAccessEntity.TASK_INSTANCE))], ) @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], ), + dependencies=[Depends(requires_access_dag(method="PUT", access_entity=DagAccessEntity.TASK_INSTANCE))], ) def patch_task_instance_dry_run( dag_id: str, @@ -808,14 +828,20 @@ def patch_task_instance_dry_run( responses=create_openapi_http_exception_doc( [status.HTTP_404_NOT_FOUND, status.HTTP_400_BAD_REQUEST, status.HTTP_409_CONFLICT], ), - dependencies=[Depends(action_logging())], + dependencies=[ + Depends(action_logging()), + Depends(requires_access_dag(method="PUT", access_entity=DagAccessEntity.TASK_INSTANCE)), + ], ) @task_instances_router.patch( task_instances_prefix + "/{task_id}/{map_index}", responses=create_openapi_http_exception_doc( [status.HTTP_404_NOT_FOUND, status.HTTP_400_BAD_REQUEST, status.HTTP_409_CONFLICT], ), - dependencies=[Depends(action_logging())], + dependencies=[ + Depends(action_logging()), + Depends(requires_access_dag(method="PUT", access_entity=DagAccessEntity.TASK_INSTANCE)), + ], ) def patch_task_instance( dag_id: str, diff --git a/airflow/ui/openapi-gen/queries/queries.ts b/airflow/ui/openapi-gen/queries/queries.ts index 853a53f26ed38..4e6f852aba7d9 100644 --- a/airflow/ui/openapi-gen/queries/queries.ts +++ b/airflow/ui/openapi-gen/queries/queries.ts @@ -3300,7 +3300,7 @@ export const useTaskInstanceServiceGetTaskInstancesBatch = < TData, TError, { - dagId: "~"; + dagId: string; dagRunId: "~"; requestBody: TaskInstancesBatchBody; }, @@ -3313,7 +3313,7 @@ export const useTaskInstanceServiceGetTaskInstancesBatch = < TData, TError, { - dagId: "~"; + dagId: string; dagRunId: "~"; requestBody: TaskInstancesBatchBody; }, diff --git a/airflow/ui/openapi-gen/requests/types.gen.ts b/airflow/ui/openapi-gen/requests/types.gen.ts index 30fc01d6ff28e..6843447af675d 100644 --- a/airflow/ui/openapi-gen/requests/types.gen.ts +++ b/airflow/ui/openapi-gen/requests/types.gen.ts @@ -2107,7 +2107,7 @@ export type GetExtraLinksData = { export type GetExtraLinksResponse = ExtraLinksResponse; export type GetTaskInstanceData = { - dagId: string; + dagId: string | null; dagRunId: string; taskId: string; }; @@ -2115,7 +2115,7 @@ export type GetTaskInstanceData = { export type GetTaskInstanceResponse = TaskInstanceResponse; export type PatchTaskInstanceData = { - dagId: string; + dagId: string | null; dagRunId: string; mapIndex?: number; requestBody: PatchTaskInstanceBody; @@ -2126,7 +2126,7 @@ export type PatchTaskInstanceData = { export type PatchTaskInstanceResponse = TaskInstanceResponse; export type GetMappedTaskInstancesData = { - dagId: string; + dagId: string | null; dagRunId: string; durationGte?: number | null; durationLte?: number | null; @@ -2154,7 +2154,7 @@ export type GetMappedTaskInstancesData = { export type GetMappedTaskInstancesResponse = TaskInstanceCollectionResponse; export type GetTaskInstanceDependenciesData = { - dagId: string; + dagId: string | null; dagRunId: string; mapIndex: number; taskId: string; @@ -2163,7 +2163,7 @@ export type GetTaskInstanceDependenciesData = { export type GetTaskInstanceDependenciesResponse = TaskDependencyCollectionResponse; export type GetTaskInstanceDependencies1Data = { - dagId: string; + dagId: string | null; dagRunId: string; mapIndex?: number; taskId: string; @@ -2172,7 +2172,7 @@ export type GetTaskInstanceDependencies1Data = { export type GetTaskInstanceDependencies1Response = TaskDependencyCollectionResponse; export type GetTaskInstanceTriesData = { - dagId: string; + dagId: string | null; dagRunId: string; mapIndex?: number; taskId: string; @@ -2181,7 +2181,7 @@ export type GetTaskInstanceTriesData = { export type GetTaskInstanceTriesResponse = TaskInstanceHistoryCollectionResponse; export type GetMappedTaskInstanceTriesData = { - dagId: string; + dagId: string | null; dagRunId: string; mapIndex: number; taskId: string; @@ -2190,7 +2190,7 @@ export type GetMappedTaskInstanceTriesData = { export type GetMappedTaskInstanceTriesResponse = TaskInstanceHistoryCollectionResponse; export type GetMappedTaskInstanceData = { - dagId: string; + dagId: string | null; dagRunId: string; mapIndex: number; taskId: string; @@ -2199,7 +2199,7 @@ export type GetMappedTaskInstanceData = { export type GetMappedTaskInstanceResponse = TaskInstanceResponse; export type PatchTaskInstance1Data = { - dagId: string; + dagId: string | null; dagRunId: string; mapIndex: number; requestBody: PatchTaskInstanceBody; @@ -2210,7 +2210,7 @@ export type PatchTaskInstance1Data = { export type PatchTaskInstance1Response = TaskInstanceResponse; export type GetTaskInstancesData = { - dagId: string; + dagId: string | null; dagRunId: string; durationGte?: number | null; durationLte?: number | null; @@ -2239,7 +2239,7 @@ export type GetTaskInstancesData = { export type GetTaskInstancesResponse = TaskInstanceCollectionResponse; export type GetTaskInstancesBatchData = { - dagId: "~"; + dagId: string | null; dagRunId: "~"; requestBody: TaskInstancesBatchBody; }; @@ -2247,7 +2247,7 @@ export type GetTaskInstancesBatchData = { export type GetTaskInstancesBatchResponse = TaskInstanceCollectionResponse; export type GetTaskInstanceTryDetailsData = { - dagId: string; + dagId: string | null; dagRunId: string; mapIndex?: number; taskId: string; @@ -2257,7 +2257,7 @@ export type GetTaskInstanceTryDetailsData = { export type GetTaskInstanceTryDetailsResponse = TaskInstanceHistoryResponse; export type GetMappedTaskInstanceTryDetailsData = { - dagId: string; + dagId: string | null; dagRunId: string; mapIndex: number; taskId: string; @@ -2267,14 +2267,14 @@ export type GetMappedTaskInstanceTryDetailsData = { export type GetMappedTaskInstanceTryDetailsResponse = TaskInstanceHistoryResponse; export type PostClearTaskInstancesData = { - dagId: string; + dagId: string | null; requestBody: ClearTaskInstancesBody; }; export type PostClearTaskInstancesResponse = TaskInstanceCollectionResponse; export type PatchTaskInstanceDryRunData = { - dagId: string; + dagId: string | null; dagRunId: string; mapIndex: number; requestBody: PatchTaskInstanceBody; @@ -2285,7 +2285,7 @@ export type PatchTaskInstanceDryRunData = { export type PatchTaskInstanceDryRunResponse = TaskInstanceCollectionResponse; export type PatchTaskInstanceDryRun1Data = { - dagId: string; + dagId: string | null; dagRunId: string; mapIndex?: number; requestBody: PatchTaskInstanceBody; 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 d7e185081c9a6..d91b7de4a1b33 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 @@ -207,6 +207,18 @@ def test_should_respond_200(self, test_client, session): "triggerer_job": None, } + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.get( + "/public/dags/example_python_operator/dagRuns/TEST_DAG_RUN_ID/taskInstances/print_the_context" + ) + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.get( + "/public/dags/example_python_operator/dagRuns/TEST_DAG_RUN_ID/taskInstances/print_the_context" + ) + assert response.status_code == 403 + @pytest.mark.parametrize( "run_id, expected_version_number", [ @@ -525,6 +537,18 @@ def test_should_respond_200_mapped_task_instance_with_rtif(self, test_client, se "triggerer_job": None, } + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.get( + "/public/dags/example_python_operator/dagRuns/TEST_DAG_RUN_ID/taskInstances/print_the_context/1", + ) + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.get( + "/public/dags/example_python_operator/dagRuns/TEST_DAG_RUN_ID/taskInstances/print_the_context/1", + ) + assert response.status_code == 403 + def test_should_respond_404_wrong_map_index(self, test_client, session): self.create_task_instances(session) @@ -665,6 +689,18 @@ def one_task_with_zero_mapped_tis(self, dag_maker, session): }, ) + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.get( + "/public/dags/mapped_tis/dagRuns/run_mapped_tis/taskInstances/task_2/listMapped", + ) + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.get( + "/public/dags/mapped_tis/dagRuns/run_mapped_tis/taskInstances/task_2/listMapped", + ) + assert response.status_code == 403 + def test_should_respond_404(self, test_client): response = test_client.get( "/public/dags/mapped_tis/dagRuns/run_mapped_tis/taskInstances/task_2/listMapped", @@ -1068,6 +1104,18 @@ def test_should_respond_200( assert response.json()["total_entries"] == expected_ti assert len(response.json()["task_instances"]) == expected_ti + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.get( + "/public/dags/example_python_operator/dagRuns/~/taskInstances", + ) + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.get( + "/public/dags/example_python_operator/dagRuns/~/taskInstances", + ) + assert response.status_code == 403 + def test_not_found(self, test_client): response = test_client.get("/public/dags/invalid/dagRuns/~/taskInstances") assert response.status_code == 404 @@ -1294,6 +1342,20 @@ def test_should_respond_dependencies_mapped(self, test_client, session): ) assert response.status_code == 200, response.text + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.get( + "/public/dags/example_python_operator/dagRuns/TEST_DAG_RUN_ID/taskInstances/" + "print_the_context/0/dependencies", + ) + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.get( + "/public/dags/example_python_operator/dagRuns/TEST_DAG_RUN_ID/taskInstances/" + "print_the_context/0/dependencies", + ) + assert response.status_code == 403 + class TestGetTaskInstancesBatch(TestTaskInstanceEndpoint): @pytest.mark.parametrize( @@ -1508,6 +1570,20 @@ def test_should_raise_400_for_no_json(self, test_client): }, ] + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.post( + "/public/dags/~/dagRuns/~/taskInstances/list", + json={}, + ) + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.post( + "/public/dags/~/dagRuns/~/taskInstances/list", + json={}, + ) + assert response.status_code == 403 + def test_should_respond_422_for_non_wildcard_path_parameters(self, test_client): response = test_client.post( "/public/dags/non_wildcard/dagRuns/~/taskInstances/list", @@ -1820,6 +1896,18 @@ def test_should_respond_200_with_task_state_in_removed(self, test_client, sessio "dag_version": None, } + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.get( + "/public/dags/example_python_operator/dagRuns/TEST_DAG_RUN_ID/taskInstances/print_the_context/tries/1", + ) + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.get( + "/public/dags/example_python_operator/dagRuns/TEST_DAG_RUN_ID/taskInstances/print_the_context/tries/1", + ) + assert response.status_code == 403 + def test_raises_404_for_nonexistent_task_instance(self, test_client, session): self.create_task_instances(session) response = test_client.get( @@ -2090,6 +2178,20 @@ def test_dag_run_with_future_or_past_flag_returns_400(self, test_client, session in response.json()["detail"] ) + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.post( + "/public/dags/dag_id/clearTaskInstances", + json={}, + ) + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.post( + "/public/dags/dag_id/clearTaskInstances", + json={}, + ) + assert response.status_code == 403 + @pytest.mark.parametrize( "main_dag, task_instances, request_dag, payload, expected_ti", [ @@ -2714,6 +2816,18 @@ def test_should_respond_200(self, test_client, session): "total_entries": 2, } + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.get( + "/public/dags/example_python_operator/dagRuns/TEST_DAG_RUN_ID/taskInstances/print_the_context/tries" + ) + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.get( + "/public/dags/example_python_operator/dagRuns/TEST_DAG_RUN_ID/taskInstances/print_the_context/tries" + ) + assert response.status_code == 403 + def test_ti_in_retry_state_not_returned(self, test_client, session): self.create_task_instances( session=session, task_instances=[{"state": State.SUCCESS}], with_ti_history=True @@ -3029,6 +3143,24 @@ def test_should_update_mapped_task_instance_state(self, test_client, session): assert response2.status_code == 200 assert response2.json()["state"] == self.NEW_STATE + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.patch( + self.ENDPOINT_URL, + json={ + "new_state": self.NEW_STATE, + }, + ) + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.patch( + self.ENDPOINT_URL, + json={ + "new_state": self.NEW_STATE, + }, + ) + assert response.status_code == 403 + @pytest.mark.parametrize( "error, code, payload", [ @@ -3529,6 +3661,20 @@ def test_should_not_update(self, test_client, session, payload): assert task_before == task_after + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.patch( + f"{self.ENDPOINT_URL}/dry_run", + json={}, + ) + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.patch( + f"{self.ENDPOINT_URL}/dry_run", + json={}, + ) + assert response.status_code == 403 + def test_should_not_update_mapped_task_instance(self, test_client, session): map_index = 1 tis = self.create_task_instances(session) From a3eef566ab43a2b3dae4e07c1a004540c3eb0c42 Mon Sep 17 00:00:00 2001 From: kalyanr Date: Sun, 2 Mar 2025 00:38:02 +0530 Subject: [PATCH 14/27] add xcom --- .../core_api/openapi/v1-generated.yaml | 24 +++++++-- .../core_api/routes/public/xcom.py | 8 ++- airflow/ui/openapi-gen/requests/types.gen.ts | 8 +-- .../core_api/routes/public/test_xcom.py | 54 +++++++++++++++++++ 4 files changed, 85 insertions(+), 9 deletions(-) diff --git a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml index f3e59c148be7f..3faa406941f6e 100644 --- a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml +++ b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml @@ -4559,12 +4559,16 @@ paths: summary: Get Xcom Entry description: Get an XCom entry. operationId: get_xcom_entry + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id - name: task_id in: path @@ -4652,12 +4656,16 @@ paths: summary: Update Xcom Entry description: Update an existing XCom entry. operationId: update_xcom_entry + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id - name: task_id in: path @@ -4731,12 +4739,16 @@ paths: This endpoint allows specifying `~` as the dag_id, dag_run_id, task_id to retrieve XCom entries for all DAGs.' operationId: get_xcom_entries + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id - name: dag_run_id in: path @@ -4826,12 +4838,16 @@ paths: summary: Create Xcom Entry description: Create an XCom entry. operationId: create_xcom_entry + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id - name: task_id in: path diff --git a/airflow/api_fastapi/core_api/routes/public/xcom.py b/airflow/api_fastapi/core_api/routes/public/xcom.py index 3da163f3e4033..572e7482540e2 100644 --- a/airflow/api_fastapi/core_api/routes/public/xcom.py +++ b/airflow/api_fastapi/core_api/routes/public/xcom.py @@ -19,9 +19,10 @@ import copy from typing import Annotated -from fastapi import HTTPException, Query, Request, status +from fastapi import Depends, HTTPException, Query, Request, status from sqlalchemy import and_, select +from airflow.api_fastapi.auth.managers.models.resource_details import DagAccessEntity from airflow.api_fastapi.common.db.common import SessionDep, paginated_select from airflow.api_fastapi.common.parameters import QueryLimit, QueryOffset from airflow.api_fastapi.common.router import AirflowRouter @@ -33,6 +34,7 @@ XComUpdateBody, ) from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc +from airflow.api_fastapi.core_api.security import requires_access_dag from airflow.exceptions import TaskNotFound from airflow.models import DAG, DagRun as DR, XCom from airflow.settings import conf @@ -50,6 +52,7 @@ status.HTTP_404_NOT_FOUND, ] ), + dependencies=[Depends(requires_access_dag(method="GET", access_entity=DagAccessEntity.XCOM))], ) def get_xcom_entry( dag_id: str, @@ -105,6 +108,7 @@ def get_xcom_entry( status.HTTP_404_NOT_FOUND, ] ), + dependencies=[Depends(requires_access_dag(method="GET", access_entity=DagAccessEntity.XCOM))], ) def get_xcom_entries( dag_id: str, @@ -155,6 +159,7 @@ def get_xcom_entries( status.HTTP_404_NOT_FOUND, ] ), + dependencies=[Depends(requires_access_dag(method="POST", access_entity=DagAccessEntity.XCOM))], ) def create_xcom_entry( dag_id: str, @@ -234,6 +239,7 @@ def create_xcom_entry( status.HTTP_404_NOT_FOUND, ] ), + dependencies=[Depends(requires_access_dag(method="PUT", access_entity=DagAccessEntity.XCOM))], ) def update_xcom_entry( dag_id: str, diff --git a/airflow/ui/openapi-gen/requests/types.gen.ts b/airflow/ui/openapi-gen/requests/types.gen.ts index 6843447af675d..4135cd252817c 100644 --- a/airflow/ui/openapi-gen/requests/types.gen.ts +++ b/airflow/ui/openapi-gen/requests/types.gen.ts @@ -2395,7 +2395,7 @@ export type GetProvidersData = { export type GetProvidersResponse = ProviderCollectionResponse; export type GetXcomEntryData = { - dagId: string; + dagId: string | null; dagRunId: string; deserialize?: boolean; mapIndex?: number; @@ -2407,7 +2407,7 @@ export type GetXcomEntryData = { export type GetXcomEntryResponse = XComResponseNative | XComResponseString; export type UpdateXcomEntryData = { - dagId: string; + dagId: string | null; dagRunId: string; requestBody: XComUpdateBody; taskId: string; @@ -2417,7 +2417,7 @@ export type UpdateXcomEntryData = { export type UpdateXcomEntryResponse = XComResponseNative; export type GetXcomEntriesData = { - dagId: string; + dagId: string | null; dagRunId: string; limit?: number; mapIndex?: number | null; @@ -2429,7 +2429,7 @@ export type GetXcomEntriesData = { export type GetXcomEntriesResponse = XComCollectionResponse; export type CreateXcomEntryData = { - dagId: string; + dagId: string | null; dagRunId: string; requestBody: XComCreateBody; taskId: string; diff --git a/tests/api_fastapi/core_api/routes/public/test_xcom.py b/tests/api_fastapi/core_api/routes/public/test_xcom.py index dd5a073c1baae..c665a5299ced3 100644 --- a/tests/api_fastapi/core_api/routes/public/test_xcom.py +++ b/tests/api_fastapi/core_api/routes/public/test_xcom.py @@ -149,6 +149,18 @@ def test_should_respond_200_native(self, test_client): "value": TEST_XCOM_VALUE, } + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.get( + f"/public/dags/{TEST_DAG_ID}/dagRuns/{run_id}/taskInstances/{TEST_TASK_ID}/xcomEntries/{TEST_XCOM_KEY}" + ) + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.get( + f"/public/dags/{TEST_DAG_ID}/dagRuns/{run_id}/taskInstances/{TEST_TASK_ID}/xcomEntries/{TEST_XCOM_KEY}" + ) + assert response.status_code == 403 + def test_should_raise_404_for_non_existent_xcom(self, test_client): response = test_client.get( f"/public/dags/{TEST_DAG_ID}/dagRuns/{run_id}/taskInstances/{TEST_TASK_ID}/xcomEntries/{TEST_XCOM_KEY_2}" @@ -437,6 +449,20 @@ def _create_xcom_entries(self, dag_id, run_id, logical_date, task_id, mapped_ti= map_index=map_index, ) + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.get( + "/public/dags/~/dagRuns/~/taskInstances/~/xcomEntries", + params={}, + ) + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.get( + "/public/dags/~/dagRuns/~/taskInstances/~/xcomEntries", + params={}, + ) + assert response.status_code == 403 + class TestPaginationGetXComEntries(TestXComEndpoint): @pytest.mark.parametrize( @@ -583,6 +609,20 @@ def test_create_xcom_entry( assert current_data["run_id"] == dag_run_id assert current_data["map_index"] == request_body.map_index + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.post( + "/public/dags/dag_id/dagRuns/dag_run_id/taskInstances/task_id/xcomEntries", + json={}, + ) + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.post( + "/public/dags/dag_id/dagRuns/dag_run_id/taskInstances/task_id/xcomEntries", + json={}, + ) + assert response.status_code == 403 + class TestPatchXComEntry(TestXComEndpoint): @pytest.mark.parametrize( @@ -623,3 +663,17 @@ def test_patch_xcom_entry(self, key, patch_body, expected_status, expected_detai assert response.json()["value"] == XCom.serialize_value(new_value) else: assert response.json()["detail"] == expected_detail + + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.patch( + f"/public/dags/{TEST_DAG_ID}/dagRuns/run_id/taskInstances/TEST_TASK_ID/xcomEntries/key", + json={}, + ) + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.patch( + f"/public/dags/{TEST_DAG_ID}/dagRuns/run_id/taskInstances/TEST_TASK_ID/xcomEntries/key", + json={}, + ) + assert response.status_code == 403 From 62b64e902b637918330cc719ee628438f7729c82 Mon Sep 17 00:00:00 2001 From: kalyanr Date: Tue, 11 Mar 2025 21:47:10 +0530 Subject: [PATCH 15/27] pr feedback --- airflow/api_fastapi/core_api/routes/public/dag_run.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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 eff8d729b18c4..01840f6c74ed8 100644 --- a/airflow/api_fastapi/core_api/routes/public/dag_run.py +++ b/airflow/api_fastapi/core_api/routes/public/dag_run.py @@ -229,7 +229,7 @@ def get_upstream_asset_events( "/{dag_run_id}/clear", responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]), dependencies=[ - Depends(requires_access_dag(method="POST", access_entity=DagAccessEntity.RUN)), + Depends(requires_access_dag(method="PUT", access_entity=DagAccessEntity.RUN)), Depends(action_logging()), ], ) @@ -407,7 +407,7 @@ def trigger_dag_run( @dag_run_router.post( "/list", responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]), - dependencies=[Depends(requires_access_dag(method="POST", access_entity=DagAccessEntity.RUN))], + dependencies=[Depends(requires_access_dag(method="GET", access_entity=DagAccessEntity.RUN))], ) def get_list_dag_runs_batch( dag_id: Literal["~"], body: DAGRunsBatchBody, session: SessionDep From 5eedfba4ef6822f6a541942928db91f6bd1243d7 Mon Sep 17 00:00:00 2001 From: kalyanr Date: Wed, 12 Mar 2025 12:51:19 +0530 Subject: [PATCH 16/27] use dag run filter --- .../core_api/routes/public/dag_run.py | 33 ++++++++++++++++--- airflow/api_fastapi/core_api/security.py | 21 ++++++++++-- 2 files changed, 47 insertions(+), 7 deletions(-) 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 01840f6c74ed8..00d38be210c63 100644 --- a/airflow/api_fastapi/core_api/routes/public/dag_run.py +++ b/airflow/api_fastapi/core_api/routes/public/dag_run.py @@ -60,7 +60,11 @@ TaskInstanceResponse, ) from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc -from airflow.api_fastapi.core_api.security import requires_access_asset, requires_access_dag +from airflow.api_fastapi.core_api.security import ( + ReadableDagRunsFilterDep, + requires_access_asset, + requires_access_dag, +) from airflow.api_fastapi.logging.decorators import action_logging from airflow.exceptions import ParamValidationError from airflow.listeners.listener import get_listener_manager @@ -313,6 +317,8 @@ def get_dag_runs( ).dynamic_depends(default="id") ), ], + # readable_dags_filter: ReadableDagsFilterDep, + readable_dag_runs_filter: ReadableDagRunsFilterDep, session: SessionDep, request: Request, ) -> DAGRunCollectionResponse: @@ -321,6 +327,9 @@ def get_dag_runs( This endpoint allows specifying `~` as the dag_id to retrieve Dag Runs for all DAGs. """ + # if readable_dags_filter.value is None: + # return DAGRunCollectionResponse(dag_runs=[], total_entries=0) + query = select(DagRun) if dag_id != "~": @@ -328,11 +337,24 @@ def get_dag_runs( if not dag: raise HTTPException(status.HTTP_404_NOT_FOUND, f"The DAG with dag_id: `{dag_id}` was not found") + # if dag_id not in readable_dags_filter.value: + # raise HTTPException(status.HTTP_403_FORBIDDEN, f"Access to DAG with dag_id: `{dag_id}` is forbidden") + query = query.filter(DagRun.dag_id == dag_id) + # else: + # query = query.filter(DagRun.dag_id.in_(readable_dags_filter.value)) dag_run_select, total_entries = paginated_select( statement=query, - filters=[run_after, logical_date, start_date_range, end_date_range, update_at_range, state], + filters=[ + run_after, + logical_date, + start_date_range, + end_date_range, + update_at_range, + state, + readable_dag_runs_filter, + ], order_by=order_by, offset=offset, limit=limit, @@ -410,7 +432,10 @@ def trigger_dag_run( dependencies=[Depends(requires_access_dag(method="GET", access_entity=DagAccessEntity.RUN))], ) def get_list_dag_runs_batch( - dag_id: Literal["~"], body: DAGRunsBatchBody, session: SessionDep + dag_id: Literal["~"], + body: DAGRunsBatchBody, + readable_dag_runs_filter: ReadableDagRunsFilterDep, + session: SessionDep, ) -> DAGRunCollectionResponse: """Get a list of DAG Runs.""" dag_ids = FilterParam(DagRun.dag_id, body.dag_ids, FilterOptionEnum.IN) @@ -455,7 +480,7 @@ def get_list_dag_runs_batch( base_query = select(DagRun) dag_runs_select, total_entries = paginated_select( statement=base_query, - filters=[dag_ids, logical_date, run_after, start_date, end_date, state], + filters=[dag_ids, logical_date, run_after, start_date, end_date, state, readable_dag_runs_filter], order_by=order_by, offset=offset, limit=limit, diff --git a/airflow/api_fastapi/core_api/security.py b/airflow/api_fastapi/core_api/security.py index 3d40f78523450..a81aea78386a5 100644 --- a/airflow/api_fastapi/core_api/security.py +++ b/airflow/api_fastapi/core_api/security.py @@ -40,7 +40,7 @@ ) from airflow.api_fastapi.core_api.base import OrmClause from airflow.configuration import conf -from airflow.models.dag import DagModel +from airflow.models.dag import DagModel, DagRun from airflow.utils.jwt_signer import JWTSigner, get_signing_key if TYPE_CHECKING: @@ -114,7 +114,16 @@ def to_orm(self, select: Select) -> Select: return select.where(DagModel.dag_id.in_(self.value)) -def permitted_dag_filter_factory(method: ResourceMethod) -> Callable[[Request, BaseUser], PermittedDagFilter]: +class PermittedDagRunFilter(PermittedDagFilter): + """A parameter that filters the permitted dag runs for the user.""" + + def to_orm(self, select: Select) -> Select: + return select.where(DagRun.dag_id.in_(self.value)) + + +def permitted_dag_filter_factory( + method: ResourceMethod, filter_class=PermittedDagFilter +) -> Callable[[Request, BaseUser], PermittedDagFilter]: """ Create a callable for Depends in FastAPI that returns a filter of the permitted dags for the user. @@ -128,13 +137,19 @@ def depends_permitted_dags_filter( ) -> PermittedDagFilter: auth_manager: BaseAuthManager = request.app.state.auth_manager permitted_dags: set[str] = auth_manager.get_permitted_dag_ids(user=user, method=method) - return PermittedDagFilter(permitted_dags) + return filter_class(permitted_dags) return depends_permitted_dags_filter EditableDagsFilterDep = Annotated[PermittedDagFilter, Depends(permitted_dag_filter_factory("PUT"))] ReadableDagsFilterDep = Annotated[PermittedDagFilter, Depends(permitted_dag_filter_factory("GET"))] +ReadableDagRunsFilterDep = Annotated[ + PermittedDagRunFilter, Depends(permitted_dag_filter_factory("GET", PermittedDagRunFilter)) +] +EditableDagRunsFilterDep = Annotated[ + PermittedDagRunFilter, Depends(permitted_dag_filter_factory("PUT", PermittedDagRunFilter)) +] def requires_access_pool(method: ResourceMethod) -> Callable[[Request, BaseUser], None]: From bca528e7a266e4e5c58c86cfd4f93a7ed3214a52 Mon Sep 17 00:00:00 2001 From: kalyanr Date: Sat, 8 Mar 2025 16:10:37 +0530 Subject: [PATCH 17/27] add dag_parsing --- .../api_fastapi/core_api/openapi/v1-generated.yaml | 10 ++++++++++ .../core_api/routes/public/dag_parsing.py | 3 ++- airflow/ui/openapi-gen/queries/queries.ts | 7 +++++-- airflow/ui/openapi-gen/requests/services.gen.ts | 4 ++++ airflow/ui/openapi-gen/requests/types.gen.ts | 1 + .../core_api/routes/public/test_dag_parsing.py | 12 ++++++++++++ 6 files changed, 34 insertions(+), 3 deletions(-) diff --git a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml index 0eab45da49794..a3bb35006a494 100644 --- a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml +++ b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml @@ -6820,6 +6820,8 @@ paths: summary: Reparse Dag File description: Request re-parsing a DAG file. operationId: reparse_dag_file + security: + - OAuth2PasswordBearer: [] parameters: - name: file_token in: path @@ -6827,6 +6829,14 @@ paths: schema: type: string title: File Token + - name: dag_id + in: query + required: false + schema: + anyOf: + - type: string + - type: 'null' + title: Dag Id responses: '201': description: Successful Response diff --git a/airflow/api_fastapi/core_api/routes/public/dag_parsing.py b/airflow/api_fastapi/core_api/routes/public/dag_parsing.py index f5b7f2c359335..a4deb3b18de02 100644 --- a/airflow/api_fastapi/core_api/routes/public/dag_parsing.py +++ b/airflow/api_fastapi/core_api/routes/public/dag_parsing.py @@ -27,6 +27,7 @@ from airflow.api_fastapi.common.db.common import SessionDep from airflow.api_fastapi.common.router import AirflowRouter from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc +from airflow.api_fastapi.core_api.security import requires_access_dag from airflow.api_fastapi.logging.decorators import action_logging from airflow.models.dag import DagModel from airflow.models.dagbag import DagPriorityParsingRequest @@ -41,7 +42,7 @@ "", responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]), status_code=status.HTTP_201_CREATED, - dependencies=[Depends(action_logging())], + dependencies=[Depends(requires_access_dag(method="PUT")), Depends(action_logging())], ) def reparse_dag_file( file_token: str, diff --git a/airflow/ui/openapi-gen/queries/queries.ts b/airflow/ui/openapi-gen/queries/queries.ts index 02b647cdfa104..d460a56a66389 100644 --- a/airflow/ui/openapi-gen/queries/queries.ts +++ b/airflow/ui/openapi-gen/queries/queries.ts @@ -3596,6 +3596,7 @@ export const useBackfillServiceCancelBackfill = < * Request re-parsing a DAG file. * @param data The data for the request. * @param data.fileToken + * @param data.dagId * @returns null Successful Response * @throws ApiError */ @@ -3609,6 +3610,7 @@ export const useDagParsingServiceReparseDagFile = < TData, TError, { + dagId?: string; fileToken: string; }, TContext @@ -3620,12 +3622,13 @@ export const useDagParsingServiceReparseDagFile = < TData, TError, { + dagId?: string; fileToken: string; }, TContext >({ - mutationFn: ({ fileToken }) => - DagParsingService.reparseDagFile({ fileToken }) as unknown as Promise, + mutationFn: ({ dagId, fileToken }) => + DagParsingService.reparseDagFile({ dagId, fileToken }) as unknown as Promise, ...options, }); /** diff --git a/airflow/ui/openapi-gen/requests/services.gen.ts b/airflow/ui/openapi-gen/requests/services.gen.ts index 11b54997040d6..3acae741015da 100644 --- a/airflow/ui/openapi-gen/requests/services.gen.ts +++ b/airflow/ui/openapi-gen/requests/services.gen.ts @@ -3388,6 +3388,7 @@ export class DagParsingService { * Request re-parsing a DAG file. * @param data The data for the request. * @param data.fileToken + * @param data.dagId * @returns null Successful Response * @throws ApiError */ @@ -3398,6 +3399,9 @@ export class DagParsingService { path: { file_token: data.fileToken, }, + query: { + dag_id: data.dagId, + }, 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 1ed2a13c95916..4b7257b9e1b2e 100644 --- a/airflow/ui/openapi-gen/requests/types.gen.ts +++ b/airflow/ui/openapi-gen/requests/types.gen.ts @@ -2492,6 +2492,7 @@ export type BulkVariablesData = { export type BulkVariablesResponse = BulkResponse; export type ReparseDagFileData = { + dagId?: string | null; fileToken: string; }; diff --git a/tests/api_fastapi/core_api/routes/public/test_dag_parsing.py b/tests/api_fastapi/core_api/routes/public/test_dag_parsing.py index 0943684e32192..00ab7dbbaf324 100644 --- a/tests/api_fastapi/core_api/routes/public/test_dag_parsing.py +++ b/tests/api_fastapi/core_api/routes/public/test_dag_parsing.py @@ -65,6 +65,18 @@ def test_201_and_400_requests(self, url_safe_serializer, session, test_client): assert parsing_requests[0].fileloc == test_dag.fileloc _check_last_log(session, dag_id=None, event="reparse_dag_file", logical_date=None) + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.put( + "/public/parseDagFile/token", headers={"Accept": "application/json"} + ) + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.put( + "/public/parseDagFile/token", headers={"Accept": "application/json"} + ) + assert response.status_code == 403 + def test_bad_file_request(self, url_safe_serializer, session, test_client): url = f"/public/parseDagFile/{url_safe_serializer.dumps('/some/random/file.py')}" response = test_client.put(url, headers={"Accept": "application/json"}) From f08f578f9c7a1228177e4b7f3827428affacc4bf Mon Sep 17 00:00:00 2001 From: kalyanr Date: Sat, 1 Mar 2025 18:09:03 +0530 Subject: [PATCH 18/27] add dag_run --- .../core_api/openapi/v1-generated.yaml | 46 ++++++++++--- .../core_api/routes/public/dag_run.py | 41 ++++++++--- airflow/ui/openapi-gen/queries/queries.ts | 8 +-- airflow/ui/openapi-gen/requests/types.gen.ts | 16 ++--- .../core_api/routes/public/test_dag_run.py | 68 +++++++++++++++++++ 5 files changed, 151 insertions(+), 28 deletions(-) diff --git a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml index a3bb35006a494..eedcac20ee719 100644 --- a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml +++ b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml @@ -2081,12 +2081,16 @@ paths: - DagRun summary: Get Dag Run operationId: get_dag_run + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id - name: dag_run_id in: path @@ -2131,12 +2135,16 @@ paths: summary: Delete Dag Run description: Delete a DAG Run entry. operationId: delete_dag_run + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id - name: dag_run_id in: path @@ -2183,12 +2191,16 @@ paths: summary: Patch Dag Run description: Modify a DAG Run. operationId: patch_dag_run + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id - name: dag_run_id in: path @@ -2264,7 +2276,9 @@ paths: in: path required: true schema: - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id - name: dag_run_id in: path @@ -2309,12 +2323,16 @@ paths: - DagRun summary: Clear Dag Run operationId: clear_dag_run + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id - name: dag_run_id in: path @@ -2373,12 +2391,16 @@ paths: This endpoint allows specifying `~` as the dag_id to retrieve Dag Runs for all DAGs.' operationId: get_dag_runs + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id - name: limit in: query @@ -2538,11 +2560,16 @@ paths: summary: Trigger Dag Run description: Trigger a DAG. operationId: trigger_dag_run + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: + anyOf: + - type: string + - type: 'null' title: Dag Id requestBody: required: true @@ -2600,13 +2627,16 @@ paths: summary: Get List Dag Runs Batch description: Get a list of DAG Runs. operationId: get_list_dag_runs_batch + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - const: '~' - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id requestBody: required: true 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 8703fe42e423a..eff8d729b18c4 100644 --- a/airflow/api_fastapi/core_api/routes/public/dag_run.py +++ b/airflow/api_fastapi/core_api/routes/public/dag_run.py @@ -29,6 +29,7 @@ set_dag_run_state_to_queued, set_dag_run_state_to_success, ) +from airflow.api_fastapi.auth.managers.models.resource_details import DagAccessEntity from airflow.api_fastapi.common.db.common import SessionDep, paginated_select from airflow.api_fastapi.common.parameters import ( FilterOptionEnum, @@ -59,7 +60,7 @@ TaskInstanceResponse, ) from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc -from airflow.api_fastapi.core_api.security import requires_access_asset +from airflow.api_fastapi.core_api.security import requires_access_asset, requires_access_dag from airflow.api_fastapi.logging.decorators import action_logging from airflow.exceptions import ParamValidationError from airflow.listeners.listener import get_listener_manager @@ -78,6 +79,7 @@ status.HTTP_404_NOT_FOUND, ] ), + dependencies=[Depends(requires_access_dag(method="GET", access_entity=DagAccessEntity.RUN))], ) def get_dag_run(dag_id: str, dag_run_id: str, session: SessionDep) -> DAGRunResponse: dag_run = session.scalar(select(DagRun).filter_by(dag_id=dag_id, run_id=dag_run_id)) @@ -99,7 +101,10 @@ def get_dag_run(dag_id: str, dag_run_id: str, session: SessionDep) -> DAGRunResp status.HTTP_404_NOT_FOUND, ], ), - dependencies=[Depends(action_logging())], + dependencies=[ + Depends(requires_access_dag(method="DELETE", access_entity=DagAccessEntity.RUN)), + Depends(action_logging()), + ], ) def delete_dag_run(dag_id: str, dag_run_id: str, session: SessionDep): """Delete a DAG Run entry.""" @@ -121,7 +126,10 @@ def delete_dag_run(dag_id: str, dag_run_id: str, session: SessionDep): status.HTTP_404_NOT_FOUND, ], ), - dependencies=[Depends(action_logging())], + dependencies=[ + Depends(requires_access_dag(method="PUT", access_entity=DagAccessEntity.RUN)), + Depends(action_logging()), + ], ) def patch_dag_run( dag_id: str, @@ -190,7 +198,10 @@ def patch_dag_run( status.HTTP_404_NOT_FOUND, ] ), - dependencies=[Depends(requires_access_asset(method="GET"))], + dependencies=[ + Depends(requires_access_asset(method="GET")), + Depends(requires_access_dag(method="GET", access_entity=DagAccessEntity.RUN)), + ], ) def get_upstream_asset_events( dag_id: str, dag_run_id: str, session: SessionDep @@ -217,7 +228,10 @@ def get_upstream_asset_events( @dag_run_router.post( "/{dag_run_id}/clear", responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]), - dependencies=[Depends(action_logging())], + dependencies=[ + Depends(requires_access_dag(method="POST", access_entity=DagAccessEntity.RUN)), + Depends(action_logging()), + ], ) def clear_dag_run( dag_id: str, @@ -263,7 +277,11 @@ def clear_dag_run( return dag_run_cleared -@dag_run_router.get("", responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND])) +@dag_run_router.get( + "", + responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]), + dependencies=[Depends(requires_access_dag(method="GET", access_entity=DagAccessEntity.RUN))], +) def get_dag_runs( dag_id: str, limit: QueryLimit, @@ -337,7 +355,10 @@ def get_dag_runs( status.HTTP_409_CONFLICT, ] ), - dependencies=[Depends(action_logging())], + dependencies=[ + Depends(requires_access_dag(method="POST", access_entity=DagAccessEntity.RUN)), + Depends(action_logging()), + ], ) def trigger_dag_run( dag_id, @@ -383,7 +404,11 @@ def trigger_dag_run( raise HTTPException(status.HTTP_400_BAD_REQUEST, str(e)) -@dag_run_router.post("/list", responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND])) +@dag_run_router.post( + "/list", + responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]), + dependencies=[Depends(requires_access_dag(method="POST", access_entity=DagAccessEntity.RUN))], +) def get_list_dag_runs_batch( dag_id: Literal["~"], body: DAGRunsBatchBody, session: SessionDep ) -> DAGRunCollectionResponse: diff --git a/airflow/ui/openapi-gen/queries/queries.ts b/airflow/ui/openapi-gen/queries/queries.ts index d460a56a66389..f6cebbf884ebe 100644 --- a/airflow/ui/openapi-gen/queries/queries.ts +++ b/airflow/ui/openapi-gen/queries/queries.ts @@ -3216,7 +3216,7 @@ export const useDagRunServiceTriggerDagRun = < TData, TError, { - dagId: unknown; + dagId: string; requestBody: TriggerDAGRunPostBody; }, TContext @@ -3228,7 +3228,7 @@ export const useDagRunServiceTriggerDagRun = < TData, TError, { - dagId: unknown; + dagId: string; requestBody: TriggerDAGRunPostBody; }, TContext @@ -3256,7 +3256,7 @@ export const useDagRunServiceGetListDagRunsBatch = < TData, TError, { - dagId: "~"; + dagId: string; requestBody: DAGRunsBatchBody; }, TContext @@ -3268,7 +3268,7 @@ export const useDagRunServiceGetListDagRunsBatch = < TData, TError, { - dagId: "~"; + dagId: string; requestBody: DAGRunsBatchBody; }, TContext diff --git a/airflow/ui/openapi-gen/requests/types.gen.ts b/airflow/ui/openapi-gen/requests/types.gen.ts index 4b7257b9e1b2e..36544aa4ad4cc 100644 --- a/airflow/ui/openapi-gen/requests/types.gen.ts +++ b/airflow/ui/openapi-gen/requests/types.gen.ts @@ -1898,21 +1898,21 @@ export type TestConnectionResponse = ConnectionTestResponse; export type CreateDefaultConnectionsResponse = void; export type GetDagRunData = { - dagId: string; + dagId: string | null; dagRunId: string; }; export type GetDagRunResponse = DAGRunResponse; export type DeleteDagRunData = { - dagId: string; + dagId: string | null; dagRunId: string; }; export type DeleteDagRunResponse = void; export type PatchDagRunData = { - dagId: string; + dagId: string | null; dagRunId: string; requestBody: DAGRunPatchBody; updateMask?: Array | null; @@ -1921,14 +1921,14 @@ export type PatchDagRunData = { export type PatchDagRunResponse = DAGRunResponse; export type GetUpstreamAssetEventsData = { - dagId: string; + dagId: string | null; dagRunId: string; }; export type GetUpstreamAssetEventsResponse = AssetEventCollectionResponse; export type ClearDagRunData = { - dagId: string; + dagId: string | null; dagRunId: string; requestBody: DAGRunClearBody; }; @@ -1936,7 +1936,7 @@ export type ClearDagRunData = { export type ClearDagRunResponse = TaskInstanceCollectionResponse | DAGRunResponse; export type GetDagRunsData = { - dagId: string; + dagId: string | null; endDateGte?: string | null; endDateLte?: string | null; limit?: number; @@ -1956,14 +1956,14 @@ export type GetDagRunsData = { export type GetDagRunsResponse = DAGRunCollectionResponse; export type TriggerDagRunData = { - dagId: unknown; + dagId: string | null; requestBody: TriggerDAGRunPostBody; }; export type TriggerDagRunResponse = DAGRunResponse; export type GetListDagRunsBatchData = { - dagId: "~"; + dagId: string | null; requestBody: DAGRunsBatchBody; }; 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 f62bff0665ec6..61199fc1f37a5 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 @@ -244,6 +244,14 @@ def test_get_dag_run_not_found(self, test_client): body = response.json() assert body["detail"] == "The DagRun with dag_id: `test_dag1` and run_id: `invalid` was not found" + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.get(f"/public/dags/{DAG1_ID}/dagRuns/invalid") + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.get(f"/public/dags/{DAG1_ID}/dagRuns/invalid") + assert response.status_code == 403 + class TestGetDagRuns: @pytest.mark.parametrize("dag_id, total_entries", [(DAG1_ID, 2), (DAG2_ID, 2), ("~", 4)]) @@ -277,6 +285,14 @@ def test_invalid_order_by_raises_400(self, test_client): == "Ordering with 'invalid' is disallowed or the attribute does not exist on the model" ) + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.get("/public/dags/test_dag1/dagRuns?order_by=invalid") + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.get("/public/dags/test_dag1/dagRuns?order_by=invalid") + assert response.status_code == 403 + @pytest.mark.parametrize( "order_by,expected_order", [ @@ -550,6 +566,14 @@ def test_list_dag_runs_return_200(self, test_client, session): expected = get_dag_run_dict(run) assert each == expected + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.post("/public/dags/~/dagRuns/list", json={}) + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.post("/public/dags/~/dagRuns/list", json={}) + assert response.status_code == 403 + def test_list_dag_runs_with_invalid_dag_id(self, test_client): response = test_client.post("/public/dags/invalid/dagRuns/list", json={}) assert response.status_code == 422 @@ -909,6 +933,14 @@ def test_patch_dag_run(self, test_client, dag_id, run_id, patch_body, response_b assert body.get("note") == response_body.get("note") _check_last_log(session, dag_id=dag_id, event="patch_dag_run", logical_date=None) + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.patch("/public/dags/dag_1/dagRuns/run_1", json={}) + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.patch("/public/dags/dag_1/dagRuns/run_1", json={}) + assert response.status_code == 403 + @pytest.mark.parametrize( "query_params, patch_body, response_body, expected_status_code", [ @@ -1008,6 +1040,14 @@ def test_delete_dag_run_not_found(self, test_client): body = response.json() assert body["detail"] == "The DagRun with dag_id: `test_dag1` and run_id: `invalid` was not found" + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.delete(f"/public/dags/{DAG1_ID}/dagRuns/invalid") + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.delete(f"/public/dags/{DAG1_ID}/dagRuns/invalid") + assert response.status_code == 403 + class TestGetDagRunAssetTriggerEvents: @pytest.mark.usefixtures("configure_git_connection_for_dag_bundle") @@ -1115,6 +1155,20 @@ def test_clear_dag_run(self, test_client, session): logical_date=None, ) + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.post( + f"/public/dags/{DAG1_ID}/dagRuns/{DAG1_RUN1_ID}/clear", + json={"dry_run": False}, + ) + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.post( + f"/public/dags/{DAG1_ID}/dagRuns/{DAG1_RUN1_ID}/clear", + json={"dry_run": False}, + ) + assert response.status_code == 403 + @pytest.mark.parametrize( "body, dag_run_id, expected_state", [ @@ -1246,6 +1300,20 @@ def test_should_respond_200( assert response.json() == expected_response_json _check_last_log(session, dag_id=DAG1_ID, event="trigger_dag_run", logical_date=None) + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.post( + f"/public/dags/{DAG1_ID}/dagRuns", + json={}, + ) + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.post( + f"/public/dags/{DAG1_ID}/dagRuns", + json={}, + ) + assert response.status_code == 403 + @pytest.mark.parametrize( "post_body, expected_detail", [ From f4d406cb4aad8398b5abb559e1c10f421c2bbfc0 Mon Sep 17 00:00:00 2001 From: kalyanr Date: Sat, 1 Mar 2025 18:40:05 +0530 Subject: [PATCH 19/27] add dag_source --- .../api_fastapi/core_api/openapi/v1-generated.yaml | 6 +++++- .../core_api/routes/public/dag_sources.py | 5 ++++- airflow/ui/openapi-gen/requests/types.gen.ts | 2 +- .../core_api/routes/public/test_dag_sources.py | 12 ++++++++++++ 4 files changed, 22 insertions(+), 3 deletions(-) diff --git a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml index eedcac20ee719..01356646a046a 100644 --- a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml +++ b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml @@ -2682,12 +2682,16 @@ paths: summary: Get Dag Source description: Get source code using file token. operationId: get_dag_source + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id - name: version_number in: query diff --git a/airflow/api_fastapi/core_api/routes/public/dag_sources.py b/airflow/api_fastapi/core_api/routes/public/dag_sources.py index 4337c92107322..fb00a4dd553b5 100644 --- a/airflow/api_fastapi/core_api/routes/public/dag_sources.py +++ b/airflow/api_fastapi/core_api/routes/public/dag_sources.py @@ -16,14 +16,16 @@ # under the License. from __future__ import annotations -from fastapi import HTTPException, Response, status +from fastapi import Depends, HTTPException, Response, status +from airflow.api_fastapi.auth.managers.models.resource_details import DagAccessEntity from airflow.api_fastapi.common.db.common import SessionDep from airflow.api_fastapi.common.headers import HeaderAcceptJsonOrText from airflow.api_fastapi.common.router import AirflowRouter from airflow.api_fastapi.common.types import Mimetype from airflow.api_fastapi.core_api.datamodels.dag_sources import DAGSourceResponse from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc +from airflow.api_fastapi.core_api.security import requires_access_dag from airflow.models.dag_version import DagVersion dag_sources_router = AirflowRouter(tags=["DagSource"], prefix="/dagSources") @@ -47,6 +49,7 @@ }, }, response_model=DAGSourceResponse, + dependencies=[Depends(requires_access_dag(method="GET", access_entity=DagAccessEntity.CODE))], ) def get_dag_source( accept: HeaderAcceptJsonOrText, diff --git a/airflow/ui/openapi-gen/requests/types.gen.ts b/airflow/ui/openapi-gen/requests/types.gen.ts index 36544aa4ad4cc..d04524adb8fa0 100644 --- a/airflow/ui/openapi-gen/requests/types.gen.ts +++ b/airflow/ui/openapi-gen/requests/types.gen.ts @@ -1971,7 +1971,7 @@ export type GetListDagRunsBatchResponse = DAGRunCollectionResponse; export type GetDagSourceData = { accept?: "application/json" | "text/plain" | "*/*"; - dagId: string; + dagId: string | null; versionNumber?: number | null; }; diff --git a/tests/api_fastapi/core_api/routes/public/test_dag_sources.py b/tests/api_fastapi/core_api/routes/public/test_dag_sources.py index 6da8902b232a0..8bd53ea2f97b9 100644 --- a/tests/api_fastapi/core_api/routes/public/test_dag_sources.py +++ b/tests/api_fastapi/core_api/routes/public/test_dag_sources.py @@ -77,6 +77,18 @@ def test_should_respond_200_text(self, test_client, test_dag): json.loads(response.content.decode()) assert response.headers["Content-Type"].startswith("text/plain") + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.get( + f"{API_PREFIX}/{TEST_DAG_ID}", headers={"Accept": "text/plain"} + ) + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.get( + f"{API_PREFIX}/{TEST_DAG_ID}", headers={"Accept": "text/plain"} + ) + assert response.status_code == 403 + @pytest.mark.parametrize( "headers", [{"Accept": "application/json"}, {"Accept": "application/json; charset=utf-8"}, {}] ) From f4ee09c7995c64c8151f9b8e8c200e84a98976b6 Mon Sep 17 00:00:00 2001 From: kalyanr Date: Sat, 1 Mar 2025 19:01:08 +0530 Subject: [PATCH 20/27] add dag_stats --- .../core_api/openapi/v1-generated.yaml | 10 ++++++++++ .../core_api/routes/public/dag_stats.py | 3 +++ airflow/ui/openapi-gen/queries/common.ts | 4 +++- airflow/ui/openapi-gen/queries/prefetch.ts | 7 +++++-- airflow/ui/openapi-gen/queries/queries.ts | 7 +++++-- airflow/ui/openapi-gen/queries/suspense.ts | 7 +++++-- .../ui/openapi-gen/requests/services.gen.ts | 2 ++ airflow/ui/openapi-gen/requests/types.gen.ts | 1 + .../core_api/routes/public/test_dag_stats.py | 20 +++++++++++++------ 9 files changed, 48 insertions(+), 13 deletions(-) diff --git a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml index 01356646a046a..3e4908ba1f57d 100644 --- a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml +++ b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml @@ -2766,7 +2766,17 @@ paths: summary: Get Dag Stats description: Get Dag statistics. operationId: get_dag_stats + security: + - OAuth2PasswordBearer: [] parameters: + - name: dag_id + in: query + required: false + schema: + anyOf: + - type: string + - type: 'null' + title: Dag Id - name: dag_ids in: query required: false diff --git a/airflow/api_fastapi/core_api/routes/public/dag_stats.py b/airflow/api_fastapi/core_api/routes/public/dag_stats.py index a6aa6063c263b..d70ef3fe6143f 100644 --- a/airflow/api_fastapi/core_api/routes/public/dag_stats.py +++ b/airflow/api_fastapi/core_api/routes/public/dag_stats.py @@ -21,6 +21,7 @@ from fastapi import Depends, status +from airflow.api_fastapi.auth.managers.models.resource_details import DagAccessEntity from airflow.api_fastapi.common.db.common import ( SessionDep, paginated_select, @@ -38,6 +39,7 @@ DagStatsStateResponse, ) from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc +from airflow.api_fastapi.core_api.security import requires_access_dag from airflow.models.dagrun import DagRun from airflow.utils.state import DagRunState @@ -52,6 +54,7 @@ status.HTTP_404_NOT_FOUND, ] ), + dependencies=[Depends(requires_access_dag(method="GET", access_entity=DagAccessEntity.RUN))], ) def get_dag_stats( session: SessionDep, diff --git a/airflow/ui/openapi-gen/queries/common.ts b/airflow/ui/openapi-gen/queries/common.ts index e16aef391704b..727dc9b6befdd 100644 --- a/airflow/ui/openapi-gen/queries/common.ts +++ b/airflow/ui/openapi-gen/queries/common.ts @@ -683,12 +683,14 @@ export type DagStatsServiceGetDagStatsQueryResult< export const useDagStatsServiceGetDagStatsKey = "DagStatsServiceGetDagStats"; export const UseDagStatsServiceGetDagStatsKeyFn = ( { + dagId, dagIds, }: { + dagId?: string; dagIds?: string[]; } = {}, queryKey?: Array, -) => [useDagStatsServiceGetDagStatsKey, ...(queryKey ?? [{ dagIds }])]; +) => [useDagStatsServiceGetDagStatsKey, ...(queryKey ?? [{ dagId, dagIds }])]; export type DagReportServiceGetDagReportsDefaultResponse = Awaited< ReturnType >; diff --git a/airflow/ui/openapi-gen/queries/prefetch.ts b/airflow/ui/openapi-gen/queries/prefetch.ts index 8f50495dbdabd..f495c57194eaa 100644 --- a/airflow/ui/openapi-gen/queries/prefetch.ts +++ b/airflow/ui/openapi-gen/queries/prefetch.ts @@ -923,6 +923,7 @@ export const prefetchUseDagSourceServiceGetDagSource = ( * Get Dag Stats * Get Dag statistics. * @param data The data for the request. + * @param data.dagId * @param data.dagIds * @returns DagStatsCollectionResponse Successful Response * @throws ApiError @@ -930,14 +931,16 @@ export const prefetchUseDagSourceServiceGetDagSource = ( export const prefetchUseDagStatsServiceGetDagStats = ( queryClient: QueryClient, { + dagId, dagIds, }: { + dagId?: string; dagIds?: string[]; } = {}, ) => queryClient.prefetchQuery({ - queryKey: Common.UseDagStatsServiceGetDagStatsKeyFn({ dagIds }), - queryFn: () => DagStatsService.getDagStats({ dagIds }), + queryKey: Common.UseDagStatsServiceGetDagStatsKeyFn({ dagId, dagIds }), + queryFn: () => DagStatsService.getDagStats({ dagId, dagIds }), }); /** * Get Dag Reports diff --git a/airflow/ui/openapi-gen/queries/queries.ts b/airflow/ui/openapi-gen/queries/queries.ts index f6cebbf884ebe..853a53f26ed38 100644 --- a/airflow/ui/openapi-gen/queries/queries.ts +++ b/airflow/ui/openapi-gen/queries/queries.ts @@ -1115,6 +1115,7 @@ export const useDagSourceServiceGetDagSource = < * Get Dag Stats * Get Dag statistics. * @param data The data for the request. + * @param data.dagId * @param data.dagIds * @returns DagStatsCollectionResponse Successful Response * @throws ApiError @@ -1125,16 +1126,18 @@ export const useDagStatsServiceGetDagStats = < TQueryKey extends Array = unknown[], >( { + dagId, dagIds, }: { + dagId?: string; dagIds?: string[]; } = {}, queryKey?: TQueryKey, options?: Omit, "queryKey" | "queryFn">, ) => useQuery({ - queryKey: Common.UseDagStatsServiceGetDagStatsKeyFn({ dagIds }, queryKey), - queryFn: () => DagStatsService.getDagStats({ dagIds }) as TData, + queryKey: Common.UseDagStatsServiceGetDagStatsKeyFn({ dagId, dagIds }, queryKey), + queryFn: () => DagStatsService.getDagStats({ dagId, dagIds }) as TData, ...options, }); /** diff --git a/airflow/ui/openapi-gen/queries/suspense.ts b/airflow/ui/openapi-gen/queries/suspense.ts index dd2c5fe4fcca9..b8c9689b4123e 100644 --- a/airflow/ui/openapi-gen/queries/suspense.ts +++ b/airflow/ui/openapi-gen/queries/suspense.ts @@ -1092,6 +1092,7 @@ export const useDagSourceServiceGetDagSourceSuspense = < * Get Dag Stats * Get Dag statistics. * @param data The data for the request. + * @param data.dagId * @param data.dagIds * @returns DagStatsCollectionResponse Successful Response * @throws ApiError @@ -1102,16 +1103,18 @@ export const useDagStatsServiceGetDagStatsSuspense = < TQueryKey extends Array = unknown[], >( { + dagId, dagIds, }: { + dagId?: string; dagIds?: string[]; } = {}, queryKey?: TQueryKey, options?: Omit, "queryKey" | "queryFn">, ) => useSuspenseQuery({ - queryKey: Common.UseDagStatsServiceGetDagStatsKeyFn({ dagIds }, queryKey), - queryFn: () => DagStatsService.getDagStats({ dagIds }) as TData, + queryKey: Common.UseDagStatsServiceGetDagStatsKeyFn({ dagId, dagIds }, queryKey), + queryFn: () => DagStatsService.getDagStats({ dagId, dagIds }) as TData, ...options, }); /** diff --git a/airflow/ui/openapi-gen/requests/services.gen.ts b/airflow/ui/openapi-gen/requests/services.gen.ts index 3acae741015da..d0cd4ede930c5 100644 --- a/airflow/ui/openapi-gen/requests/services.gen.ts +++ b/airflow/ui/openapi-gen/requests/services.gen.ts @@ -1579,6 +1579,7 @@ export class DagStatsService { * Get Dag Stats * Get Dag statistics. * @param data The data for the request. + * @param data.dagId * @param data.dagIds * @returns DagStatsCollectionResponse Successful Response * @throws ApiError @@ -1588,6 +1589,7 @@ export class DagStatsService { method: "GET", url: "/public/dagStats", query: { + dag_id: data.dagId, dag_ids: data.dagIds, }, errors: { diff --git a/airflow/ui/openapi-gen/requests/types.gen.ts b/airflow/ui/openapi-gen/requests/types.gen.ts index d04524adb8fa0..b6f27d782a2d7 100644 --- a/airflow/ui/openapi-gen/requests/types.gen.ts +++ b/airflow/ui/openapi-gen/requests/types.gen.ts @@ -1978,6 +1978,7 @@ export type GetDagSourceData = { export type GetDagSourceResponse = DAGSourceResponse; export type GetDagStatsData = { + dagId?: string | null; dagIds?: Array; }; diff --git a/tests/api_fastapi/core_api/routes/public/test_dag_stats.py b/tests/api_fastapi/core_api/routes/public/test_dag_stats.py index 8a0dae3604c1d..e7264bd48d145 100644 --- a/tests/api_fastapi/core_api/routes/public/test_dag_stats.py +++ b/tests/api_fastapi/core_api/routes/public/test_dag_stats.py @@ -129,7 +129,7 @@ def teardown_method(self) -> None: class TestGetDagStats(TestDagStatsEndpoint): """Unit tests for Get DAG Stats.""" - def test_should_respond_200(self, client, session): + def test_should_respond_200(self, test_client, session): self._create_dag_and_runs(session) exp_payload = { "dags": [ @@ -179,13 +179,21 @@ def test_should_respond_200(self, client, session): "total_entries": 2, } - response = client().get(f"{API_PREFIX}?dag_ids={DAG1_ID}&dag_ids={DAG2_ID}") + response = test_client.get(f"{API_PREFIX}?dag_ids={DAG1_ID}&dag_ids={DAG2_ID}") assert response.status_code == 200 res_json = response.json() assert res_json["total_entries"] == len(res_json["dags"]) assert res_json == exp_payload - def test_all_dags_should_respond_200(self, client, session): + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.get(f"{API_PREFIX}?dag_ids={DAG1_ID}&dag_ids={DAG2_ID}") + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.get(f"{API_PREFIX}?dag_ids={DAG1_ID}&dag_ids={DAG2_ID}") + assert response.status_code == 403 + + def test_all_dags_should_respond_200(self, test_client, session): self._create_dag_and_runs(session) exp_payload = { "dags": [ @@ -256,7 +264,7 @@ def test_all_dags_should_respond_200(self, client, session): "total_entries": 3, } - response = client().get(API_PREFIX) + response = test_client.get(API_PREFIX) assert response.status_code == 200 res_json = response.json() assert res_json["total_entries"] == len(res_json["dags"]) @@ -403,9 +411,9 @@ def test_all_dags_should_respond_200(self, client, session): ), ], ) - def test_single_dag_in_dag_ids(self, client, session, url, params, exp_payload): + def test_single_dag_in_dag_ids(self, test_client, session, url, params, exp_payload): self._create_dag_and_runs(session) - response = client().get(url, params=params) + response = test_client.get(url, params=params) assert response.status_code == 200 res_json = response.json() assert res_json["total_entries"] == len(res_json["dags"]) From 217a6e7fa08ba48843d9968678f34c03520e270a Mon Sep 17 00:00:00 2001 From: kalyanr Date: Sat, 1 Mar 2025 23:18:08 +0530 Subject: [PATCH 21/27] add dag_warning --- airflow/api_fastapi/core_api/openapi/v1-generated.yaml | 2 ++ airflow/api_fastapi/core_api/routes/public/dag_warning.py | 3 +++ .../core_api/routes/public/test_dag_warning.py | 8 ++++++++ 3 files changed, 13 insertions(+) diff --git a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml index 3e4908ba1f57d..50985f3db37ca 100644 --- a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml +++ b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml @@ -3034,6 +3034,8 @@ paths: summary: List Dag Warnings description: Get a list of DAG warnings. operationId: list_dag_warnings + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: query diff --git a/airflow/api_fastapi/core_api/routes/public/dag_warning.py b/airflow/api_fastapi/core_api/routes/public/dag_warning.py index 2c964efce7e45..69b3354a26783 100644 --- a/airflow/api_fastapi/core_api/routes/public/dag_warning.py +++ b/airflow/api_fastapi/core_api/routes/public/dag_warning.py @@ -22,6 +22,7 @@ from fastapi import Depends from sqlalchemy import select +from airflow.api_fastapi.auth.managers.models.resource_details import DagAccessEntity from airflow.api_fastapi.common.db.common import ( SessionDep, paginated_select, @@ -37,6 +38,7 @@ from airflow.api_fastapi.core_api.datamodels.dag_warning import ( DAGWarningCollectionResponse, ) +from airflow.api_fastapi.core_api.security import requires_access_dag from airflow.models.dagwarning import DagWarning, DagWarningType dag_warning_router = AirflowRouter(tags=["DagWarning"]) @@ -44,6 +46,7 @@ @dag_warning_router.get( "/dagWarnings", + dependencies=[Depends(requires_access_dag(method="GET", access_entity=DagAccessEntity.WARNING))], ) def list_dag_warnings( dag_id: Annotated[FilterParam[str | None], Depends(filter_param_factory(DagWarning.dag_id, str | None))], diff --git a/tests/api_fastapi/core_api/routes/public/test_dag_warning.py b/tests/api_fastapi/core_api/routes/public/test_dag_warning.py index 61237bd10299a..c3dfc304cd8a3 100644 --- a/tests/api_fastapi/core_api/routes/public/test_dag_warning.py +++ b/tests/api_fastapi/core_api/routes/public/test_dag_warning.py @@ -78,6 +78,14 @@ def test_get_dag_warnings(self, test_client, query_params, expected_total_entrie assert len(response_json["dag_warnings"]) == len(expected_messages) assert [dag_warning["message"] for dag_warning in response_json["dag_warnings"]] == expected_messages + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.get("/public/dagWarnings", params={}) + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.get("/public/dagWarnings", params={}) + assert response.status_code == 403 + def test_get_dag_warnings_bad_request(self, test_client): response = test_client.get("/public/dagWarnings", params={"warning_type": "invalid"}) response_json = response.json() From 3fdafa39fe4064b82e32363b7aae106da100a6b9 Mon Sep 17 00:00:00 2001 From: kalyanr Date: Sun, 2 Mar 2025 00:07:55 +0530 Subject: [PATCH 22/27] add ti --- .../core_api/openapi/v1-generated.yaml | 97 ++++++++++-- .../core_api/routes/public/task_instances.py | 34 +++- airflow/ui/openapi-gen/queries/queries.ts | 4 +- airflow/ui/openapi-gen/requests/types.gen.ts | 32 ++-- .../routes/public/test_task_instances.py | 146 ++++++++++++++++++ 5 files changed, 274 insertions(+), 39 deletions(-) diff --git a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml index 50985f3db37ca..1443a15b35df7 100644 --- a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml +++ b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml @@ -4899,12 +4899,16 @@ paths: summary: Get Task Instance description: Get task instance. operationId: get_task_instance + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id - name: dag_run_id in: path @@ -4955,12 +4959,16 @@ paths: summary: Patch Task Instance description: Update a task instance. operationId: patch_task_instance + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id - name: dag_run_id in: path @@ -5047,12 +5055,16 @@ paths: summary: Get Mapped Task Instances description: Get list of mapped task instances. operationId: get_mapped_task_instances + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id - name: dag_run_id in: path @@ -5273,12 +5285,16 @@ paths: summary: Get Task Instance Dependencies description: Get dependencies blocking task from getting scheduled. operationId: get_task_instance_dependencies + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id - name: dag_run_id in: path @@ -5336,12 +5352,16 @@ paths: summary: Get Task Instance Dependencies description: Get dependencies blocking task from getting scheduled. operationId: get_task_instance_dependencies + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id - name: dag_run_id in: path @@ -5400,12 +5420,16 @@ paths: summary: Get Task Instance Tries description: Get list of task instances history. operationId: get_task_instance_tries + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id - name: dag_run_id in: path @@ -5463,12 +5487,16 @@ paths: - Task Instance summary: Get Mapped Task Instance Tries operationId: get_mapped_task_instance_tries + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id - name: dag_run_id in: path @@ -5526,12 +5554,16 @@ paths: summary: Get Mapped Task Instance description: Get task instance. operationId: get_mapped_task_instance + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id - name: dag_run_id in: path @@ -5588,12 +5620,16 @@ paths: summary: Patch Task Instance description: Update a task instance. operationId: patch_task_instance + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id - name: dag_run_id in: path @@ -5685,12 +5721,16 @@ paths: and DAG runs.' operationId: get_task_instances + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id - name: dag_run_id in: path @@ -5921,13 +5961,16 @@ paths: summary: Get Task Instances Batch description: Get list of task instances. operationId: get_task_instances_batch + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - const: '~' - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id - name: dag_run_id in: path @@ -5980,12 +6023,16 @@ paths: summary: Get Task Instance Try Details description: Get task instance details by try number. operationId: get_task_instance_try_details + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id - name: dag_run_id in: path @@ -6049,12 +6096,16 @@ paths: - Task Instance summary: Get Mapped Task Instance Try Details operationId: get_mapped_task_instance_try_details + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id - name: dag_run_id in: path @@ -6118,12 +6169,16 @@ paths: summary: Post Clear Task Instances description: Clear task instances. operationId: post_clear_task_instances + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id requestBody: required: true @@ -6169,12 +6224,16 @@ paths: summary: Patch Task Instance Dry Run description: Update a task instance dry_run mode. operationId: patch_task_instance_dry_run + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id - name: dag_run_id in: path @@ -6254,12 +6313,16 @@ paths: summary: Patch Task Instance Dry Run description: Update a task instance dry_run mode. operationId: patch_task_instance_dry_run + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id - name: dag_run_id in: path 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 7d130bc1e37e9..89bba0aafbc74 100644 --- a/airflow/api_fastapi/core_api/routes/public/task_instances.py +++ b/airflow/api_fastapi/core_api/routes/public/task_instances.py @@ -27,6 +27,7 @@ from sqlalchemy.orm import joinedload from sqlalchemy.sql.selectable import Select +from airflow.api_fastapi.auth.managers.models.resource_details import DagAccessEntity from airflow.api_fastapi.common.db.common import SessionDep, paginated_select from airflow.api_fastapi.common.parameters import ( FilterOptionEnum, @@ -60,6 +61,7 @@ TaskInstancesBatchBody, ) from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc +from airflow.api_fastapi.core_api.security import requires_access_dag from airflow.api_fastapi.logging.decorators import action_logging from airflow.exceptions import TaskNotFound from airflow.models import Base, DagRun @@ -78,6 +80,7 @@ @task_instances_router.get( task_instances_prefix + "/{task_id}", responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]), + dependencies=[Depends(requires_access_dag(method="GET", access_entity=DagAccessEntity.TASK_INSTANCE))], ) def get_task_instance( dag_id: str, dag_run_id: str, task_id: str, session: SessionDep @@ -108,6 +111,7 @@ def get_task_instance( @task_instances_router.get( task_instances_prefix + "/{task_id}/listMapped", responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]), + dependencies=[Depends(requires_access_dag(method="GET", access_entity=DagAccessEntity.TASK_INSTANCE))], ) def get_mapped_task_instances( dag_id: str, @@ -211,10 +215,12 @@ def get_mapped_task_instances( @task_instances_router.get( task_instances_prefix + "/{task_id}/dependencies", responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]), + dependencies=[Depends(requires_access_dag(method="GET", access_entity=DagAccessEntity.TASK_INSTANCE))], ) @task_instances_router.get( task_instances_prefix + "/{task_id}/{map_index}/dependencies", responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]), + dependencies=[Depends(requires_access_dag(method="GET", access_entity=DagAccessEntity.TASK_INSTANCE))], ) def get_task_instance_dependencies( dag_id: str, @@ -265,6 +271,7 @@ def get_task_instance_dependencies( @task_instances_router.get( task_instances_prefix + "/{task_id}/tries", responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]), + dependencies=[Depends(requires_access_dag(method="GET", access_entity=DagAccessEntity.TASK_INSTANCE))], ) def get_task_instance_tries( dag_id: str, @@ -308,6 +315,7 @@ 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]), + dependencies=[Depends(requires_access_dag(method="GET", access_entity=DagAccessEntity.TASK_INSTANCE))], ) def get_mapped_task_instance_tries( dag_id: str, @@ -328,6 +336,7 @@ def get_mapped_task_instance_tries( @task_instances_router.get( task_instances_prefix + "/{task_id}/{map_index}", responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]), + dependencies=[Depends(requires_access_dag(method="GET", access_entity=DagAccessEntity.TASK_INSTANCE))], ) def get_mapped_task_instance( dag_id: str, @@ -358,6 +367,7 @@ def get_mapped_task_instance( @task_instances_router.get( task_instances_prefix, responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]), + dependencies=[Depends(requires_access_dag(method="GET", access_entity=DagAccessEntity.TASK_INSTANCE))], ) def get_task_instances( dag_id: str, @@ -464,7 +474,10 @@ def get_task_instances( @task_instances_router.post( task_instances_prefix + "/list", responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]), - dependencies=[Depends(action_logging())], + dependencies=[ + Depends(action_logging()), + Depends(requires_access_dag(method="GET", access_entity=DagAccessEntity.TASK_INSTANCE)), + ], ) def get_task_instances_batch( dag_id: Literal["~"], @@ -546,6 +559,7 @@ def get_task_instances_batch( @task_instances_router.get( task_instances_prefix + "/{task_id}/tries/{task_try_number}", responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]), + dependencies=[Depends(requires_access_dag(method="GET", access_entity=DagAccessEntity.TASK_INSTANCE))], ) def get_task_instance_try_details( dag_id: str, @@ -581,6 +595,7 @@ def _query(orm_object: Base) -> TI | TIH | None: @task_instances_router.get( task_instances_prefix + "/{task_id}/{map_index}/tries/{task_try_number}", responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]), + dependencies=[Depends(requires_access_dag(method="GET", access_entity=DagAccessEntity.TASK_INSTANCE))], ) def get_mapped_task_instance_try_details( dag_id: str, @@ -603,7 +618,10 @@ def get_mapped_task_instance_try_details( @task_instances_router.post( "/clearTaskInstances", responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]), - dependencies=[Depends(action_logging())], + dependencies=[ + Depends(action_logging()), + Depends(requires_access_dag(method="PUT", access_entity=DagAccessEntity.TASK_INSTANCE)), + ], ) def post_clear_task_instances( dag_id: str, @@ -748,12 +766,14 @@ def _patch_ti_validate_request( responses=create_openapi_http_exception_doc( [status.HTTP_404_NOT_FOUND, status.HTTP_400_BAD_REQUEST], ), + dependencies=[Depends(requires_access_dag(method="PUT", access_entity=DagAccessEntity.TASK_INSTANCE))], ) @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], ), + dependencies=[Depends(requires_access_dag(method="PUT", access_entity=DagAccessEntity.TASK_INSTANCE))], ) def patch_task_instance_dry_run( dag_id: str, @@ -808,14 +828,20 @@ def patch_task_instance_dry_run( responses=create_openapi_http_exception_doc( [status.HTTP_404_NOT_FOUND, status.HTTP_400_BAD_REQUEST, status.HTTP_409_CONFLICT], ), - dependencies=[Depends(action_logging())], + dependencies=[ + Depends(action_logging()), + Depends(requires_access_dag(method="PUT", access_entity=DagAccessEntity.TASK_INSTANCE)), + ], ) @task_instances_router.patch( task_instances_prefix + "/{task_id}/{map_index}", responses=create_openapi_http_exception_doc( [status.HTTP_404_NOT_FOUND, status.HTTP_400_BAD_REQUEST, status.HTTP_409_CONFLICT], ), - dependencies=[Depends(action_logging())], + dependencies=[ + Depends(action_logging()), + Depends(requires_access_dag(method="PUT", access_entity=DagAccessEntity.TASK_INSTANCE)), + ], ) def patch_task_instance( dag_id: str, diff --git a/airflow/ui/openapi-gen/queries/queries.ts b/airflow/ui/openapi-gen/queries/queries.ts index 853a53f26ed38..4e6f852aba7d9 100644 --- a/airflow/ui/openapi-gen/queries/queries.ts +++ b/airflow/ui/openapi-gen/queries/queries.ts @@ -3300,7 +3300,7 @@ export const useTaskInstanceServiceGetTaskInstancesBatch = < TData, TError, { - dagId: "~"; + dagId: string; dagRunId: "~"; requestBody: TaskInstancesBatchBody; }, @@ -3313,7 +3313,7 @@ export const useTaskInstanceServiceGetTaskInstancesBatch = < TData, TError, { - dagId: "~"; + dagId: string; dagRunId: "~"; requestBody: TaskInstancesBatchBody; }, diff --git a/airflow/ui/openapi-gen/requests/types.gen.ts b/airflow/ui/openapi-gen/requests/types.gen.ts index b6f27d782a2d7..b7fac3a104e34 100644 --- a/airflow/ui/openapi-gen/requests/types.gen.ts +++ b/airflow/ui/openapi-gen/requests/types.gen.ts @@ -2107,7 +2107,7 @@ export type GetExtraLinksData = { export type GetExtraLinksResponse = ExtraLinksResponse; export type GetTaskInstanceData = { - dagId: string; + dagId: string | null; dagRunId: string; taskId: string; }; @@ -2115,7 +2115,7 @@ export type GetTaskInstanceData = { export type GetTaskInstanceResponse = TaskInstanceResponse; export type PatchTaskInstanceData = { - dagId: string; + dagId: string | null; dagRunId: string; mapIndex?: number; requestBody: PatchTaskInstanceBody; @@ -2126,7 +2126,7 @@ export type PatchTaskInstanceData = { export type PatchTaskInstanceResponse = TaskInstanceResponse; export type GetMappedTaskInstancesData = { - dagId: string; + dagId: string | null; dagRunId: string; durationGte?: number | null; durationLte?: number | null; @@ -2154,7 +2154,7 @@ export type GetMappedTaskInstancesData = { export type GetMappedTaskInstancesResponse = TaskInstanceCollectionResponse; export type GetTaskInstanceDependenciesData = { - dagId: string; + dagId: string | null; dagRunId: string; mapIndex: number; taskId: string; @@ -2163,7 +2163,7 @@ export type GetTaskInstanceDependenciesData = { export type GetTaskInstanceDependenciesResponse = TaskDependencyCollectionResponse; export type GetTaskInstanceDependencies1Data = { - dagId: string; + dagId: string | null; dagRunId: string; mapIndex?: number; taskId: string; @@ -2172,7 +2172,7 @@ export type GetTaskInstanceDependencies1Data = { export type GetTaskInstanceDependencies1Response = TaskDependencyCollectionResponse; export type GetTaskInstanceTriesData = { - dagId: string; + dagId: string | null; dagRunId: string; mapIndex?: number; taskId: string; @@ -2181,7 +2181,7 @@ export type GetTaskInstanceTriesData = { export type GetTaskInstanceTriesResponse = TaskInstanceHistoryCollectionResponse; export type GetMappedTaskInstanceTriesData = { - dagId: string; + dagId: string | null; dagRunId: string; mapIndex: number; taskId: string; @@ -2190,7 +2190,7 @@ export type GetMappedTaskInstanceTriesData = { export type GetMappedTaskInstanceTriesResponse = TaskInstanceHistoryCollectionResponse; export type GetMappedTaskInstanceData = { - dagId: string; + dagId: string | null; dagRunId: string; mapIndex: number; taskId: string; @@ -2199,7 +2199,7 @@ export type GetMappedTaskInstanceData = { export type GetMappedTaskInstanceResponse = TaskInstanceResponse; export type PatchTaskInstance1Data = { - dagId: string; + dagId: string | null; dagRunId: string; mapIndex: number; requestBody: PatchTaskInstanceBody; @@ -2210,7 +2210,7 @@ export type PatchTaskInstance1Data = { export type PatchTaskInstance1Response = TaskInstanceResponse; export type GetTaskInstancesData = { - dagId: string; + dagId: string | null; dagRunId: string; durationGte?: number | null; durationLte?: number | null; @@ -2239,7 +2239,7 @@ export type GetTaskInstancesData = { export type GetTaskInstancesResponse = TaskInstanceCollectionResponse; export type GetTaskInstancesBatchData = { - dagId: "~"; + dagId: string | null; dagRunId: "~"; requestBody: TaskInstancesBatchBody; }; @@ -2247,7 +2247,7 @@ export type GetTaskInstancesBatchData = { export type GetTaskInstancesBatchResponse = TaskInstanceCollectionResponse; export type GetTaskInstanceTryDetailsData = { - dagId: string; + dagId: string | null; dagRunId: string; mapIndex?: number; taskId: string; @@ -2257,7 +2257,7 @@ export type GetTaskInstanceTryDetailsData = { export type GetTaskInstanceTryDetailsResponse = TaskInstanceHistoryResponse; export type GetMappedTaskInstanceTryDetailsData = { - dagId: string; + dagId: string | null; dagRunId: string; mapIndex: number; taskId: string; @@ -2267,14 +2267,14 @@ export type GetMappedTaskInstanceTryDetailsData = { export type GetMappedTaskInstanceTryDetailsResponse = TaskInstanceHistoryResponse; export type PostClearTaskInstancesData = { - dagId: string; + dagId: string | null; requestBody: ClearTaskInstancesBody; }; export type PostClearTaskInstancesResponse = TaskInstanceCollectionResponse; export type PatchTaskInstanceDryRunData = { - dagId: string; + dagId: string | null; dagRunId: string; mapIndex: number; requestBody: PatchTaskInstanceBody; @@ -2285,7 +2285,7 @@ export type PatchTaskInstanceDryRunData = { export type PatchTaskInstanceDryRunResponse = TaskInstanceCollectionResponse; export type PatchTaskInstanceDryRun1Data = { - dagId: string; + dagId: string | null; dagRunId: string; mapIndex?: number; requestBody: PatchTaskInstanceBody; 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 d7e185081c9a6..d91b7de4a1b33 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 @@ -207,6 +207,18 @@ def test_should_respond_200(self, test_client, session): "triggerer_job": None, } + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.get( + "/public/dags/example_python_operator/dagRuns/TEST_DAG_RUN_ID/taskInstances/print_the_context" + ) + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.get( + "/public/dags/example_python_operator/dagRuns/TEST_DAG_RUN_ID/taskInstances/print_the_context" + ) + assert response.status_code == 403 + @pytest.mark.parametrize( "run_id, expected_version_number", [ @@ -525,6 +537,18 @@ def test_should_respond_200_mapped_task_instance_with_rtif(self, test_client, se "triggerer_job": None, } + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.get( + "/public/dags/example_python_operator/dagRuns/TEST_DAG_RUN_ID/taskInstances/print_the_context/1", + ) + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.get( + "/public/dags/example_python_operator/dagRuns/TEST_DAG_RUN_ID/taskInstances/print_the_context/1", + ) + assert response.status_code == 403 + def test_should_respond_404_wrong_map_index(self, test_client, session): self.create_task_instances(session) @@ -665,6 +689,18 @@ def one_task_with_zero_mapped_tis(self, dag_maker, session): }, ) + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.get( + "/public/dags/mapped_tis/dagRuns/run_mapped_tis/taskInstances/task_2/listMapped", + ) + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.get( + "/public/dags/mapped_tis/dagRuns/run_mapped_tis/taskInstances/task_2/listMapped", + ) + assert response.status_code == 403 + def test_should_respond_404(self, test_client): response = test_client.get( "/public/dags/mapped_tis/dagRuns/run_mapped_tis/taskInstances/task_2/listMapped", @@ -1068,6 +1104,18 @@ def test_should_respond_200( assert response.json()["total_entries"] == expected_ti assert len(response.json()["task_instances"]) == expected_ti + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.get( + "/public/dags/example_python_operator/dagRuns/~/taskInstances", + ) + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.get( + "/public/dags/example_python_operator/dagRuns/~/taskInstances", + ) + assert response.status_code == 403 + def test_not_found(self, test_client): response = test_client.get("/public/dags/invalid/dagRuns/~/taskInstances") assert response.status_code == 404 @@ -1294,6 +1342,20 @@ def test_should_respond_dependencies_mapped(self, test_client, session): ) assert response.status_code == 200, response.text + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.get( + "/public/dags/example_python_operator/dagRuns/TEST_DAG_RUN_ID/taskInstances/" + "print_the_context/0/dependencies", + ) + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.get( + "/public/dags/example_python_operator/dagRuns/TEST_DAG_RUN_ID/taskInstances/" + "print_the_context/0/dependencies", + ) + assert response.status_code == 403 + class TestGetTaskInstancesBatch(TestTaskInstanceEndpoint): @pytest.mark.parametrize( @@ -1508,6 +1570,20 @@ def test_should_raise_400_for_no_json(self, test_client): }, ] + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.post( + "/public/dags/~/dagRuns/~/taskInstances/list", + json={}, + ) + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.post( + "/public/dags/~/dagRuns/~/taskInstances/list", + json={}, + ) + assert response.status_code == 403 + def test_should_respond_422_for_non_wildcard_path_parameters(self, test_client): response = test_client.post( "/public/dags/non_wildcard/dagRuns/~/taskInstances/list", @@ -1820,6 +1896,18 @@ def test_should_respond_200_with_task_state_in_removed(self, test_client, sessio "dag_version": None, } + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.get( + "/public/dags/example_python_operator/dagRuns/TEST_DAG_RUN_ID/taskInstances/print_the_context/tries/1", + ) + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.get( + "/public/dags/example_python_operator/dagRuns/TEST_DAG_RUN_ID/taskInstances/print_the_context/tries/1", + ) + assert response.status_code == 403 + def test_raises_404_for_nonexistent_task_instance(self, test_client, session): self.create_task_instances(session) response = test_client.get( @@ -2090,6 +2178,20 @@ def test_dag_run_with_future_or_past_flag_returns_400(self, test_client, session in response.json()["detail"] ) + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.post( + "/public/dags/dag_id/clearTaskInstances", + json={}, + ) + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.post( + "/public/dags/dag_id/clearTaskInstances", + json={}, + ) + assert response.status_code == 403 + @pytest.mark.parametrize( "main_dag, task_instances, request_dag, payload, expected_ti", [ @@ -2714,6 +2816,18 @@ def test_should_respond_200(self, test_client, session): "total_entries": 2, } + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.get( + "/public/dags/example_python_operator/dagRuns/TEST_DAG_RUN_ID/taskInstances/print_the_context/tries" + ) + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.get( + "/public/dags/example_python_operator/dagRuns/TEST_DAG_RUN_ID/taskInstances/print_the_context/tries" + ) + assert response.status_code == 403 + def test_ti_in_retry_state_not_returned(self, test_client, session): self.create_task_instances( session=session, task_instances=[{"state": State.SUCCESS}], with_ti_history=True @@ -3029,6 +3143,24 @@ def test_should_update_mapped_task_instance_state(self, test_client, session): assert response2.status_code == 200 assert response2.json()["state"] == self.NEW_STATE + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.patch( + self.ENDPOINT_URL, + json={ + "new_state": self.NEW_STATE, + }, + ) + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.patch( + self.ENDPOINT_URL, + json={ + "new_state": self.NEW_STATE, + }, + ) + assert response.status_code == 403 + @pytest.mark.parametrize( "error, code, payload", [ @@ -3529,6 +3661,20 @@ def test_should_not_update(self, test_client, session, payload): assert task_before == task_after + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.patch( + f"{self.ENDPOINT_URL}/dry_run", + json={}, + ) + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.patch( + f"{self.ENDPOINT_URL}/dry_run", + json={}, + ) + assert response.status_code == 403 + def test_should_not_update_mapped_task_instance(self, test_client, session): map_index = 1 tis = self.create_task_instances(session) From bfffb39d8614fbee73ec9464a0e31a166920782e Mon Sep 17 00:00:00 2001 From: kalyanr Date: Sun, 2 Mar 2025 00:38:02 +0530 Subject: [PATCH 23/27] add xcom --- .../core_api/openapi/v1-generated.yaml | 24 +++++++-- .../core_api/routes/public/xcom.py | 8 ++- airflow/ui/openapi-gen/requests/types.gen.ts | 8 +-- .../core_api/routes/public/test_xcom.py | 54 +++++++++++++++++++ 4 files changed, 85 insertions(+), 9 deletions(-) diff --git a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml index 1443a15b35df7..26ec52ad6ba38 100644 --- a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml +++ b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml @@ -4563,12 +4563,16 @@ paths: summary: Get Xcom Entry description: Get an XCom entry. operationId: get_xcom_entry + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id - name: task_id in: path @@ -4656,12 +4660,16 @@ paths: summary: Update Xcom Entry description: Update an existing XCom entry. operationId: update_xcom_entry + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id - name: task_id in: path @@ -4735,12 +4743,16 @@ paths: This endpoint allows specifying `~` as the dag_id, dag_run_id, task_id to retrieve XCom entries for all DAGs.' operationId: get_xcom_entries + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id - name: dag_run_id in: path @@ -4830,12 +4842,16 @@ paths: summary: Create Xcom Entry description: Create an XCom entry. operationId: create_xcom_entry + security: + - OAuth2PasswordBearer: [] parameters: - name: dag_id in: path required: true schema: - type: string + anyOf: + - type: string + - type: 'null' title: Dag Id - name: task_id in: path diff --git a/airflow/api_fastapi/core_api/routes/public/xcom.py b/airflow/api_fastapi/core_api/routes/public/xcom.py index 3da163f3e4033..572e7482540e2 100644 --- a/airflow/api_fastapi/core_api/routes/public/xcom.py +++ b/airflow/api_fastapi/core_api/routes/public/xcom.py @@ -19,9 +19,10 @@ import copy from typing import Annotated -from fastapi import HTTPException, Query, Request, status +from fastapi import Depends, HTTPException, Query, Request, status from sqlalchemy import and_, select +from airflow.api_fastapi.auth.managers.models.resource_details import DagAccessEntity from airflow.api_fastapi.common.db.common import SessionDep, paginated_select from airflow.api_fastapi.common.parameters import QueryLimit, QueryOffset from airflow.api_fastapi.common.router import AirflowRouter @@ -33,6 +34,7 @@ XComUpdateBody, ) from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc +from airflow.api_fastapi.core_api.security import requires_access_dag from airflow.exceptions import TaskNotFound from airflow.models import DAG, DagRun as DR, XCom from airflow.settings import conf @@ -50,6 +52,7 @@ status.HTTP_404_NOT_FOUND, ] ), + dependencies=[Depends(requires_access_dag(method="GET", access_entity=DagAccessEntity.XCOM))], ) def get_xcom_entry( dag_id: str, @@ -105,6 +108,7 @@ def get_xcom_entry( status.HTTP_404_NOT_FOUND, ] ), + dependencies=[Depends(requires_access_dag(method="GET", access_entity=DagAccessEntity.XCOM))], ) def get_xcom_entries( dag_id: str, @@ -155,6 +159,7 @@ def get_xcom_entries( status.HTTP_404_NOT_FOUND, ] ), + dependencies=[Depends(requires_access_dag(method="POST", access_entity=DagAccessEntity.XCOM))], ) def create_xcom_entry( dag_id: str, @@ -234,6 +239,7 @@ def create_xcom_entry( status.HTTP_404_NOT_FOUND, ] ), + dependencies=[Depends(requires_access_dag(method="PUT", access_entity=DagAccessEntity.XCOM))], ) def update_xcom_entry( dag_id: str, diff --git a/airflow/ui/openapi-gen/requests/types.gen.ts b/airflow/ui/openapi-gen/requests/types.gen.ts index b7fac3a104e34..6c3c5400af17d 100644 --- a/airflow/ui/openapi-gen/requests/types.gen.ts +++ b/airflow/ui/openapi-gen/requests/types.gen.ts @@ -2395,7 +2395,7 @@ export type GetProvidersData = { export type GetProvidersResponse = ProviderCollectionResponse; export type GetXcomEntryData = { - dagId: string; + dagId: string | null; dagRunId: string; deserialize?: boolean; mapIndex?: number; @@ -2407,7 +2407,7 @@ export type GetXcomEntryData = { export type GetXcomEntryResponse = XComResponseNative | XComResponseString; export type UpdateXcomEntryData = { - dagId: string; + dagId: string | null; dagRunId: string; requestBody: XComUpdateBody; taskId: string; @@ -2417,7 +2417,7 @@ export type UpdateXcomEntryData = { export type UpdateXcomEntryResponse = XComResponseNative; export type GetXcomEntriesData = { - dagId: string; + dagId: string | null; dagRunId: string; limit?: number; mapIndex?: number | null; @@ -2429,7 +2429,7 @@ export type GetXcomEntriesData = { export type GetXcomEntriesResponse = XComCollectionResponse; export type CreateXcomEntryData = { - dagId: string; + dagId: string | null; dagRunId: string; requestBody: XComCreateBody; taskId: string; diff --git a/tests/api_fastapi/core_api/routes/public/test_xcom.py b/tests/api_fastapi/core_api/routes/public/test_xcom.py index dd5a073c1baae..c665a5299ced3 100644 --- a/tests/api_fastapi/core_api/routes/public/test_xcom.py +++ b/tests/api_fastapi/core_api/routes/public/test_xcom.py @@ -149,6 +149,18 @@ def test_should_respond_200_native(self, test_client): "value": TEST_XCOM_VALUE, } + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.get( + f"/public/dags/{TEST_DAG_ID}/dagRuns/{run_id}/taskInstances/{TEST_TASK_ID}/xcomEntries/{TEST_XCOM_KEY}" + ) + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.get( + f"/public/dags/{TEST_DAG_ID}/dagRuns/{run_id}/taskInstances/{TEST_TASK_ID}/xcomEntries/{TEST_XCOM_KEY}" + ) + assert response.status_code == 403 + def test_should_raise_404_for_non_existent_xcom(self, test_client): response = test_client.get( f"/public/dags/{TEST_DAG_ID}/dagRuns/{run_id}/taskInstances/{TEST_TASK_ID}/xcomEntries/{TEST_XCOM_KEY_2}" @@ -437,6 +449,20 @@ def _create_xcom_entries(self, dag_id, run_id, logical_date, task_id, mapped_ti= map_index=map_index, ) + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.get( + "/public/dags/~/dagRuns/~/taskInstances/~/xcomEntries", + params={}, + ) + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.get( + "/public/dags/~/dagRuns/~/taskInstances/~/xcomEntries", + params={}, + ) + assert response.status_code == 403 + class TestPaginationGetXComEntries(TestXComEndpoint): @pytest.mark.parametrize( @@ -583,6 +609,20 @@ def test_create_xcom_entry( assert current_data["run_id"] == dag_run_id assert current_data["map_index"] == request_body.map_index + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.post( + "/public/dags/dag_id/dagRuns/dag_run_id/taskInstances/task_id/xcomEntries", + json={}, + ) + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.post( + "/public/dags/dag_id/dagRuns/dag_run_id/taskInstances/task_id/xcomEntries", + json={}, + ) + assert response.status_code == 403 + class TestPatchXComEntry(TestXComEndpoint): @pytest.mark.parametrize( @@ -623,3 +663,17 @@ def test_patch_xcom_entry(self, key, patch_body, expected_status, expected_detai assert response.json()["value"] == XCom.serialize_value(new_value) else: assert response.json()["detail"] == expected_detail + + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.patch( + f"/public/dags/{TEST_DAG_ID}/dagRuns/run_id/taskInstances/TEST_TASK_ID/xcomEntries/key", + json={}, + ) + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.patch( + f"/public/dags/{TEST_DAG_ID}/dagRuns/run_id/taskInstances/TEST_TASK_ID/xcomEntries/key", + json={}, + ) + assert response.status_code == 403 From 8b8923a53316588a426cb13b0776d4c7adbbb775 Mon Sep 17 00:00:00 2001 From: kalyanr Date: Tue, 11 Mar 2025 21:47:10 +0530 Subject: [PATCH 24/27] pr feedback --- airflow/api_fastapi/core_api/routes/public/dag_run.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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 eff8d729b18c4..01840f6c74ed8 100644 --- a/airflow/api_fastapi/core_api/routes/public/dag_run.py +++ b/airflow/api_fastapi/core_api/routes/public/dag_run.py @@ -229,7 +229,7 @@ def get_upstream_asset_events( "/{dag_run_id}/clear", responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]), dependencies=[ - Depends(requires_access_dag(method="POST", access_entity=DagAccessEntity.RUN)), + Depends(requires_access_dag(method="PUT", access_entity=DagAccessEntity.RUN)), Depends(action_logging()), ], ) @@ -407,7 +407,7 @@ def trigger_dag_run( @dag_run_router.post( "/list", responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]), - dependencies=[Depends(requires_access_dag(method="POST", access_entity=DagAccessEntity.RUN))], + dependencies=[Depends(requires_access_dag(method="GET", access_entity=DagAccessEntity.RUN))], ) def get_list_dag_runs_batch( dag_id: Literal["~"], body: DAGRunsBatchBody, session: SessionDep From db8ac76022509c0d763972018161a8df0dab6b6f Mon Sep 17 00:00:00 2001 From: kalyanr Date: Wed, 12 Mar 2025 12:51:19 +0530 Subject: [PATCH 25/27] use dag run filter --- .../core_api/routes/public/dag_run.py | 33 ++++++++++++++++--- airflow/api_fastapi/core_api/security.py | 21 ++++++++++-- 2 files changed, 47 insertions(+), 7 deletions(-) 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 01840f6c74ed8..00d38be210c63 100644 --- a/airflow/api_fastapi/core_api/routes/public/dag_run.py +++ b/airflow/api_fastapi/core_api/routes/public/dag_run.py @@ -60,7 +60,11 @@ TaskInstanceResponse, ) from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc -from airflow.api_fastapi.core_api.security import requires_access_asset, requires_access_dag +from airflow.api_fastapi.core_api.security import ( + ReadableDagRunsFilterDep, + requires_access_asset, + requires_access_dag, +) from airflow.api_fastapi.logging.decorators import action_logging from airflow.exceptions import ParamValidationError from airflow.listeners.listener import get_listener_manager @@ -313,6 +317,8 @@ def get_dag_runs( ).dynamic_depends(default="id") ), ], + # readable_dags_filter: ReadableDagsFilterDep, + readable_dag_runs_filter: ReadableDagRunsFilterDep, session: SessionDep, request: Request, ) -> DAGRunCollectionResponse: @@ -321,6 +327,9 @@ def get_dag_runs( This endpoint allows specifying `~` as the dag_id to retrieve Dag Runs for all DAGs. """ + # if readable_dags_filter.value is None: + # return DAGRunCollectionResponse(dag_runs=[], total_entries=0) + query = select(DagRun) if dag_id != "~": @@ -328,11 +337,24 @@ def get_dag_runs( if not dag: raise HTTPException(status.HTTP_404_NOT_FOUND, f"The DAG with dag_id: `{dag_id}` was not found") + # if dag_id not in readable_dags_filter.value: + # raise HTTPException(status.HTTP_403_FORBIDDEN, f"Access to DAG with dag_id: `{dag_id}` is forbidden") + query = query.filter(DagRun.dag_id == dag_id) + # else: + # query = query.filter(DagRun.dag_id.in_(readable_dags_filter.value)) dag_run_select, total_entries = paginated_select( statement=query, - filters=[run_after, logical_date, start_date_range, end_date_range, update_at_range, state], + filters=[ + run_after, + logical_date, + start_date_range, + end_date_range, + update_at_range, + state, + readable_dag_runs_filter, + ], order_by=order_by, offset=offset, limit=limit, @@ -410,7 +432,10 @@ def trigger_dag_run( dependencies=[Depends(requires_access_dag(method="GET", access_entity=DagAccessEntity.RUN))], ) def get_list_dag_runs_batch( - dag_id: Literal["~"], body: DAGRunsBatchBody, session: SessionDep + dag_id: Literal["~"], + body: DAGRunsBatchBody, + readable_dag_runs_filter: ReadableDagRunsFilterDep, + session: SessionDep, ) -> DAGRunCollectionResponse: """Get a list of DAG Runs.""" dag_ids = FilterParam(DagRun.dag_id, body.dag_ids, FilterOptionEnum.IN) @@ -455,7 +480,7 @@ def get_list_dag_runs_batch( base_query = select(DagRun) dag_runs_select, total_entries = paginated_select( statement=base_query, - filters=[dag_ids, logical_date, run_after, start_date, end_date, state], + filters=[dag_ids, logical_date, run_after, start_date, end_date, state, readable_dag_runs_filter], order_by=order_by, offset=offset, limit=limit, diff --git a/airflow/api_fastapi/core_api/security.py b/airflow/api_fastapi/core_api/security.py index 3d40f78523450..a81aea78386a5 100644 --- a/airflow/api_fastapi/core_api/security.py +++ b/airflow/api_fastapi/core_api/security.py @@ -40,7 +40,7 @@ ) from airflow.api_fastapi.core_api.base import OrmClause from airflow.configuration import conf -from airflow.models.dag import DagModel +from airflow.models.dag import DagModel, DagRun from airflow.utils.jwt_signer import JWTSigner, get_signing_key if TYPE_CHECKING: @@ -114,7 +114,16 @@ def to_orm(self, select: Select) -> Select: return select.where(DagModel.dag_id.in_(self.value)) -def permitted_dag_filter_factory(method: ResourceMethod) -> Callable[[Request, BaseUser], PermittedDagFilter]: +class PermittedDagRunFilter(PermittedDagFilter): + """A parameter that filters the permitted dag runs for the user.""" + + def to_orm(self, select: Select) -> Select: + return select.where(DagRun.dag_id.in_(self.value)) + + +def permitted_dag_filter_factory( + method: ResourceMethod, filter_class=PermittedDagFilter +) -> Callable[[Request, BaseUser], PermittedDagFilter]: """ Create a callable for Depends in FastAPI that returns a filter of the permitted dags for the user. @@ -128,13 +137,19 @@ def depends_permitted_dags_filter( ) -> PermittedDagFilter: auth_manager: BaseAuthManager = request.app.state.auth_manager permitted_dags: set[str] = auth_manager.get_permitted_dag_ids(user=user, method=method) - return PermittedDagFilter(permitted_dags) + return filter_class(permitted_dags) return depends_permitted_dags_filter EditableDagsFilterDep = Annotated[PermittedDagFilter, Depends(permitted_dag_filter_factory("PUT"))] ReadableDagsFilterDep = Annotated[PermittedDagFilter, Depends(permitted_dag_filter_factory("GET"))] +ReadableDagRunsFilterDep = Annotated[ + PermittedDagRunFilter, Depends(permitted_dag_filter_factory("GET", PermittedDagRunFilter)) +] +EditableDagRunsFilterDep = Annotated[ + PermittedDagRunFilter, Depends(permitted_dag_filter_factory("PUT", PermittedDagRunFilter)) +] def requires_access_pool(method: ResourceMethod) -> Callable[[Request, BaseUser], None]: From 1f0fa5d8389862adfd5d2efe7418f0938e75bc0e Mon Sep 17 00:00:00 2001 From: kalyanr Date: Wed, 12 Mar 2025 17:55:15 +0530 Subject: [PATCH 26/27] ci failures --- .../core_api/openapi/v1-generated.yaml | 133 ++++-------------- airflow/ui/openapi-gen/queries/common.ts | 4 +- airflow/ui/openapi-gen/queries/prefetch.ts | 7 +- airflow/ui/openapi-gen/queries/queries.ts | 26 ++-- airflow/ui/openapi-gen/queries/suspense.ts | 7 +- .../ui/openapi-gen/requests/services.gen.ts | 6 - airflow/ui/openapi-gen/requests/types.gen.ts | 60 ++++---- 7 files changed, 74 insertions(+), 169 deletions(-) diff --git a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml index 26ec52ad6ba38..35230f3417559 100644 --- a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml +++ b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml @@ -2088,9 +2088,7 @@ paths: in: path required: true schema: - anyOf: - - type: string - - type: 'null' + type: string title: Dag Id - name: dag_run_id in: path @@ -2142,9 +2140,7 @@ paths: in: path required: true schema: - anyOf: - - type: string - - type: 'null' + type: string title: Dag Id - name: dag_run_id in: path @@ -2198,9 +2194,7 @@ paths: in: path required: true schema: - anyOf: - - type: string - - type: 'null' + type: string title: Dag Id - name: dag_run_id in: path @@ -2276,9 +2270,7 @@ paths: in: path required: true schema: - anyOf: - - type: string - - type: 'null' + type: string title: Dag Id - name: dag_run_id in: path @@ -2330,9 +2322,7 @@ paths: in: path required: true schema: - anyOf: - - type: string - - type: 'null' + type: string title: Dag Id - name: dag_run_id in: path @@ -2398,9 +2388,7 @@ paths: in: path required: true schema: - anyOf: - - type: string - - type: 'null' + type: string title: Dag Id - name: limit in: query @@ -2567,9 +2555,6 @@ paths: in: path required: true schema: - anyOf: - - type: string - - type: 'null' title: Dag Id requestBody: required: true @@ -2634,9 +2619,8 @@ paths: in: path required: true schema: - anyOf: - - type: string - - type: 'null' + const: '~' + type: string title: Dag Id requestBody: required: true @@ -2689,9 +2673,7 @@ paths: in: path required: true schema: - anyOf: - - type: string - - type: 'null' + type: string title: Dag Id - name: version_number in: query @@ -2769,14 +2751,6 @@ paths: security: - OAuth2PasswordBearer: [] parameters: - - name: dag_id - in: query - required: false - schema: - anyOf: - - type: string - - type: 'null' - title: Dag Id - name: dag_ids in: query required: false @@ -4570,9 +4544,7 @@ paths: in: path required: true schema: - anyOf: - - type: string - - type: 'null' + type: string title: Dag Id - name: task_id in: path @@ -4667,9 +4639,7 @@ paths: in: path required: true schema: - anyOf: - - type: string - - type: 'null' + type: string title: Dag Id - name: task_id in: path @@ -4750,9 +4720,7 @@ paths: in: path required: true schema: - anyOf: - - type: string - - type: 'null' + type: string title: Dag Id - name: dag_run_id in: path @@ -4849,9 +4817,7 @@ paths: in: path required: true schema: - anyOf: - - type: string - - type: 'null' + type: string title: Dag Id - name: task_id in: path @@ -4922,9 +4888,7 @@ paths: in: path required: true schema: - anyOf: - - type: string - - type: 'null' + type: string title: Dag Id - name: dag_run_id in: path @@ -4982,9 +4946,7 @@ paths: in: path required: true schema: - anyOf: - - type: string - - type: 'null' + type: string title: Dag Id - name: dag_run_id in: path @@ -5078,9 +5040,7 @@ paths: in: path required: true schema: - anyOf: - - type: string - - type: 'null' + type: string title: Dag Id - name: dag_run_id in: path @@ -5308,9 +5268,7 @@ paths: in: path required: true schema: - anyOf: - - type: string - - type: 'null' + type: string title: Dag Id - name: dag_run_id in: path @@ -5375,9 +5333,7 @@ paths: in: path required: true schema: - anyOf: - - type: string - - type: 'null' + type: string title: Dag Id - name: dag_run_id in: path @@ -5443,9 +5399,7 @@ paths: in: path required: true schema: - anyOf: - - type: string - - type: 'null' + type: string title: Dag Id - name: dag_run_id in: path @@ -5510,9 +5464,7 @@ paths: in: path required: true schema: - anyOf: - - type: string - - type: 'null' + type: string title: Dag Id - name: dag_run_id in: path @@ -5577,9 +5529,7 @@ paths: in: path required: true schema: - anyOf: - - type: string - - type: 'null' + type: string title: Dag Id - name: dag_run_id in: path @@ -5643,9 +5593,7 @@ paths: in: path required: true schema: - anyOf: - - type: string - - type: 'null' + type: string title: Dag Id - name: dag_run_id in: path @@ -5744,9 +5692,7 @@ paths: in: path required: true schema: - anyOf: - - type: string - - type: 'null' + type: string title: Dag Id - name: dag_run_id in: path @@ -5984,9 +5930,8 @@ paths: in: path required: true schema: - anyOf: - - type: string - - type: 'null' + const: '~' + type: string title: Dag Id - name: dag_run_id in: path @@ -6046,9 +5991,7 @@ paths: in: path required: true schema: - anyOf: - - type: string - - type: 'null' + type: string title: Dag Id - name: dag_run_id in: path @@ -6119,9 +6062,7 @@ paths: in: path required: true schema: - anyOf: - - type: string - - type: 'null' + type: string title: Dag Id - name: dag_run_id in: path @@ -6192,9 +6133,7 @@ paths: in: path required: true schema: - anyOf: - - type: string - - type: 'null' + type: string title: Dag Id requestBody: required: true @@ -6247,9 +6186,7 @@ paths: in: path required: true schema: - anyOf: - - type: string - - type: 'null' + type: string title: Dag Id - name: dag_run_id in: path @@ -6336,9 +6273,7 @@ paths: in: path required: true schema: - anyOf: - - type: string - - type: 'null' + type: string title: Dag Id - name: dag_run_id in: path @@ -6954,14 +6889,6 @@ paths: schema: type: string title: File Token - - name: dag_id - in: query - required: false - schema: - anyOf: - - type: string - - type: 'null' - title: Dag Id responses: '201': description: Successful Response diff --git a/airflow/ui/openapi-gen/queries/common.ts b/airflow/ui/openapi-gen/queries/common.ts index 727dc9b6befdd..e16aef391704b 100644 --- a/airflow/ui/openapi-gen/queries/common.ts +++ b/airflow/ui/openapi-gen/queries/common.ts @@ -683,14 +683,12 @@ export type DagStatsServiceGetDagStatsQueryResult< export const useDagStatsServiceGetDagStatsKey = "DagStatsServiceGetDagStats"; export const UseDagStatsServiceGetDagStatsKeyFn = ( { - dagId, dagIds, }: { - dagId?: string; dagIds?: string[]; } = {}, queryKey?: Array, -) => [useDagStatsServiceGetDagStatsKey, ...(queryKey ?? [{ dagId, dagIds }])]; +) => [useDagStatsServiceGetDagStatsKey, ...(queryKey ?? [{ dagIds }])]; export type DagReportServiceGetDagReportsDefaultResponse = Awaited< ReturnType >; diff --git a/airflow/ui/openapi-gen/queries/prefetch.ts b/airflow/ui/openapi-gen/queries/prefetch.ts index f495c57194eaa..8f50495dbdabd 100644 --- a/airflow/ui/openapi-gen/queries/prefetch.ts +++ b/airflow/ui/openapi-gen/queries/prefetch.ts @@ -923,7 +923,6 @@ export const prefetchUseDagSourceServiceGetDagSource = ( * Get Dag Stats * Get Dag statistics. * @param data The data for the request. - * @param data.dagId * @param data.dagIds * @returns DagStatsCollectionResponse Successful Response * @throws ApiError @@ -931,16 +930,14 @@ export const prefetchUseDagSourceServiceGetDagSource = ( export const prefetchUseDagStatsServiceGetDagStats = ( queryClient: QueryClient, { - dagId, dagIds, }: { - dagId?: string; dagIds?: string[]; } = {}, ) => queryClient.prefetchQuery({ - queryKey: Common.UseDagStatsServiceGetDagStatsKeyFn({ dagId, dagIds }), - queryFn: () => DagStatsService.getDagStats({ dagId, dagIds }), + queryKey: Common.UseDagStatsServiceGetDagStatsKeyFn({ dagIds }), + queryFn: () => DagStatsService.getDagStats({ dagIds }), }); /** * Get Dag Reports diff --git a/airflow/ui/openapi-gen/queries/queries.ts b/airflow/ui/openapi-gen/queries/queries.ts index 4e6f852aba7d9..02b647cdfa104 100644 --- a/airflow/ui/openapi-gen/queries/queries.ts +++ b/airflow/ui/openapi-gen/queries/queries.ts @@ -1115,7 +1115,6 @@ export const useDagSourceServiceGetDagSource = < * Get Dag Stats * Get Dag statistics. * @param data The data for the request. - * @param data.dagId * @param data.dagIds * @returns DagStatsCollectionResponse Successful Response * @throws ApiError @@ -1126,18 +1125,16 @@ export const useDagStatsServiceGetDagStats = < TQueryKey extends Array = unknown[], >( { - dagId, dagIds, }: { - dagId?: string; dagIds?: string[]; } = {}, queryKey?: TQueryKey, options?: Omit, "queryKey" | "queryFn">, ) => useQuery({ - queryKey: Common.UseDagStatsServiceGetDagStatsKeyFn({ dagId, dagIds }, queryKey), - queryFn: () => DagStatsService.getDagStats({ dagId, dagIds }) as TData, + queryKey: Common.UseDagStatsServiceGetDagStatsKeyFn({ dagIds }, queryKey), + queryFn: () => DagStatsService.getDagStats({ dagIds }) as TData, ...options, }); /** @@ -3219,7 +3216,7 @@ export const useDagRunServiceTriggerDagRun = < TData, TError, { - dagId: string; + dagId: unknown; requestBody: TriggerDAGRunPostBody; }, TContext @@ -3231,7 +3228,7 @@ export const useDagRunServiceTriggerDagRun = < TData, TError, { - dagId: string; + dagId: unknown; requestBody: TriggerDAGRunPostBody; }, TContext @@ -3259,7 +3256,7 @@ export const useDagRunServiceGetListDagRunsBatch = < TData, TError, { - dagId: string; + dagId: "~"; requestBody: DAGRunsBatchBody; }, TContext @@ -3271,7 +3268,7 @@ export const useDagRunServiceGetListDagRunsBatch = < TData, TError, { - dagId: string; + dagId: "~"; requestBody: DAGRunsBatchBody; }, TContext @@ -3300,7 +3297,7 @@ export const useTaskInstanceServiceGetTaskInstancesBatch = < TData, TError, { - dagId: string; + dagId: "~"; dagRunId: "~"; requestBody: TaskInstancesBatchBody; }, @@ -3313,7 +3310,7 @@ export const useTaskInstanceServiceGetTaskInstancesBatch = < TData, TError, { - dagId: string; + dagId: "~"; dagRunId: "~"; requestBody: TaskInstancesBatchBody; }, @@ -3599,7 +3596,6 @@ export const useBackfillServiceCancelBackfill = < * Request re-parsing a DAG file. * @param data The data for the request. * @param data.fileToken - * @param data.dagId * @returns null Successful Response * @throws ApiError */ @@ -3613,7 +3609,6 @@ export const useDagParsingServiceReparseDagFile = < TData, TError, { - dagId?: string; fileToken: string; }, TContext @@ -3625,13 +3620,12 @@ export const useDagParsingServiceReparseDagFile = < TData, TError, { - dagId?: string; fileToken: string; }, TContext >({ - mutationFn: ({ dagId, fileToken }) => - DagParsingService.reparseDagFile({ dagId, fileToken }) as unknown as Promise, + mutationFn: ({ fileToken }) => + DagParsingService.reparseDagFile({ fileToken }) as unknown as Promise, ...options, }); /** diff --git a/airflow/ui/openapi-gen/queries/suspense.ts b/airflow/ui/openapi-gen/queries/suspense.ts index b8c9689b4123e..dd2c5fe4fcca9 100644 --- a/airflow/ui/openapi-gen/queries/suspense.ts +++ b/airflow/ui/openapi-gen/queries/suspense.ts @@ -1092,7 +1092,6 @@ export const useDagSourceServiceGetDagSourceSuspense = < * Get Dag Stats * Get Dag statistics. * @param data The data for the request. - * @param data.dagId * @param data.dagIds * @returns DagStatsCollectionResponse Successful Response * @throws ApiError @@ -1103,18 +1102,16 @@ export const useDagStatsServiceGetDagStatsSuspense = < TQueryKey extends Array = unknown[], >( { - dagId, dagIds, }: { - dagId?: string; dagIds?: string[]; } = {}, queryKey?: TQueryKey, options?: Omit, "queryKey" | "queryFn">, ) => useSuspenseQuery({ - queryKey: Common.UseDagStatsServiceGetDagStatsKeyFn({ dagId, dagIds }, queryKey), - queryFn: () => DagStatsService.getDagStats({ dagId, dagIds }) as TData, + queryKey: Common.UseDagStatsServiceGetDagStatsKeyFn({ dagIds }, queryKey), + queryFn: () => DagStatsService.getDagStats({ dagIds }) as TData, ...options, }); /** diff --git a/airflow/ui/openapi-gen/requests/services.gen.ts b/airflow/ui/openapi-gen/requests/services.gen.ts index d0cd4ede930c5..11b54997040d6 100644 --- a/airflow/ui/openapi-gen/requests/services.gen.ts +++ b/airflow/ui/openapi-gen/requests/services.gen.ts @@ -1579,7 +1579,6 @@ export class DagStatsService { * Get Dag Stats * Get Dag statistics. * @param data The data for the request. - * @param data.dagId * @param data.dagIds * @returns DagStatsCollectionResponse Successful Response * @throws ApiError @@ -1589,7 +1588,6 @@ export class DagStatsService { method: "GET", url: "/public/dagStats", query: { - dag_id: data.dagId, dag_ids: data.dagIds, }, errors: { @@ -3390,7 +3388,6 @@ export class DagParsingService { * Request re-parsing a DAG file. * @param data The data for the request. * @param data.fileToken - * @param data.dagId * @returns null Successful Response * @throws ApiError */ @@ -3401,9 +3398,6 @@ export class DagParsingService { path: { file_token: data.fileToken, }, - query: { - dag_id: data.dagId, - }, 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 6c3c5400af17d..1ed2a13c95916 100644 --- a/airflow/ui/openapi-gen/requests/types.gen.ts +++ b/airflow/ui/openapi-gen/requests/types.gen.ts @@ -1898,21 +1898,21 @@ export type TestConnectionResponse = ConnectionTestResponse; export type CreateDefaultConnectionsResponse = void; export type GetDagRunData = { - dagId: string | null; + dagId: string; dagRunId: string; }; export type GetDagRunResponse = DAGRunResponse; export type DeleteDagRunData = { - dagId: string | null; + dagId: string; dagRunId: string; }; export type DeleteDagRunResponse = void; export type PatchDagRunData = { - dagId: string | null; + dagId: string; dagRunId: string; requestBody: DAGRunPatchBody; updateMask?: Array | null; @@ -1921,14 +1921,14 @@ export type PatchDagRunData = { export type PatchDagRunResponse = DAGRunResponse; export type GetUpstreamAssetEventsData = { - dagId: string | null; + dagId: string; dagRunId: string; }; export type GetUpstreamAssetEventsResponse = AssetEventCollectionResponse; export type ClearDagRunData = { - dagId: string | null; + dagId: string; dagRunId: string; requestBody: DAGRunClearBody; }; @@ -1936,7 +1936,7 @@ export type ClearDagRunData = { export type ClearDagRunResponse = TaskInstanceCollectionResponse | DAGRunResponse; export type GetDagRunsData = { - dagId: string | null; + dagId: string; endDateGte?: string | null; endDateLte?: string | null; limit?: number; @@ -1956,14 +1956,14 @@ export type GetDagRunsData = { export type GetDagRunsResponse = DAGRunCollectionResponse; export type TriggerDagRunData = { - dagId: string | null; + dagId: unknown; requestBody: TriggerDAGRunPostBody; }; export type TriggerDagRunResponse = DAGRunResponse; export type GetListDagRunsBatchData = { - dagId: string | null; + dagId: "~"; requestBody: DAGRunsBatchBody; }; @@ -1971,14 +1971,13 @@ export type GetListDagRunsBatchResponse = DAGRunCollectionResponse; export type GetDagSourceData = { accept?: "application/json" | "text/plain" | "*/*"; - dagId: string | null; + dagId: string; versionNumber?: number | null; }; export type GetDagSourceResponse = DAGSourceResponse; export type GetDagStatsData = { - dagId?: string | null; dagIds?: Array; }; @@ -2107,7 +2106,7 @@ export type GetExtraLinksData = { export type GetExtraLinksResponse = ExtraLinksResponse; export type GetTaskInstanceData = { - dagId: string | null; + dagId: string; dagRunId: string; taskId: string; }; @@ -2115,7 +2114,7 @@ export type GetTaskInstanceData = { export type GetTaskInstanceResponse = TaskInstanceResponse; export type PatchTaskInstanceData = { - dagId: string | null; + dagId: string; dagRunId: string; mapIndex?: number; requestBody: PatchTaskInstanceBody; @@ -2126,7 +2125,7 @@ export type PatchTaskInstanceData = { export type PatchTaskInstanceResponse = TaskInstanceResponse; export type GetMappedTaskInstancesData = { - dagId: string | null; + dagId: string; dagRunId: string; durationGte?: number | null; durationLte?: number | null; @@ -2154,7 +2153,7 @@ export type GetMappedTaskInstancesData = { export type GetMappedTaskInstancesResponse = TaskInstanceCollectionResponse; export type GetTaskInstanceDependenciesData = { - dagId: string | null; + dagId: string; dagRunId: string; mapIndex: number; taskId: string; @@ -2163,7 +2162,7 @@ export type GetTaskInstanceDependenciesData = { export type GetTaskInstanceDependenciesResponse = TaskDependencyCollectionResponse; export type GetTaskInstanceDependencies1Data = { - dagId: string | null; + dagId: string; dagRunId: string; mapIndex?: number; taskId: string; @@ -2172,7 +2171,7 @@ export type GetTaskInstanceDependencies1Data = { export type GetTaskInstanceDependencies1Response = TaskDependencyCollectionResponse; export type GetTaskInstanceTriesData = { - dagId: string | null; + dagId: string; dagRunId: string; mapIndex?: number; taskId: string; @@ -2181,7 +2180,7 @@ export type GetTaskInstanceTriesData = { export type GetTaskInstanceTriesResponse = TaskInstanceHistoryCollectionResponse; export type GetMappedTaskInstanceTriesData = { - dagId: string | null; + dagId: string; dagRunId: string; mapIndex: number; taskId: string; @@ -2190,7 +2189,7 @@ export type GetMappedTaskInstanceTriesData = { export type GetMappedTaskInstanceTriesResponse = TaskInstanceHistoryCollectionResponse; export type GetMappedTaskInstanceData = { - dagId: string | null; + dagId: string; dagRunId: string; mapIndex: number; taskId: string; @@ -2199,7 +2198,7 @@ export type GetMappedTaskInstanceData = { export type GetMappedTaskInstanceResponse = TaskInstanceResponse; export type PatchTaskInstance1Data = { - dagId: string | null; + dagId: string; dagRunId: string; mapIndex: number; requestBody: PatchTaskInstanceBody; @@ -2210,7 +2209,7 @@ export type PatchTaskInstance1Data = { export type PatchTaskInstance1Response = TaskInstanceResponse; export type GetTaskInstancesData = { - dagId: string | null; + dagId: string; dagRunId: string; durationGte?: number | null; durationLte?: number | null; @@ -2239,7 +2238,7 @@ export type GetTaskInstancesData = { export type GetTaskInstancesResponse = TaskInstanceCollectionResponse; export type GetTaskInstancesBatchData = { - dagId: string | null; + dagId: "~"; dagRunId: "~"; requestBody: TaskInstancesBatchBody; }; @@ -2247,7 +2246,7 @@ export type GetTaskInstancesBatchData = { export type GetTaskInstancesBatchResponse = TaskInstanceCollectionResponse; export type GetTaskInstanceTryDetailsData = { - dagId: string | null; + dagId: string; dagRunId: string; mapIndex?: number; taskId: string; @@ -2257,7 +2256,7 @@ export type GetTaskInstanceTryDetailsData = { export type GetTaskInstanceTryDetailsResponse = TaskInstanceHistoryResponse; export type GetMappedTaskInstanceTryDetailsData = { - dagId: string | null; + dagId: string; dagRunId: string; mapIndex: number; taskId: string; @@ -2267,14 +2266,14 @@ export type GetMappedTaskInstanceTryDetailsData = { export type GetMappedTaskInstanceTryDetailsResponse = TaskInstanceHistoryResponse; export type PostClearTaskInstancesData = { - dagId: string | null; + dagId: string; requestBody: ClearTaskInstancesBody; }; export type PostClearTaskInstancesResponse = TaskInstanceCollectionResponse; export type PatchTaskInstanceDryRunData = { - dagId: string | null; + dagId: string; dagRunId: string; mapIndex: number; requestBody: PatchTaskInstanceBody; @@ -2285,7 +2284,7 @@ export type PatchTaskInstanceDryRunData = { export type PatchTaskInstanceDryRunResponse = TaskInstanceCollectionResponse; export type PatchTaskInstanceDryRun1Data = { - dagId: string | null; + dagId: string; dagRunId: string; mapIndex?: number; requestBody: PatchTaskInstanceBody; @@ -2395,7 +2394,7 @@ export type GetProvidersData = { export type GetProvidersResponse = ProviderCollectionResponse; export type GetXcomEntryData = { - dagId: string | null; + dagId: string; dagRunId: string; deserialize?: boolean; mapIndex?: number; @@ -2407,7 +2406,7 @@ export type GetXcomEntryData = { export type GetXcomEntryResponse = XComResponseNative | XComResponseString; export type UpdateXcomEntryData = { - dagId: string | null; + dagId: string; dagRunId: string; requestBody: XComUpdateBody; taskId: string; @@ -2417,7 +2416,7 @@ export type UpdateXcomEntryData = { export type UpdateXcomEntryResponse = XComResponseNative; export type GetXcomEntriesData = { - dagId: string | null; + dagId: string; dagRunId: string; limit?: number; mapIndex?: number | null; @@ -2429,7 +2428,7 @@ export type GetXcomEntriesData = { export type GetXcomEntriesResponse = XComCollectionResponse; export type CreateXcomEntryData = { - dagId: string | null; + dagId: string; dagRunId: string; requestBody: XComCreateBody; taskId: string; @@ -2493,7 +2492,6 @@ export type BulkVariablesData = { export type BulkVariablesResponse = BulkResponse; export type ReparseDagFileData = { - dagId?: string | null; fileToken: string; }; From 4dcdfeaa920eb941b332fdc4e4609c5d9aec2042 Mon Sep 17 00:00:00 2001 From: kalyanr Date: Wed, 12 Mar 2025 18:50:06 +0530 Subject: [PATCH 27/27] add dag_id filters --- .../core_api/routes/public/dag_run.py | 9 ----- .../core_api/routes/public/dag_stats.py | 5 +-- .../core_api/routes/public/dag_warning.py | 5 +-- .../core_api/routes/public/task_instances.py | 6 +++- .../core_api/routes/public/xcom.py | 4 ++- airflow/api_fastapi/core_api/security.py | 34 +++++++++++++++++-- .../routes/public/test_task_instances.py | 1 - 7 files changed, 46 insertions(+), 18 deletions(-) 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 00d38be210c63..4e98f0b5aa1cf 100644 --- a/airflow/api_fastapi/core_api/routes/public/dag_run.py +++ b/airflow/api_fastapi/core_api/routes/public/dag_run.py @@ -317,7 +317,6 @@ def get_dag_runs( ).dynamic_depends(default="id") ), ], - # readable_dags_filter: ReadableDagsFilterDep, readable_dag_runs_filter: ReadableDagRunsFilterDep, session: SessionDep, request: Request, @@ -327,9 +326,6 @@ def get_dag_runs( This endpoint allows specifying `~` as the dag_id to retrieve Dag Runs for all DAGs. """ - # if readable_dags_filter.value is None: - # return DAGRunCollectionResponse(dag_runs=[], total_entries=0) - query = select(DagRun) if dag_id != "~": @@ -337,12 +333,7 @@ def get_dag_runs( if not dag: raise HTTPException(status.HTTP_404_NOT_FOUND, f"The DAG with dag_id: `{dag_id}` was not found") - # if dag_id not in readable_dags_filter.value: - # raise HTTPException(status.HTTP_403_FORBIDDEN, f"Access to DAG with dag_id: `{dag_id}` is forbidden") - query = query.filter(DagRun.dag_id == dag_id) - # else: - # query = query.filter(DagRun.dag_id.in_(readable_dags_filter.value)) dag_run_select, total_entries = paginated_select( statement=query, diff --git a/airflow/api_fastapi/core_api/routes/public/dag_stats.py b/airflow/api_fastapi/core_api/routes/public/dag_stats.py index d70ef3fe6143f..124221571e1c0 100644 --- a/airflow/api_fastapi/core_api/routes/public/dag_stats.py +++ b/airflow/api_fastapi/core_api/routes/public/dag_stats.py @@ -39,7 +39,7 @@ DagStatsStateResponse, ) from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc -from airflow.api_fastapi.core_api.security import requires_access_dag +from airflow.api_fastapi.core_api.security import ReadableDagRunsFilterDep, requires_access_dag from airflow.models.dagrun import DagRun from airflow.utils.state import DagRunState @@ -57,6 +57,7 @@ dependencies=[Depends(requires_access_dag(method="GET", access_entity=DagAccessEntity.RUN))], ) def get_dag_stats( + readable_dag_runs_filter: ReadableDagRunsFilterDep, session: SessionDep, dag_ids: Annotated[ FilterParam[list[str]], @@ -66,7 +67,7 @@ def get_dag_stats( """Get Dag statistics.""" dagruns_select, _ = paginated_select( statement=dagruns_select_with_state_count, - filters=[dag_ids], + filters=[dag_ids, readable_dag_runs_filter], session=session, return_total_entries=False, ) diff --git a/airflow/api_fastapi/core_api/routes/public/dag_warning.py b/airflow/api_fastapi/core_api/routes/public/dag_warning.py index 69b3354a26783..cf82d2c8dde11 100644 --- a/airflow/api_fastapi/core_api/routes/public/dag_warning.py +++ b/airflow/api_fastapi/core_api/routes/public/dag_warning.py @@ -38,7 +38,7 @@ from airflow.api_fastapi.core_api.datamodels.dag_warning import ( DAGWarningCollectionResponse, ) -from airflow.api_fastapi.core_api.security import requires_access_dag +from airflow.api_fastapi.core_api.security import ReadableDagWarningsFilterDep, requires_access_dag from airflow.models.dagwarning import DagWarning, DagWarningType dag_warning_router = AirflowRouter(tags=["DagWarning"]) @@ -60,12 +60,13 @@ def list_dag_warnings( SortParam, Depends(SortParam(["dag_id", "warning_type", "message", "timestamp"], DagWarning).dynamic_depends()), ], + readable_dag_warning_filter: ReadableDagWarningsFilterDep, session: SessionDep, ) -> DAGWarningCollectionResponse: """Get a list of DAG warnings.""" dag_warnings_select, total_entries = paginated_select( statement=select(DagWarning), - filters=[warning_type, dag_id], + filters=[warning_type, dag_id, readable_dag_warning_filter], order_by=order_by, offset=offset, limit=limit, 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 89bba0aafbc74..3e9cbc4d795bc 100644 --- a/airflow/api_fastapi/core_api/routes/public/task_instances.py +++ b/airflow/api_fastapi/core_api/routes/public/task_instances.py @@ -61,7 +61,7 @@ TaskInstancesBatchBody, ) from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc -from airflow.api_fastapi.core_api.security import requires_access_dag +from airflow.api_fastapi.core_api.security import ReadableTIFilterDep, requires_access_dag from airflow.api_fastapi.logging.decorators import action_logging from airflow.exceptions import TaskNotFound from airflow.models import Base, DagRun @@ -416,6 +416,7 @@ def get_task_instances( ).dynamic_depends(default="map_index") ), ], + readable_ti_filter: ReadableTIFilterDep, session: SessionDep, ) -> TaskInstanceCollectionResponse: """ @@ -457,6 +458,7 @@ def get_task_instances( task_id, task_display_name_pattern, version_number, + readable_ti_filter, ], order_by=order_by, offset=offset, @@ -483,6 +485,7 @@ def get_task_instances_batch( dag_id: Literal["~"], dag_run_id: Literal["~"], body: TaskInstancesBatchBody, + readable_ti_filter: ReadableTIFilterDep, session: SessionDep, ) -> TaskInstanceCollectionResponse: """Get list of task instances.""" @@ -538,6 +541,7 @@ def get_task_instances_batch( pool, queue, executor, + readable_ti_filter, ], order_by=order_by, offset=offset, diff --git a/airflow/api_fastapi/core_api/routes/public/xcom.py b/airflow/api_fastapi/core_api/routes/public/xcom.py index 572e7482540e2..c5a4028eef59b 100644 --- a/airflow/api_fastapi/core_api/routes/public/xcom.py +++ b/airflow/api_fastapi/core_api/routes/public/xcom.py @@ -34,7 +34,7 @@ XComUpdateBody, ) from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc -from airflow.api_fastapi.core_api.security import requires_access_dag +from airflow.api_fastapi.core_api.security import ReadableXComFilterDep, requires_access_dag from airflow.exceptions import TaskNotFound from airflow.models import DAG, DagRun as DR, XCom from airflow.settings import conf @@ -116,6 +116,7 @@ def get_xcom_entries( task_id: str, limit: QueryLimit, offset: QueryOffset, + readable_xcom_filter: ReadableXComFilterDep, session: SessionDep, xcom_key: Annotated[str | None, Query()] = None, map_index: Annotated[int | None, Query(ge=-1)] = None, @@ -141,6 +142,7 @@ def get_xcom_entries( query, total_entries = paginated_select( statement=query, + filters=[readable_xcom_filter], offset=offset, limit=limit, session=session, diff --git a/airflow/api_fastapi/core_api/security.py b/airflow/api_fastapi/core_api/security.py index a81aea78386a5..8c3d11d5425db 100644 --- a/airflow/api_fastapi/core_api/security.py +++ b/airflow/api_fastapi/core_api/security.py @@ -41,6 +41,9 @@ from airflow.api_fastapi.core_api.base import OrmClause from airflow.configuration import conf from airflow.models.dag import DagModel, DagRun +from airflow.models.dagwarning import DagWarning +from airflow.models.taskinstance import TaskInstance as TI +from airflow.models.xcom import XCom from airflow.utils.jwt_signer import JWTSigner, get_signing_key if TYPE_CHECKING: @@ -121,6 +124,27 @@ def to_orm(self, select: Select) -> Select: return select.where(DagRun.dag_id.in_(self.value)) +class PermittedDagWarningFilter(PermittedDagFilter): + """A parameter that filters the permitted dag warnings for the user.""" + + def to_orm(self, select: Select) -> Select: + return select.where(DagWarning.dag_id.in_(self.value)) + + +class PermittedTIFilter(PermittedDagFilter): + """A parameter that filters the permitted task instances for the user.""" + + def to_orm(self, select: Select) -> Select: + return select.where(TI.dag_id.in_(self.value)) + + +class PermittedXComFilter(PermittedDagFilter): + """A parameter that filters the permitted XComs for the user.""" + + def to_orm(self, select: Select) -> Select: + return select.where(XCom.dag_id.in_(self.value)) + + def permitted_dag_filter_factory( method: ResourceMethod, filter_class=PermittedDagFilter ) -> Callable[[Request, BaseUser], PermittedDagFilter]: @@ -147,8 +171,14 @@ def depends_permitted_dags_filter( ReadableDagRunsFilterDep = Annotated[ PermittedDagRunFilter, Depends(permitted_dag_filter_factory("GET", PermittedDagRunFilter)) ] -EditableDagRunsFilterDep = Annotated[ - PermittedDagRunFilter, Depends(permitted_dag_filter_factory("PUT", PermittedDagRunFilter)) +ReadableDagWarningsFilterDep = Annotated[ + PermittedDagWarningFilter, Depends(permitted_dag_filter_factory("GET", PermittedDagWarningFilter)) +] +ReadableTIFilterDep = Annotated[ + PermittedTIFilter, Depends(permitted_dag_filter_factory("GET", PermittedTIFilter)) +] +ReadableXComFilterDep = Annotated[ + PermittedXComFilter, Depends(permitted_dag_filter_factory("GET", PermittedXComFilter)) ] 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 d91b7de4a1b33..4403bbc1428ce 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 @@ -1133,7 +1133,6 @@ def test_bad_state(self, test_client): == f"Invalid value for state. Valid values are {', '.join(TaskInstanceState)}" ) - @pytest.mark.xfail(reason="permissions not implemented yet.") def test_return_TI_only_from_readable_dags(self, test_client, session): task_instances = { "example_python_operator": 1,