Simplify s3 functions: create_s3_bucket.

src/toil/common.py

@@ -69,7 +69,7 @@
                       QueueSizeMessage,
                       gen_message_bus_path)
 from toil.fileStores import FileID
-from toil.lib.aws import zone_to_region, build_tag_dict_from_env
+from toil.lib.aws import zone_to_region
 from toil.lib.compatibility import deprecated
 from toil.lib.conversions import bytes2human, human2bytes
 from toil.lib.io import try_path

src/toil/jobStores/aws/jobStore.py

@@ -13,16 +13,13 @@
 # limitations under the License.
 import hashlib
 import itertools
-import json
 import logging
 import os
 import pickle
 import re
 import reprlib
 import stat
 import time
-import urllib.error
-import urllib.request
 import uuid
 from contextlib import contextmanager
 from io import BytesIO
@@ -35,7 +32,6 @@
 from botocore.exceptions import ClientError
 
 import toil.lib.encryption as encryption
-from toil.lib.aws import build_tag_dict_from_env
 from toil.fileStores import FileID
 from toil.jobStores.abstractJobStore import (AbstractJobStore,
                                              ConcurrentFileModificationException,
@@ -57,10 +53,8 @@
                                   ReadableTransformingPipe,
                                   WritablePipe)
 from toil.lib.aws.session import establish_boto3_session
-from toil.lib.aws.utils import (create_s3_bucket,
-                                enable_public_objects,
-                                flatten_tags,
-                                get_bucket_region,
+from toil.lib.aws.s3 import create_s3_bucket
+from toil.lib.aws.utils import (get_bucket_region,
                                 get_object_for_url,
                                 list_objects_for_url,
                                 retry_s3,
@@ -721,85 +715,58 @@ def bucket_retry_predicate(error):
             return False
 
         bucketExisted = True
-        for attempt in retry_s3(predicate=bucket_retry_predicate):
-            with attempt:
-                try:
-                    # the head_bucket() call makes sure that the bucket exists and the user can access it
-                    self.s3_client.head_bucket(Bucket=bucket_name)
-
-                    bucket = self.s3_resource.Bucket(bucket_name)
-                except ClientError as e:
-                    error_http_status = get_error_status(e)
-                    if error_http_status == 404:
-                        bucketExisted = False
-                        logger.debug("Bucket '%s' does not exist.", bucket_name)
-                        if create:
-                            bucket = create_s3_bucket(
-                                self.s3_resource, bucket_name, self.region
-                            )
-                            # Wait until the bucket exists before checking the region and adding tags
-                            bucket.wait_until_exists()
-
-                            # It is possible for create_bucket to return but
-                            # for an immediate request for the bucket region to
-                            # produce an S3ResponseError with code
-                            # NoSuchBucket. We let that kick us back up to the
-                            # main retry loop.
-                            assert (
-                                get_bucket_region(bucket_name) == self.region
-                            ), f"bucket_name: {bucket_name}, {get_bucket_region(bucket_name)} != {self.region}"
-
-                            tags = build_tag_dict_from_env()
-
-                            if tags:
-                                flat_tags = flatten_tags(tags)
-                                bucket_tagging = self.s3_resource.BucketTagging(bucket_name)
-                                bucket_tagging.put(Tagging={'TagSet': flat_tags})
-
-                            # Configure bucket so that we can make objects in
-                            # it public, which was the historical default.
-                            enable_public_objects(bucket_name)
-                        elif block:
-                            raise
-                        else:
-                            return None
-                    elif error_http_status == 301:
-                        # This is raised if the user attempts to get a bucket in a region outside
-                        # the specified one, if the specified one is not `us-east-1`.  The us-east-1
-                        # server allows a user to use buckets from any region.
-                        raise BucketLocationConflictException(get_bucket_region(bucket_name))
-                    else:
-                        raise
-                else:
-                    bucketRegion = get_bucket_region(bucket_name)
-                    if bucketRegion != self.region:
-                        raise BucketLocationConflictException(bucketRegion)
-
-                if versioning and not bucketExisted:
-                    # only call this method on bucket creation
-                    bucket.Versioning().enable()
-                    # Now wait until versioning is actually on. Some uploads
-                    # would come back with no versions; maybe they were
-                    # happening too fast and this setting isn't sufficiently
-                    # consistent?
-                    time.sleep(1)
-                    while not self._getBucketVersioning(bucket_name):
-                        logger.warning(f"Waiting for versioning activation on bucket '{bucket_name}'...")
-                        time.sleep(1)
-                elif check_versioning_consistency:
-                    # now test for versioning consistency
-                    # we should never see any of these errors since 'versioning' should always be true
-                    bucket_versioning = self._getBucketVersioning(bucket_name)
-                    if bucket_versioning != versioning:
-                        assert False, 'Cannot modify versioning on existing bucket'
-                    elif bucket_versioning is None:
-                        assert False, 'Cannot use a bucket with versioning suspended'
-                if bucketExisted:
-                    logger.debug(f"Using pre-existing job store bucket '{bucket_name}'.")
+        try:
+            # the head_bucket() call makes sure that the bucket exists and the user can access it
+            self.s3_client.head_bucket(Bucket=bucket_name)
+            bucket = self.s3_resource.Bucket(bucket_name)
+        except ClientError as e:
+            error_http_status = get_error_status(e)
+            if error_http_status == 404:
+                bucketExisted = False
+                logger.debug("Bucket '%s' does not exist.", bucket_name)
+                if create:
+                    bucket = create_s3_bucket(self.s3_resource, bucket_name, self.region)
+                elif block:
+                    raise
                 else:
-                    logger.debug(f"Created new job store bucket '{bucket_name}' with versioning state {versioning}.")
+                    return None
+            elif error_http_status == 301:
+                # This is raised if the user attempts to get a bucket in a region outside
+                # the specified one, if the specified one is not `us-east-1`.  The us-east-1
+                # server allows a user to use buckets from any region.
+                raise BucketLocationConflictException(get_bucket_region(bucket_name))
+            else:
+                raise
+        else:
+            bucketRegion = get_bucket_region(bucket_name)
+            if bucketRegion != self.region:
+                raise BucketLocationConflictException(bucketRegion)
+
+        if versioning and not bucketExisted:
+            # only call this method on bucket creation
+            bucket.Versioning().enable()
+            # Now wait until versioning is actually on. Some uploads
+            # would come back with no versions; maybe they were
+            # happening too fast and this setting isn't sufficiently
+            # consistent?
+            time.sleep(1)
+            while not self._getBucketVersioning(bucket_name):
+                logger.warning(f"Waiting for versioning activation on bucket '{bucket_name}'...")
+                time.sleep(1)
+        elif check_versioning_consistency:
+            # now test for versioning consistency
+            # we should never see any of these errors since 'versioning' should always be true
+            bucket_versioning = self._getBucketVersioning(bucket_name)
+            if bucket_versioning != versioning:
+                assert False, 'Cannot modify versioning on existing bucket'
+            elif bucket_versioning is None:
+                assert False, 'Cannot use a bucket with versioning suspended'
+        if bucketExisted:
+            logger.debug(f"Using pre-existing job store bucket '{bucket_name}'.")
+        else:
+            logger.debug(f"Created new job store bucket '{bucket_name}' with versioning state {versioning}.")
 
-                return bucket
+        return bucket
 
     def _bindDomain(self, domain_name, create=False, block=True):
         """

src/toil/lib/aws/__init__.py

@@ -168,29 +168,25 @@ def file_begins_with(path, prefix):
     except (URLError, socket.timeout, HTTPException):
         return False
 
+
 def running_on_ecs() -> bool:
     """
     Return True if we are currently running on Amazon ECS, and false otherwise.
     """
     # We only care about relatively current ECS
     return 'ECS_CONTAINER_METADATA_URI_V4' in os.environ
 
-def build_tag_dict_from_env(environment: MutableMapping[str, str] = os.environ) -> Dict[str, str]:
-    tags = dict()
-    owner_tag = environment.get('TOIL_OWNER_TAG')
+
+def tags_from_env() -> Dict[str, str]:
+    try:
+        tags = json.loads(os.environ.get('TOIL_AWS_TAGS', '{}'))
+    except json.decoder.JSONDecodeError:
+        logger.error('TOIL_AWS_TAGS must be in JSON format: {"key" : "value", ...}')
+        exit(1)
+
+    # TODO: Remove TOIL_OWNER_TAG and only use TOIL_AWS_TAGS .
+    owner_tag = os.environ.get('TOIL_OWNER_TAG')
     if owner_tag:
         tags.update({'Owner': owner_tag})
 
-    user_tags = environment.get('TOIL_AWS_TAGS')
-    if user_tags:
-        try:
-            json_user_tags = json.loads(user_tags)
-            if isinstance(json_user_tags, dict):
-                tags.update(json.loads(user_tags))
-            else:
-                logger.error('TOIL_AWS_TAGS must be in JSON format: {"key" : "value", ...}')
-                exit(1)
-        except json.decoder.JSONDecodeError:
-            logger.error('TOIL_AWS_TAGS must be in JSON format: {"key" : "value", ...}')
-            exit(1)
     return tags

src/toil/lib/aws/s3.py

@@ -0,0 +1,84 @@
+# Copyright (C) 2015-2023 Regents of the University of California
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import logging
+import sys
+
+from typing import (Any,
+                    Dict,
+                    Optional,
+                    Union)
+
+from toil.lib.retry import retry
+from toil.lib.aws import tags_from_env
+from toil.lib.aws.utils import enable_public_objects, flatten_tags
+
+if sys.version_info >= (3, 8):
+    from typing import Literal
+else:
+    from typing_extensions import Literal
+
+try:
+    from boto.exception import BotoServerError, S3ResponseError
+    from botocore.exceptions import ClientError
+    from mypy_boto3_iam import IAMClient, IAMServiceResource
+    from mypy_boto3_s3 import S3Client, S3ServiceResource
+    from mypy_boto3_s3.literals import BucketLocationConstraintType
+    from mypy_boto3_s3.service_resource import Bucket, Object
+    from mypy_boto3_sdb import SimpleDBClient
+except ImportError:
+    BotoServerError = Exception  # type: ignore
+    S3ResponseError = Exception  # type: ignore
+    ClientError = Exception  # type: ignore
+    # AWS/boto extra is not installed
+
+
+logger = logging.getLogger(__name__)
+
+
+@retry(errors=[BotoServerError, S3ResponseError, ClientError])
+def create_s3_bucket(
+    s3_resource: "S3ServiceResource",
+    bucket_name: str,
+    region: Union["BucketLocationConstraintType", Literal["us-east-1"]],
+    tags: Optional[Dict[str, str]] = None,
+    public: bool = True
+) -> "Bucket":
+    """
+    Create an AWS S3 bucket, using the given Boto3 S3 session, with the
+    given name, in the given region.
+
+    Supports the us-east-1 region, where bucket creation is special.
+
+    *ALL* S3 bucket creation should use this function.
+    """
+    logger.debug("Creating bucket '%s' in region %s.", bucket_name, region)
+    if region == "us-east-1":  # see https://github.com/boto/boto3/issues/125
+        bucket = s3_resource.create_bucket(Bucket=bucket_name)
+    else:
+        bucket = s3_resource.create_bucket(
+            Bucket=bucket_name,
+            CreateBucketConfiguration={"LocationConstraint": region},
+        )
+    # wait until the bucket exists before adding tags
+    bucket.wait_until_exists()
+
+    tags = tags_from_env() if tags is None else tags
+    bucket_tagging = s3_resource.BucketTagging(bucket_name)
+    bucket_tagging.put(Tagging={'TagSet': flatten_tags(tags)})  # type: ignore
+
+    # enabling public objects is the historical default
+    if public:
+        enable_public_objects(bucket_name)
+
+    return bucket

src/toil/lib/aws/utils.py

@@ -195,29 +195,6 @@ def delete_s3_bucket(
         printq(f'\n * S3 bucket no longer exists: {bucket}\n\n', quiet)
 
 
-def create_s3_bucket(
-    s3_resource: "S3ServiceResource",
-    bucket_name: str,
-    region: Union["BucketLocationConstraintType", Literal["us-east-1"]],
-) -> "Bucket":
-    """
-    Create an AWS S3 bucket, using the given Boto3 S3 session, with the
-    given name, in the given region.
-
-    Supports the us-east-1 region, where bucket creation is special.
-
-    *ALL* S3 bucket creation should use this function.
-    """
-    logger.debug("Creating bucket '%s' in region %s.", bucket_name, region)
-    if region == "us-east-1":  # see https://github.com/boto/boto3/issues/125
-        bucket = s3_resource.create_bucket(Bucket=bucket_name)
-    else:
-        bucket = s3_resource.create_bucket(
-            Bucket=bucket_name,
-            CreateBucketConfiguration={"LocationConstraint": region},
-        )
-    return bucket
-
 @retry(errors=[ClientError])
 def enable_public_objects(bucket_name: str) -> None:
     """

src/toil/provisioners/aws/awsProvisioner.py

@@ -50,7 +50,7 @@
                               get_policy_permissions,
                               policy_permissions_allow)
 from toil.lib.aws.session import AWSConnectionManager
-from toil.lib.aws.utils import create_s3_bucket
+from toil.lib.aws.s3 import create_s3_bucket
 from toil.lib.conversions import human2bytes
 from toil.lib.ec2 import (a_short_time,
                           create_auto_scaling_group,
@@ -259,14 +259,8 @@ def _write_file_to_cloud(self, key: str, contents: bytes) -> str:
             bucket = s3.Bucket(bucket_name)
         except ClientError as err:
             if get_error_status(err) == 404:
-                bucket = create_s3_bucket(s3, bucket_name=bucket_name, region=self._region)
-                bucket.wait_until_exists()
+                bucket = create_s3_bucket(s3, bucket_name, self._region)
                 bucket.Versioning().enable()
-
-                owner_tag = os.environ.get('TOIL_OWNER_TAG')
-                if owner_tag:
-                    bucket_tagging = s3.BucketTagging(bucket_name)
-                    bucket_tagging.put(Tagging={'TagSet': [{'Key': 'Owner', 'Value': owner_tag}]})
             else:
                 raise