diff --git a/driver/src/client.rs b/driver/src/client.rs index 0bf06130..eeca70b2 100644 --- a/driver/src/client.rs +++ b/driver/src/client.rs @@ -144,6 +144,18 @@ impl Connection { self.inner.close().await } + pub async fn begin(&self) -> Result<()> { + self.inner.begin().await + } + + pub async fn commit(&self) -> Result<()> { + self.inner.commit().await + } + + pub async fn rollback(&self) -> Result<()> { + self.inner.rollback().await + } + pub fn last_query_id(&self) -> Option { self.inner.last_query_id() } @@ -670,4 +682,158 @@ mod tests { assert!(statement_is_dml(update.as_ref())); assert!(!statement_is_dml(query.as_ref())); } + + struct MockConnection { + executed: std::sync::Arc>>, + failed_sql: Option, + } + + #[async_trait::async_trait] + impl IConnection for MockConnection { + async fn info(&self) -> ConnectionInfo { + ConnectionInfo { + handler: "mock".to_string(), + host: "localhost".to_string(), + port: 0, + user: "root".to_string(), + catalog: None, + database: None, + warehouse: None, + } + } + + fn last_query_id(&self) -> Option { + None + } + + async fn exec(&self, sql: &str) -> Result { + self.executed + .lock() + .map_err(|e| Error::IO(e.to_string()))? + .push(sql.to_string()); + if self.failed_sql.as_deref() == Some(sql) { + return Err(Error::BadArgument(format!("failed to execute {sql}"))); + } + Ok(0) + } + + async fn kill_query(&self, _query_id: &str) -> Result<()> { + Ok(()) + } + + async fn query_iter(&self, _sql: &str) -> Result { + Ok(RowIterator::new( + std::sync::Arc::new(databend_client::schema::Schema::from_vec(vec![])), + Box::pin(tokio_stream::empty()), + )) + } + + async fn query_iter_ext(&self, _sql: &str) -> Result { + Ok(RowStatsIterator::new( + std::sync::Arc::new(databend_client::schema::Schema::from_vec(vec![])), + Box::pin(tokio_stream::empty()), + )) + } + + async fn upload_to_stage(&self, _stage: &str, _data: Reader, _size: u64) -> Result<()> { + Ok(()) + } + + async fn load_data( + &self, + _sql: &str, + _data: Reader, + _size: u64, + _method: LoadMethod, + ) -> Result { + Ok(ServerStats::default()) + } + + async fn load_file( + &self, + _sql: &str, + _fp: &Path, + _method: LoadMethod, + ) -> Result { + Ok(ServerStats::default()) + } + + async fn load_file_with_options( + &self, + _sql: &str, + _fp: &Path, + _file_format_options: Option>, + _copy_options: Option>, + ) -> Result { + Ok(ServerStats::default()) + } + + async fn stream_load( + &self, + _sql: &str, + _data: Vec>, + _method: LoadMethod, + ) -> Result { + Ok(ServerStats::default()) + } + + fn set_warehouse(&self, _warehouse: &str) -> Result<()> { + Ok(()) + } + + fn set_database(&self, _database: &str) -> Result<()> { + Ok(()) + } + + fn set_role(&self, _role: &str) -> Result<()> { + Ok(()) + } + + fn set_session(&self, _key: &str, _value: &str) -> Result<()> { + Ok(()) + } + } + + fn mock_connection( + failed_sql: Option<&str>, + ) -> (Connection, std::sync::Arc>>) { + let executed = std::sync::Arc::new(std::sync::Mutex::new(Vec::new())); + let connection = Connection { + inner: Box::new(MockConnection { + executed: executed.clone(), + failed_sql: failed_sql.map(str::to_string), + }), + }; + (connection, executed) + } + + #[tokio::test] + async fn transaction_methods_execute_control_statements() { + let (connection, executed) = mock_connection(None); + + assert!(connection.begin().await.is_ok()); + assert!(connection.commit().await.is_ok()); + assert!(connection.rollback().await.is_ok()); + + let executed = executed + .lock() + .expect("mock execution log should not be poisoned"); + assert_eq!(executed.as_slice(), ["BEGIN", "COMMIT", "ROLLBACK"]); + } + + #[tokio::test] + async fn transaction_methods_propagate_execution_errors() { + let (connection, executed) = mock_connection(Some("ROLLBACK")); + + let error = connection.rollback().await; + assert!(matches!( + error, + Err(Error::BadArgument(message)) if message == "failed to execute ROLLBACK" + )); + + let executed = executed + .lock() + .expect("mock execution log should not be poisoned"); + assert_eq!(executed.as_slice(), ["ROLLBACK"]); + } } diff --git a/driver/src/conn.rs b/driver/src/conn.rs index cfc20a21..d4f8d4c2 100644 --- a/driver/src/conn.rs +++ b/driver/src/conn.rs @@ -72,6 +72,22 @@ pub trait IConnection: Send + Sync { } async fn exec(&self, sql: &str) -> Result; + + async fn begin(&self) -> Result<()> { + self.exec("BEGIN").await?; + Ok(()) + } + + async fn commit(&self) -> Result<()> { + self.exec("COMMIT").await?; + Ok(()) + } + + async fn rollback(&self) -> Result<()> { + self.exec("ROLLBACK").await?; + Ok(()) + } + async fn kill_query(&self, query_id: &str) -> Result<()>; async fn query_iter(&self, sql: &str) -> Result; async fn query_iter_ext(&self, sql: &str) -> Result;