From ba0b6e2b2845e5c09d40eb0a196091267b1ac0fe Mon Sep 17 00:00:00 2001 From: Yang Yang Date: Wed, 14 Aug 2024 18:33:58 +0800 Subject: [PATCH] feat: add KCL based customized health check --- pkg/apis/api.kusion.io/v1/types.go | 2 + pkg/cmd/apply/apply.go | 24 ++++- pkg/cmd/apply/apply_test.go | 2 +- pkg/engine/printers/generate.go | 9 ++ .../app_configurations_generator.go | 11 +- pkg/util/kcl/kcl.go | 45 ++++++++ pkg/util/kcl/kcl_test.go | 102 ++++++++++++++++++ 7 files changed, 190 insertions(+), 5 deletions(-) create mode 100644 pkg/util/kcl/kcl.go create mode 100644 pkg/util/kcl/kcl_test.go diff --git a/pkg/apis/api.kusion.io/v1/types.go b/pkg/apis/api.kusion.io/v1/types.go index e62a41f3..957f82e3 100644 --- a/pkg/apis/api.kusion.io/v1/types.go +++ b/pkg/apis/api.kusion.io/v1/types.go @@ -265,6 +265,8 @@ const ( EnvViettelCloudProjectID = "VIETTEL_CLOUD_PROJECT_ID" FieldImportedResources = "importedResources" + FieldHealthPolicy = "healthPolicy" + FieldKCLHealthCheckKCL = "health.kcl" ) // BackendConfigs contains the configuration of multiple backends and the current backend. diff --git a/pkg/cmd/apply/apply.go b/pkg/cmd/apply/apply.go index 4f44d978..5407e11a 100644 --- a/pkg/cmd/apply/apply.go +++ b/pkg/cmd/apply/apply.go @@ -34,6 +34,8 @@ import ( "k8s.io/cli-runtime/pkg/genericiooptions" "k8s.io/kubectl/pkg/util/templates" + "gopkg.in/yaml.v3" + apiv1 "kusionstack.io/kusion/pkg/apis/api.kusion.io/v1" v1 "kusionstack.io/kusion/pkg/apis/status/v1" "kusionstack.io/kusion/pkg/cmd/generate" @@ -48,6 +50,7 @@ import ( runtimeinit "kusionstack.io/kusion/pkg/engine/runtime/init" "kusionstack.io/kusion/pkg/log" "kusionstack.io/kusion/pkg/util/i18n" + "kusionstack.io/kusion/pkg/util/kcl" "kusionstack.io/kusion/pkg/util/pretty" "kusionstack.io/kusion/pkg/util/signal" "kusionstack.io/kusion/pkg/util/terminal" @@ -758,7 +761,7 @@ func Watch( // Setup a go-routine to concurrently watch K8s and TF resources. if res.Type == apiv1.Kubernetes { - go watchK8sResources(id, w.Watchers, table, tables, dryRun) + go watchK8sResources(id, w.Watchers, table, tables, dryRun, res.Extensions[apiv1.FieldHealthPolicy]) } else if res.Type == apiv1.Terraform { go watchTFResources(id, w.TFWatcher, table, dryRun) } else { @@ -889,6 +892,7 @@ func watchK8sResources( table *printers.Table, tables map[string]*printers.Table, dryRun bool, + healthPolicy interface{}, ) { defer func() { var err error @@ -930,7 +934,23 @@ func watchK8sResources( } else { // Restore to actual type target := printers.Convert(o) - detail, ready = printers.Generate(target) + + // Check reconcile status with customized health policy + if healthPolicy != nil { + if code, ok := kcl.ConvertKCLCode(healthPolicy); ok { + resByte, err := yaml.Marshal(o.Object) + if err != nil { + log.Error(err) + return + } + detail, ready = printers.PrintCustomizedHealthCheck(code, resByte) + } else { + detail, ready = printers.Generate(target) + } + // Check reconcile status with default setup + } else { + detail, ready = printers.Generate(target) + } } // Mark ready for breaking loop diff --git a/pkg/cmd/apply/apply_test.go b/pkg/cmd/apply/apply_test.go index 7641a8d7..dc6bfc66 100644 --- a/pkg/cmd/apply/apply_test.go +++ b/pkg/cmd/apply/apply_test.go @@ -415,7 +415,7 @@ func TestWatchK8sResources(t *testing.T) { id: table, } - watchK8sResources(id, chs, table, tables, true) + watchK8sResources(id, chs, table, tables, true, nil) assert.Equal(t, true, table.AllCompleted()) }) diff --git a/pkg/engine/printers/generate.go b/pkg/engine/printers/generate.go index b208a6bb..a82f03a3 100644 --- a/pkg/engine/printers/generate.go +++ b/pkg/engine/printers/generate.go @@ -4,6 +4,7 @@ import ( "k8s.io/apimachinery/pkg/runtime" "kusionstack.io/kusion/pkg/engine/printers/printer" + "kusionstack.io/kusion/pkg/util/kcl" ) var tg = printer.NewTableGenerator() @@ -15,3 +16,11 @@ func init() { func Generate(obj runtime.Object) (string, bool) { return tg.GenerateTable(obj) } + +func PrintCustomizedHealthCheck(healthPolicyCode string, resource []byte) (string, bool) { + err := kcl.RunKCLHealthCheck(healthPolicyCode, resource) + if err != nil { + return "Reconciling...", false + } + return "Reconciled", true +} diff --git a/pkg/modules/generators/app_configurations_generator.go b/pkg/modules/generators/app_configurations_generator.go index 9ef410f2..d0cd6ce9 100644 --- a/pkg/modules/generators/app_configurations_generator.go +++ b/pkg/modules/generators/app_configurations_generator.go @@ -427,7 +427,7 @@ func (g *appConfigurationGenerator) callModules(projectModuleConfigs map[string] if err != nil { return nil, nil, nil, err } - + healthPolicy := config.platformConfig[v1.FieldHealthPolicy] // parse module result // if only one resource exists in the workload module, it is the workload if workloadKey == t && len(response.Resources) == 1 { @@ -436,6 +436,10 @@ func (g *appConfigurationGenerator) callModules(projectModuleConfigs map[string] if err != nil { return nil, nil, nil, err } + // Add healthPolicy to resource extensions + if healthPolicy != nil { + workload.Extensions[v1.FieldHealthPolicy] = healthPolicy + } } else { for _, res := range response.Resources { temp := &v1.Resource{} @@ -443,7 +447,10 @@ func (g *appConfigurationGenerator) callModules(projectModuleConfigs map[string] if err != nil { return nil, nil, nil, err } - + // Add healthPolicy to resource extensions + if healthPolicy != nil { + temp.Extensions[v1.FieldHealthPolicy] = healthPolicy + } // filter out workload if workloadKey == t && temp.Extensions[isWorkload] == "true" { workload = temp diff --git a/pkg/util/kcl/kcl.go b/pkg/util/kcl/kcl.go new file mode 100644 index 00000000..cf646b2e --- /dev/null +++ b/pkg/util/kcl/kcl.go @@ -0,0 +1,45 @@ +package kcl + +import ( + "fmt" + + "kcl-lang.io/kcl-go/pkg/kcl" + "kcl-lang.io/kcl-go/pkg/tools/format" + v1 "kusionstack.io/kusion/pkg/apis/api.kusion.io/v1" +) + +const ref = "res" + +// Assemble and format the whole KCL code with yaml integration. +func assembleKCLHealthCheck(healthPolicyCode string, resource []byte) (string, error) { + yamlStr := fmt.Sprintf(` + import yaml + %s = yaml.decode(%q) + `, ref, resource) + kclCode := yamlStr + healthPolicyCode + kclFormatted, err := format.FormatCode(kclCode) + if err != nil { + return "", err + } + return string(kclFormatted), nil +} + +// Run health check with KCL health policy during apply. +func RunKCLHealthCheck(healthPolicyCode string, resource []byte) error { + kclCode, err := assembleKCLHealthCheck(healthPolicyCode, resource) + if err != nil { + return err + } + _, err = kcl.Run("", kcl.WithCode(kclCode)) + return err +} + +// Get KCL code from extensions of the resource in Spec. +func ConvertKCLCode(healthPolicy interface{}) (string, bool) { + if hp, ok := healthPolicy.(v1.GenericConfig); ok { + if code, ok := hp[v1.FieldKCLHealthCheckKCL].(string); ok { + return code, true + } + } + return "", false +} diff --git a/pkg/util/kcl/kcl_test.go b/pkg/util/kcl/kcl_test.go new file mode 100644 index 00000000..9ce34ec9 --- /dev/null +++ b/pkg/util/kcl/kcl_test.go @@ -0,0 +1,102 @@ +package kcl + +import ( + "testing" + + v1 "kusionstack.io/kusion/pkg/apis/api.kusion.io/v1" +) + +func TestAssembleKCLHealthCheck(t *testing.T) { + tests := []struct { + name string + hpCode string + resource []byte + want string + expectError bool + }{ + { + name: "Valid input", + hpCode: "assert res.a == res.b", + resource: []byte("this is resource"), + want: "import yaml\n\nres = yaml.decode(\"this is resource\")\nassert res.a == res.b\n", + expectError: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, err := assembleKCLHealthCheck(tt.hpCode, tt.resource) + if (err != nil) != tt.expectError { + t.Errorf("assembleKCLHealthCheck() error = %v, expectError %v", err, tt.expectError) + return + } + if got != tt.want { + t.Errorf("assembleKCLHealthCheck() got = %v, want %v", got, tt.want) + } + }) + } +} + +func TestRunKCLHealthCheck(t *testing.T) { + tests := []struct { + name string + hpCode string + resource []byte + expectError bool + }{ + { + name: "Valid input", + hpCode: "a = \"this is health policy\"", + resource: []byte("this is resource"), + expectError: false, + }, + // Add more test cases if needed + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := RunKCLHealthCheck(tt.hpCode, tt.resource) + if (err != nil) != tt.expectError { + t.Errorf("RunKCLHealthCheck() error = %v, expectError %v", err, tt.expectError) + } + }) + } +} + +func TestConvertKCLCode(t *testing.T) { + tests := []struct { + name string + healthPolicy interface{} + want string + expectOk bool + }{ + { + name: "Valid KCL code in health policy", + healthPolicy: v1.GenericConfig{ + v1.FieldKCLHealthCheckKCL: "assert res.a == res.b", + }, + want: "assert res.a == res.b", + expectOk: true, + }, + { + name: "No KCL code in health policy", + healthPolicy: v1.GenericConfig{ + "other_field": "other_value", + }, + want: "", + expectOk: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, ok := ConvertKCLCode(tt.healthPolicy) + if got != tt.want { + t.Errorf("ConvertKCLCode() got = %v, want %v", got, tt.want) + } + if ok != tt.expectOk { + t.Errorf("ConvertKCLCode() ok = %v, expectOk %v", ok, tt.expectOk) + } + }) + } +}