Spring 註解動態資料來源設計實踐

Soochow發表於2021-05-25

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 註解即可。許多開源元件,會在現有的基礎上增加一個擴充套件功能,比如路由策略等等。

相關文章