From ba5eaddca82c9b66be5db07ee448873ed6d93725 Mon Sep 17 00:00:00 2001 From: aaron-congo Date: Thu, 7 Nov 2024 15:48:52 -0800 Subject: [PATCH 1/4] fix: use the ConfigurationProfile AwsCredentialsProviderHandler when provided --- .../src/main/java/software/amazon/jdbc/PluginService.java | 3 +++ .../main/java/software/amazon/jdbc/PluginServiceImpl.java | 6 ++++++ .../jdbc/plugin/AwsSecretsManagerConnectionPlugin.java | 6 ++++++ .../jdbc/plugin/customendpoint/CustomEndpointPlugin.java | 6 ++++++ .../amazon/jdbc/plugin/iam/IamAuthConnectionPlugin.java | 6 ++++++ 5 files changed, 27 insertions(+) diff --git a/wrapper/src/main/java/software/amazon/jdbc/PluginService.java b/wrapper/src/main/java/software/amazon/jdbc/PluginService.java index c6e08a980..ed5ff2a6a 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/PluginService.java +++ b/wrapper/src/main/java/software/amazon/jdbc/PluginService.java @@ -27,6 +27,7 @@ import software.amazon.jdbc.dialect.Dialect; import software.amazon.jdbc.exceptions.ExceptionHandler; import software.amazon.jdbc.hostavailability.HostAvailability; +import software.amazon.jdbc.profile.ConfigurationProfile; import software.amazon.jdbc.states.SessionStateService; import software.amazon.jdbc.targetdriverdialect.TargetDriverDialect; import software.amazon.jdbc.util.telemetry.TelemetryFactory; @@ -231,6 +232,8 @@ HostSpec getHostSpecByStrategy(List hosts, HostRole role, String strat String getDriverProtocol(); + @Nullable ConfigurationProfile getConfigurationProfile(); + Properties getProperties(); TelemetryFactory getTelemetryFactory(); diff --git a/wrapper/src/main/java/software/amazon/jdbc/PluginServiceImpl.java b/wrapper/src/main/java/software/amazon/jdbc/PluginServiceImpl.java index 887135f4e..7375da441 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/PluginServiceImpl.java +++ b/wrapper/src/main/java/software/amazon/jdbc/PluginServiceImpl.java @@ -244,6 +244,12 @@ public String getDriverProtocol() { return this.driverProtocol; } + @Override + @Nullable + public ConfigurationProfile getConfigurationProfile() { + return this.configurationProfile; + } + @Override public void setCurrentConnection( final @NonNull Connection connection, final @NonNull HostSpec hostSpec) throws SQLException { diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/AwsSecretsManagerConnectionPlugin.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/AwsSecretsManagerConnectionPlugin.java index 4032588b7..5d5d6d9ad 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/AwsSecretsManagerConnectionPlugin.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/AwsSecretsManagerConnectionPlugin.java @@ -49,6 +49,7 @@ import software.amazon.jdbc.PluginService; import software.amazon.jdbc.PropertyDefinition; import software.amazon.jdbc.authentication.AwsCredentialsManager; +import software.amazon.jdbc.profile.ConfigurationProfile; import software.amazon.jdbc.util.Messages; import software.amazon.jdbc.util.RegionUtils; import software.amazon.jdbc.util.StringUtils; @@ -175,7 +176,12 @@ public AwsSecretsManagerConnectionPlugin(final PluginService pluginService, fina this.secretKey = Pair.of(secretId, region); + ConfigurationProfile profile = this.pluginService.getConfigurationProfile(); + if (profile != null && profile.getAwsCredentialsProviderHandler() != null) { + AwsCredentialsManager.setCustomHandler(profile.getAwsCredentialsProviderHandler()); + } this.secretsManagerClientFunc = secretsManagerClientFunc; + this.getSecretValueRequestFunc = getSecretValueRequestFunc; this.fetchCredentialsCounter = this.pluginService.getTelemetryFactory() .createCounter(TELEMETRY_FETCH_CREDENTIALS_COUNTER); diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/customendpoint/CustomEndpointPlugin.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/customendpoint/CustomEndpointPlugin.java index a039c6aa6..18d2ca5f6 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/customendpoint/CustomEndpointPlugin.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/customendpoint/CustomEndpointPlugin.java @@ -34,6 +34,7 @@ import software.amazon.jdbc.PropertyDefinition; import software.amazon.jdbc.authentication.AwsCredentialsManager; import software.amazon.jdbc.plugin.AbstractConnectionPlugin; +import software.amazon.jdbc.profile.ConfigurationProfile; import software.amazon.jdbc.util.Messages; import software.amazon.jdbc.util.RdsUtils; import software.amazon.jdbc.util.RegionUtils; @@ -145,6 +146,11 @@ public CustomEndpointPlugin( final BiFunction rdsClientFunc) { this.pluginService = pluginService; this.props = props; + + ConfigurationProfile profile = this.pluginService.getConfigurationProfile(); + if (profile != null && profile.getAwsCredentialsProviderHandler() != null) { + AwsCredentialsManager.setCustomHandler(profile.getAwsCredentialsProviderHandler()); + } this.rdsClientFunc = rdsClientFunc; this.shouldWaitForInfo = WAIT_FOR_CUSTOM_ENDPOINT_INFO.getBoolean(this.props); diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/iam/IamAuthConnectionPlugin.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/iam/IamAuthConnectionPlugin.java index 29613da32..c2f8390a7 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/iam/IamAuthConnectionPlugin.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/iam/IamAuthConnectionPlugin.java @@ -36,6 +36,7 @@ import software.amazon.jdbc.authentication.AwsCredentialsManager; import software.amazon.jdbc.plugin.AbstractConnectionPlugin; import software.amazon.jdbc.plugin.TokenInfo; +import software.amazon.jdbc.profile.ConfigurationProfile; import software.amazon.jdbc.util.IamAuthUtils; import software.amazon.jdbc.util.Messages; import software.amazon.jdbc.util.RdsUtils; @@ -98,6 +99,11 @@ public IamAuthConnectionPlugin(final @NonNull PluginService pluginService) { this.telemetryFactory = pluginService.getTelemetryFactory(); this.cacheSizeGauge = telemetryFactory.createGauge("iam.tokenCache.size", () -> (long) tokenCache.size()); this.fetchTokenCounter = telemetryFactory.createCounter("iam.fetchToken.count"); + + ConfigurationProfile profile = this.pluginService.getConfigurationProfile(); + if (profile != null && profile.getAwsCredentialsProviderHandler() != null) { + AwsCredentialsManager.setCustomHandler(profile.getAwsCredentialsProviderHandler()); + } } @Override From 62c3ae61bb026e657d9a2e61a608446efea29bfc Mon Sep 17 00:00:00 2001 From: aaron-congo Date: Thu, 7 Nov 2024 15:52:04 -0800 Subject: [PATCH 2/4] fix: use the AwsCredentialsProviderHandler from the ConfigurationProfile when it is defined --- CHANGELOG.md | 3 +++ 1 file changed, 3 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 419fd1759..7b0f08bd7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -3,6 +3,9 @@ All notable changes to this project will be documented in this file. The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), and this project adheres to [Semantic Versioning](https://semver.org/#semantic-versioning-200). +### :bug: Fixed +- Use the AwsCredentialsProviderHandler from the ConfigurationProfile when it is defined ([PR #1183](https://github.com/aws/aws-advanced-jdbc-wrapper/pull/1183)). + ## [2.5.2] - 2024-11-4 ### :bug: Fixed - Limitless Connection Plugin to reduce extra connections made during new connection creation ([PR #1174](https://github.com/aws/aws-advanced-jdbc-wrapper/pull/1174)). From c7803ef3adbdec63678bdd637ad11a17a785a809 Mon Sep 17 00:00:00 2001 From: aaron-congo Date: Thu, 7 Nov 2024 17:04:03 -0800 Subject: [PATCH 3/4] Fix broken test compilation --- .../software/amazon/jdbc/plugin/efm/ConcurrencyTests.java | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/wrapper/src/test/java/software/amazon/jdbc/plugin/efm/ConcurrencyTests.java b/wrapper/src/test/java/software/amazon/jdbc/plugin/efm/ConcurrencyTests.java index 1685f326b..25f707762 100644 --- a/wrapper/src/test/java/software/amazon/jdbc/plugin/efm/ConcurrencyTests.java +++ b/wrapper/src/test/java/software/amazon/jdbc/plugin/efm/ConcurrencyTests.java @@ -71,6 +71,7 @@ import software.amazon.jdbc.dialect.UnknownDialect; import software.amazon.jdbc.hostavailability.HostAvailability; import software.amazon.jdbc.hostavailability.SimpleHostAvailabilityStrategy; +import software.amazon.jdbc.profile.ConfigurationProfile; import software.amazon.jdbc.states.SessionStateService; import software.amazon.jdbc.targetdriverdialect.PgTargetDriverDialect; import software.amazon.jdbc.targetdriverdialect.TargetDriverDialect; @@ -643,6 +644,11 @@ public String getDriverProtocol() { return null; } + @Override + public @Nullable ConfigurationProfile getConfigurationProfile() { + return null; + } + @Override public Properties getProperties() { return null; From 266b61b8e6f205be4bd39ce7576121a07c681515 Mon Sep 17 00:00:00 2001 From: aaron-congo Date: Thu, 14 Nov 2024 15:21:30 -0800 Subject: [PATCH 4/4] Simplify logic --- wrapper/src/main/java/software/amazon/jdbc/Driver.java | 4 ++++ .../src/main/java/software/amazon/jdbc/PluginService.java | 3 --- .../main/java/software/amazon/jdbc/PluginServiceImpl.java | 6 ------ .../software/amazon/jdbc/ds/AwsWrapperDataSource.java | 8 ++++---- .../jdbc/plugin/AwsSecretsManagerConnectionPlugin.java | 6 ------ .../jdbc/plugin/customendpoint/CustomEndpointPlugin.java | 6 ------ .../amazon/jdbc/plugin/iam/IamAuthConnectionPlugin.java | 6 ------ .../software/amazon/jdbc/plugin/efm/ConcurrencyTests.java | 6 ------ 8 files changed, 8 insertions(+), 37 deletions(-) diff --git a/wrapper/src/main/java/software/amazon/jdbc/Driver.java b/wrapper/src/main/java/software/amazon/jdbc/Driver.java index f679fb57d..5c01ba61f 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/Driver.java +++ b/wrapper/src/main/java/software/amazon/jdbc/Driver.java @@ -31,6 +31,7 @@ import java.util.logging.Logger; import org.checkerframework.checker.nullness.qual.NonNull; import org.checkerframework.checker.nullness.qual.Nullable; +import software.amazon.jdbc.authentication.AwsCredentialsManager; import software.amazon.jdbc.profile.ConfigurationProfile; import software.amazon.jdbc.profile.DriverConfigurationProfiles; import software.amazon.jdbc.states.ResetSessionStateOnCloseCallable; @@ -125,6 +126,9 @@ public Connection connect(final String url, final Properties info) throws SQLExc configurationProfile = DriverConfigurationProfiles.getProfileConfiguration(profileName); if (configurationProfile != null) { PropertyUtils.addProperties(props, configurationProfile.getProperties()); + if (configurationProfile.getAwsCredentialsProviderHandler() != null) { + AwsCredentialsManager.setCustomHandler(configurationProfile.getAwsCredentialsProviderHandler()); + } } else { throw new SQLException( Messages.get( diff --git a/wrapper/src/main/java/software/amazon/jdbc/PluginService.java b/wrapper/src/main/java/software/amazon/jdbc/PluginService.java index ed5ff2a6a..c6e08a980 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/PluginService.java +++ b/wrapper/src/main/java/software/amazon/jdbc/PluginService.java @@ -27,7 +27,6 @@ import software.amazon.jdbc.dialect.Dialect; import software.amazon.jdbc.exceptions.ExceptionHandler; import software.amazon.jdbc.hostavailability.HostAvailability; -import software.amazon.jdbc.profile.ConfigurationProfile; import software.amazon.jdbc.states.SessionStateService; import software.amazon.jdbc.targetdriverdialect.TargetDriverDialect; import software.amazon.jdbc.util.telemetry.TelemetryFactory; @@ -232,8 +231,6 @@ HostSpec getHostSpecByStrategy(List hosts, HostRole role, String strat String getDriverProtocol(); - @Nullable ConfigurationProfile getConfigurationProfile(); - Properties getProperties(); TelemetryFactory getTelemetryFactory(); diff --git a/wrapper/src/main/java/software/amazon/jdbc/PluginServiceImpl.java b/wrapper/src/main/java/software/amazon/jdbc/PluginServiceImpl.java index 7375da441..887135f4e 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/PluginServiceImpl.java +++ b/wrapper/src/main/java/software/amazon/jdbc/PluginServiceImpl.java @@ -244,12 +244,6 @@ public String getDriverProtocol() { return this.driverProtocol; } - @Override - @Nullable - public ConfigurationProfile getConfigurationProfile() { - return this.configurationProfile; - } - @Override public void setCurrentConnection( final @NonNull Connection connection, final @NonNull HostSpec hostSpec) throws SQLException { diff --git a/wrapper/src/main/java/software/amazon/jdbc/ds/AwsWrapperDataSource.java b/wrapper/src/main/java/software/amazon/jdbc/ds/AwsWrapperDataSource.java index e6a33a728..82739c115 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/ds/AwsWrapperDataSource.java +++ b/wrapper/src/main/java/software/amazon/jdbc/ds/AwsWrapperDataSource.java @@ -22,15 +22,11 @@ import java.io.PrintWriter; import java.io.Serializable; import java.sql.Connection; -import java.sql.DriverManager; import java.sql.SQLException; import java.sql.SQLFeatureNotSupportedException; -import java.util.Collections; -import java.util.List; import java.util.Map; import java.util.Properties; import java.util.logging.Logger; -import java.util.stream.Collectors; import javax.naming.NamingException; import javax.naming.Reference; import javax.naming.Referenceable; @@ -45,6 +41,7 @@ import software.amazon.jdbc.HostSpec; import software.amazon.jdbc.PropertyDefinition; import software.amazon.jdbc.TargetDriverHelper; +import software.amazon.jdbc.authentication.AwsCredentialsManager; import software.amazon.jdbc.profile.ConfigurationProfile; import software.amazon.jdbc.profile.DriverConfigurationProfiles; import software.amazon.jdbc.targetdriverdialect.TargetDriverDialect; @@ -111,6 +108,9 @@ public Connection getConnection(final String username, final String password) th configurationProfile = DriverConfigurationProfiles.getProfileConfiguration(profileName); if (configurationProfile != null) { PropertyUtils.addProperties(props, configurationProfile.getProperties()); + if (configurationProfile.getAwsCredentialsProviderHandler() != null) { + AwsCredentialsManager.setCustomHandler(configurationProfile.getAwsCredentialsProviderHandler()); + } } else { throw new SQLException( Messages.get( diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/AwsSecretsManagerConnectionPlugin.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/AwsSecretsManagerConnectionPlugin.java index 5d5d6d9ad..4032588b7 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/AwsSecretsManagerConnectionPlugin.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/AwsSecretsManagerConnectionPlugin.java @@ -49,7 +49,6 @@ import software.amazon.jdbc.PluginService; import software.amazon.jdbc.PropertyDefinition; import software.amazon.jdbc.authentication.AwsCredentialsManager; -import software.amazon.jdbc.profile.ConfigurationProfile; import software.amazon.jdbc.util.Messages; import software.amazon.jdbc.util.RegionUtils; import software.amazon.jdbc.util.StringUtils; @@ -176,12 +175,7 @@ public AwsSecretsManagerConnectionPlugin(final PluginService pluginService, fina this.secretKey = Pair.of(secretId, region); - ConfigurationProfile profile = this.pluginService.getConfigurationProfile(); - if (profile != null && profile.getAwsCredentialsProviderHandler() != null) { - AwsCredentialsManager.setCustomHandler(profile.getAwsCredentialsProviderHandler()); - } this.secretsManagerClientFunc = secretsManagerClientFunc; - this.getSecretValueRequestFunc = getSecretValueRequestFunc; this.fetchCredentialsCounter = this.pluginService.getTelemetryFactory() .createCounter(TELEMETRY_FETCH_CREDENTIALS_COUNTER); diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/customendpoint/CustomEndpointPlugin.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/customendpoint/CustomEndpointPlugin.java index 18d2ca5f6..a039c6aa6 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/customendpoint/CustomEndpointPlugin.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/customendpoint/CustomEndpointPlugin.java @@ -34,7 +34,6 @@ import software.amazon.jdbc.PropertyDefinition; import software.amazon.jdbc.authentication.AwsCredentialsManager; import software.amazon.jdbc.plugin.AbstractConnectionPlugin; -import software.amazon.jdbc.profile.ConfigurationProfile; import software.amazon.jdbc.util.Messages; import software.amazon.jdbc.util.RdsUtils; import software.amazon.jdbc.util.RegionUtils; @@ -146,11 +145,6 @@ public CustomEndpointPlugin( final BiFunction rdsClientFunc) { this.pluginService = pluginService; this.props = props; - - ConfigurationProfile profile = this.pluginService.getConfigurationProfile(); - if (profile != null && profile.getAwsCredentialsProviderHandler() != null) { - AwsCredentialsManager.setCustomHandler(profile.getAwsCredentialsProviderHandler()); - } this.rdsClientFunc = rdsClientFunc; this.shouldWaitForInfo = WAIT_FOR_CUSTOM_ENDPOINT_INFO.getBoolean(this.props); diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/iam/IamAuthConnectionPlugin.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/iam/IamAuthConnectionPlugin.java index c2f8390a7..29613da32 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/iam/IamAuthConnectionPlugin.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/iam/IamAuthConnectionPlugin.java @@ -36,7 +36,6 @@ import software.amazon.jdbc.authentication.AwsCredentialsManager; import software.amazon.jdbc.plugin.AbstractConnectionPlugin; import software.amazon.jdbc.plugin.TokenInfo; -import software.amazon.jdbc.profile.ConfigurationProfile; import software.amazon.jdbc.util.IamAuthUtils; import software.amazon.jdbc.util.Messages; import software.amazon.jdbc.util.RdsUtils; @@ -99,11 +98,6 @@ public IamAuthConnectionPlugin(final @NonNull PluginService pluginService) { this.telemetryFactory = pluginService.getTelemetryFactory(); this.cacheSizeGauge = telemetryFactory.createGauge("iam.tokenCache.size", () -> (long) tokenCache.size()); this.fetchTokenCounter = telemetryFactory.createCounter("iam.fetchToken.count"); - - ConfigurationProfile profile = this.pluginService.getConfigurationProfile(); - if (profile != null && profile.getAwsCredentialsProviderHandler() != null) { - AwsCredentialsManager.setCustomHandler(profile.getAwsCredentialsProviderHandler()); - } } @Override diff --git a/wrapper/src/test/java/software/amazon/jdbc/plugin/efm/ConcurrencyTests.java b/wrapper/src/test/java/software/amazon/jdbc/plugin/efm/ConcurrencyTests.java index 25f707762..1685f326b 100644 --- a/wrapper/src/test/java/software/amazon/jdbc/plugin/efm/ConcurrencyTests.java +++ b/wrapper/src/test/java/software/amazon/jdbc/plugin/efm/ConcurrencyTests.java @@ -71,7 +71,6 @@ import software.amazon.jdbc.dialect.UnknownDialect; import software.amazon.jdbc.hostavailability.HostAvailability; import software.amazon.jdbc.hostavailability.SimpleHostAvailabilityStrategy; -import software.amazon.jdbc.profile.ConfigurationProfile; import software.amazon.jdbc.states.SessionStateService; import software.amazon.jdbc.targetdriverdialect.PgTargetDriverDialect; import software.amazon.jdbc.targetdriverdialect.TargetDriverDialect; @@ -644,11 +643,6 @@ public String getDriverProtocol() { return null; } - @Override - public @Nullable ConfigurationProfile getConfigurationProfile() { - return null; - } - @Override public Properties getProperties() { return null;