prefect-client 3.1.6__py3-none-any.whl → 3.1.8__py3-none-any.whl

prefect/deployments/runner.py

@@ -33,7 +33,7 @@ import importlib
 import tempfile
 from datetime import datetime, timedelta
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Union
+from typing import TYPE_CHECKING, Any, Iterable, List, Optional, Union
 from uuid import UUID
 
 from pydantic import (
@@ -160,7 +160,7 @@ class RunnerDeployment(BaseModel):
     paused: Optional[bool] = Field(
         default=None, description="Whether or not the deployment is paused."
     )
-    parameters: Dict[str, Any] = Field(default_factory=dict)
+    parameters: dict[str, Any] = Field(default_factory=dict)
     entrypoint: Optional[str] = Field(
         default=None,
         description=(
@@ -198,7 +198,7 @@ class RunnerDeployment(BaseModel):
             " the deployment is registered with a built runner."
         ),
     )
-    job_variables: Dict[str, Any] = Field(
+    job_variables: dict[str, Any] = Field(
         default_factory=dict,
         description=(
             "Job variables used to override the default values of a work pool"
@@ -280,7 +280,7 @@ class RunnerDeployment(BaseModel):
         async with get_client() as client:
             flow_id = await client.create_flow_from_name(self.flow_name)
 
-            create_payload = dict(
+            create_payload: dict[str, Any] = dict(
                 flow_id=flow_id,
                 name=self.name,
                 work_queue_name=self.work_queue_name,
@@ -428,7 +428,7 @@ class RunnerDeployment(BaseModel):
         else:
             return [create_deployment_schedule_create(schedule)]
 
-    def _set_defaults_from_flow(self, flow: "Flow"):
+    def _set_defaults_from_flow(self, flow: "Flow[..., Any]"):
         self._parameter_openapi_schema = parameter_schema(flow)
 
         if not self.version:
@@ -439,7 +439,7 @@ class RunnerDeployment(BaseModel):
     @classmethod
     def from_flow(
         cls,
-        flow: "Flow",
+        flow: "Flow[..., Any]",
         name: str,
         interval: Optional[
             Union[Iterable[Union[int, float, timedelta]], int, float, timedelta]
@@ -449,7 +449,7 @@ class RunnerDeployment(BaseModel):
         paused: Optional[bool] = None,
         schedules: Optional["FlexibleScheduleList"] = None,
         concurrency_limit: Optional[Union[int, ConcurrencyLimitConfig, None]] = None,
-        parameters: Optional[dict] = None,
+        parameters: Optional[dict[str, Any]] = None,
         triggers: Optional[List[Union[DeploymentTriggerTypes, TriggerTypes]]] = None,
         description: Optional[str] = None,
         tags: Optional[List[str]] = None,
@@ -457,7 +457,7 @@ class RunnerDeployment(BaseModel):
         enforce_parameter_schema: bool = True,
         work_pool_name: Optional[str] = None,
         work_queue_name: Optional[str] = None,
-        job_variables: Optional[Dict[str, Any]] = None,
+        job_variables: Optional[dict[str, Any]] = None,
         entrypoint_type: EntrypointType = EntrypointType.FILE_PATH,
     ) -> "RunnerDeployment":
         """
@@ -588,7 +588,7 @@ class RunnerDeployment(BaseModel):
         paused: Optional[bool] = None,
         schedules: Optional["FlexibleScheduleList"] = None,
         concurrency_limit: Optional[Union[int, ConcurrencyLimitConfig, None]] = None,
-        parameters: Optional[dict] = None,
+        parameters: Optional[dict[str, Any]] = None,
         triggers: Optional[List[Union[DeploymentTriggerTypes, TriggerTypes]]] = None,
         description: Optional[str] = None,
         tags: Optional[List[str]] = None,
@@ -596,7 +596,7 @@ class RunnerDeployment(BaseModel):
         enforce_parameter_schema: bool = True,
         work_pool_name: Optional[str] = None,
         work_queue_name: Optional[str] = None,
-        job_variables: Optional[Dict[str, Any]] = None,
+        job_variables: Optional[dict[str, Any]] = None,
     ) -> "RunnerDeployment":
         """
         Configure a deployment for a given flow located at a given entrypoint.
@@ -689,7 +689,7 @@ class RunnerDeployment(BaseModel):
         paused: Optional[bool] = None,
         schedules: Optional["FlexibleScheduleList"] = None,
         concurrency_limit: Optional[Union[int, ConcurrencyLimitConfig, None]] = None,
-        parameters: Optional[dict] = None,
+        parameters: Optional[dict[str, Any]] = None,
         triggers: Optional[List[Union[DeploymentTriggerTypes, TriggerTypes]]] = None,
         description: Optional[str] = None,
         tags: Optional[List[str]] = None,
@@ -697,7 +697,7 @@ class RunnerDeployment(BaseModel):
         enforce_parameter_schema: bool = True,
         work_pool_name: Optional[str] = None,
         work_queue_name: Optional[str] = None,
-        job_variables: Optional[Dict[str, Any]] = None,
+        job_variables: Optional[dict[str, Any]] = None,
     ):
         """
         Create a RunnerDeployment from a flow located at a given entrypoint and stored in a
@@ -945,8 +945,8 @@ async def deploy(
 
         console.print(f"Successfully pushed image {image.reference!r}", style="green")
 
-    deployment_exceptions = []
-    deployment_ids = []
+    deployment_exceptions: list[dict[str, Any]] = []
+    deployment_ids: list[UUID] = []
     image_ref = image.reference if image else None
     for deployment in track(
         deployments,

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/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

@@ -20,7 +20,6 @@ from pydantic import (
     RootModel,
     model_validator,
 )
-from pydantic_extra_types.pendulum_dt import DateTime
 from typing_extensions import Annotated, Self
 
 from prefect._internal.schemas.bases import PrefectBaseModel
@@ -28,6 +27,7 @@ from prefect.logging import get_logger
 from prefect.settings import (
     PREFECT_EVENTS_MAXIMUM_LABELS_PER_RESOURCE,
 )
+from prefect.types import DateTime
 
 from .labelling import Labelled
 

prefect/events/utilities.py

@@ -24,6 +24,7 @@ def emit_event(
     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.
@@ -62,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