Enabling option to disable response payload on writes

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,11 @@ 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 item after replace went through or if response payload on write is disabled
+            None.
         :rtype: Dict[str, Any]
         """
         if pre_trigger_include is not None:
@@ -233,6 +238,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 +250,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 self.__unwrap_response(result)
 
     @distributed_trace_async
     async def read_item(
@@ -552,6 +559,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 +582,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 replace went through or if response payload on write is disabled
+            None.
         :rtype: Dict[str, Any]
         """
         if pre_trigger_include is not None:
@@ -591,6 +603,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 +616,13 @@ async def upsert_item(
             options=request_options,
             **kwargs
         )
-        return result
+        return self.__unwrap_response(result)
+
+    def __unwrap_response(self, response: Optional[Dict[str, Any]]) -> Dict[str, Any]:
+        if response is None:
+            return {}
+
+        return response
 
     @distributed_trace_async
     async def replace_item(
@@ -617,6 +637,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 +662,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 or if response payload on write is disabled
+            None.
         :rtype: Dict[str, Any]
         """
         item_link = self._get_document_link(item)
@@ -659,6 +684,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 +694,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 self.__unwrap_response(result)
 
     @distributed_trace_async
     async def patch_item(
@@ -683,6 +710,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 +735,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 replace went through or if response payload on write is disabled
+            None.
         :raises ~azure.cosmos.exceptions.CosmosHttpResponseError: The patch operations failed or the item with
             given id does not exist.
         :rtype: dict[str, Any]
@@ -724,6 +756,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 +767,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 self.__unwrap_response(result)
 
     @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('response_payload_on_write_disabled', 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 response_payload_on_write_disabled: Indicates whether service should be instructed to skip sending 
+        response payloads for write operations by default unless specified differently per operation.
 
     .. admonition:: Example: