in sync

Author: sanderegg

packages/service-library/src/servicelib/fastapi/long_running_tasks/client.py

@@ -3,6 +3,7 @@
 """
 
 import asyncio
+import logging
 from collections.abc import AsyncGenerator, Coroutine
 from dataclasses import dataclass
 from typing import Any, Final, TypeAlias
@@ -14,6 +15,7 @@
 from tenacity import (
     AsyncRetrying,
     TryAgain,
+    before_sleep_log,
     retry,
     retry_if_exception_type,
     stop_after_delay,
@@ -33,20 +35,23 @@
 from ._client import DEFAULT_HTTP_REQUESTS_TIMEOUT, Client, setup
 from ._context_manager import periodic_task_result
 
+_logger = logging.getLogger(__name__)
+
 RequestBody: TypeAlias = Any
 
 _MINUTE: Final[int] = 60  # in secs
 _HOUR: Final[int] = 60 * _MINUTE  # in secs
 _DEFAULT_POLL_INTERVAL_S: Final[float] = 1
-_DEFAULT_AIOHTTP_RETRY_POLICY: dict[str, Any] = {
+_DEFAULT_FASTAPI_RETRY_POLICY: dict[str, Any] = {
     "retry": retry_if_exception_type(httpx.RequestError),
     "wait": wait_random_exponential(max=20),
     "stop": stop_after_delay(60),
     "reraise": True,
+    "before_sleep": before_sleep_log(_logger, logging.INFO),
 }
 
 
-@retry(**_DEFAULT_AIOHTTP_RETRY_POLICY)
+@retry(**_DEFAULT_FASTAPI_RETRY_POLICY)
 async def _start(
     session: httpx.AsyncClient, url: URL, json: RequestBody | None
 ) -> TaskGet:
@@ -56,7 +61,7 @@ async def _start(
     return TaskGet.model_validate(data)
 
 
-@retry(**_DEFAULT_AIOHTTP_RETRY_POLICY)
+@retry(**_DEFAULT_FASTAPI_RETRY_POLICY)
 async def _wait_for_completion(
     session: httpx.AsyncClient,
     task_id: TaskId,
@@ -68,6 +73,7 @@ async def _wait_for_completion(
             stop=stop_after_delay(client_timeout),
             reraise=True,
             retry=retry_if_exception_type(TryAgain),
+            before_sleep=before_sleep_log(_logger, logging.DEBUG),
         ):
             with attempt:
                 response = await session.get(f"{status_url}")
@@ -93,7 +99,7 @@ async def _wait_for_completion(
         raise TimeoutError(msg) from exc
 
 
-@retry(**_DEFAULT_AIOHTTP_RETRY_POLICY)
+@retry(**_DEFAULT_FASTAPI_RETRY_POLICY)
 async def _task_result(session: httpx.AsyncClient, result_url: URL) -> Any:
     response = await session.get(f"{result_url}", params={"return_exception": True})
     response.raise_for_status()
@@ -102,7 +108,7 @@ async def _task_result(session: httpx.AsyncClient, result_url: URL) -> Any:
     return None
 
 
-@retry(**_DEFAULT_AIOHTTP_RETRY_POLICY)
+@retry(**_DEFAULT_FASTAPI_RETRY_POLICY)
 async def _abort_task(session: httpx.AsyncClient, abort_url: URL) -> None:
     response = await session.delete(f"{abort_url}")
     response.raise_for_status()