Skip to content

Commit dfd53fb

Browse files
ysaito1001ysaito1001
authored
Make Rpc V2 CBOR awsQuery compatible (#4186)
## Description This PR makes Rpc V2 CBOR a compatible protocol for `awsQuery` using `awsQueryCompatible` trait, as described in the `Important` section in [this page](https://smithy.io/2.0/aws/protocols/aws-query-protocol.html#aws-protocols-awsquerycompatible-trait). Previously, the implementation for `awsQueryCompatible` was tightly coupled to `awsJson1_0`. This PR makes the implementation a bit more abstract so that the implementation can support more target protocols generically. ## Testing - CI - Made `AwsQueryCompatibleTest` parameterized tests to support RpcV2Cbor and verified against modified cloudwatch service model (with `awsQueryCompatible` trait applied). ~Tests for RpcV2Cbor is commented out till #4185 is merged~. ## Checklist <!--- If a checkbox below is not applicable, then please DELETE it rather than leaving it unchecked --> - [x] For changes to the smithy-rs codegen or runtime crates, I have created a changelog entry Markdown file in the `.changelog` directory, specifying "client," "server," or both in the `applies_to` key. - [x] For changes to the AWS SDK, generated SDK code, or SDK runtime crates, I have created a changelog entry Markdown file in the `.changelog` directory, specifying "aws-sdk-rust" in the `applies_to` key. ---- _By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice._
1 parent 8b32dee commit dfd53fb

File tree

5 files changed

+181
-57
lines changed

5 files changed

+181
-57
lines changed

‎.changelog/1750773066.md

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
---
2+
applies_to:
3+
- aws-sdk-rust
4+
- client
5+
authors:
6+
- ysaito1001
7+
references:
8+
- smithy-rs#4186
9+
breaking: false
10+
new_feature: false
11+
bug_fix: false
12+
---
13+
Make Rpc V2 CBOR a compatible protocol for `awsQuery` using `awsQueryCompatible` trait

‎codegen-client/build.gradle.kts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@ dependencies {
2424
implementation(project(":codegen-core"))
2525
implementation(kotlin("stdlib-jdk8"))
2626
api("software.amazon.smithy:smithy-codegen-core:$smithyVersion")
27+
implementation("com.fasterxml.jackson.dataformat:jackson-dataformat-cbor:2.13.0")
2728
implementation("software.amazon.smithy:smithy-aws-traits:$smithyVersion")
2829
implementation("software.amazon.smithy:smithy-protocol-test-traits:$smithyVersion")
2930
implementation("software.amazon.smithy:smithy-waiters:$smithyVersion")

‎codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/protocols/ClientProtocolLoader.kt

Lines changed: 45 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,13 +16,17 @@ import software.amazon.smithy.model.shapes.ServiceShape
1616
import software.amazon.smithy.protocol.traits.Rpcv2CborTrait
1717
import software.amazon.smithy.rust.codegen.client.smithy.ClientCodegenContext
1818
import software.amazon.smithy.rust.codegen.client.smithy.generators.OperationGenerator
19+
import software.amazon.smithy.rust.codegen.core.rustlang.rustTemplate
20+
import software.amazon.smithy.rust.codegen.core.rustlang.writable
1921
import software.amazon.smithy.rust.codegen.core.smithy.CodegenContext
22+
import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType
2023
import software.amazon.smithy.rust.codegen.core.smithy.generators.protocol.ProtocolSupport
2124
import software.amazon.smithy.rust.codegen.core.smithy.protocols.AwsJson
2225
import software.amazon.smithy.rust.codegen.core.smithy.protocols.AwsJsonVersion
2326
import software.amazon.smithy.rust.codegen.core.smithy.protocols.AwsQueryCompatible
2427
import software.amazon.smithy.rust.codegen.core.smithy.protocols.AwsQueryProtocol
2528
import software.amazon.smithy.rust.codegen.core.smithy.protocols.Ec2QueryProtocol
29+
import software.amazon.smithy.rust.codegen.core.smithy.protocols.ParseErrorMetadataParams
2630
import software.amazon.smithy.rust.codegen.core.smithy.protocols.Protocol
2731
import software.amazon.smithy.rust.codegen.core.smithy.protocols.ProtocolGeneratorFactory
2832
import software.amazon.smithy.rust.codegen.core.smithy.protocols.ProtocolLoader
@@ -67,7 +71,23 @@ private class ClientAwsJsonFactory(private val version: AwsJsonVersion) :
6771
ProtocolGeneratorFactory<OperationGenerator, ClientCodegenContext> {
6872
override fun protocol(codegenContext: ClientCodegenContext): Protocol =
6973
if (compatibleWithAwsQuery(codegenContext.serviceShape, version)) {
70-
AwsQueryCompatible(codegenContext, AwsJson(codegenContext, version))
74+
AwsQueryCompatible(
75+
codegenContext, AwsJson(codegenContext, version),
76+
ParseErrorMetadataParams(
77+
RuntimeType.smithyJson(codegenContext.runtimeConfig)
78+
.resolve("deserialize::error::DeserializeError"),
79+
writable {
80+
rustTemplate(
81+
"""
82+
#{parse_error_metadata}(response_body, response_headers)?
83+
""",
84+
"parse_error_metadata" to
85+
RuntimeType.jsonErrors(codegenContext.runtimeConfig)
86+
.resolve("parse_error_metadata"),
87+
)
88+
},
89+
),
90+
)
7191
} else {
7292
AwsJson(codegenContext, version)
7393
}
@@ -122,10 +142,33 @@ class ClientRestXmlFactory(
122142
}
123143

124144
class ClientRpcV2CborFactory : ProtocolGeneratorFactory<OperationGenerator, ClientCodegenContext> {
125-
override fun protocol(codegenContext: ClientCodegenContext): Protocol = RpcV2Cbor(codegenContext)
145+
override fun protocol(codegenContext: ClientCodegenContext): Protocol =
146+
if (compatibleWithAwsQuery(codegenContext.serviceShape)) {
147+
AwsQueryCompatible(
148+
codegenContext, RpcV2Cbor(codegenContext),
149+
ParseErrorMetadataParams(
150+
RuntimeType.smithyCbor(codegenContext.runtimeConfig)
151+
.resolve("decode::DeserializeError"),
152+
writable {
153+
rustTemplate(
154+
"""
155+
#{parse_error_metadata}(_response_status, response_headers, response_body)?
156+
""",
157+
"parse_error_metadata" to
158+
RuntimeType.cborErrors(codegenContext.runtimeConfig)
159+
.resolve("parse_error_metadata"),
160+
)
161+
},
162+
),
163+
)
164+
} else {
165+
RpcV2Cbor(codegenContext)
166+
}
126167

127168
override fun buildProtocolGenerator(codegenContext: ClientCodegenContext): OperationGenerator =
128169
OperationGenerator(codegenContext, protocol(codegenContext))
129170

130171
override fun support(): ProtocolSupport = CLIENT_PROTOCOL_SUPPORT
172+
173+
private fun compatibleWithAwsQuery(serviceShape: ServiceShape) = serviceShape.hasTrait<AwsQueryCompatibleTrait>()
131174
}

‎codegen-client/src/test/kotlin/software/amazon/smithy/rust/codegen/client/smithy/protocols/AwsQueryCompatibleTest.kt

Lines changed: 88 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -5,26 +5,89 @@
55

66
package software.amazon.smithy.rust.codegen.client.smithy.protocols
77

8-
import org.junit.jupiter.api.Test
8+
import com.fasterxml.jackson.databind.ObjectMapper
9+
import com.fasterxml.jackson.dataformat.cbor.CBORFactory
10+
import org.junit.jupiter.api.extension.ExtensionContext
11+
import org.junit.jupiter.params.ParameterizedTest
12+
import org.junit.jupiter.params.provider.Arguments
13+
import org.junit.jupiter.params.provider.ArgumentsProvider
14+
import org.junit.jupiter.params.provider.ArgumentsSource
915
import software.amazon.smithy.rust.codegen.client.testutil.clientIntegrationTest
1016
import software.amazon.smithy.rust.codegen.core.rustlang.CargoDependency
17+
import software.amazon.smithy.rust.codegen.core.rustlang.Writable
18+
import software.amazon.smithy.rust.codegen.core.rustlang.rust
1119
import software.amazon.smithy.rust.codegen.core.rustlang.rustTemplate
20+
import software.amazon.smithy.rust.codegen.core.rustlang.writable
1221
import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType
1322
import software.amazon.smithy.rust.codegen.core.testutil.asSmithyModel
1423
import software.amazon.smithy.rust.codegen.core.testutil.testModule
1524
import software.amazon.smithy.rust.codegen.core.testutil.tokioTest
1625
import software.amazon.smithy.rust.codegen.core.util.letIf
26+
import java.util.stream.Stream
27+
28+
data class AwsQueryCompatibleTestInput(
29+
val protocolAnnotation: String,
30+
val payload: Writable,
31+
)
32+
33+
class AwsQueryCompatibleTestInputProvider : ArgumentsProvider {
34+
private fun jsonStringToBytesArrayString(jsonString: String): String {
35+
val jsonMapper = ObjectMapper()
36+
val cborMapper = ObjectMapper(CBORFactory())
37+
// Parse JSON string to a generic type.
38+
val jsonData = jsonMapper.readValue(jsonString, Any::class.java)
39+
// Convert the parsed data to CBOR.
40+
val bytes = cborMapper.writeValueAsBytes(jsonData)
41+
return bytes
42+
.joinToString(
43+
prefix = "&[",
44+
postfix = "]",
45+
transform = { "0x${it.toUByte().toString(16).padStart(2, '0')}u8" },
46+
)
47+
}
48+
49+
override fun provideArguments(context: ExtensionContext?): Stream<out Arguments> =
50+
listOf(
51+
AwsQueryCompatibleTestInput(
52+
"@awsJson1_0",
53+
writable {
54+
rust(
55+
"""
56+
r##"{
57+
"__type": "com.amazonaws.sqs##QueueDoesNotExist",
58+
"message": "Some user-visible message"
59+
}"##
60+
""",
61+
)
62+
},
63+
),
64+
AwsQueryCompatibleTestInput(
65+
"@rpcv2Cbor",
66+
writable {
67+
val bytesArray =
68+
jsonStringToBytesArrayString(
69+
"""
70+
{
71+
"__type": "com.amazonaws.sqs#QueueDoesNotExist",
72+
"message": "Some user-visible message"
73+
}
74+
""",
75+
)
76+
rust("#T::from_static($bytesArray)", RuntimeType.Bytes)
77+
},
78+
),
79+
).map { Arguments.of(it) }.stream()
80+
}
1781

1882
class AwsQueryCompatibleTest {
1983
companion object {
2084
const val prologue = """
2185
namespace test
86+
use smithy.protocols#rpcv2Cbor
2287
use aws.protocols#awsJson1_0
2388
use aws.protocols#awsQueryCompatible
2489
use aws.protocols#awsQueryError
2590
"""
26-
27-
const val awsjson10Trait = "@awsJson1_0"
2891
const val awsQueryCompatibleTrait = "@awsQueryCompatible"
2992

3093
fun testService(withAwsQueryError: Boolean = true) =
@@ -63,10 +126,13 @@ class AwsQueryCompatibleTest {
63126
}
64127
}
65128

66-
@Test
67-
fun `aws-query-compatible json with aws query error should allow for retrieving error code and type from custom header`() {
129+
@ParameterizedTest
130+
@ArgumentsSource(AwsQueryCompatibleTestInputProvider::class)
131+
fun `aws-query-compatible json with aws query error should allow for retrieving error code and type from custom header`(
132+
testInput: AwsQueryCompatibleTestInput,
133+
) {
68134
val model =
69-
(prologue + awsQueryCompatibleTrait + awsjson10Trait + testService()).asSmithyModel(
135+
(prologue + awsQueryCompatibleTrait + testInput.protocolAnnotation + testService()).asSmithyModel(
70136
smithyVersion = "2",
71137
)
72138
clientIntegrationTest(model) { context, rustCrate ->
@@ -82,12 +148,7 @@ class AwsQueryCompatibleTest {
82148
)
83149
.status(400)
84150
.body(
85-
#{SdkBody}::from(
86-
r##"{
87-
"__type": "com.amazonaws.sqs##QueueDoesNotExist",
88-
"message": "Some user-visible message"
89-
}"##
90-
)
151+
#{SdkBody}::from(#{payload:W})
91152
)
92153
.unwrap()
93154
};
@@ -105,6 +166,7 @@ class AwsQueryCompatibleTest {
105166
assert_eq!(#{Some}("Sender"), error.meta().extra("type"));
106167
""",
107168
*RuntimeType.preludeScope,
169+
"payload" to testInput.payload,
108170
"SdkBody" to RuntimeType.sdkBody(context.runtimeConfig),
109171
"infallible_client_fn" to
110172
CargoDependency.smithyHttpClientTestUtil(context.runtimeConfig)
@@ -116,10 +178,13 @@ class AwsQueryCompatibleTest {
116178
}
117179
}
118180

119-
@Test
120-
fun `aws-query-compatible json without aws query error should allow for retrieving error code from payload`() {
181+
@ParameterizedTest
182+
@ArgumentsSource(AwsQueryCompatibleTestInputProvider::class)
183+
fun `aws-query-compatible json without aws query error should allow for retrieving error code from payload`(
184+
testInput: AwsQueryCompatibleTestInput,
185+
) {
121186
val model =
122-
(prologue + awsQueryCompatibleTrait + awsjson10Trait + testService(withAwsQueryError = false)).asSmithyModel(
187+
(prologue + awsQueryCompatibleTrait + testInput.protocolAnnotation + testService(withAwsQueryError = false)).asSmithyModel(
123188
smithyVersion = "2",
124189
)
125190
clientIntegrationTest(model) { context, rustCrate ->
@@ -131,12 +196,7 @@ class AwsQueryCompatibleTest {
131196
#{http_1x}::Response::builder()
132197
.status(400)
133198
.body(
134-
#{SdkBody}::from(
135-
r##"{
136-
"__type": "com.amazonaws.sqs##QueueDoesNotExist",
137-
"message": "Some user-visible message"
138-
}"##,
139-
)
199+
#{SdkBody}::from(#{payload:W})
140200
)
141201
.unwrap()
142202
};
@@ -152,6 +212,7 @@ class AwsQueryCompatibleTest {
152212
""",
153213
*RuntimeType.preludeScope,
154214
"SdkBody" to RuntimeType.sdkBody(context.runtimeConfig),
215+
"payload" to testInput.payload,
155216
"infallible_client_fn" to
156217
CargoDependency.smithyHttpClientTestUtil(context.runtimeConfig)
157218
.toType().resolve("test_util::infallible_client_fn"),
@@ -162,10 +223,13 @@ class AwsQueryCompatibleTest {
162223
}
163224
}
164225

165-
@Test
166-
fun `request header should include x-amzn-query-mode when the service has the awsQueryCompatible trait`() {
226+
@ParameterizedTest
227+
@ArgumentsSource(AwsQueryCompatibleTestInputProvider::class)
228+
fun `request header should include x-amzn-query-mode when the service has the awsQueryCompatible trait`(
229+
testInput: AwsQueryCompatibleTestInput,
230+
) {
167231
val model =
168-
(prologue + awsQueryCompatibleTrait + awsjson10Trait + testService()).asSmithyModel(
232+
(prologue + awsQueryCompatibleTrait + testInput.protocolAnnotation + testService()).asSmithyModel(
169233
smithyVersion = "2",
170234
)
171235
clientIntegrationTest(model) { context, rustCrate ->

0 commit comments

Comments
 (0)