datafusion_federation/sql/
table_reference.rs

1use std::sync::Arc;
2
3use datafusion::{
4    common::TableReference,
5    error::DataFusionError,
6    sql::sqlparser::{
7        self,
8        ast::{FunctionArg, ObjectNamePart},
9        dialect::{Dialect, GenericDialect},
10        tokenizer::Token,
11    },
12};
13
14/// A multipart identifier to a remote table, view or parameterized view.
15///
16/// RemoteTableRef can be created by parsing from a string representing a table object with optional
17/// ```rust
18/// use datafusion_federation::sql::RemoteTableRef;
19/// use datafusion::sql::sqlparser::dialect::PostgreSqlDialect;
20///
21/// RemoteTableRef::try_from("myschema.table");
22/// RemoteTableRef::try_from(r#"myschema."Table""#);
23/// RemoteTableRef::try_from("myschema.view('obj')");
24///
25/// RemoteTableRef::parse_with_dialect("myschema.view(name = 'obj')", &PostgreSqlDialect {});
26/// ```
27#[derive(Debug, Clone, PartialEq, Eq, Hash)]
28pub struct RemoteTableRef {
29    pub table_ref: TableReference,
30    pub args: Option<Arc<[FunctionArg]>>,
31}
32
33impl RemoteTableRef {
34    /// Get quoted_string representation for the table it is referencing, this is same as calling to_quoted_string on the inner table reference.
35    pub fn to_quoted_string(&self) -> String {
36        self.table_ref.to_quoted_string()
37    }
38
39    /// Create new using general purpose dialect. Prefer [`Self::parse_with_dialect`] if the dialect is known beforehand
40    pub fn parse_with_default_dialect(s: &str) -> Result<Self, DataFusionError> {
41        Self::parse_with_dialect(s, &GenericDialect {})
42    }
43
44    /// Create new using a specific instance of dialect.
45    pub fn parse_with_dialect(s: &str, dialect: &dyn Dialect) -> Result<Self, DataFusionError> {
46        let mut parser = sqlparser::parser::Parser::new(dialect).try_with_sql(s)?;
47        let name = parser.parse_object_name(true)?;
48        let args = if parser.consume_token(&Token::LParen) {
49            parser.parse_optional_args()?
50        } else {
51            vec![]
52        };
53
54        let table_ref = match (name.0.first(), name.0.get(1), name.0.get(2)) {
55            (
56                Some(ObjectNamePart::Identifier(catalog)),
57                Some(ObjectNamePart::Identifier(schema)),
58                Some(ObjectNamePart::Identifier(table)),
59            ) => TableReference::full(
60                catalog.value.clone(),
61                schema.value.clone(),
62                table.value.clone(),
63            ),
64            (
65                Some(ObjectNamePart::Identifier(schema)),
66                Some(ObjectNamePart::Identifier(table)),
67                None,
68            ) => TableReference::partial(schema.value.clone(), table.value.clone()),
69            (Some(ObjectNamePart::Identifier(table)), None, None) => {
70                TableReference::bare(table.value.clone())
71            }
72            _ => {
73                return Err(DataFusionError::NotImplemented(
74                    "Unable to parse string into TableReference".to_string(),
75                ))
76            }
77        };
78
79        if !args.is_empty() {
80            Ok(RemoteTableRef {
81                table_ref,
82                args: Some(args.into()),
83            })
84        } else {
85            Ok(RemoteTableRef {
86                table_ref,
87                args: None,
88            })
89        }
90    }
91
92    pub fn table_ref(&self) -> &TableReference {
93        &self.table_ref
94    }
95
96    pub fn args(&self) -> Option<&[FunctionArg]> {
97        self.args.as_deref()
98    }
99}
100
101impl From<TableReference> for RemoteTableRef {
102    fn from(table_ref: TableReference) -> Self {
103        RemoteTableRef {
104            table_ref,
105            args: None,
106        }
107    }
108}
109
110impl From<RemoteTableRef> for TableReference {
111    fn from(remote_table_ref: RemoteTableRef) -> Self {
112        remote_table_ref.table_ref
113    }
114}
115
116impl From<&RemoteTableRef> for TableReference {
117    fn from(remote_table_ref: &RemoteTableRef) -> Self {
118        remote_table_ref.table_ref.clone()
119    }
120}
121
122impl From<(TableReference, Vec<FunctionArg>)> for RemoteTableRef {
123    fn from((table_ref, args): (TableReference, Vec<FunctionArg>)) -> Self {
124        RemoteTableRef {
125            table_ref,
126            args: Some(args.into()),
127        }
128    }
129}
130
131impl TryFrom<&str> for RemoteTableRef {
132    type Error = DataFusionError;
133    fn try_from(s: &str) -> Result<Self, Self::Error> {
134        Self::parse_with_default_dialect(s)
135    }
136}
137
138impl TryFrom<String> for RemoteTableRef {
139    type Error = DataFusionError;
140    fn try_from(s: String) -> Result<Self, Self::Error> {
141        Self::parse_with_default_dialect(&s)
142    }
143}
144
145impl TryFrom<&String> for RemoteTableRef {
146    type Error = DataFusionError;
147    fn try_from(s: &String) -> Result<Self, Self::Error> {
148        Self::parse_with_default_dialect(s)
149    }
150}
151
152#[cfg(test)]
153mod tests {
154    use super::*;
155    use sqlparser::{
156        ast::{self, Expr, FunctionArgOperator, Ident, Value},
157        dialect,
158    };
159
160    #[test]
161    fn bare_table_reference() {
162        let table_ref = RemoteTableRef::parse_with_default_dialect("table").unwrap();
163        let expected = RemoteTableRef::from(TableReference::bare("table"));
164        assert_eq!(table_ref, expected);
165
166        let table_ref = RemoteTableRef::parse_with_default_dialect("Table").unwrap();
167        let expected = RemoteTableRef::from(TableReference::bare("Table"));
168        assert_eq!(table_ref, expected);
169    }
170
171    #[test]
172    fn bare_table_reference_with_args() {
173        let table_ref = RemoteTableRef::parse_with_default_dialect("table(1, 2)").unwrap();
174        let expected = RemoteTableRef::from((
175            TableReference::bare("table"),
176            vec![
177                FunctionArg::Unnamed(Expr::value(Value::Number("1".to_string(), false)).into()),
178                FunctionArg::Unnamed(Expr::value(Value::Number("2".to_string(), false)).into()),
179            ],
180        ));
181        assert_eq!(table_ref, expected);
182
183        let table_ref = RemoteTableRef::parse_with_default_dialect("Table(1, 2)").unwrap();
184        let expected = RemoteTableRef::from((
185            TableReference::bare("Table"),
186            vec![
187                FunctionArg::Unnamed(Expr::value(Value::Number("1".to_string(), false)).into()),
188                FunctionArg::Unnamed(Expr::value(Value::Number("2".to_string(), false)).into()),
189            ],
190        ));
191        assert_eq!(table_ref, expected);
192    }
193
194    #[test]
195    fn bare_table_reference_with_args_and_whitespace() {
196        let table_ref = RemoteTableRef::parse_with_default_dialect("table (1, 2)").unwrap();
197        let expected = RemoteTableRef::from((
198            TableReference::bare("table"),
199            vec![
200                FunctionArg::Unnamed(Expr::value(Value::Number("1".to_string(), false)).into()),
201                FunctionArg::Unnamed(Expr::value(Value::Number("2".to_string(), false)).into()),
202            ],
203        ));
204        assert_eq!(table_ref, expected);
205
206        let table_ref = RemoteTableRef::parse_with_default_dialect("Table (1, 2)").unwrap();
207        let expected = RemoteTableRef::from((
208            TableReference::bare("Table"),
209            vec![
210                FunctionArg::Unnamed(Expr::value(Value::Number("1".to_string(), false)).into()),
211                FunctionArg::Unnamed(Expr::value(Value::Number("2".to_string(), false)).into()),
212            ],
213        ));
214        assert_eq!(table_ref, expected);
215    }
216
217    #[test]
218    fn multi_table_reference_with_no_args() {
219        let table_ref = RemoteTableRef::parse_with_default_dialect("schema.table").unwrap();
220        let expected = RemoteTableRef::from(TableReference::partial("schema", "table"));
221        assert_eq!(table_ref, expected);
222
223        let table_ref = RemoteTableRef::parse_with_default_dialect("schema.Table").unwrap();
224        let expected = RemoteTableRef::from(TableReference::partial("schema", "Table"));
225        assert_eq!(table_ref, expected);
226    }
227
228    #[test]
229    fn multi_table_reference_with_args() {
230        let table_ref = RemoteTableRef::parse_with_default_dialect("schema.table(1, 2)").unwrap();
231        let expected = RemoteTableRef::from((
232            TableReference::partial("schema", "table"),
233            vec![
234                FunctionArg::Unnamed(Expr::value(Value::Number("1".to_string(), false)).into()),
235                FunctionArg::Unnamed(Expr::value(Value::Number("2".to_string(), false)).into()),
236            ],
237        ));
238        assert_eq!(table_ref, expected);
239
240        let table_ref = RemoteTableRef::parse_with_default_dialect("schema.Table(1, 2)").unwrap();
241        let expected = RemoteTableRef::from((
242            TableReference::partial("schema", "Table"),
243            vec![
244                FunctionArg::Unnamed(Expr::value(Value::Number("1".to_string(), false)).into()),
245                FunctionArg::Unnamed(Expr::value(Value::Number("2".to_string(), false)).into()),
246            ],
247        ));
248        assert_eq!(table_ref, expected);
249    }
250
251    #[test]
252    fn multi_table_reference_with_args_and_whitespace() {
253        let table_ref = RemoteTableRef::parse_with_default_dialect("schema.table (1, 2)").unwrap();
254        let expected = RemoteTableRef::from((
255            TableReference::partial("schema", "table"),
256            vec![
257                FunctionArg::Unnamed(Expr::value(Value::Number("1".to_string(), false)).into()),
258                FunctionArg::Unnamed(Expr::value(Value::Number("2".to_string(), false)).into()),
259            ],
260        ));
261        assert_eq!(table_ref, expected);
262    }
263
264    #[test]
265    fn bare_reference_with_named_args() {
266        let table_ref = RemoteTableRef::parse_with_dialect(
267            "Table (user_id => 1, age => 2)",
268            &dialect::PostgreSqlDialect {},
269        )
270        .unwrap();
271        let expected = RemoteTableRef::from((
272            TableReference::bare("Table"),
273            vec![
274                FunctionArg::ExprNamed {
275                    name: ast::Expr::Identifier(Ident::new("user_id")),
276                    arg: Expr::value(Value::Number("1".to_string(), false)).into(),
277                    operator: FunctionArgOperator::RightArrow,
278                },
279                FunctionArg::ExprNamed {
280                    name: ast::Expr::Identifier(Ident::new("age")),
281                    arg: Expr::value(Value::Number("2".to_string(), false)).into(),
282                    operator: FunctionArgOperator::RightArrow,
283                },
284            ],
285        ));
286        assert_eq!(table_ref, expected);
287    }
288}
datafusion_federation/sql/table_reference.rs

datafusion_federation/sql/
table_reference.rs
