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

prefect/engine.py

@@ -1410,7 +1410,9 @@ def enter_task_run_engine(
         task_runner=task_runner,
     )
 
-    if task.isasync and flow_run_context.flow.isasync:
+    if task.isasync and (
+        flow_run_context.flow is None or flow_run_context.flow.isasync
+    ):
         # return a coro for the user to await if an async task in an async flow
         return from_async.wait_for_call_in_loop_thread(begin_run)
     else:

prefect/events/__init__.py

@@ -2,7 +2,9 @@ from .schemas.events import Event, ReceivedEvent
 from .schemas.events import Resource, RelatedResource, ResourceSpecification
 from .schemas.automations import (
     Automation,
+    AutomationCore,
     Posture,
+    TriggerTypes,
     Trigger,
     ResourceTrigger,
     EventTrigger,
@@ -20,6 +22,27 @@ from .schemas.deployment_triggers import (
     DeploymentCompoundTrigger,
     DeploymentSequenceTrigger,
 )
+from .actions import (
+    ActionTypes,
+    Action,
+    DoNothing,
+    RunDeployment,
+    PauseDeployment,
+    ResumeDeployment,
+    ChangeFlowRunState,
+    CancelFlowRun,
+    SuspendFlowRun,
+    CallWebhook,
+    SendNotification,
+    PauseWorkPool,
+    ResumeWorkPool,
+    PauseWorkQueue,
+    ResumeWorkQueue,
+    PauseAutomation,
+    ResumeAutomation,
+    DeclareIncident,
+)
+from .clients import get_events_client, get_events_subscriber
 from .utilities import emit_event
 
 __all__ = [
@@ -29,7 +52,9 @@ __all__ = [
     "RelatedResource",
     "ResourceSpecification",
     "Automation",
+    "AutomationCore",
     "Posture",
+    "TriggerTypes",
     "Trigger",
     "ResourceTrigger",
     "EventTrigger",
@@ -44,5 +69,25 @@ __all__ = [
     "DeploymentMetricTrigger",
     "DeploymentCompoundTrigger",
     "DeploymentSequenceTrigger",
+    "ActionTypes",
+    "Action",
+    "DoNothing",
+    "RunDeployment",
+    "PauseDeployment",
+    "ResumeDeployment",
+    "ChangeFlowRunState",
+    "CancelFlowRun",
+    "SuspendFlowRun",
+    "CallWebhook",
+    "SendNotification",
+    "PauseWorkPool",
+    "ResumeWorkPool",
+    "PauseWorkQueue",
+    "ResumeWorkQueue",
+    "PauseAutomation",
+    "ResumeAutomation",
+    "DeclareIncident",
     "emit_event",
+    "get_events_client",
+    "get_events_subscriber",
 ]

prefect/events/actions.py

@@ -1,36 +1,71 @@
+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  # type: ignore
 
 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 +73,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"
+    )
+    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"
     )
-    body: str = Field(..., description="Notification body")
-    subject: Optional[str] = Field(None, description="Notification subject")
+
+    @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,201 @@
+"""
+Command line interface for working with automations.
+"""
+
+import functools
+from typing import Optional
+
+import orjson
+import typer
+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(
+    name: Optional[str] = typer.Argument(None, help="An automation's name"),
+    id: Optional[str] = typer.Option(None, "--id", help="An automation's id"),
+):
+    """Delete an automation.
+
+    Arguments:
+        name: the name of the automation to delete
+        id: the id of the automation to delete
+
+    Examples:
+        $ prefect automation delete "my-automation"
+        $ prefect automation delete --id "aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa"
+    """
+
+    async with get_client() as client:
+        if not id and not name:
+            exit_with_error("Please provide either a name or an id.")
+
+        if id:
+            automation = await client.read_automation(id)
+            if not automation:
+                exit_with_error(f"Automation with id {id!r} not found.")
+            if not typer.confirm(
+                (f"Are you sure you want to delete automation with id {id!r}?"),
+                default=False,
+            ):
+                exit_with_error("Deletion aborted.")
+            await client.delete_automation(id)
+            exit_with_success(f"Deleted automation with id {id!r}")
+
+        elif name:
+            automation = await client.read_automations_by_name(name=name)
+            if not automation:
+                exit_with_error(
+                    f"Automation {name!r} not found. You can also specify an id with the `--id` flag."
+                )
+            elif len(automation) > 1:
+                exit_with_error(
+                    f"Multiple automations found with name {name!r}. Please specify an id with the `--id` flag instead."
+                )
+            if not typer.confirm(
+                (f"Are you sure you want to delete automation with name {name!r}?"),
+                default=False,
+            ):
+                exit_with_error("Deletion aborted.")
+            await client.delete_automation(automation[0].id)
+            exit_with_success(f"Deleted automation with name {name!r}")