From 6567ecadc995d32b096ce2be8943e4ce96ed84d5 Mon Sep 17 00:00:00 2001
From: catalinii <catalinii@users.noreply.github.com>
Date: Mon, 5 Oct 2020 16:27:29 -0700
Subject: [PATCH] Set labels in workflow metadata (#129)

---
 .../pkg/manager/impl/execution_manager.go     | 32 +++++++++++++++++++
 .../manager/impl/execution_manager_test.go    | 13 ++++++++
 2 files changed, 45 insertions(+)

diff --git a/flyteadmin/pkg/manager/impl/execution_manager.go b/flyteadmin/pkg/manager/impl/execution_manager.go
index 2020505799..b77f67e720 100644
--- a/flyteadmin/pkg/manager/impl/execution_manager.go
+++ b/flyteadmin/pkg/manager/impl/execution_manager.go
@@ -467,6 +467,11 @@ func (m *ExecutionManager) launchSingleTaskExecution(
 	if request.Spec.Labels != nil {
 		executeTaskInputs.Labels = request.Spec.Labels.Values
 	}
+	executeTaskInputs.Labels, err = m.addProjectLabels(ctx, request.Project, executeTaskInputs.Labels)
+	if err != nil {
+		return nil, nil, err
+	}
+
 	if request.Spec.Annotations != nil {
 		executeTaskInputs.Annotations = request.Spec.Annotations.Values
 	}
@@ -632,6 +637,11 @@ func (m *ExecutionManager) launchExecutionAndPrepareModel(
 	if err != nil {
 		return nil, nil, err
 	}
+	executeWorkflowInputs.Labels, err = m.addProjectLabels(ctx, request.Project, executeWorkflowInputs.Labels)
+	if err != nil {
+		return nil, nil, err
+	}
+
 	err = m.addPluginOverrides(ctx, &workflowExecutionID, launchPlan.GetSpec().WorkflowId.Name, launchPlan.Id.Name,
 		&executeWorkflowInputs)
 	if err != nil {
@@ -1346,3 +1356,25 @@ func NewExecutionManager(
 		qualityOfServiceAllocator: executions.NewQualityOfServiceAllocator(config, resourceManager),
 	}
 }
+
+// Adds project labels with higher precedence to workflow labels. Project labels are ignored if a corresponding label is set on the workflow.
+func (m *ExecutionManager) addProjectLabels(ctx context.Context, projectName string, initialLabels map[string]string) (map[string]string, error) {
+	project, err := m.db.ProjectRepo().Get(ctx, projectName)
+	if err != nil {
+		logger.Errorf(ctx, "Failed to get project for [%+v] with error: %v", project, err)
+		return nil, err
+	}
+	// passing nil domain as not needed to retrieve labels
+	projectLabels := transformers.FromProjectModel(project, nil).Labels.GetValues()
+
+	if initialLabels == nil {
+		initialLabels = make(map[string]string)
+	}
+
+	for k, v := range projectLabels {
+		if _, ok := initialLabels[k]; !ok {
+			initialLabels[k] = v
+		}
+	}
+	return initialLabels, nil
+}
diff --git a/flyteadmin/pkg/manager/impl/execution_manager_test.go b/flyteadmin/pkg/manager/impl/execution_manager_test.go
index c3d2a90a73..56fd12d0e0 100644
--- a/flyteadmin/pkg/manager/impl/execution_manager_test.go
+++ b/flyteadmin/pkg/manager/impl/execution_manager_test.go
@@ -35,6 +35,7 @@ import (
 	"github.com/lyft/flyteadmin/pkg/repositories/interfaces"
 	repositoryMocks "github.com/lyft/flyteadmin/pkg/repositories/mocks"
 	"github.com/lyft/flyteadmin/pkg/repositories/models"
+	"github.com/lyft/flyteadmin/pkg/repositories/transformers"
 	runtimeInterfaces "github.com/lyft/flyteadmin/pkg/runtime/interfaces"
 	runtimeIFaceMocks "github.com/lyft/flyteadmin/pkg/runtime/interfaces/mocks"
 	runtimeMocks "github.com/lyft/flyteadmin/pkg/runtime/mocks"
@@ -198,6 +199,17 @@ func getMockRepositoryForExecTest() repositories.RepositoryInterface {
 
 func TestCreateExecution(t *testing.T) {
 	repository := getMockRepositoryForExecTest()
+	labels := admin.Labels{
+		Values: map[string]string{
+			"label3": "3",
+			"label2": "1", // common label, will be dropped
+		}}
+	repository.ProjectRepo().(*repositoryMocks.MockProjectRepo).GetFunction = func(
+		ctx context.Context, projectID string) (models.Project, error) {
+		return transformers.CreateProjectModel(&admin.Project{
+			Labels: &labels}), nil
+	}
+
 	principal := "principal"
 	repository.ExecutionRepo().(*repositoryMocks.MockExecutionRepo).SetCreateCallback(
 		func(ctx context.Context, input models.Execution) error {
@@ -214,6 +226,7 @@ func TestCreateExecution(t *testing.T) {
 			assert.EqualValues(t, map[string]string{
 				"label1": "1",
 				"label2": "2",
+				"label3": "3",
 			}, inputs.Labels)
 			assert.EqualValues(t, map[string]string{
 				"annotation3": "3",