Merge pull request #14 from the-virtual-brain/EBR-89

EBR-89: extend/use classes from ebrains_drive

tvb_ext_bucket/bucket_api/bucket_api.py

@@ -1,63 +1,14 @@
-from ebrains_drive.exceptions import TokenExpired
+from ebrains_drive.client import BucketApiClient
 
-from tvb_ext_bucket.bucket_api.buckets import Buckets
-from ebrains_drive.client import ClientBase
-from ebrains_drive.utils import on_401_raise_unauthorized
+from tvb_ext_bucket.bucket_api.buckets import ExtendedBuckets
 
-import base64
-import json
-import time
 
-
-class BucketApiClient(ClientBase):
+class ExtendedBucketApiClient(BucketApiClient):
 
     def __init__(self, username=None, password=None, token=None, env="") -> None:
-        if env != "":
-            raise NotImplementedError("non prod environment for dataproxy access has not yet been implemented.")
-        self._set_env(env)
-
         super().__init__(username, password, token, env)
-
-        self.server = "https://data-proxy.ebrains.eu/api"
-
-        self.buckets = Buckets(self)
+        self.buckets = ExtendedBuckets(self)
 
     @property
     def token(self):
         return self._token
-
-    @on_401_raise_unauthorized(
-        "Failed. Note: BucketApiClient.create_new needs to have clb.drive:write as a part of scope.")
-    def create_new(self, bucket_name: str, title=None, description="Created by ebrains_drive"):
-        # attempt to create new collab
-        self.send_request("POST", "https://wiki.ebrains.eu/rest/v1/collabs", json={
-            "name": bucket_name,
-            "title": title or bucket_name,
-            "description": description,
-            "drive": True,
-            "chat": True,
-            "public": False
-        }, expected=201)
-
-        # activate the bucket for the said collab
-        self.send_request("POST", "/v1/buckets", json={
-            "bucket_name": bucket_name
-        }, expected=201)
-
-    @on_401_raise_unauthorized(
-        "Failed. Note: BucketApiClient.create_new needs to have clb.drive:write as a part of scope.")
-    def delete_bucket(self, bucket_name: str):
-        self.send_request("DELETE", f"/v1/buckets/{bucket_name}")
-
-    def send_request(self, method: str, url: str, *args, **kwargs):
-        _, info, _ = self._token.split('.')
-        info_json = base64.b64decode(info + '==').decode('utf-8')
-
-        # https://www.rfc-editor.org/rfc/rfc7519#section-2
-        exp_utc_seconds = json.loads(info_json).get('exp')
-        now_tc_seconds = time.time()
-
-        if now_tc_seconds > exp_utc_seconds:
-            raise TokenExpired
-
-        return super().send_request(method, url, *args, **kwargs)

tvb_ext_bucket/bucket_api/buckets.py

@@ -1,9 +1,6 @@
-from ebrains_drive.exceptions import ClientHttpError, Unauthorized
-from ebrains_drive.utils import on_401_raise_unauthorized
+from ebrains_drive.buckets import Buckets
 from tvb_ext_bucket.exceptions import BucketDTOError
-from tvb_ext_bucket.bucket_api.bucket import Bucket, Endpoint
 from tvb_ext_bucket.logger.builder import get_logger
-from time import sleep
 from dataclasses import dataclass
 from typing import List
 
@@ -17,48 +14,13 @@ class BucketDTO:
     is_public: bool
 
 
-class Buckets:
+class ExtendedBuckets(Buckets):
     BUCKETS_ENDPOINT = '/v1/buckets'
     DATASETS_ENDPOINT = '/v1/datasets'
 
     def __init__(self, client):
+        super().__init__(client)
         self._available_buckets: List[BucketDTO] = []
-        self.client = client
-
-    @on_401_raise_unauthorized(
-        '401 response. Check you/your token have access right and/or the bucket name has been spelt correctly.')
-    def get_bucket(self, bucket_name: str, *, public: bool = False) -> Bucket:
-        """
-        Get the specified bucket according name.
-        """
-        LOGGER.info(f'Trying to retrieve bucket {bucket_name}. (Public: {public})')
-        resp = self.client.get(f"{self.BUCKETS_ENDPOINT}/{bucket_name}/stat")
-        return Bucket.from_json(self.client, resp.json(), public=public, target=Endpoint.BUCKETS)
-
-    def get_dataset(self, dataset_id: str, *, public: bool = False, request_access: bool = False):
-        request_sent = False
-        attempt_no = 0
-        while True:
-            try:
-                resp = self.client.get(f"{self.DATASETS_ENDPOINT}/{dataset_id}/stat")
-                return Bucket.from_json(self.client, resp.json(), public=public, target=Endpoint.DATASETS,
-                                        dataset_id=dataset_id)
-            except ClientHttpError as e:
-                if e.code != 401:
-                    raise e
-
-                if not request_access:
-                    raise Unauthorized(
-                        "You do not have access to this dataset. "
-                        "If this is a private dataset, try to set request_access flag to true. "
-                        "We can start the procedure of requesting access for you.")
-                if not request_sent:
-                    self.client.post(f"/v1/datasets/{dataset_id}", expected=(200, 201))
-                    request_sent = True
-                    print("Request sent. Please check the mail box associated with the token.")
-                sleep(5)
-                attempt_no = attempt_no + 1
-                print(f"Checking permission, attempt {attempt_no}")
 
     def list_buckets(self):
         # type: () -> List[BucketDTO]
@@ -76,13 +38,3 @@ def list_buckets(self):
         except KeyError as e:
             LOGGER.error(f'Received unexpected Bucket structure! {str(e)}')
             raise BucketDTOError('Unexpected response structure from server!')
-
-    def _check_bucket_availability(self, bucket_name):
-        # type: (str) -> bool
-        """
-        Checks if the bucket is available for current user.
-        """
-        for bucket in self.list_buckets():
-            if bucket.name == bucket_name:
-                return True
-        return False

tvb_ext_bucket/ebrains_drive_wrapper.py

@@ -9,13 +9,13 @@
 
 from ebrains_drive.files import DataproxyFile
 from ebrains_drive.exceptions import Unauthorized
+from ebrains_drive.bucket import Bucket
 
 from tvb_ext_bucket.logger.builder import get_logger
 from tvb_ext_bucket.exceptions import CollabTokenError, CollabAccessError, DataproxyFileNotFound
+from tvb_ext_bucket.bucket_api.bucket_api import ExtendedBucketApiClient
 import os
 
-from tvb_ext_bucket.bucket_api.bucket_api import BucketApiClient
-from tvb_ext_bucket.bucket_api.bucket import Bucket
 import pathlib
 
 LOGGER = get_logger(__name__)
@@ -100,7 +100,7 @@ def get_files_in_bucket(self, bucket_name):
 
     @staticmethod
     def get_client():
-        # type: () -> BucketApiClient
+        # type: () -> ExtendedBucketApiClient
         """
         Get an instance of the BucketApiClient
         Returns
@@ -118,7 +118,7 @@ def get_client():
             raise CollabTokenError(f"Cannot connect to EBRAINS HPC without an auth token! Either run this on "
                                    f"Collab, or define the {TOKEN_ENV_VAR} environment variable!")
         LOGGER.info('Token retrieved successfully!')
-        return BucketApiClient(token=token)
+        return ExtendedBucketApiClient(token=token)
 
     def download_file(self, file_path, bucket_name, location):
         # type: (str, str, str) -> bool

tvb_ext_bucket/tests/test_bucket.py

@@ -2,9 +2,8 @@
 import tempfile
 
 import pytest
+from ebrains_drive.bucket import Bucket
 from ebrains_drive.exceptions import DoesNotExist
-
-from tvb_ext_bucket.bucket_api.bucket import Bucket
 from requests import Response
 from tempfile import TemporaryFile
 
@@ -82,7 +81,7 @@ def test_bucket_instance():
     assert bucket.last_modified == last_modified
     assert bucket.public is False
     assert bucket.dataproxy_entity_name == name
-    assert bucket.is_initialised is False
+    assert bucket.is_initialized is None
     assert bucket.target == 'buckets'
     assert bucket.role is None
 
@@ -98,7 +97,7 @@ def test_bucket_instance_from_json():
     assert bucket.last_modified == BUCKET_STAT_JSON['last_modified']
     assert bucket.public == BUCKET_STAT_JSON['is_public']
     assert bucket.dataproxy_entity_name == BUCKET_STAT_JSON['name']
-    assert bucket.is_initialised == BUCKET_STAT_JSON['is_initialized']
+    assert bucket.is_initialized == BUCKET_STAT_JSON['is_initialized']
     assert bucket.role == BUCKET_STAT_JSON['role']
 
 

tvb_ext_bucket/tests/test_buckets.py

@@ -1,4 +1,4 @@
-from tvb_ext_bucket.bucket_api.buckets import Buckets
+from tvb_ext_bucket.bucket_api.buckets import ExtendedBuckets
 from requests import Response
 
 
@@ -20,13 +20,13 @@ def get(self, _url):
 
 def test_buckets_instance():
     fake_client = MockClient()
-    buckets = Buckets(fake_client)
+    buckets = ExtendedBuckets(fake_client)
     assert buckets.client == fake_client
     assert buckets._available_buckets == []
 
 
 def test_get_bucket():
     fake_client = MockClient()
-    buckets = Buckets(fake_client)
+    buckets = ExtendedBuckets(fake_client)
     bucket = buckets.get_bucket('tvb-widgets')
     assert bucket.name == 'tvb-widgets'