工具基本实现write功能

cmd/saastool/convert.go

@@ -0,0 +1,189 @@
+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
+}