refactor: add buffer to iam token expiry and create token expiry time…

… earlier (#706)
aws · Oct 31, 2023 · d6bb76f · d6bb76f
1 parent 2b2c6e4
commit d6bb76f
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 4 deletions.
diff --git a/docs/using-the-jdbc-driver/using-plugins/UsingTheIamAuthenticationPlugin.md b/docs/using-the-jdbc-driver/using-plugins/UsingTheIamAuthenticationPlugin.md
@@ -33,7 +33,7 @@ IAM database authentication use is limited to certain database engines. For more
 | `iamDefaultPort`  | String  |    No    | This property will override the default port that is used to generate the IAM token. The default port is determined based on the underlying driver protocol. For now, there is support for `jdbc:postgresql:` and `jdbc:mysql:`. Target drivers with different protocols will require users to provide a default port. | `1234`                                              |
 | `iamHost`         | String  |    No    | This property will override the default hostname that is used to generate the IAM token. The default hostname is derived from the connection string. This parameter is required when users are connecting with custom endpoints.                                                                                       | `database.cluster-hash.us-east-1.rds.amazonaws.com` |
 | `iamRegion`       | String  |    No    | This property will override the default region that is used to generate the IAM token. The default region is parsed from the connection string.                                                                                                                                                                        | `us-east-2`                                         |
-| `iamExpiration`   | Integer |    No    | This property will override the default expiration time that is assigned to the generated IAM token. The default expiration time is set to be 15 minutes.                                                                                                                                                              | `600`                                               |
+| `iamExpiration`   | Integer |    No    | This property determines how long an IAM token is kept in the driver cache before a new one is generated. The default expiration time is set to be 14 minutes and 30 seconds. Note that IAM database authentication tokens have a lifetime of 15 minutes.                                                              | `600`                                               |
 
 ## Sample code
 [AwsIamAuthenticationPostgresqlExample.java](../../../examples/AWSDriverExample/src/main/java/software/amazon/AwsIamAuthenticationPostgresqlExample.java)<br>

diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/IamAuthConnectionPlugin.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/IamAuthConnectionPlugin.java
@@ -57,7 +57,7 @@ public class IamAuthConnectionPlugin extends AbstractConnectionPlugin {
         }
       });
   static final ConcurrentHashMap<String, TokenInfo> tokenCache = new ConcurrentHashMap<>();
-  private static final int DEFAULT_TOKEN_EXPIRATION_SEC = 15 * 60;
+  private static final int DEFAULT_TOKEN_EXPIRATION_SEC = 15 * 60 - 30;
 
   public static final AwsWrapperProperty IAM_HOST = new AwsWrapperProperty(
       "iamHost", null,
@@ -144,6 +144,7 @@ private Connection connectInternal(String driverProtocol, HostSpec hostSpec, Pro
               new Object[] {tokenInfo.getToken()}));
       PropertyDefinition.PASSWORD.set(props, tokenInfo.getToken());
     } else {
+      final Instant tokenExpiry = Instant.now().plus(tokenExpirationSec, ChronoUnit.SECONDS);
       final String token = generateAuthenticationToken(
           hostSpec,
           props,
@@ -157,7 +158,7 @@ private Connection connectInternal(String driverProtocol, HostSpec hostSpec, Pro
       PropertyDefinition.PASSWORD.set(props, token);
       tokenCache.put(
           cacheKey,
-          new TokenInfo(token, Instant.now().plus(tokenExpirationSec, ChronoUnit.SECONDS)));
+          new TokenInfo(token, tokenExpiry));
     }
 
     try {
@@ -176,6 +177,7 @@ private Connection connectInternal(String driverProtocol, HostSpec hostSpec, Pro
       // Login unsuccessful with cached token
       // Try to generate a new token and try to connect again
 
+      final Instant tokenExpiry = Instant.now().plus(tokenExpirationSec, ChronoUnit.SECONDS);
       final String token = generateAuthenticationToken(
           hostSpec,
           props,
@@ -189,7 +191,7 @@ private Connection connectInternal(String driverProtocol, HostSpec hostSpec, Pro
       PropertyDefinition.PASSWORD.set(props, token);
       tokenCache.put(
           cacheKey,
-          new TokenInfo(token, Instant.now().plus(tokenExpirationSec, ChronoUnit.SECONDS)));
+          new TokenInfo(token, tokenExpiry));
 
       return connectFunc.call();