jason810496 commented on code in PR #54597:
URL: https://github.com/apache/airflow/pull/54597#discussion_r2297953990
##########
airflow-core/src/airflow/api_fastapi/core_api/services/public/common.py:
##########
@@ -72,3 +73,71 @@ def handle_bulk_update(self, action: BulkUpdateAction[T],
results: BulkActionRes
def handle_bulk_delete(self, action: BulkDeleteAction[T], results:
BulkActionResponse) -> None:
"""Bulk delete entities."""
raise NotImplementedError
+
+
+class PatchUtil:
+ """
+ Utility class for applying patch operations with support for update masks.
+
+ This helper is used to update only selected fields of a Pydantic model
instance,
+ while ensuring:
+ - Only fields present in the update_mask are modified.
+ - Non-updatable fields can be excluded explicitly.
+ - Both raw field names and aliases are handled correctly.
+ """
+
+ @staticmethod
+ def apply_patch_with_update_mask(
+ model,
+ patch_body,
+ update_mask: list[str] | None,
+ non_update_fields: set[str] | None = None,
+ ):
+ """
+ Apply a patch to the given model using the provided update mask.
+
+ Args:
+ model: The model instance to update.
+ patch_body: Pydantic model containing patch data.
+ update_mask (list[str] | None): Optional list of fields to update.
+ non_update_fields (set[str] | None): Fields that should not be
updated.
+
+ Returns:
+ The updated model instance.
+
+ Raises:
+ HTTPException: If invalid fields are provided in update_mask.
+ """
+ print("Inside patch")
+
+ # Always dump without aliases for internal validation
+ raw_data = patch_body.model_dump(by_alias=False)
+
+ fields_set = set(patch_body.model_fields_set)
+ allowed_fields = set(patch_body.model_fields.keys())
+
+ print(fields_set, update_mask)
+ if update_mask:
+ invalid = set(update_mask) - allowed_fields
+ if invalid:
+ raise HTTPException(
+ status_code=status.HTTP_400_BAD_REQUEST,
+ detail=f"Invalid fields in update_mask: {invalid}",
+ )
+ fields_set = fields_set.intersection(update_mask)
+
+ if non_update_fields:
+ fields_set = fields_set - non_update_fields
+
+ # Validate only the subset of fields we want to update
+ print("Validating data with fields_set:", fields_set, raw_data)
+ validated_data = {key: raw_data[key] for key in fields_set if key in
raw_data}
+
+ data = patch_body.model_dump(include=set(validated_data.keys()),
by_alias=True)
+
+ # Update the model with the validated data
+ print("Updating model with data:", data)
Review Comment:
```suggestion
```
##########
airflow-core/src/airflow/api_fastapi/core_api/services/public/common.py:
##########
@@ -72,3 +73,71 @@ def handle_bulk_update(self, action: BulkUpdateAction[T],
results: BulkActionRes
def handle_bulk_delete(self, action: BulkDeleteAction[T], results:
BulkActionResponse) -> None:
"""Bulk delete entities."""
raise NotImplementedError
+
+
+class PatchUtil:
+ """
+ Utility class for applying patch operations with support for update masks.
+
+ This helper is used to update only selected fields of a Pydantic model
instance,
+ while ensuring:
+ - Only fields present in the update_mask are modified.
+ - Non-updatable fields can be excluded explicitly.
+ - Both raw field names and aliases are handled correctly.
+ """
+
+ @staticmethod
+ def apply_patch_with_update_mask(
+ model,
+ patch_body,
+ update_mask: list[str] | None,
+ non_update_fields: set[str] | None = None,
+ ):
+ """
+ Apply a patch to the given model using the provided update mask.
+
+ Args:
+ model: The model instance to update.
+ patch_body: Pydantic model containing patch data.
+ update_mask (list[str] | None): Optional list of fields to update.
+ non_update_fields (set[str] | None): Fields that should not be
updated.
+
+ Returns:
+ The updated model instance.
+
+ Raises:
+ HTTPException: If invalid fields are provided in update_mask.
+ """
+ print("Inside patch")
+
+ # Always dump without aliases for internal validation
+ raw_data = patch_body.model_dump(by_alias=False)
+
+ fields_set = set(patch_body.model_fields_set)
+ allowed_fields = set(patch_body.model_fields.keys())
+
+ print(fields_set, update_mask)
Review Comment:
```suggestion
fields_to_update = patch_body.model_fields_set
```
##########
airflow-core/src/airflow/api_fastapi/core_api/services/public/pools.py:
##########
@@ -31,11 +31,50 @@
)
from airflow.api_fastapi.core_api.datamodels.pools import (
PoolBody,
+ PoolPatchBody,
)
-from airflow.api_fastapi.core_api.services.public.common import BulkService
+from airflow.api_fastapi.core_api.services.public.common import BulkService,
PatchUtil
from airflow.models.pool import Pool
+def update_orm_from_pydantic(
+ old_pool: Pool,
+ patch_body: PoolBody | PoolPatchBody,
+ update_mask: list[str] | None,
+) -> Pool:
+ """
+ Patch an existing Pool instance with provided update fields.
+
+ Args:
+ old_pool (Pool): The existing Pool ORM model instance to be updated.
+ patch_body (PoolBody): Pydantic model containing the fields to update.
+ update_mask (list[str] | None): Specific fields to update. If None,
all provided fields will be considered.
+
+ Returns:
+ Pool: The updated Pool instance.
+
+ Raises:
+ HTTPException: If attempting to update disallowed fields on
`default_pool`.
+ """
+ # Special restriction: default pool only allows limited fields to be
patched
+ if old_pool.pool == Pool.DEFAULT_POOL_NAME:
+ if update_mask and all(mask.strip() in {"slots", "include_deferred"}
for mask in update_mask):
+ pass
+ else:
+ raise HTTPException(
+ status.HTTP_400_BAD_REQUEST,
+ "Only slots and include_deferred can be modified on Default
Pool",
+ )
Review Comment:
Like these validation are copy from the router.
##########
airflow-core/src/airflow/api_fastapi/core_api/services/public/pools.py:
##########
@@ -90,6 +129,7 @@ def handle_bulk_update(self, action:
BulkUpdateAction[PoolBody], results: BulkAc
"""Bulk Update pools."""
to_update_pool_names = {pool.pool for pool in action.entities}
_, matched_pool_names, not_found_pool_names =
self.categorize_pools(to_update_pool_names)
+ print(to_update_pool_names)
Review Comment:
```suggestion
```
##########
airflow-core/src/airflow/api_fastapi/core_api/services/public/common.py:
##########
@@ -72,3 +73,71 @@ def handle_bulk_update(self, action: BulkUpdateAction[T],
results: BulkActionRes
def handle_bulk_delete(self, action: BulkDeleteAction[T], results:
BulkActionResponse) -> None:
"""Bulk delete entities."""
raise NotImplementedError
+
+
+class PatchUtil:
+ """
+ Utility class for applying patch operations with support for update masks.
+
+ This helper is used to update only selected fields of a Pydantic model
instance,
+ while ensuring:
+ - Only fields present in the update_mask are modified.
+ - Non-updatable fields can be excluded explicitly.
+ - Both raw field names and aliases are handled correctly.
+ """
+
+ @staticmethod
+ def apply_patch_with_update_mask(
+ model,
+ patch_body,
+ update_mask: list[str] | None,
+ non_update_fields: set[str] | None = None,
+ ):
+ """
+ Apply a patch to the given model using the provided update mask.
+
+ Args:
+ model: The model instance to update.
+ patch_body: Pydantic model containing patch data.
+ update_mask (list[str] | None): Optional list of fields to update.
+ non_update_fields (set[str] | None): Fields that should not be
updated.
+
+ Returns:
+ The updated model instance.
+
+ Raises:
+ HTTPException: If invalid fields are provided in update_mask.
+ """
+ print("Inside patch")
+
+ # Always dump without aliases for internal validation
+ raw_data = patch_body.model_dump(by_alias=False)
+
+ fields_set = set(patch_body.model_fields_set)
+ allowed_fields = set(patch_body.model_fields.keys())
+
+ print(fields_set, update_mask)
+ if update_mask:
+ invalid = set(update_mask) - allowed_fields
+ if invalid:
+ raise HTTPException(
+ status_code=status.HTTP_400_BAD_REQUEST,
+ detail=f"Invalid fields in update_mask: {invalid}",
+ )
+ fields_set = fields_set.intersection(update_mask)
+
+ if non_update_fields:
+ fields_set = fields_set - non_update_fields
+
+ # Validate only the subset of fields we want to update
+ print("Validating data with fields_set:", fields_set, raw_data)
+ validated_data = {key: raw_data[key] for key in fields_set if key in
raw_data}
Review Comment:
There is not need to raise error for invalid fields, having `intersection`
will ensure there will only be valid field.
We should align the behavior here:
https://github.com/apache/airflow/blob/d9f5d4723cfdb0733c4623fbd518c5d4bd3a7e31/airflow-core/src/airflow/api_fastapi/core_api/routes/public/variables.py#L151-L159
##########
airflow-core/src/airflow/api_fastapi/core_api/services/public/pools.py:
##########
@@ -31,11 +31,50 @@
)
from airflow.api_fastapi.core_api.datamodels.pools import (
PoolBody,
+ PoolPatchBody,
)
-from airflow.api_fastapi.core_api.services.public.common import BulkService
+from airflow.api_fastapi.core_api.services.public.common import BulkService,
PatchUtil
from airflow.models.pool import Pool
+def update_orm_from_pydantic(
+ old_pool: Pool,
+ patch_body: PoolBody | PoolPatchBody,
+ update_mask: list[str] | None,
+) -> Pool:
Review Comment:
For these `update_orm_from_pydantic` functions. How about adding `session`
in the function and select the entity in the function.
Since there are still common validation between service level and router
level, and we should consolidate them in the same function to reuse.
##########
airflow-core/src/airflow/api_fastapi/core_api/services/public/common.py:
##########
@@ -72,3 +73,71 @@ def handle_bulk_update(self, action: BulkUpdateAction[T],
results: BulkActionRes
def handle_bulk_delete(self, action: BulkDeleteAction[T], results:
BulkActionResponse) -> None:
"""Bulk delete entities."""
raise NotImplementedError
+
+
+class PatchUtil:
+ """
+ Utility class for applying patch operations with support for update masks.
+
+ This helper is used to update only selected fields of a Pydantic model
instance,
+ while ensuring:
+ - Only fields present in the update_mask are modified.
+ - Non-updatable fields can be excluded explicitly.
+ - Both raw field names and aliases are handled correctly.
+ """
+
+ @staticmethod
+ def apply_patch_with_update_mask(
+ model,
+ patch_body,
+ update_mask: list[str] | None,
+ non_update_fields: set[str] | None = None,
+ ):
+ """
+ Apply a patch to the given model using the provided update mask.
+
+ Args:
+ model: The model instance to update.
+ patch_body: Pydantic model containing patch data.
+ update_mask (list[str] | None): Optional list of fields to update.
+ non_update_fields (set[str] | None): Fields that should not be
updated.
+
+ Returns:
+ The updated model instance.
+
+ Raises:
+ HTTPException: If invalid fields are provided in update_mask.
+ """
+ print("Inside patch")
Review Comment:
```suggestion
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
