Add UpdateProject functionality to Flyte control plane

go.mod

@@ -22,7 +22,7 @@ require (
 	github.com/jinzhu/gorm v1.9.12
 	github.com/kelseyhightower/envconfig v1.4.0 // indirect
 	github.com/lib/pq v1.3.0
-	github.com/lyft/flyteidl v0.18.2
+	github.com/lyft/flyteidl v0.18.3
 	github.com/lyft/flytepropeller v0.3.7
 	github.com/lyft/flytestdlib v0.3.9
 	github.com/magiconair/properties v1.8.1
@@ -33,6 +33,7 @@ require (
 	github.com/spf13/pflag v1.0.5
 	github.com/stretchr/testify v1.6.1
 	golang.org/x/oauth2 v0.0.0-20200107190931-bf48bf16ab8d
+	golang.org/x/tools v0.0.0-20200818005847-188abfa75333 // indirect
 	google.golang.org/grpc v1.28.0
 	gopkg.in/gormigrate.v1 v1.6.0
 	k8s.io/api v0.17.3

pkg/manager/impl/project_manager.go

@@ -60,6 +60,27 @@ func (m *ProjectManager) ListProjects(ctx context.Context, request admin.Project
 	}, nil
 }
 
+func (m *ProjectManager) UpdateProject(ctx context.Context, projectUpdate admin.Project) (*admin.ProjectUpdateResponse, error) {
+	var response *admin.ProjectUpdateResponse
+	projectRepo := m.db.ProjectRepo()
+
+	// Fetch the existing project if exists. If not, return err and do not update.
+	_, err := projectRepo.Get(ctx, projectUpdate.Id)
+	if err != nil {
+		return nil, err
+	}
+
+	// Transform the provided project into a model and apply to the DB.
+	projectUpdateModel := transformers.CreateProjectModel(&projectUpdate)
+	err = projectRepo.UpdateProject(ctx, projectUpdateModel)
+
+	if err != nil {
+		return nil, err
+	}
+
+	return response, nil
+}
+
 func NewProjectManager(db repositories.RepositoryInterface, config runtimeInterfaces.Configuration) interfaces.ProjectInterface {
 	return &ProjectManager{
 		db:     db,

pkg/manager/interfaces/project.go

@@ -10,4 +10,5 @@ import (
 type ProjectInterface interface {
 	CreateProject(ctx context.Context, request admin.ProjectRegisterRequest) (*admin.ProjectRegisterResponse, error)
 	ListProjects(ctx context.Context, request admin.ProjectListRequest) (*admin.Projects, error)
+	UpdateProject(ctx context.Context, request admin.Project) (*admin.ProjectUpdateResponse, error)
 }

pkg/manager/mocks/project.go

@@ -8,24 +8,32 @@ import (
 
 type CreateProjectFunc func(ctx context.Context, request admin.ProjectRegisterRequest) (*admin.ProjectRegisterResponse, error)
 type ListProjectFunc func(ctx context.Context, request admin.ProjectListRequest) (*admin.Projects, error)
+type UpdateProjectFunc func(ctx context.Context, request admin.Project) (*admin.ProjectUpdateResponse, error)
 
 type MockProjectManager struct {
 	listProjectFunc   ListProjectFunc
 	createProjectFunc CreateProjectFunc
+	updateProjectFunc UpdateProjectFunc
 }
 
 func (m *MockProjectManager) SetCreateProject(createProjectFunc CreateProjectFunc) {
 	m.createProjectFunc = createProjectFunc
 }
 
-func (m *MockProjectManager) CreateProject(ctx context.Context, request admin.ProjectRegisterRequest) (
-	*admin.ProjectRegisterResponse, error) {
+func (m *MockProjectManager) CreateProject(ctx context.Context, request admin.ProjectRegisterRequest) (*admin.ProjectRegisterResponse, error) {
 	if m.createProjectFunc != nil {
 		return m.createProjectFunc(ctx, request)
 	}
 	return nil, nil
 }
 
+func (m *MockProjectManager) UpdateProject(ctx context.Context, request admin.Project) (*admin.ProjectUpdateResponse, error) {
+	if m.updateProjectFunc != nil {
+		return m.updateProjectFunc(ctx, request)
+	}
+	return nil, nil
+}
+
 func (m *MockProjectManager) SetListCallback(listProjectFunc ListProjectFunc) {
 	m.listProjectFunc = listProjectFunc
 }

pkg/repositories/config/migrations.go

@@ -251,4 +251,13 @@ var Migrations = []*gormigrate.Migration{
 			return tx.Model(&models.NodeExecution{}).DropColumn("parent_id").DropColumn("node_execution_metadata").Error
 		},
 	},
+	{
+		ID: "2020-08-17-labels-addition",
+		Migrate: func(tx *gorm.DB) error {
+			return tx.AutoMigrate(&models.NodeExecution{}).Error
+		},
+		Rollback: func(tx *gorm.DB) error {
+			return tx.Model(&models.NodeExecution{}).DropColumn("labels").Error
+		},
+	},
 }

pkg/repositories/gormimpl/project_repo.go

@@ -70,3 +70,15 @@ func NewProjectRepo(db *gorm.DB, errorTransformer errors.ErrorTransformer,
 		metrics:          metrics,
 	}
 }
+
+func (r *ProjectRepo) UpdateProject(ctx context.Context, projectUpdate models.Project) error {
+	// Use gorm client to update the two fields that are changed.
+	writeTx := r.db.Model(&projectUpdate).Update("Description", "Labels")
+
+	// Return error if applies.
+	if writeTx != nil {
+		_ = r.errorTransformer.ToFlyteAdminError(writeTx.Error)
+	}
+
+	return nil
+}

pkg/repositories/gormimpl/project_repo_test.go

@@ -19,7 +19,7 @@ func TestCreateProject(t *testing.T) {
 
 	query := GlobalMock.NewMock()
 	query.WithQuery(
-		`INSERT INTO "projects" ("created_at","updated_at","deleted_at","identifier","name","description") VALUES (?,?,?,?,?,?)`)
+		`INSERT INTO "projects" ("created_at","updated_at","deleted_at","identifier","name","description","labels") VALUES (?,?,?,?,?,?,?)`)
 
 	err := projectRepo.Create(context.Background(), models.Project{
 		Identifier:  "proj",

pkg/repositories/interfaces/project_repo.go

@@ -14,4 +14,8 @@ type ProjectRepoInterface interface {
 	Get(ctx context.Context, projectID string) (models.Project, error)
 	// Lists unique projects registered as namespaces
 	ListAll(ctx context.Context, sortParameter common.SortParameter) ([]models.Project, error)
+	// Given a project that exists in the DB and a partial set of fields to update
+	// as a second project (projectUpdate), updates the original project which already
+	// exists in the DB.
+	UpdateProject(ctx context.Context, projectUpdate models.Project) error
 }

pkg/repositories/mocks/project_repo.go

@@ -11,11 +11,13 @@ import (
 type CreateProjectFunction func(ctx context.Context, project models.Project) error
 type GetProjectFunction func(ctx context.Context, projectID string) (models.Project, error)
 type ListProjectsFunction func(ctx context.Context, sortParameter common.SortParameter) ([]models.Project, error)
+type UpdateProjectFunction func(ctx context.Context, projectUpdate models.Project) error
 
 type MockProjectRepo struct {
-	CreateFunction       CreateProjectFunction
-	GetFunction          GetProjectFunction
-	ListProjectsFunction ListProjectsFunction
+	CreateFunction        CreateProjectFunction
+	GetFunction           GetProjectFunction
+	ListProjectsFunction  ListProjectsFunction
+	UpdateProjectFunction UpdateProjectFunction
 }
 
 func (r *MockProjectRepo) Create(ctx context.Context, project models.Project) error {
@@ -39,6 +41,13 @@ func (r *MockProjectRepo) ListAll(ctx context.Context, sortParameter common.Sort
 	return make([]models.Project, 0), nil
 }
 
+func (r *MockProjectRepo) UpdateProject(ctx context.Context, projectUpdate models.Project) error {
+	if r.UpdateProjectFunction != nil {
+		return r.UpdateProjectFunction(ctx, projectUpdate)
+	}
+	return nil
+}
+
 func NewMockProjectRepo() interfaces.ProjectRepoInterface {
 	return &MockProjectRepo{}
 }

pkg/repositories/models/project.go

@@ -5,4 +5,5 @@ type Project struct {
 	Identifier  string `gorm:"primary_key"`
 	Name        string // Human-readable name, not a unique identifier.
 	Description string `gorm:"type:varchar(300)"`
+	Labels      []byte
 }

pkg/repositories/transformers/project.go

@@ -1,6 +1,7 @@
 package transformers
 
 import (
+	"github.com/golang/protobuf/proto"
 	"github.com/lyft/flyteadmin/pkg/repositories/models"
 	"github.com/lyft/flyteidl/gen/pb-go/flyteidl/admin"
 )
@@ -12,18 +13,29 @@ type CreateProjectModelInput struct {
 }
 
 func CreateProjectModel(project *admin.Project) models.Project {
+	projectBytes, err := proto.Marshal(project)
+	if err != nil {
+		return models.Project{}
+	}
 	return models.Project{
 		Identifier:  project.Id,
 		Name:        project.Name,
 		Description: project.Description,
+		Labels:      projectBytes,
 	}
 }
 
 func FromProjectModel(projectModel models.Project, domains []*admin.Domain) admin.Project {
+	projectDeserialized := &admin.Project{}
+	err := proto.Unmarshal(projectModel.Labels, projectDeserialized)
+	if err != nil {
+		return admin.Project{}
+	}
 	project := admin.Project{
 		Id:          projectModel.Identifier,
 		Name:        projectModel.Name,
 		Description: projectModel.Description,
+		Labels:      projectDeserialized.Labels,
 	}
 	project.Domains = domains
 	return project

pkg/repositories/transformers/project_test.go

@@ -22,6 +22,11 @@ func TestCreateProjectModel(t *testing.T) {
 		Identifier:  "project_id",
 		Name:        "project_name",
 		Description: "project_description",
+		Labels: []uint8{
+			0xa, 0xa, 0x70, 0x72, 0x6f, 0x6a, 0x65, 0x63, 0x74, 0x5f, 0x69, 0x64, 0x12, 0xc, 0x70, 0x72, 0x6f,
+			0x6a, 0x65, 0x63, 0x74, 0x5f, 0x6e, 0x61, 0x6d, 0x65, 0x22, 0x13, 0x70, 0x72, 0x6f, 0x6a, 0x65,
+			0x63, 0x74, 0x5f, 0x64, 0x65, 0x73, 0x63, 0x72, 0x69, 0x70, 0x74, 0x69, 0x6f, 0x6e,
+		},
 	}, projectModel)
 }
 

pkg/rpc/adminservice/metrics.go

@@ -54,6 +54,7 @@ type projectEndpointMetrics struct {
 
 	register util.RequestMetrics
 	list     util.RequestMetrics
+	update   util.RequestMetrics
 }
 
 type attributeEndpointMetrics struct {
@@ -154,6 +155,7 @@ func InitMetrics(adminScope promutils.Scope) AdminMetrics {
 			scope:    adminScope,
 			register: util.NewRequestMetrics(adminScope, "register_project"),
 			list:     util.NewRequestMetrics(adminScope, "list_projects"),
+			update:   util.NewRequestMetrics(adminScope, "update_project"),
 		},
 		projectAttributesEndpointMetrics: attributeEndpointMetrics{
 			scope:  adminScope,

pkg/rpc/adminservice/project.go

@@ -63,3 +63,30 @@ func (m *AdminService) ListProjects(ctx context.Context, request *admin.ProjectL
 	m.Metrics.projectEndpointMetrics.list.Success()
 	return response, nil
 }
+
+func (m *AdminService) UpdateProject(ctx context.Context, request *admin.Project) (
+	*admin.ProjectUpdateResponse, error) {
+	defer m.interceptPanic(ctx, request)
+	requestedAt := time.Now()
+	if request == nil {
+		return nil, status.Errorf(codes.InvalidArgument, "Incorrect request, nil requests not allowed")
+	}
+	var response *admin.ProjectUpdateResponse
+	var err error
+	m.Metrics.projectEndpointMetrics.register.Time(func() {
+		response, err = m.ProjectManager.UpdateProject(ctx, *request)
+	})
+	audit.NewLogBuilder().WithAuthenticatedCtx(ctx).WithRequest(
+		"UpdateProject",
+		map[string]string{
+			audit.Project: request.Id,
+		},
+		audit.ReadWrite,
+		requestedAt,
+	).WithResponse(time.Now(), err).Log(ctx)
+	if err != nil {
+		return nil, util.TransformAndRecordError(err, &m.Metrics.projectEndpointMetrics.update)
+	}
+
+	return response, nil
+}

tests/project.go

@@ -43,3 +43,108 @@ func TestCreateProject(t *testing.T) {
 	}
 	assert.True(t, sawNewProject)
 }
+
+func TestUpdateProjectDescription(t *testing.T) {
+	ctx := context.Background()
+	client, conn := GetTestAdminServiceClient()
+	defer conn.Close()
+
+	// Create a new project.
+	req := admin.ProjectRegisterRequest{
+		Project: &admin.Project{
+			Id:   "potato",
+			Name: "spud",
+			Labels: &admin.Labels{
+				Values: map[string]string{
+					"foo": "bar",
+					"bar": "baz",
+				},
+			},
+		},
+	}
+	_, err := client.RegisterProject(ctx, &req)
+	assert.Nil(t, err)
+
+	// Verify the project has been registered.
+	projects, err := client.ListProjects(ctx, &admin.ProjectListRequest{})
+	assert.Nil(t, err)
+	assert.NotEmpty(t, projects.Projects)
+
+	// Attempt to modify the name of the Project. Modifying the Name should be a
+	// no-op, while the Description is modified. Labels should be a no-op.
+	_, err = client.UpdateProject(ctx, &admin.Project{
+		Id:          "potato",
+		Name:        "foobar",
+		Description: "a-new-description",
+		Labels:      &admin.Labels{},
+	})
+
+	// Fetch updated projects.
+	projectsUpdated, err := client.ListProjects(ctx, &admin.ProjectListRequest{})
+	assert.Nil(t, err)
+	assert.NotEmpty(t, projectsUpdated.Projects)
+
+	// Verify that the project's Name has not been modified but the Description has.
+	updatedProject := projectsUpdated.Projects[0]
+	assert.Equal(t, updatedProject.Id, "potato")
+	assert.Equal(t, updatedProject.Name, "spud")                     // unchanged
+	assert.Equal(t, updatedProject.Description, "a-new-description") // changed
+
+	// Verify that project labels are not removed.
+	labelsMap := updatedProject.Labels
+	fooVal, fooExists := labelsMap.Values["foo"]
+	barVal, barExists := labelsMap.Values["bar"]
+	assert.Equal(t, fooExists, true)
+	assert.Equal(t, fooVal, "bar")
+	assert.Equal(t, barExists, true)
+	assert.Equal(t, barVal, "baz")
+}
+
+func TestUpdateProjectLabels(t *testing.T) {
+	ctx := context.Background()
+	client, conn := GetTestAdminServiceClient()
+	defer conn.Close()
+
+	// Create a new project.
+	req := admin.ProjectRegisterRequest{
+		Project: &admin.Project{
+			Id:   "potato",
+			Name: "spud",
+		},
+	}
+	_, err := client.RegisterProject(ctx, &req)
+	assert.Nil(t, err)
+
+	// Verify the project has been registered.
+	projects, err := client.ListProjects(ctx, &admin.ProjectListRequest{})
+	assert.Nil(t, err)
+	assert.NotEmpty(t, projects.Projects)
+
+	// Attempt to modify the name of the Project. Modifying the Name should be a
+	// no-op, while the Labels are modified.
+	_, err = client.UpdateProject(ctx, &admin.Project{
+		Id:   "potato",
+		Name: "foobar",
+		Labels: &admin.Labels{
+			Values: map[string]string{
+				"foo": "bar",
+				"bar": "baz",
+			},
+		},
+	})
+
+	// Fetch updated projects.
+	projectsUpdated, err := client.ListProjects(ctx, &admin.ProjectListRequest{})
+	assert.Nil(t, err)
+	assert.NotEmpty(t, projectsUpdated.Projects)
+
+	// Verify that the expected labels have been added to the project.
+	updatedProject := projectsUpdated.Projects[0]
+	labelsMap := updatedProject.Labels
+	fooVal, fooExists := labelsMap.Values["foo"]
+	barVal, barExists := labelsMap.Values["bar"]
+	assert.Equal(t, fooExists, true)
+	assert.Equal(t, fooVal, "bar")
+	assert.Equal(t, barExists, true)
+	assert.Equal(t, barVal, "baz")
+}