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

cognite_toolkit/_cdf_tk/client/api/charts.py

@@ -0,0 +1,125 @@
+from collections.abc import Sequence
+from typing import overload
+from urllib.parse import urljoin
+
+from cognite.client._api_client import APIClient
+from cognite.client.utils._identifier import IdentifierSequence
+from cognite.client.utils.useful_types import SequenceNotStr
+
+from cognite_toolkit._cdf_tk.client.data_classes.charts import (
+    Chart,
+    ChartList,
+    ChartWrite,
+    Visibility,
+)
+from cognite_toolkit._cdf_tk.utils.collection import chunker
+
+
+class ChartsAPI(APIClient):
+    _RESOURCE_PATH = "/storage/charts/charts"
+
+    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 Charts API, we need a different path structure.
+        """
+        base_path = ""
+        if self._api_version:
+            base_path = f"/apps/{self._api_version}/projects/{self._config.project}"
+        return urljoin(self._config.base_url, base_path)
+
+    @overload
+    def upsert(self, items: ChartWrite) -> Chart: ...
+
+    @overload
+    def upsert(self, items: Sequence[ChartWrite]) -> ChartList: ...
+
+    def upsert(self, items: ChartWrite | Sequence[ChartWrite]) -> Chart | ChartList:
+        """Upsert one or more charts in CDF.
+
+        Args:
+            items (ChartWrite | Sequence[ChartWrite]): Chart(s) to upsert. Can be a single ChartWrite or a sequence of ChartWrite.
+
+        Returns:
+            Chart | ChartList: List of upserted charts if multiple items are provided, otherwise a single Chart.
+        """
+        item_sequence = items if isinstance(items, Sequence) else [items]
+        result = ChartList([])
+        # We are avoiding concurrency here as the Charts backend it not necessarily designed for it.
+        for chunk in chunker(item_sequence, self._CREATE_LIMIT):
+            body = {"items": [chunk_item.as_write().dump() for chunk_item in chunk]}
+            response = self._put(
+                url_path=self._RESOURCE_PATH,
+                json=body,
+            )
+            result.extend([Chart._load(item, cognite_client=self._cognite_client) for item in response.json()["items"]])
+
+        if not result:
+            raise ValueError("No charts were upserted. This may indicate an issue with the upsert endpoint.")
+
+        if isinstance(items, ChartWrite):
+            if len(result) != 1:
+                raise ValueError(
+                    "Expected a single chart to be returned, but multiple charts were returned. "
+                    "This may indicate an issue with the upsert operation."
+                )
+            return result[0]
+        elif len(result) != len(items):
+            raise ValueError(
+                "The number of upserted charts does not match the number of input charts. "
+                "This may indicate an issue with the upsert endpoint."
+            )
+        return result
+
+    @overload
+    def retrieve(self, external_id: str) -> Chart: ...
+
+    @overload
+    def retrieve(self, external_id: SequenceNotStr[str]) -> ChartList: ...
+
+    def retrieve(self, external_id: str | SequenceNotStr[str]) -> Chart | ChartList:
+        """Retrieve one or more charts by their external ID.
+
+        Args:
+            external_id (str | Sequence[str]): External ID(s) of the chart(s) to retrieve.
+
+        Returns:
+            Chart | ChartList: A single Chart if a single external ID is provided, otherwise a ChartList.
+        """
+        return self._retrieve_multiple(
+            list_cls=ChartList, resource_cls=Chart, identifiers=IdentifierSequence.load(external_ids=external_id)
+        )
+
+    def delete(self, external_id: str | SequenceNotStr[str]) -> None:
+        """Delete one or more charts by their external ID.
+
+        Args:
+            external_id (str | Sequence[str]): External ID(s) of the chart(s) to delete.
+        """
+        self._delete_multiple(
+            identifiers=IdentifierSequence.load(external_ids=external_id),
+            wrap_ids=True,
+        )
+
+    def list(self, visibility: Visibility | None = None, is_owned: bool | None = None) -> ChartList:
+        """List charts based on visibility and ownership.
+
+        Args:
+            visibility (Visibility): Visibility of the charts to list, either 'PUBLIC' or 'PRIVATE'.
+            is_owned (bool): Whether to list only owned charts.
+
+        Returns:
+            ChartList: List of charts matching the criteria.
+        """
+        filter_: dict[str, str | bool] = {}
+        if visibility is not None:
+            filter_["visibility"] = visibility.upper()
+        if is_owned is not None:
+            filter_["isOwned"] = is_owned
+
+        response = self._post(
+            url_path=self._RESOURCE_PATH + "/list",
+            json={"filter": filter_} if filter_ else {},
+        )
+        return ChartList._load(response.json()["items"], cognite_client=self._cognite_client)

cognite_toolkit/_cdf_tk/client/api/extended_data_modeling.py

@@ -9,25 +9,22 @@ from cognite.client._api.data_modeling.instances import InstancesAPI
 from cognite.client._api_client import T
 from cognite.client._cognite_client import ClientConfig, CogniteClient
 from cognite.client._http_client import HTTPClient, HTTPClientConfig, get_global_requests_session
-from cognite.client.data_classes.data_modeling import InstanceApply
+from cognite.client.data_classes.data_modeling import EdgeId, InstanceApply, NodeId
 from cognite.client.exceptions import CogniteConnectionError, CogniteReadTimeout
 from cognite.client.utils import _json
 from cognite.client.utils._concurrency import execute_tasks
 from requests import Response
 
-from cognite_toolkit._cdf_tk.client._constants import DATA_MODELING_MAX_WRITE_WORKERS
+from cognite_toolkit._cdf_tk.client._constants import DATA_MODELING_MAX_DELETE_WORKERS, DATA_MODELING_MAX_WRITE_WORKERS
 from cognite_toolkit._cdf_tk.client.data_classes.instances import InstancesApplyResultList
 from cognite_toolkit._cdf_tk.client.utils._concurrency import ToolkitConcurrencySettings
 from cognite_toolkit._cdf_tk.client.utils._http_client import ToolkitRetryTracker
 
-from .statistics import StatisticsAPI
-
 
 class ExtendedDataModelingAPI(DataModelingAPI):
     def __init__(self, config: ClientConfig, api_version: str | None, cognite_client: CogniteClient) -> None:
         super().__init__(config, api_version, cognite_client)
-        self.statistics = StatisticsAPI(config, api_version, cognite_client)
-        self.instances = ExtendedInstancesAPI(config, api_version, cognite_client)
+        self.instances: ExtendedInstancesAPI = ExtendedInstancesAPI(config, api_version, cognite_client)
 
 
 class ExtendedInstancesAPI(InstancesAPI):
@@ -97,26 +94,85 @@ class ExtendedInstancesAPI(InstancesAPI):
             executor=ToolkitConcurrencySettings.get_data_modeling_write_executor(),
         )
 
-        def unwrap_element(el: T) -> InstanceApply | T:
+        def unwrap_element(el: T) -> InstanceApply | NodeId | EdgeId | T:
             if isinstance(el, dict):
-                return type(items[0])._load(el, cognite_client=self._cognite_client)
+                instance_type = el.get("instanceType")
+                if instance_type == "node":
+                    return NodeId.load(el)
+                elif instance_type == "edge":
+                    return EdgeId.load(el)
+                else:
+                    return InstanceApply._load(el)
             else:
                 return el
 
-        def str_format_element(el: T) -> str | T:
-            if isinstance(el, InstanceApply):
-                return f"{el.space}:{el.external_id}"
-            return el
-
         summary.raise_compound_exception_if_failed_tasks(
             task_unwrap_fn=lambda task: task[1]["items"],
             task_list_element_unwrap_fn=unwrap_element,
-            str_format_element_fn=str_format_element,
         )
         created_resources = summary.joined_results(lambda res: res.json()["items"])
 
         return InstancesApplyResultList._load(created_resources)
 
+    def delete_fast(
+        self,
+        instance_ids: Sequence[NodeId | EdgeId],
+    ) -> list[NodeId | EdgeId]:
+        """`Delete one or more instances <https://developer.cognite.com/api#tag/Instances/operation/deleteBulk>`_
+
+        Args:
+            instance_ids (Sequence[NodeId | EdgeId]): A sequence of NodeId or EdgeId instances to delete.
+
+        Returns:
+            list[NodeId | EdgeId]: A list of NodeId or EdgeId instances that were successfully deleted. An empty
+                list is returned if no instances were deleted.
+        """
+        tasks = [
+            (f"{self._RESOURCE_PATH}/delete", task_items, ToolkitRetryTracker(self._http_client_with_retry.config))
+            for task_items in self._prepare_item_chunks(
+                [
+                    {"space": item.space, "externalId": item.external_id, "instanceType": item._instance_type}
+                    for item in instance_ids
+                ],
+                self._DELETE_LIMIT,
+                None,
+            )
+        ]
+        summary = execute_tasks(
+            self._post_with_item_reduction_retry,
+            tasks,
+            max_workers=min(self._config.max_workers, DATA_MODELING_MAX_DELETE_WORKERS),
+            executor=ToolkitConcurrencySettings.get_data_modeling_delete_executor(),
+        )
+
+        def unwrap_element(el: T) -> InstanceApply | T:
+            if isinstance(el, dict) and "instanceType" in el:
+                if el["instanceType"] == "node":
+                    return NodeId.load(el)
+                elif el["instanceType"] == "edge":
+                    return EdgeId.load(el)
+                return el
+            else:
+                return el
+
+        summary.raise_compound_exception_if_failed_tasks(
+            task_unwrap_fn=lambda task: task[1]["items"],
+            task_list_element_unwrap_fn=unwrap_element,
+        )
+        deleted_resources = summary.joined_results(lambda res: res.json()["items"])
+        result: list[NodeId | EdgeId] = []
+        for resource in deleted_resources:
+            if "instanceType" not in resource:
+                raise ValueError("Resource must contain 'instanceType' key.")
+            instance_type = resource.get("instanceType")
+            if instance_type == "node":
+                result.append(NodeId.load(resource))
+            elif instance_type == "edge":
+                result.append(EdgeId.load(resource))
+            else:
+                raise TypeError(f"Resource must be a NodeId or EdgeId, not {instance_type}.")
+        return result
+
     def _post_with_item_reduction_retry(
         self,
         url_path: str,

cognite_toolkit/_cdf_tk/client/api/extended_files.py

@@ -0,0 +1,227 @@
+from collections.abc import Sequence
+from typing import overload
+
+from cognite.client import ClientConfig, CogniteClient
+from cognite.client._api.files import FilesAPI
+from cognite.client.data_classes.data_modeling import NodeId
+from cognite.client.utils._auxiliary import split_into_chunks, unpack_items_in_payload
+from cognite.client.utils._concurrency import execute_tasks
+from cognite.client.utils._identifier import IdentifierSequence
+from cognite.client.utils.useful_types import SequenceNotStr
+
+from cognite_toolkit._cdf_tk.client.data_classes.extended_filemetadata import (
+    ExtendedFileMetadata,
+    ExtendedFileMetadataList,
+)
+from cognite_toolkit._cdf_tk.client.data_classes.pending_instances_ids import PendingInstanceId
+
+
+class ExtendedFileMetadataAPI(FilesAPI):
+    """Extended FileMetadata 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
+
+    @overload
+    def set_pending_ids(
+        self, instance_id: NodeId | tuple[str, str], id: int | None = None, external_id: str | None = None
+    ) -> ExtendedFileMetadata: ...
+
+    @overload
+    def set_pending_ids(self, instance_id: Sequence[PendingInstanceId]) -> ExtendedFileMetadataList: ...
+
+    def set_pending_ids(
+        self,
+        instance_id: NodeId | tuple[str, str] | Sequence[PendingInstanceId],
+        id: int | None = None,
+        external_id: str | None = None,
+    ) -> ExtendedFileMetadata | ExtendedFileMetadataList:
+        """Set a pending identifier for one or more filemetadata.
+
+        Args:
+            instance_id: The pending instance ID to set.
+            id: The ID of the files.
+            external_id: The external ID of the files.
+
+        Returns:
+            ExtendedFileMetadata: If a single instance ID is provided, returns the updated ExtendedFileMetadata object.
+        """
+        if isinstance(instance_id, NodeId) or (
+            isinstance(instance_id, tuple)
+            and len(instance_id) == 2
+            and isinstance(instance_id[0], str)
+            and isinstance(instance_id[1], str)
+        ):
+            return self._set_pending_ids([PendingInstanceId(NodeId.load(instance_id), id=id, external_id=external_id)])[  # type: ignore[return-value]
+                0
+            ]
+        elif isinstance(instance_id, Sequence) and all(isinstance(item, PendingInstanceId) for item in instance_id):
+            return self._set_pending_ids(instance_id)  # type: ignore[arg-type]
+        else:
+            raise TypeError(
+                "instance_id must be a NodeId, a tuple of (str, str), or a sequence of PendingIdentifier objects."
+            )
+
+    def _set_pending_ids(self, identifiers: Sequence[PendingInstanceId]) -> ExtendedFileMetadataList:
+        tasks = [
+            {
+                "url_path": f"{self._RESOURCE_PATH}/set-pending-instance-ids",
+                "json": {
+                    "items": [identifier.dump(camel_case=True) for identifier in id_chunk],
+                },
+                "api_subversion": "alpha",
+            }
+            for id_chunk in split_into_chunks(list(identifiers), self._PENDING_IDS_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"])
+
+        return ExtendedFileMetadataList._load(retrieved_items, cognite_client=self._cognite_client)
+
+    def retrieve(
+        self, id: int | None = None, external_id: str | None = None, instance_id: NodeId | None = None
+    ) -> ExtendedFileMetadata | None:
+        """`Retrieve a single file metadata by id. <https://developer.cognite.com/api#tag/Files/operation/getFileByInternalId>`_
+
+        Args:
+            id (int | None): ID
+            external_id (str | None): External ID
+            instance_id (NodeId | None): Instance ID
+
+        Returns:
+            FileMetadata | None: Requested file metadata or None if it does not exist.
+
+        Examples:
+
+            Get file metadata by id:
+
+                >>> from cognite.client import CogniteClient
+                >>> client = CogniteClient()
+                >>> res = client.files.retrieve(id=1)
+
+            Get file metadata by external id:
+
+                >>> res = client.files.retrieve(external_id="1")
+        """
+        identifiers = IdentifierSequence.load(ids=id, external_ids=external_id, instance_ids=instance_id).as_singleton()
+        return self._retrieve_multiple(
+            list_cls=ExtendedFileMetadataList,
+            resource_cls=ExtendedFileMetadata,
+            identifiers=identifiers,
+            api_subversion="alpha",
+        )
+
+    def retrieve_multiple(
+        self,
+        ids: Sequence[int] | None = None,
+        external_ids: SequenceNotStr[str] | None = None,
+        instance_ids: Sequence[NodeId] | None = None,
+        ignore_unknown_ids: bool = False,
+    ) -> ExtendedFileMetadataList:
+        """`Retrieve multiple file metadatas by id. <https://developer.cognite.com/api#tag/Files/operation/byIdsFiles>`_
+
+        Args:
+            ids (Sequence[int] | None): IDs
+            external_ids (SequenceNotStr[str] | None): External IDs
+            instance_ids (Sequence[NodeId] | None): Instance IDs
+            ignore_unknown_ids (bool): Ignore IDs and external IDs that are not found rather than throw an exception.
+
+        Returns:
+            FileMetadataList: The requested file metadatas.
+
+        Examples:
+
+            Get file metadatas by id:
+
+                >>> from cognite.client import CogniteClient
+                >>> client = CogniteClient()
+                >>> res = client.files.retrieve_multiple(ids=[1, 2, 3])
+
+            Get file_metadatas by external id:
+
+                >>> res = client.files.retrieve_multiple(external_ids=["abc", "def"])
+        """
+        identifiers = IdentifierSequence.load(ids=ids, external_ids=external_ids, instance_ids=instance_ids)
+        return self._retrieve_multiple(
+            list_cls=ExtendedFileMetadataList,
+            resource_cls=ExtendedFileMetadata,
+            identifiers=identifiers,
+            ignore_unknown_ids=ignore_unknown_ids,
+            api_subversion="alpha",
+        )
+
+    @overload
+    def unlink_instance_ids(
+        self,
+        id: int | None = None,
+        external_id: str | None = None,
+    ) -> ExtendedFileMetadata | None: ...
+
+    @overload
+    def unlink_instance_ids(
+        self,
+        id: Sequence[int] | None = None,
+        external_id: SequenceNotStr[str] | None = None,
+    ) -> ExtendedFileMetadataList: ...
+
+    def unlink_instance_ids(
+        self,
+        id: int | Sequence[int] | None = None,
+        external_id: str | SequenceNotStr[str] | None = None,
+    ) -> ExtendedFileMetadata | ExtendedFileMetadataList | None:
+        """Unlink instance IDs from files.
+
+        This allows a CogniteFile node in Data Modeling to be deleted without deleting the underlying file content.
+
+        Args:
+            id (int | Sequence[int] | None): The ID(s) of the files.
+            external_id (str | SequenceNotStr[str] | None): The external ID(s) of the files.
+
+        Returns:
+            ExtendedFileMetadata | ExtendedFileMetadataList | None: The updated file metadata object(s). For single item requests, returns `None` if the file is not found.
+        """
+        if id is None and external_id is None:
+            raise ValueError("At least one of id or external_id must be provided.")
+        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_list = isinstance(id, Sequence) or (isinstance(external_id, Sequence) and not 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(), 1000)
+        ]
+        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 = ExtendedFileMetadataList._load(retrieved_items, cognite_client=self._cognite_client)
+        if is_list:
+            return result
+        if len(result) == 0:
+            return None
+        if len(result) > 1:
+            raise ValueError("Expected a single file, but multiple were returned.")
+        return result[0]  # type: ignore[return-value]

cognite_toolkit/_cdf_tk/client/api/extended_functions.py

@@ -0,0 +1,80 @@
+from cognite.client import CogniteClient
+from cognite.client._api.functions import FunctionsAPI
+from cognite.client.config import global_config
+from cognite.client.data_classes import Function, FunctionWrite
+from cognite.client.utils.useful_types import SequenceNotStr
+from rich.console import Console
+
+from cognite_toolkit._cdf_tk.client.config import ToolkitClientConfig
+from cognite_toolkit._cdf_tk.utils.collection import chunker
+from cognite_toolkit._cdf_tk.utils.http_client import HTTPClient, RequestMessage2
+from cognite_toolkit._cdf_tk.utils.useful_types import JsonVal
+
+
+class ExtendedFunctionsAPI(FunctionsAPI):
+    def __init__(
+        self,
+        config: ToolkitClientConfig,
+        api_version: str | None,
+        cognite_client: CogniteClient,
+        console: Console | None = None,
+    ) -> None:
+        """
+        Extended Functions API to include custom headers and payload preparation.
+        """
+        super().__init__(config, api_version, cognite_client)
+        self._toolkit_config = config
+        self._toolkit_http_client = HTTPClient(config, max_retries=global_config.max_retries, console=console)
+
+    def create_with_429_retry(self, function: FunctionWrite) -> Function:
+        """Create a function with manual retry handling for 429 Too Many Requests responses.
+
+        This method is a workaround for scenarios where the function creation API is temporarily unavailable
+        due to a full queue on the server side.
+
+        Args:
+            function (FunctionWrite): The function to create.
+
+        Returns:
+            Function: The created function object.
+        """
+        result = self._toolkit_http_client.request_single_retries(
+            message=RequestMessage2(
+                endpoint_url=self._toolkit_config.create_api_url("/functions"),
+                method="POST",
+                body_content={"items": [function.dump(camel_case=True)]},
+            )
+        )
+        success = result.get_success_or_raise()
+        # We assume the API response is one item on a successful creation
+        return Function._load(success.body_json["items"][0], cognite_client=self._cognite_client)
+
+    def delete_with_429_retry(self, external_id: SequenceNotStr[str], ignore_unknown_ids: bool = False) -> None:
+        """Delete one or more functions with retry handling for 429 Too Many Requests responses.
+
+        This method is an improvement over the standard delete method in the FunctionsAPI.
+
+
+        Args:
+            external_id (SequenceNotStr[str]): The external IDs of the functions to delete.
+            ignore_unknown_ids (bool): Ignore IDs and external IDs that are not found.
+            console (Console | None): The rich console to use for printing warnings.
+
+        Returns:
+            None
+        """
+        for chunk in chunker(external_id, self._DELETE_LIMIT):
+            body_content: dict[str, JsonVal] = {
+                "items": [{"externalId": eid} for eid in chunk],
+            }
+            if ignore_unknown_ids:
+                body_content["ignoreUnknownIds"] = True
+            result = self._toolkit_http_client.request_single_retries(
+                message=RequestMessage2(
+                    endpoint_url=self._toolkit_config.create_api_url("/functions/delete"),
+                    method="POST",
+                    body_content=body_content,
+                )
+            )
+            result.get_success_or_raise()
+        return None

cognite_toolkit/_cdf_tk/client/api/extended_raw.py

@@ -0,0 +1,49 @@
+from cognite.client._api.raw import RawAPI
+
+from cognite_toolkit._cdf_tk.client.data_classes.raw import RawProfileResults, RawTable
+
+
+class ExtendedRawAPI(RawAPI):
+    DEFAULT_PROFILE_LIMIT = 1000
+    MAX_PROFILE_LIMIT = 1_000_000
+
+    def profile(
+        self,
+        database: str | RawTable,
+        table: str | None = None,
+        limit: int = DEFAULT_PROFILE_LIMIT,
+        timeout_seconds: int | None = None,
+    ) -> RawProfileResults:
+        """Profiles a table in the specified database and returns the results.
+
+        This is a hidden endpoint that is not part of the official CDF API. However, it is used by the Fusion UI
+        to profile tables in the database. This is implemented internally in Cognite Toolkit as Toolkit offers
+        profiling of raw tables. This is used to show how data flows into CDF resources.
+
+        Args:
+            database (str): The name of the database to profile.
+            table (str): The name of the table to profile.
+            limit (int, optional): The maximum number of rows to profile. Defaults to DEFAULT_PROFILE_LIMIT.
+            timeout_seconds (int, optional): The timeout for the profiling operation in seconds. Defaults to global_config.timeout_seconds.
+
+        Returns:
+            RawProfileResults: The results of the profiling operation.
+
+        """
+        if limit <= 0 or limit > self.MAX_PROFILE_LIMIT:
+            raise ValueError(f"Limit must be between 1 and {self.MAX_PROFILE_LIMIT}, got {limit}.")
+        if isinstance(database, RawTable):
+            db_name = database.db_name
+            table_name = database.table_name
+        elif table is None:
+            raise ValueError(f"Table name must be provided for {database}.")
+        else:
+            db_name = database
+            table_name = table
+        response = self._do_request(
+            "POST",
+            "/profiler/raw",
+            json={"database": db_name, "table": table_name, "limit": limit},
+            timeout=timeout_seconds if timeout_seconds is not None else self._config.timeout,
+        )
+        return RawProfileResults._load(response.json())