From 12797403663e96eef4a51339388424c7c25d040d Mon Sep 17 00:00:00 2001 From: "terry.hung" Date: Fri, 25 Oct 2024 14:28:39 +0800 Subject: [PATCH] feat: add support for seedProjects with descriptions via seedProjectsWithDetails (#5864) * feat: add support for seedProjects with descriptions via seedProjectsWithDetails Signed-off-by: terryhung * fix: update the logging for seeding projects Signed-off-by: terry.hung * fix: the helm lint error and add more unit test Signed-off-by: terry.hung * feat: add more unitest and fix the helm diff Signed-off-by: terry.hung --------- Signed-off-by: terryhung Signed-off-by: terry.hung --- 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 | 5 +- flyteadmin/cmd/entrypoints/migrate.go | 4 +- .../pkg/repositories/config/seed_data.go | 55 +++- .../pkg/repositories/config/seed_data_test.go | 290 ++++++++++++++++++ flyteadmin/pkg/server/initialize.go | 2 +- 9 files changed, 371 insertions(+), 15 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 e702a7fcd7..e7df1018db 100644 --- a/charts/flyte-binary/README.md +++ b/charts/flyte-binary/README.md @@ -116,6 +116,8 @@ Chart for basic single Flyte executable deployment | flyte-core-components.admin.disableClusterResourceManager | bool | `false` | | | flyte-core-components.admin.disableScheduler | bool | `false` | | | flyte-core-components.admin.disabled | bool | `false` | | +| flyte-core-components.admin.seedProjectsWithDetails[0].description | string | `"Default project setup."` | | +| flyte-core-components.admin.seedProjectsWithDetails[0].name | string | `"flytesnacks"` | | | flyte-core-components.admin.seedProjects[0] | string | `"flytesnacks"` | | | flyte-core-components.dataCatalog.disabled | bool | `false` | | | flyte-core-components.propeller.disableWebhook | bool | `false` | | diff --git a/charts/flyte-binary/values.yaml b/charts/flyte-binary/values.yaml index 38509251a2..eee01d16c6 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 adbabe7ae5..28cdfdafc0 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 1683fad4e1..f9e38a9626 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", seedProjects) + 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 2704a2eac6..6fafa61550 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: + - description: Default project setup. + name: flytesnacks dataCatalog: disabled: false propeller: @@ -360,7 +363,7 @@ spec: app.kubernetes.io/instance: flyte app.kubernetes.io/component: flyte-binary annotations: - checksum/configuration: faaefbd3b3b2ddfd4e718bd77c02c632c75e7111dad0a6e25dc415dc88add73f + checksum/configuration: 886440a42b3eeec802cfe60d37885f69e35ffd83e53e625b3c877da5e8c7eb38 checksum/configuration-secret: d5d93f4e67780b21593dc3799f0f6682aab0765e708e4020939975d14d44f929 checksum/cluster-resource-templates: 7dfa59f3d447e9c099b8f8ffad3af466fecbc9cf9f8c97295d9634254a55d4ae spec: diff --git a/flyteadmin/cmd/entrypoints/migrate.go b/flyteadmin/cmd/entrypoints/migrate.go index c7ad6058c5..7ee28150f0 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 b4433548de..bf09be0ff5 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 0000000000..b937b30592 --- /dev/null +++ b/flyteadmin/pkg/repositories/config/seed_data_test.go @@ -0,0 +1,290 @@ +package config + +import ( + "testing" + + mocket "github.com/Selvatico/go-mocket" + "github.com/stretchr/testify/assert" + "gorm.io/gorm" +) + +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) + } + } + }) + } +} + +func TestUniqueProjectsFromNames(t *testing.T) { + tests := []struct { + name string + names []string + want []SeedProject + }{ + { + name: "Empty input", + names: []string{}, + want: []SeedProject{}, + }, + { + name: "Single name", + names: []string{"project1"}, + want: []SeedProject{ + { + Name: "project1", + Description: "project1 description", + }, + }, + }, + { + name: "Multiple unique names", + names: []string{"project1", "project2", "project3"}, + want: []SeedProject{ + { + Name: "project1", + Description: "project1 description", + }, + { + Name: "project2", + Description: "project2 description", + }, + { + Name: "project3", + Description: "project3 description", + }, + }, + }, + { + name: "Duplicate names", + names: []string{"project1", "project1", "project2", "project2", "project3"}, + want: []SeedProject{ + { + Name: "project1", + Description: "project1 description", + }, + { + Name: "project2", + Description: "project2 description", + }, + { + Name: "project3", + Description: "project3 description", + }, + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := UniqueProjectsFromNames(tt.names) + + 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 + } + + // Compare contents + 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) + } + } + }) + } +} + +func TestSeedProjects(t *testing.T) { + gormDb := GetDbForTest(t) + defer mocket.Catcher.Reset() + + mocket.Catcher.Reset() + mocket.Catcher.NewMock().WithQuery(`SELECT * FROM "projects"`).WithReply([]map[string]interface{}{}) + + projects := []SeedProject{ + { + Name: "Project 1", + Description: "New Description", + }, + } + + // Execute + err := SeedProjects(gormDb, projects) + + // Assert + assert.NoError(t, err) +} + +func TestSeedProjectsWithDuplicateKey(t *testing.T) { + gormDb := GetDbForTest(t) + defer mocket.Catcher.Reset() + + // Mock the SELECT query for existence check + mocket.Catcher.Reset() + mocket.Catcher.NewMock().WithQuery(`INSERT INTO "projects"`).WithError(gorm.ErrDuplicatedKey) + + projects := []SeedProject{ + { + Name: "Project 1", + Description: "New Description", + }, + } + + // Execute + err := SeedProjects(gormDb, projects) + + // Assert + assert.Error(t, err) + +} diff --git a/flyteadmin/pkg/server/initialize.go b/flyteadmin/pkg/server/initialize.go index 2d2d7cacc1..42e5271961 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)