Protobuf 和代码生成辅助方法

Protobuf

在元数据里面，说过 Protobuf 这种代码生成 的，无法利用 Tag 来指定列名
我们希望能够达成图二这种效果，而不是图一那种。

Protobuf 的局限性

Protobuf 虽然暴露了插件机制，但是插件并不能 修改生成的 Go 代码，插件只能自己额外生成一些 代码。
所以实际上不能利用 protobuf 的插件机制

修改 protobuf Go 插件

• clone 原本的 protobuf 仓库
• 修改 protobuf 仓库里的核心代码
• 安装修改后的 Go 插件
• 执行 protoc 命令

protobuf 必然会生成 json 的标签， 所以只需要生成找到 json 标签的位置，然后插入我们的代码。
这种是侵入式的修改方案，不过我们别无选择。 如果在公司可以维护一个自己定制过的 protobuf Go 插件 仓库

例子

实际上，最开始考虑过 Google 的 field option 扩展，但是依旧用不了，只能用这种方案。

代码生成辅助方法

Predicate 设计有一些缺陷，这些缺陷 是可以改进的：

• 生成字段名的常量
• 生成 Predicate

type User struct {
    Name     string
    Age      *int
    NickName *sql.NullString
    Picture  []byte
}

type UserDetail struct {
    Address string
}

**基本上就是 AST + 模板编程 **:

• AST 读取 Go 源文件内容，解析每个类型及其
• 生成 import 内容，并且将 orm 依赖添加进去
• 生成 const 内容
• 生成辅助方法

具体实现

在 gen/orm_gen/ 目录下

package {{ .Package }}

import (
"gitee.com/geektime-geekbang/geektime-go/orm"
{{range $idx, $import := .Imports }}
    {{$import}}
{{end -}}
)
{{- $ops := .Ops -}}
{{range $i, $type := .Types }}

    const (
    {{- range $j, $field := .Fields}}
        {{$type.Name }}{{$field.Name}} = "{{$field.Name}}"
    {{- end}}
    )

    {{range $j, $field := .Fields}}
        {{- range $k, $op := $ops}}
            func {{$type.Name }}{{$field.Name}}{{$op}}(val {{$field.Type}}) orm.Predicate {
            return orm.C("{{$field.Name}}").{{$op}}(val)
            }
        {{end}}
    {{- end}}
{{- end}}

type SingleFileEntryVisitor struct {
	file *fileVisitor
}

func (s *SingleFileEntryVisitor) Get() File {
	if s.file != nil {
		return s.file.Get()
	}
	return File{}
}

func (s *SingleFileEntryVisitor) Visit(node ast.Node) ast.Visitor {
	n, ok := node.(*ast.File)
	if ok {
		s.file = &fileVisitor{
			pkg: n.Name.String(),
		}
		return s.file
	}
	return s
}

type fileVisitor struct {
	pkg     string
	imports []string
	types   []*typeVisitor
}

func (f *fileVisitor) Visit(node ast.Node) ast.Visitor {
	switch n := node.(type) {
	case *ast.TypeSpec:
		res := &typeVisitor{
			name:   n.Name.String(),
			fields: make([]Field, 0),
		}
		if f.types == nil {
			f.types = make([]*typeVisitor, 0)
		}
		f.types = append(f.types, res)
		return res
	case *ast.ImportSpec:
		path := n.Path.Value
		if n.Name != nil && n.Name.String() != "" {
			path = n.Name.String() + " " + path
		}
		if f.imports == nil {
			f.imports = make([]string, 0)
		}
		f.imports = append(f.imports, path)
	}
	return f
}

func (f *fileVisitor) Get() File {
	types := make([]Type, 0, len(f.types))
	for _, t := range f.types {
		types = append(types, t.Get())
	}
	return File{
		Package: f.pkg,
		Imports: f.imports,
		Types:   types,
	}
}

type typeVisitor struct {
	name   string
	fields []Field
}

func (t *typeVisitor) Visit(node ast.Node) ast.Visitor {
	fd, ok := node.(*ast.Field)
	if ok {
		// 所以实际上我们在这里并没有处理 map, channel 之类的类型
		var typName string
		switch fdType := fd.Type.(type) {
		case *ast.Ident:
			typName = fdType.String()
		case *ast.StarExpr:
			switch expr := fdType.X.(type) {
			case *ast.Ident:
				typName = fmt.Sprintf("*%s", expr.String())
			case *ast.SelectorExpr:
				x := expr.X.(*ast.Ident).String()
				name := expr.Sel.String()
				typName = fmt.Sprintf("*%s.%s", x, name)
			}
		case *ast.SelectorExpr:
			x := fdType.X.(*ast.Ident).String()
			name := fdType.Sel.String()
			typName = fmt.Sprintf("%s.%s", x, name)
		case *ast.ArrayType:
			// 其它类型我们都不能处理处理，本来在 ORM 框架里面我们也没有支持
			switch ele := fdType.Elt.(type) {
			case *ast.Ident:
				typName = fmt.Sprintf("[]%s", ele)
			}
		}
		t.fields = append(t.fields, Field{
			Type: typName,
			Name: fd.Names[0].String(),
		})
		return nil
	}
	return t
}

func (t *typeVisitor) Get() Type {
	return Type{
		Name:   t.name,
		Fields: t.fields,
	}
}

type File struct {
	Package string
	Imports []string
	Types   []Type
}

type Type struct {
	Name   string
	Fields []Field
}

type Field struct {
	Name string
	Type string
}

func main() {
	// 用户必须输入一个 src，限制为文件
	// 然后我们会在同目录下生成代码
	src := os.Args[1]
	// Dir返回路径的除最后一个元素之外的所有元素，通常是路径的目录。
	dstDir := filepath.Dir(src)
	// Base返回路径的最后一个元素
	fileName := filepath.Base(src)
	// LastIndexByte返回s中c的最后一个实例的索引，如果s中不存在c，则返回-1。
	idx := strings.LastIndexByte(fileName, '.')
	dst := filepath.Join(dstDir, fileName[:idx]+"_gen.go")
	f, err := os.Create(dst)
	if err != nil {
		fmt.Println(err)
		return
	}
	err = gen(f, src)
	if err != nil {
		fmt.Println(err)
		return
	}
	fmt.Println("生成成功")

}

// Go 会读取 tpl.gohtml 里面的内容填充到变量 tpl 里面
//go:embed tpl.gohtml
var genOrm string

type OrmFile struct {
	File
	Ops []string
}

func gen(writer io.Writer, srcFile string) error {
	fset := token.NewFileSet()
	f, err := parser.ParseFile(fset, srcFile, nil, parser.ParseComments)
	if err != nil {
		return err
	}
	tv := &SingleFileEntryVisitor{}
	ast.Walk(tv, f)
	file := tv.Get()

	tpl := template.New("gen_orm")
	tpl, err = tpl.Parse(genOrm)
	if err != nil {
		return err
	}
	return tpl.Execute(writer, OrmFile{
		File: file,
		Ops:  []string{"LT", "GT", "EQ"},
	})
}

单元测试

func TestFileVisitor_Get(t *testing.T) {
	testCases := []struct {
		src  string
		want File
	}{
		{
			src: `
package orm_gen
import (
	"fmt"
    "database/sql"
) 

import (
	dri "database/sql/driver"
)
type (
	StructType struct {
		// Public is a field
		// @type string
		Public string
        Ptr *sql.NullString
		Struct sql.NullInt64
		Age *int8
		Slice []byte
	}
)
`,
			want: File{
				Package: "orm_gen",
				Imports: []string{`"fmt"`, `"database/sql"`, `dri "database/sql/driver"`},
				Types: []Type{
					{
						Name: "StructType",
						Fields: []Field{
							{
								Name: "Public",
								Type: "string",
							},
							{
								Name: "Ptr",
								Type: "*sql.NullString",
							},
							{
								Name: "Struct",
								Type: "sql.NullInt64",
							},
							{
								Name: "Age",
								Type: "*int8",
							},
							{
								Name: "Slice",
								Type: "[]byte",
							},
						},
					},
				},
			},
		},
	}
	for _, tc := range testCases {
		fset := token.NewFileSet()
		f, err := parser.ParseFile(fset, "src.go", tc.src, parser.ParseComments)
		if err != nil {
			t.Fatal(err)
		}
		tv := &SingleFileEntryVisitor{}
		ast.Walk(tv, f)
		file := tv.Get()
		assert.Equal(t, tc.want, file)
	}
}

func TestGen(t *testing.T) {
	bs := &bytes.Buffer{}
	err := gen(bs, "testdata/user.go")
	require.NoError(t, err)
	assert.Equal(t, `package testdata

import (
    "gitee.com/geektime-geekbang/geektime-go/orm"

    "database/sql"
)

const (
    UserName = "Name"
    UserAge = "Age"
    UserNickName = "NickName"
    UserPicture = "Picture"
)


func UserNameLT(val string) orm.Predicate {
    return orm.C("Name").LT(val)
}

func UserNameGT(val string) orm.Predicate {
    return orm.C("Name").GT(val)
}

func UserNameEQ(val string) orm.Predicate {
    return orm.C("Name").EQ(val)
}

func UserAgeLT(val *int) orm.Predicate {
    return orm.C("Age").LT(val)
}

func UserAgeGT(val *int) orm.Predicate {
    return orm.C("Age").GT(val)
}

func UserAgeEQ(val *int) orm.Predicate {
    return orm.C("Age").EQ(val)
}

func UserNickNameLT(val *sql.NullString) orm.Predicate {
    return orm.C("NickName").LT(val)
}

func UserNickNameGT(val *sql.NullString) orm.Predicate {
    return orm.C("NickName").GT(val)
}

func UserNickNameEQ(val *sql.NullString) orm.Predicate {
    return orm.C("NickName").EQ(val)
}

func UserPictureLT(val []byte) orm.Predicate {
    return orm.C("Picture").LT(val)
}

func UserPictureGT(val []byte) orm.Predicate {
    return orm.C("Picture").GT(val)
}

func UserPictureEQ(val []byte) orm.Predicate {
    return orm.C("Picture").EQ(val)
}


const (
    UserDetailAddress = "Address"
)


func UserDetailAddressLT(val string) orm.Predicate {
    return orm.C("Address").LT(val)
}

func UserDetailAddressGT(val string) orm.Predicate {
    return orm.C("Address").GT(val)
}

func UserDetailAddressEQ(val string) orm.Predicate {
    return orm.C("Address").EQ(val)
}
`, bs.String())
}

总结

• 怎么修改 protobuf 生成的代码？生成好的代码从实践上来说是不应该修改的，而很不幸的 是，protobuf 对应不同插件生成的目标代码，只能改插件的源码 ；
• 怎么为 protobuf 的字段添加额外的属性？可以通过 field option 来增加额外的属性，但是 这种新的属性需要你自己写代码解析 ；
• 代码生成的常用场景？一般来说，样板代码都可以考虑使用代码生成来替换掉，比如典型 的利用代码生成来生成数据库查询（如ENT），生成增删改查的代码，生成前端代码 ；