Exposing MaxBufferedItemCount to the Spring SDK. (Azure#30921)

* Exposing MaxBufferedItemCount to the Spring SDK.

* Updating the changelog.

* Update application.properties

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

Co-authored-by: Kushagra Thapar <[email protected]>

* Update README.md

* Update README.md

* Update CosmosConfig.java

* Update application.properties

Co-authored-by: Kushagra Thapar <[email protected]>

Author: trande4884

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

@@ -1115,6 +1115,22 @@ public void queryWithMaxDegreeOfParallelism() throws ClassNotFoundException {
         assertEquals((int) ReflectionTestUtils.getField(maxDegreeOfParallelismCosmosTemplate, "maxDegreeOfParallelism"), 20);
     }
 
+    @Test
+    public void queryWithMaxBufferedItemCount() throws ClassNotFoundException {
+        final CosmosConfig config = CosmosConfig.builder()
+            .maxBufferedItemCount(500)
+            .build();
+        final CosmosTemplate maxBufferedItemCountCosmosTemplate = 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 = maxBufferedItemCountCosmosTemplate.count(query, containerName);
+
+        assertEquals((int) ReflectionTestUtils.getField(maxBufferedItemCountCosmosTemplate, "maxBufferedItemCount"), 500);
+    }
+
     @Test
     public void queryDatabaseWithQueryMerticsEnabled() throws ClassNotFoundException {
         final CosmosConfig config = CosmosConfig.builder()

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

@@ -595,6 +595,26 @@ public void queryWithMaxDegreeOfParallelism() throws ClassNotFoundException {
         assertEquals((int) ReflectionTestUtils.getField(maxDegreeOfParallelismCosmosTemplate, "maxDegreeOfParallelism"), 20);
     }
 
+    @Test
+    public void queryWithMaxBufferedItemCount() throws ClassNotFoundException {
+        final CosmosConfig config = CosmosConfig.builder()
+            .maxBufferedItemCount(500)
+            .build();
+        final ReactiveCosmosTemplate maxBufferedItemCountCosmosTemplate = 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 = maxBufferedItemCountCosmosTemplate.runQuery(sqlQuerySpec, AuditableEntity.class, AuditableEntity.class);
+
+        StepVerifier.create(flux).expectNextCount(1).verifyComplete();
+        assertEquals((int) ReflectionTestUtils.getField(maxBufferedItemCountCosmosTemplate, "maxBufferedItemCount"), 500);
+    }
+
     @Test
     public void queryWithQueryMerticsEnabled() throws ClassNotFoundException {
         final CosmosConfig config = CosmosConfig.builder()

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

@@ -37,6 +37,9 @@ public class SecondaryTestRepositoryConfig {
     @Value("${cosmos.secondary.maxDegreeOfParallelism}")
     private int maxDegreeOfParallelism;
 
+    @Value("${cosmos.secondary.maxBufferedItemCount}")
+    private int maxBufferedItemCount;
+
     @Bean
     public CosmosClientBuilder secondaryCosmosClientBuilder() {
         return new CosmosClientBuilder()
@@ -61,6 +64,7 @@ public ReactiveCosmosTemplate secondaryReactiveCosmosTemplate(@Qualifier("second
             CosmosConfig config =  CosmosConfig.builder()
                 .enableQueryMetrics(queryMetricsEnabled)
                 .maxDegreeOfParallelism(maxDegreeOfParallelism)
+                .maxBufferedItemCount(maxBufferedItemCount)
                 .build();
 
             return new ReactiveCosmosTemplate(new CosmosFactory(client, getFirstDatabase()), config, mappingCosmosConverter);
@@ -78,6 +82,7 @@ public ReactiveCosmosTemplate secondaryReactiveCosmosTemplate1(@Qualifier("secon
             CosmosConfig config =  CosmosConfig.builder()
                 .enableQueryMetrics(queryMetricsEnabled)
                 .maxDegreeOfParallelism(maxDegreeOfParallelism)
+                .maxBufferedItemCount(maxBufferedItemCount)
                 .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

@@ -41,6 +41,9 @@ public class TestRepositoryConfig extends AbstractCosmosConfiguration {
     @Value("${cosmos.maxDegreeOfParallelism}")
     private int maxDegreeOfParallelism;
 
+    @Value("${cosmos.maxBufferedItemCount}")
+    private int maxBufferedItemCount;
+
     @Bean
     public ResponseDiagnosticsTestUtils responseDiagnosticsTestUtils() {
         return new ResponseDiagnosticsTestUtils();
@@ -60,6 +63,7 @@ public CosmosConfig cosmosConfig() {
         return CosmosConfig.builder()
                            .enableQueryMetrics(queryMetricsEnabled)
                            .maxDegreeOfParallelism(maxDegreeOfParallelism)
+                           .maxBufferedItemCount(maxBufferedItemCount)
                            .responseDiagnosticsProcessor(responseDiagnosticsTestUtils().getResponseDiagnosticsProcessor())
                            .build();
     }

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

@@ -7,6 +7,8 @@ dynamic.collection.name=spel-property-collection
 cosmos.queryMetricsEnabled=true
 # Max Degree of Parallelism allowed
 cosmos.maxDegreeOfParallelism=0
+# Max number of items to buffer
+cosmos.maxBufferedItemCount=0
 
 # Secondary DataSource Config
 cosmos.secondary.uri=${NEW_ACCOUNT_HOST}
@@ -17,5 +19,7 @@ cosmos.secondary.secondaryKey=${NEW_SECONDARY_ACCOUNT_KEY}
 cosmos.secondary.queryMetricsEnabled=true
 # Max Degree of Parallelism allowed
 cosmos.secondary.maxDegreeOfParallelism=0
+# Max number of items to buffer
+cosmos.secondary.maxBufferedItemCount=0
 
 

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

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

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

@@ -121,6 +121,8 @@ SLF4J is only needed if you plan to use logging, please also download an SLF4J b
 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.
+Set `maxBufferedItemCount` flag to an integer in application.properties to allow the user to set the max number of items that can be buffered during parallel query execution; if set to less than 0, the system automatically decides the number of items to buffer.
+NOTE: Setting this to a very high value can result in high memory consumption.
 
 ```java readme-sample-AppConfiguration
 @Configuration
@@ -146,6 +148,9 @@ public class AppConfiguration extends AbstractCosmosConfiguration {
 
     @Value("${azure.cosmos.maxDegreeOfParallelism}")
     private int maxDegreeOfParallelism;
+    
+    @Value("${azure.cosmos.maxBufferedItemCount}")
+    private int maxBufferedItemCount;
 
     private AzureKeyCredential azureKeyCredential;
 
@@ -165,6 +170,7 @@ public class AppConfiguration extends AbstractCosmosConfiguration {
         return CosmosConfig.builder()
                            .enableQueryMetrics(queryMetricsEnabled)
                            .maxDegreeOfParallelism(maxDegreeOfParallelism)
+                           .maxBufferedItemCount(maxBufferedItemCount)
                            .responseDiagnosticsProcessor(new ResponseDiagnosticsProcessorImplementation())
                            .build();
     }
@@ -207,6 +213,7 @@ public CosmosConfig cosmosConfig() {
     return CosmosConfig.builder()
                        .enableQueryMetrics(queryMetricsEnabled)
                        .maxDegreeOfParallelism(maxDegreeOfParallelism)
+                       .maxBufferedItemCount(maxBufferedItemCount)
                        .responseDiagnosticsProcessor(new ResponseDiagnosticsProcessorImplementation())
                        .build();
 }
@@ -684,6 +691,7 @@ public class SecondaryDatasourceConfiguration {
         return CosmosConfig.builder()
             .enableQueryMetrics(true)
             .maxDegreeOfParallelism(0)
+            .maxBufferedItemCount(0)
             .responseDiagnosticsProcessor(new ResponseDiagnosticsProcessorImplementation())
             .build();
     }
@@ -720,19 +728,21 @@ public CosmosConfig getCosmosConfig() {
     return CosmosConfig.builder()
         .enableQueryMetrics(true)
         .maxDegreeOfParallelism(0)
+        .maxBufferedItemCount(0)
         .responseDiagnosticsProcessor(new ResponseDiagnosticsProcessorImplementation())
         .build();
 }
 ```
 
-- Besides, if you want to define `queryMetricsEnabled`, `ResponseDiagnosticsProcessor` or `maxDegreeOfParallelism` , you can create the `CosmosConfig` for your cosmos template.
+- Besides, if you want to define `queryMetricsEnabled`, `ResponseDiagnosticsProcessor`, `maxDegreeOfParallelism` or `maxBufferedItemCount` , you can create the `CosmosConfig` for your cosmos template.
 
 ```java
 @Bean("secondaryCosmosConfig")
 public CosmosConfig getCosmosConfig() {
     return CosmosConfig.builder()
         .enableQueryMetrics(true)
         .maxDegreeOfParallelism(0)
+        .maxBufferedItemCount(0)
         .responseDiagnosticsProcessor(new ResponseDiagnosticsProcessorImplementation())
         .build();
 }

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

@@ -19,6 +19,8 @@ public class CosmosConfig {
 
     private final int maxDegreeOfParallelism;
 
+    private final int maxBufferedItemCount;
+
     /**
      * Initialization
      *
@@ -46,6 +48,7 @@ public CosmosConfig(ResponseDiagnosticsProcessor responseDiagnosticsProcessor,
         this.databaseThroughputConfig = databaseThroughputConfig;
         this.queryMetricsEnabled = queryMetricsEnabled;
         this.maxDegreeOfParallelism = 0;
+        this.maxBufferedItemCount = 0;
     }
 
     /**
@@ -65,6 +68,29 @@ public CosmosConfig(ResponseDiagnosticsProcessor responseDiagnosticsProcessor,
         this.databaseThroughputConfig = databaseThroughputConfig;
         this.queryMetricsEnabled = queryMetricsEnabled;
         this.maxDegreeOfParallelism = maxDegreeOfParallelism;
+        this.maxBufferedItemCount = 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}
+     * @param maxBufferedItemCount must not be {@literal null}
+     */
+    @ConstructorProperties({"responseDiagnosticsProcessor", "databaseThroughputConfig", "queryMetricsEnabled", "maxDegreeOfParallelism", "maxBufferedItemCount"})
+    CosmosConfig(ResponseDiagnosticsProcessor responseDiagnosticsProcessor,
+                 DatabaseThroughputConfig databaseThroughputConfig,
+                 boolean queryMetricsEnabled,
+                 int maxDegreeOfParallelism,
+                 int maxBufferedItemCount) {
+        this.responseDiagnosticsProcessor = responseDiagnosticsProcessor;
+        this.databaseThroughputConfig = databaseThroughputConfig;
+        this.queryMetricsEnabled = queryMetricsEnabled;
+        this.maxDegreeOfParallelism = maxDegreeOfParallelism;
+        this.maxBufferedItemCount = maxBufferedItemCount;
     }
 
     /**
@@ -94,6 +120,15 @@ public int getMaxDegreeOfParallelism() {
         return maxDegreeOfParallelism;
     }
 
+    /**
+     * Gets the value of maxBufferedItemCount
+     *
+     * @return int, value of maxBufferedItemCount
+     */
+    public int getMaxBufferedItemCount() {
+        return maxBufferedItemCount;
+    }
+
     /**
      * Gets the database throughput configuration.
      *
@@ -120,6 +155,7 @@ public static class CosmosConfigBuilder {
         private DatabaseThroughputConfig databaseThroughputConfig;
         private boolean queryMetricsEnabled;
         private int maxDegreeOfParallelism;
+        private int maxBufferedItemCount;
         CosmosConfigBuilder() {
         }
 
@@ -158,6 +194,17 @@ public CosmosConfigBuilder maxDegreeOfParallelism(int maxDegreeOfParallelism) {
             return this;
         }
 
+        /**
+         * Set maxBufferedItemCount
+         *
+         * @param maxBufferedItemCount value to initialize
+         * @return CosmosConfigBuilder
+         */
+        public CosmosConfigBuilder maxBufferedItemCount(int maxBufferedItemCount) {
+            this.maxBufferedItemCount = maxBufferedItemCount;
+            return this;
+        }
+
         /**
          * Enable database throughput
          *
@@ -176,7 +223,8 @@ public CosmosConfigBuilder enableDatabaseThroughput(boolean autoscale, int reque
          * @return CosmosConfig
          */
         public CosmosConfig build() {
-            return new CosmosConfig(this.responseDiagnosticsProcessor, this.databaseThroughputConfig, this.queryMetricsEnabled, this.maxDegreeOfParallelism);
+            return new CosmosConfig(this.responseDiagnosticsProcessor, this.databaseThroughputConfig, this.queryMetricsEnabled,
+                this.maxDegreeOfParallelism, this.maxBufferedItemCount);
         }
 
         @Override
@@ -186,6 +234,7 @@ public String toString() {
                 + ", databaseThroughputConfig=" + databaseThroughputConfig
                 + ", queryMetricsEnabled=" + queryMetricsEnabled
                 + ", maxDegreeOfParallelism=" + maxDegreeOfParallelism
+                + ", maxBufferedItemCount=" + maxBufferedItemCount
                 + '}';
         }
     }

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

@@ -78,6 +78,7 @@ public class CosmosTemplate implements CosmosOperations, ApplicationContextAware
     private final ResponseDiagnosticsProcessor responseDiagnosticsProcessor;
     private final boolean queryMetricsEnabled;
     private final int maxDegreeOfParallelism;
+    private final int maxBufferedItemCount;
     private final CosmosAsyncClient cosmosAsyncClient;
     private final DatabaseThroughputConfig databaseThroughputConfig;
 
@@ -132,6 +133,7 @@ public CosmosTemplate(CosmosFactory cosmosFactory,
         this.responseDiagnosticsProcessor = cosmosConfig.getResponseDiagnosticsProcessor();
         this.queryMetricsEnabled = cosmosConfig.isQueryMetricsEnabled();
         this.maxDegreeOfParallelism = cosmosConfig.getMaxDegreeOfParallelism();
+        this.maxBufferedItemCount = cosmosConfig.getMaxBufferedItemCount();
         this.databaseThroughputConfig = cosmosConfig.getDatabaseThroughputConfig();
     }
 
@@ -289,6 +291,7 @@ public <T> T findById(String containerName, Object id, Class<T> domainType) {
         final CosmosQueryRequestOptions options = new CosmosQueryRequestOptions();
         options.setQueryMetricsEnabled(this.queryMetricsEnabled);
         options.setMaxDegreeOfParallelism(this.maxDegreeOfParallelism);
+        options.setMaxBufferedItemCount(this.maxBufferedItemCount);
         return cosmosAsyncClient
             .getDatabase(this.databaseName)
             .getContainer(containerName)
@@ -414,6 +417,7 @@ public <T> Iterable<T> findAll(PartitionKey partitionKey, final Class<T> domainT
         cosmosQueryRequestOptions.setPartitionKey(partitionKey);
         cosmosQueryRequestOptions.setQueryMetricsEnabled(this.queryMetricsEnabled);
         cosmosQueryRequestOptions.setMaxDegreeOfParallelism(this.maxDegreeOfParallelism);
+        cosmosQueryRequestOptions.setMaxBufferedItemCount(this.maxBufferedItemCount);
 
         return cosmosAsyncClient
             .getDatabase(this.databaseName)
@@ -740,6 +744,7 @@ private <T> Slice<T> sliceQuery(SqlQuerySpec querySpec,
         final CosmosQueryRequestOptions cosmosQueryRequestOptions = new CosmosQueryRequestOptions();
         cosmosQueryRequestOptions.setQueryMetricsEnabled(this.queryMetricsEnabled);
         cosmosQueryRequestOptions.setMaxDegreeOfParallelism(this.maxDegreeOfParallelism);
+        cosmosQueryRequestOptions.setMaxBufferedItemCount(this.maxBufferedItemCount);
         partitionKeyValue.ifPresent(o -> {
             LOGGER.debug("Setting partition key {}", o);
             cosmosQueryRequestOptions.setPartitionKey(new PartitionKey(o));
@@ -886,6 +891,7 @@ private Long getCountValue(SqlQuerySpec querySpec, String containerName) {
         final CosmosQueryRequestOptions options = new CosmosQueryRequestOptions();
         options.setQueryMetricsEnabled(this.queryMetricsEnabled);
         options.setMaxDegreeOfParallelism(this.maxDegreeOfParallelism);
+        options.setMaxBufferedItemCount(this.maxBufferedItemCount);
 
         return executeQuery(querySpec, containerName, options)
             .publishOn(Schedulers.parallel())
@@ -915,6 +921,7 @@ private <T> Flux<JsonNode> findItemsAsFlux(@NonNull CosmosQuery query,
         final CosmosQueryRequestOptions cosmosQueryRequestOptions = new CosmosQueryRequestOptions();
         cosmosQueryRequestOptions.setQueryMetricsEnabled(this.queryMetricsEnabled);
         cosmosQueryRequestOptions.setMaxDegreeOfParallelism(this.maxDegreeOfParallelism);
+        cosmosQueryRequestOptions.setMaxBufferedItemCount(this.maxBufferedItemCount);
         Optional<Object> partitionKeyValue = query.getPartitionKeyValue(domainType);
         partitionKeyValue.ifPresent(o -> {
             LOGGER.debug("Setting partition key {}", o);
@@ -943,6 +950,7 @@ private Flux<JsonNode> getJsonNodeFluxFromQuerySpec(
         final CosmosQueryRequestOptions cosmosQueryRequestOptions = new CosmosQueryRequestOptions();
         cosmosQueryRequestOptions.setQueryMetricsEnabled(this.queryMetricsEnabled);
         cosmosQueryRequestOptions.setMaxDegreeOfParallelism(this.maxDegreeOfParallelism);
+        cosmosQueryRequestOptions.setMaxBufferedItemCount(this.maxBufferedItemCount);
 
         return cosmosAsyncClient
                    .getDatabase(this.databaseName)