Add support of table function WITH ORDINALITY modifier for Postgre Parser (#1337)

This commit is contained in:
hulk 2024-07-10 05:43:22 +08:00 committed by GitHub
parent 285f492589
commit 9108bffc9a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
16 changed files with 129 additions and 4 deletions

View file

@ -913,6 +913,10 @@ pub enum TableFactor {
/// Optional version qualifier to facilitate table time-travel, as
/// supported by BigQuery and MSSQL.
version: Option<TableVersion>,
// Optional table function modifier to generate the ordinality for column.
/// For example, `SELECT * FROM generate_series(1, 10) WITH ORDINALITY AS t(a, b);`
/// [WITH ORDINALITY](https://www.postgresql.org/docs/current/functions-srf.html), supported by Postgres.
with_ordinality: bool,
/// [Partition selection](https://dev.mysql.com/doc/refman/8.0/en/partitioning-selection.html), supported by MySQL.
partitions: Vec<Ident>,
},
@ -948,6 +952,7 @@ pub enum TableFactor {
array_exprs: Vec<Expr>,
with_offset: bool,
with_offset_alias: Option<Ident>,
with_ordinality: bool,
},
/// The `JSON_TABLE` table-valued function.
/// Part of the SQL standard, but implemented only by MySQL, Oracle, and DB2.
@ -1293,6 +1298,7 @@ impl fmt::Display for TableFactor {
with_hints,
version,
partitions,
with_ordinality,
} => {
write!(f, "{name}")?;
if !partitions.is_empty() {
@ -1301,6 +1307,9 @@ impl fmt::Display for TableFactor {
if let Some(args) = args {
write!(f, "({})", display_comma_separated(args))?;
}
if *with_ordinality {
write!(f, " WITH ORDINALITY")?;
}
if let Some(alias) = alias {
write!(f, " AS {alias}")?;
}
@ -1354,9 +1363,14 @@ impl fmt::Display for TableFactor {
array_exprs,
with_offset,
with_offset_alias,
with_ordinality,
} => {
write!(f, "UNNEST({})", display_comma_separated(array_exprs))?;
if *with_ordinality {
write!(f, " WITH ORDINALITY")?;
}
if let Some(alias) = alias {
write!(f, " AS {alias}")?;
}

View file

@ -518,6 +518,7 @@ define_keywords!(
OR,
ORC,
ORDER,
ORDINALITY,
OUT,
OUTER,
OUTPUTFORMAT,

View file

@ -9209,6 +9209,7 @@ impl<'a> Parser<'a> {
let array_exprs = self.parse_comma_separated(Parser::parse_expr)?;
self.expect_token(&Token::RParen)?;
let with_ordinality = self.parse_keywords(&[Keyword::WITH, Keyword::ORDINALITY]);
let alias = match self.parse_optional_table_alias(keywords::RESERVED_FOR_TABLE_ALIAS) {
Ok(Some(alias)) => Some(alias),
Ok(None) => None,
@ -9235,6 +9236,7 @@ impl<'a> Parser<'a> {
array_exprs,
with_offset,
with_offset_alias,
with_ordinality,
})
} else if self.parse_keyword_with_tokens(Keyword::JSON_TABLE, &[Token::LParen]) {
let json_expr = self.parse_expr()?;
@ -9273,6 +9275,8 @@ impl<'a> Parser<'a> {
None
};
let with_ordinality = self.parse_keywords(&[Keyword::WITH, Keyword::ORDINALITY]);
let alias = self.parse_optional_table_alias(keywords::RESERVED_FOR_TABLE_ALIAS)?;
// MSSQL-specific table hints:
@ -9294,6 +9298,7 @@ impl<'a> Parser<'a> {
with_hints,
version,
partitions,
with_ordinality,
};
while let Some(kw) = self.parse_one_of_keywords(&[Keyword::PIVOT, Keyword::UNPIVOT]) {

View file

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

View file

@ -224,6 +224,7 @@ fn parse_delete_statement() {
with_hints: vec![],
version: None,
partitions: vec![],
with_ordinality: false,
},
from[0].relation
);
@ -1353,6 +1354,7 @@ fn parse_table_identifiers() {
with_hints: vec![],
version: None,
partitions: vec![],
with_ordinality: false,
},
joins: vec![]
},]
@ -1525,6 +1527,7 @@ fn parse_table_time_travel() {
Value::SingleQuotedString(version)
))),
partitions: vec![],
with_ordinality: false,
},
joins: vec![]
},]
@ -1551,7 +1554,8 @@ fn parse_join_constraint_unnest_alias() {
Ident::new("a")
])],
with_offset: false,
with_offset_alias: None
with_offset_alias: None,
with_ordinality: false,
},
join_operator: JoinOperator::Inner(JoinConstraint::On(Expr::BinaryOp {
left: Box::new(Expr::Identifier("c1".into())),
@ -1620,6 +1624,7 @@ fn parse_merge() {
with_hints: Default::default(),
version: Default::default(),
partitions: Default::default(),
with_ordinality: false,
},
table
);
@ -1634,6 +1639,7 @@ fn parse_merge() {
with_hints: Default::default(),
version: Default::default(),
partitions: Default::default(),
with_ordinality: false,
},
source
);

View file

@ -59,6 +59,7 @@ fn parse_map_access_expr() {
with_hints: vec![],
version: None,
partitions: vec![],
with_ordinality: false,
},
joins: vec![],
}],
@ -162,6 +163,7 @@ fn parse_delimited_identifiers() {
args,
with_hints,
version,
with_ordinality: _,
partitions: _,
} => {
assert_eq!(vec![Ident::with_quote('"', "a table")], name.0);

View file

@ -359,6 +359,7 @@ fn parse_update_set_from() {
with_hints: vec![],
version: None,
partitions: vec![],
with_ordinality: false,
},
joins: vec![],
},
@ -387,6 +388,7 @@ fn parse_update_set_from() {
with_hints: vec![],
version: None,
partitions: vec![],
with_ordinality: false,
},
joins: vec![],
}],
@ -463,6 +465,7 @@ fn parse_update_with_table_alias() {
with_hints: vec![],
version: None,
partitions: vec![],
with_ordinality: false,
},
joins: vec![],
},
@ -530,6 +533,7 @@ fn parse_select_with_table_alias() {
with_hints: vec![],
version: None,
partitions: vec![],
with_ordinality: false,
},
joins: vec![],
}]
@ -566,6 +570,7 @@ fn parse_delete_statement() {
with_hints: vec![],
version: None,
partitions: vec![],
with_ordinality: false,
},
from[0].relation
);
@ -612,6 +617,7 @@ fn parse_delete_statement_for_multi_tables() {
with_hints: vec![],
version: None,
partitions: vec![],
with_ordinality: false,
},
from[0].relation
);
@ -623,6 +629,7 @@ fn parse_delete_statement_for_multi_tables() {
with_hints: vec![],
version: None,
partitions: vec![],
with_ordinality: false,
},
from[0].joins[0].relation
);
@ -648,6 +655,7 @@ fn parse_delete_statement_for_multi_tables_with_using() {
with_hints: vec![],
version: None,
partitions: vec![],
with_ordinality: false,
},
from[0].relation
);
@ -659,6 +667,7 @@ fn parse_delete_statement_for_multi_tables_with_using() {
with_hints: vec![],
version: None,
partitions: vec![],
with_ordinality: false,
},
from[1].relation
);
@ -670,6 +679,7 @@ fn parse_delete_statement_for_multi_tables_with_using() {
with_hints: vec![],
version: None,
partitions: vec![],
with_ordinality: false,
},
using[0].relation
);
@ -681,6 +691,7 @@ fn parse_delete_statement_for_multi_tables_with_using() {
with_hints: vec![],
version: None,
partitions: vec![],
with_ordinality: false,
},
using[0].joins[0].relation
);
@ -711,6 +722,7 @@ fn parse_where_delete_statement() {
with_hints: vec![],
version: None,
partitions: vec![],
with_ordinality: false,
},
from[0].relation,
);
@ -755,6 +767,7 @@ fn parse_where_delete_with_alias_statement() {
with_hints: vec![],
version: None,
partitions: vec![],
with_ordinality: false,
},
from[0].relation,
);
@ -770,6 +783,7 @@ fn parse_where_delete_with_alias_statement() {
with_hints: vec![],
version: None,
partitions: vec![],
with_ordinality: false,
},
joins: vec![],
}]),
@ -4551,6 +4565,7 @@ fn test_parse_named_window() {
with_hints: vec![],
version: None,
partitions: vec![],
with_ordinality: false,
},
joins: vec![],
}],
@ -4933,6 +4948,7 @@ fn parse_interval_and_or_xor() {
with_hints: vec![],
version: None,
partitions: vec![],
with_ordinality: false,
},
joins: vec![],
}],
@ -5286,6 +5302,7 @@ fn parse_unnest_in_from_clause() {
array_exprs: vec![Expr::Identifier(Ident::new("expr"))],
with_offset: true,
with_offset_alias: None,
with_ordinality: false,
},
joins: vec![],
}],
@ -5303,6 +5320,7 @@ fn parse_unnest_in_from_clause() {
array_exprs: vec![Expr::Identifier(Ident::new("expr"))],
with_offset: false,
with_offset_alias: None,
with_ordinality: false,
},
joins: vec![],
}],
@ -5320,6 +5338,7 @@ fn parse_unnest_in_from_clause() {
array_exprs: vec![Expr::Identifier(Ident::new("expr"))],
with_offset: true,
with_offset_alias: None,
with_ordinality: false,
},
joins: vec![],
}],
@ -5340,6 +5359,7 @@ fn parse_unnest_in_from_clause() {
array_exprs: vec![Expr::Identifier(Ident::new("expr"))],
with_offset: false,
with_offset_alias: None,
with_ordinality: false,
},
joins: vec![],
}],
@ -5364,6 +5384,7 @@ fn parse_unnest_in_from_clause() {
)],
with_offset: false,
with_offset_alias: None,
with_ordinality: false,
},
joins: vec![],
}],
@ -5394,6 +5415,7 @@ fn parse_unnest_in_from_clause() {
],
with_offset: false,
with_offset_alias: None,
with_ordinality: false,
},
joins: vec![],
}],
@ -5503,6 +5525,7 @@ fn parse_implicit_join() {
with_hints: vec![],
version: None,
partitions: vec![],
with_ordinality: false,
},
joins: vec![],
},
@ -5514,6 +5537,7 @@ fn parse_implicit_join() {
with_hints: vec![],
version: None,
partitions: vec![],
with_ordinality: false,
},
joins: vec![],
},
@ -5533,6 +5557,7 @@ fn parse_implicit_join() {
with_hints: vec![],
version: None,
partitions: vec![],
with_ordinality: false,
},
joins: vec![Join {
relation: TableFactor::Table {
@ -5542,6 +5567,7 @@ fn parse_implicit_join() {
with_hints: vec![],
version: None,
partitions: vec![],
with_ordinality: false,
},
join_operator: JoinOperator::Inner(JoinConstraint::Natural),
}],
@ -5554,6 +5580,7 @@ fn parse_implicit_join() {
with_hints: vec![],
version: None,
partitions: vec![],
with_ordinality: false,
},
joins: vec![Join {
relation: TableFactor::Table {
@ -5563,6 +5590,7 @@ fn parse_implicit_join() {
with_hints: vec![],
version: None,
partitions: vec![],
with_ordinality: false,
},
join_operator: JoinOperator::Inner(JoinConstraint::Natural),
}],
@ -5585,6 +5613,7 @@ fn parse_cross_join() {
with_hints: vec![],
version: None,
partitions: vec![],
with_ordinality: false,
},
join_operator: JoinOperator::CrossJoin,
},
@ -5607,6 +5636,7 @@ fn parse_joins_on() {
with_hints: vec![],
version: None,
partitions: vec![],
with_ordinality: false,
},
join_operator: f(JoinConstraint::On(Expr::BinaryOp {
left: Box::new(Expr::Identifier("c1".into())),
@ -5678,6 +5708,7 @@ fn parse_joins_using() {
with_hints: vec![],
version: None,
partitions: vec![],
with_ordinality: false,
},
join_operator: f(JoinConstraint::Using(vec!["c1".into()])),
}
@ -5741,6 +5772,7 @@ fn parse_natural_join() {
with_hints: vec![],
version: None,
partitions: vec![],
with_ordinality: false,
},
join_operator: f(JoinConstraint::Natural),
}
@ -6008,6 +6040,7 @@ fn parse_derived_tables() {
with_hints: vec![],
version: None,
partitions: vec![],
with_ordinality: false,
},
join_operator: JoinOperator::Inner(JoinConstraint::Natural),
}],
@ -6905,6 +6938,7 @@ fn lateral_function() {
with_hints: vec![],
version: None,
partitions: vec![],
with_ordinality: false,
},
joins: vec![Join {
relation: TableFactor::Function {
@ -7613,6 +7647,7 @@ fn parse_merge() {
with_hints: vec![],
version: None,
partitions: vec![],
with_ordinality: false,
}
);
assert_eq!(table, table_no_into);
@ -7638,6 +7673,7 @@ fn parse_merge() {
with_hints: vec![],
version: None,
partitions: vec![],
with_ordinality: false,
},
joins: vec![],
}],
@ -8700,6 +8736,7 @@ fn parse_pivot_table() {
with_hints: vec![],
version: None,
partitions: vec![],
with_ordinality: false,
}),
aggregate_functions: vec![
expected_function("a", None),
@ -8769,6 +8806,7 @@ fn parse_unpivot_table() {
with_hints: vec![],
version: None,
partitions: vec![],
with_ordinality: false,
}),
value: Ident {
value: "quantity".to_string(),
@ -8835,6 +8873,7 @@ fn parse_pivot_unpivot_table() {
with_hints: vec![],
version: None,
partitions: vec![],
with_ordinality: false,
}),
value: Ident {
value: "population".to_string(),
@ -9159,6 +9198,7 @@ fn parse_unload() {
with_hints: vec![],
version: None,
partitions: vec![],
with_ordinality: false,
},
joins: vec![],
}],
@ -9304,6 +9344,7 @@ fn parse_connect_by() {
with_hints: vec![],
version: None,
partitions: vec![],
with_ordinality: false,
},
joins: vec![],
}],
@ -9389,6 +9430,7 @@ fn parse_connect_by() {
with_hints: vec![],
version: None,
partitions: vec![],
with_ordinality: false,
},
joins: vec![],
}],
@ -9548,6 +9590,7 @@ fn test_match_recognize() {
with_hints: vec![],
version: None,
partitions: vec![],
with_ordinality: false,
};
fn check(options: &str, expect: TableFactor) {

View file

@ -177,7 +177,8 @@ fn test_values_clause() {
args: None,
with_hints: vec![],
version: None,
partitions: vec![]
partitions: vec![],
with_ordinality: false,
}),
query
.body

View file

@ -166,6 +166,7 @@ fn test_select_union_by_name() {
with_hints: vec![],
version: None,
partitions: vec![],
with_ordinality: false,
},
joins: vec![],
}],
@ -205,6 +206,7 @@ fn test_select_union_by_name() {
with_hints: vec![],
version: None,
partitions: vec![],
with_ordinality: false,
},
joins: vec![],
}],

View file

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

View file

@ -64,6 +64,7 @@ fn parse_table_time_travel() {
Value::SingleQuotedString(version)
))),
partitions: vec![],
with_ordinality: false,
},
joins: vec![]
},]
@ -335,6 +336,7 @@ fn parse_delimited_identifiers() {
args,
with_hints,
version,
with_ordinality: _,
partitions: _,
} => {
assert_eq!(vec![Ident::with_quote('"', "a table")], name.0);
@ -526,6 +528,7 @@ fn parse_substring_in_select() {
with_hints: vec![],
version: None,
partitions: vec![],
with_ordinality: false,
},
joins: vec![]
}],

View file

@ -1728,6 +1728,7 @@ fn parse_select_with_numeric_prefix_column_name() {
with_hints: vec![],
version: None,
partitions: vec![],
with_ordinality: false,
},
joins: vec![]
}],
@ -1782,6 +1783,7 @@ fn parse_select_with_concatenation_of_exp_number_and_numeric_prefix_column() {
with_hints: vec![],
version: None,
partitions: vec![],
with_ordinality: false,
},
joins: vec![]
}],
@ -1847,6 +1849,7 @@ fn parse_update_with_joins() {
with_hints: vec![],
version: None,
partitions: vec![],
with_ordinality: false,
},
joins: vec![Join {
relation: TableFactor::Table {
@ -1859,6 +1862,7 @@ fn parse_update_with_joins() {
with_hints: vec![],
version: None,
partitions: vec![],
with_ordinality: false,
},
join_operator: JoinOperator::Inner(JoinConstraint::On(Expr::BinaryOp {
left: Box::new(Expr::CompoundIdentifier(vec![
@ -2282,6 +2286,7 @@ fn parse_substring_in_select() {
with_hints: vec![],
version: None,
partitions: vec![],
with_ordinality: false,
},
joins: vec![]
}],

View file

@ -3501,6 +3501,7 @@ fn parse_delimited_identifiers() {
args,
with_hints,
version,
with_ordinality: _,
partitions: _,
} => {
assert_eq!(vec![Ident::with_quote('"', "a table")], name.0);
@ -4054,7 +4055,8 @@ fn parse_join_constraint_unnest_alias() {
Ident::new("a")
])],
with_offset: false,
with_offset_alias: None
with_offset_alias: None,
with_ordinality: false,
},
join_operator: JoinOperator::Inner(JoinConstraint::On(Expr::BinaryOp {
left: Box::new(Expr::Identifier("c1".into())),
@ -4362,3 +4364,36 @@ fn parse_create_table_with_options() {
_ => unreachable!(),
}
}
#[test]
fn test_table_function_with_ordinality() {
let from = pg_and_generic()
.verified_only_select("SELECT * FROM generate_series(1, 10) WITH ORDINALITY AS t")
.from;
assert_eq!(1, from.len());
match from[0].relation {
TableFactor::Table {
ref name,
with_ordinality: true,
..
} => {
assert_eq!("generate_series", name.to_string().as_str());
}
_ => panic!("Expecting TableFactor::Table with ordinality"),
}
}
#[test]
fn test_table_unnest_with_ordinality() {
let from = pg_and_generic()
.verified_only_select("SELECT * FROM UNNEST([10, 20, 30]) WITH ORDINALITY AS t")
.from;
assert_eq!(1, from.len());
match from[0].relation {
TableFactor::UNNEST {
with_ordinality: true,
..
} => {}
_ => panic!("Expecting TableFactor::UNNEST with ordinality"),
}
}

View file

@ -48,6 +48,7 @@ fn test_square_brackets_over_db_schema_table_name() {
with_hints: vec![],
version: None,
partitions: vec![],
with_ordinality: false,
},
joins: vec![],
}
@ -94,6 +95,7 @@ fn test_double_quotes_over_db_schema_table_name() {
with_hints: vec![],
version: None,
partitions: vec![],
with_ordinality: false,
},
joins: vec![],
}
@ -114,6 +116,7 @@ fn parse_delimited_identifiers() {
args,
with_hints,
version,
with_ordinality: _,
partitions: _,
} => {
assert_eq!(vec![Ident::with_quote('"', "a table")], name.0);

View file

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

View file

@ -399,7 +399,8 @@ fn parse_update_tuple_row_values() {
args: None,
with_hints: vec![],
version: None,
partitions: vec![]
partitions: vec![],
with_ordinality: false,
},
joins: vec![],
},