From 9b5709d473444f4a9cae5c768bdafbd204ff32f5 Mon Sep 17 00:00:00 2001
From: poonai <rbalajis25@gmail.com>
Date: Sat, 12 Mar 2022 14:50:00 +0530
Subject: [PATCH] add support for SAVEPOINT statement

---
 src/ast/mod.rs              | 34 ++++++++++++++++++++++++++++------
 src/parser.rs               |  6 ++++++
 tests/sqlparser_postgres.rs | 10 ++++++++++
 3 files changed, 44 insertions(+), 6 deletions(-)

diff --git a/src/ast/mod.rs b/src/ast/mod.rs
index 232a506f4..fc233271b 100644
--- a/src/ast/mod.rs
+++ b/src/ast/mod.rs
@@ -851,7 +851,9 @@ pub enum Statement {
     /// SHOW <variable>
     ///
     /// Note: this is a PostgreSQL-specific statement.
-    ShowVariable { variable: Vec<Ident> },
+    ShowVariable {
+        variable: Vec<Ident>,
+    },
     /// SHOW CREATE TABLE
     ///
     /// Note: this is a MySQL-specific statement.
@@ -869,7 +871,9 @@ pub enum Statement {
         filter: Option<ShowStatementFilter>,
     },
     /// `{ BEGIN [ TRANSACTION | WORK ] | START TRANSACTION } ...`
-    StartTransaction { modes: Vec<TransactionMode> },
+    StartTransaction {
+        modes: Vec<TransactionMode>,
+    },
     /// `SET TRANSACTION ...`
     SetTransaction {
         modes: Vec<TransactionMode>,
@@ -885,9 +889,13 @@ pub enum Statement {
         comment: Option<String>,
     },
     /// `COMMIT [ TRANSACTION | WORK ] [ AND [ NO ] CHAIN ]`
-    Commit { chain: bool },
+    Commit {
+        chain: bool,
+    },
     /// `ROLLBACK [ TRANSACTION | WORK ] [ AND [ NO ] CHAIN ]`
-    Rollback { chain: bool },
+    Rollback {
+        chain: bool,
+    },
     /// CREATE SCHEMA
     CreateSchema {
         schema_name: ObjectName,
@@ -924,11 +932,17 @@ pub enum Statement {
     /// `DEALLOCATE [ PREPARE ] { name | ALL }`
     ///
     /// Note: this is a PostgreSQL-specific statement.
-    Deallocate { name: Ident, prepare: bool },
+    Deallocate {
+        name: Ident,
+        prepare: bool,
+    },
     /// `EXECUTE name [ ( parameter [, ...] ) ]`
     ///
     /// Note: this is a PostgreSQL-specific statement.
-    Execute { name: Ident, parameters: Vec<Expr> },
+    Execute {
+        name: Ident,
+        parameters: Vec<Expr>,
+    },
     /// `PREPARE name [ ( data_type [, ...] ) ] AS statement`
     ///
     /// Note: this is a PostgreSQL-specific statement.
@@ -956,6 +970,10 @@ pub enum Statement {
         /// A SQL query that specifies what to explain
         statement: Box<Statement>,
     },
+    /// SAVEPOINT -- define a new savepoint within the current transaction
+    Savepoint {
+        name: Ident,
+    },
 }
 
 impl fmt::Display for Statement {
@@ -1606,6 +1624,10 @@ impl fmt::Display for Statement {
                     write!(f, "NULL")
                 }
             }
+            Statement::Savepoint { name } => {
+                write!(f, "SAVEPOINT ")?;
+                write!(f, "{}", name)
+            }
         }
     }
 }
diff --git a/src/parser.rs b/src/parser.rs
index 6ccd4980b..401fe01f0 100644
--- a/src/parser.rs
+++ b/src/parser.rs
@@ -179,6 +179,7 @@ impl<'a> Parser<'a> {
                 // standard `START TRANSACTION` statement. It is supported
                 // by at least PostgreSQL and MySQL.
                 Keyword::BEGIN => Ok(self.parse_begin()?),
+                Keyword::SAVEPOINT => Ok(self.parse_savepoint()?),
                 Keyword::COMMIT => Ok(self.parse_commit()?),
                 Keyword::ROLLBACK => Ok(self.parse_rollback()?),
                 Keyword::ASSERT => Ok(self.parse_assert()?),
@@ -366,6 +367,11 @@ impl<'a> Parser<'a> {
         Ok(Statement::Assert { condition, message })
     }
 
+    pub fn parse_savepoint(&mut self) -> Result<Statement, ParserError> {
+        let name = self.parse_identifier()?;
+        Ok(Statement::Savepoint { name })
+    }
+
     /// Parse an expression prefix
     pub fn parse_prefix(&mut self) -> Result<Expr, ParserError> {
         // PostgreSQL allows any string literal to be preceded by a type name, indicating that the
diff --git a/tests/sqlparser_postgres.rs b/tests/sqlparser_postgres.rs
index 750b49869..b0bda1438 100644
--- a/tests/sqlparser_postgres.rs
+++ b/tests/sqlparser_postgres.rs
@@ -871,6 +871,16 @@ fn test_transaction_statement() {
     );
 }
 
+#[test]
+fn test_savepoint() {
+    match pg().verified_stmt("SAVEPOINT test1") {
+        Statement::Savepoint { name } => {
+            assert_eq!(Ident::new("test1"), name);
+        }
+        _ => unreachable!(),
+    }
+}
+
 #[test]
 fn parse_comments() {
     match pg().verified_stmt("COMMENT ON COLUMN tab.name IS 'comment'") {