Stream results of autocommit queries...

When executing queries in autocommit mode, results are no longer
materialized before printing. This was causing issues for big queries
that would materialize as OOM crashes or just generally slow execution.

For --format=PLAIN this means each records is printed directly as it
arrived in the shell.

For --format=VERBOSE, we cannot simply stream, because we need to know
the widths of the columns before we can print even the column header. To
achieve this we materialize the 1000 first rows, and compute the column
widths from this sample. Any remaining rows will be streamed. If any
remaining row contains a value that overflows it's column width, that
value will either be wrapped to the next row (--wrap=true), or truncated
(--wrap=false)

Author: fickludd

cypher-shell/src/main/java/org/neo4j/shell/CypherShell.java

@@ -34,7 +34,7 @@ public class CypherShell implements StatementExecuter, Connector, TransactionHan
     protected CommandHelper commandHelper;
 
     public CypherShell(@Nonnull Logger logger) {
-        this(logger, new BoltStateHandler(), new PrettyPrinter(logger.getFormat()));
+        this(logger, new BoltStateHandler(), new PrettyPrinter(logger.getFormat(), logger.getWrap(), logger.getNumSampleRows()));
     }
 
     protected CypherShell(@Nonnull Logger logger,
@@ -83,7 +83,7 @@ public void execute(@Nonnull final String cmdString) throws ExitException, Comma
      */
     protected void executeCypher(@Nonnull final String cypher) throws CommandException {
         final Optional<BoltResult> result = boltStateHandler.runCypher(cypher, allParameterValues());
-        result.ifPresent(boltResult -> logger.printOut(prettyPrinter.format(boltResult)));
+        result.ifPresent(boltResult -> prettyPrinter.format(boltResult, logger::printOut));
     }
 
     @Override
@@ -138,7 +138,7 @@ public void beginTransaction() throws CommandException {
     @Override
     public Optional<List<BoltResult>> commitTransaction() throws CommandException {
         Optional<List<BoltResult>> results = boltStateHandler.commitTransaction();
-        results.ifPresent(boltResult -> boltResult.forEach(result -> logger.printOut(prettyPrinter.format(result))));
+        results.ifPresent(boltResult -> boltResult.forEach(result -> prettyPrinter.format(result, logger::printOut)));
         return results;
     }
 

cypher-shell/src/main/java/org/neo4j/shell/Main.java

@@ -93,6 +93,8 @@ private Logger instantiateLogger(@Nonnull CliArgs cliArgs) {
         if (cliArgs.isStringShell() && Format.AUTO.equals(cliArgs.getFormat())) {
             logger.setFormat(Format.PLAIN);
         }
+        logger.setWrap(cliArgs.getWrap());
+        logger.setNumSampleRows(cliArgs.getNumSampleRows());
         return logger;
     }
 

cypher-shell/src/main/java/org/neo4j/shell/cli/CliArgHelper.java

@@ -5,11 +5,7 @@
 import net.sourceforge.argparse4j.impl.action.StoreTrueArgumentAction;
 import net.sourceforge.argparse4j.impl.choice.CollectionArgumentChoice;
 import net.sourceforge.argparse4j.impl.type.BooleanArgumentType;
-import net.sourceforge.argparse4j.inf.ArgumentGroup;
-import net.sourceforge.argparse4j.inf.ArgumentParser;
-import net.sourceforge.argparse4j.inf.ArgumentParserException;
-import net.sourceforge.argparse4j.inf.MutuallyExclusiveGroup;
-import net.sourceforge.argparse4j.inf.Namespace;
+import net.sourceforge.argparse4j.inf.*;
 
 import java.io.PrintWriter;
 import java.util.regex.Matcher;
@@ -88,6 +84,10 @@ public static CliArgs parse(@Nonnull String... args) {
 
         cliArgs.setNonInteractive(ns.getBoolean("force-non-interactive"));
 
+        cliArgs.setWrap(ns.getBoolean("wrap"));
+
+        cliArgs.setNumSampleRows(ns.getInt("sample-rows"));
+
         cliArgs.setVersion(ns.getBoolean("version"));
 
         cliArgs.setDriverVersion(ns.getBoolean("driver-version"));
@@ -167,6 +167,17 @@ private static ArgumentParser setupParser()
                 .dest("force-non-interactive")
               .action(new StoreTrueArgumentAction());
 
+        parser.addArgument("--sample-rows")
+                .help("number of rows sampled to compute table widths (only for format=VERBOSE)")
+                .type(new PositiveIntegerType())
+                .dest("sample-rows")
+                .setDefault(CliArgs.DEFAULT_NUM_SAMPLE_ROWS);
+
+        parser.addArgument("--wrap")
+                .help("wrap table colum values if column is too narrow (only for format=VERBOSE)")
+                .type(new BooleanArgumentType())
+                .setDefault(true);
+
         parser.addArgument("-v", "--version")
                 .help("print version of cypher-shell and exit")
                 .action(new StoreTrueArgumentAction());
@@ -183,5 +194,16 @@ private static ArgumentParser setupParser()
         return parser;
     }
 
-
+    private static class PositiveIntegerType implements ArgumentType<Integer> {
+        @Override
+        public Integer convert(ArgumentParser parser, Argument arg, String value) throws ArgumentParserException {
+            try {
+                int result = Integer.parseInt(value);
+                if (result < 1) throw new NumberFormatException(value);
+                return result;
+            } catch (NumberFormatException nfe) {
+                throw new ArgumentParserException("Invalid value: "+value, parser);
+            }
+        }
+    }
 }

cypher-shell/src/main/java/org/neo4j/shell/cli/CliArgs.java

@@ -5,9 +5,14 @@
 import java.util.Optional;
 
 public class CliArgs {
-    private String scheme = "bolt://";
-    private String host = "localhost";
-    private int port = 7687;
+    private static final String DEFAULT_SCHEME = "bolt://";
+    private static final String DEFAULT_HOST = "localhost";
+    private static final int DEFAULT_PORT = 7687;
+    static final int DEFAULT_NUM_SAMPLE_ROWS = 1000;
+
+    private String scheme = DEFAULT_SCHEME;
+    private String host = DEFAULT_HOST;
+    private int port = DEFAULT_PORT;
     private String username = "";
     private String password = "";
     private FailBehavior failBehavior = FailBehavior.FAIL_FAST;
@@ -18,6 +23,8 @@ public class CliArgs {
     private boolean nonInteractive = false;
     private boolean version = false;
     private boolean driverVersion = false;
+    private int numSampleRows = DEFAULT_NUM_SAMPLE_ROWS;
+    private boolean wrap = true;
 
     /**
      * Set the scheme to the primary value, or if null, the fallback value.
@@ -106,7 +113,6 @@ public String getHost() {
         return host;
     }
 
-    @Nonnull
     public int getPort() {
         return port;
     }
@@ -167,4 +173,22 @@ public void setDriverVersion(boolean version) {
     public boolean isStringShell() {
         return cypher.isPresent();
     }
+
+    public boolean getWrap() {
+        return wrap;
+    }
+
+    public void setWrap(boolean wrap) {
+        this.wrap = wrap;
+    }
+
+    public int getNumSampleRows() {
+        return numSampleRows;
+    }
+
+    public void setNumSampleRows(Integer numSampleRows) {
+        if (numSampleRows != null && numSampleRows > 0) {
+            this.numSampleRows = numSampleRows;
+        }
+    }
 }

cypher-shell/src/main/java/org/neo4j/shell/log/AnsiLogger.java

@@ -23,6 +23,8 @@ public class AnsiLogger implements Logger {
     private final PrintStream err;
     private final boolean debug;
     private Format format;
+    private int numSampleRows;
+    private boolean wrap = true;
 
     public AnsiLogger(final boolean debug) {
         this(debug, Format.VERBOSE, System.out, System.err);
@@ -66,6 +68,26 @@ private static boolean isOutputInteractive() {
         return 1 == isatty(STDOUT_FILENO) && 1 == isatty(STDERR_FILENO);
     }
 
+    @Override
+    public int getNumSampleRows() {
+        return numSampleRows;
+    }
+
+    @Override
+    public void setNumSampleRows(int numSampleRows) {
+        this.numSampleRows = numSampleRows;
+    }
+
+    @Override
+    public boolean getWrap() {
+        return wrap;
+    }
+
+    @Override
+    public void setWrap(boolean wrap) {
+        this.wrap = wrap;
+    }
+
     @Nonnull
     @Override
     public PrintStream getOutputStream() {

cypher-shell/src/main/java/org/neo4j/shell/log/Logger.java

@@ -92,4 +92,12 @@ default void printIfPlain(@Nonnull String text) {
             printOut(text);
         }
     }
+
+    boolean getWrap();
+
+    void setWrap(boolean wrap);
+
+    int getNumSampleRows();
+
+    void setNumSampleRows(int numSampleRows);
 }

cypher-shell/src/main/java/org/neo4j/shell/prettyprint/LinePrinter.java

@@ -0,0 +1,9 @@
+package org.neo4j.shell.prettyprint;
+
+/**
+ * Prints lines.
+ */
+@FunctionalInterface
+public interface LinePrinter {
+    void println( String line );
+}

cypher-shell/src/main/java/org/neo4j/shell/prettyprint/OutputFormatter.java

@@ -13,11 +13,7 @@
 import org.neo4j.shell.state.BoltResult;
 
 import javax.annotation.Nonnull;
-import java.util.Arrays;
-import java.util.ArrayList;
-import java.util.LinkedHashMap;
-import java.util.List;
-import java.util.Map;
+import java.util.*;
 import java.util.stream.Collectors;
 
 import static java.util.Arrays.asList;
@@ -26,15 +22,18 @@
 
 public interface OutputFormatter {
 
+    enum Capablities {info, plan, result, footer, statistics}
+
     String COMMA_SEPARATOR = ", ";
     String COLON_SEPARATOR = ": ";
     String COLON = ":";
     String SPACE = " ";
     String NEWLINE =  System.getProperty("line.separator");
 
-    @Nonnull String format(@Nonnull BoltResult result);
+    void format(@Nonnull BoltResult result, @Nonnull LinePrinter linePrinter);
 
-    @Nonnull default String formatValue(@Nonnull final Value value) {
+    @Nonnull default String formatValue(final Value value) {
+        if (value == null) return "";
         TypeRepresentation type = (TypeRepresentation) value.type();
         switch (type.constructor()) {
             case LIST:
@@ -101,7 +100,7 @@ default String pathAsString(@Nonnull Path path) {
             }
         }
 
-        return list.stream().collect(Collectors.joining());
+        return String.join("", list);
     }
 
     @Nonnull default String relationshipAsString(@Nonnull Relationship relationship) {
@@ -187,6 +186,7 @@ static boolean isNotBlank(String string) {
         return "";
     }
 
+    Set<Capablities> capabilities();
 
     List<String> INFO = asList("Version", "Planner", "Runtime");
 

cypher-shell/src/main/java/org/neo4j/shell/prettyprint/PrettyPrinter.java

@@ -5,7 +5,7 @@
 
 import javax.annotation.Nonnull;
 
-import static java.util.Arrays.asList;
+import java.util.Set;
 
 /**
  * Print the result from neo4j in a intelligible fashion.
@@ -14,17 +14,26 @@ public class PrettyPrinter {
     private final StatisticsCollector statisticsCollector;
     private final OutputFormatter outputFormatter;
 
-    public PrettyPrinter(@Nonnull Format format) {
+    public PrettyPrinter(@Nonnull Format format, boolean wrap, int numSampleRows) {
         this.statisticsCollector = new StatisticsCollector(format);
-        this.outputFormatter = format == Format.VERBOSE ? new TableOutputFormatter() : new SimpleOutputFormatter();
+        this.outputFormatter = format == Format.VERBOSE ? new TableOutputFormatter(wrap, numSampleRows) : new SimpleOutputFormatter();
     }
 
-    public String format(@Nonnull final BoltResult result) {
-        String infoOutput = outputFormatter.formatInfo(result.getSummary());
-        String planOutput = outputFormatter.formatPlan(result.getSummary());
-        String statistics = statisticsCollector.collect(result.getSummary());
-        String resultOutput = outputFormatter.format(result);
-        String footer = outputFormatter.formatFooter(result);
-        return OutputFormatter.joinNonBlanks(OutputFormatter.NEWLINE, asList(infoOutput, planOutput, resultOutput, footer, statistics));
+    public void format(@Nonnull final BoltResult result, LinePrinter linePrinter) {
+        Set<OutputFormatter.Capablities> capabilities = outputFormatter.capabilities();
+
+        if (capabilities.contains(OutputFormatter.Capablities.result)) outputFormatter.format(result, linePrinter);
+
+        if (capabilities.contains(OutputFormatter.Capablities.info)) linePrinter.println(outputFormatter.formatInfo(result.getSummary()));
+        if (capabilities.contains(OutputFormatter.Capablities.plan)) linePrinter.println(outputFormatter.formatPlan(result.getSummary()));
+        if (capabilities.contains(OutputFormatter.Capablities.footer)) linePrinter.println(outputFormatter.formatFooter(result));
+        if (capabilities.contains(OutputFormatter.Capablities.statistics)) linePrinter.println(statisticsCollector.collect(result.getSummary()));
+    }
+
+    // Helper for testing
+    String format(@Nonnull final BoltResult result) {
+        StringBuilder sb = new StringBuilder();
+        format(result, line -> {if (line!=null && !line.trim().isEmpty()) sb.append(line).append(OutputFormatter.NEWLINE);});
+        return sb.toString();
     }
 }

cypher-shell/src/main/java/org/neo4j/shell/prettyprint/SimpleOutputFormatter.java

@@ -5,25 +5,26 @@
 import org.neo4j.driver.v1.summary.ResultSummary;
 import org.neo4j.shell.state.BoltResult;
 
-import java.util.List;
+import javax.annotation.Nonnull;
+import java.util.EnumSet;
+import java.util.Iterator;
 import java.util.Map;
+import java.util.Set;
 import java.util.stream.Collectors;
 
-import javax.annotation.Nonnull;
-
 public class SimpleOutputFormatter implements OutputFormatter {
 
     @Override
-    @Nonnull
-    public String format(@Nonnull final BoltResult result) {
-        StringBuilder sb = new StringBuilder();
-        List<Record> records = result.getRecords();
-        if (!records.isEmpty()) {
-            sb.append(records.get(0).keys().stream().collect(Collectors.joining(COMMA_SEPARATOR)));
-            sb.append("\n");
-            sb.append(records.stream().map(this::formatRecord).collect(Collectors.joining("\n")));
+    public void format(@Nonnull BoltResult result, @Nonnull LinePrinter output) {
+        Iterator<Record> records = result.iterate();
+        if (records.hasNext()) {
+            Record firstRow = records.next();
+            output.println(String.join(COMMA_SEPARATOR, firstRow.keys()));
+            output.println(formatRecord(firstRow));
+            while (records.hasNext()) {
+                output.println(formatRecord(records.next()));
+            }
         }
-        return sb.toString();
     }
 
     @Nonnull
@@ -41,4 +42,9 @@ public String formatInfo(@Nonnull ResultSummary summary) {
         return info.entrySet().stream()
                 .map( e -> String.format("%s: %s",e.getKey(),e.getValue())).collect(Collectors.joining(NEWLINE));
     }
+
+    @Override
+    public Set<Capablities> capabilities() {
+        return EnumSet.of(Capablities.info, Capablities.statistics, Capablities.result);
+    }
 }