diff --git a/CHANGELOG.md b/CHANGELOG.md index 0218f15..cf0c2d0 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 +- Optional post-migration validation (`--validate` / `MIGRATION_VALIDATE`). After migrating, validates the destination against the source: object parity (counts + the specific items missing in the destination) for project-scoped resources, plus event count-parity for datasets, experiments, and logs. Validation mirrors migration policy (e.g. bundle-backed code functions, which are intentionally skipped, are not reported as missing), honors the same `created_after`/`created_before` window, and is reported per-resource in the console and in `migration_report.json`. Note: for high-volume event resources we compare counts only — enumerating *which* events are missing does not scale; counts are reported per dataset/experiment/project so a discrepancy can still be localized. (roles/groups/ai_secrets/ACLs are out of scope.) ### Changed diff --git a/README.md b/README.md index 8f31f97..23b2eb7 100644 --- a/README.md +++ b/README.md @@ -402,6 +402,20 @@ The tool uses **two-level checkpointing** for streaming resources (logs, experim On resume: skips 1-30 (done), resumes experiment 31 from saved `_pagination_key`, continues with 32-100. +### Validation + +Pass `--validate` (or `MIGRATION_VALIDATE=true`) to validate the destination against the source as a final phase after migrating. It runs per project and checks: + +- **Object parity** for project-scoped resources (datasets, experiments, prompts, functions, project_scores, views, project_tags, span_iframes): compares source vs destination and reports the **specific items missing** in the destination. +- **Event count-parity** for datasets, experiments, and logs: compares a cheap `count` of events on each side (per dataset/experiment, and per project for logs). + +Results are summarized per resource in the console (mismatches are listed, with a pointer to the full report) and recorded under each project's `validation` block in `migration_report.json`. + +Notes and limitations: +- Validation **mirrors migration policy** — intentionally skipped items (e.g. bundle-backed code functions) are **not** reported as missing — and honors the same `--created-after`/`--created-before` window. +- For high-volume event resources it compares **counts only**. Enumerating *which* events are missing requires diffing full id sets and does not scale; counts are reported per dataset/experiment/project so a discrepancy can still be localized to a specific object. +- Org-scoped resources (roles, groups, ai_secrets) and ACLs are **out of scope** for validation. + ## Parallelization The migration tool currently uses **two active levels of concurrency**. The env vars in [Parallelization Tuning](#parallelization-tuning) still matter, but the within-project resource-type DAG concurrency described below has not been implemented yet. diff --git a/braintrust_migrate/cli.py b/braintrust_migrate/cli.py index d3e08a6..aaf2a10 100644 --- a/braintrust_migrate/cli.py +++ b/braintrust_migrate/cli.py @@ -227,6 +227,18 @@ def migrate( envvar="MIGRATION_GROUP_AUTO_INVITE_USERS", ), ] = None, + validate: Annotated[ + bool | None, + typer.Option( + "--validate/--no-validate", + help=( + "After migrating, validate the destination against the source " + "(object counts + which items are missing; event count-parity " + "for datasets/experiments/logs)." + ), + envvar="MIGRATION_VALIDATE", + ), + ] = None, ) -> None: """Migrate resources from source to destination Braintrust organization. @@ -259,6 +271,7 @@ def migrate( acl_auto_invite_users, group_map_users, group_auto_invite_users, + validate, ) ) @@ -280,6 +293,7 @@ async def _migrate_main( acl_auto_invite_users: bool | None, group_map_users: bool | None, group_auto_invite_users: bool | None, + validate: bool | None = None, ) -> None: """Async implementation of the migrate command.""" setup_logging(log_level, log_format) @@ -352,6 +366,8 @@ async def _migrate_main( config.migration.group_map_users = group_map_users if group_auto_invite_users is not None: config.migration.group_auto_invite_users = group_auto_invite_users + if validate is not None: + config.migration.validate_migration = validate logger.info( "Starting migration", @@ -914,6 +930,52 @@ def _display_results(results: dict) -> None: f" ... and {len(summary['errors']) - MAX_ERRORS_TO_DISPLAY} more errors" ) + # Validation summary (only when --validate was used). + validations = [ + (pname, pdata["validation"]) + for pname, pdata in results.get("projects", {}).items() + if pdata.get("validation") + ] + if validations: + total_checks = sum(len(v.get("resources", [])) for _, v in validations) + failures = [ + (pname, r) + for pname, v in validations + for r in v.get("resources", []) + if not r.get("ok") + ] + console.print("\n[bold]Validation[/bold]") + if not failures: + console.print( + f" [green]✅ all {total_checks} validated resource(s) match[/green]" + ) + else: + console.print( + f" [red]❌ {len(failures)} of {total_checks} validated " + f"resource(s) have mismatches[/red]" + ) + for pname, r in failures: + console.print(f" [red]• {pname}/{r['resource_type']}[/red]") + if r.get("error"): + console.print(f" error: {r['error']}") + for check in r.get("checks", []): + if check.get("ok"): + continue + if check.get("kind") == "object" and check.get("missing"): + miss = check["missing"] + shown = ", ".join(miss[:MAX_ERRORS_TO_DISPLAY]) + more = ( + f" (+{len(miss) - MAX_ERRORS_TO_DISPLAY} more)" + if len(miss) > MAX_ERRORS_TO_DISPLAY + else "" + ) + console.print(f" missing {len(miss)}: {shown}{more}") + elif check.get("kind") == "events": + console.print( + f" {check['scope']}: " + f"source={check['source_count']} dest={check['dest_count']}" + ) + # Always point to the full per-item report so users can drill into exactly # what was migrated / skipped / failed (and why) — counts above, detail here. report_path = results.get("report_path") diff --git a/braintrust_migrate/config.py b/braintrust_migrate/config.py index 43910c0..0c795e1 100644 --- a/braintrust_migrate/config.py +++ b/braintrust_migrate/config.py @@ -168,6 +168,15 @@ class MigrationConfig(BaseModel): description="Use a SQLite seen-id store to prevent older versions overwriting newer ones during pagination", ) + validate_migration: bool = Field( + default=False, + description=( + "After migrating, validate the destination against the source: object " + "counts (and which specific items are missing) for project-scoped " + "resources, plus event count-parity for datasets/experiments/logs." + ), + ) + created_after: str | None = Field( default=None, description=( @@ -403,6 +412,13 @@ def _get_bool(specific_key: str, unified_key: str, default: str) -> bool: logs_use_seen_db = _get_bool( "MIGRATION_LOGS_USE_SEEN_DB", "MIGRATION_EVENTS_USE_SEEN_DB", "true" ) + validate = os.getenv("MIGRATION_VALIDATE", "false").lower() in { + "1", + "true", + "yes", + "y", + "on", + } # Optional time filter (applies to logs and experiments) created_after = os.getenv("MIGRATION_CREATED_AFTER") @@ -507,6 +523,7 @@ 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, + validate_migration=validate, created_after=created_after, created_before=created_before, experiment_events_fetch_limit=experiment_events_fetch_limit, diff --git a/braintrust_migrate/orchestration.py b/braintrust_migrate/orchestration.py index 57b2aa4..c20d34c 100644 --- a/braintrust_migrate/orchestration.py +++ b/braintrust_migrate/orchestration.py @@ -11,6 +11,8 @@ from braintrust_migrate.client import BraintrustClient, create_client_pair from braintrust_migrate.config import Config +from braintrust_migrate.resources.base import list_resources +from braintrust_migrate.validation import validate_resource from braintrust_migrate.resources import ( ACLMigrator, AISecretMigrator, @@ -858,6 +860,29 @@ async def _migrate_project( } ) + # Optional post-migration validation (source vs destination parity). + if self.config.migration.validate_migration: + try: + project_results["validation"] = await self._validate_project( + source_client=source_client, + dest_client=dest_client, + source_project_id=source_project_id, + dest_project_id=dest_project_id, + migrated_resource_types=list(project_results["resources"].keys()), + id_mapping=global_id_mappings, + ) + except Exception as e: + self._logger.error( + "Project validation phase failed", + project=project_name, + error=str(e), + ) + project_results["validation"] = { + "ok": False, + "error": str(e), + "resources": [], + } + self._logger.info( f"Completed project migration: {project_name}", total_resources=project_results["total_resources"], @@ -868,6 +893,71 @@ async def _migrate_project( return project_results + async def _validate_project( + self, + *, + source_client: BraintrustClient, + dest_client: BraintrustClient, + source_project_id: str, + dest_project_id: str, + migrated_resource_types: list[str], + id_mapping: dict[str, str], + ) -> dict[str, Any]: + """Validate a project's migrated resources against the source. + + Object parity (counts + which items are missing) plus event count-parity + for datasets/experiments/logs. Out-of-scope resource types are skipped + (validate_resource returns None). + """ + results: list[dict[str, Any]] = [] + for resource_type in migrated_resource_types: + try: + validation = await validate_resource( + resource_type, + source_client=source_client, + dest_client=dest_client, + source_project_id=source_project_id, + dest_project_id=dest_project_id, + list_fn=list_resources, + id_mapping=id_mapping, + created_after=self.config.migration.created_after, + created_before=self.config.migration.created_before, + ) + except Exception as e: + self._logger.error( + "Validation failed for resource", + resource=resource_type, + error=str(e), + ) + results.append( + { + "resource_type": resource_type, + "ok": False, + "error": str(e), + "checks": [], + } + ) + continue + + if validation is None: + continue + results.append(validation.to_dict()) + if not validation.ok: + self._logger.warning( + "Validation mismatch", + resource=resource_type, + checks=[c.to_dict() for c in validation.checks if not c.ok], + ) + + ok = all(r.get("ok", False) for r in results) + self._logger.info( + "Project validation complete", + project_id=dest_project_id, + ok=ok, + validated=len(results), + ) + return {"ok": ok, "resources": results} + def _get_organization_resources_to_migrate(self) -> list[str]: """Get list of organization-scoped resource types to migrate based on configuration. @@ -1073,6 +1163,11 @@ def _generate_migration_report( } ) + # Include post-migration validation results when present. + validation = project_data.get("validation") + if validation is not None: + project_summary["validation"] = validation + detailed_report["projects"][project_name] = project_summary # Write JSON report to file diff --git a/braintrust_migrate/resources/base.py b/braintrust_migrate/resources/base.py index 7747a4d..b687ae2 100644 --- a/braintrust_migrate/resources/base.py +++ b/braintrust_migrate/resources/base.py @@ -22,6 +22,77 @@ DEFAULT_LIST_PAGE_SIZE: int = 1000 +async def list_resources( + client: BraintrustClient, + resource_type: str, + project_id: str | None = None, + *, + additional_params: dict | None = None, + client_side_filter_field: str | None = None, + log: Any = logger, +) -> list[dict[str, Any]]: + """List all resources of a type from a client via the raw paginated API. + + Standalone so both ``ResourceMigrator`` and post-migration validation can + share one lister (``GET /v1/`` with `starting_after` pagination). + """ + try: + params: dict[str, Any] = {} + if project_id and not client_side_filter_field: + params["project_id"] = project_id + if additional_params: + params.update(additional_params) + + api_path = resource_type.rstrip("s") + all_resources: list[dict[str, Any]] = [] + starting_after = None + page_num = 0 + + while True: + page_num += 1 + page_params = {**params, "limit": DEFAULT_LIST_PAGE_SIZE} + if starting_after is not None: + page_params["starting_after"] = starting_after + + response = await client.with_retry( + f"list_{resource_type}", + lambda p=page_params: client.raw_request( + "GET", f"/v1/{api_path}", params=p + ), + ) + + page_resources = response.get("objects", []) + if page_resources: + all_resources.extend(page_resources) + log.info( + f"Fetched {resource_type} page", + page_num=page_num, + page_size=len(page_resources), + total_fetched=len(all_resources), + project_id=project_id, + ) + starting_after = page_resources[-1].get("id") + if len(page_resources) < page_params["limit"]: + break + else: + break + + if project_id and client_side_filter_field: + all_resources = [ + r + for r in all_resources + if r.get(client_side_filter_field) == project_id + ] + + return all_resources + + except Exception as e: + log.error( + f"Failed to list {resource_type}", error=str(e), project_id=project_id + ) + raise + + @dataclass(slots=True) class MigrationResult: """Result of a resource migration operation.""" @@ -339,83 +410,14 @@ async def _list_resources_with_client( Returns: List of resources as dicts from raw API """ - try: - # Build base parameters - params = {} - if project_id and not client_side_filter_field: - # Use server-side filtering if no client-side field specified - params["project_id"] = project_id - if additional_params: - params.update(additional_params) - - # Convert resource_type to API path (e.g., 'datasets' -> 'dataset') - api_path = resource_type.rstrip("s") - - # Paginate through all results - all_resources = [] - starting_after = None - page_num = 0 - - while True: - page_num += 1 - page_params = {**params} - - # Set a large limit to minimize number of pages - page_params["limit"] = DEFAULT_LIST_PAGE_SIZE - - if starting_after is not None: - page_params["starting_after"] = starting_after - - # Make the API call using raw_request - # Use a default parameter to capture page_params at definition time - response = await client.with_retry( - f"list_{resource_type}", - lambda p=page_params: client.raw_request( - "GET", - f"/v1/{api_path}", - params=p, - ), - ) - - # Extract objects from response - page_resources = response.get("objects", []) - - if page_resources: - all_resources.extend(page_resources) - self._logger.info( - f"Fetched {resource_type} page", - page_num=page_num, - page_size=len(page_resources), - total_fetched=len(all_resources), - project_id=project_id, - ) - - # Get the last item's ID for pagination - last_item = page_resources[-1] - starting_after = last_item.get("id") - - # If we got fewer items than the limit, we're done - if len(page_resources) < page_params["limit"]: - break - else: - # No more results - break - - # Apply client-side filtering if needed (for dicts) - if project_id and client_side_filter_field: - all_resources = [ - resource - for resource in all_resources - if resource.get(client_side_filter_field) == project_id - ] - - return all_resources - - except Exception as e: - self._logger.error( - f"Failed to list {resource_type}", error=str(e), project_id=project_id - ) - raise + return await list_resources( + client, + resource_type, + project_id, + additional_params=additional_params, + client_side_filter_field=client_side_filter_field, + log=self._logger, + ) async def get_dependencies(self, resource: T) -> list[str]: """Get list of resource IDs that this resource depends on. diff --git a/braintrust_migrate/resources/functions.py b/braintrust_migrate/resources/functions.py index cc84eb0..9b0443c 100644 --- a/braintrust_migrate/resources/functions.py +++ b/braintrust_migrate/resources/functions.py @@ -49,7 +49,7 @@ async def migrate_batch( source_id = self.get_resource_id(resource) name = resource.get("name") slug = resource.get("slug") - self._logger.warning( + self._logger.info( "⏭️ Skipping bundled code function (its code bundle cannot be " "migrated; re-push it to the destination)", source_id=source_id, diff --git a/braintrust_migrate/validation.py b/braintrust_migrate/validation.py new file mode 100644 index 0000000..b4e8cee --- /dev/null +++ b/braintrust_migrate/validation.py @@ -0,0 +1,308 @@ +"""Post-migration validation: confirm what was migrated exists in the destination. + +Two check kinds, deliberately asymmetric (see the team discussion): + +- **Object parity** — for project-scoped object resources (datasets, experiments, + prompts, functions, project_scores, views, project_tags, span_iframes): list + source vs dest, compare, and report the *specific* items missing in the + destination. Cheap (small N, list endpoints). + +- **Event count-parity** — for high-volume event resources (dataset events, + experiment events, logs): compare a cheap ``count`` of source vs dest. We do + NOT enumerate *which* events are missing — that would require materializing and + diffing both full id sets (≈ re-reading the whole migration) and does not scale. + Counts are reported per parent object (per dataset / experiment / project) so a + discrepancy can still be localized to a specific object. + +Validation mirrors migration policy (e.g. it excludes bundle-backed code +functions, which are intentionally skipped) so deliberate skips are not reported +as missing. + +Org-scoped resources (roles, groups, ai_secrets) and ACLs are intentionally out +of scope for now. +""" + +from __future__ import annotations + +from collections.abc import Awaitable, Callable +from dataclasses import dataclass, field +from typing import Any + +from braintrust_migrate.btql import btql_quote +from braintrust_migrate.resources.functions import FunctionMigrator + +# Lists resources of a type from a client for a project: (client, type, project_id). +ListFn = Callable[[Any, str, str | None], Awaitable[list[dict[str, Any]]]] + + +@dataclass +class CheckResult: + kind: str # "object" | "events" + scope: str # e.g. "datasets" or "datasets:my-ds/events" + source_count: int + dest_count: int + ok: bool + missing: list[str] = field(default_factory=list) # object kind only + + def to_dict(self) -> dict[str, Any]: + out: dict[str, Any] = { + "kind": self.kind, + "scope": self.scope, + "source_count": self.source_count, + "dest_count": self.dest_count, + "ok": self.ok, + } + if self.missing: + out["missing"] = self.missing + return out + + +@dataclass +class ResourceValidation: + resource_type: str + ok: bool + checks: list[CheckResult] + + def to_dict(self) -> dict[str, Any]: + return { + "resource_type": self.resource_type, + "ok": self.ok, + "checks": [c.to_dict() for c in self.checks], + } + + +def _default_key(resource: dict[str, Any]) -> str | None: + """Stable match key: prefer slug (functions/prompts), else name.""" + return resource.get("slug") or resource.get("name") + + +def _dataset_from_expr(parent_id: str) -> str: + return f"dataset('{btql_quote(parent_id)}') spans" + + +def _experiment_from_expr(parent_id: str) -> str: + return f"experiment('{btql_quote(parent_id)}') spans" + + +def _logs_from_expr(parent_id: str) -> str: + return f"project_logs('{btql_quote(parent_id)}') spans" + + +@dataclass(frozen=True) +class ValidationSpec: + # Listable resource type for object parity (None for logs, which has no object layer). + object_type: str | None = None + # Builds a BTQL `from` expression for the event count, given a parent id. + event_from_expr: Callable[[str], str] | None = None + # True: events are counted per listed object; False: project-level (logs). + event_per_object: bool = False + key_fn: Callable[[dict[str, Any]], str | None] = _default_key + # Mirrors migration policy: keep only items the migrator would have migrated. + expected_filter: Callable[[dict[str, Any]], bool] | None = None + + +VALIDATION_SPECS: dict[str, ValidationSpec] = { + "datasets": ValidationSpec( + object_type="datasets", + event_from_expr=_dataset_from_expr, + event_per_object=True, + ), + "experiments": ValidationSpec( + object_type="experiments", + event_from_expr=_experiment_from_expr, + event_per_object=True, + ), + "logs": ValidationSpec( + object_type=None, + event_from_expr=_logs_from_expr, + event_per_object=False, + ), + "functions": ValidationSpec( + object_type="functions", + expected_filter=lambda r: not FunctionMigrator._is_code_bundle_function(r), + ), + "prompts": ValidationSpec(object_type="prompts"), + "project_scores": ValidationSpec(object_type="project_scores"), + "views": ValidationSpec(object_type="views"), + "project_tags": ValidationSpec(object_type="project_tags"), + "span_iframes": ValidationSpec(object_type="span_iframes"), +} + + +async def count_events( + client: Any, + from_expr: str, + *, + created_after: str | None = None, + created_before: str | None = None, +) -> int: + """Cheap event count via a single BTQL aggregation (no row scan).""" + conditions: list[str] = [] + if isinstance(created_after, str) and created_after: + conditions.append(f"created >= '{btql_quote(created_after)}'") + if isinstance(created_before, str) and created_before: + conditions.append(f"created < '{btql_quote(created_before)}'") + filter_clause = f"filter: {' and '.join(conditions)}\n" if conditions else "" + query = f"select: count(1) as n\nfrom: {from_expr}\n{filter_clause}" + + resp = await client.with_retry( + "validate_count_events", + lambda: client.raw_request("POST", "/btql", json={"query": query}), + ) + data = resp.get("data") if isinstance(resp, dict) else None + if isinstance(data, list) and data and isinstance(data[0], dict): + row = data[0] + value = row.get("n") + if value is None: + value = next( + (v for v in row.values() if isinstance(v, (int, float))), 0 + ) + return int(value) + return 0 + + +async def _validate_object_parity( + *, + source_client: Any, + dest_client: Any, + spec: ValidationSpec, + resource_name: str, + source_project_id: str, + dest_project_id: str, + list_fn: ListFn, +) -> tuple[CheckResult, list[dict[str, Any]], list[dict[str, Any]]]: + assert spec.object_type is not None + source = await list_fn(source_client, spec.object_type, source_project_id) + dest = await list_fn(dest_client, spec.object_type, dest_project_id) + + expected = ( + [r for r in source if spec.expected_filter(r)] + if spec.expected_filter + else list(source) + ) + expected_keys = {k for r in expected if (k := spec.key_fn(r))} + dest_keys = {k for r in dest if (k := spec.key_fn(r))} + + missing = sorted(expected_keys - dest_keys) + check = CheckResult( + kind="object", + scope=resource_name, + source_count=len(expected_keys), + dest_count=len(expected_keys & dest_keys), + ok=not missing, + missing=missing, + ) + return check, expected, dest + + +async def _validate_event_parity( + *, + scope: str, + source_client: Any, + dest_client: Any, + source_from_expr: str, + dest_from_expr: str, + created_after: str | None, + created_before: str | None, +) -> CheckResult: + source_count = await count_events( + source_client, + source_from_expr, + created_after=created_after, + created_before=created_before, + ) + dest_count = await count_events( + dest_client, + dest_from_expr, + created_after=created_after, + created_before=created_before, + ) + return CheckResult( + kind="events", + scope=scope, + source_count=source_count, + dest_count=dest_count, + ok=source_count == dest_count, + ) + + +async def validate_resource( + resource_name: str, + *, + source_client: Any, + dest_client: Any, + source_project_id: str, + dest_project_id: str, + list_fn: ListFn, + id_mapping: dict[str, str] | None = None, + created_after: str | None = None, + created_before: str | None = None, +) -> ResourceValidation | None: + """Validate one resource type for a project. Returns None if not in scope.""" + spec = VALIDATION_SPECS.get(resource_name) + if spec is None: + return None + + id_mapping = id_mapping or {} + checks: list[CheckResult] = [] + expected_objects: list[dict[str, Any]] = [] + dest_objects: list[dict[str, Any]] = [] + + if spec.object_type is not None: + obj_check, expected_objects, dest_objects = await _validate_object_parity( + source_client=source_client, + dest_client=dest_client, + spec=spec, + resource_name=resource_name, + source_project_id=source_project_id, + dest_project_id=dest_project_id, + list_fn=list_fn, + ) + checks.append(obj_check) + + if spec.event_from_expr is not None: + if spec.event_per_object: + dest_id_by_key = { + spec.key_fn(r): r.get("id") for r in dest_objects if r.get("id") + } + for src in expected_objects: + source_id = src.get("id") + key = spec.key_fn(src) + if not isinstance(source_id, str): + continue + # Prefer the exact migration mapping; fall back to name/slug match. + dest_id = id_mapping.get(source_id) or dest_id_by_key.get(key) + if not isinstance(dest_id, str): + # The object itself is missing; already reported by the object + # check. Don't emit a misleading event check. + continue + checks.append( + await _validate_event_parity( + scope=f"{resource_name}:{key}/events", + source_client=source_client, + dest_client=dest_client, + source_from_expr=spec.event_from_expr(source_id), + dest_from_expr=spec.event_from_expr(dest_id), + created_after=created_after, + created_before=created_before, + ) + ) + else: + # Project-level events (logs). + checks.append( + await _validate_event_parity( + scope=f"{resource_name}/events", + source_client=source_client, + dest_client=dest_client, + source_from_expr=spec.event_from_expr(source_project_id), + dest_from_expr=spec.event_from_expr(dest_project_id), + created_after=created_after, + created_before=created_before, + ) + ) + + return ResourceValidation( + resource_type=resource_name, + ok=all(c.ok for c in checks), + checks=checks, + ) diff --git a/tests/unit/test_validation.py b/tests/unit/test_validation.py new file mode 100644 index 0000000..2c57f65 --- /dev/null +++ b/tests/unit/test_validation.py @@ -0,0 +1,202 @@ +"""Unit tests for post-migration validation (object parity + event count-parity).""" + +from __future__ import annotations + +from typing import Any + +import pytest + +from braintrust_migrate.validation import count_events, validate_resource + + +class _FakeClient: + """Stub client: serves resource lists and BTQL count queries.""" + + def __init__( + self, + *, + lists: dict[str, list[dict[str, Any]]] | None = None, + counts: dict[str, int] | None = None, + ) -> None: + self._lists = lists or {} + self._counts = counts or {} # parent_id -> event count + + async def with_retry(self, _op, coro_func, **kwargs): + res = coro_func() + return await res if hasattr(res, "__await__") else res + + async def raw_request(self, method, path, *, json=None, params=None, **kwargs): + assert path == "/btql" + query = json["query"] + for parent_id, n in self._counts.items(): + if f"'{parent_id}'" in query: + return {"data": [{"n": n}]} + return {"data": [{"n": 0}]} + + +async def _list_fn(client, resource_type, project_id): + return client._lists.get(resource_type, []) + + +async def test_out_of_scope_resource_returns_none(): + src = _FakeClient() + dst = _FakeClient() + result = await validate_resource( + "roles", + source_client=src, + dest_client=dst, + source_project_id="sp", + dest_project_id="dp", + list_fn=_list_fn, + ) + assert result is None + + +async def test_object_parity_reports_specific_missing(): + src = _FakeClient( + lists={ + "prompts": [ + {"id": "s1", "name": "A"}, + {"id": "s2", "name": "B"}, + {"id": "s3", "name": "C"}, + ] + } + ) + dst = _FakeClient(lists={"prompts": [{"id": "d1", "name": "A"}, {"id": "d2", "name": "B"}]}) + + result = await validate_resource( + "prompts", + source_client=src, + dest_client=dst, + source_project_id="sp", + dest_project_id="dp", + list_fn=_list_fn, + ) + assert result is not None and result.ok is False + (check,) = result.checks + assert check.kind == "object" + assert check.source_count == 3 + assert check.dest_count == 2 + assert check.missing == ["C"] + + +async def test_function_bundle_excluded_from_expected(): + # Source has an inline fn (migratable) and a bundle fn (intentionally skipped). + src = _FakeClient( + lists={ + "functions": [ + { + "id": "s1", + "slug": "inline-fn", + "function_data": {"type": "code", "data": {"type": "inline", "code": "x"}}, + }, + { + "id": "s2", + "slug": "bundle-fn", + "function_data": {"type": "code", "data": {"type": "bundle"}}, + }, + ] + } + ) + # Dest only has the inline one — the bundle one is correctly absent. + dst = _FakeClient(lists={"functions": [{"id": "d1", "slug": "inline-fn"}]}) + + result = await validate_resource( + "functions", + source_client=src, + dest_client=dst, + source_project_id="sp", + dest_project_id="dp", + list_fn=_list_fn, + ) + assert result is not None and result.ok is True + (check,) = result.checks + assert check.source_count == 1 # only the inline fn is expected + assert check.missing == [] + + +async def test_dataset_object_and_event_parity_match(): + src = _FakeClient( + lists={"datasets": [{"id": "sd1", "name": "D1"}]}, counts={"sd1": 50} + ) + dst = _FakeClient( + lists={"datasets": [{"id": "dd1", "name": "D1"}]}, counts={"dd1": 50} + ) + + result = await validate_resource( + "datasets", + source_client=src, + dest_client=dst, + source_project_id="sp", + dest_project_id="dp", + list_fn=_list_fn, + id_mapping={"sd1": "dd1"}, + ) + assert result is not None and result.ok is True + kinds = {c.kind for c in result.checks} + assert kinds == {"object", "events"} + event_check = next(c for c in result.checks if c.kind == "events") + assert event_check.scope == "datasets:D1/events" + assert event_check.source_count == 50 and event_check.dest_count == 50 + + +async def test_dataset_event_count_mismatch_flags_failure(): + src = _FakeClient( + lists={"datasets": [{"id": "sd1", "name": "D1"}]}, counts={"sd1": 50} + ) + dst = _FakeClient( + lists={"datasets": [{"id": "dd1", "name": "D1"}]}, counts={"dd1": 48} + ) + + result = await validate_resource( + "datasets", + source_client=src, + dest_client=dst, + source_project_id="sp", + dest_project_id="dp", + list_fn=_list_fn, + id_mapping={"sd1": "dd1"}, + ) + assert result is not None and result.ok is False + event_check = next(c for c in result.checks if c.kind == "events") + assert event_check.source_count == 50 and event_check.dest_count == 48 + assert event_check.ok is False + + +async def test_logs_event_parity_project_level(): + src = _FakeClient(counts={"sp": 1000}) + dst = _FakeClient(counts={"dp": 1000}) + + result = await validate_resource( + "logs", + source_client=src, + dest_client=dst, + source_project_id="sp", + dest_project_id="dp", + list_fn=_list_fn, + ) + assert result is not None and result.ok is True + (check,) = result.checks + assert check.kind == "events" and check.scope == "logs/events" + assert check.source_count == 1000 and check.dest_count == 1000 + + +async def test_count_events_honors_created_window(): + captured: dict[str, str] = {} + + class _CaptureClient(_FakeClient): + async def raw_request(self, method, path, *, json=None, params=None, **kwargs): + captured["query"] = json["query"] + return {"data": [{"n": 7}]} + + client = _CaptureClient() + n = await count_events( + client, + "project_logs('p') spans", + created_after="2026-01-01T00:00:00Z", + created_before="2026-02-01T00:00:00Z", + ) + assert n == 7 + assert "created >= '2026-01-01T00:00:00Z'" in captured["query"] + assert "created < '2026-02-01T00:00:00Z'" in captured["query"] + assert "count(1)" in captured["query"] diff --git a/tests/unit/test_validation_wiring.py b/tests/unit/test_validation_wiring.py new file mode 100644 index 0000000..5c0bf74 --- /dev/null +++ b/tests/unit/test_validation_wiring.py @@ -0,0 +1,199 @@ +"""Wiring tests: orchestrator validation phase + console rendering.""" + +from __future__ import annotations + +from typing import Any + +import pytest +from pydantic import HttpUrl +from rich.console import Console + +import braintrust_migrate.cli as cli_module +from braintrust_migrate.config import BraintrustOrgConfig, Config, MigrationConfig +from braintrust_migrate.orchestration import MigrationOrchestrator + + +class _OrchStub: + """Serves /v1/ list pages and /btql counts via raw_request.""" + + def __init__( + self, + *, + lists: dict[str, list[dict[str, Any]]] | None = None, + counts: dict[str, int] | None = None, + ) -> None: + self._lists = lists or {} + self._counts = counts or {} + + async def with_retry(self, _op, coro_func, **kwargs): + res = coro_func() + return await res if hasattr(res, "__await__") else res + + async def raw_request(self, method, path, *, params=None, json=None, **kwargs): + if method == "GET" and path.startswith("/v1/"): + resource_type = path[len("/v1/") :] + "s" # "dataset" -> "datasets" + return {"objects": self._lists.get(resource_type, [])} + if method == "POST" and path == "/btql": + query = json["query"] + for parent_id, n in self._counts.items(): + if f"'{parent_id}'" in query: + return {"data": [{"n": n}]} + return {"data": [{"n": 0}]} + raise AssertionError(f"unexpected request: {method} {path}") + + +def _config() -> Config: + return Config( + source=BraintrustOrgConfig(api_key="s", url=HttpUrl("https://src.dev")), + destination=BraintrustOrgConfig(api_key="d", url=HttpUrl("https://dst.dev")), + migration=MigrationConfig(validate_migration=True), + ) + + +@pytest.mark.asyncio +async def test_validate_project_aggregates_object_and_event_checks(): + orch = MigrationOrchestrator(_config()) + source = _OrchStub( + lists={ + "datasets": [{"id": "sd1", "name": "D1"}], + "prompts": [{"id": "sp1", "name": "P1"}, {"id": "sp2", "name": "P2"}], + }, + counts={"sd1": 50, "sproj": 1000}, + ) + dest = _OrchStub( + lists={ + "datasets": [{"id": "dd1", "name": "D1"}], + "prompts": [{"id": "dp1", "name": "P1"}], # P2 missing + }, + counts={"dd1": 50, "dproj": 1000}, + ) + + result = await orch._validate_project( + source_client=source, + dest_client=dest, + source_project_id="sproj", + dest_project_id="dproj", + migrated_resource_types=["datasets", "prompts", "logs"], + id_mapping={"sd1": "dd1"}, + ) + + assert result["ok"] is False # prompts P2 missing + by_type = {r["resource_type"]: r for r in result["resources"]} + assert set(by_type) == {"datasets", "prompts", "logs"} + + assert by_type["datasets"]["ok"] is True # object + 50==50 events + assert by_type["logs"]["ok"] is True # 1000==1000 project events + + prompts = by_type["prompts"] + assert prompts["ok"] is False + obj_check = next(c for c in prompts["checks"] if c["kind"] == "object") + assert obj_check["missing"] == ["P2"] + + +def test_display_results_renders_validation_failures(monkeypatch): + rec = Console(record=True, width=200) + monkeypatch.setattr(cli_module, "console", rec) + + results = { + "summary": { + "total_projects": 1, + "total_resources": 3, + "migrated_resources": 3, + "skipped_resources": 0, + "failed_resources": 0, + "errors": [], + }, + "projects": { + "ProjA": { + "total_resources": 3, + "migrated_resources": 3, + "skipped_resources": 0, + "failed_resources": 0, + "project_id": "p", + "validation": { + "ok": False, + "resources": [ + {"resource_type": "datasets", "ok": True, "checks": []}, + { + "resource_type": "prompts", + "ok": False, + "checks": [ + { + "kind": "object", + "scope": "prompts", + "source_count": 2, + "dest_count": 1, + "ok": False, + "missing": ["P2"], + } + ], + }, + { + "resource_type": "logs", + "ok": False, + "checks": [ + { + "kind": "events", + "scope": "logs/events", + "source_count": 1000, + "dest_count": 990, + "ok": False, + } + ], + }, + ], + }, + } + }, + } + + cli_module._display_results(results) + out = rec.export_text() + + assert "Validation" in out + assert "2 of 3 validated" in out # 2 failing resources of 3 + assert "ProjA/prompts" in out + assert "missing 1: P2" in out + assert "logs/events: source=1000 dest=990" in out + + +def test_display_results_validation_all_ok(monkeypatch): + rec = Console(record=True, width=200) + monkeypatch.setattr(cli_module, "console", rec) + results = { + "summary": { + "total_projects": 1, + "total_resources": 1, + "migrated_resources": 1, + "skipped_resources": 0, + "failed_resources": 0, + "errors": [], + }, + "projects": { + "ProjA": { + "total_resources": 1, + "migrated_resources": 1, + "skipped_resources": 0, + "failed_resources": 0, + "project_id": "p", + "validation": { + "ok": True, + "resources": [ + {"resource_type": "datasets", "ok": True, "checks": []} + ], + }, + } + }, + } + cli_module._display_results(results) + out = rec.export_text() + assert "all 1 validated resource(s) match" in out + + +def test_config_from_env_validate_flag(monkeypatch): + monkeypatch.setenv("BT_SOURCE_API_KEY", "s") + monkeypatch.setenv("BT_DEST_API_KEY", "d") + monkeypatch.setenv("MIGRATION_VALIDATE", "true") + assert Config.from_env().migration.validate_migration is True + monkeypatch.setenv("MIGRATION_VALIDATE", "false") + assert Config.from_env().migration.validate_migration is False