More endpoints

Author: matthewmturner

.claude/settings.local.json

@@ -21,7 +21,8 @@
       "Bash(cat:*)",
       "Bash(cargo clippy:*)",
       "Bash(find:*)",
-      "WebFetch(domain:arrow.apache.org)"
+      "WebFetch(domain:arrow.apache.org)",
+      "Bash(timeout 90 cargo test:*)"
     ],
     "deny": [],
     "ask": []

crates/datafusion-app/src/flightsql.rs

@@ -19,7 +19,10 @@ use std::sync::Arc;
 
 use arrow_flight::{
     decode::FlightRecordBatchStream,
-    sql::{client::FlightSqlServiceClient, CommandGetDbSchemas, CommandGetTables},
+    sql::{
+        client::FlightSqlServiceClient, CommandGetDbSchemas, CommandGetSqlInfo,
+        CommandGetTableTypes, CommandGetTables, CommandGetXdbcTypeInfo,
+    },
     FlightInfo,
 };
 #[cfg(feature = "flightsql")]
@@ -287,6 +290,62 @@ impl FlightSQLContext {
         }
     }
 
+    pub async fn get_table_types_flight_info(&self) -> DFResult<FlightInfo> {
+        let client = Arc::clone(&self.client);
+        let mut guard = client.lock().await;
+        if let Some(client) = guard.as_mut() {
+            client
+                .get_table_types()
+                .await
+                .map_err(|e| DataFusionError::ArrowError(Box::new(e), None))
+        } else {
+            Err(DataFusionError::External(
+                "No FlightSQL client configured.  Add one in `~/.config/dft/config.toml`".into(),
+            ))
+        }
+    }
+
+    pub async fn get_sql_info_flight_info(&self, info: Option<Vec<u32>>) -> DFResult<FlightInfo> {
+        let client = Arc::clone(&self.client);
+        let mut guard = client.lock().await;
+        if let Some(client) = guard.as_mut() {
+            use arrow_flight::sql::SqlInfo;
+            // Convert u32 IDs to SqlInfo enum variants if needed
+            let sql_info_list: Vec<SqlInfo> = info
+                .unwrap_or_default()
+                .into_iter()
+                .filter_map(|id| SqlInfo::try_from(id as i32).ok())
+                .collect();
+            client
+                .get_sql_info(sql_info_list)
+                .await
+                .map_err(|e| DataFusionError::ArrowError(Box::new(e), None))
+        } else {
+            Err(DataFusionError::External(
+                "No FlightSQL client configured.  Add one in `~/.config/dft/config.toml`".into(),
+            ))
+        }
+    }
+
+    pub async fn get_xdbc_type_info_flight_info(
+        &self,
+        data_type: Option<i32>,
+    ) -> DFResult<FlightInfo> {
+        let client = Arc::clone(&self.client);
+        let mut guard = client.lock().await;
+        if let Some(client) = guard.as_mut() {
+            let cmd = CommandGetXdbcTypeInfo { data_type };
+            client
+                .get_xdbc_type_info(cmd)
+                .await
+                .map_err(|e| DataFusionError::ArrowError(Box::new(e), None))
+        } else {
+            Err(DataFusionError::External(
+                "No FlightSQL client configured.  Add one in `~/.config/dft/config.toml`".into(),
+            ))
+        }
+    }
+
     pub async fn do_get(&self, flight_info: FlightInfo) -> DFResult<Vec<FlightRecordBatchStream>> {
         let client = Arc::clone(&self.client);
         let mut guard = client.lock().await;

src/args.rs

@@ -171,6 +171,20 @@ pub enum FlightSqlCommand {
         #[clap(long)]
         table_types: Option<Vec<String>>,
     },
+    /// Executes `CommandGetTableTypes` and `DoGet` to return supported table types
+    GetTableTypes,
+    /// Executes `CommandGetSqlInfo` and `DoGet` to return server SQL capabilities
+    GetSqlInfo {
+        /// Specific SQL info IDs to retrieve (if not provided, returns all)
+        #[clap(long)]
+        info: Option<Vec<u32>>,
+    },
+    /// Executes `CommandGetXdbcTypeInfo` and `DoGet` to return type information
+    GetXdbcTypeInfo {
+        /// Optional data type to filter by
+        #[clap(long)]
+        data_type: Option<i32>,
+    },
 }
 
 #[derive(Clone, Debug, Subcommand)]

src/cli/mod.rs

@@ -155,6 +155,51 @@ impl CliApp {
                 self.print_any_stream(flight_batch_stream).await;
                 Ok(())
             }
+            FlightSqlCommand::GetTableTypes => {
+                let flight_info = self
+                    .app_execution
+                    .flightsql_ctx()
+                    .get_table_types_flight_info()
+                    .await?;
+                let streams = self
+                    .app_execution
+                    .flightsql_ctx()
+                    .do_get(flight_info)
+                    .await?;
+                let flight_batch_stream = stream::select_all(streams);
+                self.print_any_stream(flight_batch_stream).await;
+                Ok(())
+            }
+            FlightSqlCommand::GetSqlInfo { info } => {
+                let flight_info = self
+                    .app_execution
+                    .flightsql_ctx()
+                    .get_sql_info_flight_info(info)
+                    .await?;
+                let streams = self
+                    .app_execution
+                    .flightsql_ctx()
+                    .do_get(flight_info)
+                    .await?;
+                let flight_batch_stream = stream::select_all(streams);
+                self.print_any_stream(flight_batch_stream).await;
+                Ok(())
+            }
+            FlightSqlCommand::GetXdbcTypeInfo { data_type } => {
+                let flight_info = self
+                    .app_execution
+                    .flightsql_ctx()
+                    .get_xdbc_type_info_flight_info(data_type)
+                    .await?;
+                let streams = self
+                    .app_execution
+                    .flightsql_ctx()
+                    .do_get(flight_info)
+                    .await?;
+                let flight_batch_stream = stream::select_all(streams);
+                self.print_any_stream(flight_batch_stream).await;
+                Ok(())
+            }
         }
     }