From 9251b8116ff89b528b6b9eaa43d4dc2d1bc0bbaf Mon Sep 17 00:00:00 2001
From: yifan-c <yc25code@gmail.com>
Date: Tue, 30 Jun 2020 00:15:18 -0700
Subject: [PATCH] Fix cqlsh output when fetching all rows in batch mode

Patch by Yifan Cai, reviewed by brandonwilliams for CASSANDRA-15905
---
 CHANGES.txt               |  1 +
 bin/cqlsh.py              | 44 +++++++++++++++++++++++++--------------
 pylib/cqlshlib/tracing.py |  2 +-
 3 files changed, 30 insertions(+), 17 deletions(-)

diff --git a/CHANGES.txt b/CHANGES.txt
index 9b4cf55b1f..d89d22b084 100644
--- a/CHANGES.txt
+++ b/CHANGES.txt
@@ -1,4 +1,5 @@
 3.11.7
+ * Fix cqlsh output when fetching all rows in batch mode (CASSANDRA-15905)
  * Upgrade Jackson to 2.9.10 (CASSANDRA-15867)
  * Fix CQL formatting of read command restrictions for slow query log (CASSANDRA-15503)
  * Allow sstableloader to use SSL on the native port (CASSANDRA-14904)
diff --git a/bin/cqlsh.py b/bin/cqlsh.py
index 2e5e2d4056..44d4d50d11 100644
--- a/bin/cqlsh.py
+++ b/bin/cqlsh.py
@@ -1079,7 +1079,7 @@ class Shell(cmd.Cmd):
         elif result:
             # CAS INSERT/UPDATE
             self.writeresult("")
-            self.print_static_result(result, self.parse_for_update_meta(statement.query_string))
+            self.print_static_result(result, self.parse_for_update_meta(statement.query_string), with_header=True, tty=self.tty)
         self.flush_output()
         return True, future
 
@@ -1087,20 +1087,30 @@ class Shell(cmd.Cmd):
         self.decoding_errors = []
 
         self.writeresult("")
-        if result.has_more_pages and self.tty:
+
+        def print_all(result, table_meta, tty):
+            # Return the number of rows in total
             num_rows = 0
+            isFirst = True
             while True:
-                if result.current_rows:
+                # Always print for the first page even it is empty
+                if result.current_rows or isFirst:
                     num_rows += len(result.current_rows)
-                    self.print_static_result(result, table_meta)
+                    with_header = isFirst or tty
+                    self.print_static_result(result, table_meta, with_header, tty)
                 if result.has_more_pages:
-                    raw_input("---MORE---")
+                    if self.shunted_query_out is None and tty:
+                        # Only pause when not capturing.
+                        raw_input("---MORE---")
                     result.fetch_next_page()
                 else:
+                    if not tty:
+                        self.writeresult("")
                     break
-        else:
-            num_rows = len(result.current_rows)
-            self.print_static_result(result, table_meta)
+                isFirst = False
+            return num_rows
+
+        num_rows = print_all(result, table_meta, self.tty)
         self.writeresult("(%d rows)" % num_rows)
 
         if self.decoding_errors:
@@ -1110,7 +1120,7 @@ class Shell(cmd.Cmd):
                 self.writeresult('%d more decoding errors suppressed.'
                                  % (len(self.decoding_errors) - 2), color=RED)
 
-    def print_static_result(self, result, table_meta):
+    def print_static_result(self, result, table_meta, with_header, tty):
         if not result.column_names and not table_meta:
             return
 
@@ -1118,7 +1128,7 @@ class Shell(cmd.Cmd):
         formatted_names = [self.myformat_colname(name, table_meta) for name in column_names]
         if not result.current_rows:
             # print header only
-            self.print_formatted_result(formatted_names, None)
+            self.print_formatted_result(formatted_names, None, with_header=True, tty=tty)
             return
 
         cql_types = []
@@ -1132,9 +1142,9 @@ class Shell(cmd.Cmd):
         if self.expand_enabled:
             self.print_formatted_result_vertically(formatted_names, formatted_values)
         else:
-            self.print_formatted_result(formatted_names, formatted_values)
+            self.print_formatted_result(formatted_names, formatted_values, with_header, tty)
 
-    def print_formatted_result(self, formatted_names, formatted_values):
+    def print_formatted_result(self, formatted_names, formatted_values, with_header, tty):
         # determine column widths
         widths = [n.displaywidth for n in formatted_names]
         if formatted_values is not None:
@@ -1143,9 +1153,10 @@ class Shell(cmd.Cmd):
                     widths[num] = max(widths[num], col.displaywidth)
 
         # print header
-        header = ' | '.join(hdr.ljust(w, color=self.color) for (hdr, w) in zip(formatted_names, widths))
-        self.writeresult(' ' + header.rstrip())
-        self.writeresult('-%s-' % '-+-'.join('-' * w for w in widths))
+        if with_header:
+            header = ' | '.join(hdr.ljust(w, color=self.color) for (hdr, w) in zip(formatted_names, widths))
+            self.writeresult(' ' + header.rstrip())
+            self.writeresult('-%s-' % '-+-'.join('-' * w for w in widths))
 
         # stop if there are no rows
         if formatted_values is None:
@@ -1157,7 +1168,8 @@ class Shell(cmd.Cmd):
             line = ' | '.join(col.rjust(w, color=self.color) for (col, w) in zip(row, widths))
             self.writeresult(' ' + line)
 
-        self.writeresult("")
+        if tty:
+            self.writeresult("")
 
     def print_formatted_result_vertically(self, formatted_names, formatted_values):
         max_col_width = max([n.displaywidth for n in formatted_names])
diff --git a/pylib/cqlshlib/tracing.py b/pylib/cqlshlib/tracing.py
index 26e228b222..5b55367c6b 100644
--- a/pylib/cqlshlib/tracing.py
+++ b/pylib/cqlshlib/tracing.py
@@ -52,7 +52,7 @@ def print_trace(shell, trace):
     shell.writeresult('Tracing session: ', color=MAGENTA, newline=False)
     shell.writeresult(trace.trace_id)
     shell.writeresult('')
-    shell.print_formatted_result(formatted_names, formatted_values)
+    shell.print_formatted_result(formatted_names, formatted_values, with_header=True, tty=shell.tty)
     shell.writeresult('')