cognite-neat 0.110.0__py3-none-any.whl → 0.111.1__py3-none-any.whl

cognite/neat/_graph/extractors/_dms.py

@@ -6,20 +6,18 @@ 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, PropertyValue
+from cognite.client.data_classes.data_modeling.instances import Instance, InstanceSort
 from cognite.client.utils.useful_types import SequenceNotStr
-from rdflib import RDF, XSD, Literal, Namespace, URIRef
+from rdflib import RDF, Literal, Namespace, URIRef
 
 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
 from cognite.neat._shared import Triple
-from cognite.neat._utils.auxiliary import string_to_ideal_type
 from cognite.neat._utils.collection_ import iterate_progress_bar
 
 from ._base import BaseExtractor
-
-DEFAULT_EMPTY_VALUES = frozenset({"nan", "null", "none", "", " ", "nil", "n/a", "na", "unknown", "undefined"})
+from ._dict import DEFAULT_EMPTY_VALUES, DMSPropertyExtractor
 
 
 class DMSExtractor(BaseExtractor):
@@ -188,39 +186,15 @@ class DMSExtractor(BaseExtractor):
 
         for view_id, properties in instance.properties.items():
             namespace = self._get_namespace(view_id.space)
-            for key, value in properties.items():
-                for predicate_str, object_ in self._get_predicate_objects_pair(key, value):
-                    yield id_, namespace[urllib.parse.quote(predicate_str)], object_
-
-    def _get_predicate_objects_pair(self, key: str, value: PropertyValue) -> Iterable[tuple[str, Literal | URIRef]]:
-        if isinstance(value, str | float | bool | int):
-            yield key, Literal(value)
-        elif isinstance(value, dict) and "space" in value and "externalId" in value:
-            yield key, self._as_uri_ref(dm.DirectRelationReference.load(value))
-        elif isinstance(value, dict) and self.unpack_json:
-            for sub_key, sub_value in value.items():
-                if isinstance(sub_value, str):
-                    if sub_value.casefold() in self.empty_values:
-                        continue
-                    if self.str_to_ideal_type:
-                        yield sub_key, Literal(string_to_ideal_type(sub_value))
-                    else:
-                        yield sub_key, Literal(sub_value)
-                elif isinstance(sub_value, int | float | bool):
-                    yield sub_key, Literal(sub_value)
-                elif isinstance(sub_value, dict):
-                    yield from self._get_predicate_objects_pair(f"{key}_{sub_key}", sub_value)
-                elif isinstance(sub_value, list):
-                    for item in sub_value:
-                        yield from self._get_predicate_objects_pair(f"{key}_{sub_key}", item)
-                else:
-                    yield sub_key, Literal(str(sub_value))
-        elif isinstance(value, dict):
-            # This object is a json object.
-            yield key, Literal(str(value), datatype=XSD._NS["json"])
-        elif isinstance(value, list):
-            for item in value:
-                yield from self._get_predicate_objects_pair(key, item)
+            yield from DMSPropertyExtractor(
+                id_,
+                properties,
+                namespace,
+                self._as_uri_ref,
+                self.empty_values,
+                self.str_to_ideal_type,
+                self.unpack_json,
+            ).extract()
 
     def _as_uri_ref(self, instance: Instance | dm.DirectRelationReference) -> URIRef:
         return self._get_namespace(instance.space)[urllib.parse.quote(instance.external_id)]
@@ -270,8 +244,16 @@ 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, instance_type="node", sources=[view_id], space=self.instance_space
+                chunk_size=None,
+                instance_type="node",
+                sources=[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)
@@ -279,7 +261,11 @@ class _ViewInstanceIterator(Iterable[Instance]):
 
         if self.view.used_for in ("edge", "all"):
             yield from self.client.data_modeling.instances(
-                chunk_size=None, instance_type="edge", sources=[view_id], space=self.instance_space
+                chunk_size=None,
+                instance_type="edge",
+                sources=[view_id],
+                space=self.instance_space,
+                sort=InstanceSort(["edge", "externalId"]),
             )
 
         for prop in self.view.properties.values():
@@ -294,6 +280,7 @@ class _ViewInstanceIterator(Iterable[Instance]):
                         ["edge", "type"], {"space": prop.type.space, "externalId": prop.type.external_id}
                     ),
                     space=self.instance_space,
+                    sort=InstanceSort(["edge", "externalId"]),
                 )
 
     @staticmethod

cognite/neat/_graph/extractors/_dms_graph.py

@@ -12,6 +12,8 @@ from cognite.neat._issues.warnings import CDFAuthWarning, ResourceNotFoundWarnin
 from cognite.neat._rules.importers import DMSImporter
 from cognite.neat._rules.models import DMSRules, InformationRules
 from cognite.neat._rules.models.data_types import Json
+from cognite.neat._rules.models.entities import UnknownEntity
+from cognite.neat._rules.models.information import InformationProperty
 from cognite.neat._rules.transformers import DMSToInformation, VerifyDMSRules
 from cognite.neat._shared import Triple
 
@@ -131,7 +133,6 @@ class DMSGraphExtractor(KnowledgeGraphExtractor):
         yield from DMSExtractor.from_views(
             self._client,
             views,
-            overwrite_namespace=self._namespace,
             instance_space=self._instance_space,
             unpack_json=self._unpack_json,
             str_to_ideal_type=self._str_to_ideal_type,
@@ -186,8 +187,12 @@ class DMSGraphExtractor(KnowledgeGraphExtractor):
                 prop
                 for prop in dms_rules.properties
                 if not (
-                    isinstance(prop.value_type, Json)
-                    or (isinstance(prop.value_type, str) and prop.value_type == json_name)
+                    (
+                        isinstance(prop.value_type, Json)
+                        or (isinstance(prop.value_type, str) and prop.value_type == json_name)
+                    )
+                    # We are not unpacking list of JSONs.
+                    and prop.is_list is not True
                 )
             ]
 
@@ -195,5 +200,27 @@ class DMSGraphExtractor(KnowledgeGraphExtractor):
             # Any errors occur will be raised and caught outside the extractor.
             verified_dms = VerifyDMSRules(client=self._client).transform(unverified_dms)
             information_rules = DMSToInformation(self._namespace).transform(verified_dms)
+
+        # We need to sync the metadata between the two rules, such that the `.sync_with_info_rules` method works.
+        information_rules.metadata.physical = verified_dms.metadata.identifier
+        verified_dms.metadata.logical = information_rules.metadata.identifier
+        verified_dms.sync_with_info_rules(information_rules)
+
+        # Adding startNode and endNode to the information rules for views that are used for edges.
+        classes_by_prefix = {cls_.class_.prefix: cls_ for cls_ in information_rules.classes}
+        for view in self._model_views:
+            if view.used_for == "edge" and view.external_id in classes_by_prefix:
+                cls_ = classes_by_prefix[view.external_id]
+                for property_ in ("startNode", "endNode"):
+                    information_rules.properties.append(
+                        InformationProperty(
+                            class_=cls_.class_,
+                            property_=property_,
+                            value_type=UnknownEntity(),
+                            min_count=0,
+                            max_count=1,
+                        )
+                    )
+
         self._issues.extend(issues)
         return information_rules, verified_dms

cognite/neat/_graph/extractors/_raw.py

@@ -0,0 +1,67 @@
+import urllib.parse
+from collections.abc import Iterable, Set
+from typing import Any, cast
+
+from cognite.client.data_classes import Row, RowList
+from cognite.client.utils.useful_types import SequenceNotStr
+from rdflib import RDF, Namespace, URIRef
+
+from cognite.neat._client import NeatClient
+from cognite.neat._constants import DEFAULT_RAW_URI
+from cognite.neat._shared import Triple
+
+from ._base import BaseExtractor
+from ._dict import DEFAULT_EMPTY_VALUES, DictExtractor
+
+
+class RAWExtractor(BaseExtractor):
+    def __init__(
+        self,
+        client: NeatClient,
+        db_name: str,
+        table_name: str,
+        table_type: str | None = None,
+        foreign_keys: str | SequenceNotStr[str] | None = None,
+        namespace: Namespace | None = None,
+        empty_values: Set[str] = DEFAULT_EMPTY_VALUES,
+        str_to_ideal_type: bool = False,
+        unpack_json: bool = False,
+    ) -> None:
+        self.client = client
+        self.db_name = db_name
+        self.table_name = table_name
+        self.table_type = table_type
+        self.foreign_keys = {foreign_keys} if isinstance(foreign_keys, str) else set(foreign_keys or [])
+        self.namespace = namespace or Namespace(DEFAULT_RAW_URI)
+        self.empty_values = empty_values
+        self.str_to_ideal_type = str_to_ideal_type
+        self.unpack_json = unpack_json
+
+    @property
+    def _rdf_type(self) -> URIRef:
+        return self.namespace[urllib.parse.quote(self.table_type or self.table_name)]
+
+    def extract(self) -> Iterable[Triple]:
+        for row in self.client.raw.rows(self.db_name, self.table_name, partitions=10, chunk_size=None):
+            if isinstance(row, Row):
+                yield from self._row2triples(row)
+            elif isinstance(row, RowList):
+                # Bug in SDK returning row list with chunk_size= None
+                for item in row:
+                    yield from self._row2triples(item)
+
+    def _row2triples(self, row: Row) -> Iterable[Triple]:
+        # The row is always set. It is just the PySDK that have it as str | None
+        key, data = cast(tuple[str, dict[str, Any]], (row.key, row.columns))
+        identifier = self.namespace[urllib.parse.quote(key)]
+        yield identifier, RDF.type, self._rdf_type
+
+        yield from DictExtractor(
+            identifier,
+            data,
+            self.namespace,
+            self.foreign_keys,
+            self.empty_values,
+            self.str_to_ideal_type,
+            self.unpack_json,
+        ).extract()

cognite/neat/_graph/loaders/_base.py

@@ -20,6 +20,11 @@ T_Output = TypeVar("T_Output")
 class _END_OF_CLASS: ...
 
 
+class _START_OF_CLASS:
+    def __init__(self, class_name: str | None = None):
+        self.class_name = class_name
+
+
 class BaseLoader(ABC, Generic[T_Output]):
     _new_line = "\n"
     _encoding = "utf-8"
@@ -33,10 +38,16 @@ class BaseLoader(ABC, Generic[T_Output]):
 
     def load(self, stop_on_exception: bool = False) -> Iterable[T_Output | NeatIssue]:
         """Load the graph with data."""
-        return (item for item in self._load(stop_on_exception) if item is not _END_OF_CLASS)  # type: ignore[misc]
+        return (
+            item  # type: ignore[misc]
+            for item in self._load(stop_on_exception)
+            if not (item is _END_OF_CLASS or isinstance(item, _START_OF_CLASS))
+        )
 
     @abstractmethod
-    def _load(self, stop_on_exception: bool = False) -> Iterable[T_Output | NeatIssue | type[_END_OF_CLASS]]:
+    def _load(
+        self, stop_on_exception: bool = False
+    ) -> Iterable[T_Output | NeatIssue | type[_END_OF_CLASS] | _START_OF_CLASS]:
         """Load the graph with data."""
         pass
 
@@ -75,21 +86,25 @@ class CDFLoader(BaseLoader[T_Output]):
 
         issues = IssueList()
         items: list[T_Output] = []
+        last_class_name: str | None = None
         for result in self._load(stop_on_exception=False):
             if isinstance(result, NeatIssue):
                 issues.append(result)
             elif result is _END_OF_CLASS:
                 ...
+            elif isinstance(result, _START_OF_CLASS):
+                last_class_name = result.class_name
+                continue
             else:
                 # MyPy does not understand that 'else' means the item will be of type T_Output
                 items.append(result)  # type: ignore[arg-type]
 
             if len(items) >= self._UPLOAD_BATCH_SIZE or result is _END_OF_CLASS:
-                yield from self._upload_to_cdf(client, items, dry_run, issues)
+                yield from self._upload_to_cdf(client, items, dry_run, issues, last_class_name)
                 issues = IssueList()
                 items = []
         if items:
-            yield from self._upload_to_cdf(client, items, dry_run, issues)
+            yield from self._upload_to_cdf(client, items, dry_run, issues, last_class_name)
 
     @abstractmethod
     def _get_required_capabilities(self) -> list[Capability]:
@@ -102,5 +117,6 @@ class CDFLoader(BaseLoader[T_Output]):
         items: list[T_Output],
         dry_run: bool,
         read_issues: IssueList,
+        class_name: str | None = None,
     ) -> Iterable[UploadResult]:
         raise NotImplementedError