From 7a34cd517639300eb906f72ee1225d90660aa933 Mon Sep 17 00:00:00 2001
From: Toby Hede <toby@cipherstash.com>
Date: Mon, 1 Jan 2024 14:54:26 +1100
Subject: [PATCH] Support for pg ADD GENERATED in ALTER COLUMN statements

Signed-off-by: Toby Hede <toby@cipherstash.com>
---
 src/ast/ddl.rs              | 33 ++++++++++++++++++++++++++++++
 src/keywords.rs             |  1 +
 src/parser/mod.rs           | 40 ++++++++++++++++++++++++++++++++-----
 tests/sqlparser_postgres.rs | 14 +++++++++++++
 4 files changed, 83 insertions(+), 5 deletions(-)

diff --git a/src/ast/ddl.rs b/src/ast/ddl.rs
index cd0cf1797c..03fc5aabd6 100644
--- a/src/ast/ddl.rs
+++ b/src/ast/ddl.rs
@@ -315,6 +315,13 @@ pub enum AlterColumnOperation {
         /// PostgreSQL specific
         using: Option<Expr>,
     },
+    /// `ADD GENERATED { ALWAYS | BY DEFAULT } AS IDENTITY [ ( sequence_options ) ]`
+    ///
+    /// Note: this is a PostgreSQL-specific operation.
+    AddGenerated {
+        generated_as: Option<GeneratedAs>,
+        sequence_options: Option<Vec<SequenceOptions>>,
+    },
 }
 
 impl fmt::Display for AlterColumnOperation {
@@ -335,6 +342,32 @@ impl fmt::Display for AlterColumnOperation {
                     write!(f, "SET DATA TYPE {data_type}")
                 }
             }
+            AlterColumnOperation::AddGenerated {
+                generated_as,
+                sequence_options,
+            } => {
+                let generated_as = match generated_as {
+                    Some(GeneratedAs::Always) => " ALWAYS",
+                    Some(GeneratedAs::ByDefault) => " BY DEFAULT",
+                    _ => "",
+                };
+
+                write!(f, "ADD GENERATED{generated_as} AS IDENTITY",)?;
+                if let Some(options) = sequence_options {
+                    if !options.is_empty() {
+                        write!(f, " (")?;
+                    }
+
+                    for sequence_option in options {
+                        write!(f, "{sequence_option}")?;
+                    }
+
+                    if !options.is_empty() {
+                        write!(f, " )")?;
+                    }
+                }
+                Ok(())
+            }
         }
     }
 }
diff --git a/src/keywords.rs b/src/keywords.rs
index cba724019a..e136f9905e 100644
--- a/src/keywords.rs
+++ b/src/keywords.rs
@@ -552,6 +552,7 @@ define_keywords!(
     REPLICATION,
     RESET,
     RESPECT,
+    RESTART,
     RESTRICT,
     RESULT,
     RETAIN,
diff --git a/src/parser/mod.rs b/src/parser/mod.rs
index 92f4cf0285..3221e2a44f 100644
--- a/src/parser/mod.rs
+++ b/src/parser/mod.rs
@@ -4841,7 +4841,11 @@ impl<'a> Parser<'a> {
             let column_name = self.parse_identifier()?;
             let is_postgresql = dialect_of!(self is PostgreSqlDialect);
 
-            let op = if self.parse_keywords(&[Keyword::SET, Keyword::NOT, Keyword::NULL]) {
+            let op: AlterColumnOperation = if self.parse_keywords(&[
+                Keyword::SET,
+                Keyword::NOT,
+                Keyword::NULL,
+            ]) {
                 AlterColumnOperation::SetNotNull {}
             } else if self.parse_keywords(&[Keyword::DROP, Keyword::NOT, Keyword::NULL]) {
                 AlterColumnOperation::DropNotNull {}
@@ -4861,11 +4865,37 @@ impl<'a> Parser<'a> {
                     None
                 };
                 AlterColumnOperation::SetDataType { data_type, using }
+            } else if self.parse_keywords(&[Keyword::ADD, Keyword::GENERATED]) {
+                let generated_as = if self.parse_keyword(Keyword::ALWAYS) {
+                    Some(GeneratedAs::Always)
+                } else if self.parse_keywords(&[Keyword::BY, Keyword::DEFAULT]) {
+                    Some(GeneratedAs::ByDefault)
+                } else {
+                    None
+                };
+
+                let _ = self.expect_keywords(&[Keyword::AS, Keyword::IDENTITY]);
+
+                let mut sequence_options: Option<Vec<SequenceOptions>> = None;
+
+                if self.peek_token().token == Token::LParen {
+                    self.expect_token(&Token::LParen)?;
+                    sequence_options = self.parse_create_sequence_options().ok();
+                    self.expect_token(&Token::RParen)?;
+                }
+
+                AlterColumnOperation::AddGenerated {
+                    generated_as,
+                    sequence_options,
+                }
             } else {
-                return self.expected(
-                    "SET/DROP NOT NULL, SET DEFAULT, SET DATA TYPE after ALTER COLUMN",
-                    self.peek_token(),
-                );
+                let message = if is_postgresql {
+                    "SET/DROP NOT NULL, SET DEFAULT, SET DATA TYPE, or ADD GENERATED after ALTER COLUMN"
+                } else {
+                    "SET/DROP NOT NULL, SET DEFAULT, SET DATA TYPE, after ALTER COLUMN"
+                };
+
+                return self.expected(message, self.peek_token());
             };
             AlterTableOperation::AlterColumn { column_name, op }
         } else if self.parse_keyword(Keyword::SWAP) {
diff --git a/tests/sqlparser_postgres.rs b/tests/sqlparser_postgres.rs
index 0c1463f92c..da3f060837 100644
--- a/tests/sqlparser_postgres.rs
+++ b/tests/sqlparser_postgres.rs
@@ -611,6 +611,20 @@ fn parse_alter_table_alter_column() {
     }
 }
 
+#[test]
+fn parse_alter_table_alter_column_add_generated() {
+    pg_and_generic()
+        .verified_stmt("ALTER TABLE t ALTER COLUMN id ADD GENERATED ALWAYS AS IDENTITY");
+    pg_and_generic()
+        .verified_stmt("ALTER TABLE t ALTER COLUMN id ADD GENERATED BY DEFAULT AS IDENTITY");
+    pg_and_generic().verified_stmt("ALTER TABLE t ALTER COLUMN id ADD GENERATED AS IDENTITY");
+    pg_and_generic().verified_stmt(
+        "ALTER TABLE t ALTER COLUMN id ADD GENERATED AS IDENTITY ( INCREMENT 1 MINVALUE 1 )",
+    );
+
+    pg_and_generic().verified_stmt("ALTER TABLE t ALTER COLUMN id ADD GENERATED AS IDENTITY ( )");
+}
+
 #[test]
 fn parse_alter_table_add_columns() {
     match pg().verified_stmt("ALTER TABLE IF EXISTS ONLY tab ADD COLUMN a TEXT, ADD COLUMN b INT") {