fix(util-dynamodb): allow marshall function to handle more input types (#3539)

* fix(util-dynamodb): allow marshall function to handle more input types

* fix(util-dynamodb): revert marshall.spec.ts and minor formatting

* fix(util-dynamodb): code spacing

Author: kuhe

lib/lib-dynamodb/src/commands/marshallInput.spec.ts

@@ -0,0 +1,92 @@
+import { marshallInput } from "./utils";
+
+describe("marshallInput and processObj", () => {
+  it("marshallInput should not ignore falsy values", () => {
+    expect(marshallInput({ Items: [0, false, null, ""] }, [{ key: "Items" }])).toEqual({
+      Items: [{ N: "0" }, { BOOL: false }, { NULL: true }, { S: "" }],
+    });
+  });
+});
+
+describe("marshallInput for commands", () => {
+  it("marshals QueryCommand input", () => {
+    const input = {
+      TableName: "TestTable",
+      KeyConditions: {
+        id: {
+          AttributeValueList: ["test"],
+          ComparisonOperator: "EQ",
+        },
+      },
+    };
+    const inputKeyNodes = [
+      {
+        key: "KeyConditions",
+        children: {
+          children: [{ key: "AttributeValueList" }],
+        },
+      },
+      {
+        key: "QueryFilter",
+        children: {
+          children: [{ key: "AttributeValueList" }],
+        },
+      },
+      { key: "ExclusiveStartKey" },
+      { key: "ExpressionAttributeValues" },
+    ];
+    const output = {
+      TableName: "TestTable",
+      KeyConditions: { id: { AttributeValueList: [{ S: "test" }], ComparisonOperator: "EQ" } },
+      QueryFilter: undefined,
+      ExclusiveStartKey: undefined,
+      ExpressionAttributeValues: undefined,
+    };
+    expect(marshallInput(input, inputKeyNodes)).toEqual(output);
+  });
+  it("marshals ExecuteStatementCommand input", () => {
+    const input = {
+      Statement: `SELECT col_1
+                        FROM some_table
+                        WHERE contains("col_1", ?)`,
+      Parameters: ["some_param"],
+    };
+    const inputKeyNodes = [{ key: "Parameters" }];
+    const output = {
+      Statement: input.Statement,
+      Parameters: [{ S: "some_param" }],
+    };
+    expect(marshallInput(input, inputKeyNodes)).toEqual(output);
+  });
+  it("marshals BatchExecuteStatementCommand input", () => {
+    const input = {
+      Statements: [
+        {
+          Statement: `
+                        UPDATE "table"
+                        SET field1=?
+                        WHERE field2 = ?
+                          AND field3 = ?
+                    `,
+          Parameters: [false, "field 2 value", 1234],
+        },
+      ],
+    };
+    const inputKeyNodes = [{ key: "Statements", children: [{ key: "Parameters" }] }];
+    const output = {
+      Statements: [
+        {
+          Statement: input.Statements[0].Statement,
+          Parameters: [
+            {
+              BOOL: false,
+            },
+            { S: "field 2 value" },
+            { N: "1234" },
+          ],
+        },
+      ],
+    };
+    expect(marshallInput(input, inputKeyNodes)).toEqual(output);
+  });
+});

lib/lib-dynamodb/src/commands/utils.ts

@@ -10,7 +10,7 @@ export type AllNodes = {
 };
 
 const processObj = (obj: any, processFunc: Function, children?: KeyNode[] | AllNodes): any => {
-  if (obj) {
+  if (obj !== undefined) {
     if (!children || (Array.isArray(children) && children.length === 0)) {
       // Leaf of KeyNode, process the object.
       return processFunc(obj);

packages/util-dynamodb/src/marshall.ts

@@ -1,7 +1,7 @@
 import { AttributeValue } from "@aws-sdk/client-dynamodb";
 
 import { convertToAttr } from "./convertToAttr";
-import { NativeAttributeValue } from "./models";
+import { NativeAttributeBinary, NativeAttributeValue } from "./models";
 
 /**
  * An optional configuration object for `marshall`
@@ -26,8 +26,39 @@ export interface marshallOptions {
  *
  * @param {any} data - The data to convert to a DynamoDB record
  * @param {marshallOptions} options - An optional configuration object for `marshall`
+ *
  */
-export const marshall = <T extends { [K in keyof T]: NativeAttributeValue }>(
-  data: T,
+export function marshall(data: Set<string>, options?: marshallOptions): AttributeValue.SSMember;
+export function marshall(data: Set<number>, options?: marshallOptions): AttributeValue.NSMember;
+export function marshall(data: Set<NativeAttributeBinary>, options?: marshallOptions): AttributeValue.BSMember;
+export function marshall<M extends { [K in keyof M]: NativeAttributeValue }>(
+  data: M,
   options?: marshallOptions
-): { [key: string]: AttributeValue } => convertToAttr(data, options).M as { [key: string]: AttributeValue };
+): Record<string, AttributeValue>;
+export function marshall<L extends NativeAttributeValue[]>(data: L, options?: marshallOptions): AttributeValue[];
+export function marshall(data: string, options?: marshallOptions): AttributeValue.SMember;
+export function marshall(data: number, options?: marshallOptions): AttributeValue.NMember;
+export function marshall(data: NativeAttributeBinary, options?: marshallOptions): AttributeValue.BMember;
+export function marshall(data: null, options?: marshallOptions): AttributeValue.NULLMember;
+export function marshall(data: boolean, options?: marshallOptions): AttributeValue.BOOLMember;
+export function marshall(data: unknown, options?: marshallOptions): AttributeValue.$UnknownMember;
+export function marshall(data: unknown, options?: marshallOptions) {
+  const attributeValue: AttributeValue = convertToAttr(data, options);
+  const [key, value] = Object.entries(attributeValue)[0];
+  switch (key) {
+    case "M":
+    case "L":
+      return value;
+    case "SS":
+    case "NS":
+    case "BS":
+    case "S":
+    case "N":
+    case "B":
+    case "NULL":
+    case "BOOL":
+    case "$unknown":
+    default:
+      return attributeValue;
+  }
+}

packages/util-dynamodb/src/marshallTypes.spec.ts

@@ -0,0 +1,67 @@
+import { convertToAttr } from "./convertToAttr";
+import { marshall } from "./marshall";
+
+describe("marshall type discernment", () => {
+  describe("behaves as convertToAttr for non-collection values or Sets", () => {
+    it("marshals string", () => {
+      const value = "hello";
+      expect(marshall(value)).toEqual(convertToAttr(value));
+    });
+
+    it("marshals number", () => {
+      const value = 1578;
+      expect(marshall(value)).toEqual(convertToAttr(value));
+    });
+
+    it("marshals binary", () => {
+      const value = new Uint8Array([0, 1, 0, 1]);
+      expect(marshall(value)).toEqual(convertToAttr(value));
+    });
+
+    it("marshals boolean", () => {
+      let value = false;
+      expect(marshall(value)).toEqual(convertToAttr(value));
+      value = true;
+      expect(marshall(value)).toEqual(convertToAttr(value));
+    });
+
+    it("marshals null", () => {
+      const value = null;
+      expect(marshall(value)).toEqual(convertToAttr(value));
+    });
+    it("marshals string set", () => {
+      const value = new Set(["a", "b"]);
+      expect(marshall(value)).toEqual(convertToAttr(value));
+    });
+
+    it("marshals number set", () => {
+      const value = new Set([1, 2]);
+      expect(marshall(value)).toEqual(convertToAttr(value));
+    });
+
+    it("marshals binary set", () => {
+      const value = new Set([new Uint8Array([1, 0]), new Uint8Array([0, 1])]);
+      expect(marshall(value)).toEqual(convertToAttr(value));
+    });
+  });
+
+  describe("unwraps one level for input data which are lists or maps", () => {
+    it("marshals and unwraps map", () => {
+      expect(marshall({ a: 1, b: { a: 2, b: [1, 2, 3] } })).toEqual({
+        a: { N: "1" },
+        b: {
+          M: {
+            a: { N: "2" },
+            b: {
+              L: [{ N: "1" }, { N: "2" }, { N: "3" }],
+            },
+          },
+        },
+      });
+    });
+
+    it("marshals and unwraps list", () => {
+      expect(marshall(["test", 2, null])).toEqual([{ S: "test" }, { N: "2" }, { NULL: true }]);
+    });
+  });
+});