cognite-neat 0.126.0__py3-none-any.whl → 0.126.1__py3-none-any.whl

cognite/neat/_data_model/exporters/_table_exporter/writer.py

@@ -0,0 +1,391 @@
+import json
+from collections import defaultdict
+from dataclasses import dataclass, field
+from typing import Any, Literal
+
+from cognite.neat._data_model.importers._table_importer.data_classes import (
+    DMSContainer,
+    DMSEnum,
+    DMSNode,
+    DMSProperty,
+    DMSView,
+    MetadataValue,
+    TableDMS,
+)
+from cognite.neat._data_model.models.dms import (
+    ContainerPropertyDefinition,
+    ContainerReference,
+    ContainerRequest,
+    DataModelRequest,
+    DataType,
+    DirectNodeRelation,
+    EnumProperty,
+    ListablePropertyTypeDefinition,
+    NodeReference,
+    RequestSchema,
+    RequiresConstraintDefinition,
+    UniquenessConstraintDefinition,
+    ViewCorePropertyRequest,
+    ViewReference,
+    ViewRequest,
+    ViewRequestProperty,
+)
+from cognite.neat._data_model.models.dms._view_property import (
+    EdgeProperty,
+    MultiEdgeProperty,
+    MultiReverseDirectRelationPropertyRequest,
+    ReverseDirectRelationProperty,
+    SingleEdgeProperty,
+    SingleReverseDirectRelationPropertyRequest,
+)
+from cognite.neat._data_model.models.entities import ParsedEntity
+
+
+@dataclass
+class ViewProperties:
+    properties: list[DMSProperty] = field(default_factory=list)
+    nodes: list[DMSNode] = field(default_factory=list)
+
+
+@dataclass
+class ContainerProperties:
+    properties_by_id: dict[tuple[ContainerReference, str], dict] = field(default_factory=dict)
+    enum_collections: list[DMSEnum] = field(default_factory=list)
+
+
+class DMSTableWriter:
+    def __init__(self, default_space: str, default_version: str) -> None:
+        self.default_space = default_space
+        self.default_version = default_version
+
+    ## Main Entry Point ###
+    def write_tables(self, schema: RequestSchema) -> TableDMS:
+        metadata = self.write_metadata(schema.data_model)
+        container_properties = self.write_container_properties(schema.containers)
+        view_properties = self.write_view_properties(schema.views, container_properties)
+        views = self.write_views(schema.views, set(schema.data_model.views or []))
+        containers = self.write_containers(schema.containers)
+
+        return TableDMS(
+            metadata=metadata,
+            properties=view_properties.properties,
+            views=views,
+            containers=containers,
+            enum=container_properties.enum_collections,
+            nodes=view_properties.nodes,
+        )
+
+    ### Metadata Sheet ###
+    @staticmethod
+    def write_metadata(data_model: DataModelRequest) -> list[MetadataValue]:
+        return [
+            MetadataValue(key=key, value=value)
+            for key, value in data_model.model_dump(
+                mode="json", by_alias=True, exclude_none=True, exclude={"views"}
+            ).items()
+        ]
+
+    ### Container Properties Sheet ###
+
+    def write_containers(self, containers: list[ContainerRequest]) -> list[DMSContainer]:
+        return [
+            DMSContainer(
+                container=self._create_container_entity(container),
+                name=container.name,
+                description=container.description,
+                constraint=self._create_container_constraints(container),
+                used_for=container.used_for,
+            )
+            for container in containers
+        ]
+
+    def write_container_properties(self, containers: list[ContainerRequest]) -> ContainerProperties:
+        indices_by_container_property = self._write_container_indices(containers)
+        constraints_by_container_property = self._write_container_property_constraints(containers)
+
+        output = ContainerProperties()
+        for container in containers:
+            for prop_id, prop in container.properties.items():
+                container_property = self._write_container_property(
+                    container.as_reference(),
+                    prop_id,
+                    prop,
+                    indices_by_container_property,
+                    constraints_by_container_property,
+                )
+                output.properties_by_id[(container.as_reference(), prop_id)] = container_property
+                if isinstance(prop.type, EnumProperty):
+                    output.enum_collections.extend(
+                        self._write_enum_collection(container.as_reference(), prop_id, prop.type)
+                    )
+        return output
+
+    def _write_container_property(
+        self,
+        container_ref: ContainerReference,
+        prop_id: str,
+        prop: ContainerPropertyDefinition,
+        indices_by_container_property: dict[tuple[ContainerReference, str], list[ParsedEntity]],
+        constraints_by_container_property: dict[tuple[ContainerReference, str], list[ParsedEntity]],
+    ) -> dict[str, Any]:
+        return dict(
+            connection=self._write_container_property_connection(prop.type),
+            value_type=self._write_container_property_value_type(prop, prop_id, container_ref),
+            min_count=0 if prop.nullable else 1,
+            max_count=self._write_container_property_max_count(prop.type),
+            immutable=prop.immutable,
+            default=json.dumps(prop.default_value) if isinstance(prop.default_value, dict) else prop.default_value,
+            auto_increment=prop.auto_increment,
+            container=self._create_container_entity(container_ref),
+            container_property=prop_id,
+            container_property_name=prop.name,
+            container_property_description=prop.description,
+            index=indices_by_container_property.get((container_ref, prop_id)),
+            constraint=constraints_by_container_property.get((container_ref, prop_id)),
+        )
+
+    def _write_container_property_connection(self, dtype: DataType) -> ParsedEntity | None:
+        if not isinstance(dtype, DirectNodeRelation):
+            return None
+        properties: dict[str, str] = {}
+        if dtype.container is not None:
+            properties["container"] = str(self._create_container_entity(dtype.container))
+        return ParsedEntity("", "direct", properties=properties)
+
+    def _write_container_property_value_type(
+        self, prop: ContainerPropertyDefinition, prop_id: str, container_ref: ContainerReference
+    ) -> ParsedEntity:
+        if isinstance(prop.type, DirectNodeRelation):
+            # Will be overwritten if the view property has source set.
+            return ParsedEntity("", "#N/A", properties={})
+        elif isinstance(prop.type, EnumProperty):
+            enum_properties = {"collection": self._enum_collection_name(container_ref, prop_id)}
+            if prop.type.unknown_value is not None:
+                enum_properties["unknownValue"] = prop.type.unknown_value
+            return ParsedEntity("", "enum", properties=enum_properties)
+        elif isinstance(prop.type, ListablePropertyTypeDefinition):
+            # List and maxListSize are included in the maxCount of the property, so we exclude them here.
+            entity_properties = prop.type.model_dump(
+                mode="json", by_alias=True, exclude={"list", "maxListSize", "type"}, exclude_none=True
+            )
+            return ParsedEntity("", prop.type.type, properties=entity_properties)
+        else:
+            # Should not happen as all types are either ListablePropertyTypeDefinition or EnumProperty.
+            return ParsedEntity("", prop.type.type, properties={})
+
+    @staticmethod
+    def _write_container_property_max_count(dtype: DataType) -> int | None:
+        if isinstance(dtype, ListablePropertyTypeDefinition) and dtype.list:
+            return dtype.max_list_size
+        return 1
+
+    @staticmethod
+    def _write_container_indices(
+        containers: list[ContainerRequest],
+    ) -> dict[tuple[ContainerReference, str], list[ParsedEntity]]:
+        """Writes container indices and groups them by (container_reference, property_id)."""
+        indices_by_id: dict[tuple[ContainerReference, str], list[ParsedEntity]] = defaultdict(list)
+        for container in containers:
+            if not container.indexes:
+                continue
+            for index_id, index in container.indexes.items():
+                for order, prop_id in enumerate(index.properties, 1):
+                    entity_properties = index.model_dump(
+                        mode="json", by_alias=True, exclude={"index_type", "properties"}, exclude_none=True
+                    )
+                    if len(index.properties) > 1:
+                        entity_properties["order"] = str(order)
+                    entity = ParsedEntity(index.index_type, index_id, properties=entity_properties)
+                    indices_by_id[(container.as_reference(), prop_id)].append(entity)
+        return indices_by_id
+
+    @staticmethod
+    def _write_container_property_constraints(
+        containers: list[ContainerRequest],
+    ) -> dict[tuple[ContainerReference, str], list[ParsedEntity]]:
+        """Writes container constraints and groups them by (container_reference, property_id).
+
+        Note this only includes uniqueness constraints, the require constraints is handled
+        in the writing of the container itself.
+        """
+        constraints_by_id: dict[tuple[ContainerReference, str], list[ParsedEntity]] = defaultdict(list)
+        for container in containers:
+            if not container.constraints:
+                continue
+            for constraint_id, constraint in container.constraints.items():
+                if not isinstance(constraint, UniquenessConstraintDefinition):
+                    continue
+                for order, prop_id in enumerate(constraint.properties, 1):
+                    entity_properties = constraint.model_dump(
+                        mode="json", by_alias=True, exclude={"constraint_type", "properties"}, exclude_none=True
+                    )
+                    if len(constraint.properties) > 1:
+                        entity_properties["order"] = str(order)
+                    entity = ParsedEntity(constraint.constraint_type, constraint_id, properties=entity_properties)
+                    constraints_by_id[(container.as_reference(), prop_id)].append(entity)
+        return constraints_by_id
+
+    def _create_container_constraints(self, container: ContainerRequest) -> list[ParsedEntity] | None:
+        if not container.constraints:
+            return None
+        output: list[ParsedEntity] = []
+        for constraint_id, constraint in container.constraints.items():
+            if not isinstance(constraint, RequiresConstraintDefinition):
+                continue
+            entity_properties = {"require": str(self._create_container_entity(constraint.require))}
+            output.append(
+                ParsedEntity(prefix=constraint.constraint_type, suffix=constraint_id, properties=entity_properties)
+            )
+        return output or None
+
+    ### Enum Sheet ###
+    @staticmethod
+    def _enum_collection_name(container_ref: ContainerReference, prop_id: str) -> str:
+        return f"{container_ref.external_id}.{prop_id}"
+
+    def _write_enum_collection(
+        self, container_ref: ContainerReference, prop_id: str, enum: EnumProperty
+    ) -> list[DMSEnum]:
+        output: list[DMSEnum] = []
+        name = self._enum_collection_name(container_ref, prop_id)
+        for value_id, value in enum.values.items():
+            output.append(
+                DMSEnum(
+                    collection=name,
+                    value=value_id,
+                    name=value.name,
+                    description=value.description,
+                )
+            )
+        return output
+
+    ### View Sheet ###
+    def write_views(self, views: list[ViewRequest], model_views: set[ViewReference]) -> list[DMSView]:
+        return [
+            DMSView(
+                view=self._create_view_entity(view),
+                name=view.name,
+                description=view.description,
+                implements=[self._create_view_entity(parent) for parent in view.implements]
+                if view.implements
+                else None,
+                filter=json.dumps(view.filter) if view.filter else None,
+                in_model=None if view.as_reference() in model_views else False,
+            )
+            for view in views
+        ]
+
+    def write_view_properties(self, views: list[ViewRequest], container: ContainerProperties) -> ViewProperties:
+        output = ViewProperties()
+        for view in views:
+            if not view.properties:
+                continue
+            for prop_id, prop in view.properties.items():
+                output.properties.append(self._write_view_property(view, prop_id, prop, container))
+                if isinstance(prop, EdgeProperty):
+                    output.nodes.append(self._write_node(prop))
+        return output
+
+    def _write_view_property(
+        self, view: ViewRequest, prop_id: str, prop: ViewRequestProperty, container: ContainerProperties
+    ) -> DMSProperty:
+        container_properties: dict[str, Any] = {}
+        if isinstance(prop, ViewCorePropertyRequest):
+            identifier = (prop.container, prop.container_property_identifier)
+            if identifier in container.properties_by_id:
+                container_properties = container.properties_by_id[identifier]
+        view_properties: dict[str, Any] = dict(
+            view=self._create_view_entity(view), view_property=prop_id, name=prop.name, description=prop.description
+        )
+        if connection := self._write_view_property_connection(prop):
+            view_properties["connection"] = connection
+        if view_value_type := self._write_view_property_value_type(prop):
+            view_properties["value_type"] = view_value_type
+        view_min_count = self._write_view_property_min_count(prop)
+        if view_min_count is not None:
+            view_properties["min_count"] = view_min_count
+        view_max_count = self._write_view_property_max_count(prop)
+        if view_max_count != "container":
+            view_properties["max_count"] = view_max_count
+
+        # Overwrite container properties with view properties where relevant.
+        args = container_properties | view_properties
+        return DMSProperty(**args)
+
+    def _write_view_property_connection(self, prop: ViewRequestProperty) -> ParsedEntity | None:
+        if isinstance(prop, ViewCorePropertyRequest):
+            # Use the container definition for connection
+            return None
+        elif isinstance(prop, EdgeProperty):
+            edge_properties: dict[str, str] = {}
+            if prop.direction != "outwards":
+                edge_properties["direction"] = prop.direction
+            if prop.edge_source is not None:
+                edge_properties["edgeSource"] = str(self._create_view_entity(prop.edge_source))
+            edge_properties["type"] = str(self._create_node_entity(prop.type))
+            return ParsedEntity("", "edge", properties=edge_properties)
+        elif isinstance(prop, ReverseDirectRelationProperty):
+            return ParsedEntity("", "reverse", properties={"property": prop.through.identifier})
+        else:
+            raise ValueError(f"Unknown view property type: {type(prop)}")
+
+    def _write_view_property_value_type(self, prop: ViewRequestProperty) -> ParsedEntity | None:
+        if isinstance(prop, ViewCorePropertyRequest):
+            if prop.source:
+                return self._create_view_entity(prop.source)
+            else:
+                # Use the container definition for value type
+                return None
+        elif isinstance(prop, ReverseDirectRelationProperty | EdgeProperty):
+            return self._create_view_entity(prop.source)
+        else:
+            raise ValueError(f"Unknown view property type: {type(prop)}")
+
+    @staticmethod
+    def _write_view_property_min_count(prop: ViewRequestProperty) -> int | None:
+        if isinstance(prop, ViewCorePropertyRequest):
+            # Use the container definition for min count
+            return None
+        # Edges and reverse relations cannot be required.
+        return 0
+
+    @staticmethod
+    def _write_view_property_max_count(prop: ViewRequestProperty) -> int | None | Literal["container"]:
+        if isinstance(prop, ViewCorePropertyRequest):
+            # Use the container definition for max count
+            return "container"
+        elif isinstance(prop, SingleEdgeProperty | SingleReverseDirectRelationPropertyRequest):
+            return 1
+        elif isinstance(prop, MultiEdgeProperty | MultiReverseDirectRelationPropertyRequest):
+            return None
+        else:
+            raise ValueError(f"Unknown view property type: {type(prop)}")
+
+    ### Node Sheet ###
+
+    def _write_node(self, prop: EdgeProperty) -> DMSNode:
+        return DMSNode(node=self._create_node_entity(prop.type))
+
+    ## Entity Helpers ###
+
+    def _create_view_entity(self, view: ViewRequest | ViewReference) -> ParsedEntity:
+        prefix = view.space
+        properties = {"version": view.version}
+        if view.space == self.default_space:
+            prefix = ""
+            if view.version == self.default_version:
+                # Only use default version if space is also default.
+                properties = {}
+        return ParsedEntity(prefix=prefix, suffix=view.external_id, properties=properties)
+
+    def _create_container_entity(self, container: ContainerRequest | ContainerReference) -> ParsedEntity:
+        prefix = container.space
+        if container.space == self.default_space:
+            prefix = ""
+        return ParsedEntity(prefix=prefix, suffix=container.external_id, properties={})
+
+    def _create_node_entity(self, node: NodeReference) -> ParsedEntity:
+        prefix = node.space
+        if node.space == self.default_space:
+            prefix = ""
+        return ParsedEntity(prefix=prefix, suffix=node.external_id, properties={})

cognite/neat/_data_model/importers/__init__.py

@@ -1,4 +1,5 @@
+from ._api_importer import DMSAPIImporter
 from ._base import DMSImporter
 from ._table_importer.importer import DMSTableImporter
 
-__all__ = ["DMSImporter", "DMSTableImporter"]
+__all__ = ["DMSAPIImporter", "DMSImporter", "DMSTableImporter"]

cognite/neat/_data_model/importers/_api_importer.py

@@ -0,0 +1,88 @@
+from typing import Any
+
+from pydantic import ValidationError
+
+from cognite.neat._client import NeatClient
+from cognite.neat._data_model.importers._base import DMSImporter
+from cognite.neat._data_model.models.dms import (
+    DataModelReference,
+    RequestSchema,
+    SpaceReference,
+)
+from cognite.neat._exceptions import CDFAPIException, DataModelImportException
+from cognite.neat._issues import ModelSyntaxError
+from cognite.neat._utils.http_client import FailedRequestMessage
+from cognite.neat._utils.text import humanize_collection
+from cognite.neat._utils.validation import humanize_validation_error
+
+
+class DMSAPIImporter(DMSImporter):
+    """Imports DMS in the API format."""
+
+    def __init__(self, schema: RequestSchema | dict[str, Any]) -> None:
+        self._schema = schema
+
+    def to_data_model(self) -> RequestSchema:
+        if isinstance(self._schema, RequestSchema):
+            return self._schema
+        try:
+            return RequestSchema.model_validate(self._schema)
+        except ValidationError as e:
+            humanized_errors = humanize_validation_error(e)
+            errors = [ModelSyntaxError(message=error) for error in humanized_errors]
+            raise DataModelImportException(errors) from None
+
+    @classmethod
+    def from_cdf(cls, data_model: DataModelReference, client: NeatClient) -> "DMSAPIImporter":
+        """Create a DMSAPIImporter from a data model in CDF."""
+        data_models = client.data_models.retrieve([data_model])
+        if not data_models:
+            raise CDFAPIException(messages=[FailedRequestMessage(message=f"Data model {data_model} not found in CDF.")])
+        data_model = data_models[0]
+        views = client.views.retrieve(data_model.views or [])
+        if missing_views := set(data_model.views or []) - {view.as_reference() for view in views}:
+            raise CDFAPIException(
+                messages=[
+                    FailedRequestMessage(
+                        message=f"Views {humanize_collection(missing_views)} not found in CDF "
+                        f"for data model {data_model}."
+                    )
+                ]
+            )
+        container_ids = list({container for view in views for container in view.mapped_containers})
+        containers = client.containers.retrieve(container_ids)
+        if missing_containers := set(container_ids) - {container.as_reference() for container in containers}:
+            raise CDFAPIException(
+                messages=[
+                    FailedRequestMessage(
+                        message=f"Containers {humanize_collection(missing_containers)} not found in CDF "
+                        f"for data model {data_model}."
+                    )
+                ]
+            )
+        node_types = [nt for view in views for nt in view.node_types]
+        space_ids = list(
+            {data_model.space}
+            | {view.space for view in views}
+            | {container.space for container in containers}
+            | {nt.space for nt in node_types}
+        )
+        spaces = client.spaces.retrieve([SpaceReference(space=space_id) for space_id in space_ids])
+        if missing_spaces := set(space_ids) - {space.space for space in spaces}:
+            raise CDFAPIException(
+                messages=[
+                    FailedRequestMessage(
+                        message=f"Spaces {humanize_collection(missing_spaces)} not found in CDF "
+                        f"for data model {data_model}."
+                    )
+                ]
+            )
+        return DMSAPIImporter(
+            RequestSchema(
+                dataModel=data_model.as_request(),
+                views=[view.as_request() for view in views],
+                containers=[container.as_request() for container in containers],
+                nodeTypes=node_types,
+                spaces=[space.as_request() for space in spaces],
+            )
+        )

cognite/neat/_data_model/importers/_table_importer/data_classes.py

@@ -1,8 +1,9 @@
 from collections.abc import Mapping
-from typing import Annotated, cast
+from typing import Annotated, Literal, cast, get_args
 
-from pydantic import AliasGenerator, BaseModel, BeforeValidator, Field, model_validator
+from pydantic import AliasGenerator, BaseModel, BeforeValidator, Field, PlainSerializer, model_validator
 from pydantic.alias_generators import to_camel
+from pydantic.fields import FieldInfo
 
 from cognite.neat._data_model.models.entities import ParsedEntity, parse_entities, parse_entity
 from cognite.neat._utils.text import title_case
@@ -10,6 +11,8 @@ from cognite.neat._utils.useful_types import CellValueType
 
 
 def parse_entity_str(v: str) -> ParsedEntity:
+    if isinstance(v, ParsedEntity):
+        return v
     try:
         return parse_entity(v)
     except ValueError as e:
@@ -17,23 +20,31 @@ def parse_entity_str(v: str) -> ParsedEntity:
 
 
 def parse_entities_str(v: str) -> list[ParsedEntity] | None:
+    if isinstance(v, list) and all(isinstance(item, ParsedEntity) for item in v):
+        return v
     try:
         return parse_entities(v)
     except ValueError as e:
         raise ValueError(f"Invalid entity list syntax: {e}") from e
 
 
-Entity = Annotated[ParsedEntity, BeforeValidator(parse_entity_str, str)]
-EntityList = Annotated[list[ParsedEntity], BeforeValidator(parse_entities_str, str)]
+Entity = Annotated[ParsedEntity, BeforeValidator(parse_entity_str, str), PlainSerializer(func=str)]
+EntityList = Annotated[
+    list[ParsedEntity],
+    BeforeValidator(parse_entities_str, str),
+    PlainSerializer(func=lambda v: ",".join([str(item) for item in v])),
+]
 
 
 class TableObj(
     BaseModel,
     extra="ignore",
     alias_generator=AliasGenerator(
+        alias=to_camel,
         validation_alias=title_case,
-        serialization_alias=to_camel,
+        serialization_alias=title_case,
     ),
+    populate_by_name=True,
 ): ...
 
 
@@ -107,6 +118,35 @@ class TableDMS(TableObj):
             return {title_case(k): v for k, v in data.items()}
         return data
 
+    @classmethod
+    def get_sheet_columns(
+        cls, sheet_id: str, sheet: FieldInfo | None = None, *, column_type: Literal["all", "required"] = "required"
+    ) -> list[str]:
+        if sheet_id not in cls.model_fields.keys():
+            raise KeyError(f"Invalid field id: {sheet_id}")
+        if sheet is None:
+            sheet = cls.model_fields[sheet_id]
+        return [
+            # We know all fields has validation_alias because of the alias_generator in TableDMS
+            cast(str, sheet_field.validation_alias)
+            # All the fields in the sheet's model are lists.
+            for sheet_field in get_args(sheet.annotation)[0].model_fields.values()
+            if sheet_field.is_required() or column_type == "all"
+        ]
+
+    @classmethod
+    def get_sheet_column_by_name(
+        cls, sheet_name: str, *, column_type: Literal["all", "required"] = "required"
+    ) -> list[str]:
+        for field_id, field_ in cls.model_fields.items():
+            if cast(str, field_.validation_alias) == sheet_name:
+                return cls.get_sheet_columns(field_id, field_, column_type=column_type)
+        raise KeyError(f"Invalid field alias: {sheet_name}")
+
+    @classmethod
+    def required_sheets(cls) -> set[str]:
+        return {cast(str, field_.validation_alias) for field_ in cls.model_fields.values() if field_.is_required()}
+
 
 DMS_API_MAPPING: Mapping[str, Mapping[str, str]] = {
     "Views": {
@@ -114,7 +154,7 @@ DMS_API_MAPPING: Mapping[str, Mapping[str, str]] = {
         "externalId": "View",
         "version": "View",
         **{
-            cast(str, field_.serialization_alias): cast(str, field_.validation_alias)
+            cast(str, field_.alias): cast(str, field_.validation_alias)
             for field_id, field_ in DMSView.model_fields.items()
             if field_id != "View"
         },
@@ -123,7 +163,7 @@ DMS_API_MAPPING: Mapping[str, Mapping[str, str]] = {
         "space": "Container",
         "externalId": "Container",
         **{
-            cast(str, field_.serialization_alias): cast(str, field_.validation_alias)
+            cast(str, field_.alias): cast(str, field_.validation_alias)
             for field_id, field_ in DMSContainer.model_fields.items()
             if field_id != "Container"
         },
@@ -133,7 +173,7 @@ DMS_API_MAPPING: Mapping[str, Mapping[str, str]] = {
         "externalId": "View",
         "property": "ViewProperty",
         **{
-            cast(str, field_.serialization_alias): cast(str, field_.validation_alias)
+            cast(str, field_.alias): cast(str, field_.validation_alias)
             for field_id, field_ in DMSProperty.model_fields.items()
             if field_id not in ("View", "ViewProperty")
         },