Implement do_put_statement_ingest handler in FlightSqlService

arrow-flight/examples/flight_sql_server.rs

@@ -46,9 +46,9 @@ use arrow_flight::sql::{
     ActionEndTransactionRequest, Any, CommandGetCatalogs, CommandGetCrossReference,
     CommandGetDbSchemas, CommandGetExportedKeys, CommandGetImportedKeys, CommandGetPrimaryKeys,
     CommandGetSqlInfo, CommandGetTableTypes, CommandGetTables, CommandGetXdbcTypeInfo,
-    CommandPreparedStatementQuery, CommandPreparedStatementUpdate, CommandStatementQuery,
-    CommandStatementSubstraitPlan, CommandStatementUpdate, Nullable, ProstMessageExt, Searchable,
-    SqlInfo, TicketStatementQuery, XdbcDataType,
+    CommandPreparedStatementQuery, CommandPreparedStatementUpdate, CommandStatementIngest,
+    CommandStatementQuery, CommandStatementSubstraitPlan, CommandStatementUpdate, Nullable,
+    ProstMessageExt, Searchable, SqlInfo, TicketStatementQuery, XdbcDataType,
 };
 use arrow_flight::utils::batches_to_flight_data;
 use arrow_flight::{
@@ -615,6 +615,14 @@ impl FlightSqlService for FlightSqlServiceImpl {
         Ok(FAKE_UPDATE_RESULT)
     }
 
+    async fn do_put_statement_ingest(
+        &self,
+        _ticket: CommandStatementIngest,
+        _request: Request<PeekableFlightDataStream>,
+    ) -> Result<i64, Status> {
+        Ok(FAKE_UPDATE_RESULT)
+    }
+
     async fn do_put_substrait_plan(
         &self,
         _ticket: CommandStatementSubstraitPlan,

arrow-flight/src/sql/mod.rs

@@ -74,6 +74,7 @@ pub use gen::CommandGetTables;
 pub use gen::CommandGetXdbcTypeInfo;
 pub use gen::CommandPreparedStatementQuery;
 pub use gen::CommandPreparedStatementUpdate;
+pub use gen::CommandStatementIngest;
 pub use gen::CommandStatementQuery;
 pub use gen::CommandStatementSubstraitPlan;
 pub use gen::CommandStatementUpdate;
@@ -250,11 +251,12 @@ prost_message_ext!(
     CommandGetXdbcTypeInfo,
     CommandPreparedStatementQuery,
     CommandPreparedStatementUpdate,
+    CommandStatementIngest,
     CommandStatementQuery,
     CommandStatementSubstraitPlan,
     CommandStatementUpdate,
-    DoPutUpdateResult,
     DoPutPreparedStatementResult,
+    DoPutUpdateResult,
     TicketStatementQuery,
 );
 

arrow-flight/src/sql/server.rs

@@ -32,9 +32,9 @@ use super::{
     CommandGetCrossReference, CommandGetDbSchemas, CommandGetExportedKeys, CommandGetImportedKeys,
     CommandGetPrimaryKeys, CommandGetSqlInfo, CommandGetTableTypes, CommandGetTables,
     CommandGetXdbcTypeInfo, CommandPreparedStatementQuery, CommandPreparedStatementUpdate,
-    CommandStatementQuery, CommandStatementSubstraitPlan, CommandStatementUpdate,
-    DoPutPreparedStatementResult, DoPutUpdateResult, ProstMessageExt, SqlInfo,
-    TicketStatementQuery,
+    CommandStatementIngest, CommandStatementQuery, CommandStatementSubstraitPlan,
+    CommandStatementUpdate, DoPutPreparedStatementResult, DoPutUpdateResult, ProstMessageExt,
+    SqlInfo, TicketStatementQuery,
 };
 use crate::{
     flight_service_server::FlightService, gen::PollInfo, Action, ActionType, Criteria, Empty,
@@ -397,6 +397,17 @@ pub trait FlightSqlService: Sync + Send + Sized + 'static {
         ))
     }
 
+    /// Execute a bulk ingestion.
+    async fn do_put_statement_ingest(
+        &self,
+        _ticket: CommandStatementIngest,
+        _request: Request<PeekableFlightDataStream>,
+    ) -> Result<i64, Status> {
+        Err(Status::unimplemented(
+            "do_put_statement_ingest has no default implementation",
+        ))
+    }
+
     /// Bind parameters to given prepared statement.
     ///
     /// Returns an opaque handle that the client should pass
@@ -713,6 +724,14 @@ where
                 })]);
                 Ok(Response::new(Box::pin(output)))
             }
+            Command::CommandStatementIngest(command) => {
+                let record_count = self.do_put_statement_ingest(command, request).await?;
+                let result = DoPutUpdateResult { record_count };
+                let output = futures::stream::iter(vec![Ok(PutResult {
+                    app_metadata: result.as_any().encode_to_vec().into(),
+                })]);
+                Ok(Response::new(Box::pin(output)))
+            }
             Command::CommandPreparedStatementQuery(command) => {
                 let result = self
                     .do_put_prepared_statement_query(command, request)