cognite-neat 0.105.2__tar.gz → 0.106.0__tar.gz

{cognite_neat-0.105.2 → cognite_neat-0.106.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: cognite-neat
-Version: 0.105.2
+Version: 0.106.0
 Summary: Knowledge graph transformation
 Home-page: https://cognite-neat.readthedocs-hosted.com/
 License: Apache-2.0
@@ -17,13 +17,11 @@ Provides-Extra: all
 Provides-Extra: docs
 Provides-Extra: google
 Provides-Extra: oxi
-Provides-Extra: service
 Requires-Dist: PyYAML
 Requires-Dist: backports.strenum (>=1.2,<2.0) ; python_version < "3.11"
 Requires-Dist: cognite-sdk (>=7.71.2,<8.0.0)
 Requires-Dist: elementpath (>=4.0.0,<5.0.0)
 Requires-Dist: exceptiongroup (>=1.1.3,<2.0.0) ; python_version < "3.11"
-Requires-Dist: fastapi (>=0,<1) ; extra == "service" or extra == "all"
 Requires-Dist: google-api-python-client ; extra == "google"
 Requires-Dist: google-auth-oauthlib ; extra == "google"
 Requires-Dist: gspread ; extra == "google"
@@ -45,20 +43,16 @@ Requires-Dist: openpyxl
 Requires-Dist: oxrdflib[oxigraph] (>=0.4.0,<0.5.0) ; extra == "oxi" or extra == "all"
 Requires-Dist: packaging (>=22.0,<25.0)
 Requires-Dist: pandas
-Requires-Dist: prometheus-client (>=0,<1) ; extra == "service" or extra == "all"
 Requires-Dist: pydantic (>=2,<3)
 Requires-Dist: pymdown-extensions ; extra == "docs"
 Requires-Dist: pyoxigraph (==0.4.3) ; extra == "oxi" or extra == "all"
-Requires-Dist: python-multipart (==0.0.9) ; extra == "service" or extra == "all"
 Requires-Dist: pyvis (>=0.3.2,<0.4.0)
 Requires-Dist: rdflib
 Requires-Dist: requests
 Requires-Dist: rich[jupyter] (>=13.7.1,<14.0.0)
-Requires-Dist: schedule (>=1,<2) ; extra == "service" or extra == "all"
 Requires-Dist: tomli (>=2.0.1,<3.0.0) ; python_version < "3.11"
 Requires-Dist: typing_extensions (>=4.8,<5.0) ; python_version < "3.11"
 Requires-Dist: urllib3 (>=2,<3)
-Requires-Dist: uvicorn[standard] (>=0,<1) ; extra == "service" or extra == "all"
 Project-URL: Documentation, https://cognite-neat.readthedocs-hosted.com/
 Project-URL: Repository, https://github.com/cognitedata/neat
 Description-Content-Type: text/markdown

cognite_neat-0.106.0/cognite/neat/_config.py

@@ -0,0 +1,11 @@
+from typing import Literal
+
+from pydantic import BaseModel
+
+
+class NeatConfig(BaseModel, validate_assignment=True):
+    progress_bar: Literal["tqdm", "rich", "tqdm-notebook", "infer"] | None = "infer"
+    use_iterate_bar_threshold: int | None = 500
+
+
+GLOBAL_CONFIG = NeatConfig()

{cognite_neat-0.105.2 → cognite_neat-0.106.0}/cognite/neat/_graph/extractors/_classic_cdf/_base.py

@@ -1,6 +1,7 @@
 import json
 import re
 import sys
+import warnings
 from abc import ABC, abstractmethod
 from collections.abc import Callable, Iterable, Sequence, Set
 from datetime import datetime, timezone
@@ -9,13 +10,16 @@ from typing import Any, Generic, TypeVar
 
 from cognite.client import CogniteClient
 from cognite.client.data_classes._base import WriteableCogniteResource
+from cognite.client.exceptions import CogniteAPIError
 from pydantic import AnyHttpUrl, ValidationError
 from rdflib import RDF, XSD, Literal, Namespace, URIRef
 
 from cognite.neat._constants import DEFAULT_NAMESPACE
 from cognite.neat._graph.extractors._base import BaseExtractor
+from cognite.neat._issues.warnings import CDFAuthWarning
 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_if_above_config_threshold
 
 T_CogniteResource = TypeVar("T_CogniteResource", bound=WriteableCogniteResource)
 
@@ -98,17 +102,11 @@ class ClassicCDFBaseExtractor(BaseExtractor, ABC, Generic[T_CogniteResource]):
 
     def extract(self) -> Iterable[Triple]:
         """Extracts an asset with the given asset_id."""
-        if self.total:
-            try:
-                from rich.progress import track
-            except ModuleNotFoundError:
-                to_iterate = self.items
-            else:
-                to_iterate = track(
-                    self.items,
-                    total=self.limit or self.total,
-                    description=f"Extracting {type(self).__name__.removesuffix('Extractor')}",
-                )
+
+        if self.total is not None and self.total > 0:
+            to_iterate = iterate_progress_bar_if_above_config_threshold(
+                self.items, self.total, f"Extracting {type(self).__name__.removesuffix('Extractor')}"
+            )
         else:
             to_iterate = self.items
         for no, asset in enumerate(to_iterate):
@@ -221,7 +219,7 @@ class ClassicCDFBaseExtractor(BaseExtractor, ABC, Generic[T_CogniteResource]):
         camel_case: bool = True,
         as_write: bool = False,
     ):
-        total, items = cls._from_dataset(client, data_set_external_id)
+        total, items = cls._handle_no_access(lambda: cls._from_dataset(client, data_set_external_id))
         return cls(items, namespace, to_type, total, limit, unpack_metadata, skip_metadata_values, camel_case, as_write)
 
     @classmethod
@@ -244,7 +242,7 @@ class ClassicCDFBaseExtractor(BaseExtractor, ABC, Generic[T_CogniteResource]):
         camel_case: bool = True,
         as_write: bool = False,
     ):
-        total, items = cls._from_hierarchy(client, root_asset_external_id)
+        total, items = cls._handle_no_access(lambda: cls._from_hierarchy(client, root_asset_external_id))
         return cls(items, namespace, to_type, total, limit, unpack_metadata, skip_metadata_values, camel_case, as_write)
 
     @classmethod
@@ -273,3 +271,18 @@ class ClassicCDFBaseExtractor(BaseExtractor, ABC, Generic[T_CogniteResource]):
     @abstractmethod
     def _from_file(cls, file_path: str | Path) -> tuple[int | None, Iterable[T_CogniteResource]]:
         raise NotImplementedError
+
+    @classmethod
+    def _handle_no_access(
+        cls, action: Callable[[], tuple[int | None, Iterable[T_CogniteResource]]]
+    ) -> tuple[int | None, Iterable[T_CogniteResource]]:
+        try:
+            return action()
+        except CogniteAPIError as e:
+            if e.code == 403:
+                warnings.warn(
+                    CDFAuthWarning(f"extract {cls.__name__.removesuffix('Extractor').casefold()}", str(e)), stacklevel=2
+                )
+                return 0, []
+            else:
+                raise e

{cognite_neat-0.105.2 → cognite_neat-0.106.0}/cognite/neat/_graph/extractors/_classic_cdf/_classic.py

@@ -226,4 +226,7 @@ class ClassicGraphExtractor(BaseExtractor):
     @staticmethod
     def _chunk(items: Sequence, description: str) -> Iterable:
         to_iterate: Iterable = chunker(items, chunk_size=1000)
-        return iterate_progress_bar(to_iterate, (len(items) // 1_000) + 1, description)
+        if items:
+            return iterate_progress_bar(to_iterate, (len(items) // 1_000) + 1, description)
+        else:
+            return to_iterate

{cognite_neat-0.105.2 → cognite_neat-0.106.0}/cognite/neat/_graph/extractors/_classic_cdf/_sequences.py

@@ -73,7 +73,7 @@ class SequencesExtractor(ClassicCDFBaseExtractor[NeatSequence]):
         as_write: bool = False,
         unpack_columns: bool = False,
     ):
-        total, items = cls._from_dataset(client, data_set_external_id)
+        total, items = cls._handle_no_access(lambda: cls._from_dataset(client, data_set_external_id))
         return cls(
             items,
             namespace,
@@ -101,7 +101,7 @@ class SequencesExtractor(ClassicCDFBaseExtractor[NeatSequence]):
         as_write: bool = False,
         unpack_columns: bool = False,
     ):
-        total, items = cls._from_hierarchy(client, root_asset_external_id)
+        total, items = cls._handle_no_access(lambda: cls._from_hierarchy(client, root_asset_external_id))
         return cls(
             items,
             namespace,

{cognite_neat-0.105.2 → cognite_neat-0.106.0}/cognite/neat/_graph/loaders/_rdf2dms.py

@@ -37,6 +37,7 @@ 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.collection_ import iterate_progress_bar_if_above_config_threshold
 from cognite.neat._utils.rdf_ import remove_namespace_from_uri
 from cognite.neat._utils.upload import UploadResult
 
@@ -157,7 +158,7 @@ class DMSLoader(CDFLoader[dm.InstanceApply]):
                 view_ids.append(f"{view_id!r} (self)")
 
         tracker = self._tracker(type(self).__name__, view_ids, "views")
-        for view_id, (view, _) in view_and_count_by_id.items():
+        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]
             yield from issues
             tracker.issue(issues)
@@ -194,7 +195,11 @@ class DMSLoader(CDFLoader[dm.InstanceApply]):
                     # this assumes no changes in the suffix of view and class
                     reader = self.graph_store.read(view.external_id)
 
-                for identifier, properties in reader:
+                instance_iterable = iterate_progress_bar_if_above_config_threshold(
+                    reader, instance_count, f"Loading {track_id}"
+                )
+
+                for identifier, properties in instance_iterable:
                     if skip_properties:
                         properties = {k: v for k, v in properties.items() if k not in skip_properties}
                     try:

{cognite_neat-0.105.2 → cognite_neat-0.106.0}/cognite/neat/_graph/transformers/_base.py

@@ -8,7 +8,7 @@ from rdflib.query import ResultRow
 
 from cognite.neat._issues.warnings import NeatValueWarning
 from cognite.neat._shared import Triple
-from cognite.neat._utils.collection_ import iterate_progress_bar
+from cognite.neat._utils.collection_ import iterate_progress_bar_if_above_config_threshold
 from cognite.neat._utils.graph_transformations_report import GraphTransformationResult
 
 To_Add_Triples: TypeAlias = list[Triple]
@@ -42,7 +42,6 @@ class BaseTransformerStandardised(ABC):
     description: str
     _use_only_once: bool = False
     _need_changes: ClassVar[frozenset[str]] = frozenset()
-    _use_iterate_bar_threshold: int = 500
 
     @abstractmethod
     def operation(self, query_result_row: ResultRow) -> RowTransformationOutput:
@@ -99,12 +98,9 @@ class BaseTransformerStandardised(ABC):
             return outcome
 
         result_iterable = graph.query(self._iterate_query())
-        if iteration_count > self._use_iterate_bar_threshold:
-            result_iterable = iterate_progress_bar(  # type: ignore[misc, assignment]
-                result_iterable,
-                total=iteration_count,
-                description=self.description,
-            )
+        result_iterable = iterate_progress_bar_if_above_config_threshold(
+            result_iterable, iteration_count, self.description
+        )
 
         for row in result_iterable:
             row = cast(ResultRow, row)

{cognite_neat-0.105.2 → cognite_neat-0.106.0}/cognite/neat/_graph/transformers/_classic_cdf.py

@@ -230,7 +230,7 @@ class AssetEventConnector(BaseAssetConnector):
 
 
 # TODO: standardise
-class AssetRelationshipConnector(BaseTransformer):
+class AssetRelationshipConnector(BaseTransformerStandardised):
     description: str = "Connects assets via relationships"
     _use_only_once: bool = True
     _need_changes = frozenset(
@@ -248,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,
@@ -262,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
 
-                    # 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.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)))
 
-                    # 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.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)))
+
+        row_output.instances_modified_count += 2
+
+        return row_output
 
 
 # TODO: standardise

{cognite_neat-0.105.2 → cognite_neat-0.106.0}/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-0.105.2 → cognite_neat-0.106.0}/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-0.105.2 → cognite_neat-0.106.0}/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,
                 _,