diff --git a/CHANGELOG.md b/CHANGELOG.md index e4d1777..af3fc40 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -18,6 +18,12 @@ The format is based on Keep a Changelog and this project follows Semantic Versio - N/A +## [0.4.1] - 2026-06-08 + +### Fixed + +- Resolve `org_id` for attachment operations from a project (`GET /v1/project`, captured opportunistically whenever the client lists or creates a project) instead of `GET /ping`, which 404s on the SaaS REST API. Fixes a 0.4.0 regression where oversize-field attachment spilling failed on SaaS with "Unable to determine org_id for attachment operations" (also repairs the `copy_attachments` and organization-scoped ACL paths, which shared the same dependency). Verified end-to-end against SaaS. + ## [0.4.0] - 2026-06-08 ### Added diff --git a/README.md b/README.md index 0c0206f..dec638e 100644 --- a/README.md +++ b/README.md @@ -55,7 +55,7 @@ This tool provides migration capabilities for Braintrust organizations, handling This project is not yet published on PyPI. Install a pinned release from GitHub: ```bash -pip install "git+https://github.com/braintrustdata/braintrust-migrate.git@v0.4.0" +pip install "git+https://github.com/braintrustdata/braintrust-migrate.git@v0.4.1" ``` > Avoid installing directly from `main` in production workflows. diff --git a/braintrust_migrate/__init__.py b/braintrust_migrate/__init__.py index dd8059f..06dd50d 100644 --- a/braintrust_migrate/__init__.py +++ b/braintrust_migrate/__init__.py @@ -3,7 +3,7 @@ A Python CLI & library for migrating Braintrust organizations with maximum fidelity. """ -__version__ = "0.4.0" +__version__ = "0.4.1" __author__ = "Braintrust Migration Tool" __email__ = "support@braintrust.dev" diff --git a/braintrust_migrate/client.py b/braintrust_migrate/client.py index fd55aff..a2f962f 100644 --- a/braintrust_migrate/client.py +++ b/braintrust_migrate/client.py @@ -174,6 +174,7 @@ async def list_projects( for obj in objs: if isinstance(obj, dict): batch.append(obj) + self._maybe_capture_org_id(obj) projects.extend(batch) @@ -202,6 +203,7 @@ async def create_project( resp = await self.raw_request("POST", "/v1/project", json=payload) if not isinstance(resp, dict): raise BraintrustAPIError(f"Unexpected create project response: {type(resp)}") + self._maybe_capture_org_id(resp) return resp async def raw_request( @@ -326,6 +328,19 @@ def _origin(u: httpx.URL) -> tuple[str, str, int | None]: resp.raise_for_status() return resp.json() + def _maybe_capture_org_id(self, obj: Any) -> None: + """Opportunistically cache org_id from any object that carries it. + + Project (and most resource) objects include the org_id of the org they + belong to, so once the client has listed or created a project we already + know the org_id without a dedicated request. + """ + if self._org_id or not isinstance(obj, dict): + return + v = obj.get("org_id") + if isinstance(v, str) and v: + self._org_id = v + async def get_org_id(self) -> str: """Best-effort: get the org_id for this API key. @@ -334,28 +349,30 @@ async def get_org_id(self) -> str: if self._org_id: return self._org_id - # The Braintrust SDK uses GET /ping (no /v1) on the api_url. - candidates = ["/ping", "/v1/ping"] - last_err: Exception | None = None - for path in candidates: - try: - resp = await self.with_retry( - "ping", lambda p=path: self.raw_request("GET", p) - ) - if isinstance(resp, dict): - for key in ("org_id", "orgId"): - v = resp.get(key) - if isinstance(v, str) and v: - self._org_id = v - return v - except Exception as e: - last_err = e - continue + # org_id is reliably present on every Project object (and is captured + # opportunistically whenever the client lists or creates a project), so + # if it hasn't been seen yet, fetch one project and read it from there. + # Every project visible to an org-scoped API key is in that org. + try: + resp = await self.with_retry( + "get_org_id_via_project", + lambda: self.raw_request("GET", "/v1/project", params={"limit": 1}), + ) + except Exception as e: + raise BraintrustAPIError( + f"Unable to determine org_id for attachment operations: {e}" + ) from e + + objects = resp.get("objects") if isinstance(resp, dict) else None + if isinstance(objects, list): + for obj in objects: + self._maybe_capture_org_id(obj) + if self._org_id: + return self._org_id raise BraintrustAPIError( - "Unable to determine org_id for attachment operations. " - "Tried /ping and /v1/ping; last error: " - f"{last_err}" + "Unable to determine org_id for attachment operations: " + "no project with an org_id was found." ) async def health_check(self) -> dict[str, Any]: diff --git a/pyproject.toml b/pyproject.toml index f3da7c3..1970b8b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "braintrust-migrate" -version = "0.4.0" +version = "0.4.1" description = "CLI & library for migrating Braintrust organizations with maximum fidelity" readme = "README.md" requires-python = ">=3.12" diff --git a/tests/unit/test_attachment_copier.py b/tests/unit/test_attachment_copier.py index e239483..3ba2f54 100644 --- a/tests/unit/test_attachment_copier.py +++ b/tests/unit/test_attachment_copier.py @@ -27,8 +27,10 @@ async def test_attachment_copier_rewrites_braintrust_attachment_reference_and_do def source_handler(request: httpx.Request) -> httpx.Response: nonlocal src_attachment_meta_calls, src_download_calls - if str(request.url) == f"{src_base}/ping": - return httpx.Response(200, json={"org_id": src_org_id}) + if request.url.path == "/v1/project": + return httpx.Response( + 200, json={"objects": [{"id": "p1", "org_id": src_org_id}]} + ) if str(request.url).startswith(f"{src_base}/attachment"): src_attachment_meta_calls += 1 assert request.method == "GET" @@ -48,8 +50,10 @@ def source_handler(request: httpx.Request) -> httpx.Response: def dest_handler(request: httpx.Request) -> httpx.Response: nonlocal dst_upload_meta_calls, dst_upload_put_calls, dst_status_calls - if str(request.url) == f"{dst_base}/ping": - return httpx.Response(200, json={"org_id": dst_org_id}) + if request.url.path == "/v1/project": + return httpx.Response( + 200, json={"objects": [{"id": "p1", "org_id": dst_org_id}]} + ) if str(request.url) == f"{dst_base}/attachment": dst_upload_meta_calls += 1 assert request.method == "POST" diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index db3d085..113b433 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -122,3 +122,65 @@ async def mock_operation(): result = await client.with_retry("test_operation", mock_operation) assert result == "success" + + +async def test_get_org_id_queries_project_when_not_cached( + org_config, migration_config +): + """When org_id hasn't been seen yet, fetch one project and read it.""" + client = BraintrustClient(org_config, migration_config, "test-org") + + async def _raw(method, path, **kwargs): + if method == "GET" and path == "/v1/project": + return {"objects": [{"id": "p1", "org_id": "org-xyz"}]} + raise AssertionError(f"unexpected request: {method} {path}") + + with patch.object(client, "raw_request", new=AsyncMock(side_effect=_raw)): + assert await client.get_org_id() == "org-xyz" + # Cached on subsequent calls (no further requests). + assert await client.get_org_id() == "org-xyz" + + +async def test_get_org_id_raises_when_no_project(org_config, migration_config): + """If no project (with an org_id) is found, raise a clear error.""" + from braintrust_migrate.client import BraintrustAPIError + + client = BraintrustClient(org_config, migration_config, "test-org") + + async def _raw(method, path, **kwargs): + if method == "GET" and path == "/v1/project": + return {"objects": []} + raise AssertionError(f"unexpected request: {method} {path}") + + with patch.object(client, "raw_request", new=AsyncMock(side_effect=_raw)): + with pytest.raises(BraintrustAPIError, match="Unable to determine org_id"): + await client.get_org_id() + + +async def test_get_org_id_captured_from_list_projects(org_config, migration_config): + """Listing a project caches org_id, so get_org_id needs no extra request.""" + client = BraintrustClient(org_config, migration_config, "test-org") + + async def _raw(method, path, **kwargs): + if method == "GET" and path == "/v1/project": + return {"objects": [{"id": "p1", "org_id": "org-from-list"}]} + raise AssertionError(f"unexpected request: {method} {path}") + + with patch.object(client, "raw_request", new=AsyncMock(side_effect=_raw)): + await client.list_projects(limit=1) + # No /ping or extra project call — org_id was captured during listing. + assert await client.get_org_id() == "org-from-list" + + +async def test_get_org_id_captured_from_create_project(org_config, migration_config): + """Creating a project caches org_id.""" + client = BraintrustClient(org_config, migration_config, "test-org") + + async def _raw(method, path, **kwargs): + if method == "POST" and path == "/v1/project": + return {"id": "p1", "org_id": "org-from-create", "name": "x"} + raise AssertionError(f"unexpected request: {method} {path}") + + with patch.object(client, "raw_request", new=AsyncMock(side_effect=_raw)): + await client.create_project(name="x") + assert await client.get_org_id() == "org-from-create" diff --git a/tests/unit/test_oversize_field_spill.py b/tests/unit/test_oversize_field_spill.py index 3adc278..da21e08 100644 --- a/tests/unit/test_oversize_field_spill.py +++ b/tests/unit/test_oversize_field_spill.py @@ -20,8 +20,11 @@ def _make_dest_client(): def dest_handler(request: httpx.Request) -> httpx.Response: url = str(request.url) - if url == f"{dst_base}/ping": - return httpx.Response(200, json={"org_id": dst_org_id}) + if request.url.path == "/v1/project": + # get_org_id() reads org_id off a project listing. + return httpx.Response( + 200, json={"objects": [{"id": "p1", "org_id": dst_org_id}]} + ) if url == f"{dst_base}/attachment": calls["meta"] += 1 assert request.method == "POST" diff --git a/uv.lock b/uv.lock index 739e1f5..9234883 100644 --- a/uv.lock +++ b/uv.lock @@ -59,7 +59,7 @@ wheels = [ [[package]] name = "braintrust-migrate" -version = "0.4.0" +version = "0.4.1" source = { editable = "." } dependencies = [ { name = "braintrust" },