mybatis 原始碼分析（七）KeyGenerator 詳解
阿新 • • 發佈：2019-09-02
一、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，然後反射設定引數即可