快捷導(dǎo)航

PostgreSQL使用COPY協(xié)議高效批量數(shù)據(jù)寫入的實戰(zhàn)指南

更新時間：2025年11月20日 09:45:40 作者：ezreal_pan

這篇文章主要介紹了PostgreSQL的COPY協(xié)議,這是一種高效批量數(shù)據(jù)導(dǎo)入導(dǎo)出的二進制協(xié)議,適用于需要高效寫入大量數(shù)據(jù)的場景,COPY協(xié)議通過流式處理、事務(wù)安全和無參數(shù)限制等優(yōu)勢,顯著提升了數(shù)據(jù)寫入性能,并結(jié)合事務(wù)管理保證了數(shù)據(jù)一致性,需要的朋友可以參考下

問題背景

在開發(fā)過程中，我們經(jīng)常會遇到需要批量寫入大量數(shù)據(jù)到 PostgreSQL 數(shù)據(jù)庫的場景。當使用傳統(tǒng)的參數(shù)化插入語句時，可能會遇到如下錯誤：

pq: got 86575 parameters but PostgreSQL only supports 65535 parameters

這是因為 PostgreSQL 對單個查詢的參數(shù)數(shù)量有限制（通常為 65535）。傳統(tǒng)的解決方案是進行數(shù)據(jù)分片，分批寫入數(shù)據(jù)庫。但這種方法存在以下問題：

需要手動管理分片邏輯
多次數(shù)據(jù)庫往返，網(wǎng)絡(luò)開銷大
事務(wù)管理復(fù)雜
性能不夠理想

COPY 協(xié)議解決方案

COPY 協(xié)議簡介

PostgreSQL 的 COPY 協(xié)議是專門為高效批量數(shù)據(jù)操作設(shè)計的二進制協(xié)議，具有以下優(yōu)勢：

高性能：避免了 SQL 解析開銷，直接使用二進制格式傳輸數(shù)據(jù)
低內(nèi)存占用：流式處理，不需要在內(nèi)存中構(gòu)建龐大的 SQL 語句
事務(wù)安全：可以在事務(wù)中執(zhí)行，保證數(shù)據(jù)一致性
無參數(shù)限制：不受 PostgreSQL 參數(shù)數(shù)量限制

二進制協(xié)議原理

COPY 協(xié)議使用 PostgreSQL 的前后端協(xié)議進行數(shù)據(jù)傳輸，其工作流程如下：

啟動 COPY 模式：客戶端發(fā)送 COPY FROM STDIN 命令
數(shù)據(jù)傳輸：使用二進制格式按行發(fā)送數(shù)據(jù)
結(jié)束傳輸：發(fā)送特定的結(jié)束標記
確認完成：服務(wù)器返回處理結(jié)果

二進制格式避免了文本解析的開銷，直接使用網(wǎng)絡(luò)字節(jié)序傳輸數(shù)據(jù)，大大提高了傳輸效率。

實戰(zhàn)實現(xiàn)

依賴庫

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

核心實現(xiàn)代碼

// BatchCreate 批量創(chuàng)建消息接收者記錄 - 使用 COPY 協(xié)議
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 的數(shù)據(jù)
    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()
    
    // 開始事務(wù)
    tx, err := sqlDB.BeginTx(ctx, nil)
    if err != nil {
        return 0, fmt.Errorf("開始事務(wù)失?。?+v", err)
    }
    defer func() {
        if err != nil {
            tx.Rollback()
        }
    }()

    // 創(chuàng)建 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()

    // 批量寫入數(shù)據(jù)
    for _, item := range validData {
        _, err = stmt.Exec(item.SendLogID, item.UserID, item.Status, item.CreatedAt, item.UpdatedAt)
        if err != nil {
            return 0, fmt.Errorf("寫入數(shù)據(jù)失?。?+v", err)
        }
    }

    // 執(zhí)行 COPY
    _, err := stmt.Exec()
    if err != nil {
        return 0, fmt.Errorf("執(zhí)行 COPY 失?。?+v", err)
    }

    // 提交事務(wù)
    if err = tx.Commit(); err != nil {
        return 0, fmt.Errorf("提交事務(wù)失?。?+v", err)
    }
    
    rowsAffected = int64(len(validData))
    return rowsAffected, nil
}

代碼說明

數(shù)據(jù)驗證：首先過濾掉 nil 數(shù)據(jù)，確保數(shù)據(jù)有效性
事務(wù)管理：使用事務(wù)確保數(shù)據(jù)一致性，出錯時自動回滾
COPY 準備：通過 pq.CopyIn 準備 COPY 語句，指定表名和列名
批量寫入：遍歷數(shù)據(jù)并執(zhí)行 Exec，但此時數(shù)據(jù)還在客戶端緩沖區(qū)
最終執(zhí)行：調(diào)用 stmt.Exec() 真正將數(shù)據(jù)發(fā)送到服務(wù)器
事務(wù)提交：提交事務(wù)，完成批量寫入

完整測試用例

// 設(shè)置測試數(shù)據(jù)庫
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()

    // 創(chuàng)建日志工廠
    logger := setupLogger()

    // 創(chuàng)建 repository 實例
    repo := NewReceiverRepository(db, logger)

    // 準備測試數(shù)據(jù) - 20000 條記錄，使用負的 send_log_id 避免污染數(shù)據(jù)
    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()

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

    // 驗證結(jié)果
    assert.NoError(t, err)
    assert.Equal(t, int64(20000), rowsAffected)

    // 驗證數(shù)據(jù)是否正確插入
    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)

    // 清理測試數(shù)據(jù)
    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)
}