Fix search anoymizer only

core/src/main/java/org/opensearch/sql/ast/dsl/AstDSL.java

@@ -112,7 +112,7 @@ public UnresolvedPlan describe(String tableName) {
   }
 
   public static UnresolvedPlan search(UnresolvedPlan input, String queryString) {
-    return new Search(input, queryString);
+    return new Search(input, queryString, null);
   }
 
   public UnresolvedPlan subqueryAlias(UnresolvedPlan child, String alias) {

core/src/main/java/org/opensearch/sql/ast/expression/SearchAnd.java

@@ -27,6 +27,11 @@ public String toQueryString() {
     return left.toQueryString() + " AND " + right.toQueryString();
   }
 
+  @Override
+  public String toAnonymizedString() {
+    return left.toAnonymizedString() + " AND " + right.toAnonymizedString();
+  }
+
   @Override
   public List<? extends UnresolvedExpression> getChild() {
     return Arrays.asList(left, right);

core/src/main/java/org/opensearch/sql/ast/expression/SearchComparison.java

@@ -63,6 +63,11 @@ public String toQueryString() {
     }
   }
 
+  @Override
+  public String toAnonymizedString() {
+    return "identifier " + operator.symbol + " ***";
+  }
+
   @Override
   public List<? extends UnresolvedExpression> getChild() {
     return Arrays.asList(field, value);

core/src/main/java/org/opensearch/sql/ast/expression/SearchExpression.java

@@ -17,6 +17,13 @@ public abstract class SearchExpression extends UnresolvedExpression {
    */
   public abstract String toQueryString();
 
+  /**
+   * Convert the search expression to anonymized string
+   *
+   * @return the anonymized string
+   */
+  public abstract String toAnonymizedString();
+
   @Override
   public <R, C> R accept(AbstractNodeVisitor<R, C> nodeVisitor, C context) {
     return nodeVisitor.visitChildren(this, context);

core/src/main/java/org/opensearch/sql/ast/expression/SearchGroup.java

@@ -26,6 +26,11 @@ public String toQueryString() {
     return "(" + expression.toQueryString() + ")";
   }
 
+  @Override
+  public String toAnonymizedString() {
+    return "(" + expression.toAnonymizedString() + ")";
+  }
+
   @Override
   public List<? extends UnresolvedExpression> getChild() {
     return Collections.singletonList(expression);

core/src/main/java/org/opensearch/sql/ast/expression/SearchIn.java

@@ -33,6 +33,11 @@ public String toQueryString() {
     return fieldName + ":( " + valueList + " )";
   }
 
+  @Override
+  public String toAnonymizedString() {
+    return "identifier IN ***";
+  }
+
   @Override
   public List<? extends UnresolvedExpression> getChild() {
     List<UnresolvedExpression> children = new ArrayList<>();

core/src/main/java/org/opensearch/sql/ast/expression/SearchLiteral.java

@@ -55,6 +55,11 @@ public String toQueryString() {
     return QueryStringUtils.escapeLuceneSpecialCharacters(text);
   }
 
+  @Override
+  public String toAnonymizedString() {
+    return "***";
+  }
+
   @Override
   public List<? extends UnresolvedExpression> getChild() {
     return Collections.singletonList(literal);

core/src/main/java/org/opensearch/sql/ast/expression/SearchNot.java

@@ -26,6 +26,11 @@ public String toQueryString() {
     return "NOT(" + expression.toQueryString() + ")";
   }
 
+  @Override
+  public String toAnonymizedString() {
+    return "NOT(" + expression.toAnonymizedString() + ")";
+  }
+
   @Override
   public List<? extends UnresolvedExpression> getChild() {
     return Collections.singletonList(expression);

core/src/main/java/org/opensearch/sql/ast/expression/SearchOr.java

@@ -27,6 +27,11 @@ public String toQueryString() {
     return left.toQueryString() + " OR " + right.toQueryString();
   }
 
+  @Override
+  public String toAnonymizedString() {
+    return left.toAnonymizedString() + " OR " + right.toAnonymizedString();
+  }
+
   @Override
   public List<? extends UnresolvedExpression> getChild() {
     return Arrays.asList(left, right);

core/src/main/java/org/opensearch/sql/ast/tree/Search.java

@@ -12,19 +12,22 @@
 import lombok.RequiredArgsConstructor;
 import lombok.ToString;
 import org.opensearch.sql.ast.AbstractNodeVisitor;
+import org.opensearch.sql.ast.expression.SearchExpression;
 
 /**
  * Logical plan node for Search operation. Represents search expressions that get converted to
  * query_string function.
  */
 @Getter
 @ToString
-@EqualsAndHashCode(callSuper = false)
+@EqualsAndHashCode(onlyExplicitlyIncluded = true, callSuper = false)
 @RequiredArgsConstructor
 public class Search extends UnresolvedPlan {
 
-  private final UnresolvedPlan child;
-  private final String queryString;
+  @EqualsAndHashCode.Include private final UnresolvedPlan child;
+  @EqualsAndHashCode.Include private final String queryString;
+
+  private final SearchExpression originalExpression;
 
   @Override
   public List<UnresolvedPlan> getChild() {
@@ -38,6 +41,6 @@ public <T, C> T accept(AbstractNodeVisitor<T, C> nodeVisitor, C context) {
 
   @Override
   public UnresolvedPlan attach(UnresolvedPlan child) {
-    return new Search(child, queryString);
+    return new Search(child, queryString, originalExpression);
   }
 }

core/src/test/java/org/opensearch/sql/analysis/AnalyzerSearchTest.java

@@ -65,7 +65,7 @@ public void testVisitSearchWithSimpleQuery() {
     // Arrange
     String queryString = "field1:value1";
     Relation relation = new Relation(AstDSL.qualifiedName("test_index"));
-    Search searchNode = new Search(relation, queryString);
+    Search searchNode = new Search(relation, queryString, null);
 
     LogicalRelation logicalRelation = new LogicalRelation("test_index", mockTable);
     FunctionExpression queryStringExpr = mock(FunctionExpression.class);
@@ -103,7 +103,7 @@ public void testVisitSearchWithComplexQuery() {
     // Arrange
     String queryString = "(field1:value1 OR field2:value2) AND NOT field3:value3";
     Relation relation = new Relation(AstDSL.qualifiedName("test_index"));
-    Search searchNode = new Search(relation, queryString);
+    Search searchNode = new Search(relation, queryString, null);
 
     FunctionExpression queryStringExpr = mock(FunctionExpression.class);
     when(queryStringExpr.type()).thenReturn(ExprCoreType.BOOLEAN);
@@ -134,7 +134,7 @@ public void testVisitSearchPreservesChildPlan() {
     // Arrange
     String queryString = "test:query";
     UnresolvedPlan mockChild = mock(UnresolvedPlan.class);
-    Search searchNode = new Search(mockChild, queryString);
+    Search searchNode = new Search(mockChild, queryString, null);
 
     LogicalPlan mockLogicalPlan = mock(LogicalPlan.class);
     FunctionExpression queryStringExpr = mock(FunctionExpression.class);
@@ -158,7 +158,7 @@ public void testVisitSearchWithEmptyQuery() {
     // Arrange
     String queryString = "";
     Relation relation = new Relation(AstDSL.qualifiedName("test_index"));
-    Search searchNode = new Search(relation, queryString);
+    Search searchNode = new Search(relation, queryString, null);
 
     FunctionExpression queryStringExpr = mock(FunctionExpression.class);
 
@@ -193,7 +193,7 @@ public void testVisitSearchCreatesCorrectQueryStringFunction() {
     // Arrange
     String queryString = "field:\"exact phrase\" AND field2:wildcard*";
     Relation relation = new Relation(AstDSL.qualifiedName("test_index"));
-    Search searchNode = new Search(relation, queryString);
+    Search searchNode = new Search(relation, queryString, null);
 
     FunctionExpression queryStringExpr = mock(FunctionExpression.class);
 

core/src/test/java/org/opensearch/sql/ast/tree/SearchTest.java

@@ -29,7 +29,7 @@ public class SearchTest {
   public void setUp() {
     mockChild = mock(UnresolvedPlan.class);
     testQueryString = "field1:value1 AND field2:value2";
-    search = new Search(mockChild, testQueryString);
+    search = new Search(mockChild, testQueryString, null);
   }
 
   @Test
@@ -69,7 +69,7 @@ public void testAccept() {
   public void testEquals() {
     UnresolvedPlan sameChild = mockChild;
     String sameQueryString = testQueryString;
-    Search sameSearch = new Search(sameChild, sameQueryString);
+    Search sameSearch = new Search(sameChild, sameQueryString, null);
 
     assertEquals(search, sameSearch);
     assertEquals(search, search);
@@ -78,14 +78,14 @@ public void testEquals() {
   @Test
   public void testNotEqualsWithDifferentChild() {
     UnresolvedPlan differentChild = mock(UnresolvedPlan.class);
-    Search differentSearch = new Search(differentChild, testQueryString);
+    Search differentSearch = new Search(differentChild, testQueryString, null);
 
     assertNotEquals(search, differentSearch);
   }
 
   @Test
   public void testNotEqualsWithDifferentQueryString() {
-    Search differentSearch = new Search(mockChild, "different:query");
+    Search differentSearch = new Search(mockChild, "different:query", null);
 
     assertNotEquals(search, differentSearch);
   }
@@ -102,14 +102,14 @@ public void testNotEqualsWithDifferentClass() {
 
   @Test
   public void testHashCode() {
-    Search sameSearch = new Search(mockChild, testQueryString);
+    Search sameSearch = new Search(mockChild, testQueryString, null);
     assertEquals(search.hashCode(), sameSearch.hashCode());
   }
 
   @Test
   public void testHashCodeWithDifferentValues() {
     UnresolvedPlan differentChild = mock(UnresolvedPlan.class);
-    Search differentSearch = new Search(differentChild, "different:query");
+    Search differentSearch = new Search(differentChild, "different:query", null);
     assertNotEquals(search.hashCode(), differentSearch.hashCode());
   }
 
@@ -123,22 +123,22 @@ public void testToString() {
 
   @Test
   public void testWithEmptyQueryString() {
-    Search emptySearch = new Search(mockChild, "");
+    Search emptySearch = new Search(mockChild, "", null);
     assertEquals("", emptySearch.getQueryString());
     assertEquals(mockChild, emptySearch.getChild().get(0));
   }
 
   @Test
   public void testWithComplexQueryString() {
     String complexQuery = "(field1:value1 OR field2:value2) AND NOT field3:value3";
-    Search complexSearch = new Search(mockChild, complexQuery);
+    Search complexSearch = new Search(mockChild, complexQuery, null);
     assertEquals(complexQuery, complexSearch.getQueryString());
   }
 
   @Test
   public void testWithSpecialCharactersInQueryString() {
     String specialCharsQuery = "field:\"value with spaces\" AND field2:value*";
-    Search specialSearch = new Search(mockChild, specialCharsQuery);
+    Search specialSearch = new Search(mockChild, specialCharsQuery, null);
     assertEquals(specialCharsQuery, specialSearch.getQueryString());
   }
 }

core/src/test/java/org/opensearch/sql/calcite/CalciteRelNodeVisitorSearchSimpleTest.java

@@ -36,7 +36,7 @@ public void testVisitSearchRequiresContext() {
     // Arrange
     String queryString = "field:value";
     Relation relation = new Relation(AstDSL.qualifiedName("test_index"));
-    Search searchNode = new Search(relation, queryString);
+    Search searchNode = new Search(relation, queryString, null);
 
     // Act & Assert - should throw NPE without proper context
     assertThrows(
@@ -53,7 +53,7 @@ public void testSearchNodeStructure() {
     Relation relation = new Relation(AstDSL.qualifiedName("test_index"));
 
     // Act
-    Search searchNode = new Search(relation, queryString);
+    Search searchNode = new Search(relation, queryString, null);
 
     // Assert
     assertNotNull(searchNode);
@@ -70,7 +70,7 @@ public void testSearchWithEmptyQuery() {
     Relation relation = new Relation(AstDSL.qualifiedName("test_index"));
 
     // Act
-    Search searchNode = new Search(relation, queryString);
+    Search searchNode = new Search(relation, queryString, null);
 
     // Assert
     assertEquals("", searchNode.getQueryString());
@@ -84,7 +84,7 @@ public void testSearchWithComplexQuery() {
     Relation relation = new Relation(AstDSL.qualifiedName("test_index"));
 
     // Act
-    Search searchNode = new Search(relation, queryString);
+    Search searchNode = new Search(relation, queryString, null);
 
     // Assert
     assertEquals(queryString, searchNode.getQueryString());
@@ -97,7 +97,7 @@ public void testSearchWithSpecialCharacters() {
     Relation relation = new Relation(AstDSL.qualifiedName("test_index"));
 
     // Act
-    Search searchNode = new Search(relation, queryString);
+    Search searchNode = new Search(relation, queryString, null);
 
     // Assert
     assertEquals(queryString, searchNode.getQueryString());

integ-test/src/test/resources/expectedOutput/calcite/explain_search_with_match_pushdown.json

@@ -0,0 +1,6 @@
+{
+  "calcite":{
+    "logical":"LogicalSystemLimit(fetch=[10000], type=[QUERY_SIZE_LIMIT])\n  LogicalProject(account_number=[$0], firstname=[$1], address=[$2], birthdate=[$3], gender=[$4], city=[$5], lastname=[$6], balance=[$7], employer=[$8], state=[$9], age=[$10], email=[$11], male=[$12])\n    LogicalFilter(condition=[match(MAP('field_name', $3), MAP('value', '\"2016\\-12\\-08 00\\:00\\:00.000000000\"':VARCHAR))])\n      CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_bank]])\n",
+    "physical":"CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_bank]], PushDownContext=[[PROJECT->[account_number, firstname, address, birthdate, gender, city, lastname, balance, employer, state, age, email, male], FILTER->match(MAP('field_name', $3), MAP('value', '\"2016\\-12\\-08 00\\:00\\:00.000000000\"':VARCHAR)), LIMIT->10000], OpenSearchRequestBuilder(sourceBuilder={\"from\":0,\"size\":10000,\"timeout\":\"1m\",\"query\":{\"match\":{\"birthdate\":{\"query\":\"\\\"2016\\\\-12\\\\-08 00\\\\:00\\\\:00.000000000\\\"\",\"operator\":\"OR\",\"prefix_length\":0,\"max_expansions\":50,\"fuzzy_transpositions\":true,\"lenient\":false,\"zero_terms_query\":\"NONE\",\"auto_generate_synonyms_phrase_query\":true,\"boost\":1.0}}},\"_source\":{\"includes\":[\"account_number\",\"firstname\",\"address\",\"birthdate\",\"gender\",\"city\",\"lastname\",\"balance\",\"employer\",\"state\",\"age\",\"email\",\"male\"],\"excludes\":[]}}, requestedTotalSize=10000, pageSize=null, startFrom=0)])\n"
+  }
+}

ppl/src/main/java/org/opensearch/sql/ppl/parser/AstBuilder.java

@@ -178,7 +178,7 @@ public UnresolvedPlan visitSearchFrom(SearchFromContext ctx) {
 
       // Create Search node with relation and query string
       Relation relation = (Relation) visitFromClause(ctx.fromClause());
-      return new Search(relation, queryString);
+      return new Search(relation, queryString, combined);
     }
   }
 

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

@@ -106,11 +106,11 @@
 /** Utility class to mask sensitive information in incoming PPL queries. */
 public class PPLQueryDataAnonymizer extends AbstractNodeVisitor<String, String> {
 
-  private static final String MASK_LITERAL = "***";
+  public static final String MASK_LITERAL = "***";
 
-  private static final String MASK_COLUMN = "identifier";
+  public static final String MASK_COLUMN = "identifier";
 
-  private static final String MASK_TABLE = "table";
+  public static final String MASK_TABLE = "table";
 
   private final AnonymizerExpressionAnalyzer expressionAnalyzer;
   private final Settings settings;
@@ -250,9 +250,7 @@ public String visitTableFunction(TableFunction node, String context) {
   @Override
   public String visitSearch(Search node, String context) {
     String source = node.getChild().get(0).accept(this, context);
-    String queryString = node.getQueryString();
-    String anonymized = queryString.replaceAll(":\\S+", ":" + MASK_LITERAL);
-    return StringUtils.format("%s %s", source, anonymized);
+    return StringUtils.format("%s %s", source, node.getOriginalExpression().toAnonymizedString());
   }
 
   @Override