From a31c34b70efff1ef3550b3e57a440cf7d5cc082d Mon Sep 17 00:00:00 2001
From: qqxhb <1252905006@qq.com>
Date: Fri, 22 Nov 2024 09:06:02 +0800
Subject: [PATCH] feat: parse table name for all method

---
 internal/generate/query.go | 15 ++++++++++-----
 1 file changed, 10 insertions(+), 5 deletions(-)

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)