AIP-84 post variable (#42948)

apache · Oct 14, 2024 · 4c5ad9c · 4c5ad9c
1 parent 4122597
commit 4c5ad9c
Show file tree

Hide file tree

Showing 8 changed files with 201 additions and 2 deletions.
diff --git a/airflow/api_connexion/endpoints/variable_endpoint.py b/airflow/api_connexion/endpoints/variable_endpoint.py
@@ -130,6 +130,7 @@ def patch_variable(
     return variable_schema.dump(variable)
 
 
+@mark_fastapi_migration_done
 @security.requires_access_variable("POST")
 @action_logging(
     event=action_event_from_permission(

diff --git a/airflow/api_fastapi/openapi/v1-generated.yaml b/airflow/api_fastapi/openapi/v1-generated.yaml
@@ -695,6 +695,44 @@ paths:
             application/json:
               schema:
                 $ref: '#/components/schemas/HTTPValidationError'
+  /public/variables/:
+    post:
+      tags:
+      - Variable
+      summary: Post Variable
+      description: Create a variable.
+      operationId: post_variable
+      requestBody:
+        content:
+          application/json:
+            schema:
+              $ref: '#/components/schemas/VariableBody'
+        required: true
+      responses:
+        '201':
+          description: Successful Response
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/VariableResponse'
+        '401':
+          description: Unauthorized
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/HTTPExceptionResponse'
+        '403':
+          description: Forbidden
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/HTTPExceptionResponse'
+        '422':
+          description: Validation Error
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/HTTPValidationError'
   /public/dags/{dag_id}/dagRuns/{dag_run_id}:
     get:
       tags:

diff --git a/airflow/api_fastapi/views/public/variables.py b/airflow/api_fastapi/views/public/variables.py
@@ -73,10 +73,23 @@ async def patch_variable(
     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)
+        data = patch_body.model_dump(include=set(update_mask) - non_update_fields)
     else:
-        data = patch_body.dict(exclude=non_update_fields)
+        data = patch_body.model_dump(exclude=non_update_fields)
     for key, val in data.items():
         setattr(variable, key, val)
     session.add(variable)
     return variable
+
+
+@variables_router.post("/", status_code=201, responses=create_openapi_http_exception_doc([401, 403]))
+async def post_variable(
+    post_body: VariableBody,
+    session: Annotated[Session, Depends(get_session)],
+) -> VariableResponse:
+    """Create a variable."""
+    Variable.set(**post_body.model_dump(), session=session)
+
+    variable = session.scalar(select(Variable).where(Variable.key == post_body.key).limit(1))
+
+    return VariableResponse.model_validate(variable, from_attributes=True)
diff --git a/airflow/ui/openapi-gen/queries/common.ts b/airflow/ui/openapi-gen/queries/common.ts
@@ -185,6 +185,9 @@ export const UseDagRunServiceGetDagRunKeyFn = (
   },
   queryKey?: Array<unknown>,
 ) => [useDagRunServiceGetDagRunKey, ...(queryKey ?? [{ dagId, dagRunId }])];
+export type VariableServicePostVariableMutationResult = Awaited<
+  ReturnType<typeof VariableService.postVariable>
+>;
 export type DagServicePatchDagsMutationResult = Awaited<
   ReturnType<typeof DagService.patchDags>
 >;

diff --git a/airflow/ui/openapi-gen/queries/queries.ts b/airflow/ui/openapi-gen/queries/queries.ts
@@ -295,6 +295,45 @@ export const useDagRunServiceGetDagRun = <
     queryFn: () => DagRunService.getDagRun({ dagId, dagRunId }) as TData,
     ...options,
   });
+/**
+ * Post Variable
+ * Create a variable.
+ * @param data The data for the request.
+ * @param data.requestBody
+ * @returns VariableResponse Successful Response
+ * @throws ApiError
+ */
+export const useVariableServicePostVariable = <
+  TData = Common.VariableServicePostVariableMutationResult,
+  TError = unknown,
+  TContext = unknown,
+>(
+  options?: Omit<
+    UseMutationOptions<
+      TData,
+      TError,
+      {
+        requestBody: VariableBody;
+      },
+      TContext
+    >,
+    "mutationFn"
+  >,
+) =>
+  useMutation<
+    TData,
+    TError,
+    {
+      requestBody: VariableBody;
+    },
+    TContext
+  >({
+    mutationFn: ({ requestBody }) =>
+      VariableService.postVariable({
+        requestBody,
+      }) as unknown as Promise<TData>,
+    ...options,
+  });
 /**
  * Patch Dags
  * Patch multiple DAGs.

diff --git a/airflow/ui/openapi-gen/requests/services.gen.ts b/airflow/ui/openapi-gen/requests/services.gen.ts
@@ -27,6 +27,8 @@ import type {
   GetVariableResponse,
   PatchVariableData,
   PatchVariableResponse,
+  PostVariableData,
+  PostVariableResponse,
   GetDagRunData,
   GetDagRunResponse,
   DeleteDagRunData,
@@ -400,6 +402,30 @@ export class VariableService {
       },
     });
   }
+
+  /**
+   * Post Variable
+   * Create a variable.
+   * @param data The data for the request.
+   * @param data.requestBody
+   * @returns VariableResponse Successful Response
+   * @throws ApiError
+   */
+  public static postVariable(
+    data: PostVariableData,
+  ): CancelablePromise<PostVariableResponse> {
+    return __request(OpenAPI, {
+      method: "POST",
+      url: "/public/variables/",
+      body: data.requestBody,
+      mediaType: "application/json",
+      errors: {
+        401: "Unauthorized",
+        403: "Forbidden",
+        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
@@ -365,6 +365,12 @@ export type PatchVariableData = {
 
 export type PatchVariableResponse = VariableResponse;
 
+export type PostVariableData = {
+  requestBody: VariableBody;
+};
+
+export type PostVariableResponse = VariableResponse;
+
 export type GetDagRunData = {
   dagId: string;
   dagRunId: string;
@@ -684,6 +690,29 @@ export type $OpenApiTs = {
       };
     };
   };
+  "/public/variables/": {
+    post: {
+      req: PostVariableData;
+      res: {
+        /**
+         * Successful Response
+         */
+        201: VariableResponse;
+        /**
+         * Unauthorized
+         */
+        401: HTTPExceptionResponse;
+        /**
+         * Forbidden
+         */
+        403: HTTPExceptionResponse;
+        /**
+         * Validation Error
+         */
+        422: HTTPValidationError;
+      };
+    };
+  };
   "/public/dags/{dag_id}/dagRuns/{dag_run_id}": {
     get: {
       req: GetDagRunData;

diff --git a/tests/api_fastapi/views/public/test_variables.py b/tests/api_fastapi/views/public/test_variables.py
@@ -223,3 +223,53 @@ def test_patch_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 TestPostVariable(TestVariableEndpoint):
+    @pytest.mark.enable_redact
+    @pytest.mark.parametrize(
+        "body, expected_response",
+        [
+            (
+                {
+                    "key": "new variable key",
+                    "value": "new variable value",
+                    "description": "new variable description",
+                },
+                {
+                    "key": "new variable key",
+                    "value": "new variable value",
+                    "description": "new variable description",
+                },
+            ),
+            (
+                {
+                    "key": "another_password",
+                    "value": "password_value",
+                    "description": "another password",
+                },
+                {
+                    "key": "another_password",
+                    "value": "***",
+                    "description": "another password",
+                },
+            ),
+            (
+                {
+                    "key": "another value with sensitive information",
+                    "value": '{"password": "new_password"}',
+                    "description": "some description",
+                },
+                {
+                    "key": "another value with sensitive information",
+                    "value": '{"password": "***"}',
+                    "description": "some description",
+                },
+            ),
+        ],
+    )
+    def test_post_should_respond_201(self, test_client, session, body, expected_response):
+        self.create_variable()
+        response = test_client.post("/public/variables/", json=body)
+        assert response.status_code == 201
+        assert response.json() == expected_response