package com.zhgd.mybatis; import com.alibaba.fastjson.JSON; import com.alibaba.fastjson.JSONObject; import com.baomidou.mybatisplus.core.parser.SqlParserHelper; import com.baomidou.mybatisplus.core.plugins.InterceptorIgnoreHelper; import com.baomidou.mybatisplus.core.toolkit.PluginUtils; import com.baomidou.mybatisplus.core.toolkit.StringPool; import com.baomidou.mybatisplus.extension.parser.JsqlParserSupport; import com.baomidou.mybatisplus.extension.plugins.inner.InnerInterceptor; import com.zhgd.annotation.DataScope; import com.zhgd.xmgl.security.util.SecurityUtils; import lombok.Setter; import lombok.extern.slf4j.Slf4j; import net.sf.jsqlparser.schema.Table; import net.sf.jsqlparser.statement.insert.Insert; import net.sf.jsqlparser.statement.select.*; import org.apache.commons.collections.MapUtils; 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.mapping.SqlCommandType; import org.apache.ibatis.session.ResultHandler; import org.apache.ibatis.session.RowBounds; import java.lang.reflect.Method; import java.sql.Connection; import java.sql.SQLException; import java.util.Arrays; import java.util.Map; import java.util.Objects; import java.util.Optional; @Slf4j public class DataScopeInterceptor extends JsqlParserSupport implements InnerInterceptor { @Setter private DataScopeHandler dataScopeHandler; public void beforePrepare(StatementHandler sh, Connection connection, Integer transactionTimeout) { PluginUtils.MPStatementHandler mpSh = PluginUtils.mpStatementHandler(sh); MappedStatement ms = mpSh.mappedStatement(); SqlCommandType sct = ms.getSqlCommandType(); try { if (sct == SqlCommandType.INSERT) { Class clazz = Class.forName(ms.getId().substring(0, ms.getId().lastIndexOf(StringPool.DOT))); //注解判断 DataScope annotation = clazz.getAnnotation(DataScope.class); if (annotation == null || annotation.type() == 2) { return; } if (InterceptorIgnoreHelper.willIgnoreTenantLine(ms.getId())) return; if (SqlParserHelper.getSqlParserInfo(ms)) return; PluginUtils.MPBoundSql mpBs = mpSh.mPBoundSql(); Map param = JSONObject.parseObject(JSON.toJSONString(mpBs.parameterObject()), Map.class); mpBs.sql(parserMulti(mpBs.sql(), MapUtils.getString(param, "engineeringSn"))); } } catch (Exception e) { log.error(e.getMessage(), e); } } @Override protected void processInsert(Insert insert, int index, String sql, Object obj) { //dataScopeHandler.addParam(insert, obj); } public void beforeQuery(Executor executor, MappedStatement ms, Object parameter, RowBounds rowBounds, ResultHandler resultHandler, BoundSql boundSql) throws SQLException { try { if (SecurityUtils.getUser() == null) { return; } Class clazz = Class.forName(ms.getId().substring(0, ms.getId().lastIndexOf(StringPool.DOT))); String methodName = ms.getId().substring(ms.getId().lastIndexOf(".") + 1); DataScope annotation = null; Method[] declaredMethods = clazz.getDeclaredMethods(); Optional dsOption = Arrays.stream(declaredMethods).filter(method -> method.getName().equals(methodName)).map(method -> method.getAnnotation(DataScope.class)).filter(Objects::nonNull).findFirst(); annotation = dsOption.orElseGet(() -> clazz.getAnnotation(DataScope.class)); if (findIgnoreDataScope(parameter, annotation)) { return; } PluginUtils.MPBoundSql mpBs = PluginUtils.mpBoundSql(boundSql); mpBs.sql(this.parserSingle(mpBs.sql(), annotation)); } catch (Exception e) { log.error(e.getMessage(), e); } } protected void processSelect(Select select, int index, String sql, Object obj) { this.processSelectBody(select.getSelectBody(), obj); } private boolean findIgnoreDataScope(Object parameter, DataScope annotation) { if (annotation == null || !annotation.enable()) { return true; } if (parameter instanceof Map) { for (Object val : ((Map) parameter).values()) { if (val instanceof String) { if (val.equals("ignoreDataScope")) { return true; } } } } if (parameter instanceof String) { return parameter.equals("ignoreDataScope"); } return false; } protected void processSelectBody(SelectBody selectBody, Object obj) { if (selectBody != null) { if (selectBody instanceof PlainSelect) { this.processPlainSelect((PlainSelect) selectBody, obj); } else if (selectBody instanceof WithItem) { WithItem withItem = (WithItem) selectBody; this.processSelectBody(withItem.getSelectBody(), obj); } else { SetOperationList operationList = (SetOperationList) selectBody; if (operationList.getSelects() != null && operationList.getSelects().size() > 0) { this.processSelectBody(operationList.getSelects().get(0), obj); } } } } protected void processPlainSelect(PlainSelect plainSelect, Object obj) { FromItem fromItem = plainSelect.getFromItem(); if (fromItem instanceof Table) { this.dataScopeHandler.getSqlSegment(plainSelect, obj); } else { processFromItem(fromItem, obj); } } protected void processFromItem(FromItem fromItem, Object obj) { if (fromItem instanceof SubSelect) { SubSelect subSelect = (SubSelect) fromItem; if (subSelect.getSelectBody() != null) { processSelectBody(subSelect.getSelectBody(), obj); } } } }