From b8420a737784b2024de449ecfa5bbf79145105f7 Mon Sep 17 00:00:00 2001
From: serge-rider <serge@jkiss.org>
Date: Sun, 28 Jun 2020 19:59:35 +0300
Subject: [PATCH] CB-153 Batch data update fixes

---
 .../service/sql/WebSQLProcessor.java          | 29 ++++++++++++-------
 .../service/sql/WebSQLQueryResults.java       |  6 ++--
 .../service/sql/WebServiceBindingSQL.java     |  2 +-
 3 files changed, 22 insertions(+), 15 deletions(-)

diff --git a/server/bundles/io.cloudbeaver.server/src/io/cloudbeaver/service/sql/WebSQLProcessor.java b/server/bundles/io.cloudbeaver.server/src/io/cloudbeaver/service/sql/WebSQLProcessor.java
index 3300ac9d2e..42ef34dee1 100644
--- a/server/bundles/io.cloudbeaver.server/src/io/cloudbeaver/service/sql/WebSQLProcessor.java
+++ b/server/bundles/io.cloudbeaver.server/src/io/cloudbeaver/service/sql/WebSQLProcessor.java
@@ -279,7 +279,7 @@ public class WebSQLProcessor {
                         updatedResultSet.setRows(new Object[][]{finalRow});
 
                         WebSQLQueryResults updateResults = new WebSQLQueryResults();
-                        updateResults.setUpdateRowCount((int) statistics.getRowsUpdated());
+                        updateResults.setUpdateRowCount(statistics.getRowsUpdated());
                         updateResults.setResultSet(updatedResultSet);
                         result.setDuration(statistics.getExecuteTime());
                         result.setResults(new WebSQLQueryResults[]{updateResults});
@@ -329,6 +329,13 @@ public class WebSQLProcessor {
                         DBDAttributeBinding[] keyAttributes = rowIdentifier.getAttributes().toArray(new DBDAttributeBinding[0]);
 
                         if (!CommonUtils.isEmpty(updatedRows)) {
+                            WebSQLQueryResultSet updatedResultSet = new WebSQLQueryResultSet();
+                            updatedResultSet.setResultsInfo(resultsInfo);
+                            updatedResultSet.setColumns(resultsInfo.getAttributes());
+                            long totalUpdateCount = 0;
+
+                            List <Object[]> resultRows = new ArrayList<>();
+
                             for (WebSQLResultsRow row : updatedRows) {
                                 Map<String, Object> updateValues = row.getUpdateValues();
                                 if (CommonUtils.isEmpty(row.getData()) || CommonUtils.isEmpty(updateValues)) {
@@ -362,20 +369,20 @@ public class WebSQLProcessor {
                                 updateBatch.add(rowValues);
                                 DBCStatistics statistics = updateBatch.execute(session);
 
-                                WebSQLQueryResultSet updatedResultSet = new WebSQLQueryResultSet();
-                                updatedResultSet.setResultsInfo(resultsInfo);
-                                updatedResultSet.setColumns(resultsInfo.getAttributes());
-                                updatedResultSet.setRows(new Object[][]{finalRow});
-
-                                WebSQLQueryResults updateResults = new WebSQLQueryResults();
-                                updateResults.setUpdateRowCount((int) statistics.getRowsUpdated());
-                                updateResults.setResultSet(updatedResultSet);
-                                queryResults.add(updateResults);
+                                totalUpdateCount += statistics.getRowsUpdated();
                                 result.setDuration(result.getDuration() + statistics.getExecuteTime());
+                                resultRows.add(finalRow);
                             }
+
+                            WebSQLQueryResults updateResults = new WebSQLQueryResults();
+                            updateResults.setUpdateRowCount(totalUpdateCount);
+                            updateResults.setResultSet(updatedResultSet);
+                            updatedResultSet.setRows(resultRows.toArray(new Object[0][]));
+
+                            queryResults.add(updateResults);
                         }
 
-                        if (!CommonUtils.isEmpty(updatedRows)) {
+                        if (!CommonUtils.isEmpty(addedRows)) {
                             throw new DBCException("New row add is not supported");
                         }
 
diff --git a/server/bundles/io.cloudbeaver.server/src/io/cloudbeaver/service/sql/WebSQLQueryResults.java b/server/bundles/io.cloudbeaver.server/src/io/cloudbeaver/service/sql/WebSQLQueryResults.java
index 9bcc9c0a9f..d02f513330 100644
--- a/server/bundles/io.cloudbeaver.server/src/io/cloudbeaver/service/sql/WebSQLQueryResults.java
+++ b/server/bundles/io.cloudbeaver.server/src/io/cloudbeaver/service/sql/WebSQLQueryResults.java
@@ -23,15 +23,15 @@ import org.jkiss.dbeaver.model.meta.Property;
  */
 public class WebSQLQueryResults {
 
-    private Integer updateRowCount;
+    private Long updateRowCount;
     private WebSQLQueryResultSet resultSet;
 
     @Property
-    public Integer getUpdateRowCount() {
+    public Long getUpdateRowCount() {
         return updateRowCount;
     }
 
-    public void setUpdateRowCount(int updateRowCount) {
+    public void setUpdateRowCount(long updateRowCount) {
         this.updateRowCount = updateRowCount;
     }
 
diff --git a/server/bundles/io.cloudbeaver.server/src/io/cloudbeaver/service/sql/WebServiceBindingSQL.java b/server/bundles/io.cloudbeaver.server/src/io/cloudbeaver/service/sql/WebServiceBindingSQL.java
index 4874e35b65..94a9437019 100644
--- a/server/bundles/io.cloudbeaver.server/src/io/cloudbeaver/service/sql/WebServiceBindingSQL.java
+++ b/server/bundles/io.cloudbeaver.server/src/io/cloudbeaver/service/sql/WebServiceBindingSQL.java
@@ -167,7 +167,7 @@ public class WebServiceBindingSQL extends WebServiceBindingBase<DBWServiceSQL> {
     }
 
     private static List<WebSQLResultsRow> getResultsRow(DataFetchingEnvironment env, String param) {
-        List<Map<String, Object>> mapList = env.getArgument("param");
+        List<Map<String, Object>> mapList = env.getArgument(param);
         if (CommonUtils.isEmpty(mapList)) {
             return null;
         }