GORM GEN 生成代码如何自定义方法和表名

GORM GEN: 是安全&友好的go orm 框架，基于gorm，支持多种数据库，mysql、sqlserver、postgres以及clickhouse等，通过代码生成的方式，生成安全&友好的orm代码，同时也支持类似java mybatis的 动态sql代码生成。更多详细介绍和教程请移步gorm.io/gen，本文主要介绍如何给生成的model添加自定义方法和自定义表名。

1. 自定义方法

相信你在日常使用gorm gen开发的过程中也有这种场景：通过gen的g.GenerateModel("users")等方式，从数据库或者sql文件生成了对应的gorm model，已经满足了基本的使用；某些时候希望model上能有一些通用的方法，比如IsEmpty()，判断model是否为空；如果生成之后在修改model，添加自定义方法，下次生成的时候会被覆盖，因此gen提供了WithMethodOption，允许你在生成代码的时候就添加自定义方法，具体使用如下：

1. 写一个通用方法的模版
   这个模版定义了两个方法，你可以把他放到所有需要的model上

package method
type CommonMethod struct {
    ID   int32
    Name *string
}

func (m *CommonMethod) IsEmpty() bool {
    if m == nil {
        return true
    }
    return m.ID == 0
}

func (m *CommonMethod) GetName() string {
    if m == nil || m.Name == nil {
        return ""
    }
    return *m.Name
}

2. 生成代码指定需要的自定义方法，可以是一个也可以是全部

• people表对应的model添加一个自定义方法IsEmpty
• user表对应的model，添加CommonMethod上的所有方法，如示例中的IsEmpty和GetName

// people表对应的model添加一个自定义方法`IsEmpty`
g.GenerateModel("people", gen.WithMethod(method.CommonMethod{}.IsEmpty))

// user表对应的model，添加CommonMethod上的所有方法，如示例中的`IsEmpty`和`GetName`
g.GenerateModel("user", gen.WithMethod(method.CommonMethod))

3. 生成之后大代码大概就是这样，截取部分

const TableNamePerson = "people"

// Person mapped from table <people>
type Person struct {
	ID             int64          `gorm:"column:id;primaryKey;autoIncrement:true" json:"id"`
	Name           string         `gorm:"column:name" json:"name"`
	Age            int32          `gorm:"column:age" json:"age"`
	//..........
}

func (p *Person) IsEmpty() bool {
  if p == nil {
    return true
  }
  return p.ID <= 0
}

// TableName Person's table name
func (*Person) TableName() string {
	return TableNamePerson
}

2. 自定义表名称

通过上面的生成代码示例，你可能也发现了，生成mode之后默认的表名称其实就是数据库的表名，但有的时候可能不满足我们的诉求。比如生成代码一般是连接的测试库，可能和线上库的表不一样，比如多个前缀或者其他的。因此gorm gen支持自定义表名称。方法很简单了，就是利用上面的自定义方法WithMethod，只不过方法的名称是TableName。

Gorm 只能识别两种签名的TableName方法，第一种是示例代码的，没有入参，TableName() string；第二种是带有gorm NamingStrategy为入参，这样就可以实现，解析model表名时使用gorm的NamingStrategy，TableName(namer schema.Namer) string。对以上两种gen都有实现，默认是第一种不需要额外的操作，就会生成。如果想用的是第二种，则需要手动指定下，方式也很简答：gen.WithMethod(gen.DefaultMethodTableWithNamer)。

gen.WithMethod 可以只用在某个表上，e.g. g.GenerateModel("user", gen.WithMethod(xxx))；也可以全局配置，所有表生效，e.g. conf.WithOpts(gen.WithMethod(xxx))

当然你也可以自定义一个实现，比如：

@@table，会被替换成真实的表名称

// TableName 
func (m CommonMethod) TableName() strng {
	if env.IsTest(){
		return "test_@@table"
	}
    return "@@table"
}

或者

// TableName table name with gorm NamingStrategy
func (m CommonMethod) TableName(namer schema.Namer) string {
    if namer == nil {
        return "@@table"
    }
    return namer.TableName("@@table")
}

！！！ ⚠️注意： 满足上面的两种TableName定义的，都只是第一次解析model生效，也就相当于同一个gorm db只会调用一次。
如果是想要实现动态表名，比如分片表，可能没次需要的表名都是动态的，这种你可以自定义一个其他方法，当然也可以命名为TableName，重要的是每次执行前，需要手动调用一次。
e.g.

func (m *CommonMethod) TableName(ID int64) string {
	shard := 10
	return fmt.Sprintf("@@table%d", ID%shard)
}

或者其他名称

func (m *CommonMethod) ShardTableName() string {
	shard := 10
	return "@@table" + fmt.Sprintf("%d", m.ID%shard)
}

使用的时候，手动调用。e.g.

    user := &model.User{xxx}
    u := query.User.Table(user.TableName(1))
    u.WithContext(ctx).Create(user)

或者

    user := &model.User{ID:1}
    u := query.User.Table(user.ShardTableName())
    u.WithContext(ctx).Create(user)

OK，本文就介绍到这里了，有任何问题，欢迎评论区讨论。