This is an automated email from the ASF dual-hosted git repository. maxyang pushed a commit to branch main in repository https://gitbox.apache.org/repos/asf/cloudberry-go-libs.git
commit 87ad28e303e586d188dafea25786a15ac882cccd Author: Karen Huddleston <[email protected]> AuthorDate: Tue Jul 16 17:45:19 2024 -0700 Add ExecuteLocalCommandWithContext Adding a version to run local commands with context to handle timeouts for long running commands --- cluster/cluster.go | 7 +++++++ cluster/cluster_test.go | 37 +++++++++++++++++++++++++++++++++++++ testhelper/structs.go | 27 +++++++++++++++++++++++++++ 3 files changed, 71 insertions(+) diff --git a/cluster/cluster.go b/cluster/cluster.go index 62ccd85..0311d0c 100644 --- a/cluster/cluster.go +++ b/cluster/cluster.go @@ -8,6 +8,7 @@ package cluster import ( "bufio" "bytes" + "context" "fmt" "os" "os/exec" @@ -24,6 +25,7 @@ import ( type Executor interface { ExecuteLocalCommand(commandStr string) (string, error) + ExecuteLocalCommandWithContext(commandStr string, ctx context.Context) (string, error) ExecuteClusterCommand(scope Scope, commandList []ShellCommand) *RemoteOutput } @@ -330,6 +332,11 @@ func (executor *GPDBExecutor) ExecuteLocalCommand(commandStr string) (string, er return string(output), err } +func (executor *GPDBExecutor) ExecuteLocalCommandWithContext(commandStr string, ctx context.Context) (string, error) { + output, err := exec.CommandContext(ctx, "bash", "-c", commandStr).CombinedOutput() + return string(output), err +} + /* * This function just executes all of the commands passed to it in parallel; it * doesn't care about the scope of the command except to pass that on to the diff --git a/cluster/cluster_test.go b/cluster/cluster_test.go index e753ac4..60a397d 100644 --- a/cluster/cluster_test.go +++ b/cluster/cluster_test.go @@ -1,12 +1,14 @@ package cluster_test import ( + "context" "database/sql/driver" "fmt" "os" "os/user" "path" "testing" + "time" sqlmock "github.com/DATA-DOG/go-sqlmock" @@ -470,6 +472,41 @@ var _ = Describe("cluster/cluster tests", func() { Expect(err.Error()).To(Equal("exit status 127")) }) }) + Describe("ExecuteLocalCommandWithContext", func() { + BeforeEach(func() { + os.MkdirAll("/tmp/gp_common_go_libs_test", 0777) + }) + AfterEach(func() { + os.RemoveAll("/tmp/gp_common_go_libs_test") + }) + It("runs the specified command", func() { + testCluster := cluster.Cluster{} + commandStr := "touch /tmp/gp_common_go_libs_test/foo" + testCluster.Executor = &cluster.GPDBExecutor{} + testCluster.ExecuteLocalCommandWithContext(commandStr, context.TODO()) + + expectPathToExist("/tmp/gp_common_go_libs_test/foo") + }) + It("returns any error generated by the specified command", func() { + testCluster := cluster.Cluster{} + commandStr := "some-non-existent-command /tmp/gp_common_go_libs_test/foo" + testCluster.Executor = &cluster.GPDBExecutor{} + output, err := testCluster.ExecuteLocalCommandWithContext(commandStr, context.TODO()) + + Expect(output).To(ContainSubstring("some-non-existent-command: command not found\n")) + Expect(err.Error()).To(Equal("exit status 127")) + }) + It("kills the command if it runs beyond the timeout", func() { + testCluster := cluster.Cluster{} + commandStr := "while true; do echo Keep running; sleep 0.1; done" + ctx, _ := context.WithTimeout(context.Background(), 200*time.Millisecond) + testCluster.Executor = &cluster.GPDBExecutor{} + output, err := testCluster.ExecuteLocalCommandWithContext(commandStr, ctx) + Expect(ctx.Err()).To(Equal(context.DeadlineExceeded)) + Expect(err).To(HaveOccurred()) + Expect(output).To(Equal("Keep running\nKeep running\n")) + }) + }) Describe("ExecuteClusterCommand", func() { BeforeEach(func() { os.MkdirAll("/tmp/gp_common_go_libs_test", 0777) diff --git a/testhelper/structs.go b/testhelper/structs.go index 6816556..41c9462 100644 --- a/testhelper/structs.go +++ b/testhelper/structs.go @@ -5,6 +5,7 @@ package testhelper */ import ( + "context" "github.com/cloudberrydb/gp-common-go-libs/cluster" "github.com/cloudberrydb/gp-common-go-libs/gplog" "github.com/jmoiron/sqlx" @@ -58,6 +59,7 @@ type TestExecutor struct { LocalError error LocalErrors []error LocalCommands []string + LocalContexts []context.Context ClusterOutput *cluster.RemoteOutput ClusterOutputs []*cluster.RemoteOutput @@ -95,6 +97,31 @@ func (executor *TestExecutor) ExecuteLocalCommand(commandStr string) (string, er return executor.LocalOutput, nil } +func (executor *TestExecutor) ExecuteLocalCommandWithContext(commandStr string, ctx context.Context) (string, error) { + executor.NumExecutions++ + executor.NumLocalExecutions++ + executor.LocalCommands = append(executor.LocalCommands, commandStr) + executor.LocalContexts = append(executor.LocalContexts, ctx) + if (executor.LocalOutputs == nil && executor.LocalErrors != nil) || (executor.LocalOutputs != nil && executor.LocalErrors == nil) { + gplog.Fatal(nil, "If one of LocalOutputs or LocalErrors is set, both must be set") + } else if executor.LocalOutputs != nil && executor.LocalErrors != nil && len(executor.LocalOutputs) != len(executor.LocalErrors) { + gplog.Fatal(nil, "Found %d LocalOutputs and %d LocalErrors, but one output and one error must be set for each call", len(executor.LocalOutputs), len(executor.LocalErrors)) + } + if executor.LocalOutputs != nil { + if executor.NumLocalExecutions <= len(executor.LocalOutputs) { + return executor.LocalOutputs[executor.NumLocalExecutions-1], executor.LocalErrors[executor.NumLocalExecutions-1] + } else if executor.UseLastOutput { + return executor.LocalOutputs[len(executor.LocalOutputs)-1], executor.LocalErrors[len(executor.LocalErrors)-1] + } else if executor.UseDefaultOutput { + return executor.LocalOutput, executor.LocalError + } + gplog.Fatal(nil, "ExecuteLocalCommandWithContext called %d times, but only %d outputs and errors provided", executor.NumLocalExecutions, len(executor.LocalOutputs)) + } else if executor.ErrorOnExecNum == 0 || executor.NumLocalExecutions == executor.ErrorOnExecNum { + return executor.LocalOutput, executor.LocalError + } + return executor.LocalOutput, nil +} + func (executor *TestExecutor) ExecuteClusterCommand(scope cluster.Scope, commandList []cluster.ShellCommand) *cluster.RemoteOutput { executor.NumExecutions++ executor.NumClusterExecutions++ --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
