Merge pull request #5 from Bit-Quill/victor/pgxAuthToken

use aws sdk auth token gen
Bit-Quill · Jan 7, 2025 · c16f518 · c16f518
2 parents 7c006a0 + 25a4b49
commit c16f518
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 34 deletions.
diff --git 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
 }
 
@@ -104,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 (
@@ -143,8 +116,6 @@ func example(clusterEndpoint string, region string) error {
 		return err
 	}
 
-	defer conn.Close(ctx)
-
 	return nil
 }
 

diff --git 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
@@ -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=