工具基本实现write功能

This commit is contained in:
2025-04-02 19:34:36 +08:00
parent e59ff65c0d
commit 0344c09ce7
20 changed files with 918 additions and 189 deletions

View File

@@ -1,2 +1,3 @@
debug/
saastool
saastool
saastool_linux

View File

@@ -0,0 +1,5 @@
package main
func RunColumnWrite(args ...string) error {
return nil
}

View File

@@ -1,25 +1,14 @@
package main
import (
"e.coding.net/rta/public/saasapi/pkg/saashttp"
"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
Auth saashttp.Auth
ApiUrls saashttp.ApiUrls
}
// LoadConfigFile 加载配置文件

189
cmd/saastool/convert.go Normal file
View File

@@ -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
}

View File

@@ -16,8 +16,13 @@ Usage: [[command] [arguments]]
The commands are:
write Write user's bytes / uint32s / flags
read Read user's bytes / uint32s / flags
write Write user's 'bytes / uint32s / flags'
read Read user's 'bytes / uint32s / flags'
columnwrite Write columns for 'deviceid / openid' users
tasklist List tasks
taskcancel Cancel task
taskdetail Show task detail
"help" is the default command.

View File

@@ -23,10 +23,22 @@ func Run(args ...string) error {
return RunHelp(args...)
case "write":
return RunWrite(args...)
//case "read":
// return RunRead(args...)
case "read":
return RunRead(args...)
case "columnwrite":
return RunColumnWrite(args...)
case "convert":
return RunConvert(args...)
case "tasklist":
return RunTaskList(args...)
case "taskcancel":
return RunTaskCancel(args...)
case "taskdetail":
return RunTaskDetail(args...)
case "verify":
return RunVerify(args...)
default:
err := fmt.Errorf(`unknown command "%s"`+"\n"+`Run 'dmptool help' for usage`, name)
err := fmt.Errorf(`unknown command "%s"`+"\n"+`Run 'saastool help' for usage`, name)
slog.Warn(err.Error())
return err
}

View File

@@ -11,8 +11,16 @@ 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 paramSourcePath(fs *flag.FlagSet) *string {
return fs.String("source", "", "Data path source for write command.")
}
func paramDestPath(fs *flag.FlagSet) *string {
return fs.String("dest", "", "Data path destination for write command.")
}
func paramAppid(fs *flag.FlagSet) *string {
return fs.String("appid", "", "Wechat appid")
}
func paramBatchSize(fs *flag.FlagSet) *uint {
@@ -22,3 +30,7 @@ func paramBatchSize(fs *flag.FlagSet) *uint {
func paramAsync(fs *flag.FlagSet) *bool {
return fs.Bool("async", false, "Async mode")
}
func paramClear(fs *flag.FlagSet) *bool {
return fs.Bool("clear", false, "Clear all data before write")
}

View File

@@ -1 +1,5 @@
package main
func RunRead(args ...string) error {
return nil
}

View File

@@ -0,0 +1,5 @@
package main
func RunTaskCancel(args ...string) error {
return nil
}

View File

@@ -0,0 +1,5 @@
package main
func RunTaskDetail(args ...string) error {
return nil
}

View File

@@ -0,0 +1,5 @@
package main
func RunTaskList(args ...string) error {
return nil
}

25
cmd/saastool/term.go Normal file
View File

@@ -0,0 +1,25 @@
package main
import (
"os"
"syscall"
"unsafe"
)
// https://man7.org/linux/man-pages/man2/TIOCSWINSZ.2const.html
// winSize console窗口大小
type winSize struct {
wsRow uint16
wsCols uint16
wsXPixels uint16
wxYPixels uint16
}
// getConsoleSize 获取控制台窗口大小
func getConsoleSize() (cols, rows int) {
var sz winSize
_, _, _ = syscall.Syscall(syscall.SYS_IOCTL,
os.Stdout.Fd(), uintptr(syscall.TIOCGWINSZ), uintptr(unsafe.Pointer(&sz)))
return int(sz.wsCols), int(sz.wsRow)
}

5
cmd/saastool/verify.go Normal file
View File

@@ -0,0 +1,5 @@
package main
func RunVerify(args ...string) error {
return nil
}

View File

@@ -5,36 +5,40 @@ import (
"flag"
"fmt"
"log/slog"
"net/http"
"os"
"path"
"strings"
"e.coding.net/rta/public/saasapi"
"google.golang.org/protobuf/proto"
"e.coding.net/rta/public/saasapi/pkg/saashttp"
"google.golang.org/protobuf/encoding/protojson"
)
type writeParams struct {
cfg *Config
targetCfg *TargetConfig
batchSize uint
dataPath string
async bool
cfg *Config
sourcePath string
appid string
batchSize uint
async bool
clear bool
saasHttp *saashttp.SaasClient
}
func RunWrite(args ...string) error {
fs := flag.NewFlagSet("write", flag.ExitOnError)
cfgFile := paramConfig(fs)
targetCfgFile := paramTargets(fs)
dataPath := paramFromPath(fs)
sourcePath := paramSourcePath(fs)
appid := paramAppid(fs)
batchSize := paramBatchSize(fs)
async := paramAsync(fs)
clear := paramClear(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 {
if fs.NArg() > 0 || len(*sourcePath) == 0 {
fs.PrintDefaults()
return nil
}
@@ -44,26 +48,25 @@ func RunWrite(args ...string) error {
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,
cfg: cfg,
sourcePath: *sourcePath,
appid: *appid,
batchSize: *batchSize,
async: *async,
clear: *clear,
saasHttp: &saashttp.SaasClient{
Client: http.Client{},
ApiUrls: cfg.ApiUrls,
Auth: cfg.Auth,
},
}
return doWrite(writeParams)
}
func doWrite(writeParams writeParams) error {
fsInfo, err := os.Stat(writeParams.dataPath)
fsInfo, err := os.Stat(writeParams.sourcePath)
if err != nil {
return err
}
@@ -74,7 +77,7 @@ func doWrite(writeParams writeParams) error {
}
// 读取目录下信息
dirEntry, err := os.ReadDir(writeParams.dataPath)
dirEntry, err := os.ReadDir(writeParams.sourcePath)
if err != nil {
return err
}
@@ -82,7 +85,7 @@ func doWrite(writeParams writeParams) error {
// 遍历目录
for _, dir := range dirEntry {
newParam := writeParams
newParam.dataPath = path.Join(writeParams.dataPath, dir.Name())
newParam.sourcePath = path.Join(writeParams.sourcePath, dir.Name())
if err = doWrite(newParam); err != nil {
return err
@@ -94,7 +97,7 @@ func doWrite(writeParams writeParams) error {
func doLoadFileToWrite(writeParams writeParams) error {
// 读取文件并按行遍历,以\t分割为两列第一列为userid第二列解析为string数组
file, err := os.Open(writeParams.dataPath)
file, err := os.Open(writeParams.sourcePath)
if err != nil {
return err
}
@@ -103,103 +106,66 @@ func doLoadFileToWrite(writeParams writeParams) error {
scaner := bufio.NewScanner(file)
saasWriteCmds := []*saasapi.WriteCmd{}
saasReq := &saasapi.SaasReq{
UseridType: saasapi.UserIdType_DEVICEID,
Cmd: &saasapi.SaasReq_Write{
Write: &saasapi.Write{},
},
}
succ := uint32(0)
succTotal := uint32(0)
total := uint32(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,
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
}
}
saasWriteCmd := &saasapi.WriteCmd{}
if err = protojson.Unmarshal([]byte(line), saasWriteCmd); err != nil {
return err
}
saasWriteCmds = append(saasWriteCmds, saasWriteCmd)
total++
if len(saasWriteCmds) == int(writeParams.batchSize) {
if err = submitWrite(saasReq, saasWriteCmds); err != nil {
if succ, _, err = submitWrite(writeParams, saasWriteCmds); err != nil {
return err
}
succTotal += succ
fmt.Printf("[%v] batch_succ = %v, succ_total = %v, total_processed = %v\n", writeParams.sourcePath, succ, succTotal, total)
saasWriteCmds = saasWriteCmds[:0]
}
}
if len(saasWriteCmds) > 0 {
return submitWrite(saasReq, saasWriteCmds)
if succ, _, err = submitWrite(writeParams, saasWriteCmds); err != nil {
return err
}
succTotal += succ
fmt.Printf("[%v] batch_succ = %v, succ_total = %v, total_processed = %v\n", writeParams.sourcePath, succ, succTotal, total)
}
return nil
}
func submitWrite(saasReq *saasapi.SaasReq, saasWriteCmds []*saasapi.WriteCmd) error {
func submitWrite(writeParams writeParams, saasWriteCmds []*saasapi.WriteCmd) (succ, total uint32, err error) {
saasReq := &saasapi.SaasReq{
Cmd: &saasapi.SaasReq_Write{
Write: &saasapi.Write{
IsClearAllFirst: writeParams.clear,
Async: writeParams.async,
},
},
}
if writeParams.appid != "" {
saasReq.UseridType = saasapi.UserIdType_OPENID
saasReq.Appid = writeParams.appid
}
saasReq.Cmd.(*saasapi.SaasReq_Write).Write.WriteCmds = saasWriteCmds
postBuf, err := proto.Marshal(saasReq)
if err != nil {
return err
}
fmt.Println(len(postBuf))
total = uint32(len(saasWriteCmds))
succ, err = writeParams.saasHttp.Write(saasReq)
return nil
return
}