cognite-toolkit 0.5.44__py3-none-any.whl → 0.7.34__py3-none-any.whl

cognite_toolkit/_cdf_tk/cruds/_resource_cruds/timeseries.py

@@ -0,0 +1,500 @@
+import json
+from collections.abc import Hashable, Iterable, Sequence
+from itertools import zip_longest
+from pathlib import Path
+from typing import Any, Literal, cast, final
+
+from cognite.client.data_classes import (
+    DatapointSubscription,
+    DatapointSubscriptionList,
+    DataPointSubscriptionUpdate,
+    DataPointSubscriptionWrite,
+    TimeSeries,
+    TimeSeriesList,
+    TimeSeriesWrite,
+)
+from cognite.client.data_classes.capabilities import (
+    Capability,
+    TimeSeriesAcl,
+    TimeSeriesSubscriptionsAcl,
+)
+from cognite.client.data_classes.data_modeling import NodeId
+from cognite.client.data_classes.datapoints_subscriptions import TimeSeriesIDList
+from cognite.client.exceptions import CogniteAPIError, CogniteNotFoundError
+from cognite.client.utils.useful_types import SequenceNotStr
+
+from cognite_toolkit._cdf_tk.constants import MAX_TIMESTAMP_MS, MIN_TIMESTAMP_MS
+from cognite_toolkit._cdf_tk.cruds._base_cruds import ResourceContainerCRUD, ResourceCRUD
+from cognite_toolkit._cdf_tk.exceptions import (
+    ToolkitRequiredValueError,
+    ToolkitValueError,
+)
+from cognite_toolkit._cdf_tk.resource_classes import DatapointSubscriptionYAML, TimeSeriesYAML
+from cognite_toolkit._cdf_tk.utils import calculate_hash
+from cognite_toolkit._cdf_tk.utils.collection import chunker
+from cognite_toolkit._cdf_tk.utils.diff_list import diff_list_hashable, diff_list_identifiable, dm_identifier
+from cognite_toolkit._cdf_tk.utils.text import suffix_description
+
+from .auth import GroupAllScopedCRUD, SecurityCategoryCRUD
+from .classic import AssetCRUD
+from .data_organization import DataSetsCRUD
+
+
+@final
+class TimeSeriesCRUD(ResourceContainerCRUD[str, TimeSeriesWrite, TimeSeries]):
+    item_name = "datapoints"
+    folder_name = "timeseries"
+    resource_cls = TimeSeries
+    resource_write_cls = TimeSeriesWrite
+    yaml_cls = TimeSeriesYAML
+    kind = "TimeSeries"
+    dependencies = frozenset({DataSetsCRUD, GroupAllScopedCRUD, AssetCRUD})
+    _doc_url = "Time-series/operation/postTimeSeries"
+
+    @property
+    def display_name(self) -> str:
+        return "time series"
+
+    @classmethod
+    def get_required_capability(
+        cls, items: Sequence[TimeSeriesWrite] | None, read_only: bool
+    ) -> Capability | list[Capability]:
+        if not items and items is not None:
+            return []
+
+        actions = [TimeSeriesAcl.Action.Read] if read_only else [TimeSeriesAcl.Action.Read, TimeSeriesAcl.Action.Write]
+
+        scope: TimeSeriesAcl.Scope.All | TimeSeriesAcl.Scope.DataSet = TimeSeriesAcl.Scope.All()  # type: ignore[valid-type]
+        if items:
+            if dataset_ids := {item.data_set_id for item in items if item.data_set_id}:
+                scope = TimeSeriesAcl.Scope.DataSet(list(dataset_ids))
+
+        return TimeSeriesAcl(actions, scope)
+
+    @classmethod
+    def get_id(cls, item: TimeSeries | TimeSeriesWrite | dict) -> str:
+        if isinstance(item, dict):
+            return item["externalId"]
+        if item.external_id is None:
+            raise ToolkitRequiredValueError("TimeSeries must have external_id set.")
+        return item.external_id
+
+    @classmethod
+    def get_internal_id(cls, item: TimeSeries | dict) -> int:
+        if isinstance(item, dict):
+            return item["id"]
+        return item.id
+
+    @classmethod
+    def dump_id(cls, id: str) -> dict[str, Any]:
+        return {"externalId": id}
+
+    @classmethod
+    def get_dependent_items(cls, item: dict) -> Iterable[tuple[type[ResourceCRUD], Hashable]]:
+        if "dataSetExternalId" in item:
+            yield DataSetsCRUD, item["dataSetExternalId"]
+        if "securityCategoryNames" in item:
+            for security_category in item["securityCategoryNames"]:
+                yield SecurityCategoryCRUD, security_category
+        if "assetExternalId" in item:
+            yield AssetCRUD, item["assetExternalId"]
+
+    def load_resource(self, resource: dict[str, Any], is_dry_run: bool = False) -> TimeSeriesWrite:
+        if ds_external_id := resource.pop("dataSetExternalId", None):
+            resource["dataSetId"] = self.client.lookup.data_sets.id(ds_external_id, is_dry_run)
+        if security_categories_names := resource.pop("securityCategoryNames", []):
+            resource["securityCategories"] = self.client.lookup.security_categories.id(
+                security_categories_names, is_dry_run
+            )
+        if asset_external_id := resource.pop("assetExternalId", None):
+            resource["assetId"] = self.client.lookup.assets.id(asset_external_id, is_dry_run)
+        return TimeSeriesWrite._load(resource)
+
+    def dump_resource(self, resource: TimeSeries, local: dict[str, Any] | None = None) -> dict[str, Any]:
+        dumped = resource.as_write().dump()
+        if data_set_id := dumped.pop("dataSetId", None):
+            dumped["dataSetExternalId"] = self.client.lookup.data_sets.external_id(data_set_id)
+        if security_categories := dumped.pop("securityCategories", []):
+            dumped["securityCategoryNames"] = self.client.lookup.security_categories.external_id(security_categories)
+        if asset_id := dumped.pop("assetId", None):
+            dumped["assetExternalId"] = self.client.lookup.assets.external_id(asset_id)
+        return dumped
+
+    def create(self, items: Sequence[TimeSeriesWrite]) -> TimeSeriesList:
+        return self.client.time_series.create(items)
+
+    def retrieve(self, ids: SequenceNotStr[str | int]) -> TimeSeriesList:
+        internal_ids, external_ids = self._split_ids(ids)
+        return self.client.time_series.retrieve_multiple(
+            ids=internal_ids, external_ids=external_ids, ignore_unknown_ids=True
+        )
+
+    def update(self, items: Sequence[TimeSeriesWrite]) -> TimeSeriesList:
+        return self.client.time_series.update(items, mode="replace")
+
+    def delete(self, ids: SequenceNotStr[str | int]) -> int:
+        existing = self.retrieve(ids)
+        if existing:
+            self.client.time_series.delete(id=existing.as_ids(), ignore_unknown_ids=True)
+        return len(existing)
+
+    def _iterate(
+        self,
+        data_set_external_id: str | None = None,
+        space: str | None = None,
+        parent_ids: list[Hashable] | None = None,
+    ) -> Iterable[TimeSeries]:
+        return iter(
+            self.client.time_series(data_set_external_ids=[data_set_external_id] if data_set_external_id else None)
+        )
+
+    def count(self, ids: str | dict[str, Any] | SequenceNotStr[str | dict[str, Any]] | None) -> int:
+        datapoints = self.client.time_series.data.retrieve(
+            external_id=ids,  # type: ignore[arg-type]
+            start=MIN_TIMESTAMP_MS,
+            end=MAX_TIMESTAMP_MS + 1,
+            aggregates="count",
+            granularity="1000d",
+            ignore_unknown_ids=True,
+        )
+        return sum(sum(data.count or []) for data in datapoints)  # type: ignore[union-attr, misc, arg-type]
+
+    def drop_data(self, ids: SequenceNotStr[str] | None) -> int:
+        count = self.count(ids)
+        existing = self.client.time_series.retrieve_multiple(
+            external_ids=cast(SequenceNotStr[str], ids), ignore_unknown_ids=True
+        ).as_external_ids()
+        for external_id in existing:
+            self.client.time_series.data.delete_range(
+                external_id=external_id, start=MIN_TIMESTAMP_MS, end=MAX_TIMESTAMP_MS + 1
+            )
+        return count
+
+
+@final
+class DatapointSubscriptionCRUD(
+    ResourceCRUD[
+        str,
+        DataPointSubscriptionWrite,
+        DatapointSubscription,
+    ]
+):
+    folder_name = "timeseries"
+    resource_cls = DatapointSubscription
+    resource_write_cls = DataPointSubscriptionWrite
+    kind = "DatapointSubscription"
+    _doc_url = "Data-point-subscriptions/operation/postSubscriptions"
+    dependencies = frozenset(
+        {
+            TimeSeriesCRUD,
+            GroupAllScopedCRUD,
+        }
+    )
+    yaml_cls = DatapointSubscriptionYAML
+
+    _hash_key = "cdf-hash"
+    _description_character_limit = 1000
+    # A datapoint subscription can hold 10,000 timeseries, but the API
+    # only supports 100 timeseries per request. Thus, if a subscription
+    # has more than 100 timeseries, we need to split it into multiple requests.
+    _TIMESERIES_ID_REQUEST_LIMIT = 100
+    _MAX_TIMESERIES_IDS = 10_000
+
+    @property
+    def display_name(self) -> str:
+        return "timeseries subscriptions"
+
+    @classmethod
+    def get_id(cls, item: DataPointSubscriptionWrite | DatapointSubscription | dict) -> str:
+        if isinstance(item, dict):
+            return item["externalId"]
+        return item.external_id
+
+    @classmethod
+    def dump_id(cls, id: str) -> dict[str, Any]:
+        return {"externalId": id}
+
+    @classmethod
+    def get_dependent_items(cls, item: dict) -> Iterable[tuple[type[ResourceCRUD], Hashable]]:
+        if "dataSetExternalId" in item:
+            yield DataSetsCRUD, item["dataSetExternalId"]
+        for timeseries_id in item.get("timeSeriesIds", []):
+            yield TimeSeriesCRUD, timeseries_id
+
+    @classmethod
+    def get_required_capability(
+        cls, items: Sequence[DataPointSubscriptionWrite] | None, read_only: bool
+    ) -> Capability | list[Capability]:
+        if not items and items is not None:
+            return []
+
+        actions = (
+            [TimeSeriesSubscriptionsAcl.Action.Read]
+            if read_only
+            else [TimeSeriesSubscriptionsAcl.Action.Read, TimeSeriesSubscriptionsAcl.Action.Write]
+        )
+
+        scope: TimeSeriesSubscriptionsAcl.Scope.All | TimeSeriesSubscriptionsAcl.Scope.DataSet = (  # type: ignore[valid-type]
+            TimeSeriesSubscriptionsAcl.Scope.All()
+        )
+        if items:
+            if data_set_ids := {item.data_set_id for item in items if item.data_set_id}:
+                scope = TimeSeriesSubscriptionsAcl.Scope.DataSet(list(data_set_ids))
+
+        return TimeSeriesSubscriptionsAcl(actions, scope)
+
+    def create(self, items: Sequence[DataPointSubscriptionWrite]) -> DatapointSubscriptionList:
+        created_list = DatapointSubscriptionList([])
+        for item in items:
+            to_create, batches = self.create_split_timeseries_ids(item)
+            created = self.client.time_series.subscriptions.create(to_create)
+            for batch_item in batches:
+                created = self.client.time_series.subscriptions.update(batch_item)
+            created_list.append(created)
+        return created_list
+
+    def retrieve(self, ids: SequenceNotStr[str]) -> DatapointSubscriptionList:
+        items = DatapointSubscriptionList([])
+        for id_ in ids:
+            retrieved = self.client.time_series.subscriptions.retrieve(id_)
+            if retrieved:
+                items.append(retrieved)
+        return items
+
+    def update(self, items: Sequence[DataPointSubscriptionWrite]) -> DatapointSubscriptionList:
+        updated_list = DatapointSubscriptionList([])
+        for item in items:
+            current = self.client.time_series.subscriptions.list_member_time_series(item.external_id, limit=-1)
+            to_update, batches = self.update_split_timeseries_ids(item, current)
+            # There are two versions of a TimeSeries Subscription, one selects timeseries based filter
+            # and the other selects timeseries based on timeSeriesIds. If we use mode='replace', we try
+            # to set timeSeriesIds to an empty list, while the filter is set. This will result in an error.
+            updated = self.client.time_series.subscriptions.update(to_update, mode="replace_ignore_null")
+            for batch_item in batches:
+                updated = self.client.time_series.subscriptions.update(batch_item)
+            updated_list.append(updated)
+
+        return updated_list
+
+    def delete(self, ids: SequenceNotStr[str]) -> int:
+        try:
+            self.client.time_series.subscriptions.delete(ids)
+        except (CogniteAPIError, CogniteNotFoundError) as e:
+            non_existing = set(e.failed or [])
+            if existing := [id_ for id_ in ids if id_ not in non_existing]:
+                self.client.time_series.subscriptions.delete(existing)
+            return len(existing)
+        else:
+            # All deleted successfully
+            return len(ids)
+
+    def _iterate(
+        self,
+        data_set_external_id: str | None = None,
+        space: str | None = None,
+        parent_ids: list[Hashable] | None = None,
+    ) -> Iterable[DatapointSubscription]:
+        return iter(self.client.time_series.subscriptions)
+
+    def load_resource_file(
+        self, filepath: Path, environment_variables: dict[str, str | None] | None = None
+    ) -> list[dict[str, Any]]:
+        resources = super().load_resource_file(filepath, environment_variables)
+        for resource in resources:
+            if "timeSeriesIds" not in resource and "instanceIds" not in resource:
+                continue
+            # If the timeSeriesIds or instanceIds is set, we need to add the auth hash to the description.
+            # such that we can detect if the subscription has changed.
+            content: dict[str, object] = {}
+            if "timeSeriesIds" in resource:
+                content["timeSeriesIds"] = resource["timeSeriesIds"]
+            if "instanceIds" in resource:
+                content["instanceIds"] = resource["instanceIds"]
+            timeseries_hash = calculate_hash(json.dumps(content), shorten=True)
+            extra_str = f"{self._hash_key}: {timeseries_hash}"
+            resource["description"] = suffix_description(
+                extra_str,
+                resource.get("description"),
+                self._description_character_limit,
+                self.get_id(resource),
+                self.display_name,
+                self.console,
+            )
+
+        return resources
+
+    def load_resource(self, resource: dict[str, Any], is_dry_run: bool = False) -> DataPointSubscriptionWrite:
+        if ds_external_id := resource.pop("dataSetExternalId", None):
+            resource["dataSetId"] = self.client.lookup.data_sets.id(ds_external_id, is_dry_run)
+        return DataPointSubscriptionWrite._load(resource)
+
+    def dump_resource(self, resource: DatapointSubscription, local: dict[str, Any] | None = None) -> dict[str, Any]:
+        if resource.filter is not None:
+            dumped = resource.as_write().dump()
+        else:
+            # If filter is not set, the subscription uses explicit timeSeriesIds, which are not returned in the
+            # response. Calling .as_write() in this case raises ValueError because either filter or
+            # timeSeriesIds must be set.
+            dumped = resource.dump()
+            for server_prop in ("createdTime", "lastUpdatedTime", "timeSeriesCount"):
+                dumped.pop(server_prop, None)
+        local = local or {}
+        if data_set_id := dumped.pop("dataSetId", None):
+            dumped["dataSetExternalId"] = self.client.lookup.data_sets.external_id(data_set_id)
+        # timeSeriesIds and instanceIds are not returned in the response, so we need to add them
+        # to the dumped resource if they are set in the local resource. If there is a discrepancy between
+        # the local and dumped resource, th hash added to the description will change.
+        if "timeSeriesIds" in local:
+            dumped["timeSeriesIds"] = local["timeSeriesIds"]
+        if "instanceIds" in local:
+            dumped["instanceIds"] = local["instanceIds"]
+        return dumped
+
+    def diff_list(
+        self, local: list[Any], cdf: list[Any], json_path: tuple[str | int, ...]
+    ) -> tuple[dict[int, int], list[int]]:
+        if json_path[0] == "filter" or json_path == ("timeSeriesIds",):
+            return diff_list_hashable(local, cdf)
+        elif json_path == ("instanceIds",):
+            return diff_list_identifiable(local, cdf, get_identifier=dm_identifier)
+        return super().diff_list(local, cdf, json_path)
+
+    @classmethod
+    def create_split_timeseries_ids(
+        cls, subscription: DataPointSubscriptionWrite
+    ) -> tuple[DataPointSubscriptionWrite, list[DataPointSubscriptionUpdate]]:
+        """Split the time series IDs into batches of 100.
+        This is needed because the API only supports 100 time series IDs per request.
+
+        Note this is the total of time series IDs and instance IDs.
+        """
+        total_timeseries = len(subscription.time_series_ids or []) + len(subscription.instance_ids or [])
+        cls._validate_total_below_limit(subscription, total_timeseries)
+        if total_timeseries <= cls._TIMESERIES_ID_REQUEST_LIMIT:
+            # If the subscription has less than or equal to 100 time series IDs, we can return it as is.
+            # No need to split into batches.
+            return subscription, []
+
+        # Serialization to create a copy of the subscription
+        to_create = DataPointSubscriptionWrite.load(subscription.dump())
+        all_timeseries_ids = to_create.time_series_ids or []
+        all_instance_ids = to_create.instance_ids or []
+
+        # Create the first batch for the create/update call, prioritizing time_series_ids
+        to_create.time_series_ids = all_timeseries_ids[: cls._TIMESERIES_ID_REQUEST_LIMIT]
+        space_in_first_batch = cls._TIMESERIES_ID_REQUEST_LIMIT - len(to_create.time_series_ids)
+        to_create.instance_ids = all_instance_ids[:space_in_first_batch]
+
+        # Prepare remaining IDs for update batches
+        remaining_timeseries_ids = all_timeseries_ids[len(to_create.time_series_ids) :]
+        remaining_instance_ids = all_instance_ids[len(to_create.instance_ids) :]
+
+        # Using a list of tuples to preserve the type of ID
+        all_remaining_ids = [("ts", id) for id in remaining_timeseries_ids] + [
+            ("instance", id) for id in remaining_instance_ids
+        ]
+
+        batches: list[DataPointSubscriptionUpdate] = []
+        for chunk in chunker(all_remaining_ids, cls._TIMESERIES_ID_REQUEST_LIMIT):
+            update = DataPointSubscriptionUpdate(external_id=subscription.external_id)
+            ts_ids_in_chunk, instance_ids_in_chunk = cls._split_ts_instance_ids(chunk)
+            if ts_ids_in_chunk:
+                update.time_series_ids.add(ts_ids_in_chunk)
+            if instance_ids_in_chunk:
+                update.instance_ids.add(instance_ids_in_chunk)
+            batches.append(update)
+        return to_create, batches
+
+    @classmethod
+    def _validate_total_below_limit(cls, subscription: DataPointSubscriptionWrite, total_timeseries: int) -> None:
+        if total_timeseries > cls._MAX_TIMESERIES_IDS:
+            raise ToolkitValueError(
+                f'Subscription "{subscription.external_id}" has {total_timeseries:,} time series, '
+                f"which is more than the limit of {cls._MAX_TIMESERIES_IDS:,}."
+            )
+
+    @classmethod
+    def _split_ts_instance_ids(
+        cls, ids: list[tuple[Literal["ts"], str] | tuple[Literal["instance"], NodeId]]
+    ) -> tuple[list[str], list[NodeId]]:
+        ts_ids, instance_ids = [], []
+        for id_type, identifier in ids:
+            if id_type == "ts":
+                ts_ids.append(identifier)
+            else:
+                instance_ids.append(identifier)
+        # MyPy fails to understand the logic above ensures that ts_ids is a list of str
+        # and instance_ids is a list of NodeId.
+        return ts_ids, instance_ids  # type: ignore[return-value]
+
+    @classmethod
+    def update_split_timeseries_ids(
+        cls, subscription: DataPointSubscriptionWrite, current_ts: TimeSeriesIDList
+    ) -> tuple[DataPointSubscriptionWrite, list[DataPointSubscriptionUpdate]]:
+        """Split the time series IDs into batches of 100.
+        This is needed because the API only supports 100 time series IDs per request.
+
+        Note this is the total of time series IDs and instance IDs.
+
+        In addition, this method will compare to the current times series IDs and
+        update the subscription with adding and removing time series IDs as needed.
+        """
+        if subscription.time_series_ids is None and subscription.instance_ids is None:
+            # The subscription is using a filter, so we can return it as is.
+            return subscription, []
+        total_timeseries = len(subscription.time_series_ids or []) + len(subscription.instance_ids or [])
+        cls._validate_total_below_limit(subscription, total_timeseries)
+
+        # Serialization to create a copy of the subscription
+        to_update = DataPointSubscriptionWrite.load(subscription.dump())
+
+        # Get desired time series IDs
+        desired_timeseries_ids = set(to_update.time_series_ids or [])
+        desired_instance_ids = set(to_update.instance_ids or [])
+
+        # Get current time series IDs from the subscription
+        current_timeseries_ids: set[str] = set()
+        current_instance_ids: set[NodeId] = set()
+        for ts in current_ts:
+            if ts.external_id and ts.instance_id is None:
+                current_timeseries_ids.add(ts.external_id)
+            elif ts.instance_id and ts.external_id is None:
+                current_instance_ids.add(ts.instance_id)
+            else:
+                raise ValueError(f"External ID and instance ID are both set for time series {ts.external_id}.")
+
+        # Calculate what needs to be added and removed
+        ts_to_add = desired_timeseries_ids - current_timeseries_ids
+        ts_to_remove = current_timeseries_ids - desired_timeseries_ids
+        instance_to_add = desired_instance_ids - current_instance_ids
+        instance_to_remove = current_instance_ids - desired_instance_ids
+
+        # Clear the time series IDs from the main update to avoid conflicts
+        # The to_update object is used to update all other properties of the subscription.
+        to_update.time_series_ids = None
+        to_update.instance_ids = None
+
+        # Create update batches for changes
+        batches: list[DataPointSubscriptionUpdate] = []
+        all_removals = [("ts", id_) for id_ in ts_to_remove] + [("instance", id_) for id_ in instance_to_remove]
+        all_additions = [("ts", id_) for id_ in ts_to_add] + [("instance", id_) for id_ in instance_to_add]
+        for removals, additions in zip_longest(
+            chunker(all_removals, cls._TIMESERIES_ID_REQUEST_LIMIT),
+            chunker(all_additions, cls._TIMESERIES_ID_REQUEST_LIMIT),
+            fillvalue=None,
+        ):
+            update = DataPointSubscriptionUpdate(external_id=subscription.external_id)
+            ts_ids_to_remove, instance_ids_to_remove = cls._split_ts_instance_ids(removals or [])
+            if ts_ids_to_remove:
+                update.time_series_ids.remove(ts_ids_to_remove)
+            if instance_ids_to_remove:
+                update.instance_ids.remove(instance_ids_to_remove)
+
+            ts_ids_to_add, instance_ids_to_add = cls._split_ts_instance_ids(additions or [])
+            if ts_ids_to_add:
+                update.time_series_ids.add(ts_ids_to_add)
+            if instance_ids_to_add:
+                update.instance_ids.add(instance_ids_to_add)
+            batches.append(update)
+
+        return to_update, batches