Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: add data-assert tool for Redis data management #1204

Merged
merged 2 commits into from
Jan 12, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 15 additions & 0 deletions cmd/data-assert/go.mod
Original file line number Diff line number Diff line change
@@ -0,0 +1,15 @@
module data-assert

go 1.23.4

require (
github.com/redis/go-redis/v9 v9.7.0
github.com/spf13/cobra v1.8.1
)

require (
github.com/cespare/xxhash/v2 v2.2.0 // indirect
github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f // indirect
github.com/inconshreveable/mousetrap v1.1.0 // indirect
github.com/spf13/pflag v1.0.5 // indirect
)
20 changes: 20 additions & 0 deletions cmd/data-assert/go.sum
Original file line number Diff line number Diff line change
@@ -0,0 +1,20 @@
github.com/bsm/ginkgo/v2 v2.12.0 h1:Ny8MWAHyOepLGlLKYmXG4IEkioBysk6GpaRTLC8zwWs=
github.com/bsm/ginkgo/v2 v2.12.0/go.mod h1:SwYbGRRDovPVboqFv0tPTcG1sN61LM1Z4ARdbAV9g4c=
github.com/bsm/gomega v1.27.10 h1:yeMWxP2pV2fG3FgAODIY8EiRE3dy0aeFYt4l7wh6yKA=
github.com/bsm/gomega v1.27.10/go.mod h1:JyEr/xRbxbtgWNi8tIEVPUYZ5Dzef52k01W3YH0H+O0=
github.com/cespare/xxhash/v2 v2.2.0 h1:DC2CZ1Ep5Y4k3ZQ899DldepgrayRUGE6BBZ/cd9Cj44=
github.com/cespare/xxhash/v2 v2.2.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
github.com/cpuguy83/go-md2man/v2 v2.0.4/go.mod h1:tgQtvFlXSQOSOSIRvRPT7W67SCa46tRHOmNcaadrF8o=
github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f h1:lO4WD4F/rVNCu3HqELle0jiPLLBs70cWOduZpkS1E78=
github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f/go.mod h1:cuUVRXasLTGF7a8hSLbxyZXjz+1KgoB3wDUb6vlszIc=
github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8=
github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw=
github.com/redis/go-redis/v9 v9.7.0 h1:HhLSs+B6O021gwzl+locl0zEDnyNkxMtf/Z3NNBMa9E=
github.com/redis/go-redis/v9 v9.7.0/go.mod h1:f6zhXITC7JUJIlPEiBOTXxJgPLdZcA93GewI7inzyWw=
github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM=
github.com/spf13/cobra v1.8.1 h1:e5/vxKd/rZsfSJMUX1agtjeTDf+qv1/JdBF8gg5k9ZM=
github.com/spf13/cobra v1.8.1/go.mod h1:wHxEcudfqmLYa8iTfL+OuZPbBZkmvliBWKIezN3kD9Y=
github.com/spf13/pflag v1.0.5 h1:iy+VFUOCP1a+8yFto/drg2CJ5u0yRoB7fZw3DKv/JXA=
github.com/spf13/pflag v1.0.5/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg=
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
201 changes: 201 additions & 0 deletions cmd/data-assert/main.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,201 @@
package main

import (
"context"
"fmt"
"os"
"strings"
"text/template"

"github.com/redis/go-redis/v9"
"github.com/spf13/cobra"
)

// go run main.go gen-resource-yaml
// go run main.go gen-redis-data --host redis-cluster-0.redis-cluster.default.svc.cluster.local --password 123456 --mode cluster/sentinel
// go run main.go chk-redis-data --host redis-cluster-0.redis-cluster.default.svc.cluster.local --password 123456 --mode cluster/sentinel

const (
hostFlag = "host"
passFlag = "password"
modeFlag = "mode"
totalKey = 1000
)

var (
host string
pass string
mode string
)

func main() {
rootCmd := &cobra.Command{
Use: "data-assert",
}
rootCmd.AddCommand(&cobra.Command{
Use: "gen-resource-yaml",
Run: genResourceYamlCmd,
})
rootCmd.AddCommand(&cobra.Command{
Use: "gen-redis-data",
Run: printFlags(genRedisDataCmd),
})
rootCmd.AddCommand(&cobra.Command{
Use: "chk-redis-data",
Run: printFlags(chkRedisDataCmd),
})

// add flags
rootCmd.PersistentFlags().StringVarP(&host, hostFlag, "H", "", "redis host")
rootCmd.PersistentFlags().StringVarP(&pass, passFlag, "P", "", "redis password")
rootCmd.PersistentFlags().StringVarP(&mode, modeFlag, "M", "", "redis mode")

rootCmd.Execute()
}

type cmdWrapperFunc func(cmd *cobra.Command, args []string)

// printFlags print flags
func printFlags(cmdWrapperFunc cmdWrapperFunc) cmdWrapperFunc {
return func(cmd *cobra.Command, args []string) {
fmt.Printf("host: %s, password: %s, mode: %s\n", host, pass, mode)
cmdWrapperFunc(cmd, args)
}
}

func genRedisDataCmd(cmd *cobra.Command, args []string) {
ctx := context.Background()
var rdb redis.UniversalClient

switch mode {
case "cluster":
rdb = redis.NewClusterClient(&redis.ClusterOptions{
Addrs: []string{host},
Password: pass,
})
case "sentinel":
rdb = redis.NewFailoverClient(&redis.FailoverOptions{
MasterName: "mymaster",
SentinelAddrs: []string{host},
Password: pass,
})
default:
fmt.Printf("unsupported redis mode: %s\n", mode)
return
}
defer rdb.Close()

// Generate and write data
for i := 0; i < totalKey; i++ {
key := fmt.Sprintf("key-%d", i)
value := fmt.Sprintf("value-%d", i)
err := rdb.Set(ctx, key, value, 0).Err()
if err != nil {
fmt.Printf("failed to set key %s: %v\n", key, err)
return
}
}
fmt.Printf("successfully generated %d keys\n", totalKey)
}

// DataError represents data consistency check errors
type DataError struct {
ExpectedCount int // Expected number of keys
ActualCount int // Actual number of keys found
}

func (e *DataError) Error() string {
return fmt.Sprintf("\nData count mismatch:\n - Expected: %d keys\n - Actual: %d keys\n - Missing: %d keys",
e.ExpectedCount, e.ActualCount, e.ExpectedCount-e.ActualCount)
}

func chkRedisDataCmd(cmd *cobra.Command, args []string) {
if err := checkRedisData(); err != nil {
if dataErr, ok := err.(*DataError); ok {
fmt.Printf("Data consistency check failed: %s\n", dataErr.Error())
os.Exit(1)
}
fmt.Printf("Error occurred during check: %v\n", err)
os.Exit(1)
}
fmt.Printf("Data consistency check passed! All %d keys exist\n", totalKey)
}

func checkRedisData() error {
ctx := context.Background()
var rdb redis.UniversalClient

switch mode {
case "cluster":
rdb = redis.NewClusterClient(&redis.ClusterOptions{
Addrs: []string{host},
Password: pass,
})
case "sentinel":
rdb = redis.NewFailoverClient(&redis.FailoverOptions{
MasterName: "mymaster",
SentinelAddrs: []string{host},
Password: pass,
})
default:
return fmt.Errorf("unsupported redis mode: %s", mode)
}
defer rdb.Close()

// Count existing keys
actualCount := 0
for i := 0; i < totalKey; i++ {
key := fmt.Sprintf("key-%d", i)
exists, err := rdb.Exists(ctx, key).Result()
if err != nil {
return fmt.Errorf("error checking key %s: %w", key, err)
}
if exists == 1 {
actualCount++
}
}

if actualCount != totalKey {
return &DataError{
ExpectedCount: totalKey,
ActualCount: actualCount,
}
}
return nil
}

func genResourceYamlCmd(cmd *cobra.Command, args []string) {
mainGoBytes, err := os.ReadFile("main.go")
if err != nil {
panic(err)
}
indentedMain := " " + strings.Join(strings.Split(string(mainGoBytes), "\n"), "\n ")

goModBytes, err := os.ReadFile("go.mod")
if err != nil {
panic(err)
}
goModContent := " " + strings.Join(strings.Split(string(goModBytes), "\n"), "\n ")

goSumBytes, err := os.ReadFile("go.sum")
if err != nil {
panic(err)
}
goSumContent := " " + strings.Join(strings.Split(string(goSumBytes), "\n"), "\n ")

outFile, err := os.Create("resources.yaml")
if err != nil {
panic(err)
}
defer outFile.Close()

err = template.Must(template.ParseFiles("resources.yaml.tmpl")).Execute(outFile, map[string]string{
"Main": indentedMain,
"GoMod": goModContent,
"GoSum": goSumContent,
"Notice": "## DO NOT MODIFY THIS FILE!!!, IT IS GENERATED FROM resources.yaml.tmpl",
})
if err != nil {
panic(err)
}
}
Loading
Loading