From b3bba167fa26152ba7987cffe08139d95c00c8d2 Mon Sep 17 00:00:00 2001 From: Curtis Galione Date: Wed, 1 Jul 2026 22:17:08 -0700 Subject: [PATCH] Add project mapping and safer dry-run output --- .env.example | 2 + README.md | 21 ++- braintrust_migrate/cli.py | 193 +++++++++++++++++++- braintrust_migrate/config.py | 97 +++++++++- braintrust_migrate/orchestration.py | 38 +++- tests/unit/test_cli_dry_run.py | 166 ++++++++++++++++- tests/unit/test_config.py | 111 ++++++++++- tests/unit/test_orchestrator_migrate_all.py | 153 +++++++++++++++- 8 files changed, 747 insertions(+), 34 deletions(-) diff --git a/.env.example b/.env.example index b84e5c2..dcb00d9 100644 --- a/.env.example +++ b/.env.example @@ -18,3 +18,5 @@ MIGRATION_RETRY_DELAY=1.0 MIGRATION_MAX_CONCURRENT=10 MIGRATION_CHECKPOINT_INTERVAL=50 MIGRATION_STATE_DIR=./checkpoints +# MIGRATION_PROJECT_MAP={"Source Project":"Destination Project"} +# MIGRATION_PROJECT_MAP_FILE=./project-map.json diff --git a/README.md b/README.md index 8f31f97..d2a2fef 100644 --- a/README.md +++ b/README.md @@ -122,7 +122,9 @@ All options can be set via environment variables or CLI flags. CLI flags take pr | Environment Variable | CLI Flag | Default | Description | |---------------------|----------|---------|-------------| | `MIGRATION_RESOURCES` | `--resources`, `-r` | `all` | Comma-separated list of resources to migrate. Options: `all`, `ai_secrets`, `roles`, `groups`, `datasets`, `project_tags`, `span_iframes`, `functions`, `prompts`, `project_scores`, `experiments`, `logs`, `views` | -| `MIGRATION_PROJECTS` | `--projects`, `-p` | *(all projects)* | Comma-separated list of project names to migrate | +| `MIGRATION_PROJECTS` | `--projects`, `-p` | *(all projects)* | Comma-separated list of source project names to migrate | +| `MIGRATION_PROJECT_MAP` | `--project-map` | *(same-name fallback)* | JSON object mapping source project names to destination project names, e.g. `{"Source Project":"Destination Project"}` | +| `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 | @@ -258,6 +260,14 @@ braintrust-migrate migrate --resources ai_secrets,datasets,prompts # Migrate specific projects only braintrust-migrate migrate --projects "Project A","Project B" + +# Migrate a source project into a differently named destination project +braintrust-migrate migrate \ + --resources logs \ + --projects "" \ + --project-map '{"":""}' \ + --created-after "" \ + --created-before "" ``` **Resume Migration:** @@ -293,6 +303,15 @@ braintrust-migrate migrate \ **Dry Run (Validation Only):** ```bash braintrust-migrate migrate --dry-run + +# Preview project mapping, destination existence, selected resources, and logs window +braintrust-migrate migrate \ + --dry-run \ + --resources logs \ + --projects "" \ + --project-map '{"":""}' \ + --created-after "" \ + --created-before "" ``` **Time-based Filtering:** diff --git a/braintrust_migrate/cli.py b/braintrust_migrate/cli.py index d3e08a6..d15db58 100644 --- a/braintrust_migrate/cli.py +++ b/braintrust_migrate/cli.py @@ -25,6 +25,8 @@ Config, canonicalize_created_after, canonicalize_created_before, + load_project_name_mapping_file, + parse_project_name_mapping_json, ) from braintrust_migrate.orchestration import MigrationOrchestrator @@ -43,6 +45,24 @@ console = Console() +def _load_project_name_mapping_override( + project_map: str | None, + project_map_file: Path | None, +) -> dict[str, str] | None: + """Load CLI/env project-name mapping overrides, if provided.""" + if project_map is not None and project_map_file is not None: + raise ValueError("Set only one of --project-map or --project-map-file") + if project_map is not None: + return parse_project_name_mapping_json(project_map, field_name="--project-map") + if project_map_file is not None: + return load_project_name_mapping_file(project_map_file) + return None + + +def _resources_include(resources: list[str], resource_name: str) -> bool: + return "all" in resources or resource_name in resources + + def setup_logging(log_level: str = "INFO", log_format: str = "json") -> None: """Setup structured logging. @@ -89,10 +109,29 @@ def migrate( typer.Option( "--projects", "-p", - help="Comma-separated list of project names to migrate (if not specified, all projects will be migrated)", + help="Comma-separated list of source project names to migrate (if not specified, all projects will be migrated)", envvar="MIGRATION_PROJECTS", ), ] = None, + project_map: Annotated[ + str | None, + typer.Option( + "--project-map", + help=( + "JSON object mapping source project names to destination project names " + '(e.g. \'{"Source Project":"Destination Project"}\')' + ), + envvar="MIGRATION_PROJECT_MAP", + ), + ] = None, + project_map_file: Annotated[ + Path | None, + typer.Option( + "--project-map-file", + help="Path to a JSON file mapping source project names to destination project names", + envvar="MIGRATION_PROJECT_MAP_FILE", + ), + ] = None, state_dir: Annotated[ Path | None, typer.Option( @@ -245,6 +284,8 @@ def migrate( _migrate_main( resources, projects, + project_map, + project_map_file, state_dir, resume_run_dir, log_level, @@ -266,6 +307,8 @@ def migrate( async def _migrate_main( resources: str, projects: str | None, + project_map: str | None, + project_map_file: Path | None, state_dir: Path | None, resume_run_dir: Path | None, log_level: str, @@ -302,6 +345,13 @@ async def _migrate_main( else: config = Config.from_env() + project_name_mapping_override = _load_project_name_mapping_override( + project_map, + project_map_file, + ) + if project_name_mapping_override is not None: + config.project_name_mapping = project_name_mapping_override + # Checkpoint normalization: allow a *single* checkpoint path (root/run/project) # to be provided via CLI or env. CLI takes precedence, but env-driven `--state-dir` # is supported as well. @@ -359,6 +409,7 @@ async def _migrate_main( dest_url=str(config.destination.url), resources=config.resources, projects=getattr(config, "project_names", None), + project_name_mapping=config.project_name_mapping, state_dir=str(config.state_dir), dry_run=dry_run, logs_fetch_limit=config.migration.logs_fetch_limit, @@ -1097,12 +1148,24 @@ async def _run_dry_run(config: Config) -> None: source_client, dest_client, config.project_names, + config.project_name_mapping, ) progress.update( validation_task, description=f"โœ… Discovered {len(projects)} projects", ) + if _resources_include(config.resources, "logs"): + log_probe_results = await _test_logs_dry_run_probe( + source_client, + projects, + config, + progress, + validation_task, + ) + else: + log_probe_results = {} + # Test resource discovery for each migrator type test_results = await _test_resource_discovery( source_client, @@ -1116,7 +1179,12 @@ async def _run_dry_run(config: Config) -> None: progress.update(validation_task, completed=1, total=1) # Display dry run results - _display_dry_run_results(projects, test_results) + _display_dry_run_results( + config, + projects, + test_results, + log_probe_results, + ) console.print("\n[green]โœ… Dry run completed successfully![/green]") console.print( @@ -1133,6 +1201,7 @@ async def _discover_projects_read_only( source_client, dest_client, project_names: list[str] | None = None, + project_name_mapping: dict[str, str] | None = None, ) -> list[dict[str, str]]: """Discover projects without mutating destination organization. @@ -1166,19 +1235,80 @@ async def _discover_projects_read_only( if selected_names and name not in selected_names: continue - dest_id_raw = dest_id_by_name.get(name) + dest_name = (project_name_mapping or {}).get(name, name) + dest_id_raw = dest_id_by_name.get(dest_name) dest_id = dest_id_raw if isinstance(dest_id_raw, str) else "" mappings.append( { "source_id": source_id, "dest_id": dest_id, "name": name, + "dest_name": dest_name, + "dest_status": "exists" if dest_id else "would create", } ) return mappings +async def _test_logs_dry_run_probe( + source_client, + projects: list[dict], + config: Config, + progress, + validation_task, +) -> dict[str, dict[str, Any]]: + """Probe whether the selected logs time window has at least one matching span.""" + from braintrust_migrate.btql import ( + btql_quote, + fetch_btql_sorted_page_with_retries, + ) + from braintrust_migrate.streaming_utils import build_btql_sorted_page_query + + results: dict[str, dict[str, Any]] = {} + for project in projects: + project_name = project["name"] + progress.update( + validation_task, + description=f"๐Ÿ” Probing logs window for {project_name}...", + ) + from_expr = f"project_logs('{btql_quote(project['source_id'])}') spans" + + def _query_text_for_limit(n: int, *, _from_expr: str = from_expr) -> str: + return build_btql_sorted_page_query( + from_expr=_from_expr, + limit=n, + last_pagination_key=None, + created_after=config.migration.created_after, + created_before=config.migration.created_before, + select="*", + ) + + try: + page = await fetch_btql_sorted_page_with_retries( + client=source_client, + query_for_limit=_query_text_for_limit, + configured_limit=1, + operation="btql_project_logs_dry_run_probe", + log_fields={"source_project_id": project["source_id"]}, + floor_limit=1, + default_500_retry_limit=1, + ) + rows = page.get("events") + match_found = isinstance(rows, list) and len(rows) > 0 + results[project_name] = { + "status": "success", + "matching_spans_visible": match_found, + } + except Exception as e: + results[project_name] = { + "status": "error", + "error": str(e), + } + + return results + + async def _test_resource_discovery( source_client, dest_client, @@ -1268,29 +1398,78 @@ async def _test_resource_discovery( def _display_dry_run_results( - projects: list[dict], test_results: dict[str, dict] + config: Config, + projects: list[dict], + test_results: dict[str, dict], + log_probe_results: dict[str, dict[str, Any]], ) -> None: """Display dry run results in a formatted table. Args: + config: Migration configuration. projects: List of discovered projects test_results: Results from resource discovery tests + log_probe_results: Read-only logs time-window probe results. """ + plan_table = Table(title="๐Ÿงญ Migration Plan") + plan_table.add_column("Setting", style="cyan") + plan_table.add_column("Value", style="yellow") + plan_table.add_row("Resources", ", ".join(config.resources)) + plan_table.add_row("Source Projects", ", ".join(config.project_names or ["all"])) + plan_table.add_row("Created After", config.migration.created_after or "none") + plan_table.add_row("Created Before", config.migration.created_before or "none") + if config.project_name_mapping: + map_summary = ", ".join( + f"{source} -> {dest}" + for source, dest in sorted(config.project_name_mapping.items()) + ) + else: + map_summary = "same-name fallback" + plan_table.add_row("Project Map", map_summary) + + console.print("\n") + console.print(plan_table) + # Projects table if projects: - projects_table = Table(title="๐Ÿ“ Discovered Projects") - projects_table.add_column("Project Name", style="cyan") + projects_table = Table(title="๐Ÿ“ Project Migration Plan") + projects_table.add_column("Source Project", style="cyan") + projects_table.add_column("Destination Project", style="yellow") projects_table.add_column("Source ID", style="blue") projects_table.add_column("Dest ID", style="green") + projects_table.add_column("Destination Status", style="magenta") for project in projects: projects_table.add_row( - project["name"], project["source_id"], project["dest_id"] + project["name"], + project.get("dest_name", project["name"]), + project["source_id"], + project["dest_id"] or "none", + project.get("dest_status", "exists" if project["dest_id"] else "would create"), ) console.print("\n") console.print(projects_table) + if log_probe_results: + logs_table = Table(title="๐Ÿงพ Logs Time-Window Probe") + logs_table.add_column("Project", style="cyan") + logs_table.add_column("Status", style="magenta") + logs_table.add_column("Result", style="yellow") + + for project_name, result in log_probe_results.items(): + if result["status"] == "success": + visible = bool(result.get("matching_spans_visible")) + status = "โœ… Success" + detail = "matching spans visible" if visible else "no matching spans" + else: + status = "โŒ Error" + detail = result.get("error", "Failed") + logs_table.add_row(project_name, status, detail) + + console.print("\n") + console.print(logs_table) + # 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 43910c0..bc9d854 100644 --- a/braintrust_migrate/config.py +++ b/braintrust_migrate/config.py @@ -1,10 +1,11 @@ """Configuration models and environment variable parsing for Braintrust migration tool.""" +import json import os import re from datetime import UTC, datetime from pathlib import Path -from typing import cast +from typing import Any, cast from dotenv import load_dotenv from pydantic import BaseModel, Field, HttpUrl, field_validator, model_validator @@ -56,6 +57,69 @@ def canonicalize_created_before(value: str) -> str: return _canonicalize_datetime(value, "created_before") +def normalize_project_name_mapping( + value: object | None, + *, + field_name: str = "project_name_mapping", +) -> dict[str, str]: + """Validate and normalize source project name -> destination project name maps.""" + if value is None: + return {} + + if isinstance(value, str): + return parse_project_name_mapping_json(value, field_name=field_name) + + if not isinstance(value, dict): + raise ValueError(f"{field_name} must be a JSON object") + + mapping: dict[str, str] = {} + for raw_source, raw_dest in value.items(): + if not isinstance(raw_source, str) or not isinstance(raw_dest, str): + raise ValueError( + f"{field_name} must map source project names to destination project names" + ) + + source = raw_source.strip() + dest = raw_dest.strip() + if not source or not dest: + raise ValueError( + f"{field_name} entries must have non-empty source and destination project names" + ) + if source in mapping: + raise ValueError(f"{field_name} contains duplicate source project {source!r}") + + mapping[source] = dest + + return mapping + + +def parse_project_name_mapping_json( + value: str, + *, + field_name: str = "project map", +) -> dict[str, str]: + """Parse a JSON project-name mapping and validate its shape.""" + try: + parsed: Any = json.loads(value) + except json.JSONDecodeError as e: + raise ValueError(f"{field_name} must be valid JSON") from e + + return normalize_project_name_mapping(parsed, field_name=field_name) + + +def load_project_name_mapping_file(path: Path) -> dict[str, str]: + """Load a project-name mapping from a JSON file.""" + if not path.exists(): + raise FileNotFoundError(f"Project map file not found: {path}") + with open(path, encoding="utf-8") as f: + try: + parsed: Any = json.load(f) + except json.JSONDecodeError as e: + raise ValueError(f"Project map file must contain valid JSON: {path}") from e + + return normalize_project_name_mapping(parsed, field_name="project map file") + + class BraintrustOrgConfig(BaseModel): """Configuration for a Braintrust organization.""" @@ -315,6 +379,13 @@ class Config(BaseModel): default=None, description="List of project names to migrate (if None, migrate all projects)", ) + project_name_mapping: dict[str, str] = Field( + default_factory=dict, + description=( + "Optional mapping from source project names to destination project names. " + "Unmapped projects use the source project name in the destination." + ), + ) class Config: """Pydantic config.""" @@ -322,6 +393,11 @@ class Config: validate_assignment = True use_enum_values = True + @field_validator("project_name_mapping", mode="before") + def validate_project_name_mapping(cls, v: object | None) -> dict[str, str]: + """Validate source project name -> destination project name mapping.""" + return normalize_project_name_mapping(v) + @classmethod def from_env(cls) -> "Config": """Create configuration from environment variables. @@ -483,6 +559,24 @@ def _get_bool(specific_key: str, unified_key: str, default: str) -> bool: # State directory state_dir = Path(os.getenv("MIGRATION_STATE_DIR", "./checkpoints")) + project_map_env = os.getenv("MIGRATION_PROJECT_MAP") + project_map_file_env = os.getenv("MIGRATION_PROJECT_MAP_FILE") + if project_map_env and project_map_file_env: + raise ValueError( + "Set only one of MIGRATION_PROJECT_MAP or MIGRATION_PROJECT_MAP_FILE" + ) + if project_map_env: + project_name_mapping = parse_project_name_mapping_json( + project_map_env, + field_name="MIGRATION_PROJECT_MAP", + ) + elif project_map_file_env: + project_name_mapping = load_project_name_mapping_file( + Path(project_map_file_env) + ) + else: + project_name_mapping = {} + return cls( source=BraintrustOrgConfig( api_key=source_api_key, @@ -526,6 +620,7 @@ def _get_bool(specific_key: str, unified_key: str, default: str) -> bool: format=log_format, ), state_dir=state_dir, + project_name_mapping=project_name_mapping, ) @classmethod diff --git a/braintrust_migrate/orchestration.py b/braintrust_migrate/orchestration.py index 57b2aa4..e540527 100644 --- a/braintrust_migrate/orchestration.py +++ b/braintrust_migrate/orchestration.py @@ -259,7 +259,8 @@ async def migrate_all( "Added project mapping", source_project_id=project["source_id"], dest_project_id=project["dest_id"], - project_name=project["name"], + source_project_name=project["name"], + dest_project_name=project.get("dest_name"), ) # STEP 1: Migrate organization-scoped resources once @@ -358,6 +359,7 @@ async def _run_project_with_hooks( project_results = { "project_id": project.get("dest_id"), "project_name": project.get("name"), + "dest_project_name": project.get("dest_name"), "resources": {}, "total_resources": 0, "migrated_resources": 0, @@ -535,12 +537,19 @@ async def _discover_projects( # Ensure projects exist in destination and get destination IDs project_mappings = [] for project in projects: - dest_project_id = await self._ensure_project_exists(project, dest_client) + source_name = cast(str, project.get("name")) + dest_name = self.config.project_name_mapping.get(source_name, source_name) + dest_project_id = await self._ensure_project_exists( + project, + dest_client, + dest_project_name=dest_name, + ) project_mappings.append( { "source_id": cast(str, project.get("id")), "dest_id": dest_project_id, - "name": cast(str, project.get("name")), + "name": source_name, + "dest_name": dest_name, "description": project.get("description"), } ) @@ -551,12 +560,15 @@ async def _ensure_project_exists( self, source_project: dict[str, Any], dest_client: BraintrustClient, + *, + dest_project_name: str, ) -> str: """Ensure a project exists in the destination organization. Args: source_project: Source project to replicate. dest_client: Destination client. + dest_project_name: Destination project name to look up or create. Returns: Destination project ID. @@ -572,20 +584,21 @@ async def _ensure_project_exists( existing_project: dict[str, Any] | None = None for dest_project in dest_projects: - if dest_project.get("name") == source_project.get("name"): + if dest_project.get("name") == dest_project_name: existing_project = dest_project break if existing_project: self._logger.debug( "Project already exists in destination", - project_name=source_project.get("name"), + source_project_name=source_project.get("name"), + dest_project_name=dest_project_name, dest_id=existing_project.get("id"), ) return cast(str, existing_project.get("id")) # Create project in destination - create_params = {"name": source_project.get("name")} + create_params = {"name": dest_project_name} description = cast(str | None, source_project.get("description")) if description: create_params["description"] = description @@ -605,7 +618,8 @@ async def _ensure_project_exists( self._logger.info( "Created project in destination", - project_name=source_project.get("name"), + source_project_name=source_project.get("name"), + dest_project_name=dest_project_name, source_id=source_project.get("id"), dest_id=new_project_id, ) @@ -615,7 +629,8 @@ async def _ensure_project_exists( except Exception as e: self._logger.error( "Failed to ensure project exists", - project_name=source_project.get("name"), + source_project_name=source_project.get("name"), + dest_project_name=dest_project_name, error=str(e), ) raise @@ -647,6 +662,7 @@ async def _migrate_project( Migration results for the project. """ project_name = project["name"] + dest_project_name = project.get("dest_name", project_name) source_project_id = project["source_id"] dest_project_id = project["dest_id"] @@ -654,6 +670,7 @@ async def _migrate_project( f"Starting migration for project: {project_name}", source_project_id=source_project_id, dest_project_id=dest_project_id, + dest_project_name=dest_project_name, ) # Create project-specific checkpoint directory @@ -663,6 +680,7 @@ async def _migrate_project( project_results = { "project_id": dest_project_id, # Use destination project ID in results "project_name": project_name, + "dest_project_name": dest_project_name, "resources": {}, "total_resources": 0, "migrated_resources": 0, @@ -1019,6 +1037,7 @@ def _generate_migration_report( for project_name, project_data in results.get("projects", {}).items(): project_summary = { "project_name": project_name, + "dest_project_name": project_data.get("dest_project_name", project_name), "project_id": project_data.get("project_id"), "total_resources": project_data.get("total_resources", 0), "migrated_resources": project_data.get("migrated_resources", 0), @@ -1145,6 +1164,9 @@ def _write_human_readable_summary( f.write("## Project Breakdown\n") for project_name, project_data in detailed_report["projects"].items(): f.write(f"\n### {project_name}\n") + dest_project_name = project_data.get("dest_project_name", project_name) + if dest_project_name != project_name: + f.write(f"Destination Project: {dest_project_name}\n") f.write(f"Project ID: {project_data['project_id']}\n") f.write(f"Resources: {project_data['total_resources']} total, ") f.write(f"{project_data['migrated_resources']} migrated, ") diff --git a/tests/unit/test_cli_dry_run.py b/tests/unit/test_cli_dry_run.py index b571fb2..ec57132 100644 --- a/tests/unit/test_cli_dry_run.py +++ b/tests/unit/test_cli_dry_run.py @@ -7,7 +7,12 @@ import pytest from pydantic import HttpUrl -from braintrust_migrate.cli import _discover_projects_read_only, _run_dry_run +from braintrust_migrate.cli import ( + _discover_projects_read_only, + _load_project_name_mapping_override, + _run_dry_run, + _test_logs_dry_run_probe, +) from braintrust_migrate.config import BraintrustOrgConfig, Config, MigrationConfig @@ -53,12 +58,162 @@ async def with_retry(_op_name, coro_func): projects = await _discover_projects_read_only(source, dest) assert projects == [ - {"source_id": "src-1", "dest_id": "dest-1", "name": "Project A"}, - {"source_id": "src-2", "dest_id": "", "name": "Project B"}, + { + "source_id": "src-1", + "dest_id": "dest-1", + "name": "Project A", + "dest_name": "Project A", + "dest_status": "exists", + }, + { + "source_id": "src-2", + "dest_id": "", + "name": "Project B", + "dest_name": "Project B", + "dest_status": "would create", + }, ] dest.create_project.assert_not_called() +@pytest.mark.asyncio +async def test_discover_projects_read_only_uses_project_name_mapping() -> None: + """Read-only discovery should resolve mapped destination project names.""" + source = Mock() + dest = Mock() + source.list_projects = AsyncMock( + return_value=[{"id": "src-1", "name": "Project A"}] + ) + dest.list_projects = AsyncMock(return_value=[{"id": "dest-1", "name": "Project Z"}]) + dest.create_project = AsyncMock() + + async def with_retry(_op_name, coro_func): + result = coro_func() + if hasattr(result, "__await__"): + return await result + return result + + source.with_retry = with_retry + dest.with_retry = with_retry + + projects = await _discover_projects_read_only( + source, + dest, + project_name_mapping={"Project A": "Project Z"}, + ) + + assert projects == [ + { + "source_id": "src-1", + "dest_id": "dest-1", + "name": "Project A", + "dest_name": "Project Z", + "dest_status": "exists", + } + ] + dest.create_project.assert_not_called() + + +@pytest.mark.asyncio +async def test_discover_projects_read_only_reports_mapped_destination_would_create() -> None: + """Dry-run project discovery should report when a mapped destination is absent.""" + source = Mock() + dest = Mock() + source.list_projects = AsyncMock( + return_value=[{"id": "src-1", "name": "Project A"}] + ) + dest.list_projects = AsyncMock(return_value=[]) + dest.create_project = AsyncMock() + + async def with_retry(_op_name, coro_func): + result = coro_func() + if hasattr(result, "__await__"): + return await result + return result + + source.with_retry = with_retry + dest.with_retry = with_retry + + projects = await _discover_projects_read_only( + source, + dest, + project_name_mapping={"Project A": "Project Z"}, + ) + + assert projects == [ + { + "source_id": "src-1", + "dest_id": "", + "name": "Project A", + "dest_name": "Project Z", + "dest_status": "would create", + } + ] + dest.create_project.assert_not_called() + + +def test_load_project_name_mapping_override_rejects_inline_and_file( + tmp_path: Path, +) -> None: + """CLI project-map override accepts either inline JSON or a file, not both.""" + project_map_file = tmp_path / "project-map.json" + project_map_file.write_text('{"Project A":"Project Z"}') + + with pytest.raises(ValueError, match="Set only one"): + _load_project_name_mapping_override( + '{"Project A":"Project Z"}', + project_map_file, + ) + + +@pytest.mark.asyncio +async def test_logs_dry_run_probe_uses_source_project_and_created_filters( + tmp_path: Path, +) -> None: + """Logs dry-run probe should be read-only and use source id plus time filters.""" + config = _make_config(tmp_path) + config.resources = ["logs"] + config.migration.created_after = "2026-01-01T00:00:00Z" + config.migration.created_before = "2026-01-02T00:00:00Z" + + source = Mock() + queries: list[str] = [] + + async def raw_request(method, path, *, json=None, **_kwargs): + assert method == "POST" + assert path == "/btql" + assert json is not None + query = json["query"] + queries.append(query) + return {"data": [{"id": "span-1", "_pagination_key": "p1"}]} + + async def with_retry(_op_name, coro_func, **_kwargs): + result = coro_func() + if hasattr(result, "__await__"): + return await result + return result + + source.raw_request = raw_request + source.with_retry = with_retry + + results = await _test_logs_dry_run_probe( + source, + [{"source_id": "src-project-id", "name": "Project A"}], + config, + Mock(), + 1, + ) + + assert results == { + "Project A": {"status": "success", "matching_spans_visible": True} + } + assert len(queries) == 1 + assert "project_logs('src-project-id') spans" in queries[0] + assert "created >= '2026-01-01T00:00:00Z'" in queries[0] + assert "created < '2026-01-02T00:00:00Z'" in queries[0] + assert "limit: 1" in queries[0] + + @pytest.mark.asyncio async def test_run_dry_run_does_not_create_destination_projects(tmp_path: Path) -> None: """Dry-run should not create missing destination projects.""" @@ -94,6 +249,11 @@ async def mock_create_client_pair(_source_cfg, _dest_cfg, _migration_cfg): new_callable=AsyncMock, return_value={}, ), + patch( + "braintrust_migrate.cli._test_logs_dry_run_probe", + new_callable=AsyncMock, + return_value={}, + ), patch("braintrust_migrate.cli._display_dry_run_results"), ): await _run_dry_run(config) diff --git a/tests/unit/test_config.py b/tests/unit/test_config.py index 367a56c..7333725 100644 --- a/tests/unit/test_config.py +++ b/tests/unit/test_config.py @@ -4,7 +4,13 @@ import pytest -from braintrust_migrate.config import BraintrustOrgConfig, Config, MigrationConfig +from braintrust_migrate.config import ( + BraintrustOrgConfig, + Config, + MigrationConfig, + load_project_name_mapping_file, + parse_project_name_mapping_json, +) # Test constants DEFAULT_BATCH_SIZE = 100 @@ -13,6 +19,9 @@ DEFAULT_CHECKPOINT_INTERVAL = 50 TEST_BATCH_SIZE = 50 TEST_RETRY_ATTEMPTS = 5 +TEST_EVENTS_FETCH_GROUP_SIZE = 17 +TEST_EVENTS_FLUSH_MAX_ROWS = 4321 +TEST_LEGACY_LOGS_INSERT_BATCH_SIZE = 3456 class TestBraintrustOrgConfig: @@ -105,6 +114,47 @@ def test_checkpoint_dir_methods(self): project_dir = config.get_checkpoint_dir("my-project") assert project_dir == Path("/tmp/test-checkpoints/my-project") + def test_project_name_mapping_validates_and_trims(self): + """Test project name mappings are normalized on config objects.""" + config = Config( + source=BraintrustOrgConfig(api_key="source-key"), + destination=BraintrustOrgConfig(api_key="dest-key"), + project_name_mapping={" Source A ": " Dest A "}, + ) + + assert config.project_name_mapping == {"Source A": "Dest A"} + + def test_project_name_mapping_rejects_empty_names(self): + """Test project name mappings require non-empty source and dest names.""" + with pytest.raises(ValueError, match="non-empty source and destination"): + Config( + source=BraintrustOrgConfig(api_key="source-key"), + destination=BraintrustOrgConfig(api_key="dest-key"), + project_name_mapping={"Source A": " "}, + ) + + def test_project_name_mapping_rejects_non_string_names(self): + """Test project name mappings must use string keys and values.""" + with pytest.raises(ValueError, match="source project names"): + Config( + source=BraintrustOrgConfig(api_key="source-key"), + destination=BraintrustOrgConfig(api_key="dest-key"), + project_name_mapping={"Source A": 123}, + ) + + def test_parse_project_name_mapping_json(self): + """Test inline JSON project map parsing.""" + mapping = parse_project_name_mapping_json('{"Source A":"Dest A"}') + + assert mapping == {"Source A": "Dest A"} + + def test_load_project_name_mapping_file(self, tmp_path: Path): + """Test JSON project map file parsing.""" + path = tmp_path / "project-map.json" + path.write_text('{"Source A":"Dest A"}') + + assert load_project_name_mapping_file(path) == {"Source A": "Dest A"} + class TestConfigFromEnv: """Test configuration loading from environment variables.""" @@ -128,11 +178,15 @@ def test_valid_env_config(self, monkeypatch): monkeypatch.setenv("BT_SOURCE_URL", "https://source.example.com") monkeypatch.setenv("BT_DEST_URL", "https://dest.example.com") monkeypatch.setenv("MIGRATION_BATCH_SIZE", "50") + monkeypatch.setenv("MIGRATION_PROJECT_MAP", '{"Source A":"Dest A"}') monkeypatch.setenv("MIGRATION_ACL_MAP_USERS", "true") monkeypatch.setenv("MIGRATION_ACL_AUTO_INVITE_USERS", "true") monkeypatch.setenv("MIGRATION_GROUP_MAP_USERS", "true") monkeypatch.setenv("MIGRATION_GROUP_AUTO_INVITE_USERS", "true") - monkeypatch.setenv("MIGRATION_EVENTS_FETCH_GROUP_SIZE", "17") + monkeypatch.setenv( + "MIGRATION_EVENTS_FETCH_GROUP_SIZE", + str(TEST_EVENTS_FETCH_GROUP_SIZE), + ) monkeypatch.setenv("LOG_LEVEL", "DEBUG") config = Config.from_env() @@ -142,31 +196,70 @@ def test_valid_env_config(self, monkeypatch): assert str(config.source.url) == "https://source.example.com/" assert str(config.destination.url) == "https://dest.example.com/" assert config.migration.batch_size == TEST_BATCH_SIZE + assert config.project_name_mapping == {"Source A": "Dest A"} assert config.migration.acl_map_users is True assert config.migration.acl_auto_invite_users is True assert config.migration.group_map_users is True assert config.migration.group_auto_invite_users is True - assert config.migration.events_fetch_group_size == 17 + assert config.migration.events_fetch_group_size == TEST_EVENTS_FETCH_GROUP_SIZE assert config.logging.level == "DEBUG" def test_unified_events_flush_max_rows_from_env(self, monkeypatch): """Test shared streaming flush threshold env var.""" monkeypatch.setenv("BT_SOURCE_API_KEY", "source-test-key") monkeypatch.setenv("BT_DEST_API_KEY", "dest-test-key") - monkeypatch.setenv("MIGRATION_EVENTS_FLUSH_MAX_ROWS", "4321") + monkeypatch.setenv( + "MIGRATION_EVENTS_FLUSH_MAX_ROWS", + str(TEST_EVENTS_FLUSH_MAX_ROWS), + ) + + config = Config.from_env() + + assert config.migration.events_flush_max_rows == TEST_EVENTS_FLUSH_MAX_ROWS + assert config.migration.logs_insert_batch_size == TEST_EVENTS_FLUSH_MAX_ROWS + + def test_project_map_file_from_env(self, monkeypatch, tmp_path: Path): + """Test loading project name mapping from env-provided file.""" + path = tmp_path / "project-map.json" + path.write_text('{"Source A":"Dest A"}') + monkeypatch.setenv("BT_SOURCE_API_KEY", "source-test-key") + monkeypatch.setenv("BT_DEST_API_KEY", "dest-test-key") + monkeypatch.setenv("MIGRATION_PROJECT_MAP_FILE", str(path)) config = Config.from_env() - assert config.migration.events_flush_max_rows == 4321 - assert config.migration.logs_insert_batch_size == 4321 + assert config.project_name_mapping == {"Source A": "Dest A"} + + def test_project_map_env_vars_are_mutually_exclusive( + self, monkeypatch, tmp_path: Path + ): + """Test inline and file project maps cannot both be set.""" + path = tmp_path / "project-map.json" + path.write_text('{"Source A":"Dest A"}') + monkeypatch.setenv("BT_SOURCE_API_KEY", "source-test-key") + monkeypatch.setenv("BT_DEST_API_KEY", "dest-test-key") + monkeypatch.setenv("MIGRATION_PROJECT_MAP", '{"Source A":"Dest A"}') + monkeypatch.setenv("MIGRATION_PROJECT_MAP_FILE", str(path)) + + with pytest.raises(ValueError, match="Set only one"): + Config.from_env() def test_legacy_logs_insert_batch_size_alias_still_works(self, monkeypatch): """Test legacy logs-only env var maps to shared flush threshold.""" monkeypatch.setenv("BT_SOURCE_API_KEY", "source-test-key") monkeypatch.setenv("BT_DEST_API_KEY", "dest-test-key") - monkeypatch.setenv("MIGRATION_LOGS_INSERT_BATCH_SIZE", "3456") + monkeypatch.setenv( + "MIGRATION_LOGS_INSERT_BATCH_SIZE", + str(TEST_LEGACY_LOGS_INSERT_BATCH_SIZE), + ) config = Config.from_env() - assert config.migration.events_flush_max_rows == 3456 - assert config.migration.logs_insert_batch_size == 3456 + assert ( + config.migration.events_flush_max_rows + == TEST_LEGACY_LOGS_INSERT_BATCH_SIZE + ) + assert ( + config.migration.logs_insert_batch_size + == TEST_LEGACY_LOGS_INSERT_BATCH_SIZE + ) diff --git a/tests/unit/test_orchestrator_migrate_all.py b/tests/unit/test_orchestrator_migrate_all.py index 2c3e3a8..995df9d 100644 --- a/tests/unit/test_orchestrator_migrate_all.py +++ b/tests/unit/test_orchestrator_migrate_all.py @@ -1,6 +1,7 @@ from __future__ import annotations import asyncio +import json from contextlib import asynccontextmanager from pathlib import Path from unittest.mock import AsyncMock, Mock @@ -10,6 +11,9 @@ from braintrust_migrate.config import Config, MigrationConfig from braintrust_migrate.orchestration import MigrationOrchestrator +EXPECTED_PROJECT_COUNT = 3 +EXPECTED_MAX_PROJECT_CONCURRENCY = 2 + def _make_config(tmp_path: Path, *, max_concurrent: int = 2) -> Config: return Config( @@ -21,6 +25,145 @@ def _make_config(tmp_path: Path, *, max_concurrent: int = 2) -> Config: ) +async def _with_retry(_op_name, coro_func): + result = coro_func() + if hasattr(result, "__await__"): + return await result + return result + + +@pytest.mark.asyncio +async def test_discover_projects_preserves_same_name_behavior(tmp_path: Path) -> None: + """Unmapped projects should still resolve destination projects by source name.""" + orchestrator = MigrationOrchestrator(_make_config(tmp_path)) + source = Mock() + dest = Mock() + source.list_projects = AsyncMock( + return_value=[{"id": "src-1", "name": "Project A"}] + ) + dest.list_projects = AsyncMock(return_value=[{"id": "dest-1", "name": "Project A"}]) + dest.create_project = AsyncMock() + source.with_retry = _with_retry + dest.with_retry = _with_retry + + projects = await orchestrator._discover_projects(source, dest) + + assert projects == [ + { + "source_id": "src-1", + "dest_id": "dest-1", + "name": "Project A", + "dest_name": "Project A", + "description": None, + } + ] + dest.create_project.assert_not_called() + + +@pytest.mark.asyncio +async def test_discover_projects_uses_existing_mapped_destination( + tmp_path: Path, +) -> None: + """Mapped projects should resolve an existing differently named destination.""" + config = _make_config(tmp_path) + config.project_name_mapping = {"Project A": "Project Z"} + orchestrator = MigrationOrchestrator(config) + source = Mock() + dest = Mock() + source.list_projects = AsyncMock( + return_value=[{"id": "src-1", "name": "Project A"}] + ) + dest.list_projects = AsyncMock(return_value=[{"id": "dest-1", "name": "Project Z"}]) + dest.create_project = AsyncMock() + source.with_retry = _with_retry + dest.with_retry = _with_retry + + projects = await orchestrator._discover_projects(source, dest) + + assert projects == [ + { + "source_id": "src-1", + "dest_id": "dest-1", + "name": "Project A", + "dest_name": "Project Z", + "description": None, + } + ] + dest.create_project.assert_not_called() + + +@pytest.mark.asyncio +async def test_discover_projects_creates_missing_mapped_destination( + tmp_path: Path, +) -> None: + """Mapped projects should create the mapped destination name when absent.""" + config = _make_config(tmp_path) + config.project_name_mapping = {"Project A": "Project Z"} + orchestrator = MigrationOrchestrator(config) + source = Mock() + dest = Mock() + source.list_projects = AsyncMock( + return_value=[{"id": "src-1", "name": "Project A", "description": "desc"}] + ) + dest.list_projects = AsyncMock(return_value=[]) + dest.create_project = AsyncMock(return_value={"id": "dest-new"}) + source.with_retry = _with_retry + dest.with_retry = _with_retry + + projects = await orchestrator._discover_projects(source, dest) + + assert projects == [ + { + "source_id": "src-1", + "dest_id": "dest-new", + "name": "Project A", + "dest_name": "Project Z", + "description": "desc", + } + ] + dest.create_project.assert_awaited_once_with(name="Project Z", description="desc") + + +def test_migration_report_includes_destination_project_name(tmp_path: Path) -> None: + """Reports should expose mapped destination project names.""" + orchestrator = MigrationOrchestrator(_make_config(tmp_path)) + checkpoint_dir = tmp_path / "run" + checkpoint_dir.mkdir() + results = { + "start_time": "2026-01-01T00:00:00", + "end_time": "2026-01-01T00:00:01", + "duration_seconds": 1.0, + "success": True, + "summary": { + "total_projects": 1, + "total_resources": 0, + "migrated_resources": 0, + "skipped_resources": 0, + "failed_resources": 0, + }, + "organization_resources": {}, + "projects": { + "Project A": { + "project_id": "dest-1", + "project_name": "Project A", + "dest_project_name": "Project Z", + "resources": {}, + "total_resources": 0, + "migrated_resources": 0, + "skipped_resources": 0, + "failed_resources": 0, + "errors": [], + } + }, + } + + report_path = orchestrator._generate_migration_report(results, checkpoint_dir) + report = json.loads(report_path.read_text()) + + assert report["projects"]["Project A"]["project_name"] == "Project A" + assert report["projects"]["Project A"]["dest_project_name"] == "Project Z" + + @pytest.mark.asyncio async def test_migrate_all_runs_projects_concurrently_and_emits_hooks( tmp_path: Path, @@ -132,9 +275,9 @@ async def fake_migrate_project( ) assert discovered == [p["name"] for p in projects] - assert len(started) == 3 - assert len(completed) == 3 - assert max_seen == 2 - assert results["summary"]["total_projects"] == 3 - assert results["summary"]["migrated_resources"] == 3 + assert len(started) == EXPECTED_PROJECT_COUNT + assert len(completed) == EXPECTED_PROJECT_COUNT + assert max_seen == EXPECTED_MAX_PROJECT_CONCURRENCY + assert results["summary"]["total_projects"] == EXPECTED_PROJECT_COUNT + assert results["summary"]["migrated_resources"] == EXPECTED_PROJECT_COUNT assert results["report_path"] == str(report_path)