<commit_before>// Copyright 2017 Apcera Inc. All rights reserved.

package server

import (
	"fmt"
	"io"
	"sync/atomic"
	"time"

	"github.com/hashicorp/raft"
	"github.com/nats-io/go-nats"
	"github.com/nats-io/go-nats-streaming/pb"
	"github.com/nats-io/nats-streaming-server/spb"
	"github.com/nats-io/nats-streaming-server/util"
)

// serverSnapshot implements the raft.FSMSnapshot interface by snapshotting
// StanServer state.
type serverSnapshot struct {
	*StanServer
}

func newServerSnapshot(s *StanServer) raft.FSMSnapshot {
	return &serverSnapshot{s}
}

// Persist should dump all necessary state to the WriteCloser 'sink',
// and call sink.Close() when finished or call sink.Cancel() on error.
func (s *serverSnapshot) Persist(sink raft.SnapshotSink) (err error) {
	defer func() {
		if err != nil {
			sink.Cancel()
		}
	}()

	snap := &spb.RaftSnapshot{}

	s.snapshotClients(snap, sink)

	if err := s.snapshotChannels(snap, sink); err != nil {
		return err
	}

	var b []byte
	for i := 0; i < 2; i++ {
		b, err = snap.Marshal()
		if err != nil {
			return err
		}
		// Raft assumes that the follower will restore the snapshot from the leader using
		// a timeout that is equal to:
		// the transport timeout (we provide 2*time.Second) * (snapshot size / TimeoutScale).
		// We can't provide an infinite timeout to the nats-log transport otherwise some
		// Raft operations will block forever.
		// However, we persist only the first/last sequence for a channel snapshot, however,
		// the follower will request the leader to send all those messages as part of the
		// restore. Since the snapshot size may be small compared to the amout of messages
		// to recover, the timeout may be too small.
		// To trick the system, we first set the transport's TimeoutScale to 1 (the default
		// is 256KB). Then, if we want an overall timeout of 1 hour (3600 seconds), we need
		// the size to be at least 1800 bytes. If it is less than that, then we add some
		// padding to the snapshot.
		if len(b) < 1800 {
			snap.Padding = make([]byte, 1800-len(b))
			continue
		}
		break
	}

	var sizeBuf [4]byte
	util.ByteOrder.PutUint32(sizeBuf[:], uint32(len(b)))
	if _, err := sink.Write(sizeBuf[:]); err != nil {
		return err
	}
	if _, err := sink.Write(b); err != nil {
		return err
	}

	return sink.Close()
}

func (s *serverSnapshot) snapshotClients(snap *spb.RaftSnapshot, sink raft.SnapshotSink) {
	s.clients.RLock()
	defer s.clients.RUnlock()

	numClients := len(s.clients.clients)
	if numClients == 0 {
		return
	}

	snap.Clients = make([]*spb.ClientInfo, numClients)
	i := 0
	for _, client := range s.clients.clients {
		// Make a copy
		info := client.info.ClientInfo
		snap.Clients[i] = &info
		i++
	}
}

func (s *serverSnapshot) snapshotChannels(snap *spb.RaftSnapshot, sink raft.SnapshotSink) error {
	s.channels.RLock()
	defer s.channels.RUnlock()

	numChannels := len(s.channels.channels)
	if numChannels == 0 {
		return nil
	}

	snapshotASub := func(sub *subState) *spb.SubscriptionSnapshot {
		// Make a copy
		state := sub.SubState
		snapSub := &spb.SubscriptionSnapshot{State: &state}
		if len(sub.acksPending) > 0 {
			snapSub.AcksPending = make([]uint64, len(sub.acksPending))
			i := 0
			for seq := range sub.acksPending {
				snapSub.AcksPending[i] = seq
				i++
			}
		}
		return snapSub
	}

	snap.Channels = make([]*spb.ChannelSnapshot, numChannels)
	numChannel := 0
	for _, c := range s.channels.channels {
		first, last, err := c.store.Msgs.FirstAndLastSequence()
		if err != nil {
			return err
		}
		snapChannel := &spb.ChannelSnapshot{
			Channel: c.name,
			First:   first,
			Last:    last,
		}
		c.ss.RLock()

		// Start with count of all plain subs...
		snapSubs := make([]*spb.SubscriptionSnapshot, len(c.ss.psubs))
		i := 0
		for _, sub := range c.ss.psubs {
			sub.RLock()
			snapSubs[i] = snapshotASub(sub)
			sub.RUnlock()
			i++
		}

		// Now need to close durables
		for _, dur := range c.ss.durables {
			dur.RLock()
			if dur.IsClosed {
				// We need to persist a SubState with a ClientID
				// so that we can reconstruct the durable key
				// on recovery. So set to the saved value here
				// and then clear it after that.
				dur.ClientID = dur.savedClientID
				snapSubs = append(snapSubs, snapshotASub(dur))
				dur.ClientID = ""
			}
			dur.RUnlock()
		}

		// Snapshot the queue subscriptions
		for _, qsub := range c.ss.qsubs {
			qsub.RLock()
			for _, sub := range qsub.subs {
				sub.RLock()
				snapSubs = append(snapSubs, snapshotASub(sub))
				sub.RUnlock()
			}
			// If all members of a durable queue group left the group,
			// we need to persist the "shadow" queue member.
			if qsub.shadow != nil {
				qsub.shadow.RLock()
				snapSubs = append(snapSubs, snapshotASub(qsub.shadow))
				qsub.shadow.RUnlock()
			}
			qsub.RUnlock()
		}
		if len(snapSubs) > 0 {
			snapChannel.Subscriptions = snapSubs
		}

		c.ss.RUnlock()
		snap.Channels[numChannel] = snapChannel
		numChannel++
	}

	return nil
}

// Release is a no-op.
func (s *serverSnapshot) Release() {}

// restoreFromSnapshot restores a server from a snapshot. This is not called
// concurrently with any other Raft commands.
func (s *StanServer) restoreFromSnapshot(snapshot io.ReadCloser) error {
	defer snapshot.Close()

	restoreFromRaftInit := atomic.LoadInt64(&s.raftNodeCreated) == 0

	// We need to drop current state. The server will recover from snapshot
	// and all newer Raft entry logs (basically the entire state is being
	// reconstructed from this point on).
	for _, c := range s.channels.getAll() {
		for _, sub := range c.ss.getAllSubs() {
			sub.RLock()
			clientID := sub.ClientID
			sub.RUnlock()
			if err := s.unsubscribeSub(c, clientID, "unsub", sub, false); err != nil {
				return err
			}
		}
	}
	for clientID := range s.clients.getClients() {
		if _, err := s.clients.unregister(clientID); err != nil {
			return err
		}
	}

	sizeBuf := make([]byte, 4)
	// Read the snapshot size.
	if _, err := io.ReadFull(snapshot, sizeBuf); err != nil {
		if err == io.EOF {
			return nil
		}
		return err
	}
	// Read the snapshot.
	size := util.ByteOrder.Uint32(sizeBuf)
	buf := make([]byte, size)
	if _, err := io.ReadFull(snapshot, buf); err != nil {
		return err
	}

	serverSnap := &spb.RaftSnapshot{}
	if err := serverSnap.Unmarshal(buf); err != nil {
		panic(err)
	}
	if err := s.restoreClientsFromSnapshot(serverSnap); err != nil {
		return err
	}
	return s.restoreChannelsFromSnapshot(serverSnap, restoreFromRaftInit)
}

func (s *StanServer) restoreClientsFromSnapshot(serverSnap *spb.RaftSnapshot) error {
	for _, sc := range serverSnap.Clients {
		if _, err := s.clients.register(sc.ID, sc.HbInbox); err != nil {
			return err
		}
	}
	return nil
}

func (s *StanServer) restoreChannelsFromSnapshot(serverSnap *spb.RaftSnapshot, restoreFromRaftInit bool) error {
	for _, sc := range serverSnap.Channels {
		c, err := s.lookupOrCreateChannel(sc.Channel)
		if err != nil {
			return err
		}
		// Do not restore messages from snapshot if the server
		// just started and is recovering from its own snapshot.
		if !restoreFromRaftInit {
			if err := s.restoreMsgsFromSnapshot(c, sc.First, sc.Last); err != nil {
				return err
			}
		}
		for _, ss := range sc.Subscriptions {
			s.recoverOneSub(c, ss.State, nil, ss.AcksPending)
		}
	}
	return nil
}

func (s *StanServer) restoreMsgsFromSnapshot(c *channel, first, last uint64) error {
	if err := c.store.Msgs.Empty(); err != nil {
		return err
	}
	inbox := nats.NewInbox()
	sub, err := c.stan.ncsr.SubscribeSync(inbox)
	if err != nil {
		return err
	}
	sub.SetPendingLimits(-1, -1)
	defer sub.Unsubscribe()

	subject := fmt.Sprintf("%s.%s.%s", defaultSnapshotPrefix, c.stan.info.ClusterID, c.name)

	var (
		reqBuf   [16]byte
		reqNext  = first
		reqStart = first
		reqEnd   uint64
	)
	for seq := first; seq <= last; seq++ {
		if seq == reqNext {
			reqEnd = reqStart + uint64(100)
			if reqEnd > last {
				reqEnd = last
			}
			util.ByteOrder.PutUint64(reqBuf[:8], reqStart)
			util.ByteOrder.PutUint64(reqBuf[8:16], reqEnd)
			if err := c.stan.ncsr.PublishRequest(subject, inbox, reqBuf[:16]); err != nil {
				return err
			}
			if reqEnd != last {
				reqNext = reqEnd - reqStart/2
				reqStart = reqEnd + 1
			}
		}
		resp, err := sub.NextMsg(2 * time.Second)
		if err != nil {
			return err
		}
		// It is possible that the leader does not have this message because of
		// channel limits. If resp.Data is empty, we are in this situation and
		// we are done recovering snapshot.
		if len(resp.Data) == 0 {
			break
		}
		msg := &pb.MsgProto{}
		if err := msg.Unmarshal(resp.Data); err != nil {
			panic(err)
		}
		if _, err := c.store.Msgs.Store(msg); err != nil {
			return err
		}
	}
	return c.store.Msgs.Flush()
}
<commit_msg>Empty message store only when needed on snapshot restore<commit_after>// Copyright 2017 Apcera Inc. All rights reserved.

package server

import (
	"fmt"
	"io"
	"sync/atomic"
	"time"

	"github.com/hashicorp/raft"
	"github.com/nats-io/go-nats"
	"github.com/nats-io/go-nats-streaming/pb"
	"github.com/nats-io/nats-streaming-server/spb"
	"github.com/nats-io/nats-streaming-server/util"
)

// serverSnapshot implements the raft.FSMSnapshot interface by snapshotting
// StanServer state.
type serverSnapshot struct {
	*StanServer
}

func newServerSnapshot(s *StanServer) raft.FSMSnapshot {
	return &serverSnapshot{s}
}

// Persist should dump all necessary state to the WriteCloser 'sink',
// and call sink.Close() when finished or call sink.Cancel() on error.
func (s *serverSnapshot) Persist(sink raft.SnapshotSink) (err error) {
	defer func() {
		if err != nil {
			sink.Cancel()
		}
	}()

	snap := &spb.RaftSnapshot{}

	s.snapshotClients(snap, sink)

	if err := s.snapshotChannels(snap, sink); err != nil {
		return err
	}

	var b []byte
	for i := 0; i < 2; i++ {
		b, err = snap.Marshal()
		if err != nil {
			return err
		}
		// Raft assumes that the follower will restore the snapshot from the leader using
		// a timeout that is equal to:
		// the transport timeout (we provide 2*time.Second) * (snapshot size / TimeoutScale).
		// We can't provide an infinite timeout to the nats-log transport otherwise some
		// Raft operations will block forever.
		// However, we persist only the first/last sequence for a channel snapshot, however,
		// the follower will request the leader to send all those messages as part of the
		// restore. Since the snapshot size may be small compared to the amout of messages
		// to recover, the timeout may be too small.
		// To trick the system, we first set the transport's TimeoutScale to 1 (the default
		// is 256KB). Then, if we want an overall timeout of 1 hour (3600 seconds), we need
		// the size to be at least 1800 bytes. If it is less than that, then we add some
		// padding to the snapshot.
		if len(b) < 1800 {
			snap.Padding = make([]byte, 1800-len(b))
			continue
		}
		break
	}

	var sizeBuf [4]byte
	util.ByteOrder.PutUint32(sizeBuf[:], uint32(len(b)))
	if _, err := sink.Write(sizeBuf[:]); err != nil {
		return err
	}
	if _, err := sink.Write(b); err != nil {
		return err
	}

	return sink.Close()
}

func (s *serverSnapshot) snapshotClients(snap *spb.RaftSnapshot, sink raft.SnapshotSink) {
	s.clients.RLock()
	defer s.clients.RUnlock()

	numClients := len(s.clients.clients)
	if numClients == 0 {
		return
	}

	snap.Clients = make([]*spb.ClientInfo, numClients)
	i := 0
	for _, client := range s.clients.clients {
		// Make a copy
		info := client.info.ClientInfo
		snap.Clients[i] = &info
		i++
	}
}

func (s *serverSnapshot) snapshotChannels(snap *spb.RaftSnapshot, sink raft.SnapshotSink) error {
	s.channels.RLock()
	defer s.channels.RUnlock()

	numChannels := len(s.channels.channels)
	if numChannels == 0 {
		return nil
	}

	snapshotASub := func(sub *subState) *spb.SubscriptionSnapshot {
		// Make a copy
		state := sub.SubState
		snapSub := &spb.SubscriptionSnapshot{State: &state}
		if len(sub.acksPending) > 0 {
			snapSub.AcksPending = make([]uint64, len(sub.acksPending))
			i := 0
			for seq := range sub.acksPending {
				snapSub.AcksPending[i] = seq
				i++
			}
		}
		return snapSub
	}

	snap.Channels = make([]*spb.ChannelSnapshot, numChannels)
	numChannel := 0
	for _, c := range s.channels.channels {
		first, last, err := c.store.Msgs.FirstAndLastSequence()
		if err != nil {
			return err
		}
		snapChannel := &spb.ChannelSnapshot{
			Channel: c.name,
			First:   first,
			Last:    last,
		}
		c.ss.RLock()

		// Start with count of all plain subs...
		snapSubs := make([]*spb.SubscriptionSnapshot, len(c.ss.psubs))
		i := 0
		for _, sub := range c.ss.psubs {
			sub.RLock()
			snapSubs[i] = snapshotASub(sub)
			sub.RUnlock()
			i++
		}

		// Now need to close durables
		for _, dur := range c.ss.durables {
			dur.RLock()
			if dur.IsClosed {
				// We need to persist a SubState with a ClientID
				// so that we can reconstruct the durable key
				// on recovery. So set to the saved value here
				// and then clear it after that.
				dur.ClientID = dur.savedClientID
				snapSubs = append(snapSubs, snapshotASub(dur))
				dur.ClientID = ""
			}
			dur.RUnlock()
		}

		// Snapshot the queue subscriptions
		for _, qsub := range c.ss.qsubs {
			qsub.RLock()
			for _, sub := range qsub.subs {
				sub.RLock()
				snapSubs = append(snapSubs, snapshotASub(sub))
				sub.RUnlock()
			}
			// If all members of a durable queue group left the group,
			// we need to persist the "shadow" queue member.
			if qsub.shadow != nil {
				qsub.shadow.RLock()
				snapSubs = append(snapSubs, snapshotASub(qsub.shadow))
				qsub.shadow.RUnlock()
			}
			qsub.RUnlock()
		}
		if len(snapSubs) > 0 {
			snapChannel.Subscriptions = snapSubs
		}

		c.ss.RUnlock()
		snap.Channels[numChannel] = snapChannel
		numChannel++
	}

	return nil
}

// Release is a no-op.
func (s *serverSnapshot) Release() {}

// restoreFromSnapshot restores a server from a snapshot. This is not called
// concurrently with any other Raft commands.
func (s *StanServer) restoreFromSnapshot(snapshot io.ReadCloser) error {
	defer snapshot.Close()

	restoreFromRaftInit := atomic.LoadInt64(&s.raftNodeCreated) == 0

	// We need to drop current state. The server will recover from snapshot
	// and all newer Raft entry logs (basically the entire state is being
	// reconstructed from this point on).
	for _, c := range s.channels.getAll() {
		for _, sub := range c.ss.getAllSubs() {
			sub.RLock()
			clientID := sub.ClientID
			sub.RUnlock()
			if err := s.unsubscribeSub(c, clientID, "unsub", sub, false); err != nil {
				return err
			}
		}
	}
	for clientID := range s.clients.getClients() {
		if _, err := s.clients.unregister(clientID); err != nil {
			return err
		}
	}

	sizeBuf := make([]byte, 4)
	// Read the snapshot size.
	if _, err := io.ReadFull(snapshot, sizeBuf); err != nil {
		if err == io.EOF {
			return nil
		}
		return err
	}
	// Read the snapshot.
	size := util.ByteOrder.Uint32(sizeBuf)
	buf := make([]byte, size)
	if _, err := io.ReadFull(snapshot, buf); err != nil {
		return err
	}

	serverSnap := &spb.RaftSnapshot{}
	if err := serverSnap.Unmarshal(buf); err != nil {
		panic(err)
	}
	if err := s.restoreClientsFromSnapshot(serverSnap); err != nil {
		return err
	}
	return s.restoreChannelsFromSnapshot(serverSnap, restoreFromRaftInit)
}

func (s *StanServer) restoreClientsFromSnapshot(serverSnap *spb.RaftSnapshot) error {
	for _, sc := range serverSnap.Clients {
		if _, err := s.clients.register(sc.ID, sc.HbInbox); err != nil {
			return err
		}
	}
	return nil
}

func (s *StanServer) restoreChannelsFromSnapshot(serverSnap *spb.RaftSnapshot, restoreFromRaftInit bool) error {
	for _, sc := range serverSnap.Channels {
		c, err := s.lookupOrCreateChannel(sc.Channel)
		if err != nil {
			return err
		}
		// Do not restore messages from snapshot if the server
		// just started and is recovering from its own snapshot.
		if !restoreFromRaftInit {
			if err := s.restoreMsgsFromSnapshot(c, sc.First, sc.Last); err != nil {
				return err
			}
		}
		for _, ss := range sc.Subscriptions {
			s.recoverOneSub(c, ss.State, nil, ss.AcksPending)
		}
	}
	return nil
}

func (s *StanServer) restoreMsgsFromSnapshot(c *channel, first, last uint64) error {
	storeFirst, storeLast, err := c.store.Msgs.FirstAndLastSequence()
	if err != nil {
		return err
	}
	// If the leader's first sequence is more than our lastSequence+1,
	// then we need to empty the store. We don't want to have gaps.
	// Same if our first is strictly greater than the leader, or our
	// last sequence is more than the leader
	if first > storeLast+1 || storeFirst > first || storeLast > last {
		if err := c.store.Msgs.Empty(); err != nil {
			return err
		}
	} else if storeLast == last {
		// We may have a message with lower sequence than the leader,
		// but our last sequence is the same, so nothing to do.
		return nil
	} else if storeLast > 0 {
		// first is less than what we already have, just started
		// at our next sequence.
		first = storeLast + 1
	}
	inbox := nats.NewInbox()
	sub, err := c.stan.ncsr.SubscribeSync(inbox)
	if err != nil {
		return err
	}
	sub.SetPendingLimits(-1, -1)
	defer sub.Unsubscribe()

	subject := fmt.Sprintf("%s.%s.%s", defaultSnapshotPrefix, c.stan.info.ClusterID, c.name)

	var (
		reqBuf   [16]byte
		reqNext  = first
		reqStart = first
		reqEnd   uint64
	)
	for seq := first; seq <= last; seq++ {
		if seq == reqNext {
			reqEnd = reqStart + uint64(100)
			if reqEnd > last {
				reqEnd = last
			}
			util.ByteOrder.PutUint64(reqBuf[:8], reqStart)
			util.ByteOrder.PutUint64(reqBuf[8:16], reqEnd)
			if err := c.stan.ncsr.PublishRequest(subject, inbox, reqBuf[:16]); err != nil {
				return err
			}
			if reqEnd != last {
				reqNext = reqEnd - reqStart/2
				reqStart = reqEnd + 1
			}
		}
		resp, err := sub.NextMsg(2 * time.Second)
		if err != nil {
			return err
		}
		// It is possible that the leader does not have this message because of
		// channel limits. If resp.Data is empty, we are in this situation and
		// we are done recovering snapshot.
		if len(resp.Data) == 0 {
			break
		}
		msg := &pb.MsgProto{}
		if err := msg.Unmarshal(resp.Data); err != nil {
			panic(err)
		}
		if _, err := c.store.Msgs.Store(msg); err != nil {
			return err
		}
	}
	return c.store.Msgs.Flush()
}
<|endoftext|>


<commit_before>package usergrid

import (
	"encoding/json"
	"strconv"
	"testing"
	"time"
)
func TestPost(t *testing.T){
	client := Client {Organization:"yourorgname",Application:"sandbox",Uri:"https://api.usergrid.com"}
	data := map[string]string{"name":"go_dog", "status":"good dog"}
	resp, err:= client.Post("dogs", nil, data)
	if(err!=nil){
		t.Logf("TestPost failed: %s\n", err)
		t.Fail()
	}
	if(t.Failed()){
		str,_:=json.MarshalIndent(resp,"","  ")
		t.Logf("RESPONSE: %s", str)
	}
}
func TestPostDuplicate(t *testing.T){
	client := Client {Organization:"yourorgname",Application:"sandbox",Uri:"https://api.usergrid.com"}
	data := map[string]string{"name":"go_dog", "status":"good dog"}
	resp, err:= client.Post("dogs", nil, data)
	if(err==nil){
		t.Logf("TestPost failed. Should have received a duplicate entity error.\n")
		t.Fail()
	}
	if(t.Failed()){
		str,_:=json.MarshalIndent(resp,"","  ")
		t.Logf("RESPONSE: %s", str)
	}
}
func TestGet(t *testing.T){
	client := Client {Organization:"yourorgname",Application:"sandbox",Uri:"https://api.usergrid.com"}
	resp, err:= client.Get("dogs/go_dog", nil)
	if(err!=nil){
		t.Logf("TestGet failed: %s\n", err)
		t.Fail()
	}else{
		if (resp["entities"]==nil) {
			t.Logf("No entities returned")
			t.Fail()
		}else if(len(resp["entities"].([]interface{}))!=1){
			t.Logf("Incorrect number of entities")
			t.Fail()
		}
	}
	if(t.Failed()){
		str,_:=json.MarshalIndent(resp,"","  ")
		t.Logf("RESPONSE: %s", str)
	}
}
func TestGetBadEntity(t *testing.T){
	client := Client {Organization:"yourorgname",Application:"sandbox",Uri:"https://api.usergrid.com"}
	resp, err:= client.Get("dogs/go_dogizzle", nil)
	if(err==nil){
		t.Logf("TestGet failed. dog should not exist: %s\n", err)
		t.Fail()
	}
	if(t.Failed()){
		str,_:=json.MarshalIndent(resp,"","  ")
		t.Logf("RESPONSE: %s", str)
	}
}
func TestPut(t *testing.T){
	client := Client {Organization:"yourorgname",Application:"sandbox",Uri:"https://api.usergrid.com"}
	data := map[string]string{"status":"flat out on the highway and drying in the sun"}
	resp, err:= client.Put("dogs/go_dog", nil, data)
	if(err!=nil){
		t.Logf("TestPut failed: %s\n", err)
		t.Fail()
	}
	if(t.Failed()){
		str,_:=json.MarshalIndent(resp,"","  ")
		t.Logf("RESPONSE: %s", str)
	}
}
func TestDelete(t *testing.T){
	client := Client {Organization:"yourorgname",Application:"sandbox",Uri:"https://api.usergrid.com"}
	resp, err:= client.Delete("dogs/go_dog", nil)
	if(err!=nil){
		t.Logf("Test failed: %s\n", err)
		t.Fail()
	}
	if(t.Failed()){
		str,_:=json.MarshalIndent(resp,"","  ")
		t.Logf("RESPONSE: %s", str)
	}
}
func TestDelete2(t *testing.T){
	client := Client {Organization:"yourorgname",Application:"sandbox",Uri:"https://api.usergrid.com"}
	resp, err:= client.Get("dogs/go_dog", nil)
	if(err==nil){
		t.Logf("Test failed. Dog should be deleted: %s\n", err)
		t.Fail()
	}
	if(t.Failed()){
		str,_:=json.MarshalIndent(resp,"","  ")
		t.Logf("RESPONSE: %s", str)
	}
}
func TestDelete3(t *testing.T){
	client := Client {Organization:"yourorgname",Application:"sandbox",Uri:"https://api.usergrid.com"}
	resp, err:= client.Delete("dogs/go_dog", nil)
	if(err==nil){
		t.Logf("Test failed. Should return service_resource_not_found error\n")
		t.Fail()
	}
	if(t.Failed()){
		str,_:=json.MarshalIndent(resp,"","  ")
		t.Logf("RESPONSE: %s", str)
	}
}

func BenchmarkPost(b *testing.B) {
 	client := Client {Organization:"yourorgname",Application:"sandbox",Uri:"https://api.usergrid.com"}
	data := map[string]string{"index":"0", "description":"golang benchmark"}
	b.ResetTimer()
    for i := 0; i < b.N; i++ {
    	data["index"]=strconv.Itoa(i);
		resp, err:= client.Post("benchmark", nil, data)
		if(err!=nil){
			b.Logf("BenchmarkPost failed: %s\n", err)
			b.Fail()
		}
		if(b.Failed()){
			str,_:=json.MarshalIndent(resp,"","  ")
			b.Logf("RESPONSE: %s", str)
		}
   }
}
func BenchmarkRawRequests(b *testing.B) {
 	client := Client {Organization:"yourorgname",Application:"sandbox",Uri:"https://api.usergrid.com"}
	params := map[string]string{"limit":strconv.Itoa(10)}
	resp, err:= client.Get("benchmark", params)
	if(err!=nil){
		b.Logf("Test failed: %s\n", err)
		b.Fail()
	}
	if (resp["entities"]==nil || len(resp["entities"].([]interface{}))==0) {
		b.Logf("Test failed: no entities\n")
		b.Fail()
	}
	entities:=resp["entities"].([]interface{})
	requests:=0
	responses:=0
	errors:=0
	var objmap interface{}
	var entity map[string]interface{}
	responseChan := make(chan []byte)
	go func(){
		for {
	        select {
	        case v := <-responseChan:
	        	if err := json.Unmarshal(v, &objmap); err == nil{
					if err:=CheckForError(&objmap); err != nil {
						errors++
					}
				}else{
					errors++
				}

	        	responses++
	        	requests--
	        case <-time.After(time.Second * 10):
	        	return
	        }
	    }
	}()
 	b.ResetTimer()
    for i := 0; i < b.N; i++ {
    	if(len(entities)==0){
			b.Logf("Test failed: we ran out of entities\n")
			b.Fail()
			continue
    	}
    	entity = entities[i % len(entities)].(map[string]interface{})
    	for ;requests>=MAX_CONCURRENT_REQUESTS;{
        	// if we outpace GOMAXPROCS, we'll run out of threads
    		time.Sleep(60 * time.Millisecond)
    	}
		client.Request("GET", "http://api.usergrid.com/yourorgname/sandbox/benchmark/"+entity["uuid"].(string), nil, nil, responseChan)
		requests++
   }
   for ;requests>0;{
		// wait for the last few responses
		time.Sleep(60 * time.Millisecond)
   }
}
func BenchmarkGet(b *testing.B) {
 	client := Client {Organization:"yourorgname",Application:"sandbox",Uri:"https://api.usergrid.com"}
	params := map[string]string{"limit":strconv.Itoa(500)}
	resp, err:= client.Get("benchmark", params)
	if(err!=nil){
		b.Logf("Test failed: %s\n", err)
		b.Fail()
	}
	if (resp["entities"]==nil || len(resp["entities"].([]interface{}))==0) {
		b.Logf("Test failed: no entities to delete\n")
		b.Fail()
	}
	entities:=resp["entities"].([]interface{})
	var entity map[string]interface{}
 	b.ResetTimer()
    for i := 0; i < b.N; i++ {
    	if(len(entities)==0){
			b.Logf("Test failed: we ran out of entities\n")
			b.Fail()
			continue
    	}
    	entity, entities = entities[len(entities)-1].(map[string]interface{}), entities[:len(entities)-1]
		_, err:= client.Get("benchmark/"+entity["uuid"].(string), nil)
		if(err!=nil){
			b.Logf("BenchmarkGet failed: %s\n", err)
			b.Fail()
		}
   }
	if(b.Failed()){
		str,_:=json.MarshalIndent(resp,"","  ")
		b.Logf("RESPONSE: %s", str)
	}
}
func BenchmarkPut(b *testing.B) {
 	client := Client {Organization:"yourorgname",Application:"sandbox",Uri:"https://api.usergrid.com"}
	params := map[string]string{"limit":strconv.Itoa(500)}
	resp, err:= client.Get("benchmark", params)
	if(err!=nil){
		b.Logf("Test failed: %s\n", err)
		b.Fail()
	}
	if (resp["entities"]==nil || len(resp["entities"].([]interface{}))==0) {
		b.Logf("Test failed: no entities to delete\n")
		b.Fail()
	}
	entities:=resp["entities"].([]interface{})
	var entity map[string]interface{}
 	b.ResetTimer()
    for i := 0; i < b.N; i++ {
    	if(len(entities)==0){
			b.Logf("Test failed: we ran out of entities\n")
			b.Fail()
			continue
    	}
    	entity, entities = entities[len(entities)-1].(map[string]interface{}), entities[:len(entities)-1]
		_, err:= client.Put("benchmark/"+entity["uuid"].(string), nil, map[string]string{"updated":"true"})
		if(err!=nil){
			b.Logf("BenchmarkPut failed: %s\n", err)
			b.Fail()
		}
   }
	if(b.Failed()){
		str,_:=json.MarshalIndent(resp,"","  ")
		b.Logf("RESPONSE: %s", str)
	}
}

func BenchmarkDelete(b *testing.B) {
 	client := Client {Organization:"yourorgname",Application:"sandbox",Uri:"https://api.usergrid.com"}
	params := map[string]string{"limit":strconv.Itoa(500)}
	resp, err:= client.Get("benchmark", params)
	if(err!=nil){
		b.Logf("Test failed: %s\n", err)
		b.Fail()
	}
	if (resp["entities"]==nil || len(resp["entities"].([]interface{}))==0) {
		b.Logf("Test failed: no entities to delete\n")
		b.Fail()
	}
	entities:=resp["entities"].([]interface{})
	var entity map[string]interface{}
 	b.ResetTimer()
    for i := 0; i < b.N; i++ {
    	if(len(entities)==0){
			b.Logf("Test failed: we ran out of entities\n")
			b.Fail()
			continue
    	}
    	entity, entities = entities[len(entities)-1].(map[string]interface{}), entities[:len(entities)-1]
		_, err:= client.Delete("benchmark/"+entity["uuid"].(string), nil)
		if(err!=nil){
			b.Logf("BenchmarkDelete failed: %s\n", err)
			b.Fail()
		}
   }
	if(b.Failed()){
		str,_:=json.MarshalIndent(resp,"","  ")
		b.Logf("RESPONSE: %s", str)
	}
}

<commit_msg>Updated tests following changes to Request.<commit_after>package usergrid

import (
	"encoding/json"
	"strconv"
	"testing"
	"time"
)
func TestPost(t *testing.T){
	var objmap interface{}
	client := Client {Organization:"yourorgname",Application:"sandbox",Uri:"https://api.usergrid.com"}
	data := map[string]string{"name":"go_dog", "status":"good dog"}
	err:= client.Post("dogs", nil, data, JSONResponseHandler(&objmap))
	if(err!=nil){
		t.Logf("TestPost failed: %s\n", err)
		t.Fail()
	}
	if(t.Failed()){
		str,_:=json.MarshalIndent(&objmap,"","  ")
		t.Logf("RESPONSE: %s", str)
	}
}
func TestPostDuplicate(t *testing.T){
	var objmap interface{}
	client := Client {Organization:"yourorgname",Application:"sandbox",Uri:"https://api.usergrid.com"}
	data := map[string]string{"name":"go_dog", "status":"good dog"}
	err:= client.Post("dogs", nil, data, JSONResponseHandler(&objmap))
	if(err==nil){
		t.Logf("TestPost failed. Should have received a duplicate entity error.\n")
		t.Fail()
	}
	if(t.Failed()){
		str,_:=json.MarshalIndent(&objmap,"","  ")
		t.Logf("RESPONSE: %s", str)
	}
}
func TestGet(t *testing.T){
	var objmap interface{}
	client := Client {Organization:"yourorgname",Application:"sandbox",Uri:"https://api.usergrid.com"}
	err:= client.Get("dogs/go_dog", nil, JSONResponseHandler(&objmap))
	if(err!=nil){
		t.Logf("TestGet failed: %s\n", err)
		t.Fail()
	}else{
		omap:=objmap.(map[string]interface{})
		if (omap["entities"]==nil) {
			t.Logf("No entities returned")
			t.Fail()
		}else if(len(omap["entities"].([]interface{}))!=1){
			t.Logf("Incorrect number of entities")
			t.Fail()
		}
	}
	if(t.Failed()){
		str,_:=json.MarshalIndent(&objmap,"","  ")
		t.Logf("RESPONSE: %s", str)
	}
}
func TestGetBadEntity(t *testing.T){
	var objmap interface{}
	client := Client {Organization:"yourorgname",Application:"sandbox",Uri:"https://api.usergrid.com"}
	err:= client.Get("dogs/go_dogizzle", nil, JSONResponseHandler(&objmap))
	if(err==nil){
		t.Logf("TestGet failed. dog should not exist: %s\n", err)
		t.Fail()
	}
	if(t.Failed()){
		str,_:=json.MarshalIndent(&objmap,"","  ")
		t.Logf("RESPONSE: %s", str)
	}
}
func TestPut(t *testing.T){
	var objmap interface{}
	client := Client {Organization:"yourorgname",Application:"sandbox",Uri:"https://api.usergrid.com"}
	data := map[string]string{"status":"flat out on the highway and drying in the sun"}
	err:= client.Put("dogs/go_dog", nil, data, JSONResponseHandler(&objmap))
	if(err!=nil){
		t.Logf("TestPut failed: %s\n", err)
		t.Fail()
	}
	if(t.Failed()){
		str,_:=json.MarshalIndent(&objmap,"","  ")
		t.Logf("RESPONSE: %s", str)
	}
}
func TestDelete(t *testing.T){
	var objmap interface{}
	client := Client {Organization:"yourorgname",Application:"sandbox",Uri:"https://api.usergrid.com"}
	err:= client.Delete("dogs/go_dog", nil, JSONResponseHandler(&objmap))
	if(err!=nil){
		t.Logf("Test failed: %s\n", err)
		t.Fail()
	}
	if(t.Failed()){
		str,_:=json.MarshalIndent(&objmap,"","  ")
		t.Logf("RESPONSE: %s", str)
	}
}
func TestDelete2(t *testing.T){
	var objmap interface{}
	client := Client {Organization:"yourorgname",Application:"sandbox",Uri:"https://api.usergrid.com"}
	err:= client.Get("dogs/go_dog", nil, JSONResponseHandler(&objmap))
	if(err==nil){
		t.Logf("Test failed. Dog should be deleted: %s\n", err)
		t.Fail()
	}
	if(t.Failed()){
		str,_:=json.MarshalIndent(&objmap,"","  ")
		t.Logf("RESPONSE: %s", str)
	}
}
func TestDelete3(t *testing.T){
	var objmap interface{}
	client := Client {Organization:"yourorgname",Application:"sandbox",Uri:"https://api.usergrid.com"}
	err:= client.Delete("dogs/go_dog", nil, JSONResponseHandler(&objmap))
	if(err==nil){
		t.Logf("Test failed. Should return service_resource_not_found error\n")
		t.Fail()
	}
	if(t.Failed()){
		str,_:=json.MarshalIndent(&objmap,"","  ")
		t.Logf("RESPONSE: %s", str)
	}
}

func BenchmarkPost(b *testing.B) {
	var objmap interface{}
 	client := Client {Organization:"yourorgname",Application:"sandbox",Uri:"https://api.usergrid.com"}
	data := map[string]string{"index":"0", "description":"golang benchmark"}
	b.ResetTimer()
    for i := 0; i < b.N; i++ {
    	data["index"]=strconv.Itoa(i);
		err:= client.Post("benchmark", nil, data, JSONResponseHandler(&objmap))
		if(err!=nil){
			b.Logf("BenchmarkPost failed: %s\n", err)
			b.Fail()
		}
		if(b.Failed()){
			str,_:=json.MarshalIndent(objmap,"","  ")
			b.Logf("RESPONSE: %s", str)
		}
   }
}
func BenchmarkRawRequests(b *testing.B) {
	var objmap interface{}
 	client := Client {Organization:"yourorgname",Application:"sandbox",Uri:"https://api.usergrid.com"}
	params := map[string]string{"limit":strconv.Itoa(10)}
	err:= client.Get("benchmark", params, JSONResponseHandler(&objmap))
	if(err!=nil){
		b.Logf("Test failed: %s\n", err)
		b.Fail()
	}
	omap:=objmap.(map[string]interface{})
	if (omap["entities"]==nil || len(omap["entities"].([]interface{}))==0) {
		b.Logf("Test failed: no entities\n")
		b.Fail()
	}
	entities:=omap["entities"].([]interface{})
	requests:=0
	responses:=0
	errors:=0
	// var objmap interface{}
	var entity map[string]interface{}
	responseChan := make(chan []byte)
	go func(){
		for {
	        select {
	        case v := <-responseChan:
	        	if err := json.Unmarshal(v, &objmap); err == nil{
					if err:=CheckForError(&objmap); err != nil {
						errors++
					}
				}else{
					errors++
				}

	        	responses++
	        	requests--
	        case <-time.After(time.Second * 10):
	        	return
	        }
	    }
	}()
 	b.ResetTimer()
    for i := 0; i < b.N; i++ {
    	if(len(entities)==0){
			b.Logf("Test failed: we ran out of entities\n")
			b.Fail()
			continue
    	}
    	entity = entities[i % len(entities)].(map[string]interface{})
    	for ;requests>=MAX_CONCURRENT_REQUESTS;{
        	// if we outpace GOMAXPROCS, we'll run out of threads
    		time.Sleep(60 * time.Millisecond)
    	}
		client.Request("GET", "http://api.usergrid.com/yourorgname/sandbox/benchmark/"+entity["uuid"].(string), nil, nil, responseChan)
		requests++
   }
   for ;requests>0;{
		// wait for the last few responses
		time.Sleep(60 * time.Millisecond)
   }
}
func BenchmarkGet(b *testing.B) {
	var objmap interface{}
 	client := Client {Organization:"yourorgname",Application:"sandbox",Uri:"https://api.usergrid.com"}
	params := map[string]string{"limit":strconv.Itoa(500)}
	err:= client.Get("benchmark", params, JSONResponseHandler(&objmap))
	if(err!=nil){
		b.Logf("Test failed: %s\n", err)
		b.Fail()
	}
	omap:=objmap.(map[string]interface{})
	if (omap["entities"]==nil || len(omap["entities"].([]interface{}))==0) {
		b.Logf("Test failed: no entities to delete\n")
		b.Fail()
	}
	entities:=omap["entities"].([]interface{})
	var entity map[string]interface{}
 	b.ResetTimer()
    for i := 0; i < b.N; i++ {
    	if(len(entities)==0){
			b.Logf("Test failed: we ran out of entities\n")
			b.Fail()
			continue
    	}
    	entity, entities = entities[len(entities)-1].(map[string]interface{}), entities[:len(entities)-1]
		err:= client.Get("benchmark/"+entity["uuid"].(string), nil, NOOPResponseHandler(&objmap))
		if(err!=nil){
			b.Logf("BenchmarkGet failed: %s\n", err)
			b.Fail()
		}
   }
	if(b.Failed()){
		str,_:=json.MarshalIndent(objmap,"","  ")
		b.Logf("RESPONSE: %s", str)
	}
}
func BenchmarkPut(b *testing.B) {
	var objmap interface{}
 	client := Client {Organization:"yourorgname",Application:"sandbox",Uri:"https://api.usergrid.com"}
	params := map[string]string{"limit":strconv.Itoa(500)}
	err:= client.Get("benchmark", params, JSONResponseHandler(&objmap))
	if(err!=nil){
		b.Logf("Test failed: %s\n", err)
		b.Fail()
	}
	omap:=objmap.(map[string]interface{})
	if (omap["entities"]==nil || len(omap["entities"].([]interface{}))==0) {
		b.Logf("Test failed: no entities to delete\n")
		b.Fail()
	}
	entities:=omap["entities"].([]interface{})
	var entity map[string]interface{}
 	b.ResetTimer()
    for i := 0; i < b.N; i++ {
    	if(len(entities)==0){
			b.Logf("Test failed: we ran out of entities\n")
			b.Fail()
			continue
    	}
    	entity, entities = entities[len(entities)-1].(map[string]interface{}), entities[:len(entities)-1]
		err:= client.Put("benchmark/"+entity["uuid"].(string), nil, map[string]string{"updated":"true"}, NOOPResponseHandler(&objmap))
		if(err!=nil){
			b.Logf("BenchmarkPut failed: %s\n", err)
			b.Fail()
		}
   }
	if(b.Failed()){
		str,_:=json.MarshalIndent(objmap,"","  ")
		b.Logf("RESPONSE: %s", str)
	}
}

func BenchmarkDelete(b *testing.B) {
	var objmap interface{}
 	client := Client {Organization:"yourorgname",Application:"sandbox",Uri:"https://api.usergrid.com"}
	params := map[string]string{"limit":strconv.Itoa(500)}
	err:= client.Get("benchmark", params, JSONResponseHandler(&objmap))
	if(err!=nil){
		b.Logf("Test failed: %s\n", err)
		b.Fail()
	}
	omap:=objmap.(map[string]interface{})
	if (omap["entities"]==nil || len(omap["entities"].([]interface{}))==0) {
		b.Logf("Test failed: no entities to delete\n")
		b.Fail()
	}
	entities:=omap["entities"].([]interface{})
	var entity map[string]interface{}
 	b.ResetTimer()
    for i := 0; i < b.N; i++ {
    	if(len(entities)==0){
			b.Logf("Test failed: we ran out of entities\n")
			b.Fail()
			continue
    	}
    	entity, entities = entities[len(entities)-1].(map[string]interface{}), entities[:len(entities)-1]
		err:= client.Delete("benchmark/"+entity["uuid"].(string), nil, NOOPResponseHandler(&objmap))
		if(err!=nil){
			b.Logf("BenchmarkDelete failed: %s\n", err)
			b.Fail()
		}
   }
	if(b.Failed()){
		str,_:=json.MarshalIndent(objmap,"","  ")
		b.Logf("RESPONSE: %s", str)
	}
}

<|endoftext|>


<commit_before>/*
Copyright (c) 2014 Ashley Jeffs

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, sub to the following conditions:

The above copyright notice and this permission notice shall be included in
all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
THE SOFTWARE.
*/

package util

import (
	"encoding/hex"
	"fmt"
	"time"

	"code.google.com/p/go-uuid/uuid"
)

/*--------------------------------------------------------------------------------------------------
 */

/*
GenerateStampedGUID - Generates a UUID and prepends a timestamp to it.
*/
func GenerateStampedUUID() string {
	tstamp := time.Now().Unix()

	return fmt.Sprintf("%v%v", tstamp, GenerateUUID())
}

/*
GenerateUUID - Generates a UUID and returns it as a hex encoded string.
*/
func GenerateUUID() string {
	return hex.EncodeToString(uuid.NewRandom())
}

/*--------------------------------------------------------------------------------------------------
 */
<commit_msg>Fixed a comment<commit_after>/*
Copyright (c) 2014 Ashley Jeffs

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, sub to the following conditions:

The above copyright notice and this permission notice shall be included in
all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
THE SOFTWARE.
*/

package util

import (
	"encoding/hex"
	"fmt"
	"time"

	"code.google.com/p/go-uuid/uuid"
)

/*--------------------------------------------------------------------------------------------------
 */

/*
GenerateStampedUUID - Generates a UUID and prepends a timestamp to it.
*/
func GenerateStampedUUID() string {
	tstamp := time.Now().Unix()

	return fmt.Sprintf("%v%v", tstamp, GenerateUUID())
}

/*
GenerateUUID - Generates a UUID and returns it as a hex encoded string.
*/
func GenerateUUID() string {
	return hex.EncodeToString(uuid.NewRandom())
}

/*--------------------------------------------------------------------------------------------------
 */
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"github.com/dynport/digo"
	"github.com/dynport/gocli"
	"os"
	"strconv"
	"time"
)

const RENAME_USAGE =  "<droplet_id> <new_name>"

func init() {
	cli.Register("droplet/rename",
		&gocli.Action{
			Handler:     RenameDropletAction,
			Description: "Describe Droplet",
			Usage: RENAME_USAGE,
		},
	)
}

func RenameDropletAction(args *gocli.Args) error {
	if len(args.Args) != 2 {
		fmt.Errorf(RENAME_USAGE)
	}
	id, newName := args.Args[0], args.Args[1]
	i, e := strconv.Atoi(id)
	if e != nil {
		return e
	}
	logger.Infof("renaming droplet %d to %s", i, newName)
	_, e = CurrentAccount().RenameDroplet(i, newName)
	if e != nil {
		return e
	}
	logger.Infof("renamed droplet %d to %s", i, newName)
	return nil
}

func init() {
	cli.Register("droplet/info",
		&gocli.Action{
			Handler:     DescribeDropletAction,
			Description: "Describe Droplet",
		},
	)
}

func DescribeDropletAction(args *gocli.Args) error {
	if len(args.Args) != 1 {
		return fmt.Errorf("USAGE: <droplet_id>")
	}
	i, e := strconv.Atoi(args.Args[0])
	if e != nil {
		return e
	}
	droplet, e := CurrentAccount().GetDroplet(i)
	if e != nil {
		return e
	}
	table := gocli.NewTable()
	table.Add("Id", fmt.Sprintf("%d", droplet.Id))
	table.Add("Name", droplet.Name)
	table.Add("Status", droplet.Status)
	table.Add("Locked", strconv.FormatBool(droplet.Locked))
	fmt.Println(table)
	return nil
}

func init() {
	cli.Register(
		"droplet/list",
		&gocli.Action{
			Handler:     ListDropletsAction,
			Description: "List active droplets",
		},
	)
}

func ListDropletsAction(args *gocli.Args) (e error) {
	logger.Debug("listing droplets")

	droplets, e := CurrentAccount().Droplets()
	if e != nil {
		return e
	}

	if _, e := CurrentAccount().CachedSizes(); e != nil {
		return e
	}

	table := gocli.NewTable()
	if len(droplets) == 0 {
		table.Add("no droplets found")
	} else {
		table.Add("Id", "Created", "Status", "Locked", "Name", "IPAddress", "Region", "Size", "Image")
		for _, droplet := range droplets {
			table.Add(
				strconv.Itoa(droplet.Id),
				droplet.CreatedAt.Format("2006-01-02T15:04"),
				droplet.Status,
				strconv.FormatBool(droplet.Locked),
				droplet.Name,
				droplet.IpAddress,
				fmt.Sprintf("%s (%d)", CurrentAccount().RegionName(droplet.RegionId), droplet.RegionId),
				fmt.Sprintf("%s (%d)", CurrentAccount().SizeName(droplet.SizeId), droplet.SizeId),
				fmt.Sprintf("%s (%d)", CurrentAccount().ImageName(droplet.ImageId), droplet.ImageId),
			)
		}
	}
	fmt.Fprintln(os.Stdout, table.String())
	return nil
}

func init() {
	args := &gocli.Args{}
	args.RegisterInt("-i", false, CurrentAccount().ImageId, "Image id for new droplet")
	args.RegisterInt("-r", false, CurrentAccount().RegionId, "Region id for new droplet")
	args.RegisterInt("-s", false, CurrentAccount().SizeId, "Size id for new droplet")
	args.RegisterInt("-k", false, CurrentAccount().SshKey, "Ssh key to be used")

	cli.Register(
		"droplet/create",
		&gocli.Action{
			Description: "Create new droplet",
			Usage:       "<name>",
			Handler:     CreateDropletAction,
			Args:        args,
		},
	)
}

func CreateDropletAction(a *gocli.Args) error {
	started := time.Now()
	logger.Debugf("would create a new droplet with %#v", a.Args)
	if len(a.Args) != 1 {
		return fmt.Errorf("USAGE: create droplet <name>")
	}
	droplet := &digo.Droplet{Name: a.Args[0]}

	var e error
	if droplet.SizeId, e = a.GetInt("-s"); e != nil {
		return e
	}

	if droplet.ImageId, e = a.GetInt("-i"); e != nil {
		return e
	}

	if droplet.RegionId, e = a.GetInt("-r"); e != nil {
		return e
	}

	if droplet.SshKey, e = a.GetInt("-k"); e != nil {
		return e
	}

	droplet, e = CurrentAccount().CreateDroplet(droplet)
	if e != nil {
		return e
	}
	droplet.Account = CurrentAccount()
	logger.Infof("created droplet with id %d", droplet.Id)
	e = digo.WaitForDroplet(droplet)
	logger.Infof("droplet %d ready, ip: %s. total_time: %.1fs", droplet.Id, droplet.IpAddress, time.Now().Sub(started).Seconds())
	return e
}

func init() {
	cli.Register(
		"droplet/destroy",
		&gocli.Action{
			Description: "Destroy droplet",
			Handler:     DestroyDropletAction,
			Usage:       "<droplet_id>",
		},
	)
}

func DestroyDropletAction(args *gocli.Args) error {
	logger.Debugf("would destroy droplet with %#v", args)
	if len(args.Args) == 0 {
		return fmt.Errorf("USAGE: droplet destroy id1,id2,id3")
	}
	for _, id := range args.Args {
		if i, e := strconv.Atoi(id); e == nil {
			logger.Prefix = fmt.Sprintf("droplet-%d", i)
			droplet, e := CurrentAccount().GetDroplet(i)
			if e != nil {
				logger.Errorf("unable to get droplet for %d", i)
				continue
			}
			logger.Infof("destroying droplet %d", droplet.Id)
			rsp, e := CurrentAccount().DestroyDroplet(droplet.Id)
			if e != nil {
				return e
			}
			logger.Debugf("got response %+v", rsp)
			started := time.Now()
			archived := false
			for i := 0; i < 300; i++ {
				droplet.Reload()
				if droplet.Status == "archive" || droplet.Status == "off" {
					archived = true
					break
				}
				logger.Debug("status " + droplet.Status)
				fmt.Print(".")
				time.Sleep(1 * time.Second)
			}
			fmt.Print("\n")
			logger.Info("droplet destroyed")
			if !archived {
				logger.Errorf("error archiving %d", droplet.Id)
			} else {
				logger.Debugf("archived in %.06f", time.Now().Sub(started).Seconds())
			}
		}
	}
	return nil
}

func init() {
	args := &gocli.Args{}
	args.RegisterInt("-i", false, 0, "Rebuild droplet")
	cli.Register(
		"droplet/rebuild",
		&gocli.Action{
			Description: "Rebuild droplet",
			Handler:     RebuildDropletAction,
			Usage:       "<droplet_id>",
			Args:        args,
		},
	)
}

func RebuildDropletAction(a *gocli.Args) error {
	if len(a.Args) != 1 {
		return fmt.Errorf("USAGE: droplet rebuild <id>")
	}
	i, e := strconv.Atoi(a.Args[0])
	if e != nil {
		return fmt.Errorf("USAGE: droplet rebuild <id>")
	}

	imageId, e := a.GetInt("-i")
	if e != nil {
		return e
	}

	rsp, e := account.RebuildDroplet(i, imageId)
	if e != nil {
		return e
	}
	logger.Debugf("got response %+v", rsp)
	droplet := &digo.Droplet{Id: i, Account: account}
	return digo.WaitForDroplet(droplet)
}
<commit_msg>fix router calls<commit_after>package main

import (
	"fmt"
	"github.com/dynport/digo"
	"github.com/dynport/gocli"
	"os"
	"strconv"
	"time"
)

const RENAME_USAGE =  "<droplet_id> <new_name>"

func init() {
	cli.Register("droplet/rename",
		&gocli.Action{
			Handler:     RenameDropletAction,
			Description: "Describe Droplet",
			Usage: RENAME_USAGE,
		},
	)
}

func RenameDropletAction(args *gocli.Args) error {
	if len(args.Args) != 2 {
		fmt.Errorf(RENAME_USAGE)
	}
	id, newName := args.Args[0], args.Args[1]
	i, e := strconv.Atoi(id)
	if e != nil {
		return e
	}
	logger.Infof("renaming droplet %d to %s", i, newName)
	_, e = CurrentAccount().RenameDroplet(i, newName)
	if e != nil {
		return e
	}
	logger.Infof("renamed droplet %d to %s", i, newName)
	return nil
}

func init() {
	cli.Register("droplet/info",
		&gocli.Action{
			Handler:     DescribeDropletAction,
			Description: "Describe Droplet",
		},
	)
}

func DescribeDropletAction(args *gocli.Args) error {
	if len(args.Args) != 1 {
		return fmt.Errorf("USAGE: <droplet_id>")
	}
	i, e := strconv.Atoi(args.Args[0])
	if e != nil {
		return e
	}
	droplet, e := CurrentAccount().GetDroplet(i)
	if e != nil {
		return e
	}
	table := gocli.NewTable()
	table.Add("Id", fmt.Sprintf("%d", droplet.Id))
	table.Add("Name", droplet.Name)
	table.Add("Status", droplet.Status)
	table.Add("Locked", strconv.FormatBool(droplet.Locked))
	fmt.Println(table)
	return nil
}

func init() {
	cli.Register(
		"droplet/list",
		&gocli.Action{
			Handler:     ListDropletsAction,
			Description: "List active droplets",
		},
	)
}

func ListDropletsAction(args *gocli.Args) (e error) {
	logger.Debug("listing droplets")

	droplets, e := CurrentAccount().Droplets()
	if e != nil {
		return e
	}

	if _, e := CurrentAccount().CachedSizes(); e != nil {
		return e
	}

	table := gocli.NewTable()
	if len(droplets) == 0 {
		table.Add("no droplets found")
	} else {
		table.Add("Id", "Created", "Status", "Locked", "Name", "IPAddress", "Region", "Size", "Image")
		for _, droplet := range droplets {
			table.Add(
				strconv.Itoa(droplet.Id),
				droplet.CreatedAt.Format("2006-01-02T15:04"),
				droplet.Status,
				strconv.FormatBool(droplet.Locked),
				droplet.Name,
				droplet.IpAddress,
				fmt.Sprintf("%s (%d)", CurrentAccount().RegionName(droplet.RegionId), droplet.RegionId),
				fmt.Sprintf("%s (%d)", CurrentAccount().SizeName(droplet.SizeId), droplet.SizeId),
				fmt.Sprintf("%s (%d)", CurrentAccount().ImageName(droplet.ImageId), droplet.ImageId),
			)
		}
	}
	fmt.Fprintln(os.Stdout, table.String())
	return nil
}

func init() {
	args := &gocli.Args{}
	args.RegisterInt("-i", "image_id", false, CurrentAccount().ImageId, "Image id for new droplet")
	args.RegisterInt("-r", "region_id", false, CurrentAccount().RegionId, "Region id for new droplet")
	args.RegisterInt("-s", "size_id", false, CurrentAccount().SizeId, "Size id for new droplet")
	args.RegisterInt("-k", "ssh_key_id", false, CurrentAccount().SshKey, "Ssh key to be used")

	cli.Register(
		"droplet/create",
		&gocli.Action{
			Description: "Create new droplet",
			Usage:       "<name>",
			Handler:     CreateDropletAction,
			Args:        args,
		},
	)
}

func CreateDropletAction(a *gocli.Args) error {
	started := time.Now()
	logger.Debugf("would create a new droplet with %#v", a.Args)
	if len(a.Args) != 1 {
		return fmt.Errorf("USAGE: create droplet <name>")
	}
	droplet := &digo.Droplet{Name: a.Args[0]}

	var e error
	if droplet.SizeId, e = a.GetInt("-s"); e != nil {
		return e
	}

	if droplet.ImageId, e = a.GetInt("-i"); e != nil {
		return e
	}

	if droplet.RegionId, e = a.GetInt("-r"); e != nil {
		return e
	}

	if droplet.SshKey, e = a.GetInt("-k"); e != nil {
		return e
	}

	droplet, e = CurrentAccount().CreateDroplet(droplet)
	if e != nil {
		return e
	}
	droplet.Account = CurrentAccount()
	logger.Infof("created droplet with id %d", droplet.Id)
	e = digo.WaitForDroplet(droplet)
	logger.Infof("droplet %d ready, ip: %s. total_time: %.1fs", droplet.Id, droplet.IpAddress, time.Now().Sub(started).Seconds())
	return e
}

func init() {
	cli.Register(
		"droplet/destroy",
		&gocli.Action{
			Description: "Destroy droplet",
			Handler:     DestroyDropletAction,
			Usage:       "<droplet_id>",
		},
	)
}

func DestroyDropletAction(args *gocli.Args) error {
	logger.Debugf("would destroy droplet with %#v", args)
	if len(args.Args) == 0 {
		return fmt.Errorf("USAGE: droplet destroy id1,id2,id3")
	}
	for _, id := range args.Args {
		if i, e := strconv.Atoi(id); e == nil {
			logger.Prefix = fmt.Sprintf("droplet-%d", i)
			droplet, e := CurrentAccount().GetDroplet(i)
			if e != nil {
				logger.Errorf("unable to get droplet for %d", i)
				continue
			}
			logger.Infof("destroying droplet %d", droplet.Id)
			rsp, e := CurrentAccount().DestroyDroplet(droplet.Id)
			if e != nil {
				return e
			}
			logger.Debugf("got response %+v", rsp)
			started := time.Now()
			archived := false
			for i := 0; i < 300; i++ {
				droplet.Reload()
				if droplet.Status == "archive" || droplet.Status == "off" {
					archived = true
					break
				}
				logger.Debug("status " + droplet.Status)
				fmt.Print(".")
				time.Sleep(1 * time.Second)
			}
			fmt.Print("\n")
			logger.Info("droplet destroyed")
			if !archived {
				logger.Errorf("error archiving %d", droplet.Id)
			} else {
				logger.Debugf("archived in %.06f", time.Now().Sub(started).Seconds())
			}
		}
	}
	return nil
}

func init() {
	args := &gocli.Args{}
	args.RegisterInt("-i", "image_id", false, 0, "Rebuild droplet")
	cli.Register(
		"droplet/rebuild",
		&gocli.Action{
			Description: "Rebuild droplet",
			Handler:     RebuildDropletAction,
			Usage:       "<droplet_id>",
			Args:        args,
		},
	)
}

func RebuildDropletAction(a *gocli.Args) error {
	if len(a.Args) != 1 {
		return fmt.Errorf("USAGE: droplet rebuild <id>")
	}
	i, e := strconv.Atoi(a.Args[0])
	if e != nil {
		return fmt.Errorf("USAGE: droplet rebuild <id>")
	}

	imageId, e := a.GetInt("-i")
	if e != nil {
		return e
	}

	rsp, e := account.RebuildDroplet(i, imageId)
	if e != nil {
		return e
	}
	logger.Debugf("got response %+v", rsp)
	droplet := &digo.Droplet{Id: i, Account: account}
	return digo.WaitForDroplet(droplet)
}
<|endoftext|>


<commit_before>package digraph

import (
	"fmt"
	"io"
)

// GenerateDot is used to emit a GraphViz compatible definition
// for a directed graph. It can be used to dump a .dot file.
func WriteDot(w io.Writer, nodes []Node) error {
	w.Write([]byte("digraph {\n"))
	defer w.Write([]byte("}\n"))

	for _, n := range nodes {
		nodeLine := fmt.Sprintf("\t\"%s\";\n", n)

		w.Write([]byte(nodeLine))

		for _, edge := range n.Edges() {
			target := edge.Tail()
			line := fmt.Sprintf("\t\"%s\" -> \"%s\" [label=\"%s\"];\n",
				n, target, edge)
			w.Write([]byte(line))
		}
	}

	return nil
}
<commit_msg>digraph: fix docs<commit_after>package digraph

import (
	"fmt"
	"io"
)

// WriteDot is used to emit a GraphViz compatible definition
// for a directed graph. It can be used to dump a .dot file.
func WriteDot(w io.Writer, nodes []Node) error {
	w.Write([]byte("digraph {\n"))
	defer w.Write([]byte("}\n"))

	for _, n := range nodes {
		nodeLine := fmt.Sprintf("\t\"%s\";\n", n)

		w.Write([]byte(nodeLine))

		for _, edge := range n.Edges() {
			target := edge.Tail()
			line := fmt.Sprintf("\t\"%s\" -> \"%s\" [label=\"%s\"];\n",
				n, target, edge)
			w.Write([]byte(line))
		}
	}

	return nil
}
<|endoftext|>


<commit_before>package pgmgr

import (
	"bytes"
	"database/sql"
	"errors"
	"fmt"
	"github.com/lib/pq"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"time"
)

type Config struct {
	// connection
	Username string
	Password string
	Database string
	Host     string
	Port     int
	Url      string

	// filepaths
	DumpFile        string	`json:"dump-file"`
	MigrationFolder string	`json:"migration-folder"`

	// options
	SeedTables	[]string	`json:"seed-tables"`
}

type Migration struct {
	Filename string
	Version  int
}

const (
	DOWN = iota
	UP   = iota
)

// Creates the database specified by the configuration.
func Create(c *Config) error {
	return sh("createdb", []string{c.Database})
}

// Drops the database specified by the configuration.
func Drop(c *Config) error {
	return sh("dropdb", []string{c.Database})
}

// Dumps the schema and contents of the database to the dump file.
func Dump(c *Config) error {
	// dump schema first...
	schema, err := shRead("pg_dump", []string{"--schema-only", c.Database})
	if err != nil {
		return err
	}

	// then selected data...
	args := []string{c.Database, "--data-only"}
	if len(c.SeedTables) > 0 {
		for _, table := range(c.SeedTables) {
		  println("pulling data for", table)
			args = append(args, "-t", table)
		}
	}
	println(strings.Join(args, ""))

	seeds, err := shRead("pg_dump", args)
	if err != nil {
		return err
	}

	// and combine into one file.
	file, err := os.OpenFile(c.DumpFile, os.O_CREATE | os.O_TRUNC | os.O_RDWR, 0600)
	if err != nil {
		return err
	}

	file.Write(*schema)
	file.Write(*seeds)
	file.Close()

	return nil
}

// Loads the database from the dump file.
func Load(c *Config) error {
	return sh("psql", []string{"-d", c.Database, "-f", c.DumpFile})
}

// Applies un-applied migrations in the specified MigrationFolder.
func Migrate(c *Config) error {
	migrations, err := migrations(c, "up")
	if err != nil {
		return err
	}

	// ensure the version table is created
	_, err = getOrInitializeVersion(c)
	if err != nil {
		return err
	}

	appliedAny := false
	for _, m := range migrations {
		if applied, _ := migrationIsApplied(c, m.Version); !applied {
			fmt.Println("== Applying", m.Filename, "==")
			t0 := time.Now()

			if err = applyMigration(c, m, UP); err != nil { // halt the migration process and return the error.
				fmt.Println(err)
				fmt.Println("")
			  fmt.Println("ERROR! Aborting the migration process.")
				return err
			}

			fmt.Println("== Completed in", time.Now().Sub(t0).Nanoseconds() / 1e6, "ms ==")
			appliedAny = true
		}
	}

	if !appliedAny {
		fmt.Println("Nothing to do; all migrations already applied.")
	}

	return nil
}

// Un-applies the latest migration, if possible.
func Rollback(c *Config) error {
	migrations, err := migrations(c, "down")
	if err != nil {
		return err
	}

	v, _ := Version(c)
	var to_rollback *Migration
	for _, m := range migrations {
		if m.Version == v {
			to_rollback = &m
			break
		}
	}

	if to_rollback == nil {
		return nil
	}

	// rollback only the last migration
	fmt.Println("== Reverting", to_rollback.Filename, "==")
	t0 := time.Now()

	if err = applyMigration(c, *to_rollback, DOWN); err != nil {
		return err
	}

	fmt.Println("== Completed in", time.Now().Sub(t0).Nanoseconds() / 1e6, "ms ==")

	return nil
}

// Returns the highest version number stored in the database. This is not
// necessarily enough info to uniquely identify the version, since there may
// be backdated migrations which have not yet applied.
func Version(c *Config) (int, error) {
	db, err := openConnection(c)
	defer db.Close()

	if err != nil {
		return -1, err
	}

	// if the table doesn't exist, we're at version -1
	var hasTable bool
	err = db.QueryRow("SELECT true FROM pg_catalog.pg_tables WHERE tablename='schema_migrations'").Scan(&hasTable)
	if hasTable != true {
		return -1, nil
	}

	var version int
	err = db.QueryRow("SELECT MAX(version) FROM schema_migrations").Scan(&version)

	return version, err
}

// Creates the schema_migrations table on what should be a new database.
func Initialize(c *Config) error {
	db, err := openConnection(c)
	defer db.Close()
	if err != nil {
		return err
	}

	_, err = db.Exec("CREATE TABLE schema_migrations (version INTEGER NOT NULL)")
	if err != nil {
		return err
	}

	return nil
}

// Creates new, blank migration files.
func CreateMigration(c *Config, name string) error {
	version := generateVersion()
	up_filepath   := filepath.Join(c.MigrationFolder, fmt.Sprint(version, "_", name, ".up.sql"))
	down_filepath := filepath.Join(c.MigrationFolder, fmt.Sprint(version, "_", name, ".down.sql"))

	err := ioutil.WriteFile(up_filepath, []byte(`-- Migration goes here.`), 0644)
	if err != nil {
		return err
	}

	err = ioutil.WriteFile(down_filepath, []byte(`-- Rollback of migration goes here. If you don't want to write it, delete this file.`), 0644)
	if err != nil {
		return err
	}

	return nil
}

func generateVersion() int {
	// TODO: guarantee no conflicts by incrementing if there is a conflict
	v := int(time.Now().Unix())
	return v
}

// need access to the original query contents in order to print it out properly,
// unfortunately.
func formatPgErr(contents *[]byte, pgerr *pq.Error) string {
	pos, _ := strconv.Atoi(pgerr.Position)
	lineNo := bytes.Count((*contents)[:pos], []byte("\n")) + 1
	columnNo := pos - bytes.LastIndex((*contents)[:pos], []byte("\n")) - 1

	return fmt.Sprint("PGERROR: line ", lineNo, " pos ", columnNo, ": ", pgerr.Message, ". ", pgerr.Detail)
}

func applyMigration(c *Config, m Migration, direction int) error {
	db, err := openConnection(c)
	defer db.Close()
	if err != nil {
		return err
	}

	tx, err := db.Begin()
	if err != nil {
		return err
	}

	contents, err := ioutil.ReadFile(filepath.Join(c.MigrationFolder, m.Filename))
	if err != nil {
		return err
	}

	if _, err = tx.Exec(string(contents)); err != nil {
		tx.Rollback()
		return errors.New(formatPgErr(&contents, err.(*pq.Error)))
	}

	if direction == UP {
		if err = insertSchemaVersion(tx, m.Version); err != nil {
			tx.Rollback()
			return errors.New(formatPgErr(&contents, err.(*pq.Error)))
		}
	} else {
		if err = deleteSchemaVersion(tx, m.Version); err != nil {
			tx.Rollback()
			return errors.New(formatPgErr(&contents, err.(*pq.Error)))
		}
	}

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

	return nil
}

func insertSchemaVersion(tx *sql.Tx, version int) error {
	_, err := tx.Exec("INSERT INTO schema_migrations (version) VALUES ($1) RETURNING version", version)
	return err
}

func deleteSchemaVersion(tx *sql.Tx, version int) error {
	_, err := tx.Exec("DELETE FROM schema_migrations WHERE version = $1", version)
	return err
}

func migrationIsApplied(c *Config, version int) (bool, error) {
	db, err := openConnection(c)
	defer db.Close()
	if err != nil {
		return false, err
	}

	var is_applied bool
	err = db.QueryRow("SELECT EXISTS(SELECT 1 FROM schema_migrations WHERE version = $1);", version).Scan(&is_applied)
	if err != nil {
		return false, err
	}

	return is_applied, nil
}

func getOrInitializeVersion(c *Config) (int, error) {
	var v int
	if v, _ = Version(c); v < 0 {
		if err := Initialize(c); err != nil {
			return -1, err
		}
	}

	return v, nil
}

func openConnection(c *Config) (*sql.DB, error) {
	db, err := sql.Open("postgres", sqlConnectionString(c))
	return db, err
}

func sqlConnectionString(c *Config) string {
	return fmt.Sprint(
		" user='", c.Username, "'",
		" dbname='", c.Database, "'",
		" password='", c.Password, "'",
		" host='", c.Host, "'",
		" sslmode=", "disable")
}

func migrations(c *Config, direction string) ([]Migration, error) {
	files, err := ioutil.ReadDir(c.MigrationFolder)
	migrations := []Migration{}
	if err != nil {
		return []Migration{}, err
	}

	for _, file := range files {
		if match, _ := regexp.MatchString("[0-9]+_.+."+direction+".sql", file.Name()); match {
			re := regexp.MustCompile("^[0-9]+")
			version, _ := strconv.Atoi(re.FindString(file.Name()))
			migrations = append(migrations, Migration{Filename: file.Name(), Version: version})
		}
	}

	return migrations, nil
}

func sh(command string, args []string) error {
	c := exec.Command(command, args...)
	output, err := c.CombinedOutput()
	fmt.Println(string(output))
	if err != nil {
		return err
	}

	return nil
}

func shRead(command string, args []string) (*[]byte, error) {
	c := exec.Command(command, args...)
	output, err := c.CombinedOutput()
	return &output, err
}
<commit_msg>Extra output to indicate what the new migration files are.<commit_after>package pgmgr

import (
	"bytes"
	"database/sql"
	"errors"
	"fmt"
	"github.com/lib/pq"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"time"
)

type Config struct {
	// connection
	Username string
	Password string
	Database string
	Host     string
	Port     int
	Url      string

	// filepaths
	DumpFile        string	`json:"dump-file"`
	MigrationFolder string	`json:"migration-folder"`

	// options
	SeedTables	[]string	`json:"seed-tables"`
}

type Migration struct {
	Filename string
	Version  int
}

const (
	DOWN = iota
	UP   = iota
)

// Creates the database specified by the configuration.
func Create(c *Config) error {
	return sh("createdb", []string{c.Database})
}

// Drops the database specified by the configuration.
func Drop(c *Config) error {
	return sh("dropdb", []string{c.Database})
}

// Dumps the schema and contents of the database to the dump file.
func Dump(c *Config) error {
	// dump schema first...
	schema, err := shRead("pg_dump", []string{"--schema-only", c.Database})
	if err != nil {
		return err
	}

	// then selected data...
	args := []string{c.Database, "--data-only"}
	if len(c.SeedTables) > 0 {
		for _, table := range(c.SeedTables) {
		  println("pulling data for", table)
			args = append(args, "-t", table)
		}
	}
	println(strings.Join(args, ""))

	seeds, err := shRead("pg_dump", args)
	if err != nil {
		return err
	}

	// and combine into one file.
	file, err := os.OpenFile(c.DumpFile, os.O_CREATE | os.O_TRUNC | os.O_RDWR, 0600)
	if err != nil {
		return err
	}

	file.Write(*schema)
	file.Write(*seeds)
	file.Close()

	return nil
}

// Loads the database from the dump file.
func Load(c *Config) error {
	return sh("psql", []string{"-d", c.Database, "-f", c.DumpFile})
}

// Applies un-applied migrations in the specified MigrationFolder.
func Migrate(c *Config) error {
	migrations, err := migrations(c, "up")
	if err != nil {
		return err
	}

	// ensure the version table is created
	_, err = getOrInitializeVersion(c)
	if err != nil {
		return err
	}

	appliedAny := false
	for _, m := range migrations {
		if applied, _ := migrationIsApplied(c, m.Version); !applied {
			fmt.Println("== Applying", m.Filename, "==")
			t0 := time.Now()

			if err = applyMigration(c, m, UP); err != nil { // halt the migration process and return the error.
				fmt.Println(err)
				fmt.Println("")
			  fmt.Println("ERROR! Aborting the migration process.")
				return err
			}

			fmt.Println("== Completed in", time.Now().Sub(t0).Nanoseconds() / 1e6, "ms ==")
			appliedAny = true
		}
	}

	if !appliedAny {
		fmt.Println("Nothing to do; all migrations already applied.")
	}

	return nil
}

// Un-applies the latest migration, if possible.
func Rollback(c *Config) error {
	migrations, err := migrations(c, "down")
	if err != nil {
		return err
	}

	v, _ := Version(c)
	var to_rollback *Migration
	for _, m := range migrations {
		if m.Version == v {
			to_rollback = &m
			break
		}
	}

	if to_rollback == nil {
		return nil
	}

	// rollback only the last migration
	fmt.Println("== Reverting", to_rollback.Filename, "==")
	t0 := time.Now()

	if err = applyMigration(c, *to_rollback, DOWN); err != nil {
		return err
	}

	fmt.Println("== Completed in", time.Now().Sub(t0).Nanoseconds() / 1e6, "ms ==")

	return nil
}

// Returns the highest version number stored in the database. This is not
// necessarily enough info to uniquely identify the version, since there may
// be backdated migrations which have not yet applied.
func Version(c *Config) (int, error) {
	db, err := openConnection(c)
	defer db.Close()

	if err != nil {
		return -1, err
	}

	// if the table doesn't exist, we're at version -1
	var hasTable bool
	err = db.QueryRow("SELECT true FROM pg_catalog.pg_tables WHERE tablename='schema_migrations'").Scan(&hasTable)
	if hasTable != true {
		return -1, nil
	}

	var version int
	err = db.QueryRow("SELECT MAX(version) FROM schema_migrations").Scan(&version)

	return version, err
}

// Creates the schema_migrations table on what should be a new database.
func Initialize(c *Config) error {
	db, err := openConnection(c)
	defer db.Close()
	if err != nil {
		return err
	}

	_, err = db.Exec("CREATE TABLE schema_migrations (version INTEGER NOT NULL)")
	if err != nil {
		return err
	}

	return nil
}

// Creates new, blank migration files.
func CreateMigration(c *Config, name string) error {
	version := generateVersion()
	up_filepath   := filepath.Join(c.MigrationFolder, fmt.Sprint(version, "_", name, ".up.sql"))
	down_filepath := filepath.Join(c.MigrationFolder, fmt.Sprint(version, "_", name, ".down.sql"))

	err := ioutil.WriteFile(up_filepath, []byte(`-- Migration goes here.`), 0644)
	if err != nil {
		return err
	}
	fmt.Println("Created", up_filepath)

	err = ioutil.WriteFile(down_filepath, []byte(`-- Rollback of migration goes here. If you don't want to write it, delete this file.`), 0644)
	if err != nil {
		return err
	}
	fmt.Println("Created", down_filepath)

	return nil
}

func generateVersion() int {
	// TODO: guarantee no conflicts by incrementing if there is a conflict
	v := int(time.Now().Unix())
	return v
}

// need access to the original query contents in order to print it out properly,
// unfortunately.
func formatPgErr(contents *[]byte, pgerr *pq.Error) string {
	pos, _ := strconv.Atoi(pgerr.Position)
	lineNo := bytes.Count((*contents)[:pos], []byte("\n")) + 1
	columnNo := pos - bytes.LastIndex((*contents)[:pos], []byte("\n")) - 1

	return fmt.Sprint("PGERROR: line ", lineNo, " pos ", columnNo, ": ", pgerr.Message, ". ", pgerr.Detail)
}

func applyMigration(c *Config, m Migration, direction int) error {
	db, err := openConnection(c)
	defer db.Close()
	if err != nil {
		return err
	}

	tx, err := db.Begin()
	if err != nil {
		return err
	}

	contents, err := ioutil.ReadFile(filepath.Join(c.MigrationFolder, m.Filename))
	if err != nil {
		return err
	}

	if _, err = tx.Exec(string(contents)); err != nil {
		tx.Rollback()
		return errors.New(formatPgErr(&contents, err.(*pq.Error)))
	}

	if direction == UP {
		if err = insertSchemaVersion(tx, m.Version); err != nil {
			tx.Rollback()
			return errors.New(formatPgErr(&contents, err.(*pq.Error)))
		}
	} else {
		if err = deleteSchemaVersion(tx, m.Version); err != nil {
			tx.Rollback()
			return errors.New(formatPgErr(&contents, err.(*pq.Error)))
		}
	}

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

	return nil
}

func insertSchemaVersion(tx *sql.Tx, version int) error {
	_, err := tx.Exec("INSERT INTO schema_migrations (version) VALUES ($1) RETURNING version", version)
	return err
}

func deleteSchemaVersion(tx *sql.Tx, version int) error {
	_, err := tx.Exec("DELETE FROM schema_migrations WHERE version = $1", version)
	return err
}

func migrationIsApplied(c *Config, version int) (bool, error) {
	db, err := openConnection(c)
	defer db.Close()
	if err != nil {
		return false, err
	}

	var is_applied bool
	err = db.QueryRow("SELECT EXISTS(SELECT 1 FROM schema_migrations WHERE version = $1);", version).Scan(&is_applied)
	if err != nil {
		return false, err
	}

	return is_applied, nil
}

func getOrInitializeVersion(c *Config) (int, error) {
	var v int
	if v, _ = Version(c); v < 0 {
		if err := Initialize(c); err != nil {
			return -1, err
		}
	}

	return v, nil
}

func openConnection(c *Config) (*sql.DB, error) {
	db, err := sql.Open("postgres", sqlConnectionString(c))
	return db, err
}

func sqlConnectionString(c *Config) string {
	return fmt.Sprint(
		" user='", c.Username, "'",
		" dbname='", c.Database, "'",
		" password='", c.Password, "'",
		" host='", c.Host, "'",
		" sslmode=", "disable")
}

func migrations(c *Config, direction string) ([]Migration, error) {
	files, err := ioutil.ReadDir(c.MigrationFolder)
	migrations := []Migration{}
	if err != nil {
		return []Migration{}, err
	}

	for _, file := range files {
		if match, _ := regexp.MatchString("[0-9]+_.+."+direction+".sql", file.Name()); match {
			re := regexp.MustCompile("^[0-9]+")
			version, _ := strconv.Atoi(re.FindString(file.Name()))
			migrations = append(migrations, Migration{Filename: file.Name(), Version: version})
		}
	}

	return migrations, nil
}

func sh(command string, args []string) error {
	c := exec.Command(command, args...)
	output, err := c.CombinedOutput()
	fmt.Println(string(output))
	if err != nil {
		return err
	}

	return nil
}

func shRead(command string, args []string) (*[]byte, error) {
	c := exec.Command(command, args...)
	output, err := c.CombinedOutput()
	return &output, err
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"os/signal"
	"strconv"
	"syscall"
	"time"

	"github.com/andelf/go-curl"
	"github.com/nu7hatch/gouuid"
)

type Config struct {
	Location        string
	ChecksUrl       string
	MeasurementsUrl string
}

type Check struct {
	Id  string `json:"id"`
	Url string `json:"url"`
}

type Measurement struct {
	Check             Check   `json:"check"`
	Id                string  `json:"id"`
	Location          string  `json:"location"`
	T                 int     `json:"t"`
	ExitStatus        int     `json:"exit_status"`
	ConnectTime       float64 `json:"connect_time,omitempty"`
	StartTransferTime float64 `json:"starttransfer_time,omitempty"`
	LocalIp           string  `json:"local_ip,omitempty"`
	PrimaryIp         string  `json:"primary_ip,omitempty"`
	TotalTime         float64 `json:"total_time,omitempty"`
	HttpStatus        int     `json:"http_status,omitempty"`
	NameLookupTime    float64 `json:"namelookup_time,omitempty"`
}

func GetEnvWithDefault(env string, def string) string {
	tmp := os.Getenv(env)

	if tmp == "" {
		return def
	}

	return tmp
}

func (c *Check) Measure(config Config) Measurement {
	var m Measurement

	id, _ := uuid.NewV4()
	m.Id = id.String()
	m.Check = *c
	m.Location = config.Location

	easy := curl.EasyInit()
	defer easy.Cleanup()

	easy.Setopt(curl.OPT_URL, c.Url)

	// dummy func for curl output
	noOut := func(buf []byte, userdata interface{}) bool {
		return true
	}

	easy.Setopt(curl.OPT_WRITEFUNCTION, noOut)
	easy.Setopt(curl.OPT_CONNECTTIMEOUT, 10)
	easy.Setopt(curl.OPT_TIMEOUT, 10)

	now := time.Now()
	m.T = int(now.Unix())

	if err := easy.Perform(); err != nil {
		if e, ok := err.(curl.CurlError); ok {
			m.ExitStatus = (int(e))
			return m
		}
		os.Exit(1)
	}

	m.ExitStatus = 0
	http_status, _ := easy.Getinfo(curl.INFO_RESPONSE_CODE)
	m.HttpStatus = http_status.(int)

	connect_time, _ := easy.Getinfo(curl.INFO_CONNECT_TIME)
	m.ConnectTime = connect_time.(float64)

	namelookup_time, _ := easy.Getinfo(curl.INFO_NAMELOOKUP_TIME)
	m.NameLookupTime = namelookup_time.(float64)

	starttransfer_time, _ := easy.Getinfo(curl.INFO_STARTTRANSFER_TIME)
	m.StartTransferTime = starttransfer_time.(float64)

	total_time, _ := easy.Getinfo(curl.INFO_TOTAL_TIME)
	m.TotalTime = total_time.(float64)

	local_ip, _ := easy.Getinfo(curl.INFO_LOCAL_IP)
	m.LocalIp = local_ip.(string)

	primary_ip, _ := easy.Getinfo(curl.INFO_PRIMARY_IP)
	m.PrimaryIp = primary_ip.(string)

	return m
}

func MeasureLoop(config Config, checks chan Check, measurements chan Measurement) {
	for {
		c := <-checks
		m := c.Measure(config)

		measurements <- m
	}
}

func RecordLoop(config Config, measurements chan Measurement) {
	payload := make([]Measurement, 0, 100)
	for {
		m := <-measurements
		payload = append(payload, m)

		s, err := json.Marshal(&payload)
		if err != nil {
			panic(err)
		}

		body := bytes.NewBuffer(s)
		req, err := http.NewRequest("POST", config.MeasurementsUrl, body)
		if err != nil {
			panic(err)
		}

		req.Header.Add("Content-Type", "application/json")
		resp, err := http.DefaultClient.Do(req)
		if err != nil {
			panic(err)
		}
		resp.Body.Close()
		payload = make([]Measurement, 0, 100)

		fmt.Println(resp)
	}
}

func GetChecks(config Config) []Check {
	url := config.ChecksUrl

	res, err := http.Get(url)
	if err != nil {
		panic(err)
	}
	defer res.Body.Close()

	body, err := ioutil.ReadAll(res.Body)
	if err != nil {
		panic(err)
	}

	var checks []Check
	err = json.Unmarshal(body, &checks)
	if err != nil {
		panic(err)
	}

	return checks
}

func ScheduleLoop(check Check, checks chan Check) {
	for {
		checks <- check
		time.Sleep(1000 * time.Millisecond)
	}
}

func main() {
	var config Config
	config.Location = GetEnvWithDefault("LOCATION", "undefined")
	config.ChecksUrl = GetEnvWithDefault("CHECKS_URL", "https://s3.amazonaws.com/canary-public-data/data.json")
	config.MeasurementsUrl = GetEnvWithDefault("MEASUREMENTS_URL", "http://localhost:5000/measurements")
	measurerCount, _ := strconv.Atoi(GetEnvWithDefault("MEASURER_COUNT", "1"))
	recorderCount, _ := strconv.Atoi(GetEnvWithDefault("RECORDER_COUNT", "1"))

	check_list := GetChecks(config)

	checks := make(chan Check)
	measurements := make(chan Measurement)

	for i := 0; i < measurerCount; i++ {
		go MeasureLoop(config, checks, measurements)
	}
	for i := 0; i < recorderCount; i++ {
		go RecordLoop(config, measurements)
	}

	for _, c := range check_list {
		go ScheduleLoop(c, checks)
	}
	sigs := make(chan os.Signal, 1)
	signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM)

	<-sigs
}
<commit_msg>checks for errors, at least<commit_after>package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"os/signal"
	"strconv"
	"syscall"
	"time"

	"github.com/andelf/go-curl"
	"github.com/nu7hatch/gouuid"
)

type Config struct {
	Location        string
	ChecksUrl       string
	MeasurementsUrl string
}

type Check struct {
	Id  string `json:"id"`
	Url string `json:"url"`
}

type Measurement struct {
	Check             Check   `json:"check"`
	Id                string  `json:"id"`
	Location          string  `json:"location"`
	T                 int     `json:"t"`
	ExitStatus        int     `json:"exit_status"`
	ConnectTime       float64 `json:"connect_time,omitempty"`
	StartTransferTime float64 `json:"starttransfer_time,omitempty"`
	LocalIp           string  `json:"local_ip,omitempty"`
	PrimaryIp         string  `json:"primary_ip,omitempty"`
	TotalTime         float64 `json:"total_time,omitempty"`
	HttpStatus        int     `json:"http_status,omitempty"`
	NameLookupTime    float64 `json:"namelookup_time,omitempty"`
}

func GetEnvWithDefault(env string, def string) string {
	tmp := os.Getenv(env)

	if tmp == "" {
		return def
	}

	return tmp
}

func (c *Check) Measure(config Config) Measurement {
	var m Measurement

	id, _ := uuid.NewV4()
	m.Id = id.String()
	m.Check = *c
	m.Location = config.Location

	easy := curl.EasyInit()
	defer easy.Cleanup()

	easy.Setopt(curl.OPT_URL, c.Url)

	// dummy func for curl output
	noOut := func(buf []byte, userdata interface{}) bool {
		return true
	}

	easy.Setopt(curl.OPT_WRITEFUNCTION, noOut)
	easy.Setopt(curl.OPT_CONNECTTIMEOUT, 10)
	easy.Setopt(curl.OPT_TIMEOUT, 10)

	now := time.Now()
	m.T = int(now.Unix())

	if err := easy.Perform(); err != nil {
		if e, ok := err.(curl.CurlError); ok {
			m.ExitStatus = (int(e))
			return m
		}
		os.Exit(1)
	}

	m.ExitStatus = 0
	http_status, _ := easy.Getinfo(curl.INFO_RESPONSE_CODE)
	m.HttpStatus = http_status.(int)

	connect_time, _ := easy.Getinfo(curl.INFO_CONNECT_TIME)
	m.ConnectTime = connect_time.(float64)

	namelookup_time, _ := easy.Getinfo(curl.INFO_NAMELOOKUP_TIME)
	m.NameLookupTime = namelookup_time.(float64)

	starttransfer_time, _ := easy.Getinfo(curl.INFO_STARTTRANSFER_TIME)
	m.StartTransferTime = starttransfer_time.(float64)

	total_time, _ := easy.Getinfo(curl.INFO_TOTAL_TIME)
	m.TotalTime = total_time.(float64)

	local_ip, _ := easy.Getinfo(curl.INFO_LOCAL_IP)
	m.LocalIp = local_ip.(string)

	primary_ip, _ := easy.Getinfo(curl.INFO_PRIMARY_IP)
	m.PrimaryIp = primary_ip.(string)

	return m
}

func MeasureLoop(config Config, checks chan Check, measurements chan Measurement) {
	for {
		c := <-checks
		m := c.Measure(config)

		measurements <- m
	}
}

func RecordLoop(config Config, measurements chan Measurement) {
	payload := make([]Measurement, 0, 100)
	for {
		m := <-measurements
		payload = append(payload, m)

		s, err := json.Marshal(&payload)
		if err != nil {
			panic(err)
		}

		body := bytes.NewBuffer(s)
		req, err := http.NewRequest("POST", config.MeasurementsUrl, body)
		if err != nil {
			panic(err)
		}

		req.Header.Add("Content-Type", "application/json")
		resp, err := http.DefaultClient.Do(req)
		if err != nil {
			panic(err)
		}
		resp.Body.Close()
		payload = make([]Measurement, 0, 100)

		fmt.Println(resp)
	}
}

func GetChecks(config Config) []Check {
	url := config.ChecksUrl

	res, err := http.Get(url)
	if err != nil {
		panic(err)
	}
	defer res.Body.Close()

	body, err := ioutil.ReadAll(res.Body)
	if err != nil {
		panic(err)
	}

	var checks []Check
	err = json.Unmarshal(body, &checks)
	if err != nil {
		panic(err)
	}

	return checks
}

func ScheduleLoop(check Check, checks chan Check) {
	for {
		checks <- check
		time.Sleep(1000 * time.Millisecond)
	}
}

func main() {
	var config Config
	config.Location = GetEnvWithDefault("LOCATION", "undefined")
	config.ChecksUrl = GetEnvWithDefault("CHECKS_URL", "https://s3.amazonaws.com/canary-public-data/data.json")
	config.MeasurementsUrl = GetEnvWithDefault("MEASUREMENTS_URL", "http://localhost:5000/measurements")

	measurerCount, err := strconv.Atoi(GetEnvWithDefault("MEASURER_COUNT", "1"))
	if err != nil {
		panic(err)
	}

	recorderCount, err := strconv.Atoi(GetEnvWithDefault("RECORDER_COUNT", "1"))
	if err != nil {
		panic(err)
	}

	check_list := GetChecks(config)

	checks := make(chan Check)
	measurements := make(chan Measurement)

	for i := 0; i < measurerCount; i++ {
		go MeasureLoop(config, checks, measurements)
	}
	for i := 0; i < recorderCount; i++ {
		go RecordLoop(config, measurements)
	}

	for _, c := range check_list {
		go ScheduleLoop(c, checks)
	}
	sigs := make(chan os.Signal, 1)
	signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM)

	<-sigs
}
<|endoftext|>


<commit_before>package ssh

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"os"

	"github.com/docker/docker/pkg/term"
	"github.com/docker/machine/log"
	"golang.org/x/crypto/ssh"
)

type Client struct {
	Config   *ssh.ClientConfig
	Hostname string
	Port     int
}

const (
	maxDialAttempts = 10
)

func NewClient(user string, host string, port int, auth *Auth) (*Client, error) {
	config, err := NewConfig(user, auth)
	if err != nil {
		return nil, err
	}

	return &Client{
		Config:   config,
		Hostname: host,
		Port:     port,
	}, nil
}

func NewConfig(user string, auth *Auth) (*ssh.ClientConfig, error) {
	var authMethods []ssh.AuthMethod

	for _, k := range auth.Keys {
		key, err := ioutil.ReadFile(k)
		if err != nil {
			return nil, err
		}

		privateKey, err := ssh.ParsePrivateKey(key)
		if err != nil {
			return nil, err
		}

		authMethods = append(authMethods, ssh.PublicKeys(privateKey))
	}

	for _, p := range auth.Passwords {
		authMethods = append(authMethods, ssh.Password(p))
	}

	return &ssh.ClientConfig{
		User: user,
		Auth: authMethods,
	}, nil
}

func (client *Client) Run(command string) (Output, error) {
	var (
		output Output
		conn   *ssh.Client
		err    error
	)

	for i := 0; ; i++ {
		conn, err = ssh.Dial("tcp", fmt.Sprintf("%s:%d", client.Hostname, client.Port), client.Config)
		if err != nil {
			log.Errorf("Error dialing TCP: %s", err)
			if i == maxDialAttempts {
				return output, errors.New("Max SSH/TCP dial attempts exceeded")
			}
		} else {
			break
		}
	}

	session, err := conn.NewSession()
	if err != nil {
		return output, fmt.Errorf("Error getting new session: %s", err)
	}

	defer session.Close()

	var stdout, stderr bytes.Buffer

	session.Stdout = &stdout
	session.Stderr = &stderr

	output = Output{
		Stdout: &stdout,
		Stderr: &stderr,
	}

	return output, session.Run(command)
}

func (client *Client) Shell() error {
	conn, err := ssh.Dial("tcp", fmt.Sprintf("%s:%d", client.Hostname, client.Port), client.Config)
	if err != nil {
		return err
	}

	session, err := conn.NewSession()
	if err != nil {
		return err
	}

	defer session.Close()

	session.Stdout = os.Stdout
	session.Stderr = os.Stderr
	session.Stdin = os.Stdin

	modes := ssh.TerminalModes{
		ssh.ECHO: 1,
	}

	var termWidth, termHeight int

	fd := os.Stdin.Fd()

	if term.IsTerminal(fd) {
		var oldState *term.State

		oldState, err = term.MakeRaw(fd)
		if err != nil {
			return err
		}

		defer term.RestoreTerminal(fd, oldState)

		winsize, err := term.GetWinsize(fd)
		if err != nil {
			termWidth = 80
			termHeight = 24
		} else {
			termWidth = int(winsize.Width)
			termHeight = int(winsize.Height)
		}
	}

	if err := session.RequestPty("xterm", termHeight, termWidth, modes); err != nil {
		return err
	}

	if err := session.Shell(); err != nil {
		return err
	}

	session.Wait()

	return nil
}

type Auth struct {
	Passwords []string
	Keys      []string
}

type Output struct {
	Stdout io.Reader
	Stderr io.Reader
}
<commit_msg>Move over to real retries method<commit_after>package ssh

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"os"

	"github.com/docker/docker/pkg/term"
	"github.com/docker/machine/log"
	"github.com/docker/machine/utils"
	"golang.org/x/crypto/ssh"
)

type Client struct {
	Config   *ssh.ClientConfig
	Hostname string
	Port     int
}

const (
	maxDialAttempts = 10
)

func NewClient(user string, host string, port int, auth *Auth) (*Client, error) {
	config, err := NewConfig(user, auth)
	if err != nil {
		return nil, err
	}

	return &Client{
		Config:   config,
		Hostname: host,
		Port:     port,
	}, nil
}

func NewConfig(user string, auth *Auth) (*ssh.ClientConfig, error) {
	var authMethods []ssh.AuthMethod

	for _, k := range auth.Keys {
		key, err := ioutil.ReadFile(k)
		if err != nil {
			return nil, err
		}

		privateKey, err := ssh.ParsePrivateKey(key)
		if err != nil {
			return nil, err
		}

		authMethods = append(authMethods, ssh.PublicKeys(privateKey))
	}

	for _, p := range auth.Passwords {
		authMethods = append(authMethods, ssh.Password(p))
	}

	return &ssh.ClientConfig{
		User: user,
		Auth: authMethods,
	}, nil
}

func dialSuccess(client *Client) func() bool {
	return func() bool {
		if _, err := ssh.Dial("tcp", fmt.Sprintf("%s:%d", client.Hostname, client.Port), client.Config); err != nil {
			log.Debugf("Error dialing TCP: %s", err)
			return false
		}
		return true
	}
}

func (client *Client) Run(command string) (Output, error) {
	var (
		output         Output
		stdout, stderr bytes.Buffer
	)

	if err := utils.WaitFor(dialSuccess(client)); err != nil {
		return output, fmt.Errorf("Error attempting SSH client dial: %s", err)
	}

	conn, err := ssh.Dial("tcp", fmt.Sprintf("%s:%d", client.Hostname, client.Port), client.Config)
	if err != nil {
		return output, fmt.Errorf("Mysterious error dialing TCP for SSH (we already succeeded at least once) : %s", err)
	}

	session, err := conn.NewSession()
	if err != nil {
		return output, fmt.Errorf("Error getting new session: %s", err)
	}

	defer session.Close()

	session.Stdout = &stdout
	session.Stderr = &stderr

	output = Output{
		Stdout: &stdout,
		Stderr: &stderr,
	}

	return output, session.Run(command)
}

func (client *Client) Shell() error {
	conn, err := ssh.Dial("tcp", fmt.Sprintf("%s:%d", client.Hostname, client.Port), client.Config)
	if err != nil {
		return err
	}

	session, err := conn.NewSession()
	if err != nil {
		return err
	}

	defer session.Close()

	session.Stdout = os.Stdout
	session.Stderr = os.Stderr
	session.Stdin = os.Stdin

	modes := ssh.TerminalModes{
		ssh.ECHO: 1,
	}

	var termWidth, termHeight int

	fd := os.Stdin.Fd()

	if term.IsTerminal(fd) {
		var oldState *term.State

		oldState, err = term.MakeRaw(fd)
		if err != nil {
			return err
		}

		defer term.RestoreTerminal(fd, oldState)

		winsize, err := term.GetWinsize(fd)
		if err != nil {
			termWidth = 80
			termHeight = 24
		} else {
			termWidth = int(winsize.Width)
			termHeight = int(winsize.Height)
		}
	}

	if err := session.RequestPty("xterm", termHeight, termWidth, modes); err != nil {
		return err
	}

	if err := session.Shell(); err != nil {
		return err
	}

	session.Wait()

	return nil
}

type Auth struct {
	Passwords []string
	Keys      []string
}

type Output struct {
	Stdout io.Reader
	Stderr io.Reader
}
<|endoftext|>


<commit_before>package models

type PrivateMessage struct {
	owner   uint
	message string
}
<commit_msg>Export the properties<commit_after>package models

type PrivateMessage struct {
	Owner   uint
	Message string
}
<|endoftext|>


<commit_before>package service

import "fmt"

//Builder is a simple implementation of ContainerBuilder
type Builder struct {
	definitions []Definition
}

//Insert a new definition into the Builder
func (b *Builder) Insert(def Definition) {
	b.definitions = append(b.definitions, def)
}

//Build builds the container once all definitions have been place in it
//dependencies need to be inserted in order for now
func (b *Builder) Build() (Container, error) {
	numDefs := len(b.definitions)
	servs := make(map[string]interface{}, numDefs)
	for _, def := range b.definitions {
		numDeps := len(def.Dependencies)
		deps := make(map[string]interface{}, numDeps)
		for _, name := range def.Dependencies {
			dep, ok := servs[name]
			if !ok {
				return nil, fmt.Errorf("service: Could not find dependency %q for service %q. Please make sure to insert them in order", name, def.Name)
			}
			deps[name] = dep
		}
		service, err := def.Initializer.Init(deps, def.Configuration)
		if err != nil {
			return nil, err
		}
		servs[def.Name] = service
	}

	return SimpleContainer{
		services: servs,
	}, nil
}
<commit_msg>builder: Split long error message across 3 lines<commit_after>package service

import "fmt"

//Builder is a simple implementation of ContainerBuilder
type Builder struct {
	definitions []Definition
}

//Insert a new definition into the Builder
func (b *Builder) Insert(def Definition) {
	b.definitions = append(b.definitions, def)
}

//Build builds the container once all definitions have been place in it
//dependencies need to be inserted in order for now
func (b *Builder) Build() (Container, error) {
	numDefs := len(b.definitions)
	servs := make(map[string]interface{}, numDefs)
	for _, def := range b.definitions {
		numDeps := len(def.Dependencies)
		deps := make(map[string]interface{}, numDeps)
		for _, name := range def.Dependencies {
			dep, ok := servs[name]
			if !ok {
				return nil, fmt.Errorf("service: Could not find"+
					"dependency %q for service %q. Please make"+
					"sure to insert them in order", name, def.Name)
			}
			deps[name] = dep
		}
		service, err := def.Initializer.Init(deps, def.Configuration)
		if err != nil {
			return nil, err
		}
		servs[def.Name] = service
	}

	return SimpleContainer{
		services: servs,
	}, nil
}
<|endoftext|>