Upstream: Add labels to published execution events (#6104)

flyteadmin/pkg/async/cloudevent/factory.go

@@ -101,7 +101,7 @@ func NewCloudEventsPublisher(ctx context.Context, db repositoryInterfaces.Reposi
 	}
 
 	if cloudEventsConfig.CloudEventVersion == runtimeInterfaces.CloudEventVersionv2 {
-		return cloudEventImplementations.NewCloudEventsWrappedPublisher(db, sender, scope, storageClient, urlData, remoteDataConfig)
+		return cloudEventImplementations.NewCloudEventsWrappedPublisher(db, sender, scope, storageClient, urlData, remoteDataConfig, cloudEventsConfig.EventsPublisherConfig)
 	}
 
 	return cloudEventImplementations.NewCloudEventsPublisher(sender, scope, cloudEventsConfig.EventsPublisherConfig.EventTypes)

flyteadmin/pkg/async/cloudevent/implementations/cloudevent_publisher.go

@@ -114,12 +114,13 @@ func (p *Publisher) shouldPublishEvent(notificationType string) bool {
 }
 
 type CloudEventWrappedPublisher struct {
-	db               repositoryInterfaces.Repository
-	sender           interfaces.Sender
-	systemMetrics    implementations.EventPublisherSystemMetrics
-	storageClient    *storage.DataStore
-	urlData          dataInterfaces.RemoteURLInterface
-	remoteDataConfig runtimeInterfaces.RemoteDataConfig
+	db                   repositoryInterfaces.Repository
+	sender               interfaces.Sender
+	systemMetrics        implementations.EventPublisherSystemMetrics
+	storageClient        *storage.DataStore
+	urlData              dataInterfaces.RemoteURLInterface
+	remoteDataConfig     runtimeInterfaces.RemoteDataConfig
+	eventPublisherConfig runtimeInterfaces.EventsPublisherConfig
 }
 
 func (c *CloudEventWrappedPublisher) TransformWorkflowExecutionEvent(ctx context.Context, rawEvent *event.WorkflowExecutionEvent) (*event.CloudEventWorkflowExecution, error) {
@@ -133,8 +134,8 @@ func (c *CloudEventWrappedPublisher) TransformWorkflowExecutionEvent(ctx context
 		return nil, fmt.Errorf("nil execution id in event [%+v]", rawEvent)
 	}
 
-	// For now, don't append any additional information unless succeeded
-	if rawEvent.GetPhase() != core.WorkflowExecution_SUCCEEDED {
+	// For now, don't append any additional information unless succeeded or otherwise configured
+	if rawEvent.GetPhase() != core.WorkflowExecution_SUCCEEDED && !c.eventPublisherConfig.EnrichAllWorkflowEventTypes {
 		return &event.CloudEventWorkflowExecution{
 			RawEvent: rawEvent,
 		}, nil
@@ -193,6 +194,7 @@ func (c *CloudEventWrappedPublisher) TransformWorkflowExecutionEvent(ctx context
 		ReferenceExecution: spec.GetMetadata().GetReferenceExecution(),
 		Principal:          spec.GetMetadata().GetPrincipal(),
 		LaunchPlanId:       spec.GetLaunchPlan(),
+		Labels:             spec.GetLabels().GetValues(),
 	}, nil
 }
 
@@ -317,6 +319,7 @@ func (c *CloudEventWrappedPublisher) TransformNodeExecutionEvent(ctx context.Con
 		ArtifactIds:     spec.GetMetadata().GetArtifactIds(),
 		Principal:       spec.GetMetadata().GetPrincipal(),
 		LaunchPlanId:    spec.GetLaunchPlan(),
+		Labels:          spec.GetLabels().GetValues(),
 	}, nil
 }
 
@@ -326,8 +329,24 @@ func (c *CloudEventWrappedPublisher) TransformTaskExecutionEvent(ctx context.Con
 		return nil, fmt.Errorf("nothing to publish, TaskExecution event is nil")
 	}
 
+	executionModel, err := c.db.ExecutionRepo().Get(ctx, repositoryInterfaces.Identifier{
+		Project: rawEvent.GetParentNodeExecutionId().GetExecutionId().GetProject(),
+		Domain:  rawEvent.GetParentNodeExecutionId().GetExecutionId().GetDomain(),
+		Name:    rawEvent.GetParentNodeExecutionId().GetExecutionId().GetName(),
+	})
+	if err != nil {
+		logger.Warningf(ctx, "couldn't find execution [%+v] for cloud event processing", rawEvent.GetParentNodeExecutionId().GetExecutionId())
+		return nil, err
+	}
+	ex, err := transformers.FromExecutionModel(ctx, executionModel, transformers.DefaultExecutionTransformerOptions)
+	if err != nil {
+		logger.Warningf(ctx, "couldn't transform execution [%+v] for cloud event processing", rawEvent.GetParentNodeExecutionId().GetExecutionId())
+		return nil, err
+	}
+
 	return &event.CloudEventTaskExecution{
 		RawEvent: rawEvent,
+		Labels:   ex.GetSpec().GetLabels().GetValues(),
 	}, nil
 }
 
@@ -472,14 +491,15 @@ func NewCloudEventsPublisher(sender interfaces.Sender, scope promutils.Scope, ev
 }
 
 func NewCloudEventsWrappedPublisher(
-	db repositoryInterfaces.Repository, sender interfaces.Sender, scope promutils.Scope, storageClient *storage.DataStore, urlData dataInterfaces.RemoteURLInterface, remoteDataConfig runtimeInterfaces.RemoteDataConfig) interfaces.Publisher {
+	db repositoryInterfaces.Repository, sender interfaces.Sender, scope promutils.Scope, storageClient *storage.DataStore, urlData dataInterfaces.RemoteURLInterface, remoteDataConfig runtimeInterfaces.RemoteDataConfig, eventPublisherConfig runtimeInterfaces.EventsPublisherConfig) interfaces.Publisher {
 
 	return &CloudEventWrappedPublisher{
-		db:               db,
-		sender:           sender,
-		systemMetrics:    implementations.NewEventPublisherSystemMetrics(scope.NewSubScope("cloudevents_publisher")),
-		storageClient:    storageClient,
-		urlData:          urlData,
-		remoteDataConfig: remoteDataConfig,
+		db:                   db,
+		sender:               sender,
+		systemMetrics:        implementations.NewEventPublisherSystemMetrics(scope.NewSubScope("cloudevents_publisher")),
+		storageClient:        storageClient,
+		urlData:              urlData,
+		remoteDataConfig:     remoteDataConfig,
+		eventPublisherConfig: eventPublisherConfig,
 	}
 }

flyteadmin/pkg/runtime/interfaces/application_configuration.go

@@ -603,6 +603,8 @@ type EventsPublisherConfig struct {
 	TopicName string `json:"topicName"`
 	// Event types: task, node, workflow executions
 	EventTypes []string `json:"eventTypes"`
+	// Whether to publish enriched events for all workflow execution events
+	EnrichAllWorkflowEventTypes bool `json:"enrichAllWorkflowEventTypes"`
 }
 
 type ExternalEventsConfig struct {