使用Go执行PostgreSQL的SQL文件无法获取全部结果

当你用pq库执行包含多个SELECT或PL/pgSQL块的SQL脚本时，默认只能拿到第一个结果集——这是因为db.Query()返回的*sql.Rows只会指向第一个结果集，后续的结果集需要手动切换才能读取。

修复后的RunQuery函数 #

func RunQuery(query string) (Results, error) {
    var myresults Results
    connectionString := "postgres://user:password@localhost:5432/dbname?sslmode=disable"
    base, err := pq.NewConnector(connectionString)
    if err != nil {
        log.Fatal(err)
    }

    // 捕获RAISE NOTICE信息
    connector := pq.ConnectorWithNoticeHandler(
        base,
        func(notice *pq.Error) {
            myresults.notices = append(myresults.notices, notice.Message)
        })

    db := sql.OpenDB(connector)
    defer db.Close()

    response, err := db.Query(query)
    if err != nil {
        log.Fatal(err)
    }
    defer response.Close() // 必须关闭结果集，避免资源泄漏

    // 遍历所有结果集
    for {
        // 获取当前结果集的列名
        cols, err := response.Columns()
        if err != nil {
            // DO块这类无返回结果的语句会触发ErrNoRows，直接跳转到下一个结果集
            if err == sql.ErrNoRows {
                if !response.NextResultSet() {
                    break
                }
                continue
            }
            fmt.Println("获取列名失败")
            return myresults, err
        }

        // 记录当前结果集的列名
        myresults.allRows = append(myresults.allRows, cols)

        // 准备行数据扫描的容器
        raw := make([][]byte, len(cols))
        dest := make([]interface{}, len(cols))
        for i := range dest {
            dest[i] = &raw[i]
        }

        // 扫描当前结果集的所有行
        for response.Next() {
            if err := response.Scan(dest...); err != nil {
                fmt.Println("扫描行失败")
                break
            }

            row := make([]string, len(cols))
            for i, v := range raw {
                row[i] = string(v)
                if v == nil {
                    row[i] = ""
                }
            }
            myresults.allRows = append(myresults.allRows, row)
        }

        // 检查当前结果集遍历是否有错误
        if err := response.Err(); err != nil {
            fmt.Println("结果集遍历出错")
            return myresults, err
        }

        // 切换到下一个结果集，没有则退出循环
        if !response.NextResultSet() {
            break
        }
    }

    return myresults, nil
}

核心修改说明 #

• 遍历所有结果集：通过外层循环+response.NextResultSet()，逐个处理SQL脚本中的每个结果生成语句（每个SELECT、返回结果的函数等）
• 处理无返回的语句：DO块这类不返回结果的语句会触发sql.ErrNoRows，直接跳过并切换到下一个结果集
• 资源释放：添加defer response.Close()确保结果集资源被正确回收，避免连接泄漏
• 独立处理每个结果集：每个结果集的列名和行数据单独收集，保证所有SELECT的结果都被存入myresults.allRows

测试效果 #

• 执行test2.sql时，会依次收集第一个select * from mytable的列名+所有行，以及第二个select name from mytable的列名+所有行
• 执行test.sql时，会收集第一个SELECT的结果、捕获DO块输出的所有NOTICE信息，最后收集第二个SELECT的结果

内容的提问来源于stack exchange，提问作者Josh Barisano