Coverage for /home/airflow/.local/lib/python3.12/site-packages/airflow/api_fastapi/core_api/datamodels/variables.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 

20import json 

21from collections.abc import Iterable 

22 

23from pydantic import Field, JsonValue, field_validator, model_validator 

24 

25from airflow._shared.secrets_masker import redact 

26from airflow.api_fastapi.core_api.base import BaseModel, StrictBaseModel, make_partial_model 

27from airflow.configuration import conf 

28from airflow.models.base import ID_LEN 

29from airflow.typing_compat import Self 

30 

31 

32class VariableResponse(BaseModel): 

33 """Variable serializer for responses.""" 

34 

35 key: str 

36 val: str | None = Field(alias="value", default=None) 

37 description: str | None 

38 is_encrypted: bool 

39 team_name: str | None 

40 

41 @model_validator(mode="after") 

42 def redact_val(self) -> Self: 

43 if self.val is None: 43 ↛ 44line 43 didn't jump to line 44 because the condition on line 43 was never true

44 return self 

45 try: 

46 val_dict = json.loads(self.val) 

47 redacted_dict = redact(val_dict, self.key) 

48 self.val = json.dumps(redacted_dict) 

49 return self 

50 except json.JSONDecodeError: 

51 # value is not a serialized string representation of a dict. 

52 self.val = str(redact(self.val, self.key)) 

53 return self 

54 

55 

56class VariableBody(StrictBaseModel): 

57 """Variable serializer for bodies.""" 

58 

59 key: str = Field(max_length=ID_LEN) 

60 value: JsonValue = Field(serialization_alias="val") 

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

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

63 

64 @field_validator("value", mode="after") 

65 @classmethod 

66 def serialize_non_string_value(cls, value: JsonValue) -> str: 

67 # Variables are stored as strings. A non-string JSON value (list, dict, number, bool, 

68 # null) must be JSON-encoded here: POST would otherwise store its Python repr, which 

69 # deserialize_json cannot read back, and PATCH would fail in the ORM Fernet step on 

70 # bytes(value, "utf-8") (see #71010). indent=2 matches Variable.set(serialize_json=True). 

71 # Not allow_nan=False: rejecting here yields a 422 whose response embeds the NaN input, 

72 # which starlette's JSONResponse cannot serialize -- the request would 500 instead. 

73 if isinstance(value, str): 

74 return value 

75 return json.dumps(value, indent=2) 

76 

77 @model_validator(mode="after") 

78 def validate_team_name(self) -> VariableBody: 

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

80 raise ValueError( 

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

82 ) 

83 return self 

84 

85 

86VariableBodyPartial = make_partial_model(VariableBody) 

87 

88 

89class VariableCollectionResponse(BaseModel): 

90 """Variable Collection serializer for responses.""" 

91 

92 variables: Iterable[VariableResponse] 

93 total_entries: int 

94 

95 

96class VariablesImportResponse(BaseModel): 

97 """Import Variables serializer for responses.""" 

98 

99 created_variable_keys: list[str] 

100 import_count: int 

101 created_count: int