如何使用pgx.CollectRows处理含重复行的嵌套结构体

首先先修正几个基础问题：

1. 结构体字段需要导出（首字母大写），否则pgx无法映射；Account的name字段类型应该是string而非uuid.UUID，和查询结果的accountname对应。
2. 原查询语句存在错误：表名group是PostgreSQL关键字，需用双引号包裹；关联条件逻辑错误（应该是groups.id = group_account_membership.group_id，而非关联到account_id）；要获取accountname需要关联accounts表。

修正后的结构体：

import "github.com/google/uuid"

type Group struct {
    ID       uuid.UUID
    Name     string
    Accounts []*Account
}

type Account struct {
    ID   uuid.UUID // 对应group_account_membership.id
    Name string    // 对应accounts.accountname
}

方案一：数据库端聚合（推荐） #

通过PostgreSQL的聚合函数直接将同一组的账户信息合并为数组，一次查询得到结构化结果，无需在Go代码中额外处理行合并。

调整后的查询语句 #

SELECT 
    g.id, 
    g.name,
    ARRAY_AGG(ROW(gam.id, a.accountname)) AS accounts
FROM "groups" g
JOIN group_account_membership gam ON g.id = gam.group_id
JOIN accounts a ON gam.account_id = a.id
WHERE g.id = $1
GROUP BY g.id, g.name;

自定义行映射函数 #

pgx默认的RowToAddrOfStructByName无法直接映射聚合的数组，需要自定义映射函数：

func rowToGroup(row pgx.Row) (*Group, error) {
    var group Group
    var accountRows []struct {
        ID   uuid.UUID
        Name string
    }

    err := row.Scan(&group.ID, &group.Name, &accountRows)
    if err != nil {
        return nil, err
    }

    // 转换为Account切片
    group.Accounts = make([]*Account, len(accountRows))
    for i, ar := range accountRows {
        group.Accounts[i] = &Account{
            ID:   ar.ID,
            Name: ar.Name,
        }
    }

    return &group, nil
}

执行查询 #

query := `
SELECT 
    g.id, 
    g.name,
    ARRAY_AGG(ROW(gam.id, a.accountname)) AS accounts
FROM "groups" g
JOIN group_account_membership gam ON g.id = gam.group_id
JOIN accounts a ON gam.account_id = a.id
WHERE g.id = $1
GROUP BY g.id, g.name;
`

rows, err := a.conn.Query(ctx, query, 1)
if err != nil {
    // 处理错误
}
defer rows.Close()

groups, err := pgx.CollectRows(rows, rowToGroup)
if err != nil {
    // 处理错误
}
// groups[0]就是合并后的组数据

方案二：Go代码端合并行 #

如果不想修改查询语句，可在获取行数据后手动合并同一组的记录：

调整后的基础查询语句 #

SELECT 
    g.id, 
    g.name,
    gam.id AS account_id,
    a.accountname AS account_name
FROM "groups" g
JOIN group_account_membership gam ON g.id = gam.group_id
JOIN accounts a ON gam.account_id = a.id
WHERE g.id = $1;

手动合并行 #

query := `
SELECT 
    g.id, 
    g.name,
    gam.id AS account_id,
    a.accountname AS account_name
FROM "groups" g
JOIN group_account_membership gam ON g.id = gam.group_id
JOIN accounts a ON gam.account_id = a.id
WHERE g.id = $1;
`

rows, err := a.conn.Query(ctx, query, 1)
if err != nil {
    // 处理错误
}
defer rows.Close()

groupMap := make(map[uuid.UUID]*Group)
for rows.Next() {
    var groupID uuid.UUID
    var groupName string
    var accountID uuid.UUID
    var accountName string

    err := rows.Scan(&groupID, &groupName, &accountID, &accountName)
    if err != nil {
        // 处理错误
    }

    // 检查组是否已存在
    group, exists := groupMap[groupID]
    if !exists {
        group = &Group{
            ID:   groupID,
            Name: groupName,
            Accounts: []*Account{},
        }
        groupMap[groupID] = group
    }

    // 添加账户到组
    group.Accounts = append(group.Accounts, &Account{
        ID:   accountID,
        Name: accountName,
    })
}

// 转换为切片（如果需要）
var groups []*Group
for _, g := range groupMap {
    groups = append(groups, g)
}

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