[Bug Fix] Query Items with none partititon key

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)

sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py

@@ -755,14 +755,13 @@ def query_items(
         kwargs["containerProperties"] = self._get_properties_with_options
 
         utils.verify_exclusive_arguments(["feed_range", "partition_key"], **kwargs)
-        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")
-            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)
@@ -773,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
         )

sdk/cosmos/azure-cosmos/azure/cosmos/container.py

@@ -935,15 +935,16 @@ 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"))
+        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
                 kwargs["prefix_partition_key_value"] = cast(_SequentialPartitionKeyType, partition_key_value)
             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")

sdk/cosmos/azure-cosmos/tests/test_query.py

@@ -469,6 +469,21 @@ def test_cross_partition_query_with_continuation_token(self):
 
         self.assertEqual(second_page['id'], second_page_fetched_with_continuation_token['id'])
 
+    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': 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
+        )
+
+        assert len(list(query_iterable)) == 1
+
     def _validate_distinct_on_different_types_and_field_orders(self, collection, query, expected_results,
                                                                get_mock_result):
         self.count = 0

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_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': None, '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]) == 1
+
     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(