Skip to content
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
75 changes: 58 additions & 17 deletions src/sumo/wrapper/sumo_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -219,7 +219,12 @@ def _handle_invalid_shared_key(self):
)

@raise_for_status
def get(self, path: str, params: Optional[Dict] = None) -> httpx.Response:
def get(
self,
path: str,
params: Optional[Dict] = None,
retry_strategy: Optional[RetryStrategy] = None,
) -> httpx.Response:
"""Performs a GET-request to the Sumo API.

Args:
Expand Down Expand Up @@ -277,8 +282,9 @@ def _get():
timeout=self._timeout,
)

retryer = self._retry_strategy.make_retryer()

retryer = (
retry_strategy if retry_strategy else self._retry_strategy
).make_retryer()
return retryer(_get)

@raise_for_status
Expand All @@ -288,6 +294,7 @@ def post(
blob: Optional[bytes] = None,
json: Optional[dict] = None,
params: Optional[dict] = None,
retry_strategy: Optional[RetryStrategy] = None,
) -> httpx.Response:
"""Performs a POST-request to the Sumo API.

Expand Down Expand Up @@ -352,7 +359,9 @@ def _post():
timeout=self._timeout,
)

retryer = self._retry_strategy.make_retryer()
retryer = (
retry_strategy if retry_strategy else self._retry_strategy
).make_retryer()

return retryer(_post)

Expand All @@ -362,6 +371,7 @@ def put(
path: str,
blob: Optional[bytes] = None,
json: Optional[dict] = None,
retry_strategy: Optional[RetryStrategy] = None,
) -> httpx.Response:
"""Performs a PUT-request to the Sumo API.

Expand Down Expand Up @@ -401,13 +411,18 @@ def _put():
timeout=self._timeout,
)

retryer = self._retry_strategy.make_retryer()
retryer = (
retry_strategy if retry_strategy else self._retry_strategy
).make_retryer()

return retryer(_put)

@raise_for_status
def delete(
self, path: str, params: Optional[dict] = None
self,
path: str,
params: Optional[dict] = None,
retry_strategy: Optional[RetryStrategy] = None,
) -> httpx.Response:
"""Performs a DELETE-request to the Sumo API.

Expand Down Expand Up @@ -441,7 +456,9 @@ def _delete():
timeout=self._timeout,
)

retryer = self._retry_strategy.make_retryer()
retryer = (
retry_strategy if retry_strategy else self._retry_strategy
).make_retryer()

return retryer(_delete)

Expand All @@ -460,7 +477,10 @@ def _get_retry_details(self, response_in) -> Tuple[str, int]:
return location, retry_after

def poll(
self, response_in: httpx.Response, timeout=None
self,
response_in: httpx.Response,
timeout=None,
retry_strategy: Optional[RetryStrategy] = None,
) -> httpx.Response:
"""Poll a specific endpoint until a result is obtained.

Expand All @@ -474,7 +494,7 @@ def poll(
expiry = time.time() + timeout if timeout is not None else None
while True:
time.sleep(retry_after)
response = self.get(location)
response = self.get(location, retry_strategy=retry_strategy)
if response.status_code != 202:
return response
if expiry is not None and time.time() > expiry:
Expand Down Expand Up @@ -539,7 +559,10 @@ def client_for_case(self, case_uuid, interactive=False):

@raise_for_status_async
async def get_async(
self, path: str, params: Optional[dict] = None
self,
path: str,
params: Optional[dict] = None,
retry_strategy: Optional[RetryStrategy] = None,
) -> httpx.Response:
"""Performs an async GET-request to the Sumo API.

Expand Down Expand Up @@ -598,7 +621,9 @@ async def _get():
timeout=self._timeout,
)

retryer = self._retry_strategy.make_retryer_async()
retryer = (
retry_strategy if retry_strategy else self._retry_strategy
).make_retryer_async()

return await retryer(_get)

Expand All @@ -609,6 +634,7 @@ async def post_async(
blob: Optional[bytes] = None,
json: Optional[dict] = None,
params: Optional[dict] = None,
retry_strategy: Optional[RetryStrategy] = None,
) -> httpx.Response:
"""Performs an async POST-request to the Sumo API.

Expand Down Expand Up @@ -674,7 +700,9 @@ async def _post():
timeout=self._timeout,
)

retryer = self._retry_strategy.make_retryer_async()
retryer = (
retry_strategy if retry_strategy else self._retry_strategy
).make_retryer_async()

return await retryer(_post)

Expand All @@ -684,6 +712,7 @@ async def put_async(
path: str,
blob: Optional[bytes] = None,
json: Optional[dict] = None,
retry_strategy: Optional[RetryStrategy] = None,
) -> httpx.Response:
"""Performs an async PUT-request to the Sumo API.

Expand Down Expand Up @@ -723,13 +752,18 @@ async def _put():
timeout=self._timeout,
)

retryer = self._retry_strategy.make_retryer_async()
retryer = (
retry_strategy if retry_strategy else self._retry_strategy
).make_retryer_async()

return await retryer(_put)

@raise_for_status_async
async def delete_async(
self, path: str, params: Optional[dict] = None
self,
path: str,
params: Optional[dict] = None,
retry_strategy: Optional[RetryStrategy] = None,
) -> httpx.Response:
"""Performs an async DELETE-request to the Sumo API.

Expand Down Expand Up @@ -763,12 +797,17 @@ async def _delete():
timeout=self._timeout,
)

retryer = self._retry_strategy.make_retryer_async()
retryer = (
retry_strategy if retry_strategy else self._retry_strategy
).make_retryer_async()

return await retryer(_delete)

async def poll_async(
self, response_in: httpx.Response, timeout=None
self,
response_in: httpx.Response,
timeout=None,
retry_strategy: Optional[RetryStrategy] = None,
) -> httpx.Response:
"""Poll a specific endpoint until a result is obtained.

Expand All @@ -782,7 +821,9 @@ async def poll_async(
expiry = time.time() + timeout if timeout is not None else None
while True:
await asyncio.sleep(retry_after)
response = await self.get_async(location)
response = await self.get_async(
location, retry_strategy=retry_strategy
)
if response.status_code != 202:
return response
if expiry is not None and time.time() > expiry:
Expand Down
Loading