generator.go

package main

import (
	"fmt"
	"github.com/martinxsliu/protoc-gen-graphql/parameters"
	"strings"

	"google.golang.org/protobuf/compiler/protogen"
	"google.golang.org/protobuf/types/pluginpb"

	"github.com/martinxsliu/protoc-gen-graphql/graphql"
	"github.com/martinxsliu/protoc-gen-graphql/mapper"
)

var header = []byte(`# DO NOT EDIT! Generated by protoc-gen-graphql.`)

type Generator struct {
	req    *pluginpb.CodeGeneratorRequest
	gen    *protogen.Plugin
	mapper *mapper.Mapper
}

func New(gen *protogen.Plugin) *Generator {
	return &Generator{
		req: gen.Request,
		gen: gen,
	}
}

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

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

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

func (g *Generator) generateFiles(params *parameters.Parameters) {
	for _, fileName := range g.req.GetFileToGenerate() {
		fileResp := &pluginpb.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 {
				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 {
				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)
		}

		genFile := g.gen.NewGeneratedFile(graphqlFileName(fileName), "github.com/not-a-real-import")

		_, _ = genFile.Write(header)
		for _, gqlType := range gqlTypes {
			_, _ = genFile.Write([]byte("\n\n"))
			_, _ = genFile.Write([]byte(graphql.TypeDef(gqlType, params)))
		}
		_, _ = genFile.Write([]byte("\n"))
	}
}

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

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