Cache the bucket if the key is limited to one

CHANGELOG.md

@@ -6,6 +6,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+* Authorizing a key for a single bucket ensures that this bucket is cached
+
 ### Infrastructure
 * Additional tests for listing files/versions
 

b2sdk/_v3/exception.py

@@ -54,6 +54,7 @@
 from b2sdk.exception import NotAllowedByAppKeyError
 from b2sdk.exception import PartSha1Mismatch
 from b2sdk.exception import RestrictedBucket
+from b2sdk.exception import RestrictedBucketMissing
 from b2sdk.exception import RetentionWriteError
 from b2sdk.exception import SSECKeyError
 from b2sdk.exception import SSECKeyIdMismatchInCopy
@@ -134,6 +135,7 @@
     'NotAllowedByAppKeyError',
     'PartSha1Mismatch',
     'RestrictedBucket',
+    'RestrictedBucketMissing',
     'RetentionWriteError',
     'ServiceError',
     'SourceReplicationConflict',

b2sdk/api.py

@@ -12,13 +12,14 @@
 from contextlib import suppress
 
 from .account_info.abstract import AbstractAccountInfo
+from .account_info.exception import MissingAccountData
 from .api_config import B2HttpApiConfig, DEFAULT_HTTP_API_CONFIG
 from .application_key import ApplicationKey, BaseApplicationKey, FullApplicationKey
 from .cache import AbstractCache
 from .bucket import Bucket, BucketFactory
 from .encryption.setting import EncryptionSetting
 from .replication.setting import ReplicationConfiguration
-from .exception import BucketIdNotFound, NonExistentBucket, RestrictedBucket
+from .exception import BucketIdNotFound, NonExistentBucket, RestrictedBucket, RestrictedBucketMissing
 from .file_lock import FileRetentionSetting, LegalHold
 from .file_version import DownloadVersionFactory, FileIdAndName, FileVersion, FileVersionFactory
 from .large_file.services import LargeFileServices
@@ -201,6 +202,7 @@ def authorize_account(self, realm, application_key_id, application_key):
         :param str application_key: user's :term:`application key`
         """
         self.session.authorize_account(realm, application_key_id, application_key)
+        self._populate_bucket_cache_from_key()
 
     def get_account_id(self):
         """
@@ -595,3 +597,23 @@ def _check_bucket_restrictions(self, key, value):
         if allowed_bucket_identifier is not None:
             if allowed_bucket_identifier != value:
                 raise RestrictedBucket(allowed_bucket_identifier)
+
+    def _populate_bucket_cache_from_key(self):
+        # If the key is restricted to the bucket, pre-populate the cache with it
+        try:
+            allowed = self.account_info.get_allowed()
+        except MissingAccountData:
+            return
+
+        allowed_bucket_id = allowed.get('bucketId')
+        if allowed_bucket_id is None:
+            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))

b2sdk/exception.py

@@ -39,7 +39,7 @@ def __init__(self, *args, **kwargs):
         # If the exception is caused by a b2 server response,
         # the server MAY have included instructions to pause the thread before issuing any more requests
         self.retry_after_seconds = None
-        super(B2Error, self).__init__(*args, **kwargs)
+        super().__init__(*args, **kwargs)
 
     @property
     def prefix(self):
@@ -82,7 +82,7 @@ class B2SimpleError(B2Error, metaclass=ABCMeta):
     """
 
     def __str__(self):
-        return '%s: %s' % (self.prefix, super(B2SimpleError, self).__str__())
+        return '%s: %s' % (self.prefix, super().__str__())
 
 
 class NotAllowedByAppKeyError(B2SimpleError, metaclass=ABCMeta):
@@ -134,7 +134,7 @@ class CapabilityNotAllowed(NotAllowedByAppKeyError):
 
 class ChecksumMismatch(TransientErrorMixin, B2Error):
     def __init__(self, checksum_type, expected, actual):
-        super(ChecksumMismatch, self).__init__()
+        super().__init__()
         self.checksum_type = checksum_type
         self.expected = expected
         self.actual = actual
@@ -168,7 +168,7 @@ def __init__(self, clock_skew_seconds):
         """
         :param int clock_skew_seconds: The difference: local_clock - server_clock
         """
-        super(ClockSkew, self).__init__()
+        super().__init__()
         self.clock_skew_seconds = clock_skew_seconds
 
     def __str__(self):
@@ -210,7 +210,7 @@ def should_retry_http(self):
 
 class DestFileNewer(B2Error):
     def __init__(self, dest_path, source_path, dest_prefix, source_prefix):
-        super(DestFileNewer, self).__init__()
+        super().__init__()
         self.dest_path = dest_path
         self.source_path = source_path
         self.dest_prefix = dest_prefix
@@ -240,7 +240,7 @@ class ResourceNotFound(B2SimpleError):
 
 class FileOrBucketNotFound(ResourceNotFound):
     def __init__(self, bucket_name=None, file_id_or_name=None):
-        super(FileOrBucketNotFound, self).__init__()
+        super().__init__()
         self.bucket_name = bucket_name
         self.file_id_or_name = file_id_or_name
 
@@ -291,7 +291,7 @@ class SSECKeyIdMismatchInCopy(InvalidMetadataDirective):
 
 class InvalidRange(B2Error):
     def __init__(self, content_length, range_):
-        super(InvalidRange, self).__init__()
+        super().__init__()
         self.content_length = content_length
         self.range_ = range_
 
@@ -310,7 +310,7 @@ class InvalidUploadSource(B2SimpleError):
 
 class BadRequest(B2Error):
     def __init__(self, message, code):
-        super(BadRequest, self).__init__()
+        super().__init__()
         self.message = message
         self.code = code
 
@@ -326,7 +326,7 @@ def __init__(self, message, code, size: int):
 
 class Unauthorized(B2Error):
     def __init__(self, message, code):
-        super(Unauthorized, self).__init__()
+        super().__init__()
         self.message = message
         self.code = code
 
@@ -350,22 +350,29 @@ class InvalidAuthToken(Unauthorized):
     """
 
     def __init__(self, message, code):
-        super(InvalidAuthToken,
-              self).__init__('Invalid authorization token. Server said: ' + message, code)
+        super().__init__('Invalid authorization token. Server said: ' + message, code)
 
 
 class RestrictedBucket(B2Error):
     def __init__(self, bucket_name):
-        super(RestrictedBucket, self).__init__()
+        super().__init__()
         self.bucket_name = bucket_name
 
     def __str__(self):
         return 'Application key is restricted to bucket: %s' % self.bucket_name
 
 
+class RestrictedBucketMissing(RestrictedBucket):
+    def __init__(self):
+        super().__init__('')
+
+    def __str__(self):
+        return 'Application key is restricted to a bucket that doesn\'t exist'
+
+
 class MaxFileSizeExceeded(B2Error):
     def __init__(self, size, max_allowed_size):
-        super(MaxFileSizeExceeded, self).__init__()
+        super().__init__()
         self.size = size
         self.max_allowed_size = max_allowed_size
 
@@ -378,7 +385,7 @@ def __str__(self):
 
 class MaxRetriesExceeded(B2Error):
     def __init__(self, limit, exception_info_list):
-        super(MaxRetriesExceeded, self).__init__()
+        super().__init__()
         self.limit = limit
         self.exception_info_list = exception_info_list
 
@@ -405,7 +412,7 @@ class FileSha1Mismatch(B2SimpleError):
 
 class PartSha1Mismatch(B2Error):
     def __init__(self, key):
-        super(PartSha1Mismatch, self).__init__()
+        super().__init__()
         self.key = key
 
     def __str__(self):
@@ -435,7 +442,7 @@ def __str__(self):
 
 class TooManyRequests(B2Error):
     def __init__(self, retry_after_seconds=None):
-        super(TooManyRequests, self).__init__()
+        super().__init__()
         self.retry_after_seconds = retry_after_seconds
 
     def __str__(self):
@@ -447,7 +454,7 @@ def should_retry_http(self):
 
 class TruncatedOutput(TransientErrorMixin, B2Error):
     def __init__(self, bytes_read, file_size):
-        super(TruncatedOutput, self).__init__()
+        super().__init__()
         self.bytes_read = bytes_read
         self.file_size = file_size
 
@@ -482,7 +489,7 @@ def __str__(self):
 
 class UploadTokenUsedConcurrently(B2Error):
     def __init__(self, token):
-        super(UploadTokenUsedConcurrently, self).__init__()
+        super().__init__()
         self.token = token
 
     def __str__(self):

b2sdk/raw_simulator.py

@@ -16,7 +16,7 @@
 import threading
 import time
 
-from contextlib import contextmanager
+from contextlib import contextmanager, suppress
 from typing import Optional
 
 from b2sdk.http_constants import FILE_INFO_HEADER_PREFIX, HEX_DIGITS_AT_END
@@ -91,6 +91,7 @@ def __init__(
         self.capabilities = capabilities
         self.expiration_timestamp_or_none = expiration_timestamp_or_none
         self.bucket_id_or_none = bucket_id_or_none
+        self.bucket_name_or_none = bucket_name_or_none
         self.name_prefix_or_none = name_prefix_or_none
 
     def as_key(self):
@@ -121,6 +122,7 @@ def get_allowed(self):
         """
         return dict(
             bucketId=self.bucket_id_or_none,
+            bucketName=self.bucket_name_or_none,
             capabilities=self.capabilities,
             namePrefix=self.name_prefix_or_none,
         )
@@ -1305,10 +1307,13 @@ def create_key(
         self.app_key_counter += 1
         application_key_id = 'appKeyId%d' % (index,)
         app_key = 'appKey%d' % (index,)
-        if bucket_id is None:
-            bucket_name_or_none = None
-        else:
-            bucket_name_or_none = self._get_bucket_by_id(bucket_id).bucket_name
+        bucket_name_or_none = None
+        if bucket_id is not None:
+            # It is possible for bucketId to be filled and bucketName to be empty.
+            # It can happen when the bucket was deleted.
+            with suppress(NonExistentBucket):
+                bucket_name_or_none = self._get_bucket_by_id(bucket_id).bucket_name
+
         key_sim = KeySimulator(
             account_id=account_id,
             name=key_name,

b2sdk/session.py

@@ -131,7 +131,7 @@ def authorize_account(self, realm, application_key_id, application_key):
             realm=realm,
             s3_api_url=response['s3ApiUrl'],
             allowed=allowed,
-            application_key_id=application_key_id
+            application_key_id=application_key_id,
         )
 
     def cancel_large_file(self, file_id):

test/unit/bucket/test_bucket.py

@@ -32,6 +32,7 @@
     InvalidRange,
     InvalidUploadSource,
     MaxRetriesExceeded,
+    RestrictedBucketMissing,
     SSECKeyError,
     SourceReplicationConflict,
     UnsatisfiableRange,
@@ -57,6 +58,7 @@
 from apiver_deps import Range
 from apiver_deps import SimpleDownloader
 from apiver_deps import UploadSourceBytes
+from apiver_deps import DummyCache, InMemoryCache
 from apiver_deps import hex_sha1_of_bytes, TempDir
 from apiver_deps import EncryptionAlgorithm, EncryptionSetting, EncryptionMode, EncryptionKey, SSE_NONE, SSE_B2_AES
 from apiver_deps import CopySource, UploadSourceLocalFile, WriteIntent
@@ -200,10 +202,13 @@ def bucket_ls(bucket, *args, show_versions=False, **kwargs):
 
 class TestCaseWithBucket(TestBase):
     RAW_SIMULATOR_CLASS = RawSimulator
+    CACHE_CLASS = DummyCache
 
     def get_api(self):
         return B2Api(
-            self.account_info, api_config=B2HttpApiConfig(_raw_api_class=self.RAW_SIMULATOR_CLASS)
+            self.account_info,
+            cache=self.CACHE_CLASS(),
+            api_config=B2HttpApiConfig(_raw_api_class=self.RAW_SIMULATOR_CLASS),
         )
 
     def setUp(self):
@@ -215,7 +220,7 @@ def setUp(self):
         self.api.authorize_account('production', self.account_id, self.master_key)
         self.api_url = self.account_info.get_api_url()
         self.account_auth_token = self.account_info.get_account_auth_token()
-        self.bucket = self.api.create_bucket('my-bucket', 'allPublic')
+        self.bucket = self.api.create_bucket(self.bucket_name, 'allPublic')
         self.bucket_id = self.bucket.id_
 
     def bucket_ls(self, *args, show_versions=False, **kwargs):
@@ -2175,9 +2180,45 @@ def test_file_info_4(self):
         assert download_version.file_name == 'test.txt%253Ffoo%253Dbar'
 
 
-# Listing where every other response returns no entries and pointer to the next file
+class TestAuthorizeForBucket(TestCaseWithBucket):
+    CACHE_CLASS = InMemoryCache
+
+    @pytest.mark.apiver(from_ver=2)
+    def test_authorize_for_bucket_ensures_cache(self):
+        key = create_key(
+            self.api,
+            key_name='singlebucket',
+            capabilities=[
+                'listBuckets',
+            ],
+            bucket_id=self.bucket_id,
+        )
+
+        self.api.authorize_account('production', key.id_, key.application_key)
 
+        # Check whether the bucket fetching performs an API call.
+        with mock.patch.object(self.api, 'list_buckets') as mock_list_buckets:
+            self.api.get_bucket_by_id(self.bucket_id)
+            mock_list_buckets.assert_not_called()
 
+            self.api.get_bucket_by_name(self.bucket_name)
+            mock_list_buckets.assert_not_called()
+
+    @pytest.mark.apiver(from_ver=2)
+    def test_authorize_for_non_existing_bucket(self):
+        key = create_key(
+            self.api,
+            key_name='singlebucket',
+            capabilities=[
+                'listBuckets',
+            ],
+            bucket_id=self.bucket_id + 'x',
+        )
+
+        with self.assertRaises(RestrictedBucketMissing):
+            self.api.authorize_account('production', key.id_, key.application_key)
+
+# Listing where every other response returns no entries and pointer to the next file
 class EmptyListBucketSimulator(BucketSimulator):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)