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/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) 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/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/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, diff --git a/braintrust_migrate/resources/datasets.py b/braintrust_migrate/resources/datasets.py index dc7a71d..5b3accd 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 @@ -19,18 +18,22 @@ 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, + make_stream_progress_hooks, 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. @@ -43,12 +46,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, @@ -79,36 +82,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 @@ -536,86 +518,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 +605,22 @@ 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=( + source_dataset_id = event.get("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,107 +658,17 @@ 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 - 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 fb42de8..9895c1b 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, @@ -18,16 +16,20 @@ 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, + make_stream_progress_hooks, stream_btql_sorted_events_buffered, ) -# HTTP status codes -HTTP_STATUS_REQUEST_ENTITY_TOO_LARGE = 413 - class ExperimentMigrator(ResourceMigrator[dict]): """Migrator for Braintrust experiments. @@ -40,12 +42,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, @@ -76,36 +78,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 @@ -652,86 +633,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 +716,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,149 +753,19 @@ 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 - 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/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/resources/logs.py b/braintrust_migrate/resources/logs.py index eacddce..3fb810d 100644 --- a/braintrust_migrate/resources/logs.py +++ b/braintrust_migrate/resources/logs.py @@ -30,16 +30,18 @@ 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, ) 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).""" @@ -108,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", @@ -162,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) @@ -204,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" @@ -361,85 +346,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 +535,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 +759,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..e175c45 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: @@ -129,6 +228,167 @@ 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 _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: progress( + 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/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 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_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 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 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 diff --git a/tests/unit/test_streaming_helpers.py b/tests/unit/test_streaming_helpers.py new file mode 100644 index 0000000..3c8b247 --- /dev/null +++ b/tests/unit/test_streaming_helpers.py @@ -0,0 +1,223 @@ +"""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 ( + EventsStreamState, + approx_event_size_bytes, + build_stream_progress, + count_attachment_refs, + dump_oversize_event_summary, + is_http_413, + make_stream_progress_hooks, +) + + +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).") + + +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"