prefect-client 3.0.10__py3-none-any.whl → 3.0.11__py3-none-any.whl

prefect/settings/models/tasks.py

@@ -0,0 +1,91 @@
+from typing import Optional, Union
+
+from pydantic import AliasChoices, AliasPath, Field
+from pydantic_settings import SettingsConfigDict
+
+from prefect.settings.base import PrefectBaseSettings
+
+
+class TasksRunnerSettings(PrefectBaseSettings):
+    model_config = SettingsConfigDict(
+        env_prefix="PREFECT_TASKS_RUNNER_", env_file=".env", extra="ignore"
+    )
+
+    thread_pool_max_workers: Optional[int] = Field(
+        default=None,
+        gt=0,
+        description="The maximum number of workers for ThreadPoolTaskRunner.",
+        validation_alias=AliasChoices(
+            AliasPath("thread_pool_max_workers"),
+            "prefect_tasks_runner_thread_pool_max_workers",
+            "prefect_task_runner_thread_pool_max_workers",
+        ),
+    )
+
+
+class TasksSchedulingSettings(PrefectBaseSettings):
+    model_config = SettingsConfigDict(
+        env_prefix="PREFECT_TASKS_SCHEDULING_", env_file=".env", extra="ignore"
+    )
+
+    default_storage_block: Optional[str] = Field(
+        default=None,
+        description="The `block-type/block-document` slug of a block to use as the default storage for autonomous tasks.",
+        validation_alias=AliasChoices(
+            AliasPath("default_storage_block"),
+            "prefect_tasks_scheduling_default_storage_block",
+            "prefect_task_scheduling_default_storage_block",
+        ),
+    )
+
+    delete_failed_submissions: bool = Field(
+        default=True,
+        description="Whether or not to delete failed task submissions from the database.",
+        validation_alias=AliasChoices(
+            AliasPath("delete_failed_submissions"),
+            "prefect_tasks_scheduling_delete_failed_submissions",
+            "prefect_task_scheduling_delete_failed_submissions",
+        ),
+    )
+
+
+class TasksSettings(PrefectBaseSettings):
+    model_config = SettingsConfigDict(
+        env_prefix="PREFECT_TASKS_", env_file=".env", extra="ignore"
+    )
+
+    refresh_cache: bool = Field(
+        default=False,
+        description="If `True`, enables a refresh of cached results: re-executing the task will refresh the cached results.",
+    )
+
+    default_retries: int = Field(
+        default=0,
+        ge=0,
+        description="This value sets the default number of retries for all tasks.",
+        validation_alias=AliasChoices(
+            AliasPath("default_retries"),
+            "prefect_tasks_default_retries",
+            "prefect_task_default_retries",
+        ),
+    )
+
+    default_retry_delay_seconds: Union[int, float, list[float]] = Field(
+        default=0,
+        description="This value sets the default retry delay seconds for all tasks.",
+        validation_alias=AliasChoices(
+            AliasPath("default_retry_delay_seconds"),
+            "prefect_tasks_default_retry_delay_seconds",
+            "prefect_task_default_retry_delay_seconds",
+        ),
+    )
+
+    runner: TasksRunnerSettings = Field(
+        default_factory=TasksRunnerSettings,
+        description="Settings for controlling task runner behavior",
+    )
+
+    scheduling: TasksSchedulingSettings = Field(
+        default_factory=TasksSchedulingSettings,
+        description="Settings for controlling client-side task scheduling behavior",
+    )

prefect/settings/models/testing.py

@@ -0,0 +1,52 @@
+from typing import Any, Optional
+
+from pydantic import AliasChoices, AliasPath, Field
+from pydantic_settings import SettingsConfigDict
+
+from prefect.settings.base import PrefectBaseSettings
+
+
+class TestingSettings(PrefectBaseSettings):
+    model_config = SettingsConfigDict(
+        env_prefix="PREFECT_TESTING_", env_file=".env", extra="ignore"
+    )
+
+    test_mode: bool = Field(
+        default=False,
+        description="If `True`, places the API in test mode. This may modify behavior to facilitate testing.",
+        validation_alias=AliasChoices(
+            AliasPath("test_mode"),
+            "prefect_testing_test_mode",
+            "prefect_test_mode",
+        ),
+    )
+
+    unit_test_mode: bool = Field(
+        default=False,
+        description="This setting only exists to facilitate unit testing. If `True`, code is executing in a unit test context. Defaults to `False`.",
+        validation_alias=AliasChoices(
+            AliasPath("unit_test_mode"),
+            "prefect_testing_unit_test_mode",
+            "prefect_unit_test_mode",
+        ),
+    )
+
+    unit_test_loop_debug: bool = Field(
+        default=True,
+        description="If `True` turns on debug mode for the unit testing event loop.",
+        validation_alias=AliasChoices(
+            AliasPath("unit_test_loop_debug"),
+            "prefect_testing_unit_test_loop_debug",
+            "prefect_unit_test_loop_debug",
+        ),
+    )
+
+    test_setting: Optional[Any] = Field(
+        default="FOO",
+        description="This setting only exists to facilitate unit testing. If in test mode, this setting will return its value. Otherwise, it returns `None`.",
+        validation_alias=AliasChoices(
+            AliasPath("test_setting"),
+            "prefect_testing_test_setting",
+            "prefect_test_setting",
+        ),
+    )

prefect/settings/models/worker.py

@@ -0,0 +1,46 @@
+from pydantic import Field
+from pydantic_settings import SettingsConfigDict
+
+from prefect.settings.base import PrefectBaseSettings
+
+
+class WorkerWebserverSettings(PrefectBaseSettings):
+    model_config = SettingsConfigDict(
+        env_prefix="PREFECT_WORKER_WEBSERVER_", env_file=".env", extra="ignore"
+    )
+
+    host: str = Field(
+        default="0.0.0.0",
+        description="The host address the worker's webserver should bind to.",
+    )
+
+    port: int = Field(
+        default=8080,
+        description="The port the worker's webserver should bind to.",
+    )
+
+
+class WorkerSettings(PrefectBaseSettings):
+    model_config = SettingsConfigDict(
+        env_prefix="PREFECT_WORKER_", env_file=".env", extra="ignore"
+    )
+
+    heartbeat_seconds: float = Field(
+        default=30,
+        description="Number of seconds a worker should wait between sending a heartbeat.",
+    )
+
+    query_seconds: float = Field(
+        default=10,
+        description="Number of seconds a worker should wait between queries for scheduled work.",
+    )
+
+    prefetch_seconds: float = Field(
+        default=10,
+        description="The number of seconds into the future a worker should query for scheduled work.",
+    )
+
+    webserver: WorkerWebserverSettings = Field(
+        default_factory=WorkerWebserverSettings,
+        description="Settings for a worker's webserver",
+    )

prefect/settings/profiles.py

@@ -0,0 +1,390 @@
+import inspect
+import warnings
+from pathlib import Path
+from typing import Annotated, Any, Dict, Iterable, List, Optional, Set, Tuple, Union
+
+import toml
+from pydantic import (
+    BaseModel,
+    BeforeValidator,
+    ConfigDict,
+    Field,
+    TypeAdapter,
+    ValidationError,
+)
+from pydantic_settings import BaseSettings
+
+from prefect.exceptions import ProfileSettingsValidationError
+from prefect.settings.constants import DEFAULT_PROFILES_PATH
+from prefect.settings.context import get_current_settings
+from prefect.settings.legacy import Setting, _get_settings_fields
+from prefect.settings.models.root import Settings
+
+
+def _cast_settings(
+    settings: Union[Dict[Union[str, Setting], Any], Any],
+) -> Dict[Setting, Any]:
+    """For backwards compatibility, allow either Settings objects as keys or string references to settings."""
+    if not isinstance(settings, dict):
+        raise ValueError("Settings must be a dictionary.")
+    casted_settings = {}
+    for k, value in settings.items():
+        try:
+            if isinstance(k, str):
+                setting = _get_settings_fields(Settings)[k]
+            else:
+                setting = k
+            casted_settings[setting] = value
+        except KeyError as e:
+            warnings.warn(f"Setting {e} is not recognized")
+            continue
+    return casted_settings
+
+
+############################################################################
+# Profiles
+
+
+class Profile(BaseModel):
+    """A user profile containing settings."""
+
+    model_config = ConfigDict(extra="ignore", arbitrary_types_allowed=True)
+
+    name: str
+    settings: Annotated[Dict[Setting, Any], BeforeValidator(_cast_settings)] = Field(
+        default_factory=dict
+    )
+    source: Optional[Path] = None
+
+    def to_environment_variables(self) -> Dict[str, str]:
+        """Convert the profile settings to a dictionary of environment variables."""
+        return {
+            setting.name: str(value)
+            for setting, value in self.settings.items()
+            if value is not None
+        }
+
+    def validate_settings(self):
+        errors: List[Tuple[Setting, ValidationError]] = []
+        for setting, value in self.settings.items():
+            try:
+                model_fields = Settings.model_fields
+                annotation = None
+                for section in setting.accessor.split("."):
+                    annotation = model_fields[section].annotation
+                    if inspect.isclass(annotation) and issubclass(
+                        annotation, BaseSettings
+                    ):
+                        model_fields = annotation.model_fields
+
+                TypeAdapter(annotation).validate_python(value)
+            except ValidationError as e:
+                errors.append((setting, e))
+        if errors:
+            raise ProfileSettingsValidationError(errors)
+
+
+class ProfilesCollection:
+    """ "
+    A utility class for working with a collection of profiles.
+
+    Profiles in the collection must have unique names.
+
+    The collection may store the name of the active profile.
+    """
+
+    def __init__(
+        self, profiles: Iterable[Profile], active: Optional[str] = None
+    ) -> None:
+        self.profiles_by_name = {profile.name: profile for profile in profiles}
+        self.active_name = active
+
+    @property
+    def names(self) -> Set[str]:
+        """
+        Return a set of profile names in this collection.
+        """
+        return set(self.profiles_by_name.keys())
+
+    @property
+    def active_profile(self) -> Optional[Profile]:
+        """
+        Retrieve the active profile in this collection.
+        """
+        if self.active_name is None:
+            return None
+        return self[self.active_name]
+
+    def set_active(self, name: Optional[str], check: bool = True):
+        """
+        Set the active profile name in the collection.
+
+        A null value may be passed to indicate that this collection does not determine
+        the active profile.
+        """
+        if check and name is not None and name not in self.names:
+            raise ValueError(f"Unknown profile name {name!r}.")
+        self.active_name = name
+
+    def update_profile(
+        self,
+        name: str,
+        settings: Dict[Setting, Any],
+        source: Optional[Path] = None,
+    ) -> Profile:
+        """
+        Add a profile to the collection or update the existing on if the name is already
+        present in this collection.
+
+        If updating an existing profile, the settings will be merged. Settings can
+        be dropped from the existing profile by setting them to `None` in the new
+        profile.
+
+        Returns the new profile object.
+        """
+        existing = self.profiles_by_name.get(name)
+
+        # Convert the input to a `Profile` to cast settings to the correct type
+        profile = Profile(name=name, settings=settings, source=source)
+
+        if existing:
+            new_settings = {**existing.settings, **profile.settings}
+
+            # Drop null keys to restore to default
+            for key, value in tuple(new_settings.items()):
+                if value is None:
+                    new_settings.pop(key)
+
+            new_profile = Profile(
+                name=profile.name,
+                settings=new_settings,
+                source=source or profile.source,
+            )
+        else:
+            new_profile = profile
+
+        self.profiles_by_name[new_profile.name] = new_profile
+
+        return new_profile
+
+    def add_profile(self, profile: Profile) -> None:
+        """
+        Add a profile to the collection.
+
+        If the profile name already exists, an exception will be raised.
+        """
+        if profile.name in self.profiles_by_name:
+            raise ValueError(
+                f"Profile name {profile.name!r} already exists in collection."
+            )
+
+        self.profiles_by_name[profile.name] = profile
+
+    def remove_profile(self, name: str) -> None:
+        """
+        Remove a profile from the collection.
+        """
+        self.profiles_by_name.pop(name)
+
+    def without_profile_source(self, path: Optional[Path]) -> "ProfilesCollection":
+        """
+        Remove profiles that were loaded from a given path.
+
+        Returns a new collection.
+        """
+        return ProfilesCollection(
+            [
+                profile
+                for profile in self.profiles_by_name.values()
+                if profile.source != path
+            ],
+            active=self.active_name,
+        )
+
+    def to_dict(self):
+        """
+        Convert to a dictionary suitable for writing to disk.
+        """
+        return {
+            "active": self.active_name,
+            "profiles": {
+                profile.name: profile.to_environment_variables()
+                for profile in self.profiles_by_name.values()
+            },
+        }
+
+    def __getitem__(self, name: str) -> Profile:
+        return self.profiles_by_name[name]
+
+    def __iter__(self):
+        return self.profiles_by_name.__iter__()
+
+    def items(self):
+        return self.profiles_by_name.items()
+
+    def __eq__(self, __o: object) -> bool:
+        if not isinstance(__o, ProfilesCollection):
+            return False
+
+        return (
+            self.profiles_by_name == __o.profiles_by_name
+            and self.active_name == __o.active_name
+        )
+
+    def __repr__(self) -> str:
+        return (
+            f"ProfilesCollection(profiles={list(self.profiles_by_name.values())!r},"
+            f" active={self.active_name!r})>"
+        )
+
+
+def _read_profiles_from(path: Path) -> ProfilesCollection:
+    """
+    Read profiles from a path into a new `ProfilesCollection`.
+
+    Profiles are expected to be written in TOML with the following schema:
+        ```
+        active = <name: Optional[str]>
+
+        [profiles.<name: str>]
+        <SETTING: str> = <value: Any>
+        ```
+    """
+    contents = toml.loads(path.read_text())
+    active_profile = contents.get("active")
+    raw_profiles = contents.get("profiles", {})
+
+    profiles = []
+    for name, settings in raw_profiles.items():
+        profiles.append(Profile(name=name, settings=settings, source=path))
+
+    return ProfilesCollection(profiles, active=active_profile)
+
+
+def _write_profiles_to(path: Path, profiles: ProfilesCollection) -> None:
+    """
+    Write profiles in the given collection to a path as TOML.
+
+    Any existing data not present in the given `profiles` will be deleted.
+    """
+    if not path.exists():
+        path.parent.mkdir(parents=True, exist_ok=True)
+        path.touch(mode=0o600)
+    path.write_text(toml.dumps(profiles.to_dict()))
+
+
+def load_profiles(include_defaults: bool = True) -> ProfilesCollection:
+    """
+    Load profiles from the current profile path. Optionally include profiles from the
+    default profile path.
+    """
+    current_settings = get_current_settings()
+    default_profiles = _read_profiles_from(DEFAULT_PROFILES_PATH)
+
+    if current_settings.profiles_path is None:
+        raise RuntimeError(
+            "No profiles path set; please ensure `PREFECT_PROFILES_PATH` is set."
+        )
+
+    if not include_defaults:
+        if not current_settings.profiles_path.exists():
+            return ProfilesCollection([])
+        return _read_profiles_from(current_settings.profiles_path)
+
+    user_profiles_path = current_settings.profiles_path
+    profiles = default_profiles
+    if user_profiles_path.exists():
+        user_profiles = _read_profiles_from(user_profiles_path)
+
+        # Merge all of the user profiles with the defaults
+        for name in user_profiles:
+            if not (source := user_profiles[name].source):
+                raise ValueError(f"Profile {name!r} has no source.")
+            profiles.update_profile(
+                name,
+                settings=user_profiles[name].settings,
+                source=source,
+            )
+
+        if user_profiles.active_name:
+            profiles.set_active(user_profiles.active_name, check=False)
+
+    return profiles
+
+
+def load_current_profile():
+    """
+    Load the current profile from the default and current profile paths.
+
+    This will _not_ include settings from the current settings context. Only settings
+    that have been persisted to the profiles file will be saved.
+    """
+    import prefect.context
+
+    profiles = load_profiles()
+    context = prefect.context.get_settings_context()
+
+    if context:
+        profiles.set_active(context.profile.name)
+
+    return profiles.active_profile
+
+
+def save_profiles(profiles: ProfilesCollection) -> None:
+    """
+    Writes all non-default profiles to the current profiles path.
+    """
+    profiles_path = get_current_settings().profiles_path
+    assert profiles_path is not None, "Profiles path is not set."
+    profiles = profiles.without_profile_source(DEFAULT_PROFILES_PATH)
+    return _write_profiles_to(profiles_path, profiles)
+
+
+def load_profile(name: str) -> Profile:
+    """
+    Load a single profile by name.
+    """
+    profiles = load_profiles()
+    try:
+        return profiles[name]
+    except KeyError:
+        raise ValueError(f"Profile {name!r} not found.")
+
+
+def update_current_profile(
+    settings: Dict[Union[str, Setting], Any],
+) -> Profile:
+    """
+    Update the persisted data for the profile currently in-use.
+
+    If the profile does not exist in the profiles file, it will be created.
+
+    Given settings will be merged with the existing settings as described in
+    `ProfilesCollection.update_profile`.
+
+    Returns:
+        The new profile.
+    """
+    import prefect.context
+
+    current_profile = prefect.context.get_settings_context().profile
+
+    if not current_profile:
+        from prefect.exceptions import MissingProfileError
+
+        raise MissingProfileError("No profile is currently in use.")
+
+    profiles = load_profiles()
+
+    # Ensure the current profile's settings are present
+    profiles.update_profile(current_profile.name, current_profile.settings)
+    # Then merge the new settings in
+    new_profile = profiles.update_profile(
+        current_profile.name, _cast_settings(settings)
+    )
+
+    new_profile.validate_settings()
+
+    save_profiles(profiles)
+
+    return profiles[current_profile.name]