prefect-client 3.1.5__py3-none-any.whl → 3.1.7__py3-none-any.whl

prefect/deployments/steps/core.py

@@ -99,7 +99,9 @@ def _get_function_for_step(
     return step_func
 
 
-async def run_step(step: Dict, upstream_outputs: Optional[Dict] = None) -> Dict:
+async def run_step(
+    step: dict[str, Any], upstream_outputs: Optional[dict[str, Any]] = None
+) -> dict[str, Any]:
     """
     Runs a step, returns the step's output.
 

prefect/deployments/steps/pull.py

@@ -6,10 +6,11 @@ import os
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Optional
 
+from prefect._internal.compatibility.async_dispatch import async_dispatch
 from prefect._internal.retries import retry_async_fn
 from prefect.logging.loggers import get_logger
 from prefect.runner.storage import BlockStorageAdapter, GitRepository, RemoteStorage
-from prefect.utilities.asyncutils import sync_compatible
+from prefect.utilities.asyncutils import run_coro_as_sync
 
 deployment_logger = get_logger("deployment")
 
@@ -17,7 +18,7 @@ if TYPE_CHECKING:
     from prefect.blocks.core import Block
 
 
-def set_working_directory(directory: str) -> dict:
+def set_working_directory(directory: str) -> dict[str, str]:
     """
     Sets the working directory; works with both absolute and relative paths.
 
@@ -37,15 +38,64 @@ def set_working_directory(directory: str) -> dict:
     base_delay=1,
     max_delay=10,
     retry_on_exceptions=(RuntimeError,),
+    operation_name="git_clone",
 )
-@sync_compatible
-async def git_clone(
+async def _pull_git_repository_with_retries(repo: GitRepository):
+    await repo.pull_code()
+
+
+async def agit_clone(
+    repository: str,
+    branch: Optional[str] = None,
+    include_submodules: bool = False,
+    access_token: Optional[str] = None,
+    credentials: Optional["Block"] = None,
+) -> dict[str, str]:
+    """
+    Asynchronously clones a git repository into the current working directory.
+
+    Args:
+        repository: the URL of the repository to clone
+        branch: the branch to clone; if not provided, the default branch will be used
+        include_submodules (bool): whether to include git submodules when cloning the repository
+        access_token: an access token to use for cloning the repository; if not provided
+            the repository will be cloned using the default git credentials
+        credentials: a GitHubCredentials, GitLabCredentials, or BitBucketCredentials block can be used to specify the
+            credentials to use for cloning the repository.
+
+    Returns:
+        dict: a dictionary containing a `directory` key of the new directory that was created
+
+    Raises:
+        subprocess.CalledProcessError: if the git clone command fails for any reason
+    """
+    if access_token and credentials:
+        raise ValueError(
+            "Please provide either an access token or credentials but not both."
+        )
+
+    _credentials = {"access_token": access_token} if access_token else credentials
+
+    storage = GitRepository(
+        url=repository,
+        credentials=_credentials,
+        branch=branch,
+        include_submodules=include_submodules,
+    )
+
+    await _pull_git_repository_with_retries(storage)
+
+    return dict(directory=str(storage.destination.relative_to(Path.cwd())))
+
+
+@async_dispatch(agit_clone)
+def git_clone(
     repository: str,
     branch: Optional[str] = None,
     include_submodules: bool = False,
     access_token: Optional[str] = None,
     credentials: Optional["Block"] = None,
-) -> dict:
+) -> dict[str, str]:
     """
     Clones a git repository into the current working directory.
 
@@ -120,20 +170,18 @@ async def git_clone(
             "Please provide either an access token or credentials but not both."
         )
 
-    credentials = {"access_token": access_token} if access_token else credentials
+    _credentials = {"access_token": access_token} if access_token else credentials
 
     storage = GitRepository(
         url=repository,
-        credentials=credentials,
+        credentials=_credentials,
         branch=branch,
         include_submodules=include_submodules,
     )
 
-    await storage.pull_code()
+    run_coro_as_sync(_pull_git_repository_with_retries(storage))
 
-    directory = str(storage.destination.relative_to(Path.cwd()))
-    deployment_logger.info(f"Cloned repository {repository!r} into {directory!r}")
-    return {"directory": directory}
+    return dict(directory=str(storage.destination.relative_to(Path.cwd())))
 
 
 async def pull_from_remote_storage(url: str, **settings: Any):
@@ -190,7 +238,7 @@ async def pull_with_block(block_document_name: str, block_type_slug: str):
 
     full_slug = f"{block_type_slug}/{block_document_name}"
     try:
-        block = await Block.load(full_slug)
+        block = await Block.aload(full_slug)
     except Exception:
         deployment_logger.exception("Unable to load block '%s'", full_slug)
         raise

prefect/docker/__init__.py

@@ -17,4 +17,4 @@ def __getattr__(name: str) -> object:
         module, attr = _public_api[name]
         return getattr(import_module(module), attr)
 
-    raise ImportError(f"module {__name__!r} has no attribute {name!r}")
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

prefect/events/clients.py

@@ -1,11 +1,13 @@
 import abc
 import asyncio
+import os
 from types import TracebackType
 from typing import (
     TYPE_CHECKING,
     Any,
     ClassVar,
     Dict,
+    Generator,
     List,
     MutableMapping,
     Optional,
@@ -13,20 +15,22 @@ from typing import (
     Type,
     cast,
 )
+from urllib.parse import urlparse
 from uuid import UUID
 
 import orjson
 import pendulum
 from cachetools import TTLCache
 from prometheus_client import Counter
+from python_socks.async_.asyncio import Proxy
 from typing_extensions import Self
 from websockets import Subprotocol
-from websockets.client import WebSocketClientProtocol, connect
 from websockets.exceptions import (
     ConnectionClosed,
     ConnectionClosedError,
     ConnectionClosedOK,
 )
+from websockets.legacy.client import Connect, WebSocketClientProtocol
 
 from prefect.events import Event
 from prefect.logging import get_logger
@@ -80,6 +84,53 @@ def events_out_socket_from_api_url(url: str):
     return http_to_ws(url) + "/events/out"
 
 
+class WebsocketProxyConnect(Connect):
+    def __init__(self: Self, uri: str, **kwargs: Any):
+        # super() is intentionally deferred to the _proxy_connect method
+        # to allow for the socket to be established first
+
+        self.uri = uri
+        self._kwargs = kwargs
+
+        u = urlparse(uri)
+        host = u.hostname
+
+        if u.scheme == "ws":
+            port = u.port or 80
+            proxy_url = os.environ.get("HTTP_PROXY")
+        elif u.scheme == "wss":
+            port = u.port or 443
+            proxy_url = os.environ.get("HTTPS_PROXY")
+            kwargs["server_hostname"] = host
+        else:
+            raise ValueError(
+                "Unsupported scheme %s. Expected 'ws' or 'wss'. " % u.scheme
+            )
+
+        self._proxy = Proxy.from_url(proxy_url) if proxy_url else None
+        self._host = host
+        self._port = port
+
+    async def _proxy_connect(self: Self) -> WebSocketClientProtocol:
+        if self._proxy:
+            sock = await self._proxy.connect(
+                dest_host=self._host,
+                dest_port=self._port,
+            )
+            self._kwargs["sock"] = sock
+
+        super().__init__(self.uri, **self._kwargs)
+        proto = await self.__await_impl__()
+        return proto
+
+    def __await__(self: Self) -> Generator[Any, None, WebSocketClientProtocol]:
+        return self._proxy_connect().__await__()
+
+
+def websocket_connect(uri: str, **kwargs: Any) -> WebsocketProxyConnect:
+    return WebsocketProxyConnect(uri, **kwargs)
+
+
 def get_events_client(
     reconnection_attempts: int = 10,
     checkpoint_every: int = 700,
@@ -265,7 +316,7 @@ class PrefectEventsClient(EventsClient):
             )
 
         self._events_socket_url = events_in_socket_from_api_url(api_url)
-        self._connect = connect(self._events_socket_url)
+        self._connect = websocket_connect(self._events_socket_url)
         self._websocket = None
         self._reconnection_attempts = reconnection_attempts
         self._unconfirmed_events = []
@@ -435,7 +486,7 @@ class PrefectCloudEventsClient(PrefectEventsClient):
             reconnection_attempts=reconnection_attempts,
             checkpoint_every=checkpoint_every,
         )
-        self._connect = connect(
+        self._connect = websocket_connect(
             self._events_socket_url,
             extra_headers={"Authorization": f"bearer {api_key}"},
         )
@@ -494,7 +545,7 @@ class PrefectEventSubscriber:
 
         logger.debug("Connecting to %s", socket_url)
 
-        self._connect = connect(
+        self._connect = websocket_connect(
             socket_url,
             subprotocols=[Subprotocol("prefect")],
         )

prefect/events/filters.py

@@ -3,9 +3,9 @@ from uuid import UUID
 
 import pendulum
 from pydantic import Field, PrivateAttr
-from pydantic_extra_types.pendulum_dt import DateTime
 
 from prefect._internal.schemas.bases import PrefectBaseModel
+from prefect.types import DateTime
 from prefect.utilities.collections import AutoEnum
 
 from .schemas.events import Event, Resource, ResourceSpecification

prefect/events/related.py

@@ -15,7 +15,8 @@ from typing import (
 from uuid import UUID
 
 import pendulum
-from pendulum.datetime import DateTime
+
+from prefect.types import DateTime
 
 from .schemas.events import RelatedResource
 

prefect/events/schemas/events.py

@@ -13,17 +13,21 @@ from typing import (
 )
 from uuid import UUID, uuid4
 
-import pendulum
-from pydantic import ConfigDict, Field, RootModel, field_validator, model_validator
-from pydantic_extra_types.pendulum_dt import DateTime
-from typing_extensions import Self
+from pydantic import (
+    AfterValidator,
+    ConfigDict,
+    Field,
+    RootModel,
+    model_validator,
+)
+from typing_extensions import Annotated, Self
 
 from prefect._internal.schemas.bases import PrefectBaseModel
 from prefect.logging import get_logger
 from prefect.settings import (
     PREFECT_EVENTS_MAXIMUM_LABELS_PER_RESOURCE,
-    PREFECT_EVENTS_MAXIMUM_RELATED_RESOURCES,
 )
+from prefect.types import DateTime
 
 from .labelling import Labelled
 
@@ -90,22 +94,34 @@ class RelatedResource(Resource):
         return self["prefect.resource.role"]
 
 
+def _validate_related_resources(value) -> List:
+    from prefect.settings import PREFECT_EVENTS_MAXIMUM_RELATED_RESOURCES
+
+    if len(value) > PREFECT_EVENTS_MAXIMUM_RELATED_RESOURCES.value():
+        raise ValueError(
+            "The maximum number of related resources "
+            f"is {PREFECT_EVENTS_MAXIMUM_RELATED_RESOURCES.value()}"
+        )
+    return value
+
+
 class Event(PrefectBaseModel):
     """The client-side view of an event that has happened to a Resource"""
 
     model_config = ConfigDict(extra="ignore")
 
     occurred: DateTime = Field(
-        default_factory=lambda: pendulum.now("UTC"),
+        default_factory=lambda: DateTime.now("UTC"),
         description="When the event happened from the sender's perspective",
     )
-    event: str = Field(
-        description="The name of the event that happened",
-    )
+    event: str = Field(description="The name of the event that happened")
     resource: Resource = Field(
         description="The primary Resource this event concerns",
     )
-    related: List[RelatedResource] = Field(
+    related: Annotated[
+        List[RelatedResource],
+        AfterValidator(_validate_related_resources),
+    ] = Field(
         default_factory=list,
         description="A list of additional Resources involved in this event",
     )
@@ -144,17 +160,6 @@ class Event(PrefectBaseModel):
             resources[related.role].append(related)
         return resources
 
-    @field_validator("related")
-    @classmethod
-    def enforce_maximum_related_resources(cls, value: List[RelatedResource]):
-        if len(value) > PREFECT_EVENTS_MAXIMUM_RELATED_RESOURCES.value():
-            raise ValueError(
-                "The maximum number of related resources "
-                f"is {PREFECT_EVENTS_MAXIMUM_RELATED_RESOURCES.value()}"
-            )
-
-        return value
-
     def find_resource_label(self, label: str) -> Optional[str]:
         """Finds the value of the given label in this event's resource or one of its
         related resources.  If the label starts with `related:<role>:`, search for the

prefect/events/utilities.py

@@ -3,7 +3,6 @@ from typing import Any, Dict, List, Optional, Union
 from uuid import UUID
 
 import pendulum
-from pydantic_extra_types.pendulum_dt import DateTime
 
 from .clients import (
     AssertingEventsClient,
@@ -20,11 +19,12 @@ TIGHT_TIMING = timedelta(minutes=5)
 def emit_event(
     event: str,
     resource: Dict[str, str],
-    occurred: Optional[DateTime] = None,
+    occurred: Optional[pendulum.DateTime] = None,
     related: Optional[Union[List[Dict[str, str]], List[RelatedResource]]] = None,
     payload: Optional[Dict[str, Any]] = None,
     id: Optional[UUID] = None,
     follows: Optional[Event] = None,
+    **kwargs: Optional[Dict[str, Any]],
 ) -> Optional[Event]:
     """
     Send an event to Prefect Cloud.
@@ -63,6 +63,7 @@ def emit_event(
     event_kwargs: Dict[str, Any] = {
         "event": event,
         "resource": resource,
+        **kwargs,
     }
 
     if occurred is None:

prefect/events/worker.py

@@ -83,6 +83,14 @@ class EventsWorker(QueueService[Event]):
         await self._client.emit(event)
 
     async def attach_related_resources_from_context(self, event: Event):
+        if "prefect.resource.lineage-group" in event.resource:
+            # We attach related resources to lineage events in `emit_lineage_event`,
+            # instead of the worker, because not all run-related resources are
+            # upstream from every lineage event (they might be downstream).
+            # The "related" field in the event schema tracks upstream resources
+            # only.
+            return
+
         exclude = {resource.id for resource in event.involved_resources}
         event.related += await related_resources_from_run_context(
             client=self._orchestration_client, exclude=exclude

prefect/filesystems.py

@@ -1,6 +1,7 @@
 import abc
 import urllib.parse
 from pathlib import Path
+from shutil import copytree
 from typing import Any, Dict, Optional
 
 import anyio
@@ -13,7 +14,6 @@ from prefect._internal.schemas.validators import (
 )
 from prefect.blocks.core import Block
 from prefect.utilities.asyncutils import run_sync_in_worker_thread, sync_compatible
-from prefect.utilities.compat import copytree
 from prefect.utilities.filesystem import filter_files
 
 from ._internal.compatibility.migration import getattr_migration
@@ -158,7 +158,7 @@ class LocalFileSystem(WritableFileSystem, WritableDeploymentStorage):
         copytree(from_path, local_path, dirs_exist_ok=True, ignore=ignore_func)
 
     async def _get_ignore_func(self, local_path: str, ignore_file: str):
-        with open(ignore_file, "r") as f:
+        with open(ignore_file) as f:
             ignore_patterns = f.readlines()
         included_files = filter_files(root=local_path, ignore_patterns=ignore_patterns)
 
@@ -348,7 +348,7 @@ class RemoteFileSystem(WritableFileSystem, WritableDeploymentStorage):
 
         included_files = None
         if ignore_file:
-            with open(ignore_file, "r") as f:
+            with open(ignore_file) as f:
                 ignore_patterns = f.readlines()
 
             included_files = filter_files(