cognite-neat 0.106.0__py3-none-any.whl → 0.108.0__py3-none-any.whl

cognite/neat/_graph/extractors/_dms_graph.py

@@ -0,0 +1,199 @@
+from collections.abc import Iterable, Sequence
+
+from cognite.client import data_modeling as dm
+from cognite.client.exceptions import CogniteAPIError
+from cognite.client.utils.useful_types import SequenceNotStr
+from rdflib import Namespace, URIRef
+
+from cognite.neat._client import NeatClient
+from cognite.neat._constants import COGNITE_SPACES, DEFAULT_NAMESPACE
+from cognite.neat._issues import IssueList, NeatIssue, catch_warnings
+from cognite.neat._issues.warnings import CDFAuthWarning, ResourceNotFoundWarning, ResourceRetrievalWarning
+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.transformers import DMSToInformation, VerifyDMSRules
+from cognite.neat._shared import Triple
+
+from ._base import KnowledgeGraphExtractor
+from ._dms import DMSExtractor
+
+
+class DMSGraphExtractor(KnowledgeGraphExtractor):
+    def __init__(
+        self,
+        data_model: dm.DataModel[dm.View],
+        client: NeatClient,
+        namespace: Namespace = DEFAULT_NAMESPACE,
+        issues: Sequence[NeatIssue] | None = None,
+        instance_space: str | SequenceNotStr[str] | None = None,
+        skip_cognite_views: bool = True,
+        unpack_json: bool = False,
+        str_to_ideal_type: bool = False,
+    ) -> None:
+        self._client = client
+        self._data_model = data_model
+        self._namespace = namespace or DEFAULT_NAMESPACE
+        self._issues = IssueList(issues)
+        self._instance_space = instance_space
+        self._skip_cognite_views = skip_cognite_views
+        self._unpack_json = unpack_json
+        self._str_to_ideal_type = str_to_ideal_type
+
+        self._views: list[dm.View] | None = None
+        self._information_rules: InformationRules | None = None
+        self._dms_rules: DMSRules | None = None
+
+    @classmethod
+    def from_data_model_id(
+        cls,
+        data_model_id: dm.DataModelIdentifier,
+        client: NeatClient,
+        namespace: Namespace = DEFAULT_NAMESPACE,
+        instance_space: str | SequenceNotStr[str] | None = None,
+        skip_cognite_views: bool = True,
+        unpack_json: bool = False,
+        str_to_ideal_type: bool = False,
+    ) -> "DMSGraphExtractor":
+        issues: list[NeatIssue] = []
+        try:
+            data_model = client.data_modeling.data_models.retrieve(data_model_id, inline_views=True)
+        except CogniteAPIError as e:
+            issues.append(CDFAuthWarning("retrieving data model", str(e)))
+            return cls(
+                cls._create_empty_model(dm.DataModelId.load(data_model_id)),
+                client,
+                namespace,
+                issues,
+                instance_space,
+                skip_cognite_views,
+                unpack_json,
+                str_to_ideal_type,
+            )
+        if not data_model:
+            issues.append(ResourceRetrievalWarning(frozenset({data_model_id}), "data model"))
+            return cls(
+                cls._create_empty_model(dm.DataModelId.load(data_model_id)),
+                client,
+                namespace,
+                issues,
+                instance_space,
+                skip_cognite_views,
+                unpack_json,
+                str_to_ideal_type,
+            )
+        return cls(
+            data_model.latest_version(),
+            client,
+            namespace,
+            issues,
+            instance_space,
+            skip_cognite_views,
+            unpack_json,
+            str_to_ideal_type,
+        )
+
+    @classmethod
+    def _create_empty_model(cls, data_model_id: dm.DataModelId) -> dm.DataModel:
+        return dm.DataModel(
+            data_model_id.space,
+            data_model_id.external_id,
+            data_model_id.version or "MISSING",
+            is_global=False,
+            last_updated_time=0,
+            created_time=0,
+            description=None,
+            name=None,
+            views=[],
+        )
+
+    @property
+    def _model_views(self) -> list[dm.View]:
+        if self._views is None:
+            self._views = self._get_views()
+        return self._views
+
+    @property
+    def description(self) -> str:
+        return "Extracts a data model with nodes and edges."
+
+    @property
+    def source_uri(self) -> URIRef:
+        space, external_id, version = self._data_model.as_id().as_tuple()
+        return DEFAULT_NAMESPACE[f"{self._client.config.project}/{space}/{external_id}/{version}"]
+
+    def extract(self) -> Iterable[Triple]:
+        """Extracts the knowledge graph from the data model."""
+        views = self._model_views
+        if self._skip_cognite_views:
+            views = [view for view in views if view.space not in COGNITE_SPACES]
+
+        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,
+        ).extract()
+
+    def _get_views(self) -> list[dm.View]:
+        view_by_id: dict[dm.ViewId, dm.View] = {}
+        if view_ids := [view_id for view_id in self._data_model.views if isinstance(view_id, dm.ViewId)]:
+            try:
+                # MyPy does not understand the isinstance check above.
+                retrieved = self._client.data_modeling.views.retrieve(ids=view_ids)  #  type: ignore[arg-type]
+            except CogniteAPIError as e:
+                self._issues.append(CDFAuthWarning("retrieving views", str(e)))
+            else:
+                view_by_id.update({view.as_id(): view for view in retrieved})
+        views: list[dm.View] = []
+        data_model_id = self._data_model.as_id()
+        for dm_view in self._data_model.views:
+            if isinstance(dm_view, dm.View):
+                views.append(dm_view)
+            elif isinstance(dm_view, dm.ViewId):
+                if view := view_by_id.get(dm_view):
+                    views.append(view)
+                else:
+                    self._issues.append(ResourceNotFoundWarning(dm_view, "view", data_model_id, "data model"))
+        return views
+
+    def get_information_rules(self) -> InformationRules:
+        """Returns the information rules that the extractor uses."""
+        if self._information_rules is None:
+            self._information_rules, self._dms_rules = self._create_rules()
+        return self._information_rules
+
+    def get_dms_rules(self) -> DMSRules:
+        """Returns the DMS rules that the extractor uses."""
+        if self._dms_rules is None:
+            self._information_rules, self._dms_rules = self._create_rules()
+        return self._dms_rules
+
+    def get_issues(self) -> IssueList:
+        """Returns the issues that occurred during the extraction."""
+        return self._issues
+
+    def _create_rules(self) -> tuple[InformationRules, DMSRules]:
+        # The DMS and Information rules must be created together to link them property.
+        importer = DMSImporter.from_data_model(self._client, self._data_model)
+        unverified_dms = importer.to_rules()
+        if self._unpack_json and (dms_rules := unverified_dms.rules):
+            # Drop the JSON properties from the DMS rules as these are no longer valid.
+            json_name = Json().name  # To avoid instantiating Json multiple times.
+            dms_rules.properties = [
+                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)
+                )
+            ]
+
+        with catch_warnings() as issues:
+            # 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)
+        self._issues.extend(issues)
+        return information_rules, verified_dms

cognite/neat/_graph/extractors/_mock_graph_generator.py

@@ -141,7 +141,7 @@ def generate_triples(
 
     # pregenerate instance ids for each remaining class
     instance_ids = {
-        key: [URIRef(namespace[f"{key.suffix}-{i+1}"]) for i in range(value)] for key, value in class_count.items()
+        key: [URIRef(namespace[f"{key.suffix}-{i + 1}"]) for i in range(value)] for key, value in class_count.items()
     }
 
     # create triple for each class instance defining its type

cognite/neat/_graph/extractors/_rdf_file.py

@@ -1,6 +1,7 @@
+import zipfile
 from collections.abc import Iterable
 from pathlib import Path
-from typing import get_args
+from typing import cast, get_args
 
 from rdflib import URIRef
 from rdflib.util import guess_format
@@ -10,6 +11,7 @@ from cognite.neat._graph._shared import RDFTypes
 from cognite.neat._graph.extractors._base import BaseExtractor
 from cognite.neat._issues._base import IssueList
 from cognite.neat._issues.errors import FileNotFoundNeatError, FileTypeUnexpectedError
+from cognite.neat._issues.errors._general import NeatValueError
 from cognite.neat._shared import Triple
 
 
@@ -24,25 +26,51 @@ class RdfFileExtractor(BaseExtractor):
 
     def __init__(
         self,
-        filepath: Path,
+        filepath: Path | zipfile.ZipExtFile,
         base_uri: URIRef = DEFAULT_BASE_URI,
         issue_list: IssueList | None = None,
     ):
         self.issue_list = issue_list or IssueList(title=f"{filepath.name}")
         self.base_uri = base_uri
         self.filepath = filepath
-        self.format = guess_format(str(self.filepath))
 
-        if not self.filepath.exists():
+        self.format = guess_format(str(self.filepath) if isinstance(self.filepath, Path) else self.filepath.name)
+
+        if isinstance(self.filepath, Path) and not self.filepath.exists():
             self.issue_list.append(FileNotFoundNeatError(self.filepath))
 
         if not self.format:
             self.issue_list.append(
                 FileTypeUnexpectedError(
-                    self.filepath,
+                    (self.filepath if isinstance(self.filepath, Path) else Path(self.filepath.name)),
                     frozenset(get_args(RDFTypes)),
                 )
             )
 
     def extract(self) -> Iterable[Triple]:
         raise NotImplementedError()
+
+    @classmethod
+    def from_zip(
+        cls,
+        filepath: Path,
+        filename: str = "neat-session/instances/instances.trig",
+        base_uri: URIRef = DEFAULT_BASE_URI,
+        issue_list: IssueList | None = None,
+    ):
+        if not filepath.exists():
+            raise FileNotFoundNeatError(filepath)
+        if filepath.suffix not in {".zip"}:
+            raise NeatValueError("Expected a zip file, got {filepath.suffix}")
+
+        with zipfile.ZipFile(filepath, "r") as zip_ref:
+            for file_info in zip_ref.infolist():
+                if file_info.filename == filename:
+                    # We need to open the file in the zip file, and close it upon
+                    # triple extraction ...
+
+                    print(file_info)
+                    file = zip_ref.open(file_info)
+                    return cls(cast(zipfile.ZipExtFile, file), base_uri, issue_list)
+
+        raise NeatValueError(f"Cannot extract {filename} from zip file {filepath}")

cognite/neat/_graph/loaders/__init__.py

@@ -20,6 +20,4 @@ def _repr_html_() -> str:
         ]
     )._repr_html_()
 
-    return (
-        "<strong>Loader</strong> A loader writes data from Neat's triple storage into a target system" f"<br />{table}"
-    )
+    return f"<strong>Loader</strong> A loader writes data from Neat's triple storage into a target system<br />{table}"

cognite/neat/_graph/loaders/_rdf2dms.py

@@ -1,5 +1,6 @@
 import itertools
 import json
+import urllib.parse
 import warnings
 from collections import defaultdict
 from collections.abc import Iterable, Sequence
@@ -32,7 +33,7 @@ from cognite.neat._issues.errors import (
 from cognite.neat._issues.warnings import PropertyDirectRelationLimitWarning, PropertyTypeNotSupportedWarning
 from cognite.neat._rules.analysis._dms import DMSAnalysis
 from cognite.neat._rules.models import DMSRules
-from cognite.neat._rules.models.data_types import _DATA_TYPE_BY_DMS_TYPE, Json
+from cognite.neat._rules.models.data_types import _DATA_TYPE_BY_DMS_TYPE, Json, String
 from cognite.neat._rules.models.entities._single_value import ViewEntity
 from cognite.neat._shared import InstanceType
 from cognite.neat._store import NeatGraphStore
@@ -70,6 +71,7 @@ class DMSLoader(CDFLoader[dm.InstanceApply]):
         tracker: type[Tracker] | None = None,
         rules: DMSRules | None = None,
         client: NeatClient | None = None,
+        unquote_external_ids: bool = False,
     ):
         super().__init__(graph_store)
         self.data_model = data_model
@@ -79,6 +81,7 @@ class DMSLoader(CDFLoader[dm.InstanceApply]):
         self._tracker: type[Tracker] = tracker or LogTracker
         self.rules = rules
         self._client = client
+        self._unquote_external_ids = unquote_external_ids
 
     @classmethod
     def from_data_model_id(
@@ -99,7 +102,12 @@ class DMSLoader(CDFLoader[dm.InstanceApply]):
 
     @classmethod
     def from_rules(
-        cls, rules: DMSRules, graph_store: NeatGraphStore, instance_space: str, client: NeatClient | None = None
+        cls,
+        rules: DMSRules,
+        graph_store: NeatGraphStore,
+        instance_space: str,
+        client: NeatClient | None = None,
+        unquote_external_ids: bool = False,
     ) -> "DMSLoader":
         issues: list[NeatIssue] = []
         data_model: dm.DataModel[dm.View] | None = None
@@ -125,6 +133,7 @@ class DMSLoader(CDFLoader[dm.InstanceApply]):
             issues,
             rules=rules,
             client=client,
+            unquote_external_ids=unquote_external_ids,
         )
 
     def _load(self, stop_on_exception: bool = False) -> Iterable[dm.InstanceApply | NeatIssue | type[_END_OF_CLASS]]:
@@ -142,7 +151,9 @@ class DMSLoader(CDFLoader[dm.InstanceApply]):
             if self.rules and self.rules.metadata.logical
             else None
         )
+
         view_and_count_by_id = self._select_views_with_instances(self.data_model.views)
+
         if self._client:
             view_and_count_by_id, properties_point_to_self = self._sort_by_direct_relation_dependencies(
                 view_and_count_by_id
@@ -159,7 +170,7 @@ class DMSLoader(CDFLoader[dm.InstanceApply]):
 
         tracker = self._tracker(type(self).__name__, view_ids, "views")
         for view_id, (view, instance_count) in view_and_count_by_id.items():
-            pydantic_cls, edge_by_type, issues = self._create_validation_classes(view)  # type: ignore[var-annotated]
+            pydantic_cls, edge_by_type, edge_by_prop_id, issues = self._create_validation_classes(view)  # type: ignore[var-annotated]
             yield from issues
             tracker.issue(issues)
 
@@ -200,20 +211,61 @@ class DMSLoader(CDFLoader[dm.InstanceApply]):
                 )
 
                 for identifier, properties in instance_iterable:
+                    start_node, end_node = self._pop_start_end_node(properties)
+                    is_edge = start_node and end_node
+                    if (is_edge and view.used_for == "node") or (not is_edge and view.used_for == "edge"):
+                        instance_type = "edge" if is_edge else "node"
+                        creation_error = ResourceCreationError(
+                            identifier,
+                            instance_type,
+                            error=f"{instance_type.capitalize()} found in {view.used_for} view",
+                        )
+                        tracker.issue(creation_error)
+                        if stop_on_exception:
+                            raise creation_error
+                        yield creation_error
+                        continue
+
                     if skip_properties:
                         properties = {k: v for k, v in properties.items() if k not in skip_properties}
-                    try:
-                        yield self._create_node(identifier, properties, pydantic_cls, view_id)
-                    except ValueError as e:
-                        error_node = ResourceCreationError(identifier, "node", error=str(e))
-                        tracker.issue(error_node)
-                        if stop_on_exception:
-                            raise error_node from e
-                        yield error_node
-                    yield from self._create_edges(identifier, properties, edge_by_type, tracker)
+
+                    if start_node and end_node:
+                        # Is an edge
+                        try:
+                            yield self._create_edge_with_properties(
+                                identifier, properties, start_node, end_node, pydantic_cls, view_id
+                            )
+                        except ValueError as e:
+                            error_edge = ResourceCreationError(identifier, "edge", error=str(e))
+                            tracker.issue(error_edge)
+                            if stop_on_exception:
+                                raise error_edge from e
+                            yield error_edge
+                    else:
+                        try:
+                            yield self._create_node(identifier, properties, pydantic_cls, view_id)
+                        except ValueError as e:
+                            error_node = ResourceCreationError(identifier, "node", error=str(e))
+                            tracker.issue(error_node)
+                            if stop_on_exception:
+                                raise error_node from e
+                            yield error_node
+                        yield from self._create_edges_without_properties(
+                            identifier, properties, edge_by_type, edge_by_prop_id, tracker
+                        )
                 tracker.finish(track_id)
                 yield _END_OF_CLASS
 
+    @staticmethod
+    def _pop_start_end_node(properties: dict[str | InstanceType, list[str]]) -> tuple[str | None, str | None]:
+        start_node = properties.pop("startNode", [None])[0]
+        if not start_node:
+            start_node = properties.pop("start_node", [None])[0]
+        end_node = properties.pop("endNode", [None])[0]
+        if not end_node:
+            end_node = properties.pop("end_node", [None])[0]
+        return start_node, end_node
+
     def write_to_file(self, filepath: Path) -> None:
         if filepath.suffix not in [".json", ".yaml", ".yml"]:
             raise ValueError(f"File format {filepath.suffix} is not supported")
@@ -298,17 +350,30 @@ class DMSLoader(CDFLoader[dm.InstanceApply]):
 
     def _create_validation_classes(
         self, view: dm.View
-    ) -> tuple[type[BaseModel], dict[str, tuple[str, dm.EdgeConnection]], NeatIssueList]:
+    ) -> tuple[
+        type[BaseModel],
+        dict[str, tuple[str, dm.EdgeConnection]],
+        dict[str, tuple[str, dm.EdgeConnection]],
+        NeatIssueList,
+    ]:
         issues = IssueList()
         field_definitions: dict[str, tuple[type, Any]] = {}
-        edge_by_property: dict[str, tuple[str, dm.EdgeConnection]] = {}
+        edge_by_type: dict[str, tuple[str, dm.EdgeConnection]] = {}
+        edge_by_prop_id: dict[str, tuple[str, dm.EdgeConnection]] = {}
         validators: dict[str, classmethod] = {}
         direct_relation_by_property: dict[str, dm.DirectRelation] = {}
         unit_properties: list[str] = []
         json_fields: list[str] = []
+        text_fields: list[str] = []
         for prop_id, prop in view.properties.items():
             if isinstance(prop, dm.EdgeConnection):
-                edge_by_property[prop_id] = prop_id, prop
+                if prop.edge_source:
+                    # Edges with properties are created separately
+                    continue
+
+                edge_by_type[prop.type.external_id] = prop_id, prop
+                edge_by_prop_id[prop_id] = prop_id, prop
+
             if isinstance(prop, dm.MappedProperty):
                 if is_readonly_property(prop.container, prop.container_property_identifier):
                     continue
@@ -334,6 +399,8 @@ class DMSLoader(CDFLoader[dm.InstanceApply]):
 
                     if data_type == Json:
                         json_fields.append(prop_id)
+                    elif data_type == String:
+                        text_fields.append(prop_id)
                     python_type = data_type.python
                 if isinstance(prop.type, ListablePropertyType) and prop.type.is_list:
                     python_type = list[python_type]
@@ -414,7 +481,7 @@ class DMSLoader(CDFLoader[dm.InstanceApply]):
             )
 
         pydantic_cls = create_model(view.external_id, __validators__=validators, **field_definitions)  # type: ignore[arg-type, call-overload]
-        return pydantic_cls, edge_by_property, issues
+        return pydantic_cls, edge_by_type, edge_by_prop_id, issues
 
     def _create_node(
         self,
@@ -425,6 +492,8 @@ class DMSLoader(CDFLoader[dm.InstanceApply]):
     ) -> dm.InstanceApply:
         type_ = properties.pop(RDF.type, [None])[0]
         created = pydantic_cls.model_validate(properties)
+        if self._unquote_external_ids:
+            identifier = urllib.parse.unquote(identifier)
 
         return dm.NodeApply(
             space=self.instance_space,
@@ -435,17 +504,49 @@ class DMSLoader(CDFLoader[dm.InstanceApply]):
             ],
         )
 
-    def _create_edges(
+    def _create_edge_with_properties(
+        self,
+        identifier: str,
+        properties: dict[str | InstanceType, list[str]],
+        start_node: str,
+        end_node: str,
+        pydantic_cls: type[BaseModel],
+        view_id: dm.ViewId,
+    ) -> dm.EdgeApply:
+        type_ = properties.pop(RDF.type, [None])[0]
+        created = pydantic_cls.model_validate(properties)
+        if type_ is None:
+            raise ValueError(f"Missing type for edge {identifier}")
+
+        if self._unquote_external_ids:
+            identifier = urllib.parse.unquote(identifier)
+
+        return dm.EdgeApply(
+            space=self.instance_space,
+            external_id=identifier,
+            type=dm.DirectRelationReference(view_id.space, view_id.external_id),
+            start_node=dm.DirectRelationReference(self.instance_space, start_node),
+            end_node=dm.DirectRelationReference(self.instance_space, end_node),
+            sources=[
+                dm.NodeOrEdgeData(source=view_id, properties=dict(created.model_dump(exclude_unset=True).items()))
+            ],
+        )
+
+    def _create_edges_without_properties(
         self,
         identifier: str,
         properties: dict[str, list[str]],
         edge_by_type: dict[str, tuple[str, dm.EdgeConnection]],
+        edge_by_prop_id: dict[str, tuple[str, dm.EdgeConnection]],
         tracker: Tracker,
     ) -> Iterable[dm.EdgeApply | NeatIssue]:
         for predicate, values in properties.items():
-            if predicate not in edge_by_type:
+            if predicate in edge_by_type:
+                prop_id, edge = edge_by_type[predicate]
+            elif predicate in edge_by_prop_id:
+                prop_id, edge = edge_by_prop_id[predicate]
+            else:
                 continue
-            prop_id, edge = edge_by_type[predicate]
             if isinstance(edge, SingleEdgeConnection) and len(values) > 1:
                 error = ResourceDuplicatedError(
                     resource_type="edge",
@@ -456,6 +557,9 @@ class DMSLoader(CDFLoader[dm.InstanceApply]):
                 yield error
             for target in values:
                 external_id = f"{identifier}.{prop_id}.{target}"
+                if self._unquote_external_ids:
+                    external_id = urllib.parse.unquote(external_id)
+
                 yield dm.EdgeApply(
                     space=self.instance_space,
                     external_id=(external_id if len(external_id) < 256 else create_sha256_hash(external_id)),