PostgreSQL: Support subscripts in assignment and insert targets

fuzz/fuzz_seeds/postgres_subscript_assignment_target.sql

@@ -0,0 +1 @@
+UPDATE u SET arr[1] = 5

src/ast/dml.rs

@@ -33,8 +33,8 @@ use super::{
     display_comma_separated, helpers::attached_token::AttachedToken, query::InputFormatClause,
     Assignment, Expr, FromTable, Ident, InsertAliases, MysqlInsertPriority, ObjectName, OnInsert,
     OptimizerHint, OrderByExpr, Query, SelectInto, SelectItem, Setting, SqliteOnConflict,
-    TableAliasWithoutColumns, TableFactor, TableObject, TableWithJoins, UpdateTableFromKind,
-    Values,
+    TableAliasWithoutColumns, TableFactor, TableObject, TableWithJoins, TargetColumn,
+    UpdateTableFromKind, Values,
 };
 
 /// The `OVERRIDING` clause of a PostgreSQL `INSERT`: whether values supplied for an identity
@@ -84,7 +84,7 @@ pub struct Insert {
     /// `table_name foo` (for Oracle)
     pub table_alias: Option<TableAliasWithoutColumns>,
     /// COLUMNS
-    pub columns: Vec<ObjectName>,
+    pub columns: Vec<TargetColumn>,
     /// `BY NAME` clause used by Databricks SQL.
     ///
     /// When present, columns from the source query are matched to columns in
@@ -784,7 +784,7 @@ pub struct MergeInsertExpr {
     /// INSERT (product, quantity) VALUES(product, quantity)
     /// INSERT (product, quantity) ROW
     /// ```
-    pub columns: Vec<ObjectName>,
+    pub columns: Vec<TargetColumn>,
     /// The token, `[VALUES | ROW]` starting `kind`.
     pub kind_token: AttachedToken,
     /// The insert type used by the statement.

src/ast/mod.rs

@@ -8112,9 +8112,9 @@ impl fmt::Display for Assignment {
 #[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
 pub enum AssignmentTarget {
     /// A single column
-    ColumnName(ObjectName),
+    ColumnName(TargetColumn),
     /// A tuple of columns
-    Tuple(Vec<ObjectName>),
+    Tuple(Vec<TargetColumn>),
 }
 
 impl fmt::Display for AssignmentTarget {
@@ -8126,6 +8126,46 @@ impl fmt::Display for AssignmentTarget {
     }
 }
 
+/// A column written by an assignment or listed in an `INSERT` column list,
+/// e.g. `col`, `arr[1]` or `arr[1:2]`.
+///
+/// [PostgreSQL](https://www.postgresql.org/docs/current/sql-update.html)
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub struct TargetColumn {
+    /// The column, e.g. `arr` in `arr[1].x`.
+    pub name: ObjectName,
+    /// Subscripts and the subfields after them, e.g. `[1]` and `.x` in `arr[1].x`.
+    pub access_chain: Vec<AccessExpr>,
+}
+
+impl From<ObjectName> for TargetColumn {
+    fn from(name: ObjectName) -> Self {
+        Self {
+            name,
+            access_chain: vec![],
+        }
+    }
+}
+
+impl From<Ident> for TargetColumn {
+    fn from(ident: Ident) -> Self {
+        ObjectName::from(ident).into()
+    }
+}
+
+impl fmt::Display for TargetColumn {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        write!(
+            f,
+            "{}{}",
+            self.name,
+            display_separated(&self.access_chain, "")
+        )
+    }
+}
+
 #[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
 #[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]

src/ast/spans.rs

@@ -49,8 +49,9 @@ use super::{
     RefreshMaterializedView, RenameSelectItem, ReplaceSelectElement, ReplaceSelectItem, Select,
     SelectInto, SelectItem, SetExpr, Setting, SqlOption, Statement, Subscript, SymbolDefinition,
     TableAlias, TableAliasColumnDef, TableConstraint, TableFactor, TableObject,
-    TableOptionsClustered, TableWithJoins, Update, UpdateTableFromKind, Use, UserMappingUser,
-    Values, ViewColumnDef, WhileStatement, WildcardAdditionalOptions, With, WithFill,
+    TableOptionsClustered, TableWithJoins, TargetColumn, Update, UpdateTableFromKind, Use,
+    UserMappingUser, Values, ViewColumnDef, WhileStatement, WildcardAdditionalOptions, With,
+    WithFill,
 };
 
 /// Given an iterator of spans, return the [Span::union] of all spans.
@@ -1620,12 +1621,20 @@ impl Spanned for Setting {
 impl Spanned for AssignmentTarget {
     fn span(&self) -> Span {
         match self {
-            AssignmentTarget::ColumnName(object_name) => object_name.span(),
+            AssignmentTarget::ColumnName(column) => column.span(),
             AssignmentTarget::Tuple(vec) => union_spans(vec.iter().map(|i| i.span())),
         }
     }
 }
 
+impl Spanned for TargetColumn {
+    fn span(&self) -> Span {
+        let TargetColumn { name, access_chain } = self;
+
+        union_spans(iter::once(name.span()).chain(access_chain.iter().map(|i| i.span())))
+    }
+}
+
 /// # partial span
 ///
 /// Most expressions are missing keywords in their spans.

src/parser/merge.rs

@@ -18,7 +18,7 @@ use alloc::{boxed::Box, format, vec, vec::Vec};
 use crate::{
     ast::{
         Merge, MergeAction, MergeClause, MergeClauseKind, MergeInsertExpr, MergeInsertKind,
-        MergeUpdateExpr, MergeUpdateKind, ObjectName, OutputClause, SetExpr,
+        MergeUpdateExpr, MergeUpdateKind, OutputClause, SetExpr, TargetColumn,
     },
     dialect::{BigQueryDialect, GenericDialect, MySqlDialect},
     keywords::Keyword,
@@ -239,8 +239,8 @@ impl Parser<'_> {
     fn parse_merge_clause_insert_columns(
         &mut self,
         allow_empty: bool,
-    ) -> Result<Vec<ObjectName>, ParserError> {
-        self.parse_parenthesized_qualified_column_list(IsOptional::Optional, allow_empty)
+    ) -> Result<Vec<TargetColumn>, ParserError> {
+        self.parse_parenthesized_target_column_list(IsOptional::Optional, allow_empty)
     }
 
     /// Parses an `OUTPUT` clause if present (MSSQL).

src/parser/mod.rs

@@ -15062,6 +15062,18 @@ impl<'a> Parser<'a> {
         })
     }
 
+    /// Parses a parenthesized comma-separated list of target columns.
+    /// For example: `(tbl.col1, arr[1], "col 2")`
+    fn parse_parenthesized_target_column_list(
+        &mut self,
+        optional: IsOptional,
+        allow_empty: bool,
+    ) -> Result<Vec<TargetColumn>, ParserError> {
+        self.parse_parenthesized_column_list_inner(optional, allow_empty, |p| {
+            p.parse_target_column(true)
+        })
+    }
+
     /// Parses a parenthesized comma-separated list of columns using
     /// the provided function to parse each element.
     fn parse_parenthesized_column_list_inner<F, T>(
@@ -19823,7 +19835,7 @@ impl<'a> Parser<'a> {
             } else {
                 let (columns, partitioned, after_columns) = if !self.peek_subquery_start() {
                     let columns =
-                        self.parse_parenthesized_qualified_column_list(Optional, is_mysql)?;
+                        self.parse_parenthesized_target_column_list(Optional, is_mysql)?;
 
                     let partitioned = self.parse_insert_partition()?;
                     by_name = self.parse_keywords(&[Keyword::BY, Keyword::NAME]);
@@ -20190,15 +20202,30 @@ impl<'a> Parser<'a> {
     /// Parse the left-hand side of an assignment, used in an UPDATE statement
     pub fn parse_assignment_target(&mut self) -> Result<AssignmentTarget, ParserError> {
         if self.consume_token(&Token::LParen) {
-            let columns = self.parse_comma_separated(|p| p.parse_object_name(false))?;
+            let columns = self.parse_comma_separated(|p| p.parse_target_column(false))?;
             self.expect_token(&Token::RParen)?;
             Ok(AssignmentTarget::Tuple(columns))
         } else {
-            let column = self.parse_object_name(false)?;
+            let column = self.parse_target_column(false)?;
             Ok(AssignmentTarget::ColumnName(column))
         }
     }
 
+    /// Parse a column written by an assignment or an `INSERT`, e.g. `col` or `arr[1].x`.
+    fn parse_target_column(&mut self, in_table_clause: bool) -> Result<TargetColumn, ParserError> {
+        let name = self.parse_object_name(in_table_clause)?;
+        let mut access_chain = vec![];
+        if !self.dialect.supports_partiql() {
+            self.parse_multi_dim_subscript(&mut access_chain)?;
+            while !access_chain.is_empty() && self.consume_token(&Token::Period) {
+                let field = self.parse_identifier()?;
+                access_chain.push(AccessExpr::Dot(Expr::Identifier(field)));
+                self.parse_multi_dim_subscript(&mut access_chain)?;
+            }
+        }
+        Ok(TargetColumn { name, access_chain })
+    }
+
     /// Parse a single function argument, handling named and unnamed variants.
     pub fn parse_function_args(&mut self) -> Result<FunctionArg, ParserError> {
         // Parse the argument expression once, then check for a named-arg

tests/sqlparser_bigquery.rs

@@ -1850,11 +1850,11 @@ fn parse_merge() {
         update_token: AttachedToken::empty(),
         kind: MergeUpdateKind::Set(vec![
             Assignment {
-                target: AssignmentTarget::ColumnName(ObjectName::from(vec![Ident::new("a")])),
+                target: AssignmentTarget::ColumnName(Ident::new("a").into()),
                 value: Expr::value(number("1")),
             },
             Assignment {
-                target: AssignmentTarget::ColumnName(ObjectName::from(vec![Ident::new("b")])),
+                target: AssignmentTarget::ColumnName(Ident::new("b").into()),
                 value: Expr::value(number("2")),
             },
         ]),

tests/sqlparser_common.rs

@@ -174,7 +174,7 @@ fn parse_insert_values() {
                 for (index, column) in columns.iter().enumerate() {
                     assert_eq!(
                         column,
-                        &ObjectName::from(Ident::new(expected_columns[index].clone()))
+                        &TargetColumn::from(Ident::new(expected_columns[index].clone()))
                     );
                 }
                 match *source.body {
@@ -470,15 +470,15 @@ fn parse_update() {
                 assignments,
                 vec![
                     Assignment {
-                        target: AssignmentTarget::ColumnName(ObjectName::from(vec!["a".into()])),
+                        target: AssignmentTarget::ColumnName(Ident::new("a").into()),
                         value: Expr::value(number("1")),
                     },
                     Assignment {
-                        target: AssignmentTarget::ColumnName(ObjectName::from(vec!["b".into()])),
+                        target: AssignmentTarget::ColumnName(Ident::new("b").into()),
                         value: Expr::value(number("2")),
                     },
                     Assignment {
-                        target: AssignmentTarget::ColumnName(ObjectName::from(vec!["c".into()])),
+                        target: AssignmentTarget::ColumnName(Ident::new("c").into()),
                         value: Expr::value(number("3")),
                     },
                 ]
@@ -530,7 +530,7 @@ fn parse_update_set_from() {
             },
             on_cluster: None,
             assignments: vec![Assignment {
-                target: AssignmentTarget::ColumnName(ObjectName::from(vec![Ident::new("name")])),
+                target: AssignmentTarget::ColumnName(Ident::new("name").into()),
                 value: Expr::CompoundIdentifier(vec![Ident::new("t2"), Ident::new("name")])
             }],
             from: Some(UpdateTableFromKind::AfterSet(vec![TableWithJoins {
@@ -651,10 +651,9 @@ fn parse_update_with_table_alias() {
             );
             assert_eq!(
                 vec![Assignment {
-                    target: AssignmentTarget::ColumnName(ObjectName::from(vec![
-                        Ident::new("u"),
-                        Ident::new("username")
-                    ])),
+                    target: AssignmentTarget::ColumnName(
+                        ObjectName::from(vec![Ident::new("u"), Ident::new("username")]).into()
+                    ),
                     value: Expr::Value(
                         (Value::SingleQuotedString("new_user".to_string())).with_empty_span()
                     ),
@@ -11327,20 +11326,20 @@ fn parse_merge() {
                             update_token: AttachedToken::empty(),
                             kind: MergeUpdateKind::Set(vec![
                                 Assignment {
-                                    target: AssignmentTarget::ColumnName(ObjectName::from(vec![
-                                        Ident::new("dest"),
-                                        Ident::new("F")
-                                    ])),
+                                    target: AssignmentTarget::ColumnName(
+                                        ObjectName::from(vec![Ident::new("dest"), Ident::new("F")])
+                                            .into()
+                                    ),
                                     value: Expr::CompoundIdentifier(vec![
                                         Ident::new("stg"),
                                         Ident::new("F"),
                                     ]),
                                 },
                                 Assignment {
-                                    target: AssignmentTarget::ColumnName(ObjectName::from(vec![
-                                        Ident::new("dest"),
-                                        Ident::new("G")
-                                    ])),
+                                    target: AssignmentTarget::ColumnName(
+                                        ObjectName::from(vec![Ident::new("dest"), Ident::new("G")])
+                                            .into()
+                                    ),
                                     value: Expr::CompoundIdentifier(vec![
                                         Ident::new("stg"),
                                         Ident::new("G"),
@@ -23516,3 +23515,46 @@ fn parse_national_character_string_types() {
         ParserError::ParserError("Expected: ), found: EOF".to_string())
     );
 }
+
+#[test]
+fn parse_subscripted_assignment_targets() {
+    let dialects = all_dialects_except(|d| {
+        d.is::<MsSqlDialect>()
+            || d.is::<SnowflakeDialect>()
+            || d.is::<SQLiteDialect>()
+            || d.is::<RedshiftSqlDialect>()
+    });
+    for sql in [
+        "UPDATE u SET arr[1] = 5, arr[2:3] = '{6,7}', arr[:1] = '{1}', m[1][2] = 4",
+        "UPDATE u SET parr[1].x = 9, p.y = 3",
+        "UPDATE u SET (arr[1], p.x) = (11, 12)",
+        "INSERT INTO u (id, arr[1], arr[2:3], p.x, parr[1].y) VALUES (1, 10, '{20,30}', 7, 8)",
+        "INSERT INTO u (id) VALUES (1) ON CONFLICT(id) DO UPDATE SET arr[1] = 99",
+        "MERGE INTO u USING v ON u.id = v.id WHEN MATCHED THEN UPDATE SET arr[1] = 5 WHEN NOT MATCHED THEN INSERT (id, arr[2]) VALUES (v.id, 6)",
+    ] {
+        dialects.verified_stmt(sql);
+    }
+
+    match dialects.verified_stmt("UPDATE u SET parr[1].x = 9") {
+        Statement::Update(Update { assignments, .. }) => assert_eq!(
+            assignments[0].target,
+            AssignmentTarget::ColumnName(TargetColumn {
+                name: ObjectName::from(vec![Ident::new("parr")]),
+                access_chain: vec![
+                    AccessExpr::Subscript(Subscript::Index {
+                        index: Expr::value(number("1")),
+                    }),
+                    AccessExpr::Dot(Expr::Identifier(Ident::new("x"))),
+                ],
+            })
+        ),
+        stmt => panic!("unexpected statement: {stmt:?}"),
+    }
+
+    assert_eq!(
+        dialects
+            .parse_sql_statements("UPDATE u SET arr[1]. = 5")
+            .unwrap_err(),
+        ParserError::ParserError("Expected: identifier, found: =".to_string())
+    );
+}