AddSupportForSdkSupportedCapabilities

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/HttpConstants.java

@@ -4,7 +4,9 @@
 package com.azure.cosmos.implementation;
 
 import com.azure.core.util.CoreUtils;
+import com.azure.cosmos.implementation.apachecommons.lang.StringUtils;
 
+import java.util.Arrays;
 import java.util.concurrent.atomic.AtomicReference;
 
 /**
@@ -274,6 +276,9 @@ public static class HttpHeaders {
         // Client Encryption Headers
         public static final String IS_CLIENT_ENCRYPTED_HEADER = "x-ms-cosmos-is-client-encrypted";
         public static final String INTENDED_COLLECTION_RID_HEADER = "x-ms-cosmos-intended-collection-rid";
+
+        // SDK supported capacities headers
+        public static final String SDK_SUPPORTED_CAPABILITIES = "x-ms-cosmos-sdk-supported-capabilities";
     }
 
     public static class A_IMHeaderValues {
@@ -394,4 +399,15 @@ public static class HeaderValues {
         public static final String PREFER_RETURN_MINIMAL = "return=minimal";
         public static final String IF_NONE_MATCH_ALL = "*";
     }
+
+    public static class SDKSupportedCapabilities {
+        private static final long None = 0; // 0
+        private static final long PartitionMerge = 1; // 1 << 0
+
+        public static final long SUPPRTED_CAPABILITIES;
+        static {
+            // TODO: check whether the current version is higher than v2018_12_31
+            SUPPRTED_CAPABILITIES = PartitionMerge;
+        }
+    }
 }

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxGatewayStoreModel.java

@@ -78,6 +78,8 @@ public RxGatewayStoreModel(
                 "no-cache");
         this.defaultHeaders.put(HttpConstants.HttpHeaders.VERSION,
                 HttpConstants.Versions.CURRENT_VERSION);
+        this.defaultHeaders.put(HttpConstants.HttpHeaders.SDK_SUPPORTED_CAPABILITIES,
+                String.valueOf(HttpConstants.SDKSupportedCapabilities.SUPPRTED_CAPABILITIES));
 
         if (apiType != null){
             this.defaultHeaders.put(HttpConstants.HttpHeaders.API_TYPE, apiType.toString());

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/GatewayAddressCache.java

@@ -140,6 +140,9 @@ public GatewayAddressCache(
 
         // Set requested API version header for version enforcement.
         defaultRequestHeaders.put(HttpConstants.HttpHeaders.VERSION, HttpConstants.Versions.CURRENT_VERSION);
+        defaultRequestHeaders.put(
+                HttpConstants.HttpHeaders.SDK_SUPPORTED_CAPABILITIES,
+                String.valueOf(HttpConstants.SDKSupportedCapabilities.SUPPRTED_CAPABILITIES));
 
         this.serverPartitionAddressToPkRangeIdMap = new ConcurrentHashMap<>();
         this.tcpConnectionEndpointRediscoveryEnabled = tcpConnectionEndpointRediscoveryEnabled;

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/HttpTransportClient.java

@@ -86,6 +86,9 @@ public HttpTransportClient(Configs configs, ConnectionPolicy connectionPolicy, U
         // Set requested API version header for version enforcement.
         this.defaultHeaders.put(HttpConstants.HttpHeaders.VERSION, HttpConstants.Versions.CURRENT_VERSION);
         this.defaultHeaders.put(HttpConstants.HttpHeaders.CACHE_CONTROL, HttpConstants.HeaderValues.NO_CACHE);
+        this.defaultHeaders.put(
+                HttpConstants.HttpHeaders.SDK_SUPPORTED_CAPABILITIES,
+                String.valueOf(HttpConstants.SDKSupportedCapabilities.SUPPRTED_CAPABILITIES));
 
         if (userAgent == null) {
             userAgent = new UserAgentContainer();

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdConstants.java

@@ -3,6 +3,7 @@
 
 package com.azure.cosmos.implementation.directconnectivity.rntbd;
 
+import com.azure.cosmos.implementation.HttpConstants;
 import com.azure.cosmos.implementation.OperationType;
 import com.azure.cosmos.implementation.ResourceType;
 import com.azure.cosmos.implementation.guava25.collect.ImmutableMap;
@@ -585,7 +586,8 @@ public enum RntbdRequestHeader implements RntbdHeader {
         PopulateIndexMetrics((short) 0x00A9, RntbdTokenType.Byte, false),
         IsClientEncrypted((short) 0x0087, RntbdTokenType.Byte, false),
         IntendedCollectionRid((short) 0x009D, RntbdTokenType.String, false),
-        CorrelatedActivityId((short) 0x00B0, RntbdTokenType.Guid, false);
+        CorrelatedActivityId((short) 0x00B0, RntbdTokenType.Guid, false),
+        SDKSupportedCapabilities((short) 0x00A2, RntbdTokenType.ULong, false);
 
         public static final ImmutableMap<Short, RntbdRequestHeader> map;
         public static final ImmutableSet<RntbdRequestHeader> set = Sets.immutableEnumSet(EnumSet.allOf(RntbdRequestHeader.class));

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdRequestHeaders.java

@@ -119,6 +119,7 @@ final class RntbdRequestHeaders extends RntbdTokenStream<RntbdRequestHeader> {
         this.addIsClientEncrypted(headers);
         this.addIntendedCollectionRid(headers);
         this.addCorrelatedActivityId(headers);
+        this.addSDKSupportedCapabilities(headers);
 
         // Normal headers (Strings, Ints, Longs, etc.)
 
@@ -161,6 +162,7 @@ final class RntbdRequestHeaders extends RntbdTokenStream<RntbdRequestHeader> {
         this.fillTokenFromHeader(headers, this::shouldBatchContinueOnError, HttpHeaders.SHOULD_BATCH_CONTINUE_ON_ERROR);
         this.fillTokenFromHeader(headers, this::isBatchOrdered, HttpHeaders.IS_BATCH_ORDERED);
         this.fillTokenFromHeader(headers, this::getCorrelatedActivityId, HttpHeaders.CORRELATED_ACTIVITY_ID);
+        this.fillTokenFromHeader(headers, this::getSDKSupportedCapabilities, HttpHeaders.SDK_SUPPORTED_CAPABILITIES);
 
         // Will be null in case of direct, which is fine - BE will use the value slice the connection context this.
         // When this is used in Gateway, the header value will be populated with the proxied HTTP request's header,
@@ -603,6 +605,10 @@ private RntbdToken isBatchOrdered() {
         return this.get(RntbdRequestHeader.IsBatchOrdered);
     }
 
+    private RntbdToken getSDKSupportedCapabilities() {
+        return this.get(RntbdRequestHeader.SDKSupportedCapabilities);
+    }
+
     private void addAimHeader(final Map<String, String> headers) {
 
         final String value = headers.get(HttpHeaders.A_IM);
@@ -1243,6 +1249,13 @@ private void addReturnPreference(final Map<String, String> headers) {
         }
     }
 
+    private void addSDKSupportedCapabilities(final Map<String, String> headers) {
+        final String value = headers.get(HttpHeaders.SDK_SUPPORTED_CAPABILITIES);
+        if (StringUtils.isNotEmpty(value)) {
+            this.getSDKSupportedCapabilities().setValue(Long.valueOf(value));
+        }
+    }
+
     private void fillTokenFromHeader(final Map<String, String> headers, final Supplier<RntbdToken> supplier, final String name) {
 
         final String value = headers.get(name);

sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/directconnectivity/HttpTransportClientTest.java

@@ -151,7 +151,8 @@ public void validateDefaultHeaders() {
         assertThat(httpRequest.headers().value(HttpConstants.HttpHeaders.CACHE_CONTROL)).isEqualTo("no-cache");
         assertThat(httpRequest.headers().value(HttpConstants.HttpHeaders.ACCEPT)).isEqualTo("application/json");
         assertThat(httpRequest.headers().value(HttpConstants.HttpHeaders.VERSION)).isEqualTo(HttpConstants.Versions.CURRENT_VERSION);
-
+        assertThat(httpRequest.headers().value(HttpConstants.HttpHeaders.SDK_SUPPORTED_CAPABILITIES))
+            .isEqualTo(String.valueOf(HttpConstants.SDKSupportedCapabilities.SUPPRTED_CAPABILITIES));
     }
 
     @DataProvider(name = "fromMockedHttpResponseToExpectedDocumentClientException")