From 7a12ad4b912cfe2d16e045249a4a0106cecd3abb Mon Sep 17 00:00:00 2001
From: Dmitrii Barnukov <29237913+Nexus6v2@users.noreply.github.com>
Date: Mon, 29 Dec 2025 12:35:47 +0100
Subject: [PATCH] dbeaver/pro#7693 added key attributes row id check (#4019)

* dbeaver/pro#7693 added key attributes row id check

* dbeaver/pro#7693 added fix for single row tables

---------

Co-authored-by: Daria Marutkina <125263541+dariamarutkina@users.noreply.github.com>
---
 .../io/cloudbeaver/service/sql/WebSQLProcessor.java | 13 +++++++++----
 1 file changed, 9 insertions(+), 4 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 868f1b14e0..8098ef2fe0 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
@@ -951,7 +951,7 @@ public class WebSQLProcessor implements WebSessionProvider {
         DBCExecutionContext executionContext = getExecutionContext(dataContainer);
         try (DBCSession session = executionContext.openSession(monitor, DBCExecutionPurpose.USER, "Generate data update batches")) {
             DBDDataFilter dataFilter = new DBDDataFilter();
-            addKeyAttributes(resultsInfo, row, dataContainer, session, dataFilter);
+            addKeyAttributes(monitor, resultsInfo, row, dataContainer, session, dataFilter);
             WebExecutionSource executionSource = new WebExecutionSource(dataContainer, executionContext, this);
             dataContainer.readData(
                 executionSource, session, dataReceiver, dataFilter,
@@ -960,16 +960,21 @@ public class WebSQLProcessor implements WebSessionProvider {
     }
 
     private void addKeyAttributes(
+        @NotNull DBRProgressMonitor monitor,
         @NotNull WebSQLResultsInfo resultsInfo,
         @NotNull WebSQLResultsRow row,
         @NotNull DBSDataContainer dataContainer,
         @NotNull DBCSession session,
         @NotNull DBDDataFilter dataFilter
     ) throws DBException {
-        DBDRowIdentifier rowIdentifier = resultsInfo.getDefaultRowIdentifier();
-        if (rowIdentifier == null || rowIdentifier.isIncomplete()) {
-            return;
+        if (resultsInfo.isSingleRow()) {
+            long rowCount = DBUtils.readRowCount(monitor, session.getExecutionContext(), dataContainer, null, this);
+            if (rowCount == 1) {
+                return;
+            }
         }
+        DBDRowIdentifier rowIdentifier = resultsInfo.getDefaultRowIdentifier();
+        checkRowIdentifier(resultsInfo, rowIdentifier);
         DBDAttributeBinding[] keyAttributes = rowIdentifier.getAttributes().toArray(new DBDAttributeBinding[0]);
         Object[] rowValues = new Object[keyAttributes.length];
         List<DBDAttributeConstraint> constraints = new ArrayList<>();