Skip to content

Commit cd21b6d

Browse files
rcohrcoh
authored
Make Generated Builders Vec-Aware (#267)
* Make Generated Builders Vec-Aware This updates the generated builders & fluent client to be vec-aware. The builder method with the same name as the field will append the field to the list, creating it if it doesn't exist. A `set_xyz` method is added which retains the original behavior. * Fix test usages of builders * Add Map builder * Fix unit tests * Add impl Into clause for builder helpers * Update another test * One more spot that needed to be updated * Delete dead code
1 parent 336670b commit cd21b6d

7 files changed

Lines changed: 155 additions & 64 deletions

File tree

‎aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/FluentClientGenerator.kt‎

Lines changed: 52 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ import software.amazon.smithy.rust.codegen.rustlang.RustModule
1515
import software.amazon.smithy.rust.codegen.rustlang.RustType
1616
import software.amazon.smithy.rust.codegen.rustlang.RustWriter
1717
import software.amazon.smithy.rust.codegen.rustlang.asType
18+
import software.amazon.smithy.rust.codegen.rustlang.contains
1819
import software.amazon.smithy.rust.codegen.rustlang.documentShape
1920
import software.amazon.smithy.rust.codegen.rustlang.render
2021
import software.amazon.smithy.rust.codegen.rustlang.rust
@@ -29,6 +30,7 @@ import software.amazon.smithy.rust.codegen.smithy.generators.LibRsSection
2930
import software.amazon.smithy.rust.codegen.smithy.generators.ProtocolConfig
3031
import software.amazon.smithy.rust.codegen.smithy.generators.builderSymbol
3132
import software.amazon.smithy.rust.codegen.smithy.generators.error.errorSymbol
33+
import software.amazon.smithy.rust.codegen.smithy.generators.setterName
3234
import software.amazon.smithy.rust.codegen.smithy.rustType
3335
import software.amazon.smithy.rust.codegen.util.inputShape
3436
import software.amazon.smithy.rust.codegen.util.outputShape
@@ -148,19 +150,61 @@ class FluentClientGenerator(protocolConfig: ProtocolConfig) {
148150
val memberSymbol = symbolProvider.toSymbol(member)
149151
val outerType = memberSymbol.rustType()
150152
val coreType = outerType.stripOuter<RustType.Option>()
151-
val signature = when (coreType) {
152-
is RustType.String,
153-
is RustType.Box -> "(mut self, inp: impl Into<${coreType.render(true)}>) -> Self"
154-
else -> "(mut self, inp: ${coreType.render(true)}) -> Self"
153+
when (coreType) {
154+
is RustType.Vec -> renderVecHelper(member, memberName, coreType)
155+
is RustType.HashMap -> renderMapHelper(member, memberName, coreType)
156+
else -> {
157+
val signature = when (coreType) {
158+
is RustType.String,
159+
is RustType.Box -> "(mut self, inp: impl Into<${coreType.render(true)}>) -> Self"
160+
else -> "(mut self, inp: ${coreType.render(true)}) -> Self"
161+
}
162+
documentShape(member, model)
163+
rustBlock("pub fn $memberName$signature") {
164+
write("self.inner = self.inner.$memberName(inp);")
165+
write("self")
166+
}
167+
}
155168
}
156-
documentShape(member, model)
157-
rustBlock("pub fn $memberName$signature") {
158-
write("self.inner = self.inner.$memberName(inp);")
159-
write("self")
169+
// pure setter
170+
rustBlock("pub fn ${member.setterName()}(mut self, inp: ${outerType.render(true)}) -> Self") {
171+
rust(
172+
"""
173+
self.inner = self.inner.${member.setterName()}(inp);
174+
self
175+
"""
176+
)
160177
}
161178
}
162179
}
163180
}
164181
}
165182
}
183+
184+
private fun RustWriter.renderMapHelper(member: MemberShape, memberName: String, coreType: RustType.HashMap) {
185+
documentShape(member, model)
186+
val k = coreType.key
187+
val v = coreType.member
188+
189+
rustBlock("pub fn $memberName(mut self, k: impl Into<${k.render()}>, v: impl Into<${v.render()}>) -> Self") {
190+
rust(
191+
"""
192+
self.inner = self.inner.$memberName(k, v);
193+
self
194+
"""
195+
)
196+
}
197+
}
198+
199+
private fun RustWriter.renderVecHelper(member: MemberShape, memberName: String, coreType: RustType.Vec) {
200+
documentShape(member, model)
201+
rustBlock("pub fn $memberName(mut self, inp: impl Into<${coreType.member.render(true)}>) -> Self") {
202+
rust(
203+
"""
204+
self.inner = self.inner.$memberName(inp);
205+
self
206+
"""
207+
)
208+
}
209+
}
166210
}

‎aws/sdk/examples/dynamo-helloworld/src/main.rs‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -17,14 +17,14 @@ async fn main() -> Result<(), dynamodb::Error> {
1717
let new_table = client
1818
.create_table()
1919
.table_name("test-table")
20-
.key_schema(vec![KeySchemaElement::builder()
20+
.key_schema(KeySchemaElement::builder()
2121
.attribute_name("k")
2222
.key_type(KeyType::Hash)
23-
.build()])
24-
.attribute_definitions(vec![AttributeDefinition::builder()
23+
.build())
24+
.attribute_definitions(AttributeDefinition::builder()
2525
.attribute_name("k")
2626
.attribute_type(ScalarAttributeType::S)
27-
.build()])
27+
.build())
2828
.provisioned_throughput(
2929
ProvisionedThroughput::builder()
3030
.write_capacity_units(10)

‎aws/sdk/examples/dynamo-movies/src/main.rs‎

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

66
use aws_http::AwsErrorRetryPolicy;
77
use aws_hyper::{SdkError, SdkSuccess};
8-
use dynamodb::error::DescribeTableError;
98
use dynamodb::client::fluent_builders::Query;
9+
use dynamodb::error::DescribeTableError;
1010
use dynamodb::input::DescribeTableInput;
1111
use dynamodb::model::{
1212
AttributeDefinition, AttributeValue, KeySchemaElement, KeyType, ProvisionedThroughput,
@@ -71,7 +71,7 @@ async fn main() {
7171
client
7272
.put_item()
7373
.table_name(table_name)
74-
.item(parse_item(value))
74+
.set_item(Some(parse_item(value)))
7575
.send()
7676
.await
7777
.expect("failed to insert item");
@@ -110,26 +110,30 @@ fn create_table(
110110
client
111111
.create_table()
112112
.table_name(table_name)
113-
.key_schema(vec![
113+
.key_schema(
114114
KeySchemaElement::builder()
115115
.attribute_name("year")
116116
.key_type(KeyType::Hash)
117117
.build(),
118+
)
119+
.key_schema(
118120
KeySchemaElement::builder()
119121
.attribute_name("title")
120122
.key_type(KeyType::Range)
121123
.build(),
122-
])
123-
.attribute_definitions(vec![
124+
)
125+
.attribute_definitions(
124126
AttributeDefinition::builder()
125127
.attribute_name("year")
126128
.attribute_type(ScalarAttributeType::N)
127129
.build(),
130+
)
131+
.attribute_definitions(
128132
AttributeDefinition::builder()
129133
.attribute_name("title")
130134
.attribute_type(ScalarAttributeType::S)
131135
.build(),
132-
])
136+
)
133137
.provisioned_throughput(
134138
ProvisionedThroughput::builder()
135139
.read_capacity_units(10)
@@ -159,16 +163,12 @@ fn value_to_item(value: Value) -> AttributeValue {
159163
}
160164

161165
fn movies_in_year(client: &dynamodb::Client, table_name: &str, year: u16) -> Query {
162-
let mut expr_attrib_names = HashMap::new();
163-
expr_attrib_names.insert("#yr".to_string(), "year".to_string());
164-
let mut expr_attrib_values = HashMap::new();
165-
expr_attrib_values.insert(":yyyy".to_string(), AttributeValue::N(year.to_string()));
166166
client
167167
.query()
168168
.table_name(table_name)
169169
.key_condition_expression("#yr = :yyyy")
170-
.expression_attribute_names(expr_attrib_names)
171-
.expression_attribute_values(expr_attrib_values)
170+
.expression_attribute_names("#yr", "year")
171+
.expression_attribute_values(":yyyy", AttributeValue::N(year.to_string()))
172172
}
173173

174174
/// Hand-written waiter to retry every second until the table is out of `Creating` state

‎aws/sdk/integration-tests/dynamodb/tests/movies.rs‎

Lines changed: 27 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -16,41 +16,45 @@ use dynamodb::model::{
1616
};
1717
use dynamodb::operation::{CreateTable, DescribeTable};
1818
use dynamodb::output::DescribeTableOutput;
19-
use dynamodb::{Config, Region, Credentials};
19+
use dynamodb::{Config, Credentials, Region};
20+
use http::header::{HeaderName, AUTHORIZATION};
2021
use http::Uri;
2122
use serde_json::Value;
22-
use smithy_http::operation::Operation;
2323
use smithy_http::body::SdkBody;
24+
use smithy_http::operation::Operation;
2425
use smithy_http::retry::ClassifyResponse;
2526
use smithy_types::retry::RetryKind;
2627
use std::collections::HashMap;
2728
use std::time::Duration;
2829
use tokio::time::Instant;
29-
use http::header::{AUTHORIZATION, HeaderName};
3030

3131
fn create_table(table_name: &str) -> create_table_input::Builder {
3232
CreateTable::builder()
3333
.table_name(table_name)
34-
.key_schema(vec![
34+
.key_schema(
3535
KeySchemaElement::builder()
3636
.attribute_name("year")
3737
.key_type(KeyType::Hash)
3838
.build(),
39+
)
40+
.key_schema(
3941
KeySchemaElement::builder()
4042
.attribute_name("title")
4143
.key_type(KeyType::Range)
4244
.build(),
43-
])
44-
.attribute_definitions(vec![
45+
)
46+
.attribute_definitions(
4547
AttributeDefinition::builder()
4648
.attribute_name("year")
4749
.attribute_type(ScalarAttributeType::N)
4850
.build(),
51+
)
52+
.attribute_definitions(
4953
AttributeDefinition::builder()
5054
.attribute_name("title")
5155
.attribute_type(ScalarAttributeType::S)
5256
.build(),
53-
])
57+
)
5458
.provisioned_throughput(
5559
ProvisionedThroughput::builder()
5660
.read_capacity_units(10)
@@ -80,7 +84,7 @@ fn add_item(table_name: impl Into<String>, item: Value) -> put_item_input::Build
8084

8185
PutItemInput::builder()
8286
.table_name(table_name)
83-
.item(attribute_value)
87+
.set_item(Some(attribute_value))
8488
}
8589

8690
fn movies_in_year(table_name: &str, year: u16) -> query_input::Builder {
@@ -91,8 +95,8 @@ fn movies_in_year(table_name: &str, year: u16) -> query_input::Builder {
9195
QueryInput::builder()
9296
.table_name(table_name)
9397
.key_condition_expression("#yr = :yyyy")
94-
.expression_attribute_names(expr_attrib_names)
95-
.expression_attribute_values(expr_attrib_values)
98+
.set_expression_attribute_names(Some(expr_attrib_names))
99+
.set_expression_attribute_values(Some(expr_attrib_values))
96100
}
97101

98102
/// Hand-written waiter to retry every second until the table is out of `Creating` state
@@ -150,7 +154,6 @@ fn wait_for_ready_table(
150154
operation.with_retry_policy(waiting_policy)
151155
}
152156

153-
154157
/// Validate that time has passed with a 5ms tolerance
155158
///
156159
/// This is to account for some non-determinism in the Tokio timer
@@ -216,8 +219,19 @@ async fn movies_it() {
216219
.await
217220
.expect("query should succeed");
218221
assert_eq!(films_2013.count, 2);
219-
let titles: Vec<AttributeValue> = films_2013.items.unwrap().into_iter().map(|mut row|row.remove("title").expect("row should have title")).collect();
220-
assert_eq!(titles, vec![AttributeValue::S("Rush".to_string()), AttributeValue::S("Turn It Down, Or Else!".to_string())]);
222+
let titles: Vec<AttributeValue> = films_2013
223+
.items
224+
.unwrap()
225+
.into_iter()
226+
.map(|mut row| row.remove("title").expect("row should have title"))
227+
.collect();
228+
assert_eq!(
229+
titles,
230+
vec![
231+
AttributeValue::S("Rush".to_string()),
232+
AttributeValue::S("Turn It Down, Or Else!".to_string())
233+
]
234+
);
221235

222236
for req in conn.requests().iter() {
223237
req.assert_matches(vec![AUTHORIZATION, HeaderName::from_static("x-amz-date")]);

‎codegen/src/main/kotlin/software/amazon/smithy/rust/codegen/rustlang/RustTypes.kt‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -101,7 +101,7 @@ sealed class RustType {
101101
data class Opaque(override val name: kotlin.String, override val namespace: kotlin.String? = null) : RustType()
102102
}
103103

104-
fun RustType.render(fullyQualified: Boolean): String {
104+
fun RustType.render(fullyQualified: Boolean = true): String {
105105
val namespace = if (fullyQualified) {
106106
this.namespace?.let { "$it::" } ?: ""
107107
} else ""

‎codegen/src/main/kotlin/software/amazon/smithy/rust/codegen/smithy/generators/BuilderGenerator.kt‎

Lines changed: 52 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -83,7 +83,8 @@ class OperationInputBuilderGenerator(
8383
).name
8484
}"
8585
rust(
86-
"""#T::MissingField { field: ${field.dq()}, details: ${detailedMessage.dq()}}""", runtimeConfig.operationBuildError()
86+
"""#T::MissingField { field: ${field.dq()}, details: ${detailedMessage.dq()}}""",
87+
runtimeConfig.operationBuildError()
8788
)
8889
}
8990

@@ -183,34 +184,69 @@ abstract class BuilderGenerator(
183184

184185
writer.rustBlock("impl $builderName") {
185186
members.forEach { member ->
186-
val memberName = symbolProvider.toMemberName(member)
187187
// All fields in the builder are optional
188188
val memberSymbol = symbolProvider.toSymbol(member)
189189
val outerType = memberSymbol.rustType()
190190
val coreType = outerType.stripOuter<RustType.Option>()
191-
val signature = when (coreType) {
192-
is RustType.String,
193-
is RustType.Box -> "(mut self, inp: impl Into<${coreType.render(true)}>) -> Self"
194-
else -> "(mut self, inp: ${coreType.render(true)}) -> Self"
195-
}
196-
writer.documentShape(member, model)
197-
writer.rustBlock("pub fn $memberName$signature") {
198-
write("self.$memberName = Some(${builderConverter(coreType)});")
199-
write("self")
191+
val memberName = symbolProvider.toMemberName(member)
192+
// Render a context-aware builder method for certain types, eg. a method for vectors that automatically
193+
// appends
194+
when (coreType) {
195+
is RustType.Vec -> renderVecHelper(memberName, coreType)
196+
is RustType.HashMap -> renderMapHelper(memberName, coreType)
197+
else -> {
198+
val signature = when (coreType) {
199+
is RustType.String,
200+
is RustType.Box -> "(mut self, inp: impl Into<${coreType.render(true)}>) -> Self"
201+
else -> "(mut self, inp: ${coreType.render(true)}) -> Self"
202+
}
203+
writer.documentShape(member, model)
204+
writer.rustBlock("pub fn $memberName$signature") {
205+
write("self.$memberName = Some(${builderConverter(coreType)});")
206+
write("self")
207+
}
208+
}
200209
}
201210

211+
// Render a `set_foo` method. This is useful as a target for code generation, because the argument type
212+
// is exactly the same as the resulting member type.
202213
writer.rustBlock("pub fn ${member.setterName()}(mut self, inp: ${outerType.render(true)}) -> Self") {
203214
val v = "inp".letIf(outerType !is RustType.Option) {
204215
"Some($it)"
205216
}
206217
rust("self.$memberName = $v; self")
207218
}
208219
}
209-
210220
buildFn(this)
211221
}
212222
}
213223

224+
private fun RustWriter.renderVecHelper(memberName: String, coreType: RustType.Vec) {
225+
rustBlock("pub fn $memberName(mut self, inp: impl Into<${coreType.member.render(true)}>) -> Self") {
226+
rust(
227+
"""
228+
let mut v = self.$memberName.unwrap_or_default();
229+
v.push(inp.into());
230+
self.$memberName = Some(v);
231+
self
232+
"""
233+
)
234+
}
235+
}
236+
237+
private fun RustWriter.renderMapHelper(memberName: String, coreType: RustType.HashMap) {
238+
rustBlock("pub fn $memberName(mut self, k: impl Into<${coreType.key.render(true)}>, v: impl Into<${coreType.member.render(true)}>) -> Self") {
239+
rust(
240+
"""
241+
let mut hash_map = self.$memberName.unwrap_or_default();
242+
hash_map.insert(k.into(), v.into());
243+
self.$memberName = Some(hash_map);
244+
self
245+
"""
246+
)
247+
}
248+
}
249+
214250
abstract fun RustWriter.missingRequiredField(field: String)
215251

216252
abstract fun buildFn(implBlockWriter: RustWriter)
@@ -236,7 +272,10 @@ abstract class BuilderGenerator(
236272
// Write the modifier
237273
when {
238274
!memberSymbol.isOptional() && default == Default.RustDefault -> rust(".unwrap_or_default()")
239-
!memberSymbol.isOptional() -> withBlock(".ok_or(", ")?") { missingRequiredField(memberName) }
275+
!memberSymbol.isOptional() -> withBlock(
276+
".ok_or(",
277+
")?"
278+
) { missingRequiredField(memberName) }
240279
memberSymbol.isOptional() && default is Default.Custom -> {
241280
withBlock(".or_else(||Some(", "))") { default.render(this) }
242281
}

0 commit comments

Comments
 (0)