main.go

package main

import (
	"errors"
	"fmt"
	"strings"

	"google.golang.org/protobuf/compiler/protogen"
	"google.golang.org/protobuf/proto"
	"google.golang.org/protobuf/reflect/protoreflect"
	"google.golang.org/protobuf/types/descriptorpb"
	"google.golang.org/protobuf/types/pluginpb"

	entitypb "github.com/utilitywarehouse/protoc-gen-uwentity/gen/uw/entity/v1"
)

var (
	errIdentifierNotFound = errors.New("identifier not found")
	errUnsupportedField   = errors.New("unsupported field type")
)

var tmpl = `
// GetEntityIdentifier returns the value from the field marked as the identifier
func (m *%s) GetEntityIdentifier() string {
	return %s
}
`

// If set, then only these directories will raise errors for messages that do not set an identifier
// By default all messages are checked with the `Event` suffix
var paramEnforceDirs []string

type identifier struct {
	// the Message that the GetEntityIdentifier method will be added to
	message string
	// the Field on the message which should be used as the identifier
	identifier string
	converter  string
}

func main() {
	protogen.Options{
		ParamFunc: func(name, value string) error {
			switch name {
			case "enforce-dir":
				paramEnforceDirs = append(paramEnforceDirs, value)
			default:
				return fmt.Errorf("invalid param: %s", name)
			}
			return nil
		},
	}.Run(func(gen *protogen.Plugin) error {
		for _, file := range gen.Files {
			if !file.Generate {
				continue
			}
			if len(file.Messages) == 0 {
				continue
			}

			gen.SupportedFeatures |= uint64(pluginpb.CodeGeneratorResponse_FEATURE_PROTO3_OPTIONAL)

			output := gen.NewGeneratedFile(file.GeneratedFilenamePrefix+".pb.uwentity.go", file.GoImportPath)
			output.P("// Code generated by protoc-gen-uwentity. DO NOT EDIT.")
			output.P("// source: ", file.Desc.Path())
			output.P("package ", file.GoPackageName)

			var idents []identifier

			for _, msg := range file.Messages {
				// Ingore messages with `uw.entity.v1.ignore = true`
				msgopts := msg.Desc.Options().(*descriptorpb.MessageOptions)
				if proto.GetExtension(msgopts, entitypb.E_Ignore).(bool) {
					continue
				}

				messageIdentifier, err := getMessageIdentifier(msg.GoIdent.GoName, msg)
				switch {
				default:
					return err
				case errors.Is(err, nil):
					idents = append(idents, messageIdentifier)
				case errors.Is(err, errIdentifierNotFound):
					// Check that each event has set the identifier
					// Pass `enforce-dir=<path>` to only check these directories & skip others
					// Set message option `(uw.entity.v1.ignore) = true` to skip an individual message
					if strings.HasSuffix(msg.GoIdent.GoName, "Event") {
						if len(paramEnforceDirs) > 0 {
							var skipEnforce bool
							for _, dir := range paramEnforceDirs {
								if !strings.HasPrefix(file.Desc.Path(), dir) {
									skipEnforce = true
									break
								}
							}
							if skipEnforce {
								continue
							}
						}

						return fmt.Errorf(
							"%s/%s: `uw.entity.v1.identifier` not set on event",
							file.Desc.Path(),
							msg.GoIdent.GoName,
						)
					}
				}
			}

			if len(idents) == 0 {
				// file contains no identifiers
				output.Skip()
				continue
			}

			var needsStrconv bool
			for _, ident := range idents {
				if ident.converter != "" {
					needsStrconv = true
					break
				}
			}

			if needsStrconv {
				output.P()
				output.P("import \"strconv\"")
			}

			// output identifier accessing methods, checking that we support the field type
			for _, ident := range idents {
				if ident.converter != "" {
					output.P(fmt.Sprintf(tmpl, ident.message, fmt.Sprintf(ident.converter, "m."+ident.identifier)))
				} else {
					output.P(fmt.Sprintf(tmpl, ident.message, "m."+ident.identifier))
				}
			}
		}

		return nil
	})
}

func getMessageIdentifier(entity string, msg *protogen.Message, fields ...string) (identifier, error) {
	for _, field := range msg.Fields {
		// Skip field that don't have `uw.entity.v1.identifer = true`
		fieldopts := field.Desc.Options().(*descriptorpb.FieldOptions)

		if !proto.GetExtension(fieldopts, entitypb.E_Identifier).(bool) {
			continue
		}

		if field.Desc.Cardinality() == protoreflect.Repeated {
			return identifier{}, fmt.Errorf("repeated fields are not supported on %s: %s: %w", entity, field.Desc.Kind(), errUnsupportedField)
		}

		var value, converter string
		switch field.Desc.Kind() {
		case protoreflect.Int64Kind:
			value = fmt.Sprintf("Get%s()", field.GoName)
			converter = `strconv.FormatInt(%s, 10)`
		case protoreflect.Int32Kind:
			value = fmt.Sprintf("Get%s()", field.GoName)
			converter = `strconv.FormatInt(int64(%s), 10)`
		case protoreflect.StringKind:
			value = fmt.Sprintf("Get%s()", field.GoName)
		case protoreflect.EnumKind:
			value = fmt.Sprintf("Get%s().String()", field.GoName)
		case protoreflect.MessageKind:
			nestedIdentifier := false
			// check if the nested message has an identifier
			for _, f := range field.Message.Fields {
				if proto.GetExtension(f.Desc.Options().(*descriptorpb.FieldOptions), entitypb.E_Identifier).(bool) {
					nestedIdentifier = true
				}
			}

			if !nestedIdentifier {
				return identifier{}, fmt.Errorf("nested messages must have an identifier on %s: %s: %w", entity, field.Desc.Kind(), errUnsupportedField)
			}
			value = fmt.Sprintf("Get%s().GetEntityIdentifier()", field.GoName)
			return identifier{
				message:    entity,
				identifier: strings.Join(append(fields, value), "."),
				converter:  converter,
			}, nil
		default:
			return identifier{}, fmt.Errorf("unsupported field type on %s: %s: %w", entity, field.Desc.Kind(), errUnsupportedField)
		}

		return identifier{
			message:    entity,
			identifier: strings.Join(append(fields, value), "."),
			converter:  converter,
		}, nil
	}

	return identifier{}, errIdentifierNotFound
}