fix: use region overrides in endpoint discovery (#2393)

.changelog/e50b13c4d0994557a1f0098adce2c23e.json

@@ -0,0 +1,11 @@
+{
+    "id": "e50b13c4-d099-4557-a1f0-098adce2c23e",
+    "type": "bugfix",
+    "description": "Respect caller region overrides in endpoint discovery.",
+    "modules": [
+        "service/dynamodb",
+        "service/internal/endpoint-discovery",
+        "service/timestreamquery",
+        "service/timestreamwrite"
+    ]
+}

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

@@ -189,6 +189,7 @@ private static void generateAddDiscoverEndpointMiddleware(
                         writer.write("EndpointDiscoveryEnableState: o.$L.$L,", ENDPOINT_DISCOVERY_OPTION, ENABLE_ENDPOINT_DISCOVERY_OPTION);
                         writer.write("EndpointDiscoveryRequired: $L,",
                                 operationRequiresEndpointDiscovery(model, service, operation));
+                        writer.write("Region: o.Region,");
                     });
                 });
         writer.write("");
@@ -216,11 +217,10 @@ private static void generateEndpointDiscoveryHandler(
         Symbol discoveryOperationInputSymbol = symbolProvider.toSymbol(discoveryOperationInput);
 
         writer.addUseImports(SmithyGoDependency.CONTEXT);
-        writer.openBlock("func (c *Client) $L(ctx context.Context, input $P, key string, opt $T) ($T, error) {", "}",
+        writer.openBlock("func (c *Client) $L(ctx context.Context, input $P, region, key string, opt $T) ($T, error) {", "}",
                 DISCOVERY_ENDPOINT_HANDLER_NAME, discoveryOperationInputSymbol, DISCOVERY_ENDPOINT_OPTIONS,
                 DISCOVERY_ENDPOINT_TYPE, () -> {
-
-                    if (serviceSupportsCustomDiscoveryEndpoint(model, service)) {
+if (serviceSupportsCustomDiscoveryEndpoint(model, service)) {
                         // check if endpoint resolver for endpoint discovery is of service-specific endpoint resolver type
                         writer.writeDocs("assert endpoint resolver interface is of expected type.");
                         writer.write("endpointResolver, ok := opt.$L.($T)", ENDPOINT_RESOLVER_USED_FOR_DISCOVERY,
@@ -238,6 +238,7 @@ private static void generateEndpointDiscoveryHandler(
                     // fetch endpoint via making discovery call
                     writer.openBlock("output, err := c.$T(ctx, input, func(o *Options) {", "})",
                             discoveryOperationSymbol, () -> {
+                                writer.write("o.Region = region").write("");
                                 writer.write("o.EndpointOptions.$L = opt.$L", DISABLE_HTTPS, DISABLE_HTTPS);
                                 writer.write("o.Logger = opt.Logger");
 
@@ -329,7 +330,7 @@ private static void generateFetchDiscoveredEndpointFunction(
 
         writer.addUseImports(SmithyGoDependency.CONTEXT);
         writer.openBlock(
-                "func (c *Client) $L(ctx context.Context, optFns ...func($P)) ($T, error) {", "}",
+                "func (c *Client) $L(ctx context.Context, region string, optFns ...func($P)) ($T, error) {", "}",
                 fetchDiscoveredEndpointFuncName, DISCOVERY_ENDPOINT_OPTIONS, DISCOVERY_ENDPOINT_WEIGHTED_ADDRESS,
                 () -> {
                     writer.write("input := getOperationInput(ctx)");
@@ -346,6 +347,10 @@ private static void generateFetchDiscoveredEndpointFunction(
                     String IDENTIFIER_MAP = "identifierMap";
                     writer.write("$L := make(map[string]string, 0)", IDENTIFIER_MAP);
 
+                    // include region at a minimum
+                    // see https://github.com/aws/aws-sdk-go-v2/issues/2163
+                    writer.write("$L[$S] = region", IDENTIFIER_MAP, "sdk#Region");
+
                     for (MemberShape member : getMembersUsedAsIdForDiscovery(model, service, operation)) {
                         String memberName = member.getMemberName();
                         Shape targetShape = model.expectShape(member.getTarget());
@@ -385,13 +390,13 @@ private static void generateFetchDiscoveredEndpointFunction(
 
                     // if discovery not required, then spin up a unblocking go routine
                     if (!operationRequiresEndpointDiscovery(model, service, operation)) {
-                        writer.write("go c.$L(ctx, $L, key, opt)", DISCOVERY_ENDPOINT_HANDLER_NAME,
+                        writer.write("go c.$L(ctx, $L, region, key, opt)", DISCOVERY_ENDPOINT_HANDLER_NAME,
                                 DISCOVERY_OPERATION_INPUT_NAME);
                         writer.write("return $T{}, nil", DISCOVERY_ENDPOINT_WEIGHTED_ADDRESS);
                         return;
                     }
 
-                    writer.write("endpoint, err := c.$L(ctx, $L, key, opt)", DISCOVERY_ENDPOINT_HANDLER_NAME,
+                    writer.write("endpoint, err := c.$L(ctx, $L, region, key, opt)", DISCOVERY_ENDPOINT_HANDLER_NAME,
                             DISCOVERY_OPERATION_INPUT_NAME);
                     writer.write("if err != nil { return $T{}, err }", DISCOVERY_ENDPOINT_WEIGHTED_ADDRESS);
                     writer.write("");