industrial-model 0.1.6__py3-none-any.whl → 0.1.7__py3-none-any.whl

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/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/cognite_adapters/utils.py

@@ -1,11 +1,30 @@
 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.models import TViewInstance
+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"]
@@ -31,3 +50,111 @@ def get_cognite_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/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/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.dist-info → industrial_model-0.1.7.dist-info}/METADATA

@@ -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.dist-info → industrial_model-0.1.7.dist-info}/RECORD

@@ -3,26 +3,26 @@ industrial_model/config.py,sha256=wzSKVKHIdGlxCRtu0PIeL3SLYvnQBR4Py46CcETxa8U,23
 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=9N4oT7tWZx5wjVr7q64Ef5FLT75Omep48xGYv1PhxV8,1917
+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=jUreXqsaLnq7hp6T_JZQNRtL8TL8hoLSFvD4SkJ4TEE,7002
+industrial_model/cognite_adapters/query_result_mapper.py,sha256=NowXR501CywuuX5X_UdVVq7cus9JcmsnOBuTsJ774qU,7013
 industrial_model/cognite_adapters/sort_mapper.py,sha256=RJUAYlZGXoYzK0PwX63cibRF_L-MUq9g2ZsC2EeNIF4,696
-industrial_model/cognite_adapters/utils.py,sha256=nBYHK8697L2yGEAcyKjQV9NEZQmWy_aSGq_iMl7gofM,907
+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=YcNLjpYnxdU_wTpn1n0AhwlZ-XHWYqX6FzfbuevG3jo,406
+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=FzreMCR2TGX89-tMRpCNdcYqHj1u9QlocRLNwE_J2rA,3366
+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.6.dist-info/METADATA,sha256=M-BS6_DyB4pIwwqPlvBp4_7XCqFDBbdver_EXwZ8WMo,4806
-industrial_model-0.1.6.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-industrial_model-0.1.6.dist-info/RECORD,,
+industrial_model-0.1.7.dist-info/METADATA,sha256=Fk1SRRo3Yg3kqBJN2-H1ZICh1Cz00E03ZEVlAksGQFA,4806
+industrial_model-0.1.7.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+industrial_model-0.1.7.dist-info/RECORD,,