From a15e0267f39c26b59bceeb5a5fa56e27553f5226 Mon Sep 17 00:00:00 2001
From: Victor Tsang <victor.tsang@macbookpro.lan>
Date: Mon, 6 Jan 2025 14:01:48 -0800
Subject: [PATCH 1/2] use aws sdk token gen

---
 go/pgx/example.go | 35 +++--------------------------------
 go/pgx/go.mod     |  1 +
 go/pgx/go.sum     |  2 ++
 3 files changed, 6 insertions(+), 32 deletions(-)

diff --git a/go/pgx/example.go b/go/pgx/example.go
index 613cae7..c55c1c7 100644
--- a/go/pgx/example.go
+++ b/go/pgx/example.go
@@ -3,14 +3,12 @@ package main
 import (
 	"context"
 	"fmt"
-	"net/http"
 	"os"
-	"strconv"
 	"strings"
-	"time"
 
 	_ "github.com/aws/aws-sdk-go-v2/aws"
-	v4 "github.com/aws/aws-sdk-go-v2/aws/signer/v4"
+	"github.com/aws/aws-sdk-go-v2/feature/dsql/auth"
+
 	"github.com/aws/aws-sdk-go-v2/config"
 	"github.com/google/uuid"
 	"github.com/jackc/pgx/v5"
@@ -32,38 +30,11 @@ func GenerateDbConnectAdminAuthToken(clusterEndpoint string, region string, acti
 		return "", err
 	}
 
-	creds, err := cfg.Credentials.Retrieve(ctx)
-	if err != nil {
-		return "", err
-	}
-
-	// The scheme is arbitrary and is only needed because validation of the URL requires one.
-	endpoint := "https://" + clusterEndpoint
-	req, err := http.NewRequest("GET", endpoint, nil)
+	token, err := auth.GenerateDBConnectAdminAuthToken(ctx, clusterEndpoint, region, cfg.Credentials)
 	if err != nil {
 		return "", err
 	}
-	values := req.URL.Query()
-	values.Set("Action", action)
-
-	// Set an expiry time for 15 minutes
-	values.Set("X-Amz-Expires", strconv.Itoa(15*60))
-	req.URL.RawQuery = values.Encode()
-
-	signer := v4.NewSigner()
-
-	// The payloadHash is the hex encoded SHA-256 hash of the request payload, and
-	// must be provided. Even if the request has no payload (aka body). If the
-	// request has no payload you should use the hex encoded SHA-256 of an empty
-	// string as the payloadHash value.
-	// e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855
-
-	uri, _, err := signer.PresignHTTP(ctx, creds, req, "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855", "dsql", region, time.Now())
-	if err != nil {
-		panic(err)
-	}
 
-	token := uri[len("https://"):]
 	return token, nil
 }
 
diff --git a/go/pgx/go.mod b/go/pgx/go.mod
index 912ec7c..8ea8cfd 100644
--- a/go/pgx/go.mod
+++ b/go/pgx/go.mod
@@ -11,6 +11,7 @@ require (
 
 require (
 	github.com/aws/aws-sdk-go-v2/credentials v1.17.48 // indirect
+	github.com/aws/aws-sdk-go-v2/feature/dsql/auth v1.0.1
 	github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.16.22 // indirect
 	github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.26 // indirect
 	github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.6.26 // indirect
diff --git a/go/pgx/go.sum b/go/pgx/go.sum
index 993f03e..f587f0d 100644
--- a/go/pgx/go.sum
+++ b/go/pgx/go.sum
@@ -4,6 +4,8 @@ github.com/aws/aws-sdk-go-v2/config v1.28.7 h1:GduUnoTXlhkgnxTD93g1nv4tVPILbdNQO
 github.com/aws/aws-sdk-go-v2/config v1.28.7/go.mod h1:vZGX6GVkIE8uECSUHB6MWAUsd4ZcG2Yq/dMa4refR3M=
 github.com/aws/aws-sdk-go-v2/credentials v1.17.48 h1:IYdLD1qTJ0zanRavulofmqut4afs45mOWEI+MzZtTfQ=
 github.com/aws/aws-sdk-go-v2/credentials v1.17.48/go.mod h1:tOscxHN3CGmuX9idQ3+qbkzrjVIx32lqDSU1/0d/qXs=
+github.com/aws/aws-sdk-go-v2/feature/dsql/auth v1.0.1 h1:Ux1HM2GlmnpxeZrxAwLlYVgLjMgsP03jm/Uu1YdzYTI=
+github.com/aws/aws-sdk-go-v2/feature/dsql/auth v1.0.1/go.mod h1:CmXEvpYo6HSxwjL6AStNox9FVc2gTw2ZplzzxU2nauE=
 github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.16.22 h1:kqOrpojG71DxJm/KDPO+Z/y1phm1JlC8/iT+5XRmAn8=
 github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.16.22/go.mod h1:NtSFajXVVL8TA2QNngagVZmUtXciyrHOt7xgz4faS/M=
 github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.26 h1:I/5wmGMffY4happ8NOCuIUEWGUvvFp5NSeQcXl9RHcI=

From 25a4b492db3680f46748982d8befceeab47a9c02 Mon Sep 17 00:00:00 2001
From: Victor Tsang <victor.tsang@improving.com>
Date: Mon, 6 Jan 2025 15:05:25 -0800
Subject: [PATCH 2/2] Update example.go

move defer closer to original caller
---
 go/pgx/example.go | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/go/pgx/example.go b/go/pgx/example.go
index c55c1c7..a5fc222 100644
--- a/go/pgx/example.go
+++ b/go/pgx/example.go
@@ -75,6 +75,8 @@ func example(clusterEndpoint string, region string) error {
 		return err
 	}
 
+	defer conn.Close(ctx)
+
 	// Create owner table
 	_, err = conn.Exec(ctx, `
 		CREATE TABLE IF NOT EXISTS owner (
@@ -114,8 +116,6 @@ func example(clusterEndpoint string, region string) error {
 		return err
 	}
 
-	defer conn.Close(ctx)
-
 	return nil
 }