saasapi/cmd/saastool/convert.go

package main

import (
	"bufio"
	"flag"
	"fmt"
	"os"
	"path"
	"strings"

	"e.coding.net/rta/public/saasapi"
	"google.golang.org/protobuf/encoding/protojson"
)

// TODO 转换加速

type convertParams struct {
	targetCfg  *TargetConfig
	sourcePath string
	destPath   string
}

func RunConvert(args ...string) error {
	fs := flag.NewFlagSet("convert", flag.ExitOnError)
	targetCfgFile := paramTargets(fs)
	sourcePath := paramSourcePath(fs)
	destPath := paramDestPath(fs)

	if err := fs.Parse(args); err != nil {
		fmt.Println("command line parse error", "err", err)
		return err
	}

	if fs.NArg() > 0 || *targetCfgFile == "" || len(*sourcePath) == 0 || len(*destPath) == 0 {
		fs.PrintDefaults()
		return nil
	}

	targetCfg, err := LoadTargetFile(*targetCfgFile)
	if err != nil {
		fmt.Println("LoadConfigFile error", "err", err)
		return err
	}

	convertParams := convertParams{
		targetCfg:  targetCfg,
		sourcePath: *sourcePath,
		destPath:   *destPath,
	}

	return doConvert(convertParams)
}

func doConvert(convertParams convertParams) error {
	fsInfo, err := os.Stat(convertParams.sourcePath)
	if err != nil {
		return err
	}

	if !fsInfo.IsDir() {
		// 如果是文件，直接写入
		return doFileConvert(convertParams)
	}

	// 读取目录下信息
	dirEntry, err := os.ReadDir(convertParams.sourcePath)
	if err != nil {
		return err
	}

	// 遍历目录
	for _, dir := range dirEntry {
		newParam := convertParams
		newParam.sourcePath = path.Join(convertParams.sourcePath, dir.Name())

		if dir.IsDir() {
			newParam.destPath = path.Join(convertParams.destPath, dir.Name())
		}

		if err = doConvert(newParam); err != nil {
			return err
		}
	}

	return nil
}

func doFileConvert(convertParams convertParams) error {
	// 读取文件并按行遍历，以\t分割为两列，第一列为userid，第二列解析为string数组
	sourceFile, err := os.Open(convertParams.sourcePath)
	if err != nil {
		return err
	}
	defer sourceFile.Close()

	if _, err = os.Stat(convertParams.destPath); os.IsNotExist(err) {
		os.MkdirAll(convertParams.destPath, os.ModePerm)
	}

	destName := path.Join(convertParams.destPath, path.Base(convertParams.sourcePath)+".converted")
	destFile, err := os.Create(destName)
	if err != nil {
		return err
	}
	defer destFile.Close()

	scaner := bufio.NewScanner(sourceFile)
	destWriter := bufio.NewWriter(destFile)
	defer destWriter.Flush()

	jasonMarshal := protojson.MarshalOptions{Multiline: false, Indent: ""}

	processedLine := 0
	for scaner.Scan() {
		line := scaner.Text()
		if line == "" {
			continue
		}

		// 按\t分割为两列
		parts := strings.Split(line, "\t")
		if len(parts) != 2 {
			continue
		}

		// 读取userid
		userid := parts[0]
		value := parts[1]
		value = strings.ReplaceAll(value, "[", "")
		value = strings.ReplaceAll(value, "]", "")
		// 第二列解析为string数组
		targets := strings.Split(value, " ")

		saasWriteCmd := &saasapi.WriteCmd{
			Userid: userid,
		}
		if len(userid) == 0 || len(targets) == 0 {
			continue
		}
		for _, target := range targets {
			if targetinfo, ok := convertParams.targetCfg.Targets[target]; ok {
				if targetinfo.WriteByte != nil {
					if saasWriteCmd.WriteBytes == nil {
						saasWriteCmd.WriteBytes = &saasapi.Bytes{}
					}
					saasWriteCmd.WriteBytes.Bytes = append(saasWriteCmd.WriteBytes.Bytes, *targetinfo.WriteByte)
					if targetinfo.WriteBytePos < 64 {
						saasWriteCmd.WriteBytes.Index_1 |= 1 << targetinfo.WriteBytePos
					} else if targetinfo.WriteBytePos < 128 {
						saasWriteCmd.WriteBytes.Index_2 |= 1 << (targetinfo.WriteBytePos - 64)
					}
				}

				if targetinfo.WriteUint32 != nil {
					if saasWriteCmd.WriteUint32S == nil {
						saasWriteCmd.WriteUint32S = &saasapi.Uint32S{}
					}
					saasWriteCmd.WriteUint32S.Uint32S = append(saasWriteCmd.WriteUint32S.Uint32S, *targetinfo.WriteUint32)
					saasWriteCmd.WriteUint32S.Index_1 |= 1 << targetinfo.WriteUint32Pos
				}

				if targetinfo.WriteFlag != nil && targetinfo.WriteExpire != nil {
					if saasWriteCmd.WriteFlagsWithExpire == nil {
						saasWriteCmd.WriteFlagsWithExpire = &saasapi.FlagsWithExpire{}
					}
					saasWriteCmd.WriteFlagsWithExpire.FlagsWithExpire = append(
						saasWriteCmd.WriteFlagsWithExpire.FlagsWithExpire, &saasapi.FlagWithExpire{
							Flag:   *targetinfo.WriteFlag,
							Expire: *targetinfo.WriteExpire,
						})
					saasWriteCmd.WriteFlagsWithExpire.Index_1 |= 1 << targetinfo.WriteFlagWithExpirePos
				}
			}
		}

		// 写入文件
		destWriter.WriteString(jasonMarshal.Format(saasWriteCmd))
		destWriter.WriteString("\n")

		processedLine++
		if processedLine%100000 == 0 {
			fmt.Printf("\rconverted records: %v [%v]", processedLine, destName)
		}

	}

	fmt.Printf("\rconverted records: %v [%v]\n", processedLine, destName)
	return nil
}