fix: panic when record not found (#6)

Co-authored-by: CangSheng Shen <67856731+shencangsheng@users.noreply.github.com>

Author: caelansar

src-tauri/src/commands/query.rs

@@ -75,20 +75,12 @@ pub async fn fetch(
                 let _ = insert_query_history(&app, &sql, "fail");
                 err
             })?;
-        let records = collect(&mut context, &new_sql).await.map_err(|err| {
+
+        let (header, records) = collect(&mut context, &new_sql).await.map_err(|err| {
             let _ = insert_query_history(&app, &sql, "fail");
             err
         })?;
-
-        let row = &records[0];
-        let width = row.columns().len();
-
-        let header: Vec<String> = row
-            .schema()
-            .fields()
-            .iter()
-            .map(|c| c.name().to_string())
-            .collect();
+        let width = header.len();
 
         // Pre-calculate the total number of rows to avoid frequent reallocation
         let total_rows: usize = records.iter().map(|r| r.num_rows()).sum();

src-tauri/src/context/context.rs

@@ -17,8 +17,6 @@ use sqlparser::ast::{
     TableFunctionArgs, Value,
 };
 use std::collections::HashMap;
-use std::future::Future;
-use std::pin::Pin;
 use std::sync::Arc;
 
 pub fn get_sql_context() -> SessionContext {
@@ -29,12 +27,23 @@ pub async fn get_data_frame(ctx: &mut SessionContext, sql: &String) -> AppResult
     ctx.sql(sql).await.map_err(AppError::from)
 }
 
-pub async fn collect(ctx: &mut SessionContext, sql: &String) -> AppResult<Vec<RecordBatch>> {
-    get_data_frame(ctx, sql)
-        .await?
-        .collect()
-        .await
-        .map_err(AppError::from)
+pub async fn collect(
+    ctx: &mut SessionContext,
+    sql: &String,
+) -> AppResult<(Vec<String>, Vec<RecordBatch>)> {
+    let df = get_data_frame(ctx, sql).await?;
+
+    // Get header from DataFrame's schema (column information can be retrieved even if there's no data)
+    let header: Vec<String> = df
+        .schema()
+        .fields()
+        .iter()
+        .map(|f| f.name().to_string())
+        .collect();
+
+    let records = df.collect().await.map_err(AppError::from)?;
+
+    Ok((header, records))
 }
 
 pub fn get_csv_read_options(
@@ -296,12 +305,10 @@ pub async fn convert_table_name(
             for join in &mut table_with_joins.joins {
                 match &mut join.relation {
                     TableFactor::Derived { subquery, .. } => {
-                        table_count =
-                            convert_table_name(ctx, subquery, table_count).await?;
+                        table_count = convert_table_name(ctx, subquery, table_count).await?;
                     }
                     relation => {
-                        table_count =
-                            register_table(ctx, &mut join.relation, table_count).await?;
+                        table_count = register_table(ctx, &mut join.relation, table_count).await?;
                     }
                 }
             }

src/pages/notebook/notebook-middle/notebook-middle-data-result.tsx

@@ -145,8 +145,16 @@ function DataResult({ data, isLoading }: DataResultProps) {
           </table>
           <div className="flex flex-col items-center justify-center gap-2 text-default-400 pt-16">
             <FontAwesomeIcon icon={faTable} size="2x" />
-            <p className="font-medium text-default-500">No data available</p>
-            <p className="text-sm">Run a query to see results here</p>
+            {data.header.length > 0 ? (
+              <p className="font-medium text-default-500">No records found</p>
+            ) : (
+              <>
+                <p className="font-medium text-default-500">
+                  No data available
+                </p>
+                <p className="text-sm">Run a query to see results here</p>
+              </>
+            )}
           </div>
         </div>
       </div>