Add begin/end_transaction methods in FlightSqlServiceClient

arrow-flight/Cargo.toml

@@ -76,6 +76,7 @@ pin-project-lite = "0.2"
 tempfile = "3.3"
 tokio-stream = { version = "0.1", features = ["net"] }
 tower = "0.4.13"
+uuid = { version = "1.10.0", features = ["v4"] }
 
 [[example]]
 name = "flight_sql_server"

arrow-flight/src/sql/client.rs

@@ -28,15 +28,19 @@ use crate::decode::FlightRecordBatchStream;
 use crate::encode::FlightDataEncoderBuilder;
 use crate::error::FlightError;
 use crate::flight_service_client::FlightServiceClient;
-use crate::sql::server::{CLOSE_PREPARED_STATEMENT, CREATE_PREPARED_STATEMENT};
+use crate::sql::gen::action_end_transaction_request::EndTransaction;
+use crate::sql::server::{
+    BEGIN_TRANSACTION, CLOSE_PREPARED_STATEMENT, CREATE_PREPARED_STATEMENT, END_TRANSACTION,
+};
 use crate::sql::{
+    ActionBeginTransactionRequest, ActionBeginTransactionResult,
     ActionClosePreparedStatementRequest, ActionCreatePreparedStatementRequest,
-    ActionCreatePreparedStatementResult, Any, CommandGetCatalogs, CommandGetCrossReference,
-    CommandGetDbSchemas, CommandGetExportedKeys, CommandGetImportedKeys, CommandGetPrimaryKeys,
-    CommandGetSqlInfo, CommandGetTableTypes, CommandGetTables, CommandGetXdbcTypeInfo,
-    CommandPreparedStatementQuery, CommandPreparedStatementUpdate, CommandStatementQuery,
-    CommandStatementUpdate, DoPutPreparedStatementResult, DoPutUpdateResult, ProstMessageExt,
-    SqlInfo,
+    ActionCreatePreparedStatementResult, ActionEndTransactionRequest, Any, CommandGetCatalogs,
+    CommandGetCrossReference, CommandGetDbSchemas, CommandGetExportedKeys, CommandGetImportedKeys,
+    CommandGetPrimaryKeys, CommandGetSqlInfo, CommandGetTableTypes, CommandGetTables,
+    CommandGetXdbcTypeInfo, CommandPreparedStatementQuery, CommandPreparedStatementUpdate,
+    CommandStatementQuery, CommandStatementUpdate, DoPutPreparedStatementResult, DoPutUpdateResult,
+    ProstMessageExt, SqlInfo,
 };
 use crate::trailers::extract_lazy_trailers;
 use crate::{
@@ -399,6 +403,54 @@ impl FlightSqlServiceClient<Channel> {
         ))
     }
 
+    /// Request to begin a transaction.
+    pub async fn begin_transaction(&mut self) -> Result<Bytes, ArrowError> {
+        let cmd = ActionBeginTransactionRequest {};
+        let action = Action {
+            r#type: BEGIN_TRANSACTION.to_string(),
+            body: cmd.as_any().encode_to_vec().into(),
+        };
+        let req = self.set_request_headers(action.into_request())?;
+        let mut result = self
+            .flight_client
+            .do_action(req)
+            .await
+            .map_err(status_to_arrow_error)?
+            .into_inner();
+        let result = result
+            .message()
+            .await
+            .map_err(status_to_arrow_error)?
+            .unwrap();
+        let any = Any::decode(&*result.body).map_err(decode_error_to_arrow_error)?;
+        let begin_result: ActionBeginTransactionResult = any.unpack()?.unwrap();
+        Ok(begin_result.transaction_id)
+    }
+
+    /// Request to commit/rollback a transaction.
+    pub async fn end_transaction(
+        &mut self,
+        transaction_id: Bytes,
+        action: EndTransaction,
+    ) -> Result<(), ArrowError> {
+        let cmd = ActionEndTransactionRequest {
+            transaction_id,
+            action: action as i32,
+        };
+        let action = Action {
+            r#type: END_TRANSACTION.to_string(),
+            body: cmd.as_any().encode_to_vec().into(),
+        };
+        let req = self.set_request_headers(action.into_request())?;
+        let _ = self
+            .flight_client
+            .do_action(req)
+            .await
+            .map_err(status_to_arrow_error)?
+            .into_inner();
+        Ok(())
+    }
+
     /// Explicitly shut down and clean up the client.
     pub async fn close(&mut self) -> Result<(), ArrowError> {
         // TODO: consume self instead of &mut self to explicitly prevent reuse?

arrow-flight/src/sql/mod.rs

@@ -49,6 +49,7 @@ mod gen {
     include!("arrow.flight.protocol.sql.rs");
 }
 
+pub use gen::action_end_transaction_request::EndTransaction;
 pub use gen::ActionBeginSavepointRequest;
 pub use gen::ActionBeginSavepointResult;
 pub use gen::ActionBeginTransactionRequest;