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

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

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
}

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.")
}