cosalib: create libs for s3 and ec2 operations

These libraries contain code to work on s3 buckets (head, list, delete)
and ec2 (deregister images and snapshots)

Author: vrutkovs

src/cmd-buildprep

@@ -9,15 +9,14 @@ import os
 import subprocess
 import sys
 import requests
-import boto3
 import shutil
-from botocore.exceptions import ClientError
 from tenacity import retry, retry_if_exception_type
 
 sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
 
 from cosalib.builds import Builds, BUILDFILES
-from cosalib.cmdlib import load_json, rm_allow_noent, retry_stop, retry_s3_exception, retry_callback  # noqa: E402
+from cosalib.cmdlib import load_json, rm_allow_noent, retry_stop, retry_callback
+from cosalib.s3 import download_file, head_bucket, head_object
 
 retry_requests_exception = (retry_if_exception_type(requests.Timeout) |
                             retry_if_exception_type(requests.ReadTimeout) |
@@ -141,6 +140,23 @@ class Fetcher(object):
         return self.exists_impl(url)
 
 
+class S3Fetcher(Fetcher):
+
+    def fetch_impl(self, url, dest):
+        assert url.startswith("s3://")
+        bucket, key = url[len("s3://"):].split('/', 1)
+        # this function does not need to be retried with the decorator as download_file would
+        # retry automatically based on s3config settings
+        download_file(bucket, key, dest)
+
+    def exists_impl(self, url):
+        assert url.startswith("s3://")
+        bucket, key = url[len("s3://"):].split('/', 1)
+        # sanity check that the bucket exists and we have access to it
+        head_bucket(bucket=bucket)
+        return head_object(bucket=bucket, key=key)
+
+
 class HTTPFetcher(Fetcher):
 
     def __init__(self, url_base):
@@ -165,37 +181,6 @@ class HTTPFetcher(Fetcher):
             raise Exception(f"Received rc {r.status_code} for {url}")
 
 
-class S3Fetcher(Fetcher):
-
-    def __init__(self, url_base):
-        super().__init__(url_base)
-        self.s3 = boto3.client('s3')
-        self.s3config = boto3.s3.transfer.TransferConfig(
-            num_download_attempts=5
-        )
-
-    def fetch_impl(self, url, dest):
-        assert url.startswith("s3://")
-        bucket, key = url[len("s3://"):].split('/', 1)
-        # this function does not need to be retried with the decorator as download_file would
-        # retry automatically based on s3config settings
-        self.s3.download_file(bucket, key, dest, Config=self.s3config)
-
-    @retry(stop=retry_stop, retry=retry_s3_exception, before_sleep=retry_callback)
-    def exists_impl(self, url):
-        assert url.startswith("s3://")
-        bucket, key = url[len("s3://"):].split('/', 1)
-        # sanity check that the bucket exists and we have access to it
-        self.s3.head_bucket(Bucket=bucket)
-        try:
-            self.s3.head_object(Bucket=bucket, Key=key)
-        except ClientError as e:
-            if e.response['Error']['Code'] == '404':
-                return False
-            raise e
-        return True
-
-
 class LocalFetcher(Fetcher):
 
     def __init__(self, url_base):

src/cmd-buildupload

@@ -22,7 +22,7 @@ CACHE_MAX_AGE_ARTIFACT = 60 * 60 * 24 * 365
 # set metadata caching to 5m
 CACHE_MAX_AGE_METADATA = 60 * 5
 from cosalib.builds import Builds, BUILDFILES
-from cosalib.cmdlib import load_json, retry_stop, retry_s3_exception, retry_callback  # noqa: E402
+from cosalib.cmdlib import load_json, retry_stop, retry_boto_exception, retry_callback  # noqa: E402
 
 
 def main():
@@ -149,7 +149,7 @@ def s3_upload_build(args, builddir, bucket, prefix):
             dry_run=args.dry_run)
 
 
-@retry(stop=retry_stop, retry=retry_s3_exception, before_sleep=retry_callback)
+@retry(stop=retry_stop, retry=retry_boto_exception, before_sleep=retry_callback)
 def s3_check_exists(bucket, key):
     print(f"Checking if bucket '{bucket}' has key '{key}'")
     s3 = boto3.client('s3')
@@ -162,7 +162,7 @@ def s3_check_exists(bucket, key):
     return True
 
 
-@retry(stop=retry_stop, retry=retry_s3_exception, retry_error_callback=retry_callback)
+@retry(stop=retry_stop, retry=retry_boto_exception, retry_error_callback=retry_callback)
 def s3_copy(src, bucket, key, max_age, acl, extra_args={}, dry_run=False):
     if key.endswith('.json') and 'ContentType' not in extra_args:
         extra_args['ContentType'] = 'application/json'

src/cosalib/aws.py

@@ -0,0 +1,19 @@
+import boto3
+from cosalib.cmdlib import (
+    retry_stop,
+    retry_boto_exception,
+    retry_callback
+)
+from tenacity import retry
+
+
+@retry(stop=retry_stop, retry=retry_boto_exception, before_sleep=retry_callback)
+def deregister_ami(ami_id, region):
+    ec2 = boto3.client('ec2', region_name=region)
+    ec2.deregister_image(ImageId=ami_id)
+
+
+@retry(stop=retry_stop, retry=retry_boto_exception, before_sleep=retry_callback)
+def delete_snapshot(snap_id, region):
+    ec2 = boto3.client('ec2', region_name=region)
+    ec2.delete_snapshot(SnapshotId=snap_id)

src/cosalib/cmdlib.py

@@ -26,7 +26,7 @@
 from datetime import datetime
 
 retry_stop = (stop_after_delay(10) | stop_after_attempt(5))
-retry_s3_exception = (retry_if_exception_type(ConnectionClosedError) |
+retry_boto_exception = (retry_if_exception_type(ConnectionClosedError) |
                       retry_if_exception_type(ConnectTimeoutError) |
                       retry_if_exception_type(IncompleteReadError) |
                       retry_if_exception_type(ReadTimeoutError))

src/cosalib/s3.py

@@ -0,0 +1,59 @@
+import boto3
+
+from botocore.exceptions import ClientError
+from cosalib.cmdlib import (
+    retry_stop,
+    retry_boto_exception,
+    retry_callback
+)
+from tenacity import retry
+
+
+s3 = boto3.client('s3')
+s3config = boto3.s3.transfer.TransferConfig(
+    num_download_attempts=5
+)
+
+
+def download_file(bucket, key, dest):
+    s3.download_file(bucket, key, dest, Config=s3config)
+
+
+@retry(stop=retry_stop, retry=retry_boto_exception, before_sleep=retry_callback)
+def head_bucket(bucket):
+    s3.head_bucket(Bucket=bucket)
+
+
+@retry(stop=retry_stop, retry=retry_boto_exception, before_sleep=retry_callback)
+def head_object(bucket, key):
+    try:
+        s3.head_object(Bucket=bucket, Key=key)
+    except ClientError as e:
+        if e.response['Error']['Code'] == '404':
+            return False
+        raise e
+    return True
+
+
+@retry(stop=retry_stop, retry=retry_boto_exception, before_sleep=retry_callback)
+def list_objects(bucket, prefix, delimiter="/", result_key='Contents'):
+    kwargs = {
+        'Bucket': bucket,
+        'Delimiter': delimiter,
+        'Prefix': prefix,
+    }
+    isTruncated = True
+    while isTruncated:
+        batch = s3.list_objects_v2(**kwargs)
+        yield from batch.get(result_key) or []
+        kwargs['ContinuationToken'] = batch.get('NextContinuationToken')
+        isTruncated = batch['IsTruncated']
+
+
+@retry(stop=retry_stop, retry=retry_boto_exception, before_sleep=retry_callback)
+def delete_object(bucket, key):
+    sub_objects = list(list_objects(bucket, key))
+    if sub_objects != []:
+        print("s3: deleting {sub_objects}")
+        s3.delete_objects(Bucket=bucket, Delete=sub_objects)
+    s3.delete_object(Bucket=bucket, Key=key)