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

datafusion_federation/sql/
table_reference.rs
