Remove unused print statement, unify bulk action types as enums and move them to common.py under datamodels

Author: bugraoz93

airflow/api_fastapi/core_api/datamodels/common.py

@@ -0,0 +1,52 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Common Data Models for Airflow REST API.
+
+:meta private:
+"""
+
+from __future__ import annotations
+
+import enum
+
+from airflow.api_fastapi.core_api.base import BaseModel
+
+
+# Common Bulk Data Models
+class BulkAction(enum.Enum):
+    """Bulk Action to be performed on the used model."""
+
+    CREATE = "create"
+    DELETE = "delete"
+    UPDATE = "update"
+
+
+class BulkActionOnExistence(enum.Enum):
+    """Bulk Action to be taken if the entity already exists or not."""
+
+    FAIL = "fail"
+    SKIP = "skip"
+    OVERWRITE = "overwrite"
+
+
+# TODO: Unify All Bulk Operation Related Base Data Models
+class BulkBaseAction(BaseModel):
+    """Base class for bulk actions."""
+
+    action: BulkAction
+    action_on_existence: BulkActionOnExistence = BulkActionOnExistence.FAIL

airflow/api_fastapi/core_api/datamodels/connections.py

@@ -18,12 +18,13 @@
 from __future__ import annotations
 
 import json
-from typing import Any, Literal
+from typing import Any
 
 from pydantic import Field, field_validator
 from pydantic_core.core_schema import ValidationInfo
 
 from airflow.api_fastapi.core_api.base import BaseModel
+from airflow.api_fastapi.core_api.datamodels.common import BulkAction, BulkBaseAction
 from airflow.utils.log.secrets_masker import redact
 
 
@@ -91,28 +92,25 @@ class ConnectionBody(BaseModel):
     extra: str | None = Field(default=None)
 
 
-class ConnectionBulkCreateAction(BaseModel):
+class ConnectionBulkCreateAction(BulkBaseAction):
     """Bulk Create Variable serializer for request bodies."""
 
-    action: Literal["create"] = "create"
+    action: BulkAction = BulkAction.CREATE
     connections: list[ConnectionBody] = Field(..., description="A list of connections to be created.")
-    action_if_exists: Literal["skip", "overwrite", "fail"] = "fail"
 
 
-class ConnectionBulkUpdateAction(BaseModel):
+class ConnectionBulkUpdateAction(BulkBaseAction):
     """Bulk Update Connection serializer for request bodies."""
 
-    action: Literal["update"] = "update"
+    action: BulkAction = BulkAction.UPDATE
     connections: list[ConnectionBody] = Field(..., description="A list of connections to be updated.")
-    action_if_not_exists: Literal["skip", "fail"] = "fail"
 
 
-class ConnectionBulkDeleteAction(BaseModel):
+class ConnectionBulkDeleteAction(BulkBaseAction):
     """Bulk Delete Connection serializer for request bodies."""
 
-    action: Literal["delete"] = "delete"
+    action: BulkAction = BulkAction.DELETE
     connection_ids: list[str] = Field(..., description="A list of connection IDs to be deleted.")
-    action_if_not_exists: Literal["skip", "fail"] = "fail"
 
 
 class ConnectionBulkBody(BaseModel):

airflow/api_fastapi/core_api/datamodels/variables.py

@@ -18,11 +18,12 @@
 from __future__ import annotations
 
 import json
-from typing import Any, Literal
+from typing import Any
 
 from pydantic import ConfigDict, Field, model_validator
 
 from airflow.api_fastapi.core_api.base import BaseModel
+from airflow.api_fastapi.core_api.datamodels.common import BulkAction, BulkBaseAction
 from airflow.models.base import ID_LEN
 from airflow.typing_compat import Self
 from airflow.utils.log.secrets_masker import redact
@@ -76,28 +77,25 @@ class VariablesImportResponse(BaseModel):
     created_count: int
 
 
-class VariableBulkCreateAction(BaseModel):
+class VariableBulkCreateAction(BulkBaseAction):
     """Bulk Create Variable serializer for request bodies."""
 
-    action: Literal["create"] = "create"
+    action: BulkAction = BulkAction.CREATE
     variables: list[VariableBody] = Field(..., description="A list of variables to be created.")
-    action_if_exists: Literal["skip", "overwrite", "fail"] = "fail"
 
 
-class VariableBulkUpdateAction(BaseModel):
+class VariableBulkUpdateAction(BulkBaseAction):
     """Bulk Update Variable serializer for request bodies."""
 
-    action: Literal["update"] = "update"
+    action: BulkAction = BulkAction.UPDATE
     variables: list[VariableBody] = Field(..., description="A list of variables to be updated.")
-    action_if_not_exists: Literal["skip", "fail"] = "fail"
 
 
-class VariableBulkDeleteAction(BaseModel):
+class VariableBulkDeleteAction(BulkBaseAction):
     """Bulk Delete Variable serializer for request bodies."""
 
-    action: Literal["delete"] = "delete"
+    action: BulkAction = BulkAction.DELETE
     keys: list[str] = Field(..., description="A list of variable keys to be deleted.")
-    action_if_not_exists: Literal["skip", "fail"] = "fail"
 
 
 class VariableBulkBody(BaseModel):

airflow/api_fastapi/core_api/openapi/v1-generated.yaml

@@ -6554,6 +6554,22 @@ components:
       - status
       title: BaseInfoResponse
       description: Base info serializer for responses.
+    BulkAction:
+      type: string
+      enum:
+      - create
+      - delete
+      - update
+      title: BulkAction
+      description: Bulk Action to be performed on the used model.
+    BulkActionOnExistence:
+      type: string
+      enum:
+      - fail
+      - skip
+      - overwrite
+      title: BulkActionOnExistence
+      description: Bulk Action to be taken if the entity already exists or not.
     ClearTaskInstancesBody:
       properties:
         dry_run:
@@ -6847,24 +6863,17 @@ components:
     ConnectionBulkCreateAction:
       properties:
         action:
-          type: string
-          const: create
-          title: Action
+          $ref: '#/components/schemas/BulkAction'
           default: create
+        action_on_existence:
+          $ref: '#/components/schemas/BulkActionOnExistence'
+          default: fail
         connections:
           items:
             $ref: '#/components/schemas/ConnectionBody'
           type: array
           title: Connections
           description: A list of connections to be created.
-        action_if_exists:
-          type: string
-          enum:
-          - skip
-          - overwrite
-          - fail
-          title: Action If Exists
-          default: fail
       type: object
       required:
       - connections
@@ -6873,23 +6882,17 @@ components:
     ConnectionBulkDeleteAction:
       properties:
         action:
-          type: string
-          const: delete
-          title: Action
+          $ref: '#/components/schemas/BulkAction'
           default: delete
+        action_on_existence:
+          $ref: '#/components/schemas/BulkActionOnExistence'
+          default: fail
         connection_ids:
           items:
             type: string
           type: array
           title: Connection Ids
           description: A list of connection IDs to be deleted.
-        action_if_not_exists:
-          type: string
-          enum:
-          - skip
-          - fail
-          title: Action If Not Exists
-          default: fail
       type: object
       required:
       - connection_ids
@@ -6931,23 +6934,17 @@ components:
     ConnectionBulkUpdateAction:
       properties:
         action:
-          type: string
-          const: update
-          title: Action
+          $ref: '#/components/schemas/BulkAction'
           default: update
+        action_on_existence:
+          $ref: '#/components/schemas/BulkActionOnExistence'
+          default: fail
         connections:
           items:
             $ref: '#/components/schemas/ConnectionBody'
           type: array
           title: Connections
           description: A list of connections to be updated.
-        action_if_not_exists:
-          type: string
-          enum:
-          - skip
-          - fail
-          title: Action If Not Exists
-          default: fail
       type: object
       required:
       - connections
@@ -9969,24 +9966,17 @@ components:
     VariableBulkCreateAction:
       properties:
         action:
-          type: string
-          const: create
-          title: Action
+          $ref: '#/components/schemas/BulkAction'
           default: create
+        action_on_existence:
+          $ref: '#/components/schemas/BulkActionOnExistence'
+          default: fail
         variables:
           items:
             $ref: '#/components/schemas/VariableBody'
           type: array
           title: Variables
           description: A list of variables to be created.
-        action_if_exists:
-          type: string
-          enum:
-          - skip
-          - overwrite
-          - fail
-          title: Action If Exists
-          default: fail
       type: object
       required:
       - variables
@@ -9995,23 +9985,17 @@ components:
     VariableBulkDeleteAction:
       properties:
         action:
-          type: string
-          const: delete
-          title: Action
+          $ref: '#/components/schemas/BulkAction'
           default: delete
+        action_on_existence:
+          $ref: '#/components/schemas/BulkActionOnExistence'
+          default: fail
         keys:
           items:
             type: string
           type: array
           title: Keys
           description: A list of variable keys to be deleted.
-        action_if_not_exists:
-          type: string
-          enum:
-          - skip
-          - fail
-          title: Action If Not Exists
-          default: fail
       type: object
       required:
       - keys
@@ -10053,23 +10037,17 @@ components:
     VariableBulkUpdateAction:
       properties:
         action:
-          type: string
-          const: update
-          title: Action
+          $ref: '#/components/schemas/BulkAction'
           default: update
+        action_on_existence:
+          $ref: '#/components/schemas/BulkActionOnExistence'
+          default: fail
         variables:
           items:
             $ref: '#/components/schemas/VariableBody'
           type: array
           title: Variables
           description: A list of variables to be updated.
-        action_if_not_exists:
-          type: string
-          enum:
-          - skip
-          - fail
-          title: Action If Not Exists
-          default: fail
       type: object
       required:
       - variables

airflow/api_fastapi/core_api/routes/public/connections.py

@@ -27,6 +27,7 @@
 from airflow.api_fastapi.common.db.common import SessionDep, paginated_select
 from airflow.api_fastapi.common.parameters import QueryLimit, QueryOffset, SortParam
 from airflow.api_fastapi.common.router import AirflowRouter
+from airflow.api_fastapi.core_api.datamodels.common import BulkAction
 from airflow.api_fastapi.core_api.datamodels.connections import (
     ConnectionBody,
     ConnectionBulkActionResponse,
@@ -152,15 +153,15 @@ def bulk_connections(
     results: dict[str, ConnectionBulkActionResponse] = {}
 
     for action in request.actions:
-        if action.action not in results:
-            results[action.action] = ConnectionBulkActionResponse()
-
-        if action.action == "create":
-            handle_bulk_create(session, action, results[action.action])
-        elif action.action == "update":
-            handle_bulk_update(session, action, results[action.action])
-        elif action.action == "delete":
-            handle_bulk_delete(session, action, results[action.action])
+        if action.action.value not in results:
+            results[action.action.value] = ConnectionBulkActionResponse()
+
+        if action.action == BulkAction.CREATE:
+            handle_bulk_create(session, action, results[action.action.value])  # type: ignore
+        elif action.action == BulkAction.UPDATE:
+            handle_bulk_update(session, action, results[action.action.value])  # type: ignore
+        elif action.action == BulkAction.DELETE:
+            handle_bulk_delete(session, action, results[action.action.value])  # type: ignore
 
     return ConnectionBulkResponse(**results)