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

cognite_toolkit/_cdf_tk/client/api/extended_timeseries.py

@@ -1,6 +1,7 @@
 from collections.abc import Sequence
-from typing import Any, overload
+from typing import Any, cast, overload
 
+from cognite.client import ClientConfig, CogniteClient
 from cognite.client._api.time_series import SortSpec, TimeSeriesAPI
 from cognite.client._constants import DEFAULT_LIMIT_READ
 from cognite.client.data_classes.data_modeling import NodeId
@@ -19,6 +20,11 @@ from cognite_toolkit._cdf_tk.client.data_classes.pending_instances_ids import Pe
 class ExtendedTimeSeriesAPI(TimeSeriesAPI):
     """Extended TimeSeriesAPI to include pending ID methods."""
 
+    def __init__(self, config: ClientConfig, api_version: str | None, cognite_client: CogniteClient) -> None:
+        super().__init__(config, api_version, cognite_client)
+        self._PENDING_IDS_LIMIT = 1000
+        self._UNLINK_LIMIT = 1000
+
     @overload
     def set_pending_ids(
         self, instance_id: NodeId | tuple[str, str], id: int | None = None, external_id: str | None = None
@@ -74,7 +80,7 @@ class ExtendedTimeSeriesAPI(TimeSeriesAPI):
                 },
                 "api_subversion": "alpha",
             }
-            for id_chunk in split_into_chunks(list(identifiers), 1000)
+            for id_chunk in split_into_chunks(list(identifiers), self._PENDING_IDS_LIMIT)
         ]
         tasks_summary = execute_tasks(
             self._post,
@@ -90,6 +96,68 @@ class ExtendedTimeSeriesAPI(TimeSeriesAPI):
 
         return ExtendedTimeSeriesList._load(retrieved_items, cognite_client=self._cognite_client)
 
+    @overload
+    def unlink_instance_ids(
+        self,
+        id: int | None = None,
+        external_id: str | None = None,
+    ) -> ExtendedTimeSeries | None: ...
+
+    @overload
+    def unlink_instance_ids(
+        self,
+        id: Sequence[int] | None = None,
+        external_id: SequenceNotStr[str] | None = None,
+    ) -> ExtendedTimeSeriesList: ...
+
+    def unlink_instance_ids(
+        self,
+        id: int | Sequence[int] | None = None,
+        external_id: str | SequenceNotStr[str] | None = None,
+    ) -> ExtendedTimeSeries | ExtendedTimeSeriesList | None:
+        """Unlink pending instance IDs from time series.
+
+        Args:
+            id (int | Sequence[int] | None): The ID(s) of the time series.
+            external_id (str | SequenceNotStr[str] | None): The external ID(s) of the time series.
+
+        """
+        if id is None and external_id is None:
+            return None
+        if isinstance(id, int) and isinstance(external_id, str):
+            raise ValueError("Cannot specify both id and external_id as single values. Use one or the other.")
+        is_single = isinstance(id, int) or isinstance(external_id, str)
+        identifiers = IdentifierSequence.load(id, external_id)
+
+        tasks = [
+            {
+                "url_path": f"{self._RESOURCE_PATH}/unlink-instance-ids",
+                "json": {"items": id_chunk},
+                "api_subversion": "alpha",
+            }
+            for id_chunk in split_into_chunks(identifiers.as_dicts(), self._UNLINK_LIMIT)
+        ]
+        tasks_summary = execute_tasks(
+            self._post,
+            tasks,
+            max_workers=self._config.max_workers,
+            fail_fast=True,
+        )
+        tasks_summary.raise_compound_exception_if_failed_tasks(
+            task_unwrap_fn=unpack_items_in_payload,
+        )
+
+        retrieved_items = tasks_summary.joined_results(lambda res: res.json()["items"])
+
+        result = ExtendedTimeSeriesList._load(retrieved_items, cognite_client=self._cognite_client)
+        if is_single:
+            if len(result) == 0:
+                return None
+            if len(result) > 1:
+                raise ValueError("Expected a single time series, but multiple were returned.")
+            return cast(ExtendedTimeSeries, result[0])
+        return result
+
     def retrieve(
         self, id: int | None = None, external_id: str | None = None, instance_id: NodeId | None = None
     ) -> ExtendedTimeSeries | None:

cognite_toolkit/_cdf_tk/client/api/fixed_transformations.py

@@ -0,0 +1,122 @@
+import gzip
+from typing import Any
+
+from cognite.client._api.transformations import TransformationsAPI
+from cognite.client.config import global_config
+from cognite.client.data_classes import TransformationPreviewResult
+from cognite.client.utils import _json
+from requests import Response
+
+
+class FixedTransformationsAPI(TransformationsAPI):
+    """This overrides the preview method in the transformations API to allow for longer timeouts and retries."""
+
+    _default_timeout_run_query = 240  # seconds, this is the maximum timeout for running queries in CDF
+
+    def preview(
+        self,
+        query: str | None = None,
+        convert_to_string: bool = False,
+        limit: int | None = 100,
+        source_limit: int | None = 100,
+        infer_schema_limit: int | None = 10_000,
+        timeout: int | None = _default_timeout_run_query,
+    ) -> TransformationPreviewResult:
+        """`Preview the result of a query. <https://developer.cognite.com/api#tag/Query/operation/runPreview>`_
+
+        This method overrides the default preview method in the TransformationsAPI to allow for longer
+        timeouts and retries.
+
+        Toolkit runs long-running queries that takes longer than the typical default of 30 seconds. In addition,
+        we do not want to retry, which typically up to 10 times, as the user will have to wait for a long time. Instead,
+        we want to fail provide the user with the error and then let the user decide whether to retry or not by
+        running the CLI command again.
+
+        Args:
+            query (str | None): SQL query to run for preview.
+            convert_to_string (bool): Stringify values in the query results, default is False.
+            limit (int | None): Maximum number of rows to return in the final result, default is 100.
+            source_limit (int | None): Maximum number of items to read from the data source or None to run without limit, default is 100.
+            infer_schema_limit (int | None): Limit for how many rows that are used for inferring result schema, default is 10 000.
+            timeout (int | None): Number of seconds to wait before cancelling a query. The default, and maximum, is 240.
+
+        Returns:
+            TransformationPreviewResult: Result of the executed query
+        """
+        request_body = {
+            "query": query,
+            "convertToString": convert_to_string,
+            "limit": limit,
+            "sourceLimit": source_limit,
+            "inferSchemaLimit": infer_schema_limit,
+            "timeout": timeout,
+        }
+        response = self._do_requests_retry(
+            "POST",
+            "/transformations/query/run",
+            json=request_body,
+            timeout=timeout or (self._default_timeout_run_query + 60),  # add a buffer to the timeout
+            retry=False,
+        )
+        return TransformationPreviewResult._load(response.json(), cognite_client=self._cognite_client)
+
+    def _do_requests_retry(
+        self,
+        method: str,
+        url_path: str,
+        accept: str = "application/json",
+        api_subversion: str | None = None,
+        retry: bool = True,
+        **kwargs: Any,
+    ) -> Response:
+        """This is a copy of the _do_requests method from the APIClient with the option to turn off retries."""
+        is_retryable, full_url = self._resolve_url(method, url_path)
+        json_payload = kwargs.pop("json", None)
+        headers = self._configure_headers(
+            accept,
+            additional_headers=self._config.headers.copy(),
+            api_subversion=api_subversion,
+        )
+        headers.update(kwargs.get("headers") or {})
+
+        if json_payload is not None:
+            try:
+                data = _json.dumps(json_payload, allow_nan=False)
+            except ValueError as e:
+                # A lot of work to give a more human friendly error message when nans and infs are present:
+                msg = "Out of range float values are not JSON compliant"
+                if msg in str(e):  # exc. might e.g. contain an extra ": nan", depending on build (_json.make_encoder)
+                    raise ValueError(f"{msg}. Make sure your data does not contain NaN(s) or +/- Inf!").with_traceback(
+                        e.__traceback__
+                    ) from None
+                raise
+            kwargs["data"] = data
+            if method in ["PUT", "POST"] and not global_config.disable_gzip:
+                kwargs["data"] = gzip.compress(data.encode())
+                headers["Content-Encoding"] = "gzip"
+
+        kwargs["headers"] = headers
+
+        # requests will by default follow redirects. This can be an SSRF-hazard if
+        # the client can be tricked to request something with an open redirect, in
+        # addition to leaking the token, as requests will send the headers to the
+        # redirected-to endpoint.
+        # If redirects are to be followed in a call, this should be opted into instead.
+        kwargs.setdefault("allow_redirects", False)
+
+        if is_retryable and retry:
+            res = self._http_client_with_retry.request(method=method, url=full_url, **kwargs)
+        else:
+            res = self._http_client.request(method=method, url=full_url, **kwargs)
+
+        match res.status_code:
+            case 200 | 201 | 202 | 204:
+                pass
+            case 401:
+                self._raise_no_project_access_error(res)
+            case _:
+                self._raise_api_error(res, payload=json_payload)
+
+        stream = kwargs.get("stream")
+        self._log_request(res, payload=json_payload, stream=stream)
+        return res

cognite_toolkit/_cdf_tk/client/api/infield.py

@@ -0,0 +1,254 @@
+from collections.abc import Sequence
+from typing import Any, cast
+
+from pydantic import TypeAdapter
+from rich.console import Console
+
+from cognite_toolkit._cdf_tk.client.data_classes.api_classes import PagedResponse, QueryResponse
+from cognite_toolkit._cdf_tk.client.data_classes.infield import (
+    DataExplorationConfig,
+    InFieldCDMLocationConfig,
+    InfieldLocationConfig,
+)
+from cognite_toolkit._cdf_tk.client.data_classes.instance_api import (
+    InstanceResponseItem,
+    InstanceResult,
+    NodeIdentifier,
+)
+from cognite_toolkit._cdf_tk.tk_warnings import HighSeverityWarning
+from cognite_toolkit._cdf_tk.utils.http_client import (
+    HTTPClient,
+    ItemsRequest,
+    ItemsRequest2,
+    RequestMessage2,
+)
+
+
+class InfieldConfigAPI:
+    ENDPOINT = "/models/instances"
+    LOCATION_REF = "locationConfig"
+    EXPLORATION_REF = "explorerConfig"
+    # We know that this key exists and it has alias set.
+    DATA_EXPLORATION_PROP_ID = cast(str, InfieldLocationConfig.model_fields["data_exploration_config"].alias)
+
+    def __init__(self, http_client: HTTPClient, console: Console) -> None:
+        self._http_client = http_client
+        self._console = console
+        self._config = http_client.config
+
+    def apply(self, items: Sequence[InfieldLocationConfig]) -> list[InstanceResult]:
+        if len(items) > 500:
+            raise ValueError("Cannot apply more than 500 InfieldLocationConfig items at once.")
+
+        request_items = (
+            [item.as_request_item()]
+            if item.data_exploration_config is None
+            else [item.as_request_item(), item.data_exploration_config.as_request_item()]
+            for item in items
+        )
+        responses = self._http_client.request_items_retries(
+            ItemsRequest2(
+                endpoint_url=self._config.create_api_url(self.ENDPOINT),
+                method="POST",
+                items=[item for sublist in request_items for item in sublist],
+            )
+        )
+        responses.raise_for_status()
+        return TypeAdapter(list[InstanceResult]).validate_python(responses.get_items())
+
+    def retrieve(self, items: Sequence[NodeIdentifier]) -> list[InfieldLocationConfig]:
+        if len(items) > 100:
+            raise ValueError("Cannot retrieve more than 100 InfieldLocationConfig items at once.")
+        if not items:
+            return []
+        response = self._http_client.request_single_retries(
+            RequestMessage2(
+                # We use the query endpoint to be able to retrieve linked DataExplorationConfig items
+                endpoint_url=self._config.create_api_url(f"{self.ENDPOINT}/query"),
+                method="POST",
+                body_content=self._retrieve_query(items),
+            )
+        )
+        success = response.get_success_or_raise()
+        parsed_response = QueryResponse[InstanceResponseItem].model_validate(success.body_json)
+        return self._parse_retrieve_response(parsed_response)
+
+    def delete(self, items: Sequence[InfieldLocationConfig]) -> list[NodeIdentifier]:
+        if len(items) > 500:
+            raise ValueError("Cannot delete more than 500 InfieldLocationConfig items at once.")
+
+        identifiers = (
+            [item.as_id()]
+            if item.data_exploration_config is None
+            else [item.as_id(), item.data_exploration_config.as_id()]
+            for item in items
+        )
+        responses = self._http_client.request_with_retries(
+            ItemsRequest(
+                endpoint_url=self._config.create_api_url(f"{self.ENDPOINT}/delete"),
+                method="POST",
+                items=[identifier for sublist in identifiers for identifier in sublist],
+            )
+        )
+        responses.raise_for_status()
+        return PagedResponse[NodeIdentifier].model_validate(responses.get_first_body()).items
+
+    @classmethod
+    def _retrieve_query(cls, items: Sequence[NodeIdentifier]) -> dict[str, Any]:
+        return {
+            "with": {
+                cls.LOCATION_REF: {
+                    "limit": len(items),
+                    "nodes": {
+                        "filter": {
+                            "instanceReferences": [
+                                {"space": item.space, "externalId": item.external_id} for item in items
+                            ]
+                        },
+                    },
+                },
+                cls.EXPLORATION_REF: {
+                    "nodes": {
+                        "from": "locationConfig",
+                        "direction": "outwards",
+                        "through": {
+                            "source": InfieldLocationConfig.VIEW_ID.dump(),
+                            "identifier": cls.DATA_EXPLORATION_PROP_ID,
+                        },
+                    }
+                },
+            },
+            "select": {
+                cls.LOCATION_REF: {
+                    "sources": [{"source": InfieldLocationConfig.VIEW_ID.dump(), "properties": ["*"]}],
+                },
+                cls.EXPLORATION_REF: {
+                    "sources": [{"source": DataExplorationConfig.VIEW_ID.dump(), "properties": ["*"]}],
+                },
+            },
+        }
+
+    def _parse_retrieve_response(
+        self, parsed_response: QueryResponse[InstanceResponseItem]
+    ) -> list[InfieldLocationConfig]:
+        data_exploration_results = (
+            DataExplorationConfig.model_validate(
+                item.get_properties_for_source(DataExplorationConfig.VIEW_ID, include_identifier=True)
+            )
+            for item in parsed_response.items[self.EXPLORATION_REF]
+        )
+        data_exploration_config_map = {(dec.space, dec.external_id): dec for dec in data_exploration_results}
+        result: list[InfieldLocationConfig] = []
+        for item in parsed_response.items[self.LOCATION_REF]:
+            properties = item.get_properties_for_source(InfieldLocationConfig.VIEW_ID, include_identifier=True)
+            data_exploration = properties.pop(self.DATA_EXPLORATION_PROP_ID, None)
+            if isinstance(data_exploration, dict):
+                space = data_exploration["space"]
+                external_id = data_exploration["externalId"]
+                if (space, external_id) not in data_exploration_config_map:
+                    HighSeverityWarning(
+                        f"{self.DATA_EXPLORATION_PROP_ID} with space '{space}' and externalId '{external_id}' referenced in InfieldLocationConfig '{properties['externalId']}' was not found in the retrieved results."
+                    ).print_warning(console=self._console)
+                else:
+                    # Pydantic allow already validated models to be assigned to fields
+                    properties[self.DATA_EXPLORATION_PROP_ID] = data_exploration_config_map[(space, external_id)]  # type: ignore[assignment,index]
+            else:
+                HighSeverityWarning(
+                    f"InfieldLocationConfig '{properties['externalId']}' is missing a valid {self.DATA_EXPLORATION_PROP_ID} reference."
+                ).print_warning(console=self._console)
+            result.append(InfieldLocationConfig.model_validate(properties))
+        return result
+
+
+class InFieldCDMConfigAPI:
+    ENDPOINT = "/models/instances"
+    LOCATION_REF = "cdmLocationConfig"
+
+    def __init__(self, http_client: HTTPClient, console: Console) -> None:
+        self._http_client = http_client
+        self._console = console
+        self._config = http_client.config
+
+    def apply(self, items: Sequence[InFieldCDMLocationConfig]) -> list[InstanceResult]:
+        if len(items) > 500:
+            raise ValueError("Cannot apply more than 500 InFieldCDMLocationConfig items at once.")
+
+        request_items = [item.as_request_item() for item in items]
+        results = self._http_client.request_items_retries(
+            ItemsRequest2(
+                endpoint_url=self._config.create_api_url(self.ENDPOINT),
+                method="POST",
+                items=request_items,
+            )
+        )
+        results.raise_for_status()
+        return TypeAdapter(list[InstanceResult]).validate_python(results.get_items())
+
+    def retrieve(self, items: Sequence[NodeIdentifier]) -> list[InFieldCDMLocationConfig]:
+        if len(items) > 100:
+            raise ValueError("Cannot retrieve more than 100 InFieldCDMLocationConfig items at once.")
+        if not items:
+            return []
+        result = self._http_client.request_single_retries(
+            RequestMessage2(
+                endpoint_url=self._config.create_api_url(f"{self.ENDPOINT}/query"),
+                method="POST",
+                body_content=self._retrieve_query(items),
+            )
+        )
+        success = result.get_success_or_raise()
+        parsed_response = QueryResponse[InstanceResponseItem].model_validate(success.body_json)
+        return self._parse_retrieve_response(parsed_response)
+
+    def delete(self, items: Sequence[InFieldCDMLocationConfig]) -> list[NodeIdentifier]:
+        if len(items) > 500:
+            raise ValueError("Cannot delete more than 500 InFieldCDMLocationConfig items at once.")
+
+        identifiers: Sequence = [item.as_id() for item in items]
+        responses = self._http_client.request_with_retries(
+            ItemsRequest(
+                endpoint_url=self._config.create_api_url(f"{self.ENDPOINT}/delete"),
+                method="POST",
+                items=identifiers,  # type: ignore[arg-type]
+            )
+        )
+        responses.raise_for_status()
+        return PagedResponse[NodeIdentifier].model_validate(responses.get_first_body()).items
+
+    @classmethod
+    def _retrieve_query(cls, items: Sequence[NodeIdentifier]) -> dict[str, Any]:
+        return {
+            "with": {
+                cls.LOCATION_REF: {
+                    "limit": len(items),
+                    "nodes": {
+                        "filter": {
+                            "instanceReferences": [
+                                {"space": item.space, "externalId": item.external_id} for item in items
+                            ]
+                        },
+                    },
+                },
+            },
+            "select": {
+                cls.LOCATION_REF: {
+                    "sources": [{"source": InFieldCDMLocationConfig.VIEW_ID.dump(), "properties": ["*"]}],
+                },
+            },
+        }
+
+    def _parse_retrieve_response(
+        self, parsed_response: QueryResponse[InstanceResponseItem]
+    ) -> list[InFieldCDMLocationConfig]:
+        result: list[InFieldCDMLocationConfig] = []
+        for item in parsed_response.items[self.LOCATION_REF]:
+            properties = item.get_properties_for_source(InFieldCDMLocationConfig.VIEW_ID, include_identifier=True)
+            result.append(InFieldCDMLocationConfig.model_validate(properties))
+        return result
+
+
+class InfieldAPI:
+    def __init__(self, http_client: HTTPClient, console: Console) -> None:
+        self._http_client = http_client
+        self.config = InfieldConfigAPI(http_client, console)
+        self.cdm_config = InFieldCDMConfigAPI(http_client, console)

cognite_toolkit/_cdf_tk/client/api/location_filters.py

@@ -1,11 +1,10 @@
-from __future__ import annotations
-
-from collections.abc import Iterator
+from collections.abc import Iterator, Sequence
 from typing import overload
 
 from cognite.client import CogniteClient
 from cognite.client._api_client import APIClient
 from cognite.client.config import ClientConfig
+from cognite.client.utils._identifier import IdentifierSequence
 
 from cognite_toolkit._cdf_tk.client.data_classes.location_filters import (
     LocationFilter,
@@ -94,3 +93,20 @@ class LocationFiltersAPI(APIClient):
     def list(self) -> LocationFilterList:
         res = self._post(url_path=self._RESOURCE_PATH + "/list", json={"flat": True})
         return LocationFilterList._load(res.json()["items"], cognite_client=self._cognite_client)
+
+    def retrieve_multiple(self, ids: Sequence[int]) -> LocationFilterList:
+        """Retrieve multiple LocationFilters using their IDs.
+
+        Args:
+            ids: The IDs of the LocationFilters
+
+        Returns:
+            LocationFilterList
+        """
+        identifiers = IdentifierSequence.load(ids=ids)
+        return self._retrieve_multiple(
+            identifiers=identifiers,
+            resource_cls=LocationFilter,
+            list_cls=LocationFilterList,
+            resource_path=self._RESOURCE_PATH,
+        )