From 66e8a72bf1b04a6b256c94708da68ddab498a5aa Mon Sep 17 00:00:00 2001
From: Jinzhu <wosmvp@gmail.com>
Date: Mon, 23 Nov 2020 11:24:07 +0800
Subject: [PATCH] Support NameReplace for NamingStrategy, close #3779

---
 schema/naming.go      | 21 +++++++++++++--------
 schema/naming_test.go | 12 ++++++++++--
 2 files changed, 23 insertions(+), 10 deletions(-)

diff --git a/schema/naming.go b/schema/naming.go
index e3b2104a..63296967 100644
--- a/schema/naming.go
+++ b/schema/naming.go
@@ -24,19 +24,20 @@ type Namer interface {
 type NamingStrategy struct {
 	TablePrefix   string
 	SingularTable bool
+	NameReplacer  *strings.Replacer
 }
 
 // TableName convert string to table name
 func (ns NamingStrategy) TableName(str string) string {
 	if ns.SingularTable {
-		return ns.TablePrefix + toDBName(str)
+		return ns.TablePrefix + ns.toDBName(str)
 	}
-	return ns.TablePrefix + inflection.Plural(toDBName(str))
+	return ns.TablePrefix + inflection.Plural(ns.toDBName(str))
 }
 
 // ColumnName convert string to column name
 func (ns NamingStrategy) ColumnName(table, column string) string {
-	return toDBName(column)
+	return ns.toDBName(column)
 }
 
 // JoinTableName convert string to join table name
@@ -46,14 +47,14 @@ func (ns NamingStrategy) JoinTableName(str string) string {
 	}
 
 	if ns.SingularTable {
-		return ns.TablePrefix + toDBName(str)
+		return ns.TablePrefix + ns.toDBName(str)
 	}
-	return ns.TablePrefix + inflection.Plural(toDBName(str))
+	return ns.TablePrefix + inflection.Plural(ns.toDBName(str))
 }
 
 // RelationshipFKName generate fk name for relation
 func (ns NamingStrategy) RelationshipFKName(rel Relationship) string {
-	return strings.Replace(fmt.Sprintf("fk_%s_%s", rel.Schema.Table, toDBName(rel.Name)), ".", "_", -1)
+	return strings.Replace(fmt.Sprintf("fk_%s_%s", rel.Schema.Table, ns.toDBName(rel.Name)), ".", "_", -1)
 }
 
 // CheckerName generate checker name
@@ -63,7 +64,7 @@ func (ns NamingStrategy) CheckerName(table, column string) string {
 
 // IndexName generate index name
 func (ns NamingStrategy) IndexName(table, column string) string {
-	idxName := fmt.Sprintf("idx_%v_%v", table, toDBName(column))
+	idxName := fmt.Sprintf("idx_%v_%v", table, ns.toDBName(column))
 	idxName = strings.Replace(idxName, ".", "_", -1)
 
 	if utf8.RuneCountInString(idxName) > 64 {
@@ -91,13 +92,17 @@ func init() {
 	commonInitialismsReplacer = strings.NewReplacer(commonInitialismsForReplacer...)
 }
 
-func toDBName(name string) string {
+func (ns NamingStrategy) toDBName(name string) string {
 	if name == "" {
 		return ""
 	} else if v, ok := smap.Load(name); ok {
 		return v.(string)
 	}
 
+	if ns.NameReplacer != nil {
+		name = ns.NameReplacer.Replace(name)
+	}
+
 	var (
 		value                          = commonInitialismsReplacer.Replace(name)
 		buf                            strings.Builder
diff --git a/schema/naming_test.go b/schema/naming_test.go
index 26b0dcf6..b7a32160 100644
--- a/schema/naming_test.go
+++ b/schema/naming_test.go
@@ -1,6 +1,7 @@
 package schema
 
 import (
+	"strings"
 	"testing"
 )
 
@@ -26,9 +27,10 @@ func TestToDBName(t *testing.T) {
 		"ThisIsActuallyATestSoWeMayBeAbleToUseThisCodeInGormPackageAlsoIdCanBeUsedAtTheEndAsID": "this_is_actually_a_test_so_we_may_be_able_to_use_this_code_in_gorm_package_also_id_can_be_used_at_the_end_as_id",
 	}
 
+	ns := NamingStrategy{}
 	for key, value := range maps {
-		if toDBName(key) != value {
-			t.Errorf("%v toName should equal %v, but got %v", key, value, toDBName(key))
+		if ns.toDBName(key) != value {
+			t.Errorf("%v toName should equal %v, but got %v", key, value, ns.toDBName(key))
 		}
 	}
 }
@@ -37,6 +39,7 @@ func TestNamingStrategy(t *testing.T) {
 	var ns = NamingStrategy{
 		TablePrefix:   "public.",
 		SingularTable: true,
+		NameReplacer:  strings.NewReplacer("CID", "Cid"),
 	}
 	idxName := ns.IndexName("public.table", "name")
 
@@ -63,4 +66,9 @@ func TestNamingStrategy(t *testing.T) {
 	if tableName != "public.company" {
 		t.Errorf("invalid table name generated, got %v", tableName)
 	}
+
+	columdName := ns.ColumnName("", "NameCID")
+	if columdName != "name_cid" {
+		t.Errorf("invalid column name generated, got %v", columdName)
+	}
 }