fix: extended query anonymizer tests with functional tests queries

confluentinc · May 10, 2021 · 753ccb9 · 753ccb9
1 parent 8f588d3
commit 753ccb9
Show file tree

Hide file tree

Showing 6 changed files with 837 additions and 8 deletions.
diff --git a/ksqldb-engine/src/main/java/io/confluent/ksql/engine/rewrite/QueryAnonymizer.java b/ksqldb-engine/src/main/java/io/confluent/ksql/engine/rewrite/QueryAnonymizer.java
@@ -111,7 +111,7 @@ public String visitStatements(final StatementsContext context) {
         final String statement = visitSingleStatement(stmtContext);
         statementList.add(statement);
       }
-      return StringUtils.join(statementList, "");
+      return StringUtils.join(statementList, "\n");
     }
 
     @Override
@@ -122,14 +122,24 @@ public String visitSingleStatement(final SingleStatementContext context) {
     @Override
     public String visitType(final TypeContext context) {
       if (context.type().isEmpty()) {
-        return context.getText();
+        return getTypeName(context);
       }
+
+      // get type name
+      String typeName = "STRUCT";
+      if (context.MAP() != null) {
+        typeName = "MAP";
+      } else if (context.ARRAY() != null) {
+        typeName = "ARRAY";
+      }
+
+      // get type list
       final List<String> typeList = new ArrayList<>();
       for (TypeContext typeContext : context.type()) {
         typeList.add(visit(typeContext));
       }
 
-      return String.format("STRUCT<%s>", StringUtils.join(typeList, ", "));
+      return String.format("%s<%s>", typeName, StringUtils.join(typeList, ", "));
     }
 
     @Override
@@ -158,9 +168,8 @@ public String visitAlterSource(final AlterSourceContext context) {
     public String visitAlterOption(final AlterOptionContext context) {
       final String columnName = context.identifier().getText();
       final String anonColumnName = getAnonColumnName(columnName);
-      final String typeName = context.type().getText();
 
-      return String.format("ADD COLUMN %1$s %2$s", anonColumnName, typeName);
+      return String.format("ADD COLUMN %1$s %2$s", anonColumnName, visit(context.type()));
     }
 
     @Override
@@ -459,8 +468,13 @@ public String visitSingleExpression(final SingleExpressionContext context) {
     public String visitBooleanDefault(final BooleanDefaultContext context) {
       final String columnName = context.getChild(0).getChild(0).getText();
       final String anonColumnName = getAnonColumnName(columnName);
-      final String anonValue = visit(context.getChild(0).getChild(1));
-      return String.format("%1$s=%2$s", anonColumnName, anonValue);
+
+      if (context.getChild(0).getChild(1) != null) {
+        final String anonValue = visit(context.getChild(0).getChild(1));
+        return String.format("%1$s=%2$s", anonColumnName, anonValue);
+      }
+
+      return anonColumnName;
     }
 
     @Override
@@ -645,10 +659,18 @@ public String visitTableElements(final TableElementsContext context) {
 
     @Override
     public String visitTableElement(final TableElementContext context) {
+      final StringBuilder stringBuilder = new StringBuilder();
       final String columnName = ParserUtil.getIdentifierText(context.identifier());
       final String newName = getAnonColumnName(columnName);
+      stringBuilder.append(String.format("%1$s %2$s", newName, visit(context.type())));
 
-      return String.format("%1$s %2$s", newName, context.type().getText());
+      if (context.PRIMARY() != null) {
+        stringBuilder.append(" PRIMARY");
+      }
+      if (context.KEY() != null) {
+        stringBuilder.append(" KEY");
+      }
+      return stringBuilder.toString();
     }
 
     @Override
@@ -732,6 +754,24 @@ public String visitDropType(final DropTypeContext context) {
       return String.format("%s;", stringBuilder);
     }
 
+    private String getTypeName(final TypeContext context) {
+      if (context.DECIMAL() != null) {
+        return "DECIMAL";
+      }
+      switch (context.getText().toUpperCase()) {
+        case ("BOOLEAN"):
+        case ("INTEGER"):
+        case ("INT"):
+        case ("BIGINT"):
+        case ("DOUBLE"):
+        case ("STRING"):
+        case ("VARCHAR"):
+          return context.getText().toUpperCase();
+        default:
+          return "CUSTOM_TYPE";
+      }
+    }
+
     private String getQuery(final QueryContext context, final boolean isStream) {
       final StringBuilder stringBuilder = new StringBuilder();