Repository: goadesign/goa
Pull request #3065 (2022-05-14 18:10:57)

Custom proto field types


Patch:
diff --git a/dsl/meta.go b/dsl/meta.go
index e3b6daeb..afe16e4b 100644
--- a/dsl/meta.go
+++ b/dsl/meta.go
@@ -76,6 +76,25 @@ import (
 //         })
 //    })
 //
+// - "struct:field:proto" overrides the generated protobuf field type. The second
+// argument is optional and if present indicates an import path for the proto file
+// defining the type.
+//
+//    var Timestamp = Type("Timestamp", func() {
+//        Description("Google timestamp compatible design")
+//        Field(1, "seconds", Int64, "Unix timestamp in seconds", func() {
+//            Meta("struct:field:proto", "int64") // Goa generates sint64 by default
+//        })
+//        Field(2, "nanos", Int32, "Unix timestamp in nanoseconds", func() {
+//            Meta("struct:field:proto", "int32") // Goa generates sint32 by default
+//        })
+//    })
+//
+//    var MyType = Type("MyType", func() {
+//        Field(1, "created_at", Timestamp, func() {
+//            Meta("struct:field:proto", "google.protobuf.Timestamp", "google/protobuf/timestamp.proto")
+//        })
+//    })
 //
 // - "struct:tag:xxx" sets a generated Go struct field tag and overrides tags
 // that Goa would otherwise set. If the metadata value is a slice then the
@@ -89,6 +108,18 @@ import (
 //        })
 //    })
 //
+// - "protoc:include" provides the list of import paths used to invoke protoc.
+// Applicable to API and service definitions only. If used on an API definition
+// the include paths are used for all services.
+//
+//    var _ = API("myapi", func() {
+//        Meta("protoc:include", "/usr/include", "/usr/local/include")
+//    })
+//
+//    var _ = Service("service1", func() {
+//        Meta("protoc:include", "/usr/local/include/google/protobuf")
+//    })
+//
 // - "swagger:generate" DEPRECATED, use "openapi:generate" instead.
 //
 // - "openapi:generate" specifies whether OpenAPI specification should be
diff --git a/grpc/codegen/proto.go b/grpc/codegen/proto.go
index 10d2bc58..83647a77 100644
--- a/grpc/codegen/proto.go
+++ b/grpc/codegen/proto.go
@@ -47,6 +47,7 @@ func protoFile(genpkg string, svc *expr.GRPCServiceExpr) *codegen.File {
 			Data: map[string]interface{}{
 				"ProtoVersion": ProtoVersion,
 				"Pkg":          pkgName(svc, svcName),
+				"Imports":      data.Imports,
 			},
 		},
 		// service definition
@@ -62,10 +63,16 @@ func protoFile(genpkg string, svc *expr.GRPCServiceExpr) *codegen.File {
 		sections = append(sections, &codegen.SectionTemplate{Name: "grpc-message", Source: messageT, Data: m})
 	}
 
+	runProtoc := func(path string) error {
+		includes := svc.ServiceExpr.Meta["protoc:include"]
+		includes = append(includes, expr.Root.API.Meta["protoc:include"]...)
+		return protoc(path, includes)
+	}
+
 	return &codegen.File{
 		Path:             path,
 		SectionTemplates: sections,
-		FinalizeFunc:     protoc,
+		FinalizeFunc:     runProtoc,
 	}
 }
 
@@ -76,11 +83,21 @@ func pkgName(svc *expr.GRPCServiceExpr, svcName string) string {
 	return codegen.SnakeCase(svcName)
 }
 
-func protoc(path string) error {
+func protoc(path string, includes []string) error {
 	dir := filepath.Dir(path)
 	os.MkdirAll(dir, 0777)
 
-	args := []string{"--proto_path", dir, "--go_out", dir, "--go-grpc_out", dir, "--go_opt=paths=source_relative", "--go-grpc_opt=paths=source_relative", path}
+	args := []string{
+		path,
+		"--proto_path", dir,
+		"--go_out", dir,
+		"--go-grpc_out", dir,
+		"--go_opt=paths=source_relative",
+		"--go-grpc_opt=paths=source_relative",
+	}
+	for _, include := range includes {
+		args = append(args, "-I", include)
+	}
 	cmd := exec.Command("protoc", args...)
 	cmd.Dir = filepath.Dir(path)
 
@@ -108,6 +125,9 @@ syntax = {{ printf "%q" .ProtoVersion }};
 package {{ .Pkg }};
 
 option go_package = "/{{ .Pkg }}pb";
+{{- range .Imports }}
+import "{{ . }}";
+{{- end }}
 `
 
 	// input: ServiceData
diff --git a/grpc/codegen/protobuf.go b/grpc/codegen/protobuf.go
index 90e87ccd..ba2c2697 100644
--- a/grpc/codegen/protobuf.go
+++ b/grpc/codegen/protobuf.go
@@ -231,17 +231,25 @@ func protoBufGoFullTypeName(att *expr.AttributeExpr, pkg string, s *codegen.Name
 	}
 }
 
+// protoType returns the protocol buffer type name for the given attribute.
+func protoType(att *expr.AttributeExpr, sd *ServiceData) string {
+	if protos := att.Meta["struct:field:proto"]; len(protos) > 0 {
+		return protos[0]
+	}
+	return protoBufMessageDef(att, sd)
+}
+
 // protoBufMessageDef returns the protocol buffer code that defines a message
 // which matches the data structure definition (the part that comes after
 // `message foo`). The message is defined using the proto3 syntax.
 func protoBufMessageDef(att *expr.AttributeExpr, sd *ServiceData) string {
 	switch actual := att.Type.(type) {
 	case expr.Primitive:
-		return protoBufNativeMessageTypeName(att.Type)
+		return protoNativeType(att.Type)
 	case *expr.Array:
-		return "repeated " + protoBufMessageDef(actual.ElemType, sd)
+		return "repeated " + protoType(actual.ElemType, sd)
 	case *expr.Map:
-		return fmt.Sprintf("map<%s, %s>", protoBufMessageDef(actual.KeyType, sd), protoBufMessageDef(actual.ElemType, sd))
+		return fmt.Sprintf("map<%s, %s>", protoType(actual.KeyType, sd), protoType(actual.ElemType, sd))
 	case *expr.Union:
 		def := "\toneof " + codegen.SnakeCase(protoBufify(actual.Name(), false, false)) + " {"
 		for _, nat := range actual.Values {
@@ -249,9 +257,9 @@ func protoBufMessageDef(att *expr.AttributeExpr, sd *ServiceData) string {
 			fnum := rpcTag(nat.Attribute)
 			var typ string
 			if prim := getPrimitive(nat.Attribute); prim != nil {
-				typ = protoBufMessageDef(prim, sd)
+				typ = protoType(prim, sd)
 			} else {
-				typ = protoBufMessageDef(nat.Attribute, sd)
+				typ = protoType(nat.Attribute, sd)
 			}
 			var desc string
 			if d := nat.Attribute.Description; d != "" {
@@ -287,9 +295,9 @@ func protoBufMessageDef(att *expr.AttributeExpr, sd *ServiceData) string {
 				fn = codegen.SnakeCase(protoBufify(nat.Name, false, false))
 				fnum = rpcTag(nat.Attribute)
 				if prim := getPrimitive(nat.Attribute); prim != nil {
-					typ = protoBufMessageDef(prim, sd)
+					typ = protoType(prim, sd)
 				} else {
-					typ = protoBufMessageDef(nat.Attribute, sd)
+					typ = protoType(nat.Attribute, sd)
 				}
 				if nat.Attribute.Description != "" {
 					desc = codegen.Comment(nat.Attribute.Description) + "\n\t"
@@ -372,10 +380,10 @@ func protoBufifyAtt(att *expr.AttributeExpr, name string, upper bool) string {
 	return protoBufify(name, upper, false)
 }
 
-// protoBufNativeMessageTypeName returns the protocol buffer built-in type
+// protoNativeType returns the protocol buffer built-in type
 // corresponding to the given primitive type. It panics if t is not a
 // primitive type.
-func protoBufNativeMessageTypeName(t expr.DataType) string {
+func protoNativeType(t expr.DataType) string {
 	switch t.Kind() {
 	case expr.BooleanKind:
 		return "bool"
diff --git a/grpc/codegen/service_data.go b/grpc/codegen/service_data.go
index b3acb687..3c34fd9f 100644
--- a/grpc/codegen/service_data.go
+++ b/grpc/codegen/service_data.go
@@ -24,6 +24,8 @@ type (
 		Service *service.Data
 		// PkgName is the name of the generated package in *.pb.go.
 		PkgName string
+		// Imports is the list of proto package imports.
+		Imports []string
 		// Name is the service name.
 		Name string
 		// Description is the service description.
@@ -466,8 +468,19 @@ func (d ServicesData) analyze(gs *expr.GRPCServiceExpr) *ServiceData {
 		}
 
 		// collect all the nested messages and return the top-level message
+		// Also collect all proto imports specified via Meta.
 		collect := func(att *expr.AttributeExpr) *service.UserTypeData {
-			msgs := collectMessages(att, sd, seen)
+			msgs, imports := collectMessages(att, sd, seen)
+			if len(imports) > 0 {
+				imported := make(map[string]struct{})
+				for _, imp := range imports {
+					if _, ok := imported[imp]; ok {
+						continue
+					}
+					imported[imp] = struct{}{}
+					sd.Imports = append(sd.Imports, imp)
+				}
+			}
 			if len(msgs) > 0 {
 				sd.Messages = append(sd.Messages, msgs...)
 				return msgs[0]
@@ -634,17 +647,23 @@ func (d ServicesData) analyze(gs *expr.GRPCServiceExpr) *ServiceData {
 }
 
 // collectMessages recurses through the attribute to gather all the messages.
-func collectMessages(at *expr.AttributeExpr, sd *ServiceData, seen map[string]struct{}) (data []*service.UserTypeData) {
-	if at == nil || expr.IsPrimitive(at.Type) {
+func collectMessages(at *expr.AttributeExpr, sd *ServiceData, seen map[string]struct{}) (data []*service.UserTypeData, imports []string) {
+	if at == nil {
 		return
 	}
-	collect := func(at *expr.AttributeExpr) []*service.UserTypeData {
+	if proto := at.Meta["struct:field:proto"]; len(proto) > 1 {
+		imports = append(imports, proto[1])
+	}
+	if expr.IsPrimitive(at.Type) {
+		return
+	}
+	collect := func(at *expr.AttributeExpr) ([]*service.UserTypeData, []string) {
 		return collectMessages(at, sd, seen)
 	}
 	switch dt := at.Type.(type) {
 	case expr.UserType:
 		if _, ok := seen[dt.Name()]; ok {
-			return nil
+			return
 		}
 		att := dt.Attribute()
 		if rt, ok := dt.(*expr.ResultTypeExpr); ok {
@@ -662,19 +681,25 @@ func collectMessages(at *expr.AttributeExpr, sd *ServiceData, seen map[string]st
 			Type:        dt,
 		})
 		seen[dt.Name()] = struct{}{}
-		data = append(data, collect(att)...)
+		d, i := collect(att)
+		data, imports = append(data, d...), append(imports, i...)
 	case *expr.Object:
 		for _, nat := range *dt {
-			data = append(data, collect(nat.Attribute)...)
+			d, i := collect(nat.Attribute)
+			data, imports = append(data, d...), append(imports, i...)
 		}
 	case *expr.Array:
-		data = append(data, collect(dt.ElemType)...)
+		d, i := collect(dt.ElemType)
+		data, imports = append(data, d...), append(imports, i...)
 	case *expr.Map:
-		data = append(data, collect(dt.KeyType)...)
-		data = append(data, collect(dt.ElemType)...)
+		dk, ik := collect(dt.KeyType)
+		data, imports = append(data, dk...), append(imports, ik...)
+		de, ie := collect(dt.ElemType)
+		data, imports = append(data, de...), append(imports, ie...)
 	case *expr.Union:
 		for _, nat := range dt.Values {
-			data = append(data, collect(nat.Attribute)...)
+			d, i := collect(nat.Attribute)
+			data, imports = append(data, d...), append(imports, i...)
 		}
 	}
 	return


Test patch:
diff --git a/grpc/codegen/proto_test.go b/grpc/codegen/proto_test.go
index a1686bbd..29419c25 100644
--- a/grpc/codegen/proto_test.go
+++ b/grpc/codegen/proto_test.go
@@ -42,7 +42,7 @@ func TestProtoFiles(t *testing.T) {
 				t.Errorf("%s: got\n%s\ngot vs. expected:\n%s", c.Name, code, codegen.Diff(t, code, c.Code))
 			}
 			fpath := codegen.CreateTempFile(t, code)
-			if err := protoc(fpath); err != nil {
+			if err := protoc(fpath, nil); err != nil {
 				t.Fatalf("error occurred when compiling proto file %q: %s", fpath, err)
 			}
 		})
@@ -83,7 +83,7 @@ func TestMessageDefSection(t *testing.T) {
 				t.Errorf("%s: got\n%s\ngot vs. expected:\n%s", c.Name, msgCode, codegen.Diff(t, msgCode, c.Code))
 			}
 			fpath := codegen.CreateTempFile(t, code+msgCode)
-			if err := protoc(fpath); err != nil {
+			if err := protoc(fpath, nil); err != nil {
 				t.Fatalf("error occurred when compiling proto file %q: %s", fpath, err)
 			}
 		})