diff --git a/.env.example b/.env.example index dcb00d9..d2f084b 100644 --- a/.env.example +++ b/.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 diff --git a/.gitignore b/.gitignore index 2d04c97..c7e8463 100644 --- a/.gitignore +++ b/.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-migrations/ + # Requirements lock file requirements.lock \ No newline at end of file diff --git a/CHANGELOG.md b/CHANGELOG.md index 0218f15..6c50bad 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/README.md b/README.md index d2a2fef..e3c19ba 100644 --- a/README.md +++ b/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. diff --git a/braintrust_migrate/btql.py b/braintrust_migrate/btql.py index edbdff7..286b206 100644 --- a/braintrust_migrate/btql.py +++ b/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, diff --git a/braintrust_migrate/cli.py b/braintrust_migrate/cli.py index d15db58..0fa10a7 100644 --- a/braintrust_migrate/cli.py +++ b/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") diff --git a/braintrust_migrate/config.py b/braintrust_migrate/config.py index bc9d854..6402b7b 100644 --- a/braintrust_migrate/config.py +++ b/braintrust_migrate/config.py @@ -252,6 +252,34 @@ class MigrationConfig(BaseModel): ), ) + logs_include_root_span_name: str | None = Field( + default=None, + description=( + "Optional trace-level routing filter. When set, project logs migration " + "only migrates traces whose root span is named this value (the root span " + "and all of its descendants). Mutually exclusive with " + "logs_exclude_root_span_name." + ), + ) + logs_exclude_root_span_name: str | None = Field( + default=None, + description=( + "Optional trace-level routing filter. When set, project logs migration " + "migrates everything except traces whose root span is named this value. " + "The exact complement of logs_include_root_span_name, so the two can be " + "used in paired runs to split one source project across two destinations." + ), + ) + logs_root_span_prepass_fetch_limit: int = Field( + default=1000, + ge=1, + le=10_000, + description=( + "Fetch page size for the BTQL prepass that collects matching root span ids. " + "Rows are tiny (two id fields), so this can be larger than logs_fetch_limit." + ), + ) + # Experiment event migration tuning (experiments can be logs-scale) experiment_events_fetch_limit: int = Field( default=1000, @@ -325,6 +353,17 @@ def validate_created_before(cls, v: str | None) -> str | None: return None return canonicalize_created_before(v) + @model_validator(mode="after") + def validate_root_span_name_filters(self) -> "MigrationConfig": + """Ensure the trace routing filters are used one at a time.""" + if self.logs_include_root_span_name and self.logs_exclude_root_span_name: + raise ValueError( + "Set only one of logs_include_root_span_name or " + "logs_exclude_root_span_name (they are complements; use two runs " + "to split a source project across two destinations)" + ) + return self + @model_validator(mode="after") def validate_acl_user_mapping_flags(self) -> "MigrationConfig": """Ensure ACL auto-invite is only enabled with user mapping.""" @@ -480,6 +519,19 @@ def _get_bool(specific_key: str, unified_key: str, default: str) -> bool: "MIGRATION_LOGS_USE_SEEN_DB", "MIGRATION_EVENTS_USE_SEEN_DB", "true" ) + # Optional trace-level routing filter (logs only) + logs_include_root_span_name = os.getenv( + "MIGRATION_LOGS_INCLUDE_ROOT_SPAN_NAME" + ) + logs_exclude_root_span_name = os.getenv( + "MIGRATION_LOGS_EXCLUDE_ROOT_SPAN_NAME" + ) + logs_root_span_prepass_fetch_limit = _get_int( + "MIGRATION_LOGS_ROOT_SPAN_PREPASS_FETCH_LIMIT", + "MIGRATION_EVENTS_FETCH_LIMIT", + "1000", + ) + # Optional time filter (applies to logs and experiments) created_after = os.getenv("MIGRATION_CREATED_AFTER") created_before = os.getenv("MIGRATION_CREATED_BEFORE") @@ -601,6 +653,9 @@ def _get_bool(specific_key: str, unified_key: str, default: str) -> bool: logs_insert_batch_size=logs_insert_batch_size, logs_use_version_snapshot=logs_use_version_snapshot, logs_use_seen_db=logs_use_seen_db, + logs_include_root_span_name=logs_include_root_span_name, + logs_exclude_root_span_name=logs_exclude_root_span_name, + logs_root_span_prepass_fetch_limit=logs_root_span_prepass_fetch_limit, created_after=created_after, created_before=created_before, experiment_events_fetch_limit=experiment_events_fetch_limit, diff --git a/braintrust_migrate/resources/logs.py b/braintrust_migrate/resources/logs.py index 3fb810d..add953d 100644 --- a/braintrust_migrate/resources/logs.py +++ b/braintrust_migrate/resources/logs.py @@ -23,6 +23,7 @@ ) from braintrust_migrate.btql import ( btql_quote, + collect_root_span_ids_for_span_name, fetch_btql_sorted_page_with_retries, find_first_pagination_key_for_created_after, ) @@ -62,6 +63,9 @@ class _LogsStreamingState: btql_last_created: str | None = None created_after: str | None = None created_before: str | None = None + include_root_span_name: str | None = None + exclude_root_span_name: str | None = None + skipped_filtered: int = 0 @classmethod def from_path(cls, path: Path) -> _LogsStreamingState: @@ -86,6 +90,9 @@ def from_path(cls, path: Path) -> _LogsStreamingState: btql_last_created=data.get("btql_last_created"), created_after=data.get("created_after"), created_before=data.get("created_before"), + include_root_span_name=data.get("include_root_span_name"), + exclude_root_span_name=data.get("exclude_root_span_name"), + skipped_filtered=int(data.get("skipped_filtered", 0)), ) def to_dict(self) -> dict[str, Any]: @@ -104,6 +111,9 @@ def to_dict(self) -> dict[str, Any]: "btql_last_created": self.btql_last_created, "created_after": self.created_after, "created_before": self.created_before, + "include_root_span_name": self.include_root_span_name, + "exclude_root_span_name": self.exclude_root_span_name, + "skipped_filtered": self.skipped_filtered, } @@ -346,6 +356,97 @@ def _extract_ids(events: list[dict[str, Any]]) -> list[str]: ids.append(event_id) return ids + @staticmethod + def _trace_id_for_event(event: dict[str, Any]) -> str | None: + """Identify which trace a span belongs to, tolerating missing root ids.""" + root_span_id = event.get("root_span_id") + if isinstance(root_span_id, str) and root_span_id: + return root_span_id + span_id = event.get("span_id") + if isinstance(span_id, str) and span_id: + return span_id + event_id = event.get("id") + return event_id if isinstance(event_id, str) and event_id else None + + async def _collect_matched_trace_ids( + self, + *, + source_project_id: str, + dest_project_id: str, + span_name: str, + include_mode: bool, + ) -> set[str]: + """Run the prepass that resolves a root span name to a set of trace ids. + + Held in memory: one id per matching trace. Callers are expected to keep + this well under a few million traces per project. + """ + prepass_limit = int( + getattr( + getattr(self.source_client, "migration_config", None), + "logs_root_span_prepass_fetch_limit", + 1000, + ) + ) + self._logger.info( + "Starting root span name prepass", + source_project_id=source_project_id, + dest_project_id=dest_project_id, + span_name=span_name, + mode="include" if include_mode else "exclude", + prepass_fetch_limit=prepass_limit, + ) + + progress_hook = self._progress_hook + + def _on_prepass_page(info: dict[str, Any]) -> None: + if progress_hook is None: + return + progress_hook( + { + "resource": "logs", + "phase": "prepass", + "source_project_id": source_project_id, + "dest_project_id": dest_project_id, + "span_name": span_name, + **info, + } + ) + + matched_trace_ids, stats = await collect_root_span_ids_for_span_name( + client=self.source_client, + from_expr=f"project_logs('{btql_quote(source_project_id)}') spans", + span_name=span_name, + page_limit=prepass_limit, + log_fields={"source_project_id": source_project_id}, + on_page=_on_prepass_page, + ) + + nested_matches = stats["matched_spans"] - stats["root_spans"] + if nested_matches: + # A nested match still pulls its whole trace across, which keeps traces + # intact but widens the selection beyond top-level-named traces. Worth + # surfacing so an unexpectedly broad split is caught before it lands. + self._logger.warning( + "Some spans matching the name are nested rather than top-level; " + "their full traces will be routed by this filter", + source_project_id=source_project_id, + span_name=span_name, + matched_spans=stats["matched_spans"], + top_level_spans=stats["root_spans"], + nested_matches=nested_matches, + ) + + self._logger.info( + "Completed root span name prepass", + source_project_id=source_project_id, + dest_project_id=dest_project_id, + span_name=span_name, + mode="include" if include_mode else "exclude", + **stats, + ) + return matched_trace_ids + async def migrate_all( self, project_id: str | None = None, max_concurrent: int | None = None ) -> dict[str, Any]: @@ -456,6 +557,57 @@ async def migrate_all( f"run={created_before_cfg!r}. Re-run with the checkpoint value or start a fresh checkpoint." ) + # Validate and persist the trace-level routing filter, then run its prepass. + include_name_cfg = getattr(mig_cfg, "logs_include_root_span_name", None) + exclude_name_cfg = getattr(mig_cfg, "logs_exclude_root_span_name", None) + for label, cfg_value, state_attr in ( + ("include_root_span_name", include_name_cfg, "include_root_span_name"), + ("exclude_root_span_name", exclude_name_cfg, "exclude_root_span_name"), + ): + state_value = getattr(self._stream_state, state_attr) + if state_value is None and cfg_value is None: + continue + if state_value is None: + setattr(self._stream_state, state_attr, cfg_value) + self._save_stream_state() + elif cfg_value is None: + raise ValueError( + f"Checkpoint includes a {label} filter but this run does not. " + f"Re-run with the same --logs-{label.replace('_', '-')} value " + "or start a fresh checkpoint." + ) + elif state_value != cfg_value: + raise ValueError( + f"{label} mismatch vs checkpoint: checkpoint={state_value!r} " + f"run={cfg_value!r}. Re-run with the checkpoint value or start " + "a fresh checkpoint." + ) + + routing_name = ( + self._stream_state.include_root_span_name + or self._stream_state.exclude_root_span_name + ) + include_mode = self._stream_state.include_root_span_name is not None + matched_trace_ids: set[str] | None = None + if routing_name: + matched_trace_ids = await self._collect_matched_trace_ids( + source_project_id=source_project_id, + dest_project_id=dest_project_id, + span_name=routing_name, + include_mode=include_mode, + ) + # Include mode with no matching traces has nothing to migrate at all. + # Exclude mode still migrates everything, so it must not short-circuit. + if include_mode and not matched_trace_ids: + self._logger.warning( + "No traces matched the root span name filter; nothing to migrate", + source_project_id=source_project_id, + dest_project_id=dest_project_id, + include_root_span_name=routing_name, + ) + self._save_stream_state() + return self.get_partial_results() + # If this is the first BTQL page and created_after is set, preflight the first # matching pagination key so we can start near the boundary. if ( @@ -520,6 +672,7 @@ async def migrate_all( pending_inserted_events = 0 pending_inserted_bytes = 0 pending_skipped_seen = 0 + pending_skipped_filtered = 0 pending_attachments_copied = 0 pending_spilled_fields = 0 pending_last_pk: str | None = None @@ -729,6 +882,7 @@ async def _flush_pending_events() -> None: nonlocal pending_inserted_events nonlocal pending_inserted_bytes nonlocal pending_skipped_seen + nonlocal pending_skipped_filtered nonlocal pending_attachments_copied nonlocal pending_spilled_fields nonlocal pending_last_pk @@ -738,6 +892,7 @@ async def _flush_pending_events() -> None: pending_fetched_events == 0 and pending_inserted_events == 0 and pending_skipped_seen == 0 + and pending_skipped_filtered == 0 and pending_attachments_copied == 0 and pending_spilled_fields == 0 and pending_last_pk is None @@ -778,6 +933,7 @@ async def _flush_pending_events() -> None: self._stream_state.fetched_events += pending_fetched_events self._stream_state.skipped_seen += pending_skipped_seen + self._stream_state.skipped_filtered += pending_skipped_filtered self._stream_state.attachments_copied += pending_attachments_copied self._stream_state.spilled_fields += pending_spilled_fields self._stream_state.btql_last_created = pending_last_created @@ -792,6 +948,7 @@ async def _flush_pending_events() -> None: pending_inserted_events = 0 pending_inserted_bytes = 0 pending_skipped_seen = 0 + pending_skipped_filtered = 0 pending_attachments_copied = 0 pending_spilled_fields = 0 pending_last_pk = None @@ -844,9 +1001,22 @@ def _query_text_for_limit(n: int) -> str: pending_fetched_events += len(page_events) + # Trace-level routing. `page_events` stays intact below so the + # pagination key and last-created bookkeeping still advance across + # pages that are entirely filtered out. + routed_events = page_events + if matched_trace_ids is not None: + routed_events = [ + event + for event in page_events + if (self._trace_id_for_event(event) in matched_trace_ids) + == include_mode + ] + pending_skipped_filtered += len(page_events) - len(routed_events) + insert_events_list = [ self._event_to_insert(event, source_project_id) - for event in page_events + for event in routed_events ] if seen_db is not None: @@ -942,6 +1112,7 @@ def _query_text_for_limit(n: int) -> str: "version": self._stream_state.version, "resume_cursor": self._stream_state.cursor, "spilled_fields": self._stream_state.spilled_fields, + "filtered_out": self._stream_state.skipped_filtered, } finally: if seen_db is not None: @@ -962,4 +1133,5 @@ def get_partial_results(self) -> dict[str, Any]: "version": self._stream_state.version, "resume_cursor": self._stream_state.cursor, "spilled_fields": self._stream_state.spilled_fields, + "filtered_out": self._stream_state.skipped_filtered, } diff --git a/braintrust_migrate/streaming_utils.py b/braintrust_migrate/streaming_utils.py index e175c45..65887be 100644 --- a/braintrust_migrate/streaming_utils.py +++ b/braintrust_migrate/streaming_utils.py @@ -398,6 +398,7 @@ def build_btql_sorted_page_query( created_after: str | None = None, created_before: str | None = None, select: str = "*", + extra_conditions: list[str] | None = None, ) -> str: """Build a native BTQL query for stable sorted paging on `_pagination_key`. @@ -412,6 +413,7 @@ def build_btql_sorted_page_query( created_after: Only include rows with created >= this value (inclusive) created_before: Only include rows with created < this value (exclusive) select: Fields to select (default "*") + extra_conditions: Additional pre-quoted BTQL filter conditions to AND in Returns: Native BTQL query string @@ -421,6 +423,9 @@ def build_btql_sorted_page_query( conditions.append(f"created >= '{btql_quote(created_after)}'") if isinstance(created_before, str) and created_before: conditions.append(f"created < '{btql_quote(created_before)}'") + for condition in extra_conditions or []: + if isinstance(condition, str) and condition: + conditions.append(condition) if isinstance(last_pagination_key, str) and last_pagination_key: op = ">=" if last_pagination_key_inclusive else ">" conditions.append(f"_pagination_key {op} '{btql_quote(last_pagination_key)}'") diff --git a/tests/unit/test_logs_root_span_filter.py b/tests/unit/test_logs_root_span_filter.py new file mode 100644 index 0000000..f9c70ab --- /dev/null +++ b/tests/unit/test_logs_root_span_filter.py @@ -0,0 +1,430 @@ +"""Tests for trace-level routing of project logs by root span name. + +The filter exists so one source project can be split across two destination +projects. The property that matters most is that include and exclude are exact +complements: every span lands in exactly one of the two runs. +""" + +from __future__ import annotations + +import re +from pathlib import Path +from typing import Any + +import pytest + +import braintrust_migrate.resources.logs as logs_module +from braintrust_migrate.config import MigrationConfig +from braintrust_migrate.resources.logs import LogsMigrator +from braintrust_migrate.streaming_utils import build_btql_sorted_page_query + +ROOT_NAME = "chat.message" + + +def _span( + event_id: str, + span_id: str, + root_span_id: str, + name: str, + pagination_key: str, + span_parents: list[str] | None = None, +) -> dict[str, Any]: + return { + "id": event_id, + "span_id": span_id, + "root_span_id": root_span_id, + "span_parents": span_parents, + "span_attributes": {"name": name}, + "_pagination_key": pagination_key, + "_xact_id": "1", + "created": "2026-01-01T00:00:00Z", + } + + +# Two matching traces (t1 with children, t3 single-span) and one non-matching +# trace (t2 with a child). Mirrors the real shape: children do not carry the +# root's name, so only the prepass can associate them. +ALL_SPANS: list[dict[str, Any]] = [ + _span("e1", "s1", "s1", ROOT_NAME, "pk1"), + _span("e2", "s2", "s1", "llm.call", "pk2", span_parents=["s1"]), + _span("e3", "s3", "s1", "tool.lookup", "pk3", span_parents=["s1"]), + _span("e4", "s4", "s4", "other-root", "pk4"), + _span("e5", "s5", "s4", "llm.call", "pk5", span_parents=["s4"]), + _span("e6", "s6", "s6", ROOT_NAME, "pk6"), +] + +# OpenTelemetry-style ingestion: `root_span_id` is a 16-byte trace id and +# `span_id` is an 8-byte span id, so they never match even for top-level spans. +# Top-level-ness is carried solely by an empty `span_parents`. A single trace +# can hold several top-level spans of the same name (one per chat turn). +OTEL_TRACE = "b3b9b3f191d3d56ccc92b4023a321c02" +OTEL_OTHER_TRACE = "a932535782befd370e423c2d7ef93af5" +OTEL_SPANS: list[dict[str, Any]] = [ + _span("o1", "0d2ddc17e6aadaff", OTEL_TRACE, ROOT_NAME, "qk1"), + _span( + "o2", + "f2d2575b5b3db975", + OTEL_TRACE, + "llm.stream", + "qk2", + span_parents=["0d2ddc17e6aadaff"], + ), + _span("o3", "9814bb66fbb565ff", OTEL_TRACE, ROOT_NAME, "qk3"), + _span( + "o4", + "b84e2143935e5bf3", + OTEL_TRACE, + "llm.stream", + "qk4", + span_parents=["9814bb66fbb565ff"], + ), + _span("o5", "94b27b8d9f5d224f", OTEL_OTHER_TRACE, "recommendation", "qk5"), + _span( + "o6", + "c87905908e5749fb", + OTEL_OTHER_TRACE, + "llm.stream", + "qk6", + span_parents=["94b27b8d9f5d224f"], + ), +] + +MATCHING_IDS = {"e1", "e2", "e3", "e6"} +NON_MATCHING_IDS = {"e4", "e5"} + + +class _StubClient: + """Serves BTQL prepass and streaming pages, and records inserts.""" + + def __init__( + self, + *, + pages: list[list[dict[str, Any]]], + migration_config: MigrationConfig, + ) -> None: + self.pages = pages + self.migration_config = migration_config + self.inserts: list[list[dict[str, Any]]] = [] + self.prepass_queries = 0 + + async def with_retry( + self, + _operation_name: str, + coro_func, + *, + non_retryable_statuses: set[int] | None = None, + ): + _ = non_retryable_statuses + res = coro_func() + if hasattr(res, "__await__"): + return await res + return res + + def _page_index_for_query(self, query: str) -> int: + match = re.search(r"_pagination_key > '([^']+)'", query) + if match is None: + return 0 + key = match.group(1) + for index, page in enumerate(self.pages): + if page and page[-1]["_pagination_key"] == key: + return index + 1 + return len(self.pages) + + async def raw_request( + self, + method: str, + path: str, + *, + params: dict[str, Any] | None = None, + json: Any | None = None, + timeout: float | None = None, + ) -> Any: + _ = params, timeout + assert method.lower() == "post" + assert path == "/btql" + assert json is not None + query = json["query"] + assert isinstance(query, str) + + if "span_attributes.name" in query: + self.prepass_queries += 1 + if "_pagination_key >" in query: + return {"data": []} + rows = [ + { + "span_id": span["span_id"], + "root_span_id": span["root_span_id"], + "span_parents": span.get("span_parents"), + "_pagination_key": span["_pagination_key"], + } + for page in self.pages + for span in page + if span["span_attributes"]["name"] == ROOT_NAME + ] + return {"data": rows} + + index = self._page_index_for_query(query) + if index >= len(self.pages): + return {"data": []} + return {"data": self.pages[index]} + + +class _FakeSDKProjectLogsWriter: + def __init__(self, dest_client: _StubClient, project_id: str) -> None: + self._dest_client = dest_client + self._project_id = project_id + + async def write_rows(self, rows: list[dict[str, Any]]) -> None: + self._dest_client.inserts.append([dict(row) for row in rows]) + + +async def _run_migration( + tmp_path: Path, + *, + pages: list[list[dict[str, Any]]], + config: MigrationConfig, +) -> tuple[list[str], LogsMigrator, _StubClient]: + source = _StubClient(pages=pages, migration_config=config) + dest = _StubClient(pages=pages, migration_config=config) + + original_writer = logs_module.SDKProjectLogsWriter + logs_module.SDKProjectLogsWriter = _FakeSDKProjectLogsWriter + try: + migrator = LogsMigrator( + source, # type: ignore[arg-type] + dest, # type: ignore[arg-type] + tmp_path, + page_limit=10, + use_seen_db=False, + ) + migrator.set_destination_project_id("dest-project-id") + await migrator.migrate_all("source-project-id") + finally: + logs_module.SDKProjectLogsWriter = original_writer + + inserted = [event["id"] for batch in dest.inserts for event in batch] + return inserted, migrator, source + + +@pytest.mark.asyncio +async def test_include_mode_migrates_whole_matching_traces(tmp_path: Path) -> None: + """Children of a matching root come along even though their names differ.""" + config = MigrationConfig(logs_include_root_span_name=ROOT_NAME) + inserted, _, _ = await _run_migration( + tmp_path, pages=[list(ALL_SPANS)], config=config + ) + + assert set(inserted) == MATCHING_IDS + # e2/e3 are named llm.call and tool.lookup; they qualify only via root_span_id. + assert "e2" in inserted + assert "e3" in inserted + + +@pytest.mark.asyncio +async def test_exclude_mode_migrates_everything_else(tmp_path: Path) -> None: + config = MigrationConfig(logs_exclude_root_span_name=ROOT_NAME) + inserted, _, _ = await _run_migration( + tmp_path, pages=[list(ALL_SPANS)], config=config + ) + + assert set(inserted) == NON_MATCHING_IDS + + +@pytest.mark.asyncio +async def test_include_and_exclude_are_exact_complements(tmp_path: Path) -> None: + """The core guarantee for a paired split: no span duplicated, none dropped.""" + include_inserted, _, _ = await _run_migration( + tmp_path / "include", + pages=[list(ALL_SPANS)], + config=MigrationConfig(logs_include_root_span_name=ROOT_NAME), + ) + exclude_inserted, _, _ = await _run_migration( + tmp_path / "exclude", + pages=[list(ALL_SPANS)], + config=MigrationConfig(logs_exclude_root_span_name=ROOT_NAME), + ) + + assert set(include_inserted) & set(exclude_inserted) == set() + assert set(include_inserted) | set(exclude_inserted) == { + span["id"] for span in ALL_SPANS + } + + +@pytest.mark.asyncio +async def test_no_filter_migrates_everything(tmp_path: Path) -> None: + inserted, _, source = await _run_migration( + tmp_path, pages=[list(ALL_SPANS)], config=MigrationConfig() + ) + + assert set(inserted) == {span["id"] for span in ALL_SPANS} + assert source.prepass_queries == 0, "prepass must not run without a filter" + + +@pytest.mark.asyncio +async def test_pagination_advances_through_fully_filtered_page( + tmp_path: Path, +) -> None: + """A page where every row is filtered out must still advance the cursor.""" + page1 = [ALL_SPANS[3], ALL_SPANS[4]] # e4, e5 -- entirely non-matching + page2 = [ALL_SPANS[0], ALL_SPANS[1]] # e1, e2 -- matching trace + + config = MigrationConfig(logs_include_root_span_name=ROOT_NAME) + inserted, migrator, _ = await _run_migration( + tmp_path, pages=[page1, page2], config=config + ) + + assert set(inserted) == {"e1", "e2"} + state = migrator._stream_state + assert state.btql_min_pagination_key == "pk2", ( + "cursor must advance past the fully filtered page" + ) + assert state.skipped_filtered == len(page1) + assert state.fetched_events == len(page1) + len(page2), ( + "fetched counts pre-filter rows" + ) + + +@pytest.mark.asyncio +async def test_include_mode_with_zero_matches_migrates_nothing( + tmp_path: Path, +) -> None: + pages = [[ALL_SPANS[3], ALL_SPANS[4]]] # no span carries the root name + config = MigrationConfig(logs_include_root_span_name="does-not-exist") + inserted, _, _ = await _run_migration(tmp_path, pages=pages, config=config) + + assert inserted == [] + + +@pytest.mark.asyncio +async def test_checkpoint_rejects_changed_filter(tmp_path: Path) -> None: + """Resuming with a different filter would corrupt the split; it must fail.""" + await _run_migration( + tmp_path, + pages=[list(ALL_SPANS)], + config=MigrationConfig(logs_include_root_span_name=ROOT_NAME), + ) + + with pytest.raises(ValueError, match="mismatch vs checkpoint"): + await _run_migration( + tmp_path, + pages=[list(ALL_SPANS)], + config=MigrationConfig(logs_include_root_span_name="some-other-name"), + ) + + +@pytest.mark.asyncio +async def test_checkpoint_rejects_dropped_filter(tmp_path: Path) -> None: + await _run_migration( + tmp_path, + pages=[list(ALL_SPANS)], + config=MigrationConfig(logs_include_root_span_name=ROOT_NAME), + ) + + with pytest.raises(ValueError, match="but this run does not"): + await _run_migration( + tmp_path, pages=[list(ALL_SPANS)], config=MigrationConfig() + ) + + +@pytest.mark.asyncio +async def test_otel_shaped_ids_route_whole_traces(tmp_path: Path) -> None: + """Under OTel ingestion, span_id never equals root_span_id. + + Routing must still group by root_span_id (the trace id), pulling both + top-level matches in the trace and their nested children. + """ + inserted, _, _ = await _run_migration( + tmp_path / "inc", + pages=[list(OTEL_SPANS)], + config=MigrationConfig(logs_include_root_span_name=ROOT_NAME), + ) + assert set(inserted) == {"o1", "o2", "o3", "o4"} + + excluded, _, _ = await _run_migration( + tmp_path / "exc", + pages=[list(OTEL_SPANS)], + config=MigrationConfig(logs_exclude_root_span_name=ROOT_NAME), + ) + assert set(excluded) == {"o5", "o6"} + + +@pytest.mark.asyncio +async def test_top_level_detection_uses_span_parents_not_id_equality( + tmp_path: Path, +) -> None: + """Regression: counting roots via span_id == root_span_id reported 0 on + OTel data, which read as "the name is never top-level" when it always was.""" + from braintrust_migrate.btql import collect_root_span_ids_for_span_name + + client = _StubClient( + pages=[list(OTEL_SPANS)], + migration_config=MigrationConfig(), + ) + _, stats = await collect_root_span_ids_for_span_name( + client=client, # type: ignore[arg-type] + from_expr="project_logs('p') spans", + span_name=ROOT_NAME, + log_fields={}, + ) + + # Both matches are top-level despite span_id != root_span_id, and both live + # in the same trace, so one trace is routed by two matches. + assert stats == {"matched_spans": 2, "root_spans": 2, "distinct_traces": 1} + + # A nested match must NOT be counted as top-level, which is what makes the + # warning about over-broad selection meaningful. + nested = [ + *OTEL_SPANS, + _span( + "o7", + "aaaa1111bbbb2222", + OTEL_OTHER_TRACE, + ROOT_NAME, + "qk7", + span_parents=["94b27b8d9f5d224f"], + ), + ] + _, nested_stats = await collect_root_span_ids_for_span_name( + client=_StubClient( # type: ignore[arg-type] + pages=[nested], migration_config=MigrationConfig() + ), + from_expr="project_logs('p') spans", + span_name=ROOT_NAME, + log_fields={}, + ) + assert nested_stats == { + "matched_spans": 3, + "root_spans": 2, + "distinct_traces": 2, + } + + +def test_trace_id_falls_back_when_root_span_id_missing() -> None: + assert ( + LogsMigrator._trace_id_for_event({"root_span_id": "r", "span_id": "s"}) == "r" + ) + assert LogsMigrator._trace_id_for_event({"span_id": "s", "id": "e"}) == "s" + assert LogsMigrator._trace_id_for_event({"id": "e"}) == "e" + assert LogsMigrator._trace_id_for_event({}) is None + + +def test_build_query_appends_extra_conditions() -> None: + query = build_btql_sorted_page_query( + from_expr="project_logs('p') spans", + limit=10, + last_pagination_key="pk1", + created_after="2026-01-01T00:00:00Z", + extra_conditions=["span_attributes.name = 'chat.message'"], + ) + + assert "span_attributes.name = 'chat.message'" in query + assert "created >= '2026-01-01T00:00:00Z'" in query + assert "_pagination_key > 'pk1'" in query + + +def test_config_rejects_both_filters() -> None: + with pytest.raises(ValueError, match="Set only one of"): + MigrationConfig( + logs_include_root_span_name="a", + logs_exclude_root_span_name="b", + )