Set longPoll ops for SQS,SWF,SFN

codegen/src/main/java/software/amazon/awssdk/codegen/customization/processors/DefaultCustomizationProcessor.java

@@ -42,7 +42,8 @@ public static CodegenCustomizationProcessor getProcessorFor(
                 new S3RemoveBucketFromUriProcessor(),
                 new S3ControlRemoveAccountIdHostPrefixProcessor(),
                 new ExplicitStringPayloadQueryProtocolProcessor(),
-                new LowercaseShapeValidatorProcessor()
+                new LowercaseShapeValidatorProcessor(),
+                new LongPollingOperationProcessor()
                 );
     }
 }

codegen/src/main/java/software/amazon/awssdk/codegen/customization/processors/LongPollingOperationProcessor.java

@@ -0,0 +1,93 @@
+/*
+ * 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.codegen.customization.processors;
+
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import software.amazon.awssdk.annotations.SdkTestInternalApi;
+import software.amazon.awssdk.codegen.customization.CodegenCustomizationProcessor;
+import software.amazon.awssdk.codegen.model.intermediate.IntermediateModel;
+import software.amazon.awssdk.codegen.model.intermediate.OperationModel;
+import software.amazon.awssdk.codegen.model.intermediate.Protocol;
+import software.amazon.awssdk.codegen.model.service.ServiceModel;
+
+// TODO: Remove this when the long polling trait is formalized as a c2j trait.
+/**
+ * Marks specific service operations as having the long polling trait.
+ */
+public class LongPollingOperationProcessor implements CodegenCustomizationProcessor {
+    private static final Logger log = LoggerFactory.getLogger(LongPollingOperationProcessor.class);
+
+    // Note: static mapping instead of exposed via CustomizationConfig to avoid exposing it for wider use unless necessary.
+    private static final Map<String, List<String>> SERVICE_ID_TO_OPERATIONS_MAP;
+
+    static {
+        Map<String, List<String>> serviceIdToOperationsMap = new HashMap<>();
+
+        serviceIdToOperationsMap.put("SQS", Collections.singletonList("ReceiveMessage"));
+        serviceIdToOperationsMap.put("SFN", Collections.singletonList("GetActivityTask"));
+        serviceIdToOperationsMap.put("SWF", Collections.unmodifiableList(Arrays.asList("PollForActivityTask",
+                                                                                       "PollForDecisionTask")));
+
+        SERVICE_ID_TO_OPERATIONS_MAP = Collections.unmodifiableMap(serviceIdToOperationsMap);
+    }
+
+    private final Map<String, List<String>> serviceIdToOperations;
+
+    public LongPollingOperationProcessor() {
+        this(SERVICE_ID_TO_OPERATIONS_MAP);
+    }
+
+    @SdkTestInternalApi
+    LongPollingOperationProcessor(Map<String, List<String>> serviceIdToOperations) {
+        this.serviceIdToOperations = serviceIdToOperations;
+    }
+
+    @Override
+    public void preprocess(ServiceModel serviceModel) {
+        // no-op
+    }
+
+    @Override
+    public void postprocess(IntermediateModel intermediateModel) {
+        String serviceId = intermediateModel.getMetadata().getServiceId();
+
+        if (!serviceIdToOperations.containsKey(serviceId)) {
+            return;
+        }
+
+        if (intermediateModel.getMetadata().getProtocol() != Protocol.AWS_JSON) {
+            throw new IllegalArgumentException("Currently only AWS-JSON services can use the longPoll trait");
+        }
+
+        List<String> longPollingOperations = serviceIdToOperations.getOrDefault(serviceId, Collections.emptyList());
+
+        for (String longPollingOperation : longPollingOperations) {
+            OperationModel opModel = intermediateModel.getOperation(longPollingOperation);
+            if (opModel != null) {
+                log.info("Setting the longPoll trait for {}#{}", serviceId, longPollingOperation);
+                opModel.setLongPolling(true);
+            } else {
+               throw new RuntimeException("Operation " + longPollingOperation + " not found for service " + serviceId);
+            }
+        }
+    }
+}

codegen/src/main/java/software/amazon/awssdk/codegen/model/intermediate/OperationModel.java

@@ -90,6 +90,8 @@ public class OperationModel extends DocumentationModel {
 
     private boolean unsignedPayload;
 
+    private boolean longPolling;
+
     public String getOperationName() {
         return operationName;
     }
@@ -381,6 +383,14 @@ public void setUnsignedPayload(boolean unsignedPayload) {
         this.unsignedPayload = unsignedPayload;
     }
 
+    public boolean isLongPolling() {
+        return longPolling;
+    }
+
+    public void setLongPolling(boolean longPolling) {
+        this.longPolling = longPolling;
+    }
+
     @Override
     public boolean equals(Object o) {
         if (o == null || getClass() != o.getClass()) {
@@ -395,7 +405,8 @@ public boolean equals(Object o) {
                && hasStringMemberAsPayload == that.hasStringMemberAsPayload && isAuthenticated == that.isAuthenticated
                && isPaginated == that.isPaginated && endpointOperation == that.endpointOperation
                && endpointCacheRequired == that.endpointCacheRequired && httpChecksumRequired == that.httpChecksumRequired
-               && unsignedPayload == that.unsignedPayload && Objects.equals(operationName, that.operationName)
+               && unsignedPayload == that.unsignedPayload && longPolling == that.longPolling
+               && Objects.equals(operationName, that.operationName)
                && Objects.equals(serviceProtocol, that.serviceProtocol)
                && Objects.equals(deprecatedMessage, that.deprecatedMessage) && Objects.equals(input, that.input)
                && Objects.equals(returnType, that.returnType) && Objects.equals(exceptions, that.exceptions)
@@ -437,6 +448,7 @@ public int hashCode() {
         result = 31 * result + Objects.hashCode(staticContextParams);
         result = 31 * result + Objects.hashCode(operationContextParams);
         result = 31 * result + Boolean.hashCode(unsignedPayload);
+        result = 31 * result + Boolean.hashCode(longPolling);
         return result;
     }
 }

codegen/src/main/java/software/amazon/awssdk/codegen/poet/client/specs/JsonProtocolSpec.java

@@ -43,6 +43,7 @@
 import software.amazon.awssdk.codegen.poet.PoetExtension;
 import software.amazon.awssdk.codegen.poet.client.traits.HttpChecksumRequiredTrait;
 import software.amazon.awssdk.codegen.poet.client.traits.HttpChecksumTrait;
+import software.amazon.awssdk.codegen.poet.client.traits.LongPollTrait;
 import software.amazon.awssdk.codegen.poet.client.traits.RequestCompressionTrait;
 import software.amazon.awssdk.codegen.poet.eventstream.EventStreamUtils;
 import software.amazon.awssdk.codegen.poet.model.EventStreamSpecHelper;
@@ -219,6 +220,7 @@ public CodeBlock executionHandler(OperationModel opModel) {
                      .add(hostPrefixExpression(opModel))
                      .add(discoveredEndpoint(opModel))
                      .add(credentialType(opModel, model))
+                     .add(LongPollTrait.executionParamSetter(opModel))
                      .add(".withRequestConfiguration(clientConfiguration)")
                      .add(".withInput($L)\n", opModel.getInput().getVariableName())
                      .add(".withMetricCollector(apiCallMetricCollector)")
@@ -290,6 +292,7 @@ public CodeBlock asyncExecutionHandler(IntermediateModel intermediateModel, Oper
                .add(".withMarshaller($L)\n", asyncMarshaller(model, opModel, marshaller, protocolFactory))
                .add(asyncRequestBody(opModel))
                .add(fullDuplex(opModel))
+               .add(LongPollTrait.executionParamSetter(opModel))
                .add(hasInitialRequestEvent(opModel, isRestJson))
                .add(".withResponseHandler($L)\n", responseHandlerName(opModel, isRestJson))
                .add(".withErrorResponseHandler(errorResponseHandler)\n")

codegen/src/main/java/software/amazon/awssdk/codegen/poet/client/traits/LongPollTrait.java

@@ -0,0 +1,35 @@
+/*
+ * 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.codegen.poet.client.traits;
+
+import com.squareup.javapoet.CodeBlock;
+import software.amazon.awssdk.codegen.model.intermediate.OperationModel;
+
+/**
+ * Helper methods for working with the long poll trait for operations.
+ */
+public final class LongPollTrait {
+    private LongPollTrait() {
+    }
+
+    public static CodeBlock executionParamSetter(OperationModel operationModel) {
+        if (operationModel.isLongPolling()) {
+            return CodeBlock.of(".withLongPolling(true)");
+        }
+        return CodeBlock.of("");
+    }
+
+}

codegen/src/test/java/software/amazon/awssdk/codegen/customization/processors/LongPollingOperationProcessTest.java

@@ -0,0 +1,70 @@
+/*
+ * 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.codegen.customization.processors;
+
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.stream.Stream;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.MethodSource;
+import software.amazon.awssdk.codegen.C2jModels;
+import software.amazon.awssdk.codegen.IntermediateModelBuilder;
+import software.amazon.awssdk.codegen.model.intermediate.IntermediateModel;
+import software.amazon.awssdk.codegen.model.intermediate.Protocol;
+import software.amazon.awssdk.codegen.model.service.ServiceMetadata;
+import software.amazon.awssdk.codegen.poet.ClientTestModels;
+
+public class LongPollingOperationProcessTest {
+
+    @ParameterizedTest
+    @MethodSource("nonJsonProtocols")
+    void postprocess_serviceInMap_serviceNotJson_throws(Protocol protocol) {
+        C2jModels c2jModels = ClientTestModels.awsJsonServiceC2jModels();
+
+        ServiceMetadata metadata = c2jModels.serviceModel().getMetadata();
+        metadata.setProtocols(Collections.singletonList(protocol.getValue()));
+
+        IntermediateModel intermediateModel = new IntermediateModelBuilder(c2jModels).build();
+
+        Map<String, List<String>> serviceToOperations = new HashMap<>();
+        serviceToOperations.put(metadata.getServiceId(), Collections.emptyList());
+        LongPollingOperationProcessor processor = new LongPollingOperationProcessor(serviceToOperations);
+
+        assertThatThrownBy(() -> processor.postprocess(intermediateModel))
+            .hasMessage("Currently only AWS-JSON services can use the longPoll trait");
+    }
+
+    @Test
+    void postprocess_operationNotFound_throws() {
+        IntermediateModel intermediateModel = ClientTestModels.awsJsonServiceModels();
+
+        Map<String, List<String>> serviceToOperations = new HashMap<>();
+        serviceToOperations.put(intermediateModel.getMetadata().getServiceId(), Collections.singletonList("SomeOperation"));
+        LongPollingOperationProcessor processor = new LongPollingOperationProcessor(serviceToOperations);
+
+        assertThatThrownBy(() -> processor.postprocess(intermediateModel))
+            .hasMessage("Operation SomeOperation not found for service Json Service");
+    }
+
+    private static Stream<Protocol> nonJsonProtocols() {
+        return Stream.of(Protocol.values()).filter(p -> p != Protocol.AWS_JSON);
+    }
+}