jnpf-common-v6x/jnpf-boot-common/jnpf-common-database/target/classes/jnpf/database/plugins/MyTenantMasterSlaveAutoRoutingPlugin.java

/*
 * Copyright © 2018 organization baomidou
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package jnpf.database.plugins;

import cn.hutool.core.text.StrPool;
import com.baomidou.dynamic.datasource.DynamicRoutingDataSource;
import com.baomidou.dynamic.datasource.enums.DdConstants;
import com.baomidou.dynamic.datasource.toolkit.DynamicDataSourceContextHolder;
import com.baomidou.dynamic.datasource.tx.TransactionContext;
import jnpf.constant.MsgCode;
import jnpf.database.util.DynamicDataSourceUtil;
import jnpf.exception.DataException;
import jnpf.util.TenantHolder;
import lombok.extern.slf4j.Slf4j;
import org.apache.ibatis.cache.CacheKey;
import org.apache.ibatis.executor.Executor;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.mapping.SqlCommandType;
import org.apache.ibatis.plugin.*;
import org.apache.ibatis.session.ResultHandler;
import org.apache.ibatis.session.RowBounds;
import org.springframework.transaction.support.TransactionSynchronizationManager;
import org.springframework.util.StringUtils;

import javax.sql.DataSource;
import java.util.Optional;

/**
 * 租户连接模式读写分离
 *
 * @author TaoYu
 * @since 2.5.1
 */
@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})})
@Slf4j
public class MyTenantMasterSlaveAutoRoutingPlugin implements Interceptor, ITenantPlugin {

    protected DynamicRoutingDataSource dynamicDataSource;

    public MyTenantMasterSlaveAutoRoutingPlugin(DataSource dataSource){
        this.dynamicDataSource = (DynamicRoutingDataSource) dataSource;
    }

    @Override
    public Object intercept(Invocation invocation) throws Throwable {
        if(TenantHolder.getLocalTenantCache() == null){
            printNoTenant(v -> log.warn("未设置租户信息, 禁止查询数据库, {}, {}, {}, {}", v.getUserId(), v.getUrl(), v.getToken(), v.getStack()));
            //未设置租户信息不允许操作数据库
            throw new DataException(MsgCode.LOG113.get());
        }
        if (!TenantHolder.getLocalTenantCache().isRemote()
                || !DynamicDataSourceUtil.isPrimaryDataSoure()) {
            return invocation.proceed();
        }
        Object[] args = invocation.getArgs();
        MappedStatement ms = (MappedStatement) args[0];
        String pushedDataSource = null;
        try {
            String tenantId = Optional.ofNullable(TenantHolder.getLocalTenantCache().getEnCode()).orElse("");
            String masterKey = tenantId + StrPool.DASHED +DdConstants.MASTER;
            String slaveKey = tenantId + StrPool.DASHED +DdConstants.SLAVE;
            // 存在事务只使用主库
            boolean hasTrans = TransactionSynchronizationManager.isActualTransactionActive();
            if (!hasTrans) {
                hasTrans = StringUtils.hasText(TransactionContext.getXID());
            }
            // 判断切库
            String dataSource = SqlCommandType.SELECT == ms.getSqlCommandType() ? slaveKey :masterKey;
            if (hasTrans || !dynamicDataSource.getGroupDataSources().containsKey(dataSource)) {
                dataSource = masterKey;
            }
            pushedDataSource = DynamicDataSourceContextHolder.push(dataSource);
            return invocation.proceed();
        } finally {
            if (pushedDataSource != null) {
                DynamicDataSourceContextHolder.poll();
            }
        }
    }
}