Move attempts logic to decorator

Author: buurro

src/fastapi_cloud_cli/commands/deploy.py

@@ -19,7 +19,7 @@
 from typing_extensions import Annotated
 
 from fastapi_cloud_cli.commands.login import login
-from fastapi_cloud_cli.utils.api import APIClient, BuildLogError
+from fastapi_cloud_cli.utils.api import APIClient, BuildLogError, TooManyRetriesError
 from fastapi_cloud_cli.utils.apps import AppConfig, get_app_config, write_app_config
 from fastapi_cloud_cli.utils.auth import is_logged_in
 from fastapi_cloud_cli.utils.cli import get_rich_toolkit, handle_http_errors
@@ -375,7 +375,7 @@ def _wait_for_deployment(
 
                     last_message_changed_at = time.monotonic()
 
-        except BuildLogError as e:
+        except (BuildLogError, TooManyRetriesError) as e:
             logger.error("Build log streaming failed: %s", e)
             toolkit.print_line()
             toolkit.print(

src/fastapi_cloud_cli/utils/api.py

@@ -2,11 +2,19 @@
 import time
 from contextlib import contextmanager
 from datetime import timedelta
-from typing import ContextManager, Generator, Literal, Optional, Union
+from functools import wraps
+from typing import (
+    Callable,
+    Generator,
+    Literal,
+    Optional,
+    TypeVar,
+    Union,
+)
 
 import httpx
 from pydantic import BaseModel, Field, TypeAdapter, ValidationError
-from typing_extensions import Annotated
+from typing_extensions import Annotated, ParamSpec
 
 from fastapi_cloud_cli import __version__
 from fastapi_cloud_cli.config import Settings
@@ -22,6 +30,10 @@ class BuildLogError(Exception):
     pass
 
 
+class TooManyRetriesError(Exception):
+    pass
+
+
 class BuildLogLineGeneric(BaseModel):
     type: Literal["complete", "failed", "timeout", "heartbeat"]
     id: Optional[str] = None
@@ -81,18 +93,39 @@ def _backoff() -> None:
             ) from error
 
 
+P = ParamSpec("P")
+T = TypeVar("T")
+
+
 def attempts(
     total_attempts: int = 3, timeout: timedelta = timedelta(minutes=5)
-) -> Generator[ContextManager[None], None, None]:
-    start = time.monotonic()
+) -> Callable[
+    [Callable[P, Generator[T, None, None]]], Callable[P, Generator[T, None, None]]
+]:
+    def decorator(
+        func: Callable[P, Generator[T, None, None]],
+    ) -> Callable[P, Generator[T, None, None]]:
+        @wraps(func)
+        def wrapper(*args: P.args, **kwargs: P.kwargs) -> Generator[T, None, None]:
+            start = time.monotonic()
 
-    for attempt_number in range(total_attempts):
-        if time.monotonic() - start > timeout.total_seconds():
-            raise TimeoutError(
-                "Build log streaming timed out after %ds", timeout.total_seconds()
-            )
+            for attempt_number in range(total_attempts):
+                if time.monotonic() - start > timeout.total_seconds():
+                    raise TimeoutError(
+                        "Build log streaming timed out after %ds",
+                        timeout.total_seconds(),
+                    )
+
+                with attempt(attempt_number):
+                    yield from func(*args, **kwargs)
+                    # If we get here without exception, the generator completed successfully
+                    return
+
+            raise TooManyRetriesError(f"Failed after {total_attempts} attempts")
 
-        yield attempt(attempt_number)
+        return wrapper
+
+    return decorator
 
 
 class APIClient(httpx.Client):
@@ -110,54 +143,47 @@ def __init__(self) -> None:
             },
         )
 
+    @attempts(BUILD_LOG_MAX_RETRIES, BUILD_LOG_TIMEOUT)
     def stream_build_logs(
         self, deployment_id: str
     ) -> Generator[BuildLogLine, None, None]:
         last_id = None
 
-        for attempt in attempts(BUILD_LOG_MAX_RETRIES, BUILD_LOG_TIMEOUT):
-            with attempt:
-                while True:
-                    params = {"last_id": last_id} if last_id else None
-
-                    with self.stream(
-                        "GET",
-                        f"/deployments/{deployment_id}/build-logs",
-                        timeout=60,
-                        params=params,
-                    ) as response:
-                        response.raise_for_status()
-
-                        for line in response.iter_lines():
-                            if not line or not line.strip():
-                                continue
+        while True:
+            params = {"last_id": last_id} if last_id else None
 
-                            if log_line := self._parse_log_line(line):
-                                if log_line.id:
-                                    last_id = log_line.id
+            with self.stream(
+                "GET",
+                f"/deployments/{deployment_id}/build-logs",
+                timeout=60,
+                params=params,
+            ) as response:
+                response.raise_for_status()
 
-                                if log_line.type == "message":
-                                    yield log_line
+                for line in response.iter_lines():
+                    if not line or not line.strip():
+                        continue
 
-                                if log_line.type in ("complete", "failed"):
-                                    yield log_line
+                    if log_line := self._parse_log_line(line):
+                        if log_line.id:
+                            last_id = log_line.id
 
-                                    return
+                        if log_line.type == "message":
+                            yield log_line
 
-                                if log_line.type == "timeout":
-                                    logger.debug("Received timeout; reconnecting")
-                                    break  # Breaks for loop to reconnect
+                        if log_line.type in ("complete", "failed"):
+                            yield log_line
+                            return
 
-                        else:  # Only triggered if the for loop is not broken
-                            logger.debug(
-                                "Connection closed by server unexpectedly; attempting to reconnect"
-                            )
-                            break
+                        if log_line.type == "timeout":
+                            logger.debug("Received timeout; reconnecting")
+                            break  # Breaks for loop to reconnect
+                else:
+                    logger.debug("Connection closed by server unexpectedly; will retry")
 
-                    time.sleep(0.5)
+                    raise httpx.NetworkError("Connection closed without terminal state")
 
-        # Exhausted retries without getting any response
-        raise BuildLogError(f"Failed after {BUILD_LOG_MAX_RETRIES} attempts")
+            time.sleep(0.5)
 
     def _parse_log_line(self, line: str) -> Optional[BuildLogLine]:
         try:

tests/test_api_client.py

@@ -13,6 +13,7 @@
     APIClient,
     BuildLogError,
     BuildLogLineMessage,
+    TooManyRetriesError,
 )
 from tests.utils import build_logs_response
 
@@ -254,7 +255,7 @@ def test_stream_build_logs_max_retries_exceeded(
 
     with patch("time.sleep"):
         with pytest.raises(
-            BuildLogError, match=f"Failed after {BUILD_LOG_MAX_RETRIES} attempts"
+            TooManyRetriesError, match=f"Failed after {BUILD_LOG_MAX_RETRIES} attempts"
         ):
             list(client.stream_build_logs(deployment_id))
 
@@ -341,7 +342,7 @@ def test_stream_build_logs_connection_closed_without_complete_failed_or_timeout(
 
     logs = client.stream_build_logs(deployment_id)
 
-    with pytest.raises(BuildLogError, match="Failed after"):
+    with patch("time.sleep"), pytest.raises(TooManyRetriesError, match="Failed after"):
         for _ in range(BUILD_LOG_MAX_RETRIES + 1):
             next(logs)
 
@@ -367,5 +368,5 @@ def responses(request: httpx.Request, route: respx.Route) -> Response:
 
     logs_route.mock(side_effect=responses)
 
-    with pytest.raises(TimeoutError, match="timed out"):
+    with patch("time.sleep"), pytest.raises(TimeoutError, match="timed out"):
         list(client.stream_build_logs(deployment_id))