datafusion-table-providers 0.11.1

Extend the capabilities of DataFusion to support additional data sources via implementations of the `TableProvider` trait.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108

use std::ops::ControlFlow;

use datafusion::sql::{
    sqlparser::ast::{
        FunctionArg, Ident, ObjectName, TableAlias, TableFactor, TableFunctionArgs, VisitorMut,
    },
    TableReference,
};

#[derive(Debug, Clone, PartialEq, Eq, Default)]
pub struct TableArgReplace {
    pub tables: Vec<(TableReference, TableFunctionArgs)>,
}

impl TableArgReplace {
    /// Constructs a new `TableArgReplace` instance.
    pub fn new(tables: Vec<(TableReference, Vec<FunctionArg>)>) -> Self {
        Self {
            tables: tables
                .into_iter()
                .map(|(table, args)| {
                    (
                        table,
                        TableFunctionArgs {
                            args,
                            settings: None,
                        },
                    )
                })
                .collect(),
        }
    }

    /// Adds a new table argument replacement.
    pub fn with(mut self, table: TableReference, args: Vec<FunctionArg>) -> Self {
        self.tables.push((
            table,
            TableFunctionArgs {
                args,
                settings: None,
            },
        ));
        self
    }

    #[cfg(feature = "federation")]
    /// Converts the `TableArgReplace` instance into an `AstAnalyzer`.
    pub fn into_analyzer(self) -> datafusion_federation::sql::AstAnalyzer {
        let mut visitor = self;
        let x = move |mut statement: datafusion::sql::sqlparser::ast::Statement| {
            let _ = datafusion::sql::sqlparser::ast::VisitMut::visit(&mut statement, &mut visitor);
            Ok(statement)
        };
        Box::new(x)
    }
}

impl VisitorMut for TableArgReplace {
    type Break = ();
    fn pre_visit_table_factor(
        &mut self,
        table_factor: &mut TableFactor,
    ) -> ControlFlow<Self::Break> {
        if let TableFactor::Table {
            name, args, alias, ..
        } = table_factor
        {
            let name_as_tableref = name_to_table_reference(name);
            if let Some((table, arg)) = self
                .tables
                .iter()
                .find(|(t, _)| t.resolved_eq(&name_as_tableref))
            {
                *args = Some(arg.clone());
                if alias.is_none() {
                    *alias = Some(TableAlias {
                        explicit: true,
                        name: Ident::new(table.table()),
                        columns: vec![],
                    })
                }
            }
        }
        ControlFlow::Continue(())
    }
}

fn name_to_table_reference(name: &ObjectName) -> TableReference {
    let first = name
        .0
        .first()
        .map(|n| n.as_ident().expect("expected Ident").value.to_string());
    let second = name
        .0
        .get(1)
        .map(|n| n.as_ident().expect("expected Ident").value.to_string());
    let third = name
        .0
        .get(2)
        .map(|n| n.as_ident().expect("expected Ident").value.to_string());

    match (first, second, third) {
        (Some(first), Some(second), Some(third)) => TableReference::full(first, second, third),
        (Some(first), Some(second), None) => TableReference::partial(first, second),
        (Some(first), None, None) => TableReference::bare(first),
        _ => panic!("Invalid table name"),
    }
}