diff --git a/cognite/client/_api/data_modeling/records.py b/cognite/client/_api/data_modeling/records.py index a5320a8766..567394580d 100644 --- a/cognite/client/_api/data_modeling/records.py +++ b/cognite/client/_api/data_modeling/records.py @@ -1,7 +1,7 @@ from __future__ import annotations import asyncio -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Literal from cognite.client._api_client import APIClient @@ -11,6 +11,8 @@ RecordId, RecordIdSequence, RecordList, + RecordsAggregate, + RecordsAggregation, RecordSourceSelector, RecordTargetUnit, RecordTargetUnits, @@ -18,6 +20,7 @@ SyncRecord, SyncRecordList, TimeRange, + _dump_aggregate_value, ) from cognite.client.data_classes.filters import Filter from cognite.client.utils._experimental import FeaturePreviewWarning @@ -235,6 +238,146 @@ async def upsert( no_response=True, ) + async def aggregate( + self, + aggregates: Mapping[str, RecordsAggregate | dict[str, Any]], + *, + stream_id: str, + last_updated_time: TimeRange | None = None, + filter: Filter | dict[str, Any] | None = None, + target_units: RecordTargetUnits | Sequence[RecordTargetUnit] | None = None, + include_typing: bool = False, + ) -> RecordsAggregation: + """`Aggregate records from a stream `_. + + Args: + aggregates (Mapping[str, RecordsAggregate | dict[str, Any]]): Aggregate request tree keyed + by client-defined aggregate IDs. + stream_id (str): External ID of the stream to aggregate from. + last_updated_time (TimeRange | None): Filter records by last-updated time. + **Required** for immutable streams (must include a lower bound). + filter (Filter | dict[str, Any] | None): Filter expression. + target_units (RecordTargetUnits | Sequence[RecordTargetUnit] | None): Unit conversion specification. + include_typing (bool): Include property type metadata in the response. + + Returns: + RecordsAggregation: Aggregate results keyed by the requested aggregate IDs. + + Examples: + + The examples below aggregate over a stream of padel game statistics records; each + example builds on the previous one. + + The property paths used below: + + >>> game_time = ["paddle", "game_statistics", "game_time"] + >>> player_name = ["paddle", "game_statistics", "player_name"] + >>> points_scored = ["paddle", "game_statistics", "points_scored"] + + + Find the average points scored across all games, using a typed helper: + + >>> from cognite.client import CogniteClient + >>> from cognite.client.data_classes.data_modeling.records import Average + >>> client = CogniteClient() + >>> res = client.data_modeling.records.aggregate( + ... stream_id="my-stream", + ... aggregates={"avg_points_scored": Average(points_scored)}, + ... ) + + + Count the total number of games, and how many of them have a recorded score, only + considering games updated after a given time: + + >>> from cognite.client.data_classes.data_modeling.records import Count, TimeRange + >>> res = client.data_modeling.records.aggregate( + ... stream_id="my-stream", + ... aggregates={ + ... "total_games": Count(), + ... "games_with_score": Count(points_scored), + ... }, + ... last_updated_time=TimeRange(gt=1759276800000), + ... ) + + + Group games by day, then by player, and for each player-day compute their total, + highest, and average points scored, alongside the single highest score across all + games: + + >>> from cognite.client.data_classes.data_modeling.records import ( + ... Average, + ... Max, + ... Sum, + ... TimeHistogram, + ... UniqueValues, + ... ) + >>> res = client.data_modeling.records.aggregate( + ... stream_id="my-stream", + ... aggregates={ + ... "my_groups_by_1d_range": TimeHistogram( + ... property=game_time, + ... calendar_interval="1d", + ... aggregates={ + ... "my_groups_by_player_name": UniqueValues( + ... property=player_name, + ... aggregates={ + ... "my_player_daily_scores_sum": Sum(points_scored), + ... "my_player_daily_scores_maximum": Max(points_scored), + ... }, + ... ), + ... "my_daily_scores_average": Average(points_scored), + ... }, + ... ), + ... "my_scores_maximum_across_all_games": Max(points_scored), + ... }, + ... ) + + + Bucket games by day and smooth the daily count with a 7-day moving average, using the + ``MovingFunctions`` enum so the pipeline function name cannot be mistyped: + + >>> from cognite.client.data_classes.data_modeling.records import ( + ... Count, + ... MovingFunction, + ... MovingFunctions, + ... TimeHistogram, + ... ) + >>> res = client.data_modeling.records.aggregate( + ... stream_id="my-stream", + ... aggregates={ + ... "games_per_day": TimeHistogram( + ... property=game_time, + ... calendar_interval="1d", + ... aggregates={ + ... "games": Count(), + ... "games_7d_avg": MovingFunction( + ... buckets_path="games", + ... window=7, + ... function=MovingFunctions.UNWEIGHTED_AVG, + ... ), + ... }, + ... ), + ... }, + ... ) + """ + self._warning.warn() + body: dict[str, Any] = {"aggregates": _dump_aggregate_value(aggregates)} + if last_updated_time is not None: + body["lastUpdatedTime"] = last_updated_time.dump() + if filter is not None: + body["filter"] = filter.dump() if isinstance(filter, Filter) else filter + if target_units is not None: + body["targetUnits"] = self._dump_target_units(target_units) + if include_typing: + body["includeTyping"] = True + + res = await self._post( + url_path=self._records_url(stream_id, "/aggregate"), + json=body, + semaphore=self._get_semaphore("read"), + ) + return RecordsAggregation._load(res.json()) + async def filter( self, stream_id: str, diff --git a/cognite/client/_sync_api/data_modeling/records.py b/cognite/client/_sync_api/data_modeling/records.py index f527e31726..ab05f7924c 100644 --- a/cognite/client/_sync_api/data_modeling/records.py +++ b/cognite/client/_sync_api/data_modeling/records.py @@ -1,14 +1,14 @@ """ =============================================================================== -b648da6cf09c84c8e1361f9c4e6c4052 +b24d174ece0a93fee285c7a4e78766b0 This file is auto-generated from the Async API modules, - do not edit manually! =============================================================================== """ from __future__ import annotations -from collections.abc import Sequence -from typing import TYPE_CHECKING, Literal +from collections.abc import Mapping, Sequence +from typing import TYPE_CHECKING, Any, Literal from cognite.client import AsyncCogniteClient from cognite.client._sync_api_client import SyncAPIClient @@ -16,6 +16,8 @@ from cognite.client.data_classes.data_modeling.records import ( RecordId, RecordList, + RecordsAggregate, + RecordsAggregation, RecordSourceSelector, RecordTargetUnit, RecordTargetUnits, @@ -160,6 +162,140 @@ def upsert( self.__async_client.data_modeling.records.upsert(items=items, stream_id=stream_id, upsert_mode=upsert_mode) ) + def aggregate( + self, + aggregates: Mapping[str, RecordsAggregate | dict[str, Any]], + *, + stream_id: str, + last_updated_time: TimeRange | None = None, + filter: Filter | dict[str, Any] | None = None, + target_units: RecordTargetUnits | Sequence[RecordTargetUnit] | None = None, + include_typing: bool = False, + ) -> RecordsAggregation: + """ + `Aggregate records from a stream `_. + + Args: + aggregates (Mapping[str, RecordsAggregate | dict[str, Any]]): Aggregate request tree keyed + by client-defined aggregate IDs. + stream_id (str): External ID of the stream to aggregate from. + last_updated_time (TimeRange | None): Filter records by last-updated time. + **Required** for immutable streams (must include a lower bound). + filter (Filter | dict[str, Any] | None): Filter expression. + target_units (RecordTargetUnits | Sequence[RecordTargetUnit] | None): Unit conversion specification. + include_typing (bool): Include property type metadata in the response. + + Returns: + RecordsAggregation: Aggregate results keyed by the requested aggregate IDs. + + Examples: + + The examples below aggregate over a stream of padel game statistics records; each + example builds on the previous one. + + The property paths used below: + + >>> game_time = ["paddle", "game_statistics", "game_time"] + >>> player_name = ["paddle", "game_statistics", "player_name"] + >>> points_scored = ["paddle", "game_statistics", "points_scored"] + + + Find the average points scored across all games, using a typed helper: + + >>> from cognite.client import CogniteClient + >>> from cognite.client.data_classes.data_modeling.records import Average + >>> client = CogniteClient() + >>> res = client.data_modeling.records.aggregate( + ... stream_id="my-stream", + ... aggregates={"avg_points_scored": Average(points_scored)}, + ... ) + + + Count the total number of games, and how many of them have a recorded score, only + considering games updated after a given time: + + >>> from cognite.client.data_classes.data_modeling.records import Count, TimeRange + >>> res = client.data_modeling.records.aggregate( + ... stream_id="my-stream", + ... aggregates={ + ... "total_games": Count(), + ... "games_with_score": Count(points_scored), + ... }, + ... last_updated_time=TimeRange(gt=1759276800000), + ... ) + + + Group games by day, then by player, and for each player-day compute their total, + highest, and average points scored, alongside the single highest score across all + games: + + >>> from cognite.client.data_classes.data_modeling.records import ( + ... Average, + ... Max, + ... Sum, + ... TimeHistogram, + ... UniqueValues, + ... ) + >>> res = client.data_modeling.records.aggregate( + ... stream_id="my-stream", + ... aggregates={ + ... "my_groups_by_1d_range": TimeHistogram( + ... property=game_time, + ... calendar_interval="1d", + ... aggregates={ + ... "my_groups_by_player_name": UniqueValues( + ... property=player_name, + ... aggregates={ + ... "my_player_daily_scores_sum": Sum(points_scored), + ... "my_player_daily_scores_maximum": Max(points_scored), + ... }, + ... ), + ... "my_daily_scores_average": Average(points_scored), + ... }, + ... ), + ... "my_scores_maximum_across_all_games": Max(points_scored), + ... }, + ... ) + + + Bucket games by day and smooth the daily count with a 7-day moving average, using the + ``MovingFunctions`` enum so the pipeline function name cannot be mistyped: + + >>> from cognite.client.data_classes.data_modeling.records import ( + ... Count, + ... MovingFunction, + ... MovingFunctions, + ... TimeHistogram, + ... ) + >>> res = client.data_modeling.records.aggregate( + ... stream_id="my-stream", + ... aggregates={ + ... "games_per_day": TimeHistogram( + ... property=game_time, + ... calendar_interval="1d", + ... aggregates={ + ... "games": Count(), + ... "games_7d_avg": MovingFunction( + ... buckets_path="games", + ... window=7, + ... function=MovingFunctions.UNWEIGHTED_AVG, + ... ), + ... }, + ... ), + ... }, + ... ) + """ + return run_sync( + self.__async_client.data_modeling.records.aggregate( + aggregates=aggregates, + stream_id=stream_id, + last_updated_time=last_updated_time, + filter=filter, + target_units=target_units, + include_typing=include_typing, + ) + ) + def filter( self, stream_id: str, diff --git a/cognite/client/data_classes/data_modeling/__init__.py b/cognite/client/data_classes/data_modeling/__init__.py index a53c1714f3..de5d82351f 100644 --- a/cognite/client/data_classes/data_modeling/__init__.py +++ b/cognite/client/data_classes/data_modeling/__init__.py @@ -124,8 +124,18 @@ RecordContainerId, RecordId, RecordList, + RecordsAggregateResult, + RecordsAggregation, + RecordsBucket, + RecordsFilterAggregateResult, + RecordsMetricAggregateResult, + RecordsMovingFunctionAggregateResult, + RecordsNumberHistogramAggregateResult, RecordSource, RecordSourceSelector, + RecordsTimeHistogramAggregateResult, + RecordsUniqueValuesAggregateResult, + RecordsUnknownAggregateResult, RecordTargetUnit, RecordTargetUnits, RecordWrite, @@ -267,6 +277,16 @@ "RecordTargetUnits", "RecordWrite", "RecordWriteList", + "RecordsAggregateResult", + "RecordsAggregation", + "RecordsBucket", + "RecordsFilterAggregateResult", + "RecordsMetricAggregateResult", + "RecordsMovingFunctionAggregateResult", + "RecordsNumberHistogramAggregateResult", + "RecordsTimeHistogramAggregateResult", + "RecordsUniqueValuesAggregateResult", + "RecordsUnknownAggregateResult", "RequiresConstraint", "RequiresConstraintApply", "ResultSetExpression", diff --git a/cognite/client/data_classes/data_modeling/records.py b/cognite/client/data_classes/data_modeling/records.py index fccef5979f..9d36fe603a 100644 --- a/cognite/client/data_classes/data_modeling/records.py +++ b/cognite/client/data_classes/data_modeling/records.py @@ -1,9 +1,11 @@ from __future__ import annotations -from collections.abc import Sequence +from abc import abstractmethod +from collections.abc import Mapping, Sequence from copy import deepcopy from dataclasses import dataclass -from typing import Any, Literal +from enum import Enum +from typing import Any, ClassVar, Literal from typing_extensions import Self @@ -16,9 +18,20 @@ from cognite.client.data_classes.data_modeling.data_types import UnitReference, UnitSystemReference from cognite.client.data_classes.data_modeling.ids import ContainerId from cognite.client.data_classes.data_modeling.instances import TypeInformation +from cognite.client.data_classes.filters import Filter from cognite.client.utils._identifier import IdentifierSequenceCore, RecordId +from cognite.client.utils._text import convert_all_keys_to_snake_case, to_snake_case +from cognite.client.utils.useful_types import SequenceNotStr __all__ = [ + "Average", + "Count", + "Filters", + "Max", + "Min", + "MovingFunction", + "MovingFunctions", + "NumberHistogram", "Record", "RecordContainerId", "RecordId", @@ -30,12 +43,325 @@ "RecordTargetUnits", "RecordWrite", "RecordWriteList", + "RecordsAggregate", + "RecordsAggregateResult", + "RecordsAggregation", + "RecordsBucket", + "RecordsFilterAggregateResult", + "RecordsMetricAggregateResult", + "RecordsMovingFunctionAggregateResult", + "RecordsNumberHistogramAggregateResult", + "RecordsTimeHistogramAggregateResult", + "RecordsUniqueValuesAggregateResult", + "RecordsUnknownAggregateResult", + "Sum", "SyncRecord", "SyncRecordList", + "TimeHistogram", "TimeRange", + "UniqueValues", + "UnknownRecordsAggregate", ] +def _dump_aggregate_value(value: Any) -> Any: + match value: + case Mapping() as m: + return {key: _dump_aggregate_value(val) for key, val in m.items()} + case [*items]: + return [_dump_aggregate_value(item) for item in items] + case RecordsAggregate() as agg: + return _dump_aggregate_value(agg.dump()) + case _: + return value + + +def _dump_aggregate_results( + aggregates: dict[str, Any], + results: dict[str, RecordsAggregateResult], + camel_case: bool, +) -> dict[str, Any]: + """Dump a map of aggregate results keyed by client-defined aggregate IDs. + + The IDs are chosen by the caller and left untouched; only each result's own payload honors + ``camel_case``. Entries without a parsed result (e.g. non-dict values) are passed through. + """ + return { + aggregate_id: ( + results[aggregate_id].dump(camel_case=camel_case) + if aggregate_id in results + else _dump_aggregate_value(value) + ) + for aggregate_id, value in aggregates.items() + } + + +class RecordsAggregate(CogniteResource): + """Base class for typed Records aggregate request builders. + + Aggregates are request bodies: they serialize via :meth:`dump` and can be loaded back from that + same representation via :meth:`load`, so an aggregate spec round-trips through a config file. + """ + + _aggregate_name: ClassVar[str] + + @abstractmethod + def _dump_body(self) -> dict[str, Any]: + raise NotImplementedError + + def dump(self, camel_case: bool = True) -> dict[str, Any]: + return {self._aggregate_name: _dump_aggregate_value(self._dump_body())} + + @classmethod + def _load(cls, resource: dict[str, Any]) -> RecordsAggregate: + # A dumped aggregate has a single top-level key naming the aggregate type; dispatch on it. + # Nested `aggregates` are kept as their dumped form (dicts), which dump() handles verbatim. + name = next(iter(resource)) + body = resource[name] + if name == "avg": + return Average(body["property"]) + if name == "min": + return Min(body["property"]) + if name == "max": + return Max(body["property"]) + if name == "sum": + return Sum(body["property"]) + if name == "count": + return Count(body.get("property")) + if name == "uniqueValues": + return UniqueValues(body["property"], aggregates=body.get("aggregates"), size=body.get("size")) + if name == "numberHistogram": + return NumberHistogram( + body["property"], + interval=body["interval"], + aggregates=body.get("aggregates"), + hard_bounds=body.get("hardBounds"), + ) + if name == "timeHistogram": + return TimeHistogram( + body["property"], + calendar_interval=body.get("calendarInterval"), + fixed_interval=body.get("fixedInterval"), + aggregates=body.get("aggregates"), + hard_bounds=body.get("hardBounds"), + ) + if name == "filters": + return Filters(filters=body["filters"], aggregates=body.get("aggregates")) + if name == "movingFunction": + return MovingFunction(buckets_path=body["bucketsPath"], window=body["window"], function=body["function"]) + return UnknownRecordsAggregate(resource) + + +class _PropertyAggregate(RecordsAggregate): + def __init__(self, property: SequenceNotStr[str]) -> None: + self.property = list(property) + + def _dump_body(self) -> dict[str, Any]: + return {"property": self.property} + + +class Average(_PropertyAggregate): + """Average aggregate over a container property.""" + + _aggregate_name = "avg" + + +class Count(RecordsAggregate): + """Count records, or non-null values when ``property`` is provided.""" + + _aggregate_name = "count" + + def __init__(self, property: SequenceNotStr[str] | None = None) -> None: + self.property = list(property) if property is not None else None + + def _dump_body(self) -> dict[str, Any]: + return {"property": self.property} if self.property is not None else {} + + +class Min(_PropertyAggregate): + """Minimum aggregate over a property.""" + + _aggregate_name = "min" + + +class Max(_PropertyAggregate): + """Maximum aggregate over a property.""" + + _aggregate_name = "max" + + +class Sum(_PropertyAggregate): + """Sum aggregate over a container property.""" + + _aggregate_name = "sum" + + +class _NestedAggregate(RecordsAggregate): + def __init__(self, aggregates: Mapping[str, RecordsAggregate | dict[str, Any]] | None = None) -> None: + self.aggregates = aggregates + + def _add_aggregates(self, body: dict[str, Any]) -> dict[str, Any]: + if self.aggregates is not None: + body["aggregates"] = self.aggregates + return body + + +class UniqueValues(_NestedAggregate): + """Bucket records by unique property values.""" + + _aggregate_name = "uniqueValues" + + def __init__( + self, + property: SequenceNotStr[str], + aggregates: Mapping[str, RecordsAggregate | dict[str, Any]] | None = None, + size: int | None = None, + ): + super().__init__(aggregates) + self.property = list(property) + self.size = size + + def _dump_body(self) -> dict[str, Any]: + body: dict[str, Any] = {"property": self.property} + self._add_aggregates(body) + if self.size is not None: + body["size"] = self.size + return body + + +class NumberHistogram(_NestedAggregate): + """Bucket numeric property values into fixed-width intervals.""" + + _aggregate_name = "numberHistogram" + + def __init__( + self, + property: SequenceNotStr[str], + interval: float, + aggregates: Mapping[str, RecordsAggregate | dict[str, Any]] | None = None, + hard_bounds: Mapping[str, float] | None = None, + ) -> None: + super().__init__(aggregates) + self.property = list(property) + self.interval = interval + self.hard_bounds = hard_bounds + + def _dump_body(self) -> dict[str, Any]: + body: dict[str, Any] = {"property": self.property, "interval": self.interval} + self._add_aggregates(body) + if self.hard_bounds is not None: + body["hardBounds"] = self.hard_bounds + return body + + +class TimeHistogram(_NestedAggregate): + """Bucket timestamp values into calendar or fixed time intervals.""" + + _aggregate_name = "timeHistogram" + + def __init__( + self, + property: SequenceNotStr[str], + *, + calendar_interval: str | None = None, + fixed_interval: str | None = None, + aggregates: Mapping[str, RecordsAggregate | dict[str, Any]] | None = None, + hard_bounds: Mapping[str, str] | None = None, + ) -> None: + if (calendar_interval is None) == (fixed_interval is None): + raise ValueError("Exactly one of calendar_interval or fixed_interval must be specified") + super().__init__(aggregates) + self.property = list(property) + self.calendar_interval = calendar_interval + self.fixed_interval = fixed_interval + self.hard_bounds = hard_bounds + + def _dump_body(self) -> dict[str, Any]: + body: dict[str, Any] = {"property": self.property} + self._add_aggregates(body) + if self.calendar_interval is not None: + body["calendarInterval"] = self.calendar_interval + if self.fixed_interval is not None: + body["fixedInterval"] = self.fixed_interval + if self.hard_bounds is not None: + body["hardBounds"] = self.hard_bounds + return body + + +class Filters(_NestedAggregate): + """Bucket records by a list of filter expressions.""" + + _aggregate_name = "filters" + + def __init__( + self, + filters: Sequence[Filter | dict[str, Any]], + aggregates: Mapping[str, RecordsAggregate | dict[str, Any]] | None = None, + ) -> None: + super().__init__(aggregates) + self.filters = filters + + def _dump_body(self) -> dict[str, Any]: + body: dict[str, Any] = { + "filters": [filter.dump() if isinstance(filter, Filter) else filter for filter in self.filters] + } + return self._add_aggregates(body) + + +class MovingFunctions(str, Enum): + """Pipeline functions available to :class:`MovingFunction`.""" + + MAX = "MovingFunctions.max" + MIN = "MovingFunctions.min" + SUM = "MovingFunctions.sum" + UNWEIGHTED_AVG = "MovingFunctions.unweightedAvg" + LINEAR_WEIGHTED_AVG = "MovingFunctions.linearWeightedAvg" + + +class MovingFunction(RecordsAggregate): + """Pipeline aggregate over a parent histogram bucket series.""" + + _aggregate_name = "movingFunction" + + def __init__( + self, + buckets_path: str, + window: int, + function: MovingFunctions + | Literal[ + "MovingFunctions.max", + "MovingFunctions.min", + "MovingFunctions.sum", + "MovingFunctions.unweightedAvg", + "MovingFunctions.linearWeightedAvg", + ], + ) -> None: + self.buckets_path = buckets_path + self.window = window + self.function = MovingFunctions(function) + + def _dump_body(self) -> dict[str, Any]: + return {"bucketsPath": self.buckets_path, "window": self.window, "function": self.function.value} + + +class UnknownRecordsAggregate(RecordsAggregate): + """Fallback for aggregate request shapes this SDK version does not model yet. + + Preserves the raw request body verbatim so an unknown or newer aggregate type still round-trips + through :meth:`dump`/:meth:`load` instead of failing. The request builders' :meth:`dump` is + always camelCase, so the payload is returned as-is regardless of ``camel_case``. + """ + + def __init__(self, raw: dict[str, Any]) -> None: + self._raw = raw + + def _dump_body(self) -> dict[str, Any]: # never called: dump is overridden + raise NotImplementedError + + def dump(self, camel_case: bool = True) -> dict[str, Any]: + return dict(self._raw) + + class RecordIdSequence(IdentifierSequenceCore[RecordId]): @classmethod def load(cls, items: RecordId | Sequence[RecordId]) -> RecordIdSequence: @@ -127,6 +453,194 @@ def as_ids(self) -> list[RecordId]: return [v.as_id() for v in self] +class RecordsAggregateResult(CogniteResource): + """Base class for typed Records aggregate results.""" + + @classmethod + def _load(cls, resource: dict[str, Any]) -> RecordsAggregateResult: + # Dispatcher: each aggregate result carries exactly one top-level key that selects the + # concrete result type. Each subclass implements its own _load returning Self. + assert len(resource) == 1, f"expected exactly one aggregate result key, got {sorted(resource)}" + key = next(iter(resource)) + if key in _METRIC_AGGREGATE_KEYS: + return RecordsMetricAggregateResult._load(resource) + if key == "fnValue": + return RecordsMovingFunctionAggregateResult._load(resource) + if (result_cls := _BUCKET_RESULT_BY_KEY.get(key)) is not None: + return result_cls._load(resource) + return RecordsUnknownAggregateResult._load(resource) + + def dump(self, camel_case: bool = True) -> dict[str, Any]: + raise NotImplementedError + + +class RecordsMetricAggregateResult(RecordsAggregateResult): + """Metric aggregate result such as ``avg``, ``count``, ``min``, ``max``, or ``sum``.""" + + def __init__(self, aggregate: str, value: Any) -> None: + self.aggregate = aggregate + self.value = value + + @classmethod + def _load(cls, resource: dict[str, Any]) -> Self: + key = next(iter(resource)) + return cls(aggregate=key, value=resource[key]) + + def dump(self, camel_case: bool = True) -> dict[str, Any]: + return {self.aggregate: self.value} + + +class RecordsMovingFunctionAggregateResult(RecordsAggregateResult): + def __init__(self, fn_value: float) -> None: + self.fn_value = fn_value + + @classmethod + def _load(cls, resource: dict[str, Any]) -> Self: + return cls(fn_value=resource["fnValue"]) + + def dump(self, camel_case: bool = True) -> dict[str, Any]: + return {"fnValue" if camel_case else "fn_value": self.fn_value} + + +class RecordsUnknownAggregateResult(RecordsAggregateResult): + """Fallback for aggregate result shapes the SDK does not model yet. + + Preserves the raw payload verbatim so nothing is lost, snake-casing the API keys on request. + """ + + def __init__(self, raw_result: dict[str, Any]) -> None: + self._raw_result = raw_result + + @classmethod + def _load(cls, resource: dict[str, Any]) -> Self: + return cls(resource) + + def dump(self, camel_case: bool = True) -> dict[str, Any]: + if camel_case: + return dict(self._raw_result) + return convert_all_keys_to_snake_case(self._raw_result) + + +class RecordsBucket(CogniteResource): + def __init__( + self, + count: int, + value: Any = None, + interval_start: float | str | None = None, + aggregates: dict[str, Any] | None = None, + ) -> None: + self.count = count + self.value = value + self.interval_start = interval_start + self.aggregates = aggregates or {} + self.results = { + aggregate_id: RecordsAggregateResult._load(result) + for aggregate_id, result in self.aggregates.items() + if isinstance(result, dict) + } + + @classmethod + def _load(cls, resource: dict[str, Any]) -> Self: + return cls( + count=resource["count"], + value=resource.get("value"), + interval_start=resource.get("intervalStart"), + aggregates=resource.get("aggregates"), + ) + + def dump(self, camel_case: bool = True) -> dict[str, Any]: + output: dict[str, Any] = {"count": self.count} + if self.value is not None: + output["value"] = self.value + if self.interval_start is not None: + output["intervalStart" if camel_case else "interval_start"] = self.interval_start + if self.aggregates: + output["aggregates"] = _dump_aggregate_results(self.aggregates, self.results, camel_case) + return output + + +class _RecordsBucketAggregateResult(RecordsAggregateResult): + _buckets_key: ClassVar[str] + + def __init__(self, buckets: Sequence[RecordsBucket]) -> None: + self._buckets = list(buckets) + + @property + def buckets(self) -> list[RecordsBucket]: + return list(self._buckets) + + @classmethod + def _load(cls, resource: dict[str, Any]) -> Self: + return cls(buckets=[RecordsBucket._load(bucket) for bucket in resource.get(cls._buckets_key, [])]) + + def dump(self, camel_case: bool = True) -> dict[str, Any]: + key = self._buckets_key if camel_case else to_snake_case(self._buckets_key) + return {key: [bucket.dump(camel_case=camel_case) for bucket in self._buckets]} + + +class RecordsUniqueValuesAggregateResult(_RecordsBucketAggregateResult): + _buckets_key = "uniqueValueBuckets" + + +class RecordsNumberHistogramAggregateResult(_RecordsBucketAggregateResult): + _buckets_key = "numberHistogramBuckets" + + +class RecordsTimeHistogramAggregateResult(_RecordsBucketAggregateResult): + _buckets_key = "timeHistogramBuckets" + + +class RecordsFilterAggregateResult(_RecordsBucketAggregateResult): + _buckets_key = "filterBuckets" + + +_METRIC_AGGREGATE_KEYS: frozenset[str] = frozenset({"avg", "count", "min", "max", "sum"}) + +_BUCKET_RESULT_BY_KEY: dict[str, type[_RecordsBucketAggregateResult]] = { + result_cls._buckets_key: result_cls + for result_cls in ( + RecordsUniqueValuesAggregateResult, + RecordsNumberHistogramAggregateResult, + RecordsTimeHistogramAggregateResult, + RecordsFilterAggregateResult, + ) +} + + +class RecordsAggregation(CogniteResource): + """Aggregate results returned from the Records aggregate endpoint. + + Args: + aggregates (dict[str, Any]): Aggregate results keyed by the client-defined aggregate IDs. + typing (TypeInformation | None): Optional property typing metadata. + """ + + def __init__(self, aggregates: dict[str, Any], typing: TypeInformation | None = None) -> None: + self.aggregates = aggregates + self.results = { + aggregate_id: RecordsAggregateResult._load(result) + for aggregate_id, result in aggregates.items() + if isinstance(result, dict) + } + self.typing = typing + + def __getitem__(self, aggregate_id: str) -> RecordsAggregateResult: + return self.results[aggregate_id] + + @classmethod + def _load(cls, resource: dict[str, Any]) -> Self: + return cls( + aggregates=resource["aggregates"], + typing=TypeInformation._load(resource["typing"]) if "typing" in resource else None, + ) + + def dump(self, camel_case: bool = True) -> dict[str, Any]: + output: dict[str, Any] = {"aggregates": _dump_aggregate_results(self.aggregates, self.results, camel_case)} + if self.typing is not None: + output["typing"] = self.typing.dump(camel_case=camel_case) + return output + + class Record(WriteableCogniteResource["RecordWrite"]): """A record returned from the stream records API. @@ -406,10 +920,6 @@ def __init__( self.has_next = has_next self.typing = typing - @classmethod - def _load_response(cls, response: dict[str, Any]) -> Self: - return cls._load_raw_api_response([response]) - @classmethod def _load_raw_api_response(cls, responses: list[dict[str, Any]]) -> Self: last_response = responses[-1] diff --git a/tests/tests_unit/test_api/test_data_modeling/test_records.py b/tests/tests_unit/test_api/test_data_modeling/test_records.py index c476636220..c0356021c6 100644 --- a/tests/tests_unit/test_api/test_data_modeling/test_records.py +++ b/tests/tests_unit/test_api/test_data_modeling/test_records.py @@ -6,21 +6,43 @@ from pytest_httpx import HTTPXMock from cognite.client import AsyncCogniteClient, CogniteClient +from cognite.client.data_classes import filters from cognite.client.data_classes.data_modeling.data_types import UnitReference from cognite.client.data_classes.data_modeling.instances import InstanceSort, TypeInformation from cognite.client.data_classes.data_modeling.records import ( + Average, + Count, + Filters, + Max, + Min, + MovingFunction, + MovingFunctions, + NumberHistogram, Record, RecordContainerId, RecordId, RecordList, + RecordsAggregate, + RecordsAggregation, + RecordsFilterAggregateResult, + RecordsMetricAggregateResult, + RecordsMovingFunctionAggregateResult, + RecordsNumberHistogramAggregateResult, RecordSource, RecordSourceSelector, + RecordsTimeHistogramAggregateResult, + RecordsUniqueValuesAggregateResult, + RecordsUnknownAggregateResult, RecordTargetUnit, RecordTargetUnits, RecordWrite, + Sum, SyncRecord, SyncRecordList, + TimeHistogram, TimeRange, + UniqueValues, + UnknownRecordsAggregate, ) from tests.utils import jsgz_load @@ -262,6 +284,480 @@ def test_upsert_chunks( assert len(jsgz_load(requests[1].content)["items"]) == 1 +class TestRecordsAPIAggregate: + def test_aggregate_posts_request_and_returns_wrapper( + self, + cognite_client: CogniteClient, + httpx_mock: HTTPXMock, + records_base_url: str, + stream_id: str, + ) -> None: + httpx_mock.add_response( + method="POST", + url=re.compile(re.escape(records_base_url) + r"/aggregate$"), + json={"aggregates": {"avg_temp": {"avg": 22.5}}}, + ) + out = cognite_client.data_modeling.records.aggregate( + stream_id=stream_id, + aggregates={"avg_temp": {"avg": {"property": ["sp", "container-x", "temp"]}}}, + last_updated_time=TimeRange(gte=1_000_000), + filter=filters.Equals(["space"], "sp"), + target_units=RecordTargetUnits(unit_system_name="SI"), + include_typing=True, + ) + + assert isinstance(out, RecordsAggregation) + assert out.aggregates == {"avg_temp": {"avg": 22.5}} + body = jsgz_load(httpx_mock.get_requests()[0].content) + assert body == { + "aggregates": {"avg_temp": {"avg": {"property": ["sp", "container-x", "temp"]}}}, + "lastUpdatedTime": {"gte": 1_000_000}, + "filter": {"equals": {"property": ["space"], "value": "sp"}}, + "targetUnits": {"unitSystemName": "SI"}, + "includeTyping": True, + } + + def test_aggregate_accepts_dict_filter( + self, + cognite_client: CogniteClient, + httpx_mock: HTTPXMock, + records_base_url: str, + stream_id: str, + ) -> None: + httpx_mock.add_response( + method="POST", + url=re.compile(re.escape(records_base_url) + r"/aggregate$"), + json={"aggregates": {"total": {"count": 7}}}, + ) + cognite_client.data_modeling.records.aggregate( + stream_id=stream_id, + aggregates={"total": {"count": {}}}, + filter={"matchAll": {}}, + ) + + body = jsgz_load(httpx_mock.get_requests()[0].content) + assert body["filter"] == {"matchAll": {}} + + def test_aggregate_accepts_mixed_typed_and_dict_aggregates( + self, + cognite_client: CogniteClient, + httpx_mock: HTTPXMock, + records_base_url: str, + stream_id: str, + ) -> None: + httpx_mock.add_response( + method="POST", + url=re.compile(re.escape(records_base_url) + r"/aggregate$"), + json={"aggregates": {"total": {"count": 7}}}, + ) + + cognite_client.data_modeling.records.aggregate( + stream_id=stream_id, + aggregates={ + "by_day": TimeHistogram( + ["sp", "container-x", "timestamp"], + calendar_interval="1d", + aggregates={ + "avg_temp": Average(["sp", "container-x", "temp"]), + "moving_count": MovingFunction( + buckets_path="_count", + window=3, + function="MovingFunctions.unweightedAvg", + ), + "raw_total": {"count": {}}, + }, + ), + "by_region": UniqueValues( + ["sp", "container-x", "region"], + aggregates={"max_temp": Max(["sp", "container-x", "temp"])}, + size=5, + ), + "salary_histogram": NumberHistogram( + ["sp", "container-x", "salary"], + interval=1000, + aggregates={"sum_salary": Sum(["sp", "container-x", "salary"])}, + hard_bounds={"min": 0, "max": 10000}, + ), + "by_filters": Filters( + filters=[ + filters.Range(["createdTime"], gte=1), + {"matchAll": {}}, + ], + aggregates={"total": Count()}, + ), + }, + ) + + body = jsgz_load(httpx_mock.get_requests()[0].content) + assert body["aggregates"] == { + "by_day": { + "timeHistogram": { + "property": ["sp", "container-x", "timestamp"], + "calendarInterval": "1d", + "aggregates": { + "avg_temp": {"avg": {"property": ["sp", "container-x", "temp"]}}, + "moving_count": { + "movingFunction": { + "bucketsPath": "_count", + "window": 3, + "function": "MovingFunctions.unweightedAvg", + } + }, + "raw_total": {"count": {}}, + }, + } + }, + "by_region": { + "uniqueValues": { + "property": ["sp", "container-x", "region"], + "aggregates": {"max_temp": {"max": {"property": ["sp", "container-x", "temp"]}}}, + "size": 5, + } + }, + "salary_histogram": { + "numberHistogram": { + "property": ["sp", "container-x", "salary"], + "interval": 1000, + "aggregates": {"sum_salary": {"sum": {"property": ["sp", "container-x", "salary"]}}}, + "hardBounds": {"min": 0, "max": 10000}, + } + }, + "by_filters": { + "filters": { + "filters": [ + {"range": {"property": ["createdTime"], "gte": 1}}, + {"matchAll": {}}, + ], + "aggregates": {"total": {"count": {}}}, + } + }, + } + + def test_records_aggregation_loads_typed_results(self) -> None: + loaded = RecordsAggregation._load( + { + "aggregates": { + "avg_temp": {"avg": 22.5}, + "by_region": { + "uniqueValueBuckets": [ + { + "value": "north", + "count": 2, + "aggregates": {"max_temp": {"max": 30.0}}, + } + ] + }, + "by_number": {"numberHistogramBuckets": [{"intervalStart": 0.0, "count": 1}]}, + "by_time": { + "timeHistogramBuckets": [ + { + "intervalStart": "2024-05-16T00:00:00Z", + "count": 3, + "aggregates": {"moving": {"fnValue": 7.5}}, + } + ] + }, + "by_filter": {"filterBuckets": [{"count": 4}]}, + "future": {"futureAggregateResult": 1}, + } + } + ) + + avg_temp = loaded["avg_temp"] + assert isinstance(avg_temp, RecordsMetricAggregateResult) + assert avg_temp.aggregate == "avg" + assert avg_temp.value == 22.5 + + by_region = loaded["by_region"] + assert isinstance(by_region, RecordsUniqueValuesAggregateResult) + assert by_region.buckets[0].value == "north" + max_temp = by_region.buckets[0].results["max_temp"] + assert isinstance(max_temp, RecordsMetricAggregateResult) + assert max_temp.value == 30.0 + + by_number = loaded["by_number"] + assert isinstance(by_number, RecordsNumberHistogramAggregateResult) + assert by_number.buckets[0].interval_start == 0.0 + + by_time = loaded["by_time"] + assert isinstance(by_time, RecordsTimeHistogramAggregateResult) + moving = by_time.buckets[0].results["moving"] + assert isinstance(moving, RecordsMovingFunctionAggregateResult) + assert moving.fn_value == 7.5 + + by_filter = loaded["by_filter"] + assert isinstance(by_filter, RecordsFilterAggregateResult) + assert by_filter.buckets[0].count == 4 + + assert isinstance(loaded["future"], RecordsUnknownAggregateResult) + assert loaded["future"].dump() == {"futureAggregateResult": 1} + + def test_aggregate_results_dump_honors_camel_case(self) -> None: + loaded = RecordsAggregation._load( + { + "aggregates": { + "avg_temp": {"avg": 22.5}, + "moving": {"fnValue": 7.5}, + "by_time": { + "timeHistogramBuckets": [ + { + "intervalStart": "2024-05-16T00:00:00Z", + "count": 3, + "aggregates": {"moving": {"fnValue": 1.5}}, + } + ] + }, + "future": {"futureAggregateResult": 1}, + } + } + ) + + # camel_case=True round-trips the API payload unchanged. + assert loaded["avg_temp"].dump() == {"avg": 22.5} + assert loaded["moving"].dump() == {"fnValue": 7.5} + assert loaded["by_time"].dump() == { + "timeHistogramBuckets": [ + { + "intervalStart": "2024-05-16T00:00:00Z", + "count": 3, + "aggregates": {"moving": {"fnValue": 1.5}}, + } + ] + } + assert loaded["future"].dump() == {"futureAggregateResult": 1} + + # camel_case=False snake-cases every API key while leaving client-defined IDs untouched. + assert loaded["moving"].dump(camel_case=False) == {"fn_value": 7.5} + assert loaded["by_time"].dump(camel_case=False) == { + "time_histogram_buckets": [ + { + "interval_start": "2024-05-16T00:00:00Z", + "count": 3, + "aggregates": {"moving": {"fn_value": 1.5}}, + } + ] + } + assert loaded["future"].dump(camel_case=False) == {"future_aggregate_result": 1} + + def test_aggregate_dump_preserves_client_defined_ids(self) -> None: + # The keys under "aggregates" are chosen by the caller (see the aggregate() examples), so + # they are user data, not API fields: dump() must echo them back verbatim regardless of + # camel_case. Only the API-defined field names inside each result get converted. + loaded = RecordsAggregation._load( + { + "aggregates": { + "myTopLevelAvg": {"avg": 22.5}, + "myTimeGroups": { + "timeHistogramBuckets": [ + { + "intervalStart": "2024-05-16T00:00:00Z", + "count": 3, + "aggregates": {"myNestedMoving": {"fnValue": 1.5}}, + } + ] + }, + "myFutureShape": {"futureAggregateResult": 1}, + } + } + ) + + assert loaded.dump(camel_case=False) == { + "aggregates": { + "myTopLevelAvg": {"avg": 22.5}, + "myTimeGroups": { + "time_histogram_buckets": [ + { + "interval_start": "2024-05-16T00:00:00Z", + "count": 3, + "aggregates": {"myNestedMoving": {"fn_value": 1.5}}, + } + ] + }, + "myFutureShape": {"future_aggregate_result": 1}, + } + } + + +class TestRecordsAggregateBuilders: + """Dump-only aggregate request builders. + + RecordsAggregate has no load/_load (aggregates are request-only bodies), so the dump output is + the contract worth pinning down. Each test asserts a builder serializes to the exact request + body the API expects. + """ + + def test_metric_aggregates_dump_name_and_property(self) -> None: + prop = ["sp", "c", "temp"] + assert Average(prop).dump() == {"avg": {"property": prop}} + assert Min(prop).dump() == {"min": {"property": prop}} + assert Max(prop).dump() == {"max": {"property": prop}} + assert Sum(prop).dump() == {"sum": {"property": prop}} + + def test_metric_aggregate_accepts_tuple_property(self) -> None: + assert Average(("sp", "c", "temp")).dump() == {"avg": {"property": ["sp", "c", "temp"]}} + + def test_count_without_property_dumps_empty_body(self) -> None: + assert Count().dump() == {"count": {}} + + def test_count_with_property(self) -> None: + assert Count(["sp", "c", "temp"]).dump() == {"count": {"property": ["sp", "c", "temp"]}} + + def test_unique_values_minimal(self) -> None: + assert UniqueValues(["sp", "c", "region"]).dump() == {"uniqueValues": {"property": ["sp", "c", "region"]}} + + def test_unique_values_with_size_and_nested_aggregates(self) -> None: + agg = UniqueValues( + ["sp", "c", "region"], + aggregates={"max_temp": Max(["sp", "c", "temp"])}, + size=5, + ) + assert agg.dump() == { + "uniqueValues": { + "property": ["sp", "c", "region"], + "aggregates": {"max_temp": {"max": {"property": ["sp", "c", "temp"]}}}, + "size": 5, + } + } + + def test_number_histogram_minimal(self) -> None: + assert NumberHistogram(["sp", "c", "salary"], interval=1000).dump() == { + "numberHistogram": {"property": ["sp", "c", "salary"], "interval": 1000} + } + + def test_number_histogram_with_bounds_and_nested_aggregates(self) -> None: + agg = NumberHistogram( + ["sp", "c", "salary"], + interval=1000, + aggregates={"sum_salary": Sum(["sp", "c", "salary"])}, + hard_bounds={"min": 0, "max": 10000}, + ) + assert agg.dump() == { + "numberHistogram": { + "property": ["sp", "c", "salary"], + "interval": 1000, + "aggregates": {"sum_salary": {"sum": {"property": ["sp", "c", "salary"]}}}, + "hardBounds": {"min": 0, "max": 10000}, + } + } + + def test_time_histogram_calendar_interval(self) -> None: + assert TimeHistogram(["sp", "c", "ts"], calendar_interval="1d").dump() == { + "timeHistogram": {"property": ["sp", "c", "ts"], "calendarInterval": "1d"} + } + + def test_time_histogram_fixed_interval_and_bounds(self) -> None: + agg = TimeHistogram( + ["sp", "c", "ts"], + fixed_interval="12h", + hard_bounds={"min": "2024-01-01", "max": "2024-02-01"}, + ) + assert agg.dump() == { + "timeHistogram": { + "property": ["sp", "c", "ts"], + "fixedInterval": "12h", + "hardBounds": {"min": "2024-01-01", "max": "2024-02-01"}, + } + } + + def test_time_histogram_requires_exactly_one_interval(self) -> None: + match = "Exactly one of calendar_interval or fixed_interval" + with pytest.raises(ValueError, match=match): + TimeHistogram(["sp", "c", "ts"]) # neither + with pytest.raises(ValueError, match=match): + TimeHistogram(["sp", "c", "ts"], calendar_interval="1d", fixed_interval="12h") # both + + def test_filters_with_filter_objects_and_raw_dicts(self) -> None: + agg = Filters( + filters=[filters.Range(["createdTime"], gte=1), {"matchAll": {}}], + aggregates={"total": Count()}, + ) + assert agg.dump() == { + "filters": { + "filters": [ + {"range": {"property": ["createdTime"], "gte": 1}}, + {"matchAll": {}}, + ], + "aggregates": {"total": {"count": {}}}, + } + } + + def test_moving_function(self) -> None: + agg = MovingFunction(buckets_path="_count", window=3, function="MovingFunctions.unweightedAvg") + assert agg.dump() == { + "movingFunction": { + "bucketsPath": "_count", + "window": 3, + "function": "MovingFunctions.unweightedAvg", + } + } + + def test_moving_function_accepts_enum_and_dumps_plain_string(self) -> None: + # The MovingFunctions enum guards against typos; either the enum or the raw literal works, + # and both dump to the plain wire string. + from_enum = MovingFunction("_count", 3, function=MovingFunctions.UNWEIGHTED_AVG) + from_str = MovingFunction("_count", 3, function="MovingFunctions.unweightedAvg") + assert from_enum.dump() == from_str.dump() + assert from_enum.dump()["movingFunction"]["function"] == "MovingFunctions.unweightedAvg" + + def test_moving_function_rejects_unknown_function(self) -> None: + with pytest.raises(ValueError): + MovingFunction("_count", 3, function="MovingFunctions.median") # type: ignore[arg-type] + + def test_nested_aggregates_accept_raw_dicts(self) -> None: + # Values under `aggregates` may be typed builders or raw dicts; both dump through. + agg = UniqueValues(["sp", "c", "region"], aggregates={"raw": {"count": {}}}) + assert agg.dump() == { + "uniqueValues": { + "property": ["sp", "c", "region"], + "aggregates": {"raw": {"count": {}}}, + } + } + + def test_dump_is_always_camel_case(self) -> None: + # Request builders always emit camelCase keys; the camel_case flag is a no-op for them. + agg = TimeHistogram(["sp", "c", "ts"], calendar_interval="1d") + expected = {"timeHistogram": {"property": ["sp", "c", "ts"], "calendarInterval": "1d"}} + assert agg.dump(camel_case=True) == expected + assert agg.dump(camel_case=False) == expected + + def test_eq(self) -> None: + assert Average(["sp", "c", "temp"]) == Average(["sp", "c", "temp"]) + assert Average(["sp", "c", "temp"]) != Average(["sp", "c", "other"]) + assert Average(["sp", "c", "temp"]) != Max(["sp", "c", "temp"]) + + def test_load_roundtrips_every_builder(self) -> None: + # v7 makes load public so request builders round-trip through config files: the reloaded + # object is the same type and dumps identically to the original. + prop = ["sp", "c", "temp"] + builders: list[RecordsAggregate] = [ + Average(prop), + Min(prop), + Max(prop), + Sum(prop), + Count(), + Count(prop), + UniqueValues(["sp", "c", "region"], aggregates={"m": Max(prop)}, size=5), + NumberHistogram(prop, interval=10, hard_bounds={"min": 0, "max": 100}), + TimeHistogram(["sp", "c", "ts"], calendar_interval="1d"), + TimeHistogram(["sp", "c", "ts"], fixed_interval="12h"), + Filters(filters=[filters.Range(["createdTime"], gte=1)], aggregates={"n": Count()}), + MovingFunction(buckets_path="games", window=7, function=MovingFunctions.UNWEIGHTED_AVG), + ] + for builder in builders: + reloaded = RecordsAggregate.load(builder.dump()) + assert type(reloaded) is type(builder) + assert reloaded.dump() == builder.dump() + + def test_load_unknown_aggregate_falls_back_to_unknown_records_aggregate(self) -> None: + # Unknown/newer aggregate types round-trip via UnknownRecordsAggregate (a real + # RecordsAggregate) rather than crashing, so the SDK can lag the API. + payload = {"futureAggregate": {"property": ["sp", "c", "x"]}} + loaded = RecordsAggregate.load(payload) + assert isinstance(loaded, UnknownRecordsAggregate) + assert isinstance(loaded, RecordsAggregate) + assert loaded.dump() == payload + + class TestRecordsAPIFilter: def test_list_returns_record_list( self, diff --git a/tests/tests_unit/test_base.py b/tests/tests_unit/test_base.py index e9d94a8384..0f1e384191 100644 --- a/tests/tests_unit/test_base.py +++ b/tests/tests_unit/test_base.py @@ -37,6 +37,14 @@ EdgeListWithCursor, NodeListWithCursor, ) +from cognite.client.data_classes.data_modeling.records import ( + Filters, + NumberHistogram, + RecordsAggregateResult, + RecordsBucket, + TimeHistogram, + UniqueValues, +) from cognite.client.data_classes.datapoints import DatapointsArray, SyntheticDatapoints from cognite.client.data_classes.datapoints_subscriptions import SubscriptionDatapoints from cognite.client.data_classes.events import EventList @@ -187,6 +195,18 @@ class TestCogniteResource: # UnknownWorkflowTaskParameters: Requires task_type which is only available in the parent object's # full response payload UnknownWorkflowTaskParameters, + # RecordsAggregateResult/RecordsBucket have non-standard constructors. The + # nested-aggregate builders (Filters/NumberHistogram/TimeHistogram/UniqueValues) + # take a `RecordsAggregate | dict` map the fake generator can't synthesize + # (abstract union), and TimeHistogram also has an interval guard; their + # load/dump round-trip is covered in test_records.py instead. + *all_concrete_subclasses(RecordsAggregateResult), + RecordsAggregateResult, + RecordsBucket, + Filters, + NumberHistogram, + TimeHistogram, + UniqueValues, }, ) ], @@ -218,6 +238,13 @@ def test_json_serialize( Agent, *all_concrete_subclasses(WorkflowTaskOutput), UnknownWorkflowTaskParameters, + *all_concrete_subclasses(RecordsAggregateResult), + RecordsAggregateResult, + RecordsBucket, + Filters, + NumberHistogram, + TimeHistogram, + UniqueValues, }, ) ], @@ -322,6 +349,13 @@ def test_writable_list_as_write( SubscriptionDatapoints, *all_concrete_subclasses(WorkflowTaskOutput), UnknownWorkflowTaskParameters, + *all_concrete_subclasses(RecordsAggregateResult), + RecordsAggregateResult, + RecordsBucket, + Filters, + NumberHistogram, + TimeHistogram, + UniqueValues, }, ) ], @@ -359,6 +393,13 @@ def test_load_has_no_side_effects( SubscriptionDatapoints, *all_concrete_subclasses(WorkflowTaskOutput), UnknownWorkflowTaskParameters, + *all_concrete_subclasses(RecordsAggregateResult), + RecordsAggregateResult, + RecordsBucket, + Filters, + NumberHistogram, + TimeHistogram, + UniqueValues, }, ) ], @@ -409,6 +450,13 @@ def _for_all_nested_dicts(obj: dict, func: Callable[[dict], None]) -> None: SubscriptionDatapoints, *all_concrete_subclasses(WorkflowTaskOutput), UnknownWorkflowTaskParameters, + *all_concrete_subclasses(RecordsAggregateResult), + RecordsAggregateResult, + RecordsBucket, + Filters, + NumberHistogram, + TimeHistogram, + UniqueValues, }, ) ],