kubevela/pkg/definition/definition.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 contains some helper functions used in vela CLI
// and vela addon mechanism
package definition

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

	ast2 "github.com/oam-dev/kubevela/pkg/definition/ast"

	"cuelang.org/go/cue"
	"cuelang.org/go/cue/ast"
	"cuelang.org/go/cue/cuecontext"
	"cuelang.org/go/cue/format"
	"cuelang.org/go/cue/parser"
	"cuelang.org/go/encoding/gocode/gocodec"
	"cuelang.org/go/tools/fix"
	"github.com/pkg/errors"
	"k8s.io/apimachinery/pkg/api/meta"
	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/runtime/schema"
	"k8s.io/client-go/rest"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/yaml"

	"github.com/kubevela/pkg/cue/cuex"

	"github.com/oam-dev/kubevela/apis/core.oam.dev/common"
	"github.com/oam-dev/kubevela/apis/core.oam.dev/v1beta1"
	velacue "github.com/oam-dev/kubevela/pkg/cue"
	"github.com/oam-dev/kubevela/pkg/oam"
	"github.com/oam-dev/kubevela/pkg/utils"
	"github.com/oam-dev/kubevela/pkg/utils/filters"
	"github.com/oam-dev/kubevela/pkg/workflow/providers"
)

const (
	// DescriptionKey the key for accessing definition description
	DescriptionKey = "definition.oam.dev/description"
	// AliasKey the key for accessing definition alias
	AliasKey = "definition.oam.dev/alias"
	// UserPrefix defines the prefix of user customized label or annotation
	UserPrefix = "custom.definition.oam.dev/"
)

// the names for different type of definition
const (
	componentDefType    = "component"
	traitDefType        = "trait"
	policyDefType       = "policy"
	workflowStepDefType = "workflow-step"
	workloadDefType     = "workload"
)

var (
	// DefinitionTemplateKeys the keys for accessing definition template
	DefinitionTemplateKeys = []string{"spec", "schematic", "cue", "template"}
	// DefinitionTypeToKind maps the definition types to corresponding kinds
	DefinitionTypeToKind = map[string]string{
		componentDefType:    v1beta1.ComponentDefinitionKind,
		traitDefType:        v1beta1.TraitDefinitionKind,
		policyDefType:       v1beta1.PolicyDefinitionKind,
		workloadDefType:     v1beta1.WorkloadDefinitionKind,
		workflowStepDefType: v1beta1.WorkflowStepDefinitionKind,
	}
	// StringToDefinitionType converts user input to DefinitionType used in DefinitionRevisions
	StringToDefinitionType = map[string]common.DefinitionType{
		// component
		componentDefType: common.ComponentType,
		// trait
		traitDefType: common.TraitType,
		// policy
		policyDefType: common.PolicyType,
		// workflow-step
		workflowStepDefType: common.WorkflowStepType,
	}
	// DefinitionKindToNameLabel records DefinitionRevision types and labels to search its name
	DefinitionKindToNameLabel = map[common.DefinitionType]string{
		common.ComponentType:    oam.LabelComponentDefinitionName,
		common.TraitType:        oam.LabelTraitDefinitionName,
		common.PolicyType:       oam.LabelPolicyDefinitionName,
		common.WorkflowStepType: oam.LabelWorkflowStepDefinitionName,
	}
	// DefinitionKindToType maps the definition kinds to a shorter type
	DefinitionKindToType = map[string]string{
		v1beta1.ComponentDefinitionKind:    componentDefType,
		v1beta1.TraitDefinitionKind:        traitDefType,
		v1beta1.PolicyDefinitionKind:       policyDefType,
		v1beta1.WorkloadDefinitionKind:     workloadDefType,
		v1beta1.WorkflowStepDefinitionKind: workflowStepDefType,
	}
)

// Definition the general struct for handling all kinds of definitions like ComponentDefinition or TraitDefinition
type Definition struct {
	unstructured.Unstructured
}

// SetGVK set the GroupVersionKind of Definition
func (def *Definition) SetGVK(kind string) {
	def.SetGroupVersionKind(schema.GroupVersionKind{
		Group:   v1beta1.Group,
		Version: v1beta1.Version,
		Kind:    kind,
	})
}

// GetType gets the type of Definition
func (def *Definition) GetType() string {
	kind := def.GetKind()
	for k, v := range DefinitionTypeToKind {
		if v == kind {
			return k
		}
	}
	return strings.ToLower(strings.TrimSuffix(kind, "Definition"))
}

// SetType sets the type of Definition
func (def *Definition) SetType(t string) error {
	kind, ok := DefinitionTypeToKind[t]
	if !ok {
		return fmt.Errorf("invalid type %s", t)
	}
	def.SetGVK(kind)
	return nil
}

// ToCUE converts Definition to CUE value (with predefined Definition's cue format)
// nolint:staticcheck
func (def *Definition) ToCUE() (*cue.Value, string, error) {
	annotations := map[string]string{}
	for key, val := range def.GetAnnotations() {
		if strings.HasPrefix(key, UserPrefix) {
			annotations[strings.TrimPrefix(key, UserPrefix)] = val
		}
	}
	alias := def.GetAnnotations()[AliasKey]
	desc := def.GetAnnotations()[DescriptionKey]
	labels := map[string]string{}
	for key, val := range def.GetLabels() {
		if strings.HasPrefix(key, UserPrefix) {
			labels[strings.TrimPrefix(key, UserPrefix)] = val
		}
	}
	spec := map[string]interface{}{}
	for key, val := range def.Object["spec"].(map[string]interface{}) {
		if key != "schematic" {
			spec[key] = val
		}
	}
	obj := map[string]interface{}{
		def.GetName(): map[string]interface{}{
			"type":        def.GetType(),
			"alias":       alias,
			"description": desc,
			"annotations": annotations,
			"labels":      labels,
			"attributes":  spec,
		},
	}
	codec := gocodec.New((*cue.Runtime)(cuecontext.New()), &gocodec.Config{})
	val, err := codec.Decode(obj)
	if err != nil {
		return nil, "", err
	}

	templateString, _, err := unstructured.NestedString(def.Object, DefinitionTemplateKeys...)
	if err != nil {
		return nil, "", err
	}
	templateString, err = formatCUEString(templateString)
	if err != nil {
		return nil, "", err
	}
	return &val, templateString, nil
}

// ToCUEString converts definition to CUE value and then encode to string
func (def *Definition) ToCUEString() (string, error) {
	val, templateString, err := def.ToCUE()
	if err != nil {
		return "", err
	}

	var metadataString string
	syntax := val.Syntax()
	if sl, ok := syntax.(*ast.StructLit); ok {
		field, err := findAndDecodeFieldByLabel(sl, def.GetName())
		if err != nil {
			return "", err
		}
		if field != nil {
			metadataString, err = formatFieldToString(field)
			if err != nil {
				return "", errors.Wrapf(err, "failed to format metadata field %s", def.GetName())
			}
		}
	}

	if err != nil {
		return "", err
	}

	f, err := parser.ParseFile("-", templateString, parser.ParseComments)
	if err != nil {
		return "", errors.Wrapf(err, "failed to parse template cue string")
	}

	// Extract imports before fix.File() clears them
	importPaths := extractImportsFromFile(f)

	f = fix.File(f)
	var templateDecls []ast.Decl
	for _, decl := range f.Decls {
		if _, ok := decl.(*ast.ImportDecl); !ok {
			templateDecls = append(templateDecls, decl)
		}
	}
	var importString string
	if len(importPaths) > 0 {
		// Reconstruct import statements from extracted paths
		var importLines []string
		for _, importPath := range importPaths {
			importLines = append(importLines, fmt.Sprintf("import %s", importPath))
		}
		importString = strings.Join(importLines, "\n") + "\n"
	}
	templateString, err = encodeDeclsToString(templateDecls)
	if err != nil {
		return "", errors.Wrapf(err, "failed to encode template decls")
	}
	templateString = fmt.Sprintf("template: {\n%s}", templateString)

	var completeCUEString string
	if importString != "" {
		completeCUEString = importString + "\n" + metadataString + "\n" + templateString
	} else {
		completeCUEString = metadataString + "\n" + templateString
	}
	if completeCUEString, err = formatCUEString(completeCUEString); err != nil {
		return "", errors.Wrapf(err, "failed to format cue format string")
	}
	return completeCUEString, nil
}

// FromCUE converts CUE value (predefined Definition's cue format) to Definition
// nolint:gocyclo,staticcheck
func (def *Definition) FromCUE(val *cue.Value, templateString string) error {
	if def.Object == nil {
		def.Object = map[string]interface{}{}
	}
	annotations := map[string]string{}
	for k, v := range def.GetAnnotations() {
		if !strings.HasPrefix(k, UserPrefix) && k != DescriptionKey {
			annotations[k] = v
		}
	}
	labels := map[string]string{}
	for k, v := range def.GetLabels() {
		if !strings.HasPrefix(k, UserPrefix) {
			labels[k] = v
		}
	}
	spec, ok := def.Object["spec"].(map[string]interface{})
	if !ok {
		spec = map[string]interface{}{}
	}
	codec := gocodec.New(&cue.Runtime{}, &gocodec.Config{})
	nameFlag := false
	fields, err := val.Fields()
	if err != nil {
		return err
	}
	for fields.Next() {
		definitionName := fields.Label()
		v := fields.Value()
		if nameFlag {
			return fmt.Errorf("duplicated definition name found, %s and %s", def.GetName(), definitionName)
		}
		nameFlag = true
		def.SetName(definitionName)
		_fields, err := v.Fields()
		if err != nil {
			return err
		}
		for _fields.Next() {
			_key := _fields.Label()
			_value := _fields.Value()
			switch _key {
			case "type":
				_type, err := _value.String()
				if err != nil {
					return err
				}
				if err = def.SetType(_type); err != nil {
					return err
				}
			case "alias":
				alias, err := _value.String()
				if err != nil {
					return err
				}
				annotations[AliasKey] = alias
			case "description":
				desc, err := _value.String()
				if err != nil {
					return err
				}
				annotations[DescriptionKey] = desc
			case "annotations":
				var _annotations map[string]string
				if err := codec.Encode(_value, &_annotations); err != nil {
					return err
				}
				for _k, _v := range _annotations {
					if strings.Contains(_k, "oam.dev") {
						annotations[_k] = _v
					} else {
						annotations[UserPrefix+_k] = _v
					}
				}
			case "labels":
				var _labels map[string]string
				if err := codec.Encode(_value, &_labels); err != nil {
					return err
				}
				for _k, _v := range _labels {
					if strings.Contains(_k, "oam.dev") {
						labels[_k] = _v
					} else {
						labels[UserPrefix+_k] = _v
					}
				}
			case "attributes":
				if err := codec.Encode(_value, &spec); err != nil {
					return err
				}
			}
		}
	}
	def.SetAnnotations(annotations)
	def.SetLabels(labels)
	if err := unstructured.SetNestedField(spec, templateString, DefinitionTemplateKeys[1:]...); err != nil {
		return err
	}
	if err = validateSpec(spec, def.GetType()); err != nil {
		return fmt.Errorf("invalid definition spec: %w", err)
	}
	def.Object["spec"] = spec
	return nil
}

func validateSpec(spec map[string]interface{}, t string) error {
	bs, err := json.Marshal(spec)
	if err != nil {
		return err
	}
	var tpl interface{}
	switch t {
	case componentDefType:
		tpl = &v1beta1.ComponentDefinitionSpec{}
	case traitDefType:
		tpl = &v1beta1.TraitDefinitionSpec{}
	case policyDefType:
		tpl = &v1beta1.PolicyDefinitionSpec{}
	case workflowStepDefType:
		tpl = &v1beta1.WorkflowStepDefinitionSpec{}
	default:
	}
	if tpl != nil {
		return utils.StrictUnmarshal(bs, tpl)
	}
	return nil
}

func encodeDeclsToString(decls []ast.Decl) (string, error) {
	bs, err := format.Node(&ast.File{Decls: decls}, format.Simplify())
	if err != nil {
		return "", fmt.Errorf("failed to encode cue: %w", err)
	}
	return strings.TrimSpace(string(bs)) + "\n", nil
}

// FromYAML converts yaml into Definition
func (def *Definition) FromYAML(data []byte) error {
	return yaml.Unmarshal(data, def)
}

// ValidDefinitionTypes return the list of valid definition types
func ValidDefinitionTypes() []string {
	var types []string
	for k := range DefinitionTypeToKind {
		types = append(types, k)
	}
	return types
}

// FromCUEString converts cue string into Definition
func (def *Definition) FromCUEString(cueString string, _ *rest.Config) error {
	// cuectx := cuecontext.New()
	f, err := parser.ParseFile("-", cueString, parser.ParseComments)
	if err != nil {
		return err
	}
	n := fix.File(f)
	var importDecls, metadataDecls, templateDecls []ast.Decl
	for _, decl := range n.Decls {
		if importDecl, ok := decl.(*ast.ImportDecl); ok {
			importDecls = append(importDecls, importDecl)
		} else if field, ok := decl.(*ast.Field); ok {
			label := ""
			switch l := field.Label.(type) {
			case *ast.Ident:
				label = l.Name
			case *ast.BasicLit:
				label = l.Value
			}
			if label == "" {
				return errors.Errorf("found unexpected decl when parsing cue: %v", label)
			}
			if label == "template" {
				if v, ok := field.Value.(*ast.StructLit); ok {
					templateDecls = append(templateDecls, v.Elts...)
				} else {
					return errors.Errorf("unexpected decl found in template: %v", decl)
				}
			} else {
				err := ast2.EncodeMetadata(field)
				if err != nil {
					return errors.Wrapf(err, "failed to parse metadata %s", label)
				}
				metadataDecls = append(metadataDecls, field)
			}
		}
	}
	if len(metadataDecls) == 0 {
		return errors.Errorf("no metadata found, invalid")
	}
	if len(templateDecls) == 0 {
		return errors.Errorf("no template found, invalid")
	}
	var importString, metadataString, templateString string
	if importString, err = encodeDeclsToString(importDecls); err != nil {
		return errors.Wrapf(err, "failed to encode import decls to string")
	}
	if metadataString, err = encodeDeclsToString(metadataDecls); err != nil {
		return errors.Wrapf(err, "failed to encode metadata decls to string")
	}
	// notice that current template decls are concatenated without any blank lines which might be inconsistent with original cue file, but it would not affect the syntax
	if templateString, err = encodeDeclsToString(templateDecls); err != nil {
		return errors.Wrapf(err, "failed to encode template decls to string")
	}

	inst, err := providers.DefaultCompiler.Get().CompileStringWithOptions(context.Background(), metadataString, cuex.DisableResolveProviderFunctions{})
	if err != nil {
		return err
	}
	templateString, err = formatCUEString(importString + templateString)
	if err != nil {
		return err
	}
	if _, err := providers.DefaultCompiler.Get().CompileStringWithOptions(context.Background(), templateString+"\n"+velacue.BaseTemplate, cuex.DisableResolveProviderFunctions{}); err != nil {
		return err
	}
	return def.FromCUE(&inst, templateString)
}

// SearchDefinition search the Definition in k8s by traversing all possible results across types or namespaces
func SearchDefinition(c client.Client, definitionType, namespace string, additionalFilters ...filters.Filter) ([]unstructured.Unstructured, error) {
	ctx := context.Background()
	var kinds []string
	if definitionType != "" {
		kind, ok := DefinitionTypeToKind[definitionType]
		if !ok {
			return nil, fmt.Errorf("invalid definition type %s", kind)
		}
		kinds = []string{kind}
	} else {
		for _, kind := range DefinitionTypeToKind {
			kinds = append(kinds, kind)
		}
	}
	var listOptions []client.ListOption
	if namespace != "" {
		listOptions = []client.ListOption{client.InNamespace(namespace)}
	}
	var definitions []unstructured.Unstructured
	for _, kind := range kinds {
		objs := unstructured.UnstructuredList{}
		objs.SetGroupVersionKind(schema.GroupVersionKind{
			Group:   v1beta1.Group,
			Version: v1beta1.Version,
			Kind:    kind + "List",
		})
		if err := c.List(ctx, &objs, listOptions...); err != nil {
			if meta.IsNoMatchError(err) {
				continue
			}
			return nil, errors.Wrapf(err, "failed to get %s", kind)
		}

		// Apply filters to the object list
		filteredList := filters.ApplyToList(objs, additionalFilters...)

		definitions = append(definitions, filteredList.Items...)
	}
	return definitions, nil
}

// SearchDefinitionRevisions finds DefinitionRevisions.
// Use defName to filter DefinitionRevisions using the name of the underlying Definition.
// Empty defName will keep everything.
// Use defType to only keep DefinitionRevisions of the specified DefinitionType.
// Empty defType will search every possible type.
// Use rev to only keep the revision you want. rev=0 will keep every revision.
func SearchDefinitionRevisions(ctx context.Context, c client.Client, namespace string,
	defName string, defType common.DefinitionType, rev int64) ([]v1beta1.DefinitionRevision, error) {
	var nameLabels []string

	if defName == "" {
		// defName="" means we don't care about the underlying definition names.
		// So, no need to add name labels, just use anything to let the loop run once.
		nameLabels = append(nameLabels, "")
	} else {
		// Since different definitions have different labels for its name, we need to
		// find the corresponding label for definition names, to match names later.
		// Empty defType will give all possible name labels of DefinitionRevisions,
		// so that we can search for DefinitionRevisions of all Definition types.
		for k, v := range DefinitionKindToNameLabel {
			if defType != "" && defType != k {
				continue
			}
			nameLabels = append(nameLabels, v)
		}
	}

	var defRev []v1beta1.DefinitionRevision

	// Search DefinitionRevisions using each possible label
	for _, l := range nameLabels {
		var listOptions []client.ListOption
		if namespace != "" {
			listOptions = append(listOptions, client.InNamespace(namespace))
		}
		// Using name label to find DefinitionRevisions with specified name.
		if defName != "" {
			listOptions = append(listOptions, client.MatchingLabels{
				l: defName,
			})
		}

		objs := v1beta1.DefinitionRevisionList{}
		objs.SetGroupVersionKind(schema.GroupVersionKind{
			Group:   v1beta1.Group,
			Version: v1beta1.Version,
			Kind:    v1beta1.DefinitionRevisionKind,
		})

		// Search for DefinitionRevisions
		if err := c.List(ctx, &objs, listOptions...); err != nil {
			return nil, errors.Wrapf(err, "failed to list DefinitionRevisions of %s", defName)
		}

		for _, dr := range objs.Items {
			// Keep only the specified type
			if defType != "" && defType != dr.Spec.DefinitionType {
				continue
			}
			// Only give the revision that the user wants
			if rev != 0 && rev != dr.Spec.Revision {
				continue
			}
			defRev = append(defRev, dr)
		}
	}

	return defRev, nil
}

// GetDefinitionFromDefinitionRevision will extract the underlying Definition from a DefinitionRevision.
func GetDefinitionFromDefinitionRevision(rev *v1beta1.DefinitionRevision) (*Definition, error) {
	var def *Definition
	var u map[string]interface{}
	var err error

	switch rev.Spec.DefinitionType {
	case common.ComponentType:
		u, err = runtime.DefaultUnstructuredConverter.ToUnstructured(&rev.Spec.ComponentDefinition)
	case common.TraitType:
		u, err = runtime.DefaultUnstructuredConverter.ToUnstructured(&rev.Spec.TraitDefinition)
	case common.PolicyType:
		u, err = runtime.DefaultUnstructuredConverter.ToUnstructured(&rev.Spec.PolicyDefinition)
	case common.WorkflowStepType:
		u, err = runtime.DefaultUnstructuredConverter.ToUnstructured(&rev.Spec.WorkflowStepDefinition)
	default:
		return nil, fmt.Errorf("unsupported definition type: %s", rev.Spec.DefinitionType)
	}

	if err != nil {
		return nil, err
	}

	def = &Definition{Unstructured: unstructured.Unstructured{Object: u}}

	return def, nil
}

// GetDefinitionDefaultSpec returns the default spec of Definition with given kind. This may be implemented with cue in the future.
func GetDefinitionDefaultSpec(kind string) map[string]interface{} {
	switch kind {
	case v1beta1.ComponentDefinitionKind:
		return map[string]interface{}{
			"workload": map[string]interface{}{
				"definition": map[string]interface{}{
					"apiVersion": "<change me> apps/v1",
					"kind":       "<change me> Deployment",
				},
			},
			"schematic": map[string]interface{}{
				"cue": map[string]interface{}{
					"template": "output: {}\nparameter: {}\n",
				},
			},
		}
	case v1beta1.TraitDefinitionKind:
		return map[string]interface{}{
			"appliesToWorkloads": []interface{}{},
			"conflictsWith":      []interface{}{},
			"workloadRefPath":    "",
			"definitionRef":      map[string]interface{}{},
			"podDisruptive":      false,
			"schematic": map[string]interface{}{
				"cue": map[string]interface{}{
					"template": "patch: {}\nparameter: {}\n",
				},
			},
		}
	}
	return map[string]interface{}{}
}

// extractImportsFromFile extracts import paths from an AST file before fix.File() clears them.
// This is necessary because fix.File() removes import declarations that are not directly used.
// Returns a slice of import paths, where named imports are formatted as "name path".
func extractImportsFromFile(f *ast.File) []string {
	var importPaths []string
	for _, decl := range f.Decls {
		if importDecl, ok := decl.(*ast.ImportDecl); ok {
			for _, spec := range importDecl.Specs {
				if spec.Path != nil {
					importPath := spec.Path.Value
					if spec.Name != nil {
						// Handle named imports
						importPaths = append(importPaths, fmt.Sprintf("%s %s", spec.Name.Name, importPath))
					} else {
						importPaths = append(importPaths, importPath)
					}
				}
			}
		}
	}
	return importPaths
}

func formatCUEString(cueString string) (string, error) {
	f, err := parser.ParseFile("-", cueString, parser.ParseComments)
	if err != nil {
		return "", errors.Wrapf(err, "failed to parse file during format cue string")
	}

	// Extract imports before fix.File() clears them
	importPaths := extractImportsFromFile(f)

	n := fix.File(f)

	// Format only non-import declarations
	var nonImportDecls []ast.Decl
	for _, decl := range n.Decls {
		if _, ok := decl.(*ast.ImportDecl); !ok {
			nonImportDecls = append(nonImportDecls, decl)
		}
	}

	var result strings.Builder

	// Add imports first
	if len(importPaths) > 0 {
		for _, importPath := range importPaths {
			result.WriteString(fmt.Sprintf("import %s\n", importPath))
		}
		result.WriteString("\n")
	}

	// Format and add other declarations
	if len(nonImportDecls) > 0 {
		b, err := format.Node(&ast.File{Decls: nonImportDecls}, format.Simplify())
		if err != nil {
			return "", errors.Wrapf(err, "failed to format node during formating cue string")
		}
		result.WriteString(string(b))
	}

	return result.String(), nil
}

func findAndDecodeFieldByLabel(slit *ast.StructLit, targetLabel string) (*ast.Field, error) {
	for _, decl := range slit.Elts {
		field, ok := decl.(*ast.Field)
		if !ok {
			continue
		}
		label := ast2.GetFieldLabel(field.Label)
		if label == targetLabel {
			if err := ast2.DecodeMetadata(field); err != nil {
				return nil, err
			}
			return field, nil
		}
	}
	return nil, nil
}

func formatFieldToString(field *ast.Field) (string, error) {
	b, err := format.Node(field)
	if err != nil {
		return "", errors.Wrap(err, "failed to format field")
	}
	return string(b), nil
}