Skip to content

Commit 4d9c592

Browse files
kuhekuhe
authored
chore(clients): add AWS retry customizations for DynamoDB, SQS, SFN, SWF (#7922)
1 parent b19357a commit 4d9c592

14 files changed

Lines changed: 363 additions & 4 deletions

File tree

‎clients/client-dynamodb-streams/src/runtimeConfig.shared.ts‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ import { NoOpLogger } from "@smithy/smithy-client";
55
import type { IdentityProviderConfig } from "@smithy/types";
66
import { parseUrl } from "@smithy/url-parser";
77
import { fromBase64, toBase64 } from "@smithy/util-base64";
8+
import { Retry, StandardRetryStrategy } from "@smithy/util-retry";
89
import { fromUtf8, toUtf8 } from "@smithy/util-utf8";
910

1011
import { defaultDynamoDBStreamsHttpAuthSchemeProvider } from "./auth/httpAuthSchemeProvider";
@@ -41,6 +42,14 @@ export const getRuntimeConfig = (config: DynamoDBStreamsClientConfig) => {
4142
version: "2012-08-10",
4243
serviceTarget: "DynamoDBStreams_20120810",
4344
},
45+
retryStrategy: config?.retryStrategy ?? (
46+
config?.maxAttempts == null && config?.retryMode == null && Retry.v2026
47+
? new StandardRetryStrategy({
48+
maxAttempts: 4,
49+
baseDelay: 25,
50+
})
51+
: undefined
52+
),
4453
serviceId: config?.serviceId ?? "DynamoDB Streams",
4554
urlParser: config?.urlParser ?? parseUrl,
4655
utf8Decoder: config?.utf8Decoder ?? fromUtf8,

‎clients/client-dynamodb/src/runtimeConfig.shared.ts‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ import { NoOpLogger } from "@smithy/smithy-client";
66
import type { IdentityProviderConfig } from "@smithy/types";
77
import { parseUrl } from "@smithy/url-parser";
88
import { fromBase64, toBase64 } from "@smithy/util-base64";
9+
import { Retry, StandardRetryStrategy } from "@smithy/util-retry";
910
import { fromUtf8, toUtf8 } from "@smithy/util-utf8";
1011

1112
import { defaultDynamoDBHttpAuthSchemeProvider } from "./auth/httpAuthSchemeProvider";
@@ -43,6 +44,14 @@ export const getRuntimeConfig = (config: DynamoDBClientConfig) => {
4344
serviceTarget: "DynamoDB_20120810",
4445
jsonCodec: new DynamoDBJsonCodec(),
4546
},
47+
retryStrategy: config?.retryStrategy ?? (
48+
config?.maxAttempts == null && config?.retryMode == null && Retry.v2026
49+
? new StandardRetryStrategy({
50+
maxAttempts: 4,
51+
baseDelay: 25,
52+
})
53+
: undefined
54+
),
4655
serviceId: config?.serviceId ?? "DynamoDB",
4756
urlParser: config?.urlParser ?? parseUrl,
4857
utf8Decoder: config?.utf8Decoder ?? fromUtf8,

‎clients/client-sfn/src/commands/GetActivityTaskCommand.ts‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
// smithy-typescript generated code
2+
import { getLongPollPlugin } from "@aws-sdk/core/client";
23
import { getEndpointPlugin } from "@smithy/middleware-endpoint";
34
import { Command as $Command } from "@smithy/smithy-client";
45
import type { MetadataBearer as __MetadataBearer } from "@smithy/types";
@@ -105,7 +106,10 @@ export class GetActivityTaskCommand extends $Command
105106
>()
106107
.ep(commonParams)
107108
.m(function (this: any, Command: any, cs: any, config: SFNClientResolvedConfig, o: any) {
108-
return [getEndpointPlugin(config, Command.getEndpointParameterInstructions())];
109+
return [
110+
getEndpointPlugin(config, Command.getEndpointParameterInstructions()),
111+
getLongPollPlugin(config),
112+
];
109113
})
110114
.s("AWSStepFunctions", "GetActivityTask", {})
111115
.n("SFNClient", "GetActivityTaskCommand")

‎clients/client-sqs/src/commands/ReceiveMessageCommand.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
// smithy-typescript generated code
2+
import { getLongPollPlugin } from "@aws-sdk/core/client";
23
import { getReceiveMessagePlugin } from "@aws-sdk/middleware-sdk-sqs";
34
import { getEndpointPlugin } from "@smithy/middleware-endpoint";
45
import { Command as $Command } from "@smithy/smithy-client";
@@ -218,6 +219,7 @@ export class ReceiveMessageCommand extends $Command
218219
.m(function (this: any, Command: any, cs: any, config: SQSClientResolvedConfig, o: any) {
219220
return [
220221
getEndpointPlugin(config, Command.getEndpointParameterInstructions()),
222+
getLongPollPlugin(config),
221223
getReceiveMessagePlugin(config),
222224
];
223225
})

‎clients/client-swf/src/commands/PollForActivityTaskCommand.ts‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
// smithy-typescript generated code
2+
import { getLongPollPlugin } from "@aws-sdk/core/client";
23
import { getEndpointPlugin } from "@smithy/middleware-endpoint";
34
import { Command as $Command } from "@smithy/smithy-client";
45
import type { MetadataBearer as __MetadataBearer } from "@smithy/types";
@@ -128,7 +129,10 @@ export class PollForActivityTaskCommand extends $Command
128129
>()
129130
.ep(commonParams)
130131
.m(function (this: any, Command: any, cs: any, config: SWFClientResolvedConfig, o: any) {
131-
return [getEndpointPlugin(config, Command.getEndpointParameterInstructions())];
132+
return [
133+
getEndpointPlugin(config, Command.getEndpointParameterInstructions()),
134+
getLongPollPlugin(config),
135+
];
132136
})
133137
.s("SimpleWorkflowService", "PollForActivityTask", {})
134138
.n("SWFClient", "PollForActivityTaskCommand")

‎clients/client-swf/src/commands/PollForDecisionTaskCommand.ts‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
// smithy-typescript generated code
2+
import { getLongPollPlugin } from "@aws-sdk/core/client";
23
import { getEndpointPlugin } from "@smithy/middleware-endpoint";
34
import { Command as $Command } from "@smithy/smithy-client";
45
import type { MetadataBearer as __MetadataBearer } from "@smithy/types";
@@ -524,7 +525,10 @@ export class PollForDecisionTaskCommand extends $Command
524525
>()
525526
.ep(commonParams)
526527
.m(function (this: any, Command: any, cs: any, config: SWFClientResolvedConfig, o: any) {
527-
return [getEndpointPlugin(config, Command.getEndpointParameterInstructions())];
528+
return [
529+
getEndpointPlugin(config, Command.getEndpointParameterInstructions()),
530+
getLongPollPlugin(config),
531+
];
528532
})
529533
.s("SimpleWorkflowService", "PollForDecisionTask", {})
530534
.n("SWFClient", "PollForDecisionTaskCommand")

‎codegen/smithy-aws-typescript-codegen/src/main/java/software/amazon/smithy/aws/typescript/codegen/AddRetryCustomizations.java‎

Lines changed: 94 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,94 @@
1+
/*
2+
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
package software.amazon.smithy.aws.typescript.codegen;
6+
7+
import java.util.Collections;
8+
import java.util.List;
9+
import java.util.Map;
10+
import java.util.function.Consumer;
11+
import software.amazon.smithy.aws.traits.ServiceTrait;
12+
import software.amazon.smithy.codegen.core.Symbol;
13+
import software.amazon.smithy.codegen.core.SymbolProvider;
14+
import software.amazon.smithy.model.Model;
15+
import software.amazon.smithy.model.shapes.OperationShape;
16+
import software.amazon.smithy.model.shapes.ServiceShape;
17+
import software.amazon.smithy.model.shapes.Shape;
18+
import software.amazon.smithy.typescript.codegen.LanguageTarget;
19+
import software.amazon.smithy.typescript.codegen.TypeScriptDependency;
20+
import software.amazon.smithy.typescript.codegen.TypeScriptSettings;
21+
import software.amazon.smithy.typescript.codegen.TypeScriptWriter;
22+
import software.amazon.smithy.typescript.codegen.integration.RuntimeClientPlugin;
23+
import software.amazon.smithy.typescript.codegen.integration.TypeScriptIntegration;
24+
import software.amazon.smithy.utils.SmithyInternalApi;
25+
26+
/**
27+
* Adds long poll designator for specific operations.
28+
* Adds custom retryStrategy for DynamoDB.
29+
*/
30+
@SmithyInternalApi
31+
public class AddRetryCustomizations implements TypeScriptIntegration {
32+
33+
@Override
34+
public List<RuntimeClientPlugin> getClientPlugins() {
35+
return List.of(
36+
RuntimeClientPlugin.builder()
37+
.pluginFunction(
38+
Symbol.builder()
39+
.namespace(AwsDependency.AWS_SDK_CORE.getPackageName() + "/client", "/")
40+
.name("getLongPollPlugin")
41+
.addDependency(AwsDependency.AWS_SDK_CORE)
42+
.build()
43+
)
44+
.operationPredicate((m, s, o) -> isLongPoll(s, o))
45+
.build()
46+
);
47+
}
48+
49+
@Override
50+
public Map<String, Consumer<TypeScriptWriter>> getRuntimeConfigWriters(
51+
TypeScriptSettings settings,
52+
Model model,
53+
SymbolProvider symbolProvider,
54+
LanguageTarget target
55+
) {
56+
if (LanguageTarget.SHARED.equals(target) && isDynamoDB(settings.getService(model))) {
57+
return Map.of(
58+
"retryStrategy",
59+
(w) -> {
60+
w.addImport("StandardRetryStrategy", null, TypeScriptDependency.UTIL_RETRY);
61+
w.addImport("Retry", null, TypeScriptDependency.UTIL_RETRY);
62+
63+
// todo(retry): Retry.v2026 condition can be removed when made permanent.
64+
w.write("""
65+
(
66+
config?.maxAttempts == null && config?.retryMode == null && Retry.v2026
67+
? new StandardRetryStrategy({
68+
maxAttempts: 4,
69+
baseDelay: 25,
70+
})
71+
: undefined
72+
)""");
73+
}
74+
);
75+
}
76+
return Collections.emptyMap();
77+
}
78+
79+
private static boolean isDynamoDB(Shape serviceShape) {
80+
String sdkId = serviceShape.getTrait(ServiceTrait.class).map(ServiceTrait::getSdkId).orElse("");
81+
return sdkId.equals("DynamoDB") || sdkId.equals("DynamoDB Streams");
82+
}
83+
84+
private static boolean isLongPoll(ServiceShape service, OperationShape operation) {
85+
String sdkId = service.getTrait(ServiceTrait.class).map(ServiceTrait::getSdkId).orElse("");
86+
String opName = operation.getId().getName(service);
87+
return switch (sdkId) {
88+
case "SQS" -> opName.equals("ReceiveMessage");
89+
case "SWF" -> opName.equals("PollForActivityTask") || opName.equals("PollForDecisionTask");
90+
case "SFN" -> opName.equals("GetActivityTask");
91+
default -> false;
92+
};
93+
}
94+
}

‎codegen/smithy-aws-typescript-codegen/src/main/resources/META-INF/services/software.amazon.smithy.typescript.codegen.integration.TypeScriptIntegration‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ software.amazon.smithy.aws.typescript.codegen.AddServiceCustomizationPlugins
66
software.amazon.smithy.aws.typescript.codegen.AddAwsAuthPlugin
77
software.amazon.smithy.aws.typescript.codegen.AddTokenAuthPlugin
88
software.amazon.smithy.aws.typescript.codegen.AddProtocols
9+
software.amazon.smithy.aws.typescript.codegen.AddRetryCustomizations
910
software.amazon.smithy.aws.typescript.codegen.AwsEndpointGeneratorIntegration
1011
software.amazon.smithy.aws.typescript.codegen.AddDefaultAwsEndpointRuleSet
1112
software.amazon.smithy.aws.typescript.codegen.AddNimbleCustomizations

‎packages-internal/core/integ/retry-customizations.integ.spec.ts‎

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,58 @@
1+
import { DynamoDB } from "@aws-sdk/client-dynamodb";
2+
import { DynamoDBStreams } from "@aws-sdk/client-dynamodb-streams";
3+
import { AdaptiveRetryStrategy, Retry, StandardRetryStrategy } from "@smithy/util-retry";
4+
import { beforeAll, describe, expect, test as it } from "vitest";
5+
6+
describe("retry customization", () => {
7+
for (const DDBCtor of [DynamoDB, DynamoDBStreams]) {
8+
describe(DDBCtor.name, () => {
9+
beforeAll(async () => {
10+
Retry.v2026 = true;
11+
});
12+
13+
it("has a custom retry strategy when no maxAttempts or retryMode are set", async () => {
14+
const streams = new DDBCtor({
15+
region: "us-west-2",
16+
credentials: {
17+
accessKeyId: "INTEG",
18+
secretAccessKey: "INTEG",
19+
},
20+
});
21+
const retryStrategy = await streams.config.retryStrategy();
22+
expect(retryStrategy).toBeInstanceOf(StandardRetryStrategy);
23+
24+
expect(await (retryStrategy as any).maxAttemptsProvider()).toBe(4);
25+
expect((retryStrategy as any).baseDelay).toBe(25);
26+
});
27+
28+
it("uses derived retryStrategy when retryMode is set", async () => {
29+
const streams = new DDBCtor({
30+
region: "us-west-2",
31+
credentials: {
32+
accessKeyId: "INTEG",
33+
secretAccessKey: "INTEG",
34+
},
35+
retryMode: "adaptive",
36+
});
37+
const retryStrategy = await streams.config.retryStrategy();
38+
expect(retryStrategy).toBeInstanceOf(AdaptiveRetryStrategy);
39+
});
40+
41+
it("uses derived retryStrategy when maxAttempts is set", async () => {
42+
const streams = new DDBCtor({
43+
region: "us-west-2",
44+
credentials: {
45+
accessKeyId: "INTEG",
46+
secretAccessKey: "INTEG",
47+
},
48+
maxAttempts: 5,
49+
});
50+
const retryStrategy = await streams.config.retryStrategy();
51+
expect(retryStrategy).toBeInstanceOf(StandardRetryStrategy);
52+
53+
expect(await (retryStrategy as any).maxAttemptsProvider()).toBe(5);
54+
expect((retryStrategy as any).baseDelay).toBe(50);
55+
});
56+
});
57+
}
58+
});

‎packages-internal/core/src/submodules/client/index.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
export * from "./emitWarningIfUnsupportedVersion";
2+
export { getLongPollPlugin } from "./longPollMiddleware";
23
export * from "./setCredentialFeature";
34
export * from "./setFeature";
45
export * from "./setTokenFeature";

0 commit comments

Comments
 (0)
⚡