cognite-neat 1.0.25__tar.gz → 1.0.26__tar.gz

{cognite_neat-1.0.25 → cognite_neat-1.0.26}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cognite-neat
-Version: 1.0.25
+Version: 1.0.26
 Summary: Knowledge graph transformation
 Author: Nikola Vasiljevic, Anders Albert
 Author-email: Nikola Vasiljevic <nikola.vasiljevic@cognite.com>, Anders Albert <anders.albert@cognite.com>

cognite_neat-1.0.26/cognite/neat/_client/api.py

@@ -0,0 +1,178 @@
+from abc import ABC, abstractmethod
+from collections.abc import Iterable, Sequence
+from dataclasses import dataclass
+from typing import Any, Generic, Literal, TypeAlias, TypeVar
+
+from pydantic import BaseModel, JsonValue, TypeAdapter
+
+from cognite.neat._client.config import NeatClientConfig
+from cognite.neat._data_model.models.dms._base import T_Resource, T_Response
+from cognite.neat._utils.collection import chunker_sequence
+from cognite.neat._utils.http_client import HTTPClient, ParametersRequest, SimpleBodyRequest, SuccessResponse
+from cognite.neat._utils.useful_types import T_Reference
+
+from .data_classes import PagedResponse
+
+_T_BaseModel = TypeVar("_T_BaseModel", bound=BaseModel)
+
+
+@dataclass(frozen=True)
+class Endpoint:
+    method: Literal["GET", "POST"]
+    path: str
+    item_limit: int
+
+
+APIMethod: TypeAlias = Literal["apply", "retrieve", "delete", "list"]
+
+
+class NeatAPI(Generic[T_Reference, T_Resource, T_Response], ABC):
+    def __init__(
+        self, neat_config: NeatClientConfig, http_client: HTTPClient, endpoint_map: dict[APIMethod, Endpoint]
+    ) -> None:
+        self._config = neat_config
+        self._http_client = http_client
+        self._method_endpoint_map = endpoint_map
+
+    @abstractmethod
+    def _validate_page_response(self, response: SuccessResponse) -> PagedResponse[T_Response]:
+        """Parse a single item response."""
+        raise NotImplementedError()
+
+    @abstractmethod
+    def _validate_id_response(self, response: SuccessResponse) -> list[T_Reference]:
+        """Parse a single item response."""
+        raise NotImplementedError()
+
+    def _make_url(self, path: str = "") -> str:
+        """Create the full URL for this resource endpoint."""
+        return self._config.create_api_url(path)
+
+    def _request_item_response(
+        self,
+        items: Sequence[BaseModel],
+        method: APIMethod,
+        extra_body: dict[str, Any] | None = None,
+    ) -> list[T_Response]:
+        response_items: list[T_Response] = []
+        for response in self._chunk_requests(items, method, extra_body):
+            response_items.extend(self._validate_page_response(response).items)
+        return response_items
+
+    def _request_id_response(
+        self,
+        items: Sequence[BaseModel],
+        method: APIMethod,
+        extra_body: dict[str, Any] | None = None,
+    ) -> list[T_Reference]:
+        response_items: list[T_Reference] = []
+        for response in self._chunk_requests(items, method, extra_body):
+            response_items.extend(self._validate_id_response(response))
+        return response_items
+
+    def _chunk_requests(
+        self,
+        items: Sequence[_T_BaseModel],
+        method: APIMethod,
+        extra_body: dict[str, Any] | None = None,
+    ) -> Iterable[SuccessResponse]:
+        """Send requests in chunks and yield responses.
+
+        Args:
+            items: The items to process.
+            method: The API method to use. This is used ot look the up the endpoint.
+            extra_body: Optional extra body content to include in the request.
+
+        Yields:
+            The successful responses from the API.
+        """
+        # Filter out None
+        endpoint = self._method_endpoint_map[method]
+
+        for chunk in chunker_sequence(items, endpoint.item_limit):
+            request = SimpleBodyRequest(
+                endpoint_url=self._make_url(endpoint.path),
+                method=endpoint.method,
+                body=TypeAdapter(dict[str, JsonValue]).dump_json(
+                    {
+                        "items": [item.model_dump(by_alias=True) for item in chunk],
+                        **(extra_body or {}),
+                    },
+                ),
+            )
+            response = self._http_client.request_with_retries(request)
+            response.raise_for_status()
+            yield response.success_response
+
+    @classmethod
+    def _filter_out_none_values(cls, params: dict[str, Any] | None) -> dict[str, Any] | None:
+        request_params: dict[str, Any] | None = None
+        if params:
+            request_params = {k: v for k, v in params.items() if v is not None}
+        return request_params
+
+    def _paginate(
+        self,
+        limit: int,
+        cursor: str | None = None,
+        params: dict[str, Any] | None = None,
+    ) -> PagedResponse[T_Response]:
+        """Fetch a single page of resources.
+
+        Args:
+            params: Query parameters for the request. Supported parameters depend on
+                the resource type but typically include:
+                - cursor: Cursor for pagination
+                - limit: Maximum number of items (defaults to list limit)
+                - space: Filter by space
+                - includeGlobal: Whether to include global resources
+            limit: Maximum number of items to return in the page.
+            cursor: Cursor for pagination.
+
+        Returns:
+            A Page containing the items and the cursor for the next page.
+        """
+        endpoint = self._method_endpoint_map["list"]
+        if not (0 < limit <= endpoint.item_limit):
+            raise ValueError(f"Limit must be between 1 and {endpoint.item_limit}, got {limit}.")
+        if endpoint.method != "GET":
+            raise NotImplementedError(f"Pagination not implemented for method {endpoint.method}.")
+        request_params = self._filter_out_none_values(params) or {}
+        request_params["limit"] = limit
+        if cursor is not None:
+            request_params["cursor"] = cursor
+        request = ParametersRequest(
+            endpoint_url=self._make_url(endpoint.path),
+            method=endpoint.method,
+            parameters=request_params,
+        )
+        result = self._http_client.request_with_retries(request)
+        result.raise_for_status()
+        return self._validate_page_response(result.success_response)
+
+    def _iterate(
+        self,
+        limit: int | None = None,
+        cursor: str | None = None,
+        params: dict[str, Any] | None = None,
+    ) -> Iterable[list[T_Response]]:
+        """Iterate over all resources, handling pagination automatically."""
+        next_cursor = cursor
+        total = 0
+        endpoint = self._method_endpoint_map["list"]
+        while True:
+            page_limit = endpoint.item_limit if limit is None else min(limit - total, endpoint.item_limit)
+            page = self._paginate(limit=page_limit, cursor=next_cursor, params=params)
+            yield page.items
+            total += len(page.items)
+            if page.next_cursor is None or (limit is not None and total >= limit):
+                break
+            next_cursor = page.next_cursor
+
+    def _list(
+        self,
+        limit: int | None = None,
+        params: dict[str, Any] | None = None,
+    ) -> list[T_Response]:
+        """List all resources, handling pagination automatically."""
+        return [item for batch in self._iterate(limit=limit, params=params) for item in batch]

cognite_neat-1.0.26/cognite/neat/_client/containers_api.py

@@ -0,0 +1,77 @@
+from collections.abc import Sequence
+
+from cognite.neat._data_model.models.dms import ContainerReference, ContainerRequest, ContainerResponse
+from cognite.neat._utils.http_client import HTTPClient, SuccessResponse
+
+from .api import Endpoint, NeatAPI
+from .config import NeatClientConfig
+from .data_classes import PagedResponse
+from .filters import ContainerFilter
+
+
+class ContainersAPI(NeatAPI):
+    def __init__(self, neat_config: NeatClientConfig, http_client: HTTPClient) -> None:
+        super().__init__(
+            neat_config,
+            http_client,
+            endpoint_map={
+                "apply": Endpoint("POST", "models/containers", item_limit=100),
+                "retrieve": Endpoint("POST", "models/containers/byids", item_limit=100),
+                "delete": Endpoint("POST", "models/containers/delete", item_limit=100),
+                "list": Endpoint("GET", "models/containers", item_limit=1000),
+            },
+        )
+
+    def _validate_page_response(self, response: SuccessResponse) -> PagedResponse[ContainerResponse]:
+        return PagedResponse[ContainerResponse].model_validate_json(response.body)
+
+    def _validate_id_response(self, response: SuccessResponse) -> list[ContainerReference]:
+        return PagedResponse[ContainerReference].model_validate_json(response.body).items
+
+    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.
+        """
+        return self._request_item_response(items, "apply")
+
+    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.
+        """
+        return self._request_item_response(items, "retrieve")
+
+    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.
+        """
+        return self._request_id_response(items, "delete")
+
+    def list(
+        self, space: str | None = None, include_global: bool = False, limit: int | None = 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. If None, return all containers.
+
+        Returns:
+            List of ContainerResponse objects.
+        """
+        filter = ContainerFilter(space=space, include_global=include_global)
+        return self._list(limit=limit, params=filter.dump())

cognite_neat-1.0.26/cognite/neat/_client/data_model_api.py

@@ -0,0 +1,87 @@
+from collections.abc import Sequence
+
+from cognite.neat._data_model.models.dms import DataModelReference, DataModelRequest, DataModelResponse
+from cognite.neat._utils.http_client import HTTPClient, SuccessResponse
+
+from .api import Endpoint, NeatAPI
+from .config import NeatClientConfig
+from .data_classes import PagedResponse
+from .filters import DataModelFilter
+
+
+class DataModelsAPI(NeatAPI):
+    def __init__(self, neat_config: NeatClientConfig, http_client: HTTPClient) -> None:
+        super().__init__(
+            neat_config,
+            http_client,
+            endpoint_map={
+                "apply": Endpoint("POST", "/models/datamodels", item_limit=100),
+                "retrieve": Endpoint("POST", "/models/datamodels/byids", item_limit=100),
+                "delete": Endpoint("POST", "/models/datamodels/delete", item_limit=100),
+                "list": Endpoint("GET", "/models/datamodels", item_limit=1000),
+            },
+        )
+
+    def _validate_page_response(self, response: SuccessResponse) -> PagedResponse[DataModelResponse]:
+        return PagedResponse[DataModelResponse].model_validate_json(response.body)
+
+    def _validate_id_response(self, response: SuccessResponse) -> list[DataModelReference]:
+        return PagedResponse[DataModelReference].model_validate_json(response.body).items
+
+    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.
+        """
+        return self._request_item_response(data_models, "apply")
+
+    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.
+        """
+        return self._request_item_response(items, "retrieve")
+
+    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.
+        """
+        return self._request_id_response(items, "delete")
+
+    def list(
+        self,
+        space: str | None = None,
+        all_versions: bool = False,
+        inline_views: bool = False,
+        include_global: bool = False,
+        limit: int | None = 10,
+    ) -> list[DataModelResponse]:
+        """List data models in CDF Project.
+
+        Args:
+            space: If specified, only data models in this space are returned.
+            all_versions: If True, return all versions. If False, only return the latest version.
+            inline_views: If True, include views inline in the response.
+            include_global: If True, include global data models.
+            limit: Maximum number of data models to return. If None, return all data models.
+
+        Returns:
+            List of DataModelResponse objects.
+        """
+        filter = DataModelFilter(
+            space=space,
+            all_versions=all_versions,
+            inline_views=inline_views,
+            include_global=include_global,
+        )
+        return self._list(limit=limit, params=filter.dump())

cognite_neat-1.0.26/cognite/neat/_client/filters.py

@@ -0,0 +1,40 @@
+from typing import Any, Literal
+
+from pydantic import BaseModel, ConfigDict
+from pydantic.alias_generators import to_camel
+
+
+class BaseModelRequest(BaseModel):
+    """Base class for all object. This includes resources and nested objects."""
+
+
+class Filter(BaseModel):
+    model_config = ConfigDict(alias_generator=to_camel, extra="ignore", populate_by_name=True)
+
+    def dump(self, camel_case: bool = True) -> dict[str, Any]:
+        """Dump the object to a dictionary.
+
+        Args:
+            camel_case (bool): Whether to use camelCase for the keys. Default is True.
+
+        """
+        return self.model_dump(mode="json", by_alias=camel_case, exclude_unset=True)
+
+
+class DataModelingFilter(Filter):
+    space: str | None = None
+    include_global: bool | None = None
+
+
+class ContainerFilter(DataModelingFilter):
+    used_for: Literal["node", "edge", "record", "all"] | None = None
+
+
+class ViewFilter(DataModelingFilter):
+    include_inherited_properties: bool | None = None
+    all_versions: bool | None = None
+
+
+class DataModelFilter(DataModelingFilter):
+    inline_views: bool | None = None
+    all_versions: bool | None = None

cognite_neat-1.0.26/cognite/neat/_client/spaces_api.py

@@ -0,0 +1,80 @@
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+from cognite.neat._data_model.models.dms import SpaceRequest, SpaceResponse
+from cognite.neat._data_model.models.dms._references import SpaceReference
+from cognite.neat._utils.http_client import HTTPClient, SuccessResponse
+
+from .api import Endpoint, NeatAPI
+from .config import NeatClientConfig
+from .data_classes import PagedResponse
+from .filters import DataModelingFilter
+
+
+class SpacesAPI(NeatAPI):
+    def __init__(self, neat_config: NeatClientConfig, http_client: HTTPClient) -> None:
+        super().__init__(
+            neat_config,
+            http_client,
+            endpoint_map={
+                "apply": Endpoint("POST", "/models/spaces", item_limit=100),
+                "retrieve": Endpoint("POST", "/models/spaces/byids", item_limit=100),
+                "delete": Endpoint("POST", "/models/spaces/delete", item_limit=100),
+                "list": Endpoint("GET", "/models/spaces", item_limit=1000),
+            },
+        )
+
+    def _validate_page_response(self, response: SuccessResponse) -> PagedResponse[SpaceResponse]:
+        return PagedResponse[SpaceResponse].model_validate_json(response.body)
+
+    def _validate_id_response(self, response: SuccessResponse) -> list[SpaceReference]:
+        return PagedResponse[SpaceReference].model_validate_json(response.body).items
+
+    def apply(self, spaces: Sequence[SpaceRequest]) -> list[SpaceResponse]:
+        """Apply (create or update) spaces in CDF.
+
+        Args:
+            spaces: List of SpaceRequest objects to apply.
+        Returns:
+            List of SpaceResponse objects.
+        """
+        return self._request_item_response(spaces, "apply")
+
+    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.
+        """
+        return self._request_item_response(spaces, "retrieve")
+
+    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.
+        """
+        return self._request_id_response(spaces, "delete")
+
+    def list(
+        self,
+        include_global: bool = False,
+        limit: int | None = 10,
+    ) -> list[SpaceResponse]:
+        """List spaces in CDF Project.
+
+        Args:
+            include_global: If True, include global spaces.
+            limit: Maximum number of spaces to return. If None, return all spaces.
+
+        Returns:
+            List of SpaceResponse objects.
+        """
+        filter = DataModelingFilter(include_global=include_global)
+        return self._list(limit=limit, params=filter.dump())

{cognite_neat-1.0.25 → cognite_neat-1.0.26}/cognite/neat/_client/statistics_api.py

@@ -1,9 +1,14 @@
-from cognite.neat._client.api import NeatAPI
-from cognite.neat._client.data_classes import StatisticsResponse
-from cognite.neat._utils.http_client import ParametersRequest
+from cognite.neat._utils.http_client import HTTPClient, ParametersRequest
 
+from .config import NeatClientConfig
+from .data_classes import StatisticsResponse
+
+
+class StatisticsAPI:
+    def __init__(self, neat_config: NeatClientConfig, http_client: HTTPClient) -> None:
+        self._config = neat_config
+        self._http_client = http_client
 
-class StatisticsAPI(NeatAPI):
     def project(self) -> StatisticsResponse:
         """Retrieve project-wide usage data and limits.
 

cognite_neat-1.0.26/cognite/neat/_client/views_api.py

@@ -0,0 +1,94 @@
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+from cognite.neat._data_model.models.dms import ViewReference, ViewRequest, ViewResponse
+from cognite.neat._utils.http_client import HTTPClient, SuccessResponse
+
+from .api import Endpoint, NeatAPI
+from .config import NeatClientConfig
+from .data_classes import PagedResponse
+from .filters import ViewFilter
+
+
+class ViewsAPI(NeatAPI):
+    def __init__(self, neat_config: NeatClientConfig, http_client: HTTPClient) -> None:
+        super().__init__(
+            neat_config,
+            http_client,
+            endpoint_map={
+                "apply": Endpoint("POST", "/models/views", item_limit=100),
+                "retrieve": Endpoint("POST", "/models/views/byids", item_limit=100),
+                "delete": Endpoint("POST", "/models/views/delete", item_limit=100),
+                "list": Endpoint("GET", "/models/views", item_limit=1000),
+            },
+        )
+
+    def _validate_page_response(self, response: SuccessResponse) -> PagedResponse[ViewResponse]:
+        return PagedResponse[ViewResponse].model_validate_json(response.body)
+
+    def _validate_id_response(self, response: SuccessResponse) -> list[ViewReference]:
+        return PagedResponse[ViewReference].model_validate_json(response.body).items
+
+    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.
+        """
+        return self._request_item_response(items, "apply")
+
+    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.
+        """
+        return self._request_item_response(
+            items, "retrieve", extra_body={"includeInheritedProperties": include_inherited_properties}
+        )
+
+    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.
+
+        Returns:
+            List of ViewReference objects representing the deleted views.
+        """
+        return self._request_id_response(items, "delete")
+
+    def list(
+        self,
+        space: str | None = None,
+        all_versions: bool = False,
+        include_inherited_properties: bool = True,
+        include_global: bool = False,
+        limit: int | None = 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. If None, return all views.
+
+        Returns:
+            List of ViewResponse objects.
+        """
+        filter = ViewFilter(
+            space=space,
+            all_versions=all_versions,
+            include_inherited_properties=include_inherited_properties,
+            include_global=include_global,
+        )
+        return self._list(limit=limit, params=filter.dump())

{cognite_neat-1.0.25 → cognite_neat-1.0.26}/cognite/neat/_data_model/models/dms/_base.py

@@ -22,6 +22,9 @@ class WriteableResource(Resource, Generic[T_Resource], ABC):
         raise NotImplementedError()
 
 
+T_Response = TypeVar("T_Response", bound=WriteableResource)
+
+
 class APIResource(Generic[T_Reference], ABC):
     """Base class for all API data modeling resources."""
 

{cognite_neat-1.0.25 → cognite_neat-1.0.26}/cognite/neat/_utils/http_client/_client.py

@@ -167,7 +167,9 @@ class HTTPClient:
         if isinstance(item, BodyRequest):
             data = item.data()
             if not global_config.disable_gzip:
-                data = gzip.compress(data.encode("utf-8"))
+                if isinstance(data, str):
+                    data = data.encode("utf-8")
+                data = gzip.compress(data)
         return self.session.request(
             method=item.method,
             url=item.endpoint_url,

{cognite_neat-1.0.25 → cognite_neat-1.0.26}/cognite/neat/_utils/http_client/_data_classes.py

@@ -118,14 +118,14 @@ class BodyRequest(ParametersRequest, ABC):
     """Base class for HTTP request messages with a body"""
 
     @abstractmethod
-    def data(self) -> str:
+    def data(self) -> str | bytes:
         raise NotImplementedError()
 
 
 class SimpleBodyRequest(BodyRequest):
-    body: str
+    body: str | bytes
 
-    def data(self) -> str:
+    def data(self) -> str | bytes:
         return self.body
 
 

cognite_neat-1.0.26/cognite/neat/_version.py

@@ -0,0 +1,2 @@
+__version__ = "1.0.26"
+__engine__ = "^2.0.4"

{cognite_neat-1.0.25 → cognite_neat-1.0.26}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "cognite-neat"
-version = "1.0.25"
+version = "1.0.26"
 readme = "README.md"
 description = "Knowledge graph transformation"
 authors = [

cognite_neat-1.0.25/cognite/neat/_client/api.py

@@ -1,8 +0,0 @@
-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