diff --git a/codegen/src/main/java/software/amazon/awssdk/codegen/customization/processors/DefaultCustomizationProcessor.java b/codegen/src/main/java/software/amazon/awssdk/codegen/customization/processors/DefaultCustomizationProcessor.java index 225cd61d2dea..1511bf6be282 100644 --- a/codegen/src/main/java/software/amazon/awssdk/codegen/customization/processors/DefaultCustomizationProcessor.java +++ b/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() ); } } diff --git a/codegen/src/main/java/software/amazon/awssdk/codegen/customization/processors/LongPollingOperationProcessor.java b/codegen/src/main/java/software/amazon/awssdk/codegen/customization/processors/LongPollingOperationProcessor.java new file mode 100644 index 000000000000..e5debd935b18 --- /dev/null +++ b/codegen/src/main/java/software/amazon/awssdk/codegen/customization/processors/LongPollingOperationProcessor.java @@ -0,0 +1,94 @@ +/* + * 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> SERVICE_ID_TO_OPERATIONS_MAP; + + static { + Map> 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> serviceIdToOperations; + + public LongPollingOperationProcessor() { + this(SERVICE_ID_TO_OPERATIONS_MAP); + } + + @SdkTestInternalApi + LongPollingOperationProcessor(Map> 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 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); + } + } + } +} diff --git a/codegen/src/main/java/software/amazon/awssdk/codegen/model/intermediate/OperationModel.java b/codegen/src/main/java/software/amazon/awssdk/codegen/model/intermediate/OperationModel.java index 6b192644da1d..e69c6f2ab883 100644 --- a/codegen/src/main/java/software/amazon/awssdk/codegen/model/intermediate/OperationModel.java +++ b/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; } } diff --git a/codegen/src/main/java/software/amazon/awssdk/codegen/poet/client/specs/JsonProtocolSpec.java b/codegen/src/main/java/software/amazon/awssdk/codegen/poet/client/specs/JsonProtocolSpec.java index b9b69d9bd8a0..894055c2db7c 100644 --- a/codegen/src/main/java/software/amazon/awssdk/codegen/poet/client/specs/JsonProtocolSpec.java +++ b/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") diff --git a/codegen/src/main/java/software/amazon/awssdk/codegen/poet/client/traits/LongPollTrait.java b/codegen/src/main/java/software/amazon/awssdk/codegen/poet/client/traits/LongPollTrait.java new file mode 100644 index 000000000000..e12adb13e6cb --- /dev/null +++ b/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(""); + } + +} diff --git a/codegen/src/test/java/software/amazon/awssdk/codegen/customization/processors/LongPollingOperationProcessTest.java b/codegen/src/test/java/software/amazon/awssdk/codegen/customization/processors/LongPollingOperationProcessTest.java new file mode 100644 index 000000000000..3951be476d81 --- /dev/null +++ b/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> 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> 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 nonJsonProtocols() { + return Stream.of(Protocol.values()).filter(p -> p != Protocol.AWS_JSON); + } +}