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/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 void setLogWriter(PrintWriter pw) throws SQLException {
 		throw new UnsupportedOperationException("setLogWriter");
 	}
 
-
-	//---------------------------------------------------------------------
-	// Implementation of JDBC 4.0's Wrapper interface
-	//---------------------------------------------------------------------
+	@Override
+	public Logger getParentLogger() {
+		return Logger.getLogger(Logger.GLOBAL_LOGGER_NAME);
+	}
 
 	@Override
 	@SuppressWarnings("unchecked")
@@ -95,14 +95,4 @@ public boolean isWrapperFor(Class<?> iface) throws SQLException {
 		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 @@
 
 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 Connection getConnection(String username, String password) throws SQLExce
 	}
 
 	@Override
-	public PrintWriter getLogWriter() throws SQLException {
-		return obtainTargetDataSource().getLogWriter();
+	public ConnectionBuilder createConnectionBuilder() throws SQLException {
+		return obtainTargetDataSource().createConnectionBuilder();
 	}
 
 	@Override
-	public void setLogWriter(PrintWriter out) throws SQLException {
-		obtainTargetDataSource().setLogWriter(out);
+	public ShardingKeyBuilder createShardingKeyBuilder() throws SQLException {
+		return obtainTargetDataSource().createShardingKeyBuilder();
 	}
 
 	@Override
@@ -124,10 +126,20 @@ public void setLoginTimeout(int seconds) throws SQLException {
 		obtainTargetDataSource().setLoginTimeout(seconds);
 	}
 
+	@Override
+	public PrintWriter getLogWriter() throws SQLException {
+		return obtainTargetDataSource().getLogWriter();
+	}
+
+	@Override
+	public void setLogWriter(PrintWriter out) throws SQLException {
+		obtainTargetDataSource().setLogWriter(out);
+	}
 
-	//---------------------------------------------------------------------
-	// Implementation of JDBC 4.0's Wrapper interface
-	//---------------------------------------------------------------------
+	@Override
+	public Logger getParentLogger() {
+		return Logger.getLogger(Logger.GLOBAL_LOGGER_NAME);
+	}
 
 	@Override
 	@SuppressWarnings("unchecked")
@@ -143,14 +155,4 @@ public boolean isWrapperFor(Class<?> iface) throws SQLException {
 		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,
- * which is used to get the sharding keys.</p>
+ * <p>This adapter needs to be configured with a {@link ShardingKeyProvider} callback which is
+ * 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}.
-	 *
-	 * @param dataSource the target DataSource to be wrapped.
+	 * Create a new instance of ShardingKeyDataSourceAdapter, wrapping the given {@link DataSource}.
+	 * @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}.
-	 *
-	 * @param dataSource the target DataSource to be wrapped.
-	 * @param shardingKeyProvider the ShardingKeyProvider used to get the shardingKeys.
+	 * Create a new instance of ShardingKeyDataSourceAdapter, wrapping the given {@link DataSource}.
+	 * @param dataSource the target DataSource to be wrapped
+	 * @param shardingKeyProvider the ShardingKeyProvider used to get the sharding keys
 	 */
 	public ShardingKeyDataSourceAdapter(DataSource dataSource, ShardingKeyProvider shardingKeyProvider) {
 		super(dataSource);
 		this.shardingkeyProvider = shardingKeyProvider;
 	}
 
+
 	/**
-	 * Sets the {@link ShardingKeyProvider} for this adapter.
-	 *
-	 * @param shardingKeyProvider the ShardingKeyProvider to set.
+	 * Set the {@link ShardingKeyProvider} for this adapter.
 	 */
 	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},
-	 * 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.
+	 * <p>The sharding key is obtained from the {@link ShardingKeyProvider}.
+	 * @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 Connection getConnection() throws SQLException {
 	}
 
 	/**
-	 * 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.
-	 *
-	 * @param username the database user on whose behalf the connection is being made.
-	 * @param password the user's password.
+	 * <p>The sharding key is obtained from the {@link ShardingKeyProvider}.
+	 * @param username the database user on whose behalf the connection is being made
+	 * @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 Connection getConnection(String username, String password) throws SQLExce
 	}
 
 	/**
-	 * 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}.
-	 *
-	 * @return a ConnectionBuilder object representing a builder for direct shard connections.
-	 * @throws SQLException if an error occurs while creating the ConnectionBuilder.
+	 * from the {@link ShardingKeyProvider}.
+	 * @return a ConnectionBuilder object representing a builder for direct shard connections
+	 * @throws SQLException if an error occurs while creating the ConnectionBuilder
 	 */
 	@Override
 	public ConnectionBuilder createConnectionBuilder() throws SQLException {
-		ConnectionBuilder connectionBuilder = obtainTargetDataSource().createConnectionBuilder();
-
-		ShardingKey shardingKey = null;
-		ShardingKey superShardingKey = null;
-
-		if (shardingkeyProvider != null) {
-			shardingKey = shardingkeyProvider.getShardingKey();
-			superShardingKey = shardingkeyProvider.getSuperShardingKey();
+		ConnectionBuilder connectionBuilder = super.createConnectionBuilder();
+		if (this.shardingkeyProvider == null) {
+			return connectionBuilder;
 		}
 
+		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;
+	}
+
+}