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 index 16e198ab3de8..fbf9edf9826c 100644 --- 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 index 26c070e6a539..14897f7083f5 100644 --- 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 variablesMap = new ImmutableMap.Builder() {{ + put("env", "qa"); + put("env_quoted", "\"qa\""); + put("env_backQuoted", "`qa`"); + }}.build(); + + final List> 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