可以实现批量插入回写功能的KeyGenerator

abel533 · abel533 · commit 0187faec9f91 · 2015-06-04T00:37:48.000+08:00
diff --git a/src/main/java/com/github/abel533/mapperhelper/MultipleJdbc3KeyGenerator.java b/src/main/java/com/github/abel533/mapperhelper/MultipleJdbc3KeyGenerator.java
@@ -0,0 +1,128 @@
+/*
+ * The MIT License (MIT)
+ *
+ * Copyright (c) 2014 abel533@gmail.com
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+ * THE SOFTWARE.
+ */
+
+package com.github.abel533.mapperhelper;
+
+import org.apache.ibatis.executor.Executor;
+import org.apache.ibatis.executor.ExecutorException;
+import org.apache.ibatis.executor.keygen.Jdbc3KeyGenerator;
+import org.apache.ibatis.mapping.MappedStatement;
+import org.apache.ibatis.reflection.MetaObject;
+import org.apache.ibatis.session.Configuration;
+import org.apache.ibatis.type.TypeHandler;
+import org.apache.ibatis.type.TypeHandlerRegistry;
+
+import java.sql.ResultSet;
+import java.sql.ResultSetMetaData;
+import java.sql.SQLException;
+import java.sql.Statement;
+import java.util.*;
+
+/**
+ * 实现批量插入ID回写
+ */
+public class MultipleJdbc3KeyGenerator extends Jdbc3KeyGenerator {
+    @Override
+    public void processAfter(Executor executor, MappedStatement ms, Statement stmt, Object parameter) {
+        processBatch(ms, stmt, getParameters(parameter));
+    }
+
+    public void processBatch(MappedStatement ms, Statement stmt, Collection<Object> parameters) {
+        ResultSet rs = null;
+        try {
+            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) {
+            throw new ExecutorException("Error getting generated key or setting result to parameter object. Cause: " + e, e);
+        } finally {
+            if (rs != null) {
+                try {
+                    rs.close();
+                } catch (Exception e) {
+                    // ignore
+                }
+            }
+        }
+    }
+
+    private Collection<Object> getParameters(Object parameter) {
+        Collection<Object> parameters = null;
+        if (parameter instanceof Collection) {
+            parameters = (Collection) parameter;
+        } else if (parameter instanceof Map) {
+            Map parameterMap = (Map) parameter;
+            if (parameterMap.containsKey("collection")) {
+                parameters = (Collection) parameterMap.get("collection");
+            } else if (parameterMap.containsKey("list")) {
+                parameters = (List) parameterMap.get("list");
+            } else if (parameterMap.containsKey("array")) {
+                parameters = Arrays.asList((Object[]) parameterMap.get("array"));
+            }
+        }
+        if (parameters == null) {
+            parameters = new ArrayList<Object>();
+            parameters.add(parameter);
+        }
+        return parameters;
+    }
+
+    private TypeHandler<?>[] getTypeHandlers(TypeHandlerRegistry typeHandlerRegistry, MetaObject metaParam, String[] keyProperties) {
+        TypeHandler<?>[] typeHandlers = new TypeHandler<?>[keyProperties.length];
+        for (int i = 0; i < keyProperties.length; i++) {
+            if (metaParam.hasSetter(keyProperties[i])) {
+                Class<?> keyPropertyType = metaParam.getSetterType(keyProperties[i]);
+                TypeHandler<?> th = typeHandlerRegistry.getTypeHandler(keyPropertyType);
+                typeHandlers[i] = th;
+            }
+        }
+        return typeHandlers;
+    }
+
+    private void populateKeys(ResultSet rs, MetaObject metaParam, String[] keyProperties, TypeHandler<?>[] typeHandlers) throws SQLException {
+        for (int i = 0; i < keyProperties.length; i++) {
+            TypeHandler<?> th = typeHandlers[i];
+            if (th != null) {
+                Object value = th.getResult(rs, i + 1);
+                metaParam.setValue(keyProperties[i], value);
+            }
+        }
+    }
+}
