Add new AsyncResponseTransformer: toPublisher()

.changes/next-release/feature-AWSSDKforJavav2-0e09833.json

@@ -0,0 +1,6 @@
+{
+    "category": "AWS SDK for Java v2", 
+    "contributor": "", 
+    "type": "feature", 
+    "description": "Add new AsyncResponseTransformer: toPublisher(). This transformer makes it more convenient for users to directly consume a streaming-response payload (i.e., S3 GetObject) with async clients. This also allows users of Reactor/RxJava to more easily consume a streaming response (e.g., via Flux#from(Publisher))."
+}

codegen/src/main/java/software/amazon/awssdk/codegen/poet/client/AsyncClientClass.java

@@ -36,6 +36,7 @@
 import java.nio.ByteBuffer;
 import java.util.Collections;
 import java.util.List;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.Executor;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.stream.Collectors;
@@ -64,6 +65,8 @@
 import software.amazon.awssdk.codegen.poet.eventstream.EventStreamUtils;
 import software.amazon.awssdk.codegen.poet.model.EventStreamSpecHelper;
 import software.amazon.awssdk.core.RequestOverrideConfiguration;
+import software.amazon.awssdk.core.async.AsyncResponseTransformer;
+import software.amazon.awssdk.core.async.AsyncResponseTransformerUtils;
 import software.amazon.awssdk.core.async.SdkPublisher;
 import software.amazon.awssdk.core.client.config.SdkAdvancedAsyncClientOption;
 import software.amazon.awssdk.core.client.config.SdkClientConfiguration;
@@ -78,6 +81,7 @@
 import software.amazon.awssdk.protocols.json.AwsJsonProtocolFactory;
 import software.amazon.awssdk.utils.CompletableFutureUtils;
 import software.amazon.awssdk.utils.FunctionalUtils;
+import software.amazon.awssdk.utils.Pair;
 
 public final class AsyncClientClass extends AsyncClientInterface {
     private final IntermediateModel model;
@@ -244,7 +248,31 @@ protected MethodSpec.Builder operationBody(MethodSpec.Builder builder, Operation
                              CoreMetric.class, "SERVICE_ID", model.getMetadata().getServiceId());
         builder.addStatement("apiCallMetricCollector.reportMetric($T.$L, $S)",
                              CoreMetric.class, "OPERATION_NAME", opModel.getOperationName());
-
+
+        if (opModel.hasStreamingOutput()) {
+            ClassName responseType = poetExtensions.getModelClass(opModel.getReturnType().getReturnType());
+
+            builder.addStatement("$T<$T<$T, ReturnT>, $T<$T>> $N = $T.wrapWithEndOfStreamFuture($N)",
+                                 Pair.class,
+                                 AsyncResponseTransformer.class,
+                                 responseType,
+                                 CompletableFuture.class,
+                                 Void.class,
+                                 "pair",
+                                 AsyncResponseTransformerUtils.class,
+                                 "asyncResponseTransformer");
+
+            builder.addStatement("$N = $N.left()",
+                                 "asyncResponseTransformer",
+                                 "pair");
+
+            builder.addStatement("$T<$T> $N = $N.right()",
+                                 CompletableFuture.class,
+                                 Void.class,
+                                 "endOfStreamFuture",
+                                 "pair");
+        }
+
         if (shouldUseAsyncWithBodySigner(opModel)) {
             builder.addCode(applyAsyncWithBodyV4SignerOverride(opModel));
         } else {
@@ -312,8 +340,14 @@ protected MethodSpec.Builder operationBody(MethodSpec.Builder builder, Operation
                .beginControlFlow("catch ($T t)", Throwable.class);
 
         // For streaming operations we also want to notify the response handler of any exception.
+        if (opModel.hasStreamingOutput()) {
+            ClassName responseType = poetExtensions.getModelClass(opModel.getReturnType().getReturnType());
+            builder.addStatement("$T<$T, ReturnT> finalAsyncResponseTransformer = asyncResponseTransformer",
+                                 AsyncResponseTransformer.class,
+                                 responseType);
+        }
         if (opModel.hasStreamingOutput() || opModel.hasEventStreamOutput()) {
-            String paramName = opModel.hasStreamingOutput() ? "asyncResponseTransformer" : "asyncResponseHandler";
+            String paramName = opModel.hasStreamingOutput() ? "finalAsyncResponseTransformer" : "asyncResponseHandler";
             builder.addStatement("runAndLogError(log, \"Exception thrown in exceptionOccurred callback, ignoring\",\n" +
                                  "() -> $N.exceptionOccurred(t))", paramName);
         }

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

@@ -45,6 +45,7 @@
 import software.amazon.awssdk.core.SdkPojoBuilder;
 import software.amazon.awssdk.core.SdkResponse;
 import software.amazon.awssdk.core.async.AsyncRequestBody;
+import software.amazon.awssdk.core.async.AsyncResponseTransformer;
 import software.amazon.awssdk.core.client.handler.AttachHttpMetadataResponseHandler;
 import software.amazon.awssdk.core.client.handler.ClientExecutionParams;
 import software.amazon.awssdk.core.http.HttpResponseHandler;
@@ -223,7 +224,6 @@ public CodeBlock asyncExecutionHandler(IntermediateModel intermediateModel, Oper
 
         boolean isStreaming = opModel.hasStreamingOutput() || opModel.hasEventStreamOutput();
         String protocolFactory = protocolFactoryLiteral(intermediateModel, opModel);
-        String customerResponseHandler = opModel.hasEventStreamOutput() ? "asyncResponseHandler" : "asyncResponseTransformer";
         TypeName responseType = opModel.hasEventStreamOutput() && !isRestJson ? ClassName.get(SdkResponse.class)
                                                                               : pojoResponseType;
         TypeName executeFutureValueType = executeFutureValueType(opModel, poetExtensions);
@@ -245,7 +245,13 @@ public CodeBlock asyncExecutionHandler(IntermediateModel intermediateModel, Oper
                .add(".withInput($L)$L);",
                     opModel.getInput().getVariableName(), asyncResponseTransformerVariable(isStreaming, isRestJson, opModel));
 
-
+        if (opModel.hasStreamingOutput()) {
+            builder.addStatement("$T<$T, ReturnT> finalAsyncResponseTransformer = asyncResponseTransformer",
+                                 AsyncResponseTransformer.class,
+                                 pojoResponseType);
+        }
+        String customerResponseHandler = opModel.hasEventStreamOutput() ?
+                                         "asyncResponseHandler" : "finalAsyncResponseTransformer";
         String whenComplete = whenCompleteBody(opModel, customerResponseHandler);
         if (!whenComplete.isEmpty()) {
             String whenCompletedFutureName = "whenCompleted";

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

@@ -154,7 +154,9 @@ default String streamingOutputWhenComplete(String responseHandlerName) {
                              + "         runAndLogError(log, \"Exception thrown in exceptionOccurred callback, ignoring\", () "
                              + "-> %s.exceptionOccurred(e));%n"
                              + "     }%n"
-                             + "%s"
+                             + "     endOfStreamFuture.whenComplete((r2, e2) -> {%n"
+                             + "         %s%n"
+                             + "     });"
                              + "})", responseHandlerName, publishMetrics());
 
     }

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

@@ -29,6 +29,7 @@
 import software.amazon.awssdk.codegen.model.intermediate.OperationModel;
 import software.amazon.awssdk.codegen.poet.PoetExtensions;
 import software.amazon.awssdk.codegen.poet.client.traits.HttpChecksumRequiredTrait;
+import software.amazon.awssdk.core.async.AsyncResponseTransformer;
 import software.amazon.awssdk.core.client.handler.ClientExecutionParams;
 import software.amazon.awssdk.core.http.HttpResponseHandler;
 import software.amazon.awssdk.protocols.query.AwsQueryProtocolFactory;
@@ -151,8 +152,11 @@ public CodeBlock asyncExecutionHandler(IntermediateModel intermediateModel, Oper
         builder.addStatement("$T $N = null", ParameterizedTypeName.get(ClassName.get(CompletableFuture.class),
                 executeFutureValueType), whenCompleteFutureName);
         if (opModel.hasStreamingOutput()) {
+            builder.addStatement("$T<$T, ReturnT> finalAsyncResponseTransformer = asyncResponseTransformer",
+                                 AsyncResponseTransformer.class,
+                                 pojoResponseType);
             builder.addStatement("$N = executeFuture$L", whenCompleteFutureName,
-                    streamingOutputWhenComplete("asyncResponseTransformer"));
+                    streamingOutputWhenComplete("finalAsyncResponseTransformer"));
         } else {
             builder.addStatement("$N = executeFuture$L", whenCompleteFutureName, publishMetricsWhenComplete());
         }

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

@@ -209,7 +209,7 @@ public CodeBlock asyncExecutionHandler(IntermediateModel intermediateModel, Oper
         s3ArnableFields(opModel, model).ifPresent(builder::add);
 
         builder.add(".withInput($L)", opModel.getInput().getVariableName());
-        if (opModel.hasStreamingOutput() || opModel.hasEventStreamOutput()) {
+        if (opModel.hasEventStreamOutput()) {
             builder.add(", $N", executionResponseTransformerName);
         }
         builder.addStatement(")");
@@ -218,7 +218,7 @@ public CodeBlock asyncExecutionHandler(IntermediateModel intermediateModel, Oper
         builder.addStatement("$T $N = null", ParameterizedTypeName.get(ClassName.get(CompletableFuture.class),
                 executeFutureValueType), whenCompleteFutureName);
 
-        if (opModel.hasStreamingOutput() || opModel.hasEventStreamOutput()) {
+        if (opModel.hasEventStreamOutput()) {
             builder.addStatement("$N = executeFuture$L", whenCompleteFutureName,
                     whenCompleteBlock(opModel, "asyncResponseHandler",
                                       eventStreamTransformFutureName));

codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-async-client-class.java

@@ -30,6 +30,7 @@
 import software.amazon.awssdk.core.SdkResponse;
 import software.amazon.awssdk.core.async.AsyncRequestBody;
 import software.amazon.awssdk.core.async.AsyncResponseTransformer;
+import software.amazon.awssdk.core.async.AsyncResponseTransformerUtils;
 import software.amazon.awssdk.core.async.SdkPublisher;
 import software.amazon.awssdk.core.client.config.SdkAdvancedAsyncClientOption;
 import software.amazon.awssdk.core.client.config.SdkClientConfiguration;
@@ -106,6 +107,7 @@
 import software.amazon.awssdk.services.json.transform.StreamingOutputOperationRequestMarshaller;
 import software.amazon.awssdk.utils.CompletableFutureUtils;
 import software.amazon.awssdk.utils.HostnameValidator;
+import software.amazon.awssdk.utils.Pair;
 
 /**
  * Internal implementation of {@link JsonAsyncClient}.
@@ -998,6 +1000,10 @@ public <ReturnT> CompletableFuture<ReturnT> streamingInputOutputOperation(
         try {
             apiCallMetricCollector.reportMetric(CoreMetric.SERVICE_ID, "Json Service");
             apiCallMetricCollector.reportMetric(CoreMetric.OPERATION_NAME, "StreamingInputOutputOperation");
+            Pair<AsyncResponseTransformer<StreamingInputOutputOperationResponse, ReturnT>, CompletableFuture<Void>> pair =
+                AsyncResponseTransformerUtils.wrapWithEndOfStreamFuture(asyncResponseTransformer);
+            asyncResponseTransformer = pair.left();
+            CompletableFuture<Void> endOfStreamFuture = pair.right();
             streamingInputOutputOperationRequest = applySignerOverride(streamingInputOutputOperationRequest,
                     Aws4UnsignedPayloadSigner.create());
             JsonOperationMetadata operationMetadata = JsonOperationMetadata.builder().hasStreamingSuccessResponse(true)
@@ -1021,18 +1027,22 @@ public <ReturnT> CompletableFuture<ReturnT> streamingInputOutputOperation(
                             .withResponseHandler(responseHandler).withErrorResponseHandler(errorResponseHandler)
                             .withMetricCollector(apiCallMetricCollector).withAsyncRequestBody(requestBody)
                             .withInput(streamingInputOutputOperationRequest), asyncResponseTransformer);
+            AsyncResponseTransformer<StreamingInputOutputOperationResponse, ReturnT> finalAsyncResponseTransformer = asyncResponseTransformer;
             CompletableFuture<ReturnT> whenCompleted = executeFuture.whenComplete((r, e) -> {
                 if (e != null) {
                     runAndLogError(log, "Exception thrown in exceptionOccurred callback, ignoring",
-                            () -> asyncResponseTransformer.exceptionOccurred(e));
+                            () -> finalAsyncResponseTransformer.exceptionOccurred(e));
                 }
-                metricPublishers.forEach(p -> p.publish(apiCallMetricCollector.collect()));
+                endOfStreamFuture.whenComplete((r2, e2) -> {
+                    metricPublishers.forEach(p -> p.publish(apiCallMetricCollector.collect()));
+                });
             });
             executeFuture = CompletableFutureUtils.forwardExceptionTo(whenCompleted, executeFuture);
             return executeFuture;
         } catch (Throwable t) {
+            AsyncResponseTransformer<StreamingInputOutputOperationResponse, ReturnT> finalAsyncResponseTransformer = asyncResponseTransformer;
             runAndLogError(log, "Exception thrown in exceptionOccurred callback, ignoring",
-                    () -> asyncResponseTransformer.exceptionOccurred(t));
+                           () -> finalAsyncResponseTransformer.exceptionOccurred(t));
             metricPublishers.forEach(p -> p.publish(apiCallMetricCollector.collect()));
             return CompletableFutureUtils.failedFuture(t);
         }
@@ -1073,6 +1083,10 @@ public <ReturnT> CompletableFuture<ReturnT> streamingOutputOperation(
         try {
             apiCallMetricCollector.reportMetric(CoreMetric.SERVICE_ID, "Json Service");
             apiCallMetricCollector.reportMetric(CoreMetric.OPERATION_NAME, "StreamingOutputOperation");
+            Pair<AsyncResponseTransformer<StreamingOutputOperationResponse, ReturnT>, CompletableFuture<Void>> pair =
+                AsyncResponseTransformerUtils.wrapWithEndOfStreamFuture(asyncResponseTransformer);
+            asyncResponseTransformer = pair.left();
+            CompletableFuture<Void> endOfStreamFuture = pair.right();
             JsonOperationMetadata operationMetadata = JsonOperationMetadata.builder().hasStreamingSuccessResponse(true)
                     .isPayloadJson(false).build();
 
@@ -1089,18 +1103,22 @@ public <ReturnT> CompletableFuture<ReturnT> streamingOutputOperation(
                             .withResponseHandler(responseHandler).withErrorResponseHandler(errorResponseHandler)
                             .withMetricCollector(apiCallMetricCollector).withInput(streamingOutputOperationRequest),
                     asyncResponseTransformer);
+            AsyncResponseTransformer<StreamingOutputOperationResponse, ReturnT> finalAsyncResponseTransformer = asyncResponseTransformer;
             CompletableFuture<ReturnT> whenCompleted = executeFuture.whenComplete((r, e) -> {
                 if (e != null) {
                     runAndLogError(log, "Exception thrown in exceptionOccurred callback, ignoring",
-                            () -> asyncResponseTransformer.exceptionOccurred(e));
+                            () -> finalAsyncResponseTransformer.exceptionOccurred(e));
                 }
-                metricPublishers.forEach(p -> p.publish(apiCallMetricCollector.collect()));
+                endOfStreamFuture.whenComplete((r2, e2) -> {
+                    metricPublishers.forEach(p -> p.publish(apiCallMetricCollector.collect()));
+                });
             });
             executeFuture = CompletableFutureUtils.forwardExceptionTo(whenCompleted, executeFuture);
             return executeFuture;
         } catch (Throwable t) {
+            AsyncResponseTransformer<StreamingOutputOperationResponse, ReturnT> finalAsyncResponseTransformer = asyncResponseTransformer;
             runAndLogError(log, "Exception thrown in exceptionOccurred callback, ignoring",
-                    () -> asyncResponseTransformer.exceptionOccurred(t));
+                           () -> finalAsyncResponseTransformer.exceptionOccurred(t));
             metricPublishers.forEach(p -> p.publish(apiCallMetricCollector.collect()));
             return CompletableFutureUtils.failedFuture(t);
         }