prefect-client 2.16.6__py3-none-any.whl → 2.16.8__py3-none-any.whl

prefect/artifacts.py

@@ -2,21 +2,173 @@
 Interface for creating and reading artifacts.
 """
 
-import json
+from __future__ import annotations
+
+import json  # noqa: I001
 import math
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, Dict, List, Optional, Tuple, Union
 from uuid import UUID
 
+from typing_extensions import Self
+
 from prefect.client.orchestration import PrefectClient
-from prefect.client.schemas.actions import ArtifactCreate
-from prefect.client.utilities import inject_client
-from prefect.context import FlowRunContext, TaskRunContext
+from prefect.client.schemas.actions import ArtifactCreate as ArtifactRequest
+from prefect.client.schemas.filters import ArtifactFilter, ArtifactFilterKey
+from prefect.client.schemas.objects import Artifact as ArtifactResponse
+from prefect.client.schemas.sorting import ArtifactSort
+from prefect.client.utilities import get_or_create_client, inject_client
 from prefect.utilities.asyncutils import sync_compatible
+from prefect.utilities.context import get_task_and_flow_run_ids
+
+
+class Artifact(ArtifactRequest):
+    """
+    An artifact is a piece of data that is created by a flow or task run.
+    https://docs.prefect.io/latest/concepts/artifacts/
+
+    Arguments:
+        type: A string identifying the type of artifact.
+        key: A user-provided string identifier.
+          The key must only contain lowercase letters, numbers, and dashes.
+        description: A user-specified description of the artifact.
+        data: A JSON payload that allows for a result to be retrieved.
+    """
+
+    @sync_compatible
+    async def create(
+        self: Self,
+        client: Optional[PrefectClient] = None,
+    ) -> ArtifactResponse:
+        """
+        A method to create an artifact.
+
+        Arguments:
+            client: The PrefectClient
+
+        Returns:
+            - The created artifact.
+        """
+        client, _ = get_or_create_client(client)
+        task_run_id, flow_run_id = get_task_and_flow_run_ids()
+        return await client.create_artifact(
+            artifact=ArtifactRequest(
+                type=self.type,
+                key=self.key,
+                description=self.description,
+                task_run_id=self.task_run_id or task_run_id,
+                flow_run_id=self.flow_run_id or flow_run_id,
+                data=await self.format(),
+            )
+        )
+
+    @classmethod
+    @sync_compatible
+    async def get(
+        cls, key: Optional[str] = None, client: Optional[PrefectClient] = None
+    ) -> Optional[ArtifactResponse]:
+        """
+        A method to get an artifact.
+
+        Arguments:
+            key (str, optional): The key of the artifact to get.
+            client (PrefectClient, optional): The PrefectClient
+
+        Returns:
+            (ArtifactResponse, optional): The artifact (if found).
+        """
+        client, _ = get_or_create_client(client)
+        return next(
+            iter(
+                await client.read_artifacts(
+                    limit=1,
+                    sort=ArtifactSort.UPDATED_DESC,
+                    artifact_filter=ArtifactFilter(key=ArtifactFilterKey(any_=[key])),
+                )
+            ),
+            None,
+        )
+
+    @classmethod
+    @sync_compatible
+    async def get_or_create(
+        cls,
+        key: Optional[str] = None,
+        description: Optional[str] = None,
+        data: Optional[Union[Dict[str, Any], Any]] = None,
+        client: Optional[PrefectClient] = None,
+        **kwargs: Any,
+    ) -> Tuple[ArtifactResponse, bool]:
+        """
+        A method to get or create an artifact.
+
+        Arguments:
+            key (str, optional): The key of the artifact to get or create.
+            description (str, optional): The description of the artifact to create.
+            data (Union[Dict[str, Any], Any], optional): The data of the artifact to create.
+            client (PrefectClient, optional): The PrefectClient
+
+        Returns:
+            (ArtifactResponse): The artifact, either retrieved or created.
+        """
+        artifact = await cls.get(key, client)
+        if artifact:
+            return artifact, False
+        else:
+            return (
+                await cls(key=key, description=description, data=data, **kwargs).create(
+                    client
+                ),
+                True,
+            )
+
+    async def format(self) -> Optional[Union[Dict[str, Any], Any]]:
+        return json.dumps(self.data)
+
+
+class LinkArtifact(Artifact):
+    link: str
+    link_text: Optional[str] = None
+    type: Optional[str] = "markdown"
+
+    async def format(self) -> str:
+        return (
+            f"[{self.link_text}]({self.link})"
+            if self.link_text
+            else f"[{self.link}]({self.link})"
+        )
+
+
+class MarkdownArtifact(Artifact):
+    markdown: str
+    type: Optional[str] = "markdown"
+
+    async def format(self) -> str:
+        return self.markdown
 
-INVALID_TABLE_TYPE_ERROR = (
-    "`create_table_artifact` requires a `table` argument of type `dict[list]` or"
-    " `list[dict]`."
-)
+
+class TableArtifact(Artifact):
+    table: Union[Dict[str, List[Any]], List[Dict[str, Any]], List[List[Any]]]
+    type: Optional[str] = "table"
+
+    @classmethod
+    def _sanitize(
+        cls, item: Union[Dict[str, Any], List[Any], float]
+    ) -> Union[Dict[str, Any], List[Any], int, float, None]:
+        """
+        Sanitize NaN values in a given item.
+        The item can be a dict, list or float.
+        """
+        if isinstance(item, list):
+            return [cls._sanitize(sub_item) for sub_item in item]
+        elif isinstance(item, dict):
+            return {k: cls._sanitize(v) for k, v in item.items()}
+        elif isinstance(item, float) and math.isnan(item):
+            return None
+        else:
+            return item
+
+    async def format(self) -> str:
+        return json.dumps(self._sanitize(self.table))
 
 
 @inject_client
@@ -41,28 +193,15 @@ async def _create_artifact(
     Returns:
         - The table artifact ID.
     """
-    artifact_args = {}
-    task_run_ctx = TaskRunContext.get()
-    flow_run_ctx = FlowRunContext.get()
-
-    if task_run_ctx:
-        artifact_args["task_run_id"] = task_run_ctx.task_run.id
-        artifact_args["flow_run_id"] = task_run_ctx.task_run.flow_run_id
-    elif flow_run_ctx:
-        artifact_args["flow_run_id"] = flow_run_ctx.flow_run.id
-
-    if key is not None:
-        artifact_args["key"] = key
-    if type is not None:
-        artifact_args["type"] = type
-    if description is not None:
-        artifact_args["description"] = description
-    if data is not None:
-        artifact_args["data"] = data
 
-    artifact = ArtifactCreate(**artifact_args)
+    artifact = await Artifact(
+        type=type,
+        key=key,
+        description=description,
+        data=data,
+    ).create(client)
 
-    return await client.create_artifact(artifact=artifact)
+    return artifact.id
 
 
 @sync_compatible
@@ -71,6 +210,7 @@ async def create_link_artifact(
     link_text: Optional[str] = None,
     key: Optional[str] = None,
     description: Optional[str] = None,
+    client: Optional[PrefectClient] = None,
 ) -> UUID:
     """
     Create a link artifact.
@@ -87,13 +227,12 @@ async def create_link_artifact(
     Returns:
         The table artifact ID.
     """
-    formatted_link = f"[{link_text}]({link})" if link_text else f"[{link}]({link})"
-    artifact = await _create_artifact(
+    artifact = await LinkArtifact(
         key=key,
-        type="markdown",
         description=description,
-        data=formatted_link,
-    )
+        link=link,
+        link_text=link_text,
+    ).create(client)
 
     return artifact.id
 
@@ -117,12 +256,11 @@ async def create_markdown_artifact(
     Returns:
         The table artifact ID.
     """
-    artifact = await _create_artifact(
+    artifact = await MarkdownArtifact(
         key=key,
-        type="markdown",
         description=description,
-        data=markdown,
-    )
+        markdown=markdown,
+    ).create()
 
     return artifact.id
 
@@ -147,39 +285,10 @@ async def create_table_artifact(
         The table artifact ID.
     """
 
-    def _sanitize_nan_values(item):
-        """
-        Sanitize NaN values in a given item. The item can be a dict, list or float.
-        """
-
-        if isinstance(item, list):
-            return [_sanitize_nan_values(sub_item) for sub_item in item]
-
-        elif isinstance(item, dict):
-            return {k: _sanitize_nan_values(v) for k, v in item.items()}
-
-        elif isinstance(item, float) and math.isnan(item):
-            return None
-
-        else:
-            return item
-
-    sanitized_table = _sanitize_nan_values(table)
-
-    if isinstance(table, dict) and all(isinstance(v, list) for v in table.values()):
-        pass
-    elif isinstance(table, list) and all(isinstance(v, (list, dict)) for v in table):
-        pass
-    else:
-        raise TypeError(INVALID_TABLE_TYPE_ERROR)
-
-    formatted_table = json.dumps(sanitized_table)
-
-    artifact = await _create_artifact(
+    artifact = await TableArtifact(
         key=key,
-        type="table",
         description=description,
-        data=formatted_table,
-    )
+        table=table,
+    ).create()
 
     return artifact.id

prefect/client/orchestration.py

@@ -128,7 +128,7 @@ from prefect.client.schemas.sorting import (
     TaskRunSort,
 )
 from prefect.deprecated.data_documents import DataDocument
-from prefect.events.schemas import Automation, ExistingAutomation
+from prefect.events.schemas.automations import Automation, ExistingAutomation
 from prefect.logging import get_logger
 from prefect.settings import (
     PREFECT_API_DATABASE_CONNECTION_URL,

prefect/client/schemas/actions.py

@@ -652,7 +652,7 @@ class GlobalConcurrencyLimitCreate(ActionBaseModel):
     limit: int = FieldFrom(objects.GlobalConcurrencyLimit)
     active: Optional[bool] = FieldFrom(objects.GlobalConcurrencyLimit)
     active_slots: Optional[int] = FieldFrom(objects.GlobalConcurrencyLimit)
-    slot_decay_per_second: Optional[int] = FieldFrom(objects.GlobalConcurrencyLimit)
+    slot_decay_per_second: Optional[float] = FieldFrom(objects.GlobalConcurrencyLimit)
 
 
 @copy_model_fields
@@ -663,4 +663,4 @@ class GlobalConcurrencyLimitUpdate(ActionBaseModel):
     limit: Optional[int] = FieldFrom(objects.GlobalConcurrencyLimit)
     active: Optional[bool] = FieldFrom(objects.GlobalConcurrencyLimit)
     active_slots: Optional[int] = FieldFrom(objects.GlobalConcurrencyLimit)
-    slot_decay_per_second: Optional[int] = FieldFrom(objects.GlobalConcurrencyLimit)
+    slot_decay_per_second: Optional[float] = FieldFrom(objects.GlobalConcurrencyLimit)

prefect/client/schemas/objects.py

@@ -27,6 +27,8 @@ from typing_extensions import Literal
 from prefect._internal.schemas.bases import ObjectBaseModel, PrefectBaseModel
 from prefect._internal.schemas.fields import CreatedBy, DateTimeTZ, UpdatedBy
 from prefect._internal.schemas.validators import (
+    get_or_create_run_name,
+    get_or_create_state_name,
     raise_on_name_alphanumeric_dashes_only,
     raise_on_name_with_banned_characters,
 )
@@ -242,13 +244,7 @@ class State(ObjectBaseModel, Generic[R]):
 
     @validator("name", always=True)
     def default_name_from_type(cls, v, *, values, **kwargs):
-        """If a name is not provided, use the type"""
-
-        # if `type` is not in `values` it means the `type` didn't pass its own
-        # validation check and an error will be raised after this function is called
-        if v is None and values.get("type"):
-            v = " ".join([v.capitalize() for v in values.get("type").value.split("_")])
-        return v
+        return get_or_create_state_name(v, values)
 
     @root_validator
     def default_scheduled_start_time(cls, values):
@@ -546,7 +542,7 @@ class FlowRun(ObjectBaseModel):
 
     @validator("name", pre=True)
     def set_default_name(cls, name):
-        return name or generate_slug(2)
+        return get_or_create_run_name(name)
 
     # These are server-side optimizations and should not be present on client models
     # TODO: Deprecate these fields
@@ -754,7 +750,7 @@ class TaskRun(ObjectBaseModel):
 
     @validator("name", pre=True)
     def set_default_name(cls, name):
-        return name or generate_slug(2)
+        return get_or_create_run_name(name)
 
 
 class Workspace(PrefectBaseModel):
@@ -831,8 +827,7 @@ class BlockType(ObjectBaseModel):
 
     @validator("name", check_fields=False)
     def validate_name_characters(cls, v):
-        raise_on_name_with_banned_characters(v)
-        return v
+        return raise_on_name_with_banned_characters(v)
 
 
 class BlockSchema(ObjectBaseModel):
@@ -890,9 +885,7 @@ class BlockDocument(ObjectBaseModel):
     def validate_name_characters(cls, v):
         # the BlockDocumentCreate subclass allows name=None
         # and will inherit this validator
-        if v is not None:
-            raise_on_name_with_banned_characters(v)
-        return v
+        return raise_on_name_with_banned_characters(v)
 
     @root_validator
     def validate_name_is_present_if_not_anonymous(cls, values):
@@ -917,8 +910,7 @@ class Flow(ObjectBaseModel):
 
     @validator("name", check_fields=False)
     def validate_name_characters(cls, v):
-        raise_on_name_with_banned_characters(v)
-        return v
+        return raise_on_name_with_banned_characters(v)
 
 
 class FlowRunnerSettings(PrefectBaseModel):
@@ -1079,8 +1071,7 @@ class Deployment(ObjectBaseModel):
 
     @validator("name", check_fields=False)
     def validate_name_characters(cls, v):
-        raise_on_name_with_banned_characters(v)
-        return v
+        return raise_on_name_with_banned_characters(v)
 
 
 class ConcurrencyLimit(ObjectBaseModel):
@@ -1267,8 +1258,7 @@ class WorkQueue(ObjectBaseModel):
 
     @validator("name", check_fields=False)
     def validate_name_characters(cls, v):
-        raise_on_name_with_banned_characters(v)
-        return v
+        return raise_on_name_with_banned_characters(v)
 
 
 class WorkQueueHealthPolicy(PrefectBaseModel):
@@ -1431,8 +1421,7 @@ class WorkPool(ObjectBaseModel):
 
     @validator("name", check_fields=False)
     def validate_name_characters(cls, v):
-        raise_on_name_with_banned_characters(v)
-        return v
+        return raise_on_name_with_banned_characters(v)
 
     @validator("default_queue_id", always=True)
     def helpful_error_for_missing_default_queue_id(cls, v):
@@ -1625,8 +1614,8 @@ class GlobalConcurrencyLimit(ObjectBaseModel):
         default=0,
         description="Number of tasks currently using a concurrency slot.",
     )
-    slot_decay_per_second: Optional[int] = Field(
-        default=0,
+    slot_decay_per_second: Optional[float] = Field(
+        default=0.0,
         description=(
             "Controls the rate at which slots are released when the concurrency limit"
             " is used as a rate limit."

prefect/client/schemas/schedules.py

@@ -8,22 +8,22 @@ from typing import Optional, Union
 import dateutil
 import dateutil.rrule
 import pendulum
-from croniter import croniter
 
 from prefect._internal.pydantic import HAS_PYDANTIC_V2
+from prefect._internal.schemas.validators import (
+    default_anchor_date,
+    default_timezone,
+    interval_schedule_must_be_positive,
+    validate_cron_string,
+    validate_rrule_string,
+    validate_rrule_timezone,
+)
 
 if HAS_PYDANTIC_V2:
     from pydantic.v1 import Field, validator
 else:
     from pydantic import Field, validator
 
-from prefect._internal.pytz import HAS_PYTZ
-
-if HAS_PYTZ:
-    import pytz
-else:
-    from prefect._internal import pytz
-
 
 from prefect._internal.schemas.bases import PrefectBaseModel
 from prefect._internal.schemas.fields import DateTimeTZ
@@ -71,41 +71,16 @@ class IntervalSchedule(PrefectBaseModel):
     timezone: Optional[str] = Field(default=None, example="America/New_York")
 
     @validator("interval")
-    def interval_must_be_positive(cls, v):
-        if v.total_seconds() <= 0:
-            raise ValueError("The interval must be positive")
-        return v
+    def validate_interval_schedule(cls, v):
+        return interval_schedule_must_be_positive(v)
 
     @validator("anchor_date", always=True)
-    def default_anchor_date(cls, v):
-        if v is None:
-            return pendulum.now("UTC")
-        return pendulum.instance(v)
+    def validate_anchor_date(cls, v):
+        return default_anchor_date(v)
 
     @validator("timezone", always=True)
-    def default_timezone(cls, v, *, values, **kwargs):
-        # pendulum.tz.timezones is a callable in 3.0 and above
-        # https://github.com/PrefectHQ/prefect/issues/11619
-        if callable(pendulum.tz.timezones):
-            timezones = pendulum.tz.timezones()
-        else:
-            timezones = pendulum.tz.timezones
-        # if was provided, make sure its a valid IANA string
-        if v and v not in timezones:
-            raise ValueError(f'Invalid timezone: "{v}"')
-
-        # otherwise infer the timezone from the anchor date
-        elif v is None and values.get("anchor_date"):
-            tz = values["anchor_date"].tz.name
-            if tz in timezones:
-                return tz
-            # sometimes anchor dates have "timezones" that are UTC offsets
-            # like "-04:00". This happens when parsing ISO8601 strings.
-            # In this case we, the correct inferred localization is "UTC".
-            else:
-                return "UTC"
-
-        return v
+    def validate_default_timezone(cls, v, values):
+        return default_timezone(v, values=values)
 
 
 class CronSchedule(PrefectBaseModel):
@@ -147,31 +122,11 @@ class CronSchedule(PrefectBaseModel):
 
     @validator("timezone")
     def valid_timezone(cls, v):
-        # pendulum.tz.timezones is a callable in 3.0 and above
-        # https://github.com/PrefectHQ/prefect/issues/11619
-        if callable(pendulum.tz.timezones):
-            timezones = pendulum.tz.timezones()
-        else:
-            timezones = pendulum.tz.timezones
-
-        if v and v not in timezones:
-            raise ValueError(
-                f'Invalid timezone: "{v}" (specify in IANA tzdata format, for example,'
-                " America/New_York)"
-            )
-        return v
+        return default_timezone(v)
 
     @validator("cron")
     def valid_cron_string(cls, v):
-        # croniter allows "random" and "hashed" expressions
-        # which we do not support https://github.com/kiorky/croniter
-        if not croniter.is_valid(v):
-            raise ValueError(f'Invalid cron string: "{v}"')
-        elif any(c for c in v.split() if c.casefold() in ["R", "H", "r", "h"]):
-            raise ValueError(
-                f'Random and Hashed expressions are unsupported, received: "{v}"'
-            )
-        return v
+        return validate_cron_string(v)
 
 
 DEFAULT_ANCHOR_DATE = pendulum.date(2020, 1, 1)
@@ -205,20 +160,7 @@ class RRuleSchedule(PrefectBaseModel):
 
     @validator("rrule")
     def validate_rrule_str(cls, v):
-        # attempt to parse the rrule string as an rrule object
-        # this will error if the string is invalid
-        try:
-            dateutil.rrule.rrulestr(v, cache=True)
-        except ValueError as exc:
-            # rrules errors are a mix of cryptic and informative
-            # so reraise to be clear that the string was invalid
-            raise ValueError(f'Invalid RRule string "{v}": {exc}')
-        if len(v) > MAX_RRULE_LENGTH:
-            raise ValueError(
-                f'Invalid RRule string "{v[:40]}..."\n'
-                f"Max length is {MAX_RRULE_LENGTH}, got {len(v)}"
-            )
-        return v
+        return validate_rrule_string(v)
 
     @classmethod
     def from_rrule(cls, rrule: dateutil.rrule.rrule):
@@ -325,11 +267,7 @@ class RRuleSchedule(PrefectBaseModel):
 
     @validator("timezone", always=True)
     def valid_timezone(cls, v):
-        if v and v not in pytz.all_timezones_set:
-            raise ValueError(f'Invalid timezone: "{v}"')
-        elif v is None:
-            return "UTC"
-        return v
+        return validate_rrule_timezone(v)
 
 
 class NoSchedule(PrefectBaseModel):