prefect-client 2.20.4__py3-none-any.whl → 3.0.0__py3-none-any.whl

prefect/events/actions.py

@@ -2,14 +2,8 @@ import abc
 from typing import Any, Dict, Optional, Union
 from uuid import UUID
 
-from typing_extensions import Literal, TypeAlias
-
-from prefect._internal.pydantic import HAS_PYDANTIC_V2
-
-if HAS_PYDANTIC_V2:
-    from pydantic.v1 import Field, root_validator
-else:
-    from pydantic import Field, root_validator  # type: ignore
+from pydantic import Field, model_validator
+from typing_extensions import Literal, Self, TypeAlias
 
 from prefect._internal.schemas.bases import PrefectBaseModel
 from prefect.client.schemas.objects import StateType
@@ -49,16 +43,16 @@ class DeploymentAction(Action):
         None, description="The identifier of the deployment"
     )
 
-    @root_validator
-    def selected_deployment_requires_id(cls, values):
-        wants_selected_deployment = values.get("source") == "selected"
-        has_deployment_id = bool(values.get("deployment_id"))
+    @model_validator(mode="after")
+    def selected_deployment_requires_id(self):
+        wants_selected_deployment = self.source == "selected"
+        has_deployment_id = bool(self.deployment_id)
         if wants_selected_deployment != has_deployment_id:
             raise ValueError(
                 "deployment_id is "
                 + ("not allowed" if has_deployment_id else "required")
             )
-        return values
+        return self
 
 
 class RunDeployment(DeploymentAction):
@@ -199,16 +193,16 @@ class WorkQueueAction(Action):
         None, description="The identifier of the work queue to pause"
     )
 
-    @root_validator
-    def selected_work_queue_requires_id(cls, values):
-        wants_selected_work_queue = values.get("source") == "selected"
-        has_work_queue_id = bool(values.get("work_queue_id"))
+    @model_validator(mode="after")
+    def selected_work_queue_requires_id(self) -> Self:
+        wants_selected_work_queue = self.source == "selected"
+        has_work_queue_id = bool(self.work_queue_id)
         if wants_selected_work_queue != has_work_queue_id:
             raise ValueError(
                 "work_queue_id is "
                 + ("not allowed" if has_work_queue_id else "required")
             )
-        return values
+        return self
 
 
 class PauseWorkQueue(WorkQueueAction):
@@ -241,16 +235,16 @@ class AutomationAction(Action):
         None, description="The identifier of the automation to act on"
     )
 
-    @root_validator
-    def selected_automation_requires_id(cls, values):
-        wants_selected_automation = values.get("source") == "selected"
-        has_automation_id = bool(values.get("automation_id"))
+    @model_validator(mode="after")
+    def selected_automation_requires_id(self) -> Self:
+        wants_selected_automation = self.source == "selected"
+        has_automation_id = bool(self.automation_id)
         if wants_selected_automation != has_automation_id:
             raise ValueError(
                 "automation_id is "
                 + ("not allowed" if has_automation_id else "required")
             )
-        return values
+        return self
 
 
 class PauseAutomation(AutomationAction):

prefect/events/cli/automations.py

@@ -3,26 +3,27 @@ Command line interface for working with automations.
 """
 
 import functools
-from typing import Optional
+from typing import Optional, Type
 from uuid import UUID
 
 import orjson
 import typer
 import yaml as pyyaml
+from pydantic import BaseModel
 from rich.pretty import Pretty
 from rich.table import Table
 from rich.text import Text
 
 from prefect.cli._types import PrefectTyper
 from prefect.cli._utilities import exit_with_error, exit_with_success
-from prefect.cli.root import app
+from prefect.cli.root import app, is_interactive
 from prefect.client.orchestration import get_client
 from prefect.events.schemas.automations import Automation
 from prefect.exceptions import PrefectHTTPStatusError
 
 automations_app = PrefectTyper(
     name="automation",
-    help="Commands for managing automations.",
+    help="Manage automations.",
 )
 app.add_typer(automations_app, aliases=["automations"])
 
@@ -148,10 +149,22 @@ async def inspect(
                 exit_with_error(f"Automation with id {id!r} not found.")
 
     if yaml or json:
+
+        def no_really_json(obj: Type[BaseModel]):
+            # Working around a weird bug where pydantic isn't rendering enums as strings
+            #
+            # automation.trigger.model_dump(mode="json")
+            # {..., 'posture': 'Reactive', ...}
+            #
+            # automation.model_dump(mode="json")
+            # {..., 'posture': Posture.Reactive, ...}
+            return orjson.loads(obj.model_dump_json())
+
         if isinstance(automation, list):
-            automation = [a.dict(json_compatible=True) for a in automation]
+            automation = [no_really_json(a) for a in automation]
         elif isinstance(automation, Automation):
-            automation = automation.dict(json_compatible=True)
+            automation = no_really_json(automation)
+
         if yaml:
             app.console.print(pyyaml.dump(automation, sort_keys=False))
         elif json:
@@ -297,7 +310,7 @@ async def delete(
             automation = await client.read_automation(id)
             if not automation:
                 exit_with_error(f"Automation with id {id!r} not found.")
-            if not typer.confirm(
+            if is_interactive() and not typer.confirm(
                 (f"Are you sure you want to delete automation with id {id!r}?"),
                 default=False,
             ):
@@ -315,7 +328,7 @@ async def delete(
                 exit_with_error(
                     f"Multiple automations found with name {name!r}. Please specify an id with the `--id` flag instead."
                 )
-            if not typer.confirm(
+            if is_interactive() and not typer.confirm(
                 (f"Are you sure you want to delete automation with name {name!r}?"),
                 default=False,
             ):

prefect/events/clients.py

@@ -15,10 +15,10 @@ from typing import (
 )
 from uuid import UUID
 
-import httpx
 import orjson
 import pendulum
 from cachetools import TTLCache
+from prometheus_client import Counter
 from typing_extensions import Self
 from websockets import Subprotocol
 from websockets.client import WebSocketClientProtocol, connect
@@ -28,25 +28,48 @@ from websockets.exceptions import (
     ConnectionClosedOK,
 )
 
-from prefect.client.base import PrefectHttpxAsyncClient
 from prefect.events import Event
 from prefect.logging import get_logger
 from prefect.settings import (
     PREFECT_API_KEY,
     PREFECT_API_URL,
     PREFECT_CLOUD_API_URL,
-    PREFECT_EXPERIMENTAL_EVENTS,
+    PREFECT_SERVER_ALLOW_EPHEMERAL_MODE,
 )
 
 if TYPE_CHECKING:
     from prefect.events.filters import EventFilter
 
+EVENTS_EMITTED = Counter(
+    "prefect_events_emitted",
+    "The number of events emitted by Prefect event clients",
+    labelnames=["client"],
+)
+EVENTS_OBSERVED = Counter(
+    "prefect_events_observed",
+    "The number of events observed by Prefect event subscribers",
+    labelnames=["client"],
+)
+EVENT_WEBSOCKET_CONNECTIONS = Counter(
+    "prefect_event_websocket_connections",
+    (
+        "The number of times Prefect event clients have connected to an event stream, "
+        "broken down by direction (in/out) and connection (initial/reconnect)"
+    ),
+    labelnames=["client", "direction", "connection"],
+)
+EVENT_WEBSOCKET_CHECKPOINTS = Counter(
+    "prefect_event_websocket_checkpoints",
+    "The number of checkpoints performed by Prefect event clients",
+    labelnames=["client"],
+)
+
 logger = get_logger(__name__)
 
 
 def get_events_client(
     reconnection_attempts: int = 10,
-    checkpoint_every: int = 20,
+    checkpoint_every: int = 700,
 ) -> "EventsClient":
     api_url = PREFECT_API_URL.value()
     if isinstance(api_url, str) and api_url.startswith(PREFECT_CLOUD_API_URL.value()):
@@ -54,20 +77,25 @@ def get_events_client(
             reconnection_attempts=reconnection_attempts,
             checkpoint_every=checkpoint_every,
         )
-    elif PREFECT_EXPERIMENTAL_EVENTS:
-        if PREFECT_API_URL:
-            return PrefectEventsClient(
-                reconnection_attempts=reconnection_attempts,
-                checkpoint_every=checkpoint_every,
-            )
-        else:
-            return PrefectEphemeralEventsClient()
+    elif api_url:
+        return PrefectEventsClient(
+            reconnection_attempts=reconnection_attempts,
+            checkpoint_every=checkpoint_every,
+        )
+    elif PREFECT_SERVER_ALLOW_EPHEMERAL_MODE:
+        from prefect.server.api.server import SubprocessASGIServer
 
-    raise RuntimeError(
-        "The current server and client configuration does not support "
-        "events.  Enable experimental events support with the "
-        "PREFECT_EXPERIMENTAL_EVENTS setting."
-    )
+        server = SubprocessASGIServer()
+        server.start()
+        return PrefectEventsClient(
+            api_url=server.api_url,
+            reconnection_attempts=reconnection_attempts,
+            checkpoint_every=checkpoint_every,
+        )
+    else:
+        raise ValueError(
+            "No Prefect API URL provided. Please set PREFECT_API_URL to the address of a running Prefect server."
+        )
 
 
 def get_events_subscriber(
@@ -75,25 +103,38 @@ def get_events_subscriber(
     reconnection_attempts: int = 10,
 ) -> "PrefectEventSubscriber":
     api_url = PREFECT_API_URL.value()
+
     if isinstance(api_url, str) and api_url.startswith(PREFECT_CLOUD_API_URL.value()):
         return PrefectCloudEventSubscriber(
             filter=filter, reconnection_attempts=reconnection_attempts
         )
-    elif PREFECT_EXPERIMENTAL_EVENTS:
+    elif api_url:
         return PrefectEventSubscriber(
             filter=filter, reconnection_attempts=reconnection_attempts
         )
+    elif PREFECT_SERVER_ALLOW_EPHEMERAL_MODE:
+        from prefect.server.api.server import SubprocessASGIServer
 
-    raise RuntimeError(
-        "The current server and client configuration does not support "
-        "events.  Enable experimental events support with the "
-        "PREFECT_EXPERIMENTAL_EVENTS setting."
-    )
+        server = SubprocessASGIServer()
+        server.start()
+        return PrefectEventSubscriber(
+            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 EventsClient(abc.ABC):
     """The abstract interface for all Prefect Events clients"""
 
+    @property
+    def client_name(self) -> str:
+        return self.__class__.__name__
+
     async def emit(self, event: Event) -> None:
         """Emit a single event"""
         if not hasattr(self, "_in_context"):
@@ -101,7 +142,11 @@ class EventsClient(abc.ABC):
                 "Events may only be emitted while this client is being used as a "
                 "context manager"
             )
-        return await self._emit(event)
+
+        try:
+            return await self._emit(event)
+        finally:
+            EVENTS_EMITTED.labels(self.client_name).inc()
 
     @abc.abstractmethod
     async def _emit(self, event: Event) -> None:  # pragma: no cover
@@ -132,7 +177,7 @@ class AssertingEventsClient(EventsClient):
     """A Prefect Events client that records all events sent to it for inspection during
     tests."""
 
-    last: ClassVar["AssertingEventsClient | None"] = None
+    last: ClassVar["Optional[AssertingEventsClient]"] = None
     all: ClassVar[List["AssertingEventsClient"]] = []
 
     args: Tuple
@@ -152,6 +197,11 @@ class AssertingEventsClient(EventsClient):
         cls.last = None
         cls.all = []
 
+    def pop_events(self) -> List[Event]:
+        events = self.events
+        self.events = []
+        return events
+
     async def _emit(self, event: Event) -> None:
         self.events.append(event)
 
@@ -175,52 +225,6 @@ def _get_api_url_and_key(
     return api_url, api_key
 
 
-class PrefectEphemeralEventsClient(EventsClient):
-    """A Prefect Events client that sends events to an ephemeral Prefect server"""
-
-    def __init__(self):
-        if not PREFECT_EXPERIMENTAL_EVENTS:
-            raise ValueError(
-                "PrefectEphemeralEventsClient can only be used when "
-                "PREFECT_EXPERIMENTAL_EVENTS is set to True"
-            )
-        if PREFECT_API_KEY.value():
-            raise ValueError(
-                "PrefectEphemeralEventsClient cannot be used when PREFECT_API_KEY is set."
-                " Please use PrefectEventsClient or PrefectCloudEventsClient instead."
-            )
-        from prefect.server.api.server import create_app
-
-        app = create_app()
-
-        self._http_client = PrefectHttpxAsyncClient(
-            transport=httpx.ASGITransport(app=app, raise_app_exceptions=False),
-            base_url="http://ephemeral-prefect/api",
-            enable_csrf_support=False,
-        )
-
-    async def __aenter__(self) -> Self:
-        await super().__aenter__()
-        await self._http_client.__aenter__()
-        return self
-
-    async def __aexit__(
-        self,
-        exc_type: Optional[Type[Exception]],
-        exc_val: Optional[Exception],
-        exc_tb: Optional[TracebackType],
-    ) -> None:
-        self._websocket = None
-        await self._http_client.__aexit__(exc_type, exc_val, exc_tb)
-        return await super().__aexit__(exc_type, exc_val, exc_tb)
-
-    async def _emit(self, event: Event) -> None:
-        await self._http_client.post(
-            "/events",
-            json=[event.dict(json_compatible=True)],
-        )
-
-
 class PrefectEventsClient(EventsClient):
     """A Prefect Events client that streams events to a Prefect server"""
 
@@ -231,7 +235,7 @@ class PrefectEventsClient(EventsClient):
         self,
         api_url: Optional[str] = None,
         reconnection_attempts: int = 10,
-        checkpoint_every: int = 20,
+        checkpoint_every: int = 700,
     ):
         """
         Args:
@@ -311,6 +315,8 @@ class PrefectEventsClient(EventsClient):
         # don't clear the list, just the ones that we are sure of.
         self._unconfirmed_events = self._unconfirmed_events[unconfirmed_count:]
 
+        EVENT_WEBSOCKET_CHECKPOINTS.labels(self.client_name).inc()
+
     async def _emit(self, event: Event) -> None:
         for i in range(self._reconnection_attempts + 1):
             try:
@@ -324,7 +330,7 @@ class PrefectEventsClient(EventsClient):
                     await self._reconnect()
                     assert self._websocket
 
-                await self._websocket.send(event.json())
+                await self._websocket.send(event.model_dump_json())
                 await self._checkpoint(event)
 
                 return
@@ -340,6 +346,47 @@ class PrefectEventsClient(EventsClient):
                     await asyncio.sleep(1)
 
 
+class AssertingPassthroughEventsClient(PrefectEventsClient):
+    """A Prefect Events client that BOTH records all events sent to it for inspection
+    during tests AND sends them to a Prefect server."""
+
+    last: ClassVar["Optional[AssertingPassthroughEventsClient]"] = None
+    all: ClassVar[List["AssertingPassthroughEventsClient"]] = []
+
+    args: Tuple
+    kwargs: Dict[str, Any]
+    events: List[Event]
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        AssertingPassthroughEventsClient.last = self
+        AssertingPassthroughEventsClient.all.append(self)
+        self.args = args
+        self.kwargs = kwargs
+
+    @classmethod
+    def reset(cls) -> None:
+        cls.last = None
+        cls.all = []
+
+    def pop_events(self) -> List[Event]:
+        events = self.events
+        self.events = []
+        return events
+
+    async def _emit(self, event: Event) -> None:
+        # actually send the event to the server
+        await super()._emit(event)
+
+        # record the event for inspection
+        self.events.append(event)
+
+    async def __aenter__(self) -> Self:
+        await super().__aenter__()
+        self.events = []
+        return self
+
+
 class PrefectCloudEventsClient(PrefectEventsClient):
     """A Prefect Events client that streams events to a Prefect Cloud Workspace"""
 
@@ -348,7 +395,7 @@ class PrefectCloudEventsClient(PrefectEventsClient):
         api_url: Optional[str] = None,
         api_key: Optional[str] = None,
         reconnection_attempts: int = 10,
-        checkpoint_every: int = 20,
+        checkpoint_every: int = 700,
     ):
         """
         Args:
@@ -412,9 +459,9 @@ class PrefectEventSubscriber:
             reconnection_attempts: When the client is disconnected, how many times
                 the client should attempt to reconnect
         """
+        self._api_key = None
         if not api_url:
             api_url = cast(str, PREFECT_API_URL.value())
-            self._api_key = None
 
         from prefect.events.filters import EventFilter
 
@@ -438,10 +485,17 @@ class PrefectEventSubscriber:
         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
-        await self._reconnect()
+        try:
+            await self._reconnect()
+        finally:
+            EVENT_WEBSOCKET_CONNECTIONS.labels(self.client_name, "out", "initial")
         return self
 
     async def _reconnect(self) -> None:
@@ -489,7 +543,7 @@ class PrefectEventSubscriber:
         logger.debug("  filtering events since %s...", self._filter.occurred.since)
         filter_message = {
             "type": "filter",
-            "filter": self._filter.dict(json_compatible=True),
+            "filter": self._filter.model_dump(mode="json"),
         }
         await self._websocket.send(orjson.dumps(filter_message).decode())
 
@@ -515,18 +569,26 @@ class PrefectEventSubscriber:
                 # Otherwise, after the first time through this loop, we're recovering
                 # from a ConnectionClosed, so reconnect now.
                 if not self._websocket or i > 0:
-                    await self._reconnect()
+                    try:
+                        await self._reconnect()
+                    finally:
+                        EVENT_WEBSOCKET_CONNECTIONS.labels(
+                            self.client_name, "out", "reconnect"
+                        )
                     assert self._websocket
 
                 while True:
                     message = orjson.loads(await self._websocket.recv())
-                    event: Event = Event.parse_obj(message["event"])
+                    event: Event = Event.model_validate(message["event"])
 
                     if event.id in self._seen_events:
                         continue
                     self._seen_events[event.id] = True
 
-                    return event
+                    try:
+                        return event
+                    finally:
+                        EVENTS_OBSERVED.labels(self.client_name).inc()
             except ConnectionClosedOK:
                 logger.debug('Connection closed with "OK" status')
                 raise StopAsyncIteration

prefect/events/filters.py

@@ -2,24 +2,19 @@ from typing import List, Optional, Tuple, cast
 from uuid import UUID
 
 import pendulum
+from pydantic import Field, PrivateAttr
+from pydantic_extra_types.pendulum_dt import DateTime
 
-from prefect._internal.pydantic import HAS_PYDANTIC_V2
 from prefect._internal.schemas.bases import PrefectBaseModel
-from prefect._internal.schemas.fields import DateTimeTZ
 from prefect.utilities.collections import AutoEnum
 
 from .schemas.events import Event, Resource, ResourceSpecification
 
-if HAS_PYDANTIC_V2:
-    from pydantic.v1 import Field, PrivateAttr
-else:
-    from pydantic import Field, PrivateAttr  # type: ignore
-
 
 class AutomationFilterCreated(PrefectBaseModel):
     """Filter by `Automation.created`."""
 
-    before_: Optional[DateTimeTZ] = Field(
+    before_: Optional[DateTime] = Field(
         default=None,
         description="Only include automations created before this datetime",
     )
@@ -46,18 +41,19 @@ class AutomationFilter(PrefectBaseModel):
 class EventDataFilter(PrefectBaseModel, extra="forbid"):  # type: ignore[call-arg]
     """A base class for filtering event data."""
 
-    _top_level_filter: "EventFilter | None" = PrivateAttr(None)
+    _top_level_filter: Optional["EventFilter"] = PrivateAttr(None)
 
     def get_filters(self) -> List["EventDataFilter"]:
-        return [
+        filters: List["EventDataFilter"] = [
             filter
             for filter in [
-                getattr(self, name)
-                for name, field in self.__fields__.items()
-                if issubclass(field.type_, EventDataFilter)
+                getattr(self, name) for name, field in self.model_fields.items()
             ]
-            if filter
+            if isinstance(filter, EventDataFilter)
         ]
+        for filter in filters:
+            filter._top_level_filter = self._top_level_filter
+        return filters
 
     def includes(self, event: Event) -> bool:
         """Does the given event match the criteria of this filter?"""
@@ -69,15 +65,15 @@ class EventDataFilter(PrefectBaseModel, extra="forbid"):  # type: ignore[call-ar
 
 
 class EventOccurredFilter(EventDataFilter):
-    since: DateTimeTZ = Field(
+    since: DateTime = Field(
         default_factory=lambda: cast(
-            DateTimeTZ,
+            DateTime,
             pendulum.now("UTC").start_of("day").subtract(days=180),
         ),
         description="Only include events after this time (inclusive)",
     )
-    until: DateTimeTZ = Field(
-        default_factory=lambda: cast(DateTimeTZ, pendulum.now("UTC")),
+    until: DateTime = Field(
+        default_factory=lambda: cast(DateTime, pendulum.now("UTC")),
         description="Only include events prior to this time (inclusive)",
     )