prefect-client 2.18.0__py3-none-any.whl → 2.18.2__py3-none-any.whl

prefect/events/clients.py

@@ -7,10 +7,11 @@ from typing import (
     ClassVar,
     Dict,
     List,
-    Mapping,
+    MutableMapping,
     Optional,
     Tuple,
     Type,
+    cast,
 )
 from uuid import UUID
 
@@ -19,6 +20,7 @@ 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,
@@ -69,7 +71,7 @@ def get_events_client(
 
 
 def get_events_subscriber(
-    filter: "EventFilter" = None,
+    filter: Optional["EventFilter"] = None,
     reconnection_attempts: int = 10,
 ) -> "PrefectEventSubscriber":
     api_url = PREFECT_API_URL.value()
@@ -105,7 +107,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
 
@@ -153,7 +155,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
@@ -227,7 +229,7 @@ class PrefectEventsClient(EventsClient):
 
     def __init__(
         self,
-        api_url: str = None,
+        api_url: Optional[str] = None,
         reconnection_attempts: int = 10,
         checkpoint_every: int = 20,
     ):
@@ -343,8 +345,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,
     ):
@@ -376,7 +378,7 @@ SEEN_EVENTS_TTL = 120
 
 class PrefectEventSubscriber:
     """
-    Subscribes to a Prefect Cloud event stream, yielding events as they occur.
+    Subscribes to a Prefect event stream, yielding events as they occur.
 
     Example:
 
@@ -393,12 +395,14 @@ class PrefectEventSubscriber:
 
     _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,
-        filter: "EventFilter" = None,
+        api_url: Optional[str] = None,
+        filter: Optional["EventFilter"] = None,
         reconnection_attempts: int = 10,
     ):
         """
@@ -409,12 +413,12 @@ class PrefectEventSubscriber:
                 the client should attempt to reconnect
         """
         if not api_url:
-            api_url = PREFECT_API_URL.value()
+            api_url = cast(str, 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 = (
@@ -427,14 +431,14 @@ class PrefectEventSubscriber:
 
         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()
@@ -498,7 +502,7 @@ class PrefectEventSubscriber:
         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:
@@ -516,7 +520,7 @@ class PrefectEventSubscriber:
 
                 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
@@ -541,14 +545,15 @@ class PrefectEventSubscriber:
                     # 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: str = None,
-        api_key: str = None,
-        filter: "EventFilter" = None,
+        api_url: Optional[str] = None,
+        api_key: Optional[str] = None,
+        filter: Optional["EventFilter"] = None,
         reconnection_attempts: int = 10,
     ):
         """
@@ -567,3 +572,31 @@ class PrefectCloudEventSubscriber(PrefectEventSubscriber):
         )
 
         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,10 +13,37 @@ 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 EventDataFilter(PrefectBaseModel, extra="forbid"):
+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"):  # type: ignore[call-arg]
     """A base class for filtering event data."""
 
     _top_level_filter: "EventFilter | None" = PrivateAttr(None)
@@ -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(

prefect/events/schemas/automations.py

@@ -10,23 +10,24 @@ from typing import (
     Optional,
     Set,
     Union,
+    cast,
 )
 from uuid import UUID
 
 from typing_extensions import TypeAlias
 
 from prefect._internal.pydantic import HAS_PYDANTIC_V2
-from prefect._internal.schemas.validators import validate_trigger_within
 
 if HAS_PYDANTIC_V2:
-    from pydantic.v1 import Field, root_validator, validator
+    from pydantic.v1 import Field, PrivateAttr, root_validator, validator
     from pydantic.v1.fields import ModelField
 else:
-    from pydantic import Field, root_validator, validator
-    from pydantic.fields import ModelField
+    from pydantic import Field, PrivateAttr, root_validator, validator  # type: ignore
+    from pydantic.fields import ModelField  # type: ignore
 
 from prefect._internal.schemas.bases import PrefectBaseModel
-from prefect.events.actions import ActionTypes
+from prefect._internal.schemas.validators import validate_trigger_within
+from prefect.events.actions import ActionTypes, RunDeployment
 from prefect.utilities.collections import AutoEnum
 
 from .events import ResourceSpecification
@@ -38,7 +39,7 @@ class Posture(AutoEnum):
     Metric = "Metric"
 
 
-class Trigger(PrefectBaseModel, abc.ABC, extra="ignore"):
+class Trigger(PrefectBaseModel, abc.ABC, extra="ignore"):  # type: ignore[call-arg]
     """
     Base class describing a set of criteria that must be satisfied in order to trigger
     an automation.
@@ -50,6 +51,50 @@ class Trigger(PrefectBaseModel, abc.ABC, extra="ignore"):
     def describe_for_cli(self, indent: int = 0) -> str:
         """Return a human-readable description of this trigger for the CLI"""
 
+    # The following allows the regular Trigger class to be used when serving or
+    # deploying flows, analogous to how the Deployment*Trigger classes work
+
+    _deployment_id: Optional[UUID] = PrivateAttr(default=None)
+
+    def set_deployment_id(self, deployment_id: UUID):
+        self._deployment_id = deployment_id
+
+    def owner_resource(self) -> Optional[str]:
+        return f"prefect.deployment.{self._deployment_id}"
+
+    def actions(self) -> List[ActionTypes]:
+        assert self._deployment_id
+        return [
+            RunDeployment(
+                source="selected",
+                deployment_id=self._deployment_id,
+                parameters=getattr(self, "parameters", None),
+                job_variables=getattr(self, "job_variables", None),
+            )
+        ]
+
+    def as_automation(self) -> "AutomationCore":
+        assert self._deployment_id
+
+        trigger: TriggerTypes = cast(TriggerTypes, self)
+
+        # This is one of the Deployment*Trigger classes, so translate it over to a
+        # plain Trigger
+        if hasattr(self, "trigger_type"):
+            trigger = self.trigger_type(**self.dict())
+
+        return AutomationCore(
+            name=(
+                getattr(self, "name", None)
+                or f"Automation for deployment {self._deployment_id}"
+            ),
+            description="",
+            enabled=getattr(self, "enabled", True),
+            trigger=trigger,
+            actions=self.actions(),
+            owner_resource=self.owner_resource(),
+        )
+
 
 class ResourceTrigger(Trigger, abc.ABC):
     """
@@ -346,7 +391,7 @@ CompoundTrigger.update_forward_refs()
 SequenceTrigger.update_forward_refs()
 
 
-class AutomationCore(PrefectBaseModel, extra="ignore"):
+class AutomationCore(PrefectBaseModel, extra="ignore"):  # type: ignore[call-arg]
     """Defines an action a user wants to take when a certain number of events
     do or don't happen to the matching resources"""