Skip to content

Commit

Permalink
rm sync_compatible from automations SDK (#16584)
Browse files Browse the repository at this point in the history
Co-authored-by: Adam Azzam <[email protected]>
  • Loading branch information
zzstoatzz and aaazzam authored Jan 3, 2025
1 parent cecd3a1 commit 5fb0a5a
Show file tree
Hide file tree
Showing 5 changed files with 353 additions and 110 deletions.
266 changes: 236 additions & 30 deletions src/prefect/automations.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
from pydantic import Field
from typing_extensions import Self

from prefect._internal.compatibility.async_dispatch import async_dispatch
from prefect.client.orchestration import get_client
from prefect.events.actions import (
CallWebhook,
Expand Down Expand Up @@ -39,7 +40,6 @@
Trigger,
)
from prefect.exceptions import PrefectHTTPStatusError
from prefect.utilities.asyncutils import sync_compatible

__all__ = [
"AutomationCore",
Expand Down Expand Up @@ -78,11 +78,13 @@
class Automation(AutomationCore):
id: Optional[UUID] = Field(default=None, description="The ID of this automation")

@sync_compatible
async def create(self: Self) -> Self:
async def acreate(self: Self) -> Self:
"""
Create a new automation.
Asynchronously create a new automation.
Examples:
```python
auto_to_create = Automation(
name="woodchonk",
trigger=EventTrigger(
Expand All @@ -97,20 +99,55 @@ async def create(self: Self) -> Self:
),
actions=[CancelFlowRun()]
)
created_automation = auto_to_create.create()
created_automation = await auto_to_create.acreate()
```
"""
async with get_client() as client:
automation = AutomationCore(**self.model_dump(exclude={"id"}))
self.id = await client.create_automation(automation=automation)
return self

@sync_compatible
async def update(self: Self):
@async_dispatch(acreate)
def create(self: Self) -> Self:
"""
Create a new automation.
Examples:
```python
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()
```
"""
with get_client(sync_client=True) as client:
automation = AutomationCore(**self.model_dump(exclude={"id"}))
self.id = client.create_automation(automation=automation)
return self

async def aupdate(self: Self):
"""
Updates an existing automation.
Examples:
```python
auto = Automation.read(id=123)
auto.name = "new name"
auto.update()
```
"""
assert self.id is not None
async with get_client() as client:
Expand All @@ -119,28 +156,51 @@ async def update(self: Self):
)
await client.update_automation(automation_id=self.id, automation=automation)

@async_dispatch(aupdate)
def update(self: Self):
"""
Updates an existing automation.
Examples:
```python
auto = Automation.read(id=123)
auto.name = "new name"
auto.update()
```
"""
assert self.id is not None
with get_client(sync_client=True) as client:
automation = AutomationCore(
**self.model_dump(exclude={"id", "owner_resource"})
)
client.update_automation(automation_id=self.id, automation=automation)

@overload
@classmethod
async def read(cls, id: UUID, name: Optional[str] = ...) -> Self:
async def aread(cls, id: UUID, name: Optional[str] = ...) -> Self:
...

@overload
@classmethod
async def read(cls, id: None = None, name: str = ...) -> Optional[Self]:
async def aread(cls, id: None = None, name: str = ...) -> Self:
...

@classmethod
@sync_compatible
async def read(
cls, id: Optional[UUID] = None, name: Optional[str] = None
) -> Optional[Self]:
async def aread(cls, id: Optional[UUID] = None, name: Optional[str] = None) -> Self:
"""
Read an automation by ID or name.
automation = Automation.read(name="woodchonk")
Asynchronously read an automation by ID or name.
or
Examples:
automation = Automation.read(id=UUID("b3514963-02b1-47a5-93d1-6eeb131041cb"))
```python
automation = await Automation.aread(name="woodchonk")
```
```python
automation = await Automation.aread(id=UUID("b3514963-02b1-47a5-93d1-6eeb131041cb"))
```
"""
if id and name:
raise ValueError("Only one of id or name can be provided")
Expand All @@ -162,15 +222,68 @@ async def read(
assert name is not None
automation = await client.read_automations_by_name(name=name)
if len(automation) > 0:
return cls(**automation[0].model_dump()) if automation else None
else:
raise ValueError(f"Automation with name {name!r} not found")
return cls(**automation[0].model_dump())
raise ValueError(f"Automation with name {name!r} not found")

@sync_compatible
async def delete(self: Self) -> bool:
@overload
@classmethod
async def read(cls, id: UUID, name: Optional[str] = ...) -> Self:
...

@overload
@classmethod
async def read(cls, id: None = None, name: str = ...) -> Self:
...

@classmethod
@async_dispatch(aread)
def read(cls, id: Optional[UUID] = None, name: Optional[str] = None) -> Self:
"""
Read an automation by ID or name.
Examples:
```python
automation = Automation.read(name="woodchonk")
```
```python
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")
with get_client(sync_client=True) as client:
if id:
try:
automation = 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")
raise
if automation is None:
raise ValueError(f"Automation with ID {id!r} not found")
return cls(**automation.model_dump())
else:
if TYPE_CHECKING:
assert name is not None
automation = client.read_automations_by_name(name=name)
if len(automation) > 0:
return cls(**automation[0].model_dump())
raise ValueError(f"Automation with name {name!r} not found")

async def adelete(self: Self) -> bool:
"""
Asynchronously delete an automation.
Examples:
```python
auto = Automation.read(id = 123)
auto.delete()
await auto.adelete()
```
"""
if self.id is None:
raise ValueError("Can't delete an automation without an id")
Expand All @@ -184,38 +297,131 @@ async def delete(self: Self) -> bool:
return False
raise

@sync_compatible
async def disable(self: Self) -> bool:
@async_dispatch(adelete)
def delete(self: Self) -> bool:
"""
Delete an automation.
Examples:
```python
auto = Automation.read(id = 123)
auto.delete()
```
"""
if self.id is None:
raise ValueError("Can't delete an automation without an id")

with get_client(sync_client=True) as client:
try:
client.delete_automation(self.id)
return True
except PrefectHTTPStatusError as exc:
if exc.response.status_code == 404:
return False
raise

async def adisable(self: Self) -> bool:
"""
Asynchronously disable an automation.
Raises:
ValueError: If the automation does not have an id
PrefectHTTPStatusError: If the automation cannot be disabled
Example:
```python
auto = await Automation.aread(id = 123)
await auto.adisable()
```
"""
if self.id is None:
raise ValueError("Can't disable an automation without an id")

async with get_client() as client:
try:
await client.pause_automation(self.id)
return True
except PrefectHTTPStatusError as exc:
if exc.response.status_code == 404:
return False
raise

@async_dispatch(adisable)
def disable(self: Self) -> bool:
"""
Disable an automation.
Raises:
ValueError: If the automation does not have an id
PrefectHTTPStatusError: If the automation cannot be disabled
Example:
```python
auto = Automation.read(id = 123)
auto.disable()
```
"""
if self.id is None:
raise ValueError("Can't disable an automation without an id")

with get_client(sync_client=True) as client:
try:
client.pause_automation(self.id)
return True
except PrefectHTTPStatusError as exc:
if exc.response.status_code == 404:
return False
raise

async def aenable(self: Self) -> bool:
"""
Asynchronously enable an automation.
Raises:
ValueError: If the automation does not have an id
PrefectHTTPStatusError: If the automation cannot be enabled
Example:
```python
auto = await Automation.aread(id = 123)
await auto.aenable()
```
"""
if self.id is None:
raise ValueError("Can't enable an automation without an id")

async with get_client() as client:
try:
await client.pause_automation(self.id)
await client.resume_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:
@async_dispatch(aenable)
def enable(self: Self) -> bool:
"""
Enable an automation.
Raises:
ValueError: If the automation does not have an id
PrefectHTTPStatusError: If the automation cannot be enabled
Example:
```python
auto = Automation.read(id = 123)
auto.enable()
```
"""
if self.id is None:
raise ValueError("Can't enable an automation without an id")

async with get_client() as client:
with get_client(sync_client=True) as client:
try:
await client.resume_automation(self.id)
client.resume_automation(self.id)
return True
except PrefectHTTPStatusError as exc:
if exc.response.status_code == 404:
Expand Down
Loading

0 comments on commit 5fb0a5a

Please sign in to comment.