prefect-client 3.4.24.dev2__tar.gz → 3.4.24.dev3__tar.gz

{prefect_client-3.4.24.dev2 → prefect_client-3.4.24.dev3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: prefect-client
-Version: 3.4.24.dev2
+Version: 3.4.24.dev3
 Summary: Workflow orchestration and management.
 Project-URL: Changelog, https://github.com/PrefectHQ/prefect/releases
 Project-URL: Documentation, https://docs.prefect.io
@@ -50,7 +50,6 @@ Requires-Dist: pydantic-extra-types<3.0.0,>=2.8.2
 Requires-Dist: pydantic-settings!=2.9.0,<3.0.0,>2.2.1
 Requires-Dist: python-dateutil<3.0.0,>=2.8.2
 Requires-Dist: python-slugify<9.0,>=5.0
-Requires-Dist: python-socks[asyncio]<3.0,>=2.5.3
 Requires-Dist: pytz<2026,>=2021.1
 Requires-Dist: pyyaml<7.0.0,>=5.4.1
 Requires-Dist: rfc3339-validator<0.2.0,>=0.1.4
@@ -60,7 +59,7 @@ Requires-Dist: sniffio<2.0.0,>=1.3.0
 Requires-Dist: toml>=0.10.0
 Requires-Dist: typing-extensions<5.0.0,>=4.10.0
 Requires-Dist: uvicorn!=0.29.0,>=0.14.0
-Requires-Dist: websockets<16.0,>=13.0
+Requires-Dist: websockets<16.0,>=15.0.1
 Requires-Dist: whenever<0.10.0,>=0.7.3; python_version >= '3.13'
 Provides-Extra: notifications
 Requires-Dist: apprise<2.0.0,>=1.1.0; extra == 'notifications'

{prefect_client-3.4.24.dev2 → prefect_client-3.4.24.dev3}/pyproject.toml

@@ -54,7 +54,6 @@ dependencies = [
     "pydantic_settings>2.2.1,<3.0.0,!=2.9.0",
     "python-dateutil>=2.8.2,<3.0.0",
     "python-slugify>=5.0,<9.0",
-    "python-socks[asyncio]>=2.5.3,<3.0",
     "pytz>=2021.1,<2026",
     "pyyaml>=5.4.1,<7.0.0",
     "rfc3339-validator>=0.1.4,<0.2.0",
@@ -64,7 +63,7 @@ dependencies = [
     "toml>=0.10.0",
     "typing_extensions>=4.10.0,<5.0.0",
     "uvicorn>=0.14.0,!=0.29.0",
-    "websockets>=13.0,<16.0",
+    "websockets>=15.0.1,<16.0",
     "whenever>=0.7.3,<0.10.0; python_version>='3.13'",
 ]
 [project.urls]

prefect_client-3.4.24.dev3/src/prefect/_build_info.py

@@ -0,0 +1,5 @@
+# Generated by versioningit
+__version__ = "3.4.24.dev3"
+__build_date__ = "2025-10-14 08:08:51.459069+00:00"
+__git_commit__ = "5b27fc1a52ea3ca42c97918911d260537231db96"
+__dirty__ = False

prefect_client-3.4.24.dev3/src/prefect/_internal/websockets.py

@@ -0,0 +1,88 @@
+"""
+Internal WebSocket proxy utilities for Prefect client connections.
+
+This module provides shared WebSocket proxy connection logic and SSL configuration
+to avoid duplication between events and logs clients.
+"""
+
+import ssl
+import warnings
+from functools import wraps
+from typing import Any, Optional
+from urllib.parse import urlparse
+
+import certifi
+from websockets.asyncio.client import connect
+
+from prefect.settings import get_current_settings
+
+
+def create_ssl_context_for_websocket(uri: str) -> Optional[ssl.SSLContext]:
+    """Create SSL context for WebSocket connections based on URI scheme."""
+    u = urlparse(uri)
+
+    if u.scheme != "wss":
+        return None
+
+    if get_current_settings().api.tls_insecure_skip_verify:
+        # Create an unverified context for insecure connections
+        ctx = ssl.create_default_context()
+        ctx.check_hostname = False
+        ctx.verify_mode = ssl.CERT_NONE
+        return ctx
+    else:
+        # Create a verified context with the certificate file
+        cert_file = get_current_settings().api.ssl_cert_file
+        if not cert_file:
+            cert_file = certifi.where()
+        return ssl.create_default_context(cafile=cert_file)
+
+
+@wraps(connect)
+def websocket_connect(uri: str, **kwargs: Any) -> connect:
+    """
+    Create a WebSocket connection with proxy and SSL support.
+
+    Proxy support is automatic via HTTP_PROXY/HTTPS_PROXY environment variables.
+    The websockets library handles proxy detection and connection automatically.
+    """
+    # Configure SSL context for HTTPS connections
+    ssl_context = create_ssl_context_for_websocket(uri)
+    if ssl_context:
+        kwargs.setdefault("ssl", ssl_context)
+
+    # Add custom headers from settings
+    custom_headers = get_current_settings().client.custom_headers
+    if custom_headers:
+        # Get existing additional_headers or create new dict
+        additional_headers = kwargs.get("additional_headers", {})
+        if not isinstance(additional_headers, dict):
+            additional_headers = {}
+
+        for header_name, header_value in custom_headers.items():
+            # Check for protected headers that shouldn't be overridden
+            if header_name.lower() in {
+                "user-agent",
+                "sec-websocket-key",
+                "sec-websocket-version",
+                "sec-websocket-extensions",
+                "sec-websocket-protocol",
+                "connection",
+                "upgrade",
+                "host",
+            }:
+                warnings.warn(
+                    f"Custom header '{header_name}' is ignored because it conflicts with "
+                    f"a protected WebSocket header. Protected headers include: "
+                    f"User-Agent, Sec-WebSocket-Key, Sec-WebSocket-Version, "
+                    f"Sec-WebSocket-Extensions, Sec-WebSocket-Protocol, Connection, "
+                    f"Upgrade, Host",
+                    UserWarning,
+                    stacklevel=2,
+                )
+            else:
+                additional_headers[header_name] = header_value
+
+        kwargs["additional_headers"] = additional_headers
+
+    return connect(uri, **kwargs)

{prefect_client-3.4.24.dev2 → prefect_client-3.4.24.dev3}/src/prefect/client/subscriptions.py

@@ -11,7 +11,7 @@ from starlette.status import WS_1008_POLICY_VIOLATION
 from typing_extensions import Self
 
 from prefect._internal.schemas.bases import IDBaseModel
-from prefect.events.clients import websocket_connect
+from prefect._internal.websockets import websocket_connect
 from prefect.logging import get_logger
 from prefect.settings import get_current_settings
 

{prefect_client-3.4.24.dev2 → prefect_client-3.4.24.dev3}/src/prefect/events/__init__.py

@@ -43,6 +43,7 @@ from .actions import (
     DeclareIncident,
 )
 from .clients import get_events_client, get_events_subscriber
+from .subscribers import FlowRunSubscriber
 from .utilities import emit_event
 
 __all__ = [
@@ -90,4 +91,5 @@ __all__ = [
     "emit_event",
     "get_events_client",
     "get_events_subscriber",
+    "FlowRunSubscriber",
 ]

prefect_client-3.4.24.dev3/src/prefect/events/subscribers.py

@@ -0,0 +1,186 @@
+"""
+Flow run subscriber that interleaves events and logs from a flow run
+"""
+
+from __future__ import annotations
+
+import asyncio
+from types import TracebackType
+from typing import TYPE_CHECKING, Any, Optional, Union
+from uuid import UUID
+
+from typing_extensions import Self
+
+from prefect.client.schemas.filters import LogFilter, LogFilterFlowRunId
+from prefect.client.schemas.objects import TERMINAL_STATES, Log, StateType
+from prefect.events import Event
+from prefect.events.clients import get_events_subscriber
+from prefect.events.filters import EventAnyResourceFilter, EventFilter
+from prefect.logging.clients import get_logs_subscriber
+
+if TYPE_CHECKING:
+    from prefect.events.clients import PrefectEventSubscriber
+    from prefect.logging.clients import PrefectLogsSubscriber
+
+
+class FlowRunSubscriber:
+    """
+    Subscribes to both events and logs for a specific flow run, yielding them
+    in an interleaved stream.
+
+    This subscriber combines the event stream and log stream for a flow run into
+    a single async iterator. When a terminal event (Completed, Failed, or Crashed)
+    is received, the event subscription stops but log subscription continues for a
+    configurable timeout to catch any straggler logs.
+
+    Example:
+        ```python
+        from prefect.events.subscribers import FlowRunSubscriber
+
+        async with FlowRunSubscriber(flow_run_id=my_flow_run_id) as subscriber:
+            async for item in subscriber:
+                if isinstance(item, Event):
+                    print(f"Event: {item.event}")
+                else:  # isinstance(item, Log)
+                    print(f"Log: {item.message}")
+        ```
+    """
+
+    _flow_run_id: UUID
+    _queue: asyncio.Queue[Union[Log, Event, None]]
+    _tasks: list[asyncio.Task[None]]
+    _flow_completed: bool
+    _straggler_timeout: int
+    _reconnection_attempts: int
+    _log_filter: LogFilter
+    _event_filter: EventFilter
+    _logs_subscriber: PrefectLogsSubscriber | Any
+    _events_subscriber: PrefectEventSubscriber | Any
+    _sentinels_received: int
+
+    def __init__(
+        self,
+        flow_run_id: UUID,
+        straggler_timeout: int = 3,
+        reconnection_attempts: int = 10,
+    ):
+        """
+        Args:
+            flow_run_id: The ID of the flow run to follow
+            straggler_timeout: After a terminal event, how long (in seconds) to wait
+                for additional logs before stopping
+            reconnection_attempts: Number of times to attempt reconnection if
+                the websocket connection is lost
+        """
+        self._flow_run_id = flow_run_id
+        self._straggler_timeout = straggler_timeout
+        self._reconnection_attempts = reconnection_attempts
+        self._queue = asyncio.Queue()
+        self._tasks = []
+        self._flow_completed = False
+        self._sentinels_received = 0
+
+        self._log_filter = LogFilter(flow_run_id=LogFilterFlowRunId(any_=[flow_run_id]))
+        self._event_filter = EventFilter(
+            any_resource=EventAnyResourceFilter(id=[f"prefect.flow-run.{flow_run_id}"])
+        )
+
+        self._logs_subscriber = None
+        self._events_subscriber = None
+
+    async def __aenter__(self) -> Self:
+        """Enter the async context manager"""
+        self._logs_subscriber = get_logs_subscriber(
+            filter=self._log_filter, reconnection_attempts=self._reconnection_attempts
+        )
+        self._events_subscriber = get_events_subscriber(
+            filter=self._event_filter, reconnection_attempts=self._reconnection_attempts
+        )
+
+        await self._logs_subscriber.__aenter__()
+        await self._events_subscriber.__aenter__()
+
+        self._tasks = [
+            asyncio.create_task(self._consume_logs()),
+            asyncio.create_task(self._consume_events()),
+        ]
+
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: Optional[type[BaseException]],
+        exc_val: Optional[BaseException],
+        exc_tb: Optional[TracebackType],
+    ) -> None:
+        """Exit the async context manager and clean up resources"""
+        for task in self._tasks:
+            task.cancel()
+
+        await asyncio.gather(*self._tasks, return_exceptions=True)
+
+        await self._logs_subscriber.__aexit__(exc_type, exc_val, exc_tb)
+        await self._events_subscriber.__aexit__(exc_type, exc_val, exc_tb)
+
+    def __aiter__(self) -> Self:
+        """Return self as an async iterator"""
+        return self
+
+    async def __anext__(self) -> Union[Log, Event]:
+        """Get the next log or event from the interleaved stream"""
+        while self._sentinels_received < len(self._tasks):
+            if self._flow_completed:
+                try:
+                    item = await asyncio.wait_for(
+                        self._queue.get(), timeout=self._straggler_timeout
+                    )
+                except asyncio.TimeoutError:
+                    raise StopAsyncIteration
+            else:
+                item = await self._queue.get()
+
+            if item is None:
+                self._sentinels_received += 1
+                continue
+
+            return item
+
+        raise StopAsyncIteration
+
+    async def _consume_logs(self) -> None:
+        """Background task to consume logs and put them in the queue"""
+        try:
+            async for log in self._logs_subscriber:
+                await self._queue.put(log)
+        except asyncio.CancelledError:
+            pass
+        except Exception:
+            pass
+        finally:
+            await self._queue.put(None)
+
+    async def _consume_events(self) -> None:
+        """Background task to consume events and put them in the queue"""
+        try:
+            async for event in self._events_subscriber:
+                await self._queue.put(event)
+
+                # Check if this is a terminal state event for our flow run
+                if event.resource.id == f"prefect.flow-run.{self._flow_run_id}":
+                    # Get state type from event resource or payload
+                    state_type_str = event.resource.get("prefect.state-type")
+                    if not state_type_str and "validated_state" in event.payload:
+                        state_type_str = event.payload["validated_state"].get("type")
+
+                    if state_type_str:
+                        try:
+                            state_type = StateType(state_type_str)
+                            if state_type in TERMINAL_STATES:
+                                self._flow_completed = True
+                                break
+                        except ValueError:
+                            pass
+        except Exception:
+            pass
+        finally:
+            await self._queue.put(None)

prefect_client-3.4.24.dev2/src/prefect/_build_info.py

@@ -1,5 +0,0 @@
-# Generated by versioningit
-__version__ = "3.4.24.dev2"
-__build_date__ = "2025-10-13 08:09:36.212484+00:00"
-__git_commit__ = "2c75d8cd6ca1f17e98e98c7335c17201c6343c43"
-__dirty__ = False

prefect_client-3.4.24.dev2/src/prefect/_internal/websockets.py

@@ -1,144 +0,0 @@
-"""
-Internal WebSocket proxy utilities for Prefect client connections.
-
-This module provides shared WebSocket proxy connection logic and SSL configuration
-to avoid duplication between events and logs clients.
-"""
-
-import os
-import ssl
-import warnings
-from typing import Any, Generator, Optional
-from urllib.parse import urlparse
-from urllib.request import proxy_bypass
-
-import certifi
-from python_socks.async_.asyncio import Proxy
-from typing_extensions import Self
-from websockets.asyncio.client import ClientConnection, connect
-
-from prefect.settings import get_current_settings
-
-
-def create_ssl_context_for_websocket(uri: str) -> Optional[ssl.SSLContext]:
-    """Create SSL context for WebSocket connections based on URI scheme."""
-    u = urlparse(uri)
-
-    if u.scheme != "wss":
-        return None
-
-    if get_current_settings().api.tls_insecure_skip_verify:
-        # Create an unverified context for insecure connections
-        ctx = ssl.create_default_context()
-        ctx.check_hostname = False
-        ctx.verify_mode = ssl.CERT_NONE
-        return ctx
-    else:
-        # Create a verified context with the certificate file
-        cert_file = get_current_settings().api.ssl_cert_file
-        if not cert_file:
-            cert_file = certifi.where()
-        return ssl.create_default_context(cafile=cert_file)
-
-
-class WebsocketProxyConnect(connect):
-    """
-    WebSocket connection class with proxy and SSL support.
-
-    Extends the websockets.asyncio.client.connect class to add HTTP/HTTPS
-    proxy support via environment variables, proxy bypass logic, and SSL
-    certificate verification.
-    """
-
-    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 not host:
-            raise ValueError(f"Invalid URI {uri}, no hostname found")
-
-        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
-            )
-
-        # Store proxy URL for deferred creation. Creating the proxy object here
-        # can bind asyncio futures to the wrong event loop when multiple WebSocket
-        # connections are initialized at different times (e.g., events + logs clients).
-        self._proxy_url = proxy_url if proxy_url and not proxy_bypass(host) else None
-        self._host = host
-        self._port = port
-
-        # Configure SSL context for HTTPS connections
-        ssl_context = create_ssl_context_for_websocket(uri)
-        if ssl_context:
-            self._kwargs.setdefault("ssl", ssl_context)
-
-        # Add custom headers from settings
-        custom_headers = get_current_settings().client.custom_headers
-        if custom_headers:
-            # Get existing additional_headers or create new dict
-            additional_headers = self._kwargs.get("additional_headers", {})
-            if not isinstance(additional_headers, dict):
-                additional_headers = {}
-
-            for header_name, header_value in custom_headers.items():
-                # Check for protected headers that shouldn't be overridden
-                if header_name.lower() in {
-                    "user-agent",
-                    "sec-websocket-key",
-                    "sec-websocket-version",
-                    "sec-websocket-extensions",
-                    "sec-websocket-protocol",
-                    "connection",
-                    "upgrade",
-                    "host",
-                }:
-                    warnings.warn(
-                        f"Custom header '{header_name}' is ignored because it conflicts with "
-                        f"a protected WebSocket header. Protected headers include: "
-                        f"User-Agent, Sec-WebSocket-Key, Sec-WebSocket-Version, "
-                        f"Sec-WebSocket-Extensions, Sec-WebSocket-Protocol, Connection, "
-                        f"Upgrade, Host",
-                        UserWarning,
-                        stacklevel=2,
-                    )
-                else:
-                    additional_headers[header_name] = header_value
-
-            self._kwargs["additional_headers"] = additional_headers
-
-    async def _proxy_connect(self: Self) -> ClientConnection:
-        if self._proxy_url:
-            # Create proxy in the current event loop context
-            proxy = Proxy.from_url(self._proxy_url)
-            sock = await 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, ClientConnection]:
-        return self._proxy_connect().__await__()
-
-
-def websocket_connect(uri: str, **kwargs: Any) -> WebsocketProxyConnect:
-    """Create a WebSocket connection with proxy and SSL support."""
-    return WebsocketProxyConnect(uri, **kwargs)