Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[DOGE-106] Instance aggregate/search limit=None means maximum #2121

Open
wants to merge 11 commits into
base: master
Choose a base branch
from
5 changes: 5 additions & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -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`
Expand Down
45 changes: 30 additions & 15 deletions cognite/client/_api/data_modeling/instances.py
Original file line number Diff line number Diff line change
Expand Up @@ -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}
Expand Down Expand Up @@ -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
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Shouldn't these be class variables instead?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hmm, that could be. Consistently in the SDK we have limits as instance
attributes. Suggest that if we move these two class variables, we do it everywhere, which should in that case be a separate PR.

self._SEARCH_LIMIT = 1000

@overload
def __call__(
self,
Expand Down Expand Up @@ -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]: ...

Expand All @@ -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]: ...

Expand All @@ -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]: ...

Expand All @@ -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]: ...

Expand All @@ -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>`_
Expand All @@ -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:
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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(
Expand All @@ -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>`_

Expand All @@ -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.
Expand All @@ -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
Expand Down
2 changes: 1 addition & 1 deletion cognite/client/_version.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
from __future__ import annotations

__version__ = "7.73.4"
__version__ = "7.73.5"

__api_subversion__ = "20230101"
2 changes: 1 addition & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
[tool.poetry]
name = "cognite-sdk"

version = "7.73.4"
version = "7.73.5"

description = "Cognite Python SDK"
readme = "README.md"
Expand Down
68 changes: 68 additions & 0 deletions tests/tests_unit/test_api/test_data_modeling/test_instances.py
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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