Skip to content

Commit 6b81ab7

Browse files
mtdowlingadwsingh
authored andcommitted
Add aws|ec2Query schema extension
Precomputes the bytes needed for AWS Query and EC2 query names so they can be fed directly into FormUrlEncodedSink.
1 parent 7fcd04c commit 6b81ab7

4 files changed

Lines changed: 130 additions & 86 deletions

File tree

Original file line numberDiff line numberDiff line change
@@ -0,0 +1,93 @@
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.aws.client.awsquery;
7+
8+
import java.nio.ByteBuffer;
9+
import software.amazon.smithy.aws.traits.protocols.Ec2QueryNameTrait;
10+
import software.amazon.smithy.java.core.schema.Schema;
11+
import software.amazon.smithy.java.core.schema.SchemaExtensionKey;
12+
import software.amazon.smithy.java.core.schema.SchemaExtensionProvider;
13+
import software.amazon.smithy.java.core.schema.TraitKey;
14+
import software.amazon.smithy.utils.SmithyInternalApi;
15+
import software.amazon.smithy.utils.StringUtils;
16+
17+
/**
18+
* Pre-computes the URL-encoded member-name bytes for AWS Query and EC2 Query protocols once per {@link Schema}.
19+
*/
20+
@SmithyInternalApi
21+
public final class AwsQuerySchemaExtensions
22+
implements SchemaExtensionProvider<AwsQuerySchemaExtensions.QueryMemberBinding> {
23+
24+
public static final SchemaExtensionKey<QueryMemberBinding> KEY = new SchemaExtensionKey<>();
25+
26+
/**
27+
* Pre-encoded member-name bytes for both query variants.
28+
*
29+
* @param awsQueryNameBytes Bytes to use as the awsQuery name. Never null.
30+
* @param ec2QueryNameBytes Bytes to use as the ec2Query name. Never null.
31+
*/
32+
public record QueryMemberBinding(byte[] awsQueryNameBytes, byte[] ec2QueryNameBytes) {}
33+
34+
@Override
35+
public SchemaExtensionKey<QueryMemberBinding> key() {
36+
return KEY;
37+
}
38+
39+
@Override
40+
public QueryMemberBinding provide(Schema schema) {
41+
if (!schema.isMember()) {
42+
return null;
43+
}
44+
45+
return new QueryMemberBinding(
46+
encodeName(resolveAwsQueryName(schema)),
47+
encodeName(resolveEc2QueryName(schema)));
48+
}
49+
50+
private static String resolveAwsQueryName(Schema schema) {
51+
var xmlName = schema.getTrait(TraitKey.XML_NAME_TRAIT);
52+
return xmlName != null ? xmlName.getValue() : schema.memberName();
53+
}
54+
55+
private static String resolveEc2QueryName(Schema schema) {
56+
var ec2Name = schema.getTrait(TraitKey.get(Ec2QueryNameTrait.class));
57+
if (ec2Name != null) {
58+
return ec2Name.getValue();
59+
}
60+
61+
var xmlName = schema.getTrait(TraitKey.XML_NAME_TRAIT);
62+
return xmlName != null
63+
? StringUtils.capitalize(xmlName.getValue())
64+
: StringUtils.capitalize(schema.memberName());
65+
}
66+
67+
@SuppressWarnings("deprecation")
68+
static byte[] encodeName(String name) {
69+
int len = name.length();
70+
71+
boolean needsEncoding = false;
72+
for (int i = 0; i < len; i++) {
73+
char c = name.charAt(i);
74+
if (!FormUrlEncodedSink.isUnreserved(c)) {
75+
needsEncoding = true;
76+
break;
77+
}
78+
}
79+
80+
if (!needsEncoding) {
81+
byte[] result = new byte[len];
82+
name.getBytes(0, len, result, 0);
83+
return result;
84+
}
85+
86+
FormUrlEncodedSink tmp = new FormUrlEncodedSink(len * 3);
87+
tmp.writeUrlEncoded(name);
88+
ByteBuffer bb = tmp.finish();
89+
byte[] result = new byte[bb.remaining()];
90+
bb.get(result);
91+
return result;
92+
}
93+
}

aws/client/aws-client-awsquery/src/main/java/software/amazon/smithy/java/aws/client/awsquery/FormUrlEncodedSink.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -108,7 +108,7 @@ ByteBuffer finish() {
108108
return ByteBuffer.wrap(bytes, 0, pos);
109109
}
110110

111-
private static boolean isUnreserved(char c) {
111+
static boolean isUnreserved(char c) {
112112
return (c >= 'A' && c <= 'Z')
113113
|| (c >= 'a' && c <= 'z')
114114
|| (c >= '0' && c <= '9')

aws/client/aws-client-awsquery/src/main/java/software/amazon/smithy/java/aws/client/awsquery/QueryFormSerializer.java

Lines changed: 35 additions & 85 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,6 @@
1212
import java.time.Instant;
1313
import java.util.Arrays;
1414
import java.util.function.BiConsumer;
15-
import software.amazon.smithy.aws.traits.protocols.Ec2QueryNameTrait;
1615
import software.amazon.smithy.java.core.schema.Schema;
1716
import software.amazon.smithy.java.core.schema.SerializableStruct;
1817
import software.amazon.smithy.java.core.schema.TraitKey;
@@ -82,17 +81,6 @@ private void writeParam(byte[] key, String value) {
8281
sink.writeUrlEncoded(value);
8382
}
8483

85-
private void writeParam(String key, String value) {
86-
sink.writeByte('&');
87-
writeCurrentPrefix();
88-
if (prefixDepth > 0) {
89-
sink.writeByte('.');
90-
}
91-
sink.writeUrlEncoded(key);
92-
sink.writeByte('=');
93-
sink.writeUrlEncoded(value);
94-
}
95-
9684
private void writeCurrentPrefix() {
9785
for (int i = 0; i < prefixDepth; i++) {
9886
if (i > 0) {
@@ -102,10 +90,6 @@ private void writeCurrentPrefix() {
10290
}
10391
}
10492

105-
private void pushPrefix(String prefix) {
106-
pushPrefix(encodePrefix(prefix));
107-
}
108-
10993
private void pushPrefix(byte[] prefix) {
11094
ensurePrefixCacheCapacity();
11195
prefixCache[prefixDepth++] = prefix;
@@ -131,15 +115,6 @@ private void popPrefix() {
131115
prefixDepth--;
132116
}
133117

134-
private byte[] encodePrefix(String prefix) {
135-
FormUrlEncodedSink tmp = new FormUrlEncodedSink(prefix.length() * 3);
136-
tmp.writeUrlEncoded(prefix);
137-
ByteBuffer bb = tmp.finish();
138-
byte[] result = new byte[bb.remaining()];
139-
bb.get(result);
140-
return result;
141-
}
142-
143118
@SuppressWarnings("deprecation")
144119
private byte[] encodeIndexedPrefix(byte[] base, int index) {
145120
String indexStr = Integer.toString(index);
@@ -160,54 +135,29 @@ private byte[] encodeIndex(int index) {
160135

161136
// --- Member name resolution (protocol-specific) ---
162137

163-
private String getMemberName(Schema schema) {
164-
return switch (variant) {
165-
case AWS_QUERY -> getAwsQueryMemberName(schema);
166-
case EC2_QUERY -> getEc2QueryMemberName(schema);
167-
};
168-
}
169-
170-
private static String getAwsQueryMemberName(Schema schema) {
171-
var xmlName = schema.getTrait(TraitKey.XML_NAME_TRAIT);
172-
if (xmlName != null) {
173-
return xmlName.getValue();
174-
}
175-
return schema.memberName();
176-
}
177-
178-
private static String getEc2QueryMemberName(Schema schema) {
179-
var ec2Name = schema.getTrait(TraitKey.get(Ec2QueryNameTrait.class));
180-
if (ec2Name != null) {
181-
return ec2Name.getValue();
182-
}
183-
var xmlName = schema.getTrait(TraitKey.XML_NAME_TRAIT);
184-
if (xmlName != null) {
185-
return capitalize(xmlName.getValue());
186-
}
187-
return capitalize(schema.memberName());
188-
}
189-
190-
private static String capitalize(String s) {
191-
if (s == null || s.isEmpty() || Character.isUpperCase(s.charAt(0))) {
192-
return s;
138+
/**
139+
* Read the pre-computed URL-encoded member-name bytes from the schema extension.
140+
*/
141+
private byte[] getMemberNameBytes(Schema schema) {
142+
var ext = schema.getExtension(AwsQuerySchemaExtensions.KEY);
143+
if (ext == null) {
144+
return null;
193145
}
194-
return Character.toUpperCase(s.charAt(0)) + s.substring(1);
146+
return variant == QueryVariant.AWS_QUERY ? ext.awsQueryNameBytes() : ext.ec2QueryNameBytes();
195147
}
196148

197149
// --- Struct ---
198150

199151
@Override
200152
public void writeStruct(Schema schema, SerializableStruct struct) {
201153
if (schema.isMember()) {
202-
String memberName = getMemberName(schema);
203-
if (memberName != null) {
204-
pushPrefix(memberName);
205-
struct.serializeMembers(this);
206-
popPrefix();
207-
return;
208-
}
154+
// Member schemas always have a non-null QueryMemberBinding (see provider).
155+
pushPrefix(getMemberNameBytes(schema));
156+
struct.serializeMembers(this);
157+
popPrefix();
158+
} else {
159+
struct.serializeMembers(this);
209160
}
210-
struct.serializeMembers(this);
211161
}
212162

213163
// --- List (protocol-specific) ---
@@ -231,7 +181,7 @@ private <T> void writeAwsQueryList(
231181
Schema memberSchema = schema.listMember();
232182

233183
if (schema.isMember()) {
234-
pushPrefix(getMemberName(schema));
184+
pushPrefix(getMemberNameBytes(schema));
235185
}
236186

237187
if (size == 0) {
@@ -261,7 +211,7 @@ private <T> void writeAwsQueryList(
261211
private <T> void writeEc2List(Schema schema, T listState, int size, BiConsumer<T, ShapeSerializer> consumer) {
262212
// EC2 Query lists are always flattened - no .member. segment
263213
if (schema.isMember()) {
264-
pushPrefix(getMemberName(schema));
214+
pushPrefix(getMemberNameBytes(schema));
265215
}
266216

267217
if (size == 0) {
@@ -443,7 +393,7 @@ public <T> void writeMap(Schema schema, T mapState, int size, BiConsumer<T, MapS
443393
Schema valueSchema = schema.mapValueMember();
444394

445395
if (schema.isMember()) {
446-
pushPrefix(getMemberName(schema));
396+
pushPrefix(getMemberNameBytes(schema));
447397
}
448398

449399
var keyXmlName = keySchema.getTrait(TraitKey.XML_NAME_TRAIT);
@@ -641,77 +591,77 @@ private void writeValueParam(String value) {
641591

642592
@Override
643593
public void writeBoolean(Schema schema, boolean value) {
644-
writeParam(getMemberName(schema), value ? "true" : "false");
594+
writeParam(getMemberNameBytes(schema), value ? "true" : "false");
645595
}
646596

647597
@Override
648598
public void writeByte(Schema schema, byte value) {
649-
writeParam(getMemberName(schema), Byte.toString(value));
599+
writeParam(getMemberNameBytes(schema), Byte.toString(value));
650600
}
651601

652602
@Override
653603
public void writeShort(Schema schema, short value) {
654-
writeParam(getMemberName(schema), Short.toString(value));
604+
writeParam(getMemberNameBytes(schema), Short.toString(value));
655605
}
656606

657607
@Override
658608
public void writeInteger(Schema schema, int value) {
659-
writeParam(getMemberName(schema), Integer.toString(value));
609+
writeParam(getMemberNameBytes(schema), Integer.toString(value));
660610
}
661611

662612
@Override
663613
public void writeLong(Schema schema, long value) {
664-
writeParam(getMemberName(schema), Long.toString(value));
614+
writeParam(getMemberNameBytes(schema), Long.toString(value));
665615
}
666616

667617
@Override
668618
public void writeFloat(Schema schema, float value) {
669-
String memberName = getMemberName(schema);
619+
byte[] memberNameBytes = getMemberNameBytes(schema);
670620
if (Float.isNaN(value)) {
671-
writeParam(memberName, "NaN");
621+
writeParam(memberNameBytes, "NaN");
672622
} else if (Float.isInfinite(value)) {
673-
writeParam(memberName, value > 0 ? "Infinity" : "-Infinity");
623+
writeParam(memberNameBytes, value > 0 ? "Infinity" : "-Infinity");
674624
} else {
675-
writeParam(memberName, Float.toString(value));
625+
writeParam(memberNameBytes, Float.toString(value));
676626
}
677627
}
678628

679629
@Override
680630
public void writeDouble(Schema schema, double value) {
681-
String memberName = getMemberName(schema);
631+
byte[] memberNameBytes = getMemberNameBytes(schema);
682632
if (Double.isNaN(value)) {
683-
writeParam(memberName, "NaN");
633+
writeParam(memberNameBytes, "NaN");
684634
} else if (Double.isInfinite(value)) {
685-
writeParam(memberName, value > 0 ? "Infinity" : "-Infinity");
635+
writeParam(memberNameBytes, value > 0 ? "Infinity" : "-Infinity");
686636
} else {
687-
writeParam(memberName, Double.toString(value));
637+
writeParam(memberNameBytes, Double.toString(value));
688638
}
689639
}
690640

691641
@Override
692642
public void writeBigInteger(Schema schema, BigInteger value) {
693-
writeParam(getMemberName(schema), value.toString());
643+
writeParam(getMemberNameBytes(schema), value.toString());
694644
}
695645

696646
@Override
697647
public void writeBigDecimal(Schema schema, BigDecimal value) {
698-
writeParam(getMemberName(schema), value.toPlainString());
648+
writeParam(getMemberNameBytes(schema), value.toPlainString());
699649
}
700650

701651
@Override
702652
public void writeString(Schema schema, String value) {
703-
writeParam(getMemberName(schema), value);
653+
writeParam(getMemberNameBytes(schema), value);
704654
}
705655

706656
@Override
707657
public void writeBlob(Schema schema, ByteBuffer value) {
708-
writeParam(getMemberName(schema), ByteBufferUtils.base64Encode(value));
658+
writeParam(getMemberNameBytes(schema), ByteBufferUtils.base64Encode(value));
709659
}
710660

711661
@Override
712662
public void writeTimestamp(Schema schema, Instant value) {
713663
TimestampFormatter formatter = TimestampFormatter.of(schema, TimestampFormatTrait.Format.DATE_TIME);
714-
writeParam(getMemberName(schema), formatter.writeString(value));
664+
writeParam(getMemberNameBytes(schema), formatter.writeString(value));
715665
}
716666

717667
@Override
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
software.amazon.smithy.java.aws.client.awsquery.AwsQuerySchemaExtensions

0 commit comments

Comments
 (0)