Merge pull request #53 from AMP-SCZ/kcho/selectiveSync

feat: option to support data transfer for protected data

lochness/transfer/__init__.py

@@ -317,6 +317,48 @@ def lochness_to_lochness_transfer_s3(Lochness, general_only: bool = True):
     logger.debug('aws rsync completed')
 
 
+def lochness_to_lochness_transfer_s3_protected(Lochness):
+    '''Lochness to Lochness transfer using aws s3 sync for protected data
+
+    Key arguments:
+        Lochness: Lochness config.load object
+
+    Requirements:
+        - AWS CLI needs to be set with the correct credentials before executing
+        this module.
+            $ aws configure
+        - s3 bucket needs to be linked to the ID
+        - The name of the s3 bucket needs to be in the config.yml
+            eg) AWS_BUCKET_NAME: ampscz-dev
+                AWS_BUCKET_PHOENIX_ROOT: TEST_PHOENIX_ROOT
+
+    '''
+
+    s3_bucket_name = Lochness['AWS_BUCKET_NAME']
+    s3_phoenix_root = Lochness['AWS_BUCKET_ROOT']
+
+    for datatype in Lochness['selective_sync']:
+        # phoenix_root / PROTECTED / site / raw / subject / datatype
+        source_directories = Path(Lochness['phoenix_root']).glob(
+                    f'PROTECTED/*/*/*/{datatype}')
+
+        # for all studies and subjects
+        for source_directory in source_directories:
+            if source_directory.is_dir():
+                s3_phoenix_root_dtype = re.sub(Lochness['phoenix_root'],
+                                               s3_phoenix_root,
+                                               str(source_directory))
+                command = f'aws s3 sync \
+                        {source_directory}/ \
+                        s3://{s3_bucket_name}/{s3_phoenix_root_dtype}'
+
+                logger.debug('Executing aws s3 sync function for '
+                             f'{source_directory}')
+                logger.debug(re.sub(r'\s+', r' ', command))
+                logger.debug(os.popen(command).read())
+                logger.debug('aws rsync completed')
+
+
 def lochness_to_lochness_transfer_receive_sftp(Lochness):
     '''Get newly transferred file and decompress to PHOENIX
 

scripts/lochness_create_template.py

@@ -310,6 +310,12 @@ def create_config_template(config_loc: Path, args: object) -> None:
         s3_lines = f'''AWS_BUCKET_NAME: ampscz-dev
 AWS_BUCKET_ROOT: TEST_PHOENIX_ROOT'''
         config_example += s3_lines
+
+    if args.s3_selective_sync:
+        # eg)
+        # s3_selective_sync: ['mri', 'actigraphy']
+        config_example += f"\ns3_selective_sync: [{','.join(args.s3_selective_sync)}]"
+
 
     if 'redcap' in args.sources:
         config_example += '\nredcap:'
@@ -488,6 +494,11 @@ def get_arguments():
                         default=False,
                         action='store_true',
                         help='Use s3 rsync in lochness to lochness transfer')
+    parser.add_argument('--s3_selective_sync',
+                        default=False,
+                        nargs='+',
+                        help='List of dtypes from protected root to transfer '
+                             'using s3 rsync')
     parser.add_argument('-lsr', '--lochness_sync_receive',
                         default=False,
                         action='store_true',

scripts/sync.py

@@ -166,7 +166,13 @@ def do(args, Lochness):
     # transfer new files after all sync attempts are done
     if args.lochness_sync_send:
         if args.s3:
+            # for data under GENERAL
             lochness_to_lochness_transfer_s3(Lochness)
+
+            # for data under PROTECTED (for selected datatypes)
+            if 's3_selective_sync' in Lochness:
+                lochness_to_lochness_transfer_s3_protected(Lochness)
+
         elif args.rsync:
             lochness_to_lochness_transfer_rsync(Lochness)
         else:

tests/lochness_test/transfer/test_transfer.py

@@ -7,6 +7,7 @@
 from lochness.transfer import decompress_transferred_file_and_copy
 from lochness.transfer import lochness_to_lochness_transfer_receive_sftp
 from lochness.transfer import lochness_to_lochness_transfer_s3
+from lochness.transfer import lochness_to_lochness_transfer_s3_protected
 
 from pathlib import Path
 
@@ -335,3 +336,43 @@ def test_s3_sync_function(Lochness):
     print(os.popen(command).read())
 
 
+def test_lochness_to_lochness_transfer_s3_protected():
+    args = Args('tmp_lochness')
+    args.rsync = True
+    args.s3_selective_sync = ['actigraphy', 'haha']
+    create_lochness_template(args)
+    k = KeyringAndEncrypt(args.outdir)
+    k.update_var_subvars('lochness_sync', 'transfer')
+
+    lochness = config_load_test('tmp_lochness/config.yml', '')
+    print(lochness)
+    # print(lochness)
+
+    lochness['AWS_BUCKET_NAME'] = 'ampscz-dev'
+    lochness['AWS_BUCKET_ROOT'] = 'TEST_PHOENIX_ROOT'
+
+    # create fake files
+    tmp_file = Path(lochness['phoenix_root']) / \
+            'PROTECTED/StudyA/raw/subject01/actigraphy/haha.txt'
+    tmp_file.parent.mkdir(parents=True, exist_ok=True)
+    tmp_file.touch()
+
+    tmp_file = Path(lochness['phoenix_root']) / \
+            'PROTECTED/StudyA/raw/subject01/actigraphy/haha2.txt'
+    tmp_file.touch()
+
+    tmp_file = Path(lochness['phoenix_root']) / \
+            'PROTECTED/StudyA/raw/subject02/actigraphy/haha2.txt'
+    tmp_file.parent.mkdir(parents=True, exist_ok=True)
+    tmp_file.touch()
+
+    tmp_file = Path(lochness['phoenix_root']) / \
+            'PROTECTED/StudyA/processed/subject02/actigraphy/haha2.txt'
+    tmp_file.parent.mkdir(parents=True, exist_ok=True)
+    tmp_file.touch()
+
+    # Lochness['AWS_BUCKET_NAME'] = 'ampscz-dev'
+    # Lochness['AWS_BUCKET_ROOT'] = 'TEST_PHOENIX_ROOT'
+    lochness_to_lochness_transfer_s3_protected(
+            lochness, lochness['s3_selective_sync'])
+

tests/test_lochness.py

@@ -59,6 +59,7 @@ def __init__(self, root_dir):
         self.pii_csv = ''
         self.rsync = False
         self.s3 = False
+        self.s3_selective_sync = False
         self.enter_passwords = False
 
 
@@ -118,7 +119,11 @@ def update_var_subvars(self, var, module_name):
         subvar_list = token.get_var_names(module_name)
 
         for subvar, keyring in zip(subvar_list, keyrings):
-            self.keyring[var][subvar] = keyring
+            if var in self.keyring:
+                self.keyring[var][subvar] = keyring
+            else:
+                self.keyring[var] = {}
+                self.keyring[var][subvar] = keyring
 
         self.write_keyring_and_encrypt()