Add x-amzn-query-mode to http header when service has @awsQueryCompatible trait

core/protocols/aws-json-protocol/src/main/java/software/amazon/awssdk/protocols/json/BaseAwsJsonProtocolFactory.java

@@ -236,6 +236,7 @@ public final ProtocolMarshaller<SdkHttpFullRequest> createProtocolMarshaller(Ope
                                             .operationInfo(operationInfo)
                                             .sendExplicitNullForPayload(false)
                                             .protocolMetadata(protocolMetadata)
+                                            .hasAwsQueryCompatible(hasAwsQueryCompatible)
                                             .build();
     }
 

core/protocols/aws-json-protocol/src/main/java/software/amazon/awssdk/protocols/json/internal/marshall/JsonProtocolMarshaller.java

@@ -72,12 +72,14 @@ public class JsonProtocolMarshaller implements ProtocolMarshaller<SdkHttpFullReq
     private final JsonMarshallerContext marshallerContext;
     private final boolean hasEventStreamingInput;
     private final boolean hasEvent;
+    private final boolean hasAwsQueryCompatible;
 
     JsonProtocolMarshaller(URI endpoint,
                            StructuredJsonGenerator jsonGenerator,
                            String contentType,
                            OperationInfo operationInfo,
-                           AwsJsonProtocolMetadata protocolMetadata) {
+                           AwsJsonProtocolMetadata protocolMetadata,
+                           boolean hasAwsQueryCompatible) {
         this.endpoint = endpoint;
         this.jsonGenerator = jsonGenerator;
         this.contentType = contentType;
@@ -88,6 +90,7 @@ public class JsonProtocolMarshaller implements ProtocolMarshaller<SdkHttpFullReq
         this.hasEventStreamingInput = operationInfo.hasEventStreamingInput();
         this.hasEvent = operationInfo.hasEvent();
         this.request = fillBasicRequestParams(operationInfo);
+        this.hasAwsQueryCompatible = hasAwsQueryCompatible;
         this.marshallerContext = JsonMarshallerContext.builder()
                                                       .jsonGenerator(jsonGenerator)
                                                       .marshallerRegistry(MARSHALLER_REGISTRY)
@@ -292,6 +295,10 @@ private SdkHttpFullRequest finishMarshalling() {
             }
         }
 
+        if (hasAwsQueryCompatible) {
+            request.putHeader("x-amzn-query-mode", "true");
+        }
+
         return request.build();
     }
 

core/protocols/aws-json-protocol/src/main/java/software/amazon/awssdk/protocols/json/internal/marshall/JsonProtocolMarshallerBuilder.java

@@ -35,6 +35,7 @@ public final class JsonProtocolMarshallerBuilder {
     private OperationInfo operationInfo;
     private boolean sendExplicitNullForPayload;
     private AwsJsonProtocolMetadata protocolMetadata;
+    private boolean hasAwsQueryCompatible = false;
 
     private JsonProtocolMarshallerBuilder() {
     }
@@ -102,6 +103,14 @@ public JsonProtocolMarshallerBuilder protocolMetadata(AwsJsonProtocolMetadata pr
         return this;
     }
 
+    /**
+     * @param hasAwsQueryCompatible True if the service is AWS Query compatible, (has the @awsQueryCompatible trait)
+     */
+    public JsonProtocolMarshallerBuilder hasAwsQueryCompatible(boolean hasAwsQueryCompatible) {
+        this.hasAwsQueryCompatible = hasAwsQueryCompatible;
+        return this;
+    }
+
     /**
      * @return New instance of {@link ProtocolMarshaller}. If {@link #sendExplicitNullForPayload} is true then the marshaller
      * will be wrapped with {@link NullAsEmptyBodyProtocolRequestMarshaller}.
@@ -111,7 +120,8 @@ public ProtocolMarshaller<SdkHttpFullRequest> build() {
                                                                                                jsonGenerator,
                                                                                                contentType,
                                                                                                operationInfo,
-                                                                                               protocolMetadata);
+                                                                                               protocolMetadata,
+                                                                                               hasAwsQueryCompatible);
         return sendExplicitNullForPayload ? protocolMarshaller
                                           : new NullAsEmptyBodyProtocolRequestMarshaller(protocolMarshaller);
     }

core/protocols/aws-json-protocol/src/test/java/software/amazon/awssdk/protocols/json/AWSQueryModeTest.java

@@ -0,0 +1,114 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package software.amazon.awssdk.protocols.json;
+
+import java.net.URI;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import org.junit.jupiter.api.Test;
+import software.amazon.awssdk.core.ClientEndpointProvider;
+import software.amazon.awssdk.core.SdkField;
+import software.amazon.awssdk.core.SdkPojo;
+import software.amazon.awssdk.core.client.config.SdkClientConfiguration;
+import software.amazon.awssdk.core.client.config.SdkClientOption;
+import software.amazon.awssdk.http.SdkHttpFullRequest;
+import software.amazon.awssdk.http.SdkHttpMethod;
+import software.amazon.awssdk.protocols.core.OperationInfo;
+import software.amazon.awssdk.protocols.core.ProtocolMarshaller;
+import static org.assertj.core.api.Assertions.assertThat;
+
+public class AWSQueryModeTest {
+
+    private static final OperationInfo EMPTY_OPERATION_INFO = OperationInfo.builder()
+                                                                           .httpMethod(SdkHttpMethod.POST)
+                                                                           .hasImplicitPayloadMembers(true)
+                                                                           .build();
+
+    /*
+     * A simple test POJO to marshall
+     */
+    private static final class TestPojo implements SdkPojo {
+
+        private TestPojo() {}
+
+        @Override
+        public List<SdkField<?>> sdkFields() {
+            return Collections.emptyList();
+        }
+
+        @Override
+        public boolean equalsBySdkFields(Object other) {
+            if (!(other instanceof TestPojo)) {
+                return false;
+            }
+            return true;
+        }
+
+        @Override
+        public Map<String, SdkField<?>> sdkFieldNameToField() {
+            return Collections.emptyMap();
+        }
+
+    }
+
+    @Test
+    public void testMarshallWithAwsQueryCompatibleTrue() {
+        SdkClientConfiguration clientConfig =
+            SdkClientConfiguration.builder()
+                                  .option(SdkClientOption.CLIENT_ENDPOINT_PROVIDER,
+                                          ClientEndpointProvider.forEndpointOverride(URI.create("http://localhost")))
+                                  .build();
+        AwsJsonProtocolFactory factory =
+            AwsJsonProtocolFactory.builder()
+                                  .clientConfiguration(clientConfig)
+                                  .protocolVersion("1.1")
+                                  .protocol(AwsJsonProtocol.AWS_JSON)
+                                  .hasAwsQueryCompatible(true)
+                                  .build();
+
+        ProtocolMarshaller<SdkHttpFullRequest> marshaller = factory.createProtocolMarshaller(EMPTY_OPERATION_INFO);
+        SdkPojo testPojo = new TestPojo();
+
+        SdkHttpFullRequest result = marshaller.marshall(testPojo);
+
+        assertThat(result.headers()).containsKey("x-amzn-query-mode");
+        assertThat(result.headers().get("x-amzn-query-mode").get(0)).isEqualTo("true");
+    }
+
+    @Test
+    public void testMarshallWithNoAwsQueryCompatible() {
+        SdkClientConfiguration clientConfig =
+            SdkClientConfiguration.builder()
+                                  .option(SdkClientOption.CLIENT_ENDPOINT_PROVIDER,
+                                          ClientEndpointProvider.forEndpointOverride(URI.create("http://localhost")))
+                                  .build();
+        AwsJsonProtocolFactory factory =
+            AwsJsonProtocolFactory.builder()
+                                  .clientConfiguration(clientConfig)
+                                  .protocolVersion("1.1")
+                                  .protocol(AwsJsonProtocol.AWS_JSON)
+                                  .build();
+
+        ProtocolMarshaller<SdkHttpFullRequest> marshaller = factory.createProtocolMarshaller(EMPTY_OPERATION_INFO);
+        SdkPojo testPojo = new TestPojo();
+
+        SdkHttpFullRequest result = marshaller.marshall(testPojo);
+
+        assertThat(result.headers()).doesNotContainKey("x-amzn-query-mode");
+    }
+
+}