Add AWS QUERY protocol support

Makefile

@@ -14,10 +14,12 @@ build-java: ## Builds the Java code generation packages.
 	cd codegen && ./gradlew clean build
 
 
-test-protocols: ## Generates and runs the restJson1 protocol tests.
-	cd codegen && ./gradlew :protocol-test:build
-	uv pip install codegen/protocol-test/build/smithyprojections/protocol-test/rest-json-1/python-client-codegen
-	uv run pytest codegen/protocol-test/build/smithyprojections/protocol-test/rest-json-1/python-client-codegen
+test-protocols: ## Generates and runs protocol tests for all supported protocols.
+	cd codegen && ./gradlew :protocol-test:clean :protocol-test:build
+	@set -e; for projection_dir in codegen/protocol-test/build/smithyprojections/protocol-test/*/python-client-codegen; do \
+		uv pip install "$$projection_dir"; \
+		uv run pytest "$$projection_dir"; \
+	done
 
 
 lint-py: ## Runs linters and formatters on the python packages.

codegen/aws/core/build.gradle.kts

@@ -12,4 +12,5 @@ extra["moduleName"] = "software.amazon.smithy.python.aws.codegen"
 dependencies {
     implementation(project(":core"))
     implementation(libs.smithy.aws.traits)
+    implementation(libs.smithy.protocol.test.traits)
 }

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

@@ -16,6 +16,6 @@
 public class AwsProtocolsIntegration implements PythonIntegration {
     @Override
     public List<ProtocolGenerator> getProtocolGenerators() {
-        return List.of();
+        return List.of(new AwsQueryProtocolGenerator());
     }
 }

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

@@ -0,0 +1,69 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package software.amazon.smithy.python.aws.codegen;
+
+import java.util.Set;
+import software.amazon.smithy.aws.traits.protocols.AwsQueryTrait;
+import software.amazon.smithy.model.shapes.ShapeId;
+import software.amazon.smithy.python.codegen.ApplicationProtocol;
+import software.amazon.smithy.python.codegen.GenerationContext;
+import software.amazon.smithy.python.codegen.HttpProtocolTestGenerator;
+import software.amazon.smithy.python.codegen.SymbolProperties;
+import software.amazon.smithy.python.codegen.generators.ProtocolGenerator;
+import software.amazon.smithy.python.codegen.writer.PythonWriter;
+import software.amazon.smithy.utils.SmithyInternalApi;
+
+@SmithyInternalApi
+public final class AwsQueryProtocolGenerator implements ProtocolGenerator {
+    private static final Set<String> TESTS_TO_SKIP = Set.of(
+            // TODO: support the request compression trait
+            // https://smithy.io/2.0/spec/behavior-traits.html#smithy-api-requestcompression-trait
+            "SDKAppliedContentEncoding_awsQuery",
+            "SDKAppendsGzipAndIgnoresHttpProvidedEncoding_awsQuery",
+
+            // TODO: support idempotency token autofill
+            "QueryProtocolIdempotencyTokenAutoFill",
+
+            // This test asserts nan == nan, which is never true.
+            // We should update the generator to make specific assertions for these.
+            "AwsQuerySupportsNaNFloatOutputs",
+
+            // TODO: support of the endpoint trait
+            "AwsQueryEndpointTraitWithHostLabel",
+            "AwsQueryEndpointTrait");
+
+    @Override
+    public ShapeId getProtocol() {
+        return AwsQueryTrait.ID;
+    }
+
+    @Override
+    public ApplicationProtocol getApplicationProtocol(GenerationContext context) {
+        return ApplicationProtocol.createDefaultHttpApplicationProtocol();
+    }
+
+    @Override
+    public void initializeProtocol(GenerationContext context, PythonWriter writer) {
+        writer.addDependency(AwsPythonDependency.SMITHY_AWS_CORE.withOptionalDependencies("xml"));
+        writer.addImport("smithy_aws_core.aio.protocols", "AwsQueryClientProtocol");
+        var service = context.settings().service(context.model());
+        var serviceSymbol = context.symbolProvider().toSymbol(service);
+        var serviceSchema = serviceSymbol.expectProperty(SymbolProperties.SCHEMA);
+        var version = service.getVersion();
+        writer.write("AwsQueryClientProtocol($T, $S)", serviceSchema, version);
+    }
+
+    @Override
+    public void generateProtocolTests(GenerationContext context) {
+        context.writerDelegator()
+                .useFileWriter("./tests/test_awsquery_protocol.py", "tests.test_awsquery_protocol", writer -> {
+                    new HttpProtocolTestGenerator(
+                            context,
+                            getProtocol(),
+                            writer,
+                            (shape, testCase) -> TESTS_TO_SKIP.contains(testCase.getId())).run();
+                });
+    }
+}

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

@@ -16,6 +16,7 @@
 import java.util.logging.Logger;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
+import software.amazon.smithy.aws.traits.auth.SigV4Trait;
 import software.amazon.smithy.codegen.core.CodegenException;
 import software.amazon.smithy.codegen.core.Symbol;
 import software.amazon.smithy.model.Model;
@@ -188,12 +189,14 @@ private void generateRequestTest(OperationShape operation, HttpRequestTestCase t
                                     endpoint_uri="https://$L/$L",
                                     transport = $T(),
                                     retry_strategy=SimpleRetryStrategy(max_attempts=1),
+                                    ${C|}
                                 )
                                 """,
                                 CodegenUtils.getConfigSymbol(context.settings()),
                                 host,
                                 path,
-                                REQUEST_TEST_ASYNC_HTTP_CLIENT_SYMBOL);
+                                REQUEST_TEST_ASYNC_HTTP_CLIENT_SYMBOL,
+                                (Runnable) this::writeSigV4TestConfig);
                     }));
 
                     // Generate the input using the expected shape and params
@@ -418,6 +421,16 @@ private void compareMediaBlob(HttpMessageTestCase testCase, PythonWriter writer)
                     """);
             return;
         }
+        if (contentType.equals("application/x-www-form-urlencoded")) {
+            writer.addStdlibImport("urllib.parse", "parse_qsl");
+            writer.write("""
+                    actual_params = sorted(parse_qsl(actual_body_content.decode()))
+                    expected_params = sorted(parse_qsl(expected_body_content.decode()))
+                    assert actual_params == expected_params
+
+                    """);
+            return;
+        }
         writer.write("assert actual_body_content == expected_body_content\n");
     }
 
@@ -437,13 +450,15 @@ private void generateResponseTest(OperationShape operation, HttpResponseTestCase
                                         headers=$J,
                                         body=b$S,
                                     ),
+                                    ${C|}
                                 )
                                 """,
                                 CodegenUtils.getConfigSymbol(context.settings()),
                                 RESPONSE_TEST_ASYNC_HTTP_CLIENT_SYMBOL,
                                 testCase.getCode(),
                                 CodegenUtils.toTuples(testCase.getHeaders()),
-                                testCase.getBody().filter(body -> !body.isEmpty()).orElse(""));
+                                testCase.getBody().filter(body -> !body.isEmpty()).orElse(""),
+                                (Runnable) this::writeSigV4TestConfig);
                     }));
                     // Create an empty input object to pass
                     var inputShape = model.expectShape(operation.getInputShape(), StructureShape.class);
@@ -490,13 +505,15 @@ private void generateErrorResponseTest(
                                         headers=$J,
                                         body=b$S,
                                     ),
+                                    ${C|}
                                 )
                                 """,
                                 CodegenUtils.getConfigSymbol(context.settings()),
                                 RESPONSE_TEST_ASYNC_HTTP_CLIENT_SYMBOL,
                                 testCase.getCode(),
                                 CodegenUtils.toTuples(testCase.getHeaders()),
-                                testCase.getBody().orElse(""));
+                                testCase.getBody().orElse(""),
+                                (Runnable) this::writeSigV4TestConfig);
                     }));
                     // Create an empty input object to pass
                     var inputShape = model.expectShape(operation.getInputShape(), StructureShape.class);
@@ -607,6 +624,19 @@ private void writeClientBlock(
         });
     }
 
+    private void writeSigV4TestConfig() {
+        if (!service.hasTrait(SigV4Trait.class)) {
+            return;
+        }
+        writer.addImport("smithy_aws_core.identity", "StaticCredentialsResolver");
+        writer.write("""
+                region="us-east-1",
+                aws_access_key_id="test-access-key-id",
+                aws_secret_access_key="test-secret-access-key",
+                aws_credentials_identity_resolver=StaticCredentialsResolver(),
+                """);
+    }
+
     private void writeUtilStubs(Symbol serviceSymbol) {
         LOGGER.fine(String.format("Writing utility stubs for %s : %s", serviceSymbol.getName(), protocol.getName()));
         writer.addDependency(SmithyPythonDependency.SMITHY_CORE);

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

@@ -58,6 +58,9 @@ public void run() {
                         }),
                         effective_auth_schemes = [
                             $8C
+                        ],
+                        error_schemas = [
+                            $9C
                         ]
                 )
                 """,
@@ -68,7 +71,8 @@ public void run() {
                 inSymbol.expectProperty(SymbolProperties.SCHEMA),
                 outSymbol.expectProperty(SymbolProperties.SCHEMA),
                 writer.consumer(this::writeErrorTypeRegistry),
-                writer.consumer(this::writeAuthSchemes));
+                writer.consumer(this::writeAuthSchemes),
+                writer.consumer(this::writeErrorSchemas));
     }
 
     private void writeErrorTypeRegistry(PythonWriter writer) {
@@ -82,6 +86,13 @@ private void writeErrorTypeRegistry(PythonWriter writer) {
         }
     }
 
+    private void writeErrorSchemas(PythonWriter writer) {
+        for (var error : shape.getErrors()) {
+            var errSymbol = symbolProvider.toSymbol(model.expectShape(error));
+            writer.write("$T,", errSymbol.expectProperty(SymbolProperties.SCHEMA));
+        }
+    }
+
     private void writeAuthSchemes(PythonWriter writer) {
         var authSchemes = ServiceIndex.of(model)
                 .getEffectiveAuthSchemes(context.settings().service(),

codegen/protocol-test/build.gradle.kts

@@ -30,5 +30,6 @@ repositories {
 
 dependencies {
     implementation(project(":core"))
+    implementation(project(":aws:core"))
     implementation(libs.smithy.aws.protocol.tests)
 }

codegen/protocol-test/smithy-build.json

@@ -22,6 +22,28 @@
                     "moduleVersion": "0.0.1"
                 }
             }
+        },
+        "aws-query": {
+            "transforms": [
+                {
+                    "name": "includeServices",
+                    "args": {
+                        "services": [
+                            "aws.protocoltests.query#AwsQuery"
+                        ]
+                    }
+                },
+                {
+                    "name": "removeUnusedShapes"
+                }
+            ],
+            "plugins": {
+                "python-client-codegen": {
+                    "service": "aws.protocoltests.query#AwsQuery",
+                    "module": "awsquery",
+                    "moduleVersion": "0.0.1"
+                }
+            }
         }
     }
 }

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

@@ -0,0 +1,4 @@
+{
+  "type": "feature",
+  "description": "Add `awsQuery` protocol support for Smithy clients."
+}

packages/smithy-aws-core/pyproject.toml

@@ -50,6 +50,9 @@ eventstream = [
 json = [
     "smithy-json~=0.2.0"
 ]
+xml = [
+    "smithy-xml~=0.1.0"
+]
 
 [tool.hatch.build]
 exclude = [

packages/smithy-aws-core/src/smithy_aws_core/_private/__init__.py

@@ -0,0 +1,2 @@
+#  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#  SPDX-License-Identifier: Apache-2.0

packages/smithy-aws-core/src/smithy_aws_core/_private/query/__init__.py

@@ -0,0 +1,10 @@
+#  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#  SPDX-License-Identifier: Apache-2.0
+
+from .errors import create_aws_query_error
+from .serializers import QueryShapeSerializer
+
+__all__ = (
+    "QueryShapeSerializer",
+    "create_aws_query_error",
+)