perf: 支持单独给rpc添加//@bot标签

TBXark · Nov 6, 2024 · 7ca34b1 · 7ca34b1
1 parent f1f6436
commit 7ca34b1
Show file tree

Hide file tree

Showing 4 changed files with 133 additions and 29 deletions.
diff --git a/api/bot/v1/counter.pb.go b/api/bot/v1/counter.pb.go
diff --git a/contrib/protoc-gen-bot/bot.go b/contrib/protoc-gen-bot/bot.go
@@ -64,7 +64,12 @@ func genService(_ *protogen.Plugin, file *protogen.File, g *protogen.GeneratedFi
 		UpdateType:  g.QualifiedGoIdent(conf.update.pkg.Ident(conf.update.model)),
 		MessageType: g.QualifiedGoIdent(conf.message.pkg.Ident(conf.message.model)),
 	}
+	allEnable := isBotEnable(service.Comments.Leading.String())
 	for _, method := range service.Methods {
+		enable := allEnable || isBotEnable(method.Comments.Leading.String())
+		if !enable {
+			continue
+		}
 		sd.Methods = append(sd.Methods, &methodDesc{
 			Name:         method.GoName,
 			OriginalName: string(method.Desc.Name()),
@@ -84,11 +89,23 @@ var botTagsMatchRegexp = regexp.MustCompile(` *//+ *@bot`)
 
 func hasBotRule(services []*protogen.Service) bool {
 	for _, service := range services {
-		for _, line := range strings.Split(service.Comments.Leading.String(), "\n") {
-			if botTagsMatchRegexp.MatchString(line) {
+		if isBotEnable(service.Comments.Leading.String()) {
+			return true
+		}
+		for _, method := range service.Methods {
+			if isBotEnable(method.Comments.Leading.String()) {
 				return true
 			}
 		}
 	}
 	return false
 }
+
+func isBotEnable(comment string) bool {
+	for _, line := range strings.Split(comment, "\n") {
+		if botTagsMatchRegexp.MatchString(line) {
+			return true
+		}
+	}
+	return false
+}
diff --git a/contrib/protoc-gen-bot/template.go.tpl b/contrib/protoc-gen-bot/template.go.tpl
@@ -11,9 +11,6 @@ const BotHandler{{$svrType}}{{.OriginalName}} = "/{{$svrName}}/{{.OriginalName}}
 
 type {{.ServiceType}}Server interface {
 {{- range .MethodSets}}
-	{{- if ne .Comment ""}}
-	{{.Comment}}
-	{{- end}}
 	{{.Name}}(context.Context, *{{.Request}}) (*{{.Reply}}, error)
 {{- end}}
 }

diff --git a/proto/bot/v1/counter.proto b/proto/bot/v1/counter.proto
@@ -2,10 +2,12 @@ syntax = "proto3";
 
 package bot.v1;
 
-// @bot
 service CounterService {
+  // @bot `comment:"/start"`
   rpc Start(StartRequest) returns (StartResponse);
+  // @bot `comment:"/count", callback_query:"count"`
   rpc Counter(CounterRequest) returns (CounterResponse);
+  rpc Unknown(UnknownRequest) returns (UnknownResponse);
 }
 
 message StartRequest {
@@ -23,4 +25,10 @@ message CounterRequest {
 
 message CounterResponse {
   int32 count = 1;
+}
+
+message UnknownRequest {
+}
+
+message UnknownResponse {
 }