#1007 Refresh RSV row after insert

plugins/org.jkiss.dbeaver.core/src/org/jkiss/dbeaver/ui/controls/resultset/ResultSetPersister.java

@@ -32,9 +32,7 @@ import org.jkiss.dbeaver.model.data.*;
 import org.jkiss.dbeaver.model.edit.DBEPersistAction;
 import org.jkiss.dbeaver.model.exec.*;
 import org.jkiss.dbeaver.model.impl.AbstractExecutionSource;
-import org.jkiss.dbeaver.model.navigator.DBNNode;
 import org.jkiss.dbeaver.model.runtime.DBRProgressMonitor;
-import org.jkiss.dbeaver.model.runtime.load.DatabaseLoadService;
 import org.jkiss.dbeaver.model.struct.DBSDataContainer;
 import org.jkiss.dbeaver.model.struct.DBSDataManipulator;
 import org.jkiss.dbeaver.model.struct.DBSEntity;
@@ -46,7 +44,6 @@ import org.jkiss.dbeaver.ui.UIUtils;
 import org.jkiss.dbeaver.utils.GeneralUtils;
 import org.jkiss.utils.CommonUtils;
 
-import java.lang.reflect.InvocationTargetException;
 import java.util.*;
 
 /**
@@ -133,7 +130,7 @@ class ResultSetPersister {
         return execute(monitor, generateScript, listener);
     }
 
-    public boolean refreshInsertedRows(@Nullable DBRProgressMonitor monitor) {
+    public boolean refreshInsertedRows(@Nullable DBRProgressMonitor monitor) throws DBCException {
         if (!viewer.getModel().isSingleSource()) {
             return false;
         }
@@ -145,6 +142,20 @@ class ResultSetPersister {
             return false;
         }
 
+        DBCExecutionContext executionContext = viewer.getContainer().getExecutionContext();
+        if (executionContext == null) {
+            throw new DBCException("No execution context");
+        }
+        RowRefreshJob job = new RowRefreshJob(executionContext, viewer.getDataContainer(), rowIdentifier, addedRows);
+        if (monitor == null) {
+            job.schedule();
+            return true;
+        } else {
+            IStatus result = job.run(monitor);
+            if (!result.isOK()) {
+                return false;
+            }
+        }
         return true;
     }
 
@@ -716,6 +727,11 @@ class ResultSetPersister {
     }
 
     class RowDataReceiver implements DBDDataReceiver {
+        private final DBDAttributeBinding[] curAttributes;
+        private Object[] rowValues;
+        public RowDataReceiver(DBDAttributeBinding[] curAttributes) {
+            this.curAttributes = curAttributes;
+        }
 
         @Override
         public void fetchStart(DBCSession session, DBCResultSet resultSet, long offset, long maxRows)
@@ -729,39 +745,27 @@ class ResultSetPersister {
             throws DBCException
         {
             DBCResultSetMetaData rsMeta = resultSet.getMeta();
+            // Compare attributes with existing model attributes
             List<DBCAttributeMetaData> attributes = rsMeta.getAttributes();
-
-/*
-            for (int i = 0; i < attributes.size(); i++) {
-                DBCAttributeMetaData attr = attributes.get(i);
-                DBDValueHandler valueHandler = DBUtils.findValueHandler(session, attr);
-                Object keyValue = valueHandler.fetchValueObject(session, resultSet, attr, i);
-                boolean updated = false;
-                if (!CommonUtils.isEmpty(attr.getName())) {
-                    DBDAttributeBinding binding = model.getAttributeBinding(statement.entity, attr.getName());
-                    if (binding != null) {
-                        // Got it. Just update column oldValue
-                        statement.updatedCells.put(binding.getOrdinalPosition(), keyValue);
-                        //curRows.get(statement.row.row)[colIndex] = keyValue;
-                        updated = true;
-                    }
-                }
-                if (!updated) {
-                    // Key not found
-                    // Try to find and update auto-increment column
-                    for (int k = 0; k < columns.length; k++) {
-                        DBDAttributeBinding column = columns[k];
-                        if (column.isAutoGenerated()) {
-                            // Got it
-                            statement.updatedCells.put(k, keyValue);
-                            //curRows.get(statement.row.row)[k] = keyValue;
-                            updated = true;
-                            break;
-                        }
-                    }
+            if (attributes.size() != curAttributes.length) {
+                log.debug("Wrong meta attributes count - can't refresh");
+                return;
+            }
+            for (int i = 0; i < curAttributes.length; i++) {
+                if (!ResultSetUtils.equalAttributes(curAttributes[i].getMetaAttribute(), attributes.get(i))) {
+                    log.debug("Attribute '" + curAttributes[i].getMetaAttribute().getName() + "' doesn't match '" + attributes.get(i).getName() + "'");
+                    return;
                 }
             }
-*/
+
+            rowValues = new Object[curAttributes.length];
+            for (int i = 0; i < curAttributes.length; i++) {
+                final DBDAttributeBinding attr = curAttributes[i];
+                DBDValueHandler valueHandler = attr.getValueHandler();
+                Object attrValue = valueHandler.fetchValueObject(session, resultSet, attr, i);
+                rowValues[i] = attrValue;
+            }
+
         }
 
         @Override
@@ -777,37 +781,69 @@ class ResultSetPersister {
         }
     }
 
-    private class RowRefreshService extends DatabaseLoadService<Collection<DBNNode>> {
+    private class RowRefreshJob extends DataSourceJob {
 
-        private DBCExecutionContext context;
         private DBSDataContainer dataContainer;
         private DBDRowIdentifier rowIdentifier;
         private List<ResultSetRow> rows;
 
-        public RowRefreshService(DBCExecutionContext context, DBSDataContainer dataContainer, DBDRowIdentifier rowIdentifier, List<ResultSetRow> rows) {
-            super("Refresh rows", context);
-            this.context = context;
+        public RowRefreshJob(DBCExecutionContext context, DBSDataContainer dataContainer, DBDRowIdentifier rowIdentifier, List<ResultSetRow> rows) {
+            super("Refresh rows", null, context);
             this.dataContainer = dataContainer;
             this.rowIdentifier = rowIdentifier;
-            this.rows = rows;
+            this.rows = new ArrayList<>(rows);
         }
 
         @Override
-        public Collection<DBNNode> evaluate()
-            throws InvocationTargetException, InterruptedException
-        {
+        protected IStatus run(DBRProgressMonitor monitor) {
             try {
-                final AbstractExecutionSource executionSource = new AbstractExecutionSource(dataContainer, context, this);
-                List<DBDAttributeConstraint> constraints = new ArrayList<>();
-                DBDDataFilter filter = new DBDDataFilter(constraints);
-                RowDataReceiver dataReceiver = new RowDataReceiver();
-                try (DBCSession session = context.openSession(getProgressMonitor(), DBCExecutionPurpose.UTIL, "Refresh row(s) after insert/update")) {
+                final Object[][] refreshValues = new Object[rows.size()][];
+
+                final DBDAttributeBinding[] curAttributes = viewer.getModel().getAttributes();
+                final AbstractExecutionSource executionSource = new AbstractExecutionSource(dataContainer, getExecutionContext(), this);
+                try (DBCSession session = getExecutionContext().openSession(monitor, DBCExecutionPurpose.UTIL, "Refresh row(s) after insert/update")) {
+                    for (int i = 0; i < rows.size(); i++) {
+                        ResultSetRow row = rows.get(i);
+                        boolean hasKey = true;
+                        List<DBDAttributeConstraint> constraints = new ArrayList<>();
+                        for (DBDAttributeBinding keyAttr : rowIdentifier.getAttributes()) {
+                            final Object keyValue = row.values[keyAttr.getOrdinalPosition()];
+                            if (DBUtils.isNullValue(keyValue)) {
+                                hasKey = false;
+                                break;
+                            }
+                            final DBDAttributeConstraint constraint = new DBDAttributeConstraint(keyAttr);
+                            constraint.setOperator(DBCLogicalOperator.EQUALS);
+                            constraint.setValue(keyValue);
+                            constraints.add(constraint);
+                        }
+                        if (!hasKey) {
+                            // No key value for this row
+                            continue;
+                        }
+                        DBDDataFilter filter = new DBDDataFilter(constraints);
 
-                    final DBCStatistics stats = dataContainer.readData(executionSource, session, dataReceiver, filter, -1, 1, DBSDataContainer.FLAG_NONE);
+                        RowDataReceiver dataReceiver = new RowDataReceiver(curAttributes);
+                        final DBCStatistics stats = dataContainer.readData(executionSource, session, dataReceiver, filter, 0, 0, DBSDataContainer.FLAG_NONE);
+                        refreshValues[i] = dataReceiver.rowValues;
+                    }
                 }
-                return null;
+
+                // Ok, now we have refreshed values. Let's update real model
+                DBeaverUI.syncExec(new Runnable() {
+                    @Override
+                    public void run() {
+                        for (int i = 0; i < rows.size(); i++) {
+                            if (refreshValues[i] != null) {
+                                rows.get(i).values = refreshValues[i];
+                            }
+                        }
+                        viewer.redrawData(true);
+                    }
+                });
+                return Status.OK_STATUS;
             } catch (Throwable ex) {
-                throw new InvocationTargetException(ex);
+                return GeneralUtils.makeExceptionStatus(ex);
             }
         }
     }

plugins/org.jkiss.dbeaver.core/src/org/jkiss/dbeaver/ui/controls/resultset/ResultSetViewer.java

@@ -2305,11 +2305,22 @@ public class ResultSetViewer extends Viewer
      * @param monitor monitor. If null then save will be executed in async job
      * @param listener finish listener (may be null)
      */
-    public boolean applyChanges(@Nullable DBRProgressMonitor monitor, @Nullable ResultSetPersister.DataUpdateListener listener)
+    public boolean applyChanges(@Nullable final DBRProgressMonitor monitor, @Nullable final ResultSetPersister.DataUpdateListener listener)
     {
         try {
-            ResultSetPersister persister = createDataPersister(false);
-            return persister.applyChanges(monitor, false, listener);
+            final ResultSetPersister persister = createDataPersister(false);
+            return persister.applyChanges(monitor, false, new ResultSetPersister.DataUpdateListener() {
+                @Override
+                public void onUpdate(boolean success) {
+                    if (listener != null) listener.onUpdate(success);
+                    try {
+                        persister.refreshInsertedRows(monitor);
+                    } catch (Throwable e) {
+                        log.error("Error refreshing rows after update", e);
+                    }
+
+                }
+            });
         } catch (DBException e) {
             UIUtils.showErrorDialog(null, "Apply changes error", "Error saving changes in database", e);
             return false;