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

cognite_toolkit/_cdf_tk/client/api/migration.py

@@ -0,0 +1,529 @@
+import warnings
+from collections.abc import Sequence
+from itertools import groupby
+from typing import Literal, TypeVar, cast, overload
+
+from cognite.client._constants import DEFAULT_LIMIT_READ
+from cognite.client.data_classes.data_modeling import (
+    Node,
+    NodeApplyResult,
+    NodeApplyResultList,
+    NodeId,
+    NodeList,
+    ViewId,
+    filters,
+    query,
+)
+from cognite.client.utils.useful_types import SequenceNotStr
+
+from cognite_toolkit._cdf_tk.client.data_classes.migration import (
+    AssetCentricId,
+    CreatedSourceSystem,
+    InstanceSource,
+    ResourceViewMapping,
+    ResourceViewMappingApply,
+    SpaceSource,
+)
+from cognite_toolkit._cdf_tk.constants import COGNITE_MIGRATION_SPACE
+from cognite_toolkit._cdf_tk.tk_warnings import HighSeverityWarning
+from cognite_toolkit._cdf_tk.utils.collection import chunker_sequence
+from cognite_toolkit._cdf_tk.utils.useful_types import AssetCentricType
+
+from .extended_data_modeling import ExtendedInstancesAPI
+
+
+class InstanceSourceAPI:
+    def __init__(self, instance_api: ExtendedInstancesAPI) -> None:
+        self._instance_api = instance_api
+        self._RETRIEVE_LIMIT = 1000
+        self._view_id = InstanceSource.get_source()
+
+    def retrieve(self, ids: Sequence[AssetCentricId]) -> NodeList[InstanceSource]:
+        """Retrieve a list of instance sources by their IDs.
+
+        Args:
+            ids (Sequence[AssetCentricId]): A sequence of AssetCentricId objects representing the IDs of the instance sources to retrieve.
+        """
+        results: NodeList[InstanceSource] = NodeList[InstanceSource]([])
+        for chunk in chunker_sequence(ids, self._RETRIEVE_LIMIT):
+            retrieve_query = query.Query(
+                with_={
+                    "instanceSource": query.NodeResultSetExpression(
+                        filter=filters.And(filters.HasData(views=[self._view_id]), self._create_dms_filter(ids)),
+                        limit=len(chunk),
+                    ),
+                },
+                select={"instanceSource": query.Select([query.SourceSelector(self._view_id, ["*"])])},
+            )
+            chunk_response = self._instance_api.query(retrieve_query)
+            results.extend([InstanceSource._load(item.dump()) for item in chunk_response.get("instanceSource", [])])
+        return results
+
+    @staticmethod
+    def _create_dms_filter(ids: Sequence[AssetCentricId]) -> filters.Filter:
+        """Create a filter that matches all the AssetCentricIds in the list."""
+        if not ids:
+            raise ValueError("Cannot create a filter from an empty AssetCentricIdList.")
+        to_or_filters: list[filters.Filter] = []
+        instance_source_view = InstanceSource.get_source()
+        for resource_type, resource_ids in groupby(
+            sorted(ids, key=lambda x: x.resource_type), key=lambda x: x.resource_type
+        ):
+            is_resource = filters.Equals(instance_source_view.as_property_ref("resourceType"), resource_type)
+            is_id = filters.In(
+                instance_source_view.as_property_ref("id"), [resource_id.id_ for resource_id in resource_ids]
+            )
+            to_or_filters.append(filters.And(is_resource, is_id))
+        return filters.Or(*to_or_filters)
+
+    def list(self, resource_type: AssetCentricType, limit: int | None = DEFAULT_LIMIT_READ) -> NodeList[InstanceSource]:
+        """List instance sources filtered by resource type"""
+        is_selected = filters.Equals(self._view_id.as_property_ref("resourceType"), resource_type)
+
+        nodes = self._instance_api.list(instance_type="node", filter=is_selected, limit=limit, sources=self._view_id)
+        return NodeList[InstanceSource]([InstanceSource._load(node.dump()) for node in nodes])
+
+
+class ResourceViewMappingAPI:
+    """API for retrieving resource view mapping from the data model."""
+
+    def __init__(self, instance_api: ExtendedInstancesAPI) -> None:
+        self._instance_api = instance_api
+        self._view_id = ResourceViewMapping.get_source()
+
+    @overload
+    def upsert(
+        self,
+        item: ResourceViewMappingApply,
+        skip_on_version_conflict: bool = False,
+        replace: bool = False,
+    ) -> NodeApplyResult: ...
+
+    @overload
+    def upsert(
+        self,
+        item: Sequence[ResourceViewMappingApply],
+        skip_on_version_conflict: bool = False,
+        replace: bool = False,
+    ) -> NodeApplyResultList: ...
+
+    def upsert(
+        self,
+        item: ResourceViewMappingApply | Sequence[ResourceViewMappingApply],
+        skip_on_version_conflict: bool = False,
+        replace: bool = False,
+    ) -> NodeApplyResult | NodeApplyResultList:
+        """Upsert one or more view sources."""
+        result = self._instance_api.apply(
+            item, skip_on_version_conflict=skip_on_version_conflict, replace=replace
+        ).nodes
+        if isinstance(item, ResourceViewMappingApply):
+            return result[0]
+        return NodeApplyResultList(result)
+
+    @overload
+    def retrieve(self, external_id: str) -> ResourceViewMapping | None: ...
+
+    @overload
+    def retrieve(self, external_id: SequenceNotStr[str]) -> NodeList[ResourceViewMapping]: ...
+
+    def retrieve(
+        self, external_id: str | SequenceNotStr[str]
+    ) -> ResourceViewMapping | NodeList[ResourceViewMapping] | None:
+        """Retrieve one or more view sources by their external IDs."""
+        if isinstance(external_id, str):
+            return self._instance_api.retrieve_nodes(
+                NodeId(COGNITE_MIGRATION_SPACE, external_id), node_cls=ResourceViewMapping
+            )
+        else:
+            nodes = self._instance_api.retrieve(
+                nodes=[NodeId(COGNITE_MIGRATION_SPACE, ext_id) for ext_id in external_id], sources=[self._view_id]
+            ).nodes
+            return self._safe_convert(nodes)
+
+    @overload
+    def delete(self, external_id: str) -> NodeId: ...
+
+    @overload
+    def delete(self, external_id: SequenceNotStr[str]) -> list[NodeId]: ...
+
+    def delete(self, external_id: str | SequenceNotStr[str]) -> NodeId | list[NodeId]:
+        """Delete a view source or a list of view sources by their external IDs."""
+        if isinstance(external_id, str):
+            return self._instance_api.delete(NodeId(COGNITE_MIGRATION_SPACE, external_id)).nodes[0]
+        else:
+            return self._instance_api.delete([NodeId(COGNITE_MIGRATION_SPACE, ext_id) for ext_id in external_id]).nodes
+
+    def list(
+        self, resource_type: str | None = None, limit: int | None = DEFAULT_LIMIT_READ
+    ) -> NodeList[ResourceViewMapping]:
+        """List view sources optionally filtered by resource type"""
+        is_selected: filters.Filter | None = None
+        if resource_type:
+            is_selected = filters.Equals(self._view_id.as_property_ref("resourceType"), resource_type)
+
+        nodes = self._instance_api.list(
+            instance_type="node", filter=is_selected, limit=limit, space=COGNITE_MIGRATION_SPACE, sources=self._view_id
+        )
+        return self._safe_convert(nodes)
+
+    @classmethod
+    def _safe_convert(cls, nodes: NodeList[Node]) -> NodeList[ResourceViewMapping]:
+        results = NodeList[ResourceViewMapping]([])
+        for node in nodes:
+            try:
+                loaded = ResourceViewMapping._load(node.dump())
+            except ValueError as e:
+                warnings.warn(
+                    HighSeverityWarning(
+                        f"Node {node.as_id()!r} is in an invalid format. Skipping it. Error: {e!s}",
+                    )
+                )
+                continue
+            results.append(loaded)
+        return results
+
+
+class CreatedSourceSystemAPI:
+    def __init__(self, instance_api: ExtendedInstancesAPI) -> None:
+        self._instance_api = instance_api
+        self._RETRIEVE_LIMIT = 1000
+        self._view_id = CreatedSourceSystem.get_source()
+
+    def retrieve(self, source: SequenceNotStr[str]) -> NodeList[CreatedSourceSystem]:
+        """Retrieve one or more view sources by their external IDs."""
+        results: NodeList[CreatedSourceSystem] = NodeList[CreatedSourceSystem]([])
+        # MyPy does not understand that SequenceNotStr is a sequence.
+        for chunk in chunker_sequence(source, self._RETRIEVE_LIMIT):  # type: ignore[type-var]
+            retrieve_query = query.Query(
+                with_={
+                    "sourceSystem": query.NodeResultSetExpression(
+                        filter=filters.And(filters.HasData(views=[self._view_id]), self._create_dms_filter(chunk)),
+                        limit=len(chunk),
+                    ),
+                },
+                select={"sourceSystem": query.Select([query.SourceSelector(self._view_id, ["*"])])},
+            )
+            chunk_response = self._instance_api.query(retrieve_query)
+            results.extend([CreatedSourceSystem._load(item.dump()) for item in chunk_response.get("sourceSystem", [])])
+        return results
+
+    def _create_dms_filter(self, source: SequenceNotStr[str]) -> filters.Filter:
+        """Create a filter that matches all CreatedSourceSystem with given source in the list."""
+        if not source:
+            raise ValueError("Cannot create a filter from an empty source list.")
+        return filters.In(self._view_id.as_property_ref("source"), list(source))
+
+    def list(self, limit: int = -1) -> NodeList[CreatedSourceSystem]:
+        """Lists all created source systems."""
+        return self._instance_api.list(instance_type=CreatedSourceSystem, limit=limit)
+
+
+_T = TypeVar("_T", bound=int | str)
+
+
+class SpaceSourceAPI:
+    def __init__(self, instance_api: ExtendedInstancesAPI) -> None:
+        self._instance_api = instance_api
+        self._RETRIEVE_LIMIT = 1000
+        self._view_id = SpaceSource.get_source()
+        # Cache for space sources by data set ID and external ID
+        self._cache_by_id: dict[int, SpaceSource] = {}
+        self._cache_by_external_id: dict[str, SpaceSource] = {}
+
+    @overload
+    def retrieve(self, data_set_id: int) -> SpaceSource | None: ...
+
+    @overload
+    def retrieve(self, data_set_id: Sequence[int]) -> NodeList[SpaceSource]: ...
+
+    @overload
+    def retrieve(self, *, data_set_external_id: str) -> SpaceSource | None: ...
+
+    @overload
+    def retrieve(self, *, data_set_external_id: SequenceNotStr[str]) -> NodeList[SpaceSource]: ...
+
+    def retrieve(
+        self,
+        data_set_id: int | Sequence[int] | None = None,
+        data_set_external_id: str | SequenceNotStr[str] | None = None,
+    ) -> SpaceSource | NodeList[SpaceSource] | None:
+        """Retrieve a space source by data set ID or external ID.
+
+        This method uses caching to avoid redundant API calls. If a space source
+        has been retrieved before, it will be returned from cache.
+        """
+        if data_set_id is not None and data_set_external_id is None:
+            return self._retrieve_with_cache(
+                value=data_set_id,
+                property_name="dataSetId",
+                cache=self._cache_by_id,
+                is_single=isinstance(data_set_id, int),
+            )
+        elif data_set_external_id is not None:
+            return self._retrieve_with_cache(
+                # MyPy is confused by SequenceNotStr
+                value=data_set_external_id,  # type: ignore[arg-type]
+                property_name="dataSetExternalId",
+                cache=self._cache_by_external_id,
+                is_single=isinstance(data_set_external_id, str),
+            )
+        else:
+            raise ValueError("Either data_set_id or data_set_external_id must be provided, but not both.")
+
+    def _retrieve_with_cache(
+        self,
+        value: _T | Sequence[_T],
+        property_name: str,
+        cache: dict[_T, SpaceSource],
+        is_single: bool,
+    ) -> SpaceSource | NodeList[SpaceSource] | None:
+        """Retrieve space sources with caching support."""
+        # MyPy does not understand that if is_single is True, value is _T, else Sequence[_T].
+        values: list[_T] = [value] if is_single else list(value)  # type: ignore[arg-type, list-item]
+
+        # Check cache for all requested values
+        cached_results: list[SpaceSource] = []
+        missing_values: list[_T] = []
+        for val in values:
+            if val in cache:
+                cached_results.append(cache[val])
+            else:
+                missing_values.append(val)
+
+        # Fetch missing values from API
+        if missing_values:
+            fetched = self._retrieve_by_property(property_name=property_name, values=missing_values)
+            for item in fetched:
+                self._add_to_cache(item)
+                cached_results.append(item)
+
+        if is_single:
+            return cached_results[0] if cached_results else None
+        return NodeList[SpaceSource](cached_results)
+
+    def _add_to_cache(self, space_source: SpaceSource) -> None:
+        """Add a space source to both caches."""
+        if space_source.data_set_id:
+            self._cache_by_id[space_source.data_set_id] = space_source
+        if space_source.data_set_external_id is not None:
+            self._cache_by_external_id[space_source.data_set_external_id] = space_source
+
+    def _retrieve_by_property(
+        self,
+        property_name: str,
+        values: Sequence[_T],
+    ) -> NodeList[SpaceSource]:
+        """Retrieve space sources by filtering on a specific property."""
+        results: NodeList[SpaceSource] = NodeList[SpaceSource]([])
+        for chunk in chunker_sequence(values, self._RETRIEVE_LIMIT):
+            retrieve_query = query.Query(
+                with_={
+                    "spaceSource": query.NodeResultSetExpression(
+                        filter=filters.And(
+                            filters.HasData(views=[self._view_id]),
+                            filters.In(self._view_id.as_property_ref(property_name), chunk),
+                        ),
+                        limit=len(chunk),
+                    ),
+                },
+                select={"spaceSource": query.Select([query.SourceSelector(self._view_id, ["*"])])},
+            )
+            chunk_response = self._instance_api.query(retrieve_query)
+            items = chunk_response.get("spaceSource", [])
+            results.extend([SpaceSource._load(item.dump()) for item in items])
+
+        return results
+
+    def list(self, limit: int = -1) -> NodeList[SpaceSource]:
+        """Lists all space sources and populates the cache.
+
+        Note: This method always fetches from the API and does not use the cache,
+        but it will populate the cache with all retrieved space sources.
+        """
+        results = self._instance_api.list(instance_type=SpaceSource, limit=limit)
+        # Populate cache with all retrieved space sources
+        for space_source in results:
+            self._add_to_cache(space_source)
+        return results
+
+
+_T_Cached = TypeVar("_T_Cached", bound=NodeId | ViewId)
+
+
+class LookupAPI:
+    def __init__(self, instance_api: ExtendedInstancesAPI, resource_type: AssetCentricType) -> None:
+        self._instance_api = instance_api
+        self._resource_type = resource_type
+        self._view_id = InstanceSource.get_source()
+        self._node_id_by_id: dict[int, NodeId | None] = {}
+        self._node_id_by_external_id: dict[str, NodeId | None] = {}
+        self._consumer_view_id_by_id: dict[int, ViewId | None] = {}
+        self._consumer_view_id_by_external_id: dict[str, ViewId | None] = {}
+        self._RETRIEVE_LIMIT = 1000
+
+    @overload
+    def __call__(self, id: int, external_id: None = None) -> NodeId | None: ...
+
+    @overload
+    def __call__(self, id: SequenceNotStr[int], external_id: None = None) -> dict[int, NodeId]: ...
+
+    @overload
+    def __call__(self, *, external_id: str) -> NodeId | None: ...
+
+    @overload
+    def __call__(self, *, external_id: SequenceNotStr[str]) -> dict[str, NodeId]: ...
+
+    def __call__(
+        self, id: int | SequenceNotStr[int] | None = None, external_id: str | SequenceNotStr[str] | None = None
+    ) -> dict[int, NodeId] | dict[str, NodeId] | NodeId | None:
+        """Lookup NodeId by either internal ID or external ID.
+
+        Args:
+            id (int | Sequence[int] | None): The internal ID(s) to lookup.
+            external_id (str | SequenceNotStr[str] | None): The external ID(s) to lookup.
+
+        Returns:
+            NodeId | dict[int, NodeId] | dict[str, NodeId] | None: The corresponding NodeId(s) if found, otherwise None.
+
+        """
+        if id is not None and external_id is None:
+            return self._lookup(
+                identifier=id,
+                cache=self._node_id_by_id,
+                property_name="id",
+                return_type=NodeId,
+                input_type=int,
+            )
+        elif external_id is not None and id is None:
+            return self._lookup(
+                identifier=external_id,
+                cache=self._node_id_by_external_id,
+                property_name="classicExternalId",
+                return_type=NodeId,
+                input_type=str,
+            )
+        else:
+            raise TypeError("Either id or external_id must be provided, but not both.")
+
+    @overload
+    def consumer_view(self, id: int, external_id: None = None) -> ViewId | None: ...
+
+    @overload
+    def consumer_view(self, id: SequenceNotStr[int], external_id: None = None) -> dict[int, ViewId]: ...
+
+    @overload
+    def consumer_view(self, *, external_id: str) -> ViewId | None: ...
+
+    @overload
+    def consumer_view(self, *, external_id: SequenceNotStr[str]) -> dict[str, ViewId]: ...
+
+    def consumer_view(
+        self, id: int | SequenceNotStr[int] | None = None, external_id: str | SequenceNotStr[str] | None = None
+    ) -> dict[int, ViewId] | dict[str, ViewId] | ViewId | None:
+        """Lookup Consumer ViewId by either internal ID or external ID.
+
+        Args:
+            id (int | Sequence[int] | None): The internal ID(s) to lookup.
+            external_id (str | SequenceNotStr[str] | None): The external ID(s) to lookup.
+        Returns:
+            ViewId | dict[int, ViewId] | dict[str, ViewId] | None
+        """
+        if id is not None and external_id is None:
+            return self._lookup(
+                identifier=id,
+                cache=self._consumer_view_id_by_id,
+                property_name="id",
+                return_type=ViewId,
+                input_type=int,
+            )
+        elif external_id is not None and id is None:
+            return self._lookup(
+                identifier=external_id,
+                cache=self._consumer_view_id_by_external_id,
+                property_name="classicExternalId",
+                return_type=ViewId,
+                input_type=str,
+            )
+        else:
+            raise TypeError("Either id or external_id must be provided, but not both.")
+
+    def _lookup(
+        self,
+        identifier: _T | SequenceNotStr[_T],
+        cache: dict[_T, _T_Cached | None],
+        property_name: Literal["id", "classicExternalId"],
+        return_type: type[_T_Cached],
+        input_type: type[_T],
+    ) -> dict[_T, _T_Cached] | _T_Cached | None:
+        """Generic lookup method for both NodeId and ViewId by id or external_id."""
+        is_single = isinstance(identifier, input_type)
+        # MyPy does not understand that if is_single is True, identifier is _T, else SequenceNotStr[_T].
+        identifiers: list[_T] = [identifier] if is_single else list(identifier)  # type: ignore[arg-type, list-item]
+
+        missing = [id_ for id_ in identifiers if id_ not in cache]
+        if missing:
+            self._fetch_and_cache(missing, by=property_name)
+
+        if is_single:
+            return cache.get(identifier)  # type: ignore[arg-type]
+
+        return {id_: value for id_ in identifiers if isinstance(value := cache.get(id_), return_type)}
+
+    def _fetch_and_cache(self, identifiers: Sequence[int | str], by: Literal["id", "classicExternalId"]) -> None:
+        for chunk in chunker_sequence(identifiers, self._RETRIEVE_LIMIT):
+            retrieve_query = query.Query(
+                with_={
+                    "instanceSource": query.NodeResultSetExpression(
+                        filter=filters.And(
+                            filters.HasData(views=[self._view_id]),
+                            filters.Equals(self._view_id.as_property_ref("resourceType"), self._resource_type),
+                            filters.In(self._view_id.as_property_ref(by), list(chunk)),
+                        ),
+                        limit=len(chunk),
+                    ),
+                },
+                select={"instanceSource": query.Select([query.SourceSelector(self._view_id, ["*"])])},
+            )
+            chunk_response = self._instance_api.query(retrieve_query)
+            items = chunk_response.get("instanceSource", [])
+            for item in items:
+                instance_source = InstanceSource._load(item.dump())
+                node_id = instance_source.as_id()
+                self._node_id_by_id[instance_source.id_] = node_id
+                self._consumer_view_id_by_id[instance_source.id_] = instance_source.consumer_view()
+                if instance_source.classic_external_id:
+                    self._node_id_by_external_id[instance_source.classic_external_id] = node_id
+                    self._consumer_view_id_by_external_id[instance_source.classic_external_id] = (
+                        instance_source.consumer_view()
+                    )
+            missing = set(chunk) - set(self._node_id_by_id.keys()) - set(self._node_id_by_external_id.keys())
+            if by == "id":
+                for missing_id in cast(set[int], missing):
+                    if missing_id not in self._node_id_by_id:
+                        self._node_id_by_id[missing_id] = None
+                    if missing_id not in self._consumer_view_id_by_id:
+                        self._consumer_view_id_by_id[missing_id] = None
+            elif by == "classicExternalId":
+                for missing_ext_id in cast(set[str], missing):
+                    if missing_ext_id not in self._node_id_by_external_id:
+                        self._node_id_by_external_id[missing_ext_id] = None
+                    if missing_ext_id not in self._consumer_view_id_by_external_id:
+                        self._consumer_view_id_by_external_id[missing_ext_id] = None
+
+
+class MigrationLookupAPI:
+    def __init__(self, instance_api: ExtendedInstancesAPI) -> None:
+        self.assets = LookupAPI(instance_api, "asset")
+        self.events = LookupAPI(instance_api, "event")
+        self.files = LookupAPI(instance_api, "file")
+        self.time_series = LookupAPI(instance_api, "timeseries")
+
+
+class MigrationAPI:
+    def __init__(self, instance_api: ExtendedInstancesAPI) -> None:
+        self.instance_source = InstanceSourceAPI(instance_api)
+        self.resource_view_mapping = ResourceViewMappingAPI(instance_api)
+        self.created_source_system = CreatedSourceSystemAPI(instance_api)
+        self.space_source = SpaceSourceAPI(instance_api)
+        self.lookup = MigrationLookupAPI(instance_api)

cognite_toolkit/_cdf_tk/client/api/project.py

@@ -0,0 +1,24 @@
+from cognite.client import CogniteClient
+
+from cognite_toolkit._cdf_tk.client.config import ToolkitClientConfig
+from cognite_toolkit._cdf_tk.client.data_classes.project import ProjectStatusList
+from cognite_toolkit._cdf_tk.utils.http_client import HTTPClient, RequestMessage2
+
+
+class ProjectAPI:
+    def __init__(self, config: ToolkitClientConfig, cognite_client: CogniteClient) -> None:
+        self._config = config
+        self._cognite_client = cognite_client
+        self._http_client = HTTPClient(config, split_items_status_codes=set())
+
+    def status(self) -> ProjectStatusList:
+        """Retrieve information about the current project."""
+        response = self._http_client.request_single_retries(
+            RequestMessage2(
+                endpoint_url=f"{self._config.base_url}/api/v1/projects",
+                method="GET",
+                parameters={"withDataModelingStatus": True},
+            )
+        )
+        success = response.get_success_or_raise()
+        return ProjectStatusList._load(success.body_json["items"], cognite_client=self._cognite_client)

cognite_toolkit/_cdf_tk/client/api/robotics/api.py

@@ -1,5 +1,3 @@
-from __future__ import annotations
-
 from cognite.client import ClientConfig, CogniteClient
 from cognite.client._api_client import APIClient
 

cognite_toolkit/_cdf_tk/client/api/robotics/capabilities.py

@@ -1,5 +1,3 @@
-from __future__ import annotations
-
 from collections.abc import Iterator, Sequence
 from typing import overload
 

cognite_toolkit/_cdf_tk/client/api/robotics/data_postprocessing.py

@@ -1,5 +1,3 @@
-from __future__ import annotations
-
 from collections.abc import Iterator, Sequence
 from typing import overload
 

cognite_toolkit/_cdf_tk/client/api/robotics/frames.py

@@ -1,5 +1,3 @@
-from __future__ import annotations
-
 from collections.abc import Iterator, Sequence
 from typing import overload
 

cognite_toolkit/_cdf_tk/client/api/robotics/locations.py

@@ -1,5 +1,3 @@
-from __future__ import annotations
-
 from collections.abc import Iterator, Sequence
 from typing import overload
 

cognite_toolkit/_cdf_tk/client/api/robotics/maps.py

@@ -1,5 +1,3 @@
-from __future__ import annotations
-
 from collections.abc import Iterator, Sequence
 from typing import overload
 

cognite_toolkit/_cdf_tk/client/api/robotics/robots.py

@@ -1,5 +1,3 @@
-from __future__ import annotations
-
 from collections.abc import Iterator, Sequence
 from typing import overload
 

cognite_toolkit/_cdf_tk/client/api/search.py

@@ -0,0 +1,13 @@
+from cognite.client import CogniteClient
+from cognite.client._api_client import APIClient
+from cognite.client.config import ClientConfig
+
+from .location_filters import LocationFiltersAPI
+from .search_config import SearchConfigurationsAPI
+
+
+class SearchAPI(APIClient):
+    def __init__(self, config: ClientConfig, api_version: str | None, cognite_client: CogniteClient) -> None:
+        super().__init__(config, api_version, cognite_client)
+        self.locations = LocationFiltersAPI(config, api_version, cognite_client)
+        self.configurations = SearchConfigurationsAPI(config, api_version, cognite_client)

cognite_toolkit/_cdf_tk/client/api/search_config.py

@@ -0,0 +1,49 @@
+from urllib.parse import urljoin
+
+from cognite.client import CogniteClient
+from cognite.client._api_client import APIClient
+from cognite.client.config import ClientConfig
+
+from cognite_toolkit._cdf_tk.client.data_classes.search_config import SearchConfig, SearchConfigList, SearchConfigWrite
+
+
+class SearchConfigurationsAPI(APIClient):
+    def __init__(self, config: ClientConfig, api_version: str | None, cognite_client: CogniteClient) -> None:
+        super().__init__(config, api_version, cognite_client)
+
+    def _get_base_url_with_base_path(self) -> str:
+        """
+        This method is overridden to provide the correct base path for the Search Configurations API.
+        This method in base class APIClient appends /api/{api_version}/ to the base URL,
+        but for Search Configurations API, we need a different path structure.
+        """
+        base_path = ""
+        if self._api_version:
+            base_path = f"/apps/{self._api_version}/projects/{self._config.project}/storage/config/apps/search/views"
+        return urljoin(self._config.base_url, base_path)
+
+    def upsert(self, configuration_update: SearchConfigWrite) -> SearchConfig:
+        """Update/Create a Configuration.
+
+        Args:
+            configuration_update: The content of the configuration to update
+
+        Returns:
+            SearchConfig
+
+        """
+        res = self._post(
+            url_path="/upsert",
+            json=configuration_update.dump(),
+        )
+        return SearchConfig._load(res.json(), cognite_client=self._cognite_client)
+
+    def list(self) -> SearchConfigList:
+        """List all Configuration.
+
+        Returns:
+            SearchConfigList
+
+        """
+        res = self._post(url_path="/list")
+        return SearchConfigList._load(res.json()["items"], cognite_client=self._cognite_client)