Lazy handling of early setClientInfo/setNetworkTimeout calls in LazyConnectionDataSourceProxy

spring-jdbc/src/main/java/org/springframework/jdbc/datasource/LazyConnectionDataSourceProxy.java

@@ -22,7 +22,9 @@
 import java.lang.reflect.Proxy;
 import java.sql.Connection;
 import java.sql.SQLException;
+import java.util.LinkedHashMap;
 import java.util.Map;
+import java.util.concurrent.Executor;
 
 import javax.sql.DataSource;
 
@@ -35,12 +37,36 @@
 /**
  * Proxy for a target DataSource, fetching actual JDBC Connections lazily,
  * i.e. not until first creation of a Statement. Connection initialization
- * properties like auto-commit mode, transaction isolation and read-only mode
- * will be kept and applied to the actual JDBC Connection as soon as an actual
- * Connection is fetched (if ever). Consequently, commit and rollback calls will
- * be ignored if no Statements have been created. As of 6.1.2, there is also
- * special support for a {@link #setReadOnlyDataSource read-only DataSource} to use
- * during a read-only transaction, in addition to the regular target DataSource.
+ * properties like auto-commit mode, transaction isolation, read-only mode,
+ * catalog, schema, holdability, client info and network timeout will be kept
+ * and applied to the actual JDBC Connection as soon as an actual Connection
+ * is fetched (if ever). Consequently, commit and rollback calls will be ignored
+ * if no Statements have been created.
+ *
+ * <p>Once a property has been set, the corresponding getter method returns the
+ * set value until the actual Connection is fetched. If the property has not been
+ * set, invoking the getter triggers a fetch of the actual connection in order to
+ * obtain the default value.
+ *
+ * <p>Although client info is listed among the deferred properties above,
+ * the following methods are exceptions to the lazy acquisition behavior and
+ * force immediate acquisition of the underlying Connection.
+ * The {@link java.sql.Connection#getClientInfo()} and
+ * {@link java.sql.Connection#getClientInfo(java.lang.String)}
+ * methods are read operations whose values cannot be reliably cached due to
+ * driver defaults, remnants from pooled connections, or external session
+ * modifications.
+ *
+ * <p>The {@link java.sql.Connection#setClientInfo(java.util.Properties)}
+ * method also forces immediate acquisition. JDBC driver implementations are
+ * inconsistent: some treat it as an overwrite, while others treat it as an
+ * append/merge. To guarantee behavior identical to that of a non-lazy DataSource
+ * across all drivers, the proxy does not cache or replay it, thereby avoiding any
+ * risk of semantic mismatch.
+ *
+ * <p>As of 6.1.2, there is also special support for a
+ * {@link #setReadOnlyDataSource read-only DataSource} to use during a
+ * read-only transaction, in addition to the regular target DataSource.
  *
  * <p>This DataSource proxy allows to avoid fetching JDBC Connections from
  * a pool unless actually necessary. JDBC transaction control can happen
@@ -86,6 +112,7 @@
  *
  * @author Juergen Hoeller
  * @author Sam Brannen
+ * @author Chengang Guan
  * @since 1.1.4
  * @see DataSourceTransactionManager
  * @see #setTargetDataSource
@@ -311,6 +338,12 @@ private class LazyConnectionInvocationHandler implements InvocationHandler {
 
 		private @Nullable Boolean autoCommit;
 
+		private @Nullable Executor networkTimeoutExecutor;
+
+		private @Nullable Integer networkTimeout;
+
+		private @Nullable Map<String, String> clientInfo;
+
 		private boolean closed = false;
 
 		private @Nullable Connection target;
@@ -434,6 +467,28 @@ public LazyConnectionInvocationHandler(String username, String password) {
 						// Ignore: no warnings to expose yet.
 						return null;
 					}
+					case "setNetworkTimeout" -> {
+						this.networkTimeoutExecutor = (Executor) args[0];
+						this.networkTimeout = (Integer) args[1];
+						return null;
+					}
+					case "getNetworkTimeout" -> {
+						if (this.networkTimeout != null) {
+							return this.networkTimeout;
+						}
+						// Else fetch actual Connection and check there.
+					}
+					case "setClientInfo" -> {
+						if (args.length == 2) {
+							if (this.clientInfo == null) {
+								this.clientInfo = new LinkedHashMap<>();
+							}
+							this.clientInfo.put((String) args[0], (String) args[1]);
+							return null;
+						}
+						// Else fetch actual Connection and check there.
+						// setClientInfo(Properties) will fall-through
+					}
 					case "close" -> {
 						// Ignore: no target connection yet.
 						this.closed = true;
@@ -530,6 +585,14 @@ private Connection getTargetConnection(Method operation) throws Throwable {
 				if (this.autoCommit != null && this.autoCommit != defaultAutoCommit()) {
 					target.setAutoCommit(this.autoCommit);
 				}
+				if (this.networkTimeout != null) {
+					target.setNetworkTimeout(this.networkTimeoutExecutor, this.networkTimeout);
+				}
+				if (this.clientInfo != null) {
+					for (Map.Entry<String, String> entry: this.clientInfo.entrySet()) {
+						target.setClientInfo(entry.getKey(), entry.getValue());
+					}
+				}
 			}
 			catch (Throwable settingsEx) {
 				logger.debug("Failed to apply transaction settings to JDBC Connection", settingsEx);

spring-jdbc/src/test/java/org/springframework/jdbc/datasource/LazyConnectionDataSourceProxyTests.java

@@ -18,11 +18,17 @@
 
 import java.lang.reflect.Field;
 import java.sql.Connection;
+import java.sql.SQLException;
 import java.util.Arrays;
 import java.util.HashSet;
+import java.util.Properties;
 import java.util.Set;
+import java.util.concurrent.Executor;
 import java.util.stream.Stream;
 
+import javax.sql.DataSource;
+
+import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 
 import org.springframework.util.ReflectionUtils;
@@ -32,19 +38,37 @@
 import static java.sql.Connection.TRANSACTION_READ_UNCOMMITTED;
 import static java.sql.Connection.TRANSACTION_REPEATABLE_READ;
 import static java.sql.Connection.TRANSACTION_SERIALIZABLE;
+import static java.sql.ResultSet.CLOSE_CURSORS_AT_COMMIT;
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatIllegalArgumentException;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.doThrow;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
 
 /**
  * Tests for {@link LazyConnectionDataSourceProxy}.
  *
  * @author Sam Brannen
+ * @author Chengang Guan
  * @since 6.1
  */
 class LazyConnectionDataSourceProxyTests {
 
 	private final LazyConnectionDataSourceProxy proxy = new LazyConnectionDataSourceProxy();
 
+	private final LazyConnectionDataSourceProxy lazyProxy = new LazyConnectionDataSourceProxy();
+
+
+	@BeforeEach
+	void setup() {
+		lazyProxy.setDefaultAutoCommit(false);
+		lazyProxy.setDefaultTransactionIsolation(TRANSACTION_READ_UNCOMMITTED);
+	}
 
 	@Test
 	void setDefaultTransactionIsolationNameToUnsupportedValues() {
@@ -94,6 +118,157 @@ void setDefaultTransactionIsolation() {
 		assertThat(proxy.defaultTransactionIsolation()).isEqualTo(TRANSACTION_SERIALIZABLE);
 	}
 
+	@Test
+	void lazyHandlingCatalog() throws SQLException {
+		DataSource mockDataSource = mock();
+		Connection physicalConnection1 = mock();
+		Connection physicalConnection2 = mock();
+		when(mockDataSource.getConnection()).thenReturn(physicalConnection1).thenReturn(physicalConnection2);
+		lazyProxy.setTargetDataSource(mockDataSource);
+
+		Connection lazyConnection1 = lazyProxy.getConnection();
+		lazyConnection1.setCatalog("catalogName");
+		assertThat(lazyConnection1.getCatalog()).isEqualTo("catalogName");
+		verify(physicalConnection1, never()).setCatalog("catalogName");
+		verify(physicalConnection1, never()).getCatalog();
+		establishPhysicalConnection(lazyConnection1);
+		verify(physicalConnection1).setCatalog("catalogName");
+
+		Connection lazyConnection2 = lazyProxy.getConnection();
+		lazyConnection2.getCatalog(); // establish physical connection immediately
+		verify(physicalConnection2).getCatalog();
+	}
+
+	@Test
+	void lazyHandlingSchema() throws SQLException {
+		DataSource mockDataSource = mock();
+		Connection physicalConnection1 = mock();
+		Connection physicalConnection2 = mock();
+		when(mockDataSource.getConnection()).thenReturn(physicalConnection1).thenReturn(physicalConnection2);
+		lazyProxy.setTargetDataSource(mockDataSource);
+
+		Connection lazyConnection1 = lazyProxy.getConnection();
+		lazyConnection1.setSchema("schemaName");
+		assertThat(lazyConnection1.getSchema()).isEqualTo("schemaName");
+		verify(physicalConnection1, never()).setSchema("schemaName");
+		verify(physicalConnection1, never()).getSchema();
+		establishPhysicalConnection(lazyConnection1);
+		verify(physicalConnection1).setSchema("schemaName");
+
+		Connection lazyConnection2 = lazyProxy.getConnection();
+		lazyConnection2.getSchema(); // establish physical connection immediately
+		verify(physicalConnection2).getSchema();
+	}
+
+	@Test
+	void lazyHandlingHoldability() throws SQLException {
+		DataSource mockDataSource = mock();
+		Connection physicalConnection1 = mock();
+		Connection physicalConnection2 = mock();
+		when(mockDataSource.getConnection()).thenReturn(physicalConnection1).thenReturn(physicalConnection2);
+		lazyProxy.setTargetDataSource(mockDataSource);
+
+		Connection lazyConnection1 = lazyProxy.getConnection();
+		lazyConnection1.setHoldability(CLOSE_CURSORS_AT_COMMIT);
+		assertThat(lazyConnection1.getHoldability()).isEqualTo(CLOSE_CURSORS_AT_COMMIT);
+		verify(physicalConnection1, never()).setHoldability(CLOSE_CURSORS_AT_COMMIT);
+		verify(physicalConnection1, never()).getHoldability();
+		establishPhysicalConnection(lazyConnection1);
+		verify(physicalConnection1).setHoldability(CLOSE_CURSORS_AT_COMMIT);
+
+		Connection lazyConnection2 = lazyProxy.getConnection();
+		lazyConnection2.getHoldability(); // establish physical connection immediately
+		verify(physicalConnection2).getHoldability();
+	}
+
+	@Test
+	void lazyHandlingTransactionIsolation() throws SQLException {
+		DataSource mockDataSource = mock();
+		Connection physicalConnection = mock();
+		when(mockDataSource.getConnection()).thenReturn(physicalConnection);
+		lazyProxy.setTargetDataSource(mockDataSource);
+
+		Connection lazyConnection = lazyProxy.getConnection();
+		lazyConnection.setTransactionIsolation(TRANSACTION_READ_COMMITTED);
+		assertThat(lazyConnection.getTransactionIsolation()).isEqualTo(TRANSACTION_READ_COMMITTED);
+		verify(physicalConnection, never()).setTransactionIsolation(TRANSACTION_READ_COMMITTED);
+		verify(physicalConnection, never()).getTransactionIsolation();
+		establishPhysicalConnection(lazyConnection);
+		verify(physicalConnection).setTransactionIsolation(TRANSACTION_READ_COMMITTED);
+	}
+
+	@Test
+	void lazyHandlingAutoCommit() throws SQLException {
+		DataSource mockDataSource = mock();
+		Connection physicalConnection = mock();
+		when(mockDataSource.getConnection()).thenReturn(physicalConnection);
+		lazyProxy.setTargetDataSource(mockDataSource);
+
+		Connection lazyConnection = lazyProxy.getConnection();
+		lazyConnection.setAutoCommit(true);
+		assertThat(lazyConnection.getAutoCommit()).isTrue();
+		verify(physicalConnection, never()).setAutoCommit(true);
+		verify(physicalConnection, never()).getAutoCommit();
+		establishPhysicalConnection(lazyConnection);
+		verify(physicalConnection).setAutoCommit(true);
+	}
+
+	@Test
+	void lazyHandlingNetworkTimeoutExecutor() throws SQLException {
+		DataSource mockDataSource = mock();
+		Connection physicalConnection1 = mock();
+		Connection physicalConnection2 = mock();
+		Executor executor = mock();
+		when(mockDataSource.getConnection()).thenReturn(physicalConnection1).thenReturn(physicalConnection2);
+		doThrow(SQLException.class).when(physicalConnection2).setNetworkTimeout(eq(null), anyInt());
+		lazyProxy.setTargetDataSource(mockDataSource);
+
+		Connection lazyConnection1 = lazyProxy.getConnection();
+		lazyConnection1.setNetworkTimeout(executor, 1000);
+		assertThat(lazyConnection1.getNetworkTimeout()).isEqualTo(1000);
+		verify(physicalConnection1, never()).setNetworkTimeout(executor, 1000);
+		verify(physicalConnection1, never()).getNetworkTimeout();
+		establishPhysicalConnection(lazyConnection1);
+		verify(physicalConnection1).setNetworkTimeout(executor, 1000);
+
+		// null executor
+		Connection lazyConnection2 = lazyProxy.getConnection();
+		lazyConnection2.setNetworkTimeout(null, 1000);
+		assertThatThrownBy(() -> establishPhysicalConnection(lazyConnection2)).isInstanceOf(SQLException.class);
+	}
+
+	@Test
+	void lazyHandlingClientInfoForKV() throws SQLException {
+		DataSource mockDataSource = mock();
+		Connection physicalConnection = mock();
+		when(mockDataSource.getConnection()).thenReturn(physicalConnection);
+		lazyProxy.setTargetDataSource(mockDataSource);
+
+		Connection lazyConnection = lazyProxy.getConnection();
+		lazyConnection.setClientInfo("k1", "v1");
+		lazyConnection.setClientInfo("k2", "v2");
+		verify(physicalConnection, never()).setClientInfo("k1", "v1");
+		verify(physicalConnection, never()).setClientInfo("k2", "v2");
+		lazyConnection.getClientInfo("k1"); // establish physical connection immediately
+		verify(physicalConnection).setClientInfo("k1", "v1");
+		verify(physicalConnection).getClientInfo("k1");
+	}
+
+	@Test
+	void nonLazyHandlingClientInfoForProperties() throws SQLException {
+		DataSource mockDataSource = mock();
+		Connection physicalConnection = mock();
+		when(mockDataSource.getConnection()).thenReturn(physicalConnection);
+		lazyProxy.setTargetDataSource(mockDataSource);
+
+		Connection lazyConnection = lazyProxy.getConnection();
+		Properties properties = new Properties();
+		properties.setProperty("k1", "v1");
+		properties.setProperty("k2", "v2");
+		lazyConnection.setClientInfo(properties); // establish physical connection immediately
+		verify(physicalConnection).setClientInfo(properties);
+	}
+
 
 	private static Stream<String> streamIsolationConstants() {
 		return Arrays.stream(Connection.class.getFields())
@@ -102,4 +277,8 @@ private static Stream<String> streamIsolationConstants() {
 				.filter(name -> name.startsWith("TRANSACTION_"));
 	}
 
+	private static void establishPhysicalConnection(Connection lazyConnection) throws SQLException {
+		lazyConnection.prepareStatement("SELECT 1");
+	}
+
 }