pythonamazon-web-servicesboto3amazon-keyspaces

How to connect to Amazon Keyspaces using boto3?


I'm trying to connect to Amazon Keyspaces leveraging the Assume role provider which refreshes the credentials the moment they expire. I've set up my aws config file as the documentation states

[profile cassandra]
role_arn=role_to_assume
source_profile=default
role_session_name=testingCassandraConnection
duration_seconds=900

then, within the code I start a session with that profile

boto_session = boto3.Session(profile_name='cassandra', region_name='us-east-1')
auth_provider = SigV4AuthProvider(boto_session) 

cluster = Cluster(
    [CASSANDRA_CLUSTER],
    ssl_context=ssl_context,
    auth_provider=auth_provider,
    port=9142
)

session = cluster.connect()

but I get the error Error from server: code=0100 [Bad credentials] message="Authentication failure: SessionId mismatch

I also tried using sts_client.assume_role and passing the credentials directly to boto3.Session() and it works this way but I won't be able to refresh credentials when they expire.

has anyone encountered this problem?


Solution

  • I found out how to do it, turns out that the botocore library has a way to deal with refreshable credentials.

    def assumed_role_session(role_arn: str,
                             base_session: botocore.session.Session = None):
        # Default session
        base_session = base_session or boto3.session.Session()._session
    
        fetcher = botocore.credentials.AssumeRoleCredentialFetcher(
            client_creator=base_session.create_client,
            source_credentials=base_session.get_credentials(),
            role_arn=role_arn,
            extra_args={
                # set this if you want something non-default
                # 'RoleSessionName': None
                # 'DurationSeconds': 3600
            }
        )
    
        creds = botocore.credentials.DeferredRefreshableCredentials(
            method='assume-role',
            refresh_using=fetcher.fetch_credentials,
            time_fetcher=lambda: datetime.datetime.now(tzlocal())
        )
    
        botocore_session = botocore.session.Session()
        botocore_session._credentials = creds
    
        return boto3.Session(botocore_session=botocore_session)
    

    first of all we need to pass the info for a client to be created with fetcher, we then use that fetcher object to specify how the credentials are going to be refreshed in creds, finally, a boto3 session is returned with the botocore session (this session already has the rotating credentials set) as a base.

    This returned session is then passed to the auth provider as follows:

    auth_provider = SigV4AuthProvider(session)

    and it works.

    function taken from here

    full botocore.credentials documentation here