add connect exception handling to IAM plugin (#456)

Co-authored-by: sergiyv-bitquill <[email protected]>

Author: sergiyvamz

wrapper/src/main/java/software/amazon/jdbc/plugin/IamAuthConnectionPlugin.java

@@ -131,8 +131,9 @@ private Connection connectInternal(String driverProtocol, HostSpec hostSpec, Pro
         port,
         region);
     final TokenInfo tokenInfo = tokenCache.get(cacheKey);
+    final boolean isCachedToken = tokenInfo != null && !tokenInfo.isExpired();
 
-    if (tokenInfo != null && !tokenInfo.isExpired()) {
+    if (isCachedToken) {
       LOGGER.finest(
           () -> Messages.get(
               "IamAuthConnectionPlugin.useCachedIamToken",
@@ -154,7 +155,47 @@ private Connection connectInternal(String driverProtocol, HostSpec hostSpec, Pro
           cacheKey,
           new TokenInfo(token, Instant.now().plus(tokenExpirationSec, ChronoUnit.SECONDS)));
     }
-    return connectFunc.call();
+
+    try {
+      return connectFunc.call();
+    } catch (final SQLException exception) {
+
+      LOGGER.finest(
+          () -> Messages.get(
+              "IamAuthConnectionPlugin.connectException",
+              new Object[] {exception}));
+
+      if (!this.pluginService.isLoginException(exception) || !isCachedToken) {
+        throw exception;
+      }
+
+      // Login unsuccessful with cached token
+      // Try to generate a new token and try to connect again
+
+      final String token = generateAuthenticationToken(
+          hostSpec,
+          props,
+          host,
+          port,
+          region);
+      LOGGER.finest(
+          () -> Messages.get(
+              "IamAuthConnectionPlugin.generatedNewIamToken",
+              new Object[] {token}));
+      PropertyDefinition.PASSWORD.set(props, token);
+      tokenCache.put(
+          cacheKey,
+          new TokenInfo(token, Instant.now().plus(tokenExpirationSec, ChronoUnit.SECONDS)));
+
+      return connectFunc.call();
+
+    } catch (final Exception exception) {
+      LOGGER.warning(
+          () -> Messages.get(
+              "IamAuthConnectionPlugin.unhandledException",
+              new Object[] {exception}));
+      throw new SQLException(exception);
+    }
   }
 
   @Override

wrapper/src/main/resources/messages.properties

@@ -154,6 +154,8 @@ IamAuthConnectionPlugin.unsupportedHostname=Unsupported AWS hostname {0}. Amazon
 IamAuthConnectionPlugin.useCachedIamToken=Use cached IAM token = ''{0}''
 IamAuthConnectionPlugin.generatedNewIamToken=Generated new IAM token = ''{0}''
 IamAuthConnectionPlugin.invalidPort=Port number: {0} is not valid. Port number should be greater than zero.
+IamAuthConnectionPlugin.unhandledException=Unhandled exception: ''{0}''
+IamAuthConnectionPlugin.connectException=Error occurred while opening a connection: ''{0}''
 
 # Log Query Connection Plugin
 LogQueryConnectionPlugin.executingQuery=[{0}] Executing query: {1}