Golang并发读取文件数据并写入数据库

需求

最近接到一个任务，要把一批文件中的十几万条JSON格式数据写入到Oracle数据库中，Oracle是企业级别的数据库向来以高性能著称，所以尽可能地利用这一特性。当时第一时间想到的就是用多线程并发读文件并操作数据库，而Golang是为并发而生的，用Golang进行并发编程非常方便，因此这里选用Golang并发读取文件并用Gorm操作数据库。然而Gorm官方并不支持Oraclc，所以要借助第三方驱动，之前写了篇文章来记录Gorm操作Oracle的踩坑，详见使用Gorm操作Oracle数据库踩坑

项目结构

data文件夹中包含数十个.out结尾的数据文件，model.go声明数据类型，main.go中编写并发逻辑和数据库操作代码

|——db_test
|    |——data
|        |——xxx.out
|        |——yyy.out
|    |——model
|        |——model.go
|    |——main.go
|    |——go.mod

获取data目录下的文件

Golang自带的os库就可以对文件、目录进行各种丰富的操作，OpenFile函数第一个参数是目录的路径，第二个参数表示只读，第三个参数os.ModeDir表示以文件夹模式打开。ReadDir传入负数表示读取目录下所有文件信息，传入n表示读取前n个文件信息。最后将所有文件名保存到字符串数组并返回。

func loadFile(path string) []string {
        // 打开指定文件夹
	f, err := os.OpenFile(path, os.O_RDONLY, os.ModeDir)
	if err != nil {
		log.Fatalln(err.Error())
		os.Exit(0)
	}
	defer f.Close()
        // 读取目录下所有文件
	fileInfo, _ := f.ReadDir(-1)

	files := make([]string, 0)
	for _, info := range fileInfo {
		files = append(files, info.Name())
	}
	return files
}

按行读取文本数据

这里使用bufio.Scanner来一行一行读取JSON格式的数据，bufio.Reader也能实现按行读取，但bufio.Scanner是go1.1后开发的模块操作起来更简单一点。

func readRecord(filename string) {
	log.Println(filename)
	f, err := os.Open(filename)
	if err != nil {
		log.Println(filename + " error")
		return
	}
	defer f.Close()
	scanner := bufio.NewScanner(f)
	for scanner.Scan() {
		line := scanner.Text() // line就是每行文本
                // 对line进行处理
	}
}

数据类型定义

还是假设数据库中有一个SHOPS表，结构体方法TableName指定该类型对应的数据表，编写如下model.go文件

package model

type ShopInfo struct {
	ShopId   string `gorm:"column:SHOPID;not null"`
	ShopName string `gorm:"column:SHOPNAME;not null"`
        // 省略剩余的字段
}
func (s *ShopInfo) TableName() string {
	return "SHOPS"
}

并发读取文件

基本逻辑是主函数读取文件夹下面的所有文件，然后用循环开启goroutine并传入文件名和数据库指针，goroutine中按行读取每个文件并将其JSON数据转换为结构体，在调用Gorm写入Oracle数据库。这里用Golang的等待组来同步主函数与goroutine。

var wg sync.WaitGroup
func main() {
        // 打开Oracle连接
	db, err := gorm.Open(oracle.Open("database/password@127.0.0.1:1521/XE"), &gorm.Config{
		Logger: logger.New(log.New(os.Stdout, "\r\n", log.LstdFlags), logger.Config{
			SlowThreshold: 1 * time.Millisecond,
			LogLevel:      logger.Error,
			Colorful:      true,
		}),
	})
	if err != nil {
		log.Fatalln(err)
	}

	if e := db.AutoMigrate(&model.ShopInfo{}); e != nil {
		log.Fatalln(e.Error())
	}

	path := "./data/"
	files := loadFile(path) // 加载所有文件名
        // 循环创建goroutine
	for i, v := range files {
		wg.Add(1)
                // 将数据库指针和文件名传给goroutine处理
		go writeRecord(db, path+v)
	}

	wg.Wait() // 等待所有goroutine执行完成
	log.Println("over")

}

将数据写入数据库

由于这些文件中可能有重复的数据，所以这里调用了Gorm的Clauses设置，当有主键重复的数据什么都不做，有些情况下主键相同但是更新了某些字段，这时可以用Clauses设置主键重复时进行更新操作。虽然主键重复时什么都不做，但是db的执行结果也会包含"unique constraint"错误，所以在错误处理时要排除主键冲突的情况，把其他错误（如字段太长或类型不匹配）记录下来。

func writeRecord(db *gorm.DB, filename string) {
	defer wg.Done() // 不要忘记等待组-1
	f, err := os.Open(filename)
	if err != nil {
		log.Println(filename + " error")
		return
	}
	defer f.Close()
	scanner := bufio.NewScanner(f)
	iter := 0 // 记录出错的行数
	for scanner.Scan() {
		var shop model.ShopInfo
		iter++
		// 调用json.Unmarshal()将文本转换为结构体
		if err = json.Unmarshal([]byte(scanner.Text()), &shop); err != nil {
			log.Println("转换错误--->" + scanner.Text())
			return
		}
		// 用clause设置当发生ID冲突时什么都不做
		res := db.Clauses(clause.OnConflict{DoNothing: true}).Create(&shop)
		// 虽然ID相同时程序不会停止，但是还是有错误返回
		// 所以这里排除ID冲突错误，将其他错误（字段冲突）打印出来
		if res.Error != nil && !strings.Contains(res.Error.Error(), "unique constraint") {
			log.Println("插入出错--->" + shop.ShopId + " 在" + filename + "第" + strconv.Itoa(iter) + "行")
			return
		}
	}
}

完整main.go代码

将上面每一步整合后得到完整的主函数代码如下：

package main

import (
	"bufio"
	"db_test/model"
	"encoding/json"
	"log"
	"os"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/cengsin/oracle"
	"gorm.io/gorm"
	"gorm.io/gorm/clause"
	"gorm.io/gorm/logger"
)

var wg sync.WaitGroup

func main() {
	log.Println("initial database connect……")
	db, err := gorm.Open(oracle.Open("database/password@127.0.0.1:1521/XE"), &gorm.Config{
		Logger: logger.New(log.New(os.Stdout, "\r\n", log.LstdFlags), logger.Config{
			SlowThreshold: 1 * time.Millisecond,
			LogLevel:      logger.Error,
			Colorful:      true,
		}),
	})
	if err != nil {
		log.Fatalln(err)
	}

	if e := db.AutoMigrate(&model.ShopInfo{}); e != nil {
		log.Fatalln(e.Error())
	}

	path := "../out1/"
	files := loadFile(path)
	time.Sleep(2 * time.Second)
	for i, v := range files {
		wg.Add(1)
		go writeRecord(db, path+v)
	}

	wg.Wait()
	log.Println("over")

}

func loadFile(path string) []string {
        // 打开指定文件夹
	f, err := os.OpenFile(path, os.O_RDONLY, os.ModeDir)
	if err != nil {
		log.Fatalln(err.Error())
		os.Exit(0)
	}
	defer f.Close()
        // 读取目录下所有文件
	fileInfo, _ := f.ReadDir(-1)

	files := make([]string, 0)
	for _, info := range fileInfo {
		files = append(files, info.Name())
	}
	return files
}

func writeRecord(db *gorm.DB, filename string) {
	defer wg.Done()
	f, err := os.Open(filename)
	if err != nil {
		log.Println(filename + " error")
		return
	}
	defer f.Close()
	scanner := bufio.NewScanner(f)
	iter := 0 // 记录出错的行数
	for scanner.Scan() {
		var shop model.ShopInfo
		iter++
		// 调用json.Unmarshal()将文本转换为结构体
		if err = json.Unmarshal([]byte(scanner.Text()), &shop); err != nil {
			log.Println("转换错误--->" + scanner.Text())
			return
		}
		// 用clause设置当发生ID冲突时什么都不做
		res := db.Clauses(clause.OnConflict{DoNothing: true}).Create(&shop)
		// 虽然ID相同时程序不会停止，但是还是有错误返回
		// 所以这里排除ID冲突错误，将其他错误（字段冲突）打印出来
		if res.Error != nil && !strings.Contains(res.Error.Error(), "unique constraint") {
			log.Println("插入出错--->" + shop.ShopId + " 在" + filename + "第" + strconv.Itoa(iter) + "行")
			return
		}
	}
}

测试运行

go run ./main.go运行过程非常快，十几万条数据几分钟就写完了，并且CPU占用率100%，证明非常有效的利用了并发优势。若是文件数量太多（上千个）的话会创建非常多goroutine，可能消耗非常多系统资源，可以在循环创建goroutine时进行限制，只创建30或50个，一个goroutine结束后再给它传入一个新的文件名。

# golang # 数据库
Logo
华为云开发者联盟

为开发者提供学习成长、分享交流、生态实践、资源工具等服务，帮助开发者快速成长。

更多推荐

cover

从原始边列表到邻接矩阵Python实现图数据处理的完整指南

avatar 华为云开发者联盟
cover

华为云云原生FinOps解决方案，释放云原生最大价值

avatar 华为云开发者联盟
cover

解锁HDC 2024之旅：从购票到报名，全程攻略

avatar 华为云开发者联盟