Skip to content

Commit 87db6bd

Browse files
committed
Add retries for asyncio
1 parent 24296eb commit 87db6bd

File tree

6 files changed

+35
-30
lines changed

6 files changed

+35
-30
lines changed

pinecone/openapi_support/rest_aiohttp.py

Lines changed: 15 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,8 @@ class AiohttpRestClient(RestClientInterface):
99
def __init__(self, configuration: Configuration) -> None:
1010
try:
1111
import aiohttp
12+
from aiohttp_retry import RetryClient
13+
from .retry_aiohttp import JitterRetry
1214
except ImportError:
1315
raise ImportError(
1416
"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:
2830
else:
2931
self._session = aiohttp.ClientSession(connector=conn)
3032

33+
jitter_retry = JitterRetry(
34+
attempts=5,
35+
start_timeout=0.1,
36+
max_timeout=3.0,
37+
statuses={500, 502, 503, 504},
38+
methods=None, # retry on all methods
39+
exceptions={aiohttp.ClientError, aiohttp.ServerDisconnectedError},
40+
)
41+
self._retry_client = RetryClient(client_session=self._session, retry_options=jitter_retry)
42+
3143
async def close(self):
32-
await self._session.close()
44+
await self._retry_client.close()
3345

3446
async def request(
3547
self,
@@ -48,7 +60,7 @@ async def request(
4860
if "application/x-ndjson" in headers.get("Content-Type", "").lower():
4961
ndjson_data = "\n".join(json.dumps(record) for record in body)
5062

51-
async with self._session.request(
63+
async with self._retry_client.request(
5264
method, url, params=query_params, headers=headers, data=ndjson_data
5365
) as resp:
5466
content = await resp.read()
@@ -57,7 +69,7 @@ async def request(
5769
)
5870

5971
else:
60-
async with self._session.request(
72+
async with self._retry_client.request(
6173
method, url, params=query_params, headers=headers, json=body
6274
) as resp:
6375
content = await resp.read()

pinecone/openapi_support/rest_urllib3.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@
88
from .rest_utils import raise_exceptions_or_return, RESTResponse, RestClientInterface
99

1010
import urllib3
11-
from .retries import JitterRetry
11+
from .retry_urllib3 import JitterRetry
1212
from .exceptions import PineconeApiException, PineconeApiValueError
1313

1414

pinecone/openapi_support/retries.py

Lines changed: 0 additions & 21 deletions
This file was deleted.

poetry.lock

Lines changed: 16 additions & 2 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

pyproject.toml

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,7 @@ protoc-gen-openapiv2 = {version = "^0.0.1", optional = true }
5757
pinecone-plugin-interface = "^0.0.7"
5858
python-dateutil = ">=2.5.3"
5959
aiohttp = { version = ">=3.9.0", optional = true }
60+
aiohttp-retry = { version = "^2.9.1", optional = true }
6061

6162
[tool.poetry.group.types]
6263
optional = true
@@ -102,10 +103,9 @@ vprof = "^0.38"
102103
tuna = "^0.5.11"
103104
python-dotenv = "^1.1.0"
104105

105-
106106
[tool.poetry.extras]
107107
grpc = ["grpcio", "googleapis-common-protos", "lz4", "protobuf", "protoc-gen-openapiv2"]
108-
asyncio = ["aiohttp"]
108+
asyncio = ["aiohttp", "aiohttp-retry"]
109109

110110
[build-system]
111111
requires = ["poetry-core"]

tests/unit/openapi_support/test_retries.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22
from unittest.mock import patch, MagicMock
33
from urllib3.exceptions import MaxRetryError
44
from urllib3.util.retry import Retry
5-
from pinecone.openapi_support.retries import JitterRetry
5+
from pinecone.openapi_support.retry_urllib3 import JitterRetry
66

77

88
def test_jitter_retry_backoff():

0 commit comments

Comments
 (0)