Separate CLI display logic

alamb · alamb · commit 44857d154d74 · 2024-09-12T09:15:32.000-04:00
diff --git a/src/app/app_execution.rs b/src/app/app_execution.rs
@@ -0,0 +1,110 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! [`AppExecution`]: Handles executing queries for the TUI application.
+
+use crate::app::state::tabs::sql::Query;
+use crate::app::AppEvent;
+use crate::execution::ExecutionContext;
+use color_eyre::eyre::Result;
+use futures::StreamExt;
+use log::{error, info};
+use std::sync::Arc;
+use std::time::Duration;
+use tokio::sync::mpsc::UnboundedSender;
+
+/// Handles executing queries for the TUI application, formatting results
+/// and sending them to the UI.
+pub(crate) struct AppExecution {
+    inner: Arc<ExecutionContext>,
+}
+
+impl AppExecution {
+    /// Create a new instance of [`AppExecution`].
+    pub fn new(inner: Arc<ExecutionContext>) -> Self {
+        Self { inner }
+    }
+
+    /// Run the sequence of SQL queries, sending the results as [`AppEvent::QueryResult`] via the sender.
+    ///
+    /// All queries except the last one will have their results discarded.
+    ///
+    /// Error handling: If an error occurs while executing a query, the error is
+    /// logged and execution continues
+    pub async fn run_sqls(&self, sqls: Vec<&str>, sender: UnboundedSender<AppEvent>) -> Result<()> {
+        // We need to filter out empty strings to correctly determine the last query for displaying
+        // results.
+        info!("Running sqls: {:?}", sqls);
+        let non_empty_sqls: Vec<&str> = sqls.into_iter().filter(|s| !s.is_empty()).collect();
+        info!("Non empty SQLs: {:?}", non_empty_sqls);
+        let statement_count = non_empty_sqls.len();
+        for (i, sql) in non_empty_sqls.into_iter().enumerate() {
+            info!("Running query {}", i);
+            let _sender = sender.clone();
+            let mut query =
+                Query::new(sql.to_string(), None, None, None, Duration::default(), None);
+            let start = std::time::Instant::now();
+            if i == statement_count - 1 {
+                info!("Executing last query and display results");
+                match self.inner.execute_sql(sql).await {
+                    Ok(mut stream) => {
+                        let mut batches = Vec::new();
+                        while let Some(maybe_batch) = stream.next().await {
+                            match maybe_batch {
+                                Ok(batch) => {
+                                    batches.push(batch);
+                                }
+                                Err(e) => {
+                                    let elapsed = start.elapsed();
+                                    query.set_error(Some(e.to_string()));
+                                    query.set_execution_time(elapsed);
+                                    break;
+                                }
+                            }
+                        }
+                        let elapsed = start.elapsed();
+                        let rows: usize = batches.iter().map(|r| r.num_rows()).sum();
+                        query.set_results(Some(batches));
+                        query.set_num_rows(Some(rows));
+                        query.set_execution_time(elapsed);
+                    }
+                    Err(e) => {
+                        error!("Error creating dataframe: {:?}", e);
+                        let elapsed = start.elapsed();
+                        query.set_error(Some(e.to_string()));
+                        query.set_execution_time(elapsed);
+                    }
+                }
+            } else {
+                match self.inner.execute_sql_and_discard_results(sql).await {
+                    Ok(_) => {
+                        let elapsed = start.elapsed();
+                        query.set_execution_time(elapsed);
+                    }
+                    Err(e) => {
+                        // We only log failed queries, we don't want to stop the execution of the
+                        // remaining queries. Perhaps there should be a configuration option for
+                        // this though in case the user wants to stop execution on the first error.
+                        error!("Error executing {sql}: {:?}", e);
+                    }
+                }
+            }
+            _sender.send(AppEvent::QueryResult(query))?; // Send the query result to the UI
+        }
+        Ok(())
+    }
+}
diff --git a/src/app/handlers/sql.rs b/src/app/handlers/sql.rs
@@ -23,6 +23,7 @@ use ratatui::crossterm::event::{KeyCode, KeyEvent, KeyModifiers};
 use tokio_stream::StreamExt;
 
 use super::App;
+use crate::app::app_execution::AppExecution;
 use crate::app::{handlers::tab_navigation_handler, state::tabs::sql::Query, AppEvent};
 use crate::execution::collect_plan_stats;
 
@@ -64,14 +65,14 @@ pub fn normal_mode_handler(app: &mut App, key: KeyEvent) {
             info!("Run query");
             let sql = app.state.sql_tab.editor().lines().join("");
             info!("SQL: {}", sql);
-            let execution = Arc::clone(&app.execution);
+            let app_execution = AppExecution::new(Arc::clone(&app.execution));
             let _event_tx = app.app_event_tx.clone();
             // TODO: Maybe this should be on a separate runtime to prevent blocking main thread /
             // runtime
             // TODO: Extract this into function to be used in both normal and editable handler
             tokio::spawn(async move {
                 let sqls: Vec<&str> = sql.split(';').collect();
-                let _ = execution.run_sqls(sqls, _event_tx).await;
+                let _ = app_execution.run_sqls(sqls, _event_tx).await;
             });
         }
         _ => {}
diff --git a/src/app/mod.rs b/src/app/mod.rs
@@ -15,6 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
+pub mod app_execution;
 pub mod config;
 pub mod handlers;
 pub mod state;
@@ -24,6 +25,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 +379,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 +408,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 +418,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}"),
+            }
+        }
+    }
 }
diff --git a/src/execution/mod.rs b/src/execution/mod.rs
@@ -18,20 +18,17 @@
 //! [`ExecutionContext`]: DataFusion based execution context for running SQL queries
 //!
 use std::sync::Arc;
-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;
+use datafusion::physical_plan::{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};
-use tokio::sync::mpsc::UnboundedSender;
+use log::info;
 use tokio_stream::StreamExt;
 #[cfg(feature = "s3")]
 use url::Url;
@@ -42,8 +39,6 @@ use {
 };
 
 use crate::app::config::ExecutionConfig;
-use crate::app::state::tabs::sql::Query;
-use crate::app::AppEvent;
 
 /// Structure for executing queries either locally or remotely (via FlightSQL)
 ///
@@ -97,7 +92,7 @@ impl ExecutionContext {
                                 }
                             }
                             Err(e) => {
-                                error!("Error creating object store: {:?}", e);
+                                log::error!("Error creating object store: {:?}", e);
                             }
                         }
                     }
@@ -145,123 +140,44 @@ impl ExecutionContext {
         &self.flightsql_client
     }
 
-    pub async fn run_sqls(&self, sqls: Vec<&str>, sender: UnboundedSender<AppEvent>) -> Result<()> {
-        // We need to filter out empty strings to correctly determine the last query for displaying
-        // results.
-        info!("Running sqls: {:?}", sqls);
-        let non_empty_sqls: Vec<&str> = sqls.into_iter().filter(|s| !s.is_empty()).collect();
-        info!("Non empty SQLs: {:?}", non_empty_sqls);
-        let statement_count = non_empty_sqls.len();
-        for (i, sql) in non_empty_sqls.into_iter().enumerate() {
-            info!("Running query {}", i);
-            let _sender = sender.clone();
-            let mut query =
-                Query::new(sql.to_string(), None, None, None, Duration::default(), None);
-            let start = std::time::Instant::now();
-            if i == statement_count - 1 {
-                info!("Executing last query and display results");
-                match self.session_ctx.sql(sql).await {
-                    Ok(df) => {
-                        if let Ok(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();
-                            let mut batches = Vec::new();
-                            while let Some(maybe_batch) = stream.next().await {
-                                match maybe_batch {
-                                    Ok(batch) => {
-                                        batches.push(batch);
-                                    }
-                                    Err(e) => {
-                                        let elapsed = start.elapsed();
-                                        query.set_error(Some(e.to_string()));
-                                        query.set_execution_time(elapsed);
-                                        break;
-                                    }
-                                }
-                            }
-                            let elapsed = start.elapsed();
-                            let rows: usize = batches.iter().map(|r| r.num_rows()).sum();
-                            query.set_results(Some(batches));
-                            query.set_num_rows(Some(rows));
-                            query.set_execution_time(elapsed);
-                        }
-                    }
-                    Err(e) => {
-                        error!("Error creating dataframe: {:?}", e);
-                        let elapsed = start.elapsed();
-                        query.set_error(Some(e.to_string()));
-                        query.set_execution_time(elapsed);
-                    }
-                }
-            } else {
-                match self.execute_sql(sql, false).await {
-                    Ok(_) => {
-                        let elapsed = start.elapsed();
-                        query.set_execution_time(elapsed);
-                    }
-                    Err(e) => {
-                        // We only log failed queries, we don't want to stop the execution of the
-                        // remaining queries. Perhaps there should be a configuration option for
-                        // this though in case the user wants to stop execution on the first error.
-                        error!("Error executing {sql}: {:?}", e);
-                    }
-                }
-            }
-            _sender.send(AppEvent::QueryResult(query)).unwrap();
+    /// Executes the specified sql string, driving it to completion but discarding any results
+    pub async fn execute_sql_and_discard_results(
+        &self,
+        sql: &str,
+    ) -> datafusion::error::Result<()> {
+        let mut stream = self.execute_sql(sql).await?;
+        // note we don't call collect() to avoid buffering data
+        while let Some(maybe_batch) = stream.next().await {
+            maybe_batch?; // check for errors
         }
         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, returning the resulting
+    /// [`SendableRecordBatchStream`] of results
+    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 and prints the result to stdout
-    pub async fn execute_and_print_statement(&self, statement: Statement) -> Result<()> {
+    /// Executes the a pre-parsed DataFusion [`Statement`], returning the
+    /// resulting [`SendableRecordBatchStream`] of results
+    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
     }
 }
 
