com.github.antelopeframework.mybatis.shard.plugin.ShardPlugin Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of antelope-mybatis-ext Show documentation
There is a newer version: 1.1.5
package com.github.antelopeframework.mybatis.shard.plugin;

import java.lang.reflect.InvocationHandler;
import java.lang.reflect.Proxy;
import java.util.Map;
import java.util.Properties;

import org.apache.ibatis.executor.statement.RoutingStatementHandler;
import org.apache.ibatis.executor.statement.StatementHandler;
import org.apache.ibatis.mapping.MappedStatement;
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 com.github.antelopeframework.mybatis.shard.ShardContext;
import com.github.antelopeframework.mybatis.shard.ShardContext.ShardOnPair;
import com.github.antelopeframework.mybatis.shard.converter.SqlConverterFactory;
import com.github.antelopeframework.mybatis.util.ReflectionUtils;

import lombok.extern.slf4j.Slf4j;

/**
 * MyBatis 数据库分表插件.
 * 
 * @author yangzhi.yzh
 *
 */
@Slf4j
@Intercepts({ @org.apache.ibatis.plugin.Signature(type = StatementHandler.class, method = "prepare", args = {java.sql.Connection.class, Integer.class}) })
public class ShardPlugin implements Interceptor {

	@Override
	public Object intercept(Invocation invocation) throws Throwable {
		//~ 没有标注分区
		Map shardOns = ShardContext.getShardOns();
		if (shardOns == null || shardOns.isEmpty()) {
			return invocation.proceed();
		}
		
		StatementHandler statementHandler = null;
		
		//~ 如果配置了多个 MyBatis 插件, 可能存在 jdk 代理类  
		if (Proxy.isProxyClass(invocation.getTarget().getClass())) {
			InvocationHandler invocationHandler = Proxy.getInvocationHandler(invocation.getTarget());
			org.apache.ibatis.plugin.Plugin plugin = (org.apache.ibatis.plugin.Plugin) invocationHandler;
			
			statementHandler = (StatementHandler) ReflectionUtils.getFieldValue(plugin, "target");
		} else {
			statementHandler = (StatementHandler) invocation.getTarget();
		}
		
		MappedStatement mappedStatement = null;
		if ((statementHandler instanceof RoutingStatementHandler)) {
			StatementHandler delegate = (StatementHandler) ReflectionUtils.getFieldValue(statementHandler, "delegate");
			mappedStatement = (MappedStatement) ReflectionUtils.getFieldValue(delegate, "mappedStatement");
		} else {
			mappedStatement = (MappedStatement) ReflectionUtils.getFieldValue(statementHandler, "mappedStatement");
		}
		
		String mapperId = mappedStatement.getId();
		String sql = statementHandler.getBoundSql().getSql();

		if (log.isDebugEnabled()) {
			log.debug("Original Sql [{}]: {}", mapperId, sql);
		}
		
		SqlConverterFactory cf = SqlConverterFactory.getInstance();
		sql = cf.convert(mapperId, sql, shardOns);
		
		if (log.isDebugEnabled()) {
			log.debug("Converted Sql [{}]: {}", mapperId, sql);
		}
		
		ReflectionUtils.setFieldValue(statementHandler.getBoundSql(), "sql", sql);
		
		return invocation.proceed();
	}
	
	@Override
	public Object plugin(Object target) {
		return Plugin.wrap(target, this);
	}

	@Override
	public void setProperties(Properties properties) {
		//DO-NOTHING
	}
}