ns_scylla_orm/
columns.rs

1use scylla::frame::response::result::Row;
2use std::collections::{hash_map::Iter, HashMap};
3
4use crate::{CqlValue, FromCqlVal, ToCqlVal};
5
6#[derive(Debug, Default, PartialEq)]
7pub struct ColumnsMap(HashMap<String, CqlValue>);
8
9impl ColumnsMap {
10    pub fn new() -> Self {
11        Self(HashMap::new())
12    }
13
14    pub fn with_capacity(capacity: usize) -> Self {
15        Self(HashMap::with_capacity(capacity))
16    }
17
18    pub fn len(&self) -> usize {
19        self.0.len()
20    }
21
22    pub fn is_empty(&self) -> bool {
23        self.0.is_empty()
24    }
25
26    pub fn has(&self, key: &str) -> bool {
27        self.0.contains_key(key)
28    }
29
30    pub fn keys(&self) -> Vec<String> {
31        self.0.keys().cloned().collect()
32    }
33
34    pub fn get(&self, key: &str) -> Option<&CqlValue> {
35        match self.0.get(key) {
36            Some(v) => Some(v),
37            None => None,
38        }
39    }
40
41    pub fn iter(&self) -> Iter<'_, String, CqlValue> {
42        self.0.iter()
43    }
44
45    pub fn get_as<T: FromCqlVal<Option<CqlValue>>>(&self, key: &str) -> anyhow::Result<T> {
46        match self.0.get(key) {
47            Some(v) => T::from_cql(Some(v.clone())).map_err(anyhow::Error::new),
48            None => T::from_cql(None).map_err(anyhow::Error::new),
49        }
50    }
51
52    pub fn set_as<T: ToCqlVal>(&mut self, key: &str, val: &T) {
53        self.0.insert(key.to_string(), val.to_cql());
54    }
55
56    pub fn append_map<T: ToCqlVal>(&mut self, map_name: &str, key: &str, val: T) {
57        let mut map: HashMap<String, CqlValue> = self.get_as(map_name).unwrap_or_default();
58
59        map.insert(key.to_string(), val.to_cql());
60        self.0.insert(map_name.to_string(), map.to_cql());
61    }
62
63    pub fn fill(&mut self, row: Row, fields: &Vec<String>) -> anyhow::Result<()> {
64        if row.columns.len() != fields.len() {
65            return Err(anyhow::Error::msg(format!(
66                "ColumnsMap::fill: row.columns.len({}) != fields.len({})",
67                row.columns.len(),
68                fields.len()
69            )));
70        }
71        for (i, val) in row.columns.iter().enumerate() {
72            if let Some(v) = val {
73                self.0.insert(fields[i].to_owned(), v.to_owned());
74            }
75        }
76        Ok(())
77    }
78}
79
80#[cfg(test)]
81mod tests {
82    use super::*;
83    use serde::{de::DeserializeOwned, Serialize};
84
85    impl ColumnsMap {
86        pub fn get_from_cbor<T: DeserializeOwned>(&self, key: &str) -> anyhow::Result<T> {
87            let data = self.get_as::<Vec<u8>>(key)?;
88            let val: T = ciborium::from_reader(&data[..])?;
89            Ok(val)
90        }
91
92        pub fn set_in_cbor<T: ?Sized + Serialize>(
93            &mut self,
94            key: &str,
95            val: &T,
96        ) -> anyhow::Result<()> {
97            let mut buf: Vec<u8> = Vec::new();
98            ciborium::into_writer(val, &mut buf)?;
99            self.0.insert(key.to_string(), CqlValue::Blob(buf));
100            Ok(())
101        }
102    }
103
104    #[test]
105    fn columns_map_works() {
106        let mut map = ColumnsMap::new();
107
108        assert_eq!(map.len(), 0);
109        assert!(!map.has("name"));
110        assert_eq!(map.get("name"), None);
111        assert!(map.get_as::<String>("name").is_err());
112
113        map.set_as("name", &"jarvis".to_string());
114        assert_eq!(map.len(), 1);
115        assert!(map.has("name"));
116        assert_eq!(map.get("name"), Some(&CqlValue::Text("jarvis".to_string())));
117        assert_eq!(map.get_as::<String>("name").unwrap(), "jarvis".to_string());
118
119        map.set_as("name", &"jarvis2".to_string());
120        assert_eq!(map.len(), 1);
121        assert!(map.has("name"));
122        assert_eq!(
123            map.get("name"),
124            Some(&CqlValue::Text("jarvis2".to_string()))
125        );
126        assert_eq!(map.get_as::<String>("name").unwrap(), "jarvis2".to_string());
127
128        assert!(!map.has("data"));
129        assert_eq!(map.get("data"), None);
130        assert!(map.get_as::<Vec<u8>>("data").is_err());
131        assert_eq!(map.get_as::<Option<Vec<u8>>>("data").unwrap(), None);
132        assert!(map.set_in_cbor("data", &vec![1i64, 2i64, 3i64]).is_ok()); // CBOR: 0x83010203
133        assert!(map.has("data"));
134        assert_eq!(map.len(), 2);
135        assert_eq!(
136            map.get_as::<Vec<u8>>("data").unwrap(),
137            vec![0x83, 0x01, 0x02, 0x03],
138        );
139        assert_eq!(
140            map.get_as::<Option<Vec<u8>>>("data").unwrap(),
141            Some(vec![0x83, 0x01, 0x02, 0x03]),
142        );
143        assert!(map.get_as::<String>("data").is_err());
144
145        let mut keys: Option<Vec<Vec<u8>>> = None;
146        assert!(!map.has("data2"));
147        assert_eq!(map.get("data2"), None);
148        assert!(map.get_as::<Vec<Vec<u8>>>("data2").is_err());
149        assert_eq!(map.get_as::<Option<Vec<Vec<u8>>>>("data2").unwrap(), None);
150        map.set_as("data2", &keys);
151        assert!(map.has("data2"));
152        assert_eq!(map.len(), 3);
153        assert!(map.get_as::<Vec<Vec<u8>>>("data2").is_err());
154        assert_eq!(map.get_as::<Option<Vec<Vec<u8>>>>("data2").unwrap(), None);
155
156        keys = Some(vec![vec![0x83, 0x01, 0x02, 0x03]]);
157        map.set_as("data2", &keys);
158        assert_eq!(map.get_as::<Option<Vec<Vec<u8>>>>("data2").unwrap(), keys);
159        assert_eq!(map.get_as::<Vec<Vec<u8>>>("data2").unwrap(), keys.unwrap());
160
161        let mut row: Row = Row {
162            columns: Vec::new(),
163        };
164
165        let mut fields: Vec<String> = Vec::new();
166        for (k, v) in map.iter() {
167            fields.push(k.to_owned());
168            row.columns.push(Some(v.to_owned()));
169        }
170
171        assert_eq!(fields.len(), 3);
172        let mut map2 = ColumnsMap::new();
173        assert!(map2
174            .fill(
175                Row {
176                    columns: Vec::new(),
177                },
178                &fields
179            )
180            .is_err());
181        assert_ne!(map2, map);
182
183        assert!(map2.fill(row, &fields).is_ok());
184        assert_eq!(map2, map);
185    }
186}