<commit_before>package commands

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"strings"

	key "github.com/ipfs/go-ipfs/blocks/key"
	cmds "github.com/ipfs/go-ipfs/commands"
	"github.com/ipfs/go-ipfs/core"
	dag "github.com/ipfs/go-ipfs/merkledag"
	path "github.com/ipfs/go-ipfs/path"
	u "gx/ipfs/QmZNVWh8LLjAavuQ2JXuFmuYH3C11xo988vSgp7UQrTRj1/go-ipfs-util"
	context "gx/ipfs/QmZy2y8t9zQH2a1b8q2ZSLKp17ATuJoCNxxyMFG5qFExpt/go-net/context"
)

// KeyList is a general type for outputting lists of keys
type KeyList struct {
	Keys []key.Key
}

// KeyListTextMarshaler outputs a KeyList as plaintext, one key per line
func KeyListTextMarshaler(res cmds.Response) (io.Reader, error) {
	output := res.Output().(*KeyList)
	buf := new(bytes.Buffer)
	for _, key := range output.Keys {
		buf.WriteString(key.B58String() + "\n")
	}
	return buf, nil
}

var RefsCmd = &cmds.Command{
	Helptext: cmds.HelpText{
		Tagline: "Lists links (references) from an object.",
		ShortDescription: `
Lists the hashes of all the links an IPFS or IPNS object(s) contains,
with the following format:

  <link base58 hash>

NOTE: List all references recursively by using the flag '-r'.
`,
	},
	Subcommands: map[string]*cmds.Command{
		"local": RefsLocalCmd,
	},
	Arguments: []cmds.Argument{
		cmds.StringArg("ipfs-path", true, true, "Path to the object(s) to list refs from.").EnableStdin(),
	},
	Options: []cmds.Option{
		cmds.StringOption("format", "Emit edges with given format. Available tokens: <src> <dst> <linkname>."),
		cmds.BoolOption("edges", "e", "Emit edge format: `<from> -> <to>`.").Default(false),
		cmds.BoolOption("unique", "u", "Omit duplicate refs from output.").Default(false),
		cmds.BoolOption("recursive", "r", "Recursively list links of child nodes.").Default(false),
	},
	Run: func(req cmds.Request, res cmds.Response) {
		ctx := req.Context()
		n, err := req.InvocContext().GetNode()
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		unique, _, err := req.Option("unique").Bool()
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		recursive, _, err := req.Option("recursive").Bool()
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		edges, _, err := req.Option("edges").Bool()
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		format, _, err := req.Option("format").String()
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		objs, err := objectsForPaths(ctx, n, req.Arguments())
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		out := make(chan interface{})
		res.SetOutput((<-chan interface{})(out))

		go func() {
			defer close(out)

			rw := RefWriter{
				out:       out,
				DAG:       n.DAG,
				Ctx:       ctx,
				Unique:    unique,
				PrintEdge: edges,
				PrintFmt:  format,
				Recursive: recursive,
			}

			for _, o := range objs {
				if _, err := rw.WriteRefs(o); err != nil {
					out <- &RefWrapper{Err: err.Error()}
					return
				}
			}
		}()
	},
	Marshalers: cmds.MarshalerMap{
		cmds.Text: func(res cmds.Response) (io.Reader, error) {
			outChan, ok := res.Output().(<-chan interface{})
			if !ok {
				return nil, u.ErrCast()
			}

			marshal := func(v interface{}) (io.Reader, error) {
				obj, ok := v.(*RefWrapper)
				if !ok {
					fmt.Println("%#v", v)
					return nil, u.ErrCast()
				}

				if obj.Err != "" {
					return nil, errors.New(obj.Err)
				}

				return strings.NewReader(obj.Ref + "\n"), nil
			}

			return &cmds.ChannelMarshaler{
				Channel:   outChan,
				Marshaler: marshal,
				Res:       res,
			}, nil
		},
	},
	Type: RefWrapper{},
}

var RefsLocalCmd = &cmds.Command{
	Helptext: cmds.HelpText{
		Tagline: "Lists all local references.",
		ShortDescription: `
Displays the hashes of all local objects.
`,
	},

	Run: func(req cmds.Request, res cmds.Response) {
		ctx := req.Context()
		n, err := req.InvocContext().GetNode()
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		// todo: make async
		allKeys, err := n.Blockstore.AllKeysChan(ctx)
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		piper, pipew := io.Pipe()

		go func() {
			defer pipew.Close()

			for k := range allKeys {
				s := k.B58String() + "\n"
				if _, err := pipew.Write([]byte(s)); err != nil {
					log.Error("pipe write error: ", err)
					return
				}
			}
		}()

		res.SetOutput(piper)
	},
}

func objectsForPaths(ctx context.Context, n *core.IpfsNode, paths []string) ([]*dag.Node, error) {
	objects := make([]*dag.Node, len(paths))
	for i, p := range paths {
		o, err := core.Resolve(ctx, n, path.Path(p))
		if err != nil {
			return nil, err
		}
		objects[i] = o
	}
	return objects, nil
}

type RefWrapper struct {
	Ref string
	Err string
}

type RefWriter struct {
	out chan interface{}
	DAG dag.DAGService
	Ctx context.Context

	Unique    bool
	Recursive bool
	PrintEdge bool
	PrintFmt  string

	seen map[key.Key]struct{}
}

// WriteRefs writes refs of the given object to the underlying writer.
func (rw *RefWriter) WriteRefs(n *dag.Node) (int, error) {
	if rw.Recursive {
		return rw.writeRefsRecursive(n)
	}
	return rw.writeRefsSingle(n)
}

func (rw *RefWriter) writeRefsRecursive(n *dag.Node) (int, error) {
	nkey, err := n.Key()
	if err != nil {
		return 0, err
	}

	var count int
	for i, ng := range dag.GetDAG(rw.Ctx, rw.DAG, n) {
		lk := key.Key(n.Links[i].Hash)
		if rw.skip(lk) {
			continue
		}

		if err := rw.WriteEdge(nkey, lk, n.Links[i].Name); err != nil {
			return count, err
		}

		nd, err := ng.Get(rw.Ctx)
		if err != nil {
			return count, err
		}

		c, err := rw.writeRefsRecursive(nd)
		count += c
		if err != nil {
			return count, err
		}
	}
	return count, nil
}

func (rw *RefWriter) writeRefsSingle(n *dag.Node) (int, error) {
	nkey, err := n.Key()
	if err != nil {
		return 0, err
	}

	if rw.skip(nkey) {
		return 0, nil
	}

	count := 0
	for _, l := range n.Links {
		lk := key.Key(l.Hash)

		if rw.skip(lk) {
			continue
		}

		if err := rw.WriteEdge(nkey, lk, l.Name); err != nil {
			return count, err
		}
		count++
	}
	return count, nil
}

// skip returns whether to skip a key
func (rw *RefWriter) skip(k key.Key) bool {
	if !rw.Unique {
		return false
	}

	if rw.seen == nil {
		rw.seen = make(map[key.Key]struct{})
	}

	_, found := rw.seen[k]
	if !found {
		rw.seen[k] = struct{}{}
	}
	return found
}

// Write one edge
func (rw *RefWriter) WriteEdge(from, to key.Key, linkname string) error {
	if rw.Ctx != nil {
		select {
		case <-rw.Ctx.Done(): // just in case.
			return rw.Ctx.Err()
		default:
		}
	}

	var s string
	switch {
	case rw.PrintFmt != "":
		s = rw.PrintFmt
		s = strings.Replace(s, "<src>", from.B58String(), -1)
		s = strings.Replace(s, "<dst>", to.B58String(), -1)
		s = strings.Replace(s, "<linkname>", linkname, -1)
	case rw.PrintEdge:
		s = from.B58String() + " -> " + to.B58String()
	default:
		s += to.B58String()
	}

	rw.out <- &RefWrapper{Ref: s}
	return nil
}
<commit_msg>Added Default logic to refs<commit_after>package commands

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"strings"

	key "github.com/ipfs/go-ipfs/blocks/key"
	cmds "github.com/ipfs/go-ipfs/commands"
	"github.com/ipfs/go-ipfs/core"
	dag "github.com/ipfs/go-ipfs/merkledag"
	path "github.com/ipfs/go-ipfs/path"
	u "gx/ipfs/QmZNVWh8LLjAavuQ2JXuFmuYH3C11xo988vSgp7UQrTRj1/go-ipfs-util"
	context "gx/ipfs/QmZy2y8t9zQH2a1b8q2ZSLKp17ATuJoCNxxyMFG5qFExpt/go-net/context"
)

// KeyList is a general type for outputting lists of keys
type KeyList struct {
	Keys []key.Key
}

// KeyListTextMarshaler outputs a KeyList as plaintext, one key per line
func KeyListTextMarshaler(res cmds.Response) (io.Reader, error) {
	output := res.Output().(*KeyList)
	buf := new(bytes.Buffer)
	for _, key := range output.Keys {
		buf.WriteString(key.B58String() + "\n")
	}
	return buf, nil
}

var RefsCmd = &cmds.Command{
	Helptext: cmds.HelpText{
		Tagline: "Lists links (references) from an object.",
		ShortDescription: `
Lists the hashes of all the links an IPFS or IPNS object(s) contains,
with the following format:

  <link base58 hash>

NOTE: List all references recursively by using the flag '-r'.
`,
	},
	Subcommands: map[string]*cmds.Command{
		"local": RefsLocalCmd,
	},
	Arguments: []cmds.Argument{
		cmds.StringArg("ipfs-path", true, true, "Path to the object(s) to list refs from.").EnableStdin(),
	},
	Options: []cmds.Option{
		cmds.StringOption("format", "Emit edges with given format. Available tokens: <src> <dst> <linkname>.").Default("<dst>"),
		cmds.BoolOption("edges", "e", "Emit edge format: `<from> -> <to>`.").Default(false),
		cmds.BoolOption("unique", "u", "Omit duplicate refs from output.").Default(false),
		cmds.BoolOption("recursive", "r", "Recursively list links of child nodes.").Default(false),
	},
	Run: func(req cmds.Request, res cmds.Response) {
		ctx := req.Context()
		n, err := req.InvocContext().GetNode()
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		unique, _, err := req.Option("unique").Bool()
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		recursive, _, err := req.Option("recursive").Bool()
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		edges, _, err := req.Option("edges").Bool()
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		format, _, err := req.Option("format").String()
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		objs, err := objectsForPaths(ctx, n, req.Arguments())
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		out := make(chan interface{})
		res.SetOutput((<-chan interface{})(out))

		go func() {
			defer close(out)

			rw := RefWriter{
				out:       out,
				DAG:       n.DAG,
				Ctx:       ctx,
				Unique:    unique,
				PrintEdge: edges,
				PrintFmt:  format,
				Recursive: recursive,
			}

			for _, o := range objs {
				if _, err := rw.WriteRefs(o); err != nil {
					out <- &RefWrapper{Err: err.Error()}
					return
				}
			}
		}()
	},
	Marshalers: cmds.MarshalerMap{
		cmds.Text: func(res cmds.Response) (io.Reader, error) {
			outChan, ok := res.Output().(<-chan interface{})
			if !ok {
				return nil, u.ErrCast()
			}

			marshal := func(v interface{}) (io.Reader, error) {
				obj, ok := v.(*RefWrapper)
				if !ok {
					fmt.Println("%#v", v)
					return nil, u.ErrCast()
				}

				if obj.Err != "" {
					return nil, errors.New(obj.Err)
				}

				return strings.NewReader(obj.Ref + "\n"), nil
			}

			return &cmds.ChannelMarshaler{
				Channel:   outChan,
				Marshaler: marshal,
				Res:       res,
			}, nil
		},
	},
	Type: RefWrapper{},
}

var RefsLocalCmd = &cmds.Command{
	Helptext: cmds.HelpText{
		Tagline: "Lists all local references.",
		ShortDescription: `
Displays the hashes of all local objects.
`,
	},

	Run: func(req cmds.Request, res cmds.Response) {
		ctx := req.Context()
		n, err := req.InvocContext().GetNode()
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		// todo: make async
		allKeys, err := n.Blockstore.AllKeysChan(ctx)
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		piper, pipew := io.Pipe()

		go func() {
			defer pipew.Close()

			for k := range allKeys {
				s := k.B58String() + "\n"
				if _, err := pipew.Write([]byte(s)); err != nil {
					log.Error("pipe write error: ", err)
					return
				}
			}
		}()

		res.SetOutput(piper)
	},
}

func objectsForPaths(ctx context.Context, n *core.IpfsNode, paths []string) ([]*dag.Node, error) {
	objects := make([]*dag.Node, len(paths))
	for i, p := range paths {
		o, err := core.Resolve(ctx, n, path.Path(p))
		if err != nil {
			return nil, err
		}
		objects[i] = o
	}
	return objects, nil
}

type RefWrapper struct {
	Ref string
	Err string
}

type RefWriter struct {
	out chan interface{}
	DAG dag.DAGService
	Ctx context.Context

	Unique    bool
	Recursive bool
	PrintEdge bool
	PrintFmt  string

	seen map[key.Key]struct{}
}

// WriteRefs writes refs of the given object to the underlying writer.
func (rw *RefWriter) WriteRefs(n *dag.Node) (int, error) {
	if rw.Recursive {
		return rw.writeRefsRecursive(n)
	}
	return rw.writeRefsSingle(n)
}

func (rw *RefWriter) writeRefsRecursive(n *dag.Node) (int, error) {
	nkey, err := n.Key()
	if err != nil {
		return 0, err
	}

	var count int
	for i, ng := range dag.GetDAG(rw.Ctx, rw.DAG, n) {
		lk := key.Key(n.Links[i].Hash)
		if rw.skip(lk) {
			continue
		}

		if err := rw.WriteEdge(nkey, lk, n.Links[i].Name); err != nil {
			return count, err
		}

		nd, err := ng.Get(rw.Ctx)
		if err != nil {
			return count, err
		}

		c, err := rw.writeRefsRecursive(nd)
		count += c
		if err != nil {
			return count, err
		}
	}
	return count, nil
}

func (rw *RefWriter) writeRefsSingle(n *dag.Node) (int, error) {
	nkey, err := n.Key()
	if err != nil {
		return 0, err
	}

	if rw.skip(nkey) {
		return 0, nil
	}

	count := 0
	for _, l := range n.Links {
		lk := key.Key(l.Hash)

		if rw.skip(lk) {
			continue
		}

		if err := rw.WriteEdge(nkey, lk, l.Name); err != nil {
			return count, err
		}
		count++
	}
	return count, nil
}

// skip returns whether to skip a key
func (rw *RefWriter) skip(k key.Key) bool {
	if !rw.Unique {
		return false
	}

	if rw.seen == nil {
		rw.seen = make(map[key.Key]struct{})
	}

	_, found := rw.seen[k]
	if !found {
		rw.seen[k] = struct{}{}
	}
	return found
}

// Write one edge
func (rw *RefWriter) WriteEdge(from, to key.Key, linkname string) error {
	if rw.Ctx != nil {
		select {
		case <-rw.Ctx.Done(): // just in case.
			return rw.Ctx.Err()
		default:
		}
	}

	var s string
	switch {
	case rw.PrintFmt != "":
		s = rw.PrintFmt
		s = strings.Replace(s, "<src>", from.B58String(), -1)
		s = strings.Replace(s, "<dst>", to.B58String(), -1)
		s = strings.Replace(s, "<linkname>", linkname, -1)
	case rw.PrintEdge:
		s = from.B58String() + " -> " + to.B58String()
	default:
		s += to.B58String()
	}

	rw.out <- &RefWrapper{Ref: s}
	return nil
}
<|endoftext|>


<commit_before>package core

import (
	"context"
	"crypto/md5"
	"encoding/hex"
	"os"
	"runtime"
	"testing"
	"time"

	"chain/core/account"
	"chain/core/asset"
	"chain/core/asset/assettest"
	"chain/core/query"
	"chain/core/txbuilder"
	"chain/core/txdb"
	"chain/database/pg"
	"chain/database/pg/pgtest"
	"chain/database/sql"
	"chain/protocol"
	"chain/protocol/bc"
	"chain/protocol/prottest"
)

// TestRecovery tests end-to-end blockchain recovery from an exit
// during landing a block.
func TestRecovery(t *testing.T) {
	if os.Getenv("LONG") == "" {
		t.Skip("skipping test; $LONG not set")
	}

	// Setup the test environment using a clean db.
	ctx := context.Background()
	dbURL, db := pgtest.NewDB(t, pgtest.SchemaPath)
	store, pool := txdb.New(db)
	setupCtx := pg.NewContext(ctx, db)
	c := prottest.NewChainWithStorage(t, store, pool)
	asset.Init(c, nil)
	account.Init(c, nil)

	// Setup the transaction query indexer to index every transaction.
	indexer := query.NewIndexer(db, c)
	indexer.RegisterAnnotator(account.AnnotateTxs)
	indexer.RegisterAnnotator(asset.AnnotateTxs)

	// Create two assets (USD & apples) and two accounts (Alice & Bob).
	var (
		usdTags = map[string]interface{}{"currency": "usd"}
		usd     = assettest.CreateAssetFixture(setupCtx, t, nil, 0, nil, "usd", usdTags)
		apple   = assettest.CreateAssetFixture(setupCtx, t, nil, 0, nil, "apple", nil)
		alice   = assettest.CreateAccountFixture(setupCtx, t, nil, 0, "alice", nil)
		bob     = assettest.CreateAccountFixture(setupCtx, t, nil, 0, "bob", nil)
	)
	// Issue some apples to Alice and a dollar to Bob.
	_ = assettest.IssueAssetsFixture(setupCtx, t, c, apple, 10, alice)
	_ = assettest.IssueAssetsFixture(setupCtx, t, c, usd, 1, bob)

	prottest.MakeBlock(setupCtx, t, c)

	// Submit a transfer between Alice and Bob but don't publish it in a block.
	assettest.Transfer(setupCtx, t, c, []txbuilder.Action{
		assettest.NewAccountControlAction(bc.AssetAmount{AssetID: usd, Amount: 1}, alice, nil),
		assettest.NewAccountControlAction(bc.AssetAmount{AssetID: apple, Amount: 1}, bob, nil),
		assettest.NewAccountSpendAction(bc.AssetAmount{AssetID: usd, Amount: 1}, bob, nil, nil, nil),
		assettest.NewAccountSpendAction(bc.AssetAmount{AssetID: apple, Amount: 1}, alice, nil, nil, nil),
	})

	err := db.Close()
	if err != nil {
		t.Fatal(err)
	}

	// Run the `generateBlock` function repeatedly, each time starting
	// from the state of the test database initialized above. Each time,
	// the SQL driver will force the goroutine to exit at a different
	// query. After each forced exit, run `generateBlock` again to recover.
	// Despite crashing at a different point on each run, the resulting
	// states should be identical.
	timestamp := time.Now()
	var databaseDumps []string
	for n := 1; ; n++ {
		// Create a new, fresh database using the schema file we created
		// above. The database will be identical.
		cloneURL, err := pgtest.CloneDB(ctx, dbURL)
		if err != nil {
			t.Fatal(err)
		}

		// Open a new handle to the same database but using a driver that
		// will call our anonymous func to simulate crashes.
		ch := make(chan error)
		calls := 0
		wrappedDB := pgtest.WrapDB(t, cloneURL, func(q string) {
			calls++

			if calls == n {
				t.Logf("crashing on query %d: %s\n", calls, q)
				close(ch)        // let main goroutine know
				runtime.Goexit() // crash this goroutine
			}
		})

		ctx := context.Background()
		ctx = pg.NewContext(ctx, wrappedDB)
		go func() {
			err := generateBlock(ctx, wrappedDB, timestamp)
			ch <- err
		}()

		// Wait for the goroutine to finish or get killed on the N-th query.
		err, ok := <-ch
		if err != nil {
			t.Fatal(err)
		}
		if ok {
			// The driver never crashed the goroutine, so n is now greater than
			// the total number of queries performed during `generateBlock`.
			databaseDumps = append(databaseDumps, pgtest.Dump(t, cloneURL, false, "pool_txs"))
			break
		}

		// We crashed at some point during block generation. Do it again,
		// without crashing.
		err = generateBlock(ctx, wrappedDB, timestamp)
		if err != nil {
			t.Fatal(err)
		}
		databaseDumps = append(databaseDumps, pgtest.Dump(t, cloneURL, false, "pool_txs"))
	}

	if len(databaseDumps) < 2 {
		t.Fatal("no database dumps; did the wrapped driver get used?")
	}
	// Compare all of the pg_dumps. They should all be equal. Only
	// print the hashes though so that the test output isn't overwhelming.
	var prevHash string
	for n, v := range databaseDumps {
		hasher := md5.New()
		hasher.Write([]byte(v))
		hash := hex.EncodeToString(hasher.Sum(nil))
		if n >= 1 && prevHash != hash {
			t.Errorf("previous run %d - %s; current run %d - %s", n, prevHash, n+1, hash)
		}
		prevHash = hash
	}
}

func generateBlock(ctx context.Context, db *sql.DB, timestamp time.Time) error {
	store, pool := txdb.New(db)
	c, err := protocol.NewChain(ctx, store, pool, nil)
	if err != nil {
		return err
	}

	asset.Init(c, nil)
	account.Init(c, nil)
	// Setup the transaction query indexer to index every transaction.
	indexer := query.NewIndexer(db, c)
	indexer.RegisterAnnotator(account.AnnotateTxs)
	indexer.RegisterAnnotator(asset.AnnotateTxs)

	block, snapshot, err := c.Recover(ctx)
	if err != nil {
		return err
	}

	b, s, err := c.GenerateBlock(ctx, block, snapshot, timestamp)
	if err != nil {
		return err
	}
	if len(b.Transactions) == 0 {
		return nil
	}
	return c.CommitBlock(ctx, b, s)
}
<commit_msg>core: ignore sequences in recovery test<commit_after>package core

import (
	"context"
	"crypto/md5"
	"encoding/hex"
	"os"
	"runtime"
	"testing"
	"time"

	"chain/core/account"
	"chain/core/asset"
	"chain/core/asset/assettest"
	"chain/core/query"
	"chain/core/txbuilder"
	"chain/core/txdb"
	"chain/database/pg"
	"chain/database/pg/pgtest"
	"chain/database/sql"
	"chain/protocol"
	"chain/protocol/bc"
	"chain/protocol/prottest"
)

// TestRecovery tests end-to-end blockchain recovery from an exit
// during landing a block.
func TestRecovery(t *testing.T) {
	if os.Getenv("LONG") == "" {
		t.Skip("skipping test; $LONG not set")
	}

	// Setup the test environment using a clean db.
	ctx := context.Background()
	dbURL, db := pgtest.NewDB(t, pgtest.SchemaPath)
	store, pool := txdb.New(db)
	setupCtx := pg.NewContext(ctx, db)
	c := prottest.NewChainWithStorage(t, store, pool)
	asset.Init(c, nil)
	account.Init(c, nil)

	// Setup the transaction query indexer to index every transaction.
	indexer := query.NewIndexer(db, c)
	indexer.RegisterAnnotator(account.AnnotateTxs)
	indexer.RegisterAnnotator(asset.AnnotateTxs)

	// Create two assets (USD & apples) and two accounts (Alice & Bob).
	var (
		usdTags = map[string]interface{}{"currency": "usd"}
		usd     = assettest.CreateAssetFixture(setupCtx, t, nil, 0, nil, "usd", usdTags)
		apple   = assettest.CreateAssetFixture(setupCtx, t, nil, 0, nil, "apple", nil)
		alice   = assettest.CreateAccountFixture(setupCtx, t, nil, 0, "alice", nil)
		bob     = assettest.CreateAccountFixture(setupCtx, t, nil, 0, "bob", nil)
	)
	// Issue some apples to Alice and a dollar to Bob.
	_ = assettest.IssueAssetsFixture(setupCtx, t, c, apple, 10, alice)
	_ = assettest.IssueAssetsFixture(setupCtx, t, c, usd, 1, bob)

	prottest.MakeBlock(setupCtx, t, c)

	// Submit a transfer between Alice and Bob but don't publish it in a block.
	assettest.Transfer(setupCtx, t, c, []txbuilder.Action{
		assettest.NewAccountControlAction(bc.AssetAmount{AssetID: usd, Amount: 1}, alice, nil),
		assettest.NewAccountControlAction(bc.AssetAmount{AssetID: apple, Amount: 1}, bob, nil),
		assettest.NewAccountSpendAction(bc.AssetAmount{AssetID: usd, Amount: 1}, bob, nil, nil, nil),
		assettest.NewAccountSpendAction(bc.AssetAmount{AssetID: apple, Amount: 1}, alice, nil, nil, nil),
	})

	err := db.Close()
	if err != nil {
		t.Fatal(err)
	}

	// Run the `generateBlock` function repeatedly, each time starting
	// from the state of the test database initialized above. Each time,
	// the SQL driver will force the goroutine to exit at a different
	// query. After each forced exit, run `generateBlock` again to recover.
	// Despite crashing at a different point on each run, the resulting
	// states should be identical.
	timestamp := time.Now()
	var databaseDumps []string
	for n := 1; ; n++ {
		// Create a new, fresh database using the schema file we created
		// above. The database will be identical.
		cloneURL, err := pgtest.CloneDB(ctx, dbURL)
		if err != nil {
			t.Fatal(err)
		}

		// Open a new handle to the same database but using a driver that
		// will call our anonymous func to simulate crashes.
		ch := make(chan error)
		calls := 0
		wrappedDB := pgtest.WrapDB(t, cloneURL, func(q string) {
			calls++

			if calls == n {
				t.Logf("crashing on query %d: %s\n", calls, q)
				close(ch)        // let main goroutine know
				runtime.Goexit() // crash this goroutine
			}
		})

		ctx := context.Background()
		ctx = pg.NewContext(ctx, wrappedDB)
		go func() {
			err := generateBlock(ctx, wrappedDB, timestamp)
			ch <- err
		}()

		// Wait for the goroutine to finish or get killed on the N-th query.
		err, ok := <-ch
		if err != nil {
			t.Fatal(err)
		}
		if ok {
			// The driver never crashed the goroutine, so n is now greater than
			// the total number of queries performed during `generateBlock`.
			databaseDumps = append(databaseDumps, pgtest.Dump(t, cloneURL, false, "pool_txs", "*_id_seq"))
			break
		}

		// We crashed at some point during block generation. Do it again,
		// without crashing.
		err = generateBlock(ctx, wrappedDB, timestamp)
		if err != nil {
			t.Fatal(err)
		}
		databaseDumps = append(databaseDumps, pgtest.Dump(t, cloneURL, false, "pool_txs", "*_id_seq"))
	}

	if len(databaseDumps) < 2 {
		t.Fatal("no database dumps; did the wrapped driver get used?")
	}
	// Compare all of the pg_dumps. They should all be equal. Only
	// print the hashes though so that the test output isn't overwhelming.
	var prevHash string
	for n, v := range databaseDumps {
		hasher := md5.New()
		hasher.Write([]byte(v))
		hash := hex.EncodeToString(hasher.Sum(nil))
		if n >= 1 && prevHash != hash {
			t.Errorf("previous run %d - %s; current run %d - %s", n, prevHash, n+1, hash)
		}
		prevHash = hash
	}
}

func generateBlock(ctx context.Context, db *sql.DB, timestamp time.Time) error {
	store, pool := txdb.New(db)
	c, err := protocol.NewChain(ctx, store, pool, nil)
	if err != nil {
		return err
	}

	asset.Init(c, nil)
	account.Init(c, nil)
	// Setup the transaction query indexer to index every transaction.
	indexer := query.NewIndexer(db, c)
	indexer.RegisterAnnotator(account.AnnotateTxs)
	indexer.RegisterAnnotator(asset.AnnotateTxs)

	block, snapshot, err := c.Recover(ctx)
	if err != nil {
		return err
	}

	b, s, err := c.GenerateBlock(ctx, block, snapshot, timestamp)
	if err != nil {
		return err
	}
	if len(b.Transactions) == 0 {
		return nil
	}
	return c.CommitBlock(ctx, b, s)
}
<|endoftext|>


<commit_before>package riak

import (
	rpbRiakDT "github.com/basho-labs/riak-go-client/rpb/riak_dt"
	"reflect"
	"testing"
	"time"
)

// UpdateCounter
// DtUpdateReq

func TestBuildDtUpdateReqCorrectlyViaUpdateCounterCommandBuilder(t *testing.T) {
	builder := NewUpdateCounterCommandBuilder().
		WithBucketType("counters").
		WithBucket("myBucket").
		WithKey("counter_1").
		WithIncrement(100).
		WithW(3).
		WithPw(1).
		WithDw(2).
		WithReturnBody(true).
		WithTimeout(time.Second * 20)
	cmd, err := builder.Build()
	if err != nil {
		t.Fatal(err.Error())
	}
	protobuf, err := cmd.constructPbRequest()
	if err != nil {
		t.Fatal(err.Error())
	}
	if protobuf == nil {
		t.FailNow()
	}
	if req, ok := protobuf.(*rpbRiakDT.DtUpdateReq); ok {
		if expected, actual := "counters", string(req.GetType()); expected != actual {
			t.Errorf("expected %v, got %v", expected, actual)
		}
		if expected, actual := "myBucket", string(req.GetBucket()); expected != actual {
			t.Errorf("expected %v, got %v", expected, actual)
		}
		if expected, actual := "counter_1", string(req.GetKey()); expected != actual {
			t.Errorf("expected %v, got %v", expected, actual)
		}
		if expected, actual := uint32(3), req.GetW(); expected != actual {
			t.Errorf("expected %v, got %v", expected, actual)
		}
		if expected, actual := uint32(1), req.GetPw(); expected != actual {
			t.Errorf("expected %v, got %v", expected, actual)
		}
		if expected, actual := uint32(2), req.GetDw(); expected != actual {
			t.Errorf("expected %v, got %v", expected, actual)
		}
		if expected, actual := true, req.GetReturnBody(); expected != actual {
			t.Errorf("expected %v, got %v", expected, actual)
		}
		op := req.Op.CounterOp
		if expected, actual := int64(100), op.GetIncrement(); expected != actual {
			t.Errorf("expected %v, got %v", expected, actual)
		}
		validateTimeout(t, time.Second*20, req.GetTimeout())
	} else {
		t.Errorf("ok: %v - could not convert %v to *rpbRiakDT.DtUpdateReq", ok, reflect.TypeOf(protobuf))
	}
}

func TestUpdateCounterParsesDtUpdateRespCorrectly(t *testing.T) {
	counterValue := int64(1234)
	generatedKey := "generated_key"
	dtUpdateResp := &rpbRiakDT.DtUpdateResp{
		CounterValue: &counterValue,
		Key:          []byte(generatedKey),
	}

	builder := NewUpdateCounterCommandBuilder().
		WithBucketType("counters").
		WithBucket("myBucket").
		WithKey("counter_1").
		WithIncrement(100)
	cmd, err := builder.Build()
	if err != nil {
		t.Fatal(err.Error())
	}
	protobuf, err := cmd.constructPbRequest()
	if err != nil {
		t.Fatal(err.Error())
	}
	if protobuf == nil {
		t.FailNow()
	}

	cmd.onSuccess(dtUpdateResp)

	if uc, ok := cmd.(*UpdateCounterCommand); ok {
		rsp := uc.Response
		if expected, actual := int64(1234), rsp.CounterValue; expected != actual {
			t.Errorf("expected %v, got %v", expected, actual)
		}
		if expected, actual := "generated_key", rsp.GeneratedKey; expected != actual {
			t.Errorf("expected %v, got %v", expected, actual)
		}
	} else {
		t.Errorf("ok: %v - could not convert %v to *UpdateCounterCommand", ok, reflect.TypeOf(cmd))
	}
}

func TestValidationOfUpdateCounterViaBuilder(t *testing.T) {
	// validate that Bucket is required
	builder := NewUpdateCounterCommandBuilder()
	_, err := builder.Build()
	if err == nil {
		t.Fatal("expected non-nil err")
	}
	if expected, actual := ErrBucketRequired.Error(), err.Error(); expected != actual {
		t.Errorf("expected %v, actual %v", expected, actual)
	}

	// validate that Key is required
	builder = NewUpdateCounterCommandBuilder()
	builder.WithBucket("bucket_name")
	_, err = builder.Build()
	if err == nil {
		t.Fatal("expected non-nil err")
	}
	if expected, actual := ErrKeyRequired.Error(), err.Error(); expected != actual {
		t.Errorf("expected %v, actual %v", expected, actual)
	}
}

// FetchCounter
// DtFetchReq

func TestBuildDtFetchReqCorrectlyViaFetchCounterCommandBuilder(t *testing.T) {
	builder := NewFetchCounterCommandBuilder().
		WithBucketType("counters").
		WithBucket("myBucket").
		WithKey("counter_1").
		WithR(3).
		WithPr(1).
		WithNotFoundOk(true).
		WithBasicQuorum(true).
		WithTimeout(time.Second * 20)
	cmd, err := builder.Build()
	if err != nil {
		t.Fatal(err.Error())
	}
	protobuf, err := cmd.constructPbRequest()
	if err != nil {
		t.Fatal(err.Error())
	}
	if protobuf == nil {
		t.FailNow()
	}
	if req, ok := protobuf.(*rpbRiakDT.DtFetchReq); ok {
		if expected, actual := "counters", string(req.GetType()); expected != actual {
			t.Errorf("expected %v, got %v", expected, actual)
		}
		if expected, actual := "myBucket", string(req.GetBucket()); expected != actual {
			t.Errorf("expected %v, got %v", expected, actual)
		}
		if expected, actual := "counter_1", string(req.GetKey()); expected != actual {
			t.Errorf("expected %v, got %v", expected, actual)
		}
		if expected, actual := uint32(3), req.GetR(); expected != actual {
			t.Errorf("expected %v, got %v", expected, actual)
		}
		if expected, actual := uint32(1), req.GetPr(); expected != actual {
			t.Errorf("expected %v, got %v", expected, actual)
		}
		if expected, actual := true, req.GetNotfoundOk(); expected != actual {
			t.Errorf("expected %v, got %v", expected, actual)
		}
		if expected, actual := true, req.GetBasicQuorum(); expected != actual {
			t.Errorf("expected %v, got %v", expected, actual)
		}
		validateTimeout(t, time.Second*20, req.GetTimeout())
	} else {
		t.Errorf("ok: %v - could not convert %v to *rpbRiakDT.DtFetchReq", ok, reflect.TypeOf(protobuf))
	}
}

func TestFetchCounterParsesDtFetchRespCorrectly(t *testing.T) {
	counterValue := int64(1234)
	dtValue := &rpbRiakDT.DtValue{
		CounterValue: &counterValue,
	}
	dtFetchResp := &rpbRiakDT.DtFetchResp{
		Value: dtValue,
	}

	builder := NewFetchCounterCommandBuilder().
		WithBucketType("counters").
		WithBucket("myBucket").
		WithKey("counter_1")
	cmd, err := builder.Build()
	if err != nil {
		t.Fatal(err.Error())
	}
	protobuf, err := cmd.constructPbRequest()
	if err != nil {
		t.Fatal(err.Error())
	}
	if protobuf == nil {
		t.FailNow()
	}

	cmd.onSuccess(dtFetchResp)

	if uc, ok := cmd.(*FetchCounterCommand); ok {
		rsp := uc.Response
		if expected, actual := counterValue, rsp.CounterValue; expected != actual {
			t.Errorf("expected %v, got %v", expected, actual)
		}
	} else {
		t.Errorf("ok: %v - could not convert %v to *FetchCounterCommand", ok, reflect.TypeOf(cmd))
	}
}

func TestValidationOfFetchCounterViaBuilder(t *testing.T) {
	// validate that Bucket is required
	builder := NewFetchCounterCommandBuilder()
	_, err := builder.Build()
	if err == nil {
		t.Fatal("expected non-nil err")
	}
	if expected, actual := ErrBucketRequired.Error(), err.Error(); expected != actual {
		t.Errorf("expected %v, actual %v", expected, actual)
	}

	// validate that Key is required
	builder = NewFetchCounterCommandBuilder()
	builder.WithBucket("bucket_name")
	_, err = builder.Build()
	if err == nil {
		t.Fatal("expected non-nil err")
	}
	if expected, actual := ErrKeyRequired.Error(), err.Error(); expected != actual {
		t.Errorf("expected %v, actual %v", expected, actual)
	}
}
<commit_msg>Fix unit test.<commit_after>package riak

import (
	rpbRiakDT "github.com/basho-labs/riak-go-client/rpb/riak_dt"
	"reflect"
	"testing"
	"time"
)

// UpdateCounter
// DtUpdateReq

func TestBuildDtUpdateReqCorrectlyViaUpdateCounterCommandBuilder(t *testing.T) {
	builder := NewUpdateCounterCommandBuilder().
		WithBucketType("counters").
		WithBucket("myBucket").
		WithKey("counter_1").
		WithIncrement(100).
		WithW(3).
		WithPw(1).
		WithDw(2).
		WithReturnBody(true).
		WithTimeout(time.Second * 20)
	cmd, err := builder.Build()
	if err != nil {
		t.Fatal(err.Error())
	}
	protobuf, err := cmd.constructPbRequest()
	if err != nil {
		t.Fatal(err.Error())
	}
	if protobuf == nil {
		t.FailNow()
	}
	if req, ok := protobuf.(*rpbRiakDT.DtUpdateReq); ok {
		if expected, actual := "counters", string(req.GetType()); expected != actual {
			t.Errorf("expected %v, got %v", expected, actual)
		}
		if expected, actual := "myBucket", string(req.GetBucket()); expected != actual {
			t.Errorf("expected %v, got %v", expected, actual)
		}
		if expected, actual := "counter_1", string(req.GetKey()); expected != actual {
			t.Errorf("expected %v, got %v", expected, actual)
		}
		if expected, actual := uint32(3), req.GetW(); expected != actual {
			t.Errorf("expected %v, got %v", expected, actual)
		}
		if expected, actual := uint32(1), req.GetPw(); expected != actual {
			t.Errorf("expected %v, got %v", expected, actual)
		}
		if expected, actual := uint32(2), req.GetDw(); expected != actual {
			t.Errorf("expected %v, got %v", expected, actual)
		}
		if expected, actual := true, req.GetReturnBody(); expected != actual {
			t.Errorf("expected %v, got %v", expected, actual)
		}
		op := req.Op.CounterOp
		if expected, actual := int64(100), op.GetIncrement(); expected != actual {
			t.Errorf("expected %v, got %v", expected, actual)
		}
		validateTimeout(t, time.Second*20, req.GetTimeout())
	} else {
		t.Errorf("ok: %v - could not convert %v to *rpbRiakDT.DtUpdateReq", ok, reflect.TypeOf(protobuf))
	}
}

func TestUpdateCounterParsesDtUpdateRespCorrectly(t *testing.T) {
	counterValue := int64(1234)
	generatedKey := "generated_key"
	dtUpdateResp := &rpbRiakDT.DtUpdateResp{
		CounterValue: &counterValue,
		Key:          []byte(generatedKey),
	}

	builder := NewUpdateCounterCommandBuilder().
		WithBucketType("counters").
		WithBucket("myBucket").
		WithKey("counter_1").
		WithIncrement(100)
	cmd, err := builder.Build()
	if err != nil {
		t.Fatal(err.Error())
	}
	protobuf, err := cmd.constructPbRequest()
	if err != nil {
		t.Fatal(err.Error())
	}
	if protobuf == nil {
		t.FailNow()
	}

	cmd.onSuccess(dtUpdateResp)

	if uc, ok := cmd.(*UpdateCounterCommand); ok {
		rsp := uc.Response
		if expected, actual := int64(1234), rsp.CounterValue; expected != actual {
			t.Errorf("expected %v, got %v", expected, actual)
		}
		if expected, actual := "generated_key", rsp.GeneratedKey; expected != actual {
			t.Errorf("expected %v, got %v", expected, actual)
		}
	} else {
		t.Errorf("ok: %v - could not convert %v to *UpdateCounterCommand", ok, reflect.TypeOf(cmd))
	}
}

func TestValidationOfUpdateCounterViaBuilder(t *testing.T) {
	// validate that Bucket is required
	builder := NewUpdateCounterCommandBuilder()
	_, err := builder.Build()
	if err == nil {
		t.Fatal("expected non-nil err")
	}
	if expected, actual := ErrBucketRequired.Error(), err.Error(); expected != actual {
		t.Errorf("expected %v, actual %v", expected, actual)
	}

	// validate that Key is NOT required
	builder = NewUpdateCounterCommandBuilder()
	builder.WithBucket("bucket_name")
	_, err = builder.Build()
	if err != nil {
		t.Fatal("expected nil err")
	}
}

// FetchCounter
// DtFetchReq

func TestBuildDtFetchReqCorrectlyViaFetchCounterCommandBuilder(t *testing.T) {
	builder := NewFetchCounterCommandBuilder().
		WithBucketType("counters").
		WithBucket("myBucket").
		WithKey("counter_1").
		WithR(3).
		WithPr(1).
		WithNotFoundOk(true).
		WithBasicQuorum(true).
		WithTimeout(time.Second * 20)
	cmd, err := builder.Build()
	if err != nil {
		t.Fatal(err.Error())
	}
	protobuf, err := cmd.constructPbRequest()
	if err != nil {
		t.Fatal(err.Error())
	}
	if protobuf == nil {
		t.FailNow()
	}
	if req, ok := protobuf.(*rpbRiakDT.DtFetchReq); ok {
		if expected, actual := "counters", string(req.GetType()); expected != actual {
			t.Errorf("expected %v, got %v", expected, actual)
		}
		if expected, actual := "myBucket", string(req.GetBucket()); expected != actual {
			t.Errorf("expected %v, got %v", expected, actual)
		}
		if expected, actual := "counter_1", string(req.GetKey()); expected != actual {
			t.Errorf("expected %v, got %v", expected, actual)
		}
		if expected, actual := uint32(3), req.GetR(); expected != actual {
			t.Errorf("expected %v, got %v", expected, actual)
		}
		if expected, actual := uint32(1), req.GetPr(); expected != actual {
			t.Errorf("expected %v, got %v", expected, actual)
		}
		if expected, actual := true, req.GetNotfoundOk(); expected != actual {
			t.Errorf("expected %v, got %v", expected, actual)
		}
		if expected, actual := true, req.GetBasicQuorum(); expected != actual {
			t.Errorf("expected %v, got %v", expected, actual)
		}
		validateTimeout(t, time.Second*20, req.GetTimeout())
	} else {
		t.Errorf("ok: %v - could not convert %v to *rpbRiakDT.DtFetchReq", ok, reflect.TypeOf(protobuf))
	}
}

func TestFetchCounterParsesDtFetchRespCorrectly(t *testing.T) {
	counterValue := int64(1234)
	dtValue := &rpbRiakDT.DtValue{
		CounterValue: &counterValue,
	}
	dtFetchResp := &rpbRiakDT.DtFetchResp{
		Value: dtValue,
	}

	builder := NewFetchCounterCommandBuilder().
		WithBucketType("counters").
		WithBucket("myBucket").
		WithKey("counter_1")
	cmd, err := builder.Build()
	if err != nil {
		t.Fatal(err.Error())
	}
	protobuf, err := cmd.constructPbRequest()
	if err != nil {
		t.Fatal(err.Error())
	}
	if protobuf == nil {
		t.FailNow()
	}

	cmd.onSuccess(dtFetchResp)

	if uc, ok := cmd.(*FetchCounterCommand); ok {
		rsp := uc.Response
		if expected, actual := counterValue, rsp.CounterValue; expected != actual {
			t.Errorf("expected %v, got %v", expected, actual)
		}
	} else {
		t.Errorf("ok: %v - could not convert %v to *FetchCounterCommand", ok, reflect.TypeOf(cmd))
	}
}

func TestValidationOfFetchCounterViaBuilder(t *testing.T) {
	// validate that Bucket is required
	builder := NewFetchCounterCommandBuilder()
	_, err := builder.Build()
	if err == nil {
		t.Fatal("expected non-nil err")
	}
	if expected, actual := ErrBucketRequired.Error(), err.Error(); expected != actual {
		t.Errorf("expected %v, actual %v", expected, actual)
	}

	// validate that Key is required
	builder = NewFetchCounterCommandBuilder()
	builder.WithBucket("bucket_name")
	_, err = builder.Build()
	if err == nil {
		t.Fatal("expected non-nil err")
	}
	if expected, actual := ErrKeyRequired.Error(), err.Error(); expected != actual {
		t.Errorf("expected %v, actual %v", expected, actual)
	}
}
<|endoftext|>


<commit_before>package mock

import (
	"github.com/icrowley/fake"
	"github.com/rbastic/dyndao/object"
	"github.com/rbastic/dyndao/schema"
)

const PeopleObjectType string = "people"
const AddressesObjectType string = "addresses"

// Basic test mock
func fieldName() *schema.Column {
	fld := schema.DefaultColumn()
	fld.Name = "Name"
	fld.DBType = "text"
	return fld
}

func fieldLookupText() *schema.Column {
	fld := schema.DefaultColumn()
	fld.Name = "LOOKUP_TEXT"
	fld.DBType = "TEXT"
	fld.AllowNull = false
	return fld
}

func fieldNullText() *schema.Column {
	fld := schema.DefaultColumn()
	fld.Name = "NullText"
	fld.DBType = "text"
	fld.AllowNull = true
	return fld
}

func fieldNullInt() *schema.Column {
	fld := schema.DefaultColumn()
	fld.IsNumber = true
	fld.Name = "NullInt"
	fld.DBType = "integer"
	fld.AllowNull = true
	return fld
}

func fieldNullFloat() *schema.Column {
	fld := schema.DefaultColumn()
	fld.IsNumber = true
	fld.Name = "NullFloat"
	fld.DBType = "float"
	fld.AllowNull = true
	return fld
}

func fieldNullVarchar() *schema.Column {
	fld := schema.DefaultColumn()
	fld.Name = "NullVarchar"
	fld.DBType = "varchar"
	fld.AllowNull = true
	fld.Length = 30
	return fld
}

func fieldNullBlob() *schema.Column {
	fld := schema.DefaultColumn()
	fld.Name = "NullBlob"
	fld.DBType = "blob"
	fld.AllowNull = true
	return fld
}

func fieldNullTimestamp() *schema.Column {
	fld := schema.DefaultColumn()
	fld.Name = "NullTimestamp"
	fld.DBType = "timestamp"
	fld.AllowNull = true
	return fld
}

func fieldAddress(n string) *schema.Column {
	fld := schema.DefaultColumn()

	// Note: Name could be 'name', 'email', etc. or 'basic.name', 'basic.email',
	// but it could not be, say, people.basic.name where people is the root namespace
	// and there is also no way to specify the root namespace
	fld.Name = n
	fld.DBType = "text"
	fld.AllowNull = true // TODO: Should this be used to accept empty keys?
	return fld
}

func primaryColumn(name string) *schema.Column {
	fld := schema.DefaultColumn()
	fld.Name = name
	fld.DBType = "integer"
	fld.IsIdentity = true
	fld.IsNumber = true
	return fld
}

// Very simple person table initially
func peopleTable() *schema.Table {
	tbl := schema.DefaultTable()
	tbl.Name = "people"
	tbl.MultiKey = false
	tbl.Primary = "PersonID"
	tbl.Columns["PersonID"] = primaryColumn("PersonID")
	tbl.Columns["Name"] = fieldName()
	tbl.Columns["NullFloat"] = fieldNullFloat()
	tbl.Columns["NullText"] = fieldNullText()
	tbl.Columns["NullInt"] = fieldNullInt()
	tbl.Columns["NullVarchar"] = fieldNullVarchar()
	tbl.Columns["NullBlob"] = fieldNullBlob()
	tbl.Columns["NullTimestamp"] = fieldNullTimestamp()

	// TODO: Why was NullText getting retrieved as NULL when we didn't
	// have it in the EssentialColumns list?
	tbl.EssentialColumns = []string{"PersonID", "Name", "NullFloat", "NullText", "NullInt", "NullVarchar", "NullBlob", "NullTimestamp"}

	return tbl
}

func fkColumn(name string) *schema.Column {
	fld := schema.DefaultColumn()
	fld.Name = name
	fld.DBType = "integer"
	fld.IsIdentity = false
	fld.IsNumber = true
	return fld
}

// Very simple address table initially
func addressTable() *schema.Table {
	tbl := schema.DefaultTable()
	tbl.Name = "addresses"
	tbl.Primary = "AddressID"
	tbl.MultiKey = true
	tbl.ForeignKeys = []string{"PersonID"}

	tbl.Columns["PersonID"] = fkColumn("PersonID")

	tbl.Columns["AddressID"] = primaryColumn("AddressID")
	tbl.Columns["Address1"] = fieldAddress("Address1")
	tbl.Columns["Address2"] = fieldAddress("Address2")
	tbl.Columns["City"] = fieldAddress("City")
	tbl.Columns["State"] = fieldAddress("State")
	tbl.Columns["Zip"] = fieldAddress("Zip")

	tbl.EssentialColumns = []string{"AddressID", "PersonID", "Address1", "Address2", "City", "State", "Zip"}

	tbl.ParentTables = []string{"people"}
	return tbl
}

// BasicSchema is the basic mock for one table
func BasicSchema() *schema.Schema {
	sch := schema.DefaultSchema()
	personTable := peopleTable()

	sch.Tables["people"] = personTable
	return sch
}

// NestedSchema is the basic mock for two tables (one table that references a foreign table)
func NestedSchema() *schema.Schema {
	sch := BasicSchema()
	personTable := sch.Tables["people"]

	childTable := schema.DefaultChildTable()
	personTable.Children["addresses"] = childTable

	addrTable := addressTable()
	sch.Tables["addresses"] = addrTable
	return sch
}

func SampleAddressObject() *object.Object {
	addr := object.New("addresses")
	addr.Set("Address1", "Test")
	addr.Set("Address2", "Test2")
	addr.Set("City", "Nowhere")
	addr.Set("State", "AZ")
	addr.Set("Zip", "02865")
	return addr
}

func setPersonNulls(obj *object.Object) {
	obj.Set("NullText", object.NewNULLValue())
	obj.Set("NullInt", object.NewNULLValue())
	obj.Set("NullFloat", object.NewNULLValue())
	obj.Set("NullVarchar", object.NewNULLValue())
	obj.Set("NullBlob", object.NewNULLValue())
	obj.Set("NullTimestamp", object.NewNULLValue())
}

func DefaultPersonWithAddress() *object.Object {
	obj := object.New(PeopleObjectType)
	obj.Set("Name", "Ryan")
	setPersonNulls(obj)

	addrObj := SampleAddressObject()
	obj.Children["addresses"] = object.NewArray(addrObj)
	return obj
}

func RandomPerson() *object.Object {
	obj := object.New(PeopleObjectType)
	obj.Set("Name", fake.FirstName())
	setPersonNulls(obj)

	return obj
}

// Lookup-related table
func lookupTable() *schema.Table {
	tbl := schema.DefaultTable()
	tbl.Name = "LOOKUP_TABLE"
	tbl.MultiKey = false
	tbl.Primary = "LOOKUP_ID"
	tbl.Columns["LOOKUP_ID"] = primaryColumn("LOOKUP_ID")
	tbl.Columns["LOOKUP_TEXT"] = fieldLookupText()
	tbl.EssentialColumns = []string{"LOOKUP_ID", "LOOKUP_TEXT"}
	return tbl
}

func LookupSchema() *schema.Schema {
	sch := schema.DefaultSchema()
	lookupTable := lookupTable()
	sch.Tables["LOOKUP_TABLE"] = lookupTable
	return sch
}

func LookupValue() *object.Object {
	obj := object.New("LOOKUP_TABLE")
	obj.Set("LOOKUP_ID", int64(42))
	obj.Set("LOOKUP_TEXT", "SomeValue")
	return obj
}
<commit_msg>lookup mock schema cleanups<commit_after>package mock

import (
	"github.com/icrowley/fake"
	"github.com/rbastic/dyndao/object"
	"github.com/rbastic/dyndao/schema"
)

const PeopleObjectType string = "people"
const AddressesObjectType string = "addresses"

// Basic test mock
func fieldName() *schema.Column {
	fld := schema.DefaultColumn()
	fld.Name = "Name"
	fld.DBType = "text"
	return fld
}

func fieldLookupText() *schema.Column {
	fld := schema.DefaultColumn()
	fld.Name = "LOOKUP_TEXT"
	fld.DBType = "TEXT"
	fld.AllowNull = false
	return fld
}

func fieldNullText() *schema.Column {
	fld := schema.DefaultColumn()
	fld.Name = "NullText"
	fld.DBType = "text"
	fld.AllowNull = true
	return fld
}

func fieldNullInt() *schema.Column {
	fld := schema.DefaultColumn()
	fld.IsNumber = true
	fld.Name = "NullInt"
	fld.DBType = "integer"
	fld.AllowNull = true
	return fld
}

func fieldNullFloat() *schema.Column {
	fld := schema.DefaultColumn()
	fld.IsNumber = true
	fld.Name = "NullFloat"
	fld.DBType = "float"
	fld.AllowNull = true
	return fld
}

func fieldNullVarchar() *schema.Column {
	fld := schema.DefaultColumn()
	fld.Name = "NullVarchar"
	fld.DBType = "varchar"
	fld.AllowNull = true
	fld.Length = 30
	return fld
}

func fieldNullBlob() *schema.Column {
	fld := schema.DefaultColumn()
	fld.Name = "NullBlob"
	fld.DBType = "blob"
	fld.AllowNull = true
	return fld
}

func fieldNullTimestamp() *schema.Column {
	fld := schema.DefaultColumn()
	fld.Name = "NullTimestamp"
	fld.DBType = "timestamp"
	fld.AllowNull = true
	return fld
}

func fieldAddress(n string) *schema.Column {
	fld := schema.DefaultColumn()

	// Note: Name could be 'name', 'email', etc. or 'basic.name', 'basic.email',
	// but it could not be, say, people.basic.name where people is the root namespace
	// and there is also no way to specify the root namespace
	fld.Name = n
	fld.DBType = "text"
	fld.AllowNull = true // TODO: Should this be used to accept empty keys?
	return fld
}

func primaryColumn(name string) *schema.Column {
	fld := schema.DefaultColumn()
	fld.Name = name
	fld.DBType = "integer"
	fld.IsIdentity = true
	fld.IsNumber = true
	return fld
}

// Very simple person table initially
func peopleTable() *schema.Table {
	tbl := schema.DefaultTable()
	tbl.Name = "people"
	tbl.MultiKey = false
	tbl.Primary = "PersonID"
	tbl.Columns["PersonID"] = primaryColumn("PersonID")
	tbl.Columns["Name"] = fieldName()
	tbl.Columns["NullFloat"] = fieldNullFloat()
	tbl.Columns["NullText"] = fieldNullText()
	tbl.Columns["NullInt"] = fieldNullInt()
	tbl.Columns["NullVarchar"] = fieldNullVarchar()
	tbl.Columns["NullBlob"] = fieldNullBlob()
	tbl.Columns["NullTimestamp"] = fieldNullTimestamp()

	// TODO: Why was NullText getting retrieved as NULL when we didn't
	// have it in the EssentialColumns list?
	tbl.EssentialColumns = []string{"PersonID", "Name", "NullFloat", "NullText", "NullInt", "NullVarchar", "NullBlob", "NullTimestamp"}

	return tbl
}

func fkColumn(name string) *schema.Column {
	fld := schema.DefaultColumn()
	fld.Name = name
	fld.DBType = "integer"
	fld.IsIdentity = false
	fld.IsNumber = true
	return fld
}

// Very simple address table initially
func addressTable() *schema.Table {
	tbl := schema.DefaultTable()
	tbl.Name = "addresses"
	tbl.Primary = "AddressID"
	tbl.MultiKey = true
	tbl.ForeignKeys = []string{"PersonID"}

	tbl.Columns["PersonID"] = fkColumn("PersonID")

	tbl.Columns["AddressID"] = primaryColumn("AddressID")
	tbl.Columns["Address1"] = fieldAddress("Address1")
	tbl.Columns["Address2"] = fieldAddress("Address2")
	tbl.Columns["City"] = fieldAddress("City")
	tbl.Columns["State"] = fieldAddress("State")
	tbl.Columns["Zip"] = fieldAddress("Zip")

	tbl.EssentialColumns = []string{"AddressID", "PersonID", "Address1", "Address2", "City", "State", "Zip"}

	tbl.ParentTables = []string{"people"}
	return tbl
}

// BasicSchema is the basic mock for one table
func BasicSchema() *schema.Schema {
	sch := schema.DefaultSchema()
	personTable := peopleTable()

	sch.Tables["people"] = personTable
	return sch
}

// NestedSchema is the basic mock for two tables (one table that references a foreign table)
func NestedSchema() *schema.Schema {
	sch := BasicSchema()
	personTable := sch.Tables["people"]

	childTable := schema.DefaultChildTable()
	personTable.Children["addresses"] = childTable

	addrTable := addressTable()
	sch.Tables["addresses"] = addrTable
	return sch
}

func SampleAddressObject() *object.Object {
	addr := object.New("addresses")
	addr.Set("Address1", "Test")
	addr.Set("Address2", "Test2")
	addr.Set("City", "Nowhere")
	addr.Set("State", "AZ")
	addr.Set("Zip", "02865")
	return addr
}

func setPersonNulls(obj *object.Object) {
	obj.Set("NullText", object.NewNULLValue())
	obj.Set("NullInt", object.NewNULLValue())
	obj.Set("NullFloat", object.NewNULLValue())
	obj.Set("NullVarchar", object.NewNULLValue())
	obj.Set("NullBlob", object.NewNULLValue())
	obj.Set("NullTimestamp", object.NewNULLValue())
}

func DefaultPersonWithAddress() *object.Object {
	obj := object.New(PeopleObjectType)
	obj.Set("Name", "Ryan")
	setPersonNulls(obj)

	addrObj := SampleAddressObject()
	obj.Children["addresses"] = object.NewArray(addrObj)
	return obj
}

func RandomPerson() *object.Object {
	obj := object.New(PeopleObjectType)
	obj.Set("Name", fake.FirstName())
	setPersonNulls(obj)

	return obj
}

const (
	MockLookupTable = "LOOKUP_TABLE"
	MockLookupID    = "LOOKUP_ID"
	MockLookupText  = "LOOKUP_TEXT"
)

// Lookup-related table
func lookupTable() *schema.Table {
	tbl := schema.DefaultTable()
	tbl.Name = MockLookupTable
	tbl.MultiKey = false
	tbl.Primary = MockLookupID
	tbl.Columns[MockLookupID] = primaryColumn(MockLookupID)
	tbl.Columns[MockLookupText] = fieldLookupText()
	tbl.EssentialColumns = []string{MockLookupID, MockLookupText}
	return tbl
}

func LookupSchema() *schema.Schema {
	sch := schema.DefaultSchema()
	lookupTable := lookupTable()
	sch.Tables[MockLookupTable] = lookupTable
	return sch
}

func LookupValue() *object.Object {
	obj := object.New(MockLookupTable)
	obj.Set(MockLookupID, int64(42))
	obj.Set(MockLookupText, "SomeValue")
	return obj
}
<|endoftext|>


<commit_before>// Copyright 2015 Google Inc. All Rights Reserved.
//
// 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.

package bigquery

import (
	"errors"
	"fmt"
	"net/http"
	"time"

	"golang.org/x/net/context"
	bq "google.golang.org/api/bigquery/v2"
)

// service provides an internal abstraction to isolate the generated
// BigQuery API; most of this package uses this interface instead.
// The single implementation, *bigqueryService, contains all the knowledge
// of the generated BigQuery API.
type service interface {
	insertJob(ctx context.Context, job *bq.Job, projectId string) (*Job, error)
	getJobType(ctx context.Context, projectId, jobID string) (jobType, error)
	jobStatus(ctx context.Context, projectId, jobID string) (*JobStatus, error)
	listTables(ctx context.Context, projectID, datasetID, pageToken string) ([]*Table, string, error)

	// readQuery reads data resulting from a query job. If the job is not
	// yet complete, an errIncompleteJob is returned. readQuery may be
	// called repeatedly to wait for results indefinitely.
	readQuery(ctx context.Context, conf *readQueryConf, pageToken string) (*readDataResult, error)

	readTabledata(ctx context.Context, conf *readTableConf, pageToken string) (*readDataResult, error)
}

type bigqueryService struct {
	s *bq.Service
}

func newBigqueryService(client *http.Client) (*bigqueryService, error) {
	s, err := bq.New(client)
	if err != nil {
		return nil, fmt.Errorf("constructing bigquery client: %v", err)
	}

	return &bigqueryService{s: s}, nil
}

// getPages calls the supplied getPage function repeatedly until there are no pages left to get.
// token is the token of the initial page to start from.  Use an empty string to start from the beginning.
func getPages(token string, getPage func(token string) (nextToken string, err error)) error {
	for {
		var err error
		token, err = getPage(token)
		if err != nil {
			return err
		}
		if token == "" {
			return nil
		}
	}
}

func (s *bigqueryService) insertJob(ctx context.Context, job *bq.Job, projectID string) (*Job, error) {
	// TODO(mcgreevy): use ctx
	res, err := s.s.Jobs.Insert(projectID, job).Do()
	if err != nil {
		return nil, err
	}
	return &Job{service: s, projectID: projectID, jobID: res.JobReference.JobId}, nil
}

type pagingConf struct {
	recordsPerRequest    int64
	setRecordsPerRequest bool

	startIndex uint64
}

type readTableConf struct {
	projectID, datasetID, tableID string
	paging                        pagingConf
}

type readDataResult struct {
	pageToken string
	rows      [][]Value
	totalRows uint64
	schema    Schema
}

type readQueryConf struct {
	projectID, jobID string
	paging           pagingConf
}

func (s *bigqueryService) readTabledata(ctx context.Context, conf *readTableConf, pageToken string) (*readDataResult, error) {
	req := s.s.Tabledata.List(conf.projectID, conf.datasetID, conf.tableID).
		PageToken(pageToken).
		StartIndex(conf.paging.startIndex)

	if conf.paging.setRecordsPerRequest {
		req = req.MaxResults(conf.paging.recordsPerRequest)
	}

	res, err := req.Do()
	if err != nil {
		return nil, err
	}

	result := &readDataResult{
		pageToken: res.PageToken,
		rows:      convertRows(res.Rows),
		totalRows: uint64(res.TotalRows),
	}
	return result, nil
}

var errIncompleteJob = errors.New("internal error: query results not available because job is not complete")

// getQueryResultsTimeout controls the maximum duration of a request to the
// BigQuery GetQueryResults endpoint.  Setting a long timeout here does not
// cause increased overall latency, as results are returned as soon as they are
// available.
const getQueryResultsTimeout = time.Minute

func (s *bigqueryService) readQuery(ctx context.Context, conf *readQueryConf, pageToken string) (*readDataResult, error) {
	req := s.s.Jobs.GetQueryResults(conf.projectID, conf.jobID).
		PageToken(pageToken).
		StartIndex(conf.paging.startIndex).
		TimeoutMs(getQueryResultsTimeout.Nanoseconds() / 1000)

	if conf.paging.setRecordsPerRequest {
		req = req.MaxResults(conf.paging.recordsPerRequest)
	}

	res, err := req.Do()
	if err != nil {
		return nil, err
	}

	if !res.JobComplete {
		return nil, errIncompleteJob
	}

	result := &readDataResult{
		pageToken: res.PageToken,
		rows:      convertRows(res.Rows),
		totalRows: res.TotalRows,
		schema:    convertTableSchema(res.Schema),
	}
	return result, nil
}

func convertRows(rows []*bq.TableRow) [][]Value {
	convertRow := func(r *bq.TableRow) []Value {
		var values []Value
		for _, cell := range r.F {
			values = append(values, cell.V)
		}
		return values
	}

	var rs [][]Value
	for _, r := range rows {
		rs = append(rs, convertRow(r))
	}
	return rs
}

type jobType int

const (
	copyJobType jobType = iota
	extractJobType
	loadJobType
	queryJobType
)

func (s *bigqueryService) getJobType(ctx context.Context, projectID, jobID string) (jobType, error) {
	// TODO(mcgreevy): use ctx
	res, err := s.s.Jobs.Get(projectID, jobID).
		Fields("configuration").
		Do()

	if err != nil {
		return 0, err
	}

	switch {
	case res.Configuration.Copy != nil:
		return copyJobType, nil
	case res.Configuration.Extract != nil:
		return extractJobType, nil
	case res.Configuration.Load != nil:
		return loadJobType, nil
	case res.Configuration.Query != nil:
		return queryJobType, nil
	default:
		return 0, errors.New("unknown job type")
	}
}

func (s *bigqueryService) jobStatus(ctx context.Context, projectID, jobID string) (*JobStatus, error) {
	// TODO(mcgreevy): use ctx
	res, err := s.s.Jobs.Get(projectID, jobID).
		Fields("status"). // Only fetch what we need.
		Do()
	if err != nil {
		return nil, err
	}
	return jobStatusFromProto(res.Status)
}

var stateMap = map[string]State{"PENDING": Pending, "RUNNING": Running, "DONE": Done}

func jobStatusFromProto(status *bq.JobStatus) (*JobStatus, error) {
	state, ok := stateMap[status.State]
	if !ok {
		return nil, fmt.Errorf("unexpected job state: %v", status.State)
	}

	newStatus := &JobStatus{
		State: state,
		err:   nil,
	}
	if err := errorFromErrorProto(status.ErrorResult); state == Done && err != nil {
		newStatus.err = err
	}

	for _, ep := range status.Errors {
		newStatus.Errors = append(newStatus.Errors, errorFromErrorProto(ep))
	}
	return newStatus, nil
}

// listTables returns a subset of tables that belong to a dataset, and a token for fetching the next subset.
func (s *bigqueryService) listTables(ctx context.Context, projectID, datasetID, pageToken string) ([]*Table, string, error) {
	// TODO(mcgreevy): use ctx
	var tables []*Table
	res, err := s.s.Tables.List(projectID, datasetID).
		PageToken(pageToken).
		Do()
	if err != nil {
		return nil, "", err
	}
	for _, t := range res.Tables {
		tables = append(tables, convertTable(t))
	}
	return tables, res.NextPageToken, nil
}

func convertTable(t *bq.TableListTables) *Table {
	return &Table{
		ProjectID: t.TableReference.ProjectId,
		DatasetID: t.TableReference.DatasetId,
		TableID:   t.TableReference.TableId,
	}
}
<commit_msg>cloud/bigquery: Only set one of startIndex or pageToken for reads.<commit_after>// Copyright 2015 Google Inc. All Rights Reserved.
//
// 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.

package bigquery

import (
	"errors"
	"fmt"
	"net/http"
	"time"

	"golang.org/x/net/context"
	bq "google.golang.org/api/bigquery/v2"
)

// service provides an internal abstraction to isolate the generated
// BigQuery API; most of this package uses this interface instead.
// The single implementation, *bigqueryService, contains all the knowledge
// of the generated BigQuery API.
type service interface {
	insertJob(ctx context.Context, job *bq.Job, projectId string) (*Job, error)
	getJobType(ctx context.Context, projectId, jobID string) (jobType, error)
	jobStatus(ctx context.Context, projectId, jobID string) (*JobStatus, error)
	listTables(ctx context.Context, projectID, datasetID, pageToken string) ([]*Table, string, error)

	// readQuery reads data resulting from a query job. If the job is not
	// yet complete, an errIncompleteJob is returned. readQuery may be
	// called repeatedly to wait for results indefinitely.
	readQuery(ctx context.Context, conf *readQueryConf, pageToken string) (*readDataResult, error)

	readTabledata(ctx context.Context, conf *readTableConf, pageToken string) (*readDataResult, error)
}

type bigqueryService struct {
	s *bq.Service
}

func newBigqueryService(client *http.Client) (*bigqueryService, error) {
	s, err := bq.New(client)
	if err != nil {
		return nil, fmt.Errorf("constructing bigquery client: %v", err)
	}

	return &bigqueryService{s: s}, nil
}

// getPages calls the supplied getPage function repeatedly until there are no pages left to get.
// token is the token of the initial page to start from.  Use an empty string to start from the beginning.
func getPages(token string, getPage func(token string) (nextToken string, err error)) error {
	for {
		var err error
		token, err = getPage(token)
		if err != nil {
			return err
		}
		if token == "" {
			return nil
		}
	}
}

func (s *bigqueryService) insertJob(ctx context.Context, job *bq.Job, projectID string) (*Job, error) {
	// TODO(mcgreevy): use ctx
	res, err := s.s.Jobs.Insert(projectID, job).Do()
	if err != nil {
		return nil, err
	}
	return &Job{service: s, projectID: projectID, jobID: res.JobReference.JobId}, nil
}

type pagingConf struct {
	recordsPerRequest    int64
	setRecordsPerRequest bool

	startIndex uint64
}

type readTableConf struct {
	projectID, datasetID, tableID string
	paging                        pagingConf
}

type readDataResult struct {
	pageToken string
	rows      [][]Value
	totalRows uint64
	schema    Schema
}

type readQueryConf struct {
	projectID, jobID string
	paging           pagingConf
}

func (s *bigqueryService) readTabledata(ctx context.Context, conf *readTableConf, pageToken string) (*readDataResult, error) {
	req := s.s.Tabledata.List(conf.projectID, conf.datasetID, conf.tableID)

	if pageToken != "" {
		req.PageToken(pageToken)
	} else {
		req.StartIndex(conf.paging.startIndex)
	}

	if conf.paging.setRecordsPerRequest {
		req.MaxResults(conf.paging.recordsPerRequest)
	}

	res, err := req.Do()
	if err != nil {
		return nil, err
	}

	result := &readDataResult{
		pageToken: res.PageToken,
		rows:      convertRows(res.Rows),
		totalRows: uint64(res.TotalRows),
	}
	return result, nil
}

var errIncompleteJob = errors.New("internal error: query results not available because job is not complete")

// getQueryResultsTimeout controls the maximum duration of a request to the
// BigQuery GetQueryResults endpoint.  Setting a long timeout here does not
// cause increased overall latency, as results are returned as soon as they are
// available.
const getQueryResultsTimeout = time.Minute

func (s *bigqueryService) readQuery(ctx context.Context, conf *readQueryConf, pageToken string) (*readDataResult, error) {
	req := s.s.Jobs.GetQueryResults(conf.projectID, conf.jobID).
		TimeoutMs(getQueryResultsTimeout.Nanoseconds() / 1e6)

	if pageToken != "" {
		req.PageToken(pageToken)
	} else {
		req.StartIndex(conf.paging.startIndex)
	}

	if conf.paging.setRecordsPerRequest {
		req.MaxResults(conf.paging.recordsPerRequest)
	}

	res, err := req.Do()
	if err != nil {
		return nil, err
	}

	if !res.JobComplete {
		return nil, errIncompleteJob
	}

	result := &readDataResult{
		pageToken: res.PageToken,
		rows:      convertRows(res.Rows),
		totalRows: res.TotalRows,
		schema:    convertTableSchema(res.Schema),
	}
	return result, nil
}

func convertRows(rows []*bq.TableRow) [][]Value {
	convertRow := func(r *bq.TableRow) []Value {
		var values []Value
		for _, cell := range r.F {
			values = append(values, cell.V)
		}
		return values
	}

	var rs [][]Value
	for _, r := range rows {
		rs = append(rs, convertRow(r))
	}
	return rs
}

type jobType int

const (
	copyJobType jobType = iota
	extractJobType
	loadJobType
	queryJobType
)

func (s *bigqueryService) getJobType(ctx context.Context, projectID, jobID string) (jobType, error) {
	// TODO(mcgreevy): use ctx
	res, err := s.s.Jobs.Get(projectID, jobID).
		Fields("configuration").
		Do()

	if err != nil {
		return 0, err
	}

	switch {
	case res.Configuration.Copy != nil:
		return copyJobType, nil
	case res.Configuration.Extract != nil:
		return extractJobType, nil
	case res.Configuration.Load != nil:
		return loadJobType, nil
	case res.Configuration.Query != nil:
		return queryJobType, nil
	default:
		return 0, errors.New("unknown job type")
	}
}

func (s *bigqueryService) jobStatus(ctx context.Context, projectID, jobID string) (*JobStatus, error) {
	// TODO(mcgreevy): use ctx
	res, err := s.s.Jobs.Get(projectID, jobID).
		Fields("status"). // Only fetch what we need.
		Do()
	if err != nil {
		return nil, err
	}
	return jobStatusFromProto(res.Status)
}

var stateMap = map[string]State{"PENDING": Pending, "RUNNING": Running, "DONE": Done}

func jobStatusFromProto(status *bq.JobStatus) (*JobStatus, error) {
	state, ok := stateMap[status.State]
	if !ok {
		return nil, fmt.Errorf("unexpected job state: %v", status.State)
	}

	newStatus := &JobStatus{
		State: state,
		err:   nil,
	}
	if err := errorFromErrorProto(status.ErrorResult); state == Done && err != nil {
		newStatus.err = err
	}

	for _, ep := range status.Errors {
		newStatus.Errors = append(newStatus.Errors, errorFromErrorProto(ep))
	}
	return newStatus, nil
}

// listTables returns a subset of tables that belong to a dataset, and a token for fetching the next subset.
func (s *bigqueryService) listTables(ctx context.Context, projectID, datasetID, pageToken string) ([]*Table, string, error) {
	// TODO(mcgreevy): use ctx
	var tables []*Table
	res, err := s.s.Tables.List(projectID, datasetID).
		PageToken(pageToken).
		Do()
	if err != nil {
		return nil, "", err
	}
	for _, t := range res.Tables {
		tables = append(tables, convertTable(t))
	}
	return tables, res.NextPageToken, nil
}

func convertTable(t *bq.TableListTables) *Table {
	return &Table{
		ProjectID: t.TableReference.ProjectId,
		DatasetID: t.TableReference.DatasetId,
		TableID:   t.TableReference.TableId,
	}
}
<|endoftext|>


<commit_before>////////////////////////////////////////////////////////////////////////////
// Porgram: EasyGen
// Purpose: Universal code/text generator that is easy to use
// Authors: Tong Sun (c) 2015, All rights reserved
////////////////////////////////////////////////////////////////////////////

////////////////////////////////////////////////////////////////////////////
// Program start

package main

import (
	"flag"
	"fmt"
	//ht "html/template"
	"io/ioutil"
	"os"
	"strings"
	tt "text/template"

	"gopkg.in/yaml.v2"
)

////////////////////////////////////////////////////////////////////////////
// Constant and data type/structure definitions

const progname = "EasyGen" // os.Args[0]

type Options struct {
	Html bool
}

////////////////////////////////////////////////////////////////////////////
// Global variables definitions

var opts Options

////////////////////////////////////////////////////////////////////////////
// Commandline definitions

func init() {
	flag.BoolVar(&opts.Html, "html", false, "Use html template instead of text")
}

func usage() {
	fmt.Fprintf(os.Stderr, "\nUsage:\n %s [flags] TemplateFileName\n\nFlags:\n\n",
		progname)
	flag.PrintDefaults()
	fmt.Fprintf(os.Stderr, "\nTemplateFileName: The name for the template and yaml file\n\tOnly the name part, without extension. Can include the path as well.\n")
	os.Exit(0)
}

////////////////////////////////////////////////////////////////////////////
// Main

func main() {
	flag.Usage = usage
	flag.Parse()

	// One mandatory non-flag arguments
	if len(flag.Args()) < 1 {
		usage()
	}
	fileName := flag.Args()[0]

	source, err := ioutil.ReadFile(fileName + ".yaml")
	checkError(err)

	m := make(map[interface{}]interface{})

	err = yaml.Unmarshal(source, &m)
	checkError(err)

	// var t interface{}
	// if opts.Html {
	// 	t := ht.New("H")
	// } else {
	// 	t := tt.New("T")
	// }

	t := tt.New("T")
	// add functions
	t = t.Funcs(tt.FuncMap{"join": Join})
	//fmt.Fprintf(os.Stderr, "] Ok\n")
	t, err = t.Parse("{{range .Colors}}{{.}}, {{end}}.")

	t, err = tt.ParseFiles(fileName + ".tmpl")
	checkError(err)

	err = t.Execute(os.Stdout, m)
	checkError(err)

}

////////////////////////////////////////////////////////////////////////////
// Function definitions

func checkError(err error) {
	if err != nil {
		fmt.Printf("[%s] Fatal error - %v", progname, err.Error())
		os.Exit(1)
	}
}

//==========================================================================
// template custom functions

var funcs = tt.FuncMap{"minus1": func(n int) int { return n - 1 }}

func Join(s []string) string { return strings.Join(s, ", ") }
<commit_msg>- [*] make it works for both text and html templates<commit_after>////////////////////////////////////////////////////////////////////////////
// Porgram: EasyGen
// Purpose: Universal code/text generator that is easy to use
// Authors: Tong Sun (c) 2015, All rights reserved
////////////////////////////////////////////////////////////////////////////

////////////////////////////////////////////////////////////////////////////
// Program start

package main

import (
	"flag"
	"fmt"
	ht "html/template"
	"io"
	"io/ioutil"
	"os"
	"strings"
	tt "text/template"

	"gopkg.in/yaml.v2"
)

////////////////////////////////////////////////////////////////////////////
// Constant and data type/structure definitions

const progname = "EasyGen" // os.Args[0]

// Commandline Options
type Options struct {
	HTML bool
}

// common type for a *(text|html).Template value
type template interface {
	Execute(wr io.Writer, data interface{}) error
	ExecuteTemplate(wr io.Writer, name string, data interface{}) error
	Name() string
}

////////////////////////////////////////////////////////////////////////////
// Global variables definitions

var opts Options

////////////////////////////////////////////////////////////////////////////
// Commandline definitions

func init() {
	flag.BoolVar(&opts.HTML, "html", false, "Use html template instead of text")
}

func usage() {
	fmt.Fprintf(os.Stderr, "\nUsage:\n %s [flags] TemplateFileName\n\nFlags:\n\n",
		progname)
	flag.PrintDefaults()
	fmt.Fprintf(os.Stderr, "\nTemplateFileName: The name for the template and yaml file\n\tOnly the name part, without extension. Can include the path as well.\n")
	os.Exit(0)
}

////////////////////////////////////////////////////////////////////////////
// Main

func main() {
	flag.Usage = usage
	flag.Parse()

	// One mandatory non-flag arguments
	if len(flag.Args()) < 1 {
		usage()
	}
	fileName := flag.Args()[0]

	source, err := ioutil.ReadFile(fileName + ".yaml")
	checkError(err)

	m := make(map[interface{}]interface{})

	err = yaml.Unmarshal(source, &m)
	checkError(err)

	t, err := parseFiles(opts.HTML, fileName+".tmpl")
	checkError(err)

	err = t.Execute(os.Stdout, m)
	checkError(err)

}

////////////////////////////////////////////////////////////////////////////
// Function definitions

// Exit if error occurs
func checkError(err error) {
	if err != nil {
		fmt.Printf("[%s] Fatal error - %v", progname, err.Error())
		os.Exit(1)
	}
}

// parseFiles, intialization. By Matt Harden @gmail.com
func parseFiles(HTML bool, filenames ...string) (template, error) {
	if HTML {
		t, err := ht.ParseFiles(filenames...)
		return t, err
	} else {
		t, err := tt.ParseFiles(filenames...)
		return t, err
	}
}

//==========================================================================
// template custom functions

var funcs = tt.FuncMap{"minus1": func(n int) int { return n - 1 }}

// Custom template function to join string array
func join(s []string) string { return strings.Join(s, ", ") }
<|endoftext|>


<commit_before>// UVa 254 - Towers of Hanoi

package main

import (
	"fmt"
	"math/big"
	"os"
)

var (
	pegs [3]int
	bin  string
)

func solve(n, from, to int) {
	if n < 0 {
		return
	}
	tmp := 3 - from - to
	if len(bin) >= n+1 && bin[len(bin)-n-1] == '1' {
		pegs[from] -= n + 1
		pegs[tmp] += n
		pegs[to]++
		solve(n-1, tmp, to)
	} else {
		solve(n-1, from, tmp)
	}
}

func main() {
	in, _ := os.Open("254.in")
	defer in.Close()
	out, _ := os.Create("254.out")
	defer out.Close()

	var n int
	var m big.Int
	var s string
	for {
		if fmt.Fscanf(in, "%d %s", &n, &s); n == 0 && s == "0" {
			break
		}
		m.SetString(s, 10)
		bin = fmt.Sprintf("%b", &m)
		pegs = [3]int{n, 0, 0}
		if n%2 == 0 {
			solve(n-1, 0, 2)
		} else {
			solve(n-1, 0, 1)
		}
		for i, peg := range pegs {
			if i > 0 {
				fmt.Fprint(out, " ")
			}
			fmt.Fprint(out, peg)
		}
		fmt.Fprintln(out)
	}
}
<commit_msg>code clean up<commit_after>// UVa 254 - Towers of Hanoi

package main

import (
	"fmt"
	"math/big"
	"os"
)

var (
	pegs [3]int
	bin  string
)

func hanoi(n, from, to int) {
	if n < 0 {
		return
	}
	tmp := 3 - from - to
	if len(bin) >= n+1 && bin[len(bin)-n-1] == '1' {
		pegs[from] -= n + 1
		pegs[tmp] += n
		pegs[to]++
		hanoi(n-1, tmp, to)
	} else {
		hanoi(n-1, from, tmp)
	}
}

func main() {
	in, _ := os.Open("254.in")
	defer in.Close()
	out, _ := os.Create("254.out")
	defer out.Close()

	var n int
	var m big.Int
	var s string
	for {
		if fmt.Fscanf(in, "%d %s", &n, &s); n == 0 && s == "0" {
			break
		}
		m.SetString(s, 10)
		bin = fmt.Sprintf("%b", &m)
		pegs = [3]int{n, 0, 0}
		if n%2 == 0 {
			hanoi(n-1, 0, 2)
		} else {
			hanoi(n-1, 0, 1)
		}
		for i, peg := range pegs {
			if i > 0 {
				fmt.Fprint(out, " ")
			}
			fmt.Fprint(out, peg)
		}
		fmt.Fprintln(out)
	}
}
<|endoftext|>


<commit_before>package gateways

import (
	"fmt"
	"io"
	"net/http"

	"github.com/hexablock/blox/block"
)

func (server *HTTPServer) handleFS(w http.ResponseWriter, r *http.Request, resourceID string) {
	var err error

	switch r.Method {
	case http.MethodGet:
		err = server.handlerFSGet(w, resourceID)

	case http.MethodPost:
		err = server.handlerFSPost(w, r, resourceID)

	default:
		w.WriteHeader(405)
		return
	}

	if err == nil {
		return
	}

	var (
		headers map[string]string
		code    int
		data    interface{}
	)

	// if strings.Contains(err.Error(), "host not in set") {
	// 	code, headers, data, err = buildRedirect(meta.PeerSet, r.URL.RequestURI())
	// }

	writeJSONResponse(w, code, headers, data, err)
}

func (server *HTTPServer) handlerFSGet(w http.ResponseWriter, resourceID string) error {
	fs := server.fids.FileSystem()

	fh, err := fs.Open(resourceID)
	if err != nil {
		var headers map[string]string
		if err == block.ErrBlockNotFound {
			writeJSONResponse(w, 404, headers, nil, err)
		} else {
			writeJSONResponse(w, 400, headers, nil, err)
		}
		return err
	}

	w.Header().Set("Content-Length", fmt.Sprintf("%d", fh.Size()))
	w.Header().Set(headerBlockSize, fmt.Sprintf("%d", fh.BlockSize()))

	_, err = io.Copy(w, fh)
	fh.Close()

	return err
}

func (server *HTTPServer) handlerFSPost(w http.ResponseWriter, r *http.Request, resourceID string) error {

	headers := map[string]string{}
	fs := server.fids.FileSystem()

	fh, err := fs.Create(resourceID)
	if err != nil {
		writeJSONResponse(w, 400, headers, nil, err)
		return err
	}

	_, err = io.Copy(fh, r.Body)
	defer r.Body.Close()

	if err != nil {
		writeJSONResponse(w, 400, headers, nil, err)
		return err
	}

	// Final response after upload completes assuming there are no errors
	code := 201

	if err = fh.Close(); err == block.ErrBlockExists {
		err = nil
		// The server has fulfilled a request for the resource, and the response is a
		// representation of the result of one or more instance-manipulations applied
		// to the current instance
		code = 226
	}

	data := fh.Sys()
	rt := fh.Runtime()
	headers[headerRuntime] = fmt.Sprintf("%v", rt)
	headers[headerBlockSize] = fmt.Sprintf("%d", fh.BlockSize())
	writeJSONResponse(w, code, headers, data, err)

	return err
}
<commit_msg>Fix http fs endpoint<commit_after>package gateways

import (
	"fmt"
	"io"
	"net/http"

	"github.com/hexablock/blox/block"
)

func (server *HTTPServer) handleFS(w http.ResponseWriter, r *http.Request, resourceID string) {
	var err error

	switch r.Method {
	case http.MethodGet:
		err = server.handlerFSGet(w, resourceID)

	case http.MethodPost:
		server.handlerFSPost(w, r, resourceID)
		return

	default:
		w.WriteHeader(405)
		return
	}

	if err == nil {
		return
	}

	var (
		headers map[string]string
		code    int
		data    interface{}
	)

	// if strings.Contains(err.Error(), "host not in set") {
	// 	code, headers, data, err = buildRedirect(meta.PeerSet, r.URL.RequestURI())
	// }

	writeJSONResponse(w, code, headers, data, err)
}

func (server *HTTPServer) handlerFSGet(w http.ResponseWriter, resourceID string) error {
	fs := server.fids.FileSystem()

	fh, err := fs.Open(resourceID)
	if err != nil {
		var headers map[string]string
		if err == block.ErrBlockNotFound {
			writeJSONResponse(w, 404, headers, nil, err)
		} else {
			writeJSONResponse(w, 400, headers, nil, err)
		}
		return err
	}

	w.Header().Set("Content-Length", fmt.Sprintf("%d", fh.Size()))
	w.Header().Set(headerBlockSize, fmt.Sprintf("%d", fh.BlockSize()))

	_, err = io.Copy(w, fh)
	fh.Close()

	return err
}

func (server *HTTPServer) handlerFSPost(w http.ResponseWriter, r *http.Request, resourceID string) error {

	headers := map[string]string{}
	fs := server.fids.FileSystem()

	fh, err := fs.Create(resourceID)
	if err != nil {
		writeJSONResponse(w, 400, headers, nil, err)
		return err
	}

	_, err = io.Copy(fh, r.Body)
	defer r.Body.Close()

	if err != nil {
		writeJSONResponse(w, 400, headers, nil, err)
		return err
	}

	// Final response after upload completes assuming there are no errors
	code := 201

	if err = fh.Close(); err == block.ErrBlockExists {
		err = nil
		// The server has fulfilled a request for the resource, and the response is a
		// representation of the result of one or more instance-manipulations applied
		// to the current instance
		code = 226
	}

	data := fh.Sys()
	rt := fh.Runtime()
	headers[headerRuntime] = fmt.Sprintf("%v", rt)
	headers[headerBlockSize] = fmt.Sprintf("%d", fh.BlockSize())
	writeJSONResponse(w, code, headers, data, err)

	return err
}
<|endoftext|>


<commit_before>package cidrutil

import (
	"fmt"
	"net"
	"strings"

	"github.com/hashicorp/vault/helper/strutil"
)

// IPBelongsToCIDR checks if the given IP is encompassed by the given CIDR block
func IPBelongsToCIDR(ipAddr string, cidr string) (bool, error) {
	if ipAddr == "" {
		return false, fmt.Errorf("missing IP address")
	}

	ip := net.ParseIP(ipAddr)
	if ip == nil {
		return false, fmt.Errorf("invalid IP address")
	}

	_, ipnet, err := net.ParseCIDR(cidr)
	if err != nil {
		return false, err
	}

	if !ipnet.Contains(ip) {
		return false, nil
	}

	return true, nil
}

// IPBelongsToCIDRBlocksString checks if the given IP is encompassed by any of
// the given CIDR blocks, when the input is a string composed by joining all
// the CIDR blocks using a separator. The input is separated based on the given
// separator and the IP is checked to be belonged by any CIDR block.
func IPBelongsToCIDRBlocksString(ipAddr string, cidrList, separator string) (bool, error) {
	if ipAddr == "" {
		return false, fmt.Errorf("missing IP address")
	}

	if cidrList == "" {
		return false, fmt.Errorf("missing CIDR list")
	}

	if separator == "" {
		return false, fmt.Errorf("missing separator")
	}

	if ip := net.ParseIP(ipAddr); ip == nil {
		return false, fmt.Errorf("invalid IP address")
	}

	return IPBelongsToCIDRBlocksSlice(ipAddr, strutil.ParseDedupAndSortStrings(cidrList, separator))
}

// IPBelongsToCIDRBlocksSlice checks if the given IP is encompassed by any of the given
// CIDR blocks
func IPBelongsToCIDRBlocksSlice(ipAddr string, cidrs []string) (bool, error) {
	if ipAddr == "" {
		return false, fmt.Errorf("missing IP address")
	}

	if len(cidrs) == 0 {
		return false, fmt.Errorf("missing CIDR blocks to be checked against")
	}

	if ip := net.ParseIP(ipAddr); ip == nil {
		return false, fmt.Errorf("invalid IP address")
	}

	for _, cidr := range cidrs {
		belongs, err := IPBelongsToCIDR(ipAddr, cidr)
		if err != nil {
			return false, err
		}
		if belongs {
			return true, nil
		}
	}

	return false, nil
}

// ValidateCIDRListString checks if the list of CIDR blocks are valid, given
// that the input is a string composed by joining all the CIDR blocks using a
// separator. The input is separated based on the given separator and validity
// of each is checked.
func ValidateCIDRListString(cidrList string, separator string) (bool, error) {
	if cidrList == "" {
		return false, fmt.Errorf("missing CIDR list that needs validation")
	}
	if separator == "" {
		return false, fmt.Errorf("missing separator")
	}

	return ValidateCIDRListSlice(strutil.ParseDedupAndSortStrings(cidrList, separator))
}

// ValidateCIDRListSlice checks if the given list of CIDR blocks are valid
func ValidateCIDRListSlice(cidrBlocks []string) (bool, error) {
	if len(cidrBlocks) == 0 {
		return false, fmt.Errorf("missing CIDR blocks that needs validation")
	}

	for _, block := range cidrBlocks {
		if _, _, err := net.ParseCIDR(strings.TrimSpace(block)); err != nil {
			return false, err
		}
	}

	return true, nil
}

// Subset checks if the IPs belonging to a given CIDR block is a subset of IPs
// belonging to another CIDR block.
func Subset(cidr1, cidr2 string) (bool, error) {
	if cidr1 == "" {
		return false, fmt.Errorf("missing CIDR to be checked against")
	}

	if cidr2 == "" {
		return false, fmt.Errorf("missing CIDR that needs to be checked")
	}

	ip1, net1, err := net.ParseCIDR(cidr1)
	if err != nil {
		return false, fmt.Errorf("failed to parse the CIDR to be checked against: %q", err)
	}

	maskLen1, _ := net1.Mask.Size()
	if ip1.To4().Equal(net.IPv4zero) && maskLen1 == 0 {
		return false, fmt.Errorf("CIDR to be checked against is not in its canonical form")
	}

	ip2, net2, err := net.ParseCIDR(cidr2)
	if err != nil {
		return false, fmt.Errorf("failed to parse the CIDR that needs to be checked: %q", err)
	}

	maskLen2, _ := net2.Mask.Size()
	if ip2.To4().Equal(net.IPv4zero) && maskLen2 == 0 {
		return false, fmt.Errorf("CIDR that needs to be checked is not in its canonical form")
	}

	// If the mask length of the CIDR that needs to be checked is smaller
	// then the mask length of the CIDR to be checked against, then the
	// former will encompass more IPs than the latter, and hence can't be a
	// subset of the latter.
	if maskLen2 < maskLen1 {
		return false, nil
	}

	belongs, err := IPBelongsToCIDR(net2.IP.String(), cidr1)
	if err != nil {
		return false, err
	}

	return belongs, nil
}

// SubsetBlocks checks if each CIDR block of a given set of CIDR blocks, is a
// subset of at least one CIDR block belonging to another set of CIDR blocks.
// First parameter is the set of CIDR blocks to check against and the second
// parameter is the set of CIDR blocks that needs to be checked.
func SubsetBlocks(cidrBlocks1, cidrBlocks2 []string) (bool, error) {
	if len(cidrBlocks1) == 0 {
		return false, fmt.Errorf("missing CIDR blocks to be checked against")
	}

	if len(cidrBlocks2) == 0 {
		return false, fmt.Errorf("missing CIDR blocks that needs to be checked")
	}

	// Check if all the elements of cidrBlocks2 is a subset of at least one
	// element of cidrBlocks1
	for _, cidrBlock2 := range cidrBlocks2 {
		isSubset := false
		for _, cidrBlock1 := range cidrBlocks1 {
			subset, err := Subset(cidrBlock1, cidrBlock2)
			if err != nil {
				return false, err
			}
			// If CIDR is a subset of any of the CIDR block, its
			// good enough. Break out.
			if subset {
				isSubset = true
				break
			}
		}
		// CIDR block was not a subset of any of the CIDR blocks in the
		// set of blocks to check against
		if !isSubset {
			return false, nil
		}
	}

	return true, nil
}
<commit_msg>Address review feedback<commit_after>package cidrutil

import (
	"fmt"
	"net"
	"strings"

	"github.com/hashicorp/vault/helper/strutil"
)

// IPBelongsToCIDR checks if the given IP is encompassed by the given CIDR block
func IPBelongsToCIDR(ipAddr string, cidr string) (bool, error) {
	if ipAddr == "" {
		return false, fmt.Errorf("missing IP address")
	}

	ip := net.ParseIP(ipAddr)
	if ip == nil {
		return false, fmt.Errorf("invalid IP address")
	}

	_, ipnet, err := net.ParseCIDR(cidr)
	if err != nil {
		return false, err
	}

	if !ipnet.Contains(ip) {
		return false, nil
	}

	return true, nil
}

// IPBelongsToCIDRBlocksString checks if the given IP is encompassed by any of
// the given CIDR blocks, when the input is a string composed by joining all
// the CIDR blocks using a separator. The input is separated based on the given
// separator and the IP is checked to be belonged by any CIDR block.
func IPBelongsToCIDRBlocksString(ipAddr string, cidrList, separator string) (bool, error) {
	if ipAddr == "" {
		return false, fmt.Errorf("missing IP address")
	}

	if cidrList == "" {
		return false, fmt.Errorf("missing CIDR list")
	}

	if separator == "" {
		return false, fmt.Errorf("missing separator")
	}

	if ip := net.ParseIP(ipAddr); ip == nil {
		return false, fmt.Errorf("invalid IP address")
	}

	return IPBelongsToCIDRBlocksSlice(ipAddr, strutil.ParseDedupAndSortStrings(cidrList, separator))
}

// IPBelongsToCIDRBlocksSlice checks if the given IP is encompassed by any of the given
// CIDR blocks
func IPBelongsToCIDRBlocksSlice(ipAddr string, cidrs []string) (bool, error) {
	if ipAddr == "" {
		return false, fmt.Errorf("missing IP address")
	}

	if len(cidrs) == 0 {
		return false, fmt.Errorf("missing CIDR blocks to be checked against")
	}

	if ip := net.ParseIP(ipAddr); ip == nil {
		return false, fmt.Errorf("invalid IP address")
	}

	for _, cidr := range cidrs {
		belongs, err := IPBelongsToCIDR(ipAddr, cidr)
		if err != nil {
			return false, err
		}
		if belongs {
			return true, nil
		}
	}

	return false, nil
}

// ValidateCIDRListString checks if the list of CIDR blocks are valid, given
// that the input is a string composed by joining all the CIDR blocks using a
// separator. The input is separated based on the given separator and validity
// of each is checked.
func ValidateCIDRListString(cidrList string, separator string) (bool, error) {
	if cidrList == "" {
		return false, fmt.Errorf("missing CIDR list that needs validation")
	}
	if separator == "" {
		return false, fmt.Errorf("missing separator")
	}

	return ValidateCIDRListSlice(strutil.ParseDedupAndSortStrings(cidrList, separator))
}

// ValidateCIDRListSlice checks if the given list of CIDR blocks are valid
func ValidateCIDRListSlice(cidrBlocks []string) (bool, error) {
	if len(cidrBlocks) == 0 {
		return false, fmt.Errorf("missing CIDR blocks that needs validation")
	}

	for _, block := range cidrBlocks {
		if _, _, err := net.ParseCIDR(strings.TrimSpace(block)); err != nil {
			return false, err
		}
	}

	return true, nil
}

// Subset checks if the IPs belonging to a given CIDR block is a subset of IPs
// belonging to another CIDR block.
func Subset(cidr1, cidr2 string) (bool, error) {
	if cidr1 == "" {
		return false, fmt.Errorf("missing CIDR to be checked against")
	}

	if cidr2 == "" {
		return false, fmt.Errorf("missing CIDR that needs to be checked")
	}

	ip1, net1, err := net.ParseCIDR(cidr1)
	if err != nil {
		return false, fmt.Errorf("failed to parse the CIDR to be checked against: %q", err)
	}

	zeroAddr := false
	if ip := ip1.To4(); ip != nil && ip.Equal(net.IPv4zero) {
		zeroAddr = true
	}
	if ip := ip1.To16(); ip != nil && ip.Equal(net.IPv6zero) {
		zeroAddr = true
	}

	maskLen1, _ := net1.Mask.Size()
	if zeroAddr && maskLen1 == 0 {
		return false, fmt.Errorf("CIDR to be checked against is not in its canonical form")
	}

	ip2, net2, err := net.ParseCIDR(cidr2)
	if err != nil {
		return false, fmt.Errorf("failed to parse the CIDR that needs to be checked: %q", err)
	}

	zeroAddr = false
	if ip := ip2.To4(); ip != nil && ip.Equal(net.IPv4zero) {
		zeroAddr = true
	}
	if ip := ip2.To16(); ip != nil && ip.Equal(net.IPv6zero) {
		zeroAddr = true
	}

	maskLen2, _ := net2.Mask.Size()
	if zeroAddr && maskLen2 == 0 {
		return false, fmt.Errorf("CIDR that needs to be checked is not in its canonical form")
	}

	// If the mask length of the CIDR that needs to be checked is smaller
	// then the mask length of the CIDR to be checked against, then the
	// former will encompass more IPs than the latter, and hence can't be a
	// subset of the latter.
	if maskLen2 < maskLen1 {
		return false, nil
	}

	belongs, err := IPBelongsToCIDR(net2.IP.String(), cidr1)
	if err != nil {
		return false, err
	}

	return belongs, nil
}

// SubsetBlocks checks if each CIDR block of a given set of CIDR blocks, is a
// subset of at least one CIDR block belonging to another set of CIDR blocks.
// First parameter is the set of CIDR blocks to check against and the second
// parameter is the set of CIDR blocks that needs to be checked.
func SubsetBlocks(cidrBlocks1, cidrBlocks2 []string) (bool, error) {
	if len(cidrBlocks1) == 0 {
		return false, fmt.Errorf("missing CIDR blocks to be checked against")
	}

	if len(cidrBlocks2) == 0 {
		return false, fmt.Errorf("missing CIDR blocks that needs to be checked")
	}

	// Check if all the elements of cidrBlocks2 is a subset of at least one
	// element of cidrBlocks1
	for _, cidrBlock2 := range cidrBlocks2 {
		isSubset := false
		for _, cidrBlock1 := range cidrBlocks1 {
			subset, err := Subset(cidrBlock1, cidrBlock2)
			if err != nil {
				return false, err
			}
			// If CIDR is a subset of any of the CIDR block, its
			// good enough. Break out.
			if subset {
				isSubset = true
				break
			}
		}
		// CIDR block was not a subset of any of the CIDR blocks in the
		// set of blocks to check against
		if !isSubset {
			return false, nil
		}
	}

	return true, nil
}
<|endoftext|>


<commit_before>// Copyright 2011 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package tls

import (
	"crypto/x509"
	"runtime"
	"testing"
)

var tlsServers = []string{
	"google.com",
	"github.com",
	"twitter.com",
}

func TestOSCertBundles(t *testing.T) {
	if testing.Short() {
		t.Logf("skipping certificate tests in short mode")
		return
	}

	for _, addr := range tlsServers {
		conn, err := Dial("tcp", addr+":443", &Config{ServerName: addr})
		if err != nil {
			t.Errorf("unable to verify %v: %v", addr, err)
			continue
		}
		err = conn.Close()
		if err != nil {
			t.Error(err)
		}
	}
}

func TestCertHostnameVerifyWindows(t *testing.T) {
	if runtime.GOOS != "windows" {
		return
	}

	if testing.Short() {
		t.Logf("skipping certificate tests in short mode")
		return
	}

	for _, addr := range tlsServers {
		cfg := &Config{ServerName: "example.com"}
		conn, err := Dial("tcp", addr+":443", cfg)
		if err == nil {
			conn.Close()
			t.Errorf("should fail to verify for example.com: %v", addr)
			continue
		}
		_, ok := err.(x509.HostnameError)
		if !ok {
			t.Errorf("error type mismatch, got: %v", err)
		}
	}
}
<commit_msg>crypto/tls: remove flakey tests<commit_after><|endoftext|>