Skip to content

Commit

Permalink
feat: add data-assert tool for Redis data management. (#1204)
Browse files Browse the repository at this point in the history
* docs

Signed-off-by: drivebyer <wuyangmuc@gmail.com>

* feat: add data-assert tool for Redis data management

- Introduced a new command-line tool `data-assert` to manage data in Redis.
- Implemented commands to generate data, check data consistency, and generate resource YAML files.
- Added `go.mod` and `go.sum` for dependency management, including Redis and Cobra libraries.
- Created a template for generating Kubernetes resources, including a ConfigMap and Pod specifications.
- Updated documentation to include information about the new data-assert functionality.

This addition enhances the ability to manage and verify data in Redis clusters effectively.

Signed-off-by: drivebyer <wuyangmuc@gmail.com>

---------

Signed-off-by: drivebyer <wuyangmuc@gmail.com>
  • Loading branch information
drivebyer authored Jan 12, 2025
1 parent 1a5a97f commit e7e6688
Show file tree
Hide file tree
Showing 6 changed files with 574 additions and 13 deletions.
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

0 comments on commit e7e6688

Please sign in to comment.