Adding the maxDegreeOfParallelism property to the Spring Code (Azure#28756)

* Initial changes for adding the maxDegreeOfParallelism property to the spring code base.

* Initial changes for adding the maxDegreeOfParallelism property to the spring code base.

* Fixing param to default to 0 and fixing application.properties files.

* Fixing more maxDegreeOfParallelism properties.

* Fixing unit test that is broken.

* Adding some integration tests.

* Changing unit tests.

* Fixing read me.

* Removing unused import.

* Fixing constructors.

* Removing unneeded code.

* Fixing unit tests to better test the functionality.

* Fixing comments:
* Changing setMaxDegreeOfParallelism to maxDegreeOfParallelism
* Removing un-needed constructor

* Fixing comments setMaxDegreeOfParallelism in spring-test directory.

* Updating the CHANGELOG.md

* Updating the CHANGELOG.md to link PR

* Updating the CHANGELOG.md to link PR

Author: trande4884

sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/CosmosTemplateIT.java

@@ -51,6 +51,8 @@
 import org.springframework.data.repository.query.parser.Part;
 import org.springframework.test.context.ContextConfiguration;
 import org.springframework.test.context.junit4.SpringJUnit4ClassRunner;
+import org.springframework.test.util.ReflectionTestUtils;
+import reactor.core.publisher.Flux;
 
 import java.lang.reflect.InvocationTargetException;
 import java.lang.reflect.Method;
@@ -708,6 +710,38 @@ public void createDatabaseWithThroughput() throws ClassNotFoundException {
         assertEquals(expectedRequestUnits, response.getProperties().getManualThroughput());
     }
 
+    @Test
+    public void queryWithMaxDegreeOfParallelism() throws ClassNotFoundException {
+        final CosmosConfig config = CosmosConfig.builder()
+            .maxDegreeOfParallelism(20)
+            .build();
+        final CosmosTemplate maxDegreeOfParallelismCosmosTemplate = createCosmosTemplate(config, TestConstants.DB_NAME);
+
+        final Criteria criteria = Criteria.getInstance(CriteriaType.IS_EQUAL, "firstName",
+            Collections.singletonList(TEST_PERSON.getFirstName()), Part.IgnoreCaseType.NEVER);
+        final CosmosQuery query = new CosmosQuery(criteria);
+
+        final long count = maxDegreeOfParallelismCosmosTemplate.count(query, containerName);
+
+        assertEquals((int) ReflectionTestUtils.getField(maxDegreeOfParallelismCosmosTemplate, "maxDegreeOfParallelism"), 20);
+    }
+
+    @Test
+    public void queryDatabaseWithQueryMerticsEnabled() throws ClassNotFoundException {
+        final CosmosConfig config = CosmosConfig.builder()
+            .enableQueryMetrics(true)
+            .build();
+        final CosmosTemplate queryMetricsEnabledCosmosTemplate = createCosmosTemplate(config, TestConstants.DB_NAME);
+
+        final Criteria criteria = Criteria.getInstance(CriteriaType.IS_EQUAL, "firstName",
+            Collections.singletonList(TEST_PERSON.getFirstName()), Part.IgnoreCaseType.NEVER);
+        final CosmosQuery query = new CosmosQuery(criteria);
+
+        final long count = queryMetricsEnabledCosmosTemplate.count(query, containerName);
+
+        assertEquals((boolean) ReflectionTestUtils.getField(queryMetricsEnabledCosmosTemplate, "queryMetricsEnabled"), true);
+    }
+
     @Test
     public void userAgentSpringDataCosmosSuffix() throws NoSuchMethodException, InvocationTargetException, IllegalAccessException {
         //  getUserAgentSuffix method from CosmosClientBuilder

sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/ReactiveCosmosTemplateIT.java

@@ -48,6 +48,7 @@
 import org.springframework.data.repository.query.parser.Part;
 import org.springframework.test.context.ContextConfiguration;
 import org.springframework.test.context.junit4.SpringJUnit4ClassRunner;
+import org.springframework.test.util.ReflectionTestUtils;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 import reactor.test.StepVerifier;
@@ -574,6 +575,46 @@ public void createDatabaseWithThroughput() throws ClassNotFoundException {
         assertEquals(expectedRequestUnits, response.getProperties().getManualThroughput());
     }
 
+    @Test
+    public void queryWithMaxDegreeOfParallelism() throws ClassNotFoundException {
+        final CosmosConfig config = CosmosConfig.builder()
+            .maxDegreeOfParallelism(20)
+            .build();
+        final ReactiveCosmosTemplate maxDegreeOfParallelismCosmosTemplate = createReactiveCosmosTemplate(config, TestConstants.DB_NAME);
+
+        final AuditableEntity entity = new AuditableEntity();
+        entity.setId(UUID.randomUUID().toString());
+
+        auditableRepository.save(entity);
+
+        Criteria equals = Criteria.getInstance(CriteriaType.IS_EQUAL, "id", Collections.singletonList(entity.getId()), Part.IgnoreCaseType.NEVER);
+        final SqlQuerySpec sqlQuerySpec = new FindQuerySpecGenerator().generateCosmos(new CosmosQuery(equals));
+        final Flux<AuditableEntity> flux = maxDegreeOfParallelismCosmosTemplate.runQuery(sqlQuerySpec, AuditableEntity.class, AuditableEntity.class);
+
+        StepVerifier.create(flux).expectNextCount(1).verifyComplete();
+        assertEquals((int) ReflectionTestUtils.getField(maxDegreeOfParallelismCosmosTemplate, "maxDegreeOfParallelism"), 20);
+    }
+
+    @Test
+    public void queryWithQueryMerticsEnabled() throws ClassNotFoundException {
+        final CosmosConfig config = CosmosConfig.builder()
+            .enableQueryMetrics(true)
+            .build();
+        final ReactiveCosmosTemplate queryMetricsEnabledCosmosTemplate = createReactiveCosmosTemplate(config, TestConstants.DB_NAME);
+
+        final AuditableEntity entity = new AuditableEntity();
+        entity.setId(UUID.randomUUID().toString());
+
+        auditableRepository.save(entity);
+
+        Criteria equals = Criteria.getInstance(CriteriaType.IS_EQUAL, "id", Collections.singletonList(entity.getId()), Part.IgnoreCaseType.NEVER);
+        final SqlQuerySpec sqlQuerySpec = new FindQuerySpecGenerator().generateCosmos(new CosmosQuery(equals));
+        final Flux<AuditableEntity> flux = queryMetricsEnabledCosmosTemplate.runQuery(sqlQuerySpec, AuditableEntity.class, AuditableEntity.class);
+
+        StepVerifier.create(flux).expectNextCount(1).verifyComplete();
+        assertEquals((boolean) ReflectionTestUtils.getField(queryMetricsEnabledCosmosTemplate, "queryMetricsEnabled"), true);
+    }
+
     @Test
     public void userAgentSpringDataCosmosSuffix() throws NoSuchMethodException, InvocationTargetException, IllegalAccessException {
         //  getUserAgentSuffix method from CosmosClientBuilder

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

@@ -34,6 +34,9 @@ public class SecondaryTestRepositoryConfig {
     @Value("${cosmos.secondary.queryMetricsEnabled}")
     private boolean queryMetricsEnabled;
 
+    @Value("${cosmos.secondary.maxDegreeOfParallelism}")
+    private int maxDegreeOfParallelism;
+
     @Bean
     public CosmosClientBuilder secondaryCosmosClientBuilder() {
         return new CosmosClientBuilder()
@@ -57,6 +60,7 @@ public ReactiveCosmosTemplate secondaryReactiveCosmosTemplate(@Qualifier("second
 
             CosmosConfig config =  CosmosConfig.builder()
                 .enableQueryMetrics(queryMetricsEnabled)
+                .maxDegreeOfParallelism(maxDegreeOfParallelism)
                 .build();
 
             return new ReactiveCosmosTemplate(new CosmosFactory(client, getFirstDatabase()), config, mappingCosmosConverter);
@@ -73,6 +77,7 @@ public ReactiveCosmosTemplate secondaryReactiveCosmosTemplate1(@Qualifier("secon
 
             CosmosConfig config =  CosmosConfig.builder()
                 .enableQueryMetrics(queryMetricsEnabled)
+                .maxDegreeOfParallelism(maxDegreeOfParallelism)
                 .build();
 
             return new ReactiveCosmosTemplate(new CosmosFactory(client, getSecondDatabase()), config, mappingCosmosConverter);

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

@@ -38,6 +38,9 @@ public class TestRepositoryConfig extends AbstractCosmosConfiguration {
     @Value("${cosmos.queryMetricsEnabled}")
     private boolean queryMetricsEnabled;
 
+    @Value("${cosmos.maxDegreeOfParallelism}")
+    private int maxDegreeOfParallelism;
+
     @Bean
     public ResponseDiagnosticsTestUtils responseDiagnosticsTestUtils() {
         return new ResponseDiagnosticsTestUtils();
@@ -56,6 +59,7 @@ public CosmosClientBuilder cosmosClientBuilder() {
     public CosmosConfig cosmosConfig() {
         return CosmosConfig.builder()
                            .enableQueryMetrics(queryMetricsEnabled)
+                           .maxDegreeOfParallelism(maxDegreeOfParallelism)
                            .responseDiagnosticsProcessor(responseDiagnosticsTestUtils().getResponseDiagnosticsProcessor())
                            .build();
     }

sdk/cosmos/azure-spring-data-cosmos-test/src/test/resources/application.properties

@@ -5,6 +5,8 @@ cosmos.secondaryKey=${SECONDARY_ACCOUNT_KEY}
 dynamic.collection.name=spel-property-collection
 # Populate query metrics
 cosmos.queryMetricsEnabled=true
+# Max Degree of Parallelism allowed
+cosmos.maxDegreeOfParallelism=0
 
 # Secondary DataSource Config
 cosmos.secondary.uri=${NEW_ACCOUNT_HOST}
@@ -13,6 +15,7 @@ cosmos.secondary.secondaryKey=${NEW_SECONDARY_ACCOUNT_KEY}
 
 # Populate query metrics
 cosmos.secondary.queryMetricsEnabled=true
-
+# Max Degree of Parallelism allowed
+cosmos.secondary.maxDegreeOfParallelism=0
 
 

sdk/cosmos/azure-spring-data-cosmos/CHANGELOG.md

@@ -3,6 +3,7 @@
 ### 3.22.0-beta.1 (Unreleased)
 
 #### Features Added
+* Exposed `maxDegreeOfParallelism` feature from CosmosQueryRequestOptions through application.properties flag - See [PR 28756](https://github.com/Azure/azure-sdk-for-java/pull/28756)
 
 #### Breaking Changes
 

sdk/cosmos/azure-spring-data-cosmos/README.md

@@ -86,9 +86,10 @@ SLF4J is only needed if you plan to use logging, please also download an SLF4J b
 ### Setup Configuration Class
 - In order to set up configuration class, you'll need to extend `AbstractCosmosConfiguration`
 
-- Azure-spring-data-cosmos also supports `Response Diagnostics String` and `Query Metrics`.
+- Azure-spring-data-cosmos also supports `Response Diagnostics String`, `Query Metrics` and `Max Degree of Parallelism`.
 Set `queryMetricsEnabled` flag to true in application.properties to enable query metrics.
 In addition to setting the flag, implement `ResponseDiagnosticsProcessor` to log diagnostics information.
+Set `maxDegreeOfParallelism` flag to an integer in application.properties to allow parallel processing; setting the value to -1 will lead to the SDK deciding the optimal value.
 
 ```java readme-sample-AppConfiguration
 @Configuration
@@ -111,6 +112,9 @@ public class AppConfiguration extends AbstractCosmosConfiguration {
 
     @Value("${azure.cosmos.queryMetricsEnabled}")
     private boolean queryMetricsEnabled;
+    
+    @Value("${azure.cosmos.maxDegreeOfParallelism}")
+    private int maxDegreeOfParallelism;
 
     private AzureKeyCredential azureKeyCredential;
 
@@ -129,6 +133,7 @@ public class AppConfiguration extends AbstractCosmosConfiguration {
     public CosmosConfig cosmosConfig() {
         return CosmosConfig.builder()
                            .enableQueryMetrics(queryMetricsEnabled)
+                           .maxDegreeOfParallelism(maxDegreeOfParallelism)
                            .responseDiagnosticsProcessor(new ResponseDiagnosticsProcessorImplementation())
                            .build();
     }
@@ -170,6 +175,7 @@ public CosmosClientBuilder getCosmosClientBuilder() {
 public CosmosConfig cosmosConfig() {
     return CosmosConfig.builder()
                        .enableQueryMetrics(queryMetricsEnabled)
+                       .maxDegreeOfParallelism(maxDegreeOfParallelism)
                        .responseDiagnosticsProcessor(new ResponseDiagnosticsProcessorImplementation())
                        .build();
 }
@@ -646,6 +652,7 @@ public class SecondaryDatasourceConfiguration {
     public CosmosConfig getCosmosConfig() {
         return CosmosConfig.builder()
             .enableQueryMetrics(true)
+            .maxDegreeOfParallelism(0)
             .responseDiagnosticsProcessor(new ResponseDiagnosticsProcessorImplementation())
             .build();
     }
@@ -681,18 +688,20 @@ public CosmosAsyncClient getCosmosAsyncClient(@Qualifier("secondary") CosmosProp
 public CosmosConfig getCosmosConfig() {
     return CosmosConfig.builder()
         .enableQueryMetrics(true)
+        .maxDegreeOfParallelism(0)
         .responseDiagnosticsProcessor(new ResponseDiagnosticsProcessorImplementation())
         .build();
 }
 ```
 
-- Besides, if you want to define `queryMetricsEnabled` or `ResponseDiagnosticsProcessor` , you can create the `CosmosConfig` for your cosmos template.
+- Besides, if you want to define `queryMetricsEnabled`, `ResponseDiagnosticsProcessor` or `maxDegreeOfParallelism` , you can create the `CosmosConfig` for your cosmos template.
 
 ```java
 @Bean("secondaryCosmosConfig")
 public CosmosConfig getCosmosConfig() {
     return CosmosConfig.builder()
         .enableQueryMetrics(true)
+        .maxDegreeOfParallelism(0)
         .responseDiagnosticsProcessor(new ResponseDiagnosticsProcessorImplementation())
         .build();
 }

sdk/cosmos/azure-spring-data-cosmos/src/main/java/com/azure/spring/data/cosmos/config/CosmosConfig.java

@@ -17,6 +17,8 @@ public class CosmosConfig {
 
     private final boolean queryMetricsEnabled;
 
+    private final int maxDegreeOfParallelism;
+
     /**
      * Initialization
      *
@@ -33,7 +35,7 @@ public CosmosConfig(ResponseDiagnosticsProcessor responseDiagnosticsProcessor,
      * Initialization
      *
      * @param responseDiagnosticsProcessor must not be {@literal null}
-     * @param databaseThroughputConfig may be @{literal null}
+     * @param databaseThroughputConfig may be {@literal null}
      * @param queryMetricsEnabled must not be {@literal null}
      */
     @ConstructorProperties({"responseDiagnosticsProcessor", "databaseThroughputConfig", "queryMetricsEnabled"})
@@ -43,6 +45,26 @@ public CosmosConfig(ResponseDiagnosticsProcessor responseDiagnosticsProcessor,
         this.responseDiagnosticsProcessor = responseDiagnosticsProcessor;
         this.databaseThroughputConfig = databaseThroughputConfig;
         this.queryMetricsEnabled = queryMetricsEnabled;
+        this.maxDegreeOfParallelism = 0;
+    }
+
+    /**
+     * Initialization
+     *
+     * @param responseDiagnosticsProcessor must not be {@literal null}
+     * @param databaseThroughputConfig may be {@literal null}
+     * @param queryMetricsEnabled must not be {@literal null}
+     * @param maxDegreeOfParallelism must not be {@literal null}
+     */
+    @ConstructorProperties({"responseDiagnosticsProcessor", "databaseThroughputConfig", "queryMetricsEnabled", "maxDegreeOfParallelism"})
+    public CosmosConfig(ResponseDiagnosticsProcessor responseDiagnosticsProcessor,
+                        DatabaseThroughputConfig databaseThroughputConfig,
+                        boolean queryMetricsEnabled,
+                        int maxDegreeOfParallelism) {
+        this.responseDiagnosticsProcessor = responseDiagnosticsProcessor;
+        this.databaseThroughputConfig = databaseThroughputConfig;
+        this.queryMetricsEnabled = queryMetricsEnabled;
+        this.maxDegreeOfParallelism = maxDegreeOfParallelism;
     }
 
     /**
@@ -63,6 +85,15 @@ public boolean isQueryMetricsEnabled() {
         return queryMetricsEnabled;
     }
 
+    /**
+     * Gets the value of maxDegreeOfParallelism
+     *
+     * @return int, value of maxDegreeOfParallelism
+     */
+    public int getMaxDegreeOfParallelism() {
+        return maxDegreeOfParallelism;
+    }
+
     /**
      * Gets the database throughput configuration.
      *
@@ -88,6 +119,7 @@ public static class CosmosConfigBuilder {
         private ResponseDiagnosticsProcessor responseDiagnosticsProcessor;
         private DatabaseThroughputConfig databaseThroughputConfig;
         private boolean queryMetricsEnabled;
+        private int maxDegreeOfParallelism;
         CosmosConfigBuilder() {
         }
 
@@ -115,6 +147,17 @@ public CosmosConfigBuilder enableQueryMetrics(boolean queryMetricsEnabled) {
             return this;
         }
 
+        /**
+         * Set maxDegreeOfParallelism
+         *
+         * @param maxDegreeOfParallelism value to initialize
+         * @return CosmosConfigBuilder
+         */
+        public CosmosConfigBuilder maxDegreeOfParallelism(int maxDegreeOfParallelism) {
+            this.maxDegreeOfParallelism = maxDegreeOfParallelism;
+            return this;
+        }
+
         /**
          * Enable database throughput
          *
@@ -133,7 +176,7 @@ public CosmosConfigBuilder enableDatabaseThroughput(boolean autoscale, int reque
          * @return CosmosConfig
          */
         public CosmosConfig build() {
-            return new CosmosConfig(this.responseDiagnosticsProcessor, this.databaseThroughputConfig, this.queryMetricsEnabled);
+            return new CosmosConfig(this.responseDiagnosticsProcessor, this.databaseThroughputConfig, this.queryMetricsEnabled, this.maxDegreeOfParallelism);
         }
 
         @Override
@@ -142,6 +185,7 @@ public String toString() {
                 + "responseDiagnosticsProcessor=" + responseDiagnosticsProcessor
                 + ", databaseThroughputConfig=" + databaseThroughputConfig
                 + ", queryMetricsEnabled=" + queryMetricsEnabled
+                + ", maxDegreeOfParallelism=" + maxDegreeOfParallelism
                 + '}';
         }
     }