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

toil/lib/aws/session.py

@@ -15,17 +15,24 @@ import collections
 import logging
 import os
 import threading
-from typing import Dict, Optional, Tuple, cast
+from typing import TYPE_CHECKING, Literal, Optional, cast, overload
 
 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
 
+if TYPE_CHECKING:
+    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__)
 
 # A note on thread safety:
@@ -44,6 +51,7 @@ logger = logging.getLogger(__name__)
 # initializing Boto3 (or Boto2) things at a time.
 _init_lock = threading.RLock()
 
+
 def _new_boto3_session(region_name: Optional[str] = None) -> Session:
     """
     This is the One True Place where new Boto3 sessions should be made, and
@@ -59,10 +67,16 @@ def _new_boto3_session(region_name: Optional[str] = None) -> Session:
 
     with _init_lock:
         botocore_session = get_session()
-        botocore_session.get_component('credential_provider').get_provider(
-            'assume-role').cache = JSONFileCache()
+        botocore_session.get_component("credential_provider").get_provider(
+            "assume-role"
+        ).cache = JSONFileCache()
+
+        return Session(
+            botocore_session=botocore_session,
+            region_name=region_name,
+            profile_name=os.environ.get("TOIL_AWS_PROFILE", None),
+        )
 
-        return Session(botocore_session=botocore_session, region_name=region_name, profile_name=os.environ.get("TOIL_AWS_PROFILE", None))
 
 class AWSConnectionManager:
     """
@@ -98,29 +112,64 @@ class AWSConnectionManager:
         """
         # This stores Boto3 sessions in .item of a thread-local storage, by
         # region.
-        self.sessions_by_region: Dict[Optional[str], threading.local] = collections.defaultdict(threading.local)
+        self.sessions_by_region: dict[Optional[str], threading.local] = (
+            collections.defaultdict(threading.local)
+        )
         # This stores Boto3 resources in .item of a thread-local storage, by
         # (region, service name, endpoint URL) tuples
-        self.resource_cache: Dict[Tuple[Optional[str], str, Optional[str]], threading.local] = collections.defaultdict(threading.local)
+        self.resource_cache: dict[
+            tuple[Optional[str], str, Optional[str]], threading.local
+        ] = collections.defaultdict(threading.local)
         # This stores Boto3 clients in .item of a thread-local storage, by
         # (region, service name, endpoint URL) tuples
-        self.client_cache: Dict[Tuple[Optional[str], str, Optional[str]], threading.local] = collections.defaultdict(threading.local)
+        self.client_cache: dict[
+            tuple[Optional[str], str, Optional[str]], threading.local
+        ] = collections.defaultdict(threading.local)
         # This stores Boto 2 connections in .item of a thread-local storage, by
         # (region, service name) tuples.
-        self.boto2_cache: Dict[Tuple[Optional[str], str], threading.local] = collections.defaultdict(threading.local)
+        self.boto2_cache: dict[tuple[Optional[str], str], threading.local] = (
+            collections.defaultdict(threading.local)
+        )
 
     def session(self, region: Optional[str]) -> boto3.session.Session:
         """
         Get the Boto3 Session to use for the given region.
         """
         storage = self.sessions_by_region[region]
-        if not hasattr(storage, 'item'):
+        if not hasattr(storage, "item"):
             # This is the first time this thread wants to talk to this region
             # through this manager
             storage.item = _new_boto3_session(region_name=region)
         return cast(boto3.session.Session, storage.item)
 
-    def resource(self, region: Optional[str], service_name: str, endpoint_url: Optional[str] = None) -> boto3.resources.base.ServiceResource:
+    @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.
 
@@ -129,7 +178,7 @@ class AWSConnectionManager:
         """
         key = (region, service_name, endpoint_url)
         storage = self.resource_cache[key]
-        if not hasattr(storage, 'item'):
+        if not hasattr(storage, "item"):
             with _init_lock:
                 # We lock inside the if check; we don't care if the memoization
                 # sometimes results in multiple different copies leaking out.
@@ -139,14 +188,69 @@ class AWSConnectionManager:
                     # The Boto3 stubs are missing an overload for `resource` that takes
                     # a non-literal string. See
                     # <https://github.com/vemel/mypy_boto3_builder/issues/121#issuecomment-1011322636>
-                    storage.item = self.session(region).resource(service_name, endpoint_url=endpoint_url) # type: ignore
+                    storage.item = self.session(region).resource(service_name, endpoint_url=endpoint_url)  # type: ignore
                 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).resource(service_name) # type: ignore
+                    storage.item = self.session(region).resource(service_name)  # type: ignore
 
         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,42 +263,34 @@ 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]
-        if not hasattr(storage, 'item'):
+        if not hasattr(storage, "item"):
             with _init_lock:
                 # We lock because we call .client()
 
                 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
-        return cast(botocore.client.BaseClient , storage.item)
+                    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()
 
+
 def establish_boto3_session(region_name: Optional[str] = None) -> Session:
     """
     Get a Boto 3 session usable by the current thread.
@@ -205,7 +301,57 @@ 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.
 
@@ -213,9 +359,36 @@ 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:
+    return _global_manager.client(
+        region_name, service_name, endpoint_url=endpoint_url, config=config
+    )
+
+
+@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.
 
@@ -223,4 +396,6 @@ def resource(service_name: str, region_name: Optional[str] = None, endpoint_url:
     """
 
     # Just use a global version of the manager. Note that we change the argument order!
-    return _global_manager.resource(region_name, service_name, endpoint_url=endpoint_url)
+    return _global_manager.resource(
+        region_name, service_name, endpoint_url=endpoint_url
+    )