基於註解的Spring多資料來源配置和使用

前一段時間研究了一下spring多資料來源的配置和使用，為了後期從多個資料來源拉取資料定時進行資料分析和報表統計做準備。由於之前做過的專案都是單資料來源的，沒有遇到這種場景，所以也一直沒有去了解過如何配置多資料來源。
後來發現其實基於spring來配置和使用多資料來源還是比較簡單的，因為spring框架已經預留了這樣的介面可以方便資料來源的切換。
先看一下spring獲取資料來源的原始碼：

可以看到AbstractRoutingDataSource獲取資料來源之前會先呼叫determineCurrentLookupKey方法查詢當前的lookupKey，這個lookupKey就是資料來源標識。
因此透過重寫這個查詢資料來源標識的方法就可以讓spring切換到指定的資料來源了。
第一步：建立一個DynamicDataSource的類，繼承AbstractRoutingDataSource並重寫determineCurrentLookupKey方法，程式碼如下：

public class DynamicDataSource extends AbstractRoutingDataSource {

    @Override
    protected Object determineCurrentLookupKey() {
        // 從自定義的位置獲取資料來源標識
        return DynamicDataSourceHolder.getDataSource();
    }

}

第二步：建立DynamicDataSourceHolder用於持有當前執行緒中使用的資料來源標識，程式碼如下：

public class DynamicDataSourceHolder {
    /**
     * 注意：資料來源標識儲存線上程變數中，避免多執行緒運算元據源時互相干擾
     */
    private static final ThreadLocal<String> THREAD_DATA_SOURCE = new ThreadLocal<String>();

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

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

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

}

第三步：配置多個資料來源和第一步裡建立的DynamicDataSource的bean，簡化的配置如下：

<!--建立資料來源1，連線資料庫db1 -->
<bean id="dataSource1" class="org.apache.commons.dbcp.BasicDataSource" destroy-method="close">
    <property name="driverClassName" value="${db1.driver}" />
    <property name="url" value="${db1.url}" />
    <property name="username" value="${db1.username}" />
    <property name="password" value="${db1.password}" />
</bean>
<!--建立資料來源2，連線資料庫db2 -->
<bean id="dataSource2" class="org.apache.commons.dbcp.BasicDataSource" destroy-method="close">
    <property name="driverClassName" value="${db2.driver}" />
    <property name="url" value="${db2.url}" />
    <property name="username" value="${db2.username}" />
    <property name="password" value="${db2.password}" />
</bean>
<!--建立資料來源3，連線資料庫db3 -->
<bean id="dataSource3" class="org.apache.commons.dbcp.BasicDataSource" destroy-method="close">
    <property name="driverClassName" value="${db3.driver}" />
    <property name="url" value="${db3.url}" />
    <property name="username" value="${db3.username}" />
    <property name="password" value="${db3.password}" />
</bean>

<bean id="dynamicDataSource" class="com.test.context.datasource.DynamicDataSource">  
    <property name="targetDataSources">  
        <map key-type="java.lang.String">
            <!-- 指定lookupKey和與之對應的資料來源 -->
            <entry key="dataSource1" value-ref="dataSource1"></entry>  
            <entry key="dataSource2" value-ref="dataSource2"></entry>  
            <entry key="dataSource3 " value-ref="dataSource3"></entry>  
        </map>  
    </property>  
    <!-- 這裡可以指定預設的資料來源 -->
    <property name="defaultTargetDataSource" ref="dataSource1" />  
</bean>  

到這裡已經可以使用多資料來源了，在運算元據庫之前只要DynamicDataSourceHolder.setDataSource("dataSource2")即可切換到資料來源2並對資料庫db2進行操作了。

示例程式碼如下：

@Service
public class DataServiceImpl implements DataService {
    @Autowired
    private DataMapper dataMapper;

    @Override
    public List<Map<String, Object>> getList1() {
        // 沒有指定，則預設使用資料來源1
        return dataMapper.getList1();
    }

    @Override
    public List<Map<String, Object>> getList2() {
        // 指定切換到資料來源2
        DynamicDataSourceHolder.setDataSource("dataSource2");
        return dataMapper.getList2();
    }

    @Override
    public List<Map<String, Object>> getList3() {
        // 指定切換到資料來源3
        DynamicDataSourceHolder.setDataSource("dataSource3");
        return dataMapper.getList3();
    }
}

--------------------------------------------------------------------------------------華麗的分割線--------------------------------------------------------------------------------------------------

但是問題來了，如果每次切換資料來源時都呼叫DynamicDataSourceHolder.setDataSource("xxx")就顯得十分繁瑣了，而且程式碼量大了很容易會遺漏，後期維護起來也比較麻煩。能不能直接透過註解的方式指定需要訪問的資料來源呢，比如在dao層使用@DataSource("xxx")就指定訪問資料來源xxx？當然可以！前提是，再加一點額外的配置^_^。
首先，我們得定義一個名為DataSource的註解，程式碼如下：

@Target({ TYPE, METHOD })
@Retention(RUNTIME)
public @interface DataSource {
    String value();
}

然後，定義AOP切面以便攔截所有帶有註解@DataSource的方法，取出註解的值作為資料來源標識放到DynamicDataSourceHolder的執行緒變數中：

public class DataSourceAspect {

    /**
     * 攔截目標方法，獲取由@DataSource指定的資料來源標識，設定到執行緒儲存中以便切換資料來源
     * 
     * @param point
     * @throws Exception
     */
    public void intercept(JoinPoint point) throws Exception {
        Class<?> target = point.getTarget().getClass();
        MethodSignature signature = (MethodSignature) point.getSignature();
        // 預設使用目標型別的註解，如果沒有則使用其實現介面的註解
        for (Class<?> clazz : target.getInterfaces()) {
            resolveDataSource(clazz, signature.getMethod());
        }
        resolveDataSource(target, signature.getMethod());
    }

    /**
     * 提取目標物件方法註解和型別註解中的資料來源標識
     * 
     * @param clazz
     * @param method
     */
    private void resolveDataSource(Class<?> clazz, Method method) {
        try {
            Class<?>[] types = method.getParameterTypes();
            // 預設使用型別註解
            if (clazz.isAnnotationPresent(DataSource.class)) {
                DataSource source = clazz.getAnnotation(DataSource.class);
                DynamicDataSourceHolder.setDataSource(source.value());
            }
            // 方法註解可以覆蓋型別註解
            Method m = clazz.getMethod(method.getName(), types);
            if (m != null && m.isAnnotationPresent(DataSource.class)) {
                DataSource source = m.getAnnotation(DataSource.class);
                DynamicDataSourceHolder.setDataSource(source.value());
            }
        } catch (Exception e) {
            System.out.println(clazz + ":" + e.getMessage());
        }
    }

}

最後在spring配置檔案中配置攔截規則就可以了，比如攔截service層或者dao層的所有方法：

<bean id="dataSourceAspect" class="com.test.context.datasource.DataSourceAspect" />
    <aop:config>
        <aop:aspect ref="dataSourceAspect">
            <!-- 攔截所有service方法 -->
            <aop:pointcut id="dataSourcePointcut" expression="execution(* com.test.*.dao.*.*(..))"/>
            <aop:before pointcut-ref="dataSourcePointcut" method="intercept" />
        </aop:aspect>
    </aop:config>
</bean>

OK，這樣就可以直接在類或者方法上使用註解@DataSource來指定資料來源，不需要每次都手動設定了。

示例程式碼如下：

@Service
// 預設DataServiceImpl下的所有方法均訪問資料來源1
@DataSource("dataSource1")
public class DataServiceImpl implements DataService {
    @Autowired
    private DataMapper dataMapper;

    @Override
    public List<Map<String, Object>> getList1() {
        // 不指定，則預設使用資料來源1
        return dataMapper.getList1();
    }

    @Override
    // 覆蓋類上指定的，使用資料來源2
    @DataSource("dataSource2")
    public List<Map<String, Object>> getList2() {
        return dataMapper.getList2();
    }

    @Override
    // 覆蓋類上指定的，使用資料來源3
    @DataSource("dataSource3")
    public List<Map<String, Object>> getList3() {
        return dataMapper.getList3();
    }
}

提示：註解@DataSource既可以加在方法上，也可以加在介面或者介面的實現類上，優先順序別：方法>實現類>介面。也就是說如果介面、介面實現類以及方法上分別加了@DataSource註解來指定資料來源，則優先以方法上指定的為準。