diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/IamAuthConnectionPlugin.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/IamAuthConnectionPlugin.java index 80139510c..47b5978de 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/IamAuthConnectionPlugin.java +++ b/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 diff --git a/wrapper/src/main/resources/messages.properties b/wrapper/src/main/resources/messages.properties index 09228321b..3e0b8e475 100644 --- a/wrapper/src/main/resources/messages.properties +++ b/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}