Attach response metadata to operation outputs and errors, escaping collisions

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

@@ -55,6 +55,18 @@ public final class CodegenUtils {
      */
     public static final int MAX_PREFERRED_LINE_LENGTH = 88;
 
+    /**
+     * The name of the attribute carrying response metadata on operation outputs
+     * and modeled errors.
+     *
+     * <p>Declared on {@code CallError} for errors and emitted onto operation
+     * outputs by the structure generator. Because a service is free to model a
+     * member of the same name, the name is reserved on the members of those
+     * shapes and escaped if it collides. Both the emitter and the escaper read it
+     * from here so that they cannot disagree.
+     */
+    public static final String RESPONSE_METADATA_MEMBER = "response_metadata";
+
     static final Set<String> ERROR_MESSAGE_MEMBER_NAMES = SetUtils.of(
             "errormessage",
             "error_message",

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

@@ -16,6 +16,7 @@
 import software.amazon.smithy.codegen.core.SymbolProvider;
 import software.amazon.smithy.codegen.core.SymbolReference;
 import software.amazon.smithy.model.Model;
+import software.amazon.smithy.model.knowledge.OperationIndex;
 import software.amazon.smithy.model.loader.Prelude;
 import software.amazon.smithy.model.shapes.BigDecimalShape;
 import software.amazon.smithy.model.shapes.BigIntegerShape;
@@ -74,6 +75,7 @@ public final class PythonSymbolProvider implements SymbolProvider, ShapeVisitor<
     private final PythonSettings settings;
     private final ServiceShape service;
     private final Set<String> allShapeNames;
+    private final OperationIndex operationIndex;
 
     public PythonSymbolProvider(Model model, PythonSettings settings) {
         this.model = model;
@@ -107,6 +109,9 @@ public PythonSymbolProvider(Model model, PythonSettings settings) {
         // Collect all shape names that will be generated as PascalCase classes in models.py.
         // Used to detect collisions with synthesized names (union variants, unknown types).
         this.allShapeNames = collectAllShapeNames();
+
+        // Built once because toMemberName runs for every member of every shape.
+        this.operationIndex = OperationIndex.of(model);
     }
 
     /**
@@ -149,12 +154,38 @@ public String toMemberName(MemberShape shape) {
         }
 
         var container = model.expectShape(shape.getContainer());
+
+        // The response metadata attribute is added to operation outputs and errors
+        // by the generator, so a modeled member of the same name is escaped to keep
+        // it from shadowing the attribute.
+        if (CodegenUtils.RESPONSE_METADATA_MEMBER.equals(memberName)
+                && carriesResponseMetadata(container)) {
+            memberName = escapeWord(memberName);
+            LOGGER.warning(() -> format(
+                    "Renamed member %s to \"%s\" because \"%s\" is reserved for response metadata.",
+                    shape.getId(),
+                    escapeWord(CodegenUtils.RESPONSE_METADATA_MEMBER),
+                    CodegenUtils.RESPONSE_METADATA_MEMBER));
+        }
+
         if (container.isEnumShape() || container.isIntEnumShape()) {
             memberName = memberName.toUpperCase(Locale.ENGLISH);
         }
         return memberName;
     }
 
+    /**
+     * Whether a shape is given a response metadata attribute, and so has a member
+     * name that must be kept clear.
+     *
+     * <p>Operation inputs and other structures are not given the attribute, so
+     * their members are left alone.
+     */
+    private boolean carriesResponseMetadata(Shape container) {
+        return container.hasTrait(ErrorTrait.class)
+                || operationIndex.isOutputStructure(container);
+    }
+
     private String getDefaultShapeName(Shape shape) {
         // Use the service-aliased name
         return StringUtils.capitalize(shape.getId().getName(service));

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

@@ -40,6 +40,12 @@ public final class RuntimeTypes {
     public static final Symbol TYPE_REGISTRY =
             createSymbol("documents", "TypeRegistry", SmithyPythonDependency.SMITHY_CORE);
 
+    // smithy_core.response
+    public static final Symbol RESPONSE_METADATA =
+            createSymbol("response", "ResponseMetadata", SmithyPythonDependency.SMITHY_CORE);
+    public static final Symbol EMPTY_RESPONSE_METADATA =
+            createSymbol("response", "EMPTY_RESPONSE_METADATA", SmithyPythonDependency.SMITHY_CORE);
+
     // smithy_core.exceptions
     public static final Symbol MODELED_ERROR =
             createSymbol("exceptions", "ModeledError", SmithyPythonDependency.SMITHY_CORE);

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

@@ -19,6 +19,7 @@
 import software.amazon.smithy.codegen.core.SymbolProvider;
 import software.amazon.smithy.model.Model;
 import software.amazon.smithy.model.knowledge.NullableIndex;
+import software.amazon.smithy.model.knowledge.OperationIndex;
 import software.amazon.smithy.model.node.Node;
 import software.amazon.smithy.model.shapes.MemberShape;
 import software.amazon.smithy.model.shapes.Shape;
@@ -104,6 +105,7 @@ private void renderStructure() {
                 class $L:
                     ${C|}
 
+                    ${C|}
                     ${C|}
 
                     ${C|}
@@ -116,6 +118,7 @@ class $L:
                 symbol.getName(),
                 writer.consumer(w -> writeClassDocs()),
                 writer.consumer(w -> writeProperties()),
+                writer.consumer(w -> writeResponseMetadataProperty()),
                 writer.consumer(w -> generateSerializeMethod()),
                 writer.consumer(w -> generateDeserializeMethod()),
                 writer.consumer(w -> generateSmithyDefaultMethod()));
@@ -177,6 +180,38 @@ private void writeClassDocs() {
         writer.writeDocs(docs, context);
     }
 
+    /**
+     * Writes the response metadata attribute onto operation outputs.
+     *
+     * <p>The attribute is not part of the service's modeled data, so it is excluded
+     * from equality and from the generated repr. Excluding it from equality also
+     * keeps generated protocol tests comparing shapes by their modeled members
+     * alone.
+     *
+     * <p>Errors receive the same attribute by inheriting it from the service error
+     * base class rather than having it written here.
+     */
+    private void writeResponseMetadataProperty() {
+        if (!OperationIndex.of(model).isOutputStructure(shape)) {
+            return;
+        }
+
+        writer.addStdlibImport("dataclasses", "field");
+        writer.write("""
+                $L: $T = field(default=$T, repr=False, compare=False)
+                $C
+                """,
+                CodegenUtils.RESPONSE_METADATA_MEMBER,
+                RuntimeTypes.RESPONSE_METADATA,
+                RuntimeTypes.EMPTY_RESPONSE_METADATA,
+                writer.consumer(w -> w.writeDocs("""
+                        Metadata about the response that produced this output.
+
+                        Use this to recover the request identifiers a service's support team \
+                        needs in order to investigate a call. Members of the metadata are \
+                        individually optional.""", context)));
+    }
+
     private void writeProperties() {
         for (MemberShape member : requiredMembers) {
             writer.pushState();

codegen/core/src/test/java/software/amazon/smithy/python/codegen/PythonSymbolProviderTest.java

@@ -108,6 +108,76 @@ public void testOperationNameCollidingWithClientMethodIsEscaped() {
                         .getName());
     }
 
+    @Test
+    public void testResponseMetadataIsEscapedOnOutputsAndErrors() {
+        Model model = loadModel(RESPONSE_METADATA_MODEL);
+        PythonSymbolProvider provider = createProvider(model);
+
+        assertEquals("response_metadata_", memberName(provider, model, "GetThingOutput$responseMetadata"));
+        assertEquals("response_metadata_", memberName(provider, model, "ThingError$responseMetadata"));
+    }
+
+    @Test
+    public void testResponseMetadataIsNotEscapedOnOtherShapes() {
+        // Only outputs and errors are given the attribute, so members elsewhere
+        // must keep their natural name.
+        Model model = loadModel(RESPONSE_METADATA_MODEL);
+        PythonSymbolProvider provider = createProvider(model);
+
+        assertEquals("response_metadata", memberName(provider, model, "GetThingInput$responseMetadata"));
+        assertEquals("response_metadata", memberName(provider, model, "Nested$responseMetadata"));
+    }
+
+    @Test
+    public void testUnrelatedMembersOnOutputsAndErrorsAreUnaffected() {
+        Model model = loadModel(RESPONSE_METADATA_MODEL);
+        PythonSymbolProvider provider = createProvider(model);
+
+        assertEquals("thing_arn", memberName(provider, model, "GetThingOutput$thingArn"));
+        assertEquals("retry_after", memberName(provider, model, "ThingError$retryAfter"));
+    }
+
+    private static final String RESPONSE_METADATA_MODEL = """
+            $version: "2"
+            namespace smithy.example
+
+            service TestService {
+                version: "2024-01-01"
+                operations: [GetThing]
+                errors: [ThingError]
+            }
+
+            operation GetThing {
+                input: GetThingInput
+                output: GetThingOutput
+            }
+
+            structure GetThingInput {
+                responseMetadata: String
+            }
+
+            structure GetThingOutput {
+                responseMetadata: String
+                thingArn: String
+                nested: Nested
+            }
+
+            structure Nested {
+                responseMetadata: String
+            }
+
+            @error("client")
+            structure ThingError {
+                responseMetadata: String
+                retryAfter: String
+            }
+            """;
+
+    private static String memberName(PythonSymbolProvider provider, Model model, String relativeId) {
+        var member = model.expectShape(ShapeId.from(NS + "#" + relativeId), MemberShape.class);
+        return provider.toMemberName(member);
+    }
+
     private static Model loadModel(String smithyIdl) {
         return Model.assembler().addUnparsedModel("test.smithy", smithyIdl).assemble().unwrap();
     }

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

@@ -0,0 +1,4 @@
+{
+  "type": "feature",
+  "description": "AWS protocols now report request identifiers on operation outputs and errors, read from the `x-amzn-requestid`, `x-amz-request-id`, and `x-amz-id-2` headers. awsQuery reads them from the response body."
+}

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

@@ -2,9 +2,11 @@
 #  SPDX-License-Identifier: Apache-2.0
 
 from .errors import create_aws_query_error
+from .metadata import parse_aws_query_request_id
 from .serializers import QueryShapeSerializer
 
 __all__ = (
     "QueryShapeSerializer",
     "create_aws_query_error",
+    "parse_aws_query_request_id",
 )

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

@@ -0,0 +1,20 @@
+#  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#  SPDX-License-Identifier: Apache-2.0
+"""Helpers for reading awsQuery response bodies that are not shape-modeled."""
+
+from xml.etree.ElementTree import Element
+
+
+def local_name(tag: str) -> str:
+    """Strip namespace URI from an element tag: {uri}local -> local."""
+    if tag.startswith("{"):
+        return tag.split("}", 1)[1]
+    return tag
+
+
+def find_child(element: Element, name: str) -> Element | None:
+    """Return the first child element whose local name matches ``name``."""
+    for child in element:
+        if local_name(child.tag) == name:
+            return child
+    return None

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

@@ -1,7 +1,7 @@
 #  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 #  SPDX-License-Identifier: Apache-2.0
 from typing import TYPE_CHECKING, Any
-from xml.etree.ElementTree import Element, ParseError, fromstring
+from xml.etree.ElementTree import ParseError, fromstring
 
 from smithy_core.documents import TypeRegistry
 from smithy_core.exceptions import (
@@ -15,6 +15,7 @@
 from smithy_core.shapes import ShapeID
 
 from ...traits import AwsQueryErrorTrait
+from ._xml import find_child, local_name
 
 try:
     from smithy_xml import XMLCodec
@@ -34,21 +35,6 @@ def _assert_xml() -> None:
         )
 
 
-def _local_name(tag: str) -> str:
-    """Strip namespace URI from an element tag: {uri}local -> local."""
-    if tag.startswith("{"):
-        return tag.split("}", 1)[1]
-    return tag
-
-
-def _find_child(element: Element, name: str) -> Element | None:
-    """Return the first child element whose local name matches ``name``."""
-    for child in element:
-        if _local_name(child.tag) == name:
-            return child
-    return None
-
-
 def _parse_aws_query_error_code(
     body: bytes, wrapper_elements: tuple[str, ...]
 ) -> str | None:
@@ -59,15 +45,15 @@ def _parse_aws_query_error_code(
         return None
 
     if wrapper_elements:
-        if _local_name(element.tag) != wrapper_elements[0]:
+        if local_name(element.tag) != wrapper_elements[0]:
             return None
         for wrapper in wrapper_elements[1:]:
-            next_element = _find_child(element, wrapper)
+            next_element = find_child(element, wrapper)
             if next_element is None:
                 return None
             element = next_element
 
-    code_element = _find_child(element, "Code")
+    code_element = find_child(element, "Code")
     return code_element.text if code_element is not None else None
 
 

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

@@ -0,0 +1,26 @@
+#  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#  SPDX-License-Identifier: Apache-2.0
+from xml.etree.ElementTree import ParseError, fromstring
+
+from ._xml import find_child
+
+
+def parse_aws_query_request_id(body: bytes) -> str | None:
+    """Parse the request ID from an awsQuery response body.
+
+    There is no request ID header. Successes nest it under ``ResponseMetadata``,
+    errors put it directly under the root.
+    """
+    try:
+        root = fromstring(body)  # noqa: S314
+    except ParseError:
+        return None
+
+    metadata = find_child(root, "ResponseMetadata")
+    if metadata is not None:
+        nested = find_child(metadata, "RequestId")
+        if nested is not None and nested.text:
+            return nested.text
+
+    direct = find_child(root, "RequestId")
+    return direct.text or None if direct is not None else None