chore | adding the support for function as an order by argument (#293)

* chore | adding the support for function as an order by argument

Author: aman-bansal

entity-service-impl/src/main/java/org/hypertrace/entity/query/service/converter/OrderByConverter.java

@@ -3,6 +3,7 @@
 import static com.google.common.base.Suppliers.memoize;
 import static java.util.Collections.unmodifiableMap;
 import static org.hypertrace.entity.query.service.v1.Expression.ValueCase.COLUMNIDENTIFIER;
+import static org.hypertrace.entity.query.service.v1.Expression.ValueCase.FUNCTION;
 import static org.hypertrace.entity.query.service.v1.SortOrder.ASC;
 import static org.hypertrace.entity.query.service.v1.SortOrder.DESC;
 
@@ -24,15 +25,18 @@
 import org.hypertrace.entity.query.service.v1.ColumnIdentifier;
 import org.hypertrace.entity.query.service.v1.Expression;
 import org.hypertrace.entity.query.service.v1.Expression.ValueCase;
+import org.hypertrace.entity.query.service.v1.Function;
 import org.hypertrace.entity.query.service.v1.OrderByExpression;
 
 @Singleton
 @AllArgsConstructor(onConstructor_ = {@Inject})
 public class OrderByConverter implements Converter<List<OrderByExpression>, Sort> {
+
   private static final Supplier<Map<org.hypertrace.entity.query.service.v1.SortOrder, SortOrder>>
       ORDER_MAP = memoize(OrderByConverter::getMapping);
   private final OneOfAccessor<Expression, ValueCase> expressionAccessor;
   private final AliasProvider<ColumnIdentifier> aliasProvider;
+  private final AliasProvider<Function> functionAliasProvider;
 
   @Override
   public Sort convert(final List<OrderByExpression> orders, final RequestContext requestContext)
@@ -50,14 +54,24 @@ public Sort convert(final List<OrderByExpression> orders, final RequestContext r
   private SortingSpec convert(final OrderByExpression orderBy) throws ConversionException {
     final Expression innerExpression = orderBy.getExpression();
 
-    final ColumnIdentifier identifier =
-        expressionAccessor.access(
-            innerExpression, innerExpression.getValueCase(), Set.of(COLUMNIDENTIFIER));
-    final String alias = aliasProvider.getAlias(identifier);
-    final IdentifierExpression identifierExpression = IdentifierExpression.of(alias);
+    Object expression =
+        this.expressionAccessor.access(
+            innerExpression, innerExpression.getValueCase(), Set.of(COLUMNIDENTIFIER, FUNCTION));
+    String alias;
+    switch (innerExpression.getValueCase()) {
+      case COLUMNIDENTIFIER:
+        alias = aliasProvider.getAlias((ColumnIdentifier) expression);
+        break;
+      case FUNCTION:
+        alias = functionAliasProvider.getAlias((Function) expression);
+        break;
+      default:
+        throw new ConversionException(
+            String.format("Unsupported expression type %s", innerExpression.getValueCase()));
+    }
 
+    final IdentifierExpression identifierExpression = IdentifierExpression.of(alias);
     final SortOrder order = ORDER_MAP.get().get(orderBy.getOrder());
-
     if (order == null) {
       throw new ConversionException(
           String.format("Cannot find sorting order %s", orderBy.getOrder()));

entity-service-impl/src/main/java/org/hypertrace/entity/query/service/converter/identifier/FunctionAliasProvider.java

@@ -0,0 +1,22 @@
+package org.hypertrace.entity.query.service.converter.identifier;
+
+import com.google.inject.Inject;
+import com.google.inject.Singleton;
+import lombok.AllArgsConstructor;
+import org.hypertrace.entity.query.service.converter.AliasProvider;
+import org.hypertrace.entity.query.service.v1.Function;
+import org.hypertrace.entity.service.util.StringUtils;
+
+@Singleton
+@AllArgsConstructor(onConstructor_ = {@Inject})
+public class FunctionAliasProvider implements AliasProvider<Function> {
+
+  @Override
+  public String getAlias(final Function function) {
+    if (StringUtils.isNotBlank(function.getAlias())) {
+      return function.getAlias();
+    }
+
+    throw new IllegalArgumentException("Function alias is missing");
+  }
+}

entity-service-impl/src/test/java/org/hypertrace/entity/query/service/converter/OrderByConverterTest.java

@@ -2,6 +2,7 @@
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.mockito.quality.Strictness.LENIENT;
 
 import java.util.List;
@@ -12,10 +13,12 @@
 import org.hypertrace.core.grpcutils.context.RequestContext;
 import org.hypertrace.entity.query.service.converter.accessor.ExpressionOneOfAccessor;
 import org.hypertrace.entity.query.service.converter.accessor.OneOfAccessor;
+import org.hypertrace.entity.query.service.converter.identifier.FunctionAliasProvider;
 import org.hypertrace.entity.query.service.converter.identifier.IdentifierAliasProvider;
 import org.hypertrace.entity.query.service.v1.ColumnIdentifier;
 import org.hypertrace.entity.query.service.v1.Expression;
 import org.hypertrace.entity.query.service.v1.Expression.ValueCase;
+import org.hypertrace.entity.query.service.v1.Function;
 import org.hypertrace.entity.query.service.v1.OrderByExpression;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
@@ -35,6 +38,7 @@ class OrderByConverterTest {
   private Converter<List<OrderByExpression>, Sort> orderByConverter;
 
   private final AliasProvider<ColumnIdentifier> aliasProvider = new IdentifierAliasProvider();
+  private final AliasProvider<Function> functionAliasProvider = new FunctionAliasProvider();
   private final ColumnIdentifier columnIdentifier =
       ColumnIdentifier.newBuilder().setColumnName("Planet_Mars").build();
   private final IdentifierExpression identifierExpression = IdentifierExpression.of("Planet_Mars");
@@ -47,7 +51,8 @@ class OrderByConverterTest {
   @BeforeEach
   void setup() {
     OneOfAccessor<Expression, ValueCase> expressionAccessor = new ExpressionOneOfAccessor();
-    orderByConverter = new OrderByConverter(expressionAccessor, aliasProvider);
+    orderByConverter =
+        new OrderByConverter(expressionAccessor, aliasProvider, functionAliasProvider);
   }
 
   @Test
@@ -57,6 +62,32 @@ void testConvert() throws ConversionException {
     assertEquals(expected, orderByConverter.convert(List.of(orderByExpression), requestContext));
   }
 
+  @Test
+  void testConvert_functionExpression() throws ConversionException {
+    IdentifierExpression aliasExpression = IdentifierExpression.of("agg_count");
+    OrderByExpression orderByExpressionWithAlias =
+        OrderByExpression.newBuilder()
+            .setExpression(
+                Expression.newBuilder()
+                    .setFunction(Function.newBuilder().setAlias("agg_count").build()))
+            .setOrder(org.hypertrace.entity.query.service.v1.SortOrder.ASC)
+            .build();
+
+    Sort expected =
+        Sort.builder().sortingSpec(SortingSpec.of(aliasExpression, SortOrder.ASC)).build();
+    assertEquals(
+        expected, orderByConverter.convert(List.of(orderByExpressionWithAlias), requestContext));
+
+    OrderByExpression orderByExpressionWithoutAlias =
+        OrderByExpression.newBuilder()
+            .setExpression(Expression.newBuilder().setFunction(Function.newBuilder().build()))
+            .setOrder(org.hypertrace.entity.query.service.v1.SortOrder.ASC)
+            .build();
+    assertThrows(
+        IllegalArgumentException.class,
+        () -> orderByConverter.convert(List.of(orderByExpressionWithoutAlias), requestContext));
+  }
+
   @ParameterizedTest
   @EnumSource(
       value = org.hypertrace.entity.query.service.v1.SortOrder.class,