AIP-84 Patch Variable (apache#42929)

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(

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.

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
 

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

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

airflow/ui/openapi-gen/queries/common.ts

@@ -191,6 +191,9 @@ export type DagServicePatchDagsMutationResult = Awaited<
 export type DagServicePatchDagMutationResult = Awaited<
   ReturnType<typeof DagService.patchDag>
 >;
+export type VariableServicePatchVariableMutationResult = Awaited<
+  ReturnType<typeof VariableService.patchVariable>
+>;
 export type ConnectionServiceDeleteConnectionMutationResult = Awaited<
   ReturnType<typeof ConnectionService.deleteConnection>
 >;

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<TData>,
     ...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<TData>,
+    ...options,
+  });
 /**
  * Delete Connection
  * Delete a connection entry.

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;