1use 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
17pub 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
50pub 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
80pub 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
93pub 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
102pub 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
113pub 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
121pub 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
130pub 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
138pub 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
149pub 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
160pub 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#[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}