Separate CLI display logic

Author: alamb

src/app/mod.rs

@@ -24,6 +24,8 @@ use crate::{cli, ui};
 use color_eyre::eyre::eyre;
 use color_eyre::Result;
 use crossterm::event as ct;
+use datafusion::arrow::util::pretty::pretty_format_batches;
+use datafusion::execution::SendableRecordBatchStream;
 use datafusion::sql::parser::DFParser;
 use datafusion::sql::sqlparser::dialect::GenericDialect;
 use futures::FutureExt;
@@ -376,9 +378,8 @@ impl CliApp {
         let dialect = GenericDialect {};
         let statements = DFParser::parse_sql_with_dialect(sql, &dialect)?;
         for statement in statements {
-            self.execution
-                .execute_and_print_statement(statement)
-                .await?;
+            let stream = self.execution.execute_statement(statement).await?;
+            self.print_stream(stream).await;
         }
         Ok(())
     }
@@ -406,7 +407,7 @@ impl CliApp {
             if line.ends_with(';') {
                 // TODO: if the query errors, should we keep trying to execute
                 // the other queries in the file? That is what datafusion-cli does...
-                self.execution.execute_and_print_stream_sql(&query).await?;
+                self.execute_and_print_sql(&query).await?;
                 query.clear();
             } else {
                 query.push('\n');
@@ -416,9 +417,29 @@ impl CliApp {
         // run the last line(s) in file if the last statement doesn't contain ‘;’
         // ignore if it only consists of '\n'
         if query.contains(|c| c != '\n') {
-            self.execution.execute_and_print_stream_sql(&query).await?;
+            self.execute_and_print_sql(&query).await?;
         }
 
         Ok(())
     }
+
+    /// executes a sql statement and prints the result to stdout
+    pub async fn execute_and_print_sql(&self, sql: &str) -> Result<()> {
+        let stream = self.execution.execute_sql(sql).await?;
+        self.print_stream(stream).await;
+        Ok(())
+    }
+
+    /// Prints the stream to stdout
+    async fn print_stream(&self, mut stream: SendableRecordBatchStream) {
+        while let Some(maybe_batch) = stream.next().await {
+            match maybe_batch {
+                Ok(batch) => match pretty_format_batches(&[batch]) {
+                    Ok(d) => println!("{}", d),
+                    Err(e) => println!("Error formatting batch: {e}"),
+                },
+                Err(e) => println!("Error executing SQL: {e}"),
+            }
+        }
+    }
 }

src/execution/mod.rs

@@ -23,11 +23,12 @@ use std::time::Duration;
 use color_eyre::eyre::Result;
 use datafusion::execution::runtime_env::RuntimeEnv;
 use datafusion::execution::session_state::SessionStateBuilder;
-use datafusion::execution::TaskContext;
-use datafusion::physical_plan::{execute_stream, visit_execution_plan, ExecutionPlanVisitor};
+use datafusion::execution::{SendableRecordBatchStream, TaskContext};
+use datafusion::physical_plan::{
+    execute_stream, visit_execution_plan, ExecutionPlan, ExecutionPlanVisitor,
+};
 use datafusion::prelude::*;
 use datafusion::sql::parser::Statement;
-use datafusion::{arrow::util::pretty::pretty_format_batches, physical_plan::ExecutionPlan};
 #[cfg(feature = "deltalake")]
 use deltalake::delta_datafusion::DeltaTableFactory;
 use log::{error, info};
@@ -197,7 +198,7 @@ impl ExecutionContext {
                     }
                 }
             } else {
-                match self.execute_sql(sql, false).await {
+                match self.execute_sql_and_discard_results(sql).await {
                     Ok(_) => {
                         let elapsed = start.elapsed();
                         query.set_execution_time(elapsed);
@@ -215,53 +216,46 @@ impl ExecutionContext {
         Ok(())
     }
 
-    /// Executes the specified parsed DataFusion statement and discards the result
-    pub async fn execute_sql(&self, sql: &str, print: bool) -> Result<()> {
-        let df = self.session_ctx.sql(sql).await?;
-        self.execute_stream_dataframe(df, print).await
+    /// Executes the specified sql string, driving it to completion but discarding the results
+    pub async fn execute_sql_and_discard_results(
+        &self,
+        sql: &str,
+    ) -> datafusion::error::Result<()> {
+        let mut stream = self.execute_sql(sql).await?;
+        while let Some(maybe_batch) = stream.next().await {
+            match maybe_batch {
+                Ok(_) => {}
+                Err(e) => {
+                    return Err(e);
+                }
+            }
+        }
+        Ok(())
     }
 
-    /// Executes the specified parsed DataFusion statement and prints the result to stdout
-    pub async fn execute_and_print_statement(&self, statement: Statement) -> Result<()> {
+    /// Executes the specified sql string, returning the resulting Stream
+    pub async fn execute_sql(
+        &self,
+        sql: &str,
+    ) -> datafusion::error::Result<SendableRecordBatchStream> {
+        self.session_ctx.sql(sql).await?.execute_stream().await
+    }
+
+    /// Executes the specified parsed DataFusion statement, returning the resulting Stream
+    pub async fn execute_statement(
+        &self,
+        statement: Statement,
+    ) -> datafusion::error::Result<SendableRecordBatchStream> {
         let plan = self
             .session_ctx
             .state()
             .statement_to_plan(statement)
             .await?;
-        let df = self.session_ctx.execute_logical_plan(plan).await?;
-        self.execute_stream_dataframe(df, true).await
-    }
-
-    /// Executes the specified query and prints the result to stdout
-    pub async fn execute_and_print_stream_sql(&self, query: &str) -> Result<()> {
-        let df = self.session_ctx.sql(query).await?;
-        self.execute_stream_dataframe(df, true).await
-    }
-
-    pub async fn execute_stream_dataframe(&self, df: DataFrame, print: bool) -> Result<()> {
-        let physical_plan = df.create_physical_plan().await?;
-        let stream_cfg = SessionConfig::default();
-        let stream_task_ctx = TaskContext::default().with_session_config(stream_cfg);
-        let mut stream = execute_stream(physical_plan, stream_task_ctx.into()).unwrap();
-
-        while let Some(maybe_batch) = stream.next().await {
-            if print {
-                let batch = maybe_batch.unwrap();
-                let d = pretty_format_batches(&[batch]).unwrap();
-                println!("{}", d);
-            } else {
-                let _ = maybe_batch.unwrap();
-                info!("Discarding batch");
-            }
-        }
-        Ok(())
-    }
-
-    pub async fn show_catalog(&self) -> Result<()> {
-        let tables = self.session_ctx.sql("SHOW tables").await?.collect().await?;
-        let formatted = pretty_format_batches(&tables).unwrap();
-        println!("{}", formatted);
-        Ok(())
+        self.session_ctx
+            .execute_logical_plan(plan)
+            .await?
+            .execute_stream()
+            .await
     }
 }