cognite-neat 0.76.3__tar.gz → 0.77.1__tar.gz

{cognite_neat-0.76.3 → cognite_neat-0.77.1}/PKG-INFO

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

cognite_neat-0.77.1/cognite/neat/_version.py

@@ -0,0 +1 @@
+__version__ = "0.77.1"

{cognite_neat-0.76.3 → cognite_neat-0.77.1}/cognite/neat/rules/exporters/_rules2dms.py

@@ -5,7 +5,14 @@ from typing import Literal, TypeAlias, cast
 
 from cognite.client import CogniteClient
 from cognite.client.data_classes._base import CogniteResource, CogniteResourceList
-from cognite.client.data_classes.data_modeling import DataModelApply, DataModelId
+from cognite.client.data_classes.data_modeling import (
+    ContainerApplyList,
+    DataModelApply,
+    DataModelApplyList,
+    DataModelId,
+    SpaceApplyList,
+    ViewApplyList,
+)
 from cognite.client.exceptions import CogniteAPIError
 
 from cognite.neat.rules import issues
@@ -114,9 +121,7 @@ class DMSExporter(CDFExporter[DMSSchema]):
             dms_rules = rules.as_dms_architect_rules()
         else:
             raise ValueError(f"{type(rules).__name__} cannot be exported to DMS")
-        return dms_rules.as_schema(
-            include_ref=True, include_pipeline=self.export_pipeline, instance_space=self.instance_space
-        )
+        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]:
         schema, to_export = self._prepare_schema_and_exporters(rules, client)
@@ -266,13 +271,13 @@ class DMSExporter(CDFExporter[DMSSchema]):
         schema = self.export(rules)
         to_export: list[tuple[CogniteResourceList, ResourceLoader]] = []
         if self.export_components.intersection({"all", "spaces"}):
-            to_export.append((schema.spaces, SpaceLoader(client)))
+            to_export.append((SpaceApplyList(schema.spaces.values()), SpaceLoader(client)))
         if self.export_components.intersection({"all", "containers"}):
-            to_export.append((schema.containers, ContainerLoader(client)))
+            to_export.append((ContainerApplyList(schema.containers.values()), ContainerLoader(client)))
         if self.export_components.intersection({"all", "views"}):
-            to_export.append((schema.views, ViewLoader(client, self.existing_handling)))
+            to_export.append((ViewApplyList(schema.views.values()), ViewLoader(client, self.existing_handling)))
         if self.export_components.intersection({"all", "data_models"}):
-            to_export.append((schema.data_models, DataModelLoader(client)))
+            to_export.append((DataModelApplyList([schema.data_model]), DataModelLoader(client)))
         if isinstance(schema, PipelineSchema):
             to_export.append((schema.databases, RawDatabaseLoader(client)))
             to_export.append((schema.raw_tables, RawTableLoader(client)))

{cognite_neat-0.76.3 → cognite_neat-0.77.1}/cognite/neat/rules/importers/_dms2rules.py

@@ -6,7 +6,7 @@ from typing import Any, Literal, cast, overload
 
 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 import DataModelId, DataModelIdentifier
 from cognite.client.data_classes.data_modeling.containers import BTreeIndex, InvertedIndex
 from cognite.client.data_classes.data_modeling.views import (
     MultiEdgeConnectionApply,
@@ -39,7 +39,6 @@ from cognite.neat.rules.models.dms import (
 from cognite.neat.rules.models.entities import (
     ClassEntity,
     ContainerEntity,
-    DataModelEntity,
     DMSUnknownEntity,
     ViewEntity,
     ViewPropertyEntity,
@@ -52,56 +51,92 @@ class DMSImporter(BaseImporter):
         schema: DMSSchema,
         read_issues: Sequence[ValidationIssue] | None = None,
         metadata: DMSMetadata | None = None,
+        ref_metadata: DMSMetadata | None = None,
     ):
         # Calling this root schema to distinguish it from
         # * User Schema
         # * Reference Schema
         self.root_schema = schema
         self.metadata = metadata
+        self.ref_metadata = ref_metadata
         self.issue_list = IssueList(read_issues)
-        self._all_containers_by_id = {container.as_id(): container for container in schema.containers}
+        self._all_containers_by_id = schema.containers.copy()
         if self.root_schema.reference:
-            self._all_containers_by_id.update(
-                {container.as_id(): container for container in self.root_schema.reference.containers}
-            )
+            self._all_containers_by_id.update(self.root_schema.reference.containers)
 
     @classmethod
-    def from_data_model_id(cls, client: CogniteClient, data_model_id: DataModelIdentifier) -> "DMSImporter":
+    def from_data_model_id(
+        cls,
+        client: CogniteClient,
+        data_model_id: DataModelIdentifier,
+        reference_model_id: DataModelIdentifier | None = None,
+    ) -> "DMSImporter":
         """Create a DMSImporter ready to convert the given data model to rules.
 
         Args:
             client: Instantiated CogniteClient to retrieve data model.
+            reference_model_id: The reference data model to retrieve. This is the data model that
+                the given data model is built on top of, typically, an enterprise data model.
             data_model_id: Data Model to retrieve.
 
         Returns:
             DMSImporter: DMSImporter instance
         """
-        data_models = client.data_modeling.data_models.retrieve(data_model_id, inline_views=True)
-        if len(data_models) == 0:
+        data_model_ids = [data_model_id, reference_model_id] if reference_model_id else [data_model_id]
+        data_models = client.data_modeling.data_models.retrieve(data_model_ids, inline_views=True)
+
+        user_models = cls._find_model_in_list(data_models, data_model_id)
+        if len(user_models) == 0:
             return cls(DMSSchema(), [issues.importing.NoDataModelError(f"Data model {data_model_id} not found")])
-        data_model = data_models.latest_version()
+        user_model = user_models.latest_version()
+
+        if reference_model_id:
+            ref_models = cls._find_model_in_list(data_models, reference_model_id)
+            if len(ref_models) == 0:
+                return cls(
+                    DMSSchema(), [issues.importing.NoDataModelError(f"Data model {reference_model_id} not found")]
+                )
+            ref_model: dm.DataModel[dm.View] | None = ref_models.latest_version()
+        else:
+            ref_model = None
 
         try:
-            schema = DMSSchema.from_data_model(client, data_model)
+            schema = DMSSchema.from_data_model(client, user_model, ref_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 = cls._create_metadata_from_model(user_model)
+        ref_metadata = cls._create_metadata_from_model(ref_model) if ref_model else None
 
-        metadata = cls._create_metadata_from_model(data_model, created, updated)
+        return cls(schema, [], metadata, ref_metadata)
 
-        return cls(schema, [], metadata)
+    @classmethod
+    def _find_model_in_list(
+        cls, data_models: dm.DataModelList[dm.View], model_id: DataModelIdentifier
+    ) -> dm.DataModelList[dm.View]:
+        identifier = DataModelId.load(model_id)
+        return dm.DataModelList[dm.View](
+            [
+                model
+                for model in data_models
+                if (model.space, model.external_id) == (identifier.space, identifier.external_id)
+            ]
+        )
 
     @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)
+
+        if isinstance(model, dm.DataModel):
+            created = ms_to_datetime(model.created_time)
+            updated = ms_to_datetime(model.last_updated_time)
+        else:
+            now = datetime.now().replace(microsecond=0)
+            created = now
+            updated = now
         return DMSMetadata(
             schema_=SchemaCompleteness.complete,
             extension=ExtensionCategory.addition,
@@ -109,8 +144,8 @@ class DMSImporter(BaseImporter):
             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,
+            updated=updated,
+            created=created,
             creator=creator,
             description=description,
         )
@@ -147,28 +182,33 @@ class DMSImporter(BaseImporter):
             # In case there were errors during the import, the to_rules method will return None
             return self._return_or_raise(self.issue_list, errors)
 
-        if len(self.root_schema.data_models) == 0:
+        if not self.root_schema.data_model:
             self.issue_list.append(issues.importing.NoDataModelError("No data model found."))
             return self._return_or_raise(self.issue_list, errors)
-
+        model = self.root_schema.data_model
         with _handle_issues(
             self.issue_list,
         ) as future:
             schema_completeness = SchemaCompleteness.complete
             data_model_type = DataModelType.enterprise
             reference: DMSRules | None = None
-            if ref_schema := self.root_schema.reference:
+            if (ref_schema := self.root_schema.reference) and (ref_model := ref_schema.data_model):
                 # Reference should always be an enterprise model.
                 reference = DMSRules(
                     **self._create_rule_components(
-                        ref_schema, self._create_default_metadata(ref_schema.views), DataModelType.enterprise
+                        ref_model,
+                        ref_schema,
+                        self.ref_metadata or self._create_default_metadata(list(ref_schema.views.values())),
+                        DataModelType.enterprise,
                     )
                 )
                 schema_completeness = SchemaCompleteness.extended
                 data_model_type = DataModelType.solution
 
             user_rules = DMSRules(
-                **self._create_rule_components(self.root_schema, self.metadata, data_model_type, schema_completeness),
+                **self._create_rule_components(
+                    model, self.root_schema, self.metadata, data_model_type, schema_completeness
+                ),
                 reference=reference,
             )
 
@@ -179,24 +219,14 @@ class DMSImporter(BaseImporter):
 
     def _create_rule_components(
         self,
+        data_model: dm.DataModelApply,
         schema: DMSSchema,
         metadata: DMSMetadata | None = None,
         data_model_type: DataModelType | None = None,
         schema_completeness: SchemaCompleteness | None = None,
     ) -> dict[str, Any]:
-        if len(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 schema.data_models]
-                )
-            )
-
-        data_model = schema.data_models[0]
-
         properties = SheetList[DMSProperty]()
-        for view in schema.views:
-            view_id = view.as_id()
+        for view_id, view in schema.views.items():
             view_entity = ViewEntity.from_id(view_id)
             class_entity = view_entity.as_class()
             for prop_id, prop in (view.properties or {}).items():
@@ -217,10 +247,13 @@ class DMSImporter(BaseImporter):
             metadata=metadata,
             properties=properties,
             containers=SheetList[DMSContainer](
-                data=[DMSContainer.from_container(container) for container in schema.containers]
+                data=[DMSContainer.from_container(container) for container in schema.containers.values()]
             ),
             views=SheetList[DMSView](
-                data=[DMSView.from_view(view, in_model=view.as_id() in data_model_view_ids) for view in schema.views]
+                data=[
+                    DMSView.from_view(view, in_model=view_id in data_model_view_ids)
+                    for view_id, view in schema.views.items()
+                ]
             ),
         )
 

{cognite_neat-0.76.3 → cognite_neat-0.77.1}/cognite/neat/rules/models/dms/_exporter.py

@@ -22,6 +22,7 @@ from cognite.neat.rules.models.entities import (
     ViewPropertyEntity,
 )
 from cognite.neat.rules.models.wrapped_entities import DMSFilter, HasDataFilter, NodeTypeFilter
+from cognite.neat.utils.cdf_classes import ContainerApplyDict, NodeApplyDict, SpaceApplyDict, ViewApplyDict
 
 from ._rules import DMSMetadata, DMSProperty, DMSRules, DMSView
 from ._schema import DMSSchema, PipelineSchema
@@ -42,18 +43,18 @@ class _DMSExporter:
     def __init__(
         self,
         rules: DMSRules,
-        include_ref: bool = True,
         include_pipeline: bool = False,
         instance_space: str | None = None,
     ):
-        self.include_ref = include_ref
         self.include_pipeline = include_pipeline
         self.instance_space = instance_space
         self.rules = rules
         self._ref_schema = rules.reference.as_schema() if rules.reference else None
         if self._ref_schema:
             # We skip version as that will always be missing in the reference
-            self._ref_views_by_id = {dm.ViewId(view.space, view.external_id): view for view in self._ref_schema.views}
+            self._ref_views_by_id = {
+                dm.ViewId(view.space, view.external_id): view for view in self._ref_schema.views.values()
+            }
         else:
             self._ref_views_by_id = {}
 
@@ -67,7 +68,7 @@ class _DMSExporter:
         views_not_in_model = {view.view.as_id() for view in rules.views if not view.in_model}
         data_model = rules.metadata.as_data_model()
         data_model.views = sorted(
-            [view_id for view_id in views.as_ids() if view_id not in views_not_in_model],
+            [view_id for view_id in views.keys() if view_id not in views_not_in_model],
             key=lambda v: v.as_tuple(),  # type: ignore[union-attr]
         )
 
@@ -75,7 +76,7 @@ class _DMSExporter:
 
         output = DMSSchema(
             spaces=spaces,
-            data_models=dm.DataModelApplyList([data_model]),
+            data_model=data_model,
             views=views,
             containers=containers,
             node_types=node_types,
@@ -91,31 +92,30 @@ class _DMSExporter:
     def _create_spaces(
         self,
         metadata: DMSMetadata,
-        containers: dm.ContainerApplyList,
-        views: dm.ViewApplyList,
+        containers: ContainerApplyDict,
+        views: ViewApplyDict,
         data_model: dm.DataModelApply,
-    ) -> dm.SpaceApplyList:
-        used_spaces = {container.space for container in containers} | {view.space for view in views}
+    ) -> SpaceApplyDict:
+        used_spaces = {container.space for container in containers.values()} | {view.space for view in views.values()}
         if len(used_spaces) == 1:
             # We skip the default space and only use this space for the data model
             data_model.space = used_spaces.pop()
-            spaces = dm.SpaceApplyList([dm.SpaceApply(space=data_model.space)])
+            spaces = SpaceApplyDict([dm.SpaceApply(space=data_model.space)])
         else:
             used_spaces.add(metadata.space)
-            spaces = dm.SpaceApplyList([dm.SpaceApply(space=space) for space in used_spaces])
-        if self.instance_space and self.instance_space not in {space.space for space in spaces}:
-            spaces.append(dm.SpaceApply(space=self.instance_space, name=self.instance_space))
+            spaces = SpaceApplyDict([dm.SpaceApply(space=space) for space in used_spaces])
+        if self.instance_space and self.instance_space not in spaces:
+            spaces[self.instance_space] = dm.SpaceApply(space=self.instance_space, name=self.instance_space)
         return spaces
 
     def _create_views_with_node_types(
         self,
         view_properties_by_id: dict[dm.ViewId, list[DMSProperty]],
-    ) -> tuple[dm.ViewApplyList, dm.NodeApplyList]:
-        views = dm.ViewApplyList([dms_view.as_view() for dms_view in self.rules.views])
+    ) -> tuple[ViewApplyDict, NodeApplyDict]:
+        views = ViewApplyDict([dms_view.as_view() for dms_view in self.rules.views])
         dms_view_by_id = {dms_view.view.as_id(): dms_view for dms_view in self.rules.views}
 
-        for view in views:
-            view_id = view.as_id()
+        for view_id, view in views.items():
             view.properties = {}
             if not (view_properties := view_properties_by_id.get(view_id)):
                 continue
@@ -126,14 +126,13 @@ class _DMSExporter:
 
         data_model_type = self.rules.metadata.data_model_type
         unique_node_types: set[dm.NodeId] = set()
-        parent_views = {parent for view in views for parent in view.implements or []}
-        for view in views:
-            dms_view = dms_view_by_id.get(view.as_id())
-            dms_properties = view_properties_by_id.get(view.as_id(), [])
+        parent_views = {parent for view in views.values() for parent in view.implements or []}
+        for view_id, view in views.items():
+            dms_view = dms_view_by_id.get(view_id)
+            dms_properties = view_properties_by_id.get(view_id, [])
             view_filter = self._create_view_filter(view, dms_view, data_model_type, dms_properties)
 
             view.filter = view_filter.as_dms_filter()
-
             if isinstance(view_filter, NodeTypeFilter):
                 unique_node_types.update(view_filter.nodes)
                 if view.as_id() in parent_views:
@@ -153,7 +152,7 @@ class _DMSExporter:
                     issues.dms.HasDataFilterOnViewWithReferencesWarning(view.as_id(), list(references)), stacklevel=2
                 )
 
-        return views, dm.NodeApplyList(
+        return views, NodeApplyDict(
             [dm.NodeApply(space=node.space, external_id=node.external_id) for node in unique_node_types]
         )
 
@@ -176,7 +175,7 @@ class _DMSExporter:
     def _create_containers(
         self,
         container_properties_by_id: dict[dm.ContainerId, list[DMSProperty]],
-    ) -> dm.ContainerApplyList:
+    ) -> ContainerApplyDict:
         containers = dm.ContainerApplyList(
             [dms_container.as_container() for dms_container in self.rules.containers or []]
         )
@@ -231,9 +230,7 @@ class _DMSExporter:
                     for name, const in container.constraints.items()
                     if not (isinstance(const, dm.RequiresConstraint) and const.require in container_to_drop)
                 }
-        return dm.ContainerApplyList(
-            [container for container in containers if container.as_id() not in container_to_drop]
-        )
+        return ContainerApplyDict([container for container in containers if container.as_id() not in container_to_drop])
 
     def _gather_properties(self) -> tuple[dict[dm.ContainerId, list[DMSProperty]], dict[dm.ViewId, list[DMSProperty]]]:
         container_properties_by_id: dict[dm.ContainerId, list[DMSProperty]] = defaultdict(list)
@@ -256,6 +253,7 @@ class _DMSExporter:
         dms_properties: list[DMSProperty],
     ) -> DMSFilter:
         selected_filter_name = (dms_view and dms_view.filter_ and dms_view.filter_.name) or ""
+
         if dms_view and dms_view.filter_ and not dms_view.filter_.is_empty:
             # Has Explicit Filter, use it
             return dms_view.filter_

{cognite_neat-0.76.3 → cognite_neat-0.77.1}/cognite/neat/rules/models/dms/_rules.py

@@ -41,7 +41,7 @@ from cognite.neat.rules.models.entities import (
     ViewEntityList,
     ViewPropertyEntity,
 )
-from cognite.neat.rules.models.wrapped_entities import HasDataFilter, NodeTypeFilter
+from cognite.neat.rules.models.wrapped_entities import HasDataFilter, NodeTypeFilter, RawFilter
 
 from ._schema import DMSSchema
 
@@ -257,7 +257,7 @@ class DMSView(SheetEntity):
     description: str | None = Field(alias="Description", default=None)
     implements: ViewEntityList | None = Field(None, alias="Implements")
     reference: URLEntity | ReferenceEntity | None = Field(alias="Reference", default=None, union_mode="left_to_right")
-    filter_: HasDataFilter | NodeTypeFilter | None = Field(None, alias="Filter")
+    filter_: HasDataFilter | NodeTypeFilter | RawFilter | None = Field(None, alias="Filter")
     in_model: bool = Field(True, alias="In Model")
     class_: ClassEntity = Field(alias="Class (linage)")
 
@@ -345,12 +345,10 @@ class DMSRules(BaseRules):
         space, version = self.metadata.space, self.metadata.version
         return _DMSRulesSerializer(info, space, version).clean(dumped)
 
-    def as_schema(
-        self, include_ref: bool = False, include_pipeline: bool = False, instance_space: str | None = None
-    ) -> DMSSchema:
+    def as_schema(self, include_pipeline: bool = False, instance_space: str | None = None) -> DMSSchema:
         from ._exporter import _DMSExporter
 
-        return _DMSExporter(self, include_ref, include_pipeline, instance_space).to_schema()
+        return _DMSExporter(self, include_pipeline, instance_space).to_schema()
 
     def as_information_architect_rules(self) -> "InformationRules":
         from ._converter import _DMSRulesConverter