cognite-neat 0.102.0__tar.gz → 0.103.0__tar.gz

{cognite_neat-0.102.0 → cognite_neat-0.103.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: cognite-neat
-Version: 0.102.0
+Version: 0.103.0
 Summary: Knowledge graph transformation
 Home-page: https://cognite-neat.readthedocs-hosted.com/
 License: Apache-2.0

{cognite_neat-0.102.0 → cognite_neat-0.103.0}/cognite/neat/__init__.py

@@ -2,4 +2,4 @@ from ._session import NeatSession
 from ._utils.auth import get_cognite_client
 from ._version import __version__
 
-__all__ = ["__version__", "get_cognite_client", "NeatSession"]
+__all__ = ["NeatSession", "__version__", "get_cognite_client"]

{cognite_neat-0.102.0 → cognite_neat-0.103.0}/cognite/neat/_app/api/routers/crud.py

@@ -70,7 +70,7 @@ async def file_upload_handler(
 
         NEAT_APP.workflow_manager.save_workflow_to_storage(workflow_name)
 
-    if "start_workflow" in action and file_type == "rules" or file_type == "staging":
+    if ("start_workflow" in action and file_type == "rules") or file_type == "staging":
         logging.info("Starting workflow after file upload")
         workflow = NEAT_APP.workflow_manager.get_workflow(workflow_name)
         if workflow is None:

{cognite_neat-0.102.0 → cognite_neat-0.103.0}/cognite/neat/_client/__init__.py

@@ -1,4 +1,4 @@
 from ._api.data_modeling_loaders import DataModelingLoader, ResourceLoader
 from ._api_client import NeatClient
 
-__all__ = ["NeatClient", "ResourceLoader", "DataModelingLoader"]
+__all__ = ["DataModelingLoader", "NeatClient", "ResourceLoader"]

{cognite_neat-0.102.0 → cognite_neat-0.103.0}/cognite/neat/_client/_api/data_modeling_loaders.py

@@ -609,6 +609,7 @@ class ViewLoader(DataModelingLoader[ViewId, ViewApply, View, ViewApplyList, View
                 elif view_id in self._items_by_id:
                     view = self._items_by_id[view_id]
                     found.append(view)
+                    found_ids.add(view_id)
                     batch_ids.update(self.get_connected_views(view, include_ancestors, include_connections, found_ids))
                 else:
                     to_retrieve_from_cdf.add(view_id)
@@ -620,7 +621,6 @@ class ViewLoader(DataModelingLoader[ViewId, ViewApply, View, ViewApplyList, View
                 found_ids.update({view.as_id() for view in retrieved_batch})
                 for view in retrieved_batch:
                     batch_ids.update(self.get_connected_views(view, include_ancestors, include_connections, found_ids))
-
             last_batch = batch_ids
 
         if self.cache is False:

{cognite_neat-0.102.0 → cognite_neat-0.103.0}/cognite/neat/_client/_api/schema.py

@@ -20,7 +20,7 @@ class SchemaAPI:
         view_ids: Sequence[dm.ViewId],
         container_ids: Sequence[dm.ContainerId],
         include_ancestors: bool = True,
-        include_connections=True,
+        include_connections: bool = True,
         data_model_id: dm.DataModelId | None = None,
     ) -> DMSSchema:
         data_model_id = data_model_id or dm.DataModelId("NEAT_LOOKUP", "NEAT_LOOKUP", "NEAT_LOOKUP")

{cognite_neat-0.102.0 → cognite_neat-0.103.0}/cognite/neat/_graph/_tracking/__init__.py

@@ -1,4 +1,4 @@
 from .base import Tracker
 from .log import LogTracker
 
-__all__ = ["Tracker", "LogTracker"]
+__all__ = ["LogTracker", "Tracker"]

{cognite_neat-0.102.0 → cognite_neat-0.103.0}/cognite/neat/_graph/extractors/_mock_graph_generator.py

@@ -232,9 +232,8 @@ def _remove_higher_occurring_sym_pair(
             index = class_linkage[
                 (class_linkage.source_class == source) & (class_linkage.target_class == target)
             ].index.values[0]
-        elif (
-            second_sym_property_occurrence is None
-            or first_sym_property_occurrence <= second_sym_property_occurrence
+        elif second_sym_property_occurrence is None or (
+            first_sym_property_occurrence <= second_sym_property_occurrence
             and second_sym_property_occurrence > first_sym_property_occurrence
         ):
             # this means that target occurrence is unbounded

{cognite_neat-0.102.0 → cognite_neat-0.103.0}/cognite/neat/_graph/loaders/_base.py

@@ -81,7 +81,7 @@ class CDFLoader(BaseLoader[T_Output]):
             elif result is _END_OF_CLASS:
                 ...
             else:
-                # MyPy does not understand else means the item will be of type T_Output
+                # 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:

{cognite_neat-0.102.0 → cognite_neat-0.103.0}/cognite/neat/_graph/loaders/_rdf2dms.py

@@ -1,6 +1,8 @@
 import itertools
 import json
+from collections import defaultdict
 from collections.abc import Iterable, Sequence
+from graphlib import TopologicalSorter
 from pathlib import Path
 from typing import Any, get_args
 
@@ -16,6 +18,7 @@ from cognite.client.exceptions import CogniteAPIError
 from pydantic import BaseModel, ValidationInfo, create_model, field_validator
 from rdflib import RDF, URIRef
 
+from cognite.neat._client import NeatClient
 from cognite.neat._graph._tracking import LogTracker, Tracker
 from cognite.neat._issues import IssueList, NeatIssue, NeatIssueList
 from cognite.neat._issues.errors import (
@@ -32,9 +35,10 @@ from cognite.neat._rules.models.entities._single_value import ViewEntity
 from cognite.neat._shared import InstanceType
 from cognite.neat._store import NeatGraphStore
 from cognite.neat._utils.auxiliary import create_sha256_hash
+from cognite.neat._utils.rdf_ import remove_namespace_from_uri
 from cognite.neat._utils.upload import UploadResult
 
-from ._base import CDFLoader
+from ._base import _END_OF_CLASS, CDFLoader
 
 
 class DMSLoader(CDFLoader[dm.InstanceApply]):
@@ -44,9 +48,13 @@ class DMSLoader(CDFLoader[dm.InstanceApply]):
         graph_store (NeatGraphStore): The graph store to load the data into.
         data_model (dm.DataModel[dm.View] | None): The data model to load.
         instance_space (str): The instance space to load the data into.
-        class_by_view_id (dict[ViewId, str] | None): A mapping from view id to class name. Defaults to None.
+        class_neat_id_by_view_id (dict[ViewId, URIRef] | None): A mapping from view id to class name. Defaults to None.
         create_issues (Sequence[NeatIssue] | None): A list of issues that occurred during reading. Defaults to None.
         tracker (type[Tracker] | None): The tracker to use. Defaults to None.
+        rules (DMSRules | None): The DMS rules used by the data model. This is used to lookup the
+            instances in the store. Defaults to None.
+        client (NeatClient | None): This is used to lookup containers such that the loader
+            creates instances in accordance with required constraints. Defaults to None.
     """
 
     def __init__(
@@ -58,6 +66,7 @@ class DMSLoader(CDFLoader[dm.InstanceApply]):
         create_issues: Sequence[NeatIssue] | None = None,
         tracker: type[Tracker] | None = None,
         rules: DMSRules | None = None,
+        client: NeatClient | None = None,
     ):
         super().__init__(graph_store)
         self.data_model = data_model
@@ -66,11 +75,12 @@ class DMSLoader(CDFLoader[dm.InstanceApply]):
         self._issues = IssueList(create_issues or [])
         self._tracker: type[Tracker] = tracker or LogTracker
         self.rules = rules
+        self._client = client
 
     @classmethod
     def from_data_model_id(
         cls,
-        client: CogniteClient,
+        client: NeatClient,
         data_model_id: dm.DataModelId,
         graph_store: NeatGraphStore,
         instance_space: str,
@@ -82,10 +92,12 @@ class DMSLoader(CDFLoader[dm.InstanceApply]):
         except Exception as e:
             issues.append(ResourceRetrievalError(data_model_id, "data model", str(e)))
 
-        return cls(graph_store, data_model, instance_space, {}, issues)
+        return cls(graph_store, data_model, instance_space, {}, issues, client=client)
 
     @classmethod
-    def from_rules(cls, rules: DMSRules, graph_store: NeatGraphStore, instance_space: str) -> "DMSLoader":
+    def from_rules(
+        cls, rules: DMSRules, graph_store: NeatGraphStore, instance_space: str, client: NeatClient | None = None
+    ) -> "DMSLoader":
         issues: list[NeatIssue] = []
         data_model: dm.DataModel[dm.View] | None = None
         try:
@@ -109,9 +121,10 @@ class DMSLoader(CDFLoader[dm.InstanceApply]):
             class_neat_id_by_view_id,
             issues,
             rules=rules,
+            client=client,
         )
 
-    def _load(self, stop_on_exception: bool = False) -> Iterable[dm.InstanceApply | NeatIssue]:
+    def _load(self, stop_on_exception: bool = False) -> Iterable[dm.InstanceApply | NeatIssue | type[_END_OF_CLASS]]:
         if self._issues.has_errors and stop_on_exception:
             raise self._issues.as_exception()
         elif self._issues.has_errors:
@@ -126,45 +139,73 @@ 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
+            )
+        else:
+            properties_point_to_self = {}
+
+        view_ids: list[str] = []
+        for view_id in view_and_count_by_id.keys():
+            view_ids.append(repr(view_id))
+            if view_id in properties_point_to_self:
+                # If the views have a dependency on themselves, we need to run it twice.
+                view_ids.append(f"{view_id!r} (self)")
 
-        view_ids = [repr(v.as_id()) for v in self.data_model.views]
         tracker = self._tracker(type(self).__name__, view_ids, "views")
-        for view in self.data_model.views:
-            view_id = view.as_id()
-            tracker.start(repr(view_id))
+        for view_id, (view, _) in view_and_count_by_id.items():
             pydantic_cls, edge_by_type, issues = self._create_validation_classes(view)  # type: ignore[var-annotated]
             yield from issues
             tracker.issue(issues)
 
-            # this assumes no changes in the suffix of view and class
-
-            if views_with_linked_properties:
-                # we need graceful exit if the view is not in the view_property_pairs
-                property_link_pairs = views_with_linked_properties.get(ViewEntity.from_id(view_id))
-
-                if class_neat_id := self.class_neat_id_by_view_id.get(view_id):
-                    reader = self.graph_store._read_via_rules_linkage(class_neat_id, property_link_pairs)
-                else:
-                    error_view = ResourceRetrievalError(view_id, "view", "View not linked to class")
-                    tracker.issue(error_view)
-                    if stop_on_exception:
-                        raise error_view
-                    yield error_view
-
+            if view_id in properties_point_to_self:
+                # If the view has a dependency on itself, we need to run it twice.
+                # First, to ensure that all nodes are created, and then to add the direct relations.
+                # This only applies if there is a require constraint on the container, if not
+                # we can create an empty node on the fly.
+                iterations = [properties_point_to_self[view_id], set()]
             else:
-                reader = self.graph_store.read(view.external_id)
+                iterations = [set()]
 
-            for identifier, properties in reader:
-                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)
-            tracker.finish(repr(view_id))
+            for skip_properties in iterations:
+                if skip_properties:
+                    track_id = f"{view_id} (self)"
+                else:
+                    track_id = repr(view_id)
+                tracker.start(track_id)
+                if views_with_linked_properties:
+                    # we need graceful exit if the view is not in the view_property_pairs
+                    property_link_pairs = views_with_linked_properties.get(ViewEntity.from_id(view_id))
+
+                    if class_neat_id := self.class_neat_id_by_view_id.get(view_id):
+                        reader = self.graph_store._read_via_rules_linkage(class_neat_id, property_link_pairs)
+                    else:
+                        error_view = ResourceRetrievalError(view_id, "view", "View not linked to class")
+                        tracker.issue(error_view)
+                        if stop_on_exception:
+                            raise error_view
+                        yield error_view
+                        continue
+                else:
+                    # this assumes no changes in the suffix of view and class
+                    reader = self.graph_store.read(view.external_id)
+
+                for identifier, properties in reader:
+                    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)
+                tracker.finish(track_id)
+                yield _END_OF_CLASS
 
     def write_to_file(self, filepath: Path) -> None:
         if filepath.suffix not in [".json", ".yaml", ".yml"]:
@@ -186,6 +227,68 @@ class DMSLoader(CDFLoader[dm.InstanceApply]):
             else:
                 yaml.safe_dump(dumped, f, sort_keys=False)
 
+    def _select_views_with_instances(self, views: list[dm.View]) -> dict[dm.ViewId, tuple[dm.View, int]]:
+        """Selects the views with data."""
+        view_and_count_by_id: dict[dm.ViewId, tuple[dm.View, int]] = {}
+        uri_by_type: dict[str, URIRef] = {
+            remove_namespace_from_uri(uri[0]): uri[0]  # type: ignore[misc]
+            for uri in self.graph_store.queries.list_types()
+        }
+        for view in views:
+            view_id = view.as_id()
+            neat_id = self.class_neat_id_by_view_id.get(view_id)
+            if neat_id is not None:
+                count = self.graph_store.count_of_id(neat_id)
+            elif view.external_id in uri_by_type:
+                count = self.graph_store.count_of_type(uri_by_type[view.external_id])
+            else:
+                continue
+            if count > 0:
+                view_and_count_by_id[view_id] = view, count
+
+        return view_and_count_by_id
+
+    def _sort_by_direct_relation_dependencies(
+        self, view_and_count_by_id: dict[dm.ViewId, tuple[dm.View, int]]
+    ) -> tuple[dict[dm.ViewId, tuple[dm.View, int]], dict[dm.ViewId, set[str]]]:
+        """Sorts the views by container constraints."""
+        if not self._client:
+            return view_and_count_by_id, {}
+        # We need to retrieve the views to ensure we get all properties, such that we can find all
+        # the containers that the view is linked to.
+        views = self._client.data_modeling.views.retrieve(
+            list(view_and_count_by_id.keys()), include_inherited_properties=True
+        )
+        container_ids_by_view_id = {view.as_id(): view.referenced_containers() for view in views}
+        referenced_containers = {
+            container for containers in container_ids_by_view_id.values() for container in containers
+        }
+        containers = self._client.data_modeling.containers.retrieve(list(referenced_containers))
+        container_by_id = {container.as_id(): container for container in containers}
+
+        dependency_on_self: dict[dm.ViewId, set[str]] = defaultdict(set)
+        view_id_by_dependencies: dict[dm.ViewId, set[dm.ViewId]] = {}
+        for view in views:
+            view_id = view.as_id()
+            dependencies = set()
+            for prop_id, prop in view.properties.items():
+                if isinstance(prop, dm.MappedProperty) and isinstance(prop.type, dm.DirectRelation) and prop.source:
+                    container = container_by_id[prop.container]
+                    has_require_constraint = any(
+                        isinstance(constraint, dm.RequiresConstraint) for constraint in container.constraints.values()
+                    )
+                    if has_require_constraint and prop.source == view_id:
+                        dependency_on_self[view_id].add(prop_id)
+                    elif has_require_constraint:
+                        dependencies.add(prop.source)
+            view_id_by_dependencies[view_id] = dependencies
+
+        ordered_view_ids = TopologicalSorter(view_id_by_dependencies).static_order()
+
+        return {
+            view_id: view_and_count_by_id[view_id] for view_id in ordered_view_ids if view_id in view_and_count_by_id
+        }, dict(dependency_on_self)
+
     def _create_validation_classes(
         self, view: dm.View
     ) -> tuple[type[BaseModel], dict[str, tuple[str, dm.EdgeConnection]], NeatIssueList]:
@@ -194,13 +297,17 @@ class DMSLoader(CDFLoader[dm.InstanceApply]):
         edge_by_property: 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] = []
-        for prop_name, prop in view.properties.items():
+        for prop_id, prop in view.properties.items():
             if isinstance(prop, dm.EdgeConnection):
-                edge_by_property[prop_name] = prop_name, prop
+                edge_by_property[prop_id] = prop_id, prop
             if isinstance(prop, dm.MappedProperty):
                 if isinstance(prop.type, dm.DirectRelation):
-                    direct_relation_by_property[prop_name] = prop.type
+                    if prop.container == dm.ContainerId("cdf_cdm", "CogniteTimeSeries") and prop_id == "unit":
+                        unit_properties.append(prop_id)
+                    else:
+                        direct_relation_by_property[prop_id] = prop.type
                     python_type: Any = dict
                 else:
                     data_type = _DATA_TYPE_BY_DMS_TYPE.get(prop.type._type)
@@ -209,14 +316,14 @@ class DMSLoader(CDFLoader[dm.InstanceApply]):
                             PropertyTypeNotSupportedWarning(
                                 view.as_id(),
                                 "view",
-                                prop_name,
+                                prop_id,
                                 prop.type._type,
                             )
                         )
                         continue
 
                     if data_type == Json:
-                        json_fields.append(prop_name)
+                        json_fields.append(prop_id)
                     python_type = data_type.python
                 if isinstance(prop.type, ListablePropertyType) and prop.type.is_list:
                     python_type = list[python_type]
@@ -226,7 +333,7 @@ class DMSLoader(CDFLoader[dm.InstanceApply]):
                 else:
                     default_value = ...
 
-                field_definitions[prop_name] = (python_type, default_value)
+                field_definitions[prop_id] = (python_type, default_value)
 
         def parse_list(cls, value: Any, info: ValidationInfo) -> list[str]:
             if isinstance(value, list) and list.__name__ not in _get_field_value_types(cls, info):
@@ -257,15 +364,26 @@ class DMSLoader(CDFLoader[dm.InstanceApply]):
             def parse_direct_relation(cls, value: list, info: ValidationInfo) -> dict | list[dict]:
                 # We validate above that we only get one value for single direct relations.
                 if list.__name__ in _get_field_value_types(cls, info):
-                    return [{"space": self.instance_space, "externalId": v} for v in value]
+                    return [{"space": self.instance_space, "externalId": remove_namespace_from_uri(v)} for v in value]
                 elif value:
-                    return {"space": self.instance_space, "externalId": value[0]}
+                    return {"space": self.instance_space, "externalId": remove_namespace_from_uri(value[0])}
                 return {}
 
             validators["parse_direct_relation"] = field_validator(*direct_relation_by_property.keys(), mode="before")(  # type: ignore[assignment]
                 parse_direct_relation  # type: ignore[arg-type]
             )
 
+        if unit_properties:
+
+            def parse_direct_relation_to_unit(cls, value: Any, info: ValidationInfo) -> dict | list[dict]:
+                if value:
+                    return {"space": "cdf_cdm_units", "externalId": remove_namespace_from_uri(value[0])}
+                return {}
+
+            validators["parse_direct_relation_to_unit"] = field_validator(*unit_properties, mode="before")(  # type: ignore[assignment]
+                parse_direct_relation_to_unit  # type: ignore[arg-type]
+            )
+
         pydantic_cls = create_model(view.external_id, __validators__=validators, **field_definitions)  # type: ignore[arg-type, call-overload]
         return pydantic_cls, edge_by_property, issues
 
@@ -371,8 +489,8 @@ class DMSLoader(CDFLoader[dm.InstanceApply]):
                 (upserted.edges, source_by_edge_id),
             ]:
                 for name, subinstances in itertools.groupby(
-                    sorted(instances, key=lambda i: ids_by_source[i.as_id()]),  # type: ignore[call-overload, index, attr-defined]
-                    key=lambda i: ids_by_source[i.as_id()],  # type: ignore[index]
+                    sorted(instances, key=lambda i: ids_by_source.get(i.as_id(), "")),  # type: ignore[call-overload, index, attr-defined]
+                    key=lambda i: ids_by_source.get(i.as_id(), ""),  # type: ignore[index, attr-defined]
                 ):
                     result = UploadResult(name=name, issues=read_issues)
                     for instance in subinstances:  # type: ignore[attr-defined]

{cognite_neat-0.102.0 → cognite_neat-0.103.0}/cognite/neat/_graph/transformers/__init__.py

@@ -15,24 +15,26 @@ from ._prune_graph import (
     PruneTypes,
 )
 from ._rdfpath import AddSelfReferenceProperty, MakeConnectionOnExactMatch
-from ._value_type import SplitMultiValueProperty
+from ._value_type import ConvertLiteral, LiteralToEntity, SplitMultiValueProperty
 
 __all__ = [
     "AddAssetDepth",
-    "AssetTimeSeriesConnector",
-    "AssetSequenceConnector",
-    "AssetFileConnector",
+    "AddSelfReferenceProperty",
     "AssetEventConnector",
+    "AssetFileConnector",
     "AssetRelationshipConnector",
-    "AddSelfReferenceProperty",
-    "SplitMultiValueProperty",
-    "RelationshipAsEdgeTransformer",
-    "MakeConnectionOnExactMatch",
+    "AssetSequenceConnector",
+    "AssetTimeSeriesConnector",
     "AttachPropertyFromTargetToSource",
+    "ConvertLiteral",
+    "LiteralToEntity",
+    "MakeConnectionOnExactMatch",
     "PruneDanglingNodes",
-    "PruneTypes",
     "PruneDeadEndEdges",
     "PruneInstancesOfUnknownType",
+    "PruneTypes",
+    "RelationshipAsEdgeTransformer",
+    "SplitMultiValueProperty",
 ]
 
 Transformers = (
@@ -51,4 +53,6 @@ Transformers = (
     | PruneTypes
     | PruneDeadEndEdges
     | PruneInstancesOfUnknownType
+    | ConvertLiteral
+    | LiteralToEntity
 )