<commit_before>// Copyright 2014 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package instancepoller

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

	"github.com/juju/errors"
	jc "github.com/juju/testing/checkers"
	gc "gopkg.in/check.v1"

	"github.com/juju/juju/environs"
	"github.com/juju/juju/instance"
	"github.com/juju/juju/network"
	"github.com/juju/juju/status"
	"github.com/juju/juju/testing"
)

type aggregateSuite struct {
	testing.BaseSuite
}

var _ = gc.Suite(&aggregateSuite{})

type testInstance struct {
	instance.Instance
	id        instance.Id
	addresses []network.Address
	status    string
	err       error
}

var _ instance.Instance = (*testInstance)(nil)

func (t *testInstance) Id() instance.Id {
	return t.id
}

func (t *testInstance) Addresses() ([]network.Address, error) {
	if t.err != nil {
		return nil, t.err
	}
	return t.addresses, nil
}

func (t *testInstance) Status() instance.InstanceStatus {
	return instance.InstanceStatus{Status: status.StatusUnknown, Message: t.status}
}

type testInstanceGetter struct {
	sync.RWMutex
	// ids is set when the Instances method is called.
	ids     []instance.Id
	results map[instance.Id]instance.Instance
	err     error
	counter int32
}

func (tig *testInstanceGetter) Instances(ids []instance.Id) (result []instance.Instance, err error) {
	tig.Lock()
	tig.ids = ids
	tig.Unlock()
	atomic.AddInt32(&tig.counter, 1)
	results := make([]instance.Instance, len(ids))
	for i, id := range ids {
		// We don't check 'ok' here, because we want the Instance{nil}
		// response for those
		results[i] = tig.results[id]
	}
	return results, tig.err
}

func (tig *testInstanceGetter) newTestInstance(id instance.Id, status string, addresses []string) *testInstance {
	if tig.results == nil {
		tig.results = make(map[instance.Id]instance.Instance)
	}
	thisInstance := &testInstance{
		id:        id,
		status:    status,
		addresses: network.NewAddresses(addresses...),
	}
	tig.results[thisInstance.Id()] = thisInstance
	return thisInstance
}

func (s *aggregateSuite) TestSingleRequest(c *gc.C) {
	testGetter := new(testInstanceGetter)
	instance1 := testGetter.newTestInstance("foo", "foobar", []string{"127.0.0.1", "192.168.1.1"})
	aggregator := newAggregator(testGetter)

	info, err := aggregator.instanceInfo("foo")
	c.Assert(err, jc.ErrorIsNil)
	c.Assert(info, gc.DeepEquals, instanceInfo{
		status:    instance.InstanceStatus{Status: status.StatusUnknown, Message: "foobar"},
		addresses: instance1.addresses,
	})
	testGetter.RLock()
	ids := testGetter.ids
	testGetter.RUnlock()
	c.Assert(ids, gc.DeepEquals, []instance.Id{"foo"})
}

func (s *aggregateSuite) TestMultipleResponseHandling(c *gc.C) {
	s.PatchValue(&gatherTime, 30*time.Millisecond)
	testGetter := new(testInstanceGetter)

	testGetter.newTestInstance("foo", "foobar", []string{"127.0.0.1", "192.168.1.1"})
	aggregator := newAggregator(testGetter)

	replyChan := make(chan instanceInfoReply)
	req := instanceInfoReq{
		reply:  replyChan,
		instId: instance.Id("foo"),
	}
	aggregator.reqc <- req
	reply := <-replyChan
	c.Assert(reply.err, gc.IsNil)

	testGetter.newTestInstance("foo2", "not foobar", []string{"192.168.1.2"})
	testGetter.newTestInstance("foo3", "ok-ish", []string{"192.168.1.3"})

	var wg sync.WaitGroup
	checkInfo := func(id instance.Id, expectStatus string) {
		info, err := aggregator.instanceInfo(id)
		c.Check(err, jc.ErrorIsNil)
		c.Check(info.status.Message, gc.Equals, expectStatus)
		wg.Done()
	}

	wg.Add(2)
	go checkInfo("foo2", "not foobar")
	go checkInfo("foo3", "ok-ish")
	wg.Wait()

	c.Assert(len(testGetter.ids), gc.DeepEquals, 2)
}

// notifyingInstanceGetter wraps testInstanceGetter, notifying via
// a channel when Instances() is called.
type notifyingInstanceGetter struct {
	testInstanceGetter
	instancesc chan bool
}

func (g *notifyingInstanceGetter) Instances(ids []instance.Id) ([]instance.Instance, error) {
	g.instancesc <- true
	return g.testInstanceGetter.Instances(ids)
}

func (s *aggregateSuite) TestDyingWhileHandlingRequest(c *gc.C) {
	// This tests a regression where the aggregator couldn't shut down
	// if the the tomb was killed while a request was being handled,
	// leaving the reply channel unread.

	s.PatchValue(&gatherTime, 30*time.Millisecond)

	// Set up the aggregator with the instance getter.
	testGetter := &notifyingInstanceGetter{instancesc: make(chan bool)}
	testGetter.newTestInstance("foo", "foobar", []string{"127.0.0.1", "192.168.1.1"})
	aggregator := newAggregator(testGetter)

	// Make a request with a reply channel that will never be read.
	req := instanceInfoReq{
		reply:  make(chan instanceInfoReply),
		instId: instance.Id("foo"),
	}
	aggregator.reqc <- req

	// Wait for Instances to be called.
	select {
	case <-testGetter.instancesc:
	case <-time.After(testing.LongWait):
		c.Fatal("Instances() not called")
	}

	// Now we know the request is being handled - kill the aggregator.
	aggregator.Kill()
	done := make(chan error)
	go func() {
		done <- aggregator.Wait()
	}()

	// The aggregator should stop.
	select {
	case err := <-done:
		c.Assert(err, jc.ErrorIsNil)
	case <-time.After(testing.LongWait):
		c.Fatal("aggregator didn't stop")
	}
}

type batchingInstanceGetter struct {
	sync.RWMutex
	testInstanceGetter
	wg         sync.WaitGroup
	aggregator *aggregator
	totalCount int
	batchSize  int
	started    int32
}

func (g *batchingInstanceGetter) Instances(ids []instance.Id) ([]instance.Instance, error) {
	insts, err := g.testInstanceGetter.Instances(ids)
	g.startRequests()
	return insts, err
}

func (g *batchingInstanceGetter) startRequests() {
	g.RLock()
	n := g.totalCount - int(g.started)
	if n > g.batchSize {
		n = g.batchSize
	}
	g.RUnlock()
	for i := 0; i < n; i++ {
		g.startRequest()
	}
}

func (g *batchingInstanceGetter) startRequest() {
	atomic.AddInt32(&g.started, 1)
	go func() {
		g.RLock()
		defer g.RUnlock()
		_, err := g.aggregator.instanceInfo("foo")
		if err != nil {
			panic(err)
		}
		g.wg.Done()
	}()
}

func (s *aggregateSuite) TestBatching(c *gc.C) {
	s.PatchValue(&gatherTime, 10*time.Millisecond)
	var testGetter batchingInstanceGetter
	testGetter.Lock()
	testGetter.aggregator = newAggregator(&testGetter)
	// We only need to inform the system about 1 instance, because all the
	// requests are for the same instance.
	testGetter.newTestInstance("foo", "foobar", []string{"127.0.0.1", "192.168.1.1"})
	testGetter.totalCount = 100
	testGetter.batchSize = 10
	testGetter.wg.Add(testGetter.totalCount)
	// startRequest will trigger one request, which ends up calling
	// Instances, which will turn around and trigger batchSize requests,
	// which should get aggregated into a single call to Instances, which
	// then should trigger another round of batchSize requests.
	testGetter.startRequest()
	testGetter.Unlock()
	testGetter.wg.Wait()
	c.Assert(testGetter.counter, gc.Equals, int32(testGetter.totalCount/testGetter.batchSize)+1)
}

func (s *aggregateSuite) TestError(c *gc.C) {
	testGetter := new(testInstanceGetter)
	ourError := fmt.Errorf("Some error")
	testGetter.err = ourError

	aggregator := newAggregator(testGetter)

	_, err := aggregator.instanceInfo("foo")
	c.Assert(err, gc.Equals, ourError)
}

func (s *aggregateSuite) TestPartialErrResponse(c *gc.C) {
	testGetter := new(testInstanceGetter)
	testGetter.err = environs.ErrPartialInstances

	aggregator := newAggregator(testGetter)
	_, err := aggregator.instanceInfo("foo")

	c.Assert(err, gc.ErrorMatches, "instance foo not found")
	c.Assert(err, jc.Satisfies, errors.IsNotFound)
}

func (s *aggregateSuite) TestAddressesError(c *gc.C) {
	testGetter := new(testInstanceGetter)
	instance1 := testGetter.newTestInstance("foo", "foobar", []string{"127.0.0.1", "192.168.1.1"})
	ourError := fmt.Errorf("gotcha")
	instance1.err = ourError

	aggregator := newAggregator(testGetter)
	_, err := aggregator.instanceInfo("foo")
	c.Assert(err, gc.Equals, ourError)
}

func (s *aggregateSuite) TestKillAndWait(c *gc.C) {
	testGetter := new(testInstanceGetter)
	aggregator := newAggregator(testGetter)
	aggregator.Kill()
	err := aggregator.Wait()
	c.Assert(err, jc.ErrorIsNil)
}
<commit_msg>Make counter fields consistent, add comment<commit_after>// Copyright 2014 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package instancepoller

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

	"github.com/juju/errors"
	jc "github.com/juju/testing/checkers"
	gc "gopkg.in/check.v1"

	"github.com/juju/juju/environs"
	"github.com/juju/juju/instance"
	"github.com/juju/juju/network"
	"github.com/juju/juju/status"
	"github.com/juju/juju/testing"
)

type aggregateSuite struct {
	testing.BaseSuite
}

var _ = gc.Suite(&aggregateSuite{})

type testInstance struct {
	instance.Instance
	id        instance.Id
	addresses []network.Address
	status    string
	err       error
}

var _ instance.Instance = (*testInstance)(nil)

func (t *testInstance) Id() instance.Id {
	return t.id
}

func (t *testInstance) Addresses() ([]network.Address, error) {
	if t.err != nil {
		return nil, t.err
	}
	return t.addresses, nil
}

func (t *testInstance) Status() instance.InstanceStatus {
	return instance.InstanceStatus{Status: status.StatusUnknown, Message: t.status}
}

type testInstanceGetter struct {
	sync.RWMutex
	// ids is set when the Instances method is called.
	ids     []instance.Id
	results map[instance.Id]instance.Instance
	err     error
	counter int32
}

func (tig *testInstanceGetter) Instances(ids []instance.Id) (result []instance.Instance, err error) {
	tig.Lock()
	tig.ids = ids
	tig.Unlock()
	atomic.AddInt32(&tig.counter, 1)
	results := make([]instance.Instance, len(ids))
	for i, id := range ids {
		// We don't check 'ok' here, because we want the Instance{nil}
		// response for those
		results[i] = tig.results[id]
	}
	return results, tig.err
}

func (tig *testInstanceGetter) newTestInstance(id instance.Id, status string, addresses []string) *testInstance {
	if tig.results == nil {
		tig.results = make(map[instance.Id]instance.Instance)
	}
	thisInstance := &testInstance{
		id:        id,
		status:    status,
		addresses: network.NewAddresses(addresses...),
	}
	tig.results[thisInstance.Id()] = thisInstance
	return thisInstance
}

func (s *aggregateSuite) TestSingleRequest(c *gc.C) {
	testGetter := new(testInstanceGetter)
	instance1 := testGetter.newTestInstance("foo", "foobar", []string{"127.0.0.1", "192.168.1.1"})
	aggregator := newAggregator(testGetter)

	info, err := aggregator.instanceInfo("foo")
	c.Assert(err, jc.ErrorIsNil)
	c.Assert(info, gc.DeepEquals, instanceInfo{
		status:    instance.InstanceStatus{Status: status.StatusUnknown, Message: "foobar"},
		addresses: instance1.addresses,
	})
	testGetter.RLock()
	ids := testGetter.ids
	testGetter.RUnlock()
	c.Assert(ids, gc.DeepEquals, []instance.Id{"foo"})
}

func (s *aggregateSuite) TestMultipleResponseHandling(c *gc.C) {
	s.PatchValue(&gatherTime, 30*time.Millisecond)
	testGetter := new(testInstanceGetter)

	testGetter.newTestInstance("foo", "foobar", []string{"127.0.0.1", "192.168.1.1"})
	aggregator := newAggregator(testGetter)

	replyChan := make(chan instanceInfoReply)
	req := instanceInfoReq{
		reply:  replyChan,
		instId: instance.Id("foo"),
	}
	aggregator.reqc <- req
	reply := <-replyChan
	c.Assert(reply.err, gc.IsNil)

	testGetter.newTestInstance("foo2", "not foobar", []string{"192.168.1.2"})
	testGetter.newTestInstance("foo3", "ok-ish", []string{"192.168.1.3"})

	var wg sync.WaitGroup
	checkInfo := func(id instance.Id, expectStatus string) {
		info, err := aggregator.instanceInfo(id)
		c.Check(err, jc.ErrorIsNil)
		c.Check(info.status.Message, gc.Equals, expectStatus)
		wg.Done()
	}

	wg.Add(2)
	go checkInfo("foo2", "not foobar")
	go checkInfo("foo3", "ok-ish")
	wg.Wait()

	c.Assert(len(testGetter.ids), gc.DeepEquals, 2)
}

// notifyingInstanceGetter wraps testInstanceGetter, notifying via
// a channel when Instances() is called.
type notifyingInstanceGetter struct {
	testInstanceGetter
	instancesc chan bool
}

func (g *notifyingInstanceGetter) Instances(ids []instance.Id) ([]instance.Instance, error) {
	g.instancesc <- true
	return g.testInstanceGetter.Instances(ids)
}

func (s *aggregateSuite) TestDyingWhileHandlingRequest(c *gc.C) {
	// This tests a regression where the aggregator couldn't shut down
	// if the the tomb was killed while a request was being handled,
	// leaving the reply channel unread.

	s.PatchValue(&gatherTime, 30*time.Millisecond)

	// Set up the aggregator with the instance getter.
	testGetter := &notifyingInstanceGetter{instancesc: make(chan bool)}
	testGetter.newTestInstance("foo", "foobar", []string{"127.0.0.1", "192.168.1.1"})
	aggregator := newAggregator(testGetter)

	// Make a request with a reply channel that will never be read.
	req := instanceInfoReq{
		reply:  make(chan instanceInfoReply),
		instId: instance.Id("foo"),
	}
	aggregator.reqc <- req

	// Wait for Instances to be called.
	select {
	case <-testGetter.instancesc:
	case <-time.After(testing.LongWait):
		c.Fatal("Instances() not called")
	}

	// Now we know the request is being handled - kill the aggregator.
	aggregator.Kill()
	done := make(chan error)
	go func() {
		done <- aggregator.Wait()
	}()

	// The aggregator should stop.
	select {
	case err := <-done:
		c.Assert(err, jc.ErrorIsNil)
	case <-time.After(testing.LongWait):
		c.Fatal("aggregator didn't stop")
	}
}

type batchingInstanceGetter struct {
	sync.RWMutex
	testInstanceGetter
	wg         sync.WaitGroup
	aggregator *aggregator
	totalCount int32
	batchSize  int32
	started    int32
}

func (g *batchingInstanceGetter) Instances(ids []instance.Id) ([]instance.Instance, error) {
	insts, err := g.testInstanceGetter.Instances(ids)
	g.startRequests()
	return insts, err
}

func (g *batchingInstanceGetter) startRequests() {
	g.RLock()
	n := g.totalCount - g.started
	if n > g.batchSize {
		n = g.batchSize
	}
	g.RUnlock()
	var i int32
	for i = 0; i < n; i++ {
		g.startRequest()
	}
}

func (g *batchingInstanceGetter) startRequest() {
	atomic.AddInt32(&g.started, 1)
	go func() {
		g.RLock()
		defer g.RUnlock()
		_, err := g.aggregator.instanceInfo("foo")
		if err != nil {
			panic(err)
		}
		g.wg.Done()
	}()
}

func (s *aggregateSuite) TestBatching(c *gc.C) {
	s.PatchValue(&gatherTime, 10*time.Millisecond)
	var testGetter batchingInstanceGetter
	testGetter.Lock()
	testGetter.aggregator = newAggregator(&testGetter)
	// We only need to inform the system about 1 instance, because all the
	// requests are for the same instance.
	testGetter.newTestInstance("foo", "foobar", []string{"127.0.0.1", "192.168.1.1"})
	testGetter.totalCount = 100
	testGetter.batchSize = 10
	testGetter.wg.Add(int(testGetter.totalCount))
	// startRequest will trigger one request, which ends up calling
	// Instances, which will turn around and trigger batchSize requests,
	// which should get aggregated into a single call to Instances, which
	// then should trigger another round of batchSize requests.
	testGetter.startRequest()
	testGetter.Unlock()
	testGetter.wg.Wait()
	// This assertion fails occasionally of-by-one.
	//     c.Assert(testGetter.counter, gc.Equals, testGetter.totalCount/testGetter.batchSize+1)
	//     ... obtained int32 = 12
	//     ... expected int32 = 11
	// It smells like a race, but isn't reported as one when if fails with the race flag set.
	c.Assert(testGetter.counter, gc.Equals, testGetter.totalCount/testGetter.batchSize+1)
}

func (s *aggregateSuite) TestError(c *gc.C) {
	testGetter := new(testInstanceGetter)
	ourError := fmt.Errorf("Some error")
	testGetter.err = ourError

	aggregator := newAggregator(testGetter)

	_, err := aggregator.instanceInfo("foo")
	c.Assert(err, gc.Equals, ourError)
}

func (s *aggregateSuite) TestPartialErrResponse(c *gc.C) {
	testGetter := new(testInstanceGetter)
	testGetter.err = environs.ErrPartialInstances

	aggregator := newAggregator(testGetter)
	_, err := aggregator.instanceInfo("foo")

	c.Assert(err, gc.ErrorMatches, "instance foo not found")
	c.Assert(err, jc.Satisfies, errors.IsNotFound)
}

func (s *aggregateSuite) TestAddressesError(c *gc.C) {
	testGetter := new(testInstanceGetter)
	instance1 := testGetter.newTestInstance("foo", "foobar", []string{"127.0.0.1", "192.168.1.1"})
	ourError := fmt.Errorf("gotcha")
	instance1.err = ourError

	aggregator := newAggregator(testGetter)
	_, err := aggregator.instanceInfo("foo")
	c.Assert(err, gc.Equals, ourError)
}

func (s *aggregateSuite) TestKillAndWait(c *gc.C) {
	testGetter := new(testInstanceGetter)
	aggregator := newAggregator(testGetter)
	aggregator.Kill()
	err := aggregator.Wait()
	c.Assert(err, jc.ErrorIsNil)
}
<|endoftext|>


<commit_before>package database

import (
	"database/sql"
	"errors"
	_ "github.com/alexbrainman/odbc"
	_ "github.com/denisenkom/go-mssqldb"
	_ "github.com/go-sql-driver/mysql"
	_ "github.com/ziutek/mymysql/godrv"
	"runtime"
	"strings"
)

type Database struct {
	inst       *sql.DB
	connString string
	driver     string
}

func (db *Database) Open(driver string, connString string) error {
	db.driver = driver
	db.connString = connString
	runtime.SetFinalizer(db, func(f interface{}) {
		f.(*Database).Close()
	})
	return db.executeOpen()
}

func (db *Database) executeOpen() error {
	var err error
	db.inst, err = sql.Open(db.driver, db.connString)
	if err != nil {
		return err
	}
	return nil
}

func (db *Database) Close() error {
	if db.inst != nil {
		return db.inst.Close()
	}
	return nil
}

type Rows struct {
	inst    *sql.Rows
	isFirst bool
	isNil   bool
	Cols    []string
}

func (db *Database) prepare(queryStr string, retCount int) (*sql.Stmt, error) {
	stmt, err := db.inst.Prepare(queryStr)
	if err != nil {
		db.Close()
		db.executeOpen()
		if retCount > 0 {
			return db.prepare(queryStr, retCount-1)
		}
		return nil, err
	}
	return stmt, err
}

func (db *Database) Query(queryStr string) (*Rows, error) {
	stmt, err := db.prepare(queryStr, 1)
	if stmt != nil {
		defer stmt.Close()
	}
	if err != nil {
		db.Close()
		db.executeOpen()
		return db.TempQuery(queryStr)
	}

	rows := &Rows{nil, true, false, make([]string, 0, 100)}
	rows.inst, err = stmt.Query()

	if err != nil {
		db.Close()
		db.executeOpen()
		return db.TempQuery(queryStr)
	}

	rows.Cols, err = rows.inst.Columns()

	if !rows.inst.Next() {
		rows.Close()
	}

	QUERYSTR := strings.ToUpper(queryStr)
	if strings.HasPrefix(QUERYSTR, "INSERT") && strings.Contains(QUERYSTR, "OUTPUT") && strings.Contains(QUERYSTR, "INSERTED.") {
		if rows.IsNil() {
			return nil, errors.New("insert.fail")
		}
	}

	runtime.SetFinalizer(rows, func(f interface{}) {
		f.(*Rows).Close()
	})

	return rows, nil
}

func (db *Database) TempQuery(queryStr string) (*Rows, error) {
	stmt, err := db.prepare(queryStr, 1)
	if stmt != nil {
		defer stmt.Close()
	}
	if err != nil {
		println("P1 : ", err.Error())
		return nil, err
	}

	rows := &Rows{nil, true, false, make([]string, 0, 100)}
	rows.inst, err = stmt.Query()

	if err != nil {
		println("P2 : ", err.Error(), "\n", queryStr)
		return nil, err
	}

	rows.Cols, err = rows.inst.Columns()

	if !rows.inst.Next() {
		rows.Close()
	}

	QUERYSTR := strings.ToUpper(queryStr)
	if strings.HasPrefix(QUERYSTR, "INSERT") && strings.Contains(QUERYSTR, "OUTPUT") && strings.Contains(QUERYSTR, "INSERTED.") {
		if rows.IsNil() {
			return nil, errors.New("insert.fail")
		}
	}

	runtime.SetFinalizer(rows, func(f interface{}) {
		f.(*Rows).Close()
	})

	return rows, nil
}

func (rows *Rows) Next() bool {
	if !rows.isNil && rows.isFirst {
		rows.isFirst = false
		return true
	}
	if !rows.inst.Next() {
		rows.Close()
	}
	return !rows.isNil
}

func (rows *Rows) FetchArray() []interface{} {
	if rows.isNil {
		return nil
	}
	cols, err := rows.inst.Columns()
	if err != nil {
		return nil
	}

	rawResult := make([]*interface{}, len(cols))
	result := make([]interface{}, len(cols))

	dest := make([]interface{}, len(cols))
	for i, _ := range rawResult {
		dest[i] = &rawResult[i]
	}
	rows.inst.Scan(dest...)
	for i, raw := range rawResult {
		if raw != nil {
			result[i] = (*raw)
		} else {
			result[i] = nil
		}
	}
	return result
}

func (rows *Rows) FetchHash() map[string]interface{} {
	if rows.isNil {
		return nil
	}
	cols, err := rows.inst.Columns()
	if err != nil {
		return nil
	}

	rawResult := make([]*interface{}, len(cols))
	result := make(map[string]interface{}, len(cols))

	dest := make([]interface{}, len(cols))
	for i, _ := range rawResult {
		dest[i] = &rawResult[i]
	}
	rows.inst.Scan(dest...)
	for i, raw := range rawResult {
		if raw != nil {
			result[cols[i]] = (*raw)
			result[strings.ToUpper(cols[i])] = (*raw)
			result[strings.ToLower(cols[i])] = (*raw)
		} else {
			result[cols[i]] = nil
			result[strings.ToUpper(cols[i])] = nil
			result[strings.ToLower(cols[i])] = nil
		}
	}
	return result
}

func (rows *Rows) Close() error {
	if rows != nil && rows.inst != nil {
		rows.isNil = true
		return rows.inst.Close()
	}
	return nil
}

func (rows *Rows) IsNil() bool {
	return rows.isNil
}
<commit_msg>prevent result set error when 'create'<commit_after>package database

import (
	"database/sql"
	"errors"
	_ "github.com/alexbrainman/odbc"
	_ "github.com/denisenkom/go-mssqldb"
	_ "github.com/go-sql-driver/mysql"
	_ "github.com/ziutek/mymysql/godrv"
	"runtime"
	"strings"
)

type Database struct {
	inst       *sql.DB
	connString string
	driver     string
}

func (db *Database) Open(driver string, connString string) error {
	db.driver = driver
	db.connString = connString
	runtime.SetFinalizer(db, func(f interface{}) {
		f.(*Database).Close()
	})
	return db.executeOpen()
}

func (db *Database) executeOpen() error {
	var err error
	db.inst, err = sql.Open(db.driver, db.connString)
	if err != nil {
		return err
	}
	return nil
}

func (db *Database) Close() error {
	if db.inst != nil {
		return db.inst.Close()
	}
	return nil
}

type Rows struct {
	inst    *sql.Rows
	isFirst bool
	isNil   bool
	Cols    []string
}

func (db *Database) prepare(queryStr string, retCount int) (*sql.Stmt, error) {
	stmt, err := db.inst.Prepare(queryStr)
	if err != nil {
		db.Close()
		db.executeOpen()
		if retCount > 0 {
			return db.prepare(queryStr, retCount-1)
		}
		return nil, err
	}
	return stmt, err
}

func (db *Database) Query(queryStr string) (*Rows, error) {
	stmt, err := db.prepare(queryStr, 1)
	if stmt != nil {
		defer stmt.Close()
	}
	if err != nil {
		db.Close()
		db.executeOpen()
		return db.TempQuery(queryStr)
	}

	rows := &Rows{nil, true, false, make([]string, 0, 100)}
	rows.inst, err = stmt.Query()

	if err != nil {
		db.Close()
		db.executeOpen()
		return db.TempQuery(queryStr)
	}

	rows.Cols, err = rows.inst.Columns()

	if !rows.inst.Next() {
		rows.Close()
	}

	QUERYSTR := strings.ToUpper(queryStr)
	if strings.HasPrefix(QUERYSTR, "INSERT") && strings.Contains(QUERYSTR, "OUTPUT") && strings.Contains(QUERYSTR, "INSERTED.") {
		if rows.IsNil() {
			return nil, errors.New("insert.fail")
		}
	}

	runtime.SetFinalizer(rows, func(f interface{}) {
		f.(*Rows).Close()
	})

	return rows, nil
}

func (db *Database) TempQuery(queryStr string) (*Rows, error) {
	stmt, err := db.prepare(queryStr, 1)
	if stmt != nil {
		defer stmt.Close()
	}
	if err != nil {
		println("P1 : ", err.Error())
		return nil, err
	}

	rows := &Rows{nil, true, false, make([]string, 0, 100)}
	rows.inst, err = stmt.Query()

	if err != nil && err.Error() != "Stmt did not create a result set" {
		println("P2 : ", err.Error(), "\n", queryStr)
		return nil, err
	}

	rows.Cols, err = rows.inst.Columns()

	if !rows.inst.Next() {
		rows.Close()
	}

	QUERYSTR := strings.ToUpper(queryStr)
	if strings.HasPrefix(QUERYSTR, "INSERT") && strings.Contains(QUERYSTR, "OUTPUT") && strings.Contains(QUERYSTR, "INSERTED.") {
		if rows.IsNil() {
			return nil, errors.New("insert.fail")
		}
	}

	runtime.SetFinalizer(rows, func(f interface{}) {
		f.(*Rows).Close()
	})

	return rows, nil
}

func (rows *Rows) Next() bool {
	if !rows.isNil && rows.isFirst {
		rows.isFirst = false
		return true
	}
	if !rows.inst.Next() {
		rows.Close()
	}
	return !rows.isNil
}

func (rows *Rows) FetchArray() []interface{} {
	if rows.isNil {
		return nil
	}
	cols, err := rows.inst.Columns()
	if err != nil {
		return nil
	}

	rawResult := make([]*interface{}, len(cols))
	result := make([]interface{}, len(cols))

	dest := make([]interface{}, len(cols))
	for i, _ := range rawResult {
		dest[i] = &rawResult[i]
	}
	rows.inst.Scan(dest...)
	for i, raw := range rawResult {
		if raw != nil {
			result[i] = (*raw)
		} else {
			result[i] = nil
		}
	}
	return result
}

func (rows *Rows) FetchHash() map[string]interface{} {
	if rows.isNil {
		return nil
	}
	cols, err := rows.inst.Columns()
	if err != nil {
		return nil
	}

	rawResult := make([]*interface{}, len(cols))
	result := make(map[string]interface{}, len(cols))

	dest := make([]interface{}, len(cols))
	for i, _ := range rawResult {
		dest[i] = &rawResult[i]
	}
	rows.inst.Scan(dest...)
	for i, raw := range rawResult {
		if raw != nil {
			result[cols[i]] = (*raw)
			result[strings.ToUpper(cols[i])] = (*raw)
			result[strings.ToLower(cols[i])] = (*raw)
		} else {
			result[cols[i]] = nil
			result[strings.ToUpper(cols[i])] = nil
			result[strings.ToLower(cols[i])] = nil
		}
	}
	return result
}

func (rows *Rows) Close() error {
	if rows != nil && rows.inst != nil {
		rows.isNil = true
		return rows.inst.Close()
	}
	return nil
}

func (rows *Rows) IsNil() bool {
	return rows.isNil
}
<|endoftext|>


<commit_before>// Copyright 2014 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package instancepoller

import (
	"fmt"
	"time"

	gc "launchpad.net/gocheck"

	"launchpad.net/juju-core/environs"
	"launchpad.net/juju-core/errors"
	"launchpad.net/juju-core/instance"
	"launchpad.net/juju-core/testing/testbase"
)

type aggregateSuite struct {
	testbase.LoggingSuite
}

var _ = gc.Suite(&aggregateSuite{})

type testInstance struct {
	instance.Instance
	addresses []instance.Address
	status    string
	err	bool
}

var _ instance.Instance = (*testInstance)(nil)

func (t *testInstance) Addresses() ([]instance.Address, error) {
	if t.err {
		return nil, fmt.Errorf("gotcha")
	}
	return t.addresses, nil
}

func (t *testInstance) Status() string {
	return t.status
}

type testInstanceGetter struct {
	ids     []instance.Id
	results []*testInstance
	err     error
}

func (i *testInstanceGetter) Instances(ids []instance.Id) (result []instance.Instance, err error) {
	i.ids = ids
	err = i.err
	for _, inst := range i.results {
		if inst == nil {
			result = append(result, nil)
		} else {
			result = append(result, inst)
		}
	}
	return
}

func newTestInstance(status string, addresses []string) *testInstance {
	thisInstance := testInstance{status: status}
	for _, address := range addresses {
		thisInstance.addresses = append(thisInstance.addresses, instance.NewAddress(address))
	}
	return &thisInstance
}

func (s *aggregateSuite) TestSingleRequest(c *gc.C) {
	testGetter := new(testInstanceGetter)
	instance1 := newTestInstance("foobar", []string{"127.0.0.1", "192.168.1.1"})
	testGetter.results = []*testInstance{instance1}
	aggregator := newAggregator(testGetter)

	replyChan := make(chan instanceInfoReply)
	req := instanceInfoReq{
		reply:  replyChan,
		instId: instance.Id("foo"),
	}
	aggregator.reqc <- req
	reply := <-replyChan
	c.Assert(reply.err, gc.IsNil)
	c.Assert(reply.info, gc.DeepEquals, instanceInfo{
		status:    "foobar",
		addresses: instance1.addresses,
	})
	c.Assert(testGetter.ids, gc.DeepEquals, []instance.Id{instance.Id("foo")})
}

func (s *aggregateSuite) TestRequestBatching(c *gc.C) {
	s.PatchValue(&GatherTime, 30*time.Millisecond)
	testGetter := new(testInstanceGetter)

	instance1 := newTestInstance("foobar", []string{"127.0.0.1", "192.168.1.1"})
	testGetter.results = []*testInstance{instance1}
	aggregator := newAggregator(testGetter)

	replyChan := make(chan instanceInfoReply)
	req := instanceInfoReq{
		reply:  replyChan,
		instId: instance.Id("foo"),
	}
	aggregator.reqc <- req
	reply := <-replyChan
	c.Assert(reply.err, gc.IsNil)

	instance2 := newTestInstance("not foobar", []string{"192.168.1.2"})
	instance3 := newTestInstance("ok-ish", []string{"192.168.1.3"})

	replyChan2 := make(chan instanceInfoReply)
	replyChan3 := make(chan instanceInfoReply)

	aggregator.reqc <- instanceInfoReq{reply: replyChan2, instId: instance.Id("foo2")}
	aggregator.reqc <- instanceInfoReq{reply: replyChan3, instId: instance.Id("foo3")}

	testGetter.results = []*testInstance{instance2, instance3}
	reply2 := <-replyChan2
	reply3 := <-replyChan3
	c.Assert(reply2.err, gc.IsNil)
	c.Assert(reply3.err, gc.IsNil)
	c.Assert(reply2.info.status, gc.Equals, "not foobar")
	c.Assert(reply3.info.status, gc.Equals, "ok-ish")

	c.Assert(testGetter.ids, gc.DeepEquals, []instance.Id{instance.Id("foo2"), instance.Id("foo3")})
}

func (s *aggregateSuite) TestError(c *gc.C) {
	testGetter := new(testInstanceGetter)
	ourError := fmt.Errorf("Some error")
	testGetter.err = ourError

	aggregator := newAggregator(testGetter)

	replyChan := make(chan instanceInfoReply)
	req := instanceInfoReq{
		reply:  replyChan,
		instId: instance.Id("foo"),
	}
	aggregator.reqc <- req
	reply := <-replyChan
	c.Assert(reply.err, gc.Equals, ourError)
}

func (s *aggregateSuite) TestPartialErrResponse(c *gc.C) {
	testGetter := new(testInstanceGetter)
	testGetter.err = environs.ErrPartialInstances
	testGetter.results = []*testInstance{nil}

	aggregator := newAggregator(testGetter)

	replyChan := make(chan instanceInfoReply)
	req := instanceInfoReq{
		reply:  replyChan,
		instId: instance.Id("foo"),
	}
	aggregator.reqc <- req
	reply := <-replyChan
	c.Assert(reply.err, gc.DeepEquals, errors.NotFoundf("instance foo"))
}

func (s *aggregateSuite) TestAddressesError(c *gc.C) {
	testGetter := new(testInstanceGetter)
	instance1 := newTestInstance("foobar", []string{"127.0.0.1", "192.168.1.1"})
	instance1.err = true
	testGetter.results = []*testInstance{instance1}
	aggregator := newAggregator(testGetter)

	replyChan := make(chan instanceInfoReply)
	req := instanceInfoReq{
		reply:  replyChan,
		instId: instance.Id("foo"),
	}
	aggregator.reqc <- req
	reply := <-replyChan
	c.Assert(reply.err, gc.DeepEquals, fmt.Errorf("gotcha"))
}

func (s *aggregateSuite) TestKillAndWait(c *gc.C) {
	testGetter := new(testInstanceGetter)
	aggregator := newAggregator(testGetter)
	aggregator.Kill()
	err := aggregator.Wait()
	c.Assert(err, gc.IsNil)
}

func (s *aggregateSuite) TestLoopDying(c *gc.C) {
	testGetter := new(testInstanceGetter)
	aggregator := newAggregator(testGetter)
	close(aggregator.reqc)
	err := aggregator.Wait()
	c.Assert(err, gc.NotNil)
}
<commit_msg>go fmt<commit_after>// Copyright 2014 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package instancepoller

import (
	"fmt"
	"time"

	gc "launchpad.net/gocheck"

	"launchpad.net/juju-core/environs"
	"launchpad.net/juju-core/errors"
	"launchpad.net/juju-core/instance"
	"launchpad.net/juju-core/testing/testbase"
)

type aggregateSuite struct {
	testbase.LoggingSuite
}

var _ = gc.Suite(&aggregateSuite{})

type testInstance struct {
	instance.Instance
	addresses []instance.Address
	status    string
	err       bool
}

var _ instance.Instance = (*testInstance)(nil)

func (t *testInstance) Addresses() ([]instance.Address, error) {
	if t.err {
		return nil, fmt.Errorf("gotcha")
	}
	return t.addresses, nil
}

func (t *testInstance) Status() string {
	return t.status
}

type testInstanceGetter struct {
	ids     []instance.Id
	results []*testInstance
	err     error
}

func (i *testInstanceGetter) Instances(ids []instance.Id) (result []instance.Instance, err error) {
	i.ids = ids
	err = i.err
	for _, inst := range i.results {
		if inst == nil {
			result = append(result, nil)
		} else {
			result = append(result, inst)
		}
	}
	return
}

func newTestInstance(status string, addresses []string) *testInstance {
	thisInstance := testInstance{status: status}
	for _, address := range addresses {
		thisInstance.addresses = append(thisInstance.addresses, instance.NewAddress(address))
	}
	return &thisInstance
}

func (s *aggregateSuite) TestSingleRequest(c *gc.C) {
	testGetter := new(testInstanceGetter)
	instance1 := newTestInstance("foobar", []string{"127.0.0.1", "192.168.1.1"})
	testGetter.results = []*testInstance{instance1}
	aggregator := newAggregator(testGetter)

	replyChan := make(chan instanceInfoReply)
	req := instanceInfoReq{
		reply:  replyChan,
		instId: instance.Id("foo"),
	}
	aggregator.reqc <- req
	reply := <-replyChan
	c.Assert(reply.err, gc.IsNil)
	c.Assert(reply.info, gc.DeepEquals, instanceInfo{
		status:    "foobar",
		addresses: instance1.addresses,
	})
	c.Assert(testGetter.ids, gc.DeepEquals, []instance.Id{instance.Id("foo")})
}

func (s *aggregateSuite) TestRequestBatching(c *gc.C) {
	s.PatchValue(&GatherTime, 30*time.Millisecond)
	testGetter := new(testInstanceGetter)

	instance1 := newTestInstance("foobar", []string{"127.0.0.1", "192.168.1.1"})
	testGetter.results = []*testInstance{instance1}
	aggregator := newAggregator(testGetter)

	replyChan := make(chan instanceInfoReply)
	req := instanceInfoReq{
		reply:  replyChan,
		instId: instance.Id("foo"),
	}
	aggregator.reqc <- req
	reply := <-replyChan
	c.Assert(reply.err, gc.IsNil)

	instance2 := newTestInstance("not foobar", []string{"192.168.1.2"})
	instance3 := newTestInstance("ok-ish", []string{"192.168.1.3"})

	replyChan2 := make(chan instanceInfoReply)
	replyChan3 := make(chan instanceInfoReply)

	aggregator.reqc <- instanceInfoReq{reply: replyChan2, instId: instance.Id("foo2")}
	aggregator.reqc <- instanceInfoReq{reply: replyChan3, instId: instance.Id("foo3")}

	testGetter.results = []*testInstance{instance2, instance3}
	reply2 := <-replyChan2
	reply3 := <-replyChan3
	c.Assert(reply2.err, gc.IsNil)
	c.Assert(reply3.err, gc.IsNil)
	c.Assert(reply2.info.status, gc.Equals, "not foobar")
	c.Assert(reply3.info.status, gc.Equals, "ok-ish")

	c.Assert(testGetter.ids, gc.DeepEquals, []instance.Id{instance.Id("foo2"), instance.Id("foo3")})
}

func (s *aggregateSuite) TestError(c *gc.C) {
	testGetter := new(testInstanceGetter)
	ourError := fmt.Errorf("Some error")
	testGetter.err = ourError

	aggregator := newAggregator(testGetter)

	replyChan := make(chan instanceInfoReply)
	req := instanceInfoReq{
		reply:  replyChan,
		instId: instance.Id("foo"),
	}
	aggregator.reqc <- req
	reply := <-replyChan
	c.Assert(reply.err, gc.Equals, ourError)
}

func (s *aggregateSuite) TestPartialErrResponse(c *gc.C) {
	testGetter := new(testInstanceGetter)
	testGetter.err = environs.ErrPartialInstances
	testGetter.results = []*testInstance{nil}

	aggregator := newAggregator(testGetter)

	replyChan := make(chan instanceInfoReply)
	req := instanceInfoReq{
		reply:  replyChan,
		instId: instance.Id("foo"),
	}
	aggregator.reqc <- req
	reply := <-replyChan
	c.Assert(reply.err, gc.DeepEquals, errors.NotFoundf("instance foo"))
}

func (s *aggregateSuite) TestAddressesError(c *gc.C) {
	testGetter := new(testInstanceGetter)
	instance1 := newTestInstance("foobar", []string{"127.0.0.1", "192.168.1.1"})
	instance1.err = true
	testGetter.results = []*testInstance{instance1}
	aggregator := newAggregator(testGetter)

	replyChan := make(chan instanceInfoReply)
	req := instanceInfoReq{
		reply:  replyChan,
		instId: instance.Id("foo"),
	}
	aggregator.reqc <- req
	reply := <-replyChan
	c.Assert(reply.err, gc.DeepEquals, fmt.Errorf("gotcha"))
}

func (s *aggregateSuite) TestKillAndWait(c *gc.C) {
	testGetter := new(testInstanceGetter)
	aggregator := newAggregator(testGetter)
	aggregator.Kill()
	err := aggregator.Wait()
	c.Assert(err, gc.IsNil)
}

func (s *aggregateSuite) TestLoopDying(c *gc.C) {
	testGetter := new(testInstanceGetter)
	aggregator := newAggregator(testGetter)
	close(aggregator.reqc)
	err := aggregator.Wait()
	c.Assert(err, gc.NotNil)
}
<|endoftext|>


<commit_before>package database

import (
	"database/sql"
	"fmt"
	"strings"
	"time"

	"github.com/omniscale/imposm3/element"
	"github.com/omniscale/imposm3/parser/changeset"
	"github.com/omniscale/imposm3/parser/diff"

	_ "github.com/lib/pq"
)

var initSql = []string{
	`CREATE SCHEMA IF NOT EXISTS "%[1]s";`,
	`CREATE TABLE IF NOT EXISTS "%[1]s".nodes (
    id BIGINT,
    add BOOLEAN,
    modify BOOLEAN,
    delete BOOLEAN,
    changeset INT,
    geometry GEOMETRY(Point, 4326),
    user_name VARCHAR,
    user_id INT,
    timestamp TIMESTAMP WITH TIME ZONE,
    version INT,
    tags HSTORE,
    PRIMARY KEY (id, version)
);`,
	`CREATE TABLE IF NOT EXISTS "%[1]s".ways (
    id INT NOT NULL,
    add BOOLEAN,
    modify BOOLEAN,
    delete BOOLEAN,
    changeset INT,
    user_name VARCHAR,
    user_id INT,
    timestamp TIMESTAMP WITH TIME ZONE,
    version INT,
    tags HSTORE,
    PRIMARY KEY (id, version)
);`,
	`CREATE TABLE IF NOT EXISTS "%[1]s".nds (
    way_id INT NOT NULL,
    way_version INT NOT NULL,
    idx INT,
    node_id BIGINT NOT NULL,
    PRIMARY KEY (way_id, way_version, idx, node_id),
    FOREIGN KEY (way_id, way_version)
        REFERENCES "%[1]s".ways (id, version)
          ON UPDATE CASCADE
          ON DELETE CASCADE
);`,
	`CREATE TABLE IF NOT EXISTS "%[1]s".relations (
    id INT NOT NULL,
    add BOOLEAN,
    modify BOOLEAN,
    delete BOOLEAN,
    changeset INT,
    geometry GEOMETRY(Point, 4326),
    user_name VARCHAR,
    user_id INT,
    timestamp TIMESTAMP WITH TIME ZONE,
    version INT,
    tags HSTORE,
    PRIMARY KEY (id, version)
);`,
	`CREATE TABLE IF NOT EXISTS "%[1]s".members (
    relation_id INT NOT NULL,
    relation_version INT,
    type VARCHAR,
    role VARCHAR,
    idx INT,
    member_node_id BIGINT,
    member_way_id INT,
    member_relation_id INT,
    FOREIGN KEY (relation_id, relation_version)
        REFERENCES "%[1]s".relations (id, version)
          ON UPDATE CASCADE
          ON DELETE CASCADE
);`,
	`CREATE TABLE IF NOT EXISTS "%[1]s".changesets (
    id INT NOT NULL,
    created_at TIMESTAMP WITH TIME ZONE,
    closed_at TIMESTAMP WITH TIME ZONE,
    num_changes INT,
    open BOOLEAN,
    user_name VARCHAR,
    user_id INT,
    tags HSTORE,
    bbox Geometry(POLYGON, 4326),
    PRIMARY KEY (id)
);`,
	`CREATE TABLE IF NOT EXISTS "%[1]s".comments (
    changeset_id BIGINT NOT NULL,
    idx INT,
    user_name VARCHAR,
    user_id INT,
    timestamp TIMESTAMP WITH TIME ZONE,
    text VARCHAR,
    PRIMARY KEY (changeset_id, idx),
    FOREIGN KEY (changeset_id)
        REFERENCES "%[1]s".changesets (id)
          ON UPDATE CASCADE
          ON DELETE CASCADE
);`,
}

type PostGIS struct {
	db               *sql.DB
	tx               *sql.Tx
	nodeStmt         *sql.Stmt
	wayStmt          *sql.Stmt
	ndsStmt          *sql.Stmt
	relStmt          *sql.Stmt
	memberStmt       *sql.Stmt
	changeStmt       *sql.Stmt
	changeUpdateStmt *sql.Stmt
	commentStmt      *sql.Stmt
	schema           string
}

func NewPostGIS(params string, schema string) (*PostGIS, error) {
	db, err := sql.Open("postgres", params)
	if err != nil {
		return nil, err
	}

	if schema == "" {
		schema = "public"
	}

	return &PostGIS{
		db:     db,
		schema: schema,
	}, nil
}

func newSqlError(err error, elem interface{}) error {
	return &sqlError{elem: elem, err: err}
}

type sqlError struct {
	elem interface{}
	err  error
}

func (s *sqlError) Error() string {
	return fmt.Sprintf("error: %s; for %#v", s.err, s.elem)

}

func (p *PostGIS) Init() error {
	tx, err := p.db.Begin()
	if err != nil {
		return err
	}
	for _, s := range initSql {
		stmt := fmt.Sprintf(s, p.schema)
		if _, err := tx.Exec(stmt); err != nil {
			tx.Rollback()
			return fmt.Errorf("error while calling %v: %v", stmt, err)
		}
	}
	return tx.Commit()
}

func (p *PostGIS) Begin() error {
	var err error
	p.tx, err = p.db.Begin()
	if err != nil {
		return err
	}
	nodeStmt, err := p.tx.Prepare(
		fmt.Sprintf(`INSERT INTO "%[1]s".nodes (
            id,
            add,
            modify,
            delete,
            geometry,
            user_name,
            user_id,
            timestamp,
            version,
            changeset,
            tags) VALUES ($1, $2, $3, $4, ST_SetSRID(ST_Point($5, $6), 4326), $7, $8, $9, $10, $11, $12)`, p.schema),
	)
	if err != nil {
		return err
	}
	p.nodeStmt = nodeStmt

	wayStmt, err := p.tx.Prepare(
		fmt.Sprintf(`INSERT INTO "%[1]s".ways (
            id,
            add,
            modify,
            delete,
            user_name,
            user_id,
            timestamp,
            version,
            changeset,
            tags
        ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)`, p.schema),
	)
	if err != nil {
		return err
	}
	p.wayStmt = wayStmt

	ndsStmt, err := p.tx.Prepare(
		fmt.Sprintf(`INSERT INTO "%[1]s".nds (
            way_id,
            way_version,
            idx,
            node_id
        ) VALUES ($1, $2, $3, $4)`, p.schema),
	)
	if err != nil {
		return err
	}
	p.ndsStmt = ndsStmt

	relStmt, err := p.tx.Prepare(
		fmt.Sprintf(`INSERT INTO "%[1]s".relations (
            id,
            add,
            modify,
            delete,
            user_name,
            user_id,
            timestamp,
            version,
            changeset,
            tags
        ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)`, p.schema),
	)
	if err != nil {
		return err
	}
	p.relStmt = relStmt

	memberStmt, err := p.tx.Prepare(
		fmt.Sprintf(`INSERT INTO "%[1]s".members (
            relation_id,
            relation_version,
            type,
            role,
            idx,
            member_node_id,
            member_way_id,
            member_relation_id
        ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8)`, p.schema),
	)
	if err != nil {
		return err
	}
	p.memberStmt = memberStmt

	changeStmt, err := p.tx.Prepare(
		fmt.Sprintf(`INSERT INTO "%[1]s".changesets (
            id,
            created_at,
            closed_at,
            open,
            num_changes,
            user_name,
            user_id,
            bbox,
            tags
        ) VALUES ($1, $2, $3, $4, $5, $6, $7, ST_GeomFromText($8), $9)`, p.schema),
	)
	if err != nil {
		return err
	}
	p.changeStmt = changeStmt

	changeUpdateStmt, err := p.tx.Prepare(
		fmt.Sprintf(`UPDATE "%[1]s".changesets SET
            created_at=$2,
            closed_at=$3,
            open=$4,
            num_changes=$5,
            user_name=$6,
            user_id=$7,
            bbox=ST_GeomFromText($8),
            tags=$9
        WHERE id = $1`, p.schema),
	)
	if err != nil {
		return err
	}
	p.changeUpdateStmt = changeUpdateStmt

	commentStmt, err := p.tx.Prepare(
		fmt.Sprintf(`INSERT INTO "%[1]s".comments (
            changeset_id,
            idx,
            user_name,
            user_id,
            timestamp,
            text
        ) VALUES ($1, $2, $3, $4, $5, $6)`, p.schema),
	)
	if err != nil {
		return err
	}
	p.commentStmt = commentStmt

	return nil
}

func (p *PostGIS) Commit() error {
	return p.tx.Commit()
}

func (p *PostGIS) ImportElem(elem diff.Element) (err error) {
	_, err = p.tx.Exec("SAVEPOINT insert")
	if err != nil {
		return
	}

	defer func() {
		if err != nil {
			p.tx.Exec("ROLLBACK TO SAVEPOINT insert")
		} else {
			_, err = p.tx.Exec("RELEASE SAVEPOINT insert")
		}
	}()
	var add, mod, del bool
	if elem.Mod {
		mod = true
	} else if elem.Add {
		add = true
	} else if elem.Del {
		del = true
	}
	if elem.Node != nil {
		nd := elem.Node
		if _, err = p.nodeStmt.Exec(
			nd.Id,
			add,
			mod,
			del,
			nd.Long, nd.Lat,
			nd.Metadata.UserName,
			nd.Metadata.UserId,
			nd.Metadata.Timestamp.UTC(),
			nd.Metadata.Version,
			nd.Metadata.Changeset,
			hstoreString(nd.Tags),
		); err != nil {
			return newSqlError(err, elem.Node)
		}
	} else if elem.Way != nil {
		w := elem.Way
		if _, err = p.wayStmt.Exec(
			w.Id,
			add,
			mod,
			del,
			w.Metadata.UserName,
			w.Metadata.UserId,
			w.Metadata.Timestamp.UTC(),
			w.Metadata.Version,
			w.Metadata.Changeset,
			hstoreString(w.Tags),
		); err != nil {
			return newSqlError(err, elem.Way)
		}
		for i, ref := range elem.Way.Refs {
			if _, err = p.ndsStmt.Exec(
				w.Id,
				w.Metadata.Version,
				i,
				ref,
			); err != nil {
				return newSqlError(err, elem.Way)
			}
		}
	} else if elem.Rel != nil {
		rel := elem.Rel
		if _, err = p.relStmt.Exec(
			rel.Id,
			add,
			mod,
			del,
			rel.Metadata.UserName,
			rel.Metadata.UserId,
			rel.Metadata.Timestamp.UTC(),
			rel.Metadata.Version,
			rel.Metadata.Changeset,
			hstoreString(rel.Tags),
		); err != nil {
			return newSqlError(err, elem.Rel)
		}
		for i, m := range elem.Rel.Members {
			var nodeId, wayId, relId interface{}
			switch m.Type {
			case element.NODE:
				nodeId = m.Id
			case element.WAY:
				wayId = m.Id
			case element.RELATION:
				relId = m.Id
			}
			if _, err = p.memberStmt.Exec(
				rel.Id,
				rel.Metadata.Version,
				m.Type,
				m.Role,
				i,
				nodeId,
				wayId,
				relId,
			); err != nil {
				return newSqlError(err, elem.Rel)
			}
		}
	}

	return nil
}

func (p *PostGIS) ImportChangeset(c changeset.Changeset) error {
	bbox := bboxPolygon(c)
	if _, err := p.tx.Exec("SAVEPOINT insert_changeset"); err != nil {
		return err
	}
	// insert null if closedAt isZero
	var closedAt *time.Time
	if !c.ClosedAt.IsZero() {
		closedUtc := c.ClosedAt.UTC()
		closedAt = &closedUtc
	}
	if _, err := p.changeStmt.Exec(
		c.Id,
		c.CreatedAt.UTC(),
		closedAt,
		c.Open,
		c.NumChanges,
		c.User,
		c.UserId,
		bbox,
		hstoreStringChangeset(c.Tags),
	); err != nil {
		if _, err := p.tx.Exec("ROLLBACK TO SAVEPOINT insert_changeset"); err != nil {
			return err
		}
		if _, err := p.changeUpdateStmt.Exec(
			c.Id,
			c.CreatedAt.UTC(),
			closedAt,
			c.Open,
			c.NumChanges,
			c.User,
			c.UserId,
			bbox,
			hstoreStringChangeset(c.Tags),
		); err != nil {
			return newSqlError(err, c)
		}
		if _, err := p.tx.Exec(fmt.Sprintf(`DELETE FROM "%[1]s".comments WHERE changeset_id = $1`, p.schema), c.Id); err != nil {
			return err
		}
	}

	for i, com := range c.Comments {
		if _, err := p.commentStmt.Exec(
			c.Id,
			i,
			com.User,
			com.UserId,
			com.Date.UTC(),
			com.Text,
		); err != nil {
			return newSqlError(err, c)
		}
	}

	if _, err := p.tx.Exec("RELEASE SAVEPOINT insert_changeset"); err != nil {
		return err
	}

	return nil
}

var hstoreReplacer = strings.NewReplacer("\\", "\\\\", "\"", "\\\"")

func hstoreString(tags element.Tags) string {
	kv := make([]string, 0, len(tags))
	for k, v := range tags {
		kv = append(kv, `"`+hstoreReplacer.Replace(k)+`"=>"`+hstoreReplacer.Replace(v)+`"`)
	}
	return strings.Join(kv, ", ")
}

func hstoreStringChangeset(tags []changeset.Tag) string {
	kv := make([]string, 0, len(tags))
	for _, t := range tags {
		kv = append(kv, `"`+hstoreReplacer.Replace(t.Key)+`"=>"`+hstoreReplacer.Replace(t.Value)+`"`)
	}
	return strings.Join(kv, ", ")
}

func bboxPolygon(c changeset.Changeset) interface{} {
	if c.MinLon != 0.0 && c.MaxLon != 0.0 && c.MinLat != 0.0 && c.MaxLat != 0.0 {
		return fmt.Sprintf(
			"SRID=4326; POLYGON((%[1]f %[2]f, %[1]f %[4]f, %[3]f %[4]f, %[3]f %[2]f, %[1]f %[2]f))",
			c.MinLon, c.MinLat, c.MaxLon, c.MaxLat,
		)
	}
	return nil
}
<commit_msg>create indices for changeset ids<commit_after>package database

import (
	"database/sql"
	"fmt"
	"strings"
	"time"

	"github.com/omniscale/imposm3/element"
	"github.com/omniscale/imposm3/parser/changeset"
	"github.com/omniscale/imposm3/parser/diff"

	_ "github.com/lib/pq"
)

var initSql = []string{
	`CREATE SCHEMA IF NOT EXISTS "%[1]s";`,
	`CREATE TABLE IF NOT EXISTS "%[1]s".nodes (
    id BIGINT,
    add BOOLEAN,
    modify BOOLEAN,
    delete BOOLEAN,
    changeset INT,
    geometry GEOMETRY(Point, 4326),
    user_name VARCHAR,
    user_id INT,
    timestamp TIMESTAMP WITH TIME ZONE,
    version INT,
    tags HSTORE,
    PRIMARY KEY (id, version)
);`,
	`CREATE TABLE IF NOT EXISTS "%[1]s".ways (
    id INT NOT NULL,
    add BOOLEAN,
    modify BOOLEAN,
    delete BOOLEAN,
    changeset INT,
    user_name VARCHAR,
    user_id INT,
    timestamp TIMESTAMP WITH TIME ZONE,
    version INT,
    tags HSTORE,
    PRIMARY KEY (id, version)
);`,
	`CREATE INDEX IF NOT EXISTS ways_changset_idx ON "%[1]s".ways USING btree (changeset);`,
	`CREATE TABLE IF NOT EXISTS "%[1]s".nds (
    way_id INT NOT NULL,
    way_version INT NOT NULL,
    idx INT,
    node_id BIGINT NOT NULL,
    PRIMARY KEY (way_id, way_version, idx, node_id),
    FOREIGN KEY (way_id, way_version)
        REFERENCES "%[1]s".ways (id, version)
          ON UPDATE CASCADE
          ON DELETE CASCADE
);`,
	`CREATE TABLE IF NOT EXISTS "%[1]s".relations (
    id INT NOT NULL,
    add BOOLEAN,
    modify BOOLEAN,
    delete BOOLEAN,
    changeset INT,
    geometry GEOMETRY(Point, 4326),
    user_name VARCHAR,
    user_id INT,
    timestamp TIMESTAMP WITH TIME ZONE,
    version INT,
    tags HSTORE,
    PRIMARY KEY (id, version)
);`,
	`CREATE INDEX IF NOT EXISTS relations_changset_idx ON "%[1]s".relations USING btree (changeset);`,
	`CREATE TABLE IF NOT EXISTS "%[1]s".members (
    relation_id INT NOT NULL,
    relation_version INT,
    type VARCHAR,
    role VARCHAR,
    idx INT,
    member_node_id BIGINT,
    member_way_id INT,
    member_relation_id INT,
    FOREIGN KEY (relation_id, relation_version)
        REFERENCES "%[1]s".relations (id, version)
          ON UPDATE CASCADE
          ON DELETE CASCADE
);`,
	`CREATE TABLE IF NOT EXISTS "%[1]s".changesets (
    id INT NOT NULL,
    created_at TIMESTAMP WITH TIME ZONE,
    closed_at TIMESTAMP WITH TIME ZONE,
    num_changes INT,
    open BOOLEAN,
    user_name VARCHAR,
    user_id INT,
    tags HSTORE,
    bbox Geometry(POLYGON, 4326),
    PRIMARY KEY (id)
);`,
	`CREATE TABLE IF NOT EXISTS "%[1]s".comments (
    changeset_id BIGINT NOT NULL,
    idx INT,
    user_name VARCHAR,
    user_id INT,
    timestamp TIMESTAMP WITH TIME ZONE,
    text VARCHAR,
    PRIMARY KEY (changeset_id, idx),
    FOREIGN KEY (changeset_id)
        REFERENCES "%[1]s".changesets (id)
          ON UPDATE CASCADE
          ON DELETE CASCADE
);`,
}

type PostGIS struct {
	db               *sql.DB
	tx               *sql.Tx
	nodeStmt         *sql.Stmt
	wayStmt          *sql.Stmt
	ndsStmt          *sql.Stmt
	relStmt          *sql.Stmt
	memberStmt       *sql.Stmt
	changeStmt       *sql.Stmt
	changeUpdateStmt *sql.Stmt
	commentStmt      *sql.Stmt
	schema           string
}

func NewPostGIS(params string, schema string) (*PostGIS, error) {
	db, err := sql.Open("postgres", params)
	if err != nil {
		return nil, err
	}

	if schema == "" {
		schema = "public"
	}

	return &PostGIS{
		db:     db,
		schema: schema,
	}, nil
}

func newSqlError(err error, elem interface{}) error {
	return &sqlError{elem: elem, err: err}
}

type sqlError struct {
	elem interface{}
	err  error
}

func (s *sqlError) Error() string {
	return fmt.Sprintf("error: %s; for %#v", s.err, s.elem)

}

func (p *PostGIS) Init() error {
	tx, err := p.db.Begin()
	if err != nil {
		return err
	}
	for _, s := range initSql {
		stmt := fmt.Sprintf(s, p.schema)
		if _, err := tx.Exec(stmt); err != nil {
			tx.Rollback()
			return fmt.Errorf("error while calling %v: %v", stmt, err)
		}
	}
	return tx.Commit()
}

func (p *PostGIS) Begin() error {
	var err error
	p.tx, err = p.db.Begin()
	if err != nil {
		return err
	}
	nodeStmt, err := p.tx.Prepare(
		fmt.Sprintf(`INSERT INTO "%[1]s".nodes (
            id,
            add,
            modify,
            delete,
            geometry,
            user_name,
            user_id,
            timestamp,
            version,
            changeset,
            tags) VALUES ($1, $2, $3, $4, ST_SetSRID(ST_Point($5, $6), 4326), $7, $8, $9, $10, $11, $12)`, p.schema),
	)
	if err != nil {
		return err
	}
	p.nodeStmt = nodeStmt

	wayStmt, err := p.tx.Prepare(
		fmt.Sprintf(`INSERT INTO "%[1]s".ways (
            id,
            add,
            modify,
            delete,
            user_name,
            user_id,
            timestamp,
            version,
            changeset,
            tags
        ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)`, p.schema),
	)
	if err != nil {
		return err
	}
	p.wayStmt = wayStmt

	ndsStmt, err := p.tx.Prepare(
		fmt.Sprintf(`INSERT INTO "%[1]s".nds (
            way_id,
            way_version,
            idx,
            node_id
        ) VALUES ($1, $2, $3, $4)`, p.schema),
	)
	if err != nil {
		return err
	}
	p.ndsStmt = ndsStmt

	relStmt, err := p.tx.Prepare(
		fmt.Sprintf(`INSERT INTO "%[1]s".relations (
            id,
            add,
            modify,
            delete,
            user_name,
            user_id,
            timestamp,
            version,
            changeset,
            tags
        ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)`, p.schema),
	)
	if err != nil {
		return err
	}
	p.relStmt = relStmt

	memberStmt, err := p.tx.Prepare(
		fmt.Sprintf(`INSERT INTO "%[1]s".members (
            relation_id,
            relation_version,
            type,
            role,
            idx,
            member_node_id,
            member_way_id,
            member_relation_id
        ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8)`, p.schema),
	)
	if err != nil {
		return err
	}
	p.memberStmt = memberStmt

	changeStmt, err := p.tx.Prepare(
		fmt.Sprintf(`INSERT INTO "%[1]s".changesets (
            id,
            created_at,
            closed_at,
            open,
            num_changes,
            user_name,
            user_id,
            bbox,
            tags
        ) VALUES ($1, $2, $3, $4, $5, $6, $7, ST_GeomFromText($8), $9)`, p.schema),
	)
	if err != nil {
		return err
	}
	p.changeStmt = changeStmt

	changeUpdateStmt, err := p.tx.Prepare(
		fmt.Sprintf(`UPDATE "%[1]s".changesets SET
            created_at=$2,
            closed_at=$3,
            open=$4,
            num_changes=$5,
            user_name=$6,
            user_id=$7,
            bbox=ST_GeomFromText($8),
            tags=$9
        WHERE id = $1`, p.schema),
	)
	if err != nil {
		return err
	}
	p.changeUpdateStmt = changeUpdateStmt

	commentStmt, err := p.tx.Prepare(
		fmt.Sprintf(`INSERT INTO "%[1]s".comments (
            changeset_id,
            idx,
            user_name,
            user_id,
            timestamp,
            text
        ) VALUES ($1, $2, $3, $4, $5, $6)`, p.schema),
	)
	if err != nil {
		return err
	}
	p.commentStmt = commentStmt

	return nil
}

func (p *PostGIS) Commit() error {
	return p.tx.Commit()
}

func (p *PostGIS) ImportElem(elem diff.Element) (err error) {
	_, err = p.tx.Exec("SAVEPOINT insert")
	if err != nil {
		return
	}

	defer func() {
		if err != nil {
			p.tx.Exec("ROLLBACK TO SAVEPOINT insert")
		} else {
			_, err = p.tx.Exec("RELEASE SAVEPOINT insert")
		}
	}()
	var add, mod, del bool
	if elem.Mod {
		mod = true
	} else if elem.Add {
		add = true
	} else if elem.Del {
		del = true
	}
	if elem.Node != nil {
		nd := elem.Node
		if _, err = p.nodeStmt.Exec(
			nd.Id,
			add,
			mod,
			del,
			nd.Long, nd.Lat,
			nd.Metadata.UserName,
			nd.Metadata.UserId,
			nd.Metadata.Timestamp.UTC(),
			nd.Metadata.Version,
			nd.Metadata.Changeset,
			hstoreString(nd.Tags),
		); err != nil {
			return newSqlError(err, elem.Node)
		}
	} else if elem.Way != nil {
		w := elem.Way
		if _, err = p.wayStmt.Exec(
			w.Id,
			add,
			mod,
			del,
			w.Metadata.UserName,
			w.Metadata.UserId,
			w.Metadata.Timestamp.UTC(),
			w.Metadata.Version,
			w.Metadata.Changeset,
			hstoreString(w.Tags),
		); err != nil {
			return newSqlError(err, elem.Way)
		}
		for i, ref := range elem.Way.Refs {
			if _, err = p.ndsStmt.Exec(
				w.Id,
				w.Metadata.Version,
				i,
				ref,
			); err != nil {
				return newSqlError(err, elem.Way)
			}
		}
	} else if elem.Rel != nil {
		rel := elem.Rel
		if _, err = p.relStmt.Exec(
			rel.Id,
			add,
			mod,
			del,
			rel.Metadata.UserName,
			rel.Metadata.UserId,
			rel.Metadata.Timestamp.UTC(),
			rel.Metadata.Version,
			rel.Metadata.Changeset,
			hstoreString(rel.Tags),
		); err != nil {
			return newSqlError(err, elem.Rel)
		}
		for i, m := range elem.Rel.Members {
			var nodeId, wayId, relId interface{}
			switch m.Type {
			case element.NODE:
				nodeId = m.Id
			case element.WAY:
				wayId = m.Id
			case element.RELATION:
				relId = m.Id
			}
			if _, err = p.memberStmt.Exec(
				rel.Id,
				rel.Metadata.Version,
				m.Type,
				m.Role,
				i,
				nodeId,
				wayId,
				relId,
			); err != nil {
				return newSqlError(err, elem.Rel)
			}
		}
	}

	return nil
}

func (p *PostGIS) ImportChangeset(c changeset.Changeset) error {
	bbox := bboxPolygon(c)
	if _, err := p.tx.Exec("SAVEPOINT insert_changeset"); err != nil {
		return err
	}
	// insert null if closedAt isZero
	var closedAt *time.Time
	if !c.ClosedAt.IsZero() {
		closedUtc := c.ClosedAt.UTC()
		closedAt = &closedUtc
	}
	if _, err := p.changeStmt.Exec(
		c.Id,
		c.CreatedAt.UTC(),
		closedAt,
		c.Open,
		c.NumChanges,
		c.User,
		c.UserId,
		bbox,
		hstoreStringChangeset(c.Tags),
	); err != nil {
		if _, err := p.tx.Exec("ROLLBACK TO SAVEPOINT insert_changeset"); err != nil {
			return err
		}
		if _, err := p.changeUpdateStmt.Exec(
			c.Id,
			c.CreatedAt.UTC(),
			closedAt,
			c.Open,
			c.NumChanges,
			c.User,
			c.UserId,
			bbox,
			hstoreStringChangeset(c.Tags),
		); err != nil {
			return newSqlError(err, c)
		}
		if _, err := p.tx.Exec(fmt.Sprintf(`DELETE FROM "%[1]s".comments WHERE changeset_id = $1`, p.schema), c.Id); err != nil {
			return err
		}
	}

	for i, com := range c.Comments {
		if _, err := p.commentStmt.Exec(
			c.Id,
			i,
			com.User,
			com.UserId,
			com.Date.UTC(),
			com.Text,
		); err != nil {
			return newSqlError(err, c)
		}
	}

	if _, err := p.tx.Exec("RELEASE SAVEPOINT insert_changeset"); err != nil {
		return err
	}

	return nil
}

var hstoreReplacer = strings.NewReplacer("\\", "\\\\", "\"", "\\\"")

func hstoreString(tags element.Tags) string {
	kv := make([]string, 0, len(tags))
	for k, v := range tags {
		kv = append(kv, `"`+hstoreReplacer.Replace(k)+`"=>"`+hstoreReplacer.Replace(v)+`"`)
	}
	return strings.Join(kv, ", ")
}

func hstoreStringChangeset(tags []changeset.Tag) string {
	kv := make([]string, 0, len(tags))
	for _, t := range tags {
		kv = append(kv, `"`+hstoreReplacer.Replace(t.Key)+`"=>"`+hstoreReplacer.Replace(t.Value)+`"`)
	}
	return strings.Join(kv, ", ")
}

func bboxPolygon(c changeset.Changeset) interface{} {
	if c.MinLon != 0.0 && c.MaxLon != 0.0 && c.MinLat != 0.0 && c.MaxLat != 0.0 {
		return fmt.Sprintf(
			"SRID=4326; POLYGON((%[1]f %[2]f, %[1]f %[4]f, %[3]f %[4]f, %[3]f %[2]f, %[1]f %[2]f))",
			c.MinLon, c.MinLat, c.MaxLon, c.MaxLat,
		)
	}
	return nil
}
<|endoftext|>


<commit_before>// Package ethwire provides low level access to the Ethereum network and allows
// you to broadcast data over the network.
package ethwire

import (
	"bytes"
	"errors"
	"fmt"
	"github.com/ethereum/eth-go/ethutil"
	"net"
	"time"
)

// Connection interface describing the methods required to implement the wire protocol.
type Conn interface {
	Write(typ MsgType, v ...interface{}) error
	Read() *Msg
}

// The magic token which should be the first 4 bytes of every message and can be used as separator between messages.
var MagicToken = []byte{34, 64, 8, 145}

type MsgType byte

const (
	// Values are given explicitly instead of by iota because these values are
	// defined by the wire protocol spec; it is easier for humans to ensure
	// correctness when values are explicit.
	MsgHandshakeTy  = 0x00
	MsgDiscTy       = 0x01
	MsgPingTy       = 0x02
	MsgPongTy       = 0x03
	MsgGetPeersTy   = 0x10
	MsgPeersTy      = 0x11
	MsgTxTy         = 0x12
	MsgBlockTy      = 0x13
	MsgGetChainTy   = 0x14
	MsgNotInChainTy = 0x15
	MsgGetTxsTy     = 0x16

	MsgTalkTy = 0xff
)

var msgTypeToString = map[MsgType]string{
	MsgHandshakeTy:  "Handshake",
	MsgDiscTy:       "Disconnect",
	MsgPingTy:       "Ping",
	MsgPongTy:       "Pong",
	MsgGetPeersTy:   "Get peers",
	MsgPeersTy:      "Peers",
	MsgTxTy:         "Transactions",
	MsgBlockTy:      "Blocks",
	MsgGetChainTy:   "Get chain",
	MsgGetTxsTy:     "Get Txs",
	MsgNotInChainTy: "Not in chain",
}

func (mt MsgType) String() string {
	return msgTypeToString[mt]
}

type Msg struct {
	Type MsgType // Specifies how the encoded data should be interpreted
	//Data []byte
	Data *ethutil.Value
}

func NewMessage(msgType MsgType, data interface{}) *Msg {
	return &Msg{
		Type: msgType,
		Data: ethutil.NewValue(data),
	}
}

type Messages []*Msg

// The connection object allows you to set up a connection to the Ethereum network.
// The Connection object takes care of all encoding and sending objects properly over
// the network.
type Connection struct {
	conn            net.Conn
	nTimeout        time.Duration
	pendingMessages Messages
}

// Create a new connection to the Ethereum network
func New(conn net.Conn) *Connection {
	return &Connection{conn: conn, nTimeout: 500}
}

// Read, reads from the network. It will block until the next message is received.
func (self *Connection) Read() *Msg {
	if len(self.pendingMessages) == 0 {
		self.readMessages()
	}

	ret := self.pendingMessages[0]
	self.pendingMessages = self.pendingMessages[1:]

	return ret

}

// Write to the Ethereum network specifying the type of the message and
// the data. Data can be of type RlpEncodable or []interface{}. Returns
// nil or if something went wrong an error.
func (self *Connection) Write(typ MsgType, v ...interface{}) error {
	var pack []byte

	slice := [][]interface{}{[]interface{}{byte(typ)}}
	for _, value := range v {
		if encodable, ok := value.(ethutil.RlpEncodable); ok {
			slice = append(slice, encodable.RlpValue())
		} else if raw, ok := value.([]interface{}); ok {
			slice = append(slice, raw)
		} else {
			panic(fmt.Sprintf("Unable to 'write' object of type %T", value))
		}
	}

	// Encode the type and the (RLP encoded) data for sending over the wire
	encoded := ethutil.NewValue(slice).Encode()
	payloadLength := ethutil.NumberToBytes(uint32(len(encoded)), 32)

	// Write magic token and payload length (first 8 bytes)
	pack = append(MagicToken, payloadLength...)
	pack = append(pack, encoded...)

	// Write to the connection
	_, err := self.conn.Write(pack)
	if err != nil {
		return err
	}

	return nil
}

func (self *Connection) readMessage(data []byte) (msg *Msg, remaining []byte, done bool, err error) {
	if len(data) == 0 {
		return nil, nil, true, nil
	}

	if len(data) <= 8 {
		return nil, remaining, false, errors.New("Invalid message")
	}

	// Check if the received 4 first bytes are the magic token
	if bytes.Compare(MagicToken, data[:4]) != 0 {
		return nil, nil, false, fmt.Errorf("MagicToken mismatch. Received %v", data[:4])
	}

	messageLength := ethutil.BytesToNumber(data[4:8])
	remaining = data[8+messageLength:]
	if int(messageLength) > len(data[8:]) {
		return nil, nil, false, fmt.Errorf("message length %d, expected %d", len(data[8:]), messageLength)
	}

	message := data[8 : 8+messageLength]
	decoder := ethutil.NewValueFromBytes(message)
	// Type of message
	t := decoder.Get(0).Uint()
	// Actual data
	d := decoder.SliceFrom(1)

	msg = &Msg{
		Type: MsgType(t),
		Data: d,
	}

	return
}

// The basic message reader waits for data on the given connection, decoding
// and doing a few sanity checks such as if there's a data type and
// unmarhals the given data
func (self *Connection) readMessages() (err error) {
	// The recovering function in case anything goes horribly wrong
	defer func() {
		if r := recover(); r != nil {
			err = fmt.Errorf("ethwire.ReadMessage error: %v", r)
		}
	}()

	// Buff for writing network message to
	//buff := make([]byte, 1440)
	var buff []byte
	var totalBytes int
	for {
		// Give buffering some time
		self.conn.SetReadDeadline(time.Now().Add(self.nTimeout * time.Millisecond))
		// Create a new temporarily buffer
		b := make([]byte, 1440)
		// Wait for a message from this peer
		n, _ := self.conn.Read(b)
		if err != nil && n == 0 {
			if err.Error() != "EOF" {
				fmt.Println("err now", err)
				return err
			} else {
				break
			}

			// Messages can't be empty
		} else if n == 0 {
			break
		}

		buff = append(buff, b[:n]...)
		totalBytes += n
	}

	// Reslice buffer
	buff = buff[:totalBytes]
	msg, remaining, done, err := self.readMessage(buff)
	for ; done != true; msg, remaining, done, err = self.readMessage(remaining) {
		//log.Println("rx", msg)

		if msg != nil {
			self.pendingMessages = append(self.pendingMessages, msg)
		}
	}

	return
}

func ReadMessage(data []byte) (msg *Msg, remaining []byte, done bool, err error) {
	if len(data) == 0 {
		return nil, nil, true, nil
	}

	if len(data) <= 8 {
		return nil, remaining, false, errors.New("Invalid message")
	}

	// Check if the received 4 first bytes are the magic token
	if bytes.Compare(MagicToken, data[:4]) != 0 {
		return nil, nil, false, fmt.Errorf("MagicToken mismatch. Received %v", data[:4])
	}

	messageLength := ethutil.BytesToNumber(data[4:8])
	remaining = data[8+messageLength:]
	if int(messageLength) > len(data[8:]) {
		return nil, nil, false, fmt.Errorf("message length %d, expected %d", len(data[8:]), messageLength)
	}

	message := data[8 : 8+messageLength]
	decoder := ethutil.NewValueFromBytes(message)
	// Type of message
	t := decoder.Get(0).Uint()
	// Actual data
	d := decoder.SliceFrom(1)

	msg = &Msg{
		Type: MsgType(t),
		Data: d,
	}

	return
}

func bufferedRead(conn net.Conn) ([]byte, error) {
	return nil, nil
}

// The basic message reader waits for data on the given connection, decoding
// and doing a few sanity checks such as if there's a data type and
// unmarhals the given data
func ReadMessages(conn net.Conn) (msgs []*Msg, err error) {
	// The recovering function in case anything goes horribly wrong
	defer func() {
		if r := recover(); r != nil {
			err = fmt.Errorf("ethwire.ReadMessage error: %v", r)
		}
	}()

	// Buff for writing network message to
	//buff := make([]byte, 1440)
	var buff []byte
	var totalBytes int
	for {
		// Give buffering some time
		conn.SetReadDeadline(time.Now().Add(200 * time.Millisecond))
		// Create a new temporarily buffer
		b := make([]byte, 1440)
		// Wait for a message from this peer
		n, _ := conn.Read(b)
		if err != nil && n == 0 {
			if err.Error() != "EOF" {
				fmt.Println("err now", err)
				return nil, err
			} else {
				break
			}

			// Messages can't be empty
		} else if n == 0 {
			break
		}

		buff = append(buff, b[:n]...)
		totalBytes += n
	}

	// Reslice buffer
	buff = buff[:totalBytes]
	msg, remaining, done, err := ReadMessage(buff)
	for ; done != true; msg, remaining, done, err = ReadMessage(remaining) {
		//log.Println("rx", msg)

		if msg != nil {
			msgs = append(msgs, msg)
		}
	}

	return
}

// The basic message writer takes care of writing data over the given
// connection and does some basic error checking
func WriteMessage(conn net.Conn, msg *Msg) error {
	var pack []byte

	// Encode the type and the (RLP encoded) data for sending over the wire
	encoded := ethutil.NewValue(append([]interface{}{byte(msg.Type)}, msg.Data.Slice()...)).Encode()
	payloadLength := ethutil.NumberToBytes(uint32(len(encoded)), 32)

	// Write magic token and payload length (first 8 bytes)
	pack = append(MagicToken, payloadLength...)
	pack = append(pack, encoded...)
	//fmt.Printf("payload %v (%v) %q\n", msg.Type, conn.RemoteAddr(), encoded)

	// Write to the connection
	_, err := conn.Write(pack)
	if err != nil {
		return err
	}

	return nil
}
<commit_msg>Increased timeout to 500ms<commit_after>// Package ethwire provides low level access to the Ethereum network and allows
// you to broadcast data over the network.
package ethwire

import (
	"bytes"
	"errors"
	"fmt"
	"github.com/ethereum/eth-go/ethutil"
	"net"
	"time"
)

// Connection interface describing the methods required to implement the wire protocol.
type Conn interface {
	Write(typ MsgType, v ...interface{}) error
	Read() *Msg
}

// The magic token which should be the first 4 bytes of every message and can be used as separator between messages.
var MagicToken = []byte{34, 64, 8, 145}

type MsgType byte

const (
	// Values are given explicitly instead of by iota because these values are
	// defined by the wire protocol spec; it is easier for humans to ensure
	// correctness when values are explicit.
	MsgHandshakeTy  = 0x00
	MsgDiscTy       = 0x01
	MsgPingTy       = 0x02
	MsgPongTy       = 0x03
	MsgGetPeersTy   = 0x10
	MsgPeersTy      = 0x11
	MsgTxTy         = 0x12
	MsgBlockTy      = 0x13
	MsgGetChainTy   = 0x14
	MsgNotInChainTy = 0x15
	MsgGetTxsTy     = 0x16

	MsgTalkTy = 0xff
)

var msgTypeToString = map[MsgType]string{
	MsgHandshakeTy:  "Handshake",
	MsgDiscTy:       "Disconnect",
	MsgPingTy:       "Ping",
	MsgPongTy:       "Pong",
	MsgGetPeersTy:   "Get peers",
	MsgPeersTy:      "Peers",
	MsgTxTy:         "Transactions",
	MsgBlockTy:      "Blocks",
	MsgGetChainTy:   "Get chain",
	MsgGetTxsTy:     "Get Txs",
	MsgNotInChainTy: "Not in chain",
}

func (mt MsgType) String() string {
	return msgTypeToString[mt]
}

type Msg struct {
	Type MsgType // Specifies how the encoded data should be interpreted
	//Data []byte
	Data *ethutil.Value
}

func NewMessage(msgType MsgType, data interface{}) *Msg {
	return &Msg{
		Type: msgType,
		Data: ethutil.NewValue(data),
	}
}

type Messages []*Msg

// The connection object allows you to set up a connection to the Ethereum network.
// The Connection object takes care of all encoding and sending objects properly over
// the network.
type Connection struct {
	conn            net.Conn
	nTimeout        time.Duration
	pendingMessages Messages
}

// Create a new connection to the Ethereum network
func New(conn net.Conn) *Connection {
	return &Connection{conn: conn, nTimeout: 500}
}

// Read, reads from the network. It will block until the next message is received.
func (self *Connection) Read() *Msg {
	if len(self.pendingMessages) == 0 {
		self.readMessages()
	}

	ret := self.pendingMessages[0]
	self.pendingMessages = self.pendingMessages[1:]

	return ret

}

// Write to the Ethereum network specifying the type of the message and
// the data. Data can be of type RlpEncodable or []interface{}. Returns
// nil or if something went wrong an error.
func (self *Connection) Write(typ MsgType, v ...interface{}) error {
	var pack []byte

	slice := [][]interface{}{[]interface{}{byte(typ)}}
	for _, value := range v {
		if encodable, ok := value.(ethutil.RlpEncodable); ok {
			slice = append(slice, encodable.RlpValue())
		} else if raw, ok := value.([]interface{}); ok {
			slice = append(slice, raw)
		} else {
			panic(fmt.Sprintf("Unable to 'write' object of type %T", value))
		}
	}

	// Encode the type and the (RLP encoded) data for sending over the wire
	encoded := ethutil.NewValue(slice).Encode()
	payloadLength := ethutil.NumberToBytes(uint32(len(encoded)), 32)

	// Write magic token and payload length (first 8 bytes)
	pack = append(MagicToken, payloadLength...)
	pack = append(pack, encoded...)

	// Write to the connection
	_, err := self.conn.Write(pack)
	if err != nil {
		return err
	}

	return nil
}

func (self *Connection) readMessage(data []byte) (msg *Msg, remaining []byte, done bool, err error) {
	if len(data) == 0 {
		return nil, nil, true, nil
	}

	if len(data) <= 8 {
		return nil, remaining, false, errors.New("Invalid message")
	}

	// Check if the received 4 first bytes are the magic token
	if bytes.Compare(MagicToken, data[:4]) != 0 {
		return nil, nil, false, fmt.Errorf("MagicToken mismatch. Received %v", data[:4])
	}

	messageLength := ethutil.BytesToNumber(data[4:8])
	remaining = data[8+messageLength:]
	if int(messageLength) > len(data[8:]) {
		return nil, nil, false, fmt.Errorf("message length %d, expected %d", len(data[8:]), messageLength)
	}

	message := data[8 : 8+messageLength]
	decoder := ethutil.NewValueFromBytes(message)
	// Type of message
	t := decoder.Get(0).Uint()
	// Actual data
	d := decoder.SliceFrom(1)

	msg = &Msg{
		Type: MsgType(t),
		Data: d,
	}

	return
}

// The basic message reader waits for data on the given connection, decoding
// and doing a few sanity checks such as if there's a data type and
// unmarhals the given data
func (self *Connection) readMessages() (err error) {
	// The recovering function in case anything goes horribly wrong
	defer func() {
		if r := recover(); r != nil {
			err = fmt.Errorf("ethwire.ReadMessage error: %v", r)
		}
	}()

	// Buff for writing network message to
	//buff := make([]byte, 1440)
	var buff []byte
	var totalBytes int
	for {
		// Give buffering some time
		self.conn.SetReadDeadline(time.Now().Add(self.nTimeout * time.Millisecond))
		// Create a new temporarily buffer
		b := make([]byte, 1440)
		// Wait for a message from this peer
		n, _ := self.conn.Read(b)
		if err != nil && n == 0 {
			if err.Error() != "EOF" {
				fmt.Println("err now", err)
				return err
			} else {
				break
			}

			// Messages can't be empty
		} else if n == 0 {
			break
		}

		buff = append(buff, b[:n]...)
		totalBytes += n
	}

	// Reslice buffer
	buff = buff[:totalBytes]
	msg, remaining, done, err := self.readMessage(buff)
	for ; done != true; msg, remaining, done, err = self.readMessage(remaining) {
		//log.Println("rx", msg)

		if msg != nil {
			self.pendingMessages = append(self.pendingMessages, msg)
		}
	}

	return
}

func ReadMessage(data []byte) (msg *Msg, remaining []byte, done bool, err error) {
	if len(data) == 0 {
		return nil, nil, true, nil
	}

	if len(data) <= 8 {
		return nil, remaining, false, errors.New("Invalid message")
	}

	// Check if the received 4 first bytes are the magic token
	if bytes.Compare(MagicToken, data[:4]) != 0 {
		return nil, nil, false, fmt.Errorf("MagicToken mismatch. Received %v", data[:4])
	}

	messageLength := ethutil.BytesToNumber(data[4:8])
	remaining = data[8+messageLength:]
	if int(messageLength) > len(data[8:]) {
		return nil, nil, false, fmt.Errorf("message length %d, expected %d", len(data[8:]), messageLength)
	}

	message := data[8 : 8+messageLength]
	decoder := ethutil.NewValueFromBytes(message)
	// Type of message
	t := decoder.Get(0).Uint()
	// Actual data
	d := decoder.SliceFrom(1)

	msg = &Msg{
		Type: MsgType(t),
		Data: d,
	}

	return
}

func bufferedRead(conn net.Conn) ([]byte, error) {
	return nil, nil
}

// The basic message reader waits for data on the given connection, decoding
// and doing a few sanity checks such as if there's a data type and
// unmarhals the given data
func ReadMessages(conn net.Conn) (msgs []*Msg, err error) {
	// The recovering function in case anything goes horribly wrong
	defer func() {
		if r := recover(); r != nil {
			err = fmt.Errorf("ethwire.ReadMessage error: %v", r)
		}
	}()

	// Buff for writing network message to
	//buff := make([]byte, 1440)
	var buff []byte
	var totalBytes int
	for {
		// Give buffering some time
		conn.SetReadDeadline(time.Now().Add(500 * time.Millisecond))
		// Create a new temporarily buffer
		b := make([]byte, 1440)
		// Wait for a message from this peer
		n, _ := conn.Read(b)
		if err != nil && n == 0 {
			if err.Error() != "EOF" {
				fmt.Println("err now", err)
				return nil, err
			} else {
				break
			}

			// Messages can't be empty
		} else if n == 0 {
			break
		}

		buff = append(buff, b[:n]...)
		totalBytes += n
	}

	// Reslice buffer
	buff = buff[:totalBytes]
	msg, remaining, done, err := ReadMessage(buff)
	for ; done != true; msg, remaining, done, err = ReadMessage(remaining) {
		//log.Println("rx", msg)

		if msg != nil {
			msgs = append(msgs, msg)
		}
	}

	return
}

// The basic message writer takes care of writing data over the given
// connection and does some basic error checking
func WriteMessage(conn net.Conn, msg *Msg) error {
	var pack []byte

	// Encode the type and the (RLP encoded) data for sending over the wire
	encoded := ethutil.NewValue(append([]interface{}{byte(msg.Type)}, msg.Data.Slice()...)).Encode()
	payloadLength := ethutil.NumberToBytes(uint32(len(encoded)), 32)

	// Write magic token and payload length (first 8 bytes)
	pack = append(MagicToken, payloadLength...)
	pack = append(pack, encoded...)
	//fmt.Printf("payload %v (%v) %q\n", msg.Type, conn.RemoteAddr(), encoded)

	// Write to the connection
	_, err := conn.Write(pack)
	if err != nil {
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>package adaptors

import (
	"net/http"

	"github.com/codegangsta/negroni"
)

func FromNegroni(handler negroni.Handler) func(http.Handler) http.Handler {
	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
			n := negroni.New()
			n.Use(handler)
			n.UseHandler(next)
			n.ServeHTTP(rw, req)
		})
	}
}

func HandlerFromNegroni(handler negroni.Handler) http.Handler {
	return http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
		n := negroni.New()
		n.Use(handler)
		n.ServeHTTP(rw, req)
	})
}
<commit_msg>Another small correction to the Negroni adapter to avoid over-aggressive teardown and reconstruction of the Negroni framework.<commit_after>package adaptors

import (
	"net/http"

	"github.com/codegangsta/negroni"
)

func FromNegroni(handler negroni.Handler) func(http.Handler) http.Handler {
	return func(next http.Handler) http.Handler {
		n := negroni.New()
		n.Use(handler)
		n.UseHandler(next)
		return http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
			n.ServeHTTP(rw, req)
		})
	}
}

func HandlerFromNegroni(handler negroni.Handler) http.Handler {
	n := negroni.New()
	n.Use(handler)
	return http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
		n.ServeHTTP(rw, req)
	})
}
<|endoftext|>


<commit_before>package githubfetch

import (
	"archive/tar"
	"compress/gzip"
	"context"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"path"
	"strings"
	"time"

	"github.com/google/go-github/github"
	"golang.org/x/oauth2"
	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer"
)

const (
	githubDownloadTimeoutSecs = 300
)

// CodeFetcher represents an object capable of fetching code and returning a
// gzip-compressed tarball io.Reader
type CodeFetcher interface {
	GetCommitSHA(tracer.Span, string, string, string) (string, error)
	Get(tracer.Span, string, string, string) (io.Reader, error)
}

// GitHubFetcher represents a github data fetcher
type GitHubFetcher struct {
	c *github.Client
}

// NewGitHubFetcher returns a new github fetcher
func NewGitHubFetcher(token string) *GitHubFetcher {
	ts := oauth2.StaticTokenSource(&oauth2.Token{AccessToken: token})
	tc := oauth2.NewClient(oauth2.NoContext, ts)
	gf := &GitHubFetcher{
		c: github.NewClient(tc),
	}
	return gf
}

// GetCommitSHA returns the commit SHA for a reference
func (gf *GitHubFetcher) GetCommitSHA(parentSpan tracer.Span, owner string, repo string, ref string) (csha string, err error) {
	span := tracer.StartSpan("github_fetcher.get_commit_sha", tracer.ChildOf(parentSpan.Context()))
	defer func() {
		span.Finish(tracer.WithError(err))
	}()
	ctx, cf := context.WithTimeout(context.Background(), githubDownloadTimeoutSecs*time.Second)
	defer cf()
	csha, _, err = gf.c.Repositories.GetCommitSHA1(ctx, owner, repo, ref, "")
	return csha, err
}

// Get fetches contents of GitHub repo and returns the processed contents as
// an in-memory io.Reader.
func (gf *GitHubFetcher) Get(parentSpan tracer.Span, owner string, repo string, ref string) (tarball io.Reader, err error) {
	span := tracer.StartSpan("github_fetcher.get", tracer.ChildOf(parentSpan.Context()))
	defer func() {
		span.Finish(tracer.WithError(err))
	}()
	opt := &github.RepositoryContentGetOptions{
		Ref: ref,
	}
	ctx, cf := context.WithTimeout(context.Background(), githubDownloadTimeoutSecs*time.Second)
	defer cf()
	url, resp, err := gf.c.Repositories.GetArchiveLink(ctx, owner, repo, github.Tarball, opt)
	if err != nil {
		return nil, fmt.Errorf("error getting archive link: %v", err)
	}
	if resp.StatusCode > 399 {
		return nil, fmt.Errorf("error status when getting archive link: %v", resp.Status)
	}
	if url == nil {
		return nil, fmt.Errorf("url is nil")
	}
	return gf.getArchive(url)
}

func (gf *GitHubFetcher) getArchive(archiveURL *url.URL) (io.Reader, error) {
	hc := http.Client{
		Timeout: githubDownloadTimeoutSecs * time.Second,
	}
	hr, err := http.NewRequest("GET", archiveURL.String(), nil)
	if err != nil {
		return nil, fmt.Errorf("error creating http request: %v", err)
	}
	resp, err := hc.Do(hr)
	if err != nil {
		return nil, fmt.Errorf("error performing archive http request: %v", err)
	}
	if resp == nil {
		return nil, fmt.Errorf("error getting archive: response is nil")
	}
	if resp.StatusCode > 299 {
		return nil, fmt.Errorf("archive http request failed: %v", resp.StatusCode)
	}
	return newTarPrefixStripper(resp.Body), nil
}

func (gf *GitHubFetcher) debugWriteTar(contents []byte) {
	f, err := ioutil.TempFile("", "output-tar")
	defer f.Close()
	log.Printf("debug: saving tar output to %v", f.Name())
	_, err = f.Write(contents)
	if err != nil {
		log.Printf("debug: error writing tar output: %v", err)
	}
}

// tarPrefixStripper removes a random path that Github prefixes its
// archives with.
type tarPrefixStripper struct {
	tarball          io.ReadCloser
	pipeReader       *io.PipeReader
	pipeWriter       *io.PipeWriter
	strippingStarted bool
}

func newTarPrefixStripper(tarball io.ReadCloser) io.Reader {
	reader, writer := io.Pipe()
	return &tarPrefixStripper{
		tarball:    tarball,
		pipeReader: reader,
		pipeWriter: writer,
	}
}

func (t *tarPrefixStripper) Read(p []byte) (n int, err error) {
	if !t.strippingStarted {
		go t.startStrippingPipe()
		t.strippingStarted = true
	}
	return t.pipeReader.Read(p)
}

func (t *tarPrefixStripper) processHeader(h *tar.Header) (bool, error) {
	// metadata file, ignore
	if h.Name == "pax_global_header" {
		return true, nil
	}
	if path.IsAbs(h.Name) {
		return true, fmt.Errorf("archive contains absolute path: %v", h.Name)
	}

	// top-level directory entry
	spath := strings.Split(h.Name, "/")
	if len(spath) == 2 && spath[1] == "" {
		return true, nil
	}
	h.Name = strings.Join(spath[1:len(spath)], "/")

	return false, nil
}

func (t *tarPrefixStripper) startStrippingPipe() {
	gzr, err := gzip.NewReader(t.tarball)
	if err != nil {
		t.pipeWriter.CloseWithError(err)
		return
	}

	tarball := tar.NewReader(gzr)
	outTarball := tar.NewWriter(t.pipeWriter)

	closeFunc := func(e error) {
		outTarball.Close()
		t.pipeWriter.CloseWithError(e)
		t.tarball.Close()
	}

	for {
		header, err := tarball.Next()
		if err == io.EOF {
			closeFunc(nil)
			return
		}
		if err != nil {
			closeFunc(err)
			return
		}

		skip, err := t.processHeader(header)
		if err != nil {
			closeFunc(err)
			return
		}
		if skip {
			continue
		}

		if err := outTarball.WriteHeader(header); err != nil {
			closeFunc(err)
			return
		}
		if _, err := io.Copy(outTarball, tarball); err != nil {
			closeFunc(err)
			return
		}
		if err := outTarball.Flush(); err != nil {
			closeFunc(err)
			return
		}
	}
}
<commit_msg>lib/githubfetch: Log dropped error<commit_after>package githubfetch

import (
	"archive/tar"
	"compress/gzip"
	"context"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"path"
	"strings"
	"time"

	"github.com/google/go-github/github"
	"golang.org/x/oauth2"
	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer"
)

const (
	githubDownloadTimeoutSecs = 300
)

// CodeFetcher represents an object capable of fetching code and returning a
// gzip-compressed tarball io.Reader
type CodeFetcher interface {
	GetCommitSHA(tracer.Span, string, string, string) (string, error)
	Get(tracer.Span, string, string, string) (io.Reader, error)
}

// GitHubFetcher represents a github data fetcher
type GitHubFetcher struct {
	c *github.Client
}

// NewGitHubFetcher returns a new github fetcher
func NewGitHubFetcher(token string) *GitHubFetcher {
	ts := oauth2.StaticTokenSource(&oauth2.Token{AccessToken: token})
	tc := oauth2.NewClient(oauth2.NoContext, ts)
	gf := &GitHubFetcher{
		c: github.NewClient(tc),
	}
	return gf
}

// GetCommitSHA returns the commit SHA for a reference
func (gf *GitHubFetcher) GetCommitSHA(parentSpan tracer.Span, owner string, repo string, ref string) (csha string, err error) {
	span := tracer.StartSpan("github_fetcher.get_commit_sha", tracer.ChildOf(parentSpan.Context()))
	defer func() {
		span.Finish(tracer.WithError(err))
	}()
	ctx, cf := context.WithTimeout(context.Background(), githubDownloadTimeoutSecs*time.Second)
	defer cf()
	csha, _, err = gf.c.Repositories.GetCommitSHA1(ctx, owner, repo, ref, "")
	return csha, err
}

// Get fetches contents of GitHub repo and returns the processed contents as
// an in-memory io.Reader.
func (gf *GitHubFetcher) Get(parentSpan tracer.Span, owner string, repo string, ref string) (tarball io.Reader, err error) {
	span := tracer.StartSpan("github_fetcher.get", tracer.ChildOf(parentSpan.Context()))
	defer func() {
		span.Finish(tracer.WithError(err))
	}()
	opt := &github.RepositoryContentGetOptions{
		Ref: ref,
	}
	ctx, cf := context.WithTimeout(context.Background(), githubDownloadTimeoutSecs*time.Second)
	defer cf()
	url, resp, err := gf.c.Repositories.GetArchiveLink(ctx, owner, repo, github.Tarball, opt)
	if err != nil {
		return nil, fmt.Errorf("error getting archive link: %v", err)
	}
	if resp.StatusCode > 399 {
		return nil, fmt.Errorf("error status when getting archive link: %v", resp.Status)
	}
	if url == nil {
		return nil, fmt.Errorf("url is nil")
	}
	return gf.getArchive(url)
}

func (gf *GitHubFetcher) getArchive(archiveURL *url.URL) (io.Reader, error) {
	hc := http.Client{
		Timeout: githubDownloadTimeoutSecs * time.Second,
	}
	hr, err := http.NewRequest("GET", archiveURL.String(), nil)
	if err != nil {
		return nil, fmt.Errorf("error creating http request: %v", err)
	}
	resp, err := hc.Do(hr)
	if err != nil {
		return nil, fmt.Errorf("error performing archive http request: %v", err)
	}
	if resp == nil {
		return nil, fmt.Errorf("error getting archive: response is nil")
	}
	if resp.StatusCode > 299 {
		return nil, fmt.Errorf("archive http request failed: %v", resp.StatusCode)
	}
	return newTarPrefixStripper(resp.Body), nil
}

func (gf *GitHubFetcher) debugWriteTar(contents []byte) {
	f, err := ioutil.TempFile("", "output-tar")
	if err != nil {
		log.Printf("debug: error creating TempFile: %v", err)
	}
	defer f.Close()
	log.Printf("debug: saving tar output to %v", f.Name())
	_, err = f.Write(contents)
	if err != nil {
		log.Printf("debug: error writing tar output: %v", err)
	}
}

// tarPrefixStripper removes a random path that Github prefixes its
// archives with.
type tarPrefixStripper struct {
	tarball          io.ReadCloser
	pipeReader       *io.PipeReader
	pipeWriter       *io.PipeWriter
	strippingStarted bool
}

func newTarPrefixStripper(tarball io.ReadCloser) io.Reader {
	reader, writer := io.Pipe()
	return &tarPrefixStripper{
		tarball:    tarball,
		pipeReader: reader,
		pipeWriter: writer,
	}
}

func (t *tarPrefixStripper) Read(p []byte) (n int, err error) {
	if !t.strippingStarted {
		go t.startStrippingPipe()
		t.strippingStarted = true
	}
	return t.pipeReader.Read(p)
}

func (t *tarPrefixStripper) processHeader(h *tar.Header) (bool, error) {
	// metadata file, ignore
	if h.Name == "pax_global_header" {
		return true, nil
	}
	if path.IsAbs(h.Name) {
		return true, fmt.Errorf("archive contains absolute path: %v", h.Name)
	}

	// top-level directory entry
	spath := strings.Split(h.Name, "/")
	if len(spath) == 2 && spath[1] == "" {
		return true, nil
	}
	h.Name = strings.Join(spath[1:len(spath)], "/")

	return false, nil
}

func (t *tarPrefixStripper) startStrippingPipe() {
	gzr, err := gzip.NewReader(t.tarball)
	if err != nil {
		t.pipeWriter.CloseWithError(err)
		return
	}

	tarball := tar.NewReader(gzr)
	outTarball := tar.NewWriter(t.pipeWriter)

	closeFunc := func(e error) {
		outTarball.Close()
		t.pipeWriter.CloseWithError(e)
		t.tarball.Close()
	}

	for {
		header, err := tarball.Next()
		if err == io.EOF {
			closeFunc(nil)
			return
		}
		if err != nil {
			closeFunc(err)
			return
		}

		skip, err := t.processHeader(header)
		if err != nil {
			closeFunc(err)
			return
		}
		if skip {
			continue
		}

		if err := outTarball.WriteHeader(header); err != nil {
			closeFunc(err)
			return
		}
		if _, err := io.Copy(outTarball, tarball); err != nil {
			closeFunc(err)
			return
		}
		if err := outTarball.Flush(); err != nil {
			closeFunc(err)
			return
		}
	}
}
<|endoftext|>