Postgresql: Add REPLICA IDENTITY operation for ALTER TABLE (#1844)

Author: MohamedAbdeen21

src/ast/ddl.rs

@@ -39,6 +39,29 @@ use crate::ast::{
 use crate::keywords::Keyword;
 use crate::tokenizer::Token;
 
+/// ALTER TABLE operation REPLICA IDENTITY values
+/// See [Postgres ALTER TABLE docs](https://www.postgresql.org/docs/current/sql-altertable.html)
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub enum ReplicaIdentity {
+    None,
+    Full,
+    Default,
+    Index(Ident),
+}
+
+impl fmt::Display for ReplicaIdentity {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match self {
+            ReplicaIdentity::None => f.write_str("NONE"),
+            ReplicaIdentity::Full => f.write_str("FULL"),
+            ReplicaIdentity::Default => f.write_str("DEFAULT"),
+            ReplicaIdentity::Index(idx) => write!(f, "USING INDEX {}", idx),
+        }
+    }
+}
+
 /// An `ALTER TABLE` (`Statement::AlterTable`) operation
 #[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
@@ -208,6 +231,13 @@ pub enum AlterTableOperation {
         old_partitions: Vec<Expr>,
         new_partitions: Vec<Expr>,
     },
+    /// REPLICA IDENTITY { DEFAULT | USING INDEX index_name | FULL | NOTHING }
+    ///
+    /// Note: this is a PostgreSQL-specific operation.
+    /// Please refer to [PostgreSQL documentation](https://www.postgresql.org/docs/current/sql-altertable.html)
+    ReplicaIdentity {
+        identity: ReplicaIdentity,
+    },
     /// Add Partitions
     AddPartitions {
         if_not_exists: bool,
@@ -729,6 +759,9 @@ impl fmt::Display for AlterTableOperation {
             AlterTableOperation::Lock { equals, lock } => {
                 write!(f, "LOCK {}{}", if *equals { "= " } else { "" }, lock)
             }
+            AlterTableOperation::ReplicaIdentity { identity } => {
+                write!(f, "REPLICA IDENTITY {identity}")
+            }
         }
     }
 }

src/ast/mod.rs

@@ -65,7 +65,7 @@ pub use self::ddl::{
     DeferrableInitial, DropBehavior, GeneratedAs, GeneratedExpressionMode, IdentityParameters,
     IdentityProperty, IdentityPropertyFormatKind, IdentityPropertyKind, IdentityPropertyOrder,
     IndexOption, IndexType, KeyOrIndexDisplay, NullsDistinctOption, Owner, Partition,
-    ProcedureParam, ReferentialAction, TableConstraint, TagsColumnOption,
+    ProcedureParam, ReferentialAction, ReplicaIdentity, TableConstraint, TagsColumnOption,
     UserDefinedTypeCompositeAttributeDef, UserDefinedTypeRepresentation, ViewColumnDef,
 };
 pub use self::dml::{CreateIndex, CreateTable, Delete, IndexColumn, Insert};

src/ast/spans.rs

@@ -1175,6 +1175,7 @@ impl Spanned for AlterTableOperation {
             AlterTableOperation::Algorithm { .. } => Span::empty(),
             AlterTableOperation::AutoIncrement { value, .. } => value.span(),
             AlterTableOperation::Lock { .. } => Span::empty(),
+            AlterTableOperation::ReplicaIdentity { .. } => Span::empty(),
         }
     }
 }

src/parser/mod.rs

@@ -8774,6 +8774,23 @@ impl<'a> Parser<'a> {
             let equals = self.consume_token(&Token::Eq);
             let value = self.parse_number_value()?;
             AlterTableOperation::AutoIncrement { equals, value }
+        } else if self.parse_keywords(&[Keyword::REPLICA, Keyword::IDENTITY]) {
+            let identity = if self.parse_keyword(Keyword::NONE) {
+                ReplicaIdentity::None
+            } else if self.parse_keyword(Keyword::FULL) {
+                ReplicaIdentity::Full
+            } else if self.parse_keyword(Keyword::DEFAULT) {
+                ReplicaIdentity::Default
+            } else if self.parse_keywords(&[Keyword::USING, Keyword::INDEX]) {
+                ReplicaIdentity::Index(self.parse_identifier()?)
+            } else {
+                return self.expected(
+                    "NONE, FULL, DEFAULT, or USING INDEX index_name after REPLICA IDENTITY",
+                    self.peek_token(),
+                );
+            };
+
+            AlterTableOperation::ReplicaIdentity { identity }
         } else {
             let options: Vec<SqlOption> =
                 self.parse_options_with_keywords(&[Keyword::SET, Keyword::TBLPROPERTIES])?;
@@ -8783,7 +8800,7 @@ impl<'a> Parser<'a> {
                 }
             } else {
                 return self.expected(
-                    "ADD, RENAME, PARTITION, SWAP, DROP, or SET TBLPROPERTIES after ALTER TABLE",
+                    "ADD, RENAME, PARTITION, SWAP, DROP, REPLICA IDENTITY, or SET TBLPROPERTIES after ALTER TABLE",
                     self.peek_token(),
                 );
             }

tests/sqlparser_postgres.rs

@@ -5959,3 +5959,30 @@ fn parse_varbit_datatype() {
         _ => unreachable!(),
     }
 }
+
+#[test]
+fn parse_alter_table_replica_identity() {
+    match pg_and_generic().verified_stmt("ALTER TABLE foo REPLICA IDENTITY FULL") {
+        Statement::AlterTable { operations, .. } => {
+            assert_eq!(
+                operations,
+                vec![AlterTableOperation::ReplicaIdentity {
+                    identity: ReplicaIdentity::Full
+                }]
+            );
+        }
+        _ => unreachable!(),
+    }
+
+    match pg_and_generic().verified_stmt("ALTER TABLE foo REPLICA IDENTITY USING INDEX foo_idx") {
+        Statement::AlterTable { operations, .. } => {
+            assert_eq!(
+                operations,
+                vec![AlterTableOperation::ReplicaIdentity {
+                    identity: ReplicaIdentity::Index("foo_idx".into())
+                }]
+            );
+        }
+        _ => unreachable!(),
+    }
+}