industrial-model 0.1.20__py3-none-any.whl → 0.1.22__py3-none-any.whl

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/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/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,9 +61,7 @@ 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_}")
         return AggregationQuery(
             view=root_view,
             metric_aggregation=metric_aggregation,

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/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/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/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/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/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/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]:
@@ -109,15 +107,13 @@ def _create_query(
     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]
+        children = nodes_children.get(cursor_key, set())
+        parent = nodes_parent.get(cursor_key, set())
 
         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}
         )
@@ -142,17 +138,16 @@ def _get_leaf_cursors(
             cursor_key == view_external_id
             or not cursor_value
             or len(query_result[cursor_key]) != MAX_LIMIT
-            or cursor_key not in nodes_children
         ):
             continue
 
-        children = nodes_children[cursor_key]
+        children = nodes_children.get(cursor_key, set())
 
         target_cursors_keys = target_cursors.keys()
         if len(children.intersection(target_cursors_keys)) > 0:
             continue
 
-        parent = nodes_parent[cursor_key]
+        parent = nodes_parent.get(cursor_key, set())
         cursor_to_remove = parent.intersection(target_cursors_keys)
         for cursor_key_to_remove in cursor_to_remove:
             target_cursors.pop(cursor_key_to_remove)

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]
 

industrial_model/engines/async_engine.py

@@ -8,7 +8,11 @@ from industrial_model.models import (
     TWritableViewInstance,
     ValidationMode,
 )
-from industrial_model.statements import AggregationStatement, Statement
+from industrial_model.statements import (
+    AggregationStatement,
+    SearchStatement,
+    Statement,
+)
 from industrial_model.utils import run_async
 
 from .engine import Engine
@@ -22,6 +26,19 @@ class AsyncEngine:
     ):
         self._engine = Engine(cognite_client, data_model_id)
 
+    async def search_async(
+        self,
+        statement: SearchStatement[TViewInstance],
+        validation_mode: ValidationMode = "raiseOnError",
+    ) -> list[TViewInstance]:
+        """
+        Note:
+            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.
+        """
+        return await run_async(self._engine.search, statement, validation_mode)
+
     async def query_async(
         self,
         statement: Statement[TViewInstance],
@@ -34,9 +51,7 @@ class AsyncEngine:
         statement: Statement[TViewInstance],
         validation_mode: ValidationMode = "raiseOnError",
     ) -> list[TViewInstance]:
-        return await run_async(
-            self._engine.query_all_pages, statement, validation_mode
-        )
+        return await run_async(self._engine.query_all_pages, statement, validation_mode)
 
     async def aggregate_async(
         self, statement: AggregationStatement[TAggregatedViewInstance]

industrial_model/engines/engine.py

@@ -12,7 +12,11 @@ from industrial_model.models import (
     TWritableViewInstance,
     ValidationMode,
 )
-from industrial_model.statements import AggregationStatement, Statement
+from industrial_model.statements import (
+    AggregationStatement,
+    SearchStatement,
+    Statement,
+)
 
 
 class Engine:
@@ -23,6 +27,20 @@ class Engine:
     ):
         self._cognite_adapter = CogniteAdapter(cognite_client, data_model_id)
 
+    def search(
+        self,
+        statement: SearchStatement[TViewInstance],
+        validation_mode: ValidationMode = "raiseOnError",
+    ) -> list[TViewInstance]:
+        """
+        Note:
+            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.
+        """
+        data = self._cognite_adapter.search(statement)
+        return self._validate_data(statement.entity, data, validation_mode)
+
     def query(
         self,
         statement: Statement[TViewInstance],

industrial_model/models/entities.py

@@ -88,9 +88,7 @@ class AggregatedViewInstance(RootModel, metaclass=DBModelMetaclass):
 
 
 TViewInstance = TypeVar("TViewInstance", bound=ViewInstance)
-TWritableViewInstance = TypeVar(
-    "TWritableViewInstance", bound=WritableViewInstance
-)
+TWritableViewInstance = TypeVar("TWritableViewInstance", bound=WritableViewInstance)
 TAggregatedViewInstance = TypeVar(
     "TAggregatedViewInstance", bound=AggregatedViewInstance
 )

industrial_model/models/schemas.py

@@ -38,8 +38,7 @@ def get_parent_and_children_nodes(
     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)
+            NESTED_SEP.join(key_parts[:i]) for i in range(len(key_parts) - 1, 0, -1)
         }
 
         valid_paths: set[str] = set()
@@ -142,17 +141,13 @@ def _flatten_dict_keys(
 ) -> list[str]:
     paths: set[str] = set()
     for key, value in data.items():
-        full_key = (
-            f"{parent_key}{nested_separator}{key}" if parent_key else key
-        )
+        full_key = f"{parent_key}{nested_separator}{key}" if parent_key else key
         paths.add(full_key)
         if isinstance(value, dict) and value:
             paths.update(_flatten_dict_keys(value, full_key, nested_separator))
         elif isinstance(value, str):
             paths.add(f"{full_key}{nested_separator}{value}")
         elif isinstance(value, list | set):
-            paths.update(
-                [f"{full_key}{nested_separator}{item}" for item in value]
-            )
+            paths.update([f"{full_key}{nested_separator}{item}" for item in value])
 
     return sorted(paths)

industrial_model/queries/models.py

@@ -5,9 +5,7 @@ from .params import NestedQueryParam, QueryParam, SortParam
 
 
 class BaseQuery(RootModel):
-    def to_statement(
-        self, entity: type[TViewInstance]
-    ) -> Statement[TViewInstance]:
+    def to_statement(self, entity: type[TViewInstance]) -> Statement[TViewInstance]:
         statement = Statement(entity)
 
         for key, item in self.__class__.model_fields.items():
@@ -27,9 +25,7 @@ class BasePaginatedQuery(BaseQuery):
     limit: int = 1000
     cursor: str | None = None
 
-    def to_statement(
-        self, entity: type[TViewInstance]
-    ) -> Statement[TViewInstance]:
+    def to_statement(self, entity: type[TViewInstance]) -> Statement[TViewInstance]:
         statement = super().to_statement(entity)
         statement.limit(self.limit)
         statement.cursor(self.cursor)

industrial_model/statements/__init__.py

@@ -15,6 +15,7 @@ from .expressions import (
 )
 
 T = TypeVar("T")
+AggregateTypes = Literal["count", "avg", "min", "max", "sum"]
 
 
 def _create_column(property: str | Column | Any) -> Column:
@@ -22,14 +23,13 @@ def _create_column(property: str | Column | Any) -> Column:
 
 
 @dataclass
-class Statement(Generic[T]):
+class BaseStatement(Generic[T]):
     entity: type[T] = field(init=True)
     where_clauses: list[Expression] = field(init=False, default_factory=list)
     sort_clauses: list[tuple[Column, SORT_DIRECTION]] = field(
         init=False, default_factory=list
     )
     limit_: int = field(init=False, default=DEFAULT_LIMIT)
-    cursor_: str | None = field(init=False, default=None)
 
     def where(self, *expressions: bool | Expression) -> Self:
         for expression in expressions:
@@ -37,15 +37,17 @@ class Statement(Generic[T]):
             self.where_clauses.append(expression)
         return self
 
+    def limit(self, limit: int) -> Self:
+        self.limit_ = limit
+        return self
+
     def asc(self, property: str | Column | Any) -> Self:
         return self.sort(property, "ascending")
 
     def desc(self, property: str | Column | Any) -> Self:
         return self.sort(property, "descending")
 
-    def sort(
-        self, property: str | Column | Any, direction: SORT_DIRECTION
-    ) -> Self:
+    def sort(self, property: str | Column | Any, direction: SORT_DIRECTION) -> Self:
         self.sort_clauses.append(
             (
                 _create_column(property),
@@ -54,16 +56,40 @@ class Statement(Generic[T]):
         )
         return self
 
-    def limit(self, limit: int) -> Self:
-        self.limit_ = limit
-        return self
+
+@dataclass
+class Statement(BaseStatement[T]):
+    cursor_: str | None = field(init=False, default=None)
 
     def cursor(self, cursor: str | None) -> Self:
         self.cursor_ = cursor
         return self
 
 
-AggregateTypes = Literal["count", "avg", "min", "max", "sum"]
+@dataclass
+class SearchStatement(BaseStatement[T]):
+    query: str | None = field(init=False, default=None)
+    query_properties: list[str] | None = field(init=False, default=None)
+
+    def query_by(
+        self,
+        query: str,
+        query_properties: list[Column | str | Any] | None = None,
+    ) -> Self:
+        self.query = query
+        self.query_properties = (
+            [
+                prop
+                if isinstance(prop, str)
+                else prop.property
+                if isinstance(prop, Column)
+                else str(prop)
+                for prop in query_properties
+            ]
+            if query_properties
+            else None
+        )
+        return self
 
 
 @dataclass
@@ -71,9 +97,7 @@ class AggregationStatement(Generic[T]):
     entity: type[T] = field(init=True)
     aggregate: AggregateTypes = field(init=True)
 
-    aggregation_property: Column = field(
-        init=False, default=Column("externalId")
-    )
+    aggregation_property: Column = field(init=False, default=Column("externalId"))
     where_clauses: list[Expression] = field(init=False, default_factory=list)
     limit_: int = field(init=False, default=-1)
 
@@ -103,11 +127,17 @@ def aggregate(
     return AggregationStatement(entity=entity, aggregate=aggregate)
 
 
+def search(entity: type[T]) -> SearchStatement[T]:
+    return SearchStatement(entity)
+
+
 __all__ = [
     "aggregate",
     "AggregationStatement",
     "Statement",
     "select",
+    "search",
+    "SearchStatement",
     "Column",
     "col",
     "Expression",

industrial_model/statements/expressions.py

@@ -9,9 +9,7 @@ from industrial_model.constants import (
 )
 
 RANGE_SUPPORTED_VALUES = str | int | float
-LIST_SUPPORTED_VALUES = (
-    list[str] | list[int] | list[float] | list[dict[str, str]]
-)
+LIST_SUPPORTED_VALUES = list[str] | list[int] | list[float] | list[dict[str, str]]
 SUPPORTED_VALUES = (
     RANGE_SUPPORTED_VALUES | bool | dict[str, str] | LIST_SUPPORTED_VALUES
 )
@@ -49,9 +47,7 @@ class LeafExpression(Expression):
 class Column:
     def __init__(self, property: Any):
         assert isinstance(property, str | Column)
-        property_ = (
-            property.property if isinstance(property, Column) else property
-        )
+        property_ = property.property if isinstance(property, Column) else property
         self.property: str = property_
 
     def __hash__(self) -> int:
@@ -120,23 +116,17 @@ class Column:
             filters=[self._compare(operator, other)],  # type: ignore
         )
 
-    def nested_(
-        self, expression: bool | LeafExpression | BoolExpression
-    ) -> bool:
+    def nested_(self, expression: bool | LeafExpression | BoolExpression) -> bool:
         if not isinstance(expression, Expression):
             raise ValueError("Invalid expression")
 
         return self._compare("nested", expression)
 
-    def _compare(
-        self, operator: LEAF_EXPRESSION_OPERATORS, value: Any
-    ) -> bool:
+    def _compare(self, operator: LEAF_EXPRESSION_OPERATORS, value: Any) -> bool:
         if isinstance(value, Column):
             raise ValueError("can not compare two columns in a graphdb")
 
-        return LeafExpression(
-            property=self.property, operator=operator, value=value
-        )  # type: ignore
+        return LeafExpression(property=self.property, operator=operator, value=value)  # type: ignore
 
 
 def _unwrap_expressions(
@@ -150,21 +140,15 @@ def _unwrap_expressions(
 
 
 def and_(*expressions: bool | LeafExpression | BoolExpression) -> bool:
-    return BoolExpression(
-        operator="and", filters=_unwrap_expressions(*expressions)
-    )  # type: ignore
+    return BoolExpression(operator="and", filters=_unwrap_expressions(*expressions))  # type: ignore
 
 
 def or_(*expressions: bool | LeafExpression | BoolExpression) -> bool:
-    return BoolExpression(
-        operator="or", filters=_unwrap_expressions(*expressions)
-    )  # type: ignore
+    return BoolExpression(operator="or", filters=_unwrap_expressions(*expressions))  # type: ignore
 
 
 def not_(*expressions: bool | LeafExpression | BoolExpression) -> bool:
-    return BoolExpression(
-        operator="not", filters=_unwrap_expressions(*expressions)
-    )  # type: ignore
+    return BoolExpression(operator="not", filters=_unwrap_expressions(*expressions))  # type: ignore
 
 
 def col(property: Any) -> Column:

{industrial_model-0.1.20.dist-info → industrial_model-0.1.22.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: industrial-model
-Version: 0.1.20
+Version: 0.1.22
 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.22.dist-info/RECORD

@@ -0,0 +1,32 @@
+industrial_model/__init__.py,sha256=dBbAmQ-0kWd2MW3Nmrz623P17pA2P7VDlspE_XhdPuo,835
+industrial_model/config.py,sha256=LseCZ8JvewFejjElPUQWdmIQQUU2wnd43McDid_Rixs,474
+industrial_model/constants.py,sha256=wtFdxkR9gojqekwXeyVCof6VUkJYuqjjgAgVVKrcxaw,450
+industrial_model/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+industrial_model/utils.py,sha256=oh4AxwxXaWgIC2uolkCbvkgo0ququHB6yAPVIXy45Ts,663
+industrial_model/cognite_adapters/__init__.py,sha256=cKPW5y-wWiyWYLPDaNjhqc5zdsgdLc40vunIlz2Debg,6246
+industrial_model/cognite_adapters/aggregation_mapper.py,sha256=w9EksY-_pNrc15DXHVRydfMdmPbvu4dzCRhzRB3F7QM,2302
+industrial_model/cognite_adapters/filter_mapper.py,sha256=xXTzxCYKU1-6TxG5W_bb8zNyWi2bfIM9RiujyOHUujg,3900
+industrial_model/cognite_adapters/models.py,sha256=2j2IS01uPkQEp9WdVk8seYzEqGcDdWFnpzXhusHB2zk,945
+industrial_model/cognite_adapters/optimizer.py,sha256=EQfCe-4mSeXhYa2kcqJerPXYsCRgMJPA1AIf1o1KNH0,2415
+industrial_model/cognite_adapters/query_mapper.py,sha256=QdFVkfmL4tWcCy_oypOluT0Fo68obPk2RjBPvVA9mUs,5950
+industrial_model/cognite_adapters/query_result_mapper.py,sha256=dzj5CdteITq-E3uQjX9s_B3WmryQV69pnwrGGHdcWDE,9139
+industrial_model/cognite_adapters/search_mapper.py,sha256=dh69Te8oiDIRPbEMTcof-ZJ4XD-xB5iiBHUR0UETzho,1504
+industrial_model/cognite_adapters/sort_mapper.py,sha256=RJUAYlZGXoYzK0PwX63cibRF_L-MUq9g2ZsC2EeNIF4,696
+industrial_model/cognite_adapters/upsert_mapper.py,sha256=2QN4LXdUTWDmnLmXGKOLnCdTnUm5Zz2sZXD80-kAeQM,4720
+industrial_model/cognite_adapters/utils.py,sha256=-9NUG84r49cKvmmzoUinxNF1zhHIlmoWFypUCD_WTeM,4800
+industrial_model/cognite_adapters/view_mapper.py,sha256=fihoxyLKvq8xa1oArdCCkukqVpOexFcBDnEyMACpHF8,1430
+industrial_model/engines/__init__.py,sha256=7aGHrUm2MxIq39vR8h0xu3i1zNOuT9H9U-q4lV3nErQ,102
+industrial_model/engines/async_engine.py,sha256=-sQv2vn93bBmTZOxY_C1r40YP7IMl1ND2gCN2GE9BqY,2188
+industrial_model/engines/engine.py,sha256=lECOpjN6fGvKt28b441isT_u4UNF3LeLiJH-zwHmkYw,3798
+industrial_model/models/__init__.py,sha256=AzJ0CyPK5PvUCX45FFtybl13tkukUvk2UAF_90s_LQ8,742
+industrial_model/models/base.py,sha256=iGhDjXqA5ULEQIFHtkMi7WYJl0nQq1wi8_zqOr-Ep78,1649
+industrial_model/models/entities.py,sha256=tHOFjS-9XsaXVrZ-x0tZR1DWL2Cc8MHe82qsIazqmnM,2811
+industrial_model/models/schemas.py,sha256=EoLK9GYdS-0DQ98myTP3wOU9YpWIJOfDSLOYZUy3xEY,4559
+industrial_model/queries/__init__.py,sha256=7aheTE5qs03rxWm9fmGWptbz_p9OIXXYD8if56cqs18,227
+industrial_model/queries/models.py,sha256=zb2riqpB5S6BoKRqnau8MHjsrU980MOCDTbyB3hjXYY,1157
+industrial_model/queries/params.py,sha256=ehgCoR5n6E-tkEuoymZ2lkLcSzMaBAx_HnyJ7sWpqz0,964
+industrial_model/statements/__init__.py,sha256=_79ZJrJNJ4XjiDJDJld-kh4DOWA3zWqY6npe6HHP8hM,3949
+industrial_model/statements/expressions.py,sha256=4ZZOcZroI5-4xRw4PXIRlufi0ARndE5zSbbxLDpR2Ec,4816
+industrial_model-0.1.22.dist-info/METADATA,sha256=UPKMtUmBdYCfPW6wi67vfiHpajTHo2WJD_oPL0tqACA,6858
+industrial_model-0.1.22.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+industrial_model-0.1.22.dist-info/RECORD,,

industrial_model-0.1.20.dist-info/RECORD

@@ -1,31 +0,0 @@
-industrial_model/__init__.py,sha256=ix_92MZpYWxzxEytdELWUINioGYd2R1UNdJIU4IO5j0,813
-industrial_model/config.py,sha256=LseCZ8JvewFejjElPUQWdmIQQUU2wnd43McDid_Rixs,474
-industrial_model/constants.py,sha256=wtFdxkR9gojqekwXeyVCof6VUkJYuqjjgAgVVKrcxaw,450
-industrial_model/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-industrial_model/utils.py,sha256=oh4AxwxXaWgIC2uolkCbvkgo0ququHB6yAPVIXy45Ts,663
-industrial_model/cognite_adapters/__init__.py,sha256=K-F0BjbF9yM_YilQ7U16t7-70aAsIWUaaU_uQOul418,5680
-industrial_model/cognite_adapters/aggregation_mapper.py,sha256=rXt2BZ2jS3afD9d3B2t8aizL9TWnMnqCCNdzTn_XR3Y,2362
-industrial_model/cognite_adapters/filter_mapper.py,sha256=X2BVp305v1Wt00DTZEnCooOG_g6Jo7bkj-NyfRtECAE,3952
-industrial_model/cognite_adapters/models.py,sha256=2j2IS01uPkQEp9WdVk8seYzEqGcDdWFnpzXhusHB2zk,945
-industrial_model/cognite_adapters/optimizer.py,sha256=BbbsA3r9ZKJt6mN1wVRZc046ov-vbml9_ESfurTNRUg,2477
-industrial_model/cognite_adapters/query_mapper.py,sha256=3fEcaLsGjLKIh-g1BbMcffQ6rp99JeCW555iJo8JW44,6260
-industrial_model/cognite_adapters/query_result_mapper.py,sha256=0IhW8YVQM_EBMQM_NXlzl27CuzSUWGX4g7K9VMUVgPw,9507
-industrial_model/cognite_adapters/sort_mapper.py,sha256=RJUAYlZGXoYzK0PwX63cibRF_L-MUq9g2ZsC2EeNIF4,696
-industrial_model/cognite_adapters/upsert_mapper.py,sha256=tWEiBJQeeNz1HDu0AoBIfCw_LL156Zg4h7ORKlZ__uw,4870
-industrial_model/cognite_adapters/utils.py,sha256=AyQxK8tDy-tAC6pji1ZzfAMsaVX2epiBhdwdCXRCPl4,4840
-industrial_model/cognite_adapters/view_mapper.py,sha256=qvT25gwIgPdAqZf-oQEh0GylUi5KEPwjLNmIl5ISQGA,1474
-industrial_model/engines/__init__.py,sha256=7aGHrUm2MxIq39vR8h0xu3i1zNOuT9H9U-q4lV3nErQ,102
-industrial_model/engines/async_engine.py,sha256=AkEHiHaO0wYiU8fwEG8H9EzNTlLpn9ygRPt-p4cQl7A,1675
-industrial_model/engines/engine.py,sha256=pZOvBbvIfa2aX2Ad8KnGntKmbZQmmvHqWzFvxEDc8xU,3224
-industrial_model/models/__init__.py,sha256=AzJ0CyPK5PvUCX45FFtybl13tkukUvk2UAF_90s_LQ8,742
-industrial_model/models/base.py,sha256=iGhDjXqA5ULEQIFHtkMi7WYJl0nQq1wi8_zqOr-Ep78,1649
-industrial_model/models/entities.py,sha256=1ZmijxOXEAWocf3HTJ9a7ci5cJEjzA3hbt0QvVP8Gb8,2817
-industrial_model/models/schemas.py,sha256=PXbCZTN5RHTLsufYYS3fwPcBzIBTJ48m9-hxBYXx-Yk,4625
-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=15LI9POLUdh6i3Zw0Aek0fVFfcTj8P8ZEa1Cli3pgXQ,3095
-industrial_model/statements/expressions.py,sha256=Sar1cIvy3sYi7tkWJN3ylHlZ252oN2mZJpZ1TX9jN3s,4940
-industrial_model-0.1.20.dist-info/METADATA,sha256=ZkydcSIJvDKweNn0nW8NeT9AeLjxHbO9iuR8Wn_Dyr4,6419
-industrial_model-0.1.20.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-industrial_model-0.1.20.dist-info/RECORD,,