<commit_before>package sql

import (
	"crypto/sha1"
	"database/sql"
	"encoding/hex"
	"errors"
	"fmt"
	"hash"
	"sync"

	"github.com/lib/pq"

	"github.com/barakmich/glog"
	"github.com/google/cayley/graph"
	"github.com/google/cayley/graph/iterator"
	"github.com/google/cayley/quad"
)

const QuadStoreType = "sql"

func init() {
	graph.RegisterQuadStore(QuadStoreType, graph.QuadStoreRegistration{
		NewFunc:           newQuadStore,
		NewForRequestFunc: nil,
		UpgradeFunc:       nil,
		InitFunc:          createSQLTables,
		IsPersistent:      true,
	})
}

var (
	hashPool = sync.Pool{
		New: func() interface{} { return sha1.New() },
	}
	hashSize = sha1.Size
)

type QuadStore struct {
	db        *sql.DB
	sqlFlavor string
	size      int64
	lru       *cache
	noSizes   bool
}

func connectSQLTables(addr string, _ graph.Options) (*sql.DB, error) {
	// TODO(barakmich): Parse options for more friendly addr, other SQLs.
	conn, err := sql.Open("postgres", addr)
	if err != nil {
		glog.Errorf("Couldn't open database at %s: %#v", addr, err)
		return nil, err
	}
	// "Open may just validate its arguments without creating a connection to the database."
	// "To verify that the data source name is valid, call Ping."
	// Source: http://golang.org/pkg/database/sql/#Open
	if err := conn.Ping(); err != nil {
		glog.Errorf("Couldn't open database at %s: %#v", addr, err)
		return nil, err
	}
	return conn, nil
}

func createSQLTables(addr string, options graph.Options) error {
	conn, err := connectSQLTables(addr, options)
	if err != nil {
		return err
	}
	tx, err := conn.Begin()
	if err != nil {
		glog.Errorf("Couldn't begin creation transaction: %s", err)
		return err
	}

	quadTable, err := tx.Exec(`
	CREATE TABLE quads (
		subject TEXT NOT NULL,
		predicate TEXT NOT NULL,
		object TEXT NOT NULL,
		label TEXT,
		horizon BIGSERIAL PRIMARY KEY,
		id BIGINT,
		ts timestamp,
		subject_hash TEXT NOT NULL,
		predicate_hash TEXT NOT NULL,
		object_hash TEXT NOT NULL,
		label_hash TEXT,
		UNIQUE(subject_hash, predicate_hash, object_hash, label_hash)
	);`)
	if err != nil {
		glog.Errorf("Cannot create quad table: %v", quadTable)
		return err
	}
	factor, factorOk, err := options.IntKey("db_fill_factor")
	if !factorOk {
		factor = 50
	}
	var index sql.Result

	index, err = tx.Exec(fmt.Sprintf(`
	CREATE INDEX spo_index ON quads (subject_hash) WITH (FILLFACTOR = %d);
	CREATE INDEX pos_index ON quads (predicate_hash) WITH (FILLFACTOR = %d);
	CREATE INDEX osp_index ON quads (object_hash) WITH (FILLFACTOR = %d);
	`, factor, factor, factor))
	if err != nil {
		glog.Errorf("Cannot create indices: %v", index)
		return err
	}
	tx.Commit()
	return nil
}

func newQuadStore(addr string, options graph.Options) (graph.QuadStore, error) {
	var qs QuadStore
	conn, err := connectSQLTables(addr, options)
	if err != nil {
		return nil, err
	}
	localOpt, localOptOk, err := options.BoolKey("local_optimize")
	if err != nil {
		return nil, err
	}
	qs.db = conn
	qs.sqlFlavor = "postgres"
	qs.size = -1
	qs.lru = newCache(1024)

	// Skip size checking by default.
	qs.noSizes = true
	if localOptOk {
		if localOpt {
			qs.noSizes = false
		}
	}
	return &qs, nil
}

func hashOf(s string) string {
	h := hashPool.Get().(hash.Hash)
	h.Reset()
	defer hashPool.Put(h)
	key := make([]byte, 0, hashSize)
	h.Write([]byte(s))
	key = h.Sum(key)
	return hex.EncodeToString(key)
}

func (qs *QuadStore) copyFrom(tx *sql.Tx, in []graph.Delta) error {
	stmt, err := tx.Prepare(pq.CopyIn("quads", "subject", "predicate", "object", "label", "id", "ts", "subject_hash", "predicate_hash", "object_hash", "label_hash"))
	if err != nil {
		return err
	}
	for _, d := range in {
		_, err := stmt.Exec(
			d.Quad.Subject,
			d.Quad.Predicate,
			d.Quad.Object,
			d.Quad.Label,
			d.ID.Int(),
			d.Timestamp,
			hashOf(d.Quad.Subject),
			hashOf(d.Quad.Predicate),
			hashOf(d.Quad.Object),
			hashOf(d.Quad.Label),
		)
		if err != nil {
			glog.Errorf("couldn't prepare COPY statement: %v", err)
			return err
		}
	}
	_, err = stmt.Exec()
	if err != nil {
		return err
	}
	return stmt.Close()
}

func (qs *QuadStore) runTxPostgres(tx *sql.Tx, in []graph.Delta, opts graph.IgnoreOpts) error {
	allAdds := true
	for _, d := range in {
		if d.Action != graph.Add {
			allAdds = false
		}
	}
	if allAdds {
		return qs.copyFrom(tx, in)
	}

	insert, err := tx.Prepare(`INSERT INTO quads(subject, predicate, object, label, id, ts, subject_hash, predicate_hash, object_hash, label_hash) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)`)
	defer insert.Close()
	if err != nil {
		glog.Errorf("Cannot prepare insert statement: %v", err)
		return err
	}
	for _, d := range in {
		switch d.Action {
		case graph.Add:
			_, err := insert.Exec(
				d.Quad.Subject,
				d.Quad.Predicate,
				d.Quad.Object,
				d.Quad.Label,
				d.ID.Int(),
				d.Timestamp,
				hashOf(d.Quad.Subject),
				hashOf(d.Quad.Predicate),
				hashOf(d.Quad.Object),
				hashOf(d.Quad.Label),
			)
			if err != nil {
				glog.Errorf("couldn't prepare INSERT statement: %v", err)
				return err
			}
		case graph.Delete:
			result, err := tx.Exec(`DELETE FROM quads WHERE subject=$1 and predicate=$2 and object=$3 and label=$4;`,
				d.Quad.Subject, d.Quad.Predicate, d.Quad.Object, d.Quad.Label)
			if err != nil {
				glog.Errorf("couldn't exec DELETE statement: %v", err)
			}
			affected, err := result.RowsAffected()
			if err != nil {
				glog.Errorf("couldn't get DELETE RowsAffected: %v", err)
			}
			if affected != 1 && !opts.IgnoreMissing {
				return errors.New("deleting non-existent triple; rolling back")
			}
		default:
			panic("unknown action")
		}
	}
	return nil
}

func (qs *QuadStore) ApplyDeltas(in []graph.Delta, opts graph.IgnoreOpts) error {
	// TODO(barakmich): Support more ignoreOpts? "ON CONFLICT IGNORE"
	tx, err := qs.db.Begin()
	if err != nil {
		glog.Errorf("couldn't begin write transaction: %v", err)
		return err
	}
	switch qs.sqlFlavor {
	case "postgres":
		err = qs.runTxPostgres(tx, in, opts)
		if err != nil {
			tx.Rollback()
			return err
		}
	default:
		panic("no support for flavor: " + qs.sqlFlavor)
	}
	return tx.Commit()
}

func (qs *QuadStore) Quad(val graph.Value) quad.Quad {
	return val.(quad.Quad)
}

func (qs *QuadStore) QuadIterator(d quad.Direction, val graph.Value) graph.Iterator {
	return NewSQLLinkIterator(qs, d, val.(string))
}

func (qs *QuadStore) NodesAllIterator() graph.Iterator {
	return NewAllIterator(qs, "nodes")
}

func (qs *QuadStore) QuadsAllIterator() graph.Iterator {
	return NewAllIterator(qs, "quads")
}

func (qs *QuadStore) ValueOf(s string) graph.Value {
	return s
}

func (qs *QuadStore) NameOf(v graph.Value) string {
	if v == nil {
		glog.V(2).Info("NameOf was nil")
		return ""
	}
	return v.(string)
}

func (qs *QuadStore) Size() int64 {
	// TODO(barakmich): Sync size with writes.
	if qs.size != -1 {
		return qs.size
	}
	c := qs.db.QueryRow("SELECT COUNT(*) FROM quads;")
	err := c.Scan(&qs.size)
	if err != nil {
		glog.Errorf("Couldn't execute COUNT: %v", err)
		return 0
	}
	return qs.size
}

func (qs *QuadStore) Horizon() graph.PrimaryKey {
	var horizon int64
	err := qs.db.QueryRow("SELECT horizon FROM quads ORDER BY horizon DESC LIMIT 1;").Scan(&horizon)
	if err != nil {
		glog.Errorf("Couldn't execute horizon: %v", err)
		return graph.NewSequentialKey(0)
	}
	return graph.NewSequentialKey(horizon)
}

func (qs *QuadStore) FixedIterator() graph.FixedIterator {
	return iterator.NewFixed(iterator.Identity)
}

func (qs *QuadStore) Close() {
	qs.db.Close()
}

func (qs *QuadStore) QuadDirection(in graph.Value, d quad.Direction) graph.Value {
	q := in.(quad.Quad)
	return q.Get(d)
}

func (qs *QuadStore) Type() string {
	return QuadStoreType
}

func (qs *QuadStore) sizeForIterator(isAll bool, dir quad.Direction, val string) int64 {
	var err error
	if isAll {
		return qs.Size()
	}
	if qs.noSizes {
		if dir == quad.Predicate {
			return (qs.Size() / 100) + 1
		}
		return (qs.Size() / 1000) + 1
	}
	if val, ok := qs.lru.Get(val + string(dir.Prefix())); ok {
		return val
	}
	var size int64
	glog.V(4).Infoln("sql: getting size for select %s, %s", dir.String(), val)
	err = qs.db.QueryRow(
		fmt.Sprintf("SELECT count(*) FROM quads WHERE %s_hash = $1;", dir.String()), hashOf(val)).Scan(&size)
	if err != nil {
		glog.Errorln("Error getting size from SQL database: %v", err)
		return 0
	}
	qs.lru.Put(val+string(dir.Prefix()), size)
	return size
}
<commit_msg>Fix NPE in SQL Quadstore<commit_after>package sql

import (
	"crypto/sha1"
	"database/sql"
	"encoding/hex"
	"errors"
	"fmt"
	"hash"
	"sync"

	"github.com/lib/pq"

	"github.com/barakmich/glog"
	"github.com/google/cayley/graph"
	"github.com/google/cayley/graph/iterator"
	"github.com/google/cayley/quad"
)

const QuadStoreType = "sql"

func init() {
	graph.RegisterQuadStore(QuadStoreType, graph.QuadStoreRegistration{
		NewFunc:           newQuadStore,
		NewForRequestFunc: nil,
		UpgradeFunc:       nil,
		InitFunc:          createSQLTables,
		IsPersistent:      true,
	})
}

var (
	hashPool = sync.Pool{
		New: func() interface{} { return sha1.New() },
	}
	hashSize = sha1.Size
)

type QuadStore struct {
	db        *sql.DB
	sqlFlavor string
	size      int64
	lru       *cache
	noSizes   bool
}

func connectSQLTables(addr string, _ graph.Options) (*sql.DB, error) {
	// TODO(barakmich): Parse options for more friendly addr, other SQLs.
	conn, err := sql.Open("postgres", addr)
	if err != nil {
		glog.Errorf("Couldn't open database at %s: %#v", addr, err)
		return nil, err
	}
	// "Open may just validate its arguments without creating a connection to the database."
	// "To verify that the data source name is valid, call Ping."
	// Source: http://golang.org/pkg/database/sql/#Open
	if err := conn.Ping(); err != nil {
		glog.Errorf("Couldn't open database at %s: %#v", addr, err)
		return nil, err
	}
	return conn, nil
}

func createSQLTables(addr string, options graph.Options) error {
	conn, err := connectSQLTables(addr, options)
	if err != nil {
		return err
	}
	tx, err := conn.Begin()
	if err != nil {
		glog.Errorf("Couldn't begin creation transaction: %s", err)
		return err
	}

	quadTable, err := tx.Exec(`
	CREATE TABLE quads (
		subject TEXT NOT NULL,
		predicate TEXT NOT NULL,
		object TEXT NOT NULL,
		label TEXT,
		horizon BIGSERIAL PRIMARY KEY,
		id BIGINT,
		ts timestamp,
		subject_hash TEXT NOT NULL,
		predicate_hash TEXT NOT NULL,
		object_hash TEXT NOT NULL,
		label_hash TEXT,
		UNIQUE(subject_hash, predicate_hash, object_hash, label_hash)
	);`)
	if err != nil {
		glog.Errorf("Cannot create quad table: %v", quadTable)
		return err
	}
	factor, factorOk, err := options.IntKey("db_fill_factor")
	if !factorOk {
		factor = 50
	}
	var index sql.Result

	index, err = tx.Exec(fmt.Sprintf(`
	CREATE INDEX spo_index ON quads (subject_hash) WITH (FILLFACTOR = %d);
	CREATE INDEX pos_index ON quads (predicate_hash) WITH (FILLFACTOR = %d);
	CREATE INDEX osp_index ON quads (object_hash) WITH (FILLFACTOR = %d);
	`, factor, factor, factor))
	if err != nil {
		glog.Errorf("Cannot create indices: %v", index)
		return err
	}
	tx.Commit()
	return nil
}

func newQuadStore(addr string, options graph.Options) (graph.QuadStore, error) {
	var qs QuadStore
	conn, err := connectSQLTables(addr, options)
	if err != nil {
		return nil, err
	}
	localOpt, localOptOk, err := options.BoolKey("local_optimize")
	if err != nil {
		return nil, err
	}
	qs.db = conn
	qs.sqlFlavor = "postgres"
	qs.size = -1
	qs.lru = newCache(1024)

	// Skip size checking by default.
	qs.noSizes = true
	if localOptOk {
		if localOpt {
			qs.noSizes = false
		}
	}
	return &qs, nil
}

func hashOf(s string) string {
	h := hashPool.Get().(hash.Hash)
	h.Reset()
	defer hashPool.Put(h)
	key := make([]byte, 0, hashSize)
	h.Write([]byte(s))
	key = h.Sum(key)
	return hex.EncodeToString(key)
}

func (qs *QuadStore) copyFrom(tx *sql.Tx, in []graph.Delta) error {
	stmt, err := tx.Prepare(pq.CopyIn("quads", "subject", "predicate", "object", "label", "id", "ts", "subject_hash", "predicate_hash", "object_hash", "label_hash"))
	if err != nil {
		return err
	}
	for _, d := range in {
		_, err := stmt.Exec(
			d.Quad.Subject,
			d.Quad.Predicate,
			d.Quad.Object,
			d.Quad.Label,
			d.ID.Int(),
			d.Timestamp,
			hashOf(d.Quad.Subject),
			hashOf(d.Quad.Predicate),
			hashOf(d.Quad.Object),
			hashOf(d.Quad.Label),
		)
		if err != nil {
			glog.Errorf("couldn't prepare COPY statement: %v", err)
			return err
		}
	}
	_, err = stmt.Exec()
	if err != nil {
		return err
	}
	return stmt.Close()
}

func (qs *QuadStore) runTxPostgres(tx *sql.Tx, in []graph.Delta, opts graph.IgnoreOpts) error {
	allAdds := true
	for _, d := range in {
		if d.Action != graph.Add {
			allAdds = false
		}
	}
	if allAdds {
		return qs.copyFrom(tx, in)
	}

	insert, err := tx.Prepare(`INSERT INTO quads(subject, predicate, object, label, id, ts, subject_hash, predicate_hash, object_hash, label_hash) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)`)
	defer insert.Close()
	if err != nil {
		glog.Errorf("Cannot prepare insert statement: %v", err)
		return err
	}
	for _, d := range in {
		switch d.Action {
		case graph.Add:
			_, err := insert.Exec(
				d.Quad.Subject,
				d.Quad.Predicate,
				d.Quad.Object,
				d.Quad.Label,
				d.ID.Int(),
				d.Timestamp,
				hashOf(d.Quad.Subject),
				hashOf(d.Quad.Predicate),
				hashOf(d.Quad.Object),
				hashOf(d.Quad.Label),
			)
			if err != nil {
				glog.Errorf("couldn't prepare INSERT statement: %v", err)
				return err
			}
		case graph.Delete:
			result, err := tx.Exec(`DELETE FROM quads WHERE subject=$1 and predicate=$2 and object=$3 and label=$4;`,
				d.Quad.Subject, d.Quad.Predicate, d.Quad.Object, d.Quad.Label)
			if err != nil {
				glog.Errorf("couldn't exec DELETE statement: %v", err)
				return err
			}
			affected, err := result.RowsAffected()
			if err != nil {
				glog.Errorf("couldn't get DELETE RowsAffected: %v", err)
				return err
			}
			if affected != 1 && !opts.IgnoreMissing {
				return errors.New("deleting non-existent triple; rolling back")
			}
		default:
			panic("unknown action")
		}
	}
	return nil
}

func (qs *QuadStore) ApplyDeltas(in []graph.Delta, opts graph.IgnoreOpts) error {
	// TODO(barakmich): Support more ignoreOpts? "ON CONFLICT IGNORE"
	tx, err := qs.db.Begin()
	if err != nil {
		glog.Errorf("couldn't begin write transaction: %v", err)
		return err
	}
	switch qs.sqlFlavor {
	case "postgres":
		err = qs.runTxPostgres(tx, in, opts)
		if err != nil {
			tx.Rollback()
			return err
		}
	default:
		panic("no support for flavor: " + qs.sqlFlavor)
	}
	return tx.Commit()
}

func (qs *QuadStore) Quad(val graph.Value) quad.Quad {
	return val.(quad.Quad)
}

func (qs *QuadStore) QuadIterator(d quad.Direction, val graph.Value) graph.Iterator {
	return NewSQLLinkIterator(qs, d, val.(string))
}

func (qs *QuadStore) NodesAllIterator() graph.Iterator {
	return NewAllIterator(qs, "nodes")
}

func (qs *QuadStore) QuadsAllIterator() graph.Iterator {
	return NewAllIterator(qs, "quads")
}

func (qs *QuadStore) ValueOf(s string) graph.Value {
	return s
}

func (qs *QuadStore) NameOf(v graph.Value) string {
	if v == nil {
		glog.V(2).Info("NameOf was nil")
		return ""
	}
	return v.(string)
}

func (qs *QuadStore) Size() int64 {
	// TODO(barakmich): Sync size with writes.
	if qs.size != -1 {
		return qs.size
	}
	c := qs.db.QueryRow("SELECT COUNT(*) FROM quads;")
	err := c.Scan(&qs.size)
	if err != nil {
		glog.Errorf("Couldn't execute COUNT: %v", err)
		return 0
	}
	return qs.size
}

func (qs *QuadStore) Horizon() graph.PrimaryKey {
	var horizon int64
	err := qs.db.QueryRow("SELECT horizon FROM quads ORDER BY horizon DESC LIMIT 1;").Scan(&horizon)
	if err != nil {
		glog.Errorf("Couldn't execute horizon: %v", err)
		return graph.NewSequentialKey(0)
	}
	return graph.NewSequentialKey(horizon)
}

func (qs *QuadStore) FixedIterator() graph.FixedIterator {
	return iterator.NewFixed(iterator.Identity)
}

func (qs *QuadStore) Close() {
	qs.db.Close()
}

func (qs *QuadStore) QuadDirection(in graph.Value, d quad.Direction) graph.Value {
	q := in.(quad.Quad)
	return q.Get(d)
}

func (qs *QuadStore) Type() string {
	return QuadStoreType
}

func (qs *QuadStore) sizeForIterator(isAll bool, dir quad.Direction, val string) int64 {
	var err error
	if isAll {
		return qs.Size()
	}
	if qs.noSizes {
		if dir == quad.Predicate {
			return (qs.Size() / 100) + 1
		}
		return (qs.Size() / 1000) + 1
	}
	if val, ok := qs.lru.Get(val + string(dir.Prefix())); ok {
		return val
	}
	var size int64
	glog.V(4).Infoln("sql: getting size for select %s, %s", dir.String(), val)
	err = qs.db.QueryRow(
		fmt.Sprintf("SELECT count(*) FROM quads WHERE %s_hash = $1;", dir.String()), hashOf(val)).Scan(&size)
	if err != nil {
		glog.Errorln("Error getting size from SQL database: %v", err)
		return 0
	}
	qs.lru.Put(val+string(dir.Prefix()), size)
	return size
}
<|endoftext|>


<commit_before>// Package spritz provides a pure Go implementation of the Spritz stream cipher
// and hash.
package spritz

type state struct {
	// these are all ints instead of bytes to allow for states > 256
	n, d             int
	s                []int
	a, i, j, k, w, z int
}

func (s *state) initialize(n int) {
	*s = state{
		s: make([]int, n),
		w: 1,
		n: n,
		d: n / 16,
	}
	for i := range s.s {
		s.s[i] = i
	}
}

func (s *state) update() {
	s.i = (s.i + s.w) % s.n
	y := (s.j + s.s[s.i]) % s.n
	s.j = (s.k + s.s[y]) % s.n
	s.k = (s.i + s.k + s.s[s.j]) % s.n
	t := s.s[s.i]
	s.s[s.i] = s.s[s.j]
	s.s[s.j] = t
}

func (s *state) output() int {
	y1 := (s.z + s.k) % s.n
	x1 := (s.i + s.s[y1]) % s.n
	y2 := (s.j + s.s[x1]) % s.n
	s.z = s.s[y2]
	return s.z
}

func (s *state) crush() {
	for i := 0; i < s.n/2; i++ {
		y := (s.n - 1) - i
		x1 := s.s[i]
		x2 := s.s[y]
		if x1 > x2 {
			s.s[i] = x2
			s.s[y] = x1
		} else {
			s.s[i] = x1
			s.s[y] = x2
		}
	}
}

func (s *state) whip() {
	r := s.n * 2
	for i := 0; i < r; i++ {
		s.update()
	}
	s.w = (s.w + 2) % s.n
}

func (s *state) shuffle() {
	s.whip()
	s.crush()
	s.whip()
	s.crush()
	s.whip()
	s.a = 0
}

func (s *state) absorbStop() {
	if s.a == s.n/2 {
		s.shuffle()
	}
	s.a = (s.a + 1) % s.n
}

func (s *state) absorbNibble(x int) {
	if s.a == s.n/2 {
		s.shuffle()
	}
	y := (s.n/2 + x) % s.n
	t := s.s[s.a]
	s.s[s.a] = s.s[y]
	s.s[y] = t
	s.a = (s.a + 1) % s.n
}

func (s *state) absorbValue(b int) {
	s.absorbNibble(b % s.d) // LOW
	s.absorbNibble(b / s.d) // HIGH
}

func (s *state) absorb(msg []byte) {
	for _, v := range msg {
		s.absorbValue(int(v))
	}
}

func (s *state) drip() int {
	if s.a > 0 {
		s.shuffle()
	}
	s.update()
	return s.output()
}

func (s *state) squeeze(out []byte) {
	if s.a > 0 {
		s.shuffle()
	}
	for i := range out {
		out[i] = byte(s.drip())
	}
}
<commit_msg>Fix nibble size for n != 256.<commit_after>// Package spritz provides a pure Go implementation of the Spritz stream cipher
// and hash.
package spritz

import "math"

type state struct {
	// these are all ints instead of bytes to allow for states > 256
	n, d             int // state size and nibble size
	s                []int
	a, i, j, k, w, z int
}

func (s *state) initialize(n int) {
	*s = state{
		s: make([]int, n),
		w: 1,
		n: n,
		d: int(math.Ceil(math.Sqrt(float64(n)))),
	}
	for i := range s.s {
		s.s[i] = i
	}
}

func (s *state) update() {
	s.i = (s.i + s.w) % s.n
	y := (s.j + s.s[s.i]) % s.n
	s.j = (s.k + s.s[y]) % s.n
	s.k = (s.i + s.k + s.s[s.j]) % s.n
	t := s.s[s.i]
	s.s[s.i] = s.s[s.j]
	s.s[s.j] = t
}

func (s *state) output() int {
	y1 := (s.z + s.k) % s.n
	x1 := (s.i + s.s[y1]) % s.n
	y2 := (s.j + s.s[x1]) % s.n
	s.z = s.s[y2]
	return s.z
}

func (s *state) crush() {
	for i := 0; i < s.n/2; i++ {
		y := (s.n - 1) - i
		x1 := s.s[i]
		x2 := s.s[y]
		if x1 > x2 {
			s.s[i] = x2
			s.s[y] = x1
		} else {
			s.s[i] = x1
			s.s[y] = x2
		}
	}
}

func (s *state) whip() {
	r := s.n * 2
	for i := 0; i < r; i++ {
		s.update()
	}
	s.w = (s.w + 2) % s.n
}

func (s *state) shuffle() {
	s.whip()
	s.crush()
	s.whip()
	s.crush()
	s.whip()
	s.a = 0
}

func (s *state) absorbStop() {
	if s.a == s.n/2 {
		s.shuffle()
	}
	s.a = (s.a + 1) % s.n
}

func (s *state) absorbNibble(x int) {
	if s.a == s.n/2 {
		s.shuffle()
	}
	y := (s.n/2 + x) % s.n
	t := s.s[s.a]
	s.s[s.a] = s.s[y]
	s.s[y] = t
	s.a = (s.a + 1) % s.n
}

func (s *state) absorbValue(b int) {
	s.absorbNibble(b % s.d) // LOW
	s.absorbNibble(b / s.d) // HIGH
}

func (s *state) absorb(msg []byte) {
	for _, v := range msg {
		s.absorbValue(int(v))
	}
}

func (s *state) drip() int {
	if s.a > 0 {
		s.shuffle()
	}
	s.update()
	return s.output()
}

func (s *state) squeeze(out []byte) {
	if s.a > 0 {
		s.shuffle()
	}
	for i := range out {
		out[i] = byte(s.drip())
	}
}
<|endoftext|>


<commit_before>// Copyright 2015 The Cockroach Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
// implied. See the License for the specific language governing
// permissions and limitations under the License.
//
// Author: Tamir Duberstein (tamird@gmail.com)
// Author: Andrei Matei (andreimatei1@gmail.com)
// Author: Nathan VanBenschoten (nvanbenschoten@gmail.com)

package sqlbase

import (
	"fmt"
	"strings"

	"github.com/cockroachdb/cockroach/sql/parser"
	"github.com/cockroachdb/cockroach/sql/pgwire/pgerror"
	"github.com/cockroachdb/cockroach/util/caller"
)

// SrcCtx contains contextual information about the source of an error.
type SrcCtx struct {
	File     string
	Line     int
	Function string
}

// MakeSrcCtx creates a SrcCtx value with contextual information about the
// caller at the requested depth.
func MakeSrcCtx(depth int) SrcCtx {
	f, l, fun := caller.Lookup(depth + 1)
	return SrcCtx{File: f, Line: l, Function: fun}
}

// ErrorWithPGCode represents errors that carries an error code to the user.
// pgwire recognizes this interfaces and extracts the code.
type ErrorWithPGCode interface {
	error
	Code() string
	SrcContext() SrcCtx
}

var _ ErrorWithPGCode = &ErrNonNullViolation{}
var _ ErrorWithPGCode = &ErrUniquenessConstraintViolation{}
var _ ErrorWithPGCode = &ErrTransactionAborted{}
var _ ErrorWithPGCode = &ErrTransactionCommitted{}
var _ ErrorWithPGCode = &ErrUndefinedDatabase{}
var _ ErrorWithPGCode = &ErrUndefinedTable{}
var _ ErrorWithPGCode = &ErrRetry{}

const (
	txnAbortedMsg = "current transaction is aborted, commands ignored " +
		"until end of transaction block"
	txnCommittedMsg = "current transaction is committed, commands ignored " +
		"until end of transaction block"
	txnRetryMsgPrefix = "restart transaction:"
)

// NewRetryError creates a ErrRetry.
func NewRetryError(cause error) error {
	return &ErrRetry{ctx: MakeSrcCtx(1), msg: fmt.Sprintf("%s %v", txnRetryMsgPrefix, cause)}
}

// ErrRetry means that the transaction can be retried. It signals to the user
// that the SQL txn entered the RESTART_WAIT state after a serialization error,
// and that a ROLLBACK TO SAVEPOINT COCKROACH_RESTART statement should be issued.
type ErrRetry struct {
	ctx SrcCtx
	msg string
}

func (e *ErrRetry) Error() string {
	return e.msg
}

// Code implements the ErrorWithPGCode interface.
func (*ErrRetry) Code() string {
	return pgerror.CodeSerializationFailureError
}

// SrcContext implements the ErrorWithPGCode interface.
func (e *ErrRetry) SrcContext() SrcCtx {
	return e.ctx
}

// NewTransactionAbortedError creates a new ErrTransactionAborted.
func NewTransactionAbortedError(customMsg string) error {
	return &ErrTransactionAborted{ctx: MakeSrcCtx(1), CustomMsg: customMsg}
}

// ErrTransactionAborted represents an error for trying to run a command in the
// context of transaction that's already aborted.
type ErrTransactionAborted struct {
	ctx       SrcCtx
	CustomMsg string
}

func (e *ErrTransactionAborted) Error() string {
	msg := txnAbortedMsg
	if e.CustomMsg != "" {
		msg += "; " + e.CustomMsg
	}
	return msg
}

// Code implements the ErrorWithPGCode interface.
func (*ErrTransactionAborted) Code() string {
	return pgerror.CodeInFailedSQLTransactionError
}

// SrcContext implements the ErrorWithPGCode interface.
func (e *ErrTransactionAborted) SrcContext() SrcCtx {
	return e.ctx
}

// NewTransactionCommittedError creates a new ErrTransactionCommitted.
func NewTransactionCommittedError() error {
	return &ErrTransactionCommitted{ctx: MakeSrcCtx(1)}
}

// ErrTransactionCommitted signals that the SQL txn is in the COMMIT_WAIT state
// and that only a COMMIT statement will be accepted.
type ErrTransactionCommitted struct {
	ctx SrcCtx
}

func (*ErrTransactionCommitted) Error() string {
	return txnCommittedMsg
}

// Code implements the ErrorWithPGCode interface.
func (*ErrTransactionCommitted) Code() string {
	return pgerror.CodeInvalidTransactionStateError
}

// SrcContext implements the ErrorWithPGCode interface.
func (e *ErrTransactionCommitted) SrcContext() SrcCtx {
	return e.ctx
}

// NewNonNullViolationError creates a new ErrNonNullViolation.
func NewNonNullViolationError(columnName string) error {
	return &ErrNonNullViolation{ctx: MakeSrcCtx(1), columnName: columnName}
}

// ErrNonNullViolation represents a violation of a non-NULL constraint.
type ErrNonNullViolation struct {
	ctx        SrcCtx
	columnName string
}

func (e *ErrNonNullViolation) Error() string {
	return fmt.Sprintf("null value in column %q violates not-null constraint", e.columnName)
}

// Code implements the ErrorWithPGCode interface.
func (*ErrNonNullViolation) Code() string {
	return pgerror.CodeNotNullViolationError
}

// SrcContext implements the ErrorWithPGCode interface.
func (e *ErrNonNullViolation) SrcContext() SrcCtx {
	return e.ctx
}

// NewUniquenessConstraintViolationError creates a new
// ErrUniquenessConstrainViolation.
func NewUniquenessConstraintViolationError(
	index *IndexDescriptor, vals []parser.Datum,
) error {
	return &ErrUniquenessConstraintViolation{
		ctx:   MakeSrcCtx(1),
		index: index,
		vals:  vals,
	}
}

// ErrUniquenessConstraintViolation represents a violation of a UNIQUE constraint.
type ErrUniquenessConstraintViolation struct {
	ctx   SrcCtx
	index *IndexDescriptor
	vals  []parser.Datum
}

// Code implements the ErrorWithPGCode interface.
func (*ErrUniquenessConstraintViolation) Code() string {
	return pgerror.CodeUniqueViolationError
}

func (e *ErrUniquenessConstraintViolation) Error() string {
	valStrs := make([]string, 0, len(e.vals))
	for _, val := range e.vals {
		valStrs = append(valStrs, val.String())
	}

	return fmt.Sprintf("duplicate key value (%s)=(%s) violates unique constraint %q",
		strings.Join(e.index.ColumnNames, ","),
		strings.Join(valStrs, ","),
		e.index.Name)
}

// SrcContext implements the ErrorWithPGCode interface.
func (e *ErrUniquenessConstraintViolation) SrcContext() SrcCtx {
	return e.ctx
}

// NewUndefinedTableError creates a new ErrUndefinedTable.
func NewUndefinedTableError(name string) error {
	return &ErrUndefinedTable{ctx: MakeSrcCtx(1), name: name}
}

// ErrUndefinedTable represents a missing database table.
type ErrUndefinedTable struct {
	ctx  SrcCtx
	name string
}

func (e *ErrUndefinedTable) Error() string {
	return fmt.Sprintf("table %q does not exist", e.name)
}

// Code implements the ErrorWithPGCode interface.
func (*ErrUndefinedTable) Code() string {
	return pgerror.CodeUndefinedTableError
}

// SrcContext implements the ErrorWithPGCode interface.
func (e *ErrUndefinedTable) SrcContext() SrcCtx {
	return e.ctx
}

// NewUndefinedDatabaseError creates a new ErrUndefinedDatabase.
func NewUndefinedDatabaseError(name string) error {
	return &ErrUndefinedDatabase{ctx: MakeSrcCtx(1), name: name}
}

// ErrUndefinedDatabase represents a missing database error.
type ErrUndefinedDatabase struct {
	ctx  SrcCtx
	name string
}

func (e *ErrUndefinedDatabase) Error() string {
	return fmt.Sprintf("database %q does not exist", e.name)
}

// Code implements the ErrorWithPGCode interface.
func (*ErrUndefinedDatabase) Code() string {
	// Postgres will return an UndefinedTable error on queries that go to a "relation"
	// that does not exist (a query to a non-existent table or database), but will
	// return an InvalidCatalogName error when connecting to a database that does
	// not exist. We've chosen to return this code for all cases where the error cause
	// is a missing database.
	return pgerror.CodeInvalidCatalogNameError
}

// SrcContext implements the ErrorWithPGCode interface.
func (e *ErrUndefinedDatabase) SrcContext() SrcCtx {
	return e.ctx
}

// NewDatabaseAlreadyExistsError creates a new ErrDatabaseAlreadyExists.
func NewDatabaseAlreadyExistsError(name string) error {
	return &ErrDatabaseAlreadyExists{ctx: MakeSrcCtx(1), name: name}
}

// ErrDatabaseAlreadyExists represents a missing database error.
type ErrDatabaseAlreadyExists struct {
	ctx  SrcCtx
	name string
}

func (e *ErrDatabaseAlreadyExists) Error() string {
	return fmt.Sprintf("database %q already exists", e.name)
}

// Code implements the ErrorWithPGCode interface.
func (*ErrDatabaseAlreadyExists) Code() string {
	return pgerror.CodeDuplicateDatabaseError
}

// SrcContext implements the ErrorWithPGCode interface.
func (e *ErrDatabaseAlreadyExists) SrcContext() SrcCtx {
	return e.ctx
}

// NewRelationAlreadyExistsError creates a new ErrRelationAlreadyExists.
func NewRelationAlreadyExistsError(name string) error {
	return &ErrRelationAlreadyExists{ctx: MakeSrcCtx(1), name: name}
}

// ErrRelationAlreadyExists represents a missing database error.
type ErrRelationAlreadyExists struct {
	ctx  SrcCtx
	name string
}

func (e *ErrRelationAlreadyExists) Error() string {
	return fmt.Sprintf("relation %q already exists", e.name)
}

// Code implements the ErrorWithPGCode interface.
func (*ErrRelationAlreadyExists) Code() string {
	return pgerror.CodeDuplicateRelationError
}

// SrcContext implements the ErrorWithPGCode interface.
func (e *ErrRelationAlreadyExists) SrcContext() SrcCtx {
	return e.ctx
}

// IsIntegrityConstraintError returns true if the error is some kind of SQL
// constraint violation.
func IsIntegrityConstraintError(err error) bool {
	switch err.(type) {
	case *ErrNonNullViolation, *ErrUniquenessConstraintViolation:
		return true
	default:
		return false
	}
}
<commit_msg>sql: Add pgerror code type for syntax errors<commit_after>// Copyright 2015 The Cockroach Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
// implied. See the License for the specific language governing
// permissions and limitations under the License.
//
// Author: Tamir Duberstein (tamird@gmail.com)
// Author: Andrei Matei (andreimatei1@gmail.com)
// Author: Nathan VanBenschoten (nvanbenschoten@gmail.com)

package sqlbase

import (
	"fmt"
	"strings"

	"github.com/cockroachdb/cockroach/sql/parser"
	"github.com/cockroachdb/cockroach/sql/pgwire/pgerror"
	"github.com/cockroachdb/cockroach/util/caller"
)

// SrcCtx contains contextual information about the source of an error.
type SrcCtx struct {
	File     string
	Line     int
	Function string
}

// MakeSrcCtx creates a SrcCtx value with contextual information about the
// caller at the requested depth.
func MakeSrcCtx(depth int) SrcCtx {
	f, l, fun := caller.Lookup(depth + 1)
	return SrcCtx{File: f, Line: l, Function: fun}
}

// ErrorWithPGCode represents errors that carries an error code to the user.
// pgwire recognizes this interfaces and extracts the code.
type ErrorWithPGCode interface {
	error
	Code() string
	SrcContext() SrcCtx
}

var _ ErrorWithPGCode = &ErrNonNullViolation{}
var _ ErrorWithPGCode = &ErrUniquenessConstraintViolation{}
var _ ErrorWithPGCode = &ErrTransactionAborted{}
var _ ErrorWithPGCode = &ErrTransactionCommitted{}
var _ ErrorWithPGCode = &ErrUndefinedDatabase{}
var _ ErrorWithPGCode = &ErrUndefinedTable{}
var _ ErrorWithPGCode = &ErrRetry{}
var _ ErrorWithPGCode = &ErrSyntax{}

const (
	txnAbortedMsg = "current transaction is aborted, commands ignored " +
		"until end of transaction block"
	txnCommittedMsg = "current transaction is committed, commands ignored " +
		"until end of transaction block"
	txnRetryMsgPrefix = "restart transaction:"
)

// NewRetryError creates a ErrRetry.
func NewRetryError(cause error) error {
	return &ErrRetry{ctx: MakeSrcCtx(1), msg: fmt.Sprintf("%s %v", txnRetryMsgPrefix, cause)}
}

// ErrRetry means that the transaction can be retried. It signals to the user
// that the SQL txn entered the RESTART_WAIT state after a serialization error,
// and that a ROLLBACK TO SAVEPOINT COCKROACH_RESTART statement should be issued.
type ErrRetry struct {
	ctx SrcCtx
	msg string
}

func (e *ErrRetry) Error() string {
	return e.msg
}

// Code implements the ErrorWithPGCode interface.
func (*ErrRetry) Code() string {
	return pgerror.CodeSerializationFailureError
}

// SrcContext implements the ErrorWithPGCode interface.
func (e *ErrRetry) SrcContext() SrcCtx {
	return e.ctx
}

// NewTransactionAbortedError creates a new ErrTransactionAborted.
func NewTransactionAbortedError(customMsg string) error {
	return &ErrTransactionAborted{ctx: MakeSrcCtx(1), CustomMsg: customMsg}
}

// ErrTransactionAborted represents an error for trying to run a command in the
// context of transaction that's already aborted.
type ErrTransactionAborted struct {
	ctx       SrcCtx
	CustomMsg string
}

func (e *ErrTransactionAborted) Error() string {
	msg := txnAbortedMsg
	if e.CustomMsg != "" {
		msg += "; " + e.CustomMsg
	}
	return msg
}

// Code implements the ErrorWithPGCode interface.
func (*ErrTransactionAborted) Code() string {
	return pgerror.CodeInFailedSQLTransactionError
}

// SrcContext implements the ErrorWithPGCode interface.
func (e *ErrTransactionAborted) SrcContext() SrcCtx {
	return e.ctx
}

// NewTransactionCommittedError creates a new ErrTransactionCommitted.
func NewTransactionCommittedError() error {
	return &ErrTransactionCommitted{ctx: MakeSrcCtx(1)}
}

// ErrTransactionCommitted signals that the SQL txn is in the COMMIT_WAIT state
// and that only a COMMIT statement will be accepted.
type ErrTransactionCommitted struct {
	ctx SrcCtx
}

func (*ErrTransactionCommitted) Error() string {
	return txnCommittedMsg
}

// Code implements the ErrorWithPGCode interface.
func (*ErrTransactionCommitted) Code() string {
	return pgerror.CodeInvalidTransactionStateError
}

// SrcContext implements the ErrorWithPGCode interface.
func (e *ErrTransactionCommitted) SrcContext() SrcCtx {
	return e.ctx
}

// NewNonNullViolationError creates a new ErrNonNullViolation.
func NewNonNullViolationError(columnName string) error {
	return &ErrNonNullViolation{ctx: MakeSrcCtx(1), columnName: columnName}
}

// ErrNonNullViolation represents a violation of a non-NULL constraint.
type ErrNonNullViolation struct {
	ctx        SrcCtx
	columnName string
}

func (e *ErrNonNullViolation) Error() string {
	return fmt.Sprintf("null value in column %q violates not-null constraint", e.columnName)
}

// Code implements the ErrorWithPGCode interface.
func (*ErrNonNullViolation) Code() string {
	return pgerror.CodeNotNullViolationError
}

// SrcContext implements the ErrorWithPGCode interface.
func (e *ErrNonNullViolation) SrcContext() SrcCtx {
	return e.ctx
}

// NewUniquenessConstraintViolationError creates a new
// ErrUniquenessConstrainViolation.
func NewUniquenessConstraintViolationError(
	index *IndexDescriptor, vals []parser.Datum,
) error {
	return &ErrUniquenessConstraintViolation{
		ctx:   MakeSrcCtx(1),
		index: index,
		vals:  vals,
	}
}

// ErrUniquenessConstraintViolation represents a violation of a UNIQUE constraint.
type ErrUniquenessConstraintViolation struct {
	ctx   SrcCtx
	index *IndexDescriptor
	vals  []parser.Datum
}

// Code implements the ErrorWithPGCode interface.
func (*ErrUniquenessConstraintViolation) Code() string {
	return pgerror.CodeUniqueViolationError
}

func (e *ErrUniquenessConstraintViolation) Error() string {
	valStrs := make([]string, 0, len(e.vals))
	for _, val := range e.vals {
		valStrs = append(valStrs, val.String())
	}

	return fmt.Sprintf("duplicate key value (%s)=(%s) violates unique constraint %q",
		strings.Join(e.index.ColumnNames, ","),
		strings.Join(valStrs, ","),
		e.index.Name)
}

// SrcContext implements the ErrorWithPGCode interface.
func (e *ErrUniquenessConstraintViolation) SrcContext() SrcCtx {
	return e.ctx
}

// NewUndefinedTableError creates a new ErrUndefinedTable.
func NewUndefinedTableError(name string) error {
	return &ErrUndefinedTable{ctx: MakeSrcCtx(1), name: name}
}

// ErrUndefinedTable represents a missing database table.
type ErrUndefinedTable struct {
	ctx  SrcCtx
	name string
}

func (e *ErrUndefinedTable) Error() string {
	return fmt.Sprintf("table %q does not exist", e.name)
}

// Code implements the ErrorWithPGCode interface.
func (*ErrUndefinedTable) Code() string {
	return pgerror.CodeUndefinedTableError
}

// SrcContext implements the ErrorWithPGCode interface.
func (e *ErrUndefinedTable) SrcContext() SrcCtx {
	return e.ctx
}

// NewUndefinedDatabaseError creates a new ErrUndefinedDatabase.
func NewUndefinedDatabaseError(name string) error {
	return &ErrUndefinedDatabase{ctx: MakeSrcCtx(1), name: name}
}

// ErrUndefinedDatabase represents a missing database error.
type ErrUndefinedDatabase struct {
	ctx  SrcCtx
	name string
}

func (e *ErrUndefinedDatabase) Error() string {
	return fmt.Sprintf("database %q does not exist", e.name)
}

// Code implements the ErrorWithPGCode interface.
func (*ErrUndefinedDatabase) Code() string {
	// Postgres will return an UndefinedTable error on queries that go to a "relation"
	// that does not exist (a query to a non-existent table or database), but will
	// return an InvalidCatalogName error when connecting to a database that does
	// not exist. We've chosen to return this code for all cases where the error cause
	// is a missing database.
	return pgerror.CodeInvalidCatalogNameError
}

// SrcContext implements the ErrorWithPGCode interface.
func (e *ErrUndefinedDatabase) SrcContext() SrcCtx {
	return e.ctx
}

// NewDatabaseAlreadyExistsError creates a new ErrDatabaseAlreadyExists.
func NewDatabaseAlreadyExistsError(name string) error {
	return &ErrDatabaseAlreadyExists{ctx: MakeSrcCtx(1), name: name}
}

// ErrDatabaseAlreadyExists represents a missing database error.
type ErrDatabaseAlreadyExists struct {
	ctx  SrcCtx
	name string
}

func (e *ErrDatabaseAlreadyExists) Error() string {
	return fmt.Sprintf("database %q already exists", e.name)
}

// Code implements the ErrorWithPGCode interface.
func (*ErrDatabaseAlreadyExists) Code() string {
	return pgerror.CodeDuplicateDatabaseError
}

// SrcContext implements the ErrorWithPGCode interface.
func (e *ErrDatabaseAlreadyExists) SrcContext() SrcCtx {
	return e.ctx
}

// NewRelationAlreadyExistsError creates a new ErrRelationAlreadyExists.
func NewRelationAlreadyExistsError(name string) error {
	return &ErrRelationAlreadyExists{ctx: MakeSrcCtx(1), name: name}
}

// ErrRelationAlreadyExists represents a missing database error.
type ErrRelationAlreadyExists struct {
	ctx  SrcCtx
	name string
}

func (e *ErrRelationAlreadyExists) Error() string {
	return fmt.Sprintf("relation %q already exists", e.name)
}

// Code implements the ErrorWithPGCode interface.
func (*ErrRelationAlreadyExists) Code() string {
	return pgerror.CodeDuplicateRelationError
}

// SrcContext implements the ErrorWithPGCode interface.
func (e *ErrRelationAlreadyExists) SrcContext() SrcCtx {
	return e.ctx
}

// NewSyntaxError creates a new ErrSyntax.
func NewSyntaxError(msg string) error {
	return &ErrSyntax{ctx: MakeSrcCtx(1), msg: msg}
}

// ErrSyntax represents a syntax error.
type ErrSyntax struct {
	ctx SrcCtx
	msg string
}

func (e *ErrSyntax) Error() string {
	return e.msg
}

// Code implements the ErrorWithPGCode interface.
func (*ErrSyntax) Code() string {
	return pgerror.CodeSyntaxError
}

// SrcContext implements the ErrorWithPGCode interface.
func (e *ErrSyntax) SrcContext() SrcCtx {
	return e.ctx
}

// IsIntegrityConstraintError returns true if the error is some kind of SQL
// constraint violation.
func IsIntegrityConstraintError(err error) bool {
	switch err.(type) {
	case *ErrNonNullViolation, *ErrUniquenessConstraintViolation:
		return true
	default:
		return false
	}
}
<|endoftext|>


<commit_before>// Copyright (C) 2014 Yasuhiro Matsumoto <mattn.jp@gmail.com>.
//
// Use of this source code is governed by an MIT-style
// license that can be found in the LICENSE file.

// +build go1.8

package sqlite3

import (
	"database/sql/driver"
	"errors"

	"context"
)

// Ping implement Pinger.
func (c *SQLiteConn) Ping(ctx context.Context) error {
	if c.db == nil {
		return errors.New("Connection was closed")
	}
	return nil
}

// QueryContext implement QueryerContext.
func (c *SQLiteConn) QueryContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Rows, error) {
	list := make([]namedValue, len(args))
	for i, nv := range args {
		list[i] = namedValue(nv)
	}
	return c.query(ctx, query, list)
}

// ExecContext implement ExecerContext.
func (c *SQLiteConn) ExecContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Result, error) {
	list := make([]namedValue, len(args))
	for i, nv := range args {
		list[i] = namedValue(nv)
	}
	return c.exec(ctx, query, list)
}

// PrepareContext implement ConnPrepareContext.
func (c *SQLiteConn) PrepareContext(ctx context.Context, query string) (driver.Stmt, error) {
	return c.prepare(ctx, query)
}

// BeginTx implement ConnBeginTx.
func (c *SQLiteConn) BeginTx(ctx context.Context) (driver.Tx, error) {
	return c.begin(ctx)
}

// QueryContext implement QueryerContext.
func (s *SQLiteStmt) QueryContext(ctx context.Context, args []driver.NamedValue) (driver.Rows, error) {
	list := make([]namedValue, len(args))
	for i, nv := range args {
		list[i] = namedValue(nv)
	}
	return s.query(ctx, list)
}

// ExecContext implement ExecerContext.
func (s *SQLiteStmt) ExecContext(ctx context.Context, args []driver.NamedValue) (driver.Result, error) {
	list := make([]namedValue, len(args))
	for i, nv := range args {
		list[i] = namedValue(nv)
	}
	return s.exec(ctx, list)
}
<commit_msg>temporary fix BeginTx<commit_after>// Copyright (C) 2014 Yasuhiro Matsumoto <mattn.jp@gmail.com>.
//
// Use of this source code is governed by an MIT-style
// license that can be found in the LICENSE file.

// +build go1.8

package sqlite3

import (
	"database/sql/driver"
	"errors"

	"context"
)

// Ping implement Pinger.
func (c *SQLiteConn) Ping(ctx context.Context) error {
	if c.db == nil {
		return errors.New("Connection was closed")
	}
	return nil
}

// QueryContext implement QueryerContext.
func (c *SQLiteConn) QueryContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Rows, error) {
	list := make([]namedValue, len(args))
	for i, nv := range args {
		list[i] = namedValue(nv)
	}
	return c.query(ctx, query, list)
}

// ExecContext implement ExecerContext.
func (c *SQLiteConn) ExecContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Result, error) {
	list := make([]namedValue, len(args))
	for i, nv := range args {
		list[i] = namedValue(nv)
	}
	return c.exec(ctx, query, list)
}

// PrepareContext implement ConnPrepareContext.
func (c *SQLiteConn) PrepareContext(ctx context.Context, query string) (driver.Stmt, error) {
	return c.prepare(ctx, query)
}

// BeginTx implement ConnBeginTx.
func (c *SQLiteConn) BeginTx(ctx context.Context, opts driver.TxOptions) (driver.Tx, error) {
	return c.begin(ctx)
}

// QueryContext implement QueryerContext.
func (s *SQLiteStmt) QueryContext(ctx context.Context, args []driver.NamedValue) (driver.Rows, error) {
	list := make([]namedValue, len(args))
	for i, nv := range args {
		list[i] = namedValue(nv)
	}
	return s.query(ctx, list)
}

// ExecContext implement ExecerContext.
func (s *SQLiteStmt) ExecContext(ctx context.Context, args []driver.NamedValue) (driver.Result, error) {
	list := make([]namedValue, len(args))
	for i, nv := range args {
		list[i] = namedValue(nv)
	}
	return s.exec(ctx, list)
}
<|endoftext|>


<commit_before>package hybrid

import (
	"testing"

	"github.com/ready-steady/assert"
)

func TestBranin(t *testing.T) {
	fixture := &fixtureBranin
	interpolator := prepare(fixture)

	surrogate := interpolator.Compute(fixture.target)

	values := interpolator.Evaluate(surrogate, fixture.points)
	assert.EqualWithin(values, fixture.values, 100.0, t)
}
<commit_msg>a/hybrid: test admissibility and uniqueness<commit_after>package hybrid

import (
	"testing"

	"github.com/ready-steady/adapt/algorithm/internal"
	"github.com/ready-steady/assert"
)

func TestBranin(t *testing.T) {
	fixture := &fixtureBranin
	interpolator := prepare(fixture)

	surrogate := interpolator.Compute(fixture.target)
	assert.Equal(internal.IsUnique(surrogate.Indices, surrogate.Inputs), true, t)
	assert.Equal(internal.IsAdmissible(surrogate.Indices, surrogate.Inputs), true, t)

	values := interpolator.Evaluate(surrogate, fixture.points)
	assert.EqualWithin(values, fixture.values, 100.0, t)
}
<|endoftext|>


<commit_before>// Package opts helps to write commands which may take multihash
// options.
package opts

import (
	"bytes"
	"errors"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"strings"

	mh "github.com/multiformats/go-multihash"
)

// package errors
var (
	ErrMatch = errors.New("multihash checksums did not match")
)

// Options is a struct used to parse cli flags.
type Options struct {
	Encoding      string
	Algorithm     string
	AlgorithmCode uint64
	Length        int

	fs *flag.FlagSet
}

// FlagValues are the values the various option flags can take.
var FlagValues = struct {
	Encodings  []string
	Algorithms []string
}{
	Encodings:  []string{"raw", "hex", "base58", "base64"},
	Algorithms: []string{"sha1", "sha2-256", "sha2-512", "sha3"},
}

// SetupFlags adds multihash related options to given flagset.
func SetupFlags(f *flag.FlagSet) *Options {
	// TODO: add arg for adding opt prefix and/or overriding opts

	o := new(Options)
	algoStr := "one of: " + strings.Join(FlagValues.Algorithms, ", ")
	f.StringVar(&o.Algorithm, "algorithm", "sha2-256", algoStr)
	f.StringVar(&o.Algorithm, "a", "sha2-256", algoStr+" (shorthand)")

	encStr := "one of: " + strings.Join(FlagValues.Encodings, ", ")
	f.StringVar(&o.Encoding, "encoding", "base58", encStr)
	f.StringVar(&o.Encoding, "e", "base58", encStr+" (shorthand)")

	lengthStr := "checksums length in bits (truncate). -1 is default"
	f.IntVar(&o.Length, "length", -1, lengthStr)
	f.IntVar(&o.Length, "l", -1, lengthStr+" (shorthand)")
	return o
}

// Parse parses the values of flags from given argument slice.
// It is equivalent to flags.Parse(args)
func (o *Options) Parse(args []string) error {
	if err := o.fs.Parse(args); err != nil {
		return err
	}
	return o.ParseError()
}

// ParseError checks the parsed options for errors.
func (o *Options) ParseError() error {
	if !strIn(o.Encoding, FlagValues.Encodings) {
		return fmt.Errorf("encoding '%s' not %s", o.Encoding, FlagValues.Encodings)
	}

	if !strIn(o.Algorithm, FlagValues.Algorithms) {
		return fmt.Errorf("algorithm '%s' not %s", o.Algorithm, FlagValues.Algorithms)
	}

	var found bool
	o.AlgorithmCode, found = mh.Names[o.Algorithm]
	if !found {
		return fmt.Errorf("algorithm '%s' not found (lib error, pls report).", o.Algorithm)
	}

	if o.Length >= 0 {
		if o.Length%8 != 0 {
			return fmt.Errorf("length must be multiple of 8")
		}
		o.Length = o.Length / 8

		if o.Length > mh.DefaultLengths[o.AlgorithmCode] {
			o.Length = mh.DefaultLengths[o.AlgorithmCode]
		}
	}
	return nil
}

// strIn checks wither string a is in set.
func strIn(a string, set []string) bool {
	for _, s := range set {
		if s == a {
			return true
		}
	}
	return false
}

// Check reads all the data in r, calculates its multihash,
// and checks it matches h1
func (o *Options) Check(r io.Reader, h1 mh.Multihash) error {
	h2, err := o.Multihash(r)
	if err != nil {
		return err
	}

	if !bytes.Equal(h1, h2) {
		return fmt.Errorf("computed checksum did not match")
	}

	return nil
}

// Multihash reads all the data in r and calculates its multihash.
func (o *Options) Multihash(r io.Reader) (mh.Multihash, error) {
	b, err := ioutil.ReadAll(r)
	if err != nil {
		return nil, err
	}

	return mh.Sum(b, o.AlgorithmCode, o.Length)
}
<commit_msg>include all hashes in algorithms list<commit_after>// Package opts helps to write commands which may take multihash
// options.
package opts

import (
	"bytes"
	"errors"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"sort"
	"strings"

	mh "github.com/multiformats/go-multihash"
)

// package errors
var (
	ErrMatch = errors.New("multihash checksums did not match")
)

// Options is a struct used to parse cli flags.
type Options struct {
	Encoding      string
	Algorithm     string
	AlgorithmCode uint64
	Length        int

	fs *flag.FlagSet
}

// FlagValues are the values the various option flags can take.
var FlagValues = struct {
	Encodings  []string
	Algorithms []string
}{
	Encodings: []string{"raw", "hex", "base58", "base64"},
	Algorithms: func() []string {
		names := make([]string, 0, len(mh.Names))
		for n := range mh.Names {
			names = append(names, n)
		}
		sort.Strings(names)
		return names
	}(),
}

// SetupFlags adds multihash related options to given flagset.
func SetupFlags(f *flag.FlagSet) *Options {
	// TODO: add arg for adding opt prefix and/or overriding opts

	o := new(Options)
	algoStr := "one of: " + strings.Join(FlagValues.Algorithms, ", ")
	f.StringVar(&o.Algorithm, "algorithm", "sha2-256", algoStr)
	f.StringVar(&o.Algorithm, "a", "sha2-256", algoStr+" (shorthand)")

	encStr := "one of: " + strings.Join(FlagValues.Encodings, ", ")
	f.StringVar(&o.Encoding, "encoding", "base58", encStr)
	f.StringVar(&o.Encoding, "e", "base58", encStr+" (shorthand)")

	lengthStr := "checksums length in bits (truncate). -1 is default"
	f.IntVar(&o.Length, "length", -1, lengthStr)
	f.IntVar(&o.Length, "l", -1, lengthStr+" (shorthand)")
	return o
}

// Parse parses the values of flags from given argument slice.
// It is equivalent to flags.Parse(args)
func (o *Options) Parse(args []string) error {
	if err := o.fs.Parse(args); err != nil {
		return err
	}
	return o.ParseError()
}

// ParseError checks the parsed options for errors.
func (o *Options) ParseError() error {
	if !strIn(o.Encoding, FlagValues.Encodings) {
		return fmt.Errorf("encoding '%s' not %s", o.Encoding, FlagValues.Encodings)
	}

	if !strIn(o.Algorithm, FlagValues.Algorithms) {
		return fmt.Errorf("algorithm '%s' not %s", o.Algorithm, FlagValues.Algorithms)
	}

	var found bool
	o.AlgorithmCode, found = mh.Names[o.Algorithm]
	if !found {
		return fmt.Errorf("algorithm '%s' not found (lib error, pls report).", o.Algorithm)
	}

	if o.Length >= 0 {
		if o.Length%8 != 0 {
			return fmt.Errorf("length must be multiple of 8")
		}
		o.Length = o.Length / 8

		if o.Length > mh.DefaultLengths[o.AlgorithmCode] {
			o.Length = mh.DefaultLengths[o.AlgorithmCode]
		}
	}
	return nil
}

// strIn checks wither string a is in set.
func strIn(a string, set []string) bool {
	for _, s := range set {
		if s == a {
			return true
		}
	}
	return false
}

// Check reads all the data in r, calculates its multihash,
// and checks it matches h1
func (o *Options) Check(r io.Reader, h1 mh.Multihash) error {
	h2, err := o.Multihash(r)
	if err != nil {
		return err
	}

	if !bytes.Equal(h1, h2) {
		return fmt.Errorf("computed checksum did not match")
	}

	return nil
}

// Multihash reads all the data in r and calculates its multihash.
func (o *Options) Multihash(r io.Reader) (mh.Multihash, error) {
	b, err := ioutil.ReadAll(r)
	if err != nil {
		return nil, err
	}

	return mh.Sum(b, o.AlgorithmCode, o.Length)
}
<|endoftext|>


<commit_before>package libcoap

/*
#cgo LDFLAGS: -lcoap-1
#include <coap/coap.h>
*/
import "C"
import "time"
import "unsafe"

type DtlsParam struct {
    CaFilename          *string
    CaPath              *string
    CertificateFilename *string
    PrivateKeyFilename  *string
}

type Context struct {
    ptr     *C.coap_context_t
    handler ResponseHandler
    dtls    *C.coap_dtls_param_t
}

var contexts = make(map[*C.coap_context_t] *Context)

func Startup() {
    C.coap_startup()
    // C.coap_dtls_set_log_level(C.LOG_DEBUG)
    // C.coap_set_log_level(C.LOG_DEBUG)
}

func Cleanup() {
    C.coap_cleanup()
}

func NewContext(addr *Address) *Context {
    var caddr *C.coap_address_t = nil
    if addr != nil {
      caddr = &addr.value
    }

    ptr := C.coap_new_context(caddr)
    if ptr != nil {
        context := &Context{ ptr, nil, nil }
        contexts[ptr] = context
        return context
    } else {
        return nil
    }
}

func NewContextDtls(addr *Address, dtls *DtlsParam) *Context {
    var caddr *C.coap_address_t = nil
    if addr != nil {
      caddr = &addr.value
    }

    var cdtls *C.coap_dtls_param_t = nil
    if dtls != nil {
        cdtls = &C.coap_dtls_param_t{};
        if dtls.CaFilename != nil {
            cdtls.ca_filename = C.CString(*dtls.CaFilename)
        }
        if dtls.CaPath != nil {
            cdtls.ca_path = C.CString(*dtls.CaPath)
        }
        if dtls.CertificateFilename != nil {
            cdtls.cert_filename = C.CString(*dtls.CertificateFilename)
        }
        if dtls.PrivateKeyFilename != nil {
            cdtls.pkey_filename = C.CString(*dtls.PrivateKeyFilename)
        }
    }

    ptr := C.coap_new_context_dtls(caddr, cdtls)
    if ptr != nil {
        context := &Context{ ptr, nil, cdtls }
        contexts[ptr] = context
        return context
    } else {
        return nil
    }
}

func (context *Context) FreeContext() {
    ptr := context.ptr

    delete(contexts, ptr)
    context.ptr = nil
    C.coap_free_context(ptr)

    if context.dtls != nil {
        if context.dtls.ca_filename != nil {
            C.free(unsafe.Pointer(context.dtls.ca_filename))
        }
        if context.dtls.ca_path != nil {
            C.free(unsafe.Pointer(context.dtls.ca_path))
        }
        if context.dtls.cert_filename != nil {
            C.free(unsafe.Pointer(context.dtls.cert_filename))
        }
        if context.dtls.pkey_filename != nil {
            C.free(unsafe.Pointer(context.dtls.pkey_filename))
        }
        context.dtls = nil
    }
}

func (context *Context) CanExit() bool {
    return 1 == C.coap_can_exit(context.ptr)
}

func (context *Context) RunOnce(timeout time.Duration) time.Duration {
    d := C.coap_run_once(context.ptr, C.uint(timeout / time.Millisecond))
    return time.Duration(d) * time.Millisecond
}
<commit_msg>Enable PKI for DTLS context<commit_after>package libcoap

/*
#cgo LDFLAGS: -lcoap-1
#include <coap/coap.h>
*/
import "C"
import "time"
import "unsafe"

type DtlsParam struct {
    CaFilename          *string
    CaPath              *string
    CertificateFilename *string
    PrivateKeyFilename  *string
}

type Context struct {
    ptr     *C.coap_context_t
    handler ResponseHandler
    dtls    *C.coap_dtls_param_t
}

var contexts = make(map[*C.coap_context_t] *Context)

func Startup() {
    C.coap_startup()
    // C.coap_dtls_set_log_level(C.LOG_DEBUG)
    // C.coap_set_log_level(C.LOG_DEBUG)
}

func Cleanup() {
    C.coap_cleanup()
}

func NewContext(addr *Address) *Context {
    var caddr *C.coap_address_t = nil
    if addr != nil {
      caddr = &addr.value
    }

    ptr := C.coap_new_context(caddr)
    if ptr != nil {
        context := &Context{ ptr, nil, nil }
        contexts[ptr] = context
        return context
    } else {
        return nil
    }
}

func NewContextDtls(addr *Address, dtls *DtlsParam) *Context {
    var caddr *C.coap_address_t = nil
    if addr != nil {
      caddr = &addr.value
    }

    var cdtls *C.coap_dtls_param_t = nil
    if dtls != nil {
        cdtls = &C.coap_dtls_param_t{};
        if dtls.CaFilename != nil {
            cdtls.ca_filename = C.CString(*dtls.CaFilename)
        }
        if dtls.CaPath != nil {
            cdtls.ca_path = C.CString(*dtls.CaPath)
        }
        if dtls.CertificateFilename != nil {
            cdtls.cert_filename = C.CString(*dtls.CertificateFilename)
        }
        if dtls.PrivateKeyFilename != nil {
            cdtls.pkey_filename = C.CString(*dtls.PrivateKeyFilename)
        }
    }

    ptr := C.coap_new_context_dtls(caddr, cdtls)
    if ptr != nil {
        // Enable PKI
        var setup_data *C.coap_dtls_pki_t = &C.coap_dtls_pki_t{}
        if dtls.CaFilename != nil {
            setup_data.ca_file  = C.CString(*dtls.CaFilename)
        }
        if dtls.CertificateFilename != nil {
            setup_data.public_cert = C.CString(*dtls.CertificateFilename)
        }
        if dtls.PrivateKeyFilename != nil {
            setup_data.private_key = C.CString(*dtls.PrivateKeyFilename)
        }
        ok := C.coap_dtls_context_set_pki(ptr, setup_data)

        if ok == 1 {
            context := &Context{ ptr, nil, cdtls }
            contexts[ptr] = context
            return context
        } else {
            return nil
        }

    } else {
        return nil
    }
}

func (context *Context) FreeContext() {
    ptr := context.ptr

    delete(contexts, ptr)
    context.ptr = nil
    C.coap_free_context(ptr)

    if context.dtls != nil {
        if context.dtls.ca_filename != nil {
            C.free(unsafe.Pointer(context.dtls.ca_filename))
        }
        if context.dtls.ca_path != nil {
            C.free(unsafe.Pointer(context.dtls.ca_path))
        }
        if context.dtls.cert_filename != nil {
            C.free(unsafe.Pointer(context.dtls.cert_filename))
        }
        if context.dtls.pkey_filename != nil {
            C.free(unsafe.Pointer(context.dtls.pkey_filename))
        }
        context.dtls = nil
    }
}

func (context *Context) CanExit() bool {
    return 1 == C.coap_can_exit(context.ptr)
}

func (context *Context) RunOnce(timeout time.Duration) time.Duration {
    d := C.coap_run_once(context.ptr, C.uint(timeout / time.Millisecond))
    return time.Duration(d) * time.Millisecond
}
<|endoftext|>


<commit_before>package logger

// Originally from github.com/motemen/ghq/utils

import (
	"fmt"
	"io"
	"os"

	colorine "github.com/motemen/go-colorine"
)

type Logger struct {
	logger *colorine.Logger
}

// New is constructor for new colorine logger
func New() Logger {
	logger := &colorine.Logger{
		Prefixes: colorine.Prefixes{
			"warning": colorine.Warn,

			"error": colorine.Error,

			"":        colorine.Info,
			"info":    colorine.Info,
			"created": colorine.Info,
			"updated": colorine.Info,
			"thrown":  colorine.Info,
			"retired": colorine.Info,
		},
	}

	// Default output
	logger.SetOutput(os.Stderr)
	return Logger{logger: logger}
}

// SetOutput sets output
func (l *Logger) SetOutput(w io.Writer) {
	l.logger.SetOutput(w)
}

// Log outputs `message` with `prefix` by go-colorine
func (l *Logger) Log(prefix, message string) {
	l.logger.Log(prefix, message)
}

// Logf outputs `message` with `prefix` by go-colorine
func (l *Logger) Logf(prefix, message string, args ...interface{}) {
	msg := fmt.Sprintf(message, args...)
	l.logger.Log(prefix, msg)
}

// ErrorIf outputs log if `err` occurs.
func (l *Logger) ErrorIf(err error) bool {
	if err != nil {
		l.Log("error", err.Error())
		return true
	}

	return false
}

// DieIf outputs log and exit(1) if `err` occurs.
func (l *Logger) DieIf(err error) {
	if err != nil {
		l.Log("error", err.Error())
		os.Exit(1)
	}
}

// PanicIf raise panic if `err` occurs.
func (l *Logger) PanicIf(err error) {
	if err != nil {
		panic(err)
	}
}

var defaultLogger = New()

// Log outputs `message` with `prefix` by go-colorine
func Log(prefix, message string) {
	defaultLogger.Log(prefix, message)
}

// Logf outputs `message` with `prefix` by go-colorine
func Logf(prefix, message string, args ...interface{}) {
	defaultLogger.Logf(prefix, message, args)
}

// ErrorIf outputs log if `err` occurs.
func ErrorIf(err error) bool {
	return defaultLogger.ErrorIf(err)
}

// DieIf outputs log and exit(1) if `err` occurs.
func DieIf(err error) {
	defaultLogger.DieIf(err)
}

// PanicIf raise panic if `err` occurs.
func PanicIf(err error) {
	defaultLogger.PanicIf(err)
}
<commit_msg>noone uses Panicf<commit_after>package logger

// Originally from github.com/motemen/ghq/utils

import (
	"fmt"
	"io"
	"os"

	colorine "github.com/motemen/go-colorine"
)

type Logger struct {
	logger *colorine.Logger
}

// New is constructor for new colorine logger
func New() Logger {
	logger := &colorine.Logger{
		Prefixes: colorine.Prefixes{
			"warning": colorine.Warn,

			"error": colorine.Error,

			"":        colorine.Info,
			"info":    colorine.Info,
			"created": colorine.Info,
			"updated": colorine.Info,
			"thrown":  colorine.Info,
			"retired": colorine.Info,
		},
	}

	// Default output
	logger.SetOutput(os.Stderr)
	return Logger{logger: logger}
}

// SetOutput sets output
func (l *Logger) SetOutput(w io.Writer) {
	l.logger.SetOutput(w)
}

// Log outputs `message` with `prefix` by go-colorine
func (l *Logger) Log(prefix, message string) {
	l.logger.Log(prefix, message)
}

// Logf outputs `message` with `prefix` by go-colorine
func (l *Logger) Logf(prefix, message string, args ...interface{}) {
	msg := fmt.Sprintf(message, args...)
	l.logger.Log(prefix, msg)
}

// ErrorIf outputs log if `err` occurs.
func (l *Logger) ErrorIf(err error) bool {
	if err != nil {
		l.Log("error", err.Error())
		return true
	}

	return false
}

// DieIf outputs log and exit(1) if `err` occurs.
func (l *Logger) DieIf(err error) {
	if err != nil {
		l.Log("error", err.Error())
		os.Exit(1)
	}
}

var defaultLogger = New()

// Log outputs `message` with `prefix` by go-colorine
func Log(prefix, message string) {
	defaultLogger.Log(prefix, message)
}

// Logf outputs `message` with `prefix` by go-colorine
func Logf(prefix, message string, args ...interface{}) {
	defaultLogger.Logf(prefix, message, args)
}

// ErrorIf outputs log if `err` occurs.
func ErrorIf(err error) bool {
	return defaultLogger.ErrorIf(err)
}

// DieIf outputs log and exit(1) if `err` occurs.
func DieIf(err error) {
	defaultLogger.DieIf(err)
}
<|endoftext|>


<commit_before>package consumers

import (
	"github.com/Shopify/sarama"
	log "github.com/Sirupsen/logrus"
	saramaC "github.com/d33d33/sarama-cluster"
	"github.com/d33d33/viper" // FIXME https://github.com/spf13/viper/pull/285

	"github.com/runabove/metronome/src/metronome/constants"
	"github.com/runabove/metronome/src/metronome/models"
	"github.com/runabove/metronome/src/metronome/pg"
)

// TaskConsumer consumed tasks messages from a Kafka topic to maintain the tasks database.
type TaskConsumer struct {
	consumer *saramaC.Consumer
}

// NewTaskConsumer returns a new task consumer.
func NewTaskConsumer() (*TaskConsumer, error) {
	brokers := viper.GetStringSlice("kafka.brokers")

	config := saramaC.NewConfig()
	config.ClientID = "metronome-aggregator"
	config.Consumer.Offsets.Initial = sarama.OffsetOldest

	consumer, err := saramaC.NewConsumer(brokers, "aggregator", []string{constants.KafkaTopicTasks}, config)
	if err != nil {
		return nil, err
	}

	tc := &TaskConsumer{
		consumer: consumer,
	}

	go func() {
		for {
			select {
			case msg, ok := <-consumer.Messages():
				if !ok { // shuting down
					return
				}
				tc.handleMsg(msg)
			}
		}
	}()

	return tc, nil
}

// Close the consumer.
func (tc *TaskConsumer) Close() error {
	return tc.consumer.Close()
}

// Handle message from Kafka.
// Apply updates to the database.
func (tc *TaskConsumer) handleMsg(msg *sarama.ConsumerMessage) {
	var t models.Task
	if err := t.FromKafka(msg); err != nil {
		log.Error(err)
		return
	}

	db := pg.DB()

	if t.Schedule == "" {
		log.Infof("DELETE task: %s", t.GUID)

		_, err := db.Model(&t).Delete()
		if err != nil {
			log.Errorf("%v", err) // TODO log for replay or not commit
		}
		return
	}
	log.Infof("UPDATE task: %s", t.GUID)

	_, err := db.Model(&t).OnConflict("(guid) DO UPDATE").Set("name = ?name").Set("urn = ?urn").Set("schedule = ?schedule").Insert()
	if err != nil {
		log.Errorf("%v", err) // TODO log for replay or not commit
	}
}
<commit_msg>feat: send task update to ws<commit_after>package consumers

import (
	"github.com/Shopify/sarama"
	log "github.com/Sirupsen/logrus"
	saramaC "github.com/d33d33/sarama-cluster"
	"github.com/d33d33/viper" // FIXME https://github.com/spf13/viper/pull/285

	"github.com/runabove/metronome/src/metronome/constants"
	"github.com/runabove/metronome/src/metronome/models"
	"github.com/runabove/metronome/src/metronome/pg"
	"github.com/runabove/metronome/src/metronome/redis"
)

// TaskConsumer consumed tasks messages from a Kafka topic to maintain the tasks database.
type TaskConsumer struct {
	consumer *saramaC.Consumer
}

// NewTaskConsumer returns a new task consumer.
func NewTaskConsumer() (*TaskConsumer, error) {
	brokers := viper.GetStringSlice("kafka.brokers")

	config := saramaC.NewConfig()
	config.ClientID = "metronome-aggregator"
	config.Consumer.Offsets.Initial = sarama.OffsetOldest

	consumer, err := saramaC.NewConsumer(brokers, "aggregator", []string{constants.KafkaTopicTasks}, config)
	if err != nil {
		return nil, err
	}

	tc := &TaskConsumer{
		consumer: consumer,
	}

	go func() {
		for {
			select {
			case msg, ok := <-consumer.Messages():
				if !ok { // shuting down
					return
				}
				tc.handleMsg(msg)
			}
		}
	}()

	return tc, nil
}

// Close the consumer.
func (tc *TaskConsumer) Close() error {
	return tc.consumer.Close()
}

// Handle message from Kafka.
// Apply updates to the database.
func (tc *TaskConsumer) handleMsg(msg *sarama.ConsumerMessage) {
	var t models.Task
	if err := t.FromKafka(msg); err != nil {
		log.Error(err)
		return
	}

	db := pg.DB()

	if t.Schedule == "" {
		log.Infof("DELETE task: %s", t.GUID)

		_, err := db.Model(&t).Delete()
		if err != nil {
			log.Errorf("%v", err) // TODO log for replay or not commit
		}
		return
	}
	log.Infof("UPDATE task: %s", t.GUID)

	_, err := db.Model(&t).OnConflict("(guid) DO UPDATE").Set("name = ?name").Set("urn = ?urn").Set("schedule = ?schedule").Insert()
	if err != nil {
		log.Errorf("%v", err) // TODO log for replay or not commit
	}
	if err := redis.DB().PublishTopic(t.UserID, "task", t.ToJSON()).Err(); err != nil {
		log.Error(err)
	}
}
<|endoftext|>


<commit_before>package orm

import (
	"errors"
	"fmt"
	"sync"

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

type Query struct {
	db    dber
	model tableModel
	err   error

	tableName  types.Q
	tableAlias string

	tables     []byte
	fields     []string
	columns    []byte
	set        []byte
	where      []byte
	join       []byte
	group      []byte
	order      []byte
	onConflict []byte
	returning  []byte
	limit      int
	offset     int
}

func NewQuery(db dber, v ...interface{}) *Query {
	q := Query{
		db: db,
	}
	switch l := len(v); {
	case l == 1:
		v0 := v[0]
		if v0 != nil {
			q.model, q.err = newTableModel(v0)
		}
	case l > 1:
		q.model, q.err = newTableModel(&v)
	}
	if q.model != nil {
		q.tableName = q.FormatQuery(q.tableName, string(q.model.Table().Name))
	}
	return &q
}

func (q *Query) copy() *Query {
	cp := *q
	return &cp
}

func (q *Query) setErr(err error) {
	if q.err == nil {
		q.err = err
	}
}

func (q *Query) Alias(alias string) *Query {
	q.tableAlias = alias
	return q
}

func (q *Query) Table(names ...string) *Query {
	for _, name := range names {
		q.tables = appendSep(q.tables, ", ")
		q.tables = types.AppendField(q.tables, name, 1)
	}
	return q
}

func (q *Query) Column(columns ...string) *Query {
loop:
	for _, column := range columns {
		if q.model != nil {
			if j := q.model.Join(column); j != nil {
				continue loop
			}
		}

		q.fields = append(q.fields, column)
		q.columns = appendSep(q.columns, ", ")
		q.columns = types.AppendField(q.columns, column, 1)
	}
	return q
}

func (q *Query) ColumnExpr(expr string, params ...interface{}) *Query {
	q.columns = appendSep(q.columns, ", ")
	q.columns = q.FormatQuery(q.columns, expr, params...)
	return q
}

func (q *Query) Set(set string, params ...interface{}) *Query {
	q.set = appendSep(q.set, ", ")
	q.set = q.FormatQuery(q.set, set, params...)
	return q
}

func (q *Query) Where(where string, params ...interface{}) *Query {
	q.where = appendSep(q.where, " AND ")
	q.where = append(q.where, '(')
	q.where = q.FormatQuery(q.where, where, params...)
	q.where = append(q.where, ')')
	return q
}

// WhereAnd joins passed conditions using AND operation.
func (q *Query) WhereAnd(conditions ...*SQL) *Query {
	q.where = appendSep(q.where, " AND ")
	q.where = append(q.where, '(')
	for i, cond := range conditions {
		q.where = cond.AppendFormat(q.where, q)
		if i != len(conditions)-1 {
			q.where = append(q.where, " AND "...)
		}
	}
	q.where = append(q.where, ')')
	return q
}

// WhereOr joins passed conditions using OR operation.
func (q *Query) WhereOr(conditions ...*SQL) *Query {
	q.where = appendSep(q.where, " AND ")
	q.where = append(q.where, '(')
	for i, cond := range conditions {
		q.where = cond.AppendFormat(q.where, q)
		if i != len(conditions)-1 {
			q.where = append(q.where, " OR "...)
		}
	}
	q.where = append(q.where, ')')
	return q
}

func (q *Query) Join(join string, params ...interface{}) *Query {
	q.join = appendSep(q.join, " ")
	q.join = q.FormatQuery(q.join, join, params...)
	return q
}

func (q *Query) Group(group string, params ...interface{}) *Query {
	q.group = appendSep(q.group, ", ")
	q.group = q.FormatQuery(q.group, group, params...)
	return q
}

func (q *Query) Order(order string, params ...interface{}) *Query {
	q.order = appendSep(q.order, ", ")
	q.order = q.FormatQuery(q.order, order, params...)
	return q
}

func (q *Query) Limit(n int) *Query {
	q.limit = n
	return q
}

func (q *Query) Offset(n int) *Query {
	q.offset = n
	return q
}

func (q *Query) OnConflict(s string, params ...interface{}) *Query {
	q.onConflict = q.FormatQuery(nil, s, params...)
	return q
}

func (q *Query) Returning(columns ...interface{}) *Query {
	for _, column := range columns {
		q.returning = appendSep(q.returning, ", ")

		switch column := column.(type) {
		case string:
			q.returning = types.AppendField(q.returning, column, 1)
		case types.ValueAppender:
			var err error
			q.returning, err = column.AppendValue(q.returning, 1)
			if err != nil {
				q.setErr(err)
			}
		default:
			q.setErr(fmt.Errorf("unsupported column type: %T", column))
		}
	}
	return q
}

// Count returns number of rows matching the query using count aggregate function.
func (q *Query) Count() (int, error) {
	if q.err != nil {
		return 0, q.err
	}

	q = q.copy()
	q.columns = types.Q("count(*)")
	q.order = nil
	q.limit = 0
	q.offset = 0

	sel := &selectQuery{
		Query: q,
	}
	var count int
	_, err := q.db.QueryOne(Scan(&count), sel, q.model)
	return count, err
}

// First selects the first row.
func (q *Query) First() error {
	b := columns(q.model.Table().Alias, "", q.model.Table().PKs)
	return q.Order(string(b)).Limit(1).Select()
}

// Last selects the last row.
func (q *Query) Last() error {
	b := columns(q.model.Table().Alias, "", q.model.Table().PKs)
	b = append(b, " DESC"...)
	return q.Order(string(b)).Limit(1).Select()
}

// Select selects the model.
func (q *Query) Select(values ...interface{}) error {
	if q.err != nil {
		return q.err
	}

	q.joinHasOne()
	sel := selectQuery{q}

	var model Model
	var err error
	if len(values) > 0 {
		model, err = NewModel(values...)
		if err != nil {
			return err
		}
	} else {
		model = q.model
	}

	if m, ok := model.(useQueryOne); ok && m.useQueryOne() {
		_, err = q.db.QueryOne(model, sel, q.model)
	} else {
		_, err = q.db.Query(model, sel, q.model)
	}
	if err != nil {
		return err
	}

	if q.model != nil {
		return selectJoins(q.db, q.model.GetJoins())
	}
	return nil
}

// SelectAndCount runs Select and Count in two separate goroutines,
// waits for them to finish and returns the result.
func (q *Query) SelectAndCount(values ...interface{}) (count int, err error) {
	var wg sync.WaitGroup
	wg.Add(2)

	go func() {
		defer wg.Done()
		if e := q.Select(values...); e != nil {
			err = e
		}
	}()

	go func() {
		defer wg.Done()
		var e error
		count, e = q.Count()
		if e != nil {
			err = e
		}
	}()

	wg.Wait()
	return count, err
}

func (q *Query) joinHasOne() {
	if q.model == nil {
		return
	}
	joins := q.model.GetJoins()
	for i := range joins {
		j := &joins[i]
		if j.Rel.One {
			j.JoinOne(q)
		}
	}
}

func selectJoins(db dber, joins []join) error {
	var err error
	for i := range joins {
		j := &joins[i]
		if j.Rel.One {
			err = selectJoins(db, j.JoinModel.GetJoins())
		} else {
			err = j.Select(db)
		}
		if err != nil {
			return err
		}
	}
	return nil
}

// Create inserts the model.
func (q *Query) Create(values ...interface{}) (*types.Result, error) {
	if q.err != nil {
		return nil, q.err
	}

	var model Model
	if len(values) > 0 {
		model = Scan(values...)
	} else {
		model = q.model
	}

	ins := insertQuery{Query: q}
	return q.db.Query(model, ins, q.model)
}

// SelectOrCreate selects the model creating one if it does not exist.
func (q *Query) SelectOrCreate(values ...interface{}) (created bool, err error) {
	if q.err != nil {
		return false, q.err
	}

	for i := 0; i < 10; i++ {
		err := q.Select(values...)
		if err == nil {
			return false, nil
		}
		if err != internal.ErrNoRows {
			return false, err
		}

		res, err := q.Create(values...)
		if err != nil {
			if pgErr, ok := err.(internal.PGError); ok {
				if pgErr.IntegrityViolation() {
					continue
				}
				if pgErr.Field('C') == "55000" {
					// Retry on "#55000 attempted to delete invisible tuple".
					continue
				}
			}
			return false, err
		}
		if res.Affected() == 1 {
			return true, nil
		}
	}

	return false, errors.New("pg: GetOrCreate does not make progress after 10 iterations")
}

// Update updates the model.
func (q *Query) Update() (*types.Result, error) {
	if q.err != nil {
		return nil, q.err
	}
	return q.db.Query(q.model, updateQuery{q}, q.model)
}

// Delete deletes the model.
func (q *Query) Delete() (*types.Result, error) {
	if q.err != nil {
		return nil, q.err
	}
	return q.db.Exec(deleteQuery{q}, q.model)
}

func (q *Query) FormatQuery(dst []byte, query string, params ...interface{}) []byte {
	params = append(params, q.model)
	return q.db.FormatQuery(dst, query, params...)
}

func (q *Query) appendTableAlias(b []byte) []byte {
	if q.tableAlias != "" {
		return types.AppendField(b, q.tableAlias, 1)
	}
	return append(b, q.model.Table().Alias...)
}

func (q *Query) appendTableNameWithAlias(b []byte) []byte {
	b = append(b, q.tableName...)
	b = append(b, " AS "...)
	b = q.appendTableAlias(b)
	return b
}

func (q *Query) haveTables() bool {
	return q.model != nil || len(q.tables) > 0
}

func (q *Query) appendTables(b []byte) []byte {
	if q.model != nil {
		b = q.appendTableNameWithAlias(b)
		if len(q.tables) > 0 {
			b = append(b, ", "...)
		}
	}
	b = append(b, q.tables...)
	return b
}

func (q *Query) appendSet(b []byte) ([]byte, error) {
	b = append(b, " SET "...)
	if len(q.set) > 0 {
		b = append(b, q.set...)
		return b, nil
	}

	if q.model == nil {
		return nil, errors.New("pg: Model(nil)")
	}

	table := q.model.Table()
	strct := q.model.Value()

	if len(q.fields) > 0 {
		for i, fieldName := range q.fields {
			field, err := table.GetField(fieldName)
			if err != nil {
				return nil, err
			}

			b = append(b, field.ColName...)
			b = append(b, " = "...)
			b = field.AppendValue(b, strct, 1)
			if i != len(q.fields)-1 {
				b = append(b, ", "...)
			}
		}
		return b, nil
	}

	start := len(b)
	for _, field := range table.Fields {
		if field.Has(PrimaryKeyFlag) {
			continue
		}

		b = append(b, field.ColName...)
		b = append(b, " = "...)
		b = field.AppendValue(b, strct, 1)
		b = append(b, ", "...)
	}
	if len(b) > start {
		b = b[:len(b)-2]
	}
	return b, nil
}

func (q *Query) appendWhere(b []byte) ([]byte, error) {
	b = append(b, " WHERE "...)
	if len(q.where) > 0 {
		b = append(b, q.where...)
		return b, nil
	}

	if q.model == nil {
		return nil, errors.New("pg: Model(nil)")
	}

	table := q.model.Table()
	if err := table.checkPKs(); err != nil {
		return nil, err
	}
	b = appendColumnAndValue(b, q.model.Value(), table, table.PKs)
	return b, nil
}
<commit_msg>Add experimental Apply API.<commit_after>package orm

import (
	"errors"
	"fmt"
	"sync"

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

type Query struct {
	db    dber
	model tableModel
	err   error

	tableName  types.Q
	tableAlias string

	tables     []byte
	fields     []string
	columns    []byte
	set        []byte
	where      []byte
	join       []byte
	group      []byte
	order      []byte
	onConflict []byte
	returning  []byte
	limit      int
	offset     int
}

func NewQuery(db dber, v ...interface{}) *Query {
	q := Query{
		db: db,
	}
	switch l := len(v); {
	case l == 1:
		v0 := v[0]
		if v0 != nil {
			q.model, q.err = newTableModel(v0)
		}
	case l > 1:
		q.model, q.err = newTableModel(&v)
	}
	if q.model != nil {
		q.tableName = q.FormatQuery(q.tableName, string(q.model.Table().Name))
	}
	return &q
}

func (q *Query) copy() *Query {
	cp := *q
	return &cp
}

func (q *Query) setErr(err error) {
	if q.err == nil {
		q.err = err
	}
}

func (q *Query) Alias(alias string) *Query {
	q.tableAlias = alias
	return q
}

func (q *Query) Table(names ...string) *Query {
	for _, name := range names {
		q.tables = appendSep(q.tables, ", ")
		q.tables = types.AppendField(q.tables, name, 1)
	}
	return q
}

func (q *Query) Column(columns ...string) *Query {
loop:
	for _, column := range columns {
		if q.model != nil {
			if j := q.model.Join(column); j != nil {
				continue loop
			}
		}

		q.fields = append(q.fields, column)
		q.columns = appendSep(q.columns, ", ")
		q.columns = types.AppendField(q.columns, column, 1)
	}
	return q
}

func (q *Query) ColumnExpr(expr string, params ...interface{}) *Query {
	q.columns = appendSep(q.columns, ", ")
	q.columns = q.FormatQuery(q.columns, expr, params...)
	return q
}

func (q *Query) Set(set string, params ...interface{}) *Query {
	q.set = appendSep(q.set, ", ")
	q.set = q.FormatQuery(q.set, set, params...)
	return q
}

func (q *Query) Where(where string, params ...interface{}) *Query {
	q.where = appendSep(q.where, " AND ")
	q.where = append(q.where, '(')
	q.where = q.FormatQuery(q.where, where, params...)
	q.where = append(q.where, ')')
	return q
}

// WhereOr joins passed conditions using OR operation.
func (q *Query) WhereOr(conditions ...*SQL) *Query {
	q.where = appendSep(q.where, " AND ")
	q.where = append(q.where, '(')
	for i, cond := range conditions {
		q.where = cond.AppendFormat(q.where, q)
		if i != len(conditions)-1 {
			q.where = append(q.where, " OR "...)
		}
	}
	q.where = append(q.where, ')')
	return q
}

func (q *Query) Join(join string, params ...interface{}) *Query {
	q.join = appendSep(q.join, " ")
	q.join = q.FormatQuery(q.join, join, params...)
	return q
}

func (q *Query) Group(group string, params ...interface{}) *Query {
	q.group = appendSep(q.group, ", ")
	q.group = q.FormatQuery(q.group, group, params...)
	return q
}

func (q *Query) Order(order string, params ...interface{}) *Query {
	q.order = appendSep(q.order, ", ")
	q.order = q.FormatQuery(q.order, order, params...)
	return q
}

func (q *Query) Limit(n int) *Query {
	q.limit = n
	return q
}

func (q *Query) Offset(n int) *Query {
	q.offset = n
	return q
}

func (q *Query) OnConflict(s string, params ...interface{}) *Query {
	q.onConflict = q.FormatQuery(nil, s, params...)
	return q
}

func (q *Query) Returning(columns ...interface{}) *Query {
	for _, column := range columns {
		q.returning = appendSep(q.returning, ", ")

		switch column := column.(type) {
		case string:
			q.returning = types.AppendField(q.returning, column, 1)
		case types.ValueAppender:
			var err error
			q.returning, err = column.AppendValue(q.returning, 1)
			if err != nil {
				q.setErr(err)
			}
		default:
			q.setErr(fmt.Errorf("unsupported column type: %T", column))
		}
	}
	return q
}

// Apply calls the fn passing the Query as an argument.
func (q *Query) Apply(fn func(*Query) *Query) *Query {
	return fn(q)
}

// Count returns number of rows matching the query using count aggregate function.
func (q *Query) Count() (int, error) {
	if q.err != nil {
		return 0, q.err
	}

	q = q.copy()
	q.columns = types.Q("count(*)")
	q.order = nil
	q.limit = 0
	q.offset = 0

	sel := &selectQuery{
		Query: q,
	}
	var count int
	_, err := q.db.QueryOne(Scan(&count), sel, q.model)
	return count, err
}

// First selects the first row.
func (q *Query) First() error {
	b := columns(q.model.Table().Alias, "", q.model.Table().PKs)
	return q.Order(string(b)).Limit(1).Select()
}

// Last selects the last row.
func (q *Query) Last() error {
	b := columns(q.model.Table().Alias, "", q.model.Table().PKs)
	b = append(b, " DESC"...)
	return q.Order(string(b)).Limit(1).Select()
}

// Select selects the model.
func (q *Query) Select(values ...interface{}) error {
	if q.err != nil {
		return q.err
	}

	q.joinHasOne()
	sel := selectQuery{q}

	var model Model
	var err error
	if len(values) > 0 {
		model, err = NewModel(values...)
		if err != nil {
			return err
		}
	} else {
		model = q.model
	}

	if m, ok := model.(useQueryOne); ok && m.useQueryOne() {
		_, err = q.db.QueryOne(model, sel, q.model)
	} else {
		_, err = q.db.Query(model, sel, q.model)
	}
	if err != nil {
		return err
	}

	if q.model != nil {
		return selectJoins(q.db, q.model.GetJoins())
	}
	return nil
}

// SelectAndCount runs Select and Count in two separate goroutines,
// waits for them to finish and returns the result.
func (q *Query) SelectAndCount(values ...interface{}) (count int, err error) {
	var wg sync.WaitGroup
	wg.Add(2)

	go func() {
		defer wg.Done()
		if e := q.Select(values...); e != nil {
			err = e
		}
	}()

	go func() {
		defer wg.Done()
		var e error
		count, e = q.Count()
		if e != nil {
			err = e
		}
	}()

	wg.Wait()
	return count, err
}

func (q *Query) joinHasOne() {
	if q.model == nil {
		return
	}
	joins := q.model.GetJoins()
	for i := range joins {
		j := &joins[i]
		if j.Rel.One {
			j.JoinOne(q)
		}
	}
}

func selectJoins(db dber, joins []join) error {
	var err error
	for i := range joins {
		j := &joins[i]
		if j.Rel.One {
			err = selectJoins(db, j.JoinModel.GetJoins())
		} else {
			err = j.Select(db)
		}
		if err != nil {
			return err
		}
	}
	return nil
}

// Create inserts the model.
func (q *Query) Create(values ...interface{}) (*types.Result, error) {
	if q.err != nil {
		return nil, q.err
	}

	var model Model
	if len(values) > 0 {
		model = Scan(values...)
	} else {
		model = q.model
	}

	ins := insertQuery{Query: q}
	return q.db.Query(model, ins, q.model)
}

// SelectOrCreate selects the model creating one if it does not exist.
func (q *Query) SelectOrCreate(values ...interface{}) (created bool, err error) {
	if q.err != nil {
		return false, q.err
	}

	for i := 0; i < 10; i++ {
		err := q.Select(values...)
		if err == nil {
			return false, nil
		}
		if err != internal.ErrNoRows {
			return false, err
		}

		res, err := q.Create(values...)
		if err != nil {
			if pgErr, ok := err.(internal.PGError); ok {
				if pgErr.IntegrityViolation() {
					continue
				}
				if pgErr.Field('C') == "55000" {
					// Retry on "#55000 attempted to delete invisible tuple".
					continue
				}
			}
			return false, err
		}
		if res.Affected() == 1 {
			return true, nil
		}
	}

	return false, errors.New("pg: GetOrCreate does not make progress after 10 iterations")
}

// Update updates the model.
func (q *Query) Update() (*types.Result, error) {
	if q.err != nil {
		return nil, q.err
	}
	return q.db.Query(q.model, updateQuery{q}, q.model)
}

// Delete deletes the model.
func (q *Query) Delete() (*types.Result, error) {
	if q.err != nil {
		return nil, q.err
	}
	return q.db.Exec(deleteQuery{q}, q.model)
}

func (q *Query) FormatQuery(dst []byte, query string, params ...interface{}) []byte {
	params = append(params, q.model)
	return q.db.FormatQuery(dst, query, params...)
}

func (q *Query) appendTableAlias(b []byte) []byte {
	if q.tableAlias != "" {
		return types.AppendField(b, q.tableAlias, 1)
	}
	return append(b, q.model.Table().Alias...)
}

func (q *Query) appendTableNameWithAlias(b []byte) []byte {
	b = append(b, q.tableName...)
	b = append(b, " AS "...)
	b = q.appendTableAlias(b)
	return b
}

func (q *Query) haveTables() bool {
	return q.model != nil || len(q.tables) > 0
}

func (q *Query) appendTables(b []byte) []byte {
	if q.model != nil {
		b = q.appendTableNameWithAlias(b)
		if len(q.tables) > 0 {
			b = append(b, ", "...)
		}
	}
	b = append(b, q.tables...)
	return b
}

func (q *Query) appendSet(b []byte) ([]byte, error) {
	b = append(b, " SET "...)
	if len(q.set) > 0 {
		b = append(b, q.set...)
		return b, nil
	}

	if q.model == nil {
		return nil, errors.New("pg: Model(nil)")
	}

	table := q.model.Table()
	strct := q.model.Value()

	if len(q.fields) > 0 {
		for i, fieldName := range q.fields {
			field, err := table.GetField(fieldName)
			if err != nil {
				return nil, err
			}

			b = append(b, field.ColName...)
			b = append(b, " = "...)
			b = field.AppendValue(b, strct, 1)
			if i != len(q.fields)-1 {
				b = append(b, ", "...)
			}
		}
		return b, nil
	}

	start := len(b)
	for _, field := range table.Fields {
		if field.Has(PrimaryKeyFlag) {
			continue
		}

		b = append(b, field.ColName...)
		b = append(b, " = "...)
		b = field.AppendValue(b, strct, 1)
		b = append(b, ", "...)
	}
	if len(b) > start {
		b = b[:len(b)-2]
	}
	return b, nil
}

func (q *Query) appendWhere(b []byte) ([]byte, error) {
	b = append(b, " WHERE "...)
	if len(q.where) > 0 {
		b = append(b, q.where...)
		return b, nil
	}

	if q.model == nil {
		return nil, errors.New("pg: Model(nil)")
	}

	table := q.model.Table()
	if err := table.checkPKs(); err != nil {
		return nil, err
	}
	b = appendColumnAndValue(b, q.model.Value(), table, table.PKs)
	return b, nil
}
<|endoftext|>