package jnpf.database.plugins; import com.baomidou.mybatisplus.core.toolkit.CollectionUtils; import com.baomidou.mybatisplus.extension.plugins.inner.TenantLineInnerInterceptor; import jnpf.constant.MsgCode; import jnpf.database.util.DynamicDataSourceUtil; import jnpf.database.util.NotTenantPluginHolder; import jnpf.exception.DataException; import jnpf.util.TenantHolder; import jodd.util.StringPool; import lombok.extern.slf4j.Slf4j; import net.sf.jsqlparser.expression.Expression; import net.sf.jsqlparser.schema.Column; import net.sf.jsqlparser.schema.Table; import net.sf.jsqlparser.statement.select.*; import org.apache.ibatis.executor.Executor; import org.apache.ibatis.executor.statement.StatementHandler; import org.apache.ibatis.mapping.BoundSql; import org.apache.ibatis.mapping.MappedStatement; import org.apache.ibatis.session.ResultHandler; import org.apache.ibatis.session.RowBounds; import java.sql.Connection; import java.sql.SQLException; import java.util.List; /** * Column模式租户插件 * @author JNPF开发平台组 * @user N * @copyright 引迈信息技术有限公司 * @date 2022/10/14 10:29 */ @Slf4j public class MyTenantLineInnerInterceptor extends TenantLineInnerInterceptor implements ITenantPlugin { @Override public void beforeQuery(Executor executor, MappedStatement ms, Object parameter, RowBounds rowBounds, ResultHandler resultHandler, BoundSql boundSql) { 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().isColumn()) { return; } if (NotTenantPluginHolder.isNotSwitch()) { NotTenantPluginHolder.clearNotSwitchFlag(); return; } //非主库不切库 if(!DynamicDataSourceUtil.isPrimaryDataSoure()){ return; } //不绑定数据源的接口不切库 /*if(NotTenantPluginHolder.isNotSwitchAlways()){ return; }*/ try { super.beforeQuery(executor, ms, parameter, rowBounds, resultHandler, boundSql); } catch (Exception e){ //特殊语句解析失败 if(log.isDebugEnabled()){ log.debug("语句解析失败", e); } } } @Override public void beforePrepare(StatementHandler sh, Connection connection, Integer transactionTimeout) { 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().isColumn()) { return; } if (NotTenantPluginHolder.isNotSwitch()) { NotTenantPluginHolder.clearNotSwitchFlag(); return; } //非主库不切库 if(!DynamicDataSourceUtil.isPrimaryDataSoure()){ return; } //不绑定数据源的接口不切库 /*if(NotTenantPluginHolder.isNotSwitchAlways()){ return; }*/ try { super.beforePrepare(sh, connection, transactionTimeout); } catch (Exception e){ //特殊语句解析失败 if(log.isDebugEnabled()){ log.debug("语句解析失败", e); } } } @Override protected Column getAliasColumn(Table table) { return getAliasColumnWithFromItem(table); } protected Column getAliasColumnWithFromItem(FromItem table) { StringBuilder column = new StringBuilder(); if (table.getAlias() != null) { column.append(table.getAlias().getName()); }else{ if(table instanceof Table){ column.append(((Table)table).getName()); } } column.append(StringPool.DOT).append(super.getTenantLineHandler().getTenantIdColumn()); return new Column(column.toString()); } protected void appendSelectItem(List> selectItems, FromItem from) { if (CollectionUtils.isEmpty(selectItems)) { return; } if (selectItems.size() == 1) { SelectItem item = selectItems.get(0); Expression expression = item.getExpression(); if (expression instanceof AllColumns) { return; } } // 改 selectItems.add(new SelectItem<>(getAliasColumnWithFromItem(from))); } @Override protected void processInsertSelect(Select selectBody, final String whereSegment) { if (selectBody instanceof PlainSelect) { PlainSelect plainSelect = (PlainSelect)selectBody; FromItem fromItem = plainSelect.getFromItem(); if (fromItem instanceof Table) { this.processPlainSelect(plainSelect, whereSegment); // 改 this.appendSelectItem(plainSelect.getSelectItems(), fromItem); } else if (fromItem instanceof Select) { Select subSelect = (Select)fromItem; // 改 this.appendSelectItem(plainSelect.getSelectItems(), fromItem); this.processInsertSelect(subSelect, whereSegment); } } else if (selectBody instanceof ParenthesedSelect) { ParenthesedSelect parenthesedSelect = (ParenthesedSelect)selectBody; this.processInsertSelect(parenthesedSelect.getSelect(), whereSegment); } } }