From 5da28aabfd6cf0e933c43864c522fe36c9b85b1a Mon Sep 17 00:00:00 2001 From: Pierre Jeambrun Date: Mon, 14 Oct 2024 23:08:15 +0800 Subject: [PATCH] AIP-84 Patch Variable (#42929) --- .../endpoints/variable_endpoint.py | 1 + airflow/api_fastapi/openapi/v1-generated.yaml | 94 ++++++++++++++++++- airflow/api_fastapi/serializers/dags.py | 2 +- airflow/api_fastapi/serializers/variables.py | 17 +++- airflow/api_fastapi/views/public/variables.py | 28 +++++- airflow/ui/openapi-gen/queries/common.ts | 3 + airflow/ui/openapi-gen/queries/queries.ts | 49 +++++++++- .../ui/openapi-gen/requests/schemas.gen.ts | 41 +++++++- .../ui/openapi-gen/requests/services.gen.ts | 36 +++++++ airflow/ui/openapi-gen/requests/types.gen.ts | 50 +++++++++- .../views/public/test_variables.py | 88 +++++++++++++++++ 11 files changed, 394 insertions(+), 15 deletions(-) diff --git a/airflow/api_connexion/endpoints/variable_endpoint.py b/airflow/api_connexion/endpoints/variable_endpoint.py index 8efddb58419c4..b1d9e2f5c8b6f 100644 --- a/airflow/api_connexion/endpoints/variable_endpoint.py +++ b/airflow/api_connexion/endpoints/variable_endpoint.py @@ -95,6 +95,7 @@ def get_variables( ) +@mark_fastapi_migration_done @security.requires_access_variable("PUT") @provide_session @action_logging( diff --git a/airflow/api_fastapi/openapi/v1-generated.yaml b/airflow/api_fastapi/openapi/v1-generated.yaml index 6c35ca212a2a6..759ab7fdd8e9b 100644 --- a/airflow/api_fastapi/openapi/v1-generated.yaml +++ b/airflow/api_fastapi/openapi/v1-generated.yaml @@ -629,6 +629,72 @@ paths: application/json: schema: $ref: '#/components/schemas/HTTPValidationError' + patch: + tags: + - Variable + summary: Patch Variable + description: Update a variable by key. + operationId: patch_variable + parameters: + - name: variable_key + in: path + required: true + schema: + type: string + title: Variable Key + - name: update_mask + in: query + required: false + schema: + anyOf: + - type: array + items: + type: string + - type: 'null' + title: Update Mask + requestBody: + required: true + content: + application/json: + schema: + $ref: '#/components/schemas/VariableBody' + responses: + '200': + description: Successful Response + content: + application/json: + schema: + $ref: '#/components/schemas/VariableResponse' + '400': + content: + application/json: + schema: + $ref: '#/components/schemas/HTTPExceptionResponse' + description: Bad Request + '401': + content: + application/json: + schema: + $ref: '#/components/schemas/HTTPExceptionResponse' + description: Unauthorized + '403': + content: + application/json: + schema: + $ref: '#/components/schemas/HTTPExceptionResponse' + description: Forbidden + '404': + content: + application/json: + schema: + $ref: '#/components/schemas/HTTPExceptionResponse' + description: Not Found + '422': + description: Validation Error + content: + application/json: + schema: + $ref: '#/components/schemas/HTTPValidationError' /public/dags/{dag_id}/dagRuns/{dag_run_id}: get: tags: @@ -1045,7 +1111,7 @@ components: required: - is_paused title: DAGPatchBody - description: Dag Serializer for updatable body. + description: Dag Serializer for updatable bodies. DAGResponse: properties: dag_id: @@ -1492,25 +1558,47 @@ components: - msg - type title: ValidationError - VariableResponse: + VariableBody: properties: key: type: string title: Key + description: + anyOf: + - type: string + - type: 'null' + title: Description value: anyOf: - type: string - type: 'null' title: Value + type: object + required: + - key + - description + - value + title: VariableBody + description: Variable serializer for bodies. + VariableResponse: + properties: + key: + type: string + title: Key description: anyOf: - type: string - type: 'null' title: Description + value: + anyOf: + - type: string + - type: 'null' + title: Value type: object required: - key - - value - description + - value title: VariableResponse description: Variable serializer for responses. diff --git a/airflow/api_fastapi/serializers/dags.py b/airflow/api_fastapi/serializers/dags.py index 9879badf25048..c9d48aac222eb 100644 --- a/airflow/api_fastapi/serializers/dags.py +++ b/airflow/api_fastapi/serializers/dags.py @@ -95,7 +95,7 @@ def file_token(self) -> str: class DAGPatchBody(BaseModel): - """Dag Serializer for updatable body.""" + """Dag Serializer for updatable bodies.""" is_paused: bool diff --git a/airflow/api_fastapi/serializers/variables.py b/airflow/api_fastapi/serializers/variables.py index ded268432b89d..1ecc87425a24f 100644 --- a/airflow/api_fastapi/serializers/variables.py +++ b/airflow/api_fastapi/serializers/variables.py @@ -25,15 +25,20 @@ from airflow.utils.log.secrets_masker import redact -class VariableResponse(BaseModel): - """Variable serializer for responses.""" +class VariableBase(BaseModel): + """Base Variable serializer.""" model_config = ConfigDict(populate_by_name=True) key: str - val: str | None = Field(alias="value") description: str | None + +class VariableResponse(VariableBase): + """Variable serializer for responses.""" + + val: str | None = Field(alias="value") + @model_validator(mode="after") def redact_val(self) -> Self: if self.val is None: @@ -47,3 +52,9 @@ def redact_val(self) -> Self: # value is not a serialized string representation of a dict. self.val = redact(self.val, self.key) return self + + +class VariableBody(VariableBase): + """Variable serializer for bodies.""" + + value: str | None diff --git a/airflow/api_fastapi/views/public/variables.py b/airflow/api_fastapi/views/public/variables.py index e6cbb136f1cac..b4c07e23de293 100644 --- a/airflow/api_fastapi/views/public/variables.py +++ b/airflow/api_fastapi/views/public/variables.py @@ -16,14 +16,14 @@ # under the License. from __future__ import annotations -from fastapi import Depends, HTTPException +from fastapi import Depends, HTTPException, Query from sqlalchemy import select from sqlalchemy.orm import Session from typing_extensions import Annotated from airflow.api_fastapi.db.common import get_session from airflow.api_fastapi.openapi.exceptions import create_openapi_http_exception_doc -from airflow.api_fastapi.serializers.variables import VariableResponse +from airflow.api_fastapi.serializers.variables import VariableBody, VariableResponse from airflow.api_fastapi.views.router import AirflowRouter from airflow.models.variable import Variable @@ -56,3 +56,27 @@ async def get_variable( raise HTTPException(404, f"The Variable with key: `{variable_key}` was not found") return VariableResponse.model_validate(variable, from_attributes=True) + + +@variables_router.patch("/{variable_key}", responses=create_openapi_http_exception_doc([400, 401, 403, 404])) +async def patch_variable( + variable_key: str, + patch_body: VariableBody, + session: Annotated[Session, Depends(get_session)], + update_mask: list[str] | None = Query(None), +) -> VariableResponse: + """Update a variable by key.""" + if patch_body.key != variable_key: + raise HTTPException(400, "Invalid body, key from request body doesn't match uri parameter") + non_update_fields = {"key"} + variable = session.scalar(select(Variable).filter_by(key=variable_key).limit(1)) + if not variable: + raise HTTPException(404, f"The Variable with key: `{variable_key}` was not found") + if update_mask: + data = patch_body.dict(include=set(update_mask) - non_update_fields) + else: + data = patch_body.dict(exclude=non_update_fields) + for key, val in data.items(): + setattr(variable, key, val) + session.add(variable) + return variable diff --git a/airflow/ui/openapi-gen/queries/common.ts b/airflow/ui/openapi-gen/queries/common.ts index 393fad520a679..e3c0ef3ab4daf 100644 --- a/airflow/ui/openapi-gen/queries/common.ts +++ b/airflow/ui/openapi-gen/queries/common.ts @@ -191,6 +191,9 @@ export type DagServicePatchDagsMutationResult = Awaited< export type DagServicePatchDagMutationResult = Awaited< ReturnType >; +export type VariableServicePatchVariableMutationResult = Awaited< + ReturnType +>; export type ConnectionServiceDeleteConnectionMutationResult = Awaited< ReturnType >; diff --git a/airflow/ui/openapi-gen/queries/queries.ts b/airflow/ui/openapi-gen/queries/queries.ts index 01e919cce53bf..b4c8cf9fea5c9 100644 --- a/airflow/ui/openapi-gen/queries/queries.ts +++ b/airflow/ui/openapi-gen/queries/queries.ts @@ -14,7 +14,7 @@ import { DashboardService, VariableService, } from "../requests/services.gen"; -import { DAGPatchBody, DagRunState } from "../requests/types.gen"; +import { DAGPatchBody, DagRunState, VariableBody } from "../requests/types.gen"; import * as Common from "./common"; /** @@ -428,6 +428,53 @@ export const useDagServicePatchDag = < }) as unknown as Promise, ...options, }); +/** + * Patch Variable + * Update a variable by key. + * @param data The data for the request. + * @param data.variableKey + * @param data.requestBody + * @param data.updateMask + * @returns VariableResponse Successful Response + * @throws ApiError + */ +export const useVariableServicePatchVariable = < + TData = Common.VariableServicePatchVariableMutationResult, + TError = unknown, + TContext = unknown, +>( + options?: Omit< + UseMutationOptions< + TData, + TError, + { + requestBody: VariableBody; + updateMask?: string[]; + variableKey: string; + }, + TContext + >, + "mutationFn" + >, +) => + useMutation< + TData, + TError, + { + requestBody: VariableBody; + updateMask?: string[]; + variableKey: string; + }, + TContext + >({ + mutationFn: ({ requestBody, updateMask, variableKey }) => + VariableService.patchVariable({ + requestBody, + updateMask, + variableKey, + }) as unknown as Promise, + ...options, + }); /** * Delete Connection * Delete a connection entry. diff --git a/airflow/ui/openapi-gen/requests/schemas.gen.ts b/airflow/ui/openapi-gen/requests/schemas.gen.ts index 18df5284651b7..e42a3f6572ca9 100644 --- a/airflow/ui/openapi-gen/requests/schemas.gen.ts +++ b/airflow/ui/openapi-gen/requests/schemas.gen.ts @@ -526,7 +526,7 @@ export const $DAGPatchBody = { type: "object", required: ["is_paused"], title: "DAGPatchBody", - description: "Dag Serializer for updatable body.", + description: "Dag Serializer for updatable bodies.", } as const; export const $DAGResponse = { @@ -1182,12 +1182,23 @@ export const $ValidationError = { title: "ValidationError", } as const; -export const $VariableResponse = { +export const $VariableBody = { properties: { key: { type: "string", title: "Key", }, + description: { + anyOf: [ + { + type: "string", + }, + { + type: "null", + }, + ], + title: "Description", + }, value: { anyOf: [ { @@ -1199,6 +1210,19 @@ export const $VariableResponse = { ], title: "Value", }, + }, + type: "object", + required: ["key", "description", "value"], + title: "VariableBody", + description: "Variable serializer for bodies.", +} as const; + +export const $VariableResponse = { + properties: { + key: { + type: "string", + title: "Key", + }, description: { anyOf: [ { @@ -1210,9 +1234,20 @@ export const $VariableResponse = { ], title: "Description", }, + value: { + anyOf: [ + { + type: "string", + }, + { + type: "null", + }, + ], + title: "Value", + }, }, type: "object", - required: ["key", "value", "description"], + required: ["key", "description", "value"], title: "VariableResponse", description: "Variable serializer for responses.", } as const; diff --git a/airflow/ui/openapi-gen/requests/services.gen.ts b/airflow/ui/openapi-gen/requests/services.gen.ts index f0cbc099370f8..43d9e8d940df4 100644 --- a/airflow/ui/openapi-gen/requests/services.gen.ts +++ b/airflow/ui/openapi-gen/requests/services.gen.ts @@ -25,6 +25,8 @@ import type { DeleteVariableResponse, GetVariableData, GetVariableResponse, + PatchVariableData, + PatchVariableResponse, GetDagRunData, GetDagRunResponse, DeleteDagRunData, @@ -364,6 +366,40 @@ export class VariableService { }, }); } + + /** + * Patch Variable + * Update a variable by key. + * @param data The data for the request. + * @param data.variableKey + * @param data.requestBody + * @param data.updateMask + * @returns VariableResponse Successful Response + * @throws ApiError + */ + public static patchVariable( + data: PatchVariableData, + ): CancelablePromise { + return __request(OpenAPI, { + method: "PATCH", + url: "/public/variables/{variable_key}", + path: { + variable_key: data.variableKey, + }, + query: { + update_mask: data.updateMask, + }, + body: data.requestBody, + mediaType: "application/json", + errors: { + 400: "Bad Request", + 401: "Unauthorized", + 403: "Forbidden", + 404: "Not Found", + 422: "Validation Error", + }, + }); + } } export class DagRunService { diff --git a/airflow/ui/openapi-gen/requests/types.gen.ts b/airflow/ui/openapi-gen/requests/types.gen.ts index 5fe3615c7d4e6..3deba451fced4 100644 --- a/airflow/ui/openapi-gen/requests/types.gen.ts +++ b/airflow/ui/openapi-gen/requests/types.gen.ts @@ -79,7 +79,7 @@ export type DAGDetailsResponse = { }; /** - * Dag Serializer for updatable body. + * Dag Serializer for updatable bodies. */ export type DAGPatchBody = { is_paused: boolean; @@ -250,13 +250,22 @@ export type ValidationError = { type: string; }; +/** + * Variable serializer for bodies. + */ +export type VariableBody = { + key: string; + description: string | null; + value: string | null; +}; + /** * Variable serializer for responses. */ export type VariableResponse = { key: string; - value: string | null; description: string | null; + value: string | null; }; export type NextRunAssetsData = { @@ -348,6 +357,14 @@ export type GetVariableData = { export type GetVariableResponse = VariableResponse; +export type PatchVariableData = { + requestBody: VariableBody; + updateMask?: Array | null; + variableKey: string; +}; + +export type PatchVariableResponse = VariableResponse; + export type GetDagRunData = { dagId: string; dagRunId: string; @@ -637,6 +654,35 @@ export type $OpenApiTs = { 422: HTTPValidationError; }; }; + patch: { + req: PatchVariableData; + res: { + /** + * Successful Response + */ + 200: VariableResponse; + /** + * Bad Request + */ + 400: HTTPExceptionResponse; + /** + * Unauthorized + */ + 401: HTTPExceptionResponse; + /** + * Forbidden + */ + 403: HTTPExceptionResponse; + /** + * Not Found + */ + 404: HTTPExceptionResponse; + /** + * Validation Error + */ + 422: HTTPValidationError; + }; + }; }; "/public/dags/{dag_id}/dagRuns/{dag_run_id}": { get: { diff --git a/tests/api_fastapi/views/public/test_variables.py b/tests/api_fastapi/views/public/test_variables.py index 7957d0bf22249..cf5c78fd562cf 100644 --- a/tests/api_fastapi/views/public/test_variables.py +++ b/tests/api_fastapi/views/public/test_variables.py @@ -135,3 +135,91 @@ def test_get_should_respond_404(self, test_client): assert response.status_code == 404 body = response.json() assert f"The Variable with key: `{TEST_VARIABLE_KEY}` was not found" == body["detail"] + + +class TestPatchVariable(TestVariableEndpoint): + @pytest.mark.enable_redact + @pytest.mark.parametrize( + "key, body, params, expected_response", + [ + ( + TEST_VARIABLE_KEY, + { + "key": TEST_VARIABLE_KEY, + "value": "The new value", + "description": "The new description", + }, + None, + { + "key": TEST_VARIABLE_KEY, + "value": "The new value", + "description": "The new description", + }, + ), + ( + TEST_VARIABLE_KEY, + { + "key": TEST_VARIABLE_KEY, + "value": "The new value", + "description": "The new description", + }, + {"update_mask": ["value"]}, + { + "key": TEST_VARIABLE_KEY, + "value": "The new value", + "description": TEST_VARIABLE_DESCRIPTION, + }, + ), + ( + TEST_VARIABLE_KEY2, + { + "key": TEST_VARIABLE_KEY2, + "value": "some_other_value", + "description": TEST_VARIABLE_DESCRIPTION2, + }, + None, + { + "key": TEST_VARIABLE_KEY2, + "value": "***", + "description": TEST_VARIABLE_DESCRIPTION2, + }, + ), + ( + TEST_VARIABLE_KEY3, + { + "key": TEST_VARIABLE_KEY3, + "value": '{"password": "new_password"}', + "description": "new description", + }, + {"update_mask": ["value", "description"]}, + { + "key": TEST_VARIABLE_KEY3, + "value": '{"password": "***"}', + "description": "new description", + }, + ), + ], + ) + def test_patch_should_respond_200(self, test_client, session, key, body, params, expected_response): + self.create_variable() + response = test_client.patch(f"/public/variables/{key}", json=body, params=params) + assert response.status_code == 200 + assert response.json() == expected_response + + def test_patch_should_respond_400(self, test_client): + response = test_client.patch( + f"/public/variables/{TEST_VARIABLE_KEY}", + json={"key": "different_key", "value": "some_value", "description": None}, + ) + assert response.status_code == 400 + body = response.json() + assert "Invalid body, key from request body doesn't match uri parameter" == body["detail"] + + def test_patch_should_respond_404(self, test_client): + response = test_client.patch( + f"/public/variables/{TEST_VARIABLE_KEY}", + json={"key": TEST_VARIABLE_KEY, "value": "some_value", "description": None}, + ) + assert response.status_code == 404 + body = response.json() + assert f"The Variable with key: `{TEST_VARIABLE_KEY}` was not found" == body["detail"]