283 lines
6.3 KiB
Go
283 lines
6.3 KiB
Go
package main
|
|
|
|
import (
|
|
"fmt"
|
|
"golang.org/x/exp/maps"
|
|
"strings"
|
|
|
|
"github.com/mozillazg/go-pinyin"
|
|
"github.com/xuri/excelize/v2"
|
|
"gorm.io/gorm"
|
|
"mongo.games.com/goserver/core/logger"
|
|
"mongo.games.com/goserver/core/mysqlx"
|
|
)
|
|
|
|
type ExcelData struct {
|
|
*excelize.File
|
|
Head []string
|
|
Index int
|
|
IndexCell int
|
|
|
|
TableName string
|
|
TableHead []string
|
|
DataName string
|
|
}
|
|
|
|
func (e *ExcelData) SetHead(head []string) *ExcelData {
|
|
if e == nil {
|
|
return nil
|
|
}
|
|
e.Index = 1
|
|
e.SetSheetRow("Sheet1", "A1", &head)
|
|
return e
|
|
}
|
|
|
|
func (e *ExcelData) SetRow(row []string) *ExcelData {
|
|
if e == nil {
|
|
return nil
|
|
}
|
|
e.Index++
|
|
e.SetSheetRow("Sheet1", "A"+fmt.Sprintf("%d", e.Index), &row)
|
|
return e
|
|
}
|
|
|
|
func (e *ExcelData) SetCell(val interface{}) *ExcelData {
|
|
if e == nil {
|
|
return nil
|
|
}
|
|
e.IndexCell++
|
|
cell := fmt.Sprintf("%c%d", 'A'+e.IndexCell-1, e.Index)
|
|
e.SetCellValue("Sheet1", cell, val)
|
|
return e
|
|
}
|
|
|
|
func (e *ExcelData) NewLine() *ExcelData {
|
|
if e == nil {
|
|
return nil
|
|
}
|
|
e.Index++
|
|
e.IndexCell = 0
|
|
return e
|
|
}
|
|
|
|
type ExcelMgr struct {
|
|
List map[int]*ExcelData
|
|
}
|
|
|
|
func NewExcelMgr() *ExcelMgr {
|
|
return &ExcelMgr{
|
|
List: make(map[int]*ExcelData),
|
|
}
|
|
}
|
|
|
|
func (e *ExcelMgr) Register(id int, head []string, filename string) *ExcelData {
|
|
e.List[id] = &ExcelData{
|
|
File: excelize.NewFile(),
|
|
Head: head,
|
|
Index: 0,
|
|
}
|
|
|
|
name := ChineseToPinYin([]string{filename})
|
|
if len(name) == 0 {
|
|
return nil
|
|
}
|
|
tableName := fmt.Sprintf("task_%s", name[0])
|
|
e.List[id].TableName = tableName
|
|
|
|
for _, v := range head {
|
|
cl := ChineseToPinYin([]string{v})[0]
|
|
e.List[id].TableHead = append(e.List[id].TableHead, cl)
|
|
}
|
|
e.List[id].DataName = filename
|
|
|
|
e.List[id].NewSheet("Sheet1")
|
|
|
|
if len(head) > 0 {
|
|
e.List[id].SetHead(head)
|
|
}
|
|
|
|
return e.List[id]
|
|
}
|
|
|
|
func (e *ExcelMgr) Get(id int) *ExcelData {
|
|
return e.List[id]
|
|
}
|
|
|
|
func (e *ExcelMgr) Save(id int, startTime, endTime string) error {
|
|
d := e.List[id]
|
|
if d == nil {
|
|
return nil
|
|
}
|
|
|
|
filename := fmt.Sprintf("%s_%s_%s.xlsx", d.DataName, startTime, endTime)
|
|
|
|
if VP.GetBool("IsDatabaseMode") {
|
|
rows, err := d.GetRows("Sheet1")
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
db, err := mysqlx.GetDatabase("1")
|
|
if err != nil {
|
|
logger.Logger.Errorf("GetDatabase error: %v", err)
|
|
return err
|
|
}
|
|
|
|
files := make(map[string]string)
|
|
index := make(map[string]string)
|
|
for _, v := range d.Head {
|
|
cl := ChineseToPinYin([]string{v})[0]
|
|
d.TableHead = append(d.TableHead, cl)
|
|
files[cl] = v
|
|
if strings.Contains(v, "*") {
|
|
index[cl] = "INDEX"
|
|
}
|
|
}
|
|
|
|
createSQL := buildCreateTableSQLWithIndices(d.TableName, d.DataName, files, index)
|
|
|
|
if err = db.Exec(createSQL).Error; err != nil {
|
|
logger.Logger.Errorf("createTable error: %v", err)
|
|
return err
|
|
}
|
|
|
|
if err = insertData(db.DB, d.TableName, maps.Keys(files), rows); err != nil {
|
|
logger.Logger.Errorf("insertData error: %v", err)
|
|
return err
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
return d.SaveAs(filename)
|
|
}
|
|
|
|
func (e *ExcelMgr) Pull(id int, startTime, endTime string) (*excelize.File, string, error) {
|
|
d := e.List[id]
|
|
if d == nil {
|
|
return nil, "", nil
|
|
}
|
|
|
|
db, err := mysqlx.GetDatabase("1")
|
|
if err != nil {
|
|
logger.Logger.Errorf("GetDatabase error: %v", err)
|
|
return nil, "", err
|
|
}
|
|
|
|
rows, err := db.Table(d.TableName).Where("ri_qi between ? and ?", startTime, endTime).Select(d.TableHead).Rows()
|
|
if err != nil {
|
|
logger.Logger.Errorf("Pull error: %v", err)
|
|
return nil, "", err
|
|
}
|
|
defer rows.Close()
|
|
|
|
f := excelize.NewFile()
|
|
f.NewSheet("Sheet1")
|
|
f.SetSheetRow("Sheet1", "A1", &d.Head)
|
|
|
|
index := 2
|
|
for rows.Next() {
|
|
data := make([]interface{}, len(d.TableHead))
|
|
scanArgs := make([]interface{}, len(d.TableHead))
|
|
for i := range data {
|
|
scanArgs[i] = &data[i]
|
|
}
|
|
err = rows.Scan(scanArgs...)
|
|
if err != nil {
|
|
logger.Logger.Errorf("Pull error: %v", err)
|
|
return nil, "", err
|
|
}
|
|
f.SetSheetRow("Sheet1", "A"+fmt.Sprintf("%d", index), &data)
|
|
index++
|
|
}
|
|
|
|
return f, fmt.Sprintf("%s_%s_%s", d.DataName, startTime, endTime), nil
|
|
}
|
|
|
|
// 构建创建表的 SQL 语句,支持中文描述和索引
|
|
func buildCreateTableSQLWithIndices(tableName string, comment string, fields map[string]string, indices map[string]string) string {
|
|
var columns []string
|
|
var indexDefs []string
|
|
|
|
// 遍历字段定义
|
|
for field, comment := range fields {
|
|
column := fmt.Sprintf("`%s` VARCHAR(255) COMMENT '%s'", field, comment)
|
|
columns = append(columns, column)
|
|
|
|
// 根据索引类型添加索引定义
|
|
if indexType, ok := indices[field]; ok {
|
|
switch strings.ToUpper(indexType) {
|
|
case "PRIMARY":
|
|
indexDefs = append(indexDefs, fmt.Sprintf("PRIMARY KEY (`%s`)", field))
|
|
case "UNIQUE":
|
|
indexDefs = append(indexDefs, fmt.Sprintf("UNIQUE KEY `idx_%s` (`%s`)", field, field))
|
|
case "INDEX":
|
|
indexDefs = append(indexDefs, fmt.Sprintf("KEY `idx_%s` (`%s`)", field, field))
|
|
}
|
|
}
|
|
}
|
|
|
|
// 拼接字段和索引部分
|
|
allDefs := append(columns, indexDefs...)
|
|
|
|
ret := fmt.Sprintf("CREATE TABLE IF NOT EXISTS `%s` (\n%s\n) COMMENT='%s';",
|
|
tableName,
|
|
strings.Join(allDefs, ",\n"), comment)
|
|
logger.Logger.Tracef("createTableSQL: %s", ret)
|
|
// 拼接最终的 SQL
|
|
return ret
|
|
}
|
|
|
|
func insertData(db *gorm.DB, tableName string, header []string, rows [][]string) error {
|
|
placeholders := strings.Repeat("?,", len(header))
|
|
placeholders = placeholders[:len(placeholders)-1] // 移除多余的逗号
|
|
|
|
insertSQL := fmt.Sprintf("INSERT INTO `%s` (%s) VALUES (%s)", tableName, "`"+strings.Join(header, "`,`")+"`", placeholders)
|
|
|
|
for _, row := range rows {
|
|
// 确保每行数据长度和表头一致
|
|
if len(row) < len(header) {
|
|
for len(row) < len(header) {
|
|
row = append(row, "") // 填充缺失列
|
|
}
|
|
}
|
|
err := db.Exec(insertSQL, interfaceSlice(row)...).Error
|
|
if err != nil {
|
|
return err
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func interfaceSlice(slice []string) []interface{} {
|
|
result := make([]interface{}, len(slice))
|
|
for i, v := range slice {
|
|
result[i] = v
|
|
}
|
|
return result
|
|
}
|
|
|
|
var args = pinyin.NewArgs()
|
|
|
|
func init() {
|
|
args.Style = pinyin.Normal
|
|
}
|
|
|
|
func joinPinyin(py [][]string) string {
|
|
var result []string
|
|
for _, word := range py {
|
|
if len(word) > 0 {
|
|
result = append(result, word[0]) // 每个字的拼音取第一个候选
|
|
}
|
|
}
|
|
return strings.Join(result, "_")
|
|
}
|
|
|
|
func ChineseToPinYin(param []string) []string {
|
|
var ret []string
|
|
for _, v := range param {
|
|
ret = append(ret, joinPinyin(pinyin.Pinyin(v, args)))
|
|
}
|
|
return ret
|
|
}
|