transport_server.go

package template

import (
	"context"

	. "github.com/dave/jennifer/jen"
	mstrings "github.com/recolabs/microgen/generator/strings"
	"github.com/recolabs/microgen/generator/write_strategy"
	"github.com/vetcher/go-astra/types"
)

type endpointsServerTemplate struct {
	info *GenerationInfo
}

func NewEndpointsServerTemplate(info *GenerationInfo) Template {
	return &endpointsServerTemplate{
		info: info,
	}
}

// Renders endpoints file.
//
//		// This file was automatically generated by "microgen" utility.
//		// DO NOT EDIT.
//		package stringsvc
//
//		import (
//			context "context"
//			endpoint "github.com/go-kit/kit/endpoint"
//		)
//
//		func CountEndpoint(svc StringService) endpoint.Endpoint {
//			return func(ctx context.Context, request interface{}) (interface{}, error) {
//				req := request.(*CountRequest)
//				count, positions := svc.Count(ctx, req.Text, req.Symbol)
//				return &CountResponse{
//					Count:     count,
//					Positions: positions,
//				}, nil
//			}
//		}
//
func (t *endpointsServerTemplate) Render(ctx context.Context) write_strategy.Renderer {
	f := NewFile("transport")
	f.HeaderComment(t.info.FileHeader)
	f.ImportAlias(t.info.ProtobufPackageImport, "pb")

	f.Add(t.allEndpoints()).Line()
	if Tags(ctx).HasAny(TracingMiddlewareTag) {
		f.Comment("TraceServerEndpoints is used for tracing endpoints on server side.")
		f.Add(t.serverTracingMiddleware()).Line()
	}
	for _, signature := range t.info.Iface.Methods {
		if t.info.OneToManyStreamMethods[signature.Name] {
			f.Add(createOneToManyStreamEndpoint(signature, t.info)).Line().Line()
			continue
		}
		if t.info.ManyToManyStreamMethods[signature.Name] {
			f.Add(createManyToManyStreamEndpoint(signature, t.info)).Line().Line()
			continue
		}
		if t.info.ManyToOneStreamMethods[signature.Name] {
			f.Add(createManyToOneStreamEndpoint(signature, t.info)).Line().Line()
			continue
		}
		if t.info.AllowedMethods[signature.Name] {
			f.Add(createEndpoint(signature, t.info)).Line().Line()
		}
	}
	return f
}

func (endpointsServerTemplate) DefaultPath() string {
	return filenameBuilder(PathTransport, "server")
}

func (t *endpointsServerTemplate) Prepare(ctx context.Context) error {
	return nil
}

func (t *endpointsServerTemplate) ChooseStrategy(ctx context.Context) (write_strategy.Strategy, error) {
	return write_strategy.NewCreateFileStrategy(t.info.OutputFilePath, t.DefaultPath()), nil
}

// Render new Endpoint body.
//
//		return func(ctx context.Context, request interface{}) (interface{}, error) {
//			req := request.(*CountRequest)
//			count, positions := svc.Count(ctx, req.Text, req.Symbol)
//			return &CountResponse{
//				Count:     count,
//				Positions: positions,
//			}, nil
//		}
//
func createEndpointBody(signature *normalizedFunction) *Statement {
	return Return(Func().Params(
		Id(firstArgName(&signature.Function)).Qual("context", "Context"),
		Id("request").Interface(),
	).Params(
		Interface(),
		Error(),
	).BlockFunc(func(g *Group) {
		methodParams := RemoveContextIfFirst(signature.parent.Args)
		if len(methodParams) > 0 {
			g.Id("req").Op(":=").Id("request").Assert(Op("*").Id(requestStructName(signature.parent)))
		}

		g.Add(paramNames(signature.Results).
			Op(":=").
			Id("svc").
			Dot(signature.Name).
			CallFunc(func(g *Group) {
				g.Add(Id(firstArgName(&signature.Function)))
				for _, field := range methodParams {
					v := Dot(mstrings.ToUpperFirst(field.Name))
					if types.IsEllipsis(field.Type) {
						v.Op("...")
					}
					g.Add(Id("req").Add(v))
				}
			}))

		g.Return(
			Op("&").Id(responseStructName(signature.parent)).Values(dictByNormalVariables(
				removeErrorIfLast(signature.parent.Results),
				removeErrorIfLast(signature.Results),
			)),
			Id(nameOfLastResultError(&signature.Function)),
		)
	}))
}

// Render full new Endpoint function.
//
//		func CountEndpoint(svc StringService) endpoint.Endpoint {
//			return func(ctx context.Context, request interface{}) (interface{}, error) {
//				req := request.(*CountRequest)
//				count, positions := svc.Count(ctx, req.Text, req.Symbol)
//				return &CountResponse{
//					Count:     count,
//					Positions: positions,
//				}, nil
//			}
//		}
//
func removeLastVar(slice []types.Variable) []types.Variable {
	if len(slice) > 0 {
		return slice[:len(slice)-1]
	}
	return slice
}

func createOneToManyStreamEndpoint(signature *types.Function, info *GenerationInfo) *Statement {
	normal := normalizeFunction(signature)
	return Func().
		Id(endpointsStructFieldName(signature.Name)).Params(Id("svc").Qual(info.SourcePackageImport, info.Iface.Name)).
		Params(Id(OneToManyStreamEndpoint)).
		Block(createOneToManyStreamEndpointBody(info, normal))
}

func streamStructName(ifaceName string, signature *types.Function) string {
	return ifaceName + "_" + signature.Name + "Server"
}

func createOneToManyStreamEndpointBody(info *GenerationInfo, signature *normalizedFunction) *Statement {
	return Return(Func().Params(
		Id("request").Interface(),
		Id("stream").Interface(),
	).Params(
		Error(),
	).BlockFunc(func(g *Group) {
		g.Id("req").Op(":=").Id("request").Assert(Op("*").Id(requestStructName(signature.parent)))
		g.Id("st").Op(":=").Id("stream").Assert(Qual(info.ProtobufPackageImport, streamStructName(info.Iface.Name, signature.parent)))

		g.Add(paramNames(signature.Results).
			Op(":=").
			Id("svc").
			Dot(signature.Name).
			CallFunc(func(g *Group) {
				for _, field := range removeLastVar(signature.parent.Args) {
					v := Dot(mstrings.ToUpperFirst(field.Name))
					if types.IsEllipsis(field.Type) {
						v.Op("...")
					}
					g.Add(Id("req").Add(v))
				}
				g.Add(Id("st"))
			}))

		g.Return(
			Id(nameOfLastResultError(&signature.Function)),
		)
	}))
}

func createManyToManyStreamEndpoint(signature *types.Function, info *GenerationInfo) *Statement {
	normal := normalizeFunction(signature)
	return Func().
		Id(endpointsStructFieldName(signature.Name)).Params(Id("svc").Qual(info.SourcePackageImport, info.Iface.Name)).
		Params(Id(ManyToManyStreamEndpoint)).
		Block(createManyToManyStreamEndpointBody(info, normal))
}

func createManyToManyStreamEndpointBody(info *GenerationInfo, signature *normalizedFunction) *Statement {
	return Return(Func().Params(
		Id("stream").Interface(),
	).Params(
		Error(),
	).BlockFunc(func(g *Group) {
		g.Id("st").Op(":=").Id("stream").Assert(Qual(info.ProtobufPackageImport, streamStructName(info.Iface.Name, signature.parent)))

		g.Add(paramNames(signature.Results).
			Op(":=").
			Id("svc").
			Dot(signature.Name).
			CallFunc(func(g *Group) {
				g.Add(Id("st"))
			}))

		g.Return(
			Id(nameOfLastResultError(&signature.Function)),
		)
	}))
}

func createManyToOneStreamEndpoint(signature *types.Function, info *GenerationInfo) *Statement {
	normal := normalizeFunction(signature)
	return Func().
		Id(endpointsStructFieldName(signature.Name)).Params(Id("svc").Qual(info.SourcePackageImport, info.Iface.Name)).
		Params(Id(ManyToOneStreamEndpoint)).
		Block(createManyToOneStreamEndpointBody(info, normal))
}

func createManyToOneStreamEndpointBody(info *GenerationInfo, signature *normalizedFunction) *Statement {
	return Return(Func().Params(
		Id("stream").Interface(),
	).Params(
		Error(),
	).BlockFunc(func(g *Group) {
		g.Id("st").Op(":=").Id("stream").Assert(Qual(info.ProtobufPackageImport, streamStructName(info.Iface.Name, signature.parent)))

		g.Add(paramNames(signature.Results).
			Op(":=").
			Id("svc").
			Dot(signature.Name).
			CallFunc(func(g *Group) {
				g.Add(Id("st"))
			}))

		g.Return(
			Id(nameOfLastResultError(&signature.Function)),
		)
	}))
}

func createEndpoint(signature *types.Function, info *GenerationInfo) *Statement {
	normal := normalizeFunction(signature)
	return Func().
		Id(endpointsStructFieldName(signature.Name)).Params(Id("svc").Qual(info.SourcePackageImport, info.Iface.Name)).Params(Qual(PackagePathGoKitEndpoint, "Endpoint")).
		Block(createEndpointBody(normal))
}

func (t *endpointsServerTemplate) allEndpoints() *Statement {
	s := &Statement{}
	s.Func().Id("Endpoints").Call(Id("svc").Qual(t.info.SourcePackageImport, t.info.Iface.Name)).Id(EndpointsSetName).BlockFunc(func(g *Group) {
		g.Return(Id(EndpointsSetName).Values(DictFunc(func(d Dict) {
			for _, signature := range t.info.Iface.Methods {
				if t.info.AllowedMethods[signature.Name] {
					d[Id(endpointsStructFieldName(signature.Name))] = Id(endpointsStructFieldName(signature.Name)).Params(Id("svc"))
				}
			}
		})))
	})
	return s
}

func (t *endpointsServerTemplate) serverTracingMiddleware() *Statement {
	s := &Statement{}
	s.Func().Id("TraceServerEndpoints").Call(Id("endpoints").Id(EndpointsSetName), Id("tracer").Qual(PackagePathOpenTracingGo, "Tracer")).Id(EndpointsSetName).BlockFunc(func(g *Group) {
		g.Return(Id(EndpointsSetName).Values(DictFunc(func(d Dict) {
			for _, signature := range t.info.Iface.Methods {
				if t.info.AllowedMethods[signature.Name] {
					d[Id(endpointsStructFieldName(signature.Name))] = Qual(PackagePathGoKitTracing, "TraceServer").Call(Id("tracer"), Lit(signature.Name)).Call(Id("endpoints").Dot(endpointsStructFieldName(signature.Name)))
				}
			}
		})))
	})
	return s
}