Implement Refreshable Chained AWS Session for Multi-Account Role Assumption

CHANGELOG.md

@@ -1,5 +1,8 @@
 # Changelog
 
+## 1.4.0
+  * Adds proxy AWS Account support [#59](https://github.com/singer-io/tap-dynamodb/pull/59)
+
 ## 1.3.0
   * Updates to run on python 3.11.7 [#57](https://github.com/singer-io/tap-dynamodb/pull/57)
 

setup.py

@@ -3,7 +3,7 @@
 
 setup(
     name="tap-dynamodb",
-    version="1.3.0",
+    version="1.4.0",
     description="Singer.io tap for extracting data",
     author="Stitch",
     url="http://singer.io",

tap_dynamodb/__init__.py

@@ -6,7 +6,7 @@
 import singer
 from singer import metadata
 from tap_dynamodb.discover import discover_streams
-from tap_dynamodb.dynamodb import setup_aws_client
+from tap_dynamodb.dynamodb import setup_aws_client, setup_aws_client_with_proxy
 from tap_dynamodb.sync import sync_stream
 
 
@@ -97,7 +97,13 @@ def main():
 
     # TODO Is this the right way to do this? It seems bad
     if not config.get('use_local_dynamo'):
-        setup_aws_client(config)
+        # Check if 'proxy_account_id' and 'proxy_role_name' are present in config
+        if config.get('proxy_account_id') and config.get('proxy_role_name'):
+            # If both are present, call setup_aws_client_with_proxy
+            setup_aws_client_with_proxy(config)
+        else:
+            # Otherwise, call setup_aws_client
+            setup_aws_client(config)
 
     if args.discover:
         do_discover(args.config)

tap_dynamodb/dynamodb.py

@@ -7,7 +7,8 @@
     AssumeRoleCredentialFetcher,
     CredentialResolver,
     DeferredRefreshableCredentials,
-    JSONFileCache
+    JSONFileCache,
+    RefreshableCredentials
 )
 from botocore.session import Session
 from botocore.config import Config
@@ -77,6 +78,56 @@ def setup_aws_client(config):
     LOGGER.info("Attempting to assume_role on RoleArn: %s", role_arn)
     boto3.setup_default_session(botocore_session=refreshable_session)
 
+
+@retry_pattern()
+def setup_aws_client_with_proxy(config):
+    proxy_role_arn = "arn:aws:iam::{}:role/{}".format(config['proxy_account_id'].replace('-', ''),config['proxy_role_name'])
+    cust_role_arn = "arn:aws:iam::{}:role/{}".format(config['account_id'].replace('-', ''),config['role_name'])
+
+    # Step 1: Assume Role in Account Proxy and set up refreshable session
+    session_proxy = Session()
+    fetcher_proxy = AssumeRoleCredentialFetcher(
+        client_creator=session_proxy.create_client,
+        source_credentials=session_proxy.get_credentials(),
+        role_arn=proxy_role_arn,
+        extra_args={
+            'DurationSeconds': 3600,
+            'RoleSessionName': 'ProxySession'
+        },
+        cache=JSONFileCache()
+    )
+
+    # Refreshable credentials for Account Proxy
+    refreshable_credentials_proxy = RefreshableCredentials.create_from_metadata(
+        metadata=fetcher_proxy.fetch_credentials(),
+        refresh_using=fetcher_proxy.fetch_credentials,
+        method="sts-assume-role"
+    )
+
+    # Step 2: Use Proxy Account's session to assume Role in Customer Account
+    session_cust = Session()
+    fetcher_cust = AssumeRoleCredentialFetcher(
+        client_creator=session_cust.create_client,
+        source_credentials=refreshable_credentials_proxy,
+        role_arn=cust_role_arn,
+        extra_args={
+            'DurationSeconds': 3600,
+            'RoleSessionName': 'TapDynamodDBCustSession',
+            'ExternalId': config['external_id']
+        },
+        cache=JSONFileCache()
+    )
+
+    # Set up refreshable session for Customer Account
+    refreshable_session_cust = Session()
+    refreshable_session_cust.register_component(
+        'credential_provider',
+        CredentialResolver([AssumeRoleProvider(fetcher_cust)])
+    )
+
+    LOGGER.info("Attempting to assume_role on RoleArn: %s", cust_role_arn)
+    boto3.setup_default_session(botocore_session=refreshable_session_cust)
+
 def get_request_timeout(config):
     # if request_timeout is other than 0,"0" or "" then use request_timeout
     request_timeout = config.get('request_timeout')