Mybatis 通过拦截器动态修改SQL

01 使用场景

当我们在多租户的项目中，编写SQL语句都要带上tenant字段，用于区分不同的租户只能操作自己的数据。

比如，像下面的SQL

select * from member where id = #{id} and tenant = #{tenant}

但是，我们的tenant属性一般是放在一个ThreadLocal的变量中的，便于我们随时获取该属性进行一些业务处理。

public class TenantContext {
    private static final ThreadLocal<String> TENANT_CONTEXT = new ThreadLocal<>();

    private TenantContext() {
    }

    public static void set(String tenant) {
        TENANT_CONTEXT.set(tenant);
    }

    public static String get() {
        return TENANT_CONTEXT.get();
    }

    public static void remove() {
        TENANT_CONTEXT.remove();
    }
}

这时，我们就可以在SQL中编写一些特殊字符串，然后用TenantContext.get()进行替换

-- 原始SQL
select * from member where id = #{id} and tenant = @tenant@
-- 替换后SQL
select * from member where id = #{id} and tenant = 'xxx'

这样就不用每个方法都要传tenant字段，然后替换#{tenant}了。

02 代码实现

首先，我们要编写一个mybatis的拦截器

import com.eddie.context.TenantContext;
import org.apache.ibatis.executor.statement.StatementHandler;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.plugin.Interceptor;
import org.apache.ibatis.plugin.Intercepts;
import org.apache.ibatis.plugin.Invocation;
import org.apache.ibatis.plugin.Signature;

import java.lang.reflect.Field;
import java.sql.Connection;
import java.util.regex.Pattern;

/**
 * 动态修改 mybatis 的 SQL
 *
 * @author eddie
 */
@Intercepts(@Signature(type = StatementHandler.class, method = "prepare", args = {Connection.class, Integer.class}))
public class MybatisSqlInterceptor implements Interceptor {

    private static final Pattern TENANT_PATTERN = Pattern.compile("@tenant@");

    @Override
    public Object intercept(Invocation invocation) throws Throwable {
        StatementHandler handler = (StatementHandler) invocation.getTarget();
        BoundSql boundSql = handler.getBoundSql();
        String sql = boundSql.getSql();
        sql = TENANT_PATTERN.matcher(sql).replaceAll("'" + TenantContext.get() + "'");
        Field sqlField = boundSql.getClass().getDeclaredField("sql");
        sqlField.setAccessible(true);
        sqlField.set(boundSql, sql);
        return invocation.proceed();
    }
}

然后，将这个拦截器放到spring容器中

import com.baomidou.mybatisplus.extension.plugins.PaginationInterceptor;
import com.baomidou.mybatisplus.extension.plugins.pagination.optimize.JsqlParserCountOptimize;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

@Configuration
public class MybatisPlusConfig {

    @Bean
    public MybatisSqlInterceptor mybatisSqlInterceptor() {
        return new MybatisSqlInterceptor();
    }
}

最后，编写我们的SQL语句，做测试即可

<select id="getById" resultType="com.eddie.bean.Member">
    select * from member where id = #{id} and tenant = @tenant@
</select>

03 测试

@RunWith(SpringRunner.class)
@SpringBootTest(classes = MybatisPlusApplication.class)
public class MemberTest {
    @Resource
    private MemberMapper memberMapper;

    @Before
    public void before() {
        TenantContext.set("t1");
    }

    @Test
    public void test_getById() {
        Member member = memberMapper.getById(1);
        System.out.println(JSON.toJSONString(member));
    }
}

测试结果