接続先DBをSQLの内容等で振り分け

概要

・負荷分散、データ分割などの理由で接続先のDBを振り分けたい
・org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource を利用
・MyBatis のSQL実行をインターセプト

環境

・SpringFramework 4
・mybatis 3.2
・mybatis-spring 1.2
・mysql (jdbc)

構成

・Source/Replica (Master/Slave)
・デフォルトは Source の方に接続、特別な場合に Replica の方へ接続するなどコードで色々設定が可能
・インターセプターで接続先を設定する

コード

DynamicDataSource

import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;

public class DynamicDataSource extends AbstractRoutingDataSource {

	@Override
	protected Object determineCurrentLookupKey() {
		return DynamicDataSourceHolder.getDbType();
	}

}

DynamicDataSourceHolder

Source/Replica でなく Master/Slave になっていますが…

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
 
public class DynamicDataSourceHolder {
 
	public static final String DB_MASTER = "master";
	public static final String DB_SLAVE = "slave";
	
	private static Logger logger = LoggerFactory.getLogger(DynamicDataSourceHolder.class);
	private static ThreadLocal<String> contextHolder = new ThreadLocal<>();
	
	public static Object getDbType() {
		String dbType = contextHolder.get();
		if (dbType == null) {
			return DB_MASTER;
		}
		return dbType;
	}
	
	public static void setDbType(String dbType) {
		contextHolder.set(dbType);
	}
	
	public static void clearDbType() {
		contextHolder.remove();
	}
}

DynamicDataSourceInterceptor

・REGEX_WRITE の場合は Source(Master) の方へ接続
・更新系を分けていますが、SELECTから始まるSQLを replica(slave) へする等色々条件は変更が可能

import java.util.Properties;

import org.apache.ibatis.cache.CacheKey;
import org.apache.ibatis.executor.Executor;
import org.apache.ibatis.executor.keygen.SelectKeyGenerator;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.mapping.SqlCommandType;
import org.apache.ibatis.plugin.Interceptor;
import org.apache.ibatis.plugin.Intercepts;
import org.apache.ibatis.plugin.Invocation;
import org.apache.ibatis.plugin.Plugin;
import org.apache.ibatis.plugin.Signature;
import org.apache.ibatis.session.ResultHandler;
import org.apache.ibatis.session.RowBounds;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.transaction.support.TransactionSynchronizationManager;

@Intercepts({
		@Signature(type = Executor.class, method = "query", args = { MappedStatement.class, Object.class,
				RowBounds.class, ResultHandler.class }),
		@Signature(type = Executor.class, method = "query", args = { MappedStatement.class, Object.class,
				RowBounds.class, ResultHandler.class, CacheKey.class, BoundSql.class }),
		@Signature(type = Executor.class, method = "update", args = { MappedStatement.class, Object.class }), })
public class DynamicDataSourceInterceptor implements Interceptor {
	
	private static Logger logger = LoggerFactory.getLogger(DynamicDataSourceInterceptor.class);
	private static final String REGEX_WRITE = ".*update\\u0020.*|.*insert\\u0020.*|.*delete\\u0020.*|.*load\\u0020.*";

	@Override
	public Object intercept(Invocation invocation) throws Throwable {
		
		boolean synchronizationActive = TransactionSynchronizationManager.isSynchronizationActive();
		Object[] args = invocation.getArgs();
		MappedStatement ms = (MappedStatement) args[0];
		String lookupKey = "";
		
		if (!synchronizationActive) {
			if (ms.getSqlCommandType().equals(SqlCommandType.SELECT)) {
				if (ms.getId().contains(SelectKeyGenerator.SELECT_KEY_SUFFIX)) {
					lookupKey = DynamicDataSourceHolder.DB_MASTER;
				}
				else {
					lookupKey = DynamicDataSourceHolder.DB_SLAVE;
				}
			}
			else {
				BoundSql boundSql = ms.getSqlSource().getBoundSql(args[1]);
				String sql = boundSql.getSql().toLowerCase().replaceAll("[\\t\\n\\r]", " ");
				if (sql.matches(REGEX_WRITE)) {
					lookupKey = DynamicDataSourceHolder.DB_MASTER;
				}
				else {
					lookupKey = DynamicDataSourceHolder.DB_SLAVE;
				}
			}
		}
		else {
			lookupKey = DynamicDataSourceHolder.DB_MASTER;
		}
		DynamicDataSourceHolder.setDbType(lookupKey);
		
		// Execution proceed
		long startTime = System.currentTimeMillis();
		Object returnValue = invocation.proceed();
		long elapsedTime = System.currentTimeMillis() - startTime;
		logger.debug(String.format("Method [%s] use [%s] Strategy, SqlType [%s], Time [%s]ms.", 
				ms.getId(), lookupKey, ms.getSqlCommandType().name(), elapsedTime));
		return returnValue;
	}

	@Override
	public Object plugin(Object target) {
		
		if (target instanceof Executor) {
			return Plugin.wrap(target, this);
		}
		return target;
	}

	@Override
	public void setProperties(Properties properties) {
		// nop
	}

}

Mapper メソッド名で分ける

・上記の DynamicDataSourceInterceptor（一部割愛）
・振り分けの割合も設定してみたり (7:3)

	private static final String[] SLAVE_METHOD_IDS = {
			"your.project.domain.SampleMapper.countSqlSample",
			"your.project.domain.SampleMapper.selectSqlSample"
	};
	
	@Override
	public Object intercept(Invocation invocation) throws Throwable {
		
		boolean synchronizationActive = TransactionSynchronizationManager.isSynchronizationActive();
		Object[] args = invocation.getArgs();
		MappedStatement ms = (MappedStatement) args[0];
		String lookupKey = "";
		
		if (!synchronizationActive) {
			// 特定のSQLのみSLAVE実行
			String methodId = ms.getId();
			for (String method: SLAVE_METHOD_IDS) {
				if (Objects.equals(method, methodId)) {
					int max = 10;
					int min = 1;
					Random r = new Random();
					int which  = r.nextInt((max - min) + 1) + min;
                    // 7
					if (which > 3) {
						lookupKey = DynamicDataSourceHolder.DB_SLAVE;
					}
                    // 3
					else {
						lookupKey = DynamicDataSourceHolder.DB_MASTER;
					}
					break;
				}
			}
		}
		else {
			lookupKey = DynamicDataSourceHolder.DB_MASTER;
		}
		DynamicDataSourceHolder.setDbType(lookupKey);
		
		// Execution proceed
		long startTime = System.currentTimeMillis();
		Object returnValue = invocation.proceed();
		long elapsedTime = System.currentTimeMillis() - startTime;
		logger.debug(String.format("Method [%s] use [%s] Strategy, SqlType [%s], Time [%s]ms.", 
				ms.getId(), lookupKey, ms.getSqlCommandType().name(), elapsedTime));
		return returnValue;
	}

jdbc 設定

・properties にユーザなどの接続設定を記載しているものとします

	<bean id="sample.dataSource.master" class="org.apache.commons.dbcp.BasicDataSource"
		destroy-method="close">
		<property name="driverClassName" value="com.mysql.jdbc.Driver" />
		<property name="url" value="${jdbc.master.url}" />
		<property name="username" value="${jdbc.master.username}" />
		<property name="password" value="${jdbc.master.password}" />
        
        ...途中割愛...
	</bean>
	<bean id="sample.dataSource.slave" class="org.apache.commons.dbcp.BasicDataSource"
		destroy-method="close">
		<property name="driverClassName" value="com.mysql.jdbc.Driver" />
		<property name="url" value="${jdbc.slave.url}" />
		<property name="username" value="${jdbc.slave.username}" />
		<property name="password" value="${jdbc.slave.password}" />

        ...途中割愛...
	</bean>

	<bean id="sample.dataSource" class="your.project.domain.springframework.data.DynamicDataSource">
		<property name="targetDataSources">
			<map key-type="java.lang.String">
				<entry key="master" value-ref="sample.dataSource.master" />
				<entry key="slave" value-ref="sample.dataSource.slave" />
			</map>
		</property>
		<property name="defaultTargetDataSource" ref="sample.dataSource.master" />
	</bean>

	<bean id="sample.sqlSessionFactory" class="org.mybatis.spring.SqlSessionFactoryBean">
		<property name="dataSource" ref="sample.dataSource" />
		<property name="plugins">
			<array>
				<bean class="your.project.domain.springframework.data.DynamicDataSourceInterceptor" />
			</array>
		</property>
	</bean>

	<bean id="sample.sqlSessionTemplate" class="org.mybatis.spring.SqlSessionTemplate" scope="prototype">
		<constructor-arg index="0" ref="sample.sqlSessionFactory" />
	</bean>

以上、お疲れ様でした！