prefect-client 3.1.9__py3-none-any.whl → 3.1.11__py3-none-any.whl

prefect/results.py

@@ -1,4 +1,3 @@
-import abc
 import inspect
 import os
 import socket
@@ -11,11 +10,11 @@ from typing import (
     Annotated,
     Any,
     Callable,
+    ClassVar,
     Dict,
     Generic,
     Optional,
     Tuple,
-    Type,
     TypeVar,
     Union,
 )
@@ -28,13 +27,10 @@ from pydantic import (
     ConfigDict,
     Discriminator,
     Field,
-    PrivateAttr,
     Tag,
     ValidationError,
-    model_serializer,
     model_validator,
 )
-from pydantic_core import PydanticUndefinedType
 from typing_extensions import ParamSpec, Self
 
 import prefect
@@ -42,10 +38,7 @@ from prefect._experimental.lineage import (
     emit_result_read_event,
     emit_result_write_event,
 )
-from prefect._internal.compatibility import deprecated
-from prefect._internal.compatibility.deprecated import deprecated_field
 from prefect.blocks.core import Block
-from prefect.client.utilities import inject_client
 from prefect.exceptions import (
     ConfigurationError,
     MissingContextError,
@@ -63,11 +56,9 @@ from prefect.settings.context import get_current_settings
 from prefect.types import DateTime
 from prefect.utilities.annotations import NotSet
 from prefect.utilities.asyncutils import sync_compatible
-from prefect.utilities.pydantic import get_dispatch_key, lookup_type, register_base_type
 
 if TYPE_CHECKING:
     from prefect import Flow, Task
-    from prefect.client.orchestration import PrefectClient
     from prefect.transactions import IsolationLevel
 
 
@@ -245,7 +236,7 @@ async def _call_explicitly_async_block_method(
     see https://github.com/PrefectHQ/prefect/issues/15008
     """
     if hasattr(block, f"a{method}"):  # explicit async method
-        return await getattr(block.__class__.__name__, f"a{method}")(*args, **kwargs)
+        return await getattr(block, f"a{method}")(*args, **kwargs)
     elif hasattr(getattr(block, method, None), "aio"):  # sync_compatible
         return await getattr(block, method).aio(block, *args, **kwargs)
     else:  # should not happen in prefect, but users can override impls
@@ -276,13 +267,6 @@ def result_storage_discriminator(x: Any) -> str:
     return "None"
 
 
-@deprecated_field(
-    "persist_result",
-    when=lambda x: x is not None,
-    when_message="use the `should_persist_result` utility function instead",
-    start_date="Sep 2024",
-    end_date="Nov 2024",
-)
 class ResultStore(BaseModel):
     """
     Manages the storage and retrieval of results.
@@ -294,13 +278,12 @@ class ResultStore(BaseModel):
             the metadata will be stored alongside the results.
         lock_manager: The lock manager to use for locking result records. If not provided,
             the store cannot be used in transactions with the SERIALIZABLE isolation level.
-        persist_result: Whether to persist results.
         cache_result_in_memory: Whether to cache results in memory.
         serializer: The serializer to use for results.
         storage_key_fn: The function to generate storage keys.
     """
 
-    model_config = ConfigDict(arbitrary_types_allowed=True)
+    model_config: ClassVar[ConfigDict] = ConfigDict(arbitrary_types_allowed=True)
 
     result_storage: Optional[WritableFileSystem] = Field(default=None)
     metadata_storage: Annotated[
@@ -317,9 +300,6 @@ class ResultStore(BaseModel):
     storage_key_fn: Callable[[], str] = Field(default=DEFAULT_STORAGE_KEY_FN)
     cache: LRUCache[str, "ResultRecord[Any]"] = Field(default_factory=default_cache)
 
-    # Deprecated fields
-    persist_result: Optional[bool] = Field(default=None)
-
     @property
     def result_storage_block_id(self) -> Optional[UUID]:
         if self.result_storage is None:
@@ -903,52 +883,6 @@ class ResultStore(BaseModel):
             )
         return await self.lock_manager.await_for_lock(key, timeout)
 
-    @deprecated.deprecated_callable(
-        start_date="Sep 2024",
-        end_date="Nov 2024",
-        help="Use `create_result_record` instead.",
-    )
-    @sync_compatible
-    async def create_result(
-        self,
-        obj: R,
-        key: Optional[str] = None,
-        expiration: Optional[DateTime] = None,
-    ) -> Union[R, "BaseResult[R]"]:
-        """
-        Create a `PersistedResult` for the given object.
-        """
-        # Null objects are "cached" in memory at no cost
-        should_cache_object = self.cache_result_in_memory or obj is None
-        should_persist_result = (
-            self.persist_result
-            if self.persist_result is not None
-            else not should_cache_object
-        )
-
-        if key:
-
-            def key_fn():
-                return key
-
-            storage_key_fn = key_fn
-        else:
-            storage_key_fn = self.storage_key_fn
-
-        if self.result_storage is None:
-            self.result_storage = await get_default_result_storage()
-
-        return await PersistedResult.create(
-            obj,
-            storage_block=self.result_storage,
-            storage_block_id=self.result_storage_block_id,
-            storage_key_fn=storage_key_fn,
-            serializer=self.serializer,
-            cache_object=should_cache_object,
-            expiration=expiration,
-            serialize_to_none=not should_persist_result,
-        )
-
     # TODO: These two methods need to find a new home
 
     @sync_compatible
@@ -1221,242 +1155,3 @@ class ResultRecord(BaseModel, Generic[R]):
         if not isinstance(other, ResultRecord):
             return False
         return self.metadata == other.metadata and self.result == other.result
-
-
-@deprecated.deprecated_class(
-    start_date="Sep 2024", end_date="Nov 2024", help="Use `ResultRecord` instead."
-)
-@register_base_type
-class BaseResult(BaseModel, abc.ABC, Generic[R]):
-    model_config = ConfigDict(extra="forbid")
-    type: str
-
-    def __init__(self, **data: Any) -> None:
-        type_string = (
-            get_dispatch_key(self) if type(self) is not BaseResult else "__base__"
-        )
-        data.setdefault("type", type_string)
-        super().__init__(**data)
-
-    def __new__(cls: Type[Self], **kwargs) -> Self:
-        if "type" in kwargs:
-            try:
-                subcls = lookup_type(cls, dispatch_key=kwargs["type"])
-            except KeyError as exc:
-                raise ValueError(f"Invalid type: {kwargs['type']}") from exc
-            return super().__new__(subcls)
-        else:
-            return super().__new__(cls)
-
-    _cache: Any = PrivateAttr(NotSet)
-
-    def _cache_object(self, obj: Any) -> None:
-        self._cache = obj
-
-    def has_cached_object(self) -> bool:
-        return self._cache is not NotSet
-
-    @abc.abstractmethod
-    @sync_compatible
-    async def get(self) -> R:
-        ...
-
-    @abc.abstractclassmethod
-    @sync_compatible
-    async def create(
-        cls: "Type[BaseResult[R]]",
-        obj: R,
-        **kwargs: Any,
-    ) -> "BaseResult[R]":
-        ...
-
-    @classmethod
-    def __dispatch_key__(cls, **kwargs):
-        default = cls.model_fields.get("type").get_default()
-        return cls.__name__ if isinstance(default, PydanticUndefinedType) else default
-
-
-@deprecated.deprecated_class(
-    start_date="Sep 2024", end_date="Nov 2024", help="Use `ResultRecord` instead."
-)
-class PersistedResult(BaseResult):
-    """
-    Result type which stores a reference to a persisted result.
-
-    When created, the user's object is serialized and stored. The format for the content
-    is defined by `ResultRecord`. This reference contains metadata necessary for retrieval
-    of the object, such as a reference to the storage block and the key where the
-    content was written.
-    """
-
-    type: str = "reference"
-
-    serializer_type: str
-    storage_key: str
-    storage_block_id: Optional[uuid.UUID] = None
-    expiration: Optional[DateTime] = None
-    serialize_to_none: bool = False
-
-    _persisted: bool = PrivateAttr(default=False)
-    _should_cache_object: bool = PrivateAttr(default=True)
-    _storage_block: WritableFileSystem = PrivateAttr(default=None)
-    _serializer: Serializer = PrivateAttr(default=None)
-
-    @model_serializer(mode="wrap")
-    def serialize_model(self, handler, info):
-        if self.serialize_to_none:
-            return None
-        return handler(self, info)
-
-    def _cache_object(
-        self,
-        obj: Any,
-        storage_block: WritableFileSystem = None,
-        serializer: Serializer = None,
-    ) -> None:
-        self._cache = obj
-        self._storage_block = storage_block
-        self._serializer = serializer
-
-    @inject_client
-    async def _get_storage_block(self, client: "PrefectClient") -> WritableFileSystem:
-        if self._storage_block is not None:
-            return self._storage_block
-        elif self.storage_block_id is not None:
-            block_document = await client.read_block_document(self.storage_block_id)
-            self._storage_block = Block._from_block_document(block_document)
-        else:
-            self._storage_block = await get_default_result_storage()
-        return self._storage_block
-
-    @sync_compatible
-    @inject_client
-    async def get(
-        self, ignore_cache: bool = False, client: "PrefectClient" = None
-    ) -> R:
-        """
-        Retrieve the data and deserialize it into the original object.
-        """
-        if self.has_cached_object() and not ignore_cache:
-            return self._cache
-
-        result_store_kwargs = {}
-        if self._serializer:
-            result_store_kwargs["serializer"] = resolve_serializer(self._serializer)
-        storage_block = await self._get_storage_block(client=client)
-        result_store = ResultStore(result_storage=storage_block, **result_store_kwargs)
-
-        record = await result_store.aread(self.storage_key)
-        self.expiration = record.expiration
-
-        if self._should_cache_object:
-            self._cache_object(record.result)
-
-        return record.result
-
-    @staticmethod
-    def _infer_path(storage_block, key) -> str:
-        """
-        Attempts to infer a path associated with a storage block key, this method will
-        defer to the block in the future
-        """
-
-        if hasattr(storage_block, "_resolve_path"):
-            return storage_block._resolve_path(key)
-        if hasattr(storage_block, "_remote_file_system"):
-            return storage_block._remote_file_system._resolve_path(key)
-
-    @sync_compatible
-    @inject_client
-    async def write(self, obj: R = NotSet, client: "PrefectClient" = None) -> None:
-        """
-        Write the result to the storage block.
-        """
-        if self._persisted or self.serialize_to_none:
-            # don't double write or overwrite
-            return
-
-        # load objects from a cache
-
-        # first the object itself
-        if obj is NotSet and not self.has_cached_object():
-            raise ValueError("Cannot write a result that has no object cached.")
-        obj = obj if obj is not NotSet else self._cache
-
-        # next, the storage block
-        storage_block = await self._get_storage_block(client=client)
-
-        # finally, the serializer
-        serializer = self._serializer
-        if serializer is None:
-            # this could error if the serializer requires kwargs
-            serializer = Serializer(type=self.serializer_type)
-
-        result_store = ResultStore(
-            result_storage=storage_block,
-            serializer=serializer,
-        )
-        await result_store.awrite(
-            obj=obj, key=self.storage_key, expiration=self.expiration
-        )
-
-        self._persisted = True
-
-        if not self._should_cache_object:
-            self._cache = NotSet
-
-    @classmethod
-    @sync_compatible
-    async def create(
-        cls: "Type[PersistedResult]",
-        obj: R,
-        storage_block: WritableFileSystem,
-        storage_key_fn: Callable[[], str],
-        serializer: Serializer,
-        storage_block_id: Optional[uuid.UUID] = None,
-        cache_object: bool = True,
-        expiration: Optional[DateTime] = None,
-        serialize_to_none: bool = False,
-    ) -> "PersistedResult[R]":
-        """
-        Create a new result reference from a user's object.
-
-        The object will be serialized and written to the storage block under a unique
-        key. It will then be cached on the returned result.
-        """
-        key = storage_key_fn()
-        if not isinstance(key, str):
-            raise TypeError(
-                f"Expected type 'str' for result storage key; got value {key!r}"
-            )
-        uri = cls._infer_path(storage_block, key)
-
-        # in this case we store an absolute path
-        if storage_block_id is None and uri is not None:
-            key = str(uri)
-
-        result = cls(
-            serializer_type=serializer.type,
-            storage_block_id=storage_block_id,
-            storage_key=key,
-            expiration=expiration,
-            serialize_to_none=serialize_to_none,
-        )
-
-        object.__setattr__(result, "_should_cache_object", cache_object)
-        # we must cache temporarily to allow for writing later
-        # the cache will be removed on write
-        result._cache_object(obj, storage_block=storage_block, serializer=serializer)
-
-        return result
-
-    def __eq__(self, other):
-        if not isinstance(other, PersistedResult):
-            return False
-        return (
-            self.type == other.type
-            and self.serializer_type == other.serializer_type
-            and self.storage_key == other.storage_key
-            and self.storage_block_id == other.storage_block_id
-            and self.expiration == other.expiration
-        )

prefect/runner/storage.py

@@ -2,7 +2,15 @@ import shutil
 import subprocess
 from copy import deepcopy
 from pathlib import Path
-from typing import Any, Dict, Optional, Protocol, TypedDict, Union, runtime_checkable
+from typing import (
+    Any,
+    Dict,
+    Optional,
+    Protocol,
+    TypedDict,
+    Union,
+    runtime_checkable,
+)
 from urllib.parse import urlparse, urlsplit, urlunparse
 from uuid import uuid4
 
@@ -87,6 +95,7 @@ class GitRepository:
         pull_interval: The interval in seconds at which to pull contents from
             remote storage to local storage. If None, remote storage will perform
             a one-time sync.
+        directories: The directories to pull from the Git repository (uses git sparse-checkout)
 
     Examples:
         Pull the contents of a private git repository to the local filesystem:
@@ -111,6 +120,7 @@ class GitRepository:
         branch: Optional[str] = None,
         include_submodules: bool = False,
         pull_interval: Optional[int] = 60,
+        directories: Optional[str] = None,
     ):
         if credentials is None:
             credentials = {}
@@ -134,6 +144,7 @@ class GitRepository:
         self._logger = get_logger(f"runner.storage.git-repository.{self._name}")
         self._storage_base_path = Path.cwd()
         self._pull_interval = pull_interval
+        self._directories = directories
 
     @property
     def destination(self) -> Path:
@@ -147,11 +158,9 @@ class GitRepository:
         return self._pull_interval
 
     @property
-    def _repository_url_with_credentials(self) -> str:
+    def _formatted_credentials(self) -> Optional[str]:
         if not self._credentials:
-            return self._url
-
-        url_components = urlparse(self._url)
+            return None
 
         credentials = (
             self._credentials.model_dump()
@@ -165,18 +174,52 @@ class GitRepository:
             elif isinstance(v, SecretStr):
                 credentials[k] = v.get_secret_value()
 
-        formatted_credentials = _format_token_from_credentials(
-            urlparse(self._url).netloc, credentials
+        return _format_token_from_credentials(urlparse(self._url).netloc, credentials)
+
+    def _add_credentials_to_url(self, url: str) -> str:
+        """Add credentials to given url if possible."""
+        components = urlparse(url)
+        credentials = self._formatted_credentials
+
+        if components.scheme != "https" or not credentials:
+            return url
+
+        return urlunparse(
+            components._replace(netloc=f"{credentials}@{components.netloc}")
         )
-        if url_components.scheme == "https" and formatted_credentials is not None:
-            updated_components = url_components._replace(
-                netloc=f"{formatted_credentials}@{url_components.netloc}"
-            )
-            repository_url = urlunparse(updated_components)
-        else:
-            repository_url = self._url
 
-        return repository_url
+    @property
+    def _repository_url_with_credentials(self) -> str:
+        return self._add_credentials_to_url(self._url)
+
+    @property
+    def _git_config(self) -> list[str]:
+        """Build a git configuration to use when running git commands."""
+        config = {}
+
+        # Submodules can be private. The url in .gitmodules
+        # will not include the credentials, we need to
+        # propagate them down here if they exist.
+        if self._include_submodules and self._formatted_credentials:
+            base_url = urlparse(self._url)._replace(path="")
+            without_auth = urlunparse(base_url)
+            with_auth = self._add_credentials_to_url(without_auth)
+            config[f"url.{with_auth}.insteadOf"] = without_auth
+
+        return ["-c", " ".join(f"{k}={v}" for k, v in config.items())] if config else []
+
+    async def is_sparsely_checked_out(self) -> bool:
+        """
+        Check if existing repo is sparsely checked out
+        """
+
+        try:
+            result = await run_process(
+                ["git", "config", "--get", "core.sparseCheckout"], cwd=self.destination
+            )
+            return result.strip().lower() == "true"
+        except Exception:
+            return False
 
     async def pull_code(self):
         """
@@ -206,9 +249,20 @@ class GitRepository:
                     f"does not match the configured repository {self._url}"
                 )
 
+            # Sparsely checkout the repository if directories are specified and the repo is not in sparse-checkout mode already
+            if self._directories and not await self.is_sparsely_checked_out():
+                await run_process(
+                    ["git", "sparse-checkout", "set"] + self._directories,
+                    cwd=self.destination,
+                )
+
             self._logger.debug("Pulling latest changes from origin/%s", self._branch)
             # Update the existing repository
-            cmd = ["git", "pull", "origin"]
+            cmd = ["git"]
+            # Add the git configuration, must be given after `git` and before the command
+            cmd += self._git_config
+            # Add the pull command and parameters
+            cmd += ["pull", "origin"]
             if self._branch:
                 cmd += [self._branch]
             if self._include_submodules:
@@ -234,17 +288,21 @@ class GitRepository:
         self._logger.debug("Cloning repository %s", self._url)
 
         repository_url = self._repository_url_with_credentials
+        cmd = ["git"]
+        # Add the git configuration, must be given after `git` and before the command
+        cmd += self._git_config
+        # Add the clone command and its parameters
+        cmd += ["clone", repository_url]
 
-        cmd = [
-            "git",
-            "clone",
-            repository_url,
-        ]
         if self._branch:
             cmd += ["--branch", self._branch]
         if self._include_submodules:
             cmd += ["--recurse-submodules"]
 
+        # This will only checkout the top-level directory
+        if self._directories:
+            cmd += ["--sparse"]
+
         # Limit git history and set path to clone to
         cmd += ["--depth", "1", str(self.destination)]
 
@@ -258,6 +316,14 @@ class GitRepository:
                 f" {exc.returncode}."
             ) from exc_chain
 
+        # Once repository is cloned and the repo is in sparse-checkout mode then grow the working directory
+        if self._directories:
+            self._logger.debug("Will add %s", self._directories)
+            await run_process(
+                ["git", "sparse-checkout", "set"] + self._directories,
+                cwd=self.destination,
+            )
+
     def __eq__(self, __value) -> bool:
         if isinstance(__value, GitRepository):
             return (