From b4c20dbf198817a06ef3c281e3b7935ee535ed86 Mon Sep 17 00:00:00 2001
From: Ainur <ainur.iagudin@dbeaver.com>
Date: Mon, 4 Dec 2023 15:05:27 +0100
Subject: [PATCH 1/2] CB-4328 get updated rows info only using key identifier
 attributes

---
 .../src/io/cloudbeaver/service/sql/WebSQLProcessor.java    | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

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 eef4bc19f6..8179af8a0e 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
@@ -58,6 +58,7 @@
 import java.nio.file.Path;
 import java.util.*;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.stream.Collectors;
 
 /**
  * Web SQL processor.
@@ -423,7 +424,11 @@ private void getUpdatedRowsInfo(
                 }
                 List<DBDAttributeConstraint> constraints = new ArrayList<>();
                 boolean hasKey = true;
-                for (DBDAttributeBinding attr : resultsInfo.getAttributes()) {
+                // get attributes only from row identifiers
+                Set<DBDAttributeBinding> idAttributes = resultsInfo.getRowIdentifiers().stream()
+                    .flatMap(r -> r.getAttributes().stream())
+                    .collect(Collectors.toSet());
+                for (DBDAttributeBinding attr : idAttributes) {
                     if (attr.getRowIdentifier() == null) {
                         continue;
                     }

From 8ca35b74d15f063bae3dd39fc7562647f8191051 Mon Sep 17 00:00:00 2001
From: Ainur <ainur.iagudin@dbeaver.com>
Date: Tue, 5 Dec 2023 11:46:09 +0100
Subject: [PATCH 2/2] CB-4328 do not refresh row for nosql db

---
 .../service/sql/WebSQLResultsInfo.java        | 12 ++++++++---
 .../service/sql/WebSQLProcessor.java          | 20 +++++++++++++++++++
 2 files changed, 29 insertions(+), 3 deletions(-)

diff --git a/server/bundles/io.cloudbeaver.model/src/io/cloudbeaver/service/sql/WebSQLResultsInfo.java b/server/bundles/io.cloudbeaver.model/src/io/cloudbeaver/service/sql/WebSQLResultsInfo.java
index 3fd68da3c2..77a3aeae1b 100644
--- a/server/bundles/io.cloudbeaver.model/src/io/cloudbeaver/service/sql/WebSQLResultsInfo.java
+++ b/server/bundles/io.cloudbeaver.model/src/io/cloudbeaver/service/sql/WebSQLResultsInfo.java
@@ -22,9 +22,7 @@
 import org.jkiss.dbeaver.model.DBUtils;
 import org.jkiss.dbeaver.model.data.DBDAttributeBinding;
 import org.jkiss.dbeaver.model.data.DBDRowIdentifier;
-import org.jkiss.dbeaver.model.struct.DBSAttributeBase;
-import org.jkiss.dbeaver.model.struct.DBSDataContainer;
-import org.jkiss.dbeaver.model.struct.DBSTypedObject;
+import org.jkiss.dbeaver.model.struct.*;
 
 import java.util.HashSet;
 import java.util.List;
@@ -126,4 +124,12 @@ public DBSTypedObject getAttributeByPosition(int pos) {
         return null;
     }
 
+    public boolean canRefreshResults() {
+        DBSEntity entity = getDefaultRowIdentifier().getEntity();
+        // FIXME: do not refresh documents for now. Can be solved by extracting document ID attributes
+        // FIXME: but it will require to provide dynamic document metadata.
+        return entity == null || entity.getDataSource() == null ||
+            (!(entity instanceof DBSDocumentContainer) && !entity.getDataSource().getInfo().isDynamicMetadata());
+    }
+
 }
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 8179af8a0e..c58d93e43a 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
@@ -418,10 +418,15 @@ private void getUpdatedRowsInfo(
             DBCExecutionPurpose.UTIL,
             "Refresh row(s) after insert/update")
         ) {
+            boolean canRefreshResults = resultsInfo.canRefreshResults();
             for (Object[] row : newResultSetRows) {
                 if (row.length == 0) {
                     continue;
                 }
+                if (!canRefreshResults) {
+                    makeWebCellRow(resultsInfo, row, dataFormat);
+                    continue;
+                }
                 List<DBDAttributeConstraint> constraints = new ArrayList<>();
                 boolean hasKey = true;
                 // get attributes only from row identifiers
@@ -444,6 +449,7 @@ private void getUpdatedRowsInfo(
                 }
                 if (!hasKey) {
                     // No key value for this row
+                    makeWebCellRow(resultsInfo, row, dataFormat);
                     continue;
                 }
                 DBDDataFilter filter = new DBDDataFilter(constraints);
@@ -462,6 +468,20 @@ private void getUpdatedRowsInfo(
         }
     }
 
+    private void makeWebCellRow(
+        @NotNull WebSQLResultsInfo resultsInfo,
+        @NotNull Object[] row,
+        @Nullable WebDataFormat dataFormat
+    ) throws DBCException {
+        for (int i = 0; i < row.length; i++) {
+            row[i] = WebSQLUtils.makeWebCellValue(
+                webSession,
+                resultsInfo.getAttributeByPosition(i),
+                row[i],
+                dataFormat);
+        }
+    }
+
     public String generateResultsDataUpdateScript(
         @NotNull DBRProgressMonitor monitor,
         @NotNull WebSQLContextInfo contextInfo,