prefect-client 2.14.9__py3-none-any.whl → 2.14.11__py3-none-any.whl

prefect/events/clients.py

@@ -1,12 +1,42 @@
 import abc
 import asyncio
 from types import TracebackType
-from typing import Any, ClassVar, Dict, List, Optional, Tuple, Type
-
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    ClassVar,
+    Dict,
+    List,
+    Mapping,
+    Optional,
+    Tuple,
+    Type,
+)
+from uuid import UUID
+
+import orjson
+import pendulum
+
+try:
+    from cachetools import TTLCache
+except ImportError:
+    pass
+from starlette.status import WS_1008_POLICY_VIOLATION
 from websockets.client import WebSocketClientProtocol, connect
-from websockets.exceptions import ConnectionClosed
+from websockets.exceptions import (
+    ConnectionClosed,
+    ConnectionClosedError,
+    ConnectionClosedOK,
+)
 
 from prefect.events import Event
+from prefect.logging import get_logger
+from prefect.settings import PREFECT_API_KEY, PREFECT_API_URL
+
+if TYPE_CHECKING:
+    from prefect.events.filters import EventFilter
+
+logger = get_logger(__name__)
 
 
 class EventsClient(abc.ABC):
@@ -79,6 +109,20 @@ class AssertingEventsClient(EventsClient):
         return self
 
 
+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
+
+
 class PrefectCloudEventsClient(EventsClient):
     """A Prefect Events client that streams Events to a Prefect Cloud Workspace"""
 
@@ -87,8 +131,8 @@ class PrefectCloudEventsClient(EventsClient):
 
     def __init__(
         self,
-        api_url: str,
-        api_key: str,
+        api_url: str = None,
+        api_key: str = None,
         reconnection_attempts: int = 10,
         checkpoint_every: int = 20,
     ):
@@ -101,6 +145,8 @@ class PrefectCloudEventsClient(EventsClient):
             checkpoint_every: How often the client should sync with the server to
                 confirm receipt of all previously sent events
         """
+        api_url, api_key = _get_api_url_and_key(api_url, api_key)
+
         socket_url = (
             api_url.replace("https://", "wss://")
             .replace("http://", "ws://")
@@ -195,3 +241,168 @@ class PrefectCloudEventsClient(EventsClient):
                     # a standard load balancer timeout, but after that, just take a
                     # beat to let things come back around.
                     await asyncio.sleep(1)
+
+
+SEEN_EVENTS_SIZE = 500_000
+SEEN_EVENTS_TTL = 120
+
+
+class PrefectCloudEventSubscriber:
+    """
+    Subscribes to a Prefect Cloud event stream, yielding events as they occur.
+
+    Example:
+
+        from prefect.events.clients import PrefectCloudEventSubscriber
+        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 for event in subscriber:
+                print(event.occurred, event.resource.id, event.event)
+
+    """
+
+    _websocket: Optional[WebSocketClientProtocol]
+    _filter: "EventFilter"
+    _seen_events: Mapping[UUID, bool]
+
+    def __init__(
+        self,
+        api_url: str = None,
+        api_key: str = None,
+        filter: "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)
+
+        from prefect.events.filters import EventFilter
+
+        self._filter = filter or EventFilter()
+        self._seen_events = TTLCache(maxsize=SEEN_EVENTS_SIZE, ttl=SEEN_EVENTS_TTL)
+
+        socket_url = (
+            api_url.replace("https://", "wss://")
+            .replace("http://", "ws://")
+            .rstrip("/")
+        ) + "/events/out"
+
+        logger.debug("Connecting to %s", socket_url)
+
+        self._api_key = api_key
+        self._connect = connect(
+            socket_url,
+            subprotocols=["prefect"],
+        )
+        self._websocket = None
+        self._reconnection_attempts = reconnection_attempts
+
+    async def __aenter__(self) -> "PrefectCloudEventSubscriber":
+        # 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()
+        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
+
+        logger.debug("  authenticating...")
+        await self._websocket.send(
+            orjson.dumps({"type": "auth", "token": self._api_key}).decode()
+        )
+
+        try:
+            message = orjson.loads(await self._websocket.recv())
+            logger.debug("  auth result %s", message)
+            assert message["type"] == "auth_success"
+        except (AssertionError, ConnectionClosedError) as e:
+            if isinstance(e, AssertionError) or e.code == WS_1008_POLICY_VIOLATION:
+                raise Exception(
+                    "Unable to authenticate to the event stream. Please ensure the "
+                    "provided api_key you are using is valid for this environment."
+                ) from e
+            raise
+
+        from prefect.events.filters import EventOccurredFilter
+
+        self._filter.occurred = EventOccurredFilter(
+            since=pendulum.now("UTC").subtract(minutes=1),
+            until=pendulum.now("UTC").add(years=1),
+        )
+
+        logger.debug("  filtering events since %s...", self._filter.occurred.since)
+        filter_message = {
+            "type": "filter",
+            "filter": self._filter.dict(json_compatible=True),
+        }
+        await self._websocket.send(orjson.dumps(filter_message).decode())
+
+    async def __aexit__(
+        self,
+        exc_type: Optional[Type[Exception]],
+        exc_val: Optional[Exception],
+        exc_tb: Optional[TracebackType],
+    ) -> None:
+        self._websocket = None
+        await self._connect.__aexit__(exc_type, exc_val, exc_tb)
+
+    def __aiter__(self) -> "PrefectCloudEventSubscriber":
+        return self
+
+    async def __anext__(self) -> Event:
+        for i in range(self._reconnection_attempts + 1):
+            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:
+                    await self._reconnect()
+                    assert self._websocket
+
+                while True:
+                    message = orjson.loads(await self._websocket.recv())
+                    event = Event.parse_obj(message["event"])
+
+                    if event.id in self._seen_events:
+                        continue
+                    self._seen_events[event.id] = True
+
+                    return event
+            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)

prefect/events/filters.py

@@ -0,0 +1,214 @@
+from typing import List, Optional, Tuple, cast
+from uuid import UUID
+
+import pendulum
+
+from prefect._internal.pydantic import HAS_PYDANTIC_V2
+from prefect.events.schemas import Event, Resource, ResourceSpecification
+from prefect.server.utilities.schemas import DateTimeTZ, PrefectBaseModel
+
+if HAS_PYDANTIC_V2:
+    from pydantic.v1 import Field
+else:
+    from pydantic import Field
+
+
+class EventDataFilter(PrefectBaseModel):
+    """A base class for filtering event data."""
+
+    class Config:
+        extra = "forbid"
+
+    def get_filters(self) -> List["EventDataFilter"]:
+        return [
+            filter
+            for filter in [
+                getattr(self, name)
+                for name, field in self.__fields__.items()
+                if issubclass(field.type_, EventDataFilter)
+            ]
+            if filter
+        ]
+
+    def includes(self, event: Event) -> bool:
+        """Does the given event match the criteria of this filter?"""
+        return all(filter.includes(event) for filter in self.get_filters())
+
+    def excludes(self, event: Event) -> bool:
+        """Would the given filter exclude this event?"""
+        return not self.includes(event)
+
+
+class EventOccurredFilter(EventDataFilter):
+    since: DateTimeTZ = Field(
+        default_factory=lambda: cast(
+            DateTimeTZ,
+            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")),
+        description="Only include events prior to this time (inclusive)",
+    )
+
+    def includes(self, event: Event) -> bool:
+        return self.since <= event.occurred <= self.until
+
+
+class EventNameFilter(EventDataFilter):
+    prefix: Optional[List[str]] = Field(
+        None, description="Only include events matching one of these prefixes"
+    )
+    exclude_prefix: Optional[List[str]] = Field(
+        None, description="Exclude events matching one of these prefixes"
+    )
+
+    name: Optional[List[str]] = Field(
+        None, description="Only include events matching one of these names exactly"
+    )
+    exclude_name: Optional[List[str]] = Field(
+        None, description="Exclude events matching one of these names exactly"
+    )
+
+    def includes(self, event: Event) -> bool:
+        if self.prefix:
+            if not any(event.event.startswith(prefix) for prefix in self.prefix):
+                return False
+
+        if self.exclude_prefix:
+            if any(event.event.startswith(prefix) for prefix in self.exclude_prefix):
+                return False
+
+        if self.name:
+            if not any(event.event == name for name in self.name):
+                return False
+
+        if self.exclude_name:
+            if any(event.event == name for name in self.exclude_name):
+                return False
+
+        return True
+
+
+class EventResourceFilter(EventDataFilter):
+    id: Optional[List[str]] = Field(
+        None, description="Only include events for resources with these IDs"
+    )
+    id_prefix: Optional[List[str]] = Field(
+        None,
+        description=(
+            "Only include events for resources with IDs starting with these prefixes."
+        ),
+    )
+    labels: Optional[ResourceSpecification] = Field(
+        None, description="Only include events for resources with these labels"
+    )
+
+    def includes(self, event: Event) -> bool:
+        if self.id:
+            if not any(event.resource.id == resource_id for resource_id in self.id):
+                return False
+
+        if self.id_prefix:
+            if not any(
+                event.resource.id.startswith(prefix) for prefix in self.id_prefix
+            ):
+                return False
+
+        if self.labels:
+            if not self.labels.matches(event.resource):
+                return False
+
+        return True
+
+
+class EventRelatedFilter(EventDataFilter):
+    id: Optional[List[str]] = Field(
+        None, description="Only include events for related resources with these IDs"
+    )
+    role: Optional[List[str]] = Field(
+        None, description="Only include events for related resources in these roles"
+    )
+    resources_in_roles: Optional[List[Tuple[str, str]]] = Field(
+        None,
+        description=(
+            "Only include events with specific related resources in specific roles"
+        ),
+    )
+    labels: Optional[ResourceSpecification] = Field(
+        None, description="Only include events for related resources with these labels"
+    )
+
+
+class EventAnyResourceFilter(EventDataFilter):
+    id: Optional[List[str]] = Field(
+        None, description="Only include events for resources with these IDs"
+    )
+    id_prefix: Optional[List[str]] = Field(
+        None,
+        description=(
+            "Only include events for resources with IDs starting with these prefixes"
+        ),
+    )
+    labels: Optional[ResourceSpecification] = Field(
+        None, description="Only include events for related resources with these labels"
+    )
+
+    def includes(self, event: Event) -> bool:
+        resources = [event.resource] + event.related
+        if not any(self._includes(resource) for resource in resources):
+            return False
+        return True
+
+    def _includes(self, resource: Resource) -> bool:
+        if self.id:
+            if not any(resource.id == resource_id for resource_id in self.id):
+                return False
+
+        if self.id_prefix:
+            if not any(resource.id.startswith(prefix) for prefix in self.id_prefix):
+                return False
+
+        if self.labels:
+            if not self.labels.matches(resource):
+                return False
+
+        return True
+
+
+class EventIDFilter(EventDataFilter):
+    id: Optional[List[UUID]] = Field(
+        None, description="Only include events with one of these IDs"
+    )
+
+    def includes(self, event: Event) -> bool:
+        if self.id:
+            if not any(event.id == id for id in self.id):
+                return False
+
+        return True
+
+
+class EventFilter(EventDataFilter):
+    occurred: EventOccurredFilter = Field(
+        default_factory=EventOccurredFilter,
+        description="Filter criteria for when the events occurred",
+    )
+    event: Optional[EventNameFilter] = Field(
+        None,
+        description="Filter criteria for the event name",
+    )
+    any_resource: Optional[EventAnyResourceFilter] = Field(
+        None, description="Filter criteria for any resource involved in the event"
+    )
+    resource: Optional[EventResourceFilter] = Field(
+        None, description="Filter criteria for the resource of the event"
+    )
+    related: Optional[EventRelatedFilter] = Field(
+        None, description="Filter criteria for the related resources of the event"
+    )
+    id: EventIDFilter = Field(
+        default_factory=EventIDFilter,
+        description="Filter criteria for the events' ID",
+    )

prefect/exceptions.py

@@ -296,6 +296,10 @@ class Pause(PrefectSignal):
     Raised when a flow run is PAUSED and needs to exit for resubmission.
     """
 
+    def __init__(self, *args, state=None, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.state = state
+
 
 class ExternalSignal(BaseException):
     """

prefect/flows.py

@@ -556,6 +556,7 @@ class Flow(Generic[P, R]):
         cron: Optional[str] = None,
         rrule: Optional[str] = None,
         schedule: Optional[SCHEDULE_TYPES] = None,
+        is_schedule_active: Optional[bool] = None,
         parameters: Optional[dict] = None,
         triggers: Optional[List[DeploymentTrigger]] = None,
         description: Optional[str] = None,
@@ -578,6 +579,9 @@ class Flow(Generic[P, R]):
             timezone: A timezone to use for the schedule. Defaults to UTC.
             triggers: A list of triggers that will kick off runs of this deployment.
             schedule: A schedule object defining when to execute runs of this deployment.
+            is_schedule_active: Whether or not to set the schedule for this deployment as active. If
+                not provided when creating a deployment, the schedule will be set as active. If not
+                provided when updating a deployment, the schedule's activation will not be changed.
             parameters: A dictionary of default parameter values to pass to runs of this deployment.
             description: A description for the created deployment. Defaults to the flow's
                 description if not provided.
@@ -623,6 +627,7 @@ class Flow(Generic[P, R]):
                 cron=cron,
                 rrule=rrule,
                 schedule=schedule,
+                is_schedule_active=is_schedule_active,
                 tags=tags,
                 triggers=triggers,
                 parameters=parameters or {},
@@ -641,6 +646,7 @@ class Flow(Generic[P, R]):
                 cron=cron,
                 rrule=rrule,
                 schedule=schedule,
+                is_schedule_active=is_schedule_active,
                 tags=tags,
                 triggers=triggers,
                 parameters=parameters or {},
@@ -660,6 +666,7 @@ class Flow(Generic[P, R]):
         cron: Optional[str] = None,
         rrule: Optional[str] = None,
         schedule: Optional[SCHEDULE_TYPES] = None,
+        is_schedule_active: Optional[bool] = None,
         triggers: Optional[List[DeploymentTrigger]] = None,
         parameters: Optional[dict] = None,
         description: Optional[str] = None,
@@ -682,6 +689,9 @@ class Flow(Generic[P, R]):
             triggers: A list of triggers that will kick off runs of this deployment.
             schedule: A schedule object defining when to execute runs of this deployment. Used to
                 define additional scheduling options like `timezone`.
+            is_schedule_active: Whether or not to set the schedule for this deployment as active. If
+                not provided when creating a deployment, the schedule will be set as active. If not
+                provided when updating a deployment, the schedule's activation will not be changed.
             parameters: A dictionary of default parameter values to pass to runs of this deployment.
             description: A description for the created deployment. Defaults to the flow's
                 description if not provided.
@@ -738,6 +748,7 @@ class Flow(Generic[P, R]):
             cron=cron,
             rrule=rrule,
             schedule=schedule,
+            is_schedule_active=is_schedule_active,
             parameters=parameters,
             description=description,
             tags=tags,
@@ -852,6 +863,7 @@ class Flow(Generic[P, R]):
         cron: Optional[str] = None,
         rrule: Optional[str] = None,
         schedule: Optional[SCHEDULE_TYPES] = None,
+        is_schedule_active: Optional[bool] = None,
         triggers: Optional[List[DeploymentTrigger]] = None,
         parameters: Optional[dict] = None,
         description: Optional[str] = None,
@@ -891,6 +903,9 @@ class Flow(Generic[P, R]):
             triggers: A list of triggers that will kick off runs of this deployment.
             schedule: A schedule object defining when to execute runs of this deployment. Used to
                 define additional scheduling options like `timezone`.
+            is_schedule_active: Whether or not to set the schedule for this deployment as active. If
+                not provided when creating a deployment, the schedule will be set as active. If not
+                provided when updating a deployment, the schedule's activation will not be changed.
             parameters: A dictionary of default parameter values to pass to runs of this deployment.
             description: A description for the created deployment. Defaults to the flow's
                 description if not provided.
@@ -956,6 +971,7 @@ class Flow(Generic[P, R]):
             cron=cron,
             rrule=rrule,
             schedule=schedule,
+            is_schedule_active=is_schedule_active,
             triggers=triggers,
             parameters=parameters,
             description=description,

prefect/infrastructure/base.py

@@ -11,22 +11,27 @@ from prefect._internal.compatibility.experimental import (
     experiment_enabled,
 )
 from prefect._internal.pydantic import HAS_PYDANTIC_V2
+from prefect.client.schemas.actions import WorkPoolCreate
+from prefect.exceptions import ObjectAlreadyExists
 
 if HAS_PYDANTIC_V2:
     import pydantic.v1 as pydantic
 else:
     import pydantic
 
+from rich.console import Console
 from typing_extensions import Self
 
 import prefect
-from prefect.blocks.core import Block
+from prefect.blocks.core import Block, BlockNotSavedError
 from prefect.logging import get_logger
 from prefect.settings import (
     PREFECT_EXPERIMENTAL_WARN,
     PREFECT_EXPERIMENTAL_WARN_ENHANCED_CANCELLATION,
+    PREFECT_UI_URL,
     get_current_settings,
 )
+from prefect.utilities.asyncutils import sync_compatible
 
 MIN_COMPAT_PREFECT_VERSION = "2.0b12"
 
@@ -66,6 +71,106 @@ class Infrastructure(Block, abc.ABC):
         description="The command to run in the infrastructure.",
     )
 
+    async def generate_work_pool_base_job_template(self):
+        if self._block_document_id is None:
+            raise BlockNotSavedError(
+                "Cannot publish as work pool, block has not been saved. Please call"
+                " `.save()` on your block before publishing."
+            )
+
+        block_schema = self.__class__.schema()
+        return {
+            "job_configuration": {"block": "{{ block }}"},
+            "variables": {
+                "type": "object",
+                "properties": {
+                    "block": {
+                        "title": "Block",
+                        "description": (
+                            "The infrastructure block to use for job creation."
+                        ),
+                        "allOf": [{"$ref": f"#/definitions/{self.__class__.__name__}"}],
+                        "default": {
+                            "$ref": {"block_document_id": str(self._block_document_id)}
+                        },
+                    }
+                },
+                "required": ["block"],
+                "definitions": {self.__class__.__name__: block_schema},
+            },
+        }
+
+    def get_corresponding_worker_type(self):
+        return "block"
+
+    @sync_compatible
+    async def publish_as_work_pool(self, work_pool_name: Optional[str] = None):
+        """
+        Creates a work pool configured to use the given block as the job creator.
+
+        Used to migrate from a agents setup to a worker setup.
+
+        Args:
+            work_pool_name: The name to give to the created work pool. If not provided, the name of the current
+                block will be used.
+        """
+
+        base_job_template = await self.generate_work_pool_base_job_template()
+        work_pool_name = work_pool_name or self._block_document_name
+
+        if work_pool_name is None:
+            raise ValueError(
+                "`work_pool_name` must be provided if the block has not been saved."
+            )
+
+        console = Console()
+
+        try:
+            async with prefect.get_client() as client:
+                work_pool = await client.create_work_pool(
+                    work_pool=WorkPoolCreate(
+                        name=work_pool_name,
+                        type=self.get_corresponding_worker_type(),
+                        base_job_template=base_job_template,
+                    )
+                )
+        except ObjectAlreadyExists:
+            console.print(
+                (
+                    f"Work pool with name {work_pool_name!r} already exists, please use"
+                    " a different name."
+                ),
+                style="red",
+            )
+            return
+
+        console.print(
+            f"Work pool {work_pool.name} created!",
+            style="green",
+        )
+        if PREFECT_UI_URL:
+            console.print(
+                "You see your new work pool in the UI at"
+                f" {PREFECT_UI_URL.value()}/work-pools/work-pool/{work_pool.name}"
+            )
+
+        deploy_script = (
+            "my_flow.deploy(work_pool_name='{work_pool.name}', image='my_image:tag')"
+        )
+        if not hasattr(self, "image"):
+            deploy_script = (
+                "my_flow.from_source(source='https://github.com/org/repo.git',"
+                f" entrypoint='flow.py:my_flow').deploy(work_pool_name='{work_pool.name}')"
+            )
+        console.print(
+            "\nYou can deploy a flow to this work pool by calling"
+            f" [blue].deploy[/]:\n\n\t{deploy_script}\n"
+        )
+        console.print(
+            "\nTo start a worker to execute flow runs in this work pool run:\n"
+        )
+        console.print(f"\t[blue]prefect worker start --pool {work_pool.name}[/]\n")
+
     @abc.abstractmethod
     async def run(
         self,