Merge pull request #133 from kernel/hypeship/audit-log-download-helper

yummybomb · web-flow · commit fa150d406208 · 2026-07-21T14:49:23.000-04:00
feat: add complete audit log download helper
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,232 @@
+from __future__ import annotations
+
+import time
+import hashlib
+import inspect
+from typing import BinaryIO, Callable, Optional, Protocol, Awaitable, ContextManager, AsyncContextManager
+from dataclasses import dataclass
+
+import anyio
+import httpx
+from anyio import to_thread
+
+from .._exceptions import KernelError
+
+_DEFAULT_MAX_TRANSFER_RETRIES = 6
+_MAX_CHUNK_ROWS = 50_000
+_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]], ContextManager[_SyncChunkResponse]]
+AsyncFetchChunk = Callable[[Optional[str]], AsyncContextManager[_AsyncChunkResponse]]
+ProgressCallback = Callable[[AuditLogDownloadProgress], None]
+AsyncProgressCallback = Callable[[AuditLogDownloadProgress], Optional[Awaitable[None]]]
+
+
+def download_audit_logs(
+    fetch_chunk: SyncFetchChunk,
+    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, 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
+        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: 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, 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)
+        await to_thread.run_sync(_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:
+            callback_result = on_progress(
+                AuditLogDownloadProgress(
+                    bytes_written=result.bytes_written,
+                    chunks=result.chunks,
+                    rows=result.rows,
+                    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, max_transfer_retries: int
+) -> tuple[bytes, httpx.Headers]:
+    for retries in range(max_transfer_retries + 1):
+        transfer_error: Exception | None = None
+        with fetch_chunk(cursor) as response:
+            try:
+                body = response.read()
+                headers = httpx.Headers(response.headers)
+                _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")
+
+
+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):
+        transfer_error: Exception | None = None
+        async with fetch_chunk(cursor) as response:
+            try:
+                body = await response.read()
+                headers = httpx.Headers(response.headers)
+                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")
+
+
+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")
+    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")
+    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")
+
+    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 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 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/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, ContextManager, AsyncContextManager
 from datetime import datetime
 from typing_extensions import Literal
 
@@ -30,6 +30,13 @@
 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,
+    AsyncProgressCallback,
+    AuditLogDownloadResult,
+    download_audit_logs,
+    async_download_audit_logs,
+)
 
 __all__ = ["AuditLogsResource", "AsyncAuditLogsResource"]
 
@@ -222,6 +229,60 @@ 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,
+        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,
+        max_transfer_retries: int = 6,
+        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 gzip-compressed JSON Lines audit log export.
+
+        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]:
+            return self.with_streaming_response.export_chunk(
+                end=end,
+                start=start,
+                auth_strategy=auth_strategy,
+                cursor=cursor if cursor is not None else omit,
+                exclude_method=exclude_method,
+                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,
+            max_transfer_retries=max_transfer_retries,
+        )
+
 
 class AsyncAuditLogsResource(AsyncAPIResource):
     """Read audit log records for the authenticated organization."""
@@ -411,6 +472,60 @@ 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,
+        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: AsyncProgressCallback | None = None,
+        max_transfer_retries: int = 6,
+        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 gzip-compressed JSON Lines audit log export.
+
+        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]:
+            return self.with_streaming_response.export_chunk(
+                end=end,
+                start=start,
+                auth_strategy=auth_strategy,
+                cursor=cursor if cursor is not None else omit,
+                exclude_method=exclude_method,
+                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,
+            max_transfer_retries=max_transfer_retries,
+        )
+
 
 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
