fix: cannot reference variables in DEFINE statement

confluentinc · Nov 4, 2020 · d44baec · d44baec
1 parent 3b1fd3d
commit d44baec
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 3 deletions.
diff --git a/ksqldb-parser/src/main/java/io/confluent/ksql/parser/VariableSubstitutor.java b/ksqldb-parser/src/main/java/io/confluent/ksql/parser/VariableSubstitutor.java
@@ -90,14 +90,18 @@ public String replace(final SqlBaseParser.SingleStatementContext singleStatement
       return StringSubstitutor.replace(statementText, sanitizedValueMap);
     }
 
-    @Override
-    public Void visitStringLiteral(final SqlBaseParser.StringLiteralContext context) {
-      final String text = unquote(context.getText(), "\'");
+    private void lookupVariables(final String text) {
       for (String variableName : VariablesLookup.lookup(text)) {
         if (valueMap.containsKey(variableName)) {
           sanitizedValueMap.putIfAbsent(variableName, sanitize(valueMap.get(variableName)));
         }
       }
+    }
+
+    @Override
+    public Void visitStringLiteral(final SqlBaseParser.StringLiteralContext context) {
+      final String text = unquote(context.getText(), "\'");
+      lookupVariables(text);
 
       return null;
     }
@@ -113,6 +117,14 @@ public Void visitVariableLiteral(final SqlBaseParser.VariableLiteralContext cont
       return null;
     }
 
+    @Override
+    public Void visitVariableValue(final SqlBaseParser.VariableValueContext context) {
+      final String text = unquote(context.getText(), "\'");
+      lookupVariables(text);
+
+      return null;
+    }
+
     @Override
     public Void visitVariableIdentifier(final SqlBaseParser.VariableIdentifierContext context) {
       final String variableRef = context.getText();

diff --git a/ksqldb-parser/src/test/java/io/confluent/ksql/parser/VariableSubstitutorTest.java b/ksqldb-parser/src/test/java/io/confluent/ksql/parser/VariableSubstitutorTest.java
@@ -18,6 +18,28 @@
 public class VariableSubstitutorTest {
   private static final KsqlParser KSQL_PARSER = new DefaultKsqlParser();
 
+  @Test
+  public void shouldSubstituteVariableOnDefine() {
+    // Given
+    final Map<String, String> variablesMap = new ImmutableMap.Builder<String, String>() {{
+      put("env", "qa");
+      put("env_quoted", "\"qa\"");
+      put("env_backQuoted", "`qa`");
+    }}.build();
+
+    final List<Pair<String, String>> statements = Arrays.asList(
+        Pair.of("DEFINE topicName = 'topic_${env}';",
+            "DEFINE topicName = 'topic_qa';"),
+        Pair.of("DEFINE topicName = 'topic_${env_quoted}';",
+            "DEFINE topicName = 'topic_\"qa\"';"),
+        Pair.of("DEFINE topicName = 'topic_${env_backQuoted}';",
+            "DEFINE topicName = 'topic_`qa`';")
+    );
+
+    // When/Then
+    assertReplacedStatements(statements, variablesMap);
+  }
+
   @Test
   public void shouldSubstituteVariableOnDescribe() {
     // Given