Add retries for asyncio

Author: jhamon

pinecone/openapi_support/rest_aiohttp.py

@@ -9,6 +9,8 @@ class AiohttpRestClient(RestClientInterface):
     def __init__(self, configuration: Configuration) -> None:
         try:
             import aiohttp
+            from aiohttp_retry import RetryClient
+            from .retry_aiohttp import JitterRetry
         except ImportError:
             raise ImportError(
                 "Additional dependencies are required to use Pinecone with asyncio. Include these extra dependencies in your project by installing `pinecone[asyncio]`."
@@ -28,8 +30,18 @@ def __init__(self, configuration: Configuration) -> None:
         else:
             self._session = aiohttp.ClientSession(connector=conn)
 
+        jitter_retry = JitterRetry(
+            attempts=5,
+            start_timeout=0.1,
+            max_timeout=3.0,
+            statuses={500, 502, 503, 504},
+            methods=None,  # retry on all methods
+            exceptions={aiohttp.ClientError, aiohttp.ServerDisconnectedError},
+        )
+        self._retry_client = RetryClient(client_session=self._session, retry_options=jitter_retry)
+
     async def close(self):
-        await self._session.close()
+        await self._retry_client.close()
 
     async def request(
         self,
@@ -48,7 +60,7 @@ async def request(
         if "application/x-ndjson" in headers.get("Content-Type", "").lower():
             ndjson_data = "\n".join(json.dumps(record) for record in body)
 
-            async with self._session.request(
+            async with self._retry_client.request(
                 method, url, params=query_params, headers=headers, data=ndjson_data
             ) as resp:
                 content = await resp.read()
@@ -57,7 +69,7 @@ async def request(
                 )
 
         else:
-            async with self._session.request(
+            async with self._retry_client.request(
                 method, url, params=query_params, headers=headers, json=body
             ) as resp:
                 content = await resp.read()

pinecone/openapi_support/rest_urllib3.py

@@ -8,7 +8,7 @@
 from .rest_utils import raise_exceptions_or_return, RESTResponse, RestClientInterface
 
 import urllib3
-from .retries import JitterRetry
+from .retry_urllib3 import JitterRetry
 from .exceptions import PineconeApiException, PineconeApiValueError
 
 

pinecone/openapi_support/retries.py

@@ -57,6 +57,7 @@ protoc-gen-openapiv2 = {version = "^0.0.1", optional = true }
 pinecone-plugin-interface = "^0.0.7"
 python-dateutil = ">=2.5.3"
 aiohttp = { version = ">=3.9.0", optional = true }
+aiohttp-retry = { version = "^2.9.1", optional = true }
 
 [tool.poetry.group.types]
 optional = true
@@ -102,10 +103,9 @@ vprof = "^0.38"
 tuna = "^0.5.11"
 python-dotenv = "^1.1.0"
 
-
 [tool.poetry.extras]
 grpc = ["grpcio", "googleapis-common-protos", "lz4", "protobuf", "protoc-gen-openapiv2"]
-asyncio = ["aiohttp"]
+asyncio = ["aiohttp", "aiohttp-retry"]
 
 [build-system]
 requires = ["poetry-core"]

poetry.lock

@@ -2,7 +2,7 @@
 from unittest.mock import patch, MagicMock
 from urllib3.exceptions import MaxRetryError
 from urllib3.util.retry import Retry
-from pinecone.openapi_support.retries import JitterRetry
+from pinecone.openapi_support.retry_urllib3 import JitterRetry
 
 
 def test_jitter_retry_backoff():