Adding sum support for annotated queries

sdk/spring/azure-spring-data-cosmos/src/main/java/com/azure/spring/data/cosmos/core/CosmosOperations.java

@@ -367,6 +367,16 @@ public interface CosmosOperations {
      */
     <T> long count(SqlQuerySpec querySpec, String containerName);
 
+    /**
+     * Sum
+     *
+     * @param querySpec the document query spec
+     * @param containerName the container name
+     * @param <T> type class of domainType
+     * @return sum result
+     */
+    <T> long sum(SqlQuerySpec querySpec, String containerName);
+
     /**
      * To get converter
      *

sdk/spring/azure-spring-data-cosmos/src/main/java/com/azure/spring/data/cosmos/core/CosmosTemplate.java

@@ -1008,7 +1008,7 @@ private <T> Page<T> paginationQuery(SqlQuerySpec querySpec, SqlQuerySpec countQu
                                         Optional<Object> partitionKeyValue) {
         containerName = getContainerNameOverride(containerName);
         Slice<T> response = sliceQuery(querySpec, pageable, sort, returnType, containerName, partitionKeyValue);
-        final long total = getCountValue(countQuerySpec, containerName);
+        final long total = getNumericValue(countQuerySpec, containerName);
         return new CosmosPageImpl<>(response.getContent(), response.getPageable(), total);
     }
 
@@ -1132,12 +1132,21 @@ public <T> long count(CosmosQuery query, String containerName) {
 
     @Override
     public <T> long count(SqlQuerySpec querySpec, String containerName) {
+        return this.numeric(querySpec, containerName);
+    }
+
+    @Override
+    public <T> long sum(SqlQuerySpec querySpec, String containerName) {
+        return this.numeric(querySpec, containerName);
+    }
+
+    private <T> long numeric(SqlQuerySpec querySpec, String containerName) {
         containerName = getContainerNameOverride(containerName);
         Assert.hasText(containerName, "container name should not be empty");
 
-        final Long count = getCountValue(querySpec, containerName);
-        assert count != null;
-        return count;
+        final Long numericResult = getNumericValue(querySpec, containerName);
+        assert numericResult != null;
+        return numericResult;
     }
 
     @Override
@@ -1167,10 +1176,10 @@ private void markAuditedIfConfigured(Object object) {
 
     private Long getCountValue(CosmosQuery query, String containerName) {
         final SqlQuerySpec querySpec = new CountQueryGenerator().generateCosmos(query);
-        return getCountValue(querySpec, containerName);
+        return getNumericValue(querySpec, containerName);
     }
 
-    private Long getCountValue(SqlQuerySpec querySpec, String containerName) {
+    private Long getNumericValue(SqlQuerySpec querySpec, String containerName) {
         final CosmosQueryRequestOptions options = new CosmosQueryRequestOptions();
         options.setQueryMetricsEnabled(this.queryMetricsEnabled);
         options.setIndexMetricsEnabled(this.indexMetricsEnabled);

sdk/spring/azure-spring-data-cosmos/src/main/java/com/azure/spring/data/cosmos/repository/support/StringBasedCosmosQuery.java

@@ -31,6 +31,7 @@
  */
 public class StringBasedCosmosQuery extends AbstractCosmosQuery {
     private static final Pattern COUNT_QUERY_PATTERN = Pattern.compile("^\\s*select\\s+value\\s+count.*", Pattern.CASE_INSENSITIVE);
+    private static final Pattern SUM_QUERY_PATTERN = Pattern.compile("^\\s*select\\s+value\\s+sum.*", Pattern.CASE_INSENSITIVE);
 
     private final String query;
 
@@ -104,9 +105,12 @@ public Object execute(final Object[] parameters) {
         } else if (isCountQuery()) {
             final String container = ((CosmosEntityMetadata<?>) getQueryMethod().getEntityInformation()).getContainerName();
             return this.operations.count(querySpec, container);
+        } else if (isSumQuery()) {
+            final String container = ((CosmosEntityMetadata<?>) getQueryMethod().getEntityInformation()).getContainerName();
+            return this.operations.sum(querySpec, container);
         } else {
             return this.operations.runQuery(querySpec, accessor.getSort(), processor.getReturnedType().getDomainType(),
-                                            processor.getReturnedType().getReturnedType());
+                processor.getReturnedType().getReturnedType());
         }
     }
 
@@ -129,15 +133,27 @@ protected boolean isCountQuery() {
         return isCountQuery(query, getQueryMethod().getReturnedObjectType());
     }
 
+    protected boolean isSumQuery() {
+        return isSumQuery(query, getQueryMethod().getReturnedObjectType());
+    }
+
     static boolean isCountQuery(String query, Class<?> returnedType) {
-        if (isCountQueryReturnType(returnedType)) {
+        if (isNumericQueryReturnType(returnedType)) {
             return COUNT_QUERY_PATTERN.matcher(query).matches();
         } else {
             return false;
         }
     }
 
-    private static boolean isCountQueryReturnType(Class<?> returnedType) {
+    static boolean isSumQuery(String query, Class<?> returnedType) {
+        if (isNumericQueryReturnType(returnedType)) {
+            return SUM_QUERY_PATTERN.matcher(query).matches();
+        } else {
+            return false;
+        }
+    }
+
+    private static boolean isNumericQueryReturnType(Class<?> returnedType) {
         return returnedType == Long.class
             || returnedType == long.class
             || returnedType == Integer.class

sdk/spring/azure-spring-data-cosmos/src/test/java/com/azure/spring/data/cosmos/repository/integration/AnnotatedQueryIT.java

@@ -248,6 +248,18 @@ public void testAnnotatedQueryWithValueAsList() {
         assertAddressPostalCodesUnordered(postalCodes, addresses);
     }
 
+    @Test
+    public void testAnnotatedQueryWithSum() {
+        Address.TEST_ADDRESS1_PARTITION1.setLongId(1L);
+        Address.TEST_ADDRESS2_PARTITION1.setLongId(2L);
+        final List<Address> addresses = Arrays.asList(Address.TEST_ADDRESS1_PARTITION1, Address.TEST_ADDRESS2_PARTITION1);
+        addressRepository.saveAll(addresses);
+
+        final Long sumResult = addressRepository.annotatedSumLongIdValuesByCity(TestConstants.CITY);
+
+        assertThat(sumResult).isEqualTo(3);
+    }
+
     @Test
     public void testAnnotatedQueryWithJsonNodeAsPage() {
         final List<Address> addresses = Arrays.asList(Address.TEST_ADDRESS1_PARTITION1, Address.TEST_ADDRESS2_PARTITION1);

sdk/spring/azure-spring-data-cosmos/src/test/java/com/azure/spring/data/cosmos/repository/repository/AddressRepository.java

@@ -59,6 +59,9 @@ public interface AddressRepository extends CosmosRepository<Address, String> {
     @Query("select DISTINCT value a.postalCode from a where a.city = @city")
     List<String> annotatedFindPostalCodeValuesByCity(@Param("city") String city);
 
+    @Query("SELECT VALUE SUM(a.longId) from a where a.city = @city")
+    Long annotatedSumLongIdValuesByCity(@Param("city") String city);
+
     @Query(value = "select * from a where a.city IN (@cities)")
     List<Address> annotatedFindByCityIn(@Param("cities") List<String> cities, Sort sort);