Skip to content

Commit e7fee8c

Browse files
adwsinghrhernandez35
adwsingh
authored and
rhernandez35
committed
Handle additionalInput without modifying the actual input model
1 parent 873f9d0 commit e7fee8c

6 files changed

Lines changed: 235 additions & 59 deletions

File tree

‎aws/aws-service-bundle/src/main/java/software/amazon/smithy/java/aws/servicebundle/provider/AwsServiceBundle.java‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -5,8 +5,9 @@
55

66
package software.amazon.smithy.java.aws.servicebundle.provider;
77

8+
import static java.util.Objects.requireNonNull;
9+
810
import java.net.URI;
9-
import java.util.Objects;
1011
import java.util.ServiceLoader;
1112
import java.util.Set;
1213
import java.util.stream.Collectors;
@@ -25,8 +26,8 @@
2526
import software.amazon.smithy.java.client.core.endpoint.EndpointResolver;
2627
import software.amazon.smithy.java.client.core.interceptors.CallHook;
2728
import software.amazon.smithy.java.client.core.interceptors.ClientInterceptor;
28-
import software.amazon.smithy.java.core.serde.document.Document;
2929
import software.amazon.smithy.java.dynamicclient.DynamicClient;
30+
import software.amazon.smithy.java.server.ProxyService;
3031
import software.amazon.smithy.model.knowledge.ServiceIndex;
3132
import software.amazon.smithy.model.shapes.ShapeId;
3233
import software.amazon.smithy.modelbundle.api.BundlePlugin;
@@ -87,12 +88,11 @@ private record AwsServiceClientInterceptor(AwsServiceMetadata serviceMetadata, A
8788

8889
@Override
8990
public ClientConfig modifyBeforeCall(CallHook<?, ?> hook) {
90-
if (!(hook.input() instanceof Document d)) {
91-
throw new IllegalArgumentException("Input must be a Document");
92-
}
91+
var d = requireNonNull(hook.config().context().get(ProxyService.PROXY_INPUT),
92+
"Expected additionalInput in the request");
9393
var input = d.asShape(PreRequest.builder());
9494

95-
var endpoint = URI.create(Objects.requireNonNull(serviceMetadata.getEndpoints().get(input.getAwsRegion()),
95+
var endpoint = URI.create(requireNonNull(serviceMetadata.getEndpoints().get(input.getAwsRegion()),
9696
"no endpoint for region " + input.getAwsRegion()));
9797

9898
var identityResolver =

‎client/client-core/src/main/java/software/amazon/smithy/java/client/core/Client.java‎

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -80,12 +80,21 @@ protected <I extends SerializableStruct, O extends SerializableStruct> Completab
8080
IdentityResolvers callIdentityResolvers = identityResolvers;
8181
ClientInterceptor callInterceptor = interceptor;
8282

83-
// First apply overrides from interceptors.
84-
ClientConfig callConfig = callInterceptor.modifyBeforeCall(new CallHook<>(operation, config, input));
85-
// Overrides given per/operation take precedence over interceptors.
83+
//If there is an override config first apply that before sending to interceptors.
84+
ClientConfig callConfig = config;
8685
if (overrideConfig != null) {
8786
callConfig = callConfig.withRequestOverride(overrideConfig);
8887
}
88+
ClientConfig afterInterceptionConfig =
89+
callInterceptor.modifyBeforeCall(new CallHook<>(operation, callConfig, input));
90+
if (afterInterceptionConfig != null && afterInterceptionConfig != callConfig) {
91+
if (overrideConfig != null) {
92+
callConfig = afterInterceptionConfig.withRequestOverride(overrideConfig);
93+
} else {
94+
callConfig = afterInterceptionConfig;
95+
}
96+
97+
}
8998

9099
// Rebuild the pipeline, resolvers, etc if the config changed.
91100
if (callConfig != config) {

‎client/client-core/src/test/java/software/amazon/smithy/java/client/core/ClientTest.java‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -169,6 +169,8 @@ public void requestOverridesPerCallTakePrecedence() throws URISyntaxException {
169169
.addPlugin(config -> config.addInterceptor(new ClientInterceptor() {
170170
@Override
171171
public ClientConfig modifyBeforeCall(CallHook<?, ?> hook) {
172+
//RequestOverrides config should be visible here.
173+
assertThat(hook.config().context().get(CallContext.APPLICATION_ID), equalTo(id));
172174
// Note that the overrides given to the call itself will override interceptors.
173175
var override = RequestOverrideConfig.builder()
174176
.putConfig(CallContext.APPLICATION_ID, "foo")

‎model-bundle/model-bundle-api/src/main/java/software/amazon/smithy/modelbundle/api/ModelBundles.java‎

Lines changed: 53 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -5,10 +5,14 @@
55

66
package software.amazon.smithy.modelbundle.api;
77

8+
import software.amazon.smithy.java.server.ProxyOperationTrait;
89
import software.amazon.smithy.java.server.ProxyService;
910
import software.amazon.smithy.java.server.Service;
1011
import software.amazon.smithy.model.Model;
1112
import software.amazon.smithy.model.loader.ModelAssembler;
13+
import software.amazon.smithy.model.shapes.MemberShape;
14+
import software.amazon.smithy.model.shapes.OperationShape;
15+
import software.amazon.smithy.model.shapes.ServiceShape;
1216
import software.amazon.smithy.model.shapes.ShapeId;
1317
import software.amazon.smithy.model.shapes.StructureShape;
1418
import software.amazon.smithy.model.traits.StreamingTrait;
@@ -49,8 +53,20 @@ private static Model getModel(SmithyBundle bundle) {
4953
}
5054
var b = model.toBuilder();
5155

56+
var serviceShapes = model.getServiceShapes();
57+
58+
if (serviceShapes.size() != 1) {
59+
throw new IllegalStateException("Expected exactly one service shape but got "
60+
+ serviceShapes.stream().map(ServiceShape::getId).toList());
61+
}
62+
63+
var service = serviceShapes.iterator().next();
64+
5265
// mix in the generic arg members
53-
for (var op : model.getOperationShapes()) {
66+
var serviceShape = service.asServiceShape().get();
67+
var serviceBuilder = serviceShape.toBuilder();
68+
for (var opId : serviceShape.getAllOperations()) {
69+
var op = model.expectShape(opId, OperationShape.class);
5470
boolean skipOperation = false;
5571
if (op.getOutput().isPresent()) {
5672
for (var member : model.expectShape(op.getOutputShape(), StructureShape.class).members()) {
@@ -67,9 +83,7 @@ private static Model getModel(SmithyBundle bundle) {
6783
}
6884

6985
if (op.getInput().isEmpty() && additionalInputShape != null) {
70-
b.addShape(op.toBuilder()
71-
.input(additionalInputShape)
72-
.build());
86+
addProxyOperationWithAdditionalInput(op, additionalInputShape, b, serviceBuilder, model);
7387
} else {
7488
var shape = model.expectShape(op.getInputShape(), StructureShape.class);
7589
for (var member : shape.members()) {
@@ -85,24 +99,46 @@ private static Model getModel(SmithyBundle bundle) {
8599
}
86100

87101
if (additionalInputShape != null) {
88-
var input = shape.toBuilder();
89-
for (var member : additionalInputShape.members()) {
90-
input.addMember(member.toBuilder()
91-
.id(ShapeId.from(input.getId().toString() + "$" + member.getMemberName()))
92-
.build());
93-
}
94-
b.addShape(input.build());
102+
addProxyOperationWithAdditionalInput(op, additionalInputShape, b, serviceBuilder, model);
95103
}
96104
}
97105
}
98106

99-
for (var service : model.getServiceShapes()) {
100-
b.addShape(service.toBuilder()
101-
// trim the endpoint rules because they're huge and we don't need them
102-
.removeTrait(ShapeId.from("smithy.rules#endpointRuleSet"))
103-
.removeTrait(ShapeId.from("smithy.rules#endpointTests"))
107+
b.addShape(serviceBuilder
108+
// trim the endpoint rules because they're huge and we don't need them
109+
.removeTrait(ShapeId.from("smithy.rules#endpointRuleSet"))
110+
.removeTrait(ShapeId.from("smithy.rules#endpointTests"))
111+
.build());
112+
return b.build();
113+
}
114+
115+
private static void addProxyOperationWithAdditionalInput(
116+
OperationShape op,
117+
StructureShape additionalInput,
118+
Model.Builder builder,
119+
ServiceShape.Builder serviceBuilder,
120+
Model model
121+
) {
122+
var input = op.getInput();
123+
StructureShape finalInput;
124+
if (op.getInput().isEmpty()) {
125+
finalInput = additionalInput;
126+
} else {
127+
var inputBuilder = model.expectShape(input.get(), StructureShape.class).toBuilder();
128+
inputBuilder.addMember(MemberShape.builder()
129+
.id(ShapeId.from(inputBuilder.getId().toString() + "$additionalInput"))
130+
.target(additionalInput.getId())
104131
.build());
132+
finalInput = inputBuilder.id(ShapeId.from(inputBuilder.getId().toString()) + "Proxy").build();
105133
}
106-
return b.build();
134+
builder.addShape(finalInput);
135+
var newOperation = op.toBuilder()
136+
.id(ShapeId.from(op.getId().toString() + "Proxy"))
137+
.input(finalInput)
138+
.output(op.getOutputShape())
139+
.addTrait(new ProxyOperationTrait(op.getId()))
140+
.build();
141+
builder.addShape(newOperation);
142+
serviceBuilder.addOperation(newOperation).build();
107143
}
108144
}

‎server/server-proxy/src/main/java/software/amazon/smithy/java/server/ProxyOperationTrait.java‎

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,62 @@
1+
/*
2+
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package software.amazon.smithy.java.server;
7+
8+
import software.amazon.smithy.model.node.Node;
9+
import software.amazon.smithy.model.shapes.ShapeId;
10+
import software.amazon.smithy.model.traits.Trait;
11+
12+
/**
13+
* A trait that marks operations as proxies to delegate operations.
14+
*
15+
* <p>Operations annotated with this trait are treated as proxies to the
16+
* delegate operation specified by this trait. Proxy operations enable
17+
* additional input processing while maintaining compatibility with the
18+
* original operation interface.</p>
19+
*
20+
* <h3>Proxy Operation Input Handling</h3>
21+
* <p>Proxy operation inputs are expected to be a superset of the original
22+
* operation input. For operations with existing input, a new input structure
23+
* is created that contains all original input members plus an
24+
* {@code additionalInput} field containing the additional parameters. For
25+
* operations with no input, the additional input structure is used directly.</p>
26+
*
27+
* <h3>Additional Input Access</h3>
28+
* <p>The additional input data can be accessed in Dynamic client interceptors
29+
* using the {@code ProxyService.PROXY_INPUT} context key. This enables
30+
* interceptors to process the extra data before the request is forwarded
31+
* to the delegate service.</p>
32+
*
33+
* <h3>Input Stripping</h3>
34+
* <p>The proxy service automatically strips out the additional input before
35+
* sending the request to the delegate service, ensuring that only the
36+
* expected input parameters are forwarded to the original operation.</p>
37+
*
38+
* @see ProxyService#PROXY_INPUT
39+
*/
40+
public final class ProxyOperationTrait implements Trait {
41+
42+
private static final ShapeId SHAPE_ID = ShapeId.from("smithy.server.api#proxyOperation");
43+
private final ShapeId delegateOperation;
44+
45+
public ProxyOperationTrait(ShapeId delegateOperation) {
46+
this.delegateOperation = delegateOperation;
47+
}
48+
49+
@Override
50+
public ShapeId toShapeId() {
51+
return SHAPE_ID;
52+
}
53+
54+
@Override
55+
public Node toNode() {
56+
return null;
57+
}
58+
59+
public ShapeId getDelegateOperation() {
60+
return delegateOperation;
61+
}
62+
}

0 commit comments

Comments
 (0)