prefect-client 2.14.20__py3-none-any.whl → 2.15.0__py3-none-any.whl

prefect/_vendor/starlette/types.py

@@ -0,0 +1,30 @@
+import typing
+
+if typing.TYPE_CHECKING:
+    from prefect._vendor.starlette.requests import Request
+    from prefect._vendor.starlette.responses import Response
+    from prefect._vendor.starlette.websockets import WebSocket
+
+AppType = typing.TypeVar("AppType")
+
+Scope = typing.MutableMapping[str, typing.Any]
+Message = typing.MutableMapping[str, typing.Any]
+
+Receive = typing.Callable[[], typing.Awaitable[Message]]
+Send = typing.Callable[[Message], typing.Awaitable[None]]
+
+ASGIApp = typing.Callable[[Scope, Receive, Send], typing.Awaitable[None]]
+
+StatelessLifespan = typing.Callable[[AppType], typing.AsyncContextManager[None]]
+StatefulLifespan = typing.Callable[
+    [AppType], typing.AsyncContextManager[typing.Mapping[str, typing.Any]]
+]
+Lifespan = typing.Union[StatelessLifespan[AppType], StatefulLifespan[AppType]]
+
+HTTPExceptionHandler = typing.Callable[
+    ["Request", Exception], typing.Union["Response", typing.Awaitable["Response"]]
+]
+WebSocketExceptionHandler = typing.Callable[
+    ["WebSocket", Exception], typing.Awaitable[None]
+]
+ExceptionHandler = typing.Union[HTTPExceptionHandler, WebSocketExceptionHandler]

prefect/_vendor/starlette/websockets.py

@@ -0,0 +1,193 @@
+import enum
+import json
+import typing
+
+from prefect._vendor.starlette.requests import HTTPConnection
+from prefect._vendor.starlette.types import Message, Receive, Scope, Send
+
+
+class WebSocketState(enum.Enum):
+    CONNECTING = 0
+    CONNECTED = 1
+    DISCONNECTED = 2
+
+
+class WebSocketDisconnect(Exception):
+    def __init__(self, code: int = 1000, reason: typing.Optional[str] = None) -> None:
+        self.code = code
+        self.reason = reason or ""
+
+
+class WebSocket(HTTPConnection):
+    def __init__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        super().__init__(scope)
+        assert scope["type"] == "websocket"
+        self._receive = receive
+        self._send = send
+        self.client_state = WebSocketState.CONNECTING
+        self.application_state = WebSocketState.CONNECTING
+
+    async def receive(self) -> Message:
+        """
+        Receive ASGI websocket messages, ensuring valid state transitions.
+        """
+        if self.client_state == WebSocketState.CONNECTING:
+            message = await self._receive()
+            message_type = message["type"]
+            if message_type != "websocket.connect":
+                raise RuntimeError(
+                    'Expected ASGI message "websocket.connect", '
+                    f"but got {message_type!r}"
+                )
+            self.client_state = WebSocketState.CONNECTED
+            return message
+        elif self.client_state == WebSocketState.CONNECTED:
+            message = await self._receive()
+            message_type = message["type"]
+            if message_type not in {"websocket.receive", "websocket.disconnect"}:
+                raise RuntimeError(
+                    'Expected ASGI message "websocket.receive" or '
+                    f'"websocket.disconnect", but got {message_type!r}'
+                )
+            if message_type == "websocket.disconnect":
+                self.client_state = WebSocketState.DISCONNECTED
+            return message
+        else:
+            raise RuntimeError(
+                'Cannot call "receive" once a disconnect message has been received.'
+            )
+
+    async def send(self, message: Message) -> None:
+        """
+        Send ASGI websocket messages, ensuring valid state transitions.
+        """
+        if self.application_state == WebSocketState.CONNECTING:
+            message_type = message["type"]
+            if message_type not in {"websocket.accept", "websocket.close"}:
+                raise RuntimeError(
+                    'Expected ASGI message "websocket.accept" or '
+                    f'"websocket.close", but got {message_type!r}'
+                )
+            if message_type == "websocket.close":
+                self.application_state = WebSocketState.DISCONNECTED
+            else:
+                self.application_state = WebSocketState.CONNECTED
+            await self._send(message)
+        elif self.application_state == WebSocketState.CONNECTED:
+            message_type = message["type"]
+            if message_type not in {"websocket.send", "websocket.close"}:
+                raise RuntimeError(
+                    'Expected ASGI message "websocket.send" or "websocket.close", '
+                    f"but got {message_type!r}"
+                )
+            if message_type == "websocket.close":
+                self.application_state = WebSocketState.DISCONNECTED
+            await self._send(message)
+        else:
+            raise RuntimeError('Cannot call "send" once a close message has been sent.')
+
+    async def accept(
+        self,
+        subprotocol: typing.Optional[str] = None,
+        headers: typing.Optional[typing.Iterable[typing.Tuple[bytes, bytes]]] = None,
+    ) -> None:
+        headers = headers or []
+
+        if self.client_state == WebSocketState.CONNECTING:
+            # If we haven't yet seen the 'connect' message, then wait for it first.
+            await self.receive()
+        await self.send(
+            {"type": "websocket.accept", "subprotocol": subprotocol, "headers": headers}
+        )
+
+    def _raise_on_disconnect(self, message: Message) -> None:
+        if message["type"] == "websocket.disconnect":
+            raise WebSocketDisconnect(message["code"], message.get("reason"))
+
+    async def receive_text(self) -> str:
+        if self.application_state != WebSocketState.CONNECTED:
+            raise RuntimeError(
+                'WebSocket is not connected. Need to call "accept" first.'
+            )
+        message = await self.receive()
+        self._raise_on_disconnect(message)
+        return typing.cast(str, message["text"])
+
+    async def receive_bytes(self) -> bytes:
+        if self.application_state != WebSocketState.CONNECTED:
+            raise RuntimeError(
+                'WebSocket is not connected. Need to call "accept" first.'
+            )
+        message = await self.receive()
+        self._raise_on_disconnect(message)
+        return typing.cast(bytes, message["bytes"])
+
+    async def receive_json(self, mode: str = "text") -> typing.Any:
+        if mode not in {"text", "binary"}:
+            raise RuntimeError('The "mode" argument should be "text" or "binary".')
+        if self.application_state != WebSocketState.CONNECTED:
+            raise RuntimeError(
+                'WebSocket is not connected. Need to call "accept" first.'
+            )
+        message = await self.receive()
+        self._raise_on_disconnect(message)
+
+        if mode == "text":
+            text = message["text"]
+        else:
+            text = message["bytes"].decode("utf-8")
+        return json.loads(text)
+
+    async def iter_text(self) -> typing.AsyncIterator[str]:
+        try:
+            while True:
+                yield await self.receive_text()
+        except WebSocketDisconnect:
+            pass
+
+    async def iter_bytes(self) -> typing.AsyncIterator[bytes]:
+        try:
+            while True:
+                yield await self.receive_bytes()
+        except WebSocketDisconnect:
+            pass
+
+    async def iter_json(self) -> typing.AsyncIterator[typing.Any]:
+        try:
+            while True:
+                yield await self.receive_json()
+        except WebSocketDisconnect:
+            pass
+
+    async def send_text(self, data: str) -> None:
+        await self.send({"type": "websocket.send", "text": data})
+
+    async def send_bytes(self, data: bytes) -> None:
+        await self.send({"type": "websocket.send", "bytes": data})
+
+    async def send_json(self, data: typing.Any, mode: str = "text") -> None:
+        if mode not in {"text", "binary"}:
+            raise RuntimeError('The "mode" argument should be "text" or "binary".')
+        text = json.dumps(data, separators=(",", ":"), ensure_ascii=False)
+        if mode == "text":
+            await self.send({"type": "websocket.send", "text": text})
+        else:
+            await self.send({"type": "websocket.send", "bytes": text.encode("utf-8")})
+
+    async def close(
+        self, code: int = 1000, reason: typing.Optional[str] = None
+    ) -> None:
+        await self.send(
+            {"type": "websocket.close", "code": code, "reason": reason or ""}
+        )
+
+
+class WebSocketClose:
+    def __init__(self, code: int = 1000, reason: typing.Optional[str] = None) -> None:
+        self.code = code
+        self.reason = reason or ""
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        await send(
+            {"type": "websocket.close", "code": self.code, "reason": self.reason}
+        )

prefect/blocks/core.py

@@ -257,9 +257,9 @@ class Block(BaseModel, ABC):
                                     type_._to_block_schema_reference_dict(),
                                 ]
                             else:
-                                refs[field.name] = (
-                                    type_._to_block_schema_reference_dict()
-                                )
+                                refs[
+                                    field.name
+                                ] = type_._to_block_schema_reference_dict()
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)

prefect/blocks/notifications.py

@@ -24,14 +24,14 @@ class AbstractAppriseNotificationBlock(NotificationBlock, ABC):
     An abstract class for sending notifications using Apprise.
     """
 
-    notify_type: Literal["prefect_default", "info", "success", "warning", "failure"] = (
-        Field(
-            default=PREFECT_NOTIFY_TYPE_DEFAULT,
-            description=(
-                "The type of notification being performed; the prefect_default "
-                "is a plain notification that does not attach an image."
-            ),
-        )
+    notify_type: Literal[
+        "prefect_default", "info", "success", "warning", "failure"
+    ] = Field(
+        default=PREFECT_NOTIFY_TYPE_DEFAULT,
+        description=(
+            "The type of notification being performed; the prefect_default "
+            "is a plain notification that does not attach an image."
+        ),
     )
 
     def __init__(self, *args, **kwargs):
@@ -714,7 +714,8 @@ class CustomWebhookNotificationBlock(NotificationBlock):
 
         # make request with httpx
         client = httpx.AsyncClient(headers={"user-agent": "Prefect Notifications"})
-        resp = await client.request(**self._build_request_args(body, subject))
+        async with client:
+            resp = await client.request(**self._build_request_args(body, subject))
         resp.raise_for_status()
 
 

prefect/client/base.py

@@ -22,7 +22,7 @@ import anyio
 import httpx
 from asgi_lifespan import LifespanManager
 from httpx import HTTPStatusError, Response
-from starlette import status
+from prefect._vendor.starlette import status
 from typing_extensions import Self
 
 from prefect.exceptions import PrefectHTTPStatusError

prefect/client/cloud.py

@@ -11,7 +11,7 @@ if HAS_PYDANTIC_V2:
 else:
     import pydantic
 
-from starlette import status
+from prefect._vendor.starlette import status
 
 import prefect.context
 import prefect.settings

prefect/client/orchestration.py

@@ -26,7 +26,7 @@ else:
     import pydantic
 
 from asgi_lifespan import LifespanManager
-from starlette import status
+from prefect._vendor.starlette import status
 
 import prefect
 import prefect.exceptions

prefect/client/schemas/objects.py

@@ -97,6 +97,14 @@ class DeploymentStatus(AutoEnum):
     NOT_READY = AutoEnum.auto()
 
 
+class WorkQueueStatus(AutoEnum):
+    """Enumeration of work queue statuses."""
+
+    READY = AutoEnum.auto()
+    NOT_READY = AutoEnum.auto()
+    PAUSED = AutoEnum.auto()
+
+
 class StateDetails(PrefectBaseModel):
     flow_run_id: UUID = None
     task_run_id: UUID = None
@@ -1222,6 +1230,9 @@ class WorkQueue(ObjectBaseModel):
     last_polled: Optional[DateTimeTZ] = Field(
         default=None, description="The last time an agent polled this queue for work."
     )
+    status: Optional[WorkQueueStatus] = Field(
+        default=None, description="The queue status."
+    )
 
     @validator("name", check_fields=False)
     def validate_name_characters(cls, v):

prefect/client/subscriptions.py

@@ -1,25 +1,30 @@
 import asyncio
-from typing import Generic, Type, TypeVar
+from typing import Generic, List, Type, TypeVar
 
 import orjson
 import websockets
 import websockets.exceptions
-from starlette.status import WS_1008_POLICY_VIOLATION
+from prefect._vendor.starlette.status import WS_1008_POLICY_VIOLATION
 from typing_extensions import Self
 
 from prefect._internal.schemas.bases import IDBaseModel
+from prefect.logging import get_logger
 from prefect.settings import PREFECT_API_KEY, PREFECT_API_URL
 
+logger = get_logger(__name__)
+
 S = TypeVar("S", bound=IDBaseModel)
 
 
 class Subscription(Generic[S]):
-    def __init__(self, model: Type[S], path: str):
+    def __init__(self, model: Type[S], path: str, keys: List[str]):
         self.model = model
 
         base_url = PREFECT_API_URL.value().replace("http", "ws", 1)
         self.subscription_url = f"{base_url}{path}"
 
+        self.keys = keys
+
         self._connect = websockets.connect(
             self.subscription_url,
             subprotocols=["prefect"],
@@ -35,11 +40,7 @@ class Subscription(Generic[S]):
                 await self._ensure_connected()
                 message = await self._websocket.recv()
 
-                message_data = orjson.loads(message)
-
-                if message_data.get("type") == "ping":
-                    await self._websocket.send(orjson.dumps({"type": "pong"}).decode())
-                    continue
+                await self._websocket.send(orjson.dumps({"type": "ack"}).decode())
 
                 return self.model.parse_raw(message)
             except (
@@ -57,13 +58,19 @@ class Subscription(Generic[S]):
 
         websocket = await self._connect.__aenter__()
 
-        await websocket.send(
-            orjson.dumps({"type": "auth", "token": PREFECT_API_KEY.value()}).decode()
-        )
-
         try:
+            await websocket.send(
+                orjson.dumps(
+                    {"type": "auth", "token": PREFECT_API_KEY.value()}
+                ).decode()
+            )
+
             auth = orjson.loads(await websocket.recv())
             assert auth["type"] == "auth_success"
+
+            await websocket.send(
+                orjson.dumps({"type": "subscribe", "keys": self.keys}).decode()
+            )
         except (
             AssertionError,
             websockets.exceptions.ConnectionClosedError,

prefect/concurrency/services.py

@@ -7,7 +7,7 @@ from typing import (
 )
 
 import httpx
-from starlette import status
+from prefect._vendor.starlette import status
 
 from prefect import get_client
 from prefect._internal.concurrency import logger

prefect/context.py

@@ -137,9 +137,9 @@ class PrefectObjectRegistry(ContextModel):
     )
 
     # Failures will be a tuple of (exception, instance, args, kwargs)
-    _instance_init_failures: Dict[Type[T], List[Tuple[Exception, T, Tuple, Dict]]] = (
-        PrivateAttr(default_factory=lambda: defaultdict(list))
-    )
+    _instance_init_failures: Dict[
+        Type[T], List[Tuple[Exception, T, Tuple, Dict]]
+    ] = PrivateAttr(default_factory=lambda: defaultdict(list))
 
     block_code_execution: bool = False
     capture_failures: bool = False
@@ -236,7 +236,7 @@ class EngineContext(RunContext):
     autonomous_task_run: Optional[TaskRun] = None
     task_runner: BaseTaskRunner
     log_prints: bool = False
-    parameters: Dict[str, Any]
+    parameters: Optional[Dict[str, Any]] = None
 
     # Result handling
     result_factory: ResultFactory

prefect/deployments/deployments.py

@@ -93,13 +93,13 @@ async def run_deployment(
             run metadata immediately. Setting `timeout` to None will allow this
             function to poll indefinitely. Defaults to None.
         poll_interval: The number of seconds between polls
-        tags: A list of tags to associate with this flow run; note that tags are used
-            only for organizational purposes.
+        tags: A list of tags to associate with this flow run; tags can be used in
+            automations and for organizational purposes.
         idempotency_key: A unique value to recognize retries of the same run, and
             prevent creating multiple flow runs.
         work_queue_name: The name of a work queue to use for this run. Defaults to
             the default work queue for the deployment.
-        as_subflow: Whether or not to link the flow run as a subflow of the current
+        as_subflow: Whether to link the flow run as a subflow of the current
             flow or task run.
     """
     if timeout is not None and timeout < 0: