From 9e930ce75a98886d9023ec0c701254da79e9e021 Mon Sep 17 00:00:00 2001 From: sergiyvamz Date: Wed, 25 Oct 2023 15:04:44 -0700 Subject: [PATCH 1/9] allow user to set up a lambda to initialize new connections --- .../jdbc/ConnectionProviderManager.java | 35 +++++++++++++++++++ .../jdbc/plugin/DefaultConnectionPlugin.java | 2 ++ 2 files changed, 37 insertions(+) diff --git a/wrapper/src/main/java/software/amazon/jdbc/ConnectionProviderManager.java b/wrapper/src/main/java/software/amazon/jdbc/ConnectionProviderManager.java index 384da5c7b..7e9d61c48 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/ConnectionProviderManager.java +++ b/wrapper/src/main/java/software/amazon/jdbc/ConnectionProviderManager.java @@ -16,10 +16,13 @@ package software.amazon.jdbc; +import java.sql.Connection; import java.sql.SQLException; import java.util.List; import java.util.Properties; import java.util.concurrent.locks.ReentrantReadWriteLock; +import org.checkerframework.checker.nullness.qual.NonNull; +import org.checkerframework.checker.nullness.qual.Nullable; import software.amazon.jdbc.cleanup.CanReleaseResources; public class ConnectionProviderManager { @@ -28,6 +31,8 @@ public class ConnectionProviderManager { private static ConnectionProvider connProvider = null; private final ConnectionProvider defaultProvider; + private static ConnectionInitFunc connectionInitFunc = null; + /** * {@link ConnectionProviderManager} constructor. * @@ -193,4 +198,34 @@ public static void releaseResources() { } } } + + public static void setConnectionInitFunc(final @NonNull ConnectionInitFunc func) { + connectionInitFunc = func; + } + + public static void resetConnectionInitFunc() { + connectionInitFunc = null; + } + + public void initConnection( + final @Nullable Connection connection, + final @NonNull String protocol, + final @NonNull HostSpec hostSpec, + final @NonNull Properties props) throws SQLException { + + final ConnectionInitFunc copy = connectionInitFunc; + if (copy == null) { + return; + } + + copy.initConnection(connection, protocol, hostSpec, props); + } + + public interface ConnectionInitFunc { + void initConnection( + final @Nullable Connection connection, + final @NonNull String protocol, + final @NonNull HostSpec hostSpec, + final @NonNull Properties props) throws SQLException; + } } diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/DefaultConnectionPlugin.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/DefaultConnectionPlugin.java index 8fcf0ec82..cf95d0936 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/DefaultConnectionPlugin.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/DefaultConnectionPlugin.java @@ -173,6 +173,8 @@ private Connection connectInternal( telemetryContext.closeContext(); } + this.connProviderManager.initConnection(conn, driverProtocol, hostSpec, props); + this.pluginService.setAvailability(hostSpec.asAliases(), HostAvailability.AVAILABLE); this.pluginService.updateDialect(conn); From 4d001b61bdc98f063271d285b7c3785a97cef8ea Mon Sep 17 00:00:00 2001 From: sergiyvamz Date: Thu, 26 Oct 2023 16:10:26 -0700 Subject: [PATCH 2/9] unit test and docs --- .../UsingTheJdbcDriver.md | 19 ++++++++++++++++++ .../jdbc/plugin/DefaultConnectionPlugin.java | 14 ++++++++++++- .../plugin/DefaultConnectionPluginTest.java | 20 ++++++++++++++++++- 3 files changed, 51 insertions(+), 2 deletions(-) diff --git a/docs/using-the-jdbc-driver/UsingTheJdbcDriver.md b/docs/using-the-jdbc-driver/UsingTheJdbcDriver.md index 9c51707c1..70fb2c279 100644 --- a/docs/using-the-jdbc-driver/UsingTheJdbcDriver.md +++ b/docs/using-the-jdbc-driver/UsingTheJdbcDriver.md @@ -100,6 +100,25 @@ DriverConfigurationProfiles.addOrReplaceProfile( CustomConnectionPluginFactory.class)); ``` +### Connection Initialization +In some cases it's necessary to configure a connection before a user application can use it. Some target drivers provides such functionality and allow to specify a configuration parameter with SQL statements that are executed when connection is established. However, not all drivers supports such functionality. Also, in some cases, additional conditions should be checked in order to identify what initialization is required for a particular connection. + +AWS JDBC Driver allows to specify a special function that can initialize a connection. It can be done with `ConnectionProviderManager.setConnectionInitFunc` method. `resetConnectionInitFunc` method is also available. + +The initialization function is called for all connections, including pre-opened connections provided by internal connection pool (see [Using Read Write Splitting Plugin Internal Connection Pooling](./using-plugins/UsingTheReadWriteSplittingPlugin.md#internal-connection-pooling)), and, thus, helping a user application to clean up connection session "contaminated" by previous use. + +> :warning: Executing CPU and network intensive code in the initialization function may cause significant performance degradation. + +```java +ConnectionProviderManager.setConnectionInitFunc((connection, protocol, hostSpec, props) -> { + // Set custom schema for connections to a test-database + if ("test-database".equals(props.getProperty("database"))) { + connection.setSchema("test-database-schema"); + } +}); +``` + + ### List of Available Plugins The AWS JDBC Driver has several built-in plugins that are available to use. Please visit the individual plugin page for more details. diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/DefaultConnectionPlugin.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/DefaultConnectionPlugin.java index cf95d0936..837adba9e 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/DefaultConnectionPlugin.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/DefaultConnectionPlugin.java @@ -67,6 +67,18 @@ public DefaultConnectionPlugin( final PluginService pluginService, final ConnectionProvider defaultConnProvider, final PluginManagerService pluginManagerService) { + this(pluginService, + defaultConnProvider, + pluginManagerService, + new ConnectionProviderManager(defaultConnProvider)); + } + + public DefaultConnectionPlugin( + final PluginService pluginService, + final ConnectionProvider defaultConnProvider, + final PluginManagerService pluginManagerService, + final ConnectionProviderManager connProviderManager) { + if (pluginService == null) { throw new IllegalArgumentException("pluginService"); } @@ -79,7 +91,7 @@ public DefaultConnectionPlugin( this.pluginService = pluginService; this.pluginManagerService = pluginManagerService; - this.connProviderManager = new ConnectionProviderManager(defaultConnProvider); + this.connProviderManager = connProviderManager; } @Override diff --git a/wrapper/src/test/java/software/amazon/jdbc/plugin/DefaultConnectionPluginTest.java b/wrapper/src/test/java/software/amazon/jdbc/plugin/DefaultConnectionPluginTest.java index 5157448b9..45c84d8ac 100644 --- a/wrapper/src/test/java/software/amazon/jdbc/plugin/DefaultConnectionPluginTest.java +++ b/wrapper/src/test/java/software/amazon/jdbc/plugin/DefaultConnectionPluginTest.java @@ -21,6 +21,7 @@ import static org.mockito.ArgumentMatchers.anyBoolean; import static org.mockito.ArgumentMatchers.anyString; import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.atLeastOnce; import static org.mockito.Mockito.never; import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; @@ -32,6 +33,7 @@ import java.util.Arrays; import java.util.Collections; import java.util.List; +import java.util.Properties; import java.util.stream.Stream; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; @@ -42,6 +44,8 @@ import org.mockito.Mock; import org.mockito.MockitoAnnotations; import software.amazon.jdbc.ConnectionProvider; +import software.amazon.jdbc.ConnectionProviderManager; +import software.amazon.jdbc.HostSpec; import software.amazon.jdbc.JdbcCallable; import software.amazon.jdbc.PluginManagerService; import software.amazon.jdbc.PluginService; @@ -59,12 +63,16 @@ class DefaultConnectionPluginTest { @Mock ConnectionProvider connectionProvider; @Mock PluginManagerService pluginManagerService; @Mock JdbcCallable mockSqlFunction; + @Mock JdbcCallable mockConnectFunction; @Mock Connection conn; @Mock Connection oldConn; @Mock private TelemetryFactory mockTelemetryFactory; @Mock TelemetryContext mockTelemetryContext; @Mock TelemetryCounter mockTelemetryCounter; @Mock TelemetryGauge mockTelemetryGauge; + @Mock ConnectionProviderManager mockConnectionProviderManager; + @Mock ConnectionProvider mockConnectionProvider; + @Mock HostSpec mockHostSpec; private AutoCloseable closeable; @@ -79,8 +87,11 @@ void setUp() { when(mockTelemetryFactory.createCounter(anyString())).thenReturn(mockTelemetryCounter); // noinspection unchecked when(mockTelemetryFactory.createGauge(anyString(), any(GaugeCallable.class))).thenReturn(mockTelemetryGauge); + when(mockConnectionProviderManager.getConnectionProvider(anyString(), any(), any())) + .thenReturn(mockConnectionProvider); - plugin = new DefaultConnectionPlugin(pluginService, connectionProvider, pluginManagerService); + plugin = new DefaultConnectionPlugin( + pluginService, connectionProvider, pluginManagerService, mockConnectionProviderManager); } @AfterEach @@ -109,6 +120,13 @@ void testExecute_closeOldConnection() throws SQLException { verify(pluginManagerService, never()).setInTransaction(anyBoolean()); } + @Test + void testConnect() throws SQLException { + plugin.connect("anyProtocol", mockHostSpec, new Properties(), true, mockConnectFunction); + verify(mockConnectionProvider, atLeastOnce()).connect(anyString(), any(), any(), any()); + verify(mockConnectionProviderManager, atLeastOnce()).initConnection(any(), anyString(), any(), any()); + } + private static Stream multiStatementQueries() { return Stream.of( Arguments.of("", new ArrayList()), From 93afc785b0b854bbe1d1b70e6375fe71c4931e16 Mon Sep 17 00:00:00 2001 From: sergiyvamz Date: Thu, 26 Oct 2023 17:11:55 -0700 Subject: [PATCH 3/9] code review --- docs/using-the-jdbc-driver/UsingTheJdbcDriver.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/using-the-jdbc-driver/UsingTheJdbcDriver.md b/docs/using-the-jdbc-driver/UsingTheJdbcDriver.md index 70fb2c279..1c38a6c0e 100644 --- a/docs/using-the-jdbc-driver/UsingTheJdbcDriver.md +++ b/docs/using-the-jdbc-driver/UsingTheJdbcDriver.md @@ -100,14 +100,14 @@ DriverConfigurationProfiles.addOrReplaceProfile( CustomConnectionPluginFactory.class)); ``` -### Connection Initialization +### Executing Custom Code When Initializing a Connection In some cases it's necessary to configure a connection before a user application can use it. Some target drivers provides such functionality and allow to specify a configuration parameter with SQL statements that are executed when connection is established. However, not all drivers supports such functionality. Also, in some cases, additional conditions should be checked in order to identify what initialization is required for a particular connection. AWS JDBC Driver allows to specify a special function that can initialize a connection. It can be done with `ConnectionProviderManager.setConnectionInitFunc` method. `resetConnectionInitFunc` method is also available. The initialization function is called for all connections, including pre-opened connections provided by internal connection pool (see [Using Read Write Splitting Plugin Internal Connection Pooling](./using-plugins/UsingTheReadWriteSplittingPlugin.md#internal-connection-pooling)), and, thus, helping a user application to clean up connection session "contaminated" by previous use. -> :warning: Executing CPU and network intensive code in the initialization function may cause significant performance degradation. +> :warning: Executing CPU and network intensive code in the initialization function may have a significant impact in the wrapper performance overall. ```java ConnectionProviderManager.setConnectionInitFunc((connection, protocol, hostSpec, props) -> { From 7c8469e96b840033bc9187f64ebde969cc85134b Mon Sep 17 00:00:00 2001 From: sergiyvamz <75754709+sergiyvamz@users.noreply.github.com> Date: Fri, 27 Oct 2023 10:57:56 -0700 Subject: [PATCH 4/9] Update docs/using-the-jdbc-driver/UsingTheJdbcDriver.md Co-authored-by: crystall-bitquill <97126568+crystall-bitquill@users.noreply.github.com> --- docs/using-the-jdbc-driver/UsingTheJdbcDriver.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/using-the-jdbc-driver/UsingTheJdbcDriver.md b/docs/using-the-jdbc-driver/UsingTheJdbcDriver.md index 1c38a6c0e..7d7e31728 100644 --- a/docs/using-the-jdbc-driver/UsingTheJdbcDriver.md +++ b/docs/using-the-jdbc-driver/UsingTheJdbcDriver.md @@ -105,7 +105,7 @@ In some cases it's necessary to configure a connection before a user application AWS JDBC Driver allows to specify a special function that can initialize a connection. It can be done with `ConnectionProviderManager.setConnectionInitFunc` method. `resetConnectionInitFunc` method is also available. -The initialization function is called for all connections, including pre-opened connections provided by internal connection pool (see [Using Read Write Splitting Plugin Internal Connection Pooling](./using-plugins/UsingTheReadWriteSplittingPlugin.md#internal-connection-pooling)), and, thus, helping a user application to clean up connection session "contaminated" by previous use. +The initialization function is called for all connections, including pre-opened connections provided by internal connection pools (see [Using Read Write Splitting Plugin and Internal Connection Pooling](./using-plugins/UsingTheReadWriteSplittingPlugin.md#internal-connection-pooling)). This helps user applications clean up connection sessions that have been altered by previous operation, as returning a connection to a pool will reset the state and retrieving it will call the initialization function again. > :warning: Executing CPU and network intensive code in the initialization function may have a significant impact in the wrapper performance overall. From cdbc418670907fb589774b6ea059070b5e750d23 Mon Sep 17 00:00:00 2001 From: sergiyvamz <75754709+sergiyvamz@users.noreply.github.com> Date: Fri, 27 Oct 2023 10:58:04 -0700 Subject: [PATCH 5/9] Update docs/using-the-jdbc-driver/UsingTheJdbcDriver.md Co-authored-by: crystall-bitquill <97126568+crystall-bitquill@users.noreply.github.com> --- docs/using-the-jdbc-driver/UsingTheJdbcDriver.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/using-the-jdbc-driver/UsingTheJdbcDriver.md b/docs/using-the-jdbc-driver/UsingTheJdbcDriver.md index 7d7e31728..37317c9d7 100644 --- a/docs/using-the-jdbc-driver/UsingTheJdbcDriver.md +++ b/docs/using-the-jdbc-driver/UsingTheJdbcDriver.md @@ -101,7 +101,7 @@ DriverConfigurationProfiles.addOrReplaceProfile( ``` ### Executing Custom Code When Initializing a Connection -In some cases it's necessary to configure a connection before a user application can use it. Some target drivers provides such functionality and allow to specify a configuration parameter with SQL statements that are executed when connection is established. However, not all drivers supports such functionality. Also, in some cases, additional conditions should be checked in order to identify what initialization is required for a particular connection. +Users may need to define a specific configuration for a connection that has just been opened by the driver before an application can use it. Not all target drivers provide this functionality, but some allow specifying a configuration parameter with SQL statements that are executed when a connection is established. Some cases may also require that additional conditions are checked in order to identify what initialization configuration is required for a particular connection. AWS JDBC Driver allows to specify a special function that can initialize a connection. It can be done with `ConnectionProviderManager.setConnectionInitFunc` method. `resetConnectionInitFunc` method is also available. From 4eb89bbb24198203fe36fb5a0b2b25c2862d5047 Mon Sep 17 00:00:00 2001 From: sergiyvamz <75754709+sergiyvamz@users.noreply.github.com> Date: Fri, 27 Oct 2023 12:20:56 -0700 Subject: [PATCH 6/9] Update docs/using-the-jdbc-driver/UsingTheJdbcDriver.md Co-authored-by: crystall-bitquill <97126568+crystall-bitquill@users.noreply.github.com> --- docs/using-the-jdbc-driver/UsingTheJdbcDriver.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/using-the-jdbc-driver/UsingTheJdbcDriver.md b/docs/using-the-jdbc-driver/UsingTheJdbcDriver.md index 37317c9d7..e1a91966f 100644 --- a/docs/using-the-jdbc-driver/UsingTheJdbcDriver.md +++ b/docs/using-the-jdbc-driver/UsingTheJdbcDriver.md @@ -103,7 +103,7 @@ DriverConfigurationProfiles.addOrReplaceProfile( ### Executing Custom Code When Initializing a Connection Users may need to define a specific configuration for a connection that has just been opened by the driver before an application can use it. Not all target drivers provide this functionality, but some allow specifying a configuration parameter with SQL statements that are executed when a connection is established. Some cases may also require that additional conditions are checked in order to identify what initialization configuration is required for a particular connection. -AWS JDBC Driver allows to specify a special function that can initialize a connection. It can be done with `ConnectionProviderManager.setConnectionInitFunc` method. `resetConnectionInitFunc` method is also available. +The AWS JDBC Driver allows specifying a special function that can initialize a connection. It can be done with `ConnectionProviderManager.setConnectionInitFunc` method. The `resetConnectionInitFunc` method is also available to remove the function. The initialization function is called for all connections, including pre-opened connections provided by internal connection pools (see [Using Read Write Splitting Plugin and Internal Connection Pooling](./using-plugins/UsingTheReadWriteSplittingPlugin.md#internal-connection-pooling)). This helps user applications clean up connection sessions that have been altered by previous operation, as returning a connection to a pool will reset the state and retrieving it will call the initialization function again. From 71b89da84cdb4e2d593456cbcc0c55d4d8098ed4 Mon Sep 17 00:00:00 2001 From: sergiyvamz <75754709+sergiyvamz@users.noreply.github.com> Date: Fri, 27 Oct 2023 12:32:51 -0700 Subject: [PATCH 7/9] Update docs/using-the-jdbc-driver/UsingTheJdbcDriver.md Co-authored-by: Karen <64801825+karenc-bq@users.noreply.github.com> --- docs/using-the-jdbc-driver/UsingTheJdbcDriver.md | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/docs/using-the-jdbc-driver/UsingTheJdbcDriver.md b/docs/using-the-jdbc-driver/UsingTheJdbcDriver.md index e1a91966f..640927451 100644 --- a/docs/using-the-jdbc-driver/UsingTheJdbcDriver.md +++ b/docs/using-the-jdbc-driver/UsingTheJdbcDriver.md @@ -101,7 +101,9 @@ DriverConfigurationProfiles.addOrReplaceProfile( ``` ### Executing Custom Code When Initializing a Connection -Users may need to define a specific configuration for a connection that has just been opened by the driver before an application can use it. Not all target drivers provide this functionality, but some allow specifying a configuration parameter with SQL statements that are executed when a connection is established. Some cases may also require that additional conditions are checked in order to identify what initialization configuration is required for a particular connection. +In some use cases you may need to define a specific configuration for a new driver connection before your application can use it. For instance: +- you might need to run some initial SQL queries when a connection is established, or; +- you might need to check for some additional conditions to determine the initialization configuration required for a particular connection. The AWS JDBC Driver allows specifying a special function that can initialize a connection. It can be done with `ConnectionProviderManager.setConnectionInitFunc` method. The `resetConnectionInitFunc` method is also available to remove the function. From 359d75d40645db4bef36ddacf4fb3305b8ada73c Mon Sep 17 00:00:00 2001 From: sergiyvamz <75754709+sergiyvamz@users.noreply.github.com> Date: Fri, 27 Oct 2023 12:33:04 -0700 Subject: [PATCH 8/9] Update docs/using-the-jdbc-driver/UsingTheJdbcDriver.md Co-authored-by: Karen <64801825+karenc-bq@users.noreply.github.com> --- docs/using-the-jdbc-driver/UsingTheJdbcDriver.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/using-the-jdbc-driver/UsingTheJdbcDriver.md b/docs/using-the-jdbc-driver/UsingTheJdbcDriver.md index 640927451..20260b1fa 100644 --- a/docs/using-the-jdbc-driver/UsingTheJdbcDriver.md +++ b/docs/using-the-jdbc-driver/UsingTheJdbcDriver.md @@ -107,7 +107,7 @@ In some use cases you may need to define a specific configuration for a new driv The AWS JDBC Driver allows specifying a special function that can initialize a connection. It can be done with `ConnectionProviderManager.setConnectionInitFunc` method. The `resetConnectionInitFunc` method is also available to remove the function. -The initialization function is called for all connections, including pre-opened connections provided by internal connection pools (see [Using Read Write Splitting Plugin and Internal Connection Pooling](./using-plugins/UsingTheReadWriteSplittingPlugin.md#internal-connection-pooling)). This helps user applications clean up connection sessions that have been altered by previous operation, as returning a connection to a pool will reset the state and retrieving it will call the initialization function again. +The initialization function is called for all connections, including connections opened by the internal connection pools (see [Using Read Write Splitting Plugin and Internal Connection Pooling](./using-plugins/UsingTheReadWriteSplittingPlugin.md#internal-connection-pooling)). This helps user applications clean up connection sessions that have been altered by previous operations, as returning a connection to a pool will reset the state and retrieving it will call the initialization function again. > :warning: Executing CPU and network intensive code in the initialization function may have a significant impact in the wrapper performance overall. From 07f9bc4408b7854bb3b2844913908e0721041bee Mon Sep 17 00:00:00 2001 From: sergiyvamz <75754709+sergiyvamz@users.noreply.github.com> Date: Fri, 27 Oct 2023 12:33:32 -0700 Subject: [PATCH 9/9] Update docs/using-the-jdbc-driver/UsingTheJdbcDriver.md Co-authored-by: Karen <64801825+karenc-bq@users.noreply.github.com> --- docs/using-the-jdbc-driver/UsingTheJdbcDriver.md | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/docs/using-the-jdbc-driver/UsingTheJdbcDriver.md b/docs/using-the-jdbc-driver/UsingTheJdbcDriver.md index 20260b1fa..797c387c5 100644 --- a/docs/using-the-jdbc-driver/UsingTheJdbcDriver.md +++ b/docs/using-the-jdbc-driver/UsingTheJdbcDriver.md @@ -109,7 +109,8 @@ The AWS JDBC Driver allows specifying a special function that can initialize a c The initialization function is called for all connections, including connections opened by the internal connection pools (see [Using Read Write Splitting Plugin and Internal Connection Pooling](./using-plugins/UsingTheReadWriteSplittingPlugin.md#internal-connection-pooling)). This helps user applications clean up connection sessions that have been altered by previous operations, as returning a connection to a pool will reset the state and retrieving it will call the initialization function again. -> :warning: Executing CPU and network intensive code in the initialization function may have a significant impact in the wrapper performance overall. +> [!WARNING]\ +> Executing CPU and network intensive code in the initialization function may significantly impact the AWS JDBC Driver's overall performance. ```java ConnectionProviderManager.setConnectionInitFunc((connection, protocol, hostSpec, props) -> {