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

// Package info provides utilities for summarizing the contents of a kzip.
package info // import "kythe.io/kythe/go/platform/kzip/info"

import (
	"fmt"
	"log"
	"strings"

	"bitbucket.org/creachadair/stringset"

	"kythe.io/kythe/go/platform/kzip"

	apb "kythe.io/kythe/proto/analysis_go_proto"
)

// KzipInfo scans the kzip in f and counts contained files and units, giving a
// breakdown by corpus and language. It also records the size (in bytes) of the
// kzip specified by fileSize in the returned KzipInfo. This is a convenience
// method and thin wrapper over the Accumulator. If you need to do more than
// just calculate KzipInfo while doing a kzip.Scan(), you should use the
// Accumulator directly.
func KzipInfo(f kzip.File, fileSize int64, scanOpts ...kzip.ScanOption) (*apb.KzipInfo, error) {
	a := NewAccumulator(fileSize)
	if err := kzip.Scan(f, func(r *kzip.Reader, unit *kzip.Unit) error {
		a.Accumulate(unit)
		return nil
	}, scanOpts...); err != nil {
		return nil, fmt.Errorf("scanning kzip: %v", err)
	}
	return a.Get(), nil
}

// Accumulator is used to build a summary of a collection of compilation units.
// Usage:
//   a := NewAccumulator(fileSize)
//   a.Accumulate(unit) // call for each compilation unit
//   info := a.Get()    // get the resulting KzipInfo
type Accumulator struct {
	*apb.KzipInfo
}

// NewAccumulator creates a new Accumulator instance given the kzip fileSize (in
// bytes).
func NewAccumulator(fileSize int64) *Accumulator {
	return &Accumulator{
		KzipInfo: &apb.KzipInfo{
			Corpora: make(map[string]*apb.KzipInfo_CorpusInfo),
			Size:    fileSize,
		},
	}
}

// Accumulate should be called for each unit in the kzip so its counts can be
// recorded.
func (a *Accumulator) Accumulate(u *kzip.Unit) {
	srcs := stringset.New(u.Proto.SourceFile...)
	cuLang := u.Proto.GetVName().GetLanguage()
	if cuLang == "" {
		msg := fmt.Sprintf("CU does not specify a language %v", u.Proto.GetVName())
		a.KzipInfo.CriticalKzipErrors = append(a.KzipInfo.CriticalKzipErrors, msg)
		return
	}

	var srcCorpora stringset.Set
	srcsWithRI := stringset.New()
	for _, ri := range u.Proto.RequiredInput {
		if strings.HasPrefix(ri.GetVName().GetPath(), "/") {
			a.KzipInfo.AbsolutePaths = append(a.KzipInfo.AbsolutePaths, ri.GetVName().GetPath())
		}

		riCorpus := requiredInputCorpus(u, ri)
		if riCorpus == "" {
			// Trim spaces to work around the fact that log("%v", proto) is inconsistent about trailing spaces in google3 vs open-source go.
			msg := strings.TrimSpace(fmt.Sprintf("unable to determine corpus for required_input %q in CU %v", ri.Info.Path, u.Proto.GetVName()))
			a.KzipInfo.CriticalKzipErrors = append(a.KzipInfo.CriticalKzipErrors, msg)
			return
		}
		requiredInputInfo(riCorpus, cuLang, a.KzipInfo).Count++
		if srcs.Contains(ri.Info.Path) {
			sourceInfo(riCorpus, cuLang, a.KzipInfo).Count++
			srcCorpora.Add(riCorpus)
			srcsWithRI.Add(ri.Info.Path)
		}
	}
	srcsWithoutRI := srcs.Diff(srcsWithRI)
	for path := range srcsWithoutRI {
		msg := fmt.Sprintf("source %q in CU %v doesn't have a required_input entry", path, u.Proto.GetVName())
		a.KzipInfo.CriticalKzipErrors = append(a.KzipInfo.CriticalKzipErrors, msg)
	}
	if srcCorpora.Len() != 1 {
		// This is a warning for now, but may become an error.
		log.Printf("Multiple corpora in unit. unit vname={%v}; src corpora=%v; srcs=%v", u.Proto.GetVName(), srcCorpora, u.Proto.SourceFile)
	}
}

// Get returns the final KzipInfo after info from each unit in the kzip has been
// accumulated.
func (a *Accumulator) Get() *apb.KzipInfo {
	return a.KzipInfo
}

// requiredInputCorpus computes the corpus for a required input. It follows the rules in the
// CompilationUnit proto comments in kythe/proto/analysis.proto that say that any
// required_input that does not set corpus in its VName should inherit corpus from the compilation
// unit's VName.
func requiredInputCorpus(u *kzip.Unit, ri *apb.CompilationUnit_FileInput) string {
	if c := ri.GetVName().GetCorpus(); c != "" {
		return c
	}
	return u.Proto.GetVName().GetCorpus()
}

// KzipInfoTotalCount returns the total CompilationUnits counts for infos split apart by language.
func KzipInfoTotalCount(infos []*apb.KzipInfo) apb.KzipInfo_CorpusInfo {
	totals := apb.KzipInfo_CorpusInfo{
		LanguageRequiredInputs: make(map[string]*apb.KzipInfo_CorpusInfo_RequiredInputs),
		LanguageSources:        make(map[string]*apb.KzipInfo_CorpusInfo_RequiredInputs),
	}
	for _, info := range infos {
		for _, i := range info.GetCorpora() {
			for lang, stats := range i.GetLanguageRequiredInputs() {
				total := totals.LanguageRequiredInputs[lang]
				if total == nil {
					total = &apb.KzipInfo_CorpusInfo_RequiredInputs{}
					totals.LanguageRequiredInputs[lang] = total
				}
				total.Count += stats.GetCount()
			}
			for lang, stats := range i.GetLanguageSources() {
				total := totals.LanguageSources[lang]
				if total == nil {
					total = &apb.KzipInfo_CorpusInfo_RequiredInputs{}
					totals.LanguageSources[lang] = total
				}
				total.Count += stats.GetCount()
			}
		}
	}
	return totals
}

// MergeKzipInfo combines the counts from multiple KzipInfos.
func MergeKzipInfo(infos []*apb.KzipInfo) *apb.KzipInfo {
	kzipInfo := &apb.KzipInfo{Corpora: make(map[string]*apb.KzipInfo_CorpusInfo)}

	for _, i := range infos {
		for corpus, cinfo := range i.GetCorpora() {
			for lang, inputs := range cinfo.GetLanguageRequiredInputs() {
				c := requiredInputInfo(corpus, lang, kzipInfo)
				c.Count += inputs.GetCount()
			}
			for lang, sources := range cinfo.GetLanguageSources() {
				c := sourceInfo(corpus, lang, kzipInfo)
				c.Count += sources.GetCount()
			}
		}
		kzipInfo.CriticalKzipErrors = append(kzipInfo.GetCriticalKzipErrors(), i.GetCriticalKzipErrors()...)
		kzipInfo.Size += i.Size
		kzipInfo.AbsolutePaths = append(kzipInfo.AbsolutePaths, i.AbsolutePaths...)
	}
	return kzipInfo
}

func requiredInputInfo(corpus, lang string, kzipInfo *apb.KzipInfo) *apb.KzipInfo_CorpusInfo_RequiredInputs {
	c := corpusInfo(corpus, kzipInfo)
	lri := c.LanguageRequiredInputs[lang]
	if lri == nil {
		lri = &apb.KzipInfo_CorpusInfo_RequiredInputs{}
		c.LanguageRequiredInputs[lang] = lri
	}
	return lri
}

func sourceInfo(corpus, lang string, kzipInfo *apb.KzipInfo) *apb.KzipInfo_CorpusInfo_RequiredInputs {
	c := corpusInfo(corpus, kzipInfo)
	ls := c.LanguageSources[lang]
	if ls == nil {
		ls = &apb.KzipInfo_CorpusInfo_RequiredInputs{}
		c.LanguageSources[lang] = ls
	}
	return ls
}

func corpusInfo(corpus string, kzipInfo *apb.KzipInfo) *apb.KzipInfo_CorpusInfo {
	i := kzipInfo.GetCorpora()[corpus]
	if i == nil {
		i = &apb.KzipInfo_CorpusInfo{
			LanguageRequiredInputs: make(map[string]*apb.KzipInfo_CorpusInfo_RequiredInputs),
			LanguageSources:        make(map[string]*apb.KzipInfo_CorpusInfo_RequiredInputs),
		}
		kzipInfo.Corpora[corpus] = i
	}
	return i
}
<commit_msg>fix(kzip info): allow paths that begin with /kythe_builtins/ (#5030)<commit_after>/*
 * Copyright 2019 The Kythe Authors. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

// Package info provides utilities for summarizing the contents of a kzip.
package info // import "kythe.io/kythe/go/platform/kzip/info"

import (
	"fmt"
	"log"
	"strings"

	"bitbucket.org/creachadair/stringset"

	"kythe.io/kythe/go/platform/kzip"

	apb "kythe.io/kythe/proto/analysis_go_proto"
)

// KzipInfo scans the kzip in f and counts contained files and units, giving a
// breakdown by corpus and language. It also records the size (in bytes) of the
// kzip specified by fileSize in the returned KzipInfo. This is a convenience
// method and thin wrapper over the Accumulator. If you need to do more than
// just calculate KzipInfo while doing a kzip.Scan(), you should use the
// Accumulator directly.
func KzipInfo(f kzip.File, fileSize int64, scanOpts ...kzip.ScanOption) (*apb.KzipInfo, error) {
	a := NewAccumulator(fileSize)
	if err := kzip.Scan(f, func(r *kzip.Reader, unit *kzip.Unit) error {
		a.Accumulate(unit)
		return nil
	}, scanOpts...); err != nil {
		return nil, fmt.Errorf("scanning kzip: %v", err)
	}
	return a.Get(), nil
}

// Accumulator is used to build a summary of a collection of compilation units.
// Usage:
//   a := NewAccumulator(fileSize)
//   a.Accumulate(unit) // call for each compilation unit
//   info := a.Get()    // get the resulting KzipInfo
type Accumulator struct {
	*apb.KzipInfo
}

// NewAccumulator creates a new Accumulator instance given the kzip fileSize (in
// bytes).
func NewAccumulator(fileSize int64) *Accumulator {
	return &Accumulator{
		KzipInfo: &apb.KzipInfo{
			Corpora: make(map[string]*apb.KzipInfo_CorpusInfo),
			Size:    fileSize,
		},
	}
}

// Accumulate should be called for each unit in the kzip so its counts can be
// recorded.
func (a *Accumulator) Accumulate(u *kzip.Unit) {
	srcs := stringset.New(u.Proto.SourceFile...)
	cuLang := u.Proto.GetVName().GetLanguage()
	if cuLang == "" {
		msg := fmt.Sprintf("CU does not specify a language %v", u.Proto.GetVName())
		a.KzipInfo.CriticalKzipErrors = append(a.KzipInfo.CriticalKzipErrors, msg)
		return
	}

	var srcCorpora stringset.Set
	srcsWithRI := stringset.New()
	for _, ri := range u.Proto.RequiredInput {
		if strings.HasPrefix(ri.GetVName().GetPath(), "/") && !strings.HasPrefix(ri.GetVName().GetPath(), "/kythe_builtins/") {
			a.KzipInfo.AbsolutePaths = append(a.KzipInfo.AbsolutePaths, ri.GetVName().GetPath())
		}

		riCorpus := requiredInputCorpus(u, ri)
		if riCorpus == "" {
			// Trim spaces to work around the fact that log("%v", proto) is inconsistent about trailing spaces in google3 vs open-source go.
			msg := strings.TrimSpace(fmt.Sprintf("unable to determine corpus for required_input %q in CU %v", ri.Info.Path, u.Proto.GetVName()))
			a.KzipInfo.CriticalKzipErrors = append(a.KzipInfo.CriticalKzipErrors, msg)
			return
		}
		requiredInputInfo(riCorpus, cuLang, a.KzipInfo).Count++
		if srcs.Contains(ri.Info.Path) {
			sourceInfo(riCorpus, cuLang, a.KzipInfo).Count++
			srcCorpora.Add(riCorpus)
			srcsWithRI.Add(ri.Info.Path)
		}
	}
	srcsWithoutRI := srcs.Diff(srcsWithRI)
	for path := range srcsWithoutRI {
		msg := fmt.Sprintf("source %q in CU %v doesn't have a required_input entry", path, u.Proto.GetVName())
		a.KzipInfo.CriticalKzipErrors = append(a.KzipInfo.CriticalKzipErrors, msg)
	}
	if srcCorpora.Len() != 1 {
		// This is a warning for now, but may become an error.
		log.Printf("Multiple corpora in unit. unit vname={%v}; src corpora=%v; srcs=%v", u.Proto.GetVName(), srcCorpora, u.Proto.SourceFile)
	}
}

// Get returns the final KzipInfo after info from each unit in the kzip has been
// accumulated.
func (a *Accumulator) Get() *apb.KzipInfo {
	return a.KzipInfo
}

// requiredInputCorpus computes the corpus for a required input. It follows the rules in the
// CompilationUnit proto comments in kythe/proto/analysis.proto that say that any
// required_input that does not set corpus in its VName should inherit corpus from the compilation
// unit's VName.
func requiredInputCorpus(u *kzip.Unit, ri *apb.CompilationUnit_FileInput) string {
	if c := ri.GetVName().GetCorpus(); c != "" {
		return c
	}
	return u.Proto.GetVName().GetCorpus()
}

// KzipInfoTotalCount returns the total CompilationUnits counts for infos split apart by language.
func KzipInfoTotalCount(infos []*apb.KzipInfo) apb.KzipInfo_CorpusInfo {
	totals := apb.KzipInfo_CorpusInfo{
		LanguageRequiredInputs: make(map[string]*apb.KzipInfo_CorpusInfo_RequiredInputs),
		LanguageSources:        make(map[string]*apb.KzipInfo_CorpusInfo_RequiredInputs),
	}
	for _, info := range infos {
		for _, i := range info.GetCorpora() {
			for lang, stats := range i.GetLanguageRequiredInputs() {
				total := totals.LanguageRequiredInputs[lang]
				if total == nil {
					total = &apb.KzipInfo_CorpusInfo_RequiredInputs{}
					totals.LanguageRequiredInputs[lang] = total
				}
				total.Count += stats.GetCount()
			}
			for lang, stats := range i.GetLanguageSources() {
				total := totals.LanguageSources[lang]
				if total == nil {
					total = &apb.KzipInfo_CorpusInfo_RequiredInputs{}
					totals.LanguageSources[lang] = total
				}
				total.Count += stats.GetCount()
			}
		}
	}
	return totals
}

// MergeKzipInfo combines the counts from multiple KzipInfos.
func MergeKzipInfo(infos []*apb.KzipInfo) *apb.KzipInfo {
	kzipInfo := &apb.KzipInfo{Corpora: make(map[string]*apb.KzipInfo_CorpusInfo)}

	for _, i := range infos {
		for corpus, cinfo := range i.GetCorpora() {
			for lang, inputs := range cinfo.GetLanguageRequiredInputs() {
				c := requiredInputInfo(corpus, lang, kzipInfo)
				c.Count += inputs.GetCount()
			}
			for lang, sources := range cinfo.GetLanguageSources() {
				c := sourceInfo(corpus, lang, kzipInfo)
				c.Count += sources.GetCount()
			}
		}
		kzipInfo.CriticalKzipErrors = append(kzipInfo.GetCriticalKzipErrors(), i.GetCriticalKzipErrors()...)
		kzipInfo.Size += i.Size
		kzipInfo.AbsolutePaths = append(kzipInfo.AbsolutePaths, i.AbsolutePaths...)
	}
	return kzipInfo
}

func requiredInputInfo(corpus, lang string, kzipInfo *apb.KzipInfo) *apb.KzipInfo_CorpusInfo_RequiredInputs {
	c := corpusInfo(corpus, kzipInfo)
	lri := c.LanguageRequiredInputs[lang]
	if lri == nil {
		lri = &apb.KzipInfo_CorpusInfo_RequiredInputs{}
		c.LanguageRequiredInputs[lang] = lri
	}
	return lri
}

func sourceInfo(corpus, lang string, kzipInfo *apb.KzipInfo) *apb.KzipInfo_CorpusInfo_RequiredInputs {
	c := corpusInfo(corpus, kzipInfo)
	ls := c.LanguageSources[lang]
	if ls == nil {
		ls = &apb.KzipInfo_CorpusInfo_RequiredInputs{}
		c.LanguageSources[lang] = ls
	}
	return ls
}

func corpusInfo(corpus string, kzipInfo *apb.KzipInfo) *apb.KzipInfo_CorpusInfo {
	i := kzipInfo.GetCorpora()[corpus]
	if i == nil {
		i = &apb.KzipInfo_CorpusInfo{
			LanguageRequiredInputs: make(map[string]*apb.KzipInfo_CorpusInfo_RequiredInputs),
			LanguageSources:        make(map[string]*apb.KzipInfo_CorpusInfo_RequiredInputs),
		}
		kzipInfo.Corpora[corpus] = i
	}
	return i
}
<|endoftext|>


<commit_before>package gridfs

import (
	"bytes"
	"crypto/md5"
	"fmt"
	"github.com/eclark/exl/bson"
	"github.com/eclark/exl/bson/bsoncompat"
	"github.com/eclark/gomongo/mongo"
	"hash"
	"os"
	"time"
)

const (
	filesSuffix  = ".files"
	chunksSuffix = ".chunks"
	defaultPrefix = "fs"
	defaultChunksize = 256 * 1024
)

type File struct {
	Id bson.Element
	Chunksize int
	Size int64
	Nchunk int

	buf []byte
	nextc int
	pos int

	md5 hash.Hash
	db *mongo.Database
	prefix string
	filename string
}

func Open(filename string, db *mongo.Database, prefix string) (*File, os.Error) {
	query, err := mongo.Marshal(map[string]string{"filename": filename})
	if err != nil {
		return nil, err
	}

	file := new(File)
	file.db = db
	if prefix == "" {
		file.prefix = defaultPrefix
	} else {
		file.prefix = prefix
	}
	filem, err := db.GetCollection(file.prefix + filesSuffix).FindOne(query)
	if err != nil {
		return nil, err
	}

	switch filem.Get("length").Kind() {
	case mongo.IntKind:
		file.Size = int64(filem.Get("length").Int())
	case mongo.LongKind:
		file.Size = filem.Get("length").Long()
	default:
		return nil, os.NewError("No length for file!")
	}

	file.Chunksize = int(filem.Get("chunkSize").Int())
	file.Nchunk = int(file.Size / int64(file.Chunksize))
	if file.Size % int64(file.Chunksize) > 0 {
		file.Nchunk++
	}

	idbuf := bytes.NewBuffer(filem.Get("id_").Bytes())

	file.Id = new(bson.ObjectId)
	file.Id.ReadFrom(idbuf)

	return file, nil
}

func (f *File) Read(b []byte) (n int, err os.Error) {
	for {
		switch {
		case len(b) == 0:
			return
		case len(f.buf) > 0:
			m := copy(b, f.buf)
			f.pos += m
			n += m
			b = b[m:]
			f.buf = f.buf[m:]
		case f.nextc < f.Nchunk:
			query, err := mongo.Marshal(map[string]interface{}{"files_id": bsoncompat.Wrap(f.Id), "n": int32(f.nextc)})
			if err != nil {
				return n, err
			}

			chunkm, err := f.db.GetCollection(f.prefix + chunksSuffix).FindOne(query)
			if err != nil {
				return n, err
			}

			f.buf = chunkm.Get("data").Binary()
			f.nextc++
		case f.nextc == f.Nchunk:
			err = os.EOF
			return
		default:
			panic("should never be reached!")
		}
	}
	return
}

func New(filename string, db *mongo.Database, prefix string) (file *File, err os.Error) {
	file = new(File)
	file.filename = filename
	file.db = db
	if prefix == "" {
		file.prefix = defaultPrefix
	} else {
		file.prefix = prefix
	}

	file.Chunksize = defaultChunksize
	file.Id, err = bson.NewObjectId() //mongo.NewOID()
	if err != nil {
		return
	}
	file.buf = make([]byte, 0, file.Chunksize)
	file.md5 = md5.New()

	return
}

func (f *File) Write(b []byte) (n int, err os.Error) {
	for {
		switch {
			case len(b) == 0:
				return
			case len(f.buf) < cap(f.buf):
				// copy some bytes from b to f.buf
				amt := cap(f.buf) - len(f.buf)
				if amt > len(b) {
					amt = len(b)
				}
				pe := len(f.buf)
				f.buf = f.buf[:pe+amt]
				m := copy(f.buf[pe:], b[:amt])
				n += m
				b = b[m:]
			case len(f.buf) == f.Chunksize:
				// insert chunk in mongo
				f.writeChunk()
		}
	}
	return
}

func (f *File) writeChunk() os.Error {
	chunk := new(bson.Document)
	chunk.Append("files_id", f.Id)
	chunk_n := bson.Int32(f.nextc)
	chunk.Append("n", &chunk_n)
	chunk.Append("data", &bson.Binary{bson.GenericType, f.buf})

	err := f.db.GetCollection(f.prefix + chunksSuffix).Insert(bsoncompat.Wrap(chunk))
	if err != nil {
		return err
	}

	f.md5.Write(f.buf)
	f.nextc++
	f.pos += len(f.buf)
	f.buf = f.buf[0:0]

	return nil
}

type filemd5Query struct {
	filemd5 mongo.BSON
	root string
}

func (f *File) Close() os.Error {
	f.writeChunk()

	md5cmddoc := new(bson.Document)
	pre := bson.String(f.prefix)
	md5cmddoc.Append("filemd5", f.Id)
	md5cmddoc.Append("root", &pre)
	testelem := bson.Int32(1)
	md5cmddoc.Append("aaaa", &testelem)

	res, err := f.db.Command(bsoncompat.Wrap(md5cmddoc))
	if err != nil {
		return err
	}
	fmt.Println(res.Bytes())
	fmt.Println(res.Get("errmsg"))

	file := new(bson.Document)
	file.Append("_id", f.Id)
	filename_e := bson.String(f.filename)
	file.Append("filename", &filename_e)
	chunksize_e := bson.Int32(f.Chunksize)
	file.Append("chunkSize", &chunksize_e)
	time_e := bson.Time(time.Nanoseconds() / 1000)
	file.Append("uploadDate", &time_e)
	md5_e := bson.String(res.Get("md5").String())

	file.Append("md5", &md5_e)
	length_e := bson.Int64(f.pos)
	file.Append("length", &length_e)

	err = f.db.GetCollection(f.prefix + filesSuffix).Insert(bsoncompat.Wrap(file))
	if err != nil {
		return err
	}

	return nil
}
<commit_msg>Remove debug junk.<commit_after>package gridfs

import (
	"bytes"
	"crypto/md5"
	"github.com/eclark/exl/bson"
	"github.com/eclark/exl/bson/bsoncompat"
	"github.com/eclark/gomongo/mongo"
	"hash"
	"os"
	"time"
)

const (
	filesSuffix  = ".files"
	chunksSuffix = ".chunks"
	defaultPrefix = "fs"
	defaultChunksize = 256 * 1024
)

type File struct {
	Id bson.Element
	Chunksize int
	Size int64
	Nchunk int

	buf []byte
	nextc int
	pos int

	md5 hash.Hash
	db *mongo.Database
	prefix string
	filename string
}

func Open(filename string, db *mongo.Database, prefix string) (*File, os.Error) {
	query, err := mongo.Marshal(map[string]string{"filename": filename})
	if err != nil {
		return nil, err
	}

	file := new(File)
	file.db = db
	if prefix == "" {
		file.prefix = defaultPrefix
	} else {
		file.prefix = prefix
	}
	filem, err := db.GetCollection(file.prefix + filesSuffix).FindOne(query)
	if err != nil {
		return nil, err
	}

	switch filem.Get("length").Kind() {
	case mongo.IntKind:
		file.Size = int64(filem.Get("length").Int())
	case mongo.LongKind:
		file.Size = filem.Get("length").Long()
	default:
		return nil, os.NewError("No length for file!")
	}

	file.Chunksize = int(filem.Get("chunkSize").Int())
	file.Nchunk = int(file.Size / int64(file.Chunksize))
	if file.Size % int64(file.Chunksize) > 0 {
		file.Nchunk++
	}

	idbuf := bytes.NewBuffer(filem.Get("id_").Bytes())

	file.Id = new(bson.ObjectId)
	file.Id.ReadFrom(idbuf)

	return file, nil
}

func (f *File) Read(b []byte) (n int, err os.Error) {
	for {
		switch {
		case len(b) == 0:
			return
		case len(f.buf) > 0:
			m := copy(b, f.buf)
			f.pos += m
			n += m
			b = b[m:]
			f.buf = f.buf[m:]
		case f.nextc < f.Nchunk:
			query, err := mongo.Marshal(map[string]interface{}{"files_id": bsoncompat.Wrap(f.Id), "n": int32(f.nextc)})
			if err != nil {
				return n, err
			}

			chunkm, err := f.db.GetCollection(f.prefix + chunksSuffix).FindOne(query)
			if err != nil {
				return n, err
			}

			f.buf = chunkm.Get("data").Binary()
			f.nextc++
		case f.nextc == f.Nchunk:
			err = os.EOF
			return
		default:
			panic("should never be reached!")
		}
	}
	return
}

func New(filename string, db *mongo.Database, prefix string) (file *File, err os.Error) {
	file = new(File)
	file.filename = filename
	file.db = db
	if prefix == "" {
		file.prefix = defaultPrefix
	} else {
		file.prefix = prefix
	}

	file.Chunksize = defaultChunksize
	file.Id, err = bson.NewObjectId() //mongo.NewOID()
	if err != nil {
		return
	}
	file.buf = make([]byte, 0, file.Chunksize)
	file.md5 = md5.New()

	return
}

func (f *File) Write(b []byte) (n int, err os.Error) {
	for {
		switch {
			case len(b) == 0:
				return
			case len(f.buf) < cap(f.buf):
				// copy some bytes from b to f.buf
				amt := cap(f.buf) - len(f.buf)
				if amt > len(b) {
					amt = len(b)
				}
				pe := len(f.buf)
				f.buf = f.buf[:pe+amt]
				m := copy(f.buf[pe:], b[:amt])
				n += m
				b = b[m:]
			case len(f.buf) == f.Chunksize:
				// insert chunk in mongo
				f.writeChunk()
		}
	}
	return
}

func (f *File) writeChunk() os.Error {
	chunk := new(bson.Document)
	chunk.Append("files_id", f.Id)
	chunk_n := bson.Int32(f.nextc)
	chunk.Append("n", &chunk_n)
	chunk.Append("data", &bson.Binary{bson.GenericType, f.buf})

	err := f.db.GetCollection(f.prefix + chunksSuffix).Insert(bsoncompat.Wrap(chunk))
	if err != nil {
		return err
	}

	f.md5.Write(f.buf)
	f.nextc++
	f.pos += len(f.buf)
	f.buf = f.buf[0:0]

	return nil
}

type filemd5Query struct {
	filemd5 mongo.BSON
	root string
}

func (f *File) Close() os.Error {
	f.writeChunk()

	md5cmddoc := new(bson.Document)
	pre := bson.String(f.prefix)
	md5cmddoc.Append("filemd5", f.Id)
	md5cmddoc.Append("root", &pre)
	testelem := bson.Int32(1)
	md5cmddoc.Append("aaaa", &testelem)

	res, err := f.db.Command(bsoncompat.Wrap(md5cmddoc))
	if err != nil {
		return err
	}
//	fmt.Println(res.Bytes())
//	fmt.Println(res.Get("errmsg"))

	file := new(bson.Document)
	file.Append("_id", f.Id)
	filename_e := bson.String(f.filename)
	file.Append("filename", &filename_e)
	chunksize_e := bson.Int32(f.Chunksize)
	file.Append("chunkSize", &chunksize_e)
	time_e := bson.Time(time.Nanoseconds() / 1000)
	file.Append("uploadDate", &time_e)
	md5_e := bson.String(res.Get("md5").String())

	file.Append("md5", &md5_e)
	length_e := bson.Int64(f.pos)
	file.Append("length", &length_e)

	err = f.db.GetCollection(f.prefix + filesSuffix).Insert(bsoncompat.Wrap(file))
	if err != nil {
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>// Copyright (c) 2015 Pagoda Box Inc
//
// This Source Code Form is subject to the terms of the Mozilla Public License, v.
// 2.0. If a copy of the MPL was not distributed with this file, You can obtain one
// at http://mozilla.org/MPL/2.0/.
//

package monitor

import (
	"github.com/golang/mock/gomock"
	"github.com/nanobox-io/yoke/config"
	"github.com/nanobox-io/yoke/state/mock"
	"os"
	"testing"
)

func TestSingle(test *testing.T) {
	ctrl := gomock.NewController(test)
	defer ctrl.Finish()

	me := mock_state.NewMockState(ctrl)
	other := mock_state.NewMockState(ctrl)

	perform := start(me, other, test)
	defer perform.Stop()

	me.EXPECT().SetDBRole("single").Return(nil)

	err := perform.Single()
	if err != nil {
		test.Log(err)
		test.FailNow()
	}

}

func TestActive(test *testing.T) {
	ctrl := gomock.NewController(test)
	defer ctrl.Finish()

	me := mock_state.NewMockState(ctrl)
	other := mock_state.NewMockState(ctrl)

	perform := start(me, other, test)
	defer perform.Stop()

	other.EXPECT().GetDataDir().Return("test", nil)
	other.EXPECT().Location().Return("127.0.0.1:1234")

	other.EXPECT().SetSynced(true).Return(nil)

	me.EXPECT().SetDBRole("active").Return(nil)

	err := perform.Active()
	if err != nil {
		test.Log(err)
		test.FailNow()
	}

}

func TestBackup(test *testing.T) {
	ctrl := gomock.NewController(test)
	defer ctrl.Finish()

	me := mock_state.NewMockState(ctrl)
	other := mock_state.NewMockState(ctrl)

	perform := start(me, other, test)
	defer perform.Stop()

	me.EXPECT().HasSynced().Return(false, nil)
	me.EXPECT().HasSynced().Return(true, nil)
	me.EXPECT().SetDBRole("backup").Return(nil)

	err := perform.Backup()
	if err != nil {
		test.Log(err)
		test.FailNow()
	}

}

func start(me, other *mock_state.MockState, test *testing.T) *performer {
	// i need to create a tmp folder
	// start a action
	// then have it transition to different states.

	// overwrite the builtin config options
	config.Conf = config.Config{
		PGPort:      4567,
		DataDir:     os.TempDir() + "postgres/",
		StatusDir:   os.TempDir() + "postgres/",
		SyncCommand: "true",
		SystemUser:  config.SystemUser(),
	}

	perform := NewPerformer(me, other, config.Conf)

	// ignore all errors that come across this way
	go func() {
		<-perform.err
	}()

	err := perform.Initialize()
	if err != nil {
		test.Log(err)
		test.FailNow()
	}

	if err := config.ConfigureHBAConf("127.0.0.1"); err != nil {
		test.Log(err)
		test.FailNow()
	}
	if err := config.ConfigurePGConf("127.0.0.1", 5432); err != nil {
		test.Log(err)
		test.FailNow()
	}

	err = perform.Start()
	if err != nil {
		test.Log(err)
		test.FailNow()
	}

	return perform
}
<commit_msg>really need to use the right port<commit_after>// Copyright (c) 2015 Pagoda Box Inc
//
// This Source Code Form is subject to the terms of the Mozilla Public License, v.
// 2.0. If a copy of the MPL was not distributed with this file, You can obtain one
// at http://mozilla.org/MPL/2.0/.
//

package monitor

import (
	"github.com/golang/mock/gomock"
	"github.com/nanobox-io/yoke/config"
	"github.com/nanobox-io/yoke/state/mock"
	"os"
	"testing"
)

func TestSingle(test *testing.T) {
	ctrl := gomock.NewController(test)
	defer ctrl.Finish()

	me := mock_state.NewMockState(ctrl)
	other := mock_state.NewMockState(ctrl)

	perform := start(me, other, test)
	defer perform.Stop()

	me.EXPECT().SetDBRole("single").Return(nil)

	err := perform.Single()
	if err != nil {
		test.Log(err)
		test.FailNow()
	}

}

func TestActive(test *testing.T) {
	ctrl := gomock.NewController(test)
	defer ctrl.Finish()

	me := mock_state.NewMockState(ctrl)
	other := mock_state.NewMockState(ctrl)

	perform := start(me, other, test)
	defer perform.Stop()

	other.EXPECT().GetDataDir().Return("test", nil)
	other.EXPECT().Location().Return("127.0.0.1:1234")

	other.EXPECT().SetSynced(true).Return(nil)

	me.EXPECT().SetDBRole("active").Return(nil)

	err := perform.Active()
	if err != nil {
		test.Log(err)
		test.FailNow()
	}

}

func TestBackup(test *testing.T) {
	ctrl := gomock.NewController(test)
	defer ctrl.Finish()

	me := mock_state.NewMockState(ctrl)
	other := mock_state.NewMockState(ctrl)

	perform := start(me, other, test)
	defer perform.Stop()

	me.EXPECT().HasSynced().Return(false, nil)
	me.EXPECT().HasSynced().Return(true, nil)
	me.EXPECT().SetDBRole("backup").Return(nil)

	err := perform.Backup()
	if err != nil {
		test.Log(err)
		test.FailNow()
	}

}

func start(me, other *mock_state.MockState, test *testing.T) *performer {
	// i need to create a tmp folder
	// start a action
	// then have it transition to different states.

	// overwrite the builtin config options
	config.Conf = config.Config{
		PGPort:      4567,
		DataDir:     os.TempDir() + "postgres/",
		StatusDir:   os.TempDir() + "postgres/",
		SyncCommand: "true",
		SystemUser:  config.SystemUser(),
	}

	perform := NewPerformer(me, other, config.Conf)

	// ignore all errors that come across this way
	go func() {
		<-perform.err
	}()

	err := perform.Initialize()
	if err != nil {
		test.Log(err)
		test.FailNow()
	}

	if err := config.ConfigureHBAConf("127.0.0.1"); err != nil {
		test.Log(err)
		test.FailNow()
	}
	if err := config.ConfigurePGConf("127.0.0.1", config.Conf.PGPort); err != nil {
		test.Log(err)
		test.FailNow()
	}

	err = perform.Start()
	if err != nil {
		test.Log(err)
		test.FailNow()
	}

	return perform
}
<|endoftext|>


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

package gcscaching_test

import (
	"testing"
	"time"

	"golang.org/x/net/context"

	"github.com/googlecloudplatform/gcsfuse/timeutil"
	"github.com/jacobsa/gcloud/gcs"
	"github.com/jacobsa/gcloud/gcs/gcscaching"
	"github.com/jacobsa/gcloud/gcs/gcsfake"
	"github.com/jacobsa/gcloud/gcs/gcsutil"
	. "github.com/jacobsa/oglematchers"
	. "github.com/jacobsa/ogletest"
)

func TestIntegration(t *testing.T) { RunTests(t) }

////////////////////////////////////////////////////////////////////////
// Boilerplate
////////////////////////////////////////////////////////////////////////

type IntegrationTest struct {
	ctx context.Context

	cache   gcscaching.StatCache
	clock   timeutil.SimulatedClock
	wrapped gcs.Bucket

	bucket gcs.Bucket
}

func init() { RegisterTestSuite(&IntegrationTest{}) }

func (t *IntegrationTest) SetUp(ti *TestInfo) {
	t.ctx = context.Background()

	// Set up a fixed, non-zero time.
	t.clock.SetTime(time.Date(2015, 4, 5, 2, 15, 0, 0, time.Local))

	// Set up dependencies.
	const cacheCapacity = 100
	t.cache = gcscaching.NewStatCache(cacheCapacity)
	t.wrapped = gcsfake.NewFakeBucket(&t.clock, "some_bucket")

	t.bucket = gcscaching.NewFastStatBucket(
		ttl,
		t.cache,
		&t.clock,
		t.wrapped)
}

func (t *IntegrationTest) stat(name string) (o *gcs.Object, err error) {
	req := &gcs.StatObjectRequest{
		Name: name,
	}

	o, err = t.bucket.StatObject(t.ctx, req)
	return
}

////////////////////////////////////////////////////////////////////////
// Test functions
////////////////////////////////////////////////////////////////////////

func (t *IntegrationTest) StatDoesntCacheNotFoundErrors() {
	const name = "taco"
	var err error

	// Stat an unknown object.
	_, err = t.stat(name)
	AssertThat(err, HasSameTypeAs(&gcs.NotFoundError{}))

	// Create the object through the back door.
	_, err = gcsutil.CreateObject(t.ctx, t.wrapped, name, "")
	AssertEq(nil, err)

	// Stat again. We should now see the object.
	o, err := t.stat(name)
	AssertEq(nil, err)
	ExpectNe(nil, o)
}

func (t *IntegrationTest) CreateInsertsIntoCache() {
	const name = "taco"
	var err error

	// Create an object.
	_, err = gcsutil.CreateObject(t.ctx, t.bucket, name, "")
	AssertEq(nil, err)

	// Delete it through the back door.
	err = t.wrapped.DeleteObject(t.ctx, name)
	AssertEq(nil, err)

	// StatObject should still see it.
	o, err := t.stat(name)
	AssertEq(nil, err)
	ExpectNe(nil, o)
}

func (t *IntegrationTest) StatInsertsIntoCache() {
	const name = "taco"
	var err error

	// Create an object through the back door.
	_, err = gcsutil.CreateObject(t.ctx, t.wrapped, name, "")
	AssertEq(nil, err)

	// Stat it so that it's in cache.
	_, err = t.stat(name)
	AssertEq(nil, err)

	// Delete it through the back door.
	err = t.wrapped.DeleteObject(t.ctx, name)
	AssertEq(nil, err)

	// StatObject should still see it.
	o, err := t.stat(name)
	AssertEq(nil, err)
	ExpectNe(nil, o)
}

func (t *IntegrationTest) ListInsertsIntoCache() {
	const name = "taco"
	var err error

	// Create an object through the back door.
	_, err = gcsutil.CreateObject(t.ctx, t.wrapped, name, "")
	AssertEq(nil, err)

	// List so that it's in cache.
	_, err = t.bucket.ListObjects(t.ctx, &gcs.ListObjectsRequest{})
	AssertEq(nil, err)

	// Delete the object through the back door.
	err = t.wrapped.DeleteObject(t.ctx, name)
	AssertEq(nil, err)

	// StatObject should still see it.
	o, err := t.stat(name)
	AssertEq(nil, err)
	ExpectNe(nil, o)
}

func (t *IntegrationTest) UpdateUpdatesCache() {
	const name = "taco"
	var err error

	// Create an object through the back door.
	_, err = gcsutil.CreateObject(t.ctx, t.wrapped, name, "")
	AssertEq(nil, err)

	// Update it, putting the new version in cache.
	updateReq := &gcs.UpdateObjectRequest{
		Name: name,
	}

	_, err = t.bucket.UpdateObject(t.ctx, updateReq)
	AssertEq(nil, err)

	// Delete the object through the back door.
	err = t.wrapped.DeleteObject(t.ctx, name)
	AssertEq(nil, err)

	// StatObject should still see it.
	o, err := t.stat(name)
	AssertEq(nil, err)
	ExpectNe(nil, o)
}

func (t *IntegrationTest) PositiveCacheExpiration() {
	const name = "taco"
	var err error

	// Create an object.
	_, err = gcsutil.CreateObject(t.ctx, t.bucket, name, "")
	AssertEq(nil, err)

	// Delete it through the back door.
	err = t.wrapped.DeleteObject(t.ctx, name)
	AssertEq(nil, err)

	// Advance time.
	t.clock.AdvanceTime(ttl + time.Millisecond)

	// StatObject should no longer see it.
	_, err = t.stat(name)
	ExpectThat(err, HasSameTypeAs(&gcs.NotFoundError{}))
}

func (t *IntegrationTest) CreateInvalidatesNegativeCache() {
	const name = "taco"
	var err error

	// Stat an unknown object, getting it into the negative cache.
	_, err = t.stat(name)
	AssertThat(err, HasSameTypeAs(&gcs.NotFoundError{}))

	// Create the object.
	_, err = gcsutil.CreateObject(t.ctx, t.bucket, name, "")
	AssertEq(nil, err)

	// Now StatObject should see it.
	o, err := t.stat(name)
	AssertEq(nil, err)
	ExpectNe(nil, o)
}

func (t *IntegrationTest) StatAddsToNegativeCache() {
	const name = "taco"
	var err error

	// Stat an unknown object, getting it into the negative cache.
	_, err = t.stat(name)
	AssertThat(err, HasSameTypeAs(&gcs.NotFoundError{}))

	// Create the object through the back door.
	_, err = gcsutil.CreateObject(t.ctx, t.wrapped, name, "")
	AssertEq(nil, err)

	// StatObject should still not see it yet.
	_, err = t.stat(name)
	ExpectThat(err, HasSameTypeAs(&gcs.NotFoundError{}))
}

func (t *IntegrationTest) ListInvalidatesNegativeCache() {
	AssertFalse(true, "TODO")
}

func (t *IntegrationTest) UpdateInvalidatesNegativeCache() {
	AssertFalse(true, "TODO")
}

func (t *IntegrationTest) DeleteAddsToNegativeCache() {
	AssertFalse(true, "TODO")
}

func (t *IntegrationTest) NegativeCacheExpiration() {
	AssertFalse(true, "TODO")
}
<commit_msg>IntegrationTest.ListInvalidatesNegativeCache<commit_after>// Copyright 2015 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package gcscaching_test

import (
	"testing"
	"time"

	"golang.org/x/net/context"

	"github.com/googlecloudplatform/gcsfuse/timeutil"
	"github.com/jacobsa/gcloud/gcs"
	"github.com/jacobsa/gcloud/gcs/gcscaching"
	"github.com/jacobsa/gcloud/gcs/gcsfake"
	"github.com/jacobsa/gcloud/gcs/gcsutil"
	. "github.com/jacobsa/oglematchers"
	. "github.com/jacobsa/ogletest"
)

func TestIntegration(t *testing.T) { RunTests(t) }

////////////////////////////////////////////////////////////////////////
// Boilerplate
////////////////////////////////////////////////////////////////////////

type IntegrationTest struct {
	ctx context.Context

	cache   gcscaching.StatCache
	clock   timeutil.SimulatedClock
	wrapped gcs.Bucket

	bucket gcs.Bucket
}

func init() { RegisterTestSuite(&IntegrationTest{}) }

func (t *IntegrationTest) SetUp(ti *TestInfo) {
	t.ctx = context.Background()

	// Set up a fixed, non-zero time.
	t.clock.SetTime(time.Date(2015, 4, 5, 2, 15, 0, 0, time.Local))

	// Set up dependencies.
	const cacheCapacity = 100
	t.cache = gcscaching.NewStatCache(cacheCapacity)
	t.wrapped = gcsfake.NewFakeBucket(&t.clock, "some_bucket")

	t.bucket = gcscaching.NewFastStatBucket(
		ttl,
		t.cache,
		&t.clock,
		t.wrapped)
}

func (t *IntegrationTest) stat(name string) (o *gcs.Object, err error) {
	req := &gcs.StatObjectRequest{
		Name: name,
	}

	o, err = t.bucket.StatObject(t.ctx, req)
	return
}

////////////////////////////////////////////////////////////////////////
// Test functions
////////////////////////////////////////////////////////////////////////

func (t *IntegrationTest) StatDoesntCacheNotFoundErrors() {
	const name = "taco"
	var err error

	// Stat an unknown object.
	_, err = t.stat(name)
	AssertThat(err, HasSameTypeAs(&gcs.NotFoundError{}))

	// Create the object through the back door.
	_, err = gcsutil.CreateObject(t.ctx, t.wrapped, name, "")
	AssertEq(nil, err)

	// Stat again. We should now see the object.
	o, err := t.stat(name)
	AssertEq(nil, err)
	ExpectNe(nil, o)
}

func (t *IntegrationTest) CreateInsertsIntoCache() {
	const name = "taco"
	var err error

	// Create an object.
	_, err = gcsutil.CreateObject(t.ctx, t.bucket, name, "")
	AssertEq(nil, err)

	// Delete it through the back door.
	err = t.wrapped.DeleteObject(t.ctx, name)
	AssertEq(nil, err)

	// StatObject should still see it.
	o, err := t.stat(name)
	AssertEq(nil, err)
	ExpectNe(nil, o)
}

func (t *IntegrationTest) StatInsertsIntoCache() {
	const name = "taco"
	var err error

	// Create an object through the back door.
	_, err = gcsutil.CreateObject(t.ctx, t.wrapped, name, "")
	AssertEq(nil, err)

	// Stat it so that it's in cache.
	_, err = t.stat(name)
	AssertEq(nil, err)

	// Delete it through the back door.
	err = t.wrapped.DeleteObject(t.ctx, name)
	AssertEq(nil, err)

	// StatObject should still see it.
	o, err := t.stat(name)
	AssertEq(nil, err)
	ExpectNe(nil, o)
}

func (t *IntegrationTest) ListInsertsIntoCache() {
	const name = "taco"
	var err error

	// Create an object through the back door.
	_, err = gcsutil.CreateObject(t.ctx, t.wrapped, name, "")
	AssertEq(nil, err)

	// List so that it's in cache.
	_, err = t.bucket.ListObjects(t.ctx, &gcs.ListObjectsRequest{})
	AssertEq(nil, err)

	// Delete the object through the back door.
	err = t.wrapped.DeleteObject(t.ctx, name)
	AssertEq(nil, err)

	// StatObject should still see it.
	o, err := t.stat(name)
	AssertEq(nil, err)
	ExpectNe(nil, o)
}

func (t *IntegrationTest) UpdateUpdatesCache() {
	const name = "taco"
	var err error

	// Create an object through the back door.
	_, err = gcsutil.CreateObject(t.ctx, t.wrapped, name, "")
	AssertEq(nil, err)

	// Update it, putting the new version in cache.
	updateReq := &gcs.UpdateObjectRequest{
		Name: name,
	}

	_, err = t.bucket.UpdateObject(t.ctx, updateReq)
	AssertEq(nil, err)

	// Delete the object through the back door.
	err = t.wrapped.DeleteObject(t.ctx, name)
	AssertEq(nil, err)

	// StatObject should still see it.
	o, err := t.stat(name)
	AssertEq(nil, err)
	ExpectNe(nil, o)
}

func (t *IntegrationTest) PositiveCacheExpiration() {
	const name = "taco"
	var err error

	// Create an object.
	_, err = gcsutil.CreateObject(t.ctx, t.bucket, name, "")
	AssertEq(nil, err)

	// Delete it through the back door.
	err = t.wrapped.DeleteObject(t.ctx, name)
	AssertEq(nil, err)

	// Advance time.
	t.clock.AdvanceTime(ttl + time.Millisecond)

	// StatObject should no longer see it.
	_, err = t.stat(name)
	ExpectThat(err, HasSameTypeAs(&gcs.NotFoundError{}))
}

func (t *IntegrationTest) CreateInvalidatesNegativeCache() {
	const name = "taco"
	var err error

	// Stat an unknown object, getting it into the negative cache.
	_, err = t.stat(name)
	AssertThat(err, HasSameTypeAs(&gcs.NotFoundError{}))

	// Create the object.
	_, err = gcsutil.CreateObject(t.ctx, t.bucket, name, "")
	AssertEq(nil, err)

	// Now StatObject should see it.
	o, err := t.stat(name)
	AssertEq(nil, err)
	ExpectNe(nil, o)
}

func (t *IntegrationTest) StatAddsToNegativeCache() {
	const name = "taco"
	var err error

	// Stat an unknown object, getting it into the negative cache.
	_, err = t.stat(name)
	AssertThat(err, HasSameTypeAs(&gcs.NotFoundError{}))

	// Create the object through the back door.
	_, err = gcsutil.CreateObject(t.ctx, t.wrapped, name, "")
	AssertEq(nil, err)

	// StatObject should still not see it yet.
	_, err = t.stat(name)
	ExpectThat(err, HasSameTypeAs(&gcs.NotFoundError{}))
}

func (t *IntegrationTest) ListInvalidatesNegativeCache() {
	const name = "taco"
	var err error

	// Stat an unknown object, getting it into the negative cache.
	_, err = t.stat(name)
	AssertThat(err, HasSameTypeAs(&gcs.NotFoundError{}))

	// Create the object through the back door.
	_, err = gcsutil.CreateObject(t.ctx, t.wrapped, name, "")
	AssertEq(nil, err)

	// List the bucket.
	_, err = t.bucket.ListObjects(t.ctx, &gcs.ListObjectsRequest{})
	AssertEq(nil, err)

	// Now StatObject should see it.
	o, err := t.stat(name)
	AssertEq(nil, err)
	ExpectNe(nil, o)
}

func (t *IntegrationTest) UpdateInvalidatesNegativeCache() {
	AssertFalse(true, "TODO")
}

func (t *IntegrationTest) DeleteAddsToNegativeCache() {
	AssertFalse(true, "TODO")
}

func (t *IntegrationTest) NegativeCacheExpiration() {
	AssertFalse(true, "TODO")
}
<|endoftext|>


<commit_before>package main

import (
	"database/sql" // MySQL Query
	"flag"         // Command line parsing
	"fmt"          // Output formatting

	_ "github.com/go-sql-driver/mysql" // MySQL connection
	"github.com/koding/logging"        // logging
	"gopkg.in/ini.v1"                  // ini file parsing
	"os"                               // to exit with exitcode
	"strconv"                          // string conversion
	"strings"                          // string manipulation
	"time"                             // timestamp logging, ticker

	"github.com/cactus/go-statsd-client/statsd" // Statsd client
	_ "github.com/go-sql-driver/mysql"          // MySQL connection
	"github.com/koding/logging"                 // logging
	"gopkg.in/ini.v1"                           // ini file parsing
)

var logger = logging.NewLogger("Mambo")

/*
  Configuration parameters, mysql & statsd
*/
type configuration struct {
	mysqlHost  string // MySQL host to connect, if empty local socket will be used
	mysqlUser  string // User to connect MySQL with
	mysqlPass  string // Password for connecting MySQL
	mysqlDb    string // Database to connect to
	mysqlPort  int    // Port to connect MySQL, if left blank, 3306 will be used as default
	statsdHost string // statsd server hostname
	statsdPort int    // statsd server port, if left blank, 8125 will be used as default
}

/*
  Commands
*/
type command struct {
	key   string // key to send statsd server (eg. mysql.slave01.bfc.kinja-ops.com.replication lag)
	query string // query to run against mysql server. The output must be an integer
	freq  int    // what frequency the query should be run in milliseconds
}

func main() {
	// command line parameter parsing
	configfile := flag.String("cfg", "mambo.cfg", "Main configuration file")
	flag.Parse()
	logger := logging.NewLogger("Mambo")
	logger.Notice("Mambo collector started")
	logger.Notice("Loading configuration from %s", *configfile)
	// The 'results' channel will recive the results of the mysqlWorker queries
	results := make(chan string)
	config, commands := configure(*configfile) // Loading configuration and commands from ini file
	for _, command := range commands {
		go controller(command, config, results) // every command will launch a command controller
	}
	logger.Notice("Data collector running")
	for {
		select {
		// every time a MySQL worker yield data to the 'results' channel we call a statsdSender and we send that data to statsdserver
		case msg := <-results:
			{
				statsdSender(config, msg)
			}
		}
	}
}

/*
  The controller reads the command frequency (rate) from the command, and sets up
  a ticker with that frequency. We wait for the tick, and when it happens, we call
  a mysqlWorker with the command
*/
func controller(cmd command, cnf *configuration, results chan string) {
	logger.Notice("Query loaded: %s", cmd.query)
	tick := time.NewTicker(time.Millisecond * time.Duration(cmd.freq)).C // I have to convert freq to time.Duration to use with ticker
	for {
		select {
		case <-tick:
			mysqlWorker(cnf, cmd, results)
		}
	}
}

/*
  Builds up the statsd connect uri from
  statsdHost and statsdPort parameters
  For example:
  statsdHost = graphstatsdPort = 8125 -> url:"graph:8125"
*/
func statsdURIBuilder(config *configuration) string {
	uri := fmt.Sprint(config.statsdHost, ":", config.statsdPort)
	return uri

}

/*
  Connects statsd server and sends the metric
*/
func statsdSender(config *configuration, msg string) {
	client, err := statsd.NewClient(statsdURIBuilder(config), "")
	if err != nil {
		logger.Error(err.Error())
	}
	defer client.Close()
	arr := strings.Split(msg, ":")
	key := arr[0]
	value, err := strconv.ParseInt(arr[1], 10, 64)
	if err != nil {
		logger.Error(err.Error())
	}
	//	logger.Info("Statsd data flushed: %s", msg)
	err = client.Inc(key, value, 1.0)
	if err != nil {
		logger.Error(err.Error())
	}
}

/*
  The mysqlWorker function connects to the database, runs the query which came from the command
  and puts the result to the results channel
*/
func mysqlWorker(config *configuration, cmd command, results chan string) {
	var result string
	connecturi := mysqlURIBuilder(config)
	db, err := sql.Open("mysql", connecturi)
	if err != nil {
		logger.Error(err.Error())
	}
	defer db.Close()
	err = db.Ping()
	if err != nil {
		logger.Error(err.Error())
	}
	stmtOut, err := db.Prepare(cmd.query)
	if err != nil {
		logger.Error(err.Error())
	}
	defer stmtOut.Close()
	err = stmtOut.QueryRow().Scan(&result)
	if err != nil {
		logger.Error(err.Error())
	}
	res := fmt.Sprint(cmd.key, ":", result)
	//	logger.Info("Data recieved from MySQL server: %s", res)
	results <- res
}

/*
   Helper function to the mysqlWorker, it builds up the connect uri based on config
   if no mysqlHost is given, it tries to connect via local socket, and ignores the
   mysqlPort option.
*/
func mysqlURIBuilder(config *configuration) string {
	uri := ""
	if config.mysqlHost == "" { // if mysqlHost is not defined, we'll connect through local socket
		uri = fmt.Sprint(config.mysqlUser, ":", config.mysqlPass, "@", "/", config.mysqlDb)
	} else { // if we use TCP we'll also need the port of mysql too
		uri = fmt.Sprint(config.mysqlUser, ":", config.mysqlPass, "@", config.mysqlHost, ":", config.mysqlPort, "/", config.mysqlDb)
	}
	return uri
}

/*
  Builds up the configuration and command structs from the config file.
  It searches the [config] section for setting up configurations, and it
  assumes, that every other section will hold commands.
*/
func configure(cfgfile string) (*configuration, []command) {
	var mysqlPortc, statsdPortc int
	var cfg configuration
	//commands := make([]command, 0)
	var commands []command
	config, err := ini.Load(cfgfile)
	if err != nil {
		logger.Critical(err.Error())
		os.Exit(1)
	}
	sections := config.Sections()
	for _, section := range sections {
		if section.Name() != "DEFAULT" { //skip unnamed section
			if section.Name() == "config" { //[config] holds the configuratuin
				mysqlHostc := section.Key("mysql_host").String()
				mysqlUserc := section.Key("mysql_user").String()
				mysqlPassc := section.Key("mysql_pass").String()
				mysqlDbc := section.Key("mysql_db").String()
				// if mysqlPort is not defined, we'll assume that the default 3306 will be used
				mysqlPortc, err = section.Key("mysql_port").Int()
				if mysqlPortc == 0 {
					mysqlPortc = 3306
				}
				statsdHostc := section.Key("statsd_host").String()
				// if statsdPort is not defined, we'll assume that the default 8125 will be used
				statsdPortc, err = section.Key("stats_port").Int()
				if statsdPortc == 0 {
					statsdPortc = 8125
				}
				cfg = configuration{
					mysqlHost:  mysqlHostc,
					mysqlUser:  mysqlUserc,
					mysqlPass:  mysqlPassc,
					mysqlPort:  mysqlPortc,
					mysqlDb:    mysqlDbc,
					statsdHost: statsdHostc,
					statsdPort: statsdPortc,
				}
			} else { // here start the command parsing
				var cmd command
				keyc := section.Key("key").String()
				queryc := section.Key("query").String()
				freqc, _ := section.Key("freq").Int()
				cmd = command{
					key:   keyc,
					query: queryc,
					freq:  freqc,
				}
				commands = append(commands, cmd)
			}
		}
	}
	return &cfg, commands
}
<commit_msg>imports...<commit_after>package main

import (
	"database/sql" // MySQL Query
	"flag"         // Command line parsing
	"fmt"          // Output formatting

	_ "github.com/go-sql-driver/mysql" // MySQL connection
	"os"                               // to exit with exitcode
	"strconv"                          // string conversion
	"strings"                          // string manipulation
	"time"                             // timestamp logging, ticker

	"github.com/cactus/go-statsd-client/statsd" // Statsd client
	_ "github.com/go-sql-driver/mysql"          // MySQL connection
	"github.com/koding/logging"                 // logging
	"gopkg.in/ini.v1"                           // ini file parsing
)

var logger = logging.NewLogger("Mambo")

/*
  Configuration parameters, mysql & statsd
*/
type configuration struct {
	mysqlHost  string // MySQL host to connect, if empty local socket will be used
	mysqlUser  string // User to connect MySQL with
	mysqlPass  string // Password for connecting MySQL
	mysqlDb    string // Database to connect to
	mysqlPort  int    // Port to connect MySQL, if left blank, 3306 will be used as default
	statsdHost string // statsd server hostname
	statsdPort int    // statsd server port, if left blank, 8125 will be used as default
}

/*
  Commands
*/
type command struct {
	key   string // key to send statsd server (eg. mysql.slave01.bfc.kinja-ops.com.replication lag)
	query string // query to run against mysql server. The output must be an integer
	freq  int    // what frequency the query should be run in milliseconds
}

func main() {
	// command line parameter parsing
	configfile := flag.String("cfg", "mambo.cfg", "Main configuration file")
	flag.Parse()
	logger := logging.NewLogger("Mambo")
	logger.Notice("Mambo collector started")
	logger.Notice("Loading configuration from %s", *configfile)
	// The 'results' channel will recive the results of the mysqlWorker queries
	results := make(chan string)
	config, commands := configure(*configfile) // Loading configuration and commands from ini file
	for _, command := range commands {
		go controller(command, config, results) // every command will launch a command controller
	}
	logger.Notice("Data collector running")
	for {
		select {
		// every time a MySQL worker yield data to the 'results' channel we call a statsdSender and we send that data to statsdserver
		case msg := <-results:
			{
				statsdSender(config, msg)
			}
		}
	}
}

/*
  The controller reads the command frequency (rate) from the command, and sets up
  a ticker with that frequency. We wait for the tick, and when it happens, we call
  a mysqlWorker with the command
*/
func controller(cmd command, cnf *configuration, results chan string) {
	logger.Notice("Query loaded: %s", cmd.query)
	tick := time.NewTicker(time.Millisecond * time.Duration(cmd.freq)).C // I have to convert freq to time.Duration to use with ticker
	for {
		select {
		case <-tick:
			mysqlWorker(cnf, cmd, results)
		}
	}
}

/*
  Builds up the statsd connect uri from
  statsdHost and statsdPort parameters
  For example:
  statsdHost = graphstatsdPort = 8125 -> url:"graph:8125"
*/
func statsdURIBuilder(config *configuration) string {
	uri := fmt.Sprint(config.statsdHost, ":", config.statsdPort)
	return uri

}

/*
  Connects statsd server and sends the metric
*/
func statsdSender(config *configuration, msg string) {
	client, err := statsd.NewClient(statsdURIBuilder(config), "")
	if err != nil {
		logger.Error(err.Error())
	}
	defer client.Close()
	arr := strings.Split(msg, ":")
	key := arr[0]
	value, err := strconv.ParseInt(arr[1], 10, 64)
	if err != nil {
		logger.Error(err.Error())
	}
	//	logger.Info("Statsd data flushed: %s", msg)
	err = client.Inc(key, value, 1.0)
	if err != nil {
		logger.Error(err.Error())
	}
}

/*
  The mysqlWorker function connects to the database, runs the query which came from the command
  and puts the result to the results channel
*/
func mysqlWorker(config *configuration, cmd command, results chan string) {
	var result string
	connecturi := mysqlURIBuilder(config)
	db, err := sql.Open("mysql", connecturi)
	if err != nil {
		logger.Error(err.Error())
	}
	defer db.Close()
	err = db.Ping()
	if err != nil {
		logger.Error(err.Error())
	}
	stmtOut, err := db.Prepare(cmd.query)
	if err != nil {
		logger.Error(err.Error())
	}
	defer stmtOut.Close()
	err = stmtOut.QueryRow().Scan(&result)
	if err != nil {
		logger.Error(err.Error())
	}
	res := fmt.Sprint(cmd.key, ":", result)
	//	logger.Info("Data recieved from MySQL server: %s", res)
	results <- res
}

/*
   Helper function to the mysqlWorker, it builds up the connect uri based on config
   if no mysqlHost is given, it tries to connect via local socket, and ignores the
   mysqlPort option.
*/
func mysqlURIBuilder(config *configuration) string {
	uri := ""
	if config.mysqlHost == "" { // if mysqlHost is not defined, we'll connect through local socket
		uri = fmt.Sprint(config.mysqlUser, ":", config.mysqlPass, "@", "/", config.mysqlDb)
	} else { // if we use TCP we'll also need the port of mysql too
		uri = fmt.Sprint(config.mysqlUser, ":", config.mysqlPass, "@", config.mysqlHost, ":", config.mysqlPort, "/", config.mysqlDb)
	}
	return uri
}

/*
  Builds up the configuration and command structs from the config file.
  It searches the [config] section for setting up configurations, and it
  assumes, that every other section will hold commands.
*/
func configure(cfgfile string) (*configuration, []command) {
	var mysqlPortc, statsdPortc int
	var cfg configuration
	//commands := make([]command, 0)
	var commands []command
	config, err := ini.Load(cfgfile)
	if err != nil {
		logger.Critical(err.Error())
		os.Exit(1)
	}
	sections := config.Sections()
	for _, section := range sections {
		if section.Name() != "DEFAULT" { //skip unnamed section
			if section.Name() == "config" { //[config] holds the configuratuin
				mysqlHostc := section.Key("mysql_host").String()
				mysqlUserc := section.Key("mysql_user").String()
				mysqlPassc := section.Key("mysql_pass").String()
				mysqlDbc := section.Key("mysql_db").String()
				// if mysqlPort is not defined, we'll assume that the default 3306 will be used
				mysqlPortc, err = section.Key("mysql_port").Int()
				if mysqlPortc == 0 {
					mysqlPortc = 3306
				}
				statsdHostc := section.Key("statsd_host").String()
				// if statsdPort is not defined, we'll assume that the default 8125 will be used
				statsdPortc, err = section.Key("stats_port").Int()
				if statsdPortc == 0 {
					statsdPortc = 8125
				}
				cfg = configuration{
					mysqlHost:  mysqlHostc,
					mysqlUser:  mysqlUserc,
					mysqlPass:  mysqlPassc,
					mysqlPort:  mysqlPortc,
					mysqlDb:    mysqlDbc,
					statsdHost: statsdHostc,
					statsdPort: statsdPortc,
				}
			} else { // here start the command parsing
				var cmd command
				keyc := section.Key("key").String()
				queryc := section.Key("query").String()
				freqc, _ := section.Key("freq").Int()
				cmd = command{
					key:   keyc,
					query: queryc,
					freq:  freqc,
				}
				commands = append(commands, cmd)
			}
		}
	}
	return &cfg, commands
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	_ "fmt"
	"os"
	"os/exec"
)

var options = struct {
	Output  string
	Name    string
	Plain   bool
	Preview bool
}{}

func init() {
	flag.StringVar(&options.Output, "output", "", "write output to file")
	flag.StringVar(&options.Name, "name", "", "set command name")
	flag.BoolVar(&options.Plain, "plain", false, "treat comments as plain text")
	flag.BoolVar(&options.Preview, "preview", false, "preview output in man")
}

func getReader() Reader {
	if options.Plain {
		return NewPlainReader()
	} else {
		return NewMarkupReader()
	}
}

func getWriter() Writer {
	return NewTroffWriter()
}

func main() {
	flag.Parse()

	builder := NewBuilder(getReader(), getWriter())
	for _, arg := range flag.Args() {
		file, err := NewFile(arg)
		if err != nil {
			panic(err)
		}
		if options.Name != "" {
			file.Name = options.Name
		}
		text, err := builder.Build(file)
		if err != nil {
			panic(err)
		}
		if options.Preview {
			cmd := exec.Command("groff", "-Wall", "-mtty-char", "-mandoc", "-Tascii")
			cmd.Stdout = os.Stdout
			cmd.Stderr = os.Stderr
			inp, err := cmd.StdinPipe()
			if err != nil {
				panic(err)
			}
			if err := cmd.Start(); err != nil {
				panic(err)
			}
			inp.Write([]byte(text))
			inp.Close()
			cmd.Wait()
		} else {
			if options.Output == "" {
				os.Stdout.Write([]byte(text))
			} else {
				dst, err := os.Create(options.Output)
				if err != nil {
					panic(err)
				}
				dst.Write([]byte(text))
				dst.Close()
			}
		}
	}
}
<commit_msg>remove unused package<commit_after>package main

import (
	"flag"
	"os"
	"os/exec"
)

var options = struct {
	Output  string
	Name    string
	Plain   bool
	Preview bool
}{}

func init() {
	flag.StringVar(&options.Output, "output", "", "write output to file")
	flag.StringVar(&options.Name, "name", "", "set command name")
	flag.BoolVar(&options.Plain, "plain", false, "treat comments as plain text")
	flag.BoolVar(&options.Preview, "preview", false, "preview output in man")
}

func getReader() Reader {
	if options.Plain {
		return NewPlainReader()
	} else {
		return NewMarkupReader()
	}
}

func getWriter() Writer {
	return NewTroffWriter()
}

func main() {
	flag.Parse()

	builder := NewBuilder(getReader(), getWriter())
	for _, arg := range flag.Args() {
		file, err := NewFile(arg)
		if err != nil {
			panic(err)
		}
		if options.Name != "" {
			file.Name = options.Name
		}
		text, err := builder.Build(file)
		if err != nil {
			panic(err)
		}
		if options.Preview {
			cmd := exec.Command("groff", "-Wall", "-mtty-char", "-mandoc", "-Tascii")
			cmd.Stdout = os.Stdout
			cmd.Stderr = os.Stderr
			inp, err := cmd.StdinPipe()
			if err != nil {
				panic(err)
			}
			if err := cmd.Start(); err != nil {
				panic(err)
			}
			inp.Write([]byte(text))
			inp.Close()
			cmd.Wait()
		} else {
			if options.Output == "" {
				os.Stdout.Write([]byte(text))
			} else {
				dst, err := os.Create(options.Output)
				if err != nil {
					panic(err)
				}
				dst.Write([]byte(text))
				dst.Close()
			}
		}
	}
}
<|endoftext|>


<commit_before>// Copyright 2019 PingCAP, Inc.
//
// 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 executor

import (
	"context"
	"errors"
	"testing"

	"github.com/pingcap/tidb/expression"
	"github.com/pingcap/tidb/parser/mysql"
	"github.com/pingcap/tidb/sessionctx/variable"
	"github.com/pingcap/tidb/types"
	"github.com/pingcap/tidb/util/chunk"
	"github.com/pingcap/tidb/util/mock"
)

var (
	_ Executor = &mockErrorOperator{}
)

type mockErrorOperator struct {
	baseExecutor
	toPanic bool
	closed  bool
}

func (e *mockErrorOperator) Open(ctx context.Context) error {
	return nil
}

func (e *mockErrorOperator) Next(ctx context.Context, req *chunk.Chunk) error {
	if e.toPanic {
		panic("next panic")
	} else {
		return errors.New("next error")
	}
}

func (e *mockErrorOperator) Close() error {
	e.closed = true
	return errors.New("close error")
}

func getColumns() []*expression.Column {
	return []*expression.Column{
		{Index: 1, RetType: types.NewFieldType(mysql.TypeLonglong)},
	}
}

// close() must be called after next() to avoid goroutines leak
func TestExplainAnalyzeInvokeNextAndClose(t *testing.T) {
	ctx := mock.NewContext()
	ctx.GetSessionVars().InitChunkSize = variable.DefInitChunkSize
	ctx.GetSessionVars().MaxChunkSize = variable.DefMaxChunkSize
	schema := expression.NewSchema(getColumns()...)
	baseExec := newBaseExecutor(ctx, schema, 0)
	explainExec := &ExplainExec{
		baseExecutor: baseExec,
		explain:      nil,
	}
	// mockErrorOperator returns errors
	mockOpr := mockErrorOperator{baseExec, false, false}
	explainExec.analyzeExec = &mockOpr
	tmpCtx := context.Background()
	_, err := explainExec.generateExplainInfo(tmpCtx)

	expectedStr := "next error, close error"
	if err != nil && (err.Error() != expectedStr || !mockOpr.closed) {
		t.Errorf(err.Error())
	}
	// mockErrorOperator panic
	explainExec = &ExplainExec{
		baseExecutor: baseExec,
		explain:      nil,
	}
	mockOpr = mockErrorOperator{baseExec, true, false}
	explainExec.analyzeExec = &mockOpr
	defer func() {
		if panicErr := recover(); panicErr == nil || !mockOpr.closed {
			t.Errorf("panic test failed: without panic or close() is not called")
		}
	}()

	_, err = explainExec.generateExplainInfo(tmpCtx)
	if err != nil {
		t.Errorf(err.Error())
	}
}
<commit_msg>executor: migrate test-infra to testify for explain_unit_test.go (#32706)<commit_after>// Copyright 2019 PingCAP, Inc.
//
// 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 executor

import (
	"context"
	"errors"
	"testing"

	"github.com/pingcap/tidb/expression"
	"github.com/pingcap/tidb/parser/mysql"
	"github.com/pingcap/tidb/sessionctx/variable"
	"github.com/pingcap/tidb/types"
	"github.com/pingcap/tidb/util/chunk"
	"github.com/pingcap/tidb/util/mock"
	"github.com/stretchr/testify/require"
)

var (
	_ Executor = &mockErrorOperator{}
)

type mockErrorOperator struct {
	baseExecutor
	toPanic bool
	closed  bool
}

func (e *mockErrorOperator) Open(_ context.Context) error {
	return nil
}

func (e *mockErrorOperator) Next(_ context.Context, _ *chunk.Chunk) error {
	if e.toPanic {
		panic("next panic")
	} else {
		return errors.New("next error")
	}
}

func (e *mockErrorOperator) Close() error {
	e.closed = true
	return errors.New("close error")
}

func getColumns() []*expression.Column {
	return []*expression.Column{
		{Index: 1, RetType: types.NewFieldType(mysql.TypeLonglong)},
	}
}

// close() must be called after next() to avoid goroutines leak
func TestExplainAnalyzeInvokeNextAndClose(t *testing.T) {
	ctx := mock.NewContext()
	ctx.GetSessionVars().InitChunkSize = variable.DefInitChunkSize
	ctx.GetSessionVars().MaxChunkSize = variable.DefMaxChunkSize
	schema := expression.NewSchema(getColumns()...)
	baseExec := newBaseExecutor(ctx, schema, 0)
	explainExec := &ExplainExec{
		baseExecutor: baseExec,
		explain:      nil,
	}
	// mockErrorOperator returns errors
	mockOpr := mockErrorOperator{baseExec, false, false}
	explainExec.analyzeExec = &mockOpr
	tmpCtx := context.Background()
	_, err := explainExec.generateExplainInfo(tmpCtx)
	require.EqualError(t, err, "next error, close error")
	require.True(t, mockOpr.closed)

	// mockErrorOperator panic
	explainExec = &ExplainExec{
		baseExecutor: baseExec,
		explain:      nil,
	}
	mockOpr = mockErrorOperator{baseExec, true, false}
	explainExec.analyzeExec = &mockOpr
	defer func() {
		panicErr := recover()
		require.NotNil(t, panicErr)
		require.True(t, mockOpr.closed)
	}()
	_, _ = explainExec.generateExplainInfo(tmpCtx)
	require.FailNow(t, "generateExplainInfo should panic")
}
<|endoftext|>


<commit_before>// Copyright (C) 2016 Kohei YOSHIDA. All rights reserved.
//
// This program is free software; you can redistribute it and/or
// modify it under the terms of The BSD 3-Clause License
// that can be found in the LICENSE file.

package uritemplate

import (
	"unicode"
	"unicode/utf8"
)

// threadList implements https://research.swtch.com/sparse.
type threadList struct {
	dense  []thread
	sparse []uint32
}

type thread struct {
	op *progOp
	pc uint32
}

type machine struct {
	prog *prog

	list1   threadList
	list2   threadList
	matched bool

	input string
}

func (m *machine) at(pos int) (rune, int) {
	if pos < len(m.input) {
		r := m.input[pos]
		if r < utf8.RuneSelf {
			return rune(r), 1
		}
		return utf8.DecodeRuneInString(m.input[pos:])
	}
	return -1, 0
}

func (m *machine) clear(list *threadList) {
	list.dense = list.dense[:0]
}

func (m *machine) add(list *threadList, pc uint32) {
	if i := list.sparse[pc]; i < uint32(len(list.dense)) && list.dense[i].pc == pc {
		return
	}

	n := len(list.dense)
	list.dense = list.dense[:n+1]
	list.sparse[pc] = uint32(n)

	t := &list.dense[n]
	t.pc = pc
	t.op = &m.prog.op[pc]
}

func (m *machine) step(clist *threadList, nlist *threadList, r rune, pos int) {
	for i := 0; i < len(clist.dense); i++ {
		t := clist.dense[i]
		op := t.op

		switch op.code {
		default:
			panic("unhandled opcode")
		case opRune:
			if op.r == r {
				m.add(nlist, t.pc+1)
			}
		case opRuneClass:
			ret := false
			if !ret && op.rc&runeClassU == runeClassU {
				ret = ret || unicode.Is(rangeUnreserved, r)
			}
			if !ret && op.rc&runeClassR == runeClassR {
				ret = ret || unicode.Is(rangeReserved, r)
			}
			if !ret && op.rc&runeClassPctE == runeClassPctE {
				ret = ret || unicode.Is(unicode.ASCII_Hex_Digit, r)
			}
			if ret {
				m.add(nlist, t.pc+1)
			}
		case opLineBegin:
			if pos == 0 {
				m.add(clist, t.pc+1)
			}
		case opLineEnd:
			if r == -1 {
				m.add(clist, t.pc+1)
			}
		case opCapStart:
			m.add(clist, t.pc+1)
		case opCapEnd:
			m.add(clist, t.pc+1)
		case opSplit:
			m.add(clist, t.pc+1)
			m.add(clist, op.i)
		case opJmp, opJmpIfNotDefined, opJmpIfNotEmpty, opJmpIfNotFirst:
			m.add(clist, t.pc+1)
			m.add(clist, op.i)
		case opEnd:
			m.matched = true
		}
	}
	m.clear(clist)
}

func (m *machine) match() bool {
	pos := 0
	clist, nlist := &m.list1, &m.list2
	for {
		if len(clist.dense) == 0 {
			if m.matched {
				break
			}
		}
		if !m.matched {
			m.add(clist, 0)
		}

		r, width := m.at(pos)
		m.step(clist, nlist, r, pos)
		if width < 1 {
			break
		}
		pos += width

		clist, nlist = nlist, clist
	}
	return m.matched
}

type Matcher struct {
	prog prog
}

func CompileMatcher(tmpl *Template) (*Matcher, error) {
	c := compiler{}
	c.init()

	c.compile(tmpl)

	m := Matcher{
		prog: *c.prog,
	}
	return &m, nil
}

func (match *Matcher) Match(expansion string, vals map[string]Value) bool {
	n := len(match.prog.op)
	m := machine{
		prog: &match.prog,
		list1: threadList{
			dense:  make([]thread, 0, n),
			sparse: make([]uint32, n),
		},
		list2: threadList{
			dense:  make([]thread, 0, n),
			sparse: make([]uint32, n),
		},
		input: expansion,
	}
	return m.match()
}
<commit_msg>follow operation priorities<commit_after>// Copyright (C) 2016 Kohei YOSHIDA. All rights reserved.
//
// This program is free software; you can redistribute it and/or
// modify it under the terms of The BSD 3-Clause License
// that can be found in the LICENSE file.

package uritemplate

import (
	"unicode"
	"unicode/utf8"
)

// threadList implements https://research.swtch.com/sparse.
type threadList struct {
	dense  []threadEntry
	sparse []uint32
}

type threadEntry struct {
	pc uint32
	t  *thread
}

type thread struct {
	op  *progOp
	cap map[string][]int
}

type machine struct {
	prog *prog

	list1   threadList
	list2   threadList
	matched bool

	input string
}

func (m *machine) at(pos int) (rune, int, bool) {
	if l := len(m.input); pos < l {
		c := m.input[pos]
		if c < utf8.RuneSelf {
			return rune(c), 1, pos+1 < l
		}
		r, size := utf8.DecodeRuneInString(m.input[pos:])
		return r, size, pos+size < l
	}
	return -1, 0, false
}

func (m *machine) add(list *threadList, pc uint32, pos int, next bool) {
	if i := list.sparse[pc]; i < uint32(len(list.dense)) && list.dense[i].pc == pc {
		return
	}

	n := len(list.dense)
	list.dense = list.dense[:n+1]
	list.sparse[pc] = uint32(n)

	e := &list.dense[n]
	e.pc = pc
	e.t = nil

	op := &m.prog.op[pc]
	switch op.code {
	default:
		panic("unhandled opcode")
	case opRune, opRuneClass, opEnd:
		e.t = &thread{
			op: &m.prog.op[pc],
		}
	case opLineBegin:
		if pos == 0 {
			m.add(list, pc+1, pos, next)
		}
	case opLineEnd:
		if !next {
			m.add(list, pc+1, pos, next)
		}
	case opCapStart, opCapEnd:
		m.add(list, pc+1, pos, next)
	case opSplit:
		m.add(list, pc+1, pos, next)
		m.add(list, op.i, pos, next)
	case opJmp:
		m.add(list, op.i, pos, next)
	case opJmpIfNotDefined:
		m.add(list, pc+1, pos, next)
		m.add(list, op.i, pos, next)
	case opJmpIfNotFirst:
		m.add(list, pc+1, pos, next)
		m.add(list, op.i, pos, next)
	case opJmpIfNotEmpty:
		m.add(list, op.i, pos, next)
		m.add(list, pc+1, pos, next)
	}
}

func (m *machine) step(clist *threadList, nlist *threadList, r rune, pos int, nextPos int, next bool) {
	for i := 0; i < len(clist.dense); i++ {
		e := clist.dense[i]
		if e.t == nil {
			continue
		}

		t := e.t
		op := t.op
		switch op.code {
		default:
			panic("unhandled opcode")
		case opRune:
			if op.r == r {
				m.add(nlist, e.pc+1, nextPos, next)
			}
		case opRuneClass:
			ret := false
			if !ret && op.rc&runeClassU == runeClassU {
				ret = ret || unicode.Is(rangeUnreserved, r)
			}
			if !ret && op.rc&runeClassR == runeClassR {
				ret = ret || unicode.Is(rangeReserved, r)
			}
			if !ret && op.rc&runeClassPctE == runeClassPctE {
				ret = ret || unicode.Is(unicode.ASCII_Hex_Digit, r)
			}
			if ret {
				m.add(nlist, e.pc+1, nextPos, next)
			}
		case opEnd:
			m.matched = true
			clist.dense = clist.dense[:0]
		}
	}
	clist.dense = clist.dense[:0]
}

func (m *machine) match() bool {
	pos := 0
	clist, nlist := &m.list1, &m.list2
	for {
		if len(clist.dense) == 0 && m.matched {
			break
		}
		r, width, next := m.at(pos)
		if !m.matched {
			m.add(clist, 0, pos, next)
		}
		m.step(clist, nlist, r, pos, pos+width, next)

		if width < 1 {
			break
		}
		pos += width

		clist, nlist = nlist, clist
	}
	return m.matched
}

type Matcher struct {
	prog prog
}

func CompileMatcher(tmpl *Template) (*Matcher, error) {
	c := compiler{}
	c.init()
	c.compile(tmpl)

	m := Matcher{
		prog: *c.prog,
	}
	return &m, nil
}

func (match *Matcher) Match(expansion string, vals map[string]Value) bool {
	n := len(match.prog.op)
	m := machine{
		prog: &match.prog,
		list1: threadList{
			dense:  make([]threadEntry, 0, n),
			sparse: make([]uint32, n),
		},
		list2: threadList{
			dense:  make([]threadEntry, 0, n),
			sparse: make([]uint32, n),
		},
		input: expansion,
	}
	return m.match()
}
<|endoftext|>


<commit_before>/*
Copyright 2016 The Kubernetes 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 gcp

import (
	"bytes"
	"encoding/json"
	"fmt"
	"net/http"
	"os/exec"
	"strings"
	"sync"
	"time"

	"github.com/golang/glog"
	"golang.org/x/net/context"
	"golang.org/x/oauth2"
	"golang.org/x/oauth2/google"
	"k8s.io/apimachinery/pkg/util/yaml"
	restclient "k8s.io/client-go/rest"
	"k8s.io/client-go/util/jsonpath"
)

func init() {
	if err := restclient.RegisterAuthProviderPlugin("gcp", newGCPAuthProvider); err != nil {
		glog.Fatalf("Failed to register gcp auth plugin: %v", err)
	}
}

// Stubbable for testing
var execCommand = exec.Command

// gcpAuthProvider is an auth provider plugin that uses GCP credentials to provide
// tokens for kubectl to authenticate itself to the apiserver. A sample json config
// is provided below with all recognized options described.
//
// {
//   'auth-provider': {
//     # Required
//     "name": "gcp",
//
//     'config': {
//       # Caching options
//
//       # Raw string data representing cached access token.
//       "access-token": "ya29.CjWdA4GiBPTt",
//       # RFC3339Nano expiration timestamp for cached access token.
//       "expiry": "2016-10-31 22:31:9.123",
//
//       # Command execution options
//       # These options direct the plugin to execute a specified command and parse
//       # token and expiry time from the output of the command.
//
//       # Command to execute for access token. Command output will be parsed as JSON.
//       # If "cmd-args" is not present, this value will be split on whitespace, with
//       # the first element interpreted as the command, remaining elements as args.
//       "cmd-path": "/usr/bin/gcloud",
//
//       # Arguments to pass to command to execute for access token.
//       "cmd-args": "config config-helper --output=json"
//
//       # JSONPath to the string field that represents the access token in
//       # command output. If omitted, defaults to "{.access_token}".
//       "token-key": "{.credential.access_token}",
//
//       # JSONPath to the string field that represents expiration timestamp
//       # of the access token in the command output. If omitted, defaults to
//       # "{.token_expiry}"
//       "expiry-key": ""{.credential.token_expiry}",
//
//       # golang reference time in the format that the expiration timestamp uses.
//       # If omitted, defaults to time.RFC3339Nano
//       "time-fmt": "2006-01-02 15:04:05.999999999"
//     }
//   }
// }
//
type gcpAuthProvider struct {
	tokenSource oauth2.TokenSource
	persister   restclient.AuthProviderConfigPersister
}

func newGCPAuthProvider(_ string, gcpConfig map[string]string, persister restclient.AuthProviderConfigPersister) (restclient.AuthProvider, error) {
	var ts oauth2.TokenSource
	var err error
	if cmd, useCmd := gcpConfig["cmd-path"]; useCmd {
		if len(cmd) == 0 {
			return nil, fmt.Errorf("missing access token cmd")
		}
		var args []string
		if cmdArgs, ok := gcpConfig["cmd-args"]; ok {
			args = strings.Fields(cmdArgs)
		} else {
			fields := strings.Fields(cmd)
			cmd = fields[0]
			args = fields[1:]
		}
		ts = newCmdTokenSource(cmd, args, gcpConfig["token-key"], gcpConfig["expiry-key"], gcpConfig["time-fmt"])
	} else {
		ts, err = google.DefaultTokenSource(context.Background(), "https://www.googleapis.com/auth/cloud-platform")
	}
	if err != nil {
		return nil, err
	}
	cts, err := newCachedTokenSource(gcpConfig["access-token"], gcpConfig["expiry"], persister, ts, gcpConfig)
	if err != nil {
		return nil, err
	}
	return &gcpAuthProvider{cts, persister}, nil
}

func (g *gcpAuthProvider) WrapTransport(rt http.RoundTripper) http.RoundTripper {
	return &oauth2.Transport{
		Source: g.tokenSource,
		Base:   rt,
	}
}

func (g *gcpAuthProvider) Login() error { return nil }

type cachedTokenSource struct {
	lk          sync.Mutex
	source      oauth2.TokenSource
	accessToken string
	expiry      time.Time
	persister   restclient.AuthProviderConfigPersister
	cache       map[string]string
}

func newCachedTokenSource(accessToken, expiry string, persister restclient.AuthProviderConfigPersister, ts oauth2.TokenSource, cache map[string]string) (*cachedTokenSource, error) {
	var expiryTime time.Time
	if parsedTime, err := time.Parse(time.RFC3339Nano, expiry); err == nil {
		expiryTime = parsedTime
	}
	if cache == nil {
		cache = make(map[string]string)
	}
	return &cachedTokenSource{
		source:      ts,
		accessToken: accessToken,
		expiry:      expiryTime,
		persister:   persister,
		cache:       cache,
	}, nil
}

func (t *cachedTokenSource) Token() (*oauth2.Token, error) {
	tok := t.cachedToken()
	if tok.Valid() && !tok.Expiry.IsZero() {
		return tok, nil
	}
	tok, err := t.source.Token()
	if err != nil {
		return nil, err
	}
	cache := t.update(tok)
	if t.persister != nil {
		if err := t.persister.Persist(cache); err != nil {
			glog.V(4).Infof("Failed to persist token: %v", err)
		}
	}
	return tok, nil
}

func (t *cachedTokenSource) cachedToken() *oauth2.Token {
	t.lk.Lock()
	defer t.lk.Unlock()
	return &oauth2.Token{
		AccessToken: t.accessToken,
		TokenType:   "Bearer",
		Expiry:      t.expiry,
	}
}

func (t *cachedTokenSource) update(tok *oauth2.Token) map[string]string {
	t.lk.Lock()
	defer t.lk.Unlock()
	t.accessToken = tok.AccessToken
	t.expiry = tok.Expiry
	ret := map[string]string{}
	for k, v := range t.cache {
		ret[k] = v
	}
	ret["access-token"] = t.accessToken
	ret["expiry"] = t.expiry.Format(time.RFC3339Nano)
	return ret
}

type commandTokenSource struct {
	cmd       string
	args      []string
	tokenKey  string
	expiryKey string
	timeFmt   string
}

func newCmdTokenSource(cmd string, args []string, tokenKey, expiryKey, timeFmt string) *commandTokenSource {
	if len(timeFmt) == 0 {
		timeFmt = time.RFC3339Nano
	}
	if len(tokenKey) == 0 {
		tokenKey = "{.access_token}"
	}
	if len(expiryKey) == 0 {
		expiryKey = "{.token_expiry}"
	}
	return &commandTokenSource{
		cmd:       cmd,
		args:      args,
		tokenKey:  tokenKey,
		expiryKey: expiryKey,
		timeFmt:   timeFmt,
	}
}

func (c *commandTokenSource) Token() (*oauth2.Token, error) {
	fullCmd := strings.Join(append([]string{c.cmd}, c.args...), " ")
	cmd := execCommand(c.cmd, c.args...)
	output, err := cmd.Output()
	if err != nil {
		return nil, fmt.Errorf("error executing access token command %q: err=%v output=%s", fullCmd, err, output)
	}
	token, err := c.parseTokenCmdOutput(output)
	if err != nil {
		return nil, fmt.Errorf("error parsing output for access token command %q: %v", fullCmd, err)
	}
	return token, nil
}

func (c *commandTokenSource) parseTokenCmdOutput(output []byte) (*oauth2.Token, error) {
	output, err := yaml.ToJSON(output)
	if err != nil {
		return nil, err
	}
	var data interface{}
	if err := json.Unmarshal(output, &data); err != nil {
		return nil, err
	}

	accessToken, err := parseJSONPath(data, "token-key", c.tokenKey)
	if err != nil {
		return nil, fmt.Errorf("error parsing token-key %q from %q: %v", c.tokenKey, string(output), err)
	}
	expiryStr, err := parseJSONPath(data, "expiry-key", c.expiryKey)
	if err != nil {
		return nil, fmt.Errorf("error parsing expiry-key %q from %q: %v", c.expiryKey, string(output), err)
	}
	var expiry time.Time
	if t, err := time.Parse(c.timeFmt, expiryStr); err != nil {
		glog.V(4).Infof("Failed to parse token expiry from %s (fmt=%s): %v", expiryStr, c.timeFmt, err)
	} else {
		expiry = t
	}

	return &oauth2.Token{
		AccessToken: accessToken,
		TokenType:   "Bearer",
		Expiry:      expiry,
	}, nil
}

func parseJSONPath(input interface{}, name, template string) (string, error) {
	j := jsonpath.New(name)
	buf := new(bytes.Buffer)
	if err := j.Parse(template); err != nil {
		return "", err
	}
	if err := j.Execute(buf, input); err != nil {
		return "", err
	}
	return buf.String(), nil
}
<commit_msg>gcp auth provider not to override the Auth header if it's already exits<commit_after>/*
Copyright 2016 The Kubernetes 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 gcp

import (
	"bytes"
	"encoding/json"
	"fmt"
	"net/http"
	"os/exec"
	"strings"
	"sync"
	"time"

	"github.com/golang/glog"
	"golang.org/x/net/context"
	"golang.org/x/oauth2"
	"golang.org/x/oauth2/google"
	"k8s.io/apimachinery/pkg/util/yaml"
	restclient "k8s.io/client-go/rest"
	"k8s.io/client-go/util/jsonpath"
)

func init() {
	if err := restclient.RegisterAuthProviderPlugin("gcp", newGCPAuthProvider); err != nil {
		glog.Fatalf("Failed to register gcp auth plugin: %v", err)
	}
}

// Stubbable for testing
var execCommand = exec.Command

// gcpAuthProvider is an auth provider plugin that uses GCP credentials to provide
// tokens for kubectl to authenticate itself to the apiserver. A sample json config
// is provided below with all recognized options described.
//
// {
//   'auth-provider': {
//     # Required
//     "name": "gcp",
//
//     'config': {
//       # Caching options
//
//       # Raw string data representing cached access token.
//       "access-token": "ya29.CjWdA4GiBPTt",
//       # RFC3339Nano expiration timestamp for cached access token.
//       "expiry": "2016-10-31 22:31:9.123",
//
//       # Command execution options
//       # These options direct the plugin to execute a specified command and parse
//       # token and expiry time from the output of the command.
//
//       # Command to execute for access token. Command output will be parsed as JSON.
//       # If "cmd-args" is not present, this value will be split on whitespace, with
//       # the first element interpreted as the command, remaining elements as args.
//       "cmd-path": "/usr/bin/gcloud",
//
//       # Arguments to pass to command to execute for access token.
//       "cmd-args": "config config-helper --output=json"
//
//       # JSONPath to the string field that represents the access token in
//       # command output. If omitted, defaults to "{.access_token}".
//       "token-key": "{.credential.access_token}",
//
//       # JSONPath to the string field that represents expiration timestamp
//       # of the access token in the command output. If omitted, defaults to
//       # "{.token_expiry}"
//       "expiry-key": ""{.credential.token_expiry}",
//
//       # golang reference time in the format that the expiration timestamp uses.
//       # If omitted, defaults to time.RFC3339Nano
//       "time-fmt": "2006-01-02 15:04:05.999999999"
//     }
//   }
// }
//
type gcpAuthProvider struct {
	tokenSource oauth2.TokenSource
	persister   restclient.AuthProviderConfigPersister
}

func newGCPAuthProvider(_ string, gcpConfig map[string]string, persister restclient.AuthProviderConfigPersister) (restclient.AuthProvider, error) {
	var ts oauth2.TokenSource
	var err error
	if cmd, useCmd := gcpConfig["cmd-path"]; useCmd {
		if len(cmd) == 0 {
			return nil, fmt.Errorf("missing access token cmd")
		}
		var args []string
		if cmdArgs, ok := gcpConfig["cmd-args"]; ok {
			args = strings.Fields(cmdArgs)
		} else {
			fields := strings.Fields(cmd)
			cmd = fields[0]
			args = fields[1:]
		}
		ts = newCmdTokenSource(cmd, args, gcpConfig["token-key"], gcpConfig["expiry-key"], gcpConfig["time-fmt"])
	} else {
		ts, err = google.DefaultTokenSource(context.Background(), "https://www.googleapis.com/auth/cloud-platform")
	}
	if err != nil {
		return nil, err
	}
	cts, err := newCachedTokenSource(gcpConfig["access-token"], gcpConfig["expiry"], persister, ts, gcpConfig)
	if err != nil {
		return nil, err
	}
	return &gcpAuthProvider{cts, persister}, nil
}

func (g *gcpAuthProvider) WrapTransport(rt http.RoundTripper) http.RoundTripper {
	return &conditionalTransport{&oauth2.Transport{Source: g.tokenSource, Base: rt}}
}

func (g *gcpAuthProvider) Login() error { return nil }

type cachedTokenSource struct {
	lk          sync.Mutex
	source      oauth2.TokenSource
	accessToken string
	expiry      time.Time
	persister   restclient.AuthProviderConfigPersister
	cache       map[string]string
}

func newCachedTokenSource(accessToken, expiry string, persister restclient.AuthProviderConfigPersister, ts oauth2.TokenSource, cache map[string]string) (*cachedTokenSource, error) {
	var expiryTime time.Time
	if parsedTime, err := time.Parse(time.RFC3339Nano, expiry); err == nil {
		expiryTime = parsedTime
	}
	if cache == nil {
		cache = make(map[string]string)
	}
	return &cachedTokenSource{
		source:      ts,
		accessToken: accessToken,
		expiry:      expiryTime,
		persister:   persister,
		cache:       cache,
	}, nil
}

func (t *cachedTokenSource) Token() (*oauth2.Token, error) {
	tok := t.cachedToken()
	if tok.Valid() && !tok.Expiry.IsZero() {
		return tok, nil
	}
	tok, err := t.source.Token()
	if err != nil {
		return nil, err
	}
	cache := t.update(tok)
	if t.persister != nil {
		if err := t.persister.Persist(cache); err != nil {
			glog.V(4).Infof("Failed to persist token: %v", err)
		}
	}
	return tok, nil
}

func (t *cachedTokenSource) cachedToken() *oauth2.Token {
	t.lk.Lock()
	defer t.lk.Unlock()
	return &oauth2.Token{
		AccessToken: t.accessToken,
		TokenType:   "Bearer",
		Expiry:      t.expiry,
	}
}

func (t *cachedTokenSource) update(tok *oauth2.Token) map[string]string {
	t.lk.Lock()
	defer t.lk.Unlock()
	t.accessToken = tok.AccessToken
	t.expiry = tok.Expiry
	ret := map[string]string{}
	for k, v := range t.cache {
		ret[k] = v
	}
	ret["access-token"] = t.accessToken
	ret["expiry"] = t.expiry.Format(time.RFC3339Nano)
	return ret
}

type commandTokenSource struct {
	cmd       string
	args      []string
	tokenKey  string
	expiryKey string
	timeFmt   string
}

func newCmdTokenSource(cmd string, args []string, tokenKey, expiryKey, timeFmt string) *commandTokenSource {
	if len(timeFmt) == 0 {
		timeFmt = time.RFC3339Nano
	}
	if len(tokenKey) == 0 {
		tokenKey = "{.access_token}"
	}
	if len(expiryKey) == 0 {
		expiryKey = "{.token_expiry}"
	}
	return &commandTokenSource{
		cmd:       cmd,
		args:      args,
		tokenKey:  tokenKey,
		expiryKey: expiryKey,
		timeFmt:   timeFmt,
	}
}

func (c *commandTokenSource) Token() (*oauth2.Token, error) {
	fullCmd := strings.Join(append([]string{c.cmd}, c.args...), " ")
	cmd := execCommand(c.cmd, c.args...)
	output, err := cmd.Output()
	if err != nil {
		return nil, fmt.Errorf("error executing access token command %q: err=%v output=%s", fullCmd, err, output)
	}
	token, err := c.parseTokenCmdOutput(output)
	if err != nil {
		return nil, fmt.Errorf("error parsing output for access token command %q: %v", fullCmd, err)
	}
	return token, nil
}

func (c *commandTokenSource) parseTokenCmdOutput(output []byte) (*oauth2.Token, error) {
	output, err := yaml.ToJSON(output)
	if err != nil {
		return nil, err
	}
	var data interface{}
	if err := json.Unmarshal(output, &data); err != nil {
		return nil, err
	}

	accessToken, err := parseJSONPath(data, "token-key", c.tokenKey)
	if err != nil {
		return nil, fmt.Errorf("error parsing token-key %q from %q: %v", c.tokenKey, string(output), err)
	}
	expiryStr, err := parseJSONPath(data, "expiry-key", c.expiryKey)
	if err != nil {
		return nil, fmt.Errorf("error parsing expiry-key %q from %q: %v", c.expiryKey, string(output), err)
	}
	var expiry time.Time
	if t, err := time.Parse(c.timeFmt, expiryStr); err != nil {
		glog.V(4).Infof("Failed to parse token expiry from %s (fmt=%s): %v", expiryStr, c.timeFmt, err)
	} else {
		expiry = t
	}

	return &oauth2.Token{
		AccessToken: accessToken,
		TokenType:   "Bearer",
		Expiry:      expiry,
	}, nil
}

func parseJSONPath(input interface{}, name, template string) (string, error) {
	j := jsonpath.New(name)
	buf := new(bytes.Buffer)
	if err := j.Parse(template); err != nil {
		return "", err
	}
	if err := j.Execute(buf, input); err != nil {
		return "", err
	}
	return buf.String(), nil
}

type conditionalTransport struct {
	oauthTransport *oauth2.Transport
}

func (t *conditionalTransport) RoundTrip(req *http.Request) (*http.Response, error) {
	if len(req.Header.Get("Authorization")) != 0 {
		return t.oauthTransport.Base.RoundTrip(req)
	}
	return t.oauthTransport.RoundTrip(req)
}
<|endoftext|>


<commit_before>// Copyright 2017 Tamás Gulácsi
//
//
//    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 grpcer

import (
	"io"
	"io/ioutil"
	"os"
	"reflect"
	"strings"

	jsoniter "github.com/json-iterator/go"
	"github.com/pkg/errors"
)

var errNewField = errors.New("new field")

type streamEncoder interface {
	WriteField(w io.Writer, name string) error
}

func mergeStreams(w io.Writer, first interface{}, recv interface {
	Recv() (interface{}, error)
},
	Log func(...interface{}) error,
) {
	if Log == nil {
		Log = func(...interface{}) error { return nil }
	}

	slice, notSlice := sliceFields(first)
	if len(slice) == 0 {
		var err error
		part := first
		enc := jsoniter.NewEncoder(w)
		for {
			if err := enc.Encode(part); err != nil {
				Log("encode", part, "error", err)
				return
			}

			part, err = recv.Recv()
			if err != nil {
				if err != io.EOF {
					Log("msg", "recv", "error", err)
				}
				break
			}
		}
		return
	}

	names := make(map[string]bool, len(slice)+len(notSlice))

	Log("slices", slice)
	w.Write([]byte("{"))
	for _, f := range notSlice {
		tw := newTrimWriter(w, "", "\n")
		jsoniter.NewEncoder(tw).Encode(f.JSONName)
		tw.Close()
		w.Write([]byte{':'})
		tw = newTrimWriter(w, "", "\n")
		jsoniter.NewEncoder(tw).Encode(f.Value)
		tw.Close()
		w.Write([]byte{','})

		names[f.Name] = false
	}
	tw := newTrimWriter(w, "", "\n")
	jsoniter.NewEncoder(tw).Encode(slice[0].JSONName)
	tw.Close()
	w.Write([]byte(":"))
	tw = newTrimWriter(w, "", "]")
	jsoniter.NewEncoder(tw).Encode(slice[0].Value)
	tw.Close()

	names[slice[0].Name] = true

	files := make(map[string]*os.File, len(slice)-1)
	for _, f := range slice[1:] {
		fh, err := ioutil.TempFile("", "merge-"+f.Name+"-")
		if err != nil {
			Log("tempFile", f.Name, "error", err)
			return
		}
		os.Remove(fh.Name())
		defer fh.Close()
		files[f.Name] = fh
		tw := newTrimWriter(fh, "", "\n")
		jsoniter.NewEncoder(tw).Encode(f.JSONName)
		tw.Close()
		io.WriteString(fh, ":[")
		tw = newTrimWriter(fh, "[", "]")
		jsoniter.NewEncoder(tw).Encode(f.Value)
		tw.Close()

		names[f.Name] = true
	}

	var part interface{}
	var err error
	for {
		part, err = recv.Recv()
		if err != nil {
			if err != io.EOF {
				Log("msg", "recv", "error", err)
			}
			break
		}

		S, nS := sliceFields(part)
		for _, f := range S {
			if isSlice, ok := names[f.Name]; !(ok && isSlice) {
				err = errors.Wrap(errNewField, f.Name)
			}
		}
		for _, f := range nS {
			if isSlice, ok := names[f.Name]; !(ok && !isSlice) {
				err = errors.Wrap(errNewField, f.Name)
			}
		}
		if err != nil {
			Log("error", err)
			//TODO(tgulacsi): close the merge and send as is
		}

		if S[0].Name == slice[0].Name {
			w.Write([]byte{','})
			tw := newTrimWriter(w, "[", "]")
			jsoniter.NewEncoder(tw).Encode(S[0].Value)
			tw.Close()
			S = S[1:]
		}
		for _, f := range S {
			fh := files[f.Name]
			if _, err := fh.Write([]byte{','}); err != nil {
				Log("write", fh.Name(), "error", err)
			}
			tw := newTrimWriter(fh, "[", "]")
			jsoniter.NewEncoder(tw).Encode(f.Value)
			tw.Close()
		}
	}
	w.Write([]byte("]"))

	for _, fh := range files {
		if _, err := fh.Seek(0, 0); err != nil {
			Log("Seek", fh.Name(), "error", err)
			continue
		}
		w.Write([]byte{','})
		io.Copy(w, fh)
		w.Write([]byte{']'})
	}
	w.Write([]byte{'}', '\n'})
}

type field struct {
	Name     string
	JSONName string
	Value    interface{}
}

func sliceFields(part interface{}) (slice, notSlice []field) {
	rv := reflect.ValueOf(part)
	t := rv.Type()
	if t.Kind() == reflect.Ptr {
		rv = rv.Elem()
		t = rv.Type()
	}
	n := t.NumField()
	for i := 0; i < n; i++ {
		f := rv.Field(i)
		tf := t.Field(i)
		fld := field{Name: tf.Name, Value: f.Interface()}
		fld.JSONName = tf.Tag.Get("json")
		if i := strings.IndexByte(fld.JSONName, ','); i >= 0 {
			fld.JSONName = fld.JSONName[:i]
		}
		if fld.JSONName == "" {
			fld.JSONName = fld.Name
		}

		if f.Type().Kind() != reflect.Slice {
			notSlice = append(notSlice, fld)
			continue
		}
		if f.IsNil() {
			continue
		}
		slice = append(slice, fld)
	}
	return slice, notSlice
}

type trimWriter struct {
	w              io.Writer
	prefix, suffix string
	buf            []byte
}

func newTrimWriter(w io.Writer, prefix, suffix string) *trimWriter {
	return &trimWriter{w: w, prefix: prefix, suffix: suffix}
}
func (tw *trimWriter) Write(p []byte) (int, error) {
	n := len(p)
	if tw.prefix != "" {
		if len(tw.prefix) >= len(p) {
			tw.prefix = tw.prefix[len(p):]
			return n, nil
		}
		p = p[len(tw.prefix):]
		tw.prefix = ""
	}

	if len(tw.buf) > 0 && len(tw.buf) >= len(tw.suffix) {
		if _, err := tw.w.Write(tw.buf); err != nil {
			return 0, err
		}
		tw.buf = tw.buf[:0]
	}
	if len(p) <= len(tw.suffix) {
		tw.buf = append(tw.buf, p...)
		return n, nil
	}
	i := len(p) - len(tw.suffix) + len(tw.buf)
	tw.buf = append(tw.buf, p[i:]...)
	_, err := tw.w.Write(p[:i])
	return n, err
}
func (tw *trimWriter) Close() error {
	if tw.suffix == string(tw.buf) {
		return nil
	}
	_, err := tw.w.Write(tw.buf)
	return err
}
<commit_msg>trim: use transform<commit_after>// Copyright 2017 Tamás Gulácsi
//
//
//    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 grpcer

import (
	"io"
	"io/ioutil"
	"os"
	"reflect"
	"strings"

	jsoniter "github.com/json-iterator/go"
	"github.com/pkg/errors"
	"github.com/tgulacsi/go/stream"
)

var errNewField = errors.New("new field")

type streamEncoder interface {
	WriteField(w io.Writer, name string) error
}

func mergeStreams(w io.Writer, first interface{}, recv interface {
	Recv() (interface{}, error)
},
	Log func(...interface{}) error,
) {
	if Log == nil {
		Log = func(...interface{}) error { return nil }
	}

	slice, notSlice := sliceFields(first)
	if len(slice) == 0 {
		var err error
		part := first
		enc := jsoniter.NewEncoder(w)
		for {
			if err := enc.Encode(part); err != nil {
				Log("encode", part, "error", err)
				return
			}

			part, err = recv.Recv()
			if err != nil {
				if err != io.EOF {
					Log("msg", "recv", "error", err)
				}
				break
			}
		}
		return
	}

	names := make(map[string]bool, len(slice)+len(notSlice))

	Log("slices", slice)
	w.Write([]byte("{"))
	for _, f := range notSlice {
		tw := stream.NewTrimSpace(w)
		jsoniter.NewEncoder(tw).Encode(f.JSONName)
		w.Write([]byte{':'})
		tw.Close()
		tw = stream.NewTrimSpace(w)
		jsoniter.NewEncoder(tw).Encode(f.Value)
		w.Write([]byte{','})
		tw.Close()

		names[f.Name] = false
	}
	tw := stream.NewTrimSpace(w)
	jsoniter.NewEncoder(tw).Encode(slice[0].JSONName)
	w.Write([]byte(":"))
	tw.Close()
	tw = stream.NewTrimFix(w, "", "]")
	jsoniter.NewEncoder(tw).Encode(slice[0].Value)
	tw.Close()

	names[slice[0].Name] = true

	files := make(map[string]*os.File, len(slice)-1)
	for _, f := range slice[1:] {
		fh, err := ioutil.TempFile("", "merge-"+f.Name+"-")
		if err != nil {
			Log("tempFile", f.Name, "error", err)
			return
		}
		os.Remove(fh.Name())
		defer fh.Close()
		files[f.Name] = fh
		tw := stream.NewTrimSpace(fh)
		jsoniter.NewEncoder(tw).Encode(f.JSONName)
		io.WriteString(fh, ":[")
		tw.Close()
		tw = stream.NewTrimFix(fh, "[", "]")
		jsoniter.NewEncoder(tw).Encode(f.Value)
		tw.Close()

		names[f.Name] = true
	}

	var part interface{}
	var err error
	for {
		part, err = recv.Recv()
		if err != nil {
			if err != io.EOF {
				Log("msg", "recv", "error", err)
			}
			break
		}

		S, nS := sliceFields(part)
		for _, f := range S {
			if isSlice, ok := names[f.Name]; !(ok && isSlice) {
				err = errors.Wrap(errNewField, f.Name)
			}
		}
		for _, f := range nS {
			if isSlice, ok := names[f.Name]; !(ok && !isSlice) {
				err = errors.Wrap(errNewField, f.Name)
			}
		}
		if err != nil {
			Log("error", err)
			//TODO(tgulacsi): close the merge and send as is
		}

		if S[0].Name == slice[0].Name {
			w.Write([]byte{','})
			tw := stream.NewTrimFix(w, "[", "]")
			jsoniter.NewEncoder(tw).Encode(S[0].Value)
			tw.Close()
			S = S[1:]
		}
		for _, f := range S {
			fh := files[f.Name]
			if _, err := fh.Write([]byte{','}); err != nil {
				Log("write", fh.Name(), "error", err)
			}
			tw := stream.NewTrimFix(fh, "[", "]")
			jsoniter.NewEncoder(tw).Encode(f.Value)
			tw.Close()
		}
	}
	w.Write([]byte("]"))

	for _, fh := range files {
		if _, err := fh.Seek(0, 0); err != nil {
			Log("Seek", fh.Name(), "error", err)
			continue
		}
		w.Write([]byte{','})
		io.Copy(w, fh)
		w.Write([]byte{']'})
	}
	w.Write([]byte{'}', '\n'})
}

type field struct {
	Name     string
	JSONName string
	Value    interface{}
}

func sliceFields(part interface{}) (slice, notSlice []field) {
	rv := reflect.ValueOf(part)
	t := rv.Type()
	if t.Kind() == reflect.Ptr {
		rv = rv.Elem()
		t = rv.Type()
	}
	n := t.NumField()
	for i := 0; i < n; i++ {
		f := rv.Field(i)
		tf := t.Field(i)
		fld := field{Name: tf.Name, Value: f.Interface()}
		fld.JSONName = tf.Tag.Get("json")
		if i := strings.IndexByte(fld.JSONName, ','); i >= 0 {
			fld.JSONName = fld.JSONName[:i]
		}
		if fld.JSONName == "" {
			fld.JSONName = fld.Name
		}

		if f.Type().Kind() != reflect.Slice {
			notSlice = append(notSlice, fld)
			continue
		}
		if f.IsNil() {
			continue
		}
		slice = append(slice, fld)
	}
	return slice, notSlice
}
<|endoftext|>