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

@@ -284,13 +284,18 @@ 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
+
+        if self.mode in ('csv', 'json'):
+            self.color = False
 
         if isinstance(auth_provider, PlainTextAuthProvider):
             self.username = auth_provider.username
@@ -956,7 +961,8 @@ def perform_simple_statement(self, statement):
     def print_result(self, result, table_meta):
         self.decoding_errors = []
 
-        self.writeresult("")
+        if self.mode not in ('csv', 'json'):
+            self.writeresult("")
 
         def print_all(result, table_meta, tty):
             # Return the number of rows in total
@@ -981,7 +987,8 @@ def print_all(result, table_meta, tty):
             return num_rows
 
         num_rows = print_all(result, table_meta, self.tty)
-        self.writeresult("(%d rows)" % num_rows)
+        if self.mode not in ('csv', 'json'):
+            self.writeresult("(%d rows)" % num_rows)
 
         if self.decoding_errors:
             for err in self.decoding_errors[:2]:
@@ -1009,7 +1016,11 @@ 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:
+        if self.mode == 'csv':
+            self.print_csv_result(formatted_names, formatted_values, with_header, tty)
+        elif self.mode == 'json':
+            self.print_json_result(formatted_names, formatted_values, with_header, tty)
+        elif 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)
@@ -1055,6 +1066,42 @@ def print_formatted_result_vertically(self, formatted_names, formatted_values, r
                 self.writeresult(' ' + " | ".join([column, value]))
             self.writeresult('')
 
+    def print_csv_result(self, formatted_names, formatted_values, with_header, tty):
+        import csv
+        writer = csv.writer(self.query_out)
+
+        if with_header:
+            headers = [hdr.strval for hdr in formatted_names]
+            writer.writerow(headers)
+
+        if formatted_values is None:
+            return
+
+        for row in formatted_values:
+            csv_row = [col.strval for col in row]
+            writer.writerow(csv_row)
+
+    def print_json_result(self, formatted_names, formatted_values, with_header, tty):
+        import json
+
+        if formatted_values is None:
+            self.writeresult("[]")
+            return
+
+        headers = [hdr.strval for hdr in formatted_names]
+        result_list = []
+
+        for row in formatted_values:
+            row_dict = {}
+            for i, col in enumerate(row):
+                row_dict[headers[i]] = col.strval
+            result_list.append(row_dict)
+
+        self.writeresult(json.dumps(result_list, indent=2))
+
+        if tty:
+            self.writeresult("")
+
     def print_warnings(self, warnings):
         if warnings is None or len(warnings) == 0:
             return
@@ -2026,6 +2073,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 +2278,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 +2407,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/test/test_cqlsh_output.py

@@ -1017,3 +1017,35 @@ def test_quoted_output_text_in_udts(self):
                                         tty=False, input=query)
         self.assertEqual(0, result)
         self.assertEqual(output.splitlines()[3].strip(), "{data: 'I''m newb'}")
+
+    def test_csv_output(self):
+        ks = get_keyspace()
+        query = "SELECT a, b FROM twenty_rows_table WHERE a IN ('1', '2');"
+
+        output, result = cqlsh_testcall(args=('--mode', 'csv'), prompt=None, env=self.default_env,
+                                        tty=False, input=query + '\n')
+        self.assertEqual(0, result)
+
+        lines = output.strip().splitlines()
+        self.assertEqual(lines[0].strip(), 'a,b')
+        self.assertIn('1,1', [l.strip() for l in lines])
+        self.assertIn('2,2', [l.strip() for l in lines])
+
+    def test_json_output(self):
+        ks = get_keyspace()
+        query = "SELECT a, b FROM twenty_rows_table WHERE a IN ('1', '2');"
+
+        output, result = cqlsh_testcall(args=('--mode', 'json'), prompt=None, env=self.default_env,
+                                        tty=False, input=query + '\n')
+        self.assertEqual(0, result)
+
+        import json
+        try:
+            parsed_json = json.loads(output)
+            self.assertEqual(len(parsed_json), 2)
+
+            results = { (item['a'], item['b']) for item in parsed_json }
+            self.assertIn(('1', '1'), results)
+            self.assertIn(('2', '2'), results)
+        except ValueError as e:
+            self.fail("Output is not valid JSON: %s\nOutput was:\n%s" % (e, output))