From 9815ba4b4eaad77ebe6c39a54f1d57b73d7b02e0 Mon Sep 17 00:00:00 2001 From: ttitsworth-fsp Date: Tue, 28 Oct 2025 15:42:32 -0700 Subject: [PATCH 1/7] create user annotations feature in flyteadmin Signed-off-by: ttitsworth-fsp --- charts/flyte-core/values.yaml | 2 + flyteadmin/flyteadmin_config.yaml | 2 + .../pkg/manager/impl/execution_manager.go | 35 ++++++++++++ .../manager/impl/execution_manager_test.go | 53 +++++++++++++++++++ .../interfaces/application_configuration.go | 14 +++++ 5 files changed, 106 insertions(+) diff --git a/charts/flyte-core/values.yaml b/charts/flyte-core/values.yaml index e5699d693c..ff22ce8dfe 100755 --- a/charts/flyte-core/values.yaml +++ b/charts/flyte-core/values.yaml @@ -975,6 +975,8 @@ configmap: - "metadata" - "admin" eventVersion: 2 + injectUserAnnotations: false + userAnnotationPrefix: "flyte.ai/user-" testing: host: http://flyteadmin diff --git a/flyteadmin/flyteadmin_config.yaml b/flyteadmin/flyteadmin_config.yaml index 693e290b2a..dc7614d603 100644 --- a/flyteadmin/flyteadmin_config.yaml +++ b/flyteadmin/flyteadmin_config.yaml @@ -61,6 +61,8 @@ flyteadmin: - "metadata" - "admin" useOffloadedWorkflowClosure: false + injectUserAnnotations: false + userAnnotationPrefix: "flyte.ai/user-" database: postgres: port: 30001 diff --git a/flyteadmin/pkg/manager/impl/execution_manager.go b/flyteadmin/pkg/manager/impl/execution_manager.go index 214dfca120..0b83251402 100644 --- a/flyteadmin/pkg/manager/impl/execution_manager.go +++ b/flyteadmin/pkg/manager/impl/execution_manager.go @@ -590,6 +590,8 @@ func (m *ExecutionManager) launchSingleTaskExecution( annotations = executionConfig.GetAnnotations().GetValues() } + annotations = m.addUserAnnotations(ctx, annotations) + var rawOutputDataConfig *admin.RawOutputDataConfig if executionConfig.GetRawOutputDataConfig() != nil { rawOutputDataConfig = executionConfig.GetRawOutputDataConfig() @@ -1025,6 +1027,9 @@ func (m *ExecutionManager) launchExecution( if err != nil { return nil, nil, nil, err } + + annotations = m.addUserAnnotations(ctx, annotations) + var rawOutputDataConfig *admin.RawOutputDataConfig if executionConfig.GetRawOutputDataConfig() != nil { rawOutputDataConfig = executionConfig.GetRawOutputDataConfig() @@ -2050,6 +2055,36 @@ func (m *ExecutionManager) addProjectLabels(ctx context.Context, projectName str return initialLabels, nil } +// addUserAnnotations automatically injects user identity information as annotations when enabled in config. +// This allows tracking which user submitted each workflow execution and enables user-based authorization. +func (m *ExecutionManager) addUserAnnotations(ctx context.Context, initialAnnotations map[string]string) map[string]string { + // Check if user annotation injection is enabled + if !m.config.ApplicationConfiguration().GetTopLevelConfig().GetInjectUserAnnotations() { + return initialAnnotations + } + + // Get user identity from authentication context + principal := getUser(ctx) + if principal == "" { + // If no user context exists, return annotations unchanged + logger.Debugf(ctx, "No user principal found in context, skipping user annotation injection") + return initialAnnotations + } + + if initialAnnotations == nil { + initialAnnotations = make(map[string]string) + } + + prefix := m.config.ApplicationConfiguration().GetTopLevelConfig().GetUserAnnotationPrefix() + principalKey := prefix + "principal" + if _, exists := initialAnnotations[principalKey]; !exists { + initialAnnotations[principalKey] = principal + logger.Debugf(ctx, "Injected user annotation %s=%s", principalKey, principal) + } + + return initialAnnotations +} + func addStateFilter(filters []common.InlineFilter) ([]common.InlineFilter, error) { var stateFilterExists bool for _, inlineFilter := range filters { diff --git a/flyteadmin/pkg/manager/impl/execution_manager_test.go b/flyteadmin/pkg/manager/impl/execution_manager_test.go index 78cb8435c5..1e68ec53c2 100644 --- a/flyteadmin/pkg/manager/impl/execution_manager_test.go +++ b/flyteadmin/pkg/manager/impl/execution_manager_test.go @@ -6332,3 +6332,56 @@ func TestQueryTemplate(t *testing.T) { assert.Error(t, err) }) } + +func TestAddUserAnnotations(t *testing.T) { + principal := "test-user@example.com" + + t.Run("enabled with user context", func(t *testing.T) { + mockConfig := runtimeMocks.NewMockConfigurationProvider( + testutils.GetApplicationConfigWithDefaultDomains(), nil, nil, nil, nil, nil) + mockConfig.ApplicationConfiguration().GetTopLevelConfig().InjectUserAnnotations = true + mockConfig.ApplicationConfiguration().GetTopLevelConfig().UserAnnotationPrefix = "flyte.ai/user-" + + manager := ExecutionManager{config: mockConfig} + + identity, err := auth.NewIdentityContext("", principal, "", time.Now(), sets.NewString(), nil, nil) + assert.NoError(t, err) + ctx := identity.WithContext(context.Background()) + + result := manager.addUserAnnotations(ctx, map[string]string{"existing": "value"}) + + assert.Equal(t, "test-user@example.com", result["flyte.ai/user-principal"]) + assert.Equal(t, "value", result["existing"]) + }) + + t.Run("disabled", func(t *testing.T) { + mockConfig := runtimeMocks.NewMockConfigurationProvider( + testutils.GetApplicationConfigWithDefaultDomains(), nil, nil, nil, nil, nil) + mockConfig.ApplicationConfiguration().GetTopLevelConfig().InjectUserAnnotations = false + + manager := ExecutionManager{config: mockConfig} + + identity, err := auth.NewIdentityContext("", principal, "", time.Now(), sets.NewString(), nil, nil) + assert.NoError(t, err) + ctx := identity.WithContext(context.Background()) + + result := manager.addUserAnnotations(ctx, map[string]string{"existing": "value"}) + + assert.NotContains(t, result, "flyte.ai/user-principal") + assert.Equal(t, "value", result["existing"]) + }) + + t.Run("no user context", func(t *testing.T) { + mockConfig := runtimeMocks.NewMockConfigurationProvider( + testutils.GetApplicationConfigWithDefaultDomains(), nil, nil, nil, nil, nil) + mockConfig.ApplicationConfiguration().GetTopLevelConfig().InjectUserAnnotations = true + + manager := ExecutionManager{config: mockConfig} + ctx := context.Background() + + result := manager.addUserAnnotations(ctx, map[string]string{"existing": "value"}) + + assert.NotContains(t, result, "flyte.ai/user-principal") + assert.Equal(t, "value", result["existing"]) + }) +} diff --git a/flyteadmin/pkg/runtime/interfaces/application_configuration.go b/flyteadmin/pkg/runtime/interfaces/application_configuration.go index 458de59381..377935e511 100644 --- a/flyteadmin/pkg/runtime/interfaces/application_configuration.go +++ b/flyteadmin/pkg/runtime/interfaces/application_configuration.go @@ -116,6 +116,9 @@ type ApplicationConfig struct { // Enabling this will instruct operator to use storage (s3/gcs/etc) to offload workflow execution inputs instead of storing them inline in the CRD. UseOffloadedInputs bool `json:"useOffloadedInputs" pflag:",Use offloaded inputs for workflows."` + + InjectUserAnnotations bool `json:"injectUserAnnotations"` + UserAnnotationPrefix string `json:"userAnnotationPrefix"` } func (a *ApplicationConfig) GetRoleNameKey() string { @@ -201,6 +204,17 @@ func (a *ApplicationConfig) GetEnvs() *admin.Envs { } } +func (a *ApplicationConfig) GetInjectUserAnnotations() bool { + return a.InjectUserAnnotations +} + +func (a *ApplicationConfig) GetUserAnnotationPrefix() string { + if a.UserAnnotationPrefix == "" { + return "flyte.ai/user-" + } + return a.UserAnnotationPrefix +} + // GetAsWorkflowExecutionConfig returns the WorkflowExecutionConfig as extracted from this object func (a *ApplicationConfig) GetAsWorkflowExecutionConfig() *admin.WorkflowExecutionConfig { // These values should always be set as their fallback values equals to their zero value or nil, From 9cc240ab9d055400256cc6bdb8c99f3d6ec9daf7 Mon Sep 17 00:00:00 2001 From: ttitsworth-fsp Date: Wed, 29 Oct 2025 08:49:33 -0700 Subject: [PATCH 2/7] modify prefix and only use user email for annotation Signed-off-by: ttitsworth-fsp --- charts/flyte-core/values.yaml | 2 +- flyteadmin/flyteadmin_config.yaml | 2 +- .../pkg/manager/impl/execution_manager.go | 19 ++++++---- .../manager/impl/execution_manager_test.go | 35 +++++++++++++++---- .../interfaces/application_configuration.go | 2 +- 5 files changed, 44 insertions(+), 16 deletions(-) diff --git a/charts/flyte-core/values.yaml b/charts/flyte-core/values.yaml index ff22ce8dfe..f1da7d156b 100755 --- a/charts/flyte-core/values.yaml +++ b/charts/flyte-core/values.yaml @@ -976,7 +976,7 @@ configmap: - "admin" eventVersion: 2 injectUserAnnotations: false - userAnnotationPrefix: "flyte.ai/user-" + userAnnotationPrefix: "flyte.ai" testing: host: http://flyteadmin diff --git a/flyteadmin/flyteadmin_config.yaml b/flyteadmin/flyteadmin_config.yaml index dc7614d603..346e0bd92e 100644 --- a/flyteadmin/flyteadmin_config.yaml +++ b/flyteadmin/flyteadmin_config.yaml @@ -62,7 +62,7 @@ flyteadmin: - "admin" useOffloadedWorkflowClosure: false injectUserAnnotations: false - userAnnotationPrefix: "flyte.ai/user-" + userAnnotationPrefix: "flyte.ai" database: postgres: port: 30001 diff --git a/flyteadmin/pkg/manager/impl/execution_manager.go b/flyteadmin/pkg/manager/impl/execution_manager.go index 0b83251402..ee36ecf1c6 100644 --- a/flyteadmin/pkg/manager/impl/execution_manager.go +++ b/flyteadmin/pkg/manager/impl/execution_manager.go @@ -2064,10 +2064,15 @@ func (m *ExecutionManager) addUserAnnotations(ctx context.Context, initialAnnota } // Get user identity from authentication context - principal := getUser(ctx) + identityContext := auth.IdentityContextFromContext(ctx) + var principal string + if identityContext.UserInfo() != nil { + principal = identityContext.UserInfo().GetEmail() + } + if principal == "" { - // If no user context exists, return annotations unchanged - logger.Debugf(ctx, "No user principal found in context, skipping user annotation injection") + // If no email is available, skip annotation injection + logger.Debugf(ctx, "No user email found in context, skipping user annotation injection") return initialAnnotations } @@ -2076,10 +2081,10 @@ func (m *ExecutionManager) addUserAnnotations(ctx context.Context, initialAnnota } prefix := m.config.ApplicationConfiguration().GetTopLevelConfig().GetUserAnnotationPrefix() - principalKey := prefix + "principal" - if _, exists := initialAnnotations[principalKey]; !exists { - initialAnnotations[principalKey] = principal - logger.Debugf(ctx, "Injected user annotation %s=%s", principalKey, principal) + userKey := prefix + "/user" + if _, exists := initialAnnotations[userKey]; !exists { + initialAnnotations[userKey] = principal + logger.Debugf(ctx, "Injected user annotation %s=%s", userKey, principal) } return initialAnnotations diff --git a/flyteadmin/pkg/manager/impl/execution_manager_test.go b/flyteadmin/pkg/manager/impl/execution_manager_test.go index 1e68ec53c2..5d341a28db 100644 --- a/flyteadmin/pkg/manager/impl/execution_manager_test.go +++ b/flyteadmin/pkg/manager/impl/execution_manager_test.go @@ -51,6 +51,7 @@ import ( "github.com/flyteorg/flyte/flyteidl/gen/pb-go/flyteidl/admin" "github.com/flyteorg/flyte/flyteidl/gen/pb-go/flyteidl/core" "github.com/flyteorg/flyte/flyteidl/gen/pb-go/flyteidl/event" + "github.com/flyteorg/flyte/flyteidl/gen/pb-go/flyteidl/service" mockScope "github.com/flyteorg/flyte/flytestdlib/promutils" "github.com/flyteorg/flyte/flytestdlib/storage" ) @@ -6340,17 +6341,18 @@ func TestAddUserAnnotations(t *testing.T) { mockConfig := runtimeMocks.NewMockConfigurationProvider( testutils.GetApplicationConfigWithDefaultDomains(), nil, nil, nil, nil, nil) mockConfig.ApplicationConfiguration().GetTopLevelConfig().InjectUserAnnotations = true - mockConfig.ApplicationConfiguration().GetTopLevelConfig().UserAnnotationPrefix = "flyte.ai/user-" + mockConfig.ApplicationConfiguration().GetTopLevelConfig().UserAnnotationPrefix = "flyte.ai" manager := ExecutionManager{config: mockConfig} - identity, err := auth.NewIdentityContext("", principal, "", time.Now(), sets.NewString(), nil, nil) + userInfo := &service.UserInfoResponse{Email: principal} + identity, err := auth.NewIdentityContext("", "user-id-123", "", time.Now(), sets.NewString(), userInfo, nil) assert.NoError(t, err) ctx := identity.WithContext(context.Background()) result := manager.addUserAnnotations(ctx, map[string]string{"existing": "value"}) - assert.Equal(t, "test-user@example.com", result["flyte.ai/user-principal"]) + assert.Equal(t, "test-user@example.com", result["flyte.ai/user"]) assert.Equal(t, "value", result["existing"]) }) @@ -6361,13 +6363,14 @@ func TestAddUserAnnotations(t *testing.T) { manager := ExecutionManager{config: mockConfig} - identity, err := auth.NewIdentityContext("", principal, "", time.Now(), sets.NewString(), nil, nil) + userInfo := &service.UserInfoResponse{Email: principal} + identity, err := auth.NewIdentityContext("", "user-id-123", "", time.Now(), sets.NewString(), userInfo, nil) assert.NoError(t, err) ctx := identity.WithContext(context.Background()) result := manager.addUserAnnotations(ctx, map[string]string{"existing": "value"}) - assert.NotContains(t, result, "flyte.ai/user-principal") + assert.NotContains(t, result, "flyte.ai/user") assert.Equal(t, "value", result["existing"]) }) @@ -6381,7 +6384,27 @@ func TestAddUserAnnotations(t *testing.T) { result := manager.addUserAnnotations(ctx, map[string]string{"existing": "value"}) - assert.NotContains(t, result, "flyte.ai/user-principal") + assert.NotContains(t, result, "flyte.ai/user") + assert.Equal(t, "value", result["existing"]) + }) + + t.Run("enabled but no email", func(t *testing.T) { + mockConfig := runtimeMocks.NewMockConfigurationProvider( + testutils.GetApplicationConfigWithDefaultDomains(), nil, nil, nil, nil, nil) + mockConfig.ApplicationConfiguration().GetTopLevelConfig().InjectUserAnnotations = true + mockConfig.ApplicationConfiguration().GetTopLevelConfig().UserAnnotationPrefix = "flyte.ai" + + manager := ExecutionManager{config: mockConfig} + + // UserInfo with no email set + userInfo := &service.UserInfoResponse{} + identity, err := auth.NewIdentityContext("", "user-id-123", "", time.Now(), sets.NewString(), userInfo, nil) + assert.NoError(t, err) + ctx := identity.WithContext(context.Background()) + + result := manager.addUserAnnotations(ctx, map[string]string{"existing": "value"}) + + assert.NotContains(t, result, "flyte.ai/user") assert.Equal(t, "value", result["existing"]) }) } diff --git a/flyteadmin/pkg/runtime/interfaces/application_configuration.go b/flyteadmin/pkg/runtime/interfaces/application_configuration.go index 377935e511..35ae4ebea2 100644 --- a/flyteadmin/pkg/runtime/interfaces/application_configuration.go +++ b/flyteadmin/pkg/runtime/interfaces/application_configuration.go @@ -210,7 +210,7 @@ func (a *ApplicationConfig) GetInjectUserAnnotations() bool { func (a *ApplicationConfig) GetUserAnnotationPrefix() string { if a.UserAnnotationPrefix == "" { - return "flyte.ai/user-" + return "flyte.ai" } return a.UserAnnotationPrefix } From be1863d16af2ef6c16854f705da8f2011814c63e Mon Sep 17 00:00:00 2001 From: ttitsworth-fsp Date: Wed, 29 Oct 2025 11:31:11 -0700 Subject: [PATCH 3/7] update tests to improve coverage Signed-off-by: ttitsworth-fsp --- .../manager/impl/execution_manager_test.go | 57 +++++++++++++++++++ 1 file changed, 57 insertions(+) diff --git a/flyteadmin/pkg/manager/impl/execution_manager_test.go b/flyteadmin/pkg/manager/impl/execution_manager_test.go index 5d341a28db..a776579675 100644 --- a/flyteadmin/pkg/manager/impl/execution_manager_test.go +++ b/flyteadmin/pkg/manager/impl/execution_manager_test.go @@ -6407,4 +6407,61 @@ func TestAddUserAnnotations(t *testing.T) { assert.NotContains(t, result, "flyte.ai/user") assert.Equal(t, "value", result["existing"]) }) + + t.Run("enabled with nil annotations map", func(t *testing.T) { + mockConfig := runtimeMocks.NewMockConfigurationProvider( + testutils.GetApplicationConfigWithDefaultDomains(), nil, nil, nil, nil, nil) + mockConfig.ApplicationConfiguration().GetTopLevelConfig().InjectUserAnnotations = true + mockConfig.ApplicationConfiguration().GetTopLevelConfig().UserAnnotationPrefix = "flyte.ai" + + manager := ExecutionManager{config: mockConfig} + + userInfo := &service.UserInfoResponse{Email: principal} + identity, err := auth.NewIdentityContext("", "user-id-123", "", time.Now(), sets.NewString(), userInfo, nil) + assert.NoError(t, err) + ctx := identity.WithContext(context.Background()) + + result := manager.addUserAnnotations(ctx, nil) + + assert.Equal(t, "test-user@example.com", result["flyte.ai/user"]) + }) + + t.Run("annotation already exists", func(t *testing.T) { + mockConfig := runtimeMocks.NewMockConfigurationProvider( + testutils.GetApplicationConfigWithDefaultDomains(), nil, nil, nil, nil, nil) + mockConfig.ApplicationConfiguration().GetTopLevelConfig().InjectUserAnnotations = true + mockConfig.ApplicationConfiguration().GetTopLevelConfig().UserAnnotationPrefix = "flyte.ai" + + manager := ExecutionManager{config: mockConfig} + + userInfo := &service.UserInfoResponse{Email: principal} + identity, err := auth.NewIdentityContext("", "user-id-123", "", time.Now(), sets.NewString(), userInfo, nil) + assert.NoError(t, err) + ctx := identity.WithContext(context.Background()) + + result := manager.addUserAnnotations(ctx, map[string]string{"flyte.ai/user": "existing@example.com"}) + + // Should preserve existing annotation value + assert.Equal(t, "existing@example.com", result["flyte.ai/user"]) + }) + + t.Run("uses default prefix when not configured", func(t *testing.T) { + mockConfig := runtimeMocks.NewMockConfigurationProvider( + testutils.GetApplicationConfigWithDefaultDomains(), nil, nil, nil, nil, nil) + mockConfig.ApplicationConfiguration().GetTopLevelConfig().InjectUserAnnotations = true + mockConfig.ApplicationConfiguration().GetTopLevelConfig().UserAnnotationPrefix = "" + + manager := ExecutionManager{config: mockConfig} + + userInfo := &service.UserInfoResponse{Email: principal} + identity, err := auth.NewIdentityContext("", "user-id-123", "", time.Now(), sets.NewString(), userInfo, nil) + assert.NoError(t, err) + ctx := identity.WithContext(context.Background()) + + result := manager.addUserAnnotations(ctx, map[string]string{"existing": "value"}) + + // Should use default prefix "flyte.ai" + assert.Equal(t, "test-user@example.com", result["flyte.ai/user"]) + assert.Equal(t, "value", result["existing"]) + }) } From 160f952401c83637cb2079051eb69f0c2bedfaaa Mon Sep 17 00:00:00 2001 From: ttitsworth-fsp Date: Thu, 30 Oct 2025 11:19:47 -0700 Subject: [PATCH 4/7] generate helm manually Signed-off-by: ttitsworth-fsp --- charts/flyte-core/README.md | 2 +- .../eks/flyte_aws_scheduler_helm_generated.yaml | 4 +++- deployment/eks/flyte_helm_controlplane_generated.yaml | 6 ++++-- deployment/eks/flyte_helm_generated.yaml | 6 ++++-- deployment/gcp/flyte_helm_controlplane_generated.yaml | 6 ++++-- deployment/gcp/flyte_helm_generated.yaml | 6 ++++-- deployment/sandbox/flyte_helm_generated.yaml | 11 +++++++---- .../sandbox-bundled/manifests/complete-connector.yaml | 4 ++-- docker/sandbox-bundled/manifests/complete.yaml | 4 ++-- docker/sandbox-bundled/manifests/dev.yaml | 4 ++-- 10 files changed, 33 insertions(+), 20 deletions(-) diff --git a/charts/flyte-core/README.md b/charts/flyte-core/README.md index cf32bbdc8f..87b7c42558 100644 --- a/charts/flyte-core/README.md +++ b/charts/flyte-core/README.md @@ -103,7 +103,7 @@ helm install gateway bitnami/contour -n flyte | common.ingress.tls | object | `{"enabled":false}` | - Ingress hostname host: | | common.ingress.webpackHMR | bool | `false` | - Enable or disable HMR route to flyteconsole. This is useful only for frontend development. | | configmap.admin | object | `{"admin":{"clientId":"{{ .Values.secrets.adminOauthClientCredentials.clientId }}","clientSecretLocation":"/etc/secrets/client_secret","endpoint":"flyteadmin:81","insecure":true},"event":{"capacity":1000,"rate":500,"type":"admin"}}` | Admin Client configuration [structure](https://pkg.go.dev/github.com/flyteorg/flytepropeller/pkg/controller/nodes/subworkflow/launchplan#AdminConfig) | -| configmap.adminServer | object | `{"auth":{"appAuth":{"thirdPartyConfig":{"flyteClient":{"clientId":"flytectl","redirectUri":"http://localhost:53593/callback","scopes":["offline","all"]}}},"authorizedUris":["https://localhost:30081","http://flyteadmin:80","http://flyteadmin.flyte.svc.cluster.local:80"],"userAuth":{"openId":{"baseUrl":"https://accounts.google.com","clientId":"657465813211-6eog7ek7li5k7i7fvgv2921075063hpe.apps.googleusercontent.com","scopes":["profile","openid"]}}},"flyteadmin":{"eventVersion":2,"metadataStoragePrefix":["metadata","admin"],"metricsScope":"flyte:","profilerPort":10254,"roleNameKey":"iam.amazonaws.com/role","testing":{"host":"http://flyteadmin"}},"server":{"grpc":{"port":8089},"httpPort":8088,"security":{"allowCors":true,"allowedHeaders":["Content-Type","flyte-authorization"],"allowedOrigins":["*"],"secure":false,"useAuth":false}}}` | FlyteAdmin server configuration | +| configmap.adminServer | object | `{"auth":{"appAuth":{"thirdPartyConfig":{"flyteClient":{"clientId":"flytectl","redirectUri":"http://localhost:53593/callback","scopes":["offline","all"]}}},"authorizedUris":["https://localhost:30081","http://flyteadmin:80","http://flyteadmin.flyte.svc.cluster.local:80"],"userAuth":{"openId":{"baseUrl":"https://accounts.google.com","clientId":"657465813211-6eog7ek7li5k7i7fvgv2921075063hpe.apps.googleusercontent.com","scopes":["profile","openid"]}}},"flyteadmin":{"eventVersion":2,"injectUserAnnotations":false,"metadataStoragePrefix":["metadata","admin"],"metricsScope":"flyte:","profilerPort":10254,"roleNameKey":"iam.amazonaws.com/role","testing":{"host":"http://flyteadmin"},"userAnnotationPrefix":"flyte.ai"},"server":{"grpc":{"port":8089},"httpPort":8088,"security":{"allowCors":true,"allowedHeaders":["Content-Type","flyte-authorization"],"allowedOrigins":["*"],"secure":false,"useAuth":false}}}` | FlyteAdmin server configuration | | configmap.adminServer.auth | object | `{"appAuth":{"thirdPartyConfig":{"flyteClient":{"clientId":"flytectl","redirectUri":"http://localhost:53593/callback","scopes":["offline","all"]}}},"authorizedUris":["https://localhost:30081","http://flyteadmin:80","http://flyteadmin.flyte.svc.cluster.local:80"],"userAuth":{"openId":{"baseUrl":"https://accounts.google.com","clientId":"657465813211-6eog7ek7li5k7i7fvgv2921075063hpe.apps.googleusercontent.com","scopes":["profile","openid"]}}}` | Authentication configuration | | configmap.adminServer.server.security.secure | bool | `false` | Controls whether to serve requests over SSL/TLS. | | configmap.adminServer.server.security.useAuth | bool | `false` | Controls whether to enforce authentication. Follow the guide in https://docs.flyte.org/ on how to setup authentication. | diff --git a/deployment/eks/flyte_aws_scheduler_helm_generated.yaml b/deployment/eks/flyte_aws_scheduler_helm_generated.yaml index c2ef7153a3..148c2bdcb9 100644 --- a/deployment/eks/flyte_aws_scheduler_helm_generated.yaml +++ b/deployment/eks/flyte_aws_scheduler_helm_generated.yaml @@ -166,6 +166,7 @@ data: - openid flyteadmin: eventVersion: 2 + injectUserAnnotations: false metadataStoragePrefix: - metadata - admin @@ -174,6 +175,7 @@ data: roleNameKey: iam.amazonaws.com/role testing: host: http://flyteadmin + userAnnotationPrefix: flyte.ai server: grpc: port: 8089 @@ -886,7 +888,7 @@ spec: template: metadata: annotations: - configChecksum: "6fd4bb5460f260b492db7ddd34b6011581292e88b28c2e4514b7da75673cd4d" + configChecksum: "71783b5be9ab6a2bbb2fa40b936c74b39c6dcf60d70979daede4d9449ce944d" labels: app.kubernetes.io/name: flyteadmin app.kubernetes.io/instance: flyte diff --git a/deployment/eks/flyte_helm_controlplane_generated.yaml b/deployment/eks/flyte_helm_controlplane_generated.yaml index dea39b2f76..6d1fe46e25 100644 --- a/deployment/eks/flyte_helm_controlplane_generated.yaml +++ b/deployment/eks/flyte_helm_controlplane_generated.yaml @@ -147,6 +147,7 @@ data: - openid flyteadmin: eventVersion: 2 + injectUserAnnotations: false metadataStoragePrefix: - metadata - admin @@ -155,6 +156,7 @@ data: roleNameKey: iam.amazonaws.com/role testing: host: http://flyteadmin + userAnnotationPrefix: flyte.ai server: grpc: port: 8089 @@ -583,7 +585,7 @@ spec: template: metadata: annotations: - configChecksum: "b1a6f6afb902bd1384515a97c5bad38985c0799ca8173efb0e664bda8eb9ca1" + configChecksum: "92eb8185e329f235bc30c58f192a9ab6a2840f64a379ef53fe2571bc468ab22" labels: app.kubernetes.io/name: flyteadmin app.kubernetes.io/instance: flyte @@ -1009,7 +1011,7 @@ spec: template: metadata: annotations: - configChecksum: "b1a6f6afb902bd1384515a97c5bad38985c0799ca8173efb0e664bda8eb9ca1" + configChecksum: "92eb8185e329f235bc30c58f192a9ab6a2840f64a379ef53fe2571bc468ab22" labels: app.kubernetes.io/name: flytescheduler app.kubernetes.io/instance: flyte diff --git a/deployment/eks/flyte_helm_generated.yaml b/deployment/eks/flyte_helm_generated.yaml index 01ff6447ee..155e5da119 100644 --- a/deployment/eks/flyte_helm_generated.yaml +++ b/deployment/eks/flyte_helm_generated.yaml @@ -178,6 +178,7 @@ data: - openid flyteadmin: eventVersion: 2 + injectUserAnnotations: false metadataStoragePrefix: - metadata - admin @@ -186,6 +187,7 @@ data: roleNameKey: iam.amazonaws.com/role testing: host: http://flyteadmin + userAnnotationPrefix: flyte.ai server: grpc: port: 8089 @@ -917,7 +919,7 @@ spec: template: metadata: annotations: - configChecksum: "b1a6f6afb902bd1384515a97c5bad38985c0799ca8173efb0e664bda8eb9ca1" + configChecksum: "92eb8185e329f235bc30c58f192a9ab6a2840f64a379ef53fe2571bc468ab22" labels: app.kubernetes.io/name: flyteadmin app.kubernetes.io/instance: flyte @@ -1343,7 +1345,7 @@ spec: template: metadata: annotations: - configChecksum: "b1a6f6afb902bd1384515a97c5bad38985c0799ca8173efb0e664bda8eb9ca1" + configChecksum: "92eb8185e329f235bc30c58f192a9ab6a2840f64a379ef53fe2571bc468ab22" labels: app.kubernetes.io/name: flytescheduler app.kubernetes.io/instance: flyte diff --git a/deployment/gcp/flyte_helm_controlplane_generated.yaml b/deployment/gcp/flyte_helm_controlplane_generated.yaml index e65deb9584..7ee7785d9b 100644 --- a/deployment/gcp/flyte_helm_controlplane_generated.yaml +++ b/deployment/gcp/flyte_helm_controlplane_generated.yaml @@ -147,6 +147,7 @@ data: - openid flyteadmin: eventVersion: 2 + injectUserAnnotations: false metadataStoragePrefix: - metadata - admin @@ -155,6 +156,7 @@ data: roleNameKey: iam.amazonaws.com/role testing: host: http://flyteadmin + userAnnotationPrefix: flyte.ai server: grpc: port: 8089 @@ -600,7 +602,7 @@ spec: template: metadata: annotations: - configChecksum: "e952d320a403549f597a6e5c264a4284fb2ae2e33b57c54e70975bf4f0f4f9a" + configChecksum: "20b338538d7cb4f2b765e3d06619b7ecb2cfc5730dd8ae986568c6e3ef303a1" labels: app.kubernetes.io/name: flyteadmin app.kubernetes.io/instance: flyte @@ -1026,7 +1028,7 @@ spec: template: metadata: annotations: - configChecksum: "e952d320a403549f597a6e5c264a4284fb2ae2e33b57c54e70975bf4f0f4f9a" + configChecksum: "20b338538d7cb4f2b765e3d06619b7ecb2cfc5730dd8ae986568c6e3ef303a1" labels: app.kubernetes.io/name: flytescheduler app.kubernetes.io/instance: flyte diff --git a/deployment/gcp/flyte_helm_generated.yaml b/deployment/gcp/flyte_helm_generated.yaml index c9551f530b..cf0ce63830 100644 --- a/deployment/gcp/flyte_helm_generated.yaml +++ b/deployment/gcp/flyte_helm_generated.yaml @@ -178,6 +178,7 @@ data: - openid flyteadmin: eventVersion: 2 + injectUserAnnotations: false metadataStoragePrefix: - metadata - admin @@ -186,6 +187,7 @@ data: roleNameKey: iam.amazonaws.com/role testing: host: http://flyteadmin + userAnnotationPrefix: flyte.ai server: grpc: port: 8089 @@ -942,7 +944,7 @@ spec: template: metadata: annotations: - configChecksum: "e952d320a403549f597a6e5c264a4284fb2ae2e33b57c54e70975bf4f0f4f9a" + configChecksum: "20b338538d7cb4f2b765e3d06619b7ecb2cfc5730dd8ae986568c6e3ef303a1" labels: app.kubernetes.io/name: flyteadmin app.kubernetes.io/instance: flyte @@ -1368,7 +1370,7 @@ spec: template: metadata: annotations: - configChecksum: "e952d320a403549f597a6e5c264a4284fb2ae2e33b57c54e70975bf4f0f4f9a" + configChecksum: "20b338538d7cb4f2b765e3d06619b7ecb2cfc5730dd8ae986568c6e3ef303a1" labels: app.kubernetes.io/name: flytescheduler app.kubernetes.io/instance: flyte diff --git a/deployment/sandbox/flyte_helm_generated.yaml b/deployment/sandbox/flyte_helm_generated.yaml index 1d5fa817de..286583ca71 100644 --- a/deployment/sandbox/flyte_helm_generated.yaml +++ b/deployment/sandbox/flyte_helm_generated.yaml @@ -298,6 +298,7 @@ data: - openid flyteadmin: eventVersion: 2 + injectUserAnnotations: false metadataStoragePrefix: - metadata - admin @@ -306,6 +307,7 @@ data: roleNameKey: iam.amazonaws.com/role testing: host: http://flyteadmin + userAnnotationPrefix: flyte.ai server: grpc: port: 8089 @@ -710,6 +712,7 @@ data: resource_manager.yaml: | propeller: resourcemanager: + redis: null type: noop storage.yaml: | storage: @@ -6730,7 +6733,7 @@ spec: template: metadata: annotations: - configChecksum: "29b249082ba3f15e213daf85d53d386f968925a8aeab291c585078d59680378" + configChecksum: "fe83495c82ad870691613547d3dcb8acaaec70e61a501843703aba7462d0afe" labels: app.kubernetes.io/name: flyteadmin app.kubernetes.io/instance: flyte @@ -7127,7 +7130,7 @@ spec: template: metadata: annotations: - configChecksum: "29b249082ba3f15e213daf85d53d386f968925a8aeab291c585078d59680378" + configChecksum: "fe83495c82ad870691613547d3dcb8acaaec70e61a501843703aba7462d0afe" labels: app.kubernetes.io/name: flytescheduler app.kubernetes.io/instance: flyte @@ -7222,7 +7225,7 @@ spec: template: metadata: annotations: - configChecksum: "671959f1f31dcfd1a93c1a484be7c7264f05b66de43df1a770f331d389787a4" + configChecksum: "ea603a73ffb9754ac5c6d9bab9a8e868050e0e7b118399be9c8b1cd4ce86cf1" prometheus.io/path: "/metrics" prometheus.io/port: "10254" labels: @@ -7300,7 +7303,7 @@ spec: app.kubernetes.io/name: flyte-pod-webhook app.kubernetes.io/version: v1.16.0 annotations: - configChecksum: "671959f1f31dcfd1a93c1a484be7c7264f05b66de43df1a770f331d389787a4" + configChecksum: "ea603a73ffb9754ac5c6d9bab9a8e868050e0e7b118399be9c8b1cd4ce86cf1" prometheus.io/path: "/metrics" prometheus.io/port: "10254" spec: diff --git a/docker/sandbox-bundled/manifests/complete-connector.yaml b/docker/sandbox-bundled/manifests/complete-connector.yaml index 14d0d72e38..4a23c5f74c 100644 --- a/docker/sandbox-bundled/manifests/complete-connector.yaml +++ b/docker/sandbox-bundled/manifests/complete-connector.yaml @@ -822,7 +822,7 @@ type: Opaque --- apiVersion: v1 data: - haSharedSecret: NUFZSjlMb000dURTeFRHcA== + haSharedSecret: NjlZbm9Bdmdjc2RIOVV6RA== proxyPassword: "" proxyUsername: "" kind: Secret @@ -1419,7 +1419,7 @@ spec: metadata: annotations: checksum/config: 8f50e768255a87f078ba8b9879a0c174c3e045ffb46ac8723d2eedbe293c8d81 - checksum/secret: 70ab3a722bbba034ed7344ef07a16053252026dfe34e71c2448c980ff47a6181 + checksum/secret: bcb59e2e39c2a39cd92acf4a6d1e1798831a8acbe0239e2ca35a09ad73aeccbd labels: app: docker-registry release: flyte-sandbox diff --git a/docker/sandbox-bundled/manifests/complete.yaml b/docker/sandbox-bundled/manifests/complete.yaml index 24f529b75a..ab2affa6c0 100644 --- a/docker/sandbox-bundled/manifests/complete.yaml +++ b/docker/sandbox-bundled/manifests/complete.yaml @@ -803,7 +803,7 @@ type: Opaque --- apiVersion: v1 data: - haSharedSecret: UUltSk9iWTdxVzlNTGZPaQ== + haSharedSecret: c1JIT3FrYjd5VzEwZm5Kag== proxyPassword: "" proxyUsername: "" kind: Secret @@ -1367,7 +1367,7 @@ spec: metadata: annotations: checksum/config: 8f50e768255a87f078ba8b9879a0c174c3e045ffb46ac8723d2eedbe293c8d81 - checksum/secret: 809f2778eadd51c6c0289457a88bedac6de365127677a48c24dba5611799c3b2 + checksum/secret: 76562a1293dc5289634fb1e2c6f8ff003a927fb6a82c4a0d952f3dd5869b9f87 labels: app: docker-registry release: flyte-sandbox diff --git a/docker/sandbox-bundled/manifests/dev.yaml b/docker/sandbox-bundled/manifests/dev.yaml index 394ff13bea..81a8853ab3 100644 --- a/docker/sandbox-bundled/manifests/dev.yaml +++ b/docker/sandbox-bundled/manifests/dev.yaml @@ -499,7 +499,7 @@ metadata: --- apiVersion: v1 data: - haSharedSecret: aHlIMkgzdHluSUU0T2t2bA== + haSharedSecret: a0R1QmhNY0V2b3k5UUtzNA== proxyPassword: "" proxyUsername: "" kind: Secret @@ -934,7 +934,7 @@ spec: metadata: annotations: checksum/config: 8f50e768255a87f078ba8b9879a0c174c3e045ffb46ac8723d2eedbe293c8d81 - checksum/secret: 5f9abc8cd74e0ef31cadbf5c8d52979e5fa091a0361f51dd898422abaeca6ed6 + checksum/secret: 300c4e683a97067435b018fa67798a7f44bae98404bd2538f816d97519ff0505 labels: app: docker-registry release: flyte-sandbox From d7c79019fc3f72b2c40367c7618d34180567b70a Mon Sep 17 00:00:00 2001 From: ttitsworth-fsp Date: Wed, 5 Nov 2025 10:38:00 -0800 Subject: [PATCH 5/7] address suggestions and add email and sub options and app identities Signed-off-by: ttitsworth-fsp --- charts/flyte-core/values.yaml | 7 +- flyteadmin/flyteadmin_config.yaml | 7 +- .../pkg/manager/impl/execution_manager.go | 87 ++++++++--- .../manager/impl/execution_manager_test.go | 136 +++++++++++++----- .../interfaces/application_configuration.go | 22 ++- 5 files changed, 187 insertions(+), 72 deletions(-) diff --git a/charts/flyte-core/values.yaml b/charts/flyte-core/values.yaml index f1da7d156b..3967376a6d 100755 --- a/charts/flyte-core/values.yaml +++ b/charts/flyte-core/values.yaml @@ -975,8 +975,11 @@ configmap: - "metadata" - "admin" eventVersion: 2 - injectUserAnnotations: false - userAnnotationPrefix: "flyte.ai" + injectIdentityAnnotations: false + identityAnnotationPrefix: "flyte.ai" + identityAnnotationKeys: + - email + - sub testing: host: http://flyteadmin diff --git a/flyteadmin/flyteadmin_config.yaml b/flyteadmin/flyteadmin_config.yaml index 346e0bd92e..584e065702 100644 --- a/flyteadmin/flyteadmin_config.yaml +++ b/flyteadmin/flyteadmin_config.yaml @@ -61,8 +61,11 @@ flyteadmin: - "metadata" - "admin" useOffloadedWorkflowClosure: false - injectUserAnnotations: false - userAnnotationPrefix: "flyte.ai" + injectIdentityAnnotations: false + identityAnnotationPrefix: "flyte.ai" + identityAnnotationKeys: + - email + - sub database: postgres: port: 30001 diff --git a/flyteadmin/pkg/manager/impl/execution_manager.go b/flyteadmin/pkg/manager/impl/execution_manager.go index ee36ecf1c6..c8e7925246 100644 --- a/flyteadmin/pkg/manager/impl/execution_manager.go +++ b/flyteadmin/pkg/manager/impl/execution_manager.go @@ -590,7 +590,7 @@ func (m *ExecutionManager) launchSingleTaskExecution( annotations = executionConfig.GetAnnotations().GetValues() } - annotations = m.addUserAnnotations(ctx, annotations) + annotations = m.addIdentityAnnotations(ctx, annotations) var rawOutputDataConfig *admin.RawOutputDataConfig if executionConfig.GetRawOutputDataConfig() != nil { @@ -1028,7 +1028,7 @@ func (m *ExecutionManager) launchExecution( return nil, nil, nil, err } - annotations = m.addUserAnnotations(ctx, annotations) + annotations = m.addIdentityAnnotations(ctx, annotations) var rawOutputDataConfig *admin.RawOutputDataConfig if executionConfig.GetRawOutputDataConfig() != nil { @@ -2055,36 +2055,77 @@ func (m *ExecutionManager) addProjectLabels(ctx context.Context, projectName str return initialLabels, nil } -// addUserAnnotations automatically injects user identity information as annotations when enabled in config. -// This allows tracking which user submitted each workflow execution and enables user-based authorization. -func (m *ExecutionManager) addUserAnnotations(ctx context.Context, initialAnnotations map[string]string) map[string]string { - // Check if user annotation injection is enabled - if !m.config.ApplicationConfiguration().GetTopLevelConfig().GetInjectUserAnnotations() { +// addIdentityAnnotations automatically injects identity information (user or app) as annotations when enabled in config. +// This allows tracking which identity submitted each workflow execution and enables identity-based authorization. +func (m *ExecutionManager) addIdentityAnnotations(ctx context.Context, initialAnnotations map[string]string) map[string]string { + // Check if identity annotation injection is enabled + if !m.config.ApplicationConfiguration().GetTopLevelConfig().GetInjectIdentityAnnotations() { return initialAnnotations } - // Get user identity from authentication context + // Get identity from authentication context identityContext := auth.IdentityContextFromContext(ctx) - var principal string - if identityContext.UserInfo() != nil { - principal = identityContext.UserInfo().GetEmail() - } - - if principal == "" { - // If no email is available, skip annotation injection - logger.Debugf(ctx, "No user email found in context, skipping user annotation injection") - return initialAnnotations - } if initialAnnotations == nil { initialAnnotations = make(map[string]string) } - prefix := m.config.ApplicationConfiguration().GetTopLevelConfig().GetUserAnnotationPrefix() - userKey := prefix + "/user" - if _, exists := initialAnnotations[userKey]; !exists { - initialAnnotations[userKey] = principal - logger.Debugf(ctx, "Injected user annotation %s=%s", userKey, principal) + prefix := m.config.ApplicationConfiguration().GetTopLevelConfig().GetIdentityAnnotationPrefix() + keys := m.config.ApplicationConfiguration().GetTopLevelConfig().GetIdentityAnnotationKeys() + + // Determine if this is an app or user identity + isAppIdentity := identityContext.AppID() != "" + isUserIdentity := identityContext.UserInfo() != nil + + if !isAppIdentity && !isUserIdentity { + logger.Debugf(ctx, "No identity information found in context, skipping identity annotation injection") + return initialAnnotations + } + + // Add annotations based on identity type + if isAppIdentity { + // Handle app-based identity + appID := identityContext.AppID() + for _, key := range keys { + annotationKey := prefix + "/app-" + key + if _, exists := initialAnnotations[annotationKey]; !exists { + var value string + switch key { + case "email", "sub", "id": + // For app identities, use the app ID for these fields + value = appID + default: + // Skip unknown keys for app identities + continue + } + if value != "" { + initialAnnotations[annotationKey] = value + logger.Debugf(ctx, "Injected app identity annotation %s=%s", annotationKey, value) + } + } + } + } else if isUserIdentity { + // Handle user-based identity + userInfo := identityContext.UserInfo() + for _, key := range keys { + annotationKey := prefix + "/user-" + key + if _, exists := initialAnnotations[annotationKey]; !exists { + var value string + switch key { + case "email": + value = userInfo.GetEmail() + case "sub": + value = userInfo.GetSubject() + default: + // Skip unknown keys + continue + } + if value != "" { + initialAnnotations[annotationKey] = value + logger.Debugf(ctx, "Injected user identity annotation %s=%s", annotationKey, value) + } + } + } } return initialAnnotations diff --git a/flyteadmin/pkg/manager/impl/execution_manager_test.go b/flyteadmin/pkg/manager/impl/execution_manager_test.go index a776579675..a53bc693b8 100644 --- a/flyteadmin/pkg/manager/impl/execution_manager_test.go +++ b/flyteadmin/pkg/manager/impl/execution_manager_test.go @@ -6334,134 +6334,194 @@ func TestQueryTemplate(t *testing.T) { }) } -func TestAddUserAnnotations(t *testing.T) { +func TestAddIdentityAnnotations(t *testing.T) { principal := "test-user@example.com" + subject := "user-123-subject" - t.Run("enabled with user context", func(t *testing.T) { + t.Run("enabled with user context and multiple keys", func(t *testing.T) { mockConfig := runtimeMocks.NewMockConfigurationProvider( testutils.GetApplicationConfigWithDefaultDomains(), nil, nil, nil, nil, nil) - mockConfig.ApplicationConfiguration().GetTopLevelConfig().InjectUserAnnotations = true - mockConfig.ApplicationConfiguration().GetTopLevelConfig().UserAnnotationPrefix = "flyte.ai" + mockConfig.ApplicationConfiguration().GetTopLevelConfig().InjectIdentityAnnotations = true + mockConfig.ApplicationConfiguration().GetTopLevelConfig().IdentityAnnotationPrefix = "flyte.ai" + mockConfig.ApplicationConfiguration().GetTopLevelConfig().IdentityAnnotationKeys = []string{"email", "sub"} manager := ExecutionManager{config: mockConfig} - userInfo := &service.UserInfoResponse{Email: principal} + userInfo := &service.UserInfoResponse{Email: principal, Subject: subject} identity, err := auth.NewIdentityContext("", "user-id-123", "", time.Now(), sets.NewString(), userInfo, nil) assert.NoError(t, err) ctx := identity.WithContext(context.Background()) - result := manager.addUserAnnotations(ctx, map[string]string{"existing": "value"}) + result := manager.addIdentityAnnotations(ctx, map[string]string{"existing": "value"}) - assert.Equal(t, "test-user@example.com", result["flyte.ai/user"]) + assert.Equal(t, "test-user@example.com", result["flyte.ai/user-email"]) + assert.Equal(t, "user-123-subject", result["flyte.ai/user-sub"]) assert.Equal(t, "value", result["existing"]) }) t.Run("disabled", func(t *testing.T) { mockConfig := runtimeMocks.NewMockConfigurationProvider( testutils.GetApplicationConfigWithDefaultDomains(), nil, nil, nil, nil, nil) - mockConfig.ApplicationConfiguration().GetTopLevelConfig().InjectUserAnnotations = false + mockConfig.ApplicationConfiguration().GetTopLevelConfig().InjectIdentityAnnotations = false manager := ExecutionManager{config: mockConfig} - userInfo := &service.UserInfoResponse{Email: principal} + userInfo := &service.UserInfoResponse{Email: principal, Subject: subject} identity, err := auth.NewIdentityContext("", "user-id-123", "", time.Now(), sets.NewString(), userInfo, nil) assert.NoError(t, err) ctx := identity.WithContext(context.Background()) - result := manager.addUserAnnotations(ctx, map[string]string{"existing": "value"}) + result := manager.addIdentityAnnotations(ctx, map[string]string{"existing": "value"}) - assert.NotContains(t, result, "flyte.ai/user") + assert.NotContains(t, result, "flyte.ai/user-email") + assert.NotContains(t, result, "flyte.ai/user-sub") assert.Equal(t, "value", result["existing"]) }) - t.Run("no user context", func(t *testing.T) { + t.Run("no identity context", func(t *testing.T) { mockConfig := runtimeMocks.NewMockConfigurationProvider( testutils.GetApplicationConfigWithDefaultDomains(), nil, nil, nil, nil, nil) - mockConfig.ApplicationConfiguration().GetTopLevelConfig().InjectUserAnnotations = true + mockConfig.ApplicationConfiguration().GetTopLevelConfig().InjectIdentityAnnotations = true manager := ExecutionManager{config: mockConfig} ctx := context.Background() - result := manager.addUserAnnotations(ctx, map[string]string{"existing": "value"}) + result := manager.addIdentityAnnotations(ctx, map[string]string{"existing": "value"}) - assert.NotContains(t, result, "flyte.ai/user") + assert.NotContains(t, result, "flyte.ai/user-email") + assert.NotContains(t, result, "flyte.ai/user-sub") assert.Equal(t, "value", result["existing"]) }) - t.Run("enabled but no email", func(t *testing.T) { + t.Run("enabled but only subject available", func(t *testing.T) { mockConfig := runtimeMocks.NewMockConfigurationProvider( testutils.GetApplicationConfigWithDefaultDomains(), nil, nil, nil, nil, nil) - mockConfig.ApplicationConfiguration().GetTopLevelConfig().InjectUserAnnotations = true - mockConfig.ApplicationConfiguration().GetTopLevelConfig().UserAnnotationPrefix = "flyte.ai" + mockConfig.ApplicationConfiguration().GetTopLevelConfig().InjectIdentityAnnotations = true + mockConfig.ApplicationConfiguration().GetTopLevelConfig().IdentityAnnotationPrefix = "flyte.ai" + mockConfig.ApplicationConfiguration().GetTopLevelConfig().IdentityAnnotationKeys = []string{"email", "sub"} manager := ExecutionManager{config: mockConfig} - // UserInfo with no email set + // UserInfo with no email set (but NewIdentityContext will populate subject from userID) userInfo := &service.UserInfoResponse{} identity, err := auth.NewIdentityContext("", "user-id-123", "", time.Now(), sets.NewString(), userInfo, nil) assert.NoError(t, err) ctx := identity.WithContext(context.Background()) - result := manager.addUserAnnotations(ctx, map[string]string{"existing": "value"}) + result := manager.addIdentityAnnotations(ctx, map[string]string{"existing": "value"}) - assert.NotContains(t, result, "flyte.ai/user") + // Email should not be present, but subject should be (auto-filled by NewIdentityContext) + assert.NotContains(t, result, "flyte.ai/user-email") + assert.Equal(t, "user-id-123", result["flyte.ai/user-sub"]) assert.Equal(t, "value", result["existing"]) }) t.Run("enabled with nil annotations map", func(t *testing.T) { mockConfig := runtimeMocks.NewMockConfigurationProvider( testutils.GetApplicationConfigWithDefaultDomains(), nil, nil, nil, nil, nil) - mockConfig.ApplicationConfiguration().GetTopLevelConfig().InjectUserAnnotations = true - mockConfig.ApplicationConfiguration().GetTopLevelConfig().UserAnnotationPrefix = "flyte.ai" + mockConfig.ApplicationConfiguration().GetTopLevelConfig().InjectIdentityAnnotations = true + mockConfig.ApplicationConfiguration().GetTopLevelConfig().IdentityAnnotationPrefix = "flyte.ai" + mockConfig.ApplicationConfiguration().GetTopLevelConfig().IdentityAnnotationKeys = []string{"email", "sub"} manager := ExecutionManager{config: mockConfig} - userInfo := &service.UserInfoResponse{Email: principal} + userInfo := &service.UserInfoResponse{Email: principal, Subject: subject} identity, err := auth.NewIdentityContext("", "user-id-123", "", time.Now(), sets.NewString(), userInfo, nil) assert.NoError(t, err) ctx := identity.WithContext(context.Background()) - result := manager.addUserAnnotations(ctx, nil) + result := manager.addIdentityAnnotations(ctx, nil) - assert.Equal(t, "test-user@example.com", result["flyte.ai/user"]) + assert.Equal(t, "test-user@example.com", result["flyte.ai/user-email"]) + assert.Equal(t, "user-123-subject", result["flyte.ai/user-sub"]) }) t.Run("annotation already exists", func(t *testing.T) { mockConfig := runtimeMocks.NewMockConfigurationProvider( testutils.GetApplicationConfigWithDefaultDomains(), nil, nil, nil, nil, nil) - mockConfig.ApplicationConfiguration().GetTopLevelConfig().InjectUserAnnotations = true - mockConfig.ApplicationConfiguration().GetTopLevelConfig().UserAnnotationPrefix = "flyte.ai" + mockConfig.ApplicationConfiguration().GetTopLevelConfig().InjectIdentityAnnotations = true + mockConfig.ApplicationConfiguration().GetTopLevelConfig().IdentityAnnotationPrefix = "flyte.ai" + mockConfig.ApplicationConfiguration().GetTopLevelConfig().IdentityAnnotationKeys = []string{"email", "sub"} manager := ExecutionManager{config: mockConfig} - userInfo := &service.UserInfoResponse{Email: principal} + userInfo := &service.UserInfoResponse{Email: principal, Subject: subject} identity, err := auth.NewIdentityContext("", "user-id-123", "", time.Now(), sets.NewString(), userInfo, nil) assert.NoError(t, err) ctx := identity.WithContext(context.Background()) - result := manager.addUserAnnotations(ctx, map[string]string{"flyte.ai/user": "existing@example.com"}) + result := manager.addIdentityAnnotations(ctx, map[string]string{"flyte.ai/user-email": "existing@example.com"}) // Should preserve existing annotation value - assert.Equal(t, "existing@example.com", result["flyte.ai/user"]) + assert.Equal(t, "existing@example.com", result["flyte.ai/user-email"]) + // Subject should still be added + assert.Equal(t, "user-123-subject", result["flyte.ai/user-sub"]) }) - t.Run("uses default prefix when not configured", func(t *testing.T) { + t.Run("uses default prefix and keys when not configured", func(t *testing.T) { mockConfig := runtimeMocks.NewMockConfigurationProvider( testutils.GetApplicationConfigWithDefaultDomains(), nil, nil, nil, nil, nil) - mockConfig.ApplicationConfiguration().GetTopLevelConfig().InjectUserAnnotations = true - mockConfig.ApplicationConfiguration().GetTopLevelConfig().UserAnnotationPrefix = "" + mockConfig.ApplicationConfiguration().GetTopLevelConfig().InjectIdentityAnnotations = true + mockConfig.ApplicationConfiguration().GetTopLevelConfig().IdentityAnnotationPrefix = "" + mockConfig.ApplicationConfiguration().GetTopLevelConfig().IdentityAnnotationKeys = nil manager := ExecutionManager{config: mockConfig} - userInfo := &service.UserInfoResponse{Email: principal} + userInfo := &service.UserInfoResponse{Email: principal, Subject: subject} identity, err := auth.NewIdentityContext("", "user-id-123", "", time.Now(), sets.NewString(), userInfo, nil) assert.NoError(t, err) ctx := identity.WithContext(context.Background()) - result := manager.addUserAnnotations(ctx, map[string]string{"existing": "value"}) + result := manager.addIdentityAnnotations(ctx, map[string]string{"existing": "value"}) - // Should use default prefix "flyte.ai" - assert.Equal(t, "test-user@example.com", result["flyte.ai/user"]) + // Should use default prefix "flyte.ai" and default keys ["email", "sub"] + assert.Equal(t, "test-user@example.com", result["flyte.ai/user-email"]) + assert.Equal(t, "user-123-subject", result["flyte.ai/user-sub"]) + assert.Equal(t, "value", result["existing"]) + }) + + t.Run("app identity with multiple keys", func(t *testing.T) { + mockConfig := runtimeMocks.NewMockConfigurationProvider( + testutils.GetApplicationConfigWithDefaultDomains(), nil, nil, nil, nil, nil) + mockConfig.ApplicationConfiguration().GetTopLevelConfig().InjectIdentityAnnotations = true + mockConfig.ApplicationConfiguration().GetTopLevelConfig().IdentityAnnotationPrefix = "flyte.ai" + mockConfig.ApplicationConfiguration().GetTopLevelConfig().IdentityAnnotationKeys = []string{"email", "sub", "id"} + + manager := ExecutionManager{config: mockConfig} + + // App identity (no userInfo, but has appID) + identity, err := auth.NewIdentityContext("", "", "app-123", time.Now(), sets.NewString(), nil, nil) + assert.NoError(t, err) + ctx := identity.WithContext(context.Background()) + + result := manager.addIdentityAnnotations(ctx, map[string]string{"existing": "value"}) + + // Should use app prefix and app ID for all keys + assert.Equal(t, "app-123", result["flyte.ai/app-email"]) + assert.Equal(t, "app-123", result["flyte.ai/app-sub"]) + assert.Equal(t, "app-123", result["flyte.ai/app-id"]) + assert.Equal(t, "value", result["existing"]) + }) + + t.Run("only email key configured", func(t *testing.T) { + mockConfig := runtimeMocks.NewMockConfigurationProvider( + testutils.GetApplicationConfigWithDefaultDomains(), nil, nil, nil, nil, nil) + mockConfig.ApplicationConfiguration().GetTopLevelConfig().InjectIdentityAnnotations = true + mockConfig.ApplicationConfiguration().GetTopLevelConfig().IdentityAnnotationPrefix = "flyte.ai" + mockConfig.ApplicationConfiguration().GetTopLevelConfig().IdentityAnnotationKeys = []string{"email"} + + manager := ExecutionManager{config: mockConfig} + + userInfo := &service.UserInfoResponse{Email: principal, Subject: subject} + identity, err := auth.NewIdentityContext("", "user-id-123", "", time.Now(), sets.NewString(), userInfo, nil) + assert.NoError(t, err) + ctx := identity.WithContext(context.Background()) + + result := manager.addIdentityAnnotations(ctx, map[string]string{"existing": "value"}) + + // Should only add email, not subject + assert.Equal(t, "test-user@example.com", result["flyte.ai/user-email"]) + assert.NotContains(t, result, "flyte.ai/user-sub") assert.Equal(t, "value", result["existing"]) }) } diff --git a/flyteadmin/pkg/runtime/interfaces/application_configuration.go b/flyteadmin/pkg/runtime/interfaces/application_configuration.go index 35ae4ebea2..9b2ae9f1a2 100644 --- a/flyteadmin/pkg/runtime/interfaces/application_configuration.go +++ b/flyteadmin/pkg/runtime/interfaces/application_configuration.go @@ -117,8 +117,9 @@ type ApplicationConfig struct { // Enabling this will instruct operator to use storage (s3/gcs/etc) to offload workflow execution inputs instead of storing them inline in the CRD. UseOffloadedInputs bool `json:"useOffloadedInputs" pflag:",Use offloaded inputs for workflows."` - InjectUserAnnotations bool `json:"injectUserAnnotations"` - UserAnnotationPrefix string `json:"userAnnotationPrefix"` + InjectIdentityAnnotations bool `json:"injectIdentityAnnotations"` + IdentityAnnotationPrefix string `json:"identityAnnotationPrefix"` + IdentityAnnotationKeys []string `json:"identityAnnotationKeys"` } func (a *ApplicationConfig) GetRoleNameKey() string { @@ -204,15 +205,22 @@ func (a *ApplicationConfig) GetEnvs() *admin.Envs { } } -func (a *ApplicationConfig) GetInjectUserAnnotations() bool { - return a.InjectUserAnnotations +func (a *ApplicationConfig) GetInjectIdentityAnnotations() bool { + return a.InjectIdentityAnnotations } -func (a *ApplicationConfig) GetUserAnnotationPrefix() string { - if a.UserAnnotationPrefix == "" { +func (a *ApplicationConfig) GetIdentityAnnotationPrefix() string { + if a.IdentityAnnotationPrefix == "" { return "flyte.ai" } - return a.UserAnnotationPrefix + return a.IdentityAnnotationPrefix +} + +func (a *ApplicationConfig) GetIdentityAnnotationKeys() []string { + if len(a.IdentityAnnotationKeys) == 0 { + return []string{"email", "sub"} + } + return a.IdentityAnnotationKeys } // GetAsWorkflowExecutionConfig returns the WorkflowExecutionConfig as extracted from this object From 965b291bb6bd5d2bdf270021cfb24f7e0764dcf4 Mon Sep 17 00:00:00 2001 From: ttitsworth-fsp Date: Thu, 6 Nov 2025 15:02:00 -0800 Subject: [PATCH 6/7] add more coverage Signed-off-by: ttitsworth-fsp --- .../manager/impl/execution_manager_test.go | 44 +++++++++++++++++++ 1 file changed, 44 insertions(+) diff --git a/flyteadmin/pkg/manager/impl/execution_manager_test.go b/flyteadmin/pkg/manager/impl/execution_manager_test.go index a53bc693b8..7f2e6394ff 100644 --- a/flyteadmin/pkg/manager/impl/execution_manager_test.go +++ b/flyteadmin/pkg/manager/impl/execution_manager_test.go @@ -6524,4 +6524,48 @@ func TestAddIdentityAnnotations(t *testing.T) { assert.NotContains(t, result, "flyte.ai/user-sub") assert.Equal(t, "value", result["existing"]) }) + + t.Run("app identity with unknown key", func(t *testing.T) { + mockConfig := runtimeMocks.NewMockConfigurationProvider( + testutils.GetApplicationConfigWithDefaultDomains(), nil, nil, nil, nil, nil) + mockConfig.ApplicationConfiguration().GetTopLevelConfig().InjectIdentityAnnotations = true + mockConfig.ApplicationConfiguration().GetTopLevelConfig().IdentityAnnotationPrefix = "flyte.ai" + mockConfig.ApplicationConfiguration().GetTopLevelConfig().IdentityAnnotationKeys = []string{"email", "unknown-key"} + + manager := ExecutionManager{config: mockConfig} + + // App identity + identity, err := auth.NewIdentityContext("", "", "app-123", time.Now(), sets.NewString(), nil, nil) + assert.NoError(t, err) + ctx := identity.WithContext(context.Background()) + + result := manager.addIdentityAnnotations(ctx, map[string]string{"existing": "value"}) + + // Should only add email (known key), not unknown-key + assert.Equal(t, "app-123", result["flyte.ai/app-email"]) + assert.NotContains(t, result, "flyte.ai/app-unknown-key") + assert.Equal(t, "value", result["existing"]) + }) + + t.Run("user identity with unknown key", func(t *testing.T) { + mockConfig := runtimeMocks.NewMockConfigurationProvider( + testutils.GetApplicationConfigWithDefaultDomains(), nil, nil, nil, nil, nil) + mockConfig.ApplicationConfiguration().GetTopLevelConfig().InjectIdentityAnnotations = true + mockConfig.ApplicationConfiguration().GetTopLevelConfig().IdentityAnnotationPrefix = "flyte.ai" + mockConfig.ApplicationConfiguration().GetTopLevelConfig().IdentityAnnotationKeys = []string{"email", "unknown-key"} + + manager := ExecutionManager{config: mockConfig} + + userInfo := &service.UserInfoResponse{Email: principal, Subject: subject} + identity, err := auth.NewIdentityContext("", "user-id-123", "", time.Now(), sets.NewString(), userInfo, nil) + assert.NoError(t, err) + ctx := identity.WithContext(context.Background()) + + result := manager.addIdentityAnnotations(ctx, map[string]string{"existing": "value"}) + + // Should only add email (known key), not unknown-key + assert.Equal(t, "test-user@example.com", result["flyte.ai/user-email"]) + assert.NotContains(t, result, "flyte.ai/user-unknown-key") + assert.Equal(t, "value", result["existing"]) + }) } From 8efdf91f0e99e0fabff36f2d5ec0bb6ac3c96d99 Mon Sep 17 00:00:00 2001 From: ttitsworth-fsp Date: Fri, 7 Nov 2025 09:19:28 -0800 Subject: [PATCH 7/7] improve coverage and fix bad var handling Signed-off-by: ttitsworth-fsp --- flyteadmin/pkg/manager/impl/execution_manager.go | 13 +++++++------ .../pkg/manager/impl/execution_manager_test.go | 1 + 2 files changed, 8 insertions(+), 6 deletions(-) diff --git a/flyteadmin/pkg/manager/impl/execution_manager.go b/flyteadmin/pkg/manager/impl/execution_manager.go index c8e7925246..e2eaafa456 100644 --- a/flyteadmin/pkg/manager/impl/execution_manager.go +++ b/flyteadmin/pkg/manager/impl/execution_manager.go @@ -2066,6 +2066,12 @@ func (m *ExecutionManager) addIdentityAnnotations(ctx context.Context, initialAn // Get identity from authentication context identityContext := auth.IdentityContextFromContext(ctx) + // Check if identity context is empty + if identityContext.IsEmpty() { + logger.Debugf(ctx, "No identity information found in context, skipping identity annotation injection") + return initialAnnotations + } + if initialAnnotations == nil { initialAnnotations = make(map[string]string) } @@ -2075,12 +2081,7 @@ func (m *ExecutionManager) addIdentityAnnotations(ctx context.Context, initialAn // Determine if this is an app or user identity isAppIdentity := identityContext.AppID() != "" - isUserIdentity := identityContext.UserInfo() != nil - - if !isAppIdentity && !isUserIdentity { - logger.Debugf(ctx, "No identity information found in context, skipping identity annotation injection") - return initialAnnotations - } + isUserIdentity := identityContext.UserInfo() != nil && !isAppIdentity // Add annotations based on identity type if isAppIdentity { diff --git a/flyteadmin/pkg/manager/impl/execution_manager_test.go b/flyteadmin/pkg/manager/impl/execution_manager_test.go index 7f2e6394ff..f84249f0fa 100644 --- a/flyteadmin/pkg/manager/impl/execution_manager_test.go +++ b/flyteadmin/pkg/manager/impl/execution_manager_test.go @@ -6568,4 +6568,5 @@ func TestAddIdentityAnnotations(t *testing.T) { assert.NotContains(t, result, "flyte.ai/user-unknown-key") assert.Equal(t, "value", result["existing"]) }) + }