Merge pull request #77 from the-virtual-brain/EBR-79

EBR-79: add access to bucket from widgets

notebooks/HeadWidget.ipynb

@@ -83,7 +83,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "config = HeadWidgetConfig(name='Cortex', color='beige')\n",
+    "config = HeadWidgetConfig(name='Cortex', color='#F5F5DC')\n",
     "wid.add_datatype(surface, config)\n",
     "config.__dict__"
    ]
@@ -92,7 +92,9 @@
    "cell_type": "code",
    "execution_count": null,
    "id": "0b6c41e9-8a03-4996-929f-3ebbafadcc49",
-   "metadata": {},
+   "metadata": {
+    "scrolled": true
+   },
    "outputs": [],
    "source": [
     "display(wid)"
@@ -172,6 +174,14 @@
    "metadata": {},
    "outputs": [],
    "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "37c413d5-23ea-4954-9eba-cfc08ce8327a",
+   "metadata": {},
+   "outputs": [],
+   "source": []
   }
  ],
  "metadata": {
@@ -190,7 +200,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.16"
+   "version": "3.9.19"
   }
  },
  "nbformat": 4,

requirements.txt

@@ -10,6 +10,7 @@ plotly==5.14.0
 pythreejs
 pyvista>=0.43.0
 ipyreact>=0.4.1
+requests
 trame
 trame-vuetify
 trame-vtk

tvbwidgets/core/bucket/bucket_api.py

@@ -0,0 +1,16 @@
+# -*- coding: utf-8 -*-
+#
+# "TheVirtualBrain - Widgets" package
+#
+# (c) 2022-2024, TVB Widgets Team
+#
+
+from ebrains_drive import BucketApiClient
+from tvbwidgets.core.bucket.buckets import ExtendedBuckets
+
+
+class ExtendedBucketApiClient(BucketApiClient):
+
+    def __init__(self, username=None, password=None, token=None, env="") -> None:
+        super().__init__(username, password, token, env)
+        self.buckets = ExtendedBuckets(self)

tvbwidgets/core/bucket/buckets.py

@@ -0,0 +1,46 @@
+# -*- coding: utf-8 -*-
+#
+# "TheVirtualBrain - Widgets" package
+#
+# (c) 2022-2024, TVB Widgets Team
+#
+
+from dataclasses import dataclass
+from typing import List
+from ebrains_drive.buckets import Buckets
+from tvbwidgets.core.exceptions import BucketDTOError
+from tvbwidgets.core.logger.builder import get_logger
+
+LOGGER = get_logger(__name__)
+
+
+@dataclass
+class BucketDTO:
+    name: str
+    role: str
+    is_public: bool
+
+
+class ExtendedBuckets(Buckets):
+    BUCKETS_ENDPOINT = '/v1/buckets'
+
+    def __init__(self, client):
+        super().__init__(client)
+        self._available_buckets: List[BucketDTO] = []
+
+    def list_buckets(self):
+        # type: () -> List[BucketDTO]
+        """
+        Queries the buckets endpoint for the available buckets for current user
+        """
+        try:
+            resp = self.client.get(self.BUCKETS_ENDPOINT)
+            json_resp = resp.json()
+            updated_available_buckets = []
+            for obj in json_resp:
+                updated_available_buckets.append(BucketDTO(**obj))
+            self._available_buckets = updated_available_buckets
+            return self._available_buckets
+        except KeyError as e:
+            LOGGER.error(f'Received unexpected Bucket structure! {str(e)}')
+            raise BucketDTOError('Unexpected response structure from server!')

tvbwidgets/core/exceptions.py

@@ -37,3 +37,9 @@ class ModelExporterNotFoundError(WidgetsException):
     """
     To be thrown when an attempt is being made to use an exporter that is not accessible
     """
+
+
+class BucketDTOError(WidgetsException):
+    """
+    Exception on bucket DTOs
+    """

tvbwidgets/core/hpc/config.py

@@ -19,16 +19,16 @@ class HPCConfig(object):
     resources: Resources
     timeout = -1
 
-    STORAGES = {'DAINT-CSCS': 'HOME',
-                'JUSUF': 'PROJECT',
-                'JUDAC': 'PROJECT'}
-
-    PYTHON_DIRS = {'DAINT-CSCS': 'python3.9',
-                   'JUSUF': 'python3.10',
-                   'JUDAC': 'python3.10'}
-    MODULES = {'DAINT-CSCS': 'cray-python',
-               'JUSUF': 'Python',
-               'JUDAC': 'Python'}
+    STORAGES = {'JUSUF': 'PROJECT',
+                'JUDAC': 'PROJECT',
+                'JUWELS': 'HOME'}
+
+    PYTHON_DIRS = {'JUSUF': 'python3.10',
+                   'JUDAC': 'python3.10',
+                   'JUWELS': 'python3.11'}
+    MODULES = {'JUSUF': 'Python',
+               'JUDAC': 'Python',
+               'JUWELS': 'Python'}
 
     def __post_init__(self):
         if self.env_dir is None:

tvbwidgets/tests/test_bucket_widget.py

@@ -0,0 +1,104 @@
+# -*- coding: utf-8 -*-
+#
+# "TheVirtualBrain - Widgets" package
+#
+# (c) 2022-2024, TVB Widgets Team
+#
+import os
+
+import pytest
+from ebrains_drive.exceptions import Unauthorized
+from tvbwidgets.core.auth import CLB_AUTH
+from tvbwidgets.ui.bucket_widget import BucketWidget
+
+DUMMY_CONTENT = b'test content'
+
+
+class MockBucketDTO:
+    def __init__(self, name, role='', is_public=True):
+        self.name = name
+        self.role = role
+        self.is_public = is_public
+
+
+class MockFile:
+    def __init__(self, name):
+        # type: (str) -> None
+        self.name = name
+
+    def get_content(self):
+        return DUMMY_CONTENT
+
+    def get_download_link(self):
+        return ''
+
+
+class MockBucket:
+    def __init__(self, files_count=2, name='test_bucket', target='buckets', dataproxy_entity_name='test_bucket'):
+        self.name = name
+        self.files = [MockFile(f'file{number}') for number in range(files_count)]
+        self.target = target
+        self.dataproxy_entity_name = dataproxy_entity_name
+
+    def ls(self, prefix=''):
+        return [f for f in self.files if f.name.startswith(prefix)]
+
+
+class MockBuckets:
+    def __init__(self):
+        self.buckets = {
+            'test_bucket': MockBucket()
+        }
+
+    def get_bucket(self, name):
+        try:
+            return self.buckets[name]
+        except KeyError:
+            raise Unauthorized('Unauthorized in tests')
+
+    def list_buckets(self):
+        return [MockBucketDTO(b) for b in self.buckets.keys()]
+
+
+class MockBucketApiClient:
+    def __init__(self, token=''):
+        self.token = token
+        self.buckets = MockBuckets()
+
+
+@pytest.fixture
+def mock_client(mocker):
+    return mocker.patch('tvbwidgets.ui.bucket_widget.ExtendedBucketApiClient', MockBucketApiClient)
+
+
+@pytest.fixture
+def mock_requests_get(mocker):
+    mock_response = mocker.Mock()
+    mock_response.content = DUMMY_CONTENT
+    return mocker.patch('requests.get', return_value=mock_response)
+
+
+def test_get_files_in_bucket(mock_client, mock_requests_get):
+    """
+    tests that client returns list of files from bucket
+    """
+    if os.environ.get(CLB_AUTH):
+        os.environ.pop(CLB_AUTH)
+
+    with pytest.raises(RuntimeError):
+        BucketWidget()
+
+    os.environ[CLB_AUTH] = "test_auth_token"
+    widget = BucketWidget()
+
+    # test observe event on buckets dropdown
+    assert widget.buckets_dropdown.value is None
+    assert widget.files_list.value is None
+    assert len(widget.files_list.options) == 0
+    widget.buckets_dropdown.value = widget.buckets_dropdown.options[0]
+    assert len(widget.files_list.options) == 2
+    widget.files_list.value = widget.files_list.options[0]
+
+    # test BucketWidget functions
+    assert widget.get_selected_file_path() == widget.buckets_dropdown.value + '/' + widget.files_list.value
+    assert widget.get_selected_file_content() == DUMMY_CONTENT

tvbwidgets/tests/test_drive_widget.py

@@ -122,8 +122,8 @@ def mockk(token):
     os.environ[CLB_AUTH] = "test_auth_token"
     widget = StorageWidget()
 
-    widget.api.repos_dropdown.value = widget.api.repos_dropdown.options['repo1']
-    widget.api.files_list.value = widget.api.files_list.options[1]
+    widget.drive_api.repos_dropdown.value = widget.drive_api.repos_dropdown.options['repo1']
+    widget.drive_api.files_list.value = widget.drive_api.files_list.options[1]
 
     assert widget.get_selected_file_name() == DUMMY_FILENAME
     assert widget.get_selected_file_content() == DUMMY_CONTENT

tvbwidgets/tests/test_head_widget.py

@@ -133,8 +133,8 @@ def mockk(token):
     widget._TVBWidgetWithBrowser__display_message('ABC')
     assert widget.message_label.value == HeadBrowser.MSG_TEMPLATE.format('ABC', HeadBrowser.MSG_COLOR)
 
-    widget.storage_widget.api.repos_dropdown.value = widget.storage_widget.api.repos_dropdown.options['repo1']
-    widget.storage_widget.api.files_list.value = widget.storage_widget.api.files_list.options[1]
+    widget.storage_widget.drive_api.repos_dropdown.value = widget.storage_widget.drive_api.repos_dropdown.options['repo1']
+    widget.storage_widget.drive_api.files_list.value = widget.storage_widget.drive_api.files_list.options[1]
 
     widget.load_selected_file(Surface)
     assert 'Only .zip' in widget.message_label.value

tvbwidgets/ui/bucket_widget.py

@@ -0,0 +1,68 @@
+# -*- coding: utf-8 -*-
+#
+# "TheVirtualBrain - Widgets" package
+#
+# (c) 2022-2024, TVB Widgets Team
+#
+
+import ipywidgets
+import requests
+from tvbwidgets.core.auth import get_current_token
+from tvbwidgets.core.bucket.bucket_api import ExtendedBucketApiClient
+from tvbwidgets.ui.base_widget import TVBWidget
+from ebrains_drive.files import DataproxyFile
+
+
+class BucketWidget(ipywidgets.VBox, TVBWidget):
+
+    def __init__(self, **kwargs):
+        TVBWidget.__init__(self, **kwargs)
+        bearer_token = get_current_token()
+        self.client = ExtendedBucketApiClient(token=bearer_token)
+
+        try:
+            list_buckets = self.client.buckets.list_buckets()
+            buckets_name = [b.name for b in list_buckets]
+        except Exception:
+            self.logger.error("Could not retrieve the list of available Buckets!")
+            buckets_name = []
+        layout = ipywidgets.Layout(width='400px')
+        self.buckets_dropdown = ipywidgets.Dropdown(description='Bucket', value=None,
+                                                    options=buckets_name, layout=layout)
+        self.files_list = ipywidgets.Select(description='Files', value=None, disabled=False, layout=layout)
+
+        self.buckets_dropdown.observe(self.select_bucket, names='value')
+
+        ipywidgets.VBox.__init__(self, [self.buckets_dropdown, self.files_list], **kwargs)
+
+    def get_selected_bucket(self):
+        return self.buckets_dropdown.value
+
+    def get_selected_file_path(self):
+        file = self.files_list.value
+        if file is None:
+            return file
+        return self.get_selected_bucket() + "/" + file
+
+    def get_selected_file_content(self):
+        file_path = self.files_list.value
+        bucket_name = self.get_selected_bucket()
+        dataproxy_file = self._get_dataproxy_file(file_path, bucket_name)
+        response = requests.get(dataproxy_file.get_download_link())
+        return response.content
+
+    def select_bucket(self, _):
+        selected_bucket = self.get_selected_bucket()
+        bucket = self.client.buckets.get_bucket(selected_bucket)
+        self.files_list.options = [f.name for f in bucket.ls()]
+
+    def _get_dataproxy_file(self, file_path, bucket_name):
+        # type: (str, str) -> DataproxyFile
+        """
+        Get the DataProxy file corresponding to the path <file_path> in bucket <bucket_name>
+        """
+        file_path = file_path.lstrip('/')
+        bucket = self.client.buckets.get_bucket(bucket_name)
+        # find first dataproxy file corresponding to provided path
+        dataproxy_file = next((f for f in bucket.ls() if f.name == file_path), None)
+        return dataproxy_file

tvbwidgets/ui/storage_widget.py

@@ -7,28 +7,35 @@
 import ipywidgets
 
 from tvbwidgets.ui.base_widget import TVBWidget
+from tvbwidgets.ui.bucket_widget import BucketWidget
 from tvbwidgets.ui.drive_widget import DriveWidget
 
 
 class StorageWidget(ipywidgets.Tab, TVBWidget):
 
-    def __init__(self, collab=None, folder=None, **kwargs):
-        tab1 = ipywidgets.VBox()
-        tab2 = DriveWidget(collab, folder)
-        tab3 = ipywidgets.VBox()
+    def __init__(self, collab=None, folder=None, selected_storage=0, **kwargs):
+        self.tab1 = DriveWidget(collab, folder)
+        self.tab2 = BucketWidget()
 
-        super().__init__([tab1, tab2, tab3], selected_index=1,
+        super().__init__([self.tab1, self.tab2], selected_index=selected_storage,
                          layout=ipywidgets.Layout(width='550px', height='200px'), **kwargs)
 
-        self.set_title(0, 'Current Selection')
-        self.set_title(1, 'Drive')
-        self.set_title(2, 'Bucket')
-        # TODO uniform API for all tabs
-        self.api = tab2
+        self.set_title(0, 'Drive')
+        self.set_title(1, 'Bucket')
+        self.drive_api = self.tab1
+        self.bucket_api = self.tab1
 
     def get_selected_file_content(self):
-        return self.api.get_selected_file_content()
+        api = self.retrieve_api()
+        return api.get_selected_file_content()
 
     def get_selected_file_name(self):
-        filename = self.api.get_selected_file_path()
+        api = self.retrieve_api()
+        filename = api.get_selected_file_path()
         return filename
+
+    def retrieve_api(self):
+        if self.selected_index == 0:
+            return self.drive_api
+        else:
+            return self.bucket_api