prefect-client 3.4.7.dev3__tar.gz → 3.4.7.dev4__tar.gz

{prefect_client-3.4.7.dev3 → prefect_client-3.4.7.dev4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: prefect-client
-Version: 3.4.7.dev3
+Version: 3.4.7.dev4
 Summary: Workflow orchestration and management.
 Project-URL: Changelog, https://github.com/PrefectHQ/prefect/releases
 Project-URL: Documentation, https://docs.prefect.io

prefect_client-3.4.7.dev4/src/prefect/AGENTS.md

@@ -0,0 +1,28 @@
+# Core Prefect SDK
+
+The foundation for building and executing workflows with Python.
+
+## Key Components
+
+- **Flows & Tasks**: Workflow definition with `@flow` and `@task` decorators
+- **States**: Execution status tracking (Pending, Running, Completed, Failed)
+- **Context**: Runtime information and dependency injection
+- **Results**: Task output persistence and retrieval
+- **Deployments**: Packaging flows for scheduled/triggered execution
+- **Blocks**: Reusable configuration for external systems
+
+## Main Modules
+
+- `flows.py` - Flow lifecycle and execution
+- `tasks.py` - Task definition and dependency resolution
+- `engine.py` - Core execution engine
+- `client/` - Server/Cloud API communication
+- `deployments/` - Deployment management
+- `blocks/` - Infrastructure and storage blocks
+
+## SDK-Specific Notes
+
+- Async-first execution model with sync support
+- Immutable flow/task definitions after creation
+- State transitions handle retries and caching
+- Backward compatibility required for public APIs

prefect_client-3.4.7.dev4/src/prefect/_build_info.py

@@ -0,0 +1,5 @@
+# Generated by versioningit
+__version__ = "3.4.7.dev4"
+__build_date__ = "2025-06-18 08:09:28.573646+00:00"
+__git_commit__ = "f2a136e5b5fb96adfde1d9aa44fdedb5b11c97a0"
+__dirty__ = False

prefect_client-3.4.7.dev4/src/prefect/_internal/websockets.py

@@ -0,0 +1,106 @@
+"""
+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
+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
+            )
+
+        self._proxy = (
+            Proxy.from_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)
+
+    async def _proxy_connect(self: Self) -> ClientConnection:
+        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, 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)

{prefect_client-3.4.7.dev3 → prefect_client-3.4.7.dev4}/src/prefect/events/clients.py

@@ -1,7 +1,5 @@
 import abc
 import asyncio
-import os
-import ssl
 from datetime import timedelta
 from types import TracebackType
 from typing import (
@@ -9,7 +7,6 @@ from typing import (
     Any,
     ClassVar,
     Dict,
-    Generator,
     List,
     MutableMapping,
     Optional,
@@ -17,18 +14,14 @@ from typing import (
     Type,
     cast,
 )
-from urllib.parse import urlparse
-from urllib.request import proxy_bypass
 from uuid import UUID
 
-import certifi
 import orjson
 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.asyncio.client import ClientConnection, connect
+from websockets.asyncio.client import ClientConnection
 from websockets.exceptions import (
     ConnectionClosed,
     ConnectionClosedError,
@@ -36,13 +29,12 @@ from websockets.exceptions import (
 )
 
 import prefect.types._datetime
+from prefect._internal.websockets import websocket_connect
 from prefect.events import Event
 from prefect.logging import get_logger
 from prefect.settings import (
     PREFECT_API_AUTH_STRING,
     PREFECT_API_KEY,
-    PREFECT_API_SSL_CERT_FILE,
-    PREFECT_API_TLS_INSECURE_SKIP_VERIFY,
     PREFECT_API_URL,
     PREFECT_CLOUD_API_URL,
     PREFECT_DEBUG_MODE,
@@ -94,72 +86,6 @@ def events_out_socket_from_api_url(url: str) -> 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 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
-            )
-
-        self._proxy = (
-            Proxy.from_url(proxy_url) if proxy_url and not proxy_bypass(host) else None
-        )
-        self._host = host
-        self._port = port
-
-        if PREFECT_API_TLS_INSECURE_SKIP_VERIFY and u.scheme == "wss":
-            # Create an unverified context for insecure connections
-            ctx = ssl.create_default_context()
-            ctx.check_hostname = False
-            ctx.verify_mode = ssl.CERT_NONE
-            self._kwargs.setdefault("ssl", ctx)
-        elif u.scheme == "wss":
-            cert_file = PREFECT_API_SSL_CERT_FILE.value()
-            if not cert_file:
-                cert_file = certifi.where()
-            # Create a verified context with the certificate file
-            ctx = ssl.create_default_context(cafile=cert_file)
-            self._kwargs.setdefault("ssl", ctx)
-
-    async def _proxy_connect(self: Self) -> ClientConnection:
-        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, ClientConnection]:
-        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,

prefect_client-3.4.7.dev4/src/prefect/logging/clients.py

@@ -0,0 +1,347 @@
+import asyncio
+from datetime import timedelta
+from types import TracebackType
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    MutableMapping,
+    Optional,
+    Tuple,
+    Type,
+    cast,
+)
+from uuid import UUID
+
+import orjson
+from cachetools import TTLCache
+from prometheus_client import Counter
+from typing_extensions import Self
+from websockets import Subprotocol
+from websockets.asyncio.client import ClientConnection
+from websockets.exceptions import (
+    ConnectionClosed,
+    ConnectionClosedError,
+    ConnectionClosedOK,
+)
+
+from prefect._internal.websockets import (
+    create_ssl_context_for_websocket,
+    websocket_connect,
+)
+from prefect.client.schemas.objects import Log
+from prefect.logging import get_logger
+from prefect.settings import (
+    PREFECT_API_AUTH_STRING,
+    PREFECT_API_KEY,
+    PREFECT_API_URL,
+    PREFECT_CLOUD_API_URL,
+    PREFECT_SERVER_ALLOW_EPHEMERAL_MODE,
+)
+from prefect.types._datetime import now
+
+if TYPE_CHECKING:
+    import logging
+
+    from prefect.client.schemas.filters import LogFilter
+
+logger: "logging.Logger" = get_logger(__name__)
+
+LOGS_OBSERVED = Counter(
+    "prefect_logs_observed",
+    "The number of logs observed by Prefect log subscribers",
+    labelnames=["client"],
+)
+LOG_WEBSOCKET_CONNECTIONS = Counter(
+    "prefect_log_websocket_connections",
+    (
+        "The number of times Prefect log clients have connected to a log stream, "
+        "broken down by direction (in/out) and connection (initial/reconnect)"
+    ),
+    labelnames=["client", "direction", "connection"],
+)
+
+SEEN_LOGS_SIZE = 500_000
+SEEN_LOGS_TTL = 120
+
+
+def http_to_ws(url: str) -> str:
+    return url.replace("https://", "wss://").replace("http://", "ws://").rstrip("/")
+
+
+def logs_out_socket_from_api_url(url: str) -> str:
+    return http_to_ws(url) + "/logs/out"
+
+
+def _get_api_url_and_key(
+    api_url: Optional[str], api_key: Optional[str]
+) -> Tuple[str, str]:
+    api_url = api_url or PREFECT_API_URL.value()
+    api_key = api_key or PREFECT_API_KEY.value()
+
+    if not api_url or not api_key:
+        raise ValueError(
+            "api_url and api_key must be provided or set in the Prefect configuration"
+        )
+
+    return api_url, api_key
+
+
+def get_logs_subscriber(
+    filter: Optional["LogFilter"] = None,
+    reconnection_attempts: int = 10,
+) -> "PrefectLogsSubscriber":
+    """
+    Get a logs subscriber based on the current Prefect configuration.
+
+    Similar to get_events_subscriber, this automatically detects whether
+    you're using Prefect Cloud or OSS and returns the appropriate subscriber.
+    """
+    api_url = PREFECT_API_URL.value()
+
+    if isinstance(api_url, str) and api_url.startswith(PREFECT_CLOUD_API_URL.value()):
+        return PrefectCloudLogsSubscriber(
+            filter=filter, reconnection_attempts=reconnection_attempts
+        )
+    elif api_url:
+        return PrefectLogsSubscriber(
+            api_url=api_url,
+            filter=filter,
+            reconnection_attempts=reconnection_attempts,
+        )
+    elif PREFECT_SERVER_ALLOW_EPHEMERAL_MODE:
+        from prefect.server.api.server import SubprocessASGIServer
+
+        server = SubprocessASGIServer()
+        server.start()
+        return PrefectLogsSubscriber(
+            api_url=server.api_url,
+            filter=filter,
+            reconnection_attempts=reconnection_attempts,
+        )
+    else:
+        raise ValueError(
+            "No Prefect API URL provided. Please set PREFECT_API_URL to the address of a running Prefect server."
+        )
+
+
+class PrefectLogsSubscriber:
+    """
+    Subscribes to a Prefect logs stream, yielding logs as they occur.
+
+    Example:
+
+        from prefect.logging.clients import PrefectLogsSubscriber
+        from prefect.client.schemas.filters import LogFilter, LogFilterLevel
+        import logging
+
+        filter = LogFilter(level=LogFilterLevel(ge_=logging.INFO))
+
+        async with PrefectLogsSubscriber(filter=filter) as subscriber:
+            async for log in subscriber:
+                print(log.timestamp, log.level, log.message)
+
+    """
+
+    _websocket: Optional[ClientConnection]
+    _filter: "LogFilter"
+    _seen_logs: MutableMapping[UUID, bool]
+
+    _api_key: Optional[str]
+    _auth_token: Optional[str]
+
+    def __init__(
+        self,
+        api_url: Optional[str] = None,
+        filter: Optional["LogFilter"] = None,
+        reconnection_attempts: int = 10,
+    ):
+        """
+        Args:
+            api_url: The base URL for a Prefect workspace
+            filter: Log filter to apply
+            reconnection_attempts: When the client is disconnected, how many times
+                the client should attempt to reconnect
+        """
+        self._api_key = None
+        self._auth_token = PREFECT_API_AUTH_STRING.value()
+
+        if not api_url:
+            api_url = cast(str, PREFECT_API_URL.value())
+
+        from prefect.client.schemas.filters import LogFilter
+
+        self._filter = filter or LogFilter()  # type: ignore[call-arg]
+        self._seen_logs = TTLCache(maxsize=SEEN_LOGS_SIZE, ttl=SEEN_LOGS_TTL)
+
+        socket_url = logs_out_socket_from_api_url(api_url)
+
+        logger.debug("Connecting to %s", socket_url)
+
+        # Configure SSL context for the connection
+        ssl_context = create_ssl_context_for_websocket(socket_url)
+        connect_kwargs: dict[str, Any] = {"subprotocols": [Subprotocol("prefect")]}
+        if ssl_context:
+            connect_kwargs["ssl"] = ssl_context
+
+        self._connect = websocket_connect(socket_url, **connect_kwargs)
+        self._websocket = None
+        self._reconnection_attempts = reconnection_attempts
+        if self._reconnection_attempts < 0:
+            raise ValueError("reconnection_attempts must be a non-negative integer")
+
+    @property
+    def client_name(self) -> str:
+        return self.__class__.__name__
+
+    async def __aenter__(self) -> Self:
+        # Don't handle any errors in the initial connection, because these are most
+        # likely a permission or configuration issue that should propagate
+        try:
+            await self._reconnect()
+        finally:
+            LOG_WEBSOCKET_CONNECTIONS.labels(self.client_name, "out", "initial").inc()
+        return self
+
+    async def _reconnect(self) -> None:
+        logger.debug("Reconnecting...")
+        if self._websocket:
+            self._websocket = None
+            await self._connect.__aexit__(None, None, None)
+
+        self._websocket = await self._connect.__aenter__()
+
+        # make sure we have actually connected
+        logger.debug("  pinging...")
+        pong = await self._websocket.ping()
+        await pong
+
+        # Send authentication message - logs WebSocket requires auth handshake
+        auth_token = self._api_key or self._auth_token
+        auth_message = {"type": "auth", "token": auth_token}
+        logger.debug("  authenticating...")
+        await self._websocket.send(orjson.dumps(auth_message).decode())
+
+        # Wait for auth response
+        try:
+            message = orjson.loads(await self._websocket.recv())
+            logger.debug("  auth result %s", message)
+            assert message["type"] == "auth_success", message.get("reason", "")
+        except AssertionError as e:
+            raise Exception(
+                "Unable to authenticate to the log stream. Please ensure the "
+                "provided api_key or auth_token you are using is valid for this environment. "
+                f"Reason: {e.args[0]}"
+            )
+        except ConnectionClosedError as e:
+            reason = getattr(e.rcvd, "reason", None)
+            msg = "Unable to authenticate to the log stream. Please ensure the "
+            msg += "provided api_key or auth_token you are using is valid for this environment. "
+            msg += f"Reason: {reason}" if reason else ""
+            raise Exception(msg) from e
+
+        from prefect.client.schemas.filters import LogFilterTimestamp
+
+        current_time = now("UTC")
+        self._filter.timestamp = LogFilterTimestamp(
+            after_=current_time - timedelta(minutes=1),  # type: ignore[arg-type]
+            before_=current_time + timedelta(days=365),  # type: ignore[arg-type]
+        )
+
+        logger.debug("  filtering logs since %s...", self._filter.timestamp.after_)
+        filter_message = {
+            "type": "filter",
+            "filter": self._filter.model_dump(mode="json"),
+        }
+        await self._websocket.send(orjson.dumps(filter_message).decode())
+
+    async def __aexit__(
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc_val: Optional[BaseException],
+        exc_tb: Optional[TracebackType],
+    ) -> None:
+        self._websocket = None
+        await self._connect.__aexit__(exc_type, exc_val, exc_tb)
+
+    def __aiter__(self) -> Self:
+        return self
+
+    async def __anext__(self) -> Log:
+        assert self._reconnection_attempts >= 0
+        for i in range(self._reconnection_attempts + 1):  # pragma: no branch
+            try:
+                # If we're here and the websocket is None, then we've had a failure in a
+                # previous reconnection attempt.
+                #
+                # Otherwise, after the first time through this loop, we're recovering
+                # from a ConnectionClosed, so reconnect now.
+                if not self._websocket or i > 0:
+                    try:
+                        await self._reconnect()
+                    finally:
+                        LOG_WEBSOCKET_CONNECTIONS.labels(
+                            self.client_name, "out", "reconnect"
+                        ).inc()
+                    assert self._websocket
+
+                while True:
+                    message = orjson.loads(await self._websocket.recv())
+                    log: Log = Log.model_validate(message["log"])
+
+                    if log.id in self._seen_logs:
+                        continue
+                    self._seen_logs[log.id] = True
+
+                    try:
+                        return log
+                    finally:
+                        LOGS_OBSERVED.labels(self.client_name).inc()
+
+            except ConnectionClosedOK:
+                logger.debug('Connection closed with "OK" status')
+                raise StopAsyncIteration
+            except ConnectionClosed:
+                logger.debug(
+                    "Connection closed with %s/%s attempts",
+                    i + 1,
+                    self._reconnection_attempts,
+                )
+                if i == self._reconnection_attempts:
+                    # this was our final chance, raise the most recent error
+                    raise
+
+                if i > 2:
+                    # let the first two attempts happen quickly in case this is just
+                    # a standard load balancer timeout, but after that, just take a
+                    # beat to let things come back around.
+                    await asyncio.sleep(1)
+        raise StopAsyncIteration
+
+
+class PrefectCloudLogsSubscriber(PrefectLogsSubscriber):
+    """Logs subscriber for Prefect Cloud"""
+
+    def __init__(
+        self,
+        api_url: Optional[str] = None,
+        api_key: Optional[str] = None,
+        filter: Optional["LogFilter"] = None,
+        reconnection_attempts: int = 10,
+    ):
+        """
+        Args:
+            api_url: The base URL for a Prefect Cloud workspace
+            api_key: The API key of an actor with the see_flows scope
+            filter: Log filter to apply
+            reconnection_attempts: When the client is disconnected, how many times
+                the client should attempt to reconnect
+        """
+        api_url, api_key = _get_api_url_and_key(api_url, api_key)
+
+        super().__init__(
+            api_url=api_url,
+            filter=filter,
+            reconnection_attempts=reconnection_attempts,
+        )
+
+        self._api_key = api_key