Skip to content

Commit

Permalink
CB-2153 update values of tables with join fix (#1961)
Browse files Browse the repository at this point in the history 
Co-authored-by: EvgeniaBzzz <[email protected]>
  • Loading branch information
@yagudin10 @EvgeniaBzzz
yagudin10 and EvgeniaBzzz authored Sep 28, 2023
1 parent 6da5f3b commit c58e62c
Show file tree
Hide file tree
Showing 2 changed files with 99 additions and 58 deletions.
14 changes: 14 additions & 0 deletions server/bundles/io.cloudbeaver.model/src/io/cloudbeaver/service/sql/WebSQLResultsInfo.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,9 @@
import org.jkiss.dbeaver.model.struct.DBSDataContainer;
import org.jkiss.dbeaver.model.struct.DBSTypedObject;

import java.util.HashSet;
import java.util.List;
import java.util.Set;

/**
* Web query results info.
Expand Down Expand Up @@ -82,6 +84,18 @@ public DBDRowIdentifier getDefaultRowIdentifier() {
return null;
}

@NotNull
public Set<DBDRowIdentifier> getRowIdentifiers() {
Set<DBDRowIdentifier> rowIdentifiers = new HashSet<>();
for (DBDAttributeBinding column : attributes) {
DBDRowIdentifier rowIdentifier = column.getRowIdentifier();
if (rowIdentifier != null) {
rowIdentifiers.add(rowIdentifier);
}
}
return rowIdentifiers;
}

public DBSAttributeBase getAttribute(String attributeName) {
DBPDataSource dataSource = dataContainer.getDataSource();

Expand Down
143 changes: 85 additions & 58 deletions server/bundles/io.cloudbeaver.server/src/io/cloudbeaver/service/sql/WebSQLProcessor.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,7 @@
import java.lang.reflect.InvocationTargetException;
import java.util.*;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.stream.Collectors;

/**
* Web SQL processor.
Expand Down Expand Up @@ -315,56 +316,67 @@ public WebSQLExecuteInfo updateResultsDataBatch(
@Nullable List<WebSQLResultsRow> addedRows,
@Nullable WebDataFormat dataFormat) throws DBException
{
Map<DBSDataManipulator.ExecuteBatch, Object[]> resultBatches = new LinkedHashMap<>();

List<Object[]> newResultSetRows = new ArrayList<>();
KeyDataReceiver keyReceiver = new KeyDataReceiver(contextInfo.getResults(resultsId));

DBSDataManipulator dataManipulator = generateUpdateResultsDataBatch(
monitor, contextInfo, resultsId, updatedRows, deletedRows, addedRows, dataFormat, resultBatches, keyReceiver);

WebSQLResultsInfo resultsInfo = contextInfo.getResults(resultsId);

Set<DBDRowIdentifier> rowIdentifierList = new HashSet<>();
// several row identifiers could be if we update result set table with join
// we can't add or delete rows from result set table with join
if (!CommonUtils.isEmpty(deletedRows) || !CommonUtils.isEmpty(addedRows)) {
rowIdentifierList.add(resultsInfo.getDefaultRowIdentifier());
} else if (!CommonUtils.isEmpty(updatedRows)) {
rowIdentifierList = resultsInfo.getRowIdentifiers();
}

long totalUpdateCount = 0;

WebSQLExecuteInfo result = new WebSQLExecuteInfo();
List<WebSQLQueryResults> queryResults = new ArrayList<>();
for (var rowIdentifier : rowIdentifierList) {
Map<DBSDataManipulator.ExecuteBatch, Object[]> resultBatches = new LinkedHashMap<>();
DBSDataManipulator dataManipulator = generateUpdateResultsDataBatch(
monitor, resultsInfo, rowIdentifier, updatedRows, deletedRows, addedRows, dataFormat, resultBatches, keyReceiver);


DBCExecutionContext executionContext = getExecutionContext(dataManipulator);
try (DBCSession session = executionContext.openSession(monitor, DBCExecutionPurpose.USER, "Update data in container")) {
DBCTransactionManager txnManager = DBUtils.getTransactionManager(executionContext);
boolean revertToAutoCommit = false;
if (txnManager != null && txnManager.isSupportsTransactions() && txnManager.isAutoCommit()) {
txnManager.setAutoCommit(monitor, false);
revertToAutoCommit = true;
}
try {
Map<String, Object> options = Collections.emptyMap();
for (Map.Entry<DBSDataManipulator.ExecuteBatch, Object[]> rb : resultBatches.entrySet()) {
DBSDataManipulator.ExecuteBatch batch = rb.getKey();
Object[] rowValues = rb.getValue();
keyReceiver.setRow(rowValues);
DBCStatistics statistics = batch.execute(session, options);

// Patch result rows (adapt to web format)
for (int i = 0; i < rowValues.length; i++) {
rowValues[i] = WebSQLUtils.makeWebCellValue(webSession, resultsInfo.getAttributeByPosition(i), rowValues[i], dataFormat);
}

DBCExecutionContext executionContext = getExecutionContext(dataManipulator);
try (DBCSession session = executionContext.openSession(monitor, DBCExecutionPurpose.USER, "Update data in container")) {
DBCTransactionManager txnManager = DBUtils.getTransactionManager(executionContext);
boolean revertToAutoCommit = false;
if (txnManager != null && txnManager.isSupportsTransactions() && txnManager.isAutoCommit()) {
txnManager.setAutoCommit(monitor, false);
revertToAutoCommit = true;
}
try {
Map<String, Object> options = Collections.emptyMap();
for (Map.Entry<DBSDataManipulator.ExecuteBatch, Object[]> rb : resultBatches.entrySet()) {
DBSDataManipulator.ExecuteBatch batch = rb.getKey();
Object[] rowValues = rb.getValue();
keyReceiver.setRow(rowValues);
DBCStatistics statistics = batch.execute(session, options);

// Patch result rows (adapt to web format)
for (int i = 0; i < rowValues.length; i++) {
rowValues[i] = WebSQLUtils.makeWebCellValue(webSession, resultsInfo.getAttributeByPosition(i), rowValues[i], dataFormat);
totalUpdateCount += statistics.getRowsUpdated();
result.setDuration(result.getDuration() + statistics.getExecuteTime());
newResultSetRows.add(rowValues);
}

totalUpdateCount += statistics.getRowsUpdated();
result.setDuration(result.getDuration() + statistics.getExecuteTime());
}

if (txnManager != null && txnManager.isSupportsTransactions()) {
txnManager.commit(session);
}
} catch (Exception e) {
if (txnManager != null && txnManager.isSupportsTransactions()) {
txnManager.rollback(session, null);
}
throw new DBCException("Error persisting data changes", e);
} finally {
if (revertToAutoCommit) {
txnManager.setAutoCommit(monitor, true);
if (txnManager != null && txnManager.isSupportsTransactions()) {
txnManager.commit(session);
}
} catch (Exception e) {
if (txnManager != null && txnManager.isSupportsTransactions()) {
txnManager.rollback(session, null);
}
throw new DBCException("Error persisting data changes", e);
} finally {
if (revertToAutoCommit) {
txnManager.setAutoCommit(monitor, true);
}
}
}
}
Expand All @@ -376,7 +388,7 @@ public WebSQLExecuteInfo updateResultsDataBatch(
WebSQLQueryResults updateResults = new WebSQLQueryResults(webSession, dataFormat);
updateResults.setUpdateRowCount(totalUpdateCount);
updateResults.setResultSet(updatedResultSet);
updatedResultSet.setRows(resultBatches.values().toArray(new Object[0][]));
updatedResultSet.setRows(newResultSetRows.toArray(new Object[0][]));

queryResults.add(updateResults);

Expand All @@ -396,26 +408,42 @@ public String generateResultsDataUpdateScript(
{
Map<DBSDataManipulator.ExecuteBatch, Object[]> resultBatches = new LinkedHashMap<>();

DBSDataManipulator dataManipulator = generateUpdateResultsDataBatch(
monitor, contextInfo, resultsId, updatedRows, deletedRows, addedRows, dataFormat, resultBatches, null);

List<DBEPersistAction> actions = new ArrayList<>();

DBCExecutionContext executionContext = getExecutionContext(dataManipulator);
try (DBCSession session = executionContext.openSession(monitor, DBCExecutionPurpose.USER, "Update data in container")) {
Map<String, Object> options = Collections.emptyMap();
for (DBSDataManipulator.ExecuteBatch batch : resultBatches.keySet()) {
batch.generatePersistActions(session, actions, options);
WebSQLResultsInfo resultsInfo = contextInfo.getResults(resultsId);
Set<DBDRowIdentifier> rowIdentifierList = new HashSet<>();
// several row identifiers could be if we update result set table with join
// we can't add or delete rows from result set table with join
if (!CommonUtils.isEmpty(deletedRows) || !CommonUtils.isEmpty(addedRows)) {
rowIdentifierList.add(resultsInfo.getDefaultRowIdentifier());
} else if (!CommonUtils.isEmpty(updatedRows)) {
rowIdentifierList = resultsInfo.getRowIdentifiers();
}
StringBuilder sqlBuilder = new StringBuilder();
for (var rowIdentifier : rowIdentifierList) {
DBSDataManipulator dataManipulator = generateUpdateResultsDataBatch(
monitor, resultsInfo, rowIdentifier, updatedRows, deletedRows, addedRows, dataFormat, resultBatches, null);

List<DBEPersistAction> actions = new ArrayList<>();

DBCExecutionContext executionContext = getExecutionContext(dataManipulator);
try (DBCSession session = executionContext.openSession(monitor, DBCExecutionPurpose.USER, "Update data in container")) {
Map<String, Object> options = Collections.emptyMap();
for (DBSDataManipulator.ExecuteBatch batch : resultBatches.keySet()) {
batch.generatePersistActions(session, actions, options);
}
}
}

return SQLUtils.generateScript(executionContext.getDataSource(), actions.toArray(new DBEPersistAction[0]), false);
sqlBuilder.append(
SQLUtils.generateScript(executionContext.getDataSource(), actions.toArray(new DBEPersistAction[0]), false)
);
}
return sqlBuilder.toString();
}

private DBSDataManipulator generateUpdateResultsDataBatch(
@NotNull DBRProgressMonitor monitor,
@NotNull WebSQLContextInfo contextInfo,
@NotNull String resultsId,
@NotNull WebSQLResultsInfo resultsInfo,
@NotNull DBDRowIdentifier rowIdentifier,
@Nullable List<WebSQLResultsRow> updatedRows,
@Nullable List<WebSQLResultsRow> deletedRows,
@Nullable List<WebSQLResultsRow> addedRows,
Expand All @@ -424,10 +452,7 @@ private DBSDataManipulator generateUpdateResultsDataBatch(
@Nullable DBDDataReceiver keyReceiver)
throws DBException
{
WebSQLResultsInfo resultsInfo = contextInfo.getResults(resultsId);

DBDRowIdentifier rowIdentifier = resultsInfo.getDefaultRowIdentifier();
checkRowIdentifier(resultsInfo, rowIdentifier);
DBSEntity dataContainer = rowIdentifier.getEntity();
checkDataEditAllowed(dataContainer);
DBSDataManipulator dataManipulator = (DBSDataManipulator) dataContainer;
Expand All @@ -448,7 +473,9 @@ private DBSDataManipulator generateUpdateResultsDataBatch(
if (!CommonUtils.isEmpty(updatedRows)) {

for (WebSQLResultsRow row : updatedRows) {
Map<String, Object> updateValues = row.getUpdateValues();
Map<String, Object> updateValues = row.getUpdateValues().entrySet().stream()
.filter(x -> CommonUtils.equalObjects(allAttributes[CommonUtils.toInt(x.getKey())].getRowIdentifier(), rowIdentifier))
.collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
if (CommonUtils.isEmpty(row.getData()) || CommonUtils.isEmpty(updateValues)) {
continue;
}
Expand Down

0 comments on commit c58e62c

Please sign in to comment.