cognite-neat 0.117.0__tar.gz → 0.117.2__tar.gz

{cognite_neat-0.117.0 → cognite_neat-0.117.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: cognite-neat
-Version: 0.117.0
+Version: 0.117.2
 Summary: Knowledge graph transformation
 License: Apache-2.0
 Author: Nikola Vasiljevic

cognite_neat-0.117.2/cognite/neat/_client/_api/neat_instances.py

@@ -0,0 +1,71 @@
+import time
+from collections.abc import Iterator
+from typing import TYPE_CHECKING, Any, Literal
+
+from cognite.client.data_classes.data_modeling import Edge, Node, ViewId
+from cognite.client.data_classes.filters import Filter
+from cognite.client.exceptions import CogniteAPIError
+from cognite.client.utils.useful_types import SequenceNotStr
+
+if TYPE_CHECKING:
+    from cognite.neat._client._api_client import NeatClient
+
+
+class NeatInstancesAPI:
+    def __init__(self, client: "NeatClient") -> None:
+        self._client = client
+
+    def iterate(
+        self,
+        instance_type: Literal["node", "edge"] = "node",
+        space: str | SequenceNotStr[str] | None = None,
+        filter_: Filter | None = None,
+        source: ViewId | None = None,
+    ) -> Iterator[Node] | Iterator[Edge]:
+        """Neat specific implementation of the client.data_modeling.instances(...) method to account for 408.
+
+        In addition, we enforce sort based on the argument below (provided by Alex B.).
+        """
+        body: dict[str, Any] = {"limit": 1_000, "cursor": None, "instanceType": instance_type}
+        # Without a sort, the sort is implicitly by the internal id, as cursoring needs a stable sort.
+        # By making the sort be on external_id, Postgres should pick the index that's on
+        # (project_id, space, external_id) WHERE deleted_at IS NULL. In other words,
+        # avoiding soft deleted instances.
+        body["sort"] = [
+            {
+                "property": [instance_type, "externalId"],
+                "direction": "ascending",
+            }
+        ]
+        url = f"/api/{self._client._API_VERSION}/projects/{self._client.config.project}/models/instances/list"
+        filter_ = self._client.data_modeling.instances._merge_space_into_filter(instance_type, space, filter_)
+        if filter_:
+            body["filter"] = filter_.dump() if isinstance(filter_, Filter) else filter_
+        if source:
+            body["sources"] = [{"source": source.dump(include_type=True, camel_case=True)}]
+        last_limit_change: float = 0.0
+        while True:
+            try:
+                response = self._client.post(url=url, json=body)
+            except CogniteAPIError as e:
+                if e.code == 408 and body["limit"] > 1:
+                    body["limit"] = body["limit"] // 2
+                    last_limit_change = time.perf_counter()
+                    continue
+                raise e
+            response_body = response.json()
+            yield from (
+                Node.load(item) if item.get("instanceType") == "node" else Edge.load(item)
+                for item in response_body["items"]
+            )
+            next_cursor = response_body.get("nextCursor")
+            if next_cursor is None:
+                break
+            body["cursor"] = next_cursor
+            # Increase the limit every 30 seconds to avoid fetching reduced
+            # instances for a long time after a 408. This limit is somewhat arbitrary,
+            # but it should be large enough to avoid many requests, but small enough
+            # to avoid fetching very slowly for a long time after a 408.
+            if (time.perf_counter() - last_limit_change) > 30.0 and body["limit"] < 1_000:
+                body["limit"] = min(int(body["limit"] * 1.5), 1_000)
+                last_limit_change = time.perf_counter()

{cognite_neat-0.117.0 → cognite_neat-0.117.2}/cognite/neat/_client/_api_client.py

@@ -3,6 +3,7 @@ from cognite.client import ClientConfig, CogniteClient
 from cognite.neat._utils.auth import _CLIENT_NAME
 
 from ._api.data_modeling_loaders import DataModelLoaderAPI
+from ._api.neat_instances import NeatInstancesAPI
 from ._api.schema import SchemaAPI
 
 
@@ -15,3 +16,4 @@ class NeatClient(CogniteClient):
             self._config.client_name = _CLIENT_NAME
         self.loaders = DataModelLoaderAPI(self)
         self.schema = SchemaAPI(self)
+        self.instances = NeatInstancesAPI(self)

{cognite_neat-0.117.0 → cognite_neat-0.117.2}/cognite/neat/_client/testing.py

@@ -7,6 +7,7 @@ from cognite.client.testing import CogniteClientMock
 from cognite.neat._client._api_client import NeatClient
 
 from ._api.data_modeling_loaders import DataModelLoaderAPI
+from ._api.neat_instances import NeatInstancesAPI
 from ._api.schema import SchemaAPI
 
 
@@ -23,6 +24,7 @@ class NeatClientMock(CogniteClientMock):
         super().__init__(*args, **kwargs)
         self.schema = SchemaAPI(self)
         self.loaders = DataModelLoaderAPI(self)
+        self.instances = NeatInstancesAPI(self)
 
 
 @contextmanager

{cognite_neat-0.117.0 → cognite_neat-0.117.2}/cognite/neat/_graph/extractors/_dms.py

@@ -1,15 +1,15 @@
+import typing
 import urllib.parse
-from collections.abc import Iterable, Iterator, Set
+from collections.abc import Callable, Iterable, Iterator, Set
 from functools import cached_property
-from typing import cast
 
-from cognite.client import CogniteClient
 from cognite.client import data_modeling as dm
 from cognite.client.data_classes.data_modeling import DataModelIdentifier
-from cognite.client.data_classes.data_modeling.instances import Instance, InstanceSort
+from cognite.client.data_classes.data_modeling.instances import Edge, Instance, Node
 from cognite.client.utils.useful_types import SequenceNotStr
 from rdflib import RDF, Literal, Namespace, URIRef
 
+from cognite.neat._client import NeatClient
 from cognite.neat._config import GLOBAL_CONFIG
 from cognite.neat._constants import DEFAULT_SPACE_URI, is_readonly_property
 from cognite.neat._issues.errors import ResourceRetrievalError
@@ -33,6 +33,10 @@ class DMSExtractor(BaseExtractor):
             considered as an empty value and skipped.
         str_to_ideal_type: If unpack_json is True, when unpacking JSON objects, if the value is a string, the extractor
             will try to convert it to the ideal type.
+        node_type: The prioritized order of the node type to use. The options are "view" and "type". "view"
+            means the externalId of the view used as type, while type is the node.type.
+        edge_type: The prioritized order of the edge type to use. The options are "view" and "type". "view"
+            means the externalId of the view used as type, while type is the edge.type.
     """
 
     def __init__(
@@ -43,6 +47,8 @@ class DMSExtractor(BaseExtractor):
         unpack_json: bool = False,
         empty_values: Set[str] = DEFAULT_EMPTY_VALUES,
         str_to_ideal_type: bool = False,
+        node_type: tuple[typing.Literal["view", "type"], ...] = ("view",),
+        edge_type: tuple[typing.Literal["view", "type"], ...] = ("view", "type"),
     ) -> None:
         self.total_instances_pair_by_view = total_instances_pair_by_view
         self.limit = limit
@@ -50,11 +56,13 @@ class DMSExtractor(BaseExtractor):
         self.unpack_json = unpack_json
         self.empty_values = empty_values
         self.str_to_ideal_type = str_to_ideal_type
+        self.node_type = node_type
+        self.edge_type = edge_type
 
     @classmethod
     def from_data_model(
         cls,
-        client: CogniteClient,
+        client: NeatClient,
         data_model: DataModelIdentifier,
         limit: int | None = None,
         overwrite_namespace: Namespace | None = None,
@@ -88,7 +96,7 @@ class DMSExtractor(BaseExtractor):
     @classmethod
     def from_views(
         cls,
-        client: CogniteClient,
+        client: NeatClient,
         views: Iterable[dm.View],
         limit: int | None = None,
         overwrite_namespace: Namespace | None = None,
@@ -154,7 +162,10 @@ class DMSExtractor(BaseExtractor):
             else:
                 # If the edge has properties, we create a node for the edge and connect it to the start and end nodes.
                 id_ = self._as_uri_ref(instance)
-                yield id_, RDF.type, self._as_uri_ref(instance.type)
+                type_ = self._create_type(
+                    instance, fallback=self._get_namespace(instance.space).Edge, type_priority=self.edge_type
+                )
+                yield id_, RDF.type, type_
                 yield (
                     id_,
                     self._as_uri_ref(dm.DirectRelationReference(instance.space, "startNode")),
@@ -168,14 +179,9 @@ class DMSExtractor(BaseExtractor):
 
         elif isinstance(instance, dm.Node):
             id_ = self._as_uri_ref(instance)
-            if instance.type:
-                type_ = self._as_uri_ref(cast(dm.DirectRelationReference, instance.type))
-            elif len(instance.properties) == 1:
-                view_id = next(iter(instance.properties.keys()))
-                type_ = self._get_namespace(view_id.space)[urllib.parse.quote(view_id.external_id)]
-            else:
-                type_ = self._get_namespace(instance.space).Node
-
+            type_ = self._create_type(
+                instance, fallback=self._get_namespace(instance.space).Node, type_priority=self.node_type
+            )
             yield id_, RDF.type, type_
         else:
             raise NotImplementedError(f"Unknown instance type {type(instance)}")
@@ -196,6 +202,31 @@ class DMSExtractor(BaseExtractor):
                 self.unpack_json,
             ).extract()
 
+    def _create_type(
+        self, instance: Node | Edge, fallback: URIRef, type_priority: tuple[typing.Literal["view", "type"], ...]
+    ) -> URIRef:
+        method_by_name: dict[str, Callable[[Node | Edge], URIRef | None]] = {
+            "view": self._view_to_rdf_type,
+            "type": self._instance_type_to_rdf,
+        }
+        for method_name in type_priority:
+            type_ = method_by_name[method_name](instance)
+            if type_:
+                return type_
+        else:
+            return fallback
+
+    def _instance_type_to_rdf(self, instance: Node | Edge) -> URIRef | None:
+        if instance.type:
+            return self._as_uri_ref(instance.type)
+        return None
+
+    def _view_to_rdf_type(self, instance: Node | Edge) -> URIRef | None:
+        view_id = next(iter((instance.properties or {}).keys()), None)
+        if view_id:
+            return self._get_namespace(view_id.space)[urllib.parse.quote(view_id.external_id)]
+        return None
+
     def _as_uri_ref(self, instance: Instance | dm.DirectRelationReference) -> URIRef:
         return self._get_namespace(instance.space)[urllib.parse.quote(instance.external_id)]
 
@@ -206,7 +237,7 @@ class DMSExtractor(BaseExtractor):
 
 
 class _ViewInstanceIterator(Iterable[Instance]):
-    def __init__(self, client: CogniteClient, view: dm.View, instance_space: str | SequenceNotStr[str] | None = None):
+    def __init__(self, client: NeatClient, view: dm.View, instance_space: str | SequenceNotStr[str] | None = None):
         self.client = client
         self.view = view
         self.instance_space = instance_space
@@ -244,28 +275,20 @@ class _ViewInstanceIterator(Iterable[Instance]):
         }
         # All nodes and edges with properties
         if self.view.used_for in ("node", "all"):
-            # Without a sort, the sort is implicitly by the internal id, as cursoring needs a stable sort.
-            # By making the sort be on external_id, Postgres should pick the index
-            # that's on (project_id, space, external_id)
-            # WHERE deleted_at IS NULL. In other words, avoiding soft deleted instances.
-            node_iterable: Iterable[Instance] = self.client.data_modeling.instances(
-                chunk_size=None,
+            node_iterable: Iterable[Instance] = self.client.instances.iterate(
                 instance_type="node",
-                sources=[view_id],
+                source=view_id,
                 space=self.instance_space,
-                sort=InstanceSort(["node", "externalId"]),
             )
             if read_only_properties:
                 node_iterable = self._remove_read_only_properties(node_iterable, read_only_properties, view_id)
             yield from node_iterable
 
         if self.view.used_for in ("edge", "all"):
-            yield from self.client.data_modeling.instances(
-                chunk_size=None,
+            yield from self.client.instances.iterate(
                 instance_type="edge",
-                sources=[view_id],
+                source=view_id,
                 space=self.instance_space,
-                sort=InstanceSort(["edge", "externalId"]),
             )
 
         for prop in self.view.properties.values():
@@ -273,14 +296,12 @@ class _ViewInstanceIterator(Iterable[Instance]):
                 if prop.edge_source:
                     # All edges with properties are extracted from the edge source
                     continue
-                yield from self.client.data_modeling.instances(
-                    chunk_size=None,
+                yield from self.client.instances.iterate(
                     instance_type="edge",
-                    filter=dm.filters.Equals(
+                    filter_=dm.filters.Equals(
                         ["edge", "type"], {"space": prop.type.space, "externalId": prop.type.external_id}
                     ),
                     space=self.instance_space,
-                    sort=InstanceSort(["edge", "externalId"]),
                 )
 
     @staticmethod

cognite_neat-0.117.2/cognite/neat/_version.py

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

{cognite_neat-0.117.0 → cognite_neat-0.117.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "cognite-neat"
-version = "0.117.0"
+version = "0.117.2"
 readme = "README.md"
 description = "Knowledge graph transformation"
 authors = [

cognite_neat-0.117.0/cognite/neat/_version.py

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