Using security context during registration (#328)

Signed-off-by: Prafulla Mahindrakar <[email protected]>
flyteorg · Jun 4, 2022 · 096e377 · 096e377
1 parent b092146
commit 096e377
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 1 deletion.
diff --git a/flytectl/cmd/register/register_util.go b/flytectl/cmd/register/register_util.go
@@ -402,8 +402,14 @@ func hydrateLaunchPlanSpec(configAssumableIamRole string, configK8sServiceAccoun
 	outputLocationPrefix := len(configOutputLocationPrefix) > 0
 	if assumableIamRole || k8sServiceAcct {
 		lpSpec.AuthRole = &admin.AuthRole{
-			AssumableIamRole:         configAssumableIamRole,
 			KubernetesServiceAccount: configK8sServiceAccount,
+			AssumableIamRole:         configAssumableIamRole,
+		}
+		lpSpec.SecurityContext = &core.SecurityContext{
+			RunAs: &core.Identity{
+				IamRole:           configAssumableIamRole,
+				K8SServiceAccount: configK8sServiceAccount,
+			},
 		}
 	}
 	if outputLocationPrefix {

diff --git a/flytectl/cmd/register/register_util_test.go b/flytectl/cmd/register/register_util_test.go
@@ -347,6 +347,7 @@ func TestHydrateLaunchPlanSpec(t *testing.T) {
 		err := hydrateLaunchPlanSpec(rconfig.DefaultFilesConfig.AssumableIamRole, rconfig.DefaultFilesConfig.K8sServiceAccount, rconfig.DefaultFilesConfig.OutputLocationPrefix, lpSpec)
 		assert.Nil(t, err)
 		assert.Equal(t, &admin.AuthRole{AssumableIamRole: "iamRole"}, lpSpec.AuthRole)
+		assert.Equal(t, &core.SecurityContext{RunAs: &core.Identity{IamRole: "iamRole"}}, lpSpec.SecurityContext)
 	})
 	t.Run("k8sService account override", func(t *testing.T) {
 		registerFilesSetup()
@@ -355,6 +356,7 @@ func TestHydrateLaunchPlanSpec(t *testing.T) {
 		err := hydrateLaunchPlanSpec(rconfig.DefaultFilesConfig.AssumableIamRole, rconfig.DefaultFilesConfig.K8sServiceAccount, rconfig.DefaultFilesConfig.OutputLocationPrefix, lpSpec)
 		assert.Nil(t, err)
 		assert.Equal(t, &admin.AuthRole{KubernetesServiceAccount: "k8Account"}, lpSpec.AuthRole)
+		assert.Equal(t, &core.SecurityContext{RunAs: &core.Identity{K8SServiceAccount: "k8Account"}}, lpSpec.SecurityContext)
 	})
 	t.Run("Both k8sService and IamRole", func(t *testing.T) {
 		registerFilesSetup()
@@ -365,6 +367,7 @@ func TestHydrateLaunchPlanSpec(t *testing.T) {
 		assert.Nil(t, err)
 		assert.Equal(t, &admin.AuthRole{AssumableIamRole: "iamRole",
 			KubernetesServiceAccount: "k8Account"}, lpSpec.AuthRole)
+		assert.Equal(t, &core.SecurityContext{RunAs: &core.Identity{IamRole: "iamRole", K8SServiceAccount: "k8Account"}}, lpSpec.SecurityContext)
 	})
 	t.Run("Output prefix", func(t *testing.T) {
 		registerFilesSetup()