dao.go

package generator

import (
	"bufio"
	"bytes"
	"fmt"
	"go/format"
	"log"
	"os"
	"path/filepath"
	"sort"
	"strings"

	"github.com/pingcap/tidb/pkg/parser"
	"github.com/pingcap/tidb/pkg/parser/ast"
	"github.com/pingcap/tidb/pkg/parser/model"
	"github.com/pingcap/tidb/pkg/parser/opcode"
	"github.com/pingcap/tidb/pkg/parser/test_driver"
	"google.golang.org/protobuf/types/descriptorpb"

	"go.f110.dev/protoc-ddl/internal/schema"
)

const GoDAOGeneratorVersion = "v0.1"

type GoDAOGenerator struct{}

func (g GoDAOGenerator) Generate(buf *bytes.Buffer, fileOpt *descriptorpb.FileOptions, messages *schema.Messages) {
	src := newBuffer()

	entityPackageName := fileOpt.GetGoPackage()
	if strings.Contains(entityPackageName, ";") {
		s := strings.SplitN(entityPackageName, ";", 2)
		entityPackageName = s[0]
	}
	entityPackageAlias := filepath.Base(entityPackageName)

	src.WriteString("package dao\n")
	src.WriteString("import (\n")
	for _, v := range []string{"context", "database/sql", "fmt", "time", "strings"} {
		src.WriteString("\"" + v + "\"\n")
	}
	src.WriteRune('\n')
	src.WriteString(fmt.Sprintf("\"%s\"\n", entityPackageName))
	src.WriteString(")\n")

	// ListOption
	src.WriteString(`
	type ListOption func(opt *listOpt)

	func Limit(limit int) func(opt *listOpt) {
		return func(opt *listOpt) {
			opt.limit = limit
		}
	}

    func Desc(opt *listOpt) {
        opt.desc = true
    }

	func Sort(col string) func(*listOpt) {
		return func(opt *listOpt) {
			opt.sort = col
		}
	}

	type listOpt struct {
		sort  string
		limit int
        desc  bool
	}

	func newListOpt(opts ...ListOption) *listOpt {
		opt := &listOpt{}
		for _, v := range opts {
			v(opt)
		}
		return opt
	}

	type ExecOption func(opt *execOpt)

	func WithTx(tx *sql.Tx) ExecOption {
		return func(opt *execOpt) {
			opt.tx = tx
		}
	}
	
	type execOpt struct {
		tx *sql.Tx
	}

	func newExecOpt(opts ...ExecOption) *execOpt {
		opt := &execOpt{}
		for _, v := range opts {
			v(opt)
		}
		return opt
	}

	type execConn interface {
		ExecContext(ctx context.Context, query string, args ...interface{}) (sql.Result, error)
	}

`)

	messages.Each(func(m *schema.Message) {
		s := &GoDAOStruct{m: m, entityPackageName: entityPackageAlias}

		rels := make(map[string]struct{})
		for f := range m.Relations {
			if f.Virtual {
				continue
			}
			s := strings.Split(f.Type, ".")
			rels[s[len(s)-1]] = struct{}{}
		}
		relations := make([]string, 0)
		for v := range rels {
			relations = append(relations, v)
		}
		sort.Strings(relations)

		src.WriteString(fmt.Sprintf("type %s struct {\nconn *sql.DB\n\n", m.Descriptor.GetName()))
		for _, r := range relations {
			src.WriteString(fmt.Sprintf("%s *%s\n", schema.ToLowerCamel(r), schema.ToCamel(r)))
		}
		src.Write("}")
		src.LineBreak()
		src.Writef("type %sInterface interface {", m.Descriptor.GetName())
		src.WriteInterface(s.Tx(nil))
		src.WriteInterface(s.PrimaryKeySelect(nil))
		if len(m.PrimaryKeys) == 1 {
			src.WriteInterface(s.PrimaryKeyMultiSelect(nil))
		}
		src.WriteInterface(s.Select(nil)...)
		src.WriteInterface(s.Create(nil), s.Update(nil), s.Delete(nil))
		src.Write("}")
		src.LineBreak()
		src.Writef("var _ %sInterface = &%s{}", m.Descriptor.GetName(), m.Descriptor.GetName())
		src.LineBreak()
		src.WriteString(fmt.Sprintf("func New%s(conn *sql.DB) *%s {\n", m.Descriptor.GetName(), m.Descriptor.GetName()))
		src.WriteString(fmt.Sprintf("return &%s{\nconn: conn,\n", m.Descriptor.GetName()))
		for _, r := range relations {
			src.WriteString(fmt.Sprintf("%s: New%s(conn),\n", schema.ToLowerCamel(r), schema.ToCamel(r)))
		}
		src.WriteString("}\n")
		src.WriteString("}\n\n")

		src.WriteFunc(s.Tx(g.tx))
		src.LineBreak()

		src.WriteFunc(s.PrimaryKeySelect(g.primaryKeySelect))
		src.WriteFunc(s.PrimaryKeyMultiSelect(g.primaryKeyMultiSelect))
		src.WriteFunc(s.Select(g.selectRowQuery)...)
		src.WriteFunc(s.Create(g.create), s.Delete(g.delete), s.Update(g.update))
	})

	buf.WriteString("// Package dao contains the data access object.\n")
	buf.WriteString("// Generated by protoc-ddl.\n")
	buf.WriteString(fmt.Sprintf("// protoc-gen-dao: %s\n", GoDAOGeneratorVersion))
	b, err := format.Source(src.Bytes())
	if err != nil {
		r := bufio.NewScanner(strings.NewReader(src.String()))
		line := 1
		for r.Scan() {
			fmt.Fprintf(os.Stderr, "%d: %s\n", line, r.Text())
			line++
		}
		log.Print(err)
		return
	}
	buf.Write(b)
}

func (g GoDAOGenerator) tx(m *schema.Message, f *goFunc) string {
	src := newBuffer()

	src.Write("tx, err := d.conn.BeginTx(ctx, nil)")
	src.Write("if err != nil {")
	src.Write("return err")
	src.Write("}")
	src.Write("if err := fn(tx); err != nil {")
	src.Write("if rErr := tx.Rollback(); rErr != nil {")
	src.Write("return rErr")
	src.Write("}")
	src.Write("return err")
	src.Write("}")
	src.LineBreak()

	src.Write("err = tx.Commit()")
	src.Write("if err != nil {")
	src.Write("return err")
	src.Write("}")
	src.Write("return nil")

	return src.String()
}

func (g GoDAOGenerator) create(m *schema.Message, f *goFunc) string {
	src := newBuffer()

	valueArg := f.Args[1]
	src.Write(`execOpts := newExecOpt(opt...)
	var conn execConn
	if execOpts.tx != nil {
		conn = execOpts.tx
	} else {
		conn = d.conn
	}`)
	src.LineBreak()
	src.Write("res, err := conn.ExecContext(\nctx,")
	cols := make([]string, 0, m.Fields.Len())
	queryArgs := make([]string, 0, m.Fields.Len())
	m.Fields.Each(func(f *schema.Field) {
		if f.Sequence {
			return
		}
		if m.WithTimestamp && (f.Name == "created_at" || f.Name == "updated_at") {
			return
		}
		cols = append(cols, "`"+f.Name+"`")
		queryArgs = append(queryArgs, fmt.Sprintf("%s.%s", valueArg.Name, schema.ToCamel(f.Name)))
	})
	if m.WithTimestamp {
		cols = append(cols, "`created_at`")
		queryArgs = append(queryArgs, "time.Now()")
	}
	args := make([]string, len(cols))
	for i := 0; i < len(args); i++ {
		args[i] = "?"
	}

	src.Writef("\"INSERT INTO `%s` (%s) VALUES (%s)\",", m.TableName, strings.Join(cols, ", "), strings.Join(args, ", "))
	src.Write(strings.Join(queryArgs, ",") + ",")
	src.Write(")")
	src.Write("if err != nil {\nreturn nil, err\n}")
	src.LineBreak()
	src.Write("if n, err := res.RowsAffected(); err != nil {")
	src.Write("return nil, err")
	src.Write("} else if n == 0 {")
	src.Write("return nil, sql.ErrNoRows")
	src.Write("}")
	src.LineBreak()
	src.Writef("%s = %s.Copy()", valueArg.Name, valueArg.Name)
	if m.PrimaryKeys[0].Sequence {
		src.Write("insertedId, err := res.LastInsertId()")
		src.Write("if err != nil {")
		src.Write("return nil, err")
		src.Write("}")
		src.Writef("%s.Id = %s(insertedId)", valueArg.Name, GoDataTypeMap[m.PrimaryKeys[0].Type])
	}
	src.LineBreak()

	src.Writef("%s.ResetMark()", valueArg.Name)
	src.Writef("return %s, nil", valueArg.Name)

	return src.String()
}

func (g GoDAOGenerator) delete(m *schema.Message, _ *goFunc, where, whereArgs []string) string {
	src := newBuffer()

	src.WriteString(`execOpts := newExecOpt(opt...)
	var conn execConn
	if execOpts.tx != nil {
		conn = execOpts.tx
	} else {
		conn = d.conn
	}`)
	src.WriteString("\n\n")
	src.WriteString(fmt.Sprintf("res, err := conn.ExecContext(ctx, \"DELETE FROM `%s` WHERE %s\", %s)\n", m.TableName, strings.Join(where, " AND "), strings.Join(whereArgs, ",")))
	src.WriteString("if err != nil {\n")
	src.WriteString("return err\n")
	src.WriteString("}\n\n")
	src.WriteString("if n, err := res.RowsAffected(); err != nil {\n")
	src.WriteString("return err\n")
	src.WriteString("} else if n == 0 {\n")
	src.WriteString("return sql.ErrNoRows\n")
	src.WriteString("}\n\n")
	src.WriteString("return nil\n")

	return src.String()
}

func (g GoDAOGenerator) update(m *schema.Message, f *goFunc) string {
	src := newBuffer()
	valueArgs := f.Args[1]

	src.Writef("if !%s.IsChanged() {", valueArgs.Name)
	src.Write("return nil")
	src.Write("}")
	src.LineBreak()

	src.Write(`execOpts := newExecOpt(opt...)
	var conn execConn
	if execOpts.tx != nil {
		conn = execOpts.tx
	} else {
		conn = d.conn
	}`)
	src.LineBreak()

	src.Writef("changedColumn := %s.ChangedColumn()", valueArgs.Name)
	src.Write("cols := make([]string, len(changedColumn)+1)")
	src.Write("values := make([]interface{}, len(changedColumn)+1)")
	src.Write("for i := range changedColumn {")
	src.Write("cols[i] = \"`\" + changedColumn[i].Name + \"` = ?\"")
	src.Write("values[i] = changedColumn[i].Value")
	src.Write("}")
	if m.WithTimestamp {
		src.Write("cols[len(cols)-1] = \"`updated_at` = ?\"")
		src.Write("values[len(values)-1] = time.Now()")
	}
	src.LineBreak()

	where := make([]string, 0)
	whereArgs := make([]string, 0)
	for _, v := range m.PrimaryKeys {
		where = append(where, fmt.Sprintf("`%s` = ?", v.Name))
		whereArgs = append(whereArgs, valueArgs.Name+"."+schema.ToCamel(v.Name))
	}
	src.Writef("query := fmt.Sprintf(\"UPDATE `%s` SET %%s WHERE %s\", strings.Join(cols, \", \"))", m.TableName, strings.Join(where, " AND "))
	src.Write("res, err := conn.ExecContext(")
	src.Write("ctx,")
	src.Write("query,")
	src.Writef("append(values, %s)...,", strings.Join(whereArgs, ","))
	src.Write(")")
	src.Write("if err != nil {")
	src.Write("return err")
	src.Write("}")

	src.Write("if n, err := res.RowsAffected(); err != nil {")
	src.Write("return err")
	src.Write("} else if n == 0 {")
	src.Write("return sql.ErrNoRows")
	src.Write("}")
	src.LineBreak()

	src.Writef("%s.ResetMark()", valueArgs.Name)
	src.Write("return nil")

	return src.String()
}

func (g GoDAOGenerator) primaryKeySelect(entityName string, m *schema.Message, _, where, whereArgs []string) string {
	src := newBuffer()
	src.Buffer.WriteString("row := d.conn.QueryRowContext(ctx,")
	src.Writef("\"SELECT * FROM `%s` WHERE %s\", %s)\n", m.TableName, strings.Join(where, " AND "), strings.Join(whereArgs, ","))
	src.LineBreak()
	src.Writef("v := &%s.%s{}", entityName, m.Descriptor.GetName())
	cols := make([]string, 0, m.Fields.Len())
	m.Fields.Each(func(f *schema.Field) {
		cols = append(cols, "&v."+schema.ToCamel(f.Name))
	})
	src.Writef("if err := row.Scan(%s); err != nil {", strings.Join(cols, ","))
	src.Write("return nil, err\n}")
	src.LineBreak()

	g.selectChildObject(src, m)

	src.Write("v.ResetMark()")
	src.Write("return v, nil")

	return src.String()
}

func (g GoDAOGenerator) primaryKeyMultiSelect(entityName string, m *schema.Message, _, _, _ []string) string {
	src := newBuffer()
	src.Writef("inCause := strings.Repeat(\"?, \", len(%s))", schema.ToLowerCamel(m.PrimaryKeys[0].Name))
	src.Writef("args := make([]any, len(%s))", schema.ToLowerCamel(m.PrimaryKeys[0].Name))
	src.Writef("for i := 0;i < len(%s);i++ {\nargs[i] = %s[i]}", schema.ToLowerCamel(m.PrimaryKeys[0].Name), schema.ToLowerCamel(m.PrimaryKeys[0].Name))
	src.Writef("rows, err := d.conn.QueryContext(ctx,fmt.Sprintf(\"SELECT * FROM `%s` WHERE `%s` IN (%%s)\", inCause[:len(inCause)-2]),args...)", m.TableName, m.PrimaryKeys[0].Name)
	src.WriteString("if err != nil {\nreturn nil, err\n}\n")
	src.LineBreak()
	src.Writef("res := make([]*%s.%s,0,len(%s))", entityName, m.Descriptor.GetName(), schema.ToLowerCamel(m.PrimaryKeys[0].Name))
	src.WriteString("for rows.Next() {\n")
	src.WriteString(fmt.Sprintf("r := &%s.%s{}\n", entityName, m.Descriptor.GetName()))
	cols := make([]string, 0)
	fields := make([]*ast.SelectField, 0)
	for _, v := range m.Fields.List() {
		cols = append(cols, v.Name)
		fields = append(fields, &ast.SelectField{
			Expr: &ast.ColumnNameExpr{
				Name: &ast.ColumnName{Name: model.NewCIStr(v.Name)},
			},
		})
	}
	scanCols := make([]string, len(cols))
	for i, v := range cols {
		scanCols[i] = "&r." + schema.ToCamel(v)
	}
	src.WriteString(fmt.Sprintf("if err := rows.Scan(%s); err != nil {\n", strings.Join(scanCols, ",")))
	src.WriteString("return nil, err\n")
	src.WriteString("}\n")
	src.WriteString("res = append(res, r)")
	src.WriteString("}\n")
	src.LineBreak()
	if len(m.Relations) > 0 {
		g.selectChildObjects(src, m, entityName)
	}
	src.WriteString("return res, nil\n")

	return src.String()
}

func (g GoDAOGenerator) selectRowQuery(m *schema.Message, name string, stmt *ast.SelectStmt, comp []*schema.Field, cols, args []string, entityName string, single bool) string {
	if single {
		return g.selectSingleRowQuery(m, name, stmt, comp, args, entityName)
	} else {
		return g.selectMultipleRowQuery(m, name, stmt, cols, comp, args, entityName)
	}
}

func (g GoDAOGenerator) selectMultipleRowQuery(m *schema.Message, name string, stmt *ast.SelectStmt, selectCols []string, comp []*schema.Field, args []string, entityName string) string {
	src := newBuffer()
	primaryKeys := make([]string, 0)
	for _, v := range m.PrimaryKeys {
		primaryKeys = append(primaryKeys, "`"+v.Name+"`")
	}
	src.WriteString("listOpts := newListOpt(opt...)\n")
	f := NewQueryFormatter(m, stmt)
	buf := new(bytes.Buffer)
	f.Format(buf)
	src.WriteString(fmt.Sprintf("query := %q\n", buf.String()))
	src.WriteString("orderCol := \"`\"+listOpts.sort+\"`\"\n")
	src.WriteString("if listOpts.sort == \"\" {\n")
	src.WriteString(fmt.Sprintf("orderCol = \"%s\"\n", strings.Join(primaryKeys, ",")))
	src.WriteString("}\n")
	src.WriteString("orderDi := \"ASC\"\n")
	src.WriteString("if listOpts.desc {\n")
	src.WriteString("orderDi = \"DESC\"\n")
	src.WriteString("}\n")
	src.WriteString("query = query + fmt.Sprintf(\" ORDER BY %s %s\",orderCol,orderDi)\n")
	src.WriteString("if listOpts.limit > 0 {\n")
	src.WriteString("query = query + fmt.Sprintf(\" LIMIT %d\", listOpts.limit)\n")
	src.WriteString("}\n")
	src.WriteString("rows, err := d.conn.QueryContext(\nctx,\nquery,\n")
	for _, a := range comp {
		src.WriteString(schema.ToLowerCamel(a.Name) + ",\n")
	}
	src.WriteString(")\n")
	src.WriteString("if err != nil {\nreturn nil, err\n}\n")
	src.WriteRune('\n')

	// Object mapping
	src.WriteString(fmt.Sprintf("res := make([]*%s.%s, 0)\n", entityName, m.Descriptor.GetName()))
	src.WriteString("for rows.Next() {\n")
	src.WriteString(fmt.Sprintf("r := &%s.%s{}\n", entityName, m.Descriptor.GetName()))
	scanCols := make([]string, len(selectCols))
	for i, v := range selectCols {
		scanCols[i] = "&r." + schema.ToCamel(v)
	}
	src.WriteString(fmt.Sprintf("if err := rows.Scan(%s); err != nil {\n", strings.Join(scanCols, ",")))
	src.WriteString("return nil, err\n")
	src.WriteString("}\n")
	src.WriteString("r.ResetMark()\n")
	src.WriteString("res = append(res, r)\n")
	src.WriteString("}\n")

	if len(m.Relations) > 0 {
		g.selectChildObjects(src, m, entityName)
	}

	src.WriteRune('\n')
	src.WriteString("return res, nil\n")

	return src.String()
}

func (g GoDAOGenerator) selectSingleRowQuery(m *schema.Message, name string, stmt *ast.SelectStmt, comp []*schema.Field, args []string, entityName string) string {
	src := newBuffer()

	primaryKeys := make([]string, 0)
	for _, v := range m.PrimaryKeys {
		primaryKeys = append(primaryKeys, "`"+v.Name+"`")
	}
	f := NewQueryFormatter(m, stmt)
	buf := new(bytes.Buffer)
	f.Format(buf)
	src.WriteString(fmt.Sprintf("row := d.conn.QueryRowContext(\nctx,\n%q,\n", buf.String()))
	for _, a := range comp {
		src.WriteString(schema.ToLowerCamel(a.Name) + ",\n")
	}
	src.WriteString(")\n")

	src.WriteString(fmt.Sprintf("v := &%s.%s{}\n", entityName, m.Descriptor.GetName()))
	cols := make([]string, 0, m.Fields.Len())
	m.Fields.Each(func(f *schema.Field) {
		cols = append(cols, "&v."+schema.ToCamel(f.Name))
	})
	src.WriteString(fmt.Sprintf("if err := row.Scan(%s); err != nil {\n", strings.Join(cols, ",")))
	src.WriteString("return nil, err\n}\n\n")

	g.selectChildObject(src, m)

	src.WriteString("v.ResetMark()\n")
	src.WriteString("return v, nil\n")

	return src.String()
}

func (g GoDAOGenerator) selectChildObjects(src *Buffer, m *schema.Message, entityName string) {
	var bulkFetchFields, serialFetchFields []*schema.Field
	for f := range m.Relations {
		if f.Virtual {
			continue
		}
		rels := m.Relations[f]
		if len(rels) == 1 {
			bulkFetchFields = append(bulkFetchFields, f)
		} else {
			serialFetchFields = append(serialFetchFields, f)
		}
	}

	src.WriteString("if len(res) > 0 {\n")
	if len(bulkFetchFields) > 0 {
		for _, f := range bulkFetchFields {
			src.Writef("%sPrimaryKeys := make([]%s, len(res))", schema.ToLowerCamel(f.Name), GoDataTypeMap[m.Relations[f][0].Type])
		}
		src.WriteString("for i, v := range res {\n")
		for _, f := range bulkFetchFields {
			src.Writef("%sPrimaryKeys[i] = v.%s", schema.ToLowerCamel(f.Name), schema.ToCamel(m.Relations[f][0].Name))
		}
		src.WriteString("}\n")
		for _, f := range bulkFetchFields {
			s := strings.Split(f.Type, ".")
			src.Writef("%sData := make(map[%s]*%s.%s)", schema.ToLowerCamel(f.Name), GoDataTypeMap[m.Relations[f][0].Type], entityName, s[len(s)-1])
			src.WriteString("{\n")
			src.Writef("rels,_ := d.%s.SelectMulti(ctx,%sPrimaryKeys...)", schema.ToLowerCamel(s[len(s)-1]), schema.ToLowerCamel(f.Name))
			src.WriteString("for _, v := range rels {\n")
			src.Writef("%sData[v.%s] = v", schema.ToLowerCamel(f.Name), schema.ToCamel(m.Relations[f][0].Descriptor.GetName()))
			src.WriteString("}\n")
			src.WriteString("}\n")
		}
		src.WriteString("for _, v := range res {\n")
		for _, f := range bulkFetchFields {
			src.Writef("v.%s = %sData[v.%s]", schema.ToCamel(f.Name), schema.ToLowerCamel(f.Name), schema.ToCamel(m.Relations[f][0].Name))
		}
		src.WriteString("}\n")
	}

	if len(serialFetchFields) > 0 {
		src.LineBreak()
		src.WriteString("for _, v := range res {\n")
		for _, f := range serialFetchFields {
			g.selectChildObjectBySerial(src, m, f)
		}
		src.WriteString("}\n")
	}
	src.WriteString("}\n")
}

func (GoDAOGenerator) selectChildObjectBySerial(src *Buffer, m *schema.Message, f *schema.Field) {
	r := make([]string, 0, len(m.Relations[f]))
	check := make([]string, 0, len(m.Relations[f]))
	for _, v := range m.Relations[f] {
		if f.Null {
			r = append(r, "*v."+schema.ToCamel(v.Name))
		} else {
			r = append(r, "v."+schema.ToCamel(v.Name))
		}
		check = append(check, "v."+schema.ToCamel(v.Name)+" != nil")
	}
	src.Write("{")
	if f.Null {
		src.Writef("if %s {", strings.Join(check, " && "))
	}
	s := strings.Split(f.Type, ".")
	src.Writef("if rel, _ := d.%s.Select(ctx, %s); rel != nil {", schema.ToLowerCamel(s[len(s)-1]), strings.Join(r, ","))
	src.Writef("v.%s = rel", schema.ToCamel(f.Name))
	src.Write("}")
	if f.Null {
		src.Write("}")
	}
	src.Write("}")
}

func (g GoDAOGenerator) selectChildObject(src *Buffer, m *schema.Message) {
	if len(m.Relations) == 0 {
		return
	}

	relFields := make([]*schema.Field, 0)
	for f := range m.Relations {
		if f.Virtual {
			continue
		}
		relFields = append(relFields, f)
	}
	sort.Slice(relFields, func(i, j int) bool {
		return relFields[i].Name < relFields[j].Name
	})

	for _, f := range relFields {
		g.selectChildObjectBySerial(src, m, f)
	}
	src.WriteRune('\n')
}

type goFunc struct {
	Receiver *field
	Name     string
	Args     fieldList
	Body     string
	Returns  fieldList
}

func (f *goFunc) String() string {
	src := newBuffer()
	src.WriteString("func")
	if f.Receiver != nil {
		src.WriteString(" (" + f.Receiver.String() + ")")
	}
	if f.Name != "" {
		src.WriteString(" " + f.Name)
	} else {
		src.WriteString(" ")
	}
	switch len(f.Args) {
	case 0:
		src.WriteString("()")
	case 1:
		src.WriteRune('(')
		src.WriteString(f.Args.String())
		src.WriteRune(')')
	default:
		src.WriteString(f.Args.String())
	}
	src.WriteRune(' ')
	src.WriteString(f.Returns.String())
	if f.Body != "" {
		src.WriteString(" {")
		src.WriteString(f.Body)
		src.WriteString("}")
	}

	return src.String()
}

func (f *goFunc) Interface() string {
	src := newBuffer()
	src.Buffer.WriteString(f.Name)
	switch len(f.Args) {
	case 0:
		src.Buffer.WriteString("()")
	case 1:
		src.Buffer.WriteRune('(')
		src.Buffer.WriteString(f.Args.String())
		src.Buffer.WriteRune(')')
	default:
		src.Buffer.WriteString(f.Args.String())
	}
	src.Buffer.WriteString(f.Returns.String())

	return src.String()
}

type field struct {
	Name     string
	Type     string
	Slice    bool
	Variable bool
	Pointer  bool
}

func (n *field) Copy() *field {
	v := &field{}
	*v = *n

	return v
}

func (n *field) String() string {
	v := ""
	if n.Name != "" {
		v += n.Name + " "
	}
	if n.Slice {
		v += "[]"
	}
	if n.Variable {
		v += "..."
	}
	if n.Pointer {
		v += "*"
	}

	v += n.Type
	return v
}

type fieldList []*field

func (n fieldList) String() string {
	s := make([]string, 0)
	for _, v := range n {
		s = append(s, v.String())
	}

	if len(s) == 1 {
		return s[0]
	}
	return "(" + strings.Join(s, ",") + ")"
}

type GoDAOStruct struct {
	m                 *schema.Message
	entityPackageName string
	daoPath           string
}

func (s *GoDAOStruct) PrimaryKeySelect(
	body func(entityPackageName string, m *schema.Message, args, where, whereArgs []string) string,
) *goFunc {
	funcArgs := fieldList{{Name: "ctx", Type: "context.Context"}}
	args := make([]string, 0)
	where := make([]string, 0)
	whereArgs := make([]string, 0)
	for _, v := range s.m.PrimaryKeys {
		funcArgs = append(funcArgs, &field{Name: schema.ToLowerCamel(v.Name), Type: GoDataTypeMap[v.Type]})

		args = append(args, fmt.Sprintf("%s %s", schema.ToLowerCamel(v.Name), GoDataTypeMap[v.Type]))
		where = append(where, fmt.Sprintf("`%s` = ?", v.Name))
		whereArgs = append(whereArgs, schema.ToLowerCamel(v.Name))
	}

	f := &goFunc{
		Name:     "Select",
		Receiver: &field{Name: "d", Type: s.m.Descriptor.GetName(), Pointer: true},
		Args:     funcArgs,
		Returns: []*field{
			{Type: s.entityPackageName + "." + s.m.Descriptor.GetName(), Pointer: true},
			{Type: "error"},
		},
	}
	if body != nil {
		f.Body = body(s.entityPackageName, s.m, args, where, whereArgs)
	}
	return f
}

func (s *GoDAOStruct) PrimaryKeyMultiSelect(
	body func(entityPackageName string, m *schema.Message, args, where, whereArgs []string) string,
) *goFunc {
	if len(s.m.PrimaryKeys) != 1 {
		return nil
	}
	funcArgs := fieldList{{Name: "ctx", Type: "context.Context"}}
	funcArgs = append(funcArgs, &field{Name: schema.ToLowerCamel(s.m.PrimaryKeys[0].Name), Type: "..." + GoDataTypeMap[s.m.PrimaryKeys[0].Type]})

	f := &goFunc{
		Name:     "SelectMulti",
		Receiver: &field{Name: "d", Type: s.m.Descriptor.GetName(), Pointer: true},
		Args:     funcArgs,
		Returns: []*field{
			{Slice: true, Type: s.entityPackageName + "." + s.m.Descriptor.GetName(), Pointer: true},
			{Type: "error"},
		},
	}
	if body != nil {
		f.Body = body(s.entityPackageName, s.m, nil, nil, nil)
	}
	return f
}

func (s *GoDAOStruct) Select(
	body func(m *schema.Message, name string, stmt *ast.SelectStmt, comp []*schema.Field, cols, args []string, entityName string, single bool) string,
) []*goFunc {
	funcs := make([]*goFunc, 0)

	p := parser.New()
	for _, q := range s.m.SelectQueries {
		stmts, _, err := p.Parse(q.Query, "", "")
		if err != nil {
			log.Printf("Failed parse query %s: %v", q.Query, err)
			continue
		}

		stmt := stmts[0]
		switch v := stmt.(type) {
		case *ast.SelectStmt:
			f := s.selectQuery(v, q.Name, body)
			if f != nil {
				funcs = append(funcs, f)
			}
		default:
			log.Printf("%q is not supported: %s", v, q)
		}
	}

	return funcs
}

func (s *GoDAOStruct) Create(body func(m *schema.Message, f *goFunc) string) *goFunc {
	execOptionType := "ExecOption"
	if s.daoPath != "" {
		execOptionType = s.daoPath + "." + execOptionType
	}
	funcArgs := fieldList{
		{Name: "ctx", Type: "context.Context"},
		{Name: schema.ToLowerCamel(s.m.Descriptor.GetName()), Type: s.entityPackageName + "." + s.m.Descriptor.GetName(), Pointer: true},
		{Name: "opt", Type: execOptionType, Variable: true},
	}

	f := &goFunc{
		Name:     "Create",
		Receiver: &field{Name: "d", Type: s.m.Descriptor.GetName(), Pointer: true},
		Args:     funcArgs,
		Returns: fieldList{
			{Type: s.entityPackageName + "." + s.m.Descriptor.GetName(), Pointer: true},
			{Type: "error"},
		},
	}
	if body != nil {
		f.Body = body(s.m, f)
	}

	return f
}

func (s *GoDAOStruct) Delete(body func(m *schema.Message, f *goFunc, where, whereArgs []string) string) *goFunc {
	funcArgs := fieldList{{Name: "ctx", Type: "context.Context"}}
	where := make([]string, 0)
	whereArgs := make([]string, 0)
	for _, v := range s.m.PrimaryKeys {
		funcArgs = append(funcArgs, &field{Name: schema.ToLowerCamel(v.Name), Type: GoDataTypeMap[v.Type]})
		where = append(where, fmt.Sprintf("`%s` = ?", v.Name))
		whereArgs = append(whereArgs, schema.ToLowerCamel(v.Name))
	}
	execOptionType := "ExecOption"
	if s.daoPath != "" {
		execOptionType = s.daoPath + "." + execOptionType
	}
	funcArgs = append(funcArgs, &field{Name: "opt", Type: execOptionType, Variable: true})

	f := &goFunc{
		Name:     "Delete",
		Receiver: &field{Name: "d", Type: s.m.Descriptor.GetName(), Pointer: true},
		Args:     funcArgs,
		Returns:  fieldList{{Type: "error"}},
	}
	if body != nil {
		f.Body = body(s.m, f, where, whereArgs)
	}

	return f
}

func (s *GoDAOStruct) Update(body func(m *schema.Message, f *goFunc) string) *goFunc {
	execOptionType := "ExecOption"
	if s.daoPath != "" {
		execOptionType = s.daoPath + "." + execOptionType
	}
	funcArgs := fieldList{
		{Name: "ctx", Type: "context.Context"},
		{Name: schema.ToLowerCamel(s.m.Descriptor.GetName()), Type: s.entityPackageName + "." + s.m.Descriptor.GetName(), Pointer: true},
		{Name: "opt", Type: execOptionType, Variable: true},
	}

	f := &goFunc{
		Name:     "Update",
		Receiver: &field{Name: "d", Type: s.m.Descriptor.GetName(), Pointer: true},
		Args:     funcArgs,
		Returns:  fieldList{{Type: "error"}},
	}
	if body != nil {
		f.Body = body(s.m, f)
	}

	return f
}

func (s *GoDAOStruct) Tx(body func(m *schema.Message, f *goFunc) string) *goFunc {
	f := &goFunc{
		Name:     "Tx",
		Receiver: &field{Name: "d", Type: s.m.Descriptor.GetName(), Pointer: true},
		Args: fieldList{
			{Name: "ctx", Type: "context.Context"},
			{
				Name: "fn",
				Type: fmt.Sprint(&goFunc{
					Args:    fieldList{{Name: "tx", Type: "sql.Tx", Pointer: true}},
					Returns: fieldList{{Type: "error"}},
				}),
			},
		},
		Returns: fieldList{{Type: "error"}},
	}
	if body != nil {
		f.Body = body(s.m, f)
	}

	return f
}

func (s *GoDAOStruct) selectQuery(
	stmt *ast.SelectStmt,
	name string,
	body func(m *schema.Message, name string, stmt *ast.SelectStmt, comp []*schema.Field, cols, args []string, entityName string, single bool) string,
) *goFunc {
	if stmt.From.TableRefs.Right != nil {
		log.Printf("Multiple tables is not supported")
		return nil
	}

	allColumn := false
	var cols []string
	for _, c := range stmt.Fields.Fields {
		if c.WildCard != nil {
			allColumn = true
			cols = nil
			break
		}

		col, ok := c.Expr.(*ast.ColumnNameExpr)
		if !ok {
			log.Printf("%v is %T", c, c)
			continue
		}
		cols = append(cols, col.Name.String())
	}
	if allColumn {
		cols = make([]string, 0)
		fields := make([]*ast.SelectField, 0)
		for _, v := range s.m.Fields.List() {
			cols = append(cols, v.Name)
			fields = append(fields, &ast.SelectField{
				Expr: &ast.ColumnNameExpr{
					Name: &ast.ColumnName{Name: model.NewCIStr(v.Name)},
				},
			})
		}
		stmt.Fields.Fields = fields
	}

	var comp []*schema.Field
	if stmt.Where != nil {
		comp = s.findArgs(s.m.Fields.List(), stmt.Where)
	}

	isReturningSingleRow := s.m.IsReturningSingleRow(comp...)
	funcArgs := make([]*field, 0)
	args := make([]string, len(comp))
	for i := range comp {
		funcArgs = append(funcArgs, &field{Name: schema.ToLowerCamel(comp[i].Name), Type: GoDataTypeMap[comp[i].Type]})
		args[i] = fmt.Sprintf("%s %s", schema.ToLowerCamel(comp[i].Name), GoDataTypeMap[comp[i].Type])
	}
	if !isReturningSingleRow {
		args = append(args, "opt ...ListOption")
	}

	if isReturningSingleRow {
		f := &goFunc{
			Name:     "Select" + name,
			Receiver: &field{Name: "d", Type: s.m.Descriptor.GetName(), Pointer: true},
			Args:     append([]*field{{Name: "ctx", Type: "context.Context"}}, funcArgs...),
			Returns: []*field{
				{Type: s.entityPackageName + "." + s.m.Descriptor.GetName(), Pointer: true},
				{Type: "error"},
			},
		}
		if body != nil {
			f.Body = body(s.m, "Select"+name, stmt, comp, cols, args, s.entityPackageName, true)
		}
		return f
	} else {
		funcArgs = append([]*field{{Name: "ctx", Type: "context.Context"}}, funcArgs...)
		listOptionType := "ListOption"
		if s.daoPath != "" {
			listOptionType = s.daoPath + "." + listOptionType
		}
		funcArgs = append(funcArgs, &field{Name: "opt", Variable: true, Type: listOptionType})

		f := &goFunc{
			Name:     "List" + name,
			Receiver: &field{Name: "d", Type: s.m.Descriptor.GetName(), Pointer: true},
			Args:     funcArgs,
			Returns: []*field{
				{Slice: true, Type: s.entityPackageName + "." + s.m.Descriptor.GetName(), Pointer: true},
				{Type: "error"},
			},
		}
		if body != nil {
			f.Body = body(s.m, "List"+name, stmt, comp, cols, args, s.entityPackageName, false)
		}
		return f
	}
}

func (s *GoDAOStruct) findArgs(fields []*schema.Field, stmt ast.ExprNode) []*schema.Field {
	fieldMap := make(map[string]*schema.Field)
	for _, v := range fields {
		fieldMap[v.Name] = v
	}

	res := s.findArgFieldFromExprIfExist(fieldMap, stmt)
	if len(res) == 0 {
		return nil
	}

	return res
}

func (s *GoDAOStruct) findArgFieldFromExprIfExist(fields map[string]*schema.Field, e ast.ExprNode) []*schema.Field {
	res := make([]*schema.Field, 0)
	stmt, ok := e.(*ast.BinaryOperationExpr)
	if !ok {
		return res
	}

	switch stmt.Op {
	case opcode.EQ, opcode.GT:
		f := s.findArgFieldFromComparisonExprIfExist(fields, stmt)
		if len(f) > 0 {
			res = append(res, f...)
		}
	case opcode.LogicAnd:
		f := s.findArgFieldFromAndExprIfExist(fields, stmt)
		if len(f) > 0 {
			res = append(res, f...)
		}
	case opcode.LogicOr:
		f := s.findArgFieldFromOrExprIfExist(fields, stmt)
		if len(f) > 0 {
			res = append(res, f...)
		}
	default:
		log.Printf("%T", stmt)
	}

	return res
}

func (*GoDAOStruct) findArgFieldFromComparisonExprIfExist(fields map[string]*schema.Field, stmt ast.ExprNode) []*schema.Field {
	res := make([]*schema.Field, 0)
	switch v := stmt.(type) {
	case *ast.BinaryOperationExpr:
		if left, ok := v.L.(*ast.ColumnNameExpr); ok {
			switch v.R.(type) {
			case *test_driver.ParamMarkerExpr:
				if f, ok := fields[left.Name.Name.String()]; ok {
					res = append(res, f)
				}
			}
		}
	}

	return res
}

func (s *GoDAOStruct) findArgFieldFromAndExprIfExist(fields map[string]*schema.Field, stmt *ast.BinaryOperationExpr) []*schema.Field {
	res := make([]*schema.Field, 0)

	for _, st := range []ast.ExprNode{stmt.L, stmt.R} {
		f := s.findArgFieldFromExprIfExist(fields, st)
		if len(f) > 0 {
			res = append(res, f...)
		}
	}

	return res
}

func (s *GoDAOStruct) findArgFieldFromOrExprIfExist(fields map[string]*schema.Field, stmt *ast.BinaryOperationExpr) []*schema.Field {
	res := make([]*schema.Field, 0)

	for _, st := range []ast.ExprNode{stmt.L, stmt.R} {
		f := s.findArgFieldFromExprIfExist(fields, st)
		if len(f) > 0 {
			res = append(res, f...)
		}
	}

	return res
}