add ability to assume role using config param role_arn (#56)

* update docs for new config param `role_arn`

update readme for new config param `role_arn`. `role_arn` can be used to specify a role that should be assumed for appropriate access to s3.

* add support for assuming a role via `role_arn`

if `role_arn` is specified, first a session is created using this role, and then an s3 client is created with the session in the assumed role.

* added `role_arn` to known config params

* add clarifying comments & an INFO log message

* fix typo: endpoint_params

* fix format: RoleSessionName

* remove hard-coded `Meltano`

* bump version to 1.2.1

* add test for assume role

CHANGELOG.md

@@ -1,5 +1,8 @@
 # Change Log
+## [1.2.1](https://github.com/ome9ax/target-s3-jsonl/tree/1.2.1) (2022-07-13)
 
+### What's Changed
+* Added optional config parameter `role_arn`, which allows assuming additional roles.
 ## [1.2.0](https://github.com/ome9ax/target-s3-jsonl/tree/1.2.0) (2022-04-11)
 
 ### What's Changed

README.md

@@ -100,6 +100,7 @@ Full list of options in `config.json`:
 | aws_session_token                   | String  |            | AWS Session token. If not provided, `AWS_SESSION_TOKEN` environment variable will be used. |
 | encryption_type                     | String  |            | (Default: 'none') The type of encryption to use. Current supported options are: 'none' and 'KMS'. |
 | encryption_key                      | String  |            | A reference to the encryption key to use for data encryption. For KMS encryption, this should be the name of the KMS encryption key ID (e.g. '1234abcd-1234-1234-1234-1234abcd1234'). This field is ignored if 'encryption_type' is none or blank. |
+| role_arn                            | String  |            | The ARN of the role to assume |
 
 ## Test
 Install the tools

config.sample.json

@@ -4,5 +4,6 @@
     "s3_bucket": "BUCKET",
     "s3_key_prefix": "SOME-PREFIX/",
     "compression": "gzip",
-    "naming_convention": "{stream}-{timestamp}.jsonl"
+    "naming_convention": "{stream}-{timestamp}.jsonl",
+    "role_arn": "arn:aws:iam::000000000000:role/my_custom_role"
 }

setup.cfg

@@ -33,7 +33,7 @@ target_s3_jsonl = logging.conf
 [options.extras_require]
 test =
     pytest-cov
-    moto[s3]
+    moto[s3,sts]
 lint = flake8
 dist = wheel
 

target_s3_jsonl/__init__.py

@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 
-__version__ = '1.2.0'
+__version__ = '1.2.1'
 
 import argparse
 import gzip
@@ -238,6 +238,7 @@ def config_file(config_path):
         'aws_session_token',
         'aws_endpoint_url',
         'aws_profile',
+        'role_arn',
         's3_bucket',
         's3_key_prefix',
         'encryption_type',

target_s3_jsonl/s3.py

@@ -1,5 +1,7 @@
 #!/usr/bin/env python3
 import os
+import re
+import time
 import backoff
 import boto3
 from botocore.exceptions import ClientError
@@ -32,6 +34,9 @@ def create_client(config):
     aws_session_token = config.get('aws_session_token') or os.environ.get('AWS_SESSION_TOKEN')
     aws_profile = config.get('aws_profile') or os.environ.get('AWS_PROFILE')
     aws_endpoint_url = config.get('aws_endpoint_url')
+    role_arn = config.get('role_arn')
+
+    endpoint_params = {'endpoint_url': aws_endpoint_url} if aws_endpoint_url else {}
 
     # AWS credentials based authentication
     if aws_access_key_id and aws_secret_access_key:
@@ -43,10 +48,20 @@ def create_client(config):
     else:
         aws_session = boto3.session.Session(profile_name=aws_profile)
 
-    if aws_endpoint_url:
-        return aws_session.client('s3', endpoint_url=aws_endpoint_url)
-    else:
-        return aws_session.client('s3')
+    # config specifies a particular role to assume
+    # we create a session & s3-client with this role
+    if role_arn:
+        role_name = role_arn.split('/', 1)[1]
+        sts = aws_session.client('sts', **endpoint_params)
+        resp = sts.assume_role(RoleArn=role_arn, RoleSessionName=f'role-name={role_name}-profile={aws_profile}')
+        credentials = {
+            "aws_access_key_id": resp["Credentials"]["AccessKeyId"],
+            "aws_secret_access_key": resp["Credentials"]["SecretAccessKey"],
+            "aws_session_token": resp["Credentials"]["SessionToken"],
+        }
+        aws_session = boto3.Session(**credentials)
+        LOGGER.info(f"Creating s3 client with role {role_name}")
+    return aws_session.client('s3', **endpoint_params)
 
 
 # pylint: disable=too-many-arguments

tests/resources/config_assume_role.json

@@ -0,0 +1,13 @@
+{
+    "add_metadata_columns": false,
+    "aws_access_key_id": "ACCESS-KEY",
+    "aws_secret_access_key": "SECRET",
+    "s3_bucket": "BUCKET",
+    "temp_dir": "tests/output",
+    "memory_buffer": 2000000,
+    "compression": "none",
+    "timezone_offset": 0,
+    "naming_convention": "{stream}-{timestamp:%Y%m%dT%H%M%S}.jsonl",
+    "role_arn": "arn:aws:iam::123456789012:role/TestAssumeRole"
+
+}

tests/test_s3.py

@@ -3,12 +3,12 @@
 from copy import deepcopy
 from pathlib import Path
 import json
-
-# Third party imports
 import os
-from pytest import fixture, raises
+import re
+
 import boto3
-from moto import mock_s3
+from moto import mock_s3, mock_sts
+from pytest import fixture, raises
 
 # Package imports
 from target_s3_jsonl.s3 import create_client, upload_file, log_backoff_attempt
@@ -22,6 +22,11 @@ def config():
         return json.load(config_file)
 
 
+@fixture
+def config_assume_role():
+    with open(Path('tests', 'resources', 'config_assume_role.json'), 'r', encoding='utf-8') as f:
+        return json.load(f)
+
 @fixture(scope='module')
 def aws_credentials():
     """Mocked AWS Credentials for moto."""
@@ -36,9 +41,16 @@ def test_log_backoff_attempt(caplog):
     '''TEST : simple upload_files call'''
 
     log_backoff_attempt({'tries': 99})
+    pat = r'INFO     root:s3.py:\d{2} Error detected communicating with Amazon, triggering backoff: 99 try\n'
+    assert re.match(pat, caplog.text)
 
-    assert caplog.text == 'INFO     root:s3.py:22 Error detected communicating with Amazon, triggering backoff: 99 try\n'
 
+@mock_sts
+@mock_s3
+def test_create_client_with_assumed_role(config_assume_role, caplog):
+    """Assert client is created with assumed role when role_arn is specified""" 
+    client = create_client(config_assume_role)
+    assert caplog.text.endswith('Creating s3 client with role TestAssumeRole\n')
 
 @mock_s3
 def test_create_client(aws_credentials, config):