Add complete audit log downloads

yummybomb · yummybomb · commit 29163ad56198 · 2026-07-20T21:03:16.000Z
diff --git 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
@@ -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
@@ -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
