excel-export/biz/db/db.go

package db

import (
	"context"
	"database/sql"
	"excel_export/biz/config"
	"excel_export/biz/export"
	"excel_export/pkg"
	"fmt"
	"gorm.io/driver/mysql"
	"gorm.io/gorm"
	"runtime/trace"
	"strconv"
	"time"
)

var _ export.DataFetcher = new(Db)

type Db struct {
	db *gorm.DB
}

func NewDb(str string) (*Db, error) {
	db, err := gorm.Open(
		mysql.Open(str+""),
		&gorm.Config{
			//Logger: logger.Discard,
		},
	)
	if err != nil {
		return nil, err
	}
	db = db.Debug()
	return &Db{
		db: db,
	}, nil
}

func (d *Db) Fetch(s string) (*export.Data, error) {
	fetchRegion := trace.StartRegion(context.Background(), "db.fetch")
	defer func() {
		fetchRegion.End()
	}()
	rows, err := d.db.Raw(s).Rows()
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	titles, err := rows.Columns()
	if err != nil {
		return nil, err
	}

	data := getData(rows, d.db, titles)
	dataMap := getDataMap(titles, data)
	return &export.Data{
		Title:   titles,
		Data:    data,
		DataMap: dataMap,
	}, nil
}

func getDataMap(titles []string, data [][]string) []map[string]string {
	result := make([]map[string]string, 0, len(data))
	for _, row := range data {
		rowMap := make(map[string]string, len(row))
		for i, value := range row {
			rowMap[titles[i]] = value
		}
		result = append(result, rowMap)
	}
	return result
}

func getData(rows *sql.Rows, db *gorm.DB, titles []string) [][]string {
	result := make([][]string, 0, 10)
	for rows.Next() {
		var row map[string]interface{}
		db.ScanRows(rows, &row)
		result = append(result, transformRow(titles, row))
	}
	return result
}

func transform(titles []string, data []map[string]interface{}) [][]string {
	result := make([][]string, len(data))
	for i, m := range data {
		result[i] = transformRow(titles, m)
	}
	return result
}

func transformRow(titles []string, data map[string]interface{}) []string {
	row := make([]string, 0, len(data))
	for _, title := range titles {
		col := data[title]
		switch v := col.(type) {
		case string:
			row = append(row, v)
		case time.Time:
			row = append(row, v.Format("2006-01-02 15:04:05"))
		case int, int8, int16, int32, int64:
			row = append(row, fmt.Sprintf("%d", v))
		case float64:
			// When formatting floats, do not use fmt.Sprintf("%v", n), this will cause numbers below 1e-4 to be printed in
			// scientific notation. Scientific notation is not a valid way to store numbers in XML.
			// Also not not use fmt.Sprintf("%f", n), this will cause numbers to be stored as X.XXXXXX. Which means that
			// numbers will lose precision and numbers with fewer significant digits such as 0 will be stored as 0.000000
			// which causes tests to fail.
			row = append(row, strconv.FormatFloat(v, 'f', -1, 64))
		case float32:
			row = append(row, strconv.FormatFloat(float64(v), 'f', -1, 32))
		case []byte:
			row = append(row, string(v))
		case nil:
			row = append(row, "")
		default:
			row = append(row, fmt.Sprintf("%v", v))
		}
	}
	return row

}

func ServerData(conf *config.Config, serverName string) (*export.Data, error) {
	configInfo, err := config.GetServer(conf, serverName)
	if err != nil {
		return nil, err
	}

	d, err := NewDb(configInfo.Db)
	if err != nil {
		return nil, err
	}
	data, err := d.Fetch(configInfo.Sql)
	return data, err
}

func ResellerData(conf *config.Config) (map[string]map[string]string, error) {
	resellerData, err := ServerData(conf, "reseller")
	if err != nil {
		return nil, err
	}
	directResellerData, err := ServerData(conf, "direct_reseller")
	if err != nil {
		return nil, err
	}
	relation := MergeReseller(resellerData, directResellerData)
	return relation, nil
}

func MergeReseller(reseller *export.Data, directReseller *export.Data) map[string]map[string]string {
	result := make(map[string]map[string]string, len(reseller.DataMap))
	directMap := make(map[string]map[string]string, len(directReseller.DataMap))
	for _, row := range directReseller.DataMap {

		directMap[row["direct_id"]] = row
	}
	for _, row := range reseller.DataMap {
		directRow, ok := directMap[row["direct_reseller_id"]]
		if ok {
			row = pkg.MergeMaps(row, directRow)
		}
		result[row["id"]] = row
	}

	return result
}