nodedb_lite/query/document_ops/writes/
batch.rs

1// SPDX-License-Identifier: Apache-2.0
2use std::collections::HashSet;
3
4use super::super::is_strict;
5use super::super::reads::msgpack_bytes_to_crdt_fields;
6use super::super::write_helpers::affected;
7use super::super::write_helpers::{decode_strict_fields, fields_to_values, strict_schema};
8use crate::engine::crdt::{CrdtRowOp, CrdtRowWrite};
9use crate::error::LiteError;
10use crate::query::engine::LiteQueryEngine;
11use crate::query::text_index::{index_strict_rows, reindex_documents};
12use crate::storage::engine::StorageEngine;
13use nodedb_types::result::QueryResult;
14use nodedb_types::value::Value;
15
16/// BatchInsert: insert N documents in a single transaction.
17///
18/// Each row is judged as a point insert judges it. A key already stored, or
19/// one the batch names twice, refuses the whole batch, and nothing is
20/// written. With `if_absent` (`ON CONFLICT DO NOTHING`) such a row is
21/// skipped instead.
22pub async fn batch_insert<S: StorageEngine>(
23    engine: &LiteQueryEngine<S>,
24    collection: &str,
25    documents: &[(String, Vec<u8>)],
26    if_absent: bool,
27) -> Result<QueryResult, LiteError> {
28    batch_insert_coordinated(engine, None, collection, documents, if_absent).await
29}
30
31pub(crate) async fn batch_insert_coordinated<S: StorageEngine>(
32    engine: &LiteQueryEngine<S>,
33    permit: Option<&crate::engine::fts::coordinator::TextMutationPermit>,
34    collection: &str,
35    documents: &[(String, Vec<u8>)],
36    if_absent: bool,
37) -> Result<QueryResult, LiteError> {
38    if let Some(permit) = permit {
39        return batch_insert_admitted(engine, permit, collection, documents, if_absent).await;
40    }
41    let guard = engine.fts_state.admit_mutation().await;
42    let result =
43        batch_insert_admitted(engine, guard.permit(), collection, documents, if_absent).await;
44    guard.finish(result)
45}
46
47pub(crate) async fn batch_insert_admitted<S: StorageEngine>(
48    engine: &LiteQueryEngine<S>,
49    permit: &crate::engine::fts::coordinator::TextMutationPermit,
50    collection: &str,
51    documents: &[(String, Vec<u8>)],
52    if_absent: bool,
53) -> Result<QueryResult, LiteError> {
54    let _permit = permit;
55    if is_strict(engine, collection) {
56        let schema = strict_schema(engine, collection)?;
57        let mut seen: HashSet<&str> = HashSet::with_capacity(documents.len());
58        let mut rows: Vec<Vec<Value>> = Vec::with_capacity(documents.len());
59        for (doc_id, value_bytes) in documents {
60            if if_absent {
61                let repeated = !seen.insert(doc_id.as_str());
62                if repeated
63                    || engine
64                        .strict
65                        .get(collection, &Value::String(doc_id.clone()))
66                        .await?
67                        .is_some()
68                {
69                    continue;
70                }
71            }
72            let fields = decode_strict_fields(value_bytes)?;
73            rows.push(fields_to_values(&fields, &schema.columns));
74        }
75        let affected_n = rows.len() as u64;
76        if !rows.is_empty() {
77            // A primary key that is stored or repeats in `rows` refuses
78            // every row.
79            engine.strict.insert_rows(collection, &rows).await?;
80            index_strict_rows(engine, collection, rows.iter().map(Vec::as_slice))?;
81        }
82        Ok(affected(affected_n, "INSERT"))
83    } else {
84        let mut decoded: Vec<(&str, Vec<(String, loro::LoroValue)>)> =
85            Vec::with_capacity(documents.len());
86        for (doc_id, value_bytes) in documents {
87            let crdt_fields = msgpack_bytes_to_crdt_fields(value_bytes)?;
88            decoded.push((doc_id.as_str(), crdt_fields));
89        }
90        let mut crdt = engine.crdt.lock().map_err(|_| LiteError::LockPoisoned)?;
91        // Every key is judged before any document is written.
92        let mut seen: HashSet<&str> = HashSet::with_capacity(decoded.len());
93        let mut kept: Vec<&(&str, Vec<(String, loro::LoroValue)>)> =
94            Vec::with_capacity(decoded.len());
95        for entry in &decoded {
96            let doc_id = entry.0;
97            let taken = !seen.insert(doc_id) || crdt.exists(collection, doc_id);
98            if !taken {
99                kept.push(entry);
100            } else if !if_absent {
101                return Err(LiteError::BadRequest {
102                    detail: format!(
103                        "duplicate key value violates unique constraint on '{collection}' (id = '{doc_id}')"
104                    ),
105                });
106            }
107        }
108        let affected_n = kept.len() as u64;
109        let slices: Vec<Vec<(&str, loro::LoroValue)>> = kept
110            .iter()
111            .map(|(_, fields)| {
112                fields
113                    .iter()
114                    .map(|(k, v)| (k.as_str(), v.clone()))
115                    .collect()
116            })
117            .collect();
118        let rows: Vec<CrdtRowOp<'_>> = kept
119            .iter()
120            .zip(&slices)
121            .map(|((doc_id, _), f)| (CrdtRowWrite::Upsert, collection, *doc_id, f.as_slice()))
122            .collect();
123        // The whole batch is checked before any document is written.
124        crdt.check_unique_writes(&rows)?;
125        for &(_, _, doc_id, fields) in &rows {
126            crdt.upsert_deferred(collection, doc_id, fields)?;
127        }
128        crdt.flush_deltas().map_err(|e| LiteError::Storage {
129            detail: e.to_string(),
130        })?;
131        drop(crdt);
132        reindex_documents(engine, collection, kept.iter().map(|(doc_id, _)| *doc_id))?;
133        Ok(affected(affected_n, "INSERT"))
134    }
135}
nodedb_lite/query/document_ops/writes/batch.rs

nodedb_lite/query/document_ops/writes/
batch.rs
