Added Tests and Samples for Paginated Queries (#43472)

andrewmathew1 · Andrew Mathew · Copilot · web-flow · commit 378dc678d67c · 2025-10-17T17:18:19.000-04:00
* added tests and samples for paginated queries

* Apply suggestions from code review

Co-authored-by: Copilot &lt;175728472+Copilot@users.noreply.github.com&gt;

* added single partition pagination sample

---------

Co-authored-by: Andrew Mathew &lt;andrewmathew@microsoft.com&gt;
Co-authored-by: Copilot &lt;175728472+Copilot@users.noreply.github.com&gt;
diff --git a/sdk/cosmos/azure-cosmos/samples/document_management.py b/sdk/cosmos/azure-cosmos/samples/document_management.py
@@ -109,6 +109,74 @@ def query_items_with_continuation_token(container):
     print('The single items in the second page are {}.'.format(second_page_items_with_continuation[0].get("id")))
 
 
+def query_items_single_partition_with_pagination(container):
+    print('\n1.5a Querying with Pagination - Demonstrating max_item_count and Counting Results\n')
+    
+    # max_item_count controls how many items are returned per page, not the total number of results
+    # This is useful for controlling memory usage and processing items in batches
+    max_items_per_page = 5
+    
+    # Specify a partition key to query within a single partition
+    partition_key_value = "SalesOrder1"  # Or any partition key value you want
+
+    query_iterable = container.query_items(
+        query="SELECT * FROM c",
+        partition_key=partition_key_value,  # Query single partition
+        max_item_count=max_items_per_page
+    )
+
+    # Iterate through pages and count both pages and total items
+    total_item_count = 0
+    page_count = 0
+    
+    item_pages = query_iterable.by_page()
+    for page in item_pages:
+        page_count += 1
+        items_in_page = list(page)
+        items_in_current_page = len(items_in_page)
+        total_item_count += items_in_current_page
+        
+        print(f'Page {page_count}: Retrieved {items_in_current_page} items (max per page: {max_items_per_page})')
+        
+        # Process items in this page
+        for item in items_in_page:
+            # Do something with each item
+            pass
+    
+    print(f'\nTotal pages processed: {page_count}')
+    print(f'Total items retrieved: {total_item_count}')
+    print(f'Note: max_item_count limits items PER PAGE, not total results\n')
+
+
+def query_items_cross_partition_with_pagination(container):
+    print('\n1.5b Cross-Partition Query with Pagination\n')
+    
+    # When querying across partitions, max_item_count still controls page size
+    # but the results are gathered from multiple partitions
+    max_items_per_page = 3
+    
+    query_iterable = container.query_items(
+        query="SELECT * FROM c ORDER BY c._ts",  # Order by timestamp across all partitions
+        enable_cross_partition_query=True,
+        max_item_count=max_items_per_page
+    )
+    
+    total_item_count = 0
+    page_count = 0
+    
+    item_pages = query_iterable.by_page()
+    for page in item_pages:
+        page_count += 1
+        items_in_page = list(page)
+        total_item_count += len(items_in_page)
+        
+        print(f'Page {page_count}: {len(items_in_page)} items from across partitions')
+    
+    print(f'\nCross-partition query completed:')
+    print(f'  - Pages: {page_count}')
+    print(f'  - Total items: {total_item_count}')
+
+
 def replace_item(container, doc_id):
     print('\n1.6 Replace an Item\n')
 
@@ -525,6 +593,8 @@ def run_sample():
         read_items(container)
         query_items(container, 'SalesOrder1')
         query_items_with_continuation_token(container)
+        query_items_single_partition_with_pagination(container)
+        query_items_cross_partition_with_pagination(container)
         replace_item(container, 'SalesOrder1')
         replace_item_using_etags(container, 'SalesOrder1')
         upsert_item(container, 'SalesOrder1')
diff --git a/sdk/cosmos/azure-cosmos/samples/document_management_async.py b/sdk/cosmos/azure-cosmos/samples/document_management_async.py
@@ -126,6 +126,73 @@ async def query_items_with_continuation_token(container):
     print('The single items in the second page are {}'.format(second_page_items_with_continuation[0].get("id")))
 
 
+async def query_items_single_partition_with_pagination(container):
+    print('\n1.5a Querying with Pagination - Demonstrating max_item_count and Counting Results\n')
+    
+    # max_item_count controls how many items are returned per page, not the total number of results
+    # This is useful for controlling memory usage and processing items in batches
+    max_items_per_page = 5
+
+    # Specify a partition key to query within a single partition
+    partition_key_value = "SalesOrder1"  # Or any partition key value you want
+    
+    query_iterable = container.query_items(
+        query="SELECT * FROM c",
+        partition_key=partition_key_value,  # Query single partition
+        max_item_count=max_items_per_page
+    )
+    
+    # Iterate through pages and count both pages and total items
+    total_item_count = 0
+    page_count = 0
+    
+    item_pages = query_iterable.by_page()
+    async for page in item_pages:
+        page_count += 1
+        items_in_page = [item async for item in page]
+        items_in_current_page = len(items_in_page)
+        total_item_count += items_in_current_page
+        
+        print(f'Page {page_count}: Retrieved {items_in_current_page} items (max per page: {max_items_per_page})')
+        
+        # Process items in this page
+        for item in items_in_page:
+            # Do something with each item
+            pass
+    
+    print(f'\nTotal pages processed: {page_count}')
+    print(f'Total items retrieved: {total_item_count}')
+    print(f'Note: max_item_count limits items PER PAGE, not total results\n')
+
+
+async def query_items_cross_partition_with_pagination(container):
+    print('\n1.5b Cross-Partition Query with Pagination\n')
+    
+    # When querying across partitions, max_item_count still controls page size
+    # but the results are gathered from multiple partitions
+    max_items_per_page = 3
+    
+    query_iterable = container.query_items(
+        query="SELECT * FROM c ORDER BY c._ts",  # Order by timestamp across all partitions
+        max_item_count=max_items_per_page
+    )
+    
+    total_item_count = 0
+    page_count = 0
+    
+    item_pages = query_iterable.by_page()
+    async for page in item_pages:
+        page_count += 1
+        items_in_page = [item async for item in page]
+        total_item_count += len(items_in_page)
+        
+        print(f'Page {page_count}: {len(items_in_page)} items from across partitions')
+    
+    print(f'\nCross-partition query completed:')
+    print(f'  - Pages: {page_count}')
+    print(f'  - Total items: {total_item_count}')
+
+
 async def replace_item(container, doc_id):
     print('\n1.6 Replace an Item\n')
 
@@ -546,6 +613,8 @@ async def run_sample():
             await read_items(container)
             await query_items(container, 'SalesOrder1')
             await query_items_with_continuation_token(container)
+            await query_items_single_partition_with_pagination(container)
+            await query_items_cross_partition_with_pagination(container)
             await replace_item(container, 'SalesOrder1')
             await replace_item_using_etags(container, 'SalesOrder1')
             await upsert_item(container, 'SalesOrder1')
diff --git a/sdk/cosmos/azure-cosmos/tests/test_query.py b/sdk/cosmos/azure-cosmos/tests/test_query.py
@@ -598,6 +598,91 @@ def test_query_request_params_none_retry_policy(self):
         retry_utility.ExecuteFunction = self.OriginalExecuteFunction
         self.created_db.delete_container(created_collection.id)
 
+    def test_query_pagination_with_max_item_count(self):
+        """Test pagination showing per-page limits and total results counting."""
+        created_collection = self.created_db.create_container(
+            "pagination_test_" + str(uuid.uuid4()),
+            PartitionKey(path="/pk"))
+        
+        # Create 20 items in a single partition
+        total_items = 20
+        partition_key_value = "test_pk"
+        for i in range(total_items):
+            document_definition = {
+                'pk': partition_key_value,
+                'id': f'item_{i}',
+                'value': i
+            }
+            created_collection.create_item(body=document_definition)
+        
+        # Test pagination with max_item_count limiting items per page
+        max_items_per_page = 7
+        query = "SELECT * FROM c WHERE c.pk = @pk ORDER BY c.value"
+        query_iterable = created_collection.query_items(
+            query=query,
+            parameters=[{"name": "@pk", "value": partition_key_value}],
+            partition_key=partition_key_value,
+            max_item_count=max_items_per_page
+        )
+        
+        # Iterate through pages and verify per-page counts
+        all_fetched_results = []
+        page_count = 0
+        item_pages = query_iterable.by_page()
+        
+        for page in item_pages:
+            page_count += 1
+            items_in_page = list(page)
+            all_fetched_results.extend(items_in_page)
+            
+            # Each page should have at most max_item_count items
+            # (last page may have fewer)
+            self.assertLessEqual(len(items_in_page), max_items_per_page)
+        
+        # Verify total results match expected count
+        self.assertEqual(len(all_fetched_results), total_items)
+        
+        # Verify we got the expected number of pages
+        # 20 items with max 7 per page = 3 pages (7, 7, 6)
+        self.assertEqual(page_count, 3)
+        
+        # Verify ordering is maintained
+        for i, item in enumerate(all_fetched_results):
+            self.assertEqual(item['value'], i)
+        
+        self.created_db.delete_container(created_collection.id)
+    
+    def test_query_pagination_without_max_item_count(self):
+        """Test pagination behavior without specifying max_item_count."""
+        created_collection = self.created_db.create_container(
+            "pagination_no_max_test_" + str(uuid.uuid4()),
+            PartitionKey(path="/pk"))
+        
+        # Create 15 items in a single partition
+        total_items = 15
+        partition_key_value = "test_pk_2"
+        for i in range(total_items):
+            document_definition = {
+                'pk': partition_key_value,
+                'id': f'item_{i}',
+                'value': i
+            }
+            created_collection.create_item(body=document_definition)
+        
+        # Query without specifying max_item_count
+        query = "SELECT * FROM c WHERE c.pk = @pk"
+        query_iterable = created_collection.query_items(
+            query=query,
+            parameters=[{"name": "@pk", "value": partition_key_value}],
+            partition_key=partition_key_value
+        )
+        
+        # Count total results
+        all_results = list(query_iterable)
+        self.assertEqual(len(all_results), total_items)
+        
+        self.created_db.delete_container(created_collection.id)
+
     def test_query_positional_args(self):
         container = self.created_db.get_container_client(self.config.TEST_MULTI_PARTITION_CONTAINER_ID)
         partition_key_value1 = "pk1"
diff --git a/sdk/cosmos/azure-cosmos/tests/test_query_async.py b/sdk/cosmos/azure-cosmos/tests/test_query_async.py
@@ -645,6 +645,91 @@ async def test_partitioned_query_response_hook_async(self):
         assert response_hook.count == 1
         await self.created_db.delete_container(created_collection.id)
 
+    async def test_query_pagination_with_max_item_count_async(self):
+        """Test pagination showing per-page limits and total results counting."""
+        created_collection = await self.created_db.create_container(
+            "pagination_test_" + str(uuid.uuid4()),
+            PartitionKey(path="/pk"))
+        
+        # Create 20 items in a single partition
+        total_items = 20
+        partition_key_value = "test_pk"
+        for i in range(total_items):
+            document_definition = {
+                'pk': partition_key_value,
+                'id': f'item_{i}',
+                'value': i
+            }
+            await created_collection.create_item(body=document_definition)
+        
+        # Test pagination with max_item_count limiting items per page
+        max_items_per_page = 7
+        query = "SELECT * FROM c WHERE c.pk = @pk ORDER BY c.value"
+        query_iterable = created_collection.query_items(
+            query=query,
+            parameters=[{"name": "@pk", "value": partition_key_value}],
+            partition_key=partition_key_value,
+            max_item_count=max_items_per_page
+        )
+        
+        # Iterate through pages and verify per-page counts
+        all_fetched_results = []
+        page_count = 0
+        item_pages = query_iterable.by_page()
+        
+        async for page in item_pages:
+            page_count += 1
+            items_in_page = [item async for item in page]
+            all_fetched_results.extend(items_in_page)
+            
+            # Each page should have at most max_item_count items
+            # (last page may have fewer)
+            assert len(items_in_page) <= max_items_per_page
+        
+        # Verify total results match expected count
+        assert len(all_fetched_results) == total_items
+        
+        # Verify we got the expected number of pages
+        # 20 items with max 7 per page = 3 pages (7, 7, 6)
+        assert page_count == 3
+        
+        # Verify ordering is maintained
+        for i, item in enumerate(all_fetched_results):
+            assert item['value'] == i
+        
+        await self.created_db.delete_container(created_collection.id)
+    
+    async def test_query_pagination_without_max_item_count_async(self):
+        """Test pagination behavior without specifying max_item_count."""
+        created_collection = await self.created_db.create_container(
+            "pagination_no_max_test_" + str(uuid.uuid4()),
+            PartitionKey(path="/pk"))
+        
+        # Create 15 items in a single partition
+        total_items = 15
+        partition_key_value = "test_pk_2"
+        for i in range(total_items):
+            document_definition = {
+                'pk': partition_key_value,
+                'id': f'item_{i}',
+                'value': i
+            }
+            await created_collection.create_item(body=document_definition)
+        
+        # Query without specifying max_item_count
+        query = "SELECT * FROM c WHERE c.pk = @pk"
+        query_iterable = created_collection.query_items(
+            query=query,
+            parameters=[{"name": "@pk", "value": partition_key_value}],
+            partition_key=partition_key_value
+        )
+        
+        # Count total results
+        all_results = [item async for item in query_iterable]
+        assert len(all_results) == total_items
+        
+        await self.created_db.delete_container(created_collection.id)
+
     async def _MockExecuteFunctionSessionRetry(self, function, *args, **kwargs):
         if args:
             if args[1].operation_type == 'SqlQuery':
diff --git a/sdk/cosmos/azure-cosmos/tests/test_query_cross_partition.py b/sdk/cosmos/azure-cosmos/tests/test_query_cross_partition.py
diff --git a/sdk/cosmos/azure-cosmos/tests/test_query_cross_partition_async.py b/sdk/cosmos/azure-cosmos/tests/test_query_cross_partition_async.py
