增加saastool功能

cmd/rtatool/main.go

@@ -1,5 +0,0 @@
-package main
-
-func main() {
-
-}

cmd/saastool/.gitignore

@@ -0,0 +1,2 @@
+debug/
+saastool

cmd/saastool/config.go

@@ -0,0 +1,30 @@
+package main
+
+import (
+	"github.com/BurntSushi/toml"
+)
+
+// Config 配置
+type Config struct {
+	Auth    Auth
+	ApiUrls ApiUrls
+}
+
+// DB 配置
+type Auth struct {
+	Account string
+	Token   string
+}
+
+type ApiUrls struct {
+	UrlBase string
+	Write   string
+	Read    string
+}
+
+// LoadConfigFile 加载配置文件
+func LoadConfigFile(filename string) (*Config, error) {
+	sc := &Config{}
+	_, err := toml.DecodeFile(filename, sc)
+	return sc, err
+}

cmd/saastool/help.go

@@ -0,0 +1,27 @@
+package main
+
+import (
+	"fmt"
+	"strings"
+)
+
+// RunHelp 帮助
+func RunHelp(args ...string) error {
+	fmt.Println(strings.TrimSpace(usage))
+	return nil
+}
+
+const usage = `
+Usage:  [[command] [arguments]]
+
+The commands are:
+
+    write              Write user's bytes / uint32s / flags
+    read               Read user's bytes / uint32s / flags
+
+"help" is the default command.
+
+Use "saastool [command] -help" for more information about a command.
+`
+
+//    strip            Stripping redundant data from redis

cmd/saastool/main.go

@@ -0,0 +1,33 @@
+package main
+
+import (
+	"fmt"
+	"log/slog"
+	"os"
+)
+
+func main() {
+	if err := Run(os.Args[1:]...); err != nil {
+		os.Exit(1)
+	}
+
+}
+
+func Run(args ...string) error {
+
+	name, args := ParseCommandName(args)
+
+	// 从参数中解析出命令
+	switch name {
+	case "", "help":
+		return RunHelp(args...)
+	case "write":
+		return RunWrite(args...)
+	//case "read":
+	//	return RunRead(args...)
+	default:
+		err := fmt.Errorf(`unknown command "%s"`+"\n"+`Run 'dmptool help' for usage`, name)
+		slog.Warn(err.Error())
+		return err
+	}
+}

cmd/saastool/params.go

@@ -0,0 +1,24 @@
+package main
+
+import (
+	"flag"
+)
+
+func paramConfig(fs *flag.FlagSet) *string {
+	return fs.String("config", "cfg.toml", "Config file.")
+}
+func paramTargets(fs *flag.FlagSet) *string {
+	return fs.String("targets", "", "target setting")
+}
+
+func paramFromPath(fs *flag.FlagSet) *string {
+	return fs.String("from", "", "Data path source for write command. (*required*)")
+}
+
+func paramBatchSize(fs *flag.FlagSet) *uint {
+	return fs.Uint("batchsize", 10000, "Batch size to sync")
+}
+
+func paramAsync(fs *flag.FlagSet) *bool {
+	return fs.Bool("async", false, "Async mode")
+}

cmd/saastool/parse.go

@@ -0,0 +1,25 @@
+package main
+
+import (
+	"strings"
+)
+
+func ParseCommandName(args []string) (string, []string) {
+	var name string
+	if len(args) > 0 {
+		if !strings.HasPrefix(args[0], "-") {
+			name = args[0]
+		} else if args[0] == "-h" || args[0] == "-help" || args[0] == "--help" {
+			name = "help"
+		}
+	}
+
+	if name == "help" && len(args) > 2 && !strings.HasPrefix(args[1], "-") {
+		return args[1], []string{"-h"}
+	}
+
+	if name != "" {
+		return name, args[1:]
+	}
+	return "", args
+}

cmd/saastool/read.go

@@ -0,0 +1 @@
+package main

cmd/saastool/target_cfg.go

@@ -0,0 +1,38 @@
+package main
+
+import (
+	"encoding/json"
+	"os"
+)
+
+// TargetConfig 配置
+type TargetConfig struct {
+	Targets map[string]*Target `json:"targets"`
+}
+
+// Target 配置
+type Target struct {
+	WriteByte              *byte   `json:"write_byte"`                 // byte值
+	WriteBytePos           int     `json:"write_byte_pos"`             // byte写入位置
+	WriteUint32            *uint32 `json:"write_uint32"`               // uint32值
+	WriteUint32Pos         int     `json:"write_uint32_pos"`           // uint32写入位置
+	WriteFlag              *bool   `json:"write_flag"`                 // 标志位
+	WriteExpire            *uint32 `json:"write_expire"`               // 过期时间
+	WriteFlagWithExpirePos int     `json:"write_flag_with_expire_pos"` // 标志与过期写入位置
+
+}
+
+// LoadConfigFile 加载配置文件
+func LoadTargetFile(filename string) (*TargetConfig, error) {
+	// 打开文件
+	f, err := os.Open(filename)
+	if err != nil {
+		return nil, err
+	}
+	defer f.Close()
+
+	sc := &TargetConfig{}
+
+	err = json.NewDecoder(f).Decode(sc)
+	return sc, err
+}

cmd/saastool/write.go

@@ -0,0 +1,205 @@
+package main
+
+import (
+	"bufio"
+	"flag"
+	"fmt"
+	"log/slog"
+	"os"
+	"path"
+	"strings"
+
+	"e.coding.net/rta/public/saasapi"
+	"google.golang.org/protobuf/proto"
+)
+
+type writeParams struct {
+	cfg       *Config
+	targetCfg *TargetConfig
+	batchSize uint
+	dataPath  string
+	async     bool
+}
+
+func RunWrite(args ...string) error {
+	fs := flag.NewFlagSet("write", flag.ExitOnError)
+	cfgFile := paramConfig(fs)
+	targetCfgFile := paramTargets(fs)
+	dataPath := paramFromPath(fs)
+	batchSize := paramBatchSize(fs)
+	async := paramAsync(fs)
+
+	if err := fs.Parse(args); err != nil {
+		fmt.Println("command line parse error", "err", err)
+		return err
+	}
+
+	if fs.NArg() > 0 || *targetCfgFile == "" || len(*dataPath) == 0 {
+		fs.PrintDefaults()
+		return nil
+	}
+
+	cfg, err := LoadConfigFile(*cfgFile)
+	if err != nil {
+		slog.Error("LoadConfigFile error", "err", err)
+		return err
+	}
+
+	targetCfg, err := LoadTargetFile(*targetCfgFile)
+	if err != nil {
+		fmt.Println("LoadConfigFile error", "err", err)
+		return err
+	}
+
+	writeParams := writeParams{
+		cfg:       cfg,
+		targetCfg: targetCfg,
+		batchSize: *batchSize,
+		dataPath:  *dataPath,
+		async:     *async,
+	}
+
+	return doWrite(writeParams)
+}
+
+func doWrite(writeParams writeParams) error {
+	fsInfo, err := os.Stat(writeParams.dataPath)
+	if err != nil {
+		return err
+	}
+
+	if !fsInfo.IsDir() {
+		// 如果是文件，直接写入
+		return doLoadFileToWrite(writeParams)
+	}
+
+	// 读取目录下信息
+	dirEntry, err := os.ReadDir(writeParams.dataPath)
+	if err != nil {
+		return err
+	}
+
+	// 遍历目录
+	for _, dir := range dirEntry {
+		newParam := writeParams
+		newParam.dataPath = path.Join(writeParams.dataPath, dir.Name())
+
+		if err = doWrite(newParam); err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+func doLoadFileToWrite(writeParams writeParams) error {
+	// 读取文件并按行遍历，以\t分割为两列，第一列为userid，第二列解析为string数组
+	file, err := os.Open(writeParams.dataPath)
+	if err != nil {
+		return err
+	}
+	defer file.Close()
+
+	scaner := bufio.NewScanner(file)
+
+	saasWriteCmds := []*saasapi.WriteCmd{}
+	saasReq := &saasapi.SaasReq{
+		UseridType: saasapi.UserIdType_DEVICEID,
+		Cmd: &saasapi.SaasReq_Write{
+			Write: &saasapi.Write{},
+		},
+	}
+
+	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,
+			IsFullOverwrite: true,
+		}
+		if len(userid) == 0 || len(targets) == 0 {
+			continue
+		}
+		for _, target := range targets {
+			if targetinfo, ok := writeParams.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
+				}
+			}
+		}
+
+		saasWriteCmds = append(saasWriteCmds, saasWriteCmd)
+
+		if len(saasWriteCmds) == int(writeParams.batchSize) {
+			if err = submitWrite(saasReq, saasWriteCmds); err != nil {
+				return err
+			}
+
+			saasWriteCmds = saasWriteCmds[:0]
+		}
+
+	}
+
+	if len(saasWriteCmds) > 0 {
+		return submitWrite(saasReq, saasWriteCmds)
+	}
+
+	return nil
+}
+
+func submitWrite(saasReq *saasapi.SaasReq, saasWriteCmds []*saasapi.WriteCmd) error {
+	saasReq.Cmd.(*saasapi.SaasReq_Write).Write.WriteCmds = saasWriteCmds
+
+	postBuf, err := proto.Marshal(saasReq)
+	if err != nil {
+		return err
+	}
+	fmt.Println(len(postBuf))
+
+	return nil
+}