diff --git a/addons/s3/utils.py b/addons/s3/utils.py index 6be76bd8290..750c81f8b38 100644 --- a/addons/s3/utils.py +++ b/addons/s3/utils.py @@ -1,4 +1,5 @@ import re +import logging from rest_framework import status as http_status from boto import exception @@ -8,6 +9,11 @@ from framework.exceptions import HTTPError from addons.base.exceptions import InvalidAuthError, InvalidFolderError from addons.s3.settings import BUCKET_LOCATIONS +from django.apps import apps +from django.db.models import F, Value +from django.db.models.functions import Concat, Replace + +logger = logging.getLogger(__name__) def connect_s3(access_key=None, secret_key=None, node_settings=None): @@ -18,7 +24,7 @@ def connect_s3(access_key=None, secret_key=None, node_settings=None): if node_settings is not None: if node_settings.external_account is not None: access_key, secret_key = node_settings.external_account.oauth_key, node_settings.external_account.oauth_secret - connection = S3Connection(access_key, secret_key) + connection = S3Connection(access_key, secret_key, calling_format=OrdinaryCallingFormat()) return connection @@ -125,7 +131,7 @@ def get_bucket_location_or_error(access_key, secret_key, bucket_name): def get_bucket_prefixes(access_key, secret_key, prefix, bucket_name): - bucket = connect_s3(access_key, secret_key).get_bucket(bucket_name) + bucket = S3Connection(access_key, secret_key).get_bucket(bucket_name) # Don't use OrdinaryCallingFormat folders = [] for key in bucket.list(delimiter='/', prefix=prefix): @@ -143,3 +149,47 @@ def get_bucket_prefixes(access_key, secret_key, prefix, bucket_name): ) return folders + + +def update_folder_names(): + NodeSettings = apps.get_model('addons_s3', 'NodeSettings') + + # Update folder_id for all records + NodeSettings.objects.exclude( + folder_name__contains=':/' + ).update( + folder_id=Concat(F('folder_id'), Value(':/')) + ) + + # Update folder_name for records containing '(' + NodeSettings.objects.filter( + folder_name__contains=' (' + ).exclude( + folder_name__contains=':/' + ).update( + folder_name=Replace(F('folder_name'), Value(' ('), Value(':/ (')) + ) + NodeSettings.objects.exclude( + folder_name__contains=':/' + ).exclude( + folder_name__contains=' (' + ).update( + folder_name=Concat(F('folder_name'), Value(':/')) + ) + logger.info('Update Folder Names/IDs complete') + + +def reverse_update_folder_names(): + NodeSettings = apps.get_model('addons_s3', 'NodeSettings') + + # Reverse update folder_id for all records + NodeSettings.objects.update(folder_id=Replace(F('folder_id'), Value(':/'), Value(''))) + + # Reverse update folder_name for records containing ':/ (' + NodeSettings.objects.filter(folder_name__contains=':/ (').update( + folder_name=Replace(F('folder_name'), Value(':/ ('), Value(' (')) + ) + NodeSettings.objects.filter(folder_name__contains=':/').update( + folder_name=Replace(F('folder_name'), Value(':/'), Value('')) + ) + logger.info('Reverse Update Folder Names/IDs complete') diff --git a/osf/management/commands/add_colon_delim_to_s3_buckets.py b/osf/management/commands/add_colon_delim_to_s3_buckets.py new file mode 100644 index 00000000000..de3e5f42413 --- /dev/null +++ b/osf/management/commands/add_colon_delim_to_s3_buckets.py @@ -0,0 +1,31 @@ +# -*- coding: utf-8 -*- +import logging + +from django.core.management.base import BaseCommand +from addons.s3.utils import update_folder_names, reverse_update_folder_names + +logger = logging.getLogger(__name__) + + +class Command(BaseCommand): + """ + Adds Colon (':') delineators to s3 buckets to separate them from them from their subfolder, so `` + becomes `:/` , the root path. Folder names will also be updated to maintain consistency. + + """ + + def add_arguments(self, parser): + super().add_arguments(parser) + parser.add_argument( + '--reverse', + action='store_true', + dest='reverse', + help='Unsets date_retraction' + ) + + def handle(self, *args, **options): + reverse = options.get('reverse', False) + if reverse: + reverse_update_folder_names() + else: + update_folder_names() diff --git a/osf_tests/test_s3_folder_migration.py b/osf_tests/test_s3_folder_migration.py new file mode 100644 index 00000000000..a74fcd84802 --- /dev/null +++ b/osf_tests/test_s3_folder_migration.py @@ -0,0 +1,41 @@ +import pytest +from addons.s3.utils import update_folder_names, reverse_update_folder_names + +@pytest.mark.django_db +class TestUpdateFolderNamesMigration: + + def test_update_folder_names_migration(self): + from addons.s3.models import NodeSettings + from addons.s3.tests.factories import S3NodeSettingsFactory + # Create sample folder names and IDs + S3NodeSettingsFactory(folder_name='Folder 1 (Location 1)', folder_id='folder1') + S3NodeSettingsFactory(folder_name='Folder 2', folder_id='folder2') + S3NodeSettingsFactory(folder_name='Folder 3 (Location 3)', folder_id='folder3') + S3NodeSettingsFactory(folder_name='Folder 4:/ (Location 4)', folder_id='folder4:/') + + update_folder_names() + + # Verify updated folder names and IDs + updated_folder_names_ids = NodeSettings.objects.values_list('folder_name', 'folder_id') + expected_updated_folder_names_ids = { + ('Folder 1:/ (Location 1)', 'folder1:/'), + ('Folder 2:/', 'folder2:/'), + ('Folder 3:/ (Location 3)', 'folder3:/'), + ('Folder 3:/ (Location 3)', 'folder3:/'), + ('Folder 4:/ (Location 4)', 'folder4:/'), + + } + assert set(updated_folder_names_ids) == expected_updated_folder_names_ids + + # Reverse the migration + reverse_update_folder_names() + + # Verify the folder names and IDs after the reverse migration + reverted_folder_names_ids = NodeSettings.objects.values_list('folder_name', 'folder_id') + expected_reverted_folder_names_ids = { + ('Folder 1 (Location 1)', 'folder1'), + ('Folder 2', 'folder2'), + ('Folder 3 (Location 3)', 'folder3'), + ('Folder 4 (Location 4)', 'folder4'), + } + assert set(reverted_folder_names_ids) == expected_reverted_folder_names_ids