[Cosmos] Improve system key computation logic (Azure#23821)

* initial commit

* Client Constructor (Azure#20310)

* Removed some stuff

* Looking at constructors

* Updated request

* Added client close

* working client creation

Co-authored-by: simorenoh <[email protected]>

* read database

database read works, but ignored exception is returned:
Fatal error on SSL transport
NoneType has no attribute 'send' (_loop._proactor.send)
RuntimeError: Event loop is closed
Unclosed connector/ connection

* Update simon_testfile.py

* with coroutine

Added methods needed to use async with when initializing client, but logs output "Exception ignored... Runtime Error: Event loop is closed"

* Update simon_testfile.py

* small changes

* async with returns no exceptions

* async read container

* async item read

* cleaning up

* create item/ database methods

* item delete working

* docs replace functionality

missing upsert and other resources

* upsert functionality

missing read_all_items and both query methods for container class

* missing query methods

* CRUD for udf, sproc, triggers

* initial query logic + container methods

* missing some execution logic and tests

* oops

* fully working queries

* small fix to query_items()

also fixed README and added examples_async

* Update _cosmos_client_connection_async.py

* Update _cosmos_client_connection.py

* documentation update

* updated MIT dates and get_user_client() description

* Update CHANGELOG.md

* Delete simon_testfile.py

* leftover retry utility

* Update README.md

* docs and removed six package

* changes based on comments

still missing discussion resolution on SSL verification and tests for async functionality under test module (apart from samples which are basically end to end tests)

* small change in type hints

* updated readme

* fixes based on conversations

* added missing type comments

* update changelog for ci pipeline

* added typehints, moved params into keywords, added decorators, made _connection_policy private

* changes based on sync with central sdk

* remove is_system_key from scripts (only used in execute_sproc)

is_system_key verifies that an empty partition key is properly dealt with if ['partitionKey']['systemKey'] exists in the container options - however, we do not allow containers to be created with empty partition key values in the python sdk, so the functionality is needless

* Revert "remove is_system_key from scripts (only used in execute_sproc)"

Reverting last commit, will find way to init is_system_key for now

* async script proxy using composition

* pylint

* capitalized constants

* Apply suggestions from code review

Clarifying comments for README

Co-authored-by: Gahl Levy <[email protected]>

* closing python code snippet

* last doc updates

* Update sdk/cosmos/azure-cosmos/CHANGELOG.md

Co-authored-by: Simon Moreno <[email protected]>

* version update

* cosmos updates for release

* fix connection string comma

* Update CHANGELOG.md

* fixing extra await keyword in sample

* Update CHANGELOG.md

* Update CHANGELOG.md

* Update cosmos_client.py

* Update container.py

* Update container.py

* changelog

Co-authored-by: annatisch <[email protected]>
Co-authored-by: Gahl Levy <[email protected]>
Co-authored-by: Travis Prescott <[email protected]>

sdk/cosmos/azure-cosmos/CHANGELOG.md

@@ -3,8 +3,11 @@
 ### 4.3.0b4 (Unreleased)
 
 #### Features Added
-- Added support for AAD authentication for the async client
-- Added support for AAD authentication for the sync client
+- Added support for AAD authentication for the async client.
+- Added support for AAD authentication for the sync client.
+
+#### Other Changes
+- Changed `_set_partition_key` return typehint in async client.
 
 ### 4.3.0b3 (2022-03-10)
 

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

@@ -22,7 +22,7 @@
 """Create, read, update and delete items in the Azure Cosmos DB SQL API service.
 """
 
-from typing import Any, Dict, List, Optional, Union, cast
+from typing import Any, Dict, List, Optional, Union, cast, Awaitable
 from azure.core.async_paging import AsyncItemPaged
 
 from azure.core.tracing.decorator import distributed_trace  # pylint: disable=unused-import
@@ -106,11 +106,12 @@ def _get_conflict_link(self, conflict_or_link):
             return u"{}/conflicts/{}".format(self.container_link, conflict_or_link)
         return conflict_or_link["_self"]
 
-    async def _set_partition_key(self, partition_key):
+    def _set_partition_key(self, partition_key) -> Union[str, Awaitable]:
         if partition_key == NonePartitionKeyValue:
-            return CosmosClientConnection._return_undefined_or_empty_partition_key(await self.is_system_key)
+            return CosmosClientConnection._return_undefined_or_empty_partition_key(self.is_system_key)
         return partition_key
 
+
     @distributed_trace_async
     async def read(
             self,
@@ -231,8 +232,7 @@ async def read_item(
         doc_link = self._get_document_link(item)
         request_options = _build_options(kwargs)
         response_hook = kwargs.pop('response_hook', None)
-        if partition_key is not None:
-            request_options["partitionKey"] = await self._set_partition_key(partition_key)
+        request_options["partitionKey"] = self._set_partition_key(partition_key)
         max_integrated_cache_staleness_in_ms = kwargs.pop('max_integrated_cache_staleness_in_ms', None)
         if max_integrated_cache_staleness_in_ms is not None:
             validate_cache_staleness_value(max_integrated_cache_staleness_in_ms)
@@ -399,7 +399,7 @@ def query_items_change_feed(
             feed_options["partitionKeyRangeId"] = partition_key_range_id
         partition_key = kwargs.pop("partitionKey", None)
         if partition_key is not None:
-            feed_options["partitionKey"] = partition_key
+            feed_options["partitionKey"] = self._set_partition_key(partition_key)
         if is_start_from_beginning is not None:
             feed_options["isStartFromBeginning"] = is_start_from_beginning
         if max_item_count is not None:
@@ -534,8 +534,7 @@ async def delete_item(
         """
         request_options = _build_options(kwargs)
         response_hook = kwargs.pop('response_hook', None)
-        if partition_key is not None:
-            request_options["partitionKey"] = await self._set_partition_key(partition_key)
+        request_options["partitionKey"] = self._set_partition_key(partition_key)
         if pre_trigger_include is not None:
             request_options["preTriggerInclude"] = pre_trigger_include
         if post_trigger_include is not None:
@@ -691,9 +690,7 @@ async def read_conflict(
         """
         request_options = _build_options(kwargs)
         response_hook = kwargs.pop('response_hook', None)
-        if partition_key is not None:
-            request_options["partitionKey"] = await self._set_partition_key(partition_key)
-
+        request_options["partitionKey"] = self._set_partition_key(partition_key)
         result = await self.client_connection.ReadConflict(
             conflict_link=self._get_conflict_link(conflict), options=request_options, **kwargs
         )
@@ -721,9 +718,7 @@ async def delete_conflict(
         """
         request_options = _build_options(kwargs)
         response_hook = kwargs.pop('response_hook', None)
-        if partition_key is not None:
-            request_options["partitionKey"] = await self._set_partition_key(partition_key)
-
+        request_options["partitionKey"] = self._set_partition_key(partition_key)
         result = await self.client_connection.DeleteConflict(
             conflict_link=self._get_conflict_link(conflict), options=request_options, **kwargs
         )

sdk/cosmos/azure-cosmos/azure/cosmos/partition_key.py

@@ -22,6 +22,7 @@
 """Create partition keys in the Azure Cosmos DB SQL API service.
 """
 
+
 class NonePartitionKeyValue(object):
     """Represents None value for partitionKey when it's missing in a container.
     """