support column type definitions in table aliases (apache#1526)

lovasoa · web-flow · commit f961efc0c92c · 2024-11-18T09:02:22.000-05:00
diff --git a/src/ast/mod.rs b/src/ast/mod.rs
@@ -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::{
diff --git a/src/ast/query.rs b/src/ast/query.rs
@@ -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))]
diff --git a/src/parser/mod.rs b/src/parser/mod.rs
@@ -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) {
diff --git a/tests/sqlparser_common.rs b/tests/sqlparser_common.rs
@@ -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(),
             }),
         }
     );
