Add CommandGetXdbcTypeInfo to Flight SQL Server (#4055)

* Add CommandGetXdbcTypeInfo to Flight SQL Server * Add CommandGetXdbcTypeInfo in a few moer places * Add get_xdbc_type_info to flight-sql client --------- Co-authored-by: Andrew Lamb <[email protected]>
apache · Apr 12, 2023 · 1e08706 · 1e08706
1 parent a35c6c5
commit 1e08706
Show file tree

Hide file tree

Showing 5 changed files with 82 additions and 12 deletions.
diff --git a/arrow-flight/examples/flight_sql_server.rs b/arrow-flight/examples/flight_sql_server.rs
@@ -44,9 +44,9 @@ use arrow_flight::{
         ActionCreatePreparedStatementRequest, CommandGetCatalogs,
         CommandGetCrossReference, CommandGetDbSchemas, CommandGetExportedKeys,
         CommandGetImportedKeys, CommandGetPrimaryKeys, CommandGetSqlInfo,
-        CommandGetTableTypes, CommandGetTables, CommandPreparedStatementQuery,
-        CommandPreparedStatementUpdate, CommandStatementQuery, CommandStatementUpdate,
-        TicketStatementQuery,
+        CommandGetTableTypes, CommandGetTables, CommandGetXdbcTypeInfo,
+        CommandPreparedStatementQuery, CommandPreparedStatementUpdate,
+        CommandStatementQuery, CommandStatementUpdate, TicketStatementQuery,
     },
     FlightDescriptor, FlightInfo,
 };
@@ -315,6 +315,16 @@ impl FlightSqlService for FlightSqlServiceImpl {
         ))
     }
 
+    async fn get_flight_info_xdbc_type_info(
+        &self,
+        _query: CommandGetXdbcTypeInfo,
+        _request: Request<FlightDescriptor>,
+    ) -> Result<Response<FlightInfo>, Status> {
+        Err(Status::unimplemented(
+            "get_flight_info_xdbc_type_info not implemented",
+        ))
+    }
+
     // do_get
     async fn do_get_statement(
         &self,
@@ -412,6 +422,16 @@ impl FlightSqlService for FlightSqlServiceImpl {
         ))
     }
 
+    async fn do_get_xdbc_type_info(
+        &self,
+        _query: CommandGetXdbcTypeInfo,
+        _request: Request<Ticket>,
+    ) -> Result<Response<<Self as FlightService>::DoGetStream>, Status> {
+        Err(Status::unimplemented(
+            "do_get_xdbc_type_info not implemented",
+        ))
+    }
+
     // do_put
     async fn do_put_statement_update(
         &self,

diff --git a/arrow-flight/src/sql/client.rs b/arrow-flight/src/sql/client.rs
@@ -31,9 +31,9 @@ use crate::sql::{
     ActionCreatePreparedStatementResult, Any, CommandGetCatalogs,
     CommandGetCrossReference, CommandGetDbSchemas, CommandGetExportedKeys,
     CommandGetImportedKeys, CommandGetPrimaryKeys, CommandGetSqlInfo,
-    CommandGetTableTypes, CommandGetTables, CommandPreparedStatementQuery,
-    CommandStatementQuery, CommandStatementUpdate, DoPutUpdateResult, ProstMessageExt,
-    SqlInfo,
+    CommandGetTableTypes, CommandGetTables, CommandGetXdbcTypeInfo,
+    CommandPreparedStatementQuery, CommandStatementQuery, CommandStatementUpdate,
+    DoPutUpdateResult, ProstMessageExt, SqlInfo,
 };
 use crate::{
     Action, FlightData, FlightDescriptor, FlightInfo, HandshakeRequest,
@@ -313,6 +313,14 @@ impl FlightSqlServiceClient<Channel> {
         self.get_flight_info_for_command(request).await
     }
 
+    /// Request XDBC SQL information.
+    pub async fn get_xdbc_type_info(
+        &mut self,
+        request: CommandGetXdbcTypeInfo,
+    ) -> Result<FlightInfo, ArrowError> {
+        self.get_flight_info_for_command(request).await
+    }
+
     /// Create a prepared statement object.
     pub async fn prepare(
         &mut self,

diff --git a/arrow-flight/src/sql/mod.rs b/arrow-flight/src/sql/mod.rs
@@ -58,6 +58,7 @@ pub use gen::CommandGetPrimaryKeys;
 pub use gen::CommandGetSqlInfo;
 pub use gen::CommandGetTableTypes;
 pub use gen::CommandGetTables;
+pub use gen::CommandGetXdbcTypeInfo;
 pub use gen::CommandPreparedStatementQuery;
 pub use gen::CommandPreparedStatementUpdate;
 pub use gen::CommandStatementQuery;
@@ -214,6 +215,7 @@ prost_message_ext!(
     CommandGetSqlInfo,
     CommandGetTableTypes,
     CommandGetTables,
+    CommandGetXdbcTypeInfo,
     CommandPreparedStatementQuery,
     CommandPreparedStatementUpdate,
     CommandStatementQuery,

diff --git a/arrow-flight/src/sql/server.rs b/arrow-flight/src/sql/server.rs
@@ -34,9 +34,9 @@ use super::{
     ActionCreatePreparedStatementResult, CommandGetCatalogs, CommandGetCrossReference,
     CommandGetDbSchemas, CommandGetExportedKeys, CommandGetImportedKeys,
     CommandGetPrimaryKeys, CommandGetSqlInfo, CommandGetTableTypes, CommandGetTables,
-    CommandPreparedStatementQuery, CommandPreparedStatementUpdate, CommandStatementQuery,
-    CommandStatementUpdate, DoPutUpdateResult, ProstMessageExt, SqlInfo,
-    TicketStatementQuery,
+    CommandGetXdbcTypeInfo, CommandPreparedStatementQuery,
+    CommandPreparedStatementUpdate, CommandStatementQuery, CommandStatementUpdate,
+    DoPutUpdateResult, ProstMessageExt, SqlInfo, TicketStatementQuery,
 };
 
 pub(crate) static CREATE_PREPARED_STATEMENT: &str = "CreatePreparedStatement";
@@ -151,6 +151,13 @@ pub trait FlightSqlService: Sync + Send + Sized + 'static {
         request: Request<FlightDescriptor>,
     ) -> Result<Response<FlightInfo>, Status>;
 
+    /// Get a FlightInfo to extract information about the supported XDBC types.
+    async fn get_flight_info_xdbc_type_info(
+        &self,
+        query: CommandGetXdbcTypeInfo,
+        request: Request<FlightDescriptor>,
+    ) -> Result<Response<FlightInfo>, Status>;
+
     // do_get
 
     /// Get a FlightDataStream containing the query results.
@@ -230,6 +237,13 @@ pub trait FlightSqlService: Sync + Send + Sized + 'static {
         request: Request<Ticket>,
     ) -> Result<Response<<Self as FlightService>::DoGetStream>, Status>;
 
+    /// Get a FlightDataStream containing the data related to the supported XDBC types.
+    async fn do_get_xdbc_type_info(
+        &self,
+        query: CommandGetXdbcTypeInfo,
+        request: Request<Ticket>,
+    ) -> Result<Response<<Self as FlightService>::DoGetStream>, Status>;
+
     // do_put
 
     /// Execute an update SQL statement.
@@ -352,6 +366,9 @@ where
             Command::CommandGetCrossReference(token) => {
                 self.get_flight_info_cross_reference(token, request).await
             }
+            Command::CommandGetXdbcTypeInfo(token) => {
+                self.get_flight_info_xdbc_type_info(token, request).await
+            }
             cmd => Err(Status::unimplemented(format!(
                 "get_flight_info: The defined request is invalid: {}",
                 cmd.type_url()
@@ -407,6 +424,9 @@ where
             Command::CommandGetCrossReference(command) => {
                 self.do_get_cross_reference(command, request).await
             }
+            Command::CommandGetXdbcTypeInfo(command) => {
+                self.do_get_xdbc_type_info(command, request).await
+            }
             cmd => self.do_get_fallback(request, cmd.into_any()).await,
         }
     }

diff --git a/arrow-flight/tests/flight_sql_client_cli.rs b/arrow-flight/tests/flight_sql_client_cli.rs
@@ -26,9 +26,9 @@ use arrow_flight::{
         CommandGetCatalogs, CommandGetCrossReference, CommandGetDbSchemas,
         CommandGetExportedKeys, CommandGetImportedKeys, CommandGetPrimaryKeys,
         CommandGetSqlInfo, CommandGetTableTypes, CommandGetTables,
-        CommandPreparedStatementQuery, CommandPreparedStatementUpdate,
-        CommandStatementQuery, CommandStatementUpdate, ProstMessageExt, SqlInfo,
-        TicketStatementQuery,
+        CommandGetXdbcTypeInfo, CommandPreparedStatementQuery,
+        CommandPreparedStatementUpdate, CommandStatementQuery, CommandStatementUpdate,
+        ProstMessageExt, SqlInfo, TicketStatementQuery,
     },
     utils::batches_to_flight_data,
     Action, FlightData, FlightDescriptor, FlightEndpoint, FlightInfo, HandshakeRequest,
@@ -302,6 +302,16 @@ impl FlightSqlService for FlightSqlServiceImpl {
         ))
     }
 
+    async fn get_flight_info_xdbc_type_info(
+        &self,
+        _query: CommandGetXdbcTypeInfo,
+        _request: Request<FlightDescriptor>,
+    ) -> Result<Response<FlightInfo>, Status> {
+        Err(Status::unimplemented(
+            "get_flight_info_xdbc_type_info not implemented",
+        ))
+    }
+
     // do_get
     async fn do_get_statement(
         &self,
@@ -399,6 +409,16 @@ impl FlightSqlService for FlightSqlServiceImpl {
         ))
     }
 
+    async fn do_get_xdbc_type_info(
+        &self,
+        _query: CommandGetXdbcTypeInfo,
+        _request: Request<Ticket>,
+    ) -> Result<Response<<Self as FlightService>::DoGetStream>, Status> {
+        Err(Status::unimplemented(
+            "do_get_xdbc_type_info not implemented",
+        ))
+    }
+
     // do_put
     async fn do_put_statement_update(
         &self,