excel 파일 내용 sql 로 작성 (Go, postgresql)

w00j00ng351·2023년 2월 7일

Go

목록 보기
2/6
package main

import (
	"fmt"
	"log"
	"os"
	"path/filepath"
	"strings"

	"github.com/360EntSecGroup-Skylar/excelize"
)

func main() {
	log.SetFlags(log.Lshortfile)

	const DIRECTORY_PATH = "./targets"
	const BATCH_SIZE = 100

	targetDirectory, err := os.Open(DIRECTORY_PATH)
	if err != nil {
		log.Fatalln(err.Error())
	}
	defer targetDirectory.Close()
	files, err := targetDirectory.ReadDir(0)
	if err != nil {
		log.Fatalln(err.Error())
	}

	for _, fDirEntry := range files {
		if fDirEntry.IsDir() {
			continue
		}
		f, err := os.Open(filepath.Join(DIRECTORY_PATH, fDirEntry.Name()))
		if err != nil {
			log.Fatalln(err.Error())
		}
		defer f.Close()

		excelFile, err := excelize.OpenReader(f)
		if err != nil {
			continue
		}

		tableName := strings.ReplaceAll(fDirEntry.Name(), ".xlsx", "")
		filePath, err := filepath.Abs(fmt.Sprintf("./%s.sql", tableName))
		if err != nil {
			log.Fatalln(err.Error())
		}
		sqlFile, err := os.OpenFile(filePath, os.O_CREATE|os.O_TRUNC|os.O_RDWR, os.ModePerm)
		if err != nil {
			log.Fatalln(err.Error())
		}
		defer sqlFile.Close()

		rowCount := 0
		sheetIndex := 0
		headers := []string{}
		sqlWriter := &strings.Builder{}
		sqlWriter.WriteString(fmt.Sprintf("TRUNCATE %s CASCADE;\n\n", tableName))
		sqlWriter.WriteString(fmt.Sprintf("ALTER SEQUENCE %s_id_seq RESTART WITH 1;", tableName))
		for _, sheetName := range excelFile.GetSheetMap() {
			rows, err := excelFile.Rows(sheetName)
			if err != nil {
				log.Fatalln(err.Error())
			}
			rowIndex := 0
			for rows.Next() {
				rowValues := rows.Columns()
				bValidRow := false
				for _, v := range rowValues {
					if v != "" {
						bValidRow = true
						break
					}
				}
				if !bValidRow {
					continue
				}
				if rowIndex == 0 || rowIndex > BATCH_SIZE {
					if rowIndex != 0 {
						sqlWriter.WriteString(fmt.Sprintf("; -- %d", rowCount))
					}
					rowIndex = 0
					sqlWriter.WriteString(fmt.Sprintf("\n\nINSERT INTO %s ", tableName))
					if sheetIndex == 0 {
						headers = rowValues
					} else if len(rowValues) != len(headers) {
						log.Fatalf("invalid header count (filePath: %s, sheetName: %s)\n", filePath, sheetName)
					}
					sqlWriter.WriteString("(")
					for columnIndex, h := range headers {
						if columnIndex != 0 {
							sqlWriter.WriteString(", ")
						}
						sqlWriter.WriteString(fmt.Sprintf("\"%s\"", h))
					}
					sqlWriter.WriteString(")")
					sqlWriter.WriteString(" VALUES \n  ")
					rowIndex++
					sheetIndex++
					continue
				}
				if rowIndex != 1 {
					sqlWriter.WriteString(fmt.Sprintf(", -- %d \n  ", rowCount))
				}
				sqlWriter.WriteString("(")
				for columnIndex, v := range rowValues {
					if columnIndex > len(headers)-1 {
						break
					}
					if columnIndex != 0 {
						sqlWriter.WriteString(", ")
					}
					sqlWriter.WriteString(fmt.Sprintf("'%s'", strings.ReplaceAll(v, "'", "''")))
				}
				sqlWriter.WriteString(")")
				rowCount++
				rowIndex++
			}
			sqlWriter.WriteString(fmt.Sprintf("; -- %d \n  ", rowCount))
			sheetIndex++
		}
		sqlWriter.WriteString(fmt.Sprintf("\nALTER SEQUENCE %s_id_seq RESTART WITH %d;", tableName, rowCount+1))
		if _, err := sqlFile.Write([]byte(sqlWriter.String())); err != nil {
			log.Fatalln(err.Error())
		}
	}
}
profile
시간이 만든 코드

0개의 댓글