generator.go

package main

import (
	"fmt"
	"strings"

	plugin "github.com/golang/protobuf/protoc-gen-go/plugin"

	"github.com/apalchys/protoc-gen-graphql/descriptor"
	"github.com/apalchys/protoc-gen-graphql/graphql"
	"github.com/apalchys/protoc-gen-graphql/mapper"
)

const header = `# DO NOT EDIT! Generated by protoc-gen-graphql.`

type Generator struct {
	req    *plugin.CodeGeneratorRequest
	resp   *plugin.CodeGeneratorResponse
	mapper *mapper.Mapper
}

func New(req *plugin.CodeGeneratorRequest, resp *plugin.CodeGeneratorResponse) *Generator {
	return &Generator{
		req:  req,
		resp: resp,
	}
}

func (g *Generator) Generate() (err error) {
	defer func() {
		if r := recover(); r != nil {
			err = fmt.Errorf("%v", r)
		}
	}()

	params, err := mapper.NewParameters(g.req.GetParameter())
	if err != nil {
		return err
	}

	g.mapper = mapper.New(g.req.GetProtoFile(), params)
	g.generateFiles()
	return nil
}

func (g *Generator) generateFiles() {
	for _, fileName := range g.req.GetFileToGenerate() {
		fileResp := &plugin.CodeGeneratorResponse_File{}
		fileResp.Name = stringPtr(graphqlFileName(fileName))

		var gqlTypes []graphql.Type
		file := g.mapper.Files[fileName]

		for _, service := range file.Services {
			m, ok := g.mapper.ServiceMappers[service.FullName]
			if !ok {
				continue // Service was skipped
			}

			if m.Queries != nil {
				if m.Queries.ExtendRootObject != nil {
					gqlTypes = append(gqlTypes, m.Queries.ExtendRootObject)
				}
				gqlTypes = append(gqlTypes, m.Queries.Object)
			}
			if m.Mutations != nil {
				if m.Mutations.ExtendRootObject != nil {
					gqlTypes = append(gqlTypes, m.Mutations.ExtendRootObject)
				}
				gqlTypes = append(gqlTypes, m.Mutations.Object)
			}
			if m.Subscriptions != nil {
				if m.Subscriptions.ExtendRootObject != nil {
					gqlTypes = append(gqlTypes, m.Subscriptions.ExtendRootObject)
				}
				gqlTypes = append(gqlTypes, m.Subscriptions.Object)
			}
		}

		for _, message := range file.Messages {
			m := g.mapper.MessageMappers[message.FullName]

			if m.Object != nil && len(m.Oneofs) == 0 {
				gqlTypes = append(gqlTypes, m.Object)
			}
			for _, oneof := range m.Oneofs {
				gqlTypes = append(gqlTypes, oneof.Union)
				for _, object := range oneof.Objects {
					gqlTypes = append(gqlTypes, object)
				}
			}

			if m.Input != nil && len(m.Oneofs) == 0 {
				gqlTypes = append(gqlTypes, m.Input)
			}
			for _, oneof := range m.Oneofs {
				if oneof.Input != nil {
					gqlTypes = append(gqlTypes, oneof.Input)
				}
			}
		}

		for _, enum := range file.Enums {
			gqlTypes = append(gqlTypes, g.mapper.EnumMappers[enum.FullName].Enum)
		}

		var b strings.Builder
		b.WriteString(header)
		for _, gqlType := range gqlTypes {
			b.WriteString("\n\n")
			b.WriteString(graphql.TypeDef(gqlType))
		}
		b.WriteString("\n")
		fileResp.Content = stringPtr(b.String())

		g.resp.File = append(g.resp.File, fileResp)
	}
}

func (g *Generator) graphqlFromMessage(message *descriptor.Message) []graphql.Type {
	var graphqlTypes []graphql.Type
	return graphqlTypes
}

func graphqlFileName(name string) string {
	return strings.TrimSuffix(name, ".proto") + "_pb.graphql"
}

func stringPtr(v string) *string {
	return &v
}