Start adding more flightsql functionality

Author: matthewmturner

crates/datafusion-app/src/flightsql.rs

@@ -17,15 +17,17 @@
 
 use std::sync::Arc;
 
-use arrow_flight::sql::client::FlightSqlServiceClient;
+use arrow_flight::{
+    decode::FlightRecordBatchStream, sql::client::FlightSqlServiceClient, FlightInfo,
+};
 #[cfg(feature = "flightsql")]
 use base64::engine::{general_purpose::STANDARD, Engine as _};
 use datafusion::{
     error::{DataFusionError, Result as DFResult},
     physical_plan::stream::RecordBatchStreamAdapter,
     sql::parser::DFParser,
 };
-use log::{error, info, warn};
+use log::{debug, error, info, warn};
 
 use color_eyre::eyre::{self, Result};
 use tokio::sync::Mutex;
@@ -203,4 +205,43 @@ impl FlightSQLContext {
             return Err(DataFusionError::External("Missing client".into()));
         }
     }
+
+    pub async fn get_catalogs_flight_info(&self) -> DFResult<FlightInfo> {
+        let client = self.client.clone();
+        let mut guard = client.lock().await;
+        if let Some(client) = guard.as_mut() {
+            client
+                .get_catalogs()
+                .await
+                .map_err(|e| DataFusionError::ArrowError(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 = self.client.clone();
+        let mut guard = client.lock().await;
+        if let Some(client) = guard.as_mut() {
+            let mut streams = Vec::new();
+            for endpoint in flight_info.endpoint {
+                if let Some(ticket) = endpoint.ticket {
+                    let stream = client
+                        .do_get(ticket.into_request())
+                        .await
+                        .map_err(|e| DataFusionError::ArrowError(e, None))?;
+                    streams.push(stream);
+                } else {
+                    debug!("No ticket for endpoint: {endpoint}");
+                }
+            }
+            Ok(streams)
+        } else {
+            Err(DataFusionError::External(
+                "No FlightSQL client configured. Add one in `~/.config/dft/config.toml`".into(),
+            ))
+        }
+    }
 }

src/args.rs

@@ -130,7 +130,9 @@ impl DftArgs {
 #[derive(Clone, Debug, Subcommand)]
 pub enum FlightSqlCommand {
     /// Executes `GetFlightInfo` and `DoGet` on the provided SQL query
-    Query { sql: String },
+    Statement { sql: String },
+    /// Executes `GetCatalogsFlightInfo` and `DoGet`
+    Catalogs,
 }
 
 #[derive(Clone, Debug, Subcommand)]

src/cli/mod.rs

@@ -41,6 +41,7 @@ use std::path::{Path, PathBuf};
 #[cfg(feature = "flightsql")]
 use {
     crate::args::FlightSqlCommand,
+    arrow_flight::FlightInfo,
     datafusion_app::{
         config::{AuthConfig, FlightSQLConfig},
         flightsql::FlightSQLContext,
@@ -87,7 +88,15 @@ impl CliApp {
     #[cfg(feature = "flightsql")]
     async fn handle_flightsql_command(&self, command: FlightSqlCommand) -> color_eyre::Result<()> {
         match command {
-            FlightSqlCommand::Query { sql } => self.exec_from_flightsql(sql, 0).await,
+            FlightSqlCommand::Statement { sql } => self.exec_from_flightsql(sql, 0).await,
+            FlightSqlCommand::Catalogs => {
+                let flight_info = self
+                    .app_execution
+                    .flightsql_ctx()
+                    .get_catalogs_flight_info()
+                    .await?;
+                Ok(())
+            }
         }
     }
 

src/server/flightsql/service.rs

@@ -20,7 +20,9 @@ use arrow_flight::encode::FlightDataEncoderBuilder;
 use arrow_flight::error::FlightError;
 use arrow_flight::flight_service_server::{FlightService, FlightServiceServer};
 use arrow_flight::sql::server::FlightSqlService;
-use arrow_flight::sql::{Any, CommandStatementQuery, SqlInfo, TicketStatementQuery};
+use arrow_flight::sql::{
+    Any, CommandGetCatalogs, CommandStatementQuery, SqlInfo, TicketStatementQuery,
+};
 use arrow_flight::{FlightDescriptor, FlightEndpoint, FlightInfo, Ticket};
 use color_eyre::Result;
 use datafusion::logical_expr::LogicalPlan;
@@ -218,6 +220,31 @@ impl FlightSqlServiceImpl {
 impl FlightSqlService for FlightSqlServiceImpl {
     type FlightService = FlightSqlServiceImpl;
 
+    async fn get_flight_info_catalogs(
+        &self,
+        _query: CommandGetCatalogs,
+        request: Request<FlightDescriptor>,
+    ) -> Result<Response<FlightInfo>, Status> {
+        counter!("requests", "endpoint" => "get_flight_info").increment(1);
+        let start = Timestamp::now();
+        let request_id = uuid::Uuid::new_v4();
+        let query = "SELECT DISTINCT table_catalog FROM information_schema.tables".to_string();
+        let res = self
+            .get_flight_info_statement_handler(query, request_id, request)
+            .await;
+
+        // TODO: Move recording to after response is sent to not impact response latency
+        self.record_request(
+            start,
+            Some(request_id.to_string()),
+            res.as_ref().err(),
+            "/get_flight_info_catalogs".to_string(),
+            "get_flight_info_catalogs_latency_ms",
+        )
+        .await;
+        res
+    }
+
     async fn get_flight_info_statement(
         &self,
         query: CommandStatementQuery,

tests/extension_cases/flightsql.rs

@@ -645,7 +645,7 @@ async fn test_flightsql_query_command() {
         Command::cargo_bin("dft")
             .unwrap()
             .arg("flightsql")
-            .arg("query")
+            .arg("statement")
             .arg(sql.clone())
             .timeout(Duration::from_secs(5))
             .assert()