fix: support for attribute expressions in selections (#129)

* added sup for attr expr in execution context

* added sup for attr expr in projection transformation

* added sup for attr expr in pinotbasedrequesthandler

* refactored

* added unit test for execution context

* added unit test in pinot based rqst handler

* changed function signature

* resolved comments

* nit

* fixed unit tests

* added template for new tests

* added integ test with attr expr for existing queries

* added functions for integ test

* added new integ tests

* nit

* added contains key filter for selection

* fix: support for attribute expressioons in selections

* test: update tests

* style: fix nit

* Apply suggestions from code review

Co-authored-by: sarthak <sarthak02singhal@gmail.com>

* added unit test for contains_key_value op

* chore: update dependencies

* chore: fix integration tests

Co-authored-by: Sarthak Singhal <sarthaksinghal@Sarthaks-MacBook-Pro.local>
Co-authored-by: sarthak <sarthak02singhal@gmail.com>
Co-authored-by: kotharironak <53209990+kotharironak@users.noreply.github.com>

Author: 4 people

query-service-api/build.gradle.kts

@@ -23,7 +23,7 @@ protobuf {
     // the identifier, which can be referred to in the "plugins"
     // container of the "generateProtoTasks" closure.
     id("grpc_java") {
-      artifact = "io.grpc:protoc-gen-grpc-java:1.42.0"
+      artifact = "io.grpc:protoc-gen-grpc-java:1.43.1"
     }
 
     if (generateLocalGoGrpcFiles) {
@@ -66,8 +66,9 @@ tasks.test {
 }
 
 dependencies {
-  api("io.grpc:grpc-protobuf:1.42.0")
-  api("io.grpc:grpc-stub:1.42.0")
+  api(platform("io.grpc:grpc-bom:1.43.1"))
+  api("io.grpc:grpc-protobuf")
+  api("io.grpc:grpc-stub")
   api("javax.annotation:javax.annotation-api:1.3.2")
 
   testImplementation("org.junit.jupiter:junit-jupiter:5.7.1")

query-service-client/build.gradle.kts

@@ -7,10 +7,10 @@ plugins {
 
 dependencies {
   api(project(":query-service-api"))
-  implementation("org.hypertrace.core.grpcutils:grpc-client-utils:0.6.2")
+  implementation("org.hypertrace.core.grpcutils:grpc-client-utils:0.7.0")
 
   // Logging
-  implementation("org.slf4j:slf4j-api:1.7.30")
+  implementation("org.slf4j:slf4j-api:1.7.32")
   // Config
   implementation("com.typesafe:config:1.4.1")
 }

query-service-impl/build.gradle.kts

@@ -10,9 +10,6 @@ tasks.test {
 
 dependencies {
   constraints {
-    implementation("com.fasterxml.jackson.core:jackson-databind:2.12.2") {
-      because("Multiple vulnerabilities")
-    }
     implementation("io.netty:netty:3.10.6.Final") {
       because("https://snyk.io/vuln/SNYK-JAVA-IONETTY-30430")
     }
@@ -28,21 +25,21 @@ dependencies {
   }
   api(project(":query-service-api"))
   api("com.typesafe:config:1.4.1")
-  implementation("org.hypertrace.core.grpcutils:grpc-context-utils:0.6.2")
-  implementation("org.hypertrace.core.grpcutils:grpc-client-utils:0.6.2")
-  implementation("org.hypertrace.core.grpcutils:grpc-server-rx-utils:0.6.2")
+  implementation("org.hypertrace.core.grpcutils:grpc-context-utils:0.7.0")
+  implementation("org.hypertrace.core.grpcutils:grpc-client-utils:0.7.0")
+  implementation("org.hypertrace.core.grpcutils:grpc-server-rx-utils:0.7.0")
   implementation("org.hypertrace.core.attribute.service:attribute-service-api:0.12.3")
   implementation("org.hypertrace.core.attribute.service:attribute-projection-registry:0.12.3")
   implementation("org.hypertrace.core.attribute.service:caching-attribute-service-client:0.12.3")
-  implementation("org.hypertrace.core.serviceframework:service-framework-spi:0.1.28")
+  implementation("org.hypertrace.core.serviceframework:service-framework-spi:0.1.33")
   implementation("com.google.inject:guice:5.0.1")
   implementation("org.apache.pinot:pinot-java-client:0.6.0") {
     // We want to use log4j2 impl so exclude the log4j binding of slf4j
     exclude("org.slf4j", "slf4j-log4j12")
   }
-  implementation("org.slf4j:slf4j-api:1.7.30")
+  implementation("org.slf4j:slf4j-api:1.7.32")
   implementation("commons-codec:commons-codec:1.15")
-  implementation("org.hypertrace.core.serviceframework:platform-metrics:0.1.28")
+  implementation("org.hypertrace.core.serviceframework:platform-metrics:0.1.33")
   implementation("com.google.protobuf:protobuf-java-util:3.15.6")
   implementation("com.google.guava:guava:30.1.1-jre")
   implementation("io.reactivex.rxjava3:rxjava:3.0.11")

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

@@ -0,0 +1,174 @@
+package org.hypertrace.core.query.service;
+
+import static io.reactivex.rxjava3.core.Single.zip;
+
+import io.reactivex.rxjava3.core.Observable;
+import io.reactivex.rxjava3.core.Single;
+import java.util.List;
+import org.hypertrace.core.query.service.api.AttributeExpression;
+import org.hypertrace.core.query.service.api.ColumnIdentifier;
+import org.hypertrace.core.query.service.api.Expression;
+import org.hypertrace.core.query.service.api.Filter;
+import org.hypertrace.core.query.service.api.Function;
+import org.hypertrace.core.query.service.api.LiteralConstant;
+import org.hypertrace.core.query.service.api.OrderByExpression;
+import org.hypertrace.core.query.service.api.QueryRequest;
+import org.slf4j.Logger;
+
+public abstract class AbstractQueryTransformation implements QueryTransformation {
+
+  protected abstract Logger getLogger();
+
+  @Override
+  public Single<QueryRequest> transform(
+      QueryRequest queryRequest, QueryTransformationContext transformationContext) {
+    return zip(
+            this.transformExpressionList(queryRequest.getSelectionList()),
+            this.transformExpressionList(queryRequest.getAggregationList()),
+            this.transformFilter(queryRequest.getFilter()),
+            this.transformExpressionList(queryRequest.getGroupByList()),
+            this.transformOrderByList(queryRequest.getOrderByList()),
+            (selections, aggregations, filter, groupBys, orderBys) ->
+                this.rebuildRequest(
+                    queryRequest, selections, aggregations, filter, groupBys, orderBys))
+        .doOnSuccess(transformed -> this.debugLogIfRequestTransformed(queryRequest, transformed));
+  }
+
+  protected QueryRequest rebuildRequest(
+      QueryRequest original,
+      List<Expression> selections,
+      List<Expression> aggregations,
+      Filter filter,
+      List<Expression> groupBys,
+      List<OrderByExpression> orderBys) {
+
+    QueryRequest.Builder builder = original.toBuilder();
+
+    if (Filter.getDefaultInstance().equals(filter)) {
+      builder.clearFilter();
+    } else {
+      builder.setFilter(filter);
+    }
+
+    return builder
+        .clearSelection()
+        .addAllSelection(selections)
+        .clearAggregation()
+        .addAllAggregation(aggregations)
+        .clearGroupBy()
+        .addAllGroupBy(groupBys)
+        .clearOrderBy()
+        .addAllOrderBy(orderBys)
+        .build();
+  }
+
+  protected Single<Expression> transformExpression(Expression expression) {
+    switch (expression.getValueCase()) {
+      case COLUMNIDENTIFIER:
+        return this.transformColumnIdentifier(expression.getColumnIdentifier());
+      case ATTRIBUTE_EXPRESSION:
+        return this.transformAttributeExpression(expression.getAttributeExpression());
+      case FUNCTION:
+        return this.transformFunction(expression.getFunction());
+      case ORDERBY:
+        return this.transformOrderBy(expression.getOrderBy())
+            .map(expression.toBuilder()::setOrderBy)
+            .map(Expression.Builder::build);
+      case LITERAL:
+        return this.transformLiteral(expression.getLiteral());
+      case VALUE_NOT_SET:
+      default:
+        return Single.just(expression);
+    }
+  }
+
+  protected Single<Expression> transformColumnIdentifier(ColumnIdentifier columnIdentifier) {
+    return Single.just(Expression.newBuilder().setColumnIdentifier(columnIdentifier).build());
+  }
+
+  protected Single<Expression> transformAttributeExpression(
+      AttributeExpression attributeExpression) {
+    return Single.just(Expression.newBuilder().setAttributeExpression(attributeExpression).build());
+  }
+
+  protected Single<Expression> transformLiteral(LiteralConstant literalConstant) {
+    return Single.just(Expression.newBuilder().setLiteral(literalConstant).build());
+  }
+
+  protected Single<Expression> transformFunction(Function function) {
+    return this.transformExpressionList(function.getArgumentsList())
+        .map(expressions -> function.toBuilder().clearArguments().addAllArguments(expressions))
+        .map(Function.Builder::build)
+        .map(Expression.newBuilder()::setFunction)
+        .map(Expression.Builder::build);
+  }
+
+  protected Single<OrderByExpression> transformOrderBy(OrderByExpression orderBy) {
+    return this.transformExpression(orderBy.getExpression())
+        .map(orderBy.toBuilder()::setExpression)
+        .map(OrderByExpression.Builder::build);
+  }
+
+  protected Single<Filter> transformFilter(Filter filter) {
+    if (filter.equals(Filter.getDefaultInstance())) {
+      return Single.just(filter);
+    }
+
+    Single<Expression> lhsSingle = this.transformExpression(filter.getLhs());
+    Single<Expression> rhsSingle = this.transformExpression(filter.getRhs());
+    Single<List<Filter>> childFilterListSingle =
+        Observable.fromIterable(filter.getChildFilterList())
+            .concatMapSingle(this::transformFilter)
+            .toList();
+    return zip(
+        lhsSingle,
+        rhsSingle,
+        childFilterListSingle,
+        (lhs, rhs, childFilterList) ->
+            this.rebuildFilterOmittingDefaults(filter, lhs, rhs, childFilterList));
+  }
+
+  private Single<List<Expression>> transformExpressionList(List<Expression> expressionList) {
+    return Observable.fromIterable(expressionList)
+        .concatMapSingle(this::transformExpression)
+        .toList();
+  }
+
+  private Single<List<OrderByExpression>> transformOrderByList(
+      List<OrderByExpression> orderByList) {
+    return Observable.fromIterable(orderByList).concatMapSingle(this::transformOrderBy).toList();
+  }
+
+  /**
+   * This doesn't change any functional behavior, but omits fields that aren't needed, shrinking the
+   * object and keeping it equivalent to the source object for equality checks.
+   */
+  private Filter rebuildFilterOmittingDefaults(
+      Filter original, Expression lhs, Expression rhs, List<Filter> childFilters) {
+    Filter.Builder builder = original.toBuilder();
+
+    if (Expression.getDefaultInstance().equals(lhs)) {
+      builder.clearLhs();
+    } else {
+      builder.setLhs(lhs);
+    }
+
+    if (Expression.getDefaultInstance().equals(rhs)) {
+      builder.clearRhs();
+    } else {
+      builder.setRhs(rhs);
+    }
+
+    return builder.clearChildFilter().addAllChildFilter(childFilters).build();
+  }
+
+  private void debugLogIfRequestTransformed(QueryRequest original, QueryRequest transformed) {
+    if (!original.equals(transformed)) {
+      getLogger()
+          .debug(
+              "Request transformation occurred. Original request: {} Transformed Request: {}",
+              original,
+              transformed);
+    }
+  }
+}

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

@@ -1,12 +1,9 @@
 package org.hypertrace.core.query.service;
 
 import static org.hypertrace.core.query.service.api.Expression.ValueCase.ATTRIBUTE_EXPRESSION;
-import static org.hypertrace.core.query.service.api.Expression.ValueCase.COLUMNIDENTIFIER;
-import static org.hypertrace.core.query.service.api.Expression.ValueCase.FUNCTION;
 
 import java.util.Optional;
 import org.hypertrace.core.query.service.api.AttributeExpression;
-import org.hypertrace.core.query.service.api.ColumnIdentifier;
 import org.hypertrace.core.query.service.api.Expression;
 import org.hypertrace.core.query.service.api.Expression.ValueCase;
 import org.hypertrace.core.query.service.api.Filter;
@@ -21,12 +18,6 @@
  */
 public class QueryRequestUtil {
 
-  public static Expression createColumnExpression(String columnName) {
-    return Expression.newBuilder()
-        .setColumnIdentifier(ColumnIdentifier.newBuilder().setColumnName(columnName))
-        .build();
-  }
-
   public static Expression createStringLiteralExpression(String value) {
     return Expression.newBuilder()
         .setLiteral(
@@ -145,10 +136,7 @@ public static Optional<String> getAlias(Expression expression) {
                 ? getLogicalColumnName(expression).get()
                 : expression.getColumnIdentifier().getAlias());
       case ATTRIBUTE_EXPRESSION:
-        return Optional.of(
-            expression.getAttributeExpression().getAlias().isBlank()
-                ? getLogicalColumnName(expression).get()
-                : expression.getAttributeExpression().getAlias());
+        return Optional.of(getAlias(expression.getAttributeExpression()));
       case FUNCTION:
         // todo: handle recursive functions max(rollup(time,50)
         // workaround is to use alias for now
@@ -160,4 +148,10 @@ public static Optional<String> getAlias(Expression expression) {
         return Optional.empty();
     }
   }
+
+  public static String getAlias(AttributeExpression attributeExpression) {
+    return attributeExpression.getAlias().isBlank()
+        ? attributeExpression.getAttributeId()
+        : attributeExpression.getAlias();
+  }
 }

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

@@ -6,6 +6,7 @@
 import javax.inject.Singleton;
 import org.hypertrace.core.attribute.service.cachingclient.CachingAttributeClient;
 import org.hypertrace.core.query.service.api.QueryServiceGrpc.QueryServiceImplBase;
+import org.hypertrace.core.query.service.attribubteexpression.AttributeExpressionModule;
 import org.hypertrace.core.query.service.pinot.PinotModule;
 import org.hypertrace.core.query.service.projection.ProjectionModule;
 import org.hypertrace.core.query.service.prometheus.PrometheusModule;
@@ -33,5 +34,6 @@ protected void configure() {
     install(new PinotModule());
     install(new ProjectionModule());
     install(new PrometheusModule());
+    install(new AttributeExpressionModule());
   }
 }

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

@@ -2,10 +2,20 @@
 
 import io.reactivex.rxjava3.core.Single;
 import org.hypertrace.core.query.service.api.QueryRequest;
+import org.jetbrains.annotations.NotNull;
 
-public interface QueryTransformation {
+public interface QueryTransformation extends Comparable<QueryTransformation> {
   Single<QueryRequest> transform(
       QueryRequest queryRequest, QueryTransformationContext transformationContext);
 
+  default int getPriority() {
+    return 10;
+  }
+
+  @Override
+  default int compareTo(@NotNull QueryTransformation other) {
+    return Integer.compare(this.getPriority(), other.getPriority());
+  }
+
   interface QueryTransformationContext {}
 }

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

@@ -24,6 +24,7 @@ Single<QueryRequest> transform(QueryRequest originalRequest, String tenantId) {
     QueryTransformationContext transformationContext =
         new DefaultQueryTransformationContext(tenantId);
     return Observable.fromIterable(transformations)
+        .sorted()
         .reduce(
             Single.just(originalRequest),
             (requestSingle, transformation) ->

query-service-impl/src/main/java/org/hypertrace/core/query/service/attribubteexpression/AttributeExpressionModule.java

@@ -0,0 +1,18 @@
+package org.hypertrace.core.query.service.attribubteexpression;
+
+import com.google.inject.AbstractModule;
+import com.google.inject.multibindings.Multibinder;
+import org.hypertrace.core.query.service.QueryTransformation;
+
+public class AttributeExpressionModule extends AbstractModule {
+
+  @Override
+  protected void configure() {
+    Multibinder<QueryTransformation> transformationMultibinder =
+        Multibinder.newSetBinder(binder(), QueryTransformation.class);
+    transformationMultibinder
+        .addBinding()
+        .to(AttributeExpressionSubpathExistsFilteringTransformation.class);
+    transformationMultibinder.addBinding().to(AttributeExpressionNormalizationTransformation.class);
+  }
+}