Skip to content

Commit

Permalink
Add ability to generate enum types from string configs (#286)
Browse files Browse the repository at this point in the history 
* Implement string enum config codegen

* Add converting string configs to proto enums when grouping

* Add flag for disabling gen proto enum when grouping

* Clean some comments

* Fix for duplicate values

* Fix lint
  • Loading branch information
@DavidSGK
DavidSGK authored Jan 18, 2024
1 parent 34bd4c3 commit d152f48
Show file tree
Hide file tree
Showing 3 changed files with 301 additions and 21 deletions.
77 changes: 72 additions & 5 deletions cmd/lekko/feature.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -340,6 +340,7 @@ func featureEval() *cobra.Command {
func configGroup() *cobra.Command {
var ns, protoPkg, outName, description string
var configNames []string
var disableGenEnum bool
cmd := &cobra.Command{
Use: "group",
Short: "group multiple configs into 1 config",
Expand Down Expand Up @@ -436,6 +437,33 @@ func configGroup() *cobra.Command {
for _, c := range configNames {
compiledList = append(compiledList, compiledMap[c])
}
// For now, it's required to statically parse configs to get metadata info
// which we use to determine if a config is an enum config
// Keep map of (enum) config names to translated enum name and values
// We'll build up the values lookup map in a later stage
enumLookupMap := make(map[string]struct {
name string
values map[string]string
})
for _, cn := range configNames {
sf, err := r.Parse(ctx, ns, cn, registry)
if err != nil {
return errors.Wrap(err, "pre-group static parsing")
}
c := compiledMap[cn]
if genEnum, ok := sf.Feature.Metadata.AsMap()["gen-enum"]; ok && c.FeatureType == eval.ConfigTypeString && !disableGenEnum {
if genEnumBool, ok := genEnum.(bool); ok && genEnumBool {
enumLookupMap[cn] = struct {
name string
values map[string]string
}{
name: strcase.ToCamel(cn),
values: make(map[string]string),
}
}
}
}

// Prompt for grouped name if necessary, using suggestion engine
if outName == "" {
// tab for suggestions?
Expand All @@ -458,13 +486,34 @@ func configGroup() *cobra.Command {
protoMsgName := strcase.ToCamel(outName)
pdf := feature.NewProtoDefBuilder(protoMsgName)
protoFieldNames := make([]string, len(configNames))
for i, c := range configNames {
pt, err := pdf.ToProtoTypeName(compiledMap[c].FeatureType)
for i, cn := range configNames {
c := compiledMap[cn]
// Handle "enum" string configs
if enumLookup, ok := enumLookupMap[cn]; ok {
pedf := feature.NewProtoEnumDefBuilder(enumLookup.name)
if defaultVal, ok := c.Value.(string); ok {
enumLookup.values[defaultVal] = pedf.AddValue(defaultVal)
}
for _, o := range c.Overrides {
if val, ok := o.Value.(string); ok {
// Check to prevent duplicate values
if _, ok := enumLookup.values[val]; !ok {
enumLookup.values[val] = pedf.AddValue(val)
}
}
}
enumDefStr := pedf.Build()
pdf.AddEnum(enumDefStr)
protoFieldNames[i] = pdf.AddField(cn, enumLookup.name, c.Description)
continue
}

pt, err := pdf.ToProtoTypeName(c.FeatureType)
if err != nil {
return err
}
// Use original description as comment in generated proto
protoFieldNames[i] = pdf.AddField(c, pt, compiledMap[c].Description)
protoFieldNames[i] = pdf.AddField(cn, pt, c.Description)
}
// Update proto file
protoPath := path.Join(rootMD.ProtoDirectory, strings.ReplaceAll(protoPkg, ".", "/"), "gen.proto")
Expand Down Expand Up @@ -521,8 +570,25 @@ func configGroup() *cobra.Command {
return errors.Wrap(err, "find message")
}
defaultValue := mt.New()
for i, c := range configNames {
orig := compiledMap[c]
for i, cn := range configNames {
orig := compiledMap[cn]
// If we converted to an enum, need to set the enum number value accordingly
if enumLookup, ok := enumLookupMap[cn]; ok {
if origVal, ok := orig.Value.(string); ok {
enumDescriptor := mt.Descriptor().Enums().ByName(protoreflect.Name(enumLookup.name))
if enumDescriptor == nil {
return errors.Errorf("missing enum descriptor for %s", enumLookup.name)
}
valueDescriptor := enumDescriptor.Values().ByName(protoreflect.Name(enumLookup.values[origVal]))
if valueDescriptor == nil {
return errors.Errorf("missing enum value for %s", enumLookup.values[origVal])
}
defaultValue.Set(mt.Descriptor().Fields().ByName(protoreflect.Name(protoFieldNames[i])), protoreflect.ValueOf(valueDescriptor.Number()))
continue
} else {
return errors.New("unexpected non-string original value")
}
}
defaultValue.Set(mt.Descriptor().Fields().ByName(protoreflect.Name(protoFieldNames[i])), protoreflect.ValueOf(orig.Value))
}
newF.Value = defaultValue
Expand Down Expand Up @@ -585,6 +651,7 @@ func configGroup() *cobra.Command {
cmd.Flags().StringSliceVarP(&configNames, "configs", "c", []string{}, "comma-separated names of configs to group together")
cmd.Flags().StringVarP(&protoPkg, "proto-pkg", "p", "default.config.v1beta1", "package for generated protobuf type(s)")
cmd.Flags().StringVarP(&description, "description", "d", "", "description for the grouped config")
cmd.Flags().BoolVar(&disableGenEnum, "disable-gen-enum", false, "whether to disable conversion of protobuf enums from string enum configs")
return cmd
}

Expand Down
96 changes: 93 additions & 3 deletions cmd/lekko/gen.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,9 @@ package main

import (
"bytes"
"context"
"fmt"
"go/format"
"os"
"os/exec"
"regexp"
Expand Down Expand Up @@ -101,7 +103,7 @@ func genGoCmd() *cobra.Command {
if err := proto.Unmarshal(fff, f); err != nil {
return err
}
codeString := try.To1(genGoForFeature(f, ns, staticCtxType))
codeString := try.To1(genGoForFeature(cmd.Context(), r, f, ns, staticCtxType))
codeStrings = append(codeStrings, codeString)
if f.Type == featurev1beta1.FeatureType_FEATURE_TYPE_PROTO {
protoImport := unpackProtoType(moduleRoot, f.Tree.Default.TypeUrl)
Expand Down Expand Up @@ -240,7 +242,7 @@ var genCmd = &cobra.Command{
Short: "generate library code from configs",
}

func genGoForFeature(f *featurev1beta1.Feature, ns string, staticCtxType *protoImport) (string, error) {
func genGoForFeature(ctx context.Context, r repo.ConfigurationRepository, f *featurev1beta1.Feature, ns string, staticCtxType *protoImport) (string, error) {
const defaultTemplateBody = `{{if $.UnSafeClient}}
// {{$.Description}}
func (c *LekkoClient) {{$.FuncName}}(ctx context.Context) ({{$.RetType}}, error) {
Expand Down Expand Up @@ -275,6 +277,25 @@ func (c *SafeLekkoClient) {{$.FuncName}}(ctx context.Context, result interface{}
c.{{$.GetFunction}}(ctx, "{{$.Namespace}}", "{{$.Key}}", result)
}
`

// Generate an enum type and const declarations
const stringEnumTemplateBody = `type {{$.EnumTypeName}} string
const (
{{range $index, $field := $.EnumFields}}{{$field.Name}} {{$.EnumTypeName}} = "{{$field.Value}}"
{{end}}
)
{{if $.UnSafeClient}}
// {{$.Description}}
func (c *LekkoClient) {{$.FuncName}}(ctx context.Context) ({{$.RetType}}, error) {
return c.{{$.GetFunction}}(ctx, "{{$.Namespace}}", "{{$.Key}}")
}
{{end}}
// {{$.Description}}
func (c *SafeLekkoClient) {{$.FuncName}}(ctx *{{$.StaticType}}) {{$.RetType}} {
{{range $.NaturalLanguage}}{{ . }}
{{end}}}`

var funcNameBuilder strings.Builder
funcNameBuilder.WriteString("Get")
for _, word := range regexp.MustCompile("[_-]+").Split(f.Key, -1) {
Expand All @@ -283,6 +304,12 @@ func (c *SafeLekkoClient) {{$.FuncName}}(ctx context.Context, result interface{}
funcName := funcNameBuilder.String()
var retType string
var getFunction string
var enumTypeName string
type EnumField struct {
Name string
Value string
}
var enumFields []EnumField
templateBody := defaultTemplateBody

type StaticContextInfo struct {
Expand All @@ -305,6 +332,37 @@ func (c *SafeLekkoClient) {{$.FuncName}}(ctx context.Context, result interface{}
case featurev1beta1.FeatureType_FEATURE_TYPE_STRING:
retType = "string"
getFunction = "GetString"
// HACK: The metadata field is only for presentation at the moment
// so is not part of the compiled object - need to statically parse
// This also means that this only works for statically parseable
// configs
sf, err := r.Parse(ctx, ns, f.Key, typeRegistry)
if err != nil {
return "", errors.Wrap(err, "static parsing")
}
fm := sf.Feature.Metadata.AsMap()
// TODO: This enum codegen does not handle possible conflicts at all
if genEnum, ok := fm["gen-enum"]; ok {
if genEnumBool, ok := genEnum.(bool); ok && genEnumBool {
enumTypeName = strcase.ToCamel(f.Key)
retType = enumTypeName
templateBody = stringEnumTemplateBody
for _, ret := range getStringRetValues(f) {
// Result of translating ret values is wrapped in quotes
ret = ret[1 : len(ret)-1]
name := enumTypeName
if ret == "" {
name += "Unspecified"
} else {
name += strcase.ToCamel(ret)
}
enumFields = append(enumFields, EnumField{
Name: name,
Value: ret,
})
}
}
}
case featurev1beta1.FeatureType_FEATURE_TYPE_JSON:
getFunction = "GetJSON"
templateBody = jsonTemplateBody
Expand Down Expand Up @@ -334,6 +392,8 @@ func (c *SafeLekkoClient) {{$.FuncName}}(ctx context.Context, result interface{}
NaturalLanguage []string
StaticType string
UnSafeClient bool
EnumTypeName string
EnumFields []EnumField
}{
f.Description,
funcName,
Expand All @@ -344,6 +404,8 @@ func (c *SafeLekkoClient) {{$.FuncName}}(ctx context.Context, result interface{}
[]string{},
"",
UnSafeClient,
enumTypeName,
enumFields,
}
if staticContextInfo != nil {
data.NaturalLanguage = staticContextInfo.Natty
Expand All @@ -355,7 +417,15 @@ func (c *SafeLekkoClient) {{$.FuncName}}(ctx context.Context, result interface{}
}
var ret bytes.Buffer
err = templ.Execute(&ret, data)
return ret.String(), err
if err != nil {
return "", err
}
// Final canonical Go format
formatted, err := format.Source(ret.Bytes())
if err != nil {
return "", errors.Wrap(err, "format")
}
return string(formatted), nil
}

type protoImport struct {
Expand Down Expand Up @@ -497,3 +567,23 @@ func translateRetValue(val *anypb.Any, protoType *protoImport) string {
})
return fmt.Sprintf("&%s.%s{%s}", protoType.PackageAlias, protoType.Type, strings.Join(lines, ", "))
}

// TODO: Generify
// Get all unique possible return values of a config
func getStringRetValues(f *featurev1beta1.Feature) []string {
if f.Type != featurev1beta1.FeatureType_FEATURE_TYPE_STRING {
return []string{}
}
valSet := make(map[string]bool)
valSet[translateRetValue(f.Tree.Default, nil)] = true
for _, constraint := range f.Tree.Constraints {
ret := translateRetValue(constraint.Value, nil)
valSet[ret] = true
}
var rets []string
for val := range valSet {
rets = append(rets, val)
}
sort.Strings(rets)
return rets
}
Loading

0 comments on commit d152f48

Please sign in to comment.