From 72bca8babe519ac4aa9a4d9d1535e3f95d2bdfde Mon Sep 17 00:00:00 2001
From: George Fu <kuhe@users.noreply.github.com>
Date: Wed, 19 Jun 2024 17:08:45 +0000
Subject: [PATCH] awsquerycompat

---
 smithy-typescript-codegen/build.gradle.kts    |  1 +
 .../protocols/SmithyProtocolUtils.java        |  3 ++-
 .../protocols/cbor/SmithyRpcV2Cbor.java       | 25 ++++++++++++++++---
 3 files changed, 25 insertions(+), 4 deletions(-)

diff --git a/smithy-typescript-codegen/build.gradle.kts b/smithy-typescript-codegen/build.gradle.kts
index c4da63eb2e0..98db5703aad 100644
--- a/smithy-typescript-codegen/build.gradle.kts
+++ b/smithy-typescript-codegen/build.gradle.kts
@@ -43,6 +43,7 @@ dependencies {
     api("software.amazon.smithy:smithy-protocol-test-traits:$smithyVersion")
     api("software.amazon.smithy:smithy-rules-engine:$smithyVersion")
     api("software.amazon.smithy:smithy-waiters:$smithyVersion")
+    api("software.amazon.smithy:smithy-aws-traits:$smithyVersion")
 }
 
 sourceSets {
diff --git a/smithy-typescript-codegen/src/main/java/software/amazon/smithy/typescript/codegen/protocols/SmithyProtocolUtils.java b/smithy-typescript-codegen/src/main/java/software/amazon/smithy/typescript/codegen/protocols/SmithyProtocolUtils.java
index f4a0b4c72a2..8e688bf8013 100644
--- a/smithy-typescript-codegen/src/main/java/software/amazon/smithy/typescript/codegen/protocols/SmithyProtocolUtils.java
+++ b/smithy-typescript-codegen/src/main/java/software/amazon/smithy/typescript/codegen/protocols/SmithyProtocolUtils.java
@@ -69,7 +69,8 @@ private static boolean filterProtocolTests(
             return true;
         }
 
-        // TODO(cbor): verify the test is declared correctly.
+        // TODO(cbor): enable this test when upgrading to a Smithy version
+        // TODO(cbor): in which it is fixed.
         if (testCase.getId().equals("RpcV2CborDeserializesDenseSetMapAndSkipsNull")) {
             return true;
         }
diff --git a/smithy-typescript-codegen/src/main/java/software/amazon/smithy/typescript/codegen/protocols/cbor/SmithyRpcV2Cbor.java b/smithy-typescript-codegen/src/main/java/software/amazon/smithy/typescript/codegen/protocols/cbor/SmithyRpcV2Cbor.java
index e60209ff71a..0286e58fa5a 100644
--- a/smithy-typescript-codegen/src/main/java/software/amazon/smithy/typescript/codegen/protocols/cbor/SmithyRpcV2Cbor.java
+++ b/smithy-typescript-codegen/src/main/java/software/amazon/smithy/typescript/codegen/protocols/cbor/SmithyRpcV2Cbor.java
@@ -7,6 +7,7 @@
 
 import java.util.Set;
 import java.util.TreeSet;
+import software.amazon.smithy.aws.traits.protocols.AwsQueryCompatibleTrait;
 import software.amazon.smithy.codegen.core.Symbol;
 import software.amazon.smithy.codegen.core.SymbolProvider;
 import software.amazon.smithy.codegen.core.SymbolReference;
@@ -137,6 +138,20 @@ public void generateSharedComponents(GenerationContext context) {
         writeSharedRequestHeaders(context);
         writer.write("");
 
+        if (context.getService().hasTrait(AwsQueryCompatibleTrait.class)) {
+            writer.addImport("HeaderBag", "__HeaderBag", TypeScriptDependency.SMITHY_TYPES);
+            writer.write("""
+                const populateBodyWithQueryCompatibility = (parsedOutput: any, headers: __HeaderBag) => {
+                  const queryErrorHeader = headers["x-amzn-query-error"];
+                  if (parsedOutput.body !== undefined && queryErrorHeader != null) {
+                    const codeAndType = queryErrorHeader.split(";");
+                    parsedOutput.body.Code = codeAndType[0];
+                    parsedOutput.body.Type = codeAndType[1];
+                  }
+                };
+                """);
+        }
+
         writer.write(
             context.getStringStore().flushVariableDeclarationCode()
         );
@@ -269,9 +284,13 @@ protected void serializeInputDocument(GenerationContext generationContext,
     protected void writeErrorCodeParser(GenerationContext generationContext) {
         TypeScriptWriter writer = generationContext.getWriter();
 
-        // TODO(cbor) handle Query-Compatible
-        // TODO(cbor) either by accepting a function from downstream
-        // TODO(cbor) or natively here.
+        if (generationContext.getService().hasTrait(AwsQueryCompatibleTrait.class)) {
+            // Populate parsedOutput.body with 'Code' and 'Type' fields
+            // "x-amzn-query-error" header is available when AwsQueryCompatibleTrait is applied to a service
+            // The header value contains query error Code and Type joined by ';'
+            // E.g. "MalformedInput;Sender" or "InternalFailure;Receiver"
+            writer.write("populateBodyWithQueryCompatibility(parsedOutput, output.headers);");
+        }
 
         writer.addSubPathImport(
             "loadSmithyRpcV2CborErrorCode", null,