From c1d3a83392ef81f0b292ef8c49e60b84a5b7b6bf Mon Sep 17 00:00:00 2001 From: Doug Guthrie Date: Mon, 8 Jun 2026 09:29:29 -0600 Subject: [PATCH 1/8] Hoist duplicated streaming helpers into streaming_utils (R2, R3) The three streaming migrators (logs, datasets, experiments) each carried byte-for-byte copies of _is_http_413, _approx_event_size_bytes, _count_attachment_refs, the HTTP 413 constant, and a near-identical _dump_oversize_event_summary method (differing only by output dir, filename prefix, dest-id key name, and log wording). Move these to single shared functions in streaming_utils (is_http_413, approx_event_size_bytes, count_attachment_refs, dump_oversize_event_summary) and have all three migrators call them. The dump function takes the per-resource out_dir / filename_prefix / event_label / dest_id_field so behavior (including the exact summary schema and log messages) is preserved. Drops the now-unused httpx import from datasets and experiments. Behavior-preserving refactor: all existing tests pass unchanged. Adds test_streaming_helpers.py to pin the shared helpers directly (previously only covered indirectly through the migration paths). Co-Authored-By: Claude Opus 4.8 (1M context) --- braintrust_migrate/resources/datasets.py | 100 ++--------------- braintrust_migrate/resources/experiments.py | 101 ++--------------- braintrust_migrate/resources/logs.py | 95 ++-------------- braintrust_migrate/streaming_utils.py | 99 +++++++++++++++++ tests/unit/test_streaming_helpers.py | 114 ++++++++++++++++++++ 5 files changed, 245 insertions(+), 264 deletions(-) create mode 100644 tests/unit/test_streaming_helpers.py diff --git a/braintrust_migrate/resources/datasets.py b/braintrust_migrate/resources/datasets.py index dc7a71d..6b9c36d 100644 --- a/braintrust_migrate/resources/datasets.py +++ b/braintrust_migrate/resources/datasets.py @@ -8,7 +8,6 @@ from pathlib import Path from typing import Any, ClassVar -import httpx import structlog from braintrust_migrate.attachments import AttachmentCopier, OversizeFieldSpiller @@ -23,14 +22,13 @@ SeenIdsDB, build_btql_sorted_page_query, coerce_int_config, + dump_oversize_event_summary, + is_http_413, stream_btql_sorted_events_buffered, ) logger = structlog.get_logger(__name__) -# HTTP status codes -HTTP_STATUS_REQUEST_ENTITY_TOO_LARGE = 413 - class DatasetMigrator(ResourceMigrator[dict]): """Migrator for Braintrust datasets. @@ -536,86 +534,6 @@ def _query_text_for_limit(n: int) -> str: timeout_seconds=120.0, ) - @staticmethod - def _is_http_413(exc: Exception) -> bool: - return ( - isinstance(exc, httpx.HTTPStatusError) - and exc.response is not None - and int(exc.response.status_code) == HTTP_STATUS_REQUEST_ENTITY_TOO_LARGE - ) - - @staticmethod - def _approx_event_size_bytes(event: dict[str, Any]) -> int | None: - try: - return len(_json.dumps(event, separators=(",", ":"), ensure_ascii=False)) - except Exception: - return None - - @staticmethod - def _count_attachment_refs(event: dict[str, Any]) -> int: - def _walk(v: Any) -> int: - if isinstance(v, dict): - if v.get("type") == "braintrust_attachment" and isinstance( - v.get("key"), str - ): - return 1 - return sum(_walk(x) for x in v.values()) - if isinstance(v, list): - return sum(_walk(x) for x in v) - return 0 - - return _walk(event) - - def _dump_oversize_event_summary( - self, - *, - events_dir: Path, - cursor: str | None, - dest_dataset_id: str, - event: dict[str, Any], - error: Exception, - ) -> None: - event_id = event.get("id") - safe_id = str(event_id) if isinstance(event_id, str) and event_id else "unknown" - root_span_id = event.get("root_span_id") - span_id = event.get("span_id") - approx_size = self._approx_event_size_bytes(event) - attachment_refs = self._count_attachment_refs(event) - path = events_dir / f"oversize_dataset_event_{safe_id}.json" - summary = { - "error": str(error), - "cursor": cursor, - "dest_dataset_id": dest_dataset_id, - "event_id": event.get("id"), - "root_span_id": root_span_id, - "span_id": span_id, - "created": event.get("created"), - "approx_size_bytes": approx_size, - "attachment_refs": attachment_refs, - "top_level_keys": sorted(list(event.keys())), - } - try: - with open(path, "w") as f: - _json.dump(summary, f, indent=2) - self._logger.error( - "Oversize dataset event isolated (413). This specific event cannot be inserted.", - summary_path=str(path), - event_id=safe_id, - root_span_id=root_span_id, - span_id=span_id, - approx_size_bytes=approx_size, - attachment_refs=attachment_refs, - cursor=cursor, - ) - except Exception: - self._logger.error( - "Oversize dataset event isolated; failed to write summary", - event_id=safe_id, - root_span_id=root_span_id, - span_id=span_id, - cursor=cursor, - ) - @staticmethod def _group_stream_basename(source_dataset_ids: list[str]) -> str: if len(source_dataset_ids) == 1: @@ -703,17 +621,21 @@ async def _fetch(n: int) -> dict[str, Any]: ) async def _on_single_413(event: dict[str, Any], err: Exception) -> None: - self._dump_oversize_event_summary( - events_dir=events_dir, - cursor=state.btql_min_pagination_key, - dest_dataset_id=( + dump_oversize_event_summary( + out_dir=events_dir, + filename_prefix="oversize_dataset_event_", + event_label="dataset event", + dest_id_field="dest_dataset_id", + dest_id_value=( source_to_dest_dataset_ids.get(source_dataset_id) if isinstance(source_dataset_id, str) else None ) or "unknown", + cursor=state.btql_min_pagination_key, event=event, error=err, + logger=self._logger, ) async def _fetch_group(n: int) -> dict[str, Any]: @@ -751,7 +673,7 @@ async def _fetch_group(n: int) -> dict[str, Any]: ), flush_max_rows=self._sdk_flush_max_rows, flush_max_bytes=self._sdk_flush_max_bytes, - is_http_413=self._is_http_413, + is_http_413=is_http_413, on_single_413=_on_single_413, hooks=None if progress is None diff --git a/braintrust_migrate/resources/experiments.py b/braintrust_migrate/resources/experiments.py index fb42de8..36f092e 100644 --- a/braintrust_migrate/resources/experiments.py +++ b/braintrust_migrate/resources/experiments.py @@ -8,8 +8,6 @@ from pathlib import Path from typing import Any, ClassVar -import httpx - from braintrust_migrate.attachments import AttachmentCopier, OversizeFieldSpiller from braintrust_migrate.btql import ( btql_quote, @@ -22,12 +20,11 @@ SeenIdsDB, build_btql_sorted_page_query, coerce_int_config, + dump_oversize_event_summary, + is_http_413, stream_btql_sorted_events_buffered, ) -# HTTP status codes -HTTP_STATUS_REQUEST_ENTITY_TOO_LARGE = 413 - class ExperimentMigrator(ResourceMigrator[dict]): """Migrator for Braintrust experiments. @@ -652,86 +649,6 @@ def _query_text_for_limit(n: int) -> str: timeout_seconds=120.0, ) - @staticmethod - def _is_http_413(exc: Exception) -> bool: - return ( - isinstance(exc, httpx.HTTPStatusError) - and exc.response is not None - and int(exc.response.status_code) == HTTP_STATUS_REQUEST_ENTITY_TOO_LARGE - ) - - @staticmethod - def _approx_event_size_bytes(event: dict[str, Any]) -> int | None: - try: - return len(_json.dumps(event, separators=(",", ":"), ensure_ascii=False)) - except Exception: - return None - - @staticmethod - def _count_attachment_refs(event: dict[str, Any]) -> int: - def _walk(v: Any) -> int: - if isinstance(v, dict): - if v.get("type") == "braintrust_attachment" and isinstance( - v.get("key"), str - ): - return 1 - return sum(_walk(x) for x in v.values()) - if isinstance(v, list): - return sum(_walk(x) for x in v) - return 0 - - return _walk(event) - - def _dump_oversize_event_summary( - self, - *, - events_dir: Path, - cursor: str | None, - dest_experiment_id: str, - event: dict[str, Any], - error: Exception, - ) -> None: - event_id = event.get("id") - safe_id = str(event_id) if isinstance(event_id, str) and event_id else "unknown" - root_span_id = event.get("root_span_id") - span_id = event.get("span_id") - approx_size = self._approx_event_size_bytes(event) - attachment_refs = self._count_attachment_refs(event) - path = events_dir / f"oversize_experiment_event_{safe_id}.json" - summary = { - "error": str(error), - "cursor": cursor, - "dest_experiment_id": dest_experiment_id, - "event_id": event.get("id"), - "root_span_id": root_span_id, - "span_id": span_id, - "created": event.get("created"), - "approx_size_bytes": approx_size, - "attachment_refs": attachment_refs, - "top_level_keys": sorted(list(event.keys())), - } - try: - with open(path, "w") as f: - _json.dump(summary, f, indent=2) - self._logger.error( - "Oversize experiment event isolated (413). This specific event cannot be inserted.", - summary_path=str(path), - event_id=safe_id, - root_span_id=root_span_id, - span_id=span_id, - approx_size_bytes=approx_size, - attachment_refs=attachment_refs, - cursor=cursor, - ) - except Exception: - self._logger.error( - "Oversize experiment event isolated; failed to write summary", - event_id=safe_id, - root_span_id=root_span_id, - span_id=span_id, - cursor=cursor, - ) - @staticmethod def _group_stream_basename(source_experiment_ids: list[str]) -> str: if len(source_experiment_ids) == 1: @@ -815,17 +732,21 @@ async def _fetch(n: int) -> dict[str, Any]: async def _on_single_413(event: dict[str, Any], err: Exception) -> None: source_experiment_id = event.get("experiment_id") - self._dump_oversize_event_summary( - events_dir=events_dir, - cursor=state.btql_min_pagination_key, - dest_experiment_id=( + dump_oversize_event_summary( + out_dir=events_dir, + filename_prefix="oversize_experiment_event_", + event_label="experiment event", + dest_id_field="dest_experiment_id", + dest_id_value=( source_to_dest_experiment_ids.get(source_experiment_id) if isinstance(source_experiment_id, str) else None ) or "unknown", + cursor=state.btql_min_pagination_key, event=event, error=err, + logger=self._logger, ) await stream_btql_sorted_events_buffered( @@ -848,7 +769,7 @@ async def _on_single_413(event: dict[str, Any], err: Exception) -> None: ), flush_max_rows=self._sdk_flush_max_rows, flush_max_bytes=self._sdk_flush_max_bytes, - is_http_413=self._is_http_413, + is_http_413=is_http_413, on_single_413=_on_single_413, hooks=None if progress is None diff --git a/braintrust_migrate/resources/logs.py b/braintrust_migrate/resources/logs.py index eacddce..987e32a 100644 --- a/braintrust_migrate/resources/logs.py +++ b/braintrust_migrate/resources/logs.py @@ -33,13 +33,12 @@ SeenIdsDB, build_btql_sorted_page_query, coerce_int_config, + dump_oversize_event_summary, + is_http_413, ) logger = structlog.get_logger(__name__) -# HTTP status codes -HTTP_STATUS_REQUEST_ENTITY_TOO_LARGE = 413 - @dataclass(slots=True) class _LogsStreamingState: """Checkpoint state for streaming logs migration (small + restartable).""" @@ -361,85 +360,6 @@ def _extract_ids(events: list[dict[str, Any]]) -> list[str]: ids.append(event_id) return ids - @staticmethod - def _is_http_413(exc: Exception) -> bool: - return ( - isinstance(exc, httpx.HTTPStatusError) - and exc.response is not None - and int(exc.response.status_code) == HTTP_STATUS_REQUEST_ENTITY_TOO_LARGE - ) - - @staticmethod - def _approx_event_size_bytes(event: dict[str, Any]) -> int | None: - try: - return len(_json.dumps(event, separators=(",", ":"), ensure_ascii=False)) - except Exception: - return None - - @staticmethod - def _count_attachment_refs(event: dict[str, Any]) -> int: - def _walk(v: Any) -> int: - if isinstance(v, dict): - if v.get("type") == "braintrust_attachment" and isinstance( - v.get("key"), str - ): - return 1 - return sum(_walk(x) for x in v.values()) - if isinstance(v, list): - return sum(_walk(x) for x in v) - return 0 - - return _walk(event) - - def _dump_oversize_event_summary( - self, - *, - cursor: str | None, - dest_project_id: str, - event: dict[str, Any], - error: Exception, - ) -> None: - event_id = event.get("id") - safe_id = str(event_id) if isinstance(event_id, str) and event_id else "unknown" - root_span_id = event.get("root_span_id") - span_id = event.get("span_id") - approx_size = self._approx_event_size_bytes(event) - attachment_refs = self._count_attachment_refs(event) - path = self.checkpoint_dir / f"oversize_project_logs_event_{safe_id}.json" - summary = { - "error": str(error), - "cursor": cursor, - "dest_project_id": dest_project_id, - "event_id": event.get("id"), - "root_span_id": root_span_id, - "span_id": span_id, - "created": event.get("created"), - "approx_size_bytes": approx_size, - "attachment_refs": attachment_refs, - "top_level_keys": sorted(list(event.keys())), - } - try: - with open(path, "w") as f: - _json.dump(summary, f, indent=2) - self._logger.error( - "Oversize event isolated (413). This specific event cannot be inserted.", - summary_path=str(path), - event_id=safe_id, - root_span_id=root_span_id, - span_id=span_id, - approx_size_bytes=approx_size, - attachment_refs=attachment_refs, - cursor=cursor, - ) - except Exception: - self._logger.error( - "Oversize event isolated; failed to write summary", - event_id=safe_id, - root_span_id=root_span_id, - span_id=span_id, - cursor=cursor, - ) - async def migrate_all( self, project_id: str | None = None, max_concurrent: int | None = None ) -> dict[str, Any]: @@ -629,11 +549,16 @@ def _save_state() -> None: self._save_stream_state() async def _on_single_413(event: dict[str, Any], err: Exception) -> None: - self._dump_oversize_event_summary( + dump_oversize_event_summary( + out_dir=self.checkpoint_dir, + filename_prefix="oversize_project_logs_event_", + event_label="event", + dest_id_field="dest_project_id", + dest_id_value=dest_project_id, cursor=self._stream_state.btql_min_pagination_key, - dest_project_id=dest_project_id, event=event, error=err, + logger=self._logger, ) def _on_fetch(info: dict[str, Any]) -> None: @@ -848,7 +773,7 @@ async def _flush_pending_events() -> None: "error": e, } ) - if len(batch) == 1 and self._is_http_413(e): + if len(batch) == 1 and is_http_413(e): await _on_single_413(batch[0], e) raise if seen_db is not None and pending_seen_ids: diff --git a/braintrust_migrate/streaming_utils.py b/braintrust_migrate/streaming_utils.py index 6ea73fc..fcc5d61 100644 --- a/braintrust_migrate/streaming_utils.py +++ b/braintrust_migrate/streaming_utils.py @@ -15,12 +15,111 @@ from pathlib import Path from typing import Any, TypedDict, cast +import httpx + from braintrust_migrate.batching import ( approx_events_insert_payload_bytes, approx_json_bytes, ) from braintrust_migrate.btql import btql_quote +# HTTP status codes +HTTP_STATUS_REQUEST_ENTITY_TOO_LARGE = 413 + + +def is_http_413(exc: Exception) -> bool: + """Whether ``exc`` is an HTTP 413 (Request Entity Too Large) error.""" + return ( + isinstance(exc, httpx.HTTPStatusError) + and exc.response is not None + and int(exc.response.status_code) == HTTP_STATUS_REQUEST_ENTITY_TOO_LARGE + ) + + +def approx_event_size_bytes(event: dict[str, Any]) -> int | None: + """Compact serialized size of an event, or None if it can't be serialized.""" + try: + return len(_json.dumps(event, separators=(",", ":"), ensure_ascii=False)) + except Exception: + return None + + +def count_attachment_refs(event: dict[str, Any]) -> int: + """Count braintrust_attachment references anywhere within an event.""" + + def _walk(v: Any) -> int: + if isinstance(v, dict): + if v.get("type") == "braintrust_attachment" and isinstance( + v.get("key"), str + ): + return 1 + return sum(_walk(x) for x in v.values()) + if isinstance(v, list): + return sum(_walk(x) for x in v) + return 0 + + return _walk(event) + + +def dump_oversize_event_summary( + *, + out_dir: Path, + filename_prefix: str, + event_label: str, + dest_id_field: str, + dest_id_value: str | None, + cursor: str | None, + event: dict[str, Any], + error: Exception, + logger: Any, +) -> None: + """Write a diagnostic summary for a single event that cannot be inserted (413). + + Pure side effect: writes a JSON file under ``out_dir`` and logs. Never touches + stream state, the seen-ids DB, or the insert flow, and never raises (a failed + write is logged and swallowed so it cannot abort the streaming loop). + """ + event_id = event.get("id") + safe_id = str(event_id) if isinstance(event_id, str) and event_id else "unknown" + root_span_id = event.get("root_span_id") + span_id = event.get("span_id") + approx_size = approx_event_size_bytes(event) + attachment_refs = count_attachment_refs(event) + path = out_dir / f"{filename_prefix}{safe_id}.json" + summary = { + "error": str(error), + "cursor": cursor, + dest_id_field: dest_id_value, + "event_id": event.get("id"), + "root_span_id": root_span_id, + "span_id": span_id, + "created": event.get("created"), + "approx_size_bytes": approx_size, + "attachment_refs": attachment_refs, + "top_level_keys": sorted(list(event.keys())), + } + try: + with open(path, "w") as f: + _json.dump(summary, f, indent=2) + logger.error( + f"Oversize {event_label} isolated (413). This specific event cannot be inserted.", + summary_path=str(path), + event_id=safe_id, + root_span_id=root_span_id, + span_id=span_id, + approx_size_bytes=approx_size, + attachment_refs=attachment_refs, + cursor=cursor, + ) + except Exception: + logger.error( + f"Oversize {event_label} isolated; failed to write summary", + event_id=safe_id, + root_span_id=root_span_id, + span_id=span_id, + cursor=cursor, + ) + @dataclass class EventsStreamState: diff --git a/tests/unit/test_streaming_helpers.py b/tests/unit/test_streaming_helpers.py new file mode 100644 index 0000000..2b73659 --- /dev/null +++ b/tests/unit/test_streaming_helpers.py @@ -0,0 +1,114 @@ +"""Unit tests for the shared streaming helpers in streaming_utils. + +These were previously duplicated as private staticmethods/methods across the +logs, dataset, and experiment migrators. They are now single shared functions; +this pins their behavior so the three call sites stay consistent. +""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any + +import httpx + +from braintrust_migrate.streaming_utils import ( + approx_event_size_bytes, + count_attachment_refs, + dump_oversize_event_summary, + is_http_413, +) + + +def _http_status_error(status: int) -> httpx.HTTPStatusError: + request = httpx.Request("POST", "https://api.example/logs3") + response = httpx.Response(status, request=request) + return httpx.HTTPStatusError("boom", request=request, response=response) + + +def test_is_http_413(): + assert is_http_413(_http_status_error(413)) is True + assert is_http_413(_http_status_error(400)) is False + assert is_http_413(ValueError("nope")) is False + + +def test_approx_event_size_bytes(): + assert approx_event_size_bytes({"a": "hello"}) == len( + json.dumps({"a": "hello"}, separators=(",", ":")) + ) + # Non-serializable values return None rather than raising. + assert approx_event_size_bytes({"a": {1, 2, 3}}) is None + + +def test_count_attachment_refs_walks_nested_structures(): + ref = {"type": "braintrust_attachment", "key": "k", "filename": "f", "content_type": "application/json"} + event = { + "input": {"a": ref, "b": [ref, {"c": ref}]}, # 3 refs + "metadata": {"type": "braintrust_attachment"}, # missing key -> not counted + "output": "plain", + } + assert count_attachment_refs(event) == 3 + + +class _RecordingLogger: + def __init__(self) -> None: + self.calls: list[tuple[str, dict[str, Any]]] = [] + + def error(self, message: str, **kwargs: Any) -> None: + self.calls.append((message, kwargs)) + + +def test_dump_oversize_event_summary_writes_file_and_logs(tmp_path: Path): + logger = _RecordingLogger() + event = { + "id": "evt-1", + "root_span_id": "rs", + "span_id": "sp", + "created": "2023-01-01T00:00:00Z", + "input": {"big": "x" * 100}, + } + dump_oversize_event_summary( + out_dir=tmp_path, + filename_prefix="oversize_dataset_event_", + event_label="dataset event", + dest_id_field="dest_dataset_id", + dest_id_value="dest-123", + cursor="cur", + event=event, + error=RuntimeError("too big"), + logger=logger, + ) + + path = tmp_path / "oversize_dataset_event_evt-1.json" + assert path.exists() + summary = json.loads(path.read_text()) + # Per-resource dest-id key name is preserved (not flattened to a generic key). + assert summary["dest_dataset_id"] == "dest-123" + assert summary["event_id"] == "evt-1" + assert summary["cursor"] == "cur" + assert summary["attachment_refs"] == 0 + assert summary["top_level_keys"] == sorted(event.keys()) + + assert len(logger.calls) == 1 + message, kwargs = logger.calls[0] + # event_label drives the message wording. + assert message.startswith("Oversize dataset event isolated (413).") + assert kwargs["event_id"] == "evt-1" + + +def test_dump_oversize_event_summary_unknown_id_and_prefix(tmp_path: Path): + logger = _RecordingLogger() + dump_oversize_event_summary( + out_dir=tmp_path, + filename_prefix="oversize_project_logs_event_", + event_label="event", + dest_id_field="dest_project_id", + dest_id_value="p1", + cursor=None, + event={"input": "x"}, # no id + error=RuntimeError("x"), + logger=logger, + ) + assert (tmp_path / "oversize_project_logs_event_unknown.json").exists() + assert logger.calls[0][0].startswith("Oversize event isolated (413).") From 1736d152fbb63c3cfb374dc7606816898c852627 Mon Sep 17 00:00:00 2001 From: Doug Guthrie Date: Mon, 8 Jun 2026 09:30:45 -0600 Subject: [PATCH 2/8] Delete dead base-class methods and fix bogus failed_at timestamp (R12) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Remove _get_client_resource_attr and _handle_api_response_to_list from ResourceMigrator; both had zero callers anywhere in the package. - record_failure stored failed_at as str(Path(__file__).stat().st_mtime) — base.py's own mtime, a static value identical for every failure. Replace with datetime.now(UTC).isoformat() so the recorded time is real. No test referenced these; full suite passes unchanged. Co-Authored-By: Claude Opus 4.8 (1M context) --- braintrust_migrate/resources/base.py | 60 +--------------------------- 1 file changed, 2 insertions(+), 58 deletions(-) diff --git a/braintrust_migrate/resources/base.py b/braintrust_migrate/resources/base.py index 6ca89e5..7747a4d 100644 --- a/braintrust_migrate/resources/base.py +++ b/braintrust_migrate/resources/base.py @@ -5,6 +5,7 @@ import json from abc import ABC, abstractmethod from dataclasses import dataclass, field +from datetime import UTC, datetime from pathlib import Path from typing import Any, Generic, TypeVar @@ -258,7 +259,7 @@ def record_failure(self, source_id: str, error: str) -> None: # Store error in metadata self.state.metadata[source_id] = { "error": error, - "failed_at": str(Path(__file__).stat().st_mtime), # Timestamp + "failed_at": datetime.now(UTC).isoformat(), } self._logger.error( @@ -316,63 +317,6 @@ async def migrate_resource(self, resource: T) -> str: """ pass - def _get_client_resource_attr(self, client, resource_type: str): - """Get the resource attribute from a client (e.g., client.datasets). - - Args: - client: Braintrust client instance - resource_type: Resource type name (e.g., 'datasets', 'experiments') - - Returns: - Resource client attribute - """ - return getattr(client.client, resource_type) - - async def _handle_api_response_to_list(self, response) -> list[T]: - """Convert various API response formats to a list. - - Handles: - - Async iterators - - Paginated responses with .objects - - Direct lists - - Args: - response: API response in various formats - - Returns: - List of resources - """ - # Handle None or empty response - if response is None: - return [] - - # Handle async iterator - if hasattr(response, "__aiter__"): - result_list = [] - async for item in response: - result_list.append(item) - return result_list - - # Handle paginated response with objects - elif hasattr(response, "objects"): - return list(response.objects) - - # Handle already a list - elif isinstance(response, list): - return response - - # Handle direct iterable (convert to list) - else: - try: - return list(response) - except (TypeError, ValueError) as e: - self._logger.warning( - "Could not convert API response to list", - response_type=type(response).__name__, - error=str(e), - ) - return [] - async def _list_resources_with_client( self, client, From 8aa01697ce8f04ab26a7d10065329a2b8bceaa2c Mon Sep 17 00:00:00 2001 From: Doug Guthrie Date: Mon, 8 Jun 2026 09:32:06 -0600 Subject: [PATCH 3/8] Route project list/create through with_retry (R11) list_projects (per-page) and create_project called raw_request directly, bypassing the adaptive retry/backoff that every other API call in the tool uses. Wrap both in with_retry so transient 429/5xx during project discovery and creation are retried consistently. Happy-path behavior is unchanged (with_retry calls the request once on success); only transient-failure behavior improves. Project methods are mocked wholesale in tests, so the suite passes unchanged. Co-Authored-By: Claude Opus 4.8 (1M context) --- braintrust_migrate/client.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/braintrust_migrate/client.py b/braintrust_migrate/client.py index a2f962f..28f677f 100644 --- a/braintrust_migrate/client.py +++ b/braintrust_migrate/client.py @@ -161,7 +161,12 @@ async def list_projects( if org_name is not None: params["org_name"] = org_name - resp = await self.raw_request("GET", "/v1/project", params=params) + resp = await self.with_retry( + "list_projects", + lambda params=params: self.raw_request( + "GET", "/v1/project", params=params + ), + ) if not isinstance(resp, dict): raise BraintrustAPIError(f"Unexpected project list response: {type(resp)}") objs = resp.get("objects") @@ -200,7 +205,10 @@ async def create_project( if description: payload["description"] = description - resp = await self.raw_request("POST", "/v1/project", json=payload) + resp = await self.with_retry( + "create_project", + lambda: self.raw_request("POST", "/v1/project", json=payload), + ) if not isinstance(resp, dict): raise BraintrustAPIError(f"Unexpected create project response: {type(resp)}") self._maybe_capture_org_id(resp) From 47615d8f06ae9915d394998c893b2d90767aeeb6 Mon Sep 17 00:00:00 2001 From: Doug Guthrie Date: Mon, 8 Jun 2026 10:01:44 -0600 Subject: [PATCH 4/8] Centralize streaming config; drop dead _insert_max_bytes block (R5) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The three streaming migrators each re-read byte-batching config with an identical try/except, and each computed self._insert_max_bytes from insert_max_request_bytes * insert_request_headroom_ratio — a value that was never read anywhere (dead), and whose fallback default (0.5) disagreed with the canonical MigrationConfig default (0.75). - Delete the dead _insert_max_bytes computation from all three migrators (removes the drift entirely rather than "fixing" an unused default). - Add StreamingConfig.resolve(source, dest) in streaming_utils as the single place that resolves the flush/fetch knobs, with module-constant defaults (STREAMING_FLUSH_MAX_ROWS/BYTES, STREAMING_MAX_EVENT_BYTES, STREAMING_EVENT_FETCH_GROUP_SIZE). The migrator ClassVars now alias those constants. Test mechanism change (not a product-behavior change): the three resume-after-insert-failure tests forced per-row flushing by monkeypatching the SDK_FLUSH_MAX_ROWS ClassVar; they now set events_flush_max_rows=1 on a MigrationConfig — the actual user-facing lever — which is both cleaner and lets resolve use plain module-constant defaults. Full suite: 293 passing. Co-Authored-By: Claude Opus 4.8 (1M context) --- braintrust_migrate/resources/datasets.py | 45 ++++++------------ braintrust_migrate/resources/experiments.py | 45 ++++++------------ braintrust_migrate/resources/logs.py | 36 +++++--------- braintrust_migrate/streaming_utils.py | 47 +++++++++++++++++++ tests/unit/test_dataset_streaming_resume.py | 6 +-- .../unit/test_experiment_streaming_resume.py | 6 +-- tests/unit/test_logs_streaming_resume.py | 7 +-- 7 files changed, 96 insertions(+), 96 deletions(-) diff --git a/braintrust_migrate/resources/datasets.py b/braintrust_migrate/resources/datasets.py index 6b9c36d..fe81afa 100644 --- a/braintrust_migrate/resources/datasets.py +++ b/braintrust_migrate/resources/datasets.py @@ -18,10 +18,14 @@ from braintrust_migrate.resources.base import MigrationResult, ResourceMigrator from braintrust_migrate.sdk_logs import SDKDatasetWriter from braintrust_migrate.streaming_utils import ( + STREAMING_EVENT_FETCH_GROUP_SIZE, + STREAMING_FLUSH_MAX_BYTES, + STREAMING_FLUSH_MAX_ROWS, + STREAMING_MAX_EVENT_BYTES, EventsStreamState, SeenIdsDB, + StreamingConfig, build_btql_sorted_page_query, - coerce_int_config, dump_oversize_event_summary, is_http_413, stream_btql_sorted_events_buffered, @@ -41,12 +45,12 @@ class DatasetMigrator(ResourceMigrator[dict]): Uses raw API requests instead of SDK to avoid model dependencies. """ - SDK_FLUSH_MAX_ROWS: ClassVar[int] = 5_000 - SDK_FLUSH_MAX_BYTES: ClassVar[int] = 25 * 1024 * 1024 - DEFAULT_EVENT_FETCH_GROUP_SIZE: ClassVar[int] = 25 + SDK_FLUSH_MAX_ROWS: ClassVar[int] = STREAMING_FLUSH_MAX_ROWS + SDK_FLUSH_MAX_BYTES: ClassVar[int] = STREAMING_FLUSH_MAX_BYTES + DEFAULT_EVENT_FETCH_GROUP_SIZE: ClassVar[int] = STREAMING_EVENT_FETCH_GROUP_SIZE # Spill individual rows above this size into JSON attachments so they fit # under Braintrust's ~20MB per-span logging upload limit (logs3/overflow). - MAX_EVENT_BYTES: ClassVar[int] = 18 * 1024 * 1024 + MAX_EVENT_BYTES: ClassVar[int] = STREAMING_MAX_EVENT_BYTES def __init__( self, @@ -77,36 +81,15 @@ def __init__( ), ) - # Byte-aware insert batching config (best-effort; falls back to count-only if missing). - cfg = getattr(self.dest_client, "migration_config", None) or getattr( - self.source_client, "migration_config", None - ) - try: - max_req = int(getattr(cfg, "insert_max_request_bytes", 6 * 1024 * 1024)) - headroom = float(getattr(cfg, "insert_request_headroom_ratio", 0.5)) - if headroom <= 0: - raise ValueError("headroom must be > 0") - self._insert_max_bytes: int | None = int(max_req * headroom) - except Exception: - self._insert_max_bytes = None - self._sdk_flush_max_rows = coerce_int_config( - cfg, - "events_flush_max_rows", - self.SDK_FLUSH_MAX_ROWS, - minimum=1, - ) - self._sdk_flush_max_bytes = int(self.SDK_FLUSH_MAX_BYTES) - self._event_fetch_group_size = coerce_int_config( - cfg, - "events_fetch_group_size", - self.DEFAULT_EVENT_FETCH_GROUP_SIZE, - minimum=1, - ) + stream_cfg = StreamingConfig.resolve(self.source_client, self.dest_client) + self._sdk_flush_max_rows = stream_cfg.sdk_flush_max_rows + self._sdk_flush_max_bytes = stream_cfg.sdk_flush_max_bytes + self._event_fetch_group_size = stream_cfg.event_fetch_group_size # Always-on oversize-field spilling (see OversizeFieldSpiller). Only hits # the network for rows that exceed the per-span upload limit. self._spiller = OversizeFieldSpiller( dest_client=self.dest_client, - max_event_bytes=int(self.MAX_EVENT_BYTES), + max_event_bytes=stream_cfg.max_event_bytes, ) @property diff --git a/braintrust_migrate/resources/experiments.py b/braintrust_migrate/resources/experiments.py index 36f092e..bda46cb 100644 --- a/braintrust_migrate/resources/experiments.py +++ b/braintrust_migrate/resources/experiments.py @@ -16,10 +16,14 @@ from braintrust_migrate.resources.base import MigrationResult, ResourceMigrator from braintrust_migrate.sdk_logs import SDKExperimentWriter from braintrust_migrate.streaming_utils import ( + STREAMING_EVENT_FETCH_GROUP_SIZE, + STREAMING_FLUSH_MAX_BYTES, + STREAMING_FLUSH_MAX_ROWS, + STREAMING_MAX_EVENT_BYTES, EventsStreamState, SeenIdsDB, + StreamingConfig, build_btql_sorted_page_query, - coerce_int_config, dump_oversize_event_summary, is_http_413, stream_btql_sorted_events_buffered, @@ -37,12 +41,12 @@ class ExperimentMigrator(ResourceMigrator[dict]): Uses raw API requests instead of SDK to avoid model dependencies. """ - SDK_FLUSH_MAX_ROWS: ClassVar[int] = 5_000 - SDK_FLUSH_MAX_BYTES: ClassVar[int] = 25 * 1024 * 1024 - DEFAULT_EVENT_FETCH_GROUP_SIZE: ClassVar[int] = 25 + SDK_FLUSH_MAX_ROWS: ClassVar[int] = STREAMING_FLUSH_MAX_ROWS + SDK_FLUSH_MAX_BYTES: ClassVar[int] = STREAMING_FLUSH_MAX_BYTES + DEFAULT_EVENT_FETCH_GROUP_SIZE: ClassVar[int] = STREAMING_EVENT_FETCH_GROUP_SIZE # Spill individual rows above this size into JSON attachments so they fit # under Braintrust's ~20MB per-span logging upload limit (logs3/overflow). - MAX_EVENT_BYTES: ClassVar[int] = 18 * 1024 * 1024 + MAX_EVENT_BYTES: ClassVar[int] = STREAMING_MAX_EVENT_BYTES def __init__( self, @@ -73,36 +77,15 @@ def __init__( ), ) - # Byte-aware insert batching config (best-effort; falls back to count-only if missing). - cfg = getattr(self.dest_client, "migration_config", None) or getattr( - self.source_client, "migration_config", None - ) - try: - max_req = int(getattr(cfg, "insert_max_request_bytes", 6 * 1024 * 1024)) - headroom = float(getattr(cfg, "insert_request_headroom_ratio", 0.5)) - if headroom <= 0: - raise ValueError("headroom must be > 0") - self._insert_max_bytes: int | None = int(max_req * headroom) - except Exception: - self._insert_max_bytes = None - self._sdk_flush_max_rows = coerce_int_config( - cfg, - "events_flush_max_rows", - self.SDK_FLUSH_MAX_ROWS, - minimum=1, - ) - self._sdk_flush_max_bytes = int(self.SDK_FLUSH_MAX_BYTES) + stream_cfg = StreamingConfig.resolve(self.source_client, self.dest_client) + self._sdk_flush_max_rows = stream_cfg.sdk_flush_max_rows + self._sdk_flush_max_bytes = stream_cfg.sdk_flush_max_bytes + self._event_fetch_group_size = stream_cfg.event_fetch_group_size # Always-on oversize-field spilling (see OversizeFieldSpiller). Only hits # the network for rows that exceed the per-span upload limit. self._spiller = OversizeFieldSpiller( dest_client=self.dest_client, - max_event_bytes=int(self.MAX_EVENT_BYTES), - ) - self._event_fetch_group_size = coerce_int_config( - cfg, - "events_fetch_group_size", - self.DEFAULT_EVENT_FETCH_GROUP_SIZE, - minimum=1, + max_event_bytes=stream_cfg.max_event_bytes, ) @property diff --git a/braintrust_migrate/resources/logs.py b/braintrust_migrate/resources/logs.py index 987e32a..3fb810d 100644 --- a/braintrust_migrate/resources/logs.py +++ b/braintrust_migrate/resources/logs.py @@ -30,9 +30,12 @@ from braintrust_migrate.resources.base import MigrationState, ResourceMigrator from braintrust_migrate.sdk_logs import SDKProjectLogsWriter from braintrust_migrate.streaming_utils import ( + STREAMING_FLUSH_MAX_BYTES, + STREAMING_FLUSH_MAX_ROWS, + STREAMING_MAX_EVENT_BYTES, SeenIdsDB, + StreamingConfig, build_btql_sorted_page_query, - coerce_int_config, dump_oversize_event_summary, is_http_413, ) @@ -107,11 +110,11 @@ def to_dict(self) -> dict[str, Any]: class LogsMigrator(ResourceMigrator[dict[str, Any]]): """Streaming migrator for Braintrust project logs.""" - SDK_FLUSH_MAX_ROWS: ClassVar[int] = 5_000 - SDK_FLUSH_MAX_BYTES: ClassVar[int] = 25 * 1024 * 1024 + SDK_FLUSH_MAX_ROWS: ClassVar[int] = STREAMING_FLUSH_MAX_ROWS + SDK_FLUSH_MAX_BYTES: ClassVar[int] = STREAMING_FLUSH_MAX_BYTES # Spill individual rows above this size into JSON attachments so they fit # under Braintrust's ~20MB per-span logging upload limit (logs3/overflow). - MAX_EVENT_BYTES: ClassVar[int] = 18 * 1024 * 1024 + MAX_EVENT_BYTES: ClassVar[int] = STREAMING_MAX_EVENT_BYTES _INSERT_FIELDS: ClassVar[set[str]] = { "input", @@ -161,16 +164,9 @@ def __init__( self._logger = logger.bind(migrator=self.__class__.__name__) self._sdk_logs_writer: SDKProjectLogsWriter | None = None - cfg = getattr(self.dest_client, "migration_config", None) or getattr( - self.source_client, "migration_config", None - ) - self._sdk_flush_max_rows = coerce_int_config( - cfg, - "events_flush_max_rows", - self.SDK_FLUSH_MAX_ROWS, - minimum=1, - ) - self._sdk_flush_max_bytes = int(self.SDK_FLUSH_MAX_BYTES) + stream_cfg = StreamingConfig.resolve(self.source_client, self.dest_client) + self._sdk_flush_max_rows = stream_cfg.sdk_flush_max_rows + self._sdk_flush_max_bytes = stream_cfg.sdk_flush_max_bytes self._stream_state_path = self.checkpoint_dir / "logs_streaming_state.json" self._stream_state = _LogsStreamingState.from_path(self._stream_state_path) @@ -203,22 +199,12 @@ def __init__( # Always-on: rows larger than the per-span upload limit get their biggest # fields spilled into JSON attachments. The spiller only touches the # network for rows that actually exceed the cap, so it is free otherwise. - self._max_event_bytes = int(self.MAX_EVENT_BYTES) + self._max_event_bytes = stream_cfg.max_event_bytes self._spiller = OversizeFieldSpiller( dest_client=self.dest_client, max_event_bytes=self._max_event_bytes, ) - # Byte-aware insert batching config (best-effort; falls back to count-only if missing). - try: - max_req = int(getattr(cfg, "insert_max_request_bytes", 6 * 1024 * 1024)) - headroom = float(getattr(cfg, "insert_request_headroom_ratio", 0.5)) - if headroom <= 0: - raise ValueError("headroom must be > 0") - self._insert_max_bytes: int | None = int(max_req * headroom) - except Exception: - self._insert_max_bytes = None - @property def resource_name(self) -> str: return "Logs" diff --git a/braintrust_migrate/streaming_utils.py b/braintrust_migrate/streaming_utils.py index fcc5d61..775052f 100644 --- a/braintrust_migrate/streaming_utils.py +++ b/braintrust_migrate/streaming_utils.py @@ -228,6 +228,53 @@ def coerce_int_config( return value +# Default streaming flush/fetch knobs. Single source of truth: the migrators +# expose these as ClassVars for back-compat and tests. +STREAMING_FLUSH_MAX_ROWS = 5_000 +STREAMING_FLUSH_MAX_BYTES = 25 * 1024 * 1024 +STREAMING_MAX_EVENT_BYTES = 18 * 1024 * 1024 +STREAMING_EVENT_FETCH_GROUP_SIZE = 25 + + +@dataclass(frozen=True) +class StreamingConfig: + """Resolved streaming flush/fetch configuration for the event migrators. + + Replaces the byte-batching config block that was duplicated (with a drifting + headroom default) across the logs, dataset, and experiment migrators. + """ + + sdk_flush_max_rows: int + sdk_flush_max_bytes: int + max_event_bytes: int + event_fetch_group_size: int + + @classmethod + def resolve(cls, source_client: Any, dest_client: Any) -> StreamingConfig: + """Resolve config from the dest client's migration_config (or source's). + + ``events_flush_max_rows`` / ``events_fetch_group_size`` are the only + env-overridable knobs; the byte caps are fixed. Falls back to the module + defaults and is tolerant of lightweight test doubles / missing config. + """ + cfg = getattr(dest_client, "migration_config", None) or getattr( + source_client, "migration_config", None + ) + return cls( + sdk_flush_max_rows=coerce_int_config( + cfg, "events_flush_max_rows", STREAMING_FLUSH_MAX_ROWS, minimum=1 + ), + sdk_flush_max_bytes=STREAMING_FLUSH_MAX_BYTES, + max_event_bytes=STREAMING_MAX_EVENT_BYTES, + event_fetch_group_size=coerce_int_config( + cfg, + "events_fetch_group_size", + STREAMING_EVENT_FETCH_GROUP_SIZE, + minimum=1, + ), + ) + + def build_btql_sorted_page_query( *, from_expr: str, diff --git a/tests/unit/test_dataset_streaming_resume.py b/tests/unit/test_dataset_streaming_resume.py index 6ef9a7f..d79bf22 100644 --- a/tests/unit/test_dataset_streaming_resume.py +++ b/tests/unit/test_dataset_streaming_resume.py @@ -6,6 +6,7 @@ import pytest import braintrust_migrate.resources.datasets as datasets_module +from braintrust_migrate.config import MigrationConfig from braintrust_migrate.resources.datasets import DatasetMigrator @@ -96,10 +97,10 @@ async def test_dataset_streaming_resume_after_insert_failure(tmp_path: Path) -> source = _StubClient(page1=page1, page2=page2) dest = _StubClient() dest.fail_on_insert_call = 2 + # Flush one row at a time via the real config lever so the second insert fails. + dest.migration_config = MigrationConfig(events_flush_max_rows=1) original_writer = datasets_module.SDKDatasetWriter - original_flush_max_rows = datasets_module.DatasetMigrator.SDK_FLUSH_MAX_ROWS datasets_module.SDKDatasetWriter = _FakeSDKDatasetWriter - datasets_module.DatasetMigrator.SDK_FLUSH_MAX_ROWS = 1 try: migrator = DatasetMigrator( @@ -126,4 +127,3 @@ async def test_dataset_streaming_resume_after_insert_failure(tmp_path: Path) -> assert inserted_all == ["a", "b"] finally: datasets_module.SDKDatasetWriter = original_writer - datasets_module.DatasetMigrator.SDK_FLUSH_MAX_ROWS = original_flush_max_rows diff --git a/tests/unit/test_experiment_streaming_resume.py b/tests/unit/test_experiment_streaming_resume.py index 1c06601..dfbd1c5 100644 --- a/tests/unit/test_experiment_streaming_resume.py +++ b/tests/unit/test_experiment_streaming_resume.py @@ -6,6 +6,7 @@ import pytest import braintrust_migrate.resources.experiments as experiments_module +from braintrust_migrate.config import MigrationConfig from braintrust_migrate.resources.experiments import ExperimentMigrator @@ -100,10 +101,10 @@ async def test_experiment_streaming_resume_after_insert_failure(tmp_path: Path) source = _StubClient(page1=page1, page2=page2) dest = _StubClient() dest.fail_on_insert_call = 2 # fail on second insert during first run + # Flush one row at a time via the real config lever so the second insert fails. + dest.migration_config = MigrationConfig(events_flush_max_rows=1) original_writer = experiments_module.SDKExperimentWriter - original_flush_max_rows = experiments_module.ExperimentMigrator.SDK_FLUSH_MAX_ROWS experiments_module.SDKExperimentWriter = _FakeSDKExperimentWriter - experiments_module.ExperimentMigrator.SDK_FLUSH_MAX_ROWS = 1 try: migrator = ExperimentMigrator( @@ -130,4 +131,3 @@ async def test_experiment_streaming_resume_after_insert_failure(tmp_path: Path) assert inserted_all == ["a", "b"] finally: experiments_module.SDKExperimentWriter = original_writer - experiments_module.ExperimentMigrator.SDK_FLUSH_MAX_ROWS = original_flush_max_rows diff --git a/tests/unit/test_logs_streaming_resume.py b/tests/unit/test_logs_streaming_resume.py index 598fb7c..c55db62 100644 --- a/tests/unit/test_logs_streaming_resume.py +++ b/tests/unit/test_logs_streaming_resume.py @@ -6,6 +6,7 @@ import pytest import braintrust_migrate.resources.logs as logs_module +from braintrust_migrate.config import MigrationConfig from braintrust_migrate.resources.logs import LogsMigrator @@ -103,10 +104,11 @@ async def test_logs_migrator_resume_after_insert_failure(tmp_path: Path) -> None source = _StubClient(page1=page1, page2=page2) dest = _StubClient(page1=page1, page2=page2) dest.fail_on_insert_call = 2 # fail on second insert (page2) during first run + # Flush one row at a time via the real config lever so the second insert + # (page2) is the one that fails. + dest.migration_config = MigrationConfig(events_flush_max_rows=1) original_writer = logs_module.SDKProjectLogsWriter - original_flush_max_rows = logs_module.LogsMigrator.SDK_FLUSH_MAX_ROWS logs_module.SDKProjectLogsWriter = _FakeSDKProjectLogsWriter - logs_module.LogsMigrator.SDK_FLUSH_MAX_ROWS = 1 try: migrator = LogsMigrator( @@ -132,4 +134,3 @@ async def test_logs_migrator_resume_after_insert_failure(tmp_path: Path) -> None assert inserted_all == ["a", "b"] finally: logs_module.SDKProjectLogsWriter = original_writer - logs_module.LogsMigrator.SDK_FLUSH_MAX_ROWS = original_flush_max_rows From 2dc98aea941d1771d558cf31f414b5c929bbbf17 Mon Sep 17 00:00:00 2001 From: Doug Guthrie Date: Mon, 8 Jun 2026 10:07:27 -0600 Subject: [PATCH 5/8] Extract shared UserResolver; stop conflating 404 with transient failures (R8) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit _get_source_user_email, _find_dest_user_id_by_email, and _invite_user_to_dest_org were duplicated near-verbatim across the ACL and group migrators (only the with_retry op-name strings differed), and each swallowed every exception and cached None permanently — so a transient auth/network/5xx failure became a permanent "no such user". Extract a single UserResolver (braintrust_migrate/user_resolver.py) used by both migrators. with_retry re-raises the original exception, so a 404 propagates as httpx.HTTPStatusError(404): the resolver now caches the negative result only on a genuine 404, and on any other failure logs a warning and does NOT cache, leaving a later attempt able to succeed. The duplicated methods and their three per-migrator caches are removed; both migrators delegate to the resolver. Op-name suffix keeps telemetry distinguishable per caller. Unit suite (293) and the ACL/group user-mapping integration flow tests pass; the 404-vs-transient change only affects the error path (happy path is unchanged). Co-Authored-By: Claude Opus 4.8 (1M context) --- braintrust_migrate/resources/acls.py | 111 ++-------------- braintrust_migrate/resources/groups.py | 112 ++-------------- braintrust_migrate/user_resolver.py | 174 +++++++++++++++++++++++++ 3 files changed, 192 insertions(+), 205 deletions(-) create mode 100644 braintrust_migrate/user_resolver.py diff --git a/braintrust_migrate/resources/acls.py b/braintrust_migrate/resources/acls.py index b2696d7..41f0942 100644 --- a/braintrust_migrate/resources/acls.py +++ b/braintrust_migrate/resources/acls.py @@ -8,6 +8,7 @@ MigrationResult, ResourceMigrator, ) +from braintrust_migrate.user_resolver import UserResolver class ACLMigrator(ResourceMigrator[dict]): @@ -32,9 +33,7 @@ def resource_name(self) -> str: def __init__(self, source_client, dest_client, checkpoint_dir, batch_size: int = 100): super().__init__(source_client, dest_client, checkpoint_dir, batch_size=batch_size) - self._source_user_email_cache: dict[str, str | None] = {} - self._dest_user_id_by_email_cache: dict[str, str | None] = {} - self._invited_user_emails: set[str] = set() + self._user_resolver = UserResolver(self.source_client, self.dest_client) @property def supported_object_types(self) -> set[str]: @@ -227,120 +226,26 @@ def _acl_auto_invite_enabled(self) -> bool: return value.strip().lower() in {"1", "true", "yes", "y", "on"} return False - async def _get_source_user_email(self, source_user_id: str) -> str | None: - """Get source user email by user ID.""" - if source_user_id in self._source_user_email_cache: - return self._source_user_email_cache[source_user_id] - - try: - response = await self.source_client.with_retry( - "get_source_user", - lambda uid=source_user_id: self.source_client.raw_request( - "GET", - f"/v1/user/{uid}", - ), - ) - email = response.get("email") if isinstance(response, dict) else None - email = email.strip().lower() if isinstance(email, str) and email.strip() else None - self._source_user_email_cache[source_user_id] = email - return email - except Exception: - self._source_user_email_cache[source_user_id] = None - return None - - async def _find_dest_user_id_by_email( - self, email: str, *, force_refresh: bool = False - ) -> str | None: - """Find destination user ID by email.""" - normalized_email = email.strip().lower() - if force_refresh: - self._dest_user_id_by_email_cache.pop(normalized_email, None) - if normalized_email in self._dest_user_id_by_email_cache: - return self._dest_user_id_by_email_cache[normalized_email] - - try: - response = await self.dest_client.with_retry( - "list_dest_users_by_email", - lambda e=normalized_email: self.dest_client.raw_request( - "GET", - "/v1/user", - params={"email": e, "limit": 100}, - ), - ) - - if isinstance(response, dict): - objects = response.get("objects", []) - elif isinstance(response, list): - objects = response - else: - objects = [] - - dest_user_id = None - for user in objects: - if not isinstance(user, dict): - continue - user_email = user.get("email") - user_id = user.get("id") - if ( - isinstance(user_email, str) - and user_email.strip().lower() == normalized_email - and isinstance(user_id, str) - and user_id - ): - dest_user_id = user_id - break - - self._dest_user_id_by_email_cache[normalized_email] = dest_user_id - return dest_user_id - except Exception: - self._dest_user_id_by_email_cache[normalized_email] = None - return None - - async def _invite_user_to_dest_org(self, email: str) -> bool: - """Invite a user to destination org via organization members API.""" - normalized_email = email.strip().lower() - if normalized_email in self._invited_user_emails: - return True - - try: - await self.dest_client.with_retry( - "invite_user_to_dest_org", - lambda e=normalized_email: self.dest_client.raw_request( - "PATCH", - "/v1/organization/members", - json={ - "invite_users": { - "emails": [e], - "send_invite_emails": False, - } - }, - ), - ) - self._invited_user_emails.add(normalized_email) - # Invalidate cache in case it was previously absent. - self._dest_user_id_by_email_cache.pop(normalized_email, None) - return True - except Exception: - return False - async def _resolve_acl_user_id(self, source_user_id: str) -> str | None: """Resolve ACL user_id by source/destination email matching.""" existing_mapping = self.state.id_mapping.get(source_user_id) if existing_mapping: return existing_mapping - source_email = await self._get_source_user_email(source_user_id) + source_email = await self._user_resolver.source_user_email(source_user_id) if not source_email: return None - dest_user_id = await self._find_dest_user_id_by_email(source_email) + dest_user_id = await self._user_resolver.find_dest_user_id_by_email( + source_email + ) if not dest_user_id and self._acl_auto_invite_enabled(): - invited = await self._invite_user_to_dest_org(source_email) + invited = await self._user_resolver.invite_user_to_dest_org(source_email) if invited: # Membership propagation may be eventual; retry lookup briefly. post_invite_attempts = 4 for attempt in range(post_invite_attempts): - dest_user_id = await self._find_dest_user_id_by_email( + dest_user_id = await self._user_resolver.find_dest_user_id_by_email( source_email, force_refresh=True, ) diff --git a/braintrust_migrate/resources/groups.py b/braintrust_migrate/resources/groups.py index a1ab7c0..bbf91a3 100644 --- a/braintrust_migrate/resources/groups.py +++ b/braintrust_migrate/resources/groups.py @@ -3,6 +3,7 @@ import asyncio from braintrust_migrate.resources.base import ResourceMigrator +from braintrust_migrate.user_resolver import UserResolver class GroupMigrator(ResourceMigrator[dict]): @@ -20,9 +21,9 @@ class GroupMigrator(ResourceMigrator[dict]): def __init__(self, source_client, dest_client, checkpoint_dir, batch_size: int = 100): super().__init__(source_client, dest_client, checkpoint_dir, batch_size=batch_size) - self._source_user_email_cache: dict[str, str | None] = {} - self._dest_user_id_by_email_cache: dict[str, str | None] = {} - self._invited_user_emails: set[str] = set() + self._user_resolver = UserResolver( + self.source_client, self.dest_client, op_suffix="_for_group_members" + ) @property def resource_name(self) -> str: @@ -53,118 +54,25 @@ def _group_auto_invite_enabled(self) -> bool: return value.strip().lower() in {"1", "true", "yes", "y", "on"} return False - async def _get_source_user_email(self, source_user_id: str) -> str | None: - """Get source user email by user ID.""" - if source_user_id in self._source_user_email_cache: - return self._source_user_email_cache[source_user_id] - - try: - response = await self.source_client.with_retry( - "get_source_user_for_group_member", - lambda uid=source_user_id: self.source_client.raw_request( - "GET", - f"/v1/user/{uid}", - ), - ) - email = response.get("email") if isinstance(response, dict) else None - email = email.strip().lower() if isinstance(email, str) and email.strip() else None - self._source_user_email_cache[source_user_id] = email - return email - except Exception: - self._source_user_email_cache[source_user_id] = None - return None - - async def _find_dest_user_id_by_email( - self, email: str, *, force_refresh: bool = False - ) -> str | None: - """Find destination user ID by email.""" - normalized_email = email.strip().lower() - if force_refresh: - self._dest_user_id_by_email_cache.pop(normalized_email, None) - if normalized_email in self._dest_user_id_by_email_cache: - return self._dest_user_id_by_email_cache[normalized_email] - - try: - response = await self.dest_client.with_retry( - "list_dest_users_by_email_for_group_members", - lambda e=normalized_email: self.dest_client.raw_request( - "GET", - "/v1/user", - params={"email": e, "limit": 100}, - ), - ) - - if isinstance(response, dict): - objects = response.get("objects", []) - elif isinstance(response, list): - objects = response - else: - objects = [] - - dest_user_id = None - for user in objects: - if not isinstance(user, dict): - continue - user_email = user.get("email") - user_id = user.get("id") - if ( - isinstance(user_email, str) - and user_email.strip().lower() == normalized_email - and isinstance(user_id, str) - and user_id - ): - dest_user_id = user_id - break - - self._dest_user_id_by_email_cache[normalized_email] = dest_user_id - return dest_user_id - except Exception: - self._dest_user_id_by_email_cache[normalized_email] = None - return None - - async def _invite_user_to_dest_org(self, email: str) -> bool: - """Invite a user to destination org via organization members API.""" - normalized_email = email.strip().lower() - if normalized_email in self._invited_user_emails: - return True - - try: - await self.dest_client.with_retry( - "invite_user_to_dest_org_for_group_members", - lambda e=normalized_email: self.dest_client.raw_request( - "PATCH", - "/v1/organization/members", - json={ - "invite_users": { - "emails": [e], - "send_invite_emails": False, - } - }, - ), - ) - self._invited_user_emails.add(normalized_email) - self._dest_user_id_by_email_cache.pop(normalized_email, None) - return True - except Exception: - return False - async def _resolve_group_member_user_id(self, source_user_id: str) -> str | None: """Resolve group member user_id by source/destination email matching.""" existing_mapping = self.state.id_mapping.get(source_user_id) if existing_mapping: return existing_mapping - source_email = await self._get_source_user_email(source_user_id) + source_email = await self._user_resolver.source_user_email(source_user_id) if not source_email: return None - dest_user_id = await self._find_dest_user_id_by_email(source_email) + dest_user_id = await self._user_resolver.find_dest_user_id_by_email( + source_email + ) if not dest_user_id and self._group_auto_invite_enabled(): - invited = await self._invite_user_to_dest_org(source_email) + invited = await self._user_resolver.invite_user_to_dest_org(source_email) if invited: post_invite_attempts = 4 for attempt in range(post_invite_attempts): - dest_user_id = await self._find_dest_user_id_by_email( + dest_user_id = await self._user_resolver.find_dest_user_id_by_email( source_email, force_refresh=True, ) diff --git a/braintrust_migrate/user_resolver.py b/braintrust_migrate/user_resolver.py new file mode 100644 index 0000000..545e3ef --- /dev/null +++ b/braintrust_migrate/user_resolver.py @@ -0,0 +1,174 @@ +"""Shared cross-org user resolution for ACL and group-member migration. + +Both the ACL and group migrators need to map a source user to the matching +destination user by email (and optionally invite missing users). This logic was +previously duplicated near-verbatim in both migrators; it now lives here. + +It also distinguishes a genuine 404 (user not found -> cache the negative result +so we don't refetch) from auth/network/5xx failures (transient -> log and do NOT +cache, so a later attempt can still succeed). The old code swallowed every +exception and cached ``None`` permanently, which turned a transient failure into +a permanent "no such user". +""" + +from __future__ import annotations + +import httpx +import structlog + +from braintrust_migrate.client import BraintrustClient + +logger = structlog.get_logger(__name__) + + +def _is_http_404(exc: Exception) -> bool: + return ( + isinstance(exc, httpx.HTTPStatusError) + and exc.response is not None + and int(exc.response.status_code) == 404 + ) + + +def _normalize_email(email: str) -> str: + return email.strip().lower() + + +class UserResolver: + """Resolve source users to destination users by email, with caching.""" + + def __init__( + self, + source_client: BraintrustClient, + dest_client: BraintrustClient, + *, + op_suffix: str = "", + ) -> None: + self.source_client = source_client + self.dest_client = dest_client + # Suffix appended to with_retry operation names so callers stay + # distinguishable in telemetry (e.g. "_for_group_members"). + self._op_suffix = op_suffix + self._source_user_email_cache: dict[str, str | None] = {} + self._dest_user_id_by_email_cache: dict[str, str | None] = {} + self._invited_user_emails: set[str] = set() + self._logger = logger.bind(component="UserResolver") + + async def source_user_email(self, source_user_id: str) -> str | None: + """Get a source user's (normalized) email by user id.""" + if source_user_id in self._source_user_email_cache: + return self._source_user_email_cache[source_user_id] + + try: + response = await self.source_client.with_retry( + f"get_source_user{self._op_suffix}", + lambda uid=source_user_id: self.source_client.raw_request( + "GET", f"/v1/user/{uid}" + ), + ) + except Exception as e: + if _is_http_404(e): + # Genuinely not found: cache the negative result. + self._source_user_email_cache[source_user_id] = None + return None + # Transient/auth/network failure: do not cache, so a retry can work. + self._logger.warning( + "Failed to fetch source user email; not caching", + source_user_id=source_user_id, + error=str(e), + ) + return None + + email = response.get("email") if isinstance(response, dict) else None + email = ( + email.strip().lower() + if isinstance(email, str) and email.strip() + else None + ) + self._source_user_email_cache[source_user_id] = email + return email + + async def find_dest_user_id_by_email( + self, email: str, *, force_refresh: bool = False + ) -> str | None: + """Find a destination user id by email.""" + normalized_email = _normalize_email(email) + if force_refresh: + self._dest_user_id_by_email_cache.pop(normalized_email, None) + if normalized_email in self._dest_user_id_by_email_cache: + return self._dest_user_id_by_email_cache[normalized_email] + + try: + response = await self.dest_client.with_retry( + f"list_dest_users_by_email{self._op_suffix}", + lambda e=normalized_email: self.dest_client.raw_request( + "GET", "/v1/user", params={"email": e, "limit": 100} + ), + ) + except Exception as e: + if _is_http_404(e): + self._dest_user_id_by_email_cache[normalized_email] = None + return None + self._logger.warning( + "Failed to list dest users by email; not caching", + email=normalized_email, + error=str(e), + ) + return None + + if isinstance(response, dict): + objects = response.get("objects", []) + elif isinstance(response, list): + objects = response + else: + objects = [] + + dest_user_id = None + for user in objects: + if not isinstance(user, dict): + continue + user_email = user.get("email") + user_id = user.get("id") + if ( + isinstance(user_email, str) + and user_email.strip().lower() == normalized_email + and isinstance(user_id, str) + and user_id + ): + dest_user_id = user_id + break + + self._dest_user_id_by_email_cache[normalized_email] = dest_user_id + return dest_user_id + + async def invite_user_to_dest_org(self, email: str) -> bool: + """Invite a user to the destination org via the organization members API.""" + normalized_email = _normalize_email(email) + if normalized_email in self._invited_user_emails: + return True + + try: + await self.dest_client.with_retry( + f"invite_user_to_dest_org{self._op_suffix}", + lambda e=normalized_email: self.dest_client.raw_request( + "PATCH", + "/v1/organization/members", + json={ + "invite_users": { + "emails": [e], + "send_invite_emails": False, + } + }, + ), + ) + except Exception as e: + self._logger.warning( + "Failed to invite user to dest org", + email=normalized_email, + error=str(e), + ) + return False + + self._invited_user_emails.add(normalized_email) + # Invalidate cache in case it was previously absent. + self._dest_user_id_by_email_cache.pop(normalized_email, None) + return True From 64a7202d24b3d1b762f35821bea1670e16534504 Mon Sep 17 00:00:00 2001 From: Doug Guthrie Date: Mon, 8 Jun 2026 10:38:41 -0600 Subject: [PATCH 6/8] Remove dead streaming_pipeline knob and correct the docs (R10a) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit streaming_pipeline (env MIGRATION_STREAMING_PIPELINE) was defined, parsed, and constructed but never read anywhere in the product code — the streaming loop is strictly sequential (fetch page, insert, fetch next). The README, however, documented it as a working, default-on "pipelined event streaming" feature ("prefetch the next BTQL page while inserting the current batch"). Remove the dead config field/env-parse/constructor arg, and correct the README (env table, the Pipelined Event Streaming section, the parallelization diagram, tuning tips, and example .env blocks) so it no longer claims a feature that doesn't exist. Drop the streaming_pipeline assertions from the unit config tests and the field from the live e2e scenario matrix (the two scenarios remain distinct via their concurrency settings; renamed to concurrent/sequential). No runtime behavior change — the knob never affected anything. Implementing real pipelined prefetch remains possible future work. Co-Authored-By: Claude Opus 4.8 (1M context) --- README.md | 25 ++++++------------- braintrust_migrate/config.py | 12 ---------- tests/integration/test_live_e2e_matrix.py | 8 ++----- tests/unit/test_parallelization_config.py | 29 ----------------------- 4 files changed, 9 insertions(+), 65 deletions(-) diff --git a/README.md b/README.md index dec638e..23d45f9 100644 --- a/README.md +++ b/README.md @@ -29,7 +29,7 @@ This tool provides migration capabilities for Braintrust organizations, handling - **Dependency-Aware Migration**: Resources are migrated in an order that respects dependencies (see below) - **Organization Scoping**: AI secrets, roles, and groups migrated once at org level - **Batch Processing**: Configurable batch sizes for optimal performance -- **Multi-Level Parallelization**: Concurrent resource types, concurrent items within a type, and pipelined event streaming (see [Parallelization](#parallelization) below) +- **Multi-Level Parallelization**: Concurrent resource types and concurrent items within a type (see [Parallelization](#parallelization) below) ### Reliability Features - **Retry Logic**: Adaptive retries with exponential backoff + jitter; respects `Retry-After` when rate-limited (429) @@ -154,7 +154,6 @@ All options can be set via environment variables or CLI flags. CLI flags take pr | Environment Variable | CLI Flag | Default | Description | |---------------------|----------|---------|-------------| | `MIGRATION_MAX_CONCURRENT_RESOURCES` | — | `5` | Max concurrent items within a resource type (e.g. 5 experiments migrating at once). Also controls concurrent event streams for datasets/experiments. Range: 1–50 | -| `MIGRATION_STREAMING_PIPELINE` | — | `true` | Prefetch the next BTQL page while inserting the current batch, overlapping source reads with destination writes | | `MIGRATION_MAX_CONCURRENT_REQUESTS` | — | `20` | Global cap on concurrent HTTP requests per client (source and destination independently). Prevents API overwhelm when multiple parallelization layers are active. Range: 1–200 | #### Streaming Migration (Logs, Experiments, Datasets) @@ -405,7 +404,7 @@ On resume: skips 1-30 (done), resumes experiment 31 from saved `_pagination_key` ## Parallelization -The migration tool currently uses **two active levels of concurrency** plus pipelined event streaming. The env vars in [Parallelization Tuning](#parallelization-tuning) still matter, but the within-project resource-type DAG concurrency described below has not been implemented yet. +The migration tool currently uses **two active levels of concurrency**. The env vars in [Parallelization Tuning](#parallelization-tuning) still matter, but the within-project resource-type DAG concurrency described below has not been implemented yet. ### How It Works @@ -424,8 +423,8 @@ The migration tool currently uses **two active levels of concurrency** plus pipe │ │ within each project. │ │ │ │ │ │ │ │ For streaming resources (logs/datasets/exps): │ │ -│ │ - Each stream can prefetch the next page while │ │ -│ │ inserting the current one │ │ +│ │ - Each stream fetches a page, inserts it, then │ │ +│ │ fetches the next (sequential) │ │ │ │ - Dataset/experiment streams are grouped for fetch │ │ │ │ efficiency, not scheduled as independent parallel │ │ │ │ DAG tasks within a project │ │ @@ -448,8 +447,7 @@ For streaming resources: - `datasets` and `experiments` group multiple ids into one BTQL fetch stream for efficiency. - `MIGRATION_MAX_CONCURRENT_RESOURCES` does not currently create multiple independent resource-type DAG tasks within a single project. -**Pipelined Event Streaming** (`MIGRATION_STREAMING_PIPELINE`, default true) -For each individual event stream (logs, dataset records, experiment events), the next BTQL page is prefetched from the source while the current page's batches are being inserted into the destination. This overlaps source reads with destination writes, reducing idle time for large migrations. +Event streams currently fetch and insert sequentially (fetch a page, insert its batches, then fetch the next page). ### Safety Mechanisms @@ -464,9 +462,9 @@ State mutations (ID mappings, checkpoint files) are protected by `asyncio.Lock` | **Small migration** (<5 projects, <100 resources) | Defaults work well. No tuning needed. | | **Many small projects** (50+ projects, small data) | Increase `MIGRATION_MAX_CONCURRENT=20` for more project-level parallelism. | | **Few projects with many resources** (e.g. 500 experiments in one project) | `MIGRATION_MAX_CONCURRENT_RESOURCES` helps only for migrators that support per-item fanout. Streaming resources within one project still run mostly as a single grouped stream. | -| **Large event streams** (TB-scale logs) | Defaults are good. Pipeline is on by default. Consider increasing `MIGRATION_MAX_CONCURRENT_REQUESTS=40` if the API can handle it. | +| **Large event streams** (TB-scale logs) | Defaults are good. Consider increasing `MIGRATION_MAX_CONCURRENT_REQUESTS=40` if the API can handle it. | | **Rate-limited API** (frequent 429s) | *Decrease* `MIGRATION_MAX_CONCURRENT_RESOURCES=2` and `MIGRATION_MAX_CONCURRENT_REQUESTS=10`. The tool handles 429s with backoff, but fewer concurrent requests reduces throttling. | -| **Debugging or sequential run** | Set `MIGRATION_MAX_CONCURRENT_RESOURCES=1` and `MIGRATION_STREAMING_PIPELINE=false` for deterministic, sequential execution. | +| **Debugging or sequential run** | Set `MIGRATION_MAX_CONCURRENT_RESOURCES=1` and `MIGRATION_MAX_CONCURRENT=1` for deterministic, sequential execution. | ### Example: Tuning for a Large Migration @@ -483,9 +481,6 @@ MIGRATION_MAX_CONCURRENT_RESOURCES=8 # Allow more HTTP connections (API can handle it) MIGRATION_MAX_CONCURRENT_REQUESTS=40 - -# Pipeline is on by default, but explicit for clarity -MIGRATION_STREAMING_PIPELINE=true ``` ```bash @@ -497,9 +492,6 @@ BT_DEST_API_KEY=... MIGRATION_MAX_CONCURRENT=5 MIGRATION_MAX_CONCURRENT_RESOURCES=2 MIGRATION_MAX_CONCURRENT_REQUESTS=10 - -# Disable pipeline for simpler debugging -MIGRATION_STREAMING_PIPELINE=false ``` ## Resource Types @@ -553,9 +545,6 @@ export MIGRATION_RETRY_DELAY=2.0 export MIGRATION_MAX_CONCURRENT_RESOURCES=2 export MIGRATION_MAX_CONCURRENT_REQUESTS=10 -# Disable pipelining for simpler debugging -export MIGRATION_STREAMING_PIPELINE=false - # Migrate incrementally braintrust-migrate migrate --resources ai_secrets,datasets braintrust-migrate migrate --resources prompts,functions diff --git a/braintrust_migrate/config.py b/braintrust_migrate/config.py index 6bbebaf..43910c0 100644 --- a/braintrust_migrate/config.py +++ b/braintrust_migrate/config.py @@ -103,10 +103,6 @@ class MigrationConfig(BaseModel): le=50, description="Maximum number of resources migrated concurrently within a batch", ) - streaming_pipeline: bool = Field( - default=True, - description="Enable pipelined page prefetch during streaming event migrations", - ) max_concurrent_requests: int = Field( default=20, ge=1, @@ -357,13 +353,6 @@ def from_env(cls) -> "Config": max_concurrent_resources = int( os.getenv("MIGRATION_MAX_CONCURRENT_RESOURCES", "5") ) - streaming_pipeline = os.getenv("MIGRATION_STREAMING_PIPELINE", "true").lower() in { - "1", - "true", - "yes", - "y", - "on", - } max_concurrent_requests = int( os.getenv("MIGRATION_MAX_CONCURRENT_REQUESTS", "20") ) @@ -509,7 +498,6 @@ def _get_bool(specific_key: str, unified_key: str, default: str) -> bool: retry_delay=retry_delay, max_concurrent=max_concurrent, max_concurrent_resources=max_concurrent_resources, - streaming_pipeline=streaming_pipeline, max_concurrent_requests=max_concurrent_requests, checkpoint_interval=checkpoint_interval, insert_max_request_bytes=insert_max_request_bytes, diff --git a/tests/integration/test_live_e2e_matrix.py b/tests/integration/test_live_e2e_matrix.py index 2a3048e..0ab1acb 100644 --- a/tests/integration/test_live_e2e_matrix.py +++ b/tests/integration/test_live_e2e_matrix.py @@ -54,7 +54,6 @@ class E2EScenario: max_concurrent: int max_concurrent_resources: int max_concurrent_requests: int - streaming_pipeline: bool @pytest.mark.asyncio @@ -66,18 +65,16 @@ class E2EScenario: "scenario", [ E2EScenario( - name="concurrent_pipeline_on", + name="concurrent", max_concurrent=3, max_concurrent_resources=6, max_concurrent_requests=30, - streaming_pipeline=True, ), E2EScenario( - name="sequential_pipeline_off", + name="sequential", max_concurrent=1, max_concurrent_resources=1, max_concurrent_requests=10, - streaming_pipeline=False, ), ], ids=lambda s: s.name, @@ -114,7 +111,6 @@ async def test_live_e2e_matrix( config.migration.max_concurrent = scenario.max_concurrent config.migration.max_concurrent_resources = scenario.max_concurrent_resources config.migration.max_concurrent_requests = scenario.max_concurrent_requests - config.migration.streaming_pipeline = scenario.streaming_pipeline config.migration.created_after = created_after config.migration.created_before = created_before diff --git a/tests/unit/test_parallelization_config.py b/tests/unit/test_parallelization_config.py index 640fbd2..b67843b 100644 --- a/tests/unit/test_parallelization_config.py +++ b/tests/unit/test_parallelization_config.py @@ -16,10 +16,6 @@ def test_max_concurrent_resources_default(self): config = MigrationConfig() assert config.max_concurrent_resources == 5 - def test_streaming_pipeline_default(self): - config = MigrationConfig() - assert config.streaming_pipeline is True - def test_max_concurrent_requests_default(self): config = MigrationConfig() assert config.max_concurrent_requests == 20 @@ -58,13 +54,6 @@ def test_max_concurrent_requests_above_max(self): with pytest.raises(ValueError): MigrationConfig(max_concurrent_requests=201) - def test_streaming_pipeline_bool(self): - config = MigrationConfig(streaming_pipeline=False) - assert config.streaming_pipeline is False - - config = MigrationConfig(streaming_pipeline=True) - assert config.streaming_pipeline is True - class TestParallelizationConfigFromEnv: """Verify env var parsing for parallelization fields.""" @@ -85,22 +74,6 @@ def test_max_concurrent_resources_from_env(self, monkeypatch): config = Config.from_env() assert config.migration.max_concurrent_resources == 10 - def test_streaming_pipeline_from_env_true(self, monkeypatch): - monkeypatch.setenv("BT_SOURCE_API_KEY", "src") - monkeypatch.setenv("BT_DEST_API_KEY", "dst") - monkeypatch.setenv("MIGRATION_STREAMING_PIPELINE", "true") - - config = Config.from_env() - assert config.migration.streaming_pipeline is True - - def test_streaming_pipeline_from_env_false(self, monkeypatch): - monkeypatch.setenv("BT_SOURCE_API_KEY", "src") - monkeypatch.setenv("BT_DEST_API_KEY", "dst") - monkeypatch.setenv("MIGRATION_STREAMING_PIPELINE", "false") - - config = Config.from_env() - assert config.migration.streaming_pipeline is False - def test_max_concurrent_requests_from_env(self, monkeypatch): monkeypatch.setenv("BT_SOURCE_API_KEY", "src") monkeypatch.setenv("BT_DEST_API_KEY", "dst") @@ -115,11 +88,9 @@ def test_defaults_when_env_not_set(self, monkeypatch): # Ensure the parallelization env vars are NOT set monkeypatch.delenv("MIGRATION_MAX_CONCURRENT", raising=False) monkeypatch.delenv("MIGRATION_MAX_CONCURRENT_RESOURCES", raising=False) - monkeypatch.delenv("MIGRATION_STREAMING_PIPELINE", raising=False) monkeypatch.delenv("MIGRATION_MAX_CONCURRENT_REQUESTS", raising=False) config = Config.from_env() assert config.migration.max_concurrent == 1 assert config.migration.max_concurrent_resources == 5 - assert config.migration.streaming_pipeline is True assert config.migration.max_concurrent_requests == 20 From 7df4e752f0f8b9d05e2b8b576c89be31c955f2f0 Mon Sep 17 00:00:00 2001 From: Doug Guthrie Date: Mon, 8 Jun 2026 11:23:13 -0600 Subject: [PATCH 7/8] Normalize streaming progress payloads behind a shared builder (R7) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The dataset and experiment migrators each defined four progress-hook lambdas (on_fetch/on_page/on_insert/on_done) that built byte-identical ~95-line dicts, differing only by the resource label and the source/dest id-field keys. Add build_stream_progress(phase, info, state, *, resource, id_fields) and make_stream_progress_hooks(...) in streaming_utils, and have both migrators pass a single descriptor instead of the four hand-written dicts. The emitted payloads are unchanged (per-phase sourcing preserved: fetch/page/done pull totals from the loop info; insert pulls committed totals from state; cursor truncation and the resource/id-field keys are identical), so the CLI progress display contract is preserved. Scope: datasets + experiments (which share the buffered loop). Logs keeps its own progress functions for now — its payload shape differs and it folds in with R1 (logs onto the shared loop). Adds build_stream_progress / make_stream_progress_hooks unit tests. Suite: 294. Co-Authored-By: Claude Opus 4.8 (1M context) --- braintrust_migrate/resources/datasets.py | 109 ++------------ braintrust_migrate/resources/experiments.py | 153 ++------------------ braintrust_migrate/streaming_utils.py | 114 +++++++++++++++ tests/unit/test_streaming_helpers.py | 109 ++++++++++++++ 4 files changed, 245 insertions(+), 240 deletions(-) diff --git a/braintrust_migrate/resources/datasets.py b/braintrust_migrate/resources/datasets.py index fe81afa..0d917e8 100644 --- a/braintrust_migrate/resources/datasets.py +++ b/braintrust_migrate/resources/datasets.py @@ -28,6 +28,7 @@ build_btql_sorted_page_query, dump_oversize_event_summary, is_http_413, + make_stream_progress_hooks, stream_btql_sorted_events_buffered, ) @@ -658,105 +659,15 @@ async def _fetch_group(n: int) -> dict[str, Any]: flush_max_bytes=self._sdk_flush_max_bytes, is_http_413=is_http_413, on_single_413=_on_single_413, - hooks=None - if progress is None - else { - "on_fetch": lambda info, _p=progress: _p( - { - "resource": "dataset_events", - "phase": "fetch", - "source_dataset_ids": source_dataset_ids, - "dest_dataset_ids": list(source_to_dest_dataset_ids.values()), - "page_num": info.get("page_num"), - "page_events": info.get("page_events"), - "fetched_total": info.get("fetched_total"), - "inserted_total": info.get("inserted_total"), - "inserted_bytes_total": info.get("inserted_bytes_total"), - "skipped_deleted_total": info.get("skipped_deleted_total"), - "skipped_seen_total": info.get("skipped_seen_total"), - "attachments_copied_total": info.get("attachments_copied_total"), - "pending_buffered_rows": info.get("pending_buffered_rows"), - "pending_buffered_bytes": info.get("pending_buffered_bytes"), - "cursor": ( - (state.btql_min_pagination_key[:16] + "…") - if isinstance(state.btql_min_pagination_key, str) - else None - ), - "next_cursor": None, - } - ), - "on_page": lambda info, _p=progress: _p( - { - "resource": "dataset_events", - "phase": "page", - "source_dataset_ids": source_dataset_ids, - "dest_dataset_ids": list(source_to_dest_dataset_ids.values()), - "page_num": info.get("page_num"), - "page_events": info.get("page_events"), - "fetched_total": info.get("fetched_total"), - "inserted_total": info.get("inserted_total"), - "inserted_bytes_total": info.get("inserted_bytes_total"), - "skipped_deleted_total": info.get("skipped_deleted_total"), - "skipped_seen_total": info.get("skipped_seen_total"), - "attachments_copied_total": info.get("attachments_copied_total"), - "pending_buffered_rows": info.get("pending_buffered_rows"), - "pending_buffered_bytes": info.get("pending_buffered_bytes"), - "cursor": ( - (state.btql_min_pagination_key[:16] + "…") - if isinstance(state.btql_min_pagination_key, str) - else None - ), - "next_cursor": None, - } - ), - "on_insert": lambda info, _p=progress: _p( - { - "resource": "dataset_events", - "phase": "insert", - "source_dataset_ids": source_dataset_ids, - "dest_dataset_ids": list(source_to_dest_dataset_ids.values()), - "page_num": None, - "page_events": None, - "inserted_last": info.get("inserted_last"), - "inserted_bytes_last": info.get("inserted_bytes_last"), - "insert_seconds": info.get("insert_seconds"), - "flush_rows": info.get("flush_rows"), - "flush_buffer_bytes": info.get("flush_buffer_bytes"), - "fetched_total": state.fetched_events, - "inserted_total": state.inserted_events, - "inserted_bytes_total": state.inserted_bytes, - "skipped_deleted_total": state.skipped_deleted, - "skipped_seen_total": state.skipped_seen, - "attachments_copied_total": state.attachments_copied, - "pending_buffered_rows": 0, - "pending_buffered_bytes": 0, - "cursor": ( - (state.btql_min_pagination_key[:16] + "…") - if isinstance(state.btql_min_pagination_key, str) - else None - ), - "next_cursor": None, - } - ), - "on_done": lambda info, _p=progress: _p( - { - "resource": "dataset_events", - "phase": "done", - "source_dataset_ids": source_dataset_ids, - "dest_dataset_ids": list(source_to_dest_dataset_ids.values()), - "fetched_total": info.get("fetched_total"), - "inserted_total": info.get("inserted_total"), - "inserted_bytes_total": info.get("inserted_bytes_total"), - "skipped_deleted_total": info.get("skipped_deleted_total"), - "skipped_seen_total": info.get("skipped_seen_total"), - "attachments_copied_total": info.get("attachments_copied_total"), - "pending_buffered_rows": info.get("pending_buffered_rows"), - "pending_buffered_bytes": info.get("pending_buffered_bytes"), - "cursor": None, - "next_cursor": None, - } - ), - }, + hooks=make_stream_progress_hooks( + progress, + state, + resource="dataset_events", + id_fields={ + "source_dataset_ids": source_dataset_ids, + "dest_dataset_ids": list(source_to_dest_dataset_ids.values()), + }, + ), ) self._logger.info( diff --git a/braintrust_migrate/resources/experiments.py b/braintrust_migrate/resources/experiments.py index bda46cb..9895c1b 100644 --- a/braintrust_migrate/resources/experiments.py +++ b/braintrust_migrate/resources/experiments.py @@ -26,6 +26,7 @@ build_btql_sorted_page_query, dump_oversize_event_summary, is_http_413, + make_stream_progress_hooks, stream_btql_sorted_events_buffered, ) @@ -754,147 +755,17 @@ async def _on_single_413(event: dict[str, Any], err: Exception) -> None: flush_max_bytes=self._sdk_flush_max_bytes, is_http_413=is_http_413, on_single_413=_on_single_413, - hooks=None - if progress is None - else { - "on_fetch": lambda info, _p=progress: _p( - { - "resource": "experiment_events", - "phase": "fetch", - "source_experiment_ids": source_experiment_ids, - "dest_experiment_ids": list( - source_to_dest_experiment_ids.values() - ), - "page_num": info.get("page_num"), - "page_events": info.get("page_events"), - "fetched_total": info.get("fetched_total"), - "inserted_total": info.get("inserted_total"), - "inserted_bytes_total": info.get( - "inserted_bytes_total" - ), - "skipped_deleted_total": info.get( - "skipped_deleted_total" - ), - "skipped_seen_total": info.get("skipped_seen_total"), - "attachments_copied_total": info.get( - "attachments_copied_total" - ), - "pending_buffered_rows": info.get( - "pending_buffered_rows" - ), - "pending_buffered_bytes": info.get( - "pending_buffered_bytes" - ), - "cursor": ( - (state.btql_min_pagination_key[:16] + "…") - if isinstance(state.btql_min_pagination_key, str) - else None - ), - "next_cursor": None, - } - ), - "on_page": lambda info, _p=progress: _p( - { - "resource": "experiment_events", - "phase": "page", - "source_experiment_ids": source_experiment_ids, - "dest_experiment_ids": list( - source_to_dest_experiment_ids.values() - ), - "page_num": info.get("page_num"), - "page_events": info.get("page_events"), - "fetched_total": info.get("fetched_total"), - "inserted_total": info.get("inserted_total"), - "inserted_bytes_total": info.get( - "inserted_bytes_total" - ), - "skipped_deleted_total": info.get( - "skipped_deleted_total" - ), - "skipped_seen_total": info.get("skipped_seen_total"), - "attachments_copied_total": info.get( - "attachments_copied_total" - ), - "pending_buffered_rows": info.get( - "pending_buffered_rows" - ), - "pending_buffered_bytes": info.get( - "pending_buffered_bytes" - ), - "cursor": ( - (state.btql_min_pagination_key[:16] + "…") - if isinstance(state.btql_min_pagination_key, str) - else None - ), - "next_cursor": None, - } - ), - "on_insert": lambda info, _p=progress: _p( - { - "resource": "experiment_events", - "phase": "insert", - "source_experiment_ids": source_experiment_ids, - "dest_experiment_ids": list( - source_to_dest_experiment_ids.values() - ), - "page_num": None, - "page_events": None, - "inserted_last": info.get("inserted_last"), - "inserted_bytes_last": info.get( - "inserted_bytes_last" - ), - "insert_seconds": info.get("insert_seconds"), - "flush_rows": info.get("flush_rows"), - "flush_buffer_bytes": info.get( - "flush_buffer_bytes" - ), - "fetched_total": state.fetched_events, - "inserted_total": state.inserted_events, - "inserted_bytes_total": state.inserted_bytes, - "skipped_deleted_total": state.skipped_deleted, - "skipped_seen_total": state.skipped_seen, - "attachments_copied_total": state.attachments_copied, - "pending_buffered_rows": 0, - "pending_buffered_bytes": 0, - "cursor": ( - (state.btql_min_pagination_key[:16] + "…") - if isinstance(state.btql_min_pagination_key, str) - else None - ), - "next_cursor": None, - } - ), - "on_done": lambda info, _p=progress: _p( - { - "resource": "experiment_events", - "phase": "done", - "source_experiment_ids": source_experiment_ids, - "dest_experiment_ids": list( - source_to_dest_experiment_ids.values() - ), - "fetched_total": info.get("fetched_total"), - "inserted_total": info.get("inserted_total"), - "inserted_bytes_total": info.get( - "inserted_bytes_total" - ), - "skipped_deleted_total": info.get( - "skipped_deleted_total" - ), - "skipped_seen_total": info.get("skipped_seen_total"), - "attachments_copied_total": info.get( - "attachments_copied_total" - ), - "pending_buffered_rows": info.get( - "pending_buffered_rows" - ), - "pending_buffered_bytes": info.get( - "pending_buffered_bytes" - ), - "cursor": None, - "next_cursor": None, - } - ), - }, + hooks=make_stream_progress_hooks( + progress, + state, + resource="experiment_events", + id_fields={ + "source_experiment_ids": source_experiment_ids, + "dest_experiment_ids": list( + source_to_dest_experiment_ids.values() + ), + }, + ), ) self._logger.info( diff --git a/braintrust_migrate/streaming_utils.py b/braintrust_migrate/streaming_utils.py index 775052f..bacf28a 100644 --- a/braintrust_migrate/streaming_utils.py +++ b/braintrust_migrate/streaming_utils.py @@ -275,6 +275,120 @@ def resolve(cls, source_client: Any, dest_client: Any) -> StreamingConfig: ) +def _truncate_cursor(state: Any) -> str | None: + """Short, display-friendly form of the current pagination key.""" + pk = state.btql_min_pagination_key + return (pk[:16] + "…") if isinstance(pk, str) else None + + +def build_stream_progress( + phase: str, + info: dict[str, Any], + state: Any, + *, + resource: str, + id_fields: dict[str, Any], +) -> dict[str, Any]: + """Build a normalized streaming progress payload for the dataset/experiment + event migrators. + + Replaces the four near-identical per-migrator hook lambdas. ``resource`` and + ``id_fields`` carry the per-resource bits (e.g. ``"dataset_events"`` and + ``{"source_dataset_ids": [...], "dest_dataset_ids": [...]}``); everything else + comes from the loop's ``info`` payload or ``state`` exactly as before. + """ + payload: dict[str, Any] = {"resource": resource, "phase": phase, **id_fields} + + if phase in ("fetch", "page"): + payload.update( + { + "page_num": info.get("page_num"), + "page_events": info.get("page_events"), + "fetched_total": info.get("fetched_total"), + "inserted_total": info.get("inserted_total"), + "inserted_bytes_total": info.get("inserted_bytes_total"), + "skipped_deleted_total": info.get("skipped_deleted_total"), + "skipped_seen_total": info.get("skipped_seen_total"), + "attachments_copied_total": info.get("attachments_copied_total"), + "pending_buffered_rows": info.get("pending_buffered_rows"), + "pending_buffered_bytes": info.get("pending_buffered_bytes"), + "cursor": _truncate_cursor(state), + "next_cursor": None, + } + ) + elif phase == "insert": + # Totals come from committed state (not the per-batch info); page fields + # are not meaningful for an insert event. + payload.update( + { + "page_num": None, + "page_events": None, + "inserted_last": info.get("inserted_last"), + "inserted_bytes_last": info.get("inserted_bytes_last"), + "insert_seconds": info.get("insert_seconds"), + "flush_rows": info.get("flush_rows"), + "flush_buffer_bytes": info.get("flush_buffer_bytes"), + "fetched_total": state.fetched_events, + "inserted_total": state.inserted_events, + "inserted_bytes_total": state.inserted_bytes, + "skipped_deleted_total": state.skipped_deleted, + "skipped_seen_total": state.skipped_seen, + "attachments_copied_total": state.attachments_copied, + "pending_buffered_rows": 0, + "pending_buffered_bytes": 0, + "cursor": _truncate_cursor(state), + "next_cursor": None, + } + ) + elif phase == "done": + payload.update( + { + "fetched_total": info.get("fetched_total"), + "inserted_total": info.get("inserted_total"), + "inserted_bytes_total": info.get("inserted_bytes_total"), + "skipped_deleted_total": info.get("skipped_deleted_total"), + "skipped_seen_total": info.get("skipped_seen_total"), + "attachments_copied_total": info.get("attachments_copied_total"), + "pending_buffered_rows": info.get("pending_buffered_rows"), + "pending_buffered_bytes": info.get("pending_buffered_bytes"), + "cursor": None, + "next_cursor": None, + } + ) + + return payload + + +def make_stream_progress_hooks( + progress: Callable[[dict[str, Any]], None] | None, + state: Any, + *, + resource: str, + id_fields: dict[str, Any], +) -> StreamHooks | None: + """Build the four streaming hooks that emit normalized progress payloads. + + Returns ``None`` when no progress callback is provided, matching the + ``hooks=None`` shape ``stream_btql_sorted_events_buffered`` expects. + """ + if progress is None: + return None + + def _hook(phase: str) -> Callable[[dict[str, Any]], None]: + return lambda info, _p=progress: _p( + build_stream_progress( + phase, info, state, resource=resource, id_fields=id_fields + ) + ) + + return { + "on_fetch": _hook("fetch"), + "on_page": _hook("page"), + "on_insert": _hook("insert"), + "on_done": _hook("done"), + } + + def build_btql_sorted_page_query( *, from_expr: str, diff --git a/tests/unit/test_streaming_helpers.py b/tests/unit/test_streaming_helpers.py index 2b73659..3c8b247 100644 --- a/tests/unit/test_streaming_helpers.py +++ b/tests/unit/test_streaming_helpers.py @@ -14,10 +14,13 @@ import httpx from braintrust_migrate.streaming_utils import ( + EventsStreamState, approx_event_size_bytes, + build_stream_progress, count_attachment_refs, dump_oversize_event_summary, is_http_413, + make_stream_progress_hooks, ) @@ -112,3 +115,109 @@ def test_dump_oversize_event_summary_unknown_id_and_prefix(tmp_path: Path): ) assert (tmp_path / "oversize_project_logs_event_unknown.json").exists() assert logger.calls[0][0].startswith("Oversize event isolated (413).") + + +def _events_state() -> EventsStreamState: + return EventsStreamState( + btql_min_pagination_key="0123456789ABCDEFGHIJ", # 20 chars + fetched_events=10, + inserted_events=8, + inserted_bytes=100, + skipped_deleted=1, + skipped_seen=2, + attachments_copied=3, + ) + + +def test_build_stream_progress_fetch_uses_info_and_truncates_cursor(): + state = _events_state() + info = { + "page_num": 2, + "page_events": 50, + "fetched_total": 7, + "inserted_total": 5, + "inserted_bytes_total": 60, + "skipped_deleted_total": 0, + "skipped_seen_total": 1, + "attachments_copied_total": 0, + "pending_buffered_rows": 3, + "pending_buffered_bytes": 30, + } + p = build_stream_progress( + "fetch", + info, + state, + resource="dataset_events", + id_fields={"source_dataset_ids": ["s"], "dest_dataset_ids": ["d"]}, + ) + assert p["resource"] == "dataset_events" + assert p["phase"] == "fetch" + assert p["source_dataset_ids"] == ["s"] + assert p["dest_dataset_ids"] == ["d"] + # fetch/page totals come from info, not state. + assert p["fetched_total"] == 7 + assert p["pending_buffered_rows"] == 3 + # cursor is the first 16 chars + ellipsis. + assert p["cursor"] == "0123456789ABCDEF…" + assert p["next_cursor"] is None + + +def test_build_stream_progress_insert_uses_state_totals(): + state = _events_state() + info = { + "inserted_last": 4, + "inserted_bytes_last": 40, + "insert_seconds": 0.5, + "flush_rows": 4, + "flush_buffer_bytes": 40, + } + p = build_stream_progress( + "insert", + info, + state, + resource="experiment_events", + id_fields={"source_experiment_ids": ["s"]}, + ) + assert p["phase"] == "insert" + # insert totals come from committed state, not the per-batch info. + assert p["fetched_total"] == state.fetched_events + assert p["inserted_total"] == state.inserted_events + assert p["inserted_bytes_total"] == state.inserted_bytes + assert p["page_num"] is None + assert p["pending_buffered_rows"] == 0 + assert p["inserted_last"] == 4 + assert p["insert_seconds"] == 0.5 + + +def test_build_stream_progress_done_nulls_cursor(): + state = _events_state() + info = {"fetched_total": 10, "inserted_total": 8, "inserted_bytes_total": 100} + p = build_stream_progress( + "done", info, state, resource="dataset_events", id_fields={} + ) + assert p["phase"] == "done" + assert p["fetched_total"] == 10 + assert p["cursor"] is None + + +def test_make_stream_progress_hooks_none_passthrough(): + assert ( + make_stream_progress_hooks(None, _events_state(), resource="x", id_fields={}) + is None + ) + + +def test_make_stream_progress_hooks_emits_via_callback(): + captured: list[dict] = [] + hooks = make_stream_progress_hooks( + captured.append, + _events_state(), + resource="dataset_events", + id_fields={"source_dataset_ids": ["s"]}, + ) + assert set(hooks) == {"on_fetch", "on_page", "on_insert", "on_done"} + hooks["on_fetch"]({"page_num": 1}) + assert captured[-1]["phase"] == "fetch" + assert captured[-1]["resource"] == "dataset_events" + hooks["on_done"]({}) + assert captured[-1]["phase"] == "done" From d140f086986c1ec7850e65b10619deb5890efca1 Mon Sep 17 00:00:00 2001 From: Doug Guthrie Date: Mon, 8 Jun 2026 15:03:10 -0600 Subject: [PATCH 8/8] Fix dataset single-event 413 NameError; tidy progress-hook lambda MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - datasets.py: `_on_single_413` referenced an unbound `source_dataset_id` (only a comprehension/loop variable elsewhere), so a single dataset event hitting HTTP 413 raised NameError — aborting the migration and skipping the oversize diagnostic instead of writing it and propagating the 413. Bind `source_dataset_id = event.get("dataset_id")` in the closure, matching the experiments migrator. (Latent bug present since before this branch; fixed here since the refactor touches this code.) - streaming_utils.py: drop the vestigial `_p=progress` default-arg in make_stream_progress_hooks' lambda — it's the loop-capture idiom, but there is no loop and `progress` is a stable captured param. Adds test_dataset_oversize_413.py, which drives a single-event 413 through the dataset streaming path (fails with NameError without the fix). Co-Authored-By: Claude Opus 4.8 (1M context) --- braintrust_migrate/resources/datasets.py | 1 + braintrust_migrate/streaming_utils.py | 2 +- tests/unit/test_dataset_oversize_413.py | 109 +++++++++++++++++++++++ 3 files changed, 111 insertions(+), 1 deletion(-) create mode 100644 tests/unit/test_dataset_oversize_413.py diff --git a/braintrust_migrate/resources/datasets.py b/braintrust_migrate/resources/datasets.py index 0d917e8..5b3accd 100644 --- a/braintrust_migrate/resources/datasets.py +++ b/braintrust_migrate/resources/datasets.py @@ -605,6 +605,7 @@ async def _fetch(n: int) -> dict[str, Any]: ) async def _on_single_413(event: dict[str, Any], err: Exception) -> None: + source_dataset_id = event.get("dataset_id") dump_oversize_event_summary( out_dir=events_dir, filename_prefix="oversize_dataset_event_", diff --git a/braintrust_migrate/streaming_utils.py b/braintrust_migrate/streaming_utils.py index bacf28a..e175c45 100644 --- a/braintrust_migrate/streaming_utils.py +++ b/braintrust_migrate/streaming_utils.py @@ -375,7 +375,7 @@ def make_stream_progress_hooks( return None def _hook(phase: str) -> Callable[[dict[str, Any]], None]: - return lambda info, _p=progress: _p( + return lambda info: progress( build_stream_progress( phase, info, state, resource=resource, id_fields=id_fields ) diff --git a/tests/unit/test_dataset_oversize_413.py b/tests/unit/test_dataset_oversize_413.py new file mode 100644 index 0000000..1a8e137 --- /dev/null +++ b/tests/unit/test_dataset_oversize_413.py @@ -0,0 +1,109 @@ +"""Regression test: a single dataset event hitting HTTP 413 must invoke the +oversize-event diagnostic without crashing. + +Previously `_on_single_413` referenced an unbound `source_dataset_id`, so a +single-event 413 raised `NameError` (aborting the migration and never writing +the diagnostic) instead of writing the summary and propagating the 413. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +import httpx +import pytest + +import braintrust_migrate.resources.datasets as datasets_module +from braintrust_migrate.config import MigrationConfig +from braintrust_migrate.resources.datasets import DatasetMigrator + + +class _SourceStub: + def __init__(self, page: list[dict[str, Any]]) -> None: + self._page = page + + async def with_retry(self, _op, coro_func, *, non_retryable_statuses=None): + _ = non_retryable_statuses + res = coro_func() + return await res if hasattr(res, "__await__") else res + + async def raw_request(self, method, path, *, params=None, json=None, timeout=None): + _ = params, timeout + assert method.lower() == "post" and path == "/btql" + q = json["query"] + # First page returns the event; subsequent (paginated) pages are empty. + if "_pagination_key >" in q: + return {"data": []} + return {"data": self._page} + + +class _DestStub: + migration_config = MigrationConfig(events_flush_max_rows=1) + + async def with_retry(self, _op, coro_func, *, non_retryable_statuses=None): + _ = non_retryable_statuses + res = coro_func() + return await res if hasattr(res, "__await__") else res + + +class _FakeWriter413: + """SDK writer stub whose insert always raises an HTTP 413.""" + + def __init__(self, dest_client: _DestStub, dataset_id: str) -> None: + self._dataset_id = dataset_id + + async def write_rows(self, rows: list[dict[str, Any]]) -> None: + req = httpx.Request("POST", "https://dest.example/logs3") + raise httpx.HTTPStatusError( + "payload too large", + request=req, + response=httpx.Response(413, request=req), + ) + + +@pytest.mark.asyncio +async def test_single_dataset_event_413_writes_summary_and_propagates( + tmp_path: Path, +) -> None: + page = [ + { + "id": "evt-a", + "_pagination_key": "p1", + "_xact_id": "10", + "created": "2023-01-01T00:00:00Z", + "input": {"big": "x" * 100}, + } + ] + source = _SourceStub(page) + dest = _DestStub() + + original_writer = datasets_module.SDKDatasetWriter + datasets_module.SDKDatasetWriter = _FakeWriter413 + try: + migrator = DatasetMigrator( + source, # type: ignore[arg-type] + dest, # type: ignore[arg-type] + tmp_path, + events_fetch_limit=1, + events_use_seen_db=False, + ) + + # The 413 is re-raised after the diagnostic is written. Crucially it must + # be the HTTPStatusError, NOT a NameError from an unbound variable. + with pytest.raises(httpx.HTTPStatusError): + await migrator._migrate_dataset_records( # type: ignore[attr-defined] + "source-dataset-id", "dest-dataset-id" + ) + + # The oversize diagnostic was written (proves _on_single_413 ran cleanly), + # with the correct per-resource dest id resolved from the event. + summary_path = tmp_path / "dataset_events" / "oversize_dataset_event_evt-a.json" + assert summary_path.exists() + import json + + summary = json.loads(summary_path.read_text()) + assert summary["dest_dataset_id"] == "dest-dataset-id" + assert summary["event_id"] == "evt-a" + finally: + datasets_module.SDKDatasetWriter = original_writer