diff --git a/README.md b/README.md index 9246681..e098501 100644 --- a/README.md +++ b/README.md @@ -70,6 +70,9 @@ export SSL_CERT_FILE=/etc/pki/tls/certs/ca-bundle.crt Optional Configuration: ```bash +# Trustify API version (defaults to v3; set to v2 for legacy deployments) +export TRUSTIFY_API_VERSION=v3 + # Set custom configuration directory (defaults to ~/.config/trustshell/) export TRUSTSHELL_SCRATCH="/path/to/custom/config/dir" diff --git a/src/trustshell/__init__.py b/src/trustshell/__init__.py index d82cbaf..54e9e4e 100644 --- a/src/trustshell/__init__.py +++ b/src/trustshell/__init__.py @@ -38,7 +38,8 @@ LOCAL_AUTH_SERVER_PORT = os.getenv("LOCAL_AUTH_SERVER_PORT", "") -TRUSTIFY_URL_PATH = "/api/v2/" +TRUSTIFY_API_VERSION = os.getenv("TRUSTIFY_API_VERSION", "v3") +TRUSTIFY_URL_PATH = f"/api/{TRUSTIFY_API_VERSION}/" if "TRUSTIFY_URL" in os.environ: url_env = os.getenv("TRUSTIFY_URL", "") parsed_url = urlparse(url_env) @@ -48,10 +49,9 @@ ) else: TRUSTIFY_URL = url_env - # Only enable authentication if AUTH_ENDPOINT is also set AUTH_ENABLED = bool(os.getenv("AUTH_ENDPOINT")) else: - TRUSTIFY_URL = "http://localhost:8080/api/v2/" + TRUSTIFY_URL = f"http://localhost:8080/api/{TRUSTIFY_API_VERSION}/" AUTH_ENABLED = False custom_theme = Theme({"warning": "magenta", "error": "bold red"}) @@ -363,36 +363,58 @@ def make_request_with_retry( return response raise + page_params_base = dict(base_params) + if TRUSTIFY_API_VERSION == "v3" and "total" not in page_params_base: + page_params_base["total"] = True + with httpx.Client() as client: # First request to get total count - query_params = {**base_params, "limit": limit, "offset": 0} + query_params = {**page_params_base, "limit": limit, "offset": 0} first_response = make_request_with_retry(client, query_params, auth_header) first_result = first_response.json() - total_available = first_result.get("total", 0) - if total_available == 0: + all_items = first_result.get("items", []) + total_available = first_result.get("total") + total_known = total_available is not None + + if not all_items and (not total_known or total_available == 0): if component_name: console.print(f"No items found for {component_name}") return {"items": [], "total": 0} - all_items = first_result.get("items", []) - total_pages = (total_available + limit - 1) // limit - - if logger.isEnabledFor(logging.DEBUG): - logger.debug( - f"Paginated request: {total_available} total items, " - f"{total_pages} page(s), page 1/{total_pages} complete" - ) + if total_known: + total_pages = (total_available + limit - 1) // limit + if logger.isEnabledFor(logging.DEBUG): + logger.debug( + f"Paginated request: {total_available} total items, " + f"{total_pages} page(s), page 1/{total_pages} complete" + ) + else: + if logger.isEnabledFor(logging.DEBUG): + logger.debug( + f"Paginated request: total unknown, " + f"page 1 returned {len(all_items)} items" + ) - # Fetch remaining pages sequentially offset = limit page_num = 2 - while offset < total_available: - page_params = {**base_params, "limit": limit, "offset": offset} + while True: + if total_known and offset >= total_available: + break + if not total_known and len(all_items) - (offset - limit) < limit: + break + + page_params = {**page_params_base, "limit": limit, "offset": offset} if logger.isEnabledFor(logging.DEBUG): - logger.debug( - f"Fetching page {page_num}/{total_pages} (offset {offset})..." - ) + if total_known: + total_pages = (total_available + limit - 1) // limit + logger.debug( + f"Fetching page {page_num}/{total_pages} (offset {offset})..." + ) + else: + logger.debug( + f"Fetching page {page_num} (offset {offset})..." + ) try: response = make_request_with_retry(client, page_params, auth_header) result = response.json() @@ -400,21 +422,24 @@ def make_request_with_retry( all_items.extend(page_items) if logger.isEnabledFor(logging.DEBUG): logger.debug( - f"Page {page_num}/{total_pages} complete " - f"({len(all_items)}/{total_available} items)" + f"Page {page_num} complete " + f"({len(all_items)} items so far)" ) + if not page_items or len(page_items) < limit: + break offset += limit page_num += 1 except Exception as e: logger.error(f"Error fetching page at offset {offset}: {e}") break + total_count = total_available if total_known else len(all_items) if component_name: console.print( - f"Retrieved {len(all_items)} items out of {total_available} total for {component_name}" + f"Retrieved {len(all_items)} items out of {total_count} total for {component_name}" ) - return {"items": all_items, "total": total_available} + return {"items": all_items, "total": total_count} def render_tree_to_string(root: Node) -> str: diff --git a/tests/test_pagination.py b/tests/test_pagination.py new file mode 100644 index 0000000..64f8893 --- /dev/null +++ b/tests/test_pagination.py @@ -0,0 +1,126 @@ +from unittest.mock import MagicMock, patch + +from trustshell import paginated_trustify_query + + +def _mock_response(json_data: dict) -> MagicMock: + response = MagicMock() + response.json.return_value = json_data + response.raise_for_status = MagicMock() + return response + + +def _item(n: int) -> dict[str, int]: + return {"id": n} + + +@patch("trustshell.AUTH_ENABLED", False) +@patch("trustshell.TRUSTIFY_API_VERSION", "v3") +@patch("trustshell.httpx.Client") +class TestPaginatedTrustifyQuery: + endpoint = "http://localhost:8080/api/v3/analysis/latest/component" + + def test_v3_requests_total_param(self, mock_client_cls: MagicMock) -> None: + mock_client = MagicMock() + mock_client_cls.return_value.__enter__.return_value = mock_client + mock_client.get.return_value = _mock_response({"items": [_item(1)], "total": 1}) + + paginated_trustify_query(self.endpoint, {"q": "purl~foo"}, {}, limit=100) + + first_call_params = mock_client.get.call_args_list[0].kwargs["params"] + assert first_call_params["total"] is True + assert first_call_params["limit"] == 100 + assert first_call_params["offset"] == 0 + + def test_known_total_fetches_all_pages(self, mock_client_cls: MagicMock) -> None: + mock_client = MagicMock() + mock_client_cls.return_value.__enter__.return_value = mock_client + mock_client.get.side_effect = [ + _mock_response({"items": [_item(i) for i in range(100)], "total": 150}), + _mock_response( + {"items": [_item(i) for i in range(100, 150)], "total": 150} + ), + ] + + result = paginated_trustify_query( + self.endpoint, {"q": "purl~foo"}, {}, limit=100 + ) + + assert len(result["items"]) == 150 + assert result["total"] == 150 + assert mock_client.get.call_count == 2 + + def test_null_total_fetches_until_short_page( + self, mock_client_cls: MagicMock + ) -> None: + mock_client = MagicMock() + mock_client_cls.return_value.__enter__.return_value = mock_client + mock_client.get.side_effect = [ + _mock_response({"items": [_item(i) for i in range(100)], "total": None}), + _mock_response( + {"items": [_item(i) for i in range(100, 125)], "total": None} + ), + ] + + result = paginated_trustify_query( + self.endpoint, {"q": "purl~foo"}, {}, limit=100 + ) + + assert len(result["items"]) == 125 + assert result["total"] == 125 + assert mock_client.get.call_count == 2 + + def test_empty_response_with_null_total(self, mock_client_cls: MagicMock) -> None: + mock_client = MagicMock() + mock_client_cls.return_value.__enter__.return_value = mock_client + mock_client.get.return_value = _mock_response({"items": [], "total": None}) + + result = paginated_trustify_query( + self.endpoint, {"q": "purl~missing"}, {}, limit=100 + ) + + assert result == {"items": [], "total": 0} + assert mock_client.get.call_count == 1 + + def test_respects_explicit_total_param(self, mock_client_cls: MagicMock) -> None: + mock_client = MagicMock() + mock_client_cls.return_value.__enter__.return_value = mock_client + mock_client.get.return_value = _mock_response({"items": [_item(1)], "total": 1}) + + paginated_trustify_query( + self.endpoint, {"q": "purl~foo", "total": False}, {}, limit=100 + ) + + first_call_params = mock_client.get.call_args_list[0].kwargs["params"] + assert first_call_params["total"] is False + + +@patch("trustshell.AUTH_ENABLED", False) +@patch("trustshell.TRUSTIFY_API_VERSION", "v2") +@patch("trustshell.httpx.Client") +class TestPaginatedTrustifyQueryV2: + endpoint = "http://localhost:8080/api/v2/analysis/latest/component" + + def test_v2_does_not_add_total_param(self, mock_client_cls: MagicMock) -> None: + mock_client = MagicMock() + mock_client_cls.return_value.__enter__.return_value = mock_client + mock_client.get.return_value = _mock_response({"items": [_item(1)], "total": 1}) + + paginated_trustify_query(self.endpoint, {"q": "purl~foo"}, {}, limit=100) + + first_call_params = mock_client.get.call_args_list[0].kwargs["params"] + assert "total" not in first_call_params + + def test_v2_numeric_total_pagination(self, mock_client_cls: MagicMock) -> None: + mock_client = MagicMock() + mock_client_cls.return_value.__enter__.return_value = mock_client + mock_client.get.side_effect = [ + _mock_response({"items": [_item(0)], "total": 2}), + _mock_response({"items": [_item(1)], "total": 2}), + ] + + result = paginated_trustify_query(self.endpoint, {"q": "purl~foo"}, {}, limit=1) + + assert len(result["items"]) == 2 + assert result["total"] == 2 + assert mock_client.get.call_count == 2