mirrors/datafusion-sqlparse
1
0
Fork 0
mirror of https://github.com/apache/datafusion-sqlparser-rs.git synced 2025-07-07 17:04:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 3

PartiQL queries in Redshift (#1534)

Browse source
This commit is contained in:
Yoav Cohen 2024-11-23 13:14:38 +01:00 • committed by GitHub
parent 10519003ed
commit 62fa8604af
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
17 changed files with 254 additions and 10 deletions
Download patch file Download diff file Expand all files Collapse all files

6
src/ast/query.rs
View file

@ -974,6 +974,8 @@ pub enum TableFactor {
with_ordinality: bool,
/// [Partition selection](https://dev.mysql.com/doc/refman/8.0/en/partitioning-selection.html), supported by MySQL.
partitions: Vec<Ident>,
/// Optional PartiQL JsonPath: <https://partiql.org/dql/from.html>
json_path: Option<JsonPath>,
},
Derived {
lateral: bool,
@ -1375,8 +1377,12 @@ impl fmt::Display for TableFactor {
version,
partitions,
with_ordinality,
json_path,
} => {
write!(f, "{name}")?;
if let Some(json_path) = json_path {
write!(f, "{json_path}")?;
}
if !partitions.is_empty() {
write!(f, "PARTITION ({})", display_comma_separated(partitions))?;
}

6
src/dialect/mod.rs
View file

@ -675,6 +675,12 @@ pub trait Dialect: Debug + Any {
fn supports_create_table_select(&self) -> bool {
false
}
/// Returns true if the dialect supports PartiQL for querying semi-structured data
/// <https://partiql.org/index.html>
fn supports_partiql(&self) -> bool {
false
}
}
/// This represents the operators for which precedence must be defined

5
src/dialect/redshift.rs
View file

@ -74,4 +74,9 @@ impl Dialect for RedshiftSqlDialect {
fn supports_top_before_distinct(&self) -> bool {
true
}
/// Redshift supports PartiQL: <https://docs.aws.amazon.com/redshift/latest/dg/super-overview.html>
fn supports_partiql(&self) -> bool {
true
}
}

21
src/parser/mod.rs
View file

@ -2936,7 +2936,7 @@ impl<'a> Parser<'a> {
} else if Token::LBracket == tok {
if dialect_of!(self is PostgreSqlDialect | DuckDbDialect | GenericDialect) {
self.parse_subscript(expr)
} else if dialect_of!(self is SnowflakeDialect) {
} else if dialect_of!(self is SnowflakeDialect) || self.dialect.supports_partiql() {
self.prev_token();
self.parse_json_access(expr)
} else {
@ -3072,6 +3072,14 @@ impl<'a> Parser<'a> {
}
fn parse_json_access(&mut self, expr: Expr) -> Result<Expr, ParserError> {
let path = self.parse_json_path()?;
Ok(Expr::JsonAccess {
value: Box::new(expr),
path,
})
}
fn parse_json_path(&mut self) -> Result<JsonPath, ParserError> {
let mut path = Vec::new();
loop {
match self.next_token().token {
@ -3095,10 +3103,7 @@ impl<'a> Parser<'a> {
}
debug_assert!(!path.is_empty());
Ok(Expr::JsonAccess {
value: Box::new(expr),
path: JsonPath { path },
})
Ok(JsonPath { path })
}
pub fn parse_map_access(&mut self, expr: Expr) -> Result<Expr, ParserError> {
@ -10338,6 +10343,11 @@ impl<'a> Parser<'a> {
} else {
let name = self.parse_object_name(true)?;
let json_path = match self.peek_token().token {
Token::LBracket if self.dialect.supports_partiql() => Some(self.parse_json_path()?),
_ => None,
};
let partitions: Vec<Ident> = if dialect_of!(self is MySqlDialect | GenericDialect)
&& self.parse_keyword(Keyword::PARTITION)
{
@ -10380,6 +10390,7 @@ impl<'a> Parser<'a> {
version,
partitions,
with_ordinality,
json_path,
};
while let Some(kw) = self.parse_one_of_keywords(&[Keyword::PIVOT, Keyword::UNPIVOT]) {

2
src/test_utils.rs
View file

@ -345,6 +345,7 @@ pub fn table(name: impl Into<String>) -> TableFactor {
version: None,
partitions: vec![],
with_ordinality: false,
json_path: None,
}
}
@ -360,6 +361,7 @@ pub fn table_with_alias(name: impl Into<String>, alias: impl Into<String>) -> Ta
version: None,
partitions: vec![],
with_ordinality: false,
json_path: None,
}
}

5
tests/sqlparser_bigquery.rs
View file

@ -229,6 +229,7 @@ fn parse_delete_statement() {
version: None,
partitions: vec![],
with_ordinality: false,
json_path: None,
},
from[0].relation
);
@ -1373,6 +1374,7 @@ fn parse_table_identifiers() {
version: None,
partitions: vec![],
with_ordinality: false,
json_path: None,
},
joins: vec![]
},]
@ -1546,6 +1548,7 @@ fn parse_table_time_travel() {
))),
partitions: vec![],
with_ordinality: false,
json_path: None,
},
joins: vec![]
},]
@ -1644,6 +1647,7 @@ fn parse_merge() {
version: Default::default(),
partitions: Default::default(),
with_ordinality: false,
json_path: None,
},
table
);
@ -1659,6 +1663,7 @@ fn parse_merge() {
version: Default::default(),
partitions: Default::default(),
with_ordinality: false,
json_path: None,
},
source
);

2
tests/sqlparser_clickhouse.rs
View file

@ -67,6 +67,7 @@ fn parse_map_access_expr() {
version: None,
partitions: vec![],
with_ordinality: false,
json_path: None,
},
joins: vec![],
}],
@ -172,6 +173,7 @@ fn parse_delimited_identifiers() {
version,
with_ordinality: _,
partitions: _,
json_path: _,
} => {
assert_eq!(vec![Ident::with_quote('"', "a table")], name.0);
assert_eq!(Ident::with_quote('"', "alias"), alias.unwrap().name);

37
tests/sqlparser_common.rs
View file

@ -364,6 +364,7 @@ fn parse_update_set_from() {
version: None,
partitions: vec![],
with_ordinality: false,
json_path: None,
},
joins: vec![],
},
@ -394,6 +395,7 @@ fn parse_update_set_from() {
version: None,
partitions: vec![],
with_ordinality: false,
json_path: None,
},
joins: vec![],
}],
@ -473,6 +475,7 @@ fn parse_update_with_table_alias() {
version: None,
partitions: vec![],
with_ordinality: false,
json_path: None,
},
joins: vec![],
},
@ -564,6 +567,7 @@ fn parse_select_with_table_alias() {
version: None,
partitions: vec![],
with_ordinality: false,
json_path: None,
},
joins: vec![],
}]
@ -601,6 +605,7 @@ fn parse_delete_statement() {
version: None,
partitions: vec![],
with_ordinality: false,
json_path: None,
},
from[0].relation
);
@ -648,6 +653,7 @@ fn parse_delete_statement_for_multi_tables() {
version: None,
partitions: vec![],
with_ordinality: false,
json_path: None,
},
from[0].relation
);
@ -660,6 +666,7 @@ fn parse_delete_statement_for_multi_tables() {
version: None,
partitions: vec![],
with_ordinality: false,
json_path: None,
},
from[0].joins[0].relation
);
@ -686,6 +693,7 @@ fn parse_delete_statement_for_multi_tables_with_using() {
version: None,
partitions: vec![],
with_ordinality: false,
json_path: None,
},
from[0].relation
);
@ -698,6 +706,7 @@ fn parse_delete_statement_for_multi_tables_with_using() {
version: None,
partitions: vec![],
with_ordinality: false,
json_path: None,
},
from[1].relation
);
@ -710,6 +719,7 @@ fn parse_delete_statement_for_multi_tables_with_using() {
version: None,
partitions: vec![],
with_ordinality: false,
json_path: None,
},
using[0].relation
);
@ -722,6 +732,7 @@ fn parse_delete_statement_for_multi_tables_with_using() {
version: None,
partitions: vec![],
with_ordinality: false,
json_path: None,
},
using[0].joins[0].relation
);
@ -753,6 +764,7 @@ fn parse_where_delete_statement() {
version: None,
partitions: vec![],
with_ordinality: false,
json_path: None,
},
from[0].relation,
);
@ -798,6 +810,7 @@ fn parse_where_delete_with_alias_statement() {
version: None,
partitions: vec![],
with_ordinality: false,
json_path: None,
},
from[0].relation,
);
@ -814,6 +827,7 @@ fn parse_where_delete_with_alias_statement() {
version: None,
partitions: vec![],
with_ordinality: false,
json_path: None,
},
joins: vec![],
}]),
@ -4718,6 +4732,7 @@ fn test_parse_named_window() {
version: None,
partitions: vec![],
with_ordinality: false,
json_path: None,
},
joins: vec![],
}],
@ -5301,6 +5316,7 @@ fn parse_interval_and_or_xor() {
version: None,
partitions: vec![],
with_ordinality: false,
json_path: None,
},
joins: vec![],
}],
@ -5912,6 +5928,7 @@ fn parse_implicit_join() {
version: None,
partitions: vec![],
with_ordinality: false,
json_path: None,
},
joins: vec![],
},
@ -5924,6 +5941,7 @@ fn parse_implicit_join() {
version: None,
partitions: vec![],
with_ordinality: false,
json_path: None,
},
joins: vec![],
},
@ -5944,6 +5962,7 @@ fn parse_implicit_join() {
version: None,
partitions: vec![],
with_ordinality: false,
json_path: None,
},
joins: vec![Join {
relation: TableFactor::Table {
@ -5954,6 +5973,7 @@ fn parse_implicit_join() {
version: None,
partitions: vec![],
with_ordinality: false,
json_path: None,
},
global: false,
join_operator: JoinOperator::Inner(JoinConstraint::Natural),
@ -5968,6 +5988,7 @@ fn parse_implicit_join() {
version: None,
partitions: vec![],
with_ordinality: false,
json_path: None,
},
joins: vec![Join {
relation: TableFactor::Table {
@ -5978,6 +5999,7 @@ fn parse_implicit_join() {
version: None,
partitions: vec![],
with_ordinality: false,
json_path: None,
},
global: false,
join_operator: JoinOperator::Inner(JoinConstraint::Natural),
@ -6002,6 +6024,7 @@ fn parse_cross_join() {
version: None,
partitions: vec![],
with_ordinality: false,
json_path: None,
},
global: false,
join_operator: JoinOperator::CrossJoin,
@ -6027,6 +6050,7 @@ fn parse_joins_on() {
version: None,
partitions: vec![],
with_ordinality: false,
json_path: None,
},
global,
join_operator: f(JoinConstraint::On(Expr::BinaryOp {
@ -6154,6 +6178,7 @@ fn parse_joins_using() {
version: None,
partitions: vec![],
with_ordinality: false,
json_path: None,
},
global: false,
join_operator: f(JoinConstraint::Using(vec!["c1".into()])),
@ -6227,6 +6252,7 @@ fn parse_natural_join() {
version: None,
partitions: vec![],
with_ordinality: false,
json_path: None,
},
global: false,
join_operator: f(JoinConstraint::Natural),
@ -6496,6 +6522,7 @@ fn parse_derived_tables() {
version: None,
partitions: vec![],
with_ordinality: false,
json_path: None,
},
global: false,
join_operator: JoinOperator::Inner(JoinConstraint::Natural),
@ -7443,6 +7470,7 @@ fn lateral_function() {
version: None,
partitions: vec![],
with_ordinality: false,
json_path: None,
},
joins: vec![Join {
relation: TableFactor::Function {
@ -8258,6 +8286,7 @@ fn parse_merge() {
version: None,
partitions: vec![],
with_ordinality: false,
json_path: None,
}
);
assert_eq!(table, table_no_into);
@ -8285,6 +8314,7 @@ fn parse_merge() {
version: None,
partitions: vec![],
with_ordinality: false,
json_path: None,
},
joins: vec![],
}],
@ -9359,6 +9389,7 @@ fn parse_pivot_table() {
version: None,
partitions: vec![],
with_ordinality: false,
json_path: None,
}),
aggregate_functions: vec![
expected_function("a", None),
@ -9432,6 +9463,7 @@ fn parse_unpivot_table() {
version: None,
partitions: vec![],
with_ordinality: false,
json_path: None,
}),
value: Ident {
value: "quantity".to_string(),
@ -9499,6 +9531,7 @@ fn parse_pivot_unpivot_table() {
version: None,
partitions: vec![],
with_ordinality: false,
json_path: None,
}),
value: Ident {
value: "population".to_string(),
@ -9910,6 +9943,7 @@ fn parse_unload() {
version: None,
partitions: vec![],
with_ordinality: false,
json_path: None,
},
joins: vec![],
}],
@ -10089,6 +10123,7 @@ fn parse_connect_by() {
version: None,
partitions: vec![],
with_ordinality: false,
json_path: None,
},
joins: vec![],
}],
@ -10176,6 +10211,7 @@ fn parse_connect_by() {
version: None,
partitions: vec![],
with_ordinality: false,
json_path: None,
},
joins: vec![],
}],
@ -10337,6 +10373,7 @@ fn test_match_recognize() {
version: None,
partitions: vec![],
with_ordinality: false,
json_path: None,
};
fn check(options: &str, expect: TableFactor) {

1
tests/sqlparser_databricks.rs
View file

@ -193,6 +193,7 @@ fn test_values_clause() {
version: None,
partitions: vec![],
with_ordinality: false,
json_path: None,
}),
query
.body

2
tests/sqlparser_duckdb.rs
View file

@ -282,6 +282,7 @@ fn test_select_union_by_name() {
version: None,
partitions: vec![],
with_ordinality: false,
json_path: None,
},
joins: vec![],
}],
@ -323,6 +324,7 @@ fn test_select_union_by_name() {
version: None,
partitions: vec![],
with_ordinality: false,
json_path: None,
},
joins: vec![],
}],

1
tests/sqlparser_hive.rs
View file

@ -457,6 +457,7 @@ fn parse_delimited_identifiers() {
version,
with_ordinality: _,
partitions: _,
json_path: _,
} => {
assert_eq!(vec![Ident::with_quote('"', "a table")], name.0);
assert_eq!(Ident::with_quote('"', "alias"), alias.unwrap().name);

18
tests/sqlparser_mssql.rs
View file

@ -70,6 +70,7 @@ fn parse_table_time_travel() {
))),
partitions: vec![],
with_ordinality: false,
json_path: None,
},
joins: vec![]
},]
@ -218,7 +219,8 @@ fn parse_mssql_openjson() {
with_hints: vec![],
version: None,
with_ordinality: false,
partitions: vec![]
partitions: vec![],
json_path: None,
},
joins: vec![Join {
relation: TableFactor::OpenJsonTable {
@ -293,7 +295,8 @@ fn parse_mssql_openjson() {
with_hints: vec![],
version: None,
with_ordinality: false,
partitions: vec![]
partitions: vec![],
json_path: None,
},
joins: vec![Join {
relation: TableFactor::OpenJsonTable {
@ -368,7 +371,8 @@ fn parse_mssql_openjson() {
with_hints: vec![],
version: None,
with_ordinality: false,
partitions: vec![]
partitions: vec![],
json_path: None,
},
joins: vec![Join {
relation: TableFactor::OpenJsonTable {
@ -443,7 +447,8 @@ fn parse_mssql_openjson() {
with_hints: vec![],
version: None,
with_ordinality: false,
partitions: vec![]
partitions: vec![],
json_path: None,
},
joins: vec![Join {
relation: TableFactor::OpenJsonTable {
@ -496,7 +501,8 @@ fn parse_mssql_openjson() {
with_hints: vec![],
version: None,
with_ordinality: false,
partitions: vec![]
partitions: vec![],
json_path: None,
},
joins: vec![Join {
relation: TableFactor::OpenJsonTable {
@ -679,6 +685,7 @@ fn parse_delimited_identifiers() {
version,
with_ordinality: _,
partitions: _,
json_path: _,
} => {
assert_eq!(vec![Ident::with_quote('"', "a table")], name.0);
assert_eq!(Ident::with_quote('"', "alias"), alias.unwrap().name);
@ -1314,6 +1321,7 @@ fn parse_substring_in_select() {
version: None,
partitions: vec![],
with_ordinality: false,
json_path: None,
},
joins: vec![]
}],

5
tests/sqlparser_mysql.rs
View file

@ -1862,6 +1862,7 @@ fn parse_select_with_numeric_prefix_column_name() {
version: None,
partitions: vec![],
with_ordinality: false,
json_path: None,
},
joins: vec![]
}],
@ -1918,6 +1919,7 @@ fn parse_select_with_concatenation_of_exp_number_and_numeric_prefix_column() {
version: None,
partitions: vec![],
with_ordinality: false,
json_path: None,
},
joins: vec![]
}],
@ -1985,6 +1987,7 @@ fn parse_update_with_joins() {
version: None,
partitions: vec![],
with_ordinality: false,
json_path: None,
},
joins: vec![Join {
relation: TableFactor::Table {
@ -1998,6 +2001,7 @@ fn parse_update_with_joins() {
version: None,
partitions: vec![],
with_ordinality: false,
json_path: None,
},
global: false,
join_operator: JoinOperator::Inner(JoinConstraint::On(Expr::BinaryOp {
@ -2428,6 +2432,7 @@ fn parse_substring_in_select() {
version: None,
partitions: vec![],
with_ordinality: false,
json_path: None,
},
joins: vec![]
}],

1
tests/sqlparser_postgres.rs
View file

@ -3511,6 +3511,7 @@ fn parse_delimited_identifiers() {
version,
with_ordinality: _,
partitions: _,
json_path: _,
} => {
assert_eq!(vec![Ident::with_quote('"', "a table")], name.0);
assert_eq!(Ident::with_quote('"', "alias"), alias.unwrap().name);

150
tests/sqlparser_redshift.rs
View file

@ -54,6 +54,7 @@ fn test_square_brackets_over_db_schema_table_name() {
version: None,
partitions: vec![],
with_ordinality: false,
json_path: None,
},
joins: vec![],
}
@ -101,6 +102,7 @@ fn test_double_quotes_over_db_schema_table_name() {
version: None,
partitions: vec![],
with_ordinality: false,
json_path: None,
},
joins: vec![],
}
@ -123,6 +125,7 @@ fn parse_delimited_identifiers() {
version,
with_ordinality: _,
partitions: _,
json_path: _,
} => {
assert_eq!(vec![Ident::with_quote('"', "a table")], name.0);
assert_eq!(Ident::with_quote('"', "alias"), alias.unwrap().name);
@ -196,3 +199,150 @@ fn test_create_view_with_no_schema_binding() {
redshift_and_generic()
.verified_stmt("CREATE VIEW myevent AS SELECT eventname FROM event WITH NO SCHEMA BINDING");
}
#[test]
fn test_redshift_json_path() {
let dialects = all_dialects_where(|d| d.supports_partiql());
let sql = "SELECT cust.c_orders[0].o_orderkey FROM customer_orders_lineitem";
let select = dialects.verified_only_select(sql);
assert_eq!(
&Expr::JsonAccess {
value: Box::new(Expr::CompoundIdentifier(vec![
Ident::new("cust"),
Ident::new("c_orders")
])),
path: JsonPath {
path: vec![
JsonPathElem::Bracket {
key: Expr::Value(Value::Number("0".parse().unwrap(), false))
},
JsonPathElem::Dot {
key: "o_orderkey".to_string(),
quoted: false
}
]
}
},
expr_from_projection(only(&select.projection))
);
let sql = "SELECT cust.c_orders[0]['id'] FROM customer_orders_lineitem";
let select = dialects.verified_only_select(sql);
assert_eq!(
&Expr::JsonAccess {
value: Box::new(Expr::CompoundIdentifier(vec![
Ident::new("cust"),
Ident::new("c_orders")
])),
path: JsonPath {
path: vec![
JsonPathElem::Bracket {
key: Expr::Value(Value::Number("0".parse().unwrap(), false))
},
JsonPathElem::Bracket {
key: Expr::Value(Value::SingleQuotedString("id".to_owned()))
}
]
}
},
expr_from_projection(only(&select.projection))
);
let sql = "SELECT db1.sc1.tbl1.col1[0]['id'] FROM customer_orders_lineitem";
let select = dialects.verified_only_select(sql);
assert_eq!(
&Expr::JsonAccess {
value: Box::new(Expr::CompoundIdentifier(vec![
Ident::new("db1"),
Ident::new("sc1"),
Ident::new("tbl1"),
Ident::new("col1")
])),
path: JsonPath {
path: vec![
JsonPathElem::Bracket {
key: Expr::Value(Value::Number("0".parse().unwrap(), false))
},
JsonPathElem::Bracket {
key: Expr::Value(Value::SingleQuotedString("id".to_owned()))
}
]
}
},
expr_from_projection(only(&select.projection))
);
}
#[test]
fn test_parse_json_path_from() {
let dialects = all_dialects_where(|d| d.supports_partiql());
let select = dialects.verified_only_select("SELECT * FROM src[0].a AS a");
match &select.from[0].relation {
TableFactor::Table {
name, json_path, ..
} => {
assert_eq!(name, &ObjectName(vec![Ident::new("src")]));
assert_eq!(
json_path,
&Some(JsonPath {
path: vec![
JsonPathElem::Bracket {
key: Expr::Value(Value::Number("0".parse().unwrap(), false))
},
JsonPathElem::Dot {
key: "a".to_string(),
quoted: false
}
]
})
);
}
_ => panic!(),
}
let select = dialects.verified_only_select("SELECT * FROM src[0].a[1].b AS a");
match &select.from[0].relation {
TableFactor::Table {
name, json_path, ..
} => {
assert_eq!(name, &ObjectName(vec![Ident::new("src")]));
assert_eq!(
json_path,
&Some(JsonPath {
path: vec![
JsonPathElem::Bracket {
key: Expr::Value(Value::Number("0".parse().unwrap(), false))
},
JsonPathElem::Dot {
key: "a".to_string(),
quoted: false
},
JsonPathElem::Bracket {
key: Expr::Value(Value::Number("1".parse().unwrap(), false))
},
JsonPathElem::Dot {
key: "b".to_string(),
quoted: false
},
]
})
);
}
_ => panic!(),
}
let select = dialects.verified_only_select("SELECT * FROM src.a.b");
match &select.from[0].relation {
TableFactor::Table {
name, json_path, ..
} => {
assert_eq!(
name,
&ObjectName(vec![Ident::new("src"), Ident::new("a"), Ident::new("b")])
);
assert_eq!(json_path, &None);
}
_ => panic!(),
}
}

1
tests/sqlparser_snowflake.rs
View file

@ -1190,6 +1190,7 @@ fn parse_delimited_identifiers() {
version,
with_ordinality: _,
partitions: _,
json_path: _,
} => {
assert_eq!(vec![Ident::with_quote('"', "a table")], name.0);
assert_eq!(Ident::with_quote('"', "alias"), alias.unwrap().name);

1
tests/sqlparser_sqlite.rs
View file

@ -486,6 +486,7 @@ fn parse_update_tuple_row_values() {
version: None,
partitions: vec![],
with_ordinality: false,
json_path: None,
},
joins: vec![],
},