industrial-model 0.1.9__py3-none-any.whl → 0.1.11__py3-none-any.whl

industrial_model/__init__.py

@@ -1,16 +1,21 @@
 from .config import DataModelId
 from .engines import AsyncEngine, Engine
 from .models import (
+    AggregationResult,
     InstanceId,
     PaginatedResult,
     TViewInstance,
+    TWritableViewInstance,
     ValidationMode,
     ViewInstance,
     ViewInstanceConfig,
+    WritableViewInstance,
 )
-from .statements import and_, col, not_, or_, select
+from .statements import aggregate, and_, col, not_, or_, select
 
 __all__ = [
+    "aggregate",
+    "AggregationResult",
     "and_",
     "or_",
     "col",
@@ -20,9 +25,11 @@ __all__ = [
     "InstanceId",
     "TViewInstance",
     "DataModelId",
+    "TWritableViewInstance",
     "ValidationMode",
     "Engine",
     "AsyncEngine",
     "PaginatedResult",
     "ViewInstanceConfig",
+    "WritableViewInstance",
 ]

industrial_model/cognite_adapters/__init__.py

@@ -1,3 +1,4 @@
+import logging
 from typing import Any
 
 from cognite.client import CogniteClient
@@ -9,15 +10,21 @@ from cognite.client.data_classes.data_modeling.query import (
     QueryResult as CogniteQueryResult,
 )
 
-from industrial_model.cognite_adapters.optimizer import QueryOptimizer
 from industrial_model.config import DataModelId
-from industrial_model.models import TViewInstance
-from industrial_model.statements import Statement
+from industrial_model.models import (
+    AggregationResult,
+    TViewInstance,
+    TWritableViewInstance,
+)
+from industrial_model.statements import AggregationStatement, Statement
 
+from .aggregation_mapper import AggregationMapper
+from .optimizer import QueryOptimizer
 from .query_mapper import QueryMapper
 from .query_result_mapper import (
     QueryResultMapper,
 )
+from .upsert_mapper import UpsertMapper
 from .utils import (
     append_nodes_and_edges,
     get_query_for_dependencies_pagination,
@@ -32,14 +39,12 @@ class CogniteAdapter:
     ):
         self._cognite_client = cognite_client
 
-        dm = cognite_client.data_modeling.data_models.retrieve(
-            ids=data_model_id.as_tuple(),
-            inline_views=True,
-        ).latest_version()
-        view_mapper = ViewMapper(dm.views)
+        view_mapper = ViewMapper(cognite_client, data_model_id)
+        self._optmizer = QueryOptimizer(cognite_client, data_model_id)
         self._query_mapper = QueryMapper(view_mapper)
         self._result_mapper = QueryResultMapper(view_mapper)
-        self._optmizer = QueryOptimizer(cognite_client)
+        self._upsert_mapper = UpsertMapper(view_mapper)
+        self._aggregation_mapper = AggregationMapper(view_mapper)
 
     def query(
         self, statement: Statement[TViewInstance], all_pages: bool
@@ -77,6 +82,59 @@ class CogniteAdapter:
             if not all_pages or last_page:
                 return data, next_cursor_
 
+    def aggregate(
+        self, statement: AggregationStatement[TViewInstance]
+    ) -> list[AggregationResult]:
+        query = self._aggregation_mapper.map(statement)
+
+        result = self._cognite_client.data_modeling.instances.aggregate(
+            view=query.view.as_id(),
+            aggregates=query.metric_aggregation,
+            filter=query.filters,
+            group_by=query.group_by_columns,
+            limit=query.limit,
+        )
+
+        return [
+            AggregationResult(
+                group=item.group,
+                value=item.aggregates[0].value,
+                aggregate=statement.aggregate_,
+            )
+            for item in result
+            if item.aggregates and item.aggregates[0].value is not None
+        ]
+
+    def upsert(
+        self, entries: list[TWritableViewInstance], replace: bool = False
+    ) -> None:
+        logger = logging.getLogger(__name__)
+        operation = self._upsert_mapper.map(entries)
+
+        for node_chunk in operation.chunk_nodes():
+            logger.info(
+                f"Upserting {len(node_chunk)} nodes (replace={replace})"
+            )
+            self._cognite_client.data_modeling.instances.apply(
+                nodes=node_chunk,
+                replace=replace,
+            )
+
+        for edge_chunk in operation.chunk_edges():
+            logger.info(
+                f"Upserting {len(edge_chunk)} edges (replace={replace})"
+            )
+            self._cognite_client.data_modeling.instances.apply(
+                edges=edge_chunk,
+                replace=replace,
+            )
+
+        for edges_to_remove_chunk in operation.chunk_edges_to_delete():
+            logger.info(f"Deleting {len(edges_to_remove_chunk)} edges")
+            self._cognite_client.data_modeling.instances.delete(
+                edges=[item.as_tuple() for item in edges_to_remove_chunk],
+            )
+
     def _query_dependencies_pages(
         self,
         cognite_query: CogniteQuery,

industrial_model/cognite_adapters/aggregation_mapper.py

@@ -0,0 +1,64 @@
+from dataclasses import dataclass
+
+import cognite.client.data_classes.filters as filters
+from cognite.client.data_classes.aggregations import Count, MetricAggregation
+from cognite.client.data_classes.data_modeling import (
+    View,
+)
+
+from industrial_model.models import TViewInstance
+from industrial_model.statements import AggregationStatement
+
+from .filter_mapper import (
+    FilterMapper,
+)
+from .view_mapper import ViewMapper
+
+
+@dataclass
+class AggregationQuery:
+    view: View
+    metric_aggregation: MetricAggregation
+    filters: filters.Filter | None
+    group_by_columns: list[str]
+    limit: int
+
+
+class AggregationMapper:
+    def __init__(self, view_mapper: ViewMapper):
+        self._view_mapper = view_mapper
+        self._filter_mapper = FilterMapper(view_mapper)
+
+    def map(
+        self, statement: AggregationStatement[TViewInstance]
+    ) -> AggregationQuery:
+        root_node = statement.entity.get_view_external_id()
+
+        root_view = self._view_mapper.get_view(root_node)
+
+        filters_ = (
+            filters.And(
+                *self._filter_mapper.map(statement.where_clauses, root_view)
+            )
+            if statement.where_clauses
+            else None
+        )
+
+        metric_aggregation = (
+            Count(statement.aggregation_property.property)
+            if statement.aggregate_ == "count"
+            else None
+        )
+        if metric_aggregation is None:
+            raise ValueError(
+                f"Unsupported aggregate function: {statement.aggregate_}"
+            )
+        return AggregationQuery(
+            view=root_view,
+            metric_aggregation=metric_aggregation,
+            filters=filters_,
+            group_by_columns=[
+                column.property for column in statement.group_by_columns
+            ],
+            limit=statement.limit_,
+        )

industrial_model/cognite_adapters/models.py

@@ -0,0 +1,35 @@
+from dataclasses import dataclass
+from typing import Any
+
+from cognite.client.data_classes.data_modeling import (
+    EdgeApply,
+    NodeApply,
+)
+
+from industrial_model.models.entities import EdgeContainer
+
+_PAGE_SIZE = 1000
+
+
+@dataclass
+class UpsertOperation:
+    nodes: list[NodeApply]
+    edges: list[EdgeApply]
+    edges_to_delete: list[EdgeContainer]
+
+    def chunk_nodes(self) -> list[list[NodeApply]]:
+        return self._chunk_list(self.nodes)
+
+    def chunk_edges(self) -> list[list[EdgeApply]]:
+        return self._chunk_list(self.edges)
+
+    def chunk_edges_to_delete(self) -> list[list[EdgeContainer]]:
+        return self._chunk_list(self.edges_to_delete)
+
+    def _chunk_list(self, entries: list[Any]) -> list[list[Any]]:
+        data: list[list[Any]] = []
+        for i in range(0, len(entries), _PAGE_SIZE):
+            start = i
+            end = i + _PAGE_SIZE
+            data.append(entries[start:end])
+        return data

industrial_model/cognite_adapters/optimizer.py

@@ -1,5 +1,8 @@
+from threading import Lock
+
 from cognite.client import CogniteClient
 
+from industrial_model.config import DataModelId
 from industrial_model.models import TViewInstance
 from industrial_model.statements import (
     BoolExpression,
@@ -14,11 +17,11 @@ SPACE_PROPERTY = "space"
 
 class QueryOptimizer:
     def __init__(
-        self,
-        cognite_client: CogniteClient,
+        self, cognite_client: CogniteClient, data_model_id: DataModelId
     ):
-        self._all_spaces: list[str] | None = None
+        self._all_spaces = data_model_id.instance_spaces
         self._cognite_client = cognite_client
+        self._lock = Lock()
 
     def optimize(self, statement: Statement[TViewInstance]) -> None:
         instance_spaces = statement.entity.view_config.get("instance_spaces")
@@ -58,7 +61,7 @@ class QueryOptimizer:
         return False
 
     def _find_spaces(self, instance_spaces_prefix: str) -> list[str]:
-        all_spaces = self._get_all_spaces()
+        all_spaces = self._load_spaces()
 
         return [
             space
@@ -66,12 +69,18 @@ class QueryOptimizer:
             if space.startswith(instance_spaces_prefix)
         ]
 
-    def _get_all_spaces(self) -> list[str]:
+    def _load_spaces(self) -> list[str]:
         all_spaces = self._all_spaces
-        if all_spaces is None:
+        if all_spaces:
+            return all_spaces
+
+        with self._lock:
+            if self._all_spaces:
+                return self._all_spaces
+
             all_spaces = self._cognite_client.data_modeling.spaces.list(
                 limit=-1
             ).as_ids()
 
-        self._all_spaces = all_spaces
-        return all_spaces
+            self._all_spaces = all_spaces
+            return all_spaces

industrial_model/cognite_adapters/query_result_mapper.py

@@ -1,5 +1,5 @@
 from collections import defaultdict
-from typing import Any
+from typing import Any, TypedDict
 
 from cognite.client.data_classes.data_modeling import (
     Edge,
@@ -14,10 +14,17 @@ from cognite.client.data_classes.data_modeling.views import (
 )
 
 from industrial_model.constants import EDGE_DIRECTION, EDGE_MARKER, NESTED_SEP
+from industrial_model.models import EdgeContainer
 
 from .view_mapper import ViewMapper
 
 
+class _PropertyMapping(TypedDict):
+    is_list: bool
+    nodes: dict[tuple[str, str], list[Node]]
+    edges: dict[tuple[str, str], list[Edge]]
+
+
 class QueryResultMapper:
     def __init__(self, view_mapper: ViewMapper):
         self._view_mapper = view_mapper
@@ -79,8 +86,10 @@ class QueryResultMapper:
 
             visited.add(identify)
             properties = node.properties.get(view_id, {})
+
+            edges_mapping: dict[str, list[EdgeContainer]] = {}
             node_id = get_node_id(node)
-            for mapping_key, (values, is_list) in mappings.items():
+            for mapping_key, mapping_value in mappings.items():
                 element = properties.get(mapping_key)
 
                 element_key: tuple[str, str] = (
@@ -89,15 +98,24 @@ class QueryResultMapper:
                     else (node.space, node.external_id)
                 )
 
-                entries = values.get(element_key)
-                if not entries:
+                mapping_nodes = mapping_value.get("nodes", {})
+                mapping_edges = mapping_value.get("edges", {})
+                is_list = mapping_value.get("is_list", False)
+
+                node_entries = mapping_nodes.get(element_key)
+                if not node_entries:
                     continue
 
-                entry_data = self._nodes_to_dict(entries)
+                entry_data = self._nodes_to_dict(node_entries)
                 properties[mapping_key] = (
                     entry_data if is_list else entry_data[0]
                 )
-
+                edge_entries = mapping_edges.get(element_key)
+                if edge_entries:
+                    edges_mapping[mapping_key] = self._edges_to_model(
+                        edge_entries
+                    )
+            properties["_edges"] = edges_mapping
             node.properties[view_id] = properties
 
             result[node_id].append(node)
@@ -109,19 +127,18 @@ class QueryResultMapper:
         key: str,
         view: View,
         query_result: dict[str, list[Node | Edge]],
-    ) -> dict[str, tuple[dict[tuple[str, str], list[Node]], bool]]:
-        mappings: dict[
-            str, tuple[dict[tuple[str, str], list[Node]], bool]
-        ] = {}
+    ) -> dict[str, _PropertyMapping]:
+        mappings: dict[str, _PropertyMapping] = {}
 
         for property_name, property in view.properties.items():
             property_key = f"{key}{NESTED_SEP}{property_name}"
 
-            entry: dict[tuple[str, str], list[Node]] | None = None
+            nodes: dict[tuple[str, str], list[Node]] | None = None
+            edges: dict[tuple[str, str], list[Edge]] | None = None
             is_list = False
 
             if isinstance(property, MappedProperty) and property.source:
-                entry = self._map_node_property(
+                nodes = self._map_node_property(
                     property_key,
                     self._view_mapper.get_view(property.source.external_id),
                     query_result,
@@ -131,7 +148,7 @@ class QueryResultMapper:
                 isinstance(property, SingleReverseDirectRelation)
                 and property.source
             ):
-                entry = self._map_node_property(
+                nodes = self._map_node_property(
                     property_key,
                     self._view_mapper.get_view(property.source.external_id),
                     query_result,
@@ -142,7 +159,7 @@ class QueryResultMapper:
                 isinstance(property, MultiReverseDirectRelation)
                 and property.source
             ):
-                entry = self._map_node_property(
+                nodes = self._map_node_property(
                     property_key,
                     self._view_mapper.get_view(property.source.external_id),
                     query_result,
@@ -151,7 +168,7 @@ class QueryResultMapper:
                 is_list = True
 
             elif isinstance(property, EdgeConnection) and property.source:
-                entry = self._map_edge_property(
+                nodes, edges = self._map_edge_property(
                     property_key,
                     self._view_mapper.get_view(property.source.external_id),
                     query_result,
@@ -159,8 +176,10 @@ class QueryResultMapper:
                 )
                 is_list = True
 
-            if entry:
-                mappings[property_name] = entry, is_list
+            if nodes:
+                mappings[property_name] = _PropertyMapping(
+                    is_list=is_list, nodes=nodes, edges=edges or {}
+                )
 
         return mappings
 
@@ -170,17 +189,25 @@ class QueryResultMapper:
         view: View,
         query_result: dict[str, list[Node | Edge]],
         edge_direction: EDGE_DIRECTION,
-    ) -> dict[tuple[str, str], list[Node]] | None:
+    ) -> tuple[
+        dict[tuple[str, str], list[Node]] | None,
+        dict[tuple[str, str], list[Edge]] | None,
+    ]:
         edge_key = f"{key}{NESTED_SEP}{EDGE_MARKER}"
         if key not in query_result or edge_key not in query_result:
-            return None
+            return None, None
 
         nodes = self._map_node_property(key, view, query_result)
         if not nodes:
-            return None
+            return None, None
 
         visited: set[tuple[str, str]] = set()
-        result: defaultdict[tuple[str, str], list[Node]] = defaultdict(list)
+        nodes_result: defaultdict[tuple[str, str], list[Node]] = defaultdict(
+            list
+        )
+        edges_result: defaultdict[tuple[str, str], list[Edge]] = defaultdict(
+            list
+        )
         for edge in query_result[edge_key]:
             identify = (edge.space, edge.external_id)
             if not isinstance(edge, Edge) or identify in visited:
@@ -195,15 +222,18 @@ class QueryResultMapper:
                 if edge_direction == "inwards"
                 else (edge.start_node.as_tuple(), edge.end_node.as_tuple())
             )
-
+            edges_result[entry_key].append(edge)
             if node_item := nodes.get(node_key):
-                result[entry_key].extend(node_item)
+                nodes_result[entry_key].extend(node_item)
 
-        return dict(result)
+        return dict(nodes_result), dict(edges_result)
 
     def _nodes_to_dict(self, nodes: list[Node]) -> list[dict[str, Any]]:
         return [self._node_to_dict(node) for node in nodes]
 
+    def _edges_to_model(self, edges: list[Edge]) -> list[EdgeContainer]:
+        return [EdgeContainer.model_validate(edge) for edge in edges]
+
     def _node_to_dict(self, node: Node) -> dict[str, Any]:
         entry = node.dump()
         properties: dict[str, dict[str, dict[str, Any]]] = (

industrial_model/cognite_adapters/upsert_mapper.py

@@ -0,0 +1,146 @@
+import datetime
+from typing import Any
+
+from cognite.client.data_classes.data_modeling import (
+    DirectRelationReference,
+    EdgeApply,
+    EdgeConnection,
+    MappedProperty,
+    NodeApply,
+    NodeOrEdgeData,
+)
+
+from industrial_model.cognite_adapters.models import UpsertOperation
+from industrial_model.models import (
+    EdgeContainer,
+    InstanceId,
+    TWritableViewInstance,
+)
+from industrial_model.utils import datetime_to_ms_iso_timestamp
+
+from .view_mapper import ViewMapper
+
+
+class UpsertMapper:
+    def __init__(self, view_mapper: ViewMapper):
+        self._view_mapper = view_mapper
+
+    def map(self, instances: list[TWritableViewInstance]) -> UpsertOperation:
+        nodes: dict[tuple[str, str], NodeApply] = {}
+        edges: dict[tuple[str, str], EdgeApply] = {}
+        edges_to_delete: dict[tuple[str, str], EdgeContainer] = {}
+
+        for instance in instances:
+            entry_nodes, entry_edges, entry_edges_to_delete = (
+                self._map_instance(instance)
+            )
+
+            nodes[instance.as_tuple()] = entry_nodes
+            edges.update(
+                {(item.space, item.external_id): item for item in entry_edges}
+            )
+            edges_to_delete.update(
+                {
+                    (item.space, item.external_id): item
+                    for item in entry_edges_to_delete
+                }
+            )
+
+        return UpsertOperation(
+            nodes=list(nodes.values()),
+            edges=list(edges.values()),
+            edges_to_delete=list(edges_to_delete.values()),
+        )
+
+    def _map_instance(
+        self, instance: TWritableViewInstance
+    ) -> tuple[NodeApply, list[EdgeApply], list[EdgeContainer]]:
+        view = self._view_mapper.get_view(instance.get_view_external_id())
+
+        edges: list[EdgeApply] = []
+        edges_to_delete: list[EdgeContainer] = []
+        properties: dict[str, Any] = {}
+        for property_name, property in view.properties.items():
+            property_key = instance.get_field_name(property_name)
+            if not property_key:
+                continue
+            entry = instance.__getattribute__(property_key)
+
+            if isinstance(property, MappedProperty):
+                properties[property_name] = (
+                    DirectRelationReference(
+                        space=entry.space, external_id=entry.external_id
+                    )
+                    if isinstance(entry, InstanceId)
+                    else datetime_to_ms_iso_timestamp(entry)
+                    if isinstance(entry, datetime.datetime)
+                    else entry
+                )
+            elif isinstance(property, EdgeConnection) and isinstance(
+                entry, list
+            ):
+                possible_entries = self._map_edges(instance, property, entry)
+
+                previous_edges = {
+                    item.as_tuple(): item
+                    for item in instance._edges.get(property_name, [])
+                }
+
+                new_entries = [
+                    edge
+                    for edge_id, edge in possible_entries.items()
+                    if edge_id not in previous_edges
+                ]
+                edges_to_delete.extend(
+                    [
+                        previous_edges[edge_id]
+                        for edge_id in previous_edges
+                        if edge_id not in possible_entries
+                    ]
+                )
+
+                edges.extend(new_entries)
+
+        node = NodeApply(
+            external_id=instance.external_id,
+            space=instance.space,
+            sources=[
+                NodeOrEdgeData(source=view.as_id(), properties=properties)
+            ],
+        )
+
+        return node, edges, edges_to_delete
+
+    def _map_edges(
+        self,
+        instance: TWritableViewInstance,
+        property: EdgeConnection,
+        values: list[Any],
+    ) -> dict[tuple[str, str], EdgeApply]:
+        edge_type = InstanceId.model_validate(property.type)
+
+        result: dict[tuple[str, str], EdgeApply] = {}
+        for value in values:
+            if not isinstance(value, InstanceId):
+                raise ValueError(
+                    f"""Invalid value for edge property {property.name}:
+                        Received {type(value)} | Expected: InstanceId"""
+                )
+
+            start_node, end_node = (
+                (instance, value)
+                if property.direction == "outwards"
+                else (value, instance)
+            )
+
+            edge_id = instance.edge_id_factory(value, edge_type)
+
+            result[edge_id.as_tuple()] = EdgeApply(
+                external_id=edge_id.external_id,
+                space=edge_id.space,
+                type=property.type,
+                start_node=start_node.as_tuple(),
+                end_node=end_node.as_tuple(),
+            )
+
+        return result

industrial_model/cognite_adapters/view_mapper.py

@@ -1,16 +1,50 @@
+from threading import Lock
+
+from cognite.client import CogniteClient
 from cognite.client.data_classes.data_modeling import (
     View,
 )
 
+from industrial_model.config import DataModelId
+
 
 class ViewMapper:
-    def __init__(self, views: list[View]):
-        self._views_as_dict = {view.external_id: view for view in views}
+    def __init__(
+        self, cognite_client: CogniteClient, data_model_id: DataModelId
+    ):
+        self._cognite_client = cognite_client
+        self._data_model_id = data_model_id
+        self._views_as_dict: dict[str, View] | None = None
+
+        if data_model_id.views:
+            self._views_as_dict = {
+                view.external_id: view for view in data_model_id.views
+            }
+
+        self._lock = Lock()
 
     def get_view(self, view_external_id: str) -> View:
-        if view_external_id not in self._views_as_dict:
+        views = self._load_views()
+        if view_external_id not in views:
             raise ValueError(
                 f"View {view_external_id} is not available in data model"
             )
 
-        return self._views_as_dict[view_external_id]
+        return views[view_external_id]
+
+    def _load_views(self) -> dict[str, View]:
+        if self._views_as_dict:
+            return self._views_as_dict
+
+        with self._lock:
+            if self._views_as_dict:
+                return self._views_as_dict
+
+            dm = self._cognite_client.data_modeling.data_models.retrieve(
+                ids=self._data_model_id.as_tuple(),
+                inline_views=True,
+            ).latest_version()
+
+            views = {view.external_id: view for view in dm.views}
+            self._views_as_dict = views
+            return views

industrial_model/config.py

@@ -1,10 +1,19 @@
-from industrial_model.models import RootModel
+from dataclasses import dataclass, field
 
+from cognite.client.data_classes.data_modeling import View
+from pydantic import BaseModel, ConfigDict
+
+
+@dataclass
+class DataModelId(BaseModel):
+    model_config = ConfigDict(arbitrary_types_allowed=True)
 
-class DataModelId(RootModel):
     external_id: str
     space: str
     version: str
 
+    views: list[View] | None = field(default=None)
+    instance_spaces: list[str] | None = field(default=None)
+
     def as_tuple(self) -> tuple[str, str, str]:
         return self.space, self.external_id, self.version

industrial_model/engines/async_engine.py

@@ -2,11 +2,13 @@ from cognite.client import CogniteClient
 
 from industrial_model.config import DataModelId
 from industrial_model.models import (
+    AggregationResult,
     PaginatedResult,
     TViewInstance,
+    TWritableViewInstance,
     ValidationMode,
 )
-from industrial_model.statements import Statement
+from industrial_model.statements import AggregationStatement, Statement
 from industrial_model.utils import run_async
 
 from .engine import Engine
@@ -35,3 +37,13 @@ class AsyncEngine:
         return await run_async(
             self._engine.query_all_pages, statement, validation_mode
         )
+
+    async def aggregate_async(
+        self, statement: AggregationStatement[TViewInstance]
+    ) -> list[AggregationResult]:
+        return await run_async(self._engine.aggregate, statement)
+
+    async def upsert_async(
+        self, entries: list[TWritableViewInstance], replace: bool = False
+    ) -> None:
+        return await run_async(self._engine.upsert, entries, replace)

industrial_model/engines/engine.py

@@ -5,11 +5,14 @@ from cognite.client import CogniteClient
 from industrial_model.cognite_adapters import CogniteAdapter
 from industrial_model.config import DataModelId
 from industrial_model.models import (
+    AggregationResult,
+    EdgeContainer,
     PaginatedResult,
     TViewInstance,
+    TWritableViewInstance,
     ValidationMode,
 )
-from industrial_model.statements import Statement
+from industrial_model.statements import AggregationStatement, Statement
 
 
 class Engine:
@@ -45,6 +48,19 @@ class Engine:
 
         return self._validate_data(statement.entity, data, validation_mode)
 
+    def aggregate(
+        self, statement: AggregationStatement[TViewInstance]
+    ) -> list[AggregationResult]:
+        return self._cognite_adapter.aggregate(statement)
+
+    def upsert(
+        self, entries: list[TWritableViewInstance], replace: bool = False
+    ) -> None:
+        if not entries:
+            return
+
+        return self._cognite_adapter.upsert(entries, replace)
+
     def _validate_data(
         self,
         entity: type[TViewInstance],
@@ -54,9 +70,25 @@ class Engine:
         result: list[TViewInstance] = []
         for item in data:
             try:
-                result.append(entity.model_validate(item))
+                validated_item = entity.model_validate(item)
+                self._include_edges(item, validated_item)
+                result.append(validated_item)
             except Exception:
                 if validation_mode == "ignoreOnError":
                     continue
                 raise
         return result
+
+    def _include_edges(
+        self, item: dict[str, Any], validated_item: TViewInstance
+    ) -> None:
+        if "_edges" not in item or not isinstance(item["_edges"], dict):
+            return
+        entries: dict[str, list[EdgeContainer]] = {}
+        for property_, edges in item["_edges"].items():
+            if not edges or not isinstance(edges, list):
+                continue
+
+            assert isinstance(edges[0], EdgeContainer)
+            entries[property_] = edges
+        validated_item._edges = entries

industrial_model/models/__init__.py

@@ -1,22 +1,30 @@
 from .base import RootModel
 from .entities import (
+    AggregationResult,
+    EdgeContainer,
     InstanceId,
     PaginatedResult,
     TViewInstance,
+    TWritableViewInstance,
     ValidationMode,
     ViewInstance,
     ViewInstanceConfig,
+    WritableViewInstance,
 )
 from .schemas import get_parent_and_children_nodes, get_schema_properties
 
 __all__ = [
+    "AggregationResult",
     "RootModel",
+    "EdgeContainer",
     "InstanceId",
     "TViewInstance",
+    "TWritableViewInstance",
     "ViewInstance",
     "ValidationMode",
     "PaginatedResult",
     "ViewInstanceConfig",
     "get_schema_properties",
     "get_parent_and_children_nodes",
+    "WritableViewInstance",
 ]

industrial_model/models/entities.py

@@ -1,3 +1,4 @@
+from abc import abstractmethod
 from typing import (
     Any,
     ClassVar,
@@ -7,6 +8,8 @@ from typing import (
     TypeVar,
 )
 
+from pydantic import PrivateAttr
+
 from .base import DBModelMetaclass, RootModel
 
 
@@ -29,6 +32,15 @@ class InstanceId(RootModel):
         return (self.space, self.external_id)
 
 
+class EdgeContainer(InstanceId):
+    type: InstanceId
+    start_node: InstanceId
+    end_node: InstanceId
+
+
+TInstanceId = TypeVar("TInstanceId", bound=InstanceId)
+
+
 class ViewInstanceConfig(TypedDict, total=False):
     view_external_id: str | None
     instance_spaces: list[str] | None
@@ -38,12 +50,38 @@ class ViewInstanceConfig(TypedDict, total=False):
 class ViewInstance(InstanceId, metaclass=DBModelMetaclass):
     view_config: ClassVar[ViewInstanceConfig] = ViewInstanceConfig()
 
+    _edges: dict[str, list[EdgeContainer]] = PrivateAttr(default_factory=dict)
+
     @classmethod
     def get_view_external_id(cls) -> str:
         return cls.view_config.get("view_external_id") or cls.__name__
 
+    def get_field_name(self, field_name_or_alias: str) -> str | None:
+        entry = self.__class__.model_fields.get(field_name_or_alias)
+        if entry:
+            return field_name_or_alias
+
+        for key, field_info in self.__class__.model_fields.items():
+            if field_info.alias == field_name_or_alias:
+                return key
+
+        return None
+
+
+class WritableViewInstance(ViewInstance):
+    @abstractmethod
+    def edge_id_factory(
+        self, target_node: TInstanceId, edge_type: InstanceId
+    ) -> InstanceId:
+        raise NotImplementedError(
+            "edge_id_factory method must be implemented in subclasses"
+        )
+
 
 TViewInstance = TypeVar("TViewInstance", bound=ViewInstance)
+TWritableViewInstance = TypeVar(
+    "TWritableViewInstance", bound=WritableViewInstance
+)
 
 
 class PaginatedResult(RootModel, Generic[TViewInstance]):
@@ -52,4 +90,10 @@ class PaginatedResult(RootModel, Generic[TViewInstance]):
     next_cursor: str | None
 
 
+class AggregationResult(RootModel):
+    group: dict[str, str | int | float | bool] | None
+    value: float
+    aggregate: str
+
+
 ValidationMode = Literal["raiseOnError", "ignoreOnError"]

industrial_model/statements/__init__.py

@@ -1,5 +1,5 @@
 from dataclasses import dataclass, field
-from typing import Any, Generic, Self, TypeVar
+from typing import Any, Generic, Literal, Self, TypeVar
 
 from industrial_model.constants import DEFAULT_LIMIT, SORT_DIRECTION
 
@@ -17,6 +17,10 @@ from .expressions import (
 T = TypeVar("T")
 
 
+def _create_column(property: str | Column | Any) -> Column:
+    return property if isinstance(property, Column) else Column(property)
+
+
 @dataclass
 class Statement(Generic[T]):
     entity: type[T] = field(init=True)
@@ -33,16 +37,21 @@ class Statement(Generic[T]):
             self.where_clauses.append(expression)
         return self
 
-    def asc(self, property: Any) -> Self:
-        self.sort_clauses.append((Column(property), "ascending"))
-        return self
+    def asc(self, property: str | Column | Any) -> Self:
+        return self.sort(property, "ascending")
 
-    def desc(self, property: Any) -> Self:
-        self.sort_clauses.append((Column(property), "descending"))
-        return self
+    def desc(self, property: str | Column | Any) -> Self:
+        return self.sort(property, "descending")
 
-    def sort(self, property: Any, direction: SORT_DIRECTION) -> Self:
-        self.sort_clauses.append((Column(property), direction))
+    def sort(
+        self, property: str | Column | Any, direction: SORT_DIRECTION
+    ) -> Self:
+        self.sort_clauses.append(
+            (
+                _create_column(property),
+                direction,
+            )
+        )
         return self
 
     def limit(self, limit: int) -> Self:
@@ -54,11 +63,52 @@ class Statement(Generic[T]):
         return self
 
 
+@dataclass
+class AggregationStatement(Generic[T]):
+    aggregate_: Literal["count"] = field(init=False, default="count")
+    aggregation_property: Column = field(
+        init=False, default=Column("externalId")
+    )
+
+    entity: type[T] = field(init=True)
+    group_by_columns: list[Column] = field(init=False, default_factory=list)
+    where_clauses: list[Expression] = field(init=False, default_factory=list)
+    limit_: int = field(init=False, default=-1)
+
+    def aggregate(self, aggregates: Literal["count"]) -> Self:
+        self.aggregate_ = aggregates
+        return self
+
+    def aggregate_by(self, property: str | Column | Any) -> Self:
+        self.aggregation_property = _create_column(property)
+        return self
+
+    def group_by(self, *property: str | Column | Any) -> Self:
+        self.group_by_columns.extend(_create_column(p) for p in property)
+        return self
+
+    def where(self, *expressions: bool | Expression) -> Self:
+        for expression in expressions:
+            assert isinstance(expression, Expression)
+            self.where_clauses.append(expression)
+        return self
+
+    def limit(self, limit: int) -> Self:
+        self.limit_ = limit
+        return self
+
+
 def select(entity: type[T]) -> Statement[T]:
     return Statement(entity)
 
 
+def aggregate(entity: type[T]) -> AggregationStatement[T]:
+    return AggregationStatement(entity)
+
+
 __all__ = [
+    "aggregate",
+    "AggregationStatement",
     "Statement",
     "select",
     "Column",

industrial_model/utils.py

@@ -1,14 +1,12 @@
-from collections.abc import Callable, Generator
+from collections.abc import Callable
 from datetime import datetime
 from typing import (
-    Any,
     ParamSpec,
     TypeVar,
 )
 
 from anyio import to_thread
 
-TAny = TypeVar("TAny")
 T_Retval = TypeVar("T_Retval")
 P = ParamSpec("P")
 
@@ -21,15 +19,6 @@ def datetime_to_ms_iso_timestamp(dt: datetime) -> str:
     return dt.isoformat(timespec="milliseconds")
 
 
-def chunk_list(
-    entries: list[TAny], chunk_size: int
-) -> Generator[list[TAny], Any, None]:
-    for i in range(0, len(entries), chunk_size):
-        start = i
-        end = i + chunk_size
-        yield entries[start:end]
-
-
 async def run_async(
     func: Callable[..., T_Retval],
     *args: object,

{industrial_model-0.1.9.dist-info → industrial_model-0.1.11.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: industrial-model
-Version: 0.1.9
+Version: 0.1.11
 Summary: Industrial Model ORM
 Author-email: Lucas Alves <lucasrosaalves@gmail.com>
 Classifier: Programming Language :: Python
@@ -181,6 +181,40 @@ statement = select(Person).where(
 )
 all_results = engine.query_all_pages(statement)
 
+
+# 7. Data Ingestion
+
+from industrial_model import (
+    WritableViewInstance  # necessary for data ingestion
+)
+
+
+class WritablePerson(WritableViewInstance):
+    name: str
+    lives_in: InstanceId
+    cars: list[InstanceId]
+
+    # You need to implement the end_id_factory so the model can build the edge ids automatically.
+    def edge_id_factory(
+        self, target_node: InstanceId, edge_type: InstanceId
+    ) -> InstanceId:
+        return InstanceId(
+            external_id=f"{self.external_id}-{target_node.external_id}-{edge_type.external_id}",
+            space=self.space,
+        )
+
+statement = select(WritablePerson).where(
+    (WritablePerson.external_id == "Lucas") 
+)
+
+person = engine.query_all_pages(statement)[0]
+
+person.lives_in = InstanceId(external_id="br", space="data-space")
+person.cars.clear() # Gonna remove all car edges from the person
+
+engine.upsert([person])
+
 ```
 
+
 ---

industrial_model-0.1.11.dist-info/RECORD

@@ -0,0 +1,31 @@
+industrial_model/__init__.py,sha256=13TekQQlw7euJr1pIncV1R1Z8cWtvCPPx5e6s0NY_C0,711
+industrial_model/config.py,sha256=bNEfjgD2vth4RPGKDwjGmh0edleQ95g__UJGHV4vix0,520
+industrial_model/constants.py,sha256=wtFdxkR9gojqekwXeyVCof6VUkJYuqjjgAgVVKrcxaw,450
+industrial_model/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+industrial_model/utils.py,sha256=oh4AxwxXaWgIC2uolkCbvkgo0ququHB6yAPVIXy45Ts,663
+industrial_model/cognite_adapters/__init__.py,sha256=a7xPiSkuWkVxvsmVfFD-lOzdOc-tTNo8JmbyL5tC-68,5461
+industrial_model/cognite_adapters/aggregation_mapper.py,sha256=wG8ottGilTAJcpf2kQvhvyRmDBpSKiT4esx0P5kmZBU,1845
+industrial_model/cognite_adapters/filter_mapper.py,sha256=NqH-OW7_iKFY9POCG8W3KjkwXUgrZP1d_yxDx1J0fXM,3859
+industrial_model/cognite_adapters/models.py,sha256=2j2IS01uPkQEp9WdVk8seYzEqGcDdWFnpzXhusHB2zk,945
+industrial_model/cognite_adapters/optimizer.py,sha256=BbbsA3r9ZKJt6mN1wVRZc046ov-vbml9_ESfurTNRUg,2477
+industrial_model/cognite_adapters/query_mapper.py,sha256=3fEcaLsGjLKIh-g1BbMcffQ6rp99JeCW555iJo8JW44,6260
+industrial_model/cognite_adapters/query_result_mapper.py,sha256=lsKz0wiqMOo54MVEw_YqLcHXa_HWYgAbSLnBbbdVupw,8451
+industrial_model/cognite_adapters/sort_mapper.py,sha256=RJUAYlZGXoYzK0PwX63cibRF_L-MUq9g2ZsC2EeNIF4,696
+industrial_model/cognite_adapters/upsert_mapper.py,sha256=tWEiBJQeeNz1HDu0AoBIfCw_LL156Zg4h7ORKlZ__uw,4870
+industrial_model/cognite_adapters/utils.py,sha256=rztCtS10ZTQmXfBv0nLgDiQIMWemhdSFK-SwrbRjVxM,4792
+industrial_model/cognite_adapters/view_mapper.py,sha256=qvT25gwIgPdAqZf-oQEh0GylUi5KEPwjLNmIl5ISQGA,1474
+industrial_model/engines/__init__.py,sha256=7aGHrUm2MxIq39vR8h0xu3i1zNOuT9H9U-q4lV3nErQ,102
+industrial_model/engines/async_engine.py,sha256=qILOHTBnq8DOwHoGnBQ6qrRh2zv2-0rjBd36hqWXM0w,1523
+industrial_model/engines/engine.py,sha256=fSJL0i7TwJsVP_-uNrN_WoxToCpwidIik58KREEw-oQ,3003
+industrial_model/models/__init__.py,sha256=klwPCssjGx-0rGi8aBLzP6GWO9kgR7zDAo93el9hQ1I,672
+industrial_model/models/base.py,sha256=jbiaICJ0R1mmxXtDqxxlVdq-tTX4RLdqnLTgs9HLm_4,1279
+industrial_model/models/entities.py,sha256=x4c-ch-9cRVVk039dx60kXSDFhsxQ3lIY96tJ4vPjYI,2506
+industrial_model/models/schemas.py,sha256=LKNPDnUy1jtMyOHDf28En9vThhdzOSswewIcjC_y-6U,4393
+industrial_model/queries/__init__.py,sha256=7aheTE5qs03rxWm9fmGWptbz_p9OIXXYD8if56cqs18,227
+industrial_model/queries/models.py,sha256=iiHQ7-cfg0nukEv5PoCx9QPF-w1gVSnoNbXBOK9Mzeo,1185
+industrial_model/queries/params.py,sha256=ehgCoR5n6E-tkEuoymZ2lkLcSzMaBAx_HnyJ7sWpqz0,964
+industrial_model/statements/__init__.py,sha256=Y57G1TJlYxCHg3Zz7ojQoAMVnCMaUCRRWXEm2IcbAWo,3332
+industrial_model/statements/expressions.py,sha256=Sar1cIvy3sYi7tkWJN3ylHlZ252oN2mZJpZ1TX9jN3s,4940
+industrial_model-0.1.11.dist-info/METADATA,sha256=IXdJyysssgu2JDxbSe6IN0Wg3US7Ax3XYu4NXUXU838,5698
+industrial_model-0.1.11.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+industrial_model-0.1.11.dist-info/RECORD,,

industrial_model-0.1.9.dist-info/RECORD

@@ -1,28 +0,0 @@
-industrial_model/__init__.py,sha256=qst1HvSXw7ytYTvi_hgg8989rSqKXF3n_8jv3wO73x8,525
-industrial_model/config.py,sha256=wzSKVKHIdGlxCRtu0PIeL3SLYvnQBR4Py46CcETxa8U,238
-industrial_model/constants.py,sha256=wtFdxkR9gojqekwXeyVCof6VUkJYuqjjgAgVVKrcxaw,450
-industrial_model/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-industrial_model/utils.py,sha256=OuuAkIqu-551axsZJPuqHu5d9iReGcfY9ROBFtHIwkY,933
-industrial_model/cognite_adapters/__init__.py,sha256=SjI-cO8tsZpd1uBSyP6nt3WvWbMHnWzf3v1X5oDTNvw,3436
-industrial_model/cognite_adapters/filter_mapper.py,sha256=NqH-OW7_iKFY9POCG8W3KjkwXUgrZP1d_yxDx1J0fXM,3859
-industrial_model/cognite_adapters/optimizer.py,sha256=G8I07jJ9tarE5GZXXUSpTMKUei6ptV-cudJSLsoykX4,2223
-industrial_model/cognite_adapters/query_mapper.py,sha256=3fEcaLsGjLKIh-g1BbMcffQ6rp99JeCW555iJo8JW44,6260
-industrial_model/cognite_adapters/query_result_mapper.py,sha256=cDQ2eEOqKRr4Ifowal7qvUlKSDLUj-xMFSo6qcR_UPE,7179
-industrial_model/cognite_adapters/sort_mapper.py,sha256=RJUAYlZGXoYzK0PwX63cibRF_L-MUq9g2ZsC2EeNIF4,696
-industrial_model/cognite_adapters/utils.py,sha256=rztCtS10ZTQmXfBv0nLgDiQIMWemhdSFK-SwrbRjVxM,4792
-industrial_model/cognite_adapters/view_mapper.py,sha256=lnv64KezSQTAr6XdcExa8d92GU5Ll9K9HfMcQGzhX6k,488
-industrial_model/engines/__init__.py,sha256=7aGHrUm2MxIq39vR8h0xu3i1zNOuT9H9U-q4lV3nErQ,102
-industrial_model/engines/async_engine.py,sha256=3Y4Ao14CDJJAZFbgTX9I6LOJUxfe8nwOOukGVeiw914,1070
-industrial_model/engines/engine.py,sha256=8-wfktRLZpad_V7_F1Vz7vtjE_bQhga_jXXjuwSgfrE,1897
-industrial_model/models/__init__.py,sha256=q2cI-WutiT-gmRlwC5CnacWqtQo0oSIr7Ohn9hK9jes,474
-industrial_model/models/base.py,sha256=jbiaICJ0R1mmxXtDqxxlVdq-tTX4RLdqnLTgs9HLm_4,1279
-industrial_model/models/entities.py,sha256=NrFV_a3ZP6Y4b1M2PFl_746qcCsexsyO2cuB0snZhkw,1319
-industrial_model/models/schemas.py,sha256=LKNPDnUy1jtMyOHDf28En9vThhdzOSswewIcjC_y-6U,4393
-industrial_model/queries/__init__.py,sha256=7aheTE5qs03rxWm9fmGWptbz_p9OIXXYD8if56cqs18,227
-industrial_model/queries/models.py,sha256=iiHQ7-cfg0nukEv5PoCx9QPF-w1gVSnoNbXBOK9Mzeo,1185
-industrial_model/queries/params.py,sha256=ehgCoR5n6E-tkEuoymZ2lkLcSzMaBAx_HnyJ7sWpqz0,964
-industrial_model/statements/__init__.py,sha256=9fD-qpNXIkrjoahxC_R6hS4DKSVelehimvRPKbpYfA0,1775
-industrial_model/statements/expressions.py,sha256=Sar1cIvy3sYi7tkWJN3ylHlZ252oN2mZJpZ1TX9jN3s,4940
-industrial_model-0.1.9.dist-info/METADATA,sha256=E1MZ0OJAQiRttTC-Dea6eEthmVdj_N0bLYQ7ccNSWXI,4806
-industrial_model-0.1.9.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-industrial_model-0.1.9.dist-info/RECORD,,