mybatis 源碼分析(七)KeyGenerator 詳解

  • 2019 年 10 月 3 日
  • 筆記

一、KeyGenerator 概述

在平時開發的時候經常會有這樣的需求,插入數據返回主鍵,或者插入數據之前需要獲取主鍵,這樣的需求在 mybatis 中也是支援的,其中主要的邏輯部分就在 KeyGenerator 中,下面是他的類圖:

其中:

  • NoKeyGenerator:默認空實現,不需要對主鍵單獨處理;
  • Jdbc3KeyGenerator:主要用於資料庫的自增主鍵,比如 MySQL、PostgreSQL;
  • SelectKeyGenerator:主要用於資料庫不支援自增主鍵的情況,比如 Oracle、DB2;

介面方法如下:

public interface KeyGenerator {    void processBefore(Executor executor, MappedStatement ms, Statement stmt, Object parameter);    void processAfter(Executor executor, MappedStatement ms, Statement stmt, Object parameter);  }

如程式碼所見 KeyGenerator 非常的簡單,主要是通過兩個攔截方法實現的:

  • Jdbc3KeyGenerator:主要基於 java.sql.Statement.getGeneratedKeys 的主鍵返回介面實現的,所以他不需要 processBefore 方法,只需要在獲取到結果後使用 processAfter 攔截,然後用反射將主鍵設置到參數中即可;
  • SelectKeyGenerator:主要是通過 XML 配置或者註解設置 selectKey ,然後單獨發出查詢語句,在返回攔截方法中使用反射設置主鍵,其中兩個攔截方法只能使用其一,在 selectKey.order 屬性中設置 AFTER|BEFORE 來確定;

攔截時機:

processBefore 是在生成 StatementHandler 的時候;

protected BaseStatementHandler(Executor executor, MappedStatement mappedStatement, Object parameterObject, RowBounds rowBounds, ResultHandler resultHandler, BoundSql boundSql) {    ...    if (boundSql == null) { // issue #435, get the key before calculating the statement      generateKeys(parameterObject);      boundSql = mappedStatement.getBoundSql(parameterObject);    }    ...  }    protected void generateKeys(Object parameter) {    KeyGenerator keyGenerator = mappedStatement.getKeyGenerator();    ErrorContext.instance().store();    keyGenerator.processBefore(executor, mappedStatement, null, parameter);    ErrorContext.instance().recall();  }

processAfter 則是在完成插入返回結果之前,但是 PreparedStatementHandler、SimpleStatementHandler、CallableStatementHandler 的程式碼稍微有一點不同,但是位置是不變的,這裡以 PreparedStatementHandler 舉例:

@Override  public int update(Statement statement) throws SQLException {    PreparedStatement ps = (PreparedStatement) statement;    ps.execute();    int rows = ps.getUpdateCount();    Object parameterObject = boundSql.getParameterObject();    KeyGenerator keyGenerator = mappedStatement.getKeyGenerator();    keyGenerator.processAfter(executor, mappedStatement, ps, parameterObject);    return rows;  }

二、Jdbc3KeyGenerator

上面也將了 Jdbc3KeyGenerator 是主要基於 java.sql.Statement.getGeneratedKeys 的主鍵返回介面實現的,但是 Statement 和 PreparedStatement 稍有不同,所以導致了 PreparedStatementHandler、SimpleStatementHandler 的 update 方法稍有不同:

// java.sql.Connection  PreparedStatement prepareStatement(String sql, int autoGeneratedKeys) throws SQLException;  PreparedStatement prepareStatement(String sql, String columnNames[]) throws SQLException;  PreparedStatement prepareStatement(String sql, int columnIndexes[]) throws SQLException;    // java.sql.Statement  boolean execute(String sql, int autoGeneratedKeys) throws SQLException;  boolean execute(String sql, int columnIndexes[]) throws SQLException;  boolean execute(String sql, String columnNames[]) throws SQLException;  // 其中 autoGenerateKeys - Statement.RETURN_GENERATED_KEYS、Statement.NO_GENERATED_KEYS

可以看到 PreparedStatement 是在實例化的時候就指定了,而 Statement 是在執行 sql 的時候才指定但實質是一樣的,這裡就以 PreparedStatement 舉例:

public void testJDBC3() {    try {      String url = "jdbc:mysql://localhost:3306/mybatis?serverTimezone=GMT";      String sql = "INSERT INTO user(username,password,address) VALUES (?,?,?)";      Class.forName("com.mysql.jdbc.Driver");      Connection conn = DriverManager.getConnection(url, "root", "root");      String[] columnNames = {"ids", "name"};      PreparedStatement stmt = conn.prepareStatement(sql, columnNames);      stmt.setString(1, "test");      stmt.setString(2, "123456");      stmt.setString(3, "test");      stmt.executeUpdate();      ResultSet rs = stmt.getGeneratedKeys();      int id = 0;      if (rs.next()) {        id = rs.getInt(1);        System.out.println("----------" + id);      }    } catch (Exception e) {      e.printStackTrace();    }  }

這裡的 User 表以 id 為主鍵,但是程式碼中我傳的 columnNames 都不符合,而結果仍然可以正確的返回主鍵,主要是因為在 mybatis 的驅動中只要 columnNames.length > 1就可以了,所以在具體使用的時候還要注意不同資料庫驅動實現不同所帶來的影響;

上面將了 Statement 和 PreparedStatement 指定返回主鍵的位置不同,在下面就能很清楚的看到:

// org.apache.ibatis.executor.statement.SimpleStatementHandler  public int update(Statement statement) throws SQLException {    String sql = boundSql.getSql();    Object parameterObject = boundSql.getParameterObject();    KeyGenerator keyGenerator = mappedStatement.getKeyGenerator();    int rows;    if (keyGenerator instanceof Jdbc3KeyGenerator) {      statement.execute(sql, Statement.RETURN_GENERATED_KEYS);      rows = statement.getUpdateCount();      keyGenerator.processAfter(executor, mappedStatement, statement, parameterObject);    } else if (keyGenerator instanceof SelectKeyGenerator) {      statement.execute(sql);      rows = statement.getUpdateCount();      keyGenerator.processAfter(executor, mappedStatement, statement, parameterObject);    } else {      //如果沒有keyGenerator,直接調用Statement.execute和Statement.getUpdateCount      statement.execute(sql);      rows = statement.getUpdateCount();    }    return rows;  }    // org.apache.ibatis.executor.statement.PreparedStatementHandler  protected Statement instantiateStatement(Connection connection) throws SQLException {    String sql = boundSql.getSql();    if (mappedStatement.getKeyGenerator() instanceof Jdbc3KeyGenerator) {      String[] keyColumnNames = mappedStatement.getKeyColumns();      if (keyColumnNames == null) {        return connection.prepareStatement(sql, PreparedStatement.RETURN_GENERATED_KEYS);      } else {        return connection.prepareStatement(sql, keyColumnNames);      }    } else if (mappedStatement.getResultSetType() != null) {      return connection.prepareStatement(sql, mappedStatement.getResultSetType().getValue(), ResultSet.CONCUR_READ_ONLY);    } else {      return connection.prepareStatement(sql);    }  }

在完成初始化後,下面來看 Jdbc3KeyGenerator 中最主要的攔截方法:

public void processAfter(Executor executor, MappedStatement ms, Statement stmt, Object parameter) {    List<Object> parameters = new ArrayList<Object>();    parameters.add(parameter);    processBatch(ms, stmt, parameters);  }    public void processBatch(MappedStatement ms, Statement stmt, List<Object> parameters) {    ResultSet rs = null;    try {      //核心是使用JDBC3的Statement.getGeneratedKeys      rs = stmt.getGeneratedKeys();      final Configuration configuration = ms.getConfiguration();      final TypeHandlerRegistry typeHandlerRegistry = configuration.getTypeHandlerRegistry();      final String[] keyProperties = ms.getKeyProperties();      final ResultSetMetaData rsmd = rs.getMetaData();      TypeHandler<?>[] typeHandlers = null;      if (keyProperties != null && rsmd.getColumnCount() >= keyProperties.length) {        for (Object parameter : parameters) {          // there should be one row for each statement (also one for each parameter)          if (!rs.next()) {            break;          }          final MetaObject metaParam = configuration.newMetaObject(parameter);          if (typeHandlers == null) {            //先取得類型處理器            typeHandlers = getTypeHandlers(typeHandlerRegistry, metaParam, keyProperties);          }          //填充鍵值          populateKeys(rs, metaParam, keyProperties, typeHandlers);        }      }    } catch (Exception e) {      ...    }  }

這裡就很清楚了,直接獲取返回的主鍵,然後一次使用反射設置到參數中;

三、SelectKeyGenerator

上面也講了 SelectKeyGenerator 主要是配置 selectKey 使用的,默認 使用 processBefore,但是可以配置 order 屬性(AFTER|BEFORE);

<insert id="insertUser2" parameterType="u" useGeneratedKeys="true" keyProperty="id">    <selectKey keyProperty="id" resultType="long" order="BEFORE">      SELECT if(max(id) is null,1,max(id)+2) as newId FROM user2    </selectKey>    INSERT INTO user2(id,username,password,address) VALUES (#{id},#{userName},#{password},#{address})  </insert>

這裡直接看源碼:

public void processBefore(Executor executor, MappedStatement ms, Statement stmt, Object parameter) {    if (executeBefore) processGeneratedKeys(executor, ms, parameter);  }    public void processAfter(Executor executor, MappedStatement ms, Statement stmt, Object parameter) {    if (!executeBefore) processGeneratedKeys(executor, ms, parameter);  }    private void processGeneratedKeys(Executor executor, MappedStatement ms, Object parameter) {    try {      if (parameter != null && keyStatement != null && keyStatement.getKeyProperties() != null) {        String[] keyProperties = keyStatement.getKeyProperties();        final Configuration configuration = ms.getConfiguration();        final MetaObject metaParam = configuration.newMetaObject(parameter);        if (keyProperties != null) {          // Do not close keyExecutor.          // The transaction will be closed by parent executor.          Executor keyExecutor = configuration.newExecutor(executor.getTransaction(), ExecutorType.SIMPLE);          List<Object> values = keyExecutor.query(keyStatement, parameter, RowBounds.DEFAULT, Executor.NO_RESULT_HANDLER);          if (values.size() == 0) {            throw new ExecutorException("SelectKey returned no data.");          } else if (values.size() > 1) {            throw new ExecutorException("SelectKey returned more than one value.");          } else {            MetaObject metaResult = configuration.newMetaObject(values.get(0));            if (keyProperties.length == 1) {              if (metaResult.hasGetter(keyProperties[0])) {                setValue(metaParam, keyProperties[0], metaResult.getValue(keyProperties[0]));              } else {                // no getter for the property - maybe just a single value object                // so try that                setValue(metaParam, keyProperties[0], values.get(0));              }            } else {              handleMultipleProperties(keyProperties, metaParam, metaResult);            }          }        }      }    } catch (ExecutorException e) {      ...    }  }

這裡程式碼也很簡單,就是用一個新的 Executor 再發一條 SQL,然後反射設置參數即可;