add interceptors (#3143)

Author: lucix-aws

.changelog/5e1864bad1c34d96a76929f148c61b52.json

@@ -1 +1 @@
-107eb1414a9c309076941ef37a6c71465ec397f6
+a22818fe19b9bc7d4926501907b154d813b309a5

SMITHY_GO_CODEGEN_VERSION

@@ -6,6 +6,7 @@ import (
 	smithybearer "github.com/aws/smithy-go/auth/bearer"
 	"github.com/aws/smithy-go/logging"
 	"github.com/aws/smithy-go/middleware"
+	smithyhttp "github.com/aws/smithy-go/transport/http"
 )
 
 // HTTPClient provides the interface to provide custom HTTPClients. Generally
@@ -192,6 +193,9 @@ type Config struct {
 	// This variable is sourced from environment variable AWS_RESPONSE_CHECKSUM_VALIDATION or
 	// the shared config profile attribute "response_checksum_validation".
 	ResponseChecksumValidation ResponseChecksumValidation
+
+	// Registry of HTTP interceptors.
+	Interceptors smithyhttp.InterceptorRegistry
 }
 
 // NewConfig returns a new Config pointer that can be chained with builder

aws/config.go

@@ -2,6 +2,6 @@ module github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream
 
 go 1.22
 
-require github.com/aws/smithy-go v1.22.4
+require github.com/aws/smithy-go v1.22.5
 
 replace github.com/aws/aws-sdk-go-v2 => ../../../

aws/protocol/eventstream/go.mod

@@ -1,2 +1,2 @@
-github.com/aws/smithy-go v1.22.4 h1:uqXzVZNuNexwc/xrh6Tb56u89WDlJY6HS+KC0S4QSjw=
-github.com/aws/smithy-go v1.22.4/go.mod h1:t1ufH5HMublsJYulve2RKmHDC15xu1f26kHCp/HgceI=
+github.com/aws/smithy-go v1.22.5 h1:P9ATCXPMb2mPjYBgueqJNCA5S9UfktsW0tTxi+a7eqw=
+github.com/aws/smithy-go v1.22.5/go.mod h1:t1ufH5HMublsJYulve2RKmHDC15xu1f26kHCp/HgceI=

aws/protocol/eventstream/go.sum

@@ -260,6 +260,12 @@ public class AddAwsConfigFields implements GoIntegration {
                     .type(SdkGoTypes.Aws.ResponseChecksumValidation)
                     .documentation("Indicates how user opt-in/out response checksum validation")
                     .servicePredicate(AwsHttpChecksumGenerator::hasOutputChecksumTrait)
+                    .build(),
+            AwsConfigField.builder()
+                    .name("Interceptors")
+                    .type(SmithyGoDependency.SMITHY_HTTP_TRANSPORT.struct("InterceptorRegistry"))
+                    .generatedOnClient(false)
+                    .awsResolveFunction(buildPackageSymbol("resolveInterceptors"))
                     .build()
     );
 
@@ -324,6 +330,15 @@ private void writeAwsDefaultResolvers(GoWriter writer) {
         writeRetryerResolvers(writer);
         writeRetryMaxAttemptsFinalizers(writer);
         writeAwsConfigEndpointResolver(writer);
+        writeInterceptorResolver(writer);
+    }
+
+    private void writeInterceptorResolver(GoWriter writer) {
+        writer.write("""
+                func resolveInterceptors(cfg aws.Config, o *Options) {
+                    o.Interceptors = cfg.Interceptors.Copy()
+                }
+                """);
     }
 
     private void writerAwsDefaultResolversTests(GoWriter writer) {

codegen/smithy-aws-go-codegen/src/main/java/software/amazon/smithy/aws/go/codegen/AddAwsConfigFields.java

@@ -0,0 +1,56 @@
+package software.amazon.smithy.aws.go.codegen.customization;
+
+import static software.amazon.smithy.go.codegen.SymbolUtils.buildPackageSymbol;
+
+import java.util.List;
+import software.amazon.smithy.go.codegen.GoCodegenContext;
+import software.amazon.smithy.go.codegen.integration.GoIntegration;
+import software.amazon.smithy.go.codegen.integration.MiddlewareRegistrar;
+import software.amazon.smithy.go.codegen.integration.RuntimeClientPlugin;
+
+// We add retry interceptors here since Smithy clients don't have a retry loop atm.
+public class RetryInterceptors implements GoIntegration {
+    private static RuntimeClientPlugin interceptor(String name) {
+        return RuntimeClientPlugin.builder()
+                .registerMiddleware(
+                        MiddlewareRegistrar.builder()
+                                .resolvedFunction(buildPackageSymbol(name))
+                                .useClientOptions()
+                                .build()
+                )
+                .build();
+    }
+
+    @Override
+    public byte getOrder() {
+        return 127;
+    }
+
+    @Override
+    public List<RuntimeClientPlugin> getClientPlugins() {
+        return List.of(
+                interceptor("addInterceptBeforeRetryLoop"),
+                interceptor("addInterceptAttempt")
+        );
+    }
+
+    @Override
+    public void writeAdditionalFiles(GoCodegenContext ctx) {
+        ctx.writerDelegator().useFileWriter("api_client.go", ctx.settings().getModuleName(), writer -> {
+            writer.write("""
+                    func addInterceptBeforeRetryLoop(stack *middleware.Stack, opts Options) error {
+                        return stack.Finalize.Insert(&smithyhttp.InterceptBeforeRetryLoop{
+                            Interceptors: opts.Interceptors.BeforeRetryLoop,
+                        }, "Retry", middleware.Before)
+                    }
+
+                    func addInterceptAttempt(stack *middleware.Stack, opts Options) error {
+                        return stack.Finalize.Insert(&smithyhttp.InterceptAttempt{
+                            BeforeAttempt: opts.Interceptors.BeforeAttempt,
+                            AfterAttempt: opts.Interceptors.AfterAttempt,
+                        }, "Retry", middleware.After)
+                    }
+                    """);
+        });
+    }
+}

codegen/smithy-aws-go-codegen/src/main/java/software/amazon/smithy/aws/go/codegen/customization/RetryInterceptors.java

@@ -91,4 +91,5 @@ software.amazon.smithy.aws.go.codegen.customization.ChecksumMetricsTracking
 software.amazon.smithy.aws.go.codegen.customization.AccountIdEndpointModeUserAgent
 software.amazon.smithy.aws.go.codegen.CredentialSourceFeatureTrackerGenerator
 software.amazon.smithy.aws.go.codegen.customization.SraOperationOrderTest
-software.amazon.smithy.aws.go.codegen.customization.service.bedrockruntime.RemoveOperations
+software.amazon.smithy.aws.go.codegen.customization.service.bedrockruntime.RemoveOperations
+software.amazon.smithy.aws.go.codegen.customization.RetryInterceptors

codegen/smithy-aws-go-codegen/src/main/resources/META-INF/services/software.amazon.smithy.go.codegen.integration.GoIntegration

@@ -89,6 +89,8 @@ var defaultAWSConfigResolvers = []awsConfigResolver{
 
 	// Sets the ResponseChecksumValidation if present in env var or shared config profile
 	resolveResponseChecksumValidation,
+
+	resolveInterceptors,
 }
 
 // A Config represents a generic configuration value or set of values. This type

config/config.go

@@ -10,7 +10,7 @@ require (
 	github.com/aws/aws-sdk-go-v2/service/sso v1.25.6
 	github.com/aws/aws-sdk-go-v2/service/ssooidc v1.30.4
 	github.com/aws/aws-sdk-go-v2/service/sts v1.34.1
-	github.com/aws/smithy-go v1.22.4
+	github.com/aws/smithy-go v1.22.5
 )
 
 require (