weaviate_community/lib.rs
1//! # weaviate-community
2//!
3//! Community client for handling Weaviate vector database transactions written in Rust, for Rust.
4//! More information on Weaviate can be found on the official Weaviate webpage.
5mod backups;
6mod batch;
7mod classification;
8pub mod collections;
9mod meta;
10mod modules;
11mod nodes;
12mod objects;
13mod oidc;
14mod query;
15mod schema;
16pub use self::backups::Backups;
17pub use self::batch::Batch;
18pub use self::classification::Classification;
19pub use self::meta::Meta;
20pub use self::modules::Modules;
21pub use self::nodes::Nodes;
22pub use self::objects::Objects;
23pub use self::oidc::Oidc;
24pub use self::query::Query;
25pub use self::schema::Schema;
26use collections::auth::{ApiKey, AuthApiKey};
27
28use std::error::Error;
29use std::sync::Arc;
30
31use reqwest::header::{HeaderMap, AUTHORIZATION};
32use reqwest::Url;
33
34/// An asynchronous `WeaviateClient` to interact with a Weaviate database.
35#[derive(Debug)]
36pub struct WeaviateClient {
37 pub base_url: Url,
38 client: Arc<reqwest::Client>,
39 pub schema: Schema,
40 pub objects: Objects,
41 pub batch: Batch,
42 pub backups: Backups,
43 pub classification: Classification,
44 pub meta: Meta,
45 pub nodes: Nodes,
46 pub oidc: Oidc,
47 pub modules: Modules,
48 pub query: Query,
49}
50
51impl WeaviateClient {
52 /// Construct a new `WeaviateClient`
53 ///
54 /// # Parameters
55 /// - url: the root url for the client
56 /// - auth_client_secret: the API authentication key
57 ///
58 /// # Example
59 /// Using the WeaviateClient
60 /// ```
61 /// use std::collections::HashMap;
62 /// use weaviate_community::WeaviateClient;
63 /// use weaviate_community::collections::auth::{AuthApiKey, ApiKey};
64 ///
65 /// #[tokio::main]
66 /// async fn main() -> Result<(), Box<dyn std::error::Error>> {
67 /// let auth = AuthApiKey::new("test-key");
68 /// let client = WeaviateClient::new(
69 /// "http://localhost:8080",
70 /// Some(auth),
71 /// Some(vec![]),
72 /// )?;
73 ///
74 /// Ok(())
75 /// }
76 /// ```
77 ///
78 /// Using the WeaviateClientBuilder
79 /// ```
80 /// use weaviate_community::WeaviateClient;
81 ///
82 /// #[tokio::main]
83 /// async fn main() -> Result<(), Box<dyn std::error::Error>> {
84 /// let client = WeaviateClient::builder("http://localhost:8080")
85 /// .with_auth_secret("test-key")
86 /// .with_api_key("X-OpenAI-Api-Key", "your-key")
87 /// .build();
88 /// Ok(())
89 /// }
90 /// ```
91 pub fn new(
92 url: &str,
93 auth_client_secret: Option<AuthApiKey>,
94 api_keys: Option<Vec<ApiKey>>,
95 ) -> Result<Self, Box<dyn Error>> {
96 let base = Url::parse(url)?;
97 let mut client_builder = reqwest::Client::builder();
98
99 let mut headers = HeaderMap::new();
100
101 // Add the authorization header to the client if it is present
102 if let Some(auth) = auth_client_secret {
103 headers.insert(AUTHORIZATION, auth.get_header_value());
104 };
105
106 // Add any of the other header keys to the client, for example, OpenAI
107 if let Some(keys) = api_keys {
108 for i in 0..keys.len() {
109 headers.insert(
110 keys.get(i).unwrap().get_header_name(),
111 keys.get(i).unwrap().get_header_value(),
112 );
113 }
114 }
115
116 client_builder = client_builder.default_headers(headers);
117
118 // Each of the endpoint categories hold a strong ref to the main client.
119 let client = Arc::new(client_builder.build()?);
120 let schema = Schema::new(&base, Arc::clone(&client))?;
121 let objects = Objects::new(&base, Arc::clone(&client))?;
122 let batch = Batch::new(&base, Arc::clone(&client))?;
123 let backups = Backups::new(&base, Arc::clone(&client))?;
124 let classification = Classification::new(&base, Arc::clone(&client))?;
125 let meta = Meta::new(&base, Arc::clone(&client))?;
126 let nodes = Nodes::new(&base, Arc::clone(&client))?;
127 let oidc = Oidc::new(&base, Arc::clone(&client))?;
128 let modules = Modules::new(&base, Arc::clone(&client))?;
129 let query = Query::new(&base, Arc::clone(&client))?;
130
131 Ok(WeaviateClient {
132 base_url: base,
133 client,
134 schema,
135 objects,
136 batch,
137 backups,
138 classification,
139 meta,
140 nodes,
141 oidc,
142 modules,
143 query,
144 })
145 }
146
147 /// Determine if the application is ready to receive traffic.
148 ///
149 /// More info on the liveness can be found [here](https://weaviate.io/developers/weaviate/api/rest/well-known#liveness)
150 ///
151 /// GET /v1/.well-known/live
152 ///
153 /// Endpoint returns HTTP status code 200 if the application is able to respond to HTTP
154 /// requests.
155 ///
156 /// # Returns
157 /// * Ok(bool) => True if 200, False otherwise
158 ///
159 /// # Errors
160 /// When there is a reqwest error
161 ///
162 /// # Example
163 /// ```
164 /// use weaviate_community::WeaviateClient;
165 ///
166 /// #[tokio::main]
167 /// async fn main() -> Result<(), Box<dyn std::error::Error>> {
168 /// let client = WeaviateClient::builder("http://localhost:8080")
169 /// .with_auth_secret("test-key")
170 /// .build()?;
171 /// let res = client.is_live().await;
172 /// Ok(())
173 /// }
174 /// ```
175 pub async fn is_live(&self) -> Result<bool, Box<dyn Error>> {
176 let endpoint = self.base_url.join("/v1/.well-known/live")?;
177 let resp = self.client.get(endpoint).send().await?;
178 match resp.status() {
179 reqwest::StatusCode::OK => Ok(true),
180 _ => Ok(false),
181 }
182 }
183
184 /// Determine if the application is ready to receive traffic.
185 ///
186 /// More info on the readiness can be found [here](https://weaviate.io/developers/weaviate/api/rest/well-known#readiness)
187 ///
188 /// GET /v1/.well-known/ready
189 ///
190 /// # Example
191 /// ```
192 /// use weaviate_community::WeaviateClient;
193 ///
194 /// #[tokio::main]
195 /// async fn main() -> Result<(), Box<dyn std::error::Error>> {
196 /// let client = WeaviateClient::builder("http://localhost:8080")
197 /// .with_auth_secret("test-key")
198 /// .build()?;
199 /// let res = client.is_ready().await;
200 /// Ok(())
201 /// }
202 /// ```
203 pub async fn is_ready(&self) -> Result<bool, Box<dyn Error>> {
204 let endpoint = self.base_url.join("/v1/.well-known/ready")?;
205 let resp = self.client.get(endpoint).send().await?;
206 match resp.status() {
207 reqwest::StatusCode::OK => Ok(true),
208 _ => Ok(false),
209 }
210 }
211
212 /// Builder for the WeaviateClient
213 ///
214 /// # Parameters
215 /// - base_url: the root url for the client
216 ///
217 /// # Examples
218 /// Anonymous
219 /// ```
220 /// use weaviate_community::WeaviateClient;
221 /// let client = WeaviateClient::builder("http://localhost:8080").build();
222 /// ```
223 ///
224 /// Authenticated with API key
225 /// ```
226 /// use weaviate_community::WeaviateClient;
227 ///
228 /// let client = WeaviateClient::builder("http://localhost:8080")
229 /// .with_auth_secret("your-key")
230 /// .build();
231 /// ```
232 pub fn builder(base_url: &str) -> WeaviateClientBuilder {
233 WeaviateClientBuilder::new(base_url)
234 }
235}
236
237/// A `WeaviateClientBuilder` can be used to create a new `WeaviateClient`.
238#[derive(Default, Debug)]
239pub struct WeaviateClientBuilder {
240 pub base_url: String,
241 pub auth_secret: Option<AuthApiKey>,
242 pub api_keys: Vec<ApiKey>,
243}
244
245impl WeaviateClientBuilder {
246 /// Construct a new `WeaviateClientBuilder`.
247 ///
248 /// # Parameters
249 /// - base_url: the root url for the client
250 ///
251 /// This is the same as `WeaviateClient::builder()`.
252 ///
253 /// # Examples
254 /// Anonymous
255 /// ```
256 /// use weaviate_community::WeaviateClientBuilder;
257 /// let client = WeaviateClientBuilder::new("http://localhost:8080").build();
258 /// ```
259 ///
260 /// Authenticated with API key
261 /// ```
262 /// use weaviate_community::WeaviateClientBuilder;
263 ///
264 /// let client = WeaviateClientBuilder::new("http://localhost:8080")
265 /// .with_auth_secret("your-key")
266 /// .build();
267 /// ```
268 pub fn new(base_url: &str) -> WeaviateClientBuilder {
269 WeaviateClientBuilder {
270 base_url: base_url.into(),
271 auth_secret: None,
272 api_keys: Vec::new(),
273 }
274 }
275
276 /// Sets the authentication token to be used by the client.
277 ///
278 /// # Parameters
279 /// - auth_secret: the AuthApiKey to set in the client
280 ///
281 /// # Example
282 /// ```
283 /// use weaviate_community::WeaviateClientBuilder;
284 ///
285 /// let client = WeaviateClientBuilder::new("http://localhost:8080")
286 /// .with_auth_secret("your-key")
287 /// .build();
288 /// ```
289 pub fn with_auth_secret(mut self, auth_secret: &str) -> WeaviateClientBuilder {
290 self.auth_secret = Some(AuthApiKey::new(auth_secret));
291 self
292 }
293
294 /// Sets a new api key to be used by the client.
295 ///
296 /// # Parameters
297 /// - header: the header to set in the client
298 /// - api_key: the api key
299 ///
300 /// # Example
301 /// ```
302 /// use weaviate_community::WeaviateClientBuilder;
303 ///
304 /// let client = WeaviateClientBuilder::new("http://localhost:8080")
305 /// .with_api_key("X-OpenAI-Api-Key", "abcdefg")
306 /// .build();
307 /// ```
308 pub fn with_api_key(mut self, header: &str, api_key: &str) -> WeaviateClientBuilder {
309 self.api_keys.push(ApiKey {
310 api_header: header.into(),
311 api_key: api_key.into(),
312 });
313 self
314 }
315
316 /// Build a `WeaviateClient` from the values set in the WeaviateClientBuilder.
317 ///
318 /// # Example
319 /// ```
320 /// use weaviate_community::WeaviateClientBuilder;
321 ///
322 /// let client = WeaviateClientBuilder::new("http://localhost:8080").build();
323 /// ```
324 pub fn build(self) -> Result<WeaviateClient, Box<dyn Error>> {
325 let client = WeaviateClient::new(&self.base_url, self.auth_secret, Some(self.api_keys))?;
326 Ok(client)
327 }
328}
329
330#[cfg(test)]
331mod tests {
332 use super::*;
333
334 fn get_test_harness() -> (mockito::ServerGuard, WeaviateClient) {
335 let mock_server = mockito::Server::new();
336 let mut host = "http://".to_string();
337 host.push_str(&mock_server.host_with_port());
338 let client = WeaviateClient::builder(&host).build().unwrap();
339 (mock_server, client)
340 }
341
342 fn mock_get(
343 server: &mut mockito::ServerGuard,
344 endpoint: &str,
345 status_code: usize,
346 body: &str,
347 ) -> mockito::Mock {
348 server
349 .mock("GET", endpoint)
350 .with_status(status_code)
351 .with_header("content-type", "application/json")
352 .with_body(body)
353 .create()
354 }
355
356 #[tokio::test]
357 async fn test_is_ready_ok() {
358 let (mut mock_server, client) = get_test_harness();
359 let mock = mock_get(&mut mock_server, "/v1/.well-known/ready", 200, "");
360 let res = client.is_ready().await;
361 mock.assert();
362 assert!(res.is_ok());
363 }
364
365 #[tokio::test]
366 async fn test_is_ready_err() {
367 let (mut mock_server, client) = get_test_harness();
368 let mock = mock_get(&mut mock_server, "/v1/.well-known/ready", 503, "");
369 let res = client.is_ready().await;
370 mock.assert();
371 assert!(res.is_ok());
372 assert_eq!(false, res.unwrap());
373 }
374
375 #[tokio::test]
376 async fn test_is_live_ok() {
377 let (mut mock_server, client) = get_test_harness();
378 let mock = mock_get(&mut mock_server, "/v1/.well-known/live", 200, "");
379 let res = client.is_live().await;
380 mock.assert();
381 assert!(res.is_ok());
382 }
383
384 #[tokio::test]
385 async fn test_is_live_err() {
386 let (mut mock_server, client) = get_test_harness();
387 let mock = mock_get(&mut mock_server, "/v1/.well-known/live", 404, "");
388 let res = client.is_live().await;
389 mock.assert();
390 assert!(res.is_ok());
391 assert_eq!(false, res.unwrap());
392 }
393}