kstone_client/
transaction.rs

1/// Remote transaction 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 transact get request builder
9pub struct RemoteTransactGetRequest {
10    keys: Vec<proto::Key>,
11}
12
13impl RemoteTransactGetRequest {
14    /// Create a new transact 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 get(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 get_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 transact get operation
38    pub async fn execute(self, client: &mut KeystoneDbClient<Channel>) -> Result<RemoteTransactGetResponse> {
39        let request = proto::TransactGetRequest {
40            keys: self.keys,
41        };
42
43        let response = client
44            .transact_get(request)
45            .await?
46            .into_inner();
47
48        // Convert protobuf response to Rust types
49        let items: Vec<Option<Item>> = response
50            .items
51            .into_iter()
52            .map(|tx_item| {
53                tx_item.item.map(|proto_item| {
54                    proto_item_to_ks(proto_item).expect("Invalid item from server")
55                })
56            })
57            .collect();
58
59        Ok(RemoteTransactGetResponse { items })
60    }
61}
62
63impl Default for RemoteTransactGetRequest {
64    fn default() -> Self {
65        Self::new()
66    }
67}
68
69/// Transact get response
70pub struct RemoteTransactGetResponse {
71    /// Items retrieved (in same order as request, None if not found)
72    pub items: Vec<Option<Item>>,
73}
74
75/// Remote transact write request builder
76pub struct RemoteTransactWriteRequest {
77    writes: Vec<proto::TransactWriteItem>,
78}
79
80impl RemoteTransactWriteRequest {
81    /// Create a new transact write request
82    pub fn new() -> Self {
83        Self { writes: Vec::new() }
84    }
85
86    /// Add a put request
87    pub fn put(mut self, pk: &[u8], item: Item) -> Self {
88        self.writes.push(proto::TransactWriteItem {
89            item: Some(proto::transact_write_item::Item::Put(proto::TransactPut {
90                partition_key: pk.to_vec(),
91                sort_key: None,
92                item: Some(ks_item_to_proto(&item)),
93                condition_expression: None,
94            })),
95        });
96        self
97    }
98
99    /// Add a put request with sort key
100    pub fn put_with_sk(mut self, pk: &[u8], sk: &[u8], item: Item) -> Self {
101        self.writes.push(proto::TransactWriteItem {
102            item: Some(proto::transact_write_item::Item::Put(proto::TransactPut {
103                partition_key: pk.to_vec(),
104                sort_key: Some(sk.to_vec()),
105                item: Some(ks_item_to_proto(&item)),
106                condition_expression: None,
107            })),
108        });
109        self
110    }
111
112    /// Add an update request
113    pub fn update(mut self, pk: &[u8], update_expression: impl Into<String>) -> Self {
114        self.writes.push(proto::TransactWriteItem {
115            item: Some(proto::transact_write_item::Item::Update(proto::TransactUpdate {
116                partition_key: pk.to_vec(),
117                sort_key: None,
118                update_expression: update_expression.into(),
119                condition_expression: None,
120            })),
121        });
122        self
123    }
124
125    /// Add a delete request
126    pub fn delete(mut self, pk: &[u8]) -> Self {
127        self.writes.push(proto::TransactWriteItem {
128            item: Some(proto::transact_write_item::Item::Delete(proto::TransactDelete {
129                partition_key: pk.to_vec(),
130                sort_key: None,
131                condition_expression: None,
132            })),
133        });
134        self
135    }
136
137    /// Add a condition check
138    pub fn condition_check(mut self, pk: &[u8], condition: impl Into<String>) -> Self {
139        self.writes.push(proto::TransactWriteItem {
140            item: Some(proto::transact_write_item::Item::ConditionCheck(proto::ConditionCheck {
141                partition_key: pk.to_vec(),
142                sort_key: None,
143                condition_expression: condition.into(),
144            })),
145        });
146        self
147    }
148
149    /// Execute the transact write operation
150    pub async fn execute(self, client: &mut KeystoneDbClient<Channel>) -> Result<()> {
151        let request = proto::TransactWriteRequest {
152            items: self.writes,
153        };
154
155        client
156            .transact_write(request)
157            .await?;
158
159        Ok(())
160    }
161}
162
163impl Default for RemoteTransactWriteRequest {
164    fn default() -> Self {
165        Self::new()
166    }
167}
kstone_client/transaction.rs

kstone_client/
transaction.rs
