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 allows users to directly consume a streaming-response payload (i.e., S3 GetObject) without having to buffer to memory or disk. 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.NotifyingAsyncResponseTransformer;
 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",
+                                 NotifyingAsyncResponseTransformer.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/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));