prefect-client 2.19.4__py3-none-any.whl → 3.0.0rc2__py3-none-any.whl

prefect/client/schemas/responses.py

@@ -2,19 +2,13 @@ import datetime
 from typing import Any, Dict, List, Optional, TypeVar, Union
 from uuid import UUID
 
-from prefect._internal.compatibility.deprecated import DeprecatedInfraOverridesField
-from prefect._internal.pydantic import HAS_PYDANTIC_V2
-
-if HAS_PYDANTIC_V2:
-    from pydantic.v1 import Field
-else:
-    from pydantic import Field
-
+from pydantic import ConfigDict, Field
+from pydantic_extra_types.pendulum_dt import DateTime
 from typing_extensions import Literal
 
 import prefect.client.schemas.objects as objects
 from prefect._internal.schemas.bases import ObjectBaseModel, PrefectBaseModel
-from prefect._internal.schemas.fields import CreatedBy, DateTimeTZ, UpdatedBy
+from prefect._internal.schemas.fields import CreatedBy, UpdatedBy
 from prefect.client.schemas.schedules import SCHEDULE_TYPES
 from prefect.utilities.collections import AutoEnum
 from prefect.utilities.names import generate_slug
@@ -120,10 +114,10 @@ class HistoryResponseState(PrefectBaseModel):
 class HistoryResponse(PrefectBaseModel):
     """Represents a history of aggregation states over an interval"""
 
-    interval_start: DateTimeTZ = Field(
+    interval_start: DateTime = Field(
         default=..., description="The start date of the interval."
     )
-    interval_end: DateTimeTZ = Field(
+    interval_end: DateTime = Field(
         default=..., description="The end date of the interval."
     )
     states: List[HistoryResponseState] = Field(
@@ -147,8 +141,7 @@ class OrchestrationResult(PrefectBaseModel):
 
 
 class WorkerFlowRunResponse(PrefectBaseModel):
-    class Config:
-        arbitrary_types_allowed = True
+    model_config = ConfigDict(arbitrary_types_allowed=True)
 
     work_pool_id: UUID
     work_queue_id: UUID
@@ -219,18 +212,18 @@ class FlowRunResponse(ObjectBaseModel):
     run_count: int = Field(
         default=0, description="The number of times the flow run was executed."
     )
-    expected_start_time: Optional[DateTimeTZ] = Field(
+    expected_start_time: Optional[DateTime] = Field(
         default=None,
         description="The flow run's expected start time.",
     )
-    next_scheduled_start_time: Optional[DateTimeTZ] = Field(
+    next_scheduled_start_time: Optional[DateTime] = Field(
         default=None,
         description="The next time the flow run is scheduled to start.",
     )
-    start_time: Optional[DateTimeTZ] = Field(
+    start_time: Optional[DateTime] = Field(
         default=None, description="The actual start time."
     )
-    end_time: Optional[DateTimeTZ] = Field(
+    end_time: Optional[DateTime] = Field(
         default=None, description="The actual end time."
     )
     total_run_time: datetime.timedelta = Field(
@@ -279,7 +272,7 @@ class FlowRunResponse(ObjectBaseModel):
     state: Optional[objects.State] = Field(
         default=None,
         description="The state of the flow run.",
-        examples=[objects.State(type=objects.StateType.COMPLETED)],
+        examples=["objects.State(type=objects.StateType.COMPLETED)"],
     )
     job_variables: Optional[dict] = Field(
         default=None, description="Job variables for the flow run."
@@ -304,13 +297,13 @@ class FlowRunResponse(ObjectBaseModel):
         """
         if isinstance(other, objects.FlowRun):
             exclude_fields = {"estimated_run_time", "estimated_start_time_delta"}
-            return self.dict(exclude=exclude_fields) == other.dict(
+            return self.model_dump(exclude=exclude_fields) == other.model_dump(
                 exclude=exclude_fields
             )
         return super().__eq__(other)
 
 
-class DeploymentResponse(DeprecatedInfraOverridesField, ObjectBaseModel):
+class DeploymentResponse(ObjectBaseModel):
     name: str = Field(default=..., description="The name of the deployment.")
     version: Optional[str] = Field(
         default=None, description="An optional version for the deployment."
@@ -357,7 +350,7 @@ class DeploymentResponse(DeprecatedInfraOverridesField, ObjectBaseModel):
             " be scheduled."
         ),
     )
-    last_polled: Optional[DateTimeTZ] = Field(
+    last_polled: Optional[DateTime] = Field(
         default=None,
         description="The last time the deployment was polled for status updates.",
     )
@@ -407,7 +400,7 @@ class DeploymentResponse(DeprecatedInfraOverridesField, ObjectBaseModel):
         ),
     )
     enforce_parameter_schema: bool = Field(
-        default=False,
+        default=True,
         description=(
             "Whether or not the deployment should enforce the parameter schema."
         ),
@@ -423,8 +416,7 @@ class DeploymentResponse(DeprecatedInfraOverridesField, ObjectBaseModel):
 
 
 class MinimalConcurrencyLimitResponse(PrefectBaseModel):
-    class Config:
-        extra = "ignore"  # 2024/4/1
+    model_config = ConfigDict(extra="ignore")
 
     id: UUID
     name: str

prefect/client/schemas/schedules.py

@@ -3,28 +3,21 @@ Schedule schemas
 """
 
 import datetime
-from typing import Optional, Union
+from typing import Annotated, Optional, Union
 
 import dateutil
 import dateutil.rrule
 import pendulum
+from pydantic import AfterValidator, ConfigDict, Field, field_validator, model_validator
+from pydantic_extra_types.pendulum_dt import DateTime
 
-from prefect._internal.pydantic import HAS_PYDANTIC_V2
 from prefect._internal.schemas.bases import PrefectBaseModel
-from prefect._internal.schemas.fields import DateTimeTZ
 from prefect._internal.schemas.validators import (
     default_anchor_date,
     default_timezone,
     validate_cron_string,
     validate_rrule_string,
-    validate_rrule_timezone,
 )
-from prefect.types import PositiveDuration
-
-if HAS_PYDANTIC_V2:
-    from pydantic.v1 import Field, validator
-else:
-    from pydantic import Field, validator
 
 MAX_ITERATIONS = 1000
 # approx. 1 years worth of RDATEs + buffer
@@ -55,26 +48,24 @@ class IntervalSchedule(PrefectBaseModel):
 
     Args:
         interval (datetime.timedelta): an interval to schedule on
-        anchor_date (DateTimeTZ, optional): an anchor date to schedule increments against;
+        anchor_date (DateTime, optional): an anchor date to schedule increments against;
             if not provided, the current timestamp will be used
         timezone (str, optional): a valid timezone string
     """
 
-    class Config:
-        extra = "forbid"
-        exclude_none = True
+    model_config = ConfigDict(extra="forbid", exclude_none=True)
 
-    interval: PositiveDuration
-    anchor_date: Optional[DateTimeTZ] = None
+    interval: datetime.timedelta = Field(gt=datetime.timedelta(0))
+    anchor_date: Annotated[DateTime, AfterValidator(default_anchor_date)] = Field(
+        default_factory=lambda: pendulum.now("UTC"),
+        examples=["2020-01-01T00:00:00Z"],
+    )
     timezone: Optional[str] = Field(default=None, examples=["America/New_York"])
 
-    @validator("anchor_date", always=True)
-    def validate_anchor_date(cls, v):
-        return default_anchor_date(v)
-
-    @validator("timezone", always=True)
-    def validate_default_timezone(cls, v, values):
-        return default_timezone(v, values=values)
+    @model_validator(mode="after")
+    def validate_timezone(self):
+        self.timezone = default_timezone(self.timezone, self.model_dump())
+        return self
 
 
 class CronSchedule(PrefectBaseModel):
@@ -102,8 +93,7 @@ class CronSchedule(PrefectBaseModel):
 
     """
 
-    class Config:
-        extra = "forbid"
+    model_config = ConfigDict(extra="forbid")
 
     cron: str = Field(default=..., examples=["0 0 * * *"])
     timezone: Optional[str] = Field(default=None, examples=["America/New_York"])
@@ -114,11 +104,13 @@ class CronSchedule(PrefectBaseModel):
         ),
     )
 
-    @validator("timezone")
+    @field_validator("timezone")
+    @classmethod
     def valid_timezone(cls, v):
         return default_timezone(v)
 
-    @validator("cron")
+    @field_validator("cron")
+    @classmethod
     def valid_cron_string(cls, v):
         return validate_cron_string(v)
 
@@ -146,13 +138,15 @@ class RRuleSchedule(PrefectBaseModel):
         timezone (str, optional): a valid timezone string
     """
 
-    class Config:
-        extra = "forbid"
+    model_config = ConfigDict(extra="forbid")
 
     rrule: str
-    timezone: Optional[str] = Field(default=None, examples=["America/New_York"])
+    timezone: Optional[str] = Field(
+        default="UTC", examples=["America/New_York"], validate_default=True
+    )
 
-    @validator("rrule")
+    @field_validator("rrule")
+    @classmethod
     def validate_rrule_str(cls, v):
         return validate_rrule_string(v)
 
@@ -259,14 +253,30 @@ class RRuleSchedule(PrefectBaseModel):
 
             return rrule
 
-    @validator("timezone", always=True)
+    @field_validator("timezone")
     def valid_timezone(cls, v):
-        return validate_rrule_timezone(v)
+        """
+        Validate that the provided timezone is a valid IANA timezone.
+
+        Unfortunately this list is slightly different from the list of valid
+        timezones in pendulum that we use for cron and interval timezone validation.
+        """
+        from prefect._internal.pytz import HAS_PYTZ
+
+        if HAS_PYTZ:
+            import pytz
+        else:
+            from prefect._internal import pytz
+
+        if v and v not in pytz.all_timezones_set:
+            raise ValueError(f'Invalid timezone: "{v}"')
+        elif v is None:
+            return "UTC"
+        return v
 
 
 class NoSchedule(PrefectBaseModel):
-    class Config:
-        extra = "forbid"
+    model_config = ConfigDict(extra="forbid")
 
 
 SCHEDULE_TYPES = Union[IntervalSchedule, CronSchedule, RRuleSchedule, NoSchedule]
@@ -308,6 +318,8 @@ def construct_schedule(
     if interval:
         if isinstance(interval, (int, float)):
             interval = datetime.timedelta(seconds=interval)
+        if not anchor_date:
+            anchor_date = DateTime.now()
         schedule = IntervalSchedule(
             interval=interval, anchor_date=anchor_date, timezone=timezone
         )

prefect/client/subscriptions.py

@@ -4,7 +4,7 @@ from typing import Any, Dict, Generic, List, Optional, Type, TypeVar
 import orjson
 import websockets
 import websockets.exceptions
-from prefect._vendor.starlette.status import WS_1008_POLICY_VIOLATION
+from starlette.status import WS_1008_POLICY_VIOLATION
 from typing_extensions import Self
 
 from prefect._internal.schemas.bases import IDBaseModel
@@ -48,7 +48,7 @@ class Subscription(Generic[S]):
 
                 await self._websocket.send(orjson.dumps({"type": "ack"}).decode())
 
-                return self.model.parse_raw(message)
+                return self.model.model_validate_json(message)
             except (
                 ConnectionRefusedError,
                 websockets.exceptions.ConnectionClosedError,

prefect/client/utilities.py

@@ -42,12 +42,15 @@ def get_or_create_client(
     if client is not None:
         return client, True
     from prefect._internal.concurrency.event_loop import get_running_loop
-    from prefect.context import FlowRunContext, TaskRunContext
+    from prefect.context import ClientContext, FlowRunContext, TaskRunContext
 
+    client_context = ClientContext.get()
     flow_run_context = FlowRunContext.get()
     task_run_context = TaskRunContext.get()
 
-    if (
+    if client_context and client_context.async_client._loop == get_running_loop():
+        return client_context.async_client, True
+    elif (
         flow_run_context
         and getattr(flow_run_context.client, "_loop", None) == get_running_loop()
     ):

prefect/concurrency/asyncio.py

@@ -11,7 +11,7 @@ except ImportError:
     # pendulum < 3
     from pendulum.period import Period as Interval  # type: ignore
 
-from prefect import get_client
+from prefect.client.orchestration import get_client
 from prefect.client.schemas.responses import MinimalConcurrencyLimitResponse
 from prefect.utilities.timeout import timeout_async
 
@@ -117,4 +117,6 @@ async def _release_concurrency_slots(
 def _response_to_minimal_concurrency_limit_response(
     response: httpx.Response,
 ) -> List[MinimalConcurrencyLimitResponse]:
-    return [MinimalConcurrencyLimitResponse.parse_obj(obj_) for obj_ in response.json()]
+    return [
+        MinimalConcurrencyLimitResponse.model_validate(obj_) for obj_ in response.json()
+    ]

prefect/concurrency/events.py

@@ -20,7 +20,7 @@ def _emit_concurrency_event(
     }
 
     related = [
-        RelatedResource.parse_obj(
+        RelatedResource.model_validate(
             {
                 "prefect.resource.id": f"prefect.concurrency-limit.{limit.id}",
                 "prefect.resource.role": "concurrency-limit",

prefect/concurrency/services.py

@@ -2,23 +2,26 @@ import asyncio
 import concurrent.futures
 from contextlib import asynccontextmanager
 from typing import (
+    TYPE_CHECKING,
     FrozenSet,
     Tuple,
 )
 
 import httpx
-from prefect._vendor.starlette import status
+from starlette import status
 
-from prefect import get_client
 from prefect._internal.concurrency import logger
 from prefect._internal.concurrency.services import QueueService
-from prefect.client.orchestration import PrefectClient
+from prefect.client.orchestration import get_client
+
+if TYPE_CHECKING:
+    from prefect.client.orchestration import PrefectClient
 
 
 class ConcurrencySlotAcquisitionService(QueueService):
     def __init__(self, concurrency_limit_names: FrozenSet[str]):
         super().__init__(concurrency_limit_names)
-        self._client: PrefectClient
+        self._client: "PrefectClient"
         self.concurrency_limit_names = sorted(list(concurrency_limit_names))
 
     @asynccontextmanager