Skip to content

Commit 6ac6022

Browse files
aranditoarandito
authored
Add AWS QUERY protocol support (#673)
1 parent 68d6412 commit 6ac6022

38 files changed

+2672
-39
lines changed

‎Makefile‎

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -14,10 +14,12 @@ build-java: ## Builds the Java code generation packages.
1414
cd codegen && ./gradlew clean build
1515

1616

17-
test-protocols: ## Generates and runs the restJson1 protocol tests.
18-
cd codegen && ./gradlew :protocol-test:build
19-
uv pip install codegen/protocol-test/build/smithyprojections/protocol-test/rest-json-1/python-client-codegen
20-
uv run pytest codegen/protocol-test/build/smithyprojections/protocol-test/rest-json-1/python-client-codegen
17+
test-protocols: ## Generates and runs protocol tests for all supported protocols.
18+
cd codegen && ./gradlew :protocol-test:clean :protocol-test:build
19+
@set -e; for projection_dir in codegen/protocol-test/build/smithyprojections/protocol-test/*/python-client-codegen; do \
20+
uv pip install "$$projection_dir"; \
21+
uv run pytest "$$projection_dir"; \
22+
done
2123

2224

2325
lint-py: ## Runs linters and formatters on the python packages.

‎codegen/aws/core/build.gradle.kts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,4 +12,5 @@ extra["moduleName"] = "software.amazon.smithy.python.aws.codegen"
1212
dependencies {
1313
implementation(project(":core"))
1414
implementation(libs.smithy.aws.traits)
15+
implementation(libs.smithy.protocol.test.traits)
1516
}

‎codegen/aws/core/src/main/java/software/amazon/smithy/python/aws/codegen/AwsProtocolsIntegration.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,6 @@
1616
public class AwsProtocolsIntegration implements PythonIntegration {
1717
@Override
1818
public List<ProtocolGenerator> getProtocolGenerators() {
19-
return List.of();
19+
return List.of(new AwsQueryProtocolGenerator());
2020
}
2121
}

‎codegen/aws/core/src/main/java/software/amazon/smithy/python/aws/codegen/AwsQueryProtocolGenerator.java‎

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,69 @@
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.python.aws.codegen;
6+
7+
import java.util.Set;
8+
import software.amazon.smithy.aws.traits.protocols.AwsQueryTrait;
9+
import software.amazon.smithy.model.shapes.ShapeId;
10+
import software.amazon.smithy.python.codegen.ApplicationProtocol;
11+
import software.amazon.smithy.python.codegen.GenerationContext;
12+
import software.amazon.smithy.python.codegen.HttpProtocolTestGenerator;
13+
import software.amazon.smithy.python.codegen.SymbolProperties;
14+
import software.amazon.smithy.python.codegen.generators.ProtocolGenerator;
15+
import software.amazon.smithy.python.codegen.writer.PythonWriter;
16+
import software.amazon.smithy.utils.SmithyInternalApi;
17+
18+
@SmithyInternalApi
19+
public final class AwsQueryProtocolGenerator implements ProtocolGenerator {
20+
private static final Set<String> TESTS_TO_SKIP = Set.of(
21+
// TODO: support the request compression trait
22+
// https://smithy.io/2.0/spec/behavior-traits.html#smithy-api-requestcompression-trait
23+
"SDKAppliedContentEncoding_awsQuery",
24+
"SDKAppendsGzipAndIgnoresHttpProvidedEncoding_awsQuery",
25+
26+
// TODO: support idempotency token autofill
27+
"QueryProtocolIdempotencyTokenAutoFill",
28+
29+
// This test asserts nan == nan, which is never true.
30+
// We should update the generator to make specific assertions for these.
31+
"AwsQuerySupportsNaNFloatOutputs",
32+
33+
// TODO: support of the endpoint trait
34+
"AwsQueryEndpointTraitWithHostLabel",
35+
"AwsQueryEndpointTrait");
36+
37+
@Override
38+
public ShapeId getProtocol() {
39+
return AwsQueryTrait.ID;
40+
}
41+
42+
@Override
43+
public ApplicationProtocol getApplicationProtocol(GenerationContext context) {
44+
return ApplicationProtocol.createDefaultHttpApplicationProtocol();
45+
}
46+
47+
@Override
48+
public void initializeProtocol(GenerationContext context, PythonWriter writer) {
49+
writer.addDependency(AwsPythonDependency.SMITHY_AWS_CORE.withOptionalDependencies("xml"));
50+
writer.addImport("smithy_aws_core.aio.protocols", "AwsQueryClientProtocol");
51+
var service = context.settings().service(context.model());
52+
var serviceSymbol = context.symbolProvider().toSymbol(service);
53+
var serviceSchema = serviceSymbol.expectProperty(SymbolProperties.SCHEMA);
54+
var version = service.getVersion();
55+
writer.write("AwsQueryClientProtocol($T, $S)", serviceSchema, version);
56+
}
57+
58+
@Override
59+
public void generateProtocolTests(GenerationContext context) {
60+
context.writerDelegator()
61+
.useFileWriter("./tests/test_awsquery_protocol.py", "tests.test_awsquery_protocol", writer -> {
62+
new HttpProtocolTestGenerator(
63+
context,
64+
getProtocol(),
65+
writer,
66+
(shape, testCase) -> TESTS_TO_SKIP.contains(testCase.getId())).run();
67+
});
68+
}
69+
}

‎codegen/core/src/main/java/software/amazon/smithy/python/codegen/HttpProtocolTestGenerator.java‎

Lines changed: 33 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
import java.util.logging.Logger;
1717
import java.util.stream.Collectors;
1818
import java.util.stream.Stream;
19+
import software.amazon.smithy.aws.traits.auth.SigV4Trait;
1920
import software.amazon.smithy.codegen.core.CodegenException;
2021
import software.amazon.smithy.codegen.core.Symbol;
2122
import software.amazon.smithy.model.Model;
@@ -188,12 +189,14 @@ private void generateRequestTest(OperationShape operation, HttpRequestTestCase t
188189
endpoint_uri="https://$L/$L",
189190
transport = $T(),
190191
retry_strategy=SimpleRetryStrategy(max_attempts=1),
192+
${C|}
191193
)
192194
""",
193195
CodegenUtils.getConfigSymbol(context.settings()),
194196
host,
195197
path,
196-
REQUEST_TEST_ASYNC_HTTP_CLIENT_SYMBOL);
198+
REQUEST_TEST_ASYNC_HTTP_CLIENT_SYMBOL,
199+
(Runnable) this::writeSigV4TestConfig);
197200
}));
198201

199202
// Generate the input using the expected shape and params
@@ -418,6 +421,16 @@ private void compareMediaBlob(HttpMessageTestCase testCase, PythonWriter writer)
418421
""");
419422
return;
420423
}
424+
if (contentType.equals("application/x-www-form-urlencoded")) {
425+
writer.addStdlibImport("urllib.parse", "parse_qsl");
426+
writer.write("""
427+
actual_params = sorted(parse_qsl(actual_body_content.decode(), keep_blank_values=True))
428+
expected_params = sorted(parse_qsl(expected_body_content.decode(), keep_blank_values=True))
429+
assert actual_params == expected_params
430+
431+
""");
432+
return;
433+
}
421434
writer.write("assert actual_body_content == expected_body_content\n");
422435
}
423436

@@ -437,13 +450,15 @@ private void generateResponseTest(OperationShape operation, HttpResponseTestCase
437450
headers=$J,
438451
body=b$S,
439452
),
453+
${C|}
440454
)
441455
""",
442456
CodegenUtils.getConfigSymbol(context.settings()),
443457
RESPONSE_TEST_ASYNC_HTTP_CLIENT_SYMBOL,
444458
testCase.getCode(),
445459
CodegenUtils.toTuples(testCase.getHeaders()),
446-
testCase.getBody().filter(body -> !body.isEmpty()).orElse(""));
460+
testCase.getBody().filter(body -> !body.isEmpty()).orElse(""),
461+
(Runnable) this::writeSigV4TestConfig);
447462
}));
448463
// Create an empty input object to pass
449464
var inputShape = model.expectShape(operation.getInputShape(), StructureShape.class);
@@ -490,13 +505,15 @@ private void generateErrorResponseTest(
490505
headers=$J,
491506
body=b$S,
492507
),
508+
${C|}
493509
)
494510
""",
495511
CodegenUtils.getConfigSymbol(context.settings()),
496512
RESPONSE_TEST_ASYNC_HTTP_CLIENT_SYMBOL,
497513
testCase.getCode(),
498514
CodegenUtils.toTuples(testCase.getHeaders()),
499-
testCase.getBody().orElse(""));
515+
testCase.getBody().orElse(""),
516+
(Runnable) this::writeSigV4TestConfig);
500517
}));
501518
// Create an empty input object to pass
502519
var inputShape = model.expectShape(operation.getInputShape(), StructureShape.class);
@@ -607,6 +624,19 @@ private void writeClientBlock(
607624
});
608625
}
609626

627+
private void writeSigV4TestConfig() {
628+
if (!service.hasTrait(SigV4Trait.class)) {
629+
return;
630+
}
631+
writer.addImport("smithy_aws_core.identity", "StaticCredentialsResolver");
632+
writer.write("""
633+
region="us-east-1",
634+
aws_access_key_id="test-access-key-id",
635+
aws_secret_access_key="test-secret-access-key",
636+
aws_credentials_identity_resolver=StaticCredentialsResolver(),
637+
""");
638+
}
639+
610640
private void writeUtilStubs(Symbol serviceSymbol) {
611641
LOGGER.fine(String.format("Writing utility stubs for %s : %s", serviceSymbol.getName(), protocol.getName()));
612642
writer.addDependency(SmithyPythonDependency.SMITHY_CORE);

‎codegen/core/src/main/java/software/amazon/smithy/python/codegen/generators/OperationGenerator.java‎

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,9 @@ public void run() {
5858
}),
5959
effective_auth_schemes = [
6060
$8C
61+
],
62+
error_schemas = [
63+
$9C
6164
]
6265
)
6366
""",
@@ -68,7 +71,8 @@ public void run() {
6871
inSymbol.expectProperty(SymbolProperties.SCHEMA),
6972
outSymbol.expectProperty(SymbolProperties.SCHEMA),
7073
writer.consumer(this::writeErrorTypeRegistry),
71-
writer.consumer(this::writeAuthSchemes));
74+
writer.consumer(this::writeAuthSchemes),
75+
writer.consumer(this::writeErrorSchemas));
7276
}
7377

7478
private void writeErrorTypeRegistry(PythonWriter writer) {
@@ -82,6 +86,13 @@ private void writeErrorTypeRegistry(PythonWriter writer) {
8286
}
8387
}
8488

89+
private void writeErrorSchemas(PythonWriter writer) {
90+
for (var error : shape.getErrors()) {
91+
var errSymbol = symbolProvider.toSymbol(model.expectShape(error));
92+
writer.write("$T,", errSymbol.expectProperty(SymbolProperties.SCHEMA));
93+
}
94+
}
95+
8596
private void writeAuthSchemes(PythonWriter writer) {
8697
var authSchemes = ServiceIndex.of(model)
8798
.getEffectiveAuthSchemes(context.settings().service(),

‎codegen/protocol-test/build.gradle.kts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,5 +30,6 @@ repositories {
3030

3131
dependencies {
3232
implementation(project(":core"))
33+
implementation(project(":aws:core"))
3334
implementation(libs.smithy.aws.protocol.tests)
3435
}

‎codegen/protocol-test/smithy-build.json‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,28 @@
2222
"moduleVersion": "0.0.1"
2323
}
2424
}
25+
},
26+
"aws-query": {
27+
"transforms": [
28+
{
29+
"name": "includeServices",
30+
"args": {
31+
"services": [
32+
"aws.protocoltests.query#AwsQuery"
33+
]
34+
}
35+
},
36+
{
37+
"name": "removeUnusedShapes"
38+
}
39+
],
40+
"plugins": {
41+
"python-client-codegen": {
42+
"service": "aws.protocoltests.query#AwsQuery",
43+
"module": "awsquery",
44+
"moduleVersion": "0.0.1"
45+
}
46+
}
2547
}
2648
}
2749
}

‎packages/smithy-aws-core/.changes/next-release/smithy-aws-core-feature-ec24e8dbe26b4ee58029f89da5b6526e.json‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
{
2+
"type": "feature",
3+
"description": "Add `awsQuery` protocol support for Smithy clients."
4+
}

‎packages/smithy-aws-core/pyproject.toml‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,9 @@ eventstream = [
5050
json = [
5151
"smithy-json~=0.2.0"
5252
]
53+
xml = [
54+
"smithy-xml~=0.0.0"
55+
]
5356

5457
[tool.hatch.build]
5558
exclude = [

0 commit comments

Comments
 (0)