#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,6 +1723,7 @@ public class SQLEditor extends SQLEditorBase implements
         if (getActivePreferenceStore().getBoolean(SQLPreferenceConstants.AUTO_SAVE_ON_EXECUTE) && isDirty()) {
             doSave(new NullProgressMonitor());
         }
+        if (!export) {
             if (getActivePreferenceStore().getBoolean(SQLPreferenceConstants.CLEAR_OUTPUT_BEFORE_EXECUTE)) {
                 outputViewer.clearOutput();
             }
@@ -1722,6 +1732,7 @@ public class SQLEditor extends SQLEditorBase implements
                 // We don't need new tab or we are executing a script - so close all extra tabs
                 closeExtraResultTabs(null);
             }
+        }
 
         if (resultTabs.getItemCount() == 0) {
             // If all tabs were closed
@@ -1740,6 +1751,7 @@ public class SQLEditor extends SQLEditorBase implements
                         getActivePreferenceStore().getBoolean(SQLPreferenceConstants.RESULT_SET_CLOSE_ON_ERROR), queryListener);
             }
         } else {
+            if (!export) {
                 // Use current tab.
                 // If current tab was pinned then use first tab
                 QueryResultsContainer firstResults = curQueryProcessor.getFirstResults();
@@ -1765,6 +1777,7 @@ public class SQLEditor extends SQLEditorBase implements
                         resultTabs.setSelection(firstResults.tabItem);
                     }
                 }
+            }
             curQueryProcessor.processQueries(queries, false, export, false, queryListener);
         }
     }
@@ -2340,32 +2353,48 @@ public class SQLEditor extends SQLEditorBase implements
                     listener.setExtListener(queryListener);
                 }
                 File localFile = EditorUtils.getLocalFileFromInput(getEditorInput());
-                final SQLQueryJob job = new SQLQueryJob(
+
+                if (export) {
+                    // Use special consumer and data containers
+                    SQLQueryResultsConsumer resultsConsumer = new SQLQueryResultsConsumer();
+                    SQLQueryJob job = new SQLQueryJob(
                         getSite(),
-                    isSingleQuery ? SQLEditorMessages.editors_sql_job_execute_query : SQLEditorMessages.editors_sql_job_execute_script,
+                        "Export from query",
                         executionContext,
                         resultsContainer,
                         queries,
                         new SQLScriptContext(globalScriptContext, SQLEditor.this, localFile, new OutputLogWriter()),
-                    this,
+                        resultsConsumer,
                         listener);
 
-                if (export || isSingleQuery) {
-                    resultsContainer.query = queries.get(0);
+                    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));
                     }
-                if (export) {
-                    // Assign current job from active query and open wizard
-                    resultsContainer.lastGoodQuery = null;
-                    curJob = job;
+
                     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) {
+                } else {
+                    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();
@@ -2383,6 +2412,7 @@ public class SQLEditor extends SQLEditorBase implements
                     }
                 }
             }
+        }
 
         public boolean isDirty() {
             for (QueryResultsContainer resultsProvider : resultContainers) {

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;
+    }
+}