root
/
spring-framework
mirror of https://github.com/spring-projects/spring-framework.git
1
0
Fork 0
Code Issues Actions 6 Packages Projects Releases Wiki Activity

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
This commit is contained in:
Juergen Hoeller 2023-12-08 23:52:22 +01:00
parent e4e2224449
commit 69bc4e2828
7 changed files with 208 additions and 151 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -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;
}

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

@ -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");
}
//---------------------------------------------------------------------
// 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 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);
}
}

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

@ -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 {
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 class DelegatingDataSource implements DataSource, InitializingBean {
obtainTargetDataSource().setLoginTimeout(seconds);
}
@Override
public PrintWriter getLogWriter() throws SQLException {
return obtainTargetDataSource().getLogWriter();
}
//---------------------------------------------------------------------
// Implementation of JDBC 4.0's Wrapper interface
//---------------------------------------------------------------------
@Override
public void setLogWriter(PrintWriter out) throws SQLException {
obtainTargetDataSource().setLogWriter(out);
}
@Override
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);
}
}

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

@ -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 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.
*
* @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 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}.
*
* @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();
}
}

61
spring-jdbc/src/main/java/org/springframework/jdbc/datasource/ShardingKeyProvider.java Normal file
View File

@ -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;
}
}

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

@ -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

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

@ -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.*;
import static org.mockito.BDDMockito.*;
import static org.assertj.core.api.Assertions.assertThat;
import static org.mockito.BDDMockito.RETURNS_DEEP_STUBS;
import static org.mockito.BDDMockito.given;
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 {
ShardingKeyDataSourceAdapter dataSourceAdapter = new ShardingKeyDataSourceAdapter(
dataSource,
shardingKeyProvider);
void getConnectionWithKeyProvider() throws SQLException {
ShardingKeyDataSourceAdapter dataSourceAdapter =
new ShardingKeyDataSourceAdapter(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 {
ShardingKeyDataSourceAdapter dataSourceAdapter = new ShardingKeyDataSourceAdapter(
dataSource,
shardingKeyProvider);
void getConnectionWithCredentialsAndKeyProvider() throws SQLException {
ShardingKeyDataSourceAdapter dataSourceAdapter =
new ShardingKeyDataSourceAdapter(dataSource, 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())
.thenReturn(userWithKeyProviderConnection);
assertThat(dataSourceAdapter.getConnection(username, password)).isEqualTo(userWithKeyProviderConnection);
assertThat(dataSourceAdapter.getConnection(username, password)).isEqualTo(connection);
}
}
}