prefect-client 3.7.9.dev7__tar.gz → 3.8.1.dev1__tar.gz

{prefect_client-3.7.9.dev7 → prefect_client-3.8.1.dev1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: prefect-client
-Version: 3.7.9.dev7
+Version: 3.8.1.dev1
 Summary: Workflow orchestration and management.
 Project-URL: Changelog, https://github.com/PrefectHQ/prefect/releases
 Project-URL: Documentation, https://docs.prefect.io

{prefect_client-3.7.9.dev7 → prefect_client-3.8.1.dev1}/src/prefect/AGENTS.md

@@ -23,6 +23,7 @@ There is no formal public/private boundary beyond the `_` prefix convention. Mod
 - **Engine ordering matters.** The engines apply features (retries, caching, result persistence, transactions) in a specific order. Changing the order or forgetting a feature in one engine path is the most common source of breakage.
 - **Set run metadata before state transitions.** Any run attribute that event subscribers observe at the moment of a state transition (e.g., custom `flow_run_name`) must be persisted to the server *before* `set_state()` is called for that transition. Metadata set after the call appears stale to subscribers of that state event.
 - **Sync and async must stay in sync.** Both `flow_engine.py` and `task_engine.py` have sync and async paths. Any behavior change must be applied to both.
+- **Event and log WebSocket subscribers are parallel implementations.** `events/clients.py::PrefectEventSubscriber` and `logging/clients.py::PrefectLogsSubscriber` duplicate the same reconnect/backoff/dedup logic, including the `reconnect_on_clean_close` option (used by `events/subscribers.py::FlowRunSubscriber` so a clean disconnect doesn't hide a terminal event). Mirror any reconnection-behavior change across both.
 - **Use `whenever` compat helpers for all datetime conversions on Python 3.13+.** `types/_datetime.py` provides `now()`, `_whenever_to_stdlib()`, `_whenever_zdt_from_py()`, and `_whenever_pdt_from_py()` to abstract over API differences between whenever 0.7.x–0.9.x and ≥ 0.10.0. The `_WHENEVER_NEW_API` flag (True when whenever ≥ 0.10.0) guards version-specific code. Never call `DateTime.now()` or `pendulum.now()` directly — use `now()` from `types/_datetime.py` instead. Avoid no-argument `datetime.datetime.now()` because it returns a tz-naive datetime that raises `TypeError` when compared to the tz-aware datetimes produced by `now()`. Explicitly timezone-aware stdlib calls such as `datetime.now(ZoneInfo("UTC"))` are allowed when needed to preserve timezone information. Never call `ZonedDateTime.from_py_datetime()`, `PlainDateTime.from_py_datetime()`, or `.py_datetime()` directly — use the helpers instead. Violations will silently break on whichever whenever version the helpers weren't written for.
 - **Flow and task engines advance state differently.** The flow engine makes blocking API calls to the server to propose and advance states. The task engine emits `prefect.task_run.*` events (delivered via WebSockets) but advances state locally through `set_state` calls with polling/backoff — do not assume the two engines work the same way.
 - **Flow state transitions go through the server.** The flow engine proposes states to the server, which accepts or rejects them via orchestration rules. The task engine, by contrast, manages state transitions locally via `set_state` and emits `prefect.task_run.*` events — it does not propose states to the server.

prefect_client-3.8.1.dev1/src/prefect/_build_info.py

@@ -0,0 +1,5 @@
+# Generated by versioningit
+__version__ = "3.8.1.dev1"
+__build_date__ = "2026-07-25 08:45:04.608753+00:00"
+__git_commit__ = "d2d873b2e1d8747a23e1719882da81d65f2400bb"
+__dirty__ = False

prefect_client-3.8.1.dev1/src/prefect/concurrency/context.py

@@ -0,0 +1,53 @@
+import logging
+from contextvars import ContextVar
+from typing import Any, ClassVar
+from uuid import UUID
+
+from typing_extensions import Self
+
+from prefect.client.orchestration import get_client
+from prefect.context import ContextModel, Field
+from prefect.logging.loggers import get_logger, get_run_logger
+
+
+def _cleanup_logger() -> "logging.Logger | logging.LoggerAdapter[logging.Logger]":
+    try:
+        # Use a run logger if available so failures reach the run logs
+        return get_run_logger()
+    except Exception:
+        return get_logger("concurrency")
+
+
+class ConcurrencyContext(ContextModel):
+    __var__: ClassVar[ContextVar[Self]] = ContextVar("concurrency")
+
+    # Track the leases that have been acquired but were not able to be released
+    # due to cancellation or some other error. These leases are revoked when
+    # the context manager exits.
+    cleanup_lease_ids: list[UUID] = Field(default_factory=lambda: [])
+
+    def __exit__(self, *exc_info: Any) -> None:
+        # Releasing these leases is best-effort: a failure for one lease must not
+        # strand the others or skip the context teardown below.
+        if self.cleanup_lease_ids:
+            logger = _cleanup_logger()
+            try:
+                with get_client(sync_client=True) as client:
+                    for lease_id in self.cleanup_lease_ids:
+                        try:
+                            client.release_concurrency_slots_with_lease(
+                                lease_id=lease_id
+                            )
+                        except Exception:
+                            logger.warning(
+                                "Failed to release concurrency lease %s during cleanup",
+                                lease_id,
+                                exc_info=True,
+                            )
+            except Exception:
+                logger.warning(
+                    "Failed to release concurrency leases during cleanup",
+                    exc_info=True,
+                )
+
+        return super().__exit__(*exc_info)

prefect_client-3.8.1.dev1/src/prefect/concurrency/v1/context.py

@@ -0,0 +1,56 @@
+import logging
+from contextvars import ContextVar
+from typing import Any, ClassVar
+from uuid import UUID
+
+from typing_extensions import Self
+
+from prefect.client.orchestration import get_client
+from prefect.context import ContextModel, Field
+from prefect.logging.loggers import get_logger, get_run_logger
+
+
+def _cleanup_logger() -> "logging.Logger | logging.LoggerAdapter[logging.Logger]":
+    try:
+        # Use a run logger if available so failures reach the run logs
+        return get_run_logger()
+    except Exception:
+        return get_logger("concurrency")
+
+
+class ConcurrencyContext(ContextModel):
+    __var__: ClassVar[ContextVar[Self]] = ContextVar("concurrency_v1")
+
+    # Track the limits that have been acquired but were not able to be released
+    # due to cancellation or some other error. These limits are released when
+    # the context manager exits.
+    cleanup_slots: list[tuple[list[str], float, UUID]] = Field(default_factory=list)
+
+    def __exit__(self, *exc_info: Any) -> None:
+        # Releasing these slots is best-effort: a failure for one entry must not
+        # strand the others or skip the context teardown below.
+        if self.cleanup_slots:
+            logger = _cleanup_logger()
+            try:
+                with get_client(sync_client=True) as client:
+                    for names, occupancy_seconds, task_run_id in self.cleanup_slots:
+                        try:
+                            client.decrement_v1_concurrency_slots(
+                                names=names,
+                                occupancy_seconds=occupancy_seconds,
+                                task_run_id=task_run_id,
+                            )
+                        except Exception:
+                            logger.warning(
+                                "Failed to release concurrency slots %s for task run %s during cleanup",
+                                names,
+                                task_run_id,
+                                exc_info=True,
+                            )
+            except Exception:
+                logger.warning(
+                    "Failed to release concurrency slots during cleanup",
+                    exc_info=True,
+                )
+
+        return super().__exit__(*exc_info)

{prefect_client-3.7.9.dev7 → prefect_client-3.8.1.dev1}/src/prefect/events/clients.py

@@ -136,16 +136,21 @@ def get_events_client(
 def get_events_subscriber(
     filter: Optional["EventFilter"] = None,
     reconnection_attempts: int = 10,
+    reconnect_on_clean_close: bool = False,
 ) -> "PrefectEventSubscriber":
     api_url = PREFECT_API_URL.value()
 
     if isinstance(api_url, str) and api_url.startswith(PREFECT_CLOUD_API_URL.value()):
         return PrefectCloudEventSubscriber(
-            filter=filter, reconnection_attempts=reconnection_attempts
+            filter=filter,
+            reconnection_attempts=reconnection_attempts,
+            reconnect_on_clean_close=reconnect_on_clean_close,
         )
     elif api_url:
         return PrefectEventSubscriber(
-            filter=filter, reconnection_attempts=reconnection_attempts
+            filter=filter,
+            reconnection_attempts=reconnection_attempts,
+            reconnect_on_clean_close=reconnect_on_clean_close,
         )
     elif PREFECT_SERVER_ALLOW_EPHEMERAL_MODE:
         from prefect.server.api.server import SubprocessASGIServer
@@ -156,6 +161,7 @@ def get_events_subscriber(
             api_url=server.api_url,
             filter=filter,
             reconnection_attempts=reconnection_attempts,
+            reconnect_on_clean_close=reconnect_on_clean_close,
         )
     else:
         raise ValueError(
@@ -677,6 +683,10 @@ class PrefectEventSubscriber:
     """
     Subscribes to a Prefect event stream, yielding events as they occur.
 
+    Abnormal disconnects are retried. Clean disconnects end iteration unless
+    `reconnect_on_clean_close` is enabled. Event backfill and ID-based
+    deduplication cover the reconnect window.
+
     Example:
 
         from prefect.events.clients import PrefectEventSubscriber
@@ -693,6 +703,7 @@ class PrefectEventSubscriber:
     _websocket: Optional[ClientConnection]
     _filter: "EventFilter"
     _seen_events: MutableMapping[UUID, bool]
+    _backfill_since: Optional[prefect.types._datetime.DateTime]
 
     _api_key: Optional[str]
     _auth_token: Optional[str]
@@ -702,6 +713,7 @@ class PrefectEventSubscriber:
         api_url: Optional[str] = None,
         filter: Optional["EventFilter"] = None,
         reconnection_attempts: int = 10,
+        reconnect_on_clean_close: bool = False,
     ):
         """
         Args:
@@ -709,6 +721,8 @@ class PrefectEventSubscriber:
             api_key: The API of an actor with the manage_events scope
             reconnection_attempts: When the client is disconnected, how many times
                 the client should attempt to reconnect
+            reconnect_on_clean_close: Whether to reconnect when the server closes
+                the connection normally instead of ending iteration
         """
         self._api_key = None
         auth_string = get_current_settings().api.auth_string
@@ -725,6 +739,7 @@ class PrefectEventSubscriber:
 
         self._filter = filter or EventFilter()  # type: ignore[call-arg]
         self._seen_events = TTLCache(maxsize=SEEN_EVENTS_SIZE, ttl=SEEN_EVENTS_TTL)
+        self._backfill_since = None
 
         socket_url = events_out_socket_from_api_url(api_url)
 
@@ -736,6 +751,7 @@ class PrefectEventSubscriber:
         )
         self._websocket = None
         self._reconnection_attempts = reconnection_attempts
+        self._reconnect_on_clean_close = reconnect_on_clean_close
         if self._reconnection_attempts < 0:
             raise ValueError("reconnection_attempts must be a non-negative integer")
 
@@ -808,9 +824,13 @@ class PrefectEventSubscriber:
 
         from prefect.events.filters import EventOccurredFilter
 
+        current_time = prefect.types._datetime.now("UTC")
+        if self._backfill_since is None:
+            self._backfill_since = current_time - timedelta(minutes=1)
+
         self._filter.occurred = EventOccurredFilter(
-            since=prefect.types._datetime.now("UTC") - timedelta(minutes=1),
-            until=prefect.types._datetime.now("UTC") + timedelta(days=365),
+            since=self._backfill_since,
+            until=current_time + timedelta(days=365),
         )
 
         logger.debug("  filtering events since %s...", self._filter.occurred.since)
@@ -841,6 +861,7 @@ class PrefectEventSubscriber:
     async def __anext__(self) -> Event:
         assert self._reconnection_attempts >= 0
         consecutive_failures = 0
+        consecutive_clean_closes = 0
         while consecutive_failures <= self._reconnection_attempts:
             try:
                 # If we're here and the websocket is None, then we've had a failure in a
@@ -865,25 +886,39 @@ class PrefectEventSubscriber:
                     if event.id in self._seen_events:
                         continue
                     self._seen_events[event.id] = True
+                    replay_since = event.occurred - timedelta(minutes=1)
+                    if (
+                        self._backfill_since is None
+                        or replay_since > self._backfill_since
+                    ):
+                        self._backfill_since = replay_since
 
                     try:
                         return event
                     finally:
                         EVENTS_OBSERVED.labels(self.client_name).inc()
-            except ConnectionClosedOK:
-                logger.debug('Connection closed with "OK" status')
-                raise StopAsyncIteration
-            except RETRYABLE_EXCEPTIONS:
+            except RETRYABLE_EXCEPTIONS as exc:
+                if isinstance(exc, ConnectionClosedOK):
+                    if not self._reconnect_on_clean_close:
+                        logger.debug('Connection closed with "OK" status')
+                        raise StopAsyncIteration
+                    consecutive_clean_closes += 1
+
                 consecutive_failures += 1
+                attempts = (
+                    consecutive_clean_closes
+                    if isinstance(exc, ConnectionClosedOK)
+                    else consecutive_failures
+                )
                 logger.debug(
                     "Retryable error with %s/%s attempts",
-                    consecutive_failures,
+                    attempts,
                     self._reconnection_attempts,
                 )
-                if consecutive_failures > self._reconnection_attempts:
+                if attempts > self._reconnection_attempts:
                     raise
 
-                if consecutive_failures > 2:
+                if attempts > 2:
                     # let the first two attempts happen quickly in case this is just
                     # a standard load balancer timeout, but after that, just take a
                     # beat to let things come back around.
@@ -898,6 +933,7 @@ class PrefectCloudEventSubscriber(PrefectEventSubscriber):
         api_key: Optional[str] = None,
         filter: Optional["EventFilter"] = None,
         reconnection_attempts: int = 10,
+        reconnect_on_clean_close: bool = False,
     ):
         """
         Args:
@@ -905,6 +941,8 @@ class PrefectCloudEventSubscriber(PrefectEventSubscriber):
             api_key: The API of an actor with the manage_events scope
             reconnection_attempts: When the client is disconnected, how many times
                 the client should attempt to reconnect
+            reconnect_on_clean_close: Whether to reconnect when the server closes
+                the connection normally instead of ending iteration
         """
         api_url, api_key = _get_api_url_and_key(api_url, api_key)
 
@@ -912,6 +950,7 @@ class PrefectCloudEventSubscriber(PrefectEventSubscriber):
             api_url=api_url,
             filter=filter,
             reconnection_attempts=reconnection_attempts,
+            reconnect_on_clean_close=reconnect_on_clean_close,
         )
 
         self._api_key = api_key
@@ -924,6 +963,7 @@ class PrefectCloudAccountEventSubscriber(PrefectCloudEventSubscriber):
         api_key: Optional[str] = None,
         filter: Optional["EventFilter"] = None,
         reconnection_attempts: int = 10,
+        reconnect_on_clean_close: bool = False,
     ):
         """
         Args:
@@ -931,6 +971,8 @@ class PrefectCloudAccountEventSubscriber(PrefectCloudEventSubscriber):
             api_key: The API of an actor with the manage_events scope
             reconnection_attempts: When the client is disconnected, how many times
                 the client should attempt to reconnect
+            reconnect_on_clean_close: Whether to reconnect when the server closes
+                the connection normally instead of ending iteration
         """
         api_url, api_key = _get_api_url_and_key(api_url, api_key)
 
@@ -940,6 +982,7 @@ class PrefectCloudAccountEventSubscriber(PrefectCloudEventSubscriber):
             api_url=account_api_url,
             filter=filter,
             reconnection_attempts=reconnection_attempts,
+            reconnect_on_clean_close=reconnect_on_clean_close,
         )
 
         self._api_key = api_key

{prefect_client-3.7.9.dev7 → prefect_client-3.8.1.dev1}/src/prefect/events/subscribers.py

@@ -33,6 +33,9 @@ class FlowRunSubscriber:
     is received, the event subscription stops but log subscription continues for a
     configurable timeout to catch any straggler logs.
 
+    The underlying event and log subscribers own reconnection behavior. If either
+    exhausts its retries or fails for another reason, that error is propagated.
+
     Example:
         ```python
         from prefect.events.subscribers import FlowRunSubscriber
@@ -47,7 +50,7 @@ class FlowRunSubscriber:
     """
 
     _flow_run_id: UUID
-    _queue: asyncio.Queue[Union[Log, Event, None]]
+    _queue: asyncio.Queue[Union[Log, Event, Exception, None]]
     _tasks: list[asyncio.Task[None]]
     _flow_completed: bool
     _straggler_timeout: int
@@ -91,10 +94,14 @@ class FlowRunSubscriber:
     async def __aenter__(self) -> Self:
         """Enter the async context manager"""
         self._logs_subscriber = get_logs_subscriber(
-            filter=self._log_filter, reconnection_attempts=self._reconnection_attempts
+            filter=self._log_filter,
+            reconnection_attempts=self._reconnection_attempts,
+            reconnect_on_clean_close=True,
         )
         self._events_subscriber = get_events_subscriber(
-            filter=self._event_filter, reconnection_attempts=self._reconnection_attempts
+            filter=self._event_filter,
+            reconnection_attempts=self._reconnection_attempts,
+            reconnect_on_clean_close=True,
         )
 
         await self._logs_subscriber.__aenter__()
@@ -122,6 +129,22 @@ class FlowRunSubscriber:
         await self._logs_subscriber.__aexit__(exc_type, exc_val, exc_tb)
         await self._events_subscriber.__aexit__(exc_type, exc_val, exc_tb)
 
+    def _is_terminal_event(self, event: Event) -> bool:
+        """Whether an event marks the flow run reaching a terminal state"""
+        if event.resource.id != f"prefect.flow-run.{self._flow_run_id}":
+            return False
+
+        state_type_str = event.resource.get("prefect.state-type")
+        if not state_type_str and "validated_state" in event.payload:
+            state_type_str = event.payload["validated_state"].get("type")
+        if not state_type_str:
+            return False
+
+        try:
+            return StateType(state_type_str) in TERMINAL_STATES
+        except ValueError:
+            return False
+
     def __aiter__(self) -> Self:
         """Return self as an async iterator"""
         return self
@@ -142,6 +165,8 @@ class FlowRunSubscriber:
             if item is None:
                 self._sentinels_received += 1
                 continue
+            if isinstance(item, Exception):
+                raise item
 
             return item
 
@@ -154,8 +179,13 @@ class FlowRunSubscriber:
                 await self._queue.put(log)
         except asyncio.CancelledError:
             pass
-        except Exception:
-            pass
+        except Exception as exc:
+            await self._queue.put(exc)
+        else:
+            if not self._flow_completed:
+                await self._queue.put(
+                    ConnectionError("Flow run logs stream closed unexpectedly")
+                )
         finally:
             await self._queue.put(None)
 
@@ -164,23 +194,17 @@ class FlowRunSubscriber:
         try:
             async for event in self._events_subscriber:
                 await self._queue.put(event)
-
-                # Check if this is a terminal state event for our flow run
-                if event.resource.id == f"prefect.flow-run.{self._flow_run_id}":
-                    # Get state type from event resource or payload
-                    state_type_str = event.resource.get("prefect.state-type")
-                    if not state_type_str and "validated_state" in event.payload:
-                        state_type_str = event.payload["validated_state"].get("type")
-
-                    if state_type_str:
-                        try:
-                            state_type = StateType(state_type_str)
-                            if state_type in TERMINAL_STATES:
-                                self._flow_completed = True
-                                break
-                        except ValueError:
-                            pass
-        except Exception:
+                if self._is_terminal_event(event):
+                    self._flow_completed = True
+                    break
+        except asyncio.CancelledError:
             pass
+        except Exception as exc:
+            await self._queue.put(exc)
+        else:
+            if not self._flow_completed:
+                await self._queue.put(
+                    ConnectionError("Flow run events stream closed unexpectedly")
+                )
         finally:
             await self._queue.put(None)

{prefect_client-3.7.9.dev7 → prefect_client-3.8.1.dev1}/src/prefect/exceptions.py

@@ -425,6 +425,10 @@ class FlowRunWaitTimeout(PrefectException):
     """Raised when a flow run takes longer than a given timeout"""
 
 
+class FlowRunWatchError(PrefectException):
+    """Raised when watching a flow run stops before it reaches a terminal state"""
+
+
 class PrefectImportError(ImportError):
     """
     An error raised when a Prefect object cannot be imported due to a move or removal.

{prefect_client-3.7.9.dev7 → prefect_client-3.8.1.dev1}/src/prefect/futures.py

@@ -610,6 +610,7 @@ def as_completed(
     unique_futures: set[PrefectFuture[R]] = set(futures)
     total_futures = len(unique_futures)
     pending = unique_futures
+    deadline: float | None = None
     try:
         with timeout_context(timeout):
             done = {f for f in unique_futures if f._final_state}  # type: ignore[privateUsage]
@@ -628,8 +629,17 @@ def as_completed(
             for future in pending:
                 _register_prefect_done_callback(future, add_to_done)
 
+            deadline = time.monotonic() + timeout if timeout is not None else None
+
             while pending:
-                finished_event.wait()
+                if timeout is None:
+                    finished_event.wait()
+                else:
+                    assert deadline is not None
+                    remaining = deadline - time.monotonic()
+                    if not finished_event.wait(timeout=max(remaining, 0.0)):
+                        raise TimeoutError
+
                 with finished_lock:
                     done = finished_futures
                     finished_futures = []

{prefect_client-3.7.9.dev7 → prefect_client-3.8.1.dev1}/src/prefect/logging/clients.py

@@ -47,6 +47,13 @@ if TYPE_CHECKING:
 
 logger: "logging.Logger" = get_logger(__name__)
 
+# Exceptions that indicate transient network issues and should trigger retries.
+_RETRYABLE_EXCEPTIONS: Tuple[Type[Exception], ...] = (
+    ConnectionClosed,
+    TimeoutError,
+    OSError,
+)
+
 LOGS_OBSERVED = Counter(
     "prefect_logs_observed",
     "The number of logs observed by Prefect log subscribers",
@@ -90,6 +97,7 @@ def _get_api_url_and_key(
 def get_logs_subscriber(
     filter: Optional["LogFilter"] = None,
     reconnection_attempts: int = 10,
+    reconnect_on_clean_close: bool = False,
 ) -> "PrefectLogsSubscriber":
     """
     Get a logs subscriber based on the current Prefect configuration.
@@ -101,13 +109,16 @@ def get_logs_subscriber(
 
     if isinstance(api_url, str) and api_url.startswith(PREFECT_CLOUD_API_URL.value()):
         return PrefectCloudLogsSubscriber(
-            filter=filter, reconnection_attempts=reconnection_attempts
+            filter=filter,
+            reconnection_attempts=reconnection_attempts,
+            reconnect_on_clean_close=reconnect_on_clean_close,
         )
     elif api_url:
         return PrefectLogsSubscriber(
             api_url=api_url,
             filter=filter,
             reconnection_attempts=reconnection_attempts,
+            reconnect_on_clean_close=reconnect_on_clean_close,
         )
     elif PREFECT_SERVER_ALLOW_EPHEMERAL_MODE:
         from prefect.server.api.server import SubprocessASGIServer
@@ -118,6 +129,7 @@ def get_logs_subscriber(
             api_url=server.api_url,
             filter=filter,
             reconnection_attempts=reconnection_attempts,
+            reconnect_on_clean_close=reconnect_on_clean_close,
         )
     else:
         raise ValueError(
@@ -129,6 +141,10 @@ class PrefectLogsSubscriber:
     """
     Subscribes to a Prefect logs stream, yielding logs as they occur.
 
+    Abnormal disconnects are retried. Clean disconnects end iteration unless
+    `reconnect_on_clean_close` is enabled. The logs stream is live-only, so logs
+    emitted while disconnected may be omitted.
+
     Example:
 
         from prefect.logging.clients import PrefectLogsSubscriber
@@ -155,6 +171,7 @@ class PrefectLogsSubscriber:
         api_url: Optional[str] = None,
         filter: Optional["LogFilter"] = None,
         reconnection_attempts: int = 10,
+        reconnect_on_clean_close: bool = False,
     ):
         """
         Args:
@@ -162,6 +179,8 @@ class PrefectLogsSubscriber:
             filter: Log filter to apply
             reconnection_attempts: When the client is disconnected, how many times
                 the client should attempt to reconnect
+            reconnect_on_clean_close: Whether to reconnect when the server closes
+                the connection normally instead of ending iteration
         """
         self._api_key = None
         self._auth_token = PREFECT_API_AUTH_STRING.value()
@@ -189,6 +208,7 @@ class PrefectLogsSubscriber:
         self._connect = websocket_connect(socket_url, **connect_kwargs)
         self._websocket = None
         self._reconnection_attempts = reconnection_attempts
+        self._reconnect_on_clean_close = reconnect_on_clean_close
         if self._reconnection_attempts < 0:
             raise ValueError("reconnection_attempts must be a non-negative integer")
 
@@ -272,14 +292,15 @@ class PrefectLogsSubscriber:
 
     async def __anext__(self) -> Log:
         assert self._reconnection_attempts >= 0
-        for i in range(self._reconnection_attempts + 1):  # pragma: no branch
+        attempts = 0
+        while attempts <= self._reconnection_attempts:
             try:
                 # If we're here and the websocket is None, then we've had a failure in a
                 # previous reconnection attempt.
                 #
                 # Otherwise, after the first time through this loop, we're recovering
                 # from a ConnectionClosed, so reconnect now.
-                if not self._websocket or i > 0:
+                if not self._websocket or attempts > 0:
                     try:
                         await self._reconnect()
                     finally:
@@ -301,20 +322,22 @@ class PrefectLogsSubscriber:
                     finally:
                         LOGS_OBSERVED.labels(self.client_name).inc()
 
-            except ConnectionClosedOK:
-                logger.debug('Connection closed with "OK" status')
-                raise StopAsyncIteration
-            except ConnectionClosed:
+            except _RETRYABLE_EXCEPTIONS as exc:
+                if isinstance(exc, ConnectionClosedOK):
+                    if not self._reconnect_on_clean_close:
+                        logger.debug('Connection closed with "OK" status')
+                        raise StopAsyncIteration
+
+                attempts += 1
                 logger.debug(
-                    "Connection closed with %s/%s attempts",
-                    i + 1,
+                    "Retryable error with %s/%s attempts",
+                    attempts,
                     self._reconnection_attempts,
                 )
-                if i == self._reconnection_attempts:
-                    # this was our final chance, raise the most recent error
+                if attempts > self._reconnection_attempts:
                     raise
 
-                if i > 2:
+                if attempts > 2:
                     # let the first two attempts happen quickly in case this is just
                     # a standard load balancer timeout, but after that, just take a
                     # beat to let things come back around.
@@ -331,6 +354,7 @@ class PrefectCloudLogsSubscriber(PrefectLogsSubscriber):
         api_key: Optional[str] = None,
         filter: Optional["LogFilter"] = None,
         reconnection_attempts: int = 10,
+        reconnect_on_clean_close: bool = False,
     ):
         """
         Args:
@@ -339,6 +363,8 @@ class PrefectCloudLogsSubscriber(PrefectLogsSubscriber):
             filter: Log filter to apply
             reconnection_attempts: When the client is disconnected, how many times
                 the client should attempt to reconnect
+            reconnect_on_clean_close: Whether to reconnect when the server closes
+                the connection normally instead of ending iteration
         """
         api_url, api_key = _get_api_url_and_key(api_url, api_key)
 
@@ -346,6 +372,7 @@ class PrefectCloudLogsSubscriber(PrefectLogsSubscriber):
             api_url=api_url,
             filter=filter,
             reconnection_attempts=reconnection_attempts,
+            reconnect_on_clean_close=reconnect_on_clean_close,
         )
 
         self._api_key = api_key

{prefect_client-3.7.9.dev7 → prefect_client-3.8.1.dev1}/src/prefect/settings/_types.py

@@ -90,6 +90,7 @@ SettingAccessor: TypeAlias = Literal[
     "server.database.echo",
     "server.database.host",
     "server.database.migrate_on_start",
+    "server.database.migration_timeout",
     "server.database.name",
     "server.database.password",
     "server.database.port",

{prefect_client-3.7.9.dev7 → prefect_client-3.8.1.dev1}/src/prefect/settings/models/server/database.py

@@ -280,6 +280,12 @@ class ServerDatabaseSettings(PrefectBaseSettings):
         ),
     )
 
+    migration_timeout: Optional[float] = Field(
+        default=None,
+        gt=0,
+        description="A statement timeout, in seconds, applied to database migrations. Schema changes such as concurrent index builds on large tables can take much longer than an ordinary API query, so migrations are not bound by `server.database.timeout`. Defaults to `None` (no timeout); set a positive value to bound migration statement duration.",
+    )
+
     # handle deprecated fields
 
     def __getattribute__(self, name: str) -> Any: