diff --git a/internal/generate/query.go b/internal/generate/query.go index c9e27e10..f25835da 100644 --- a/internal/generate/query.go +++ b/internal/generate/query.go @@ -186,16 +186,12 @@ func (b *QueryStructMeta) ReviseDIYMethod() error { } method.Receiver.Package = "" method.Receiver.Type = b.ModelStructName + b.pasreTableName(method) methods = append(methods, method) methodMap[method.MethodName] = true } if tableName == nil { methods = append(methods, parser.DefaultMethodTableName(b.ModelStructName)) - } else { - // e.g. return "@@table" => return TableNameUser - tableName.Body = strings.ReplaceAll(tableName.Body, "\"@@table\"", "TableName"+b.ModelStructName) - // e.g. return "t_@@table" => return "t_user" - tableName.Body = strings.ReplaceAll(tableName.Body, "@@table", b.TableName) } b.ModelMethods = methods @@ -204,7 +200,16 @@ func (b *QueryStructMeta) ReviseDIYMethod() error { } return nil } +func (b *QueryStructMeta) pasreTableName(method *parser.Method) { + if method == nil || method.Body == "" || !strings.Contains(method.Body, "@@table") { + return + } + // e.g. return "@@table" => return TableNameUser + method.Body = strings.ReplaceAll(method.Body, "\"@@table\"", "TableName"+b.ModelStructName) + // e.g. return "t_@@table" => return "t_user" + method.Body = strings.ReplaceAll(method.Body, "@@table", b.TableName) +} func (b *QueryStructMeta) addMethodFromAddMethodOpt(methods ...interface{}) *QueryStructMeta { for _, method := range methods { modelMethods, err := parser.GetModelMethod(method)