kstone_client/
batch.rs

1/// Remote batch operations
2use crate::convert::*;
3use crate::error::Result;
4use kstone_core::Item;
5use kstone_proto::{self as proto, keystone_db_client::KeystoneDbClient};
6use tonic::transport::Channel;
7
8/// Remote batch get request builder
9pub struct RemoteBatchGetRequest {
10    keys: Vec<proto::Key>,
11}
12
13impl RemoteBatchGetRequest {
14    /// Create a new batch get request
15    pub fn new() -> Self {
16        Self { keys: Vec::new() }
17    }
18
19    /// Add a key with partition key only
20    pub fn add_key(mut self, pk: &[u8]) -> Self {
21        self.keys.push(proto::Key {
22            partition_key: pk.to_vec(),
23            sort_key: None,
24        });
25        self
26    }
27
28    /// Add a key with partition key and sort key
29    pub fn add_key_with_sk(mut self, pk: &[u8], sk: &[u8]) -> Self {
30        self.keys.push(proto::Key {
31            partition_key: pk.to_vec(),
32            sort_key: Some(sk.to_vec()),
33        });
34        self
35    }
36
37    /// Execute the batch get operation
38    pub async fn execute(self, client: &mut KeystoneDbClient<Channel>) -> Result<RemoteBatchGetResponse> {
39        let request = proto::BatchGetRequest {
40            keys: self.keys,
41        };
42
43        let response = client
44            .batch_get(request)
45            .await?
46            .into_inner();
47
48        // Convert protobuf response to Rust types
49        let items: Vec<Item> = response
50            .items
51            .into_iter()
52            .map(|proto_item| proto_item_to_ks(proto_item).expect("Invalid item from server"))
53            .collect();
54
55        Ok(RemoteBatchGetResponse {
56            items,
57            count: response.count as usize,
58        })
59    }
60}
61
62impl Default for RemoteBatchGetRequest {
63    fn default() -> Self {
64        Self::new()
65    }
66}
67
68/// Batch get response
69pub struct RemoteBatchGetResponse {
70    /// Items retrieved
71    pub items: Vec<Item>,
72    /// Number of items returned
73    pub count: usize,
74}
75
76/// Remote batch write request builder
77pub struct RemoteBatchWriteRequest {
78    writes: Vec<proto::WriteRequest>,
79}
80
81impl RemoteBatchWriteRequest {
82    /// Create a new batch write request
83    pub fn new() -> Self {
84        Self { writes: Vec::new() }
85    }
86
87    /// Add a put request with partition key
88    pub fn put(mut self, pk: &[u8], item: Item) -> Self {
89        self.writes.push(proto::WriteRequest {
90            request: Some(proto::write_request::Request::Put(proto::PutItem {
91                partition_key: pk.to_vec(),
92                sort_key: None,
93                item: Some(ks_item_to_proto(&item)),
94            })),
95        });
96        self
97    }
98
99    /// Add a put request with partition key and sort key
100    pub fn put_with_sk(mut self, pk: &[u8], sk: &[u8], item: Item) -> Self {
101        self.writes.push(proto::WriteRequest {
102            request: Some(proto::write_request::Request::Put(proto::PutItem {
103                partition_key: pk.to_vec(),
104                sort_key: Some(sk.to_vec()),
105                item: Some(ks_item_to_proto(&item)),
106            })),
107        });
108        self
109    }
110
111    /// Add a delete request with partition key
112    pub fn delete(mut self, pk: &[u8]) -> Self {
113        self.writes.push(proto::WriteRequest {
114            request: Some(proto::write_request::Request::Delete(proto::DeleteKey {
115                partition_key: pk.to_vec(),
116                sort_key: None,
117            })),
118        });
119        self
120    }
121
122    /// Add a delete request with partition key and sort key
123    pub fn delete_with_sk(mut self, pk: &[u8], sk: &[u8]) -> Self {
124        self.writes.push(proto::WriteRequest {
125            request: Some(proto::write_request::Request::Delete(proto::DeleteKey {
126                partition_key: pk.to_vec(),
127                sort_key: Some(sk.to_vec()),
128            })),
129        });
130        self
131    }
132
133    /// Execute the batch write operation
134    pub async fn execute(self, client: &mut KeystoneDbClient<Channel>) -> Result<RemoteBatchWriteResponse> {
135        let request = proto::BatchWriteRequest {
136            writes: self.writes,
137        };
138
139        let response = client
140            .batch_write(request)
141            .await?
142            .into_inner();
143
144        Ok(RemoteBatchWriteResponse {
145            success: response.success,
146        })
147    }
148}
149
150impl Default for RemoteBatchWriteRequest {
151    fn default() -> Self {
152        Self::new()
153    }
154}
155
156/// Batch write response
157pub struct RemoteBatchWriteResponse {
158    /// Whether the batch write succeeded
159    pub success: bool,
160}
kstone_client/batch.rs

kstone_client/
batch.rs
