From 18aa8c925b3f7699c30a6202685447432df2da6c Mon Sep 17 00:00:00 2001 From: yummybomb <19238148+yummybomb@users.noreply.github.com> Date: Mon, 20 Jul 2026 20:59:38 +0000 Subject: [PATCH 1/6] feat: add complete audit log downloads --- src/kernel/__init__.py | 8 ++ src/kernel/lib/audit_log_download.py | 205 +++++++++++++++++++++++++++ src/kernel/resources/audit_logs.py | 106 +++++++++++++- tests/test_audit_log_download.py | 136 ++++++++++++++++++ 4 files changed, 454 insertions(+), 1 deletion(-) create mode 100644 src/kernel/lib/audit_log_download.py create mode 100644 tests/test_audit_log_download.py diff --git a/src/kernel/__init__.py b/src/kernel/__init__.py index 333d7018..af9462fb 100644 --- a/src/kernel/__init__.py +++ b/src/kernel/__init__.py @@ -51,6 +51,11 @@ app_registry, export_registry, ) +from .lib.audit_log_download import ( + AuditLogDownloadError, + AuditLogDownloadResult, + AuditLogDownloadProgress, +) __all__ = [ "types", @@ -105,6 +110,9 @@ "App", "app_registry", "export_registry", + "AuditLogDownloadError", + "AuditLogDownloadProgress", + "AuditLogDownloadResult", ] if not _t.TYPE_CHECKING: diff --git a/src/kernel/lib/audit_log_download.py b/src/kernel/lib/audit_log_download.py new file mode 100644 index 00000000..7f739c8e --- /dev/null +++ b/src/kernel/lib/audit_log_download.py @@ -0,0 +1,205 @@ +from __future__ import annotations + +import time +import hashlib +from typing import BinaryIO, Callable, Optional, Protocol, Awaitable +from dataclasses import dataclass + +import anyio +import httpx + +from .._exceptions import KernelError, APIStatusError + +_DOWNLOAD_ATTEMPTS = 7 +_MAX_RETRY_DELAY = 8.0 + + +class AuditLogDownloadError(KernelError): + pass + + +@dataclass(frozen=True) +class AuditLogDownloadResult: + bytes_written: int + chunks: int + rows: int + + +@dataclass(frozen=True) +class AuditLogDownloadProgress(AuditLogDownloadResult): + chunk_rows: int + + +class _SyncChunkResponse(Protocol): + @property + def headers(self) -> httpx.Headers: ... + + def read(self) -> bytes: ... + + def close(self) -> None: ... + + +class _AsyncChunkResponse(Protocol): + @property + def headers(self) -> httpx.Headers: ... + + async def read(self) -> bytes: ... + + async def close(self) -> None: ... + + +SyncFetchChunk = Callable[[Optional[str]], _SyncChunkResponse] +AsyncFetchChunk = Callable[[Optional[str]], Awaitable[_AsyncChunkResponse]] +ProgressCallback = Callable[[AuditLogDownloadProgress], None] + + +def download_audit_logs( + fetch_chunk: SyncFetchChunk, + destination: BinaryIO, + *, + on_progress: ProgressCallback | None = None, +) -> AuditLogDownloadResult: + if not callable(getattr(destination, "write", None)): + raise TypeError("audit log download destination must provide write()") + + cursor: str | None = None + result = AuditLogDownloadResult(bytes_written=0, chunks=0, rows=0) + while True: + body, headers = _fetch_verified_chunk(fetch_chunk, cursor) + chunk_rows, next_cursor, has_more = _parse_chunk_headers(headers, cursor) + _write_chunk(destination, body) + + cursor = next_cursor + result = AuditLogDownloadResult( + bytes_written=result.bytes_written + len(body), + chunks=result.chunks + 1, + rows=result.rows + chunk_rows, + ) + if on_progress is not None: + on_progress( + AuditLogDownloadProgress( + bytes_written=result.bytes_written, + chunks=result.chunks, + rows=result.rows, + chunk_rows=chunk_rows, + ) + ) + if not has_more: + return result + + +async def async_download_audit_logs( + fetch_chunk: AsyncFetchChunk, + destination: BinaryIO, + *, + on_progress: ProgressCallback | None = None, +) -> AuditLogDownloadResult: + if not callable(getattr(destination, "write", None)): + raise TypeError("audit log download destination must provide write()") + + cursor: str | None = None + result = AuditLogDownloadResult(bytes_written=0, chunks=0, rows=0) + while True: + body, headers = await _async_fetch_verified_chunk(fetch_chunk, cursor) + chunk_rows, next_cursor, has_more = _parse_chunk_headers(headers, cursor) + _write_chunk(destination, body) + + cursor = next_cursor + result = AuditLogDownloadResult( + bytes_written=result.bytes_written + len(body), + chunks=result.chunks + 1, + rows=result.rows + chunk_rows, + ) + if on_progress is not None: + on_progress( + AuditLogDownloadProgress( + bytes_written=result.bytes_written, + chunks=result.chunks, + rows=result.rows, + chunk_rows=chunk_rows, + ) + ) + if not has_more: + return result + + +def _fetch_verified_chunk(fetch_chunk: SyncFetchChunk, cursor: str | None) -> tuple[bytes, httpx.Headers]: + for attempt in range(1, _DOWNLOAD_ATTEMPTS + 1): + try: + response = fetch_chunk(cursor) + try: + body = response.read() + headers = httpx.Headers(response.headers) + finally: + response.close() + _verify_checksum(body, headers) + return body, headers + except Exception as error: + if attempt == _DOWNLOAD_ATTEMPTS or not _is_retryable(error): + raise + time.sleep(min(2 ** (attempt - 1), _MAX_RETRY_DELAY)) + raise AssertionError("unreachable") + + +async def _async_fetch_verified_chunk(fetch_chunk: AsyncFetchChunk, cursor: str | None) -> tuple[bytes, httpx.Headers]: + for attempt in range(1, _DOWNLOAD_ATTEMPTS + 1): + try: + response = await fetch_chunk(cursor) + try: + body = await response.read() + headers = httpx.Headers(response.headers) + finally: + await response.close() + _verify_checksum(body, headers) + return body, headers + except Exception as error: + if attempt == _DOWNLOAD_ATTEMPTS or not _is_retryable(error): + raise + await anyio.sleep(min(2 ** (attempt - 1), _MAX_RETRY_DELAY)) + raise AssertionError("unreachable") + + +def _verify_checksum(body: bytes, headers: httpx.Headers) -> None: + expected = headers.get("x-content-sha256") + if not expected: + raise AuditLogDownloadError("response missing X-Content-Sha256 header") + actual = hashlib.sha256(body).hexdigest() + if actual != expected: + raise AuditLogDownloadError(f"audit log chunk checksum mismatch (got {actual}, want {expected})") + + +def _parse_chunk_headers(headers: httpx.Headers, current_cursor: str | None) -> tuple[int, str | None, bool]: + has_more_value = headers.get("x-has-more") + if has_more_value not in {"true", "false"}: + raise AuditLogDownloadError("response missing or invalid X-Has-More header") + has_more = has_more_value == "true" + + row_count = headers.get("x-row-count") + try: + rows = int(row_count) if row_count is not None else -1 + except ValueError as error: + raise AuditLogDownloadError("response missing or invalid X-Row-Count header") from error + if rows < 0: + raise AuditLogDownloadError("response missing or invalid X-Row-Count header") + + next_cursor = headers.get("x-next-cursor") or None + if has_more and (not next_cursor or next_cursor == current_cursor): + raise AuditLogDownloadError("response has invalid X-Next-Cursor header") + if not has_more and next_cursor: + raise AuditLogDownloadError("response returned a cursor after the final chunk") + return rows, next_cursor, has_more + + +def _write_chunk(destination: BinaryIO, body: bytes) -> None: + remaining = memoryview(body) + while remaining: + written = destination.write(remaining) + if written <= 0 or written > len(remaining): + raise AuditLogDownloadError("audit log download destination performed a short write") + remaining = remaining[written:] + + +def _is_retryable(error: Exception) -> bool: + if isinstance(error, APIStatusError): + return error.status_code == 429 or error.status_code >= 500 + return True diff --git a/src/kernel/resources/audit_logs.py b/src/kernel/resources/audit_logs.py index 2970e153..5b50ad59 100644 --- a/src/kernel/resources/audit_logs.py +++ b/src/kernel/resources/audit_logs.py @@ -2,7 +2,7 @@ from __future__ import annotations -from typing import Union +from typing import Union, BinaryIO from datetime import datetime from typing_extensions import Literal @@ -30,6 +30,12 @@ from ..pagination import SyncPageTokenPagination, AsyncPageTokenPagination from .._base_client import AsyncPaginator, make_request_options from ..types.audit_log_entry import AuditLogEntry +from ..lib.audit_log_download import ( + ProgressCallback, + AuditLogDownloadResult, + download_audit_logs, + async_download_audit_logs, +) __all__ = ["AuditLogsResource", "AsyncAuditLogsResource"] @@ -222,6 +228,55 @@ def export_chunk( cast_to=BinaryAPIResponse, ) + def download( + self, + *, + to: BinaryIO, + end: Union[str, datetime], + start: Union[str, datetime], + auth_strategy: str | Omit = omit, + exclude_method: SequenceNotStr[str] | Omit = omit, + format: Literal["jsonl", "jsonl.gz"] | Omit = omit, + limit: int | Omit = omit, + method: str | Omit = omit, + search: str | Omit = omit, + search_user_id: SequenceNotStr[str] | Omit = omit, + service: str | Omit = omit, + on_progress: ProgressCallback | None = None, + extra_headers: Headers | None = None, + extra_query: Query | None = None, + extra_body: Body | None = None, + timeout: float | httpx.Timeout | None | NotGiven = not_given, + ) -> AuditLogDownloadResult: + """Download a complete audit log export to a writable binary destination. + + The SDK verifies every chunk and retries transient transfer failures. It + does not close the destination. + """ + + resource = self._client.with_options(max_retries=0).audit_logs + + def fetch_chunk(cursor: str | None) -> BinaryAPIResponse: + return resource.export_chunk( + end=end, + start=start, + auth_strategy=auth_strategy, + cursor=cursor if cursor is not None else omit, + exclude_method=exclude_method, + format=format, + limit=limit, + method=method, + search=search, + search_user_id=search_user_id, + service=service, + extra_headers=extra_headers, + extra_query=extra_query, + extra_body=extra_body, + timeout=timeout, + ) + + return download_audit_logs(fetch_chunk, to, on_progress=on_progress) + class AsyncAuditLogsResource(AsyncAPIResource): """Read audit log records for the authenticated organization.""" @@ -411,6 +466,55 @@ async def export_chunk( cast_to=AsyncBinaryAPIResponse, ) + async def download( + self, + *, + to: BinaryIO, + end: Union[str, datetime], + start: Union[str, datetime], + auth_strategy: str | Omit = omit, + exclude_method: SequenceNotStr[str] | Omit = omit, + format: Literal["jsonl", "jsonl.gz"] | Omit = omit, + limit: int | Omit = omit, + method: str | Omit = omit, + search: str | Omit = omit, + search_user_id: SequenceNotStr[str] | Omit = omit, + service: str | Omit = omit, + on_progress: ProgressCallback | None = None, + extra_headers: Headers | None = None, + extra_query: Query | None = None, + extra_body: Body | None = None, + timeout: float | httpx.Timeout | None | NotGiven = not_given, + ) -> AuditLogDownloadResult: + """Download a complete audit log export to a writable binary destination. + + The SDK verifies every chunk and retries transient transfer failures. It + does not close the destination. + """ + + resource = self._client.with_options(max_retries=0).audit_logs + + async def fetch_chunk(cursor: str | None) -> AsyncBinaryAPIResponse: + return await resource.export_chunk( + end=end, + start=start, + auth_strategy=auth_strategy, + cursor=cursor if cursor is not None else omit, + exclude_method=exclude_method, + format=format, + limit=limit, + method=method, + search=search, + search_user_id=search_user_id, + service=service, + extra_headers=extra_headers, + extra_query=extra_query, + extra_body=extra_body, + timeout=timeout, + ) + + return await async_download_audit_logs(fetch_chunk, to, on_progress=on_progress) + class AuditLogsResourceWithRawResponse: def __init__(self, audit_logs: AuditLogsResource) -> None: diff --git a/tests/test_audit_log_download.py b/tests/test_audit_log_download.py new file mode 100644 index 00000000..9c1a8c59 --- /dev/null +++ b/tests/test_audit_log_download.py @@ -0,0 +1,136 @@ +from __future__ import annotations + +import hashlib +from io import BytesIO + +import httpx +import pytest +from respx import MockRouter + +from kernel import Kernel, AsyncKernel, AuditLogDownloadError +from kernel.lib import audit_log_download + + +def chunk_response(body: bytes, *, rows: int, has_more: bool, next_cursor: str | None = None) -> httpx.Response: + headers = { + "x-content-sha256": hashlib.sha256(body).hexdigest(), + "x-has-more": str(has_more).lower(), + "x-row-count": str(rows), + } + if next_cursor is not None: + headers["x-next-cursor"] = next_cursor + return httpx.Response(200, content=body, headers=headers) + + +def test_download_writes_verified_chunks(client: Kernel, respx_mock: MockRouter) -> None: + responses = iter( + [ + chunk_response(b"first", rows=2, has_more=True, next_cursor="next"), + chunk_response(b"second", rows=1, has_more=False), + ] + ) + cursors: list[str | None] = [] + + def respond(request: httpx.Request) -> httpx.Response: + cursors.append(request.url.params.get("cursor")) + return next(responses) + + respx_mock.get("/audit-logs/export/chunk").mock(side_effect=respond) + destination = BytesIO() + progress: list[tuple[int, int, int, int]] = [] + + result = client.audit_logs.download( + to=destination, + start="2026-06-01T00:00:00Z", + end="2026-06-02T00:00:00Z", + format="jsonl.gz", + on_progress=lambda update: progress.append( + (update.bytes_written, update.chunks, update.rows, update.chunk_rows) + ), + ) + + assert destination.getvalue() == b"firstsecond" + assert result.bytes_written == 11 + assert result.chunks == 2 + assert result.rows == 3 + assert progress == [(5, 1, 2, 2), (11, 2, 3, 1)] + assert cursors == [None, "next"] + + +async def test_async_download_writes_verified_chunks(async_client: AsyncKernel, respx_mock: MockRouter) -> None: + respx_mock.get("/audit-logs/export/chunk").mock( + side_effect=[ + chunk_response(b"first", rows=2, has_more=True, next_cursor="next"), + chunk_response(b"second", rows=1, has_more=False), + ] + ) + destination = BytesIO() + + result = await async_client.audit_logs.download( + to=destination, + start="2026-06-01T00:00:00Z", + end="2026-06-02T00:00:00Z", + ) + + assert destination.getvalue() == b"firstsecond" + assert result.bytes_written == 11 + assert result.chunks == 2 + assert result.rows == 3 + + +def test_download_retries_checksum_mismatch( + client: Kernel, respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + def no_sleep(_: float) -> None: + pass + + monkeypatch.setattr(audit_log_download.time, "sleep", no_sleep) + bad = chunk_response(b"bad", rows=1, has_more=False) + bad.headers["x-content-sha256"] = hashlib.sha256(b"good").hexdigest() + route = respx_mock.get("/audit-logs/export/chunk").mock( + side_effect=[bad, chunk_response(b"good", rows=1, has_more=False)] + ) + destination = BytesIO() + + client.audit_logs.download( + to=destination, + start="2026-06-01T00:00:00Z", + end="2026-06-02T00:00:00Z", + ) + + assert destination.getvalue() == b"good" + assert route.call_count == 2 + + +def test_download_owns_http_retries(client: Kernel, respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch) -> None: + delays: list[float] = [] + monkeypatch.setattr(audit_log_download.time, "sleep", delays.append) + route = respx_mock.get("/audit-logs/export/chunk").mock( + side_effect=[ + httpx.Response(500, json={"message": "temporary failure"}), + chunk_response(b"good", rows=1, has_more=False), + ] + ) + + client.audit_logs.download( + to=BytesIO(), + start="2026-06-01T00:00:00Z", + end="2026-06-02T00:00:00Z", + ) + + assert route.call_count == 2 + assert delays == [1] + + +def test_download_rejects_invalid_cursor_before_writing(client: Kernel, respx_mock: MockRouter) -> None: + respx_mock.get("/audit-logs/export/chunk").mock(return_value=chunk_response(b"chunk", rows=1, has_more=True)) + destination = BytesIO() + + with pytest.raises(AuditLogDownloadError, match="response has invalid X-Next-Cursor header"): + client.audit_logs.download( + to=destination, + start="2026-06-01T00:00:00Z", + end="2026-06-02T00:00:00Z", + ) + + assert destination.getvalue() == b"" From 8174c85167700e283e254d0bcae75f8482abb159 Mon Sep 17 00:00:00 2001 From: yummybomb <19238148+yummybomb@users.noreply.github.com> Date: Tue, 21 Jul 2026 14:44:45 +0000 Subject: [PATCH 2/6] Harden audit log downloads --- src/kernel/lib/audit_log_download.py | 70 +++++++++++++++--------- src/kernel/resources/audit_logs.py | 35 ++++++++---- tests/test_audit_log_download.py | 80 +++++++++++++++++++++++++--- 3 files changed, 144 insertions(+), 41 deletions(-) diff --git a/src/kernel/lib/audit_log_download.py b/src/kernel/lib/audit_log_download.py index 7f739c8e..783317e5 100644 --- a/src/kernel/lib/audit_log_download.py +++ b/src/kernel/lib/audit_log_download.py @@ -2,15 +2,17 @@ import time import hashlib +import inspect from typing import BinaryIO, Callable, Optional, Protocol, Awaitable from dataclasses import dataclass import anyio import httpx +from anyio import to_thread -from .._exceptions import KernelError, APIStatusError +from .._exceptions import KernelError -_DOWNLOAD_ATTEMPTS = 7 +_DEFAULT_MAX_TRANSFER_RETRIES = 6 _MAX_RETRY_DELAY = 8.0 @@ -51,6 +53,7 @@ async def close(self) -> None: ... SyncFetchChunk = Callable[[Optional[str]], _SyncChunkResponse] AsyncFetchChunk = Callable[[Optional[str]], Awaitable[_AsyncChunkResponse]] ProgressCallback = Callable[[AuditLogDownloadProgress], None] +AsyncProgressCallback = Callable[[AuditLogDownloadProgress], Optional[Awaitable[None]]] def download_audit_logs( @@ -58,15 +61,22 @@ def download_audit_logs( destination: BinaryIO, *, on_progress: ProgressCallback | None = None, + max_transfer_retries: int = _DEFAULT_MAX_TRANSFER_RETRIES, ) -> AuditLogDownloadResult: if not callable(getattr(destination, "write", None)): raise TypeError("audit log download destination must provide write()") + _validate_max_transfer_retries(max_transfer_retries) cursor: str | None = None result = AuditLogDownloadResult(bytes_written=0, chunks=0, rows=0) + seen_cursors: set[str] = set() while True: - body, headers = _fetch_verified_chunk(fetch_chunk, cursor) + body, headers = _fetch_verified_chunk(fetch_chunk, cursor, max_transfer_retries) chunk_rows, next_cursor, has_more = _parse_chunk_headers(headers, cursor) + if has_more and next_cursor is not None: + if next_cursor in seen_cursors: + raise AuditLogDownloadError("response repeated X-Next-Cursor header") + seen_cursors.add(next_cursor) _write_chunk(destination, body) cursor = next_cursor @@ -92,17 +102,24 @@ async def async_download_audit_logs( fetch_chunk: AsyncFetchChunk, destination: BinaryIO, *, - on_progress: ProgressCallback | None = None, + on_progress: AsyncProgressCallback | None = None, + max_transfer_retries: int = _DEFAULT_MAX_TRANSFER_RETRIES, ) -> AuditLogDownloadResult: if not callable(getattr(destination, "write", None)): raise TypeError("audit log download destination must provide write()") + _validate_max_transfer_retries(max_transfer_retries) cursor: str | None = None result = AuditLogDownloadResult(bytes_written=0, chunks=0, rows=0) + seen_cursors: set[str] = set() while True: - body, headers = await _async_fetch_verified_chunk(fetch_chunk, cursor) + body, headers = await _async_fetch_verified_chunk(fetch_chunk, cursor, max_transfer_retries) chunk_rows, next_cursor, has_more = _parse_chunk_headers(headers, cursor) - _write_chunk(destination, body) + if has_more and next_cursor is not None: + if next_cursor in seen_cursors: + raise AuditLogDownloadError("response repeated X-Next-Cursor header") + seen_cursors.add(next_cursor) + await to_thread.run_sync(_write_chunk, destination, body) cursor = next_cursor result = AuditLogDownloadResult( @@ -111,7 +128,7 @@ async def async_download_audit_logs( rows=result.rows + chunk_rows, ) if on_progress is not None: - on_progress( + callback_result = on_progress( AuditLogDownloadProgress( bytes_written=result.bytes_written, chunks=result.chunks, @@ -119,14 +136,18 @@ async def async_download_audit_logs( chunk_rows=chunk_rows, ) ) + if inspect.isawaitable(callback_result): + await callback_result if not has_more: return result -def _fetch_verified_chunk(fetch_chunk: SyncFetchChunk, cursor: str | None) -> tuple[bytes, httpx.Headers]: - for attempt in range(1, _DOWNLOAD_ATTEMPTS + 1): +def _fetch_verified_chunk( + fetch_chunk: SyncFetchChunk, cursor: str | None, max_transfer_retries: int +) -> tuple[bytes, httpx.Headers]: + for retries in range(max_transfer_retries + 1): + response = fetch_chunk(cursor) try: - response = fetch_chunk(cursor) try: body = response.read() headers = httpx.Headers(response.headers) @@ -134,28 +155,30 @@ def _fetch_verified_chunk(fetch_chunk: SyncFetchChunk, cursor: str | None) -> tu response.close() _verify_checksum(body, headers) return body, headers - except Exception as error: - if attempt == _DOWNLOAD_ATTEMPTS or not _is_retryable(error): + except Exception: + if retries == max_transfer_retries: raise - time.sleep(min(2 ** (attempt - 1), _MAX_RETRY_DELAY)) + time.sleep(min(2**retries, _MAX_RETRY_DELAY)) raise AssertionError("unreachable") -async def _async_fetch_verified_chunk(fetch_chunk: AsyncFetchChunk, cursor: str | None) -> tuple[bytes, httpx.Headers]: - for attempt in range(1, _DOWNLOAD_ATTEMPTS + 1): +async def _async_fetch_verified_chunk( + fetch_chunk: AsyncFetchChunk, cursor: str | None, max_transfer_retries: int +) -> tuple[bytes, httpx.Headers]: + for retries in range(max_transfer_retries + 1): + response = await fetch_chunk(cursor) try: - response = await fetch_chunk(cursor) try: body = await response.read() headers = httpx.Headers(response.headers) finally: await response.close() - _verify_checksum(body, headers) + await to_thread.run_sync(_verify_checksum, body, headers) return body, headers - except Exception as error: - if attempt == _DOWNLOAD_ATTEMPTS or not _is_retryable(error): + except Exception: + if retries == max_transfer_retries: raise - await anyio.sleep(min(2 ** (attempt - 1), _MAX_RETRY_DELAY)) + await anyio.sleep(min(2**retries, _MAX_RETRY_DELAY)) raise AssertionError("unreachable") @@ -199,7 +222,6 @@ def _write_chunk(destination: BinaryIO, body: bytes) -> None: remaining = remaining[written:] -def _is_retryable(error: Exception) -> bool: - if isinstance(error, APIStatusError): - return error.status_code == 429 or error.status_code >= 500 - return True +def _validate_max_transfer_retries(max_transfer_retries: int) -> None: + if max_transfer_retries < 0: + raise ValueError("max_transfer_retries must be non-negative") diff --git a/src/kernel/resources/audit_logs.py b/src/kernel/resources/audit_logs.py index 5b50ad59..25e48993 100644 --- a/src/kernel/resources/audit_logs.py +++ b/src/kernel/resources/audit_logs.py @@ -32,6 +32,7 @@ from ..types.audit_log_entry import AuditLogEntry from ..lib.audit_log_download import ( ProgressCallback, + AsyncProgressCallback, AuditLogDownloadResult, download_audit_logs, async_download_audit_logs, @@ -243,6 +244,7 @@ def download( search_user_id: SequenceNotStr[str] | Omit = omit, service: str | Omit = omit, on_progress: ProgressCallback | None = None, + max_transfer_retries: int = 6, extra_headers: Headers | None = None, extra_query: Query | None = None, extra_body: Body | None = None, @@ -251,13 +253,13 @@ def download( """Download a complete audit log export to a writable binary destination. The SDK verifies every chunk and retries transient transfer failures. It - does not close the destination. + does not close the destination. If the download fails, the destination + may contain a partial export; use a temporary file and atomic rename + when the completed export must be published atomically. """ - resource = self._client.with_options(max_retries=0).audit_logs - def fetch_chunk(cursor: str | None) -> BinaryAPIResponse: - return resource.export_chunk( + return self.export_chunk( end=end, start=start, auth_strategy=auth_strategy, @@ -275,7 +277,12 @@ def fetch_chunk(cursor: str | None) -> BinaryAPIResponse: timeout=timeout, ) - return download_audit_logs(fetch_chunk, to, on_progress=on_progress) + return download_audit_logs( + fetch_chunk, + to, + on_progress=on_progress, + max_transfer_retries=max_transfer_retries, + ) class AsyncAuditLogsResource(AsyncAPIResource): @@ -480,7 +487,8 @@ async def download( search: str | Omit = omit, search_user_id: SequenceNotStr[str] | Omit = omit, service: str | Omit = omit, - on_progress: ProgressCallback | None = None, + on_progress: AsyncProgressCallback | None = None, + max_transfer_retries: int = 6, extra_headers: Headers | None = None, extra_query: Query | None = None, extra_body: Body | None = None, @@ -489,13 +497,13 @@ async def download( """Download a complete audit log export to a writable binary destination. The SDK verifies every chunk and retries transient transfer failures. It - does not close the destination. + does not close the destination. If the download fails, the destination + may contain a partial export; use a temporary file and atomic rename + when the completed export must be published atomically. """ - resource = self._client.with_options(max_retries=0).audit_logs - async def fetch_chunk(cursor: str | None) -> AsyncBinaryAPIResponse: - return await resource.export_chunk( + return await self.export_chunk( end=end, start=start, auth_strategy=auth_strategy, @@ -513,7 +521,12 @@ async def fetch_chunk(cursor: str | None) -> AsyncBinaryAPIResponse: timeout=timeout, ) - return await async_download_audit_logs(fetch_chunk, to, on_progress=on_progress) + return await async_download_audit_logs( + fetch_chunk, + to, + on_progress=on_progress, + max_transfer_retries=max_transfer_retries, + ) class AuditLogsResourceWithRawResponse: diff --git a/tests/test_audit_log_download.py b/tests/test_audit_log_download.py index 9c1a8c59..ccf95e07 100644 --- a/tests/test_audit_log_download.py +++ b/tests/test_audit_log_download.py @@ -1,13 +1,15 @@ from __future__ import annotations import hashlib +import threading from io import BytesIO +from typing import BinaryIO import httpx import pytest from respx import MockRouter -from kernel import Kernel, AsyncKernel, AuditLogDownloadError +from kernel import Kernel, AsyncKernel, InternalServerError, AuditLogDownloadError from kernel.lib import audit_log_download @@ -57,25 +59,43 @@ def respond(request: httpx.Request) -> httpx.Response: assert cursors == [None, "next"] -async def test_async_download_writes_verified_chunks(async_client: AsyncKernel, respx_mock: MockRouter) -> None: +async def test_async_download_writes_verified_chunks( + async_client: AsyncKernel, respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: respx_mock.get("/audit-logs/export/chunk").mock( side_effect=[ chunk_response(b"first", rows=2, has_more=True, next_cursor="next"), chunk_response(b"second", rows=1, has_more=False), ] ) + thread_ids: list[int] = [] + write_chunk = audit_log_download._write_chunk + + def record_write(destination: BinaryIO, body: bytes) -> None: + thread_ids.append(threading.get_ident()) + write_chunk(destination, body) + + monkeypatch.setattr(audit_log_download, "_write_chunk", record_write) destination = BytesIO() + progress_called = False + + async def on_progress(_: object) -> None: + nonlocal progress_called + progress_called = True result = await async_client.audit_logs.download( to=destination, start="2026-06-01T00:00:00Z", end="2026-06-02T00:00:00Z", + on_progress=on_progress, ) assert destination.getvalue() == b"firstsecond" assert result.bytes_written == 11 assert result.chunks == 2 assert result.rows == 3 + assert progress_called + assert thread_ids and all(thread_id != threading.get_ident() for thread_id in thread_ids) def test_download_retries_checksum_mismatch( @@ -102,9 +122,7 @@ def no_sleep(_: float) -> None: assert route.call_count == 2 -def test_download_owns_http_retries(client: Kernel, respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch) -> None: - delays: list[float] = [] - monkeypatch.setattr(audit_log_download.time, "sleep", delays.append) +def test_download_uses_client_http_retries(client: Kernel, respx_mock: MockRouter) -> None: route = respx_mock.get("/audit-logs/export/chunk").mock( side_effect=[ httpx.Response(500, json={"message": "temporary failure"}), @@ -119,7 +137,57 @@ def test_download_owns_http_retries(client: Kernel, respx_mock: MockRouter, monk ) assert route.call_count == 2 - assert delays == [1] + + +def test_download_respects_disabled_http_retries(client: Kernel, respx_mock: MockRouter) -> None: + route = respx_mock.get("/audit-logs/export/chunk").mock( + return_value=httpx.Response(500, json={"message": "temporary failure"}) + ) + + with pytest.raises(InternalServerError, match="500"): + client.with_options(max_retries=0).audit_logs.download( + to=BytesIO(), + start="2026-06-01T00:00:00Z", + end="2026-06-02T00:00:00Z", + ) + + assert route.call_count == 1 + + +def test_download_respects_transfer_retry_limit(client: Kernel, respx_mock: MockRouter) -> None: + bad = chunk_response(b"bad", rows=1, has_more=False) + bad.headers["x-content-sha256"] = hashlib.sha256(b"good").hexdigest() + route = respx_mock.get("/audit-logs/export/chunk").mock(return_value=bad) + + with pytest.raises(AuditLogDownloadError, match="checksum mismatch"): + client.audit_logs.download( + to=BytesIO(), + start="2026-06-01T00:00:00Z", + end="2026-06-02T00:00:00Z", + max_transfer_retries=0, + ) + + assert route.call_count == 1 + + +def test_download_rejects_cursor_cycle(client: Kernel, respx_mock: MockRouter) -> None: + respx_mock.get("/audit-logs/export/chunk").mock( + side_effect=[ + chunk_response(b"first", rows=1, has_more=True, next_cursor="a"), + chunk_response(b"second", rows=1, has_more=True, next_cursor="b"), + chunk_response(b"duplicate", rows=1, has_more=True, next_cursor="a"), + ] + ) + destination = BytesIO() + + with pytest.raises(AuditLogDownloadError, match="response repeated X-Next-Cursor header"): + client.audit_logs.download( + to=destination, + start="2026-06-01T00:00:00Z", + end="2026-06-02T00:00:00Z", + ) + + assert destination.getvalue() == b"firstsecond" def test_download_rejects_invalid_cursor_before_writing(client: Kernel, respx_mock: MockRouter) -> None: From 39e5555ba9d411869e3f5ceb48495f613c00c2ba Mon Sep 17 00:00:00 2001 From: yummybomb <19238148+yummybomb@users.noreply.github.com> Date: Tue, 21 Jul 2026 15:31:47 +0000 Subject: [PATCH 3/6] Harden audit log download validation --- src/kernel/lib/audit_log_download.py | 16 ++-- tests/test_audit_log_download.py | 106 ++++++++++++++++++++++++++- 2 files changed, 113 insertions(+), 9 deletions(-) diff --git a/src/kernel/lib/audit_log_download.py b/src/kernel/lib/audit_log_download.py index 783317e5..cda07dab 100644 --- a/src/kernel/lib/audit_log_download.py +++ b/src/kernel/lib/audit_log_download.py @@ -13,6 +13,7 @@ from .._exceptions import KernelError _DEFAULT_MAX_TRANSFER_RETRIES = 6 +_MAX_CHUNK_ROWS = 50_000 _MAX_RETRY_DELAY = 8.0 @@ -198,11 +199,10 @@ def _parse_chunk_headers(headers: httpx.Headers, current_cursor: str | None) -> has_more = has_more_value == "true" row_count = headers.get("x-row-count") - try: - rows = int(row_count) if row_count is not None else -1 - except ValueError as error: - raise AuditLogDownloadError("response missing or invalid X-Row-Count header") from error - if rows < 0: + if row_count is None or not row_count.isascii() or not row_count.isdecimal(): + raise AuditLogDownloadError("response missing or invalid X-Row-Count header") + rows = int(row_count) + if rows > _MAX_CHUNK_ROWS: raise AuditLogDownloadError("response missing or invalid X-Row-Count header") next_cursor = headers.get("x-next-cursor") or None @@ -217,11 +217,11 @@ def _write_chunk(destination: BinaryIO, body: bytes) -> None: remaining = memoryview(body) while remaining: written = destination.write(remaining) - if written <= 0 or written > len(remaining): + if type(written) is not int or written <= 0 or written > len(remaining): raise AuditLogDownloadError("audit log download destination performed a short write") remaining = remaining[written:] def _validate_max_transfer_retries(max_transfer_retries: int) -> None: - if max_transfer_retries < 0: - raise ValueError("max_transfer_retries must be non-negative") + if type(max_transfer_retries) is not int or max_transfer_retries < 0: + raise ValueError("max_transfer_retries must be a non-negative integer") diff --git a/tests/test_audit_log_download.py b/tests/test_audit_log_download.py index ccf95e07..262c32a8 100644 --- a/tests/test_audit_log_download.py +++ b/tests/test_audit_log_download.py @@ -3,7 +3,7 @@ import hashlib import threading from io import BytesIO -from typing import BinaryIO +from typing import BinaryIO, cast import httpx import pytest @@ -98,6 +98,110 @@ async def on_progress(_: object) -> None: assert thread_ids and all(thread_id != threading.get_ident() for thread_id in thread_ids) +async def test_async_download_retries_checksum_mismatch( + async_client: AsyncKernel, respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + async def no_sleep(_: float) -> None: + pass + + monkeypatch.setattr(audit_log_download.anyio, "sleep", no_sleep) + bad = chunk_response(b"bad", rows=1, has_more=False) + bad.headers["x-content-sha256"] = hashlib.sha256(b"good").hexdigest() + route = respx_mock.get("/audit-logs/export/chunk").mock( + side_effect=[bad, chunk_response(b"good", rows=1, has_more=False)] + ) + destination = BytesIO() + + await async_client.audit_logs.download( + to=destination, + start="2026-06-01T00:00:00Z", + end="2026-06-02T00:00:00Z", + ) + + assert destination.getvalue() == b"good" + assert route.call_count == 2 + + +async def test_async_download_uses_client_http_retries(async_client: AsyncKernel, respx_mock: MockRouter) -> None: + route = respx_mock.get("/audit-logs/export/chunk").mock( + side_effect=[ + httpx.Response(500, json={"message": "temporary failure"}), + chunk_response(b"good", rows=1, has_more=False), + ] + ) + + await async_client.audit_logs.download( + to=BytesIO(), + start="2026-06-01T00:00:00Z", + end="2026-06-02T00:00:00Z", + ) + + assert route.call_count == 2 + + +async def test_async_download_respects_disabled_http_retries(async_client: AsyncKernel, respx_mock: MockRouter) -> None: + route = respx_mock.get("/audit-logs/export/chunk").mock( + return_value=httpx.Response(500, json={"message": "temporary failure"}) + ) + + with pytest.raises(InternalServerError, match="500"): + await async_client.with_options(max_retries=0).audit_logs.download( + to=BytesIO(), + start="2026-06-01T00:00:00Z", + end="2026-06-02T00:00:00Z", + ) + + assert route.call_count == 1 + + +@pytest.mark.parametrize("row_count", ["", "1.0", "50001"]) +def test_download_rejects_invalid_row_count(row_count: str, client: Kernel, respx_mock: MockRouter) -> None: + response = chunk_response(b"chunk", rows=1, has_more=False) + response.headers["x-row-count"] = row_count + respx_mock.get("/audit-logs/export/chunk").mock(return_value=response) + destination = BytesIO() + + with pytest.raises(AuditLogDownloadError, match="response missing or invalid X-Row-Count header"): + client.audit_logs.download( + to=destination, + start="2026-06-01T00:00:00Z", + end="2026-06-02T00:00:00Z", + ) + + assert destination.getvalue() == b"" + + +class InvalidWriteDestination: + def __init__(self, result: object) -> None: + self.result = result + + def write(self, _: object) -> object: + return self.result + + +@pytest.mark.parametrize("write_result", [float("nan"), 0.5, True]) +def test_download_rejects_invalid_write_result(write_result: object, client: Kernel, respx_mock: MockRouter) -> None: + respx_mock.get("/audit-logs/export/chunk").mock(return_value=chunk_response(b"chunk", rows=1, has_more=False)) + + with pytest.raises(AuditLogDownloadError, match="audit log download destination performed a short write"): + client.audit_logs.download( + to=cast(BinaryIO, InvalidWriteDestination(write_result)), + start="2026-06-01T00:00:00Z", + end="2026-06-02T00:00:00Z", + ) + + +@pytest.mark.parametrize("max_transfer_retries", [True, 1.5]) +def test_download_rejects_invalid_transfer_retry_count(max_transfer_retries: object, client: Kernel) -> None: + with pytest.raises(ValueError, match="max_transfer_retries must be a non-negative integer"): + client.audit_logs.download( + to=BytesIO(), + start="2026-06-01T00:00:00Z", + end="2026-06-02T00:00:00Z", + max_transfer_retries=cast(int, max_transfer_retries), + ) + + def test_download_retries_checksum_mismatch( client: Kernel, respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch ) -> None: From e25457909c27de921570c872a70e38434adca250 Mon Sep 17 00:00:00 2001 From: yummybomb <19238148+yummybomb@users.noreply.github.com> Date: Tue, 21 Jul 2026 16:13:20 +0000 Subject: [PATCH 4/6] Retry streamed audit log transfers --- src/kernel/lib/audit_log_download.py | 55 ++++++++++-------- src/kernel/resources/audit_logs.py | 10 ++-- tests/test_audit_log_download.py | 85 ++++++++++++++++++++++++++-- 3 files changed, 117 insertions(+), 33 deletions(-) diff --git a/src/kernel/lib/audit_log_download.py b/src/kernel/lib/audit_log_download.py index cda07dab..4cfdcc26 100644 --- a/src/kernel/lib/audit_log_download.py +++ b/src/kernel/lib/audit_log_download.py @@ -3,7 +3,7 @@ import time import hashlib import inspect -from typing import BinaryIO, Callable, Optional, Protocol, Awaitable +from typing import BinaryIO, Callable, Optional, Protocol, Awaitable, ContextManager, AsyncContextManager from dataclasses import dataclass import anyio @@ -51,8 +51,8 @@ async def read(self) -> bytes: ... async def close(self) -> None: ... -SyncFetchChunk = Callable[[Optional[str]], _SyncChunkResponse] -AsyncFetchChunk = Callable[[Optional[str]], Awaitable[_AsyncChunkResponse]] +SyncFetchChunk = Callable[[Optional[str]], ContextManager[_SyncChunkResponse]] +AsyncFetchChunk = Callable[[Optional[str]], AsyncContextManager[_AsyncChunkResponse]] ProgressCallback = Callable[[AuditLogDownloadProgress], None] AsyncProgressCallback = Callable[[AuditLogDownloadProgress], Optional[Awaitable[None]]] @@ -147,19 +147,20 @@ def _fetch_verified_chunk( fetch_chunk: SyncFetchChunk, cursor: str | None, max_transfer_retries: int ) -> tuple[bytes, httpx.Headers]: for retries in range(max_transfer_retries + 1): - response = fetch_chunk(cursor) - try: + transfer_error: Exception | None = None + with fetch_chunk(cursor) as response: try: body = response.read() headers = httpx.Headers(response.headers) - finally: - response.close() - _verify_checksum(body, headers) - return body, headers - except Exception: - if retries == max_transfer_retries: - raise - time.sleep(min(2**retries, _MAX_RETRY_DELAY)) + _verify_checksum(body, headers) + except Exception as error: + transfer_error = error + else: + return body, headers + assert transfer_error is not None + if retries == max_transfer_retries: + raise transfer_error + time.sleep(min(2**retries, _MAX_RETRY_DELAY)) raise AssertionError("unreachable") @@ -167,19 +168,20 @@ async def _async_fetch_verified_chunk( fetch_chunk: AsyncFetchChunk, cursor: str | None, max_transfer_retries: int ) -> tuple[bytes, httpx.Headers]: for retries in range(max_transfer_retries + 1): - response = await fetch_chunk(cursor) - try: + transfer_error: Exception | None = None + async with fetch_chunk(cursor) as response: try: body = await response.read() headers = httpx.Headers(response.headers) - finally: - await response.close() - await to_thread.run_sync(_verify_checksum, body, headers) - return body, headers - except Exception: - if retries == max_transfer_retries: - raise - await anyio.sleep(min(2**retries, _MAX_RETRY_DELAY)) + await to_thread.run_sync(_verify_checksum, body, headers) + except Exception as error: + transfer_error = error + else: + return body, headers + assert transfer_error is not None + if retries == max_transfer_retries: + raise transfer_error + await anyio.sleep(min(2**retries, _MAX_RETRY_DELAY)) raise AssertionError("unreachable") @@ -199,7 +201,12 @@ def _parse_chunk_headers(headers: httpx.Headers, current_cursor: str | None) -> has_more = has_more_value == "true" row_count = headers.get("x-row-count") - if row_count is None or not row_count.isascii() or not row_count.isdecimal(): + if ( + row_count is None + or not row_count.isascii() + or not row_count.isdecimal() + or len(row_count) > len(str(_MAX_CHUNK_ROWS)) + ): raise AuditLogDownloadError("response missing or invalid X-Row-Count header") rows = int(row_count) if rows > _MAX_CHUNK_ROWS: diff --git a/src/kernel/resources/audit_logs.py b/src/kernel/resources/audit_logs.py index 25e48993..a5491697 100644 --- a/src/kernel/resources/audit_logs.py +++ b/src/kernel/resources/audit_logs.py @@ -2,7 +2,7 @@ from __future__ import annotations -from typing import Union, BinaryIO +from typing import Union, BinaryIO, ContextManager, AsyncContextManager from datetime import datetime from typing_extensions import Literal @@ -258,8 +258,8 @@ def download( when the completed export must be published atomically. """ - def fetch_chunk(cursor: str | None) -> BinaryAPIResponse: - return self.export_chunk( + def fetch_chunk(cursor: str | None) -> ContextManager[StreamedBinaryAPIResponse]: + return self.with_streaming_response.export_chunk( end=end, start=start, auth_strategy=auth_strategy, @@ -502,8 +502,8 @@ async def download( when the completed export must be published atomically. """ - async def fetch_chunk(cursor: str | None) -> AsyncBinaryAPIResponse: - return await self.export_chunk( + def fetch_chunk(cursor: str | None) -> AsyncContextManager[AsyncStreamedBinaryAPIResponse]: + return self.with_streaming_response.export_chunk( end=end, start=start, auth_strategy=auth_strategy, diff --git a/tests/test_audit_log_download.py b/tests/test_audit_log_download.py index 262c32a8..0419ca9e 100644 --- a/tests/test_audit_log_download.py +++ b/tests/test_audit_log_download.py @@ -3,7 +3,8 @@ import hashlib import threading from io import BytesIO -from typing import BinaryIO, cast +from typing import BinaryIO, Iterator, AsyncIterator, cast +from typing_extensions import override import httpx import pytest @@ -13,7 +14,7 @@ from kernel.lib import audit_log_download -def chunk_response(body: bytes, *, rows: int, has_more: bool, next_cursor: str | None = None) -> httpx.Response: +def chunk_headers(body: bytes, *, rows: int, has_more: bool, next_cursor: str | None = None) -> dict[str, str]: headers = { "x-content-sha256": hashlib.sha256(body).hexdigest(), "x-has-more": str(has_more).lower(), @@ -21,7 +22,29 @@ def chunk_response(body: bytes, *, rows: int, has_more: bool, next_cursor: str | } if next_cursor is not None: headers["x-next-cursor"] = next_cursor - return httpx.Response(200, content=body, headers=headers) + return headers + + +def chunk_response(body: bytes, *, rows: int, has_more: bool, next_cursor: str | None = None) -> httpx.Response: + return httpx.Response( + 200, + content=body, + headers=chunk_headers(body, rows=rows, has_more=has_more, next_cursor=next_cursor), + ) + + +class BrokenSyncByteStream(httpx.SyncByteStream): + @override + def __iter__(self) -> Iterator[bytes]: + yield b"partial" + raise httpx.ReadError("truncated response body") + + +class BrokenAsyncByteStream(httpx.AsyncByteStream): + @override + async def __aiter__(self) -> AsyncIterator[bytes]: + yield b"partial" + raise httpx.ReadError("truncated response body") def test_download_writes_verified_chunks(client: Kernel, respx_mock: MockRouter) -> None: @@ -98,6 +121,33 @@ async def on_progress(_: object) -> None: assert thread_ids and all(thread_id != threading.get_ident() for thread_id in thread_ids) +async def test_async_download_retries_body_read_failure( + async_client: AsyncKernel, respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + async def no_sleep(_: float) -> None: + pass + + monkeypatch.setattr(audit_log_download.anyio, "sleep", no_sleep) + headers = chunk_headers(b"good", rows=1, has_more=False) + route = respx_mock.get("/audit-logs/export/chunk").mock( + side_effect=[ + httpx.Response(200, headers=headers, stream=BrokenAsyncByteStream()), + chunk_response(b"good", rows=1, has_more=False), + ] + ) + destination = BytesIO() + + await async_client.with_options(max_retries=0).audit_logs.download( + to=destination, + start="2026-06-01T00:00:00Z", + end="2026-06-02T00:00:00Z", + max_transfer_retries=1, + ) + + assert destination.getvalue() == b"good" + assert route.call_count == 2 + + async def test_async_download_retries_checksum_mismatch( async_client: AsyncKernel, respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch ) -> None: @@ -154,7 +204,7 @@ async def test_async_download_respects_disabled_http_retries(async_client: Async assert route.call_count == 1 -@pytest.mark.parametrize("row_count", ["", "1.0", "50001"]) +@pytest.mark.parametrize("row_count", ["", "1.0", "50001", "9" * 5000]) def test_download_rejects_invalid_row_count(row_count: str, client: Kernel, respx_mock: MockRouter) -> None: response = chunk_response(b"chunk", rows=1, has_more=False) response.headers["x-row-count"] = row_count @@ -202,6 +252,33 @@ def test_download_rejects_invalid_transfer_retry_count(max_transfer_retries: obj ) +def test_download_retries_body_read_failure( + client: Kernel, respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + def no_sleep(_: float) -> None: + pass + + monkeypatch.setattr(audit_log_download.time, "sleep", no_sleep) + headers = chunk_headers(b"good", rows=1, has_more=False) + route = respx_mock.get("/audit-logs/export/chunk").mock( + side_effect=[ + httpx.Response(200, headers=headers, stream=BrokenSyncByteStream()), + chunk_response(b"good", rows=1, has_more=False), + ] + ) + destination = BytesIO() + + client.with_options(max_retries=0).audit_logs.download( + to=destination, + start="2026-06-01T00:00:00Z", + end="2026-06-02T00:00:00Z", + max_transfer_retries=1, + ) + + assert destination.getvalue() == b"good" + assert route.call_count == 2 + + def test_download_retries_checksum_mismatch( client: Kernel, respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch ) -> None: From 0467576d093f437f1200202d067f216312ea59f5 Mon Sep 17 00:00:00 2001 From: yummybomb <19238148+yummybomb@users.noreply.github.com> Date: Tue, 21 Jul 2026 16:19:22 +0000 Subject: [PATCH 5/6] Accept padded audit log row counts --- src/kernel/lib/audit_log_download.py | 12 +++++------- tests/test_audit_log_download.py | 14 ++++++++++++++ 2 files changed, 19 insertions(+), 7 deletions(-) diff --git a/src/kernel/lib/audit_log_download.py b/src/kernel/lib/audit_log_download.py index 4cfdcc26..f334d6ae 100644 --- a/src/kernel/lib/audit_log_download.py +++ b/src/kernel/lib/audit_log_download.py @@ -201,14 +201,12 @@ def _parse_chunk_headers(headers: httpx.Headers, current_cursor: str | None) -> has_more = has_more_value == "true" row_count = headers.get("x-row-count") - if ( - row_count is None - or not row_count.isascii() - or not row_count.isdecimal() - or len(row_count) > len(str(_MAX_CHUNK_ROWS)) - ): + if row_count is None or not row_count.isascii() or not row_count.isdecimal(): raise AuditLogDownloadError("response missing or invalid X-Row-Count header") - rows = int(row_count) + normalized_row_count = row_count.lstrip("0") or "0" + if len(normalized_row_count) > len(str(_MAX_CHUNK_ROWS)): + raise AuditLogDownloadError("response missing or invalid X-Row-Count header") + rows = int(normalized_row_count) if rows > _MAX_CHUNK_ROWS: raise AuditLogDownloadError("response missing or invalid X-Row-Count header") diff --git a/tests/test_audit_log_download.py b/tests/test_audit_log_download.py index 0419ca9e..c400ac41 100644 --- a/tests/test_audit_log_download.py +++ b/tests/test_audit_log_download.py @@ -221,6 +221,20 @@ def test_download_rejects_invalid_row_count(row_count: str, client: Kernel, resp assert destination.getvalue() == b"" +def test_download_accepts_zero_padded_row_count(client: Kernel, respx_mock: MockRouter) -> None: + response = chunk_response(b"chunk", rows=1, has_more=False) + response.headers["x-row-count"] = "0" * 5000 + "1" + respx_mock.get("/audit-logs/export/chunk").mock(return_value=response) + + result = client.audit_logs.download( + to=BytesIO(), + start="2026-06-01T00:00:00Z", + end="2026-06-02T00:00:00Z", + ) + + assert result.rows == 1 + + class InvalidWriteDestination: def __init__(self, result: object) -> None: self.result = result From 3e2fcf923865a32d0c3768dbfde720f073e47b5e Mon Sep 17 00:00:00 2001 From: yummybomb <19238148+yummybomb@users.noreply.github.com> Date: Tue, 21 Jul 2026 18:27:20 +0000 Subject: [PATCH 6/6] Default audit log downloads to gzip --- src/kernel/resources/audit_logs.py | 26 ++++++++++++-------------- tests/test_audit_log_download.py | 16 +++++++++++++--- 2 files changed, 25 insertions(+), 17 deletions(-) diff --git a/src/kernel/resources/audit_logs.py b/src/kernel/resources/audit_logs.py index a5491697..aac5d71c 100644 --- a/src/kernel/resources/audit_logs.py +++ b/src/kernel/resources/audit_logs.py @@ -237,7 +237,6 @@ def download( start: Union[str, datetime], auth_strategy: str | Omit = omit, exclude_method: SequenceNotStr[str] | Omit = omit, - format: Literal["jsonl", "jsonl.gz"] | Omit = omit, limit: int | Omit = omit, method: str | Omit = omit, search: str | Omit = omit, @@ -250,12 +249,13 @@ def download( extra_body: Body | None = None, timeout: float | httpx.Timeout | None | NotGiven = not_given, ) -> AuditLogDownloadResult: - """Download a complete audit log export to a writable binary destination. + """Download a complete gzip-compressed JSON Lines audit log export. - The SDK verifies every chunk and retries transient transfer failures. It - does not close the destination. If the download fails, the destination - may contain a partial export; use a temporary file and atomic rename - when the completed export must be published atomically. + The SDK writes the export to a writable binary destination, verifies + every chunk, and retries transient transfer failures. It does not close + the destination. If the download fails, the destination may contain a + partial export; use a temporary file and atomic rename when the completed + export must be published atomically. """ def fetch_chunk(cursor: str | None) -> ContextManager[StreamedBinaryAPIResponse]: @@ -265,7 +265,6 @@ def fetch_chunk(cursor: str | None) -> ContextManager[StreamedBinaryAPIResponse] auth_strategy=auth_strategy, cursor=cursor if cursor is not None else omit, exclude_method=exclude_method, - format=format, limit=limit, method=method, search=search, @@ -481,7 +480,6 @@ async def download( start: Union[str, datetime], auth_strategy: str | Omit = omit, exclude_method: SequenceNotStr[str] | Omit = omit, - format: Literal["jsonl", "jsonl.gz"] | Omit = omit, limit: int | Omit = omit, method: str | Omit = omit, search: str | Omit = omit, @@ -494,12 +492,13 @@ async def download( extra_body: Body | None = None, timeout: float | httpx.Timeout | None | NotGiven = not_given, ) -> AuditLogDownloadResult: - """Download a complete audit log export to a writable binary destination. + """Download a complete gzip-compressed JSON Lines audit log export. - The SDK verifies every chunk and retries transient transfer failures. It - does not close the destination. If the download fails, the destination - may contain a partial export; use a temporary file and atomic rename - when the completed export must be published atomically. + The SDK writes the export to a writable binary destination, verifies + every chunk, and retries transient transfer failures. It does not close + the destination. If the download fails, the destination may contain a + partial export; use a temporary file and atomic rename when the completed + export must be published atomically. """ def fetch_chunk(cursor: str | None) -> AsyncContextManager[AsyncStreamedBinaryAPIResponse]: @@ -509,7 +508,6 @@ def fetch_chunk(cursor: str | None) -> AsyncContextManager[AsyncStreamedBinaryAP auth_strategy=auth_strategy, cursor=cursor if cursor is not None else omit, exclude_method=exclude_method, - format=format, limit=limit, method=method, search=search, diff --git a/tests/test_audit_log_download.py b/tests/test_audit_log_download.py index c400ac41..e254066e 100644 --- a/tests/test_audit_log_download.py +++ b/tests/test_audit_log_download.py @@ -55,9 +55,11 @@ def test_download_writes_verified_chunks(client: Kernel, respx_mock: MockRouter) ] ) cursors: list[str | None] = [] + formats: list[str | None] = [] def respond(request: httpx.Request) -> httpx.Response: cursors.append(request.url.params.get("cursor")) + formats.append(request.url.params.get("format")) return next(responses) respx_mock.get("/audit-logs/export/chunk").mock(side_effect=respond) @@ -68,7 +70,6 @@ def respond(request: httpx.Request) -> httpx.Response: to=destination, start="2026-06-01T00:00:00Z", end="2026-06-02T00:00:00Z", - format="jsonl.gz", on_progress=lambda update: progress.append( (update.bytes_written, update.chunks, update.rows, update.chunk_rows) ), @@ -80,17 +81,25 @@ def respond(request: httpx.Request) -> httpx.Response: assert result.rows == 3 assert progress == [(5, 1, 2, 2), (11, 2, 3, 1)] assert cursors == [None, "next"] + assert formats == [None, None] async def test_async_download_writes_verified_chunks( async_client: AsyncKernel, respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch ) -> None: - respx_mock.get("/audit-logs/export/chunk").mock( - side_effect=[ + responses = iter( + [ chunk_response(b"first", rows=2, has_more=True, next_cursor="next"), chunk_response(b"second", rows=1, has_more=False), ] ) + formats: list[str | None] = [] + + def respond(request: httpx.Request) -> httpx.Response: + formats.append(request.url.params.get("format")) + return next(responses) + + respx_mock.get("/audit-logs/export/chunk").mock(side_effect=respond) thread_ids: list[int] = [] write_chunk = audit_log_download._write_chunk @@ -119,6 +128,7 @@ async def on_progress(_: object) -> None: assert result.rows == 3 assert progress_called assert thread_ids and all(thread_id != threading.get_ident() for thread_id in thread_ids) + assert formats == [None, None] async def test_async_download_retries_body_read_failure(