Merge pull request #36 from kuzxnia/config-generation

added config generator
kuzxnia · Mar 24, 2024 · 7bbfb33 · 7bbfb33
2 parents 0ca972f + 086081c
commit 7bbfb33
Show file tree

Hide file tree

Showing 7 changed files with 298 additions and 87 deletions.
diff --git a/cli/cli.go b/cli/cli.go
@@ -25,6 +25,7 @@ func New(version string, commit string, date string) *cobra.Command {
 	cmd.AddCommand(provideWorkloadCommands()...)
 	cmd.AddGroup(&WorkloadGroup)
 	cmd.AddCommand(provideOrchiestrationCommands()...)
+	cmd.AddGroup(&OrchiestrationGroup)
 	cmd.Root().CompletionOptions.HiddenDefaultCmd = true
 
 	return &cmd
@@ -43,11 +44,12 @@ var WorkloadGroup = cobra.Group{
 const (
 	WorkloadRootCommand = "workload"
 
-	CommandStartWorkload    = "start"
-	CommandStopWorkload     = "stop"
-	CommandWatchWorkload    = "watch"
-	CommandProgressWorkload = "progress"
-	CommandConfigWorkload   = "config"
+	CommandStartWorkload          = "start"
+	CommandStopWorkload           = "stop"
+	CommandWatchWorkload          = "watch"
+	CommandProgressWorkload       = "progress"
+	CommandConfigWorkload         = "config"
+	CommandGenerateConfigWorkload = "generate-config"
 
 	// config args
 	ConfigFile = "config-file"
@@ -166,7 +168,7 @@ func provideWorkloadCommands() []*cobra.Command {
 
 	configCommand := cobra.Command{
 		Use:               CommandConfigWorkload,
-		Short:             "Config",
+		Short:             "Get or set workload config",
 		GroupID:           WorkloadGroup.ID,
 		PersistentPreRunE: persistentPreRunE,
 		PersistentPostRun: persistentPostRun,
@@ -176,23 +178,31 @@ func provideWorkloadCommands() []*cobra.Command {
 			stdin, _ := flags.GetBool(StdIn)
 
 			if configFile == "" && stdin == false {
-				return workload.GetConfigWorkload(Conn)
+				return workload.GetWorkloadConfig(Conn)
 			}
 
 			config, err := ParseConfigFile(configFile, stdin)
 			if err != nil {
 				return err
 			}
 
-			return workload.SetConfigWorkload(Conn, config)
+			return workload.SetWorkloadConfig(Conn, config)
 		},
 	}
 	configCommandFlags := configCommand.Flags()
 	configCommandFlags.StringP(ConfigFile, "f", "", "file with workload configuration")
 	configCommandFlags.Bool(StdIn, false, "get workload configuration from stdin")
 	configCommandFlags.StringP(AgentUri, "u", "127.0.0.1:1234", "loadbot agent uri (default: 127.0.0.1:1234)")
 
-	return []*cobra.Command{&startCommand, &stopCommand, &configCommand, &progressCommand}
+	generateConfigCommand := cobra.Command{
+		Use:   CommandGenerateConfigWorkload,
+		Short: "Generate sample workload config",
+		RunE: func(cmd *cobra.Command, args []string) (err error) {
+      return workload.GenerateConfigWorkload()
+		},
+	}
+
+	return []*cobra.Command{&startCommand, &stopCommand, &configCommand, &generateConfigCommand, &progressCommand}
 }
 
 var AgentGroup = cobra.Group{
@@ -277,12 +287,18 @@ const (
 	FlagHelmSetFile   = "helm-set-file"
 )
 
+var OrchiestrationGroup = cobra.Group{
+	ID:    "orchiestrator",
+	Title: "Orchiestration Commands:",
+}
+
 func provideOrchiestrationCommands() []*cobra.Command {
 	installationCommand := cobra.Command{
 		Use:     CommandInstall + " <name>",
 		Aliases: []string{"i"},
 		Short:   "Install workload driver with helm charts on k8s or only with docker locally",
 		Args:    cobra.ExactArgs(1),
+		GroupID: OrchiestrationGroup.ID,
 		RunE: func(cmd *cobra.Command, args []string) (err error) {
 			flags := cmd.Flags()
 
@@ -332,9 +348,10 @@ func provideOrchiestrationCommands() []*cobra.Command {
 	flags.StringP(FlagWorkloadConfig, "f", "", "set additional Helm values by a YAML file or a URL (can specify multiple)")
 
 	upgradeCommand := cobra.Command{
-		Use:   CommandUpgrade + " <name>",
-		Short: "Upgrade workload driver with helm charts on k8s or only with docker locally",
-		Args:  cobra.ExactArgs(1),
+		Use:     CommandUpgrade + " <name>",
+		Short:   "Upgrade workload driver with helm charts on k8s or only with docker locally",
+		Args:    cobra.ExactArgs(1),
+		GroupID: OrchiestrationGroup.ID,
 		RunE: func(cmd *cobra.Command, args []string) (err error) {
 			flags := cmd.Flags()
 
@@ -384,9 +401,10 @@ func provideOrchiestrationCommands() []*cobra.Command {
 
 	unInstallationCommand := cobra.Command{
 		// todo: where to keep configuration? there will be couple workloads at the same time
-		Use:   CommandUnInstall,
-		Short: "Uninstall workload driver",
-		Args:  cobra.ExactArgs(1),
+		Use:     CommandUnInstall,
+		Short:   "Uninstall workload driver",
+		Args:    cobra.ExactArgs(1),
+		GroupID: OrchiestrationGroup.ID,
 		RunE: func(cmd *cobra.Command, args []string) (err error) {
 			flags := cmd.Flags()
 
@@ -419,8 +437,9 @@ func provideOrchiestrationCommands() []*cobra.Command {
 
 	listCommand := cobra.Command{
 		// todo: where to keep configuration? there will be couple workloads at the same time
-		Use:   CommandList,
-		Short: "List workloads",
+		Use:     CommandList,
+		Short:   "List workloads",
+		GroupID: OrchiestrationGroup.ID,
 		RunE: func(cmd *cobra.Command, args []string) (err error) {
 			flags := cmd.Flags()
 

diff --git a/cli/workload/config.go b/cli/workload/config.go
@@ -8,9 +8,19 @@ package workload
 
 import (
 	"context"
+	"encoding/json"
+	"errors"
 	"fmt"
+	"io/ioutil"
+	"os"
+	"strconv"
+	"time"
 
+	"github.com/cqroot/prompt"
+	"github.com/cqroot/prompt/choose"
+	"github.com/cqroot/prompt/input"
 	"github.com/kuzxnia/loadbot/lbot"
+	"github.com/kuzxnia/loadbot/lbot/config"
 	"github.com/kuzxnia/loadbot/lbot/proto"
 	"google.golang.org/grpc"
 	"google.golang.org/protobuf/encoding/prototext"
@@ -19,7 +29,7 @@ import (
 
 // checks if process is running in local system
 // here should be cli config request - not lbot one
-func SetConfigWorkload(conn grpc.ClientConnInterface, parsedConfig *lbot.ConfigRequest) (err error) {
+func SetWorkloadConfig(conn grpc.ClientConnInterface, parsedConfig *lbot.ConfigRequest) (err error) {
 	requestConfig := BuildConfigRequest(parsedConfig)
 
 	fmt.Println("🚀 Setting new config")
@@ -34,7 +44,7 @@ func SetConfigWorkload(conn grpc.ClientConnInterface, parsedConfig *lbot.ConfigR
 	return
 }
 
-func GetConfigWorkload(conn grpc.ClientConnInterface) (err error) {
+func GetWorkloadConfig(conn grpc.ClientConnInterface) (err error) {
 	client := proto.NewConfigServiceClient(conn)
 	cfg, err := client.GetConfig(context.TODO(), &emptypb.Empty{})
 	if err != nil {
@@ -46,6 +56,143 @@ func GetConfigWorkload(conn grpc.ClientConnInterface) (err error) {
 	return
 }
 
+func GenerateConfigWorkload() (err error) {
+	cfg := lbot.ConfigRequest{}
+	cfg.ConnectionString = GetStringInput("Please provide connection string to database:", "mongodb://admin:pass@127.0.0.1:27017")
+
+	metrics := GetSelectWithDescriptionInput(
+		"Do you want to export metrics?",
+		[]choose.Choice{
+			{Text: "no", Note: "Metrics won't be exported"},
+			{Text: "on port", Note: "Export metrics on provided port"},
+			{Text: "push", Note: "Push metrics to prometheus url with interval"},
+		},
+	)
+	cfg.Agent = &lbot.AgentRequest{}
+	switch metrics {
+	case "no": // just skip
+	case "on port":
+		cfg.Agent.Name = GetStringInput("Provide agent name(used for metrics labeling):", "")
+		cfg.Agent.MetricsExportPort = GetStringInput("Provide metrics export port:", "6060")
+	case "push":
+		cfg.Agent.Name = GetStringInput("Provide agent name(used for metrics labeling):", "")
+		cfg.Agent.MetricsExportUrl = GetStringInput("Provide metrics import url:", "http://victoria-metrics:8428/api/v1/import/prometheus")
+		cfg.Agent.MetricsExportIntervalSeconds = uint64(GetDurationInput("Provide metrics push interval:", "5s").Seconds())
+	}
+
+	job := lbot.JobRequest{}
+	cfg.Jobs = []*lbot.JobRequest{&job}
+	job.Type = GetSelectInput(
+		"Choose job type",
+		[]string{string(config.Write), string(config.BulkWrite), string(config.Read), string(config.Update)},
+		choose.WithTheme(choose.ThemeLine),
+		choose.WithKeyMap(choose.HorizontalKeyMap),
+	)
+	if job.Type == string(config.BulkWrite) {
+		job.BatchSize = uint64(GetNumberInput("Provide batch size:", "100"))
+	}
+	job.Name = GetStringInput("Provide workload name:", "Test workload")
+	job.Database = GetStringInput("Provide database name:", "tmp-database")
+	job.Collection = GetStringInput("Provide collection name:", "tmp-collection")
+	job.Pace = uint64(GetNumberInput("How many requests per second?", "100"))
+	job.Connections = uint64(GetNumberInput("Provide workload concurency:", "100"))
+	job.DataSize = uint64(GetNumberInput("How big document should be? (whole document will be <data_size> + 33 in bytes):", "250"))
+
+	limitType := GetSelectWithDescriptionInput(
+		"Choose workload limit type:",
+		[]choose.Choice{
+			{Text: "operations", Note: "Executes number of operations"},
+			{Text: "duration", Note: "Runs workload until time ends"},
+			{Text: "infinite", Note: "Runs workload endlessly"},
+		},
+	)
+	switch limitType {
+	case "operations":
+		job.Operations = uint64(GetNumberInput("How many operations to perform?", "100"))
+	case "duration":
+		job.Duration = *GetDurationInput("How long workload should take?", "10m30s")
+	default:
+		// set noting to duration or operations
+	}
+
+	filename := GetStringInput("Where to save config file:", "workload_config.json")
+	if _, error := os.Stat(filename); !os.IsNotExist(error) {
+		overwrite := GetBooleanInput("File exists! Do you want to overwrite?")
+		if !overwrite {
+			fmt.Println("Exiting...")
+			os.Exit(1)
+		}
+	}
+
+	data, err := json.MarshalIndent(cfg, "", "\t")
+	if err != nil {
+		return err
+	}
+	err = ioutil.WriteFile(filename, data, 0o644)
+	if err != nil {
+		return err
+	}
+
+	return
+}
+
+func CheckPromptErr(err error) {
+	if err != nil {
+		if errors.Is(err, prompt.ErrUserQuit) {
+			fmt.Fprintln(os.Stderr, "Error:", err)
+			os.Exit(1)
+		} else {
+			panic(err)
+		}
+	}
+}
+
+func GetStringInput(label string, defaultValue string) string {
+	result, err := prompt.New().Ask(label).Input(defaultValue, input.WithHelp(true))
+	CheckPromptErr(err)
+	return result
+}
+
+func GetBooleanInput(label string) bool {
+	result := GetSelectInput(
+		label,
+		[]string{"t", "F"},
+		choose.WithTheme(choose.ThemeLine),
+		choose.WithKeyMap(choose.HorizontalKeyMap),
+	)
+	resultBool, err := strconv.ParseBool(result)
+	CheckPromptErr(err)
+	return resultBool
+}
+
+func GetNumberInput(label string, defaultValue string) int {
+	result, err := prompt.New().Ask(label).Input(defaultValue, input.WithInputMode(input.InputInteger), input.WithHelp(true))
+	CheckPromptErr(err)
+	resultInt, err := strconv.Atoi(result)
+	CheckPromptErr(err)
+	return resultInt
+}
+
+func GetSelectInput(label string, options []string, ops ...choose.Option) string {
+	result, err := prompt.New().Ask(label).Choose(options, ops...)
+	CheckPromptErr(err)
+	return result
+}
+
+func GetSelectWithDescriptionInput(label string, options []choose.Choice, ops ...choose.Option) string {
+	result, err := prompt.New().Ask(label).AdvancedChoose(options, ops...)
+	CheckPromptErr(err)
+	return result
+}
+
+func GetDurationInput(label string, defaultValue string) *time.Duration {
+	result, err := prompt.New().Ask(label).Input(defaultValue, input.WithHelp(true))
+	CheckPromptErr(err)
+	resultDuration, err := time.ParseDuration(result)
+	CheckPromptErr(err)
+	return &resultDuration
+}
+
 func BuildConfigRequest(request *lbot.ConfigRequest) *proto.ConfigRequest {
 	cfg := &proto.ConfigRequest{
 		ConnectionString: request.ConnectionString,

diff --git a/go.mod b/go.mod
@@ -6,6 +6,7 @@ require (
 	github.com/VictoriaMetrics/metrics v1.31.0
 	github.com/benbjohnson/clock v1.3.0
 	github.com/cheggaaa/pb/v3 v3.1.5
+	github.com/cqroot/prompt v0.9.3
 	github.com/fsnotify/fsnotify v1.7.0
 	github.com/go-faker/faker/v4 v4.2.0
 	github.com/google/uuid v1.3.1
@@ -38,11 +39,18 @@ require (
 	github.com/Microsoft/hcsshim v0.11.4 // indirect
 	github.com/VividCortex/ewma v1.2.0 // indirect
 	github.com/asaskevich/govalidator v0.0.0-20200428143746-21a406dcc535 // indirect
+	github.com/atotto/clipboard v0.1.4 // indirect
+	github.com/aymanbagabas/go-osc52/v2 v2.0.1 // indirect
 	github.com/beorn7/perks v1.0.1 // indirect
 	github.com/cespare/xxhash/v2 v2.2.0 // indirect
 	github.com/chai2010/gettext-go v1.0.2 // indirect
+	github.com/charmbracelet/bubbles v0.16.1 // indirect
+	github.com/charmbracelet/bubbletea v0.24.2 // indirect
+	github.com/charmbracelet/lipgloss v0.9.1 // indirect
+	github.com/containerd/console v1.0.4-0.20230313162750-1ae8d489ac81 // indirect
 	github.com/containerd/containerd v1.7.12 // indirect
 	github.com/containerd/log v0.1.0 // indirect
+	github.com/cqroot/multichoose v0.1.1 // indirect
 	github.com/cyphar/filepath-securejoin v0.2.4 // indirect
 	github.com/davecgh/go-spew v1.1.1 // indirect
 	github.com/docker/cli v24.0.6+incompatible // indirect
@@ -90,9 +98,11 @@ require (
 	github.com/lann/ps v0.0.0-20150810152359-62de8c46ede0 // indirect
 	github.com/lib/pq v1.10.9 // indirect
 	github.com/liggitt/tabwriter v0.0.0-20181228230101-89fcab3d43de // indirect
+	github.com/lucasb-eyer/go-colorful v1.2.0 // indirect
 	github.com/mailru/easyjson v0.7.7 // indirect
 	github.com/mattn/go-colorable v0.1.13 // indirect
 	github.com/mattn/go-isatty v0.0.20 // indirect
+	github.com/mattn/go-localereader v0.0.1 // indirect
 	github.com/mattn/go-runewidth v0.0.15 // indirect
 	github.com/matttproud/golang_protobuf_extensions v1.0.4 // indirect
 	github.com/mitchellh/copystructure v1.2.0 // indirect
@@ -106,6 +116,10 @@ require (
 	github.com/monochromegane/go-gitignore v0.0.0-20200626010858-205db1a8cc00 // indirect
 	github.com/montanaflynn/stats v0.7.1 // indirect
 	github.com/morikuni/aec v1.0.0 // indirect
+	github.com/muesli/ansi v0.0.0-20211018074035-2e021307bc4b // indirect
+	github.com/muesli/cancelreader v0.2.2 // indirect
+	github.com/muesli/reflow v0.3.0 // indirect
+	github.com/muesli/termenv v0.15.2 // indirect
 	github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect
 	github.com/mxk/go-flowrate v0.0.0-20140419014527-cca7078d478f // indirect
 	github.com/opencontainers/go-digest v1.0.0 // indirect