industrial-model 0.1.8__tar.gz → 0.1.10__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (45) hide show
  1. {industrial_model-0.1.8 → industrial_model-0.1.10}/PKG-INFO +35 -1
  2. {industrial_model-0.1.8 → industrial_model-0.1.10}/README.md +34 -0
  3. {industrial_model-0.1.8 → industrial_model-0.1.10}/industrial_model/__init__.py +4 -0
  4. {industrial_model-0.1.8 → industrial_model-0.1.10}/industrial_model/cognite_adapters/__init__.py +35 -2
  5. industrial_model-0.1.10/industrial_model/cognite_adapters/models.py +35 -0
  6. {industrial_model-0.1.8 → industrial_model-0.1.10}/industrial_model/cognite_adapters/query_result_mapper.py +55 -25
  7. industrial_model-0.1.10/industrial_model/cognite_adapters/upsert_mapper.py +146 -0
  8. {industrial_model-0.1.8 → industrial_model-0.1.10}/industrial_model/engines/async_engine.py +6 -0
  9. {industrial_model-0.1.8 → industrial_model-0.1.10}/industrial_model/engines/engine.py +29 -1
  10. {industrial_model-0.1.8 → industrial_model-0.1.10}/industrial_model/models/__init__.py +6 -0
  11. {industrial_model-0.1.8 → industrial_model-0.1.10}/industrial_model/models/entities.py +38 -0
  12. {industrial_model-0.1.8 → industrial_model-0.1.10}/industrial_model/utils.py +1 -12
  13. {industrial_model-0.1.8 → industrial_model-0.1.10}/pyproject.toml +1 -1
  14. industrial_model-0.1.10/scripts/format.sh +2 -0
  15. industrial_model-0.1.10/scripts/lint.sh +3 -0
  16. industrial_model-0.1.8/tests/tests_adapter.py → industrial_model-0.1.10/tests/models.py +27 -22
  17. industrial_model-0.1.10/tests/test_upsert_mapper.py +39 -0
  18. industrial_model-0.1.10/tests/tests_adapter.py +27 -0
  19. {industrial_model-0.1.8 → industrial_model-0.1.10}/uv.lock +1 -1
  20. industrial_model-0.1.8/scripts/format.sh +0 -2
  21. industrial_model-0.1.8/scripts/lint.sh +0 -3
  22. {industrial_model-0.1.8 → industrial_model-0.1.10}/.gitignore +0 -0
  23. {industrial_model-0.1.8 → industrial_model-0.1.10}/.python-version +0 -0
  24. {industrial_model-0.1.8 → industrial_model-0.1.10}/.vscode/settings.json +0 -0
  25. {industrial_model-0.1.8 → industrial_model-0.1.10}/industrial_model/cognite_adapters/filter_mapper.py +0 -0
  26. {industrial_model-0.1.8 → industrial_model-0.1.10}/industrial_model/cognite_adapters/optimizer.py +0 -0
  27. {industrial_model-0.1.8 → industrial_model-0.1.10}/industrial_model/cognite_adapters/query_mapper.py +0 -0
  28. {industrial_model-0.1.8 → industrial_model-0.1.10}/industrial_model/cognite_adapters/sort_mapper.py +0 -0
  29. {industrial_model-0.1.8 → industrial_model-0.1.10}/industrial_model/cognite_adapters/utils.py +0 -0
  30. {industrial_model-0.1.8 → industrial_model-0.1.10}/industrial_model/cognite_adapters/view_mapper.py +0 -0
  31. {industrial_model-0.1.8 → industrial_model-0.1.10}/industrial_model/config.py +0 -0
  32. {industrial_model-0.1.8 → industrial_model-0.1.10}/industrial_model/constants.py +0 -0
  33. {industrial_model-0.1.8 → industrial_model-0.1.10}/industrial_model/engines/__init__.py +0 -0
  34. {industrial_model-0.1.8 → industrial_model-0.1.10}/industrial_model/models/base.py +0 -0
  35. {industrial_model-0.1.8 → industrial_model-0.1.10}/industrial_model/models/schemas.py +0 -0
  36. {industrial_model-0.1.8 → industrial_model-0.1.10}/industrial_model/py.typed +0 -0
  37. {industrial_model-0.1.8 → industrial_model-0.1.10}/industrial_model/queries/__init__.py +0 -0
  38. {industrial_model-0.1.8 → industrial_model-0.1.10}/industrial_model/queries/models.py +0 -0
  39. {industrial_model-0.1.8 → industrial_model-0.1.10}/industrial_model/queries/params.py +0 -0
  40. {industrial_model-0.1.8 → industrial_model-0.1.10}/industrial_model/statements/__init__.py +0 -0
  41. {industrial_model-0.1.8 → industrial_model-0.1.10}/industrial_model/statements/expressions.py +0 -0
  42. {industrial_model-0.1.8 → industrial_model-0.1.10}/scripts/build.sh +0 -0
  43. {industrial_model-0.1.8 → industrial_model-0.1.10}/tests/__init__.py +0 -0
  44. {industrial_model-0.1.8 → industrial_model-0.1.10}/tests/cognite-sdk-config.yaml +0 -0
  45. {industrial_model-0.1.8 → industrial_model-0.1.10}/tests/hubs.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: industrial-model
3
- Version: 0.1.8
3
+ Version: 0.1.10
4
4
  Summary: Industrial Model ORM
5
5
  Author-email: Lucas Alves <lucasrosaalves@gmail.com>
6
6
  Classifier: Programming Language :: Python
@@ -181,6 +181,40 @@ statement = select(Person).where(
181
181
  )
182
182
  all_results = engine.query_all_pages(statement)
183
183
 
184
+
185
+ # 7. Data Ingestion
186
+
187
+ from industrial_model import (
188
+ WritableViewInstance # necessary for data ingestion
189
+ )
190
+
191
+
192
+ class WritablePerson(WritableViewInstance):
193
+ name: str
194
+ lives_in: InstanceId
195
+ cars: list[InstanceId]
196
+
197
+ # You need to implement the end_id_factory so the model can build the edge ids automatically.
198
+ def edge_id_factory(
199
+ self, target_node: InstanceId, edge_type: InstanceId
200
+ ) -> InstanceId:
201
+ return InstanceId(
202
+ external_id=f"{self.external_id}-{target_node.external_id}-{edge_type.external_id}",
203
+ space=self.space,
204
+ )
205
+
206
+ statement = select(WritablePerson).where(
207
+ (WritablePerson.external_id == "Lucas")
208
+ )
209
+
210
+ person = engine.query_all_pages(statement)[0]
211
+
212
+ person.lives_in = InstanceId(external_id="br", space="data-space")
213
+ person.cars.clear() # Gonna remove all car edges from the person
214
+
215
+ engine.upsert([person])
216
+
184
217
  ```
185
218
 
219
+
186
220
  ---
@@ -161,6 +161,40 @@ statement = select(Person).where(
161
161
  )
162
162
  all_results = engine.query_all_pages(statement)
163
163
 
164
+
165
+ # 7. Data Ingestion
166
+
167
+ from industrial_model import (
168
+ WritableViewInstance # necessary for data ingestion
169
+ )
170
+
171
+
172
+ class WritablePerson(WritableViewInstance):
173
+ name: str
174
+ lives_in: InstanceId
175
+ cars: list[InstanceId]
176
+
177
+ # You need to implement the end_id_factory so the model can build the edge ids automatically.
178
+ def edge_id_factory(
179
+ self, target_node: InstanceId, edge_type: InstanceId
180
+ ) -> InstanceId:
181
+ return InstanceId(
182
+ external_id=f"{self.external_id}-{target_node.external_id}-{edge_type.external_id}",
183
+ space=self.space,
184
+ )
185
+
186
+ statement = select(WritablePerson).where(
187
+ (WritablePerson.external_id == "Lucas")
188
+ )
189
+
190
+ person = engine.query_all_pages(statement)[0]
191
+
192
+ person.lives_in = InstanceId(external_id="br", space="data-space")
193
+ person.cars.clear() # Gonna remove all car edges from the person
194
+
195
+ engine.upsert([person])
196
+
164
197
  ```
165
198
 
199
+
166
200
  ---
@@ -4,9 +4,11 @@ from .models import (
4
4
  InstanceId,
5
5
  PaginatedResult,
6
6
  TViewInstance,
7
+ TWritableViewInstance,
7
8
  ValidationMode,
8
9
  ViewInstance,
9
10
  ViewInstanceConfig,
11
+ WritableViewInstance,
10
12
  )
11
13
  from .statements import and_, col, not_, or_, select
12
14
 
@@ -20,9 +22,11 @@ __all__ = [
20
22
  "InstanceId",
21
23
  "TViewInstance",
22
24
  "DataModelId",
25
+ "TWritableViewInstance",
23
26
  "ValidationMode",
24
27
  "Engine",
25
28
  "AsyncEngine",
26
29
  "PaginatedResult",
27
30
  "ViewInstanceConfig",
31
+ "WritableViewInstance",
28
32
  ]
@@ -1,3 +1,4 @@
1
+ import logging
1
2
  from typing import Any
2
3
 
3
4
  from cognite.client import CogniteClient
@@ -9,15 +10,16 @@ from cognite.client.data_classes.data_modeling.query import (
9
10
  QueryResult as CogniteQueryResult,
10
11
  )
11
12
 
12
- from industrial_model.cognite_adapters.optimizer import QueryOptimizer
13
13
  from industrial_model.config import DataModelId
14
- from industrial_model.models import TViewInstance
14
+ from industrial_model.models import TViewInstance, TWritableViewInstance
15
15
  from industrial_model.statements import Statement
16
16
 
17
+ from .optimizer import QueryOptimizer
17
18
  from .query_mapper import QueryMapper
18
19
  from .query_result_mapper import (
19
20
  QueryResultMapper,
20
21
  )
22
+ from .upsert_mapper import UpsertMapper
21
23
  from .utils import (
22
24
  append_nodes_and_edges,
23
25
  get_query_for_dependencies_pagination,
@@ -39,6 +41,7 @@ class CogniteAdapter:
39
41
  view_mapper = ViewMapper(dm.views)
40
42
  self._query_mapper = QueryMapper(view_mapper)
41
43
  self._result_mapper = QueryResultMapper(view_mapper)
44
+ self._upsert_mapper = UpsertMapper(view_mapper)
42
45
  self._optmizer = QueryOptimizer(cognite_client)
43
46
 
44
47
  def query(
@@ -77,6 +80,36 @@ class CogniteAdapter:
77
80
  if not all_pages or last_page:
78
81
  return data, next_cursor_
79
82
 
83
+ def upsert(
84
+ self, entries: list[TWritableViewInstance], replace: bool = False
85
+ ) -> None:
86
+ logger = logging.getLogger(__name__)
87
+ operation = self._upsert_mapper.map(entries)
88
+
89
+ for node_chunk in operation.chunk_nodes():
90
+ logger.info(
91
+ f"Upserting {len(node_chunk)} nodes (replace={replace})"
92
+ )
93
+ self._cognite_client.data_modeling.instances.apply(
94
+ nodes=node_chunk,
95
+ replace=replace,
96
+ )
97
+
98
+ for edge_chunk in operation.chunk_edges():
99
+ logger.info(
100
+ f"Upserting {len(edge_chunk)} edges (replace={replace})"
101
+ )
102
+ self._cognite_client.data_modeling.instances.apply(
103
+ edges=edge_chunk,
104
+ replace=replace,
105
+ )
106
+
107
+ for edges_to_remove_chunk in operation.chunk_edges_to_delete():
108
+ logger.info(f"Deleting {len(edges_to_remove_chunk)} edges")
109
+ self._cognite_client.data_modeling.instances.delete(
110
+ edges=[item.as_tuple() for item in edges_to_remove_chunk],
111
+ )
112
+
80
113
  def _query_dependencies_pages(
81
114
  self,
82
115
  cognite_query: CogniteQuery,
@@ -0,0 +1,35 @@
1
+ from dataclasses import dataclass
2
+ from typing import Any
3
+
4
+ from cognite.client.data_classes.data_modeling import (
5
+ EdgeApply,
6
+ NodeApply,
7
+ )
8
+
9
+ from industrial_model.models.entities import EdgeContainer
10
+
11
+ _PAGE_SIZE = 1000
12
+
13
+
14
+ @dataclass
15
+ class UpsertOperation:
16
+ nodes: list[NodeApply]
17
+ edges: list[EdgeApply]
18
+ edges_to_delete: list[EdgeContainer]
19
+
20
+ def chunk_nodes(self) -> list[list[NodeApply]]:
21
+ return self._chunk_list(self.nodes)
22
+
23
+ def chunk_edges(self) -> list[list[EdgeApply]]:
24
+ return self._chunk_list(self.edges)
25
+
26
+ def chunk_edges_to_delete(self) -> list[list[EdgeContainer]]:
27
+ return self._chunk_list(self.edges_to_delete)
28
+
29
+ def _chunk_list(self, entries: list[Any]) -> list[list[Any]]:
30
+ data: list[list[Any]] = []
31
+ for i in range(0, len(entries), _PAGE_SIZE):
32
+ start = i
33
+ end = i + _PAGE_SIZE
34
+ data.append(entries[start:end])
35
+ return data
@@ -1,5 +1,5 @@
1
1
  from collections import defaultdict
2
- from typing import Any
2
+ from typing import Any, TypedDict
3
3
 
4
4
  from cognite.client.data_classes.data_modeling import (
5
5
  Edge,
@@ -14,10 +14,17 @@ from cognite.client.data_classes.data_modeling.views import (
14
14
  )
15
15
 
16
16
  from industrial_model.constants import EDGE_DIRECTION, EDGE_MARKER, NESTED_SEP
17
+ from industrial_model.models import EdgeContainer
17
18
 
18
19
  from .view_mapper import ViewMapper
19
20
 
20
21
 
22
+ class _PropertyMapping(TypedDict):
23
+ is_list: bool
24
+ nodes: dict[tuple[str, str], list[Node]]
25
+ edges: dict[tuple[str, str], list[Edge]]
26
+
27
+
21
28
  class QueryResultMapper:
22
29
  def __init__(self, view_mapper: ViewMapper):
23
30
  self._view_mapper = view_mapper
@@ -79,25 +86,36 @@ class QueryResultMapper:
79
86
 
80
87
  visited.add(identify)
81
88
  properties = node.properties.get(view_id, {})
89
+
90
+ edges_mapping: dict[str, list[EdgeContainer]] = {}
82
91
  node_id = get_node_id(node)
83
- for mapping_key, (values, is_list) in mappings.items():
92
+ for mapping_key, mapping_value in mappings.items():
84
93
  element = properties.get(mapping_key)
85
94
 
86
95
  element_key: tuple[str, str] = (
87
96
  (element.get("space", ""), element.get("externalId", ""))
88
97
  if isinstance(element, dict)
89
- else node_id
98
+ else (node.space, node.external_id)
90
99
  )
91
100
 
92
- entries = values.get(element_key)
93
- if not entries:
101
+ mapping_nodes = mapping_value.get("nodes", {})
102
+ mapping_edges = mapping_value.get("edges", {})
103
+ is_list = mapping_value.get("is_list", False)
104
+
105
+ node_entries = mapping_nodes.get(element_key)
106
+ if not node_entries:
94
107
  continue
95
108
 
96
- entry_data = self._nodes_to_dict(entries)
109
+ entry_data = self._nodes_to_dict(node_entries)
97
110
  properties[mapping_key] = (
98
111
  entry_data if is_list else entry_data[0]
99
112
  )
100
-
113
+ edge_entries = mapping_edges.get(element_key)
114
+ if edge_entries:
115
+ edges_mapping[mapping_key] = self._edges_to_model(
116
+ edge_entries
117
+ )
118
+ properties["_edges"] = edges_mapping
101
119
  node.properties[view_id] = properties
102
120
 
103
121
  result[node_id].append(node)
@@ -109,19 +127,18 @@ class QueryResultMapper:
109
127
  key: str,
110
128
  view: View,
111
129
  query_result: dict[str, list[Node | Edge]],
112
- ) -> dict[str, tuple[dict[tuple[str, str], list[Node]], bool]]:
113
- mappings: dict[
114
- str, tuple[dict[tuple[str, str], list[Node]], bool]
115
- ] = {}
130
+ ) -> dict[str, _PropertyMapping]:
131
+ mappings: dict[str, _PropertyMapping] = {}
116
132
 
117
133
  for property_name, property in view.properties.items():
118
134
  property_key = f"{key}{NESTED_SEP}{property_name}"
119
135
 
120
- entry: dict[tuple[str, str], list[Node]] | None = None
136
+ nodes: dict[tuple[str, str], list[Node]] | None = None
137
+ edges: dict[tuple[str, str], list[Edge]] | None = None
121
138
  is_list = False
122
139
 
123
140
  if isinstance(property, MappedProperty) and property.source:
124
- entry = self._map_node_property(
141
+ nodes = self._map_node_property(
125
142
  property_key,
126
143
  self._view_mapper.get_view(property.source.external_id),
127
144
  query_result,
@@ -131,7 +148,7 @@ class QueryResultMapper:
131
148
  isinstance(property, SingleReverseDirectRelation)
132
149
  and property.source
133
150
  ):
134
- entry = self._map_node_property(
151
+ nodes = self._map_node_property(
135
152
  property_key,
136
153
  self._view_mapper.get_view(property.source.external_id),
137
154
  query_result,
@@ -142,7 +159,7 @@ class QueryResultMapper:
142
159
  isinstance(property, MultiReverseDirectRelation)
143
160
  and property.source
144
161
  ):
145
- entry = self._map_node_property(
162
+ nodes = self._map_node_property(
146
163
  property_key,
147
164
  self._view_mapper.get_view(property.source.external_id),
148
165
  query_result,
@@ -151,7 +168,7 @@ class QueryResultMapper:
151
168
  is_list = True
152
169
 
153
170
  elif isinstance(property, EdgeConnection) and property.source:
154
- entry = self._map_edge_property(
171
+ nodes, edges = self._map_edge_property(
155
172
  property_key,
156
173
  self._view_mapper.get_view(property.source.external_id),
157
174
  query_result,
@@ -159,8 +176,10 @@ class QueryResultMapper:
159
176
  )
160
177
  is_list = True
161
178
 
162
- if entry:
163
- mappings[property_name] = entry, is_list
179
+ if nodes:
180
+ mappings[property_name] = _PropertyMapping(
181
+ is_list=is_list, nodes=nodes, edges=edges or {}
182
+ )
164
183
 
165
184
  return mappings
166
185
 
@@ -170,17 +189,25 @@ class QueryResultMapper:
170
189
  view: View,
171
190
  query_result: dict[str, list[Node | Edge]],
172
191
  edge_direction: EDGE_DIRECTION,
173
- ) -> dict[tuple[str, str], list[Node]] | None:
192
+ ) -> tuple[
193
+ dict[tuple[str, str], list[Node]] | None,
194
+ dict[tuple[str, str], list[Edge]] | None,
195
+ ]:
174
196
  edge_key = f"{key}{NESTED_SEP}{EDGE_MARKER}"
175
197
  if key not in query_result or edge_key not in query_result:
176
- return None
198
+ return None, None
177
199
 
178
200
  nodes = self._map_node_property(key, view, query_result)
179
201
  if not nodes:
180
- return None
202
+ return None, None
181
203
 
182
204
  visited: set[tuple[str, str]] = set()
183
- result: defaultdict[tuple[str, str], list[Node]] = defaultdict(list)
205
+ nodes_result: defaultdict[tuple[str, str], list[Node]] = defaultdict(
206
+ list
207
+ )
208
+ edges_result: defaultdict[tuple[str, str], list[Edge]] = defaultdict(
209
+ list
210
+ )
184
211
  for edge in query_result[edge_key]:
185
212
  identify = (edge.space, edge.external_id)
186
213
  if not isinstance(edge, Edge) or identify in visited:
@@ -195,15 +222,18 @@ class QueryResultMapper:
195
222
  if edge_direction == "inwards"
196
223
  else (edge.start_node.as_tuple(), edge.end_node.as_tuple())
197
224
  )
198
-
225
+ edges_result[entry_key].append(edge)
199
226
  if node_item := nodes.get(node_key):
200
- result[entry_key].extend(node_item)
227
+ nodes_result[entry_key].extend(node_item)
201
228
 
202
- return dict(result)
229
+ return dict(nodes_result), dict(edges_result)
203
230
 
204
231
  def _nodes_to_dict(self, nodes: list[Node]) -> list[dict[str, Any]]:
205
232
  return [self._node_to_dict(node) for node in nodes]
206
233
 
234
+ def _edges_to_model(self, edges: list[Edge]) -> list[EdgeContainer]:
235
+ return [EdgeContainer.model_validate(edge) for edge in edges]
236
+
207
237
  def _node_to_dict(self, node: Node) -> dict[str, Any]:
208
238
  entry = node.dump()
209
239
  properties: dict[str, dict[str, dict[str, Any]]] = (
@@ -0,0 +1,146 @@
1
+ import datetime
2
+ from typing import Any
3
+
4
+ from cognite.client.data_classes.data_modeling import (
5
+ DirectRelationReference,
6
+ EdgeApply,
7
+ EdgeConnection,
8
+ MappedProperty,
9
+ NodeApply,
10
+ NodeOrEdgeData,
11
+ )
12
+
13
+ from industrial_model.cognite_adapters.models import UpsertOperation
14
+ from industrial_model.models import (
15
+ EdgeContainer,
16
+ InstanceId,
17
+ TWritableViewInstance,
18
+ )
19
+ from industrial_model.utils import datetime_to_ms_iso_timestamp
20
+
21
+ from .view_mapper import ViewMapper
22
+
23
+
24
+ class UpsertMapper:
25
+ def __init__(self, view_mapper: ViewMapper):
26
+ self._view_mapper = view_mapper
27
+
28
+ def map(self, instances: list[TWritableViewInstance]) -> UpsertOperation:
29
+ nodes: dict[tuple[str, str], NodeApply] = {}
30
+ edges: dict[tuple[str, str], EdgeApply] = {}
31
+ edges_to_delete: dict[tuple[str, str], EdgeContainer] = {}
32
+
33
+ for instance in instances:
34
+ entry_nodes, entry_edges, entry_edges_to_delete = (
35
+ self._map_instance(instance)
36
+ )
37
+
38
+ nodes[instance.as_tuple()] = entry_nodes
39
+ edges.update(
40
+ {(item.space, item.external_id): item for item in entry_edges}
41
+ )
42
+ edges_to_delete.update(
43
+ {
44
+ (item.space, item.external_id): item
45
+ for item in entry_edges_to_delete
46
+ }
47
+ )
48
+
49
+ return UpsertOperation(
50
+ nodes=list(nodes.values()),
51
+ edges=list(edges.values()),
52
+ edges_to_delete=list(edges_to_delete.values()),
53
+ )
54
+
55
+ def _map_instance(
56
+ self, instance: TWritableViewInstance
57
+ ) -> tuple[NodeApply, list[EdgeApply], list[EdgeContainer]]:
58
+ view = self._view_mapper.get_view(instance.get_view_external_id())
59
+
60
+ edges: list[EdgeApply] = []
61
+ edges_to_delete: list[EdgeContainer] = []
62
+ properties: dict[str, Any] = {}
63
+ for property_name, property in view.properties.items():
64
+ property_key = instance.get_field_name(property_name)
65
+ if not property_key:
66
+ continue
67
+ entry = instance.__getattribute__(property_key)
68
+
69
+ if isinstance(property, MappedProperty):
70
+ properties[property_name] = (
71
+ DirectRelationReference(
72
+ space=entry.space, external_id=entry.external_id
73
+ )
74
+ if isinstance(entry, InstanceId)
75
+ else datetime_to_ms_iso_timestamp(entry)
76
+ if isinstance(entry, datetime.datetime)
77
+ else entry
78
+ )
79
+ elif isinstance(property, EdgeConnection) and isinstance(
80
+ entry, list
81
+ ):
82
+ possible_entries = self._map_edges(instance, property, entry)
83
+
84
+ previous_edges = {
85
+ item.as_tuple(): item
86
+ for item in instance._edges.get(property_name, [])
87
+ }
88
+
89
+ new_entries = [
90
+ edge
91
+ for edge_id, edge in possible_entries.items()
92
+ if edge_id not in previous_edges
93
+ ]
94
+ edges_to_delete.extend(
95
+ [
96
+ previous_edges[edge_id]
97
+ for edge_id in previous_edges
98
+ if edge_id not in possible_entries
99
+ ]
100
+ )
101
+
102
+ edges.extend(new_entries)
103
+
104
+ node = NodeApply(
105
+ external_id=instance.external_id,
106
+ space=instance.space,
107
+ sources=[
108
+ NodeOrEdgeData(source=view.as_id(), properties=properties)
109
+ ],
110
+ )
111
+
112
+ return node, edges, edges_to_delete
113
+
114
+ def _map_edges(
115
+ self,
116
+ instance: TWritableViewInstance,
117
+ property: EdgeConnection,
118
+ values: list[Any],
119
+ ) -> dict[tuple[str, str], EdgeApply]:
120
+ edge_type = InstanceId.model_validate(property.type)
121
+
122
+ result: dict[tuple[str, str], EdgeApply] = {}
123
+ for value in values:
124
+ if not isinstance(value, InstanceId):
125
+ raise ValueError(
126
+ f"""Invalid value for edge property {property.name}:
127
+ Received {type(value)} | Expected: InstanceId"""
128
+ )
129
+
130
+ start_node, end_node = (
131
+ (instance, value)
132
+ if property.direction == "outwards"
133
+ else (value, instance)
134
+ )
135
+
136
+ edge_id = instance.edge_id_factory(value, edge_type)
137
+
138
+ result[edge_id.as_tuple()] = EdgeApply(
139
+ external_id=edge_id.external_id,
140
+ space=edge_id.space,
141
+ type=property.type,
142
+ start_node=start_node.as_tuple(),
143
+ end_node=end_node.as_tuple(),
144
+ )
145
+
146
+ return result
@@ -6,6 +6,7 @@ from industrial_model.models import (
6
6
  TViewInstance,
7
7
  ValidationMode,
8
8
  )
9
+ from industrial_model.models.entities import TWritableViewInstance
9
10
  from industrial_model.statements import Statement
10
11
  from industrial_model.utils import run_async
11
12
 
@@ -35,3 +36,8 @@ class AsyncEngine:
35
36
  return await run_async(
36
37
  self._engine.query_all_pages, statement, validation_mode
37
38
  )
39
+
40
+ async def upsert_async(
41
+ self, entries: list[TWritableViewInstance], replace: bool = False
42
+ ) -> None:
43
+ return await run_async(self._engine.upsert, entries, replace)
@@ -9,6 +9,10 @@ from industrial_model.models import (
9
9
  TViewInstance,
10
10
  ValidationMode,
11
11
  )
12
+ from industrial_model.models.entities import (
13
+ EdgeContainer,
14
+ TWritableViewInstance,
15
+ )
12
16
  from industrial_model.statements import Statement
13
17
 
14
18
 
@@ -45,6 +49,14 @@ class Engine:
45
49
 
46
50
  return self._validate_data(statement.entity, data, validation_mode)
47
51
 
52
+ def upsert(
53
+ self, entries: list[TWritableViewInstance], replace: bool = False
54
+ ) -> None:
55
+ if not entries:
56
+ return
57
+
58
+ return self._cognite_adapter.upsert(entries, replace)
59
+
48
60
  def _validate_data(
49
61
  self,
50
62
  entity: type[TViewInstance],
@@ -54,9 +66,25 @@ class Engine:
54
66
  result: list[TViewInstance] = []
55
67
  for item in data:
56
68
  try:
57
- result.append(entity.model_validate(item))
69
+ validated_item = entity.model_validate(item)
70
+ self._include_edges(item, validated_item)
71
+ result.append(validated_item)
58
72
  except Exception:
59
73
  if validation_mode == "ignoreOnError":
60
74
  continue
61
75
  raise
62
76
  return result
77
+
78
+ def _include_edges(
79
+ self, item: dict[str, Any], validated_item: TViewInstance
80
+ ) -> None:
81
+ if "_edges" not in item or not isinstance(item["_edges"], dict):
82
+ return
83
+ entries: dict[str, list[EdgeContainer]] = {}
84
+ for property_, edges in item["_edges"].items():
85
+ if not edges or not isinstance(edges, list):
86
+ continue
87
+
88
+ assert isinstance(edges[0], EdgeContainer)
89
+ entries[property_] = edges
90
+ validated_item._edges = entries
@@ -1,22 +1,28 @@
1
1
  from .base import RootModel
2
2
  from .entities import (
3
+ EdgeContainer,
3
4
  InstanceId,
4
5
  PaginatedResult,
5
6
  TViewInstance,
7
+ TWritableViewInstance,
6
8
  ValidationMode,
7
9
  ViewInstance,
8
10
  ViewInstanceConfig,
11
+ WritableViewInstance,
9
12
  )
10
13
  from .schemas import get_parent_and_children_nodes, get_schema_properties
11
14
 
12
15
  __all__ = [
13
16
  "RootModel",
17
+ "EdgeContainer",
14
18
  "InstanceId",
15
19
  "TViewInstance",
20
+ "TWritableViewInstance",
16
21
  "ViewInstance",
17
22
  "ValidationMode",
18
23
  "PaginatedResult",
19
24
  "ViewInstanceConfig",
20
25
  "get_schema_properties",
21
26
  "get_parent_and_children_nodes",
27
+ "WritableViewInstance",
22
28
  ]
@@ -1,3 +1,4 @@
1
+ from abc import abstractmethod
1
2
  from typing import (
2
3
  Any,
3
4
  ClassVar,
@@ -7,6 +8,8 @@ from typing import (
7
8
  TypeVar,
8
9
  )
9
10
 
11
+ from pydantic import PrivateAttr
12
+
10
13
  from .base import DBModelMetaclass, RootModel
11
14
 
12
15
 
@@ -29,6 +32,15 @@ class InstanceId(RootModel):
29
32
  return (self.space, self.external_id)
30
33
 
31
34
 
35
+ class EdgeContainer(InstanceId):
36
+ type: InstanceId
37
+ start_node: InstanceId
38
+ end_node: InstanceId
39
+
40
+
41
+ TInstanceId = TypeVar("TInstanceId", bound=InstanceId)
42
+
43
+
32
44
  class ViewInstanceConfig(TypedDict, total=False):
33
45
  view_external_id: str | None
34
46
  instance_spaces: list[str] | None
@@ -38,12 +50,38 @@ class ViewInstanceConfig(TypedDict, total=False):
38
50
  class ViewInstance(InstanceId, metaclass=DBModelMetaclass):
39
51
  view_config: ClassVar[ViewInstanceConfig] = ViewInstanceConfig()
40
52
 
53
+ _edges: dict[str, list[EdgeContainer]] = PrivateAttr(default_factory=dict)
54
+
41
55
  @classmethod
42
56
  def get_view_external_id(cls) -> str:
43
57
  return cls.view_config.get("view_external_id") or cls.__name__
44
58
 
59
+ def get_field_name(self, field_name_or_alias: str) -> str | None:
60
+ entry = self.__class__.model_fields.get(field_name_or_alias)
61
+ if entry:
62
+ return field_name_or_alias
63
+
64
+ for key, field_info in self.__class__.model_fields.items():
65
+ if field_info.alias == field_name_or_alias:
66
+ return key
67
+
68
+ return None
69
+
70
+
71
+ class WritableViewInstance(ViewInstance):
72
+ @abstractmethod
73
+ def edge_id_factory(
74
+ self, target_node: TInstanceId, edge_type: InstanceId
75
+ ) -> InstanceId:
76
+ raise NotImplementedError(
77
+ "edge_id_factory method must be implemented in subclasses"
78
+ )
79
+
45
80
 
46
81
  TViewInstance = TypeVar("TViewInstance", bound=ViewInstance)
82
+ TWritableViewInstance = TypeVar(
83
+ "TWritableViewInstance", bound=WritableViewInstance
84
+ )
47
85
 
48
86
 
49
87
  class PaginatedResult(RootModel, Generic[TViewInstance]):
@@ -1,14 +1,12 @@
1
- from collections.abc import Callable, Generator
1
+ from collections.abc import Callable
2
2
  from datetime import datetime
3
3
  from typing import (
4
- Any,
5
4
  ParamSpec,
6
5
  TypeVar,
7
6
  )
8
7
 
9
8
  from anyio import to_thread
10
9
 
11
- TAny = TypeVar("TAny")
12
10
  T_Retval = TypeVar("T_Retval")
13
11
  P = ParamSpec("P")
14
12
 
@@ -21,15 +19,6 @@ def datetime_to_ms_iso_timestamp(dt: datetime) -> str:
21
19
  return dt.isoformat(timespec="milliseconds")
22
20
 
23
21
 
24
- def chunk_list(
25
- entries: list[TAny], chunk_size: int
26
- ) -> Generator[list[TAny], Any, None]:
27
- for i in range(0, len(entries), chunk_size):
28
- start = i
29
- end = i + chunk_size
30
- yield entries[start:end]
31
-
32
-
33
22
  async def run_async(
34
23
  func: Callable[..., T_Retval],
35
24
  *args: object,
@@ -1,6 +1,6 @@
1
1
  [project]
2
2
  name = "industrial-model"
3
- version = "0.1.8"
3
+ version = "0.1.10"
4
4
  description = "Industrial Model ORM"
5
5
  readme = "README.md"
6
6
  requires-python = ">=3.11"
@@ -0,0 +1,2 @@
1
+ ruff check industrial_model tests --fix
2
+ ruff format industrial_model tests
@@ -0,0 +1,3 @@
1
+ mypy industrial_model tests
2
+ ruff check industrial_model tests
3
+ ruff format industrial_model tests --check
@@ -1,12 +1,14 @@
1
1
  import datetime
2
- import json
3
2
  from typing import Annotated
4
3
 
5
4
  from pydantic import Field
6
5
 
7
- from industrial_model import ViewInstance, ViewInstanceConfig, col, select
8
-
9
- from .hubs import generate_engine
6
+ from industrial_model import (
7
+ InstanceId,
8
+ ViewInstance,
9
+ ViewInstanceConfig,
10
+ WritableViewInstance,
11
+ )
10
12
 
11
13
 
12
14
  class DescribableEntity(ViewInstance):
@@ -32,7 +34,8 @@ class EventDetail(ViewInstance):
32
34
 
33
35
  class Event(ViewInstance):
34
36
  view_config = ViewInstanceConfig(
35
- view_external_id="OEEEvent", instance_spaces_prefix="OEE-"
37
+ view_external_id="OEEEvent",
38
+ instance_spaces_prefix="OEE-",
36
39
  )
37
40
 
38
41
  start_date_time: datetime.datetime | None = None
@@ -49,22 +52,24 @@ class Event(ViewInstance):
49
52
  ]
50
53
 
51
54
 
52
- adapter = generate_engine()
53
-
54
- filter = (
55
- col(Event.start_date_time).gt_(datetime.datetime(2025, 3, 1))
56
- & col(Event.ref_site).nested_(DescribableEntity.external_id == "STS-CLK")
57
- & (Event.start_date_time < datetime.datetime(2025, 6, 1))
58
- )
59
-
55
+ class WritableEvent(WritableViewInstance):
56
+ view_config = ViewInstanceConfig(
57
+ view_external_id="OEEEvent",
58
+ instance_spaces_prefix="OEE-",
59
+ )
60
60
 
61
- statement = select(Event).limit(50).where(filter)
61
+ start_date_time: datetime.datetime | None = None
62
+ ref_site: ReportingSite | None = None
63
+ ref_reporting_line: DescribableEntity | None = None
64
+ ref_oee_event_detail: Annotated[
65
+ list[EventDetail],
66
+ Field(default_factory=list, alias="refOEEEventDetail"),
67
+ ]
62
68
 
63
- result = [
64
- item.model_dump(mode="json") for item in adapter.query_all_pages(statement)
65
- ]
66
- print(len(result))
67
- json.dump(
68
- result,
69
- open("test.json", "w"),
70
- )
69
+ def edge_id_factory(
70
+ self, target_node: InstanceId, edge_type: InstanceId
71
+ ) -> InstanceId:
72
+ return InstanceId(
73
+ external_id=f"{self.external_id}-{target_node.external_id}-{edge_type.external_id}",
74
+ space=self.space,
75
+ )
@@ -0,0 +1,39 @@
1
+ import json
2
+
3
+ from industrial_model.cognite_adapters.upsert_mapper import UpsertMapper
4
+ from industrial_model.statements import select
5
+ from tests.hubs import generate_engine
6
+ from tests.models import EventDetail, WritableEvent
7
+
8
+ if __name__ == "__main__":
9
+ engine = generate_engine()
10
+
11
+ view_mapper = engine._cognite_adapter._query_mapper._view_mapper
12
+
13
+ upsert_mapper = UpsertMapper(view_mapper)
14
+
15
+ statement = select(WritableEvent).limit(1)
16
+
17
+ item = engine.query(statement).data[0]
18
+
19
+ item.ref_oee_event_detail.clear()
20
+ item.ref_oee_event_detail.append(
21
+ EventDetail(external_id="test", space="test")
22
+ )
23
+
24
+ operation = upsert_mapper.map([item])
25
+
26
+ data = {
27
+ "nodes": [node.dump() for node in operation.nodes],
28
+ "edges": [edge.dump() for edge in operation.edges],
29
+ "edges_to_delete": [
30
+ edge_to_delete.model_dump(mode="json")
31
+ for edge_to_delete in operation.edges_to_delete
32
+ ],
33
+ }
34
+
35
+ json.dump(
36
+ data,
37
+ open("upsert.json", "w"),
38
+ indent=4,
39
+ )
@@ -0,0 +1,27 @@
1
+ import datetime
2
+ import json
3
+
4
+ from industrial_model import col, select
5
+
6
+ from .hubs import generate_engine
7
+ from .models import DescribableEntity, Event
8
+
9
+ if __name__ == "__main__":
10
+ adapter = generate_engine()
11
+
12
+ filter = (
13
+ col(Event.start_date_time).gt_(datetime.datetime(2025, 3, 1))
14
+ & col(Event.ref_site).nested_(
15
+ DescribableEntity.external_id == "STS-CLK"
16
+ )
17
+ & (col(Event.start_date_time) < datetime.datetime(2025, 6, 1))
18
+ )
19
+
20
+ statement = select(Event).limit(50).where(filter)
21
+
22
+ result = [
23
+ item.model_dump(mode="json")
24
+ for item in adapter.query_all_pages(statement)
25
+ ]
26
+ print(len(result))
27
+ json.dump(result, open("entities.json", "w"), indent=2)
@@ -215,7 +215,7 @@ wheels = [
215
215
 
216
216
  [[package]]
217
217
  name = "industrial-model"
218
- version = "0.1.8"
218
+ version = "0.1.10"
219
219
  source = { editable = "." }
220
220
  dependencies = [
221
221
  { name = "anyio" },
@@ -1,2 +0,0 @@
1
- ruff check industrial_model --fix
2
- ruff format industrial_model
@@ -1,3 +0,0 @@
1
- mypy industrial_model
2
- ruff check industrial_model
3
- ruff format industrial_model --check