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())
}
}
}