Coverage for /home/airflow/.local/lib/python3.12/site-packages/airflow/api_fastapi/core_api/datamodels/pools.py: 96%

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 collections.abc import Callable, Iterable 

21from typing import Annotated 

22 

23from pydantic import BeforeValidator, Field, model_validator 

24 

25from airflow.api_fastapi.core_api.base import BaseModel, StrictBaseModel 

26from airflow.configuration import conf 

27 

28 

29def _call_function(function: Callable[[], int]) -> int: 

30 """ 

31 Call the given function. 

32 

33 Used for the BeforeValidator to get the actual values from the bound method. 

34 """ 

35 return function() 

36 

37 

38PoolSlots = Annotated[ 

39 int, 

40 Field(ge=-1, description="Number of slots. Use -1 for unlimited."), 

41] 

42 

43 

44class BasePool(BaseModel): 

45 """Base serializer for Pool.""" 

46 

47 pool: str = Field(serialization_alias="name") 

48 slots: PoolSlots 

49 description: str | None = Field(default=None) 

50 include_deferred: bool 

51 

52 

53def _sanitize_open_slots(value) -> int: 

54 if isinstance(value, float) and value == float("inf"): 54 ↛ 55line 54 didn't jump to line 55 because the condition on line 54 was never true

55 return -1 

56 return value 

57 

58 

59class PoolResponse(BasePool): 

60 """Pool serializer for responses.""" 

61 

62 occupied_slots: Annotated[int, BeforeValidator(_call_function)] 

63 running_slots: Annotated[int, BeforeValidator(_call_function)] 

64 queued_slots: Annotated[int, BeforeValidator(_call_function)] 

65 scheduled_slots: Annotated[int, BeforeValidator(_call_function)] 

66 open_slots: Annotated[int, BeforeValidator(lambda v: _sanitize_open_slots(_call_function(v)))] 

67 deferred_slots: Annotated[int, BeforeValidator(_call_function)] 

68 team_name: str | None 

69 

70 

71class PoolCollectionResponse(BaseModel): 

72 """Pool Collection serializer for responses.""" 

73 

74 pools: Iterable[PoolResponse] 

75 total_entries: int 

76 

77 

78class PoolPatchBody(StrictBaseModel): 

79 """Pool serializer for patch bodies.""" 

80 

81 name: str | None = Field(default=None, alias="pool") 

82 slots: PoolSlots | None = None 

83 description: str | None = None 

84 include_deferred: bool | None = None 

85 team_name: str | None = Field(max_length=50, default=None) 

86 

87 @model_validator(mode="after") 

88 def validate_team_name(self) -> PoolPatchBody: 

89 if self.team_name is not None and not conf.getboolean("core", "multi_team"): 

90 raise ValueError( 

91 "team_name cannot be set when multi_team mode is disabled. Please contact your administrator." 

92 ) 

93 return self 

94 

95 

96class PoolBody(BasePool, StrictBaseModel): 

97 """Pool serializer for post bodies.""" 

98 

99 pool: str = Field(alias="name", max_length=256) 

100 description: str | None = None 

101 include_deferred: bool = False 

102 team_name: str | None = Field(max_length=50, default=None) 

103 

104 @model_validator(mode="after") 

105 def validate_team_name(self) -> PoolBody: 

106 if self.team_name is not None and not conf.getboolean("core", "multi_team"): 

107 raise ValueError( 

108 "team_name cannot be set when multi_team mode is disabled. Please contact your administrator." 

109 ) 

110 return self