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::{AnyDocumentId, OwnedDocument},
8
    schema::{
9
        view::map::MappedDocuments, Collection, Map, MappedValue, Nameable, Schema, SchemaName,
10
        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
    async fn delete_user<'user, U: Nameable<'user, u64> + Send + Sync>(
90
        &self,
91
        user: U,
92
    ) -> Result<(), bonsaidb_core::Error> {
93
        match self {
94
            Self::Local(server) => server.delete_user(user).await,
95
            Self::Networked(client) => client.delete_user(user).await,
96
        }
97
    }
98

            
99
    #[cfg(feature = "password-hashing")]
100
    async fn set_user_password<'user, U: Nameable<'user, u64> + Send + Sync>(
101
        &self,
102
        user: U,
103
        password: bonsaidb_core::connection::SensitiveString,
104
    ) -> Result<(), bonsaidb_core::Error> {
105
        match self {
106
            Self::Local(server) => server.set_user_password(user, password).await,
107
            Self::Networked(client) => client.set_user_password(user, password).await,
108
        }
109
    }
110

            
111
    #[cfg(feature = "password-hashing")]
112
    async fn authenticate<'user, U: Nameable<'user, u64> + Send + Sync>(
113
        &self,
114
        user: U,
115
        authentication: Authentication,
116
    ) -> Result<Authenticated, bonsaidb_core::Error> {
117
        match self {
118
            Self::Local(server) => server.authenticate(user, authentication).await,
119
            Self::Networked(client) => client.authenticate(user, authentication).await,
120
        }
121
    }
122

            
123
    async fn add_permission_group_to_user<
124
        'user,
125
        'group,
126
        U: Nameable<'user, u64> + Send + Sync,
127
        G: Nameable<'group, u64> + Send + Sync,
128
    >(
129
        &self,
130
        user: U,
131
        permission_group: G,
132
    ) -> Result<(), bonsaidb_core::Error> {
133
        match self {
134
            Self::Local(server) => {
135
                server
136
                    .add_permission_group_to_user(user, permission_group)
137
                    .await
138
            }
139
            Self::Networked(client) => {
140
                client
141
                    .add_permission_group_to_user(user, permission_group)
142
                    .await
143
            }
144
        }
145
    }
146

            
147
    async fn remove_permission_group_from_user<
148
        'user,
149
        'group,
150
        U: Nameable<'user, u64> + Send + Sync,
151
        G: Nameable<'group, u64> + Send + Sync,
152
    >(
153
        &self,
154
        user: U,
155
        permission_group: G,
156
    ) -> Result<(), bonsaidb_core::Error> {
157
        match self {
158
            Self::Local(server) => {
159
                server
160
                    .remove_permission_group_from_user(user, permission_group)
161
                    .await
162
            }
163
            Self::Networked(client) => {
164
                client
165
                    .remove_permission_group_from_user(user, permission_group)
166
                    .await
167
            }
168
        }
169
    }
170

            
171
    async fn add_role_to_user<
172
        'user,
173
        'role,
174
        U: Nameable<'user, u64> + Send + Sync,
175
        R: Nameable<'role, u64> + Send + Sync,
176
    >(
177
        &self,
178
        user: U,
179
        role: R,
180
    ) -> Result<(), bonsaidb_core::Error> {
181
        match self {
182
            Self::Local(server) => server.add_role_to_user(user, role).await,
183
            Self::Networked(client) => client.add_role_to_user(user, role).await,
184
        }
185
    }
186

            
187
    async fn remove_role_from_user<
188
        'user,
189
        'role,
190
        U: Nameable<'user, u64> + Send + Sync,
191
        R: Nameable<'role, u64> + Send + Sync,
192
    >(
193
        &self,
194
        user: U,
195
        role: R,
196
    ) -> Result<(), bonsaidb_core::Error> {
197
        match self {
198
            Self::Local(server) => server.remove_role_from_user(user, role).await,
199
            Self::Networked(client) => client.remove_role_from_user(user, role).await,
200
        }
201
    }
202
}
203

            
204
/// A database connection that can be either from a local server or a server
205
/// over a network connection.
206
pub enum AnyDatabase<B: Backend = NoBackend> {
207
    /// A local database.
208
    Local(ServerDatabase<B>),
209
    /// A networked database accessed with a [`Client`].
210
    Networked(RemoteDatabase<B::CustomApi>),
211
}
212

            
213
#[async_trait]
214
impl<B: Backend> Connection for AnyDatabase<B> {
215
12252
    async fn get<C, PrimaryKey>(
216
12252
        &self,
217
12252
        id: PrimaryKey,
218
12252
    ) -> Result<Option<OwnedDocument>, bonsaidb_core::Error>
219
12252
    where
220
12252
        C: Collection,
221
12252
        PrimaryKey: Into<AnyDocumentId<C::PrimaryKey>> + Send,
222
12252
    {
223
12252
        match self {
224
6126
            Self::Local(server) => server.get::<C, _>(id).await,
225
6126
            Self::Networked(client) => client.get::<C, _>(id).await,
226
        }
227
24504
    }
228

            
229
    async fn get_multiple<C, PrimaryKey, DocumentIds, I>(
230
        &self,
231
        ids: DocumentIds,
232
    ) -> Result<Vec<OwnedDocument>, bonsaidb_core::Error>
233
    where
234
        C: Collection,
235
        DocumentIds: IntoIterator<Item = PrimaryKey, IntoIter = I> + Send + Sync,
236
        I: Iterator<Item = PrimaryKey> + Send + Sync,
237
        PrimaryKey: Into<AnyDocumentId<C::PrimaryKey>> + Send + Sync,
238
    {
239
        match self {
240
            Self::Local(server) => server.get_multiple::<C, _, _, _>(ids).await,
241
            Self::Networked(client) => client.get_multiple::<C, _, _, _>(ids).await,
242
        }
243
    }
244

            
245
    async fn list<C, R, PrimaryKey>(
246
        &self,
247
        ids: R,
248
        order: Sort,
249
        limit: Option<usize>,
250
    ) -> Result<Vec<OwnedDocument>, bonsaidb_core::Error>
251
    where
252
        C: Collection,
253
        R: Into<Range<PrimaryKey>> + Send,
254
        PrimaryKey: Into<AnyDocumentId<C::PrimaryKey>> + Send,
255
    {
256
        match self {
257
            Self::Local(server) => server.list::<C, _, _>(ids, order, limit).await,
258
            Self::Networked(client) => client.list::<C, _, _>(ids, order, limit).await,
259
        }
260
    }
261

            
262
    async fn count<C, R, PrimaryKey>(&self, ids: R) -> Result<u64, bonsaidb_core::Error>
263
    where
264
        C: Collection,
265
        R: Into<Range<PrimaryKey>> + Send,
266
        PrimaryKey: Into<AnyDocumentId<C::PrimaryKey>> + Send,
267
    {
268
        match self {
269
            Self::Local(server) => server.count::<C, _, _>(ids).await,
270
            Self::Networked(client) => client.count::<C, _, _>(ids).await,
271
        }
272
    }
273

            
274
    async fn query<V: SerializedView>(
275
        &self,
276
        key: Option<QueryKey<V::Key>>,
277
        order: Sort,
278
        limit: Option<usize>,
279
        access_policy: AccessPolicy,
280
    ) -> Result<Vec<Map<V::Key, V::Value>>, bonsaidb_core::Error>
281
    where
282
        Self: Sized,
283
    {
284
        match self {
285
            Self::Local(server) => server.query::<V>(key, order, limit, access_policy).await,
286
            Self::Networked(client) => client.query::<V>(key, order, limit, access_policy).await,
287
        }
288
    }
289

            
290
7732
    async fn query_with_docs<V: SerializedView>(
291
7732
        &self,
292
7732
        key: Option<QueryKey<V::Key>>,
293
7732
        order: Sort,
294
7732
        limit: Option<usize>,
295
7732
        access_policy: AccessPolicy,
296
7732
    ) -> Result<MappedDocuments<OwnedDocument, V>, bonsaidb_core::Error>
297
7732
    where
298
7732
        Self: Sized,
299
7732
    {
300
7732
        match self {
301
3866
            Self::Local(server) => {
302
3866
                server
303
4315
                    .query_with_docs::<V>(key, order, limit, access_policy)
304
4315
                    .await
305
            }
306
3866
            Self::Networked(client) => {
307
3866
                client
308
3866
                    .query_with_docs::<V>(key, order, limit, access_policy)
309
3866
                    .await
310
            }
311
        }
312
15464
    }
313

            
314
15917
    async fn reduce<V: SerializedView>(
315
15917
        &self,
316
15917
        key: Option<QueryKey<V::Key>>,
317
15917
        access_policy: AccessPolicy,
318
15917
    ) -> Result<V::Value, bonsaidb_core::Error>
319
15917
    where
320
15917
        Self: Sized,
321
15917
    {
322
15916
        match self {
323
7957
            Self::Local(server) => server.reduce::<V>(key, access_policy).await,
324
7959
            Self::Networked(client) => client.reduce::<V>(key, access_policy).await,
325
        }
326
31833
    }
327

            
328
    async fn reduce_grouped<V: SerializedView>(
329
        &self,
330
        key: Option<QueryKey<V::Key>>,
331
        access_policy: AccessPolicy,
332
    ) -> Result<Vec<MappedValue<V::Key, V::Value>>, bonsaidb_core::Error>
333
    where
334
        Self: Sized,
335
    {
336
        match self {
337
            Self::Local(server) => server.reduce_grouped::<V>(key, access_policy).await,
338
            Self::Networked(client) => client.reduce_grouped::<V>(key, access_policy).await,
339
        }
340
    }
341

            
342
    async fn delete_docs<V: SerializedView>(
343
        &self,
344
        key: Option<QueryKey<V::Key>>,
345
        access_policy: AccessPolicy,
346
    ) -> Result<u64, bonsaidb_core::Error>
347
    where
348
        Self: Sized,
349
    {
350
        match self {
351
            Self::Local(server) => server.delete_docs::<V>(key, access_policy).await,
352
            Self::Networked(client) => client.delete_docs::<V>(key, access_policy).await,
353
        }
354
    }
355

            
356
6614
    async fn apply_transaction(
357
6614
        &self,
358
6614
        transaction: Transaction,
359
6614
    ) -> Result<Vec<OperationResult>, bonsaidb_core::Error> {
360
6614
        match self {
361
3307
            Self::Local(server) => server.apply_transaction(transaction).await,
362
3307
            Self::Networked(client) => client.apply_transaction(transaction).await,
363
        }
364
13228
    }
365

            
366
    async fn list_executed_transactions(
367
        &self,
368
        starting_id: Option<u64>,
369
        result_limit: Option<usize>,
370
    ) -> Result<Vec<Executed>, bonsaidb_core::Error> {
371
        match self {
372
            Self::Local(server) => {
373
                server
374
                    .list_executed_transactions(starting_id, result_limit)
375
                    .await
376
            }
377
            Self::Networked(client) => {
378
                client
379
                    .list_executed_transactions(starting_id, result_limit)
380
                    .await
381
            }
382
        }
383
    }
384

            
385
    async fn last_transaction_id(&self) -> Result<Option<u64>, bonsaidb_core::Error> {
386
        match self {
387
            Self::Local(server) => server.last_transaction_id().await,
388
            Self::Networked(client) => client.last_transaction_id().await,
389
        }
390
    }
391

            
392
    async fn compact_collection<C: Collection>(&self) -> Result<(), bonsaidb_core::Error> {
393
        match self {
394
            Self::Local(server) => server.compact_collection::<C>().await,
395
            Self::Networked(client) => client.compact_collection::<C>().await,
396
        }
397
    }
398

            
399
    async fn compact(&self) -> Result<(), bonsaidb_core::Error> {
400
        match self {
401
            Self::Local(server) => server.compact().await,
402
            Self::Networked(client) => client.compact().await,
403
        }
404
    }
405

            
406
    async fn compact_key_value_store(&self) -> Result<(), bonsaidb_core::Error> {
407
        match self {
408
            Self::Local(server) => server.compact_key_value_store().await,
409
            Self::Networked(client) => client.compact_key_value_store().await,
410
        }
411
    }
412
}