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

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

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

cognite_neat-0.77.0/cognite/neat/_version.py

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

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

@@ -5,7 +5,7 @@ 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 DataModelApply, DataModelApplyList, DataModelId
 from cognite.client.exceptions import CogniteAPIError
 
 from cognite.neat.rules import issues
@@ -114,9 +114,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)
@@ -272,7 +270,7 @@ class DMSExporter(CDFExporter[DMSSchema]):
         if self.export_components.intersection({"all", "views"}):
             to_export.append((schema.views, 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.0}/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,12 +51,14 @@ 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}
         if self.root_schema.reference:
@@ -66,42 +67,78 @@ class DMSImporter(BaseImporter):
             )
 
     @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 +146,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 +184,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(ref_schema.views),
+                        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,21 +221,12 @@ 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()

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

@@ -42,11 +42,9 @@ 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
@@ -75,7 +73,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,

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

@@ -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

{cognite_neat-0.76.3 → cognite_neat-0.77.0}/cognite/neat/rules/models/dms/_schema.py

@@ -3,6 +3,7 @@ import sys
 import warnings
 import zipfile
 from collections import Counter, defaultdict
+from collections.abc import Sequence
 from dataclasses import Field, dataclass, field, fields
 from pathlib import Path
 from typing import Any, ClassVar, cast
@@ -43,12 +44,12 @@ else:
 
 @dataclass
 class DMSSchema:
-    data_models: dm.DataModelApplyList = field(default_factory=lambda: dm.DataModelApplyList([]))
+    data_model: dm.DataModelApply | None = None
     spaces: dm.SpaceApplyList = field(default_factory=lambda: dm.SpaceApplyList([]))
     views: dm.ViewApplyList = field(default_factory=lambda: dm.ViewApplyList([]))
     containers: dm.ContainerApplyList = field(default_factory=lambda: dm.ContainerApplyList([]))
     node_types: dm.NodeApplyList = field(default_factory=lambda: dm.NodeApplyList([]))
-    # The last schema is the previous version of the data model. In the case, extension=additio, this
+    # The last schema is the previous version of the data model. In the case, extension=addition, this
     # should not be modified.
     last: "DMSSchema | None" = None
     # Reference is typically the Enterprise model, while this is the solution model.
@@ -57,7 +58,7 @@ class DMSSchema:
     _FIELD_NAME_BY_RESOURCE_TYPE: ClassVar[dict[str, str]] = {
         "container": "containers",
         "view": "views",
-        "datamodel": "data_models",
+        "datamodel": "data_model",
         "space": "spaces",
         "node": "node_types",
     }
@@ -95,18 +96,49 @@ class DMSSchema:
         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.
@@ -121,11 +153,41 @@ class DMSSchema:
         # 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]),
-            data_models=dm.DataModelApplyList([data_model_write]),
+            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
@@ -187,11 +249,10 @@ class DMSSchema:
                 (data_models / f"{space.space}.space.yaml").write_text(
                     space.dump_yaml(), newline=new_line, encoding=encoding
                 )
-        if "data_models" not in exclude_set:
-            for model in self.data_models:
-                (data_models / f"{model.external_id}.datamodel.yaml").write_text(
-                    model.dump_yaml(), newline=new_line, encoding=encoding
-                )
+        if "data_models" not in exclude_set and self.data_model:
+            (data_models / f"{self.data_model.external_id}.datamodel.yaml").write_text(
+                self.data_model.dump_yaml(), newline=new_line, encoding=encoding
+            )
         if "views" not in exclude_set and self.views:
             view_dir = data_models / "views"
             view_dir.mkdir(parents=True, exist_ok=True)
@@ -265,9 +326,10 @@ class DMSSchema:
             if "spaces" not in exclude_set:
                 for space in self.spaces:
                     zip_ref.writestr(f"data_models/{space.space}.space.yaml", space.dump_yaml())
-            if "data_models" not in exclude_set:
-                for model in self.data_models:
-                    zip_ref.writestr(f"data_models/{model.external_id}.datamodel.yaml", model.dump_yaml())
+            if "data_models" not in exclude_set and self.data_model:
+                zip_ref.writestr(
+                    f"data_models/{self.data_model.external_id}.datamodel.yaml", self.data_model.dump_yaml()
+                )
             if "views" not in exclude_set:
                 for view in self.views:
                     zip_ref.writestr(f"data_models/views/{view.external_id}.view.yaml", view.dump_yaml())
@@ -308,10 +370,30 @@ class DMSSchema:
         loaded: dict[str, Any] = {}
         for attr in fields(cls):
             if items := data_dict.get(attr.name) or data_dict.get(to_camel(attr.name)):
-                try:
-                    loaded[attr.name] = attr.type.load(items)
-                except Exception as e:
-                    loaded[attr.name] = cls._load_individual_resources(items, attr, str(e), context.get(attr.name, []))
+                if attr.name == "data_model":
+                    if isinstance(items, list) and len(items) > 1:
+                        warnings.warn(
+                            issues.importing.MultipleDataModelsWarning(
+                                [item.get("externalId", "Unknown") for item in items]
+                            ),
+                            stacklevel=2,
+                        )
+                    item = items[0] if isinstance(items, list) else items
+                    try:
+                        loaded[attr.name] = dm.DataModelApply.load(item)
+                    except Exception as e:
+                        data_model_file = context.get(attr.name, [Path("UNKNOWN")])[0]
+                        warnings.warn(
+                            issues.fileread.FailedLoadWarning(data_model_file, dm.DataModelApply.__name__, str(e)),
+                            stacklevel=2,
+                        )
+                else:
+                    try:
+                        loaded[attr.name] = attr.type.load(items)
+                    except Exception as e:
+                        loaded[attr.name] = cls._load_individual_resources(
+                            items, attr, str(e), context.get(attr.name, [])
+                        )
         return cls(**loaded)
 
     @classmethod
@@ -353,9 +435,12 @@ class DMSSchema:
         cls_fields = sorted(fields(self), key=lambda f: f.name) if sort else fields(self)
         for attr in cls_fields:
             if items := getattr(self, attr.name):
-                items = sorted(items, key=self._to_sortable_identifier) if sort else items
                 key = to_camel(attr.name) if camel_case else attr.name
-                output[key] = [item.dump(camel_case=camel_case) for item in items]
+                if isinstance(items, Sequence):
+                    items = sorted(items, key=self._to_sortable_identifier) if sort else items
+                    output[key] = [item.dump(camel_case=camel_case) for item in items]
+                else:
+                    output[key] = items.dump(camel_case=camel_case)
         return output
 
     @classmethod
@@ -458,7 +543,8 @@ class DMSSchema:
                         )
                     )
 
-        for model in self.data_models:
+        if self.data_model:
+            model = self.data_model
             if model.space not in defined_spaces:
                 errors.add(MissingSpaceError(space=model.space, referred_by=model.as_id()))
 
@@ -506,16 +592,25 @@ class DMSSchema:
             )
         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}
-        referenced_spaces |= {model.space for model in self.data_models}
-        referenced_spaces |= {view.space for model in self.data_models for view in model.views or []}
+        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
 
 
@@ -625,9 +720,9 @@ class PipelineSchema(DMSSchema):
 
     @classmethod
     def from_dms(cls, schema: DMSSchema, instance_space: str | None = None) -> "PipelineSchema":
-        if not schema.data_models:
+        if not schema.data_model:
             raise ValueError("PipelineSchema must contain at least one data model")
-        first_data_model = schema.data_models[0]
+        first_data_model = schema.data_model
         # The database name is limited to 32 characters
         database_name = first_data_model.external_id[:32]
         instance_space = instance_space or first_data_model.space
@@ -666,7 +761,7 @@ class PipelineSchema(DMSSchema):
 
         return cls(
             spaces=schema.spaces,
-            data_models=schema.data_models,
+            data_model=schema.data_model,
             views=schema.views,
             containers=schema.containers,
             transformations=transformations,

{cognite_neat-0.76.3 → cognite_neat-0.77.0}/cognite/neat/rules/models/dms/_validation.py

@@ -164,7 +164,7 @@ class DMSPostValidation:
             # Everything is allowed
             return None
         # Is an extension of an existing model.
-        user_schema = self.rules.as_schema(include_ref=False)
+        user_schema = self.rules.as_schema()
         ref_schema = self.rules.reference.as_schema()
         new_containers = {container.as_id(): container for container in user_schema.containers}
         existing_containers = {container.as_id(): container for container in ref_schema.containers}

{cognite_neat-0.76.3 → cognite_neat-0.77.0}/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 = [

cognite_neat-0.76.3/cognite/neat/_version.py

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