Merge pull request #538 from reef-technologies/multi-bucket-keys

Add support for multi-bucket application keys

Author: olzhasar-reef

b2sdk/_internal/__init__.py

@@ -10,5 +10,5 @@
 """
 b2sdk._internal package contains internal modules, and should not be used directly.
 
-Please use chosen apiver package instead, e.g. b2sdk.v2
+Please use chosen apiver package instead, e.g. b2sdk.v3
 """

b2sdk/_internal/account_info/abstract.py

@@ -31,8 +31,7 @@ class AbstractAccountInfo(metaclass=B2TraceMetaAbstract):
 
     # The 'allowed' structure to use for old account info that was saved without 'allowed'.
     DEFAULT_ALLOWED = dict(
-        bucketId=None,
-        bucketName=None,
+        buckets=None,
         capabilities=ALL_CAPABILITIES,
         namePrefix=None,
     )
@@ -318,7 +317,7 @@ def set_auth_data(
         """
         if allowed is None:
             allowed = self.DEFAULT_ALLOWED
-        assert self.allowed_is_valid(allowed)
+        assert self.allowed_is_valid(allowed), allowed
 
         self._set_auth_data(
             account_id,
@@ -337,22 +336,15 @@ def set_auth_data(
     @classmethod
     def allowed_is_valid(cls, allowed):
         """
-        Make sure that all of the required fields are present, and that
-        bucketId is set if bucketName is.
+        Make sure that all of the required fields are present
 
         If the bucketId is for a bucket that no longer exists, or the
         capabilities do not allow for listBuckets, then we will not have a bucketName.
 
         :param dict allowed: the structure to use for old account info that was saved without 'allowed'
         :rtype: bool
         """
-        return (
-            ('bucketId' in allowed)
-            and ('bucketName' in allowed)
-            and ((allowed['bucketId'] is not None) or (allowed['bucketName'] is None))
-            and ('capabilities' in allowed)
-            and ('namePrefix' in allowed)
-        )
+        return ('buckets' in allowed) and ('capabilities' in allowed) and ('namePrefix' in allowed)
 
     @abstractmethod
     def _set_auth_data(

b2sdk/_internal/account_info/sqlite_account_info.py

@@ -280,7 +280,7 @@ def _create_tables(self, conn, last_upgrade_to_run):
         """
         )
         # By default, we run all the upgrades
-        last_upgrade_to_run = 4 if last_upgrade_to_run is None else last_upgrade_to_run
+        last_upgrade_to_run = 5 if last_upgrade_to_run is None else last_upgrade_to_run
         # Add the 'allowed' column if it hasn't been yet.
         if 1 <= last_upgrade_to_run:
             self._ensure_update(1, ['ALTER TABLE account ADD COLUMN allowed TEXT;'])
@@ -384,26 +384,67 @@ def _create_tables(self, conn, last_upgrade_to_run):
                 ],
             )
 
+        if 5 <= last_upgrade_to_run:
+            self._migrate_allowed_to_multi_bucket()
+
+    def _migrate_allowed_to_multi_bucket(self):
+        """
+        Migrate existing allowed json dict to a new multi-bucket keys format
+        """
+        if self._get_update_count(5) > 0:
+            return
+
+        try:
+            allowed_json = self._get_account_info_or_raise('allowed')
+        except MissingAccountData:
+            allowed_json = None
+
+        if allowed_json is None:
+            self._perform_update(5, [])
+            return
+
+        allowed = json.loads(allowed_json)
+
+        bucket_id = allowed.pop('bucketId')
+        bucket_name = allowed.pop('bucketName')
+
+        if bucket_id is not None:
+            allowed['buckets'] = [{'id': bucket_id, 'name': bucket_name}]
+        else:
+            allowed['buckets'] = None
+
+        allowed_text = json.dumps(allowed)
+        stmt = f"UPDATE account SET allowed = ('{allowed_text}');"
+
+        self._perform_update(5, [stmt])
+
     def _ensure_update(self, update_number, update_commands: list[str]):
         """
         Run the update with the given number if it hasn't been done yet.
 
         Does the update and stores the number as a single transaction,
         so they will always be in sync.
         """
+        update_count = self._get_update_count(update_number)
+        if update_count > 0:
+            return
+
+        self._perform_update(update_number, update_commands)
+
+    def _get_update_count(self, update_number: int):
         with self._get_connection() as conn:
-            conn.execute('BEGIN')
             cursor = conn.execute(
                 'SELECT COUNT(*) AS count FROM update_done WHERE update_number = ?;',
                 (update_number,),
             )
-            update_count = cursor.fetchone()[0]
-            if update_count == 0:
-                for command in update_commands:
-                    conn.execute(command)
-                conn.execute(
-                    'INSERT INTO update_done (update_number) VALUES (?);', (update_number,)
-                )
+            return cursor.fetchone()[0]
+
+    def _perform_update(self, update_number, update_commands: list[str]):
+        with self._get_connection() as conn:
+            conn.execute('BEGIN')
+            for command in update_commands:
+                conn.execute(command)
+            conn.execute('INSERT INTO update_done (update_number) VALUES (?);', (update_number,))
 
     def clear(self):
         """
@@ -551,8 +592,7 @@ def get_allowed(self):
         .. code-block:: python
 
             {
-                "bucketId": null,
-                "bucketName": null,
+                "buckets": null,
                 "capabilities": [
                     "listKeys",
                     "writeKeys"
@@ -567,8 +607,8 @@ def get_allowed(self):
         allowed_json = self._get_account_info_or_raise('allowed')
         if allowed_json is None:
             return self.DEFAULT_ALLOWED
-        else:
-            return json.loads(allowed_json)
+
+        return json.loads(allowed_json)
 
     def get_s3_api_url(self):
         result = self._get_account_info_or_raise('s3_api_url')

b2sdk/_internal/api.py

@@ -37,7 +37,8 @@
 )
 from .large_file.services import LargeFileServices
 from .progress import AbstractProgressListener
-from .raw_api import API_VERSION, LifecycleRule
+from .raw_api import API_VERSION as RAW_API_VERSION
+from .raw_api import LifecycleRule
 from .replication.setting import ReplicationConfiguration
 from .session import B2Session
 from .transfer import (
@@ -52,21 +53,6 @@
 logger = logging.getLogger(__name__)
 
 
-def url_for_api(info, api_name):
-    """
-    Return URL for an API endpoint.
-
-    :param info: account info
-    :param str api_nam:
-    :rtype: str
-    """
-    if api_name in ['b2_download_file_by_id']:
-        base = info.get_download_url()
-    else:
-        base = info.get_api_url()
-    return f'{base}/b2api/{API_VERSION}/{api_name}'
-
-
 class Services:
     """Gathers objects that provide high level logic over raw api usage."""
 
@@ -141,7 +127,10 @@ class handles several things that simplify the task of uploading
     FILE_VERSION_FACTORY_CLASS = staticmethod(FileVersionFactory)
     DOWNLOAD_VERSION_FACTORY_CLASS = staticmethod(DownloadVersionFactory)
     SERVICES_CLASS = staticmethod(Services)
+    APPLICATION_KEY_CLASS = ApplicationKey
+    FULL_APPLICATION_KEY_CLASS = FullApplicationKey
     DEFAULT_LIST_KEY_COUNT = 1000
+    API_VERSION = RAW_API_VERSION
 
     def __init__(
         self,
@@ -484,14 +473,29 @@ def delete_file_version(
         response = self.session.delete_file_version(file_id, file_name, bypass_governance)
         return FileIdAndName.from_cancel_or_delete_response(response)
 
+    @classmethod
+    def _get_url_for_api(cls, info, api_name) -> str:
+        """
+        Return URL for an API endpoint.
+
+        :param info: account info
+        :param str api_nam:
+        :rtype: str
+        """
+        if api_name in ['b2_download_file_by_id']:
+            base = info.get_download_url()
+        else:
+            base = info.get_api_url()
+        return f'{base}/b2api/{cls.API_VERSION}/{api_name}'
+
     # download
     def get_download_url_for_fileid(self, file_id):
         """
         Return a URL to download the given file by ID.
 
         :param str file_id: a file ID
         """
-        url = url_for_api(self.account_info, 'b2_download_file_by_id')
+        url = self._get_url_for_api(self.account_info, 'b2_download_file_by_id')
         return f'{url}?fileId={file_id}'
 
     def get_download_url_for_file_name(self, bucket_name, file_name):
@@ -512,7 +516,7 @@ def create_key(
         capabilities: list[str],
         key_name: str,
         valid_duration_seconds: int | None = None,
-        bucket_id: str | None = None,
+        bucket_ids: list[str] | None = None,
         name_prefix: str | None = None,
     ) -> FullApplicationKey:
         """
@@ -531,14 +535,14 @@ def create_key(
             capabilities=capabilities,
             key_name=key_name,
             valid_duration_seconds=valid_duration_seconds,
-            bucket_id=bucket_id,
+            bucket_ids=bucket_ids,
             name_prefix=name_prefix,
         )
 
         assert set(response['capabilities']) == set(capabilities)
         assert response['keyName'] == key_name
 
-        return FullApplicationKey.from_create_response(response)
+        return self.FULL_APPLICATION_KEY_CLASS.from_create_response(response)
 
     def delete_key(self, application_key: BaseApplicationKey):
         """
@@ -557,7 +561,7 @@ def delete_key_by_id(self, application_key_id: str) -> ApplicationKey:
         """
 
         response = self.session.delete_key(application_key_id=application_key_id)
-        return ApplicationKey.from_api_response(response)
+        return self.APPLICATION_KEY_CLASS.from_api_response(response)
 
     def list_keys(
         self, start_application_key_id: str | None = None
@@ -576,7 +580,7 @@ def list_keys(
                 start_application_key_id=start_application_key_id,
             )
             for entry in response['keys']:
-                yield ApplicationKey.from_api_response(entry)
+                yield self.APPLICATION_KEY_CLASS.from_api_response(entry)
 
             next_application_key_id = response['nextApplicationKeyId']
             if next_application_key_id is None:
@@ -599,6 +603,8 @@ def get_key(self, key_id: str) -> ApplicationKey | None:
             if key.id_ == key_id:
                 return key
 
+        return None
+
     # other
     def get_file_info(self, file_id: str) -> FileVersion:
         """
@@ -630,7 +636,7 @@ def check_bucket_name_restrictions(self, bucket_name: str):
 
         :raises b2sdk.v2.exception.RestrictedBucket: if the account is not allowed to use this bucket
         """
-        self._check_bucket_restrictions('bucketName', bucket_name)
+        self._check_bucket_restrictions('name', bucket_name)
 
     def check_bucket_id_restrictions(self, bucket_id: str):
         """
@@ -641,15 +647,21 @@ def check_bucket_id_restrictions(self, bucket_id: str):
 
         :raises b2sdk.v2.exception.RestrictedBucket: if the account is not allowed to use this bucket
         """
-        self._check_bucket_restrictions('bucketId', bucket_id)
+        self._check_bucket_restrictions('id', bucket_id)
 
     def _check_bucket_restrictions(self, key, value):
-        allowed = self.account_info.get_allowed()
-        allowed_bucket_identifier = allowed[key]
+        buckets = self.account_info.get_allowed()['buckets']
 
-        if allowed_bucket_identifier is not None:
-            if allowed_bucket_identifier != value:
-                raise RestrictedBucket(allowed_bucket_identifier)
+        if not buckets:
+            return
+
+        for item in buckets:
+            if item[key] == value:
+                return
+
+        msg = str([b['name'] for b in buckets])
+
+        raise RestrictedBucket(msg)
 
     def _populate_bucket_cache_from_key(self):
         # If the key is restricted to the bucket, pre-populate the cache with it
@@ -658,15 +670,16 @@ def _populate_bucket_cache_from_key(self):
         except MissingAccountData:
             return
 
-        allowed_bucket_id = allowed.get('bucketId')
-        if allowed_bucket_id is None:
+        allowed_buckets = allowed.get('buckets')
+        if not allowed_buckets:
             return
 
-        allowed_bucket_name = allowed.get('bucketName')
-
         # If we have bucketId set we still need to check bucketName. If the bucketName is None,
         # it means that the bucketId belongs to a bucket that was already removed.
-        if allowed_bucket_name is None:
-            raise RestrictedBucketMissing()
 
-        self.cache.save_bucket(self.BUCKET_CLASS(self, allowed_bucket_id, name=allowed_bucket_name))
+        for item in allowed_buckets:
+            if item['name'] is None:
+                raise RestrictedBucketMissing
+
+        for item in allowed_buckets:
+            self.cache.save_bucket(self.BUCKET_CLASS(self, item['id'], name=item['name']))

b2sdk/_internal/application_key.py

@@ -20,7 +20,7 @@ def __init__(
         capabilities: list[str],
         account_id: str,
         expiration_timestamp_millis: int | None = None,
-        bucket_id: str | None = None,
+        bucket_ids: list[str] | None = None,
         name_prefix: str | None = None,
         options: list[str] | None = None,
     ):
@@ -39,7 +39,7 @@ def __init__(
         self.capabilities = capabilities
         self.account_id = account_id
         self.expiration_timestamp_millis = expiration_timestamp_millis
-        self.bucket_id = bucket_id
+        self.bucket_ids = bucket_ids
         self.name_prefix = name_prefix
         self.options = options
 
@@ -54,7 +54,7 @@ def parse_response_dict(cls, response: dict):
 
         optional_args = {
             'expiration_timestamp_millis': response.get('expirationTimestamp'),
-            'bucket_id': response.get('bucketId'),
+            'bucket_ids': response.get('bucketIds'),
             'name_prefix': response.get('namePrefix'),
             'options': response.get('options'),
         }
@@ -77,7 +77,7 @@ def as_dict(self):
         }
         optional_keys = {
             'expirationTimestamp': self.expiration_timestamp_millis,
-            'bucketId': self.bucket_id,
+            'bucketIds': self.bucket_ids,
             'namePrefix': self.name_prefix,
             'options': self.options,
         }
@@ -107,7 +107,7 @@ def __init__(
         capabilities: list[str],
         account_id: str,
         expiration_timestamp_millis: int | None = None,
-        bucket_id: str | None = None,
+        bucket_ids: list[str] | None = None,
         name_prefix: str | None = None,
         options: list[str] | None = None,
     ):
@@ -129,7 +129,7 @@ def __init__(
             capabilities=capabilities,
             account_id=account_id,
             expiration_timestamp_millis=expiration_timestamp_millis,
-            bucket_id=bucket_id,
+            bucket_ids=bucket_ids,
             name_prefix=name_prefix,
             options=options,
         )