MybatisPlus對租戶模式的支援（二）——重寫TenantSqlParser

前言

上期說到，假設我是系統管理者，想看到租戶A資料的同時也能看到租戶B的資料，這樣就不能給系統管理者的賬号設定租戶id。

但不設定租戶id的話，當系統管理者進行查詢時，sql 的 where 條件會加上 ”tenant_id = null”，這個條件會導緻管理者連 租戶A 和 租戶B 的資料都看不到。

我們可以在加上 ”tenant_id = ？” 條件前判斷目前使用者是否為系統管理者，如果是系統管理者，則直接傳回，不處理為 sql 語句添加租戶判斷條件的代碼邏輯。

通過檢視 MybatisPlus 源碼，得知處理租戶條件的 sql 解析器位置如下：

com.baomidou.mybatisplus.extension.plugins.tenant.TenantSqlParser           

那麼我們可以參考和重寫它，讓 MybatisPlus 使用我們自定義的租戶 sql 解析器。

具體操作

首先，在 MyTenantHandler 加一下 doUserFilter() 接口，用于判斷目前使用者是否為系統管理者。

package com.rjkj.quickboot.base.tenant;

import com.baomidou.mybatisplus.extension.plugins.tenant.TenantHandler;

public interface MyTenantHandler extends TenantHandler {
    // 判斷目前使用者是否為系統管理者
    boolean doUserFilter();

    void setTenantId(String tenantId);
}           

然後，參考 TenantSqlParser 建立我們自定義的租戶 sql 解析器，對 CRUD 的處理方法進行重寫，代碼可以直接複制，在方法開始前額外加上以下判斷：

MyTenantHandler tenantHandler = (MyTenantHandler) getTenantHandler();
// 過濾系統管理者
if (tenantHandler.doUserFilter()){
    return;
}           

整體代碼如下：

package com.rjkj.quickboot.base.tenant;

import com.baomidou.mybatisplus.core.toolkit.Assert;
import com.baomidou.mybatisplus.core.toolkit.ExceptionUtils;
import com.baomidou.mybatisplus.extension.plugins.tenant.TenantSqlParser;
import lombok.Data;
import lombok.EqualsAndHashCode;
import lombok.experimental.Accessors;
import net.sf.jsqlparser.expression.operators.relational.ExpressionList;
import net.sf.jsqlparser.expression.operators.relational.ItemsList;
import net.sf.jsqlparser.expression.operators.relational.MultiExpressionList;
import net.sf.jsqlparser.schema.Column;
import net.sf.jsqlparser.schema.Table;
import net.sf.jsqlparser.statement.delete.Delete;
import net.sf.jsqlparser.statement.insert.Insert;
import net.sf.jsqlparser.statement.select.*;
import net.sf.jsqlparser.statement.update.Update;

import javax.annotation.PostConstruct;
import java.util.List;

@Accessors(chain = true)
@EqualsAndHashCode(callSuper = true)
public class MyTenantSqlParser extends TenantSqlParser {

    /**
     * select 語句處理
     */
    @Override
    public void processSelectBody(SelectBody selectBody) {
        MyTenantHandler tenantHandler = (MyTenantHandler) getTenantHandler();
        if (tenantHandler.doUserFilter()){
            // 過濾系統管理者
            return;
        }
        if (selectBody instanceof PlainSelect) {
            processPlainSelect((PlainSelect) selectBody);
        } else if (selectBody instanceof WithItem) {
            WithItem withItem = (WithItem) selectBody;
            if (withItem.getSelectBody() != null) {
                processSelectBody(withItem.getSelectBody());
            }
        } else {
            SetOperationList operationList = (SetOperationList) selectBody;
            if (operationList.getSelects() != null && operationList.getSelects().size() > 0) {
                operationList.getSelects().forEach(this::processSelectBody);
            }
        }
    }

    /**
     * insert 語句處理
     */
    @Override
    public void processInsert(Insert insert) {
        MyTenantHandler tenantHandler = (MyTenantHandler) getTenantHandler();
        if (tenantHandler.doTableFilter(insert.getTable().getName())) {
            // 過濾退出執行
            return;
        }
        if (tenantHandler.doUserFilter()){
            // 過濾系統管理者
            return;
        }
        insert.getColumns().add(new Column(tenantHandler.getTenantIdColumn()));
        if (insert.getSelect() != null) {
            processPlainSelect((PlainSelect) insert.getSelect().getSelectBody(), true);
        } else if (insert.getItemsList() != null) {
            // fixed github pull/295
            ItemsList itemsList = insert.getItemsList();
            if (itemsList instanceof MultiExpressionList) {
                ((MultiExpressionList) itemsList).getExprList().forEach(el -> el.getExpressions().add(tenantHandler.getTenantId()));
            } else {
                ((ExpressionList) insert.getItemsList()).getExpressions().add(tenantHandler.getTenantId());
            }
        } else {
            throw ExceptionUtils.mpe("Failed to process multiple-table update, please exclude the tableName or statementId");
        }
    }

    /**
     * update 語句處理
     */
    @Override
    public void processUpdate(Update update) {
        MyTenantHandler tenantHandler = (MyTenantHandler) getTenantHandler();
        if (tenantHandler.doUserFilter()){
            // 過濾系統管理者
            return;
        }
        List<Table> tableList = update.getTables();
        Assert.isTrue(null != tableList && tableList.size() < 2,
                "Failed to process multiple-table update, please exclude the statementId");
        Table table = tableList.get(0);
        if (tenantHandler.doTableFilter(table.getName())) {
            // 過濾退出執行
            return;
        }
        update.setWhere(this.andExpression(table, update.getWhere()));
    }

    /**
     * delete 語句處理
     */
    @Override
    public void processDelete(Delete delete) {
        MyTenantHandler tenantHandler = (MyTenantHandler) getTenantHandler();
        if (tenantHandler.doUserFilter()){
            // 過濾系統管理者
            return;
        }
        if (tenantHandler.doTableFilter(delete.getTable().getName())) {
            // 過濾退出執行
            return;
        }
        delete.setWhere(this.andExpression(delete.getTable(), delete.getWhere()));
    }

}
           

同時對上期文章中的部分代碼進行修改，主要是對處理類進行類型轉換。（為了簡單顯示，隻列印需要改動的代碼部分）

MybatisPlusConfig.java

@Configuration
public class MybatisPlusConfig {

    ...

    @Bean
    public PaginationInterceptor paginationInterceptor() {
        PaginationInterceptor paginationInterceptor = new PaginationInterceptor();
        // 自定義租戶sql解析器
        MyTenantSqlParser tenantSqlParser = new MyTenantSqlParser();
        // 自定義租戶sql處理類
        tenantSqlParser.setTenantHandler(new MyTenantHandler() {
            @Override
            public boolean doUserFilter() {
                // 如果是系統管理者，它的租戶id為null
                return ObjectUtil.isNull(getTenantId());
            }
            ...
        }
    }
    ...
           

TenantInterceptor.java

@Component
public class TenantInterceptor implements HandlerInterceptor {

    @Autowired
    private PaginationInterceptor pi;

    @Override
    public boolean preHandle(HttpServletRequest httpServletRequest, HttpServletResponse httpServletResponse, Object o) throws Exception {
        // 自定義租戶sql解析器
        MyTenantSqlParser tenantSqlParser = (MyTenantSqlParser) pi.getSqlParserList().get(0);
        MyTenantHandler myTenantHandler = (MyTenantHandler) tenantSqlParser.getTenantHandler();

        //擷取目前使用者
        LoginUser currentUser = getCurrentUser();
        myTenantHandler.setTenantId(currentUser.getTenantId());
        return true;
    }
    ...
}           

相比 Mybatis，Mybatis-Plus 的支援更加豐富，除了本文描述的對租戶支援，還内置了各種插件，如性能分析、分頁插件等，并且使用強大的條件構造器，足以替代手寫 sql，滿足大部分的查詢需求。

下次有機會再繼續研究 Mybatis-Plus 的其他的特性，希望本文對大家有所幫助。

以上就是本期分享，如果大家對此感興趣，歡迎各位關注、留言，大家的支援就是我的動力！