nodedb_lite/query/document_ops/sets/
copy.rs

1// SPDX-License-Identifier: Apache-2.0
2use super::super::is_strict;
3use super::super::reads::loro_value_to_ndb_value;
4use super::super::writes::batch_insert_admitted;
5use crate::error::LiteError;
6use crate::query::engine::LiteQueryEngine;
7use crate::query::value_utils::value_to_string;
8use crate::storage::engine::StorageEngine;
9use nodedb_types::result::QueryResult;
10use nodedb_types::value::Value;
11use std::collections::HashMap;
12/// InsertSelect: copy documents from source to target collection.
13///
14/// Reads the documents of `source_collection` up to `source_limit`, then
15/// batch-inserts them into `target_collection`. Every document is copied
16/// unchanged. A key the target already stores refuses the whole copy.
17pub async fn insert_select<S: StorageEngine>(
18    engine: &LiteQueryEngine<S>,
19    target_collection: &str,
20    source_collection: &str,
21    source_limit: usize,
22) -> Result<QueryResult, LiteError> {
23    insert_select_coordinated(
24        engine,
25        None,
26        target_collection,
27        source_collection,
28        source_limit,
29    )
30    .await
31}
32
33pub(crate) async fn insert_select_coordinated<S: StorageEngine>(
34    engine: &LiteQueryEngine<S>,
35    permit: Option<&crate::engine::fts::coordinator::TextMutationPermit>,
36    target_collection: &str,
37    source_collection: &str,
38    source_limit: usize,
39) -> Result<QueryResult, LiteError> {
40    if let Some(permit) = permit {
41        return insert_select_admitted(
42            engine,
43            permit,
44            target_collection,
45            source_collection,
46            source_limit,
47        )
48        .await;
49    }
50    let guard = engine.fts_state.admit_mutation().await;
51    let result = insert_select_admitted(
52        engine,
53        guard.permit(),
54        target_collection,
55        source_collection,
56        source_limit,
57    )
58    .await;
59    guard.finish(result)
60}
61
62pub(crate) async fn insert_select_admitted<S: StorageEngine>(
63    engine: &LiteQueryEngine<S>,
64    permit: &crate::engine::fts::coordinator::TextMutationPermit,
65    target_collection: &str,
66    source_collection: &str,
67    source_limit: usize,
68) -> Result<QueryResult, LiteError> {
69    let _permit = permit;
70    let documents: Vec<(String, Vec<u8>)> = if is_strict(engine, source_collection) {
71        let schema = engine
72            .strict
73            .schema(source_collection)
74            .ok_or_else(|| LiteError::collection_not_found("strict", source_collection))?;
75        let pk_idx = schema
76            .columns
77            .iter()
78            .position(|c| c.primary_key)
79            .ok_or_else(|| LiteError::BadRequest {
80                detail: format!(
81                    "strict source collection '{source_collection}' has no primary key"
82                ),
83            })?;
84        let columns = schema.columns;
85        let all_rows = engine.strict.list_rows(source_collection).await?;
86        let mut docs = Vec::with_capacity(all_rows.len().min(source_limit));
87        for row in all_rows.into_iter().take(source_limit) {
88            let pk = row.get(pk_idx).ok_or_else(|| LiteError::Storage {
89                detail: format!("strict collection '{source_collection}' row omits primary key at column {pk_idx}"),
90            })?;
91            let pk = value_to_string(pk);
92            let map: HashMap<String, Value> = columns
93                .iter()
94                .enumerate()
95                .filter_map(|(i, col)| row.get(i).map(|value| (col.name.clone(), value.clone())))
96                .collect();
97            let bytes = zerompk::to_msgpack_vec(&Value::Object(map)).map_err(|e| {
98                LiteError::Serialization {
99                    detail: format!("serialize source row: {e}"),
100                }
101            })?;
102            docs.push((pk, bytes));
103        }
104        docs
105    } else {
106        let crdt = engine.crdt.lock().map_err(|_| LiteError::LockPoisoned)?;
107        let ids = crdt.list_ids(source_collection);
108        let mut docs = Vec::with_capacity(ids.len().min(source_limit));
109        for id in ids.into_iter().take(source_limit) {
110            if let Some(val) = crdt.read(source_collection, &id) {
111                let ndb_val = loro_value_to_ndb_value(&val);
112                let bytes =
113                    zerompk::to_msgpack_vec(&ndb_val).map_err(|e| LiteError::Serialization {
114                        detail: format!("serialize crdt source row: {e}"),
115                    })?;
116                docs.push((id, bytes));
117            }
118        }
119        drop(crdt);
120        docs
121    };
122
123    batch_insert_admitted(engine, permit, target_collection, &documents, false).await
124}
125
126#[cfg(test)]
127mod tests {
128    #[tokio::test]
129    async fn insert_select_borrows_existing_admission() -> Result<(), Box<dyn std::error::Error>> {
130        use nodedb_types::value::Value;
131        let storage = crate::PagedbStorageMem::open_in_memory().await?;
132        let db = crate::NodeDbLite::open(storage).await?;
133        let bytes = zerompk::to_msgpack_vec(&Value::Object(std::collections::HashMap::from([(
134            "body".into(),
135            Value::String("copied".into()),
136        )])))?;
137        crate::query::document_ops::writes::point_put(&db.query_engine, "source", "one", &bytes)
138            .await?;
139        let guard = db.query_engine.fts_state.admit_mutation().await;
140        let result = tokio::time::timeout(
141            std::time::Duration::from_secs(5),
142            super::insert_select_coordinated(
143                &db.query_engine,
144                Some(guard.permit()),
145                "target",
146                "source",
147                10,
148            ),
149        )
150        .await?;
151        let result = guard.finish(result)?;
152        assert_eq!(result.rows_affected, 1);
153        assert!(
154            db.query_engine
155                .crdt
156                .lock()
157                .map_err(|_| crate::error::LiteError::LockPoisoned)?
158                .exists("target", "one")
159        );
160        Ok(())
161    }
162}
nodedb_lite/query/document_ops/sets/copy.rs

nodedb_lite/query/document_ops/sets/
copy.rs
