Update example and README.md

Author: imforster

go/pgx/README.md

@@ -1,20 +1,33 @@
-# Aurora DSQL pgx code examples
+# PGX with Aurora DSQL
 
 ## Overview
 
-The code examples in this topic show you how to use DSQL with Go pgx.
+This code example demonstrates how to use the `pgx` driver with Amazon Aurora DSQL. The example shows you how to connect to an Aurora DSQL cluster and perform database operations using IAM authentication.
 
-## Features
+Aurora DSQL is a distributed SQL database service that provides high availability and scalability for your PostgreSQL-compatible applications. `pgx` is a pure Go driver and toolkit for PostgreSQL that offers robust features including automatic connection pool management and authentication token refresh.
 
-- Uses `dsql.NewFromConfig` to create DSQL clients from AWS configuration
-- Supports automatic token refresh to maintain database connectivity
-- Implements connection pooling with pgx
-- Thread-safe token refresh mechanism
-- Demonstrates that connections before and after refresh are different
+
+## About the code example
+
+The example demonstrates a flexible connection approach using IAM authentication:
+
+- Implements automatic token refresh mechanism to maintain continuous database connectivity
+- Handles secure IAM-based authentication token generation
+- Provides connection pooling and management
+- Demonstrates best practices for Aurora DSQL connectivity in Go applications
+
+## ⚠️ Important
+
+- Running this code might result in charges to your AWS account.
+- We recommend that you grant your code least privilege. At most, grant only the
+  minimum permissions required to perform the task. For more information, see
+  Grant least privilege in the AWS IAM User Guide.
+- This code is not tested in every AWS Region. For more information, see
+  AWS Regional Services.
 
 ## Configuration
 
-The following environment variables can be used to configure the connection:
+The following environment variables can be used to configure the connection parameter in this example:
 
 - `CLUSTER_ENDPOINT`: Your Aurora cluster endpoint (required)
 - `REGION`: AWS region where your cluster is located (required)
@@ -31,21 +44,40 @@ The following environment variables can be used to configure the connection:
 
 ### Prerequisites
 
-* Go version >= 1.21
-* AWS credentials file is configured
+- You must have an AWS account, and have your default credentials and AWS Region
+  configured as described in the
+  [Globally configuring AWS SDKs and tools](https://docs.aws.amazon.com/credref/latest/refdocs/creds-config-files.html)
+  guide.
+-You must have an Aurora DSQL cluster. For information about creating an Aurora DSQL cluster, see the
+      [Getting started with Aurora DSQL](https://docs.aws.amazon.com/aurora-dsql/latest/userguide/getting-started.html)
+      guide.
+- If connecting as a non-admin user, ensure the user is linked to an IAM role and is granted access to the `myschema`
+  schema. See the
+  [Using database roles with IAM roles](https://docs.aws.amazon.com/aurora-dsql/latest/userguide/using-database-and-iam-roles.html)
+  guide.
+- Go version >= 1.21
 
 ### Setup test running environment
 
 Ensure you are authenticated with AWS credentials. No other setup is needed besides having Go installed.
 
+### Environment Variables
+Set the following required environment variables:
+
+```shell
+# Your cluster endpoint (e.g., "cluster-name.cluster-xxx.region.rds.amazonaws.com")
+export CLUSTER_ENDPOINT="<your cluster endpoint>"
+
+# Your AWS region (e.g., "us-east-1")
+export REGION="<your cluster region>"
+```
+
 ### Run the example tests
 
 In a terminal run the following commands:
 
 ```sh
 # Use the account credentials dedicated for golang
-export CLUSTER_ENDPOINT="<your cluster endpoint>"
-export REGION="<your cluster region>"
 go env -w GOPROXY=direct
 go test
 
@@ -56,31 +88,19 @@ go build
 
 ## Token Refresh
 
+Token Refresh
 The implementation includes an automatic token refresh mechanism that:
 
 1. Creates a new token before the current one expires (default: every 15 minutes)
-2. Updates the connection pool with the new token
-3. Ensures all new connections use the refreshed token
-4. Handles the refresh process in a thread-safe manner
-
-You can disable token refresh by setting the environment variable `DB_REFRESH_TOKEN=false`.
-
-### Connection Refresh Verification
 
-The code now includes functionality to verify that connections before and after token refresh are different:
+1. Ensures continuous database connectivity
 
-1. The `TestTokenRefresh` test captures connection IDs before and after refresh and verifies they are different
-2. The `TestMultipleConnectionsRefresh` test checks multiple connections to ensure all are replaced
-3. The `TestComprehensiveConnectionRefresh` test provides detailed statistics about the connection pool before and after refresh
-4. The example application includes a demonstration of connection refresh that shows the connection IDs changing
+1. Handles token generation and rotation securelye fresh connections with valid authentication tokens.
 
-When the token is refreshed, the following happens:
-- The old connection pool is reset
-- A new connection pool is created with the new token
-- All existing connections are closed and replaced with new connections
-- The new connections have different process IDs on the database server
+## Additional resources
 
-This ensures that all database operations after a token refresh use fresh connections with valid authentication tokens.
+- [Amazon Aurora DSQL Documentation](https://docs.aws.amazon.com/aurora-dsql/latest/userguide/what-is-aurora-dsql.html)
+- [AWS SDK for Go Documentation](https://docs.aws.amazon.com/sdk-for-go/)
 
 ---
 

go/pgx/example.go

@@ -63,14 +63,23 @@ func NewDSQLClient(ctx context.Context, region string) (*dsql.Client, error) {
 	return dsqlClient, nil
 }
 
-// GenerateDbConnectAdminAuthToken generates an authentication token for database connection
-func GenerateDbConnectAdminAuthToken(ctx context.Context, clusterEndpoint, region string) (string, error) {
+// GenerateDbConnectAuthToken generates an authentication token for database connection
+func GenerateDbConnectAuthToken(ctx context.Context, clusterEndpoint, region, user string) (string, error) {
 	cfg, err := config.LoadDefaultConfig(ctx)
 	if err != nil {
 		return "", err
 	}
 
-	token, err := auth.GenerateDBConnectAdminAuthToken(ctx, clusterEndpoint, region, cfg.Credentials)
+	if user == "admin" {
+		token, err := auth.GenerateDBConnectAdminAuthToken(ctx, clusterEndpoint, region, cfg.Credentials)
+		if err != nil {
+			return "", err
+		}
+
+		return token, nil
+	}
+
+	token, err := auth.GenerateDbConnectAuthToken(ctx, clusterEndpoint, region, cfg.Credentials)
 	if err != nil {
 		return "", err
 	}
@@ -122,7 +131,7 @@ func NewPool(ctx context.Context, clusterEndpoint string, region string) (*Pool,
 	}
 
 	// Generate initial token
-	token, err := GenerateDbConnectAdminAuthToken(poolCtx, clusterEndpoint, region)
+	token, err := GenerateDbConnectAuthToken(poolCtx, clusterEndpoint, region, dbConfig.User)
 	if err != nil {
 		cancel()
 		return nil, fmt.Errorf("failed to generate auth token: %v", err)
@@ -205,7 +214,7 @@ func (p *Pool) refreshToken() error {
 	defer p.mu.Unlock()
 
 	// Generate new token
-	token, err := GenerateDbConnectAdminAuthToken(p.ctx, p.clusterEndpoint, p.config.Region)
+	token, err := GenerateDbConnectAuthToken(p.ctx, p.clusterEndpoint, p.config.Region, p.config.User)
 	if err != nil {
 		return fmt.Errorf("failed to refresh auth token: %v", err)
 	}

go/pgx/example_test.go

@@ -40,7 +40,7 @@ func TestGenerateDbConnectAdminAuthToken(t *testing.T) {
 	ctx := context.Background()
 
 	// Test token generation
-	token, err := GenerateDbConnectAdminAuthToken(ctx, clusterEndpoint, region)
+	token, err := GenerateDbConnectAuthToken(ctx, clusterEndpoint, region, "admin")
 	if err != nil {
 		t.Errorf("Error generating auth token: %v\n", err)
 	}