#9941 Allow to select database for query during task creation

Former-commit-id: 40135350

plugins/org.jkiss.dbeaver.data.transfer.ui/src/org/jkiss/dbeaver/tools/transfer/ui/wizard/DataTransferTaskConfigurator.java

@@ -17,7 +17,6 @@
 
 package org.jkiss.dbeaver.tools.transfer.ui.wizard;
 
-import org.eclipse.jface.dialogs.IDialogConstants;
 import org.eclipse.osgi.util.NLS;
 import org.eclipse.swt.SWT;
 import org.eclipse.swt.events.SelectionAdapter;
@@ -29,20 +28,21 @@ import org.jkiss.dbeaver.DBException;
 import org.jkiss.dbeaver.Log;
 import org.jkiss.dbeaver.model.DBPDataSource;
 import org.jkiss.dbeaver.model.DBPDataSourceContainer;
+import org.jkiss.dbeaver.model.DBPObject;
+import org.jkiss.dbeaver.model.DBUtils;
 import org.jkiss.dbeaver.model.app.DBPProject;
+import org.jkiss.dbeaver.model.data.DBDCollection;
 import org.jkiss.dbeaver.model.impl.DataSourceContextProvider;
-import org.jkiss.dbeaver.model.navigator.DBNDatabaseNode;
-import org.jkiss.dbeaver.model.navigator.DBNNode;
-import org.jkiss.dbeaver.model.navigator.DBNProjectDatabases;
+import org.jkiss.dbeaver.model.navigator.*;
 import org.jkiss.dbeaver.model.runtime.DBRRunnableContext;
 import org.jkiss.dbeaver.model.sql.SQLQuery;
+import org.jkiss.dbeaver.model.sql.SQLScriptContext;
 import org.jkiss.dbeaver.model.sql.data.SQLQueryDataContainer;
-import org.jkiss.dbeaver.model.struct.DBSDataContainer;
-import org.jkiss.dbeaver.model.struct.DBSDataManipulator;
-import org.jkiss.dbeaver.model.struct.DBSObject;
-import org.jkiss.dbeaver.model.struct.DBSObjectContainer;
+import org.jkiss.dbeaver.model.struct.*;
+import org.jkiss.dbeaver.model.struct.rdb.DBSCatalog;
 import org.jkiss.dbeaver.model.task.DBTTask;
 import org.jkiss.dbeaver.model.task.DBTTaskType;
+import org.jkiss.dbeaver.model.task.DBTaskUtils;
 import org.jkiss.dbeaver.runtime.DBWorkbench;
 import org.jkiss.dbeaver.runtime.ui.UIServiceSQL;
 import org.jkiss.dbeaver.tasks.ui.DBTTaskConfigPanel;
@@ -60,9 +60,9 @@ import org.jkiss.dbeaver.ui.UIIcon;
 import org.jkiss.dbeaver.ui.UIUtils;
 import org.jkiss.dbeaver.ui.navigator.NavigatorUtils;
 import org.jkiss.dbeaver.ui.navigator.dialogs.ObjectBrowserDialog;
-import org.jkiss.dbeaver.ui.navigator.dialogs.SelectDataSourceDialog;
 import org.jkiss.utils.CommonUtils;
 
+import java.io.PrintWriter;
 import java.lang.reflect.InvocationTargetException;
 import java.util.ArrayList;
 import java.util.List;
@@ -163,19 +163,27 @@ public class DataTransferTaskConfigurator implements DBTTaskConfigurator {
                 UIUtils.createDialogButton(buttonsPanel, DTUIMessages.data_transfer_task_configurator_dialog_button_label_add_query, new SelectionAdapter() {
                     @Override
                     public void widgetSelected(SelectionEvent e) {
-                        DBPDataSource lastDataSource = getLastDataSource();
-                        if (lastDataSource == null) {
-                            SelectDataSourceDialog dsDialog = new SelectDataSourceDialog(group.getShell(), currentProject, null);
-                            if (dsDialog.open() == IDialogConstants.OK_ID) {
-                                DBPDataSourceContainer dataSource = dsDialog.getDataSource();
-                                if (dataSource == null) {
-                                    return;
-                                }
-                                if (!dataSource.isConnected()) {
+                        DBSObject dataSourceObject = null;
+                        DBPDataSource dataSource = null;
+
+                        DBNProjectDatabases rootNode = DBWorkbench.getPlatform().getNavigatorModel().getRoot().getProjectNode(currentProject).getDatabases();
+                        DBNNode node = ObjectBrowserDialog.selectObject(
+                                group.getShell(),
+                                DTUIMessages.data_transfer_task_configurator_tables_title_choose_source,
+                                rootNode,
+                                rootNode,
+                                new Class[]{DBPDataSource.class, DBSObjectContainer.class, DBSCatalog.class},
+                                null,
+                                new Class[]{DBPDataSource.class, DBSObjectContainer.class, DBSCatalog.class});
+
+                        if (node != null) {
+                            if (node instanceof DBNDataSource) {
+                                DBPDataSourceContainer container = ((DBNDataSource) node).getDataSourceContainer();
+                                if (!container.isConnected()) {
                                     try {
                                         runnableContext.run(true, true, monitor -> {
                                             try {
-                                                dataSource.connect(monitor, true, true);
+                                                container.connect(monitor, true, true);
                                             } catch (DBException ex) {
                                                 throw new InvocationTargetException(ex);
                                             }
@@ -188,22 +196,26 @@ public class DataTransferTaskConfigurator implements DBTTaskConfigurator {
                                         return;
                                     }
                                 }
-                                lastDataSource = dataSource.getDataSource();
+                                dataSourceObject = container;
+                                dataSource = container.getDataSource();
+                            } else if (node instanceof DBNDatabaseItem) {
+                                dataSourceObject = ((DBNDatabaseItem) node).getObject();
+                                dataSource = dataSourceObject.getDataSource();
+                            } else {
+                                log.debug("Unhandled node type: " + node);
+                                return;
                             }
                         }
 
-                        if (lastDataSource != null) {
+                        if (dataSource != null) {
                             UIServiceSQL serviceSQL = DBWorkbench.getService(UIServiceSQL.class);
                             if (serviceSQL != null) {
-                                DataSourceContextProvider contextProvider = new DataSourceContextProvider(lastDataSource);
+                                DataSourceContextProvider contextProvider = new DataSourceContextProvider(dataSourceObject);
                                 String query = serviceSQL.openSQLEditor(contextProvider, DTUIMessages.data_transfer_task_configurator_sql_query_title, UIIcon.SQL_SCRIPT, "");
                                 if (query != null) {
-
-                                    DataTransferPipe pipe = new DataTransferPipe(
-                                        new DatabaseTransferProducer(
-                                            new SQLQueryDataContainer(contextProvider, new SQLQuery(lastDataSource, query), null, log)),
-                                        null);
-
+                                    SQLScriptContext scriptContext = new SQLScriptContext(null, contextProvider, null, new PrintWriter(System.err, true), null);
+                                    SQLQueryDataContainer container = new SQLQueryDataContainer(contextProvider, new SQLQuery(dataSource, query), scriptContext, log);
+                                    DataTransferPipe pipe = new DataTransferPipe(new DatabaseTransferProducer(container), null);
                                     addPipeToTable(pipe);
                                     updateSettings(propertyChangeListener);
                                 }