cognite-neat 0.75.9__py3-none-any.whl → 0.76.1__py3-none-any.whl

cognite/neat/rules/importers/_dms2rules.py

@@ -1,3 +1,6 @@
+from collections import Counter
+from collections.abc import Sequence
+from datetime import datetime
 from pathlib import Path
 from typing import Literal, cast, overload
 
@@ -5,21 +8,29 @@ from cognite.client import CogniteClient
 from cognite.client import data_modeling as dm
 from cognite.client.data_classes.data_modeling import DataModelIdentifier
 from cognite.client.data_classes.data_modeling.containers import BTreeIndex, InvertedIndex
+from cognite.client.data_classes.data_modeling.views import (
+    MultiEdgeConnectionApply,
+    MultiReverseDirectRelationApply,
+    SingleEdgeConnectionApply,
+    SingleReverseDirectRelationApply,
+    ViewPropertyApply,
+)
 from cognite.client.utils import ms_to_datetime
 
 from cognite.neat.rules import issues
-from cognite.neat.rules.importers._base import BaseImporter, Rules
-from cognite.neat.rules.issues import IssueList
+from cognite.neat.rules.importers._base import BaseImporter, Rules, _handle_issues
+from cognite.neat.rules.issues import IssueList, ValidationIssue
 from cognite.neat.rules.models.data_types import DataType
 from cognite.neat.rules.models.entities import (
     ClassEntity,
     ContainerEntity,
+    DataModelEntity,
     DMSUnknownEntity,
     ViewEntity,
     ViewPropertyEntity,
 )
 from cognite.neat.rules.models.rules import DMSRules, DMSSchema, RoleTypes
-from cognite.neat.rules.models.rules._base import ExtensionCategory, SchemaCompleteness
+from cognite.neat.rules.models.rules._base import DataModelType, ExtensionCategory, SchemaCompleteness
 from cognite.neat.rules.models.rules._dms_architect_rules import (
     DMSContainer,
     DMSMetadata,
@@ -30,9 +41,16 @@ from cognite.neat.rules.models.rules._dms_architect_rules import (
 
 
 class DMSImporter(BaseImporter):
-    def __init__(self, schema: DMSSchema, metadata: DMSMetadata | None = None):
+    def __init__(
+        self,
+        schema: DMSSchema,
+        read_issues: Sequence[ValidationIssue] | None = None,
+        metadata: DMSMetadata | None = None,
+    ):
         self.schema = schema
         self.metadata = metadata
+        self.issue_list = IssueList(read_issues)
+        self._container_by_id = {container.as_id(): container for container in schema.containers}
 
     @classmethod
     def from_data_model_id(cls, client: CogniteClient, data_model_id: DataModelIdentifier) -> "DMSImporter":
@@ -47,38 +65,59 @@ class DMSImporter(BaseImporter):
         """
         data_models = client.data_modeling.data_models.retrieve(data_model_id, inline_views=True)
         if len(data_models) == 0:
-            raise ValueError(f"Data model {data_model_id} not found")
+            return cls(DMSSchema(), [issues.importing.NoDataModelError(f"Data model {data_model_id} not found")])
         data_model = data_models.latest_version()
-        schema = DMSSchema.from_data_model(client, data_model)
-        description, creator = DMSMetadata._get_description_and_creator(data_model.description)
+
+        try:
+            schema = DMSSchema.from_data_model(client, data_model)
+        except Exception as e:
+            return cls(DMSSchema(), [issues.importing.APIError(str(e))])
 
         created = ms_to_datetime(data_model.created_time)
         updated = ms_to_datetime(data_model.last_updated_time)
 
-        metadata = DMSMetadata(
+        metadata = cls._create_metadata_from_model(data_model, created, updated)
+
+        return cls(schema, [], metadata)
+
+    @classmethod
+    def _create_metadata_from_model(
+        cls,
+        model: dm.DataModel[dm.View] | dm.DataModelApply,
+        created: datetime | None = None,
+        updated: datetime | None = None,
+    ) -> DMSMetadata:
+        description, creator = DMSMetadata._get_description_and_creator(model.description)
+        now = datetime.now().replace(microsecond=0)
+        return DMSMetadata(
             schema_=SchemaCompleteness.complete,
             extension=ExtensionCategory.addition,
-            space=data_model.space,
-            external_id=data_model.external_id,
-            name=data_model.name or data_model.external_id,
-            version=data_model.version or "0.1.0",
-            updated=updated,
-            created=created,
+            space=model.space,
+            external_id=model.external_id,
+            name=model.name or model.external_id,
+            version=model.version or "0.1.0",
+            updated=updated or now,
+            created=created or now,
             creator=creator,
             description=description,
-            default_view_version=data_model.version or "0.1.0",
         )
-        return cls(schema, metadata)
 
     @classmethod
     def from_directory(cls, directory: str | Path) -> "DMSImporter":
-        return cls(DMSSchema.from_directory(directory))
+        issue_list = IssueList()
+        with _handle_issues(issue_list) as _:
+            schema = DMSSchema.from_directory(directory)
+        # If there were errors during the import, the to_rules
+        return cls(schema, issue_list)
 
     @classmethod
     def from_zip_file(cls, zip_file: str | Path) -> "DMSImporter":
         if Path(zip_file).suffix != ".zip":
-            raise ValueError("File extension is not .zip")
-        return cls(DMSSchema.from_zip(zip_file))
+            return cls(DMSSchema(), [issues.fileread.InvalidFileFormatError(Path(zip_file), [".zip"])])
+        issue_list = IssueList()
+        with _handle_issues(issue_list) as _:
+            schema = DMSSchema.from_zip(zip_file)
+        return cls(schema, issue_list)
 
     @overload
     def to_rules(self, errors: Literal["raise"], role: RoleTypes | None = None) -> Rules: ...
@@ -91,126 +130,281 @@ class DMSImporter(BaseImporter):
     def to_rules(
         self, errors: Literal["raise", "continue"] = "continue", role: RoleTypes | None = None
     ) -> tuple[Rules | None, IssueList] | Rules:
-        if role is RoleTypes.domain_expert:
-            raise ValueError(f"Role {role} is not supported for DMSImporter")
-        issue_list = IssueList()
-        data_model = self.schema.data_models[0]
+        if self.issue_list.has_errors:
+            # In case there were errors during the import, the to_rules method will return None
+            return self._return_or_raise(self.issue_list, errors)
 
-        container_by_id = {container.as_id(): container for container in self.schema.containers}
+        if len(self.schema.data_models) == 0:
+            self.issue_list.append(issues.importing.NoDataModelError("No data model found."))
+            return self._return_or_raise(self.issue_list, errors)
+
+        if len(self.schema.data_models) > 2:
+            # Creating a DataModelEntity to convert the data model id to a string.
+            self.issue_list.append(
+                issues.importing.MultipleDataModelsWarning(
+                    [str(DataModelEntity.from_id(model.as_id())) for model in self.schema.data_models]
+                )
+            )
+
+        data_model = self.schema.data_models[0]
 
         properties = SheetList[DMSProperty]()
+        ref_properties = SheetList[DMSProperty]()
         for view in self.schema.views:
-            class_entity = ClassEntity(prefix=view.space, suffix=view.external_id, version=view.version)
+            view_id = view.as_id()
+            view_entity = ViewEntity.from_id(view_id)
+            class_entity = view_entity.as_class()
             for prop_id, prop in (view.properties or {}).items():
-                if isinstance(prop, dm.MappedPropertyApply):
-                    if prop.container not in container_by_id:
-                        raise ValueError(f"Container {prop.container} not found")
-                    container = container_by_id[prop.container]
-                    if prop.container_property_identifier not in container.properties:
-                        raise ValueError(
-                            f"Property {prop.container_property_identifier} not found "
-                            f"in container {container.external_id}"
-                        )
-                    container_prop = container.properties[prop.container_property_identifier]
-
-                    index: list[str] = []
-                    for index_name, index_obj in (container.indexes or {}).items():
-                        if isinstance(index_obj, BTreeIndex | InvertedIndex) and prop_id in index_obj.properties:
-                            index.append(index_name)
-                    unique_constraints: list[str] = []
-                    for constraint_name, constraint_obj in (container.constraints or {}).items():
-                        if isinstance(constraint_obj, dm.RequiresConstraint):
-                            # This is handled in the .from_container method of DMSContainer
-                            continue
-                        elif (
-                            isinstance(constraint_obj, dm.UniquenessConstraint) and prop_id in constraint_obj.properties
-                        ):
-                            unique_constraints.append(constraint_name)
-                        elif isinstance(constraint_obj, dm.UniquenessConstraint):
-                            # This does not apply to this property
-                            continue
-                        else:
-                            raise NotImplementedError(f"Constraint type {type(constraint_obj)} not implemented")
-
-                    if isinstance(container_prop.type, dm.DirectRelation):
-                        direct_value_type: str | ViewEntity | DataType | DMSUnknownEntity
-                        if prop.source is None:
-                            issue_list.append(
-                                issues.importing.UnknownValueTypeWarning(class_entity.versioned_id, prop_id)
-                            )
-                            direct_value_type = DMSUnknownEntity()
-                        else:
-                            direct_value_type = ViewEntity.from_id(prop.source)
-
-                        dms_property = DMSProperty(
-                            class_=class_entity,
-                            property_=prop_id,
-                            description=prop.description,
-                            name=prop.name,
-                            value_type=direct_value_type,
-                            relation="direct",
-                            nullable=container_prop.nullable,
-                            default=container_prop.default_value,
-                            is_list=False,
-                            container=ContainerEntity.from_id(container.as_id()),
-                            container_property=prop.container_property_identifier,
-                            view=ViewEntity.from_id(view.as_id()),
-                            view_property=prop_id,
-                            index=index or None,
-                            constraint=unique_constraints or None,
-                        )
+                dms_property = self._create_dms_property(prop_id, prop, view_entity, class_entity)
+                if dms_property is not None:
+                    if view_id in self.schema.frozen_ids:
+                        ref_properties.append(dms_property)
                     else:
-                        dms_property = DMSProperty(
-                            class_=ClassEntity(prefix=view.space, suffix=view.external_id, version=view.version),
-                            property_=prop_id,
-                            description=prop.description,
-                            name=prop.name,
-                            value_type=cast(ViewPropertyEntity | DataType, container_prop.type._type),
-                            nullable=container_prop.nullable,
-                            is_list=container_prop.type.is_list,
-                            default=container_prop.default_value,
-                            container=ContainerEntity.from_id(container.as_id()),
-                            container_property=prop.container_property_identifier,
-                            view=ViewEntity.from_id(view.as_id()),
-                            view_property=prop_id,
-                            index=index or None,
-                            constraint=unique_constraints or None,
-                        )
-                elif isinstance(prop, dm.MultiEdgeConnectionApply):
-                    view_entity = ViewEntity.from_id(prop.source)
-                    dms_property = DMSProperty(
-                        class_=ClassEntity(prefix=view.space, suffix=view.external_id, version=view.version),
-                        property_=prop_id,
-                        relation="multiedge",
-                        description=prop.description,
-                        name=prop.name,
-                        value_type=view_entity,
-                        view=ViewEntity.from_id(view.as_id()),
-                        view_property=prop_id,
-                    )
-                else:
-                    raise NotImplementedError(f"Property type {type(prop)} not implemented")
-
-                properties.append(dms_property)
+                        properties.append(dms_property)
 
         data_model_view_ids: set[dm.ViewId] = {
             view.as_id() if isinstance(view, dm.View | dm.ViewApply) else view for view in data_model.views or []
         }
 
-        dms_rules = DMSRules(
-            metadata=self.metadata or DMSMetadata.from_data_model(data_model),
+        metadata = self.metadata or DMSMetadata.from_data_model(data_model)
+        metadata.data_model_type = self._infer_data_model_type(metadata.space)
+        if ref_properties:
+            metadata.schema_ = SchemaCompleteness.extended
+
+        with _handle_issues(
+            self.issue_list,
+        ) as future:
+            user_rules = DMSRules(
+                metadata=metadata,
+                properties=properties,
+                containers=SheetList[DMSContainer](
+                    data=[
+                        DMSContainer.from_container(container)
+                        for container in self.schema.containers
+                        if container.as_id() not in self.schema.frozen_ids
+                    ]
+                ),
+                views=SheetList[DMSView](
+                    data=[
+                        DMSView.from_view(view, in_model=view.as_id() in data_model_view_ids)
+                        for view in self.schema.views
+                        if view.as_id() not in self.schema.frozen_ids
+                    ]
+                ),
+                reference=self._create_reference_rules(ref_properties),
+            )
+
+        if future.result == "failure" or self.issue_list.has_errors:
+            return self._return_or_raise(self.issue_list, errors)
+
+        return self._to_output(user_rules, self.issue_list, errors, role)
+
+    def _create_reference_rules(self, properties: SheetList[DMSProperty]) -> DMSRules | None:
+        if not properties:
+            return None
+
+        if len(self.schema.data_models) == 2:
+            data_model = self.schema.data_models[1]
+            data_model_view_ids: set[dm.ViewId] = {
+                view.as_id() if isinstance(view, dm.View | dm.ViewApply) else view for view in data_model.views or []
+            }
+            metadata = self._create_metadata_from_model(data_model)
+        else:
+            data_model_view_ids = set()
+            now = datetime.now().replace(microsecond=0)
+            space = Counter(prop.view.space for prop in properties).most_common(1)[0][0]
+            metadata = DMSMetadata(
+                schema_=SchemaCompleteness.complete,
+                extension=ExtensionCategory.addition,
+                space=space,
+                external_id="Unknown",
+                version="0.1.0",
+                creator=["Unknown"],
+                created=now,
+                updated=now,
+            )
+
+        metadata.data_model_type = DataModelType.enterprise
+        return DMSRules(
+            metadata=metadata,
             properties=properties,
+            views=SheetList[DMSView](
+                data=[
+                    DMSView.from_view(view, in_model=not data_model_view_ids or (view.as_id() in data_model_view_ids))
+                    for view in self.schema.views
+                    if view.as_id() in self.schema.frozen_ids
+                ]
+            ),
             containers=SheetList[DMSContainer](
-                data=[DMSContainer.from_container(container) for container in self.schema.containers]
+                data=[
+                    DMSContainer.from_container(container)
+                    for container in self.schema.containers
+                    if container.as_id() in self.schema.frozen_ids
+                ]
             ),
-            views=SheetList[DMSView](data=[DMSView.from_view(view, data_model_view_ids) for view in self.schema.views]),
+            reference=None,
+        )
+
+    def _infer_data_model_type(self, space: str) -> DataModelType:
+        if self.schema.referenced_spaces() - {space}:
+            # If the data model has containers, views, node types in another space
+            # we assume it is a solution model.
+            return DataModelType.solution
+        else:
+            # All containers, views, node types are in the same space as the data model
+            return DataModelType.enterprise
+
+    def _create_dms_property(
+        self, prop_id: str, prop: ViewPropertyApply, view_entity: ViewEntity, class_entity: ClassEntity
+    ) -> DMSProperty | None:
+        if isinstance(prop, dm.MappedPropertyApply) and prop.container not in self._container_by_id:
+            self.issue_list.append(
+                issues.importing.MissingContainerWarning(
+                    view_id=str(view_entity),
+                    property_=prop_id,
+                    container_id=str(ContainerEntity.from_id(prop.container)),
+                )
+            )
+            return None
+        if (
+            isinstance(prop, dm.MappedPropertyApply)
+            and prop.container_property_identifier not in self._container_by_id[prop.container].properties
+        ):
+            self.issue_list.append(
+                issues.importing.MissingContainerPropertyWarning(
+                    view_id=str(view_entity),
+                    property_=prop_id,
+                    container_id=str(ContainerEntity.from_id(prop.container)),
+                )
+            )
+            return None
+        if not isinstance(
+            prop,
+            dm.MappedPropertyApply
+            | SingleEdgeConnectionApply
+            | MultiEdgeConnectionApply
+            | SingleReverseDirectRelationApply
+            | MultiReverseDirectRelationApply,
+        ):
+            self.issue_list.append(
+                issues.importing.UnknownPropertyTypeWarning(view_entity.versioned_id, prop_id, type(prop).__name__)
+            )
+            return None
+
+        value_type = self._get_value_type(prop, view_entity, prop_id)
+        if value_type is None:
+            return None
+
+        return DMSProperty(
+            class_=class_entity,
+            property_=prop_id,
+            description=prop.description,
+            name=prop.name,
+            connection=self._get_relation_type(prop),
+            value_type=value_type,
+            is_list=self._get_is_list(prop),
+            nullable=self._get_nullable(prop),
+            default=self._get_default(prop),
+            container=ContainerEntity.from_id(prop.container) if isinstance(prop, dm.MappedPropertyApply) else None,
+            container_property=prop.container_property_identifier if isinstance(prop, dm.MappedPropertyApply) else None,
+            view=view_entity,
+            view_property=prop_id,
+            index=self._get_index(prop, prop_id),
+            constraint=self._get_constraint(prop, prop_id),
         )
-        output_rules: Rules
-        if role is RoleTypes.information_architect:
-            output_rules = dms_rules.as_information_architect_rules()
+
+    def _container_prop_unsafe(self, prop: dm.MappedPropertyApply) -> dm.ContainerProperty:
+        """This method assumes you have already checked that the container with property exists."""
+        return self._container_by_id[prop.container].properties[prop.container_property_identifier]
+
+    def _get_relation_type(self, prop: ViewPropertyApply) -> Literal["edge", "reverse", "direct"] | None:
+        if isinstance(prop, SingleEdgeConnectionApply | MultiEdgeConnectionApply) and prop.direction == "outwards":
+            return "edge"
+        elif isinstance(prop, SingleEdgeConnectionApply | MultiEdgeConnectionApply) and prop.direction == "inwards":
+            return "reverse"
+        elif isinstance(prop, SingleReverseDirectRelationApply | MultiReverseDirectRelationApply):
+            return "reverse"
+        elif isinstance(prop, dm.MappedPropertyApply) and isinstance(
+            self._container_prop_unsafe(prop).type, dm.DirectRelation
+        ):
+            return "direct"
         else:
-            output_rules = dms_rules
-        if errors == "raise":
-            return output_rules
+            return None
+
+    def _get_value_type(
+        self, prop: ViewPropertyApply, view_entity: ViewEntity, prop_id
+    ) -> DataType | ViewEntity | ViewPropertyEntity | DMSUnknownEntity | None:
+        if isinstance(prop, SingleEdgeConnectionApply | MultiEdgeConnectionApply) and prop.direction == "outwards":
+            return ViewEntity.from_id(prop.source)
+        elif isinstance(prop, SingleReverseDirectRelationApply | MultiReverseDirectRelationApply):
+            return ViewPropertyEntity.from_id(prop.through)
+        elif isinstance(prop, SingleEdgeConnectionApply | MultiEdgeConnectionApply) and prop.direction == "inwards":
+            return ViewEntity.from_id(prop.source)
+        elif isinstance(prop, dm.MappedPropertyApply):
+            container_prop = self._container_prop_unsafe(cast(dm.MappedPropertyApply, prop))
+            if isinstance(container_prop.type, dm.DirectRelation):
+                if prop.source is None:
+                    # The warning is issued when the DMS Rules are created.
+                    return DMSUnknownEntity()
+                else:
+                    return ViewEntity.from_id(prop.source)
+            else:
+                return DataType.load(container_prop.type._type)
+        else:
+            self.issue_list.append(issues.importing.FailedToInferValueTypeWarning(str(view_entity), prop_id))
+            return None
+
+    def _get_nullable(self, prop: ViewPropertyApply) -> bool | None:
+        if isinstance(prop, dm.MappedPropertyApply):
+            return self._container_prop_unsafe(prop).nullable
+        else:
+            return None
+
+    def _get_is_list(self, prop: ViewPropertyApply) -> bool | None:
+        if isinstance(prop, dm.MappedPropertyApply):
+            return self._container_prop_unsafe(prop).type.is_list
+        elif isinstance(prop, MultiEdgeConnectionApply | MultiReverseDirectRelationApply):
+            return True
+        elif isinstance(prop, SingleEdgeConnectionApply | SingleReverseDirectRelationApply):
+            return False
         else:
-            return output_rules, issue_list
+            return None
+
+    def _get_default(self, prop: ViewPropertyApply) -> str | None:
+        if isinstance(prop, dm.MappedPropertyApply):
+            default = self._container_prop_unsafe(prop).default_value
+            if default is not None:
+                return str(default)
+        return None
+
+    def _get_index(self, prop: ViewPropertyApply, prop_id) -> list[str] | None:
+        if not isinstance(prop, dm.MappedPropertyApply):
+            return None
+        container = self._container_by_id[prop.container]
+        index: list[str] = []
+        for index_name, index_obj in (container.indexes or {}).items():
+            if isinstance(index_obj, BTreeIndex | InvertedIndex) and prop_id in index_obj.properties:
+                index.append(index_name)
+        return index or None
+
+    def _get_constraint(self, prop: ViewPropertyApply, prop_id: str) -> list[str] | None:
+        if not isinstance(prop, dm.MappedPropertyApply):
+            return None
+        container = self._container_by_id[prop.container]
+        unique_constraints: list[str] = []
+        for constraint_name, constraint_obj in (container.constraints or {}).items():
+            if isinstance(constraint_obj, dm.RequiresConstraint):
+                # This is handled in the .from_container method of DMSContainer
+                continue
+            elif isinstance(constraint_obj, dm.UniquenessConstraint) and prop_id in constraint_obj.properties:
+                unique_constraints.append(constraint_name)
+            elif isinstance(constraint_obj, dm.UniquenessConstraint):
+                # This does not apply to this property
+                continue
+            else:
+                self.issue_list.append(
+                    issues.importing.UnknownContainerConstraintWarning(
+                        str(ContainerEntity.from_id(prop.container)), prop_id, type(constraint_obj).__name__
+                    )
+                )
+        return unique_constraints or None

cognite/neat/rules/importers/_spreadsheet2rules.py

@@ -22,7 +22,15 @@ from cognite.neat.utils.spreadsheet import SpreadsheetRead, read_individual_shee
 from ._base import BaseImporter, Rules, _handle_issues
 
 SOURCE_SHEET__TARGET_FIELD__HEADERS = [
-    ("Properties", "Properties", "Class"),
+    (
+        "Properties",
+        "Properties",
+        {
+            RoleTypes.domain_expert: "Property",
+            RoleTypes.information_architect: "Property",
+            RoleTypes.dms_architect: "View Property",
+        },
+    ),
     ("Classes", "Classes", "Class"),
     ("Containers", "Containers", "Container"),
     ("Views", "Views", "View"),
@@ -131,11 +139,15 @@ class SpreadsheetReader:
             )
             return None, read_info_by_sheet
 
-        for source_sheet_name, target_sheet_name, headers in SOURCE_SHEET__TARGET_FIELD__HEADERS:
+        for source_sheet_name, target_sheet_name, headers_input in SOURCE_SHEET__TARGET_FIELD__HEADERS:
             source_sheet_name = self.to_reference_sheet(source_sheet_name) if self._is_reference else source_sheet_name
 
             if source_sheet_name not in excel_file.sheet_names:
                 continue
+            if isinstance(headers_input, dict):
+                headers = headers_input[metadata.role]
+            else:
+                headers = headers_input
 
             try:
                 sheets[target_sheet_name], read_info_by_sheet[source_sheet_name] = read_individual_sheet(
@@ -229,12 +241,6 @@ class ExcelImporter(BaseImporter):
             role=role,
         )
 
-    @classmethod
-    def _return_or_raise(cls, issue_list: IssueList, errors: Literal["raise", "continue"]) -> tuple[None, IssueList]:
-        if errors == "raise":
-            raise issue_list.as_errors()
-        return None, issue_list
-
 
 class GoogleSheetImporter(BaseImporter):
     def __init__(self, sheet_id: str, skiprows: int = 1):

cognite/neat/rules/issues/base.py

@@ -78,6 +78,9 @@ class NeatValidationError(ValidationIssue, ABC):
                 all_errors.append(DefaultPydanticError.from_pydantic_error(error))
         return all_errors
 
+    def as_exception(self) -> Exception:
+        return ValueError(self.message())
+
 
 @dataclass(frozen=True)
 class DefaultPydanticError(NeatValidationError):