progress

sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdTokenStreamTests.java

@@ -0,0 +1,34 @@
+package com.azure.cosmos.implementation.directconnectivity.rntbd;
+
+import io.netty.buffer.ByteBuf;
+import io.netty.buffer.Unpooled;
+import org.testng.annotations.Test;
+
+import java.util.EnumSet;
+import java.util.Map;
+
+import static com.azure.cosmos.implementation.directconnectivity.WFConstants.BackendHeaders.EFFECTIVE_PARTITION_KEY;
+
+public class RntbdTokenStreamTests {
+    // Created this test for thin client testing
+    @Test(groups = { "unit" })
+    public void testThinClientSpecialCasing() {
+        RntbdContextRequest.Headers headers = new RntbdContextRequest.Headers(Unpooled.EMPTY_BUFFER);
+
+        RntbdTokenStream<RntbdConstants.RntbdRequestHeader> rntbdTokenStream =
+            new TestRntbdTokenStream(
+                RntbdConstants.RntbdRequestHeader.set,
+                RntbdConstants.RntbdRequestHeader.map,
+                Unpooled.EMPTY_BUFFER,
+                RntbdConstants.RntbdRequestHeader.class);
+
+        final ByteBuf out = Unpooled.buffer();
+        headers.encode(out);
+    }
+
+    final class TestRntbdTokenStream extends RntbdTokenStream<RntbdConstants.RntbdRequestHeader> {
+        TestRntbdTokenStream(EnumSet<RntbdConstants.RntbdRequestHeader> headers, Map<Short, RntbdConstants.RntbdRequestHeader> ids, ByteBuf in, Class<RntbdConstants.RntbdRequestHeader> classType) {
+            super(headers, ids, in, classType);
+        }
+    }
+}

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

@@ -283,6 +283,8 @@ public static class HttpHeaders {
         // Priority Level for throttling
         public static final String PRIORITY_LEVEL = "x-ms-cosmos-priority-level";
 
+        public static final String GLOBAL_DATABASE_ACCOUNT_NAME = "GlobalDatabaseAccountName";
+
         // Thinclient headers
         public static final String THINCLIENT_PROXY_OPERATION_TYPE = "x-ms-thinclient-proxy-operation-type";
         public static final String THINCLIENT_PROXY_RESOURCE_TYPE = "x-ms-thinclient-proxy-resource-type";

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

@@ -591,7 +591,8 @@ public enum RntbdRequestHeader implements RntbdHeader {
         CorrelatedActivityId((short) 0x00B0, RntbdTokenType.Guid, false),
         SDKSupportedCapabilities((short) 0x00A2, RntbdTokenType.ULong, false),
         ChangeFeedWireFormatVersion((short) 0x00B2, RntbdTokenType.String, false),
-        PriorityLevel((short) 0x00BF, RntbdTokenType.Byte, false);
+        PriorityLevel((short) 0x00BF, RntbdTokenType.Byte, false),
+        GlobalDatabaseAccountName((short) 0x00CE, RntbdTokenType.Byte, false);
 
         public static final Map<Short, RntbdRequestHeader> map;
         public static final EnumSet<RntbdRequestHeader> set = EnumSet.allOf(RntbdRequestHeader.class);

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

@@ -101,7 +101,7 @@ public String toString() {
         }
     }
 
-    private static final class Headers extends RntbdTokenStream<RntbdContextRequestHeader> {
+    static final class Headers extends RntbdTokenStream<RntbdContextRequestHeader> {
 
         private static final byte[] ClientVersion = Versions.CURRENT_VERSION.getBytes(StandardCharsets.UTF_8);
 
@@ -121,7 +121,7 @@ private static final class Headers extends RntbdTokenStream<RntbdContextRequestH
             this.protocolVersion.setValue(CURRENT_PROTOCOL_VERSION);
         }
 
-        private Headers(ByteBuf in) {
+        Headers(ByteBuf in) {
 
             super(RntbdContextRequestHeader.set, RntbdContextRequestHeader.map, in, RntbdContextRequestHeader.class);
 

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

@@ -74,7 +74,6 @@ final class RntbdRequestHeaders extends RntbdTokenStream<RntbdRequestHeader> {
         final Map<String, String> headers = request.getHeaders();
 
         // Special-case headers
-
         this.addAimHeader(headers);
         this.addAllowScanOnQuery(headers);
         this.addBinaryIdIfPresent(headers);
@@ -124,6 +123,7 @@ final class RntbdRequestHeaders extends RntbdTokenStream<RntbdRequestHeader> {
         this.addSDKSupportedCapabilities(headers);
         this.addChangeFeedWireFormatVersion(headers);
         this.addPriorityLevel(headers);
+        this.addGlobalDatabaseAccountName(headers);
 
         // Normal headers (Strings, Ints, Longs, etc.)
 
@@ -283,6 +283,8 @@ private RntbdToken getCorrelatedActivityId() {
 
     private RntbdToken getPriorityLevel() { return this.get(RntbdRequestHeader.PriorityLevel); }
 
+    private RntbdToken getGlobalDatabaseAccountName() { return this.get(RntbdRequestHeader.GlobalDatabaseAccountName); }
+
     private RntbdToken getDatabaseName() {
         return this.get(RntbdRequestHeader.DatabaseName);
     }
@@ -776,6 +778,15 @@ private void addPriorityLevel(final Map<String, String> headers)
         }
     }
 
+    private void addGlobalDatabaseAccountName(final Map<String, String> headers)
+    {
+        final String value = headers.get(HttpHeaders.GLOBAL_DATABASE_ACCOUNT_NAME);
+
+        if (StringUtils.isNotEmpty(value)) {
+            this.getGlobalDatabaseAccountName().setValue(value);
+        }
+    }
+
     private void addDateHeader(final Map<String, String> headers) {
 
         // Since the HTTP date header is overridden by some proxies/http client libraries, we support an additional date

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

@@ -11,6 +11,8 @@
 import java.util.EnumSet;
 import java.util.Map;
 
+import static com.azure.cosmos.implementation.HttpConstants.HttpHeaders.GLOBAL_DATABASE_ACCOUNT_NAME;
+import static com.azure.cosmos.implementation.directconnectivity.WFConstants.BackendHeaders.EFFECTIVE_PARTITION_KEY;
 import static com.azure.cosmos.implementation.directconnectivity.rntbd.RntbdConstants.RntbdHeader;
 import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull;
 import static com.azure.cosmos.implementation.guava27.Strings.lenientFormat;
@@ -28,7 +30,7 @@ abstract class RntbdTokenStream<T extends Enum<T> & RntbdHeader> implements Refe
         checkNotNull(ids, "expected non-null ids");
         checkNotNull(in, "expected non-null in");
 
-        this.tokens = new EnumMap<T, RntbdToken>(classType);
+        this.tokens = new EnumMap<>(classType);
         headers.stream().forEach(h -> tokens.put(h, RntbdToken.create(h)));
         this.headers = ids;
         this.in = in;
@@ -89,6 +91,16 @@ static <T extends RntbdTokenStream<?>> T decode(final T stream) {
     }
 
     final void encode(final ByteBuf out) {
+        // TODO: special casing for thin client. need to revisit perf implications.
+        if (this.tokens.containsKey(EFFECTIVE_PARTITION_KEY)) {
+            this.tokens.get(EFFECTIVE_PARTITION_KEY).encode(out);
+            this.tokens.remove(EFFECTIVE_PARTITION_KEY);
+        }
+        if (this.tokens.containsKey(GLOBAL_DATABASE_ACCOUNT_NAME)) {
+            this.tokens.get(GLOBAL_DATABASE_ACCOUNT_NAME).encode(out);
+            this.tokens.remove(GLOBAL_DATABASE_ACCOUNT_NAME);
+        }
+
         for (final RntbdToken token : this.tokens.values()) {
             token.encode(out);
         }