Merge pull request #4136 from HubSpot/fix-batch-upsert-noinsert

Fix getGeneratedKeys when no insert during batch upsert
vitessio · Sep 1, 2018 · e55032e · e55032e
2 parents faa1a87 + 8ff0d4b
commit e55032e
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 6 deletions.
diff --git a/java/jdbc/src/main/java/io/vitess/jdbc/VitessStatement.java b/java/jdbc/src/main/java/io/vitess/jdbc/VitessStatement.java
@@ -591,7 +591,7 @@ protected void checkSQLNullOrEmpty(String sql) throws SQLException {
     protected int[] generateBatchUpdateResult(List<CursorWithError> cursorWithErrorList, List<String> batchedArgs)
         throws BatchUpdateException {
         int[] updateCounts = new int[cursorWithErrorList.size()];
-        long[][] generatedKeys = new long[cursorWithErrorList.size()][2];
+        ArrayList<long[]> generatedKeys = new ArrayList<long[]>();
 
         Vtrpc.RPCError rpcError = null;
         String batchCommand = null;
@@ -603,35 +603,38 @@ protected int[] generateBatchUpdateResult(List<CursorWithError> cursorWithErrorL
                 try {
                     long rowsAffected = cursorWithError.getCursor().getRowsAffected();
                     int truncatedUpdateCount;
+                    boolean queryBatchUpsert = false;
                     if (rowsAffected > Integer.MAX_VALUE) {
                         truncatedUpdateCount = Integer.MAX_VALUE;
                     } else {
                         if (sqlIsUpsert(batchCommand)) {
                             // mimicking mysql-connector-j here.
                             // but it would fail for: insert into t1 values ('a'), ('b') on duplicate key update ts = now();
                             truncatedUpdateCount = 1;
+                            queryBatchUpsert = true;
                         } else {
                             truncatedUpdateCount = (int) rowsAffected;
                         }
                     }
                     updateCounts[i] = truncatedUpdateCount;
-                    if (this.retrieveGeneratedKeys) {
-                        generatedKeys[i] = new long[]{cursorWithError.getCursor().getInsertId(), truncatedUpdateCount};
+                    long insertId = cursorWithError.getCursor().getInsertId();
+                    if (this.retrieveGeneratedKeys && (!queryBatchUpsert || insertId > 0)) {
+                        generatedKeys.add(new long[]{insertId, truncatedUpdateCount});
                     }
                 } catch (SQLException ex) {
                         /* This case should not happen as API has returned cursor and not error.
                          * Handling by Statement.SUCCESS_NO_INFO
                          */
                     updateCounts[i] = Statement.SUCCESS_NO_INFO;
                     if (this.retrieveGeneratedKeys) {
-                        generatedKeys[i] = new long[]{Statement.SUCCESS_NO_INFO, Statement.SUCCESS_NO_INFO};
+                        generatedKeys.add(new long[]{Statement.SUCCESS_NO_INFO, Statement.SUCCESS_NO_INFO});
                     }
                 }
             } else {
                 rpcError = cursorWithError.getError();
                 updateCounts[i] = Statement.EXECUTE_FAILED;
                 if (this.retrieveGeneratedKeys) {
-                    generatedKeys[i] = new long[]{Statement.EXECUTE_FAILED, Statement.EXECUTE_FAILED};
+                    generatedKeys.add(new long[]{Statement.EXECUTE_FAILED, Statement.EXECUTE_FAILED});
                 }
             }
         }
@@ -642,7 +645,7 @@ protected int[] generateBatchUpdateResult(List<CursorWithError> cursorWithErrorL
             throw new BatchUpdateException(rpcError.toString(), sqlState, errno, updateCounts);
         }
         if (this.retrieveGeneratedKeys) {
-            this.batchGeneratedKeys = generatedKeys;
+            this.batchGeneratedKeys = generatedKeys.toArray(new long[generatedKeys.size()][2]);
         }
         return updateCounts;
     }

diff --git a/java/jdbc/src/test/java/io/vitess/jdbc/VitessStatementTest.java b/java/jdbc/src/test/java/io/vitess/jdbc/VitessStatementTest.java
@@ -742,5 +742,15 @@ private void testExecute(int fetchSize, boolean simpleExecute, boolean shouldRun
             Assert.assertEquals(i, 0); // we should only have one
             i++;
         }
+
+        VitessStatement noUpdate = new VitessStatement(mockConn);
+        PowerMockito.when(mockCursor.getInsertId()).thenReturn(0L);
+        PowerMockito.when(mockCursor.getRowsAffected()).thenReturn(1L);
+
+        noUpdate.addBatch(sqlUpsert);
+        noUpdate.executeBatch();
+
+        ResultSet empty = noUpdate.getGeneratedKeys();
+        Assert.assertFalse(empty.next());
     }
 }