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()); 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}