Add merge support (#42924)

* Add merge support

* Updated CHANGELOG.md

* Move header tests to `test_headers.py`

* Add Partition Merge header's async test

Author: allenkim0129

sdk/cosmos/azure-cosmos/CHANGELOG.md

@@ -28,6 +28,7 @@ This version and all future versions will require Python 3.9+.
 #### Features Added
 * Added read_items API to provide an efficient method for retrieving multiple items in a single request. See [PR 42167](https://github.com/Azure/azure-sdk-for-python/pull/42167).
 * Added ability to replace a container's indexing policy if a vector embedding policy was present. See [PR 42810](https://github.com/Azure/azure-sdk-for-python/pull/42810).
+* Added merge support. See [PR 42924](https://github.com/Azure/azure-sdk-for-python/pull/42924).
 
 #### Bugs Fixed
 * Improved the resilience of Database Account Read metadata operation against short-lived network issues by increasing number of retries. See [PR 42525](https://github.com/Azure/azure-sdk-for-python/pull/42525).

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

@@ -147,6 +147,10 @@ def GetHeaders(  # pylint: disable=too-many-statements,too-many-branches
     headers = dict(default_headers)
     options = options or {}
 
+    # SDK supported capabilities header for partition merge support
+    headers[http_constants.HttpHeaders.SDKSupportedCapabilities] = \
+        http_constants.SDKSupportedCapabilities.PARTITION_MERGE
+
     # Generate a new activity ID for each request client side.
     headers[http_constants.HttpHeaders.ActivityId] = GenerateGuidId()
     if cosmos_client_connection.UseMultipleWriteLocations:

sdk/cosmos/azure-cosmos/azure/cosmos/http_constants.py

@@ -373,6 +373,11 @@ class _ErrorCodes:
     # Linux Error Codes
     LinuxConnectionReset = 131
 
+class SDKSupportedCapabilities:
+    """Constants of SDK supported capabilities.
+    """
+    NONE = '0'
+    PARTITION_MERGE = '1'
 
 class StatusCodes:
     """HTTP status codes returned by the REST operations

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

@@ -23,6 +23,12 @@ def request_raw_response_hook(response):
         assert (response.http_request.headers[http_constants.HttpHeaders.ThroughputBucket]
                 == str(request_throughput_bucket_number))
 
+def partition_merge_support_response_hook(raw_response):
+    header = raw_response.http_request.headers
+    assert http_constants.HttpHeaders.SDKSupportedCapabilities in header
+    assert header[http_constants.HttpHeaders.SDKSupportedCapabilities] == \
+           http_constants.SDKSupportedCapabilities.PARTITION_MERGE
+
 @pytest.mark.cosmosEmulator
 class TestHeaders(unittest.TestCase):
     database: DatabaseProxy = None
@@ -279,5 +285,10 @@ def test_container_read_item_negative_throughput_bucket(self):
 
     """
 
+    def test_partition_merge_support_header(self):
+        # This test only runs read API to verify if the header was set correctly, because all APIs are using the same
+        # base method to set the header(GetHeaders).
+        self.container.read(raw_response_hook=partition_merge_support_response_hook)
+
 if __name__ == "__main__":
     unittest.main()

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

@@ -12,6 +12,7 @@
 from azure.cosmos import http_constants
 from azure.cosmos.aio import CosmosClient, _retry_utility_async, DatabaseProxy
 from azure.cosmos.partition_key import PartitionKey
+from test_headers import partition_merge_support_response_hook
 
 client_throughput_bucket_number = 2
 request_throughput_bucket_number = 3
@@ -230,6 +231,10 @@ async def test_client_id(self):
         finally:
             cosmos_client_connection._CosmosClientConnection__Get = original_connection_get
 
+    async def test_partition_merge_support_header(self):
+        # This test only runs read API to verify if the header was set correctly, because all APIs are using the same
+        # base method to set the header(GetHeaders).
+        await self.container.read(raw_response_hook=partition_merge_support_response_hook)
 
 if __name__ == "__main__":
     unittest.main()