kubevela/pkg/cue/definition/template.go

/*
Copyright 2021 The KubeVela Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package definition

import (
	"context"
	"encoding/json"
	"fmt"
	"sort"
	"strings"
	"time"

	"k8s.io/apiserver/pkg/util/feature"
	"k8s.io/klog/v2"

	"github.com/oam-dev/kubevela/pkg/cue/definition/health"
	"github.com/oam-dev/kubevela/pkg/features"

	velacuex "github.com/oam-dev/kubevela/pkg/cue/cuex"

	"cuelang.org/go/cue"
	cueerrors "cuelang.org/go/cue/errors"
	"github.com/kubevela/pkg/multicluster"

	"github.com/pkg/errors"
	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
	"sigs.k8s.io/controller-runtime/pkg/client"

	"github.com/kubevela/workflow/pkg/cue/model"
	"github.com/kubevela/workflow/pkg/cue/model/sets"
	"github.com/kubevela/workflow/pkg/cue/model/value"
	"github.com/kubevela/workflow/pkg/cue/process"

	velaprocess "github.com/oam-dev/kubevela/pkg/cue/process"
	"github.com/oam-dev/kubevela/pkg/cue/task"
	"github.com/oam-dev/kubevela/pkg/cue/upgrade"
	"github.com/oam-dev/kubevela/pkg/oam"
	"github.com/oam-dev/kubevela/pkg/oam/util"
)

const (
	// OutputFieldName is the name of the struct contains the CR data
	OutputFieldName = velaprocess.OutputFieldName
	// OutputsFieldName is the name of the struct contains the map[string]CR data
	OutputsFieldName = velaprocess.OutputsFieldName
	// PatchFieldName is the name of the struct contains the patch of CR data
	PatchFieldName = "patch"
	// PatchOutputsFieldName is the name of the struct contains the patch of outputs CR data
	PatchOutputsFieldName = "patchOutputs"
	// ErrsFieldName check if errors contained in the cue
	ErrsFieldName = "errs"
	// TemplateContextPrefix is the base prefix for storing templates in context
	TemplateContextPrefix = "template-context-"
)

// GetWorkloadTemplateKey returns the context key for storing workload templates
func GetWorkloadTemplateKey(name string) string {
	return TemplateContextPrefix + "workload-" + name
}

// GetTraitTemplateKey returns the context key for storing trait templates
func GetTraitTemplateKey(name string) string {
	return TemplateContextPrefix + "trait-" + name
}

const (
	// AuxiliaryWorkload defines the extra workload obj from a workloadDefinition,
	// e.g. a workload composed by deployment and service, the service will be marked as AuxiliaryWorkload
	AuxiliaryWorkload = "AuxiliaryWorkload"
)

// AbstractEngine defines Definition's Render interface
type AbstractEngine interface {
	Complete(ctx process.Context, abstractTemplate string, params interface{}) error
	Status(templateContext map[string]interface{}, request *health.StatusRequest) (*health.StatusResult, error)
	GetTemplateContext(ctx process.Context, cli client.Client, accessor util.NamespaceAccessor) (map[string]interface{}, error)
}

type def struct {
	name string
}

type workloadDef struct {
	def
}

// NewWorkloadAbstractEngine create Workload Definition AbstractEngine
func NewWorkloadAbstractEngine(name string) AbstractEngine {
	return &workloadDef{
		def: def{
			name: name,
		},
	}
}

// Complete do workload definition's rendering
func (wd *workloadDef) Complete(ctx process.Context, abstractTemplate string, params interface{}) (retErr error) {
	start := time.Now()
	defer func() {
		status := "ok"
		if retErr != nil {
			status = "error"
		}
		CUERenderDuration.WithLabelValues(string(upgrade.ComponentKind), status).Observe(time.Since(start).Seconds())
	}()

	var paramFile = velaprocess.ParameterFieldName + ": {}"
	if params != nil {
		bt, err := json.Marshal(params)
		if err != nil {
			return errors.WithMessagef(err, "marshal parameter of workload %s", wd.name)
		}
		if string(bt) != "null" {
			paramFile = fmt.Sprintf("%s: %s", velaprocess.ParameterFieldName, string(bt))
		}
	}

	c, err := ctx.BaseContextFile()
	if err != nil {
		return err
	}

	abstractTemplate, _ = upgrade.EnsureCueVersionCompatibility(abstractTemplate, wd.name, upgrade.ComponentKind, upgrade.TemplateAreaMain)

	val, err := velacuex.WorkloadCompiler.Get().CompileString(ctx.GetCtx(), strings.Join([]string{
		renderTemplate(abstractTemplate), paramFile, c,
	}, "\n"))
	if err != nil {
		return errors.WithMessagef(err, "failed to compile workload %s after merge parameter and context", wd.name)
	}

	var userErrors []string
	if errs := val.LookupPath(value.FieldPath(ErrsFieldName)); errs.Exists() {
		if err := errs.Decode(&userErrors); err != nil {
			klog.Warningf("Workload definition '%s' has malformed 'errs' field (expected []string): %v. Custom error reporting will be skipped.", wd.name, err)
		}
	}

	validationErr := val.Validate()

	if validationErr != nil || len(userErrors) > 0 {
		var result strings.Builder
		result.WriteString(fmt.Sprintf("validation failed for workload %s:", wd.name))

		if len(userErrors) > 0 {
			result.WriteString("\n\nUser Errors:\n")
			for _, e := range userErrors {
				result.WriteString(fmt.Sprintf("  %s\n", e))
			}
		}

		if validationErr != nil {
			if fmtErr := FormatCUEError(validationErr, "validation failed for", "workload", wd.name, &val); fmtErr != nil {
				errMsg := fmtErr.Error()
				errMsg = strings.TrimPrefix(errMsg, fmt.Sprintf("validation failed for workload %s:", wd.name))
				result.WriteString(errMsg)
			}
		}

		return errors.New(strings.TrimRight(result.String(), "\n"))
	}
	output := val.LookupPath(value.FieldPath(OutputFieldName))

	base, err := model.NewBase(output)
	if err != nil {
		return errors.WithMessagef(err, "invalid output of workload %s", wd.name)
	}
	if err := ctx.SetBase(base); err != nil {
		return err
	}

	// Store template for error context (use workload-specific key to avoid pollution)
	ctx.PushData(GetWorkloadTemplateKey(wd.name), val)

	// we will support outputs for workload composition, and it will become trait in AppConfig.
	outputs := val.LookupPath(value.FieldPath(OutputsFieldName))
	if !outputs.Exists() {
		return nil
	}

	iter, err := outputs.Fields(cue.Definitions(true), cue.Hidden(true), cue.All())
	if err != nil {
		return errors.WithMessagef(err, "invalid outputs of workload %s", wd.name)
	}
	for iter.Next() {
		if iter.Selector().IsDefinition() || iter.Selector().PkgPath() != "" || iter.IsOptional() {
			continue
		}
		other, err := model.NewOther(iter.Value())
		name := util.GetIteratorLabel(*iter)
		if err != nil {
			return errors.WithMessagef(err, "invalid outputs(%s) of workload %s", name, wd.name)
		}
		if err := ctx.AppendAuxiliaries(process.Auxiliary{Ins: other, Type: AuxiliaryWorkload, Name: name}); err != nil {
			return err
		}
	}
	return nil
}

func withCluster(ctx context.Context, o client.Object) context.Context {
	if cluster := oam.GetCluster(o); cluster != "" {
		return multicluster.WithCluster(ctx, cluster)
	}
	return ctx
}

func (wd *workloadDef) getTemplateContext(ctx process.Context, cli client.Reader, accessor util.NamespaceAccessor) (map[string]interface{}, error) {
	baseLabels := GetBaseContextLabels(ctx)
	var root = initRoot(baseLabels)
	var commonLabels = GetCommonLabels(baseLabels)

	base, assists := ctx.Output()
	componentWorkload, err := base.Unstructured()
	if err != nil {
		return nil, err
	}
	// workload main resource will have a unique label("app.oam.dev/resourceType"="WORKLOAD") in per component/app level
	_ctx := withCluster(ctx.GetCtx(), componentWorkload)
	object, err := getResourceFromObj(_ctx, ctx, componentWorkload, cli, accessor.For(componentWorkload), util.MergeMapOverrideWithDst(map[string]string{
		oam.LabelOAMResourceType: oam.ResourceTypeWorkload,
	}, commonLabels), "")
	if err != nil {
		return nil, err
	}
	root[OutputFieldName] = object
	outputs := make(map[string]interface{})
	for _, assist := range assists {
		if assist.Type != AuxiliaryWorkload {
			continue
		}
		if assist.Name == "" {
			return nil, errors.New("the auxiliary of workload must have a name with format 'outputs.<my-name>'")
		}
		traitRef, err := assist.Ins.Unstructured()
		if err != nil {
			return nil, err
		}
		// AuxiliaryWorkload will have a unique label("trait.oam.dev/resource"="name of outputs") in per component/app level
		_ctx := withCluster(ctx.GetCtx(), traitRef)
		object, err := getResourceFromObj(_ctx, ctx, traitRef, cli, accessor.For(traitRef), util.MergeMapOverrideWithDst(map[string]string{
			oam.TraitTypeLabel: AuxiliaryWorkload,
		}, commonLabels), assist.Name)
		if err != nil {
			return nil, err
		}
		outputs[assist.Name] = object
	}
	if len(outputs) > 0 {
		root[OutputsFieldName] = outputs
	}
	return root, nil
}

// Status get workload status by customStatusTemplate
func (wd *workloadDef) Status(templateContext map[string]interface{}, request *health.StatusRequest) (*health.StatusResult, error) {
	return health.GetStatus(templateContext, request)
}

func (wd *workloadDef) GetTemplateContext(ctx process.Context, cli client.Client, accessor util.NamespaceAccessor) (map[string]interface{}, error) {
	return wd.getTemplateContext(ctx, cli, accessor)
}

type traitDef struct {
	def
}

// NewTraitAbstractEngine create Trait Definition AbstractEngine
func NewTraitAbstractEngine(name string) AbstractEngine {
	return &traitDef{
		def: def{
			name: name,
		},
	}
}

// Complete do trait definition's rendering
// nolint:gocyclo
func (td *traitDef) Complete(ctx process.Context, abstractTemplate string, params interface{}) (retErr error) {
	start := time.Now()
	defer func() {
		status := "ok"
		if retErr != nil {
			status = "error"
		}
		CUERenderDuration.WithLabelValues(string(upgrade.TraitKind), status).Observe(time.Since(start).Seconds())
	}()

	abstractTemplate, _ = upgrade.EnsureCueVersionCompatibility(abstractTemplate, td.name, upgrade.TraitKind, upgrade.TemplateAreaMain)
	buff := abstractTemplate + "\n"
	if params != nil {
		bt, err := json.Marshal(params)
		if err != nil {
			return errors.WithMessagef(err, "marshal parameter of trait %s", td.name)
		}
		if string(bt) != "null" {
			buff += fmt.Sprintf("%s: %s\n", velaprocess.ParameterFieldName, string(bt))
		}
	}

	multiStageEnabled := feature.DefaultMutableFeatureGate.Enabled(features.MultiStageComponentApply)
	var statusBytes []byte
	if multiStageEnabled {
		statusBytes = outputStatusBytes(ctx)
	}

	c, err := ctx.BaseContextFile()
	if err != nil {
		return err
	}

	// When multi-stage is enabled, merge the existing output.status from ctx into the
	// base context so downstream CUE can reference it deterministically.
	if multiStageEnabled {
		c = injectOutputStatusIntoBaseContext(ctx, c, statusBytes)
	}

	buff += c

	val, err := velacuex.WorkloadCompiler.Get().CompileString(ctx.GetCtx(), buff)

	if err != nil {
		return errors.WithMessagef(err, "failed to compile trait %s after merge parameter and context", td.name)
	}

	var userErrors []string
	if errs := val.LookupPath(value.FieldPath(ErrsFieldName)); errs.Exists() {
		if err := errs.Decode(&userErrors); err != nil {
			klog.Warningf("Trait definition '%s' has malformed 'errs' field (expected []string): %v. Custom error reporting will be skipped.", td.name, err)
		}
	}

	validationErr := val.Validate()

	if validationErr != nil || len(userErrors) > 0 {
		var result strings.Builder
		result.WriteString(fmt.Sprintf("validation failed for trait %s:", td.name))

		if len(userErrors) > 0 {
			result.WriteString("\n\nUser Errors:\n")
			for _, e := range userErrors {
				result.WriteString(fmt.Sprintf("  %s\n", e))
			}
		}

		if validationErr != nil {
			if fmtErr := FormatCUEError(validationErr, "validation failed for", "trait", td.name, &val); fmtErr != nil {
				errMsg := fmtErr.Error()
				errMsg = strings.TrimPrefix(errMsg, fmt.Sprintf("validation failed for trait %s:", td.name))
				result.WriteString(errMsg)
			}
		}

		return errors.New(strings.TrimRight(result.String(), "\n"))
	}

	processing := val.LookupPath(value.FieldPath("processing"))
	if processing.Exists() {
		if val, err = task.Process(val); err != nil {
			return errors.WithMessagef(err, "invalid process of trait %s", td.name)
		}
	}
	outputs := val.LookupPath(value.FieldPath(OutputsFieldName))
	if outputs.Exists() {

		iter, err := outputs.Fields(cue.Definitions(true), cue.Hidden(true), cue.All())
		if err != nil {
			return errors.WithMessagef(err, "invalid outputs of trait %s", td.name)
		}
		for iter.Next() {
			if iter.Selector().IsDefinition() || iter.Selector().PkgPath() != "" || iter.IsOptional() {
				continue
			}
			other, err := model.NewOther(iter.Value())
			name := util.GetIteratorLabel(*iter)
			if err != nil {
				return errors.WithMessagef(err, "invalid outputs(resource=%s) of trait %s", name, td.name)
			}
			if err := ctx.AppendAuxiliaries(process.Auxiliary{Ins: other, Type: td.name, Name: name}); err != nil {
				return err
			}
		}
	}

	patcher := val.LookupPath(value.FieldPath(PatchFieldName))
	base, auxiliaries := ctx.Output()
	if patcher.Exists() {
		if base == nil {
			return fmt.Errorf("patch trait %s into an invalid workload", td.name)
		}
		if err := base.Unify(patcher, sets.CreateUnifyOptionsForPatcher(patcher)...); err != nil {
			return errors.WithMessagef(err, "invalid patch trait %s into workload", td.name)
		}
	}
	outputsPatcher := val.LookupPath(value.FieldPath(PatchOutputsFieldName))
	if outputsPatcher.Exists() {
		for _, auxiliary := range auxiliaries {
			target := outputsPatcher.LookupPath(value.FieldPath(auxiliary.Name))
			if !target.Exists() {
				continue
			}
			if err = auxiliary.Ins.Unify(target); err != nil {
				return errors.WithMessagef(err, "trait=%s, to=%s, invalid patch trait into auxiliary workload", td.name, auxiliary.Name)
			}
		}
	}

	return nil
}

func outputStatusBytes(ctx process.Context) []byte {
	var statusBytes []byte
	var outputMap map[string]interface{}
	if output := ctx.GetData(OutputFieldName); output != nil {
		if m, ok := output.(map[string]interface{}); ok {
			outputMap = m
		} else if ptr, ok := output.(*interface{}); ok && ptr != nil {
			if m, ok := (*ptr).(map[string]interface{}); ok {
				outputMap = m
			}
		}

		if outputMap != nil {
			if status, ok := outputMap["status"]; ok {
				if b, err := json.Marshal(status); err == nil {
					statusBytes = b
				}
			}
		}
	}
	return statusBytes
}

func injectOutputStatusIntoBaseContext(ctx process.Context, c string, statusBytes []byte) string {
	if len(statusBytes) > 0 {
		// If output is an empty object, replace it with only the status field without trailing comma.
		emptyOutputMarker := "\"output\":{}"
		if strings.Contains(c, emptyOutputMarker) {
			replacement := fmt.Sprintf("\"output\":{\"status\":%s}", string(statusBytes))
			c = strings.Replace(c, emptyOutputMarker, replacement, 1)
		} else {
			// Otherwise, insert status as the first field and keep the comma to separate from existing fields.
			replacement := fmt.Sprintf("\"output\":{\"status\":%s,", string(statusBytes))
			c = strings.Replace(c, "\"output\":{", replacement, 1)
		}

		// Restore the status field to the current output in ctx.data
		var status interface{}
		if err := json.Unmarshal(statusBytes, &status); err == nil {
			if currentOutput := ctx.GetData(OutputFieldName); currentOutput != nil {
				if currentMap, ok := currentOutput.(map[string]interface{}); ok {
					currentMap["status"] = status
					ctx.PushData(OutputFieldName, currentMap)
				}
			}
		}
	}
	return c
}

// GetCommonLabels will convert context based labels to OAM standard labels
func GetCommonLabels(contextLabels map[string]string) map[string]string {
	var commonLabels = map[string]string{}
	for k, v := range contextLabels {
		switch k {
		case velaprocess.ContextAppName:
			commonLabels[oam.LabelAppName] = v
		case velaprocess.ContextName:
			commonLabels[oam.LabelAppComponent] = v
		case velaprocess.ContextAppRevision:
			commonLabels[oam.LabelAppRevision] = v
		case velaprocess.ContextReplicaKey:
			commonLabels[oam.LabelReplicaKey] = v

		}
	}
	return commonLabels
}

// GetBaseContextLabels get base context labels
func GetBaseContextLabels(ctx process.Context) map[string]string {
	baseLabels := ctx.BaseContextLabels()
	baseLabels[velaprocess.ContextAppName] = ctx.GetData(velaprocess.ContextAppName).(string)
	baseLabels[velaprocess.ContextAppRevision] = ctx.GetData(velaprocess.ContextAppRevision).(string)

	return baseLabels
}

func initRoot(contextLabels map[string]string) map[string]interface{} {
	var root = map[string]interface{}{}
	for k, v := range contextLabels {
		root[k] = v
	}
	return root
}

func renderTemplate(templ string) string {
	return templ + `
context: _
parameter: _
`
}

func (td *traitDef) getTemplateContext(ctx process.Context, cli client.Reader, accessor util.NamespaceAccessor) (map[string]interface{}, error) {
	baseLabels := GetBaseContextLabels(ctx)
	var root = initRoot(baseLabels)
	var commonLabels = GetCommonLabels(baseLabels)
	_, assists := ctx.Output()

	outputs := make(map[string]interface{})
	for _, assist := range assists {
		if assist.Type != td.name {
			continue
		}
		traitRef, err := assist.Ins.Unstructured()
		if err != nil {
			return nil, err
		}
		_ctx := withCluster(ctx.GetCtx(), traitRef)
		object, err := getResourceFromObj(_ctx, ctx, traitRef, cli, accessor.For(traitRef), util.MergeMapOverrideWithDst(map[string]string{
			oam.TraitTypeLabel: assist.Type,
		}, commonLabels), assist.Name)
		if err != nil {
			return nil, err
		}
		outputs[assist.Name] = object
	}
	if len(outputs) > 0 {
		root[OutputsFieldName] = outputs
	}
	return root, nil
}

// Status get trait status by customStatusTemplate
func (td *traitDef) Status(templateContext map[string]interface{}, request *health.StatusRequest) (*health.StatusResult, error) {
	return health.GetStatus(templateContext, request)
}

func (td *traitDef) GetTemplateContext(ctx process.Context, cli client.Client, accessor util.NamespaceAccessor) (map[string]interface{}, error) {
	return td.getTemplateContext(ctx, cli, accessor)
}

func getResourceFromObj(ctx context.Context, pctx process.Context, obj *unstructured.Unstructured, client client.Reader, namespace string, labels map[string]string, outputsResource string) (map[string]interface{}, error) {
	if outputsResource != "" {
		labels[oam.TraitResource] = outputsResource
	}
	if obj.GetName() != "" {
		u, err := util.GetObjectGivenGVKAndName(ctx, client, obj.GroupVersionKind(), namespace, obj.GetName())
		if err != nil {
			return nil, err
		}
		return u.Object, nil
	}
	if ctxName := pctx.GetData(model.ContextName).(string); ctxName != "" {
		u, err := util.GetObjectGivenGVKAndName(ctx, client, obj.GroupVersionKind(), namespace, ctxName)
		if err == nil {
			return u.Object, nil
		}
	}
	list, err := util.GetObjectsGivenGVKAndLabels(ctx, client, obj.GroupVersionKind(), namespace, labels)
	if err != nil {
		return nil, err
	}
	if len(list.Items) == 1 {
		return list.Items[0].Object, nil
	}
	for _, v := range list.Items {
		if v.GetLabels()[oam.TraitResource] == outputsResource {
			return v.Object, nil
		}
	}
	return nil, errors.Errorf("no resources found gvk(%v) labels(%v)", obj.GroupVersionKind(), labels)
}

// FormatCUEError formats CUE errors in a user-friendly grouped format
func FormatCUEError(err error, messagePrefix string, entityType, entityName string, val ...*cue.Value) error {
	var allParamErrors = make(map[string]bool)
	var allTemplateErrors = make(map[string]bool)

	if err != nil {
		errList := cueerrors.Errors(err)
		for _, e := range errList {
			errMsg := e.Error()
			if strings.HasPrefix(errMsg, "parameter.") {
				allParamErrors[errMsg] = true
			} else {
				allTemplateErrors[errMsg] = true
			}
		}

		if len(val) > 0 && val[0] != nil {
			if concreteErr := val[0].Validate(cue.Concrete(true)); concreteErr != nil {
				concreteErrList := cueerrors.Errors(concreteErr)
				for _, e := range concreteErrList {
					errMsg := e.Error()
					if strings.HasPrefix(errMsg, "parameter.") {
						allParamErrors[errMsg] = true
					} else {
						allTemplateErrors[errMsg] = true
					}
				}
			}
		}
	}

	if len(allParamErrors) == 0 && len(allTemplateErrors) == 0 {
		return nil
	}

	var result strings.Builder
	result.WriteString(fmt.Sprintf("%s %s %s:", messagePrefix, entityType, entityName))

	if len(allParamErrors) > 0 {
		result.WriteString("\n\nParameter errors:\n")
		// Sort errors for deterministic output
		paramErrs := make([]string, 0, len(allParamErrors))
		for errMsg := range allParamErrors {
			paramErrs = append(paramErrs, errMsg)
		}
		sort.Strings(paramErrs)
		for _, errMsg := range paramErrs {
			result.WriteString("  " + errMsg + "\n")
		}
	}

	if len(allTemplateErrors) > 0 {
		result.WriteString("\n\nTemplate errors:\n")
		templateErrs := make([]string, 0, len(allTemplateErrors))
		for errMsg := range allTemplateErrors {
			templateErrs = append(templateErrs, errMsg)
		}
		sort.Strings(templateErrs)
		for _, errMsg := range templateErrs {
			result.WriteString("  " + errMsg + "\n")
		}
	}

	return fmt.Errorf("%s", strings.TrimRight(result.String(), "\n"))
}