From f5da3c5d7177ee17e3c08b20af0f6579707f7b06 Mon Sep 17 00:00:00 2001 From: tvaron3 Date: Wed, 27 Aug 2025 20:13:16 -0700 Subject: [PATCH 1/3] fix query items with none partition key --- .../azure-cosmos/azure/cosmos/aio/_container.py | 5 ++--- .../azure-cosmos/azure/cosmos/container.py | 5 +++-- sdk/cosmos/azure-cosmos/tests/test_query.py | 16 ++++++++++++++++ .../azure-cosmos/tests/test_query_async.py | 14 ++++++++++++++ 4 files changed, 35 insertions(+), 5 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py index af4c83694ea3..55b50b45bc72 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py @@ -755,10 +755,9 @@ def query_items( kwargs["containerProperties"] = self._get_properties_with_options utils.verify_exclusive_arguments(["feed_range", "partition_key"], **kwargs) - partition_key = None + partition_key = kwargs.pop("partition_key", None) # If 'partition_key' is provided, set 'partitionKey' in 'feed_options' - if utils.valid_key_value_exist(kwargs, "partition_key"): - partition_key = kwargs.pop("partition_key") + if partition_key: feed_options["partitionKey"] = self._set_partition_key(partition_key) # If 'partition_key' or 'feed_range' is not provided, set 'enableCrossPartitionQuery' to True elif not utils.valid_key_value_exist(kwargs, "feed_range"): diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py index 364f9dd5ddbc..bc1b462d166a 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py @@ -935,8 +935,9 @@ def query_items( # pylint:disable=docstring-missing-param # Set range filters for a query. Options are either 'feed_range' or 'partition_key' utils.verify_exclusive_arguments(["feed_range", "partition_key"], **kwargs) - if utils.valid_key_value_exist(kwargs, "partition_key"): - partition_key_value = self._set_partition_key(kwargs.pop("partition_key")) + pk_value = kwargs.pop("partition_key", None) + if pk_value: + partition_key_value = self._set_partition_key(pk_value) partition_key_obj = _build_partition_key_from_properties(container_properties) if partition_key_obj._is_prefix_partition_key(partition_key_value): kwargs["prefix_partition_key_object"] = partition_key_obj diff --git a/sdk/cosmos/azure-cosmos/tests/test_query.py b/sdk/cosmos/azure-cosmos/tests/test_query.py index 365720454ae7..d397cf0b9da3 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_query.py +++ b/sdk/cosmos/azure-cosmos/tests/test_query.py @@ -469,6 +469,22 @@ def test_cross_partition_query_with_continuation_token(self): self.assertEqual(second_page['id'], second_page_fetched_with_continuation_token['id']) + def test_cross_partition_query_with_none_partition_key(self): + created_collection = self.created_db.get_container_client(self.config.TEST_MULTI_PARTITION_CONTAINER_ID) + document_definition = {'pk': 'pk1', 'id': str(uuid.uuid4())} + created_collection.create_item(body=document_definition) + document_definition = {'pk': 'pk2', 'id': str(uuid.uuid4())} + created_collection.create_item(body=document_definition) + + query = 'SELECT * from c' + query_iterable = created_collection.query_items( + query=query, + partition_key=None, + enable_cross_partition_query=True + ) + + assert len(list(query_iterable)) > 0 + def _validate_distinct_on_different_types_and_field_orders(self, collection, query, expected_results, get_mock_result): self.count = 0 diff --git a/sdk/cosmos/azure-cosmos/tests/test_query_async.py b/sdk/cosmos/azure-cosmos/tests/test_query_async.py index a3f2edc40c44..d0d2d302b96e 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_query_async.py +++ b/sdk/cosmos/azure-cosmos/tests/test_query_async.py @@ -465,6 +465,20 @@ async def test_cross_partition_query_with_continuation_token_async(self): assert second_page['id'] == second_page_fetched_with_continuation_token['id'] + async def test_cross_partition_query_with_none_partition_key_async(self): + created_collection = self.created_db.get_container_client(self.config.TEST_MULTI_PARTITION_CONTAINER_ID) + document_definition = {'pk': 'pk1', 'id': str(uuid.uuid4())} + await created_collection.create_item(body=document_definition) + document_definition = {'pk': 'pk2', 'id': str(uuid.uuid4())} + await created_collection.create_item(body=document_definition) + + query = 'SELECT * from c' + query_iterable = created_collection.query_items( + query=query, + partition_key=None) + + assert len([item async for item in query_iterable]) > 0 + async def test_value_max_query_results_async(self): container = self.created_db.get_container_client(self.config.TEST_MULTI_PARTITION_CONTAINER_ID) await container.create_item( From 95de46d9a3a677699fb0c56179aa6613763081d5 Mon Sep 17 00:00:00 2001 From: tvaron3 Date: Wed, 27 Aug 2025 20:59:39 -0700 Subject: [PATCH 2/3] add changelog --- sdk/cosmos/azure-cosmos/CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/sdk/cosmos/azure-cosmos/CHANGELOG.md b/sdk/cosmos/azure-cosmos/CHANGELOG.md index f683405407cb..16c173ae423b 100644 --- a/sdk/cosmos/azure-cosmos/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos/CHANGELOG.md @@ -12,6 +12,7 @@ * Fixed bug where during health checks read regions were marked as unavailable for write operations. See [PR 42525](https://github.com/Azure/azure-sdk-for-python/pull/42525). * Fixed bug where containers named with spaces or special characters using session consistency would fall back to eventual consistency. See [PR 42608](https://github.com/Azure/azure-sdk-for-python/pull/42608) * Fixed bug where `excluded_locations` was not being honored for some metadata calls. See [PR 42266](https://github.com/Azure/azure-sdk-for-python/pull/42266). +* Fixed bug where `partition_key` set to None was not properly handled for query operations. See [PR 42747](https://github.com/Azure/azure-sdk-for-python/pull/42747) #### Other Changes * Added session token false progress merge logic. See [42393](https://github.com/Azure/azure-sdk-for-python/pull/42393) From 116b3b3fe2ef8e6c335265b8734949d8e2514f88 Mon Sep 17 00:00:00 2001 From: tvaron3 Date: Wed, 27 Aug 2025 23:05:00 -0700 Subject: [PATCH 3/3] fix query behavior for null partition key --- sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py | 8 ++++---- sdk/cosmos/azure-cosmos/azure/cosmos/container.py | 6 +++--- sdk/cosmos/azure-cosmos/tests/test_query.py | 9 ++++----- sdk/cosmos/azure-cosmos/tests/test_query_async.py | 6 +++--- 4 files changed, 14 insertions(+), 15 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py index 55b50b45bc72..edf72f13dc06 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py @@ -755,13 +755,13 @@ def query_items( kwargs["containerProperties"] = self._get_properties_with_options utils.verify_exclusive_arguments(["feed_range", "partition_key"], **kwargs) - partition_key = kwargs.pop("partition_key", None) # If 'partition_key' is provided, set 'partitionKey' in 'feed_options' - if partition_key: - feed_options["partitionKey"] = self._set_partition_key(partition_key) + if "partition_key" in kwargs: + feed_options["partitionKey"] = self._set_partition_key(kwargs["partition_key"]) # If 'partition_key' or 'feed_range' is not provided, set 'enableCrossPartitionQuery' to True elif not utils.valid_key_value_exist(kwargs, "feed_range"): feed_options["enableCrossPartitionQuery"] = True + kwargs.pop("partition_key", None) # Set 'response_hook' response_hook = kwargs.pop("response_hook", None) @@ -772,7 +772,7 @@ def query_items( database_or_container_link=self.container_link, query=query, options=feed_options, - partition_key=partition_key, + partition_key=feed_options.get("partitionKey"), response_hook=response_hook, **kwargs ) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py index bc1b462d166a..d93d7830bc74 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py @@ -935,9 +935,8 @@ def query_items( # pylint:disable=docstring-missing-param # Set range filters for a query. Options are either 'feed_range' or 'partition_key' utils.verify_exclusive_arguments(["feed_range", "partition_key"], **kwargs) - pk_value = kwargs.pop("partition_key", None) - if pk_value: - partition_key_value = self._set_partition_key(pk_value) + if "partition_key" in kwargs: + partition_key_value = self._set_partition_key(kwargs["partition_key"]) partition_key_obj = _build_partition_key_from_properties(container_properties) if partition_key_obj._is_prefix_partition_key(partition_key_value): kwargs["prefix_partition_key_object"] = partition_key_obj @@ -945,6 +944,7 @@ def query_items( # pylint:disable=docstring-missing-param else: # Add to feed_options, only when feed_range not given and partition_key was not prefixed partition_key feed_options["partitionKey"] = partition_key_value + kwargs.pop("partition_key", None) # Set 'partition_key' for QueryItems method. This can be 'None' if feed range or prefix partition key was set partition_key = feed_options.get("partitionKey") diff --git a/sdk/cosmos/azure-cosmos/tests/test_query.py b/sdk/cosmos/azure-cosmos/tests/test_query.py index d397cf0b9da3..7132b07a6fba 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_query.py +++ b/sdk/cosmos/azure-cosmos/tests/test_query.py @@ -469,21 +469,20 @@ def test_cross_partition_query_with_continuation_token(self): self.assertEqual(second_page['id'], second_page_fetched_with_continuation_token['id']) - def test_cross_partition_query_with_none_partition_key(self): + def test_query_with_none_partition_key(self): created_collection = self.created_db.get_container_client(self.config.TEST_MULTI_PARTITION_CONTAINER_ID) document_definition = {'pk': 'pk1', 'id': str(uuid.uuid4())} created_collection.create_item(body=document_definition) - document_definition = {'pk': 'pk2', 'id': str(uuid.uuid4())} + document_definition = {'pk': None, 'id': str(uuid.uuid4())} created_collection.create_item(body=document_definition) query = 'SELECT * from c' query_iterable = created_collection.query_items( query=query, - partition_key=None, - enable_cross_partition_query=True + partition_key=None ) - assert len(list(query_iterable)) > 0 + assert len(list(query_iterable)) == 1 def _validate_distinct_on_different_types_and_field_orders(self, collection, query, expected_results, get_mock_result): diff --git a/sdk/cosmos/azure-cosmos/tests/test_query_async.py b/sdk/cosmos/azure-cosmos/tests/test_query_async.py index d0d2d302b96e..773d5f8d57e1 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_query_async.py +++ b/sdk/cosmos/azure-cosmos/tests/test_query_async.py @@ -465,11 +465,11 @@ async def test_cross_partition_query_with_continuation_token_async(self): assert second_page['id'] == second_page_fetched_with_continuation_token['id'] - async def test_cross_partition_query_with_none_partition_key_async(self): + async def test_query_with_none_partition_key_async(self): created_collection = self.created_db.get_container_client(self.config.TEST_MULTI_PARTITION_CONTAINER_ID) document_definition = {'pk': 'pk1', 'id': str(uuid.uuid4())} await created_collection.create_item(body=document_definition) - document_definition = {'pk': 'pk2', 'id': str(uuid.uuid4())} + document_definition = {'pk': None, 'id': str(uuid.uuid4())} await created_collection.create_item(body=document_definition) query = 'SELECT * from c' @@ -477,7 +477,7 @@ async def test_cross_partition_query_with_none_partition_key_async(self): query=query, partition_key=None) - assert len([item async for item in query_iterable]) > 0 + assert len([item async for item in query_iterable]) == 1 async def test_value_max_query_results_async(self): container = self.created_db.get_container_client(self.config.TEST_MULTI_PARTITION_CONTAINER_ID)