Test working

Author: matthewmturner

crates/datafusion-app/src/local.rs

@@ -27,7 +27,7 @@ use futures::TryFutureExt;
 use log::{debug, error, info};
 
 use crate::catalog::create_app_catalog;
-use crate::config::ExecutionConfig;
+use crate::config::{ExecutionConfig, ObservabilityConfig};
 use crate::{ExecOptions, ExecResult};
 use color_eyre::eyre::{self, Result};
 use datafusion::common::Result as DFResult;
@@ -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

@@ -132,7 +132,7 @@ pub enum FlightSqlCommand {
     /// Executes `GetFlightInfo` and `DoGet` on the provided SQL query
     StatementQuery { sql: String },
     /// Executes `GetCatalogsFlightInfo` and `DoGet`
-    Catalogs,
+    GetCatalogs,
 }
 
 #[derive(Clone, Debug, Subcommand)]

src/cli/mod.rs

@@ -16,7 +16,6 @@
 // under the License.
 //! [`CliApp`]: Command Line User Interface
 
-use crate::args::Command;
 use crate::config::AppConfig;
 use crate::db::register_db;
 use crate::{args::DftArgs, execution::AppExecution};
@@ -40,8 +39,7 @@ use std::io::Write;
 use std::path::{Path, PathBuf};
 #[cfg(feature = "flightsql")]
 use {
-    crate::args::FlightSqlCommand,
-    arrow_flight::FlightInfo,
+    crate::args::{Command, FlightSqlCommand},
     datafusion_app::{
         config::{AuthConfig, FlightSQLConfig},
         flightsql::FlightSQLContext,
@@ -87,14 +85,24 @@ impl CliApp {
 
     #[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::Catalogs => {
+            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(())
             }
         }

src/main.rs

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

src/server/flightsql/service.rs

@@ -229,6 +229,7 @@ impl FlightSqlService for FlightSqlServiceImpl {
         let start = Timestamp::now();
         let request_id = uuid::Uuid::new_v4();
         let query = "SELECT DISTINCT table_catalog FROM information_schema.tables".to_string();
+        // 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;

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},
@@ -665,3 +670,34 @@ async fn test_flightsql_query_command() {
 
     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;
+}