gotypes.go

package jennies

import (
	"context"
	"fmt"
	"go/format"
	"path"
	"strings"

	"cuelang.org/go/cue"
	"github.com/grafana/codejen"
	"github.com/grafana/cog"

	"github.com/grafana/grafana-app-sdk/codegen"
	"github.com/grafana/grafana-app-sdk/codegen/templates"
)

const GoTypesMaxDepth = 5

// GoTypes is a Jenny for turning a codegen.AppManifest into go types according to its codegen settings.
type GoTypes struct {
	// GenerateOnlyCurrent should be set to true if you only want to generate code for the kind.Properties().Current version.
	// This will affect the package and path(s) of the generated file(s).
	GenerateOnlyCurrent bool

	// Depth represents the tree depth for creating go types from fields. A Depth of 0 will return one go type
	// (plus any definitions used by that type), a Depth of 1 will return a file with a go type for each top-level field
	// (plus any definitions encompassed by each type), etc. Note that types are _not_ generated for fields above the Depth
	// level--i.e. a Depth of 1 will generate go types for each field within the KindVersion.Schema, but not a type for the
	// Schema itself. Because Depth results in recursive calls, the highest value is bound to a max of GoTypesMaxDepth.
	Depth int

	// NamingDepth determines how types are named in relation to Depth. If Depth <= NamingDepth, the go types are named
	// using the field name of the type. Otherwise, Names used are prefixed by field names between Depth and NamingDepth.
	// Typically, a value of 0 is "safest" for NamingDepth, as it prevents overlapping names for types.
	// However, if you know that your fields have unique names up to a certain depth, you may configure this to be higher.
	NamingDepth int

	// AddKubernetesCodegen toggles whether kubernetes codegen comments are added to the go types generated by this jenny.
	// The codegen comments can then be used with the OpenAPI jenny, or with the kubernetes codegen tooling.
	AddKubernetesCodegen bool

	// GroupByKind determines whether kinds are grouped by GroupVersionKind or just GroupVersion.
	// If GroupByKind is true, generated paths are <kind>/<version>/<file>, instead of the default <version>/<file>.
	// When GroupByKind is false, subresource types (such as spec and status) are prefixed with the kind name,
	// i.e. generating FooSpec instead of Spec for kind.Name() = "Foo" and Depth=1
	GroupByKind bool

	// AnyAsInterface determines whether to use `interface{}` instead of `any` in generated go code.
	// If true, `interface{}` will be used instead of `any`.
	AnyAsInterface bool

	// ExcludeFields will exclude fields with the provided name at the provided Depth value from having types generated for them.
	// This only applies to fields at the Depth where types will be generated, and will not be applied to field names
	// within a type already having a go type generated for it. This only applies at Depth > 0, as at Depth = 0,
	// only one type is generated for the entire CUE object.
	// Practically, since current implementation is to generate one go file per field at Depth, this excludes
	// a go file for the type from being generated.
	// ex:
	//
	//  {
	//     Foo: string
	//     Bar: string
	//  }
	// With an ExcludeFields of `Bar` will only generate a file and go type for `Foo`.
	ExcludeFields []string

	// OpenAPINamer will be called on each generated type's name to create an OpenAPIModelName() function
	// for the type to implement k8s.io/kube-openapi/pkg/util/OpenAPIModelNamer.
	// If nil, types will not implement OpenAPIModelNamer, which will cause problems with
	// apiservers if using the default apiserver.AppInstaller.
	OpenAPINamer func(OpenAPINamerInfo) string
}

type OpenAPINamerInfo struct {
	TypeName   string
	FullGroup  string
	ShortGroup string
	Version    string
	Kind       string
}

func (*GoTypes) JennyName() string {
	return "GoTypes"
}

func (g *GoTypes) Generate(appManifest codegen.AppManifest) (codejen.Files, error) {
	files := make(codejen.Files, 0)
	for version, kind := range codegen.VersionedKinds(appManifest) {
		if g.GenerateOnlyCurrent && appManifest.Properties().PreferredVersion != version.Name() {
			continue
		}

		genCfg := goTypesGenerateFilesConfig{
			VersionName:   version.Name(),
			KindName:      kind.Kind,
			MachineName:   kind.MachineName,
			Group:         appManifest.Properties().Group,
			PackageName:   ToPackageName(version.Name()),
			PathPrefix:    GetGeneratedGoTypePath(g.GroupByKind, appManifest.Properties().Group, version.Name(), kind.MachineName),
			ExcludeFields: g.ExcludeFields,
		}
		if g.Depth > 0 {
			if !g.GroupByKind {
				genCfg.NamePrefix = exportField(kind.Kind)
			}
			generated, err := g.generateFilesAtDepth(kind.Schema, kind.Schema.Path(), 0, genCfg)
			if err != nil {
				return nil, err
			}
			files = append(files, generated...)
			continue
		}

		codegenPipeline := cog.TypesFromSchema().
			CUEValue(genCfg.PackageName, kind.Schema, cog.ForceEnvelope(kind.Kind)).
			Golang(cog.GoConfig{
				AnyAsInterface: g.AnyAsInterface,
			})

		if g.AddKubernetesCodegen {
			codegenPipeline = codegenPipeline.SchemaTransformations(cog.AppendCommentToObjects("+k8s:openapi-gen=true"))
		}

		generated, err := codegenPipeline.Run(context.Background())
		if err != nil {
			return nil, err
		}

		if len(generated) != 1 {
			return nil, fmt.Errorf("expected one file to be generated, got %d", len(generated))
		}

		formatted, err := format.Source(generated[0].Data)
		if err != nil {
			return nil, err
		}

		files = append(files, codejen.File{
			Data:         formatted,
			RelativePath: fmt.Sprintf(path.Join(genCfg.PathPrefix, "%s_gen.go"), strings.ToLower(genCfg.MachineName)),
			From:         []codejen.NamedJenny{g},
		})
	}
	return files, nil
}

type goTypesGenerateFilesConfig struct {
	VersionName   string
	KindName      string
	MachineName   string
	PackageName   string
	PathPrefix    string
	NamePrefix    string
	Group         string
	ExcludeFields []string
}

//nolint:goconst
func (g *GoTypes) generateFilesAtDepth(v cue.Value, schemaPath cue.Path, currDepth int, cfg goTypesGenerateFilesConfig) (codejen.Files, error) {
	if currDepth == g.Depth {
		selectors := TrimPathPrefix(v.Path(), schemaPath).Selectors()
		fieldName := make([]string, 0, len(selectors))
		for _, s := range selectors {
			fieldName = append(fieldName, s.String())
		}
		exclude := false
		for _, s := range cfg.ExcludeFields {
			// Check if the exclude name matches either the final element of the path, or the joined path
			if len(fieldName) > 0 && strings.EqualFold(fieldName[len(fieldName)-1], s) {
				exclude = true
				break
			}
			if strings.EqualFold(strings.Join(fieldName, ""), s) {
				exclude = true
				break
			}
		}
		if exclude {
			return codejen.Files{}, nil
		}

		var namerFunc func(string) string
		if g.OpenAPINamer != nil {
			namerFunc = func(name string) string {
				return g.OpenAPINamer(OpenAPINamerInfo{
					TypeName:   name,
					ShortGroup: cfg.Group,
					Version:    cfg.VersionName,
					Kind:       cfg.KindName,
				})
			}
		}

		goBytes, err := GoTypesFromCUE(v, CUEGoConfig{
			PackageName:                    cfg.PackageName,
			Name:                           exportField(strings.Join(fieldName, "")),
			NamePrefix:                     cfg.NamePrefix,
			AddKubernetesOpenAPIGenComment: g.AddKubernetesCodegen && (len(fieldName) != 1 || fieldName[0] != "metadata"),
			AnyAsInterface:                 g.AnyAsInterface,
		}, len(v.Path().Selectors())-(g.Depth-g.NamingDepth), namerFunc)
		if err != nil {
			return nil, fmt.Errorf("error converting schema path %s.%s to go type: %w", schemaPath.String(), fieldName, err)
		}

		return codejen.Files{codejen.File{
			Data:         goBytes,
			RelativePath: fmt.Sprintf(path.Join(cfg.PathPrefix, "%s_%s_gen.go"), strings.ToLower(cfg.MachineName), strings.Join(fieldName, "_")),
			From:         []codejen.NamedJenny{g},
		}}, nil
	}

	it, err := v.Fields()
	if err != nil {
		return nil, err
	}

	files := make(codejen.Files, 0)
	for it.Next() {
		f, err := g.generateFilesAtDepth(it.Value(), schemaPath, currDepth+1, cfg)
		if err != nil {
			return nil, err
		}
		files = append(files, f...)
	}
	return files, nil
}

type CUEGoConfig struct {
	PackageName string
	Name        string

	AddKubernetesOpenAPIGenComment bool
	AnyAsInterface                 bool

	// NamePrefix prefixes all generated types with the provided NamePrefix
	NamePrefix string
}

func GoTypesFromCUE(v cue.Value, cfg CUEGoConfig, maxNamingDepth int, namerFunc func(string) string) ([]byte, error) {
	nameFunc := func(_ cue.Value, definitionPath cue.Path) string {
		prefix := ""
		if typePrefix := getTypePrefix(v); typePrefix != "" {
			prefix = typePrefix
		}
		i := 0
		for ; i < len(definitionPath.Selectors()) && i < len(v.Path().Selectors()); i++ {
			if maxNamingDepth > 0 && i >= maxNamingDepth {
				break
			}
			if !SelEq(definitionPath.Selectors()[i], v.Path().Selectors()[i]) {
				break
			}
		}
		if i > 0 {
			definitionPath = cue.MakePath(definitionPath.Selectors()[i:]...)
		}
		name := strings.Trim(definitionPath.String(), "?#")
		if prefix != "" {
			return fmt.Sprintf("%s.%s", prefix, name)
		}
		return name
	}
	// Force CUE to evaluate the value tree without setting isData
	// This speeds up subsequent Subsume() and Equals() calls being done by cog, as otherwise this jenny's cost can become prohibitive.
	// TODO: this should probably be a fix in cog rather than here
	err := v.Validate(cue.All())
	if err != nil {
		return nil, fmt.Errorf("invalid CUE value at %s: %w", v.Path(), err)
	}
	codegenPipeline := cog.TypesFromSchema().
		CUEValue(cfg.PackageName, v, cog.ForceEnvelope(cfg.Name), cog.NameFunc(nameFunc)).
		SchemaTransformations(cog.PrefixObjectsNames(cfg.NamePrefix)).
		Golang(cog.GoConfig{
			AnyAsInterface:    cfg.AnyAsInterface,
			CustomTemplatesFS: templates.GetCogTemplates(),
			CustomTemplatesFuncs: map[string]any{
				"namerFunc": namerFunc,
			},
		})

	if cfg.AddKubernetesOpenAPIGenComment {
		codegenPipeline = codegenPipeline.SchemaTransformations(cog.AppendCommentToObjects("+k8s:openapi-gen=true"))
	}

	files, err := codegenPipeline.Run(context.Background())
	if err != nil {
		return nil, err
	}

	if len(files) != 1 {
		return nil, fmt.Errorf("expected one file to be generated, got %d", len(files))
	}

	formatted, err := format.Source(files[0].Data)
	if err != nil {
		return nil, err
	}

	return formatted, nil
}

// SanitizeLabelString strips characters from a string that are not allowed for
// use in a CUE label.
func sanitizeLabelString(s string) string {
	return strings.Map(func(r rune) rune {
		switch {
		case r >= 'a' && r <= 'z':
			fallthrough
		case r >= 'A' && r <= 'Z':
			fallthrough
		case r >= '0' && r <= '9':
			fallthrough
		case r == '_':
			return r
		default:
			return -1
		}
	}, s)
}

// exportField makes a field name exported
func exportField(field string) string {
	if len(field) > 0 {
		return strings.ToUpper(field[:1]) + field[1:]
	}
	return strings.ToUpper(field)
}