datafusion_table_providers/util/
mod.rs

1use snafu::prelude::*;
2use std::hash::Hash;
3
4use datafusion::common::DataFusionError;
5use std::collections::HashMap;
6
7use crate::UnsupportedTypeAction;
8
9pub mod column_reference;
10pub mod constraints;
11pub mod indexes;
12pub mod ns_lookup;
13pub mod on_conflict;
14pub mod retriable_error;
15
16#[cfg(any(feature = "sqlite", feature = "duckdb", feature = "postgres"))]
17pub mod schema;
18pub mod secrets;
19pub mod test;
20
21#[derive(Debug, Snafu)]
22pub enum Error {
23    #[snafu(display("Unable to generate SQL: {source}"))]
24    UnableToGenerateSQL {
25        source: datafusion::error::DataFusionError,
26    },
27}
28
29#[must_use]
30pub fn hashmap_from_option_string<K, V>(hashmap_option_str: &str) -> HashMap<K, V>
31where
32    K: for<'a> From<&'a str> + Eq + Hash,
33    V: for<'a> From<&'a str> + Default,
34{
35    hashmap_option_str
36        .split(';')
37        .map(|index| {
38            let parts: Vec<&str> = index.split(':').collect();
39            if parts.len() == 2 {
40                (K::from(parts[0]), V::from(parts[1]))
41            } else {
42                (K::from(index), V::default())
43            }
44        })
45        .collect()
46}
47
48#[must_use]
49pub fn remove_prefix_from_hashmap_keys<V>(
50    hashmap: HashMap<String, V>,
51    prefix: &str,
52) -> HashMap<String, V> {
53    hashmap
54        .into_iter()
55        .map(|(key, value)| {
56            let new_key = key
57                .strip_prefix(prefix)
58                .map_or(key.clone(), |s| s.to_string());
59            (new_key, value)
60        })
61        .collect()
62}
63
64#[must_use]
65pub fn to_datafusion_error<E>(error: E) -> DataFusionError
66where
67    E: std::error::Error + Send + Sync + 'static,
68{
69    DataFusionError::External(Box::new(error))
70}
71
72/// If the `UnsupportedTypeAction` is `Error` or `String`, the function will return an error.
73/// If the `UnsupportedTypeAction` is `Warn`, the function will log a warning.
74/// If the `UnsupportedTypeAction` is `Ignore`, the function will do nothing.
75///
76/// Used in the handling of unsupported schemas to determine if columns are removed or if the function should return an error.
77///
78/// Components that want to handle `UnsupportedTypeAction::String` via the component's own logic should handle it before calling this function.
79///
80/// # Errors
81///
82/// If the `UnsupportedTypeAction` is `Error` or `String` the function will return an error.
83pub fn handle_unsupported_type_error<E>(
84    unsupported_type_action: UnsupportedTypeAction,
85    error: E,
86) -> Result<(), E>
87where
88    E: std::error::Error + Send + Sync,
89{
90    match unsupported_type_action {
91        UnsupportedTypeAction::Error | UnsupportedTypeAction::String => return Err(error),
92        UnsupportedTypeAction::Warn => {
93            tracing::warn!("{error}");
94        }
95        UnsupportedTypeAction::Ignore => {}
96    };
97
98    Ok(())
99}
100
101#[cfg(test)]
102mod tests {
103    use super::*;
104    use std::collections::HashMap;
105
106    #[test]
107    fn test_remove_prefix() {
108        let mut hashmap = HashMap::new();
109        hashmap.insert("prefix_key1".to_string(), "value1".to_string());
110        hashmap.insert("prefix_key2".to_string(), "value2".to_string());
111        hashmap.insert("key3".to_string(), "value3".to_string());
112
113        let result = remove_prefix_from_hashmap_keys(hashmap, "prefix_");
114
115        let mut expected = HashMap::new();
116        expected.insert("key1".to_string(), "value1".to_string());
117        expected.insert("key2".to_string(), "value2".to_string());
118        expected.insert("key3".to_string(), "value3".to_string());
119
120        assert_eq!(result, expected);
121    }
122
123    #[test]
124    fn test_no_prefix() {
125        let mut hashmap = HashMap::new();
126        hashmap.insert("key1".to_string(), "value1".to_string());
127        hashmap.insert("key2".to_string(), "value2".to_string());
128
129        let result = remove_prefix_from_hashmap_keys(hashmap, "prefix_");
130
131        let mut expected = HashMap::new();
132        expected.insert("key1".to_string(), "value1".to_string());
133        expected.insert("key2".to_string(), "value2".to_string());
134
135        assert_eq!(result, expected);
136    }
137
138    #[test]
139    fn test_empty_hashmap() {
140        let hashmap: HashMap<String, String> = HashMap::new();
141
142        let result = remove_prefix_from_hashmap_keys(hashmap, "prefix_");
143
144        let expected: HashMap<String, String> = HashMap::new();
145
146        assert_eq!(result, expected);
147    }
148
149    #[test]
150    fn test_full_prefix() {
151        let mut hashmap = HashMap::new();
152        hashmap.insert("prefix_".to_string(), "value1".to_string());
153        hashmap.insert("prefix_key2".to_string(), "value2".to_string());
154
155        let result = remove_prefix_from_hashmap_keys(hashmap, "prefix_");
156
157        let mut expected = HashMap::new();
158        expected.insert("".to_string(), "value1".to_string());
159        expected.insert("key2".to_string(), "value2".to_string());
160
161        assert_eq!(result, expected);
162    }
163}
datafusion_table_providers/util/mod.rs

datafusion_table_providers/util/
mod.rs
