Golang中使用CopyFrom向PostgreSQL插入数据与性能对比

一、背景与意义

PostgreSQL 中的 CopyFrom 是一种性能较快的插入数据的方式，文本给出 Go 语言下的 CopyFrom 使用示例，并与常规的批量插入的方式做性能对比。但在实测中，未发现两种方法的性能存在显著差异。故在实际应用中，如果想用 CopyFrom 提升性能，要以实测数据为准。

二、代码示例

创建一个main.go文件，其内容如下：

package main

import (
    "context"
    "fmt"
    "log"
    "math/rand"
    "time"

    "github.com/jackc/pgx/v5"
    "gorm.io/driver/postgres"
    "gorm.io/gorm"
)

const dsn = "host=127.0.0.1 user=postgres password=123456 dbname=postgres port=5432 sslmode=disable TimeZone=Asia/Shanghai"

// TestTbl1 定义表结构
type TestTbl1 struct {
    Col1 string `gorm:"column:col1"`
    Col2 string `gorm:"column:col2"`
    Col3 string `gorm:"column:col3"`
}

// TableName 指定表名
func (TestTbl1) TableName() string {
    return "test_tbl1"
}

// randomString 生成指定长度的随机字符串
func randomString(length int) string {
    const charset = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
    b := make([]byte, length)
    for i := range b {
       b[i] = charset[rand.Intn(len(charset))]
    }
    return string(b)
}

func createTestTblData(size int) []TestTbl1 {
    var list []TestTbl1
    for i := 0; i < size; i++ {
       list = append(list, TestTbl1{
          Col1: randomString(32),
          Col2: randomString(32),
          Col3: randomString(32),
       })
    }
    return list
}

// 使用常规的方式插入数据
func testBatchInsert() error {
    db, err := gorm.Open(postgres.Open(dsn), &gorm.Config{})
    if err != nil {
       log.Fatalf("连接数据库失败: %v", err)
       return err
    }

    // 插入 100 万条数据
    totalRecords := 1000000
    batchSize := 1000
    fmt.Printf("开始插入 %d 条数据...\n", totalRecords)

    startTime := time.Now()

    for i := 0; i < totalRecords; i += batchSize {
       batch := createTestTblData(batchSize)

       // 批量插入
       if err := db.Create(&batch).Error; err != nil {
          log.Fatalf("插入数据失败: %v", err)
       }

       // 打印进度
       if (i+batchSize)%100000 == 0 || (i+batchSize) >= totalRecords {
          fmt.Printf("已插入 %d 条记录...\n", i+batchSize)
       }
    }

    elapsed := time.Since(startTime)
    fmt.Printf("插入完成！总耗时: %s\n", elapsed)
    return nil
}

// 使用CopyFrom的方式插入数据
func testCopyFrom() error {
    startTime := time.Now()

    fmt.Println("数据准备完成，开始执行 CopyFrom...")

    pgxConn, err := pgx.Connect(context.Background(), dsn)
    if err != nil {
       panic(err)
    }
    defer pgxConn.Close(context.Background())

    totalRecords := 1000000
    batchSize := 1000

    // 执行 CopyFrom
    for i := 0; i < totalRecords; i += batchSize {
       rows := make([][]interface{}, batchSize)
       list := createTestTblData(batchSize)
       for i := 0; i < batchSize; i++ {
          rows[i] = []interface{}{
             list[i].Col1,
             list[i].Col2,
             list[i].Col3,
          }
       }
       _, err = pgxConn.CopyFrom(
          context.Background(),
          pgx.Identifier{"test_tbl1"},
          []string{"col1", "col2", "col3"},
          pgx.CopyFromRows(rows),
       )
       if err != nil {
          return err
       }

       // 打印进度
       if (i+batchSize)%100000 == 0 || (i+batchSize) >= totalRecords {
          fmt.Printf("已插入 %d 条记录...\n", i+batchSize)
       }
    }

    elapsed := time.Since(startTime)
    fmt.Printf("GORM + CopyFrom 完成！插入 %d 条记录，总耗时: %s\n", totalRecords, elapsed)
    return nil
}

func main() {
    var err error
    err = testBatchInsert()
    if err != nil {
       log.Fatalf("BatchInsert 测试失败: %v", err)
    }

    err = testCopyFrom()
    if err != nil {
       log.Fatalf("CoppyFrom 测试失败: %v", err)
    }
}

三、代码说明

前面的代码中，有 testBatchInsert 和 testCopyFrom 两个方法，分别以常规方式和 CopyFrom 方式往数据库表 test_tbl1 插入100万条数据。其中 test_tb1 包含 col1、col2、col3 三个字符串字段，且有一个包含这三个字段的索引。

四、测试结果

先清空 test_tbl1 表的数据，执行 testBatchInsert 方法，记录方法耗时，然后再清空 test_tbl1 的数据，执行 testCopyFrom 方法，记录方法耗时。
两种方法插入100万条数据的耗时均为9~10秒左右。故在当前场景下，两个方法并没有显著差异，但使用常规方式插入数据的testBatchInsert方法代码结构更清晰和简洁。