@@ -58,6 +58,7 @@ type Extension struct {
58
58
entc.DefaultExtension
59
59
protoDir string
60
60
skipGenFile bool
61
+ goPkg string // inspired from: https://github.com/ent/contrib/pull/616
61
62
}
62
63
63
64
// WithProtoDir sets the directory where the generated .proto files will be written.
@@ -67,6 +68,13 @@ func WithProtoDir(dir string) ExtensionOption {
67
68
}
68
69
}
69
70
71
+ // WithProtoDir sets the directory where the generated .proto files will be written.
72
+ func WithGoPkg (pkg string ) ExtensionOption {
73
+ return func (e * Extension ) {
74
+ e .protoDir = pkg
75
+ }
76
+ }
77
+
70
78
// SkipGenFile skips the generation of a generate.go file next to each .proto file.
71
79
func SkipGenFile () ExtensionOption {
72
80
return func (e * Extension ) {
@@ -197,7 +205,7 @@ func (e *Extension) generate(g *gen.Graph) error {
197
205
if err != nil {
198
206
return err
199
207
}
200
- contents := protocGenerateGo (fd , toSchema , toEnt , g .Config .Package )
208
+ contents := e . protocGenerateGo (fd , toSchema , toEnt , g .Config .Package )
201
209
if err := os .WriteFile (genGoPath , []byte (contents ), 0600 ); err != nil {
202
210
return fmt .Errorf ("entproto: failed generating generate.go file for %q: %w" , protoFilePath , err )
203
211
}
@@ -216,7 +224,7 @@ func fileExists(fpath string) bool {
216
224
return true
217
225
}
218
226
219
- func protocGenerateGo (fd * desc.FileDescriptor , toSchemaDir , entPath , entPackage string ) string {
227
+ func ( e * Extension ) protocGenerateGo (fd * desc.FileDescriptor , toSchemaDir , entPath , entPackage string ) string {
220
228
levelsUp := len (strings .Split (fd .GetPackage (), "." ))
221
229
toProtoBase := ""
222
230
for i := 0 ; i < levelsUp ; i ++ {
@@ -235,5 +243,8 @@ func protocGenerateGo(fd *desc.FileDescriptor, toSchemaDir, entPath, entPackage
235
243
}
236
244
goGen := fmt .Sprintf ("//go:generate %s" , strings .Join (protocCmd , " " ))
237
245
goPkgName := extractLastFqnPart (fd .GetPackage ())
246
+ if e .goPkg != "" {
247
+ goPkgName = e .goPkg
248
+ }
238
249
return fmt .Sprintf ("package %s\n %s\n " , goPkgName , goGen )
239
250
}
0 commit comments