Add anonymize test to chart command

Signed-off-by: Yuanchun Shen <[email protected]>

Author: yuancu

ppl/src/main/java/org/opensearch/sql/ppl/utils/PPLQueryDataAnonymizer.java

@@ -56,6 +56,7 @@
 import org.opensearch.sql.ast.tree.Append;
 import org.opensearch.sql.ast.tree.AppendCol;
 import org.opensearch.sql.ast.tree.Bin;
+import org.opensearch.sql.ast.tree.Chart;
 import org.opensearch.sql.ast.tree.CountBin;
 import org.opensearch.sql.ast.tree.Dedupe;
 import org.opensearch.sql.ast.tree.DefaultBin;
@@ -520,6 +521,42 @@ public String visitTimechart(Timechart node, String context) {
     return StringUtils.format("%s%s", child, timechartCommand.toString());
   }
 
+  @Override
+  public String visitChart(Chart node, String context) {
+    String child = node.getChild().get(0).accept(this, context);
+    StringBuilder chartCommand = new StringBuilder();
+    chartCommand.append(" | chart");
+
+    for (Argument arg : node.getArguments()) {
+      String argName = arg.getArgName();
+      // Skip the auto-generated "top" parameter that's added when limit is specified
+      if ("top".equals(argName)) {
+        continue;
+      }
+      if ("limit".equals(argName) || "useother".equals(argName) || "usenull".equals(argName)) {
+        chartCommand.append(" ").append(argName).append("=").append(MASK_LITERAL);
+      } else if ("otherstr".equals(argName) || "nullstr".equals(argName)) {
+        chartCommand.append(" ").append(argName).append("=").append(MASK_LITERAL);
+      }
+    }
+
+    chartCommand.append(" ").append(visitExpression(node.getAggregationFunction()));
+
+    if (node.getRowSplit() != null && node.getColumnSplit() != null) {
+      chartCommand
+          .append(" by ")
+          .append(visitExpression(node.getRowSplit()))
+          .append(" ")
+          .append(visitExpression(node.getColumnSplit()));
+    } else if (node.getRowSplit() != null) {
+      chartCommand.append(" by ").append(visitExpression(node.getRowSplit()));
+    } else if (node.getColumnSplit() != null) {
+      chartCommand.append(" by ").append(visitExpression(node.getColumnSplit()));
+    }
+
+    return StringUtils.format("%s%s", child, chartCommand.toString());
+  }
+
   public String visitRex(Rex node, String context) {
     String child = node.getChild().get(0).accept(this, context);
     String field = visitExpression(node.getField());

ppl/src/test/java/org/opensearch/sql/ppl/utils/PPLQueryDataAnonymizerTest.java

@@ -232,6 +232,34 @@ public void testTimechartCommand() {
         anonymize("source=t | timechart count() by host"));
   }
 
+  @Test
+  public void testChartCommand() {
+    assertEquals(
+        "source=table | chart count(identifier) by identifier identifier",
+        anonymize("source=t | chart count(age) by gender country"));
+  }
+
+  @Test
+  public void testChartCommandWithParameters() {
+    assertEquals(
+        "source=table | chart limit=*** useother=*** avg(identifier) by identifier",
+        anonymize("source=t | chart limit=5 useother=false avg(balance) by state"));
+  }
+
+  @Test
+  public void testChartCommandOver() {
+    assertEquals(
+        "source=table | chart avg(identifier) by identifier",
+        anonymize("source=t | chart avg(balance) over gender"));
+  }
+
+  @Test
+  public void testChartCommandOverBy() {
+    assertEquals(
+        "source=table | chart sum(identifier) by identifier identifier",
+        anonymize("source=t | chart sum(amount) over gender by age"));
+  }
+
   // todo, sort order is ignored, it doesn't impact the log analysis.
   @Test
   public void testSortCommandWithOptions() {