Redshift: PartiQL AT <index> (#2303)

Author: yoavcloud

src/ast/query.rs

@@ -2533,6 +2533,12 @@ pub struct TableAlias {
     pub name: Ident,
     /// Optional column aliases declared in parentheses after the table alias.
     pub columns: Vec<TableAliasColumnDef>,
+    /// Optional PartiQL index alias declared with `AT`. For example:
+    /// ```sql
+    /// SELECT element, index FROM bar AS b, b.data.scalar_array AS element AT index
+    /// ```
+    /// See: <https://docs.aws.amazon.com/redshift/latest/dg/query-super.html>
+    pub at: Option<Ident>,
 }
 
 impl fmt::Display for TableAlias {
@@ -2541,6 +2547,9 @@ impl fmt::Display for TableAlias {
         if !self.columns.is_empty() {
             write!(f, " ({})", display_comma_separated(&self.columns))?;
         }
+        if let Some(at) = &self.at {
+            write!(f, " AT {at}")?;
+        }
         Ok(())
     }
 }

src/ast/spans.rs

@@ -2186,8 +2186,13 @@ impl Spanned for TableAlias {
             explicit: _,
             name,
             columns,
+            at,
         } = self;
-        union_spans(core::iter::once(name.span).chain(columns.iter().map(Spanned::span)))
+        union_spans(
+            core::iter::once(name.span)
+                .chain(columns.iter().map(Spanned::span))
+                .chain(at.iter().map(|at| at.span)),
+        )
     }
 }
 

src/parser/mod.rs

@@ -12822,10 +12822,16 @@ impl<'a> Parser<'a> {
         match self.parse_optional_alias_inner(None, validator)? {
             Some(name) => {
                 let columns = self.parse_table_alias_column_defs()?;
+                let at = if self.dialect.supports_partiql() && self.parse_keyword(Keyword::AT) {
+                    Some(self.parse_identifier()?)
+                } else {
+                    None
+                };
                 Ok(Some(TableAlias {
                     explicit,
                     name,
                     columns,
+                    at,
                 }))
             }
             None => Ok(None),
@@ -14474,6 +14480,7 @@ impl<'a> Parser<'a> {
                         explicit: false,
                         name,
                         columns: vec![],
+                        at: None,
                     },
                     query,
                     from: None,
@@ -14518,6 +14525,7 @@ impl<'a> Parser<'a> {
                 explicit: false,
                 name,
                 columns,
+                at: None,
             },
             query,
             from: None,

src/test_utils.rs

@@ -374,6 +374,7 @@ pub fn table_alias(explicit: bool, name: impl Into<String>) -> Option<TableAlias
         explicit,
         name: Ident::new(name),
         columns: vec![],
+        at: None,
     })
 }
 

tests/sqlparser_common.rs

@@ -669,6 +669,7 @@ fn parse_select_with_table_alias() {
                         TableAliasColumnDef::from_name("B"),
                         TableAliasColumnDef::from_name("C"),
                     ],
+                    at: None,
                 }),
                 args: None,
                 with_hints: vec![],
@@ -7861,6 +7862,7 @@ fn parse_recursive_cte() {
                 span: Span::empty(),
             },
             columns: vec![TableAliasColumnDef::from_name("val")],
+            at: None,
         },
         query: Box::new(cte_query),
         from: None,
@@ -11347,6 +11349,7 @@ fn parse_pivot_table() {
                     TableAliasColumnDef::from_name("c"),
                     TableAliasColumnDef::from_name("d"),
                 ],
+                at: None,
             })
         }
     );
@@ -11485,6 +11488,7 @@ fn parse_unpivot_table() {
                 .into_iter()
                 .map(TableAliasColumnDef::from_name)
                 .collect(),
+            at: None,
         }),
     };
     pretty_assertions::assert_eq!(verified_only_select(sql).from[0].relation, base_unpivot);

tests/sqlparser_mssql.rs

@@ -106,7 +106,8 @@ fn parse_mssql_delimited_identifiers() {
                         &Some(TableAlias {
                             explicit: false,
                             name: Ident::with_quote('[', "WHERE"),
-                            columns: vec![]
+                            columns: vec![],
+                            at: None,
                         })
                     );
                 }

tests/sqlparser_redshift.rs

@@ -517,3 +517,28 @@ fn test_null_treatment_inside_and_outside_window_function() {
     redshift().verified_stmt("SELECT FIRST_VALUE(1 IGNORE NULLS) OVER () FROM (SELECT 1) t");
     redshift().verified_stmt("SELECT FIRST_VALUE(1) IGNORE NULLS OVER () FROM (SELECT 1) t");
 }
+
+#[test]
+fn test_partiql_from_alias_with_at_index() {
+    let dialects = all_dialects_where(|d| d.supports_partiql());
+    dialects.verified_stmt("SELECT * FROM lineitem AS l (a, b, c) AT idx");
+
+    let sql =
+        "SELECT index, val FROM (SELECT array('AAA', 'BBB') AS val) AS b, b.val AS val AT index";
+    let select = dialects.verified_only_select(sql);
+
+    match &select.from[1].relation {
+        TableFactor::Table { name, alias, .. } => {
+            assert_eq!(
+                name,
+                &ObjectName::from(vec![Ident::new("b"), Ident::new("val")])
+            );
+            assert_eq!(alias.as_ref().map(|a| &a.name), Some(&Ident::new("val")));
+            assert_eq!(
+                alias.as_ref().and_then(|a| a.at.as_ref()),
+                Some(&Ident::new("index"))
+            );
+        }
+        _ => panic!("expected table factor"),
+    }
+}