Implement continuation token size limit (#30600)

* implement continuation token size limit

this adds continuation token size limit tot he python sdk. The new option for item queries accepts an integer value greater than 1 and is represented in terms of KB, ie 2 is 2kb or 2048 Bytes.

* update gitignore

* added samples

added samples for using the new feature of limiting continuation token size.

* Pylint Error Fix

This should prevent CSpell error for necessary string literal.

* additional pylint fixes

Ignore line too long

* added provisional keyword

Added provisional for the new keyword we are adding with this feature.

* Additional context info

* Added basic test for continuation token size limit

For now this test verifies that the continuation token limit size was passed in the request headers.

* pylint fixes

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

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

---------

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

.gitignore

@@ -148,4 +148,4 @@ sdk/cosmos/azure-cosmos/test/test_config.py
 *_python.json
 
 # temporary folder to refresh SDK with cadl
-TempTypeSpecFiles/
+TempTypeSpecFiles/

sdk/cosmos/azure-cosmos/CHANGELOG.md

@@ -14,6 +14,7 @@
 
 #### Features Added
 * Added conditional patching for Patch operations. See [PR 30455](https://github.com/Azure/azure-sdk-for-python/pull/30455).
+* Added **provisional** ability to limit Continuation Token size when querying for items. See [PR 30600](https://github.com/Azure/azure-sdk-for-python/pull/30600)
 
 #### Bugs Fixed
 * Fixed bug with non english locales causing an error with the RFC 1123 Date Format. See [PR 30125](https://github.com/Azure/azure-sdk-for-python/pull/30125).

sdk/cosmos/azure-cosmos/azure/cosmos/_base.py

@@ -223,6 +223,9 @@ def GetHeaders(  # pylint: disable=too-many-statements,too-many-branches
     if options.get("populateQueryMetrics"):
         headers[http_constants.HttpHeaders.PopulateQueryMetrics] = options["populateQueryMetrics"]
 
+    if options.get("responseContinuationTokenLimitInKb"):
+        headers[http_constants.HttpHeaders.ResponseContinuationTokenLimitInKb] = options["responseContinuationTokenLimitInKb"] # pylint: disable=line-too-long
+
     if cosmos_client_connection.master_key:
         #formatedate guarantees RFC 1123 date format regardless of current locale
         headers[http_constants.HttpHeaders.XDate] = formatdate(timeval=None, localtime=False, usegmt=True)

sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py

@@ -320,6 +320,9 @@ def query_items(
         :keyword dict[str, str] initial_headers: Initial headers to be sent as part of the request.
         :keyword response_hook: A callable invoked with the response metadata.
         :paramtype response_hook: Callable[[Dict[str, str], AsyncItemPaged[Dict[str, Any]]], None]
+        :keyword int response_continuation_token_limit_in_kb: **provisional keyword** The size limit in kb of the
+        response continuation token in the query response. Valid values are positive integers.
+        A value of 0 is the same as not passing a value (default no limit).
         :keyword int max_integrated_cache_staleness_in_ms: The max cache staleness for the integrated cache in
             milliseconds. For accounts configured to use the integrated cache, using Session or Eventual consistency,
             responses are guaranteed to be no staler than this value.
@@ -366,9 +369,13 @@ def query_items(
             feed_options["maxIntegratedCacheStaleness"] = max_integrated_cache_staleness_in_ms
         correlated_activity_id = GenerateGuidId()
         feed_options["correlatedActivityId"] = correlated_activity_id
+        response_continuation_token_limit_in_kb = kwargs.pop("response_continuation_token_limit_in_kb", None)
+        if response_continuation_token_limit_in_kb is not None:
+            feed_options["responseContinuationTokenLimitInKb"] = response_continuation_token_limit_in_kb
         if hasattr(response_hook, "clear"):
             response_hook.clear()
 
+
         parameters = kwargs.pop('parameters', None)
         items = self.client_connection.QueryItems(
             database_or_container_link=self.container_link,

sdk/cosmos/azure-cosmos/azure/cosmos/container.py

@@ -354,6 +354,9 @@ def query_items(
         :keyword str session_token: Token for use with Session consistency.
         :keyword dict[str,str] initial_headers: Initial headers to be sent as part of the request.
         :keyword Callable response_hook: A callable invoked with the response metadata.
+        :keyword int response_continuation_token_limit_in_kb: **provisional keyword** The size limit in kb of the
+        response continuation token in the query response. Valid values are positive integers.
+        A value of 0 is the same as not passing a value (default no limit).
         :keyword int max_integrated_cache_staleness_in_ms: The max cache staleness for the integrated cache in
             milliseconds. For accounts configured to use the integrated cache, using Session or Eventual consistency,
             responses are guaranteed to be no staler than this value.
@@ -396,6 +399,9 @@ def query_items(
             feed_options["maxIntegratedCacheStaleness"] = max_integrated_cache_staleness_in_ms
         correlated_activity_id = GenerateGuidId()
         feed_options["correlatedActivityId"] = correlated_activity_id
+        response_continuation_token_limit_in_kb = kwargs.pop("response_continuation_token_limit_in_kb", None)
+        if response_continuation_token_limit_in_kb is not None:
+            feed_options["responseContinuationTokenLimitInKb"] = response_continuation_token_limit_in_kb
         if hasattr(response_hook, "clear"):
             response_hook.clear()
 

sdk/cosmos/azure-cosmos/azure/cosmos/http_constants.py

@@ -99,6 +99,7 @@ class HttpHeaders(object):
     # Our custom DocDB headers
     Continuation = "x-ms-continuation"
     PageSize = "x-ms-max-item-count"
+    ResponseContinuationTokenLimitInKb = "x-ms-documentdb-responsecontinuationtokenlimitinkb"  # cspell:disable-line
 
     # Request sender generated. Simply echoed by backend.
     ActivityId = "x-ms-activity-id"

sdk/cosmos/azure-cosmos/samples/document_management.py

@@ -191,6 +191,23 @@ def delete_all_items_by_partition_key(db, partitionkey):
     for doc in item_list:
         print('Item Id: {0}; Partition Key: {1}'.format(doc.get('id'), doc.get("company")))
 
+
+def query_items_with_continuation_token_size_limit(container, doc_id):
+    print('\n1.11 Query Items With Continuation Token Size Limit.\n')
+
+    size_limit_in_kb = 8
+    sales_order = get_sales_order(doc_id)
+    container.create_item(body=sales_order)
+
+    # set response_continuation_token_limit_in_kb to 8 to limit size to 8KB
+    items = list(container.query_items(
+        query="SELECT * FROM r",
+        partition_key=doc_id,
+        response_continuation_token_limit_in_kb=size_limit_in_kb
+    ))
+
+    print('Continuation Token size has been limited to {}KB.'.format(size_limit_in_kb))
+
 def get_sales_order(item_id):
     order1 = {'id' : item_id,
             'account_number' : 'Account1',
@@ -259,6 +276,7 @@ def run_sample():
         patch_item(container, 'SalesOrder1')
         delete_item(container, 'SalesOrder1')
         delete_all_items_by_partition_key(db, "CompanyA")
+        query_items_with_continuation_token_size_limit(container, 'SalesOrder1')
 
         # cleanup database after sample
         try:

sdk/cosmos/azure-cosmos/samples/document_management_async.py

@@ -208,6 +208,23 @@ async def delete_all_items_by_partition_key(db, partitionkey):
     for doc in item_list:
         print('Item Id: {0}; Partition Key: {1}'.format(doc.get('id'), doc.get("company")))
 
+
+async def query_items_with_continuation_token_size_limit(container, doc_id):
+    print('\n1.11 Query Items With Continuation Token Size Limit.\n')
+
+    size_limit_in_kb = 8
+    sales_order = get_sales_order(doc_id)
+    await container.create_item(body=sales_order)
+
+    # set response_continuation_token_limit_in_kb to 8 to limit size to 8KB
+    items = container.query_items(
+        query="SELECT * FROM r",
+        partition_key=doc_id,
+        response_continuation_token_limit_in_kb=size_limit_in_kb
+    )
+
+    print('Continuation Token size has been limited to {}KB.'.format(size_limit_in_kb))
+
 def get_sales_order(item_id):
     order1 = {'id' : item_id,
             'account_number' : 'Account1',
@@ -278,6 +295,7 @@ async def run_sample():
             await patch_item(container, 'SalesOrder1')
             await delete_item(container, 'SalesOrder1')
             await delete_all_items_by_partition_key(db, "CompanyA")
+            await query_items_with_continuation_token_size_limit(container, 'SalesOrder1')
 
             # cleanup database after sample
             try:

sdk/cosmos/azure-cosmos/test/test_query.py

@@ -10,6 +10,8 @@
 import pytest
 import collections
 import test_config
+from unittest.mock import MagicMock
+from azure.cosmos import http_constants
 
 pytestmark = pytest.mark.cosmosEmulator
 
@@ -828,6 +830,23 @@ def test_value_max_query(self):
 
         self.assertListEqual(list(query_results), [None])
 
+    def side_effect_continuation_token_size_limit(self, *args, **kwargs):
+        # Extract request headers from args
+        self.assertTrue(args[2][http_constants.HttpHeaders.ResponseContinuationTokenLimitInKb] is 8)
+        raise StopIteration
+
+    def test_continuation_token_size_limit_query(self):
+        container = self.created_db.create_container_if_not_exists(
+            self.config.TEST_COLLECTION_MULTI_PARTITION_WITH_CUSTOM_PK_ID, PartitionKey(path="/pk"))
+        cosmos_client_connection = container.client_connection
+        cosmos_client_connection._CosmosClientConnection__Get = MagicMock(
+            side_effect=self.side_effect_continuation_token_size_limit)
+        try:
+            query = "Select * from c"
+            container.query_items(query, response_continuation_token_limit_in_kb=8)
+        except StopIteration:
+            pass
+
     def _MockNextFunction(self):
         if self.count < len(self.payloads):
             item, result = self.get_mock_result(self.payloads, self.count)