feat: add transaction related method

Author: hantmac

driver/src/conn.rs

@@ -190,6 +190,10 @@ pub trait Connection: DynClone + Send + Sync {
         ))
     }
 
+    async fn begin(&self) -> Result<()>;
+    async fn commit(&self) -> Result<()>;
+    async fn rollback(&self) -> Result<()>;
+
     async fn get_files(&self, stage: &str, local_file: &str) -> Result<RowStatsIterator> {
         let mut total_count: usize = 0;
         let mut total_size: usize = 0;

driver/src/flight_sql.rs

@@ -143,6 +143,21 @@ impl Connection for FlightSQLConnection {
             "STREAM LOAD unavailable for FlightSQL".to_string(),
         ))
     }
+
+    async fn begin(&self) -> Result<()> {
+        self.exec("BEGIN").await.unwrap();
+        Ok(())
+    }
+
+    async fn commit(&self) -> Result<()> {
+        self.exec("COMMIT").await.unwrap();
+        Ok(())
+    }
+
+    async fn rollback(&self) -> Result<()> {
+        self.exec("ROLLBACK").await.unwrap();
+        Ok(())
+    }
 }
 
 impl FlightSQLConnection {
@@ -273,7 +288,7 @@ impl Args {
                         return Err(Error::BadArgument(format!(
                             "Invalid value for sslmode: {}",
                             v.as_ref()
-                        )))
+                        )));
                     }
                 },
                 "tls_ca_file" => args.tls_ca_file = Some(v.to_string()),

driver/src/rest_api.rs

@@ -187,6 +187,21 @@ impl Connection for RestAPIConnection {
         let stats = self.load_data(sql, reader, size, None, None).await?;
         Ok(stats)
     }
+
+    async fn begin(&self) -> Result<()> {
+        self.exec("BEGIN").await.unwrap();
+        Ok(())
+    }
+
+    async fn commit(&self) -> Result<()> {
+        self.exec("COMMIT").await.unwrap();
+        Ok(())
+    }
+
+    async fn rollback(&self) -> Result<()> {
+        self.exec("ROLLBACK").await.unwrap();
+        Ok(())
+    }
 }
 
 impl<'o> RestAPIConnection {
@@ -219,16 +234,16 @@ impl<'o> RestAPIConnection {
             ("record_delimiter", "\n"),
             ("skip_header", "0"),
         ]
-        .into_iter()
-        .collect()
+            .into_iter()
+            .collect()
     }
 
     fn default_copy_options() -> BTreeMap<&'o str, &'o str> {
         vec![("purge", "true")].into_iter().collect()
     }
 }
 
-type PageFut = Pin<Box<dyn Future<Output = Result<QueryResponse>> + Send>>;
+type PageFut = Pin<Box<dyn Future<Output=Result<QueryResponse>> + Send>>;
 
 pub struct RestAPIRows {
     client: APIClient,

driver/tests/driver/main.rs

@@ -19,3 +19,4 @@ mod load;
 mod select_iter;
 mod select_simple;
 mod session;
+mod transaction;

driver/tests/driver/transaction.rs

@@ -0,0 +1,55 @@
+// Copyright 2021 Datafuse Labs
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use databend_driver::Client;
+
+use crate::common::DEFAULT_DSN;
+
+
+#[tokio::test]
+async fn test_commit() {
+    let dsn = option_env!("TEST_DATABEND_DSN").unwrap_or(DEFAULT_DSN);
+    let client = Client::new(dsn.to_string());
+    let conn = client.get_conn().await.unwrap();
+
+    conn.exec("CREATE OR REPLACE TABLE t(c int);").await.unwrap();
+    conn.begin().await.unwrap();
+    conn.exec("INSERT INTO t VALUES(1);").await.unwrap();
+    let row = conn.query_row("SELECT * FROM t").await.unwrap();
+    let row = row.unwrap();
+    let (val, ): (i32, ) = row.try_into().unwrap();
+    assert_eq!(val, 1);
+    conn.commit().await.unwrap();
+}
+
+#[tokio::test]
+async fn test_rollback() {
+    let dsn = option_env!("TEST_DATABEND_DSN").unwrap_or(DEFAULT_DSN);
+    let client = Client::new(dsn.to_string());
+    let conn = client.get_conn().await.unwrap();
+
+    conn.exec("CREATE OR REPLACE TABLE t(c int);").await.unwrap();
+    conn.begin().await.unwrap();
+    conn.exec("INSERT INTO t VALUES(1);").await.unwrap();
+    let row = conn.query_row("SELECT * FROM t").await.unwrap();
+    let row = row.unwrap();
+    let (val, ): (i32, ) = row.try_into().unwrap();
+    assert_eq!(val, 1);
+    conn.rollback().await.unwrap();
+    let row = conn.query_row("SELECT * FROM t").await.unwrap();
+    let row = row.unwrap();
+    let (val, ): (Option<i32>, ) = row.try_into().unwrap();
+    assert_eq!(val, None)
+}
+