wisdomisite-java/src/main/java/com/zhgd/mybatis/DataScopeInterceptor.java

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 lombok.Setter;
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.sql.Connection;
import java.sql.SQLException;
import java.util.Map;

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) {
            throw new RuntimeException(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 {
            Class<?> clazz = Class.forName(ms.getId().substring(0, ms.getId().lastIndexOf(StringPool.DOT)));
            DataScope annotation = clazz.getAnnotation(DataScope.class);
            if (findIgnoreDataScope(parameter) || annotation == null) {
                return;
            }
            PluginUtils.MPBoundSql mpBs = PluginUtils.mpBoundSql(boundSql);
            mpBs.sql(this.parserSingle(mpBs.sql(), annotation.type()));
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    protected void processSelect(Select select, int index, String sql, Object obj) {
        this.processSelectBody(select.getSelectBody(), obj);
    }

    private boolean findIgnoreDataScope(Object parameter) {
        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);
            }
        }
    }
}