gobuf/gobuf.go

package gobuf

import (
	"bytes"
	"fmt"
	"io"
	"os"

	"github.com/emicklei/proto"
	"gitter.top/sync/proto-contrib/pkg/protofmt"
)

type Parser struct {
	filename string
	rawData  io.Reader
	proto    *proto.Proto
}

func NewParser(file string) (*Parser, error) {
	reader, err := os.Open(file)
	if err != nil {
		return nil, err
	}
	defer reader.Close()

	parser := proto.NewParser(reader)
	definition, err := parser.Parse()
	if err != nil {
		return nil, err
	}
	return &Parser{
		filename: file,
		rawData:  reader,
		proto:    definition,
	}, nil
}

func (parser *Parser) ExistService(serviceName string) bool {
	var result bool
	proto.Walk(parser.proto, proto.WithService(func(service *proto.Service) {
		if service.Name == serviceName {
			result = true
		}
	}))
	return result
}

func (parser *Parser) ExistMessage(messageName string) bool {
	var result bool
	proto.Walk(parser.proto, proto.WithMessage(func(message *proto.Message) {
		if message.Name == messageName {
			result = true
		}
	}))
	return result
}

func (parser *Parser) ExistRPC(serviceName, rpcName string) bool {
	var result bool
	proto.Walk(parser.proto, proto.WithService(func(service *proto.Service) {
		if service.Name != serviceName {
			return
		}
		for _, element := range service.Elements {
			if rpc, ok := element.(*proto.RPC); ok {
				if rpc.Name == rpcName {
					result = true
					return
				}
			}
		}
	}))
	return result
}

func (parser *Parser) AddRPC(serviceName, rpcName string) error {
	proto.Walk(parser.proto, proto.WithService(func(service *proto.Service) {
		if service.Name != serviceName {
			return
		}
		service.Elements = append(service.Elements, &proto.RPC{
			Comment: &proto.Comment{
				Lines: []string{
					" @desc: 描述",
					" @author: 作者",
					" @method: GET",
					" @api: /" + calm2Case(rpcName),
				},
			},
			Name:        rpcName,
			RequestType: toUpperCamlCase(serviceName) + toUpperCamlCase(rpcName) + "Req",
			ReturnsType: toUpperCamlCase(serviceName) + toUpperCamlCase(rpcName) + "Resp",
			Parent:      service,
		})
	}))

	parser.proto.Elements = append(parser.proto.Elements, &proto.Message{
		Name:   toUpperCamlCase(serviceName) + toUpperCamlCase(rpcName) + "Req",
		Parent: parser.proto,
	})
	parser.proto.Elements = append(parser.proto.Elements, &proto.Message{
		Name:   toUpperCamlCase(serviceName) + toUpperCamlCase(rpcName) + "Resp",
		Parent: parser.proto,
	})

	return parser.writeSync()
}

func (parser *Parser) AddService(serviceName string) error {
	if parser.ExistService(serviceName) {
		return fmt.Errorf("service name exist")
	}

	parser.proto.Elements = append(parser.proto.Elements, &proto.Service{
		Comment: &proto.Comment{
			Lines: []string{
				" @route_group: true",
				" @base_url: /v1/" + calm2Case(serviceName),
				" @gen_to: ./services/controller/v1/" + calm2Case(serviceName) + "_controller.go",
			},
		},
		Name:   serviceName,
		Parent: parser.proto,
	})

	return parser.writeSync()
}

func (parser *Parser) writeSync() error {
	var buf = new(bytes.Buffer)

	protofmt.NewFormatter(buf, "    ").Format(parser.proto) // 1 tab

	// write back to input
	if err := os.WriteFile(parser.filename, buf.Bytes(), os.ModePerm); err != nil {
		return err
	}

	return nil
}