Add --flightsql-host param to cli args (#246)

Author: matthewmturner

README.md

@@ -92,7 +92,7 @@ $ dft -c "SELECT 1+2"
 
 #### FlightSQL
 
-Both of the commands above support the `--flightsql` parameter to run the SQL with your configured FlightSQL client.
+Both of the commands above support the `--flightsql` parameter to run the SQL with your configured FlightSQL client.  You can also configure the host used for creating FlightSQL client per command with `--flightsql-host` - for example `--flightsql-host "http://127.0.0.1:50052"`.
 
 #### DDL
 

src/args.rs

@@ -90,6 +90,10 @@ pub struct DftArgs {
 
     #[clap(short = 'n', help = "Set the number of benchmark iterations to run")]
     pub benchmark_iterations: Option<usize>,
+
+    #[cfg(feature = "flightsql")]
+    #[clap(long, help = "Set the host and port to be used for FlightSQL")]
+    pub flightsql_host: Option<String>,
 }
 
 impl DftArgs {

src/execution/flightsql.rs

@@ -49,8 +49,9 @@ impl FlightSQLContext {
     }
 
     /// Create FlightSQL client from users FlightSQL config
-    pub async fn create_client(&self) -> Result<()> {
-        let url = Box::leak(self.config.connection_url.clone().into_boxed_str());
+    pub async fn create_client(&self, cli_host: Option<String>) -> Result<()> {
+        let final_url = cli_host.unwrap_or(self.config.connection_url.clone());
+        let url = Box::leak(final_url.into_boxed_str());
         info!("Connecting to FlightSQL host: {}", url);
         let channel = Channel::from_static(url).connect().await;
         match channel {

src/main.rs

@@ -46,11 +46,13 @@ async fn main() -> Result<()> {
         {
             if cli.flightsql {
                 let flightsql_ctx = FlightSQLContext::new(state.config.flightsql.clone());
-                flightsql_ctx.create_client().await?;
+                flightsql_ctx
+                    .create_client(cli.flightsql_host.clone())
+                    .await?;
                 app_execution.with_flightsql_ctx(flightsql_ctx);
             }
         }
-        let app = CliApp::new(app_execution, cli);
+        let app = CliApp::new(app_execution, cli.clone());
         app.execute_files_or_commands().await?;
     } else if cli.serve {
         #[cfg(feature = "experimental-flightsql-server")]
@@ -66,7 +68,12 @@ async fn main() -> Result<()> {
             }
             let app_execution = AppExecution::new(execution_ctx);
             let server = FlightSqlServiceImpl::new(app_execution);
-            let app = FlightSqlApp::new(server.service(), DEFAULT_SERVER_ADDRESS).await;
+            let app = FlightSqlApp::new(
+                server.service(),
+                &cli.flightsql_host
+                    .unwrap_or(DEFAULT_SERVER_ADDRESS.to_string()),
+            )
+            .await;
             app.run_app().await;
         }
 
@@ -82,7 +89,7 @@ async fn main() -> Result<()> {
         let state = state::initialize(cli.config_path());
         let execution_ctx = ExecutionContext::try_new(&state.config.execution, AppType::Tui)?;
         let app_execution = AppExecution::new(execution_ctx);
-        let app = App::new(state, app_execution);
+        let app = App::new(state, cli, app_execution);
         app.run_app().await?;
     }
 

src/tui/execution.rs

@@ -403,10 +403,9 @@ impl TuiExecution {
 
     // TODO: Maybe just expose `inner` and use that rather than re-implementing the same
     // functions here.
-
     #[cfg(feature = "flightsql")]
-    pub async fn create_flightsql_client(&self) -> Result<()> {
-        self.inner.flightsql_ctx().create_client().await
+    pub async fn create_flightsql_client(&self, cli_host: Option<String>) -> Result<()> {
+        self.inner.flightsql_ctx().create_client(cli_host).await
     }
 
     #[cfg(feature = "flightsql")]

src/tui/handlers/mod.rs

@@ -245,8 +245,9 @@ pub fn app_event_handler(app: &mut App, event: AppEvent) -> Result<()> {
         AppEvent::FlightSQLEstablishConnection => {
             let execution = Arc::clone(&app.execution);
             let _event_tx = app.event_tx.clone();
+            let cli_url = app.args.flightsql_host.clone();
             tokio::spawn(async move {
-                if let Err(e) = execution.create_flightsql_client().await {
+                if let Err(e) = execution.create_flightsql_client(cli_url).await {
                     error!("Error creating FlightSQL client: {:?}", e);
                     if let Err(e) = _event_tx.send(AppEvent::FlightSQLFailedToConnect) {
                         error!("Error sending FlightSQLFailedToConnect message: {e}");

src/tui/mod.rs

@@ -40,6 +40,7 @@ use tokio_util::sync::CancellationToken;
 
 use self::execution::{ExecutionError, ExecutionResultsBatch, TuiExecution};
 use self::handlers::{app_event_handler, crossterm_event_handler};
+use crate::args::DftArgs;
 use crate::execution::sql_utils::clean_sql;
 use crate::execution::AppExecution;
 
@@ -92,18 +93,19 @@ pub struct App<'app> {
     cancellation_token: CancellationToken,
     task: JoinHandle<()>,
     ddl_task: Option<JoinHandle<()>>,
+    args: DftArgs,
 }
 
 impl<'app> App<'app> {
-    pub fn new(state: state::AppState<'app>, execution: AppExecution) -> Self {
+    pub fn new(state: state::AppState<'app>, args: DftArgs, execution: AppExecution) -> Self {
         let (event_tx, event_rx) = mpsc::unbounded_channel();
         let cancellation_token = CancellationToken::new();
         let task = tokio::spawn(async {});
-        // let ddl_task = tokio::spawn(async {});
         let app_execution = Arc::new(TuiExecution::new(Arc::new(execution)));
 
         Self {
             state,
+            args,
             task,
             event_rx,
             event_tx,

tests/extension_cases/flightsql.rs

@@ -495,3 +495,34 @@ SELECT 1
     assert.stdout(contains_str(expected));
     fixture.shutdown_and_wait().await;
 }
+
+#[tokio::test]
+pub async fn test_execute_custom_port() {
+    let test_server = TestFlightSqlServiceImpl::new();
+    let fixture = TestFixture::new(test_server.service(), "127.0.0.1:50052").await;
+
+    let assert = tokio::task::spawn_blocking(|| {
+        Command::cargo_bin("dft")
+            .unwrap()
+            .arg("-c")
+            .arg("SELECT 1 + 2;")
+            .arg("--flightsql")
+            .arg("--flightsql-host")
+            .arg("http://localhost:50052")
+            .timeout(Duration::from_secs(5))
+            .assert()
+            .success()
+    })
+    .await
+    .unwrap();
+
+    let expected = r##"
++---------------------+
+| Int64(1) + Int64(2) |
++---------------------+
+| 3                   |
++---------------------+
+    "##;
+    assert.stdout(contains_str(expected));
+    fixture.shutdown_and_wait().await;
+}

tests/tui_cases/mod.rs

@@ -24,6 +24,7 @@ mod quit;
 use datafusion::arrow::array::RecordBatch;
 use datafusion::common::Result;
 use dft::{
+    args::DftArgs,
     execution::{local::ExecutionContext, AppExecution, AppType},
     tui::{state::initialize, App, AppEvent},
 };
@@ -51,7 +52,8 @@ impl<'app> TestApp<'app> {
         let execution_ctx =
             ExecutionContext::try_new(&state.config.execution, AppType::Tui).unwrap();
         let app_execution = AppExecution::new(execution_ctx);
-        let mut app = App::new(state, app_execution);
+        let args = DftArgs::default();
+        let mut app = App::new(state, args, app_execution);
         app.enter(false).unwrap();
         Self { config_path, app }
     }

tests/tui_cases/quit.rs

@@ -17,6 +17,7 @@
 
 //! Tests for the TUI (e.g. user application with keyboard commands)
 
+use dft::args::DftArgs;
 use dft::execution::{local::ExecutionContext, AppExecution, AppType};
 use dft::tui::state::initialize;
 use dft::tui::{App, AppEvent};
@@ -105,8 +106,9 @@ impl<'app> TestApp<'app> {
         let config_path = tempdir().unwrap();
         let state = initialize(config_path.path().to_path_buf());
         let execution = ExecutionContext::try_new(&state.config.execution, AppType::Tui).unwrap();
+        let args = DftArgs::default();
         let app_execution = AppExecution::new(execution);
-        let app = App::new(state, app_execution);
+        let app = App::new(state, args, app_execution);
         Self { config_path, app }
     }