1
use bonsaidb_client::{Client, RemoteDatabase};
2
#[cfg(feature = "password-hashing")]
3
use bonsaidb_core::connection::{Authenticated, Authentication};
4
use bonsaidb_core::{
5
    async_trait::async_trait,
6
    connection::{self, AccessPolicy, Connection, QueryKey, Range, Sort, StorageConnection},
7
    document::OwnedDocument,
8
    schema::{
9
        view::map::MappedDocuments, Collection, Map, MappedValue, NamedReference, Schema,
10
        SchemaName, SerializedView,
11
    },
12
    transaction::{Executed, OperationResult, Transaction},
13
};
14
use bonsaidb_server::{Backend, CustomServer, NoBackend, ServerDatabase};
15

            
16
/// A local server or a server over a network connection.
17
pub enum AnyServerConnection<B: Backend> {
18
    /// A local server.
19
    Local(CustomServer<B>),
20
    /// A server accessed with a [`Client`].
21
    Networked(Client<B::CustomApi>),
22
}
23

            
24
#[async_trait]
25
impl<B: Backend> StorageConnection for AnyServerConnection<B> {
26
    type Database = AnyDatabase<B>;
27

            
28
3
    async fn database<DB: Schema>(
29
3
        &self,
30
3
        name: &str,
31
3
    ) -> Result<Self::Database, bonsaidb_core::Error> {
32
3
        match self {
33
1
            Self::Local(server) => server.database::<DB>(name).await.map(AnyDatabase::Local),
34
2
            Self::Networked(client) => client
35
2
                .database::<DB>(name)
36
                .await
37
2
                .map(AnyDatabase::Networked),
38
        }
39
6
    }
40

            
41
    async fn create_database_with_schema(
42
        &self,
43
        name: &str,
44
        schema: SchemaName,
45
        only_if_needed: bool,
46
    ) -> Result<(), bonsaidb_core::Error> {
47
        match self {
48
            Self::Local(server) => {
49
                server
50
                    .create_database_with_schema(name, schema, only_if_needed)
51
                    .await
52
            }
53
            Self::Networked(client) => {
54
                client
55
                    .create_database_with_schema(name, schema, only_if_needed)
56
                    .await
57
            }
58
        }
59
    }
60

            
61
    async fn delete_database(&self, name: &str) -> Result<(), bonsaidb_core::Error> {
62
        match self {
63
            Self::Local(server) => server.delete_database(name).await,
64
            Self::Networked(client) => client.delete_database(name).await,
65
        }
66
    }
67

            
68
    async fn list_databases(&self) -> Result<Vec<connection::Database>, bonsaidb_core::Error> {
69
        match self {
70
            Self::Local(server) => server.list_databases().await,
71
            Self::Networked(client) => client.list_databases().await,
72
        }
73
    }
74

            
75
    async fn list_available_schemas(&self) -> Result<Vec<SchemaName>, bonsaidb_core::Error> {
76
        match self {
77
            Self::Local(server) => server.list_available_schemas().await,
78
            Self::Networked(client) => client.list_available_schemas().await,
79
        }
80
    }
81

            
82
    async fn create_user(&self, username: &str) -> Result<u64, bonsaidb_core::Error> {
83
        match self {
84
            Self::Local(server) => server.create_user(username).await,
85
            Self::Networked(client) => client.create_user(username).await,
86
        }
87
    }
88

            
89
    #[cfg(feature = "password-hashing")]
90
    async fn set_user_password<'user, U: Into<NamedReference<'user>> + Send + Sync>(
91
        &self,
92
        user: U,
93
        password: bonsaidb_core::connection::SensitiveString,
94
    ) -> Result<(), bonsaidb_core::Error> {
95
        match self {
96
            Self::Local(server) => server.set_user_password(user, password).await,
97
            Self::Networked(client) => client.set_user_password(user, password).await,
98
        }
99
    }
100

            
101
    #[cfg(feature = "password-hashing")]
102
    async fn authenticate<'user, U: Into<NamedReference<'user>> + Send + Sync>(
103
        &self,
104
        user: U,
105
        authentication: Authentication,
106
    ) -> Result<Authenticated, bonsaidb_core::Error> {
107
        match self {
108
            Self::Local(server) => server.authenticate(user, authentication).await,
109
            Self::Networked(client) => client.authenticate(user, authentication).await,
110
        }
111
    }
112

            
113
    async fn add_permission_group_to_user<
114
        'user,
115
        'group,
116
        U: Into<NamedReference<'user>> + Send + Sync,
117
        G: Into<NamedReference<'group>> + Send + Sync,
118
    >(
119
        &self,
120
        user: U,
121
        permission_group: G,
122
    ) -> Result<(), bonsaidb_core::Error> {
123
        match self {
124
            Self::Local(server) => {
125
                server
126
                    .add_permission_group_to_user(user, permission_group)
127
                    .await
128
            }
129
            Self::Networked(client) => {
130
                client
131
                    .add_permission_group_to_user(user, permission_group)
132
                    .await
133
            }
134
        }
135
    }
136

            
137
    async fn remove_permission_group_from_user<
138
        'user,
139
        'group,
140
        U: Into<NamedReference<'user>> + Send + Sync,
141
        G: Into<NamedReference<'group>> + Send + Sync,
142
    >(
143
        &self,
144
        user: U,
145
        permission_group: G,
146
    ) -> Result<(), bonsaidb_core::Error> {
147
        match self {
148
            Self::Local(server) => {
149
                server
150
                    .remove_permission_group_from_user(user, permission_group)
151
                    .await
152
            }
153
            Self::Networked(client) => {
154
                client
155
                    .remove_permission_group_from_user(user, permission_group)
156
                    .await
157
            }
158
        }
159
    }
160

            
161
    async fn add_role_to_user<
162
        'user,
163
        'role,
164
        U: Into<NamedReference<'user>> + Send + Sync,
165
        R: Into<NamedReference<'role>> + Send + Sync,
166
    >(
167
        &self,
168
        user: U,
169
        role: R,
170
    ) -> Result<(), bonsaidb_core::Error> {
171
        match self {
172
            Self::Local(server) => server.add_role_to_user(user, role).await,
173
            Self::Networked(client) => client.add_role_to_user(user, role).await,
174
        }
175
    }
176

            
177
    async fn remove_role_from_user<
178
        'user,
179
        'role,
180
        U: Into<NamedReference<'user>> + Send + Sync,
181
        R: Into<NamedReference<'role>> + Send + Sync,
182
    >(
183
        &self,
184
        user: U,
185
        role: R,
186
    ) -> Result<(), bonsaidb_core::Error> {
187
        match self {
188
            Self::Local(server) => server.remove_role_from_user(user, role).await,
189
            Self::Networked(client) => client.remove_role_from_user(user, role).await,
190
        }
191
    }
192
}
193

            
194
/// A database connection that can be either from a local server or a server
195
/// over a network connection.
196
pub enum AnyDatabase<B: Backend = NoBackend> {
197
    /// A local database.
198
    Local(ServerDatabase<B>),
199
    /// A networked database accessed with a [`Client`].
200
    Networked(RemoteDatabase<B::CustomApi>),
201
}
202

            
203
#[async_trait]
204
impl<B: Backend> Connection for AnyDatabase<B> {
205
9057
    async fn get<C: Collection>(
206
9057
        &self,
207
9057
        id: u64,
208
9057
    ) -> Result<Option<OwnedDocument>, bonsaidb_core::Error> {
209
9057
        match self {
210
3019
            Self::Local(server) => server.get::<C>(id).await,
211
6038
            Self::Networked(client) => client.get::<C>(id).await,
212
        }
213
18114
    }
214

            
215
    async fn get_multiple<C: Collection>(
216
        &self,
217
        ids: &[u64],
218
    ) -> Result<Vec<OwnedDocument>, bonsaidb_core::Error> {
219
        match self {
220
            Self::Local(server) => server.get_multiple::<C>(ids).await,
221
            Self::Networked(client) => client.get_multiple::<C>(ids).await,
222
        }
223
    }
224

            
225
    async fn list<C: Collection, R: Into<Range<u64>> + Send>(
226
        &self,
227
        ids: R,
228
        order: Sort,
229
        limit: Option<usize>,
230
    ) -> Result<Vec<OwnedDocument>, bonsaidb_core::Error> {
231
        match self {
232
            Self::Local(server) => server.list::<C, R>(ids, order, limit).await,
233
            Self::Networked(client) => client.list::<C, R>(ids, order, limit).await,
234
        }
235
    }
236

            
237
    async fn query<V: SerializedView>(
238
        &self,
239
        key: Option<QueryKey<V::Key>>,
240
        order: Sort,
241
        limit: Option<usize>,
242
        access_policy: AccessPolicy,
243
    ) -> Result<Vec<Map<V::Key, V::Value>>, bonsaidb_core::Error>
244
    where
245
        Self: Sized,
246
    {
247
        match self {
248
            Self::Local(server) => server.query::<V>(key, order, limit, access_policy).await,
249
            Self::Networked(client) => client.query::<V>(key, order, limit, access_policy).await,
250
        }
251
    }
252

            
253
6029
    async fn query_with_docs<V: SerializedView>(
254
6029
        &self,
255
6029
        key: Option<QueryKey<V::Key>>,
256
6029
        order: Sort,
257
6029
        limit: Option<usize>,
258
6029
        access_policy: AccessPolicy,
259
6029
    ) -> Result<MappedDocuments<OwnedDocument, V>, bonsaidb_core::Error>
260
6029
    where
261
6029
        Self: Sized,
262
6030
    {
263
6030
        match self {
264
2010
            Self::Local(server) => {
265
2010
                server
266
2063
                    .query_with_docs::<V>(key, order, limit, access_policy)
267
2063
                    .await
268
            }
269
4020
            Self::Networked(client) => {
270
4020
                client
271
4020
                    .query_with_docs::<V>(key, order, limit, access_policy)
272
4020
                    .await
273
            }
274
        }
275
12060
    }
276

            
277
11989
    async fn reduce<V: SerializedView>(
278
11989
        &self,
279
11989
        key: Option<QueryKey<V::Key>>,
280
11989
        access_policy: AccessPolicy,
281
11989
    ) -> Result<V::Value, bonsaidb_core::Error>
282
11989
    where
283
11989
        Self: Sized,
284
11989
    {
285
11989
        match self {
286
3996
            Self::Local(server) => server.reduce::<V>(key, access_policy).await,
287
7993
            Self::Networked(client) => client.reduce::<V>(key, access_policy).await,
288
        }
289
23978
    }
290

            
291
    async fn reduce_grouped<V: SerializedView>(
292
        &self,
293
        key: Option<QueryKey<V::Key>>,
294
        access_policy: AccessPolicy,
295
    ) -> Result<Vec<MappedValue<V::Key, V::Value>>, bonsaidb_core::Error>
296
    where
297
        Self: Sized,
298
    {
299
        match self {
300
            Self::Local(server) => server.reduce_grouped::<V>(key, access_policy).await,
301
            Self::Networked(client) => client.reduce_grouped::<V>(key, access_policy).await,
302
        }
303
    }
304

            
305
    async fn delete_docs<V: SerializedView>(
306
        &self,
307
        key: Option<QueryKey<V::Key>>,
308
        access_policy: AccessPolicy,
309
    ) -> Result<u64, bonsaidb_core::Error>
310
    where
311
        Self: Sized,
312
    {
313
        match self {
314
            Self::Local(server) => server.delete_docs::<V>(key, access_policy).await,
315
            Self::Networked(client) => client.delete_docs::<V>(key, access_policy).await,
316
        }
317
    }
318

            
319
5027
    async fn apply_transaction(
320
5027
        &self,
321
5027
        transaction: Transaction,
322
5027
    ) -> Result<Vec<OperationResult>, bonsaidb_core::Error> {
323
5027
        match self {
324
1676
            Self::Local(server) => server.apply_transaction(transaction).await,
325
3351
            Self::Networked(client) => client.apply_transaction(transaction).await,
326
        }
327
10053
    }
328

            
329
    async fn list_executed_transactions(
330
        &self,
331
        starting_id: Option<u64>,
332
        result_limit: Option<usize>,
333
    ) -> Result<Vec<Executed>, bonsaidb_core::Error> {
334
        match self {
335
            Self::Local(server) => {
336
                server
337
                    .list_executed_transactions(starting_id, result_limit)
338
                    .await
339
            }
340
            Self::Networked(client) => {
341
                client
342
                    .list_executed_transactions(starting_id, result_limit)
343
                    .await
344
            }
345
        }
346
    }
347

            
348
    async fn last_transaction_id(&self) -> Result<Option<u64>, bonsaidb_core::Error> {
349
        match self {
350
            Self::Local(server) => server.last_transaction_id().await,
351
            Self::Networked(client) => client.last_transaction_id().await,
352
        }
353
    }
354

            
355
    async fn compact_collection<C: Collection>(&self) -> Result<(), bonsaidb_core::Error> {
356
        match self {
357
            Self::Local(server) => server.compact_collection::<C>().await,
358
            Self::Networked(client) => client.compact_collection::<C>().await,
359
        }
360
    }
361

            
362
    async fn compact(&self) -> Result<(), bonsaidb_core::Error> {
363
        match self {
364
            Self::Local(server) => server.compact().await,
365
            Self::Networked(client) => client.compact().await,
366
        }
367
    }
368

            
369
    async fn compact_key_value_store(&self) -> Result<(), bonsaidb_core::Error> {
370
        match self {
371
            Self::Local(server) => server.compact_key_value_store().await,
372
            Self::Networked(client) => client.compact_key_value_store().await,
373
        }
374
    }
375
}