tcvectordb_rust/
database.rs

1use serde::{Deserialize, Serialize};
2use serde_json::Value;
3
4use crate::client::VectorDBClient;
5use crate::collection::Collection;
6use crate::error::{Result, VectorDBError};
7use crate::index::Index;
8
9#[derive(Debug, Clone)]
10pub struct Database {
11    client: VectorDBClient,
12    name: String,
13}
14
15#[derive(Debug, Serialize)]
16struct CreateDatabaseRequest {
17    database: String,
18}
19
20#[derive(Debug, Serialize)]
21struct CreateCollectionRequest {
22    database: String,
23    collection: String,
24    shard: u32,
25    replicas: u32,
26    #[serde(skip_serializing_if = "Option::is_none")]
27    description: Option<String>,
28    #[serde(skip_serializing_if = "Option::is_none")]
29    index: Option<Index>,
30    #[serde(skip_serializing_if = "Option::is_none")]
31    embedding: Option<Value>, // TODO: Define proper Embedding struct
32    #[serde(skip_serializing_if = "Option::is_none")]
33    ttl_config: Option<Value>,
34    #[serde(skip_serializing_if = "Option::is_none")]
35    filter_index_config: Option<Value>,
36}
37
38#[derive(Debug, Deserialize)]
39struct ApiResponse<T> {
40    code: i32,
41    msg: String,
42    #[serde(default)]
43    data: Option<T>,
44}
45
46impl Database {
47    pub fn new(client: VectorDBClient, name: String) -> Self {
48        Self { client, name }
49    }
50
51    pub fn name(&self) -> &str {
52        &self.name
53    }
54
55    pub fn client(&self) -> &VectorDBClient {
56        &self.client
57    }
58
59    /// Create the database
60    pub async fn create(&self) -> Result<Value> {
61        let request = CreateDatabaseRequest {
62            database: self.name.clone(),
63        };
64
65        let response = self.client
66            .post("/database/create")
67            .await?
68            .json(&request)
69            .send()
70            .await?;
71
72        let api_response: ApiResponse<Value> = self.client.handle_response(response).await?;
73        
74        if api_response.code != 0 {
75            return Err(VectorDBError::server_error(api_response.code, api_response.msg));
76        }
77
78        Ok(api_response.data.unwrap_or(Value::Null))
79    }
80
81    /// Drop the database
82    pub async fn drop(&self) -> Result<Value> {
83        let path = format!("/database/drop?database={}", self.name);
84        let response = self.client
85            .delete(&path)
86            .await?
87            .send()
88            .await?;
89
90        let api_response: ApiResponse<Value> = self.client.handle_response(response).await?;
91        
92        if api_response.code != 0 {
93            return Err(VectorDBError::server_error(api_response.code, api_response.msg));
94        }
95
96        Ok(api_response.data.unwrap_or(Value::Null))
97    }
98
99    /// Check if a collection exists
100    pub async fn exists_collection(&self, collection_name: &str) -> Result<bool> {
101        let collections = self.list_collections().await?;
102        Ok(collections.iter().any(|c| c.name() == collection_name))
103    }
104
105    /// Create a new collection
106    pub async fn create_collection(
107        &self,
108        name: impl Into<String>,
109        shard: u32,
110        replicas: u32,
111        description: Option<String>,
112        index: Option<Index>,
113        embedding: Option<Value>,
114        ttl_config: Option<Value>,
115    ) -> Result<Collection> {
116        let collection_name = name.into();
117        let request = CreateCollectionRequest {
118            database: self.name.clone(),
119            collection: collection_name.clone(),
120            shard,
121            replicas,
122            description,
123            index,
124            embedding,
125            ttl_config,
126            filter_index_config: None,
127        };
128
129        let response = self.client
130            .post("/collection/create")
131            .await?
132            .json(&request)
133            .send()
134            .await?;
135
136        let api_response: ApiResponse<Value> = self.client.handle_response(response).await?;
137        
138        if api_response.code != 0 {
139            return Err(VectorDBError::server_error(api_response.code, api_response.msg));
140        }
141
142        Ok(Collection::new(self.clone(), collection_name))
143    }
144
145    /// Create a collection if it doesn't exist
146    pub async fn create_collection_if_not_exists(
147        &self,
148        name: impl Into<String>,
149        shard: u32,
150        replicas: u32,
151        description: Option<String>,
152        index: Option<Index>,
153        embedding: Option<Value>,
154        ttl_config: Option<Value>,
155    ) -> Result<Collection> {
156        let collection_name = name.into();
157        if self.exists_collection(&collection_name).await? {
158            Ok(Collection::new(self.clone(), collection_name))
159        } else {
160            self.create_collection(collection_name, shard, replicas, description, index, embedding, ttl_config).await
161        }
162    }
163
164    /// List all collections in the database
165    pub async fn list_collections(&self) -> Result<Vec<Collection>> {
166        let path = format!("/collection/list?database={}", self.name);
167        let response = self.client
168            .get(&path)
169            .await?
170            .send()
171            .await?;
172
173        let api_response: ApiResponse<Value> = self.client.handle_response(response).await?;
174        
175        if api_response.code != 0 {
176            return Err(VectorDBError::server_error(api_response.code, api_response.msg));
177        }
178
179        let collections_data = api_response.data.unwrap_or(Value::Array(vec![]));
180        let collection_names: Vec<String> = serde_json::from_value(collections_data)?;
181        
182        Ok(collection_names
183            .into_iter()
184            .map(|name| Collection::new(self.clone(), name))
185            .collect())
186    }
187
188    /// Get a collection by name
189    pub async fn collection(&self, collection_name: impl Into<String>) -> Result<Collection> {
190        let collection_name = collection_name.into();
191        if self.exists_collection(&collection_name).await? {
192            Ok(Collection::new(self.clone(), collection_name))
193        } else {
194            Err(VectorDBError::param_error(
195                14100,
196                format!("Collection not exist: {}", collection_name),
197            ))
198        }
199    }
200
201    /// Describe a collection
202    pub async fn describe_collection(&self, collection_name: impl Into<String>) -> Result<Collection> {
203        let collection_name = collection_name.into();
204        let path = format!("/collection/describe?database={}&collection={}", self.name, collection_name);
205        let response = self.client
206            .get(&path)
207            .await?
208            .send()
209            .await?;
210
211        let api_response: ApiResponse<Value> = self.client.handle_response(response).await?;
212        
213        if api_response.code != 0 {
214            return Err(VectorDBError::server_error(api_response.code, api_response.msg));
215        }
216
217        // TODO: Parse collection details from response
218        Ok(Collection::new(self.clone(), collection_name))
219    }
220
221    /// Drop a collection
222    pub async fn drop_collection(&self, collection_name: impl Into<String>) -> Result<Value> {
223        let collection_name = collection_name.into();
224        let path = format!("/collection/drop?database={}&collection={}", self.name, collection_name);
225        let response = self.client
226            .delete(&path)
227            .await?
228            .send()
229            .await?;
230
231        let api_response: ApiResponse<Value> = self.client.handle_response(response).await?;
232        
233        if api_response.code != 0 {
234            return Err(VectorDBError::server_error(api_response.code, api_response.msg));
235        }
236
237        Ok(api_response.data.unwrap_or(Value::Null))
238    }
239
240    /// Truncate a collection (clear all data)
241    pub async fn truncate_collection(&self, collection_name: impl Into<String>) -> Result<Value> {
242        let collection_name = collection_name.into();
243        let path = format!("/collection/truncate?database={}&collection={}", self.name, collection_name);
244        let response = self.client
245            .post(&path)
246            .await?
247            .send()
248            .await?;
249
250        let api_response: ApiResponse<Value> = self.client.handle_response(response).await?;
251        
252        if api_response.code != 0 {
253            return Err(VectorDBError::server_error(api_response.code, api_response.msg));
254        }
255
256        Ok(api_response.data.unwrap_or(Value::Null))
257    }
258
259    /// Set alias for a collection
260    pub async fn set_alias(&self, collection_name: impl Into<String>, alias: impl Into<String>) -> Result<Value> {
261        let collection_name = collection_name.into();
262        let alias = alias.into();
263        let path = format!("/collection/alias/set?database={}&collection={}&alias={}", 
264                          self.name, collection_name, alias);
265        let response = self.client
266            .post(&path)
267            .await?
268            .send()
269            .await?;
270
271        let api_response: ApiResponse<Value> = self.client.handle_response(response).await?;
272        
273        if api_response.code != 0 {
274            return Err(VectorDBError::server_error(api_response.code, api_response.msg));
275        }
276
277        Ok(api_response.data.unwrap_or(Value::Null))
278    }
279
280    /// Delete alias
281    pub async fn delete_alias(&self, alias: impl Into<String>) -> Result<Value> {
282        let alias = alias.into();
283        let path = format!("/collection/alias/delete?database={}&alias={}", self.name, alias);
284        let response = self.client
285            .delete(&path)
286            .await?
287            .send()
288            .await?;
289
290        let api_response: ApiResponse<Value> = self.client.handle_response(response).await?;
291        
292        if api_response.code != 0 {
293            return Err(VectorDBError::server_error(api_response.code, api_response.msg));
294        }
295
296        Ok(api_response.data.unwrap_or(Value::Null))
297    }
298}
tcvectordb_rust/database.rs

tcvectordb_rust/
database.rs
