port-ocean 0.28.15__tar.gz → 0.28.17__tar.gz

{port_ocean-0.28.15 → port_ocean-0.28.17}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: port-ocean
-Version: 0.28.15
+Version: 0.28.17
 Summary: Port Ocean is a CLI tool for managing your Port projects.
 Home-page: https://app.getport.io
 Keywords: ocean,port-ocean,port

{port_ocean-0.28.15 → port_ocean-0.28.17}/port_ocean/clients/port/mixins/entities.py

@@ -1,24 +1,24 @@
 import asyncio
+import json
 from typing import Any, Literal
 from urllib.parse import quote_plus
-import json
 
 import httpx
 from loguru import logger
-from port_ocean.context.ocean import ocean
+from starlette import status
+
 from port_ocean.clients.port.authentication import PortAuthentication
 from port_ocean.clients.port.types import RequestOptions, UserAgentType
 from port_ocean.clients.port.utils import (
-    handle_port_status_code,
     PORT_HTTP_MAX_CONNECTIONS_LIMIT,
+    handle_port_status_code,
 )
+from port_ocean.context.ocean import ocean
 from port_ocean.core.models import (
     BulkUpsertResponse,
     Entity,
     PortAPIErrorMessage,
 )
-from starlette import status
-
 from port_ocean.helpers.metric.metric import MetricPhase, MetricType
 
 ENTITIES_BULK_SAMPLES_SIZE = 10
@@ -484,54 +484,89 @@ class EntityClientMixin:
         parameters_to_include: list[str] | None = None,
     ) -> list[Entity]:
         if query is None:
-            datasource_prefix = f"port-ocean/{self.auth.integration_type}/"
-            datasource_suffix = (
-                f"/{self.auth.integration_identifier}/{user_agent_type.value}"
-            )
-            logger.info(
-                f"Searching entities with datasource prefix: {datasource_prefix} and suffix: {datasource_suffix}"
-            )
+            return await self._search_entities_by_datasource_paginated(user_agent_type)
+
+        return await self._search_entities_by_query(
+            user_agent_type=user_agent_type,
+            query=query,
+            parameters_to_include=parameters_to_include,
+        )
+
+    async def _search_entities_by_datasource_paginated(
+        self, user_agent_type: UserAgentType
+    ) -> list[Entity]:
+        datasource_prefix = f"port-ocean/{self.auth.integration_type}/"
+        datasource_suffix = (
+            f"/{self.auth.integration_identifier}/{user_agent_type.value}"
+        )
+        logger.info(
+            f"Searching entities with datasource prefix: {datasource_prefix} and suffix: {datasource_suffix}"
+        )
+
+        next_from: str | None = None
+        aggregated_entities: list[Entity] = []
+        while True:
+            request_body: dict[str, Any] = {
+                "datasource_prefix": datasource_prefix,
+                "datasource_suffix": datasource_suffix,
+            }
+
+            if next_from:
+                request_body["from"] = next_from
 
             response = await self.client.post(
                 f"{self.auth.api_url}/blueprints/entities/datasource-entities",
-                json={
-                    "datasource_prefix": datasource_prefix,
-                    "datasource_suffix": datasource_suffix,
-                },
+                json=request_body,
                 headers=await self.auth.headers(user_agent_type),
                 extensions={"retryable": True},
             )
-        else:
-            default_query = {
-                "combinator": "and",
-                "rules": [
-                    {
-                        "property": "$datasource",
-                        "operator": "contains",
-                        "value": f"port-ocean/{self.auth.integration_type}/",
-                    },
-                    {
-                        "property": "$datasource",
-                        "operator": "contains",
-                        "value": f"/{self.auth.integration_identifier}/{user_agent_type.value}",
-                    },
-                ],
-            }
+            handle_port_status_code(response)
+            response_json = response.json()
+            aggregated_entities.extend(
+                Entity.parse_obj(result) for result in response_json.get("entities", [])
+            )
+            next_from = response_json.get("next")
+            if not next_from:
+                break
 
-            if query.get("rules"):
-                query["rules"].extend(default_query["rules"])
+        return aggregated_entities
 
-            logger.info(f"Searching entities with custom query: {query}")
-            response = await self.client.post(
-                f"{self.auth.api_url}/entities/search",
-                json=query,
-                headers=await self.auth.headers(user_agent_type),
-                params={
-                    "exclude_calculated_properties": "true",
-                    "include": parameters_to_include or ["blueprint", "identifier"],
+    async def _search_entities_by_query(
+        self,
+        user_agent_type: UserAgentType,
+        query: dict[Any, Any],
+        parameters_to_include: list[str] | None,
+    ) -> list[Entity]:
+        default_query = {
+            "combinator": "and",
+            "rules": [
+                {
+                    "property": "$datasource",
+                    "operator": "contains",
+                    "value": f"port-ocean/{self.auth.integration_type}/",
                 },
-                extensions={"retryable": True},
-            )
+                {
+                    "property": "$datasource",
+                    "operator": "contains",
+                    "value": f"/{self.auth.integration_identifier}/{user_agent_type.value}",
+                },
+            ],
+        }
+
+        if query.get("rules"):
+            query["rules"].extend(default_query["rules"])
+
+        logger.info(f"Searching entities with custom query: {query}")
+        response = await self.client.post(
+            f"{self.auth.api_url}/entities/search",
+            json=query,
+            headers=await self.auth.headers(user_agent_type),
+            params={
+                "exclude_calculated_properties": "true",
+                "include": parameters_to_include or ["blueprint", "identifier"],
+            },
+            extensions={"retryable": True},
+        )
 
         handle_port_status_code(response)
         return [Entity.parse_obj(result) for result in response.json()["entities"]]

{port_ocean-0.28.15 → port_ocean-0.28.17}/port_ocean/helpers/retry.py

@@ -359,8 +359,11 @@ class RetryTransport(httpx.AsyncBaseTransport, httpx.BaseTransport):
         else:
             # If no Content-Length header, try to get actual content size
             try:
-                actual_size = len(await response.aread())
+                content = await response.aread()
+                actual_size = len(content)
                 size_info = actual_size
+                # Restore the cached body so downstream code can still use .json()/.text/.content
+                response._content = content  # httpx convention
             except Exception as e:
                 cast(logging.Logger, self._logger).error(
                     f"Error getting response size: {e}"
@@ -383,8 +386,11 @@ class RetryTransport(httpx.AsyncBaseTransport, httpx.BaseTransport):
         else:
             # If no Content-Length header, try to get actual content size
             try:
-                actual_size = len(response.read())
+                content = response.read()
+                actual_size = len(content)
                 size_info = actual_size
+                # Restore the cached body so downstream code can still use .json()/.text/.content
+                response._content = content  # httpx convention
             except Exception as e:
                 cast(logging.Logger, self._logger).error(
                     f"Error getting response size: {e}"

{port_ocean-0.28.15 → port_ocean-0.28.17}/port_ocean/tests/clients/port/mixins/test_entities.py

@@ -179,7 +179,7 @@ async def test_search_entities_uses_datasource_route_when_query_is_none(
 ) -> None:
     """Test that search_entities uses datasource route when query is None"""
     mock_response = MagicMock()
-    mock_response.json.return_value = {"entities": []}
+    mock_response.json.return_value = {"entities": [], "next": None}
     mock_response.is_error = False
     mock_response.status_code = 200
     mock_response.headers = {}
@@ -208,8 +208,74 @@ async def test_search_entities_uses_datasource_route_when_query_is_none(
         == "https://api.getport.io/v1/blueprints/entities/datasource-entities"
     )
 
-    expected_json = {
-        "datasource_prefix": "port-ocean/test-integration/",
-        "datasource_suffix": "/test-identifier/sync",
+    sent_json = call_args[1]["json"]
+    assert sent_json["datasource_prefix"] == "port-ocean/test-integration/"
+    assert sent_json["datasource_suffix"] == "/test-identifier/sync"
+
+
+async def test_search_entities_uses_datasource_route_when_query_is_none_two_pages(
+    entity_client: EntityClientMixin,
+) -> None:
+    """Test that search_entities uses datasource route when query is None"""
+    # First response with pagination token
+    mock_response_first = MagicMock()
+    mock_response_first.json.return_value = {
+        "entities": [
+            Entity(identifier="entity_1", blueprint="entity_1"),
+            Entity(identifier="entity_2", blueprint="entity_2"),
+        ],
+        "next": "next_page_token",
+    }
+    mock_response_first.is_error = False
+    mock_response_first.status_code = 200
+    mock_response_first.headers = {}
+
+    # Second response without pagination token (end of pagination)
+    mock_response_second = MagicMock()
+    mock_response_second.json.return_value = {
+        "entities": [Entity(identifier="entity_3", blueprint="entity_3")],
+        "next": None,
     }
-    assert call_args[1]["json"] == expected_json
+    mock_response_second.is_error = False
+    mock_response_second.status_code = 200
+    mock_response_second.headers = {}
+
+    # Mock the client to return different responses for each call
+    entity_client.client.post = AsyncMock(side_effect=[mock_response_first, mock_response_second])  # type: ignore
+    entity_client.auth.headers = AsyncMock(return_value={"Authorization": "Bearer test"})  # type: ignore
+
+    entity_client.auth.integration_type = "test-integration"
+    entity_client.auth.integration_identifier = "test-identifier"
+    entity_client.auth.api_url = "https://api.getport.io/v1"
+
+    mock_user_agent_type = MagicMock()
+    mock_user_agent_type.value = "sync"
+
+    entities = await entity_client.search_entities(
+        user_agent_type=mock_user_agent_type,
+        query=None,
+    )
+
+    # Should call the datasource-entities endpoint exactly twice for pagination
+    assert entity_client.client.post.await_count == 2
+    assert len(entities) == 3
+
+    # Check first call
+    first_call_args = entity_client.client.post.call_args_list[0]
+    assert (
+        first_call_args[0][0]
+        == "https://api.getport.io/v1/blueprints/entities/datasource-entities"
+    )
+    first_sent_json = first_call_args[1]["json"]
+    assert first_sent_json["datasource_prefix"] == "port-ocean/test-integration/"
+    assert first_sent_json["datasource_suffix"] == "/test-identifier/sync"
+
+    # Check second call
+    second_call_args = entity_client.client.post.call_args_list[1]
+    assert (
+        second_call_args[0][0]
+        == "https://api.getport.io/v1/blueprints/entities/datasource-entities"
+    )
+    second_sent_json = second_call_args[1]["json"]
+    assert second_sent_json["datasource_prefix"] == "port-ocean/test-integration/"
+    assert second_sent_json["datasource_suffix"] == "/test-identifier/sync"