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

cognite/neat/_graph/transformers/_classic_cdf.py

@@ -1,4 +1,3 @@
-import textwrap
 import warnings
 from abc import ABC
 from collections.abc import Callable, Iterable
@@ -15,6 +14,7 @@ from cognite.neat._utils.collection_ import iterate_progress_bar
 from cognite.neat._utils.rdf_ import (
     Triple,
     add_triples_in_batch,
+    get_namespace,
     remove_instance_ids_in_batch,
     remove_namespace_from_uri,
 )
@@ -72,91 +72,165 @@ class AddAssetDepth(BaseTransformerStandardised):
         return row_output
 
 
-# TODO: standardise
-class BaseAssetConnector(BaseTransformer, ABC):
-    _asset_type: URIRef = DEFAULT_NAMESPACE.Asset
-    _item_type: URIRef
-    _default_attribute: URIRef
-    _connection_type: URIRef
+class BaseAssetConnector(BaseTransformerStandardised, ABC):
+    description: str = "Connects assets to other cognite resources, thus forming bi-directional connection"
+    _use_only_once: bool = True
 
-    _select_item_ids = "SELECT DISTINCT ?item_id WHERE {{?item_id a <{item_type}>}}"
-    _select_connected_assets: str = textwrap.dedent("""SELECT ?asset_id WHERE {{
-                              <{item_id}> <{attribute}> ?asset_id .
-                              ?asset_id a <{asset_type}>}}""")
+    def _count_query(self) -> str:
+        query = """SELECT (COUNT(?asset) as ?count)
+                   WHERE {{
+                        ?resource a <{resource_type}> .
+                        ?resource <{connection}> ?asset .
+                        ?asset a <{asset_type}> .
+                        }}"""
 
-    def __init__(self, attribute: URIRef | None = None) -> None:
-        self._attribute = attribute or self._default_attribute
+        return query.format(
+            asset_type=self.asset_type,
+            resource_type=self.resource_type,
+            connection=self.resource_to_asset_connection,
+        )
 
-    def transform(self, graph: Graph) -> None:
-        for item_id, *_ in graph.query(self._select_item_ids.format(item_type=self._item_type)):  # type: ignore[misc]
-            triples: list[Triple] = []
-            for asset_id, *_ in graph.query(  # type: ignore[misc]
-                self._select_connected_assets.format(
-                    item_id=item_id, attribute=self._attribute, asset_type=self._asset_type
-                )
-            ):
-                triples.append((asset_id, self._connection_type, item_id))  # type: ignore[arg-type]
-            add_triples_in_batch(graph, triples)
+    def _iterate_query(self) -> str:
+        query = """SELECT ?asset ?resource
+                   WHERE {{
+                        ?resource a <{resource_type}> .
+                        ?resource <{connection}> ?asset .
+                        ?asset a <{asset_type}> .
+                        }}"""
+
+        return query.format(
+            asset_type=self.asset_type,
+            resource_type=self.resource_type,
+            connection=self.resource_to_asset_connection,
+        )
+
+    def __init__(
+        self,
+        resource_to_asset_connection: URIRef,
+        resource_type: URIRef,
+        asset_to_resource_connection: URIRef | None = None,
+        asset_type: URIRef | None = None,
+    ) -> None:
+        self.asset_type = asset_type or DEFAULT_NAMESPACE.Asset
+        self.resource_to_asset_connection = resource_to_asset_connection
+        self.resource_type = resource_type
+
+        if asset_to_resource_connection:
+            self.asset_to_resource_connection = asset_to_resource_connection
+        else:
+            namespace = Namespace(get_namespace(resource_type))
+            type_ = remove_namespace_from_uri(resource_type)
+            self.asset_to_resource_connection = namespace[type_[0].lower() + type_[1:]]
+
+    def operation(self, query_result_row: ResultRow) -> RowTransformationOutput:
+        row_output = RowTransformationOutput()
+        subject, object = query_result_row
+
+        row_output.add_triples.append(cast(Triple, (subject, self.asset_to_resource_connection, object)))
+
+        row_output.instances_modified_count += 1
+
+        return row_output
 
 
 class AssetTimeSeriesConnector(BaseAssetConnector):
     description: str = "Connects assets to timeseries, thus forming bi-directional connection"
-    _use_only_once: bool = True
     _need_changes = frozenset(
         {
             str(extractors.AssetsExtractor.__name__),
             str(extractors.TimeSeriesExtractor.__name__),
         }
     )
-    _item_type = DEFAULT_NAMESPACE.TimeSeries
-    _default_attribute = DEFAULT_NAMESPACE.assetId
-    _connection_type = DEFAULT_NAMESPACE.timeSeries
+
+    def __init__(
+        self,
+        resource_to_asset_connection: URIRef | None = None,
+        resource_type: URIRef | None = None,
+        asset_to_resource_connection: URIRef | None = None,
+        asset_type: URIRef | None = None,
+    ):
+        super().__init__(
+            resource_to_asset_connection=resource_to_asset_connection or DEFAULT_NAMESPACE.assetId,
+            resource_type=resource_type or DEFAULT_NAMESPACE.TimeSeries,
+            asset_to_resource_connection=asset_to_resource_connection or DEFAULT_NAMESPACE.timeSeries,
+            asset_type=asset_type or DEFAULT_NAMESPACE.Asset,
+        )
 
 
 class AssetSequenceConnector(BaseAssetConnector):
     description: str = "Connects assets to sequences, thus forming bi-directional connection"
-    _use_only_once: bool = True
     _need_changes = frozenset(
         {
             str(extractors.AssetsExtractor.__name__),
             str(extractors.SequencesExtractor.__name__),
         }
     )
-    _item_type = DEFAULT_NAMESPACE.Sequence
-    _default_attribute = DEFAULT_NAMESPACE.assetId
-    _connection_type = DEFAULT_NAMESPACE.sequence
+
+    def __init__(
+        self,
+        resource_to_asset_connection: URIRef | None = None,
+        resource_type: URIRef | None = None,
+        asset_to_resource_connection: URIRef | None = None,
+        asset_type: URIRef | None = None,
+    ):
+        super().__init__(
+            resource_to_asset_connection=resource_to_asset_connection or DEFAULT_NAMESPACE.assetId,
+            resource_type=resource_type or DEFAULT_NAMESPACE.Sequence,
+            asset_to_resource_connection=asset_to_resource_connection or DEFAULT_NAMESPACE.sequence,
+            asset_type=asset_type or DEFAULT_NAMESPACE.Asset,
+        )
 
 
 class AssetFileConnector(BaseAssetConnector):
     description: str = "Connects assets to files, thus forming bi-directional connection"
-    _use_only_once: bool = True
     _need_changes = frozenset(
         {
             str(extractors.AssetsExtractor.__name__),
             str(extractors.FilesExtractor.__name__),
         }
     )
-    _item_type = DEFAULT_NAMESPACE.File
-    _default_attribute = DEFAULT_NAMESPACE.assetIds
-    _connection_type = DEFAULT_NAMESPACE.file
+
+    def __init__(
+        self,
+        resource_to_asset_connection: URIRef | None = None,
+        resource_type: URIRef | None = None,
+        asset_to_resource_connection: URIRef | None = None,
+        asset_type: URIRef | None = None,
+    ):
+        super().__init__(
+            resource_to_asset_connection=resource_to_asset_connection or DEFAULT_NAMESPACE.assetIds,
+            resource_type=resource_type or DEFAULT_NAMESPACE.File,
+            asset_to_resource_connection=asset_to_resource_connection or DEFAULT_NAMESPACE.file,
+            asset_type=asset_type or DEFAULT_NAMESPACE.Asset,
+        )
 
 
 class AssetEventConnector(BaseAssetConnector):
     description: str = "Connects assets to events, thus forming bi-directional connection"
-    _use_only_once: bool = True
     _need_changes = frozenset(
         {
             str(extractors.AssetsExtractor.__name__),
             str(extractors.EventsExtractor.__name__),
         }
     )
-    _item_type = DEFAULT_NAMESPACE.Event
-    _default_attribute = DEFAULT_NAMESPACE.assetIds
-    _connection_type = DEFAULT_NAMESPACE.event
+
+    def __init__(
+        self,
+        resource_to_asset_connection: URIRef | None = None,
+        resource_type: URIRef | None = None,
+        asset_to_resource_connection: URIRef | None = None,
+        asset_type: URIRef | None = None,
+    ):
+        super().__init__(
+            resource_to_asset_connection=resource_to_asset_connection or DEFAULT_NAMESPACE.assetIds,
+            resource_type=resource_type or DEFAULT_NAMESPACE.Event,
+            asset_to_resource_connection=asset_to_resource_connection or DEFAULT_NAMESPACE.event,
+            asset_type=asset_type or DEFAULT_NAMESPACE.Asset,
+        )
 
 
 # TODO: standardise
-class AssetRelationshipConnector(BaseTransformer):
+class AssetRelationshipConnector(BaseTransformerStandardised):
     description: str = "Connects assets via relationships"
     _use_only_once: bool = True
     _need_changes = frozenset(
@@ -174,6 +248,44 @@ class AssetRelationshipConnector(BaseTransformer):
                               ?target <{asset_xid_property}> ?target_xid .
                               ?target a <{asset_type}> .}}"""
 
+    def _count_query(self) -> str:
+        query = """SELECT (COUNT(?target) as ?count) WHERE {{
+                   ?relationship a <{relationship_type}> .
+                   ?relationship <{relationship_source_xid_prop}> ?source_xid .
+                   ?source <{asset_xid_property}> ?source_xid .
+                   ?source a <{asset_type}> .
+
+                   ?relationship <{relationship_target_xid_prop}> ?target_xid .
+                   ?target <{asset_xid_property}> ?target_xid .
+                   ?target a <{asset_type}> .}}"""
+
+        return query.format(
+            relationship_type=self.relationship_type,
+            relationship_source_xid_prop=self.relationship_source_xid_prop,
+            relationship_target_xid_prop=self.relationship_target_xid_prop,
+            asset_xid_property=self.asset_xid_property,
+            asset_type=self.asset_type,
+        )
+
+    def _iterate_query(self) -> str:
+        query = """SELECT ?source ?relationship ?target WHERE {{
+                   ?relationship a <{relationship_type}> .
+                   ?relationship <{relationship_source_xid_prop}> ?source_xid .
+                   ?source <{asset_xid_property}> ?source_xid .
+                   ?source a <{asset_type}> .
+
+                   ?relationship <{relationship_target_xid_prop}> ?target_xid .
+                   ?target <{asset_xid_property}> ?target_xid .
+                   ?target a <{asset_type}> .}}"""
+
+        return query.format(
+            relationship_type=self.relationship_type,
+            relationship_source_xid_prop=self.relationship_source_xid_prop,
+            relationship_target_xid_prop=self.relationship_target_xid_prop,
+            asset_xid_property=self.asset_xid_property,
+            asset_type=self.asset_type,
+        )
+
     def __init__(
         self,
         asset_type: URIRef | None = None,
@@ -188,48 +300,20 @@ class AssetRelationshipConnector(BaseTransformer):
         self.relationship_target_xid_prop = relationship_target_xid_prop or DEFAULT_NAMESPACE.targetExternalId
         self.asset_xid_property = asset_xid_property or DEFAULT_NAMESPACE.externalId
 
-    def transform(self, graph: Graph) -> None:
-        for relationship_id_result in graph.query(
-            f"SELECT DISTINCT ?relationship_id WHERE {{?relationship_id a <{self.relationship_type}>}}"
-        ):
-            relationship_id: URIRef = cast(tuple, relationship_id_result)[0]
-
-            if assets_id_res := list(
-                graph.query(
-                    self._asset_template.format(
-                        relationship_id=relationship_id,
-                        asset_xid_property=self.asset_xid_property,
-                        relationship_source_xid_prop=self.relationship_source_xid_prop,
-                        relationship_target_xid_prop=self.relationship_target_xid_prop,
-                        asset_type=self.asset_type,
-                    )
-                )
-            ):
-                # files can be connected to multiple assets in the graph
-                for source_asset_id, target_asset_id in cast(list[tuple], assets_id_res):
-                    # create a relationship between the two assets
-                    graph.add(
-                        (
-                            source_asset_id,
-                            DEFAULT_NAMESPACE.relationship,
-                            relationship_id,
-                        )
-                    )
-                    graph.add(
-                        (
-                            target_asset_id,
-                            DEFAULT_NAMESPACE.relationship,
-                            relationship_id,
-                        )
-                    )
+    def operation(self, query_result_row: ResultRow) -> RowTransformationOutput:
+        row_output = RowTransformationOutput()
+        source, relationship, target = query_result_row
+
+        row_output.add_triples.append(cast(Triple, (source, DEFAULT_NAMESPACE.relationship, target)))
+        row_output.add_triples.append(cast(Triple, (relationship, DEFAULT_NAMESPACE.source, source)))
+        row_output.add_triples.append(cast(Triple, (relationship, DEFAULT_NAMESPACE.target, target)))
 
-                    # add source and target to the relationship
-                    graph.add((relationship_id, DEFAULT_NAMESPACE.source, source_asset_id))
-                    graph.add((relationship_id, DEFAULT_NAMESPACE.target, target_asset_id))
+        row_output.remove_triples.append(cast(Triple, (relationship, self.relationship_source_xid_prop, None)))
+        row_output.remove_triples.append(cast(Triple, (relationship, self.relationship_target_xid_prop, None)))
 
-                    # remove properties that are not needed, specifically the external ids
-                    graph.remove((relationship_id, self.relationship_source_xid_prop, None))
-                    graph.remove((relationship_id, self.relationship_target_xid_prop, None))
+        row_output.instances_modified_count += 2
+
+        return row_output
 
 
 # TODO: standardise

cognite/neat/_graph/transformers/_rdfpath.py

@@ -59,7 +59,7 @@ class AddSelfReferenceProperty(BaseTransformer):
 
 class MakeConnectionOnExactMatch(BaseTransformerStandardised):
     description: str = "Adds property that contains id of reference to all references of given class in Rules"
-    _use_only_once: bool = True
+    _use_only_once: bool = False
     _need_changes = frozenset({})
 
     def __init__(

cognite/neat/_issues/warnings/_external.py

@@ -42,7 +42,7 @@ class FileItemNotSupportedWarning(NeatWarning):
 
 @dataclass(unsafe_hash=True)
 class CDFAuthWarning(NeatWarning):
-    """Failed to {action} due to {reason}"""
+    """Failed to {action}: {reason}"""
 
     action: str
     reason: str

cognite/neat/_rules/importers/_rdf/_inference2rules.py

@@ -17,6 +17,7 @@ from cognite.neat._rules.models.information import (
 )
 from cognite.neat._store import NeatGraphStore
 from cognite.neat._store._provenance import INSTANCES_ENTITY
+from cognite.neat._utils.collection_ import iterate_progress_bar
 from cognite.neat._utils.rdf_ import remove_namespace_from_uri, uri_to_short_form
 
 from ._base import DEFAULT_NON_EXISTING_NODE_TYPE, BaseRDFImporter
@@ -27,7 +28,6 @@ ORDERED_CLASSES_QUERY = """SELECT ?class (count(?s) as ?instances )
                            WHERE { ?s a ?class . }
                            group by ?class order by DESC(?instances)"""
 
-
 INSTANCES_OF_CLASS_QUERY = """SELECT ?s ?propertyCount WHERE { ?s a <class> . BIND ('Unknown' as ?propertyCount) }"""
 
 
@@ -171,8 +171,10 @@ class InferenceImporter(BaseRDFImporter):
             INSTANCES_OF_CLASS_QUERY if self.max_number_of_instance == -1 else INSTANCES_OF_CLASS_RICHNESS_ORDERED_QUERY
         )
 
+        classes_iterable = iterate_progress_bar(classes.items(), len(classes), "Inferring classes")
+
         # Infers all the properties of the class
-        for class_id, class_definition in classes.items():
+        for class_id, class_definition in classes_iterable:
             for (  # type: ignore[misc]
                 instance,
                 _,