Skip to main content

kstone_client/
convert.rs

1/// Type conversions between protobuf and KeystoneDB types
2///
3/// This module provides bidirectional conversions for all data types
4/// used in the gRPC protocol and KeystoneDB internal representation.
5///
6/// Due to Rust's orphan rules, we use conversion functions instead of
7/// trait implementations.
8
9use bytes::Bytes;
10use kstone_core::{Key, Value as KsValue};
11use kstone_proto::{self as proto, value::Value as ProtoValueEnum};
12use std::collections::HashMap;
13use tonic::Status;
14
15pub type Item = HashMap<String, KsValue>;
16
17// ============================================================================
18// Value Conversions
19// ============================================================================
20
21/// Convert protobuf Value to KeystoneDB Value
22pub fn proto_value_to_ks(value: proto::Value) -> Result<KsValue, Status> {
23    let value_enum = value
24        .value
25        .ok_or_else(|| Status::invalid_argument("Value field is missing"))?;
26
27    match value_enum {
28        ProtoValueEnum::StringValue(s) => Ok(KsValue::S(s)),
29        ProtoValueEnum::NumberValue(n) => Ok(KsValue::N(n)),
30        ProtoValueEnum::BinaryValue(b) => Ok(KsValue::B(Bytes::from(b))),
31        ProtoValueEnum::BoolValue(b) => Ok(KsValue::Bool(b)),
32        ProtoValueEnum::NullValue(_) => Ok(KsValue::Null),
33        ProtoValueEnum::ListValue(list) => {
34            let items: Result<Vec<KsValue>, Status> =
35                list.items.into_iter().map(proto_value_to_ks).collect();
36            Ok(KsValue::L(items?))
37        }
38        ProtoValueEnum::MapValue(map) => {
39            let mut kv_map = HashMap::new();
40            for (k, v) in map.fields {
41                kv_map.insert(k, proto_value_to_ks(v)?);
42            }
43            Ok(KsValue::M(kv_map))
44        }
45        ProtoValueEnum::VectorValue(vec) => Ok(KsValue::VecF32(vec.values)),
46        ProtoValueEnum::TimestampValue(ts) => Ok(KsValue::Ts(ts as i64)),
47    }
48}
49
50/// Convert KeystoneDB Value to protobuf Value
51pub fn ks_value_to_proto(value: &KsValue) -> proto::Value {
52    let value_enum = match value {
53        KsValue::S(s) => ProtoValueEnum::StringValue(s.clone()),
54        KsValue::N(n) => ProtoValueEnum::NumberValue(n.clone()),
55        KsValue::B(b) => ProtoValueEnum::BinaryValue(b.to_vec()),
56        KsValue::Bool(b) => ProtoValueEnum::BoolValue(*b),
57        KsValue::Null => ProtoValueEnum::NullValue(proto::NullValue::NullValue as i32),
58        KsValue::L(items) => {
59            let proto_items: Vec<proto::Value> = items.iter().map(ks_value_to_proto).collect();
60            ProtoValueEnum::ListValue(proto::ListValue { items: proto_items })
61        }
62        KsValue::M(map) => {
63            let mut proto_map = HashMap::new();
64            for (k, v) in map {
65                proto_map.insert(k.clone(), ks_value_to_proto(v));
66            }
67            ProtoValueEnum::MapValue(proto::MapValue { fields: proto_map })
68        }
69        KsValue::VecF32(vec) => ProtoValueEnum::VectorValue(proto::VectorValue {
70            values: vec.clone(),
71        }),
72        KsValue::Ts(ts) => ProtoValueEnum::TimestampValue(*ts as u64),
73    };
74
75    proto::Value {
76        value: Some(value_enum),
77    }
78}
79
80// ============================================================================
81// Item Conversions
82// ============================================================================
83
84/// Convert protobuf Item to KeystoneDB Item (HashMap)
85pub fn proto_item_to_ks(item: proto::Item) -> Result<Item, Status> {
86    let mut kv_map = HashMap::new();
87    for (k, v) in item.attributes {
88        kv_map.insert(k, proto_value_to_ks(v)?);
89    }
90    Ok(kv_map)
91}
92
93/// Convert KeystoneDB Item to protobuf Item
94pub fn ks_item_to_proto(item: &Item) -> proto::Item {
95    let mut attributes = HashMap::new();
96    for (k, v) in item {
97        attributes.insert(k.clone(), ks_value_to_proto(v));
98    }
99    proto::Item { attributes }
100}
101
102// ============================================================================
103// Key Conversions
104// ============================================================================
105
106/// Convert protobuf Key to (partition_key, optional sort_key)
107pub fn proto_key_to_ks(key: proto::Key) -> (Bytes, Option<Bytes>) {
108    let pk = Bytes::from(key.partition_key);
109    let sk = key.sort_key.map(Bytes::from);
110    (pk, sk)
111}
112
113/// Convert (partition_key, optional sort_key) to protobuf Key
114pub fn ks_key_to_proto(pk: impl Into<Vec<u8>>, sk: Option<impl Into<Vec<u8>>>) -> proto::Key {
115    proto::Key {
116        partition_key: pk.into(),
117        sort_key: sk.map(|s| s.into()),
118    }
119}
120
121/// Convert protobuf Key to kstone_core Key
122pub fn proto_key_to_core_key(key: proto::Key) -> Key {
123    if let Some(sk) = key.sort_key {
124        Key::with_sk(Bytes::from(key.partition_key), Bytes::from(sk))
125    } else {
126        Key::new(Bytes::from(key.partition_key))
127    }
128}
129
130/// Convert kstone_core Key to protobuf Key
131pub fn core_key_to_proto(key: &Key) -> proto::Key {
132    proto::Key {
133        partition_key: key.pk.to_vec(),
134        sort_key: key.sk.as_ref().map(|s| s.to_vec()),
135    }
136}
137
138// ============================================================================
139// LastKey Conversions
140// ============================================================================
141
142/// Convert protobuf LastKey to (partition_key, optional sort_key)
143pub fn proto_last_key_to_ks(key: proto::LastKey) -> (Bytes, Option<Bytes>) {
144    let pk = Bytes::from(key.partition_key);
145    let sk = key.sort_key.map(Bytes::from);
146    (pk, sk)
147}
148
149/// Convert (partition_key, optional sort_key) to protobuf LastKey
150pub fn ks_last_key_to_proto(
151    pk: impl Into<Vec<u8>>,
152    sk: Option<impl Into<Vec<u8>>>,
153) -> proto::LastKey {
154    proto::LastKey {
155        partition_key: pk.into(),
156        sort_key: sk.map(|s| s.into()),
157    }
158}
159
160// ============================================================================
161// Helper Functions for Option<LastKey>
162// ============================================================================
163
164/// Convert Option<(Bytes, Option<Bytes>)> to Option<proto::LastKey>
165pub fn ks_last_key_opt_to_proto(last_key: Option<(Bytes, Option<Bytes>)>) -> Option<proto::LastKey> {
166    last_key.map(|(pk, sk)| ks_last_key_to_proto(pk, sk))
167}
168
169// ============================================================================
170// Tests
171// ============================================================================
172
173#[cfg(test)]
174mod tests {
175    use super::*;
176
177    #[test]
178    fn test_value_string_roundtrip() {
179        let ks_value = KsValue::S("hello".to_string());
180        let proto_value = ks_value_to_proto(&ks_value);
181        let converted = proto_value_to_ks(proto_value).unwrap();
182        assert_eq!(ks_value, converted);
183    }
184
185    #[test]
186    fn test_value_number_roundtrip() {
187        let ks_value = KsValue::N("42.5".to_string());
188        let proto_value = ks_value_to_proto(&ks_value);
189        let converted = proto_value_to_ks(proto_value).unwrap();
190        assert_eq!(ks_value, converted);
191    }
192
193    #[test]
194    fn test_value_bool_roundtrip() {
195        let ks_value = KsValue::Bool(true);
196        let proto_value = ks_value_to_proto(&ks_value);
197        let converted = proto_value_to_ks(proto_value).unwrap();
198        assert_eq!(ks_value, converted);
199    }
200
201    #[test]
202    fn test_value_null_roundtrip() {
203        let ks_value = KsValue::Null;
204        let proto_value = ks_value_to_proto(&ks_value);
205        let converted = proto_value_to_ks(proto_value).unwrap();
206        assert_eq!(ks_value, converted);
207    }
208
209    #[test]
210    fn test_value_binary_roundtrip() {
211        let ks_value = KsValue::B(Bytes::from(vec![1, 2, 3, 4]));
212        let proto_value = ks_value_to_proto(&ks_value);
213        let converted = proto_value_to_ks(proto_value).unwrap();
214        assert_eq!(ks_value, converted);
215    }
216
217    #[test]
218    fn test_value_list_roundtrip() {
219        let ks_value = KsValue::L(vec![
220            KsValue::S("a".to_string()),
221            KsValue::N("123".to_string()),
222            KsValue::Bool(false),
223        ]);
224        let proto_value = ks_value_to_proto(&ks_value);
225        let converted = proto_value_to_ks(proto_value).unwrap();
226        assert_eq!(ks_value, converted);
227    }
228
229    #[test]
230    fn test_value_map_roundtrip() {
231        let mut map = HashMap::new();
232        map.insert("name".to_string(), KsValue::S("Alice".to_string()));
233        map.insert("age".to_string(), KsValue::N("30".to_string()));
234
235        let ks_value = KsValue::M(map);
236        let proto_value = ks_value_to_proto(&ks_value);
237        let converted = proto_value_to_ks(proto_value).unwrap();
238        assert_eq!(ks_value, converted);
239    }
240
241    #[test]
242    fn test_value_vector_roundtrip() {
243        let ks_value = KsValue::VecF32(vec![1.0, 2.0, 3.0]);
244        let proto_value = ks_value_to_proto(&ks_value);
245        let converted = proto_value_to_ks(proto_value).unwrap();
246        assert_eq!(ks_value, converted);
247    }
248
249    #[test]
250    fn test_value_timestamp_roundtrip() {
251        let ks_value = KsValue::Ts(1234567890);
252        let proto_value = ks_value_to_proto(&ks_value);
253        let converted = proto_value_to_ks(proto_value).unwrap();
254        assert_eq!(ks_value, converted);
255    }
256
257    #[test]
258    fn test_item_roundtrip() {
259        let mut item = HashMap::new();
260        item.insert("name".to_string(), KsValue::S("Alice".to_string()));
261        item.insert("age".to_string(), KsValue::N("30".to_string()));
262        item.insert("active".to_string(), KsValue::Bool(true));
263
264        let proto_item = ks_item_to_proto(&item);
265        let converted = proto_item_to_ks(proto_item).unwrap();
266        assert_eq!(item, converted);
267    }
268
269    #[test]
270    fn test_key_with_sort_key() {
271        let proto_key = ks_key_to_proto(b"pk123".to_vec(), Some(b"sk456".to_vec()));
272        let (pk, sk) = proto_key_to_ks(proto_key);
273        assert_eq!(pk, Bytes::from("pk123"));
274        assert_eq!(sk, Some(Bytes::from("sk456")));
275    }
276
277    #[test]
278    fn test_key_without_sort_key() {
279        let proto_key = ks_key_to_proto(b"pk123".to_vec(), None::<Vec<u8>>);
280        let (pk, sk) = proto_key_to_ks(proto_key);
281        assert_eq!(pk, Bytes::from("pk123"));
282        assert_eq!(sk, None);
283    }
284
285    #[test]
286    fn test_empty_item() {
287        let item = HashMap::new();
288        let proto_item = ks_item_to_proto(&item);
289        let converted = proto_item_to_ks(proto_item).unwrap();
290        assert_eq!(item, converted);
291    }
292
293    #[test]
294    fn test_nested_map() {
295        let mut inner_map = HashMap::new();
296        inner_map.insert("city".to_string(), KsValue::S("NYC".to_string()));
297
298        let mut item = HashMap::new();
299        item.insert("name".to_string(), KsValue::S("Alice".to_string()));
300        item.insert("address".to_string(), KsValue::M(inner_map));
301
302        let proto_item = ks_item_to_proto(&item);
303        let converted = proto_item_to_ks(proto_item).unwrap();
304        assert_eq!(item, converted);
305    }
306}