1pub use anyhow;
2use aws_sdk_athena::types::ResultSet;
3pub use from_athena_derive::FromAthena;
4pub use std::collections::HashMap;
5
6pub trait FromAthena: Sized {
21 fn from_athena(values: HashMap<String, String>) -> anyhow::Result<Self, anyhow::Error>;
31}
32
33pub fn build_map(result_set: ResultSet) -> Vec<HashMap<String, String>> {
58 if let Some(meta) = result_set.result_set_metadata() {
59 let columns: Vec<String> = meta
60 .column_info()
61 .iter()
62 .map(|c| c.name().to_string())
63 .collect();
64
65 let rows: Vec<HashMap<String, String>> = result_set
66 .rows()
67 .iter()
68 .map(|r| {
69 r.data()
70 .iter()
71 .map(|d| d.var_char_value().unwrap_or("").to_string())
72 .zip(columns.iter())
73 .map(|(val, col)| (col.clone(), val.clone()))
74 .collect::<HashMap<String, String>>()
75 })
76 .collect();
77
78 rows
79 } else {
80 vec![]
81 }
82}
83
84#[cfg(test)]
85mod test {
86 use super::*;
87 use aws_sdk_athena::types::{ColumnInfo, Datum, ResultSetMetadata, Row};
88
89 #[derive(from_athena_derive::FromAthena)]
90 struct Testing {
91 pub test: i64,
92 }
93
94 #[derive(from_athena_derive::FromAthena)]
95 #[allow(dead_code)]
96 struct BadTesting {
97 pub no_exist: String,
98 }
99
100 #[derive(from_athena_derive::FromAthena)]
101 struct LargeStruct {
102 pub test1: i64,
103 pub test2: i32,
104 pub test3: String,
105 pub test4: String,
106 pub test5: f64,
107 pub test6: bool,
108 }
109
110 #[test]
111 fn convert_result_set_to_map() {
112 let column = ColumnInfo::builder()
113 .name("test")
114 .r#type("bigint")
115 .build()
116 .unwrap();
117 let metadata = ResultSetMetadata::builder().column_info(column).build();
118 let data = Datum::builder()
119 .set_var_char_value(Some("100".to_string()))
120 .build();
121 let row = Row::builder().set_data(Some(vec![data])).build();
122 let result_set = ResultSet::builder()
123 .result_set_metadata(metadata)
124 .set_rows(Some(vec![row]))
125 .build();
126
127 let res = build_map(result_set);
128 assert!(res.len() == 1);
129 assert!(res[0].get("test").is_some());
130 assert_eq!(res[0].get("test").unwrap(), "100");
131 }
132
133 #[test]
134 fn converted_results_to_struct() {
135 let column = ColumnInfo::builder()
136 .name("test")
137 .r#type("bigint")
138 .build()
139 .unwrap();
140 let metadata = ResultSetMetadata::builder().column_info(column).build();
141 let data = Datum::builder()
142 .set_var_char_value(Some("100".to_string()))
143 .build();
144 let row = Row::builder().set_data(Some(vec![data])).build();
145 let result_set = ResultSet::builder()
146 .result_set_metadata(metadata)
147 .set_rows(Some(vec![row]))
148 .build();
149
150 let res: Vec<Testing> = build_map(result_set)
151 .iter()
152 .flat_map(|x| Testing::from_athena(x.clone()))
153 .collect();
154
155 assert_eq!(res[0].test, 100);
156 }
157
158 #[test]
159 fn converted_results_to_large_struct() {
160 let columns = [
161 ("test1", "bigint"),
162 ("test2", "integer"),
163 ("test3", "varchar"),
164 ("test4", "varchar"),
165 ("test5", "double"),
166 ("test6", "boolean"),
167 ]
168 .iter()
169 .map(|i| {
170 ColumnInfo::builder()
171 .name(i.0.to_string())
172 .r#type(i.1.to_string())
173 .build()
174 .unwrap()
175 })
176 .collect();
177
178 let metadata = ResultSetMetadata::builder()
179 .set_column_info(Some(columns))
180 .build();
181
182 let data: Vec<Datum> = ["1000", "100", "test", "test", "100.0", "true"]
183 .iter()
184 .map(|v| {
185 Datum::builder()
186 .set_var_char_value(Some(v.to_string()))
187 .build()
188 })
189 .collect();
190
191 let row = Row::builder().set_data(Some(data)).build();
192
193 let result_set = ResultSet::builder()
194 .result_set_metadata(metadata)
195 .set_rows(Some(vec![row]))
196 .build();
197
198 let res: Vec<Result<LargeStruct, anyhow::Error>> = build_map(result_set)
199 .iter()
200 .map(|x| LargeStruct::from_athena(x.clone()))
201 .collect();
202
203 for res in res {
204 match res {
205 Ok(r) => {
206 assert_eq!(r.test1, 1000);
207 assert_eq!(r.test2, 100);
208 assert_eq!(r.test3, "test".to_string());
209 assert_eq!(r.test4, "test");
210 assert_eq!(r.test5, 100.0);
211 assert!(r.test6);
212 }
213 Err(e) => {
214 panic!("{:#}", e);
215 }
216 }
217 }
218 }
219
220 #[test]
221 fn error_convert_results_to_invalid_struct() {
222 let column = ColumnInfo::builder()
223 .name("test")
224 .r#type("bigint")
225 .build()
226 .unwrap();
227 let metadata = ResultSetMetadata::builder().column_info(column).build();
228 let data = Datum::builder()
229 .set_var_char_value(Some("100".to_string()))
230 .build();
231 let row = Row::builder().set_data(Some(vec![data])).build();
232 let result_set = ResultSet::builder()
233 .result_set_metadata(metadata)
234 .set_rows(Some(vec![row]))
235 .build();
236
237 let res: Vec<Result<BadTesting, anyhow::Error>> = build_map(result_set)
238 .iter()
239 .map(|x| BadTesting::from_athena(x.clone()))
240 .collect();
241
242 assert!(res[0].is_err());
243 assert_eq!(
244 res[0].as_ref().err().unwrap().to_string(),
245 "Missing field within result set. `no_exist` was not found!".to_string()
246 );
247 }
248}