Coverage for /home/airflow/.local/lib/python3.12/site-packages/airflow/api_fastapi/core_api/services/public/common.py: 82%

1# Licensed to the Apache Software Foundation (ASF) under one 

2# or more contributor license agreements. See the NOTICE file 

3# distributed with this work for additional information 

4# regarding copyright ownership. The ASF licenses this file 

5# to you under the Apache License, Version 2.0 (the 

6# "License"); you may not use this file except in compliance 

7# with the License. You may obtain a copy of the License at 

8# 

9# http://www.apache.org/licenses/LICENSE-2.0 

10# 

11# Unless required by applicable law or agreed to in writing, 

12# software distributed under the License is distributed on an 

13# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY 

14# KIND, either express or implied. See the License for the 

15# specific language governing permissions and limitations 

16# under the License. 

17 

18from __future__ import annotations 

19 

20from abc import ABC, abstractmethod 

21from typing import Generic 

22 

23from fastapi import HTTPException, status 

24from pydantic import BaseModel 

25from sqlalchemy.ext.declarative import DeclarativeMeta 

26from sqlalchemy.orm import Session 

27 

28from airflow.api_fastapi.core_api.datamodels.common import ( 

29 BulkAction, 

30 BulkActionResponse, 

31 BulkBody, 

32 BulkCreateAction, 

33 BulkDeleteAction, 

34 BulkResponse, 

35 BulkUpdateAction, 

36 T, 

37) 

38 

39 

40class BulkService(Generic[T], ABC): 

41 """Base class for bulk service operations.""" 

42 

43 def __init__(self, session: Session, request: BulkBody[T]): 

44 self.session = session 

45 self.request = request 

46 

47 def handle_request(self) -> BulkResponse: 

48 """Handle request for bulk actions.""" 

49 results: dict[str, BulkActionResponse] = {} 

50 

51 for action in self.request.actions: 

52 if action.action.value not in results: 

53 results[action.action.value] = BulkActionResponse() 

54 

55 if action.action == BulkAction.CREATE: 

56 self.handle_bulk_create(action, results[action.action.value]) 

57 elif action.action == BulkAction.UPDATE: 

58 self.handle_bulk_update(action, results[action.action.value]) 

59 elif action.action == BulkAction.DELETE: 59 ↛ 51line 59 didn't jump to line 51 because the condition on line 59 was always true

60 self.handle_bulk_delete(action, results[action.action.value]) 

61 

62 return BulkResponse(**results) 

63 

64 @abstractmethod 

65 def handle_bulk_create(self, action: BulkCreateAction[T], results: BulkActionResponse) -> None: 

66 """Bulk create entities.""" 

67 raise NotImplementedError 

68 

69 @abstractmethod 

70 def handle_bulk_update(self, action: BulkUpdateAction[T], results: BulkActionResponse) -> None: 

71 """Bulk update entities.""" 

72 raise NotImplementedError 

73 

74 @abstractmethod 

75 def handle_bulk_delete(self, action: BulkDeleteAction[T], results: BulkActionResponse) -> None: 

76 """Bulk delete entities.""" 

77 raise NotImplementedError 

78 

79 @staticmethod 

80 def apply_patch_with_update_mask( 

81 model: DeclarativeMeta, 

82 patch_body: BaseModel, 

83 update_mask: list[str] | None, 

84 non_update_fields: set[str] | None = None, 

85 ) -> DeclarativeMeta: 

86 """ 

87 Apply a patch to the given model using the provided update mask. 

88 

89 :param model: The SQLAlchemy model instance to update. 

90 :param patch_body: Pydantic model containing patch data. 

91 :param update_mask: Optional list of fields to update. 

92 :param non_update_fields: Fields that should not be updated. 

93 :return: The updated SQLAlchemy model instance. 

94 :raises HTTPException: If invalid fields are provided in update_mask. 

95 """ 

96 # Always dump without aliases for internal validation 

97 raw_data = patch_body.model_dump(by_alias=False) 

98 fields_to_update = set(patch_body.model_fields_set) 

99 

100 non_update_fields = non_update_fields or set() 

101 

102 if update_mask: 102 ↛ 103line 102 didn't jump to line 103 because the condition on line 102 was never true

103 restricted_in_mask = set(update_mask).intersection(non_update_fields) 

104 if restricted_in_mask: 

105 raise HTTPException( 

106 status_code=status.HTTP_400_BAD_REQUEST, 

107 detail=f"Update not allowed: the following fields are immutable and cannot be modified: {restricted_in_mask}", 

108 ) 

109 fields_to_update = fields_to_update.intersection(update_mask) 

110 

111 if non_update_fields: 111 ↛ 114line 111 didn't jump to line 114 because the condition on line 111 was always true

112 fields_to_update = fields_to_update - non_update_fields 

113 

114 validated_data = {key: raw_data[key] for key in fields_to_update if key in raw_data} 

115 

116 data = patch_body.model_dump(include=set(validated_data.keys()), by_alias=True) 

117 

118 # Update the model with the validated data 

119 for key, value in data.items(): 

120 setattr(model, key, value) 

121 

122 return model