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

prefect/events/actions.py

@@ -1,36 +1,70 @@
+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
+    from pydantic.v1 import Field, root_validator
 else:
-    from pydantic import Field
-
-from typing_extensions import Literal
-
+    from pydantic import Field, root_validator
 from prefect._internal.schemas.bases import PrefectBaseModel
+from prefect.client.schemas.objects import StateType
 
 
-class Action(PrefectBaseModel):
+class Action(PrefectBaseModel, abc.ABC):
     """An Action that may be performed when an Automation is triggered"""
 
     type: str
 
+    def describe_for_cli(self) -> str:
+        """A human-readable description of the action"""
+        return self.type.replace("-", " ").capitalize()
+
 
 class DoNothing(Action):
-    """Do nothing, which may be helpful for testing automations"""
+    """Do nothing when an Automation is triggered"""
 
     type: Literal["do-nothing"] = "do-nothing"
 
 
-class RunDeployment(Action):
-    """Run the given deployment with the given parameters"""
+class DeploymentAction(Action):
+    """Base class for Actions that operate on Deployments and need to infer them from
+    events"""
+
+    source: Literal["selected", "inferred"] = Field(
+        "selected",
+        description=(
+            "Whether this Action applies to a specific selected "
+            "deployment (given by `deployment_id`), or to a deployment that is "
+            "inferred from the triggering event.  If the source is 'inferred', "
+            "the `deployment_id` may not be set.  If the source is 'selected', the "
+            "`deployment_id` must be set."
+        ),
+    )
+    deployment_id: Optional[UUID] = Field(
+        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"))
+        if wants_selected_deployment != has_deployment_id:
+            raise ValueError(
+                "deployment_id is "
+                + ("not allowed" if has_deployment_id else "required")
+            )
+        return values
+
+
+class RunDeployment(DeploymentAction):
+    """Runs the given deployment with the given parameters"""
 
     type: Literal["run-deployment"] = "run-deployment"
 
-    source: Literal["selected"] = "selected"
     parameters: Optional[Dict[str, Any]] = Field(
         None,
         description=(
@@ -38,26 +72,223 @@ class RunDeployment(Action):
             "deployment's default parameters"
         ),
     )
-    deployment_id: UUID = Field(..., description="The identifier of the deployment")
     job_variables: Optional[Dict[str, Any]] = Field(
         None,
         description=(
-            "Job variables to pass to the run, or None to use the "
-            "deployment's default job variables"
+            "The job variables to pass to the created flow run, or None "
+            "to use the deployment's default job variables"
         ),
     )
 
 
+class PauseDeployment(DeploymentAction):
+    """Pauses the given Deployment"""
+
+    type: Literal["pause-deployment"] = "pause-deployment"
+
+
+class ResumeDeployment(DeploymentAction):
+    """Resumes the given Deployment"""
+
+    type: Literal["resume-deployment"] = "resume-deployment"
+
+
+class ChangeFlowRunState(Action):
+    """Changes the state of a flow run associated with the trigger"""
+
+    type: Literal["change-flow-run-state"] = "change-flow-run-state"
+
+    name: Optional[str] = Field(
+        None,
+        description="The name of the state to change the flow run to",
+    )
+    state: StateType = Field(
+        ...,
+        description="The type of the state to change the flow run to",
+    )
+    message: Optional[str] = Field(
+        None,
+        description="An optional message to associate with the state change",
+    )
+
+
+class CancelFlowRun(Action):
+    """Cancels a flow run associated with the trigger"""
+
+    type: Literal["cancel-flow-run"] = "cancel-flow-run"
+
+
+class SuspendFlowRun(Action):
+    """Suspends a flow run associated with the trigger"""
+
+    type: Literal["suspend-flow-run"] = "suspend-flow-run"
+
+
+class CallWebhook(Action):
+    """Call a webhook when an Automation is triggered."""
+
+    type: Literal["call-webhook"] = "call-webhook"
+    block_document_id: UUID = Field(
+        description="The identifier of the webhook block to use"
+    )
+    payload: str = Field(
+        default="",
+        description="An optional templatable payload to send when calling the webhook.",
+    )
+
+
 class SendNotification(Action):
-    """Send a notification with the given parameters"""
+    """Send a notification when an Automation is triggered"""
 
     type: Literal["send-notification"] = "send-notification"
-
     block_document_id: UUID = Field(
-        ..., description="The identifier of the notification block"
+        description="The identifier of the notification block to use"
     )
-    body: str = Field(..., description="Notification body")
-    subject: Optional[str] = Field(None, description="Notification subject")
+    subject: str = Field("Prefect automated notification")
+    body: str = Field(description="The text of the notification to send")
+
+
+class WorkPoolAction(Action):
+    """Base class for Actions that operate on Work Pools and need to infer them from
+    events"""
+
+    source: Literal["selected", "inferred"] = Field(
+        "selected",
+        description=(
+            "Whether this Action applies to a specific selected "
+            "work pool (given by `work_pool_id`), or to a work pool that is "
+            "inferred from the triggering event.  If the source is 'inferred', "
+            "the `work_pool_id` may not be set.  If the source is 'selected', the "
+            "`work_pool_id` must be set."
+        ),
+    )
+    work_pool_id: Optional[UUID] = Field(
+        None,
+        description="The identifier of the work pool to pause",
+    )
+
+
+class PauseWorkPool(WorkPoolAction):
+    """Pauses a Work Pool"""
+
+    type: Literal["pause-work-pool"] = "pause-work-pool"
+
+
+class ResumeWorkPool(WorkPoolAction):
+    """Resumes a Work Pool"""
+
+    type: Literal["resume-work-pool"] = "resume-work-pool"
+
+
+class WorkQueueAction(Action):
+    """Base class for Actions that operate on Work Queues and need to infer them from
+    events"""
+
+    source: Literal["selected", "inferred"] = Field(
+        "selected",
+        description=(
+            "Whether this Action applies to a specific selected "
+            "work queue (given by `work_queue_id`), or to a work queue that is "
+            "inferred from the triggering event.  If the source is 'inferred', "
+            "the `work_queue_id` may not be set.  If the source is 'selected', the "
+            "`work_queue_id` must be set."
+        ),
+    )
+    work_queue_id: Optional[UUID] = Field(
+        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"))
+        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
+
+
+class PauseWorkQueue(WorkQueueAction):
+    """Pauses a Work Queue"""
+
+    type: Literal["pause-work-queue"] = "pause-work-queue"
+
+
+class ResumeWorkQueue(WorkQueueAction):
+    """Resumes a Work Queue"""
+
+    type: Literal["resume-work-queue"] = "resume-work-queue"
+
+
+class AutomationAction(Action):
+    """Base class for Actions that operate on Automations and need to infer them from
+    events"""
+
+    source: Literal["selected", "inferred"] = Field(
+        "selected",
+        description=(
+            "Whether this Action applies to a specific selected "
+            "automation (given by `automation_id`), or to an automation that is "
+            "inferred from the triggering event.  If the source is 'inferred', "
+            "the `automation_id` may not be set.  If the source is 'selected', the "
+            "`automation_id` must be set."
+        ),
+    )
+    automation_id: Optional[UUID] = Field(
+        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"))
+        if wants_selected_automation != has_automation_id:
+            raise ValueError(
+                "automation_id is "
+                + ("not allowed" if has_automation_id else "required")
+            )
+        return values
+
+
+class PauseAutomation(AutomationAction):
+    """Pauses a Work Queue"""
+
+    type: Literal["pause-automation"] = "pause-automation"
+
+
+class ResumeAutomation(AutomationAction):
+    """Resumes a Work Queue"""
+
+    type: Literal["resume-automation"] = "resume-automation"
+
+
+class DeclareIncident(Action):
+    """Declares an incident for the triggering event.  Only available on Prefect Cloud"""
+
+    type: Literal["declare-incident"] = "declare-incident"
 
 
-ActionTypes = Union[DoNothing, RunDeployment, SendNotification]
+# The actual action types that we support.  It's important to update this
+# Union when adding new subclasses of Action so that they are available for clients
+# and in the OpenAPI docs
+ActionTypes: TypeAlias = Union[
+    DoNothing,
+    RunDeployment,
+    PauseDeployment,
+    ResumeDeployment,
+    CancelFlowRun,
+    ChangeFlowRunState,
+    PauseWorkQueue,
+    ResumeWorkQueue,
+    SendNotification,
+    CallWebhook,
+    PauseAutomation,
+    ResumeAutomation,
+    SuspendFlowRun,
+    PauseWorkPool,
+    ResumeWorkPool,
+    # Prefect Cloud only
+    DeclareIncident,
+]

prefect/events/cli/automations.py

@@ -0,0 +1,163 @@
+"""
+Command line interface for working with automations.
+"""
+
+import functools
+
+import orjson
+import yaml as pyyaml
+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.client.orchestration import get_client
+
+automations_app = PrefectTyper(
+    name="automation",
+    help="Commands for managing automations.",
+)
+app.add_typer(automations_app, aliases=["automations"])
+
+
+def requires_automations(func):
+    @functools.wraps(func)
+    async def wrapper(*args, **kwargs):
+        try:
+            return await func(*args, **kwargs)
+        except RuntimeError as exc:
+            if "Enable experimental" in str(exc):
+                exit_with_error(str(exc))
+            raise
+
+    return wrapper
+
+
+@automations_app.command()
+@requires_automations
+async def ls():
+    """List all automations."""
+    async with get_client() as client:
+        automations = await client.read_automations()
+
+    table = Table(title="Automations", show_lines=True)
+
+    table.add_column("Automation")
+    table.add_column("Enabled")
+    table.add_column("Trigger")
+    table.add_column("Actions")
+
+    for automation in automations:
+        identifier_column = Text()
+
+        identifier_column.append(automation.name, style="white bold")
+        identifier_column.append("\n")
+
+        identifier_column.append(str(automation.id), style="cyan")
+        identifier_column.append("\n")
+
+        if automation.description:
+            identifier_column.append(automation.description, style="white")
+            identifier_column.append("\n")
+
+        if automation.actions_on_trigger or automation.actions_on_resolve:
+            actions = (
+                [
+                    f"(trigger) {action.describe_for_cli()}"
+                    for action in automation.actions
+                ]
+                + [
+                    f"(trigger) {action.describe_for_cli()}"
+                    for action in automation.actions_on_trigger
+                ]
+                + [
+                    f"(resolve) {action.describe_for_cli()}"
+                    for action in automation.actions
+                ]
+                + [
+                    f"(resolve) {action.describe_for_cli()}"
+                    for action in automation.actions_on_resolve
+                ]
+            )
+        else:
+            actions = [action.describe_for_cli() for action in automation.actions]
+
+        table.add_row(
+            identifier_column,
+            str(automation.enabled),
+            automation.trigger.describe_for_cli(),
+            "\n".join(actions),
+        )
+
+    app.console.print(table)
+
+
+@automations_app.command()
+@requires_automations
+async def inspect(id_or_name: str, yaml: bool = False, json: bool = False):
+    """Inspect an automation."""
+    async with get_client() as client:
+        automation = await client.find_automation(id_or_name)
+        if not automation:
+            exit_with_error(f"Automation {id_or_name!r} not found.")
+
+    if yaml:
+        app.console.print(
+            pyyaml.dump(automation.dict(json_compatible=True), sort_keys=False)
+        )
+    elif json:
+        app.console.print(
+            orjson.dumps(
+                automation.dict(json_compatible=True), option=orjson.OPT_INDENT_2
+            ).decode()
+        )
+    else:
+        app.console.print(Pretty(automation))
+
+
+@automations_app.command(aliases=["enable"])
+@requires_automations
+async def resume(id_or_name: str):
+    """Resume an automation."""
+    async with get_client() as client:
+        automation = await client.find_automation(id_or_name)
+        if not automation:
+            exit_with_error(f"Automation {id_or_name!r} not found.")
+
+    async with get_client() as client:
+        await client.resume_automation(automation.id)
+
+    exit_with_success(f"Resumed automation {automation.name!r} ({automation.id})")
+
+
+@automations_app.command(aliases=["disable"])
+@requires_automations
+async def pause(id_or_name: str):
+    """Pause an automation."""
+    async with get_client() as client:
+        automation = await client.find_automation(id_or_name)
+        if not automation:
+            exit_with_error(f"Automation {id_or_name!r} not found.")
+
+    async with get_client() as client:
+        await client.pause_automation(automation.id)
+
+    exit_with_success(f"Paused automation {automation.name!r} ({automation.id})")
+
+
+@automations_app.command()
+@requires_automations
+async def delete(id_or_name: str):
+    """Delete an automation."""
+    async with get_client() as client:
+        automation = await client.find_automation(id_or_name)
+
+    if not automation:
+        exit_with_success(f"Automation {id_or_name!r} not found.")
+
+    async with get_client() as client:
+        await client.delete_automation(automation.id)
+
+    exit_with_success(f"Deleted automation {automation.name!r} ({automation.id})")

prefect/events/clients.py

@@ -14,6 +14,7 @@ from typing import (
 )
 from uuid import UUID
 
+import httpx
 import orjson
 import pendulum
 from cachetools import TTLCache
@@ -25,9 +26,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 +42,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: "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"""
 
@@ -119,6 +173,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"""
 
@@ -274,18 +374,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)
 
@@ -298,7 +398,6 @@ class PrefectCloudEventSubscriber:
     def __init__(
         self,
         api_url: str = None,
-        api_key: str = None,
         filter: "EventFilter" = None,
         reconnection_attempts: int = 10,
     ):
@@ -309,7 +408,9 @@ 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
 
@@ -324,7 +425,6 @@ class PrefectCloudEventSubscriber:
 
         logger.debug("Connecting to %s", socket_url)
 
-        self._api_key = api_key
         self._connect = connect(
             socket_url,
             subprotocols=["prefect"],
@@ -441,3 +541,29 @@ class PrefectCloudEventSubscriber:
                     # a standard load balancer timeout, but after that, just take a
                     # beat to let things come back around.
                     await asyncio.sleep(1)
+
+
+class PrefectCloudEventSubscriber(PrefectEventSubscriber):
+    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)
+
+        super().__init__(
+            api_url=api_url,
+            filter=filter,
+            reconnection_attempts=reconnection_attempts,
+        )
+
+        self._api_key = api_key