1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157
|
package gengateway
import (
"errors"
"fmt"
"go/format"
"path"
"github.com/golang/glog"
"github.com/grpc-ecosystem/grpc-gateway/v2/internal/descriptor"
gen "github.com/grpc-ecosystem/grpc-gateway/v2/internal/generator"
"google.golang.org/protobuf/proto"
"google.golang.org/protobuf/types/pluginpb"
)
var (
errNoTargetService = errors.New("no target service defined in the file")
)
type generator struct {
reg *descriptor.Registry
baseImports []descriptor.GoPackage
useRequestContext bool
registerFuncSuffix string
allowPatchFeature bool
standalone bool
}
// New returns a new generator which generates grpc gateway files.
func New(reg *descriptor.Registry, useRequestContext bool, registerFuncSuffix string,
allowPatchFeature, standalone bool) gen.Generator {
var imports []descriptor.GoPackage
for _, pkgpath := range []string{
"context",
"io",
"net/http",
"github.com/grpc-ecosystem/grpc-gateway/v2/runtime",
"github.com/grpc-ecosystem/grpc-gateway/v2/utilities",
"google.golang.org/protobuf/proto",
"google.golang.org/grpc",
"google.golang.org/grpc/codes",
"google.golang.org/grpc/grpclog",
"google.golang.org/grpc/metadata",
"google.golang.org/grpc/status",
} {
pkg := descriptor.GoPackage{
Path: pkgpath,
Name: path.Base(pkgpath),
}
if err := reg.ReserveGoPackageAlias(pkg.Name, pkg.Path); err != nil {
for i := 0; ; i++ {
alias := fmt.Sprintf("%s_%d", pkg.Name, i)
if err := reg.ReserveGoPackageAlias(alias, pkg.Path); err != nil {
continue
}
pkg.Alias = alias
break
}
}
imports = append(imports, pkg)
}
return &generator{
reg: reg,
baseImports: imports,
useRequestContext: useRequestContext,
registerFuncSuffix: registerFuncSuffix,
allowPatchFeature: allowPatchFeature,
standalone: standalone,
}
}
func (g *generator) Generate(targets []*descriptor.File) ([]*descriptor.ResponseFile, error) {
var files []*descriptor.ResponseFile
for _, file := range targets {
glog.V(1).Infof("Processing %s", file.GetName())
code, err := g.generate(file)
if err == errNoTargetService {
glog.V(1).Infof("%s: %v", file.GetName(), err)
continue
}
if err != nil {
return nil, err
}
formatted, err := format.Source([]byte(code))
if err != nil {
glog.Errorf("%v: %s", err, code)
return nil, err
}
files = append(files, &descriptor.ResponseFile{
GoPkg: file.GoPkg,
CodeGeneratorResponse_File: &pluginpb.CodeGeneratorResponse_File{
Name: proto.String(file.GeneratedFilenamePrefix + ".pb.gw.go"),
Content: proto.String(string(formatted)),
},
})
}
return files, nil
}
func (g *generator) generate(file *descriptor.File) (string, error) {
pkgSeen := make(map[string]bool)
var imports []descriptor.GoPackage
for _, pkg := range g.baseImports {
pkgSeen[pkg.Path] = true
imports = append(imports, pkg)
}
if g.standalone {
imports = append(imports, file.GoPkg)
}
for _, svc := range file.Services {
for _, m := range svc.Methods {
imports = append(imports, g.addEnumPathParamImports(file, m, pkgSeen)...)
pkg := m.RequestType.File.GoPkg
if len(m.Bindings) == 0 ||
pkg == file.GoPkg || pkgSeen[pkg.Path] {
continue
}
pkgSeen[pkg.Path] = true
imports = append(imports, pkg)
}
}
params := param{
File: file,
Imports: imports,
UseRequestContext: g.useRequestContext,
RegisterFuncSuffix: g.registerFuncSuffix,
AllowPatchFeature: g.allowPatchFeature,
}
if g.reg != nil {
params.OmitPackageDoc = g.reg.GetOmitPackageDoc()
}
return applyTemplate(params, g.reg)
}
// addEnumPathParamImports handles adding import of enum path parameter go packages
func (g *generator) addEnumPathParamImports(file *descriptor.File, m *descriptor.Method, pkgSeen map[string]bool) []descriptor.GoPackage {
var imports []descriptor.GoPackage
for _, b := range m.Bindings {
for _, p := range b.PathParams {
e, err := g.reg.LookupEnum("", p.Target.GetTypeName())
if err != nil {
continue
}
pkg := e.File.GoPkg
if pkg == file.GoPkg || pkgSeen[pkg.Path] {
continue
}
pkgSeen[pkg.Path] = true
imports = append(imports, pkg)
}
}
return imports
}
|