diff --git a/query-service-impl/src/main/java/org/hypertrace/core/query/service/ExecutionContext.java b/query-service-impl/src/main/java/org/hypertrace/core/query/service/ExecutionContext.java index 26559593..49761ca2 100644 --- a/query-service-impl/src/main/java/org/hypertrace/core/query/service/ExecutionContext.java +++ b/query-service-impl/src/main/java/org/hypertrace/core/query/service/ExecutionContext.java @@ -1,5 +1,11 @@ package org.hypertrace.core.query.service; +import static org.hypertrace.core.query.service.QueryRequestUtil.getAlias; +import static org.hypertrace.core.query.service.QueryRequestUtil.getLogicalColumnName; +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 com.google.common.base.Supplier; import com.google.common.base.Suppliers; import java.time.Duration; @@ -15,7 +21,6 @@ import java.util.Optional; import java.util.Set; import java.util.concurrent.TimeUnit; -import org.hypertrace.core.query.service.api.ColumnIdentifier; import org.hypertrace.core.query.service.api.ColumnMetadata; import org.hypertrace.core.query.service.api.Expression; import org.hypertrace.core.query.service.api.Expression.ValueCase; @@ -137,26 +142,10 @@ private ColumnMetadata toColumnMetadata(Expression expression) { ValueCase valueCase = expression.getValueCase(); switch (valueCase) { case COLUMNIDENTIFIER: - ColumnIdentifier columnIdentifier = expression.getColumnIdentifier(); - String alias = columnIdentifier.getAlias(); - if (alias != null && alias.trim().length() > 0) { - builder.setColumnName(alias); - } else { - builder.setColumnName(columnIdentifier.getColumnName()); - } - builder.setValueType(ValueType.STRING); - builder.setIsRepeated(false); - break; + case ATTRIBUTE_EXPRESSION: case FUNCTION: - Function function = expression.getFunction(); - alias = function.getAlias(); - if (alias != null && alias.trim().length() > 0) { - builder.setColumnName(alias); - } else { - // todo: handle recursive functions max(rollup(time,50) - // workaround is to use alias for now - builder.setColumnName(function.getFunctionName()); - } + String alias = getAlias(expression).orElseThrow(IllegalArgumentException::new); + builder.setColumnName(alias); builder.setValueType(ValueType.STRING); builder.setIsRepeated(false); break; @@ -172,8 +161,10 @@ private void extractColumns(List columns, Expression expression) { ValueCase valueCase = expression.getValueCase(); switch (valueCase) { case COLUMNIDENTIFIER: - ColumnIdentifier columnIdentifier = expression.getColumnIdentifier(); - columns.add(columnIdentifier.getColumnName()); + case ATTRIBUTE_EXPRESSION: + String logicalColumnName = + getLogicalColumnName(expression).orElseThrow(IllegalArgumentException::new); + columns.add(logicalColumnName); break; case LITERAL: // no columns @@ -233,7 +224,7 @@ private Optional buildQueryTimeRange(Filter filter, String timeF } private boolean isMatchingFilter(Filter filter, String column, Collection operators) { - return column.equals(filter.getLhs().getColumnIdentifier().getColumnName()) + return getLogicalColumnName(filter.getLhs()).map(column::equals).orElse(false) && (operators.stream() .anyMatch(operator -> Objects.equals(operator, filter.getOperator()))); } diff --git a/query-service-impl/src/main/java/org/hypertrace/core/query/service/QueryRequestUtil.java b/query-service-impl/src/main/java/org/hypertrace/core/query/service/QueryRequestUtil.java index 0fb9cfa9..f63db9ee 100644 --- a/query-service-impl/src/main/java/org/hypertrace/core/query/service/QueryRequestUtil.java +++ b/query-service-impl/src/main/java/org/hypertrace/core/query/service/QueryRequestUtil.java @@ -2,7 +2,9 @@ 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; @@ -124,19 +126,38 @@ public static boolean isDateTimeFunction(Expression expression) { && expression.getFunction().getFunctionName().equals("dateTimeConvert"); } - public static String getLogicalColumnName(Expression expression) { + public static Optional getLogicalColumnName(Expression expression) { switch (expression.getValueCase()) { case COLUMNIDENTIFIER: - return expression.getColumnIdentifier().getColumnName(); + return Optional.of(expression.getColumnIdentifier().getColumnName()); case ATTRIBUTE_EXPRESSION: - return expression.getAttributeExpression().getAttributeId(); + return Optional.of(expression.getAttributeExpression().getAttributeId()); default: - throw new IllegalArgumentException( - "Supports " - + ATTRIBUTE_EXPRESSION - + " and " - + COLUMNIDENTIFIER - + " expression type only"); + return Optional.empty(); + } + } + + public static Optional getAlias(Expression expression) { + switch (expression.getValueCase()) { + case COLUMNIDENTIFIER: + return Optional.of( + expression.getColumnIdentifier().getAlias().isBlank() + ? getLogicalColumnName(expression).get() + : expression.getColumnIdentifier().getAlias()); + case ATTRIBUTE_EXPRESSION: + return Optional.of( + expression.getAttributeExpression().getAlias().isBlank() + ? getLogicalColumnName(expression).get() + : expression.getAttributeExpression().getAlias()); + case FUNCTION: + // todo: handle recursive functions max(rollup(time,50) + // workaround is to use alias for now + return Optional.of( + expression.getFunction().getAlias().isBlank() + ? expression.getFunction().getFunctionName() + : expression.getFunction().getAlias()); + default: + return Optional.empty(); } } } diff --git a/query-service-impl/src/main/java/org/hypertrace/core/query/service/pinot/PinotBasedRequestHandler.java b/query-service-impl/src/main/java/org/hypertrace/core/query/service/pinot/PinotBasedRequestHandler.java index 278ed173..12d37ed4 100644 --- a/query-service-impl/src/main/java/org/hypertrace/core/query/service/pinot/PinotBasedRequestHandler.java +++ b/query-service-impl/src/main/java/org/hypertrace/core/query/service/pinot/PinotBasedRequestHandler.java @@ -1,6 +1,7 @@ package org.hypertrace.core.query.service.pinot; import static org.hypertrace.core.query.service.ConfigUtils.optionallyGet; +import static org.hypertrace.core.query.service.QueryRequestUtil.getLogicalColumnName; import com.google.common.base.Preconditions; import com.google.common.base.Stopwatch; @@ -225,10 +226,8 @@ && rhsHasLongValue(filter.getRhs())) { } private boolean lhsIsStartTimeAttribute(Expression lhs) { - return lhs.hasColumnIdentifier() - && startTimeAttributeName - .map(attributeName -> attributeName.equals(lhs.getColumnIdentifier().getColumnName())) - .orElse(false); + return startTimeAttributeName.isPresent() + && startTimeAttributeName.equals(getLogicalColumnName(lhs)); } private boolean rhsHasLongValue(Expression rhs) { @@ -245,7 +244,7 @@ private Set getMatchingViewFilterColumns( // return it. if (filter.getChildFilterCount() == 0) { return doesSingleViewFilterMatchLeafQueryFilter(viewFilterMap, filter) - ? Set.of(filter.getLhs().getColumnIdentifier().getColumnName()) + ? Set.of(getLogicalColumnName(filter.getLhs()).orElseThrow(IllegalArgumentException::new)) : Set.of(); } else { // 2. Internal filter node. Recursively get the matching nodes from children. @@ -274,15 +273,13 @@ private Set getMatchingViewFilterColumns( */ private boolean doesSingleViewFilterMatchLeafQueryFilter( Map viewFilterMap, Filter queryFilter) { - if (queryFilter.getLhs().getValueCase() != ValueCase.COLUMNIDENTIFIER) { - return false; - } + if (queryFilter.getOperator() != Operator.IN && queryFilter.getOperator() != Operator.EQ) { return false; } - String columnName = queryFilter.getLhs().getColumnIdentifier().getColumnName(); - ViewColumnFilter viewColumnFilter = viewFilterMap.get(columnName); + ViewColumnFilter viewColumnFilter = + viewFilterMap.get(getLogicalColumnName(queryFilter.getLhs()).orElse(null)); if (viewColumnFilter == null) { return false; } @@ -469,7 +466,8 @@ private Filter removeViewColumnFilter( private Filter rewriteLeafFilter( Filter queryFilter, Map columnFilterMap) { ViewColumnFilter viewColumnFilter = - columnFilterMap.get(queryFilter.getLhs().getColumnIdentifier().getColumnName()); + columnFilterMap.get( + getLogicalColumnName(queryFilter.getLhs()).orElseThrow(IllegalArgumentException::new)); // If the RHS of both the view filter and query filter match, return empty filter. if (viewColumnFilter != null && isEquals(viewColumnFilter.getValues(), queryFilter.getRhs())) { return Filter.getDefaultInstance(); diff --git a/query-service-impl/src/main/java/org/hypertrace/core/query/service/pinot/QueryRequestToPinotSQLConverter.java b/query-service-impl/src/main/java/org/hypertrace/core/query/service/pinot/QueryRequestToPinotSQLConverter.java index b405a901..8b219716 100644 --- a/query-service-impl/src/main/java/org/hypertrace/core/query/service/pinot/QueryRequestToPinotSQLConverter.java +++ b/query-service-impl/src/main/java/org/hypertrace/core/query/service/pinot/QueryRequestToPinotSQLConverter.java @@ -211,7 +211,7 @@ private Expression handleValueConversionForLiteralExpression(Expression lhs, Exp return rhs; } - String lhsColumnName = getLogicalColumnName(lhs); + String lhsColumnName = getLogicalColumnName(lhs).orElseThrow(IllegalArgumentException::new); try { Value value = DestinationColumnValueConverter.INSTANCE.convert( @@ -272,7 +272,8 @@ private String convertExpressionToString( case COLUMNIDENTIFIER: // this takes care of the Map Type where it's split into 2 columns List columnNames = - viewDefinition.getPhysicalColumnNames(getLogicalColumnName(expression)); + viewDefinition.getPhysicalColumnNames( + getLogicalColumnName(expression).orElseThrow(IllegalArgumentException::new)); return joiner.join(columnNames); case ATTRIBUTE_EXPRESSION: if (isAttributeExpressionWithSubpath(expression)) { @@ -292,7 +293,9 @@ private String convertExpressionToString( valCol); } else { // this takes care of the Map Type where it's split into 2 columns - columnNames = viewDefinition.getPhysicalColumnNames(getLogicalColumnName(expression)); + columnNames = + viewDefinition.getPhysicalColumnNames( + getLogicalColumnName(expression).orElseThrow(IllegalArgumentException::new)); return joiner.join(columnNames); } case LITERAL: @@ -313,7 +316,9 @@ private String convertExpressionToString( } private String convertExpressionToMapKeyColumn(Expression expression) { - String col = viewDefinition.getKeyColumnNameForMap(getLogicalColumnName(expression)); + String col = + viewDefinition.getKeyColumnNameForMap( + getLogicalColumnName(expression).orElseThrow(IllegalArgumentException::new)); if (col != null && col.length() > 0) { return col; } @@ -321,7 +326,9 @@ private String convertExpressionToMapKeyColumn(Expression expression) { } private String convertExpressionToMapValueColumn(Expression expression) { - String col = viewDefinition.getValueColumnNameForMap(getLogicalColumnName(expression)); + String col = + viewDefinition.getValueColumnNameForMap( + getLogicalColumnName(expression).orElseThrow(IllegalArgumentException::new)); if (col != null && col.length() > 0) { return col; } diff --git a/query-service-impl/src/main/java/org/hypertrace/core/query/service/projection/ProjectionTransformation.java b/query-service-impl/src/main/java/org/hypertrace/core/query/service/projection/ProjectionTransformation.java index 558afd96..f1e9b819 100644 --- a/query-service-impl/src/main/java/org/hypertrace/core/query/service/projection/ProjectionTransformation.java +++ b/query-service-impl/src/main/java/org/hypertrace/core/query/service/projection/ProjectionTransformation.java @@ -17,6 +17,7 @@ import java.util.List; import java.util.Optional; import java.util.stream.Collectors; +import java.util.stream.Stream; import javax.inject.Inject; import org.hypertrace.core.attribute.service.cachingclient.CachingAttributeClient; import org.hypertrace.core.attribute.service.projection.AttributeProjection; @@ -30,6 +31,7 @@ import org.hypertrace.core.query.service.QueryFunctionConstants; import org.hypertrace.core.query.service.QueryRequestUtil; import org.hypertrace.core.query.service.QueryTransformation; +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; @@ -79,6 +81,8 @@ private Single 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()) .map(expression.toBuilder()::setFunction) @@ -96,10 +100,19 @@ private Single transformExpression(Expression expression) { private Single transformColumnIdentifier(ColumnIdentifier columnIdentifier) { return this.projectAttributeIfPossible(columnIdentifier.getColumnName()) - .map(expression -> this.aliasToMatchOriginal(columnIdentifier, expression)) + .map(expression -> this.aliasToMatchOriginal(getOriginalKey(columnIdentifier), expression)) .defaultIfEmpty(Expression.newBuilder().setColumnIdentifier(columnIdentifier).build()); } + private Single transformAttributeExpression(AttributeExpression attributeExpression) { + return this.projectAttributeIfPossible(attributeExpression.getAttributeId()) + .map( + expression -> + this.aliasToMatchOriginal(getOriginalKey(attributeExpression), expression)) + .defaultIfEmpty( + Expression.newBuilder().setAttributeExpression(attributeExpression).build()); + } + private Single transformFunction(Function function) { return this.transformExpressionList(function.getArgumentsList()) .map(expressions -> function.toBuilder().clearArguments().addAllArguments(expressions)) @@ -267,15 +280,18 @@ private Single convertOperator(ProjectionOperator operator) { } } - private Expression aliasToMatchOriginal(ColumnIdentifier original, Expression newExpression) { - String originalKey = - original.getAlias().isEmpty() ? original.getColumnName() : original.getAlias(); + private Expression aliasToMatchOriginal(String originalKey, Expression newExpression) { switch (newExpression.getValueCase()) { case COLUMNIDENTIFIER: return newExpression.toBuilder() .setColumnIdentifier( newExpression.getColumnIdentifier().toBuilder().setAlias(originalKey)) .build(); + case ATTRIBUTE_EXPRESSION: + return newExpression.toBuilder() + .setAttributeExpression( + newExpression.getAttributeExpression().toBuilder().setAlias(originalKey)) + .build(); case FUNCTION: return newExpression.toBuilder() .setFunction(newExpression.getFunction().toBuilder().setAlias(originalKey)) @@ -329,7 +345,8 @@ private QueryRequest rebuildRequest( List orderBys) { QueryRequest.Builder builder = original.toBuilder(); - Filter updatedFilter = rebuildFilterForComplexAttributeExpression(originalFilter, orderBys); + Filter updatedFilter = + rebuildFilterForComplexAttributeExpression(originalFilter, orderBys, selections); if (Filter.getDefaultInstance().equals(updatedFilter)) { builder.clearFilter(); @@ -350,16 +367,20 @@ private QueryRequest rebuildRequest( } /* - * We need the CONTAINS_KEY filter in all filters and order bys dealing with complex + * We need the CONTAINS_KEY filter in all filters, selections and order bys dealing with complex * attribute expressions as Pinot gives error if particular key is absent. Rest all work fine. - * To handle order bys, we add the corresponding filter at the top and 'AND' it with the main filter. + * To handle order bys and selections, we add the corresponding filter at the top and 'AND' it with the main filter. * To handle filter, we modify each filter (say filter1) as : "CONTAINS_KEY AND filter1". */ private Filter rebuildFilterForComplexAttributeExpression( - Filter originalFilter, List orderBys) { + Filter originalFilter, List orderBys, List selections) { Filter updatedFilter = updateFilterForComplexAttributeExpressionFromFilter(originalFilter); - List filterList = createFilterForComplexAttributeExpressionFromOrderBy(orderBys); + List filterList = + Stream.concat( + createFilterForComplexAttributeExpressionFromOrderBy(orderBys), + createFilterForComplexAttributeExpressionFromSelection(selections)) + .collect(Collectors.toList()); if (filterList.isEmpty()) { return updatedFilter; @@ -408,13 +429,40 @@ private Filter updateFilterForComplexAttributeExpressionFromFilter(Filter origin } } - private List createFilterForComplexAttributeExpressionFromOrderBy( + private Stream createFilterForComplexAttributeExpressionFromOrderBy( List orderByExpressionList) { return orderByExpressionList.stream() .map(OrderByExpression::getExpression) .filter(QueryRequestUtil::isAttributeExpressionWithSubpath) .map(Expression::getAttributeExpression) - .map(QueryRequestUtil::createContainsKeyFilter) - .collect(Collectors.toList()); + .map(QueryRequestUtil::createContainsKeyFilter); + } + + private Stream createFilterForComplexAttributeExpressionFromSelection( + List selections) { + return selections.stream() + .flatMap(this::getAnyAttributeExpression) + .map(QueryRequestUtil::createContainsKeyFilter); + } + + private Stream getAnyAttributeExpression(Expression selection) { + if (selection.hasFunction()) { + return selection.getFunction().getArgumentsList().stream() + .flatMap(this::getAnyAttributeExpression); + } else { + return Stream.of(selection) + .filter(QueryRequestUtil::isAttributeExpressionWithSubpath) + .map(Expression::getAttributeExpression); + } + } + + private String getOriginalKey(AttributeExpression attributeExpression) { + String alias = attributeExpression.getAlias(); + return alias.isEmpty() ? attributeExpression.getAttributeId() : alias; + } + + private String getOriginalKey(ColumnIdentifier columnIdentifier) { + String alias = columnIdentifier.getAlias(); + return alias.isEmpty() ? columnIdentifier.getColumnName() : alias; } } diff --git a/query-service-impl/src/main/java/org/hypertrace/core/query/service/prometheus/FilterToPromqlConverter.java b/query-service-impl/src/main/java/org/hypertrace/core/query/service/prometheus/FilterToPromqlConverter.java index b596d652..6cb99fc3 100644 --- a/query-service-impl/src/main/java/org/hypertrace/core/query/service/prometheus/FilterToPromqlConverter.java +++ b/query-service-impl/src/main/java/org/hypertrace/core/query/service/prometheus/FilterToPromqlConverter.java @@ -28,7 +28,8 @@ void convertFilterToString( } } else { if (QueryRequestUtil.isSimpleAttributeExpression(filter.getLhs()) - && timeFilterColumn.equals(getLogicalColumnName(filter.getLhs()))) { + && timeFilterColumn.equals( + getLogicalColumnName(filter.getLhs()).orElseThrow(IllegalArgumentException::new))) { return; } StringBuilder builder = new StringBuilder(); diff --git a/query-service-impl/src/main/java/org/hypertrace/core/query/service/prometheus/PrometheusBasedRequestHandler.java b/query-service-impl/src/main/java/org/hypertrace/core/query/service/prometheus/PrometheusBasedRequestHandler.java index a4320e8a..d2522a5a 100644 --- a/query-service-impl/src/main/java/org/hypertrace/core/query/service/prometheus/PrometheusBasedRequestHandler.java +++ b/query-service-impl/src/main/java/org/hypertrace/core/query/service/prometheus/PrometheusBasedRequestHandler.java @@ -138,7 +138,8 @@ private List prepareSelectionColumnSet( switch (valueCase) { case ATTRIBUTE_EXPRESSION: case COLUMNIDENTIFIER: - return QueryRequestUtil.getLogicalColumnName(expression); + return QueryRequestUtil.getLogicalColumnName(expression) + .orElseThrow(IllegalArgumentException::new); case FUNCTION: if (QueryRequestUtil.isDateTimeFunction(expression)) { return executionContext.getTimeFilterColumn(); diff --git a/query-service-impl/src/main/java/org/hypertrace/core/query/service/prometheus/PrometheusUtils.java b/query-service-impl/src/main/java/org/hypertrace/core/query/service/prometheus/PrometheusUtils.java index 1f231cfb..fb3053f1 100644 --- a/query-service-impl/src/main/java/org/hypertrace/core/query/service/prometheus/PrometheusUtils.java +++ b/query-service-impl/src/main/java/org/hypertrace/core/query/service/prometheus/PrometheusUtils.java @@ -7,7 +7,8 @@ class PrometheusUtils { static String generateAliasForMetricFunction(Expression functionExpression) { String functionName = functionExpression.getFunction().getFunctionName().toUpperCase(); String columnName = - QueryRequestUtil.getLogicalColumnName(functionExpression.getFunction().getArguments(0)); + QueryRequestUtil.getLogicalColumnName(functionExpression.getFunction().getArguments(0)) + .orElseThrow(IllegalArgumentException::new); return String.join(":", functionName, columnName); } } diff --git a/query-service-impl/src/main/java/org/hypertrace/core/query/service/prometheus/QueryRequestEligibilityValidator.java b/query-service-impl/src/main/java/org/hypertrace/core/query/service/prometheus/QueryRequestEligibilityValidator.java index 2968d711..0ab9cf34 100644 --- a/query-service-impl/src/main/java/org/hypertrace/core/query/service/prometheus/QueryRequestEligibilityValidator.java +++ b/query-service-impl/src/main/java/org/hypertrace/core/query/service/prometheus/QueryRequestEligibilityValidator.java @@ -4,6 +4,7 @@ import com.google.common.base.Preconditions; import java.util.List; +import java.util.Optional; import java.util.Set; import java.util.function.Predicate; import java.util.stream.Collectors; @@ -89,12 +90,14 @@ private boolean selectionAndGroupByOnDifferentColumn( Set selections = selectionList.stream() .map(QueryRequestUtil::getLogicalColumnName) + .map(Optional::orElseThrow) .collect(Collectors.toSet()); Set groupBys = groupByList.stream() .filter(Predicate.not(QueryRequestUtil::isDateTimeFunction)) .map(QueryRequestUtil::getLogicalColumnName) + .map(Optional::orElseThrow) .collect(Collectors.toSet()); return !selections.equals(groupBys); } @@ -113,7 +116,8 @@ private boolean areAggregationsNotSupported(List aggregationList) { return true; } Expression functionArgument = function.getArgumentsList().get(0); - String attributeId = getLogicalColumnName(functionArgument); + String attributeId = + getLogicalColumnName(functionArgument).orElseThrow(IllegalArgumentException::new); if (!PrometheusFunctionConverter.supportedFunctions.contains( function.getFunctionName())) { return true; diff --git a/query-service-impl/src/main/java/org/hypertrace/core/query/service/prometheus/QueryRequestToPromqlConverter.java b/query-service-impl/src/main/java/org/hypertrace/core/query/service/prometheus/QueryRequestToPromqlConverter.java index 8a833e8d..d1ff8c35 100644 --- a/query-service-impl/src/main/java/org/hypertrace/core/query/service/prometheus/QueryRequestToPromqlConverter.java +++ b/query-service-impl/src/main/java/org/hypertrace/core/query/service/prometheus/QueryRequestToPromqlConverter.java @@ -157,11 +157,12 @@ private String buildQuery( private MetricConfig getMetricConfigForFunction(Expression functionSelection) { return prometheusViewDefinition.getMetricConfigForLogicalMetricName( - getLogicalColumnName(functionSelection.getFunction().getArgumentsList().get(0))); + getLogicalColumnName(functionSelection.getFunction().getArgumentsList().get(0)) + .orElseThrow(IllegalArgumentException::new)); } private String convertColumnAttributeToString(Expression expression) { return prometheusViewDefinition.getPhysicalColumnNameForLogicalColumnName( - getLogicalColumnName(expression)); + getLogicalColumnName(expression).orElseThrow(IllegalArgumentException::new)); } } diff --git a/query-service-impl/src/test/java/org/hypertrace/core/query/service/QueryRequestBuilderUtils.java b/query-service-impl/src/test/java/org/hypertrace/core/query/service/QueryRequestBuilderUtils.java index 7d90b323..51b739e1 100644 --- a/query-service-impl/src/test/java/org/hypertrace/core/query/service/QueryRequestBuilderUtils.java +++ b/query-service-impl/src/test/java/org/hypertrace/core/query/service/QueryRequestBuilderUtils.java @@ -157,6 +157,13 @@ public static Filter createSimpleAttributeFilter( return createFilter(createSimpleAttributeExpression(column).build(), operator, expression); } + public static Filter createSimpleAttributeFilter(String column, Operator operator, String value) { + return createFilter( + createSimpleAttributeExpression(column).build(), + operator, + createStringLiteralValueExpression(value)); + } + public static Filter createFilter( Expression columnExpression, Operator operator, Expression expression) { return Filter.newBuilder() diff --git a/query-service-impl/src/test/java/org/hypertrace/core/query/service/pinot/ExecutionContextTest.java b/query-service-impl/src/test/java/org/hypertrace/core/query/service/pinot/ExecutionContextTest.java index 1a3513aa..b550977a 100644 --- a/query-service-impl/src/test/java/org/hypertrace/core/query/service/pinot/ExecutionContextTest.java +++ b/query-service-impl/src/test/java/org/hypertrace/core/query/service/pinot/ExecutionContextTest.java @@ -3,6 +3,7 @@ import static org.hypertrace.core.query.service.QueryRequestBuilderUtils.createFilter; import static org.hypertrace.core.query.service.QueryRequestBuilderUtils.createTimeColumnGroupByExpression; import static org.hypertrace.core.query.service.QueryRequestBuilderUtils.createTimeFilter; +import static org.hypertrace.core.query.service.QueryRequestUtil.createSimpleAttributeExpression; import static org.hypertrace.core.query.service.QueryRequestUtil.createStringLiteralValueExpression; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertNotNull; @@ -320,6 +321,78 @@ public void testSelectionsLinkedHashSet() { Expression.newBuilder().setFunction(minFunction).build(), selectionsIterator.next()); } + @Test + public void testSelectionsLinkedHashSetWithAttributeExpression() { + Builder builder = QueryRequest.newBuilder(); + // agg function with alias + Function count = + Function.newBuilder() + .setFunctionName("Count") + .setAlias("myCountAlias") + .addArguments(createSimpleAttributeExpression("Trace.id")) + .build(); + builder.addAggregation(Expression.newBuilder().setFunction(count)); + + // agg function without alias + Function minFunction = + Function.newBuilder() + .setFunctionName("MIN") + .addArguments(createSimpleAttributeExpression("Trace.duration")) + .build(); + builder.addAggregation(Expression.newBuilder().setFunction(minFunction)); + + // Add some selections + builder.addSelection(createSimpleAttributeExpression("Trace.transaction_name")); + builder.addSelection(createSimpleAttributeExpression("Trace.id")); + + // A function added into selections list is treated as a selection + Function avg = + Function.newBuilder() + .setFunctionName("AVG") + .setAlias("myAvgAlias") + .addArguments(createSimpleAttributeExpression("Trace.duration")) + .build(); + builder.addSelection(Expression.newBuilder().setFunction(avg)); + + // Add some group bys + builder.addGroupBy(createSimpleAttributeExpression("Trace.api_name")); + builder.addGroupBy(createSimpleAttributeExpression("Trace.service_name")); + QueryRequest queryRequest = builder.build(); + + ExecutionContext context = new ExecutionContext("test", queryRequest); + + // The order in resultSetMetadata.getColumnMetadataList() and selections is group bys, + // selections then aggregations + ResultSetMetadata resultSetMetadata = context.getResultSetMetadata(); + + assertNotNull(resultSetMetadata); + assertEquals(7, resultSetMetadata.getColumnMetadataCount()); + assertEquals("Trace.api_name", resultSetMetadata.getColumnMetadata(0).getColumnName()); + assertEquals("Trace.service_name", resultSetMetadata.getColumnMetadata(1).getColumnName()); + assertEquals("Trace.transaction_name", resultSetMetadata.getColumnMetadata(2).getColumnName()); + assertEquals("Trace.id", resultSetMetadata.getColumnMetadata(3).getColumnName()); + assertEquals("myAvgAlias", resultSetMetadata.getColumnMetadata(4).getColumnName()); + assertEquals("myCountAlias", resultSetMetadata.getColumnMetadata(5).getColumnName()); + assertEquals("MIN", resultSetMetadata.getColumnMetadata(6).getColumnName()); + + // Selections should correspond in size and order to the + // resultSetMetadata.getColumnMetadataList() + assertEquals(7, context.getAllSelections().size()); + Iterator selectionsIterator = context.getAllSelections().iterator(); + assertEquals( + createSimpleAttributeExpression("Trace.api_name").build(), selectionsIterator.next()); + assertEquals( + createSimpleAttributeExpression("Trace.service_name").build(), selectionsIterator.next()); + assertEquals( + createSimpleAttributeExpression("Trace.transaction_name").build(), + selectionsIterator.next()); + assertEquals(createSimpleAttributeExpression("Trace.id").build(), selectionsIterator.next()); + assertEquals(Expression.newBuilder().setFunction(avg).build(), selectionsIterator.next()); + assertEquals(Expression.newBuilder().setFunction(count).build(), selectionsIterator.next()); + assertEquals( + Expression.newBuilder().setFunction(minFunction).build(), selectionsIterator.next()); + } + @Test public void testSetTimeSeriesPeriod() { diff --git a/query-service-impl/src/test/java/org/hypertrace/core/query/service/pinot/PinotBasedRequestHandlerTest.java b/query-service-impl/src/test/java/org/hypertrace/core/query/service/pinot/PinotBasedRequestHandlerTest.java index 54845bd7..a6bb46ae 100644 --- a/query-service-impl/src/test/java/org/hypertrace/core/query/service/pinot/PinotBasedRequestHandlerTest.java +++ b/query-service-impl/src/test/java/org/hypertrace/core/query/service/pinot/PinotBasedRequestHandlerTest.java @@ -1,6 +1,10 @@ package org.hypertrace.core.query.service.pinot; import static org.hypertrace.core.query.service.QueryRequestBuilderUtils.createComplexAttributeExpression; +import static org.hypertrace.core.query.service.QueryRequestBuilderUtils.createIntLiteralValueExpression; +import static org.hypertrace.core.query.service.QueryRequestBuilderUtils.createLongLiteralValueExpression; +import static org.hypertrace.core.query.service.QueryRequestUtil.createBooleanLiteralExpression; +import static org.hypertrace.core.query.service.QueryRequestUtil.createSimpleAttributeExpression; import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; @@ -466,6 +470,236 @@ public void testCanHandleWithEqViewFilter() { } } + @Test + public void testCanHandleWithMultipleViewFiltersWithAttributeExpression() { + for (Config config : serviceConfig.getConfigList("queryRequestHandlersConfig")) { + if (!isPinotConfig(config)) { + continue; + } + + PinotBasedRequestHandler handler = + new PinotBasedRequestHandler( + config.getString("name"), config.getConfig("requestHandlerInfo")); + + // Verify that the entry span view handler can handle the query which has the filter + // on the column which has a view filter. + if (config.getString("name").equals("error-entry-span-view-handler")) { + // Positive case, straight forward. + QueryRequest request = + QueryRequest.newBuilder() + .setDistinctSelections(true) + .addSelection(createSimpleAttributeExpression("EVENT.startTime")) + .addSelection(createSimpleAttributeExpression("EVENT.id")) + .addSelection(createSimpleAttributeExpression("EVENT.traceId")) + .setFilter( + Filter.newBuilder() + .setOperator(Operator.AND) + .addChildFilter( + QueryRequestBuilderUtils.createSimpleAttributeFilter( + "EVENT.isEntrySpan", Operator.EQ, "true")) + .addChildFilter( + QueryRequestBuilderUtils.createSimpleAttributeFilter( + "EVENT.statusCode", Operator.EQ, "401"))) + .build(); + + ExecutionContext executionContext = new ExecutionContext("__default", request); + QueryCost cost = handler.canHandle(request, executionContext); + Assertions.assertTrue(cost.getCost() >= 0.0d && cost.getCost() < 1.0d); + + // Positive case but the filters are AND'ed in two different child filters. + Filter filter = + Filter.newBuilder() + .setOperator(Operator.AND) + .addChildFilter( + QueryRequestBuilderUtils.createSimpleAttributeFilter( + "EVENT.isEntrySpan", Operator.EQ, "true")) + .addChildFilter( + QueryRequestBuilderUtils.createSimpleAttributeFilter( + "EVENT.startTime", + Operator.GT, + QueryRequestBuilderUtils.createLongLiteralValueExpression( + System.currentTimeMillis()))) + .build(); + request = + QueryRequest.newBuilder() + .setDistinctSelections(true) + .addSelection(createSimpleAttributeExpression("EVENT.startTime")) + .addSelection(createSimpleAttributeExpression("EVENT.id")) + .addSelection(createSimpleAttributeExpression("EVENT.traceId")) + .setFilter( + Filter.newBuilder() + .setOperator(Operator.AND) + .addChildFilter(filter) + .addChildFilter( + QueryRequestBuilderUtils.createSimpleAttributeFilter( + "EVENT.statusCode", + Operator.EQ, + createIntLiteralValueExpression(401)))) + .build(); + + executionContext = new ExecutionContext("__default", request); + + cost = handler.canHandle(request, executionContext); + Assertions.assertTrue(cost.getCost() >= 0.0d && cost.getCost() < 1.0d); + + // Negative case. Query has only one leaf filter and it matches only one of the view + // filters + request = + QueryRequest.newBuilder() + .setDistinctSelections(true) + .addSelection(createSimpleAttributeExpression("EVENT.startTime")) + .addSelection(createSimpleAttributeExpression("EVENT.id")) + .addSelection(createSimpleAttributeExpression("EVENT.traceId")) + .setFilter( + QueryRequestBuilderUtils.createSimpleAttributeFilter( + "EVENT.isEntrySpan", Operator.EQ, createBooleanLiteralExpression(true))) + .build(); + + executionContext = new ExecutionContext("__default", request); + cost = handler.canHandle(request, executionContext); + Assertions.assertFalse(cost.getCost() >= 0.0d && cost.getCost() < 1.0d); + + // Negative case. Only one view filter is present in the query filters + request = + QueryRequest.newBuilder() + .setDistinctSelections(true) + .addSelection(createSimpleAttributeExpression("EVENT.startTime")) + .addSelection(createSimpleAttributeExpression("EVENT.id")) + .addSelection(createSimpleAttributeExpression("EVENT.traceId")) + .setFilter( + Filter.newBuilder() + .setOperator(Operator.AND) + .addChildFilter( + QueryRequestBuilderUtils.createSimpleAttributeFilter( + "EVENT.isEntrySpan", + Operator.EQ, + createBooleanLiteralExpression(true))) + .addChildFilter( + QueryRequestBuilderUtils.createSimpleAttributeFilter( + "EVENT.startTime", + Operator.GT, + QueryRequestBuilderUtils.createLongLiteralValueExpression( + System.currentTimeMillis()))) + .build()) + .build(); + + executionContext = new ExecutionContext("__default", request); + cost = handler.canHandle(request, executionContext); + Assertions.assertFalse(cost.getCost() >= 0.0d && cost.getCost() < 1.0d); + + // Negative case with correct filters but 'OR' operation. + request = + QueryRequest.newBuilder() + .setDistinctSelections(true) + .addSelection(createSimpleAttributeExpression("EVENT.startTime")) + .addSelection(createSimpleAttributeExpression("EVENT.id")) + .addSelection(createSimpleAttributeExpression("EVENT.traceId")) + .setFilter( + Filter.newBuilder() + .setOperator(Operator.OR) + .addChildFilter( + QueryRequestBuilderUtils.createSimpleAttributeFilter( + "EVENT.isEntrySpan", + Operator.EQ, + createBooleanLiteralExpression(true))) + .addChildFilter( + QueryRequestBuilderUtils.createSimpleAttributeFilter( + "EVENT.statusCode", + Operator.EQ, + createLongLiteralValueExpression(401L)))) + .build(); + + executionContext = new ExecutionContext("__default", request); + cost = handler.canHandle(request, executionContext); + Assertions.assertFalse(cost.getCost() >= 0.0d && cost.getCost() < 1.0d); + + // Negative case with a complex filter but 'OR' at the root level, hence shouldn't match. + filter = + Filter.newBuilder() + .setOperator(Operator.AND) + .addChildFilter( + QueryRequestBuilderUtils.createSimpleAttributeFilter( + "EVENT.isEntrySpan", Operator.EQ, createBooleanLiteralExpression(true))) + .addChildFilter( + QueryRequestBuilderUtils.createSimpleAttributeFilter( + "EVENT.statusCode", Operator.EQ, "401")) + .build(); + request = + QueryRequest.newBuilder() + .setDistinctSelections(true) + .addSelection(createSimpleAttributeExpression("EVENT.startTime")) + .addSelection(createSimpleAttributeExpression("EVENT.id")) + .addSelection(createSimpleAttributeExpression("EVENT.traceId")) + .setFilter( + Filter.newBuilder() + .setOperator(Operator.OR) + .addChildFilter( + QueryRequestBuilderUtils.createSimpleAttributeFilter( + "EVENT.isEntrySpan", + Operator.EQ, + createBooleanLiteralExpression(true))) + .addChildFilter(filter)) + .build(); + + executionContext = new ExecutionContext("__default", request); + cost = handler.canHandle(request, executionContext); + Assertions.assertFalse(cost.getCost() >= 0.0d && cost.getCost() < 1.0d); + + // Negative case. Value in query filter is different from the value in view filter + request = + QueryRequest.newBuilder() + .setDistinctSelections(true) + .addSelection(createSimpleAttributeExpression("EVENT.startTime")) + .addSelection(createSimpleAttributeExpression("EVENT.id")) + .addSelection(createSimpleAttributeExpression("EVENT.traceId")) + .setFilter( + Filter.newBuilder() + .setOperator(Operator.AND) + .addChildFilter( + QueryRequestBuilderUtils.createSimpleAttributeFilter( + "EVENT.isEntrySpan", Operator.EQ, "true")) + .addChildFilter( + QueryRequestBuilderUtils.createSimpleAttributeFilter( + "EVENT.statusCode", Operator.EQ, "200"))) + .build(); + + executionContext = new ExecutionContext("__default", request); + cost = handler.canHandle(request, executionContext); + Assertions.assertFalse(cost.getCost() >= 0.0d && cost.getCost() < 1.0d); + + // Negative case. Unsupported operator in the query filter. + request = + QueryRequest.newBuilder() + .setDistinctSelections(true) + .addSelection(createSimpleAttributeExpression("EVENT.startTime")) + .addSelection(createSimpleAttributeExpression("EVENT.id")) + .addSelection(createSimpleAttributeExpression("EVENT.traceId")) + .setFilter( + QueryRequestBuilderUtils.createSimpleAttributeFilter( + "EVENT.isEntrySpan", + Operator.IN, + QueryRequestUtil.createStringLiteralValueExpression("dummy"))) + .build(); + + executionContext = new ExecutionContext("__default", request); + cost = handler.canHandle(request, executionContext); + Assertions.assertFalse(cost.getCost() >= 0.0d && cost.getCost() < 1.0d); + + // Negative case. Any query without filter should not be handled. + request = + QueryRequest.newBuilder() + .setDistinctSelections(true) + .addSelection(createSimpleAttributeExpression("EVENT.startTime")) + .addSelection(createSimpleAttributeExpression("EVENT.id")) + .addSelection(createSimpleAttributeExpression("EVENT.traceId")) + .build(); + executionContext = new ExecutionContext("__default", request); + QueryCost negativeCost = handler.canHandle(request, executionContext); + Assertions.assertFalse(negativeCost.getCost() >= 0.0d && negativeCost.getCost() < 1.0d); + } + } + } + @Test public void testCanHandleWithMultipleViewFilters() { for (Config config : serviceConfig.getConfigList("queryRequestHandlersConfig")) { @@ -1215,6 +1449,61 @@ public void testNullTenantIdQueryRequestContextThrowsNPE() { .blockingSubscribe()); } + @Test + public void + testGroupBysAndAggregationsMixedWithSelectionsThrowsExceptionWhenDistinctSelectionIsSpecifiedWithAttributeExpression() { + // Setting distinct selections and mixing selections and group bys should throw exception + QueryRequest request = + QueryRequest.newBuilder() + .setDistinctSelections(true) + .addSelection(createSimpleAttributeExpression("col1")) + .addSelection(createSimpleAttributeExpression("col2")) + .addGroupBy(createSimpleAttributeExpression("col3")) + .build(); + Assertions.assertThrows( + IllegalArgumentException.class, + () -> + pinotBasedRequestHandler + .handleRequest(request, new ExecutionContext("test-tenant-id", request)) + .blockingSubscribe()); + + // Setting distinct selections and mixing selections and aggregations should throw exception + QueryRequest request2 = + QueryRequest.newBuilder() + .setDistinctSelections(true) + .addSelection(createSimpleAttributeExpression("col1")) + .addSelection(createSimpleAttributeExpression("col2")) + .addAggregation( + QueryRequestBuilderUtils.createAliasedFunctionExpressionWithSimpleAttribute( + "AVG", "duration", "avg_duration")) + .build(); + Assertions.assertThrows( + IllegalArgumentException.class, + () -> + pinotBasedRequestHandler + .handleRequest(request2, new ExecutionContext("test-tenant-id", request2)) + .blockingSubscribe()); + + // Setting distinct selections and mixing selections, group bys and aggregations should throw + // exception + QueryRequest request3 = + QueryRequest.newBuilder() + .setDistinctSelections(true) + .addSelection(createSimpleAttributeExpression("col1")) + .addSelection(createSimpleAttributeExpression("col2")) + .addGroupBy(createSimpleAttributeExpression("col3")) + .addAggregation( + QueryRequestBuilderUtils.createAliasedFunctionExpressionWithSimpleAttribute( + "AVG", "duration", "avg_duration")) + .build(); + Assertions.assertThrows( + IllegalArgumentException.class, + () -> + pinotBasedRequestHandler + .handleRequest(request3, new ExecutionContext("test-tenant-id", request3)) + .blockingSubscribe()); + } + @Test public void testWithMockPinotClient() throws IOException { for (Config config : serviceConfig.getConfigList("queryRequestHandlersConfig")) { diff --git a/query-service-impl/src/test/java/org/hypertrace/core/query/service/projection/ProjectionTransformationTest.java b/query-service-impl/src/test/java/org/hypertrace/core/query/service/projection/ProjectionTransformationTest.java index 374a2811..4ded70c5 100644 --- a/query-service-impl/src/test/java/org/hypertrace/core/query/service/projection/ProjectionTransformationTest.java +++ b/query-service-impl/src/test/java/org/hypertrace/core/query/service/projection/ProjectionTransformationTest.java @@ -81,6 +81,7 @@ void beforeEach() { @Test void transQueryWithComplexAttributeExpression_SingleFilter() { this.mockAttribute("server", AttributeMetadata.getDefaultInstance()); + this.mockAttribute("Span.tags", AttributeMetadata.getDefaultInstance()); Expression spanTags = createComplexAttributeExpression("Span.tags", "span.kind").build(); Filter filter = @@ -112,6 +113,7 @@ void transQueryWithComplexAttributeExpression_SingleFilter() { void transQueryWithComplexAttributeExpression_MultipleFilter() { this.mockAttribute("server", AttributeMetadata.getDefaultInstance()); this.mockAttribute("0", AttributeMetadata.getDefaultInstance()); + this.mockAttribute("Span.tags", AttributeMetadata.getDefaultInstance()); Expression spanTags1 = createComplexAttributeExpression("Span.tags", "FLAGS").build(); Expression spanTags2 = createComplexAttributeExpression("Span.tags", "span.kind").build(); @@ -155,6 +157,7 @@ void transQueryWithComplexAttributeExpression_HierarchicalFilter() { this.mockAttribute("server", AttributeMetadata.getDefaultInstance()); this.mockAttribute("0", AttributeMetadata.getDefaultInstance()); this.mockAttribute(SIMPLE_ATTRIBUTE_ID, AttributeMetadata.getDefaultInstance()); + this.mockAttribute("Span.tags", AttributeMetadata.getDefaultInstance()); Expression spanTags1 = createComplexAttributeExpression("Span.tags", "FLAGS").build(); Expression spanTags2 = createComplexAttributeExpression("Span.tags", "span.kind").build(); @@ -202,6 +205,8 @@ void transQueryWithComplexAttributeExpression_HierarchicalFilter() { @Test void transQueryWithComplexAttributeExpression_OrderByAndFilter() { this.mockAttribute("server", AttributeMetadata.getDefaultInstance()); + this.mockAttribute("Span.tags", AttributeMetadata.getDefaultInstance()); + Expression.Builder spanTag = createComplexAttributeExpression("Span.tags", "span.kind"); Filter filter = @@ -235,8 +240,31 @@ void transQueryWithComplexAttributeExpression_OrderByAndFilter() { .blockingGet()); } + @Test + void transQueryWithComplexAttributeExpression_SingleSelection() { + this.mockAttribute("Span.tags", AttributeMetadata.getDefaultInstance()); + + Expression.Builder spanTag = createComplexAttributeExpression("Span.tags", "span.kind"); + + QueryRequest originalRequest = QueryRequest.newBuilder().addSelection(spanTag).build(); + + QueryRequest expectedTransform = + QueryRequest.newBuilder() + .addSelection(spanTag) + .setFilter(createContainsKeyFilter("Span.tags", "span.kind")) + .build(); + + assertEquals( + expectedTransform, + this.projectionTransformation + .transform(originalRequest, mockTransformationContext) + .blockingGet()); + } + @Test void transQueryWithComplexAttributeExpression_SingleOrderBy() { + this.mockAttribute("Span.tags", AttributeMetadata.getDefaultInstance()); + Expression.Builder spanTag = createComplexAttributeExpression("Span.tags", "span.kind"); QueryRequest originalRequest = @@ -259,6 +287,8 @@ void transQueryWithComplexAttributeExpression_SingleOrderBy() { @Test void transQueryWithComplexAttributeExpression_MultipleOrderBy() { + this.mockAttribute("Span.tags", AttributeMetadata.getDefaultInstance()); + Expression.Builder spanTag1 = createComplexAttributeExpression("Span.tags", "span.kind"); Expression.Builder spanTag2 = createComplexAttributeExpression("Span.tags", "FLAGS"); diff --git a/query-service/build.gradle.kts b/query-service/build.gradle.kts index 51d72641..edec2e29 100644 --- a/query-service/build.gradle.kts +++ b/query-service/build.gradle.kts @@ -26,6 +26,7 @@ dependencies { } } + integrationTestImplementation("com.google.protobuf:protobuf-java-util:3.17.3") integrationTestImplementation("org.junit.jupiter:junit-jupiter-api:5.7.1") integrationTestImplementation("org.junit.jupiter:junit-jupiter-params:5.7.1") integrationTestImplementation("org.junit.jupiter:junit-jupiter-engine:5.7.1") diff --git a/query-service/src/integrationTest/java/org/hypertrace/core/query/service/QueryServiceTestUtils.java b/query-service/src/integrationTest/java/org/hypertrace/core/query/service/QueryServiceTestUtils.java index eb69b8f9..66cb6a61 100644 --- a/query-service/src/integrationTest/java/org/hypertrace/core/query/service/QueryServiceTestUtils.java +++ b/query-service/src/integrationTest/java/org/hypertrace/core/query/service/QueryServiceTestUtils.java @@ -1,6 +1,13 @@ package org.hypertrace.core.query.service; -import org.hypertrace.core.query.service.api.AttributeExpression; +import com.google.protobuf.InvalidProtocolBufferException; +import com.google.protobuf.util.JsonFormat; +import java.io.BufferedReader; +import java.io.IOException; +import java.io.InputStream; +import java.io.InputStreamReader; +import java.io.Reader; +import org.apache.kafka.common.requests.RequestContext; import org.hypertrace.core.query.service.api.ColumnIdentifier; import org.hypertrace.core.query.service.api.Expression; import org.hypertrace.core.query.service.api.Filter; @@ -8,12 +15,15 @@ import org.hypertrace.core.query.service.api.LiteralConstant; import org.hypertrace.core.query.service.api.Operator; import org.hypertrace.core.query.service.api.OrderByExpression; +import org.hypertrace.core.query.service.api.QueryRequest; import org.hypertrace.core.query.service.api.SortOrder; import org.hypertrace.core.query.service.api.Value; import org.hypertrace.core.query.service.api.ValueType; public class QueryServiceTestUtils { + private static final String REQUESTS_DIR = "attribute-expression-test-queries"; + public static Filter createFilter( String columnName, Operator op, ValueType valueType, Object valueObject) { ColumnIdentifier startTimeColumn = @@ -103,10 +113,38 @@ public static Expression createStringLiteralValueExpression(String value) { .build(); } - public static Expression.Builder createComplexAttributeExpression( - String attributeId, String subPath) { - return Expression.newBuilder() - .setAttributeExpression( - AttributeExpression.newBuilder().setAttributeId(attributeId).setSubpath(subPath)); + public static QueryRequest getAttributeExpressionQuery(QueryRequest originalQueryRequest) + throws InvalidProtocolBufferException { + // Serialize into json. + String json = + JsonFormat.printer().omittingInsignificantWhitespace().print(originalQueryRequest); + System.out.println(json); + + // Change for attribute Expression + json = json.replaceAll("columnIdentifier", "attributeExpression"); + json = json.replaceAll("columnName", "attributeId"); + System.out.println(json); + + // Deserialize and return + QueryRequest.Builder newBuilder = QueryRequest.newBuilder(); + JsonFormat.parser().merge(json, newBuilder); + return newBuilder.build(); + } + + public static QueryRequest buildQueryFromJsonFile(String filename) throws IOException { + String resourceFileName = REQUESTS_DIR + "/" + filename; + Reader requestJsonStr = readResourceFile(resourceFileName); + QueryRequest.Builder builder = QueryRequest.newBuilder(); + try { + JsonFormat.parser().merge(requestJsonStr, builder); + } catch (InvalidProtocolBufferException e) { + e.printStackTrace(); + } + return builder.build(); + } + + private static Reader readResourceFile(String fileName) { + InputStream in = RequestContext.class.getClassLoader().getResourceAsStream(fileName); + return new BufferedReader(new InputStreamReader(in)); } } diff --git a/query-service/src/integrationTest/java/org/hypertrace/core/query/service/htqueries/ExplorerQueries.java b/query-service/src/integrationTest/java/org/hypertrace/core/query/service/htqueries/ExplorerQueries.java index 21b31688..5a00dee1 100644 --- a/query-service/src/integrationTest/java/org/hypertrace/core/query/service/htqueries/ExplorerQueries.java +++ b/query-service/src/integrationTest/java/org/hypertrace/core/query/service/htqueries/ExplorerQueries.java @@ -1,10 +1,6 @@ package org.hypertrace.core.query.service.htqueries; -import static org.hypertrace.core.query.service.QueryRequestUtil.createContainsKeyFilter; -import static org.hypertrace.core.query.service.QueryServiceTestUtils.createComplexAttributeExpression; import static org.hypertrace.core.query.service.QueryServiceTestUtils.createFilter; -import static org.hypertrace.core.query.service.QueryServiceTestUtils.createOrderByExpression; -import static org.hypertrace.core.query.service.QueryServiceTestUtils.createStringLiteralValueExpression; import java.time.Duration; import org.hypertrace.core.query.service.api.ColumnIdentifier; @@ -15,7 +11,6 @@ import org.hypertrace.core.query.service.api.Operator; import org.hypertrace.core.query.service.api.QueryRequest; import org.hypertrace.core.query.service.api.QueryRequest.Builder; -import org.hypertrace.core.query.service.api.SortOrder; import org.hypertrace.core.query.service.api.Value; import org.hypertrace.core.query.service.api.ValueType; @@ -98,46 +93,4 @@ static QueryRequest buildQuery1() { builder.addGroupBy(Expression.newBuilder().setFunction(dateTimeConvert).build()); return builder.build(); } - - static QueryRequest buildQuery2() { - Builder builder = QueryRequest.newBuilder(); - - Expression apiTraceTags = - createComplexAttributeExpression("API_TRACE.tags", "span.kind").build(); - builder.addSelection(apiTraceTags); - - Filter equalFilter = - Filter.newBuilder() - .setOperator(Operator.EQ) - .setLhs(apiTraceTags) - .setRhs(createStringLiteralValueExpression("client")) - .build(); - - Filter startTimeFilter = - createFilter( - "API_TRACE.startTime", - Operator.GE, - ValueType.LONG, - System.currentTimeMillis() - Duration.ofHours(1).toMillis()); - - Filter endTimeFilter = - createFilter( - "API_TRACE.startTime", - Operator.LT, - ValueType.LONG, - System.currentTimeMillis() + Duration.ofHours(1).toMillis()); - - builder.setFilter( - Filter.newBuilder() - .setOperator(Operator.AND) - .addChildFilter(startTimeFilter) - .addChildFilter(endTimeFilter) - .addChildFilter(equalFilter) - .addChildFilter(createContainsKeyFilter("API_TRACE.tags", "span.kind")) - .build()); - - builder.addOrderBy(createOrderByExpression(apiTraceTags, SortOrder.DESC)); - - return builder.build(); - } } diff --git a/query-service/src/integrationTest/java/org/hypertrace/core/query/service/htqueries/HTPinotQueriesTest.java b/query-service/src/integrationTest/java/org/hypertrace/core/query/service/htqueries/HTPinotQueriesTest.java index f38b3f70..957aa6d8 100644 --- a/query-service/src/integrationTest/java/org/hypertrace/core/query/service/htqueries/HTPinotQueriesTest.java +++ b/query-service/src/integrationTest/java/org/hypertrace/core/query/service/htqueries/HTPinotQueriesTest.java @@ -1,16 +1,20 @@ package org.hypertrace.core.query.service.htqueries; import static com.github.stefanbirkner.systemlambda.SystemLambda.withEnvironmentVariable; +import static org.hypertrace.core.query.service.QueryServiceTestUtils.buildQueryFromJsonFile; +import static org.hypertrace.core.query.service.QueryServiceTestUtils.getAttributeExpressionQuery; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertTrue; import com.google.common.collect.ImmutableMap; import com.google.common.collect.Maps; import com.google.common.collect.Streams; +import com.google.protobuf.InvalidProtocolBufferException; import com.typesafe.config.ConfigFactory; import io.grpc.ManagedChannel; import io.grpc.ManagedChannelBuilder; import java.io.File; +import java.io.IOException; import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; @@ -19,6 +23,7 @@ import java.util.Map; import java.util.UUID; import java.util.stream.Collectors; +import java.util.stream.Stream; import org.apache.avro.file.DataFileReader; import org.apache.avro.specific.SpecificDatumReader; import org.apache.kafka.clients.admin.AdminClient; @@ -34,6 +39,7 @@ import org.hypertrace.core.attribute.service.v1.AttributeMetadataFilter; import org.hypertrace.core.datamodel.StructuredTrace; import org.hypertrace.core.kafkastreams.framework.serdes.AvroSerde; +import org.hypertrace.core.query.service.api.QueryRequest; import org.hypertrace.core.query.service.api.ResultSetChunk; import org.hypertrace.core.query.service.api.Row; import org.hypertrace.core.query.service.client.QueryServiceClient; @@ -42,6 +48,9 @@ import org.junit.jupiter.api.AfterAll; import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import org.testcontainers.containers.GenericContainer; @@ -313,11 +322,12 @@ private static void validateRows(List rows, double divisor) { }); } - @Test - public void testServicesQueries() { + @ParameterizedTest + @MethodSource("provideQueryRequestForServiceQueries") + public void testServicesQueries(QueryRequest queryRequest) { LOG.info("Services queries"); Iterator itr = - queryServiceClient.executeQuery(ServicesQueries.buildQuery1(), TENANT_ID_MAP, 10000); + queryServiceClient.executeQuery(queryRequest, TENANT_ID_MAP, 10000); List list = Streams.stream(itr).collect(Collectors.toList()); List rows = list.get(0).getRowList(); assertEquals(4, rows.size()); @@ -369,11 +379,12 @@ public void testServicesQueriesForAvgRateWithTimeAggregation() { validateRows(rows, 15); } - @Test - public void testBackendsQueries() { + @ParameterizedTest + @MethodSource("provideQueryRequestForBackendQueries") + public void testBackendsQueries(QueryRequest queryRequest) { LOG.info("Backends queries"); Iterator itr = - queryServiceClient.executeQuery(BackendsQueries.buildQuery1(), TENANT_ID_MAP, 10000); + queryServiceClient.executeQuery(queryRequest, TENANT_ID_MAP, 10000); List list = Streams.stream(itr).collect(Collectors.toList()); List rows = list.get(0).getRowList(); assertEquals(1, rows.size()); @@ -382,11 +393,12 @@ public void testBackendsQueries() { assertTrue(backendNames.isEmpty()); } - @Test - public void testExplorerQueries() { + @ParameterizedTest + @MethodSource("provideQueryRequestForExplorerQueries") + public void testExplorerQueries(QueryRequest queryRequest) { LOG.info("Explorer queries"); Iterator itr = - queryServiceClient.executeQuery(ExplorerQueries.buildQuery1(), TENANT_ID_MAP, 10000); + queryServiceClient.executeQuery(queryRequest, TENANT_ID_MAP, 10000); List list = Streams.stream(itr).collect(Collectors.toList()); List rows = list.get(0).getRowList(); assertEquals(1, rows.size()); @@ -394,14 +406,47 @@ public void testExplorerQueries() { assertEquals("13", rows.get(0).getColumn(1).getString()); } - @Test - public void testExplorerQueriesForAttributeExpression() { - LOG.info("Explorer queries for attribute expression"); + @ParameterizedTest + @MethodSource("provideQueryRequestForAttributeExpressionQueries") + public void testAttributeExpressionQueries( + QueryRequest queryRequest, int rowSize, String expectedValue) { + LOG.info("Attribute Expression queries"); Iterator itr = - queryServiceClient.executeQuery(ExplorerQueries.buildQuery2(), TENANT_ID_MAP, 10000); + queryServiceClient.executeQuery(queryRequest, TENANT_ID_MAP, 10000); List list = Streams.stream(itr).collect(Collectors.toList()); List rows = list.get(0).getRowList(); - assertEquals(10, rows.size()); - assertEquals("client", rows.get(0).getColumn(0).getString()); + assertEquals(rowSize, rows.size()); + assertEquals(expectedValue, rows.get(0).getColumn(0).getString()); + } + + private static Stream provideQueryRequestForServiceQueries() + throws InvalidProtocolBufferException { + QueryRequest queryRequest1 = ServicesQueries.buildQuery1(); + return Stream.of( + Arguments.arguments(queryRequest1), + Arguments.arguments(getAttributeExpressionQuery(queryRequest1))); + } + + private static Stream provideQueryRequestForBackendQueries() + throws InvalidProtocolBufferException { + QueryRequest queryRequest1 = BackendsQueries.buildQuery1(); + return Stream.of( + Arguments.arguments(queryRequest1), + Arguments.arguments(getAttributeExpressionQuery(queryRequest1))); + } + + private static Stream provideQueryRequestForExplorerQueries() + throws InvalidProtocolBufferException { + QueryRequest queryRequest1 = ExplorerQueries.buildQuery1(); + return Stream.of( + Arguments.arguments(queryRequest1), + Arguments.arguments(getAttributeExpressionQuery(queryRequest1))); + } + + private static Stream provideQueryRequestForAttributeExpressionQueries() + throws IOException { + return Stream.of( + Arguments.arguments(buildQueryFromJsonFile("query1.json"), 10, "server"), + Arguments.arguments(buildQueryFromJsonFile("query2.json"), 2, "server")); } } diff --git a/query-service/src/integrationTest/resources/attribute-expression-test-queries/query1.json b/query-service/src/integrationTest/resources/attribute-expression-test-queries/query1.json new file mode 100644 index 00000000..5134a4e1 --- /dev/null +++ b/query-service/src/integrationTest/resources/attribute-expression-test-queries/query1.json @@ -0,0 +1,49 @@ +{ + "filter": { + "childFilter": [{ + "lhs": { + "attributeExpression": { + "attributeId": "EVENT.spanTags" + } + }, + "operator": "CONTAINS_KEY", + "rhs": { + "literal": { + "value": { + "string": "span.kind" + } + } + } + }, { + "lhs": { + "attributeExpression": { + "attributeId": "EVENT.spanTags", + "subpath": "span.kind" + } + }, + "operator": "GE", + "rhs": { + "literal": { + "value": { + "string": "client" + } + } + } + }] + }, + "selection": [{ + "attributeExpression": { + "attributeId": "EVENT.spanTags", + "subpath": "span.kind" + } + }], + "orderBy": [{ + "expression": { + "attributeExpression": { + "attributeId": "EVENT.spanTags", + "subpath": "span.kind" + } + }, + "order": "DESC" + }] +} \ No newline at end of file diff --git a/query-service/src/integrationTest/resources/attribute-expression-test-queries/query2.json b/query-service/src/integrationTest/resources/attribute-expression-test-queries/query2.json new file mode 100644 index 00000000..5b746cf9 --- /dev/null +++ b/query-service/src/integrationTest/resources/attribute-expression-test-queries/query2.json @@ -0,0 +1,85 @@ +{ + "filter": { + "childFilter": [{ + "lhs": { + "attributeExpression": { + "attributeId": "EVENT.startTime" + } + }, + "operator": "GT", + "rhs": { + "literal": { + "value": { + "valueType": "LONG", + "long": "1570658506605" + } + } + } + }, { + "lhs": { + "attributeExpression": { + "attributeId": "EVENT.startTime" + } + }, + "operator": "LT", + "rhs": { + "literal": { + "value": { + "valueType": "LONG", + "long": "2570744906673" + } + } + } + }, { + "lhs": { + "attributeExpression": { + "attributeId": "EVENT.spanTags" + } + }, + "operator": "CONTAINS_KEY", + "rhs": { + "literal": { + "value": { + "string": "span.kind" + } + } + } + }, { + "lhs": { + "attributeExpression": { + "attributeId": "EVENT.spanTags", + "subpath": "span.kind" + } + }, + "operator": "NEQ", + "rhs": { + "literal": { + "value": { + } + } + } + }] + }, + "selection": [{ + "function": { + "functionName": "AVG", + "arguments": [{ + "attributeExpression": { + "attributeId": "EVENT.duration" + } + }], + "alias": "avg_duration" + } + }, { + "attributeExpression": { + "attributeId": "EVENT.spanTags", + "subpath": "span.kind" + } + }], + "groupBy": [{ + "attributeExpression": { + "attributeId": "EVENT.spanTags", + "subpath": "span.kind" + } + }] +} \ No newline at end of file