Go中如何无需重复执行SET ROLE即可运行多条SELECT INTO OUTFILE S3语句

Go的sql.DB本质是连接池，每次调用Query/Exec时会从池中随机获取空闲连接。SET ROLE是会话级命令，仅对执行该命令的当前连接生效。第二条导出语句可能使用了未执行过SET ROLE的新连接，因此触发权限错误。

以下三种方法可避免重复执行SET ROLE：

1. 使用单个连接执行所有操作 #

从连接池中获取固定连接，在该连接上完成所有操作，确保会话状态一致：

func main() {
    cfg := mysql.Config{
        User:                 "mydb_user",
        Passwd:               "mydb_pass",
        Net:                  "tcp",
        Addr:                 "database-1.cluster-falexpegpmn8.us-west-2.rds.amazonaws.com:3306",
        DBName:               "mydb",
        AllowNativePasswords: true,
    }

    db, err := sql.Open("mysql", cfg.FormatDSN())
    if err != nil {
        // 处理错误，如log.Fatal(err)
    }
    defer db.Close()

    // 获取单个连接
    conn, err := db.Conn(context.Background())
    if err != nil {
        // 处理错误
    }
    defer conn.Close()

    // 在当前连接执行SET ROLE
    _, err = conn.ExecContext(context.Background(), "SET ROLE AWS_SELECT_S3_ACCESS")
    if err != nil {
        // 处理错误
    }

    // 第一条导出语句
    _, err = conn.QueryContext(context.Background(), "SELECT * FROM items_20231026 INTO OUTFILE S3 "+
        "'s3://falex-aurora-backups/mysqldumps/2023/10/26/mydb-items_20231026-AAA' "+
        "FIELDS TERMINATED BY ',' LINES TERMINATED BY '\n' MANIFEST ON")
    if err != nil {
        // 处理错误
    }

    // 第二条导出语句，无需重复SET ROLE
    _, err = conn.QueryContext(context.Background(), "SELECT * FROM items_20231026 INTO OUTFILE S3 "+
        "'s3://falex-aurora-backups/mysqldumps/2023/10/26/mydb-items_20231026-BBB' "+
        "FIELDS TERMINATED BY ',' LINES TERMINATED BY '\n' MANIFEST ON")
    if err != nil {
        // 处理错误
    }
}

2. 在DSN中配置初始化命令 #

利用mysql驱动的InitCmd参数，让每个新连接建立时自动执行SET ROLE：

func main() {
    cfg := mysql.Config{
        User:                 "mydb_user",
        Passwd:               "mydb_pass",
        Net:                  "tcp",
        Addr:                 "database-1.cluster-falexpegpmn8.us-west-2.rds.amazonaws.com:3306",
        DBName:               "mydb",
        AllowNativePasswords: true,
        // 每个新连接建立后自动执行该命令
        InitCmd: []string{"SET ROLE AWS_SELECT_S3_ACCESS"},
    }

    db, err := sql.Open("mysql", cfg.FormatDSN())
    if err != nil {
        // 处理错误
    }
    defer db.Close()

    // 第一条导出语句
    _, err = db.Query("SELECT * FROM items_20231026 INTO OUTFILE S3 "+
        "'s3://falex-aurora-backups/mysqldumps/2023/10/26/mydb-items_20231026-AAA' "+
        "FIELDS TERMINATED BY ',' LINES TERMINATED BY '\n' MANIFEST ON")
    if err != nil {
        // 处理错误
    }

    // 第二条导出语句，无需手动SET ROLE
    _, err = db.Query("SELECT * FROM items_20231026 INTO OUTFILE S3 "+
        "'s3://falex-aurora-backups/mysqldumps/2023/10/26/mydb-items_20231026-BBB' "+
        "FIELDS TERMINATED BY ',' LINES TERMINATED BY '\n' MANIFEST ON")
    if err != nil {
        // 处理错误
    }
}

3. 配置用户默认角色（数据库端操作） #

直接在数据库中把AWS_SELECT_S3_ACCESS设为用户的默认角色，连接后自动生效：

ALTER USER 'mydb_user'@'%' DEFAULT ROLE AWS_SELECT_S3_ACCESS;

执行上述SQL后，代码中无需再执行SET ROLE，直接运行导出语句即可。

内容的提问来源于stack exchange，提问作者hermit.crab