1
use std::collections::HashMap;
2
use std::marker::PhantomData;
3
use std::sync::Arc;
4
use std::time::Duration;
5

            
6
use bonsaidb_core::api;
7
use bonsaidb_core::api::ApiName;
8
use bonsaidb_core::networking::CURRENT_PROTOCOL_VERSION;
9
#[cfg(not(target_arch = "wasm32"))]
10
use fabruic::Certificate;
11
#[cfg(not(target_arch = "wasm32"))]
12
use tokio::runtime::Handle;
13
use url::Url;
14

            
15
use crate::client::{AnyApiCallback, ApiCallback};
16
#[cfg(not(target_arch = "wasm32"))]
17
use crate::BlockingClient;
18
use crate::{AsyncClient, Error};
19

            
20
pub struct Async;
21
#[cfg(not(target_arch = "wasm32"))]
22
pub struct Blocking;
23

            
24
/// Builder for a [`BlockingClient`] or an [`AsyncClient`].
25
#[must_use]
26
pub struct Builder<AsyncMode> {
27
    url: Url,
28
    protocol_version: &'static str,
29
    custom_apis: HashMap<ApiName, Option<Arc<dyn AnyApiCallback>>>,
30
    connect_timeout: Option<Duration>,
31
    request_timeout: Option<Duration>,
32
    #[cfg(not(target_arch = "wasm32"))]
33
    certificate: Option<fabruic::Certificate>,
34
    #[cfg(not(target_arch = "wasm32"))]
35
    tokio: Option<Handle>,
36
    mode: PhantomData<AsyncMode>,
37
}
38

            
39
impl<AsyncMode> Builder<AsyncMode> {
40
    /// Creates a new builder for a client connecting to `url`.
41
1818
    pub(crate) fn new(url: Url) -> Self {
42
1818
        Self {
43
1818
            url,
44
1818
            protocol_version: CURRENT_PROTOCOL_VERSION,
45
1818
            custom_apis: HashMap::new(),
46
1818
            request_timeout: None,
47
1818
            connect_timeout: None,
48
1818
            #[cfg(not(target_arch = "wasm32"))]
49
1818
            certificate: None,
50
1818
            #[cfg(not(target_arch = "wasm32"))]
51
1818
            tokio: None,
52
1818
            mode: PhantomData,
53
1818
        }
54
1818
    }
55

            
56
    /// Specifies the tokio runtime this client should use for its async tasks.
57
    /// If not specified, `Client` will try to acquire a handle via
58
    /// `tokio::runtime::Handle::try_current()`.
59
    #[cfg(not(target_arch = "wasm32"))]
60
    #[allow(clippy::missing_const_for_fn)]
61
    pub fn with_runtime(mut self, handle: Handle) -> Self {
62
        self.tokio = Some(handle);
63
        self
64
    }
65

            
66
    /// Enables using a [`Api`](api::Api) with this client. If you want to
67
    /// receive out-of-band API requests, set a callback using
68
    /// `with_custom_api_callback` instead.
69
1
    pub fn with_api<Api: api::Api>(mut self) -> Self {
70
1
        self.custom_apis.insert(Api::name(), None);
71
1
        self
72
1
    }
73

            
74
    /// Enables using a [`Api`](api::Api) with this client. `callback` will be
75
    /// invoked when custom API responses are received from the server.
76
    pub fn with_api_callback<Api: api::Api>(mut self, callback: ApiCallback<Api>) -> Self {
77
        self.custom_apis
78
            .insert(Api::name(), Some(Arc::new(callback)));
79
        self
80
    }
81

            
82
    /// Connects to a server using a pinned `certificate`. Only supported with BonsaiDb protocol-based connections.
83
    #[cfg(not(target_arch = "wasm32"))]
84
    #[allow(clippy::missing_const_for_fn)]
85
72
    pub fn with_certificate(mut self, certificate: Certificate) -> Self {
86
72
        self.certificate = Some(certificate);
87
72
        self
88
72
    }
89

            
90
    /// Overrides the protocol version. Only for testing purposes.
91
    #[cfg(feature = "test-util")]
92
    #[allow(clippy::missing_const_for_fn)]
93
2
    pub fn with_protocol_version(mut self, version: &'static str) -> Self {
94
2
        self.protocol_version = version;
95
2
        self
96
2
    }
97

            
98
    /// Sets the request timeout for the client.
99
    ///
100
    /// If not specified, requests will time out after 60 seconds.
101
4
    pub fn with_request_timeout(mut self, timeout: impl Into<Duration>) -> Self {
102
4
        self.request_timeout = Some(timeout.into());
103
4
        self
104
4
    }
105

            
106
    /// Sets the connection timeout for the client.
107
    ///
108
    /// If not specified, the client will time out after 60 seconds if a
109
    /// connection cannot be established.
110
4
    pub fn with_connect_timeout(mut self, timeout: impl Into<Duration>) -> Self {
111
4
        self.connect_timeout = Some(timeout.into());
112
4
        self
113
4
    }
114

            
115
1818
    fn finish_internal(self) -> Result<AsyncClient, Error> {
116
1818
        AsyncClient::new_from_parts(
117
1818
            self.url,
118
1818
            self.protocol_version,
119
1818
            self.custom_apis,
120
1818
            self.connect_timeout,
121
1818
            self.request_timeout,
122
1818
            #[cfg(not(target_arch = "wasm32"))]
123
1818
            self.certificate,
124
1818
            #[cfg(not(target_arch = "wasm32"))]
125
1818
            self.tokio.or_else(|| Handle::try_current().ok()),
126
1818
        )
127
1818
    }
128
}
129

            
130
#[cfg(not(target_arch = "wasm32"))]
131
impl Builder<Blocking> {
132
    /// Finishes building the client for use in a blocking (not async) context.
133
72
    pub fn build(self) -> Result<BlockingClient, Error> {
134
72
        self.finish_internal().map(BlockingClient)
135
72
    }
136
}
137

            
138
impl Builder<Async> {
139
    /// Finishes building the client for use in a tokio async context.
140
1746
    pub fn build(self) -> Result<AsyncClient, Error> {
141
1746
        self.finish_internal()
142
1746
    }
143
}