dbeaver/pro#2743 AI custom scope dialog refactor (#34379)

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

@@ -140,7 +140,7 @@ public class DataTransferTaskConfigurator implements DBTTaskConfigurator, DBTTas
                         group.getShell(),
                         isExport ? DTUIMessages.data_transfer_task_configurator_tables_title_choose_source : DTUIMessages.data_transfer_task_configurator_tables_title_choose_target,
                         rootNode,
-                        selNode,
+                        CommonUtils.singletonOrEmpty(selNode),
                         new Class[]{DBSInstance.class, DBSObjectContainer.class, tableClass},
                         new Class[]{tableClass},
                         null);

plugins/org.jkiss.dbeaver.ext.athena.ui/src/org/jkiss/dbeaver/ext/athena/ui/views/AthenaConnectionPage.java

@@ -129,9 +129,9 @@ public class AthenaConnectionPage extends ConnectionPageWithAuth implements IDia
                         ObjectBrowserDialogBase dialog = new ObjectBrowserDialogBase(
                             s3LocationText.getShell(), "S3 browser",
                             fsRootNode,
-                            selectedNode,
-                            true)
-                        {
+                            CommonUtils.singletonOrEmpty(selectedNode),
+                            true
+                        ) {
                             @Override
                             protected boolean matchesResultNode(DBNNode node) {
                                 return node instanceof DBNPathBase &&

plugins/org.jkiss.dbeaver.ui.editors.sql.ai/src/org/jkiss/dbeaver/ui/editors/sql/ai/controls/ScopeConfigDialog.java

@@ -1,211 +0,0 @@
-/*
- * DBeaver - Universal Database Manager
- * Copyright (C) 2010-2024 DBeaver Corp and others
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *     http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.jkiss.dbeaver.ui.editors.sql.ai.controls;
-
-import org.eclipse.core.runtime.IStatus;
-import org.eclipse.core.runtime.Status;
-import org.eclipse.swt.SWT;
-import org.eclipse.swt.events.SelectionAdapter;
-import org.eclipse.swt.events.SelectionEvent;
-import org.eclipse.swt.layout.GridData;
-import org.eclipse.swt.widgets.Composite;
-import org.eclipse.swt.widgets.Shell;
-import org.eclipse.swt.widgets.Tree;
-import org.eclipse.swt.widgets.TreeItem;
-import org.jkiss.code.NotNull;
-import org.jkiss.dbeaver.DBException;
-import org.jkiss.dbeaver.Log;
-import org.jkiss.dbeaver.model.DBIcon;
-import org.jkiss.dbeaver.model.DBPDataSource;
-import org.jkiss.dbeaver.model.DBUtils;
-import org.jkiss.dbeaver.model.navigator.DBNDatabaseNode;
-import org.jkiss.dbeaver.model.navigator.DBNUtils;
-import org.jkiss.dbeaver.model.runtime.AbstractJob;
-import org.jkiss.dbeaver.model.runtime.DBRProgressMonitor;
-import org.jkiss.dbeaver.model.struct.DBSEntity;
-import org.jkiss.dbeaver.model.struct.DBSObject;
-import org.jkiss.dbeaver.model.struct.DBSObjectContainer;
-import org.jkiss.dbeaver.model.struct.DBSStructContainer;
-import org.jkiss.dbeaver.ui.DBeaverIcons;
-import org.jkiss.dbeaver.ui.UIUtils;
-import org.jkiss.dbeaver.ui.dialogs.BaseDialog;
-import org.jkiss.dbeaver.utils.GeneralUtils;
-import org.jkiss.utils.ArrayUtils;
-
-import java.util.*;
-
-public class ScopeConfigDialog extends BaseDialog {
-    private static final Log log = Log.getLog(ScopeConfigDialog.class);
-
-    private Tree objectTree;
-    private final Set<String> checkedObjectIds;
-    private final DBPDataSource dataSource;
-
-    public ScopeConfigDialog(@NotNull Shell shell, @NotNull Set<String> checkedIds, @NotNull DBPDataSource dataSource) {
-        super(shell, "Customize scope", DBIcon.AI);
-        this.checkedObjectIds = new LinkedHashSet<>(checkedIds);
-        this.dataSource = dataSource;
-    }
-
-    @Override
-    protected Composite createDialogArea(Composite parent) {
-        Composite composite = super.createDialogArea(parent);
-
-        objectTree = new Tree(composite, SWT.CHECK);
-        GridData gd = new GridData(GridData.FILL_BOTH);
-        gd.widthHint = 400;
-        gd.heightHint = 300;
-        objectTree.setLayoutData(gd);
-        objectTree.addSelectionListener(new SelectionAdapter() {
-            @Override
-            public void widgetSelected(SelectionEvent e) {
-                if (e.detail != SWT.CHECK) {
-                    return;
-                }
-                TreeItem item = (TreeItem) e.item;
-                if (item.getData() instanceof DBSStructContainer) {
-                    checkTreeItems(item.getItems(), item.getChecked());
-                }
-            }
-
-            private void checkTreeItems(TreeItem[] items, boolean check) {
-                for (TreeItem child : items) {
-                    child.setChecked(check);
-                    if (!(child.getData() instanceof DBSEntity)) {
-                        TreeItem[] children = child.getItems();
-                        if (!ArrayUtils.isEmpty(children)) {
-                            checkTreeItems(children, check);
-                        }
-                    }
-                }
-            }
-        });
-
-        loadObjects(objectTree, dataSource);
-
-        return composite;
-    }
-
-    @NotNull
-    public Set<String> getCheckedObjectIds() {
-        return checkedObjectIds;
-    }
-
-    private void loadObjects(Tree objectTree, DBPDataSource ds) {
-        new AbstractJob("Load database structure") {
-            @Override
-            protected IStatus run(DBRProgressMonitor monitor) {
-                if (ds instanceof DBSObjectContainer) {
-                    try {
-                        loadContainer(monitor, objectTree, null, (DBSObjectContainer) ds, checkedObjectIds);
-                    } catch (Exception e) {
-                        return GeneralUtils.makeExceptionStatus(e);
-                    }
-                    UIUtils.syncExec(() -> {
-                        if (objectTree != null && !objectTree.isDisposed()) {
-                            for (TreeItem item : objectTree.getItems()) {
-                                item.setExpanded(true);
-                            }
-                        }
-                    });
-                }
-                return Status.OK_STATUS;
-            }
-        }.schedule();
-    }
-
-    private void loadContainer(
-        DBRProgressMonitor monitor,
-        Tree objectTree,
-        TreeItem parentItem,
-        DBSObjectContainer objectContainer,
-        Set<String> checkedObjectIds
-    ) throws DBException {
-        Collection<? extends DBSObject> children;
-        try {
-            children = objectContainer.getChildren(monitor);
-        } catch (Exception e) {
-            log.debug("Error loading container '" + objectContainer.getName() + "' contents: " + e.getMessage());
-            return;
-        }
-        if (children == null) {
-            return;
-        }
-        Map<TreeItem, DBSObjectContainer> addedContainers = new LinkedHashMap<>();
-        UIUtils.syncExec(() -> {
-            for (DBSObject child : children) {
-                if (monitor.isCanceled() || objectTree == null || objectTree.isDisposed()) {
-                    return;
-                }
-                if (!(child instanceof DBSStructContainer) && !(child instanceof DBSEntity)) {
-                    continue;
-                }
-                DBNDatabaseNode node = DBNUtils.getNodeByObject(monitor, child, false);
-                if (node == null) {
-                    continue;
-                }
-                TreeItem item = parentItem == null ?
-                    new TreeItem(objectTree, SWT.NONE) : new TreeItem(parentItem, SWT.NONE);
-                item.setData(child);
-                item.setImage(DBeaverIcons.getImage(node.getNodeIconDefault()));
-                item.setText(node.getNodeDisplayName());
-                String objectId = DBUtils.getObjectFullId(child);
-                if (checkedObjectIds.contains(objectId)) {
-                    item.setChecked(true);
-                    if (parentItem != null && !parentItem.getExpanded()) {
-                        parentItem.setExpanded(true);
-                    }
-                }
-                if (child instanceof DBSObjectContainer) {
-                    addedContainers.put(item, (DBSObjectContainer) child);
-                }
-            }
-        });
-        if (monitor.isCanceled()) {
-            return;
-        }
-        for (Map.Entry<TreeItem, DBSObjectContainer> contItem : addedContainers.entrySet()) {
-            if (monitor.isCanceled()) {
-                break;
-            }
-            DBSObjectContainer object = contItem.getValue();
-            loadContainer(monitor, objectTree, contItem.getKey(), object, checkedObjectIds);
-        }
-    }
-
-    @Override
-    protected void okPressed() {
-        checkedObjectIds.clear();
-        collectCheckedObjects(objectTree.getItems());
-
-        super.okPressed();
-    }
-
-    private void collectCheckedObjects(TreeItem[] items) {
-        for (TreeItem item : items) {
-            if (item.getChecked()) {
-                if (item.getData() instanceof DBSEntity) {
-                    checkedObjectIds.add(DBUtils.getObjectFullId((DBSEntity) item.getData()));
-                }
-            }
-            TreeItem[] children = item.getItems();
-            if (!ArrayUtils.isEmpty(children)) {
-                collectCheckedObjects(children);
-            }
-        }
-    }
-}

plugins/org.jkiss.dbeaver.ui.editors.sql.ai/src/org/jkiss/dbeaver/ui/editors/sql/ai/controls/ScopeSelectorControl.java

@@ -16,7 +16,6 @@
  */
 package org.jkiss.dbeaver.ui.editors.sql.ai.controls;
 
-import org.eclipse.jface.dialogs.IDialogConstants;
 import org.eclipse.jface.layout.GridLayoutFactory;
 import org.eclipse.swt.SWT;
 import org.eclipse.swt.events.SelectionAdapter;
@@ -25,6 +24,7 @@ import org.eclipse.swt.events.SelectionListener;
 import org.eclipse.swt.layout.GridData;
 import org.eclipse.swt.widgets.*;
 import org.jkiss.code.NotNull;
+import org.jkiss.code.Nullable;
 import org.jkiss.dbeaver.DBException;
 import org.jkiss.dbeaver.Log;
 import org.jkiss.dbeaver.model.DBPDataSource;
@@ -34,15 +34,24 @@ import org.jkiss.dbeaver.model.ai.completion.DAICompletionSettings;
 import org.jkiss.dbeaver.model.app.DBPProject;
 import org.jkiss.dbeaver.model.exec.DBCExecutionContext;
 import org.jkiss.dbeaver.model.logical.DBSLogicalDataSource;
+import org.jkiss.dbeaver.model.navigator.DBNDatabaseNode;
+import org.jkiss.dbeaver.model.navigator.DBNModel;
+import org.jkiss.dbeaver.model.navigator.DBNNode;
 import org.jkiss.dbeaver.model.runtime.DBRProgressMonitor;
+import org.jkiss.dbeaver.model.runtime.DBRRunnableContext;
 import org.jkiss.dbeaver.model.struct.DBSEntity;
+import org.jkiss.dbeaver.model.struct.DBSInstance;
+import org.jkiss.dbeaver.model.struct.DBSObjectContainer;
 import org.jkiss.dbeaver.ui.UIIcon;
 import org.jkiss.dbeaver.ui.UIUtils;
+import org.jkiss.dbeaver.ui.navigator.dialogs.ObjectBrowserDialog;
 import org.jkiss.utils.ArrayUtils;
 import org.jkiss.utils.CommonUtils;
 
+import java.lang.reflect.InvocationTargetException;
 import java.util.List;
 import java.util.*;
+import java.util.stream.Collectors;
 
 public class ScopeSelectorControl extends Composite {
     private static final Log log = Log.getLog(ScopeSelectorControl.class);
@@ -184,6 +193,53 @@ public class ScopeSelectorControl extends Composite {
         layout(true, true);
     }
 
+    @Nullable
+    public static Set<String> chooseCustomEntities(
+        @NotNull Shell shell,
+        @NotNull DBRRunnableContext context,
+        @NotNull DBPDataSource dataSource,
+        @NotNull Set<String> ids
+    ) {
+        DBNModel navigator = Objects.requireNonNull(dataSource.getContainer().getProject().getNavigatorModel());
+        List<DBNDatabaseNode> nodes = new ArrayList<>();
+
+        try {
+            // Find nodes of already selected objects
+            context.run(true, true, monitor -> {
+                for (DBSEntity entity : loadCustomEntities(monitor, dataSource, ids)) {
+                    DBNDatabaseNode node = navigator.getNodeByObject(monitor, entity, true);
+                    if (node != null) {
+                        nodes.add(node);
+                    }
+                }
+            });
+        } catch (InvocationTargetException | InterruptedException e) {
+            log.warn("Error loading custom entities", e);
+        }
+
+        // Select custom objects
+        List<DBNNode> selected = ObjectBrowserDialog.selectObjects(
+            shell,
+            "Select objects to include in completion scope",
+            navigator.getNodeByObject(dataSource),
+            nodes,
+            new Class[]{DBSInstance.class, DBSObjectContainer.class, DBSEntity.class},
+            new Class[]{DBSEntity.class},
+            new Class[]{DBSEntity.class}
+        );
+
+        if (selected == null) {
+            return null;
+        }
+
+        return selected.stream()
+            .map(DBNDatabaseNode.class::cast)
+            .map(DBNDatabaseNode::getValueObject)
+            .map(DBSEntity.class::cast)
+            .map(DBUtils::getObjectFullId)
+            .collect(Collectors.toSet());
+    }
+
     @NotNull
     public static List<DBSEntity> loadCustomEntities(
         @NotNull DBRProgressMonitor monitor,
@@ -221,13 +277,19 @@ public class ScopeSelectorControl extends Composite {
 
     public void changeScope(@NotNull DAICompletionScope scope) {
         if (scope == DAICompletionScope.CUSTOM) {
-            final ScopeConfigDialog dialog = new ScopeConfigDialog(getShell(), checkedObjectIds, executionContext.getDataSource());
-            if (dialog.open() != IDialogConstants.OK_ID) {
+            Set<String> ids = chooseCustomEntities(
+                getShell(),
+                UIUtils.getDefaultRunnableContext(),
+                executionContext.getDataSource(),
+                checkedObjectIds
+            );
+
+            if (ids == null) {
                 return;
             }
 
             checkedObjectIds.clear();
-            checkedObjectIds.addAll(dialog.getCheckedObjectIds());
+            checkedObjectIds.addAll(ids);
         }
 
         currentScope = scope;

plugins/org.jkiss.dbeaver.ui.navigator/src/org/jkiss/dbeaver/ui/navigator/dialogs/ObjectBrowserDialog.java

@@ -20,6 +20,8 @@ import org.eclipse.jface.dialogs.IDialogConstants;
 import org.eclipse.jface.viewers.Viewer;
 import org.eclipse.jface.viewers.ViewerFilter;
 import org.eclipse.swt.widgets.Shell;
+import org.jkiss.code.NotNull;
+import org.jkiss.code.Nullable;
 import org.jkiss.dbeaver.model.navigator.*;
 import org.jkiss.dbeaver.model.navigator.fs.DBNFileSystem;
 import org.jkiss.dbeaver.model.navigator.fs.DBNPathBase;
@@ -29,6 +31,7 @@ import org.jkiss.dbeaver.model.struct.DBSObjectContainer;
 import org.jkiss.dbeaver.model.struct.DBSWrapper;
 import org.jkiss.dbeaver.ui.navigator.database.DatabaseNavigatorTreeFilter;
 import org.jkiss.dbeaver.ui.navigator.database.load.TreeNodeSpecial;
+import org.jkiss.utils.CommonUtils;
 
 import java.util.List;
 
@@ -38,21 +41,21 @@ import java.util.List;
  * @author Serge Rider
  */
 public class ObjectBrowserDialog extends ObjectBrowserDialogBase {
-    private Class<?>[] allowedTypes;
-    private Class<?>[] resultTypes;
-    private Class<?>[] leafTypes;
+    private final Class<?>[] allowedTypes;
+    private final Class<?>[] resultTypes;
+    private final Class<?>[] leafTypes;
 
     private ObjectBrowserDialog(
-        Shell parentShell,
-        String title,
-        DBNNode rootNode,
-        DBNNode selectedNode,
+        @NotNull Shell parentShell,
+        @NotNull String title,
+        @NotNull DBNNode rootNode,
+        @NotNull List<? extends DBNNode> selectedNodes,
         boolean singleSelection,
-        Class<?>[] allowedTypes,
-        Class<?>[] resultTypes,
-        Class<?>[] leafTypes)
-    {
-        super(parentShell, title, rootNode, selectedNode, singleSelection);
+        @NotNull Class<?>[] allowedTypes,
+        @Nullable Class<?>[] resultTypes,
+        @Nullable Class<?>[] leafTypes
+    ) {
+        super(parentShell, title, rootNode, selectedNodes, singleSelection);
         this.allowedTypes = allowedTypes;
         this.resultTypes = resultTypes == null ? allowedTypes : resultTypes;
         this.leafTypes = leafTypes;
@@ -147,9 +150,26 @@ public class ObjectBrowserDialog extends ObjectBrowserDialogBase {
         return false;
     }
 
-    public static DBNNode selectObject(Shell parentShell, String title, DBNNode rootNode, DBNNode selectedNode, Class<?>[] allowedTypes, Class<?>[] resultTypes, Class<?>[] leafTypes)
-    {
-        ObjectBrowserDialog scDialog = new ObjectBrowserDialog(parentShell, title, rootNode, selectedNode, true, allowedTypes, resultTypes, leafTypes);
+    @Nullable
+    public static DBNNode selectObject(
+        @NotNull Shell parentShell,
+        @NotNull String title,
+        @NotNull DBNNode rootNode,
+        @Nullable DBNNode selectedNode,
+        @NotNull Class<?>[] allowedTypes,
+        @Nullable Class<?>[] resultTypes,
+        @Nullable Class<?>[] leafTypes
+    ) {
+        ObjectBrowserDialog scDialog = new ObjectBrowserDialog(
+            parentShell,
+            title,
+            rootNode,
+            CommonUtils.singletonOrEmpty(selectedNode),
+            true,
+            allowedTypes,
+            resultTypes,
+            leafTypes
+        );
         if (scDialog.open() == IDialogConstants.OK_ID) {
             List<DBNNode> result = scDialog.getSelectedObjects();
             return result.isEmpty() ? null : result.get(0);
@@ -158,9 +178,27 @@ public class ObjectBrowserDialog extends ObjectBrowserDialogBase {
         }
     }
 
-    public static List<DBNNode> selectObjects(Shell parentShell, String title, DBNNode rootNode, DBNNode selectedNode, Class<?>[] allowedTypes, Class<?>[] resultTypes, Class<?>[] leafTypes)
+    @Nullable
+    public static List<DBNNode> selectObjects(
+        @NotNull Shell parentShell,
+        @NotNull String title,
+        @NotNull DBNNode rootNode,
+        @NotNull List<? extends DBNNode> selectedNodes,
+        @NotNull Class<?>[] allowedTypes,
+        @Nullable Class<?>[] resultTypes,
+        @Nullable Class<?>[] leafTypes
+    )
     {
-        ObjectBrowserDialog scDialog = new ObjectBrowserDialog(parentShell, title, rootNode, selectedNode, false, allowedTypes, resultTypes, leafTypes);
+        ObjectBrowserDialog scDialog = new ObjectBrowserDialog(
+            parentShell,
+            title,
+            rootNode,
+            selectedNodes,
+            false,
+            allowedTypes,
+            resultTypes,
+            leafTypes
+        );
         if (scDialog.open() == IDialogConstants.OK_ID) {
             return scDialog.getSelectedObjects();
         } else {

plugins/org.jkiss.dbeaver.ui.navigator/src/org/jkiss/dbeaver/ui/navigator/dialogs/ObjectBrowserDialogBase.java

@@ -30,6 +30,7 @@ import org.eclipse.swt.widgets.Button;
 import org.eclipse.swt.widgets.Composite;
 import org.eclipse.swt.widgets.Control;
 import org.eclipse.swt.widgets.Shell;
+import org.jkiss.code.NotNull;
 import org.jkiss.dbeaver.model.DBPDataSourceContainer;
 import org.jkiss.dbeaver.model.DBUtils;
 import org.jkiss.dbeaver.model.navigator.DBNContainer;
@@ -45,6 +46,7 @@ import org.jkiss.dbeaver.ui.navigator.database.DatabaseNavigatorTreeFilterObject
 import org.jkiss.dbeaver.ui.navigator.database.load.TreeNodeSpecial;
 
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
 
 /**
@@ -58,7 +60,7 @@ public abstract class ObjectBrowserDialogBase extends Dialog {
 
     private final String title;
     private final DBNNode rootNode;
-    private final DBNNode selectedNode;
+    private final DBNNode[] selectedNodes;
     private final boolean singleSelection;
     private final List<DBNNode> selectedObjects = new ArrayList<>();
     private TreeNodeSpecial specialNode;
@@ -67,16 +69,16 @@ public abstract class ObjectBrowserDialogBase extends Dialog {
     private static boolean showConnected;
 
     protected ObjectBrowserDialogBase(
-        Shell parentShell,
-        String title,
-        DBNNode rootNode,
-        DBNNode selectedNode,
-        boolean singleSelection)
-    {
+        @NotNull Shell parentShell,
+        @NotNull String title,
+        @NotNull DBNNode rootNode,
+        @NotNull List<? extends DBNNode> selectedNodes,
+        boolean singleSelection
+    ) {
         super(parentShell);
         this.title = title;
         this.rootNode = rootNode;
-        this.selectedNode = selectedNode;
+        this.selectedNodes = selectedNodes.toArray(DBNNode[]::new);
         this.singleSelection = singleSelection;
     }
 
@@ -113,12 +115,15 @@ public abstract class ObjectBrowserDialogBase extends Dialog {
         if (viewerFilter != null) {
             treeViewer.addFilter(viewerFilter);
         }
-        if (selectedNode != null) {
-            treeViewer.setSelection(new StructuredSelection(selectedNode));
-            if (!(selectedNode instanceof DBNDataSource) || ((DBNDataSource) selectedNode).getDataSourceContainer().isConnected()) {
-                treeViewer.expandToLevel(selectedNode, 1);
+        if (selectedNodes.length > 0) {
+            treeViewer.setSelection(new StructuredSelection(selectedNodes));
+            Collections.addAll(selectedObjects, selectedNodes);
+
+            for (DBNNode node : selectedNodes) {
+                if (!(node instanceof DBNDataSource dataSource) || dataSource.getDataSourceContainer().isConnected()) {
+                    treeViewer.expandToLevel(selectedNodes, 1);
+                }
             }
-            selectedObjects.add(selectedNode);
         }
         treeViewer.addSelectionChangedListener(event -> {
             selectedObjects.clear();
@@ -177,15 +182,12 @@ public abstract class ObjectBrowserDialogBase extends Dialog {
                 return true;
             }
             object = DBUtils.getAdapter(DBSObject.class, object);
-            if (object != null && matchesType(object, true)) {
-                return true;
-            }
+            return object != null && matchesType(object, true);
         } else if (node instanceof DBNObjectNode) {
             return matchesType(((DBNObjectNode) node).getNodeObject(), true);
         } else {
             return matchesType(node, true);
         }
-        return false;
     }
 
     protected boolean matchesType(Object object, boolean result) {