This is an automated email from the ASF dual-hosted git repository.

johnthuss pushed a commit to branch genpkbatch
in repository https://gitbox.apache.org/repos/asf/cayenne.git

commit a1019c802b9f1d0df22850aeeca91606d861072f
Author: John Huss <[email protected]>
AuthorDate: Tue Feb 11 15:44:09 2020 -0600

    CAY-2650 Support using generated primary keys along with batch inserts
---
 RELEASE-NOTES.txt                                  |  1 +
 .../cayenne/access/DataDomainFlushObserver.java    | 91 +++++++++++-----------
 .../access/DataDomainLegacyQueryAction.java        |  4 +-
 .../cayenne/access/DataDomainQueryAction.java      |  2 +-
 .../apache/cayenne/access/DataNodeQueryAction.java |  4 +-
 .../apache/cayenne/access/OperationObserver.java   |  2 +-
 .../apache/cayenne/access/flush/FlushObserver.java | 87 +++++++++++----------
 .../apache/cayenne/access/jdbc/BatchAction.java    | 64 +++++++++++++--
 .../access/util/DefaultOperationObserver.java      |  2 +-
 .../access/util/DoNothingOperationObserver.java    |  2 +-
 .../org/apache/cayenne/dba/JdbcPkGenerator.java    |  2 +-
 .../dba/sqlserver/SQLServerProcedureAction.java    |  4 +-
 .../cayenne/access/MockOperationObserver.java      |  2 +-
 13 files changed, 160 insertions(+), 107 deletions(-)

diff --git a/RELEASE-NOTES.txt b/RELEASE-NOTES.txt
index 1f2b31b..fa40183 100644
--- a/RELEASE-NOTES.txt
+++ b/RELEASE-NOTES.txt
@@ -55,6 +55,7 @@ CAY-2610 Align methods in ObjectSelect and SQLSelect
 CAY-2611 Exclude system catalogs and schemas when run dbImport without config
 CAY-2612 Modeler: add lazy-loading to dbImport tab
 CAY-2645 Modeler: DbImport tree highlight improvement
+CAY-2650 Support using generated primary keys along with batch inserts
 
 Bug Fixes:
 
diff --git 
a/cayenne-server/src/main/java/org/apache/cayenne/access/DataDomainFlushObserver.java
 
b/cayenne-server/src/main/java/org/apache/cayenne/access/DataDomainFlushObserver.java
index 0e94332..e99a546 100644
--- 
a/cayenne-server/src/main/java/org/apache/cayenne/access/DataDomainFlushObserver.java
+++ 
b/cayenne-server/src/main/java/org/apache/cayenne/access/DataDomainFlushObserver.java
@@ -67,7 +67,7 @@ class DataDomainFlushObserver implements OperationObserver {
      */
     @Override
     @SuppressWarnings({ "rawtypes", "unchecked" })
-    public void nextGeneratedRows(Query query, ResultIterator keysIterator, 
ObjectId idToUpdate) {
+    public void nextGeneratedRows(Query query, ResultIterator<?> keysIterator, 
List<ObjectId> idsToUpdate) {
 
         // read and close the iterator before doing anything else
         List<DataRow> keys;
@@ -81,53 +81,50 @@ class DataDomainFlushObserver implements OperationObserver {
             throw new CayenneRuntimeException("Generated keys only supported 
for InsertBatchQuery, instead got %s", query);
         }
 
-        if (idToUpdate == null || !idToUpdate.isTemporary()) {
-            // why would this happen?
-            return;
+        if (keys.size() != idsToUpdate.size()) {
+            throw new CayenneRuntimeException("Mismatching number of generated 
PKs: expected %d, instead got %d", idsToUpdate.size(), keys.size());
         }
-
-        if (keys.size() != 1) {
-            throw new CayenneRuntimeException("One and only one PK row is 
expected, instead got %d",  keys.size());
-        }
-
-        DataRow key = keys.get(0);
-
-        // empty key?
-        if (key.size() == 0) {
-            throw new CayenneRuntimeException("Empty key generated.");
-        }
-
-        // determine DbAttribute name...
-
-        // As of now (01/2005) all tested drivers don't provide decent
-        // descriptors of
-        // identity result sets, so a data row will contain garbage labels. 
Also
-        // most
-        // DBs only support one autogenerated key per table... So here we will
-        // have to
-        // infer the key name and currently will only support a single 
column...
-        if (key.size() > 1) {
-            throw new CayenneRuntimeException("Only a single column 
autogenerated PK is supported. "
-                    + "Generated key: %s", key);
-        }
-
-        BatchQuery batch = (BatchQuery) query;
-        for (DbAttribute attribute : 
batch.getDbEntity().getGeneratedAttributes()) {
-
-            // batch can have generated attributes that are not PKs, e.g.
-            // columns with
-            // DB DEFAULT values. Ignore those.
-            if (attribute.isPrimaryKey()) {
-                Object value = key.values().iterator().next();
-
-                // Log the generated PK
-                logger.logGeneratedKey(attribute, value);
-
-                // I guess we should override any existing value,
-                // as generated key is the latest thing that exists in the DB.
-                idToUpdate.getReplacementIdMap().put(attribute.getName(), 
value);
-                break;
-            }
+        
+        for (int i = 0; i < keys.size(); i++) {
+               DataRow key = keys.get(i);
+       
+               // empty key?
+               if (key.size() == 0) {
+                   throw new CayenneRuntimeException("Empty key generated.");
+               }
+       
+               ObjectId idToUpdate = idsToUpdate.get(i);
+               if (idToUpdate == null || !idToUpdate.isTemporary()) {
+                   // why would this happen?
+                   return;
+               }
+
+               BatchQuery batch = (BatchQuery) query;
+               for (DbAttribute attribute : 
batch.getDbEntity().getGeneratedAttributes()) {
+       
+                   // batch can have generated attributes that are not PKs, 
e.g.
+                   // columns with
+                   // DB DEFAULT values. Ignore those.
+                   if (attribute.isPrimaryKey()) {
+                       
+                       Object value = key.get(attribute.getName());
+                       
+                       // As of now (01/2005) many tested drivers don't 
provide decent
+                       // descriptors of
+                       // identity result sets, so a data row may contain 
garbage labels.
+                       if (value == null) {
+                               value = key.values().iterator().next();
+                       }
+                       
+                       // Log the generated PK
+                       logger.logGeneratedKey(attribute, value);
+       
+                       // I guess we should override any existing value,
+                       // as generated key is the latest thing that exists in 
the DB.
+                       
idToUpdate.getReplacementIdMap().put(attribute.getName(), value);
+                       break;
+                   }
+               }
         }
     }
 
diff --git 
a/cayenne-server/src/main/java/org/apache/cayenne/access/DataDomainLegacyQueryAction.java
 
b/cayenne-server/src/main/java/org/apache/cayenne/access/DataDomainLegacyQueryAction.java
index b2ff3c5..31145ef 100644
--- 
a/cayenne-server/src/main/java/org/apache/cayenne/access/DataDomainLegacyQueryAction.java
+++ 
b/cayenne-server/src/main/java/org/apache/cayenne/access/DataDomainLegacyQueryAction.java
@@ -170,8 +170,8 @@ class DataDomainLegacyQueryAction implements QueryRouter, 
OperationObserver {
     }
 
     @Override
-    public void nextGeneratedRows(Query query, ResultIterator keys, ObjectId 
idToUpdate) {
-        callback.nextGeneratedRows(queryForExecutedQuery(query), keys, 
idToUpdate);
+    public void nextGeneratedRows(Query query, ResultIterator<?> keys, 
List<ObjectId> idsToUpdate) {
+        callback.nextGeneratedRows(queryForExecutedQuery(query), keys, 
idsToUpdate);
     }
 
     @Override
diff --git 
a/cayenne-server/src/main/java/org/apache/cayenne/access/DataDomainQueryAction.java
 
b/cayenne-server/src/main/java/org/apache/cayenne/access/DataDomainQueryAction.java
index 80a861f..f95c2df 100644
--- 
a/cayenne-server/src/main/java/org/apache/cayenne/access/DataDomainQueryAction.java
+++ 
b/cayenne-server/src/main/java/org/apache/cayenne/access/DataDomainQueryAction.java
@@ -607,7 +607,7 @@ class DataDomainQueryAction implements QueryRouter, 
OperationObserver {
     }
 
     @Override
-    public void nextGeneratedRows(Query query, ResultIterator<?> keys, 
ObjectId idToUpdate) {
+    public void nextGeneratedRows(Query query, ResultIterator<?> keys, 
List<ObjectId> idsToUpdate) {
         if (keys != null) {
             try {
                 nextRows(query, keys.allRows());
diff --git 
a/cayenne-server/src/main/java/org/apache/cayenne/access/DataNodeQueryAction.java
 
b/cayenne-server/src/main/java/org/apache/cayenne/access/DataNodeQueryAction.java
index 001e869..a7787b7 100644
--- 
a/cayenne-server/src/main/java/org/apache/cayenne/access/DataNodeQueryAction.java
+++ 
b/cayenne-server/src/main/java/org/apache/cayenne/access/DataNodeQueryAction.java
@@ -73,8 +73,8 @@ class DataNodeQueryAction {
             }
             
             @Override
-            public void nextGeneratedRows(Query query, ResultIterator keys, 
ObjectId idToUpdate) {
-                observer.nextGeneratedRows(originalQuery, keys, idToUpdate);
+            public void nextGeneratedRows(Query query, ResultIterator<?> keys, 
List<ObjectId> idsToUpdate) {
+                observer.nextGeneratedRows(originalQuery, keys, idsToUpdate);
             }
 
             @Override
diff --git 
a/cayenne-server/src/main/java/org/apache/cayenne/access/OperationObserver.java 
b/cayenne-server/src/main/java/org/apache/cayenne/access/OperationObserver.java
index 0a52cf1..d24ba90 100644
--- 
a/cayenne-server/src/main/java/org/apache/cayenne/access/OperationObserver.java
+++ 
b/cayenne-server/src/main/java/org/apache/cayenne/access/OperationObserver.java
@@ -63,7 +63,7 @@ public interface OperationObserver extends OperationHints {
      * 
      * @since 4.0
      */
-    void nextGeneratedRows(Query query, ResultIterator<?> keys, ObjectId 
idToUpdate);
+    void nextGeneratedRows(Query query, ResultIterator<?> keys, List<ObjectId> 
idsToUpdate);
 
     /**
      * Callback method invoked on exceptions that happen during an execution 
of a specific
diff --git 
a/cayenne-server/src/main/java/org/apache/cayenne/access/flush/FlushObserver.java
 
b/cayenne-server/src/main/java/org/apache/cayenne/access/flush/FlushObserver.java
index 0968cb0..ca00510 100644
--- 
a/cayenne-server/src/main/java/org/apache/cayenne/access/flush/FlushObserver.java
+++ 
b/cayenne-server/src/main/java/org/apache/cayenne/access/flush/FlushObserver.java
@@ -59,7 +59,7 @@ class FlushObserver implements OperationObserver {
      */
     @Override
     @SuppressWarnings("unchecked")
-    public void nextGeneratedRows(Query query, ResultIterator<?> keysIterator, 
ObjectId idToUpdate) {
+    public void nextGeneratedRows(Query query, ResultIterator<?> keysIterator, 
List<ObjectId> idsToUpdate) {
 
         // read and close the iterator before doing anything else
         List<DataRow> keys;
@@ -73,49 +73,50 @@ class FlushObserver implements OperationObserver {
             throw new CayenneRuntimeException("Generated keys only supported 
for InsertBatchQuery, instead got %s", query);
         }
 
-        if (idToUpdate == null || !idToUpdate.isTemporary()) {
-            // why would this happen?
-            return;
+        if (keys.size() != idsToUpdate.size()) {
+            throw new CayenneRuntimeException("Mismatching number of generated 
PKs: expected %d, instead got %d", idsToUpdate.size(), keys.size());
         }
-
-        if (keys.size() != 1) {
-            throw new CayenneRuntimeException("One and only one PK row is 
expected, instead got %d",  keys.size());
-        }
-
-        DataRow key = keys.get(0);
-
-        // empty key?
-        if (key.size() == 0) {
-            throw new CayenneRuntimeException("Empty key generated.");
-        }
-
-        // determine DbAttribute name...
-
-        // As of now (01/2005) all tested drivers don't provide decent
-        // descriptors of identity result sets, so a data row will contain 
garbage labels.
-        // Also most DBs only support one autogenerated key per table...
-        // So here we will have to infer the key name and currently will only 
support a single column...
-        if (key.size() > 1) {
-            throw new CayenneRuntimeException("Only a single column 
autogenerated PK is supported. "
-                    + "Generated key: %s", key);
-        }
-
-        BatchQuery batch = (BatchQuery) query;
-        for (DbAttribute attribute : 
batch.getDbEntity().getGeneratedAttributes()) {
-
-            // batch can have generated attributes that are not PKs, e.g.
-            // columns with DB DEFAULT values. Ignore those.
-            if (attribute.isPrimaryKey()) {
-                Object value = key.values().iterator().next();
-
-                // Log the generated PK
-                logger.logGeneratedKey(attribute, value);
-
-                // I guess we should override any existing value,
-                // as generated key is the latest thing that exists in the DB.
-                idToUpdate.getReplacementIdMap().put(attribute.getName(), 
value);
-                break;
-            }
+        
+        for (int i = 0; i < keys.size(); i++) {
+               DataRow key = keys.get(i);
+       
+               // empty key?
+               if (key.size() == 0) {
+                   throw new CayenneRuntimeException("Empty key generated.");
+               }
+       
+               ObjectId idToUpdate = idsToUpdate.get(i);
+               if (idToUpdate == null || !idToUpdate.isTemporary()) {
+                   // why would this happen?
+                   return;
+               }
+
+               BatchQuery batch = (BatchQuery) query;
+               for (DbAttribute attribute : 
batch.getDbEntity().getGeneratedAttributes()) {
+       
+                   // batch can have generated attributes that are not PKs, 
e.g.
+                   // columns with
+                   // DB DEFAULT values. Ignore those.
+                   if (attribute.isPrimaryKey()) {
+                       
+                       Object value = key.get(attribute.getName());
+                       
+                       // As of now (01/2005) many tested drivers don't 
provide decent
+                       // descriptors of
+                       // identity result sets, so a data row may contain 
garbage labels.
+                       if (value == null) {
+                               value = key.values().iterator().next();
+                       }
+                       
+                       // Log the generated PK
+                       logger.logGeneratedKey(attribute, value);
+       
+                       // I guess we should override any existing value,
+                       // as generated key is the latest thing that exists in 
the DB.
+                       
idToUpdate.getReplacementIdMap().put(attribute.getName(), value);
+                       break;
+                   }
+               }
         }
     }
 
diff --git 
a/cayenne-server/src/main/java/org/apache/cayenne/access/jdbc/BatchAction.java 
b/cayenne-server/src/main/java/org/apache/cayenne/access/jdbc/BatchAction.java
index b04dd6c..3bcea41 100644
--- 
a/cayenne-server/src/main/java/org/apache/cayenne/access/jdbc/BatchAction.java
+++ 
b/cayenne-server/src/main/java/org/apache/cayenne/access/jdbc/BatchAction.java
@@ -38,10 +38,14 @@ import org.apache.cayenne.query.InsertBatchQuery;
 import java.sql.Connection;
 import java.sql.PreparedStatement;
 import java.sql.ResultSet;
+import java.sql.ResultSetMetaData;
 import java.sql.SQLException;
 import java.sql.Statement;
 import java.util.Collection;
 import java.util.Collections;
+import java.util.List;
+import java.util.function.Consumer;
+import java.util.stream.Collectors;
 
 /**
  * @since 1.2
@@ -84,8 +88,8 @@ public class BatchAction extends BaseSQLAction {
                BatchTranslator translator = createTranslator();
                boolean generatesKeys = hasGeneratedKeys();
 
-               if (runningAsBatch && !generatesKeys) {
-                       runAsBatch(connection, translator, observer);
+               if (runningAsBatch) {
+                       runAsBatch(connection, translator, observer, 
generatesKeys);
                } else {
                        runAsIndividualQueries(connection, translator, 
observer, generatesKeys);
                }
@@ -95,7 +99,7 @@ public class BatchAction extends BaseSQLAction {
                return dataNode.batchTranslator(query, null);
        }
 
-       protected void runAsBatch(Connection con, BatchTranslator translator, 
OperationObserver delegate)
+       protected void runAsBatch(Connection con, BatchTranslator translator, 
OperationObserver delegate, boolean generatesKeys)
                        throws SQLException, Exception {
 
                String sql = translator.getSql();
@@ -109,7 +113,7 @@ public class BatchAction extends BaseSQLAction {
 
                DbAdapter adapter = dataNode.getAdapter();
 
-               try (PreparedStatement statement = con.prepareStatement(sql)) {
+               try (PreparedStatement statement = con.prepareStatement(sql, 
PreparedStatement.RETURN_GENERATED_KEYS)) {
                        for (BatchQueryRow row : query.getRows()) {
 
                                DbAttributeBinding[] bindings = 
translator.updateBindings(row);
@@ -123,6 +127,10 @@ public class BatchAction extends BaseSQLAction {
                        int[] results = statement.executeBatch();
                        delegate.nextBatchCount(query, results);
 
+                       if (generatesKeys) {
+                               processGeneratedKeys(statement, delegate, 
query.getRows());
+                       }
+                       
                        if (isLoggable) {
                                int totalUpdateCount = 0;
                                for (int result : results) {
@@ -263,6 +271,52 @@ public class BatchAction extends BaseSQLAction {
                                Collections.<ObjAttribute, ColumnDescriptor> 
emptyMap());
                ResultIterator iterator = new JDBCResultIterator(null, keysRS, 
rowReader);
 
-               observer.nextGeneratedRows(query, iterator, row.getObjectId());
+               observer.nextGeneratedRows(query, iterator, 
Collections.singletonList(row.getObjectId()));
+       }
+       
+       @SuppressWarnings({ "rawtypes", "unchecked" })
+       protected void processGeneratedKeys(Statement statement, 
OperationObserver observer, List<BatchQueryRow> rows)
+                       throws SQLException {
+
+               ResultSet keysRS = statement.getGeneratedKeys();
+
+               // TODO: andrus, 7/4/2007 - (1) get the type of meaningful PK's 
from
+               // their
+               // ObjAttributes; (2) use a different form of Statement.execute 
-
+               // "execute(String,String[])" to be able to map generated 
column names
+               // (this way
+               // we can support multiple columns.. although need to check how 
well
+               // this works
+               // with most common drivers)
+
+               RowDescriptorBuilder builder = new RowDescriptorBuilder();
+
+               if (this.keyRowDescriptor == null) {
+                       // attempt to figure out the right descriptor from the 
mapping...
+                       Collection<DbAttribute> generated = 
query.getDbEntity().getGeneratedAttributes();
+                       if (generated.size() == 1 && 
keysRS.getMetaData().getColumnCount() == 1) {
+                               DbAttribute key = generated.iterator().next();
+
+                               ColumnDescriptor[] columns = new 
ColumnDescriptor[1];
+
+                               // use column name from result set, but type 
and Java class from
+                               // DB
+                               // attribute
+                               columns[0] = new 
ColumnDescriptor(keysRS.getMetaData(), 1);
+                               columns[0].setJdbcType(key.getType());
+                               
columns[0].setJavaClass(TypesMapping.getJavaBySqlType(key.getType()));
+                               builder.setColumns(columns);
+                       } else {
+                               builder.setResultSet(keysRS);
+                       }
+
+                       this.keyRowDescriptor = 
builder.getDescriptor(dataNode.getAdapter().getExtendedTypes());
+               }
+
+               RowReader<?> rowReader = dataNode.rowReader(keyRowDescriptor, 
query.getMetaData(dataNode.getEntityResolver()),
+                               Collections.<ObjAttribute, ColumnDescriptor> 
emptyMap());
+               ResultIterator iterator = new JDBCResultIterator(null, keysRS, 
rowReader);
+
+               observer.nextGeneratedRows(query, iterator, rows.stream().map(r 
-> r.getObjectId()).collect(Collectors.toList()));
        }
 }
diff --git 
a/cayenne-server/src/main/java/org/apache/cayenne/access/util/DefaultOperationObserver.java
 
b/cayenne-server/src/main/java/org/apache/cayenne/access/util/DefaultOperationObserver.java
index 75b2de5..8add835 100644
--- 
a/cayenne-server/src/main/java/org/apache/cayenne/access/util/DefaultOperationObserver.java
+++ 
b/cayenne-server/src/main/java/org/apache/cayenne/access/util/DefaultOperationObserver.java
@@ -126,7 +126,7 @@ public class DefaultOperationObserver implements 
OperationObserver {
      * 
      * @since 4.0
      */
-    public void nextGeneratedRows(Query query, ResultIterator keys, 
org.apache.cayenne.ObjectId idToUpdate) {
+    public void nextGeneratedRows(Query query, ResultIterator<?> keys, 
List<org.apache.cayenne.ObjectId> idsToUpdate) {
         if (keys != null) {
             keys.close();
         }
diff --git 
a/cayenne-server/src/main/java/org/apache/cayenne/access/util/DoNothingOperationObserver.java
 
b/cayenne-server/src/main/java/org/apache/cayenne/access/util/DoNothingOperationObserver.java
index 0423a7a..bc90d0b 100644
--- 
a/cayenne-server/src/main/java/org/apache/cayenne/access/util/DoNothingOperationObserver.java
+++ 
b/cayenne-server/src/main/java/org/apache/cayenne/access/util/DoNothingOperationObserver.java
@@ -64,7 +64,7 @@ public class DoNothingOperationObserver implements 
OperationObserver {
        }
 
        @Override
-       public void nextGeneratedRows(Query query, ResultIterator<?> keys, 
ObjectId idToUpdate) {
+       public void nextGeneratedRows(Query query, ResultIterator<?> keys, 
List<ObjectId> idsToUpdate) {
                // do
        }
 
diff --git 
a/cayenne-server/src/main/java/org/apache/cayenne/dba/JdbcPkGenerator.java 
b/cayenne-server/src/main/java/org/apache/cayenne/dba/JdbcPkGenerator.java
index 9fa758d..6423370 100644
--- a/cayenne-server/src/main/java/org/apache/cayenne/dba/JdbcPkGenerator.java
+++ b/cayenne-server/src/main/java/org/apache/cayenne/dba/JdbcPkGenerator.java
@@ -364,7 +364,7 @@ public class JdbcPkGenerator implements PkGenerator {
         }
 
         @Override
-        public void nextGeneratedRows(Query query, ResultIterator keys, 
ObjectId idToUpdate) {
+        public void nextGeneratedRows(Query query, ResultIterator<?> keys, 
List<ObjectId> idsToUpdate) {
         }
 
         public void nextRows(Query q, ResultIterator it) {
diff --git 
a/cayenne-server/src/main/java/org/apache/cayenne/dba/sqlserver/SQLServerProcedureAction.java
 
b/cayenne-server/src/main/java/org/apache/cayenne/dba/sqlserver/SQLServerProcedureAction.java
index 1b4355b..2f5022a 100644
--- 
a/cayenne-server/src/main/java/org/apache/cayenne/dba/sqlserver/SQLServerProcedureAction.java
+++ 
b/cayenne-server/src/main/java/org/apache/cayenne/dba/sqlserver/SQLServerProcedureAction.java
@@ -171,8 +171,8 @@ public class SQLServerProcedureAction extends 
ProcedureAction {
                }
 
                @Override
-               public void nextGeneratedRows(Query query, ResultIterator keys, 
ObjectId idToUpdate) {
-                       observer.nextGeneratedRows(query, keys, idToUpdate);
+               public void nextGeneratedRows(Query query, ResultIterator<?> 
keys, List<ObjectId> idsToUpdate) {
+                       observer.nextGeneratedRows(query, keys, idsToUpdate);
                }
 
                @Override
diff --git 
a/cayenne-server/src/test/java/org/apache/cayenne/access/MockOperationObserver.java
 
b/cayenne-server/src/test/java/org/apache/cayenne/access/MockOperationObserver.java
index 202ffb7..bc30278 100644
--- 
a/cayenne-server/src/test/java/org/apache/cayenne/access/MockOperationObserver.java
+++ 
b/cayenne-server/src/test/java/org/apache/cayenne/access/MockOperationObserver.java
@@ -74,7 +74,7 @@ public class MockOperationObserver implements 
OperationObserver {
     }
 
     @Override
-    public void nextGeneratedRows(Query query, ResultIterator<?> keys, 
ObjectId idToUpdate) {
+    public void nextGeneratedRows(Query query, ResultIterator<?> keys, 
List<ObjectId> idsToUpdate) {
     }
 
     public boolean isIteratedResult() {

Reply via email to