Support multiple key properties in @SelectKey

Author: jeffgbutler

src/main/java/org/apache/ibatis/annotations/SelectKey.java

@@ -29,6 +29,8 @@
 
   String keyProperty();
 
+  String keyColumn() default "";
+
   boolean before();
 
   Class<?> resultType();

src/main/java/org/apache/ibatis/builder/annotation/MapperAnnotationBuilder.java

@@ -564,6 +564,7 @@ private KeyGenerator handleSelectKeyAnnotation(SelectKey selectKeyAnnotation, St
     Class<?> resultTypeClass = selectKeyAnnotation.resultType();
     StatementType statementType = selectKeyAnnotation.statementType();
     String keyProperty = selectKeyAnnotation.keyProperty();
+    String keyColumn = selectKeyAnnotation.keyColumn();
     boolean executeBefore = selectKeyAnnotation.before();
 
     // defaults
@@ -581,7 +582,7 @@ private KeyGenerator handleSelectKeyAnnotation(SelectKey selectKeyAnnotation, St
 
     assistant.addMappedStatement(id, sqlSource, statementType, sqlCommandType, fetchSize, timeout, parameterMap, parameterTypeClass, resultMap, resultTypeClass, resultSetTypeEnum,
         flushCache, useCache, false,
-        keyGenerator, keyProperty, null, null, languageDriver, null);
+        keyGenerator, keyProperty, keyColumn, null, languageDriver, null);
 
     id = assistant.applyCurrentNamespace(id, false);
 

src/main/java/org/apache/ibatis/executor/keygen/SelectKeyGenerator.java

@@ -52,10 +52,10 @@ public void processAfter(Executor executor, MappedStatement ms, Statement stmt,
   private void processGeneratedKeys(Executor executor, MappedStatement ms, Object parameter) {
     try {
       if (parameter != null && keyStatement != null && keyStatement.getKeyProperties() != null) {
-        String keyProperty = keyStatement.getKeyProperties()[0]; // just one key property is supported
+        String[] keyProperties = keyStatement.getKeyProperties();
         final Configuration configuration = ms.getConfiguration();
         final MetaObject metaParam = configuration.newMetaObject(parameter);
-        if (keyProperty != null && metaParam.hasSetter(keyProperty)) {
+        if (keyProperties != null) {
           // Do not close keyExecutor.
           // The transaction will be closed by parent executor.
           Executor keyExecutor = configuration.newExecutor(executor.getTransaction(), ExecutorType.SIMPLE);
@@ -65,10 +65,13 @@ private void processGeneratedKeys(Executor executor, MappedStatement ms, Object
           } else if (values.size() > 1) {
             throw new ExecutorException("SelectKey returned more than one value.");
           } else {
-            metaParam.setValue(keyProperty, values.get(0));
+            if (keyProperties.length == 1) {
+              setValue(metaParam, keyProperties[0], values.get(0));
+            } else {
+              MetaObject metaResult = configuration.newMetaObject(values.get(0));
+              handleMultipleProperties(keyProperties, metaParam, metaResult);
+            }
           }
-        } else {
-          throw new ExecutorException("No setter found for the keyProperty '" + keyProperty + "' in " + metaParam.getOriginalObject().getClass().getName() + ".");
         }
       }
     } catch (ExecutorException e) {
@@ -78,4 +81,30 @@ private void processGeneratedKeys(Executor executor, MappedStatement ms, Object
     }
   }
 
+  private void handleMultipleProperties(String[] keyProperties,
+      MetaObject metaParam, MetaObject metaResult) {
+    String[] keyColumns = keyStatement.getKeyColumns();
+      
+    if (keyColumns == null || keyColumns.length == 0) {
+      // no key columns specified, just use the property names
+      for (int i = 0; i < keyProperties.length; i++) {
+        setValue(metaParam, keyProperties[i], metaResult.getValue(keyProperties[i]));
+      }
+    } else {
+      if (keyColumns.length != keyProperties.length) {
+        throw new ExecutorException("If SelectKey has key columns, the number must match the number of key properties.");
+      }
+      for (int i = 0; i < keyProperties.length; i++) {
+        setValue(metaParam, keyProperties[i], metaResult.getValue(keyColumns[i]));
+      }
+    }
+  }
+
+  private void setValue(MetaObject metaParam, String property, Object value) {
+    if (metaParam.hasSetter(property)) {
+      metaParam.setValue(property, value);
+    } else {
+      throw new ExecutorException("No setter found for the keyProperty '" + property + "' in " + metaParam.getOriginalObject().getClass().getName() + ".");
+    }
+  }
 }

src/test/java/org/apache/ibatis/submitted/selectkey/AnnotatedMapper.java

@@ -15,6 +15,8 @@
  */
 package org.apache.ibatis.submitted.selectkey;
 
+import java.util.Map;
+
 import org.apache.ibatis.annotations.Insert;
 import org.apache.ibatis.annotations.InsertProvider;
 import org.apache.ibatis.annotations.Options;
@@ -30,6 +32,14 @@ public interface AnnotatedMapper {
     @Options(useGeneratedKeys=true, keyProperty="nameId,generatedName", keyColumn="ID,NAME_FRED")
     int insertTable2WithOptions(Name name);
 
+    @Insert("insert into table2 (name) values(#{name})")
+    @SelectKey(statement="select id, name_fred from table2 where id = identity()", keyProperty="nameId,generatedName", keyColumn="ID,NAME_FRED", before=false, resultType=Map.class)
+    int insertTable2WithSelectKeyWithKeyMap(Name name);
+
+    @Insert("insert into table2 (name) values(#{name})")
+    @SelectKey(statement="select id as nameId, name_fred as generatedName from table2 where id = identity()", keyProperty="nameId,generatedName", before=false, resultType=Name.class)
+    int insertTable2WithSelectKeyWithKeyObject(Name name);
+
     @Insert("insert into table3 (id, name) values(#{nameId}, #{name})")
     @SelectKey(statement="call next value for TestSequence", keyProperty="nameId", before=true, resultType=int.class)
     int insertTable3(Name name);

src/test/java/org/apache/ibatis/submitted/selectkey/SelectKeyTest.java

@@ -169,6 +169,40 @@ public void testAnnotatedInsertTable2WithOptions() {
       }
   }
 
+  @Test
+  public void testAnnotatedInsertTable2WithSelectKeyWithKeyMap() {
+      SqlSession sqlSession = sqlSessionFactory.openSession();
+
+      try {
+        Name name = new Name();
+        name.setName("barney");
+        AnnotatedMapper mapper = sqlSession.getMapper(AnnotatedMapper.class);
+        int rows = mapper.insertTable2WithSelectKeyWithKeyMap(name);
+        assertEquals(1, rows);
+        assertEquals(22, name.getNameId());
+        assertEquals("barney_fred", name.getGeneratedName());
+      } finally {
+        sqlSession.close();
+      }
+  }
+
+  @Test
+  public void testAnnotatedInsertTable2WithSelectKeyWithKeyObject() {
+      SqlSession sqlSession = sqlSessionFactory.openSession();
+
+      try {
+        Name name = new Name();
+        name.setName("barney");
+        AnnotatedMapper mapper = sqlSession.getMapper(AnnotatedMapper.class);
+        int rows = mapper.insertTable2WithSelectKeyWithKeyObject(name);
+        assertEquals(1, rows);
+        assertEquals(22, name.getNameId());
+        assertEquals("barney_fred", name.getGeneratedName());
+      } finally {
+        sqlSession.close();
+      }
+  }
+
   @Test
   public void testAnnotatedInsertTable3() {
       SqlSession sqlSession = sqlSessionFactory.openSession();