aws_athena_parser/
lib.rs

1pub use anyhow;
2use aws_sdk_athena::types::ResultSet;
3pub use from_athena_derive::FromAthena;
4pub use std::collections::HashMap;
5
6/// A trait for converting data from an Athena query result into a specified type.
7///
8/// This trait defines a method `from_athena` which converts a HashMap of string
9/// key-value pairs representing data retrieved from an Athena query into an instance
10/// of the implementing type. The conversion process may involve parsing, validation,
11/// or any other necessary transformation.
12///
13/// # Errors
14///
15/// If the conversion process fails due to invalid or missing data, an error is returned.
16///
17/// # Examples
18///
19/// Implementing `FromAthena` for a custom struct:
20pub trait FromAthena: Sized {
21    /// Converts a HashMap of string key-value pairs into an instance of the implementing type.
22    ///
23    /// # Arguments
24    ///
25    /// * `values` - A HashMap containing the data to be converted.
26    ///
27    /// # Returns
28    ///
29    /// Result containing the converted instance of the implementing type or an error if conversion fails.
30    fn from_athena(values: HashMap<String, String>) -> anyhow::Result<Self, anyhow::Error>;
31}
32
33/// Builds a vector of hash maps representing the rows of the given ResultSet.
34///
35/// This function takes a ResultSet as input and returns a vector of hash maps,
36/// where each hash map represents a row in the ResultSet. If the ResultSet contains
37/// no data or metadata, an empty vector is returned.
38///
39/// # Arguments
40///
41/// * `result_set` - A ResultSet containing the data to be converted into hash maps.
42///
43/// # Returns
44///
45/// A vector of hash maps, where each hash map represents a row in the ResultSet.
46///
47/// # Examples
48///
49/// ```
50/// use aws_sdk_athena::types::ResultSet;
51/// use aws_sdk_athena::build_map;
52///
53/// let result_set = ResultSet::new(/* Some initialization */);
54/// let mapped_data = build_map(result_set);
55/// // Use mapped_data for further processing
56/// ```
57pub 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}