cognite-neat 0.76.0__py3-none-any.whl → 0.76.2__py3-none-any.whl

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

@@ -0,0 +1,379 @@
+import math
+import re
+import sys
+import warnings
+from collections.abc import Callable
+from datetime import datetime
+from typing import TYPE_CHECKING, Any, ClassVar, Literal, cast
+
+from cognite.client import data_modeling as dm
+from pydantic import Field, field_serializer, field_validator, model_serializer, model_validator
+from pydantic_core.core_schema import SerializationInfo, ValidationInfo
+
+from cognite.neat.rules import issues
+from cognite.neat.rules.issues import MultiValueError
+from cognite.neat.rules.models._base import (
+    BaseMetadata,
+    BaseRules,
+    DataModelType,
+    ExtensionCategory,
+    RoleTypes,
+    SchemaCompleteness,
+    SheetEntity,
+    SheetList,
+)
+from cognite.neat.rules.models._types import (
+    ExternalIdType,
+    PropertyType,
+    StrListType,
+    VersionType,
+)
+from cognite.neat.rules.models.data_types import DataType
+from cognite.neat.rules.models.domain import DomainRules
+from cognite.neat.rules.models.entities import (
+    ClassEntity,
+    ContainerEntity,
+    ContainerEntityList,
+    DMSUnknownEntity,
+    ReferenceEntity,
+    URLEntity,
+    ViewEntity,
+    ViewEntityList,
+    ViewPropertyEntity,
+)
+from cognite.neat.rules.models.wrapped_entities import HasDataFilter, NodeTypeFilter
+
+from ._schema import DMSSchema
+
+if TYPE_CHECKING:
+    from cognite.neat.rules.models.information._rules import InformationRules
+
+if sys.version_info >= (3, 11):
+    from typing import Self
+else:
+    from typing_extensions import Self
+
+_DEFAULT_VERSION = "1"
+
+
+class DMSMetadata(BaseMetadata):
+    role: ClassVar[RoleTypes] = RoleTypes.dms_architect
+    data_model_type: DataModelType = Field(DataModelType.solution, alias="dataModelType")
+    schema_: SchemaCompleteness = Field(alias="schema")
+    extension: ExtensionCategory = ExtensionCategory.addition
+    space: ExternalIdType
+    name: str | None = Field(
+        None,
+        description="Human readable name of the data model",
+        min_length=1,
+        max_length=255,
+    )
+    description: str | None = Field(None, min_length=1, max_length=1024)
+    external_id: ExternalIdType = Field(alias="externalId")
+    version: VersionType
+    creator: StrListType
+    created: datetime = Field(
+        description=("Date of the data model creation"),
+    )
+    updated: datetime = Field(
+        description=("Date of the data model update"),
+    )
+
+    @field_validator("*", mode="before")
+    def strip_string(cls, value: Any) -> Any:
+        if isinstance(value, str):
+            return value.strip()
+        return value
+
+    @field_serializer("schema_", "extension", "data_model_type", when_used="always")
+    @staticmethod
+    def as_string(value: SchemaCompleteness | ExtensionCategory | DataModelType) -> str:
+        return str(value)
+
+    @field_validator("schema_", mode="plain")
+    def as_enum_schema(cls, value: str) -> SchemaCompleteness:
+        return SchemaCompleteness(value)
+
+    @field_validator("extension", mode="plain")
+    def as_enum_extension(cls, value: str) -> ExtensionCategory:
+        return ExtensionCategory(value)
+
+    @field_validator("data_model_type", mode="plain")
+    def as_enum_model_type(cls, value: str) -> DataModelType:
+        return DataModelType(value)
+
+    @field_validator("description", mode="before")
+    def nan_as_none(cls, value):
+        if isinstance(value, float) and math.isnan(value):
+            return None
+        return value
+
+    def as_space(self) -> dm.SpaceApply:
+        return dm.SpaceApply(
+            space=self.space,
+        )
+
+    def as_data_model_id(self) -> dm.DataModelId:
+        return dm.DataModelId(space=self.space, external_id=self.external_id, version=self.version)
+
+    def as_data_model(self) -> dm.DataModelApply:
+        suffix = f"Creator: {', '.join(self.creator)}"
+        if self.description:
+            description = f"{self.description} Creator: {', '.join(self.creator)}"
+        else:
+            description = suffix
+
+        return dm.DataModelApply(
+            space=self.space,
+            external_id=self.external_id,
+            name=self.name or None,
+            version=self.version or "missing",
+            description=description,
+            views=[],
+        )
+
+    @classmethod
+    def _get_description_and_creator(cls, description_raw: str | None) -> tuple[str | None, list[str]]:
+        if description_raw and (description_match := re.search(r"Creator: (.+)", description_raw)):
+            creator = description_match.group(1).split(", ")
+            description = description_raw.replace(description_match.string, "").strip() or None
+        elif description_raw:
+            creator = ["MISSING"]
+            description = description_raw
+        else:
+            creator = ["MISSING"]
+            description = None
+        return description, creator
+
+    @classmethod
+    def from_data_model(cls, data_model: dm.DataModelApply) -> "DMSMetadata":
+        description, creator = cls._get_description_and_creator(data_model.description)
+        return cls(
+            schema_=SchemaCompleteness.complete,
+            space=data_model.space,
+            name=data_model.name or None,
+            description=description,
+            external_id=data_model.external_id,
+            version=data_model.version,
+            creator=creator,
+            created=datetime.now(),
+            updated=datetime.now(),
+        )
+
+
+class DMSProperty(SheetEntity):
+    view: ViewEntity = Field(alias="View")
+    view_property: str = Field(alias="View Property")
+    name: str | None = Field(alias="Name", default=None)
+    description: str | None = Field(alias="Description", default=None)
+    connection: Literal["direct", "edge", "reverse"] | None = Field(None, alias="Connection")
+    value_type: DataType | ViewPropertyEntity | ViewEntity | DMSUnknownEntity = Field(alias="Value Type")
+    nullable: bool | None = Field(default=None, alias="Nullable")
+    is_list: bool | None = Field(default=None, alias="Is List")
+    default: str | int | dict | None = Field(None, alias="Default")
+    reference: URLEntity | ReferenceEntity | None = Field(default=None, alias="Reference", union_mode="left_to_right")
+    container: ContainerEntity | None = Field(None, alias="Container")
+    container_property: str | None = Field(None, alias="Container Property")
+    index: StrListType | None = Field(None, alias="Index")
+    constraint: StrListType | None = Field(None, alias="Constraint")
+    class_: ClassEntity = Field(alias="Class (linage)")
+    property_: PropertyType = Field(alias="Property (linage)")
+
+    @field_validator("nullable")
+    def direct_relation_must_be_nullable(cls, value: Any, info: ValidationInfo) -> None:
+        if info.data.get("connection") == "direct" and value is False:
+            raise ValueError("Direct relation must be nullable")
+        return value
+
+    @field_validator("value_type", mode="after")
+    def connections_value_type(
+        cls, value: ViewPropertyEntity | ViewEntity | DMSUnknownEntity, info: ValidationInfo
+    ) -> DataType | ViewPropertyEntity | ViewEntity | DMSUnknownEntity:
+        if (connection := info.data.get("connection")) is None:
+            return value
+        if connection == "direct" and not isinstance(value, ViewEntity | DMSUnknownEntity):
+            raise ValueError(f"Direct relation must have a value type that points to a view, got {value}")
+        elif connection == "edge" and not isinstance(value, ViewEntity):
+            raise ValueError(f"Edge connection must have a value type that points to a view, got {value}")
+        elif connection == "reverse" and not isinstance(value, ViewPropertyEntity | ViewEntity):
+            raise ValueError(
+                f"Reverse connection must have a value type that points to a view or view property, got {value}"
+            )
+        return value
+
+    @field_serializer("value_type", when_used="always")
+    @staticmethod
+    def as_dms_type(value_type: DataType | ViewPropertyEntity | ViewEntity) -> str:
+        if isinstance(value_type, DataType):
+            return value_type.dms._type
+        else:
+            return str(value_type)
+
+
+class DMSContainer(SheetEntity):
+    container: ContainerEntity = Field(alias="Container")
+    name: str | None = Field(alias="Name", default=None)
+    description: str | None = Field(alias="Description", default=None)
+    reference: URLEntity | ReferenceEntity | None = Field(alias="Reference", default=None, union_mode="left_to_right")
+    constraint: ContainerEntityList | None = Field(None, alias="Constraint")
+    class_: ClassEntity = Field(alias="Class (linage)")
+
+    def as_container(self) -> dm.ContainerApply:
+        container_id = self.container.as_id()
+        constraints: dict[str, dm.Constraint] = {}
+        for constraint in self.constraint or []:
+            requires = dm.RequiresConstraint(constraint.as_id())
+            constraints[f"{constraint.space}_{constraint.external_id}"] = requires
+
+        return dm.ContainerApply(
+            space=container_id.space,
+            external_id=container_id.external_id,
+            name=self.name or None,
+            description=self.description,
+            constraints=constraints or None,
+            properties={},
+        )
+
+    @classmethod
+    def from_container(cls, container: dm.ContainerApply) -> "DMSContainer":
+        constraints: list[ContainerEntity] = []
+        for _, constraint_obj in (container.constraints or {}).items():
+            if isinstance(constraint_obj, dm.RequiresConstraint):
+                constraints.append(ContainerEntity.from_id(constraint_obj.require))
+            # UniquenessConstraint it handled in the properties
+        container_entity = ContainerEntity.from_id(container.as_id())
+        return cls(
+            class_=container_entity.as_class(),
+            container=container_entity,
+            name=container.name or None,
+            description=container.description,
+            constraint=constraints or None,
+        )
+
+
+class DMSView(SheetEntity):
+    view: ViewEntity = Field(alias="View")
+    name: str | None = Field(alias="Name", default=None)
+    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")
+    in_model: bool = Field(True, alias="In Model")
+    class_: ClassEntity = Field(alias="Class (linage)")
+
+    def as_view(self) -> dm.ViewApply:
+        view_id = self.view.as_id()
+        return dm.ViewApply(
+            space=view_id.space,
+            external_id=view_id.external_id,
+            version=view_id.version or _DEFAULT_VERSION,
+            name=self.name or None,
+            description=self.description,
+            implements=[parent.as_id() for parent in self.implements or []] or None,
+            properties={},
+        )
+
+    @classmethod
+    def from_view(cls, view: dm.ViewApply, in_model: bool) -> "DMSView":
+        view_entity = ViewEntity.from_id(view.as_id())
+        class_entity = view_entity.as_class(skip_version=True)
+
+        return cls(
+            class_=class_entity,
+            view=view_entity,
+            description=view.description,
+            name=view.name,
+            implements=[ViewEntity.from_id(parent, _DEFAULT_VERSION) for parent in view.implements] or None,
+            in_model=in_model,
+        )
+
+
+class DMSRules(BaseRules):
+    metadata: DMSMetadata = Field(alias="Metadata")
+    properties: SheetList[DMSProperty] = Field(alias="Properties")
+    views: SheetList[DMSView] = Field(alias="Views")
+    containers: SheetList[DMSContainer] | None = Field(None, alias="Containers")
+    last: "DMSRules | None" = Field(None, alias="Last", description="The previous version of the data model")
+    reference: "DMSRules | None" = Field(None, alias="Reference")
+
+    @field_validator("reference")
+    def check_reference_of_reference(cls, value: "DMSRules | None", info: ValidationInfo) -> "DMSRules | None":
+        if value is None:
+            return None
+        if value.reference is not None:
+            raise ValueError("Reference rules cannot have a reference")
+        if value.metadata.data_model_type == DataModelType.solution and (metadata := info.data.get("metadata")):
+            warnings.warn(
+                issues.dms.SolutionOnTopOfSolutionModelWarning(
+                    metadata.as_data_model_id(), value.metadata.as_data_model_id()
+                ),
+                stacklevel=2,
+            )
+        return value
+
+    @field_validator("views")
+    def matching_version_and_space(cls, value: SheetList[DMSView], info: ValidationInfo) -> SheetList[DMSView]:
+        if not (metadata := info.data.get("metadata")):
+            return value
+        model_version = metadata.version
+        if different_version := [view.view.as_id() for view in value if view.view.version != model_version]:
+            warnings.warn(issues.dms.ViewModelVersionNotMatchingWarning(different_version, model_version), stacklevel=2)
+        if different_space := [view.view.as_id() for view in value if view.view.space != metadata.space]:
+            warnings.warn(issues.dms.ViewModelSpaceNotMatchingWarning(different_space, metadata.space), stacklevel=2)
+        return value
+
+    @model_validator(mode="after")
+    def post_validation(self) -> "DMSRules":
+        from ._validation import DMSPostValidation
+
+        issue_list = DMSPostValidation(self).validate()
+        if issue_list.warnings:
+            issue_list.trigger_warnings()
+        if issue_list.has_errors:
+            raise MultiValueError([error for error in issue_list if isinstance(error, issues.NeatValidationError)])
+        return self
+
+    @model_serializer(mode="wrap", when_used="always")
+    def dms_rules_serialization(
+        self,
+        handler: Callable,
+        info: SerializationInfo,
+    ) -> dict[str, Any]:
+        from ._serializer import _DMSRulesSerializer
+
+        dumped = cast(dict[str, Any], handler(self, info))
+        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:
+        from ._exporter import _DMSExporter
+
+        return _DMSExporter(self, include_ref, include_pipeline, instance_space).to_schema()
+
+    def as_information_architect_rules(self) -> "InformationRules":
+        from ._converter import _DMSRulesConverter
+
+        return _DMSRulesConverter(self).as_information_architect_rules()
+
+    def as_domain_expert_rules(self) -> DomainRules:
+        from ._converter import _DMSRulesConverter
+
+        return _DMSRulesConverter(self).as_domain_rules()
+
+    def reference_self(self) -> Self:
+        new_rules = self.model_copy(deep=True)
+        for prop in new_rules.properties:
+            prop.reference = ReferenceEntity(
+                prefix=prop.view.prefix, suffix=prop.view.suffix, version=prop.view.version, property=prop.property_
+            )
+        view: DMSView
+        for view in new_rules.views:
+            view.reference = ReferenceEntity(
+                prefix=view.view.prefix, suffix=view.view.suffix, version=view.view.version
+            )
+        container: DMSContainer
+        for container in new_rules.containers or []:
+            container.reference = ReferenceEntity(prefix=container.container.prefix, suffix=container.container.suffix)
+        return new_rules

cognite/neat/rules/models/{rules/_dms_rules_write.py → dms/_rules_input.py}

@@ -5,6 +5,7 @@ from typing import Any, Literal, cast, overload
 
 from pydantic import BaseModel
 
+from cognite.neat.rules.models._base import DataModelType, ExtensionCategory, SchemaCompleteness
 from cognite.neat.rules.models.data_types import DataType
 from cognite.neat.rules.models.entities import (
     ClassEntity,
@@ -15,12 +16,11 @@ from cognite.neat.rules.models.entities import (
     ViewPropertyEntity,
 )
 
-from ._base import DataModelType, ExtensionCategory, SchemaCompleteness
-from ._dms_architect_rules import DMSContainer, DMSMetadata, DMSProperty, DMSRules, DMSView
+from ._rules import DMSContainer, DMSMetadata, DMSProperty, DMSRules, DMSView
 
 
 @dataclass
-class DMSMetadataWrite:
+class DMSMetadataInput:
     schema_: Literal["complete", "partial", "extended"]
     space: str
     external_id: str
@@ -34,7 +34,7 @@ class DMSMetadataWrite:
     updated: datetime | str | None = None
 
     @classmethod
-    def load(cls, data: dict[str, Any] | None) -> "DMSMetadataWrite | None":
+    def load(cls, data: dict[str, Any] | None) -> "DMSMetadataInput | None":
         if data is None:
             return None
         _add_alias(data, DMSMetadata)
@@ -69,7 +69,7 @@ class DMSMetadataWrite:
 
 
 @dataclass
-class DMSPropertyWrite:
+class DMSPropertyInput:
     view: str
     view_property: str | None
     value_type: str
@@ -93,16 +93,16 @@ class DMSPropertyWrite:
 
     @classmethod
     @overload
-    def load(cls, data: dict[str, Any]) -> "DMSPropertyWrite": ...
+    def load(cls, data: dict[str, Any]) -> "DMSPropertyInput": ...
 
     @classmethod
     @overload
-    def load(cls, data: list[dict[str, Any]]) -> list["DMSPropertyWrite"]: ...
+    def load(cls, data: list[dict[str, Any]]) -> list["DMSPropertyInput"]: ...
 
     @classmethod
     def load(
         cls, data: dict[str, Any] | list[dict[str, Any]] | None
-    ) -> "DMSPropertyWrite | list[DMSPropertyWrite] | None":
+    ) -> "DMSPropertyInput | list[DMSPropertyInput] | None":
         if data is None:
             return None
         if isinstance(data, list) or (isinstance(data, dict) and isinstance(data.get("data"), list)):
@@ -146,9 +146,9 @@ class DMSPropertyWrite:
             "View Property": self.view_property,
             "Value Type": value_type,
             "Property (linage)": self.property_ or self.view_property,
-            "Class (linage)": ClassEntity.load(self.class_, prefix=default_space, version=default_version)
-            if self.class_
-            else None,
+            "Class (linage)": (
+                ClassEntity.load(self.class_, prefix=default_space, version=default_version) if self.class_ else None
+            ),
             "Name": self.name,
             "Description": self.description,
             "Connection": self.connection,
@@ -156,9 +156,11 @@ class DMSPropertyWrite:
             "Is List": self.is_list,
             "Default": self.default,
             "Reference": self.reference,
-            "Container": ContainerEntity.load(self.container, space=default_space, version=default_version)
-            if self.container
-            else None,
+            "Container": (
+                ContainerEntity.load(self.container, space=default_space, version=default_version)
+                if self.container
+                else None
+            ),
             "Container Property": self.container_property,
             "Index": self.index,
             "Constraint": self.constraint,
@@ -166,7 +168,7 @@ class DMSPropertyWrite:
 
 
 @dataclass
-class DMSContainerWrite:
+class DMSContainerInput:
     container: str
     class_: str | None = None
     name: str | None = None
@@ -180,16 +182,16 @@ class DMSContainerWrite:
 
     @classmethod
     @overload
-    def load(cls, data: dict[str, Any]) -> "DMSContainerWrite": ...
+    def load(cls, data: dict[str, Any]) -> "DMSContainerInput": ...
 
     @classmethod
     @overload
-    def load(cls, data: list[dict[str, Any]]) -> list["DMSContainerWrite"]: ...
+    def load(cls, data: list[dict[str, Any]]) -> list["DMSContainerInput"]: ...
 
     @classmethod
     def load(
         cls, data: dict[str, Any] | list[dict[str, Any]] | None
-    ) -> "DMSContainerWrite | list[DMSContainerWrite] | None":
+    ) -> "DMSContainerInput | list[DMSContainerInput] | None":
         if data is None:
             return None
         if isinstance(data, list) or (isinstance(data, dict) and isinstance(data.get("data"), list)):
@@ -210,23 +212,25 @@ class DMSContainerWrite:
         container = ContainerEntity.load(self.container, space=default_space)
         return {
             "Container": container,
-            "Class (linage)": ClassEntity.load(self.class_, prefix=default_space)
-            if self.class_
-            else container.as_class(),
+            "Class (linage)": (
+                ClassEntity.load(self.class_, prefix=default_space) if self.class_ else container.as_class()
+            ),
             "Name": self.name,
             "Description": self.description,
             "Reference": self.reference,
-            "Constraint": [
-                ContainerEntity.load(constraint.strip(), space=default_space)
-                for constraint in self.constraint.split(",")
-            ]
-            if self.constraint
-            else None,
+            "Constraint": (
+                [
+                    ContainerEntity.load(constraint.strip(), space=default_space)
+                    for constraint in self.constraint.split(",")
+                ]
+                if self.constraint
+                else None
+            ),
         }
 
 
 @dataclass
-class DMSViewWrite:
+class DMSViewInput:
     view: str
     class_: str | None = None
     name: str | None = None
@@ -242,14 +246,14 @@ class DMSViewWrite:
 
     @classmethod
     @overload
-    def load(cls, data: dict[str, Any]) -> "DMSViewWrite": ...
+    def load(cls, data: dict[str, Any]) -> "DMSViewInput": ...
 
     @classmethod
     @overload
-    def load(cls, data: list[dict[str, Any]]) -> list["DMSViewWrite"]: ...
+    def load(cls, data: list[dict[str, Any]]) -> list["DMSViewInput"]: ...
 
     @classmethod
-    def load(cls, data: dict[str, Any] | list[dict[str, Any]] | None) -> "DMSViewWrite | list[DMSViewWrite] | None":
+    def load(cls, data: dict[str, Any] | list[dict[str, Any]] | None) -> "DMSViewInput | list[DMSViewInput] | None":
         if data is None:
             return None
         if isinstance(data, list) or (isinstance(data, dict) and isinstance(data.get("data"), list)):
@@ -272,17 +276,21 @@ class DMSViewWrite:
         view = ViewEntity.load(self.view, space=default_space, version=default_version)
         return {
             "View": view,
-            "Class (linage)": ClassEntity.load(self.class_, prefix=default_space, version=default_version)
-            if self.class_
-            else view.as_class(),
+            "Class (linage)": (
+                ClassEntity.load(self.class_, prefix=default_space, version=default_version)
+                if self.class_
+                else view.as_class()
+            ),
             "Name": self.name,
             "Description": self.description,
-            "Implements": [
-                ViewEntity.load(implement, space=default_space, version=default_version)
-                for implement in self.implements.split(",")
-            ]
-            if self.implements
-            else None,
+            "Implements": (
+                [
+                    ViewEntity.load(implement, space=default_space, version=default_version)
+                    for implement in self.implements.split(",")
+                ]
+                if self.implements
+                else None
+            ),
             "Reference": self.reference,
             "Filter": self.filter_,
             "In Model": self.in_model,
@@ -290,46 +298,46 @@ class DMSViewWrite:
 
 
 @dataclass
-class DMSRulesWrite:
-    metadata: DMSMetadataWrite
-    properties: Sequence[DMSPropertyWrite]
-    views: Sequence[DMSViewWrite]
-    containers: Sequence[DMSContainerWrite] | None = None
-    reference: "DMSRulesWrite | DMSRules | None" = None
+class DMSRulesInput:
+    metadata: DMSMetadataInput
+    properties: Sequence[DMSPropertyInput]
+    views: Sequence[DMSViewInput]
+    containers: Sequence[DMSContainerInput] | None = None
+    reference: "DMSRulesInput | DMSRules | None" = None
 
     @classmethod
     @overload
-    def load(cls, data: dict[str, Any]) -> "DMSRulesWrite": ...
+    def load(cls, data: dict[str, Any]) -> "DMSRulesInput": ...
 
     @classmethod
     @overload
     def load(cls, data: None) -> None: ...
 
     @classmethod
-    def load(cls, data: dict | None) -> "DMSRulesWrite | None":
+    def load(cls, data: dict | None) -> "DMSRulesInput | None":
         if data is None:
             return None
         _add_alias(data, DMSRules)
         return cls(
-            metadata=DMSMetadataWrite.load(data.get("metadata")),  # type: ignore[arg-type]
-            properties=DMSPropertyWrite.load(data.get("properties")),  # type: ignore[arg-type]
-            views=DMSViewWrite.load(data.get("views")),  # type: ignore[arg-type]
-            containers=DMSContainerWrite.load(data.get("containers")) or [],
-            reference=DMSRulesWrite.load(data.get("reference")),
+            metadata=DMSMetadataInput.load(data.get("metadata")),  # type: ignore[arg-type]
+            properties=DMSPropertyInput.load(data.get("properties")),  # type: ignore[arg-type]
+            views=DMSViewInput.load(data.get("views")),  # type: ignore[arg-type]
+            containers=DMSContainerInput.load(data.get("containers")) or [],
+            reference=DMSRulesInput.load(data.get("reference")),
         )
 
-    def as_read(self) -> DMSRules:
+    def as_rules(self) -> DMSRules:
         return DMSRules.model_validate(self.dump())
 
     def dump(self) -> dict[str, Any]:
         default_space = self.metadata.space
         default_version = self.metadata.version
         reference: dict[str, Any] | None = None
-        if isinstance(self.reference, DMSRulesWrite):
+        if isinstance(self.reference, DMSRulesInput):
             reference = self.reference.dump()
         elif isinstance(self.reference, DMSRules):
-            # We need to load through the DMSRulesWrite to set the correct default space and version
-            reference = DMSRulesWrite.load(self.reference.model_dump()).dump()
+            # We need to load through the DMSRulesInput to set the correct default space and version
+            reference = DMSRulesInput.load(self.reference.model_dump()).dump()
 
         return dict(
             Metadata=self.metadata.dump(),

cognite/neat/rules/models/{rules/_dms_schema.py → dms/_schema.py}

@@ -41,14 +41,16 @@ else:
 
 @dataclass
 class DMSSchema:
-    spaces: dm.SpaceApplyList = field(default_factory=lambda: dm.SpaceApplyList([]))
     data_models: dm.DataModelApplyList = field(default_factory=lambda: dm.DataModelApplyList([]))
+    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 frozen ids are parts of the schema that should not be modified or deleted.
-    # This is used the exporting the schema.
-    frozen_ids: set[dm.ViewId | dm.ContainerId | dm.NodeId] = field(default_factory=set)
+    # The last schema is the previous version of the data model. In the case, extension=additio, this
+    # should not be modified.
+    last: "DMSSchema | None" = None
+    # Reference is typically the Enterprise model, while this is the solution model.
+    reference: "DMSSchema | None" = None
 
     _FIELD_NAME_BY_RESOURCE_TYPE: ClassVar[dict[str, str]] = {
         "container": "containers",
@@ -351,6 +353,10 @@ class DMSSchema:
         defined_spaces = {space.space for space in self.spaces}
         defined_containers = {container.as_id(): container for container in self.containers}
         defined_views = {view.as_id() for view in self.views}
+        if self.reference:
+            defined_spaces |= {space.space for space in self.reference.spaces}
+            defined_containers |= {container.as_id(): container for container in self.reference.containers}
+            defined_views |= {view.as_id() for view in self.reference.views}
 
         for container in self.containers:
             if container.space not in defined_spaces: