[DOGE-106] Change behavior of aggregate (#2121)

Author: doctrino

CHANGELOG.md

@@ -22,6 +22,11 @@ Changes are grouped as follows
 - Support for the `/simulators/models` and `/simulators/models/revisions` API endpoints.
 - Support for the `/simulators` and `/simulators/integration` API endpoints.
 
+## [7.73.5] - 2025-02-26
+### Fixed
+- The `client.data_modeling.instances.aggregate/search()` methods now correctly returns maximum, 1000, results when setting
+  the `limit` parameter to `None`, `-1`, or `math.inf`.
+
 ## [7.73.4] - 2025-02-24
 ### Fixed
 - An issue with `DatapointsAPI.retrieve_latest` and usage of `instance_id` when using `ignore_unknown_ids=True`

cognite/client/_api/data_modeling/instances.py

@@ -71,15 +71,15 @@
 )
 from cognite.client.data_classes.data_modeling.views import View
 from cognite.client.data_classes.filters import _BASIC_FILTERS, Filter, _validate_filter
-from cognite.client.utils._auxiliary import load_yaml_or_json
+from cognite.client.utils._auxiliary import is_unlimited, load_yaml_or_json
 from cognite.client.utils._concurrency import ConcurrencySettings
 from cognite.client.utils._identifier import DataModelingIdentifierSequence
 from cognite.client.utils._retry import Backoff
 from cognite.client.utils._text import random_string
 from cognite.client.utils.useful_types import SequenceNotStr
 
 if TYPE_CHECKING:
-    from cognite.client import CogniteClient
+    from cognite.client import ClientConfig, CogniteClient
 
 _FILTERS_SUPPORTED: frozenset[type[Filter]] = _BASIC_FILTERS.union(
     {filters.Nested, filters.HasData, filters.MatchAll, filters.Overlaps, filters.InstanceReferences}
@@ -164,6 +164,11 @@ def _load(data: dict, cognite_client: CogniteClient | None = None) -> NodeApply
 class InstancesAPI(APIClient):
     _RESOURCE_PATH = "/models/instances"
 
+    def __init__(self, config: ClientConfig, api_version: str | None, cognite_client: CogniteClient) -> None:
+        super().__init__(config, api_version, cognite_client)
+        self._AGGREGATE_LIMIT = 1000
+        self._SEARCH_LIMIT = 1000
+
     @overload
     def __call__(
         self,
@@ -1042,7 +1047,7 @@ def search(
         space: str | SequenceNotStr[str] | None = None,
         filter: Filter | dict[str, Any] | None = None,
         include_typing: bool = False,
-        limit: int = DEFAULT_LIMIT_READ,
+        limit: int | None = DEFAULT_LIMIT_READ,
         sort: Sequence[InstanceSort | dict] | InstanceSort | dict | None = None,
     ) -> NodeList[Node]: ...
 
@@ -1058,7 +1063,7 @@ def search(
         space: str | SequenceNotStr[str] | None = None,
         filter: Filter | dict[str, Any] | None = None,
         include_typing: bool = False,
-        limit: int = DEFAULT_LIMIT_READ,
+        limit: int | None = DEFAULT_LIMIT_READ,
         sort: Sequence[InstanceSort | dict] | InstanceSort | dict | None = None,
     ) -> EdgeList[Edge]: ...
 
@@ -1074,7 +1079,7 @@ def search(
         space: str | SequenceNotStr[str] | None = None,
         filter: Filter | dict[str, Any] | None = None,
         include_typing: bool = False,
-        limit: int = DEFAULT_LIMIT_READ,
+        limit: int | None = DEFAULT_LIMIT_READ,
         sort: Sequence[InstanceSort | dict] | InstanceSort | dict | None = None,
     ) -> NodeList[T_Node]: ...
 
@@ -1090,7 +1095,7 @@ def search(
         space: str | SequenceNotStr[str] | None = None,
         filter: Filter | dict[str, Any] | None = None,
         include_typing: bool = False,
-        limit: int = DEFAULT_LIMIT_READ,
+        limit: int | None = DEFAULT_LIMIT_READ,
         sort: Sequence[InstanceSort | dict] | InstanceSort | dict | None = None,
     ) -> EdgeList[T_Edge]: ...
 
@@ -1104,7 +1109,7 @@ def search(
         space: str | SequenceNotStr[str] | None = None,
         filter: Filter | dict[str, Any] | None = None,
         include_typing: bool = False,
-        limit: int = DEFAULT_LIMIT_READ,
+        limit: int | None = DEFAULT_LIMIT_READ,
         sort: Sequence[InstanceSort | dict] | InstanceSort | dict | None = None,
     ) -> NodeList[T_Node] | EdgeList[T_Edge]:
         """`Search instances <https://developer.cognite.com/api/v1/#tag/Instances/operation/searchInstances>`_
@@ -1118,7 +1123,8 @@ def search(
             space (str | SequenceNotStr[str] | None): Restrict instance search to the given space (or list of spaces).
             filter (Filter | dict[str, Any] | None): Advanced filtering of instances.
             include_typing (bool): Whether to include typing information.
-            limit (int): Maximum number of instances to return. Defaults to 25.
+            limit (int | None): Maximum number of instances to return. Defaults to 25. Will return the maximum number
+                of results (1000) if set to None, -1, or math.inf.
             sort (Sequence[InstanceSort | dict] | InstanceSort | dict | None): How you want the listed instances information ordered.
 
         Returns:
@@ -1166,7 +1172,11 @@ def search(
         else:
             raise ValueError(f"Invalid instance type: {instance_type}")
 
-        body: dict[str, Any] = {"view": view.dump(camel_case=True), "instanceType": instance_type_str, "limit": limit}
+        body: dict[str, Any] = {
+            "view": view.dump(camel_case=True),
+            "instanceType": instance_type_str,
+            "limit": self._SEARCH_LIMIT if is_unlimited(limit) else limit,
+        }
         if query:
             body["query"] = query
         if properties:
@@ -1202,7 +1212,7 @@ def aggregate(
         target_units: list[TargetUnit] | None = None,
         space: str | SequenceNotStr[str] | None = None,
         filter: Filter | dict[str, Any] | None = None,
-        limit: int = DEFAULT_LIMIT_READ,
+        limit: int | None = DEFAULT_LIMIT_READ,
     ) -> AggregatedNumberedValue: ...
 
     @overload
@@ -1217,7 +1227,7 @@ def aggregate(
         target_units: list[TargetUnit] | None = None,
         space: str | SequenceNotStr[str] | None = None,
         filter: Filter | dict[str, Any] | None = None,
-        limit: int = DEFAULT_LIMIT_READ,
+        limit: int | None = DEFAULT_LIMIT_READ,
     ) -> list[AggregatedNumberedValue]: ...
 
     @overload
@@ -1232,7 +1242,7 @@ def aggregate(
         target_units: list[TargetUnit] | None = None,
         space: str | SequenceNotStr[str] | None = None,
         filter: Filter | dict[str, Any] | None = None,
-        limit: int = DEFAULT_LIMIT_READ,
+        limit: int | None = DEFAULT_LIMIT_READ,
     ) -> InstanceAggregationResultList: ...
 
     def aggregate(
@@ -1246,7 +1256,7 @@ def aggregate(
         target_units: list[TargetUnit] | None = None,
         space: str | SequenceNotStr[str] | None = None,
         filter: Filter | dict[str, Any] | None = None,
-        limit: int = DEFAULT_LIMIT_READ,
+        limit: int | None = DEFAULT_LIMIT_READ,
     ) -> AggregatedNumberedValue | list[AggregatedNumberedValue] | InstanceAggregationResultList:
         """`Aggregate data across nodes/edges <https://developer.cognite.com/api/v1/#tag/Instances/operation/aggregateInstances>`_
 
@@ -1260,7 +1270,8 @@ def aggregate(
             target_units (list[TargetUnit] | None): Properties to convert to another unit. The API can only convert to another unit if a unit has been defined as part of the type on the underlying container being queried.
             space (str | SequenceNotStr[str] | None): Restrict instance aggregate query to the given space (or list of spaces).
             filter (Filter | dict[str, Any] | None): Advanced filtering of instances.
-            limit (int): Maximum number of instances to return. Defaults to 25.
+            limit (int | None): Maximum number of instances to return. Defaults to 25. Will return the maximum number
+                of results (1000) if set to None, -1, or math.inf.
 
         Returns:
             AggregatedNumberedValue | list[AggregatedNumberedValue] | InstanceAggregationResultList: Node or edge aggregation results.
@@ -1282,7 +1293,11 @@ def aggregate(
 
         self._validate_filter(filter)
         filter = self._merge_space_into_filter(instance_type, space, filter)
-        body: dict[str, Any] = {"view": view.dump(camel_case=True), "instanceType": instance_type, "limit": limit}
+        body: dict[str, Any] = {
+            "view": view.dump(camel_case=True),
+            "instanceType": instance_type,
+            "limit": self._AGGREGATE_LIMIT if is_unlimited(limit) else limit,
+        }
         is_single = isinstance(aggregates, (dict, MetricAggregation))
         aggregate_seq: Sequence[MetricAggregation | dict] = (
             [aggregates] if isinstance(aggregates, (dict, MetricAggregation)) else aggregates

cognite/client/_version.py

@@ -1,5 +1,5 @@
 from __future__ import annotations
 
-__version__ = "7.73.4"
+__version__ = "7.73.5"
 
 __api_subversion__ = "20230101"

pyproject.toml

@@ -1,7 +1,7 @@
 [tool.poetry]
 name = "cognite-sdk"
 
-version = "7.73.4"
+version = "7.73.5"
 
 description = "Cognite Python SDK"
 readme = "README.md"

tests/tests_unit/test_api/test_data_modeling/test_instances.py

@@ -1,7 +1,14 @@
 from __future__ import annotations
 
+import json
+import math
+import re
+
 import pytest
+from responses import RequestsMock
 
+from cognite.client import CogniteClient
+from cognite.client.data_classes.aggregations import Count
 from cognite.client.data_classes.data_modeling.ids import ViewId
 from cognite.client.data_classes.data_modeling.query import SourceSelector
 from tests.tests_unit.test_api.test_data_modeling.conftest import make_test_view
@@ -36,3 +43,64 @@ def test_instances_api_dump_instance_source(self, sources, expected):
         # ViewIdentifier = Union[ViewId, Tuple[str, str], Tuple[str, str, str]]
         # ViewIdentifier | Sequence[ViewIdentifier] | View | Sequence[View]
         assert expected == [source.dump() for source in SourceSelector._load_list(sources)]
+
+
+class TestAggregate:
+    @pytest.mark.usefixtures("disable_gzip")
+    @pytest.mark.parametrize("limit", [None, -1, math.inf])
+    def test_aggregate_maximum(
+        self, limit: int | float | None, rsps: RequestsMock, cognite_client: CogniteClient
+    ) -> None:
+        url = re.compile(r".*/models/instances/aggregate$")
+        response = {
+            "items": [
+                {
+                    "instanceType": "node",
+                    "group": {"site": "MyLocation"},
+                    "aggregates": [
+                        {
+                            "aggregate": "count",
+                            "property": "site",
+                            "value": 42,
+                        }
+                    ],
+                },
+            ]
+        }
+        rsps.add(rsps.POST, url, status=200, json=response)
+
+        _ = cognite_client.data_modeling.instances.aggregate(
+            ViewId("my_space", "MyView", "v1"), Count("externalId"), group_by="site", limit=limit
+        )
+        assert len(rsps.calls) == 1
+        call = rsps.calls[0]
+        body = json.loads(call.request.body)
+        assert "limit" in body
+        assert body["limit"] == cognite_client.data_modeling.instances._AGGREGATE_LIMIT
+
+
+class TestSearch:
+    @pytest.mark.usefixtures("disable_gzip")
+    @pytest.mark.parametrize("limit", [None, -1, math.inf])
+    def test_search_maximum(self, limit: int | float | None, rsps: RequestsMock, cognite_client: CogniteClient) -> None:
+        url = re.compile(r".*/models/instances/search$")
+        response = {
+            "items": [
+                {
+                    "instanceType": "node",
+                    "version": 1,
+                    "space": "my_instance_space",
+                    "externalId": "my_instance_id",
+                    "createdTime": 0,
+                    "lastUpdatedTime": 0,
+                },
+            ]
+        }
+        rsps.add(rsps.POST, url, status=200, json=response)
+
+        _ = cognite_client.data_modeling.instances.search(ViewId("my_space", "MyView", "v1"), "dummy text", limit=limit)
+        assert len(rsps.calls) == 1
+        call = rsps.calls[0]
+        body = json.loads(call.request.body)
+        assert "limit" in body
+        assert body["limit"] == cognite_client.data_modeling.instances._SEARCH_LIMIT