Store resolved security context in execution spec model (flyteorg#400)

Signed-off-by: Katrina Rogan <[email protected]>
eapolinario · Apr 12, 2022 · a7f65ea · a7f65ea
1 parent 0730c0e
commit a7f65ea
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 0 deletions.
diff --git a/flyteadmin/pkg/manager/impl/execution_manager.go b/flyteadmin/pkg/manager/impl/execution_manager.go
@@ -743,6 +743,7 @@ func (m *ExecutionManager) launchSingleTaskExecution(
 		Cluster:               execInfo.Cluster,
 		InputsURI:             inputsURI,
 		UserInputsURI:         userInputsURI,
+		SecurityContext:       resolvedSecurityCtx,
 	})
 	if err != nil {
 		logger.Infof(ctx, "Failed to create execution model in transformer for id: [%+v] with err: %v",
@@ -983,6 +984,7 @@ func (m *ExecutionManager) launchExecutionAndPrepareModel(
 		Cluster:               execInfo.Cluster,
 		InputsURI:             inputsURI,
 		UserInputsURI:         userInputsURI,
+		SecurityContext:       resolvedSecurityCtx,
 	})
 	if err != nil {
 		logger.Infof(ctx, "Failed to create execution model in transformer for id: [%+v] with err: %v",

diff --git a/flyteadmin/pkg/repositories/transformers/execution.go b/flyteadmin/pkg/repositories/transformers/execution.go
@@ -39,6 +39,7 @@ type CreateExecutionModelInput struct {
 	Cluster               string
 	InputsURI             storage.DataReference
 	UserInputsURI         storage.DataReference
+	SecurityContext       *core.SecurityContext
 }
 
 // CreateExecutionModel transforms a ExecutionCreateRequest to a Execution model
@@ -50,6 +51,7 @@ func CreateExecutionModel(input CreateExecutionModelInput) (*models.Execution, e
 	requestSpec.Metadata.SystemMetadata = &admin.SystemMetadata{
 		ExecutionCluster: input.Cluster,
 	}
+	requestSpec.SecurityContext = input.SecurityContext
 	spec, err := proto.Marshal(requestSpec)
 	if err != nil {
 		return nil, errors.NewFlyteAdminErrorf(codes.Internal, "Failed to serialize execution spec: %v", err)

diff --git a/flyteadmin/pkg/repositories/transformers/execution_test.go b/flyteadmin/pkg/repositories/transformers/execution_test.go
@@ -68,6 +68,11 @@ func TestCreateExecutionModel(t *testing.T) {
 	}
 
 	cluster := "cluster"
+	securityCtx := &core.SecurityContext{
+		RunAs: &core.Identity{
+			IamRole: "iam_role",
+		},
+	}
 	execution, err := CreateExecutionModel(CreateExecutionModelInput{
 		WorkflowExecutionID: core.WorkflowExecutionIdentifier{
 			Project: "project",
@@ -83,6 +88,7 @@ func TestCreateExecutionModel(t *testing.T) {
 		ParentNodeExecutionID: nodeID,
 		SourceExecutionID:     sourceID,
 		Cluster:               cluster,
+		SecurityContext:       securityCtx,
 	})
 	assert.NoError(t, err)
 	assert.Equal(t, "project", execution.Project)
@@ -100,6 +106,7 @@ func TestCreateExecutionModel(t *testing.T) {
 	expectedSpec.Metadata.SystemMetadata = &admin.SystemMetadata{
 		ExecutionCluster: cluster,
 	}
+	expectedSpec.SecurityContext = securityCtx
 	expectedSpecBytes, _ := proto.Marshal(expectedSpec)
 	assert.Equal(t, expectedSpecBytes, execution.Spec)
 	assert.Equal(t, execution.User, principal)