def set_folder(self, folder_id, auth):
if not bucket_exists(self.external_account.oauth_key, self.external_account.oauth_secret, folder_id):
error_message = ('We are having trouble connecting to that bucket. '
'Try a different one.')
raise exceptions.InvalidFolderError(error_message)
self.folder_id = str(folder_id)
bucket_location = get_bucket_location_or_error(
self.external_account.oauth_key,
self.external_account.oauth_secret,
folder_id
)
try:
bucket_location = BUCKET_LOCATIONS[bucket_location]
except KeyError:
# Unlisted location, S3 may have added it recently.
# Default to the key. When hit, add mapping to settings
pass
self.folder_name = '{} ({})'.format(folder_id, bucket_location)
self.save()
self.nodelogger.log(action='bucket_linked', extra={'bucket': str(folder_id)}, save=True)
def migrate(dry_run=False):
bucket_name_location_map = {}
sns_collection = database['s3nodesettings']
logger.info('Migrating all {} S3NodeSettings documents: {}'.format(
sns_collection.count(), [s['_id'] for s in sns_collection.find()]
))
for document in sns_collection.find():
sns_collection.find_and_modify(
{'_id': document['_id']},
{
'$set': {'folder_id': document['bucket'], 'folder_name': document['bucket']},
'$unset': {'bucket': ''}
}
)
allowance = 2
last_call = time.time()
for node_settings in S3NodeSettings.find(Q('folder_id', 'ne', None)):
if node_settings.folder_id in bucket_name_location_map:
# See if this bucket is cached
node_settings.folder_name = '{} ({})'.format(
node_settings.folder_id,
bucket_name_location_map[node_settings.folder_id]
)
else:
# Attempt to determine bucket location, default to just bucket name.
node_settings.folder_name = node_settings.folder_id
if allowance < 1:
try:
time.sleep(1 - (time.time() - last_call))
except (ValueError, IOError):
pass # ValueError/IOError indicates a negative sleep time
allowance = 2
allowance -= 1
last_call = time.time()
bucket_location = None
try:
bucket_location = get_bucket_location_or_error(
node_settings.external_account.oauth_key,
node_settings.external_account.oauth_secret,
node_settings.folder_id
)
except InvalidAuthError:
logger.info('Found S3NodeSettings {} with invalid credentials.'.format(node_settings._id))
except InvalidFolderError:
logger.info('Found S3NodeSettings {} with invalid bucket linked.'.format(node_settings._id))
except Exception as e:
logger.info('Found S3NodeSettings {} throwing unknown error. Likely configured improperly; with a bucket but no credentials'.format(node_settings._id))
logger.exception(e)
else:
try:
bucket_location = BUCKET_LOCATIONS[bucket_location]
except KeyError:
# Unlisted location, S3 may have added it recently.
# Default to the key. When hit, add mapping to settings
logger.info('Found unknown location key: {}'.format(bucket_location))
node_settings.folder_name = '{} ({})'.format(node_settings.folder_id, bucket_location)
bucket_name_location_map[node_settings.folder_id] = bucket_location
if not bucket_location:
node_settings.folder_name = node_settings.folder_id
node_settings.save()
if dry_run:
raise RuntimeError('Dry run, transaction rolled back.')