Skip to main content

linkrs_storage/engine/sync_wrapper/
write.rs

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    /// Sync failure branch: the writer already discarded its write scope and
9    /// aborted the timestamp before returning the error; aborting the
10    /// transaction here runs the undo log over already applied rows.
11    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}