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
    #[cfg(feature = "password-hashing")]
90
    async fn set_user_password<'user, U: Nameable<'user, u64> + 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: Nameable<'user, u64> + 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: Nameable<'user, u64> + Send + Sync,
117
        G: Nameable<'group, u64> + 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: Nameable<'user, u64> + Send + Sync,
141
        G: Nameable<'group, u64> + 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: Nameable<'user, u64> + Send + Sync,
165
        R: Nameable<'role, u64> + 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: Nameable<'user, u64> + Send + Sync,
181
        R: Nameable<'role, u64> + 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
11736
    async fn get<C, PrimaryKey>(
206
11736
        &self,
207
11736
        id: PrimaryKey,
208
11736
    ) -> Result<Option<OwnedDocument>, bonsaidb_core::Error>
209
11736
    where
210
11736
        C: Collection,
211
11736
        PrimaryKey: Into<AnyDocumentId<C::PrimaryKey>> + Send,
212
11736
    {
213
11736
        match self {
214
5868
            Self::Local(server) => server.get::<C, _>(id).await,
215
5868
            Self::Networked(client) => client.get::<C, _>(id).await,
216
        }
217
23472
    }
218

            
219
    async fn get_multiple<C, PrimaryKey, DocumentIds, I>(
220
        &self,
221
        ids: DocumentIds,
222
    ) -> Result<Vec<OwnedDocument>, bonsaidb_core::Error>
223
    where
224
        C: Collection,
225
        DocumentIds: IntoIterator<Item = PrimaryKey, IntoIter = I> + Send + Sync,
226
        I: Iterator<Item = PrimaryKey> + Send + Sync,
227
        PrimaryKey: Into<AnyDocumentId<C::PrimaryKey>> + Send + Sync,
228
    {
229
        match self {
230
            Self::Local(server) => server.get_multiple::<C, _, _, _>(ids).await,
231
            Self::Networked(client) => client.get_multiple::<C, _, _, _>(ids).await,
232
        }
233
    }
234

            
235
    async fn list<C, R, PrimaryKey>(
236
        &self,
237
        ids: R,
238
        order: Sort,
239
        limit: Option<usize>,
240
    ) -> Result<Vec<OwnedDocument>, bonsaidb_core::Error>
241
    where
242
        C: Collection,
243
        R: Into<Range<PrimaryKey>> + Send,
244
        PrimaryKey: Into<AnyDocumentId<C::PrimaryKey>> + Send,
245
    {
246
        match self {
247
            Self::Local(server) => server.list::<C, _, _>(ids, order, limit).await,
248
            Self::Networked(client) => client.list::<C, _, _>(ids, order, limit).await,
249
        }
250
    }
251

            
252
    async fn query<V: SerializedView>(
253
        &self,
254
        key: Option<QueryKey<V::Key>>,
255
        order: Sort,
256
        limit: Option<usize>,
257
        access_policy: AccessPolicy,
258
    ) -> Result<Vec<Map<V::Key, V::Value>>, bonsaidb_core::Error>
259
    where
260
        Self: Sized,
261
    {
262
        match self {
263
            Self::Local(server) => server.query::<V>(key, order, limit, access_policy).await,
264
            Self::Networked(client) => client.query::<V>(key, order, limit, access_policy).await,
265
        }
266
    }
267

            
268
7656
    async fn query_with_docs<V: SerializedView>(
269
7656
        &self,
270
7656
        key: Option<QueryKey<V::Key>>,
271
7656
        order: Sort,
272
7656
        limit: Option<usize>,
273
7656
        access_policy: AccessPolicy,
274
7656
    ) -> Result<MappedDocuments<OwnedDocument, V>, bonsaidb_core::Error>
275
7656
    where
276
7656
        Self: Sized,
277
7656
    {
278
7656
        match self {
279
3828
            Self::Local(server) => {
280
3828
                server
281
4049
                    .query_with_docs::<V>(key, order, limit, access_policy)
282
4049
                    .await
283
            }
284
3828
            Self::Networked(client) => {
285
3828
                client
286
3828
                    .query_with_docs::<V>(key, order, limit, access_policy)
287
3828
                    .await
288
            }
289
        }
290
15312
    }
291

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

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

            
320
    async fn delete_docs<V: SerializedView>(
321
        &self,
322
        key: Option<QueryKey<V::Key>>,
323
        access_policy: AccessPolicy,
324
    ) -> Result<u64, bonsaidb_core::Error>
325
    where
326
        Self: Sized,
327
    {
328
        match self {
329
            Self::Local(server) => server.delete_docs::<V>(key, access_policy).await,
330
            Self::Networked(client) => client.delete_docs::<V>(key, access_policy).await,
331
        }
332
    }
333

            
334
6550
    async fn apply_transaction(
335
6550
        &self,
336
6550
        transaction: Transaction,
337
6550
    ) -> Result<Vec<OperationResult>, bonsaidb_core::Error> {
338
6550
        match self {
339
3275
            Self::Local(server) => server.apply_transaction(transaction).await,
340
3275
            Self::Networked(client) => client.apply_transaction(transaction).await,
341
        }
342
13100
    }
343

            
344
    async fn list_executed_transactions(
345
        &self,
346
        starting_id: Option<u64>,
347
        result_limit: Option<usize>,
348
    ) -> Result<Vec<Executed>, bonsaidb_core::Error> {
349
        match self {
350
            Self::Local(server) => {
351
                server
352
                    .list_executed_transactions(starting_id, result_limit)
353
                    .await
354
            }
355
            Self::Networked(client) => {
356
                client
357
                    .list_executed_transactions(starting_id, result_limit)
358
                    .await
359
            }
360
        }
361
    }
362

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

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

            
377
    async fn compact(&self) -> Result<(), bonsaidb_core::Error> {
378
        match self {
379
            Self::Local(server) => server.compact().await,
380
            Self::Networked(client) => client.compact().await,
381
        }
382
    }
383

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