def test_migrate(self): migration.migrate(dry_run=False) assert_equal( S3UserSettings.find().count(), len(self.linked_user_settings + self.unlinked_user_settings) ) assert_equal( S3NodeSettings.find().count(), len(self.node_settings_documents + self.node_settings_no_encrypt) ) for user_settings in S3UserSettings.find(): assert_is_not_none(user_settings.owner) assert_false(hasattr(user_settings, 'access_key')) for node_settings in S3NodeSettings.find(): assert_is_not_none(node_settings.owner) assert_is_not_none(node_settings.external_account)
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.')