CASSANDRA-19985: Enhance CQLSH to support machine-readable output formatting (csv, json)

Author: arvindKandpal-ksolves

conf/cqlshrc.sample

@@ -37,6 +37,9 @@
 ; version = None
 
 [ui]
+;; The format of the output. Valid values are tabular, csv, and json.
+; mode = tabular
+
 ;; Whether or not to display query results with colors
 ; color = on
 

doc/modules/cassandra/pages/managing/tools/cqlsh.adoc

@@ -98,6 +98,8 @@ Options:
   Collect coverage data
 `--encoding=ENCODING`::
   Specify a non-default encoding for output. (Default: utf-8)
+`--mode=MODE`::
+  Specify the output display format. Valid values are `tabular` (default), `csv`, and `json`.
 `--cqlshrc=CQLSHRC`::
   Specify an alternative cqlshrc file location.
 `--credentials=CREDENTIALS`::

pylib/cqlshlib/cqlshmain.py

@@ -46,7 +46,8 @@
 from cqlshlib import cql3handling, pylexotron, sslhandling, cqlshhandling, authproviderhandling
 from cqlshlib.copyutil import ExportTask, ImportTask
 from cqlshlib.displaying import (ANSI_RESET, BLUE, COLUMN_NAME_COLORS, CYAN,
-                                 RED, WHITE, FormattedValue, colorme)
+                                 RED, WHITE, FormattedValue, colorme,
+                                 TablePrinter, TabularTablePrinter, CsvTablePrinter, JsonTablePrinter)
 from cqlshlib.formatting import (DEFAULT_DATE_FORMAT, DEFAULT_NANOTIME_FORMAT,
                                  DEFAULT_TIMESTAMP_FORMAT, CqlType, DateTimeFormat,
                                  format_by_type)
@@ -284,13 +285,15 @@ def __init__(self, hostname, port, config_file, color=False,
                  connect_timeout=DEFAULT_CONNECT_TIMEOUT_SECONDS,
                  is_subshell=False,
                  auth_provider=None,
-                 disable_history=False):
+                 disable_history=False,
+                 mode='tabular'):
         cmd.Cmd.__init__(self, completekey=completekey)
         self.hostname = hostname
         self.port = port
         self.auth_provider = auth_provider
         self.username = username
         self.config_file = config_file
+        self.mode = mode.lower()
 
         if isinstance(auth_provider, PlainTextAuthProvider):
             self.username = auth_provider.username
@@ -329,6 +332,8 @@ def __init__(self, hostname, port, config_file, color=False,
         self.browser = browser
         self.docspath = docspath
         self.color = color
+        if self.mode in ('csv', 'json'):
+            self.color = False
 
         self.display_nanotime_format = display_nanotime_format
         self.display_timestamp_format = display_timestamp_format
@@ -946,42 +951,55 @@ def perform_simple_statement(self, statement):
             self.print_result(result, self.get_table_meta('system_auth', 'generated_values'))
         elif result:
             # CAS INSERT/UPDATE
-            self.writeresult("")
-            self.print_static_result(result, self.parse_for_update_meta(statement.query_string), with_header=True, tty=self.tty)
+            if self.mode not in ('csv', 'json'):
+                self.writeresult("")
+            cas_printer = TablePrinter.factory(self.mode, self)
+            self.print_static_result(result, self.parse_for_update_meta(statement.query_string),
+                                     with_header=True, tty=self.tty,
+                                     printer=cas_printer)
+            cas_printer.finish()
         if self.elapsed_enabled:
-            self.writeresult("(%dms elapsed)" % elapsed)
+            elapsed_msg = "(%dms elapsed)" % elapsed
+            if self.mode in ('csv', 'json'):
+                self.printerr(elapsed_msg)
+            else:
+                self.writeresult(elapsed_msg)
         self.flush_output()
         return True, future
 
     def print_result(self, result, table_meta):
         self.decoding_errors = []
 
-        self.writeresult("")
+        if self.mode not in ('csv', 'json'):
+            self.writeresult("")
+        printer = TablePrinter.factory(self.mode, self)
 
-        def print_all(result, table_meta, tty):
-            # Return the number of rows in total
+        def print_all(result, table_meta, tty, printer):
+            machine_mode = self.mode in ('csv', 'json')
+            effective_tty = tty and not machine_mode
             num_rows = 0
             is_first = True
             while True:
-                # Always print for the first page even it is empty
                 if result.current_rows or is_first:
-                    with_header = is_first or tty
-                    self.print_static_result(result, table_meta, with_header, tty, num_rows)
+                    with_header = is_first or effective_tty
+                    self.print_static_result(result, table_meta, with_header, effective_tty,
+                                             num_rows, printer)
                     num_rows += len(result.current_rows)
                 if result.has_more_pages:
-                    if self.shunted_query_out is None and tty:
-                        # Only pause when not capturing.
+                    if self.shunted_query_out is None and effective_tty:
                         input("---MORE---")
                     result.fetch_next_page()
                 else:
-                    if not tty:
+                    if not effective_tty and not machine_mode:
                         self.writeresult("")
                     break
                 is_first = False
             return num_rows
 
-        num_rows = print_all(result, table_meta, self.tty)
-        self.writeresult("(%d rows)" % num_rows)
+        num_rows = print_all(result, table_meta, self.tty, printer)
+        printer.finish()
+        if self.mode not in ('csv', 'json'):
+            self.writeresult("(%d rows)" % num_rows)
 
         if self.decoding_errors:
             for err in self.decoding_errors[:2]:
@@ -990,15 +1008,16 @@ def print_all(result, table_meta, tty):
                 self.writeresult('%d more decoding errors suppressed.'
                                  % (len(self.decoding_errors) - 2), color=RED)
 
-    def print_static_result(self, result, table_meta, with_header, tty, row_count_offset=0):
+    def print_static_result(self, result, table_meta, with_header, tty, row_count_offset=0, printer=None):
         if not result.column_names and not table_meta:
             return
 
         column_names = result.column_names or list(table_meta.columns.keys())
         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, with_header=True, tty=tty)
+            if with_header:
+                printer.print_header(formatted_names)
             return
 
         cql_types = []
@@ -1009,10 +1028,9 @@ def print_static_result(self, result, table_meta, with_header, tty, row_count_of
 
         formatted_values = [list(map(self.myformat_value, [row[c] for c in column_names], cql_types)) for row in result.current_rows]
 
-        if self.expand_enabled:
-            self.print_formatted_result_vertically(formatted_names, formatted_values, row_count_offset)
-        else:
-            self.print_formatted_result(formatted_names, formatted_values, with_header, tty)
+        if with_header:
+            printer.print_header(formatted_names)
+        printer.print_rows(formatted_names, formatted_values)
 
     def print_formatted_result(self, formatted_names, formatted_values, with_header, tty):
         # determine column widths
@@ -2026,6 +2044,7 @@ def read_options(cmdlineargs, parser, config_file, cql_dir, environment=os.envir
     argvalues.completekey = option_with_default(configs.get, 'ui', 'completekey',
                                                 DEFAULT_COMPLETEKEY)
     argvalues.color = option_with_default(configs.getboolean, 'ui', 'color')
+    argvalues.mode = option_with_default(configs.get, 'ui', 'mode', 'tabular')
     argvalues.time_format = raw_option_with_default(configs, 'ui', 'time_format',
                                                     DEFAULT_TIMESTAMP_FORMAT)
     argvalues.nanotime_format = raw_option_with_default(configs, 'ui', 'nanotime_format',
@@ -2230,6 +2249,8 @@ def main(cmdline, pkgpath):
                         help='Force tty mode (command prompt).')
     parser.add_argument('--disable-history', default=False, action='store_true',
                         help='Disable saving of history (existing history will still be loaded)')
+    parser.add_argument('--mode', choices=['tabular', 'csv', 'json'],
+                        help='Specify the output format (tabular, csv, json). Default is tabular.')
 
     # This is a hidden option to suppress the warning when the -p/--password command line option is used.
     # Power users may use this option if they know no other people has access to the system where cqlsh is run or don't care about security.
@@ -2357,6 +2378,7 @@ def main(cmdline, pkgpath):
                       display_double_precision=options.double_precision,
                       display_timezone=timezone,
                       max_trace_wait=options.max_trace_wait,
+                      mode=options.mode,
                       ssl=options.ssl,
                       single_statement=options.execute,
                       request_timeout=options.request_timeout,

pylib/cqlshlib/displaying.py

@@ -126,3 +126,101 @@ def color_ljust(self, width, fill=' '):
                                  )
 
 NO_COLOR_MAP = dict()
+
+class TablePrinter:
+    def print_header(self, formatted_names):
+        raise NotImplementedError
+
+    def print_rows(self, formatted_names, formatted_values):
+        raise NotImplementedError
+
+    def finish(self):
+        pass
+
+    @staticmethod
+    def factory(format_type, shell):
+        format_map = {'csv': CsvTablePrinter, 'json': JsonTablePrinter, 'tabular': TabularTablePrinter}
+        printer_cls = format_map.get(format_type.lower(), TabularTablePrinter)
+        return printer_cls(shell) if format_type.lower() != 'tabular' else printer_cls(shell, shell.tty)
+
+class TabularTablePrinter(TablePrinter):
+    def __init__(self, shell, tty, row_count_offset=0):
+        self._shell = shell
+        self._tty = tty
+        self._row_count_offset = row_count_offset
+        self._pending_header = None
+
+    def print_header(self, formatted_names):
+        # Store only — cannot render yet because column widths depend on
+        # data values. print_rows will render header+data together.
+        # Empty-result case is handled in finish().
+        self._pending_header = formatted_names
+
+    def print_rows(self, formatted_names, formatted_values):
+        # with_header=True only when print_header was called for this page.
+        with_header = self._pending_header is not None
+        self._pending_header = None
+        if self._shell.expand_enabled:
+            self._shell.print_formatted_result_vertically(
+                formatted_names, formatted_values, self._row_count_offset)
+        else:
+            self._shell.print_formatted_result(
+                formatted_names, formatted_values, with_header, self._tty)
+        if formatted_values:
+            self._row_count_offset += len(formatted_values)
+
+    def finish(self):
+        if self._pending_header is not None:
+            self._shell.print_formatted_result(
+                self._pending_header, None, with_header=True, tty=self._tty)
+            self._pending_header = None
+
+class CsvTablePrinter(TablePrinter):
+    def __init__(self, shell):
+        import csv
+        self._writer = csv.writer(shell.query_out)
+        self._header_written = False
+        self._colnames = None
+
+    def print_header(self, formatted_names):
+        self._colnames = [n.strval for n in formatted_names]
+
+    def print_rows(self, formatted_names, formatted_values):
+        if not self._header_written:
+            self._writer.writerow(self._colnames)
+            self._header_written = True
+        if formatted_values is None:
+            return
+        for row in formatted_values:
+            self._writer.writerow([col.strval for col in row])
+
+    def finish(self):
+        if self._colnames is not None and not self._header_written:
+            self._writer.writerow(self._colnames)
+            self._header_written = True
+
+class JsonTablePrinter(TablePrinter):
+    def __init__(self, shell):
+        self._shell = shell
+        self._colnames = None
+        self._first_row = True
+
+    def print_header(self, formatted_names):
+        self._colnames = [n.strval for n in formatted_names]
+        self._shell.writeresult('[')
+
+    def print_rows(self, formatted_names, formatted_values):
+        import json
+        if formatted_values is None:
+            return
+        for row in formatted_values:
+            row_dict = {self._colnames[i]: col.strval for i, col in enumerate(row)}
+            serialized = json.dumps(row_dict)
+            if self._first_row:
+                self._shell.writeresult('  ' + serialized, newline=False)
+                self._first_row = False
+            else:
+                self._shell.writeresult(',\n  ' + serialized, newline=False)
+
+    def finish(self):
+        self._shell.writeresult('\n]')