feat: new syntax to interact with session variables (define/undefine/show variables)

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

@@ -25,10 +25,13 @@
 import io.confluent.ksql.parser.KsqlParser;
 import io.confluent.ksql.parser.KsqlParser.ParsedStatement;
 import io.confluent.ksql.parser.SqlBaseParser;
+import io.confluent.ksql.parser.SqlBaseParser.DefineVariableContext;
+import io.confluent.ksql.parser.SqlBaseParser.ListVariablesContext;
 import io.confluent.ksql.parser.SqlBaseParser.PrintTopicContext;
 import io.confluent.ksql.parser.SqlBaseParser.QueryStatementContext;
 import io.confluent.ksql.parser.SqlBaseParser.SetPropertyContext;
 import io.confluent.ksql.parser.SqlBaseParser.StatementContext;
+import io.confluent.ksql.parser.SqlBaseParser.UndefineVariableContext;
 import io.confluent.ksql.parser.SqlBaseParser.UnsetPropertyContext;
 import io.confluent.ksql.reactive.BaseSubscriber;
 import io.confluent.ksql.rest.Errors;
@@ -42,6 +45,7 @@
 import io.confluent.ksql.rest.entity.KsqlEntityList;
 import io.confluent.ksql.rest.entity.ServerInfo;
 import io.confluent.ksql.rest.entity.StreamedRow;
+import io.confluent.ksql.rest.entity.VariablesList;
 import io.confluent.ksql.util.AppInfo;
 import io.confluent.ksql.util.ErrorMessageUtil;
 import io.confluent.ksql.util.HandlerMaps;
@@ -53,14 +57,19 @@
 import io.vertx.core.VertxException;
 import java.io.Closeable;
 import java.io.PrintWriter;
+import java.util.Collections;
 import java.util.List;
+import java.util.Map;
 import java.util.Objects;
 import java.util.Optional;
+import java.util.TreeMap;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.TimeUnit;
 import java.util.function.BiFunction;
 import java.util.function.Supplier;
+import java.util.stream.Collectors;
+
 import org.jline.reader.EndOfFileException;
 import org.jline.reader.UserInterruptException;
 import org.jline.terminal.Terminal;
@@ -84,6 +93,9 @@ public class Cli implements KsqlRequestExecutor, Closeable {
           .put(PrintTopicContext.class, Cli::handlePrintedTopic)
           .put(SetPropertyContext.class, Cli::setPropertyFromCtxt)
           .put(UnsetPropertyContext.class, Cli::unsetPropertyFromCtxt)
+          .put(DefineVariableContext.class, Cli::defineVariableFromCtxt)
+          .put(UndefineVariableContext.class, Cli::undefineVariableFromCtxt)
+          .put(ListVariablesContext.class, Cli::listVariablesFromCtxt)
           .build();
 
   private final Long streamedQueryRowLimit;
@@ -93,6 +105,8 @@ public class Cli implements KsqlRequestExecutor, Closeable {
   private final Console terminal;
   private final RemoteServerState remoteServerState;
 
+  private final Map<String, String> sessionVariables;
+
   public static Cli build(
       final Long streamedQueryRowLimit,
       final Long streamedQueryTimeoutMs,
@@ -117,6 +131,7 @@ public static Cli build(
     this.restClient = restClient;
     this.terminal = terminal;
     this.remoteServerState = new RemoteServerState();
+    this.sessionVariables = new TreeMap<>(String.CASE_INSENSITIVE_ORDER);
 
     final Supplier<String> versionSuppler =
         () -> restClient.getServerInfo().getResponse().getVersion();
@@ -480,6 +495,44 @@ private void unsetProperty(final String property) {
     terminal.flush();
   }
 
+  @SuppressWarnings("unused")
+  private void defineVariableFromCtxt(
+      final String ignored,
+      final DefineVariableContext context
+  ) {
+    final String variableName = context.variableName().getText();
+    final String variableValue = ParserUtil.unquote(context.variableValue().getText(), "'");
+    sessionVariables.put(variableName, variableValue);
+  }
+
+  @SuppressWarnings("unused")
+  private void undefineVariableFromCtxt(
+      final String ignored,
+      final UndefineVariableContext context
+  ) {
+    final String variableName = context.variableName().getText();
+    if (sessionVariables.remove(variableName) == null) {
+      // Print only (no throws exception) to keep it as a warning message (like VariableExecutor)
+      terminal.writer()
+          .printf("Cannot undefine variable '%s' which was never defined.%n", variableName);
+      terminal.flush();
+    }
+  }
+
+  @SuppressWarnings("unused")
+  private void listVariablesFromCtxt(
+      final String ignored,
+      final ListVariablesContext listVariablesContext
+  ) {
+    final List<VariablesList.Variable> variables = sessionVariables.entrySet().stream()
+        .map(e -> new VariablesList.Variable(e.getKey(), e.getValue()))
+        .collect(Collectors.toList());
+
+    terminal.printKsqlEntityList(Collections.singletonList(
+        new VariablesList(listVariablesContext.getText(),variables)
+    ));
+  }
+
   private static boolean isSequenceNumberTimeout(final RestResponse<?> response) {
     return response.isErroneous()
         && (response.getErrorMessage().getErrorCode()

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

@@ -36,6 +36,7 @@
 import io.confluent.ksql.cli.console.table.builder.ExecutionPlanTableBuilder;
 import io.confluent.ksql.cli.console.table.builder.FunctionNameListTableBuilder;
 import io.confluent.ksql.cli.console.table.builder.KafkaTopicsListTableBuilder;
+import io.confluent.ksql.cli.console.table.builder.ListVariablesTableBuilder;
 import io.confluent.ksql.cli.console.table.builder.PropertiesListTableBuilder;
 import io.confluent.ksql.cli.console.table.builder.QueriesTableBuilder;
 import io.confluent.ksql.cli.console.table.builder.StreamsListTableBuilder;
@@ -83,6 +84,7 @@
 import io.confluent.ksql.rest.entity.TablesList;
 import io.confluent.ksql.rest.entity.TopicDescription;
 import io.confluent.ksql.rest.entity.TypeList;
+import io.confluent.ksql.rest.entity.VariablesList;
 import io.confluent.ksql.rest.entity.WarningEntity;
 import io.confluent.ksql.schema.ksql.LogicalSchema;
 import io.confluent.ksql.util.CmdLineUtil;
@@ -178,6 +180,8 @@ public class Console implements Closeable {
               tablePrinter(ErrorEntity.class, ErrorEntityTableBuilder::new))
           .put(WarningEntity.class,
               tablePrinter(WarningEntity.class, WarningEntityTableBuilder::new))
+          .put(VariablesList.class,
+              tablePrinter(VariablesList.class, ListVariablesTableBuilder::new))
           .build();
 
   private static <T extends KsqlEntity> Handler1<KsqlEntity, Console> tablePrinter(

ksqldb-cli/src/main/java/io/confluent/ksql/cli/console/table/builder/ListVariablesTableBuilder.java

@@ -0,0 +1,45 @@
+/*
+ * Copyright 2020 Confluent Inc.
+ *
+ * Licensed under the Confluent Community License (the "License"); you may not use
+ * this file except in compliance with the License.  You may obtain a copy of the
+ * License at
+ *
+ * http://www.confluent.io/confluent-community-license
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OF ANY KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations under the License.
+ */
+
+package io.confluent.ksql.cli.console.table.builder;
+
+import com.google.common.collect.ImmutableList;
+import io.confluent.ksql.cli.console.table.Table;
+import io.confluent.ksql.rest.entity.VariablesList;
+
+import java.util.Comparator;
+import java.util.List;
+import java.util.stream.Collectors;
+
+public class ListVariablesTableBuilder implements TableBuilder<VariablesList> {
+  private static final List<String> HEADERS =
+      ImmutableList.of("Variable Name", "Value");
+
+  @Override
+  public Table buildTable(final VariablesList entity) {
+    return new Table.Builder()
+        .withColumnHeaders(HEADERS)
+        .withRows(defRowValues(entity.getVariables()))
+        .build();
+  }
+
+  private static List<List<String>> defRowValues(final List<VariablesList.Variable> variables) {
+    return variables.stream()
+        .sorted(Comparator.comparing(var -> var.getName()))
+        .map(var -> ImmutableList.of(
+            var.getName(), var.getValue()))
+        .collect(Collectors.toList());
+  }
+}

ksqldb-cli/src/test/java/io/confluent/ksql/cli/CliTest.java

@@ -431,6 +431,26 @@ public void shouldPrintTopicWithDelimitedValue() {
         containsString(", key: ITEM_1, value: home cinema"));
   }
 
+  @Test
+  public void testVariableDefineUndefine() {
+    assertRunCommand("define var1 = '1';", is(EMPTY_RESULT));
+    assertRunCommand("define var2 = '2';", is(EMPTY_RESULT));
+    assertRunCommand("define var3 = '3';", is(EMPTY_RESULT));
+
+    assertRunCommand("undefine var1;", is(EMPTY_RESULT));
+
+    assertRunListCommand("variables", hasRows(
+        row(
+            "var2",
+            "2"
+        ),
+        row(
+            "var3",
+            "3"
+        )
+    ));
+  }
+
   @Test
   public void testPropertySetUnset() {
     assertRunCommand("set 'auto.offset.reset' = 'latest';", is(EMPTY_RESULT));

ksqldb-parser/src/main/antlr4/io/confluent/ksql/parser/SqlBase.g4

@@ -53,6 +53,7 @@ statement
     | (LIST | SHOW) FUNCTIONS                                               #listFunctions
     | (LIST | SHOW) (SOURCE | SINK)? CONNECTORS                             #listConnectors
     | (LIST | SHOW) TYPES                                                   #listTypes
+    | (LIST | SHOW) VARIABLES                                               #listVariables
     | DESCRIBE EXTENDED? sourceName                                         #showColumns
     | DESCRIBE FUNCTION identifier                                          #describeFunction
     | DESCRIBE CONNECTOR identifier                                         #describeConnector
@@ -62,6 +63,8 @@ statement
     | TERMINATE ALL                                                         #terminateQuery
     | SET STRING EQ STRING                                                  #setProperty
     | UNSET STRING                                                          #unsetProperty
+    | DEFINE variableName EQ variableValue                                  #defineVariable
+    | UNDEFINE variableName                                                 #undefineVariable
     | CREATE (OR REPLACE)? STREAM (IF NOT EXISTS)? sourceName
                 (tableElements)?
                 (WITH tableProperties)?                                     #createStream
@@ -337,6 +340,15 @@ identifier
     | DIGIT_IDENTIFIER       #digitIdentifier
     ;
 
+variableName
+    : IDENTIFIER
+    | nonReserved
+    ;
+
+variableValue
+    : STRING
+    ;
+
 sourceName
     : identifier
     ;
@@ -480,6 +492,8 @@ RENAME: 'RENAME';
 ARRAY: 'ARRAY';
 MAP: 'MAP';
 SET: 'SET';
+DEFINE: 'DEFINE';
+UNDEFINE: 'UNDEFINE';
 RESET: 'RESET';
 SESSION: 'SESSION';
 SAMPLE: 'SAMPLE';
@@ -504,6 +518,7 @@ REPLACE: 'REPLACE';
 ASSERT: 'ASSERT';
 ADD: 'ADD';
 ALTER: 'ALTER';
+VARIABLES: 'VARIABLES';
 
 IF: 'IF';
 
@@ -566,6 +581,10 @@ TIMESTAMP_WITH_TIME_ZONE
     : 'TIMESTAMP' WS 'WITH' WS 'TIME' WS 'ZONE'
     ;
 
+VARIABLE
+    : '${' IDENTIFIER '}'
+    ;
+
 fragment EXPONENT
     : 'E' [+-]? DIGIT+
     ;

ksqldb-parser/src/main/java/io/confluent/ksql/parser/AstBuilder.java

@@ -108,6 +108,7 @@
 import io.confluent.ksql.parser.tree.CreateStreamAsSelect;
 import io.confluent.ksql.parser.tree.CreateTable;
 import io.confluent.ksql.parser.tree.CreateTableAsSelect;
+import io.confluent.ksql.parser.tree.DefineVariable;
 import io.confluent.ksql.parser.tree.DescribeConnector;
 import io.confluent.ksql.parser.tree.DescribeFunction;
 import io.confluent.ksql.parser.tree.DropConnector;
@@ -130,6 +131,7 @@
 import io.confluent.ksql.parser.tree.ListTables;
 import io.confluent.ksql.parser.tree.ListTopics;
 import io.confluent.ksql.parser.tree.ListTypes;
+import io.confluent.ksql.parser.tree.ListVariables;
 import io.confluent.ksql.parser.tree.PartitionBy;
 import io.confluent.ksql.parser.tree.PrintTopic;
 import io.confluent.ksql.parser.tree.Query;
@@ -147,6 +149,7 @@
 import io.confluent.ksql.parser.tree.TableElement.Namespace;
 import io.confluent.ksql.parser.tree.TableElements;
 import io.confluent.ksql.parser.tree.TerminateQuery;
+import io.confluent.ksql.parser.tree.UndefineVariable;
 import io.confluent.ksql.parser.tree.UnsetProperty;
 import io.confluent.ksql.parser.tree.WindowExpression;
 import io.confluent.ksql.parser.tree.WithinExpression;
@@ -730,6 +733,11 @@ public Node visitListProperties(final SqlBaseParser.ListPropertiesContext contex
       return new ListProperties(getLocation(context));
     }
 
+    @Override
+    public Node visitListVariables(final SqlBaseParser.ListVariablesContext context) {
+      return new ListVariables(getLocation(context));
+    }
+
     @Override
     public Node visitSetProperty(final SqlBaseParser.SetPropertyContext context) {
       final String propertyName = ParserUtil.unquote(context.STRING(0).getText(), "'");
@@ -743,6 +751,19 @@ public Node visitUnsetProperty(final SqlBaseParser.UnsetPropertyContext context)
       return new UnsetProperty(getLocation(context), propertyName);
     }
 
+    @Override
+    public Node visitDefineVariable(final SqlBaseParser.DefineVariableContext context) {
+      final String variableName = context.variableName().getText();
+      final String variableValue = ParserUtil.unquote(context.variableValue().getText(), "'");
+      return new DefineVariable(getLocation(context), variableName, variableValue);
+    }
+
+    @Override
+    public Node visitUndefineVariable(final SqlBaseParser.UndefineVariableContext context) {
+      final String variableName = context.variableName().getText();
+      return new UndefineVariable(getLocation(context), variableName);
+    }
+
     @Override
     public Node visitPrintTopic(final SqlBaseParser.PrintTopicContext context) {
       final boolean fromBeginning = context.printClause().FROM() != null;

ksqldb-parser/src/main/java/io/confluent/ksql/parser/SqlFormatter.java

@@ -38,6 +38,7 @@
 import io.confluent.ksql.parser.tree.CreateStreamAsSelect;
 import io.confluent.ksql.parser.tree.CreateTable;
 import io.confluent.ksql.parser.tree.CreateTableAsSelect;
+import io.confluent.ksql.parser.tree.DefineVariable;
 import io.confluent.ksql.parser.tree.DropStatement;
 import io.confluent.ksql.parser.tree.DropStream;
 import io.confluent.ksql.parser.tree.DropTable;
@@ -52,6 +53,7 @@
 import io.confluent.ksql.parser.tree.ListFunctions;
 import io.confluent.ksql.parser.tree.ListStreams;
 import io.confluent.ksql.parser.tree.ListTables;
+import io.confluent.ksql.parser.tree.ListVariables;
 import io.confluent.ksql.parser.tree.PartitionBy;
 import io.confluent.ksql.parser.tree.Query;
 import io.confluent.ksql.parser.tree.RegisterType;
@@ -65,6 +67,7 @@
 import io.confluent.ksql.parser.tree.TableElement;
 import io.confluent.ksql.parser.tree.TableElements;
 import io.confluent.ksql.parser.tree.TerminateQuery;
+import io.confluent.ksql.parser.tree.UndefineVariable;
 import io.confluent.ksql.parser.tree.UnsetProperty;
 import io.confluent.ksql.query.QueryId;
 import io.confluent.ksql.schema.utils.FormatOptions;
@@ -443,6 +446,32 @@ protected Void visitSetProperty(final SetProperty node, final Integer context) {
       return null;
     }
 
+    @Override
+    protected Void visitDefineVariable(final DefineVariable node, final Integer context) {
+      builder.append("DEFINE ");
+      builder.append(node.getVariableName());
+      builder.append("='");
+      builder.append(node.getVariableValue());
+      builder.append("'");
+
+      return null;
+    }
+
+    @Override
+    public Void visitListVariables(final ListVariables node, final Integer context) {
+      builder.append("SHOW VARIABLES");
+
+      return null;
+    }
+
+    @Override
+    protected Void visitUndefineVariable(final UndefineVariable node, final Integer context) {
+      builder.append("UNDEFINE ");
+      builder.append(node.getVariableName());
+
+      return null;
+    }
+
     private void visitExtended() {
       builder.append(" EXTENDED");
     }