Add force flag when terminating execution (#6)

Signed-off-by: Rafael Raposo <[email protected]>
RRap0so · Oct 10, 2024 · d30b87e · d30b87e
1 parent f758d76
commit d30b87e
Show file tree

Hide file tree

Showing 3 changed files with 54 additions and 8 deletions.
diff --git a/flytectl/cmd/delete/delete.go b/flytectl/cmd/delete/delete.go
@@ -1,6 +1,7 @@
 package delete
 
 import (
+	"context"
 	"github.com/flyteorg/flyte/flytectl/cmd/config/subcommand/clusterresourceattribute"
 	"github.com/flyteorg/flyte/flytectl/cmd/config/subcommand/execution"
 	"github.com/flyteorg/flyte/flytectl/cmd/config/subcommand/executionclusterlabel"
@@ -19,7 +20,7 @@ const (
 Delete a resource; if an execution:
 ::
 
- flytectl delete execution kxd1i72850  -d development  -p flytesnacks
+ flytectl delete execution kxd1i72850  -d development  -p flytesnacks -f true
 `
 )
 
@@ -30,9 +31,18 @@ func RemoteDeleteCommand() *cobra.Command {
 		Short: deleteCmdShort,
 		Long:  deleteCmdLong,
 	}
+	var force bool
+	deleteCmd.PersistentFlags().BoolVarP(&force, "force", "f", false, "Force deletion without confirmation")
 	terminateResourcesFuncs := map[string]cmdcore.CommandEntry{
-		"execution": {CmdFunc: terminateExecutionFunc, Aliases: []string{"executions"}, Short: execCmdShort,
-			Long: execCmdLong, PFlagProvider: execution.DefaultExecDeleteConfig},
+		"execution": {
+			CmdFunc: func(ctx context.Context, args []string, cmdCtx cmdcore.CommandContext) error {
+				return terminateExecutionFunc(ctx, args, cmdCtx, force)
+			},
+			Aliases:       []string{"executions"},
+			Short:         execCmdShort,
+			Long:          execCmdLong,
+			PFlagProvider: execution.DefaultExecDeleteConfig,
+		},
 		"task-resource-attribute": {CmdFunc: deleteTaskResourceAttributes, Aliases: []string{"task-resource-attributes"},
 			Short: taskResourceAttributesShort,
 			Long:  taskResourceAttributesLong, PFlagProvider: taskresourceattribute.DefaultDelConfig, ProjectDomainNotRequired: true},

diff --git a/flytectl/cmd/delete/execution.go b/flytectl/cmd/delete/execution.go
@@ -2,7 +2,6 @@ package delete
 
 import (
 	"context"
-
 	"github.com/flyteorg/flyte/flytectl/cmd/config"
 	"github.com/flyteorg/flyte/flytectl/cmd/config/subcommand/execution"
 	cmdCore "github.com/flyteorg/flyte/flytectl/cmd/core"
@@ -22,6 +21,11 @@ Terminate a single execution with its name:
 
  flytectl delete execution c6a51x2l9e  -d development  -p flytesnacks
 
+Force an execution termination:
+::
+
+ flytectl delete execution ds3vgtgc4  -d development  -p flytesnacks -f true
+
 .. note::
     The terms execution/executions are interchangeable in these commands.
 
@@ -60,7 +64,7 @@ Usage
 `
 )
 
-func terminateExecutionFunc(ctx context.Context, args []string, cmdCtx cmdCore.CommandContext) error {
+func terminateExecutionFunc(ctx context.Context, args []string, cmdCtx cmdCore.CommandContext, force bool) error {
 	for i := 0; i < len(args); i++ {
 		name := args[i]
 		logger.Infof(ctx, "Terminating execution of %v execution ", name)
@@ -73,6 +77,7 @@ func terminateExecutionFunc(ctx context.Context, args []string, cmdCtx cmdCore.C
 					Domain:  config.GetConfig().Domain,
 					Name:    name,
 				},
+				Force: force,
 			})
 			if err != nil {
 				logger.Errorf(ctx, "Failed to terminate execution of %v execution due to %v ", name, err)

diff --git a/flytectl/cmd/delete/execution_test.go b/flytectl/cmd/delete/execution_test.go
@@ -31,13 +31,44 @@ func terminateExecutionSetup() {
 	}
 }
 
+func forceTerminateExecutionSetup() {
+	args = append(args, "exec1", "exec2")
+	terminateExecRequests = []*admin.ExecutionTerminateRequest{
+		{Id: &core.WorkflowExecutionIdentifier{
+			Name:    "exec1",
+			Project: config.GetConfig().Project,
+			Domain:  config.GetConfig().Domain,
+		},
+			Force: true},
+		{Id: &core.WorkflowExecutionIdentifier{
+			Name:    "exec2",
+			Project: config.GetConfig().Project,
+			Domain:  config.GetConfig().Domain,
+		},
+			Force: true},
+	}
+}
+
 func TestTerminateExecutionFunc(t *testing.T) {
 	s := setup()
 	terminateExecutionSetup()
 	terminateExecResponse := &admin.ExecutionTerminateResponse{}
 	s.MockAdminClient.OnTerminateExecutionMatch(s.Ctx, terminateExecRequests[0]).Return(terminateExecResponse, nil)
 	s.MockAdminClient.OnTerminateExecutionMatch(s.Ctx, terminateExecRequests[1]).Return(terminateExecResponse, nil)
-	err := terminateExecutionFunc(s.Ctx, args, s.CmdCtx)
+	err := terminateExecutionFunc(s.Ctx, args, s.CmdCtx, false)
+	assert.Nil(t, err)
+	s.MockAdminClient.AssertCalled(t, "TerminateExecution", s.Ctx, terminateExecRequests[0])
+	s.MockAdminClient.AssertCalled(t, "TerminateExecution", s.Ctx, terminateExecRequests[1])
+	s.TearDownAndVerify(t, "")
+}
+
+func TestForceTerminateExecutionFunc(t *testing.T) {
+	s := setup()
+	forceTerminateExecutionSetup()
+	terminateExecResponse := &admin.ExecutionTerminateResponse{}
+	s.MockAdminClient.OnTerminateExecutionMatch(s.Ctx, terminateExecRequests[0]).Return(terminateExecResponse, nil)
+	s.MockAdminClient.OnTerminateExecutionMatch(s.Ctx, terminateExecRequests[1]).Return(terminateExecResponse, nil)
+	err := terminateExecutionFunc(s.Ctx, args, s.CmdCtx, true)
 	assert.Nil(t, err)
 	s.MockAdminClient.AssertCalled(t, "TerminateExecution", s.Ctx, terminateExecRequests[0])
 	s.MockAdminClient.AssertCalled(t, "TerminateExecution", s.Ctx, terminateExecRequests[1])
@@ -50,7 +81,7 @@ func TestTerminateExecutionFuncWithError(t *testing.T) {
 	terminateExecResponse := &admin.ExecutionTerminateResponse{}
 	s.MockAdminClient.OnTerminateExecutionMatch(s.Ctx, terminateExecRequests[0]).Return(nil, errors.New("failed to terminate"))
 	s.MockAdminClient.OnTerminateExecutionMatch(s.Ctx, terminateExecRequests[1]).Return(terminateExecResponse, nil)
-	err := terminateExecutionFunc(s.Ctx, args, s.CmdCtx)
+	err := terminateExecutionFunc(s.Ctx, args, s.CmdCtx, false)
 	assert.Equal(t, errors.New("failed to terminate"), err)
 	s.MockAdminClient.AssertCalled(t, "TerminateExecution", s.Ctx, terminateExecRequests[0])
 	s.MockAdminClient.AssertNotCalled(t, "TerminateExecution", s.Ctx, terminateExecRequests[1])
@@ -63,7 +94,7 @@ func TestTerminateExecutionFuncWithPartialSuccess(t *testing.T) {
 	terminateExecResponse := &admin.ExecutionTerminateResponse{}
 	s.MockAdminClient.OnTerminateExecutionMatch(s.Ctx, terminateExecRequests[0]).Return(terminateExecResponse, nil)
 	s.MockAdminClient.OnTerminateExecutionMatch(s.Ctx, terminateExecRequests[1]).Return(nil, errors.New("failed to terminate"))
-	err := terminateExecutionFunc(s.Ctx, args, s.CmdCtx)
+	err := terminateExecutionFunc(s.Ctx, args, s.CmdCtx, false)
 	assert.Equal(t, errors.New("failed to terminate"), err)
 	s.MockAdminClient.AssertCalled(t, "TerminateExecution", s.Ctx, terminateExecRequests[0])
 	s.MockAdminClient.AssertCalled(t, "TerminateExecution", s.Ctx, terminateExecRequests[1])