Code-smell quick wins: dedupe streaming helpers, dead-code cleanup, project retry

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

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)

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,

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,
                     )

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,