fix(entproto): allow specifying import path

Basically a merge of: ent#616

Author: kdevo

entproto/adapter.go

@@ -57,12 +57,13 @@ var (
 )
 
 // LoadAdapter takes a *gen.Graph and parses it into protobuf file descriptors
-func LoadAdapter(graph *gen.Graph) (*Adapter, error) {
+func LoadAdapter(graph *gen.Graph, goPkg string) (*Adapter, error) {
 	a := &Adapter{
 		graph:            graph,
 		descriptors:      make(map[string]*desc.FileDescriptor),
 		schemaProtoFiles: make(map[string]string),
 		errors:           make(map[string]error),
+		goPkg:            goPkg,
 	}
 	if err := a.parse(); err != nil {
 		return nil, err
@@ -76,6 +77,7 @@ type Adapter struct {
 	descriptors      map[string]*desc.FileDescriptor
 	schemaProtoFiles map[string]string
 	errors           map[string]error
+	goPkg            string
 }
 
 // AllFileDescriptors returns a file descriptor per proto package for each package that contains
@@ -204,10 +206,15 @@ func (a *Adapter) parse() error {
 }
 
 func (a *Adapter) goPackageName(protoPkgName string) string {
-	// TODO(rotemtam): make this configurable from an annotation
-	entBase := a.graph.Config.Package
-	slashed := strings.ReplaceAll(protoPkgName, ".", "/")
-	return path.Join(entBase, "proto", slashed)
+	// TODO(kdevo): maybe better to make this configurable from an annotation
+	slashedProtoPkg := strings.ReplaceAll(protoPkgName, ".", "/")
+	if a.goPkg == "" {
+		entBase := a.graph.Config.Package
+		return path.Join(entBase, "proto", slashedProtoPkg)
+	} else {
+		slashed := strings.ReplaceAll(protoPkgName, ".", "/")
+		return path.Join(a.goPkg, slashed)
+	}
 }
 
 // GetFileDescriptor returns the proto file descriptor containing the transformed proto message descriptor for

entproto/extension.go

@@ -58,6 +58,7 @@ type Extension struct {
 	entc.DefaultExtension
 	protoDir    string
 	skipGenFile bool
+	goPkg       string
 }
 
 // WithProtoDir sets the directory where the generated .proto files will be written.
@@ -67,6 +68,13 @@ func WithProtoDir(dir string) ExtensionOption {
 	}
 }
 
+// WithProtoDir sets the directory where the generated .proto files will be written.
+func WithGoPkg(pkg string) ExtensionOption {
+	return func(e *Extension) {
+		e.goPkg = pkg
+	}
+}
+
 // SkipGenFile skips the generation of a generate.go file next to each .proto file.
 func SkipGenFile() ExtensionOption {
 	return func(e *Extension) {
@@ -124,10 +132,11 @@ func (e *Extension) generate(g *gen.Graph) error {
 	if e.protoDir != "" {
 		entProtoDir = e.protoDir
 	}
-	adapter, err := LoadAdapter(g)
+	adapter, err := LoadAdapter(g, e.goPkg)
 	if err != nil {
 		return fmt.Errorf("entproto: failed parsing ent graph: %w", err)
 	}
+
 	var errs error
 	for _, schema := range g.Schemas {
 		name := schema.Name
@@ -197,7 +206,7 @@ func (e *Extension) generate(g *gen.Graph) error {
 				if err != nil {
 					return err
 				}
-				contents := protocGenerateGo(fd, toSchema, toEnt, g.Config.Package)
+				contents := e.protocGenerateGo(fd, toSchema, toEnt, g.Config.Package)
 				if err := os.WriteFile(genGoPath, []byte(contents), 0600); err != nil {
 					return fmt.Errorf("entproto: failed generating generate.go file for %q: %w", protoFilePath, err)
 				}
@@ -216,7 +225,7 @@ func fileExists(fpath string) bool {
 	return true
 }
 
-func protocGenerateGo(fd *desc.FileDescriptor, toSchemaDir, entPath, entPackage string) string {
+func (e *Extension) protocGenerateGo(fd *desc.FileDescriptor, toSchemaDir, entPath, entPackage string) string {
 	levelsUp := len(strings.Split(fd.GetPackage(), "."))
 	toProtoBase := ""
 	for i := 0; i < levelsUp; i++ {