add client opts extension point for waiters/paginators (#518)

aws · Jun 18, 2024 · 5ca24b4 · 5ca24b4
1 parent e421cf2
commit 5ca24b4
Show file tree

Hide file tree

Showing 3 changed files with 35 additions and 3 deletions.
diff --git a/...hy-go-codegen/src/main/java/software/amazon/smithy/go/codegen/integration/Paginators.java b/...hy-go-codegen/src/main/java/software/amazon/smithy/go/codegen/integration/Paginators.java
@@ -15,7 +15,11 @@
 
 package software.amazon.smithy.go.codegen.integration;
 
+import static java.util.Collections.emptySet;
+import static software.amazon.smithy.go.codegen.GoWriter.goTemplate;
+
 import java.util.Optional;
+import java.util.Set;
 import java.util.function.Consumer;
 import software.amazon.smithy.codegen.core.Symbol;
 import software.amazon.smithy.codegen.core.SymbolProvider;
@@ -42,6 +46,10 @@
  * Implements support for PaginatedTrait.
  */
 public class Paginators implements GoIntegration {
+    public Set<Symbol> getAdditionalClientOptions() {
+        return emptySet();
+    }
+
     @Override
     public void writeAdditionalFiles(
             GoSettings settings,
@@ -217,13 +225,21 @@ private void writePaginator(
                         }
                     });
 
+                    var optFns = GoWriter.ChainWritable.of(
+                            getAdditionalClientOptions().stream()
+                                    .map(it -> goTemplate("$T,", it))
+                                    .toList()
+                    ).compose(false);
                     writer.write("""
+                                 optFns = append([]func(*Options) {
+                                     $W
+                                 }, optFns...)
                                  result, err := p.client.$L(ctx, &params, optFns...)
                                  if err != nil {
                                      return nil, err
                                  }
                                  p.firstPage = false
-                                 """, operationSymbol.getName());
+                                 """, optFns, operationSymbol.getName());
 
                     var outputMemberPath = paginationInfo.getOutputTokenMemberPath();
                     var tokenMember = outputMemberPath.get(outputMemberPath.size() - 1);

diff --git a/...mithy-go-codegen/src/main/java/software/amazon/smithy/go/codegen/integration/Waiters.java b/...mithy-go-codegen/src/main/java/software/amazon/smithy/go/codegen/integration/Waiters.java
@@ -15,11 +15,13 @@
 
 package software.amazon.smithy.go.codegen.integration;
 
+import static java.util.Collections.emptySet;
 import static software.amazon.smithy.go.codegen.GoWriter.autoDocTemplate;
 import static software.amazon.smithy.go.codegen.GoWriter.goTemplate;
 
 import java.util.Map;
 import java.util.Optional;
+import java.util.Set;
 import software.amazon.smithy.codegen.core.CodegenException;
 import software.amazon.smithy.codegen.core.Symbol;
 import software.amazon.smithy.codegen.core.SymbolProvider;
@@ -53,6 +55,10 @@ public class Waiters implements GoIntegration {
     private static final String WAITER_INVOKER_FUNCTION_NAME = "Wait";
     private static final String WAITER_INVOKER_WITH_OUTPUT_FUNCTION_NAME = "WaitForOutput";
 
+    public Set<Symbol> getAdditionalClientOptions() {
+        return emptySet();
+    }
+
     @Override
     public void writeAdditionalFiles(
             GoSettings settings,
@@ -440,10 +446,22 @@ private void generateWaiterInvokerWithOutput(
                         }).write("");
 
                         // make a request
+                        var baseOpts = GoWriter.ChainWritable.of(
+                                getAdditionalClientOptions().stream()
+                                        .map(it -> goTemplate("$T,", it))
+                                        .toList()
+                        ).compose(false);
                         writer.openBlock("out, err := w.client.$T(ctx, params, func (o *Options) { ", "})",
                                 operationSymbol, () -> {
+                                    writer.write("""
+                                            baseOpts := []func(*Options) {
+                                                $W
+                                            }""", baseOpts);
                                     writer.write("o.APIOptions = append(o.APIOptions, apiOptions...)");
                                     writer.write("""
+                                            for _, opt := range baseOpts {
+                                                opt(o)
+                                            }
                                             for _, opt := range options.ClientOptions {
                                                 opt(o)
                                             }""");

diff --git a/...n/resources/META-INF/services/software.amazon.smithy.go.codegen.integration.GoIntegration b/...n/resources/META-INF/services/software.amazon.smithy.go.codegen.integration.GoIntegration
@@ -4,8 +4,6 @@ software.amazon.smithy.go.codegen.integration.AddChecksumRequiredMiddleware
 software.amazon.smithy.go.codegen.integration.RequiresLengthTraitSupport
 software.amazon.smithy.go.codegen.integration.EndpointHostPrefixMiddleware
 software.amazon.smithy.go.codegen.integration.OperationInterfaceGenerator
-software.amazon.smithy.go.codegen.integration.Paginators
-software.amazon.smithy.go.codegen.integration.Waiters
 software.amazon.smithy.go.codegen.integration.ClientLogger
 software.amazon.smithy.go.codegen.endpoints.EndpointClientPluginsGenerator