Skip to content

Commit

Permalink
Add support of table function WITH ORDINALITY modifier for Postgre Pa…
Browse files Browse the repository at this point in the history
…rser (#1337)
  • Loading branch information
git-hulk authored Jul 9, 2024
1 parent 285f492 commit 9108bff
Show file tree
Hide file tree
Showing 16 changed files with 129 additions and 4 deletions.
14 changes: 14 additions & 0 deletions src/ast/query.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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>,
},
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -1293,6 +1298,7 @@ impl fmt::Display for TableFactor {
with_hints,
version,
partitions,
with_ordinality,
} => {
write!(f, "{name}")?;
if !partitions.is_empty() {
Expand All @@ -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}")?;
}
Expand Down Expand Up @@ -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}")?;
}
Expand Down
1 change: 1 addition & 0 deletions src/keywords.rs
Original file line number Diff line number Diff line change
Expand Up @@ -518,6 +518,7 @@ define_keywords!(
OR,
ORC,
ORDER,
ORDINALITY,
OUT,
OUTER,
OUTPUTFORMAT,
Expand Down
5 changes: 5 additions & 0 deletions src/parser/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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()?;
Expand Down Expand Up @@ -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:
Expand All @@ -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]) {
Expand Down
2 changes: 2 additions & 0 deletions src/test_utils.rs
Original file line number Diff line number Diff line change
Expand Up @@ -309,6 +309,7 @@ pub fn table(name: impl Into<String>) -> TableFactor {
with_hints: vec![],
version: None,
partitions: vec![],
with_ordinality: false,
}
}

Expand All @@ -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,
}
}

Expand Down
8 changes: 7 additions & 1 deletion tests/sqlparser_bigquery.rs
Original file line number Diff line number Diff line change
Expand Up @@ -224,6 +224,7 @@ fn parse_delete_statement() {
with_hints: vec![],
version: None,
partitions: vec![],
with_ordinality: false,
},
from[0].relation
);
Expand Down Expand Up @@ -1353,6 +1354,7 @@ fn parse_table_identifiers() {
with_hints: vec![],
version: None,
partitions: vec![],
with_ordinality: false,
},
joins: vec![]
},]
Expand Down Expand Up @@ -1525,6 +1527,7 @@ fn parse_table_time_travel() {
Value::SingleQuotedString(version)
))),
partitions: vec![],
with_ordinality: false,
},
joins: vec![]
},]
Expand All @@ -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())),
Expand Down Expand Up @@ -1620,6 +1624,7 @@ fn parse_merge() {
with_hints: Default::default(),
version: Default::default(),
partitions: Default::default(),
with_ordinality: false,
},
table
);
Expand All @@ -1634,6 +1639,7 @@ fn parse_merge() {
with_hints: Default::default(),
version: Default::default(),
partitions: Default::default(),
with_ordinality: false,
},
source
);
Expand Down
2 changes: 2 additions & 0 deletions tests/sqlparser_clickhouse.rs
Original file line number Diff line number Diff line change
Expand Up @@ -59,6 +59,7 @@ fn parse_map_access_expr() {
with_hints: vec![],
version: None,
partitions: vec![],
with_ordinality: false,
},
joins: vec![],
}],
Expand Down Expand Up @@ -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);
Expand Down
Loading

0 comments on commit 9108bff

Please sign in to comment.