Merge pull request #90 from DenisaCG/include_exclude_config

Update backend configuration for including and excluding buckets

Author: afshin

README.md

@@ -54,6 +54,8 @@ c.DrivesConfig.access_key_id = "<Drives Access Key ID / IAM Access Key ID>"
 c.DrivesConfig.secret_access_key = "<Drives Secret Access Key / IAM Secret>"
 c.DrivesConfig.session_token = "<Drives Session Token / IAM Session Token (optional)>"
 c.DrivesConfig.provider = "<Drives provider e.g.: s3, gcs>"
+c.DrivesConfig.included_drives = "<List of the drives to be shown in the drives filebrowser (optional)>"
+c.DrivesConfig.excluded_drives = "<List of the drives to be excluded from the drives filebrowser listing (optional)>"
 ```
 
 ### Custom credentials file path
@@ -76,6 +78,8 @@ export JP_DRIVES_ACCESS_KEY_ID="<Drives Access Key ID>"
 export JP_DRIVES_SECRET_ACCESS_KEY="<Drives Secret Access Key>"
 export JP_DRIVES_SESSION_TOKEN="<Drives Session Token (optional)>"
 export JP_DRIVES_CUSTOM_CREDENTIALS_PATH="<Path to local file which contains credentials (optional)>"
+export JP_DRIVES_INCLUDED_DRIVES ="<Names of the drives to be shown in the drives filebrowser, separated by spaces (optional)>"
+export JP_DRIVES_EXCLUDED_DRIVES="<Names of the drives to be excluded from the drives filebrowser listing, separated by spaces (optional)>"
 ```
 
 ## Uninstall

jupyter_drives/base.py

@@ -1,7 +1,7 @@
 import os
 from sys import platform
 import entrypoints
-from traitlets import Enum, Unicode, default
+from traitlets import Enum, Unicode, default, Set
 from traitlets.config import Configurable
 import boto3
 
@@ -71,10 +71,27 @@ def set_default_api_base_url(self):
         help="The source control provider.",
     )
 
+    excluded_drives = Set(
+        trait=Unicode(),
+        config=True,
+        help="List of drives that should be excluded from drive browser listing. Drive names should be separated by spaces.",
+    )
+
+    included_drives = Set(
+        trait=Unicode(),
+        config=True,
+        help="List of drives that should be included in drive browser listing. Drive names should be separated by spaces.",
+    )
+
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
         # check if credentials were already set in jupyter_notebook_config.py
         self.credentials_already_set = self.access_key_id is not None and self.secret_access_key is not None
+
+        # check list of excluded and included drives
+        self.check_excluded_and_included_drives()
+
+        # load credentials
         self.load_credentials()
 
     def load_credentials(self):
@@ -111,4 +128,26 @@ def load_credentials(self):
         self.access_key_id = c.access_key
         self.secret_access_key = c.secret_key
         self.region_name = s.region_name
-        self.session_token = c.token
+        self.session_token = c.token
+
+    def check_excluded_and_included_drives(self):
+        # list of drives to exclude was provided
+        if len(self.excluded_drives) != 0: 
+            return
+        
+        # list of what drives to include was provided
+        if len(self.included_drives) != 0: 
+            return
+        
+        # check environment variables for excluded or included drives configuration, only one is taken into account
+        if "JP_DRIVES_EXCLUDED_DRIVES" in os.environ:
+            drives =  os.environ["JP_DRIVES_EXCLUDED_DRIVES"]
+            self.excluded_drives = drives.split(' ')
+            return
+        
+        if "JP_DRIVES_INCLUDED_DRIVES" in os.environ:
+            drives =  os.environ["JP_DRIVES_INCLUDED_DRIVES"]
+            self.included_drives = drives.split(' ')
+            return
+
+        return

jupyter_drives/manager.py

@@ -53,8 +53,9 @@ def __init__(self, config: traitlets.config.Config) -> None:
         self._max_files_listed = 1025
         self._drives = None
         self._external_drives = {}
-        self._excluded_drives = set()
-        
+        self._excluded_drives = self._config.excluded_drives if len(self._config.excluded_drives) != 0 else set()
+        self._included_drives =  self._config.included_drives if len(self._config.included_drives) != 0 else set()
+
         # instate fsspec file system
         self._file_system = fsspec.filesystem(self._config.provider, asynchronous=True)
 
@@ -255,7 +256,14 @@ async def list_drives(self):
                         reason=f"The following error occured when listing drives: {e}",
                     )
 
-            for result in results:
+            if len(self._included_drives) != 0:
+                for result in results: 
+                    if result.name not in self._included_drives:
+                        self._excluded_drives.add(result.name)
+                # clear list once initialized
+                self._included_drives.clear()
+            
+            for result in results:        
                 if result.name not in self._excluded_drives:
                     data.append(
                         {