cognite-neat 0.117.1__py3-none-any.whl → 0.117.2__py3-none-any.whl

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/_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/_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/_graph/extractors/_dms.py

@@ -3,13 +3,13 @@ import urllib.parse
 from collections.abc import Callable, Iterable, Iterator, Set
 from functools import cached_property
 
-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 Edge, Instance, InstanceSort, Node
+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
@@ -62,7 +62,7 @@ class DMSExtractor(BaseExtractor):
     @classmethod
     def from_data_model(
         cls,
-        client: CogniteClient,
+        client: NeatClient,
         data_model: DataModelIdentifier,
         limit: int | None = None,
         overwrite_namespace: Namespace | None = None,
@@ -96,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,
@@ -237,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
@@ -275,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():
@@ -304,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/_version.py

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

{cognite_neat-0.117.1.dist-info → cognite_neat-0.117.2.dist-info}/METADATA

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

{cognite_neat-0.117.1.dist-info → cognite_neat-0.117.2.dist-info}/RECORD

@@ -3,13 +3,14 @@ cognite/neat/_alpha.py,sha256=AVw3YiSwpOiV3fSoTbwTxd-p7WGlS-zgf4TPZCT_9vE,1163
 cognite/neat/_client/__init__.py,sha256=RQ7MwL8mwGqGHokRzsPqO3XStDzmI4-c12_gz1UPJ74,177
 cognite/neat/_client/_api/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 cognite/neat/_client/_api/data_modeling_loaders.py,sha256=LsB8k2erMZ3BJ_wLMIBhav7ykimLkV2QJjnbeC2xikE,39790
+cognite/neat/_client/_api/neat_instances.py,sha256=BxFmErL5FX1FVupHWsZRBsMMMT0BtExpnioiM5XOSHo,3349
 cognite/neat/_client/_api/schema.py,sha256=K8iNuTKHayHGVldMWfY8N6KWG71hOpNIklWGsAkCPMY,6925
-cognite/neat/_client/_api_client.py,sha256=6cNMizDuqMJZiOqiNRLX46BEtlCB-BpgGLyypksRVYU,616
+cognite/neat/_client/_api_client.py,sha256=wsNK3abwoIppleWzVZmoXe9Ia_wrU4QqOxnxxls3ipY,714
 cognite/neat/_client/data_classes/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 cognite/neat/_client/data_classes/data_modeling.py,sha256=RvpUp9ygd-yffQFJ7O2mQqMLDPIa-dmip5zPb8QVIiw,6672
 cognite/neat/_client/data_classes/neat_sequence.py,sha256=QZWSfWnwk6KlYJvsInco4Wdwc1U8DnOQKWmHebArbQY,10830
 cognite/neat/_client/data_classes/schema.py,sha256=jPFTIKrZI0qZtwu1Bc1m3aLu_VbxIEkhZCneHpIbNkc,23205
-cognite/neat/_client/testing.py,sha256=M5WUzi3YbtZAN22TXas0SjvMSxmjPFT5m3lNwR0MVkI,1077
+cognite/neat/_client/testing.py,sha256=JnzLQegw2f6SATWRNDQ8Fui6qBYyz7vFgA5myjioohY,1175
 cognite/neat/_config.py,sha256=WT1BS8uADcFvGoUYOOfwFOVq_VBl472TisdoA3wLick,280
 cognite/neat/_constants.py,sha256=WcrInu37p56BP7gPRIvTVScu6D9Di4RQG-OiYg5NAek,8250
 cognite/neat/_graph/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -36,7 +37,7 @@ cognite/neat/_graph/extractors/_classic_cdf/_sequences.py,sha256=HxN0mnm5GrPAfcN
 cognite/neat/_graph/extractors/_classic_cdf/_timeseries.py,sha256=6CmmxWWG2IErfNKOPhsjQ5wSOTUZZMjulpaRbHj0Q-g,1560
 cognite/neat/_graph/extractors/_dexpi.py,sha256=SFWnKXYpQToWag9aoX8fxISNa9b8KlqjZnkwI18BzNY,9431
 cognite/neat/_graph/extractors/_dict.py,sha256=CmXvYkzhRjWDPEaXZAb_YyWCwJDeNzdWBkmFyc2K90s,4321
-cognite/neat/_graph/extractors/_dms.py,sha256=rXcwJnkzAP2NcxghC7cWEi-ljpq0c5-8hgeMmaGOWMc,14473
+cognite/neat/_graph/extractors/_dms.py,sha256=fpr9jmJlnB9IS9sdDA7RV50e_RedF7AEVE25qLlgtp4,13810
 cognite/neat/_graph/extractors/_dms_graph.py,sha256=b_OpXsLZ1xCGEWAqsF5Qk6EknZzxDbcfaZ5NOZGaYrE,9441
 cognite/neat/_graph/extractors/_iodd.py,sha256=sV5xKZi2ZT7MQsUQTVji9vxH34pD5yAvUjtT-gb8mts,18473
 cognite/neat/_graph/extractors/_mock_graph_generator.py,sha256=k6qVslbxwE2WpU1L5NLgHJExk8Sho96Ag6-_bcCiLwA,15436
@@ -79,7 +80,7 @@ cognite/neat/_rules/analysis/__init__.py,sha256=0Y7KiXpzrwOIU7GVlC3zlizj1TXe6sod
 cognite/neat/_rules/analysis/_base.py,sha256=t9NBxZoM95_7UqOa9coWzsH0Kv6neKcb0dtebTV2KnI,23434
 cognite/neat/_rules/catalog/__init__.py,sha256=dwDB8b-5GKZuwVyPuiwsH0EaK2FY9-wJrkTjKoL8KE4,250
 cognite/neat/_rules/catalog/classic_model.xlsx,sha256=YkocpkKypizjsWYwOdn5yzIz_BSl8T8SQLxgm4GIjLQ,15014
-cognite/neat/_rules/catalog/hello_world_pump.xlsx,sha256=4Fv9qyv7ARHBXSIh-PiOUHcYrtgznKyR2PZCdEBlA0o,17071
+cognite/neat/_rules/catalog/hello_world_pump.xlsx,sha256=E63t5U1PQLIoUfXp1mEuhuq8I2TGKovZlEfIhO5bevw,23322
 cognite/neat/_rules/catalog/info-rules-imf.xlsx,sha256=vrE5g8vBtsGpwJqygxG3t9I3x4SUAyQsi1vtWfZ8QW4,53682
 cognite/neat/_rules/exporters/__init__.py,sha256=IYBa0DIYlx8cFItgYRw9W4FY_LmVEjuaqMz3JORZZX0,1204
 cognite/neat/_rules/exporters/_base.py,sha256=VkNMy8wsH-x4tAjS44cXgzzNH0CM2k_4RhkMwK50J7g,2284
@@ -176,10 +177,10 @@ cognite/neat/_utils/text.py,sha256=BFJoEOQBFgpelysL92FdF0OVRVFl0q9tRNoz-oRanNc,7
 cognite/neat/_utils/time_.py,sha256=O30LUiDH9TdOYz8_a9pFqTtJdg8vEjC3qHCk8xZblG8,345
 cognite/neat/_utils/upload.py,sha256=xWtM6mFuD2QYQHaZ7zCAuGptbEpPIxcH-raWQu93-Ug,5845
 cognite/neat/_utils/xml_.py,sha256=FQkq84u35MUsnKcL6nTMJ9ajtG9D5i1u4VBnhGqP2DQ,1710
-cognite/neat/_version.py,sha256=162Al0wYAoMRqw3d_VE0sPIzASFvLI9jYjqPsyL1Ht0,46
+cognite/neat/_version.py,sha256=QPLtZ3sS71bmoKNu3TyTYKy7lRyMCDh2l6Ur8ITOlz8,46
 cognite/neat/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-cognite_neat-0.117.1.dist-info/LICENSE,sha256=W8VmvFia4WHa3Gqxq1Ygrq85McUNqIGDVgtdvzT-XqA,11351
-cognite_neat-0.117.1.dist-info/METADATA,sha256=DrUQPta5Ya3Dhco5Ent40MgVts-aNe4dpB8xk90YHIk,5361
-cognite_neat-0.117.1.dist-info/WHEEL,sha256=XbeZDeTWKc1w7CSIyre5aMDU_-PohRwTQceYnisIYYY,88
-cognite_neat-0.117.1.dist-info/entry_points.txt,sha256=SsQlnl8SNMSSjE3acBI835JYFtsIinLSbVmHmMEXv6E,51
-cognite_neat-0.117.1.dist-info/RECORD,,
+cognite_neat-0.117.2.dist-info/LICENSE,sha256=W8VmvFia4WHa3Gqxq1Ygrq85McUNqIGDVgtdvzT-XqA,11351
+cognite_neat-0.117.2.dist-info/METADATA,sha256=ekCAFZpHqHGP0JjfiP9NOa7ETMffALadml3HTGnfWvQ,5361
+cognite_neat-0.117.2.dist-info/WHEEL,sha256=XbeZDeTWKc1w7CSIyre5aMDU_-PohRwTQceYnisIYYY,88
+cognite_neat-0.117.2.dist-info/entry_points.txt,sha256=SsQlnl8SNMSSjE3acBI835JYFtsIinLSbVmHmMEXv6E,51
+cognite_neat-0.117.2.dist-info/RECORD,,