support for session transaction and transaction snapshot. (#379)

* add support for snapshot id in set transaction

Signed-off-by: poonai <[email protected]>

* add support for default session transaction characteristics

Signed-off-by: poonai <[email protected]>

* add additional assertion for parse_set_transaction test

Signed-off-by: poonai <[email protected]>

* Fix clippy

Co-authored-by: Andrew Lamb <[email protected]>

Author: poonai

src/ast/mod.rs

@@ -786,7 +786,11 @@ pub enum Statement {
     /// `{ BEGIN [ TRANSACTION | WORK ] | START TRANSACTION } ...`
     StartTransaction { modes: Vec<TransactionMode> },
     /// `SET TRANSACTION ...`
-    SetTransaction { modes: Vec<TransactionMode> },
+    SetTransaction {
+        modes: Vec<TransactionMode>,
+        snapshot: Option<Value>,
+        session: bool,
+    },
     /// `COMMIT [ TRANSACTION | WORK ] [ AND [ NO ] CHAIN ]`
     Commit { chain: bool },
     /// `ROLLBACK [ TRANSACTION | WORK ] [ AND [ NO ] CHAIN ]`
@@ -1369,11 +1373,22 @@ impl fmt::Display for Statement {
                 }
                 Ok(())
             }
-            Statement::SetTransaction { modes } => {
-                write!(f, "SET TRANSACTION")?;
+            Statement::SetTransaction {
+                modes,
+                snapshot,
+                session,
+            } => {
+                if *session {
+                    write!(f, "SET SESSION CHARACTERISTICS AS TRANSACTION")?;
+                } else {
+                    write!(f, "SET TRANSACTION")?;
+                }
                 if !modes.is_empty() {
                     write!(f, " {}", display_comma_separated(modes))?;
                 }
+                if let Some(snapshot_id) = snapshot {
+                    write!(f, " SNAPSHOT {}", snapshot_id)?;
+                }
                 Ok(())
             }
             Statement::Commit { chain } => {

src/keywords.rs

@@ -412,6 +412,7 @@ define_keywords!(
     SHOW,
     SIMILAR,
     SMALLINT,
+    SNAPSHOT,
     SOME,
     SORT,
     SPECIFIC,

src/parser.rs

@@ -2670,9 +2670,26 @@ impl<'a> Parser<'a> {
                     value: values,
                 });
             }
+        } else if variable.value == "CHARACTERISTICS" {
+            self.expect_keywords(&[Keyword::AS, Keyword::TRANSACTION])?;
+            Ok(Statement::SetTransaction {
+                modes: self.parse_transaction_modes()?,
+                snapshot: None,
+                session: true,
+            })
         } else if variable.value == "TRANSACTION" && modifier.is_none() {
+            if self.parse_keyword(Keyword::SNAPSHOT) {
+                let snaphot_id = self.parse_value()?;
+                return Ok(Statement::SetTransaction {
+                    modes: vec![],
+                    snapshot: Some(snaphot_id),
+                    session: false,
+                });
+            }
             Ok(Statement::SetTransaction {
                 modes: self.parse_transaction_modes()?,
+                snapshot: None,
+                session: false,
             })
         } else {
             self.expected("equals sign or TO", self.peek_token())

tests/sqlparser_common.rs

@@ -3586,14 +3586,22 @@ fn parse_set_transaction() {
     // TRANSACTION, so no need to duplicate the tests here. We just do a quick
     // sanity check.
     match verified_stmt("SET TRANSACTION READ ONLY, READ WRITE, ISOLATION LEVEL SERIALIZABLE") {
-        Statement::SetTransaction { modes } => assert_eq!(
+        Statement::SetTransaction {
             modes,
-            vec![
-                TransactionMode::AccessMode(TransactionAccessMode::ReadOnly),
-                TransactionMode::AccessMode(TransactionAccessMode::ReadWrite),
-                TransactionMode::IsolationLevel(TransactionIsolationLevel::Serializable),
-            ]
-        ),
+            session,
+            snapshot,
+        } => {
+            assert_eq!(
+                modes,
+                vec![
+                    TransactionMode::AccessMode(TransactionAccessMode::ReadOnly),
+                    TransactionMode::AccessMode(TransactionAccessMode::ReadWrite),
+                    TransactionMode::IsolationLevel(TransactionIsolationLevel::Serializable),
+                ]
+            );
+            assert!(!session);
+            assert_eq!(snapshot, None);
+        }
         _ => unreachable!(),
     }
 }

tests/sqlparser_postgres.rs

@@ -723,6 +723,32 @@ fn parse_map_access_expr() {
     );
 }
 
+#[test]
+fn test_transaction_statement() {
+    let statement = pg().verified_stmt("SET TRANSACTION SNAPSHOT '000003A1-1'");
+    assert_eq!(
+        statement,
+        Statement::SetTransaction {
+            modes: vec![],
+            snapshot: Some(Value::SingleQuotedString(String::from("000003A1-1"))),
+            session: false
+        }
+    );
+    let statement = pg().verified_stmt("SET SESSION CHARACTERISTICS AS TRANSACTION READ ONLY, READ WRITE, ISOLATION LEVEL SERIALIZABLE");
+    assert_eq!(
+        statement,
+        Statement::SetTransaction {
+            modes: vec![
+                TransactionMode::AccessMode(TransactionAccessMode::ReadOnly),
+                TransactionMode::AccessMode(TransactionAccessMode::ReadWrite),
+                TransactionMode::IsolationLevel(TransactionIsolationLevel::Serializable),
+            ],
+            snapshot: None,
+            session: true
+        }
+    );
+}
+
 fn pg() -> TestedDialects {
     TestedDialects {
         dialects: vec![Box::new(PostgreSqlDialect {})],