industrial-model 0.1.0__py3-none-any.whl

industrial_model/__init__.py

@@ -0,0 +1,27 @@
+from .config import DataModelId
+from .engines import AsyncEngine, Engine
+from .models import (
+    InstanceId,
+    PaginatedResult,
+    TViewInstance,
+    ValidationMode,
+    ViewInstance,
+    ViewInstanceConfig,
+)
+from .statements import and_, col, or_, select
+
+__all__ = [
+    "and_",
+    "or_",
+    "col",
+    "select",
+    "ViewInstance",
+    "InstanceId",
+    "TViewInstance",
+    "DataModelId",
+    "ValidationMode",
+    "Engine",
+    "AsyncEngine",
+    "PaginatedResult",
+    "ViewInstanceConfig",
+]

industrial_model/cognite_adapters/__init__.py

@@ -0,0 +1,56 @@
+from typing import Any
+
+from cognite.client import CogniteClient
+
+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 .query_mapper import QueryMapper
+from .query_result_mapper import (
+    QueryResultMapper,
+)
+from .view_mapper import ViewMapper
+
+
+class CogniteAdapter:
+    def __init__(
+        self, cognite_client: CogniteClient, data_model_id: DataModelId
+    ):
+        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)
+        self._query_mapper = QueryMapper(view_mapper)
+        self._result_mapper = QueryResultMapper(view_mapper)
+        self._optmizer = QueryOptimizer(cognite_client)
+
+    def query(
+        self, statement: Statement[TViewInstance], all_pages: bool
+    ) -> tuple[list[dict[str, Any]], str | None]:
+        self._optmizer.optimize(statement)
+        cognite_query = self._query_mapper.map(statement)
+        view_external_id = statement.entity.get_view_external_id()
+
+        data: list[dict[str, Any]] = []
+        while True:
+            query_result = self._cognite_client.data_modeling.instances.query(
+                cognite_query
+            )
+
+            page_result, next_cursor = self._result_mapper.map_nodes(
+                view_external_id, query_result
+            )
+
+            data.extend(page_result)
+
+            last_page = len(page_result) < statement.limit_ or not next_cursor
+            next_cursor_ = None if last_page else next_cursor
+            cognite_query.cursors = {view_external_id: next_cursor_}
+
+            if not all_pages or last_page:
+                return data, next_cursor

industrial_model/cognite_adapters/filter_mapper.py

@@ -0,0 +1,99 @@
+from datetime import datetime
+
+import cognite.client.data_classes.filters as cdf_filters
+from cognite.client.data_classes.data_modeling import MappedProperty, View
+
+from industrial_model.cognite_adapters.utils import get_property_ref
+from industrial_model.statements import (
+    BoolExpression,
+    Expression,
+    LeafExpression,
+)
+from industrial_model.utils import datetime_to_ms_iso_timestamp
+
+from .view_mapper import ViewMapper
+
+
+class FilterMapper:
+    def __init__(self, view_mapper: ViewMapper):
+        self._view_mapper = view_mapper
+
+    def map(
+        self, expressions: list[Expression], root_view: View
+    ) -> list[cdf_filters.Filter]:
+        result: list[cdf_filters.Filter] = []
+        for expression in expressions:
+            if isinstance(expression, BoolExpression):
+                result.append(self.to_cdf_filter_bool(expression, root_view))
+            elif isinstance(expression, LeafExpression):
+                result.append(self.to_cdf_filter_leaf(expression, root_view))
+            else:
+                cls_name = expression.__class__.__name__
+                raise ValueError(f"Expression not implemented {cls_name}")
+        return result
+
+    def to_cdf_filter_bool(
+        self, expression: BoolExpression, root_view: View
+    ) -> cdf_filters.Filter:
+        arguments = self.map(expression.filters, root_view)
+
+        if expression.operator == "and":
+            return cdf_filters.And(*arguments)
+        elif expression.operator == "or":
+            return cdf_filters.Or(*arguments)
+        elif expression.operator == "not":
+            return cdf_filters.Not(*arguments)
+
+        raise NotImplementedError(f"Operator {self.operator} not implemented")
+
+    def to_cdf_filter_leaf(
+        self,
+        expression: LeafExpression,
+        root_view: View,
+    ) -> cdf_filters.Filter:
+        property_ref = get_property_ref(expression.property, root_view)
+
+        value_ = expression.value
+        if isinstance(value_, datetime):
+            value_ = datetime_to_ms_iso_timestamp(value_)
+
+        if expression.operator == "==":
+            return cdf_filters.Equals(property_ref, value_)
+        elif expression.operator == "in":
+            return cdf_filters.In(property_ref, value_)
+        elif expression.operator == ">":
+            return cdf_filters.Range(property_ref, gt=value_)
+        elif expression.operator == ">=":
+            return cdf_filters.Range(property_ref, gte=value_)
+        elif expression.operator == "<":
+            return cdf_filters.Range(property_ref, lt=value_)
+        elif expression.operator == "<=":
+            return cdf_filters.Range(property_ref, lte=value_)
+        elif expression.operator == "nested":
+            target_view = self._get_nested_target_view(
+                expression.property, root_view
+            )
+
+            assert isinstance(value_, Expression)
+
+            return cdf_filters.Nested(
+                property_ref,
+                self.map([value_], target_view)[0],
+            )
+        elif expression.operator == "exists":
+            return cdf_filters.Exists(property_ref)
+        elif expression.operator == "prefix":
+            return cdf_filters.Prefix(property_ref, value_)
+        elif expression.operator == "containsAll":
+            return cdf_filters.ContainsAll(property_ref, value_)
+        elif expression.operator == "containsAny":
+            return cdf_filters.ContainsAny(property_ref, value_)
+        raise NotImplementedError(
+            f"Operator {expression.operator} not implemented"
+        )
+
+    def _get_nested_target_view(self, property: str, root_view: View) -> View:
+        view_definiton = root_view.properties[property]
+        assert isinstance(view_definiton, MappedProperty)
+        assert view_definiton.source
+        return self._view_mapper.get_view(view_definiton.source.external_id)

industrial_model/cognite_adapters/optimizer.py

@@ -0,0 +1,74 @@
+from cognite.client import CogniteClient
+
+from industrial_model.models import TViewInstance
+from industrial_model.statements import (
+    BoolExpression,
+    Expression,
+    LeafExpression,
+    Statement,
+    col,
+)
+
+SPACE_PROPERTY = "space"
+
+
+class QueryOptimizer:
+    def __init__(
+        self,
+        cognite_client: CogniteClient,
+    ):
+        self._all_spaces: list[str] | None = None
+        self._cognite_client = cognite_client
+
+    def optimize(self, statement: Statement[TViewInstance]) -> None:
+        instance_spaces = statement.entity.view_config.get("instance_spaces")
+        instance_spaces_prefix = statement.entity.view_config.get(
+            "instance_spaces_prefix"
+        )
+
+        if not instance_spaces and not instance_spaces_prefix:
+            return
+
+        if self._has_space_filter(statement.where_clauses):
+            return
+
+        filter_spaces = (
+            self._find_spaces(instance_spaces_prefix)
+            if instance_spaces_prefix
+            else []
+        )
+        if instance_spaces:
+            filter_spaces.extend(instance_spaces)
+
+        if filter_spaces:
+            statement.where(col(SPACE_PROPERTY).in_(filter_spaces))
+
+    def _has_space_filter(self, where_clauses: list[Expression]) -> bool:
+        for where_clause in where_clauses:
+            if isinstance(
+                where_clause, BoolExpression
+            ) and self._has_space_filter(where_clause.filters):
+                return True
+            elif isinstance(where_clause, LeafExpression):
+                return where_clause.property == SPACE_PROPERTY
+
+        return False
+
+    def _find_spaces(self, instance_spaces_prefix: str) -> list[str]:
+        all_spaces = self._get_all_spaces()
+
+        return [
+            space
+            for space in all_spaces
+            if space.startswith(instance_spaces_prefix)
+        ]
+
+    def _get_all_spaces(self) -> list[str]:
+        all_spaces = self._all_spaces
+        if all_spaces is None:
+            all_spaces = self._cognite_client.data_modeling.spaces.list(
+                limit=-1
+            ).as_ids()
+
+        self._all_spaces = all_spaces
+        return all_spaces

industrial_model/cognite_adapters/query_mapper.py

@@ -0,0 +1,181 @@
+import cognite.client.data_classes.filters as filters
+from cognite.client.data_classes.data_modeling import (
+    EdgeConnection,
+    MappedProperty,
+    View,
+    ViewId,
+)
+from cognite.client.data_classes.data_modeling.query import (
+    EdgeResultSetExpression,
+    NodeResultSetExpression,
+    ResultSetExpression,
+    Select,
+    SourceSelector,
+)
+from cognite.client.data_classes.data_modeling.query import (
+    Query as CogniteQuery,
+)
+from cognite.client.data_classes.data_modeling.views import (
+    MultiReverseDirectRelation,
+    SingleReverseDirectRelation,
+)
+
+from industrial_model.constants import EDGE_MARKER, MAX_LIMIT, NESTED_SEP
+from industrial_model.models import TViewInstance
+from industrial_model.statements import Statement
+
+from .filter_mapper import (
+    FilterMapper,
+)
+from .schemas import get_schema_properties
+from .sort_mapper import SortMapper
+from .view_mapper import ViewMapper
+
+
+class QueryMapper:
+    def __init__(self, view_mapper: ViewMapper):
+        self._view_mapper = view_mapper
+        self._filter_mapper = FilterMapper(view_mapper)
+        self._sort_mapper = SortMapper()
+
+    def map(self, statement: Statement[TViewInstance]) -> CogniteQuery:
+        root_node = statement.entity.get_view_external_id()
+
+        root_view = self._view_mapper.get_view(root_node)
+        root_view_id = root_view.as_id()
+
+        filters_: list[filters.Filter] = [
+            filters.HasData(views=[root_view_id])
+        ]
+
+        filters_.extend(
+            self._filter_mapper.map(statement.where_clauses, root_view)
+        )
+
+        with_: dict[str, ResultSetExpression] = {
+            root_node: NodeResultSetExpression(
+                filter=filters.And(*filters_),
+                sort=self._sort_mapper.map(statement.sort_clauses, root_view),
+                limit=statement.limit_,
+            )
+        }
+        select_: dict[str, Select] = {}
+
+        relations = get_schema_properties(
+            statement.entity, NESTED_SEP, root_node
+        )
+
+        properties = self._include_statements(
+            root_node, root_view, relations, with_, select_
+        )
+
+        select_[root_node] = self._get_select(root_view_id, properties)
+
+        return CogniteQuery(
+            with_=with_, select=select_, cursors={root_node: statement.cursor_}
+        )
+
+    def _get_select(self, view_id: ViewId, properties: list[str]) -> Select:
+        return (
+            Select(
+                sources=[SourceSelector(source=view_id, properties=properties)]
+            )
+            if properties
+            else Select()
+        )
+
+    def _include_statements(
+        self,
+        key: str,
+        view: View,
+        relations_to_include: list[str] | None,
+        with_: dict[str, ResultSetExpression],
+        select_: dict[str, Select],
+    ) -> list[str]:
+        if not relations_to_include:
+            return []
+
+        select_properties: list[str] = []
+        for property_name, property in view.properties.items():
+            property_key = f"{key}{NESTED_SEP}{property_name}"
+            if property_key not in relations_to_include:
+                continue
+
+            if isinstance(property, MappedProperty) and not property.source:
+                select_properties.append(property_name)
+            elif isinstance(property, MappedProperty) and property.source:
+                select_properties.append(property_name)
+
+                props = self._include_statements(
+                    property_key,
+                    self._view_mapper.get_view(property.source.external_id),
+                    relations_to_include,
+                    with_,
+                    select_,
+                )
+                if props:
+                    with_[property_key] = NodeResultSetExpression(
+                        from_=key,
+                        through=view.as_property_ref(property_name),
+                        limit=MAX_LIMIT,
+                    )
+                    select_[property_key] = self._get_select(
+                        property.source, props
+                    )
+
+            elif (
+                isinstance(property, MultiReverseDirectRelation)
+                or isinstance(property, SingleReverseDirectRelation)
+                and property.source
+            ):
+                props = self._include_statements(
+                    property_key,
+                    self._view_mapper.get_view(property.source.external_id),
+                    relations_to_include,
+                    with_,
+                    select_,
+                )
+                if not props:
+                    with_[property_key] = NodeResultSetExpression(
+                        from_=key,
+                        direction="inwards",
+                        through=property.source.as_property_ref(
+                            property.through.property
+                        ),
+                        limit=MAX_LIMIT,
+                    )
+
+                    select_[property_key] = self._get_select(
+                        property.source, props
+                    )
+            elif isinstance(property, EdgeConnection) and property.source:
+                edge_property_key = f"{property_key}{NESTED_SEP}{EDGE_MARKER}"
+
+                with_[edge_property_key] = EdgeResultSetExpression(
+                    from_=key,
+                    max_distance=1,
+                    filter=filters.Equals(
+                        ["edge", "type"], property.type.dump()
+                    ),
+                    direction=property.direction,
+                    limit=MAX_LIMIT,
+                )
+                with_[property_key] = NodeResultSetExpression(
+                    from_=edge_property_key,
+                    limit=MAX_LIMIT,
+                )
+
+                select_[edge_property_key] = Select()
+
+                props = self._include_statements(
+                    property_key,
+                    self._view_mapper.get_view(property.source.external_id),
+                    relations_to_include,
+                    with_,
+                    select_,
+                )
+                select_[property_key] = self._get_select(
+                    property.source, props
+                )
+
+        return select_properties

industrial_model/cognite_adapters/query_result_mapper.py

@@ -0,0 +1,216 @@
+from collections import defaultdict
+from typing import Any
+
+from cognite.client.data_classes.data_modeling import (
+    Edge,
+    EdgeConnection,
+    MappedProperty,
+    Node,
+    View,
+)
+from cognite.client.data_classes.data_modeling.query import (
+    QueryResult,
+)
+from cognite.client.data_classes.data_modeling.views import (
+    MultiReverseDirectRelation,
+    SingleReverseDirectRelation,
+)
+
+from industrial_model.constants import EDGE_DIRECTION, EDGE_MARKER, NESTED_SEP
+
+from .view_mapper import ViewMapper
+
+
+class QueryResultMapper:
+    def __init__(self, view_mapper: ViewMapper):
+        self._view_mapper = view_mapper
+
+    def map_nodes(
+        self, root_node: str, query_result: QueryResult
+    ) -> tuple[list[dict[str, Any]], str | None]:
+        if root_node not in query_result:
+            raise ValueError(
+                f"{root_node} is not available in the query result"
+            )
+
+        root_view = self._view_mapper.get_view(root_node)
+
+        values = self._map_node_property(root_node, root_view, query_result)
+        if not values:
+            return [], None
+
+        data = [
+            node
+            for nodes in values.values()
+            for node in self._nodes_to_dict(nodes)
+        ]
+
+        return data, query_result.cursors.get(root_node)
+
+    def _map_node_property(
+        self,
+        key: str,
+        view: View,
+        query_result: QueryResult,
+        result_property_key: str | None = None,
+    ) -> dict[tuple[str, str], list[Node]] | None:
+        if key not in query_result:
+            return None
+
+        mappings = self._get_property_mappings(key, view, query_result)
+
+        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)
+
+            entry = properties.get(result_property_key)
+            if not isinstance(entry, dict):
+                raise ValueError(
+                    f"Invalid result property key {result_property_key}"
+                )
+
+            return entry.get("space", ""), entry.get("externalId", "")
+
+        result: defaultdict[tuple[str, str], list[Node]] = defaultdict(list)
+        for node in nodes:
+            properties = node.properties.get(view_id, {})
+            node_id = get_node_id(node)
+            for mapping_key, (values, is_list) in mappings.items():
+                element = properties.get(mapping_key)
+
+                element_key: tuple[str, str] = (
+                    (element.get("space", ""), element.get("externalId", ""))
+                    if isinstance(element, dict)
+                    else node_id
+                )
+
+                entries = values.get(element_key)
+                if not entries:
+                    continue
+
+                entry_data = self._nodes_to_dict(entries)
+                properties[mapping_key] = (
+                    entry_data if is_list else entry_data[0]
+                )
+
+            node.properties[view_id] = properties
+
+            result[node_id].append(node)
+
+        return dict(result)
+
+    def _get_property_mappings(
+        self, key: str, view: View, query_result: QueryResult
+    ) -> dict[str, tuple[dict[tuple[str, str], list[Node]], bool]]:
+        mappings: dict[
+            str, tuple[dict[tuple[str, str], list[Node]], bool]
+        ] = {}
+
+        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
+            is_list = False
+
+            if isinstance(property, MappedProperty) and property.source:
+                entry = self._map_node_property(
+                    property_key,
+                    self._view_mapper.get_view(property.source.external_id),
+                    query_result,
+                )
+                is_list = False
+            elif (
+                isinstance(property, SingleReverseDirectRelation)
+                and property.source
+            ):
+                entry = self._map_node_property(
+                    property_key,
+                    self._view_mapper.get_view(property.source.external_id),
+                    query_result,
+                    property.through.property,
+                )
+                is_list = False
+            elif (
+                isinstance(property, MultiReverseDirectRelation)
+                and property.source
+            ):
+                entry = self._map_node_property(
+                    property_key,
+                    self._view_mapper.get_view(property.source.external_id),
+                    query_result,
+                    property.through.property,
+                )
+                is_list = True
+
+            elif isinstance(property, EdgeConnection) and property.source:
+                entry = self._map_edge_property(
+                    property_key,
+                    self._view_mapper.get_view(property.source.external_id),
+                    query_result,
+                    property.direction,
+                )
+                is_list = True
+
+            if entry:
+                mappings[property_name] = entry, is_list
+
+        return mappings
+
+    def _map_edge_property(
+        self,
+        key: str,
+        view: View,
+        query_result: QueryResult,
+        edge_direction: EDGE_DIRECTION,
+    ) -> dict[tuple[str, str], list[Node]] | None:
+        edge_key = f"{key}{NESTED_SEP}{EDGE_MARKER}"
+        if key not in query_result or edge_key not in query_result:
+            return None
+
+        nodes = self._map_node_property(key, view, query_result)
+        if not nodes:
+            return None
+
+        result: defaultdict[tuple[str, str], list[Node]] = defaultdict(list)
+
+        edges = (
+            edge
+            for edge in query_result.get_edges(edge_key)
+            if isinstance(edge, Edge)
+        )
+        for edge in edges:
+            entry_key, node_key = (
+                (
+                    edge.end_node.as_tuple(),
+                    edge.start_node.as_tuple(),
+                )
+                if edge_direction == "inwards"
+                else (edge.start_node.as_tuple(), edge.end_node.as_tuple())
+            )
+
+            if node_item := nodes.get(node_key):
+                result[entry_key].extend(node_item)
+
+        return dict(result)
+
+    def _nodes_to_dict(self, nodes: list[Node]) -> list[dict[str, Any]]:
+        return [self._node_to_dict(node) for node in nodes]
+
+    def _node_to_dict(self, node: Node) -> dict[str, Any]:
+        entry = node.dump()
+        properties: dict[str, dict[str, dict[str, Any]]] = (
+            entry.pop("properties") or {}
+        )
+        for space_mapping in properties.values():
+            for view_mapping in space_mapping.values():
+                entry.update(view_mapping)
+
+        return entry