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

datafusion_table_providers/util/
mod.rs
