EBR-79: extend files from ebrains_drive for bucket implementation

requirements.txt

@@ -16,7 +16,6 @@ trame-vuetify
 trame-vtk
 tvb-library
 tvb-framework
-tvb-ext-bucket
 plotly-resampler
 pyunicore>=1.0.0
 traitlets>=5.7.1

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/tests/test_bucket_widget.py

@@ -8,6 +8,7 @@
 
 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'
@@ -67,10 +68,7 @@ def __init__(self, token=''):
 
 @pytest.fixture
 def mock_client(mocker):
-    def mock_get_client(_):
-        return MockBucketApiClient()
-
-    mocker.patch('tvb_ext_bucket.ebrains_drive_wrapper.BucketWrapper.get_client', mock_get_client)
+    return mocker.patch('tvbwidgets.ui.bucket_widget.ExtendedBucketApiClient', MockBucketApiClient)
 
 
 @pytest.fixture
@@ -84,6 +82,13 @@ 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

tvbwidgets/ui/bucket_widget.py

@@ -7,44 +7,62 @@
 
 import ipywidgets
 import requests
-from tvb_ext_bucket.ebrains_drive_wrapper import BucketWrapper
+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)
-        self.client = BucketWrapper()
+        bearer_token = get_current_token()
+        self.client = ExtendedBucketApiClient(token=bearer_token)
 
         try:
-            all_buckets = self.client.list_buckets()
+            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!")
-            all_buckets = []
+            buckets_name = []
         layout = ipywidgets.Layout(width='400px')
         self.buckets_dropdown = ipywidgets.Dropdown(description='Bucket', value=None,
-                                                    options=all_buckets, layout=layout)
+                                                    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')
 
-        self._parent_dir = None
-        self._map_names_to_files = dict()
         ipywidgets.VBox.__init__(self, [self.buckets_dropdown, self.files_list], **kwargs)
 
-    def get_chosen_bucket(self):
+    def get_selected_bucket(self):
         return self.buckets_dropdown.value
 
     def get_selected_file_path(self):
-        return self.buckets_dropdown.value + "/" + self.files_list.value
+        file = self.files_list.value
+        if file is None:
+            return file
+        return self.get_selected_bucket() + "/" + file
 
     def get_selected_file_content(self):
-        path = self.files_list.value
-        downloadable_url = self.client.get_download_url(path, self.buckets_dropdown.value)
-        response = requests.get(downloadable_url)
+        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.buckets_dropdown.value
-        self.files_list.options = self.client.get_files_in_bucket(selected_bucket)
+        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