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

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

prefect/flow_engine.py

@@ -24,10 +24,8 @@ from uuid import UUID
 
 from anyio import CancelScope
 from opentelemetry import propagate, trace
-from opentelemetry.trace import Tracer, get_tracer
 from typing_extensions import ParamSpec
 
-import prefect
 from prefect import Task
 from prefect.client.orchestration import PrefectClient, SyncPrefectClient, get_client
 from prefect.client.schemas import FlowRun, TaskRun
@@ -72,7 +70,12 @@ from prefect.states import (
     exception_to_failed_state,
     return_value_to_state,
 )
-from prefect.telemetry.run_telemetry import OTELSetter
+from prefect.telemetry.run_telemetry import (
+    LABELS_TRACEPARENT_KEY,
+    TRACEPARENT_KEY,
+    OTELSetter,
+    RunTelemetry,
+)
 from prefect.types import KeyValueLabels
 from prefect.utilities._engine import get_hook_name, resolve_custom_flow_run_name
 from prefect.utilities.annotations import NotSet
@@ -95,8 +98,6 @@ from prefect.utilities.urls import url_for
 
 P = ParamSpec("P")
 R = TypeVar("R")
-LABELS_TRACEPARENT_KEY = "__OTEL_TRACEPARENT"
-TRACEPARENT_KEY = "traceparent"
 
 
 class FlowRunTimeoutError(TimeoutError):
@@ -136,10 +137,7 @@ class BaseFlowRunEngine(Generic[P, R]):
     _is_started: bool = False
     short_circuit: bool = False
     _flow_run_name_set: bool = False
-    _tracer: Tracer = field(
-        default_factory=lambda: get_tracer("prefect", prefect.__version__)
-    )
-    _span: Optional[trace.Span] = None
+    _telemetry: RunTelemetry = field(default_factory=RunTelemetry)
 
     def __post_init__(self):
         if self.flow is None and self.flow_run_id is None:
@@ -152,21 +150,6 @@ class BaseFlowRunEngine(Generic[P, R]):
     def state(self) -> State:
         return self.flow_run.state  # type: ignore
 
-    def _end_span_on_success(self):
-        if not self._span:
-            return
-        self._span.set_status(trace.Status(trace.StatusCode.OK))
-        self._span.end(time.time_ns())
-        self._span = None
-
-    def _end_span_on_error(self, exc: BaseException, description: Optional[str]):
-        if not self._span:
-            return
-        self._span.record_exception(exc)
-        self._span.set_status(trace.Status(trace.StatusCode.ERROR, description))
-        self._span.end(time.time_ns())
-        self._span = None
-
     def is_running(self) -> bool:
         if getattr(self, "flow_run", None) is None:
             return False
@@ -185,6 +168,7 @@ class BaseFlowRunEngine(Generic[P, R]):
         self, span: trace.Span, client: Union[SyncPrefectClient, PrefectClient]
     ):
         parent_flow_run_ctx = FlowRunContext.get()
+
         if parent_flow_run_ctx and parent_flow_run_ctx.flow_run:
             if traceparent := parent_flow_run_ctx.flow_run.labels.get(
                 LABELS_TRACEPARENT_KEY
@@ -194,6 +178,7 @@ class BaseFlowRunEngine(Generic[P, R]):
                     carrier={TRACEPARENT_KEY: traceparent},
                     setter=OTELSetter(),
                 )
+
             else:
                 carrier: KeyValueLabels = {}
                 propagate.get_global_textmap().inject(
@@ -315,16 +300,7 @@ class FlowRunEngine(BaseFlowRunEngine[P, R]):
         self.flow_run.state_name = state.name  # type: ignore
         self.flow_run.state_type = state.type  # type: ignore
 
-        if self._span:
-            self._span.add_event(
-                state.name or state.type,
-                {
-                    "prefect.state.message": state.message or "",
-                    "prefect.state.type": state.type,
-                    "prefect.state.name": state.name or state.type,
-                    "prefect.state.id": str(state.id),
-                },
-            )
+        self._telemetry.update_state(state)
         return state
 
     def result(self, raise_on_failure: bool = True) -> "Union[R, State, None]":
@@ -374,7 +350,7 @@ class FlowRunEngine(BaseFlowRunEngine[P, R]):
         self.set_state(terminal_state)
         self._return_value = resolved_result
 
-        self._end_span_on_success()
+        self._telemetry.end_span_on_success()
 
         return result
 
@@ -406,8 +382,8 @@ class FlowRunEngine(BaseFlowRunEngine[P, R]):
             )
             state = self.set_state(Running())
         self._raised = exc
-
-        self._end_span_on_error(exc, state.message)
+        self._telemetry.record_exception(exc)
+        self._telemetry.end_span_on_failure(state.message)
 
         return state
 
@@ -426,8 +402,8 @@ class FlowRunEngine(BaseFlowRunEngine[P, R]):
         )
         self.set_state(state)
         self._raised = exc
-
-        self._end_span_on_error(exc, message)
+        self._telemetry.record_exception(exc)
+        self._telemetry.end_span_on_failure(message)
 
     def handle_crash(self, exc: BaseException) -> None:
         state = run_coro_as_sync(exception_to_crashed_state(exc))
@@ -435,8 +411,8 @@ class FlowRunEngine(BaseFlowRunEngine[P, R]):
         self.logger.debug("Crash details:", exc_info=exc)
         self.set_state(state, force=True)
         self._raised = exc
-
-        self._end_span_on_error(exc, state.message if state else "")
+        self._telemetry.record_exception(exc)
+        self._telemetry.end_span_on_failure(state.message if state else None)
 
     def load_subflow_run(
         self,
@@ -681,19 +657,12 @@ class FlowRunEngine(BaseFlowRunEngine[P, R]):
                     empirical_policy=self.flow_run.empirical_policy,
                 )
 
-            span = self._tracer.start_span(
-                name=self.flow_run.name,
-                attributes={
-                    **self.flow_run.labels,
-                    "prefect.run.type": "flow",
-                    "prefect.run.id": str(self.flow_run.id),
-                    "prefect.tags": self.flow_run.tags,
-                    "prefect.flow.name": self.flow.name,
-                },
+            self._telemetry.start_span(
+                name=self.flow.name,
+                run=self.flow_run,
+                client=self.client,
+                parameters=self.parameters,
             )
-            self._update_otel_labels(span, self.client)
-
-            self._span = span
 
             try:
                 yield self
@@ -736,7 +705,9 @@ class FlowRunEngine(BaseFlowRunEngine[P, R]):
     @contextmanager
     def start(self) -> Generator[None, None, None]:
         with self.initialize_run():
-            with trace.use_span(self._span) if self._span else nullcontext():
+            with trace.use_span(
+                self._telemetry.span
+            ) if self._telemetry.span else nullcontext():
                 self.begin_run()
 
                 if self.state.is_running():
@@ -892,16 +863,7 @@ class AsyncFlowRunEngine(BaseFlowRunEngine[P, R]):
         self.flow_run.state_name = state.name  # type: ignore
         self.flow_run.state_type = state.type  # type: ignore
 
-        if self._span:
-            self._span.add_event(
-                state.name or state.type,
-                {
-                    "prefect.state.message": state.message or "",
-                    "prefect.state.type": state.type,
-                    "prefect.state.name": state.name or state.type,
-                    "prefect.state.id": str(state.id),
-                },
-            )
+        self._telemetry.update_state(state)
         return state
 
     async def result(self, raise_on_failure: bool = True) -> "Union[R, State, None]":
@@ -949,7 +911,7 @@ class AsyncFlowRunEngine(BaseFlowRunEngine[P, R]):
         await self.set_state(terminal_state)
         self._return_value = resolved_result
 
-        self._end_span_on_success()
+        self._telemetry.end_span_on_success()
 
         return result
 
@@ -979,8 +941,8 @@ class AsyncFlowRunEngine(BaseFlowRunEngine[P, R]):
             )
             state = await self.set_state(Running())
         self._raised = exc
-
-        self._end_span_on_error(exc, state.message)
+        self._telemetry.record_exception(exc)
+        self._telemetry.end_span_on_failure(state.message)
 
         return state
 
@@ -1000,7 +962,8 @@ class AsyncFlowRunEngine(BaseFlowRunEngine[P, R]):
         await self.set_state(state)
         self._raised = exc
 
-        self._end_span_on_error(exc, message)
+        self._telemetry.record_exception(exc)
+        self._telemetry.end_span_on_failure(message)
 
     async def handle_crash(self, exc: BaseException) -> None:
         # need to shield from asyncio cancellation to ensure we update the state
@@ -1012,7 +975,8 @@ class AsyncFlowRunEngine(BaseFlowRunEngine[P, R]):
             await self.set_state(state, force=True)
             self._raised = exc
 
-            self._end_span_on_error(exc, state.message)
+            self._telemetry.record_exception(exc)
+            self._telemetry.end_span_on_failure(state.message)
 
     async def load_subflow_run(
         self,
@@ -1255,18 +1219,12 @@ class AsyncFlowRunEngine(BaseFlowRunEngine[P, R]):
                     empirical_policy=self.flow_run.empirical_policy,
                 )
 
-            span = self._tracer.start_span(
-                name=self.flow_run.name,
-                attributes={
-                    **self.flow_run.labels,
-                    "prefect.run.type": "flow",
-                    "prefect.run.id": str(self.flow_run.id),
-                    "prefect.tags": self.flow_run.tags,
-                    "prefect.flow.name": self.flow.name,
-                },
+            await self._telemetry.async_start_span(
+                name=self.flow.name,
+                run=self.flow_run,
+                client=self.client,
+                parameters=self.parameters,
             )
-            self._update_otel_labels(span, self.client)
-            self._span = span
 
             try:
                 yield self
@@ -1309,7 +1267,9 @@ class AsyncFlowRunEngine(BaseFlowRunEngine[P, R]):
     @asynccontextmanager
     async def start(self) -> AsyncGenerator[None, None]:
         async with self.initialize_run():
-            with trace.use_span(self._span) if self._span else nullcontext():
+            with trace.use_span(
+                self._telemetry.span
+            ) if self._telemetry.span else nullcontext():
                 await self.begin_run()
 
                 if self.state.is_running():

prefect/flow_runs.py

@@ -1,6 +1,6 @@
 from typing import (
     TYPE_CHECKING,
-    Dict,
+    Any,
     Optional,
     Type,
     TypeVar,
@@ -430,7 +430,9 @@ async def suspend_flow_run(
 
 
 @sync_compatible
-async def resume_flow_run(flow_run_id, run_input: Optional[Dict] = None):
+async def resume_flow_run(
+    flow_run_id: UUID, run_input: Optional[dict[str, Any]] = None
+) -> None:
     """
     Resumes a paused flow.
 

prefect/flows.py

@@ -564,14 +564,12 @@ class Flow(Generic[P, R]):
                 "Cannot mix Pydantic v1 and v2 types as arguments to a flow."
             )
 
+        validated_fn_kwargs = dict(arbitrary_types_allowed=True)
+
         if has_v1_models:
-            validated_fn = V1ValidatedFunction(
-                self.fn, config={"arbitrary_types_allowed": True}
-            )
+            validated_fn = V1ValidatedFunction(self.fn, config=validated_fn_kwargs)
         else:
-            validated_fn = V2ValidatedFunction(
-                self.fn, config=pydantic.ConfigDict(arbitrary_types_allowed=True)
-            )
+            validated_fn = V2ValidatedFunction(self.fn, config=validated_fn_kwargs)
 
         try:
             with warnings.catch_warnings():