用 Go 语言并发处理 CSV 文件到数据库-51CTO.COM

问题背景

假设你拥有一个包含大量联系人信息的 CSV 文件，需要将这些信息迁移到数据库中。这些联系人信息可能包含姓名、电话号码、邮箱地址等。如果使用传统的单线程方式，逐条处理数据，迁移过程可能会非常缓慢，尤其是在数据量很大时。

在处理大量的 CSV 文件数据并迁移到数据库时，使用并发可以显著提升处理效率。Go 语言的 goroutine 和通道（channel）非常适合用来并发地处理数据。

下面我将给出一个示例，展示如何使用 Go 语言并发地处理 CSV 文件，并将数据插入到数据库中。

主要思路：

读取 CSV 文件：使用 encoding/csv 包来解析 CSV 文件。
并发处理数据：将 CSV 文件的数据分批次发送到多个 goroutine 中进行并发处理。
数据库插入：每个 goroutine 从通道中接收数据并将其插入到数据库中。
同步控制：使用 sync.WaitGroup 来等待所有 goroutine 完成任务。

假设我们的数据库是 MySQL，使用 github.com/jinzhu/gorm 作为 ORM 库来处理数据库插入。我们会定义一个 Contact 结构体来映射数据库中的表，并用并发的方式将每一行 CSV 数据插入到数据库。

示例代码

1. 安装必要的依赖

首先，你需要安装 gorm 和 csv 相关的包：

go get github.com/jinzhu/gorm
go get github.com/jinzhu/gorm/dialects/mysql
go get encoding/csv

2. 数据库模型定义

我们先定义一个 Contact 结构体，它会对应数据库中的联系人表。

package main

import (
	"github.com/jinzhu/gorm"
	_ "github.com/jinzhu/gorm/dialects/mysql"
	"fmt"
)

// Contact 是数据库中表的模型
type Contact struct {
	ID        uint   `gorm:"primary_key"`
	Name      string `gorm:"size:255"`
	Phone     string `gorm:"size:255"`
	Email     string `gorm:"size:255"`
}

func initDB() (*gorm.DB, error) {
	// 使用 MySQL 数据库
	db, err := gorm.Open("mysql", "user:password@/dbname?charset=utf8&parseTime=True&loc=Local")
	if err != nil {
		return nil, err
	}

	// 自动迁移表结构
	db.AutoMigrate(&Contact{})
	return db, nil
}

3. 读取 CSV 文件并处理

接下来，我们需要读取 CSV 文件并将每一行数据并发地插入到数据库中。

package main

import (
	"encoding/csv"
	"fmt"
	"os"
	"strings"
	"sync"
)

// 处理 CSV 文件并将数据插入数据库
func processCSV(filePath string, db *gorm.DB) error {
	// 打开 CSV 文件
	file, err := os.Open(filePath)
	if err != nil {
		return err
	}
	defer file.Close()

	// 创建 CSV 阅读器
	reader := csv.NewReader(file)

	// 读取所有行
	records, err := reader.ReadAll()
	if err != nil {
		return err
	}

	// 使用 WaitGroup 来同步所有的 goroutine
	var wg sync.WaitGroup

	// 通道用于发送每行数据
	ch := make(chan Contact, len(records))

	// 启动多个 goroutine 来并发处理 CSV 数据
	for i := 1; i < len(records); i++ { // 从 1 开始，跳过标题行
		wg.Add(1)
		go func(record []string) {
			defer wg.Done()
			// 将 CSV 行转换为 Contact 实例
			contact := Contact{
				Name:  record[0],
				Phone: record[1],
				Email: record[2],
			}
			ch <- contact // 发送数据到通道
		}(records[i])
	}

	// 启动一个 goroutine 来将通道中的数据插入到数据库
	go func() {
		for contact := range ch {
			if err := db.Create(&contact).Error; err != nil {
				fmt.Println("Error inserting record:", err)
			}
		}
	}()

	// 等待所有 goroutine 完成
	wg.Wait()

	// 关闭通道
	close(ch)

	return nil
}

func main() {
	// 初始化数据库
	db, err := initDB()
	if err != nil {
		fmt.Println("Failed to connect to database:", err)
		return
	}
	defer db.Close()

	// 处理 CSV 文件并将数据迁移到数据库
	err = processCSV("contacts.csv", db)
	if err != nil {
		fmt.Println("Error processing CSV file:", err)
		return
	}

	fmt.Println("CSV data successfully migrated to the database.")
}

4. 代码说明

a.初始化数据库：

initDB 函数用于初始化 MySQL 数据库连接并进行自动迁移。
我们使用 gorm 来处理数据库操作，模型 Contact 映射到数据库中的 contacts 表。

b.读取 CSV 文件：

processCSV 函数打开并读取 CSV 文件。然后，它读取所有的记录，并将每条记录通过 goroutine 异步发送到通道中。
每个 goroutine 都会将一条记录从 CSV 转换为 Contact 对象，并将其发送到通道。

c.并发处理数据：

sync.WaitGroup 被用来确保所有的 goroutine 完成任务。wg.Add(1) 在启动每个 goroutine 时调用，wg.Done() 在每个 goroutine 完成时调用。
使用 chan Contact 通道来将数据从多个 goroutine 传递到数据库插入部分。一个单独的 goroutine 从通道中接收数据并将其插入到数据库。

d.并发插入数据库：

每个 goroutine 向通道发送数据，然后另一个 goroutine 从通道中读取数据并将其插入数据库。通过这种方式，多个数据库插入操作是并发进行的。

e.关闭通道与等待：

在所有数据都发送到通道后，使用 wg.Wait() 等待所有 goroutine 完成处理。
关闭通道以确保数据库插入操作可以顺利结束。

5. 性能优化

在这个例子中，我们并发地读取 CSV 文件并将数据插入数据库，显著提高了处理速度。但是，对于大型数据集，还可以做更多的性能优化：

批量插入：可以将多个数据条目批量插入数据库，而不是每次插入一条记录。批量插入可以显著减少数据库的 I/O 操作，提升性能。
控制并发数：通过 semacphore 或者限制通道缓冲区大小，可以控制并发数，避免数据库被过多并发请求压垮。
数据库连接池：确保数据库连接池的配置合理，避免过多的并发连接造成数据库连接耗尽。

6. 总结

通过并发处理，我们能够大大提升 CSV 文件迁移到数据库的速度。Go 的 goroutines 和通道非常适合这种类型的任务，可以高效地处理 I/O 密集型的操作。

在处理大型 CSV 文件时，使用并发处理可以显著提升性能，减少总体处理时间。