更新 GORM 数据库连接的方法使用新的 AWS 令牌

在使用 GORM 连接数据库时，若采用 AWS 令牌身份验证，令牌每 15 分钟过期，需要定期刷新。传统的刷新方法，如在 cron 作业中执行 gorm.open()，存在线程安全问题和旧连接保持打开的问题。本解决方案提供了一种优雅的方法，通过 pgx 驱动程序的 beforeconnect 函数，在每次连接之前更新令牌，从而避免了上述问题。该方法利用了 tokenStruct 结构，包含令牌和过期时间，并通过 refreshAccessToken() 函数在令牌过期前获取新令牌。这种方法确保了 GORM 连接始终使用有效的 AWS 令牌，解决了令牌刷新问题。

我有一个 gorm 连接，最初是通过传递每 15 分钟过期的 aws 身份验证令牌创建的。该服务将能够连接到数据库 15 分钟。我有一些用于连接到数据库、创建新令牌以及使用 cron 库“刷新”连接的函数。我一直不知道如何向 gorm 提供新的代币。我虽然只是在 cron 作业中执行 gorm.open() ，但被告知这会很糟糕，因为它不是线程安全的，并且会使旧连接保持打开状态。我被引导到这个，但我不知道该怎么办： https://golang.org/pkg/database/sql/driver/#connector

这里有一个例子： https://github.com/aws/aws-sdk-go/issues/3043#issuecomment-581931580

但我不确定要替换 gorm 示例中的驱动程序和连接器代码。无论如何。

这是我到目前为止的代码。首先这个连接函数包裹着 gorm.open：

func connect(conf *config.dbconfig) (*gorm.db, error) {
    host, port, err := net.splithostport(conf.host)
    if err != nil {
        return nil, err
    }
    sslmode := "require"
    if conf.disabletls {
        sslmode = "disable"
    }
    return gorm.open(conf.dialect,
        fmt.sprintf("host=%s port=%s user=%s password=%s dbname=%s sslmode=%s",
            host, port, conf.username, conf.password, conf.name, sslmode))
}

这在 main 中被调用，如下所示：

//if we want to use aws token, replace pw in config with token
    if conf.db.useawstoken {
        authtoken, err := field.createauthtoken(&conf.db)
        if err != nil {
            logger.fatal(err)
        }
        conf.db.password = authtoken
    }

       //connect to db
    db, err := field.connect(&conf.db)
    if err != nil {
        logger.fatal(err)
    }
    defer db.close()

      //initialize cron job that runs in separate thread to refresh connection
       if conf.db.useawstoken {
        processorcron := cron.new()

        croninterval := "@every 10m"
        if conf.db.croninterval > 0 {
            croninterval = fmt.sprintf("@every %s", conf.db.croninterval)
        }

        processorcron.addfunc(croninterval, func() {
            repo.renewdb(&conf.db)
        })
        processorcron.start()
    }

我被 renewdb 函数的实际代码所困扰的部分，在 todo 的位置​​：

type Repo struct {
    db *gorm.DB
}

// FUNCTION 
func (r *Repo) RenewDB(dbConf *config.DbConfig) {
    var err error
    if dbConf.Password, err = CreateAuthToken(dbConf); err != nil {
        panic(fmt.Sprintf("failed at createAuthToken: %s\n", err.Error()))
    }

    //TODO: HOW DO I UPDATE GORM TO USE THE NEW AUTH TOKEN THAT I INSERTED
       // INTO dbConf?
}

//AUTH TOKEN GENERATOR
func CreateAuthToken(dbConf *config.DbConfig) (string, error) {
    awsSession, err := session.NewSession()
    if err != nil {
        return "", err
    }

    authToken, errToken := rdsutils.BuildAuthToken(dbConf.Host, dbConf.AwsRegion, dbConf.Username, awsSession.Config.Credentials)
    if errToken != nil {
        return "", errToken
    }

    return authToken, nil
}

解决方案

只要您可以使用 postgres 和 pgx 作为驱动程序，就可以通过 beforeconnect 函数非常优雅地实现：

package main

import (
    "context"
    "fmt"
    "github.com/jackc/pgx/v5"
    "github.com/jackc/pgx/v5/stdlib"
    "gorm.io/driver/postgres"
    "gorm.io/gorm"
    "log/slog"
    "os"
    "time"
)

type tokenStruct struct {
    Token string
    ExpiresOn time.Time
}

func main() {

    dsn := fmt.Sprintf("host=%s user=%s dbname=%s port=%s", host, username, databaseName, port)

    config, err := pgx.ParseConfig(dsn)
    if err != nil {
        slog.Error("error while parsing db config", "error", err)
        os.Exit(1)
    }

    db := stdlib.OpenDB(*config, stdlib.OptionBeforeConnect(refreshAccessToken()))

    gormDB, err := gorm.Open(postgres.New(postgres.Config{Conn: db}), &gorm.Config{})

    ...
}

func refreshAccessToken() func(ctx context.Context, config *pgx.ConnConfig) error {

    accessToken := tokenStruct{}

    return func(ctx context.Context, config *pgx.ConnConfig) error {
        if isExpired(accessToken) {
            var err error

            accessToken, err = getNewToken(ctx)
            if err != nil {
                return fmt.Errorf("failed to get new token: %w", err)
            }
            slog.Info("acquired new token for postgres", "expiry", accessToken.ExpiresOn)
        }

        config.Password = accessToken.Token
        return nil
    }
}

func getNewToken(ctx context.Context) (tokenStruct, error) {
    // get a fresh token
}

func isExpired(accessToken tokenStruct) bool {
    // expire the token a bit earlier, just to be safe
    oneMinuteInFuture := time.Now().Add(time.Minute * 1)
    return accessToken.ExpiresOn.Before(oneMinuteInFuture)
}

今天关于《更新 GORM 数据库连接的方法使用新的 AWS 令牌》的内容介绍就到此结束，如果有什么疑问或者建议，可以在golang学习网公众号下多多回复交流；文中若有不正之处，也希望回复留言以告知！