Delegation support for JDBC 4.3 ConnectionBuilder and ShardingKeyBuilder

Also moves ShardingKeyProvider to datasource package and declares getSuperShardingKey as default method.

Closes gh-31795
See gh-31506

spring-jdbc/src/main/java/org/springframework/jdbc/core/ShardingKeyProvider.java

@@ -1,38 +0,0 @@
-package org.springframework.jdbc.core;
-
-import java.sql.SQLException;
-import java.sql.ShardingKey;
-
-import org.springframework.lang.Nullable;
-
-/**
- * Interface defines methods for retrieving sharding keys, which are used to establish
- * direct shard connections (in the context of sharded databases). This is used as a
- * way of providing the sharding key in
- * {@link org.springframework.jdbc.datasource.ShardingKeyDataSourceAdapter}.
- *
- * @author Mohamed Lahyane (Anir)
- */
-
-
-public interface ShardingKeyProvider {
-	/**
-	 * Retrieves the sharding key. This method returns the sharding key relevant to the current context,
-	 * which will be used to obtain a direct shard connection.
-	 *
-	 * @return The sharding key, or null if it is not available or cannot be determined.
-	 * @throws SQLException If an error occurs while obtaining the sharding key.
-	 */
-	@Nullable
-	ShardingKey getShardingKey() throws SQLException;
-
-	/**
-	 * Retrieves the super sharding key. This method returns the super sharding key relevant to the
-	 * current context, which will be used to obtain a direct shard connection.
-	 *
-	 * @return The super sharding key, or null if it is not available or cannot be determined.
-	 * @throws SQLException If an error occurs while obtaining the super sharding key.
-	 */
-	@Nullable
-	ShardingKey getSuperShardingKey() throws SQLException;
-}

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2018 the original author or authors.
+ * Copyright 2002-2023 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -75,10 +75,10 @@ public abstract class AbstractDataSource implements DataSource {
 		throw new UnsupportedOperationException("setLogWriter");
 	}
 
-
+	@Override
-	//---------------------------------------------------------------------
+	public Logger getParentLogger() {
-	// Implementation of JDBC 4.0's Wrapper interface
+		return Logger.getLogger(Logger.GLOBAL_LOGGER_NAME);
-	//---------------------------------------------------------------------
+	}
 
 	@Override
 	@SuppressWarnings("unchecked")
@@ -95,14 +95,4 @@ public abstract class AbstractDataSource implements DataSource {
 		return iface.isInstance(this);
 	}
 
-
-	//---------------------------------------------------------------------
-	// Implementation of JDBC 4.1's getParentLogger method
-	//---------------------------------------------------------------------
-
-	@Override
-	public Logger getParentLogger() {
-		return Logger.getLogger(Logger.GLOBAL_LOGGER_NAME);
-	}
-
 }

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2017 the original author or authors.
+ * Copyright 2002-2023 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -18,7 +18,9 @@ package org.springframework.jdbc.datasource;
 
 import java.io.PrintWriter;
 import java.sql.Connection;
+import java.sql.ConnectionBuilder;
 import java.sql.SQLException;
+import java.sql.ShardingKeyBuilder;
 import java.util.logging.Logger;
 
 import javax.sql.DataSource;
@@ -105,13 +107,13 @@ public class DelegatingDataSource implements DataSource, InitializingBean {
 	}
 
 	@Override
-	public PrintWriter getLogWriter() throws SQLException {
+	public ConnectionBuilder createConnectionBuilder() throws SQLException {
-		return obtainTargetDataSource().getLogWriter();
+		return obtainTargetDataSource().createConnectionBuilder();
 	}
 
 	@Override
-	public void setLogWriter(PrintWriter out) throws SQLException {
+	public ShardingKeyBuilder createShardingKeyBuilder() throws SQLException {
-		obtainTargetDataSource().setLogWriter(out);
+		return obtainTargetDataSource().createShardingKeyBuilder();
 	}
 
 	@Override
@@ -124,10 +126,20 @@ public class DelegatingDataSource implements DataSource, InitializingBean {
 		obtainTargetDataSource().setLoginTimeout(seconds);
 	}
 
+	@Override
+	public PrintWriter getLogWriter() throws SQLException {
+		return obtainTargetDataSource().getLogWriter();
+	}
+
+	@Override
+	public void setLogWriter(PrintWriter out) throws SQLException {
+		obtainTargetDataSource().setLogWriter(out);
+	}
 
-	//---------------------------------------------------------------------
+	@Override
-	// Implementation of JDBC 4.0's Wrapper interface
+	public Logger getParentLogger() {
-	//---------------------------------------------------------------------
+		return Logger.getLogger(Logger.GLOBAL_LOGGER_NAME);
+	}
 
 	@Override
 	@SuppressWarnings("unchecked")
@@ -143,14 +155,4 @@ public class DelegatingDataSource implements DataSource, InitializingBean {
 		return (iface.isInstance(this) || obtainTargetDataSource().isWrapperFor(iface));
 	}
 
-
-	//---------------------------------------------------------------------
-	// Implementation of JDBC 4.1's getParentLogger method
-	//---------------------------------------------------------------------
-
-	@Override
-	public Logger getParentLogger() {
-		return Logger.getLogger(Logger.GLOBAL_LOGGER_NAME);
-	}
-
 }

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

@@ -1,77 +1,95 @@
+/*
+ * Copyright 2002-2023 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
 package org.springframework.jdbc.datasource;
 
 import java.sql.Connection;
 import java.sql.ConnectionBuilder;
 import java.sql.SQLException;
 import java.sql.ShardingKey;
-import java.sql.ShardingKeyBuilder;
 
 import javax.sql.DataSource;
 
-import org.springframework.core.NamedThreadLocal;
-import org.springframework.jdbc.core.ShardingKeyProvider;
 import org.springframework.lang.Nullable;
-import org.springframework.util.Assert;
 
 /**
  * An adapter for a target {@link DataSource}, designed to apply sharding keys, if specified,
- * to every standard {@code getConnection()} call, returning a direct connection to the shard
+ * to every standard {@code #getConnection} call, returning a direct connection to the shard
  * corresponding to the specified sharding key value. All other methods are simply delegated
  * to the corresponding methods of the target DataSource.
  *
- * <p>The target {@link DataSource} must implement the {@code createConnectionBuilder()} method;
+ * <p>The target {@link DataSource} must implement the {@link #createConnectionBuilder} method;
  * otherwise, a {@link java.sql.SQLFeatureNotSupportedException} will be thrown when attempting
- * to acquire shard connections.</p>
+ * to acquire shard connections.
  *
- * <p>This proxy datasource takes a {@link ShardingKeyProvider} object as an attribute,
+ * <p>This adapter needs to be configured with a {@link ShardingKeyProvider} callback which is
- * which is used to get the sharding keys.</p>
+ * used to get the current sharding keys for every {@code #getConnection} call, for example:
+ *
+ * <pre class="code">
+ * ShardingKeyDataSourceAdapter dataSourceAdapter = new ShardingKeyDataSourceAdapter(dataSource);
+ * dataSourceAdapter.setShardingKeyProvider(() -> dataSource.createShardingKeyBuilder()
+ *     .subkey(SecurityContextHolder.getContext().getAuthentication().getName(), JDBCType.VARCHAR).build());
+ * </pre>
  *
  * @author Mohamed Lahyane (Anir)
+ * @author Juergen Hoeller
+ * @since 6.1.2
  * @see #getConnection
  * @see #createConnectionBuilder()
  * @see UserCredentialsDataSourceAdapter
  */
 public class ShardingKeyDataSourceAdapter extends DelegatingDataSource {
+
 	@Nullable
 	private ShardingKeyProvider shardingkeyProvider;
 
+
 	/**
-	 * Creates a new instance of ShardingKeyDataSourceAdapter, wrapping the given {@link DataSource}.
+	 * Create a new instance of ShardingKeyDataSourceAdapter, wrapping the given {@link DataSource}.
-	 *
+	 * @param dataSource the target DataSource to be wrapped
-	 * @param dataSource the target DataSource to be wrapped.
 	 */
 	public ShardingKeyDataSourceAdapter(DataSource dataSource) {
 		super(dataSource);
 	}
 
 	/**
-	 * Creates a new instance of ShardingKeyDataSourceAdapter, wrapping the given {@link DataSource}.
+	 * Create a new instance of ShardingKeyDataSourceAdapter, wrapping the given {@link DataSource}.
-	 *
+	 * @param dataSource the target DataSource to be wrapped
-	 * @param dataSource the target DataSource to be wrapped.
+	 * @param shardingKeyProvider the ShardingKeyProvider used to get the sharding keys
-	 * @param shardingKeyProvider the ShardingKeyProvider used to get the shardingKeys.
 	 */
 	public ShardingKeyDataSourceAdapter(DataSource dataSource, ShardingKeyProvider shardingKeyProvider) {
 		super(dataSource);
 		this.shardingkeyProvider = shardingKeyProvider;
 	}
 
+
 	/**
-	 * Sets the {@link ShardingKeyProvider} for this adapter.
+	 * Set the {@link ShardingKeyProvider} for this adapter.
-	 *
-	 * @param shardingKeyProvider the ShardingKeyProvider to set.
 	 */
 	public void setShardingKeyProvider(ShardingKeyProvider shardingKeyProvider) {
 		this.shardingkeyProvider = shardingKeyProvider;
 	}
 
+
 	/**
-	 * Obtains a connection to the database shard using the provided sharding key
+	 * Obtain a connection to the database shard using the provided sharding key
 	 * and super sharding key (if available).
-	 * <p>the sharding key is obtained from the thread local storage, if is {@code null},
+	 * <p>The sharding key is obtained from the {@link ShardingKeyProvider}.
-	 * it is obtained from the {@link ShardingKeyProvider}.</p>
+	 * @return a Connection object representing a direct shard connection
-	 *
+	 * @throws SQLException if an error occurs while creating the connection
-	 * @return a Connection object representing a direct shard connection.
-	 * @throws SQLException if an error occurs while creating the connection.
 	 * @see #createConnectionBuilder()
 	 */
 	@Override
@@ -80,11 +98,11 @@ public class ShardingKeyDataSourceAdapter extends DelegatingDataSource {
 	}
 
 	/**
-	 * Obtains a connection to the database shard using the provided username and password,
+	 * Obtain a connection to the database shard using the provided username and password,
 	 * considering the sharding keys (if available) and the given credentials.
-	 *
+	 * <p>The sharding key is obtained from the {@link ShardingKeyProvider}.
-	 * @param username the database user on whose behalf the connection is being made.
+	 * @param username the database user on whose behalf the connection is being made
-	 * @param password the user's password.
+	 * @param password the user's password
 	 * @return a Connection object representing a direct shard connection.
 	 * @throws SQLException if an error occurs while creating the connection.
 	 */
@@ -94,37 +112,22 @@ public class ShardingKeyDataSourceAdapter extends DelegatingDataSource {
 	}
 
 	/**
-	 * Creates a new instance of {@link ConnectionBuilder} using the target DataSource's
+	 * Create a new instance of {@link ConnectionBuilder} using the target DataSource's
 	 * {@code createConnectionBuilder()} method, and sets the appropriate sharding keys
-	 * from the thread-local storage or the {@link ShardingKeyProvider}.
+	 * from the {@link ShardingKeyProvider}.
-	 *
+	 * @return a ConnectionBuilder object representing a builder for direct shard connections
-	 * @return a ConnectionBuilder object representing a builder for direct shard connections.
+	 * @throws SQLException if an error occurs while creating the ConnectionBuilder
-	 * @throws SQLException if an error occurs while creating the ConnectionBuilder.
 	 */
 	@Override
 	public ConnectionBuilder createConnectionBuilder() throws SQLException {
-		ConnectionBuilder connectionBuilder = obtainTargetDataSource().createConnectionBuilder();
+		ConnectionBuilder connectionBuilder = super.createConnectionBuilder();
-
+		if (this.shardingkeyProvider == null) {
-		ShardingKey shardingKey = null;
+			return connectionBuilder;
-		ShardingKey superShardingKey = null;
-
-		if (shardingkeyProvider != null) {
-			shardingKey = shardingkeyProvider.getShardingKey();
-			superShardingKey = shardingkeyProvider.getSuperShardingKey();
 		}
 
+		ShardingKey shardingKey = this.shardingkeyProvider.getShardingKey();
+		ShardingKey superShardingKey = this.shardingkeyProvider.getSuperShardingKey();
 		return connectionBuilder.shardingKey(shardingKey).superShardingKey(superShardingKey);
 	}
 
-	/**
-	 * Creates a new instance of {@link ShardingKeyBuilder} using the target DataSource's
-	 * {@code createShardingKeyBuilder()} method.
-	 *
-	 * @return a ShardingKeyBuilder object representing a builder for sharding keys.
-	 * @throws SQLException if an error occurs while creating the ShardingKeyBuilder.
-	 */
-	@Override
-	public ShardingKeyBuilder createShardingKeyBuilder() throws SQLException {
-		return obtainTargetDataSource().createShardingKeyBuilder();
-	}
 }

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

@@ -0,0 +1,61 @@
+/*
+ * Copyright 2002-2023 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.jdbc.datasource;
+
+import java.sql.SQLException;
+import java.sql.ShardingKey;
+
+import org.springframework.lang.Nullable;
+
+/**
+ * Strategy interface for determining sharding keys which are used to establish direct
+ * shard connections in the context of sharded databases. This is used as a callback
+ * for providing the current sharding key (plus optionally a super sharding key) in
+ * {@link org.springframework.jdbc.datasource.ShardingKeyDataSourceAdapter}.
+ *
+ * <p>Can be used as a functional interface (e.g. with a lambda expression) for a simple
+ * sharding key, or as a two-method interface when including a super sharding key as well.
+ *
+ * @author Mohamed Lahyane (Anir)
+ * @author Juergen Hoeller
+ * @since 6.1.2
+ * @see ShardingKeyDataSourceAdapter#setShardingKeyProvider
+ */
+public interface ShardingKeyProvider {
+
+	/**
+	 * Determine the sharding key. This method returns the sharding key relevant to the current
+	 * context which will be used to obtain a direct shard connection.
+	 * @return the sharding key, or {@code null} if it is not available or cannot be determined
+	 * @throws SQLException if an error occurs while obtaining the sharding key
+	 */
+	@Nullable
+	ShardingKey getShardingKey() throws SQLException;
+
+	/**
+	 * Determine the super sharding key, if any. This method returns the super sharding key
+	 * relevant to the current context which will be used to obtain a direct shard connection.
+	 * @return the super sharding key, or {@code null} if it is not available or cannot be
+	 * determined (the default)
+	 * @throws SQLException if an error occurs while obtaining the super sharding key
+	 */
+	@Nullable
+	default ShardingKey getSuperShardingKey() throws SQLException {
+		return null;
+	}
+
+}

spring-jdbc/src/main/java/org/springframework/jdbc/datasource/lookup/AbstractRoutingDataSource.java

@@ -17,7 +17,9 @@
 package org.springframework.jdbc.datasource.lookup;
 
 import java.sql.Connection;
+import java.sql.ConnectionBuilder;
 import java.sql.SQLException;
+import java.sql.ShardingKeyBuilder;
 import java.util.Collections;
 import java.util.Map;
 
@@ -216,6 +218,16 @@ public abstract class AbstractRoutingDataSource extends AbstractDataSource imple
 		return determineTargetDataSource().getConnection(username, password);
 	}
 
+	@Override
+	public ConnectionBuilder createConnectionBuilder() throws SQLException {
+		return determineTargetDataSource().createConnectionBuilder();
+	}
+
+	@Override
+	public ShardingKeyBuilder createShardingKeyBuilder() throws SQLException {
+		return determineTargetDataSource().createShardingKeyBuilder();
+	}
+
 	@Override
 	@SuppressWarnings("unchecked")
 	public <T> T unwrap(Class<T> iface) throws SQLException {
@@ -230,6 +242,7 @@ public abstract class AbstractRoutingDataSource extends AbstractDataSource imple
 		return (iface.isInstance(this) || determineTargetDataSource().isWrapperFor(iface));
 	}
 
+
 	/**
 	 * Retrieve the current target DataSource. Determines the
 	 * {@link #determineCurrentLookupKey() current lookup key}, performs

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

@@ -1,3 +1,19 @@
+/*
+ * Copyright 2002-2023 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
 package org.springframework.jdbc.datasource;
 
 import java.sql.Connection;
@@ -10,41 +26,58 @@ import javax.sql.DataSource;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 
-import org.springframework.jdbc.core.ShardingKeyProvider;
+import static org.assertj.core.api.Assertions.assertThat;
-
+import static org.mockito.BDDMockito.RETURNS_DEEP_STUBS;
-import static org.assertj.core.api.Assertions.*;
+import static org.mockito.BDDMockito.given;
-import static org.mockito.BDDMockito.*;
+import static org.mockito.BDDMockito.mock;
-
+import static org.mockito.BDDMockito.when;
+
+/**
+ * Tests for {@link ShardingKeyDataSourceAdapter}.
+ *
+ * @author Mohamed Lahyane (Anir)
+ * @author Juergen Hoeller
+ * @since 6.1.2
+ */
 public class ShardingKeyDataSourceAdapterTests {
+
 	private final Connection connection = mock();
+
 	private final Connection shardConnection = mock();
+
 	private final DataSource dataSource = mock();
+
 	private final ConnectionBuilder connectionBuilder = mock(ConnectionBuilder.class, RETURNS_DEEP_STUBS);
+
 	private final ConnectionBuilder shardConnectionBuilder = mock(ConnectionBuilder.class, RETURNS_DEEP_STUBS);
+
 	private final ShardingKey shardingKey = mock();
+
 	private final ShardingKey superShardingKey = mock();
+
 	private final ShardingKeyProvider shardingKeyProvider = new ShardingKeyProvider() {
 		@Override
 		public ShardingKey getShardingKey() throws SQLException {
 			return shardingKey;
 		}
-
 		@Override
 		public ShardingKey getSuperShardingKey() throws SQLException {
 			return superShardingKey;
 		}
 	};
 
+
 	@BeforeEach
-	public void setUp() throws SQLException {
+	void setup() throws SQLException {
 		given(dataSource.createConnectionBuilder()).willReturn(connectionBuilder);
 		when(connectionBuilder.shardingKey(null).superShardingKey(null)).thenReturn(connectionBuilder);
 		when(connectionBuilder.shardingKey(shardingKey).superShardingKey(superShardingKey))
 				.thenReturn(shardConnectionBuilder);
 	}
 
+
 	@Test
-	public void testGetConnectionNoKeyProvider() throws SQLException {
+	void getConnectionNoKeyProvider() throws SQLException {
 		ShardingKeyDataSourceAdapter dataSourceAdapter = new ShardingKeyDataSourceAdapter(dataSource);
 
 		when(connectionBuilder.build()).thenReturn(connection);
@@ -53,11 +86,9 @@ public class ShardingKeyDataSourceAdapterTests {
 	}
 
 	@Test
-	public void testGetConnectionWithKeyProvider() throws SQLException {
+	void getConnectionWithKeyProvider() throws SQLException {
-
+		ShardingKeyDataSourceAdapter dataSourceAdapter =
-		ShardingKeyDataSourceAdapter dataSourceAdapter = new ShardingKeyDataSourceAdapter(
+				new ShardingKeyDataSourceAdapter(dataSource, shardingKeyProvider);
-				dataSource,
-				shardingKeyProvider);
 
 		when(shardConnectionBuilder.build()).thenReturn(shardConnection);
 
@@ -65,33 +96,28 @@ public class ShardingKeyDataSourceAdapterTests {
 	}
 
 	@Test
-	public void testGetConnectionWithCredentialsNoKeyProvider() throws SQLException {
+	void getConnectionWithCredentialsNoKeyProvider() throws SQLException {
 		ShardingKeyDataSourceAdapter dataSourceAdapter = new ShardingKeyDataSourceAdapter(dataSource);
 
 		String username = "Anir";
 		String password = "spring";
 
-		Connection userConnection = mock();
+		when(connectionBuilder.user(username).password(password).build()).thenReturn(connection);
-
-		when(connectionBuilder.user(username).password(password).build()).thenReturn(userConnection);
 
-		assertThat(dataSourceAdapter.getConnection(username, password)).isEqualTo(userConnection);
+		assertThat(dataSourceAdapter.getConnection(username, password)).isEqualTo(connection);
 	}
 
 	@Test
-	public void testGetConnectionWithCredentialsAndKeyProvider() throws SQLException {
+	void getConnectionWithCredentialsAndKeyProvider() throws SQLException {
-		ShardingKeyDataSourceAdapter dataSourceAdapter = new ShardingKeyDataSourceAdapter(
+		ShardingKeyDataSourceAdapter dataSourceAdapter =
-				dataSource,
+				new ShardingKeyDataSourceAdapter(dataSource, shardingKeyProvider);
-				shardingKeyProvider);
 
 		String username = "mbekraou";
 		String password = "jdbc";
 
-		Connection userWithKeyProviderConnection = mock();
+		when(shardConnectionBuilder.user(username).password(password).build()).thenReturn(connection);
 
-		when(shardConnectionBuilder.user(username).password(password).build())
+		assertThat(dataSourceAdapter.getConnection(username, password)).isEqualTo(connection);
-				.thenReturn(userWithKeyProviderConnection);
-
-		assertThat(dataSourceAdapter.getConnection(username, password)).isEqualTo(userWithKeyProviderConnection);
 	}
-}
+
+}