用 Go 语言并发处理 CSV 文件到数据库

文章精选推荐

1 JetBrains Ai assistant 编程工具让你的工作效率翻倍
2 Extra Icons：JetBrains IDE的图标增强神器
3 IDEA插件推荐-SequenceDiagram，自动生成时序图
4 BashSupport Pro 这个ides插件主要是用来干嘛的？
5 IDEA必装的插件：Spring Boot Helper的使用与功能特点
6 Ai assistant ,又是一个写代码神器

文章正文

问题背景

假设你拥有一个包含大量联系人信息的 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 mainimport ("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 mainimport ("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 来同步所有的 goroutinevar 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. 代码说明

初始化数据库：
- initDB 函数用于初始化 MySQL 数据库连接并进行自动迁移。
- 我们使用 gorm 来处理数据库操作，模型 Contact 映射到数据库中的 contacts 表。
读取 CSV 文件：
- processCSV 函数打开并读取 CSV 文件。然后，它读取所有的记录，并将每条记录通过 goroutine 异步发送到通道中。
- 每个 goroutine 都会将一条记录从 CSV 转换为 Contact 对象，并将其发送到通道。
并发处理数据：
- sync.WaitGroup 被用来确保所有的 goroutine 完成任务。wg.Add(1) 在启动每个 goroutine 时调用，wg.Done() 在每个 goroutine 完成时调用。
- 使用 chan Contact 通道来将数据从多个 goroutine 传递到数据库插入部分。一个单独的 goroutine 从通道中接收数据并将其插入到数据库。
并发插入数据库：
- 每个 goroutine 向通道发送数据，然后另一个 goroutine 从通道中读取数据并将其插入数据库。通过这种方式，多个数据库插入操作是并发进行的。
关闭通道与等待：
- 在所有数据都发送到通道后，使用 wg.Wait() 等待所有 goroutine 完成处理。
- 关闭通道以确保数据库插入操作可以顺利结束。