diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_execution_context/document_producer.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_execution_context/document_producer.py index f6f7db7e911e..0d4c47aef3df 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_execution_context/document_producer.py +++ b/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 diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_execution_context/multi_execution_aggregator.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_execution_context/multi_execution_aggregator.py index f0bfb8887fd1..8d2d2dc304ce 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_execution_context/multi_execution_aggregator.py +++ b/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): diff --git a/sdk/cosmos/azure-cosmos/test/query_tests.py b/sdk/cosmos/azure-cosmos/test/query_tests.py index 4a313a2e1785..ca6c8a377fc0 100644 --- a/sdk/cosmos/azure-cosmos/test/query_tests.py +++ b/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() \ No newline at end of file