|
| 1 | +package main |
| 2 | + |
| 3 | +import ( |
| 4 | + "fmt" |
| 5 | + "slices" |
| 6 | + "sort" |
| 7 | + "strings" |
| 8 | + |
| 9 | + nexusannotationsv1 "github.com/nexus-rpc/nexus-proto-annotations/go/nexusannotations/v1" |
| 10 | + "google.golang.org/protobuf/compiler/protogen" |
| 11 | + "google.golang.org/protobuf/proto" |
| 12 | + "google.golang.org/protobuf/reflect/protoreflect" |
| 13 | + "google.golang.org/protobuf/types/descriptorpb" |
| 14 | + "gopkg.in/yaml.v3" |
| 15 | +) |
| 16 | + |
| 17 | +// params holds the parsed protoc plugin options. |
| 18 | +// Passed via --nexus-rpc-yaml_opt=key=value (multiple opts are comma-joined by protoc). |
| 19 | +// |
| 20 | +// - nexus-rpc_langs_out: optional. Output path for the langs YAML. |
| 21 | +// If empty, nothing is written. |
| 22 | +// Example: "nexus/temporal-proto-models-nexusrpc.yaml" |
| 23 | +// |
| 24 | +// - python_package_prefix: optional. Dot-separated package prefix for $pythonRef. |
| 25 | +// The last two path segments of the go_package ({service}/v{n}) are appended. |
| 26 | +// Example: "temporalio.api" → "temporalio.api.workflowservice.v1.TypeName" |
| 27 | +// If empty, $pythonRef is omitted. |
| 28 | +// |
| 29 | +// - typescript_package_prefix: optional. Scoped package prefix for $typescriptRef. |
| 30 | +// The last two path segments of the go_package ({service}/v{n}) are appended. |
| 31 | +// Example: "@temporalio/api" → "@temporalio/api/workflowservice/v1.TypeName" |
| 32 | +// If empty, $typescriptRef is omitted. |
| 33 | +// |
| 34 | +// - include_operation_tags: optional, repeatable. Only include operations whose tags |
| 35 | +// contain at least one of these values. If empty, all annotated operations are included |
| 36 | +// (subject to exclude_operation_tags). Specify multiple times for multiple tags. |
| 37 | +// Example: include_operation_tags=exposed |
| 38 | +// |
| 39 | +// - exclude_operation_tags: optional, repeatable. Exclude operations whose tags contain |
| 40 | +// any of these values. Applied after include_operation_tags. |
| 41 | +// Example: exclude_operation_tags=internal |
| 42 | +type params struct { |
| 43 | + nexusRpcLangsOut string |
| 44 | + pythonPackagePrefix string |
| 45 | + typescriptPackagePrefix string |
| 46 | + includeOperationTags []string |
| 47 | + excludeOperationTags []string |
| 48 | +} |
| 49 | + |
| 50 | +// parseParams parses the comma-separated key=value parameter string provided by protoc. |
| 51 | +func parseParams(raw string) (params, error) { |
| 52 | + var p params |
| 53 | + if raw == "" { |
| 54 | + return p, nil |
| 55 | + } |
| 56 | + for kv := range strings.SplitSeq(raw, ",") { |
| 57 | + key, value, ok := strings.Cut(kv, "=") |
| 58 | + if !ok { |
| 59 | + return p, fmt.Errorf("invalid parameter %q: expected key=value", kv) |
| 60 | + } |
| 61 | + switch key { |
| 62 | + case "nexus-rpc_langs_out": |
| 63 | + p.nexusRpcLangsOut = value |
| 64 | + case "python_package_prefix": |
| 65 | + p.pythonPackagePrefix = value |
| 66 | + case "typescript_package_prefix": |
| 67 | + p.typescriptPackagePrefix = value |
| 68 | + case "include_operation_tags": |
| 69 | + p.includeOperationTags = append(p.includeOperationTags, value) |
| 70 | + case "exclude_operation_tags": |
| 71 | + p.excludeOperationTags = append(p.excludeOperationTags, value) |
| 72 | + default: |
| 73 | + return p, fmt.Errorf("unknown parameter %q", key) |
| 74 | + } |
| 75 | + } |
| 76 | + return p, nil |
| 77 | +} |
| 78 | + |
| 79 | +// shouldIncludeOperation returns true if the method's nexus operation tags pass |
| 80 | +// the include/exclude filters. Mirrors the logic from protoc-gen-go-nexus: |
| 81 | +// 1. Method must have the nexus operation extension set. |
| 82 | +// 2. If includeOperationTags is non-empty, at least one of the method's tags must match. |
| 83 | +// 3. If excludeOperationTags is non-empty, none of the method's tags may match. |
| 84 | +func shouldIncludeOperation(p params, m *protogen.Method) bool { |
| 85 | + opts, ok := m.Desc.Options().(*descriptorpb.MethodOptions) |
| 86 | + if !ok || opts == nil { |
| 87 | + return false |
| 88 | + } |
| 89 | + if !proto.HasExtension(opts, nexusannotationsv1.E_Operation) { |
| 90 | + return false |
| 91 | + } |
| 92 | + tags := proto.GetExtension(opts, nexusannotationsv1.E_Operation).(*nexusannotationsv1.OperationOptions).GetTags() |
| 93 | + if len(p.includeOperationTags) > 0 && !slices.ContainsFunc(p.includeOperationTags, func(t string) bool { |
| 94 | + return slices.Contains(tags, t) |
| 95 | + }) { |
| 96 | + return false |
| 97 | + } |
| 98 | + return !slices.ContainsFunc(p.excludeOperationTags, func(t string) bool { |
| 99 | + return slices.Contains(tags, t) |
| 100 | + }) |
| 101 | +} |
| 102 | + |
| 103 | +func generate(gen *protogen.Plugin) error { |
| 104 | + p, err := parseParams(gen.Request.GetParameter()) |
| 105 | + if err != nil { |
| 106 | + return err |
| 107 | + } |
| 108 | + |
| 109 | + langsDoc := newDoc() |
| 110 | + hasOps := false |
| 111 | + |
| 112 | + for _, f := range gen.Files { |
| 113 | + if !f.Generate { |
| 114 | + continue |
| 115 | + } |
| 116 | + for _, svc := range f.Services { |
| 117 | + for _, m := range svc.Methods { |
| 118 | + if !shouldIncludeOperation(p, m) { |
| 119 | + continue |
| 120 | + } |
| 121 | + svcName := string(svc.Desc.Name()) |
| 122 | + methodName := string(m.Desc.Name()) |
| 123 | + hasOps = true |
| 124 | + addOperation(langsDoc, svcName, methodName, |
| 125 | + langRefs(p, f.Desc, m.Input.Desc), |
| 126 | + langRefs(p, f.Desc, m.Output.Desc), |
| 127 | + ) |
| 128 | + } |
| 129 | + } |
| 130 | + } |
| 131 | + |
| 132 | + if !hasOps { |
| 133 | + return nil |
| 134 | + } |
| 135 | + if p.nexusRpcLangsOut != "" { |
| 136 | + return writeFile(gen, p.nexusRpcLangsOut, langsDoc) |
| 137 | + } |
| 138 | + return nil |
| 139 | +} |
| 140 | + |
| 141 | +// langRefs builds the map of language-specific type refs for a message. |
| 142 | +// |
| 143 | +// Go, Java, dotnet, and Ruby refs are derived from proto file-level package options. |
| 144 | +// Python and TypeScript refs require the corresponding prefix params to be set; if |
| 145 | +// empty they are omitted. Both use the last two path segments of go_package |
| 146 | +// ({service}/v{n}), dropping any intermediate grouping directory. |
| 147 | +func langRefs(p params, file protoreflect.FileDescriptor, msg protoreflect.MessageDescriptor) map[string]string { |
| 148 | + opts, ok := file.Options().(*descriptorpb.FileOptions) |
| 149 | + if !ok || opts == nil { |
| 150 | + return nil |
| 151 | + } |
| 152 | + name := string(msg.Name()) |
| 153 | + refs := make(map[string]string) |
| 154 | + |
| 155 | + if pkg := opts.GetGoPackage(); pkg != "" { |
| 156 | + // strip the ";alias" suffix (e.g. "go.temporal.io/api/workflowservice/v1;workflowservice") |
| 157 | + pkg = strings.SplitN(pkg, ";", 2)[0] |
| 158 | + refs["$goRef"] = pkg + "." + name |
| 159 | + |
| 160 | + segments := strings.Split(pkg, "/") |
| 161 | + if len(segments) >= 2 { |
| 162 | + tail := segments[len(segments)-2] + "/" + segments[len(segments)-1] |
| 163 | + if p.pythonPackagePrefix != "" { |
| 164 | + dotTail := strings.ReplaceAll(tail, "/", ".") |
| 165 | + refs["$pythonRef"] = p.pythonPackagePrefix + "." + dotTail + "." + name |
| 166 | + } |
| 167 | + if p.typescriptPackagePrefix != "" { |
| 168 | + refs["$typescriptRef"] = p.typescriptPackagePrefix + "/" + tail + "." + name |
| 169 | + } |
| 170 | + } |
| 171 | + } |
| 172 | + if pkg := opts.GetJavaPackage(); pkg != "" { |
| 173 | + refs["$javaRef"] = pkg + "." + name |
| 174 | + } |
| 175 | + if pkg := opts.GetRubyPackage(); pkg != "" { |
| 176 | + refs["$rubyRef"] = pkg + "::" + name |
| 177 | + } |
| 178 | + if pkg := opts.GetCsharpNamespace(); pkg != "" { |
| 179 | + refs["$dotnetRef"] = pkg + "." + name |
| 180 | + } |
| 181 | + if len(refs) == 0 { |
| 182 | + return nil |
| 183 | + } |
| 184 | + return refs |
| 185 | +} |
| 186 | + |
| 187 | +// newDoc creates a yaml.Node document with the "nexusrpc: 1.0.0" header |
| 188 | +// and an empty "services" mapping node. |
| 189 | +func newDoc() *yaml.Node { |
| 190 | + doc := &yaml.Node{Kind: yaml.DocumentNode} |
| 191 | + root := &yaml.Node{Kind: yaml.MappingNode, Tag: "!!map"} |
| 192 | + doc.Content = []*yaml.Node{root} |
| 193 | + root.Content = append(root.Content, |
| 194 | + scalarNode("nexusrpc"), |
| 195 | + scalarNode("1.0.0"), |
| 196 | + scalarNode("services"), |
| 197 | + &yaml.Node{Kind: yaml.MappingNode, Tag: "!!map"}, |
| 198 | + ) |
| 199 | + return doc |
| 200 | +} |
| 201 | + |
| 202 | +// servicesNode returns the "services" mapping node from a doc created by newDoc. |
| 203 | +func servicesNode(doc *yaml.Node) *yaml.Node { |
| 204 | + root := doc.Content[0] |
| 205 | + for i := 0; i < len(root.Content)-1; i += 2 { |
| 206 | + if root.Content[i].Value == "services" { |
| 207 | + return root.Content[i+1] |
| 208 | + } |
| 209 | + } |
| 210 | + panic("services node not found") |
| 211 | +} |
| 212 | + |
| 213 | +// addOperation inserts a service → operation → {input, output} entry into doc. |
| 214 | +// Services and operations are inserted in the order first encountered. |
| 215 | +func addOperation(doc *yaml.Node, svcName, methodName string, input, output map[string]string) { |
| 216 | + svcs := servicesNode(doc) |
| 217 | + |
| 218 | + var svcOps *yaml.Node |
| 219 | + for i := 0; i < len(svcs.Content)-1; i += 2 { |
| 220 | + if svcs.Content[i].Value == svcName { |
| 221 | + svcMap := svcs.Content[i+1] |
| 222 | + for j := 0; j < len(svcMap.Content)-1; j += 2 { |
| 223 | + if svcMap.Content[j].Value == "operations" { |
| 224 | + svcOps = svcMap.Content[j+1] |
| 225 | + } |
| 226 | + } |
| 227 | + } |
| 228 | + } |
| 229 | + if svcOps == nil { |
| 230 | + svcMap := &yaml.Node{Kind: yaml.MappingNode, Tag: "!!map"} |
| 231 | + svcOps = &yaml.Node{Kind: yaml.MappingNode, Tag: "!!map"} |
| 232 | + svcMap.Content = append(svcMap.Content, scalarNode("operations"), svcOps) |
| 233 | + svcs.Content = append(svcs.Content, scalarNode(svcName), svcMap) |
| 234 | + } |
| 235 | + |
| 236 | + opNode := &yaml.Node{Kind: yaml.MappingNode, Tag: "!!map"} |
| 237 | + if len(input) > 0 { |
| 238 | + opNode.Content = append(opNode.Content, scalarNode("input"), mapNode(input)) |
| 239 | + } |
| 240 | + if len(output) > 0 { |
| 241 | + opNode.Content = append(opNode.Content, scalarNode("output"), mapNode(output)) |
| 242 | + } |
| 243 | + svcOps.Content = append(svcOps.Content, scalarNode(methodName), opNode) |
| 244 | +} |
| 245 | + |
| 246 | +// mapNode serializes a map[string]string as a yaml mapping node with keys in sorted order. |
| 247 | +func mapNode(m map[string]string) *yaml.Node { |
| 248 | + keys := make([]string, 0, len(m)) |
| 249 | + for k := range m { |
| 250 | + keys = append(keys, k) |
| 251 | + } |
| 252 | + sort.Strings(keys) |
| 253 | + node := &yaml.Node{Kind: yaml.MappingNode, Tag: "!!map"} |
| 254 | + for _, k := range keys { |
| 255 | + node.Content = append(node.Content, scalarNode(k), scalarNode(m[k])) |
| 256 | + } |
| 257 | + return node |
| 258 | +} |
| 259 | + |
| 260 | +func scalarNode(value string) *yaml.Node { |
| 261 | + return &yaml.Node{Kind: yaml.ScalarNode, Tag: "!!str", Value: value} |
| 262 | +} |
| 263 | + |
| 264 | +func writeFile(gen *protogen.Plugin, name string, doc *yaml.Node) error { |
| 265 | + f := gen.NewGeneratedFile(name, "") |
| 266 | + enc := yaml.NewEncoder(f) |
| 267 | + enc.SetIndent(2) |
| 268 | + if err := enc.Encode(doc); err != nil { |
| 269 | + return err |
| 270 | + } |
| 271 | + return enc.Close() |
| 272 | +} |
0 commit comments