Start adding flightsql command (#325)

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(),
+            ))
+        }
+    }
 }

crates/datafusion-app/src/local.rs

@@ -43,7 +43,7 @@ use super::stats::{ExecutionDurationStats, ExecutionStats};
 #[cfg(feature = "udfs-wasm")]
 use super::wasm::create_wasm_udfs;
 #[cfg(feature = "observability")]
-use crate::observability::ObservabilityContext;
+use {crate::config::ObservabilityConfig, crate::observability::ObservabilityContext};
 
 /// Structure for executing queries
 ///
@@ -72,6 +72,24 @@ pub struct ExecutionContext {
     observability: ObservabilityContext,
 }
 
+impl Default for ExecutionContext {
+    fn default() -> Self {
+        let cfg = SessionConfig::new().with_information_schema(true);
+        let session_ctx = SessionContext::new_with_config(cfg);
+        #[cfg(feature = "observability")]
+        let observability =
+            ObservabilityContext::try_new(ObservabilityConfig::default(), "test").unwrap();
+        Self {
+            config: ExecutionConfig::default(),
+            session_ctx,
+            ddl_path: None,
+            executor: None,
+            #[cfg(feature = "observability")]
+            observability,
+        }
+    }
+}
+
 impl std::fmt::Debug for ExecutionContext {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         f.debug_struct("ExecutionContext").finish()

src/args.rs

@@ -127,6 +127,14 @@ impl DftArgs {
     }
 }
 
+#[derive(Clone, Debug, Subcommand)]
+pub enum FlightSqlCommand {
+    /// Executes `GetFlightInfo` and `DoGet` on the provided SQL query
+    StatementQuery { sql: String },
+    /// Executes `GetCatalogsFlightInfo` and `DoGet`
+    GetCatalogs,
+}
+
 #[derive(Clone, Debug, Subcommand)]
 pub enum Command {
     /// Start a HTTP server
@@ -139,6 +147,13 @@ pub enum Command {
         #[clap(long, help = "Set the port to be used for serving metrics")]
         metrics_addr: Option<SocketAddr>,
     },
+    /// Make a request to a FlightSQL server
+    #[cfg(feature = "flightsql")]
+    #[command(name = "flightsql")]
+    FlightSql {
+        #[clap(subcommand)]
+        command: FlightSqlCommand,
+    },
     /// Start a FlightSQL server
     #[cfg(feature = "flightsql")]
     #[command(name = "serve-flightsql")]

src/cli/mod.rs

@@ -39,6 +39,7 @@ use std::io::Write;
 use std::path::{Path, PathBuf};
 #[cfg(feature = "flightsql")]
 use {
+    crate::args::{Command, FlightSqlCommand},
     datafusion_app::{
         config::{AuthConfig, FlightSQLConfig},
         flightsql::FlightSQLContext,
@@ -82,6 +83,31 @@ impl CliApp {
         Ok(())
     }
 
+    #[cfg(feature = "flightsql")]
+    async fn handle_flightsql_command(&self, command: FlightSqlCommand) -> color_eyre::Result<()> {
+        use futures::stream;
+
+        match command {
+            FlightSqlCommand::StatementQuery { sql } => self.exec_from_flightsql(sql, 0).await,
+            FlightSqlCommand::GetCatalogs => {
+                let flight_info = self
+                    .app_execution
+                    .flightsql_ctx()
+                    .get_catalogs_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(())
+            }
+        }
+    }
+
     /// Execute the provided sql, which was passed as an argument from CLI.
     ///
     /// Optionally, use the FlightSQL client for execution.
@@ -92,6 +118,11 @@ impl CliApp {
 
         self.validate_args()?;
 
+        #[cfg(feature = "flightsql")]
+        if let Some(Command::FlightSql { command }) = &self.args.command {
+            return self.handle_flightsql_command(command.clone()).await;
+        };
+
         #[cfg(not(feature = "flightsql"))]
         match (
             self.args.files.is_empty(),
@@ -603,7 +634,7 @@ pub async fn try_run(cli: DftArgs, config: AppConfig) -> Result<()> {
     let mut app_execution = AppExecution::new(execution_ctx);
     #[cfg(feature = "flightsql")]
     {
-        if cli.flightsql {
+        if cli.flightsql || matches!(cli.command, Some(Command::FlightSql { .. })) {
             let auth = AuthConfig {
                 basic_auth: config.flightsql_client.auth.basic_auth,
                 bearer_token: config.flightsql_client.auth.bearer_token,

src/main.rs

@@ -67,10 +67,16 @@ async fn app_entry_point(cli: DftArgs) -> Result<()> {
     }
 
     #[cfg(feature = "flightsql")]
-    if let Some(Command::ServeFlightSql { .. }) = cli.command {
-        server::flightsql::try_run(cli.clone(), cfg.clone()).await?;
-        return Ok(());
+    {
+        if matches!(cli.command, Some(Command::FlightSql { .. })) {
+            cli::try_run(cli, cfg).await?;
+            return Ok(());
+        } else if let Some(Command::ServeFlightSql { .. }) = cli.command {
+            server::flightsql::try_run(cli.clone(), cfg.clone()).await?;
+            return Ok(());
+        }
     }
+
     #[cfg(feature = "http")]
     {
         if let Some(Command::ServeHttp { .. }) = cli.command {

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/cli_cases/tpch.rs

@@ -73,7 +73,7 @@ async fn test_custom_config_with_s3() {
     );
     let config = config_builder.build("my_config.toml");
 
-    let a = Command::cargo_bin("dft")
+    Command::cargo_bin("dft")
         .unwrap()
         .arg("--config")
         .arg(config.path)

tests/extension_cases/flightsql.rs

@@ -18,7 +18,12 @@
 use std::{io::Read, time::Duration};
 
 use assert_cmd::Command;
-use datafusion_dft::test_utils::fixture::{TestFixture, TestFlightSqlServiceImpl};
+use datafusion_app::local::ExecutionContext;
+use datafusion_dft::{
+    execution::AppExecution,
+    server::flightsql::service::FlightSqlServiceImpl,
+    test_utils::fixture::{TestFixture, TestFlightSqlServiceImpl},
+};
 
 use crate::{
     cli_cases::{contains_str, sql_in_file},
@@ -603,24 +608,34 @@ async fn test_output_parquet() {
 
     let cloned_path = path.clone();
 
-    let sql = "SELECT 1".to_string();
-    Command::cargo_bin("dft")
-        .unwrap()
-        .arg("-c")
-        .arg(sql.clone())
-        .arg("-o")
-        .arg(cloned_path)
-        .assert()
-        .success();
+    tokio::task::spawn_blocking(|| {
+        let sql = "SELECT 1".to_string();
+        Command::cargo_bin("dft")
+            .unwrap()
+            .arg("-c")
+            .arg(sql.clone())
+            .arg("--flightsql")
+            .arg("-o")
+            .arg(cloned_path)
+            .timeout(Duration::from_secs(5))
+            .assert()
+            .success();
+    })
+    .await
+    .unwrap();
 
     let read_sql = format!("SELECT * FROM '{}'", path.to_str().unwrap());
 
-    let assert = Command::cargo_bin("dft")
-        .unwrap()
-        .arg("-c")
-        .arg(read_sql)
-        .assert()
-        .success();
+    let assert = tokio::task::spawn_blocking(|| {
+        Command::cargo_bin("dft")
+            .unwrap()
+            .arg("-c")
+            .arg(read_sql)
+            .assert()
+            .success()
+    })
+    .await
+    .unwrap();
 
     let expected = r#"
 +----------+
@@ -633,3 +648,65 @@ async fn test_output_parquet() {
 
     fixture.shutdown_and_wait().await;
 }
+
+#[tokio::test]
+async fn test_flightsql_query_command() {
+    let test_server = TestFlightSqlServiceImpl::new();
+    let fixture = TestFixture::new(test_server.service(), "127.0.0.1:50051").await;
+
+    let assert = tokio::task::spawn_blocking(|| {
+        let sql = "SELECT 1".to_string();
+        Command::cargo_bin("dft")
+            .unwrap()
+            .arg("flightsql")
+            .arg("statement-query")
+            .arg(sql.clone())
+            .timeout(Duration::from_secs(5))
+            .assert()
+            .success()
+    })
+    .await
+    .unwrap();
+
+    let expected = r#"
++----------+
+| Int64(1) |
++----------+
+| 1        |
++----------+"#;
+
+    assert.stdout(contains_str(expected));
+
+    fixture.shutdown_and_wait().await;
+}
+
+#[tokio::test]
+async fn test_flightsql_get_catalogs() {
+    let ctx = ExecutionContext::default();
+    let exec = AppExecution::new(ctx);
+    let test_server = FlightSqlServiceImpl::new(exec);
+    let fixture = TestFixture::new(test_server.service(), "127.0.0.1:50051").await;
+
+    let assert = tokio::task::spawn_blocking(|| {
+        Command::cargo_bin("dft")
+            .unwrap()
+            .arg("flightsql")
+            .arg("get-catalogs")
+            .timeout(Duration::from_secs(5))
+            .assert()
+            .success()
+    })
+    .await
+    .unwrap();
+
+    let expected = r#"
++---------------+
+| table_catalog |
++---------------+
+| datafusion    |
++---------------+"#;
+
+    assert.stdout(contains_str(expected));
+
+    fixture.shutdown_and_wait().await;
+}