From 2c8d0ed0d6de2060bf328bc4ed3035e881c86289 Mon Sep 17 00:00:00 2001
From: Luc Talatinian <lucix@amazon.com>
Date: Fri, 14 Jun 2024 17:16:32 -0400
Subject: [PATCH] add client opts extension point for waiters/paginators

---
 .../go/codegen/integration/Paginators.java     | 18 +++++++++++++++++-
 .../smithy/go/codegen/integration/Waiters.java | 18 ++++++++++++++++++
 ...smithy.go.codegen.integration.GoIntegration |  2 --
 3 files changed, 35 insertions(+), 3 deletions(-)

diff --git a/codegen/smithy-go-codegen/src/main/java/software/amazon/smithy/go/codegen/integration/Paginators.java b/codegen/smithy-go-codegen/src/main/java/software/amazon/smithy/go/codegen/integration/Paginators.java
index 4f35befde..e9bd6e58c 100644
--- a/codegen/smithy-go-codegen/src/main/java/software/amazon/smithy/go/codegen/integration/Paginators.java
+++ b/codegen/smithy-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/codegen/smithy-go-codegen/src/main/java/software/amazon/smithy/go/codegen/integration/Waiters.java b/codegen/smithy-go-codegen/src/main/java/software/amazon/smithy/go/codegen/integration/Waiters.java
index ebbd18082..40a01c847 100644
--- a/codegen/smithy-go-codegen/src/main/java/software/amazon/smithy/go/codegen/integration/Waiters.java
+++ b/codegen/smithy-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/codegen/smithy-go-codegen/src/main/resources/META-INF/services/software.amazon.smithy.go.codegen.integration.GoIntegration b/codegen/smithy-go-codegen/src/main/resources/META-INF/services/software.amazon.smithy.go.codegen.integration.GoIntegration
index 4781a49a2..04ed5213c 100644
--- a/codegen/smithy-go-codegen/src/main/resources/META-INF/services/software.amazon.smithy.go.codegen.integration.GoIntegration
+++ b/codegen/smithy-go-codegen/src/main/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