Cosmos async query_items and query_conflicts break if partition_key is specified

[stevesimmons]

• Package Name: azure-cosmos
• Package Version: 4.3.0b2
• Operating System: Linux
• Python Version: 3.9.7

Describe the bug

azure.cosmos.aio.container.query_items(sql, partition_key=x) fails with the error

RuntimeWarning: coroutine 'ContainerProxy._set_partition_key' was never awaited

when run on a single partition.

It runs correctly when no partition key is given and instead enable_cross_partition_query is set.

The same error is present in query_conflicts(), though I haven't tested it.

To reproduce

Run the example for async container.query_items() in <https://github.com/Azure/azure-sdk-for-python/blob/main/sdk/cosmos/azure-cosmos/samples/nonpartitioned_container_operations_async.py`.

This executes the query with:

[item async for item in container.query_items(sql, enable_cross_partition_query=True)]

and returns the list of items, as expected.

However when trying the same query for a single partition:

[item async for item in container.query_items(sql, partition_key="test")]

the result is an async error while setting the partition key:

RuntimeWarning: coroutine 'ContainerProxy._set_partition_key' was never awaited

Root cause

The problem arises in lines 348-349 of azure/cosmos/aio/container.py, in the function query_items:

        if partition_key is not None:
            feed_options["partitionKey"] = self._set_partition_key(partition_key)

This _set_partition_key() function is actually an async function. So the two lines should be:

        if partition_key is not None:
            feed_options["partitionKey"] = await self._set_partition_key(partition_key)

This version of the code is already present in read_item() and delete_item(). It is wrong in query_items() and also query_conflicts().

This in turn makes the query_items() function async. Thus line 285 needs the async keyword added:

    async def query_items(

The same changes need to be made to query_conflicts() at lines 619 and 644.

Finally, that means the examples now need an await as well:

# From
[item async for item in container.query_items(sql, partition_key="test")]
[item async for item in container.query_items(sql, enable_cross_partition_query=True)]
# To
[item async for item in await container.query_items(sql, partition_key="test")]
[item async for item in await container.query_items(sql, enable_cross_partition_query=True)]

[swathipil]

Hi @stevesimmons - thanks for your question! We'll investigate asap!