Coverage for /usr/local/lib/python3.12/site-packages/prefect/settings/profiles.py: 37%
164 statements
« prev ^ index » next coverage.py v7.10.6, created at 2025-12-05 13:38 +0000
1from __future__ import annotations 1a
3import warnings 1a
4from pathlib import Path 1a
5from typing import ( 1a
6 Annotated,
7 Any,
8 ClassVar,
9 Iterable,
10 Iterator,
11 Optional,
12)
14import toml 1a
15from pydantic import ( 1a
16 BaseModel,
17 BeforeValidator,
18 ConfigDict,
19 Field,
20 ValidationError,
21)
23from prefect.exceptions import ProfileSettingsValidationError 1a
24from prefect.settings.constants import DEFAULT_PROFILES_PATH 1a
25from prefect.settings.context import get_current_settings 1a
26from prefect.settings.legacy import Setting, _get_settings_fields 1a
27from prefect.settings.models.root import Settings 1a
28from prefect.utilities.collections import set_in_dict 1a
31def _cast_settings( 1a
32 settings: dict[str | Setting, Any] | Any,
33) -> dict[Setting, Any]:
34 """For backwards compatibility, allow either Settings objects as keys or string references to settings."""
35 if not isinstance(settings, dict): 35 ↛ 36line 35 didn't jump to line 36 because the condition on line 35 was never true1a
36 raise ValueError("Settings must be a dictionary.")
37 casted_settings = {} 1a
38 for k, value in settings.items(): 1a
39 try: 1a
40 if isinstance(k, str): 40 ↛ 43line 40 didn't jump to line 43 because the condition on line 40 was always true1a
41 setting = _get_settings_fields(Settings)[k] 1a
42 else:
43 setting = k
44 casted_settings[setting] = value 1a
45 except KeyError as e:
46 warnings.warn(f"Setting {e} is not recognized")
47 continue
48 return casted_settings 1a
51############################################################################
52# Profiles
55class Profile(BaseModel): 1a
56 """A user profile containing settings."""
58 model_config: ClassVar[ConfigDict] = ConfigDict( 1a
59 extra="ignore", arbitrary_types_allowed=True
60 )
62 name: str 1a
63 settings: Annotated[dict[Setting, Any], BeforeValidator(_cast_settings)] = Field( 1a
64 default_factory=dict
65 )
66 source: Optional[Path] = None 1a
68 def to_environment_variables(self) -> dict[str, str]: 1a
69 """Convert the profile settings to a dictionary of environment variables."""
70 return {
71 setting.name: str(value)
72 for setting, value in self.settings.items()
73 if value is not None
74 }
76 def validate_settings(self) -> None: 1a
77 """
78 Validate all settings in this profile by creating a partial Settings object
79 with the nested structure properly constructed using accessor paths.
80 """
81 if not self.settings:
82 return
84 nested_settings: dict[str, Any] = {}
86 for setting, value in self.settings.items():
87 set_in_dict(nested_settings, setting.accessor, value)
89 try:
90 Settings.model_validate(nested_settings)
91 except ValidationError as e:
92 errors: list[tuple[Setting, ValidationError]] = []
94 for error in e.errors():
95 error_path = ".".join(str(loc) for loc in error["loc"])
97 for setting in self.settings.keys():
98 if setting.accessor == error_path:
99 errors.append(
100 (
101 setting,
102 ValidationError.from_exception_data(
103 "ValidationError", [error]
104 ),
105 )
106 )
107 break
109 if errors:
110 raise ProfileSettingsValidationError(errors)
113class ProfilesCollection: 1a
114 """ "
115 A utility class for working with a collection of profiles.
117 Profiles in the collection must have unique names.
119 The collection may store the name of the active profile.
120 """
122 def __init__(self, profiles: Iterable[Profile], active: str | None = None) -> None: 1a
123 self.profiles_by_name: dict[str, Profile] = { 1a
124 profile.name: profile for profile in profiles
125 }
126 self.active_name = active 1a
128 @property 1a
129 def names(self) -> set[str]: 1a
130 """
131 Return a set of profile names in this collection.
132 """
133 return set(self.profiles_by_name.keys()) 1a
135 @property 1a
136 def active_profile(self) -> Profile | None: 1a
137 """
138 Retrieve the active profile in this collection.
139 """
140 if self.active_name is None:
141 return None
142 return self[self.active_name]
144 def set_active(self, name: str | None, check: bool = True) -> None: 1a
145 """
146 Set the active profile name in the collection.
148 A null value may be passed to indicate that this collection does not determine
149 the active profile.
150 """
151 if check and name is not None and name not in self.names:
152 raise ValueError(f"Unknown profile name {name!r}.")
153 self.active_name = name
155 def update_profile( 1a
156 self,
157 name: str,
158 settings: dict[Setting, Any],
159 source: Path | None = None,
160 ) -> Profile:
161 """
162 Add a profile to the collection or update the existing on if the name is already
163 present in this collection.
165 If updating an existing profile, the settings will be merged. Settings can
166 be dropped from the existing profile by setting them to `None` in the new
167 profile.
169 Returns the new profile object.
170 """
171 existing = self.profiles_by_name.get(name)
173 # Convert the input to a `Profile` to cast settings to the correct type
174 profile = Profile(name=name, settings=settings, source=source)
176 if existing:
177 new_settings = {**existing.settings, **profile.settings}
179 # Drop null keys to restore to default
180 for key, value in tuple(new_settings.items()):
181 if value is None:
182 new_settings.pop(key)
184 new_profile = Profile(
185 name=profile.name,
186 settings=new_settings,
187 source=source or profile.source,
188 )
189 else:
190 new_profile = profile
192 self.profiles_by_name[new_profile.name] = new_profile
194 return new_profile
196 def add_profile(self, profile: Profile) -> None: 1a
197 """
198 Add a profile to the collection.
200 If the profile name already exists, an exception will be raised.
201 """
202 if profile.name in self.profiles_by_name:
203 raise ValueError(
204 f"Profile name {profile.name!r} already exists in collection."
205 )
207 self.profiles_by_name[profile.name] = profile
209 def remove_profile(self, name: str) -> None: 1a
210 """
211 Remove a profile from the collection.
212 """
213 self.profiles_by_name.pop(name)
215 def without_profile_source(self, path: Path | None) -> "ProfilesCollection": 1a
216 """
217 Remove profiles that were loaded from a given path.
219 Returns a new collection.
220 """
221 return ProfilesCollection(
222 [
223 profile
224 for profile in self.profiles_by_name.values()
225 if profile.source != path
226 ],
227 active=self.active_name,
228 )
230 def to_dict(self) -> dict[str, Any]: 1a
231 """
232 Convert to a dictionary suitable for writing to disk.
233 """
234 return {
235 "active": self.active_name,
236 "profiles": {
237 profile.name: profile.to_environment_variables()
238 for profile in self.profiles_by_name.values()
239 },
240 }
242 def __getitem__(self, name: str) -> Profile: 1a
243 return self.profiles_by_name[name] 1a
245 def __iter__(self) -> Iterator[str]: 1a
246 return self.profiles_by_name.__iter__()
248 def items(self) -> list[tuple[str, Profile]]: 1a
249 return list(self.profiles_by_name.items())
251 def __eq__(self, __o: object) -> bool: 1a
252 if not isinstance(__o, ProfilesCollection):
253 return False
255 return (
256 self.profiles_by_name == __o.profiles_by_name
257 and self.active_name == __o.active_name
258 )
260 def __repr__(self) -> str: 1a
261 return (
262 f"ProfilesCollection(profiles={list(self.profiles_by_name.values())!r},"
263 f" active={self.active_name!r})>"
264 )
267def _read_profiles_from(path: Path) -> ProfilesCollection: 1a
268 """
269 Read profiles from a path into a new `ProfilesCollection`.
271 Profiles are expected to be written in TOML with the following schema:
272 ```
273 active = <name: Optional[str]>
275 [profiles.<name: str>]
276 <SETTING: str> = <value: Any>
277 ```
278 """
279 contents = toml.loads(path.read_text()) 1a
280 active_profile = contents.get("active") 1a
281 raw_profiles = contents.get("profiles", {}) 1a
283 profiles = [] 1a
284 for name, settings in raw_profiles.items(): 1a
285 profiles.append(Profile(name=name, settings=settings, source=path)) 1a
287 return ProfilesCollection(profiles, active=active_profile) 1a
290def _write_profiles_to(path: Path, profiles: ProfilesCollection) -> None: 1a
291 """
292 Write profiles in the given collection to a path as TOML.
294 Any existing data not present in the given `profiles` will be deleted.
295 """
296 if not path.exists():
297 path.parent.mkdir(parents=True, exist_ok=True)
298 path.touch(mode=0o600)
299 path.write_text(toml.dumps(profiles.to_dict()))
302def load_profiles(include_defaults: bool = True) -> ProfilesCollection: 1a
303 """
304 Load profiles from the current profile path. Optionally include profiles from the
305 default profile path.
306 """
307 current_settings = get_current_settings() 1a
308 default_profiles = _read_profiles_from(DEFAULT_PROFILES_PATH) 1a
310 if current_settings.profiles_path is None: 310 ↛ 311line 310 didn't jump to line 311 because the condition on line 310 was never true1a
311 raise RuntimeError(
312 "No profiles path set; please ensure `PREFECT_PROFILES_PATH` is set."
313 )
315 if not include_defaults: 315 ↛ 316line 315 didn't jump to line 316 because the condition on line 315 was never true1a
316 if not current_settings.profiles_path.exists():
317 return ProfilesCollection([])
318 return _read_profiles_from(current_settings.profiles_path)
320 user_profiles_path = current_settings.profiles_path 1a
321 profiles = default_profiles 1a
322 if user_profiles_path.exists(): 322 ↛ 323line 322 didn't jump to line 323 because the condition on line 322 was never true1a
323 user_profiles = _read_profiles_from(user_profiles_path)
325 # Merge all of the user profiles with the defaults
326 for name in user_profiles:
327 if not (source := user_profiles[name].source):
328 raise ValueError(f"Profile {name!r} has no source.")
329 profiles.update_profile(
330 name,
331 settings=user_profiles[name].settings,
332 source=source,
333 )
335 if user_profiles.active_name:
336 profiles.set_active(user_profiles.active_name, check=False)
338 return profiles 1a
341def load_current_profile() -> Profile: 1a
342 """
343 Load the current profile from the default and current profile paths.
345 This will _not_ include settings from the current settings context. Only settings
346 that have been persisted to the profiles file will be saved.
347 """
348 import prefect.context
350 profiles = load_profiles()
351 context = prefect.context.get_settings_context()
353 if context:
354 profiles.set_active(context.profile.name)
356 return profiles.active_profile
359def save_profiles(profiles: ProfilesCollection) -> None: 1a
360 """
361 Writes all non-default profiles to the current profiles path.
362 """
363 profiles_path = get_current_settings().profiles_path
364 assert profiles_path is not None, "Profiles path is not set."
365 profiles = profiles.without_profile_source(DEFAULT_PROFILES_PATH)
366 return _write_profiles_to(profiles_path, profiles)
369def load_profile(name: str) -> Profile: 1a
370 """
371 Load a single profile by name.
372 """
373 profiles = load_profiles()
374 try:
375 return profiles[name]
376 except KeyError:
377 raise ValueError(f"Profile {name!r} not found.")
380def update_current_profile( 1a
381 settings: dict[str | Setting, Any],
382) -> Profile:
383 """
384 Update the persisted data for the profile currently in-use.
386 If the profile does not exist in the profiles file, it will be created.
388 Given settings will be merged with the existing settings as described in
389 `ProfilesCollection.update_profile`.
391 Returns:
392 The new profile.
393 """
394 import prefect.context
396 current_profile = prefect.context.get_settings_context().profile
398 if not current_profile:
399 from prefect.exceptions import MissingProfileError
401 raise MissingProfileError("No profile is currently in use.")
403 profiles = load_profiles()
405 # Ensure the current profile's settings are present
406 profiles.update_profile(current_profile.name, current_profile.settings)
407 # Then merge the new settings in
408 new_profile = profiles.update_profile(
409 current_profile.name, _cast_settings(settings)
410 )
412 new_profile.validate_settings()
414 save_profiles(profiles)
416 return profiles[current_profile.name]