Securely Connecting to Amazon Keyspaces from SageMaker Using Server-Side Credentials

Amazon Keyspaces (for Apache Cassandra) provides a scalable, highly available, and managed Apache Cassandra-compatible database service. When working with SageMaker notebooks, it's crucial to establish secure connections using proper authentication methods. Amazon Keyspaces supports two distinct authentication methods for programmatic access: service-specific credentials and AWS Signature Version 4 (SigV4) authentication plugin (see more at Seamless Amazon Keyspaces Integration with SageMaker Using SigV4 Authentication )

This post demonstrates how to connect to Amazon Keyspaces using service-specific credentials stored securely in AWS Secrets Manager.

Service-specific credentials are specialized authentication mechanisms that function similarly to traditional Cassandra username and password combinations. These credentials are long-term credentials permanently associated with a specific IAM user and can only be used for Amazon Keyspaces - they cannot be used across other AWS services. In contrast, the SigV4 authentication plugin enables authentication using standard IAM access keys (access key ID and secret access key), which can be used across multiple AWS services. The SigV4 method also supports temporary credentials through AWS Security Token Service (STS), allowing IAM roles and federated identities to access Amazon Keyspaces resources .

Why Use Service-Specific Credentials?

Service-specific credentials provide a secure way to authenticate with Amazon Keyspaces without exposing your primary AWS credentials. These credentials are:

• Scoped specifically to the Keyspaces service using more traditional username and password authentication method
• A simpler, more traditional username and password authentication method that is limited to Amazon Keyspaces
• Rotatable independently of your main AWS credentials
• More secure than embedding credentials directly in code

Prerequisites

Before starting, ensure you have:

• An AWS account with appropriate IAM permissions
• A SageMaker notebook instance
• SecretsManagerReadWrite policy attached to your execution role
• Generated service-specific credentials for Keyspaces

Step 1: Generate Service-Specific Credentials

First, generate service-specific credentials for Amazon Keyspaces through the AWS Console or CLI. Your credentials will look similar to this:

{
    "ServiceSpecificCredential": {
        "CreateDate": "2019-10-09T16:12:04Z",
        "ServiceName": "cassandra.amazonaws.com",
        "ServiceUserName": "keyspace-user-at-11122223333",
        "ServicePassword": "wAbCdEUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY",
        "ServiceSpecificCredentialId": "ABCDXYFI33SINPGJEBYESF",
        "UserName": "keyspace-user",
        "Status": "Active"
    }
}

Step 2: Store Credentials in AWS Secrets Manager

Store your ServiceUserName and ServicePassword in AWS Secrets Manager. Create a secret named Keyspaces_Server_Generated_credential with the following structure:

{
    "keyspaces_generated_id": "your-service-username",
    "keyspaces_generated_pw": "your-service-password"
}

Step 3: Implement the Connection Code

Here's the complete implementation for connecting to Amazon Keyspaces from your SageMaker notebook:

# Import required libraries
from cassandra.cluster import Cluster
from ssl import SSLContext, PROTOCOL_TLSv1_2, CERT_REQUIRED
from cassandra.auth import PlainTextAuthProvider
import boto3
import base64
from botocore.exceptions import ClientError
import json
from pandas import DataFrame

# Set up SSL context
ssl_context = SSLContext(PROTOCOL_TLSv1_2)

# Download and load the required certificate
!curl https://certs.secureserver.net/repository/sf-class2-root.crt -O 
ssl_context.load_verify_locations('sf-class2-root.crt')

# Retrieve credentials from Secrets Manager
secret_name = "Keyspaces_Server_Generated_credential"
region_name = "us-west-1"

# Create Secrets Manager client
session = boto3.session.Session()
client = session.client(
    service_name='secretsmanager',
    region_name=region_name
)

# Get secret value
get_secret_value_response = client.get_secret_value(SecretId=secret_name)

# Extract credentials
if 'SecretString' in get_secret_value_response:
    secret = json.loads(get_secret_value_response['SecretString'])
    ServiceUserName = secret['keyspaces_generated_id']
    ServicePassword = secret['keyspaces_generated_pw']
    print(f"Retrieved credentials for user: {ServiceUserName}")
else:
    decoded_binary_secret = base64.b64decode(get_secret_value_response['SecretBinary'])

# Establish connection to Keyspaces
auth_provider = PlainTextAuthProvider(username=ServiceUserName, password=ServicePassword)
cluster = Cluster(
    ['cassandra.us-west-1.amazonaws.com'], 
    ssl_context=ssl_context, 
    auth_provider=auth_provider, 
    port=9142
)
session = cluster.connect()

# Test the connection
cql = "SELECT * FROM system_schema.keyspaces"
result = session.execute(cql)
print("Available keyspaces:")
print(result.current_rows)

# Convert results to pandas DataFrame for analysis
df = DataFrame(result)
print("\nKeyspaces DataFrame:")
print(df)

Key Security Benefits

This approach provides several security advantages:

1. No hardcoded credentials: Sensitive information is stored securely in Secrets Manager
2. Encryption in transit: SSL/TLS encryption protects data during transmission
3. IAM integration: Leverages AWS IAM for access control
4. Credential rotation: Service-specific credentials can be rotated without affecting other services

Best Practices

• Always use SSL/TLS connections when connecting to Keyspaces
• Store credentials in AWS Secrets Manager rather than in notebook code
• Use service-specific credentials instead of root AWS credentials
• Regularly rotate your service-specific credentials
• Implement proper error handling for production applications
• Use appropriate IAM policies with least privilege access

Conclusion

By following this approach, you can securely connect to Amazon Keyspaces from SageMaker while maintaining security best practices. The combination of service-specific credentials and AWS Secrets Manager ensures that your database connections are both secure and manageable at scale.

This pattern can be extended to other AWS services and provides a solid foundation for building secure, scalable data science workflows in SageMaker.