#2855 Export from multiple queries

plugins/org.jkiss.dbeaver.ui.editors.sql/src/org/jkiss/dbeaver/ui/editors/sql/SQLEditor.java

@@ -47,7 +47,6 @@ import org.eclipse.swt.widgets.*;
 import org.eclipse.ui.*;
 import org.eclipse.ui.actions.CompoundContributionItem;
 import org.eclipse.ui.ide.FileStoreEditorInput;
-import org.eclipse.ui.menus.CommandContributionItem;
 import org.eclipse.ui.texteditor.DefaultRangeIndicator;
 import org.eclipse.ui.texteditor.ITextEditorActionConstants;
 import org.eclipse.ui.texteditor.rulers.IColumnSupport;
@@ -1624,13 +1623,23 @@ public class SQLEditor extends SQLEditorBase implements
 
     public void exportDataFromQuery()
     {
-        SQLScriptElement sqlQuery = extractActiveQuery();
-        if (sqlQuery instanceof SQLQuery) {
-            processQueries(Collections.singletonList(sqlQuery), false, true, true, null);
+        List<SQLScriptElement> elements;
+        ITextSelection selection = (ITextSelection) getSelectionProvider().getSelection();
+        if (selection.getLength() > 1) {
+            elements = extractScriptQueries(selection.getOffset(), selection.getLength(), true, false, true);
+        } else {
+            elements = new ArrayList<>();
+            elements.add(extractActiveQuery());
+        }
+
+        elements.removeIf(element -> !(element instanceof SQLQuery));
+
+        if (!elements.isEmpty()) {
+            processQueries(elements, false, true, true, null);
         } else {
             DBWorkbench.getPlatformUI().showError(
                     "Extract data",
-                    "Can't extract data from control command");
+                    "Choose one or more queries to export from");
         }
     }
 
@@ -1714,13 +1723,15 @@ public class SQLEditor extends SQLEditorBase implements
         if (getActivePreferenceStore().getBoolean(SQLPreferenceConstants.AUTO_SAVE_ON_EXECUTE) && isDirty()) {
             doSave(new NullProgressMonitor());
         }
-        if (getActivePreferenceStore().getBoolean(SQLPreferenceConstants.CLEAR_OUTPUT_BEFORE_EXECUTE)) {
-            outputViewer.clearOutput();
-        }
+        if (!export) {
+            if (getActivePreferenceStore().getBoolean(SQLPreferenceConstants.CLEAR_OUTPUT_BEFORE_EXECUTE)) {
+                outputViewer.clearOutput();
+            }
 
-        if (!newTab || !isSingleQuery) {
-            // We don't need new tab or we are executing a script - so close all extra tabs
-            closeExtraResultTabs(null);
+            if (!newTab || !isSingleQuery) {
+                // We don't need new tab or we are executing a script - so close all extra tabs
+                closeExtraResultTabs(null);
+            }
         }
 
         if (resultTabs.getItemCount() == 0) {
@@ -1740,29 +1751,31 @@ public class SQLEditor extends SQLEditorBase implements
                         getActivePreferenceStore().getBoolean(SQLPreferenceConstants.RESULT_SET_CLOSE_ON_ERROR), queryListener);
             }
         } else {
-            // Use current tab.
-            // If current tab was pinned then use first tab
-            QueryResultsContainer firstResults = curQueryProcessor.getFirstResults();
-            if (firstResults.isPinned()) {
-                curQueryProcessor = queryProcessors.get(0);
-                firstResults = curQueryProcessor.getFirstResults();
+            if (!export) {
+                // Use current tab.
+                // If current tab was pinned then use first tab
+                QueryResultsContainer firstResults = curQueryProcessor.getFirstResults();
                 if (firstResults.isPinned()) {
-                    // The very first tab is also pinned
-                    // Well, let's create a new tab
-                    curQueryProcessor = createQueryProcessor(true, true);
-                    // Make new tab the default
+                    curQueryProcessor = queryProcessors.get(0);
                     firstResults = curQueryProcessor.getFirstResults();
                     if (firstResults.isPinned()) {
-                        firstResults.tabItem.setShowClose(false);
+                        // The very first tab is also pinned
+                        // Well, let's create a new tab
+                        curQueryProcessor = createQueryProcessor(true, true);
+                        // Make new tab the default
+                        firstResults = curQueryProcessor.getFirstResults();
+                        if (firstResults.isPinned()) {
+                            firstResults.tabItem.setShowClose(false);
+                        }
                     }
                 }
-            }
-            closeExtraResultTabs(curQueryProcessor);
-            if (firstResults.tabItem != null) {
-                // Do not switch tab if Output tab is active
-                CTabItem selectedTab = resultTabs.getSelection();
-                if (selectedTab == null || selectedTab.getData() != outputViewer) {
-                    resultTabs.setSelection(firstResults.tabItem);
+                closeExtraResultTabs(curQueryProcessor);
+                if (firstResults.tabItem != null) {
+                    // Do not switch tab if Output tab is active
+                    CTabItem selectedTab = resultTabs.getSelection();
+                    if (selectedTab == null || selectedTab.getData() != outputViewer) {
+                        resultTabs.setSelection(firstResults.tabItem);
+                    }
                 }
             }
             curQueryProcessor.processQueries(queries, false, export, false, queryListener);
@@ -2340,46 +2353,63 @@ public class SQLEditor extends SQLEditorBase implements
                     listener.setExtListener(queryListener);
                 }
                 File localFile = EditorUtils.getLocalFileFromInput(getEditorInput());
-                final SQLQueryJob job = new SQLQueryJob(
-                    getSite(),
-                    isSingleQuery ? SQLEditorMessages.editors_sql_job_execute_query : SQLEditorMessages.editors_sql_job_execute_script,
-                    executionContext,
-                    resultsContainer,
-                    queries,
-                    new SQLScriptContext(globalScriptContext, SQLEditor.this, localFile, new OutputLogWriter()),
-                    this,
-                    listener);
 
-                if (export || isSingleQuery) {
-                    resultsContainer.query = queries.get(0);
-                }
                 if (export) {
-                    // Assign current job from active query and open wizard
-                    resultsContainer.lastGoodQuery = null;
-                    curJob = job;
+                    // Use special consumer and data containers
+                    SQLQueryResultsConsumer resultsConsumer = new SQLQueryResultsConsumer();
+                    SQLQueryJob job = new SQLQueryJob(
+                        getSite(),
+                        "Export from query",
+                        executionContext,
+                        resultsContainer,
+                        queries,
+                        new SQLScriptContext(globalScriptContext, SQLEditor.this, localFile, new OutputLogWriter()),
+                        resultsConsumer,
+                        listener);
+
+                    List<IDataTransferProducer> producers = new ArrayList<>();
+                    for (int i = 0; i < queries.size(); i++) {
+                        SQLScriptElement element = queries.get(i);
+                        SQLQuery query = (SQLQuery) element;
+                        producers.add(new DatabaseTransferProducer(new SQLQueryDataContainer(job, resultsConsumer, query, i), null));
+                    }
+
                     ActiveWizardDialog dialog = new ActiveWizardDialog(
                         getSite().getWorkbenchWindow(),
                         new DataTransferWizard(
-                            new IDataTransferProducer[] {
-                                new DatabaseTransferProducer(resultsContainer, null)},
+                            producers.toArray(new IDataTransferProducer[0]),
                             null),
                         new StructuredSelection(this));
                     dialog.open();
-                } else if (isSingleQuery) {
-                    closeJob();
-                    curJob = job;
-                    ResultSetViewer rsv = resultsContainer.getResultSetController();
-                    if (rsv != null) {
-                        rsv.resetDataFilter(false);
-                        rsv.resetHistory();
-                        rsv.refresh();
-                    }
                 } else {
-                    if (fetchResults) {
-                        job.setFetchResultSets(true);
+                    final SQLQueryJob job = new SQLQueryJob(
+                        getSite(),
+                        isSingleQuery ? SQLEditorMessages.editors_sql_job_execute_query : SQLEditorMessages.editors_sql_job_execute_script,
+                        executionContext,
+                        resultsContainer,
+                        queries,
+                        new SQLScriptContext(globalScriptContext, SQLEditor.this, localFile, new OutputLogWriter()),
+                        this,
+                        listener);
+
+                    if (isSingleQuery) {
+                        resultsContainer.query = queries.get(0);
+
+                        closeJob();
+                        curJob = job;
+                        ResultSetViewer rsv = resultsContainer.getResultSetController();
+                        if (rsv != null) {
+                            rsv.resetDataFilter(false);
+                            rsv.resetHistory();
+                            rsv.refresh();
+                        }
+                    } else {
+                        if (fetchResults) {
+                            job.setFetchResultSets(true);
+                        }
+                        job.schedule();
+                        curJob = job;
                     }
-                    job.schedule();
-                    curJob = job;
                 }
             }
         }

plugins/org.jkiss.dbeaver.ui.editors.sql/src/org/jkiss/dbeaver/ui/editors/sql/SQLQueryDataContainer.java

+package org.jkiss.dbeaver.ui.editors.sql;
+
+import org.jkiss.code.NotNull;
+import org.jkiss.code.Nullable;
+import org.jkiss.dbeaver.model.DBPContextProvider;
+import org.jkiss.dbeaver.model.DBPDataSource;
+import org.jkiss.dbeaver.model.DBPDataSourceContainer;
+import org.jkiss.dbeaver.model.IDataSourceContainerProvider;
+import org.jkiss.dbeaver.model.data.DBDDataFilter;
+import org.jkiss.dbeaver.model.data.DBDDataReceiver;
+import org.jkiss.dbeaver.model.exec.*;
+import org.jkiss.dbeaver.model.sql.SQLQuery;
+import org.jkiss.dbeaver.model.sql.SQLQueryContainer;
+import org.jkiss.dbeaver.model.sql.SQLScriptElement;
+import org.jkiss.dbeaver.model.struct.DBSDataContainer;
+import org.jkiss.dbeaver.model.struct.DBSObject;
+import org.jkiss.dbeaver.ui.editors.sql.execute.SQLQueryJob;
+import org.jkiss.dbeaver.ui.editors.sql.internal.SQLEditorMessages;
+
+class SQLQueryDataContainer implements DBSDataContainer, IDataSourceContainerProvider, SQLQueryContainer, DBPContextProvider {
+
+    private SQLQueryJob queryJob;
+    private SQLQueryResultsConsumer resultsConsumer;
+    private SQLQuery query = null;
+    private int resultSetNumber;
+
+    private DBDDataReceiver dataReceiver;
+
+    SQLQueryDataContainer(SQLQueryJob queryJob, SQLQueryResultsConsumer resultsConsumer, SQLQuery query, int resultSetNumber)
+    {
+        this.queryJob = queryJob;
+        this.resultsConsumer = resultsConsumer;
+        this.query = query;
+        this.resultSetNumber = resultSetNumber;
+    }
+
+    @Override
+    public DBCExecutionContext getExecutionContext() {
+        return queryJob.getExecutionContext();
+    }
+
+    @Override
+    public int getSupportedFeatures() {
+        return DATA_SELECT;
+    }
+
+    @NotNull
+    @Override
+    public DBCStatistics readData(@NotNull DBCExecutionSource source, @NotNull DBCSession session, @NotNull DBDDataReceiver dataReceiver, DBDDataFilter dataFilter, long firstRow, long maxRows, long flags) throws DBCException
+    {
+        SQLQuery query = (SQLQuery) this.query;
+        if (query.getResultsMaxRows() >= 0) {
+            firstRow = query.getResultsOffset();
+            maxRows = query.getResultsMaxRows();
+        }
+
+        resultsConsumer.setDataReceiver(dataReceiver);
+
+        // Count number of results for this query. If > 1 then we will refresh them all at once
+        if (resultSetNumber > 0) {
+            queryJob.setFetchResultSetNumber(resultSetNumber);
+        } else {
+            queryJob.setFetchResultSetNumber(-1);
+        }
+        queryJob.setResultSetLimit(firstRow, maxRows);
+        queryJob.setReadFlags(flags);
+        queryJob.setDataFilter(dataFilter);
+
+        queryJob.extractData(session, this.query, resultSetNumber);
+
+        return queryJob.getStatistics();
+    }
+
+    @Override
+    public long countData(@NotNull DBCExecutionSource source, @NotNull DBCSession session, DBDDataFilter dataFilter, long flags)
+        throws DBCException
+    {
+        return -1;
+    }
+
+    @Nullable
+    @Override
+    public String getDescription()
+    {
+        return SQLEditorMessages.editors_sql_description;
+    }
+
+    @Nullable
+    @Override
+    public DBSObject getParentObject()
+    {
+        return getDataSource();
+    }
+
+    @Nullable
+    @Override
+    public DBPDataSource getDataSource()
+    {
+        return queryJob.getDataSourceContainer().getDataSource();
+    }
+
+    @Override
+    public boolean isPersisted() {
+        return false;
+    }
+
+    @NotNull
+    @Override
+    public String getName()
+    {
+        String name = query.getOriginalText();
+        if (name == null) {
+            name = "SQL";
+        }
+        return name;
+    }
+
+    @Nullable
+    @Override
+    public DBPDataSourceContainer getDataSourceContainer() {
+        return queryJob.getDataSourceContainer();
+    }
+
+    @Override
+    public String toString() {
+        return query.getOriginalText();
+    }
+
+    @Override
+    public SQLScriptElement getQuery() {
+        return query;
+    }
+
+}

plugins/org.jkiss.dbeaver.ui.editors.sql/src/org/jkiss/dbeaver/ui/editors/sql/SQLQueryResultsConsumer.java

+package org.jkiss.dbeaver.ui.editors.sql;
+
+import org.jkiss.dbeaver.model.data.DBDDataReceiver;
+import org.jkiss.dbeaver.model.sql.SQLQuery;
+import org.jkiss.dbeaver.runtime.sql.SQLResultsConsumer;
+
+class SQLQueryResultsConsumer implements SQLResultsConsumer {
+
+    private DBDDataReceiver dataReceiver;
+
+    public void setDataReceiver(DBDDataReceiver dataReceiver) {
+        this.dataReceiver = dataReceiver;
+    }
+
+    @Override
+    public DBDDataReceiver getDataReceiver(SQLQuery statement, int resultSetNumber) {
+        return dataReceiver;
+    }
+}