cognite-neat 0.98.0__py3-none-any.whl → 0.99.1__py3-none-any.whl

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

@@ -1,36 +1,50 @@
-from collections import defaultdict
-from typing import Any, ClassVar, cast
+import warnings
+from collections import Counter, defaultdict
+from typing import ClassVar
 
 from cognite.client import data_modeling as dm
+from cognite.client.data_classes.data_modeling import ContainerList, ViewId, ViewList
+from cognite.client.data_classes.data_modeling.views import (
+    ReverseDirectRelation,
+    ReverseDirectRelationApply,
+    ViewProperty,
+    ViewPropertyApply,
+)
 
-from cognite.neat._constants import COGNITE_MODELS, DMS_CONTAINER_PROPERTY_SIZE_LIMIT
-from cognite.neat._issues import IssueList, NeatError, NeatIssue, NeatIssueList
+from cognite.neat._client import NeatClient
+from cognite.neat._client.data_classes.data_modeling import ViewApplyDict
+from cognite.neat._client.data_classes.schema import DMSSchema
+from cognite.neat._constants import COGNITE_MODELS, DMS_CONTAINER_PROPERTY_SIZE_LIMIT, DMS_VIEW_CONTAINER_SIZE_LIMIT
+from cognite.neat._issues import IssueList, NeatError, NeatIssueList
 from cognite.neat._issues.errors import (
+    CDFMissingClientError,
     PropertyDefinitionDuplicatedError,
-    ResourceNotDefinedError,
+    PropertyMappingDuplicatedError,
+    PropertyNotFoundError,
+    ResourceDuplicatedError,
+    ResourceNotFoundError,
+    ReversedConnectionNotFeasibleError,
 )
-from cognite.neat._issues.errors._properties import ReversedConnectionNotFeasibleError
 from cognite.neat._issues.warnings import (
     NotSupportedHasDataFilterLimitWarning,
     NotSupportedViewContainerLimitWarning,
     UndefinedViewWarning,
 )
 from cognite.neat._issues.warnings.user_modeling import (
+    ContainerPropertyLimitWarning,
+    DirectRelationMissingSourceWarning,
     NotNeatSupportedFilterWarning,
-    ViewPropertyLimitWarning,
 )
 from cognite.neat._rules.models.data_types import DataType
 from cognite.neat._rules.models.entities import ContainerEntity, RawFilter
 from cognite.neat._rules.models.entities._single_value import (
-    ReverseConnectionEntity,
     ViewEntity,
 )
 
 from ._rules import DMSProperty, DMSRules
-from ._schema import DMSSchema
 
 
-class DMSPostValidation:
+class DMSValidation:
     """This class does all the validation of the DMS rules that have dependencies between
     components."""
 
@@ -38,32 +52,132 @@ class DMSPostValidation:
     # For example, changing the filter is allowed, but changing the properties is not.
     changeable_view_attributes: ClassVar[set[str]] = {"filter"}
 
-    def __init__(self, rules: DMSRules):
-        self.rules = rules
-        self.metadata = rules.metadata
-        self.properties = rules.properties
-        self.containers = rules.containers
-        self.views = rules.views
-        self.issue_list = IssueList()
+    def __init__(self, rules: DMSRules, client: NeatClient | None = None) -> None:
+        self._rules = rules
+        self._client = client
+        self._metadata = rules.metadata
+        self._properties = rules.properties
+        self._containers = rules.containers
+        self._views = rules.views
+
+    def imported_views_and_containers_ids(
+        self, include_views_with_no_properties: bool = True
+    ) -> tuple[set[ViewEntity], set[ContainerEntity]]:
+        existing_views = {view.view for view in self._views}
+        imported_views: set[ViewEntity] = set()
+        for view in self._views:
+            for parent in view.implements or []:
+                if parent not in existing_views:
+                    imported_views.add(parent)
+        existing_containers = {container.container for container in self._containers or []}
+        imported_containers: set[ContainerEntity] = set()
+        view_with_properties: set[ViewEntity] = set()
+        for prop in self._properties:
+            if prop.container and prop.container not in existing_containers:
+                imported_containers.add(prop.container)
+            if prop.view not in existing_views:
+                imported_views.add(prop.view)
+            view_with_properties.add(prop.view)
+
+        if include_views_with_no_properties:
+            extra_views = existing_views - view_with_properties
+            imported_views.update({view for view in extra_views})
+
+        return imported_views, imported_containers
 
     def validate(self) -> NeatIssueList:
-        self._validate_raw_filter()
-        self._consistent_container_properties()
-        self._validate_value_type_existence()
-        self._validate_reverse_connections()
+        imported_views, imported_containers = self.imported_views_and_containers_ids(
+            include_views_with_no_properties=False
+        )
+        if (imported_views or imported_containers) and self._client is None:
+            raise CDFMissingClientError(
+                f"{self._rules.metadata.as_data_model_id()} has imported views and/or container: "
+                f"{imported_views}, {imported_containers}."
+            )
+        referenced_views = ViewList([])
+        referenced_containers = ContainerList([])
+        if self._client:
+            referenced_views = self._client.loaders.views.retrieve(
+                list(imported_views), include_connected=True, include_ancestor=True
+            )
+            referenced_containers = self._client.loaders.containers.retrieve(
+                list(imported_containers), include_connected=True
+            )
+
+        # Setup data structures for validation
+        dms_schema = self._rules.as_schema()
+        ref_view_by_id = {view.as_id(): view for view in referenced_views}
+        ref_container_by_id = {container.as_id(): container for container in referenced_containers}
+        all_containers_by_id: dict[dm.ContainerId, dm.ContainerApply | dm.Container] = {
+            **dict(dms_schema.containers.items()),
+            **ref_container_by_id,
+        }
+        all_views_by_id: dict[dm.ViewId, dm.ViewApply | dm.View] = {**dict(dms_schema.views.items()), **ref_view_by_id}
+        properties_by_ids = self._as_properties_by_ids(dms_schema, ref_view_by_id)
+        view_properties_by_id: dict[dm.ViewId, list[tuple[str, ViewProperty | ViewPropertyApply]]] = defaultdict(list)
+        for (view_id, prop_id), prop in properties_by_ids.items():
+            view_properties_by_id[view_id].append((prop_id, prop))
+
+        issue_list = IssueList()
+        # Neat DMS classes Validation
+        # These are errors that can only happen due to the format of the Neat DMS classes
+        issue_list.extend(self._validate_raw_filter())
+        issue_list.extend(self._consistent_container_properties())
+        issue_list.extend(self._validate_value_type_existence())
+        issue_list.extend(
+            self._validate_property_referenced_views_and_containers_exists(all_views_by_id, all_containers_by_id)
+        )
+
+        # SDK classes validation
+        issue_list.extend(self._containers_are_proper_size(dms_schema))
+        issue_list.extend(self._validate_reverse_connections(properties_by_ids, all_containers_by_id))
+        issue_list.extend(self._validate_schema(dms_schema, all_views_by_id, all_containers_by_id))
+        issue_list.extend(self._validate_referenced_container_limits(dms_schema.views, view_properties_by_id))
+        return issue_list
+
+    @staticmethod
+    def _as_properties_by_ids(
+        dms_schema: DMSSchema, ref_view_by_id: dict[dm.ViewId, dm.View]
+    ) -> dict[tuple[ViewId, str], ViewPropertyApply | ViewProperty]:
+        # Priority DMS schema properties.
+        # No need to do long lookups in ref_views as these already contain all ancestor properties.
+        properties_by_id: dict[tuple[ViewId, str], ViewPropertyApply | ViewProperty] = {}
+        for view in dms_schema.views.values():
+            view_id = view.as_id()
+            for prop_id, prop in (view.properties or {}).items():
+                properties_by_id[(view_id, prop_id)] = prop
+            if view.implements:
+                to_check = view.implements.copy()
+                while to_check:
+                    parent_id = to_check.pop()
+                    if parent_id in dms_schema.views:
+                        # Priority DMS Schema properties
+                        parent_view = dms_schema.views[parent_id]
+                        for prop_id, prop in (parent_view.properties or {}).items():
+                            if (view_id, prop_id) not in properties_by_id:
+                                properties_by_id[(view_id, prop_id)] = prop
+                        to_check.extend(parent_view.implements or [])
+                    elif parent_id in ref_view_by_id:
+                        # SDK properties
+                        parent_read_view = ref_view_by_id[parent_id]
+                        for prop_id, read_prop in parent_read_view.properties.items():
+                            if (view_id, prop_id) not in properties_by_id:
+                                properties_by_id[(view_id, prop_id)] = read_prop
+                        # Read format of views already includes all ancestor properties
+                        # so no need to check further
+                    else:
+                        # Missing views are caught else where
+                        continue
 
-        self._referenced_views_and_containers_are_existing_and_proper_size()
-        dms_schema = self.rules.as_schema()
-        self._validate_performance(dms_schema)
-        return self.issue_list
+        return properties_by_id
 
-    def _consistent_container_properties(self) -> None:
+    def _consistent_container_properties(self) -> IssueList:
         container_properties_by_id: dict[tuple[ContainerEntity, str], list[tuple[int, DMSProperty]]] = defaultdict(list)
-        for prop_no, prop in enumerate(self.properties):
+        for prop_no, prop in enumerate(self._properties):
             if prop.container and prop.container_property:
                 container_properties_by_id[(prop.container, prop.container_property)].append((prop_no, prop))
 
-        errors: list[NeatError] = []
+        errors = IssueList()
         for (container, prop_name), properties in container_properties_by_id.items():
             if len(properties) == 1:
                 continue
@@ -150,132 +264,290 @@ class DMSPostValidation:
                     )
                 )
 
-        self.issue_list.extend(errors)
+        return errors
 
-    def _referenced_views_and_containers_are_existing_and_proper_size(self) -> None:
-        # TODO: Split this method and keep only validation that should be independent of
-        # whether view and/or container exist in the pydantic model instance
-        # other validation should be done through NeatSession.verify()
-        defined_views = {view.view.as_id() for view in self.views}
-
-        property_count_by_view: dict[dm.ViewId, int] = defaultdict(int)
-        errors: list[NeatIssue] = []
-        for prop_no, prop in enumerate(self.properties):
-            view_id = prop.view.as_id()
-            if view_id not in defined_views:
-                errors.append(
-                    ResourceNotDefinedError[dm.ViewId](
-                        identifier=view_id,
-                        resource_type="view",
-                        location="Views Sheet",
-                        column_name="View",
-                        row_number=prop_no,
-                        sheet_name="Properties",
-                    ),
-                )
-            else:
-                property_count_by_view[view_id] += 1
-        for view_id, count in property_count_by_view.items():
+    @staticmethod
+    def _containers_are_proper_size(dms_schema: DMSSchema) -> IssueList:
+        errors = IssueList()
+        for container_id, container in dms_schema.containers.items():
+            count = len(container.properties or {})
             if count > DMS_CONTAINER_PROPERTY_SIZE_LIMIT:
-                errors.append(ViewPropertyLimitWarning(view_id, count))
+                errors.append(ContainerPropertyLimitWarning(container_id, count))
 
-        self.issue_list.extend(errors)
+        return errors
 
-    def _validate_performance(self, dms_schema: DMSSchema) -> None:
-        for view_id, view in dms_schema.views.items():
-            mapped_containers = dms_schema._get_mapped_container_from_view(view_id)
+    @staticmethod
+    def _validate_referenced_container_limits(
+        views: ViewApplyDict, view_properties_by_id: dict[dm.ViewId, list[tuple[str, ViewProperty | ViewPropertyApply]]]
+    ) -> IssueList:
+        issue_list = IssueList()
+        for view_id, view in views.items():
+            view_properties = view_properties_by_id.get(view_id, [])
+            mapped_containers = {
+                prop.container
+                for _, prop in view_properties
+                if isinstance(prop, dm.MappedPropertyApply | dm.MappedProperty)
+            }
 
-            if mapped_containers and len(mapped_containers) > 10:
-                self.issue_list.append(
+            if mapped_containers and len(mapped_containers) > DMS_VIEW_CONTAINER_SIZE_LIMIT:
+                issue_list.append(
                     NotSupportedViewContainerLimitWarning(
                         view_id,
                         len(mapped_containers),
                     )
                 )
-                if (
-                    view.filter
-                    and isinstance(view.filter, dm.filters.HasData)
-                    and len(view.filter.dump()["hasData"]) > 10
-                ):
-                    self.issue_list.append(
-                        NotSupportedHasDataFilterLimitWarning(
-                            view_id,
-                            len(view.filter.dump()["hasData"]),
-                        )
+
+            if view.filter and isinstance(view.filter, dm.filters.HasData) and len(view.filter.dump()["hasData"]) > 10:
+                issue_list.append(
+                    NotSupportedHasDataFilterLimitWarning(
+                        view_id,
+                        len(view.filter.dump()["hasData"]),
                     )
+                )
+        return issue_list
 
-    def _validate_raw_filter(self) -> None:
-        for view in self.views:
+    def _validate_raw_filter(self) -> IssueList:
+        issue_list = IssueList()
+        for view in self._views:
             if view.filter_ and isinstance(view.filter_, RawFilter):
-                self.issue_list.append(
+                issue_list.append(
                     NotNeatSupportedFilterWarning(view.view.as_id()),
                 )
+        return issue_list
 
-    def _validate_value_type_existence(self) -> None:
-        views = {prop_.view for prop_ in self.properties}.union({view_.view for view_ in self.views})
-
-        for prop_ in self.properties:
+    def _validate_value_type_existence(self) -> IssueList:
+        views = {prop_.view for prop_ in self._properties}.union({view_.view for view_ in self._views})
+        issue_list = IssueList()
+        for prop_ in self._properties:
             if isinstance(prop_.value_type, ViewEntity) and prop_.value_type not in views:
-                self.issue_list.append(
+                issue_list.append(
                     UndefinedViewWarning(
                         str(prop_.view),
                         str(prop_.value_type),
                         prop_.view_property,
                     )
                 )
+        return issue_list
 
-    def _validate_reverse_connections(self) -> None:
+    def _validate_property_referenced_views_and_containers_exists(
+        self,
+        view_by_id: dict[dm.ViewId, dm.ViewApply | dm.View],
+        containers_by_id: dict[dm.ContainerId, dm.ContainerApply | dm.Container],
+    ) -> IssueList:
+        issue_list = IssueList()
+        for prop in self._properties:
+            if prop.container:
+                container_id = prop.container.as_id()
+                if container_id not in containers_by_id:
+                    issue_list.append(
+                        ResourceNotFoundError(
+                            container_id,
+                            "container",
+                            prop.view,
+                            "view",
+                        )
+                    )
+                elif (
+                    prop.container_property and prop.container_property not in containers_by_id[container_id].properties
+                ):
+                    issue_list.append(
+                        PropertyNotFoundError(
+                            prop.container,
+                            "container property",
+                            prop.container_property,
+                            dm.PropertyId(prop.view.as_id(), prop.view_property),
+                            "view property",
+                        )
+                    )
+
+            if prop.view.as_id() not in view_by_id:
+                issue_list.append(
+                    ResourceNotFoundError(
+                        prop.view,
+                        "view",
+                        prop.view_property,
+                        "property",
+                    )
+                )
+
+        return issue_list
+
+    def _validate_reverse_connections(
+        self,
+        properties_by_ids: dict[tuple[dm.ViewId, str], ViewPropertyApply | ViewProperty],
+        containers_by_id: dict[dm.ContainerId, dm.ContainerApply | dm.Container],
+    ) -> IssueList:
+        issue_list = IssueList()
         # do not check for reverse connections in Cognite models
-        if self.metadata.as_data_model_id() in COGNITE_MODELS:
-            return None
-
-        properties_by_ids = {f"{prop_.view!s}.{prop_.view_property}": prop_ for prop_ in self.properties}
-        reversed_by_ids = {
-            id_: prop_
-            for id_, prop_ in properties_by_ids.items()
-            if prop_.connection and isinstance(prop_.connection, ReverseConnectionEntity)
-        }
+        if self._metadata.as_data_model_id() in COGNITE_MODELS:
+            return issue_list
 
-        for id_, prop_ in reversed_by_ids.items():
-            source_id = f"{prop_.value_type!s}." f"{cast(ReverseConnectionEntity, prop_.connection).property_}"
+        for (view_id, prop_id), prop_ in properties_by_ids.items():
+            if not isinstance(prop_, ReverseDirectRelationApply | ReverseDirectRelation):
+                continue
+            source_id = prop_.through.source, prop_.through.property
             if source_id not in properties_by_ids:
-                print(f"source_id: {source_id}, first issue")
-                self.issue_list.append(
+                issue_list.append(
                     ReversedConnectionNotFeasibleError(
-                        id_,
+                        view_id,
                         "reversed connection",
-                        prop_.view_property,
-                        str(prop_.view),
-                        str(prop_.value_type),
-                        cast(ReverseConnectionEntity, prop_.connection).property_,
+                        prop_id,
+                        f"The {prop_.through.source} {prop_.through.property} does not exist",
                     )
                 )
+                continue
+            if isinstance(source_id[0], dm.ContainerId):
+                # Todo: How to handle this case? Should not happen if you created the model with Neat
+                continue
 
-            elif source_id in properties_by_ids and properties_by_ids[source_id].value_type != prop_.view:
-                print(f"source_id: {source_id}, second issue")
-                self.issue_list.append(
+            target_property = properties_by_ids[(source_id[0], source_id[1])]
+            # Validate that the target is a direct relation pointing to the view_id
+            is_direct_relation = False
+            if isinstance(target_property, dm.MappedProperty) and isinstance(target_property.type, dm.DirectRelation):
+                is_direct_relation = True
+            elif isinstance(target_property, dm.MappedPropertyApply):
+                container = containers_by_id[target_property.container]
+                if target_property.container_property_identifier in container.properties:
+                    container_property = container.properties[target_property.container_property_identifier]
+                    if isinstance(container_property.type, dm.DirectRelation):
+                        is_direct_relation = True
+            if not is_direct_relation:
+                issue_list.append(
                     ReversedConnectionNotFeasibleError(
-                        id_,
-                        "view property",
-                        prop_.view_property,
-                        str(prop_.view),
-                        str(prop_.value_type),
-                        cast(ReverseConnectionEntity, prop_.connection).property_,
+                        view_id,
+                        "reversed connection",
+                        prop_id,
+                        f"{prop_.through.source} {prop_.through.property} is not a direct relation",
                     )
                 )
-
-            else:
                 continue
+            if not (
+                isinstance(target_property, dm.MappedPropertyApply | dm.MappedProperty)
+                and target_property.source == view_id
+            ):
+                issue_list.append(
+                    ReversedConnectionNotFeasibleError(
+                        view_id,
+                        "reversed connection",
+                        prop_id,
+                        f"{prop_.through.source} {prop_.through.property} is not pointing to {view_id}",
+                    )
+                )
+        return issue_list
 
     @staticmethod
-    def _changed_attributes_and_properties(
-        new_dumped: dict[str, Any], existing_dumped: dict[str, Any]
-    ) -> tuple[list[str], list[str]]:
-        """Helper method to find the changed attributes and properties between two containers or views."""
-        new_attributes = {key: value for key, value in new_dumped.items() if key != "properties"}
-        existing_attributes = {key: value for key, value in existing_dumped.items() if key != "properties"}
-        changed_attributes = [key for key in new_attributes if new_attributes[key] != existing_attributes.get(key)]
-        new_properties = new_dumped.get("properties", {})
-        existing_properties = existing_dumped.get("properties", {})
-        changed_properties = [prop for prop in new_properties if new_properties[prop] != existing_properties.get(prop)]
-        return changed_attributes, changed_properties
+    def _validate_schema(
+        schema: DMSSchema,
+        view_by_id: dict[dm.ViewId, dm.ViewApply | dm.View],
+        containers_by_id: dict[dm.ContainerId, dm.ContainerApply | dm.Container],
+    ) -> IssueList:
+        errors: set[NeatError] = set()
+        defined_spaces = schema.spaces.copy()
+
+        for container_id, container in schema.containers.items():
+            if container.space not in defined_spaces:
+                errors.add(ResourceNotFoundError(container.space, "space", container_id, "container"))
+            for constraint in container.constraints.values():
+                if isinstance(constraint, dm.RequiresConstraint) and constraint.require not in containers_by_id:
+                    errors.add(ResourceNotFoundError(constraint.require, "container", container_id, "container"))
+
+        for view_id, view in schema.views.items():
+            if view.space not in defined_spaces:
+                errors.add(ResourceNotFoundError(view.space, "space", view_id, "view"))
+
+            for parent in view.implements or []:
+                if parent not in view_by_id:
+                    errors.add(PropertyNotFoundError(parent, "view", "implements", view_id, "view"))
+
+            for prop_name, prop in (view.properties or {}).items():
+                if isinstance(prop, dm.MappedPropertyApply):
+                    ref_container = containers_by_id.get(prop.container)
+                    if ref_container is None:
+                        errors.add(ResourceNotFoundError(prop.container, "container", view_id, "view"))
+                    elif prop.container_property_identifier not in ref_container.properties:
+                        errors.add(
+                            PropertyNotFoundError(
+                                prop.container,
+                                "container",
+                                prop.container_property_identifier,
+                                view_id,
+                                "view",
+                            )
+                        )
+                    else:
+                        container_property = ref_container.properties[prop.container_property_identifier]
+
+                        if isinstance(container_property.type, dm.DirectRelation) and prop.source is None:
+                            warnings.warn(
+                                DirectRelationMissingSourceWarning(view_id, prop_name),
+                                stacklevel=2,
+                            )
+
+                if (
+                    isinstance(prop, dm.EdgeConnectionApply | ReverseDirectRelationApply)
+                    and prop.source not in view_by_id
+                ):
+                    errors.add(PropertyNotFoundError(prop.source, "view", prop_name, view_id, "view"))
+
+                if (
+                    isinstance(prop, dm.EdgeConnectionApply)
+                    and prop.edge_source is not None
+                    and prop.edge_source not in view_by_id
+                ):
+                    errors.add(PropertyNotFoundError(prop.edge_source, "view", prop_name, view_id, "view"))
+
+            # This allows for multiple view properties to be mapped to the same container property,
+            # as long as they have different external_id, otherwise this will lead to raising
+            # error ContainerPropertyUsedMultipleTimesError
+            property_count = Counter(
+                (prop.container, prop.container_property_identifier, view_property_identifier)
+                for view_property_identifier, prop in (view.properties or {}).items()
+                if isinstance(prop, dm.MappedPropertyApply)
+            )
+
+            for (
+                container_id,
+                container_property_identifier,
+                _,
+            ), count in property_count.items():
+                if count > 1:
+                    view_properties = [
+                        prop_name
+                        for prop_name, prop in (view.properties or {}).items()
+                        if isinstance(prop, dm.MappedPropertyApply)
+                        and (prop.container, prop.container_property_identifier)
+                        == (container_id, container_property_identifier)
+                    ]
+                    errors.add(
+                        PropertyMappingDuplicatedError(
+                            container_id,
+                            "container",
+                            container_property_identifier,
+                            frozenset({dm.PropertyId(view_id, prop_name) for prop_name in view_properties}),
+                            "view property",
+                        )
+                    )
+
+        if schema.data_model:
+            model = schema.data_model
+            if model.space not in defined_spaces:
+                errors.add(ResourceNotFoundError(model.space, "space", model.as_id(), "data model"))
+
+            view_counts: dict[dm.ViewId, int] = defaultdict(int)
+            for view_id_or_class in model.views or []:
+                view_id = view_id_or_class if isinstance(view_id_or_class, dm.ViewId) else view_id_or_class.as_id()
+                if view_id not in view_by_id:
+                    errors.add(ResourceNotFoundError(view_id, "view", model.as_id(), "data model"))
+                view_counts[view_id] += 1
+
+            for view_id, count in view_counts.items():
+                if count > 1:
+                    errors.add(
+                        ResourceDuplicatedError(
+                            view_id,
+                            "view",
+                            repr(model.as_id()),
+                        )
+                    )
+
+        return IssueList(list(errors))

cognite/neat/_rules/models/entities/_single_value.py

@@ -3,7 +3,7 @@ import sys
 from abc import ABC, abstractmethod
 from functools import total_ordering
 from types import UnionType
-from typing import Any, ClassVar, Generic, Literal, TypeVar, Union, cast, get_args, get_origin
+from typing import Any, ClassVar, Generic, Literal, TypeVar, Union, cast, get_args, get_origin, overload
 
 from cognite.client.data_classes.data_modeling import DirectRelationReference
 from cognite.client.data_classes.data_modeling.data_types import UnitReference
@@ -22,6 +22,7 @@ from pydantic import (
     model_validator,
 )
 
+from cognite.neat._issues.errors import NeatValueError
 from cognite.neat._utils.text import replace_non_alphanumeric_with_underscore
 
 if sys.version_info >= (3, 11):
@@ -56,14 +57,29 @@ class Entity(BaseModel, extra="ignore"):
     suffix: str
 
     @classmethod
-    def load(cls: "type[T_Entity]", data: Any, **defaults) -> "T_Entity | UnknownEntity":
+    @overload
+    def load(cls: "type[T_Entity]", data: Any, strict: Literal[True], **defaults) -> "T_Entity": ...
+
+    @classmethod
+    @overload
+    def load(
+        cls: "type[T_Entity]", data: Any, strict: Literal[False] = False, **defaults
+    ) -> "T_Entity | UnknownEntity": ...
+
+    @classmethod
+    def load(cls: "type[T_Entity]", data: Any, strict: bool = False, **defaults) -> "T_Entity | UnknownEntity":
         if isinstance(data, cls):
             return data
         elif isinstance(data, str) and data == str(Unknown):
+            if strict:
+                raise NeatValueError(f"Failed to load entity {data!s}")
             return UnknownEntity(prefix=Undefined, suffix=Unknown)
         if defaults and isinstance(defaults, dict):
             # This is a trick to pass in default values
-            return cls.model_validate({_PARSE: data, "defaults": defaults})
+            try:
+                return cls.model_validate({_PARSE: data, "defaults": defaults})
+            except ValueError:
+                raise
         else:
             return cls.model_validate(data)
 
@@ -291,9 +307,21 @@ class DMSEntity(Entity, Generic[T_ID], ABC):
             defaults["prefix"] = defaults.pop("space")
         return super().dump(**defaults)
 
+    @classmethod  # type: ignore[override]
+    @overload
+    def load(cls: "type[T_DMSEntity]", data: Any, strict: Literal[True], **defaults) -> "T_DMSEntity": ...
+
+    @classmethod
+    @overload
+    def load(
+        cls: "type[T_DMSEntity]", data: Any, strict: Literal[False] = False, **defaults
+    ) -> "T_DMSEntity | DMSUnknownEntity": ...
+
     @classmethod
-    def load(cls: "type[T_DMSEntity]", data: Any, **defaults) -> "T_DMSEntity | DMSUnknownEntity":  # type: ignore[override]
+    def load(cls: "type[T_DMSEntity]", data: Any, strict: bool = False, **defaults) -> "T_DMSEntity | DMSUnknownEntity":  # type: ignore[override]
         if isinstance(data, str) and data == str(Unknown):
+            if strict:
+                raise NeatValueError(f"Failed to load entity {data!s}")
             return DMSUnknownEntity.from_id(None)
         return cast(T_DMSEntity, super().load(data, **defaults))