1
use std::collections::HashMap;
2
#[cfg(feature = "encryption")]
3
use std::sync::Arc;
4

            
5
use bonsaidb_core::{
6
    document::KeyId,
7
    schema::{CollectionName, Schematic},
8
};
9
use nebari::{
10
    io::fs::StdFile,
11
    tree::{AnyTreeRoot, Root, Unversioned, Versioned},
12
};
13

            
14
#[cfg(feature = "encryption")]
15
use crate::vault::{TreeVault, Vault};
16
use crate::{
17
    database::document_tree_name,
18
    views::{
19
        view_document_map_tree_name, view_entries_tree_name, view_invalidated_docs_tree_name,
20
        view_omitted_docs_tree_name,
21
    },
22
    Error,
23
};
24

            
25
300420
#[derive(Default)]
26
pub(crate) struct OpenTrees {
27
    pub trees: Vec<Box<dyn AnyTreeRoot<StdFile>>>,
28
    pub trees_index_by_name: HashMap<String, usize>,
29
}
30

            
31
impl OpenTrees {
32
    #[cfg_attr(not(feature = "encryption"), allow(unused_mut))]
33
    #[cfg_attr(feature = "encryption", allow(clippy::unnecessary_wraps))]
34
1401507
    pub fn open_tree<R: Root>(
35
1401507
        &mut self,
36
1401507
        name: &str,
37
1401507
        encryption_key: Option<&KeyId>,
38
1401507
        #[cfg(feature = "encryption")] vault: &Arc<Vault>,
39
1401507
    ) -> Result<(), Error> {
40
1401507
        if !self.trees_index_by_name.contains_key(name) {
41
1038436
            self.trees_index_by_name
42
1038436
                .insert(name.to_string(), self.trees.len());
43
1038436
            let mut tree = R::tree(name.to_string());
44

            
45
            #[cfg(feature = "encryption")]
46
1038436
            if let Some(encryption_key) = encryption_key {
47
10210
                tree = tree.with_vault(TreeVault {
48
10210
                    key: encryption_key.clone(),
49
10210
                    vault: vault.clone(),
50
10210
                });
51
1028269
            }
52

            
53
            #[cfg(not(feature = "encryption"))]
54
            if encryption_key.is_some() {
55
                return Err(Error::EncryptionDisabled);
56
            }
57

            
58
1038479
            self.trees.push(Box::new(tree));
59
363071
        }
60
1401550
        Ok(())
61
1401550
    }
62

            
63
    pub fn open_trees_for_document_change(
64
        &mut self,
65
        collection: &CollectionName,
66
        schema: &Schematic,
67
        encryption_key: Option<&KeyId>,
68
        #[cfg(feature = "encryption")] vault: &Arc<Vault>,
69
    ) -> Result<(), Error> {
70
500003
        self.open_tree::<Versioned>(
71
500003
            &document_tree_name(collection),
72
500003
            encryption_key,
73
500003
            #[cfg(feature = "encryption")]
74
500003
            vault,
75
500003
        )?;
76

            
77
500003
        if let Some(views) = schema.views_in_collection(collection) {
78
1041324
            for view in views {
79
774280
                let view_name = view.view_name();
80
774280
                if view.unique() {
81
63601
                    self.open_tree::<Unversioned>(
82
63601
                        &view_omitted_docs_tree_name(&view_name),
83
63601
                        encryption_key,
84
63601
                        #[cfg(feature = "encryption")]
85
63601
                        vault,
86
63601
                    )?;
87
63601
                    self.open_tree::<Unversioned>(
88
63601
                        &view_document_map_tree_name(&view_name),
89
63601
                        encryption_key,
90
63601
                        #[cfg(feature = "encryption")]
91
63601
                        vault,
92
63601
                    )?;
93
63601
                    self.open_tree::<Unversioned>(
94
63601
                        &view_entries_tree_name(&view_name),
95
63601
                        encryption_key,
96
63601
                        #[cfg(feature = "encryption")]
97
63601
                        vault,
98
63601
                    )?;
99
                } else {
100
710679
                    self.open_tree::<Unversioned>(
101
710679
                        &view_invalidated_docs_tree_name(&view_name),
102
710679
                        encryption_key,
103
710679
                        #[cfg(feature = "encryption")]
104
710679
                        vault,
105
710679
                    )?;
106
                }
107
            }
108
232959
        }
109

            
110
500003
        Ok(())
111
500003
    }
112
}