test environment from bm

lsst-uk · Nov 19, 2024 · 2f1cbbf · 2f1cbbf
1 parent f28b20a
commit 2f1cbbf
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 27 deletions.
diff --git a/bucket_manager/bucket_manager.py b/bucket_manager/bucket_manager.py
@@ -65,9 +65,12 @@ def get_resource(access_key: str, secret_key: str, s3_host: str):
     Returns:
     An S3 resource object.
     """
-    creds = {'access_key': os.environ['S3_ACCESS_KEY'],
-             'secret_key': os.environ['S3_SECRET_KEY'],
-             'host_url': os.environ['S3_HOST_URL']}
+    try:
+        creds = {'access_key': os.environ['S3_ACCESS_KEY'],
+            'secret_key': os.environ['S3_SECRET_KEY'],
+            'host_url': os.environ['S3_HOST_URL']}
+    except KeyError as e:
+        raise KeyError('Set S3_ACCESS_KEY, S3_SECRET_KEY and S3_HOST_URL environment variables.')
     session = boto3.Session(
         aws_access_key_id=creds['access_key'],
         aws_secret_access_key=creds['secret_key']
@@ -90,9 +93,12 @@ def get_client():
     Returns:
     An S3 client object.
     """
-    creds = {'access_key': os.environ['S3_ACCESS_KEY'],
-             'secret_key': os.environ['S3_SECRET_KEY'],
-             'host_url': os.environ['S3_HOST_URL']}
+    try:
+        creds = {'access_key': os.environ['S3_ACCESS_KEY'],
+            'secret_key': os.environ['S3_SECRET_KEY'],
+            'host_url': os.environ['S3_HOST_URL']}
+    except KeyError as e:
+        raise KeyError('Set S3_ACCESS_KEY, S3_SECRET_KEY and S3_HOST_URL environment variables.')
     session = boto3.Session(
         aws_access_key_id=creds['access_key'],
         aws_secret_access_key=creds['secret_key']
@@ -237,9 +243,12 @@ def get_conn_swift() -> swiftclient.Connection:
     Returns:
     A Swift connection object.
     """
-    creds = {'user': os.environ['ST_USER'],
-             'key': os.environ['ST_SECRET'],
-             'authurl': os.environ['ST_AUTH']}
+    try:
+        creds = {'user': os.environ['ST_USER'],
+            'key': os.environ['ST_SECRET'],
+            'authurl': os.environ['ST_AUTH']}
+    except KeyError as e:
+        raise KeyError('Set ST_USER, ST_SECRET and ST_AUTH environment variables.')
     return swiftclient.Connection(
         user=creds['user'],
         key=creds['key'],
@@ -258,9 +267,12 @@ def get_service_swift() -> swiftclient.service.SwiftService:
     Returns:
     A Swift service object.
     """
-    creds = {'user': os.environ['ST_USER'],
-             'key': os.environ['ST_SECRET'],
+    try:
+        creds = {'user': os.environ['ST_USER'],
+             'key': os.environ['ST_KEY'],
              'authurl': os.environ['ST_AUTH']}
+    except KeyError as e:
+        raise KeyError('Set ST_USER, ST_SECRET and ST_AUTH environment variables.')
     return swiftclient.service.SwiftService(
         {
             'auth_version': '1',

diff --git a/csd3-side/scripts/lsst-backup.py b/csd3-side/scripts/lsst-backup.py
@@ -1502,31 +1502,33 @@ def error(self, message):
                 logfile.write('LOCAL_FOLDER,LOCAL_PATH,FILE_SIZE,BUCKET_NAME,DESTINATION_KEY,CHECKSUM,ZIP_CONTENTS\n')
 
     # Setup bucket
-    if api == 's3':
-        s3_host = 'https://echo.stfc.ac.uk'
-    elif api == 'swift':
-        s3_host = 'https://s3.echo.stfc.ac.uk/auth/1.0'
-    print(f'Using {api.capitalize()} API with host {s3_host}')
     try:
-        keys = bm.get_keys(api)
+        if bm.check_keys(api):
+            if api == 's3':
+                s3_host = 'https://echo.stfc.ac.uk'
+                access_key = os.environ['S3_ACCESS_KEY']
+                secret_key = os.environ['S3_ACCESS_KEY']
+            elif api == 'swift':
+                access_key = os.environ['ST_USER']
+                secret_key = os.environ['ST_KEY']
+                os.environ['ST_AUTH'] = s3_host
+                os.environ['ST_USER'] = access_key
+                os.environ['ST_KEY'] = secret_key
+                s3_host = 'https://s3.echo.stfc.ac.uk/auth/1.0'
+    except AssertionError as e:
+        print(f'AssertionError {e}', file=sys.stderr)
+        sys.exit()
     except KeyError as e:
         print(f'KeyError {e}', file=sys.stderr)
         sys.exit()
     except ValueError as e:
         print(f'ValueError {e}', file=sys.stderr)
         sys.exit()
-    if api == 's3':
-        access_key = keys['access_key']
-        secret_key = keys['secret_key']
-    elif api == 'swift':
-        access_key = keys['user']
-        secret_key = keys['secret_key']
-        os.environ['ST_AUTH'] = s3_host
-        os.environ['ST_USER'] = access_key
-        os.environ['ST_KEY'] = secret_key
+
+    print(f'Using {api.capitalize()} API with host {s3_host}')
 
     if api == 's3':
-        s3 = bm.get_resource(access_key, secret_key, s3_host)
+        s3 = bm.get_resource()
         bucket_list = bm.bucket_list(s3)
     elif api == 'swift':
         s3 = bm.get_conn_swift(access_key, secret_key, s3_host)