From bae534f601042300bc8b3f93cb93b1d82b883502 Mon Sep 17 00:00:00 2001 From: Bhaskarjyoti Bora Date: Tue, 3 Sep 2024 20:18:44 +0530 Subject: [PATCH] drtprod: yaml config support This implementation of drtprod enables you to define the series of commands in a YAML file and execute the same. This way you can define a sequence of commands per cluster. Fixes: #125381 Epic: None --- pkg/BUILD.bazel | 7 + pkg/cmd/dev/build.go | 1 + pkg/cmd/drtprod/BUILD.bazel | 15 + pkg/cmd/drtprod/README.md | 8 + pkg/cmd/drtprod/cli/BUILD.bazel | 17 + pkg/cmd/drtprod/cli/commands/BUILD.bazel | 24 ++ pkg/cmd/drtprod/cli/commands/rootcmd.go | 34 ++ pkg/cmd/drtprod/cli/commands/yamlprocessor.go | 311 ++++++++++++++++++ .../cli/commands/yamlprocessor_test.go | 157 +++++++++ pkg/cmd/drtprod/cli/handlers.go | 57 ++++ pkg/cmd/drtprod/cli/registry.go | 26 ++ pkg/cmd/drtprod/configs/drt_test.yaml | 39 +++ pkg/cmd/drtprod/helpers/BUILD.bazel | 12 + pkg/cmd/drtprod/helpers/utils.go | 91 +++++ pkg/cmd/drtprod/main.go | 22 ++ 15 files changed, 821 insertions(+) create mode 100644 pkg/cmd/drtprod/BUILD.bazel create mode 100644 pkg/cmd/drtprod/README.md create mode 100644 pkg/cmd/drtprod/cli/BUILD.bazel create mode 100644 pkg/cmd/drtprod/cli/commands/BUILD.bazel create mode 100644 pkg/cmd/drtprod/cli/commands/rootcmd.go create mode 100644 pkg/cmd/drtprod/cli/commands/yamlprocessor.go create mode 100644 pkg/cmd/drtprod/cli/commands/yamlprocessor_test.go create mode 100644 pkg/cmd/drtprod/cli/handlers.go create mode 100644 pkg/cmd/drtprod/cli/registry.go create mode 100644 pkg/cmd/drtprod/configs/drt_test.yaml create mode 100644 pkg/cmd/drtprod/helpers/BUILD.bazel create mode 100644 pkg/cmd/drtprod/helpers/utils.go create mode 100644 pkg/cmd/drtprod/main.go diff --git a/pkg/BUILD.bazel b/pkg/BUILD.bazel index ab6917f9bcdf..8686b3769bd3 100644 --- a/pkg/BUILD.bazel +++ b/pkg/BUILD.bazel @@ -143,6 +143,7 @@ ALL_TESTS = [ "//pkg/cmd/dev:dev_test", "//pkg/cmd/docgen/extract:extract_test", "//pkg/cmd/docs-issue-generation:docs-issue-generation_test", + "//pkg/cmd/drtprod/cli/commands:commands_test", "//pkg/cmd/github-pull-request-make:github-pull-request-make_test", "//pkg/cmd/label-merged-pr:label-merged-pr_test", "//pkg/cmd/mirror/go:go_test", @@ -1122,6 +1123,12 @@ GO_TARGETS = [ "//pkg/cmd/docs-issue-generation:docs-issue-generation_test", "//pkg/cmd/drt-run:drt-run", "//pkg/cmd/drt-run:drt_run_lib", + "//pkg/cmd/drtprod/cli/commands:commands", + "//pkg/cmd/drtprod/cli/commands:commands_test", + "//pkg/cmd/drtprod/cli:cli", + "//pkg/cmd/drtprod/helpers:helpers", + "//pkg/cmd/drtprod:drtprod", + "//pkg/cmd/drtprod:drtprod_lib", "//pkg/cmd/fuzz:fuzz", "//pkg/cmd/fuzz:fuzz_lib", "//pkg/cmd/generate-acceptance-tests:generate-acceptance-tests", diff --git a/pkg/cmd/dev/build.go b/pkg/cmd/dev/build.go index 3931c23b802d..6574a2468af3 100644 --- a/pkg/cmd/dev/build.go +++ b/pkg/cmd/dev/build.go @@ -101,6 +101,7 @@ var buildTargetMapping = map[string]string{ "optfmt": "//pkg/sql/opt/optgen/cmd/optfmt:optfmt", "oss": cockroachTargetOss, "reduce": "//pkg/cmd/reduce:reduce", + "drtprod": "//pkg/cmd/drtprod:drtprod", "roachprod": "//pkg/cmd/roachprod:roachprod", "roachprod-stress": "//pkg/cmd/roachprod-stress:roachprod-stress", "roachprod-microbench": "//pkg/cmd/roachprod-microbench:roachprod-microbench", diff --git a/pkg/cmd/drtprod/BUILD.bazel b/pkg/cmd/drtprod/BUILD.bazel new file mode 100644 index 000000000000..6299c2c2e98f --- /dev/null +++ b/pkg/cmd/drtprod/BUILD.bazel @@ -0,0 +1,15 @@ +load("@io_bazel_rules_go//go:def.bzl", "go_binary", "go_library") + +go_library( + name = "drtprod_lib", + srcs = ["main.go"], + importpath = "github.com/cockroachdb/cockroach/pkg/cmd/drtprod", + visibility = ["//visibility:private"], + deps = ["//pkg/cmd/drtprod/cli"], +) + +go_binary( + name = "drtprod", + embed = [":drtprod_lib"], + visibility = ["//visibility:public"], +) diff --git a/pkg/cmd/drtprod/README.md b/pkg/cmd/drtprod/README.md new file mode 100644 index 000000000000..20964139905a --- /dev/null +++ b/pkg/cmd/drtprod/README.md @@ -0,0 +1,8 @@ +# DRT Prod + +drtprod is a tool for manipulating drt clusters using roachprod, +allowing easy creating, destruction, controls and configurations of clusters. + +Commands include:
+execute: executes the commands in the YAML
+*: any other command is passed to roachprod, potentially with flags added diff --git a/pkg/cmd/drtprod/cli/BUILD.bazel b/pkg/cmd/drtprod/cli/BUILD.bazel new file mode 100644 index 000000000000..7c9ff2c0ae90 --- /dev/null +++ b/pkg/cmd/drtprod/cli/BUILD.bazel @@ -0,0 +1,17 @@ +load("@io_bazel_rules_go//go:def.bzl", "go_library") + +go_library( + name = "cli", + srcs = [ + "handlers.go", + "registry.go", + ], + importpath = "github.com/cockroachdb/cockroach/pkg/cmd/drtprod/cli", + visibility = ["//visibility:public"], + deps = [ + "//pkg/cmd/drtprod/cli/commands", + "//pkg/cmd/drtprod/helpers", + "//pkg/roachprod", + "@com_github_spf13_cobra//:cobra", + ], +) diff --git a/pkg/cmd/drtprod/cli/commands/BUILD.bazel b/pkg/cmd/drtprod/cli/commands/BUILD.bazel new file mode 100644 index 000000000000..1a4f6c534a32 --- /dev/null +++ b/pkg/cmd/drtprod/cli/commands/BUILD.bazel @@ -0,0 +1,24 @@ +load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test") + +go_library( + name = "commands", + srcs = [ + "rootcmd.go", + "yamlprocessor.go", + ], + importpath = "github.com/cockroachdb/cockroach/pkg/cmd/drtprod/cli/commands", + visibility = ["//visibility:public"], + deps = [ + "//pkg/build", + "//pkg/cmd/drtprod/helpers", + "@com_github_spf13_cobra//:cobra", + "@in_gopkg_yaml_v2//:yaml_v2", + ], +) + +go_test( + name = "commands_test", + srcs = ["yamlprocessor_test.go"], + embed = [":commands"], + deps = ["@com_github_stretchr_testify//require"], +) diff --git a/pkg/cmd/drtprod/cli/commands/rootcmd.go b/pkg/cmd/drtprod/cli/commands/rootcmd.go new file mode 100644 index 000000000000..dd86942bf9e2 --- /dev/null +++ b/pkg/cmd/drtprod/cli/commands/rootcmd.go @@ -0,0 +1,34 @@ +// Copyright 2024 The Cockroach Authors. +// +// Use of this software is governed by the Business Source License +// included in the file licenses/BSL.txt. +// +// As of the Change Date specified in that file, in accordance with +// the Business Source License, use of this software will be governed +// by the Apache License, Version 2.0, included in the file +// licenses/APL.txt. + +package commands + +import ( + "context" + + "github.com/cockroachdb/cockroach/pkg/build" + "github.com/spf13/cobra" +) + +// GetRootCommand returns the root command +func GetRootCommand(_ context.Context) *cobra.Command { + return &cobra.Command{ + Use: "drtprod [command] (flags)", + Short: "drtprod runs roachprod commands against DRT clusters", + Long: `drtprod is a tool for manipulating drt clusters using roachprod, +allowing easy creating, destruction, controls and configurations of clusters. + +Commands include: + execute: executes the commands as per the YAML configuration. Refer to pkg/cmd/drtprod/configs/drt_test.yaml as an example + *: any other command is passed to roachprod, potentially with flags added +`, + Version: "details:\n" + build.GetInfo().Long(), + } +} diff --git a/pkg/cmd/drtprod/cli/commands/yamlprocessor.go b/pkg/cmd/drtprod/cli/commands/yamlprocessor.go new file mode 100644 index 000000000000..793602c87732 --- /dev/null +++ b/pkg/cmd/drtprod/cli/commands/yamlprocessor.go @@ -0,0 +1,311 @@ +// Copyright 2024 The Cockroach Authors. +// +// Use of this software is governed by the Business Source License +// included in the file licenses/BSL.txt. +// +// As of the Change Date specified in that file, in accordance with +// the Business Source License, use of this software will be governed +// by the Apache License, Version 2.0, included in the file +// licenses/APL.txt. + +package commands + +import ( + "context" + "fmt" + "os" + "sync" + + "github.com/cockroachdb/cockroach/pkg/cmd/drtprod/helpers" + "github.com/spf13/cobra" + "gopkg.in/yaml.v2" +) + +// commandExecutor is responsible for executing the shell commands +var commandExecutor = helpers.ExecuteCmd + +// GetYamlProcessor creates a new Cobra command for processing a YAML file. +// The command expects a YAML file as an argument and runs the commands defined in it. +func GetYamlProcessor(ctx context.Context) *cobra.Command { + displayOnly := false + targets := make([]string, 0) + cobraCmd := &cobra.Command{ + Use: "execute [flags]", + Short: "Executes the commands in sequence as specified in the YAML", + Long: `Executes the commands in sequence as specified in the YAML. +You can also specify the rollback commands in case of a step failure. +`, + Args: cobra.ExactArgs(1), + // Wraps the command execution with additional error handling + Run: helpers.Wrap(func(cmd *cobra.Command, args []string) (retErr error) { + yamlFileLocation := args[0] + // Read the YAML file from the specified location + yamlContent, err := os.ReadFile(yamlFileLocation) + if err != nil { + return err + } + return processYaml(ctx, yamlContent, displayOnly, targets) + }), + } + cobraCmd.Flags().BoolVarP(&displayOnly, + "display-only", "d", false, "displays the commands that will be executed without running them") + cobraCmd.Flags().StringArrayVarP(&targets, + "targets", "t", nil, "the targets to execute. executes all if not mentioned.") + return cobraCmd +} + +// step represents an individual step in the YAML configuration. +// It can include an ActionStep and additional information for error handling and rollback. +type step struct { + Command string `yaml:"command"` // The command to execute + Script string `yaml:"script"` // The script to execute + + Args []string `yaml:"args"` // Arguments to pass to the command or script + Flags map[string]string `yaml:"flags"` // Flags to pass to the command or script + ContinueOnFailure bool `yaml:"continue_on_failure"` // Whether to continue on failure + OnRollback []step `yaml:"on_rollback"` // Steps to execute if rollback is needed +} + +// target defines a target cluster with associated steps to be executed. +type target struct { + TargetName string `yaml:"target_name"` // Name of the target cluster + Steps []step `yaml:"steps"` // Steps to execute on the target cluster +} + +// yamlConfig represents the structure of the entire YAML configuration file. +type yamlConfig struct { + Environment map[string]string `yaml:"environment"` // Environment variables to set + Targets []target `yaml:"targets"` // List of target clusters with their steps +} + +// command is a simplified representation of a shell command that needs to be executed. +type command struct { + name string // Command name + args []string // Command arguments + continueOnFailure bool // Whether to continue on failure + rollbackCmds []*command // Rollback commands to execute in case of failure +} + +// String returns the command as a string for easy printing. +func (c *command) String() string { + cmdStr := c.name + for _, arg := range c.args { + cmdStr = fmt.Sprintf("%s %s", cmdStr, arg) + } + return cmdStr +} + +// processYaml reads the YAML file, parses it, sets the environment variables, and processes the targets. +func processYaml( + ctx context.Context, yamlContent []byte, displayOnly bool, targets []string, +) (err error) { + + // Unmarshal the YAML content into the yamlConfig struct + var config yamlConfig + if err = yaml.UnmarshalStrict(yamlContent, &config); err != nil { + return err + } + + // Set the environment variables specified in the YAML + if err = setEnv(config.Environment, displayOnly); err != nil { + return err + } + + // Process the targets defined in the YAML + if err = processTargets(ctx, config.Targets, displayOnly, targets); err != nil { + return err + } + + return nil +} + +// setEnv sets the environment variables as defined in the YAML configuration. +func setEnv(environment map[string]string, displayOnly bool) error { + for key, value := range environment { + if displayOnly { + fmt.Printf("export %s=%s\n", key, value) + } else { + fmt.Printf("Setting env %s to %s\n", key, value) + } + // setting the environment for display only as well. This is because + // the environment will be used in the yaml as well. + err := os.Setenv(key, value) + if err != nil { + return err + } + } + return nil +} + +// processTargets processes each target defined in the YAML configuration. +// It generates commands for each target and executes them concurrently. +func processTargets( + ctx context.Context, targets []target, displayOnly bool, targetNames []string, +) error { + targetNameMap := make(map[string]struct{}) + targetMap := make(map[string][]*command) + for _, tn := range targetNames { + targetNameMap[tn] = struct{}{} + } + for i := 0; i < len(targets); i++ { + targets[i].TargetName = os.ExpandEnv(targets[i].TargetName) + t := targets[i] + if _, ok := targetNameMap[t.TargetName]; len(targetNames) > 0 && !ok { + fmt.Printf("Ignoring execution for target %s\n", t.TargetName) + continue + } + // Generate the commands for each target's steps + targetSteps, err := generateCmdsFromSteps(t.TargetName, t.Steps) + if err != nil { + return err + } + targetMap[t.TargetName] = targetSteps + } + + // Use a WaitGroup to execute commands concurrently + wg := sync.WaitGroup{} + for targetName, cmds := range targetMap { + if displayOnly { + displayCommands(targetName, cmds) + continue + } + wg.Add(1) + go func(tn string, commands []*command) { + err := executeCommands(ctx, tn, commands) + if err != nil { + fmt.Printf("%s: Error executing commands: %v\n", tn, err) + } + wg.Done() + }(targetName, cmds) + } + wg.Wait() + return nil +} + +// displayCommands prints the commands in stdout +func displayCommands(name string, cmds []*command) { + fmt.Printf("For target <%s>:\n", name) + for _, cmd := range cmds { + fmt.Printf("|-> %s\n", cmd) + for _, rCmd := range cmd.rollbackCmds { + fmt.Printf(" |-> (Rollback) %s\n", rCmd) + } + } + +} + +// executeCommands runs the list of commands for a specific target. +// It handles output streaming and error management. +func executeCommands(ctx context.Context, logPrefix string, cmds []*command) error { + // rollbackCmds maintains a list of commands to be executed in case of a failure + rollbackCmds := make([]*command, 0) + + // Defer rollback execution if any rollback commands are added + defer func() { + if len(rollbackCmds) > 0 { + _ = executeCommands(ctx, fmt.Sprintf("%s:Rollback", logPrefix), rollbackCmds) + } + }() + + for _, cmd := range cmds { + fmt.Printf("[%s] Starting <%v>\n", logPrefix, cmd) + err := commandExecutor(ctx, logPrefix, cmd.name, cmd.args...) + if err != nil { + if !cmd.continueOnFailure { + // Return the error if not configured to continue on failure + return err + } + // Log the failure and continue if configured to do so + fmt.Printf("[%s] Failed <%v>, Error Ignored: %v\n", logPrefix, cmd, err) + } else { + fmt.Printf("[%s] Completed <%v>\n", logPrefix, cmd) + } + + // Add rollback commands if specified + if len(cmd.rollbackCmds) > 0 { + for i := 0; i < len(cmd.rollbackCmds); i++ { + // rollback command failures are ignored + cmd.rollbackCmds[i].continueOnFailure = true + } + rollbackCmds = append(cmd.rollbackCmds, rollbackCmds...) + } + } + // Clear rollback commands if all commands executed successfully + rollbackCmds = make([]*command, 0) + return nil +} + +// generateCmdsFromSteps generates the commands to be executed for a given cluster and steps. +func generateCmdsFromSteps(clusterName string, steps []step) ([]*command, error) { + cmds := make([]*command, 0) + for _, s := range steps { + // Generate a command from each step + cmd, err := generateStepCmd(clusterName, s) + if err != nil { + return nil, err + } + if cmd == nil { + continue + } + cmds = append(cmds, cmd) + } + return cmds, nil +} + +// generateStepCmd generates a command for a given step within a target. +// It handles both command-based and script-based steps. +func generateStepCmd(clusterName string, s step) (*command, error) { + var cmd *command + var err error + + // Generate the command based on whether it's a command or a script + if s.Command != "" { + cmd, err = generateCmdFromCommand(s, clusterName) + } else if s.Script != "" { + cmd, err = generateCmdFromScript(s, clusterName) + } + + if err != nil { + return nil, err + } + + // Generate rollback commands if specified + if len(s.OnRollback) > 0 { + cmd.rollbackCmds, err = generateCmdsFromSteps(clusterName, s.OnRollback) + if err != nil { + return nil, err + } + } + return cmd, err +} + +// generateCmdFromCommand creates a command from a step that uses a command. +func generateCmdFromCommand(s step, _ string) (*command, error) { + // Prepend the cluster name to the command arguments + s.Args = append([]string{s.Command}, s.Args...) + return getCommand(s, "roachprod") +} + +// generateCmdFromScript creates a command from a step that uses a script. +func generateCmdFromScript(s step, _ string) (*command, error) { + return getCommand(s, s.Script) +} + +// getCommand constructs the final command with all arguments and flags. +func getCommand(step step, name string) (*command, error) { + args := make([]string, 0) + for _, arg := range step.Args { + args = append(args, os.ExpandEnv(arg)) + } + + // Append flags to the command arguments + for key, value := range step.Flags { + args = append(args, fmt.Sprintf("--%s=%s", key, os.ExpandEnv(value))) + } + + return &command{ + name: name, + args: args, + continueOnFailure: step.ContinueOnFailure, + }, nil +} diff --git a/pkg/cmd/drtprod/cli/commands/yamlprocessor_test.go b/pkg/cmd/drtprod/cli/commands/yamlprocessor_test.go new file mode 100644 index 000000000000..6bc4ef633e67 --- /dev/null +++ b/pkg/cmd/drtprod/cli/commands/yamlprocessor_test.go @@ -0,0 +1,157 @@ +// Copyright 2024 The Cockroach Authors. +// +// Use of this software is governed by the Business Source License +// included in the file licenses/BSL.txt. +// +// As of the Change Date specified in that file, in accordance with +// the Business Source License, use of this software will be governed +// by the Apache License, Version 2.0, included in the file +// licenses/APL.txt. + +package commands + +import ( + "context" + "fmt" + "strings" + "testing" + + "github.com/stretchr/testify/require" +) + +func Test_processYaml(t *testing.T) { + ctx := context.Background() + // setting to nil as a precaution that the command execution does not invoke an + // actual command + commandExecutor = nil + t.Run("expect unmarshall to fail", func(t *testing.T) { + err := processYaml(ctx, []byte("invalid"), false, nil) + require.NotNil(t, err) + require.Contains(t, err.Error(), "cannot unmarshal") + }) + t.Run("expect failure due to unwanted field", func(t *testing.T) { + err := processYaml(ctx, []byte(` +unwanted: value +environment: + NAME_1: name_value1 + NAME_2: name_value2 +`), false, nil) + require.NotNil(t, err) + require.Contains(t, err.Error(), "field unwanted not found in type commands.yamlConfig") + }) + t.Run("expect no command execution on display-only=true", func(t *testing.T) { + require.Nil(t, processYaml(ctx, getTestYaml(), true, nil)) + }) + t.Run("expect partial failure and rollback", func(t *testing.T) { + name1Commands := make([]string, 0) + name2Commands := make([]string, 0) + commandExecutor = func(ctx context.Context, logPrefix string, cmd string, args ...string) error { + if strings.HasPrefix(logPrefix, "name_value1") { + name1Commands = append(name1Commands, (&command{name: cmd, args: args}).String()) + } else if strings.HasPrefix(logPrefix, "name_value2") { + name2Commands = append(name2Commands, (&command{name: cmd, args: args}).String()) + } + if cmd == "dummy_script1" || cmd == "script33" || args[0] == "rb_dummy2" { + return fmt.Errorf("error while processing script %s", cmd) + } + return nil + } + require.Nil(t, processYaml(ctx, getTestYaml(), false, nil)) + require.Equal(t, 8, len(name1Commands)) + require.Equal(t, 1, len(name2Commands)) + // the flags are maintained as map and can be in any sequence + require.True(t, strings.HasPrefix(name1Commands[0], "roachprod dummy1 name_value1 arg11")) + require.True(t, strings.Contains(name1Commands[0], "--clouds=gce")) + require.True(t, strings.Contains(name1Commands[0], "--nodes=1")) + require.Equal(t, []string{ + "dummy_script1", "dummy_script2 arg11", "roachprod dummy2", "script33", + }, name1Commands[1:5]) + // rollback + require.True(t, strings.HasPrefix(name1Commands[5], "roachprod rb_dummy2 arg1 arg2")) + require.True(t, strings.Contains(name1Commands[5], "--flag1=value1")) + require.True(t, strings.Contains(name1Commands[5], "--flag2=value2")) + require.True(t, strings.HasPrefix(name1Commands[6], "dummy_script22")) + require.True(t, strings.Contains(name1Commands[6], "--f1=\\\"v1 v2\\\"")) + require.Equal(t, "roachprod rb_dummy1", name1Commands[7]) + require.Equal(t, []string{ + "roachprod dummy2 name_value2 arg12", + }, name2Commands) + }) + t.Run("expect no failure", func(t *testing.T) { + name1Commands := make([]string, 0) + name2Commands := make([]string, 0) + commandExecutor = func(ctx context.Context, logPrefix string, cmd string, args ...string) error { + if strings.HasPrefix(logPrefix, "name_value1") { + name1Commands = append(name1Commands, (&command{name: cmd, args: args}).String()) + } else if strings.HasPrefix(logPrefix, "name_value2") { + name2Commands = append(name2Commands, (&command{name: cmd, args: args}).String()) + } + return nil + } + require.Nil(t, processYaml(ctx, getTestYaml(), false, nil)) + require.Equal(t, 6, len(name1Commands)) + require.Equal(t, 1, len(name2Commands)) + // the flags are maintained as map and can be in any sequence + require.True(t, strings.HasPrefix(name1Commands[0], "roachprod dummy1 name_value1 arg11")) + require.True(t, strings.Contains(name1Commands[0], "--clouds=gce")) + require.True(t, strings.Contains(name1Commands[0], "--nodes=1")) + require.Equal(t, []string{ + "dummy_script1", "dummy_script2 arg11", "roachprod dummy2", "script33", "last_script", + }, name1Commands[1:]) + require.Equal(t, []string{ + "roachprod dummy2 name_value2 arg12", + }, name2Commands) + }) +} + +func getTestYaml() []byte { + return []byte(` +environment: + NAME_1: name_value1 + NAME_2: name_value2 + +targets: + - target_name: $NAME_1 + steps: + - command: dummy1 + args: + - $NAME_1 + - arg11 + flags: + clouds: gce + nodes: 1 + on_rollback: + - command: rb_dummy1 + - script: "dummy_script1" + continue_on_failure: True + - script: "dummy_script2" + args: + - arg11 + - command: dummy2 + on_rollback: + - command: rb_dummy2 + flags: + flag1: value1 + flag2: value2 + args: + - arg1 + - arg2 + - script: "dummy_script22" + flags: + f1: \"v1 v2\" + - script: "script33" + on_rollback: + - command: script33_rb + - script: "last_script" + on_rollback: + - command: rb_last + - target_name: $NAME_2 + steps: + - command: dummy2 + args: + - $NAME_2 + - arg12 + + +`) +} diff --git a/pkg/cmd/drtprod/cli/handlers.go b/pkg/cmd/drtprod/cli/handlers.go new file mode 100644 index 000000000000..28a06f0e79bb --- /dev/null +++ b/pkg/cmd/drtprod/cli/handlers.go @@ -0,0 +1,57 @@ +// Copyright 2024 The Cockroach Authors. +// +// Use of this software is governed by the Business Source License +// included in the file licenses/BSL.txt. +// +// As of the Change Date specified in that file, in accordance with +// the Business Source License, use of this software will be governed +// by the Apache License, Version 2.0, included in the file +// licenses/APL.txt. + +package cli + +import ( + "context" + "os" + "strings" + + "github.com/cockroachdb/cockroach/pkg/cmd/drtprod/cli/commands" + "github.com/cockroachdb/cockroach/pkg/cmd/drtprod/helpers" + "github.com/cockroachdb/cockroach/pkg/roachprod" + "github.com/spf13/cobra" +) + +// Initialize sets up the environment and initializes the command-line interface. +func Initialize(ctx context.Context) { + // Set environment variables for the GCE project and DNS configurations. + _ = os.Setenv("ROACHPROD_DNS", "drt.crdb.io") + _ = os.Setenv("ROACHPROD_GCE_DNS_DOMAIN", "drt.crdb.io") + _ = os.Setenv("ROACHPROD_GCE_DNS_ZONE", "drt") + _ = os.Setenv("ROACHPROD_GCE_DEFAULT_PROJECT", "cockroach-drt") + // Initialize cloud providers for roachprod. + _ = roachprod.InitProviders() + + // Disable command sorting in Cobra (command-line parser). + cobra.EnableCommandSorting = false + + // Create the root command and add subcommands. + rootCommand := commands.GetRootCommand(ctx) + rootCommand.AddCommand(register(ctx)...) + + // Check if the command is found in drtprod; if not, redirect to roachprod. + _, _, err := rootCommand.Find(os.Args[1:]) + if err != nil { + if strings.Contains(err.Error(), "unknown command") { + // Command not found, execute it in roachprod instead. + _ = helpers.ExecuteCmd(ctx, "roachprod", "roachprod", os.Args[1:]...) + return + } + // If another error occurs, exit with a failure status. + os.Exit(1) + } + + // Execute the root command, exit if an error occurs. + if err := rootCommand.Execute(); err != nil { + os.Exit(1) + } +} diff --git a/pkg/cmd/drtprod/cli/registry.go b/pkg/cmd/drtprod/cli/registry.go new file mode 100644 index 000000000000..2d1446a7d046 --- /dev/null +++ b/pkg/cmd/drtprod/cli/registry.go @@ -0,0 +1,26 @@ +// Copyright 2024 The Cockroach Authors. +// +// Use of this software is governed by the Business Source License +// included in the file licenses/BSL.txt. +// +// As of the Change Date specified in that file, in accordance with +// the Business Source License, use of this software will be governed +// by the Apache License, Version 2.0, included in the file +// licenses/APL.txt. + +package cli + +import ( + "context" + + "github.com/cockroachdb/cockroach/pkg/cmd/drtprod/cli/commands" + "github.com/spf13/cobra" +) + +// register registers all drtprod subcommands. +// Add your commands here +func register(ctx context.Context) []*cobra.Command { + return []*cobra.Command{ + commands.GetYamlProcessor(ctx), + } +} diff --git a/pkg/cmd/drtprod/configs/drt_test.yaml b/pkg/cmd/drtprod/configs/drt_test.yaml new file mode 100644 index 000000000000..bbf173e2b6c6 --- /dev/null +++ b/pkg/cmd/drtprod/configs/drt_test.yaml @@ -0,0 +1,39 @@ +# Yaml to create a test-cluster. Please make sure that you change the cluster names to avoid conflicts. +environment: + GCE_PROJECT: cockroach-ephemeral + CLUSTER: drt-test + WORKLOAD_CLUSTER: workload-test + +targets: + - target_name: $CLUSTER + steps: + - command: create + args: + - $CLUSTER + flags: + clouds: gce + nodes: 1 + username: drt + gce-machine-type: n2-standard-2 + gce-use-spot: true + on_rollback: + - command: destroy + args: + - $CLUSTER + - command: stage + args: + - $CLUSTER + - cockroach + - command: start + args: + - $CLUSTER + - target_name: $WORKLOAD_CLUSTER + steps: + - command: create + args: + - $WORKLOAD_CLUSTER + flags: + nodes: 1 + username: workload + on_rollback: + - command: destroy diff --git a/pkg/cmd/drtprod/helpers/BUILD.bazel b/pkg/cmd/drtprod/helpers/BUILD.bazel new file mode 100644 index 000000000000..77e4cb652820 --- /dev/null +++ b/pkg/cmd/drtprod/helpers/BUILD.bazel @@ -0,0 +1,12 @@ +load("@io_bazel_rules_go//go:def.bzl", "go_library") + +go_library( + name = "helpers", + srcs = ["utils.go"], + importpath = "github.com/cockroachdb/cockroach/pkg/cmd/drtprod/helpers", + visibility = ["//visibility:public"], + deps = [ + "//pkg/roachprod/errors", + "@com_github_spf13_cobra//:cobra", + ], +) diff --git a/pkg/cmd/drtprod/helpers/utils.go b/pkg/cmd/drtprod/helpers/utils.go new file mode 100644 index 000000000000..b853acf8284a --- /dev/null +++ b/pkg/cmd/drtprod/helpers/utils.go @@ -0,0 +1,91 @@ +// Copyright 2024 The Cockroach Authors. +// +// Use of this software is governed by the Business Source License +// included in the file licenses/BSL.txt. +// +// As of the Change Date specified in that file, in accordance with +// the Business Source License, use of this software will be governed +// by the Apache License, Version 2.0, included in the file +// licenses/APL.txt. + +package helpers + +import ( + "bufio" + "context" + "fmt" + "os" + "os/exec" + + rperrors "github.com/cockroachdb/cockroach/pkg/roachprod/errors" + "github.com/spf13/cobra" +) + +// Wrap provide `cobra.Command` functions with a standard return code handler. +// Exit codes come from rperrors.Error.ExitCode(). +// +// If the wrapped error tree of an error does not contain an instance of +// rperrors.Error, the error will automatically be wrapped with +// rperrors.Unclassified. +func Wrap(f func(cmd *cobra.Command, args []string) error) func(cmd *cobra.Command, args []string) { + return func(cmd *cobra.Command, args []string) { + var err error + err = f(cmd, args) + if err != nil { + drtprodError, ok := rperrors.AsError(err) + if !ok { + drtprodError = rperrors.Unclassified{Err: err} + err = drtprodError + } + + cmd.Printf("Error: %+v\n", err) + + os.Exit(drtprodError.ExitCode()) + } + } +} + +// ExecuteCmd runs a shell command with the given arguments and streams the output. +func ExecuteCmd(ctx context.Context, logPrefix string, cmd string, args ...string) error { + // Create a command with the given context and arguments. + c := exec.CommandContext(ctx, cmd, args...) + + // Set up pipes to capture the command's stdout and stderr. + stdout, err := c.StdoutPipe() + if err != nil { + return err + } + stderr, err := c.StderrPipe() + if err != nil { + return err + } + + // Start the command execution + err = c.Start() + if err != nil { + return err + } + + // Stream stdout output + outScanner := bufio.NewScanner(stdout) + outScanner.Split(bufio.ScanLines) + go func() { + for outScanner.Scan() { + m := outScanner.Text() + fmt.Printf("[%s] %s\n", logPrefix, m) + } + }() + + // Stream stderr output + errScanner := bufio.NewScanner(stderr) + errScanner.Split(bufio.ScanLines) + go func() { + for errScanner.Scan() { + m := errScanner.Text() + fmt.Printf("[%s] %s\n", logPrefix, m) + } + }() + + // Wait for the command to complete and return any errors encountered. + return c.Wait() +} diff --git a/pkg/cmd/drtprod/main.go b/pkg/cmd/drtprod/main.go new file mode 100644 index 000000000000..6303a4f4c5bf --- /dev/null +++ b/pkg/cmd/drtprod/main.go @@ -0,0 +1,22 @@ +// Copyright 2024 The Cockroach Authors. +// +// Use of this software is governed by the Business Source License +// included in the file licenses/BSL.txt. +// +// As of the Change Date specified in that file, in accordance with +// the Business Source License, use of this software will be governed +// by the Apache License, Version 2.0, included in the file +// licenses/APL.txt. + +package main + +import ( + "context" + + "github.com/cockroachdb/cockroach/pkg/cmd/drtprod/cli" +) + +func main() { + // Initialize and register all commands + cli.Initialize(context.Background()) +}