prefect-client 2.17.1__py3-none-any.whl → 2.18.1__py3-none-any.whl

prefect/events/clients.py

@@ -7,17 +7,19 @@ from typing import (
     ClassVar,
     Dict,
     List,
-    Mapping,
+    MutableMapping,
     Optional,
     Tuple,
     Type,
 )
 from uuid import UUID
 
+import httpx
 import orjson
 import pendulum
 from cachetools import TTLCache
 from typing_extensions import Self
+from websockets import Subprotocol
 from websockets.client import WebSocketClientProtocol, connect
 from websockets.exceptions import (
     ConnectionClosed,
@@ -25,9 +27,15 @@ from websockets.exceptions import (
     ConnectionClosedOK,
 )
 
+from prefect.client.base import PrefectHttpxClient
 from prefect.events import Event
 from prefect.logging import get_logger
-from prefect.settings import PREFECT_API_KEY, PREFECT_API_URL
+from prefect.settings import (
+    PREFECT_API_KEY,
+    PREFECT_API_URL,
+    PREFECT_CLOUD_API_URL,
+    PREFECT_EXPERIMENTAL_EVENTS,
+)
 
 if TYPE_CHECKING:
     from prefect.events.filters import EventFilter
@@ -35,6 +43,53 @@ if TYPE_CHECKING:
 logger = get_logger(__name__)
 
 
+def get_events_client(
+    reconnection_attempts: int = 10,
+    checkpoint_every: int = 20,
+) -> "EventsClient":
+    api_url = PREFECT_API_URL.value()
+    if isinstance(api_url, str) and api_url.startswith(PREFECT_CLOUD_API_URL.value()):
+        return PrefectCloudEventsClient(
+            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()
+
+    raise RuntimeError(
+        "The current server and client configuration does not support "
+        "events.  Enable experimental events support with the "
+        "PREFECT_EXPERIMENTAL_EVENTS setting."
+    )
+
+
+def get_events_subscriber(
+    filter: Optional["EventFilter"] = None,
+    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:
+        return PrefectEventSubscriber(
+            filter=filter, reconnection_attempts=reconnection_attempts
+        )
+
+    raise RuntimeError(
+        "The current server and client configuration does not support "
+        "events.  Enable experimental events support with the "
+        "PREFECT_EXPERIMENTAL_EVENTS setting."
+    )
+
+
 class EventsClient(abc.ABC):
     """The abstract interface for all Prefect Events clients"""
 
@@ -51,7 +106,7 @@ class EventsClient(abc.ABC):
     async def _emit(self, event: Event) -> None:  # pragma: no cover
         ...
 
-    async def __aenter__(self) -> "EventsClient":
+    async def __aenter__(self) -> Self:
         self._in_context = True
         return self
 
@@ -99,7 +154,7 @@ class AssertingEventsClient(EventsClient):
     async def _emit(self, event: Event) -> None:
         self.events.append(event)
 
-    async def __aenter__(self) -> "AssertingEventsClient":
+    async def __aenter__(self) -> Self:
         await super().__aenter__()
         self.events = []
         return self
@@ -119,6 +174,52 @@ 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 = PrefectHttpxClient(
+            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"""
 
@@ -127,7 +228,7 @@ class PrefectEventsClient(EventsClient):
 
     def __init__(
         self,
-        api_url: str = None,
+        api_url: Optional[str] = None,
         reconnection_attempts: int = 10,
         checkpoint_every: int = 20,
     ):
@@ -243,8 +344,8 @@ class PrefectCloudEventsClient(PrefectEventsClient):
 
     def __init__(
         self,
-        api_url: str = None,
-        api_key: str = None,
+        api_url: Optional[str] = None,
+        api_key: Optional[str] = None,
         reconnection_attempts: int = 10,
         checkpoint_every: int = 20,
     ):
@@ -274,18 +375,18 @@ SEEN_EVENTS_SIZE = 500_000
 SEEN_EVENTS_TTL = 120
 
 
-class PrefectCloudEventSubscriber:
+class PrefectEventSubscriber:
     """
     Subscribes to a Prefect Cloud event stream, yielding events as they occur.
 
     Example:
 
-        from prefect.events.clients import PrefectCloudEventSubscriber
+        from prefect.events.clients import PrefectEventSubscriber
         from prefect.events.filters import EventFilter, EventNameFilter
 
         filter = EventFilter(event=EventNameFilter(prefix=["prefect.flow-run."]))
 
-        async with PrefectCloudEventSubscriber(api_url, api_key, filter) as subscriber:
+        async with PrefectEventSubscriber(filter=filter) as subscriber:
             async for event in subscriber:
                 print(event.occurred, event.resource.id, event.event)
 
@@ -293,13 +394,14 @@ class PrefectCloudEventSubscriber:
 
     _websocket: Optional[WebSocketClientProtocol]
     _filter: "EventFilter"
-    _seen_events: Mapping[UUID, bool]
+    _seen_events: MutableMapping[UUID, bool]
+
+    _api_key: Optional[str]
 
     def __init__(
         self,
-        api_url: str = None,
-        api_key: str = None,
-        filter: "EventFilter" = None,
+        api_url: Optional[str] = None,
+        filter: Optional["EventFilter"] = None,
         reconnection_attempts: int = 10,
     ):
         """
@@ -309,11 +411,13 @@ class PrefectCloudEventSubscriber:
             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)
+        if not api_url:
+            api_url = PREFECT_API_URL.value()
+            self._api_key = None
 
         from prefect.events.filters import EventFilter
 
-        self._filter = filter or EventFilter()
+        self._filter = filter or EventFilter()  # type: ignore[call-arg]
         self._seen_events = TTLCache(maxsize=SEEN_EVENTS_SIZE, ttl=SEEN_EVENTS_TTL)
 
         socket_url = (
@@ -324,17 +428,16 @@ class PrefectCloudEventSubscriber:
 
         logger.debug("Connecting to %s", socket_url)
 
-        self._api_key = api_key
         self._connect = connect(
             socket_url,
-            subprotocols=["prefect"],
+            subprotocols=[Subprotocol("prefect")],
         )
         self._websocket = None
         self._reconnection_attempts = reconnection_attempts
         if self._reconnection_attempts < 0:
             raise ValueError("reconnection_attempts must be a non-negative integer")
 
-    async def __aenter__(self) -> "PrefectCloudEventSubscriber":
+    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()
@@ -398,7 +501,7 @@ class PrefectCloudEventSubscriber:
         self._websocket = None
         await self._connect.__aexit__(exc_type, exc_val, exc_tb)
 
-    def __aiter__(self) -> "PrefectCloudEventSubscriber":
+    def __aiter__(self) -> Self:
         return self
 
     async def __anext__(self) -> Event:
@@ -416,7 +519,7 @@ class PrefectCloudEventSubscriber:
 
                 while True:
                     message = orjson.loads(await self._websocket.recv())
-                    event = Event.parse_obj(message["event"])
+                    event: Event = Event.parse_obj(message["event"])
 
                     if event.id in self._seen_events:
                         continue
@@ -441,3 +544,58 @@ class PrefectCloudEventSubscriber:
                     # 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 PrefectCloudEventSubscriber(PrefectEventSubscriber):
+    def __init__(
+        self,
+        api_url: Optional[str] = None,
+        api_key: Optional[str] = None,
+        filter: Optional["EventFilter"] = None,
+        reconnection_attempts: int = 10,
+    ):
+        """
+        Args:
+            api_url: The base URL for a Prefect Cloud workspace
+            api_key: The API of an actor with the manage_events scope
+            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
+
+
+class PrefectCloudAccountEventSubscriber(PrefectCloudEventSubscriber):
+    def __init__(
+        self,
+        api_url: Optional[str] = None,
+        api_key: Optional[str] = None,
+        filter: Optional["EventFilter"] = None,
+        reconnection_attempts: int = 10,
+    ):
+        """
+        Args:
+            api_url: The base URL for a Prefect Cloud workspace
+            api_key: The API of an actor with the manage_events scope
+            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)
+
+        account_api_url, _, _ = api_url.partition("/workspaces/")
+
+        super().__init__(
+            api_url=account_api_url,
+            filter=filter,
+            reconnection_attempts=reconnection_attempts,
+        )
+
+        self._api_key = api_key

prefect/events/filters.py

@@ -13,7 +13,34 @@ from .schemas.events import Event, Resource, ResourceSpecification
 if HAS_PYDANTIC_V2:
     from pydantic.v1 import Field, PrivateAttr
 else:
-    from pydantic import Field, PrivateAttr
+    from pydantic import Field, PrivateAttr  # type: ignore
+
+
+class AutomationFilterCreated(PrefectBaseModel):
+    """Filter by `Automation.created`."""
+
+    before_: Optional[DateTimeTZ] = Field(
+        default=None,
+        description="Only include automations created before this datetime",
+    )
+
+
+class AutomationFilterName(PrefectBaseModel):
+    """Filter by `Automation.created`."""
+
+    any_: Optional[List[str]] = Field(
+        default=None,
+        description="Only include automations with names that match any of these strings",
+    )
+
+
+class AutomationFilter(PrefectBaseModel):
+    name: Optional[AutomationFilterName] = Field(
+        default=None, description="Filter criteria for `Automation.name`"
+    )
+    created: Optional[AutomationFilterCreated] = Field(
+        default=None, description="Filter criteria for `Automation.created`"
+    )
 
 
 class EventDataFilter(PrefectBaseModel, extra="forbid"):
@@ -203,7 +230,7 @@ class EventOrder(AutoEnum):
 
 class EventFilter(EventDataFilter):
     occurred: EventOccurredFilter = Field(
-        default_factory=EventOccurredFilter,
+        default_factory=lambda: EventOccurredFilter(),
         description="Filter criteria for when the events occurred",
     )
     event: Optional[EventNameFilter] = Field(
@@ -220,7 +247,7 @@ class EventFilter(EventDataFilter):
         None, description="Filter criteria for the related resources of the event"
     )
     id: EventIDFilter = Field(
-        default_factory=EventIDFilter,
+        default_factory=lambda: EventIDFilter(id=[]),
         description="Filter criteria for the events' ID",
     )
 

prefect/events/instrument.py

@@ -10,7 +10,9 @@ from typing import (
     Set,
     Tuple,
     Type,
+    TypeVar,
     Union,
+    cast,
 )
 
 from prefect.events import emit_event
@@ -41,45 +43,45 @@ def emit_instance_method_called_event(
     )
 
 
-def instrument_instance_method_call():
-    def instrument(function):
-        if is_instrumented(function):
-            return function
+F = TypeVar("F", bound=Callable)
 
-        if inspect.iscoroutinefunction(function):
 
-            @functools.wraps(function)
-            async def inner(self, *args, **kwargs):
-                success = True
-                try:
-                    return await function(self, *args, **kwargs)
-                except Exception as exc:
-                    success = False
-                    raise exc
-                finally:
-                    emit_instance_method_called_event(
-                        instance=self, method_name=function.__name__, successful=success
-                    )
+def instrument_instance_method_call(function: F) -> F:
+    if is_instrumented(function):
+        return function
 
-        else:
+    if inspect.iscoroutinefunction(function):
 
-            @functools.wraps(function)
-            def inner(self, *args, **kwargs):
-                success = True
-                try:
-                    return function(self, *args, **kwargs)
-                except Exception as exc:
-                    success = False
-                    raise exc
-                finally:
-                    emit_instance_method_called_event(
-                        instance=self, method_name=function.__name__, successful=success
-                    )
+        @functools.wraps(function)
+        async def inner(self, *args, **kwargs):
+            success = True
+            try:
+                return await function(self, *args, **kwargs)
+            except Exception as exc:
+                success = False
+                raise exc
+            finally:
+                emit_instance_method_called_event(
+                    instance=self, method_name=function.__name__, successful=success
+                )
 
-        setattr(inner, "__events_instrumented__", True)
-        return inner
+    else:
 
-    return instrument
+        @functools.wraps(function)
+        def inner(self, *args, **kwargs):
+            success = True
+            try:
+                return function(self, *args, **kwargs)
+            except Exception as exc:
+                success = False
+                raise exc
+            finally:
+                emit_instance_method_called_event(
+                    instance=self, method_name=function.__name__, successful=success
+                )
+
+    setattr(inner, "__events_instrumented__", True)
+    return cast(F, inner)
 
 
 def is_instrumented(function: Callable) -> bool:
@@ -119,17 +121,15 @@ def instrument_method_calls_on_class_instances(cls: Type) -> Type:
     """
 
     required_events_methods = ["_event_kind", "_event_method_called_resources"]
-    for method in required_events_methods:
-        if not hasattr(cls, method):
+    for method_name in required_events_methods:
+        if not hasattr(cls, method_name):
             raise RuntimeError(
-                f"Unable to instrument class {cls}. Class must define {method!r}."
+                f"Unable to instrument class {cls}. Class must define {method_name!r}."
             )
 
-    decorator = instrument_instance_method_call()
-
-    for name, method in instrumentable_methods(
+    for method_name, method in instrumentable_methods(
         cls,
         exclude_methods=getattr(cls, "_events_excluded_methods", []),
     ):
-        setattr(cls, name, decorator(method))
+        setattr(cls, method_name, instrument_instance_method_call(method))
     return cls

prefect/events/related.py

@@ -57,6 +57,7 @@ async def related_resources_from_run_context(
     exclude: Optional[Set[str]] = None,
 ) -> List[RelatedResource]:
     from prefect.client.orchestration import get_client
+    from prefect.client.schemas.objects import FlowRun
     from prefect.context import FlowRunContext, TaskRunContext
 
     if exclude is None:
@@ -111,7 +112,7 @@ async def related_resources_from_run_context(
 
         flow_run = related_objects[0]["object"]
 
-        if flow_run:
+        if isinstance(flow_run, FlowRun):
             related_objects += list(
                 await asyncio.gather(
                     _get_and_cache_related_object(