Skip to content

Commit 48c0ee6

Browse files
authored
feat: enable endpoint discovery (and Timestream Query and Timestream Write) (#976)
1 parent 5b5fdd7 commit 48c0ee6

File tree

5 files changed

+136
-27
lines changed

5 files changed

+136
-27
lines changed
Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
{
2+
"id": "59b05d2f-0ec8-457f-b08a-b3f902c654e5",
3+
"type": "feature",
4+
"description": "Enable **Timestream Query** and **Timestream Write** service clients",
5+
"issues": [
6+
"awslabs/smithy-kotlin#146",
7+
"awslabs/aws-sdk-kotlin#800"
8+
]
9+
}

codegen/sdk/build.gradle.kts

Lines changed: 0 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -94,13 +94,6 @@ data class AwsService(
9494

9595
)
9696

97-
val disabledServices = setOf(
98-
// timestream requires endpoint discovery
99-
// https://github.com/awslabs/smithy-kotlin/issues/146
100-
"timestream-write",
101-
"timestream-query",
102-
)
103-
10497
// Manually create the projections rather than using the extension to avoid unnecessary configuration evaluation.
10598
// Otherwise we would be reading the models from disk on every gradle invocation for unrelated projects/tasks
10699
fun awsServiceProjections(): Provider<List<SmithyProjection>> {
@@ -197,11 +190,6 @@ fun fileToService(applyFilters: Boolean): (File) -> AwsService? = { file: File -
197190
null
198191
}
199192

200-
applyFilters && filename in disabledServices -> {
201-
logger.warn("skipping ${file.absolutePath}, it is explicitly disabled")
202-
null
203-
}
204-
205193
else -> {
206194
logger.info("discovered service: ${serviceTrait.sdkId}")
207195
AwsService(

codegen/smithy-aws-kotlin-codegen/src/test/kotlin/aws/sdk/kotlin/codegen/PresignerGeneratorTest.kt

Lines changed: 1 addition & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -6,17 +6,11 @@ package aws.sdk.kotlin.codegen
66

77
import aws.sdk.kotlin.codegen.model.traits.Presignable
88
import software.amazon.smithy.build.MockManifest
9-
import software.amazon.smithy.codegen.core.SymbolProvider
10-
import software.amazon.smithy.kotlin.codegen.KotlinSettings
11-
import software.amazon.smithy.kotlin.codegen.core.CodegenContext
12-
import software.amazon.smithy.kotlin.codegen.integration.KotlinIntegration
139
import software.amazon.smithy.kotlin.codegen.model.expectShape
14-
import software.amazon.smithy.kotlin.codegen.rendering.protocol.ProtocolGenerator
15-
import software.amazon.smithy.kotlin.codegen.test.TestContext
1610
import software.amazon.smithy.kotlin.codegen.test.newTestContext
1711
import software.amazon.smithy.kotlin.codegen.test.shouldContainOnlyOnceWithDiff
12+
import software.amazon.smithy.kotlin.codegen.test.toCodegenContext
1813
import software.amazon.smithy.kotlin.codegen.test.toSmithyModel
19-
import software.amazon.smithy.model.Model
2014
import software.amazon.smithy.model.shapes.OperationShape
2115
import kotlin.test.Test
2216
import kotlin.test.assertTrue
@@ -291,11 +285,3 @@ class PresignerGeneratorTest {
291285
actual.shouldContainOnlyOnceWithDiff(expected)
292286
}
293287
}
294-
295-
private fun TestContext.toCodegenContext() = object : CodegenContext {
296-
override val model: Model = generationCtx.model
297-
override val symbolProvider: SymbolProvider = generationCtx.symbolProvider
298-
override val settings: KotlinSettings = generationCtx.settings
299-
override val protocolGenerator: ProtocolGenerator = generator
300-
override val integrations: List<KotlinIntegration> = generationCtx.integrations
301-
}
Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,63 @@
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.services.timestreamquery
6+
7+
import aws.sdk.kotlin.services.timestreamquery.model.DescribeEndpointsResponse
8+
import aws.sdk.kotlin.services.timestreamquery.model.ListScheduledQueriesResponse
9+
import aws.smithy.kotlin.runtime.client.ResponseInterceptorContext
10+
import aws.smithy.kotlin.runtime.client.operationName
11+
import aws.smithy.kotlin.runtime.http.interceptors.HttpInterceptor
12+
import aws.smithy.kotlin.runtime.http.request.HttpRequest
13+
import aws.smithy.kotlin.runtime.http.response.HttpResponse
14+
import kotlinx.coroutines.runBlocking
15+
import org.junit.jupiter.api.Assertions.assertEquals
16+
import org.junit.jupiter.api.Assertions.assertNotNull
17+
import org.junit.jupiter.api.Test
18+
import org.junit.jupiter.api.TestInstance
19+
20+
@TestInstance(TestInstance.Lifecycle.PER_CLASS)
21+
class EndpointDiscoveryTest {
22+
@Test
23+
fun testEndpointDiscovery() = runBlocking {
24+
val operationLog = mutableListOf<String>()
25+
var discoveredHost: String? = null
26+
27+
val interceptor = object : HttpInterceptor {
28+
override fun readAfterExecution(
29+
context: ResponseInterceptorContext<Any, Any, HttpRequest?, HttpResponse?>,
30+
) {
31+
operationLog += context.executionContext.operationName!!
32+
33+
when (val response = context.response.getOrNull()) {
34+
is DescribeEndpointsResponse -> discoveredHost = response.endpoints!!.first().address
35+
36+
is ListScheduledQueriesResponse -> {
37+
// Make sure every request _except_ DescribeEndpoints uses the discovered endpoint
38+
assertNotNull(discoveredHost)
39+
assertEquals(discoveredHost, context.protocolRequest!!.url.host.toString())
40+
}
41+
42+
else -> error("Unexpected response ${context.response}")
43+
}
44+
}
45+
}
46+
47+
TimestreamQueryClient {
48+
region = "us-west-2"
49+
interceptors += interceptor
50+
}.use { tsq ->
51+
tsq.listScheduledQueries()
52+
53+
// Have to discover the endpoint the first time
54+
assertEquals(listOf("DescribeEndpoints", "ListScheduledQueries"), operationLog)
55+
operationLog.clear()
56+
57+
tsq.listScheduledQueries()
58+
59+
// Don't have to discover the endpoint again because it's cached
60+
assertEquals(listOf("ListScheduledQueries"), operationLog)
61+
}
62+
}
63+
}
Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,63 @@
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.services.timestreamwrite
6+
7+
import aws.sdk.kotlin.services.timestreamwrite.model.DescribeEndpointsResponse
8+
import aws.sdk.kotlin.services.timestreamwrite.model.ListDatabasesResponse
9+
import aws.smithy.kotlin.runtime.client.ResponseInterceptorContext
10+
import aws.smithy.kotlin.runtime.client.operationName
11+
import aws.smithy.kotlin.runtime.http.interceptors.HttpInterceptor
12+
import aws.smithy.kotlin.runtime.http.request.HttpRequest
13+
import aws.smithy.kotlin.runtime.http.response.HttpResponse
14+
import kotlinx.coroutines.runBlocking
15+
import org.junit.jupiter.api.Assertions.assertEquals
16+
import org.junit.jupiter.api.Assertions.assertNotNull
17+
import org.junit.jupiter.api.Test
18+
import org.junit.jupiter.api.TestInstance
19+
20+
@TestInstance(TestInstance.Lifecycle.PER_CLASS)
21+
class EndpointDiscoveryTest {
22+
@Test
23+
fun testEndpointDiscovery() = runBlocking {
24+
val operationLog = mutableListOf<String>()
25+
var discoveredHost: String? = null
26+
27+
val interceptor = object : HttpInterceptor {
28+
override fun readAfterExecution(
29+
context: ResponseInterceptorContext<Any, Any, HttpRequest?, HttpResponse?>,
30+
) {
31+
operationLog += context.executionContext.operationName!!
32+
33+
when (val response = context.response.getOrNull()) {
34+
is DescribeEndpointsResponse -> discoveredHost = response.endpoints!!.first().address
35+
36+
is ListDatabasesResponse -> {
37+
// Make sure every request _except_ DescribeEndpoints uses the discovered endpoint
38+
assertNotNull(discoveredHost)
39+
assertEquals(discoveredHost, context.protocolRequest!!.url.host.toString())
40+
}
41+
42+
else -> error("Unexpected response ${context.response}")
43+
}
44+
}
45+
}
46+
47+
TimestreamWriteClient {
48+
region = "us-west-2"
49+
interceptors += interceptor
50+
}.use { tsw ->
51+
tsw.listDatabases()
52+
53+
// Have to discover the endpoint the first time
54+
assertEquals(listOf("DescribeEndpoints", "ListDatabases"), operationLog)
55+
operationLog.clear()
56+
57+
tsw.listDatabases()
58+
59+
// Don't have to discover the endpoint again because it's cached
60+
assertEquals(listOf("ListDatabases"), operationLog)
61+
}
62+
}
63+
}

0 commit comments

Comments
 (0)