toil 6.1.0a1__py3-none-any.whl → 7.0.0__py3-none-any.whl

toil/lib/aws/session.py

@@ -15,16 +15,21 @@ import collections
 import logging
 import os
 import threading
-from typing import Dict, Optional, Tuple, cast
+from typing import Dict, Optional, Tuple, cast, Union, Literal, overload, TypeVar
 
 import boto3
 import boto3.resources.base
-import boto.connection
 import botocore
 from boto3 import Session
 from botocore.client import Config
 from botocore.session import get_session
 from botocore.utils import JSONFileCache
+from mypy_boto3_autoscaling import AutoScalingClient
+from mypy_boto3_ec2 import EC2Client, EC2ServiceResource
+from mypy_boto3_iam import IAMClient, IAMServiceResource
+from mypy_boto3_s3 import S3Client, S3ServiceResource
+from mypy_boto3_sdb import SimpleDBClient
+from mypy_boto3_sts import STSClient
 
 logger = logging.getLogger(__name__)
 
@@ -120,6 +125,13 @@ class AWSConnectionManager:
             storage.item = _new_boto3_session(region_name=region)
         return cast(boto3.session.Session, storage.item)
 
+    @overload
+    def resource(self, region: Optional[str], service_name: Literal["s3"], endpoint_url: Optional[str] = None) -> S3ServiceResource: ...
+    @overload
+    def resource(self, region: Optional[str], service_name: Literal["iam"], endpoint_url: Optional[str] = None) -> IAMServiceResource: ...
+    @overload
+    def resource(self, region: Optional[str], service_name: Literal["ec2"], endpoint_url: Optional[str] = None) -> EC2ServiceResource: ...
+
     def resource(self, region: Optional[str], service_name: str, endpoint_url: Optional[str] = None) -> boto3.resources.base.ServiceResource:
         """
         Get the Boto3 Resource to use with the given service (like 'ec2') in the given region.
@@ -146,7 +158,28 @@ class AWSConnectionManager:
 
         return cast(boto3.resources.base.ServiceResource, storage.item)
 
-    def client(self, region: Optional[str], service_name: str, endpoint_url: Optional[str] = None, config: Optional[Config] = None) -> botocore.client.BaseClient:
+    @overload
+    def client(self, region: Optional[str], service_name: Literal["ec2"], endpoint_url: Optional[str] = None,
+               config: Optional[Config] = None) -> EC2Client: ...
+    @overload
+    def client(self, region: Optional[str], service_name: Literal["iam"], endpoint_url: Optional[str] = None,
+               config: Optional[Config] = None) -> IAMClient: ...
+    @overload
+    def client(self, region: Optional[str], service_name: Literal["s3"], endpoint_url: Optional[str] = None,
+               config: Optional[Config] = None) -> S3Client: ...
+    @overload
+    def client(self, region: Optional[str], service_name: Literal["sts"], endpoint_url: Optional[str] = None,
+               config: Optional[Config] = None) -> STSClient: ...
+    @overload
+    def client(self, region: Optional[str], service_name: Literal["sdb"], endpoint_url: Optional[str] = None,
+               config: Optional[Config] = None) -> SimpleDBClient: ...
+    @overload
+    def client(self, region: Optional[str], service_name: Literal["autoscaling"], endpoint_url: Optional[str] = None,
+               config: Optional[Config] = None) -> AutoScalingClient: ...
+
+
+    def client(self, region: Optional[str], service_name: Literal["ec2", "iam", "s3", "sts", "sdb", "autoscaling"], endpoint_url: Optional[str] = None,
+               config: Optional[Config] = None) -> botocore.client.BaseClient:
         """
         Get the Boto3 Client to use with the given service (like 'ec2') in the given region.
 
@@ -159,9 +192,9 @@ class AWSConnectionManager:
             # Don't try and memoize if a custom config is used
             with _init_lock:
                 if endpoint_url is not None:
-                    return self.session(region).client(service_name, endpoint_url=endpoint_url, config=config) # type: ignore
+                    return self.session(region).client(service_name, endpoint_url=endpoint_url, config=config)
                 else:
-                    return self.session(region).client(service_name, config=config) # type: ignore
+                    return self.session(region).client(service_name, config=config)
 
         key = (region, service_name, endpoint_url)
         storage = self.client_cache[key]
@@ -172,25 +205,12 @@ class AWSConnectionManager:
                 if endpoint_url is not None:
                     # The Boto3 stubs are probably missing an overload here too. See:
                     # <https://github.com/vemel/mypy_boto3_builder/issues/121#issuecomment-1011322636>
-                    storage.item = self.session(region).client(service_name, endpoint_url=endpoint_url) # type: ignore
+                    storage.item = self.session(region).client(service_name, endpoint_url=endpoint_url)
                 else:
                     # We might not be able to pass None to Boto3 and have it be the same as no argument.
-                    storage.item = self.session(region).client(service_name) # type: ignore
+                    storage.item = self.session(region).client(service_name)
         return cast(botocore.client.BaseClient , storage.item)
 
-    def boto2(self, region: Optional[str], service_name: str) -> boto.connection.AWSAuthConnection:
-        """
-        Get the connected boto2 connection for the given region and service.
-        """
-        if service_name == 'iam':
-            # IAM connections are regionless
-            region = 'universal'
-        key = (region, service_name)
-        storage = self.boto2_cache[key]
-        if not hasattr(storage, 'item'):
-            with _init_lock:
-                storage.item = getattr(boto, service_name).connect_to_region(region, profile_name=os.environ.get("TOIL_AWS_PROFILE", None))
-        return cast(boto.connection.AWSAuthConnection, storage.item)
 
 # If you don't want your own AWSConnectionManager, we have a global one and some global functions
 _global_manager = AWSConnectionManager()
@@ -205,7 +225,20 @@ def establish_boto3_session(region_name: Optional[str] = None) -> Session:
     # Just use a global version of the manager. Note that we change the argument order!
     return _global_manager.session(region_name)
 
-def client(service_name: str, region_name: Optional[str] = None, endpoint_url: Optional[str] = None, config: Optional[Config] = None) -> botocore.client.BaseClient:
+@overload
+def client(service_name: Literal["ec2"], region_name: Optional[str] = None, endpoint_url: Optional[str] = None, config: Optional[Config] = None) -> EC2Client: ...
+@overload
+def client(service_name: Literal["iam"], region_name: Optional[str] = None, endpoint_url: Optional[str] = None, config: Optional[Config] = None) -> IAMClient: ...
+@overload
+def client(service_name: Literal["s3"], region_name: Optional[str] = None, endpoint_url: Optional[str] = None, config: Optional[Config] = None) -> S3Client: ...
+@overload
+def client(service_name: Literal["sts"], region_name: Optional[str] = None, endpoint_url: Optional[str] = None, config: Optional[Config] = None) -> STSClient: ...
+@overload
+def client(service_name: Literal["sdb"], region_name: Optional[str] = None, endpoint_url: Optional[str] = None, config: Optional[Config] = None) -> SimpleDBClient: ...
+@overload
+def client(service_name: Literal["autoscaling"], region_name: Optional[str] = None, endpoint_url: Optional[str] = None, config: Optional[Config] = None) -> AutoScalingClient: ...
+
+def client(service_name: Literal["ec2", "iam", "s3", "sts", "sdb", "autoscaling"], region_name: Optional[str] = None, endpoint_url: Optional[str] = None, config: Optional[Config] = None) -> botocore.client.BaseClient:
     """
     Get a Boto 3 client for a particular AWS service, usable by the current thread.
 
@@ -215,7 +248,14 @@ def client(service_name: str, region_name: Optional[str] = None, endpoint_url: O
     # Just use a global version of the manager. Note that we change the argument order!
     return _global_manager.client(region_name, service_name, endpoint_url=endpoint_url, config=config)
 
-def resource(service_name: str, region_name: Optional[str] = None, endpoint_url: Optional[str] = None) -> boto3.resources.base.ServiceResource:
+@overload
+def resource(service_name: Literal["s3"], region_name: Optional[str] = None, endpoint_url: Optional[str] = None) -> S3ServiceResource: ...
+@overload
+def resource(service_name: Literal["iam"], region_name: Optional[str] = None, endpoint_url: Optional[str] = None) -> IAMServiceResource: ...
+@overload
+def resource(service_name: Literal["ec2"], region_name: Optional[str] = None, endpoint_url: Optional[str] = None) -> EC2ServiceResource: ...
+
+def resource(service_name: Literal["s3", "iam", "ec2"], region_name: Optional[str] = None, endpoint_url: Optional[str] = None) -> boto3.resources.base.ServiceResource:
     """
     Get a Boto 3 resource for a particular AWS service, usable by the current thread.
 

toil/lib/aws/utils.py

@@ -15,7 +15,6 @@ import errno
 import logging
 import os
 import socket
-import sys
 from typing import (Any,
                     Callable,
                     ContextManager,
@@ -25,35 +24,29 @@ from typing import (Any,
                     List,
                     Optional,
                     Set,
-                    Union,
                     cast)
 from urllib.parse import ParseResult
 
-from toil.lib.aws import session
+from mypy_boto3_sdb.type_defs import AttributeTypeDef
+from toil.lib.aws import session, AWSRegionName, AWSServerErrors
 from toil.lib.misc import printq
 from toil.lib.retry import (DEFAULT_DELAYS,
                             DEFAULT_TIMEOUT,
                             get_error_code,
                             get_error_status,
                             old_retry,
-                            retry)
-
-if sys.version_info >= (3, 8):
-    from typing import Literal
-else:
-    from typing_extensions import Literal
+                            retry, ErrorCondition)
 
 try:
-    from boto.exception import BotoServerError, S3ResponseError
-    from botocore.exceptions import ClientError
+    from botocore.exceptions import ClientError, EndpointConnectionError
     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 = None  # type: ignore
     ClientError = None  # type: ignore
+    EndpointConnectionError = None  # type: ignore
     # AWS/boto extra is not installed
 
 logger = logging.getLogger(__name__)
@@ -77,12 +70,10 @@ THROTTLED_ERROR_CODES = [
         'EC2ThrottledException',
 ]
 
-@retry(errors=[BotoServerError])
+@retry(errors=[AWSServerErrors])
 def delete_iam_role(
     role_name: str, region: Optional[str] = None, quiet: bool = True
 ) -> None:
-    from boto.iam.connection import IAMConnection
-
     # TODO: the Boto3 type hints are a bit oversealous here; they want hundreds
     # of overloads of the client-getting methods to exist based on the literal
     # string passed in, to return exactly the right kind of client or resource.
@@ -92,9 +83,8 @@ def delete_iam_role(
     # we wanted MyPy to be able to understand us. So at some point we should
     # consider revising our API here to be less annoying to explain to the type
     # checker.
-    iam_client = cast(IAMClient, session.client('iam', region_name=region))
-    iam_resource = cast(IAMServiceResource, session.resource('iam', region_name=region))
-    boto_iam_connection = IAMConnection()
+    iam_client = session.client('iam', region_name=region)
+    iam_resource = session.resource('iam', region_name=region)
     role = iam_resource.Role(role_name)
     # normal policies
     for attached_policy in role.attached_policies.all():
@@ -103,17 +93,16 @@ def delete_iam_role(
     # inline policies
     for inline_policy in role.policies.all():
         printq(f'Deleting inline policy: {inline_policy.policy_name} from role {role.name}', quiet)
-        # couldn't find an easy way to remove inline policies with boto3; use boto
-        boto_iam_connection.delete_role_policy(role.name, inline_policy.policy_name)
+        iam_client.delete_role_policy(RoleName=role.name, PolicyName=inline_policy.policy_name)
     iam_client.delete_role(RoleName=role_name)
     printq(f'Role {role_name} successfully deleted.', quiet)
 
 
-@retry(errors=[BotoServerError])
+@retry(errors=[AWSServerErrors])
 def delete_iam_instance_profile(
     instance_profile_name: str, region: Optional[str] = None, quiet: bool = True
 ) -> None:
-    iam_resource = cast(IAMServiceResource, session.resource("iam", region_name=region))
+    iam_resource = session.resource("iam", region_name=region)
     instance_profile = iam_resource.InstanceProfile(instance_profile_name)
     if instance_profile.roles is not None:
         for role in instance_profile.roles:
@@ -123,11 +112,11 @@ def delete_iam_instance_profile(
     printq(f'Instance profile "{instance_profile_name}" successfully deleted.', quiet)
 
 
-@retry(errors=[BotoServerError])
+@retry(errors=[AWSServerErrors])
 def delete_sdb_domain(
     sdb_domain_name: str, region: Optional[str] = None, quiet: bool = True
 ) -> None:
-    sdb_client = cast(SimpleDBClient, session.client("sdb", region_name=region))
+    sdb_client = session.client("sdb", region_name=region)
     sdb_client.delete_domain(DomainName=sdb_domain_name)
     printq(f'SBD Domain: "{sdb_domain_name}" successfully deleted.', quiet)
 
@@ -141,16 +130,24 @@ def connection_reset(e: Exception) -> bool:
     # errno is listed as 104. To be safe, we check for both:
     return isinstance(e, socket.error) and e.errno in (errno.ECONNRESET, 104)
 
+def connection_error(e: Exception) -> bool:
+    """
+    Return True if an error represents a failure to make a network connection.
+    """
+    return (connection_reset(e)
+            or isinstance(e, EndpointConnectionError))
+
+
 # TODO: Replace with: @retry and ErrorCondition
 def retryable_s3_errors(e: Exception) -> bool:
     """
     Return true if this is an error from S3 that looks like we ought to retry our request.
     """
-    return (connection_reset(e)
-            or (isinstance(e, BotoServerError) and e.status in (429, 500))
-            or (isinstance(e, BotoServerError) and e.code in THROTTLED_ERROR_CODES)
+    return (connection_error(e)
+            or (isinstance(e, ClientError) and get_error_status(e) in (429, 500))
+            or (isinstance(e, ClientError) and get_error_code(e) in THROTTLED_ERROR_CODES)
             # boto3 errors
-            or (isinstance(e, (S3ResponseError, ClientError)) and get_error_code(e) in THROTTLED_ERROR_CODES)
+            or (isinstance(e, ClientError) and get_error_code(e) in THROTTLED_ERROR_CODES)
             or (isinstance(e, ClientError) and 'BucketNotEmpty' in str(e))
             or (isinstance(e, ClientError) and e.response.get('ResponseMetadata', {}).get('HTTPStatusCode') == 409 and 'try again' in str(e))
             or (isinstance(e, ClientError) and e.response.get('ResponseMetadata', {}).get('HTTPStatusCode') in (404, 429, 500, 502, 503, 504)))
@@ -162,7 +159,7 @@ def retry_s3(delays: Iterable[float] = DEFAULT_DELAYS, timeout: float = DEFAULT_
     """
     return old_retry(delays=delays, timeout=timeout, predicate=predicate)
 
-@retry(errors=[BotoServerError])
+@retry(errors=[AWSServerErrors])
 def delete_s3_bucket(
     s3_resource: "S3ServiceResource",
     bucket: str,
@@ -195,7 +192,7 @@ def delete_s3_bucket(
 def create_s3_bucket(
     s3_resource: "S3ServiceResource",
     bucket_name: str,
-    region: Union["BucketLocationConstraintType", Literal["us-east-1"]],
+    region: AWSRegionName,
 ) -> "Bucket":
     """
     Create an AWS S3 bucket, using the given Boto3 S3 session, with the
@@ -238,7 +235,7 @@ def enable_public_objects(bucket_name: str) -> None:
     would be a very awkward way to do it. So we restore the old behavior.
     """
 
-    s3_client = cast(S3Client, session.client('s3'))
+    s3_client = session.client('s3')
 
     # Even though the new default is for public access to be prohibited, this
     # is implemented by adding new things attached to the bucket. If we remove
@@ -261,7 +258,7 @@ def get_bucket_region(bucket_name: str, endpoint_url: Optional[str] = None, only
     :param only_strategies: For testing, use only strategies with 1-based numbers in this set.
     """
 
-    s3_client = cast(S3Client, session.client('s3', endpoint_url=endpoint_url))
+    s3_client = session.client('s3', endpoint_url=endpoint_url)
 
     def attempt_get_bucket_location() -> Optional[str]:
         """
@@ -283,7 +280,7 @@ def get_bucket_region(bucket_name: str, endpoint_url: Optional[str] = None, only
         # It could also be because AWS open data buckets (which we tend to
         # encounter this problem for) tend to actually themselves be in
         # us-east-1.
-        backup_s3_client = cast(S3Client, session.client('s3', region_name='us-east-1'))
+        backup_s3_client = session.client('s3', region_name='us-east-1')
         return backup_s3_client.get_bucket_location(Bucket=bucket_name).get('LocationConstraint', None)
 
     def attempt_head_bucket() -> Optional[str]:
@@ -368,11 +365,11 @@ def get_object_for_url(url: ParseResult, existing: Optional[bool] = None) -> "Ob
         try:
             # Get the bucket's region to avoid a redirect per request
             region = get_bucket_region(bucket_name, endpoint_url=endpoint_url)
-            s3 = cast(S3ServiceResource, session.resource('s3', region_name=region, endpoint_url=endpoint_url))
+            s3 = session.resource('s3', region_name=region, endpoint_url=endpoint_url)
         except ClientError:
             # Probably don't have permission.
             # TODO: check if it is that
-            s3 = cast(S3ServiceResource, session.resource('s3', endpoint_url=endpoint_url))
+            s3 = session.resource('s3', endpoint_url=endpoint_url)
 
         obj = s3.Object(bucket_name, key_name)
         objExists = True
@@ -394,7 +391,7 @@ def get_object_for_url(url: ParseResult, existing: Optional[bool] = None) -> "Ob
         return obj
 
 
-@retry(errors=[BotoServerError])
+@retry(errors=[AWSServerErrors])
 def list_objects_for_url(url: ParseResult) -> List[str]:
         """
         Extracts a key (object) from a given parsed s3:// URL. The URL will be
@@ -419,7 +416,7 @@ def list_objects_for_url(url: ParseResult) -> List[str]:
         if host:
             endpoint_url = f'{protocol}://{host}' + f':{port}' if port else ''
 
-        client = cast(S3Client, session.client('s3', endpoint_url=endpoint_url))
+        client = session.client('s3', endpoint_url=endpoint_url)
 
         listing = []
 
@@ -444,3 +441,42 @@ def flatten_tags(tags: Dict[str, str]) -> List[Dict[str, str]]:
     Convert tags from a key to value dict into a list of 'Key': xxx, 'Value': xxx dicts.
     """
     return [{'Key': k, 'Value': v} for k, v in tags.items()]
+
+
+def boto3_pager(requestor_callable: Callable[..., Any], result_attribute_name: str,
+                **kwargs: Any) -> Iterable[Any]:
+    """
+    Yield all the results from calling the given Boto 3 method with the
+    given keyword arguments, paging through the results using the Marker or
+    NextToken, and fetching out and looping over the list in the response
+    with the given attribute name.
+    """
+
+    # Recover the Boto3 client, and the name of the operation
+    client = requestor_callable.__self__  # type: ignore[attr-defined]
+    op_name = requestor_callable.__name__
+
+    # grab a Boto 3 built-in paginator. See
+    # <https://boto3.amazonaws.com/v1/documentation/api/latest/guide/paginators.html>
+    paginator = client.get_paginator(op_name)
+
+    for page in paginator.paginate(**kwargs):
+        # Invoke it and go through the pages, yielding from them
+        yield from page.get(result_attribute_name, [])
+
+
+def get_item_from_attributes(attributes: List[AttributeTypeDef], name: str) -> Any:
+    """
+    Given a list of attributes, find the attribute associated with the name and return its corresponding value.
+
+    The `attribute_list` will be a list of TypedDict's (which boto3 SDB functions commonly return),
+    where each TypedDict has a "Name" and "Value" key value pair.
+    This function grabs the value out of the associated TypedDict.
+
+    If the attribute with the name does not exist, the function will return None.
+
+    :param attributes: list of attributes as List[AttributeTypeDef]
+    :param name: name of the attribute
+    :return: value of the attribute
+    """
+    return next((attribute["Value"] for attribute in attributes if attribute["Name"] == name), None)

toil/lib/conversions.py

@@ -4,7 +4,7 @@ Also contains general conversion functions
 """
 
 import math
-from typing import SupportsInt, Tuple, Union
+from typing import SupportsInt, Tuple, Union, Optional
 
 # See https://en.wikipedia.org/wiki/Binary_prefix
 BINARY_PREFIXES = ['ki', 'mi', 'gi', 'ti', 'pi', 'ei', 'kib', 'mib', 'gib', 'tib', 'pib', 'eib']
@@ -128,3 +128,26 @@ def hms_duration_to_seconds(hms: str) -> float:
     seconds += float(vals_to_convert[2]) 
 
     return seconds
+
+
+def strtobool(val: str) -> bool:
+    """
+    Make a human-readable string into a bool.
+    
+    Convert a string along the lines of "y", "1", "ON", "TrUe", or
+    "Yes" to True, and the corresponding false-ish values to False.
+    """
+    # We only track prefixes, so "y" covers "y", "yes",
+    # and "yeah no" and makes them all True.
+    TABLE = {True: ["1", "on", "y", "t"], False: ["0", "off", "n", "f"]}
+    lowered = val.lower()
+    for result, prefixes in TABLE.items():
+        for prefix in prefixes:
+            if lowered.startswith(prefix):
+                return result
+    raise ValueError(f"Cannot convert \"{val}\" to a bool")
+
+
+def opt_strtobool(b: Optional[str]) -> Optional[bool]:
+    """Convert an optional string representation of bool to None or bool"""
+    return b if b is None else strtobool(b)