datafusion_table_providers/util/
mod.rs

1use datafusion::logical_expr::Expr;
2use snafu::prelude::*;
3use std::collections::HashMap;
4use std::hash::Hash;
5use std::sync::Arc;
6
7use crate::sql::sql_provider_datafusion::Engine;
8use datafusion::common::DataFusionError;
9use datafusion::{
10    error::Result as DataFusionResult,
11    sql::unparser::{dialect::DefaultDialect, Unparser},
12};
13
14pub mod column_reference;
15pub mod constraints;
16pub mod indexes;
17pub mod ns_lookup;
18pub mod on_conflict;
19pub mod retriable_error;
20pub mod secrets;
21pub mod test;
22
23#[derive(Debug, Snafu)]
24pub enum Error {
25    #[snafu(display("Unable to generate SQL: {source}"))]
26    UnableToGenerateSQL {
27        source: datafusion::error::DataFusionError,
28    },
29}
30
31pub fn filters_to_sql(filters: &[Expr], engine: Option<Engine>) -> Result<String, Error> {
32    let dialect = engine
33        .map(|e| e.dialect())
34        .unwrap_or(Arc::new(DefaultDialect {}));
35    Ok(filters
36        .iter()
37        .map(|f| {
38            Unparser::new(dialect.as_ref())
39                .expr_to_sql(f)
40                .map(|e| e.to_string())
41        })
42        .collect::<DataFusionResult<Vec<String>>>()
43        .context(UnableToGenerateSQLSnafu)?
44        .join(" AND "))
45}
46
47#[must_use]
48pub fn hashmap_from_option_string<K, V>(hashmap_option_str: &str) -> HashMap<K, V>
49where
50    K: for<'a> From<&'a str> + Eq + Hash,
51    V: for<'a> From<&'a str> + Default,
52{
53    hashmap_option_str
54        .split(';')
55        .map(|index| {
56            let parts: Vec<&str> = index.split(':').collect();
57            if parts.len() == 2 {
58                (K::from(parts[0]), V::from(parts[1]))
59            } else {
60                (K::from(index), V::default())
61            }
62        })
63        .collect()
64}
65
66#[must_use]
67pub fn remove_prefix_from_hashmap_keys<V>(
68    hashmap: HashMap<String, V>,
69    prefix: &str,
70) -> HashMap<String, V> {
71    hashmap
72        .into_iter()
73        .map(|(key, value)| {
74            let new_key = key
75                .strip_prefix(prefix)
76                .map_or(key.clone(), |s| s.to_string());
77            (new_key, value)
78        })
79        .collect()
80}
81
82#[must_use]
83pub fn to_datafusion_error<E>(error: E) -> DataFusionError
84where
85    E: std::error::Error + Send + Sync + 'static,
86{
87    DataFusionError::External(Box::new(error))
88}
89
90#[cfg(test)]
91mod tests {
92    use super::*;
93    use std::collections::HashMap;
94
95    #[test]
96    fn test_remove_prefix() {
97        let mut hashmap = HashMap::new();
98        hashmap.insert("prefix_key1".to_string(), "value1".to_string());
99        hashmap.insert("prefix_key2".to_string(), "value2".to_string());
100        hashmap.insert("key3".to_string(), "value3".to_string());
101
102        let result = remove_prefix_from_hashmap_keys(hashmap, "prefix_");
103
104        let mut expected = HashMap::new();
105        expected.insert("key1".to_string(), "value1".to_string());
106        expected.insert("key2".to_string(), "value2".to_string());
107        expected.insert("key3".to_string(), "value3".to_string());
108
109        assert_eq!(result, expected);
110    }
111
112    #[test]
113    fn test_no_prefix() {
114        let mut hashmap = HashMap::new();
115        hashmap.insert("key1".to_string(), "value1".to_string());
116        hashmap.insert("key2".to_string(), "value2".to_string());
117
118        let result = remove_prefix_from_hashmap_keys(hashmap, "prefix_");
119
120        let mut expected = HashMap::new();
121        expected.insert("key1".to_string(), "value1".to_string());
122        expected.insert("key2".to_string(), "value2".to_string());
123
124        assert_eq!(result, expected);
125    }
126
127    #[test]
128    fn test_empty_hashmap() {
129        let hashmap: HashMap<String, String> = HashMap::new();
130
131        let result = remove_prefix_from_hashmap_keys(hashmap, "prefix_");
132
133        let expected: HashMap<String, String> = HashMap::new();
134
135        assert_eq!(result, expected);
136    }
137
138    #[test]
139    fn test_full_prefix() {
140        let mut hashmap = HashMap::new();
141        hashmap.insert("prefix_".to_string(), "value1".to_string());
142        hashmap.insert("prefix_key2".to_string(), "value2".to_string());
143
144        let result = remove_prefix_from_hashmap_keys(hashmap, "prefix_");
145
146        let mut expected = HashMap::new();
147        expected.insert("".to_string(), "value1".to_string());
148        expected.insert("key2".to_string(), "value2".to_string());
149
150        assert_eq!(result, expected);
151    }
152}
datafusion_table_providers/util/mod.rs

datafusion_table_providers/util/
mod.rs
