cognite-neat 0.78.3__py3-none-any.whl → 0.78.5__py3-none-any.whl

cognite/neat/graph/loaders/_rdf2dms.py

@@ -0,0 +1,280 @@
+import itertools
+import json
+from collections import defaultdict
+from collections.abc import Iterable, Sequence
+from pathlib import Path
+from typing import Any
+
+import yaml
+from cognite.client import CogniteClient
+from cognite.client import data_modeling as dm
+from cognite.client.data_classes.capabilities import Capability, DataModelInstancesAcl
+from cognite.client.data_classes.data_modeling import ViewId
+from cognite.client.data_classes.data_modeling.views import SingleEdgeConnection
+from cognite.client.exceptions import CogniteAPIError
+from pydantic import ValidationInfo, create_model, field_validator
+from pydantic.main import Model
+
+from cognite.neat.graph._tracking import LogTracker, Tracker
+from cognite.neat.graph.issues import loader as loader_issues
+from cognite.neat.graph.stores import NeatGraphStoreBase
+from cognite.neat.issues import NeatIssue, NeatIssueList
+from cognite.neat.rules.models import DMSRules
+from cognite.neat.rules.models.data_types import _DATA_TYPE_BY_DMS_TYPE
+from cognite.neat.utils.upload import UploadDiffsID
+
+from ._base import CDFLoader
+
+
+class DMSLoader(CDFLoader[dm.InstanceApply]):
+    def __init__(
+        self,
+        graph_store: NeatGraphStoreBase,
+        data_model: dm.DataModel[dm.View] | None,
+        instance_space: str,
+        class_by_view_id: dict[ViewId, str] | None = None,
+        creat_issues: Sequence[NeatIssue] | None = None,
+        tracker: type[Tracker] | None = None,
+    ):
+        super().__init__(graph_store)
+        self.data_model = data_model
+        self.instance_space = instance_space
+        self.class_by_view_id = class_by_view_id or {}
+        self._issues = NeatIssueList[NeatIssue](creat_issues or [])
+        self._tracker: type[Tracker] = tracker or LogTracker
+
+    @classmethod
+    def from_data_model_id(
+        cls,
+        client: CogniteClient,
+        data_model_id: dm.DataModelId,
+        graph_store: NeatGraphStoreBase,
+        instance_space: str,
+    ) -> "DMSLoader":
+        issues: list[NeatIssue] = []
+        data_model: dm.DataModel[dm.View] | None = None
+        try:
+            data_model = client.data_modeling.data_models.retrieve(data_model_id, inline_views=True).latest_version()
+        except Exception as e:
+            issues.append(loader_issues.MissingDataModelError(identifier=repr(data_model_id), reason=str(e)))
+
+        return cls(graph_store, data_model, instance_space, {}, issues)
+
+    @classmethod
+    def from_rules(cls, rules: DMSRules, graph_store: NeatGraphStoreBase, instance_space: str) -> "DMSLoader":
+        issues: list[NeatIssue] = []
+        data_model: dm.DataModel[dm.View] | None = None
+        try:
+            data_model = rules.as_schema().as_read_model()
+        except Exception as e:
+            issues.append(
+                loader_issues.FailedConvertError(
+                    identifier=rules.metadata.as_identifier(), target_format="read DMS model", reason=str(e)
+                )
+            )
+        return cls(graph_store, data_model, instance_space, {}, issues)
+
+    def _load(self, stop_on_exception: bool = False) -> Iterable[dm.InstanceApply | NeatIssue]:
+        if self._issues.has_errors and stop_on_exception:
+            raise self._issues.as_exception()
+        elif self._issues.has_errors:
+            yield from self._issues
+            return
+        if not self.data_model:
+            # There should already be an error in this case.
+            return
+        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))
+            pydantic_cls, edge_by_properties, issues = self._create_validation_classes(view)  # type: ignore[var-annotated]
+            yield from issues
+            tracker.issue(issues)
+            class_name = self.class_by_view_id.get(view.as_id(), view.external_id)
+            triples = self.graph_store.queries.triples_of_type_instances(class_name)
+            for identifier, properties in _triples2dictionary(triples).items():
+                try:
+                    yield self._create_node(identifier, properties, pydantic_cls, view_id)
+                except ValueError as e:
+                    error = loader_issues.InvalidInstanceError(type_="node", identifier=identifier, reason=str(e))
+                    tracker.issue(error)
+                    if stop_on_exception:
+                        raise error.as_exception() from e
+                    yield error
+                yield from self._create_edges(identifier, properties, edge_by_properties, tracker)
+            tracker.finish(repr(view_id))
+
+    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")
+        dumped: dict[str, list] = {"nodes": [], "edges": [], "issues": []}
+        for item in self.load(stop_on_exception=False):
+            key = {
+                dm.NodeApply: "nodes",
+                dm.EdgeApply: "edges",
+                NeatIssue: "issues",
+            }.get(type(item))
+            if key is None:
+                # This should never happen, and is a bug in neat
+                raise ValueError(f"Item {item} is not supported. This is a bug in neat please report it.")
+            dumped[key].append(item.dump())
+        with filepath.open("w", encoding=self._encoding, newline=self._new_line) as f:
+            if filepath.suffix == ".json":
+                json.dump(dumped, f, indent=2)
+            else:
+                yaml.safe_dump(dumped, f, sort_keys=False)
+
+    def _create_validation_classes(
+        self, view: dm.View
+    ) -> tuple[type[Model], dict[str, dm.EdgeConnection], NeatIssueList]:
+        issues = NeatIssueList[NeatIssue]()
+        field_definitions: dict[str, tuple[type, Any]] = {}
+        edge_by_property: dict[str, dm.EdgeConnection] = {}
+        direct_relation_by_property: dict[str, dm.DirectRelation] = {}
+        for prop_name, prop in view.properties.items():
+            if isinstance(prop, dm.EdgeConnection):
+                edge_by_property[prop_name] = prop
+            if isinstance(prop, dm.MappedProperty):
+                if isinstance(prop.type, dm.DirectRelation):
+                    direct_relation_by_property[prop_name] = prop.type
+                    python_type: Any = dict
+                else:
+                    data_type = _DATA_TYPE_BY_DMS_TYPE.get(prop.type._type)
+                    if not data_type:
+                        issues.append(
+                            loader_issues.InvalidClassWarning(
+                                class_name=repr(view.as_id()),
+                                reason=f"Unknown data type for property {prop_name}: {prop.type._type}",
+                            )
+                        )
+                        continue
+                    python_type = data_type.python
+                if prop.type.is_list:
+                    python_type = list[python_type]
+                default_value: Any = prop.default_value
+                if prop.nullable:
+                    python_type = python_type | None
+                else:
+                    default_value = ...
+
+                field_definitions[prop_name] = (python_type, default_value)
+
+        def parse_list(cls, value: Any, info: ValidationInfo) -> list[str]:
+            if isinstance(value, list) and cls.model_fields[info.field_name].annotation is not list:
+                if len(value) == 1:
+                    return value[0]
+                raise ValueError(f"Got multiple values for {info.field_name}: {value}")
+            return value
+
+        validators: dict[str, classmethod] = {"parse_list": field_validator("*", mode="before")(parse_list)}  # type: ignore[dict-item,arg-type]
+        if direct_relation_by_property:
+
+            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 cls.model_fields[info.field_name].annotation is list:
+                    return [{"space": self.instance_space, "externalId": v} for v in value]
+                elif value:
+                    return {"space": self.instance_space, "externalId": 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]
+            )
+
+        pydantic_cls = create_model(view.external_id, __validators__=validators, **field_definitions)  # type: ignore[arg-type, call-overload]
+        return pydantic_cls, edge_by_property, issues
+
+    def _create_node(
+        self, identifier: str, properties: dict[str, list[str]], pydantic_cls: type[Model], view_id: dm.ViewId
+    ) -> dm.InstanceApply:
+        created = pydantic_cls.model_validate(properties)
+
+        return dm.NodeApply(
+            space=self.instance_space,
+            external_id=identifier,
+            # type=#RDF type
+            sources=[dm.NodeOrEdgeData(source=view_id, properties=dict(created.model_dump().items()))],
+        )
+
+    def _create_edges(
+        self,
+        identifier: str,
+        properties: dict[str, list[str]],
+        edge_by_properties: dict[str, dm.EdgeConnection],
+        tracker: Tracker,
+    ) -> Iterable[dm.EdgeApply | NeatIssue]:
+        for prop, values in properties.items():
+            if prop not in edge_by_properties:
+                continue
+            edge = edge_by_properties[prop]
+            if isinstance(edge, SingleEdgeConnection) and len(values) > 1:
+                error = loader_issues.InvalidInstanceError(
+                    type_="edge",
+                    identifier=identifier,
+                    reason=f"Multiple values for single edge {edge}. Expected only one.",
+                )
+                tracker.issue(error)
+                yield error
+            for target in values:
+                yield dm.EdgeApply(
+                    space=self.instance_space,
+                    external_id=f"{identifier}.{prop}.{target}",
+                    type=edge.type,
+                    start_node=dm.DirectRelationReference(self.instance_space, identifier),
+                    end_node=dm.DirectRelationReference(self.instance_space, target),
+                )
+
+    def _get_required_capabilities(self) -> list[Capability]:
+        return [
+            DataModelInstancesAcl(
+                actions=[
+                    DataModelInstancesAcl.Action.Write,
+                    DataModelInstancesAcl.Action.Write_Properties,
+                    DataModelInstancesAcl.Action.Read,
+                ],
+                scope=DataModelInstancesAcl.Scope.SpaceID([self.instance_space]),
+            )
+        ]
+
+    def _upload_to_cdf(
+        self,
+        client: CogniteClient,
+        items: list[dm.InstanceApply],
+        return_diffs: bool,
+        dry_run: bool,
+        read_issues: NeatIssueList,
+    ) -> UploadDiffsID:
+        result = UploadDiffsID(name=type(self).__name__, issues=read_issues)
+        try:
+            nodes = [item for item in items if isinstance(item, dm.NodeApply)]
+            edges = [item for item in items if isinstance(item, dm.EdgeApply)]
+            upserted = client.data_modeling.instances.apply(
+                nodes,
+                edges,
+                auto_create_end_nodes=True,
+                auto_create_start_nodes=True,
+                skip_on_version_conflict=True,
+            )
+        except CogniteAPIError as e:
+            result.error_messages.append(str(e))
+            result.failed.append([repr(instance.as_id()) for instance in items])  # type: ignore[arg-type, attr-defined]
+        else:
+            for instance in itertools.chain(upserted.nodes, upserted.edges):
+                if instance.was_modified and instance.created_time == instance.last_updated_time:
+                    result.created.append(repr(instance.as_id()))
+                elif instance.was_modified:
+                    result.changed.append(repr(instance.as_id()))
+                else:
+                    result.unchanged.append(repr(instance.as_id()))
+        return result if return_diffs else result.as_upload_result_ids()  # type: ignore[return-value]
+
+
+def _triples2dictionary(
+    triples: Iterable[tuple[str, str, str]],
+) -> dict[str, dict[str, list[str]]]:
+    """Converts list of triples to dictionary"""
+    values_by_property_by_identifier: dict[str, dict[str, list[str]]] = defaultdict(lambda: defaultdict(list))
+    for id_, property_, value in triples:
+        values_by_property_by_identifier[id_][property_].append(value)
+    return values_by_property_by_identifier

cognite/neat/graph/stores/_base.py

@@ -8,12 +8,13 @@ from typing import Literal, TypeAlias, cast
 
 import pandas as pd
 from prometheus_client import Gauge, Summary
-from rdflib import Graph, Namespace, URIRef
+from rdflib import RDF, Graph, Namespace, URIRef
 from rdflib.query import Result, ResultRow
 
 from cognite.neat.constants import DEFAULT_NAMESPACE, PREFIXES
 from cognite.neat.graph.models import Triple
 from cognite.neat.graph.stores._rdf_to_graph import rdf_file_to_graph
+from cognite.neat.utils import remove_namespace
 
 if sys.version_info >= (3, 11):
     pass
@@ -109,9 +110,9 @@ class NeatGraphStoreBase(ABC):
 
         if base_prefix:
             self.base_prefix = base_prefix
-
-        self.graph.bind(self.base_prefix, self.namespace)
-        logging.info("Adding prefix %s with namespace %s", self.base_prefix, self.namespace)
+        if self.base_prefix:
+            self.graph.bind(self.base_prefix, self.namespace)
+            logging.info("Adding prefix %s with namespace %s", self.base_prefix, self.namespace)
         logging.info("Graph initialized")
 
     def reinitialize_graph(self):
@@ -361,3 +362,17 @@ class _Queries:
         logging.info(query)
         # Select queries gives an iterable of result rows
         return cast(list[ResultRow], list(self.store.query(query)))
+
+    def triples_of_type_instances(self, rdf_type: str) -> list[tuple[str, str, str]]:
+        """Get all triples of a given type.
+
+        This method assumes the graph has been transformed into the default namespace.
+        """
+        query = (
+            f"SELECT ?instance ?prop ?value "
+            f"WHERE {{ ?instance a <{self.store.namespace[rdf_type]}> . ?instance ?prop ?value . }} order by ?instance"
+        )
+        result = self.store.query(query)
+
+        # We cannot include the RDF.type in case there is a neat:type property
+        return [remove_namespace(*triple) for triple in result if triple[1] != RDF.type]  # type: ignore[misc, index]

cognite/neat/issues.py

@@ -0,0 +1,150 @@
+import sys
+import warnings
+from abc import ABC, abstractmethod
+from collections import UserList
+from collections.abc import Sequence
+from dataclasses import dataclass
+from functools import total_ordering
+from typing import Any, ClassVar, TypeVar
+from warnings import WarningMessage
+
+import pandas as pd
+
+if sys.version_info < (3, 11):
+    from exceptiongroup import ExceptionGroup
+    from typing_extensions import Self
+else:
+    from typing import Self
+
+
+@total_ordering
+@dataclass(frozen=True)
+class NeatIssue(ABC):
+    description: ClassVar[str]
+    fix: ClassVar[str]
+
+    def message(self) -> str:
+        """Return a human-readable message for the issue.
+
+        This is the default implementation, which returns the description.
+        It is recommended to override this method in subclasses with a more
+        specific message.
+        """
+        return self.description
+
+    @abstractmethod
+    def dump(self) -> dict[str, Any]:
+        """Return a dictionary representation of the issue."""
+        raise NotImplementedError()
+
+    def __lt__(self, other: "NeatIssue") -> bool:
+        if not isinstance(other, NeatIssue):
+            return NotImplemented
+        return (type(self).__name__, self.message()) < (type(other).__name__, other.message())
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, NeatIssue):
+            return NotImplemented
+        return (type(self).__name__, self.message()) == (type(other).__name__, other.message())
+
+
+@dataclass(frozen=True)
+class NeatError(NeatIssue, ABC):
+    def dump(self) -> dict[str, Any]:
+        return {"errorType": type(self).__name__}
+
+    def as_exception(self) -> ValueError:
+        return ValueError(self.message())
+
+
+@dataclass(frozen=True)
+class NeatWarning(NeatIssue, ABC, UserWarning):
+    def dump(self) -> dict[str, Any]:
+        return {"warningType": type(self).__name__}
+
+    @classmethod
+    def from_warning(cls, warning: WarningMessage) -> "NeatWarning":
+        return DefaultWarning.from_warning_message(warning)
+
+
+@dataclass(frozen=True)
+class DefaultWarning(NeatWarning):
+    description = "A warning was raised during validation."
+    fix = "No fix is available."
+
+    warning: str | Warning
+    category: type[Warning]
+    source: str | None = None
+
+    def dump(self) -> dict[str, Any]:
+        output = super().dump()
+        output["msg"] = str(self.warning)
+        output["category"] = self.category.__name__
+        output["source"] = self.source
+        return output
+
+    @classmethod
+    def from_warning_message(cls, warning: WarningMessage) -> NeatWarning:
+        if isinstance(warning.message, NeatWarning):
+            return warning.message
+
+        return cls(
+            warning=warning.message,
+            category=warning.category,
+            source=warning.source,
+        )
+
+    def message(self) -> str:
+        return str(self.warning)
+
+
+T_NeatIssue = TypeVar("T_NeatIssue", bound=NeatIssue)
+
+
+class NeatIssueList(UserList[T_NeatIssue], ABC):
+    def __init__(self, issues: Sequence[T_NeatIssue] | None = None, title: str | None = None):
+        super().__init__(issues or [])
+        self.title = title
+
+    @property
+    def errors(self) -> Self:
+        return type(self)([issue for issue in self if isinstance(issue, NeatError)])  # type: ignore[misc]
+
+    @property
+    def has_errors(self) -> bool:
+        return any(isinstance(issue, NeatError) for issue in self)
+
+    @property
+    def warnings(self) -> Self:
+        return type(self)([issue for issue in self if isinstance(issue, NeatWarning)])  # type: ignore[misc]
+
+    def as_errors(self) -> ExceptionGroup:
+        return ExceptionGroup(
+            "Operation failed",
+            [ValueError(issue.message()) for issue in self if isinstance(issue, NeatError)],
+        )
+
+    def trigger_warnings(self) -> None:
+        for warning in [issue for issue in self if isinstance(issue, NeatWarning)]:
+            warnings.warn(warning, stacklevel=2)
+
+    def to_pandas(self) -> pd.DataFrame:
+        return pd.DataFrame([issue.dump() for issue in self])
+
+    def _repr_html_(self) -> str | None:
+        return self.to_pandas()._repr_html_()  # type: ignore[operator]
+
+    def as_exception(self) -> "MultiValueError":
+        return MultiValueError(self.errors)
+
+
+class MultiValueError(ValueError):
+    """This is a container for multiple errors.
+
+    It is used in the pydantic field_validator/model_validator to collect multiple errors, which
+    can then be caught in a try-except block and returned as an IssueList.
+
+    """
+
+    def __init__(self, errors: Sequence[T_NeatIssue]):
+        self.errors = list(errors)

cognite/neat/rules/exporters/_base.py

@@ -7,8 +7,7 @@ from cognite.client import CogniteClient
 
 from cognite.neat.rules._shared import Rules
 from cognite.neat.rules.models import DMSRules, InformationRules, RoleTypes
-
-from ._models import UploadResult
+from cognite.neat.utils.upload import UploadDiffsCount
 
 T_Export = TypeVar("T_Export")
 
@@ -38,5 +37,5 @@ class BaseExporter(ABC, Generic[T_Export]):
 
 class CDFExporter(BaseExporter[T_Export]):
     @abstractmethod
-    def export_to_cdf(self, rules: Rules, client: CogniteClient, dry_run: bool = False) -> Iterable[UploadResult]:
+    def export_to_cdf(self, rules: Rules, client: CogniteClient, dry_run: bool = False) -> Iterable[UploadDiffsCount]:
         raise NotImplementedError

cognite/neat/rules/exporters/_rules2dms.py

@@ -31,9 +31,9 @@ from cognite.neat.utils.cdf_loaders import (
     TransformationLoader,
     ViewLoader,
 )
+from cognite.neat.utils.upload import UploadDiffsCount
 
 from ._base import CDFExporter
-from ._models import UploadResult
 
 Component: TypeAlias = Literal["all", "spaces", "data_models", "views", "containers", "node_types"]
 
@@ -123,7 +123,7 @@ class DMSExporter(CDFExporter[DMSSchema]):
             raise ValueError(f"{type(rules).__name__} cannot be exported to DMS")
         return dms_rules.as_schema(include_pipeline=self.export_pipeline, instance_space=self.instance_space)
 
-    def delete_from_cdf(self, rules: Rules, client: CogniteClient, dry_run: bool = False) -> Iterable[UploadResult]:
+    def delete_from_cdf(self, rules: Rules, client: CogniteClient, dry_run: bool = False) -> Iterable[UploadDiffsCount]:
         schema, to_export = self._prepare_schema_and_exporters(rules, client)
 
         # we need to reverse order in which we are picking up the items to delete
@@ -161,7 +161,7 @@ class DMSExporter(CDFExporter[DMSSchema]):
                         deleted -= failed_deleted
                         error_messages.append(f"Failed delete: {e.message}")
 
-            yield UploadResult(
+            yield UploadDiffsCount(
                 name=loader.resource_name,
                 deleted=deleted,
                 skipped=0,
@@ -169,7 +169,7 @@ class DMSExporter(CDFExporter[DMSSchema]):
                 error_messages=error_messages,
             )
 
-    def export_to_cdf(self, rules: Rules, client: CogniteClient, dry_run: bool = False) -> Iterable[UploadResult]:
+    def export_to_cdf(self, rules: Rules, client: CogniteClient, dry_run: bool = False) -> Iterable[UploadDiffsCount]:
         schema, to_export = self._prepare_schema_and_exporters(rules, client)
 
         # The conversion from DMS to GraphQL does not seem to be triggered even if the views
@@ -250,7 +250,7 @@ class DMSExporter(CDFExporter[DMSSchema]):
                         changed -= failed_changed
                         error_messages.append(e.message)
 
-            yield UploadResult(
+            yield UploadDiffsCount(
                 name=loader.resource_name,
                 created=created,
                 changed=changed,

cognite/neat/rules/importers/_base.py

@@ -118,4 +118,4 @@ def _handle_issues(
             future_result._result = "success"
         finally:
             if warning_logger:
-                issues.extend([warning_cls.from_warning(warning) for warning in warning_logger])
+                issues.extend([warning_cls.from_warning(warning) for warning in warning_logger])  # type: ignore[misc]

cognite/neat/rules/issues/__init__.py

@@ -1,9 +1,9 @@
+from cognite.neat.issues import MultiValueError
+
 from . import dms, fileread, formatters, importing, spreadsheet, spreadsheet_file
 from .base import (
     DefaultPydanticError,
-    DefaultWarning,
     IssueList,
-    MultiValueError,
     NeatValidationError,
     ValidationIssue,
     ValidationWarning,
@@ -11,7 +11,6 @@ from .base import (
 
 __all__ = [
     "DefaultPydanticError",
-    "DefaultWarning",
     "MultiValueError",
     "IssueList",
     "NeatValidationError",