Merge pull request #154 from boegel/ingestion_staging_buckets

update ingestion script to support multiple S3 buckets

Author: bedroge

scripts/automated_ingestion/automated_ingestion.cfg.example

@@ -10,7 +10,7 @@ ingestion_script = /absolute/path/to/ingest-tarball.sh
 metadata_file_extension = .meta.txt
 
 [aws]
-staging_bucket = eessi-staging
+staging_buckets = eessi-staging, eessi-staging-2023.06
 
 [cvmfs]
 ingest_as_root = yes

scripts/automated_ingestion/automated_ingestion.py

@@ -17,7 +17,7 @@
 REQUIRED_CONFIG = {
     'secrets': ['aws_secret_access_key', 'aws_access_key_id', 'github_pat'],
     'paths': ['download_dir', 'ingestion_script', 'metadata_file_extension'],
-    'aws': ['staging_bucket'],
+    'aws': ['staging_buckets'],
     'github': ['staging_repo', 'failed_ingestion_issue_body', 'pr_body'],
 }
 
@@ -39,10 +39,9 @@ def error(msg, code=1):
 def find_tarballs(s3, bucket, extension='.tar.gz', metadata_extension='.meta.txt'):
     """Return a list of all tarballs in an S3 bucket that have a metadata file with the given extension (and same filename)."""
     # TODO: list_objects_v2 only returns up to 1000 objects
-    files = [
-        object['Key']
-        for object in s3.list_objects_v2(Bucket=bucket)['Contents']
-    ]
+    s3_objects = s3.list_objects_v2(Bucket=bucket).get('Contents', [])
+    files = [obj['Key'] for obj in s3_objects]
+
     tarballs = [
         file
         for file in files
@@ -100,15 +99,16 @@ def main():
         aws_secret_access_key=config['secrets']['aws_secret_access_key'],
     )
 
-    tarballs = find_tarballs(s3, config['aws']['staging_bucket'])
-    if args.list_only:
-        for num, tarball in enumerate(tarballs):
-            print(f'{num}: {tarball}')
-        sys.exit(0)
-
-    for tarball in tarballs:
-        tar = EessiTarball(tarball, config, gh, s3)
-        tar.run_handler()
+    buckets = [x.strip() for x in config['aws']['staging_buckets'].split(',')]
+    for bucket in buckets:
+        tarballs = find_tarballs(s3, bucket)
+        if args.list_only:
+            for num, tarball in enumerate(tarballs):
+                print(f'{num}: {tarball}')
+        else:
+            for tarball in tarballs:
+                tar = EessiTarball(tarball, config, gh, s3, bucket)
+                tar.run_handler()
 
 
 if __name__ == '__main__':

scripts/automated_ingestion/eessitarball.py

@@ -19,17 +19,18 @@ class EessiTarball:
     for which it interfaces with the S3 bucket, GitHub, and CVMFS.
     """
 
-    def __init__(self, object_name, config, github, s3):
+    def __init__(self, object_name, config, github, s3, bucket):
         """Initialize the tarball object."""
         self.config = config
         self.github = github
         self.git_repo = github.get_repo(config['github']['staging_repo'])
         self.metadata_file = object_name + config['paths']['metadata_file_extension']
         self.object = object_name
         self.s3 = s3
+        self.bucket = bucket
         self.local_path = os.path.join(config['paths']['download_dir'], os.path.basename(object_name))
         self.local_metadata_path = self.local_path + config['paths']['metadata_file_extension']
-        self.url = f'https://{config["aws"]["staging_bucket"]}.s3.amazonaws.com/{object_name}'
+        self.url = f'https://{bucket}.s3.amazonaws.com/{object_name}'
 
         self.states = {
             'new': {'handler': self.mark_new_tarball_as_staged, 'next_state': 'staged'},
@@ -47,21 +48,20 @@ def download(self, force=False):
         """
         Download this tarball and its corresponding metadata file, if this hasn't been already done.
         """
-        bucket = self.config['aws']['staging_bucket']
         if force or not os.path.exists(self.local_path):
             try:
-                self.s3.download_file(bucket, self.object, self.local_path)
+                self.s3.download_file(self.bucket, self.object, self.local_path)
             except:
                 logging.error(
-                    f'Failed to download tarball {self.object} from {bucket} to {self.local_path}.'
+                    f'Failed to download tarball {self.object} from {self.bucket} to {self.local_path}.'
                 )
                 self.local_path = None
         if force or not os.path.exists(self.local_metadata_path):
             try:
-                self.s3.download_file(bucket, self.metadata_file, self.local_metadata_path)
+                self.s3.download_file(self.bucket, self.metadata_file, self.local_metadata_path)
             except:
                 logging.error(
-                    f'Failed to download metadata file {self.metadata_file} from {bucket} to {self.local_metadata_path}.'
+                    f'Failed to download metadata file {self.metadata_file} from {self.bucket} to {self.local_metadata_path}.'
                 )
                 self.local_metadata_path = None