Support passing an iam role AND a k8s service account (flyteorg#198)

go.mod

@@ -16,7 +16,7 @@ require (
 	github.com/coreos/go-oidc v2.2.1+incompatible
 	github.com/dgrijalva/jwt-go v3.2.0+incompatible
 	github.com/evanphx/json-patch v4.9.0+incompatible
-	github.com/flyteorg/flyteidl v0.18.40
+	github.com/flyteorg/flyteidl v0.18.50
 	github.com/flyteorg/flyteplugins v0.5.38
 	github.com/flyteorg/flytepropeller v0.7.8
 	github.com/flyteorg/flytestdlib v0.3.22

go.sum

@@ -86,6 +86,7 @@ github.com/Azure/go-autorest/logger v0.2.1/go.mod h1:T9E3cAhj2VqvPOtCYAvby9aBXkZ
 github.com/Azure/go-autorest/tracing v0.5.0/go.mod h1:r/s2XiOKccPW3HrqB+W0TQzfbtp2fGCgRFtBroKn4Dk=
 github.com/Azure/go-autorest/tracing v0.6.0 h1:TYi4+3m5t6K48TGI9AUdb+IzbnSxvnvUMfuitfgcfuo=
 github.com/Azure/go-autorest/tracing v0.6.0/go.mod h1:+vhtPC754Xsa23ID7GlGsrdKBpUA79WCAKPPZVC2DeU=
+github.com/BurntSushi/toml v0.3.1 h1:WXkYYl6Yr3qBf1K79EBnL4mak0OimBfB0XUf9Vl28OQ=
 github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03qcyfWMU=
 github.com/BurntSushi/xgb v0.0.0-20160522181843-27f122750802/go.mod h1:IVnqGOEym/WlBOVXweHU+Q+/VP0lqqI8lqeDx9IjBqo=
 github.com/DATA-DOG/go-sqlmock v1.3.3/go.mod h1:f/Ixk793poVmq4qj/V1dPUg2JEAKC73Q5eFN3EC/SaM=
@@ -306,8 +307,8 @@ github.com/felixge/httpsnoop v1.0.1 h1:lvB5Jl89CsZtGIWuTcDM1E/vkVs49/Ml7JJe07l8S
 github.com/felixge/httpsnoop v1.0.1/go.mod h1:m8KPJKqk1gH5J9DgRY2ASl2lWCfGKXixSwevea8zH2U=
 github.com/flyteorg/flyteidl v0.18.17/go.mod h1:b5Fq4Z8a5b0mF6pEwTd48ufvikUGVkWSjZiMT0ZtqKI=
 github.com/flyteorg/flyteidl v0.18.24/go.mod h1:b5Fq4Z8a5b0mF6pEwTd48ufvikUGVkWSjZiMT0ZtqKI=
-github.com/flyteorg/flyteidl v0.18.40 h1:YuLBNpIotOFwyLSXSs0aj3B9N9vwPhzLRAQTWxYSI/w=
-github.com/flyteorg/flyteidl v0.18.40/go.mod h1:IJD02cc/95QMkGDBJNibsr5aWd6V7TlQiJ8Iz5mVZ28=
+github.com/flyteorg/flyteidl v0.18.50 h1:L1fMj6QEXoKin+cPQn9sfwJ1x14tlChdz1mG1WaaIW4=
+github.com/flyteorg/flyteidl v0.18.50/go.mod h1:576W2ViEyjTpT+kEVHAGbrTP3HARNUZ/eCwrNPmdx9U=
 github.com/flyteorg/flyteplugins v0.5.38 h1:xAQ1J23cRxzwNDgzbmRuuvflq2PFetntRCjuM5RBfTw=
 github.com/flyteorg/flyteplugins v0.5.38/go.mod h1:CxerBGWWEmNYmPxSMHnwQEr9cc1Fbo/g5fcABazU6Jo=
 github.com/flyteorg/flytepropeller v0.7.8 h1:O441kDHJUayS/2rebTj7VG4e1LowrweazQhzTaZ97m4=
@@ -1390,6 +1391,7 @@ go.uber.org/atomic v1.5.1/go.mod h1:sABNBOSYdrvTF6hTgEIbc7YasKWGhgEQZyfxyTvoXHQ=
 go.uber.org/atomic v1.6.0/go.mod h1:sABNBOSYdrvTF6hTgEIbc7YasKWGhgEQZyfxyTvoXHQ=
 go.uber.org/atomic v1.7.0 h1:ADUqmZGgLDDfbSL9ZmPxKTybcoEYHgpYfELNoN+7hsw=
 go.uber.org/atomic v1.7.0/go.mod h1:fEN4uk6kAWBTFdckzkM89CLk9XfWZrxpCo0nPH17wJc=
+go.uber.org/goleak v1.1.10 h1:z+mqJhf6ss6BSfSM671tgKyZBFPTTJM+HLxnhPC3wu0=
 go.uber.org/goleak v1.1.10/go.mod h1:8a7PlsEVH3e/a/GLqe5IIrQx6GzcnRmZEufDUTk4A7A=
 go.uber.org/multierr v1.1.0/go.mod h1:wR5kodmAFQ0UK8QlbwjlSNy0Z68gJhDJUG5sjR94q/0=
 go.uber.org/multierr v1.3.0/go.mod h1:VgVr7evmIr6uPjLBxg28wmKNXyqE9akIJ5XnfpiKl+4=
@@ -2029,6 +2031,7 @@ honnef.co/go/tools v0.0.0-20190418001031-e561f6794a2a/go.mod h1:rf3lG4BRIbNafJWh
 honnef.co/go/tools v0.0.0-20190523083050-ea95bdfd59fc/go.mod h1:rf3lG4BRIbNafJWhAfAdb/ePZxsR/4RtNHQocxwk9r4=
 honnef.co/go/tools v0.0.1-2019.2.3/go.mod h1:a3bituU0lyd329TUQxRnasdCoJDkEUEAqEt0JzvZhAg=
 honnef.co/go/tools v0.0.1-2020.1.3/go.mod h1:X/FiERA/W4tHapMX5mGpAtMSVEeEUOyHaw9vFzvIQ3k=
+honnef.co/go/tools v0.0.1-2020.1.4 h1:UoveltGrhghAA7ePc+e+QYDHXrBps2PqFZiHkGR/xK8=
 honnef.co/go/tools v0.0.1-2020.1.4/go.mod h1:X/FiERA/W4tHapMX5mGpAtMSVEeEUOyHaw9vFzvIQ3k=
 howett.net/plist v0.0.0-20181124034731-591f970eefbb/go.mod h1:vMygbs4qMhSZSc4lCUl2OEE+rDiIIJAIdR4m7MiMcm0=
 k8s.io/api v0.0.0-20210217171935-8e2decd92398/go.mod h1:60tmSUpHxGPFerNHbo/ayI2lKxvtrhbxFyXuEIWJd78=

pkg/manager/impl/execution_manager_test.go

@@ -2877,9 +2877,7 @@ func TestCreateSingleTaskExecution(t *testing.T) {
 				ResourceType: core.ResourceType_TASK,
 			},
 			AuthRole: &admin.AuthRole{
-				Method: &admin.AuthRole_KubernetesServiceAccount{
-					KubernetesServiceAccount: "foo",
-				},
+				KubernetesServiceAccount: "foo",
 			},
 		},
 		Inputs: &core.LiteralMap{

pkg/manager/impl/shared/constants.go

@@ -34,5 +34,4 @@ const (
 	MatchingAttributes    = "matching_attributes"
 	// Parent of a node execution in the node executions table
 	ParentID = "parent_id"
-	AuthRole = "auth_role"
 )

pkg/manager/impl/util/single_task_execution_test.go

@@ -223,7 +223,7 @@ func TestCreateOrGetLaunchPlan(t *testing.T) {
 	spec := admin.ExecutionSpec{
 		LaunchPlan: taskIdentifier,
 		AuthRole: &admin.AuthRole{
-			Method: &admin.AuthRole_AssumableIamRole{AssumableIamRole: "assumable_role"},
+			AssumableIamRole: "assumable_role",
 		},
 	}
 	launchPlan, err := CreateOrGetLaunchPlan(

pkg/manager/impl/validation/execution_validator.go

@@ -59,11 +59,6 @@ func ValidateExecutionRequest(ctx context.Context, request admin.ExecutionCreate
 			"Invalid reference entity resource type [%v], only [%+v] allowed",
 			request.Spec.LaunchPlan.ResourceType, acceptedReferenceLaunchTypes)
 	}
-	if request.Spec.LaunchPlan.ResourceType == core.ResourceType_TASK {
-		if err := validateLaunchSingleTaskExecutionReq(request); err != nil {
-			return err
-		}
-	}
 	if err := validateLiteralMap(request.Inputs, shared.Inputs); err != nil {
 		return err
 	}
@@ -162,12 +157,3 @@ func ValidateWorkflowExecutionIdentifier(identifier *core.WorkflowExecutionIdent
 	}
 	return nil
 }
-
-// Because single task executions don't use launch plans, some parameters that are optional overrides for conventional
-// ExecutionCreateRequests are actually mandatory.
-func validateLaunchSingleTaskExecutionReq(request admin.ExecutionCreateRequest) error {
-	if request.Spec.AuthRole == nil || request.Spec.AuthRole.GetMethod() == nil {
-		return shared.GetMissingArgumentError(shared.AuthRole)
-	}
-	return nil
-}

pkg/manager/impl/validation/execution_validator_test.go

@@ -65,22 +65,6 @@ func TestValidateExecInvalidProjectAndDomain(t *testing.T) {
 	assert.EqualError(t, err, "failed to validate that project [project] and domain [domain] are registered, err: [foo]")
 }
 
-func TestValidateSingleTaskExecution(t *testing.T) {
-	request := testutils.GetExecutionRequest()
-	request.Spec.LaunchPlan.ResourceType = core.ResourceType_TASK
-
-	err := ValidateExecutionRequest(context.Background(), request, testutils.GetRepoWithDefaultProject(), execConfig)
-	assert.EqualError(t, err, "missing auth_role")
-
-	request.Spec.AuthRole = &admin.AuthRole{
-		Method: &admin.AuthRole_KubernetesServiceAccount{
-			KubernetesServiceAccount: "foo",
-		},
-	}
-	err = ValidateExecutionRequest(context.Background(), request, testutils.GetRepoWithDefaultProject(), execConfig)
-	assert.Nil(t, err)
-}
-
 func TestGetExecutionInputs(t *testing.T) {
 	executionRequest := testutils.GetExecutionRequest()
 	lpRequest := testutils.GetLaunchPlanRequest()

pkg/workflowengine/impl/propeller_executor.go

@@ -79,7 +79,8 @@ func (c *FlytePropeller) addPermissions(launchPlan admin.LaunchPlan, flyteWf *v1
 	} else if len(launchPlan.GetSpec().GetRole()) > 0 {
 		// Although deprecated, older launch plans may reference the role field instead of the Auth AssumableIamRole.
 		role = launchPlan.GetSpec().GetRole()
-	} else if launchPlan.GetSpec().GetAuthRole() != nil && len(launchPlan.GetSpec().GetAuthRole().GetKubernetesServiceAccount()) > 0 {
+	}
+	if launchPlan.GetSpec().GetAuthRole() != nil && len(launchPlan.GetSpec().GetAuthRole().GetKubernetesServiceAccount()) > 0 {
 		flyteWf.ServiceAccountName = launchPlan.GetSpec().GetAuthRole().GetKubernetesServiceAccount()
 	} else if launchPlan.GetSpec().GetAuth() != nil && len(launchPlan.GetSpec().GetAuth().GetKubernetesServiceAccount()) > 0 {
 		flyteWf.ServiceAccountName = launchPlan.GetSpec().GetAuth().GetKubernetesServiceAccount()
@@ -217,7 +218,8 @@ func (c *FlytePropeller) ExecuteTask(ctx context.Context, input interfaces.Execu
 			flyteWf.Annotations = map[string]string{}
 		}
 		flyteWf.Annotations[c.roleNameKey] = role
-	} else if input.Auth != nil && len(input.Auth.GetKubernetesServiceAccount()) > 0 {
+	}
+	if input.Auth != nil && len(input.Auth.GetKubernetesServiceAccount()) > 0 {
 		flyteWf.ServiceAccountName = input.Auth.GetKubernetesServiceAccount()
 	}
 

pkg/workflowengine/impl/propeller_executor_test.go

@@ -435,9 +435,7 @@ func TestAddPermissions(t *testing.T) {
 	propeller.addPermissions(admin.LaunchPlan{
 		Spec: &admin.LaunchPlanSpec{
 			Auth: &admin.Auth{
-				Method: &admin.Auth_AssumableIamRole{
-					AssumableIamRole: "rollie-pollie",
-				},
+				AssumableIamRole: "rollie-pollie",
 			},
 			Role: "ignore-me",
 		},
@@ -460,12 +458,13 @@ func TestAddPermissions(t *testing.T) {
 	propeller.addPermissions(admin.LaunchPlan{
 		Spec: &admin.LaunchPlanSpec{
 			Auth: &admin.Auth{
-				Method: &admin.Auth_KubernetesServiceAccount{
-					KubernetesServiceAccount: "service-account",
-				},
+				KubernetesServiceAccount: "service-account",
+				AssumableIamRole:         "rollie-pollie",
 			},
 		},
 	}, &flyteWf)
 	assert.Equal(t, "service-account", flyteWf.ServiceAccountName)
-	assert.Empty(t, flyteWf.Annotations)
+	assert.EqualValues(t, flyteWf.Annotations, map[string]string{
+		roleNameKey: "rollie-pollie",
+	})
 }