Added partition key support to spring data single partition queries (Azure#22483)

* Added partition key support to spring data single partition queries

* Spot bug fixes

* Check style fixes

* Updated logic for IN clause and ignoreCase

* Added logging for partition key in spring repository queries

Author: kushagraThapar

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

@@ -96,6 +96,27 @@ public void testFindByPartitionedCity() {
         assertThat(result.get(1).getCity()).isEqualTo(city);
     }
 
+    @Test
+    public void testFindByPartitionedCityIn() {
+        final String city = TEST_ADDRESS1_PARTITION1.getCity();
+        final List<Address> result = TestUtils.toList(repository.findByCityIn(Lists.newArrayList(city)));
+
+        assertThat(result.size()).isEqualTo(2);
+        assertThat(result.get(0).getCity()).isEqualTo(city);
+        assertThat(result.get(1).getCity()).isEqualTo(city);
+    }
+
+    @Test
+    public void testFindByPostalCodeAndCityIn() {
+        final String city = TEST_ADDRESS1_PARTITION1.getCity();
+        final List<String> postalCodes = Lists.newArrayList(TEST_ADDRESS1_PARTITION1.getPostalCode(),
+            TEST_ADDRESS2_PARTITION1.getPostalCode());
+        final List<Address> result = TestUtils.toList(repository.findByPostalCodeInAndCity(postalCodes, city));
+
+        assertThat(result.size()).isEqualTo(2);
+        assertThat(result).isEqualTo(Lists.newArrayList(TEST_ADDRESS1_PARTITION1, TEST_ADDRESS2_PARTITION1));
+    }
+
     @Test
     public void testFindByStreetOrCity() {
         final String city = TEST_ADDRESS1_PARTITION1.getCity();

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

@@ -389,6 +389,10 @@ public void testFindByIn() {
         projects = TestUtils.toList(repository.findByCreatorIn(Arrays.asList(CREATOR_0, FAKE_CREATOR)));
 
         assertProjectListEquals(projects, Arrays.asList(PROJECT_0, PROJECT_4));
+
+        projects = TestUtils.toList(repository.findByCreatorIn(Collections.singletonList(CREATOR_1)));
+
+        assertProjectListEquals(projects, Collections.singletonList(PROJECT_1));
     }
 
     @Test
@@ -411,6 +415,24 @@ public void testFindByInWithAnd() {
         assertProjectListEquals(projects, Arrays.asList(PROJECT_0, PROJECT_1, PROJECT_2, PROJECT_4));
     }
 
+    @Test
+    public void testFindByInWithOr() {
+        List<Project> projects = TestUtils.toList(repository.findByCreatorInOrStarCount(Arrays.asList(CREATOR_0,
+            CREATOR_1), STAR_COUNT_2));
+
+        assertProjectListEquals(projects, Arrays.asList(PROJECT_0, PROJECT_4, PROJECT_1, PROJECT_2));
+
+        projects = TestUtils.toList(repository.findByCreatorInOrStarCount(Collections.singletonList(CREATOR_1),
+            STAR_COUNT_2));
+
+        assertProjectListEquals(projects, Arrays.asList(PROJECT_1, PROJECT_2));
+
+        projects = TestUtils.toList(repository.findByCreatorInOrStarCount(Collections.singletonList(CREATOR_0),
+            STAR_COUNT_0));
+
+        assertProjectListEquals(projects, Arrays.asList(PROJECT_0, PROJECT_4));
+    }
+
     @Test
     public void testFindByNotIn() {
         List<Project> projects = TestUtils.toList(repository.findByCreatorNotIn(

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

@@ -24,8 +24,12 @@ public interface AddressRepository extends CosmosRepository<Address, String> {
 
     Iterable<Address> findByCity(String city);
 
+    Iterable<Address> findByCityIn(List<String> cities);
+
     Iterable<Address> findByPostalCode(String postalCode);
 
+    Iterable<Address> findByPostalCodeInAndCity(List<String> postalCodes, String city);
+
     Iterable<Address> findByStreetOrCity(String street, String city);
 
     @Query("select * from a where a.city = @city")

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

@@ -71,6 +71,8 @@ Iterable<Project> findByNameOrCreatorAndForkCountOrStarCount(String name, String
 
     Iterable<Project> findByCreatorInAndStarCountIn(Collection<String> creators, Collection<Long> starCounts);
 
+    Iterable<Project> findByCreatorInOrStarCount(Collection<String> creators, Long starCount);
+
     Iterable<Project> findByCreatorNotIn(Collection<String> creators);
 
     Iterable<Project> findByCreatorInAndStarCountNotIn(Collection<String> creators, Collection<Long> starCounts);

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

@@ -55,6 +55,7 @@
 import java.util.Collections;
 import java.util.Iterator;
 import java.util.List;
+import java.util.Optional;
 import java.util.UUID;
 import java.util.stream.Collectors;
 
@@ -632,7 +633,8 @@ public <T> Iterable<T> delete(@NonNull CosmosQuery query, @NonNull Class<T> doma
         Assert.notNull(domainType, "domainType should not be null.");
         Assert.hasText(containerName, "container should not be null, empty or only whitespaces");
 
-        final List<JsonNode> results = findItems(query, containerName);
+        final List<JsonNode> results = findItemsAsFlux(query, containerName, domainType).collectList().block();
+        assert results != null;
         return results.stream()
                       .map(item -> deleteItem(item, containerName, domainType))
                       .collect(Collectors.toList());
@@ -654,39 +656,49 @@ public <T> Page<T> runPaginationQuery(SqlQuerySpec querySpec, Pageable pageable,
         final String containerName = getContainerName(domainType);
         final SqlQuerySpec sortedQuerySpec = NativeQueryGenerator.getInstance().generateSortedQuery(querySpec, pageable.getSort());
         final SqlQuerySpec countQuerySpec = NativeQueryGenerator.getInstance().generateCountQuery(querySpec);
-        return paginationQuery(sortedQuerySpec, countQuerySpec, pageable, pageable.getSort(), returnType, containerName);
+        return paginationQuery(sortedQuerySpec, countQuerySpec, pageable,
+            pageable.getSort(), returnType, containerName, Optional.empty());
     }
 
     @Override
     public <T> Page<T> paginationQuery(CosmosQuery query, Class<T> domainType, String containerName) {
         final SqlQuerySpec querySpec = new FindQuerySpecGenerator().generateCosmos(query);
         final SqlQuerySpec countQuerySpec = new CountQueryGenerator().generateCosmos(query);
-        return paginationQuery(querySpec, countQuerySpec, query.getPageable(), query.getSort(), domainType, containerName);
+        Optional<Object> partitionKeyValue = query.getPartitionKeyValue(domainType);
+        return paginationQuery(querySpec, countQuerySpec, query.getPageable(),
+            query.getSort(), domainType, containerName, partitionKeyValue);
     }
 
     @Override
     public <T> Slice<T> sliceQuery(CosmosQuery query, Class<T> domainType, String containerName) {
         final SqlQuerySpec querySpec = new FindQuerySpecGenerator().generateCosmos(query);
-        return sliceQuery(querySpec, query.getPageable(), query.getSort(), domainType, containerName);
+        Optional<Object> partitionKeyValue = query.getPartitionKeyValue(domainType);
+        return sliceQuery(querySpec, query.getPageable(), query.getSort(), domainType, containerName, partitionKeyValue);
     }
 
     private <T> Page<T> paginationQuery(SqlQuerySpec querySpec, SqlQuerySpec countQuerySpec,
                                        Pageable pageable, Sort sort,
-                                       Class<T> returnType, String containerName) {
-        Slice<T> response = sliceQuery(querySpec, pageable, sort, returnType, containerName);
+                                       Class<T> returnType, String containerName,
+                                        Optional<Object> partitionKeyValue) {
+        Slice<T> response = sliceQuery(querySpec, pageable, sort, returnType, containerName, partitionKeyValue);
         final long total = getCountValue(countQuerySpec, containerName);
         return new CosmosPageImpl<>(response.getContent(), response.getPageable(), total);
     }
 
     private <T> Slice<T> sliceQuery(SqlQuerySpec querySpec,
                                    Pageable pageable, Sort sort,
-                                   Class<T> returnType, String containerName) {
+                                   Class<T> returnType, String containerName,
+                                    Optional<Object> partitionKeyValue) {
         Assert.isTrue(pageable.getPageSize() > 0,
             "pageable should have page size larger than 0");
         Assert.hasText(containerName, "container should not be null, empty or only whitespaces");
 
         final CosmosQueryRequestOptions cosmosQueryRequestOptions = new CosmosQueryRequestOptions();
         cosmosQueryRequestOptions.setQueryMetricsEnabled(this.queryMetricsEnabled);
+        partitionKeyValue.ifPresent(o -> {
+            LOGGER.debug("Setting partition key {}", o);
+            cosmosQueryRequestOptions.setPartitionKey(new PartitionKey(o));
+        });
 
         CosmosAsyncContainer container =
             cosmosAsyncClient.getDatabase(this.databaseName).getContainer(containerName);
@@ -836,11 +848,17 @@ private Flux<FeedResponse<JsonNode>> executeQuery(SqlQuerySpec sqlQuerySpec,
                                 .byPage();
     }
 
-    private Flux<JsonNode> findItemsAsFlux(@NonNull CosmosQuery query,
-                                           @NonNull String containerName) {
+    private <T> Flux<JsonNode> findItemsAsFlux(@NonNull CosmosQuery query,
+                                               @NonNull String containerName,
+                                               @NonNull Class<T> domainType) {
         final SqlQuerySpec sqlQuerySpec = new FindQuerySpecGenerator().generateCosmos(query);
         final CosmosQueryRequestOptions cosmosQueryRequestOptions = new CosmosQueryRequestOptions();
         cosmosQueryRequestOptions.setQueryMetricsEnabled(this.queryMetricsEnabled);
+        Optional<Object> partitionKeyValue = query.getPartitionKeyValue(domainType);
+        partitionKeyValue.ifPresent(o -> {
+            LOGGER.debug("Setting partition key {}", o);
+            cosmosQueryRequestOptions.setPartitionKey(new PartitionKey(o));
+        });
 
         return cosmosAsyncClient
             .getDatabase(this.databaseName)
@@ -879,17 +897,10 @@ private Flux<JsonNode> getJsonNodeFluxFromQuerySpec(
                                       CosmosExceptionUtils.exceptionHandler("Failed to find items", throwable));
     }
 
-    private List<JsonNode> findItems(@NonNull CosmosQuery query,
-                                     @NonNull String containerName) {
-        return findItemsAsFlux(query, containerName)
-            .collectList()
-            .block();
-    }
-
     private <T> Iterable<T> findItems(@NonNull CosmosQuery query,
                                       @NonNull String containerName,
                                       @NonNull Class<T> domainType) {
-        return findItemsAsFlux(query, containerName)
+        return findItemsAsFlux(query, containerName, domainType)
             .map(jsonNode -> toDomainObject(domainType, jsonNode))
             .toIterable();
     }

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

@@ -28,6 +28,8 @@
 import com.azure.spring.data.cosmos.exception.CosmosExceptionUtils;
 import com.azure.spring.data.cosmos.repository.support.CosmosEntityInformation;
 import com.fasterxml.jackson.databind.JsonNode;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 import org.springframework.beans.BeansException;
 import org.springframework.context.ApplicationContext;
 import org.springframework.context.ApplicationContextAware;
@@ -40,6 +42,7 @@
 import reactor.core.publisher.Mono;
 import reactor.core.scheduler.Schedulers;
 
+import java.util.Optional;
 import java.util.UUID;
 
 /**
@@ -48,6 +51,8 @@
 @SuppressWarnings("unchecked")
 public class ReactiveCosmosTemplate implements ReactiveCosmosOperations, ApplicationContextAware {
 
+    private static final Logger LOGGER = LoggerFactory.getLogger(ReactiveCosmosTemplate.class);
+
     private final MappingCosmosConverter mappingCosmosConverter;
     private final String databaseName;
     private final ResponseDiagnosticsProcessor responseDiagnosticsProcessor;
@@ -529,7 +534,7 @@ public <T> Flux<T> delete(CosmosQuery query, Class<T> domainType, String contain
         Assert.notNull(domainType, "domainType should not be null.");
         Assert.hasText(containerName, "container name should not be null, empty or only whitespaces");
 
-        final Flux<JsonNode> results = findItems(query, containerName);
+        final Flux<JsonNode> results = findItems(query, containerName, domainType);
 
         return results.flatMap(d -> deleteItem(d, containerName, domainType));
     }
@@ -544,7 +549,7 @@ public <T> Flux<T> delete(CosmosQuery query, Class<T> domainType, String contain
      */
     @Override
     public <T> Flux<T> find(CosmosQuery query, Class<T> domainType, String containerName) {
-        return findItems(query, containerName)
+        return findItems(query, containerName, domainType)
             .map(cosmosItemProperties -> toDomainObject(domainType, cosmosItemProperties));
     }
 
@@ -710,11 +715,17 @@ private void markAuditedIfConfigured(Object object) {
         }
     }
 
-    private Flux<JsonNode> findItems(@NonNull CosmosQuery query,
-                                     @NonNull String containerName) {
+    private <T> Flux<JsonNode> findItems(@NonNull CosmosQuery query,
+                                     @NonNull String containerName,
+                                     @NonNull Class<T> domainType) {
         final SqlQuerySpec sqlQuerySpec = new FindQuerySpecGenerator().generateCosmos(query);
         final CosmosQueryRequestOptions cosmosQueryRequestOptions = new CosmosQueryRequestOptions();
         cosmosQueryRequestOptions.setQueryMetricsEnabled(this.queryMetricsEnabled);
+        Optional<Object> partitionKeyValue = query.getPartitionKeyValue(domainType);
+        partitionKeyValue.ifPresent(o -> {
+            LOGGER.debug("Setting partition key {}", o);
+            cosmosQueryRequestOptions.setPartitionKey(new PartitionKey(o));
+        });
 
         return cosmosAsyncClient
             .getDatabase(this.databaseName)

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

@@ -2,11 +2,15 @@
 // Licensed under the MIT License.
 package com.azure.spring.data.cosmos.core.query;
 
+import com.azure.spring.data.cosmos.repository.support.CosmosEntityInformation;
 import org.springframework.data.domain.Pageable;
 import org.springframework.data.domain.Sort;
+import org.springframework.data.repository.query.parser.Part;
 import org.springframework.lang.NonNull;
 import org.springframework.util.Assert;
 
+import java.util.Collection;
+import java.util.Collections;
 import java.util.List;
 import java.util.Optional;
 
@@ -111,14 +115,37 @@ private boolean isCrossPartitionQuery(@NonNull String keyName) {
 
         final Optional<Criteria> criteria = this.getSubjectCriteria(this.criteria, keyName);
 
-        return criteria.map(criteria1 -> criteria1.getType() != CriteriaType.IS_EQUAL).orElse(true);
+        return criteria.map(criteria1 -> {
+            //  If there is equal criteria, then it is a single partition query
+            if (isEqualCriteria(criteria1)) {
+                return false;
+            }
+            //  IN is a special case, where we want to first check if the partition key is used with IN clause
+            if (criteria1.getType() == CriteriaType.IN && criteria1.getSubjectValues().size() == 1) {
+                @SuppressWarnings("unchecked")
+                Collection<Object> collection = (Collection<Object>) criteria1.getSubjectValues().get(0);
+                //  IN query types can have multiple values,
+                //  so we are checking the internal collection of the criteria
+                return collection.size() != 1;
+            }
+            return !hasKeywordAnd();
+        }).orElse(true);
     }
 
     private boolean hasKeywordOr() {
         // If there is OR keyword in DocumentQuery, the top node of Criteria must be OR type.
         return this.criteria.getType() == CriteriaType.OR;
     }
 
+    private boolean hasKeywordAnd() {
+        // If there is AND keyword in DocumentQuery, the top node of Criteria must be AND type.
+        return this.criteria.getType() == CriteriaType.AND;
+    }
+
+    private boolean isEqualCriteria(Criteria criteria) {
+        return criteria.getType() == CriteriaType.IS_EQUAL;
+    }
+
     /**
      * Indicate if DocumentQuery should enable cross partition query.
      *
@@ -136,6 +163,57 @@ public boolean isCrossPartitionQuery(@NonNull List<String> partitionKeys) {
                             .orElse(hasKeywordOr());
     }
 
+    /**
+     * Returns true if this criteria or sub-criteria has partition key field present as one of the subjects.
+     * @param partitionKeyFieldName partition key field name
+     * @return returns true if this criteria or sub criteria has partition key field present as one of the subjects.
+     */
+    public boolean hasPartitionKeyCriteria(@NonNull String partitionKeyFieldName) {
+        if (partitionKeyFieldName.isEmpty()) {
+            return false;
+        }
+
+        final Optional<Criteria> criteria = this.getSubjectCriteria(this.criteria, partitionKeyFieldName);
+        return criteria.isPresent();
+    }
+
+    /**
+     * Returns partition key value based on the criteria.
+     * @param domainType domain type
+     * @param <T> entity class type
+     * @return Optional of partition key value
+     */
+    public <T> Optional<Object> getPartitionKeyValue(@NonNull Class<T> domainType) {
+        CosmosEntityInformation<?, ?> instance = CosmosEntityInformation.getInstance(domainType);
+        String partitionKeyFieldName = instance.getPartitionKeyFieldName();
+        if (partitionKeyFieldName == null
+            || partitionKeyFieldName.isEmpty()
+            || isCrossPartitionQuery(Collections.singletonList(partitionKeyFieldName))) {
+            return Optional.empty();
+        }
+
+        final Optional<Criteria> criteria = this.getSubjectCriteria(this.criteria, partitionKeyFieldName);
+        return criteria.map(criteria1 -> {
+            //  If the criteria has ignoreCase, then we cannot set the partition key
+            //  because of case sensitivity of partition key
+            if (!criteria1.getIgnoreCase().equals(Part.IgnoreCaseType.NEVER)) {
+                return null;
+            }
+            if (criteria1.getType() == CriteriaType.IN && criteria1.getSubjectValues().size() == 1) {
+                @SuppressWarnings("unchecked")
+                Collection<Object> collection = (Collection<Object>) criteria1.getSubjectValues().get(0);
+                //  IN query types can have multiple values,
+                //  so we are checking the internal collection of the criteria
+                if (collection.size() == 1) {
+                    return collection.iterator().next();
+                } else {
+                    return null;
+                }
+            }
+            return criteria1.getSubjectValues().get(0);
+        });
+    }
+
     /**
      * To get criteria by type
      *

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

@@ -252,6 +252,10 @@ public Object getPartitionKeyFieldValue(T entity) {
         return partitionKeyField == null ? null : ReflectionUtils.getField(partitionKeyField, entity);
     }
 
+    public String getPartitionKeyFieldName() {
+        return partitionKeyField == null ? null : partitionKeyField.getName();
+    }
+
     /**
      * Check if auto creating container is allowed
      *