やっつけ

pg_test.go

package pg_test

import (
	"fmt"
	"reflect"
	"time"

	"gopkg.in/pg.v4"
	"gopkg.in/pg.v4/orm"
	"gopkg.in/pg.v4/types"
)

type User struct {
	Id        int64
	Name      string
	Emails    []string
	CreatedAt time.Time
}

func (u User) String() string {
	return fmt.Sprintf("User<%d %s %v>", u.Id, u.Name, u.Emails)
}

type Story struct {
	Id     int64
	Title  string
	UserId int64
	User   *User
}

func (s Story) String() string {
	return fmt.Sprintf("Story<%d %s %s>", s.Id, s.Title, s.User)
}

func createSchema(db *pg.DB) error {
	queries := []string{
		`CREATE TEMP TABLE users (id serial, name text, emails jsonb, created_at timestamp)`,
		`CREATE TEMP TABLE stories (id serial, title text, user_id bigint)`,
	}
	for _, q := range queries {
		_, err := db.Exec(q)
		if err != nil {
			return err
		}
	}
	return nil
}

func MyCreate(db *pg.DB, v interface{}) error {
	model, err := orm.NewTableModel(v)
	if err != nil {
		return err
	}
	_, err = db.Query(model, insert{TableModel: model})
	return err
}

type insert struct {
	*orm.TableModel
}

func (ins insert) AppendQuery(b []byte, params ...interface{}) ([]byte, error) {
	strct := ins.Value()

	b = append(b, "INSERT INTO "...)
	b = types.AppendField(b, ins.Table.Name, true)

	b = append(b, " ("...)
	for i, field := range ins.Table.Fields {
		if field.Has(orm.PrimaryKeyFlag) && field.IsEmpty(strct) {
			continue
		}
		b = types.AppendField(b, field.SQLName, true)
		if i != len(ins.Table.Fields)-1 {
			b = append(b, ", "...)
		}
	}

	b = append(b, ") VALUES ("...)

	for i, field := range ins.Table.Fields {
		if field.Has(orm.PrimaryKeyFlag) && field.IsEmpty(strct) {
			continue
		}
		b = field.AppendValue(b, strct, true)
		if i != len(ins.Table.Fields)-1 {
			b = append(b, ", "...)
		}
	}
	b = append(b, ")"...)

	var fs []*orm.Field
	fs = append(fs, ins.Table.PKs...)
	for _, field := range ins.Table.Fields {
		if field.SQLName == "created_at" || field.SQLName == "updated_at" {
			fs = append(fs, field)
		}
	}
	b = appendReturning(b, strct, fs)
	return b, nil
}

func appendReturning(b []byte, v reflect.Value, fields []*orm.Field) []byte {
	var hasReturning bool
	for i, f := range fields {
		if !f.IsEmpty(v) {
			//continue
		}
		if !hasReturning {
			b = append(b, " RETURNING "...)
			hasReturning = true
		}
		b = types.AppendField(b, f.SQLName, true)
		if i != len(fields)-1 {
			b = append(b, ", "...)
		}
	}
	return b
}

func ExampleDB_Query() {
	db := pg.Connect(&pg.Options{
		User:     "relax",
		Password: "relax",
		Database: "relaxdb-test",
	})

	err := createSchema(db)
	if err != nil {
		panic(err)
	}

	user1 := &User{
		Name:   "admin",
		Emails: []string{"admin1@admin", "admin2@admin"},
	}
	err = MyCreate(db, user1)
	if err != nil {
		panic(err)
	}

	err = db.Create(&User{
		Name:   "root",
		Emails: []string{"root1@root", "root2@root"},
	})
	if err != nil {
		panic(err)
	}

	story1 := &Story{
		Title:  "Cool story",
		UserId: user1.Id,
	}
	err = db.Create(story1)

	var user User
	err = db.Model(&user).Where("id = ?", user1.Id).Select()
	if err != nil {
		panic(err)
	}

	var users []User
	err = db.Model(&users).Select()
	if err != nil {
		panic(err)
	}

	var story Story
	err = db.Model(&story).
		Columns("story.*", "User").
		Where("story.id = ?", story1.Id).
		Select()
	if err != nil {
		panic(err)
	}

	fmt.Println(user)
	fmt.Println(users[0], users[1])
	fmt.Println(story)
	// Output: User<1 admin [admin1@admin admin2@admin]>
	// User<1 admin [admin1@admin admin2@admin]> User<2 root [root1@root root2@root]>
	// Story<1 Cool story User<1 admin [admin1@admin admin2@admin]>>
}