Make CLI requests wait for last command sequence number

ksql-cli/src/main/java/io/confluent/ksql/cli/Cli.java

@@ -19,7 +19,9 @@
 import io.confluent.ksql.cli.console.Console;
 import io.confluent.ksql.cli.console.KsqlTerminal.StatusClosable;
 import io.confluent.ksql.cli.console.OutputFormat;
+import io.confluent.ksql.cli.console.cmd.CliCommandRegisterUtil;
 import io.confluent.ksql.cli.console.cmd.RemoteServerSpecificCommand;
+import io.confluent.ksql.cli.console.cmd.RequestPipeliningCommand;
 import io.confluent.ksql.ddl.DdlConfig;
 import io.confluent.ksql.parser.AstBuilder;
 import io.confluent.ksql.parser.KsqlParser;
@@ -34,6 +36,7 @@
 import io.confluent.ksql.rest.entity.KsqlEntity;
 import io.confluent.ksql.rest.entity.KsqlEntityList;
 import io.confluent.ksql.rest.entity.StreamedRow;
+import io.confluent.ksql.rest.server.resources.Errors;
 import io.confluent.ksql.util.CliUtils;
 import io.confluent.ksql.util.ErrorMessageUtil;
 import io.confluent.ksql.util.KsqlConstants;
@@ -57,12 +60,14 @@
 import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
+import java.util.function.BiFunction;
 import org.jline.reader.EndOfFileException;
 import org.jline.reader.UserInterruptException;
 import org.jline.terminal.Terminal;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+
 public class Cli implements Closeable {
 
   private static final Logger LOGGER = LoggerFactory.getLogger(Cli.class);
@@ -74,14 +79,15 @@ public class Cli implements Closeable {
 
   private final KsqlRestClient restClient;
   private final Console terminal;
+  private final RemoteServerState remoteServerState;
 
   public static Cli build(
       final Long streamedQueryRowLimit,
       final Long streamedQueryTimeoutMs,
       final OutputFormat outputFormat,
       final KsqlRestClient restClient
   ) {
-    final Console console = Console.build(outputFormat, restClient);
+    final Console console = Console.build(outputFormat);
     return new Cli(streamedQueryRowLimit, streamedQueryTimeoutMs, restClient, console);
   }
 
@@ -99,9 +105,14 @@ public static Cli build(
     this.restClient = restClient;
     this.terminal = terminal;
     this.queryStreamExecutorService = Executors.newSingleThreadExecutor();
-
-    terminal
-        .registerCliSpecificCommand(new RemoteServerSpecificCommand(restClient, terminal.writer()));
+    this.remoteServerState = new RemoteServerState();
+
+    CliCommandRegisterUtil.registerDefaultCliCommands(
+        terminal,
+        restClient,
+        remoteServerState::reset,
+        remoteServerState::getRequestPipelining,
+        remoteServerState::setRequestPipelining);
   }
 
   public void runInteractively() {
@@ -239,8 +250,7 @@ private void handleStatements(final String line)
     }
     if (consecutiveStatements.length() != 0) {
       printKsqlResponse(
-          restClient.makeKsqlRequest(consecutiveStatements.toString())
-      );
+          makeKsqlRequest(consecutiveStatements.toString(), restClient::makeKsqlRequest));
     }
   }
 
@@ -275,7 +285,7 @@ private void runScript(
     }
     setProperty(KsqlConstants.RUN_SCRIPT_STATEMENTS_CONTENT, fileContent);
     printKsqlResponse(
-        restClient.makeKsqlRequest(statementText)
+        makeKsqlRequest(statementText, restClient::makeKsqlRequest)
     );
   }
 
@@ -285,7 +295,8 @@ private StringBuilder printOrDisplayQueryResults(
       final String statementText
   ) throws InterruptedException, IOException, ExecutionException {
     if (consecutiveStatements.length() != 0) {
-      printKsqlResponse(restClient.makeKsqlRequest(consecutiveStatements.toString()));
+      printKsqlResponse(
+          makeKsqlRequest(consecutiveStatements.toString(), restClient::makeKsqlRequest));
       consecutiveStatements.setLength(0);
     }
     if (statementContext.statement() instanceof SqlBaseParser.QuerystatementContext) {
@@ -297,7 +308,8 @@ private StringBuilder printOrDisplayQueryResults(
   }
 
   private void listProperties(final String statementText) throws IOException {
-    final KsqlEntityList ksqlEntityList = restClient.makeKsqlRequest(statementText).getResponse();
+    final KsqlEntityList ksqlEntityList =
+        makeKsqlRequest(statementText, restClient::makeKsqlRequest).getResponse();
     terminal.printKsqlEntityList(ksqlEntityList);
   }
 
@@ -327,7 +339,7 @@ private void printKsqlResponse(final RestResponse<KsqlEntityList> response) thro
   private void handleStreamedQuery(final String query) throws IOException {
 
     final RestResponse<KsqlRestClient.QueryStream> queryResponse =
-        restClient.makeQueryRequest(query);
+        makeKsqlRequest(query, restClient::makeQueryRequest);
 
     LOGGER.debug("Handling streamed query");
 
@@ -392,7 +404,7 @@ private boolean limitNotReached(final long rowsRead) {
   private void handlePrintedTopic(final String printTopic)
       throws InterruptedException, ExecutionException, IOException {
     final RestResponse<InputStream> topicResponse =
-        restClient.makePrintTopicRequest(printTopic);
+        makeKsqlRequest(printTopic, restClient::makePrintTopicRequest);
 
     if (topicResponse.isSuccessful()) {
       try (Scanner topicStreamScanner = new Scanner(topicResponse.getResponse(), UTF_8.name());
@@ -461,8 +473,7 @@ private StringBuilder unsetProperty(
   ) throws IOException {
     if (consecutiveStatements.length() != 0) {
       printKsqlResponse(
-          restClient.makeKsqlRequest(consecutiveStatements.toString())
-      );
+          makeKsqlRequest(consecutiveStatements.toString(), restClient::makeKsqlRequest));
       consecutiveStatements.setLength(0);
     }
     final SqlBaseParser.UnsetPropertyContext unsetPropertyContext =
@@ -481,5 +492,76 @@ private void unsetProperty(final String property) {
 
     terminal.writer()
         .printf("Successfully unset local property '%s' (value was '%s').%n", property, oldValue);
+    terminal.flush();
+  }
+
+  private <R> RestResponse<R> makeKsqlRequest(
+      final String ksql,
+      final BiFunction<String, Long, RestResponse<R>> requestIssuer) {
+    final Long commandSequenceNumberToWaitFor = remoteServerState.getRequestPipelining()
+        ? null
+        : remoteServerState.getLastCommandSequenceNumber();
+    final RestResponse<R> response = requestIssuer.apply(ksql, commandSequenceNumberToWaitFor);
+
+    if (isSequenceNumberTimeout(response)) {
+      terminal.writer().printf(
+          "Error: command not executed since the server timed out "
+          + "while waiting for prior commands to finish executing.%n"
+          + "If you wish to execute new commands without waiting for "
+          + "prior commands to finish, run the command '%s ON'.%n",
+          RequestPipeliningCommand.NAME);
+    } else if (isKsqlEntityList(response)) {
+      updateLastCommandSequenceNumber((KsqlEntityList)response.getResponse());
+    }
+    return response;
+  }
+
+  private static boolean isSequenceNumberTimeout(final RestResponse<?> response) {
+    return response.isErroneous()
+        && (response.getErrorMessage().getErrorCode()
+            == Errors.ERROR_CODE_COMMAND_QUEUE_CATCHUP_TIMEOUT);
+  }
+
+  private static boolean isKsqlEntityList(final RestResponse<?> response) {
+    return response.isSuccessful() && response.getResponse() instanceof KsqlEntityList;
+  }
+
+  private void updateLastCommandSequenceNumber(final KsqlEntityList entities) {
+    entities.stream()
+        .filter(entity -> entity instanceof CommandStatusEntity)
+        .map(entity -> (CommandStatusEntity)entity)
+        .mapToLong(CommandStatusEntity::getCommandSequenceNumber)
+        .max()
+        .ifPresent(remoteServerState::setLastCommandSequenceNumber);
+  }
+
+  private static final class RemoteServerState {
+    private long lastCommandSequenceNumber;
+    private boolean requestPipelining;
+
+    private RemoteServerState() {
+      reset();
+    }
+
+    private void reset() {
+      lastCommandSequenceNumber = -1L;
+      requestPipelining = false;
+    }
+
+    private long getLastCommandSequenceNumber() {
+      return lastCommandSequenceNumber;
+    }
+
+    private boolean getRequestPipelining() {
+      return requestPipelining;
+    }
+
+    private void setLastCommandSequenceNumber(final long seqNum) {
+      lastCommandSequenceNumber = seqNum;
+    }
+
+    private void setRequestPipelining(final boolean newSetting) {
+      requestPipelining = newSetting;
+    }
   }
 }

ksql-cli/src/main/java/io/confluent/ksql/cli/console/Console.java

@@ -34,7 +34,6 @@
 import io.confluent.ksql.cli.console.table.builder.TableBuilder;
 import io.confluent.ksql.cli.console.table.builder.TablesListTableBuilder;
 import io.confluent.ksql.cli.console.table.builder.TopicDescriptionTableBuilder;
-import io.confluent.ksql.rest.client.KsqlRestClient;
 import io.confluent.ksql.rest.entity.CommandStatusEntity;
 import io.confluent.ksql.rest.entity.ExecutionPlan;
 import io.confluent.ksql.rest.entity.FieldInfo;
@@ -151,7 +150,7 @@ public interface RowCaptor {
     void addRows(List<List<String>> fields);
   }
 
-  public static Console build(final OutputFormat outputFormat, final KsqlRestClient restClient) {
+  public static Console build(final OutputFormat outputFormat) {
     final AtomicReference<Console> consoleRef = new AtomicReference<>();
     final Predicate<String> isCliCommand = line -> {
       final Console theConsole = consoleRef.get();
@@ -166,19 +165,15 @@ public static Console build(final OutputFormat outputFormat, final KsqlRestClien
 
     final KsqlTerminal terminal = new JLineTerminal(isCliCommand, historyFilePath);
 
-    final Supplier<String> versionSuppler =
-        () -> restClient.getServerInfo().getResponse().getVersion();
-
     final Console console = new Console(
-        outputFormat, versionSuppler, terminal, new NoOpRowCaptor());
+        outputFormat, terminal, new NoOpRowCaptor());
 
     consoleRef.set(console);
     return console;
   }
 
   public Console(
       final OutputFormat outputFormat,
-      final Supplier<String> versionSuppler,
       final KsqlTerminal terminal,
       final RowCaptor rowCaptor
   ) {
@@ -188,7 +183,7 @@ public Console(
     this.cliSpecificCommands = Maps.newLinkedHashMap();
     this.objectMapper = new ObjectMapper().disable(JsonGenerator.Feature.AUTO_CLOSE_TARGET);
 
-    CliCommandRegisterUtil.registerDefaultCommands(this, versionSuppler);
+    CliCommandRegisterUtil.registerDefaultConsoleCommands(this);
   }
 
   public PrintWriter writer() {

ksql-cli/src/main/java/io/confluent/ksql/cli/console/cmd/CliCommandRegisterUtil.java

@@ -15,20 +15,24 @@
 package io.confluent.ksql.cli.console.cmd;
 
 import io.confluent.ksql.cli.console.Console;
+import io.confluent.ksql.rest.client.KsqlRestClient;
+import io.confluent.ksql.util.Event;
+import java.util.function.Consumer;
 import java.util.function.Supplier;
 
 /**
  * God class for registering Cli commands
  */
+// CHECKSTYLE_RULES.OFF: ClassDataAbstractionCoupling
 public final class CliCommandRegisterUtil {
+  // CHECKSTYLE_RULES.ON: ClassDataAbstractionCoupling
 
   private CliCommandRegisterUtil() {
   }
 
-  public static void registerDefaultCommands(
-      final Console console,
-      final Supplier<String> versionSuppler) {
-
+  public static void registerDefaultConsoleCommands(
+      final Console console
+  ) {
     console.registerCliSpecificCommand(new Help(console));
 
     console.registerCliSpecificCommand(new Clear(console));
@@ -37,8 +41,24 @@ public static void registerDefaultCommands(
 
     console.registerCliSpecificCommand(new History(console));
 
+    console.registerCliSpecificCommand(new Exit(console));
+  }
+
+  public static void registerDefaultCliCommands(
+      final Console console,
+      final KsqlRestClient restClient,
+      final Event resetCliForNewServer,
+      final Supplier<Boolean> requestPipeliningSupplier,
+      final Consumer<Boolean> requestPipeliningConsumer
+  ) {
+    final Supplier<String> versionSuppler =
+        () -> restClient.getServerInfo().getResponse().getVersion();
     console.registerCliSpecificCommand(new Version(console, versionSuppler));
 
-    console.registerCliSpecificCommand(new Exit(console));
+    console.registerCliSpecificCommand(new RemoteServerSpecificCommand(
+        restClient, console.writer(), resetCliForNewServer));
+
+    console.registerCliSpecificCommand(new RequestPipeliningCommand(
+        console.writer(), requestPipeliningSupplier, requestPipeliningConsumer));
   }
 }

ksql-cli/src/main/java/io/confluent/ksql/cli/console/cmd/RemoteServerSpecificCommand.java

@@ -22,6 +22,7 @@
 import io.confluent.ksql.rest.entity.KsqlErrorMessage;
 import io.confluent.ksql.rest.server.resources.Errors;
 import io.confluent.ksql.util.ErrorMessageUtil;
+import io.confluent.ksql.util.Event;
 import java.io.PrintWriter;
 import java.util.Objects;
 import javax.ws.rs.ProcessingException;
@@ -30,12 +31,16 @@ public class RemoteServerSpecificCommand implements CliSpecificCommand {
 
   private final KsqlRestClient restClient;
   private final PrintWriter writer;
+  private final Event resetCliForNewServer;
 
-  public RemoteServerSpecificCommand(
+  RemoteServerSpecificCommand(
       final KsqlRestClient restClient,
-      final PrintWriter writer) {
+      final PrintWriter writer,
+      final Event resetCliForNewServer) {
     this.restClient = Objects.requireNonNull(restClient, "restClient");
     this.writer = Objects.requireNonNull(writer, "writer");
+    this.resetCliForNewServer =
+        Objects.requireNonNull(resetCliForNewServer, "resetCliForNewServer");
   }
 
   @Override
@@ -60,6 +65,7 @@ public void execute(final String command) {
       final String serverAddress = command.trim();
       restClient.setServerAddress(serverAddress);
       writer.write("Server now: " + command);
+      resetCliForNewServer.fire();
     }
 
     validateClient(writer, restClient);