industrial-model 0.1.21__tar.gz → 0.1.23__tar.gz

{industrial_model-0.1.21 → industrial_model-0.1.23}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: industrial-model
-Version: 0.1.21
+Version: 0.1.23
 Summary: Industrial Model ORM
 Author-email: Lucas Alves <lucasrosaalves@gmail.com>
 Classifier: Programming Language :: Python
@@ -63,6 +63,7 @@ from industrial_model import (
     col,
     or_,
     select,
+    search,
 )
 
 # Define entities (view instances)
@@ -245,6 +246,25 @@ person = engine.query_all_pages(statement)[0]
 
 engine.delete([person])
 
+
+# 10. Search
+
+#  Notes:
+#     External ID searches work as prefix searches.
+#     This method does not include edges or direct relations in the result model.
+#     Filter does not support nested properties.
+
+class Entity(ViewInstance):
+    view_config = ViewInstanceConfig(
+        view_external_id="Person"
+    )
+    name: str
+
+
+statement = search(Entity).query_by("Lucas", [Entity.name])
+person = engine.search(statement)
+
+
 ```
 
 ---

{industrial_model-0.1.21 → industrial_model-0.1.23}/README.md

@@ -43,6 +43,7 @@ from industrial_model import (
     col,
     or_,
     select,
+    search,
 )
 
 # Define entities (view instances)
@@ -225,6 +226,25 @@ person = engine.query_all_pages(statement)[0]
 
 engine.delete([person])
 
+
+# 10. Search
+
+#  Notes:
+#     External ID searches work as prefix searches.
+#     This method does not include edges or direct relations in the result model.
+#     Filter does not support nested properties.
+
+class Entity(ViewInstance):
+    view_config = ViewInstanceConfig(
+        view_external_id="Person"
+    )
+    name: str
+
+
+statement = search(Entity).query_by("Lucas", [Entity.name])
+person = engine.search(statement)
+
+
 ```
 
 ---

{industrial_model-0.1.21 → industrial_model-0.1.23}/industrial_model/__init__.py

@@ -13,7 +13,7 @@ from .models import (
     ViewInstanceConfig,
     WritableViewInstance,
 )
-from .statements import aggregate, and_, col, not_, or_, select
+from .statements import aggregate, and_, col, not_, or_, search, select
 
 __all__ = [
     "aggregate",
@@ -23,6 +23,7 @@ __all__ = [
     "col",
     "not_",
     "select",
+    "search",
     "ViewInstance",
     "InstanceId",
     "TViewInstance",

{industrial_model-0.1.21 → industrial_model-0.1.23}/industrial_model/cognite_adapters/__init__.py

@@ -16,7 +16,11 @@ from industrial_model.models import (
     TViewInstance,
     TWritableViewInstance,
 )
-from industrial_model.statements import AggregationStatement, Statement
+from industrial_model.statements import (
+    AggregationStatement,
+    SearchStatement,
+    Statement,
+)
 
 from .aggregation_mapper import AggregationMapper
 from .optimizer import QueryOptimizer
@@ -24,6 +28,7 @@ from .query_mapper import QueryMapper
 from .query_result_mapper import (
     QueryResultMapper,
 )
+from .search_mapper import SearchMapper
 from .upsert_mapper import UpsertMapper
 from .utils import (
     append_nodes_and_edges,
@@ -34,9 +39,7 @@ from .view_mapper import ViewMapper
 
 
 class CogniteAdapter:
-    def __init__(
-        self, cognite_client: CogniteClient, data_model_id: DataModelId
-    ):
+    def __init__(self, cognite_client: CogniteClient, data_model_id: DataModelId):
         self._cognite_client = cognite_client
 
         view_mapper = ViewMapper(cognite_client, data_model_id)
@@ -45,6 +48,20 @@ class CogniteAdapter:
         self._result_mapper = QueryResultMapper(view_mapper)
         self._upsert_mapper = UpsertMapper(view_mapper)
         self._aggregation_mapper = AggregationMapper(view_mapper)
+        self._search_mapper = SearchMapper(view_mapper)
+
+    def search(self, statement: SearchStatement[TViewInstance]) -> list[dict[str, Any]]:
+        search_query = self._search_mapper.map(statement)
+        data = self._cognite_client.data_modeling.instances.search(
+            view=search_query.view.as_id(),
+            query=search_query.query,
+            filter=search_query.filter,
+            properties=search_query.query_properties,
+            limit=search_query.limit,
+            sort=search_query.sort,
+        )
+
+        return self._result_mapper.nodes_to_dict(data)
 
     def query(
         self, statement: Statement[TViewInstance], all_pages: bool
@@ -111,18 +128,14 @@ class CogniteAdapter:
         operation = self._upsert_mapper.map(entries)
 
         for node_chunk in operation.chunk_nodes():
-            logger.debug(
-                f"Upserting {len(node_chunk)} nodes (replace={replace})"
-            )
+            logger.debug(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.debug(
-                f"Upserting {len(edge_chunk)} edges (replace={replace})"
-            )
+            logger.debug(f"Upserting {len(edge_chunk)} edges (replace={replace})")
             self._cognite_client.data_modeling.instances.apply(
                 edges=edge_chunk,
                 replace=replace,
@@ -151,9 +164,7 @@ class CogniteAdapter:
         if not new_query:
             return None
 
-        new_query_result = self._cognite_client.data_modeling.instances.query(
-            new_query
-        )
+        new_query_result = self._cognite_client.data_modeling.instances.query(new_query)
 
         result = map_nodes_and_edges(new_query_result, new_query)
 

{industrial_model-0.1.21 → industrial_model-0.1.23}/industrial_model/cognite_adapters/aggregation_mapper.py

@@ -44,9 +44,7 @@ class AggregationMapper:
         root_view = self._view_mapper.get_view(root_node)
 
         filters_ = (
-            filters.And(
-                *self._filter_mapper.map(statement.where_clauses, root_view)
-            )
+            filters.And(*self._filter_mapper.map(statement.where_clauses, root_view))
             if statement.where_clauses
             else None
         )
@@ -63,13 +61,17 @@ class AggregationMapper:
             metric_aggregation = Count(statement.aggregation_property.property)
 
         if metric_aggregation is None:
-            raise ValueError(
-                f"Unsupported aggregate function: {statement.aggregate_}"
-            )
+            raise ValueError(f"Unsupported aggregate function: {statement.aggregate_}")
+
+        group_by_columns = (
+            [prop.property for prop in statement.group_by_properties]
+            if statement.group_by_properties is not None
+            else statement.entity.get_group_by_fields()
+        )
         return AggregationQuery(
             view=root_view,
             metric_aggregation=metric_aggregation,
             filters=filters_,
-            group_by_columns=statement.entity.get_group_by_fields(),
+            group_by_columns=group_by_columns,
             limit=statement.limit_,
         )

{industrial_model-0.1.21 → industrial_model-0.1.23}/industrial_model/cognite_adapters/filter_mapper.py

@@ -73,9 +73,7 @@ class FilterMapper:
         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
-            )
+            target_view = self._get_nested_target_view(expression.property, root_view)
 
             assert isinstance(value_, Expression)
 
@@ -91,9 +89,7 @@ class FilterMapper:
             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"
-        )
+        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]

{industrial_model-0.1.21 → industrial_model-0.1.23}/industrial_model/cognite_adapters/optimizer.py

@@ -16,9 +16,7 @@ SPACE_PROPERTY = "space"
 
 
 class QueryOptimizer:
-    def __init__(
-        self, cognite_client: CogniteClient, data_model_id: DataModelId
-    ):
+    def __init__(self, cognite_client: CogniteClient, data_model_id: DataModelId):
         self._all_spaces = data_model_id.instance_spaces
         self._cognite_client = cognite_client
         self._lock = Lock()
@@ -36,9 +34,7 @@ class QueryOptimizer:
             return
 
         filter_spaces = (
-            self._find_spaces(instance_spaces_prefix)
-            if instance_spaces_prefix
-            else []
+            self._find_spaces(instance_spaces_prefix) if instance_spaces_prefix else []
         )
         if instance_spaces:
             filter_spaces.extend(instance_spaces)
@@ -48,9 +44,9 @@ class QueryOptimizer:
 
     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):
+            if isinstance(where_clause, BoolExpression) and self._has_space_filter(
+                where_clause.filters
+            ):
                 return True
             elif (
                 isinstance(where_clause, LeafExpression)
@@ -64,9 +60,7 @@ class QueryOptimizer:
         all_spaces = self._load_spaces()
 
         return [
-            space
-            for space in all_spaces
-            if space.startswith(instance_spaces_prefix)
+            space for space in all_spaces if space.startswith(instance_spaces_prefix)
         ]
 
     def _load_spaces(self) -> list[str]:

{industrial_model-0.1.21 → industrial_model-0.1.23}/industrial_model/cognite_adapters/query_mapper.py

@@ -43,13 +43,9 @@ class QueryMapper:
         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_: list[filters.Filter] = [filters.HasData(views=[root_view_id])]
 
-        filters_.extend(
-            self._filter_mapper.map(statement.where_clauses, root_view)
-        )
+        filters_.extend(self._filter_mapper.map(statement.where_clauses, root_view))
 
         with_: dict[str, ResultSetExpression] = {
             root_node: NodeResultSetExpression(
@@ -60,9 +56,7 @@ class QueryMapper:
         }
         select_: dict[str, Select] = {}
 
-        relations = get_schema_properties(
-            statement.entity, NESTED_SEP, root_node
-        )
+        relations = get_schema_properties(statement.entity, NESTED_SEP, root_node)
 
         properties = self._include_statements(
             root_node, root_view, relations, with_, select_
@@ -76,9 +70,7 @@ class QueryMapper:
 
     def _get_select(self, view_id: ViewId, properties: list[str]) -> Select:
         return (
-            Select(
-                sources=[SourceSelector(source=view_id, properties=properties)]
-            )
+            Select(sources=[SourceSelector(source=view_id, properties=properties)])
             if properties
             else Select()
         )
@@ -118,9 +110,7 @@ class QueryMapper:
                         through=view.as_property_ref(property_name),
                         limit=MAX_LIMIT,
                     )
-                    select_[property_key] = self._get_select(
-                        property.source, props
-                    )
+                    select_[property_key] = self._get_select(property.source, props)
 
             elif (
                 isinstance(property, MultiReverseDirectRelation)
@@ -138,27 +128,21 @@ class QueryMapper:
                 with_[property_key] = NodeResultSetExpression(
                     from_=key,
                     direction="inwards",
-                    through=property.source.as_property_ref(
-                        property.through.property
-                    ),
+                    through=property.source.as_property_ref(property.through.property),
                     limit=MAX_LIMIT,
                 )
 
                 if property.through.property not in props:
                     props.append(property.through.property)
 
-                select_[property_key] = self._get_select(
-                    property.source, props
-                )
+                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()
-                    ),
+                    filter=filters.Equals(["edge", "type"], property.type.dump()),
                     direction=property.direction,
                     limit=MAX_LIMIT,
                 )
@@ -176,8 +160,6 @@ class QueryMapper:
                     with_,
                     select_,
                 )
-                select_[property_key] = self._get_select(
-                    property.source, props
-                )
+                select_[property_key] = self._get_select(property.source, props)
 
         return select_properties

{industrial_model-0.1.21 → industrial_model-0.1.23}/industrial_model/cognite_adapters/query_result_mapper.py

@@ -7,6 +7,7 @@ from cognite.client.data_classes.data_modeling import (
     EdgeConnection,
     MappedProperty,
     Node,
+    NodeList,
     View,
 )
 from cognite.client.data_classes.data_modeling.views import (
@@ -41,9 +42,7 @@ class QueryResultMapper:
         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"
-            )
+            raise ValueError(f"{root_node} is not available in the query result")
 
         root_view = self._view_mapper.get_view(root_node)
 
@@ -51,14 +50,13 @@ class QueryResultMapper:
         if not values:
             return []
 
-        data = [
-            node
-            for nodes in values.values()
-            for node in self._nodes_to_dict(nodes)
-        ]
+        data = [node for nodes in values.values() for node in self.nodes_to_dict(nodes)]
 
         return data
 
+    def nodes_to_dict(self, nodes: list[Node] | NodeList[Node]) -> list[dict[str, Any]]:
+        return [self._node_to_dict(node) for node in nodes]
+
     def _map_node_property(
         self,
         key: str,
@@ -79,9 +77,7 @@ class QueryResultMapper:
 
             entry = properties.get(result_property_key)
             if not isinstance(entry, dict):
-                raise ValueError(
-                    f"Invalid result property key {result_property_key}"
-                )
+                raise ValueError(f"Invalid result property key {result_property_key}")
 
             return entry.get("space", ""), entry.get("externalId", "")
 
@@ -124,15 +120,11 @@ class QueryResultMapper:
                 if not node_entries:
                     continue
 
-                entry_data = self._nodes_to_dict(node_entries)
-                properties[mapping_key] = (
-                    entry_data if is_list else entry_data[0]
-                )
+                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
-                    )
+                    edges_mapping[mapping_key] = self._edges_to_model(edge_entries)
             properties["_edges"] = edges_mapping
 
             node.properties[view_id] = properties
@@ -155,9 +147,7 @@ class QueryResultMapper:
             nodes: dict[tuple[str, str], list[Node]] | None = None
             edges: dict[tuple[str, str], list[Edge]] | None = None
             is_list = False
-            connection_type: ConnectionTypeEnum = (
-                ConnectionTypeEnum.DIRECT_RELATION
-            )
+            connection_type: ConnectionTypeEnum = ConnectionTypeEnum.DIRECT_RELATION
 
             if isinstance(property, MappedProperty) and property.source:
                 nodes = self._map_node_property(
@@ -167,10 +157,7 @@ class QueryResultMapper:
                 )
                 is_list = False
                 connection_type = ConnectionTypeEnum.DIRECT_RELATION
-            elif (
-                isinstance(property, SingleReverseDirectRelation)
-                and property.source
-            ):
+            elif isinstance(property, SingleReverseDirectRelation) and property.source:
                 nodes = self._map_node_property(
                     property_key,
                     self._view_mapper.get_view(property.source.external_id),
@@ -179,10 +166,7 @@ class QueryResultMapper:
                 )
                 is_list = False
                 connection_type = ConnectionTypeEnum.REVERSE_DIRECT_RELATION
-            elif (
-                isinstance(property, MultiReverseDirectRelation)
-                and property.source
-            ):
+            elif isinstance(property, MultiReverseDirectRelation) and property.source:
                 nodes = self._map_node_property(
                     property_key,
                     self._view_mapper.get_view(property.source.external_id),
@@ -231,12 +215,8 @@ class QueryResultMapper:
             return None, None
 
         visited: set[tuple[str, str]] = set()
-        nodes_result: defaultdict[tuple[str, str], list[Node]] = defaultdict(
-            list
-        )
-        edges_result: defaultdict[tuple[str, str], list[Edge]] = 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:
@@ -257,17 +237,12 @@ class QueryResultMapper:
 
         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]]] = (
-            entry.pop("properties") or {}
-        )
+        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)

industrial_model-0.1.23/industrial_model/cognite_adapters/search_mapper.py

@@ -0,0 +1,50 @@
+from dataclasses import dataclass
+
+import cognite.client.data_classes.filters as filters
+from cognite.client.data_classes.data_modeling import InstanceSort, View
+
+from industrial_model.models import TViewInstance
+from industrial_model.statements import SearchStatement
+
+from .filter_mapper import (
+    FilterMapper,
+)
+from .sort_mapper import SortMapper
+from .view_mapper import ViewMapper
+
+
+@dataclass
+class SearchQuery:
+    view: View
+    filter: filters.Filter | None
+    query: str | None
+    query_properties: list[str] | None
+    limit: int
+    sort: list[InstanceSort]
+
+
+class SearchMapper:
+    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: SearchStatement[TViewInstance]) -> SearchQuery:
+        root_node = statement.entity.get_view_external_id()
+
+        root_view = self._view_mapper.get_view(root_node)
+
+        filters_ = self._filter_mapper.map(statement.where_clauses, root_view)
+
+        sort_clauses = self._sort_mapper.map(statement.sort_clauses, root_view)
+        for item in sort_clauses:
+            item.nulls_first = None
+
+        return SearchQuery(
+            view=root_view,
+            filter=filters.And(*filters_) if filters_ else None,
+            query=statement.query,
+            query_properties=statement.query_properties,
+            limit=statement.limit_,
+            sort=sort_clauses,
+        )

{industrial_model-0.1.21 → industrial_model-0.1.23}/industrial_model/cognite_adapters/upsert_mapper.py

@@ -31,19 +31,14 @@ class UpsertMapper:
         edges_to_delete: dict[tuple[str, str], EdgeContainer] = {}
 
         for instance in instances:
-            entry_nodes, entry_edges, entry_edges_to_delete = (
-                self._map_instance(instance)
+            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.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
-                }
+                {(item.space, item.external_id): item for item in entry_edges_to_delete}
             )
 
         return UpsertOperation(
@@ -76,9 +71,7 @@ class UpsertMapper:
                     if isinstance(entry, datetime.datetime)
                     else entry
                 )
-            elif isinstance(property, EdgeConnection) and isinstance(
-                entry, list
-            ):
+            elif isinstance(property, EdgeConnection) and isinstance(entry, list):
                 possible_entries = self._map_edges(instance, property, entry)
 
                 previous_edges = {
@@ -104,9 +97,7 @@ class UpsertMapper:
         node = NodeApply(
             external_id=instance.external_id,
             space=instance.space,
-            sources=[
-                NodeOrEdgeData(source=view.as_id(), properties=properties)
-            ],
+            sources=[NodeOrEdgeData(source=view.as_id(), properties=properties)],
         )
 
         return node, edges, edges_to_delete

{industrial_model-0.1.21 → industrial_model-0.1.23}/industrial_model/cognite_adapters/utils.py

@@ -43,9 +43,7 @@ def get_property_ref(
 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
-    ]
+    return [get_cognite_instance_id(instance_id) for instance_id in instance_ids]
 
 
 def get_cognite_instance_id(instance_id: TViewInstance) -> dict[str, str]:
@@ -115,9 +113,7 @@ def _create_query(
         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}
-        )
+        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}
         )

{industrial_model-0.1.21 → industrial_model-0.1.23}/industrial_model/cognite_adapters/view_mapper.py

@@ -9,9 +9,7 @@ from industrial_model.config import DataModelId
 
 
 class ViewMapper:
-    def __init__(
-        self, cognite_client: CogniteClient, data_model_id: DataModelId
-    ):
+    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
@@ -26,9 +24,7 @@ class ViewMapper:
     def get_view(self, view_external_id: str) -> View:
         views = self._load_views()
         if view_external_id not in views:
-            raise ValueError(
-                f"View {view_external_id} is not available in data model"
-            )
+            raise ValueError(f"View {view_external_id} is not available in data model")
 
         return views[view_external_id]