fix: regex escaping special characters

query-service-client/src/main/java/org/hypertrace/core/query/service/client/QueryServiceClient.java

@@ -3,14 +3,18 @@
 import io.grpc.Deadline;
 import io.grpc.ManagedChannel;
 import io.grpc.ManagedChannelBuilder;
+import io.grpc.stub.StreamObserver;
 import java.util.Iterator;
 import java.util.Map;
+import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
 import org.hypertrace.core.grpcutils.client.GrpcClientRequestContextUtil;
 import org.hypertrace.core.grpcutils.client.RequestContextClientCallCredsProviderFactory;
 import org.hypertrace.core.query.service.api.QueryRequest;
 import org.hypertrace.core.query.service.api.QueryServiceGrpc;
 import org.hypertrace.core.query.service.api.QueryServiceGrpc.QueryServiceBlockingStub;
+import org.hypertrace.core.query.service.api.QueryServiceGrpc.QueryServiceFutureStub;
+import org.hypertrace.core.query.service.api.QueryServiceGrpc.QueryServiceStub;
 import org.hypertrace.core.query.service.api.ResultSetChunk;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -25,6 +29,8 @@ public class QueryServiceClient {
   public static final int DEFAULT_QUERY_SERVICE_GROUP_BY_LIMIT = 10000;
 
   private final QueryServiceBlockingStub queryServiceClient;
+  private final QueryServiceFutureStub queryServiceFutureStub;
+  private final QueryServiceStub queryServiceStub;
 
   public QueryServiceClient(QueryServiceConfig queryServiceConfig) {
     ManagedChannel managedChannel =
@@ -36,6 +42,14 @@ public QueryServiceClient(QueryServiceConfig queryServiceConfig) {
         QueryServiceGrpc.newBlockingStub(managedChannel)
             .withCallCredentials(
                 RequestContextClientCallCredsProviderFactory.getClientCallCredsProvider().get());
+    queryServiceFutureStub =
+        QueryServiceGrpc.newFutureStub(managedChannel)
+            .withCallCredentials(
+                RequestContextClientCallCredsProviderFactory.getClientCallCredsProvider().get());
+    queryServiceStub =
+        QueryServiceGrpc.newStub(managedChannel)
+            .withCallCredentials(
+                RequestContextClientCallCredsProviderFactory.getClientCallCredsProvider().get());
   }
 
   public Iterator<ResultSetChunk> executeQuery(
@@ -49,4 +63,16 @@ public Iterator<ResultSetChunk> executeQuery(
                 .withDeadline(Deadline.after(timeoutMillis, TimeUnit.MILLISECONDS))
                 .execute(request));
   }
+
+  public void executeQueryAsync(
+      QueryRequest request, Map<String, String> context, int timeoutMillis, StreamObserver<ResultSetChunk> resultSetChunkStreamObserver) {
+    LOG.debug(
+        "Sending query to query service with timeout: {}, and request: {}", timeoutMillis, request);
+    GrpcClientRequestContextUtil.executeWithHeadersContext(
+        context,
+        () ->
+            queryServiceStub
+                .withDeadline(Deadline.after(timeoutMillis, TimeUnit.MILLISECONDS))
+                .execute(request, resultSetChunkStreamObserver));
+  }
 }

query-service-impl/build.gradle.kts

@@ -39,6 +39,8 @@ dependencies {
   implementation("org.hypertrace.core.serviceframework:platform-metrics:0.1.8")
   implementation("com.google.protobuf:protobuf-java-util:3.12.2")
 
+  implementation("io.grpc:grpc-netty:1.31.1")
+
   testImplementation(project(":query-service-api"))
   testImplementation("org.junit.jupiter:junit-jupiter:5.6.2")
   testImplementation("org.mockito:mockito-core:3.3.3")

query-service-impl/src/main/java/org/hypertrace/core/query/service/pinot/PinotBasedRequestHandler.java

@@ -13,6 +13,7 @@
 import java.util.Map;
 import java.util.Map.Entry;
 import java.util.Set;
+import java.util.concurrent.Future;
 import org.apache.pinot.client.ResultSet;
 import org.apache.pinot.client.ResultSetGroup;
 import org.hypertrace.core.query.service.QueryContext;
@@ -157,15 +158,18 @@ public void handleRequest(
     validateQueryRequest(queryContext, request);
     Entry<String, Params> pql =
         request2PinotSqlConverter.toSQL(queryContext, request, requestAnalyzer.getAllSelections());
-    if (LOG.isDebugEnabled()) {
-      LOG.debug("Trying to execute PQL: [ {} ] by RequestHandler: [ {} ]", pql, this.getName());
-    }
+//    if (LOG.isDebugEnabled()) {
+//      LOG.debug("Trying to execute PQL: [ {} ] by RequestHandler: [ {} ]", pql, this.getName());
+//    }
+    LOG.info("Trying to execute PQL: [ {} ] by RequestHandler: [ {} ]", pql, this.getName());
     final PinotClient pinotClient = pinotClientFactory.getPinotClient(this.getName());
 
     final ResultSetGroup resultSetGroup;
     try {
-      resultSetGroup = pinotQueryExecutionTimer.recordCallable(
-          () -> pinotClient.executeQuery(pql.getKey(), pql.getValue()));
+      resultSetGroup = pinotQueryExecutionTimer.recordCallable(() -> {
+        Future<ResultSetGroup> resultSetGroupFuture = pinotClient.executeQueryAsync(pql.getKey(), pql.getValue());
+        return resultSetGroupFuture.get();
+      });
     } catch (Exception ex) {
       // Catch this exception to log the Pinot SQL query that caused the issue
       LOG.error("An error occurred while executing: {}", pql.getKey(), ex);
@@ -187,6 +191,7 @@ public void handleRequest(
       } catch (InvalidProtocolBufferException ignore) {
       }
     }
+    LOG.info("Query Execution time: {} millis", requestTimeMs);
   }
 
   void convert(

query-service-impl/src/main/java/org/hypertrace/core/query/service/pinot/QueryRequestToPinotSQLConverter.java

@@ -8,6 +8,8 @@
 import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map.Entry;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
 import org.apache.commons.codec.binary.Hex;
 import org.hypertrace.core.query.service.QueryContext;
 import org.hypertrace.core.query.service.api.Expression;
@@ -43,6 +45,15 @@ class QueryRequestToPinotSQLConverter {
   private final String percentileAggFunction;
   private final Joiner joiner = Joiner.on(", ").skipNulls();
 
+  final String regExSpecialChars = "<([{\\^-=$!|]})?*+.>";
+  final String regExSpecialCharsRE = regExSpecialChars.replaceAll( ".", "\\\\$0");
+  final Pattern reCharsREP = Pattern.compile( "[" + regExSpecialCharsRE + "]");
+
+  String quoteRegExSpecialChars( String s) {
+    Matcher m = reCharsREP.matcher( s);
+    return m.replaceAll( "\\\\$0");
+  }
+
   QueryRequestToPinotSQLConverter(ViewDefinition viewDefinition, String percentileAggFunc) {
     this.viewDefinition = viewDefinition;
     this.percentileAggFunction = percentileAggFunc;
@@ -117,6 +128,7 @@ Entry<String, Params> toSQL(
     if (LOG.isDebugEnabled()) {
       LOG.debug("Converted QueryRequest to Pinot SQL: {}", pqlBuilder);
     }
+    LOG.info("Converted QueryRequest to Pinot SQL: {}", pqlBuilder);
     return new SimpleEntry<>(pqlBuilder.toString(), paramsBuilder.build());
   }
 
@@ -138,10 +150,21 @@ private String convertFilter2String(Filter filter, Params.Builder paramsBuilder)
         case LIKE:
           // The like operation in PQL looks like `regexp_like(lhs, rhs)`
           Expression rhs = handleValueConversionForLiteralExpression(filter.getLhs(), filter.getRhs());
+          //rhs = escapeLiteralStringForLikeOperator(rhs);
+          LOG.info("Escaped expression: {}", rhs);
           builder.append(operator);
           builder.append("(");
           builder.append(convertExpression2String(filter.getLhs(), paramsBuilder));
           builder.append(",");
+//          if (!rhs.hasLiteral()) {
+//            throw new RuntimeException("Cannot escape any other type of expression except LiteralExpression");
+//          }
+//          if (rhs.getLiteral().getValue().getValueType() != ValueType.STRING) {
+//            throw new RuntimeException("rhs for LIKE should be a String");
+//          }
+//          String rhsVal = rhs.getLiteral().getValue().getString();
+//          paramsBuilder.addStringParam(rhsVal);
+//          builder.append(QUESTION_MARK);
           builder.append(convertExpression2String(rhs, paramsBuilder));
           builder.append(")");
           break;
@@ -185,6 +208,24 @@ private String convertFilter2String(Filter filter, Params.Builder paramsBuilder)
     return builder.toString();
   }
 
+  private Expression escapeLiteralStringForLikeOperator(Expression expression) {
+    if (!expression.hasLiteral() ||
+        expression.getLiteral().getValue().getValueType() != ValueType.STRING) {
+      return expression;
+    }
+
+    return Expression.newBuilder()
+        .setLiteral(
+            LiteralConstant.newBuilder()
+                .setValue(
+                    Value.newBuilder()
+                        .setValueType(ValueType.STRING)
+                        .setString(quoteRegExSpecialChars(expression.getLiteral().getValue().getString()))
+                )
+        )
+        .build();
+  }
+
   /**
    * Handles value conversion of a literal expression based on its associated column.
    *
@@ -304,6 +345,36 @@ private String convertExpression2String(Expression expression, Params.Builder pa
     return "";
   }
 
+//  private String convertExpression2EscapedString(Expression expression, Params.Builder paramsBuilder) {
+//    switch (expression.getValueCase()) {
+//      case LITERAL:
+//        return convertLiteralToString(expression.getLiteral(), paramsBuilder);
+//      case FUNCTION:
+//        Function function = expression.getFunction();
+//        String functionName = function.getFunctionName();
+//        // For COUNT(column_name), Pinot sql format converts it to COUNT(*) and even only works with
+//        // COUNT(*) for ORDER BY
+//        if (functionName.equalsIgnoreCase("COUNT")) {
+//          return functionName + "(*)";
+//        } else if (functionName.startsWith(PERCENTILE_PREFIX) && !PERCENTILE_PREFIX.equals(functionName)) {
+//          functionName = functionName.replaceFirst(PERCENTILE_PREFIX, percentileAggFunction);
+//        }
+//        List<Expression> argumentsList = function.getArgumentsList();
+//        String[] args = new String[argumentsList.size()];
+//        for (int i = 0; i < argumentsList.size(); i++) {
+//          Expression expr = argumentsList.get(i);
+//          args[i] = convertExpression2String(expr, paramsBuilder);
+//        }
+//        return functionName + "(" + joiner.join(args) + ")";
+//      case ORDERBY:
+//        OrderByExpression orderBy = expression.getOrderBy();
+//        return convertExpression2String(orderBy.getExpression(), paramsBuilder);
+//      default:
+//        throw new RuntimeException("Cannot escape any other type of expression except LiteralExpression");
+//    }
+//    return "";
+//  }
+
   private String convertExpressionToMapKeyColumn(Expression expression) {
     if (expression.getValueCase() == COLUMNIDENTIFIER) {
       String logicalColumnName = expression.getColumnIdentifier().getColumnName();

query-service-impl/src/test/java/org/hypertrace/core/query/service/QueryServiceImplTest.java

@@ -83,6 +83,53 @@ public void testGrpc() {
     }
   }
 
+  @Test
+  public void testLike() {
+    ManagedChannel managedChannel =
+        ManagedChannelBuilder.forAddress("localhost", 8090).usePlaintext().build();
+    QueryServiceBlockingStub QueryServiceBlockingStub =
+        QueryServiceGrpc.newBlockingStub(managedChannel);
+
+    QueryRequest queryRequest = buildSimpleLikeQuery();
+
+    LOGGER.info("Trying to send request {}", queryRequest);
+    Iterator<ResultSetChunk> resultSetChunkIterator =
+        QueryServiceBlockingStub.withDeadline(Deadline.after(15, TimeUnit.SECONDS))
+            .execute(queryRequest);
+    LOGGER.info("Got response back: {}", resultSetChunkIterator);
+    while (resultSetChunkIterator.hasNext()) {
+      LOGGER.info("{}", resultSetChunkIterator.next());
+    }
+  }
+
+  private QueryRequest buildSimpleLikeQuery() {
+    Builder builder = QueryRequest.newBuilder();
+    ColumnIdentifier spanId = ColumnIdentifier.newBuilder().setColumnName("EVENT.id").build();
+    builder.addSelection(Expression.newBuilder().setColumnIdentifier(spanId).build());
+    builder.addSelection(createSelection("EVENT.serviceName"));
+
+//    Filter startTimeFilter =
+//        createTimeFilter(
+//            "EVENT.start_time_millis",
+//            Operator.GT,
+//            System.currentTimeMillis() - 1000 * 60 * 60 * 24);
+//    Filter endTimeFilter =
+//        createTimeFilter("EVENT.end_time_millis", Operator.LT, System.currentTimeMillis());
+
+    Filter likeFilter = createStringColumnFilter("EVENT.serviceName", Operator.LIKE, "\\!\\(");
+
+    Filter andFilter =
+        Filter.newBuilder()
+            .setOperator(Operator.AND)
+            .addChildFilter(likeFilter)
+            //.addChildFilter(endTimeFilter)
+            .build();
+    builder.setFilter(andFilter);
+    builder.setLimit(5);
+
+    return builder.build();
+  }
+
   @Disabled
   public void testGrpcMap() {
     ManagedChannel managedChannel =

query-service/src/main/resources/configs/common/application.conf

@@ -5,8 +5,8 @@ service.admin.port = 8091
 service.config = {
   clients = [
     {
-      type = zookeeper
-      connectionString = "localhost:2181/pinot/org-views"
+      type = broker
+      connectionString = "localhost:8099"
       connectionString = ${?ZK_CONNECT_STR}
     }
   ]
@@ -15,7 +15,7 @@ service.config = {
     {
       name = trace-view-handler
       type = pinot
-      clientConfig = zookeeper
+      clientConfig = broker
       requestHandlerInfo = {
         tenantColumnName = tenant_id
         viewDefinition = {
@@ -38,7 +38,7 @@ service.config = {
     {
       name = span-event-view-handler
       type = pinot
-      clientConfig = zookeeper
+      clientConfig = broker
       requestHandlerInfo = {
         tenantColumnName = tenant_id
         viewDefinition = {
@@ -74,7 +74,7 @@ service.config = {
     {
       name = service-handler
       type = pinot
-      clientConfig = zookeeper
+      clientConfig = broker
       requestHandlerInfo = {
         tenantColumnName = tenant_id
         viewDefinition = {
@@ -100,7 +100,7 @@ service.config = {
     {
       name = api-traces-view-handler
       type = pinot
-      clientConfig = zookeeper
+      clientConfig = broker
       requestHandlerInfo = {
         tenantColumnName = tenant_id
         viewDefinition = {
@@ -149,7 +149,7 @@ service.config = {
     {
       name = backend-traces-view-handler
       type = pinot
-      clientConfig = zookeeper
+      clientConfig = broker
       requestHandlerInfo = {
         tenantColumnName = tenant_id
         viewDefinition = {
@@ -181,7 +181,7 @@ service.config = {
     {
       name = interactions-handler
       type = pinot
-      clientConfig = zookeeper
+      clientConfig = broker
       requestHandlerInfo = {
         tenantColumnName = tenant_id
         viewDefinition = {
@@ -205,7 +205,7 @@ service.config = {
     {
       name = backend-entity-handler
       type = pinot
-      clientConfig = zookeeper
+      clientConfig = broker
       requestHandlerInfo = {
         tenantColumnName = tenant_id
         viewDefinition = {
@@ -231,7 +231,7 @@ service.config = {
     {
       name = api-handler
       type = pinot
-      clientConfig = zookeeper
+      clientConfig = broker
       requestHandlerInfo = {
         tenantColumnName = tenant_id
         viewDefinition = {