fix(lib-dynamodb): make command middleware useable, turn marshalling into middleware (#3808)

* fix(lib-dynamodb): make command middleware useable, turn marshalling into middleware

* fix(lib-dynamodb): fix unit test

* fix(lib-dynamodb): move inner command construction to ctor body

* fix(lib-dynamodb): remove command name field

* fix(lib-dynamodb): fix ctor code order

Author: kuhe

codegen/smithy-aws-typescript-codegen/src/main/java/software/amazon/smithy/aws/typescript/codegen/DocumentClientCommandGenerator.java

@@ -15,8 +15,6 @@
 
 package software.amazon.smithy.aws.typescript.codegen;
 
-import static software.amazon.smithy.aws.typescript.codegen.propertyaccess.PropertyAccessor.getFrom;
-
 import java.nio.file.Paths;
 import java.util.ArrayList;
 import java.util.HashSet;
@@ -55,10 +53,14 @@ final class DocumentClientCommandGenerator implements Runnable {
     private final TypeScriptWriter writer;
     private final Symbol symbol;
     private final OperationIndex operationIndex;
+    private final String originalInputTypeName;
     private final String inputTypeName;
     private final List<MemberShape> inputMembersWithAttr;
+    private final String originalOutputTypeName;
     private final String outputTypeName;
     private final List<MemberShape> outputMembersWithAttr;
+    private final String clientCommandClassName;
+    private final String clientCommandLocalName;
 
     DocumentClientCommandGenerator(
             TypeScriptSettings settings,
@@ -75,14 +77,21 @@ final class DocumentClientCommandGenerator implements Runnable {
 
         symbol = symbolProvider.toSymbol(operation);
         operationIndex = OperationIndex.of(model);
+        String inputType = symbol.expectProperty("inputType", Symbol.class).getName();
+        originalInputTypeName = inputType;
         inputTypeName = DocumentClientUtils.getModifiedName(
-            symbol.expectProperty("inputType", Symbol.class).getName()
+            inputType
         );
         inputMembersWithAttr = getStructureMembersWithAttr(operationIndex.getInput(operation));
+        String outputType = symbol.expectProperty("outputType", Symbol.class).getName();
+        originalOutputTypeName = outputType;
         outputTypeName = DocumentClientUtils.getModifiedName(
-            symbol.expectProperty("outputType", Symbol.class).getName()
+            outputType
         );
         outputMembersWithAttr = getStructureMembersWithAttr(operationIndex.getOutput(operation));
+
+        clientCommandClassName = symbol.getName();
+        clientCommandLocalName = "__" + clientCommandClassName;
     }
 
     @Override
@@ -92,44 +101,65 @@ public void run() {
 
         // Add required imports.
         writer.addImport(configType, configType, servicePath);
-        writer.addImport("Command", "$Command", "@aws-sdk/smithy-client");
+        writer.addImport(
+            "DynamoDBDocumentClientCommand",
+            "DynamoDBDocumentClientCommand",
+            "./baseCommand/DynamoDBDocumentClientCommand"
+        );
 
         generateInputAndOutputTypes();
 
+        String ioTypes = String.join(", ", new String[]{
+            inputTypeName,
+            outputTypeName,
+            "__" + originalInputTypeName,
+            "__" + originalOutputTypeName
+        });
+
         String name = DocumentClientUtils.getModifiedName(symbol.getName());
         writer.writeDocs(DocumentClientUtils.getCommandDocs(symbol.getName()));
-        writer.openBlock("export class $L extends $$Command<$L, $L, $L> {", "}",
-                name, inputTypeName, outputTypeName, configType, () -> {
-
-            // Section for adding custom command properties.
-            writer.pushState(COMMAND_PROPERTIES_SECTION);
-            if (!inputMembersWithAttr.isEmpty()) {
-                writer.openBlock("private readonly $L = [", "];", COMMAND_INPUT_KEYNODES, () -> {
+        writer.openBlock(
+            "export class $L extends DynamoDBDocumentClientCommand<" + ioTypes + ", $L> {",
+            "}",
+            name,
+            configType,
+            () -> {
+                // Section for adding custom command properties.
+                writer.pushState(COMMAND_PROPERTIES_SECTION);
+                writer.openBlock("protected readonly $L = [", "];", COMMAND_INPUT_KEYNODES, () -> {
                     writeKeyNodes(inputMembersWithAttr);
                 });
-            }
-            if (!outputMembersWithAttr.isEmpty()) {
-                writer.openBlock("private readonly $L = [", "];", COMMAND_OUTPUT_KEYNODES, () -> {
+                writer.openBlock("protected readonly $L = [", "];", COMMAND_OUTPUT_KEYNODES, () -> {
                     writeKeyNodes(outputMembersWithAttr);
                 });
-            }
-            writer.popState();
-            writer.write("");
+                writer.popState();
+                writer.write("");
+
+                writer.write("protected readonly clientCommand: $L;", clientCommandLocalName);
+                writer.write(
+                    "public readonly middlewareStack: MiddlewareStack<$L>;",
+                    inputTypeName + " | __" + originalInputTypeName
+                        + ", \n" + outputTypeName + " | __" + originalOutputTypeName
+                );
+                writer.write("");
 
-            generateCommandConstructor();
-            writer.write("");
-            generateCommandMiddlewareResolver(configType);
+                generateCommandConstructor();
+                writer.write("");
+                generateCommandMiddlewareResolver(configType);
 
-            // Hook for adding more methods to the command.
-            writer.pushState(COMMAND_BODY_EXTRA_SECTION).popState();
-        });
+                // Hook for adding more methods to the command.
+                writer.pushState(COMMAND_BODY_EXTRA_SECTION).popState();
+            }
+        );
     }
 
     private void generateCommandConstructor() {
         writer.openBlock("constructor(readonly input: $L) {", "}", inputTypeName, () -> {
             // The constructor can be intercepted and changed.
             writer.pushState(COMMAND_CONSTRUCTOR_SECTION)
                     .write("super();")
+                    .write("this.clientCommand = new $L(this.input as any);", clientCommandLocalName)
+                    .write("this.middlewareStack = this.clientCommand.middlewareStack;")
                     .popState();
         });
     }
@@ -155,52 +185,27 @@ private void generateCommandMiddlewareResolver(String configType) {
                 .write("options?: $T", ApplicationProtocol.createDefaultHttpApplicationProtocol().getOptionsType())
                 .dedent();
         writer.openBlock("): $L<$L, $L> {", "}", handler, inputTypeName, outputTypeName, () -> {
-            String marshallOptions = DocumentClientUtils.CLIENT_MARSHALL_OPTIONS;
-            String unmarshallOptions = DocumentClientUtils.CLIENT_UNMARSHALL_OPTIONS;
-
-            writer.write("const { $L, $L } = configuration.$L || {};", marshallOptions, unmarshallOptions,
-                    DocumentClientUtils.CLIENT_TRANSLATE_CONFIG_KEY);
-
-            writer.addImport(symbol.getName(), "__" + symbol.getName(), "@aws-sdk/client-dynamodb");
-
-            String marshallInput = "marshallInput";
-            String unmarshallOutput = "unmarshallOutput";
-            String utilsFileLocation = String.format("./%s/%s",
-                DocumentClientUtils.CLIENT_COMMANDS_FOLDER, DocumentClientUtils.CLIENT_UTILS_FILE);
-            writer.addImport(marshallInput, marshallInput, utilsFileLocation);
-            writer.addImport(unmarshallOutput, unmarshallOutput, utilsFileLocation);
-
-            String commandVarName = "command";
-            writer.openBlock("const $L = new $L(", ");", commandVarName, "__" + symbol.getName(),
-                () -> {
-                    if (inputMembersWithAttr.isEmpty()) {
-                        writer.write("this.input,");
-                    } else {
-                        writer.openBlock("$L(", ")", marshallInput, () -> {
-                            writer.write("this.input,");
-                            writer.write(getFrom("this", COMMAND_INPUT_KEYNODES) + ",");
-                            writer.write("$L,", marshallOptions);
-                        });
-                    }
-                });
+
+            writer.addImport(clientCommandClassName, clientCommandLocalName, "@aws-sdk/client-dynamodb");
+
+            String commandVarName = "this.clientCommand";
+
+            // marshall middlewares
+            writer.openBlock("this.addMarshallingMiddleware(", ");", () -> {
+                writer.write("configuration");
+            });
+
+            writer.write("const stack = clientStack.concat(this.middlewareStack as typeof clientStack);");
+
             String handlerVarName = "handler";
-            writer.write("const $L = $L.resolveMiddleware(clientStack, configuration, options);",
+            writer.write("const $L = $L.resolveMiddleware(stack, configuration, options);",
                 handlerVarName, commandVarName);
             writer.write("");
 
             if (outputMembersWithAttr.isEmpty()) {
                 writer.write("return $L;", handlerVarName);
             } else {
-                writer.openBlock("return async () => {", "};", () -> {
-                    String dataVarName = "data";
-                    String outputVarName = "output";
-                    writer.write("const $L = await $L($L);", dataVarName, handlerVarName, commandVarName);
-                    writer.openBlock("return {", "};", () -> {
-                        writer.write("...$L,", dataVarName);
-                        writer.write("$1L: $2L($3L.$1L, this.$4L, $5L),", outputVarName, unmarshallOutput,
-                            dataVarName, COMMAND_OUTPUT_KEYNODES, unmarshallOptions);
-                    });
-                });
+                writer.write("return async () => handler($L)", commandVarName);
             }
         });
     }
@@ -259,10 +264,8 @@ private void writeStructureKeyNode(StructureShape structureTarget) {
 
     private void generateInputAndOutputTypes() {
         writer.write("");
-        String originalInputTypeName = symbol.expectProperty("inputType", Symbol.class).getName();
         writeType(inputTypeName, originalInputTypeName, operationIndex.getInput(operation), inputMembersWithAttr);
         writer.write("");
-        String originalOutputTypeName = symbol.expectProperty("outputType", Symbol.class).getName();
         writeType(outputTypeName, originalOutputTypeName, operationIndex.getOutput(operation), outputMembersWithAttr);
         writer.write("");
     }

lib/lib-dynamodb/README.md

@@ -160,6 +160,79 @@ await ddbDocClient.put({
 });
 ```
 
+### Client and Command middleware stacks
+
+As with other AWS SDK for JavaScript v3 clients, you can apply middleware functions
+both on the client itself and individual `Command`s.
+
+For individual `Command`s, here are examples of how to add middleware before and after
+both marshalling and unmarshalling. We will use `QueryCommand` as an example.
+Others follow the same pattern.
+
+```js
+import { DynamoDBDocumentClient, QueryCommand } from "@aws-sdk/lib-dynamodb";
+
+const client = new DynamoDBClient({
+  /*...*/
+});
+const doc = DynamoDBDocumentClient.from(client);
+const command = new QueryCommand({
+  /*...*/
+});
+```
+
+Before and after marshalling:
+
+```js
+command.middlewareStack.addRelativeTo(
+  (next) => async (args) => {
+    console.log("pre-marshall", args.input);
+    return next(args);
+  },
+  {
+    relation: "before",
+    toMiddleware: "DocumentMarshall",
+  }
+);
+command.middlewareStack.addRelativeTo(
+  (next) => async (args) => {
+    console.log("post-marshall", args.input);
+    return next(args);
+  },
+  {
+    relation: "after",
+    toMiddleware: "DocumentMarshall",
+  }
+);
+```
+
+Before and after unmarshalling:
+
+```js
+command.middlewareStack.addRelativeTo(
+  (next) => async (args) => {
+    const result = await next(args);
+    console.log("pre-unmarshall", result.output.Items);
+    return result;
+  },
+  {
+    relation: "after", // <- after for pre-unmarshall
+    toMiddleware: "DocumentUnmarshall",
+  }
+);
+command.middlewareStack.addRelativeTo(
+  (next) => async (args) => {
+    const result = await next(args);
+    console.log("post-unmarshall", result.output.Items);
+    return result;
+  },
+  {
+    relation: "before", // <- before for post-unmarshall
+    toMiddleware: "DocumentUnmarshall",
+  }
+);
+```
+
 ### Destroying document client
 
 The `destroy()` call on document client is a no-op as document client does not

lib/lib-dynamodb/src/baseCommand/DynamoDBDocumentClientCommand.spec.ts

@@ -0,0 +1,57 @@
+import { Handler, MiddlewareStack } from "@aws-sdk/types";
+
+import { KeyNode } from "../commands/utils";
+import { DynamoDBDocumentClientCommand } from "./DynamoDBDocumentClientCommand";
+
+class AnyCommand extends DynamoDBDocumentClientCommand<{}, {}, {}, {}, {}> {
+  public middlewareStack: MiddlewareStack<{}, {}>;
+  public input: {};
+  protected inputKeyNodes: KeyNode[] = [];
+  protected outputKeyNodes: KeyNode[] = [];
+
+  public argCaptor: [Function, object][] = [];
+
+  protected readonly clientCommand = {
+    middlewareStack: {
+      argCaptor: this.argCaptor,
+      add(fn, config) {
+        this.argCaptor.push([fn, config]);
+      },
+    },
+  } as any;
+  protected readonly clientCommandName = "AnyCommand";
+
+  public constructor() {
+    super();
+  }
+
+  public resolveMiddleware(clientStack: MiddlewareStack<any, any>, configuration: {}, options: any): Handler<{}, {}> {
+    this.addMarshallingMiddleware({} as any);
+    return null as any;
+  }
+}
+
+describe("DynamoDBDocumentClientCommand", () => {
+  it("should not allow usage of the default middlewareStack", () => {
+    const command = new AnyCommand();
+    command.resolveMiddleware(null as any, null as any, null as any);
+    {
+      const [middleware, options] = command.argCaptor[0];
+      expect(middleware.toString()).toContain(`marshallInput`);
+      expect(options).toEqual({
+        name: "DocumentMarshall",
+        override: true,
+        step: "initialize",
+      });
+    }
+    {
+      const [middleware, options] = command.argCaptor[1];
+      expect(middleware.toString()).toContain(`unmarshallOutput`);
+      expect(options).toEqual({
+        name: "DocumentUnmarshall",
+        override: true,
+        step: "deserialize",
+      });
+    }
+  });
+});