From 698caa2c555f8be7f2c7832f2f8e05616e055506 Mon Sep 17 00:00:00 2001
From: Serge Rider <serge@jkiss.org>
Date: Thu, 29 Aug 2019 14:52:16 +0200
Subject: [PATCH] #6637 Data transfer from custom query results: reuse data
 container's execution context

Former-commit-id: 6803f7ca8bee627b36fdeefe1588308e03662f03
---
 .../handlers/OpenSpreadsheetHandler.java       |  2 +-
 .../database/DatabaseTransferProducer.java     |  9 ++++++---
 .../src/org/jkiss/dbeaver/model/DBUtils.java   |  3 +++
 .../resultset/ResultSetDataContainer.java      | 18 ++++++++++++++----
 .../resultset/ResultSetHandlerCopyAs.java      |  2 +-
 .../resultset/ResultSetHandlerMain.java        |  2 +-
 .../resultset/ResultSetHandlerOpenWith.java    |  2 +-
 .../ui/controls/resultset/ResultSetViewer.java |  4 ++--
 8 files changed, 29 insertions(+), 13 deletions(-)

diff --git a/plugins/org.jkiss.dbeaver.data.office/src/org/jkiss/dbeaver/data/office/handlers/OpenSpreadsheetHandler.java b/plugins/org.jkiss.dbeaver.data.office/src/org/jkiss/dbeaver/data/office/handlers/OpenSpreadsheetHandler.java
index 235fbdea2b..25adcd381f 100644
--- a/plugins/org.jkiss.dbeaver.data.office/src/org/jkiss/dbeaver/data/office/handlers/OpenSpreadsheetHandler.java
+++ b/plugins/org.jkiss.dbeaver.data.office/src/org/jkiss/dbeaver/data/office/handlers/OpenSpreadsheetHandler.java
@@ -77,7 +77,7 @@ public class OpenSpreadsheetHandler extends AbstractHandler
             options.setSelectedRows(selectedRows);
             options.setSelectedColumns(selectedAttributes);
         }
-        ResultSetDataContainer dataContainer = new ResultSetDataContainer(resultSet.getDataContainer(), resultSet.getModel(), options);
+        ResultSetDataContainer dataContainer = new ResultSetDataContainer(resultSet, options);
         if (dataContainer.getDataSource() == null) {
             DBeaverUI.getInstance().showError("Open Excel", ModelMessages.error_not_connected_to_database);
             return null;
diff --git a/plugins/org.jkiss.dbeaver.data.transfer/src/org/jkiss/dbeaver/tools/transfer/database/DatabaseTransferProducer.java b/plugins/org.jkiss.dbeaver.data.transfer/src/org/jkiss/dbeaver/tools/transfer/database/DatabaseTransferProducer.java
index bfdb80511a..53a2016734 100644
--- a/plugins/org.jkiss.dbeaver.data.transfer/src/org/jkiss/dbeaver/tools/transfer/database/DatabaseTransferProducer.java
+++ b/plugins/org.jkiss.dbeaver.data.transfer/src/org/jkiss/dbeaver/tools/transfer/database/DatabaseTransferProducer.java
@@ -128,9 +128,12 @@ public class DatabaseTransferProducer implements IDataTransferProducer<DatabaseP
 
         boolean newConnection = settings.isOpenNewConnections() && !getDatabaseObject().getDataSource().getContainer().getDriver().isEmbedded();
         boolean forceDataReadTransactions = Boolean.TRUE.equals(dataSource.getDataSourceFeature(DBConstants.FEATURE_LOB_REQUIRE_TRANSACTIONS));
-        DBCExecutionContext context = !selectiveExportFromUI && newConnection ?
-            DBUtils.getObjectOwnerInstance(getDatabaseObject()).openIsolatedContext(monitor, "Data transfer producer") :
-            DBUtils.getDefaultContext(getDatabaseObject(), false);
+        DBCExecutionContext context;
+        if (!selectiveExportFromUI && newConnection) {
+            context = DBUtils.getObjectOwnerInstance(getDatabaseObject()).openIsolatedContext(monitor, "Data transfer producer");
+        } else {
+            context = DBUtils.getDefaultContext(dataContainer, false);
+        }
         try (DBCSession session = context.openSession(monitor, DBCExecutionPurpose.UTIL, contextTask)) {
             try {
                 AbstractExecutionSource transferSource = new AbstractExecutionSource(dataContainer, context, consumer);
diff --git a/plugins/org.jkiss.dbeaver.model/src/org/jkiss/dbeaver/model/DBUtils.java b/plugins/org.jkiss.dbeaver.model/src/org/jkiss/dbeaver/model/DBUtils.java
index b595cf9ccc..1d4cfa7419 100644
--- a/plugins/org.jkiss.dbeaver.model/src/org/jkiss/dbeaver/model/DBUtils.java
+++ b/plugins/org.jkiss.dbeaver.model/src/org/jkiss/dbeaver/model/DBUtils.java
@@ -1811,6 +1811,9 @@ public final class DBUtils {
         if (object == null) {
             return null;
         }
+        if (object instanceof DBPContextProvider) {
+            return ((DBPContextProvider) object).getExecutionContext();
+        }
         DBSInstance instance = getObjectOwnerInstance(object);
         return instance == null ? null : instance.getDefaultContext(meta);
     }
diff --git a/plugins/org.jkiss.dbeaver.ui.editors.data/src/org/jkiss/dbeaver/ui/controls/resultset/ResultSetDataContainer.java b/plugins/org.jkiss.dbeaver.ui.editors.data/src/org/jkiss/dbeaver/ui/controls/resultset/ResultSetDataContainer.java
index 24c73c3a00..31110a61f8 100644
--- a/plugins/org.jkiss.dbeaver.ui.editors.data/src/org/jkiss/dbeaver/ui/controls/resultset/ResultSetDataContainer.java
+++ b/plugins/org.jkiss.dbeaver.ui.editors.data/src/org/jkiss/dbeaver/ui/controls/resultset/ResultSetDataContainer.java
@@ -18,7 +18,9 @@ package org.jkiss.dbeaver.ui.controls.resultset;
 
 import org.eclipse.core.runtime.IAdaptable;
 import org.jkiss.code.NotNull;
+import org.jkiss.code.Nullable;
 import org.jkiss.dbeaver.Log;
+import org.jkiss.dbeaver.model.DBPContextProvider;
 import org.jkiss.dbeaver.model.DBPDataSource;
 import org.jkiss.dbeaver.model.DBUtils;
 import org.jkiss.dbeaver.model.data.DBDAttributeBinding;
@@ -38,17 +40,19 @@ import java.util.List;
  * Client-side data container.
  * Wraps RSV model and original data container.
  */
-public class ResultSetDataContainer implements DBSDataContainer, IAdaptable {
+public class ResultSetDataContainer implements DBSDataContainer, DBPContextProvider, IAdaptable {
 
     private static final Log log = Log.getLog(ResultSetDataContainer.class);
 
+    private final IResultSetController controller;
     private final DBSDataContainer dataContainer;
     private final ResultSetModel model;
     private ResultSetDataContainerOptions options;
 
-    public ResultSetDataContainer(DBSDataContainer dataContainer, ResultSetModel model, ResultSetDataContainerOptions options) {
-        this.dataContainer = dataContainer;
-        this.model = model;
+    public ResultSetDataContainer(IResultSetController controller, ResultSetDataContainerOptions options) {
+        this.controller = controller;
+        this.dataContainer = controller.getDataContainer();
+        this.model = controller.getModel();
         this.options = options;
     }
 
@@ -149,6 +153,12 @@ public class ResultSetDataContainer implements DBSDataContainer, IAdaptable {
         return null;
     }
 
+    @Nullable
+    @Override
+    public DBCExecutionContext getExecutionContext() {
+        return controller.getExecutionContext();
+    }
+
     private class ModelResultSet implements DBCResultSet {
 
         private final DBCSession session;
diff --git a/plugins/org.jkiss.dbeaver.ui.editors.data/src/org/jkiss/dbeaver/ui/controls/resultset/ResultSetHandlerCopyAs.java b/plugins/org.jkiss.dbeaver.ui.editors.data/src/org/jkiss/dbeaver/ui/controls/resultset/ResultSetHandlerCopyAs.java
index b8d138e271..71ff6ffcae 100644
--- a/plugins/org.jkiss.dbeaver.ui.editors.data/src/org/jkiss/dbeaver/ui/controls/resultset/ResultSetHandlerCopyAs.java
+++ b/plugins/org.jkiss.dbeaver.ui.editors.data/src/org/jkiss/dbeaver/ui/controls/resultset/ResultSetHandlerCopyAs.java
@@ -112,7 +112,7 @@ public class ResultSetHandlerCopyAs extends AbstractHandler implements IElementU
             options.setSelectedRows(selectedRows);
             options.setSelectedColumns(selectedAttributes);
         }
-        ResultSetDataContainer dataContainer = new ResultSetDataContainer(resultSet.getDataContainer(), resultSet.getModel(), options);
+        ResultSetDataContainer dataContainer = new ResultSetDataContainer(resultSet, options);
         if (dataContainer.getDataSource() == null) {
             DBWorkbench.getPlatformUI().showError("Copy As " + processor.getName(), ModelMessages.error_not_connected_to_database);
             return;
diff --git a/plugins/org.jkiss.dbeaver.ui.editors.data/src/org/jkiss/dbeaver/ui/controls/resultset/ResultSetHandlerMain.java b/plugins/org.jkiss.dbeaver.ui.editors.data/src/org/jkiss/dbeaver/ui/controls/resultset/ResultSetHandlerMain.java
index 7bda839022..39ad238029 100644
--- a/plugins/org.jkiss.dbeaver.ui.editors.data/src/org/jkiss/dbeaver/ui/controls/resultset/ResultSetHandlerMain.java
+++ b/plugins/org.jkiss.dbeaver.ui.editors.data/src/org/jkiss/dbeaver/ui/controls/resultset/ResultSetHandlerMain.java
@@ -509,7 +509,7 @@ public class ResultSetHandlerMain extends AbstractHandler {
                 options.setSelectedRows(selectedRows);
                 options.setSelectedColumns(selectedAttributes);
 
-                ResultSetDataContainer dataContainer = new ResultSetDataContainer(rsv.getDataContainer(), rsv.getModel(), options);
+                ResultSetDataContainer dataContainer = new ResultSetDataContainer(rsv, options);
                 ActiveWizardDialog dialog = new ActiveWizardDialog(
                     HandlerUtil.getActiveWorkbenchWindow(event),
                     new DataTransferWizard(
diff --git a/plugins/org.jkiss.dbeaver.ui.editors.data/src/org/jkiss/dbeaver/ui/controls/resultset/ResultSetHandlerOpenWith.java b/plugins/org.jkiss.dbeaver.ui.editors.data/src/org/jkiss/dbeaver/ui/controls/resultset/ResultSetHandlerOpenWith.java
index 086c7a7860..8c024d24dd 100644
--- a/plugins/org.jkiss.dbeaver.ui.editors.data/src/org/jkiss/dbeaver/ui/controls/resultset/ResultSetHandlerOpenWith.java
+++ b/plugins/org.jkiss.dbeaver.ui.editors.data/src/org/jkiss/dbeaver/ui/controls/resultset/ResultSetHandlerOpenWith.java
@@ -131,7 +131,7 @@ public class ResultSetHandlerOpenWith extends AbstractHandler implements IElemen
             options.setSelectedRows(selectedRows);
             options.setSelectedColumns(selectedAttributes);
         }
-        ResultSetDataContainer dataContainer = new ResultSetDataContainer(resultSet.getDataContainer(), resultSet.getModel(), options);
+        ResultSetDataContainer dataContainer = new ResultSetDataContainer(resultSet, options);
         if (dataContainer.getDataSource() == null) {
             DBWorkbench.getPlatformUI().showError("Open " + processor.getAppName(), ModelMessages.error_not_connected_to_database);
             return;
diff --git a/plugins/org.jkiss.dbeaver.ui.editors.data/src/org/jkiss/dbeaver/ui/controls/resultset/ResultSetViewer.java b/plugins/org.jkiss.dbeaver.ui.editors.data/src/org/jkiss/dbeaver/ui/controls/resultset/ResultSetViewer.java
index 00237ec61d..fe3e555865 100644
--- a/plugins/org.jkiss.dbeaver.ui.editors.data/src/org/jkiss/dbeaver/ui/controls/resultset/ResultSetViewer.java
+++ b/plugins/org.jkiss.dbeaver.ui.editors.data/src/org/jkiss/dbeaver/ui/controls/resultset/ResultSetViewer.java
@@ -1221,7 +1221,7 @@ public class ResultSetViewer extends Viewer
     void fillOpenWithMenu(IContributionManager openWithMenu) {
 
         ResultSetDataContainerOptions options = new ResultSetDataContainerOptions();
-        ResultSetDataContainer dataContainer = new ResultSetDataContainer(getDataContainer(), getModel(), options);
+        ResultSetDataContainer dataContainer = new ResultSetDataContainer(this, options);
 
         List<DataTransferProcessorDescriptor> appProcessors = new ArrayList<>();
 
@@ -1262,7 +1262,7 @@ public class ResultSetViewer extends Viewer
         copyAsMenu.add(new Separator());
 
         ResultSetDataContainerOptions options = new ResultSetDataContainerOptions();
-        ResultSetDataContainer dataContainer = new ResultSetDataContainer(getDataContainer(), getModel(), options);
+        ResultSetDataContainer dataContainer = new ResultSetDataContainer(this, options);
 
         List<DataTransferProcessorDescriptor> appProcessors = new ArrayList<>();
 
-- 
GitLab