Enabling option to disable response payload on writes (Azure#37365)

* Initial draft

* Adding tests

* Renaming parameter

* Update container.py

* Renaming test file

* Fixing LINT issues

* Update container.py

* Update _base.py

* Update _base.py

* Fixing tests

* Fixing tests

* Adding support to disable response payload on write for AIO

* Update CHANGELOG.md

* Update _cosmos_client.py

* Reacting to code review comments

* Addressing code review feedback

* Addressed CR feedback

* Fixing pyLint errors

* Fixing pylint errors

* Update test_crud.py

* Fixing svc regression

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

Co-authored-by: Anna Tisch <[email protected]>

* Reacting to code review feedback.

* Update container.py

* Update test_query_vector_similarity.py

---------

Co-authored-by: Anna Tisch <[email protected]>

Author: 2 people

sdk/cosmos/azure-cosmos/CHANGELOG.md

@@ -4,6 +4,7 @@
 
 #### Features Added
 * Added Retry Policy for Container Recreate in the Python SDK. See [PR 36043](https://github.com/Azure/azure-sdk-for-python/pull/36043)
+* Added option to disable write payload on writes. See [PR 37365](https://github.com/Azure/azure-sdk-for-python/pull/37365)
 
 #### Breaking Changes
 

sdk/cosmos/azure-cosmos/azure/cosmos/_base.py

@@ -60,7 +60,8 @@
     'is_query_plan_request': 'isQueryPlanRequest',
     'supported_query_features': 'supportedQueryFeatures',
     'query_version': 'queryVersion',
-    'priority': 'priorityLevel'
+    'priority': 'priorityLevel',
+    'no_response': 'responsePayloadOnWriteDisabled'
 }
 
 # Cosmos resource ID validation regex breakdown:
@@ -318,6 +319,15 @@ def GetHeaders(  # pylint: disable=too-many-statements,too-many-branches
     if options.get("correlatedActivityId"):
         headers[http_constants.HttpHeaders.CorrelatedActivityId] = options["correlatedActivityId"]
 
+    if resource_type == "docs" and verb != "get":
+        if "responsePayloadOnWriteDisabled" in options:
+            responsePayloadOnWriteDisabled = options["responsePayloadOnWriteDisabled"]
+        else:
+            responsePayloadOnWriteDisabled = cosmos_client_connection.connection_policy.ResponsePayloadOnWriteDisabled
+
+        if responsePayloadOnWriteDisabled:
+            headers[http_constants.HttpHeaders.Prefer] = "return=minimal"
+
     # If it is an operation at the container level, verify the rid of the container to see if the cache needs to be
     # refreshed.
     if resource_type != 'dbs' and options.get("containerRID"):

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

@@ -190,6 +190,7 @@ async def create_item(
         etag: Optional[str] = None,
         match_condition: Optional[MatchConditions] = None,
         priority: Optional[Literal["High", "Low"]] = None,
+        no_response: Optional[bool] = None,
         **kwargs: Any
     ) -> Dict[str, Any]:
         """Create an item in the container.
@@ -216,7 +217,10 @@ async def create_item(
             request. Once the user has reached their provisioned throughput, low priority requests are throttled
             before high priority requests start getting throttled. Feature must first be enabled at the account level.
         :raises ~azure.cosmos.exceptions.CosmosHttpResponseError: Item with the given ID already exists.
-        :returns: A dict representing the new item.
+        :keyword bool no_response: Indicates whether service should be instructed to skip
+            sending response payloads. When not specified explicitly here, the default value will be determined from 
+            client-level options.  
+        :returns: A dict representing the new item. The dict will be empty if `no_response` is specified.
         :rtype: Dict[str, Any]
         """
         if pre_trigger_include is not None:
@@ -233,6 +237,8 @@ async def create_item(
             kwargs['etag'] = etag
         if match_condition is not None:
             kwargs['match_condition'] = match_condition
+        if no_response is not None:
+            kwargs['no_response'] = no_response
         request_options = _build_options(kwargs)
         request_options["disableAutomaticIdGeneration"] = not enable_automatic_id_generation
         if indexing_directive is not None:
@@ -243,7 +249,7 @@ async def create_item(
         result = await self.client_connection.CreateItem(
             database_or_container_link=self.container_link, document=body, options=request_options, **kwargs
         )
-        return result
+        return result or {}
 
     @distributed_trace_async
     async def read_item(
@@ -552,6 +558,7 @@ async def upsert_item(
         etag: Optional[str] = None,
         match_condition: Optional[MatchConditions] = None,
         priority: Optional[Literal["High", "Low"]] = None,
+        no_response: Optional[bool] = None,
         **kwargs: Any
     ) -> Dict[str, Any]:
         """Insert or update the specified item.
@@ -574,7 +581,11 @@ async def upsert_item(
             request. Once the user has reached their provisioned throughput, low priority requests are throttled
             before high priority requests start getting throttled. Feature must first be enabled at the account level.
         :raises ~azure.cosmos.exceptions.CosmosHttpResponseError: The given item could not be upserted.
-        :returns: A dict representing the upserted item.
+        :keyword bool no_response: Indicates whether service should be instructed to skip
+            sending response payloads. When not specified explicitly here, the default value will be determined from 
+            client-level options.  
+        :returns: A dict representing the item after the upsert operation went through. The dict will be empty if 
+            `no_response` is specified.
         :rtype: Dict[str, Any]
         """
         if pre_trigger_include is not None:
@@ -591,6 +602,8 @@ async def upsert_item(
             kwargs['etag'] = etag
         if match_condition is not None:
             kwargs['match_condition'] = match_condition
+        if no_response is not None:
+            kwargs['no_response'] = no_response
         request_options = _build_options(kwargs)
         request_options["disableAutomaticIdGeneration"] = True
         if self.container_link in self.__get_client_container_caches():
@@ -602,7 +615,7 @@ async def upsert_item(
             options=request_options,
             **kwargs
         )
-        return result
+        return result or {}
 
     @distributed_trace_async
     async def replace_item(
@@ -617,6 +630,7 @@ async def replace_item(
         etag: Optional[str] = None,
         match_condition: Optional[MatchConditions] = None,
         priority: Optional[Literal["High", "Low"]] = None,
+        no_response: Optional[bool] = None,
         **kwargs: Any
     ) -> Dict[str, Any]:
         """Replaces the specified item if it exists in the container.
@@ -641,7 +655,11 @@ async def replace_item(
             before high priority requests start getting throttled. Feature must first be enabled at the account level.
         :raises ~azure.cosmos.exceptions.CosmosHttpResponseError: The replace failed or the item with
             given id does not exist.
-        :returns: A dict representing the item after replace went through.
+        :keyword bool no_response: Indicates whether service should be instructed to skip
+            sending response payloads. When not specified explicitly here, the default value will be determined from 
+            client-level options.  
+        :returns: A dict representing the item after replace went through. The dict will be empty if `no_response` 
+            is specified.
         :rtype: Dict[str, Any]
         """
         item_link = self._get_document_link(item)
@@ -659,6 +677,8 @@ async def replace_item(
             kwargs['etag'] = etag
         if match_condition is not None:
             kwargs['match_condition'] = match_condition
+        if no_response is not None:
+            kwargs['no_response'] = no_response
         request_options = _build_options(kwargs)
         request_options["disableAutomaticIdGeneration"] = True
         if self.container_link in self.__get_client_container_caches():
@@ -667,7 +687,7 @@ async def replace_item(
         result = await self.client_connection.ReplaceItem(
             document_link=item_link, new_document=body, options=request_options, **kwargs
         )
-        return result
+        return result or {}
 
     @distributed_trace_async
     async def patch_item(
@@ -683,6 +703,7 @@ async def patch_item(
         etag: Optional[str] = None,
         match_condition: Optional[MatchConditions] = None,
         priority: Optional[Literal["High", "Low"]] = None,
+        no_response: Optional[bool] = None,
         **kwargs: Any
     ) -> Dict[str, Any]:
         """ Patches the specified item with the provided operations if it
@@ -707,7 +728,11 @@ async def patch_item(
         :keyword Literal["High", "Low"] priority: Priority based execution allows users to set a priority for each
             request. Once the user has reached their provisioned throughput, low priority requests are throttled
             before high priority requests start getting throttled. Feature must first be enabled at the account level.
-        :returns: A dict representing the item after the patch operations went through.
+        :keyword bool no_response: Indicates whether service should be instructed to skip
+            sending response payloads. When not specified explicitly here, the default value will be determined from 
+            client-level options.  
+        :returns: A dict representing the item after the patch operation went through. The dict will be empty if 
+            `no_response` is specified.
         :raises ~azure.cosmos.exceptions.CosmosHttpResponseError: The patch operations failed or the item with
             given id does not exist.
         :rtype: dict[str, Any]
@@ -724,6 +749,8 @@ async def patch_item(
             kwargs['etag'] = etag
         if match_condition is not None:
             kwargs['match_condition'] = match_condition
+        if no_response is not None:
+            kwargs['no_response'] = no_response
         request_options = _build_options(kwargs)
         request_options["disableAutomaticIdGeneration"] = True
         request_options["partitionKey"] = await self._set_partition_key(partition_key)
@@ -733,8 +760,9 @@ async def patch_item(
             request_options["containerRID"] = self.__get_client_container_caches()[self.container_link]["_rid"]
 
         item_link = self._get_document_link(item)
-        return await self.client_connection.PatchItem(
+        result = await self.client_connection.PatchItem(
             document_link=item_link, operations=patch_operations, options=request_options, **kwargs)
+        return result or {}
 
     @distributed_trace_async
     async def delete_item(

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

@@ -113,6 +113,7 @@ def _build_connection_policy(kwargs: Dict[str, Any]) -> ConnectionPolicy:
             retry_backoff_factor=kwargs.pop('retry_backoff_factor', 0.8),
         )
     policy.ConnectionRetryConfiguration = connection_retry
+    policy.ResponsePayloadOnWriteDisabled = kwargs.pop('no_response_on_write', False)
     return policy
 
 
@@ -156,6 +157,8 @@ class CosmosClient:  # pylint: disable=client-accepts-api-version-keyword
         Must be used along with a logger to work.
     :keyword ~logging.Logger logger: Logger to be used for collecting request diagnostics. Can be passed in at client
         level (to log all requests) or at a single request level. Requests will be logged at INFO level.
+    :keyword bool no_response_on_write: Indicates whether service should be instructed to skip sending 
+        response payloads for write operations on items by default unless specified differently per operation.
 
     .. admonition:: Example: