Query annotation support (Azure#14661)

* This PR adds support for query annotations. This enables end users to add annotated queries to their repositories using @query(value = "<query>") which opens up all the query capabilities of underlying cosmos java sdk to the cosmos spring driver.
AzureSDKAutomation · Sep 9, 2020 · 32f9ae0 · 32f9ae0
1 parent 2521ad6
commit 32f9ae0
Show file tree

Hide file tree

Showing 18 changed files with 477 additions and 53 deletions.
diff --git a/...ng-data-cosmos-core/src/main/java/com/azure/spring/data/cosmos/core/CosmosOperations.java b/...ng-data-cosmos-core/src/main/java/com/azure/spring/data/cosmos/core/CosmosOperations.java
@@ -5,6 +5,7 @@
 
 import com.azure.cosmos.models.CosmosContainerProperties;
 import com.azure.cosmos.models.PartitionKey;
+import com.azure.cosmos.models.SqlQuerySpec;
 import com.azure.spring.data.cosmos.core.convert.MappingCosmosConverter;
 import com.azure.spring.data.cosmos.core.query.CosmosQuery;
 import com.azure.spring.data.cosmos.repository.support.CosmosEntityInformation;
@@ -274,4 +275,15 @@ public interface CosmosOperations {
      * @return MappingCosmosConverter
      */
     MappingCosmosConverter getConverter();
+
+    /**
+     * Run the query.
+     *
+     * @param <T> the type parameter
+     * @param querySpec the query spec
+     * @param domainType the domain type
+     * @param returnType the return type
+     * @return the Iterable
+     */
+    <T> Iterable<T> runQuery(SqlQuerySpec querySpec, Class<?> domainType, Class<T> returnType);
 }
diff --git a/...ring-data-cosmos-core/src/main/java/com/azure/spring/data/cosmos/core/CosmosTemplate.java b/...ring-data-cosmos-core/src/main/java/com/azure/spring/data/cosmos/core/CosmosTemplate.java
@@ -720,6 +720,13 @@ public MappingCosmosConverter getConverter() {
         return this.mappingCosmosConverter;
     }
 
+    @Override
+    public <T> Iterable<T> runQuery(SqlQuerySpec querySpec, Class<?> domainType, Class<T> returnType) {
+        return getJsonNodeFluxFromQuerySpec(domainType.getSimpleName(), querySpec, returnType)
+                   .collectList()
+                   .block();
+    }
+
     private JsonNode prepareToPersistAndConvertToItemProperties(Object object) {
         if (cosmosAuditingHandler != null) {
             cosmosAuditingHandler.markAudited(object);
@@ -766,13 +773,35 @@ private Flux<JsonNode> findItemsAsFlux(@NonNull CosmosQuery query,
             .publishOn(Schedulers.parallel())
             .flatMap(cosmosItemFeedResponse -> {
                 CosmosUtils.fillAndProcessResponseDiagnostics(this.responseDiagnosticsProcessor,
-                    cosmosItemFeedResponse.getCosmosDiagnostics(), cosmosItemFeedResponse);
+                                                              cosmosItemFeedResponse.getCosmosDiagnostics(),
+                                                              cosmosItemFeedResponse);
                 return Flux.fromIterable(cosmosItemFeedResponse.getResults());
             })
             .onErrorResume(throwable ->
                 CosmosExceptionUtils.exceptionHandler("Failed to find items", throwable));
     }
 
+    private <T> Flux<T> getJsonNodeFluxFromQuerySpec(
+        @NonNull String containerName, SqlQuerySpec sqlQuerySpec, Class<T> classType) {
+        final CosmosQueryRequestOptions cosmosQueryRequestOptions = new CosmosQueryRequestOptions();
+        cosmosQueryRequestOptions.setQueryMetricsEnabled(this.queryMetricsEnabled);
+
+        return cosmosAsyncClient
+                   .getDatabase(this.databaseName)
+                   .getContainer(containerName)
+                   .queryItems(sqlQuerySpec, cosmosQueryRequestOptions, classType)
+                   .byPage()
+                   .publishOn(Schedulers.parallel())
+                   .flatMap(cosmosItemFeedResponse -> {
+                       CosmosUtils.fillAndProcessResponseDiagnostics(this.responseDiagnosticsProcessor,
+                                                                     cosmosItemFeedResponse.getCosmosDiagnostics(),
+                                                                     cosmosItemFeedResponse);
+                       return Flux.fromIterable(cosmosItemFeedResponse.getResults());
+                   })
+                   .onErrorResume(throwable ->
+                                      CosmosExceptionUtils.exceptionHandler("Failed to find items", throwable));
+    }
+
     private List<JsonNode> findItems(@NonNull CosmosQuery query,
                                      @NonNull String containerName) {
         return findItemsAsFlux(query, containerName)

diff --git a/...cosmos-core/src/main/java/com/azure/spring/data/cosmos/core/ReactiveCosmosOperations.java b/...cosmos-core/src/main/java/com/azure/spring/data/cosmos/core/ReactiveCosmosOperations.java
@@ -5,6 +5,7 @@
 
 import com.azure.cosmos.models.CosmosContainerResponse;
 import com.azure.cosmos.models.PartitionKey;
+import com.azure.cosmos.models.SqlQuerySpec;
 import com.azure.spring.data.cosmos.core.convert.MappingCosmosConverter;
 import com.azure.spring.data.cosmos.core.query.CosmosQuery;
 import com.azure.spring.data.cosmos.repository.support.CosmosEntityInformation;
@@ -242,4 +243,15 @@ public interface ReactiveCosmosOperations {
      * @return MappingCosmosConverter
      */
     MappingCosmosConverter getConverter();
+
+    /**
+     * Run the query.
+     *
+     * @param <T> the type parameter
+     * @param querySpec the query spec
+     * @param domainType the domain type
+     * @param returnType the return type
+     * @return the flux
+     */
+    <T> Flux<T> runQuery(SqlQuerySpec querySpec, Class<?> domainType, Class<T> returnType);
 }
diff --git a/...a-cosmos-core/src/main/java/com/azure/spring/data/cosmos/core/ReactiveCosmosTemplate.java b/...a-cosmos-core/src/main/java/com/azure/spring/data/cosmos/core/ReactiveCosmosTemplate.java
@@ -578,6 +578,26 @@ public MappingCosmosConverter getConverter() {
         return mappingCosmosConverter;
     }
 
+    @Override
+    public <T> Flux<T> runQuery(SqlQuerySpec querySpec, Class<?> domainType, Class<T> returnType) {
+        String containerName = domainType.getSimpleName();
+        CosmosQueryRequestOptions options = new CosmosQueryRequestOptions();
+        return cosmosAsyncClient.getDatabase(this.databaseName)
+                   .getContainer(containerName)
+                   .queryItems(querySpec, options, returnType)
+                   .byPage()
+                   .publishOn(Schedulers.parallel())
+                   .flatMap(cosmosItemFeedResponse -> {
+                       CosmosUtils
+                           .fillAndProcessResponseDiagnostics(this.responseDiagnosticsProcessor,
+                                                              cosmosItemFeedResponse.getCosmosDiagnostics(),
+                                                              cosmosItemFeedResponse);
+                       return Flux.fromIterable(cosmosItemFeedResponse.getResults());
+                   })
+                   .onErrorResume(throwable ->
+                                      CosmosExceptionUtils.exceptionHandler("Failed to find items", throwable));
+    }
+
     private Mono<Long> getCountValue(CosmosQuery query, String containerName) {
         final SqlQuerySpec querySpec = new CountQueryGenerator().generateCosmos(query);
         final CosmosQueryRequestOptions options = new CosmosQueryRequestOptions();

diff --git a/...-spring-data-cosmos-core/src/main/java/com/azure/spring/data/cosmos/repository/Query.java b/...-spring-data-cosmos-core/src/main/java/com/azure/spring/data/cosmos/repository/Query.java
@@ -0,0 +1,24 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+package com.azure.spring.data.cosmos.repository;
+import org.springframework.data.annotation.QueryAnnotation;
+
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+
+/**
+ *  Annotation to declare finder queries directly on repository methods.
+ */
+@Retention(RetentionPolicy.RUNTIME)
+@Target({ElementType.METHOD, ElementType.ANNOTATION_TYPE})
+@QueryAnnotation
+public @interface Query {
+    /**
+     * value of the query
+     *
+     * @return the value
+     */
+    String value() default "";
+}
diff --git a/...core/src/main/java/com/azure/spring/data/cosmos/repository/query/AbstractCosmosQuery.java b/...core/src/main/java/com/azure/spring/data/cosmos/repository/query/AbstractCosmosQuery.java
@@ -14,7 +14,7 @@
 public abstract class AbstractCosmosQuery implements RepositoryQuery {
 
     private final CosmosQueryMethod method;
-    private final CosmosOperations operations;
+    protected final CosmosOperations operations;
 
     /**
      * Initialization
@@ -33,6 +33,7 @@ public AbstractCosmosQuery(CosmosQueryMethod method, CosmosOperations operations
      * @param parameters must not be {@literal null}.
      * @return execution result. Can be {@literal null}.
      */
+    @Override
     public Object execute(Object[] parameters) {
         final CosmosParameterAccessor accessor = new CosmosParameterParameterAccessor(method, parameters);
         final CosmosQuery query = createQuery(accessor);

diff --git a/.../main/java/com/azure/spring/data/cosmos/repository/query/AbstractReactiveCosmosQuery.java b/.../main/java/com/azure/spring/data/cosmos/repository/query/AbstractReactiveCosmosQuery.java
@@ -15,7 +15,7 @@
 public abstract class AbstractReactiveCosmosQuery implements RepositoryQuery {
 
     private final ReactiveCosmosQueryMethod method;
-    private final ReactiveCosmosOperations operations;
+    protected final ReactiveCosmosOperations operations;
 
     /**
      * Initialization
@@ -35,6 +35,7 @@ public AbstractReactiveCosmosQuery(ReactiveCosmosQueryMethod method,
      * @param parameters must not be {@literal null}.
      * @return execution result. Can be {@literal null}.
      */
+    @Override
     public Object execute(Object[] parameters) {
         final ReactiveCosmosParameterAccessor accessor =
             new ReactiveCosmosParameterParameterAccessor(method, parameters);

diff --git a/...s-core/src/main/java/com/azure/spring/data/cosmos/repository/query/CosmosQueryMethod.java b/...s-core/src/main/java/com/azure/spring/data/cosmos/repository/query/CosmosQueryMethod.java
@@ -2,20 +2,26 @@
 // Licensed under the MIT License.
 package com.azure.spring.data.cosmos.repository.query;
 
+import com.azure.spring.data.cosmos.repository.Query;
 import com.azure.spring.data.cosmos.repository.support.CosmosEntityInformation;
+import org.springframework.core.annotation.AnnotatedElementUtils;
 import org.springframework.data.projection.ProjectionFactory;
 import org.springframework.data.repository.core.EntityMetadata;
 import org.springframework.data.repository.core.RepositoryMetadata;
 import org.springframework.data.repository.query.QueryMethod;
+import org.springframework.lang.Nullable;
+import org.springframework.util.StringUtils;
 
 import java.lang.reflect.Method;
+import java.util.Optional;
 
 /**
  * Inherit QueryMethod class to generate a method that is designated to execute a finder query.
  */
 public class CosmosQueryMethod extends QueryMethod {
 
     private CosmosEntityMetadata<?> metadata;
+    private final String annotatedQueryValue;
 
     /**
      * Creates a new {@link CosmosQueryMethod} from the given parameters. Looks up the correct query to use
@@ -27,6 +33,7 @@ public class CosmosQueryMethod extends QueryMethod {
      */
     public CosmosQueryMethod(Method method, RepositoryMetadata metadata, ProjectionFactory factory) {
         super(method, metadata, factory);
+        this.annotatedQueryValue = findAnnotatedQuery(method).orElse(null);
     }
 
     @Override
@@ -39,4 +46,32 @@ public EntityMetadata<?> getEntityInformation() {
         this.metadata = new SimpleCosmosEntityMetadata<Object>(domainType, entityInformation);
         return this.metadata;
     }
+
+    /**
+     * Returns whether the method has an annotated query.
+     *
+     * @return if the query method has an annotated query
+     */
+    public boolean hasAnnotatedQuery() {
+        return annotatedQueryValue != null;
+    }
+
+    /**
+     * Returns the query string declared in a {@link Query} annotation or {@literal null} if neither the annotation
+     * found
+     * nor the attribute was specified.
+     *
+     * @return the query string or null
+     */
+    @Nullable
+    public String getQueryAnnotation() {
+        return annotatedQueryValue;
+    }
+
+    private Optional<String> findAnnotatedQuery(Method method) {
+        return Optional.ofNullable(AnnotatedElementUtils.findMergedAnnotation(method, Query.class))
+                   .map(Query::value)
+                   .filter(StringUtils::hasText);
+    }
+
 }
diff --git a/...rc/main/java/com/azure/spring/data/cosmos/repository/query/ReactiveCosmosQueryMethod.java b/...rc/main/java/com/azure/spring/data/cosmos/repository/query/ReactiveCosmosQueryMethod.java
@@ -2,15 +2,20 @@
 // Licensed under the MIT License.
 package com.azure.spring.data.cosmos.repository.query;
 
+import com.azure.spring.data.cosmos.repository.Query;
 import com.azure.spring.data.cosmos.repository.support.CosmosEntityInformation;
+import org.springframework.core.annotation.AnnotatedElementUtils;
 import org.springframework.data.projection.ProjectionFactory;
 import org.springframework.data.repository.core.EntityMetadata;
 import org.springframework.data.repository.core.RepositoryMetadata;
 import org.springframework.data.repository.query.QueryMethod;
+import org.springframework.lang.Nullable;
+import org.springframework.util.StringUtils;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 
 import java.lang.reflect.Method;
+import java.util.Optional;
 
 /**
  * Inherit from QueryMethod class to execute a finder query.
@@ -19,6 +24,7 @@ public class ReactiveCosmosQueryMethod extends QueryMethod {
 
     private ReactiveCosmosEntityMetadata<?> metadata;
     private final Method method;
+    private final String annotatedQueryValue;
 
     /**
      * Creates a new {@link QueryMethod} from the given parameters. Looks up the correct query to use for following
@@ -31,6 +37,7 @@ public class ReactiveCosmosQueryMethod extends QueryMethod {
     public ReactiveCosmosQueryMethod(Method method, RepositoryMetadata metadata, ProjectionFactory factory) {
         super(method, metadata, factory);
         this.method = method;
+        this.annotatedQueryValue = findAnnotatedQuery(method).orElse(null);
     }
 
     @Override
@@ -56,4 +63,28 @@ public Class<?> getReactiveWrapper() {
     private static boolean isReactiveWrapperClass(Class<?> clazz) {
         return clazz.equals(Flux.class) || clazz.equals(Mono.class);
     }
+
+    /**
+     * Returns whether the method has an annotated query.
+     * @return if the query method has an annotated query
+     */
+    public boolean hasAnnotatedQuery() {
+        return annotatedQueryValue != null;
+    }
+
+    /**
+     * Gets the annotated query or returns null
+     * @return the annotated query String or null
+     */
+    @Nullable
+    public String getQueryAnnotation() {
+        return annotatedQueryValue;
+    }
+
+    private Optional<String> findAnnotatedQuery(Method method) {
+        return Optional.ofNullable(AnnotatedElementUtils.findMergedAnnotation(method, Query.class))
+                   .map(Query::value)
+                   .filter(StringUtils::hasText);
+    }
+
 }
diff --git a/...rc/main/java/com/azure/spring/data/cosmos/repository/support/CosmosRepositoryFactory.java b/...rc/main/java/com/azure/spring/data/cosmos/repository/support/CosmosRepositoryFactory.java
@@ -74,7 +74,12 @@ public RepositoryQuery resolveQuery(Method method, RepositoryMetadata metadata,
 
             Assert.notNull(queryMethod, "queryMethod must not be null!");
             Assert.notNull(dbOperations, "dbOperations must not be null!");
-            return new PartTreeCosmosQuery(queryMethod, dbOperations);
+
+            if (queryMethod.hasAnnotatedQuery()) {
+                return new StringBasedCosmosQuery(queryMethod, dbOperations);
+            } else {
+                return new PartTreeCosmosQuery(queryMethod, dbOperations);
+            }
 
         }
     }

diff --git a/...java/com/azure/spring/data/cosmos/repository/support/ReactiveCosmosRepositoryFactory.java b/...java/com/azure/spring/data/cosmos/repository/support/ReactiveCosmosRepositoryFactory.java
@@ -77,8 +77,11 @@ public RepositoryQuery resolveQuery(Method method, RepositoryMetadata metadata,
 
             Assert.notNull(queryMethod, "queryMethod must not be null!");
             Assert.notNull(cosmosOperations, "dbOperations must not be null!");
-            return new PartTreeReactiveCosmosQuery(queryMethod, cosmosOperations);
-
+            if (queryMethod.hasAnnotatedQuery()) {
+                return new StringBasedReactiveCosmosQuery(queryMethod, cosmosOperations);
+            } else {
+                return new PartTreeReactiveCosmosQuery(queryMethod, cosmosOperations);
+            }
         }
     }