Adding responseContinuationTokenLimitInKb support to the Spring SDK. (Azure#30980)

* Adding responseContinuationTokenLimitInKb support to the Spring SDK.

* Updating the changelog.

* Update README.md

Author: trande4884

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

@@ -940,6 +940,24 @@ public void queryWithMaxBufferedItemCount() throws ClassNotFoundException {
         assertEquals((int) ReflectionTestUtils.getField(maxBufferedItemCountCosmosTemplate, "maxBufferedItemCount"), 500);
     }
 
+    @Test
+    public void queryWithResponseContinuationTokenLimitInKb() throws ClassNotFoundException {
+        final CosmosConfig config = CosmosConfig.builder()
+            .responseContinuationTokenLimitInKb(2000)
+            .build();
+        final CosmosTemplate responseContinuationTokenLimitInKbCosmosTemplate =
+            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 = responseContinuationTokenLimitInKbCosmosTemplate.count(query, containerName);
+
+        assertEquals((int) ReflectionTestUtils.getField(responseContinuationTokenLimitInKbCosmosTemplate,
+            "responseContinuationTokenLimitInKb"), 2000);
+    }
+
     @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

@@ -615,6 +615,28 @@ public void queryWithMaxBufferedItemCount() throws ClassNotFoundException {
         assertEquals((int) ReflectionTestUtils.getField(maxBufferedItemCountCosmosTemplate, "maxBufferedItemCount"), 500);
     }
 
+    @Test
+    public void queryWithResponseContinuationTokenLimitInKb() throws ClassNotFoundException {
+        final CosmosConfig config = CosmosConfig.builder()
+            .responseContinuationTokenLimitInKb(2000)
+            .build();
+        final ReactiveCosmosTemplate responseContinuationTokenLimitInKbCosmosTemplate =
+            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 = responseContinuationTokenLimitInKbCosmosTemplate.runQuery(sqlQuerySpec, AuditableEntity.class, AuditableEntity.class);
+
+        StepVerifier.create(flux).expectNextCount(1).verifyComplete();
+        assertEquals((int) ReflectionTestUtils.getField(responseContinuationTokenLimitInKbCosmosTemplate,
+            "responseContinuationTokenLimitInKb"), 2000);
+    }
+
     @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

@@ -40,6 +40,9 @@ public class SecondaryTestRepositoryConfig {
     @Value("${cosmos.secondary.maxBufferedItemCount}")
     private int maxBufferedItemCount;
 
+    @Value("${cosmos.secondary.responseContinuationTokenLimitInKb}")
+    private int responseContinuationTokenLimitInKb;
+
     @Bean
     public CosmosClientBuilder secondaryCosmosClientBuilder() {
         return new CosmosClientBuilder()
@@ -65,6 +68,7 @@ public ReactiveCosmosTemplate secondaryReactiveCosmosTemplate(@Qualifier("second
                 .enableQueryMetrics(queryMetricsEnabled)
                 .maxDegreeOfParallelism(maxDegreeOfParallelism)
                 .maxBufferedItemCount(maxBufferedItemCount)
+                .responseContinuationTokenLimitInKb(responseContinuationTokenLimitInKb)
                 .build();
 
             return new ReactiveCosmosTemplate(new CosmosFactory(client, getFirstDatabase()), config, mappingCosmosConverter);
@@ -83,6 +87,7 @@ public ReactiveCosmosTemplate secondaryReactiveCosmosTemplate1(@Qualifier("secon
                 .enableQueryMetrics(queryMetricsEnabled)
                 .maxDegreeOfParallelism(maxDegreeOfParallelism)
                 .maxBufferedItemCount(maxBufferedItemCount)
+                .responseContinuationTokenLimitInKb(responseContinuationTokenLimitInKb)
                 .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

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

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

@@ -9,6 +9,8 @@ cosmos.queryMetricsEnabled=true
 cosmos.maxDegreeOfParallelism=0
 # Max number of items to buffer
 cosmos.maxBufferedItemCount=0
+# Max size of the response continuation token
+cosmos.responseContinuationTokenLimitInKb=0
 
 # Secondary DataSource Config
 cosmos.secondary.uri=${NEW_ACCOUNT_HOST}
@@ -21,5 +23,5 @@ cosmos.secondary.queryMetricsEnabled=true
 cosmos.secondary.maxDegreeOfParallelism=0
 # Max number of items to buffer
 cosmos.secondary.maxBufferedItemCount=0
-
-
+# Max size of the response continuation token
+cosmos.secondary.responseContinuationTokenLimitInKb=0

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

@@ -14,6 +14,7 @@
 
 #### Features Added
 * Exposed `maxBufferedItemCount` feature from `CosmosQueryRequestOptions` through `application.properties` flag - See [PR 30921](https://github.com/Azure/azure-sdk-for-java/pull/30921)
+* Exposed `responseContinuationTokenLimitInKb` feature from `CosmosQueryRequestOptions` through `application.properties` flag - See [PR 30980](https://github.com/Azure/azure-sdk-for-java/pull/30980)
 
 #### Bugs Fixed
 * Fixing pagination bug when performing a cross-partition query to fill every page and fix the total page count reporting. - See [PR 30694](https://github.com/Azure/azure-sdk-for-java/pull/30694)

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

@@ -123,6 +123,7 @@ In addition to setting the flag, implement `ResponseDiagnosticsProcessor` to log
 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.
+Set `responseContinuationTokenLimitInKb` flag to an integer in application.properties to allow the user to limit the length of the continuation token in the query response. The continuation token contains both required and optional fields. The required fields are necessary for resuming the execution from where it was stoped. The optional fields may contain serialized index lookup work that was done but not yet utilized. This avoids redoing the work again in subsequent continuations and hence improve the query performance. Setting the maximum continuation size to 1KB, the Azure Cosmos DB service will only serialize required fields. Starting from 2KB, the Azure Cosmos DB service would serialize as much as it could fit till it reaches the maximum specified size.
 
 ```java readme-sample-AppConfiguration
 @Configuration
@@ -151,6 +152,9 @@ public class AppConfiguration extends AbstractCosmosConfiguration {
 
     @Value("${azure.cosmos.maxBufferedItemCount}")
     private int maxBufferedItemCount;
+    
+    @Value("${azure.cosmos.responseContinuationTokenLimitInKb}")
+    private int responseContinuationTokenLimitInKb;
 
     private AzureKeyCredential azureKeyCredential;
 
@@ -171,6 +175,7 @@ public class AppConfiguration extends AbstractCosmosConfiguration {
                            .enableQueryMetrics(queryMetricsEnabled)
                            .maxDegreeOfParallelism(maxDegreeOfParallelism)
                            .maxBufferedItemCount(maxBufferedItemCount)
+                           .responseContinuationTokenLimitInKb(responseContinuationTokenLimitInKb)
                            .responseDiagnosticsProcessor(new ResponseDiagnosticsProcessorImplementation())
                            .build();
     }
@@ -214,6 +219,7 @@ public CosmosConfig cosmosConfig() {
                        .enableQueryMetrics(queryMetricsEnabled)
                        .maxDegreeOfParallelism(maxDegreeOfParallelism)
                        .maxBufferedItemCount(maxBufferedItemCount)
+                       .responseContinuationTokenLimitInKb(responseContinuationTokenLimitInKb)
                        .responseDiagnosticsProcessor(new ResponseDiagnosticsProcessorImplementation())
                        .build();
 }
@@ -692,6 +698,7 @@ public class SecondaryDatasourceConfiguration {
             .enableQueryMetrics(true)
             .maxDegreeOfParallelism(0)
             .maxBufferedItemCount(0)
+            .responseContinuationTokenLimitInKb(0)
             .responseDiagnosticsProcessor(new ResponseDiagnosticsProcessorImplementation())
             .build();
     }
@@ -729,12 +736,13 @@ public CosmosConfig getCosmosConfig() {
         .enableQueryMetrics(true)
         .maxDegreeOfParallelism(0)
         .maxBufferedItemCount(0)
+        .responseContinuationTokenLimitInKb(0)
         .responseDiagnosticsProcessor(new ResponseDiagnosticsProcessorImplementation())
         .build();
 }
 ```
 
-- Besides, if you want to define `queryMetricsEnabled`, `ResponseDiagnosticsProcessor`, `maxDegreeOfParallelism` or `maxBufferedItemCount` , you can create the `CosmosConfig` for your cosmos template.
+- Besides, if you want to define `queryMetricsEnabled`, `ResponseDiagnosticsProcessor`, `maxDegreeOfParallelism`, `maxBufferedItemCount` or `responseContinuationTokenLimitInKb` , you can create the `CosmosConfig` for your cosmos template.
 
 ```java
 @Bean("secondaryCosmosConfig")
@@ -743,6 +751,7 @@ public CosmosConfig getCosmosConfig() {
         .enableQueryMetrics(true)
         .maxDegreeOfParallelism(0)
         .maxBufferedItemCount(0)
+        .responseContinuationTokenLimitInKb(0)
         .responseDiagnosticsProcessor(new ResponseDiagnosticsProcessorImplementation())
         .build();
 }

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

@@ -21,6 +21,8 @@ public class CosmosConfig {
 
     private final int maxBufferedItemCount;
 
+    private final int responseContinuationTokenLimitInKb;
+
     /**
      * Initialization
      *
@@ -49,6 +51,7 @@ public CosmosConfig(ResponseDiagnosticsProcessor responseDiagnosticsProcessor,
         this.queryMetricsEnabled = queryMetricsEnabled;
         this.maxDegreeOfParallelism = 0;
         this.maxBufferedItemCount = 0;
+        this.responseContinuationTokenLimitInKb = 0;
     }
 
     /**
@@ -69,6 +72,7 @@ public CosmosConfig(ResponseDiagnosticsProcessor responseDiagnosticsProcessor,
         this.queryMetricsEnabled = queryMetricsEnabled;
         this.maxDegreeOfParallelism = maxDegreeOfParallelism;
         this.maxBufferedItemCount = 0;
+        this.responseContinuationTokenLimitInKb = 0;
     }
 
     /**
@@ -91,6 +95,32 @@ public CosmosConfig(ResponseDiagnosticsProcessor responseDiagnosticsProcessor,
         this.queryMetricsEnabled = queryMetricsEnabled;
         this.maxDegreeOfParallelism = maxDegreeOfParallelism;
         this.maxBufferedItemCount = maxBufferedItemCount;
+        this.responseContinuationTokenLimitInKb = 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}
+     * @param responseContinuationTokenLimitInKb must not be {@literal null}
+     */
+    @ConstructorProperties({"responseDiagnosticsProcessor", "databaseThroughputConfig", "queryMetricsEnabled", "maxDegreeOfParallelism", "maxBufferedItemCount", "responseContinuationTokenLimitInKb"})
+    CosmosConfig(ResponseDiagnosticsProcessor responseDiagnosticsProcessor,
+                 DatabaseThroughputConfig databaseThroughputConfig,
+                 boolean queryMetricsEnabled,
+                 int maxDegreeOfParallelism,
+                 int maxBufferedItemCount,
+                 int responseContinuationTokenLimitInKb) {
+        this.responseDiagnosticsProcessor = responseDiagnosticsProcessor;
+        this.databaseThroughputConfig = databaseThroughputConfig;
+        this.queryMetricsEnabled = queryMetricsEnabled;
+        this.maxDegreeOfParallelism = maxDegreeOfParallelism;
+        this.maxBufferedItemCount = maxBufferedItemCount;
+        this.responseContinuationTokenLimitInKb = responseContinuationTokenLimitInKb;
     }
 
     /**
@@ -129,6 +159,15 @@ public int getMaxBufferedItemCount() {
         return maxBufferedItemCount;
     }
 
+    /**
+     * Gets the value of responseContinuationTokenLimitInKb
+     *
+     * @return int, value of responseContinuationTokenLimitInKb
+     */
+    public int getResponseContinuationTokenLimitInKb() {
+        return responseContinuationTokenLimitInKb;
+    }
+
     /**
      * Gets the database throughput configuration.
      *
@@ -156,6 +195,7 @@ public static class CosmosConfigBuilder {
         private boolean queryMetricsEnabled;
         private int maxDegreeOfParallelism;
         private int maxBufferedItemCount;
+        private int responseContinuationTokenLimitInKb;
         CosmosConfigBuilder() {
         }
 
@@ -205,6 +245,17 @@ public CosmosConfigBuilder maxBufferedItemCount(int maxBufferedItemCount) {
             return this;
         }
 
+        /**
+         * Set responseContinuationTokenLimitInKb
+         *
+         * @param responseContinuationTokenLimitInKb value to initialize
+         * @return CosmosConfigBuilder
+         */
+        public CosmosConfigBuilder responseContinuationTokenLimitInKb(int responseContinuationTokenLimitInKb) {
+            this.responseContinuationTokenLimitInKb = responseContinuationTokenLimitInKb;
+            return this;
+        }
+
         /**
          * Enable database throughput
          *
@@ -224,7 +275,7 @@ public CosmosConfigBuilder enableDatabaseThroughput(boolean autoscale, int reque
          */
         public CosmosConfig build() {
             return new CosmosConfig(this.responseDiagnosticsProcessor, this.databaseThroughputConfig, this.queryMetricsEnabled,
-                this.maxDegreeOfParallelism, this.maxBufferedItemCount);
+                this.maxDegreeOfParallelism, this.maxBufferedItemCount, this.responseContinuationTokenLimitInKb);
         }
 
         @Override
@@ -235,6 +286,7 @@ public String toString() {
                 + ", queryMetricsEnabled=" + queryMetricsEnabled
                 + ", maxDegreeOfParallelism=" + maxDegreeOfParallelism
                 + ", maxBufferedItemCount=" + maxBufferedItemCount
+                + ", responseContinuationTokenLimitInKb=" + responseContinuationTokenLimitInKb
                 + '}';
         }
     }

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

@@ -79,6 +79,7 @@ public class CosmosTemplate implements CosmosOperations, ApplicationContextAware
     private final boolean queryMetricsEnabled;
     private final int maxDegreeOfParallelism;
     private final int maxBufferedItemCount;
+    private final int responseContinuationTokenLimitInKb;
     private final CosmosAsyncClient cosmosAsyncClient;
     private final DatabaseThroughputConfig databaseThroughputConfig;
 
@@ -134,6 +135,7 @@ public CosmosTemplate(CosmosFactory cosmosFactory,
         this.queryMetricsEnabled = cosmosConfig.isQueryMetricsEnabled();
         this.maxDegreeOfParallelism = cosmosConfig.getMaxDegreeOfParallelism();
         this.maxBufferedItemCount = cosmosConfig.getMaxBufferedItemCount();
+        this.responseContinuationTokenLimitInKb = cosmosConfig.getResponseContinuationTokenLimitInKb();
         this.databaseThroughputConfig = cosmosConfig.getDatabaseThroughputConfig();
     }
 
@@ -292,6 +294,7 @@ public <T> T findById(String containerName, Object id, Class<T> domainType) {
         options.setQueryMetricsEnabled(this.queryMetricsEnabled);
         options.setMaxDegreeOfParallelism(this.maxDegreeOfParallelism);
         options.setMaxBufferedItemCount(this.maxBufferedItemCount);
+        options.setResponseContinuationTokenLimitInKb(this.responseContinuationTokenLimitInKb);
         return cosmosAsyncClient
             .getDatabase(this.databaseName)
             .getContainer(containerName)
@@ -418,6 +421,7 @@ public <T> Iterable<T> findAll(PartitionKey partitionKey, final Class<T> domainT
         cosmosQueryRequestOptions.setQueryMetricsEnabled(this.queryMetricsEnabled);
         cosmosQueryRequestOptions.setMaxDegreeOfParallelism(this.maxDegreeOfParallelism);
         cosmosQueryRequestOptions.setMaxBufferedItemCount(this.maxBufferedItemCount);
+        cosmosQueryRequestOptions.setResponseContinuationTokenLimitInKb(this.responseContinuationTokenLimitInKb);
 
         return cosmosAsyncClient
             .getDatabase(this.databaseName)
@@ -745,6 +749,7 @@ private <T> Slice<T> sliceQuery(SqlQuerySpec querySpec,
         cosmosQueryRequestOptions.setQueryMetricsEnabled(this.queryMetricsEnabled);
         cosmosQueryRequestOptions.setMaxDegreeOfParallelism(this.maxDegreeOfParallelism);
         cosmosQueryRequestOptions.setMaxBufferedItemCount(this.maxBufferedItemCount);
+        cosmosQueryRequestOptions.setResponseContinuationTokenLimitInKb(this.responseContinuationTokenLimitInKb);
         partitionKeyValue.ifPresent(o -> {
             LOGGER.debug("Setting partition key {}", o);
             cosmosQueryRequestOptions.setPartitionKey(new PartitionKey(o));
@@ -892,6 +897,7 @@ private Long getCountValue(SqlQuerySpec querySpec, String containerName) {
         options.setQueryMetricsEnabled(this.queryMetricsEnabled);
         options.setMaxDegreeOfParallelism(this.maxDegreeOfParallelism);
         options.setMaxBufferedItemCount(this.maxBufferedItemCount);
+        options.setResponseContinuationTokenLimitInKb(this.responseContinuationTokenLimitInKb);
 
         return executeQuery(querySpec, containerName, options)
             .publishOn(Schedulers.parallel())
@@ -922,6 +928,7 @@ private <T> Flux<JsonNode> findItemsAsFlux(@NonNull CosmosQuery query,
         cosmosQueryRequestOptions.setQueryMetricsEnabled(this.queryMetricsEnabled);
         cosmosQueryRequestOptions.setMaxDegreeOfParallelism(this.maxDegreeOfParallelism);
         cosmosQueryRequestOptions.setMaxBufferedItemCount(this.maxBufferedItemCount);
+        cosmosQueryRequestOptions.setResponseContinuationTokenLimitInKb(this.responseContinuationTokenLimitInKb);
         Optional<Object> partitionKeyValue = query.getPartitionKeyValue(domainType);
         partitionKeyValue.ifPresent(o -> {
             LOGGER.debug("Setting partition key {}", o);
@@ -951,6 +958,7 @@ private Flux<JsonNode> getJsonNodeFluxFromQuerySpec(
         cosmosQueryRequestOptions.setQueryMetricsEnabled(this.queryMetricsEnabled);
         cosmosQueryRequestOptions.setMaxDegreeOfParallelism(this.maxDegreeOfParallelism);
         cosmosQueryRequestOptions.setMaxBufferedItemCount(this.maxBufferedItemCount);
+        cosmosQueryRequestOptions.setResponseContinuationTokenLimitInKb(this.responseContinuationTokenLimitInKb);
 
         return cosmosAsyncClient
                    .getDatabase(this.databaseName)