<commit_before>package reader

import (
	"context"
	"encoding/binary"
	"io"
	"reflect"
	"strings"
	"sync"

	"github.com/apache/thrift/lib/go/thrift"
	"github.com/xitongsys/parquet-go/common"
	"github.com/xitongsys/parquet-go/layout"
	"github.com/xitongsys/parquet-go/marshal"
	"github.com/xitongsys/parquet-go/parquet"
	"github.com/xitongsys/parquet-go/schema"
	"github.com/xitongsys/parquet-go/source"
)

type ParquetReader struct {
	SchemaHandler *schema.SchemaHandler
	NP            int64 //parallel number
	Footer        *parquet.FileMetaData
	PFile         source.ParquetFile

	ColumnBuffers map[string]*ColumnBufferType

	//One reader can only read one type objects
	ObjType        reflect.Type
	ObjPartialType reflect.Type
}

//Create a parquet reader: obj is a object with schema tags or a JSON schema string
func NewParquetReader(pFile source.ParquetFile, obj interface{}, np int64) (*ParquetReader, error) {
	var err error
	res := new(ParquetReader)
	res.NP = np
	res.PFile = pFile
	if err = res.ReadFooter(); err != nil {
		return nil, err
	}
	res.ColumnBuffers = make(map[string]*ColumnBufferType)

	if obj != nil {
		if sa, ok := obj.(string); ok {
			err = res.SetSchemaHandlerFromJSON(sa)
			return res, err

		} else if sa, ok := obj.([]*parquet.SchemaElement); ok {
			res.SchemaHandler = schema.NewSchemaHandlerFromSchemaList(sa)

		} else {
			if res.SchemaHandler, err = schema.NewSchemaHandlerFromStruct(obj); err != nil {
				return res, err
			}
		}

	} else {
		res.SchemaHandler = schema.NewSchemaHandlerFromSchemaList(res.Footer.Schema)
	}

	res.RenameSchema()
	for i := 0; i < len(res.SchemaHandler.SchemaElements); i++ {
		schema := res.SchemaHandler.SchemaElements[i]
		if schema.GetNumChildren() == 0 {
			pathStr := res.SchemaHandler.IndexMap[int32(i)]
			if res.ColumnBuffers[pathStr], err = NewColumnBuffer(pFile, res.Footer, res.SchemaHandler, pathStr); err != nil {
				return res, err
			}
		}
	}

	return res, nil
}

func (self *ParquetReader) SetSchemaHandlerFromJSON(jsonSchema string) error {
	var err error

	if self.SchemaHandler, err = schema.NewSchemaHandlerFromJSON(jsonSchema); err != nil {
		return err
	}

	self.RenameSchema()
	for i := 0; i < len(self.SchemaHandler.SchemaElements); i++ {
		schemaElement := self.SchemaHandler.SchemaElements[i]
		if schemaElement.GetNumChildren() == 0 {
			pathStr := self.SchemaHandler.IndexMap[int32(i)]
			if self.ColumnBuffers[pathStr], err = NewColumnBuffer(self.PFile, self.Footer, self.SchemaHandler, pathStr); err != nil {
				return err
			}
		}
	}
	return nil
}

//Rename schema name to inname
func (self *ParquetReader) RenameSchema() {
	for i := 0; i < len(self.SchemaHandler.Infos); i++ {
		self.Footer.Schema[i].Name = self.SchemaHandler.Infos[i].InName
	}
	for _, rowGroup := range self.Footer.RowGroups {
		for _, chunk := range rowGroup.Columns {
			exPath := make([]string, 0)
			exPath = append(exPath, self.SchemaHandler.GetRootExName())
			exPath = append(exPath, chunk.MetaData.GetPathInSchema()...)
			exPathStr := common.PathToStr(exPath)

			inPathStr := self.SchemaHandler.ExPathToInPath[exPathStr]
			inPath := common.StrToPath(inPathStr)[1:]
			chunk.MetaData.PathInSchema = inPath
		}
	}
}

func (self *ParquetReader) GetNumRows() int64 {
	return self.Footer.GetNumRows()
}

//Get the footer size
func (self *ParquetReader) GetFooterSize() (uint32, error) {
	var err error
	buf := make([]byte, 4)
	if _, err = self.PFile.Seek(-8, io.SeekEnd); err != nil {
		return 0, err
	}
	if _, err = io.ReadFull(self.PFile, buf); err != nil {
		return 0, err
	}
	size := binary.LittleEndian.Uint32(buf)
	return size, err
}

//Read footer from parquet file
func (self *ParquetReader) ReadFooter() error {
	size, err := self.GetFooterSize()
	if err != nil {
		return err
	}
	if _, err = self.PFile.Seek(-(int64)(8+size), io.SeekEnd); err != nil {
		return err
	}
	self.Footer = parquet.NewFileMetaData()
	pf := thrift.NewTCompactProtocolFactory()
	protocol := pf.GetProtocol(thrift.NewStreamTransportR(self.PFile))
	return self.Footer.Read(context.TODO(), protocol)
}

//Skip rows of parquet file
func (self *ParquetReader) SkipRows(num int64) error {
	var err error
	if num <= 0 {
		return nil
	}
	doneChan := make(chan int, self.NP)
	taskChan := make(chan string, len(self.SchemaHandler.ValueColumns))
	stopChan := make(chan int)

	for _, pathStr := range self.SchemaHandler.ValueColumns {
		if _, ok := self.ColumnBuffers[pathStr]; !ok {
			if self.ColumnBuffers[pathStr], err = NewColumnBuffer(self.PFile, self.Footer, self.SchemaHandler, pathStr); err != nil {
				return err
			}
		}
	}

	for i := int64(0); i < self.NP; i++ {
		go func() {
			for {
				select {
				case <-stopChan:
					return
				case pathStr := <-taskChan:
					cb := self.ColumnBuffers[pathStr]
					cb.SkipRows(int64(num))
					doneChan <- 0
				}
			}
		}()
	}

	for key, _ := range self.ColumnBuffers {
		taskChan <- key
	}

	for i := 0; i < len(self.ColumnBuffers); i++ {
		<-doneChan
	}
	for i := int64(0); i < self.NP; i++ {
		stopChan <- 0
	}
	return err
}

//Read rows of parquet file and unmarshal all to dst
func (self *ParquetReader) Read(dstInterface interface{}) error {
	return self.read(dstInterface, "")
}

// Read maxReadNumber objects
func (self *ParquetReader) ReadByNumber(maxReadNumber int) ([]interface{}, error) {
	var err error
	if self.ObjType == nil {
		if self.ObjType, err = self.SchemaHandler.GetType(self.SchemaHandler.GetRootInName()); err != nil {
			return nil, err
		}
	}

	vs := reflect.MakeSlice(reflect.SliceOf(self.ObjType), maxReadNumber, maxReadNumber)
	res := reflect.New(vs.Type())
	res.Elem().Set(vs)

	if err = self.Read(res.Interface()); err != nil {
		return nil, err
	}

	ln := res.Elem().Len()
	ret := make([]interface{}, ln)
	for i := 0; i < ln; i++ {
		ret[i] = res.Elem().Index(i).Interface()
	}

	return ret, nil
}

//Read rows of parquet file and unmarshal all to dst
func (self *ParquetReader) ReadPartial(dstInterface interface{}, prefixPath string) error {
	prefixPath, err := self.SchemaHandler.ConvertToInPathStr(prefixPath)
	if err != nil {
		return err
	}

	return self.read(dstInterface, prefixPath)
}

// Read maxReadNumber partial objects
func (self *ParquetReader) ReadPartialByNumber(maxReadNumber int, prefixPath string) ([]interface{}, error) {
	var err error
	if self.ObjPartialType == nil {
		if self.ObjPartialType, err = self.SchemaHandler.GetType(prefixPath); err != nil {
			return nil, err
		}
	}

	vs := reflect.MakeSlice(reflect.SliceOf(self.ObjPartialType), maxReadNumber, maxReadNumber)
	res := reflect.New(vs.Type())
	res.Elem().Set(vs)

	if err = self.ReadPartial(res.Interface(), prefixPath); err != nil {
		return nil, err
	}

	ln := res.Elem().Len()
	ret := make([]interface{}, ln)
	for i := 0; i < ln; i++ {
		ret[i] = res.Elem().Index(i).Interface()
	}

	return ret, nil
}

//Read rows of parquet file with a prefixPath
func (self *ParquetReader) read(dstInterface interface{}, prefixPath string) error {
	var err error
	tmap := make(map[string]*layout.Table)
	locker := new(sync.Mutex)
	ot := reflect.TypeOf(dstInterface).Elem().Elem()
	num := reflect.ValueOf(dstInterface).Elem().Len()
	if num <= 0 {
		return nil
	}

	doneChan := make(chan int, self.NP)
	taskChan := make(chan string, len(self.ColumnBuffers))
	stopChan := make(chan int)

	for i := int64(0); i < self.NP; i++ {
		go func() {
			for {
				select {
				case <-stopChan:
					return
				case pathStr := <-taskChan:
					cb := self.ColumnBuffers[pathStr]
					table, _ := cb.ReadRows(int64(num))
					locker.Lock()
					if _, ok := tmap[pathStr]; ok {
						tmap[pathStr].Merge(table)
					} else {
						tmap[pathStr] = layout.NewTableFromTable(table)
						tmap[pathStr].Merge(table)
					}
					locker.Unlock()
					doneChan <- 0
				}
			}
		}()
	}

	readNum := 0
	for key, _ := range self.ColumnBuffers {
		if strings.HasPrefix(key, prefixPath) {
			taskChan <- key
			readNum++
		}
	}
	for i := 0; i < readNum; i++ {
		<-doneChan
	}

	for i := int64(0); i < self.NP; i++ {
		stopChan <- 0
	}

	dstList := make([]interface{}, self.NP)
	delta := (int64(num) + self.NP - 1) / self.NP

	var wg sync.WaitGroup
	for c := int64(0); c < self.NP; c++ {
		bgn := c * delta
		end := bgn + delta
		if end > int64(num) {
			end = int64(num)
		}
		if bgn >= int64(num) {
			bgn, end = int64(num), int64(num)
		}
		wg.Add(1)
		go func(b, e, index int) {
			defer func() {
				wg.Done()
			}()

			dstList[index] = reflect.New(reflect.SliceOf(ot)).Interface()
			if err2 := marshal.Unmarshal(&tmap, b, e, dstList[index], self.SchemaHandler, prefixPath); err2 != nil {
				err = err2
			}
		}(int(bgn), int(end), int(c))
	}

	wg.Wait()

	resTmp := reflect.MakeSlice(reflect.SliceOf(ot), 0, num)
	for _, dst := range dstList {
		resTmp = reflect.AppendSlice(resTmp, reflect.ValueOf(dst).Elem())
	}
	reflect.ValueOf(dstInterface).Elem().Set(resTmp)
	return err
}

//Stop Read
func (self *ParquetReader) ReadStop() {
	for _, cb := range self.ColumnBuffers {
		if cb != nil {
			cb.PFile.Close()
		}
	}
}
<commit_msg>reuse dst list<commit_after>package reader

import (
	"context"
	"encoding/binary"
	"io"
	"reflect"
	"strings"
	"sync"

	"github.com/apache/thrift/lib/go/thrift"
	"github.com/xitongsys/parquet-go/common"
	"github.com/xitongsys/parquet-go/layout"
	"github.com/xitongsys/parquet-go/marshal"
	"github.com/xitongsys/parquet-go/parquet"
	"github.com/xitongsys/parquet-go/schema"
	"github.com/xitongsys/parquet-go/source"
)

type ParquetReader struct {
	SchemaHandler *schema.SchemaHandler
	NP            int64 //parallel number
	Footer        *parquet.FileMetaData
	PFile         source.ParquetFile

	ColumnBuffers map[string]*ColumnBufferType

	//One reader can only read one type objects
	ObjType        reflect.Type
	ObjPartialType reflect.Type
}

//Create a parquet reader: obj is a object with schema tags or a JSON schema string
func NewParquetReader(pFile source.ParquetFile, obj interface{}, np int64) (*ParquetReader, error) {
	var err error
	res := new(ParquetReader)
	res.NP = np
	res.PFile = pFile
	if err = res.ReadFooter(); err != nil {
		return nil, err
	}
	res.ColumnBuffers = make(map[string]*ColumnBufferType)

	if obj != nil {
		if sa, ok := obj.(string); ok {
			err = res.SetSchemaHandlerFromJSON(sa)
			return res, err

		} else if sa, ok := obj.([]*parquet.SchemaElement); ok {
			res.SchemaHandler = schema.NewSchemaHandlerFromSchemaList(sa)

		} else {
			if res.SchemaHandler, err = schema.NewSchemaHandlerFromStruct(obj); err != nil {
				return res, err
			}
		}

	} else {
		res.SchemaHandler = schema.NewSchemaHandlerFromSchemaList(res.Footer.Schema)
	}

	res.RenameSchema()
	for i := 0; i < len(res.SchemaHandler.SchemaElements); i++ {
		schema := res.SchemaHandler.SchemaElements[i]
		if schema.GetNumChildren() == 0 {
			pathStr := res.SchemaHandler.IndexMap[int32(i)]
			if res.ColumnBuffers[pathStr], err = NewColumnBuffer(pFile, res.Footer, res.SchemaHandler, pathStr); err != nil {
				return res, err
			}
		}
	}

	return res, nil
}

func (self *ParquetReader) SetSchemaHandlerFromJSON(jsonSchema string) error {
	var err error

	if self.SchemaHandler, err = schema.NewSchemaHandlerFromJSON(jsonSchema); err != nil {
		return err
	}

	self.RenameSchema()
	for i := 0; i < len(self.SchemaHandler.SchemaElements); i++ {
		schemaElement := self.SchemaHandler.SchemaElements[i]
		if schemaElement.GetNumChildren() == 0 {
			pathStr := self.SchemaHandler.IndexMap[int32(i)]
			if self.ColumnBuffers[pathStr], err = NewColumnBuffer(self.PFile, self.Footer, self.SchemaHandler, pathStr); err != nil {
				return err
			}
		}
	}
	return nil
}

//Rename schema name to inname
func (self *ParquetReader) RenameSchema() {
	for i := 0; i < len(self.SchemaHandler.Infos); i++ {
		self.Footer.Schema[i].Name = self.SchemaHandler.Infos[i].InName
	}
	for _, rowGroup := range self.Footer.RowGroups {
		for _, chunk := range rowGroup.Columns {
			exPath := make([]string, 0)
			exPath = append(exPath, self.SchemaHandler.GetRootExName())
			exPath = append(exPath, chunk.MetaData.GetPathInSchema()...)
			exPathStr := common.PathToStr(exPath)

			inPathStr := self.SchemaHandler.ExPathToInPath[exPathStr]
			inPath := common.StrToPath(inPathStr)[1:]
			chunk.MetaData.PathInSchema = inPath
		}
	}
}

func (self *ParquetReader) GetNumRows() int64 {
	return self.Footer.GetNumRows()
}

//Get the footer size
func (self *ParquetReader) GetFooterSize() (uint32, error) {
	var err error
	buf := make([]byte, 4)
	if _, err = self.PFile.Seek(-8, io.SeekEnd); err != nil {
		return 0, err
	}
	if _, err = io.ReadFull(self.PFile, buf); err != nil {
		return 0, err
	}
	size := binary.LittleEndian.Uint32(buf)
	return size, err
}

//Read footer from parquet file
func (self *ParquetReader) ReadFooter() error {
	size, err := self.GetFooterSize()
	if err != nil {
		return err
	}
	if _, err = self.PFile.Seek(-(int64)(8+size), io.SeekEnd); err != nil {
		return err
	}
	self.Footer = parquet.NewFileMetaData()
	pf := thrift.NewTCompactProtocolFactory()
	protocol := pf.GetProtocol(thrift.NewStreamTransportR(self.PFile))
	return self.Footer.Read(context.TODO(), protocol)
}

//Skip rows of parquet file
func (self *ParquetReader) SkipRows(num int64) error {
	var err error
	if num <= 0 {
		return nil
	}
	doneChan := make(chan int, self.NP)
	taskChan := make(chan string, len(self.SchemaHandler.ValueColumns))
	stopChan := make(chan int)

	for _, pathStr := range self.SchemaHandler.ValueColumns {
		if _, ok := self.ColumnBuffers[pathStr]; !ok {
			if self.ColumnBuffers[pathStr], err = NewColumnBuffer(self.PFile, self.Footer, self.SchemaHandler, pathStr); err != nil {
				return err
			}
		}
	}

	for i := int64(0); i < self.NP; i++ {
		go func() {
			for {
				select {
				case <-stopChan:
					return
				case pathStr := <-taskChan:
					cb := self.ColumnBuffers[pathStr]
					cb.SkipRows(int64(num))
					doneChan <- 0
				}
			}
		}()
	}

	for key, _ := range self.ColumnBuffers {
		taskChan <- key
	}

	for i := 0; i < len(self.ColumnBuffers); i++ {
		<-doneChan
	}
	for i := int64(0); i < self.NP; i++ {
		stopChan <- 0
	}
	return err
}

//Read rows of parquet file and unmarshal all to dst
func (self *ParquetReader) Read(dstInterface interface{}) error {
	return self.read(dstInterface, "")
}

// Read maxReadNumber objects
func (self *ParquetReader) ReadByNumber(maxReadNumber int) ([]interface{}, error) {
	var err error
	if self.ObjType == nil {
		if self.ObjType, err = self.SchemaHandler.GetType(self.SchemaHandler.GetRootInName()); err != nil {
			return nil, err
		}
	}

	vs := reflect.MakeSlice(reflect.SliceOf(self.ObjType), maxReadNumber, maxReadNumber)
	res := reflect.New(vs.Type())
	res.Elem().Set(vs)

	if err = self.Read(res.Interface()); err != nil {
		return nil, err
	}

	ln := res.Elem().Len()
	ret := make([]interface{}, ln)
	for i := 0; i < ln; i++ {
		ret[i] = res.Elem().Index(i).Interface()
	}

	return ret, nil
}

//Read rows of parquet file and unmarshal all to dst
func (self *ParquetReader) ReadPartial(dstInterface interface{}, prefixPath string) error {
	prefixPath, err := self.SchemaHandler.ConvertToInPathStr(prefixPath)
	if err != nil {
		return err
	}

	return self.read(dstInterface, prefixPath)
}

// Read maxReadNumber partial objects
func (self *ParquetReader) ReadPartialByNumber(maxReadNumber int, prefixPath string) ([]interface{}, error) {
	var err error
	if self.ObjPartialType == nil {
		if self.ObjPartialType, err = self.SchemaHandler.GetType(prefixPath); err != nil {
			return nil, err
		}
	}

	vs := reflect.MakeSlice(reflect.SliceOf(self.ObjPartialType), maxReadNumber, maxReadNumber)
	res := reflect.New(vs.Type())
	res.Elem().Set(vs)

	if err = self.ReadPartial(res.Interface(), prefixPath); err != nil {
		return nil, err
	}

	ln := res.Elem().Len()
	ret := make([]interface{}, ln)
	for i := 0; i < ln; i++ {
		ret[i] = res.Elem().Index(i).Interface()
	}

	return ret, nil
}

//Read rows of parquet file with a prefixPath
func (self *ParquetReader) read(dstInterface interface{}, prefixPath string) error {
	var err error
	tmap := make(map[string]*layout.Table)
	locker := new(sync.Mutex)
	ot := reflect.TypeOf(dstInterface).Elem().Elem()
	num := reflect.ValueOf(dstInterface).Elem().Len()
	if num <= 0 {
		return nil
	}

	doneChan := make(chan int, self.NP)
	taskChan := make(chan string, len(self.ColumnBuffers))
	stopChan := make(chan int)

	for i := int64(0); i < self.NP; i++ {
		go func() {
			for {
				select {
				case <-stopChan:
					return
				case pathStr := <-taskChan:
					cb := self.ColumnBuffers[pathStr]
					table, _ := cb.ReadRows(int64(num))
					locker.Lock()
					if _, ok := tmap[pathStr]; ok {
						tmap[pathStr].Merge(table)
					} else {
						tmap[pathStr] = layout.NewTableFromTable(table)
						tmap[pathStr].Merge(table)
					}
					locker.Unlock()
					doneChan <- 0
				}
			}
		}()
	}

	readNum := 0
	for key, _ := range self.ColumnBuffers {
		if strings.HasPrefix(key, prefixPath) {
			taskChan <- key
			readNum++
		}
	}
	for i := 0; i < readNum; i++ {
		<-doneChan
	}

	for i := int64(0); i < self.NP; i++ {
		stopChan <- 0
	}

	dstList := make([]interface{}, self.NP)
	delta := (int64(num) + self.NP - 1) / self.NP

	var wg sync.WaitGroup
	for c := int64(0); c < self.NP; c++ {
		bgn := c * delta
		end := bgn + delta
		if end > int64(num) {
			end = int64(num)
		}
		if bgn >= int64(num) {
			bgn, end = int64(num), int64(num)
		}
		wg.Add(1)
		go func(b, e, index int) {
			defer func() {
				wg.Done()
			}()

			dstList[index] = reflect.New(reflect.SliceOf(ot)).Interface()
			if err2 := marshal.Unmarshal(&tmap, b, e, dstList[index], self.SchemaHandler, prefixPath); err2 != nil {
				err = err2
			}
		}(int(bgn), int(end), int(c))
	}

	wg.Wait()

	dstValue := reflect.ValueOf(dstInterface).Elem()
	dstValue.SetLen(0)
	for _, dst := range dstList {
		dstValue.Set(reflect.AppendSlice(dstValue, reflect.ValueOf(dst).Elem()))
	}

	return err
}

//Stop Read
func (self *ParquetReader) ReadStop() {
	for _, cb := range self.ColumnBuffers {
		if cb != nil {
			cb.PFile.Close()
		}
	}
}
<|endoftext|>


<commit_before><commit_msg>Generated from da48a54cf0068154a6ac389f781c008887731b92<commit_after><|endoftext|>


<commit_before>package pgwp

import (
	"log"
	"testing"
	"time"

	"gopkg.in/ory-am/dockertest.v2"

	"database/sql"
	_ "github.com/lib/pq"
)

type Item struct {
	ID   string         `db:"id"`
	Name sql.NullString `db:"name"`
}

func TestConnect(t *testing.T) {
	// must fail: workers, queue len
	{
		expected := "Could not set up PostgreSQL container."
		var err error
		var x *Pool
		c, err := dockertest.ConnectToPostgreSQL(0, time.Second, func(url string) bool {
			// Check if postgres is responsive...
			x, err = Connect("postgres", url+"bad", 10, 10)
			return err == nil
		})
		if err == nil {
			log.Printf("expected [%s] actual [%s]", expected, err)
			t.Fail()
		}
		c.KillRemove()
	}
	// must fail: bad connection
	{
		expected := "Could not set up PostgreSQL container."
		var err error
		var x *Pool
		c, err := dockertest.ConnectToPostgreSQL(0, time.Second, func(url string) bool {
			// Check if postgres is responsive...
			x, err = Connect("postgres", url, -1, -1)
			return err == nil
		})
		if err == nil {
			log.Printf("expected [%s] actual [%s]", expected, err)
			t.Fail()
		}
		c.KillRemove()
	}
	// Get and Select
	{
		var err error
		var x *Pool
		c, err := dockertest.ConnectToPostgreSQL(2, time.Second, func(url string) bool {
			// Check if postgres is responsive...
			x, err = Connect("postgres", url, 10, 10)
			return err == nil
		})
		if err != nil {
			log.Fatalf("Could not connect to database: %s", err)
			t.Fail()
		}
		prepare(x)

		var one Item
		err = x.Get(&one, "SELECT * FROM mockdata WHERE name=$1", "two")
		if err != nil {
			log.Printf("Get err [%s]", err)
			t.Fail()
		}

		var list []*Item
		err = x.Select(&list, "SELECT * FROM mockdata LIMIT 2")
		if err != nil {
			log.Printf("Select err [%s]", err)
			t.Fail()
		}
		if len(list) != 2 {
			t.Fail()
		}
		log.Printf("Get result [%#v] Select result [%#v]", one, list)
		x.Close()
		c.KillRemove()
	}
}

// prepare generate mock data.
func prepare(pool *Pool) {
	qr := pool.MustExec(`
	CREATE EXTENSION IF NOT EXISTS "uuid-ossp";
	CREATE TABLE mockdata
	(
		id uuid NOT NULL DEFAULT uuid_generate_v1mc(),
		name text,
		CONSTRAINT mockdata_pkey PRIMARY KEY (id)
	)
	WITH (
		OIDS=FALSE
	);
	INSERT INTO mockdata (name) VALUES('one');
	INSERT INTO mockdata (name) VALUES('two');
	INSERT INTO mockdata (name) VALUES('three');
	`)
	_, err := qr.LastInsertId()
	if err != nil {
		log.Printf("prepare :  last insert id err [%s]", err)
	}
	_, err = qr.RowsAffected()
	if err != nil {
		log.Printf("prepare : rows affected err [%s]", err)
	}
}
<commit_msg>build #6<commit_after>package pgwp

import (
	"log"
	"testing"
	"time"

	"gopkg.in/ory-am/dockertest.v2"

	"database/sql"

	_ "github.com/lib/pq"
)

// Item struct.
type Item struct {
	ID   string         `db:"id"`
	Name sql.NullString `db:"name"`
}

func TestConnect(t *testing.T) {
	log.SetFlags(log.Lshortfile)
	// must fail: workers, queue len
	{
		expected := "Could not set up PostgreSQL container."
		var err error
		var x *Pool
		c, err := dockertest.ConnectToPostgreSQL(0, time.Second, func(url string) bool {
			// Check if postgres is responsive...
			x, err = Connect("postgres", url+"bad", 10, 10)
			return err == nil
		})
		if err == nil {
			log.Printf("expected [%s] actual [%s]", expected, err)
			t.Fail()
		}
		c.KillRemove()
	}
	// must fail: bad connection
	{
		expected := "Could not set up PostgreSQL container."
		var err error
		var x *Pool
		c, err := dockertest.ConnectToPostgreSQL(0, time.Second, func(url string) bool {
			// Check if postgres is responsive...
			x, err = Connect("postgres", url, -1, -1)
			return err == nil
		})
		if err == nil {
			log.Printf("expected [%s] actual [%s]", expected, err)
			t.Fail()
		}
		c.KillRemove()
	}
	// Get and Select
	{
		var err error
		var x *Pool
		c, err := dockertest.ConnectToPostgreSQL(2, time.Second, func(url string) bool {
			// Check if postgres is responsive...
			x, err = Connect("postgres", url, 10, 10)
			return err == nil
		})
		if err != nil {
			log.Fatalf("Could not connect to database: %s", err)
			t.Fail()
		}
		prepare(x)

		var one Item
		err = x.Get(&one, "SELECT * FROM mockdata WHERE name=$1", "two")
		if err != nil {
			log.Printf("Get err [%s]", err)
			t.Fail()
		}

		var list []*Item
		err = x.Select(&list, "SELECT * FROM mockdata LIMIT 2")
		if err != nil {
			log.Printf("Select err [%s]", err)
			t.Fail()
		}
		if len(list) != 2 {
			t.Fail()
		}
		log.Printf("Get result [%#v] Select result [%#v]", one, list)
		x.Close()
		c.KillRemove()
	}
}

// prepare generate mock data.
func prepare(pool *Pool) {
	qr := pool.MustExec(`
	CREATE EXTENSION IF NOT EXISTS "uuid-ossp";
	CREATE TABLE mockdata
	(
		id uuid NOT NULL DEFAULT uuid_generate_v1mc(),
		name text,
		CONSTRAINT mockdata_pkey PRIMARY KEY (id)
	)
	WITH (
		OIDS=FALSE
	);
	INSERT INTO mockdata (name) VALUES('one');
	INSERT INTO mockdata (name) VALUES('two');
	INSERT INTO mockdata (name) VALUES('three');
	`)
	_, err := qr.LastInsertId()
	if err != nil {
		log.Printf("prepare :  last insert id err [%s]", err)
	}
	_, err = qr.RowsAffected()
	if err != nil {
		log.Printf("prepare : rows affected err [%s]", err)
	}
}
<|endoftext|>


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

package bootstrap

import (
	"crypto/tls"
	"fmt"
	"time"

	"istio.io/istio/pilot/pkg/features"
	kubecontroller "istio.io/istio/pilot/pkg/serviceregistry/kube/controller"
	"istio.io/istio/pkg/config/constants"
	"istio.io/istio/pkg/keepalive"
	"istio.io/pkg/ctrlz"
	"istio.io/pkg/env"
)

// RegistryOptions provide configuration options for the configuration controller. If FileDir is set, that directory will
// be monitored for CRD yaml files and will update the controller as those files change (This is used for testing
// purposes). Otherwise, a CRD client is created based on the configuration.
type RegistryOptions struct {
	// If FileDir is set, the below kubernetes options are ignored
	FileDir string

	Registries []string

	// Kubernetes controller options
	KubeOptions kubecontroller.Options
	// ClusterRegistriesNamespace specifies where the multi-cluster secret resides
	ClusterRegistriesNamespace string
	KubeConfig                 string

	// DistributionTracking control
	DistributionCacheRetention time.Duration

	// DistributionTracking control
	DistributionTrackingEnabled bool
}

// PilotArgs provides all of the configuration parameters for the Pilot discovery service.
type PilotArgs struct {
	ServerOptions      DiscoveryServerOptions
	InjectionOptions   InjectionOptions
	PodName            string
	Namespace          string
	Revision           string
	MeshConfigFile     string
	NetworksConfigFile string
	RegistryOptions    RegistryOptions
	CtrlZOptions       *ctrlz.Options
	Plugins            []string
	KeepaliveOptions   *keepalive.Options
	ShutdownDuration   time.Duration
	JwtRule            string
}

// DiscoveryServerOptions contains options for create a new discovery server instance.
type DiscoveryServerOptions struct {
	// The listening address for HTTP (debug). If the port in the address is empty or "0" (as in "127.0.0.1:" or "[::1]:0")
	// a port number is automatically chosen.
	HTTPAddr string

	// The listening address for HTTPS (webhooks). If the port in the address is empty or "0" (as in "127.0.0.1:" or "[::1]:0")
	// a port number is automatically chosen.
	// If the address is empty, the secure port is disabled, and the
	// webhooks are registered on the HTTP port - a gateway in front will
	// terminate TLS instead.
	HTTPSAddr string

	// The listening address for gRPC. If the port in the address is empty or "0" (as in "127.0.0.1:" or "[::1]:0")
	// a port number is automatically chosen.
	GRPCAddr string

	// The listening address for the monitoring port. If the port in the address is empty or "0" (as in "127.0.0.1:" or "[::1]:0")
	// a port number is automatically chosen.
	MonitoringAddr string

	EnableProfiling bool

	// Optional TLS configuration
	TLSOptions TLSOptions

	// The listening address for secured gRPC. If the port in the address is empty or "0" (as in "127.0.0.1:" or "[::1]:0")
	// a port number is automatically chosen.
	SecureGRPCAddr string
}

type InjectionOptions struct {
	// Directory of injection related config files.
	InjectionDirectory string
}

// TLSOptions is optional TLS parameters for Istiod server.
type TLSOptions struct {
	CaCertFile      string
	CertFile        string
	KeyFile         string
	TLSCipherSuites []string
	CipherSuits     []uint16 // This is the parsed cipher suites
}

var (
	PodNamespace = env.RegisterStringVar("POD_NAMESPACE", constants.IstioSystemNamespace, "").Get()
	PodName      = env.RegisterStringVar("POD_NAME", "", "").Get()
	JwtRule      = env.RegisterStringVar("JWT_RULE", "",
		"The JWT rule used by istiod authentication").Get()
)

// Revision is the value of the Istio control plane revision, e.g. "canary",
// and is the value used by the "istio.io/rev" label.
var Revision = env.RegisterStringVar("REVISION", "", "").Get()

// NewPilotArgs constructs pilotArgs with default values.
func NewPilotArgs(initFuncs ...func(*PilotArgs)) *PilotArgs {
	p := &PilotArgs{}

	// Apply Default Values.
	p.applyDefaults()

	// Apply custom initialization functions.
	for _, fn := range initFuncs {
		fn(p)
	}

	// Set the ClusterRegistries namespace based on the selected namespace.
	if p.Namespace != "" {
		p.RegistryOptions.ClusterRegistriesNamespace = p.Namespace
	} else {
		p.RegistryOptions.ClusterRegistriesNamespace = constants.IstioSystemNamespace
	}

	return p
}

// Apply default value to PilotArgs
func (p *PilotArgs) applyDefaults() {
	p.Namespace = PodNamespace
	p.PodName = PodName
	p.Revision = Revision
	p.JwtRule = JwtRule
	p.KeepaliveOptions = keepalive.DefaultOption()
	p.RegistryOptions.DistributionTrackingEnabled = features.EnableDistributionTracking
	p.RegistryOptions.DistributionCacheRetention = features.DistributionHistoryRetention
}

func (p *PilotArgs) Complete() error {
	cipherSuits, err := TLSCipherSuites(p.ServerOptions.TLSOptions.TLSCipherSuites)
	if err != nil {
		return err
	}
	p.ServerOptions.TLSOptions.CipherSuits = cipherSuits
	return nil
}

func allCiphers() map[string]uint16 {
	acceptedCiphers := make(map[string]uint16, len(tls.CipherSuites())+len(tls.InsecureCipherSuites()))
	for _, cipher := range tls.InsecureCipherSuites() {
		acceptedCiphers[cipher.Name] = cipher.ID
	}
	for _, cipher := range tls.CipherSuites() {
		acceptedCiphers[cipher.Name] = cipher.ID
	}
	return acceptedCiphers
}

// TLSCipherSuites returns a list of cipher suite IDs from the cipher suite names passed.
func TLSCipherSuites(cipherNames []string) ([]uint16, error) {
	if len(cipherNames) == 0 {
		return nil, nil
	}
	ciphersIntSlice := make([]uint16, 0)
	possibleCiphers := allCiphers()
	for _, cipher := range cipherNames {
		intValue, ok := possibleCiphers[cipher]
		if !ok {
			return nil, fmt.Errorf("cipher suite %s not supported or doesn't exist", cipher)
		}
		ciphersIntSlice = append(ciphersIntSlice, intValue)
	}
	return ciphersIntSlice, nil
}
<commit_msg>optimize codes, the logic that apply pilot default args should be placed together (#35913)<commit_after>// Copyright Istio Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package bootstrap

import (
	"crypto/tls"
	"fmt"
	"time"

	"istio.io/istio/pilot/pkg/features"
	kubecontroller "istio.io/istio/pilot/pkg/serviceregistry/kube/controller"
	"istio.io/istio/pkg/config/constants"
	"istio.io/istio/pkg/keepalive"
	"istio.io/pkg/ctrlz"
	"istio.io/pkg/env"
)

// RegistryOptions provide configuration options for the configuration controller. If FileDir is set, that directory will
// be monitored for CRD yaml files and will update the controller as those files change (This is used for testing
// purposes). Otherwise, a CRD client is created based on the configuration.
type RegistryOptions struct {
	// If FileDir is set, the below kubernetes options are ignored
	FileDir string

	Registries []string

	// Kubernetes controller options
	KubeOptions kubecontroller.Options
	// ClusterRegistriesNamespace specifies where the multi-cluster secret resides
	ClusterRegistriesNamespace string
	KubeConfig                 string

	// DistributionTracking control
	DistributionCacheRetention time.Duration

	// DistributionTracking control
	DistributionTrackingEnabled bool
}

// PilotArgs provides all of the configuration parameters for the Pilot discovery service.
type PilotArgs struct {
	ServerOptions      DiscoveryServerOptions
	InjectionOptions   InjectionOptions
	PodName            string
	Namespace          string
	Revision           string
	MeshConfigFile     string
	NetworksConfigFile string
	RegistryOptions    RegistryOptions
	CtrlZOptions       *ctrlz.Options
	Plugins            []string
	KeepaliveOptions   *keepalive.Options
	ShutdownDuration   time.Duration
	JwtRule            string
}

// DiscoveryServerOptions contains options for create a new discovery server instance.
type DiscoveryServerOptions struct {
	// The listening address for HTTP (debug). If the port in the address is empty or "0" (as in "127.0.0.1:" or "[::1]:0")
	// a port number is automatically chosen.
	HTTPAddr string

	// The listening address for HTTPS (webhooks). If the port in the address is empty or "0" (as in "127.0.0.1:" or "[::1]:0")
	// a port number is automatically chosen.
	// If the address is empty, the secure port is disabled, and the
	// webhooks are registered on the HTTP port - a gateway in front will
	// terminate TLS instead.
	HTTPSAddr string

	// The listening address for gRPC. If the port in the address is empty or "0" (as in "127.0.0.1:" or "[::1]:0")
	// a port number is automatically chosen.
	GRPCAddr string

	// The listening address for the monitoring port. If the port in the address is empty or "0" (as in "127.0.0.1:" or "[::1]:0")
	// a port number is automatically chosen.
	MonitoringAddr string

	EnableProfiling bool

	// Optional TLS configuration
	TLSOptions TLSOptions

	// The listening address for secured gRPC. If the port in the address is empty or "0" (as in "127.0.0.1:" or "[::1]:0")
	// a port number is automatically chosen.
	SecureGRPCAddr string
}

type InjectionOptions struct {
	// Directory of injection related config files.
	InjectionDirectory string
}

// TLSOptions is optional TLS parameters for Istiod server.
type TLSOptions struct {
	CaCertFile      string
	CertFile        string
	KeyFile         string
	TLSCipherSuites []string
	CipherSuits     []uint16 // This is the parsed cipher suites
}

var (
	PodNamespace = env.RegisterStringVar("POD_NAMESPACE", constants.IstioSystemNamespace, "").Get()
	PodName      = env.RegisterStringVar("POD_NAME", "", "").Get()
	JwtRule      = env.RegisterStringVar("JWT_RULE", "",
		"The JWT rule used by istiod authentication").Get()
)

// Revision is the value of the Istio control plane revision, e.g. "canary",
// and is the value used by the "istio.io/rev" label.
var Revision = env.RegisterStringVar("REVISION", "", "").Get()

// NewPilotArgs constructs pilotArgs with default values.
func NewPilotArgs(initFuncs ...func(*PilotArgs)) *PilotArgs {
	p := &PilotArgs{}

	// Apply Default Values.
	p.applyDefaults()

	// Apply custom initialization functions.
	for _, fn := range initFuncs {
		fn(p)
	}

	return p
}

// Apply default value to PilotArgs
func (p *PilotArgs) applyDefaults() {
	p.Namespace = PodNamespace
	p.PodName = PodName
	p.Revision = Revision
	p.JwtRule = JwtRule
	p.KeepaliveOptions = keepalive.DefaultOption()
	p.RegistryOptions.DistributionTrackingEnabled = features.EnableDistributionTracking
	p.RegistryOptions.DistributionCacheRetention = features.DistributionHistoryRetention
	p.RegistryOptions.ClusterRegistriesNamespace = p.Namespace
}

func (p *PilotArgs) Complete() error {
	cipherSuits, err := TLSCipherSuites(p.ServerOptions.TLSOptions.TLSCipherSuites)
	if err != nil {
		return err
	}
	p.ServerOptions.TLSOptions.CipherSuits = cipherSuits
	return nil
}

func allCiphers() map[string]uint16 {
	acceptedCiphers := make(map[string]uint16, len(tls.CipherSuites())+len(tls.InsecureCipherSuites()))
	for _, cipher := range tls.InsecureCipherSuites() {
		acceptedCiphers[cipher.Name] = cipher.ID
	}
	for _, cipher := range tls.CipherSuites() {
		acceptedCiphers[cipher.Name] = cipher.ID
	}
	return acceptedCiphers
}

// TLSCipherSuites returns a list of cipher suite IDs from the cipher suite names passed.
func TLSCipherSuites(cipherNames []string) ([]uint16, error) {
	if len(cipherNames) == 0 {
		return nil, nil
	}
	ciphersIntSlice := make([]uint16, 0)
	possibleCiphers := allCiphers()
	for _, cipher := range cipherNames {
		intValue, ok := possibleCiphers[cipher]
		if !ok {
			return nil, fmt.Errorf("cipher suite %s not supported or doesn't exist", cipher)
		}
		ciphersIntSlice = append(ciphersIntSlice, intValue)
	}
	return ciphersIntSlice, nil
}
<|endoftext|>


<commit_before>package backend

import (
	"bytes"
	"log"
	"sync"

	"github.com/coreos/etcd/Godeps/_workspace/src/github.com/boltdb/bolt"
)

type BatchTx interface {
	Lock()
	Unlock()
	UnsafeCreateBucket(name []byte)
	UnsafePut(bucketName []byte, key []byte, value []byte)
	UnsafeRange(bucketName []byte, key, endKey []byte, limit int64) (keys [][]byte, vals [][]byte)
	UnsafeDelete(bucketName []byte, key []byte)
	Commit()
	CommitAndStop()
}

type batchTx struct {
	sync.Mutex
	tx      *bolt.Tx
	backend *backend
	pending int
}

func (t *batchTx) UnsafeCreateBucket(name []byte) {
	_, err := t.tx.CreateBucket(name)
	if err != nil && err != bolt.ErrBucketExists {
		log.Fatalf("storage: cannot create bucket %s (%v)", string(name), err)
	}
}

// before calling unsafePut, the caller MUST hold the lock on tx.
func (t *batchTx) UnsafePut(bucketName []byte, key []byte, value []byte) {
	bucket := t.tx.Bucket(bucketName)
	if bucket == nil {
		log.Fatalf("storage: bucket %s does not exist", string(bucketName))
	}
	if err := bucket.Put(key, value); err != nil {
		log.Fatalf("storage: cannot put key into bucket (%v)", err)
	}
	t.pending++
	if t.pending > t.backend.batchLimit {
		t.commit(false)
		t.pending = 0
	}
}

// before calling unsafeRange, the caller MUST hold the lock on tx.
func (t *batchTx) UnsafeRange(bucketName []byte, key, endKey []byte, limit int64) (keys [][]byte, vs [][]byte) {
	bucket := t.tx.Bucket(bucketName)
	if bucket == nil {
		log.Fatalf("storage: bucket %s does not exist", string(bucketName))
	}

	if len(endKey) == 0 {
		if v := bucket.Get(key); v == nil {
			return keys, vs
		} else {
			return append(keys, key), append(vs, v)
		}
	}

	c := bucket.Cursor()
	for ck, cv := c.Seek(key); ck != nil && bytes.Compare(ck, endKey) < 0; ck, cv = c.Next() {
		vs = append(vs, cv)
		keys = append(keys, ck)
	}

	return keys, vs
}

// before calling unsafeDelete, the caller MUST hold the lock on tx.
func (t *batchTx) UnsafeDelete(bucketName []byte, key []byte) {
	bucket := t.tx.Bucket(bucketName)
	if bucket == nil {
		log.Fatalf("storage: bucket %s does not exist", string(bucketName))
	}
	err := bucket.Delete(key)
	if err != nil {
		log.Fatalf("storage: cannot delete key from bucket (%v)", err)
	}
	t.pending++
	if t.pending > t.backend.batchLimit {
		t.commit(false)
		t.pending = 0
	}
}

// Commit commits a previous tx and begins a new writable one.
func (t *batchTx) Commit() {
	t.Lock()
	defer t.Unlock()
	t.commit(false)
}

// CommitAndStop commits the previous tx and do not create a new one.
func (t *batchTx) CommitAndStop() {
	t.Lock()
	defer t.Unlock()
	t.commit(true)
}

func (t *batchTx) commit(stop bool) {
	var err error
	// commit the last tx
	if t.tx != nil {
		err = t.tx.Commit()
		if err != nil {
			log.Fatalf("storage: cannot commit tx (%s)", err)
		}
	}

	if stop {
		return
	}

	// begin a new tx
	t.tx, err = t.backend.db.Begin(true)
	if err != nil {
		log.Fatalf("storage: cannot begin tx (%s)", err)
	}
}
<commit_msg>storage/backend: fix limit doesn't effect in range<commit_after>package backend

import (
	"bytes"
	"log"
	"sync"

	"github.com/coreos/etcd/Godeps/_workspace/src/github.com/boltdb/bolt"
)

type BatchTx interface {
	Lock()
	Unlock()
	UnsafeCreateBucket(name []byte)
	UnsafePut(bucketName []byte, key []byte, value []byte)
	UnsafeRange(bucketName []byte, key, endKey []byte, limit int64) (keys [][]byte, vals [][]byte)
	UnsafeDelete(bucketName []byte, key []byte)
	Commit()
	CommitAndStop()
}

type batchTx struct {
	sync.Mutex
	tx      *bolt.Tx
	backend *backend
	pending int
}

func (t *batchTx) UnsafeCreateBucket(name []byte) {
	_, err := t.tx.CreateBucket(name)
	if err != nil && err != bolt.ErrBucketExists {
		log.Fatalf("storage: cannot create bucket %s (%v)", string(name), err)
	}
}

// before calling unsafePut, the caller MUST hold the lock on tx.
func (t *batchTx) UnsafePut(bucketName []byte, key []byte, value []byte) {
	bucket := t.tx.Bucket(bucketName)
	if bucket == nil {
		log.Fatalf("storage: bucket %s does not exist", string(bucketName))
	}
	if err := bucket.Put(key, value); err != nil {
		log.Fatalf("storage: cannot put key into bucket (%v)", err)
	}
	t.pending++
	if t.pending > t.backend.batchLimit {
		t.commit(false)
		t.pending = 0
	}
}

// before calling unsafeRange, the caller MUST hold the lock on tx.
func (t *batchTx) UnsafeRange(bucketName []byte, key, endKey []byte, limit int64) (keys [][]byte, vs [][]byte) {
	bucket := t.tx.Bucket(bucketName)
	if bucket == nil {
		log.Fatalf("storage: bucket %s does not exist", string(bucketName))
	}

	if len(endKey) == 0 {
		if v := bucket.Get(key); v == nil {
			return keys, vs
		} else {
			return append(keys, key), append(vs, v)
		}
	}

	c := bucket.Cursor()
	for ck, cv := c.Seek(key); ck != nil && bytes.Compare(ck, endKey) < 0; ck, cv = c.Next() {
		vs = append(vs, cv)
		keys = append(keys, ck)
		if limit > 0 && limit == int64(len(keys)) {
			break
		}
	}

	return keys, vs
}

// before calling unsafeDelete, the caller MUST hold the lock on tx.
func (t *batchTx) UnsafeDelete(bucketName []byte, key []byte) {
	bucket := t.tx.Bucket(bucketName)
	if bucket == nil {
		log.Fatalf("storage: bucket %s does not exist", string(bucketName))
	}
	err := bucket.Delete(key)
	if err != nil {
		log.Fatalf("storage: cannot delete key from bucket (%v)", err)
	}
	t.pending++
	if t.pending > t.backend.batchLimit {
		t.commit(false)
		t.pending = 0
	}
}

// Commit commits a previous tx and begins a new writable one.
func (t *batchTx) Commit() {
	t.Lock()
	defer t.Unlock()
	t.commit(false)
}

// CommitAndStop commits the previous tx and do not create a new one.
func (t *batchTx) CommitAndStop() {
	t.Lock()
	defer t.Unlock()
	t.commit(true)
}

func (t *batchTx) commit(stop bool) {
	var err error
	// commit the last tx
	if t.tx != nil {
		err = t.tx.Commit()
		if err != nil {
			log.Fatalf("storage: cannot commit tx (%s)", err)
		}
	}

	if stop {
		return
	}

	// begin a new tx
	t.tx, err = t.backend.db.Begin(true)
	if err != nil {
		log.Fatalf("storage: cannot begin tx (%s)", err)
	}
}
<|endoftext|>


<commit_before>package core

import (
	"context"
	"encoding/json"
	"fmt"
	"time"

	"chain/core/query"
	"chain/core/query/filter"
	"chain/errors"
	"chain/net/http/httpjson"
	"chain/protocol/bc"
)

// These types enforce the ordering of JSON fields in API output.
type (
	txinResp struct {
		Action          interface{} `json:"action"`
		AssetID         interface{} `json:"asset_id"`
		AssetAlias      interface{} `json:"asset_alias"`
		AssetTags       interface{} `json:"asset_tags"`
		Amount          interface{} `json:"amount"`
		IssuanceProgram interface{} `json:"issuance_program"`
		SpentOutput     interface{} `json:"spent_output"`
		AccountID       interface{} `json:"account_id"`
		AccountAlias    interface{} `json:"account_alias"`
		AccountTags     interface{} `json:"account_tags"`
		ReferenceData   interface{} `json:"reference_data"`
		AssetDefinition interface{} `json:"asset_definition"`
	}
	txoutResp struct {
		Action         interface{} `json:"action"`
		Position       interface{} `json:"position"`
		AssetID        interface{} `json:"asset_id"`
		AssetAlias     interface{} `json:"asset_alias"`
		AssetTags      interface{} `json:"asset_tags"`
		Amount         interface{} `json:"asset_amount"`
		AccountID      interface{} `json:"account_id"`
		AccountAlias   interface{} `json:"account_alias"`
		AccountTags    interface{} `json:"account_tags"`
		ControlProgram interface{} `json:"control_program"`
		ReferenceData  interface{} `json:"reference_data"`
	}
	txResp struct {
		ID            interface{} `json:"id"`
		Timestamp     interface{} `json:"timestamp"`
		BlockID       interface{} `json:"block_id"`
		BlockHeight   interface{} `json:"block_height"`
		Position      interface{} `json:"position"`
		ReferenceData interface{} `json:"reference_data"`
		Inputs        interface{} `json:"inputs"`
		Outputs       interface{} `json:"outputs"`
	}
)

// listTransactions is an http handler for listing transactions matching
// an index or an ad-hoc filter.
//
// POST /list-transactions
func (a *api) listTransactions(ctx context.Context, in requestQuery) (result page, err error) {
	if in.EndTimeMS == 0 {
		in.EndTimeMS = bc.Millis(time.Now())
	}

	var (
		p     filter.Predicate
		after query.TxAfter
	)

	// Build the filter predicate.
	p, err = filter.Parse(in.Filter)
	if err != nil {
		return result, err
	}

	// Either parse the provided `after` or look one up for the time range.
	if in.After != "" {
		after, err = query.DecodeTxAfter(in.After)
		if err != nil {
			return result, errors.Wrap(err, "decoding `after`")
		}
	} else {
		after, err = a.indexer.LookupTxAfter(ctx, in.StartTimeMS, in.EndTimeMS)
		if err != nil {
			return result, err
		}
	}

	var asc bool
	if in.Order == "asc" {
		asc = true
	}

	limit := defGenericPageSize
	txns, nextAfter, err := a.indexer.Transactions(ctx, p, in.FilterParams, after, limit, asc)
	if err != nil {
		return result, errors.Wrap(err, "running tx query")
	}

	resp := make([]*txResp, 0, len(txns))
	for _, t := range txns {
		tjson, ok := t.(*json.RawMessage)
		if !ok {
			return result, fmt.Errorf("unexpected type %T in Indexer.Transactions output", t)
		}
		if tjson == nil {
			return result, fmt.Errorf("unexpected nil in Indexer.Transactions output")
		}
		var tx map[string]interface{}
		err = json.Unmarshal(*tjson, &tx)
		if err != nil {
			return result, errors.Wrap(err, "decoding Indexer.Transactions output")
		}
		inputs, ok := tx["inputs"].([]map[string]interface{})
		if !ok {
			return result, fmt.Errorf("unexpected type %T for inputs in Indexer.Transactions output", tx["inputs"])
		}
		outputs, ok := tx["outputs"].([]map[string]interface{})
		if !ok {
			return result, fmt.Errorf("unexpected type %T for outputs in Indexer.Transactions output", tx["outputs"])
		}
		inResps := make([]*txinResp, 0, len(inputs))
		for _, in := range inputs {
			r := &txinResp{
				Action:          in["action"],
				AssetID:         in["asset_id"],
				AssetAlias:      in["asset_alias"],
				AssetTags:       in["asset_tags"],
				Amount:          in["amount"],
				IssuanceProgram: in["issuance_program"],
				SpentOutput:     in["spent_output"],
				AccountID:       in["account_id"],
				AccountAlias:    in["account_alias"],
				AccountTags:     in["account_tags"],
				ReferenceData:   in["reference_data"],
				AssetDefinition: in["asset_definition"],
			}
			inResps = append(inResps, r)
		}
		outResps := make([]*txoutResp, 0, len(outputs))
		for _, out := range outputs {
			r := &txoutResp{
				Action:         out["action"],
				Position:       out["position"],
				AssetID:        out["asset_id"],
				AssetAlias:     out["asset_alias"],
				AssetTags:      out["asset_tags"],
				Amount:         out["asset_amount"],
				AccountID:      out["account_id"],
				AccountAlias:   out["account_alias"],
				AccountTags:    out["account_tags"],
				ControlProgram: out["control_program"],
				ReferenceData:  out["reference_data"],
			}
			outResps = append(outResps, r)
		}
		r := &txResp{
			ID:            tx["id"],
			Timestamp:     tx["timestamp"],
			BlockID:       tx["block_id"],
			BlockHeight:   tx["block_height"],
			Position:      tx["position"],
			ReferenceData: tx["reference_data"],
			Inputs:        inResps,
			Outputs:       outResps,
		}
		resp = append(resp, r)
	}

	out := in
	out.After = nextAfter.String()
	return page{
		Items:    httpjson.Array(resp),
		LastPage: len(resp) < limit,
		Next:     out,
	}, nil
}

// listAccounts is an http handler for listing accounts matching
// an index or an ad-hoc filter.
//
// POST /list-accounts
func (a *api) listAccounts(ctx context.Context, in requestQuery) (page, error) {
	limit := defGenericPageSize

	// Build the filter predicate.
	p, err := filter.Parse(in.Filter)
	if err != nil {
		return page{}, errors.Wrap(err, "parsing acc query")
	}
	after := in.After

	// Use the filter engine for querying account tags.
	accounts, after, err := a.indexer.Accounts(ctx, p, in.FilterParams, after, limit)
	if err != nil {
		return page{}, errors.Wrap(err, "running acc query")
	}

	result := make([]*accountResponse, 0, len(accounts))
	for _, a := range accounts {
		r := &accountResponse{
			ID:     a["id"],
			Alias:  a["alias"],
			XPubs:  a["xpubs"],
			Quorum: a["quorum"],
			Tags:   a["tags"],
		}
		result = append(result, r)
	}

	// Pull in the accounts by the IDs
	out := in
	out.After = after
	return page{
		Items:    httpjson.Array(result),
		LastPage: len(result) < limit,
		Next:     out,
	}, nil
}

// POST /list-balances
func (a *api) listBalances(ctx context.Context, in requestQuery) (result page, err error) {
	if in.TimestampMS == 0 {
		in.TimestampMS = bc.Millis(time.Now())
	}

	var p filter.Predicate
	var sumBy []filter.Field
	p, err = filter.Parse(in.Filter)
	if err != nil {
		return result, err
	}

	for _, field := range in.SumBy {
		f, err := filter.ParseField(field)
		if err != nil {
			return result, err
		}
		sumBy = append(sumBy, f)
	}

	// TODO(jackson): paginate this endpoint.
	balances, err := a.indexer.Balances(ctx, p, in.FilterParams, sumBy, in.TimestampMS)
	if err != nil {
		return result, err
	}

	result.Items = httpjson.Array(balances)
	result.LastPage = true
	result.Next = in
	return result, nil
}

// This type enforces the ordering of JSON fields in API output.
type utxoResp struct {
	Action         interface{} `json:"action"`
	TransactionID  interface{} `json:"transaction_id"`
	Position       interface{} `json:"position"`
	AssetID        interface{} `json:"asset_id"`
	AssetAlias     interface{} `json:"asset_alias"`
	AssetTags      interface{} `json:"asset_tags"`
	Amount         interface{} `json:"asset_amount"`
	AccountID      interface{} `json:"account_id"`
	AccountAlias   interface{} `json:"account_alias"`
	AccountTags    interface{} `json:"account_tags"`
	ControlProgram interface{} `json:"control_program"`
	ReferenceData  interface{} `json:"reference_data"`
}

// POST /list-unspent-outputs
func (a *api) listUnspentOutputs(ctx context.Context, in requestQuery) (result page, err error) {
	if in.TimestampMS == 0 {
		in.TimestampMS = bc.Millis(time.Now())
	}
	var p filter.Predicate
	p, err = filter.Parse(in.Filter)
	if err != nil {
		return result, err
	}

	var after *query.OutputsAfter
	if in.After != "" {
		after, err = query.DecodeOutputsAfter(in.After)
		if err != nil {
			return result, errors.Wrap(err, "decoding `after`")
		}
	}

	limit := defGenericPageSize
	outputs, nextAfter, err := a.indexer.Outputs(ctx, p, in.FilterParams, in.TimestampMS, after, limit)
	if err != nil {
		return result, errors.Wrap(err, "querying outputs")
	}

	resp := make([]*utxoResp, 0, len(outputs))
	for _, o := range outputs {
		ojson, ok := o.(*json.RawMessage)
		if !ok {
			return result, fmt.Errorf("unexpected type %T in Indexer.Outputs output", o)
		}
		if ojson == nil {
			return result, fmt.Errorf("unexpected nil in Indexer.Outputs output")
		}
		var out map[string]interface{}
		err = json.Unmarshal(*ojson, &out)
		if err != nil {
			return result, errors.Wrap(err, "decoding Indexer.Outputs output")
		}
		r := &utxoResp{
			Action:         out["action"],
			TransactionID:  out["transaction_id"],
			Position:       out["position"],
			AssetID:        out["asset_id"],
			AssetAlias:     out["asset_alias"],
			AssetTags:      out["asset_tags"],
			Amount:         out["asset_amount"],
			AccountID:      out["account_id"],
			AccountAlias:   out["account_alias"],
			AccountTags:    out["account_tags"],
			ControlProgram: out["control_program"],
			ReferenceData:  out["reference_data"],
		}
		resp = append(resp, r)
	}

	outQuery := in
	outQuery.After = nextAfter.String()
	return page{
		Items:    resp,
		LastPage: len(resp) < limit,
		Next:     outQuery,
	}, nil
}

// listAssets is an http handler for listing assets matching
// an index or an ad-hoc filter.
//
// POST /list-assets
func (a *api) listAssets(ctx context.Context, in requestQuery) (page, error) {
	limit := defGenericPageSize

	// Build the filter predicate.
	p, err := filter.Parse(in.Filter)
	if err != nil {
		return page{}, err
	}
	after := in.After

	// Use the query engine for querying asset tags.
	var assets []map[string]interface{}
	assets, after, err = a.indexer.Assets(ctx, p, in.FilterParams, after, limit)
	if err != nil {
		return page{}, errors.Wrap(err, "running asset query")
	}

	result := make([]*assetResponse, 0, len(assets))
	for _, a := range assets {
		r := &assetResponse{
			ID:              a["id"],
			IssuanceProgram: a["issuance_program"],
			XPubs:           a["xpubs"],
			Quorum:          a["quorum"],
			Definition:      a["definition"],
			Tags:            a["tags"],
			Origin:          a["origin"],
		}
		if alias, ok := a["alias"].(string); ok && alias != "" {
			r.Alias = &alias
		}
		result = append(result, r)
	}

	out := in
	out.After = after
	return page{
		Items:    httpjson.Array(result),
		LastPage: len(result) < limit,
		Next:     out,
	}, nil
}
<commit_msg>core: fix type assertions in /list-transactions<commit_after>package core

import (
	"context"
	"encoding/json"
	"fmt"
	"time"

	"chain/core/query"
	"chain/core/query/filter"
	"chain/errors"
	"chain/net/http/httpjson"
	"chain/protocol/bc"
)

// These types enforce the ordering of JSON fields in API output.
type (
	txinResp struct {
		Action          interface{} `json:"action"`
		AssetID         interface{} `json:"asset_id"`
		AssetAlias      interface{} `json:"asset_alias"`
		AssetTags       interface{} `json:"asset_tags"`
		Amount          interface{} `json:"amount"`
		IssuanceProgram interface{} `json:"issuance_program"`
		SpentOutput     interface{} `json:"spent_output"`
		AccountID       interface{} `json:"account_id"`
		AccountAlias    interface{} `json:"account_alias"`
		AccountTags     interface{} `json:"account_tags"`
		ReferenceData   interface{} `json:"reference_data"`
		AssetDefinition interface{} `json:"asset_definition"`
	}
	txoutResp struct {
		Action         interface{} `json:"action"`
		Position       interface{} `json:"position"`
		AssetID        interface{} `json:"asset_id"`
		AssetAlias     interface{} `json:"asset_alias"`
		AssetTags      interface{} `json:"asset_tags"`
		Amount         interface{} `json:"asset_amount"`
		AccountID      interface{} `json:"account_id"`
		AccountAlias   interface{} `json:"account_alias"`
		AccountTags    interface{} `json:"account_tags"`
		ControlProgram interface{} `json:"control_program"`
		ReferenceData  interface{} `json:"reference_data"`
	}
	txResp struct {
		ID            interface{} `json:"id"`
		Timestamp     interface{} `json:"timestamp"`
		BlockID       interface{} `json:"block_id"`
		BlockHeight   interface{} `json:"block_height"`
		Position      interface{} `json:"position"`
		ReferenceData interface{} `json:"reference_data"`
		Inputs        interface{} `json:"inputs"`
		Outputs       interface{} `json:"outputs"`
	}
)

// listTransactions is an http handler for listing transactions matching
// an index or an ad-hoc filter.
//
// POST /list-transactions
func (a *api) listTransactions(ctx context.Context, in requestQuery) (result page, err error) {
	if in.EndTimeMS == 0 {
		in.EndTimeMS = bc.Millis(time.Now())
	}

	var (
		p     filter.Predicate
		after query.TxAfter
	)

	// Build the filter predicate.
	p, err = filter.Parse(in.Filter)
	if err != nil {
		return result, err
	}

	// Either parse the provided `after` or look one up for the time range.
	if in.After != "" {
		after, err = query.DecodeTxAfter(in.After)
		if err != nil {
			return result, errors.Wrap(err, "decoding `after`")
		}
	} else {
		after, err = a.indexer.LookupTxAfter(ctx, in.StartTimeMS, in.EndTimeMS)
		if err != nil {
			return result, err
		}
	}

	var asc bool
	if in.Order == "asc" {
		asc = true
	}

	limit := defGenericPageSize
	txns, nextAfter, err := a.indexer.Transactions(ctx, p, in.FilterParams, after, limit, asc)
	if err != nil {
		return result, errors.Wrap(err, "running tx query")
	}

	resp := make([]*txResp, 0, len(txns))
	for _, t := range txns {
		tjson, ok := t.(*json.RawMessage)
		if !ok {
			return result, fmt.Errorf("unexpected type %T in Indexer.Transactions output", t)
		}
		if tjson == nil {
			return result, fmt.Errorf("unexpected nil in Indexer.Transactions output")
		}
		var tx map[string]interface{}
		err = json.Unmarshal(*tjson, &tx)
		if err != nil {
			return result, errors.Wrap(err, "decoding Indexer.Transactions output")
		}

		inp, ok := tx["inputs"].([]interface{})
		if !ok {
			return result, fmt.Errorf("unexpected type %T for inputs in Indexer.Transactions output", tx["inputs"])
		}

		var inputs []map[string]interface{}
		for i, in := range inp {
			input, ok := in.(map[string]interface{})
			if !ok {
				return result, fmt.Errorf("unexpected type %T for input %d in Indexer.Transactions output", in, i)
			}
			inputs = append(inputs, input)
		}

		outp, ok := tx["outputs"].([]interface{})
		if !ok {
			return result, fmt.Errorf("unexpected type %T for outputs in Indexer.Transactions output", tx["outputs"])
		}

		var outputs []map[string]interface{}
		for i, out := range outp {
			output, ok := out.(map[string]interface{})
			if !ok {
				return result, fmt.Errorf("unexpected type %T for output %d in Indexer.Transactions output", out, i)
			}
			outputs = append(outputs, output)
		}

		inResps := make([]*txinResp, 0, len(inputs))
		for _, in := range inputs {
			r := &txinResp{
				Action:          in["action"],
				AssetID:         in["asset_id"],
				AssetAlias:      in["asset_alias"],
				AssetTags:       in["asset_tags"],
				Amount:          in["amount"],
				IssuanceProgram: in["issuance_program"],
				SpentOutput:     in["spent_output"],
				AccountID:       in["account_id"],
				AccountAlias:    in["account_alias"],
				AccountTags:     in["account_tags"],
				ReferenceData:   in["reference_data"],
				AssetDefinition: in["asset_definition"],
			}
			inResps = append(inResps, r)
		}
		outResps := make([]*txoutResp, 0, len(outputs))
		for _, out := range outputs {
			r := &txoutResp{
				Action:         out["action"],
				Position:       out["position"],
				AssetID:        out["asset_id"],
				AssetAlias:     out["asset_alias"],
				AssetTags:      out["asset_tags"],
				Amount:         out["asset_amount"],
				AccountID:      out["account_id"],
				AccountAlias:   out["account_alias"],
				AccountTags:    out["account_tags"],
				ControlProgram: out["control_program"],
				ReferenceData:  out["reference_data"],
			}
			outResps = append(outResps, r)
		}
		r := &txResp{
			ID:            tx["id"],
			Timestamp:     tx["timestamp"],
			BlockID:       tx["block_id"],
			BlockHeight:   tx["block_height"],
			Position:      tx["position"],
			ReferenceData: tx["reference_data"],
			Inputs:        inResps,
			Outputs:       outResps,
		}
		resp = append(resp, r)
	}

	out := in
	out.After = nextAfter.String()
	return page{
		Items:    httpjson.Array(resp),
		LastPage: len(resp) < limit,
		Next:     out,
	}, nil
}

// listAccounts is an http handler for listing accounts matching
// an index or an ad-hoc filter.
//
// POST /list-accounts
func (a *api) listAccounts(ctx context.Context, in requestQuery) (page, error) {
	limit := defGenericPageSize

	// Build the filter predicate.
	p, err := filter.Parse(in.Filter)
	if err != nil {
		return page{}, errors.Wrap(err, "parsing acc query")
	}
	after := in.After

	// Use the filter engine for querying account tags.
	accounts, after, err := a.indexer.Accounts(ctx, p, in.FilterParams, after, limit)
	if err != nil {
		return page{}, errors.Wrap(err, "running acc query")
	}

	result := make([]*accountResponse, 0, len(accounts))
	for _, a := range accounts {
		r := &accountResponse{
			ID:     a["id"],
			Alias:  a["alias"],
			XPubs:  a["xpubs"],
			Quorum: a["quorum"],
			Tags:   a["tags"],
		}
		result = append(result, r)
	}

	// Pull in the accounts by the IDs
	out := in
	out.After = after
	return page{
		Items:    httpjson.Array(result),
		LastPage: len(result) < limit,
		Next:     out,
	}, nil
}

// POST /list-balances
func (a *api) listBalances(ctx context.Context, in requestQuery) (result page, err error) {
	if in.TimestampMS == 0 {
		in.TimestampMS = bc.Millis(time.Now())
	}

	var p filter.Predicate
	var sumBy []filter.Field
	p, err = filter.Parse(in.Filter)
	if err != nil {
		return result, err
	}

	for _, field := range in.SumBy {
		f, err := filter.ParseField(field)
		if err != nil {
			return result, err
		}
		sumBy = append(sumBy, f)
	}

	// TODO(jackson): paginate this endpoint.
	balances, err := a.indexer.Balances(ctx, p, in.FilterParams, sumBy, in.TimestampMS)
	if err != nil {
		return result, err
	}

	result.Items = httpjson.Array(balances)
	result.LastPage = true
	result.Next = in
	return result, nil
}

// This type enforces the ordering of JSON fields in API output.
type utxoResp struct {
	Action         interface{} `json:"action"`
	TransactionID  interface{} `json:"transaction_id"`
	Position       interface{} `json:"position"`
	AssetID        interface{} `json:"asset_id"`
	AssetAlias     interface{} `json:"asset_alias"`
	AssetTags      interface{} `json:"asset_tags"`
	Amount         interface{} `json:"asset_amount"`
	AccountID      interface{} `json:"account_id"`
	AccountAlias   interface{} `json:"account_alias"`
	AccountTags    interface{} `json:"account_tags"`
	ControlProgram interface{} `json:"control_program"`
	ReferenceData  interface{} `json:"reference_data"`
}

// POST /list-unspent-outputs
func (a *api) listUnspentOutputs(ctx context.Context, in requestQuery) (result page, err error) {
	if in.TimestampMS == 0 {
		in.TimestampMS = bc.Millis(time.Now())
	}
	var p filter.Predicate
	p, err = filter.Parse(in.Filter)
	if err != nil {
		return result, err
	}

	var after *query.OutputsAfter
	if in.After != "" {
		after, err = query.DecodeOutputsAfter(in.After)
		if err != nil {
			return result, errors.Wrap(err, "decoding `after`")
		}
	}

	limit := defGenericPageSize
	outputs, nextAfter, err := a.indexer.Outputs(ctx, p, in.FilterParams, in.TimestampMS, after, limit)
	if err != nil {
		return result, errors.Wrap(err, "querying outputs")
	}

	resp := make([]*utxoResp, 0, len(outputs))
	for _, o := range outputs {
		ojson, ok := o.(*json.RawMessage)
		if !ok {
			return result, fmt.Errorf("unexpected type %T in Indexer.Outputs output", o)
		}
		if ojson == nil {
			return result, fmt.Errorf("unexpected nil in Indexer.Outputs output")
		}
		var out map[string]interface{}
		err = json.Unmarshal(*ojson, &out)
		if err != nil {
			return result, errors.Wrap(err, "decoding Indexer.Outputs output")
		}
		r := &utxoResp{
			Action:         out["action"],
			TransactionID:  out["transaction_id"],
			Position:       out["position"],
			AssetID:        out["asset_id"],
			AssetAlias:     out["asset_alias"],
			AssetTags:      out["asset_tags"],
			Amount:         out["asset_amount"],
			AccountID:      out["account_id"],
			AccountAlias:   out["account_alias"],
			AccountTags:    out["account_tags"],
			ControlProgram: out["control_program"],
			ReferenceData:  out["reference_data"],
		}
		resp = append(resp, r)
	}

	outQuery := in
	outQuery.After = nextAfter.String()
	return page{
		Items:    resp,
		LastPage: len(resp) < limit,
		Next:     outQuery,
	}, nil
}

// listAssets is an http handler for listing assets matching
// an index or an ad-hoc filter.
//
// POST /list-assets
func (a *api) listAssets(ctx context.Context, in requestQuery) (page, error) {
	limit := defGenericPageSize

	// Build the filter predicate.
	p, err := filter.Parse(in.Filter)
	if err != nil {
		return page{}, err
	}
	after := in.After

	// Use the query engine for querying asset tags.
	var assets []map[string]interface{}
	assets, after, err = a.indexer.Assets(ctx, p, in.FilterParams, after, limit)
	if err != nil {
		return page{}, errors.Wrap(err, "running asset query")
	}

	result := make([]*assetResponse, 0, len(assets))
	for _, a := range assets {
		r := &assetResponse{
			ID:              a["id"],
			IssuanceProgram: a["issuance_program"],
			XPubs:           a["xpubs"],
			Quorum:          a["quorum"],
			Definition:      a["definition"],
			Tags:            a["tags"],
			Origin:          a["origin"],
		}
		if alias, ok := a["alias"].(string); ok && alias != "" {
			r.Alias = &alias
		}
		result = append(result, r)
	}

	out := in
	out.After = after
	return page{
		Items:    httpjson.Array(result),
		LastPage: len(result) < limit,
		Next:     out,
	}, nil
}
<|endoftext|>


<commit_before>package common

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"regexp"
	"strconv"
	"strings"
)

// MaxInt gives the maximum value of the machine-dependent default integer type. (Standard library
// constants are specific to machine-independent types.)
const MaxInt = int(^uint(0) >> 1)

// Major, Minor, and Patch define group indices of corresponding version number components in versionNumber.
const (
	_ = iota
	Major
	Minor
	Patch
)

// versionNumber groups each of the three numbers of a three-part version number '#.#.#'.
var versionNumber = regexp.MustCompile(`(\d+)\.(\d+)\.(\d+)`)

// Bump increments the specified component of the first three-part version number '#.#.#' found in
// the input, returning the incremented version string alone.
func Bump(versioned string, component int) (bumped string, err error) {
	nums := versionNumber.FindStringSubmatch(versioned)
	if nums == nil {
		err = errors.New("No existing version number '#.#.#' found")
		return
	}
	if component < Major || component > Patch {
		err = fmt.Errorf("Invalid component %v selected for increment of version %v", component, nums[0])
		return
	}
	i, err := strconv.Atoi(nums[component])
	if err != nil {
		err = fmt.Errorf("Error parsing component %v of version %v: %v", component, nums[0], err)
		return
	}
	if i == MaxInt {
		err = fmt.Errorf("Integer overflow incrementing component %v of version %v", component, nums[0])
	}
	nums[component] = strconv.Itoa(i + 1)
	bumped = strings.Join(nums[Major:], ".")
	return
}

// UpdateFile rewrites the named file by applying the given function to its contents, returning the
// modified contents along with any auxiliary information.
func UpdateFile(name string, update func([]byte) ([]byte, interface{}, error)) (info interface{}, err error) {
	stat, err := os.Stat(name)
	if err != nil {
		err = fmt.Errorf("Error finding file %s: %v", name, err)
		return
	}
	contents, err := ioutil.ReadFile(name)
	if err != nil {
		err = fmt.Errorf("Error reading file %s: %v", name, err)
		return
	}
	changed, info, err := update(contents)
	if err != nil {
		err = fmt.Errorf("Error updating file %s: %v", name, err)
		return
	}
	err = ioutil.WriteFile(name, changed, stat.Mode())
	if err != nil {
		err = fmt.Errorf("Error writing file %s: %v", name, err)
		return
	}
	return
}

// ReplaceValue replaces the value of a top-level field in a JSON object with the given changed
// value, returning the modified object.
//
// This implementation works around the lack of order-preserving (un)marshaling in the standard
// library `json` package, by the most straightforward means available for our purposes: It uses the
// standard library decoder to locate the boundaries of the existing value, replacing it with an
// encoding of the changed value with no extraneous formatting, and leaving the remainder of the
// object encoding unchanged.
//
// Note that Go developers have rejected order-preserving capabilities in the standard library with
// assertions that JSON is not a human-readable format (the standard library's support for
// indentation notwithstanding). Compare the second sentence of http://www.json.org, which states
// that "[JSON] is easy for humans to read and write."
func ReplaceValue(in []byte, field string, changed interface{}) (out []byte, err error) {
	malformed := func(err error) error {
		if err != nil {
			return fmt.Errorf("Error parsing JSON object: %v", err)
		} else {
			return errors.New("Error parsing JSON object")
		}
	}

	serial, err := json.Marshal(changed)
	if err != nil {
		err = fmt.Errorf("Error encoding %#v to JSON: %v", changed, err)
		return
	}

	var behind bytes.Buffer
	ahead := io.TeeReader(bytes.NewReader(in), &behind)
	parse := json.NewDecoder(ahead)

	// Enter top-level object
	open, err := parse.Token()
	if err != nil || open != json.Delim('{') {
		err = malformed(err)
		return
	}

	var token json.Token
	var buffer, tail []byte
	for parse.More() {

		// Read next top-level name
		token, err = parse.Token()
		name, ok := token.(string)
		if err != nil || !ok {
			err = malformed(err)
			return
		}
		hit := (name == field)
		if hit {
			buffer, err = ioutil.ReadAll(parse.Buffered())
			if err != nil {
				err = fmt.Errorf("Error reading JSON decoder buffer: %v", err)
				return
			}
			out = bytes.TrimSuffix(behind.Bytes(), buffer)
			out = append(out, ": "...)
			out = append(out, serial...)
		}

		// Read next top-level value
		var value json.RawMessage
		err = parse.Decode(&value)
		if err != nil {
			err = malformed(err)
			return
		}
		if hit {
			buffer, err = ioutil.ReadAll(parse.Buffered())
			if err != nil {
				err = fmt.Errorf("Error encoding %#v to JSON: %v", changed, err)
				return
			}
			out = append(out, buffer...)
			tail, err = ioutil.ReadAll(ahead)
			if err != nil {
				err = fmt.Errorf("Error reading JSON object: %v", err)
			}
			out = append(out, tail...)
			return
		}
	}
	err = fmt.Errorf("Top-level field %#v not found in JSON object", field)
	return
}
<commit_msg>Address PR feedback<commit_after>package common

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"regexp"
	"strconv"
	"strings"
)

// MaxInt gives the maximum value of the machine-dependent default integer type. (Standard library
// constants are specific to machine-independent types.)
const MaxInt = int(^uint(0) >> 1)

// Major, Minor, and Patch define group indices of corresponding version number components in
// versionNumber. Indices are one-based to match subgroup indices in regexp capture.
const (
	_ = iota
	Major
	Minor
	Patch
)

// versionNumber groups each of the three numbers of a three-part version number '#.#.#'.
var versionNumber = regexp.MustCompile(`(\d+)\.(\d+)\.(\d+)`)

// Bump increments the specified component of the first three-part version number '#.#.#' found in
// the input, returning the incremented version string alone.
func Bump(versioned string, component int) (bumped string, err error) {
	nums := versionNumber.FindStringSubmatch(versioned)
	if nums == nil {
		err = errors.New("No existing version number '#.#.#' found")
		return
	}
	if component < Major || component > Patch {
		err = fmt.Errorf("Invalid component %v selected for increment of version %v", component, nums[0])
		return
	}
	i, err := strconv.Atoi(nums[component])
	if err != nil {
		err = fmt.Errorf("Error parsing component %v of version %v: %v", component, nums[0], err)
		return
	}
	if i == MaxInt {
		err = fmt.Errorf("Integer overflow incrementing component %v of version %v", component, nums[0])
	}
	nums[component] = strconv.Itoa(i + 1)
	bumped = strings.Join(nums[Major:], ".")
	return
}

// UpdateFile rewrites the named file by applying the given update function to its contents,
// returning the modified contents along with any auxiliary information returned by the update
// function.
func UpdateFile(name string, update func([]byte) ([]byte, interface{}, error)) (info interface{}, err error) {
	stat, err := os.Stat(name)
	if err != nil {
		err = fmt.Errorf("Error finding file %s: %v", name, err)
		return
	}
	contents, err := ioutil.ReadFile(name)
	if err != nil {
		err = fmt.Errorf("Error reading file %s: %v", name, err)
		return
	}
	changed, info, err := update(contents)
	if err != nil {
		err = fmt.Errorf("Error updating file %s: %v", name, err)
		return
	}
	err = ioutil.WriteFile(name, changed, stat.Mode())
	if err != nil {
		err = fmt.Errorf("Error writing file %s: %v", name, err)
		return
	}
	return
}

// ReplaceValue replaces the value of a top-level field in a JSON object with the given changed
// value, returning the modified object.
//
// This implementation works around the lack of order-preserving (un)marshaling in the standard
// library `json` package, by the most straightforward means available for our purposes: It uses the
// standard library decoder to locate the boundaries of the existing value, replacing it with an
// encoding of the changed value with no extraneous formatting, and leaving the remainder of the
// object encoding unchanged.
//
// Note that Go developers have rejected order-preserving capabilities in the standard library with
// assertions that JSON is not a human-readable format (the standard library's support for
// indentation notwithstanding). Compare the second sentence of http://www.json.org, which states
// that "[JSON] is easy for humans to read and write."
func ReplaceValue(in []byte, field string, changed interface{}) (out []byte, err error) {
	malformed := func(err error) error {
		if err != nil {
			return fmt.Errorf("Error parsing JSON object: %v", err)
		} else {
			return errors.New("Error parsing JSON object")
		}
	}

	serial, err := json.Marshal(changed)
	if err != nil {
		err = fmt.Errorf("Error encoding %#v to JSON: %v", changed, err)
		return
	}

	var behind bytes.Buffer
	ahead := io.TeeReader(bytes.NewReader(in), &behind)
	parse := json.NewDecoder(ahead)

	// Enter top-level object
	open, err := parse.Token()
	if err != nil || open != json.Delim('{') {
		err = malformed(err)
		return
	}

	var token json.Token
	var buffer, tail []byte
	for parse.More() {

		// Read next top-level name
		token, err = parse.Token()
		name, ok := token.(string)
		if err != nil || !ok {
			err = malformed(err)
			return
		}
		hit := (name == field)
		if hit {
			buffer, err = ioutil.ReadAll(parse.Buffered())
			if err != nil {
				err = fmt.Errorf("Error reading JSON decoder buffer: %v", err)
				return
			}
			out = bytes.TrimSuffix(behind.Bytes(), buffer)
			out = append(out, ": "...)
			out = append(out, serial...)
		}

		// Read next top-level value
		var value json.RawMessage
		err = parse.Decode(&value)
		if err != nil {
			err = malformed(err)
			return
		}
		if hit {
			buffer, err = ioutil.ReadAll(parse.Buffered())
			if err != nil {
				err = fmt.Errorf("Error encoding %#v to JSON: %v", changed, err)
				return
			}
			out = append(out, buffer...)
			tail, err = ioutil.ReadAll(ahead)
			if err != nil {
				err = fmt.Errorf("Error reading JSON object: %v", err)
			}
			out = append(out, tail...)
			return
		}
	}
	err = fmt.Errorf("Top-level field %#v not found in JSON object", field)
	return
}
<|endoftext|>


<commit_before>// Copyright 2016 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// TODO(jba): document in CONTRIBUTING.md that service account must be given "Logs Configuration Writer" IAM role for sink tests to pass.
// TODO(jba): [cont] (1) From top left menu, go to IAM & Admin. (2) In Roles dropdown for acct, select Logging > Logs Configuration Writer. (3) Save.
// TODO(jba): Also, cloud-logs@google.com must have Owner permission on the GCS bucket named for the test project.

package logadmin

import (
	"context"
	"log"
	"testing"
	"time"

	"cloud.google.com/go/internal/testutil"
	"cloud.google.com/go/internal/uid"
	ltest "cloud.google.com/go/logging/internal/testing"
	"cloud.google.com/go/storage"
	"google.golang.org/api/iterator"
	"google.golang.org/api/option"
)

var sinkIDs = uid.NewSpace("GO-CLIENT-TEST-SINK", nil)

const testFilter = ""

var testSinkDestination string

// Called just before TestMain calls m.Run.
// Returns a cleanup function to be called after the tests finish.
func initSinks(ctx context.Context) func() {
	// Create a unique GCS bucket so concurrent tests don't interfere with each other.
	bucketIDs := uid.NewSpace(testProjectID+"-log-sink", nil)
	testBucket := bucketIDs.New()
	testSinkDestination = "storage.googleapis.com/" + testBucket
	var storageClient *storage.Client
	if integrationTest {
		// Create a unique bucket as a sink destination, and give the cloud logging account
		// owner right.
		ts := testutil.TokenSource(ctx, storage.ScopeFullControl)
		var err error
		storageClient, err = storage.NewClient(ctx, option.WithTokenSource(ts))
		if err != nil {
			log.Fatalf("new storage client: %v", err)
		}
		bucket := storageClient.Bucket(testBucket)
		if err := bucket.Create(ctx, testProjectID, nil); err != nil {
			log.Fatalf("creating storage bucket %q: %v", testBucket, err)
		}
		log.Printf("successfully created bucket %s", testBucket)
		if err := bucket.ACL().Set(ctx, "group-cloud-logs@google.com", storage.RoleOwner); err != nil {
			log.Fatalf("setting owner role: %v", err)
		}
	}
	// Clean up from aborted tests.
	it := client.Sinks(ctx)
	for {
		s, err := it.Next()
		if err == iterator.Done {
			break
		}
		if err != nil {
			log.Printf("listing sinks: %v", err)
			break
		}
		if sinkIDs.Older(s.ID, time.Hour) {
			client.DeleteSink(ctx, s.ID) // ignore error
		}
	}
	if integrationTest {
		for _, bn := range bucketNames(ctx, storageClient) {
			if bucketIDs.Older(bn, 24*time.Hour) {
				storageClient.Bucket(bn).Delete(ctx) // ignore error
			}
		}
		return func() {
			if err := storageClient.Bucket(testBucket).Delete(ctx); err != nil {
				log.Printf("deleting %q: %v", testBucket, err)
			}
			storageClient.Close()
		}
	}
	return func() {}
}

// Collect the name of all buckets for the test project.
func bucketNames(ctx context.Context, client *storage.Client) []string {
	var names []string
	it := client.Buckets(ctx, testProjectID)
loop:
	for {
		b, err := it.Next()
		switch err {
		case nil:
			names = append(names, b.Name)
		case iterator.Done:
			break loop
		default:
			log.Printf("listing buckets: %v", err)
			break loop
		}
	}
	return names
}

func TestCreateSink(t *testing.T) {
	ctx := context.Background()
	sink := &Sink{
		ID:              sinkIDs.New(),
		Destination:     testSinkDestination,
		Filter:          testFilter,
		IncludeChildren: true,
	}
	got, err := client.CreateSink(ctx, sink)
	if err != nil {
		t.Fatal(err)
	}
	sink.WriterIdentity = ltest.SharedServiceAccount
	if want := sink; !testutil.Equal(got, want) {
		t.Errorf("got %+v, want %+v", got, want)
	}
	got, err = client.Sink(ctx, sink.ID)
	if err != nil {
		t.Fatal(err)
	}
	if want := sink; !testutil.Equal(got, want) {
		t.Errorf("got %+v, want %+v", got, want)
	}

	// UniqueWriterIdentity
	sink.ID = sinkIDs.New()
	got, err = client.CreateSinkOpt(ctx, sink, SinkOptions{UniqueWriterIdentity: true})
	if err != nil {
		t.Fatal(err)
	}
	// The WriterIdentity should be different.
	if got.WriterIdentity == sink.WriterIdentity {
		t.Errorf("got %s, want something different", got.WriterIdentity)
	}
}

func TestUpdateSink(t *testing.T) {
	ctx := context.Background()
	sink := &Sink{
		ID:              sinkIDs.New(),
		Destination:     testSinkDestination,
		Filter:          testFilter,
		IncludeChildren: true,
		WriterIdentity:  ltest.SharedServiceAccount,
	}

	if _, err := client.CreateSink(ctx, sink); err != nil {
		t.Fatal(err)
	}
	got, err := client.UpdateSink(ctx, sink)
	if err != nil {
		t.Fatal(err)
	}
	if want := sink; !testutil.Equal(got, want) {
		t.Errorf("got\n%+v\nwant\n%+v", got, want)
	}
	got, err = client.Sink(ctx, sink.ID)
	if err != nil {
		t.Fatal(err)
	}
	if want := sink; !testutil.Equal(got, want) {
		t.Errorf("got\n%+v\nwant\n%+v", got, want)
	}

	// Updating an existing sink changes it.
	sink.Filter = ""
	sink.IncludeChildren = false
	if _, err := client.UpdateSink(ctx, sink); err != nil {
		t.Fatal(err)
	}
	got, err = client.Sink(ctx, sink.ID)
	if err != nil {
		t.Fatal(err)
	}
	if want := sink; !testutil.Equal(got, want) {
		t.Errorf("got\n%+v\nwant\n%+v", got, want)
	}
}

func TestUpdateSinkOpt(t *testing.T) {
	ctx := context.Background()
	id := sinkIDs.New()
	origSink := &Sink{
		ID:              id,
		Destination:     testSinkDestination,
		Filter:          testFilter,
		IncludeChildren: true,
		WriterIdentity:  ltest.SharedServiceAccount,
	}

	if _, err := client.CreateSink(ctx, origSink); err != nil {
		t.Fatal(err)
	}

	// Updating with empty options is an error.
	_, err := client.UpdateSinkOpt(ctx, &Sink{ID: id, Destination: testSinkDestination}, SinkOptions{})
	if err == nil {
		t.Errorf("got %v, want nil", err)
	}

	// Update selected fields.
	got, err := client.UpdateSinkOpt(ctx, &Sink{ID: id}, SinkOptions{
		UpdateFilter:          true,
		UpdateIncludeChildren: true,
	})
	if err != nil {
		t.Fatal(err)
	}
	want := *origSink
	want.Filter = ""
	want.IncludeChildren = false
	if !testutil.Equal(got, &want) {
		t.Errorf("got\n%+v\nwant\n%+v", got, want)
	}

	// Update writer identity.
	got, err = client.UpdateSinkOpt(ctx, &Sink{ID: id, Filter: "foo"},
		SinkOptions{UniqueWriterIdentity: true})
	if err != nil {
		t.Fatal(err)
	}
	if got.WriterIdentity == want.WriterIdentity {
		t.Errorf("got %s, want something different", got.WriterIdentity)
	}
	want.WriterIdentity = got.WriterIdentity
	if !testutil.Equal(got, &want) {
		t.Errorf("got\n%+v\nwant\n%+v", got, want)
	}
}

func TestListSinks(t *testing.T) {
	ctx := context.Background()
	var sinks []*Sink
	want := map[string]*Sink{}
	for i := 0; i < 4; i++ {
		s := &Sink{
			ID:             sinkIDs.New(),
			Destination:    testSinkDestination,
			Filter:         testFilter,
			WriterIdentity: "serviceAccount:cloud-logs@system.gserviceaccount.com",
		}
		sinks = append(sinks, s)
		want[s.ID] = s
	}
	for _, s := range sinks {
		if _, err := client.CreateSink(ctx, s); err != nil {
			t.Fatalf("Create(%q): %v", s.ID, err)
		}
	}

	got := map[string]*Sink{}
	it := client.Sinks(ctx)
	for {
		s, err := it.Next()
		if err == iterator.Done {
			break
		}
		if err != nil {
			t.Fatal(err)
		}
		// If tests run simultaneously, we may have more sinks than we
		// created. So only check for our own.
		if _, ok := want[s.ID]; ok {
			got[s.ID] = s
		}
	}
	if !testutil.Equal(got, want) {
		t.Errorf("got %+v, want %+v", got, want)
	}
}
<commit_msg>logadmin: Delete test buckets only after 24 hours<commit_after>// Copyright 2016 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// TODO(jba): document in CONTRIBUTING.md that service account must be given "Logs Configuration Writer" IAM role for sink tests to pass.
// TODO(jba): [cont] (1) From top left menu, go to IAM & Admin. (2) In Roles dropdown for acct, select Logging > Logs Configuration Writer. (3) Save.
// TODO(jba): Also, cloud-logs@google.com must have Owner permission on the GCS bucket named for the test project.

package logadmin

import (
	"context"
	"log"
	"testing"
	"time"

	"cloud.google.com/go/internal/testutil"
	"cloud.google.com/go/internal/uid"
	ltest "cloud.google.com/go/logging/internal/testing"
	"cloud.google.com/go/storage"
	"google.golang.org/api/iterator"
	"google.golang.org/api/option"
)

var sinkIDs = uid.NewSpace("GO-CLIENT-TEST-SINK", nil)

const testFilter = ""

var testSinkDestination string

// Called just before TestMain calls m.Run.
// Returns a cleanup function to be called after the tests finish.
func initSinks(ctx context.Context) func() {
	// Create a unique GCS bucket so concurrent tests don't interfere with each other.
	bucketIDs := uid.NewSpace(testProjectID+"-log-sink", nil)
	testBucket := bucketIDs.New()
	testSinkDestination = "storage.googleapis.com/" + testBucket
	var storageClient *storage.Client
	if integrationTest {
		// Create a unique bucket as a sink destination, and give the cloud logging account
		// owner right.
		ts := testutil.TokenSource(ctx, storage.ScopeFullControl)
		var err error
		storageClient, err = storage.NewClient(ctx, option.WithTokenSource(ts))
		if err != nil {
			log.Fatalf("new storage client: %v", err)
		}
		bucket := storageClient.Bucket(testBucket)
		if err := bucket.Create(ctx, testProjectID, nil); err != nil {
			log.Fatalf("creating storage bucket %q: %v", testBucket, err)
		}
		log.Printf("successfully created bucket %s", testBucket)
		if err := bucket.ACL().Set(ctx, "group-cloud-logs@google.com", storage.RoleOwner); err != nil {
			log.Fatalf("setting owner role: %v", err)
		}
	}
	// Clean up from aborted tests.
	it := client.Sinks(ctx)
	for {
		s, err := it.Next()
		if err == iterator.Done {
			break
		}
		if err != nil {
			log.Printf("listing sinks: %v", err)
			break
		}
		if sinkIDs.Older(s.ID, time.Hour) {
			client.DeleteSink(ctx, s.ID) // ignore error
		}
	}
	if integrationTest {
		for _, bn := range bucketNames(ctx, storageClient) {
			if bucketIDs.Older(bn, 24*time.Hour) {
				storageClient.Bucket(bn).Delete(ctx) // ignore error
			}
		}
		return func() {
			storageClient.Close()
		}
	}
	return func() {}
}

// Collect the name of all buckets for the test project.
func bucketNames(ctx context.Context, client *storage.Client) []string {
	var names []string
	it := client.Buckets(ctx, testProjectID)
loop:
	for {
		b, err := it.Next()
		switch err {
		case nil:
			names = append(names, b.Name)
		case iterator.Done:
			break loop
		default:
			log.Printf("listing buckets: %v", err)
			break loop
		}
	}
	return names
}

func TestCreateSink(t *testing.T) {
	ctx := context.Background()
	sink := &Sink{
		ID:              sinkIDs.New(),
		Destination:     testSinkDestination,
		Filter:          testFilter,
		IncludeChildren: true,
	}
	got, err := client.CreateSink(ctx, sink)
	if err != nil {
		t.Fatal(err)
	}
	sink.WriterIdentity = ltest.SharedServiceAccount
	if want := sink; !testutil.Equal(got, want) {
		t.Errorf("got %+v, want %+v", got, want)
	}
	got, err = client.Sink(ctx, sink.ID)
	if err != nil {
		t.Fatal(err)
	}
	if want := sink; !testutil.Equal(got, want) {
		t.Errorf("got %+v, want %+v", got, want)
	}

	// UniqueWriterIdentity
	sink.ID = sinkIDs.New()
	got, err = client.CreateSinkOpt(ctx, sink, SinkOptions{UniqueWriterIdentity: true})
	if err != nil {
		t.Fatal(err)
	}
	// The WriterIdentity should be different.
	if got.WriterIdentity == sink.WriterIdentity {
		t.Errorf("got %s, want something different", got.WriterIdentity)
	}
}

func TestUpdateSink(t *testing.T) {
	ctx := context.Background()
	sink := &Sink{
		ID:              sinkIDs.New(),
		Destination:     testSinkDestination,
		Filter:          testFilter,
		IncludeChildren: true,
		WriterIdentity:  ltest.SharedServiceAccount,
	}

	if _, err := client.CreateSink(ctx, sink); err != nil {
		t.Fatal(err)
	}
	got, err := client.UpdateSink(ctx, sink)
	if err != nil {
		t.Fatal(err)
	}
	if want := sink; !testutil.Equal(got, want) {
		t.Errorf("got\n%+v\nwant\n%+v", got, want)
	}
	got, err = client.Sink(ctx, sink.ID)
	if err != nil {
		t.Fatal(err)
	}
	if want := sink; !testutil.Equal(got, want) {
		t.Errorf("got\n%+v\nwant\n%+v", got, want)
	}

	// Updating an existing sink changes it.
	sink.Filter = ""
	sink.IncludeChildren = false
	if _, err := client.UpdateSink(ctx, sink); err != nil {
		t.Fatal(err)
	}
	got, err = client.Sink(ctx, sink.ID)
	if err != nil {
		t.Fatal(err)
	}
	if want := sink; !testutil.Equal(got, want) {
		t.Errorf("got\n%+v\nwant\n%+v", got, want)
	}
}

func TestUpdateSinkOpt(t *testing.T) {
	ctx := context.Background()
	id := sinkIDs.New()
	origSink := &Sink{
		ID:              id,
		Destination:     testSinkDestination,
		Filter:          testFilter,
		IncludeChildren: true,
		WriterIdentity:  ltest.SharedServiceAccount,
	}

	if _, err := client.CreateSink(ctx, origSink); err != nil {
		t.Fatal(err)
	}

	// Updating with empty options is an error.
	_, err := client.UpdateSinkOpt(ctx, &Sink{ID: id, Destination: testSinkDestination}, SinkOptions{})
	if err == nil {
		t.Errorf("got %v, want nil", err)
	}

	// Update selected fields.
	got, err := client.UpdateSinkOpt(ctx, &Sink{ID: id}, SinkOptions{
		UpdateFilter:          true,
		UpdateIncludeChildren: true,
	})
	if err != nil {
		t.Fatal(err)
	}
	want := *origSink
	want.Filter = ""
	want.IncludeChildren = false
	if !testutil.Equal(got, &want) {
		t.Errorf("got\n%+v\nwant\n%+v", got, want)
	}

	// Update writer identity.
	got, err = client.UpdateSinkOpt(ctx, &Sink{ID: id, Filter: "foo"},
		SinkOptions{UniqueWriterIdentity: true})
	if err != nil {
		t.Fatal(err)
	}
	if got.WriterIdentity == want.WriterIdentity {
		t.Errorf("got %s, want something different", got.WriterIdentity)
	}
	want.WriterIdentity = got.WriterIdentity
	if !testutil.Equal(got, &want) {
		t.Errorf("got\n%+v\nwant\n%+v", got, want)
	}
}

func TestListSinks(t *testing.T) {
	ctx := context.Background()
	var sinks []*Sink
	want := map[string]*Sink{}
	for i := 0; i < 4; i++ {
		s := &Sink{
			ID:             sinkIDs.New(),
			Destination:    testSinkDestination,
			Filter:         testFilter,
			WriterIdentity: "serviceAccount:cloud-logs@system.gserviceaccount.com",
		}
		sinks = append(sinks, s)
		want[s.ID] = s
	}
	for _, s := range sinks {
		if _, err := client.CreateSink(ctx, s); err != nil {
			t.Fatalf("Create(%q): %v", s.ID, err)
		}
	}

	got := map[string]*Sink{}
	it := client.Sinks(ctx)
	for {
		s, err := it.Next()
		if err == iterator.Done {
			break
		}
		if err != nil {
			t.Fatal(err)
		}
		// If tests run simultaneously, we may have more sinks than we
		// created. So only check for our own.
		if _, ok := want[s.ID]; ok {
			got[s.ID] = s
		}
	}
	if !testutil.Equal(got, want) {
		t.Errorf("got %+v, want %+v", got, want)
	}
}
<|endoftext|>


<commit_before>package ingress

import (
	"crypto/md5"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"strings"

	"github.com/rancher/norman/store/transform"
	"github.com/rancher/norman/types"
	"github.com/rancher/norman/types/convert"
	"github.com/rancher/norman/types/values"
	"github.com/sirupsen/logrus"
)

func Wrap(store types.Store) types.Store {
	modify := &Store{
		store,
	}
	return New(modify)
}

type Store struct {
	types.Store
}

func (p *Store) Create(apiContext *types.APIContext, schema *types.Schema, data map[string]interface{}) (map[string]interface{}, error) {
	formatData(data, false)
	data, err := p.Store.Create(apiContext, schema, data)
	return data, err
}

func formatData(data map[string]interface{}, forFrontend bool) {
	oldState := getState(data)
	newState := map[string]string{}

	// transform default backend
	if target, ok := values.GetValue(data, "defaultBackend"); ok {
		updateRule(convert.ToMapInterface(target), "/", forFrontend, data, oldState, newState)
	}

	// transform rules
	if paths, ok := getPaths(data); ok {
		for hostpath, target := range paths {
			updateRule(target, hostpath, forFrontend, data, oldState, newState)
		}
	}

	setState(data, newState)

}
func updateRule(target map[string]interface{}, hostpath string, forFrontend bool, data map[string]interface{}, oldState map[string]string, newState map[string]string) {
	targetData := convert.ToMapInterface(target)
	port, _ := targetData["targetPort"]
	serviceID, _ := targetData["serviceId"].(string)
	stateKey := getStateKey(hostpath, convert.ToString(port))
	if forFrontend {
		isService := true
		if serviceValue, ok := oldState[stateKey]; ok && !convert.IsEmpty(serviceValue) {
			targetData["workloadIds"] = strings.Split(serviceValue, "/")
			isService = false
		}

		if isService {
			targetData["serviceId"] = fmt.Sprintf("%s/%s", data["namespaceId"].(string), serviceID)
		} else {
			delete(targetData, "serviceId")
		}
	} else {
		workloadIDs := convert.ToStringSlice(targetData["workloadIds"])
		if serviceID != "" {
			splitted := strings.Split(serviceID, ":")
			if len(splitted) > 1 {
				serviceID = splitted[1]
			}
		} else {
			serviceID = getServiceID(stateKey)
		}
		newState[stateKey] = strings.Join(workloadIDs, "/")
		targetData["serviceId"] = serviceID
	}
}

func getServiceID(stateKey string) string {
	bytes, err := base64.URLEncoding.DecodeString(stateKey)
	if err != nil {
		return ""
	}

	sum := md5.Sum(bytes)
	hex := "ingress-" + hex.EncodeToString(sum[:])

	return hex
}

func getStateKey(hostpath string, port string) string {
	key := fmt.Sprintf("%s/%s", hostpath, port)
	return base64.URLEncoding.EncodeToString([]byte(key))
}

func getPaths(data map[string]interface{}) (map[string]map[string]interface{}, bool) {
	v, ok := values.GetValue(data, "rules")
	if !ok {
		return nil, false
	}

	result := make(map[string]map[string]interface{})
	for _, rule := range convert.ToMapSlice(v) {
		converted := convert.ToMapInterface(rule)
		paths, ok := converted["paths"]
		if ok {
			for path, target := range convert.ToMapInterface(paths) {
				result[fmt.Sprintf("%s/%s", convert.ToString(converted["host"]), path)] = convert.ToMapInterface(target)
			}
		}
	}

	return result, true
}

func setState(data map[string]interface{}, stateMap map[string]string) {
	content, err := json.Marshal(stateMap)
	if err != nil {
		logrus.Errorf("failed to save state on ingress: %v", data["id"])
		return
	}

	values.PutValue(data, string(content), "annotations", "ingress.cattle.io/state")
}

func getState(data map[string]interface{}) map[string]string {
	state := make(map[string]string)

	v, ok := values.GetValue(data, "annotations", "ingress.cattle.io/state")
	if ok {
		json.Unmarshal([]byte(convert.ToString(v)), &state)
	}

	return state
}

func New(store types.Store) types.Store {
	return &transform.Store{
		Store: store,
		Transformer: func(apiContext *types.APIContext, data map[string]interface{}, opt *types.QueryOptions) (map[string]interface{}, error) {
			formatData(data, true)
			return data, nil
		},
	}
}
<commit_msg>Ingress fixes<commit_after>package ingress

import (
	"crypto/md5"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"strings"

	"github.com/rancher/norman/store/transform"
	"github.com/rancher/norman/types"
	"github.com/rancher/norman/types/convert"
	"github.com/rancher/norman/types/values"
	"github.com/sirupsen/logrus"
)

func Wrap(store types.Store) types.Store {
	modify := &Store{
		store,
	}
	return New(modify)
}

type Store struct {
	types.Store
}

func (p *Store) Create(apiContext *types.APIContext, schema *types.Schema, data map[string]interface{}) (map[string]interface{}, error) {
	formatData(data, false)
	data, err := p.Store.Create(apiContext, schema, data)
	return data, err
}

func (p *Store) Update(apiContext *types.APIContext, schema *types.Schema, data map[string]interface{}, id string) (map[string]interface{}, error) {
	formatData(data, false)
	data, err := p.Store.Create(apiContext, schema, data)
	return data, err
}

func formatData(data map[string]interface{}, forFrontend bool) {
	oldState := getState(data)
	newState := map[string]string{}

	// transform default backend
	if target, ok := values.GetValue(data, "defaultBackend"); ok {
		updateRule(convert.ToMapInterface(target), "/", forFrontend, data, oldState, newState)
	}

	// transform rules
	if paths, ok := getPaths(data); ok {
		for hostpath, target := range paths {
			updateRule(target, hostpath, forFrontend, data, oldState, newState)
		}
	}

	setState(data, newState)

}
func updateRule(target map[string]interface{}, hostpath string, forFrontend bool, data map[string]interface{}, oldState map[string]string, newState map[string]string) {
	targetData := convert.ToMapInterface(target)
	port, _ := targetData["targetPort"]
	serviceID, _ := targetData["serviceId"].(string)
	stateKey := getStateKey(hostpath, convert.ToString(port))
	if forFrontend {
		isService := true
		if serviceValue, ok := oldState[stateKey]; ok && !convert.IsEmpty(serviceValue) {
			targetData["workloadIds"] = strings.Split(serviceValue, "/")
			isService = false
		}

		if isService {
			targetData["serviceId"] = fmt.Sprintf("%s/%s", data["namespaceId"].(string), serviceID)
		} else {
			delete(targetData, "serviceId")
		}
	} else {
		workloadIDs := convert.ToStringSlice(targetData["workloadIds"])
		if serviceID != "" {
			splitted := strings.Split(serviceID, ":")
			if len(splitted) > 1 {
				serviceID = splitted[1]
			}
		} else {
			serviceID = getServiceID(stateKey)
		}
		newState[stateKey] = strings.Join(workloadIDs, "/")
		targetData["serviceId"] = serviceID
	}
}

func getServiceID(stateKey string) string {
	bytes, err := base64.URLEncoding.DecodeString(stateKey)
	if err != nil {
		return ""
	}

	sum := md5.Sum(bytes)
	hex := "ingress-" + hex.EncodeToString(sum[:])

	return hex
}

func getStateKey(hostpath string, port string) string {
	key := fmt.Sprintf("%s/%s", hostpath, port)
	return base64.URLEncoding.EncodeToString([]byte(key))
}

func getPaths(data map[string]interface{}) (map[string]map[string]interface{}, bool) {
	v, ok := values.GetValue(data, "rules")
	if !ok {
		return nil, false
	}

	result := make(map[string]map[string]interface{})
	for _, rule := range convert.ToMapSlice(v) {
		converted := convert.ToMapInterface(rule)
		paths, ok := converted["paths"]
		if ok {
			for path, target := range convert.ToMapInterface(paths) {
				result[fmt.Sprintf("%s/%s", convert.ToString(converted["host"]), path)] = convert.ToMapInterface(target)
			}
		}
	}

	return result, true
}

func setState(data map[string]interface{}, stateMap map[string]string) {
	content, err := json.Marshal(stateMap)
	if err != nil {
		logrus.Errorf("failed to save state on ingress: %v", data["id"])
		return
	}

	values.PutValue(data, string(content), "annotations", "ingress.cattle.io/state")
}

func getState(data map[string]interface{}) map[string]string {
	state := make(map[string]string)

	v, ok := values.GetValue(data, "annotations", "ingress.cattle.io/state")
	if ok {
		json.Unmarshal([]byte(convert.ToString(v)), &state)
	}

	return state
}

func New(store types.Store) types.Store {
	return &transform.Store{
		Store: store,
		Transformer: func(apiContext *types.APIContext, data map[string]interface{}, opt *types.QueryOptions) (map[string]interface{}, error) {
			formatData(data, true)
			return data, nil
		},
	}
}
<|endoftext|>