industrial-model 0.1.6__tar.gz → 0.1.7__tar.gz

{industrial_model-0.1.6 → industrial_model-0.1.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: industrial-model
-Version: 0.1.6
+Version: 0.1.7
 Summary: Industrial Model ORM
 Author-email: Lucas Alves <lucasrosaalves@gmail.com>
 Classifier: Programming Language :: Python

{industrial_model-0.1.6 → industrial_model-0.1.7}/industrial_model/cognite_adapters/__init__.py

@@ -1,6 +1,13 @@
 from typing import Any
 
 from cognite.client import CogniteClient
+from cognite.client.data_classes.data_modeling import Edge, Node
+from cognite.client.data_classes.data_modeling.query import (
+    Query as CogniteQuery,
+)
+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
@@ -11,6 +18,11 @@ from .query_mapper import QueryMapper
 from .query_result_mapper import (
     QueryResultMapper,
 )
+from .utils import (
+    append_nodes_and_edges,
+    get_query_for_dependencies_pagination,
+    map_nodes_and_edges,
+)
 from .view_mapper import ViewMapper
 
 
@@ -42,10 +54,20 @@ class CogniteAdapter:
                 cognite_query
             )
 
-            page_result, next_cursor = self._result_mapper.map_nodes(
-                view_external_id, query_result
+            dependencies_data = self._query_dependencies_pages(
+                cognite_query, query_result, view_external_id
+            )
+
+            query_result_data = append_nodes_and_edges(
+                map_nodes_and_edges(query_result, cognite_query),
+                dependencies_data,
             )
 
+            page_result = self._result_mapper.map_nodes(
+                view_external_id,
+                query_result_data,
+            )
+            next_cursor = query_result.cursors.get(view_external_id)
             data.extend(page_result)
 
             last_page = len(page_result) < statement.limit_ or not next_cursor
@@ -53,4 +75,27 @@ class CogniteAdapter:
             cognite_query.cursors = {view_external_id: next_cursor_}
 
             if not all_pages or last_page:
-                return data, next_cursor
+                return data, next_cursor_
+
+    def _query_dependencies_pages(
+        self,
+        cognite_query: CogniteQuery,
+        query_result: CogniteQueryResult,
+        view_external_id: str,
+    ) -> dict[str, list[Node | Edge]] | None:
+        new_query = get_query_for_dependencies_pagination(
+            cognite_query, query_result, view_external_id
+        )
+        if not new_query:
+            return None
+
+        new_query_result = self._cognite_client.data_modeling.instances.query(
+            new_query
+        )
+
+        result = map_nodes_and_edges(new_query_result, new_query)
+
+        nested_results = self._query_dependencies_pages(
+            new_query, new_query_result, view_external_id
+        )
+        return append_nodes_and_edges(result, nested_results)

{industrial_model-0.1.6 → industrial_model-0.1.7}/industrial_model/cognite_adapters/query_result_mapper.py

@@ -8,9 +8,6 @@ from cognite.client.data_classes.data_modeling import (
     Node,
     View,
 )
-from cognite.client.data_classes.data_modeling.query import (
-    QueryResult,
-)
 from cognite.client.data_classes.data_modeling.views import (
     MultiReverseDirectRelation,
     SingleReverseDirectRelation,
@@ -26,8 +23,8 @@ class QueryResultMapper:
         self._view_mapper = view_mapper
 
     def map_nodes(
-        self, root_node: str, query_result: QueryResult
-    ) -> tuple[list[dict[str, Any]], str | None]:
+        self, root_node: str, query_result: dict[str, list[Node | Edge]]
+    ) -> list[dict[str, Any]]:
         if root_node not in query_result:
             raise ValueError(
                 f"{root_node} is not available in the query result"
@@ -37,7 +34,7 @@ class QueryResultMapper:
 
         values = self._map_node_property(root_node, root_view, query_result)
         if not values:
-            return [], None
+            return []
 
         data = [
             node
@@ -45,13 +42,13 @@ class QueryResultMapper:
             for node in self._nodes_to_dict(nodes)
         ]
 
-        return data, query_result.cursors.get(root_node)
+        return data
 
     def _map_node_property(
         self,
         key: str,
         view: View,
-        query_result: QueryResult,
+        query_result: dict[str, list[Node | Edge]],
         result_property_key: str | None = None,
     ) -> dict[tuple[str, str], list[Node]] | None:
         if key not in query_result:
@@ -61,12 +58,6 @@ class QueryResultMapper:
 
         view_id = view.as_id()
 
-        nodes = (
-            node
-            for node in query_result.get_nodes(key)
-            if isinstance(node, Node)
-        )
-
         def get_node_id(node: Node) -> tuple[str, str]:
             if not result_property_key:
                 return (node.space, node.external_id)
@@ -79,8 +70,13 @@ class QueryResultMapper:
 
             return entry.get("space", ""), entry.get("externalId", "")
 
+        nodes = {
+            (node.space, node.external_id): node
+            for node in query_result[key]
+            if isinstance(node, Node)
+        }
         result: defaultdict[tuple[str, str], list[Node]] = defaultdict(list)
-        for node in nodes:
+        for node in nodes.values():
             properties = node.properties.get(view_id, {})
             node_id = get_node_id(node)
             for mapping_key, (values, is_list) in mappings.items():
@@ -108,7 +104,10 @@ class QueryResultMapper:
         return dict(result)
 
     def _get_property_mappings(
-        self, key: str, view: View, query_result: QueryResult
+        self,
+        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]
@@ -168,7 +167,7 @@ class QueryResultMapper:
         self,
         key: str,
         view: View,
-        query_result: QueryResult,
+        query_result: dict[str, list[Node | Edge]],
         edge_direction: EDGE_DIRECTION,
     ) -> dict[tuple[str, str], list[Node]] | None:
         edge_key = f"{key}{NESTED_SEP}{EDGE_MARKER}"
@@ -181,12 +180,12 @@ class QueryResultMapper:
 
         result: defaultdict[tuple[str, str], list[Node]] = defaultdict(list)
 
-        edges = (
-            edge
-            for edge in query_result.get_edges(edge_key)
+        edges = {
+            (edge.space, edge.external_id): edge
+            for edge in query_result[edge_key]
             if isinstance(edge, Edge)
-        )
-        for edge in edges:
+        }
+        for edge in edges.values():
             entry_key, node_key = (
                 (
                     edge.end_node.as_tuple(),

industrial_model-0.1.7/industrial_model/cognite_adapters/utils.py

@@ -0,0 +1,160 @@
+from typing import Literal
+
+from cognite.client.data_classes.data_modeling import (
+    Edge,
+    Node,
+    NodeListWithCursor,
+    View,
+    ViewId,
+)
+from cognite.client.data_classes.data_modeling.query import (
+    Query as CogniteQuery,
+)
+from cognite.client.data_classes.data_modeling.query import (
+    QueryResult as CogniteQueryResult,
+)
+from cognite.client.data_classes.data_modeling.query import (
+    ResultSetExpression,
+)
+from cognite.client.data_classes.data_modeling.query import (
+    Select as CogniteSelect,
+)
+
+from industrial_model.constants import MAX_LIMIT
+from industrial_model.models import (
+    TViewInstance,
+    get_parent_and_children_nodes,
+)
+
+NODE_PROPERTIES = {"externalId", "space", "createdTime", "deletedTime"}
+INTANCE_TYPE = Literal["node", "edge"]
+
+
+def get_property_ref(
+    property: str, view: View | ViewId, instance_type: INTANCE_TYPE = "node"
+) -> tuple[str, str, str] | tuple[str, str]:
+    return (
+        view.as_property_ref(property)
+        if property not in NODE_PROPERTIES
+        else (instance_type, property)
+    )
+
+
+def get_cognite_instance_ids(
+    instance_ids: list[TViewInstance],
+) -> list[dict[str, str]]:
+    return [
+        get_cognite_instance_id(instance_id) for instance_id in instance_ids
+    ]
+
+
+def get_cognite_instance_id(instance_id: TViewInstance) -> dict[str, str]:
+    return {"space": instance_id.space, "externalId": instance_id.external_id}
+
+
+def get_query_for_dependencies_pagination(
+    query: CogniteQuery,
+    query_result: CogniteQueryResult,
+    view_external_id: str,
+) -> CogniteQuery | None:
+    nodes_parent, nodes_children = get_parent_and_children_nodes(
+        set(query_result.cursors.keys())
+    )
+
+    leaf_cursors = _get_leaf_cursors(
+        query_result, view_external_id, nodes_parent, nodes_children
+    )
+
+    if not leaf_cursors:
+        return None
+
+    return _create_query(query, nodes_parent, nodes_children, leaf_cursors)
+
+
+def map_nodes_and_edges(
+    query_result: CogniteQueryResult, query: CogniteQuery
+) -> dict[str, list[Node | Edge]]:
+    result_schema = query.instance_type_by_result_expression()
+
+    return {
+        key: query_result.get_nodes(key).data
+        if entity_type is NodeListWithCursor
+        else query_result.get_edges(key).data
+        for key, entity_type in result_schema.items()
+        if entity_type is not None
+    }
+
+
+def append_nodes_and_edges(
+    initial_dataset: dict[str, list[Node | Edge]],
+    additional_dataset: dict[str, list[Node | Edge]] | None,
+) -> dict[str, list[Node | Edge]]:
+    if not additional_dataset:
+        return initial_dataset
+    for key, additional_data in additional_dataset.items():
+        if key not in initial_dataset:
+            initial_dataset[key] = []
+        initial_dataset[key].extend(additional_data)
+    return initial_dataset
+
+
+def _create_query(
+    previous_query: CogniteQuery,
+    nodes_parent: dict[str, set[str]],
+    nodes_children: dict[str, set[str]],
+    leaf_cursors: dict[str, str],
+) -> CogniteQuery:
+    with_: dict[str, ResultSetExpression] = {}
+    select_: dict[str, CogniteSelect] = {}
+    final_cursors: dict[str, str | None] = {}
+
+    for cursor_key, cursor_value in leaf_cursors.items():
+        children = nodes_children[cursor_key]
+        parent = nodes_parent[cursor_key]
+
+        valid_keys = parent.union(children)
+        valid_keys.add(cursor_key)
+
+        with_.update(
+            {k: v for k, v in previous_query.with_.items() if k in valid_keys}
+        )
+        select_.update(
+            {k: v for k, v in previous_query.select.items() if k in valid_keys}
+        )
+
+        final_cursors.update(
+            {k: v for k, v in previous_query.cursors.items() if k in parent}
+        )
+        final_cursors[cursor_key] = cursor_value
+
+    return CogniteQuery(with_=with_, select=select_, cursors=final_cursors)
+
+
+def _get_leaf_cursors(
+    query_result: CogniteQueryResult,
+    view_external_id: str,
+    nodes_parent: dict[str, set[str]],
+    nodes_children: dict[str, set[str]],
+) -> dict[str, str]:
+    target_cursors: dict[str, str] = {}
+    for cursor_key, cursor_value in query_result.cursors.items():
+        if (
+            cursor_key == view_external_id
+            or not cursor_value
+            or len(query_result[cursor_key]) != MAX_LIMIT
+        ):
+            continue
+
+        children = nodes_children[cursor_key]
+
+        target_cursors_keys = target_cursors.keys()
+        if len(children.intersection(target_cursors_keys)) > 0:
+            continue
+
+        parent = nodes_parent[cursor_key]
+        cursor_to_remove = parent.intersection(target_cursors_keys)
+        for cursor_key_to_remove in cursor_to_remove:
+            target_cursors.pop(cursor_key_to_remove)
+
+        target_cursors[cursor_key] = cursor_value
+    return target_cursors

{industrial_model-0.1.6 → industrial_model-0.1.7}/industrial_model/models/__init__.py

@@ -7,7 +7,7 @@ from .entities import (
     ViewInstance,
     ViewInstanceConfig,
 )
-from .schemas import get_schema_properties
+from .schemas import get_parent_and_children_nodes, get_schema_properties
 
 __all__ = [
     "RootModel",
@@ -18,4 +18,5 @@ __all__ = [
     "PaginatedResult",
     "ViewInstanceConfig",
     "get_schema_properties",
+    "get_parent_and_children_nodes",
 ]

{industrial_model-0.1.6 → industrial_model-0.1.7}/industrial_model/models/schemas.py

@@ -12,6 +12,8 @@ from typing import (
 
 from pydantic import BaseModel
 
+from industrial_model.constants import EDGE_MARKER, NESTED_SEP
+
 TBaseModel = TypeVar("TBaseModel", bound=BaseModel)
 
 
@@ -28,6 +30,32 @@ def get_schema_properties(
     return [f"{prefix}{nested_separator}{key}" for key in keys]
 
 
+def get_parent_and_children_nodes(
+    keys: set[str],
+) -> tuple[dict[str, set[str]], dict[str, set[str]]]:
+    nodes_parent: dict[str, set[str]] = {}
+    nodes_children: defaultdict[str, set[str]] = defaultdict(set)
+    for key in keys:
+        key_parts = key.split(NESTED_SEP)
+        parent_paths = {
+            NESTED_SEP.join(key_parts[:i])
+            for i in range(len(key_parts) - 1, 0, -1)
+        }
+
+        valid_paths: set[str] = set()
+        for parent_path in parent_paths:
+            parent_path_edge_marker = parent_path + NESTED_SEP + EDGE_MARKER
+            if parent_path_edge_marker in keys:
+                valid_paths.add(parent_path_edge_marker)
+                nodes_children[parent_path_edge_marker].add(key)
+            if parent_path in keys:
+                valid_paths.add(parent_path)
+                nodes_children[parent_path].add(key)
+
+        nodes_parent[key] = valid_paths
+    return nodes_parent, dict(nodes_children)
+
+
 def _get_type_properties(
     cls: type[BaseModel], visited_count: defaultdict[type, int] | None = None
 ) -> dict[str, Any] | None:

{industrial_model-0.1.6 → industrial_model-0.1.7}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "industrial-model"
-version = "0.1.6"
+version = "0.1.7"
 description = "Industrial Model ORM"
 readme = "README.md"
 requires-python = ">=3.11"

{industrial_model-0.1.6 → industrial_model-0.1.7}/uv.lock

@@ -215,7 +215,7 @@ wheels = [
 
 [[package]]
 name = "industrial-model"
-version = "0.1.6"
+version = "0.1.7"
 source = { editable = "." }
 dependencies = [
     { name = "anyio" },

industrial_model-0.1.6/industrial_model/cognite_adapters/utils.py

@@ -1,33 +0,0 @@
-from typing import Literal
-
-from cognite.client.data_classes.data_modeling import (
-    View,
-    ViewId,
-)
-
-from industrial_model.models import TViewInstance
-
-NODE_PROPERTIES = {"externalId", "space", "createdTime", "deletedTime"}
-INTANCE_TYPE = Literal["node", "edge"]
-
-
-def get_property_ref(
-    property: str, view: View | ViewId, instance_type: INTANCE_TYPE = "node"
-) -> tuple[str, str, str] | tuple[str, str]:
-    return (
-        view.as_property_ref(property)
-        if property not in NODE_PROPERTIES
-        else (instance_type, property)
-    )
-
-
-def get_cognite_instance_ids(
-    instance_ids: list[TViewInstance],
-) -> list[dict[str, str]]:
-    return [
-        get_cognite_instance_id(instance_id) for instance_id in instance_ids
-    ]
-
-
-def get_cognite_instance_id(instance_id: TViewInstance) -> dict[str, str]:
-    return {"space": instance_id.space, "externalId": instance_id.external_id}