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

prefect/automations.py

@@ -0,0 +1,162 @@
+from typing import Optional
+from uuid import UUID
+
+from pydantic import Field
+from typing_extensions import Self
+
+from prefect.client.utilities import get_or_create_client
+from prefect.events.schemas.automations import (
+    AutomationCore,
+    CompositeTrigger,
+    CompoundTrigger,
+    EventTrigger,
+    MetricTrigger,
+    MetricTriggerOperator,
+    MetricTriggerQuery,
+    Posture,
+    PrefectMetric,
+    ResourceSpecification,
+    ResourceTrigger,
+    SequenceTrigger,
+    Trigger,
+)
+from prefect.exceptions import PrefectHTTPStatusError
+from prefect.utilities.asyncutils import sync_compatible
+
+__all__ = [
+    "AutomationCore",
+    "EventTrigger",
+    "ResourceTrigger",
+    "Posture",
+    "Trigger",
+    "ResourceSpecification",
+    "MetricTriggerOperator",
+    "MetricTrigger",
+    "PrefectMetric",
+    "CompositeTrigger",
+    "SequenceTrigger",
+    "CompoundTrigger",
+    "MetricTriggerQuery",
+]
+
+
+class Automation(AutomationCore):
+    id: Optional[UUID] = Field(default=None, description="The ID of this automation")
+
+    @sync_compatible
+    async def create(self: Self) -> Self:
+        """
+        Create a new automation.
+
+        auto_to_create = Automation(
+            name="woodchonk",
+            trigger=EventTrigger(
+                expect={"animal.walked"},
+                match={
+                    "genus": "Marmota",
+                    "species": "monax",
+                },
+                posture="Reactive",
+                threshold=3,
+                within=timedelta(seconds=10),
+            ),
+            actions=[CancelFlowRun()]
+        )
+        created_automation = auto_to_create.create()
+        """
+        client, _ = get_or_create_client()
+        automation = AutomationCore(**self.dict(exclude={"id"}))
+        self.id = await client.create_automation(automation=automation)
+        return self
+
+    @sync_compatible
+    async def update(self: Self):
+        """
+        Updates an existing automation.
+        auto = Automation.read(id=123)
+        auto.name = "new name"
+        auto.update()
+        """
+
+        client, _ = get_or_create_client()
+        automation = AutomationCore(**self.dict(exclude={"id", "owner_resource"}))
+        await client.update_automation(automation_id=self.id, automation=automation)
+
+    @classmethod
+    @sync_compatible
+    async def read(
+        cls: Self, id: Optional[UUID] = None, name: Optional[str] = None
+    ) -> Self:
+        """
+        Read an automation by ID or name.
+        automation = Automation.read(name="woodchonk")
+
+        or
+
+        automation = Automation.read(id=UUID("b3514963-02b1-47a5-93d1-6eeb131041cb"))
+        """
+        if id and name:
+            raise ValueError("Only one of id or name can be provided")
+        if not id and not name:
+            raise ValueError("One of id or name must be provided")
+        client, _ = get_or_create_client()
+        if id:
+            try:
+                automation = await client.read_automation(automation_id=id)
+            except PrefectHTTPStatusError as exc:
+                if exc.response.status_code == 404:
+                    raise ValueError(f"Automation with ID {id!r} not found")
+            return Automation(**automation.dict())
+        else:
+            automation = await client.read_automations_by_name(name=name)
+            if len(automation) > 0:
+                return Automation(**automation[0].dict()) if automation else None
+            else:
+                raise ValueError(f"Automation with name {name!r} not found")
+
+    @sync_compatible
+    async def delete(self: Self) -> bool:
+        """
+        auto = Automation.read(id = 123)
+        auto.delete()
+        """
+        try:
+            client, _ = get_or_create_client()
+            await client.delete_automation(self.id)
+            return True
+        except PrefectHTTPStatusError as exc:
+            if exc.response.status_code == 404:
+                return False
+            raise
+
+    @sync_compatible
+    async def disable(self: Self) -> bool:
+        """
+        Disable an automation.
+        auto = Automation.read(id = 123)
+        auto.disable()
+        """
+        try:
+            client, _ = get_or_create_client()
+            await client.pause_automation(self.id)
+            return True
+        except PrefectHTTPStatusError as exc:
+            if exc.response.status_code == 404:
+                return False
+            raise
+
+    @sync_compatible
+    async def enable(self: Self) -> bool:
+        """
+        Enable an automation.
+        auto = Automation.read(id = 123)
+        auto.enable()
+        """
+        try:
+            client, _ = get_or_create_client()
+            await client.resume_automation("asd")
+            return True
+        except PrefectHTTPStatusError as exc:
+            if exc.response.status_code == 404:
+                return False
+            raise

prefect/client/orchestration.py

@@ -623,12 +623,12 @@ class PrefectClient:
     async def create_flow_run(
         self,
         flow: "FlowObject",
-        name: str = None,
+        name: Optional[str] = None,
         parameters: Optional[Dict[str, Any]] = None,
         context: Optional[Dict[str, Any]] = None,
-        tags: Iterable[str] = None,
-        parent_task_run_id: UUID = None,
-        state: "prefect.states.State" = None,
+        tags: Optional[Iterable[str]] = None,
+        parent_task_run_id: Optional[UUID] = None,
+        state: Optional["prefect.states.State"] = None,
     ) -> FlowRun:
         """
         Create a flow run for a flow.
@@ -3161,6 +3161,16 @@ class PrefectClient:
 
         return UUID(response.json()["id"])
 
+    async def update_automation(self, automation_id: UUID, automation: AutomationCore):
+        """Updates an automation in Prefect Cloud."""
+        if not self.server_type.supports_automations():
+            self._raise_for_unsupported_automations()
+        response = await self._client.put(
+            f"/automations/{automation_id}",
+            json=automation.dict(json_compatible=True, exclude_unset=True),
+        )
+        response.raise_for_status
+
     async def read_automations(self) -> List[Automation]:
         if not self.server_type.supports_automations():
             self._raise_for_unsupported_automations()
@@ -3169,16 +3179,24 @@ class PrefectClient:
         response.raise_for_status()
         return pydantic.parse_obj_as(List[Automation], response.json())
 
-    async def find_automation(self, id_or_name: str) -> Optional[Automation]:
-        try:
-            id = UUID(id_or_name)
-        except ValueError:
-            id = None
+    async def find_automation(
+        self, id_or_name: Union[str, UUID], exit_if_not_found: bool = True
+    ) -> Optional[Automation]:
+        if isinstance(id_or_name, str):
+            try:
+                id = UUID(id_or_name)
+            except ValueError:
+                id = None
+        elif isinstance(id_or_name, UUID):
+            id = id_or_name
 
         if id:
-            automation = await self.read_automation(id)
-            if automation:
+            try:
+                automation = await self.read_automation(id)
                 return automation
+            except prefect.exceptions.HTTPStatusError as e:
+                if e.response.status_code == status.HTTP_404_NOT_FOUND:
+                    raise prefect.exceptions.ObjectNotFound(http_exc=e) from e
 
         automations = await self.read_automations()
 

prefect/client/schemas/objects.py

@@ -53,6 +53,7 @@ if TYPE_CHECKING:
     from prefect.deprecated.data_documents import DataDocument
     from prefect.results import BaseResult
 
+
 R = TypeVar("R")
 
 
@@ -120,20 +121,20 @@ class WorkQueueStatus(AutoEnum):
 
 
 class StateDetails(PrefectBaseModel):
-    flow_run_id: UUID = None
-    task_run_id: UUID = None
+    flow_run_id: Optional[UUID] = None
+    task_run_id: Optional[UUID] = None
     # for task runs that represent subflows, the subflow's run ID
-    child_flow_run_id: UUID = None
+    child_flow_run_id: Optional[UUID] = None
     scheduled_time: DateTimeTZ = None
-    cache_key: str = None
+    cache_key: Optional[str] = None
     cache_expiration: DateTimeTZ = None
     untrackable_result: bool = False
     pause_timeout: DateTimeTZ = None
     pause_reschedule: bool = False
-    pause_key: str = None
+    pause_key: Optional[str] = None
     run_input_keyset: Optional[Dict[str, str]] = None
-    refresh_cache: bool = None
-    retriable: bool = None
+    refresh_cache: Optional[bool] = None
+    retriable: Optional[bool] = None
     transition_id: Optional[UUID] = None
     task_parameters_id: Optional[UUID] = None
 
@@ -160,7 +161,9 @@ class State(ObjectBaseModel, Generic[R]):
     def result(self: "State[R]", raise_on_failure: bool = False) -> Union[R, Exception]:
         ...
 
-    def result(self, raise_on_failure: bool = True, fetch: Optional[bool] = None):
+    def result(
+        self, raise_on_failure: bool = True, fetch: Optional[bool] = None
+    ) -> Union[R, Exception]:
         """
         Retrieve the result attached to this state.
 

prefect/events/cli/automations.py

@@ -4,6 +4,7 @@ Command line interface for working with automations.
 
 import functools
 from typing import Optional
+from uuid import UUID
 
 import orjson
 import typer
@@ -16,6 +17,8 @@ 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
+from prefect.events.schemas.automations import Automation
+from prefect.exceptions import PrefectHTTPStatusError
 
 automations_app = PrefectTyper(
     name="automation",
@@ -98,55 +101,175 @@ async def ls():
 
 @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.")
+async def inspect(
+    name: Optional[str] = typer.Argument(None, help="An automation's name"),
+    id: Optional[str] = typer.Option(None, "--id", help="An automation's id"),
+    yaml: bool = typer.Option(False, "--yaml", help="Output as YAML"),
+    json: bool = typer.Option(False, "--json", help="Output as JSON"),
+):
+    """
+    Inspect an automation.
 
-    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()
-        )
+    Arguments:
+
+        name: the name of the automation to inspect
+
+        id: the id of the automation to inspect
+
+        yaml: output as YAML
+
+        json: output as JSON
+
+    Examples:
+
+        $ prefect automation inspect "my-automation"
+
+        $ prefect automation inspect --id "aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa"
+
+        $ prefect automation inspect "my-automation" --yaml
+
+        $ prefect automation inspect "my-automation" --json
+    """
+    if not id and not name:
+        exit_with_error("Please provide either a name or an id.")
+
+    if name:
+        async with get_client() as client:
+            automation = await client.read_automations_by_name(name=name)
+            if not automation:
+                exit_with_error(f"Automation {name!r} not found.")
+
+    elif id:
+        async with get_client() as client:
+            try:
+                uuid_id = UUID(id)
+                automation = await client.read_automation(uuid_id)
+            except (PrefectHTTPStatusError, ValueError):
+                exit_with_error(f"Automation with id {id!r} not found.")
+
+    if yaml or json:
+        if isinstance(automation, list):
+            automation = [a.dict(json_compatible=True) for a in automation]
+        elif isinstance(automation, Automation):
+            automation = automation.dict(json_compatible=True)
+        if yaml:
+            app.console.print(pyyaml.dump(automation, sort_keys=False))
+        elif json:
+            app.console.print(
+                orjson.dumps(automation, 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 def resume(
+    name: Optional[str] = typer.Argument(None, help="An automation's name"),
+    id: Optional[str] = typer.Option(None, "--id", help="An automation's id"),
+):
+    """
+    Resume an automation.
 
-    async with get_client() as client:
-        await client.resume_automation(automation.id)
+    Arguments:
 
-    exit_with_success(f"Resumed automation {automation.name!r} ({automation.id})")
+            name: the name of the automation to resume
+
+            id: the id of the automation to resume
+
+    Examples:
+
+            $ prefect automation resume "my-automation"
+
+            $ prefect automation resume --id "aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa"
+    """
+    if not id and not name:
+        exit_with_error("Please provide either a name or an id.")
+
+    if name:
+        async with get_client() as client:
+            automation = await client.read_automations_by_name(name=name)
+            if not automation:
+                exit_with_error(
+                    f"Automation with name {name!r} not found. You can also specify an id with the `--id` flag."
+                )
+            if len(automation) > 1:
+                if not typer.confirm(
+                    f"Multiple automations found with name {name!r}. Do you want to resume all of them?",
+                    default=False,
+                ):
+                    exit_with_error("Resume aborted.")
+
+            for a in automation:
+                await client.resume_automation(a.id)
+            exit_with_success(
+                f"Resumed automation(s) with name {name!r} and id(s) {', '.join([repr(str(a.id)) for a in automation])}."
+            )
+
+    elif id:
+        async with get_client() as client:
+            try:
+                uuid_id = UUID(id)
+                automation = await client.read_automation(uuid_id)
+            except (PrefectHTTPStatusError, ValueError):
+                exit_with_error(f"Automation with id {id!r} not found.")
+            await client.resume_automation(automation.id)
+            exit_with_success(f"Resumed automation with id {str(automation.id)!r}.")
 
 
 @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 def pause(
+    name: Optional[str] = typer.Argument(None, help="An automation's name"),
+    id: Optional[str] = typer.Option(None, "--id", help="An automation's id"),
+):
+    """
+    Pause an automation.
 
-    async with get_client() as client:
-        await client.pause_automation(automation.id)
+    Arguments:
+
+            name: the name of the automation to pause
+
+            id: the id of the automation to pause
+
+    Examples:
+
+        $ prefect automation pause "my-automation"
+
+        $ prefect automation pause --id "aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa"
+    """
+    if not id and not name:
+        exit_with_error("Please provide either a name or an id.")
+
+    if name:
+        async with get_client() as client:
+            automation = await client.read_automations_by_name(name=name)
+            if not automation:
+                exit_with_error(
+                    f"Automation with name {name!r} not found. You can also specify an id with the `--id` flag."
+                )
+            if len(automation) > 1:
+                if not typer.confirm(
+                    f"Multiple automations found with name {name!r}. Do you want to pause all of them?",
+                    default=False,
+                ):
+                    exit_with_error("Pause aborted.")
+
+            for a in automation:
+                await client.pause_automation(a.id)
+            exit_with_success(
+                f"Paused automation(s) with name {name!r} and id(s) {', '.join([repr(str(a.id)) for a in automation])}."
+            )
 
-    exit_with_success(f"Paused automation {automation.name!r} ({automation.id})")
+    elif id:
+        async with get_client() as client:
+            try:
+                uuid_id = UUID(id)
+                automation = await client.read_automation(uuid_id)
+            except (PrefectHTTPStatusError, ValueError):
+                exit_with_error(f"Automation with id {id!r} not found.")
+            await client.pause_automation(automation.id)
+            exit_with_success(f"Paused automation with id {str(automation.id)!r}.")
 
 
 @automations_app.command()

prefect/events/clients.py

@@ -11,6 +11,7 @@ from typing import (
     Optional,
     Tuple,
     Type,
+    cast,
 )
 from uuid import UUID
 
@@ -377,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:
 
@@ -412,7 +413,7 @@ 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

prefect/events/filters.py

@@ -43,7 +43,7 @@ class AutomationFilter(PrefectBaseModel):
     )
 
 
-class EventDataFilter(PrefectBaseModel, extra="forbid"):
+class EventDataFilter(PrefectBaseModel, extra="forbid"):  # type: ignore[call-arg]
     """A base class for filtering event data."""
 
     _top_level_filter: "EventFilter | None" = PrivateAttr(None)

prefect/events/schemas/automations.py

@@ -39,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.
@@ -391,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"""
 

prefect/events/schemas/deployment_triggers.py

@@ -48,7 +48,7 @@ from .automations import (
 from .events import ResourceSpecification
 
 
-class BaseDeploymentTrigger(PrefectBaseModel, abc.ABC, extra="ignore"):
+class BaseDeploymentTrigger(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.

prefect/events/schemas/events.py

@@ -192,15 +192,22 @@ class ReceivedEvent(Event):
 
 def matches(expected: str, value: Optional[str]) -> bool:
     """Returns true if the given value matches the expected string, which may
-    include wildcards"""
+    include a a negation prefix ("!this-value") or a wildcard suffix
+    ("any-value-starting-with*")"""
     if value is None:
         return False
 
-    # TODO: handle wildcards/globs better than this
+    positive = True
+    if expected.startswith("!"):
+        expected = expected[1:]
+        positive = False
+
     if expected.endswith("*"):
-        return value.startswith(expected[:-1])
+        match = value.startswith(expected[:-1])
+    else:
+        match = value == expected
 
-    return value == expected
+    return match if positive else not match
 
 
 class ResourceSpecification(PrefectBaseModel):

prefect/events/schemas/labelling.py

@@ -71,7 +71,7 @@ class LabelDiver:
             raise AttributeError
 
 
-class Labelled(PrefectBaseModel, extra="ignore"):
+class Labelled(PrefectBaseModel, extra="ignore"):  # type: ignore[call-arg]
     """An object defined by string labels and values"""
 
     __root__: Dict[str, str]

prefect/flows.py

@@ -733,7 +733,7 @@ class Flow(Generic[P, R]):
     @sync_compatible
     async def serve(
         self,
-        name: str,
+        name: Optional[str] = None,
         interval: Optional[
             Union[
                 Iterable[Union[int, float, datetime.timedelta]],
@@ -764,7 +764,7 @@ class Flow(Generic[P, R]):
         Creates a deployment for this flow and starts a runner to monitor for scheduled work.
 
         Args:
-            name: The name to give the created deployment.
+            name: The name to give the created deployment. Defaults to the name of the flow.
             interval: An interval on which to execute the deployment. Accepts a number or a
                 timedelta object to create a single schedule. If a number is given, it will be
                 interpreted as seconds. Also accepts an iterable of numbers or timedelta to create
@@ -827,10 +827,13 @@ class Flow(Generic[P, R]):
         """
         from prefect.runner import Runner
 
-        # Handling for my_flow.serve(__file__)
-        # Will set name to name of file where my_flow.serve() without the extension
-        # Non filepath strings will pass through unchanged
-        name = Path(name).stem
+        if not name:
+            name = self.name
+        else:
+            # Handling for my_flow.serve(__file__)
+            # Will set name to name of file where my_flow.serve() without the extension
+            # Non filepath strings will pass through unchanged
+            name = Path(name).stem
 
         runner = Runner(name=name, pause_on_shutdown=pause_on_shutdown, limit=limit)
         deployment_id = await runner.add_flow(
@@ -1226,19 +1229,19 @@ class Flow(Generic[P, R]):
             return track_viz_task(self.isasync, self.name, parameters)
 
         if PREFECT_EXPERIMENTAL_ENABLE_NEW_ENGINE.value():
-            from prefect.new_flow_engine import run_flow
-            from prefect.utilities.asyncutils import run_sync
+            from prefect.new_flow_engine import run_flow, run_flow_sync
 
-            awaitable = run_flow(
+            run_kwargs = dict(
                 flow=self,
                 parameters=parameters,
                 wait_for=wait_for,
                 return_type=return_type,
             )
             if self.isasync:
-                return awaitable
+                # this returns an awaitable coroutine
+                return run_flow(**run_kwargs)
             else:
-                return run_sync(awaitable)
+                return run_flow_sync(**run_kwargs)
 
         return enter_flow_run_engine_from_flow_call(
             self,

prefect/input/run_input.py

@@ -582,7 +582,9 @@ def receive_input(
     # the signature is the same as here:
     #   Union[Type[R], Type[T], pydantic.BaseModel],
     # Seems like a possible mypy bug, so we'll ignore the type check here.
-    input_cls = run_input_subclass_from_type(input_type)  # type: ignore[arg-type]
+    input_cls: Union[
+        Type[AutomaticRunInput[T]], Type[R]
+    ] = run_input_subclass_from_type(input_type)  # type: ignore[arg-type]
 
     if issubclass(input_cls, AutomaticRunInput):
         return input_cls.receive(