[NEAT-236] DMS Exporter Support downloading Solution with Enterprise 🧑‍🧒

Makefile

@@ -1,6 +1,6 @@
 .PHONY: run-explorer run-tests run-linters build-ui build-python build-docker run-docker compose-up
 
-version="0.76.3"
+version="0.77.0"
 run-explorer:
 	@echo "Running explorer API server..."
 	# open "http://localhost:8000/static/index.html" || true

cognite/neat/_version.py

@@ -1 +1 @@
-__version__ = "0.76.3"
+__version__ = "0.77.0"

cognite/neat/rules/importers/_dms2rules.py

@@ -6,7 +6,7 @@
 
 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,
@@ -51,12 +51,14 @@ def __init__(
         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}
         if self.root_schema.reference:
@@ -65,51 +67,87 @@ def __init__(
             )
 
     @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,
             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,
+            updated=updated,
+            created=created,
             creator=creator,
             description=description,
         )
@@ -160,7 +198,10 @@ def to_rules(
                 # Reference should always be an enterprise model.
                 reference = DMSRules(
                     **self._create_rule_components(
-                        ref_model, ref_schema, self._create_default_metadata(ref_schema.views), DataModelType.enterprise
+                        ref_model,
+                        ref_schema,
+                        self.ref_metadata or self._create_default_metadata(ref_schema.views),
+                        DataModelType.enterprise,
                     )
                 )
                 schema_completeness = SchemaCompleteness.extended

cognite/neat/rules/models/dms/_schema.py

@@ -96,18 +96,49 @@ def from_model_id(cls, client: CogniteClient, data_model_id: dm.DataModelIdentif
         return cls.from_data_model(client, data_model)
 
     @classmethod
-    def from_data_model(cls, client: CogniteClient, data_model: dm.DataModel) -> "DMSSchema":
+    def from_data_model(
+        cls,
+        client: CogniteClient,
+        data_model: dm.DataModel[dm.View],
+        reference_model: dm.DataModel[dm.View] | None = None,
+    ) -> "DMSSchema":
+        """Create a schema from a data model.
+
+        If a reference model is provided, the schema will include a reference schema. To determine which views,
+        and containers to put in the reference schema, the following rule is applied:
+
+            If a view or container space is different from the data model space,
+            it will be included in the reference schema.*
+
+        *One exception to this rule is if a view is directly referenced by the data model. In this case, the view will
+        be included in the data model schema, even if the space is different.
+
+        Args:
+            client: The Cognite client used for retrieving components referenced by the data model.
+            data_model: The data model to create the schema from.
+            reference_model: (Optional) The reference model to include in the schema.
+                This is typically the Enterprise model.
+
+        Returns:
+            DMSSchema: The schema created from the data model.
+        """
         views = dm.ViewList(data_model.views)
+
+        data_model_write = data_model.as_write()
+        data_model_write.views = list(views.as_ids())
+
+        if reference_model:
+            views.extend(reference_model.views)
+
         container_ids = views.referenced_containers()
         containers = client.data_modeling.containers.retrieve(list(container_ids))
         cls._append_referenced_containers(client, containers)
 
-        space_read = client.data_modeling.spaces.retrieve(data_model.space)
-        if space_read is None:
-            raise ValueError(f"Space {data_model.space} not found")
-        space = space_read.as_write()
-        data_model_write = data_model.as_write()
-        data_model_write.views = list(views.as_write())
+        space_ids = [data_model.space, reference_model.space] if reference_model else [data_model.space]
+        space_read = client.data_modeling.spaces.retrieve(space_ids)
+        if len(space_read) != len(space_ids):
+            raise ValueError(f"Space(s) {space_read} not found")
+        space_write = space_read.as_write()
 
         view_loader = ViewLoader(client)
         # We need to include parent views in the schema to make sure that the schema is valid.
@@ -122,11 +153,41 @@ def from_data_model(cls, client: CogniteClient, data_model: dm.DataModel) -> "DM
         # The ViewLoader as_write method looks up parents and remove properties from them.
         view_write = dm.ViewApplyList([view_loader.as_write(view) for view in views])
 
+        container_write = containers.as_write()
+        user_space = data_model.space
+        if reference_model:
+            user_model_view_ids = set(data_model_write.views)
+            ref_model_write = reference_model.as_write()
+            ref_model_write.views = [view.as_id() for view in reference_model.views]
+
+            ref_views = dm.ViewApplyList(
+                [view for view in view_write if (view.space != user_space) or (view.as_id() not in user_model_view_ids)]
+            )
+            view_write = dm.ViewApplyList(
+                [view for view in view_write if view.space == user_space or view.as_id() in user_model_view_ids]
+            )
+
+            ref_containers = dm.ContainerApplyList(
+                [container for container in container_write if container.space != user_space]
+            )
+            container_write = dm.ContainerApplyList(
+                [container for container in container_write if container.space == user_space]
+            )
+
+            ref_schema: DMSSchema | None = cls(
+                spaces=dm.SpaceApplyList([s for s in space_write if s.space != user_space]),
+                data_model=ref_model_write,
+                views=ref_views,
+                containers=ref_containers,
+            )
+        else:
+            ref_schema = None
         return cls(
-            spaces=dm.SpaceApplyList([space]),
+            spaces=dm.SpaceApplyList([s for s in space_write if s.space == user_space]),
             data_model=data_model_write,
             views=view_write,
-            containers=containers.as_write(),
+            containers=container_write,
+            reference=ref_schema,
         )
 
     @classmethod
@@ -531,17 +592,25 @@ def _append_referenced_containers(cls, client: CogniteClient, containers: dm.Con
             )
         return None
 
-    def referenced_spaces(self) -> set[str]:
-        referenced_spaces = {container.space for container in self.containers}
-        referenced_spaces |= {view.space for view in self.views}
-        referenced_spaces |= {container.space for view in self.views for container in view.referenced_containers()}
-        referenced_spaces |= {parent.space for view in self.views for parent in view.implements or []}
+    def referenced_spaces(self, include_indirect_references: bool = True) -> set[str]:
+        """Get the spaces referenced by the schema.
+
+        Args:
+            include_indirect_references (bool): If True, the spaces referenced by as view.implements, and
+                view.referenced_containers will be included in the output.
+        Returns:
+            set[str]: The spaces referenced by the schema.
+        """
+        referenced_spaces = {view.space for view in self.views}
+        referenced_spaces |= {container.space for container in self.containers}
+        if include_indirect_references:
+            referenced_spaces |= {container.space for view in self.views for container in view.referenced_containers()}
+            referenced_spaces |= {parent.space for view in self.views for parent in view.implements or []}
         referenced_spaces |= {node.space for node in self.node_types}
         if self.data_model:
             referenced_spaces |= {self.data_model.space}
             referenced_spaces |= {view.space for view in self.data_model.views or []}
         referenced_spaces |= {s.space for s in self.spaces}
-
         return referenced_spaces
 
 

docs/CHANGELOG.md

@@ -15,7 +15,7 @@ Changes are grouped as follows:
 - `Fixed` for any bug fixes.
 - `Security` in case of vulnerabilities.
 
-## TBD
+## [0.77.0] - 13-05-24
 ### Changed
 - [BREAKING] The subpackage `cognite.neat.rules.models` is reorganized. All imports using this subpackage must be
   updated.
@@ -24,6 +24,8 @@ Changes are grouped as follows:
 - Support for exporting/importing `Last` spreadsheets in the `ExcelExporter` and `ExcelImporter`.
 - [BREAKING] As a result of the above, in the `ExcelExporter` the parameter `is_reference` is replaced by `dump_as`.
   To continue using the old behavior, set `dump_as='reference'`.
+- In the `DMSImporter.from_data_model_id`, now supports setting `reference_model_id` to download a solution model
+  with a reference model.
 
 ## [0.76.3] - 10-05-24
 ### Added

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "cognite-neat"
-version = "0.76.3"
+version = "0.77.0"
 readme = "README.md"
 description = "Knowledge graph transformation"
 authors = [

tests/tests_integration/test_rules/test_importers/test_dms_importers.py

@@ -21,10 +21,40 @@ def alice_data_model_id(alice_rules: DMSRules) -> DataModelId:
     return alice_rules.metadata.as_data_model_id()
 
 
+@pytest.fixture(scope="session")
+def olav_rules() -> DMSRules:
+    filepath = DOC_RULES / "dms-analytics-olav.xlsx"
+
+    excel_importer = ExcelImporter(filepath)
+
+    return excel_importer.to_rules(errors="raise", role=RoleTypes.dms_architect)
+
+
+@pytest.fixture(scope="session")
+def olav_data_model_id(olav_rules: DMSRules) -> DataModelId:
+    return olav_rules.metadata.as_data_model_id()
+
+
 class TestDMSImporter:
-    def test_import_from_cdf(self, cognite_client: CogniteClient, alice_data_model_id: DataModelId):
+    def test_import_alice_from_cdf(self, cognite_client: CogniteClient, alice_data_model_id: DataModelId):
         dms_exporter = DMSImporter.from_data_model_id(cognite_client, alice_data_model_id)
 
         rules = dms_exporter.to_rules(errors="raise", role=RoleTypes.information_architect)
 
         assert isinstance(rules, InformationRules)
+
+    def test_import_olav_from_cdf(
+        self, cognite_client: CogniteClient, olav_data_model_id: DataModelId, alice_data_model_id: DataModelId
+    ):
+        dms_exporter = DMSImporter.from_data_model_id(cognite_client, olav_data_model_id, alice_data_model_id)
+
+        assert dms_exporter.root_schema.referenced_spaces(include_indirect_references=False) == {
+            olav_data_model_id.space
+        }, "The direct reference should be the data model space."
+
+        rules = dms_exporter.to_rules(errors="raise", role=RoleTypes.dms_architect)
+
+        assert isinstance(rules, DMSRules)
+        assert rules.metadata.as_data_model_id() == olav_data_model_id
+        assert isinstance(rules.reference, DMSRules)
+        assert rules.reference.metadata.as_data_model_id() == alice_data_model_id