Merge pull request #186 from mercari/support-overload-plugin

Support overload for cel plugin functions

_examples/15_cel_plugin/.gitignore

@@ -1 +1,2 @@
 grpc/federation
+*.wasm

_examples/15_cel_plugin/Makefile

@@ -13,7 +13,7 @@ lint:
 	@$(GOBIN)/grpc-federation-linter -Iproto -Iproto_deps ./proto/federation/federation.proto
 
 .PHONY: test
-test:
+test: build/wasm
 	go test -race ./ -count=1
 
 .PHONY: grpc-federation/generate

_examples/16_code_gen_plugin/.gitignore

@@ -1 +1,4 @@
 grpc/federation
+*.wasm
+grpc-federation.yaml
+buf.gen.yaml

_examples/16_code_gen_plugin/Makefile

@@ -5,9 +5,12 @@ PATH := $(GOBIN):$(PATH)
 JAEGER_IMAGE := jaegertracing/all-in-one:latest
 
 .PHONY: generate
-generate:
+generate: generate/config
 	$(GOBIN)/buf generate
 
+generate/config: build/wasm
+	go run ./cmd/config/main.go
+
 .PHONY: lint
 lint:
 	@$(GOBIN)/grpc-federation-linter -Iproto -Iproto_deps ./proto/federation/federation.proto
@@ -17,7 +20,7 @@ test:
 	go test -race ./ -count=1
 
 .PHONY: grpc-federation/generate
-grpc-federation/generate:
+grpc-federation/generate: generate/config
 	@$(GOBIN)/grpc-federation-generator ./proto/federation/federation.proto
 
 .PHONY: grpc-federation/watch

_examples/16_code_gen_plugin/cmd/config/main.go

@@ -0,0 +1,91 @@
+package main
+
+import (
+	"crypto/sha256"
+	"encoding/hex"
+	"fmt"
+	"log"
+	"os"
+	"path/filepath"
+	"runtime"
+)
+
+var bufGenYAML = `
+# Code generated by cmd/config/main.go. DO NOT EDIT!
+version: v1
+managed:
+  enabled: true
+plugins:
+  - plugin: go
+    out: .
+    opt: paths=source_relative
+  - plugin: go-grpc
+    out: .
+    opt: paths=source_relative
+  - plugin: grpc-federation
+    out: .
+    opt:
+      - paths=source_relative
+      - plugins=file://plugin.wasm:%s
+`
+
+var grpcFederationYAML = `
+# Code generated by cmd/config/main.go. DO NOT EDIT!
+imports:
+  - proto
+src:
+  - proto
+out: .
+plugins:
+  - plugin: go
+    opt: paths=source_relative
+  - plugin: go-grpc
+    opt: paths=source_relative
+  - plugin: grpc-federation
+    opt:
+      - paths=source_relative
+      - plugins=file://plugin.wasm:%s
+`
+
+func main() {
+	if err := run(); err != nil {
+		log.Fatalf("%+v", err)
+	}
+}
+
+func run() error {
+	wasm, err := os.ReadFile(filepath.Join(root(), "plugin.wasm"))
+	if err != nil {
+		return err
+	}
+	hash := toSha256(wasm)
+	if err := os.WriteFile(
+		filepath.Join(root(), "buf.gen.yaml"),
+		[]byte(fmt.Sprintf(bufGenYAML, hash)),
+		0o600,
+	); err != nil {
+		return err
+	}
+	if err := os.WriteFile(
+		filepath.Join(root(), "grpc-federation.yaml"),
+		[]byte(fmt.Sprintf(grpcFederationYAML, hash)),
+		0o600,
+	); err != nil {
+		return err
+	}
+	return nil
+}
+
+func toSha256(v []byte) string {
+	hash := sha256.Sum256(v)
+	return hex.EncodeToString(hash[:])
+}
+
+func root() string {
+	return filepath.Join(curDir(), "..", "..")
+}
+
+func curDir() string {
+	_, file, _, _ := runtime.Caller(0) //nolint:dogsled
+	return filepath.Dir(file)
+}

_examples/18_load/.gitignore

@@ -1 +1,2 @@
 grpc/federation
+*.wasm

_examples/18_load/Makefile

@@ -13,7 +13,7 @@ lint:
 	@$(GOBIN)/grpc-federation-linter -Iproto -Iproto_deps ./proto/federation/federation.proto
 
 .PHONY: test
-test:
+test: build/wasm
 	go test -race ./ -count=1
 
 .PHONY: grpc-federation/generate

_examples/18_load/cmd/plugin/main.go

@@ -22,6 +22,18 @@ func (_ *plugin) Example_Account_GetId(ctx context.Context) (string, error) {
 	return values[0], nil
 }
 
+func (_ *plugin) Example_Account_GetId2(ctx context.Context, ext string) (string, error) {
+	md, ok := metadata.FromIncomingContext(ctx)
+	if !ok {
+		return "", fmt.Errorf("failed to get metadata")
+	}
+	values := md["id"]
+	if len(values) == 0 {
+		return "", fmt.Errorf("failed to find id key in metadata")
+	}
+	return values[0], nil
+}
+
 func main() {
 	pluginpb.RegisterAccountPlugin(&plugin{})
 }

_examples/18_load/proto/plugin/plugin.proto

@@ -13,6 +13,12 @@ option (grpc.federation.plugin).export = {
       name: "get_id"
       desc: "get id text from incoming metadata"
       return { type: "string" }
+    },
+    {
+      name: "get_id"
+      desc: "overload method for get_id"
+      args { name: "ext" type: "string" }
+      return { type: "string" }
     }
   ]
 };

generator/code_generator.go

@@ -97,13 +97,26 @@ func (p *CELPlugin) Functions() []*CELFunction {
 type CELFunction struct {
 	file *File
 	*resolver.CELFunction
+	overloadIdx int
+}
+
+func (f *CELFunction) GoName() string {
+	funcName := f.CELFunction.Name
+	if f.overloadIdx != 0 {
+		funcName += fmt.Sprint(f.overloadIdx + 1)
+	}
+	if f.Receiver != nil {
+		return protoFQDNToPublicGoName(fmt.Sprintf("%s.%s", f.Receiver.FQDN(), funcName))
+	}
+	return protoFQDNToPublicGoName(funcName)
 }
 
 func (f *CELFunction) Name() string {
+	funcName := f.CELFunction.Name
 	if f.Receiver != nil {
-		return protoFQDNToPublicGoName(fmt.Sprintf("%s.%s", f.Receiver.FQDN(), f.CELFunction.Name))
+		return protoFQDNToPublicGoName(fmt.Sprintf("%s.%s", f.Receiver.FQDN(), funcName))
 	}
-	return protoFQDNToPublicGoName(f.CELFunction.Name)
+	return protoFQDNToPublicGoName(funcName)
 }
 
 func (f *CELFunction) IsMethod() bool {
@@ -230,11 +243,14 @@ func (r *CELFunctionReturn) Converter() string {
 
 func (p *CELPlugin) PluginFunctions() []*CELFunction {
 	ret := make([]*CELFunction, 0, len(p.CELPlugin.Functions))
+	overloadCount := make(map[string]int)
 	for _, fn := range p.CELPlugin.Functions {
 		ret = append(ret, &CELFunction{
 			file:        p.file,
 			CELFunction: fn,
+			overloadIdx: overloadCount[fn.Name],
 		})
+		overloadCount[fn.Name]++
 	}
 	return ret
 }

generator/templates/plugin.go.tmpl

@@ -4,7 +4,7 @@
 {{ $pluginFunctions := .PluginFunctions }}
 type {{ $pluginName }} interface {
 	{{- range $pluginFunctions }}
-	{{ .Name }}(context.Context, {{- range .Args }}{{ .Type }},{{- end }}) ({{- .Return.Type }}, error)
+	{{ .GoName }}(context.Context, {{- range .Args }}{{ .Type }},{{- end }}) ({{- .Return.Type }}, error)
 	{{- end }}
 }
 
@@ -70,7 +70,7 @@ func handle{{ $pluginName }}(content []byte, plug {{ $pluginName }}) (*grpcfed.C
 			return nil, err
 		}
 		{{- end }}
-		ret, err := plug.{{ .Name }}(ctx, {{- range $idx, $arg := .Args }}arg{{ $idx }},{{- end }})
+		ret, err := plug.{{ .GoName }}(ctx, {{- range $idx, $arg := .Args }}arg{{ $idx }},{{- end }})
 		if err != nil {
 			return nil, err
 		}