Added customizable fields during registration (flyteorg#80)

Signed-off-by: Prafulla Mahindrakar <[email protected]>

flytectl/cmd/config/subcommand/register/files_config.go

@@ -0,0 +1,20 @@
+package register
+
+//go:generate pflags FilesConfig --default-var DefaultFilesConfig --bind-default-var
+
+var (
+	DefaultFilesConfig = &FilesConfig{
+		Version:         "v1",
+		ContinueOnError: false,
+	}
+)
+
+// FilesConfig containing flags used for registration
+type FilesConfig struct {
+	Version              string `json:"version" pflag:",version of the entity to be registered with flyte."`
+	ContinueOnError      bool   `json:"continueOnError" pflag:",continue on error when registering files."`
+	Archive              bool   `json:"archive" pflag:",pass in archive file either an http link or local path."`
+	AssumableIamRole     string `json:"assumableIamRole" pflag:", custom assumable iam auth role to register launch plans with."`
+	K8ServiceAccount     string `json:"k8ServiceAccount" pflag:", custom kubernetes service account auth role to register launch plans with."`
+	OutputLocationPrefix string `json:"outputLocationPrefix" pflag:", custom output location prefix for offloaded types (files/schemas)."`
+}

flytectl/cmd/create/execution_util.go

@@ -57,12 +57,11 @@ func createExecutionRequestForTask(ctx context.Context, taskName string, project
 		Literals: variableLiterals,
 	}
 	var authRole *admin.AuthRole
-	if executionConfig.KubeServiceAcct != "" {
-		authRole = &admin.AuthRole{Method: &admin.AuthRole_KubernetesServiceAccount{
-			KubernetesServiceAccount: executionConfig.KubeServiceAcct}}
-	} else {
-		authRole = &admin.AuthRole{Method: &admin.AuthRole_AssumableIamRole{
-			AssumableIamRole: executionConfig.IamRoleARN}}
+	if executionConfig.KubeServiceAcct != "" || executionConfig.IamRoleARN != "" {
+		authRole = &admin.AuthRole{
+			AssumableIamRole:         executionConfig.IamRoleARN,
+			KubernetesServiceAccount: executionConfig.KubeServiceAcct,
+		}
 	}
 	ID := &core.Identifier{
 		ResourceType: core.ResourceType_TASK,

flytectl/cmd/register/files.go

@@ -5,26 +5,12 @@ import (
 	"encoding/json"
 	"os"
 
+	rconfig "github.com/flyteorg/flytectl/cmd/config/subcommand/register"
 	cmdCore "github.com/flyteorg/flytectl/cmd/core"
 	"github.com/flyteorg/flytectl/pkg/printer"
 	"github.com/flyteorg/flytestdlib/logger"
 )
 
-//go:generate pflags FilesConfig
-var (
-	filesConfig = &FilesConfig{
-		Version:         "v1",
-		ContinueOnError: false,
-	}
-)
-
-// FilesConfig
-type FilesConfig struct {
-	Version         string `json:"version" pflag:",version of the entity to be registered with flyte."`
-	ContinueOnError bool   `json:"continueOnError" pflag:",continue on error when registering files."`
-	Archive         bool   `json:"archive" pflag:",pass in archive file either an http link or local path."`
-}
-
 const (
 	registerFilesShort = "Registers file resources"
 	registerFilesLong  = `
@@ -70,6 +56,24 @@ Change the o/p format has not effect on registration. The O/p is currently avail
 
  bin/flytectl register file  _pb_output/* -d development  -p flytesnacks -c -o yaml
 
+Override IamRole during registration.
+
+::
+
+ bin/flytectl register file  _pb_output/* -d development  -p flytesnacks -c -v v2 -i "arn:aws:iam::123456789:role/dummy"
+
+Override Kubernetes service account during registration.
+
+::
+
+ bin/flytectl register file  _pb_output/* -d development  -p flytesnacks -c -v v2 -k "kubernetes-service-account"
+
+Override Output location prefix during registration.
+
+::
+
+ bin/flytectl register file  _pb_output/* -d development  -p flytesnacks -c -v v2 -l "s3://dummy/prefix"
+
 Usage
 `
 )
@@ -81,7 +85,7 @@ func registerFromFilesFunc(ctx context.Context, args []string, cmdCtx cmdCore.Co
 		return _err
 	}
 	logger.Infof(ctx, "Parsing files... Total(%v)", len(dataRefs))
-	fastFail := !filesConfig.ContinueOnError
+	fastFail := !rconfig.DefaultFilesConfig.ContinueOnError
 	var registerResults []Result
 	for i := 0; i < len(dataRefs) && !(fastFail && _err != nil); i++ {
 		registerResults, _err = registerFile(ctx, dataRefs[i], registerResults, cmdCtx)

flytectl/cmd/register/files_test.go

@@ -0,0 +1,32 @@
+package register
+
+import (
+	"testing"
+
+	rconfig "github.com/flyteorg/flytectl/cmd/config/subcommand/register"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+)
+
+func TestRegisterFromFiles(t *testing.T) {
+	t.Run("Valid registration", func(t *testing.T) {
+		setup()
+		registerFilesSetup()
+		rconfig.DefaultFilesConfig.Archive = true
+		args = []string{"testdata/valid-parent-folder-register.tar"}
+		mockAdminClient.OnCreateTaskMatch(mock.Anything, mock.Anything).Return(nil, nil)
+		mockAdminClient.OnCreateWorkflowMatch(mock.Anything, mock.Anything).Return(nil, nil)
+		mockAdminClient.OnCreateLaunchPlanMatch(mock.Anything, mock.Anything).Return(nil, nil)
+		err := registerFromFilesFunc(ctx, args, cmdCtx)
+		assert.Nil(t, err)
+	})
+	t.Run("Invalid registration file", func(t *testing.T) {
+		setup()
+		registerFilesSetup()
+		rconfig.DefaultFilesConfig.Archive = true
+		args = []string{"testdata/invalid.tar"}
+		err := registerFromFilesFunc(ctx, args, cmdCtx)
+		assert.NotNil(t, err)
+	})
+}

flytectl/cmd/register/register.go

@@ -1,7 +1,9 @@
 package register
 
 import (
+	rconfig "github.com/flyteorg/flytectl/cmd/config/subcommand/register"
 	cmdcore "github.com/flyteorg/flytectl/cmd/core"
+
 	"github.com/spf13/cobra"
 )
 
@@ -24,8 +26,8 @@ func RemoteRegisterCommand() *cobra.Command {
 		Long:  registercmdLong,
 	}
 	registerResourcesFuncs := map[string]cmdcore.CommandEntry{
-		"files": {CmdFunc: registerFromFilesFunc, Aliases: []string{"file"}, PFlagProvider: filesConfig, Short: registerFilesShort,
-			Long: registerFilesLong},
+		"files": {CmdFunc: registerFromFilesFunc, Aliases: []string{"file"}, PFlagProvider: rconfig.DefaultFilesConfig,
+			Short: registerFilesShort, Long: registerFilesLong},
 	}
 	cmdcore.AddCommands(registerCmd, registerResourcesFuncs)
 	return registerCmd

flytectl/cmd/register/register_test.go

@@ -1,13 +1,29 @@
 package register
 
 import (
+	"context"
 	"fmt"
+	"net/http"
 	"sort"
 	"testing"
 
+	cmdCore "github.com/flyteorg/flytectl/cmd/core"
+	u "github.com/flyteorg/flytectl/cmd/testutils"
+	"github.com/flyteorg/flyteidl/clients/go/admin/mocks"
+
 	"github.com/stretchr/testify/assert"
 )
 
+var (
+	ctx             context.Context
+	mockAdminClient *mocks.AdminServiceClient
+	cmdCtx          cmdCore.CommandContext
+	args            []string
+	GetDoFunc       func(req *http.Request) (*http.Response, error)
+)
+
+var setup = u.Setup
+
 func TestRegisterCommand(t *testing.T) {
 	registerCommand := RemoteRegisterCommand()
 	assert.Equal(t, registerCommand.Use, "register")