fixed bug in maxItemCount propagation for Order by queries (#6608)

* fixed bug with maxItemCount in order by queries

* fixed merge conflicts

sdk/cosmos/azure-cosmos/azure/cosmos/_execution_context/document_producer.py

@@ -35,12 +35,11 @@ class _DocumentProducer(object):
     When handling an orderby query, MultiExecutionContextAggregator instantiates one instance of this class
     per target partition key range and aggregates the result of each.
     '''
-    def __init__(self, partition_key_target_range, client, collection_link, query, document_producer_comp):
+    def __init__(self, partition_key_target_range, client, collection_link, query, document_producer_comp, options):
         '''
         Constructor
         '''
-        # TODO: is that fine we build the options dict and we don't inherit it?
-        self._options = {}
+        self._options = options
         self._partition_key_target_range = partition_key_target_range
         self._doc_producer_comp = document_producer_comp
         self._client = client

sdk/cosmos/azure-cosmos/azure/cosmos/_execution_context/multi_execution_aggregator.py

@@ -142,7 +142,7 @@ def _createTargetPartitionQueryExecutionContext(self, partition_key_target_range
         else:
             query = self._query            
 
-        return document_producer._DocumentProducer(partition_key_target_range, self._client, self._resource_link, query, self._document_producer_comparator)
+        return document_producer._DocumentProducer(partition_key_target_range, self._client, self._resource_link, query, self._document_producer_comparator, self._options)
 
     def _get_target_parition_key_range(self):
 

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

@@ -1,5 +1,7 @@
 import unittest
+import uuid
 import azure.cosmos.cosmos_client as cosmos_client
+import azure.cosmos._retry_utility as retry_utility
 import pytest
 import test_config
 
@@ -172,6 +174,46 @@ def test_populate_query_metrics(self):
         self.assertTrue(len(metrics) > 1)
         self.assertTrue(all(['=' in x for x in metrics]))
 
+    def test_max_item_count_honored_in_order_by_query(self):
+        created_collection = self.config.create_multi_partition_collection_with_custom_pk_if_not_exist(self.client)
+        docs = []
+        for i in range(10):
+            document_definition = {'pk': 'pk', 'id': 'myId' + str(uuid.uuid4())}
+            docs.append(created_collection.create_item(body=document_definition))
+
+        query = 'SELECT * from c ORDER BY c._ts'
+        query_iterable = created_collection.query_items(
+            query=query,
+            max_item_count=1,
+            enable_cross_partition_query=True
+        )
+        # 1 call to get query plans, 1 call to get pkr, 10 calls to one partion with the documents, 1 call each to other 4 partitions
+        self.validate_query_requests_count(query_iterable, 16 * 2)
+
+        query_iterable = created_collection.query_items(
+            query=query,
+            max_item_count=100,
+            enable_cross_partition_query=True
+        )
+
+        # 1 call to get query plan 1 calls to one partition with the documents, 1 call each to other 4 partitions
+        self.validate_query_requests_count(query_iterable, 6 * 2)
+
+    def validate_query_requests_count(self, query_iterable, expected_count):
+        self.count = 0
+        self.OriginalExecuteFunction = retry_utility.ExecuteFunction
+        retry_utility.ExecuteFunction = self._MockExecuteFunction
+        block = query_iterable.fetch_next_block()
+        while block:
+            block = query_iterable.fetch_next_block()
+        retry_utility.ExecuteFunction = self.OriginalExecuteFunction
+        self.assertEquals(self.count, expected_count)
+        self.count = 0
+
+    def _MockExecuteFunction(self, function, *args, **kwargs):
+        self.count += 1
+        return self.OriginalExecuteFunction(function, *args, **kwargs)
+
 
 if __name__ == "__main__":
     unittest.main()