datafusion_federation/sql/
table_reference.rs

1use std::sync::Arc;
2
3use datafusion::{
4    error::DataFusionError,
5    sql::{
6        sqlparser::{
7            self,
8            ast::FunctionArg,
9            dialect::{Dialect, GenericDialect},
10            tokenizer::Token,
11        },
12        TableReference,
13    },
14};
15
16/// A multipart identifier to a remote table, view or parameterized view.
17///
18/// RemoteTableRef can be created by parsing from a string represeting a table obbject with optional
19/// ```rust
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')", &PostgresSqlDialect {});
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 specfic 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            (Some(catalog), Some(schema), Some(table)) => TableReference::full(
56                catalog.value.clone(),
57                schema.value.clone(),
58                table.value.clone(),
59            ),
60            (Some(schema), Some(table), None) => {
61                TableReference::partial(schema.value.clone(), table.value.clone())
62            }
63            (Some(table), None, None) => TableReference::bare(table.value.clone()),
64            _ => {
65                return Err(DataFusionError::NotImplemented(
66                    "Unable to parse string into TableReference".to_string(),
67                ))
68            }
69        };
70
71        if !args.is_empty() {
72            Ok(RemoteTableRef {
73                table_ref,
74                args: Some(args.into()),
75            })
76        } else {
77            Ok(RemoteTableRef {
78                table_ref,
79                args: None,
80            })
81        }
82    }
83
84    pub fn table_ref(&self) -> &TableReference {
85        &self.table_ref
86    }
87
88    pub fn args(&self) -> Option<&[FunctionArg]> {
89        self.args.as_deref()
90    }
91}
92
93impl From<TableReference> for RemoteTableRef {
94    fn from(table_ref: TableReference) -> Self {
95        RemoteTableRef {
96            table_ref,
97            args: None,
98        }
99    }
100}
101
102impl From<RemoteTableRef> for TableReference {
103    fn from(remote_table_ref: RemoteTableRef) -> Self {
104        remote_table_ref.table_ref
105    }
106}
107
108impl From<&RemoteTableRef> for TableReference {
109    fn from(remote_table_ref: &RemoteTableRef) -> Self {
110        remote_table_ref.table_ref.clone()
111    }
112}
113
114impl From<(TableReference, Vec<FunctionArg>)> for RemoteTableRef {
115    fn from((table_ref, args): (TableReference, Vec<FunctionArg>)) -> Self {
116        RemoteTableRef {
117            table_ref,
118            args: Some(args.into()),
119        }
120    }
121}
122
123impl TryFrom<&str> for RemoteTableRef {
124    type Error = DataFusionError;
125    fn try_from(s: &str) -> Result<Self, Self::Error> {
126        Self::parse_with_default_dialect(s)
127    }
128}
129
130impl TryFrom<String> for RemoteTableRef {
131    type Error = DataFusionError;
132    fn try_from(s: String) -> Result<Self, Self::Error> {
133        Self::parse_with_default_dialect(&s)
134    }
135}
136
137impl TryFrom<&String> for RemoteTableRef {
138    type Error = DataFusionError;
139    fn try_from(s: &String) -> Result<Self, Self::Error> {
140        Self::parse_with_default_dialect(s)
141    }
142}
143
144#[cfg(test)]
145mod tests {
146    use super::*;
147    use sqlparser::{
148        ast::{self, Expr, FunctionArgOperator, Ident, Value},
149        dialect,
150    };
151
152    #[test]
153    fn bare_table_reference() {
154        let table_ref = RemoteTableRef::parse_with_default_dialect("table").unwrap();
155        let expected = RemoteTableRef::from(TableReference::bare("table"));
156        assert_eq!(table_ref, expected);
157
158        let table_ref = RemoteTableRef::parse_with_default_dialect("Table").unwrap();
159        let expected = RemoteTableRef::from(TableReference::bare("Table"));
160        assert_eq!(table_ref, expected);
161    }
162
163    #[test]
164    fn bare_table_reference_with_args() {
165        let table_ref = RemoteTableRef::parse_with_default_dialect("table(1, 2)").unwrap();
166        let expected = RemoteTableRef::from((
167            TableReference::bare("table"),
168            vec![
169                FunctionArg::Unnamed(Expr::Value(Value::Number("1".to_string(), false)).into()),
170                FunctionArg::Unnamed(Expr::Value(Value::Number("2".to_string(), false)).into()),
171            ],
172        ));
173        assert_eq!(table_ref, expected);
174
175        let table_ref = RemoteTableRef::parse_with_default_dialect("Table(1, 2)").unwrap();
176        let expected = RemoteTableRef::from((
177            TableReference::bare("Table"),
178            vec![
179                FunctionArg::Unnamed(Expr::Value(Value::Number("1".to_string(), false)).into()),
180                FunctionArg::Unnamed(Expr::Value(Value::Number("2".to_string(), false)).into()),
181            ],
182        ));
183        assert_eq!(table_ref, expected);
184    }
185
186    #[test]
187    fn bare_table_reference_with_args_and_whitespace() {
188        let table_ref = RemoteTableRef::parse_with_default_dialect("table (1, 2)").unwrap();
189        let expected = RemoteTableRef::from((
190            TableReference::bare("table"),
191            vec![
192                FunctionArg::Unnamed(Expr::Value(Value::Number("1".to_string(), false)).into()),
193                FunctionArg::Unnamed(Expr::Value(Value::Number("2".to_string(), false)).into()),
194            ],
195        ));
196        assert_eq!(table_ref, expected);
197
198        let table_ref = RemoteTableRef::parse_with_default_dialect("Table (1, 2)").unwrap();
199        let expected = RemoteTableRef::from((
200            TableReference::bare("Table"),
201            vec![
202                FunctionArg::Unnamed(Expr::Value(Value::Number("1".to_string(), false)).into()),
203                FunctionArg::Unnamed(Expr::Value(Value::Number("2".to_string(), false)).into()),
204            ],
205        ));
206        assert_eq!(table_ref, expected);
207    }
208
209    #[test]
210    fn multi_table_reference_with_no_args() {
211        let table_ref = RemoteTableRef::parse_with_default_dialect("schema.table").unwrap();
212        let expected = RemoteTableRef::from(TableReference::partial("schema", "table"));
213        assert_eq!(table_ref, expected);
214
215        let table_ref = RemoteTableRef::parse_with_default_dialect("schema.Table").unwrap();
216        let expected = RemoteTableRef::from(TableReference::partial("schema", "Table"));
217        assert_eq!(table_ref, expected);
218    }
219
220    #[test]
221    fn multi_table_reference_with_args() {
222        let table_ref = RemoteTableRef::parse_with_default_dialect("schema.table(1, 2)").unwrap();
223        let expected = RemoteTableRef::from((
224            TableReference::partial("schema", "table"),
225            vec![
226                FunctionArg::Unnamed(Expr::Value(Value::Number("1".to_string(), false)).into()),
227                FunctionArg::Unnamed(Expr::Value(Value::Number("2".to_string(), false)).into()),
228            ],
229        ));
230        assert_eq!(table_ref, expected);
231
232        let table_ref = RemoteTableRef::parse_with_default_dialect("schema.Table(1, 2)").unwrap();
233        let expected = RemoteTableRef::from((
234            TableReference::partial("schema", "Table"),
235            vec![
236                FunctionArg::Unnamed(Expr::Value(Value::Number("1".to_string(), false)).into()),
237                FunctionArg::Unnamed(Expr::Value(Value::Number("2".to_string(), false)).into()),
238            ],
239        ));
240        assert_eq!(table_ref, expected);
241    }
242
243    #[test]
244    fn multi_table_reference_with_args_and_whitespace() {
245        let table_ref = RemoteTableRef::parse_with_default_dialect("schema.table (1, 2)").unwrap();
246        let expected = RemoteTableRef::from((
247            TableReference::partial("schema", "table"),
248            vec![
249                FunctionArg::Unnamed(Expr::Value(Value::Number("1".to_string(), false)).into()),
250                FunctionArg::Unnamed(Expr::Value(Value::Number("2".to_string(), false)).into()),
251            ],
252        ));
253        assert_eq!(table_ref, expected);
254    }
255
256    #[test]
257    fn bare_reference_with_named_args() {
258        let table_ref = RemoteTableRef::parse_with_dialect(
259            "Table (user_id => 1, age => 2)",
260            &dialect::PostgreSqlDialect {},
261        )
262        .unwrap();
263        let expected = RemoteTableRef::from((
264            TableReference::bare("Table"),
265            vec![
266                FunctionArg::ExprNamed {
267                    name: ast::Expr::Identifier(Ident::new("user_id")),
268                    arg: Expr::Value(Value::Number("1".to_string(), false)).into(),
269                    operator: FunctionArgOperator::RightArrow,
270                },
271                FunctionArg::ExprNamed {
272                    name: ast::Expr::Identifier(Ident::new("age")),
273                    arg: Expr::Value(Value::Number("2".to_string(), false)).into(),
274                    operator: FunctionArgOperator::RightArrow,
275                },
276            ],
277        ));
278        assert_eq!(table_ref, expected);
279    }
280}
datafusion_federation/sql/table_reference.rs

datafusion_federation/sql/
table_reference.rs
