Test for query command

Author: matthewmturner

src/cli/mod.rs

@@ -85,7 +85,11 @@ impl CliApp {
     }
 
     #[cfg(feature = "flightsql")]
-    fn handle_flightsql_command(&self, command: FlightSqlCommand) {}
+    async fn handle_flightsql_command(&self, command: FlightSqlCommand) -> color_eyre::Result<()> {
+        match command {
+            FlightSqlCommand::Query { sql } => self.exec_from_flightsql(sql, 0).await,
+        }
+    }
 
     /// Execute the provided sql, which was passed as an argument from CLI.
     ///
@@ -98,9 +102,9 @@ impl CliApp {
         self.validate_args()?;
 
         #[cfg(feature = "flightsql")]
-        if let Some(Command::FlightSql { command }) = self.args.command {
-            self.handle_flightsql_command(command)
-        }
+        if let Some(Command::FlightSql { command }) = &self.args.command {
+            return self.handle_flightsql_command(command.clone()).await;
+        };
 
         #[cfg(not(feature = "flightsql"))]
         match (
@@ -613,7 +617,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

@@ -91,7 +91,10 @@ async fn app_entry_point(cli: DftArgs) -> Result<()> {
         }
     }
 
-    if !cli.files.is_empty() || !cli.commands.is_empty() {
+    if !cli.files.is_empty()
+        || !cli.commands.is_empty()
+        || matches!(cli.command, Some(Command::FlightSql { .. }))
+    {
         cli::try_run(cli, cfg).await?;
     } else {
         tui::try_run(cli, cfg).await?;

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

@@ -608,6 +608,7 @@ async fn test_output_parquet() {
         .unwrap()
         .arg("-c")
         .arg(sql.clone())
+        .arg("--flightsql")
         .arg("-o")
         .arg(cloned_path)
         .assert()
@@ -633,3 +634,34 @@ 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("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;
+}