fix buttom_rows != 0 cause panic

databendlabs · Jun 16, 2023 · 5b226cd · 5b226cd
1 parent a8b54d7
commit 5b226cd
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 7 deletions.
diff --git a/cli/Cargo.toml b/cli/Cargo.toml
@@ -39,6 +39,7 @@ tokio = { version = "1.28.1", features = [
     "parking_lot",
 ] }
 toml = "0.7.3"
+unicode-segmentation = "1.10.1"
 url = { version = "2.3.1", default-features = false }
 
 [build-dependencies]

diff --git a/cli/src/display.rs b/cli/src/display.rs
@@ -14,6 +14,7 @@
 
 use std::collections::HashSet;
 use std::fmt::Write;
+use unicode_segmentation::UnicodeSegmentation;
 
 use anyhow::Result;
 use comfy_table::{Cell, CellAlignment, Table};
@@ -417,7 +418,7 @@ fn create_table(
         (top_rows, rows_to_render - top_rows)
     };
 
-    let mut res_vec: Vec<Vec<String>> = Vec::with_capacity(top_rows + bottom_rows);
+    let mut res_vec: Vec<Vec<String>> = vec![];
     for row in results.iter().take(top_rows) {
         let values = row.values();
         let mut v = vec![];
@@ -428,11 +429,13 @@ fn create_table(
     }
 
     if bottom_rows != 0 {
-        for (idx, row) in results.iter().skip(row_count - bottom_rows).enumerate() {
+        for row in results.iter().skip(row_count - bottom_rows) {
             let values = row.values();
+            let mut v = vec![];
             for value in values {
-                res_vec[idx].push(value.to_string());
+                v.push(value.to_string());
             }
+            res_vec.push(v);
         }
     }
 
@@ -475,7 +478,16 @@ fn create_table(
                 } else {
                     let mut value = values[*col_index as usize].clone();
                     if value.len() > widths[idx] {
-                        value = value[0..widths[idx] - 3].to_string() + "..."
+                        value = String::from_utf8(
+                            value
+                                .graphemes(true)
+                                .take(widths[idx] - 3)
+                                .flat_map(|g| g.as_bytes().iter())
+                                .copied() // copied converts &u8 into u8
+                                .chain(b"...".iter().copied())
+                                .collect::<Vec<u8>>(),
+                        )
+                        .unwrap();
                     }
                     let cell = Cell::new(value).set_alignment(aligns[idx]);
                     cells.push(cell);
@@ -490,6 +502,7 @@ fn create_table(
     if bottom_rows != 0 {
         // first render the divider
         let mut cells: Vec<Cell> = Vec::new();
+        let display_res_len = res_vec.len();
         for align in aligns.iter() {
             let cell = Cell::new("·").set_alignment(*align);
             cells.push(cell);
@@ -499,7 +512,7 @@ fn create_table(
             table.add_row(cells.clone());
         }
         if column_map.is_empty() {
-            for values in res_vec.iter().skip(row_count - bottom_rows) {
+            for values in res_vec.iter().skip(display_res_len - bottom_rows) {
                 let mut cells = Vec::new();
                 for (idx, align) in aligns.iter().enumerate() {
                     let cell = Cell::new(&values[idx]).set_alignment(*align);
@@ -508,7 +521,7 @@ fn create_table(
                 table.add_row(cells);
             }
         } else {
-            for values in res_vec.iter().skip(row_count - bottom_rows) {
+            for values in res_vec.iter().skip(display_res_len - bottom_rows) {
                 let mut cells = Vec::new();
                 for (idx, col_index) in column_map.iter().enumerate() {
                     if *col_index == -1 {
@@ -517,7 +530,16 @@ fn create_table(
                     } else {
                         let mut value = values[*col_index as usize].clone();
                         if value.len() > widths[idx] {
-                            value = value[0..widths[idx] - 3].to_string() + "...";
+                            value = String::from_utf8(
+                                value
+                                    .graphemes(true)
+                                    .take(widths[idx] - 3)
+                                    .flat_map(|g| g.as_bytes().iter())
+                                    .copied() // copied converts &u8 into u8
+                                    .chain(b"...".iter().copied())
+                                    .collect::<Vec<u8>>(),
+                            )
+                            .unwrap();
                         }
                         let cell = Cell::new(value).set_alignment(aligns[idx]);
                         cells.push(cell);