Fix FlightSQL client to display all batches from multi-partition queries

Previously, the FlightSQL client would only display partial results when
queries returned multiple batches (common with GROUP BY without ORDER BY
due to partitioned execution). This happened because:

1. `run_flightsqls()` only fetched one batch using `match streams.next().await`
   instead of iterating through all batches
2. `take_record_batches()` only used the first batch when multiple batches
   existed, ignoring the rest

This fix:
- Changes `run_flightsqls()` to use `while let` loop to collect all batches
  from the FlightSQL stream
- Updates `take_record_batches()` to concatenate all batches before applying
  pagination indices, ensuring all rows are visible

Known limitations for future consideration:
- All batches are collected eagerly into memory before display, which could
  be problematic for very large result sets
- `concat_batches()` creates a temporary copy when concatenating, briefly
  doubling memory usage during pagination
- Unlike the SQL tab which uses lazy batch loading, FlightSQL now loads
  all data upfront (this matches expected behavior for aggregation queries
  but may warrant lazy loading for large streaming results in the future)

Fixes #346

Author: Matthew Agostinelli

src/tui/execution.rs

@@ -279,37 +279,42 @@ impl TuiExecution {
                                             if let Some(streams) =
                                                 self.flightsql_result_stream.lock().await.as_mut()
                                             {
-                                                match streams.next().await {
-                                                    Some((ticket, Ok(batch))) => {
-                                                        info!("Received batch for {ticket}");
-                                                        let duration = start.elapsed();
-                                                        let results = ExecutionResultsBatch {
-                                                            batch,
-                                                            duration,
-                                                            query: sql.to_string(),
-                                                        };
-                                                        sender.send(
-                                                            AppEvent::FlightSQLExecutionResultsNextBatch(
-                                                                results,
-                                                            ),
-                                                        )?;
+                                                // Collect all batches from the stream
+                                                while let Some((ticket, result)) =
+                                                    streams.next().await
+                                                {
+                                                    match result {
+                                                        Ok(batch) => {
+                                                            info!("Received batch for {ticket}");
+                                                            let duration = start.elapsed();
+                                                            let results = ExecutionResultsBatch {
+                                                                batch,
+                                                                duration,
+                                                                query: sql.to_string(),
+                                                            };
+                                                            sender.send(
+                                                                AppEvent::FlightSQLExecutionResultsNextBatch(
+                                                                    results,
+                                                                ),
+                                                            )?;
+                                                        }
+                                                        Err(e) => {
+                                                            error!(
+                                                                "Error executing stream for ticket {ticket}: {:?}",
+                                                                e
+                                                            );
+                                                            let elapsed = start.elapsed();
+                                                            let e = ExecutionError {
+                                                                query: sql.to_string(),
+                                                                error: e.to_string(),
+                                                                duration: elapsed,
+                                                            };
+                                                            sender.send(
+                                                                AppEvent::FlightSQLExecutionResultsError(e),
+                                                            )?;
+                                                            break;
+                                                        }
                                                     }
-                                                    Some((ticket, Err(e))) => {
-                                                        error!(
-                                                            "Error executing stream for ticket {ticket}: {:?}",
-                                                            e
-                                                        );
-                                                        let elapsed = start.elapsed();
-                                                        let e = ExecutionError {
-                                                            query: sql.to_string(),
-                                                            error: e.to_string(),
-                                                            duration: elapsed,
-                                                        };
-                                                        sender.send(
-                                                            AppEvent::FlightSQLExecutionResultsError(e),
-                                                        )?;
-                                                    }
-                                                    None => {}
                                                 }
                                             }
                                         }

src/tui/state/tabs/flightsql.rs

@@ -21,7 +21,7 @@ use std::sync::Arc;
 use color_eyre::Result;
 use datafusion::arrow::{
     array::{Array, RecordBatch, UInt32Array},
-    compute::take_record_batch,
+    compute::{concat_batches, take_record_batch},
     datatypes::Schema,
     error::ArrowError,
 };
@@ -316,7 +316,11 @@ fn take_record_batches(
     match batches.len() {
         0 => Ok(RecordBatch::new_empty(Arc::new(Schema::empty()))),
         1 => take_record_batch(&batches[0], indices),
-        // For now we just get the first batch
-        _ => take_record_batch(&batches[0], indices),
+        _ => {
+            // Concatenate all batches into a single batch before taking indices
+            let schema = batches[0].schema();
+            let concatenated = concat_batches(&schema, batches)?;
+            take_record_batch(&concatenated, indices)
+        }
     }
 }