cognite-neat 0.117.1__tar.gz → 0.117.3__tar.gz

{cognite_neat-0.117.1 → cognite_neat-0.117.3}/PKG-INFO

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

cognite_neat-0.117.3/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.1 → cognite_neat-0.117.3}/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.1 → cognite_neat-0.117.3}/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.1 → cognite_neat-0.117.3}/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-0.117.1 → cognite_neat-0.117.3}/cognite/neat/_utils/spreadsheet.py

@@ -1,5 +1,7 @@
+import warnings
 from dataclasses import dataclass, field
 from typing import Any, Literal, cast, overload
+from warnings import catch_warnings
 
 import pandas as pd
 from openpyxl import load_workbook
@@ -67,12 +69,22 @@ def read_individual_sheet(
     expected_headers: list[str] | None = None,
 ) -> tuple[list[dict], SpreadsheetRead] | list[dict]:
     if expected_headers:
-        target_row = _get_row_number(cast(Worksheet, load_workbook(excel_file)[sheet_name]), expected_headers)
+        with catch_warnings():
+            # When reading spreadsheets produced by neat, they contain dropdowns. These
+            # are not supported by openpyxl and will raise a warning as openpyxl cannot validate these.
+            # We ignore these warnings as Neat will do the same checks.
+            warnings.simplefilter("ignore")
+            target_row = _get_row_number(cast(Worksheet, load_workbook(excel_file)[sheet_name]), expected_headers)
         skiprows = target_row - 1 if target_row is not None else 0
     else:
         skiprows = 0
 
-    raw = pd.read_excel(excel_file, sheet_name, skiprows=skiprows)
+    with catch_warnings():
+        # When reading spreadsheets produced by neat, they contain dropdowns. These
+        # are not supported by openpyxl and will raise a warning as openpyxl cannot validate these.
+        # We ignore these warnings as Neat will do the same checks.
+        warnings.simplefilter("ignore")
+        raw = pd.read_excel(excel_file, sheet_name, skiprows=skiprows)
     is_na = raw.isnull().all(axis=1)
     skip_rows = _find_rows_to_skip(raw)
     empty_rows = is_na[is_na].index.tolist()

cognite_neat-0.117.3/cognite/neat/_version.py

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

{cognite_neat-0.117.1 → cognite_neat-0.117.3}/pyproject.toml

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

cognite_neat-0.117.1/cognite/neat/_version.py

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