Add support for including and excluding drives from listing

jupyter_drives/handlers.py

@@ -49,10 +49,20 @@ class ConfigJupyterDrivesHandler(JupyterDrivesAPIHandler):
     def initialize(self, logger: logging.Logger, manager: JupyterDrivesManager):
         return super().initialize(logger, manager)
 
+    @tornado.web.authenticated
+    async def get(self):
+        result = self._manager.get_excluded_drives()
+        self.finish(json.dumps(result["data"]))
+
     @tornado.web.authenticated
     async def post(self):
         body = self.get_json_body()
-        result = self._manager.set_listing_limit(**body)
+        if 'new_limit' in body:
+            result = self._manager.set_listing_limit(**body)
+        if 'exclude_drive_name' in body:
+            result = self._manager.exclude_drive(**body)
+        if 'include_drive_name' in body:
+            result = self._manager.include_drive(**body)
         self.finish(result)
 
 class ListJupyterDrivesHandler(JupyterDrivesAPIHandler):

jupyter_drives/manager.py

@@ -53,6 +53,7 @@ def __init__(self, config: traitlets.config.Config) -> None:
         self._max_files_listed = 1025
         self._drives = None
         self._external_drives = {}
+        self._excluded_drives = set()
 
         # instate fsspec file system
         self._file_system = fsspec.filesystem(self._config.provider, asynchronous=True)
@@ -171,6 +172,65 @@ def set_listing_limit(self, new_limit):
 
         return
 
+    def exclude_drive(self, exclude_drive_name):
+        """Exclude drive from listing.
+
+        Args:
+            exclude_bucket_name: drive to exclude
+        """
+        try:
+            self._excluded_drives.add(exclude_drive_name);
+        except Exception as e:
+            raise tornado.web.HTTPError(
+            status_code= httpx.codes.BAD_REQUEST,
+            reason= f"The following error occured when excluding the drive: {e}"
+            )
+
+        return
+
+    def get_excluded_drives(self):
+        """Get list of excluded drives.
+
+        Returns: 
+            List of excluded drives and their properties.
+        """
+        data = []
+        for drive in self._excluded_drives:
+            try:
+                data.append({
+                    "name": drive,
+                    "region": self._config.region_name if drive not in self._content_managers else self._content_managers[drive]["location"],
+                    "creationDate": datetime.now().isoformat(timespec='milliseconds').replace('+00:00', 'Z'),
+                    "mounted": False if drive not in self._content_managers else True,
+                    "provider": self._config.provider
+                })
+            except Exception as e:
+                raise tornado.web.HTTPError(
+                    status_code=httpx.codes.BAD_REQUEST,
+                    reason=f"The following error occured when listing excluded drives: {e}",
+                )
+
+        response = {
+            "data": data
+        }
+        return response
+
+    def include_drive(self, include_drive_name):
+        """Include drive in listing.
+
+        Args:
+            include_bucket_name: drive to include in listing
+        """
+        try:
+            self._excluded_drives.remove(include_drive_name);
+        except Exception as e:
+            raise tornado.web.HTTPError(
+            status_code= httpx.codes.BAD_REQUEST,
+            reason= f"The following error occured when including the drive: {e}"
+            )
+
+        return
+
     async def list_drives(self): 
         """Get list of available drives.
 
@@ -196,26 +256,28 @@ async def list_drives(self):
                     )
 
             for result in results:
-                data.append(
-                    {
-                        "name": result.name,
-                        "region": self._config.region_name,
-                        "creationDate": result.extra["creation_date"],
-                        "mounted": False if result.name not in self._content_managers else True,
-                        "provider": self._config.provider
-                    }
-                )
+                if result.name not in self._excluded_drives:
+                    data.append(
+                        {
+                            "name": result.name,
+                            "region": self._config.region_name if result.name not in self._content_managers else self._content_managers[result.name]["location"],
+                            "creationDate": result.extra["creation_date"],
+                            "mounted": False if result.name not in self._content_managers else True,
+                            "provider": self._config.provider
+                        }
+                    )
 
             if len(self._external_drives) != 0:
                 for drive in self._external_drives.values():
                     try:
-                        data.append({
-                            "name": drive['url'],
-                            "region": self._config.region_name,
-                            "creationDate": datetime.now().isoformat(timespec='milliseconds').replace('+00:00', 'Z'),
-                            "mounted": False if result.name not in self._content_managers else True,
-                            "provider": self._config.provider
-                        })
+                        if drive['url'] not in self._excluded_drives:
+                            data.append({
+                                "name": drive['url'],
+                                "region": self._config.region_name if drive['url'] not in self._content_managers else self._content_managers[drive['url']]["location"],
+                                "creationDate": datetime.now().isoformat(timespec='milliseconds').replace('+00:00', 'Z'),
+                                "mounted": False if result.name not in self._content_managers else True,
+                                "provider": self._config.provider
+                            })
                     except Exception as e:
                         raise tornado.web.HTTPError(
                             status_code=httpx.codes.BAD_REQUEST,

src/contents.ts

@@ -22,7 +22,9 @@ import {
   copyObjects,
   presignedLink,
   createDrive,
-  getDrivesList
+  getDrivesList,
+  excludeDrive,
+  includeDrive
 } from './requests';
 
 let data: Contents.IModel = {
@@ -683,6 +685,46 @@ export class Drive implements Contents.IDrive {
     return data;
   }
 
+  /**
+   * Exclude drive from browser.
+   *
+   * @param driveName: The name of drive to exclude.
+   *
+   * @returns A promise which resolves with the contents model.
+   */
+  async excludeDrive(driveName: string): Promise<Contents.IModel> {
+    data = await excludeDrive(driveName);
+
+    Contents.validateContentsModel(data);
+    this._fileChanged.emit({
+      type: 'delete',
+      oldValue: data,
+      newValue: null
+    });
+
+    return data;
+  }
+
+  /**
+   * Include drive in browser listing.
+   *
+   * @param driveName: The name of drive to include.
+   *
+   * @returns A promise which resolves with the contents model.
+   */
+  async includeDrive(driveName: string): Promise<Contents.IModel> {
+    data = await includeDrive(driveName);
+
+    Contents.validateContentsModel(data);
+    this._fileChanged.emit({
+      type: 'new',
+      oldValue: null,
+      newValue: data
+    });
+
+    return data;
+  }
+
   /**
    * Create a checkpoint for a file.
    *

src/icons.ts

@@ -1,7 +1,19 @@
 import { LabIcon } from '@jupyterlab/ui-components';
 import driveBrowserSvg from '../style/driveIconFileBrowser.svg';
+import addIconSvg from '../style/addIcon.svg';
+import removeIconSvg from '../style/removeIcon.svg';
 
 export const driveBrowserIcon = new LabIcon({
   name: 'jupyter-drives:drive-browser',
   svgstr: driveBrowserSvg
 });
+
+export const addIcon = new LabIcon({
+  name: 'jupyter-drives:add-drive',
+  svgstr: addIconSvg
+});
+
+export const removeIcon = new LabIcon({
+  name: 'jupyter-drives:remove-drive',
+  svgstr: removeIconSvg
+});

src/plugins/driveBrowserPlugin.ts

@@ -33,7 +33,7 @@ import { PageConfig, PathExt } from '@jupyterlab/coreutils';
 import { CommandRegistry } from '@lumino/commands';
 import { Widget } from '@lumino/widgets';
 
-import { driveBrowserIcon } from '../icons';
+import { addIcon, driveBrowserIcon, removeIcon } from '../icons';
 import { Drive } from '../contents';
 import { setListingLimit } from '../requests';
 import { CommandIDs } from '../token';
@@ -423,7 +423,7 @@ namespace Private {
     });
 
     app.commands.addCommand(CommandIDs.addPublicDrive, {
-      isEnabled: () => {
+      isVisible: () => {
         return browser.model.path === 's3:';
       },
       execute: async () => {
@@ -542,5 +542,65 @@ namespace Private {
       icon: fileIcon.bindprops({ stylesheet: 'menuItem' }),
       label: 'Copy Path'
     });
+
+    app.commands.addCommand(CommandIDs.excludeDrive, {
+      isVisible: () => {
+        return browser.model.path === 's3:';
+      },
+      execute: async () => {
+        const widget = tracker.currentWidget;
+        if (!widget) {
+          return;
+        }
+        const item = widget.selectedItems().next();
+        if (item.done) {
+          return;
+        }
+
+        const driveName: string = item.value.name;
+        await drive.excludeDrive(driveName);
+      },
+      label: 'Exclude Drive',
+      icon: removeIcon.bindprops({ stylesheet: 'menuItem' })
+    });
+
+    app.contextMenu.addItem({
+      command: CommandIDs.excludeDrive,
+      selector:
+        '#drive-file-browser.jp-SidePanel .jp-DirListing-content .jp-DirListing-item[data-isdir]',
+      rank: 110
+    });
+
+    app.commands.addCommand(CommandIDs.includeDrive, {
+      isVisible: () => {
+        return browser.model.path === 's3:';
+      },
+      execute: async () => {
+        return showDialog({
+          title: 'Include Drive',
+          body: new Private.AddPublicDriveHandler(drive.name),
+          focusNodeSelector: 'input',
+          buttons: [
+            Dialog.cancelButton(),
+            Dialog.okButton({
+              label: 'Add',
+              ariaLabel: 'Include Drive'
+            })
+          ]
+        }).then(async result => {
+          if (result.value) {
+            await drive.includeDrive(result.value);
+          }
+        });
+      },
+      label: 'Include Drive',
+      icon: addIcon.bindprops({ stylesheet: 'menuItem' })
+    });
+
+    app.contextMenu.addItem({
+      command: CommandIDs.includeDrive,
+      selector: '#drive-file-browser.jp-SidePanel .jp-DirListing-content',
+      rank: 110
+    });
   }
 }