Add SmithyContextCodeSection to CommandGenerator (#957)

smithy-lang · Sep 25, 2023 · 0b3b4cb · 0b3b4cb
1 parent 2503655
commit 0b3b4cb
Show file tree

Hide file tree

Showing 2 changed files with 61 additions and 0 deletions.
diff --git a/...ipt-codegen/src/main/java/software/amazon/smithy/typescript/codegen/CommandGenerator.java b/...ipt-codegen/src/main/java/software/amazon/smithy/typescript/codegen/CommandGenerator.java
@@ -50,6 +50,7 @@
 import software.amazon.smithy.typescript.codegen.endpointsV2.RuleSetParameterFinder;
 import software.amazon.smithy.typescript.codegen.integration.ProtocolGenerator;
 import software.amazon.smithy.typescript.codegen.integration.RuntimeClientPlugin;
+import software.amazon.smithy.typescript.codegen.sections.SmithyContextCodeSection;
 import software.amazon.smithy.typescript.codegen.validation.SensitiveDataFinder;
 import software.amazon.smithy.utils.OptionalUtils;
 import software.amazon.smithy.utils.SmithyInternalApi;
@@ -359,8 +360,13 @@ private void generateCommandMiddlewareResolver(String configType) {
                             () -> writer.writeInline("(_: any) => _"));
                 });
                 writer.openBlock("[SMITHY_CONTEXT_KEY]: {", "},", () -> {
+                    writer.pushState(SmithyContextCodeSection.builder()
+                        .service(service)
+                        .operation(operation)
+                        .build());
                     writer.write("service: $S,", service.toShapeId().getName());
                     writer.write("operation: $S,", operation.toShapeId().getName());
+                    writer.popState();
                 });
             });
             writer.write("const { requestHandler } = configuration;");

diff --git a/...ain/java/software/amazon/smithy/typescript/codegen/sections/SmithyContextCodeSection.java b/...ain/java/software/amazon/smithy/typescript/codegen/sections/SmithyContextCodeSection.java
@@ -0,0 +1,55 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package software.amazon.smithy.typescript.codegen.sections;
+
+import software.amazon.smithy.model.shapes.OperationShape;
+import software.amazon.smithy.model.shapes.ServiceShape;
+import software.amazon.smithy.utils.CodeSection;
+import software.amazon.smithy.utils.SmithyBuilder;
+import software.amazon.smithy.utils.SmithyUnstableApi;
+
+@SmithyUnstableApi
+public final class SmithyContextCodeSection implements CodeSection {
+    private final ServiceShape service;
+    private final OperationShape operation;
+
+    private SmithyContextCodeSection(Builder builder) {
+        service = SmithyBuilder.requiredState("service", builder.service);
+        operation = SmithyBuilder.requiredState("operation", builder.operation);
+    }
+
+    public ServiceShape getService() {
+        return service;
+    }
+
+    public OperationShape getOperation() {
+        return operation;
+    }
+
+    public static Builder builder() {
+        return new Builder();
+    }
+
+    public static class Builder implements SmithyBuilder<SmithyContextCodeSection> {
+        private ServiceShape service;
+        private OperationShape operation;
+
+        @Override
+        public SmithyContextCodeSection build() {
+            return new SmithyContextCodeSection(this);
+        }
+
+        public Builder service(ServiceShape service) {
+            this.service = service;
+            return this;
+        }
+
+        public Builder operation(OperationShape operation) {
+            this.operation = operation;
+            return this;
+        }
+    }
+}