fix(core/protocols): $unknown union member support (#7593)

kuhe · web-flow · commit 596fc405be9f · 2025-12-18T18:54:17.000-05:00
diff --git a/packages/core/src/submodules/protocols/UnionSerde.ts b/packages/core/src/submodules/protocols/UnionSerde.ts
@@ -0,0 +1,38 @@
+/**
+ * Helper for identifying unknown union members during deserialization.
+ */
+export class UnionSerde {
+  private keys: Set<string>;
+
+  public constructor(private from: any, private to: any) {
+    this.keys = new Set(Object.keys(this.from).filter((k) => k !== "__type"));
+  }
+
+  /**
+   * Marks the key as being a known member.
+   * @param key - to mark.
+   */
+  public mark(key: string): void {
+    this.keys.delete(key);
+  }
+
+  /**
+   * @returns whether only one key remains unmarked and nothing has been written,
+   * implying the object is a union.
+   */
+  public hasUnknown(): boolean {
+    return this.keys.size === 1 && Object.keys(this.to).length === 0;
+  }
+
+  /**
+   * Writes the unknown key-value pair, if present, into the $unknown property
+   * of the union object.
+   */
+  public writeUnknown(): void {
+    if (this.hasUnknown()) {
+      const k = this.keys.values().next().value as string;
+      const v = this.from[k];
+      this.to.$unknown = [k, v];
+    }
+  }
+}
diff --git a/packages/core/src/submodules/protocols/json/JsonShapeDeserializer.spec.ts b/packages/core/src/submodules/protocols/json/JsonShapeDeserializer.spec.ts
@@ -2,7 +2,7 @@ import { NumericValue } from "@smithy/core/serde";
 import type { TimestampEpochSecondsSchema } from "@smithy/types";
 import { describe, expect, test as it } from "vitest";
 
-import { createNestingWidget, nestingWidget, widget } from "../test-schema.spec";
+import { createNestingWidget, nestingWidget, unionStruct, unionStructControl, widget } from "../test-schema.spec";
 import { JsonShapeDeserializer } from "./JsonShapeDeserializer";
 import { JsonShapeSerializer } from "./JsonShapeSerializer";
 
@@ -155,6 +155,30 @@ describe(JsonShapeDeserializer.name, () => {
     expect(await deserializer.read(widget, JSON.stringify({ scalar: "NaN" }))).toEqual({ scalar: NaN });
   });
 
+  it("deserializes $unknown union members", async () => {
+    const json = `{"union":{"unknownKey":{"timestamp":0,"blob":"AAECAw=="}}}`;
+    {
+      const deserialization = await deserializer.read(unionStruct, json);
+      expect(deserialization).toEqual({
+        union: {
+          $unknown: [
+            "unknownKey",
+            {
+              blob: "AAECAw==",
+              timestamp: 0,
+            },
+          ],
+        },
+      });
+    }
+    {
+      const deserialization = await deserializer.read(unionStructControl, json);
+      expect(deserialization).toEqual({
+        union: {},
+      });
+    }
+  });
+
   describe("performance baseline indicator", () => {
     const serializer = new JsonShapeSerializer({
       jsonName: true,
diff --git a/packages/core/src/submodules/protocols/json/JsonShapeDeserializer.ts b/packages/core/src/submodules/protocols/json/JsonShapeDeserializer.ts
@@ -19,6 +19,7 @@ import { fromBase64 } from "@smithy/util-base64";
 
 import { SerdeContextConfig } from "../ConfigurableSerdeContext";
 import { deserializingStructIterator } from "../structIterator";
+import { UnionSerde } from "../UnionSerde";
 import { JsonSettings } from "./JsonCodec";
 import { jsonReviver } from "./jsonReviver";
 import { parseJsonBody } from "./parseJsonBody";
@@ -49,18 +50,28 @@ export class JsonShapeDeserializer extends SerdeContextConfig implements ShapeDe
 
     if (isObject) {
       if (ns.isStructSchema()) {
+        const union = ns.isUnionSchema();
         const out = {} as any;
+        let unionSerde: UnionSerde;
+        if (union) {
+          unionSerde = new UnionSerde(value, out);
+        }
         for (const [memberName, memberSchema] of deserializingStructIterator(
           ns,
           value,
           this.settings.jsonName ? "jsonName" : false
         )) {
           const fromKey = this.settings.jsonName ? memberSchema.getMergedTraits().jsonName ?? memberName : memberName;
-          const deserializedValue = this._read(memberSchema, (value as any)[fromKey]);
-          if (deserializedValue != null) {
-            out[memberName] = deserializedValue;
+          if (union) {
+            unionSerde!.mark(fromKey);
+          }
+          if ((value as any)[fromKey] != null) {
+            out[memberName] = this._read(memberSchema, (value as any)[fromKey]);
           }
         }
+        if (union) {
+          unionSerde!.writeUnknown();
+        }
         return out;
       }
       if (Array.isArray(value) && ns.isListSchema()) {
diff --git a/packages/core/src/submodules/protocols/json/JsonShapeSerializer.spec.ts b/packages/core/src/submodules/protocols/json/JsonShapeSerializer.spec.ts
@@ -2,7 +2,7 @@ import { NumericValue } from "@smithy/core/serde";
 import type { TimestampEpochSecondsSchema } from "@smithy/types";
 import { describe, expect, test as it } from "vitest";
 
-import { createNestingWidget, nestingWidget, widget } from "../test-schema.spec";
+import { createNestingWidget, nestingWidget, unionStruct, widget } from "../test-schema.spec";
 import { SinglePassJsonShapeSerializer } from "./experimental/SinglePassJsonShapeSerializer";
 import { JsonShapeSerializer } from "./JsonShapeSerializer";
 
@@ -31,6 +31,22 @@ describe(JsonShapeSerializer.name, () => {
     );
   });
 
+  it("serializes $unknown union members", () => {
+    serializer1.write(unionStruct, {
+      union: {
+        $unknown: [
+          "unknownKey",
+          {
+            timestamp: new Date(0),
+            blob: new Uint8Array([0, 1, 2, 3]),
+          },
+        ],
+      },
+    });
+    const serialization = serializer1.flush();
+    expect(serialization).toEqual(`{"union":{"unknownKey":{"timestamp":0,"blob":"AAECAw=="}}}`);
+  });
+
   describe("performance baseline indicator", () => {
     for (const serializer of [serializer1, serializer2]) {
       it("should serialize objects", () => {
diff --git a/packages/core/src/submodules/protocols/json/JsonShapeSerializer.ts b/packages/core/src/submodules/protocols/json/JsonShapeSerializer.ts
@@ -2,6 +2,7 @@ import { determineTimestampFormat } from "@smithy/core/protocols";
 import { NormalizedSchema } from "@smithy/core/schema";
 import { dateToUtcString, generateIdempotencyToken, LazyJsonString, NumericValue } from "@smithy/core/serde";
 import type {
+  DocumentSchema,
   Schema,
   ShapeSerializer,
   TimestampDateTimeSchema,
@@ -81,6 +82,13 @@ export class JsonShapeSerializer extends SerdeContextConfig implements ShapeSeri
             out[targetKey] = serializableValue;
           }
         }
+        if (ns.isUnionSchema() && Object.keys(out).length === 0) {
+          const { $unknown } = value as any;
+          if (Array.isArray($unknown)) {
+            const [k, v] = $unknown;
+            out[k] = this._write(15 satisfies DocumentSchema, v);
+          }
+        }
         return out;
       }
 
diff --git a/packages/core/src/submodules/protocols/json/experimental/SinglePassJsonShapeSerializer.ts b/packages/core/src/submodules/protocols/json/experimental/SinglePassJsonShapeSerializer.ts
@@ -2,6 +2,7 @@ import { determineTimestampFormat } from "@smithy/core/protocols";
 import { NormalizedSchema } from "@smithy/core/schema";
 import { dateToUtcString, generateIdempotencyToken, LazyJsonString, NumericValue } from "@smithy/core/serde";
 import type {
+  DocumentSchema,
   Schema,
   ShapeSerializer,
   TimestampDateTimeSchema,
@@ -73,15 +74,24 @@ export class SinglePassJsonShapeSerializer extends SerdeContextConfig implements
       }
     } else if (ns.isStructSchema()) {
       b += "{";
+      let didWriteMember = false;
       for (const [name, member] of serializingStructIterator(ns, value)) {
         const item = (value as any)[name];
         const targetKey = this.settings.jsonName ? member.getMergedTraits().jsonName ?? name : name;
         const serializableValue = this.writeValue(member, item);
         if (item != null || member.isIdempotencyToken()) {
+          didWriteMember = true;
           b += `"${targetKey}":${serializableValue}`;
           b += ",";
         }
       }
+      if (!didWriteMember && ns.isUnionSchema()) {
+        const { $unknown } = value as any;
+        if (Array.isArray($unknown)) {
+          const [k, v] = $unknown;
+          b += `"${k}":${this.writeValue(15 satisfies DocumentSchema, v)}`;
+        }
+      }
     } else if (ns.isMapSchema() || ns.isDocumentSchema()) {
       b += "{";
       for (const [k, v] of Object.entries(value)) {
diff --git a/packages/core/src/submodules/protocols/query/QueryShapeSerializer.spec.ts b/packages/core/src/submodules/protocols/query/QueryShapeSerializer.spec.ts
@@ -2,20 +2,15 @@ import { NumericValue } from "@smithy/core/serde";
 import type { TimestampDateTimeSchema } from "@smithy/types";
 import { describe, expect, test as it } from "vitest";
 
-import { widget } from "../test-schema.spec";
+import { unionStruct, unionStructControl, widget } from "../test-schema.spec";
 import { QueryShapeSerializer } from "./QueryShapeSerializer";
 
 describe(QueryShapeSerializer.name, () => {
-  it("serializes data to Query", async () => {
-    const serializer = new QueryShapeSerializer({
-      timestampFormat: { default: 5 satisfies TimestampDateTimeSchema, useTrait: true },
-    });
-    serializer.setSerdeContext({
-      base64Encoder: (input: Uint8Array) => {
-        return Buffer.from(input).toString("base64");
-      },
-    } as any);
+  const serializer = new QueryShapeSerializer({
+    timestampFormat: { default: 5 satisfies TimestampDateTimeSchema, useTrait: true },
+  });
 
+  it("serializes data to Query", async () => {
     const data = {
       timestamp: new Date(0),
       bigint: 10000000000000000000000054321n,
@@ -28,4 +23,41 @@ describe(QueryShapeSerializer.name, () => {
       `&blob=AAAAAQ%3D%3D&timestamp=0&bigint=10000000000000000000000054321&bigdecimal=0.10000000000000000000000054321`
     );
   });
+
+  it("serializes $unknown union members", () => {
+    {
+      serializer.write(unionStruct, {
+        union: {
+          $unknown: [
+            "unknownKey",
+            {
+              timestamp: new Date(0),
+              blob: new Uint8Array([0, 1, 2, 3]),
+            },
+          ],
+        },
+      });
+      const serialization = String(serializer.flush()).replaceAll(/&/g, "\n&");
+      expect(serialization).toEqual(`
+&union.unknownKey.entry.1.key=timestamp
+&union.unknownKey.entry.1.value=1970-01-01T00%3A00%3A00Z
+&union.unknownKey.entry.2.key=blob
+&union.unknownKey.entry.2.value=AAECAw%3D%3D`);
+    }
+    {
+      serializer.write(unionStructControl, {
+        union: {
+          $unknown: [
+            "unknownKey",
+            {
+              timestamp: new Date(0),
+              blob: new Uint8Array([0, 1, 2, 3]),
+            },
+          ],
+        },
+      });
+      const serialization = serializer.flush();
+      expect(serialization).toEqual(``);
+    }
+  });
 });
diff --git a/packages/core/src/submodules/protocols/query/QueryShapeSerializer.ts b/packages/core/src/submodules/protocols/query/QueryShapeSerializer.ts
@@ -3,9 +3,15 @@ import { NormalizedSchema } from "@smithy/core/schema";
 import { generateIdempotencyToken, NumericValue } from "@smithy/core/serde";
 import { dateToUtcString } from "@smithy/smithy-client";
 import type {
+  BlobSchema,
+  DocumentSchema,
+  ListSchemaModifier,
+  MapSchemaModifier,
   Schema,
   ShapeSerializer,
+  StringSchema,
   TimestampDateTimeSchema,
+  TimestampDefaultSchema,
   TimestampEpochSecondsSchema,
   TimestampHttpDateSchema,
 } from "@smithy/types";
@@ -74,7 +80,18 @@ export class QueryShapeSerializer extends SerdeContextConfig implements ShapeSer
         }
       }
     } else if (ns.isDocumentSchema()) {
-      throw new Error(`@aws-sdk/core/protocols - QuerySerializer unsupported document type ${ns.getName(true)}`);
+      if (Array.isArray(value)) {
+        this.write((64 satisfies ListSchemaModifier) | (15 satisfies DocumentSchema), value, prefix);
+      } else if (value instanceof Date) {
+        this.write(4 satisfies TimestampDefaultSchema, value, prefix);
+      } else if (value instanceof Uint8Array) {
+        this.write(21 satisfies BlobSchema, value, prefix);
+      } else if (value && typeof value === "object") {
+        this.write((128 satisfies MapSchemaModifier) | (15 satisfies DocumentSchema), value, prefix);
+      } else {
+        this.writeKey(prefix);
+        this.writeValue(String(value));
+      }
     } else if (ns.isListSchema()) {
       if (Array.isArray(value)) {
         if (value.length === 0) {
@@ -120,13 +137,23 @@ export class QueryShapeSerializer extends SerdeContextConfig implements ShapeSer
       }
     } else if (ns.isStructSchema()) {
       if (value && typeof value === "object") {
+        let didWriteMember = false;
         for (const [memberName, member] of serializingStructIterator(ns, value)) {
           if ((value as any)[memberName] == null && !member.isIdempotencyToken()) {
             continue;
           }
           const suffix = this.getKey(memberName, member.getMergedTraits().xmlName);
           const key = `${prefix}${suffix}`;
           this.write(member, (value as any)[memberName], key);
+          didWriteMember = true;
+        }
+        if (!didWriteMember && ns.isUnionSchema()) {
+          const { $unknown } = value as any;
+          if (Array.isArray($unknown)) {
+            const [k, v] = $unknown;
+            const key = `${prefix}${k}`;
+            this.write(15 satisfies DocumentSchema, v, key);
+          }
         }
       }
     } else if (ns.isUnitSchema()) {
diff --git a/packages/core/src/submodules/protocols/test-schema.spec.ts b/packages/core/src/submodules/protocols/test-schema.spec.ts
@@ -7,6 +7,7 @@ import type {
   StaticListSchema,
   StaticOperationSchema,
   StaticStructureSchema,
+  StaticUnionSchema,
   StringSchema,
   TimestampDefaultSchema,
   TimestampEpochSecondsSchema,
@@ -122,3 +123,21 @@ export const context = {
     };
   },
 } as any;
+
+export const unionStruct = [
+  3,
+  "ns",
+  "UnionStruct",
+  0,
+  ["union"],
+  [[4, "ns", "Union", 0, ["string", "timestamp", "blob"], [0, 7, 21]] satisfies StaticUnionSchema],
+] satisfies StaticStructureSchema;
+
+export const unionStructControl = [
+  3,
+  "ns",
+  "UnionStruct",
+  0,
+  ["union"],
+  [[3, "ns", "Union", 0, ["string", "timestamp", "blob"], [0, 7, 21]] satisfies StaticStructureSchema],
+] satisfies StaticStructureSchema;
diff --git a/packages/core/src/submodules/protocols/xml/XmlShapeDeserializer.spec.ts b/packages/core/src/submodules/protocols/xml/XmlShapeDeserializer.spec.ts
diff --git a/packages/core/src/submodules/protocols/xml/XmlShapeDeserializer.ts b/packages/core/src/submodules/protocols/xml/XmlShapeDeserializer.ts
diff --git a/packages/core/src/submodules/protocols/xml/XmlShapeSerializer.spec.ts b/packages/core/src/submodules/protocols/xml/XmlShapeSerializer.spec.ts
diff --git a/packages/core/src/submodules/protocols/xml/XmlShapeSerializer.ts b/packages/core/src/submodules/protocols/xml/XmlShapeSerializer.ts
