Merge pull request #905 from kaialang/presentation-zanzibar-v1

Allow proto files to have a separate configurable idl root

codegen/client.go

@@ -252,7 +252,7 @@ func newGRPCClientSpec(
 	instance *ModuleInstance,
 	h *PackageHelper,
 ) (*ClientSpec, error) {
-	protoFile := filepath.Join(h.IdlPath(), h.GetModuleIdlSubDir(false), config.IDLFile)
+	protoFile := filepath.Join(h.ProtoIdlPath(), h.GetModuleIdlSubDir(false), config.IDLFile)
 	protoSpec, err := NewProtoModuleSpec(protoFile, false, h)
 	if err != nil {
 		return nil, errors.Wrapf(

codegen/client_test.go

@@ -438,6 +438,7 @@ func newTestPackageHelper(t *testing.T) *PackageHelper {
 		},
 		TraceKey:        "trace-key",
 		ModuleIdlSubDir: map[string]string{"endpoints": "endpoints-idl", "default": "clients-idl"},
+		ProtoRelIdlRootDir: "./bazel-out/idl",
 	}
 
 	h, err := NewPackageHelper(
@@ -540,8 +541,8 @@ func TestGRPCClientNewClientSpec(t *testing.T) {
 		},
 	}
 	h := newTestPackageHelper(t)
-
-	idlFile := filepath.Join(h.IdlPath(), h.GetModuleIdlSubDir(false), "clients/echo/echo.proto")
+	assert.Equal(t, "/go/src/github.com/uber/zanzibar/examples/example-gateway/bazel-out/idl", h.ProtoIdlPath())
+	idlFile := filepath.Join(h.ProtoIdlPath(), h.GetModuleIdlSubDir(false), "clients/echo/echo.proto")
 	expectedSpec := &ClientSpec{
 		ModuleSpec:         nil,
 		YAMLFile:           instance.YAMLFileName,

codegen/package.go

@@ -1,4 +1,4 @@
-// Copyright (c) 2022 Uber Technologies, Inc.
+// Copyright (c) 2023 Uber Technologies, Inc.
 //
 // Permission is hereby granted, free of charge, to any person obtaining a copy
 // of this software and associated documentation files (the "Software"), to deal
@@ -45,6 +45,8 @@ type PackageHelper struct {
 	configRoot string
 	// The absolute root directory containing idl files
 	idlRootDir string
+	// The absolute root directory containing proto idl files
+	protoIdlRootDir string
 	// moduleIdlSubDir defines subdir for idl per module
 	moduleIdlSubDir map[string]string
 	// The map of idl type to go package name of where the generated structs are
@@ -86,6 +88,8 @@ func NewDefaultPackageHelperOptions() *PackageHelperOptions {
 type PackageHelperOptions struct {
 	// relative path to the idl dir, defaults to "./idl"
 	RelIdlRootDir string
+	// relative path to the proto idl dir, defaults to RelIdlRootDir, if neither of them is set, defaults to "./idl"
+	ProtoRelIdlRootDir string
 	// subdir for idl per module
 	ModuleIdlSubDir map[string]string
 	// relative path to the target dir that will contain generated code, defaults to "./build"
@@ -138,6 +142,14 @@ func (p *PackageHelperOptions) relIdlRootDir() string {
 	return "./idl"
 }
 
+func (p *PackageHelperOptions) protoRelIdlRootDir() string {
+	if p.ProtoRelIdlRootDir != "" {
+		return p.ProtoRelIdlRootDir
+	}
+
+	return p.relIdlRootDir()
+}
+
 func (p *PackageHelperOptions) moduleIdlSubDir() map[string]string {
 	return p.ModuleIdlSubDir
 }
@@ -232,6 +244,7 @@ func NewPackageHelper(
 		packageRoot:            packageRoot,
 		configRoot:             absConfigRoot,
 		idlRootDir:             filepath.Join(absConfigRoot, options.relIdlRootDir()),
+		protoIdlRootDir:        filepath.Join(absConfigRoot, options.protoRelIdlRootDir()),
 		genCodePackage:         options.genCodePackage(packageRoot),
 		goGatewayNamespace:     goGatewayNamespace,
 		targetGenDir:           filepath.Join(absConfigRoot, options.relTargetGenDir()),
@@ -279,6 +292,7 @@ func (p PackageHelper) TypeImportPath(idlFile string) (string, error) {
 	if !strings.HasSuffix(idlFile, thriftExtension) && !strings.HasSuffix(idlFile, protoExtension) {
 		return "", errors.Errorf("idl file %s is not %s or %s", idlFile, thriftExtension, protoExtension)
 	}
+
 	var suffix string
 	// for a filepath: a/b/c.(thrift|proto)
 	// - thrift generates code in path: a/b/c/c.go
@@ -290,10 +304,11 @@ func (p PackageHelper) TypeImportPath(idlFile string) (string, error) {
 	}
 
 	idx := strings.Index(idlFile, p.idlRootDir)
-	if idx == -1 {
+	idxProto := strings.Index(idlFile, p.protoIdlRootDir)
+	if idx == -1 && idxProto == -1 {
 		return "", errors.Errorf(
-			"file %s is not in IDL dir (%s)",
-			idlFile, p.idlRootDir,
+			"file %s is not in IDL dir (%s) or (%s)",
+			idlFile, p.idlRootDir, p.protoIdlRootDir,
 		)
 	}
 
@@ -302,6 +317,14 @@ func (p PackageHelper) TypeImportPath(idlFile string) (string, error) {
 	if !ok {
 		return "", errors.Errorf("genCodePackage for %q idl file is not configured in build.yaml", ext)
 	}
+
+	if strings.HasSuffix(idlFile, protoExtension) {
+		return path.Join(
+			genCodePkg,
+			idlFile[idxProto+len(p.protoIdlRootDir):len(suffix)],
+		), nil
+	}
+
 	return path.Join(
 		genCodePkg,
 		idlFile[idx+len(p.idlRootDir):len(suffix)],
@@ -318,6 +341,11 @@ func (p PackageHelper) IdlPath() string {
 	return p.idlRootDir
 }
 
+// ProtoIdlPath returns the file path to the proto idl folder
+func (p PackageHelper) ProtoIdlPath() string {
+	return p.protoIdlRootDir
+}
+
 // CodeGenTargetPath returns the file path where the code should
 // be generated.
 func (p PackageHelper) CodeGenTargetPath() string {
@@ -330,22 +358,26 @@ func (p PackageHelper) TypePackageName(idlFile string) (string, error) {
 		return "", errors.Errorf("file %s is not %s or %s", idlFile, thriftExtension, protoExtension)
 	}
 	idx := strings.Index(idlFile, p.idlRootDir)
-	if idx == -1 {
+	idxProto := strings.Index(idlFile, p.protoIdlRootDir)
+	if idx == -1 && idxProto == -1 {
 		return "", errors.Errorf(
-			"file %s is not in IDL dir (%s)",
-			idlFile, p.idlRootDir,
+			"file %s is not in IDL dir (%s) or (%s)",
+			idlFile, p.idlRootDir, p.protoIdlRootDir,
 		)
 	}
+
 	var prefix string
+	var segment string
 	if strings.HasSuffix(idlFile, thriftExtension) {
 		prefix = strings.TrimSuffix(idlFile, thriftExtension)
+		// Strip the leading / and the trailing .thrift suffix.
+		segment = idlFile[idx+len(p.idlRootDir)+1 : len(prefix)]
 	} else {
 		prefix = strings.TrimSuffix(idlFile, protoExtension)
+		// Strip the leading / and the trailing .proto suffix.
+		segment = idlFile[idxProto+len(p.protoIdlRootDir)+1 : len(prefix)]
 	}
 
-	// Strip the leading / and the trailing .thrift/.proto suffix.
-	segment := idlFile[idx+len(p.idlRootDir)+1 : len(prefix)]
-
 	packageName := strings.Replace(segment, "/", "_", -1)
 	return CamelCase(packageName), nil
 }
@@ -355,12 +387,18 @@ func (p PackageHelper) getRelativeFileName(idlFile string) (string, error) {
 		return "", errors.Errorf("file %s is not %s or %s", idlFile, thriftExtension, protoExtension)
 	}
 	idx := strings.Index(idlFile, p.idlRootDir)
-	if idx == -1 {
+	idxProto := strings.Index(idlFile, p.protoIdlRootDir)
+	if idx == -1 && idxProto == -1 {
 		return "", errors.Errorf(
-			"file %s is not in IDL dir (%s)",
-			idlFile, p.idlRootDir,
+			"file %s is not in IDL dir (%s) or (%s)",
+			idlFile, p.idlRootDir, p.protoIdlRootDir,
 		)
 	}
+
+	if strings.HasSuffix(idlFile, protoExtension) {
+		return idlFile[idxProto+len(p.protoIdlRootDir):], nil
+	}
+
 	return idlFile[idx+len(p.idlRootDir):], nil
 }
 

codegen/package_test.go

@@ -38,6 +38,12 @@ var fooThrift = filepath.Join(
 	"examples/example-gateway/idl/",
 	"clients-idl/clients/foo/foo.thrift")
 
+var fooProto = filepath.Join(
+	os.Getenv("GOPATH"),
+	"/src/github.com/uber/zanzibar/",
+	"examples/example-gateway/bazel-out/idl/",
+	"clients-idl/clients/foo/foo.proto")
+
 var testCopyrightHeader = `// Copyright (c) 2018 Uber Technologies, Inc.
 //
 // Permission is hereby granted, free of charge, to any person obtaining a copy
@@ -74,6 +80,7 @@ func newPackageHelper(t *testing.T) *codegen.PackageHelper {
 			".proto":  packageRoot + "/build/gen-code",
 		},
 		TraceKey: "trace-key",
+		ProtoRelIdlRootDir: "./bazel-out/idl",
 	}
 
 	h, err := codegen.NewPackageHelper(
@@ -87,14 +94,39 @@ func newPackageHelper(t *testing.T) *codegen.PackageHelper {
 	return h
 }
 
+func TestProtoRelIdlRootDirInitializeToIDL(t *testing.T) {
+	relativeGatewayPath := "../examples/example-gateway"
+	absGatewayPath, err := filepath.Abs(relativeGatewayPath)
+	assert.NoError(t, err)
+
+	packageRoot := "foo/"
+	options := &codegen.PackageHelperOptions{}
+
+	h, err := codegen.NewPackageHelper(
+		packageRoot,
+		absGatewayPath,
+		options,
+	)
+
+	assert.NoError(t, err)
+	assert.Equal(t, "/go/src/github.com/uber/zanzibar/examples/example-gateway/idl", h.ProtoIdlPath())
+}
+
 func TestImportPath(t *testing.T) {
 	h := newPackageHelper(t)
 	p, err := h.TypeImportPath(fooThrift)
 	assert.Nil(t, err, "should not return error")
 	assert.Equal(t, "github.com/uber/zanzibar/examples/example-gateway/build/gen-code/clients-idl/clients/foo/foo",
 		p, "wrong type import path")
+
+	p, err = h.TypeImportPath(fooProto)
+	assert.Nil(t, err, "should not return error")
+	assert.Equal(t, "github.com/uber/zanzibar/examples/example-gateway/build/gen-code/clients-idl/clients/foo",
+		p, "wrong type import path")
+
 	_, err = h.TypeImportPath("/Users/xxx/go/src/github.com/uber/zanzibar/examples/example-gateway/build/idl/github.com/uber/zanzibar/clients/foo/foo.go")
 	assert.Error(t, err, "should return error for not a thrift file")
+
 	_, err = h.TypeImportPath("/Users/xxx/go/src/github.com/uber/zanzibar/examples/example-gateway/build/zanzibar/clients/foo/foo.thrift")
 	assert.Error(t, err, "should return error for not in IDL dir")
 }
@@ -104,6 +136,11 @@ func TestTypePackageName(t *testing.T) {
 	packageName, err := h.TypePackageName(fooThrift)
 	assert.Nil(t, err, "should not return error")
 	assert.Equal(t, "clientsIDlClientsFooFoo", packageName, "wrong package name")
+
+	packageName, err = h.TypePackageName(fooProto)
+	assert.Nil(t, err, "should not return error")
+	assert.Equal(t, "clientsIDlClientsFooFoo", packageName, "wrong package name")
+
 	_, err = h.TypeImportPath("/Users/xxx/go/src/github.com/uber/zanzibar/examples/example-gateway/build/idl/github.com/uber/zanzibar/clients/foo/foo.txt")
 	assert.Error(t, err, "should return error for not a thrift file")
 }

codegen/service_test.go

@@ -43,13 +43,13 @@ func TestModuleSpecNoThriftExist(t *testing.T) {
 }
 
 func TestProtoModuleSpec(t *testing.T) {
-	echoProto := "../examples/example-gateway/idl/clients-idl/clients/echo/echo.proto"
+	echoProto := "../examples/example-gateway/bazel-out/idl/clients-idl/clients/echo/echo.proto"
 	_, err := codegen.NewProtoModuleSpec(echoProto, false, newPackageHelper(t))
 	assert.NoError(t, err, "unable to parse the proto file")
 }
 
 func TestProtoModuleSpecParseError(t *testing.T) {
-	tmpFile, err := ioutil.TempFile("../examples/example-gateway/idl/clients-idl/clients/", "temp*.proto")
+	tmpFile, err := ioutil.TempFile("../examples/example-gateway/bazel-out/idl/clients-idl/clients/", "temp*.proto")
 	assert.NoError(t, err, "failed to create temp file")
 	defer func(name string) {
 		_ = os.Remove(name)

examples/example-gateway/build.yaml

@@ -14,6 +14,7 @@ defaultHeaders:
 packageRoot: github.com/uber/zanzibar/examples/example-gateway
 targetGenDir: ./build
 idlRootDir: ./idl
+protoIdlRootDir: ./bazel-out/idl
 moduleIdlSubDir:
   endpoints: endpoints-idl
   default: clients-idl