From 5cfe8565c8c3fc84aeb6961d1249ed01329f72de Mon Sep 17 00:00:00 2001 From: terryhung Date: Sat, 19 Oct 2024 00:11:49 +0800 Subject: [PATCH] feat: add support for seedProjects with descriptions via seedProjectsWithDetails Signed-off-by: terryhung --- charts/flyte-binary/README.md | 2 + charts/flyte-binary/values.yaml | 8 + cmd/single/config.go | 14 +- cmd/single/start.go | 6 +- .../flyte_sandbox_binary_helm_generated.yaml | 3 + flyteadmin/cmd/entrypoints/migrate.go | 4 +- .../pkg/repositories/config/seed_data.go | 55 ++++++- .../pkg/repositories/config/seed_data_test.go | 145 ++++++++++++++++++ flyteadmin/pkg/server/initialize.go | 2 +- 9 files changed, 225 insertions(+), 14 deletions(-) create mode 100644 flyteadmin/pkg/repositories/config/seed_data_test.go diff --git a/charts/flyte-binary/README.md b/charts/flyte-binary/README.md index e702a7fcd7d..1d0cfd04c69 100644 --- a/charts/flyte-binary/README.md +++ b/charts/flyte-binary/README.md @@ -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` | | diff --git a/charts/flyte-binary/values.yaml b/charts/flyte-binary/values.yaml index 38509251a27..d4c210e9e6c 100644 --- a/charts/flyte-binary/values.yaml +++ b/charts/flyte-binary/values.yaml @@ -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 diff --git a/cmd/single/config.go b/cmd/single/config.go index adbabe7ae58..28cdfdafc08 100644 --- a/cmd/single/config.go +++ b/cmd/single/config.go @@ -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 @@ -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 { diff --git a/cmd/single/start.go b/cmd/single/start.go index 1683fad4e1a..e951aea55cd 100644 --- a/cmd/single/start.go +++ b/cmd/single/start.go @@ -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" @@ -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 } diff --git a/deployment/sandbox-binary/flyte_sandbox_binary_helm_generated.yaml b/deployment/sandbox-binary/flyte_sandbox_binary_helm_generated.yaml index 2704a2eac63..9d5165aa09b 100644 --- a/deployment/sandbox-binary/flyte_sandbox_binary_helm_generated.yaml +++ b/deployment/sandbox-binary/flyte_sandbox_binary_helm_generated.yaml @@ -92,6 +92,9 @@ data: disabled: false seedProjects: - flytesnacks + seedProjectsWithDetails: + - name: flytesnacks + description: Default project setup dataCatalog: disabled: false propeller: diff --git a/flyteadmin/cmd/entrypoints/migrate.go b/flyteadmin/cmd/entrypoints/migrate.go index c7ad6058c51..7ee28150f04 100644 --- a/flyteadmin/cmd/entrypoints/migrate.go +++ b/flyteadmin/cmd/entrypoints/migrate.go @@ -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" ) @@ -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) }, } diff --git a/flyteadmin/pkg/repositories/config/seed_data.go b/flyteadmin/pkg/repositories/config/seed_data.go index b4433548de7..bf09be0ff53 100644 --- a/flyteadmin/pkg/repositories/config/seed_data.go +++ b/flyteadmin/pkg/repositories/config/seed_data.go @@ -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 diff --git a/flyteadmin/pkg/repositories/config/seed_data_test.go b/flyteadmin/pkg/repositories/config/seed_data_test.go new file mode 100644 index 00000000000..3f077dc866d --- /dev/null +++ b/flyteadmin/pkg/repositories/config/seed_data_test.go @@ -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) + } + } + }) + } +} diff --git a/flyteadmin/pkg/server/initialize.go b/flyteadmin/pkg/server/initialize.go index 2d2d7cacc14..42e52719612 100644 --- a/flyteadmin/pkg/server/initialize.go +++ b/flyteadmin/pkg/server/initialize.go @@ -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)