Skip to content

Commit 273b846

Browse files
matthewmturnermatthewmturner
committed
Fix
1 parent 6f1d93e commit 273b846

File tree

8 files changed

+524
-42
lines changed

8 files changed

+524
-42
lines changed

‎src/tui/handlers/mod.rs‎

Lines changed: 38 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -28,8 +28,6 @@ use tui_logger::TuiWidgetEvent;
2828
use crate::tui::state::tabs::flightsql::FlightSQLConnectionStatus;
2929
use crate::tui::state::tabs::history::Context;
3030
use crate::tui::ExecutionResultsBatch;
31-
32-
#[cfg(feature = "flightsql")]
3331
use std::sync::Arc;
3432

3533
use super::App;
@@ -222,13 +220,47 @@ pub fn app_event_handler(app: &mut App, event: AppEvent) -> Result<()> {
222220
duration,
223221
batch,
224222
} = r;
223+
224+
let is_first_batch = app.state.sql_tab.current_page().is_none();
225+
225226
app.state.sql_tab.add_batch(batch);
227+
228+
if is_first_batch {
229+
app.state.sql_tab.next_page();
230+
app.state.sql_tab.refresh_query_results_state();
231+
232+
let history_query =
233+
HistoryQuery::new(Context::Local, query.to_string(), duration, None, None);
234+
app.state.history_tab.add_to_history(history_query);
235+
app.state.history_tab.refresh_history_table_state();
236+
} else {
237+
app.state.sql_tab.refresh_query_results_state();
238+
239+
// Check if we have enough data for the next page now
240+
// If not, automatically fetch another batch
241+
if let Some(current_page) = app.state.sql_tab.current_page() {
242+
let next_page = current_page + 1;
243+
if app.state.sql_tab.needs_more_batches_for_page(next_page) {
244+
info!("Still need more batches for page {}, fetching next batch", next_page);
245+
let execution = Arc::clone(&app.execution);
246+
let sql = query.clone();
247+
let _event_tx = app.event_tx();
248+
tokio::spawn(async move {
249+
execution.next_batch(sql, _event_tx).await;
250+
});
251+
} else {
252+
// We now have enough data, advance to the page
253+
info!("Sufficient data loaded, advancing to page {}", next_page);
254+
if let Err(e) = app.event_tx().send(AppEvent::ExecutionResultsNextPage) {
255+
error!("Error advancing to next page: {e}");
256+
}
257+
}
258+
}
259+
}
260+
}
261+
AppEvent::ExecutionResultsNextPage => {
226262
app.state.sql_tab.next_page();
227263
app.state.sql_tab.refresh_query_results_state();
228-
let history_query =
229-
HistoryQuery::new(Context::Local, query.to_string(), duration, None, None);
230-
app.state.history_tab.add_to_history(history_query);
231-
app.state.history_tab.refresh_history_table_state();
232264
}
233265
#[cfg(feature = "flightsql")]
234266
AppEvent::FlightSQLExecutionResultsNextBatch(r) => {

‎src/tui/handlers/sql.rs‎

Lines changed: 21 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -92,25 +92,29 @@ pub fn normal_mode_handler(app: &mut App, key: KeyEvent) {
9292
}
9393
},
9494
(KeyCode::Right, KeyModifiers::NONE) => {
95-
let _event_tx = app.event_tx().clone();
96-
if let (Some(p), c) = (
97-
app.state().sql_tab.current_page(),
98-
app.state().sql_tab.batches_count(),
99-
) {
100-
// We don't need to fetch the next batch if moving forward a page and we're not
101-
// on the last page since we would have already fetched it.
102-
if p < c - 1 {
103-
app.state.sql_tab.next_page();
104-
app.state.sql_tab.refresh_query_results_state();
105-
return;
95+
let _event_tx = app.event_tx();
96+
97+
if let Some(current_page) = app.state.sql_tab.current_page() {
98+
let next_page = current_page + 1;
99+
100+
// Check if we need more batches for the next page
101+
if app.state.sql_tab.needs_more_batches_for_page(next_page) {
102+
info!("Fetching more batches for page {}", next_page);
103+
104+
if let Some(last_query) = app.state.history_tab.history().last() {
105+
let execution = Arc::clone(&app.execution);
106+
let sql = last_query.sql().clone();
107+
tokio::spawn(async move {
108+
execution.next_batch(sql, _event_tx).await;
109+
});
110+
}
111+
return; // Wait for batch to load before advancing page
106112
}
107113
}
108-
if let Some(p) = app.state.history_tab.history().last() {
109-
let execution = Arc::clone(&app.execution);
110-
let sql = p.sql().clone();
111-
tokio::spawn(async move {
112-
execution.next_batch(sql, _event_tx).await;
113-
});
114+
115+
// Sufficient data available, advance page
116+
if let Err(e) = app.event_tx().send(AppEvent::ExecutionResultsNextPage) {
117+
error!("Error going to next SQL results page: {e}");
114118
}
115119
}
116120
(KeyCode::Left, KeyModifiers::NONE) => {

‎src/tui/mod.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -72,6 +72,7 @@ pub enum AppEvent {
7272
// Query Execution
7373
NewExecution,
7474
ExecutionResultsNextBatch(ExecutionResultsBatch),
75+
ExecutionResultsNextPage,
7576
ExecutionResultsPreviousPage,
7677
ExecutionResultsError(ExecutionError),
7778
// FlightSQL

‎src/tui/state/tabs/sql.rs‎

Lines changed: 24 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@ use tokio::task::JoinHandle;
2929
use tui_textarea::TextArea;
3030

3131
use crate::config::AppConfig;
32+
use crate::tui::pagination::{extract_page, has_sufficient_rows, PAGE_SIZE};
3233
use crate::tui::ExecutionError;
3334

3435
pub fn get_keywords() -> Vec<String> {
@@ -241,19 +242,33 @@ impl SQLTabState<'_> {
241242
}
242243
}
243244

244-
pub fn current_batch(&self) -> Option<&RecordBatch> {
245+
pub fn current_page_results(&self) -> Option<RecordBatch> {
246+
use std::sync::Arc;
247+
use datafusion::arrow::datatypes::Schema;
248+
245249
match (self.current_page, self.result_batches.as_ref()) {
246-
(Some(page), Some(batches)) => batches.get(page),
247-
_ => None,
250+
(Some(page), Some(batches)) => {
251+
match extract_page(batches, page, PAGE_SIZE) {
252+
Ok(batch) => Some(batch),
253+
Err(err) => {
254+
log::error!("Error getting page {}: {}", page, err);
255+
None
256+
}
257+
}
258+
}
259+
_ => Some(RecordBatch::new_empty(Arc::new(Schema::empty()))),
248260
}
249261
}
250262

251-
pub fn batches_count(&self) -> usize {
252-
if let Some(batches) = &self.result_batches {
253-
batches.len()
254-
} else {
255-
0
256-
}
263+
pub fn total_loaded_rows(&self) -> usize {
264+
self.result_batches
265+
.as_ref()
266+
.map(|batches| batches.iter().map(|b| b.num_rows()).sum())
267+
.unwrap_or(0)
268+
}
269+
270+
pub fn needs_more_batches_for_page(&self, page: usize) -> bool {
271+
!has_sufficient_rows(self.total_loaded_rows(), page, PAGE_SIZE)
257272
}
258273

259274
pub fn execution_error(&self) -> &Option<ExecutionError> {

‎src/tui/ui/tabs/sql.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -58,7 +58,7 @@ pub fn render_sql_results(area: Rect, buf: &mut Buffer, app: &App) {
5858
// TODO: Change this to a match on state and batch
5959
let sql_tab = &app.state.sql_tab;
6060
match (
61-
sql_tab.current_batch(),
61+
sql_tab.current_page_results(),
6262
sql_tab.current_page(),
6363
sql_tab.query_results_state(),
6464
sql_tab.execution_error(),
@@ -68,7 +68,7 @@ pub fn render_sql_results(area: Rect, buf: &mut Buffer, app: &App) {
6868
.title(" Results ")
6969
.borders(Borders::ALL)
7070
.title_top(Line::from(format!(" Page {p} ")).right_aligned());
71-
let batches = vec![batch];
71+
let batches = vec![&batch];
7272
let maybe_table = record_batches_to_table(&batches);
7373

7474
let block = block.title_bottom("Stats").fg(tailwind::ORANGE.c500);

‎tests/tui_cases/mod.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@ mod keyboard_modifiers;
2626
mod pagination;
2727
mod quit;
2828
mod sql_execution;
29+
mod sql_pagination;
2930

3031
use datafusion::arrow::array::RecordBatch;
3132
use datafusion::common::Result;

‎tests/tui_cases/pagination.rs‎

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -55,11 +55,11 @@ async fn single_page() {
5555
let page = state.sql_tab.current_page().unwrap();
5656
assert_eq!(page, 0);
5757

58-
let batch = state.sql_tab.current_batch();
58+
let batch = state.sql_tab.current_page_results();
5959
assert!(batch.is_some());
6060

6161
let batch = batch.unwrap();
62-
let batches = vec![batch.clone()];
62+
let batches = vec![batch];
6363
let expected = [
6464
"+---+---+",
6565
"| a | b |",
@@ -103,11 +103,11 @@ async fn multiple_pages_forward_and_back() {
103103

104104
{
105105
let state = test_app.state();
106-
let batch = state.sql_tab.current_batch();
106+
let batch = state.sql_tab.current_page_results();
107107
assert!(batch.is_some());
108108

109109
let batch = batch.unwrap();
110-
let batches = vec![batch.clone()];
110+
let batches = vec![batch];
111111
let expected = [
112112
"+---+---+",
113113
"| a | b |",
@@ -134,11 +134,11 @@ async fn multiple_pages_forward_and_back() {
134134

135135
{
136136
let state = test_app.state();
137-
let batch = state.sql_tab.current_batch();
137+
let batch = state.sql_tab.current_page_results();
138138
assert!(batch.is_some());
139139

140140
let batch = batch.unwrap();
141-
let batches = vec![batch.clone()];
141+
let batches = vec![batch];
142142
let expected = [
143143
"+---+---+",
144144
"| a | b |",
@@ -194,11 +194,11 @@ async fn multiple_pages_forward_and_back_and_forward() {
194194

195195
{
196196
let state = test_app.state();
197-
let batch = state.sql_tab.current_batch();
197+
let batch = state.sql_tab.current_page_results();
198198
assert!(batch.is_some());
199199

200200
let batch = batch.unwrap();
201-
let batches = vec![batch.clone()];
201+
let batches = vec![batch];
202202
let expected = [
203203
"+---+---+",
204204
"| a | b |",

0 commit comments

Comments
 (0)