PostgreSQL使用COPY协议高效批量数据写入的实战指南

pq: got 86575 parameters but PostgreSQL only supports 65535 parameters

import (
    "github.com/lib/pq"
    "gorm.io/gorm"
)

// BatchCreate 批量创建消息接收者记录 - 使用 COPY 协议
func (r *receiverRepo) BatchCreate(ctx context.Context, db *gorm.DB, data []*define.WecomMsgReceiver) (rowsAffected int64, err error) {
    db = r.WithTrace(ctx, db)
    db = db.Table(r.TableName())

    if len(data) == 0 {
        return 0, nil
    }

    // 过滤掉 nil 的数据
    validData := make([]*define.WecomMsgReceiver, 0, len(data))
    for _, item := range data {
        if item != nil {
            validData = append(validData, item)
        }
    }
    if len(validData) == 0 {
        return 0, nil
    }

    // 获取底层 sql.DB
    sqlDB := db.DB()
    
    // 开始事务
    tx, err := sqlDB.BeginTx(ctx, nil)
    if err != nil {
        return 0, fmt.Errorf("开始事务失败：%+v", err)
    }
    defer func() {
        if err != nil {
            tx.Rollback()
        }
    }()

    // 创建 COPY writer
    stmt, err := tx.Prepare(pq.CopyIn(r.TableName(), "send_log_id", "user_id", "status", "created_at", "updated_at"))
    if err != nil {
        return 0, fmt.Errorf("准备 COPY 语句失败：%+v", err)
    }
    defer stmt.Close()

    // 批量写入数据
    for _, item := range validData {
        _, err = stmt.Exec(item.SendLogID, item.UserID, item.Status, item.CreatedAt, item.UpdatedAt)
        if err != nil {
            return 0, fmt.Errorf("写入数据失败：%+v", err)
        }
    }

    // 执行 COPY
    _, err := stmt.Exec()
    if err != nil {
        return 0, fmt.Errorf("执行 COPY 失败：%+v", err)
    }

    // 提交事务
    if err = tx.Commit(); err != nil {
        return 0, fmt.Errorf("提交事务失败：%+v", err)
    }
    
    rowsAffected = int64(len(validData))
    return rowsAffected, nil
}

// 设置测试数据库
func setupTestDB() (*gorm.DB, error) {
    ctx := context.Background()
    postgres, err := infrastructure.DialPostgres(ctx, infrastructure.PostgresConfig{
        Host:     "host",
        Port:     5432,
        Username: "postgres",
        Password: "xxxxx",
        Database: "xxxxx",
    })
    if err != nil {
        return nil, err
    }

    return postgres, nil
}

func setupLogger() factory.LogFactory {
    logger, _ := factory.NewJsonFactory(factory.NewLevel("info"), factory.NewZapOption(factory.AddCallerSkip(0)))
    return logger
}

func TestReceiverRepo_BatchCreate(t *testing.T) {
    db, err := setupTestDB()
    require.NoError(t, err)
    defer db.Close()

    // 创建日志工厂
    logger := setupLogger()

    // 创建 repository 实例
    repo := NewReceiverRepository(db, logger)

    // 准备测试数据 - 20000 条记录，使用负的 send_log_id 避免污染数据
    testData := make([]*define.WecomMsgReceiver, 0, 20000)
    now := time.Now()
    negativeSendLogID := int64(-100000) // 使用负的 send_log_id

    for i := 0; i < 20000; i++ {
        testData = append(testData, &define.WecomMsgReceiver{
            SendLogID: negativeSendLogID,
            UserID:    "test_user_" + fmt.Sprint(i),
            Status:    1,
            CreatedAt: now,
            UpdatedAt: now,
        })
    }

    ctx := context.Background()

    // 执行批量插入
    rowsAffected, err := repo.BatchCreate(ctx, db, testData)

    // 验证结果
    assert.NoError(t, err)
    assert.Equal(t, int64(20000), rowsAffected)

    // 验证数据是否正确插入
    var count int64
    query := "SELECT COUNT(*) FROM wecom_msg_receiver WHERE send_log_id < 0 AND send_log_id >= ?"
    err = db.Raw(query, negativeSendLogID).Count(&count).Error
    assert.NoError(t, err)
    assert.Equal(t, int64(20000), count)

    // 清理测试数据
    deleteQuery := "DELETE FROM wecom_msg_receiver WHERE send_log_id < 0 AND send_log_id >= ?"
    result := db.Exec(deleteQuery, negativeSendLogID)
    assert.NoError(t, result.Error)
    assert.Equal(t, int64(20000), result.RowsAffected)

    // 验证清理是否成功
    err = db.Raw(query, negativeSendLogID).Count(&count).Error
    assert.NoError(t, err)
    assert.Equal(t, int64(0), count)
}