diff --git a/src/unstract/clone/client.py b/src/unstract/clone/client.py index 3f79499..fb60e97 100644 --- a/src/unstract/clone/client.py +++ b/src/unstract/clone/client.py @@ -70,7 +70,32 @@ def _request( files: dict[str, Any] | None = None, data: dict[str, Any] | None = None, ) -> Any: - url = self._url(path) + return self._send( + method, + self._url(path), + path, + params=params, + json=json, + files=files, + data=data, + ) + + def _send( + self, + method: str, + url: str, + label: str, + *, + params: dict[str, Any] | None = None, + json: Any = None, + files: dict[str, Any] | None = None, + data: dict[str, Any] | None = None, + ) -> Any: + """Issue a request against an already-built URL. + + Split out from ``_request`` so pagination can follow the absolute + ``next`` links DRF returns, which are not org-path-relative. + """ # Redact secrets from logs: only entity path + method, never body. logger.debug("%s %s", method, url) resp = self._session.request( @@ -85,7 +110,7 @@ def _request( ) if not 200 <= resp.status_code < 300: raise PlatformAPIError( - f"{method} {path} returned {resp.status_code}", + f"{method} {label} returned {resp.status_code}", status_code=resp.status_code, body=resp.text[:2000], ) @@ -93,6 +118,41 @@ def _request( return None return resp.json() + def _paginate( + self, path: str, params: dict[str, Any] | None = None + ) -> list[dict[str, Any]]: + """GET a list endpoint, following DRF pagination to exhaustion. + + A bare list is returned unchanged, so the same client works against + deployments where the endpoint is not paginated. A short read is + raised rather than returned: a clone that silently copies the first + page is worse than one that fails. + """ + result = self._request("GET", path, params=dict(params or {})) + if isinstance(result, list): + return result + if not isinstance(result, dict) or "results" not in result: + raise PlatformAPIError(f"GET {path} returned an unrecognised list payload") + + rows: list[dict[str, Any]] = [] + expected = result.get("count") + seen: set[str] = set() + while True: + rows.extend(result.get("results") or []) + next_url = result.get("next") + if not next_url: + break + if next_url in seen: + raise PlatformAPIError(f"GET {path} pagination looped at {next_url}") + seen.add(next_url) + result = self._send("GET", next_url, path) + + if expected is not None and len(rows) != expected: + raise PlatformAPIError( + f"GET {path} returned {len(rows)} rows but reported count={expected}" + ) + return rows + def get_post_schema(self, entity_path: str) -> frozenset[str]: """Return the set of fields the backend's POST accepts. @@ -137,8 +197,7 @@ def list_users(self) -> list[dict[str, Any]]: def list_groups(self) -> list[dict[str, Any]]: """List org groups; no server-side name filter — callers match in memory.""" - result = self._request("GET", "groups/") - return result if isinstance(result, list) else (result or {}).get("results", []) + return self._paginate("groups/") def create_group(self, payload: dict[str, Any]) -> dict[str, Any]: """Create a group; response has no ``id`` — re-list to learn the pk.""" @@ -146,8 +205,7 @@ def create_group(self, payload: dict[str, Any]) -> dict[str, Any]: def list_group_members(self, group_id: Any) -> list[dict[str, Any]]: """List a group's member rows (each carries ``email``).""" - result = self._request("GET", f"groups/{group_id}/members/") - return result if isinstance(result, list) else result.get("results", []) + return self._paginate(f"groups/{group_id}/members/") def add_group_members(self, group_id: Any, user_ids: list[int]) -> Any: """Bulk-add members by user pk; idempotent server-side.""" @@ -175,9 +233,7 @@ def list_adapters( params["adapter_name"] = name if adapter_type is not None: params["adapter_type"] = adapter_type - result = self._request("GET", "adapter/", params=params) - # DRF ModelViewSet.list returns a bare list (no pagination on this endpoint). - return result if isinstance(result, list) else result.get("results", []) + return self._paginate("adapter/", params) def get_adapter(self, adapter_pk: str) -> dict[str, Any]: return self._request("GET", f"adapter/{adapter_pk}/") @@ -199,8 +255,7 @@ def list_connectors( params["connector_name"] = name if connector_type is not None: params["connector_type"] = connector_type - result = self._request("GET", "connector/", params=params) - return result if isinstance(result, list) else result.get("results", []) + return self._paginate("connector/", params) def get_connector(self, connector_pk: str) -> dict[str, Any]: return self._request("GET", f"connector/{connector_pk}/") @@ -215,9 +270,7 @@ def list_tags(self, *, name: str | None = None) -> list[dict[str, Any]]: params: dict[str, Any] = {} if name is not None: params["name"] = name - result = self._request("GET", "tags/", params=params) - # Tags endpoint uses pagination — accept either bare list or paginated envelope. - return result if isinstance(result, list) else result.get("results", []) + return self._paginate("tags/", params) def create_tag(self, payload: dict[str, Any]) -> dict[str, Any]: return self._request("POST", "tags/", json=payload) @@ -226,8 +279,7 @@ def create_tag(self, payload: dict[str, Any]) -> dict[str, Any]: def list_custom_tools(self) -> list[dict[str, Any]]: """List all prompt-studio projects in this org. No name filter.""" - result = self._request("GET", "prompt-studio/") - return result if isinstance(result, list) else result.get("results", []) + return self._paginate("prompt-studio/") def get_custom_tool(self, tool_id: str) -> dict[str, Any]: """Fetch a single prompt-studio project. @@ -249,8 +301,7 @@ def list_profiles(self, tool_id: str) -> list[dict[str, Any]]: default profile's adapter UUIDs so they can be remapped to target adapter ids for ``import_project``. """ - result = self._request("GET", f"prompt-studio/prompt-studio-profile/{tool_id}/") - return result if isinstance(result, list) else result.get("results", []) + return self._paginate(f"prompt-studio/prompt-studio-profile/{tool_id}/") def list_prompts(self, tool_id: str) -> list[dict[str, Any]]: """List a tool's prompts (``prompt_id`` + ``prompt_key`` per row). @@ -259,10 +310,7 @@ def list_prompts(self, tool_id: str) -> list[dict[str, Any]]: ``import_project`` / ``sync_prompts`` (matched by ``prompt_key``), so prompt-scoped cloud config can remap its FKs. """ - result = self._request( - "GET", "prompt-studio/prompt/", params={"tool_id": tool_id} - ) - return result if isinstance(result, list) else result.get("results", []) + return self._paginate("prompt-studio/prompt/", {"tool_id": tool_id}) def export_project(self, tool_id: str) -> dict[str, Any]: """Export a prompt-studio project as a portable JSON blob. @@ -338,10 +386,7 @@ def list_prompt_documents(self, tool_id: str) -> list[dict[str, Any]]: enumeration. Response items carry ``document_id``, ``document_name``, and ``tool``. """ - result = self._request( - "GET", "prompt-studio/prompt-document/", params={"tool_id": tool_id} - ) - return result if isinstance(result, list) else result.get("results", []) + return self._paginate("prompt-studio/prompt-document/", {"tool_id": tool_id}) def download_prompt_file(self, tool_id: str, document_id: str) -> dict[str, Any]: """GET a Prompt Studio document by tool + document id. @@ -398,8 +443,7 @@ def list_workflows(self, *, name: str | None = None) -> list[dict[str, Any]]: params: dict[str, Any] = {} if name is not None: params["workflow_name"] = name - result = self._request("GET", "workflow/", params=params) - return result if isinstance(result, list) else result.get("results", []) + return self._paginate("workflow/", params) def get_workflow(self, workflow_id: str) -> dict[str, Any]: return self._request("GET", f"workflow/{workflow_id}/") @@ -420,8 +464,7 @@ def list_registries( params: dict[str, Any] = {} if custom_tool is not None: params["custom_tool"] = custom_tool - result = self._request("GET", "prompt-studio/registry/", params=params) - return result if isinstance(result, list) else result.get("results", []) + return self._paginate("prompt-studio/registry/", params) # ----- tool instances ----- @@ -432,8 +475,7 @@ def list_tool_instances( params: dict[str, Any] = {} if workflow_id is not None: params["workflow"] = workflow_id - result = self._request("GET", "tool_instance/", params=params) - return result if isinstance(result, list) else result.get("results", []) + return self._paginate("tool_instance/", params) def create_tool_instance(self, payload: dict[str, Any]) -> dict[str, Any]: """Create a tool instance (max 1 per workflow). The created row comes @@ -466,8 +508,7 @@ def list_workflow_endpoints( params: dict[str, Any] = {} if workflow_id is not None: params["workflow"] = workflow_id - result = self._request("GET", "workflow/endpoint/", params=params) - return result if isinstance(result, list) else result.get("results", []) + return self._paginate("workflow/endpoint/", params) def update_workflow_endpoint( self, endpoint_id: str, payload: dict[str, Any] @@ -490,8 +531,7 @@ def list_pipelines( params["pipeline_name"] = name if pipeline_type is not None: params["type"] = pipeline_type - result = self._request("GET", "pipeline/", params=params) - return result if isinstance(result, list) else result.get("results", []) + return self._paginate("pipeline/", params) def get_pipeline(self, pipeline_id: str) -> dict[str, Any]: return self._request("GET", f"pipeline/{pipeline_id}/") @@ -518,8 +558,7 @@ def list_api_deployments( params: dict[str, Any] = {} if api_name is not None: params["api_name"] = api_name - result = self._request("GET", "api/deployment/", params=params) - return result if isinstance(result, list) else result.get("results", []) + return self._paginate("api/deployment/", params) def get_api_deployment(self, deployment_id: str) -> dict[str, Any]: return self._request("GET", f"api/deployment/{deployment_id}/") @@ -539,13 +578,11 @@ def update_api_deployment( def list_pipeline_keys(self, pipeline_id: str) -> list[dict[str, Any]]: """List API keys belonging to a pipeline.""" - result = self._request("GET", f"api/keys/pipeline/{pipeline_id}/") - return result if isinstance(result, list) else result.get("results", []) + return self._paginate(f"api/keys/pipeline/{pipeline_id}/") def list_api_deployment_keys(self, deployment_id: str) -> list[dict[str, Any]]: """List API keys belonging to an API deployment.""" - result = self._request("GET", f"api/keys/api/{deployment_id}/") - return result if isinstance(result, list) else result.get("results", []) + return self._paginate(f"api/keys/api/{deployment_id}/") def create_api_key(self, payload: dict[str, Any]) -> dict[str, Any]: """Create an extra API key tied to a pipeline or deployment. @@ -560,8 +597,7 @@ def create_api_key(self, payload: dict[str, Any]) -> dict[str, Any]: def list_lookup_definitions(self) -> list[dict[str, Any]]: """List lookup definitions in this org. Also the capability-probe path.""" - result = self._request("GET", "lookups/definitions/") - return result if isinstance(result, list) else (result or {}).get("results", []) + return self._paginate("lookups/definitions/") def get_lookup_definition(self, lookup_id: str) -> dict[str, Any]: """Fetch a lookup definition's detail. @@ -607,8 +643,7 @@ def list_lookup_files(self, lookup_id: str) -> list[dict[str, Any]]: """List a lookup's draft reference files (rows carry ``file_id``, ``file_name``, ``file_size``). """ - result = self._request("GET", f"lookups/definitions/{lookup_id}/files/") - return result if isinstance(result, list) else (result or {}).get("results", []) + return self._paginate(f"lookups/definitions/{lookup_id}/files/") def download_lookup_file(self, lookup_id: str, file_id: str) -> bytes: """Download a reference file's original bytes. @@ -648,8 +683,7 @@ def list_lookup_assignments(self) -> list[dict[str, Any]]: ``version`` (source lookup-version uuid), ``lookup_definition`` (source lookup_id), ``is_draft_version``, and ``variable_mappings``. """ - result = self._request("GET", "lookups/assignments/") - return result if isinstance(result, list) else (result or {}).get("results", []) + return self._paginate("lookups/assignments/") def create_lookup_assignment(self, payload: dict[str, Any]) -> dict[str, Any]: """Create a prompt-lookup assignment. @@ -668,9 +702,7 @@ def update_lookup_share( ``payload`` carries ``shared_to_org`` + ``shared_users`` (target user pks). Lookups expose no group-sharing axis, so no ``shared_groups``. """ - return self._request( - "PATCH", f"lookups/definitions/{lookup_id}/", json=payload - ) + return self._request("PATCH", f"lookups/definitions/{lookup_id}/", json=payload) def list_lookup_versions(self, lookup_id: str) -> list[dict[str, Any]]: """List a lookup's versions (draft + published). @@ -678,17 +710,13 @@ def list_lookup_versions(self, lookup_id: str) -> list[dict[str, Any]]: Rows carry ``version_id``, ``is_draft``, ``version_number``, ``version_name``; the detail (``get_lookup_version``) inlines content. """ - result = self._request( - "GET", f"lookups/definitions/{lookup_id}/versions/" - ) + result = self._request("GET", f"lookups/definitions/{lookup_id}/versions/") if isinstance(result, list): return result # This endpoint wraps rows as {"versions": [...], "next_version_number"}. return (result or {}).get("versions", (result or {}).get("results", [])) - def get_lookup_version( - self, lookup_id: str, version_id: str - ) -> dict[str, Any]: + def get_lookup_version(self, lookup_id: str, version_id: str) -> dict[str, Any]: """Fetch a version's detail (``prompt_template``, adapters, files).""" return self._request( "GET", f"lookups/definitions/{lookup_id}/versions/{version_id}/" @@ -796,12 +824,9 @@ def list_auto_approval_settings(self) -> list[dict[str, Any]]: Returns 200 bare with no query params, so it doubles as the manual-review capability probe path. """ - result = self._request("GET", "manual_review/auto_approval_settings/") - return result if isinstance(result, list) else (result or {}).get("results", []) + return self._paginate("manual_review/auto_approval_settings/") - def create_auto_approval_settings( - self, payload: dict[str, Any] - ) -> dict[str, Any]: + def create_auto_approval_settings(self, payload: dict[str, Any]) -> dict[str, Any]: """Create org-level auto-approval settings. Writable: ``auto_approved_document_classes``, ``auto_approved_users``. @@ -813,8 +838,7 @@ def create_auto_approval_settings( def list_review_api_keys(self) -> list[dict[str, Any]]: """List review API keys in this org.""" - result = self._request("GET", "manual_review/api/keys/") - return result if isinstance(result, list) else (result or {}).get("results", []) + return self._paginate("manual_review/api/keys/") def create_review_api_key(self, payload: dict[str, Any]) -> dict[str, Any]: """Create a review API key. The ``api_key`` secret is server-minted @@ -833,8 +857,7 @@ def list_agentic_projects(self) -> list[dict[str, Any]]: ``lightweight_llm_connector_id`` / ``text_extractor_connector_id``), and ``canary_fields``. """ - result = self._request("GET", "agentic/projects/") - return result if isinstance(result, list) else (result or {}).get("results", []) + return self._paginate("agentic/projects/") def create_agentic_project(self, payload: dict[str, Any]) -> dict[str, Any]: """Create an agentic project. Returns the created row (carries ``id``).""" @@ -851,8 +874,7 @@ def list_agentic_prompt_versions( params: dict[str, Any] = {} if project_id is not None: params["project_id"] = project_id - result = self._request("GET", "agentic/prompt-versions/", params=params) - return result if isinstance(result, list) else (result or {}).get("results", []) + return self._paginate("agentic/prompt-versions/", params) def create_agentic_prompt_version(self, payload: dict[str, Any]) -> dict[str, Any]: """Create an agentic prompt version (flat endpoint, ``project`` in body).""" @@ -869,8 +891,7 @@ def list_agentic_schemas( params: dict[str, Any] = {} if project_id is not None: params["project_id"] = project_id - result = self._request("GET", "agentic/schemas/", params=params) - return result if isinstance(result, list) else (result or {}).get("results", []) + return self._paginate("agentic/schemas/", params) def create_agentic_schema(self, payload: dict[str, Any]) -> dict[str, Any]: """Create an agentic schema (flat endpoint, ``project`` in body).""" @@ -878,8 +899,7 @@ def create_agentic_schema(self, payload: dict[str, Any]) -> dict[str, Any]: def list_agentic_settings(self) -> list[dict[str, Any]]: """List agentic settings. Org-wide key/value rows (no project FK).""" - result = self._request("GET", "agentic/settings/") - return result if isinstance(result, list) else (result or {}).get("results", []) + return self._paginate("agentic/settings/") def create_agentic_setting(self, payload: dict[str, Any]) -> dict[str, Any]: """Create an org-wide agentic setting.""" @@ -919,18 +939,14 @@ def list_agentic_registries( params: dict[str, Any] = {} if agentic_project is not None: params["agentic_project"] = agentic_project - result = self._request("GET", "agentic-studio-registry/", params=params) - return result if isinstance(result, list) else (result or {}).get("results", []) + return self._paginate("agentic-studio-registry/", params) def list_agentic_documents(self, project_id: str) -> list[dict[str, Any]]: """List a project's uploaded documents. Rows carry ``id`` and ``original_filename``. Agentic docs are a store of their own, distinct from Prompt Studio ``prompt-document`` rows. """ - result = self._request( - "GET", "agentic/documents/", params={"project_id": project_id} - ) - return result if isinstance(result, list) else (result or {}).get("results", []) + return self._paginate("agentic/documents/", {"project_id": project_id}) def download_agentic_document(self, document_id: str) -> bytes: """Download an agentic document's original bytes. @@ -968,10 +984,7 @@ def list_agentic_verified_data(self, project_id: str) -> list[dict[str, Any]]: """List a project's verified (ground-truth) data rows. Each carries ``document_name``, ``document``, and ``data`` (the curated JSON). """ - result = self._request( - "GET", "agentic/verified-data/", params={"project_id": project_id} - ) - return result if isinstance(result, list) else (result or {}).get("results", []) + return self._paginate("agentic/verified-data/", {"project_id": project_id}) def create_agentic_verified_data(self, payload: dict[str, Any]) -> dict[str, Any]: """Create a verified-data row (``project``, ``document``, ``data``). diff --git a/tests/clone/test_client.py b/tests/clone/test_client.py index 561c91a..c11f12f 100644 --- a/tests/clone/test_client.py +++ b/tests/clone/test_client.py @@ -7,6 +7,7 @@ - 204 / empty body returns ``None`` instead of raising on .json(). - ``get_post_schema`` parses DRF ``actions.POST`` and caches per path. - ``close()`` shuts the underlying session; context manager works. +- ``_paginate`` follows ``next`` to exhaustion and refuses to return a short read. """ from __future__ import annotations @@ -157,3 +158,43 @@ def test_get_review_settings_reraises_non_500(): with pytest.raises(PlatformAPIError) as exc_info: client.get_review_settings("wf-1") assert exc_info.value.status_code == 403 + + +def _client_with_pages(*payloads) -> tuple[PlatformClient, MagicMock]: + """Client whose session returns each payload in turn, one per request.""" + client = PlatformClient(_endpoint()) + mock_request = MagicMock( + side_effect=[_fake_response(200, p) for p in payloads], + ) + client._session.request = mock_request + return client, mock_request + + +def test_paginate_follows_next_across_pages(): + page1 = { + "count": 3, + "next": "https://api.example.com/next?page=2", + "results": [1, 2], + } + page2 = {"count": 3, "next": None, "results": [3]} + client, mock_request = _client_with_pages(page1, page2) + + assert client.list_tags() == [1, 2, 3] + # Second hop must GET the absolute ``next`` URL verbatim, not an org path. + assert mock_request.call_args.args[1] == "https://api.example.com/next?page=2" + + +def test_paginate_raises_on_short_read(): + # A page set that doesn't add up means rows were dropped; a clone that + # silently copies a subset is worse than one that fails. + truncated = {"count": 9, "next": None, "results": [1, 2]} + client, _ = _client_with_pages(truncated) + with pytest.raises(PlatformAPIError, match="count=9"): + client.list_tags() + + +def test_paginate_raises_on_cyclic_next(): + looping = {"count": 2, "next": "https://api.example.com/loop", "results": [1]} + client, _ = _client_with_pages(looping, looping, looping) + with pytest.raises(PlatformAPIError, match="looped"): + client.list_tags()