Spring 註解動態數據源設計實踐

Spring 動態數據源

動態數據源是什麼？解決了什麼問題？

在實際的開發中，同一個項目中使用多個數據源是很常見的場景。比如，一個讀寫分離的項目存在主數據源與讀數據源。
所謂動態數據源，就是通過Spring的一些配置來自動控制某段數據操作邏輯是走哪一個數據源。舉個讀寫分離的例子，項目中引用了兩個數據源，master、slave。通過Spring配置或擴展能力來使得一個接口中調用了查詢方法會自動使用slave數據源。

一般實現這種效果可以通過：

1. 使用@MapperScan註解指定某個包下的所有方法走固定的數據源（這個比較死板些,會產生冗餘代碼，到也可以達到效果，可以作爲臨時方案使用）;
2. 使用註解+AOP+AbstractRoutingDataSource的形式來指定某個方法下的數據庫操作是走那個數據源。

關鍵核心類

這裏主要介紹通過註解+AOP+AbstractRoutingDataSource的聯動來實現動態數據源的方式。

一切的起點是AbstractRoutingDataSource這個類，此類實現了 DataSource 接口

public abstract class AbstractRoutingDataSource extends AbstractDataSource implements InitializingBean {

    // .... 省略 ... 
    
    @Nullable
	private Map<Object, Object> targetDataSources;
	
	@Nullable
	private Map<Object, DataSource> resolvedDataSources;


	public void setTargetDataSources(Map<Object, Object> targetDataSources) {
		this.targetDataSources = targetDataSources;
	}
	
	public void setDefaultTargetDataSource(Object defaultTargetDataSource) {
		this.defaultTargetDataSource = defaultTargetDataSource;
	}
	
	@Override
	public void afterPropertiesSet() {
	    
	    // 初始化 targetDataSources、resolvedDataSources
		if (this.targetDataSources == null) {
			throw new IllegalArgumentException("Property 'targetDataSources' is required");
		}
		this.resolvedDataSources = new HashMap<>(this.targetDataSources.size());
		this.targetDataSources.forEach((key, value) -> {
			Object lookupKey = resolveSpecifiedLookupKey(key);
			DataSource dataSource = resolveSpecifiedDataSource(value);
			this.resolvedDataSources.put(lookupKey, dataSource);
		});
		if (this.defaultTargetDataSource != null) {
			this.resolvedDefaultDataSource = resolveSpecifiedDataSource(this.defaultTargetDataSource);
		}
	}

    
	@Override
	public Connection getConnection() throws SQLException {
		return determineTargetDataSource().getConnection();
	}

	@Override
	public Connection getConnection(String username, String password) throws SQLException {
		return determineTargetDataSource().getConnection(username, password);
	}


	/**
	 * Retrieve the current target DataSource. Determines the
	 * {@link #determineCurrentLookupKey() current lookup key}, performs
	 * a lookup in the {@link #setTargetDataSources targetDataSources} map,
	 * falls back to the specified
	 * {@link #setDefaultTargetDataSource default target DataSource} if necessary.
	 * @see #determineCurrentLookupKey()
	 */
	protected DataSource determineTargetDataSource() {
		Assert.notNull(this.resolvedDataSources, "DataSource router not initialized");
	
	    // @1 start
		Object lookupKey = determineCurrentLookupKey();
		DataSource dataSource = this.resolvedDataSources.get(lookupKey);
		// @1 end
		
		if (dataSource == null && (this.lenientFallback || lookupKey == null)) {
			dataSource = this.resolvedDefaultDataSource;
		}
		if (dataSource == null) {
			throw new IllegalStateException("Cannot determine target DataSource for lookup key [" + lookupKey + "]");
		}
		return dataSource;
	}

	/**
	 * 返回一個key，這個key用來從 resolvedDataSources 數據源中獲取具體的數據源對象 見 @1
	 */
	@Nullable
	protected abstract Object determineCurrentLookupKey();

}

可以看到 AbstractRoutingDataSource 中有個可擴展抽象方法 determineCurrentLookupKey()，利用這個方法可以來實現動態數據源效果。

從0寫一個簡單動態數據源組件

從上一個part我們知道可以通過實現AbstractRoutingDataSource的 determineCurrentLookupKey() 方法動態設置一個key，然後
在配置類下通過setTargetDataSources()方法設置我們提前準備好的DataSource Map。

註解、常量定義、ThreadLocal 準備

/**
 * @author axin
 * @Summary 動態數據源註解定義
 */
public @interface MyDS {
    String value() default "default";
}

/**
 * @author axin
 * @Summary 動態數據源常量
 */
public interface DSConst {

    String 默認 = "default";

    String 主庫 = "master";

    String 從庫 = "slave";

    String 統計 = "stat";
}

/**
 * @author axin
 * @Summary 動態數據源 ThreadLocal 工具
 */
public class DynamicDataSourceHolder {

    //保存當前線程所指定的DataSource
    private static final ThreadLocal<String> THREAD_DATA_SOURCE = new ThreadLocal<>();

    public static String getDataSource() {
        return THREAD_DATA_SOURCE.get();
    }

    public static void setDataSource(String dataSource) {
        THREAD_DATA_SOURCE.set(dataSource);
    }

    public static void removeDataSource() {
        THREAD_DATA_SOURCE.remove();
    }
}

自定一個 AbstractRoutingDataSource 類

/**
 * @author axin
 * @Summary 動態數據源
 */
public class DynamicDataSource extends AbstractRoutingDataSource {

    /**
     * 從數據源中獲取目標數據源的key
     * @return
     */
    @Override
    protected Object determineCurrentLookupKey() {
        // 從ThreadLocal中獲取key
        String dataSourceKey = DynamicDataSourceHolder.getDataSource();
        if (StringUtils.isEmpty(dataSourceKey)) {
            return DSConst.默認;
        }
        return dataSourceKey;
    }
}

AOP 實現

/**
 * @author axin
 * @Summary 數據源切換AOP
 */
@Slf4j
@Aspect
@Service
public class DynamicDataSourceAOP {

    public DynamicDataSourceAOP() {
        log.info("/*---------------------------------------*/");
        log.info("/*----------                   ----------*/");
        log.info("/*---------- 動態數據源初始化... ----------*/");
        log.info("/*----------                   ----------*/");
        log.info("/*---------------------------------------*/");
    }

    /**
     * 切點
     */
    @Pointcut(value = "@annotation(xxx.xxx.MyDS)")
    private void method(){}

    /**
     * 方法執行前，切換到指定的數據源
     * @param point
     */
    @Before("method()")
    public void before(JoinPoint point) {
        MethodSignature methodSignature = (MethodSignature) point.getSignature();
        //獲取被代理的方法對象
        Method targetMethod = methodSignature.getMethod();
        //獲取被代理方法的註解信息
        CultureDS cultureDS = AnnotationUtils.findAnnotation(targetMethod, CultureDS.class);

        // 方法鏈條最外層的動態數據源註解優先級最高
        String key = DynamicDataSourceHolder.getDataSource();

        if (!StringUtils.isEmpty(key)) {
            log.warn("提醒：動態數據源註解調用鏈上出現覆蓋場景，請確認是否無問題");
            return;
        }

        if (cultureDS != null ) {
            //設置數據庫標誌
            DynamicDataSourceHolder.setDataSource(MyDS.value());
        }
    }

    /**
     * 釋放數據源
     */
    @AfterReturning("method()")
    public void doAfter() {
        DynamicDataSourceHolder.removeDataSource();
    }
}

DataSourceConfig 配置

通過以下代碼來將動態數據源配置到 SqlSession 中去

/**
 * 數據源的一些配置，主要是配置讀寫分離的sqlsession，這裏沒有使用mybatis annotation
 *
@Configuration
@EnableTransactionManagement
@EnableAspectJAutoProxy
class DataSourceConfig {
   
    /** 可讀寫的SQL Session */
    public static final String BEANNAME_SQLSESSION_COMMON = "sqlsessionCommon";
    /** 事務管理器的名稱，如果有多個事務管理器時，需要指定beanName */
    public static final String BEANNAME_TRANSACTION_MANAGER = "transactionManager";

    /** 主數據源，必須配置，spring啓動時會執行初始化數據操作（無論是否真的需要），選擇查找DataSource class類型的數據源 配置通用數據源，可讀寫，連接的是主庫 */
    @Bean
    @Primary
    @ConfigurationProperties(prefix = "datasource.common")
    public DataSource datasourceCommon() {
        // 數據源配置 可更換爲其他實現方式
        return DataSourceBuilder.create().build();
    }

    /**
     * 動態數據源
     * @returnr
     */
    @Bean
    public DynamicDataSource dynamicDataSource() {
        DynamicDataSource dynamicDataSource = new DynamicDataSource();
        LinkedHashMap<Object, Object> hashMap = Maps.newLinkedHashMap();
        hashMap.put(DSConst.默認, datasourceCommon());
        hashMap.put(DSConst.主庫, datasourceCommon());
        hashMap.put(DSConst.從庫, datasourceReadOnly());
        hashMap.put(DSConst.統計, datasourceStat());
        
        // 初始化數據源 Map
        dynamicDataSource.setTargetDataSources(hashMap);
        dynamicDataSource.setDefaultTargetDataSource(datasourceCommon());
        return dynamicDataSource;
    }

    /**
     * 配置事務管理器
     */
    @Bean(name = BEANNAME_TRANSACTION_MANAGER)
    public DataSourceTransactionManager createDataSourceTransactionManager() {
        DataSource dataSource = this.datasourceCommon();
        DataSourceTransactionManager manager = new DataSourceTransactionManager(dataSource);
        return manager;
    }

    /**
     * 配置讀寫sqlsession
     */
    @Primary
    @Bean(name = BEANNAME_SQLSESSION_COMMON)
    public SqlSession readWriteSqlSession() throws Exception {
        SqlSessionFactoryBean factory = new SqlSessionFactoryBean();

        // 設置動態數據源
        factory.setDataSource(this.dynamicDataSource());
        PathMatchingResourcePatternResolver resolver = new PathMatchingResourcePatternResolver();
        factory.setConfigLocation(resolver.getResource("mybatis/mybatis-config.xml"));
        factory.setMapperLocations(resolver.getResources("mybatis/mappers/**/*.xml"));
        return new SqlSessionTemplate(factory.getObject());
    }

}

總結

綜上，實現了一個簡單的Spring動態數據源功能，使用的時候，僅需要在目標方法上加上 @MyDS 註解即可。許多開源組件，會在現有的基礎上增加一個擴展功能，比如路由策略等等。