industrial-model 0.1.9__tar.gz → 0.1.11__tar.gz

{industrial_model-0.1.9 → industrial_model-0.1.11}/PKG-INFO

@@ -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.9 → industrial_model-0.1.11}/README.md

@@ -161,6 +161,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.9 → industrial_model-0.1.11}/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-0.1.9 → industrial_model-0.1.11}/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-0.1.11/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-0.1.11/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-0.1.9 → industrial_model-0.1.11}/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-0.1.9 → industrial_model-0.1.11}/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]]] = (