1use super::SyncWrapper;
2use crate::{StorageClient, StorageWriter};
3use linkrs_core::types::{InsertEdgeInfo, InsertVertexInfo, UpdateInfo, UpdateOp, VertexId};
4use linkrs_core::{Edge, EdgeDeleteKey, StorageError, Value, Vertex};
5use linkrs_sync::types::ChangeType;
6
7impl<S: StorageClient + 'static> SyncWrapper<S> {
8 fn reject_staged_write(&self, error: StorageError) -> StorageError {
12 if let Some(transaction_id) = self.get_current_txn_id() {
13 let _ = self.abort_transaction_fact(transaction_id);
14 }
15 error
16 }
17
18 fn stage_vertex_data_change(
19 &self,
20 info: &InsertVertexInfo,
21 change_type: ChangeType,
22 ) -> Result<(), StorageError> {
23 if !self.enabled {
24 return Ok(());
25 }
26 let Some(manager) = self.get_sync_manager() else {
27 return Ok(());
28 };
29 let transaction_id = self.get_current_txn_id().ok_or_else(|| {
30 StorageError::db_error(
31 "Synchronized writes require an operation transaction context".to_string(),
32 )
33 })?;
34 manager
35 .on_vertex_change_with_txn(
36 transaction_id,
37 info.space_id,
38 &info.tag_name,
39 &info.vertex_id,
40 &info.props,
41 change_type,
42 )
43 .map_err(|error| StorageError::db_error(error.to_string()))
44 }
45}
46
47impl<S: StorageClient + 'static> StorageWriter for SyncWrapper<S> {
48 fn insert_vertex(&mut self, space: &str, vertex: Vertex) -> Result<VertexId, StorageError> {
49 let result = self.inner.insert_vertex(space, vertex.clone())?;
50 if let Err(error) = self.sync_insert_vertex(space, &vertex) {
51 return Err(self.reject_staged_write(error));
52 }
53 self.commit_auto_transaction()?;
54 Ok(result)
55 }
56
57 fn update_vertex(&mut self, space: &str, vertex: Vertex) -> Result<(), StorageError> {
58 let old_vertex = self
59 .inner
60 .get_vertex(space, &vertex.tag.name, &vertex.vid)?
61 .ok_or_else(|| StorageError::node_not_found(vertex.vid))?;
62
63 self.inner.update_vertex(space, vertex.clone())?;
64 if let Err(error) = self.sync_update_vertex(space, &old_vertex, &vertex) {
65 return Err(self.reject_staged_write(error));
66 }
67 self.commit_auto_transaction()?;
68 Ok(())
69 }
70
71 fn update_vertex_replace(&mut self, space: &str, vertex: Vertex) -> Result<(), StorageError> {
72 let old_vertex = self
73 .inner
74 .get_vertex(space, &vertex.tag.name, &vertex.vid)?
75 .ok_or_else(|| StorageError::node_not_found(vertex.vid))?;
76 self.inner.update_vertex_replace(space, vertex.clone())?;
77 if let Err(error) = self.sync_replace_vertex(space, &old_vertex, &vertex) {
78 return Err(self.reject_staged_write(error));
79 }
80 self.commit_auto_transaction()?;
81 Ok(())
82 }
83
84 fn delete_vertex(&mut self, space: &str, tag: &str, id: &VertexId) -> Result<(), StorageError> {
85 let vertex = self
86 .inner
87 .get_vertex(space, tag, id)?
88 .ok_or_else(|| StorageError::node_not_found(*id))?;
89
90 StorageWriter::delete_vertex(&mut self.inner, space, tag, id)?;
91 if let Err(error) = self.sync_delete_vertex(space, &vertex.vid, &vertex) {
92 return Err(self.reject_staged_write(error));
93 }
94 self.commit_auto_transaction()?;
95 Ok(())
96 }
97
98 fn delete_vertex_with_edges(
99 &mut self,
100 space: &str,
101 tag: &str,
102 id: &VertexId,
103 ) -> Result<(), StorageError> {
104 let vertex = self
105 .inner
106 .get_vertex(space, tag, id)?
107 .ok_or_else(|| StorageError::node_not_found(*id))?;
108
109 StorageWriter::delete_vertex_with_edges(&mut self.inner, space, tag, id)?;
110 if let Err(error) = self.sync_delete_vertex(space, id, &vertex) {
111 return Err(self.reject_staged_write(error));
112 }
113 self.commit_auto_transaction()?;
114 Ok(())
115 }
116
117 fn batch_delete_vertices_with_edges(
118 &mut self,
119 space: &str,
120 tag: &str,
121 ids: &[VertexId],
122 ) -> Result<usize, StorageError> {
123 let mut vertices: Vec<Vertex> = Vec::with_capacity(ids.len());
124 for id in ids {
125 if let Some(v) = self.inner.get_vertex(space, tag, id)? {
126 vertices.push(v);
127 }
128 }
129 let result =
130 StorageWriter::batch_delete_vertices_with_edges(&mut self.inner, space, tag, ids)?;
131 for (id, vertex) in ids.iter().zip(vertices) {
132 if let Err(error) = self.sync_delete_vertex(space, id, &vertex) {
133 return Err(self.reject_staged_write(error));
134 }
135 }
136 self.commit_auto_transaction()?;
137 Ok(result)
138 }
139
140 fn batch_insert_vertices(
141 &mut self,
142 space: &str,
143 vertices: Vec<Vertex>,
144 ) -> Result<Vec<VertexId>, StorageError> {
145 let results = self.inner.batch_insert_vertices(space, vertices.clone())?;
146 if let Err(error) = self.sync_batch_insert_vertices(space, &vertices) {
147 return Err(self.reject_staged_write(error));
148 }
149 self.commit_auto_transaction()?;
150 Ok(results)
151 }
152
153 fn insert_edge(&mut self, space: &str, edge: Edge) -> Result<(), StorageError> {
154 let result = self.inner.insert_edge(space, edge.clone());
155 if result.is_ok() {
156 if let Err(error) = self.sync_insert_edge(space, &edge) {
157 return Err(self.reject_staged_write(error));
158 }
159 }
160 result?;
161 self.commit_auto_transaction()?;
162 Ok(())
163 }
164
165 fn update_edge(&mut self, space: &str, edge: Edge) -> Result<(), StorageError> {
166 let result = self.inner.update_edge(space, edge.clone());
167 if result.is_ok() {
168 if let Err(error) =
169 self.sync_delete_edge(space, &edge.src, &edge.dst, &edge.edge_type, edge.ranking)
170 {
171 return Err(self.reject_staged_write(error));
172 }
173 if let Err(error) = self.sync_insert_edge(space, &edge) {
174 return Err(self.reject_staged_write(error));
175 }
176 }
177 result?;
178 self.commit_auto_transaction()?;
179 Ok(())
180 }
181
182 fn update_edge_replace(&mut self, space: &str, edge: Edge) -> Result<(), StorageError> {
183 let result = self.inner.update_edge_replace(space, edge.clone());
184 if result.is_ok() {
185 if let Err(error) =
186 self.sync_delete_edge(space, &edge.src, &edge.dst, &edge.edge_type, edge.ranking)
187 {
188 return Err(self.reject_staged_write(error));
189 }
190 if let Err(error) = self.sync_insert_edge(space, &edge) {
191 return Err(self.reject_staged_write(error));
192 }
193 }
194 result?;
195 self.commit_auto_transaction()?;
196 Ok(())
197 }
198
199 fn delete_edge(
200 &mut self,
201 space: &str,
202 src: &VertexId,
203 dst: &VertexId,
204 edge_type: &str,
205 rank: i64,
206 ) -> Result<(), StorageError> {
207 let result = StorageWriter::delete_edge(&mut self.inner, space, src, dst, edge_type, rank);
208 if result.is_ok() {
209 if let Err(error) = self.sync_delete_edge(space, src, dst, edge_type, rank) {
210 return Err(self.reject_staged_write(error));
211 }
212 }
213 result?;
214 self.commit_auto_transaction()?;
215 Ok(())
216 }
217
218 fn batch_insert_edges(&mut self, space: &str, edges: Vec<Edge>) -> Result<(), StorageError> {
219 let result = self.inner.batch_insert_edges(space, edges.clone());
220 if result.is_ok() {
221 if let Err(error) = self.sync_batch_insert_edges(space, &edges) {
222 return Err(self.reject_staged_write(error));
223 }
224 }
225 result?;
226 self.commit_auto_transaction()?;
227 Ok(())
228 }
229
230 fn batch_delete_edges(
231 &mut self,
232 space: &str,
233 deletes: &[EdgeDeleteKey],
234 ) -> Result<usize, StorageError> {
235 let result = self.inner.batch_delete_edges(space, deletes);
236 if result.is_ok() {
237 if let Err(error) = self.sync_batch_delete_edges(space, deletes) {
238 return Err(self.reject_staged_write(error));
239 }
240 }
241 let deleted = result?;
242 self.commit_auto_transaction()?;
243 Ok(deleted)
244 }
245
246 fn insert_vertex_data(
247 &mut self,
248 space: &str,
249 info: &InsertVertexInfo,
250 ) -> Result<bool, StorageError> {
251 let result = self.inner.insert_vertex_data(space, info)?;
252 if result {
253 if let Err(error) = self.stage_vertex_data_change(info, ChangeType::Insert) {
254 return Err(self.reject_staged_write(error));
255 }
256 }
257 self.commit_auto_transaction()?;
258 Ok(result)
259 }
260
261 fn delete_vertex_data(
262 &mut self,
263 space: &str,
264 tag: &str,
265 vertex_id: &str,
266 ) -> Result<bool, StorageError> {
267 let vid_type = self
268 .inner
269 .get_space(space)?
270 .ok_or_else(|| StorageError::not_found(format!("Space {} not found", space)))?
271 .vid_type;
272 let raw = if let Ok(parsed) = vertex_id.parse::<i64>() {
273 VertexId::try_from_int64(parsed)
274 .or_else(|_| VertexId::try_from_string(vertex_id))
275 .map_err(StorageError::invalid_input)?
276 } else {
277 VertexId::try_from_string(vertex_id).map_err(StorageError::invalid_input)?
278 };
279 let parsed_id = VertexId::normalize_for_vid_type(&vid_type, raw)?;
280 let previous = self.inner.get_vertex(space, tag, &parsed_id)?;
281 let result = self.inner.delete_vertex_data(space, tag, vertex_id)?;
282 if result {
283 if let Some(vertex) = previous {
284 if let Err(error) = self.sync_delete_vertex(space, &vertex.vid, &vertex) {
285 return Err(self.reject_staged_write(error));
286 }
287 }
288 }
289 self.commit_auto_transaction()?;
290 Ok(result)
291 }
292
293 fn insert_edge_data(
294 &mut self,
295 space: &str,
296 info: &InsertEdgeInfo,
297 ) -> Result<bool, StorageError> {
298 let edge = Edge {
299 src: VertexId::try_from(&info.src_vertex_id)
300 .map_err(|error| StorageError::invalid_input(error.to_string()))?,
301 dst: VertexId::try_from(&info.dst_vertex_id)
302 .map_err(|error| StorageError::invalid_input(error.to_string()))?,
303 edge_type: info.edge_name.clone(),
304 ranking: info.rank,
305 props: info.props.iter().cloned().collect(),
306 };
307 let result = self.inner.insert_edge_data(space, info)?;
308 if result {
309 if let Err(error) = self.sync_insert_edge(space, &edge) {
310 return Err(self.reject_staged_write(error));
311 }
312 }
313 self.commit_auto_transaction()?;
314 Ok(result)
315 }
316
317 fn delete_edge_data(
318 &mut self,
319 space: &str,
320 src: &str,
321 dst: &str,
322 rank: i64,
323 ) -> Result<bool, StorageError> {
324 let vid_type = self
325 .inner
326 .get_space(space)?
327 .ok_or_else(|| StorageError::not_found(format!("Space {} not found", space)))?
328 .vid_type;
329 let source = VertexId::normalize_for_vid_type(
330 &vid_type,
331 if let Ok(parsed) = src.parse::<i64>() {
332 VertexId::try_from_int64(parsed)
333 .or_else(|_| VertexId::try_from_string(src))
334 .map_err(StorageError::invalid_input)?
335 } else {
336 VertexId::try_from_string(src).map_err(StorageError::invalid_input)?
337 },
338 )?;
339 let destination = VertexId::normalize_for_vid_type(
340 &vid_type,
341 if let Ok(parsed) = dst.parse::<i64>() {
342 VertexId::try_from_int64(parsed)
343 .or_else(|_| VertexId::try_from_string(dst))
344 .map_err(StorageError::invalid_input)?
345 } else {
346 VertexId::try_from_string(dst).map_err(StorageError::invalid_input)?
347 },
348 )?;
349 let previous = self
350 .inner
351 .scan_all_edges(space)?
352 .into_iter()
353 .filter(|edge| edge.src == source && edge.dst == destination && edge.ranking == rank)
354 .collect::<Vec<_>>();
355 let result = self.inner.delete_edge_data(space, src, dst, rank)?;
356 if result {
357 for edge in previous {
358 if let Err(error) = self.sync_delete_edge(
359 space,
360 &edge.src,
361 &edge.dst,
362 &edge.edge_type,
363 edge.ranking,
364 ) {
365 return Err(self.reject_staged_write(error));
366 }
367 }
368 }
369 self.commit_auto_transaction()?;
370 Ok(result)
371 }
372
373 fn update_data(
374 &mut self,
375 space: &str,
376 space_id: u64,
377 info: &UpdateInfo,
378 ) -> Result<bool, StorageError> {
379 let target = &info.update_target;
380 let vertex_id = VertexId::try_from(&target.id)
381 .map_err(|error| StorageError::invalid_input(error.to_string()))?;
382 let previous = self.inner.get_vertex(space, &target.label, &vertex_id)?;
383 let current = previous
384 .as_ref()
385 .and_then(|vertex| vertex.get_property(&target.prop))
386 .cloned();
387 let updated_value = match (&info.update_op, current) {
388 (UpdateOp::Add, Some(Value::Int(current))) => match &info.value {
389 Value::Int(delta) => Value::Int(current + *delta),
390 _ => info.value.clone(),
391 },
392 (UpdateOp::Subtract, Some(Value::Int(current))) => match &info.value {
393 Value::Int(delta) => Value::Int(current - *delta),
394 _ => info.value.clone(),
395 },
396 _ => info.value.clone(),
397 };
398 let result = self.inner.update_data(space, space_id, info)?;
399 if result && self.enabled && self.get_sync_manager().is_some() {
400 let transaction_id = self
401 .get_current_txn_id()
402 .ok_or_else(|| {
403 StorageError::db_error(
404 "Synchronized writes require an operation transaction context".to_string(),
405 )
406 })
407 .map_err(|error| self.reject_staged_write(error))?;
408 if let Some(manager) = self.get_sync_manager() {
409 if let Err(error) = manager.on_vertex_change_with_txn(
410 transaction_id,
411 space_id,
412 &target.label,
413 &target.id,
414 &[(target.prop.clone(), updated_value)],
415 ChangeType::Update,
416 ) {
417 return Err(self.reject_staged_write(StorageError::db_error(error.to_string())));
418 }
419 }
420 }
421 self.commit_auto_transaction()?;
422 Ok(result)
423 }
424}