Merge pull request #19 from lewiszlw/fix-endpoint-fetch

Fetch data from each endpoint

Cargo.toml

@@ -7,7 +7,7 @@ repository = "https://github.com/sundy-li/arrow_cli"
 edition = "2021"
 license = "Apache-2.0"
 name = "arrow_cli"
-version = "0.1.1"
+version = "0.1.2"
 
 
 

src/helper.rs

@@ -66,7 +66,7 @@ impl Highlighter for CliHelper {
                     );
                     pos = start;
                 }
-                Span::None(start, _) => {
+                Span::None(start) => {
                     pos = start;
                 }
             }
@@ -104,12 +104,12 @@ impl Highlighter for CliHelper {
 enum Span {
     Keyword(usize, usize),
     Literal(usize, usize),
-    None(usize, usize),
+    None(usize),
 }
 
 fn find_last_word(line: &str, pos: usize) -> Span {
     if line.is_empty() {
-        return Span::None(0, 0);
+        return Span::None(0);
     }
     let mut pos = pos;
     if pos >= line.len() {
@@ -143,7 +143,7 @@ fn find_last_word(line: &str, pos: usize) -> Span {
     {
         Span::Literal(pos, end)
     } else {
-        Span::None(pos, end)
+        Span::None(pos)
     }
 }
 

src/session.rs

@@ -124,13 +124,16 @@ impl Session {
         let start = Instant::now();
         let mut stmt = self.client.prepare(query.to_string(), None).await?;
         let flight_info = stmt.execute().await?;
-        let ticket = flight_info.endpoint[0]
-            .ticket
-            .as_ref()
-            .ok_or_else(|| ArrowError::IpcError("Ticket is emtpy".to_string()))?;
-
-        let flight_data = self.client.do_get(ticket.clone()).await?;
-        let batches: Vec<RecordBatch> = flight_data.try_collect().await.unwrap();
+        let mut batches: Vec<RecordBatch> = Vec::new();
+        for endpoint in flight_info.endpoint {
+            let ticket = endpoint
+                .ticket
+                .as_ref()
+                .ok_or_else(|| ArrowError::IpcError("Ticket is emtpy".to_string()))?;
+            let flight_data = self.client.do_get(ticket.clone()).await?;
+            let result: Vec<RecordBatch> = flight_data.try_collect().await.unwrap();
+            batches.extend(result);
+        }
         if is_repl {
             let res = pretty_format_batches(batches.as_slice())?;