Merge branch 'main' into trino_func_types

satish-mittal · satish-mittal · commit dc27836f7b23 · 2023-10-17T15:55:40.000+05:30
diff --git a/helm/templates/deployment.yaml b/helm/templates/deployment.yaml
@@ -90,6 +90,13 @@ spec:
                   name: {{ .Values.queryServiceConfig.data.postgresPasswordSecretName }}
                   key: {{ .Values.queryServiceConfig.data.postgresqlPasswordSecretKey }}
             {{- end }}
+            {{- if eq .Values.queryServiceConfig.interactiveClientType "trino" }}
+            - name: {{ .Values.queryServiceConfig.data.trinoPasswordEnvVariable }}
+              valueFrom:
+                secretKeyRef:
+                  name: {{ .Values.queryServiceConfig.data.trinoPasswordSecretName }}
+                  key: {{ .Values.queryServiceConfig.data.trinoPasswordSecretKey }}
+            {{- end }}
           volumeMounts:
             - name: service-config
               mountPath: /app/resources/configs/{{ .Chart.Name }}/application.conf
diff --git a/helm/templates/query-service-config.yaml b/helm/templates/query-service-config.yaml
@@ -67,9 +67,9 @@ data:
       {{- end }}
       {{- if eq .Values.queryServiceConfig.interactiveClientType "trino" }}
       {{- range .Values.trinoHandlers }}
-      {
+        {
 {{ tpl . $ | indent 10 }}
-      }
+        }
       {{- end }}
       {{- end }}
       ]
diff --git a/owasp-suppressions.xml b/owasp-suppressions.xml
@@ -39,7 +39,7 @@
     <notes><![CDATA[
    file name: netty-handler-4.1.94.Final.jar
    ]]></notes>
-    <packageUrl regex="true">^pkg:maven/io\.netty/netty\-handler@.*$</packageUrl>
+    <packageUrl regex="true">^pkg:maven/io\.netty/netty.*@.*$</packageUrl>
     <vulnerabilityName>CVE-2023-4586</vulnerabilityName>
   </suppress>
 </suppressions>
diff --git a/query-service-impl/src/main/java/org/hypertrace/core/query/service/trino/QueryRequestToTrinoSQLConverter.java b/query-service-impl/src/main/java/org/hypertrace/core/query/service/trino/QueryRequestToTrinoSQLConverter.java
@@ -13,6 +13,7 @@
 import java.util.stream.Stream;
 import org.hypertrace.core.query.service.ExecutionContext;
 import org.hypertrace.core.query.service.api.Expression;
+import org.hypertrace.core.query.service.api.Filter;
 import org.hypertrace.core.query.service.api.OrderByExpression;
 import org.hypertrace.core.query.service.api.QueryRequest;
 import org.hypertrace.core.query.service.api.SortOrder;
@@ -31,11 +32,13 @@ class QueryRequestToTrinoSQLConverter {
 
   private final org.hypertrace.core.query.service.trino.TableDefinition tableDefinition;
   private final TrinoFunctionConverter functionConverter;
+  private final TrinoFilterHandler trinoFilterHandler;
 
   QueryRequestToTrinoSQLConverter(
       TableDefinition tableDefinition, TrinoFunctionConverter functionConverter) {
     this.tableDefinition = tableDefinition;
     this.functionConverter = functionConverter;
+    this.trinoFilterHandler = new TrinoFilterHandler();
   }
 
   Entry<String, Params> toSQL(
@@ -61,10 +64,13 @@ Entry<String, Params> toSQL(
 
     paramsBuilder.addStringParam(trinoExecutionContext.getExecutionContext().getTenantId());
     if (request.hasFilter()) {
-      String filterClause =
-          columnRequestConverter.convertFilterClause(
-              request.getFilter(), paramsBuilder, trinoExecutionContext);
-      trinoExecutionContext.addResolvedFilterColumnQuery(filterClause);
+      Filter filter = trinoFilterHandler.skipTrinoAttributeFilter(request.getFilter());
+      if (!filter.equals(Filter.getDefaultInstance())) {
+        String filterClause =
+            columnRequestConverter.convertFilterClause(
+                filter, paramsBuilder, trinoExecutionContext);
+        trinoExecutionContext.addResolvedFilterColumnQuery(filterClause);
+      }
     }
     trinoExecutionContext.addAllFilterTableColumnNames(
         trinoExecutionContext.getActualTableColumnNames());
diff --git a/query-service-impl/src/main/java/org/hypertrace/core/query/service/trino/TrinoBasedRequestHandler.java b/query-service-impl/src/main/java/org/hypertrace/core/query/service/trino/TrinoBasedRequestHandler.java
@@ -43,6 +43,7 @@ public class TrinoBasedRequestHandler implements RequestHandler {
   private static final String START_TIME_ATTRIBUTE_NAME_CONFIG_KEY = "startTimeAttributeName";
   private static final String SLOW_QUERY_THRESHOLD_MS_CONFIG = "slowQueryThresholdMs";
   private static final String MIN_REQUEST_DURATION_KEY = "minRequestDuration";
+  private static final String IS_TRINO_ATTRIBUTE = "EVENT.isTrino";
 
   private static final int DEFAULT_SLOW_QUERY_THRESHOLD_MS = 3000;
   private static final Set<Operator> GTE_OPERATORS = Set.of(Operator.GE, Operator.GT, Operator.EQ);
@@ -67,6 +68,7 @@ public class TrinoBasedRequestHandler implements RequestHandler {
   private Optional<String> startTimeAttributeName;
   private QueryRequestToTrinoSQLConverter request2TrinoSqlConverter;
   private final TrinoClientFactory trinoClientFactory;
+  private final TrinoFilterHandler trinoFilterHandler;
 
   private final JsonFormat.Printer protoJsonPrinter =
       JsonFormat.printer().omittingInsignificantWhitespace();
@@ -79,6 +81,7 @@ public class TrinoBasedRequestHandler implements RequestHandler {
     this.name = name;
     this.trinoClientFactory = trinoClientFactory;
     this.processConfig(config);
+    this.trinoFilterHandler = new TrinoFilterHandler();
   }
 
   @Override
@@ -96,7 +99,16 @@ public QueryCost canHandle(QueryRequest request, ExecutionContext executionConte
     Set<String> referencedColumns = executionContext.getReferencedColumns();
 
     Preconditions.checkArgument(!referencedColumns.isEmpty());
+
+    // query must contain isTrino attribute filter
+    if (!trinoFilterHandler.containsAttributeFilter(request)) {
+      return QueryCost.UNSUPPORTED;
+    }
+
     for (String referencedColumn : referencedColumns) {
+      if (referencedColumn.equalsIgnoreCase(IS_TRINO_ATTRIBUTE)) {
+        continue;
+      }
       if (!tableDefinition.containsColumn(referencedColumn)) {
         return QueryCost.UNSUPPORTED;
       }
diff --git a/query-service-impl/src/main/java/org/hypertrace/core/query/service/trino/TrinoFilterHandler.java b/query-service-impl/src/main/java/org/hypertrace/core/query/service/trino/TrinoFilterHandler.java
@@ -0,0 +1,58 @@
+package org.hypertrace.core.query.service.trino;
+
+import static org.hypertrace.core.query.service.QueryRequestUtil.getLogicalColumnName;
+
+import java.util.List;
+import java.util.Optional;
+import org.hypertrace.core.query.service.api.Filter;
+import org.hypertrace.core.query.service.api.QueryRequest;
+
+public class TrinoFilterHandler {
+  private static final String IS_TRINO_ATTRIBUTE = "EVENT.isTrino";
+
+  public boolean containsAttributeFilter(QueryRequest request) {
+    return request.hasFilter() && containsAttributeFilter(request.getFilter());
+  }
+
+  public Filter skipTrinoAttributeFilter(Filter filter) {
+    return skipAttributeFilterIfPresent(filter);
+  }
+
+  private boolean containsAttributeFilter(Filter filter) {
+    if (filter.getChildFilterCount() > 0) {
+      for (Filter childFilter : filter.getChildFilterList()) {
+        if (containsAttributeFilter(childFilter)) {
+          return true;
+        }
+      }
+    }
+    return isTrinoAttributeFilter(filter);
+  }
+
+  private Filter skipAttributeFilterIfPresent(Filter filter) {
+    if (filter.getChildFilterCount() > 0) {
+      Filter.Builder builder = filter.toBuilder();
+      List<Filter> childFilters = filter.getChildFilterList();
+      builder.clearChildFilter();
+      for (Filter childFilter : childFilters) {
+        Filter skippedFilter = skipAttributeFilterIfPresent(childFilter);
+        if (!skippedFilter.equals(Filter.getDefaultInstance())) {
+          builder.addChildFilter(skippedFilter);
+        }
+      }
+      if (builder.getChildFilterList().isEmpty()) {
+        return Filter.getDefaultInstance();
+      }
+      return builder.build();
+    } else if (isTrinoAttributeFilter(filter)) {
+      return Filter.getDefaultInstance();
+    }
+    return filter;
+  }
+
+  private boolean isTrinoAttributeFilter(Filter filter) {
+    // filter must contain Event.isTrino attribute
+    Optional<String> mayBeColumn = getLogicalColumnName(filter.getLhs());
+    return mayBeColumn.isPresent() && mayBeColumn.get().equalsIgnoreCase(IS_TRINO_ATTRIBUTE);
+  }
+}
diff --git a/query-service-impl/src/test/java/org/hypertrace/core/query/service/trino/QueryRequestToTrinoSQLConverterTest.java b/query-service-impl/src/test/java/org/hypertrace/core/query/service/trino/QueryRequestToTrinoSQLConverterTest.java
@@ -218,6 +218,96 @@ void testQueryWithBooleanFilter() {
   }
 
   @Test
+  void testQueryWithIsTrinoFilterOnly() {
+    QueryRequest queryRequest =
+        buildSimpleQueryWithFilter(createEqualsFilter("Event.isTrino", true));
+    TableDefinition tableDefinition = getDefaultTableDefinition();
+    defaultMockingForExecutionContext();
+
+    assertSQLQuery(
+        queryRequest,
+        "Select lower(to_hex(span_id)) FROM span-event-view WHERE "
+            + tableDefinition.getTenantIdColumn()
+            + " = '"
+            + TENANT_ID
+            + "'",
+        tableDefinition,
+        executionContext);
+  }
+
+  @Test
+  void testQueryWithIsTrinoFilterInBeginning() {
+    Filter isTrinoFilter = createEqualsFilter("EVENT.isTrino", true);
+    Filter isEntryFilter = createEqualsFilter("Span.is_entry", true);
+    Filter isBareFilter = createEqualsFilter("Span.isBare", false);
+    QueryRequest queryRequest =
+        buildSimpleQueryWithFilter(
+            createCompositeFilter(Operator.AND, isTrinoFilter, isEntryFilter, isBareFilter)
+                .build());
+    TableDefinition tableDefinition = getDefaultTableDefinition();
+    defaultMockingForExecutionContext();
+
+    assertSQLQuery(
+        queryRequest,
+        "Select lower(to_hex(span_id)) FROM span-event-view WHERE "
+            + tableDefinition.getTenantIdColumn()
+            + " = '"
+            + TENANT_ID
+            + "' "
+            + "AND ( is_entry = true AND is_bare = false )",
+        tableDefinition,
+        executionContext);
+  }
+
+  @Test
+  void testQueryWithIsTrinoFilterInMiddle() {
+    Filter isEntryFilter = createEqualsFilter("Span.is_entry", true);
+    Filter isTrinoFilter = createEqualsFilter("EVENT.isTrino", true);
+    Filter isBareFilter = createEqualsFilter("Span.isBare", false);
+    QueryRequest queryRequest =
+        buildSimpleQueryWithFilter(
+            createCompositeFilter(Operator.AND, isEntryFilter, isTrinoFilter, isBareFilter)
+                .build());
+    TableDefinition tableDefinition = getDefaultTableDefinition();
+    defaultMockingForExecutionContext();
+
+    assertSQLQuery(
+        queryRequest,
+        "Select lower(to_hex(span_id)) FROM span-event-view WHERE "
+            + tableDefinition.getTenantIdColumn()
+            + " = '"
+            + TENANT_ID
+            + "' "
+            + "AND ( is_entry = true AND is_bare = false )",
+        tableDefinition,
+        executionContext);
+  }
+
+  @Test
+  void testQueryWithIsTrinoFilterInEnd() {
+    Filter isEntryFilter = createEqualsFilter("Span.is_entry", true);
+    Filter isBareFilter = createEqualsFilter("Span.isBare", false);
+    Filter isTrinoFilter = createEqualsFilter("EVENT.isTrino", true);
+    QueryRequest queryRequest =
+        buildSimpleQueryWithFilter(
+            createCompositeFilter(Operator.AND, isEntryFilter, isBareFilter, isTrinoFilter)
+                .build());
+    TableDefinition tableDefinition = getDefaultTableDefinition();
+    defaultMockingForExecutionContext();
+
+    assertSQLQuery(
+        queryRequest,
+        "Select lower(to_hex(span_id)) FROM span-event-view WHERE "
+            + tableDefinition.getTenantIdColumn()
+            + " = '"
+            + TENANT_ID
+            + "' "
+            + "AND ( is_entry = true AND is_bare = false )",
+        tableDefinition,
+        executionContext);
+  }
+
+  // @Test
   void testQueryWithDoubleFilter() {
     QueryRequest queryRequest =
         buildSimpleQueryWithFilter(createEqualsFilter("Span.metrics.duration_millis", 1.2));
diff --git a/query-service/build.gradle.kts b/query-service/build.gradle.kts
@@ -11,7 +11,7 @@ plugins {
 dependencies {
   implementation(project(":query-service-factory"))
   implementation("org.hypertrace.core.grpcutils:grpc-server-utils:0.12.1")
-  implementation("org.hypertrace.core.serviceframework:platform-grpc-service-framework:0.1.60")
+  implementation("org.hypertrace.core.serviceframework:platform-grpc-service-framework:0.1.61")
   implementation("org.slf4j:slf4j-api:1.7.32")
   implementation("com.typesafe:config:1.4.1")
 
@@ -22,7 +22,7 @@ dependencies {
   integrationTestImplementation("org.testcontainers:testcontainers:1.16.2")
   integrationTestImplementation("org.testcontainers:junit-jupiter:1.16.2")
   integrationTestImplementation("org.testcontainers:kafka:1.16.2")
-  integrationTestImplementation("org.hypertrace.core.serviceframework:integrationtest-service-framework:0.1.60")
+  integrationTestImplementation("org.hypertrace.core.serviceframework:integrationtest-service-framework:0.1.61")
   integrationTestImplementation("com.github.stefanbirkner:system-lambda:1.2.0")
 
   integrationTestImplementation("org.apache.kafka:kafka-clients:7.2.1-ccs")
