Coverage for /usr/local/lib/python3.12/site-packages/prefect/settings/base.py: 75%
96 statements
« prev ^ index » next coverage.py v7.10.6, created at 2025-12-05 10:48 +0000
1from __future__ import annotations 1a
3import inspect 1a
4import json 1a
5from functools import partial 1a
6from typing import Any 1a
8from pydantic import ( 1a
9 AliasChoices,
10 AliasPath,
11 SerializationInfo,
12 SerializerFunctionWrapHandler,
13 model_serializer,
14)
15from pydantic_settings import ( 1a
16 BaseSettings,
17 PydanticBaseSettingsSource,
18 SettingsConfigDict,
19)
21from prefect.settings.sources import ( 1a
22 EnvFilterSettingsSource,
23 FilteredDotEnvSettingsSource,
24 PrefectTomlConfigSettingsSource,
25 ProfileSettingsTomlLoader,
26 PyprojectTomlConfigSettingsSource,
27)
28from prefect.utilities.collections import visit_collection 1a
31class PrefectBaseSettings(BaseSettings): 1a
32 @classmethod 1a
33 def settings_customise_sources( 1a
34 cls,
35 settings_cls: type[BaseSettings],
36 init_settings: PydanticBaseSettingsSource,
37 env_settings: PydanticBaseSettingsSource,
38 dotenv_settings: PydanticBaseSettingsSource,
39 file_secret_settings: PydanticBaseSettingsSource,
40 ) -> tuple[PydanticBaseSettingsSource, ...]:
41 """
42 Define an order for Prefect settings sources.
44 The order of the returned callables decides the priority of inputs; first item is the highest priority.
46 See https://docs.pydantic.dev/latest/concepts/pydantic_settings/#customise-settings-sources
47 """
48 env_filter: set[str] = set() 1a
49 for field_name, field in settings_cls.model_fields.items(): 1a
50 if field.validation_alias is not None and isinstance( 1a
51 field.validation_alias, AliasChoices
52 ):
53 for alias in field.validation_alias.choices: 1a
54 if (
55 isinstance(alias, AliasPath)
56 and len(alias.path) > 0
57 and isinstance(alias.path[0], str)
58 ):
59 env_filter.add(alias.path[0]) 1a
60 env_filter.add(field_name) 1a
61 return ( 1a
62 init_settings,
63 EnvFilterSettingsSource(
64 settings_cls,
65 case_sensitive=cls.model_config.get("case_sensitive"),
66 env_prefix=cls.model_config.get("env_prefix"),
67 env_nested_delimiter=cls.model_config.get("env_nested_delimiter"),
68 env_ignore_empty=cls.model_config.get("env_ignore_empty"),
69 env_parse_none_str=cls.model_config.get("env_parse_none_str"),
70 env_parse_enums=cls.model_config.get("env_parse_enums"),
71 env_filter=list(env_filter),
72 ),
73 FilteredDotEnvSettingsSource(
74 settings_cls,
75 env_file=cls.model_config.get("env_file"),
76 env_file_encoding=cls.model_config.get("env_file_encoding"),
77 case_sensitive=cls.model_config.get("case_sensitive"),
78 env_prefix=cls.model_config.get("env_prefix"),
79 env_nested_delimiter=cls.model_config.get("env_nested_delimiter"),
80 env_ignore_empty=cls.model_config.get("env_ignore_empty"),
81 env_parse_none_str=cls.model_config.get("env_parse_none_str"),
82 env_parse_enums=cls.model_config.get("env_parse_enums"),
83 env_blacklist=list(env_filter),
84 ),
85 file_secret_settings,
86 PrefectTomlConfigSettingsSource(settings_cls),
87 PyprojectTomlConfigSettingsSource(settings_cls),
88 ProfileSettingsTomlLoader(settings_cls),
89 )
91 @staticmethod 1a
92 def _settings_warn_unused_config_keys( 1a
93 sources: tuple[object, ...],
94 model_config: SettingsConfigDict,
95 ) -> None:
96 """
97 Override PrefectBaseSettings._settings_warn_unused_config_keys to ensure
98 Prefect’s TOML sources satisfy pydantic-settings 2.11’s unused-config check
99 without renaming public config keys. This method suppresses warnings when
100 Prefect sources are configured.
102 In the future pydantic-settings's init might stop invoking this method, so we
103 should remove this method.
104 """
105 warn_method = getattr(BaseSettings, "_settings_warn_unused_config_keys", None) 1a
106 if warn_method is None: 106 ↛ 107line 106 didn't jump to line 107 because the condition on line 106 was never true1a
107 return
109 adjusted_config = dict(model_config) 1a
111 if any( 111 ↛ 116line 111 didn't jump to line 116 because the condition on line 111 was always true1a
112 isinstance(source, PrefectTomlConfigSettingsSource) for source in sources
113 ):
114 adjusted_config["toml_file"] = None 1a
116 if any( 116 ↛ 122line 116 didn't jump to line 122 because the condition on line 116 was always true1a
117 isinstance(source, PyprojectTomlConfigSettingsSource) for source in sources
118 ):
119 adjusted_config["pyproject_toml_table_header"] = None 1a
120 adjusted_config["pyproject_toml_depth"] = None 1a
122 warn_method(sources, adjusted_config) 1a
124 def to_environment_variables( 1a
125 self,
126 exclude_unset: bool = False,
127 include_secrets: bool = True,
128 include_aliases: bool = False,
129 ) -> dict[str, str]:
130 """Convert the settings object to a dictionary of environment variables."""
131 env: dict[str, Any] = self.model_dump( 1abcd
132 exclude_unset=exclude_unset,
133 mode="json",
134 context={"include_secrets": include_secrets},
135 )
136 env_variables: dict[str, str] = {} 1abcd
137 for key in type(self).model_fields.keys(): 1abcd
138 if isinstance(child_settings := getattr(self, key), PrefectBaseSettings): 1abcd
139 child_env = child_settings.to_environment_variables( 1abcd
140 exclude_unset=exclude_unset,
141 include_secrets=include_secrets,
142 include_aliases=include_aliases,
143 )
144 env_variables.update(child_env) 1abcd
145 elif (value := env.get(key)) is not None: 1abcd
146 validation_alias = type(self).model_fields[key].validation_alias 1abcd
147 if include_aliases and validation_alias is not None: 1abcd
148 if isinstance(validation_alias, AliasChoices): 148 ↛ 154line 148 didn't jump to line 154 because the condition on line 148 was always true1abcd
149 for alias in validation_alias.choices: 1abcd
150 if isinstance(alias, str): 1abcd
151 env_variables[alias.upper()] = ( 1abcd
152 _to_environment_variable_value(value)
153 )
154 elif isinstance(validation_alias, str):
155 env_variables[validation_alias.upper()] = (
156 _to_environment_variable_value(value)
157 )
158 else:
159 env_variables[ 1abcd
160 f"{self.model_config.get('env_prefix')}{key.upper()}"
161 ] = _to_environment_variable_value(value)
162 return env_variables 1abcd
164 @model_serializer( 1a
165 mode="wrap", when_used="always"
166 ) # TODO: reconsider `when_used` default for more control
167 def ser_model( 1a
168 self, handler: SerializerFunctionWrapHandler, info: SerializationInfo
169 ) -> Any:
170 jsonable_self: dict[str, Any] = handler(self) 1aebcd
171 # iterate over fields to ensure child models that have been updated are also included
172 for key in type(self).model_fields.keys(): 1aebcd
173 if info.exclude and key in info.exclude: 173 ↛ 174line 173 didn't jump to line 174 because the condition on line 173 was never true1aebcd
174 continue
175 if info.include and key not in info.include: 175 ↛ 176line 175 didn't jump to line 176 because the condition on line 175 was never true1aebcd
176 continue
178 child_include = None 1aebcd
179 child_exclude = None 1aebcd
180 if info.include and key in info.include and isinstance(info.include, dict): 180 ↛ 181line 180 didn't jump to line 181 because the condition on line 180 was never true1aebcd
181 child_include = info.include[key]
183 if isinstance(child_settings := getattr(self, key), PrefectBaseSettings): 1aebcd
184 child_jsonable = child_settings.model_dump( 1aebcd
185 mode=info.mode,
186 include=child_include, # type: ignore
187 exclude=child_exclude,
188 exclude_unset=info.exclude_unset,
189 context=info.context,
190 by_alias=info.by_alias,
191 exclude_none=info.exclude_none,
192 exclude_defaults=info.exclude_defaults,
193 serialize_as_any=info.serialize_as_any,
194 )
195 if child_jsonable: 1aebcd
196 jsonable_self[key] = child_jsonable 1aebcd
197 if info.context and info.context.get("include_secrets") is True: 1aebcd
198 from prefect.utilities.pydantic import handle_secret_render 1abcd
200 jsonable_self.update( 1abcd
201 {
202 field_name: visit_collection(
203 expr=getattr(self, field_name),
204 visit_fn=partial(handle_secret_render, context=info.context),
205 return_data=True,
206 )
207 for field_name in set(jsonable_self.keys()) # type: ignore
208 }
209 )
211 return jsonable_self 1aebcd
214class PrefectSettingsConfigDict(SettingsConfigDict, total=False): 1a
215 """
216 Configuration for the behavior of Prefect settings models.
217 """
219 prefect_toml_table_header: tuple[str, ...] 1a
220 """ 1a
221 Header of the TOML table within a prefect.toml file to use when filling variables.
222 This is supplied as a `tuple[str, ...]` instead of a `str` to accommodate for headers
223 containing a `.`.
225 To use the root table, exclude this config setting or provide an empty tuple.
226 """
229def _add_environment_variables( 1a
230 schema: dict[str, Any], model: type[PrefectBaseSettings]
231) -> None:
232 for property in schema["properties"]:
233 env_vars: list[str] = []
234 schema["properties"][property]["supported_environment_variables"] = env_vars
235 field = model.model_fields[property]
236 if inspect.isclass(field.annotation) and issubclass(
237 field.annotation, PrefectBaseSettings
238 ):
239 continue
240 elif field.validation_alias:
241 if isinstance(field.validation_alias, AliasChoices):
242 for alias in field.validation_alias.choices:
243 if isinstance(alias, str):
244 env_vars.append(alias.upper())
245 else:
246 env_vars.append(f"{model.model_config.get('env_prefix')}{property.upper()}")
249def build_settings_config( 1a
250 path: tuple[str, ...] = tuple(), frozen: bool = False
251) -> PrefectSettingsConfigDict:
252 env_prefix = f"PREFECT_{'_'.join(path).upper()}_" if path else "PREFECT_" 1a
253 return PrefectSettingsConfigDict( 1a
254 env_prefix=env_prefix,
255 env_file=".env",
256 extra="ignore",
257 toml_file="prefect.toml",
258 prefect_toml_table_header=path,
259 pyproject_toml_table_header=("tool", "prefect", *path),
260 json_schema_extra=_add_environment_variables, # type: ignore
261 frozen=frozen,
262 )
265_build_settings_config = build_settings_config # noqa # TODO: remove once all usage updated 1a
268def _to_environment_variable_value( 1a
269 value: list[object] | set[object] | tuple[object] | Any,
270) -> str:
271 if isinstance(value, (list, set, tuple)): 1abcd
272 return ",".join(str(v) for v in sorted(value, key=str)) 1abcd
273 if isinstance(value, dict): 1abcd
274 return json.dumps(value) 1abcd
275 return str(value) 1abcd