fix: compatible with add method opt in main (go-gorm#586)

a631807682 · Aug 16, 2022 · acb2bb9 · acb2bb9
1 parent c4a646f
commit acb2bb9
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 4 deletions.
diff --git a/internal/generate/export.go b/internal/generate/export.go
@@ -35,6 +35,7 @@ func GetQueryStructMeta(db *gorm.DB, conf *model.Config) (*QueryStructMeta, erro
 	if err != nil {
 		return nil, err
 	}
+
 	return (&QueryStructMeta{
 		db:              db,
 		Source:          model.Table,
@@ -47,7 +48,7 @@ func GetQueryStructMeta(db *gorm.DB, conf *model.Config) (*QueryStructMeta, erro
 		StructInfo:      parser.Param{Type: structName, Package: conf.ModelPkg},
 		ImportPkgPaths:  conf.ImportPkgPaths,
 		Fields:          getFields(db, conf, columns),
-	}).AddMethod(conf.GetModelMethods()...), nil
+	}).addMethodFromAddMethodOpt(conf.GetModelMethods()...), nil
 }
 
 // GetQueryStructMetaFromObject generate base struct from object

diff --git a/internal/generate/query.go b/internal/generate/query.go
@@ -161,7 +161,23 @@ func (b *QueryStructMeta) ReviseDIYMethod() error {
 // eg: g.GenerateModel("users").AddMethod(user.IsEmpty,user.GetName) or g.GenerateModel("users").AddMethod(model.User)
 func (b *QueryStructMeta) AddMethod(methods ...interface{}) *QueryStructMeta {
 	for _, method := range methods {
-		modelMethods, err := parser.GetModelMethod(method)
+		modelMethods, err := parser.GetModelMethod(method, 2)
+		if err != nil {
+			panic("add diy method err:" + err.Error())
+		}
+		b.ModelMethods = append(b.ModelMethods, modelMethods.Methods...)
+	}
+
+	err := b.ReviseDIYMethod()
+	if err != nil {
+		b.db.Logger.Warn(context.Background(), err.Error())
+	}
+	return b
+}
+
+func (b *QueryStructMeta) addMethodFromAddMethodOpt(methods ...interface{}) *QueryStructMeta {
+	for _, method := range methods {
+		modelMethods, err := parser.GetModelMethod(method, 4)
 		if err != nil {
 			panic("add diy method err:" + err.Error())
 		}

diff --git a/internal/parser/export.go b/internal/parser/export.go
@@ -76,7 +76,7 @@ func fileExists(path string) bool {
 }
 
 // GetModelMethod get diy methods
-func GetModelMethod(v interface{}) (method *DIYMethods, err error) {
+func GetModelMethod(v interface{}, skip int) (method *DIYMethods, err error) {
 	method = new(DIYMethods)
 
 	// get diy method info by input value, must input a function or a struct
@@ -100,7 +100,7 @@ func GetModelMethod(v interface{}) (method *DIYMethods, err error) {
 	// if struct in main file
 	ctx := build.Default
 	if method.pkgPath == "main" {
-		_, file, _, _ := runtime.Caller(2)
+		_, file, _, _ := runtime.Caller(skip)
 		p, err = ctx.ImportDir(filepath.Dir(file), build.ImportComment)
 	} else {
 		p, err = ctx.Import(method.pkgPath, "", build.ImportComment)