spring boot动态多数据源demo

package com.springboot.multiplyds.datasource; import com.zaxxer.hikari.HikariConfig; import com.zaxxer.hikari.HikariDataSource; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import org.springframework.beans.factory.InitializingBean; import org.springframework.boot.jdbc.DataSourceBuilder; import javax.sql.DataSource; import java.sql.Connection; import java.sql.SQLException; import java.util.List; import java.util.Map; import java.util.Set; import java.util.concurrent.ConcurrentHashMap; import java.util.stream.Collectors; /** * 自定义多数据源类 */ public class MultiplyDataSource extends HikariDataSource implements InitializingBean { private static final Logger LOGGER = LoggerFactory.getLogger(MultiplyDataSource.class); /** * 所有数据源 */ private final Map<String, HikariDataSource> allDataSource = new ConcurrentHashMap<>(); /** * 其他需要切换控制的数据源 */ private final Map<String, String> dsInfoIdKeyMap = new ConcurrentHashMap<>(); public MultiplyDataSource() { super(); } public MultiplyDataSource(HikariConfig configuration) { super(configuration); } public Map<String, HikariDataSource> getAllDataSource() { return allDataSource; } /** * 获取数据源，如果多线程下设置了数据源就取设置的，如果没有取到就用默认的 * * @return 数据源 */ public DataSource getMultiplyDataSource() throws SQLException { String key = MultiplyDataSourceKeyHolder.getKey(); DataSource dataSource; if (key == null) { dataSource = this; } else { dataSource = this.allDataSource.get(dsInfoIdKeyMap.get(key)); } if (dataSource == null) { throw new SQLException("未找到和key[" + key + "]匹配的数据源。"); } return dataSource; } /** * 如果不存在，添加一个新数据源 * * @param key 数据连接信息 * @return 数据源 * @throws SQLException */ private HikariDataSource createIfNotExistsDataSource(DataSourceInfo key) throws SQLException { HikariDataSource dataSource = allDataSource.get(key.toString()); String keyToString = key.toString(); if (dataSource == null) { synchronized (allDataSource) { if (allDataSource.get(keyToString) == null) { HikariDataSource newDataSource = DataSourceBuilder.derivedFrom(this).type(HikariDataSource.class).driverClassName(key.getDriverClassName()).url(key.getUrl()).username(key.getUsername()).password(key.getPassword()).build(); newDataSource.setLoginTimeout(this.getLoginTimeout()); newDataSource.setValidationTimeout(this.getValidationTimeout()); newDataSource.setConnectionTimeout(this.getConnectionTimeout()); newDataSource.setIdleTimeout(this.getIdleTimeout()); newDataSource.setMaxLifetime(this.getMaxLifetime()); newDataSource.setMaximumPoolSize(this.getMaximumPoolSize()); newDataSource.setMinimumIdle(this.getMinimumIdle()); allDataSource.put(keyToString, newDataSource); dataSource = newDataSource; } else { dataSource = allDataSource.get(keyToString); } } } return dataSource; } @Override public Connection getConnection() throws SQLException { DataSource multiplyDataSource = this.getMultiplyDataSource(); if (this.equals(multiplyDataSource)) { return super.getConnection(); } return multiplyDataSource.getConnection(); } @Override public Connection getConnection(String username, String password) throws SQLException { DataSource multiplyDataSource = this.getMultiplyDataSource(); if (this.equals(multiplyDataSource)) { return super.getConnection(username, password); } return this.getMultiplyDataSource().getConnection(username, password); } /** * 更新数据源信息 * * @param newDatasourceList 数据源列表 */ public void updateMultiplyDataSources(List<DataSourceInfo> newDatasourceList) { if (newDatasourceList == null) { //清理之前的数据源 if (dsInfoIdKeyMap.size() > 0) { dsInfoIdKeyMap.values().forEach(row -> allDataSource.remove(row).close()); dsInfoIdKeyMap.clear(); } } else { newDatasourceList.forEach(row -> { try { allDataSource.putIfAbsent(row.toString(), createIfNotExistsDataSource(row)); } catch (SQLException e) { e.printStackTrace(); LOGGER.error("MultiplyDataSource init datasource exception:{}", e.getMessage()); } }); //通过toString()为key,删除为用到的,数据源信息可能会被修改，id未该，所以不通过id删除。 if (dsInfoIdKeyMap.size() > 0) { Set<String> keySet = newDatasourceList.stream().map(DataSourceInfo::toString).collect(Collectors.toSet()); dsInfoIdKeyMap.values().forEach(key -> { if (!keySet.contains(key)) { allDataSource.remove(key).close(); } }); } dsInfoIdKeyMap.clear(); newDatasourceList.forEach(row -> dsInfoIdKeyMap.put(row.getId(), row.toString())); } } @Override public void afterPropertiesSet() { allDataSource.put(DataSourceInfo.of(this.getDriverClassName(), this.getJdbcUrl(), this.getUsername(), this.getPassword()).toString(), this); LOGGER.info("MultiplyDataSource init,datasource.size:{}", getAllDataSource().size()); } }