Mybatis Generator Plugin悲觀鎖實現

raledong發表於2021-09-28
MyBatisPlugin

前言

Mybatis Generator外掛可以快速的實現基礎的資料庫CRUD操作，它同時支援JAVA語言和Kotlin語言，將程式設計師從重複的Mapper和Dao層程式碼編寫中釋放出來。Mybatis Generator可以自動生成大部分的SQL程式碼，如update，updateSelectively，insert，insertSelectively，select語句等。但是，當程式中需要SQL不在自動生成的SQL範圍內時，就需要使用自定義Mapper來實現，即手動編寫DAO層和Mapper檔案（這裡有一個小坑，當資料庫實體增加欄位時，對應的自定義Mapper也要及時手動更新）。拋開復雜的定製化SQL如join，group by等，其實還是有一些比較常用的SQL在基礎的Mybatis Generator工具中沒有自動生成，比如分頁能力，悲觀鎖，樂觀鎖等，而Mybatis Generator也為這些訴求提供了Plugin的能力。通過自定義實現Plugin可以改變Mybatis Generator在生成Mapper和Dao檔案時的行為。本文將從悲觀鎖為例，讓你快速瞭解如何實現Mybatis Generator Plugin。

實現背景：
資料庫：MYSQL
mybatis generator runtime：MyBatis3

<!-- more -->

實現Mybatis悲觀鎖

當業務出現需要保證強一致的場景時，可以通過在事務中對資料行上悲觀鎖後再進行操作來實現，這就是經典的”一鎖二判三更新“。在交易或是支付系統中，這種訴求非常普遍。Mysql提供了Select...For Update語句來實現對資料行上悲觀鎖。本文將不對Select...For Update進行詳細的介紹，有興趣的同學可以檢視其它文章深入瞭解。

Mybatis Generator Plugin為這種具有通用性的SQL提供了很好的支援。通過繼承org.mybatis.generator.api.PluginAdapter類即可自定義SQL生成邏輯並在在配置檔案中使用。PluginAdapter是Plugin介面的實現類，提供了Plugin的預設實現，本文將介紹其中比較重要的幾個方法：

public interface Plugin {
    /**
    * 將Mybatis Generator配置檔案中的上下文資訊傳遞到Plugin實現類中
    * 這些資訊包括資料庫連結，型別對映配置等
    */
    void setContext(Context context);

    /**
    * 配置檔案中的所有properties標籤
    **/
    void setProperties(Properties properties);

    /**
    * 校驗該Plugin是否執行，如果返回false，則該外掛不會執行
    **/
    boolean validate(List<String> warnings);

    /**
    * 當DAO檔案完成生成後會觸發該方法，可以通過實現該方法在DAO檔案中新增方法或屬性
    **/
    boolean clientGenerated(Interface interfaze, TopLevelClass topLevelClass,
            IntrospectedTable introspectedTable);

    /**
    * 當SQL XML 檔案生成後會呼叫該方法，可以通過實現該方法在MAPPER XML檔案中新增XML定義
    **/
    boolean sqlMapDocumentGenerated(Document document, IntrospectedTable introspectedTable);
}

這裡結合Mybatis Generator的配置檔案和生成的DAO（也稱為Client檔案）和Mapper XML檔案可以更好的理解。Mybatis Generator配置檔案樣例如下，其中包含了主要的一些配置資訊，如用於描述資料庫連結的<jdbcConnection>標籤，用於定義資料庫和Java型別轉換的<javaTypeResolver>標籤等。

<?xml version="1.0" encoding="UTF-8"?>
<!DOCTYPE generatorConfiguration
  PUBLIC "-//mybatis.org//DTD MyBatis Generator Configuration 1.0//EN"
  "http://mybatis.org/dtd/mybatis-generator-config_1_0.dtd">

<generatorConfiguration>
  <classPathEntry location="/Program Files/IBM/SQLLIB/java/db2java.zip" />

  <context id="DB2Tables" targetRuntime="MyBatis3">
    <jdbcConnection driverClass="COM.ibm.db2.jdbc.app.DB2Driver"
        connectionURL="jdbc:db2:TEST"
        userId="db2admin"
        password="db2admin">
    </jdbcConnection>

    <javaTypeResolver >
      <property name="forceBigDecimals" value="false" />
    </javaTypeResolver>

    <javaModelGenerator targetPackage="test.model" targetProject="\MBGTestProject\src">
      <property name="enableSubPackages" value="true" />
      <property name="trimStrings" value="true" />
    </javaModelGenerator>

    <sqlMapGenerator targetPackage="test.xml"  targetProject="\MBGTestProject\src">
      <property name="enableSubPackages" value="true" />
    </sqlMapGenerator>

    <javaClientGenerator type="XMLMAPPER" targetPackage="test.dao"  targetProject="\MBGTestProject\src">
      <property name="enableSubPackages" value="true" />
    </javaClientGenerator>

    <property name="printLog" value="true"/>

    <table schema="DB2ADMIN" tableName="ALLTYPES" domainObjectName="Customer" >
      <property name="useActualColumnNames" value="true"/>
      <generatedKey column="ID" sqlStatement="DB2" identity="true" />
      <columnOverride column="DATE_FIELD" property="startDate" />
      <ignoreColumn column="FRED" />
      <columnOverride column="LONG_VARCHAR_FIELD" jdbcType="VARCHAR" />
    </table>

  </context>
</generatorConfiguration>

這些都被對映成Context物件，並通過setContext(Context context)方法傳遞到具體的Plugin實現中：

public class Context extends PropertyHolder{

    /**
    * <context>標籤的id屬性
    */
    private String id;

    /**
    * jdbc連結資訊，對應<jdbcConnection>標籤中的資訊
    */
    private JDBCConnectionConfiguration jdbcConnectionConfiguration;

    /**
    * 型別對映配置，對應<javaTypeResolver>
    */
    private JavaTypeResolverConfiguration javaTypeResolverConfiguration;

    /**
    * ...其它標籤對應的配置資訊
    */
}

setProperties則將context下的<properties>標籤收集起來並對映成Properties類，它實際上是一個Map容器，正如Properties類本身就繼承了Hashtable。以上文中的配置檔案為例，可以通過properties.get("printLog")獲得值"true"。

validate方法則代表了這個Plugin是否執行，它通常進行一些非常基礎的校驗，比如是否相容對應的資料庫驅動或者是Mybatis版本：

    public boolean validate(List<String> warnings) {
        if (StringUtility.stringHasValue(this.getContext().getTargetRuntime()) && !"MyBatis3".equalsIgnoreCase(this.getContext().getTargetRuntime())) {
            logger.warn("itfsw:外掛" + this.getClass().getTypeName() + "要求執行targetRuntime必須為MyBatis3！");
            return false;
        } else {
            return true;
        }
    }

如果validate方法返回false，則無論什麼場景下都不會執行這個Plugin。

接著是最重要的兩個方法，分別是用於在DAO中生成新的方法clientGenerated和在XML檔案中生成新的SQL sqlMapDocumentGenerated。

先說clientGenerated，這個方法共有三個引數，interfaze是當前已經生成的客戶端Dao介面，topLevelClass是指生成的實現類，這個類可能為空，introspectedTable是指當前處理的資料表，這裡包含了從資料庫中獲取的關於表的各種資訊，包括列名稱，列型別等。這裡可以看一下introspectedTable中幾個比較重要的方法：

public abstract class IntrospectedTable {
    /**
    * 該方法可以獲得配置檔案中該表對應<table>標籤下的配置資訊，包括對映成的Mapper名稱，PO名稱等
    * 也可以在table標籤下自定義<property>標籤並通過getProperty方法獲得值
    */
    public TableConfiguration getTableConfiguration() {
        return tableConfiguration;
    }

    /**
    * 這個方法中定義了預設的生成規則，可以通過calculateAllFieldsClass獲得返回型別
    */
    public Rules getRules() {
        return rules;
    }
}

悲觀鎖的clientGenerated方法如下：

    // Plugin配置，是否要生成selectForUpdate語句
    private static final String CONFIG_XML_KEY = "implementSelectForUpdate";

    @Override
    public boolean clientGenerated(Interface interfaze, TopLevelClass topLevelClass, IntrospectedTable introspectedTable) {
        String implementUpdate = introspectedTable.getTableConfiguration().getProperty(CONFIG_XML_KEY);
        if (StringUtility.isTrue(implementUpdate)) {
            Method method = new Method(METHOD_NAME);
            FullyQualifiedJavaType returnType = introspectedTable.getRules().calculateAllFieldsClass();
            method.setReturnType(returnType);
            method.addParameter(new Parameter(new FullyQualifiedJavaType("java.lang.Long"), "id"));
            String docComment = "/**\n" +
                    "      * 使用id對資料行上悲觀鎖\n" +
                    "      */";
            method.addJavaDocLine(docComment);
            interfaze.addMethod(method);
            log.debug("(悲觀鎖外掛):" + interfaze.getType().getShortName() + "增加" + METHOD_NAME + "方法。");
        }

        return super.clientGenerated(interfaze, topLevelClass, introspectedTable);
    }

這裡可以通過在對應table下新增property標籤來決定是否要為這張表生成對應的悲觀鎖方法，配置樣例如下：

 <table tableName="demo" domainObjectName="DemoPO" mapperName="DemoMapper"
               enableCountByExample="true"
               enableUpdateByExample="true"
               enableDeleteByExample="true"
               enableSelectByExample="true"
               enableInsert="true"
               selectByExampleQueryId="true">
    <property name="implementUpdateWithCAS" value="true"/>
 </table>

程式碼中通過mybatis提供的Method方法，定義了方法的名稱，引數，返回型別等，並使用interfaze.addMethod方法將方法新增到客戶端的介面中。

再到sqlMapDocumentGenerated這個方法，這個方法中傳入了Document物件，它對應生成的XML檔案，並通過XmlElement來對映XML檔案中的元素。通過document.getRootElement().addElement可以將自定義的XML元素插入到Mapper檔案中。自定義XML元素就是指拼接XmlElement，XmlElement的addAttribute方法可以為XML元素設定屬性，addElement則可以為XML標籤新增子元素。有兩種型別的子元素，分別是TextElement和XmlElement本身，TextElement則直接填充標籤中的內容，而XmlElement則對應新的標籤，如<where> <include>等。悲觀鎖的SQL生成邏輯如下：

    // Plugin配置，是否要生成selectForUpdate語句
    private static final String CONFIG_XML_KEY = "implementSelectForUpdate";

    @Override
    public boolean sqlMapDocumentGenerated(Document document, IntrospectedTable introspectedTable) {
        String implementUpdate = introspectedTable.getTableConfiguration().getProperty(CONFIG_XML_KEY);
        if (!StringUtility.isTrue(implementUpdate)) {
            return super.sqlMapDocumentGenerated(document, introspectedTable);
        }

        XmlElement selectForUpdate = new XmlElement("select");
        selectForUpdate.addAttribute(new Attribute("id", METHOD_NAME));
        StringBuilder sb;

        String resultMapId = introspectedTable.hasBLOBColumns() ? introspectedTable.getResultMapWithBLOBsId() : introspectedTable.getBaseResultMapId();
        selectForUpdate.addAttribute(new Attribute("resultMap", resultMapId));
        selectForUpdate.addAttribute(new Attribute("parameterType", introspectedTable.getExampleType()));
        selectForUpdate.addElement(new TextElement("select"));

        sb = new StringBuilder();
        if (StringUtility.stringHasValue(introspectedTable.getSelectByExampleQueryId())) {
            sb.append('\'');
            sb.append(introspectedTable.getSelectByExampleQueryId());
            sb.append("' as QUERYID,");
            selectForUpdate.addElement(new TextElement(sb.toString()));
        }

        XmlElement baseColumn = new XmlElement("include");
        baseColumn.addAttribute(new Attribute("refid", introspectedTable.getBaseColumnListId()));
        selectForUpdate.addElement(baseColumn);
        if (introspectedTable.hasBLOBColumns()) {
            selectForUpdate.addElement(new TextElement(","));
            XmlElement blobColumns = new XmlElement("include");
            blobColumns.addAttribute(new Attribute("refid", introspectedTable.getBaseColumnListId()));
            selectForUpdate.addElement(blobColumns);
        }

        sb.setLength(0);
        sb.append("from ");
        sb.append(introspectedTable.getAliasedFullyQualifiedTableNameAtRuntime());
        selectForUpdate.addElement(new TextElement(sb.toString()));
        TextElement whereXml = new TextElement("where id = #{id} for update");
        selectForUpdate.addElement(whereXml);

        document.getRootElement().addElement(selectForUpdate);
        log.debug("(悲觀鎖外掛):" + introspectedTable.getMyBatis3XmlMapperFileName() + "增加" + METHOD_NAME + "方法(" + (introspectedTable.hasBLOBColumns() ? "有" : "無") + "Blob型別))。");
        return super.sqlMapDocumentGenerated(document, introspectedTable);
    }

完整程式碼

@Slf4j
public class SelectForUpdatePlugin extends PluginAdapter {

    private static final String CONFIG_XML_KEY = "implementSelectForUpdate";

    private static final String METHOD_NAME = "selectByIdForUpdate";

    @Override
    public boolean validate(List<String> list) {
        return true;
    }

    @Override
    public boolean clientGenerated(Interface interfaze, TopLevelClass topLevelClass, IntrospectedTable introspectedTable) {
        String implementUpdate = introspectedTable.getTableConfiguration().getProperty(CONFIG_XML_KEY);
        if (StringUtility.isTrue(implementUpdate)) {
            Method method = new Method(METHOD_NAME);
            FullyQualifiedJavaType returnType = introspectedTable.getRules().calculateAllFieldsClass();
            method.setReturnType(returnType);
            method.addParameter(new Parameter(new FullyQualifiedJavaType("java.lang.Long"), "id"));
            String docComment = "/**\n" +
                    "      * 使用id對資料行上悲觀鎖\n" +
                    "      */";
            method.addJavaDocLine(docComment);
            interfaze.addMethod(method);
            log.debug("(悲觀鎖外掛):" + interfaze.getType().getShortName() + "增加" + METHOD_NAME + "方法。");
        }

        return super.clientGenerated(interfaze, topLevelClass, introspectedTable);
    }

    @Override
    public boolean sqlMapDocumentGenerated(Document document, IntrospectedTable introspectedTable) {
        String implementUpdate = introspectedTable.getTableConfiguration().getProperty(CONFIG_XML_KEY);
        if (!StringUtility.isTrue(implementUpdate)) {
            return super.sqlMapDocumentGenerated(document, introspectedTable);
        }

        XmlElement selectForUpdate = new XmlElement("select");
        selectForUpdate.addAttribute(new Attribute("id", METHOD_NAME));
        StringBuilder sb;

        String resultMapId = introspectedTable.hasBLOBColumns() ? introspectedTable.getResultMapWithBLOBsId() : introspectedTable.getBaseResultMapId();
        selectForUpdate.addAttribute(new Attribute("resultMap", resultMapId));
        selectForUpdate.addAttribute(new Attribute("parameterType", introspectedTable.getExampleType()));
        selectForUpdate.addElement(new TextElement("select"));

        sb = new StringBuilder();
        if (StringUtility.stringHasValue(introspectedTable.getSelectByExampleQueryId())) {
            sb.append('\'');
            sb.append(introspectedTable.getSelectByExampleQueryId());
            sb.append("' as QUERYID,");
            selectForUpdate.addElement(new TextElement(sb.toString()));
        }

        XmlElement baseColumn = new XmlElement("include");
        baseColumn.addAttribute(new Attribute("refid", introspectedTable.getBaseColumnListId()));
        selectForUpdate.addElement(baseColumn);
        if (introspectedTable.hasBLOBColumns()) {
            selectForUpdate.addElement(new TextElement(","));
            XmlElement blobColumns = new XmlElement("include");
            blobColumns.addAttribute(new Attribute("refid", introspectedTable.getBaseColumnListId()));
            selectForUpdate.addElement(blobColumns);
        }

        sb.setLength(0);
        sb.append("from ");
        sb.append(introspectedTable.getAliasedFullyQualifiedTableNameAtRuntime());
        selectForUpdate.addElement(new TextElement(sb.toString()));
        TextElement whereXml = new TextElement("where id = #{id} for update");
        selectForUpdate.addElement(whereXml);

        document.getRootElement().addElement(selectForUpdate);
        log.debug("(悲觀鎖外掛):" + introspectedTable.getMyBatis3XmlMapperFileName() + "增加" + METHOD_NAME + "方法(" + (introspectedTable.hasBLOBColumns() ? "有" : "無") + "Blob型別))。");
        return super.sqlMapDocumentGenerated(document, introspectedTable);
    }

}

相關文章