Skip to content

Commit fafaa12

Browse files
lauzadislauzadis
authored
misc: add x-amzn-query-mode customization (#1455)
1 parent 78b53b6 commit fafaa12

File tree

4 files changed

+164
-0
lines changed

4 files changed

+164
-0
lines changed

‎.changes/20e287a5-cee0-4a5f-9d08-4022dcdee843.json‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
{
2+
"id": "20e287a5-cee0-4a5f-9d08-4022dcdee843",
3+
"type": "misc",
4+
"description": "Send x-amzn-query-mode=true for services with query-compatible trait"
5+
}

‎codegen/aws-sdk-codegen/src/main/kotlin/aws/sdk/kotlin/codegen/customization/AwsQueryModeCustomization.kt‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,28 @@
1+
/*
2+
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
package aws.sdk.kotlin.codegen.customization
6+
7+
import software.amazon.smithy.aws.traits.protocols.AwsQueryCompatibleTrait
8+
import software.amazon.smithy.kotlin.codegen.KotlinSettings
9+
import software.amazon.smithy.kotlin.codegen.integration.KotlinIntegration
10+
import software.amazon.smithy.kotlin.codegen.model.hasTrait
11+
import software.amazon.smithy.kotlin.codegen.rendering.protocol.MutateHeadersMiddleware
12+
import software.amazon.smithy.kotlin.codegen.rendering.protocol.ProtocolGenerator
13+
import software.amazon.smithy.kotlin.codegen.rendering.protocol.ProtocolMiddleware
14+
import software.amazon.smithy.model.Model
15+
16+
/**
17+
* Send an extra `x-amzn-query-mode` header with a value of `true` for services which have the [AwsQueryCompatibleTrait] applied.
18+
*/
19+
class AwsQueryModeCustomization : KotlinIntegration {
20+
override fun enabledForService(model: Model, settings: KotlinSettings): Boolean =
21+
model
22+
.getShape(settings.service)
23+
.get()
24+
.hasTrait<AwsQueryCompatibleTrait>()
25+
26+
override fun customizeMiddleware(ctx: ProtocolGenerator.GenerationContext, resolved: List<ProtocolMiddleware>): List<ProtocolMiddleware> =
27+
resolved + MutateHeadersMiddleware(extraHeaders = mapOf("x-amzn-query-mode" to "true"))
28+
}

‎codegen/aws-sdk-codegen/src/main/resources/META-INF/services/software.amazon.smithy.kotlin.codegen.integration.KotlinIntegration‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,3 +43,4 @@ aws.sdk.kotlin.codegen.customization.s3.express.S3ExpressIntegration
4343
aws.sdk.kotlin.codegen.customization.s3.S3ExpiresIntegration
4444
aws.sdk.kotlin.codegen.BusinessMetricsIntegration
4545
aws.sdk.kotlin.codegen.smoketests.SmokeTestsDenyListIntegration
46+
aws.sdk.kotlin.codegen.customization.AwsQueryModeCustomization

‎codegen/aws-sdk-codegen/src/test/kotlin/aws/sdk/kotlin/codegen/customization/AwsQueryModeCustomizationTest.kt‎

Lines changed: 130 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,130 @@
1+
/*
2+
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
package aws.sdk.kotlin.codegen.customization
6+
7+
import aws.sdk.kotlin.codegen.testutil.lines
8+
import software.amazon.smithy.kotlin.codegen.test.*
9+
import kotlin.test.Test
10+
import kotlin.test.assertFalse
11+
import kotlin.test.assertTrue
12+
13+
class AwsQueryModeCustomizationTest {
14+
private val queryCompatibleModel = """
15+
namespace com.test
16+
17+
use aws.protocols#awsJson1_0
18+
use aws.protocols#awsQueryCompatible
19+
use aws.api#service
20+
21+
@awsJson1_0
22+
@awsQueryCompatible
23+
@service(sdkId: "QueryCompatible")
24+
service QueryCompatible {
25+
version: "1.0.0",
26+
operations: [GetFoo]
27+
}
28+
29+
@http(method: "POST", uri: "/foo")
30+
operation GetFoo {
31+
input: GetFooInput
32+
}
33+
34+
structure GetFooInput {
35+
payload: String
36+
}
37+
"""
38+
.trimIndent()
39+
.toSmithyModel()
40+
41+
private val nonQueryCompatibleModel = """
42+
namespace com.test
43+
44+
use aws.protocols#awsJson1_0
45+
use aws.protocols#awsQueryCompatible
46+
use aws.api#service
47+
48+
@awsJson1_0
49+
@service(sdkId: "NonQueryCompatible")
50+
service NonQueryCompatible {
51+
version: "1.0.0",
52+
operations: [GetFoo]
53+
}
54+
55+
@http(method: "POST", uri: "/foo")
56+
operation GetFoo {
57+
input: GetFooInput
58+
}
59+
60+
structure GetFooInput {
61+
payload: String
62+
}
63+
"""
64+
.trimIndent()
65+
.toSmithyModel()
66+
67+
@Test
68+
fun testEnabledForQueryCompatibleModel() {
69+
assertTrue {
70+
AwsQueryModeCustomization()
71+
.enabledForService(queryCompatibleModel, queryCompatibleModel.defaultSettings())
72+
}
73+
}
74+
75+
@Test
76+
fun testNotExpectedForNonQueryCompatibleModel() {
77+
assertFalse {
78+
AwsQueryModeCustomization()
79+
.enabledForService(nonQueryCompatibleModel, nonQueryCompatibleModel.defaultSettings())
80+
}
81+
}
82+
83+
@Test
84+
fun `x-amzn-query-mode applied`() {
85+
val ctx = queryCompatibleModel.newTestContext("QueryCompatible", integrations = listOf(AwsQueryModeCustomization()))
86+
val generator = MockHttpProtocolGenerator(queryCompatibleModel)
87+
generator.generateProtocolClient(ctx.generationCtx)
88+
89+
ctx.generationCtx.delegator.finalize()
90+
ctx.generationCtx.delegator.flushWriters()
91+
92+
val actual = ctx.manifest.expectFileString("/src/main/kotlin/com/test/DefaultTestClient.kt")
93+
94+
val getFooMethod = actual.lines(" override suspend fun getFoo(input: GetFooRequest): GetFooResponse {", " }")
95+
96+
val expectedHeaderMutation = """
97+
op.install(
98+
MutateHeaders().apply {
99+
append("x-amzn-query-mode", "true")
100+
}
101+
)
102+
""".replaceIndent(" ")
103+
104+
getFooMethod.shouldContainOnlyOnceWithDiff(expectedHeaderMutation)
105+
}
106+
107+
@Test
108+
fun `x-amzn-query-mode NOT applied`() {
109+
val ctx = nonQueryCompatibleModel.newTestContext("NonQueryCompatible", integrations = listOf())
110+
val generator = MockHttpProtocolGenerator(nonQueryCompatibleModel)
111+
generator.generateProtocolClient(ctx.generationCtx)
112+
113+
ctx.generationCtx.delegator.finalize()
114+
ctx.generationCtx.delegator.flushWriters()
115+
116+
val actual = ctx.manifest.expectFileString("/src/main/kotlin/com/test/DefaultTestClient.kt")
117+
118+
val getFooMethod = actual.lines(" override suspend fun getFoo(input: GetFooRequest): GetFooResponse {", " }")
119+
120+
val unexpectedHeaderMutation = """
121+
op.install(
122+
MutateHeaders().apply {
123+
append("x-amzn-query-mode", "true")
124+
}
125+
)
126+
""".replaceIndent(" ")
127+
128+
getFooMethod.shouldNotContainOnlyOnceWithDiff(unexpectedHeaderMutation)
129+
}
130+
}

0 commit comments

Comments
 (0)