Add trace-level routing filters for project logs

.env.example

@@ -20,3 +20,8 @@ MIGRATION_CHECKPOINT_INTERVAL=50
 MIGRATION_STATE_DIR=./checkpoints
 # MIGRATION_PROJECT_MAP={"Source Project":"Destination Project"}
 # MIGRATION_PROJECT_MAP_FILE=./project-map.json
+
+# Trace-level routing for project logs (mutually exclusive).
+# Paired runs split one source project across two destination projects.
+# MIGRATION_LOGS_INCLUDE_ROOT_SPAN_NAME=my-root-span-name
+# MIGRATION_LOGS_EXCLUDE_ROOT_SPAN_NAME=my-root-span-name

.gitignore

@@ -178,5 +178,9 @@ checkpoints/
 *.checkpoint
 migration_state.json
 
+# Per-customer, one-off migration runbooks and configs (never committed).
+# One subdirectory per customer, e.g. customer-migrations/<customer>/
+customer-migrations/
+
 # Requirements lock file
 requirements.lock 

CHANGELOG.md

@@ -8,7 +8,7 @@ The format is based on Keep a Changelog and this project follows Semantic Versio
 
 ### Added
 
-- N/A
+- Trace-level routing filters for project logs: `--logs-include-root-span-name` / `MIGRATION_LOGS_INCLUDE_ROOT_SPAN_NAME` and `--logs-exclude-root-span-name` / `MIGRATION_LOGS_EXCLUDE_ROOT_SPAN_NAME`. The two are exact complements, so a paired run splits one source project across two destination projects with every span landing in exactly one of them. Because child spans do not carry their root's name, a one-time BTQL prepass collects the `root_span_id` of every span matching the name, and the streaming loop routes each span by its `root_span_id` — so whole traces (root plus all descendants) move together. The prepass is a full scan of the source project selecting only two id fields, is deliberately not constrained by `--created-after` / `--created-before` (a trace can straddle the boundary), and holds matched trace ids in memory (~150 bytes per trace). `--dry-run` runs the prepass too and reports matched span, root, and trace counts per project. The filter value is recorded in the logs checkpoint and a resume with a different value is rejected. Page size is tunable via `MIGRATION_LOGS_ROOT_SPAN_PREPASS_FETCH_LIMIT` (default `1000`).
 
 ### Changed
 

README.md

@@ -127,6 +127,37 @@ All options can be set via environment variables or CLI flags. CLI flags take pr
 | `MIGRATION_PROJECT_MAP_FILE` | `--project-map-file` | *(none)* | Path to a JSON file mapping source project names to destination project names. Mutually exclusive with `--project-map` / `MIGRATION_PROJECT_MAP` |
 | `MIGRATION_CREATED_AFTER` | `--created-after` | *(none)* | Only applies to resources that support created-time filtering. Currently this affects project logs event streaming and experiment listing. Migrates items with `created >=` this value (**inclusive**). Format: `YYYY-MM-DD` or ISO-8601 |
 | `MIGRATION_CREATED_BEFORE` | `--created-before` | *(none)* | Only applies to resources that support created-time filtering. Currently this affects project logs event streaming and experiment listing. Migrates items with `created <` this value (**exclusive**). Format: `YYYY-MM-DD` or ISO-8601 |
+| `MIGRATION_LOGS_INCLUDE_ROOT_SPAN_NAME` | `--logs-include-root-span-name` | *(none)* | Project logs only. Migrate **only** traces whose root span has this name (root span plus all descendants). Mutually exclusive with the exclude form |
+| `MIGRATION_LOGS_EXCLUDE_ROOT_SPAN_NAME` | `--logs-exclude-root-span-name` | *(none)* | Project logs only. Migrate every trace **except** those whose root span has this name. Exact complement of the include form |
+
+#### Splitting one source project across two destinations
+
+`--logs-include-root-span-name` and `--logs-exclude-root-span-name` are complements, so a paired run routes every span to exactly one destination:
+
+```bash
+# Run 1: everything except the "my-root-span-name" traces -> Dest A (all resources)
+braintrust-migrate migrate \
+  --projects "Src Project" \
+  --logs-exclude-root-span-name "my-root-span-name" \
+  --state-dir ./checkpoints/split-a
+
+# Run 2: only the "my-root-span-name" traces -> Dest B (logs only)
+braintrust-migrate migrate \
+  --projects "Src Project" \
+  --project-map '{"Src Project":"Dest B"}' \
+  --resources logs \
+  --logs-include-root-span-name "my-root-span-name" \
+  --state-dir ./checkpoints/split-b
+```
+
+How it works, and what it costs:
+
+- **Trace-level, not span-level.** Child spans do not carry their root's name, so a one-time BTQL prepass scans for spans matching the name and collects their `root_span_id`s. The streaming loop then routes each span by its `root_span_id`, keeping whole traces intact on both sides.
+- **The prepass is a full scan** of the source project's logs (selecting only two id fields) and is **not** constrained by `--created-after` / `--created-before` — a trace can straddle a time boundary, and a partial id set would misroute spans. It also runs during `--dry-run`, which reports matched span, root, and trace counts per project.
+- **The matched trace ids are held in memory** — roughly 150 bytes per trace. Fine into the low millions; beyond that the prepass should be moved to a disk-backed store.
+- **Both runs page through all spans**, since routing is applied client-side. Expect two full passes over the source project.
+- **If the name also appears mid-trace** (not just as a root), those full traces are routed too and the run logs a warning with the root-vs-non-root counts.
+- **Use a separate `--state-dir` per run.** The per-project checkpoint dir is keyed by *source* project name, so paired runs would otherwise collide — and the checkpoint refuses to resume if the filter value changed.
 
 #### Logging
 
@@ -171,6 +202,7 @@ These settings control BTQL-based streaming for high-volume resources.
 | `MIGRATION_EVENTS_USE_SEEN_DB` | — | `true` | Use SQLite store for deduplication |
 | `MIGRATION_LOGS_FETCH_LIMIT` | `--logs-fetch-limit` | *(inherits)* | Override fetch limit for logs only |
 | `MIGRATION_LOGS_INSERT_BATCH_SIZE` | `--logs-insert-batch-size` | `5000` | Deprecated alias for `MIGRATION_EVENTS_FLUSH_MAX_ROWS` |
+| `MIGRATION_LOGS_ROOT_SPAN_PREPASS_FETCH_LIMIT` | — | `1000` | BTQL page size for the root-span-name prepass. Rows carry only two id fields, so this can exceed the logs fetch limit |
 
 Resource-specific overrides follow the pattern `MIGRATION_{RESOURCE}_FETCH_LIMIT` and `MIGRATION_{RESOURCE}_USE_SEEN_DB` where `{RESOURCE}` is `LOGS`, `EXPERIMENT_EVENTS`, or `DATASET_EVENTS`. Logs additionally support `MIGRATION_LOGS_USE_VERSION_SNAPSHOT`. The older `MIGRATION_LOGS_INSERT_BATCH_SIZE` name is still accepted as a compatibility alias for the shared flush threshold.
 

braintrust_migrate/btql.py

@@ -69,6 +69,115 @@ def _query_text_for_limit(n: int) -> str:
     return lp if isinstance(lp, str) and lp else None
 
 
+async def collect_root_span_ids_for_span_name(
+    *,
+    client: BraintrustClient,
+    from_expr: str,
+    span_name: str,
+    page_limit: int = 1000,
+    operation: str = "btql_root_span_id_prepass",
+    log_fields: dict[str, Any],
+    on_page: Callable[[dict[str, Any]], None] | None = None,
+    timeout_seconds: float = 120.0,
+) -> tuple[set[str], dict[str, int]]:
+    """Collect the `root_span_id` of every span named `span_name`.
+
+    This is the prepass behind trace-level routing filters. Child spans do not
+    carry their root's name, so there is no single predicate that selects "every
+    span whose trace root is named X". Instead we scan for the named spans once
+    (selecting only two small fields) and build the set of matching trace ids;
+    the streaming loop then routes each span by its `root_span_id`.
+
+    The scan is deliberately *not* constrained by any created_after/created_before
+    window. A trace can straddle the window boundary, and a partial id set would
+    misroute spans whose root falls outside it. Extra ids are harmless — they
+    simply never match a streamed span.
+
+    Returns:
+        (root_span_ids, stats) where stats counts `matched_spans`, `root_spans`
+        (matches that are top-level, i.e. have no `span_parents`), and
+        `distinct_traces`. Note that one trace can contain several matching
+        top-level spans, so `root_spans` is often greater than `distinct_traces`.
+    """
+    # Imported here to avoid a circular import at module load time.
+    from braintrust_migrate.streaming_utils import build_btql_sorted_page_query
+
+    root_span_ids: set[str] = set()
+    matched_spans = 0
+    root_spans = 0
+    last_pk: str | None = None
+    page_num = 0
+
+    name_condition = f"span_attributes.name = '{btql_quote(span_name)}'"
+
+    while True:
+        page_num += 1
+
+        def _query_text_for_limit(n: int, *, _last_pk: str | None = last_pk) -> str:
+            return build_btql_sorted_page_query(
+                from_expr=from_expr,
+                limit=n,
+                last_pagination_key=_last_pk,
+                select="span_id, root_span_id, span_parents, _pagination_key",
+                extra_conditions=[name_condition],
+            )
+
+        page = await fetch_btql_sorted_page_with_retries(
+            client=client,
+            query_for_limit=_query_text_for_limit,
+            configured_limit=int(page_limit),
+            operation=operation,
+            log_fields={**log_fields, "span_name": span_name},
+            timeout_seconds=timeout_seconds,
+        )
+
+        rows = cast(list[dict[str, Any]], page.get("events") or [])
+        if not rows:
+            break
+
+        for row in rows:
+            matched_spans += 1
+            span_id = row.get("span_id")
+            root_span_id = row.get("root_span_id")
+            # A span is top-level when it has no parent. Do NOT infer this from
+            # `span_id == root_span_id`: under OTel-style ingestion `root_span_id`
+            # holds the 16-byte trace id while `span_id` is an 8-byte span id, so
+            # they never match and every span would look non-root.
+            span_parents = row.get("span_parents")
+            if not (isinstance(span_parents, list) and span_parents):
+                root_spans += 1
+            # Fall back to span_id so a root span with no explicit root_span_id
+            # still routes its own trace.
+            trace_id = (
+                root_span_id
+                if isinstance(root_span_id, str) and root_span_id
+                else span_id
+            )
+            if isinstance(trace_id, str) and trace_id:
+                root_span_ids.add(trace_id)
+
+        if on_page is not None:
+            on_page(
+                {
+                    "page_num": page_num,
+                    "page_rows": len(rows),
+                    "matched_spans": matched_spans,
+                    "distinct_traces": len(root_span_ids),
+                }
+            )
+
+        next_pk = cast(str | None, page.get("btql_last_pagination_key"))
+        if not next_pk or next_pk == last_pk:
+            break
+        last_pk = next_pk
+
+    return root_span_ids, {
+        "matched_spans": matched_spans,
+        "root_spans": root_spans,
+        "distinct_traces": len(root_span_ids),
+    }
+
+
 async def fetch_btql_sorted_page_with_retries(
     *,
     client: BraintrustClient,

braintrust_migrate/cli.py

@@ -226,6 +226,30 @@ def migrate(
             envvar="MIGRATION_CREATED_BEFORE",
         ),
     ] = None,
+    logs_include_root_span_name: Annotated[
+        str | None,
+        typer.Option(
+            "--logs-include-root-span-name",
+            help=(
+                "Only migrate logs traces whose root span has this name (the root span "
+                "and all of its descendants move together). Mutually exclusive with "
+                "--logs-exclude-root-span-name."
+            ),
+            envvar="MIGRATION_LOGS_INCLUDE_ROOT_SPAN_NAME",
+        ),
+    ] = None,
+    logs_exclude_root_span_name: Annotated[
+        str | None,
+        typer.Option(
+            "--logs-exclude-root-span-name",
+            help=(
+                "Migrate every logs trace except those whose root span has this name. "
+                "The exact complement of --logs-include-root-span-name, so paired runs "
+                "split one source project across two destination projects."
+            ),
+            envvar="MIGRATION_LOGS_EXCLUDE_ROOT_SPAN_NAME",
+        ),
+    ] = None,
     acl_map_users: Annotated[
         bool | None,
         typer.Option(
@@ -296,6 +320,8 @@ def migrate(
             logs_insert_batch_size,
             created_after,
             created_before,
+            logs_include_root_span_name,
+            logs_exclude_root_span_name,
             acl_map_users,
             acl_auto_invite_users,
             group_map_users,
@@ -319,6 +345,8 @@ async def _migrate_main(
     logs_insert_batch_size: int | None,
     created_after: str | None,
     created_before: str | None,
+    logs_include_root_span_name: str | None,
+    logs_exclude_root_span_name: str | None,
     acl_map_users: bool | None,
     acl_auto_invite_users: bool | None,
     group_map_users: bool | None,
@@ -394,6 +422,18 @@ async def _migrate_main(
             config.migration.created_after = canonicalize_created_after(created_after)
         if created_before is not None:
             config.migration.created_before = canonicalize_created_before(created_before)
+        if logs_include_root_span_name is not None:
+            config.migration.logs_include_root_span_name = logs_include_root_span_name
+        if logs_exclude_root_span_name is not None:
+            config.migration.logs_exclude_root_span_name = logs_exclude_root_span_name
+        if (
+            config.migration.logs_include_root_span_name
+            and config.migration.logs_exclude_root_span_name
+        ):
+            raise ValueError(
+                "Set only one of --logs-include-root-span-name or "
+                "--logs-exclude-root-span-name"
+            )
         if acl_map_users is not None:
             config.migration.acl_map_users = acl_map_users
         if acl_auto_invite_users is not None:
@@ -1261,6 +1301,7 @@ async def _test_logs_dry_run_probe(
     """Probe whether the selected logs time window has at least one matching span."""
     from braintrust_migrate.btql import (
         btql_quote,
+        collect_root_span_ids_for_span_name,
         fetch_btql_sorted_page_with_retries,
     )
     from braintrust_migrate.streaming_utils import build_btql_sorted_page_query
@@ -1305,6 +1346,42 @@ def _query_text_for_limit(n: int, *, _from_expr: str = from_expr) -> str:
                 "status": "error",
                 "error": str(e),
             }
+            continue
+
+        routing_name = (
+            config.migration.logs_include_root_span_name
+            or config.migration.logs_exclude_root_span_name
+        )
+        if not routing_name:
+            continue
+
+        progress.update(
+            validation_task,
+            description=f"🔍 Resolving root span filter for {project_name}...",
+        )
+        try:
+            _, stats = await collect_root_span_ids_for_span_name(
+                client=source_client,
+                from_expr=from_expr,
+                span_name=routing_name,
+                page_limit=config.migration.logs_root_span_prepass_fetch_limit,
+                log_fields={"source_project_id": project["source_id"]},
+            )
+            results[project_name].update(
+                {
+                    "root_span_filter_name": routing_name,
+                    "root_span_filter_mode": (
+                        "include"
+                        if config.migration.logs_include_root_span_name
+                        else "exclude"
+                    ),
+                    "matched_spans": stats["matched_spans"],
+                    "matched_root_spans": stats["root_spans"],
+                    "matched_traces": stats["distinct_traces"],
+                }
+            )
+        except Exception as e:
+            results[project_name]["root_span_filter_error"] = str(e)
 
     return results
 
@@ -1470,6 +1547,54 @@ def _display_dry_run_results(
         console.print("\n")
         console.print(logs_table)
 
+        has_filter = any(
+            "root_span_filter_name" in result or "root_span_filter_error" in result
+            for result in log_probe_results.values()
+        )
+        if has_filter:
+            filter_table = Table(title="🧬 Logs Root Span Filter")
+            filter_table.add_column("Project", style="cyan")
+            filter_table.add_column("Mode", style="magenta")
+            filter_table.add_column("Matched spans", justify="right", style="blue")
+            filter_table.add_column("Top-level", justify="right", style="blue")
+            filter_table.add_column("Traces routed", justify="right", style="green")
+
+            for project_name, result in log_probe_results.items():
+                if "root_span_filter_error" in result:
+                    filter_table.add_row(
+                        project_name,
+                        "❌ error",
+                        result["root_span_filter_error"],
+                        "-",
+                        "-",
+                    )
+                    continue
+                if "root_span_filter_name" not in result:
+                    continue
+                mode = result["root_span_filter_mode"]
+                matched_spans = int(result["matched_spans"])
+                matched_roots = int(result["matched_root_spans"])
+                traces = int(result["matched_traces"])
+                roots_cell = str(matched_roots)
+                if matched_roots != matched_spans:
+                    roots_cell = f"[yellow]{matched_roots}[/yellow]"
+                filter_table.add_row(
+                    project_name,
+                    f"{mode} '{result['root_span_filter_name']}'",
+                    str(matched_spans),
+                    roots_cell,
+                    str(traces) if mode == "include" else f"all traces except {traces}",
+                )
+
+            console.print("\n")
+            console.print(filter_table)
+            console.print(
+                "[dim]'Top-level' counts matches with no parent span. If it is lower "
+                "than 'Matched spans', the name also appears nested mid-trace, and "
+                "those full traces are routed too. One trace can hold several "
+                "top-level matches, so 'Top-level' may exceed 'Traces routed'.[/dim]"
+            )
+
     # Resource discovery results
     if test_results:
         resources_table = Table(title="🔍 Resource Discovery Test Results")