kstone_client/client.rs
1/// KeystoneDB gRPC client implementation
2use crate::error::{ClientError, Result};
3use kstone_core::Item;
4use kstone_proto::{self as proto, keystone_db_client::KeystoneDbClient};
5use tonic::transport::Channel;
6
7/// KeystoneDB remote client
8pub struct Client {
9 inner: KeystoneDbClient<Channel>,
10}
11
12impl Client {
13 /// Connect to a KeystoneDB server
14 ///
15 /// # Arguments
16 /// * `addr` - Server address (e.g., "http://127.0.0.1:50051")
17 ///
18 /// # Example
19 /// ```no_run
20 /// # use kstone_client::Client;
21 /// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
22 /// let client = Client::connect("http://localhost:50051").await?;
23 /// # Ok(())
24 /// # }
25 /// ```
26 pub async fn connect(addr: impl Into<String>) -> Result<Self> {
27 let addr = addr.into();
28 let channel = Channel::from_shared(addr)
29 .map_err(|e| ClientError::ConnectionError(format!("Invalid address: {}", e)))?
30 .connect()
31 .await
32 .map_err(|e| ClientError::ConnectionError(format!("Failed to connect: {}", e)))?;
33
34 let inner = KeystoneDbClient::new(channel);
35 Ok(Self { inner })
36 }
37
38 /// Put an item with a simple partition key
39 ///
40 /// # Arguments
41 /// * `pk` - Partition key
42 /// * `item` - Item to store
43 ///
44 /// # Example
45 /// ```no_run
46 /// # use kstone_client::Client;
47 /// # use std::collections::HashMap;
48 /// # use kstone_core::Value;
49 /// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
50 /// let mut client = Client::connect("http://localhost:50051").await?;
51 ///
52 /// let mut item = HashMap::new();
53 /// item.insert("name".to_string(), Value::S("Alice".to_string()));
54 /// item.insert("age".to_string(), Value::N("30".to_string()));
55 ///
56 /// client.put(b"user#123", item).await?;
57 /// # Ok(())
58 /// # }
59 /// ```
60 pub async fn put(&mut self, pk: &[u8], item: Item) -> Result<()> {
61 let request = proto::PutRequest {
62 partition_key: pk.to_vec(),
63 sort_key: None,
64 item: Some(crate::convert::ks_item_to_proto(&item)),
65 condition_expression: None,
66 expression_values: std::collections::HashMap::new(),
67 };
68
69 self.inner
70 .put(request)
71 .await
72 .map_err(|e| e.into())
73 .map(|_| ())
74 }
75
76 /// Put an item with partition key and sort key
77 ///
78 /// # Arguments
79 /// * `pk` - Partition key
80 /// * `sk` - Sort key
81 /// * `item` - Item to store
82 pub async fn put_with_sk(&mut self, pk: &[u8], sk: &[u8], item: Item) -> Result<()> {
83 let request = proto::PutRequest {
84 partition_key: pk.to_vec(),
85 sort_key: Some(sk.to_vec()),
86 item: Some(crate::convert::ks_item_to_proto(&item)),
87 condition_expression: None,
88 expression_values: std::collections::HashMap::new(),
89 };
90
91 self.inner
92 .put(request)
93 .await
94 .map_err(|e| e.into())
95 .map(|_| ())
96 }
97
98 /// Put an item with a condition expression
99 ///
100 /// # Arguments
101 /// * `pk` - Partition key
102 /// * `item` - Item to store
103 /// * `condition` - Condition expression (e.g., "attribute_not_exists(pk)")
104 /// * `values` - Expression attribute values
105 pub async fn put_conditional(
106 &mut self,
107 pk: &[u8],
108 item: Item,
109 condition: impl Into<String>,
110 values: std::collections::HashMap<String, kstone_core::Value>,
111 ) -> Result<()> {
112 let proto_values: std::collections::HashMap<String, proto::Value> = values
113 .iter()
114 .map(|(k, v)| (k.clone(), crate::convert::ks_value_to_proto(v)))
115 .collect();
116
117 let request = proto::PutRequest {
118 partition_key: pk.to_vec(),
119 sort_key: None,
120 item: Some(crate::convert::ks_item_to_proto(&item)),
121 condition_expression: Some(condition.into()),
122 expression_values: proto_values,
123 };
124
125 self.inner
126 .put(request)
127 .await
128 .map_err(|e| e.into())
129 .map(|_| ())
130 }
131
132 /// Get an item with a simple partition key
133 ///
134 /// # Arguments
135 /// * `pk` - Partition key
136 ///
137 /// # Returns
138 /// The item if found, None otherwise
139 pub async fn get(&mut self, pk: &[u8]) -> Result<Option<Item>> {
140 let request = proto::GetRequest {
141 partition_key: pk.to_vec(),
142 sort_key: None,
143 };
144
145 let response = self
146 .inner
147 .get(request)
148 .await
149 .map_err(|e| ClientError::from(e))?
150 .into_inner();
151
152 Ok(response.item.map(|proto_item| {
153 crate::convert::proto_item_to_ks(proto_item)
154 .expect("Server returned invalid item")
155 }))
156 }
157
158 /// Get an item with partition key and sort key
159 ///
160 /// # Arguments
161 /// * `pk` - Partition key
162 /// * `sk` - Sort key
163 ///
164 /// # Returns
165 /// The item if found, None otherwise
166 pub async fn get_with_sk(&mut self, pk: &[u8], sk: &[u8]) -> Result<Option<Item>> {
167 let request = proto::GetRequest {
168 partition_key: pk.to_vec(),
169 sort_key: Some(sk.to_vec()),
170 };
171
172 let response = self
173 .inner
174 .get(request)
175 .await
176 .map_err(|e| ClientError::from(e))?
177 .into_inner();
178
179 Ok(response.item.map(|proto_item| {
180 crate::convert::proto_item_to_ks(proto_item)
181 .expect("Server returned invalid item")
182 }))
183 }
184
185 /// Delete an item with a simple partition key
186 ///
187 /// # Arguments
188 /// * `pk` - Partition key
189 pub async fn delete(&mut self, pk: &[u8]) -> Result<()> {
190 let request = proto::DeleteRequest {
191 partition_key: pk.to_vec(),
192 sort_key: None,
193 condition_expression: None,
194 expression_values: std::collections::HashMap::new(),
195 };
196
197 self.inner
198 .delete(request)
199 .await
200 .map_err(|e| e.into())
201 .map(|_| ())
202 }
203
204 /// Delete an item with partition key and sort key
205 ///
206 /// # Arguments
207 /// * `pk` - Partition key
208 /// * `sk` - Sort key
209 pub async fn delete_with_sk(&mut self, pk: &[u8], sk: &[u8]) -> Result<()> {
210 let request = proto::DeleteRequest {
211 partition_key: pk.to_vec(),
212 sort_key: Some(sk.to_vec()),
213 condition_expression: None,
214 expression_values: std::collections::HashMap::new(),
215 };
216
217 self.inner
218 .delete(request)
219 .await
220 .map_err(|e| e.into())
221 .map(|_| ())
222 }
223
224 /// Delete an item with a condition expression
225 ///
226 /// # Arguments
227 /// * `pk` - Partition key
228 /// * `condition` - Condition expression (e.g., "attribute_exists(pk)")
229 /// * `values` - Expression attribute values
230 pub async fn delete_conditional(
231 &mut self,
232 pk: &[u8],
233 condition: impl Into<String>,
234 values: std::collections::HashMap<String, kstone_core::Value>,
235 ) -> Result<()> {
236 let proto_values: std::collections::HashMap<String, proto::Value> = values
237 .iter()
238 .map(|(k, v)| (k.clone(), crate::convert::ks_value_to_proto(v)))
239 .collect();
240
241 let request = proto::DeleteRequest {
242 partition_key: pk.to_vec(),
243 sort_key: None,
244 condition_expression: Some(condition.into()),
245 expression_values: proto_values,
246 };
247
248 self.inner
249 .delete(request)
250 .await
251 .map_err(|e| e.into())
252 .map(|_| ())
253 }
254
255 /// Execute a query operation
256 ///
257 /// # Arguments
258 /// * `query` - Query builder with conditions
259 ///
260 /// # Example
261 /// ```no_run
262 /// # use kstone_client::{Client, RemoteQuery};
263 /// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
264 /// let mut client = Client::connect("http://localhost:50051").await?;
265 ///
266 /// let query = RemoteQuery::new(b"user#org1")
267 /// .sk_begins_with(b"USER#")
268 /// .limit(10);
269 ///
270 /// let response = client.query(query).await?;
271 /// println!("Found {} items", response.count);
272 /// # Ok(())
273 /// # }
274 /// ```
275 pub async fn query(&mut self, query: crate::query::RemoteQuery) -> Result<crate::query::RemoteQueryResponse> {
276 query.execute(&mut self.inner).await
277 }
278
279 /// Execute a scan operation
280 ///
281 /// # Arguments
282 /// * `scan` - Scan builder with options
283 ///
284 /// # Example
285 /// ```no_run
286 /// # use kstone_client::{Client, RemoteScan};
287 /// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
288 /// let mut client = Client::connect("http://localhost:50051").await?;
289 ///
290 /// let scan = RemoteScan::new()
291 /// .limit(100);
292 ///
293 /// let response = client.scan(scan).await?;
294 /// println!("Scanned {} items", response.count);
295 /// # Ok(())
296 /// # }
297 /// ```
298 pub async fn scan(&mut self, scan: crate::scan::RemoteScan) -> Result<crate::scan::RemoteScanResponse> {
299 scan.execute(&mut self.inner).await
300 }
301
302 /// Execute a batch get operation
303 ///
304 /// # Arguments
305 /// * `request` - Batch get request with keys
306 ///
307 /// # Example
308 /// ```no_run
309 /// # use kstone_client::{Client, RemoteBatchGetRequest};
310 /// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
311 /// let mut client = Client::connect("http://localhost:50051").await?;
312 ///
313 /// let batch = RemoteBatchGetRequest::new()
314 /// .add_key(b"user#1")
315 /// .add_key(b"user#2");
316 ///
317 /// let response = client.batch_get(batch).await?;
318 /// println!("Retrieved {} items", response.count);
319 /// # Ok(())
320 /// # }
321 /// ```
322 pub async fn batch_get(&mut self, request: crate::batch::RemoteBatchGetRequest) -> Result<crate::batch::RemoteBatchGetResponse> {
323 request.execute(&mut self.inner).await
324 }
325
326 /// Execute a batch write operation
327 ///
328 /// # Arguments
329 /// * `request` - Batch write request with puts/deletes
330 ///
331 /// # Example
332 /// ```no_run
333 /// # use kstone_client::{Client, RemoteBatchWriteRequest};
334 /// # use std::collections::HashMap;
335 /// # use kstone_core::Value;
336 /// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
337 /// let mut client = Client::connect("http://localhost:50051").await?;
338 ///
339 /// let mut item = HashMap::new();
340 /// item.insert("name".to_string(), Value::S("Alice".to_string()));
341 ///
342 /// let batch = RemoteBatchWriteRequest::new()
343 /// .put(b"user#1", item.clone())
344 /// .delete(b"user#old");
345 ///
346 /// let response = client.batch_write(batch).await?;
347 /// println!("Batch write success: {}", response.success);
348 /// # Ok(())
349 /// # }
350 /// ```
351 pub async fn batch_write(&mut self, request: crate::batch::RemoteBatchWriteRequest) -> Result<crate::batch::RemoteBatchWriteResponse> {
352 request.execute(&mut self.inner).await
353 }
354
355 /// Execute a transactional get operation
356 ///
357 /// # Arguments
358 /// * `request` - TransactGet request with keys to retrieve
359 ///
360 /// # Example
361 /// ```no_run
362 /// # use kstone_client::{Client, RemoteTransactGetRequest};
363 /// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
364 /// let mut client = Client::connect("http://localhost:50051").await?;
365 ///
366 /// let request = RemoteTransactGetRequest::new()
367 /// .get(b"user#1")
368 /// .get(b"user#2");
369 ///
370 /// let response = client.transact_get(request).await?;
371 /// println!("Retrieved {} items", response.items.len());
372 /// # Ok(())
373 /// # }
374 /// ```
375 pub async fn transact_get(&mut self, request: crate::transaction::RemoteTransactGetRequest) -> Result<crate::transaction::RemoteTransactGetResponse> {
376 request.execute(&mut self.inner).await
377 }
378
379 /// Execute a transactional write operation
380 ///
381 /// # Arguments
382 /// * `request` - TransactWrite request with operations
383 ///
384 /// # Example
385 /// ```no_run
386 /// # use kstone_client::{Client, RemoteTransactWriteRequest};
387 /// # use std::collections::HashMap;
388 /// # use kstone_core::Value;
389 /// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
390 /// let mut client = Client::connect("http://localhost:50051").await?;
391 ///
392 /// let mut item = HashMap::new();
393 /// item.insert("name".to_string(), Value::S("Alice".to_string()));
394 ///
395 /// let request = RemoteTransactWriteRequest::new()
396 /// .put(b"user#1", item);
397 ///
398 /// client.transact_write(request).await?;
399 /// # Ok(())
400 /// # }
401 /// ```
402 pub async fn transact_write(&mut self, request: crate::transaction::RemoteTransactWriteRequest) -> Result<()> {
403 request.execute(&mut self.inner).await
404 }
405
406 /// Update an item using update expression
407 ///
408 /// # Arguments
409 /// * `request` - Update request with expression
410 ///
411 /// # Example
412 /// ```no_run
413 /// # use kstone_client::{Client, RemoteUpdate};
414 /// # use kstone_core::Value;
415 /// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
416 /// let mut client = Client::connect("http://localhost:50051").await?;
417 ///
418 /// let update = RemoteUpdate::new(b"user#1")
419 /// .expression("SET age = age + :inc")
420 /// .value(":inc", Value::N("1".to_string()));
421 ///
422 /// let response = client.update(update).await?;
423 /// println!("Updated item: {:?}", response.item);
424 /// # Ok(())
425 /// # }
426 /// ```
427 pub async fn update(&mut self, request: crate::update::RemoteUpdate) -> Result<crate::update::RemoteUpdateResponse> {
428 request.execute(&mut self.inner).await
429 }
430
431 /// Execute a PartiQL statement
432 ///
433 /// # Arguments
434 /// * `statement` - PartiQL SQL statement
435 ///
436 /// # Example
437 /// ```no_run
438 /// # use kstone_client::Client;
439 /// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
440 /// let mut client = Client::connect("http://localhost:50051").await?;
441 ///
442 /// let response = client.execute_statement(
443 /// "SELECT * FROM users WHERE pk = 'user#123'"
444 /// ).await?;
445 ///
446 /// println!("Query result: {:?}", response);
447 /// # Ok(())
448 /// # }
449 /// ```
450 pub async fn execute_statement(&mut self, statement: impl Into<String>) -> Result<crate::partiql::RemoteExecuteStatementResponse> {
451 let statement = statement.into();
452 let request = kstone_proto::ExecuteStatementRequest { statement };
453
454 let response = self.inner
455 .execute_statement(request)
456 .await?
457 .into_inner();
458
459 crate::partiql::parse_execute_statement_response(response)
460 }
461
462 /// Get a reference to the underlying gRPC client
463 pub(crate) fn inner_mut(&mut self) -> &mut KeystoneDbClient<Channel> {
464 &mut self.inner
465 }
466}