Golang新手问题：使用database/sql库时连接失败如何从Vault获取密码

在Go中处理数据库密码轮换时，可以通过实现连接重试逻辑和动态更新连接字符串来解决。以下是几种可行的方案：

方案1：使用自定义连接器（Driver.Connector）

package main

import (
    "database/sql"
    "database/sql/driver"
    "fmt"
    "net"
    "sync"
    "time"
    
    _ "github.com/go-sql-driver/mysql"
)

var db *sql.DB
const port = "port number"

type VaultConnector struct {
    dsnTemplate string // 不包含密码的连接字符串模板
    mu          sync.RWMutex
    currentDSN  string
}

func NewVaultConnector(dsnTemplate string) *VaultConnector {
    vc := &VaultConnector{
        dsnTemplate: dsnTemplate,
    }
    vc.refreshDSN()
    return vc
}

func (vc *VaultConnector) refreshDSN() {
    password := fetchPasswordFromVault()
    vc.mu.Lock()
    vc.currentDSN = fmt.Sprintf(vc.dsnTemplate, password)
    vc.mu.Unlock()
}

func (vc *VaultConnector) Connect(ctx context.Context) (driver.Conn, error) {
    vc.mu.RLock()
    dsn := vc.currentDSN
    vc.mu.RUnlock()
    
    // 使用标准mysql驱动创建连接
    driver := mysql.MySQLDriver{}
    conn, err := driver.Open(dsn)
    if err != nil {
        // 如果连接失败，尝试刷新密码重连
        vc.refreshDSN()
        vc.mu.RLock()
        dsn = vc.currentDSN
        vc.mu.RUnlock()
        return driver.Open(dsn)
    }
    return conn, nil
}

func (vc *VaultConnector) Driver() driver.Driver {
    return mysql.MySQLDriver{}
}

func fetchPasswordFromVault() string {
    // 实现从Vault获取密码的逻辑
    return "new_password_from_vault"
}

func main() {
    dsnTemplate := "user:%s@tcp(localhost:3306)/dbname"
    connector := NewVaultConnector(dsnTemplate)
    
    var err error
    db, err = sql.OpenDB(connector)
    if err != nil {
        panic(err)
    }
    
    db.SetMaxOpenConns(100)
    
    listener, err := net.Listen("tcp", ":"+port)
    if err != nil {
        panic(err)
    }
    
    for {
        conn, err := listener.Accept()
        if err != nil {
            continue
        }
        go handleConnection(conn)
    }
}

方案2：包装sql.DB并实现重试逻辑

package main

import (
    "database/sql"
    "net"
    "sync"
    "time"
    
    _ "github.com/go-sql-driver/mysql"
)

type VaultDB struct {
    db           *sql.DB
    dsnTemplate  string
    mu           sync.RWMutex
    currentDSN   string
}

func NewVaultDB(dsnTemplate string) *VaultDB {
    vdb := &VaultDB{
        dsnTemplate: dsnTemplate,
    }
    vdb.refreshConnection()
    return vdb
}

func (vdb *VaultDB) refreshConnection() {
    password := fetchPasswordFromVault()
    vdb.mu.Lock()
    defer vdb.mu.Unlock()
    
    vdb.currentDSN = fmt.Sprintf(vdb.dsnTemplate, password)
    
    if vdb.db != nil {
        vdb.db.Close()
    }
    
    var err error
    vdb.db, err = sql.Open("mysql", vdb.currentDSN)
    if err != nil {
        panic(err)
    }
    vdb.db.SetMaxOpenConns(100)
}

func (vdb *VaultDB) Query(query string, args ...interface{}) (*sql.Rows, error) {
    vdb.mu.RLock()
    db := vdb.db
    vdb.mu.RUnlock()
    
    rows, err := db.Query(query, args...)
    if err != nil && isAuthenticationError(err) {
        vdb.refreshConnection()
        vdb.mu.RLock()
        db = vdb.db
        vdb.mu.RUnlock()
        return db.Query(query, args...)
    }
    return rows, err
}

func (vdb *VaultDB) Exec(query string, args ...interface{}) (sql.Result, error) {
    vdb.mu.RLock()
    db := vdb.db
    vdb.mu.RUnlock()
    
    result, err := db.Exec(query, args...)
    if err != nil && isAuthenticationError(err) {
        vdb.refreshConnection()
        vdb.mu.RLock()
        db = vdb.db
        vdb.mu.RUnlock()
        return db.Exec(query, args...)
    }
    return result, err
}

func isAuthenticationError(err error) bool {
    // 检查错误是否为认证错误
    return err != nil && (err.Error() == "Access denied" || 
                         err.Error() == "Invalid authentication")
}

var vdb *VaultDB

func main() {
    dsnTemplate := "user:%s@tcp(localhost:3306)/dbname"
    vdb = NewVaultDB(dsnTemplate)
    
    listener, err := net.Listen("tcp", ":"+port)
    if err != nil {
        panic(err)
    }
    
    for {
        conn, err := listener.Accept()
        if err != nil {
            continue
        }
        go handleConnection(conn)
    }
}

func handleConnection(conn net.Conn) {
    defer conn.Close()
    
    // 使用包装的数据库接口
    rows, err := vdb.Query("SELECT * FROM some_table")
    if err != nil {
        // 处理错误
        return
    }
    defer rows.Close()
    
    // 处理查询结果
}

方案3：定期刷新密码

package main

import (
    "database/sql"
    "net"
    "sync"
    "time"
    
    _ "github.com/go-sql-driver/mysql"
)

type PasswordManager struct {
    currentPassword string
    mu              sync.RWMutex
}

func (pm *PasswordManager) GetPassword() string {
    pm.mu.RLock()
    defer pm.mu.RUnlock()
    return pm.currentPassword
}

func (pm *PasswordManager) RefreshPassword() {
    newPassword := fetchPasswordFromVault()
    pm.mu.Lock()
    pm.currentPassword = newPassword
    pm.mu.Unlock()
}

func (pm *PasswordManager) StartRefreshLoop(interval time.Duration) {
    ticker := time.NewTicker(interval)
    go func() {
        for range ticker.C {
            pm.RefreshPassword()
        }
    }()
}

var passwordManager *PasswordManager
var db *sql.DB

func main() {
    passwordManager = &PasswordManager{}
    passwordManager.RefreshPassword()
    
    // 每12小时刷新一次密码
    passwordManager.StartRefreshLoop(12 * time.Hour)
    
    // 创建初始数据库连接
    refreshDBConnection()
    
    listener, err := net.Listen("tcp", ":"+port)
    if err != nil {
        panic(err)
    }
    
    for {
        conn, err := listener.Accept()
        if err != nil {
            continue
        }
        go handleConnection(conn)
    }
}

func refreshDBConnection() {
    password := passwordManager.GetPassword()
    dsn := fmt.Sprintf("user:%s@tcp(localhost:3306)/dbname", password)
    
    newDB, err := sql.Open("mysql", dsn)
    if err != nil {
        panic(err)
    }
    newDB.SetMaxOpenConns(100)
    
    if db != nil {
        db.Close()
    }
    db = newDB
}

这些方案提供了不同的处理方式：方案1通过自定义连接器在连接级别处理认证失败，方案2通过包装sql.DB在操作级别重试，方案3通过定期刷新密码预防连接失败。选择哪种方案取决于具体的应用场景和需求。