Skip to content

Commit b670bec

Browse files
create_from_paths: Allow complete retry
- Introduced `send_request_once` and `run_with_retry` for consistent request execution and retry handling. - Implemented retries for `create_from_paths` in both sync and async clients, covering 429, transport, and timeout errors. - Added concurrency control using semaphores for async file info fetches after uploads. - Updated tests to validate retry behavior, including various error scenarios and proper retry counts. Assisted-by: Codex
1 parent b45a81f commit b670bec

2 files changed

Lines changed: 381 additions & 38 deletions

File tree

src/pdfrest/client.py

Lines changed: 120 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -896,6 +896,22 @@ def _send_request(self, request: _RequestModel) -> Any:
896896
should_continue=stream_retry_checker,
897897
)
898898

899+
def send_request_once(self, request: _RequestModel) -> Any:
900+
return self._perform_request(self._client, request)
901+
902+
def run_with_retry(
903+
self,
904+
func: Callable[[], ReturnType],
905+
*,
906+
operation: str,
907+
should_continue: Callable[[PdfRestError], bool] | None = None,
908+
) -> ReturnType:
909+
return self._execute_with_retry(
910+
func,
911+
operation=operation,
912+
should_continue=should_continue,
913+
)
914+
899915
def _perform_request(
900916
self, http_client: httpx.Client, request: _RequestModel
901917
) -> Any:
@@ -1143,6 +1159,22 @@ async def _send_request(self, request: _RequestModel) -> Any:
11431159
should_continue=stream_retry_checker,
11441160
)
11451161

1162+
async def send_request_once(self, request: _RequestModel) -> Any:
1163+
return await self._perform_request(self._client, request)
1164+
1165+
async def run_with_retry(
1166+
self,
1167+
func: Callable[[], Awaitable[ReturnType]],
1168+
*,
1169+
operation: str,
1170+
should_continue: Callable[[PdfRestError], bool] | None = None,
1171+
) -> ReturnType:
1172+
return await self._execute_with_retry(
1173+
func,
1174+
operation=operation,
1175+
should_continue=should_continue,
1176+
)
1177+
11461178
async def _perform_request(
11471179
self, http_client: httpx.AsyncClient, request: _RequestModel
11481180
) -> Any:
@@ -1441,25 +1473,44 @@ def create_from_paths(
14411473
closed once the request completes.
14421474
"""
14431475
normalized_paths = _normalize_path_inputs(file_paths)
1444-
1445-
with ExitStack() as stack:
1446-
upload_specs: list[FileTypes] = []
1447-
for spec in normalized_paths:
1448-
path, content_type, headers = _parse_path_spec(spec)
1449-
file_obj = stack.enter_context(path.open("rb"))
1450-
filename = path.name
1451-
if headers:
1452-
upload_specs.append((filename, file_obj, content_type, headers))
1453-
elif content_type is not None:
1454-
upload_specs.append((filename, file_obj, content_type))
1455-
else:
1456-
upload_specs.append((filename, file_obj))
1457-
return self.create(
1458-
upload_specs,
1459-
extra_query=extra_query,
1460-
extra_headers=extra_headers,
1461-
timeout=timeout,
1462-
)
1476+
path_specs = [_parse_path_spec(spec) for spec in normalized_paths]
1477+
1478+
def attempt() -> list[PdfRestFile]:
1479+
with ExitStack() as stack:
1480+
upload_specs: list[tuple[str, FileTypes]] = []
1481+
for path, content_type, headers in path_specs:
1482+
file_obj = stack.enter_context(path.open("rb"))
1483+
filename = path.name
1484+
normalized_spec: FileTypes
1485+
if headers:
1486+
normalized_spec = (filename, file_obj, content_type, headers)
1487+
elif content_type is not None:
1488+
normalized_spec = (filename, file_obj, content_type)
1489+
else:
1490+
normalized_spec = (filename, file_obj)
1491+
upload_specs.append((FILE_UPLOAD_FIELD_NAME, normalized_spec))
1492+
1493+
request = self._client.prepare_request(
1494+
"POST",
1495+
"/upload",
1496+
files=upload_specs,
1497+
extra_query=extra_query,
1498+
extra_headers=extra_headers,
1499+
timeout=timeout,
1500+
)
1501+
payload = self._client.send_request_once(request)
1502+
file_ids = _extract_uploaded_file_ids(payload)
1503+
return [
1504+
self._client.fetch_file_info(
1505+
file_id,
1506+
extra_query=extra_query,
1507+
extra_headers=extra_headers,
1508+
timeout=timeout,
1509+
)
1510+
for file_id in file_ids
1511+
]
1512+
1513+
return self._client.run_with_retry(attempt, operation="POST /upload (paths)")
14631514

14641515
def create_from_urls(
14651516
self,
@@ -1681,25 +1732,56 @@ async def create_from_paths(
16811732
closed once the request completes.
16821733
"""
16831734
normalized_paths = _normalize_path_inputs(file_paths)
1684-
1685-
with ExitStack() as stack:
1686-
upload_specs: list[FileTypes] = []
1687-
for spec in normalized_paths:
1688-
path, content_type, headers = _parse_path_spec(spec)
1689-
file_obj = stack.enter_context(path.open("rb"))
1690-
filename = path.name
1691-
if headers:
1692-
upload_specs.append((filename, file_obj, content_type, headers))
1693-
elif content_type is not None:
1694-
upload_specs.append((filename, file_obj, content_type))
1695-
else:
1696-
upload_specs.append((filename, file_obj))
1697-
return await self.create(
1698-
upload_specs,
1699-
extra_query=extra_query,
1700-
extra_headers=extra_headers,
1701-
timeout=timeout,
1702-
)
1735+
path_specs = [_parse_path_spec(spec) for spec in normalized_paths]
1736+
1737+
async def attempt() -> list[PdfRestFile]:
1738+
with ExitStack() as stack:
1739+
upload_specs: list[tuple[str, FileTypes]] = []
1740+
for path, content_type, headers in path_specs:
1741+
file_obj = stack.enter_context(path.open("rb"))
1742+
filename = path.name
1743+
normalized_spec: FileTypes
1744+
if headers:
1745+
normalized_spec = (filename, file_obj, content_type, headers)
1746+
elif content_type is not None:
1747+
normalized_spec = (filename, file_obj, content_type)
1748+
else:
1749+
normalized_spec = (filename, file_obj)
1750+
upload_specs.append((FILE_UPLOAD_FIELD_NAME, normalized_spec))
1751+
1752+
request = self._client.prepare_request(
1753+
"POST",
1754+
"/upload",
1755+
files=upload_specs,
1756+
extra_query=extra_query,
1757+
extra_headers=extra_headers,
1758+
timeout=timeout,
1759+
)
1760+
payload = await self._client.send_request_once(request)
1761+
file_ids = _extract_uploaded_file_ids(payload)
1762+
results: list[PdfRestFile] = []
1763+
semaphore = asyncio.Semaphore(self._concurrency_limit)
1764+
1765+
async def throttled_fetch(file_id: str) -> PdfRestFile:
1766+
async with semaphore:
1767+
return await self._client.fetch_file_info(
1768+
file_id,
1769+
extra_query=extra_query,
1770+
extra_headers=extra_headers,
1771+
timeout=timeout,
1772+
)
1773+
1774+
if file_ids:
1775+
results = list(
1776+
await asyncio.gather(
1777+
*(throttled_fetch(fid) for fid in file_ids)
1778+
)
1779+
)
1780+
return results
1781+
1782+
return await self._client.run_with_retry(
1783+
attempt, operation="POST /upload (paths)"
1784+
)
17031785

17041786
async def create_from_urls(
17051787
self,

0 commit comments

Comments
 (0)