Skip to content

Commit

Permalink
feat: add support for seedProjects with descriptions via seedProjects…
Browse files Browse the repository at this point in the history
…WithDetails

Signed-off-by: terryhung <[email protected]>
  • Loading branch information
Terryhung committed Oct 18, 2024
1 parent 38af4ec commit 5cfe856
Show file tree
Hide file tree
Showing 9 changed files with 225 additions and 14 deletions.
2 changes: 2 additions & 0 deletions charts/flyte-binary/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -117,6 +117,8 @@ Chart for basic single Flyte executable deployment
| flyte-core-components.admin.disableScheduler | bool | `false` | |
| flyte-core-components.admin.disabled | bool | `false` | |
| flyte-core-components.admin.seedProjects[0] | string | `"flytesnacks"` | |
| flyte-core-components.admin.seedProjectsWithDetails.name | string | `"flytesnacks"` | |
| flyte-core-components.admin.seedProjectsWithDetails.description | string | `"Default project setup"` | |
| flyte-core-components.dataCatalog.disabled | bool | `false` | |
| flyte-core-components.propeller.disableWebhook | bool | `false` | |
| flyte-core-components.propeller.disabled | bool | `false` | |
Expand Down
8 changes: 8 additions & 0 deletions charts/flyte-binary/values.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,14 @@ flyte-core-components:
# seedProjects flyte projects to create by default
seedProjects:
- flytesnacks
# seedProjectsWithDetails flyte projects to create by default with description
# If there is an overlap between seedProjects and seedProjectsWithDetails,
# the description provided in seedProjectsWithDetails will take precedence.
# For seedProjects without a corresponding description in seedProjectsWithDetails,
# a default description will be auto-generated for the project.
seedProjectsWithDetails:
- name: flytesnacks
description: Default project setup."
# propeller Configuration to disable propeller or any of its components
propeller:
# disabled Disables flytepropeller
Expand Down
14 changes: 9 additions & 5 deletions cmd/single/config.go
Original file line number Diff line number Diff line change
@@ -1,6 +1,9 @@
package single

import "github.com/flyteorg/flyte/flytestdlib/config"
import (
adminRepositoriesConfig "github.com/flyteorg/flyte/flyteadmin/pkg/repositories/config"
"github.com/flyteorg/flyte/flytestdlib/config"
)

//go:generate pflags Config --default-var=DefaultConfig

Expand All @@ -21,10 +24,11 @@ type Propeller struct {
}

type Admin struct {
Disabled bool `json:"disabled" pflag:",Disables flyteadmin in the single binary mode"`
DisableScheduler bool `json:"disableScheduler" pflag:",Disables Native scheduler in the single binary mode"`
DisableClusterResourceManager bool `json:"disableClusterResourceManager" pflag:",Disables Cluster resource manager"`
SeedProjects []string `json:"seedProjects" pflag:",flyte projects to create by default."`
Disabled bool `json:"disabled" pflag:",Disables flyteadmin in the single binary mode"`
DisableScheduler bool `json:"disableScheduler" pflag:",Disables Native scheduler in the single binary mode"`
DisableClusterResourceManager bool `json:"disableClusterResourceManager" pflag:",Disables Cluster resource manager"`
SeedProjects []string `json:"seedProjects" pflag:",flyte projects to create by default."`
SeedProjectsWithDetails []adminRepositoriesConfig.SeedProject `json:"seedProjectsWithDetails" pflag:",,Detailed configuration for Flyte projects to be created by default."`
}

type DataCatalog struct {
Expand Down
6 changes: 4 additions & 2 deletions cmd/single/start.go
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@ import (
datacatalog "github.com/flyteorg/flyte/datacatalog/pkg/rpc/datacatalogservice"
"github.com/flyteorg/flyte/flyteadmin/pkg/clusterresource"
"github.com/flyteorg/flyte/flyteadmin/pkg/common"
adminRepositoriesConfig "github.com/flyteorg/flyte/flyteadmin/pkg/repositories/config"
"github.com/flyteorg/flyte/flyteadmin/pkg/runtime"
adminServer "github.com/flyteorg/flyte/flyteadmin/pkg/server"
"github.com/flyteorg/flyte/flyteadmin/plugins"
Expand Down Expand Up @@ -75,8 +76,9 @@ func startAdmin(ctx context.Context, cfg Admin) error {
if len(cfg.SeedProjects) != 0 {
projects = cfg.SeedProjects
}
logger.Infof(ctx, "Seeding default projects...", projects)
if err := adminServer.SeedProjects(ctx, projects); err != nil {
seedProjects := adminRepositoriesConfig.MergeSeedProjectsWithUniqueNames(projects, cfg.SeedProjectsWithDetails)
logger.Infof(ctx, "Seeding default projects... %v", projects)
if err := adminServer.SeedProjects(ctx, seedProjects); err != nil {
return err
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -92,6 +92,9 @@ data:
disabled: false
seedProjects:
- flytesnacks
seedProjectsWithDetails:
- name: flytesnacks
description: Default project setup
dataCatalog:
disabled: false
propeller:
Expand Down
4 changes: 3 additions & 1 deletion flyteadmin/cmd/entrypoints/migrate.go
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ import (
"github.com/spf13/cobra"
_ "gorm.io/driver/postgres" // Required to import database driver.

"github.com/flyteorg/flyte/flyteadmin/pkg/repositories/config"
"github.com/flyteorg/flyte/flyteadmin/pkg/server"
)

Expand Down Expand Up @@ -40,7 +41,8 @@ var seedProjectsCmd = &cobra.Command{
Short: "Seed projects in the database.",
RunE: func(cmd *cobra.Command, args []string) error {
ctx := context.Background()
return server.SeedProjects(ctx, args)
seedProjects := config.UniqueProjectsFromNames(args)
return server.SeedProjects(ctx, seedProjects)
},
}

Expand Down
55 changes: 50 additions & 5 deletions flyteadmin/pkg/repositories/config/seed_data.go
Original file line number Diff line number Diff line change
Expand Up @@ -10,16 +10,61 @@ import (
"github.com/flyteorg/flyte/flytestdlib/logger"
)

type SeedProject struct {
Name string `json:"name" pflag:",Name of flyte project to create"`
Description string `json:"description" pflag:",Description of flyte project to create"`
}

func UniqueProjectsFromNames(names []string) []SeedProject {
return uniqueProjects(names, nil)
}

// MergeSeedProjectsWithUniqueNames merges seed projects from names and details while maintaining uniqueness
func MergeSeedProjectsWithUniqueNames(seedProjects []string, seedProjectsWithDetails []SeedProject) []SeedProject {
return uniqueProjects(seedProjects, seedProjectsWithDetails)
}

func uniqueProjects(seedProjects []string, seedProjectsWithDetails []SeedProject) []SeedProject {
// Track unique project names
seen := make(map[string]struct{})

// Create the final result slice
var combinedProjects []SeedProject

// First, add all projects from SeedProjectsWithDetails to the map
for _, project := range seedProjectsWithDetails {
// Handle the duplication
if _, exists := seen[project.Name]; !exists {
seen[project.Name] = struct{}{}
combinedProjects = append(combinedProjects, project)
}
}

// Process SeedProjects
for _, projectName := range seedProjects {
// Check if project not exists in SeedProjectsWithDetails
if _, exists := seen[projectName]; !exists {
seen[projectName] = struct{}{}
combinedProjects = append(combinedProjects, SeedProject{
Name: projectName,
Description: fmt.Sprintf("%s description", projectName),
})
}
}

return combinedProjects
}

// Returns a function to seed the database with default values.
func SeedProjects(db *gorm.DB, projects []string) error {
func SeedProjects(db *gorm.DB, projects []SeedProject) error {
tx := db.Begin()
for _, project := range projects {
projectModel := models.Project{
Identifier: project,
Name: project,
Description: fmt.Sprintf("%s description", project),
Identifier: project.Name,
Name: project.Name,
Description: project.Description,
}
if err := tx.Where(models.Project{Identifier: project}).Omit("id").FirstOrCreate(&projectModel).Error; err != nil {
if err := tx.Where(models.Project{Identifier: project.Name}).Omit("id").FirstOrCreate(&projectModel).Error; err != nil {
logger.Warningf(context.Background(), "failed to save project [%s]", project)
tx.Rollback()
return err
Expand Down
145 changes: 145 additions & 0 deletions flyteadmin/pkg/repositories/config/seed_data_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,145 @@
package config

import "testing"

func TestMergeSeedProjectsWithUniqueNames(t *testing.T) {
tests := []struct {
name string
seedProjects []string
seedProjectsWithDetails []SeedProject
want []SeedProject
}{
{
name: "Empty inputs",
seedProjects: []string{},
seedProjectsWithDetails: []SeedProject{},
want: []SeedProject{},
},
{
name: "Empty inputs",
seedProjects: []string{},
seedProjectsWithDetails: nil,
want: []SeedProject{},
},
{
name: "Only seedProjects",
seedProjects: []string{"project1", "project2"},
seedProjectsWithDetails: nil,
want: []SeedProject{
{Name: "project1", Description: "project1 description"},
{Name: "project2", Description: "project2 description"},
},
},
{
name: "Only seedProjectsWithDetails",
seedProjects: []string{},
seedProjectsWithDetails: []SeedProject{
{Name: "project1", Description: "custom description 1"},
{Name: "project2", Description: "custom description 2"},
},
want: []SeedProject{
{Name: "project1", Description: "custom description 1"},
{Name: "project2", Description: "custom description 2"},
},
},
{
name: "Mixed with no overlaps",
seedProjects: []string{"project1", "project2"},
seedProjectsWithDetails: []SeedProject{
{Name: "project3", Description: "custom description 3"},
{Name: "project4", Description: "custom description 4"},
},
want: []SeedProject{
{Name: "project3", Description: "custom description 3"},
{Name: "project4", Description: "custom description 4"},
{Name: "project1", Description: "project1 description"},
{Name: "project2", Description: "project2 description"},
},
},
{
name: "Mixed with overlaps",
seedProjects: []string{"project1", "project2", "project3"},
seedProjectsWithDetails: []SeedProject{
{Name: "project2", Description: "custom description 2"},
{Name: "project3", Description: "custom description 3"},
},
want: []SeedProject{
{Name: "project2", Description: "custom description 2"},
{Name: "project3", Description: "custom description 3"},
{Name: "project1", Description: "project1 description"},
},
},
{
name: "Duplicates in seedProjects",
seedProjects: []string{"project1", "project1", "project2"},
seedProjectsWithDetails: []SeedProject{
{Name: "project3", Description: "custom description 3"},
},
want: []SeedProject{
{Name: "project3", Description: "custom description 3"},
{Name: "project1", Description: "project1 description"},
{Name: "project2", Description: "project2 description"},
},
},
{
name: "Duplicates in seedProjectsWithDetails",
seedProjects: []string{"project1"},
seedProjectsWithDetails: []SeedProject{
{Name: "project2", Description: "custom description 2"},
{Name: "project2", Description: "duplicate description 2"},
},
want: []SeedProject{
{Name: "project2", Description: "custom description 2"},
{Name: "project1", Description: "project1 description"},
},
},
{
name: "All duplicates",
seedProjects: []string{"project1", "project1", "project2"},
seedProjectsWithDetails: []SeedProject{
{Name: "project1", Description: "custom description 1"},
{Name: "project2", Description: "custom description 2"},
{Name: "project2", Description: "duplicate description 2"},
},
want: []SeedProject{
{Name: "project1", Description: "custom description 1"},
{Name: "project2", Description: "custom description 2"},
},
},
}

for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got := MergeSeedProjectsWithUniqueNames(tt.seedProjects, tt.seedProjectsWithDetails)

// Check length
if len(got) != len(tt.want) {
t.Errorf("length mismatch: got %d projects, want %d projects", len(got), len(tt.want))
return
}

gotMap := make(map[string]string)
for _, project := range got {
gotMap[project.Name] = project.Description
}
wantMap := make(map[string]string)
for _, project := range tt.want {
wantMap[project.Name] = project.Description
}

for name, wantDesc := range wantMap {
if gotDesc, exists := gotMap[name]; !exists {
t.Errorf("missing project %q in result", name)
} else if gotDesc != wantDesc {
t.Errorf("project %q description mismatch: got %q, want %q", name, gotDesc, wantDesc)
}
}

for name := range gotMap {
if _, exists := wantMap[name]; !exists {
t.Errorf("unexpected project %q in result", name)
}
}
})
}
}
2 changes: 1 addition & 1 deletion flyteadmin/pkg/server/initialize.go
Original file line number Diff line number Diff line change
Expand Up @@ -67,7 +67,7 @@ func Rollback(ctx context.Context) error {
}

// SeedProjects creates a set of given projects in the DB
func SeedProjects(ctx context.Context, projects []string) error {
func SeedProjects(ctx context.Context, projects []config.SeedProject) error {
return withDB(ctx, func(db *gorm.DB) error {
if err := config.SeedProjects(db, projects); err != nil {
return fmt.Errorf("could not add projects to database with err: %v", err)
Expand Down

0 comments on commit 5cfe856

Please sign in to comment.