support column type definitions in table aliases (#1526)

This commit is contained in:
Ophir LOJKINE 2024-11-18 15:02:22 +01:00 committed by GitHub
parent 4c629e8520
commit f961efc0c9
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 105 additions and 14 deletions

View file

@ -60,8 +60,8 @@ pub use self::query::{
OrderBy, OrderByExpr, PivotValueSource, ProjectionSelect, Query, RenameSelectItem,
RepetitionQuantifier, ReplaceSelectElement, ReplaceSelectItem, RowsPerMatch, Select,
SelectInto, SelectItem, SetExpr, SetOperator, SetQuantifier, Setting, SymbolDefinition, Table,
TableAlias, TableFactor, TableFunctionArgs, TableVersion, TableWithJoins, Top, TopQuantity,
ValueTableMode, Values, WildcardAdditionalOptions, With, WithFill,
TableAlias, TableAliasColumnDef, TableFactor, TableFunctionArgs, TableVersion, TableWithJoins,
Top, TopQuantity, ValueTableMode, Values, WildcardAdditionalOptions, With, WithFill,
};
pub use self::trigger::{

View file

@ -1597,7 +1597,7 @@ impl fmt::Display for TableFactor {
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
pub struct TableAlias {
pub name: Ident,
pub columns: Vec<Ident>,
pub columns: Vec<TableAliasColumnDef>,
}
impl fmt::Display for TableAlias {
@ -1610,6 +1610,41 @@ impl fmt::Display for TableAlias {
}
}
/// SQL column definition in a table expression alias.
/// Most of the time, the data type is not specified.
/// But some table-valued functions do require specifying the data type.
///
/// See <https://www.postgresql.org/docs/17/queries-table-expressions.html#QUERIES-TABLEFUNCTIONS>
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
pub struct TableAliasColumnDef {
/// Column name alias
pub name: Ident,
/// Some table-valued functions require specifying the data type in the alias.
pub data_type: Option<DataType>,
}
impl TableAliasColumnDef {
/// Create a new table alias column definition with only a name and no type
pub fn from_name<S: Into<String>>(name: S) -> Self {
TableAliasColumnDef {
name: Ident::new(name),
data_type: None,
}
}
}
impl fmt::Display for TableAliasColumnDef {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
write!(f, "{}", self.name)?;
if let Some(ref data_type) = self.data_type {
write!(f, " {}", data_type)?;
}
Ok(())
}
}
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]

View file

@ -8270,7 +8270,7 @@ impl<'a> Parser<'a> {
) -> Result<Option<TableAlias>, ParserError> {
match self.parse_optional_alias(reserved_kwds)? {
Some(name) => {
let columns = self.parse_parenthesized_column_list(Optional, false)?;
let columns = self.parse_table_alias_column_defs()?;
Ok(Some(TableAlias { name, columns }))
}
None => Ok(None),
@ -8607,6 +8607,21 @@ impl<'a> Parser<'a> {
}
}
/// Parse a parenthesized comma-separated list of table alias column definitions.
fn parse_table_alias_column_defs(&mut self) -> Result<Vec<TableAliasColumnDef>, ParserError> {
if self.consume_token(&Token::LParen) {
let cols = self.parse_comma_separated(|p| {
let name = p.parse_identifier(false)?;
let data_type = p.maybe_parse(|p| p.parse_data_type())?;
Ok(TableAliasColumnDef { name, data_type })
})?;
self.expect_token(&Token::RParen)?;
Ok(cols)
} else {
Ok(vec![])
}
}
pub fn parse_precision(&mut self) -> Result<u64, ParserError> {
self.expect_token(&Token::LParen)?;
let n = self.parse_literal_uint()?;
@ -9174,7 +9189,7 @@ impl<'a> Parser<'a> {
materialized: is_materialized,
}
} else {
let columns = self.parse_parenthesized_column_list(Optional, false)?;
let columns = self.parse_table_alias_column_defs()?;
self.expect_keyword(Keyword::AS)?;
let mut is_materialized = None;
if dialect_of!(self is PostgreSqlDialect) {

View file

@ -553,7 +553,11 @@ fn parse_select_with_table_alias() {
name: ObjectName(vec![Ident::new("lineitem")]),
alias: Some(TableAlias {
name: Ident::new("l"),
columns: vec![Ident::new("A"), Ident::new("B"), Ident::new("C"),],
columns: vec![
TableAliasColumnDef::from_name("A"),
TableAliasColumnDef::from_name("B"),
TableAliasColumnDef::from_name("C"),
],
}),
args: None,
with_hints: vec![],
@ -5597,6 +5601,40 @@ fn parse_table_function() {
);
}
#[test]
fn parse_select_with_alias_and_column_defs() {
let sql = r#"SELECT * FROM jsonb_to_record('{"a": "x", "b": 2}'::JSONB) AS x (a TEXT, b INT)"#;
let select = verified_only_select(sql);
match only(&select.from) {
TableWithJoins {
relation: TableFactor::Table {
alias: Some(alias), ..
},
..
} => {
assert_eq!(alias.name.value, "x");
assert_eq!(
alias.columns,
vec![
TableAliasColumnDef {
name: Ident::new("a"),
data_type: Some(DataType::Text),
},
TableAliasColumnDef {
name: Ident::new("b"),
data_type: Some(DataType::Int(None)),
},
]
);
}
_ => unreachable!(
"Expecting only TableWithJoins with TableFactor::Table, got {:#?}",
select.from
),
}
}
#[test]
fn parse_unnest() {
let sql = "SELECT UNNEST(make_array(1, 2, 3))";
@ -6372,7 +6410,10 @@ fn parse_cte_renamed_columns() {
let sql = "WITH cte (col1, col2) AS (SELECT foo, bar FROM baz) SELECT * FROM cte";
let query = all_dialects().verified_query(sql);
assert_eq!(
vec![Ident::new("col1"), Ident::new("col2")],
vec![
TableAliasColumnDef::from_name("col1"),
TableAliasColumnDef::from_name("col2")
],
query
.with
.unwrap()
@ -6401,10 +6442,7 @@ fn parse_recursive_cte() {
value: "nums".to_string(),
quote_style: None,
},
columns: vec![Ident {
value: "val".to_string(),
quote_style: None,
}],
columns: vec![TableAliasColumnDef::from_name("val")],
},
query: Box::new(cte_query),
from: None,
@ -9347,7 +9385,10 @@ fn parse_pivot_table() {
value: "p".to_string(),
quote_style: None
},
columns: vec![Ident::new("c"), Ident::new("d")],
columns: vec![
TableAliasColumnDef::from_name("c"),
TableAliasColumnDef::from_name("d"),
],
}),
}
);
@ -9408,8 +9449,8 @@ fn parse_unpivot_table() {
name: Ident::new("u"),
columns: ["product", "quarter", "quantity"]
.into_iter()
.map(Ident::new)
.collect()
.map(TableAliasColumnDef::from_name)
.collect(),
}),
}
);