cognite-neat 0.70.1__py3-none-any.whl → 0.127.19__py3-none-any.whl

cognite/neat/__init__.py

@@ -1,3 +1,6 @@
+from cognite.neat.v0.core._utils.auth import get_cognite_client
+from cognite.neat.v0.session import NeatSession
+
 from ._version import __version__
 
-__all__ = ["__version__"]
+__all__ = ["NeatSession", "__version__", "get_cognite_client"]

cognite/neat/_client/__init__.py

@@ -0,0 +1,4 @@
+from .client import NeatClient
+from .config import NeatClientConfig
+
+__all__ = ["NeatClient", "NeatClientConfig"]

cognite/neat/_client/api.py

@@ -0,0 +1,8 @@
+from cognite.neat._client.config import NeatClientConfig
+from cognite.neat._utils.http_client import HTTPClient
+
+
+class NeatAPI:
+    def __init__(self, neat_config: NeatClientConfig, http_client: HTTPClient) -> None:
+        self._config = neat_config
+        self._http_client = http_client

cognite/neat/_client/client.py

@@ -0,0 +1,21 @@
+from cognite.client import ClientConfig, CogniteClient
+
+from cognite.neat._utils.http_client import HTTPClient
+
+from .config import NeatClientConfig
+from .containers_api import ContainersAPI
+from .data_model_api import DataModelsAPI
+from .spaces_api import SpacesAPI
+from .statistics_api import StatisticsAPI
+from .views_api import ViewsAPI
+
+
+class NeatClient:
+    def __init__(self, cognite_client_or_config: CogniteClient | ClientConfig) -> None:
+        self.config = NeatClientConfig(cognite_client_or_config)
+        self.http_client = HTTPClient(self.config)
+        self.data_models = DataModelsAPI(self.config, self.http_client)
+        self.views = ViewsAPI(self.config, self.http_client)
+        self.containers = ContainersAPI(self.config, self.http_client)
+        self.spaces = SpacesAPI(self.config, self.http_client)
+        self.statistics = StatisticsAPI(self.config, self.http_client)

cognite/neat/_client/config.py

@@ -0,0 +1,40 @@
+from cognite.client import ClientConfig, CogniteClient
+
+
+class NeatClientConfig(ClientConfig):
+    def __init__(self, cognite_client_or_config: CogniteClient | ClientConfig):
+        config = (
+            cognite_client_or_config.config
+            if isinstance(cognite_client_or_config, CogniteClient)
+            else cognite_client_or_config
+        )
+        super().__init__(
+            client_name=config.client_name,
+            project=config.project,
+            credentials=config.credentials,
+            api_subversion=config.api_subversion,
+            base_url=config.base_url,
+            max_workers=config.max_workers,
+            headers=config.headers,
+            timeout=config.timeout,
+            file_transfer_timeout=config.file_transfer_timeout,
+            debug=config.debug,
+        )
+
+    def create_api_url(self, endpoint: str) -> str:
+        """Create a full API URL for the given endpoint.
+
+        Args:
+            endpoint (str): The API endpoint to append to the base URL.
+
+        Returns:
+            str: The full API URL.
+
+        Examples:
+            >>> config = NeatClientConfig(cluster="bluefield", project="my_project", ...)
+            >>> config.create_api_url("/models/instances")
+            "https://bluefield.cognitedata.com/api/v1/my_project/models/instances"
+        """
+        if not endpoint.startswith("/"):
+            endpoint = f"/{endpoint}"
+        return f"{self.base_url}/api/v1/projects/{self.project}{endpoint}"

cognite/neat/_client/containers_api.py

@@ -0,0 +1,125 @@
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+from cognite.neat._data_model.models.dms import ContainerReference, ContainerRequest, ContainerResponse, DataModelBody
+from cognite.neat._utils.http_client import ItemIDBody, ItemsRequest, ParametersRequest
+from cognite.neat._utils.useful_types import PrimitiveType
+
+from .api import NeatAPI
+from .data_classes import PagedResponse
+
+
+class ContainersAPI(NeatAPI):
+    ENDPOINT = "/models/containers"
+
+    def apply(self, items: Sequence[ContainerRequest]) -> list[ContainerResponse]:
+        """Apply (create or update) containers in CDF.
+
+        Args:
+            items: List of ContainerRequest objects to apply.
+        Returns:
+            List of ContainerResponse objects.
+        """
+        if not items:
+            return []
+        if len(items) > 100:
+            raise ValueError("Cannot apply more than 100 containers at once.")
+        result = self._http_client.request_with_retries(
+            ItemsRequest(
+                endpoint_url=self._config.create_api_url(self.ENDPOINT),
+                method="POST",
+                body=DataModelBody(items=items),
+            )
+        )
+        result.raise_for_status()
+        result = PagedResponse[ContainerResponse].model_validate_json(result.success_response.body)
+        return result.items
+
+    def retrieve(
+        self,
+        items: list[ContainerReference],
+    ) -> list[ContainerResponse]:
+        """Retrieve containers by their identifiers.
+
+        Args:
+            items: List of (space, external_id) tuples identifying the containers to retrieve.
+
+        Returns:
+            List of ContainerResponse objects.
+        """
+        if not items:
+            return []
+        if len(items) > 1000:
+            raise ValueError("Cannot retrieve more than 1000 containers at once.")
+
+        result = self._http_client.request_with_retries(
+            ItemsRequest(
+                endpoint_url=self._config.create_api_url(f"{self.ENDPOINT}/byids"),
+                method="POST",
+                body=ItemIDBody(items=items),
+            )
+        )
+        result.raise_for_status()
+        result = PagedResponse[ContainerResponse].model_validate_json(result.success_response.body)
+        return result.items
+
+    def delete(self, items: list[ContainerReference]) -> list[ContainerReference]:
+        """Delete containers by their identifiers.
+
+        Args:
+            items: List of ContainerReference objects identifying the containers to delete.
+
+        Returns:
+            List of ContainerReference objects representing the deleted containers.
+        """
+        if not items:
+            return []
+        if len(items) > 100:
+            raise ValueError("Cannot delete more than 100 containers at once.")
+
+        result = self._http_client.request_with_retries(
+            ItemsRequest(
+                endpoint_url=self._config.create_api_url(f"{self.ENDPOINT}/delete"),
+                method="POST",
+                body=ItemIDBody(items=items),
+            )
+        )
+        result.raise_for_status()
+        result = PagedResponse[ContainerReference].model_validate_json(result.success_response.body)
+        return result.items
+
+    def list(
+        self,
+        space: str | None = None,
+        include_global: bool = False,
+        limit: int = 10,
+    ) -> list[ContainerResponse]:
+        """List containers in CDF Project.
+
+        Args:
+            space: If specified, only containers in this space are returned.
+            include_global: If True, include global containers.
+            limit: Maximum number of containers to return. Max is 1000.
+
+        Returns:
+            List of ContainerResponse objects.
+        """
+        if limit > 1000:
+            raise ValueError("Pagination is not (yet) supported for listing containers. The maximum limit is 1000.")
+        parameters: dict[str, PrimitiveType] = {
+            "includeGlobal": include_global,
+            "limit": limit,
+        }
+        if space is not None:
+            parameters["space"] = space
+        result = self._http_client.request_with_retries(
+            ParametersRequest(
+                endpoint_url=self._config.create_api_url(self.ENDPOINT),
+                method="GET",
+                parameters=parameters,
+            )
+        )
+        result.raise_for_status()
+        result = PagedResponse[ContainerResponse].model_validate_json(result.success_response.body)
+        return result.items

cognite/neat/_client/data_classes.py

@@ -0,0 +1,44 @@
+from typing import Generic, TypeVar
+
+from pydantic import BaseModel, Field
+
+T = TypeVar("T", bound=BaseModel)
+
+
+class PagedResponse(BaseModel, Generic[T]):
+    items: list[T]
+    next_cursor: str | None = Field(None, alias="nextCursor")
+
+
+class ResourceLimit(BaseModel):
+    """Model for resources with count and limit."""
+
+    count: int
+    limit: int
+
+
+class InstancesDetail(BaseModel, populate_by_name=True):
+    """Model for instances with detailed metrics."""
+
+    edges: int
+    soft_deleted_edges: int = Field(alias="softDeletedEdges")
+    nodes: int
+    soft_deleted_nodes: int = Field(alias="softDeletedNodes")
+    instances: int
+    instances_limit: int = Field(alias="instancesLimit")
+    soft_deleted_instances: int = Field(alias="softDeletedInstances")
+    soft_deleted_instances_limit: int = Field(alias="softDeletedInstancesLimit")
+
+
+class StatisticsResponse(BaseModel, populate_by_name=True):
+    """Main API response model."""
+
+    spaces: ResourceLimit
+    containers: ResourceLimit
+    views: ResourceLimit
+    data_models: ResourceLimit = Field(alias="dataModels")
+    container_properties: ResourceLimit = Field(alias="containerProperties")
+    instances: InstancesDetail
+    concurrent_read_limit: int = Field(alias="concurrentReadLimit")
+    concurrent_write_limit: int = Field(alias="concurrentWriteLimit")
+    concurrent_delete_limit: int = Field(alias="concurrentDeleteLimit")

cognite/neat/_client/data_model_api.py

@@ -0,0 +1,115 @@
+from collections.abc import Sequence
+
+from cognite.neat._data_model.models.dms import DataModelBody, DataModelReference, DataModelRequest, DataModelResponse
+from cognite.neat._utils.http_client import ItemIDBody, ItemsRequest, ParametersRequest
+from cognite.neat._utils.useful_types import PrimitiveType
+
+from .api import NeatAPI
+from .data_classes import PagedResponse
+
+
+class DataModelsAPI(NeatAPI):
+    ENDPOINT = "/models/datamodels"
+
+    def apply(self, data_models: Sequence[DataModelRequest]) -> list[DataModelResponse]:
+        """Apply (create or update) data models in CDF.
+
+        Args:
+            data_models: List of DataModelRequest objects to apply.
+        Returns:
+            List of DataModelResponse objects.
+        """
+        if not data_models:
+            return []
+        if len(data_models) > 100:
+            raise ValueError("Cannot apply more than 100 data models at once.")
+
+        result = self._http_client.request_with_retries(
+            ItemsRequest(
+                endpoint_url=self._config.create_api_url(self.ENDPOINT),
+                method="POST",
+                body=DataModelBody(items=data_models),
+            )
+        )
+        result.raise_for_status()
+        result = PagedResponse[DataModelResponse].model_validate_json(result.success_response.body)
+        return result.items
+
+    def retrieve(
+        self,
+        items: list[DataModelReference],
+    ) -> list[DataModelResponse]:
+        """Retrieve data models by their identifiers.
+
+        Args:
+            items: List of data models references identifying the data models to retrieve.
+        Returns:
+            List of DataModelResponse objects.
+        """
+        if not items:
+            return []
+        if len(items) > 100:
+            raise ValueError("Cannot retrieve more than 1000 containers at once.")
+
+        result = self._http_client.request_with_retries(
+            ItemsRequest(
+                endpoint_url=self._config.create_api_url(f"{self.ENDPOINT}/byids"),
+                method="POST",
+                body=ItemIDBody(items=items),
+            )
+        )
+        result.raise_for_status()
+        result = PagedResponse[DataModelResponse].model_validate_json(result.success_response.body)
+        return result.items
+
+    def delete(self, items: list[DataModelReference]) -> list[DataModelReference]:
+        """Delete data models by their identifiers.
+
+        Args:
+            items: List of data model references identifying the data models to delete.
+        Returns:
+            List of DataModelReference objects representing the deleted data models.
+        """
+        if not items:
+            return []
+        if len(items) > 100:
+            raise ValueError("Cannot delete more than 100 data models at once.")
+
+        result = self._http_client.request_with_retries(
+            ItemsRequest(
+                endpoint_url=self._config.create_api_url(f"{self.ENDPOINT}/delete"),
+                method="POST",
+                body=ItemIDBody(items=items),
+            )
+        )
+        result.raise_for_status()
+        result = PagedResponse[DataModelReference].model_validate_json(result.success_response.body)
+        return result.items
+
+    def list(
+        self,
+        space: str | None = None,
+        all_versions: bool = False,
+        include_global: bool = False,
+        limit: int = 10,
+    ) -> list[DataModelResponse]:
+        """List data models in CDF Project."""
+        if limit > 1000:
+            raise ValueError("Pagination is not (yet) supported for listing data models. The maximum limit is 1000.")
+        parameters: dict[str, PrimitiveType] = {
+            "allVersions": all_versions,
+            "includeGlobal": include_global,
+            "limit": limit,
+        }
+        if space is not None:
+            parameters["space"] = space
+        result = self._http_client.request_with_retries(
+            ParametersRequest(
+                endpoint_url=self._config.create_api_url(self.ENDPOINT),
+                method="GET",
+                parameters=parameters,
+            )
+        )
+        result.raise_for_status()
+        result = PagedResponse[DataModelResponse].model_validate_json(result.success_response.body)
+        return result.items

cognite/neat/_client/spaces_api.py

@@ -0,0 +1,115 @@
+from __future__ import annotations
+
+from cognite.neat._data_model.models.dms import DataModelBody, SpaceRequest, SpaceResponse
+from cognite.neat._data_model.models.dms._references import SpaceReference
+from cognite.neat._utils.http_client import ItemIDBody, ItemsRequest, ParametersRequest
+from cognite.neat._utils.useful_types import PrimitiveType
+
+from .api import NeatAPI
+from .data_classes import PagedResponse
+
+
+class SpacesAPI(NeatAPI):
+    ENDPOINT = "/models/spaces"
+
+    def apply(self, spaces: list[SpaceRequest]) -> list[SpaceResponse]:
+        """Apply (create or update) spaces in CDF.
+
+        Args:
+            spaces: List of SpaceRequest objects to apply.
+        Returns:
+            List of SpaceResponse objects.
+        """
+        if not spaces:
+            return []
+        if len(spaces) > 100:
+            raise ValueError("Cannot apply more than 100 spaces at once.")
+        result = self._http_client.request_with_retries(
+            ItemsRequest(
+                endpoint_url=self._config.create_api_url(self.ENDPOINT),
+                method="POST",
+                body=DataModelBody(items=spaces),
+            )
+        )
+        result.raise_for_status()
+        result = PagedResponse[SpaceResponse].model_validate_json(result.success_response.body)
+        return result.items
+
+    def retrieve(self, spaces: list[SpaceReference]) -> list[SpaceResponse]:
+        """Retrieve spaces by their identifiers.
+
+        Args:
+            spaces: List of space identifiers to retrieve.
+
+        Returns:
+            List of SpaceResponse objects.
+        """
+        if not spaces:
+            return []
+        if len(spaces) > 1000:
+            raise ValueError("Cannot retrieve more than 1000 spaces at once.")
+
+        result = self._http_client.request_with_retries(
+            ItemsRequest(
+                endpoint_url=self._config.create_api_url(f"{self.ENDPOINT}/byids"),
+                method="POST",
+                body=ItemIDBody(items=spaces),
+            )
+        )
+        result.raise_for_status()
+        result = PagedResponse[SpaceResponse].model_validate_json(result.success_response.body)
+        return result.items
+
+    def delete(self, spaces: list[SpaceReference]) -> list[SpaceReference]:
+        """Delete spaces by their identifiers.
+
+        Args:
+            spaces: List of space identifiers to delete.
+        Returns:
+            List of SpaceReference objects representing the deleted spaces.
+        """
+        if not spaces:
+            return []
+        if len(spaces) > 100:
+            raise ValueError("Cannot delete more than 100 spaces at once.")
+        result = self._http_client.request_with_retries(
+            ItemsRequest(
+                endpoint_url=self._config.create_api_url(f"{self.ENDPOINT}/delete"),
+                method="POST",
+                body=ItemIDBody(items=spaces),
+            )
+        )
+        result.raise_for_status()
+        result = PagedResponse[SpaceReference].model_validate_json(result.success_response.body)
+        return result.items
+
+    def list(
+        self,
+        include_global: bool = False,
+        limit: int = 10,
+    ) -> list[SpaceResponse]:
+        """List spaces in CDF Project.
+
+        Args:
+            include_global: If True, include global spaces.
+            limit: Maximum number of spaces to return. Max is 1000.
+
+        Returns:
+            List of SpaceResponse objects.
+        """
+        if limit > 1000:
+            raise ValueError("Pagination is not (yet) supported for listing spaces. The maximum limit is 1000.")
+        parameters: dict[str, PrimitiveType] = {
+            "includeGlobal": include_global,
+            "limit": limit,
+        }
+        result = self._http_client.request_with_retries(
+            ParametersRequest(
+                endpoint_url=self._config.create_api_url(self.ENDPOINT),
+                method="GET",
+                parameters=parameters,
+            )
+        )
+        result.raise_for_status()
+        result = PagedResponse[SpaceResponse].model_validate_json(result.success_response.body)
+        return result.items

cognite/neat/_client/statistics_api.py

@@ -0,0 +1,24 @@
+from cognite.neat._client.api import NeatAPI
+from cognite.neat._client.data_classes import StatisticsResponse
+from cognite.neat._utils.http_client import ParametersRequest
+
+
+class StatisticsAPI(NeatAPI):
+    def project(self) -> StatisticsResponse:
+        """Retrieve project-wide usage data and limits.
+
+        Returns:
+            StatisticsResponse object.
+        """
+
+        result = self._http_client.request_with_retries(
+            ParametersRequest(
+                endpoint_url=self._config.create_api_url("/models/statistics"),
+                method="GET",
+                parameters=None,
+            )
+        )
+
+        result.raise_for_status()
+        result = StatisticsResponse.model_validate_json(result.success_response.body)
+        return result

cognite/neat/_client/views_api.py

@@ -0,0 +1,129 @@
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+from cognite.neat._data_model.models.dms import DataModelBody, ViewReference, ViewRequest, ViewResponse
+from cognite.neat._utils.collection import chunker_sequence
+from cognite.neat._utils.http_client import ItemIDBody, ItemsRequest, ParametersRequest
+from cognite.neat._utils.useful_types import PrimitiveType
+
+from .api import NeatAPI
+from .data_classes import PagedResponse
+
+
+class ViewsAPI(NeatAPI):
+    ENDPOINT = "/models/views"
+
+    def apply(self, items: Sequence[ViewRequest]) -> list[ViewResponse]:
+        """Create or update views in CDF Project.
+        Args:
+            items: List of ViewRequest objects to create or update.
+        Returns:
+            List of ViewResponse objects.
+        """
+        if not items:
+            return []
+        if len(items) > 100:
+            raise ValueError("Cannot apply more than 100 views at once.")
+        result = self._http_client.request_with_retries(
+            ItemsRequest(
+                endpoint_url=self._config.create_api_url(self.ENDPOINT),
+                method="POST",
+                body=DataModelBody(items=items),
+            )
+        )
+        result.raise_for_status()
+        result = PagedResponse[ViewResponse].model_validate_json(result.success_response.body)
+        return result.items
+
+    def retrieve(
+        self,
+        items: list[ViewReference],
+        include_inherited_properties: bool = True,
+    ) -> list[ViewResponse]:
+        """Retrieve views by their identifiers.
+
+        Args:
+            items: List of (space, external_id, version) tuples identifying the views to retrieve.
+            include_inherited_properties: If True, include properties inherited from parent views.
+
+        Returns:
+            List of ViewResponse objects.
+        """
+        results: list[ViewResponse] = []
+        for chunk in chunker_sequence(items, 100):
+            batch = self._http_client.request_with_retries(
+                ItemsRequest(
+                    endpoint_url=self._config.create_api_url(f"{self.ENDPOINT}/byids"),
+                    method="POST",
+                    body=ItemIDBody(items=chunk),
+                    parameters={"includeInheritedProperties": include_inherited_properties},
+                )
+            )
+            batch.raise_for_status()
+            result = PagedResponse[ViewResponse].model_validate_json(batch.success_response.body)
+            results.extend(result.items)
+        return results
+
+    def delete(self, items: list[ViewReference]) -> list[ViewReference]:
+        """Delete views by their identifiers.
+
+        Args:
+            items: List of (space, external_id, version) tuples identifying the views to delete.
+        """
+        if not items:
+            return []
+        if len(items) > 100:
+            raise ValueError("Cannot delete more than 100 views at once.")
+
+        result = self._http_client.request_with_retries(
+            ItemsRequest(
+                endpoint_url=self._config.create_api_url(f"{self.ENDPOINT}/delete"),
+                method="POST",
+                body=ItemIDBody(items=items),
+            )
+        )
+        result.raise_for_status()
+        result = PagedResponse[ViewReference].model_validate_json(result.success_response.body)
+        return result.items
+
+    def list(
+        self,
+        space: str | None = None,
+        all_versions: bool = False,
+        include_inherited_properties: bool = True,
+        include_global: bool = False,
+        limit: int = 10,
+    ) -> list[ViewResponse]:
+        """List views in CDF Project.
+
+        Args:
+            space: If specified, only views in this space are returned.
+            all_versions: If True, return all versions. If False, only return the latest version.
+            include_inherited_properties: If True, include properties inherited from parent views.
+            include_global: If True, include global views.
+            limit: Maximum number of views to return. Max is 1000.
+
+        Returns:
+            List of ViewResponse objects.
+        """
+        if limit > 1000:
+            raise ValueError("Pagination is not (yet) supported for listing views. The maximum limit is 1000.")
+        parameters: dict[str, PrimitiveType] = {
+            "allVersions": all_versions,
+            "includeInheritedProperties": include_inherited_properties,
+            "includeGlobal": include_global,
+            "limit": limit,
+        }
+        if space is not None:
+            parameters["space"] = space
+        result = self._http_client.request_with_retries(
+            ParametersRequest(
+                endpoint_url=self._config.create_api_url(self.ENDPOINT),
+                method="GET",
+                parameters=parameters,
+            )
+        )
+        result.raise_for_status()
+        result = PagedResponse[ViewResponse].model_validate_json(result.success_response.body)
+        return result.items