Support gsql output in JSON format

I've talked to at least a few teams that are directly querying
the database over gsql and scraping its output.  To mirror our
gerrit stream-events output gsql can now produce its records as
JSON objects, making them more suitable for machine consumption.

Change-Id: Ib2812b60a5d77824a48d511c50f6d8c2b23c4190
Signed-off-by: Shawn O. Pearce <sop@google.com>

Documentation/cmd-gsql.txt

@@ -8,7 +8,9 @@ gerrit gsql - Administrative interface to active database
 SYNOPSIS
 --------
 [verse]
-'ssh' -p <port> <host> 'gerrit gsql'
+'ssh' -p <port> <host> 'gerrit gsql' \
+[\--format \{PRETTY | JSON\}] \
+[\-c QUERY]
 
 DESCRIPTION
 -----------
@@ -16,13 +18,25 @@ Provides interactive query support directly against the underlying
 SQL database used by the host Gerrit server.  All SQL statements
 are supported, including SELECT, UPDATE, INSERT, DELETE and ALTER.
 
+OPTIONS
+-------
+\--format::
+	Set the format records are output in.  In PRETTY (the
+	default) records are displayed in a tabular output suitable
+	for reading by a human on a sufficiently wide terminal.
+	In JSON mode records are output as JSON objects using the
+	column names as the property names, one object per line.
+
+-c::
+	Execute the single query statement supplied, and then exit.
+
 ACCESS
 ------
 Caller must be a member of the privileged 'Administrators' group.
 
 SCRIPTING
 ---------
-Intended for interactive use only.
+Intended for interactive use only, unless format is JSON.
 
 EXAMPLES
 --------

gerrit-pgm/src/main/java/com/google/gerrit/pgm/Gsql.java

@@ -24,6 +24,8 @@ import com.google.gerrit.sshd.commands.QueryShell;
 import com.google.gerrit.sshd.commands.QueryShell.Factory;
 import com.google.inject.Injector;
 
+import org.kohsuke.args4j.Option;
+
 import java.io.IOException;
 
 /** Run Gerrit's SQL query tool */
@@ -31,6 +33,12 @@ public class Gsql extends SiteProgram {
   private final LifecycleManager manager = new LifecycleManager();
   private Injector dbInjector;
 
+  @Option(name = "--format", usage = "Set output format")
+  private QueryShell.OutputFormat format = QueryShell.OutputFormat.PRETTY;
+
+  @Option(name = "-c", metaVar = "SQL QUERY", usage = "Query to execute")
+  private String query;
+
   @Override
   public int run() throws Exception {
     mustHaveValidSite();
@@ -47,7 +55,13 @@ public class Gsql extends SiteProgram {
         manager.stop();
       }
     });
-    shellFactory().create(System.in, System.out).run();
+    final QueryShell shell = shellFactory().create(System.in, System.out);
+    shell.setOutputFormat(format);
+    if (query != null) {
+      shell.execute(query);
+    } else {
+      shell.run();
+    }
     return 0;
   }
 

gerrit-sshd/src/main/java/com/google/gerrit/sshd/commands/AdminQueryShell.java

@@ -19,6 +19,7 @@ import com.google.gerrit.sshd.BaseCommand;
 import com.google.inject.Inject;
 
 import org.apache.sshd.server.Environment;
+import org.kohsuke.args4j.Option;
 
 /** Opens a query processor. */
 @AdminCommand
@@ -26,13 +27,25 @@ final class AdminQueryShell extends BaseCommand {
   @Inject
   private QueryShell.Factory factory;
 
+  @Option(name = "--format", usage = "Set output format")
+  private QueryShell.OutputFormat format = QueryShell.OutputFormat.PRETTY;
+
+  @Option(name = "-c", metaVar = "SQL QUERY", usage = "Query to execute")
+  private String query;
+
   @Override
   public void start(final Environment env) {
     startThread(new CommandRunnable() {
       @Override
       public void run() throws Exception {
         parseCommandLine();
-        factory.create(in, out).run();
+        final QueryShell shell = factory.create(in, out);
+        shell.setOutputFormat(format);
+        if (query != null) {
+          shell.execute(query);
+        } else {
+          shell.run();
+        }
       }
     });
   }

gerrit-sshd/src/main/java/com/google/gerrit/sshd/commands/QueryShell.java

@@ -16,6 +16,7 @@ package com.google.gerrit.sshd.commands;
 
 import com.google.gerrit.common.Version;
 import com.google.gerrit.reviewdb.ReviewDb;
+import com.google.gson.JsonObject;
 import com.google.gwtorm.client.OrmException;
 import com.google.gwtorm.client.SchemaFactory;
 import com.google.gwtorm.jdbc.JdbcSchema;
@@ -47,9 +48,14 @@ public class QueryShell {
     QueryShell create(@Assisted InputStream in, @Assisted OutputStream out);
   }
 
+  public static enum OutputFormat {
+    PRETTY, JSON;
+  }
+
   private final BufferedReader in;
   private final PrintWriter out;
   private final SchemaFactory<ReviewDb> dbFactory;
+  private OutputFormat outputFormat = OutputFormat.PRETTY;
 
   private ReviewDb db;
   private Connection connection;
@@ -65,6 +71,10 @@ public class QueryShell {
     this.out = new PrintWriter(new OutputStreamWriter(out, "UTF-8"));
   }
 
+  public void setOutputFormat(OutputFormat fmt) {
+    outputFormat = fmt;
+  }
+
   public void run() {
     try {
       db = dbFactory.open();
@@ -94,11 +104,41 @@ public class QueryShell {
     }
   }
 
+  public void execute(String query) {
+    try {
+      db = dbFactory.open();
+      try {
+        connection = ((JdbcSchema) db).getConnection();
+        connection.setAutoCommit(true);
+
+        statement = connection.createStatement();
+        try {
+          executeStatement(query);
+        } finally {
+          statement.close();
+          statement = null;
+        }
+      } finally {
+        db.close();
+        db = null;
+      }
+    } catch (OrmException err) {
+      out.println("fatal: Cannot open connection: " + err.getMessage());
+
+    } catch (SQLException err) {
+      out.println("fatal: Cannot open connection: " + err.getMessage());
+    } finally {
+      out.flush();
+    }
+  }
+
   private void readEvalPrintLoop() {
     final StringBuilder buffer = new StringBuilder();
     boolean executed = false;
     for (;;) {
+      if (outputFormat == OutputFormat.PRETTY) {
         print(buffer.length() == 0 || executed ? "gerrit> " : "     -> ");
+      }
       String line = readLine();
       if (line == null) {
         return;
@@ -112,7 +152,9 @@ public class QueryShell {
           showHelp();
 
         } else if (line.equals("q")) {
+          if (outputFormat == OutputFormat.PRETTY) {
             println("Bye");
+          }
           return;
 
         } else if (line.equals("r")) {
@@ -135,9 +177,22 @@ public class QueryShell {
           showTable(line.substring(2).trim());
 
         } else {
-          println("ERROR: '\\" + line + "' not supported");
+          final String msg = "'\\" + line + "' not supported";
+          switch (outputFormat) {
+            case JSON: {
+              final JsonObject err = new JsonObject();
+              err.addProperty("type", "error");
+              err.addProperty("message", msg);
+              println(err.toString());
+              break;
+            }
+            case PRETTY:
+            default:
+              println("ERROR: " + msg);
               println("");
               showHelp();
+              break;
+          }
         }
         continue;
       }
@@ -171,7 +226,9 @@ public class QueryShell {
       final String[] types = {"TABLE", "VIEW"};
       ResultSet rs = meta.getTables(null, null, null, types);
       try {
+        if (outputFormat == OutputFormat.PRETTY) {
           println("                     List of relations");
+        }
         showResultSet(rs, false, //
             Identity.create(rs, "TABLE_SCHEM"), //
             Identity.create(rs, "TABLE_NAME"), //
@@ -208,7 +265,9 @@ public class QueryShell {
           throw new SQLException("Table " + tableName + " not found");
         }
 
+        if (outputFormat == OutputFormat.PRETTY) {
           println("                     Table " + tableName);
+        }
         showResultSet(rs, true, //
             Identity.create(rs, "COLUMN_NAME"), //
             new Function("TYPE") {
@@ -275,11 +334,13 @@ public class QueryShell {
           }
         }
 
+        if (outputFormat == OutputFormat.PRETTY) {
           println("");
           println("Indexes on " + tableName + ":");
           for (IndexInfo def : indexes.values()) {
             println("  " + def);
           }
+        }
       } finally {
         rs.close();
       }
@@ -307,9 +368,22 @@ public class QueryShell {
         try {
           final int rowCount = showResultSet(rs, false);
           final long ms = System.currentTimeMillis() - start;
+          switch (outputFormat) {
+            case JSON: {
+              final JsonObject tail = new JsonObject();
+              tail.addProperty("type", "query-stats");
+              tail.addProperty("rowCount", rowCount);
+              tail.addProperty("runTimeMilliseconds", ms);
+              println(tail.toString());
+              break;
+            }
+
+            case PRETTY:
+            default:
               println("(" + rowCount + (rowCount == 1 ? " row" : " rows") //
                   + "; " + ms + " ms)");
-          println("");
+              break;
+          }
         } finally {
           rs.close();
         }
@@ -317,7 +391,21 @@ public class QueryShell {
       } else {
         final int updateCount = statement.getUpdateCount();
         final long ms = System.currentTimeMillis() - start;
+        switch (outputFormat) {
+          case JSON: {
+            final JsonObject tail = new JsonObject();
+            tail.addProperty("type", "update-stats");
+            tail.addProperty("rowCount", updateCount);
+            tail.addProperty("runTimeMilliseconds", ms);
+            println(tail.toString());
+            break;
+          }
+
+          case PRETTY:
+          default:
             println("UPDATE " + updateCount + "; " + ms + " ms");
+            break;
+        }
       }
     } catch (SQLException e) {
       error(e);
@@ -326,6 +414,55 @@ public class QueryShell {
 
   private int showResultSet(final ResultSet rs, boolean alreadyOnRow,
       Function... show) throws SQLException {
+    switch (outputFormat) {
+      case JSON:
+        return showResultSetJson(rs, alreadyOnRow, show);
+      case PRETTY:
+      default:
+        return showResultSetPretty(rs, alreadyOnRow, show);
+    }
+  }
+
+  private int showResultSetJson(final ResultSet rs, boolean alreadyOnRow,
+      Function... show) throws SQLException {
+    final ResultSetMetaData meta = rs.getMetaData();
+    final Function[] columnMap;
+    if (show != null && 0 < show.length) {
+      columnMap = show;
+
+    } else {
+      final int colCnt = meta.getColumnCount();
+      columnMap = new Function[colCnt];
+      for (int colId = 0; colId < colCnt; colId++) {
+        final int p = colId + 1;
+        final String name = meta.getColumnLabel(p);
+        columnMap[colId] = new Identity(p, name);
+      }
+    }
+
+    int rowCnt = 0;
+    final int colCnt = columnMap.length;
+    while (alreadyOnRow || rs.next()) {
+      final JsonObject row = new JsonObject();
+      final JsonObject cols = new JsonObject();
+      for (int c = 0; c < colCnt; c++) {
+        String v = columnMap[c].apply(rs);
+        if (v == null) {
+          continue;
+        }
+        cols.addProperty(columnMap[c].name.toLowerCase(), v);
+      }
+      row.addProperty("type", "row");
+      row.add("columns", cols);
+      println(row.toString());
+      alreadyOnRow = false;
+      rowCnt++;
+    }
+    return rowCnt;
+  }
+
+  private int showResultSetPretty(final ResultSet rs, boolean alreadyOnRow,
+      Function... show) throws SQLException {
     final ResultSetMetaData meta = rs.getMetaData();
 
     final Function[] columnMap;
@@ -427,11 +564,37 @@ public class QueryShell {
   }
 
   private void warning(final String msg) {
+    switch (outputFormat) {
+      case JSON: {
+        final JsonObject obj = new JsonObject();
+        obj.addProperty("type", "warning");
+        obj.addProperty("message", msg);
+        println(obj.toString());
+        break;
+      }
+
+      case PRETTY:
+      default:
         println("WARNING: " + msg);
+        break;
+    }
   }
 
   private void error(final SQLException err) {
+    switch (outputFormat) {
+      case JSON: {
+        final JsonObject obj = new JsonObject();
+        obj.addProperty("type", "error");
+        obj.addProperty("message", err.getMessage());
+        println(obj.toString());
+        break;
+      }
+
+      case PRETTY:
+      default:
         println("ERROR: " + err.getMessage());
+        break;
+    }
   }
 
   private void print(String s) {
@@ -454,6 +617,7 @@ public class QueryShell {
   }
 
   private void showBanner() {
+    if (outputFormat == OutputFormat.PRETTY) {
       println("Welcome to Gerrit Code Review " + Version.getVersion());
       try {
         print("(");
@@ -468,6 +632,7 @@ public class QueryShell {
       println("Type '\\h' for help.  Type '\\r' to clear the buffer.");
       println("");
     }
+  }
 
   private void showHelp() {
     final StringBuilder help = new StringBuilder();