Skip to content

Commit a3cb42a

Browse files
git-hulkgit-hulk
committed
Fix review comments
1 parent b58cb88 commit a3cb42a

File tree

5 files changed

+65
-46
lines changed

5 files changed

+65
-46
lines changed

‎src/ast/query.rs‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -944,6 +944,7 @@ pub enum TableFactor {
944944
array_exprs: Vec<Expr>,
945945
with_offset: bool,
946946
with_offset_alias: Option<Ident>,
947+
with_ordinality: bool,
947948
},
948949
/// The `JSON_TABLE` table-valued function.
949950
/// Part of the SQL standard, but implemented only by MySQL, Oracle, and DB2.
@@ -1354,9 +1355,14 @@ impl fmt::Display for TableFactor {
13541355
array_exprs,
13551356
with_offset,
13561357
with_offset_alias,
1358+
with_ordinality,
13571359
} => {
13581360
write!(f, "UNNEST({})", display_comma_separated(array_exprs))?;
13591361

1362+
if *with_ordinality {
1363+
write!(f, " WITH ORDINALITY")?;
1364+
}
1365+
13601366
if let Some(alias) = alias {
13611367
write!(f, " AS {alias}")?;
13621368
}

‎src/parser/mod.rs‎

Lines changed: 22 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -7331,6 +7331,25 @@ impl<'a> Parser<'a> {
73317331
}
73327332
}
73337333

7334+
/// Parse `WITH ORDINALITY` for the table function in the `FROM` clause,
7335+
/// like in `... FROM generate_series(1, 10) WITH ORDINALITY`. This is an
7336+
/// optional clause that can be used to include a column of row numbers in
7337+
/// the output.
7338+
///
7339+
/// [PostgreSQL](https://www.postgresql.org/docs/current/functions-srf.html)
7340+
pub fn parse_optional_with_ordinality(&mut self) -> bool {
7341+
if !self.parse_keyword(Keyword::WITH) {
7342+
return false;
7343+
}
7344+
if self.parse_keyword(Keyword::ORDINALITY) {
7345+
true
7346+
} else {
7347+
// Rewind, the WITH keyword might belong to the cte or hint clause
7348+
self.prev_token();
7349+
false
7350+
}
7351+
}
7352+
73347353
/// Parse a possibly qualified, possibly quoted identifier, e.g.
73357354
/// `foo` or `myschema."table"
73367355
///
@@ -9152,6 +9171,7 @@ impl<'a> Parser<'a> {
91529171
let array_exprs = self.parse_comma_separated(Parser::parse_expr)?;
91539172
self.expect_token(&Token::RParen)?;
91549173

9174+
let with_ordinality = self.parse_optional_with_ordinality();
91559175
let alias = match self.parse_optional_table_alias(keywords::RESERVED_FOR_TABLE_ALIAS) {
91569176
Ok(Some(alias)) => Some(alias),
91579177
Ok(None) => None,
@@ -9178,6 +9198,7 @@ impl<'a> Parser<'a> {
91789198
array_exprs,
91799199
with_offset,
91809200
with_offset_alias,
9201+
with_ordinality,
91819202
})
91829203
} else if self.parse_keyword_with_tokens(Keyword::JSON_TABLE, &[Token::LParen]) {
91839204
let json_expr = self.parse_expr()?;
@@ -9216,17 +9237,7 @@ impl<'a> Parser<'a> {
92169237
None
92179238
};
92189239

9219-
let mut with_ordinality = false;
9220-
if dialect_of!(self is PostgreSqlDialect|GenericDialect)
9221-
&& self.parse_keyword(Keyword::WITH)
9222-
{
9223-
if self.parse_keyword(Keyword::ORDINALITY) {
9224-
with_ordinality = true;
9225-
} else {
9226-
// rewind, as WITH may belong to the next statement's CTE or hints
9227-
self.prev_token();
9228-
}
9229-
}
9240+
let with_ordinality = self.parse_optional_with_ordinality();
92309241

92319242
let alias = self.parse_optional_table_alias(keywords::RESERVED_FOR_TABLE_ALIAS)?;
92329243

‎tests/sqlparser_bigquery.rs‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1554,7 +1554,8 @@ fn parse_join_constraint_unnest_alias() {
15541554
Ident::new("a")
15551555
])],
15561556
with_offset: false,
1557-
with_offset_alias: None
1557+
with_offset_alias: None,
1558+
with_ordinality: false,
15581559
},
15591560
join_operator: JoinOperator::Inner(JoinConstraint::On(Expr::BinaryOp {
15601561
left: Box::new(Expr::Identifier("c1".into())),

‎tests/sqlparser_common.rs‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5298,6 +5298,7 @@ fn parse_unnest_in_from_clause() {
52985298
array_exprs: vec![Expr::Identifier(Ident::new("expr"))],
52995299
with_offset: true,
53005300
with_offset_alias: None,
5301+
with_ordinality: false,
53015302
},
53025303
joins: vec![],
53035304
}],
@@ -5315,6 +5316,7 @@ fn parse_unnest_in_from_clause() {
53155316
array_exprs: vec![Expr::Identifier(Ident::new("expr"))],
53165317
with_offset: false,
53175318
with_offset_alias: None,
5319+
with_ordinality: false,
53185320
},
53195321
joins: vec![],
53205322
}],
@@ -5332,6 +5334,7 @@ fn parse_unnest_in_from_clause() {
53325334
array_exprs: vec![Expr::Identifier(Ident::new("expr"))],
53335335
with_offset: true,
53345336
with_offset_alias: None,
5337+
with_ordinality: false,
53355338
},
53365339
joins: vec![],
53375340
}],
@@ -5352,6 +5355,7 @@ fn parse_unnest_in_from_clause() {
53525355
array_exprs: vec![Expr::Identifier(Ident::new("expr"))],
53535356
with_offset: false,
53545357
with_offset_alias: None,
5358+
with_ordinality: false,
53555359
},
53565360
joins: vec![],
53575361
}],
@@ -5376,6 +5380,7 @@ fn parse_unnest_in_from_clause() {
53765380
)],
53775381
with_offset: false,
53785382
with_offset_alias: None,
5383+
with_ordinality: false,
53795384
},
53805385
joins: vec![],
53815386
}],
@@ -5406,6 +5411,7 @@ fn parse_unnest_in_from_clause() {
54065411
],
54075412
with_offset: false,
54085413
with_offset_alias: None,
5414+
with_ordinality: false,
54095415
},
54105416
joins: vec![],
54115417
}],

‎tests/sqlparser_postgres.rs‎

Lines changed: 29 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -18,8 +18,6 @@
1818
mod test_utils;
1919
use test_utils::*;
2020

21-
use sqlparser::ast::FunctionArg::Unnamed;
22-
use sqlparser::ast::Value::Number;
2321
use sqlparser::ast::*;
2422
use sqlparser::dialect::{GenericDialect, PostgreSqlDialect};
2523
use sqlparser::parser::ParserError;
@@ -3881,7 +3879,8 @@ fn parse_join_constraint_unnest_alias() {
38813879
Ident::new("a")
38823880
])],
38833881
with_offset: false,
3884-
with_offset_alias: None
3882+
with_offset_alias: None,
3883+
with_ordinality: false,
38853884
},
38863885
join_operator: JoinOperator::Inner(JoinConstraint::On(Expr::BinaryOp {
38873886
left: Box::new(Expr::Identifier("c1".into())),
@@ -4189,37 +4188,33 @@ fn parse_create_table_with_options() {
41894188

41904189
#[test]
41914190
fn test_table_function_with_ordinality() {
4192-
let sql = "SELECT * FROM generate_series(1, 10) WITH ORDINALITY AS t";
4193-
match pg_and_generic().verified_stmt(sql) {
4194-
Statement::Query(query) => {
4195-
assert_eq!(
4196-
query.body.as_select().unwrap().from,
4197-
vec![TableWithJoins {
4198-
relation: TableFactor::Table {
4199-
name: ObjectName(vec![Ident::new("generate_series")]),
4200-
args: Some(vec![
4201-
Unnamed(FunctionArgExpr::Expr(Expr::Value(Number(
4202-
"1".parse().unwrap(),
4203-
false
4204-
)))),
4205-
Unnamed(FunctionArgExpr::Expr(Expr::Value(Number(
4206-
"10".parse().unwrap(),
4207-
false
4208-
)))),
4209-
]),
4210-
alias: Some(TableAlias {
4211-
name: Ident::new("t"),
4212-
columns: vec![],
4213-
}),
4214-
with_hints: vec![],
4215-
version: None,
4216-
partitions: vec![],
4217-
with_ordinality: true,
4218-
},
4219-
joins: vec![],
4220-
}]
4221-
);
4191+
let from = pg_and_generic()
4192+
.verified_only_select("SELECT * FROM generate_series(1, 10) WITH ORDINALITY AS t")
4193+
.from;
4194+
assert_eq!(1, from.len());
4195+
match from[0].relation {
4196+
TableFactor::Table {
4197+
ref name,
4198+
with_ordinality: true,
4199+
..
4200+
} => {
4201+
assert_eq!("generate_series", name.to_string().as_str());
42224202
}
4223-
_ => unreachable!(),
4203+
_ => panic!("Expecting TableFactor::Table with ordinality"),
4204+
}
4205+
}
4206+
4207+
#[test]
4208+
fn test_table_unnest_with_ordinality() {
4209+
let from = pg_and_generic()
4210+
.verified_only_select("SELECT * FROM UNNEST([10, 20, 30]) WITH ORDINALITY AS t")
4211+
.from;
4212+
assert_eq!(1, from.len());
4213+
match from[0].relation {
4214+
TableFactor::UNNEST {
4215+
with_ordinality: true,
4216+
..
4217+
} => {}
4218+
_ => panic!("Expecting TableFactor::UNNEST with ordinality"),
42244219
}
42254220
}

0 commit comments

Comments
 (0)