Refactor based on upstream changes, simplify parameter addition, chan…

…ge `add` to `put`
awslabs · Jan 9, 2025 · 1c391b8 · 1c391b8
1 parent cc6961f
commit 1c391b8
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 13 deletions.
diff --git a/services/dsql/common/src/aws/sdk/kotlin/services/dsql/DsqlAuthTokenGenerator.kt b/services/dsql/common/src/aws/sdk/kotlin/services/dsql/DsqlAuthTokenGenerator.kt
@@ -28,7 +28,7 @@ public class DsqlAuthTokenGenerator(
     public val signer: AwsSigner = DefaultAwsSigner,
     public val clock: Clock = Clock.System,
 ) {
-    private val generator = AuthTokenGenerator("dsql", credentialsProvider, credentialsRefreshBuffer, signer, clock)
+    private val generator = AuthTokenGenerator("dsql", credentialsProvider, signer, clock)
 
     /**
      * Generates an auth token for the DbConnect action.
@@ -38,11 +38,7 @@ public class DsqlAuthTokenGenerator(
      */
     public suspend fun generateDbConnectAuthToken(endpoint: Url, region: String, expiration: Duration = 900.seconds): String {
         val dbConnectEndpoint = endpoint.toBuilder().apply {
-            parameters.apply {
-                decodedParameters {
-                    add("Action", "DbConnect")
-                }
-            }
+            parameters.decodedParameters.put("Action", "DbConnect")
         }.build()
 
         return generator.generateAuthToken(dbConnectEndpoint, region, expiration)
@@ -56,11 +52,7 @@ public class DsqlAuthTokenGenerator(
      */
     public suspend fun generateDbConnectAdminAuthToken(endpoint: Url, region: String, expiration: Duration = 900.seconds): String {
         val dbConnectAdminEndpoint = endpoint.toBuilder().apply {
-            parameters.apply {
-                decodedParameters {
-                    add("Action", "DbConnectAdmin")
-                }
-            }
+            parameters.decodedParameters.put("Action", "DbConnectAdmin")
         }.build()
 
         return generator.generateAuthToken(dbConnectAdminEndpoint, region, expiration)

diff --git a/services/rds/common/src/aws/sdk/kotlin/services/rds/RdsAuthTokenGenerator.kt b/services/rds/common/src/aws/sdk/kotlin/services/rds/RdsAuthTokenGenerator.kt
@@ -29,7 +29,7 @@ public class RdsAuthTokenGenerator(
     public val signer: AwsSigner = DefaultAwsSigner,
     public val clock: Clock = Clock.System,
 ) {
-    private val generator = AuthTokenGenerator("rds-db", credentialsProvider, credentialsRefreshBuffer, signer, clock)
+    private val generator = AuthTokenGenerator("rds-db", credentialsProvider, signer, clock)
 
     /**
      * Generates an auth token for the `connect` action.

diff --git a/services/rds/common/test/RdsAuthTokenGeneratorTest.kt b/services/rds/common/test/RdsAuthTokenGeneratorTest.kt
@@ -20,7 +20,6 @@ class RdsAuthTokenGeneratorTest {
     @Test
     fun testGenerateAuthToken() = runTest {
         val clock = ManualClock(Instant.fromEpochSeconds(1724716800))
-        println(clock.now())
 
         val credentials = Credentials("akid", "secret")
         val credentialsProvider = StaticCredentialsProvider(credentials)