From 6052f8c4e782e3d881382d5d47d114fa72363f30 Mon Sep 17 00:00:00 2001 From: Srinath Narayanan Date: Thu, 1 Aug 2019 10:34:21 -0700 Subject: [PATCH 1/2] fixed bug with maxItemCount in order by queries --- .../execution_context/document_producer.py | 5 +-- .../multi_execution_aggregator.py | 2 +- sdk/cosmos/azure-cosmos/test/query_tests.py | 42 +++++++++++++++++++ 3 files changed, 45 insertions(+), 4 deletions(-) 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 c74a60e3bade..8dbdd2ce5f77 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 63a2dfb3b706..6dc1df35dd79 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..97913e411b83 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, 12) + + 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 From fd93ec5494d9c7c1797317c90794d6cab50ec103 Mon Sep 17 00:00:00 2001 From: Srinath Narayanan Date: Thu, 1 Aug 2019 10:50:52 -0700 Subject: [PATCH 2/2] fixed merge conflicts --- .../cosmos/_execution_context/document_producer.py | 5 ++--- .../_execution_context/multi_execution_aggregator.py | 2 +- sdk/cosmos/azure-cosmos/test/query_tests.py | 10 +++++----- 3 files changed, 8 insertions(+), 9 deletions(-) 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 97913e411b83..ca6c8a377fc0 100644 --- a/sdk/cosmos/azure-cosmos/test/query_tests.py +++ b/sdk/cosmos/azure-cosmos/test/query_tests.py @@ -1,7 +1,7 @@ import unittest import uuid import azure.cosmos.cosmos_client as cosmos_client -import azure.cosmos.retry_utility as retry_utility +import azure.cosmos._retry_utility as retry_utility import pytest import test_config @@ -197,16 +197,16 @@ def test_max_item_count_honored_in_order_by_query(self): ) # 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, 12) + 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 + 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 + retry_utility.ExecuteFunction = self.OriginalExecuteFunction self.assertEquals(self.count, expected_count) self.count = 0