<commit_before>//  Copyright (c) 2014 Marty Schoch
//  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 gouchstore

import (
	"os"
	"sort"
)

// Document represents a document stored in the database.
type Document struct {
	ID   string
	Body []byte
}

// DocumentInfo is document meta-data.
type DocumentInfo struct {
	ID           string `json:"id"`           // document identifier
	Seq          uint64 `json:"seq"`          // sequence number in database
	Rev          uint64 `json:"rev"`          // revision number of document
	RevMeta      []byte `json:"revMeta"`      // additional revision meta-data (uninterpreted by Gouchstore)
	ContentMeta  uint8  `json:"contentMeta"`  // content meta-data flags
	Deleted      bool   `json:"deleted"`      // is the revision deleted?
	Size         uint64 `json:"size"`         // size of document data in bytes
	bodyPosition uint64 `json:"bodyPosition"` // byte offset of document body in file
}

// Compressed returns whether or not this document has been compressed for storage.
func (di *DocumentInfo) compressed() bool {
	if di.ContentMeta&gs_DOC_IS_COMPRESSED != 0 {
		return true
	}
	return false
}

// DatabaseInfo describes the database as a whole.
type DatabaseInfo struct {
	FileName       string `json:"fileName"`      // filesystem path
	LastSeq        uint64 `json:"lastSeq"`       // last sequence number allocated
	DocumentCount  uint64 `json:"documentCount"` // total number of (non-deleted) documents
	DeletedCount   uint64 `json:"deletedCount"`  // total number of deleted documents
	SpaceUsed      uint64 `json:"spaceUsed"`     // disk space actively used by docs
	FileSize       uint64 `json:fileSize"`       // total disk space used by database
	HeaderPosition uint64 `json:headerPosition"` // file offset of current header
}

// DocumentInfoCallback is a function definition which is used for document iteration.
// For example, when iterating all documents with the AllDocuments() method,
// the provided callback will be invoked for each document in the database.
type DocumentInfoCallback func(gouchstore *Gouchstore, documentInfo *DocumentInfo, userContext interface{})

// Gouchstore gives access to a couchstore database file.
type Gouchstore struct {
	file   *os.File
	pos    int64
	header *header
}

// Open attemps to open an existing couchstore file.
// Since write operations are not yet supported, all files are opened read-only.
//
// All Gouchstore files successfully opened should be closed with the Close() method.
func Open(filename string) (*Gouchstore, error) {
	file, err := os.Open(filename)
	if err != nil {
		return nil, err
	}
	rv := Gouchstore{
		file: file,
	}
	err = rv.findLastHeader()
	if err != nil {
		return nil, err
	}
	return &rv, nil
}

func gouchstoreFetchCallback(g *Gouchstore, docInfo *DocumentInfo, userContext interface{}) {
	resultList := userContext.(*([]*DocumentInfo))
	*resultList = append(*resultList, docInfo)
}

// DocumentInfoById returns DocumentInfo for a single document with the specified ID.
func (g *Gouchstore) DocumentInfoById(id string) (*DocumentInfo, error) {
	current := 0
	resultList := make([]*DocumentInfo, 0)
	err := g.btree_lookup(
		g.header.byIdRoot,
		gs_INDEX_TYPE_BY_ID,
		&current,
		[][]byte{[]byte(id)},
		gouchstoreIdComparator,
		gouchstoreFetchCallback,
		&resultList)
	if err != nil {
		return nil, err
	}
	if len(resultList) > 0 {
		return resultList[0], nil
	}
	return nil, gs_ERROR_DOCUMENT_NOT_FOUND
}

// DocumentInfosByIds returns DocumentInfo objects for the specified document IDs.
// This will be more efficient than making consecutive calls to DocumentInfoById().
//
// NOTE: contents of the result slice will be in ascending ID order, not the order they
// appeared in the argument list.
func (g *Gouchstore) DocumentInfosByIds(identifiers []string) ([]*DocumentInfo, error) {
	ids := sort.StringSlice(identifiers)
	// we need the ids in sorted order
	sort.Sort(ids)

	// convert it to a slice of byte slices
	idsBytes := make([][]byte, len(ids))
	for i, id := range ids {
		idsBytes[i] = []byte(id)
	}

	current := 0
	resultList := make([]*DocumentInfo, 0)
	err := g.btree_lookup(
		g.header.byIdRoot,
		gs_INDEX_TYPE_BY_ID,
		&current,
		idsBytes,
		gouchstoreIdComparator,
		gouchstoreFetchCallback,
		&resultList)
	if err != nil {
		return nil, err
	}
	return resultList, nil
}

// DocumentInfoBySeq returns DocumentInfo for a single document with the specified sequence number.
func (g *Gouchstore) DocumentInfoBySeq(seq uint64) (*DocumentInfo, error) {
	current := 0
	resultList := make([]*DocumentInfo, 0)
	err := g.btree_lookup(
		g.header.bySeqRoot,
		gs_INDEX_TYPE_BY_SEQ,
		&current,
		[][]byte{encode_raw48(seq)},
		gouchstoreSeqComparator,
		gouchstoreFetchCallback,
		&resultList)
	if err != nil {
		return nil, err
	}
	if len(resultList) > 0 {
		return resultList[0], nil
	}
	return nil, gs_ERROR_DOCUMENT_NOT_FOUND
}

// DocumentInfosBySeqs returns DocumentInfo objects for the specified document sequence numbers.
// This will be more efficient than making consecutive calls to DocumentInfoBySeq().
//
// NOTE: contents of the result slice will be in ascending sequence order, not the order they
// appeared in the argument list.
func (g *Gouchstore) DocumentInfosBySeqs(sequences []uint64) ([]*DocumentInfo, error) {

	seqs := seqList(sequences)
	// we need the ids in sorted order
	sort.Sort(seqs)

	// convert it to a slice of byte slices
	idsBytes := make([][]byte, len(seqs))
	for i, seq := range seqs {
		idsBytes[i] = encode_raw48(seq)
	}

	current := 0
	resultList := make([]*DocumentInfo, 0)
	err := g.btree_lookup(
		g.header.bySeqRoot,
		gs_INDEX_TYPE_BY_SEQ,
		&current,
		idsBytes,
		gouchstoreSeqComparator,
		gouchstoreFetchCallback,
		&resultList)
	if err != nil {
		return nil, err
	}
	return resultList, nil
}

// AllDocuments will iterate through all documents in the database in ascending ID order,
// from startId (inclusive) through endId (inclusive).
// For each document, the provided DocumentInfoCallback will be invoked.  A user specified context can be included,
// and this will be passed to each invocation of the callback.
//
// If startId is the empty string, the iteration will start with the first document.
//
// If endId is the empty string, the iteration will continue to the last document.
func (g *Gouchstore) AllDocuments(startId, endId string, cb DocumentInfoCallback, userContext interface{}) error {
	active := 0
	err := g.btree_range(
		g.header.byIdRoot,
		gs_INDEX_TYPE_BY_ID,
		&active,
		[]byte(startId),
		[]byte(endId),
		gouchstoreIdComparator,
		cb,
		userContext)
	if err != nil {
		return err
	}
	return nil
}

// ChangesSince will iterate through all documents in the database in ascending sequence number order,
// from since (inclusive) through till (inclusive).
// For each document, the provided DocumentInfoCallback will be invoked.  A user specified context can be included,
// and this will be passed to each invocation of the callback.
//
// If since is 0, the iteration will start with the first document.
//
// If endId is 0, the iteration will continue to the last document.
func (g *Gouchstore) ChangesSince(since uint64, till uint64, cb DocumentInfoCallback, userContext interface{}) error {
	// treat till of 0 to mean no end boundary
	var endId []byte
	if till == 0 {
		endId = []byte{}
	} else {
		endId = encode_raw48(till)
	}
	active := 0
	err := g.btree_range(
		g.header.bySeqRoot,
		gs_INDEX_TYPE_BY_SEQ,
		&active,
		encode_raw48(since),
		endId,
		gouchstoreSeqComparator,
		cb,
		userContext)
	if err != nil {
		return err
	}
	return nil
}

// DocumentById returns the Document with the specified identifier.
func (g *Gouchstore) DocumentById(id string) (*Document, error) {
	docInfo, err := g.DocumentInfoById(id)
	if err != nil {
		return nil, err
	}
	return g.DocumentByDocumentInfo(docInfo)
}

// DocumentByDocumentInfo returns the Document using the provided DocumentInfo.
// The provided DocumentInfo should be valid, such as one received by one of the
// DocumentInfo*() methods, on the current couchstore file.
func (g *Gouchstore) DocumentByDocumentInfo(docInfo *DocumentInfo) (*Document, error) {
	var rv Document
	var err error
	if docInfo.compressed() {
		rv.Body, err = g.readCompressedDataChunkAt(int64(docInfo.bodyPosition))
		if err != nil {
			return nil, err
		}
	} else {
		rv.Body, err = g.readChunkAt(int64(docInfo.bodyPosition), false)
		if err != nil {
			return nil, err
		}
	}
	rv.ID = docInfo.ID
	return &rv, nil
}

// DatabaseInfo returns information describing the database itself.
func (g *Gouchstore) DatabaseInfo() (*DatabaseInfo, error) {
	rv := DatabaseInfo{
		FileName:       g.file.Name(),
		LastSeq:        g.header.updateSeq,
		FileSize:       uint64(g.pos),
		HeaderPosition: g.header.position,
	}
	if g.header.byIdRoot != nil {
		rv.DocumentCount = decode_raw40(g.header.byIdRoot.reducedValue[0:5])
		rv.DeletedCount = decode_raw40(g.header.byIdRoot.reducedValue[5:10])
		rv.SpaceUsed = decode_raw48(g.header.byIdRoot.reducedValue[10:16])
		rv.SpaceUsed += g.header.byIdRoot.subtreeSize
	}
	if g.header.bySeqRoot != nil {
		rv.SpaceUsed += g.header.bySeqRoot.subtreeSize
	}
	if g.header.localDocsRoot != nil {
		rv.SpaceUsed += g.header.localDocsRoot.subtreeSize
	}
	return &rv, nil
}

// Close will close the underlying file handle and release any resources associated with the Gouchstore object.
func (g *Gouchstore) Close() error {
	return g.file.Close()
}

const (
	gs_DOC_IS_JSON               = 0
	gs_DOC_INVALID_JSON          = 1
	gs_DOC_INVALID_JSON_KEY      = 2
	gs_DOC_NON_JSON_MODE         = 3
	gs_DOC_IS_COMPRESSED    byte = 128
)

type seqList []uint64

func (s seqList) Len() int           { return len(s) }
func (s seqList) Swap(i, j int)      { s[i], s[j] = s[j], s[i] }
func (s seqList) Less(i, j int) bool { return s[i] < s[j] }
<commit_msg>fix issues identified by go vet<commit_after>//  Copyright (c) 2014 Marty Schoch
//  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 gouchstore

import (
	"os"
	"sort"
)

// Document represents a document stored in the database.
type Document struct {
	ID   string
	Body []byte
}

// DocumentInfo is document meta-data.
type DocumentInfo struct {
	ID           string `json:"id"`           // document identifier
	Seq          uint64 `json:"seq"`          // sequence number in database
	Rev          uint64 `json:"rev"`          // revision number of document
	RevMeta      []byte `json:"revMeta"`      // additional revision meta-data (uninterpreted by Gouchstore)
	ContentMeta  uint8  `json:"contentMeta"`  // content meta-data flags
	Deleted      bool   `json:"deleted"`      // is the revision deleted?
	Size         uint64 `json:"size"`         // size of document data in bytes
	bodyPosition uint64 `json:"bodyPosition"` // byte offset of document body in file
}

// Compressed returns whether or not this document has been compressed for storage.
func (di *DocumentInfo) compressed() bool {
	if di.ContentMeta&gs_DOC_IS_COMPRESSED != 0 {
		return true
	}
	return false
}

// DatabaseInfo describes the database as a whole.
type DatabaseInfo struct {
	FileName       string `json:"fileName"`       // filesystem path
	LastSeq        uint64 `json:"lastSeq"`        // last sequence number allocated
	DocumentCount  uint64 `json:"documentCount"`  // total number of (non-deleted) documents
	DeletedCount   uint64 `json:"deletedCount"`   // total number of deleted documents
	SpaceUsed      uint64 `json:"spaceUsed"`      // disk space actively used by docs
	FileSize       uint64 `json:"fileSize"`       // total disk space used by database
	HeaderPosition uint64 `json:"headerPosition"` // file offset of current header
}

// DocumentInfoCallback is a function definition which is used for document iteration.
// For example, when iterating all documents with the AllDocuments() method,
// the provided callback will be invoked for each document in the database.
type DocumentInfoCallback func(gouchstore *Gouchstore, documentInfo *DocumentInfo, userContext interface{})

// Gouchstore gives access to a couchstore database file.
type Gouchstore struct {
	file   *os.File
	pos    int64
	header *header
}

// Open attemps to open an existing couchstore file.
// Since write operations are not yet supported, all files are opened read-only.
//
// All Gouchstore files successfully opened should be closed with the Close() method.
func Open(filename string) (*Gouchstore, error) {
	file, err := os.Open(filename)
	if err != nil {
		return nil, err
	}
	rv := Gouchstore{
		file: file,
	}
	err = rv.findLastHeader()
	if err != nil {
		return nil, err
	}
	return &rv, nil
}

func gouchstoreFetchCallback(g *Gouchstore, docInfo *DocumentInfo, userContext interface{}) {
	resultList := userContext.(*([]*DocumentInfo))
	*resultList = append(*resultList, docInfo)
}

// DocumentInfoById returns DocumentInfo for a single document with the specified ID.
func (g *Gouchstore) DocumentInfoById(id string) (*DocumentInfo, error) {
	current := 0
	resultList := make([]*DocumentInfo, 0)
	err := g.btree_lookup(
		g.header.byIdRoot,
		gs_INDEX_TYPE_BY_ID,
		&current,
		[][]byte{[]byte(id)},
		gouchstoreIdComparator,
		gouchstoreFetchCallback,
		&resultList)
	if err != nil {
		return nil, err
	}
	if len(resultList) > 0 {
		return resultList[0], nil
	}
	return nil, gs_ERROR_DOCUMENT_NOT_FOUND
}

// DocumentInfosByIds returns DocumentInfo objects for the specified document IDs.
// This will be more efficient than making consecutive calls to DocumentInfoById().
//
// NOTE: contents of the result slice will be in ascending ID order, not the order they
// appeared in the argument list.
func (g *Gouchstore) DocumentInfosByIds(identifiers []string) ([]*DocumentInfo, error) {
	ids := sort.StringSlice(identifiers)
	// we need the ids in sorted order
	sort.Sort(ids)

	// convert it to a slice of byte slices
	idsBytes := make([][]byte, len(ids))
	for i, id := range ids {
		idsBytes[i] = []byte(id)
	}

	current := 0
	resultList := make([]*DocumentInfo, 0)
	err := g.btree_lookup(
		g.header.byIdRoot,
		gs_INDEX_TYPE_BY_ID,
		&current,
		idsBytes,
		gouchstoreIdComparator,
		gouchstoreFetchCallback,
		&resultList)
	if err != nil {
		return nil, err
	}
	return resultList, nil
}

// DocumentInfoBySeq returns DocumentInfo for a single document with the specified sequence number.
func (g *Gouchstore) DocumentInfoBySeq(seq uint64) (*DocumentInfo, error) {
	current := 0
	resultList := make([]*DocumentInfo, 0)
	err := g.btree_lookup(
		g.header.bySeqRoot,
		gs_INDEX_TYPE_BY_SEQ,
		&current,
		[][]byte{encode_raw48(seq)},
		gouchstoreSeqComparator,
		gouchstoreFetchCallback,
		&resultList)
	if err != nil {
		return nil, err
	}
	if len(resultList) > 0 {
		return resultList[0], nil
	}
	return nil, gs_ERROR_DOCUMENT_NOT_FOUND
}

// DocumentInfosBySeqs returns DocumentInfo objects for the specified document sequence numbers.
// This will be more efficient than making consecutive calls to DocumentInfoBySeq().
//
// NOTE: contents of the result slice will be in ascending sequence order, not the order they
// appeared in the argument list.
func (g *Gouchstore) DocumentInfosBySeqs(sequences []uint64) ([]*DocumentInfo, error) {

	seqs := seqList(sequences)
	// we need the ids in sorted order
	sort.Sort(seqs)

	// convert it to a slice of byte slices
	idsBytes := make([][]byte, len(seqs))
	for i, seq := range seqs {
		idsBytes[i] = encode_raw48(seq)
	}

	current := 0
	resultList := make([]*DocumentInfo, 0)
	err := g.btree_lookup(
		g.header.bySeqRoot,
		gs_INDEX_TYPE_BY_SEQ,
		&current,
		idsBytes,
		gouchstoreSeqComparator,
		gouchstoreFetchCallback,
		&resultList)
	if err != nil {
		return nil, err
	}
	return resultList, nil
}

// AllDocuments will iterate through all documents in the database in ascending ID order,
// from startId (inclusive) through endId (inclusive).
// For each document, the provided DocumentInfoCallback will be invoked.  A user specified context can be included,
// and this will be passed to each invocation of the callback.
//
// If startId is the empty string, the iteration will start with the first document.
//
// If endId is the empty string, the iteration will continue to the last document.
func (g *Gouchstore) AllDocuments(startId, endId string, cb DocumentInfoCallback, userContext interface{}) error {
	active := 0
	err := g.btree_range(
		g.header.byIdRoot,
		gs_INDEX_TYPE_BY_ID,
		&active,
		[]byte(startId),
		[]byte(endId),
		gouchstoreIdComparator,
		cb,
		userContext)
	if err != nil {
		return err
	}
	return nil
}

// ChangesSince will iterate through all documents in the database in ascending sequence number order,
// from since (inclusive) through till (inclusive).
// For each document, the provided DocumentInfoCallback will be invoked.  A user specified context can be included,
// and this will be passed to each invocation of the callback.
//
// If since is 0, the iteration will start with the first document.
//
// If endId is 0, the iteration will continue to the last document.
func (g *Gouchstore) ChangesSince(since uint64, till uint64, cb DocumentInfoCallback, userContext interface{}) error {
	// treat till of 0 to mean no end boundary
	var endId []byte
	if till == 0 {
		endId = []byte{}
	} else {
		endId = encode_raw48(till)
	}
	active := 0
	err := g.btree_range(
		g.header.bySeqRoot,
		gs_INDEX_TYPE_BY_SEQ,
		&active,
		encode_raw48(since),
		endId,
		gouchstoreSeqComparator,
		cb,
		userContext)
	if err != nil {
		return err
	}
	return nil
}

// DocumentById returns the Document with the specified identifier.
func (g *Gouchstore) DocumentById(id string) (*Document, error) {
	docInfo, err := g.DocumentInfoById(id)
	if err != nil {
		return nil, err
	}
	return g.DocumentByDocumentInfo(docInfo)
}

// DocumentByDocumentInfo returns the Document using the provided DocumentInfo.
// The provided DocumentInfo should be valid, such as one received by one of the
// DocumentInfo*() methods, on the current couchstore file.
func (g *Gouchstore) DocumentByDocumentInfo(docInfo *DocumentInfo) (*Document, error) {
	var rv Document
	var err error
	if docInfo.compressed() {
		rv.Body, err = g.readCompressedDataChunkAt(int64(docInfo.bodyPosition))
		if err != nil {
			return nil, err
		}
	} else {
		rv.Body, err = g.readChunkAt(int64(docInfo.bodyPosition), false)
		if err != nil {
			return nil, err
		}
	}
	rv.ID = docInfo.ID
	return &rv, nil
}

// DatabaseInfo returns information describing the database itself.
func (g *Gouchstore) DatabaseInfo() (*DatabaseInfo, error) {
	rv := DatabaseInfo{
		FileName:       g.file.Name(),
		LastSeq:        g.header.updateSeq,
		FileSize:       uint64(g.pos),
		HeaderPosition: g.header.position,
	}
	if g.header.byIdRoot != nil {
		rv.DocumentCount = decode_raw40(g.header.byIdRoot.reducedValue[0:5])
		rv.DeletedCount = decode_raw40(g.header.byIdRoot.reducedValue[5:10])
		rv.SpaceUsed = decode_raw48(g.header.byIdRoot.reducedValue[10:16])
		rv.SpaceUsed += g.header.byIdRoot.subtreeSize
	}
	if g.header.bySeqRoot != nil {
		rv.SpaceUsed += g.header.bySeqRoot.subtreeSize
	}
	if g.header.localDocsRoot != nil {
		rv.SpaceUsed += g.header.localDocsRoot.subtreeSize
	}
	return &rv, nil
}

// Close will close the underlying file handle and release any resources associated with the Gouchstore object.
func (g *Gouchstore) Close() error {
	return g.file.Close()
}

const (
	gs_DOC_IS_JSON               = 0
	gs_DOC_INVALID_JSON          = 1
	gs_DOC_INVALID_JSON_KEY      = 2
	gs_DOC_NON_JSON_MODE         = 3
	gs_DOC_IS_COMPRESSED    byte = 128
)

type seqList []uint64

func (s seqList) Len() int           { return len(s) }
func (s seqList) Swap(i, j int)      { s[i], s[j] = s[j], s[i] }
func (s seqList) Less(i, j int) bool { return s[i] < s[j] }
<|endoftext|>


<commit_before>package db

import (
	"fmt"
	"github.com/lunny/xorm"
	r "github.com/robfig/revel"
)

var (
	Engine *xorm.Engine
)

func ensureOption(option string) string {
	value, found := r.Config.String(option)
	if !found {
		r.ERROR.Fatalf("Option %v not found", option)
	}
	return value
}

func calcSpec() string {
	host := ensureOption("db.host")
	name := ensureOption("db.name")
	username := ensureOption("db.username")
	password := r.Config.StringDefault("db.password", "")
	return fmt.Sprintf("%v:%v@tcp(%v)/%v", username, password, host, name)
}

func Init() {
	var found bool
	var driver string
	if driver, found = r.Config.String("db.driver"); !found {
		r.ERROR.Fatal("No db.driver found.")
	}
	var spec string
	if spec, found = r.Config.String("db.spec"); !found {
		spec = calcSpec()
	}
	engine, err := xorm.NewEngine(driver, spec)
	if err != nil {
		panic(err)
	}
	Engine = engine
}

func init() {
	r.OnAppStart(Init)
}
<commit_msg>db.go: print the spec we are using to connect to mysql<commit_after>package db

import (
	"fmt"
	"github.com/lunny/xorm"
	r "github.com/robfig/revel"
)

var (
	Engine *xorm.Engine
)

func ensureOption(option string) string {
	value, found := r.Config.String(option)
	if !found {
		r.ERROR.Fatalf("Option %v not found", option)
	}
	return value
}

func calcSpec() string {
	host := ensureOption("db.host")
	name := ensureOption("db.name")
	username := ensureOption("db.username")
	password := r.Config.StringDefault("db.password", "")
	return fmt.Sprintf("%v:%v@tcp(%v)/%v", username, password, host, name)
}

func Init() {
	var found bool
	var driver string
	if driver, found = r.Config.String("db.driver"); !found {
		r.ERROR.Fatal("No db.driver found.")
	}
	var spec string
	if spec, found = r.Config.String("db.spec"); !found {
		spec = calcSpec()
	}
	r.INFO.Printf("Connecting to mysql at %v", spec)
	engine, err := xorm.NewEngine(driver, spec)
	if err != nil {
		panic(err)
	}
	Engine = engine
}

func init() {
	r.OnAppStart(Init)
}
<|endoftext|>


<commit_before>// Package xmltree converts XML documents as a tree of Go structs.
//
// The xmltree package provides routines for accessing an XML document
// as a tree, along with functionality to resolve namespace-prefixed
// strings at any point in the tree.
package xmltree // import "aqwari.net/xml/xmltree"

import (
	"bytes"
	"encoding/xml"
	"errors"
	"fmt"
	"strings"
)

const (
	xmlNamespaceURI = "http://www.w3.org/2000/xmlns/"
	xmlLangURI      = "http://www.w3.org/XML/1998/namespace"
	recursionLimit  = 3000
)

var errDeepXML = errors.New("xmltree: xml document too deeply nested")

// An Element represents a single element in an XML document. Elements
// may have zero or more children. The byte array used by the Content
// field is shared among all elements in the document, and should not
// be modified. An Element also captures xml namespace prefixes, so
// that arbitrary QNames in attribute values can be resolved.
type Element struct {
	xml.StartElement
	// The XML namespace scope at this element's location in the
	// document.
	Scope
	// The raw content contained within this element's start and
	// end tags. Uses the underlying byte array passed to Parse.
	Content []byte
	// Sub-elements contained within this element.
	Children []Element
}

// Attr gets the value of the first attribute whose name matches the
// space and local arguments. If space is the empty string, only
// attributes' local names are considered when looking for a match.
// If an attribute could not be found, the empty string is returned.
func (el *Element) Attr(space, local string) string {
	for _, v := range el.StartElement.Attr {
		if v.Name.Local != local {
			continue
		}
		if space == "" || space == v.Name.Space {
			return v.Value
		}
	}
	return ""
}

// The JoinScope method joins two Scopes together. When resolving
// prefixes using the returned scope, the prefix list in the argument
// Scope is searched before that of the receiver Scope.
func (outer *Scope) JoinScope(inner *Scope) *Scope {
	return &Scope{append(outer.ns, inner.ns...)}
}

// Unmarshal parses the XML encoding of the Element and stores the result
// in the value pointed to by v. Unmarshal follows the same rules as
// xml.Unmarshal, but only parses the portion of the XML document
// contained by the Element.
func (el *Element) Unmarshal(v interface{}) error {
	return xml.Unmarshal(Marshal(el), v)
}

// A Scope represents the xml namespace scope at a given position in
// the document.
type Scope struct {
	ns []xml.Name
}

// Resolve translates an XML QName (namespace-prefixed string) to an
// xml.Name with a canonicalized namespace in its Space field.  This can
// be used when working with XSD documents, which put QNames in attribute
// values. If qname does not have a prefix, the default namespace is used.If
// a namespace prefix cannot be resolved, the returned value's Space field
// will be the unresolved prefix. Use the ResolveNS function to detect when
// a namespace prefix cannot be resolved.
func (scope *Scope) Resolve(qname string) xml.Name {
	name, _ := scope.ResolveNS(qname)
	return name
}

// The ResolveNS method is like Resolve, but returns false for its second
// return value if a namespace prefix cannot be resolved.
func (scope *Scope) ResolveNS(qname string) (xml.Name, bool) {
	var prefix, local string
	parts := strings.SplitN(qname, ":", 2)
	if len(parts) == 2 {
		prefix, local = parts[0], parts[1]
	} else {
		prefix, local = "", parts[0]
	}
	switch prefix {
	case "xml":
		return xml.Name{Space: xmlLangURI, Local: local}, true
	case "xmlns":
		return xml.Name{Space: xmlNamespaceURI, Local: local}, true
	}
	for i := len(scope.ns) - 1; i >= 0; i-- {
		if scope.ns[i].Local == prefix {
			return xml.Name{Space: scope.ns[i].Space, Local: local}, true
		}
	}
	return xml.Name{Space: prefix, Local: local}, false
}

// ResolveDefault is like Resolve, but allows for the default namespace to
// be overridden. The namespace of strings without a namespace prefix
// (known as an NCName in XML terminology) will be defaultns.
func (scope *Scope) ResolveDefault(qname, defaultns string) xml.Name {
	if defaultns == "" || strings.Contains(qname, ":") {
		return scope.Resolve(qname)
	}
	return xml.Name{defaultns, qname}
}

// Prefix is the inverse of Resolve. It uses the closest prefix
// defined for a namespace to create a string of the form
// prefix:local. If the namespace cannot be found, or is the
// default namespace, an unqualified name is returned.
func (scope *Scope) Prefix(name xml.Name) (qname string) {
	// print("LOOKUP ", name.Space, " ", name.Local)
	switch name.Space {
	case "":
		// println(" ->", name.Local)
		return name.Local
	case xmlLangURI:
		// println(" ->", "xml:"+name.Local)
		return "xml:" + name.Local
	case xmlNamespaceURI:
		// println(" ->", "xmlns:"+name.Local)
		return "xmlns:" + name.Local
	}
	for i := len(scope.ns) - 1; i >= 0; i-- {
		if scope.ns[i].Space == name.Space {
			if scope.ns[i].Local == "" {
				// println(" ->", name.Local, "(default)")
				return name.Local
			}
			// println(" ->", scope.ns[i].Local+":"+name.Local)
			return scope.ns[i].Local + ":" + name.Local
		}
	}
	// println(" ->", name.Local)
	return name.Local
}

func (scope *Scope) pushNS(tag xml.StartElement) {
	var ns []xml.Name
	for _, attr := range tag.Attr {
		if attr.Name.Space == "xmlns" {
			ns = append(ns, xml.Name{attr.Value, attr.Name.Local})
		} else if attr.Name.Local == "xmlns" {
			ns = append(ns, xml.Name{attr.Value, ""})
		} else {
			continue
		}
	}
	if len(ns) > 0 {
		scope.ns = append(scope.ns, ns...)
		// Ensure that future additions to the scope create
		// a new backing array. This prevents the scope from
		// being clobbered during parsing.
		scope.ns = scope.ns[:len(scope.ns):len(scope.ns)]
	}
}

// Save some typing when scanning xml
type scanner struct {
	*xml.Decoder
	tok xml.Token
	err error
}

func (s *scanner) scan() bool {
	if s.err != nil {
		return false
	}
	s.tok, s.err = s.Token()
	return s.err == nil
}

// Parse builds a tree of Elements by reading an XML document.  The
// byte slice passed to Parse is expected to be a valid XML document
// with a single root element.
func Parse(doc []byte) (*Element, error) {
	d := xml.NewDecoder(bytes.NewReader(doc))
	scanner := scanner{Decoder: d}
	root := new(Element)

	for scanner.scan() {
		if start, ok := scanner.tok.(xml.StartElement); ok {
			root.StartElement = start
			break
		}
	}
	if scanner.err != nil {
		return nil, scanner.err
	}
	if err := root.parse(&scanner, doc, 0); err != nil {
		return nil, err
	}
	return root, nil
}

func (el *Element) parse(scanner *scanner, data []byte, depth int) error {
	if depth > recursionLimit {
		return errDeepXML
	}
	el.pushNS(el.StartElement)

	begin := scanner.InputOffset()
	end := begin
walk:
	for scanner.scan() {
		switch tok := scanner.tok.(type) {
		case xml.StartElement:
			child := Element{StartElement: tok.Copy(), Scope: el.Scope}
			if err := child.parse(scanner, data, depth+1); err != nil {
				return err
			}
			el.Children = append(el.Children, child)
		case xml.EndElement:
			if tok.Name != el.Name {
				return fmt.Errorf("Expecting </%s>, got </%s>", el.Prefix(el.Name), el.Prefix(tok.Name))
			}
			el.Content = data[int(begin):int(end)]
			break walk
		}
		end = scanner.InputOffset()
	}
	return scanner.err
}

// The walk method calls the walkFunc for each of the Element's children.
// If the WalkFunc returns a non-nil error, Walk will return it
// immediately.
func (el *Element) walk(fn walkFunc) error {
	for i := 0; i < len(el.Children); i++ {
		fn(&el.Children[i])
	}
	return nil
}

// SetAttr adds an XML attribute to an Element's existing Attributes.
// If the attribute already exists, it is replaced.
func (el *Element) SetAttr(space, local, value string) {
	for i, a := range el.StartElement.Attr {
		if a.Name.Local != local {
			continue
		}
		if space == "" || a.Name.Space == space {
			el.StartElement.Attr[i].Value = value
			return
		}
	}
	el.StartElement.Attr = append(el.StartElement.Attr, xml.Attr{
		Name:  xml.Name{space, local},
		Value: value,
	})
}

// walkFunc is the type of the function called for each of an Element's
// children.
type walkFunc func(*Element)

// SearchFunc traverses the Element tree in depth-first order and returns
// a slice of Elements for which the function fn returns true. Note that
// SearchFunc does not search the children of Elements that match the search;
// there is no parent-child relationship between the Elements returned in
// the result.
func (root *Element) SearchFunc(fn func(*Element) bool) []*Element {
	var results []*Element
	var search func(el *Element)

	search = func(el *Element) {
		if fn(el) {
			results = append(results, el)
		}
		el.walk(search)
	}
	root.walk(search)
	return results
}

// Search searches the Element tree for Elements with an xml tag
// matching the name and xml namespace. If space is the empty string,
// any namespace is matched.
func (root *Element) Search(space, local string) []*Element {
	return root.SearchFunc(func(el *Element) bool {
		if local != el.Name.Local {
			return false
		}
		return space == "" || space == el.Name.Space
	})
}
<commit_msg>Favor default namespace when generating QNames<commit_after>// Package xmltree converts XML documents as a tree of Go structs.
//
// The xmltree package provides routines for accessing an XML document
// as a tree, along with functionality to resolve namespace-prefixed
// strings at any point in the tree.
package xmltree // import "aqwari.net/xml/xmltree"

import (
	"bytes"
	"encoding/xml"
	"errors"
	"fmt"
	"strings"
)

const (
	xmlNamespaceURI = "http://www.w3.org/2000/xmlns/"
	xmlLangURI      = "http://www.w3.org/XML/1998/namespace"
	recursionLimit  = 3000
)

var errDeepXML = errors.New("xmltree: xml document too deeply nested")

// An Element represents a single element in an XML document. Elements
// may have zero or more children. The byte array used by the Content
// field is shared among all elements in the document, and should not
// be modified. An Element also captures xml namespace prefixes, so
// that arbitrary QNames in attribute values can be resolved.
type Element struct {
	xml.StartElement
	// The XML namespace scope at this element's location in the
	// document.
	Scope
	// The raw content contained within this element's start and
	// end tags. Uses the underlying byte array passed to Parse.
	Content []byte
	// Sub-elements contained within this element.
	Children []Element
}

// Attr gets the value of the first attribute whose name matches the
// space and local arguments. If space is the empty string, only
// attributes' local names are considered when looking for a match.
// If an attribute could not be found, the empty string is returned.
func (el *Element) Attr(space, local string) string {
	for _, v := range el.StartElement.Attr {
		if v.Name.Local != local {
			continue
		}
		if space == "" || space == v.Name.Space {
			return v.Value
		}
	}
	return ""
}

// The JoinScope method joins two Scopes together. When resolving
// prefixes using the returned scope, the prefix list in the argument
// Scope is searched before that of the receiver Scope.
func (outer *Scope) JoinScope(inner *Scope) *Scope {
	return &Scope{append(outer.ns, inner.ns...)}
}

// Unmarshal parses the XML encoding of the Element and stores the result
// in the value pointed to by v. Unmarshal follows the same rules as
// xml.Unmarshal, but only parses the portion of the XML document
// contained by the Element.
func (el *Element) Unmarshal(v interface{}) error {
	return xml.Unmarshal(Marshal(el), v)
}

// A Scope represents the xml namespace scope at a given position in
// the document.
type Scope struct {
	ns []xml.Name
}

// Resolve translates an XML QName (namespace-prefixed string) to an
// xml.Name with a canonicalized namespace in its Space field.  This can
// be used when working with XSD documents, which put QNames in attribute
// values. If qname does not have a prefix, the default namespace is used.If
// a namespace prefix cannot be resolved, the returned value's Space field
// will be the unresolved prefix. Use the ResolveNS function to detect when
// a namespace prefix cannot be resolved.
func (scope *Scope) Resolve(qname string) xml.Name {
	name, _ := scope.ResolveNS(qname)
	return name
}

// The ResolveNS method is like Resolve, but returns false for its second
// return value if a namespace prefix cannot be resolved.
func (scope *Scope) ResolveNS(qname string) (xml.Name, bool) {
	var prefix, local string
	parts := strings.SplitN(qname, ":", 2)
	if len(parts) == 2 {
		prefix, local = parts[0], parts[1]
	} else {
		prefix, local = "", parts[0]
	}
	switch prefix {
	case "xml":
		return xml.Name{Space: xmlLangURI, Local: local}, true
	case "xmlns":
		return xml.Name{Space: xmlNamespaceURI, Local: local}, true
	}
	for i := len(scope.ns) - 1; i >= 0; i-- {
		if scope.ns[i].Local == prefix {
			return xml.Name{Space: scope.ns[i].Space, Local: local}, true
		}
	}
	return xml.Name{Space: prefix, Local: local}, false
}

// ResolveDefault is like Resolve, but allows for the default namespace to
// be overridden. The namespace of strings without a namespace prefix
// (known as an NCName in XML terminology) will be defaultns.
func (scope *Scope) ResolveDefault(qname, defaultns string) xml.Name {
	if defaultns == "" || strings.Contains(qname, ":") {
		return scope.Resolve(qname)
	}
	return xml.Name{defaultns, qname}
}

// Prefix is the inverse of Resolve. It uses the closest prefix
// defined for a namespace to create a string of the form
// prefix:local. If the namespace cannot be found, or is the
// default namespace, an unqualified name is returned.
func (scope *Scope) Prefix(name xml.Name) (qname string) {
	switch name.Space {
	case "":
		return name.Local
	case xmlLangURI:
		return "xml:" + name.Local
	case xmlNamespaceURI:
		return "xmlns:" + name.Local
	}
	for i := len(scope.ns) - 1; i >= 0; i-- {
		if scope.ns[i].Space == name.Space {
			if scope.ns[i].Local == "" {
				// Favor default NS if there is an extra
				// qualified NS declaration
				qname = name.Local
			} else if len(qname) == 0 {
				qname = scope.ns[i].Local + ":" + name.Local
			}
		}
	}
	return qname
}

func (scope *Scope) pushNS(tag xml.StartElement) {
	var ns []xml.Name
	for _, attr := range tag.Attr {
		if attr.Name.Space == "xmlns" {
			ns = append(ns, xml.Name{attr.Value, attr.Name.Local})
		} else if attr.Name.Local == "xmlns" {
			ns = append(ns, xml.Name{attr.Value, ""})
		} else {
			continue
		}
	}
	if len(ns) > 0 {
		scope.ns = append(scope.ns, ns...)
		// Ensure that future additions to the scope create
		// a new backing array. This prevents the scope from
		// being clobbered during parsing.
		scope.ns = scope.ns[:len(scope.ns):len(scope.ns)]
	}
}

// Save some typing when scanning xml
type scanner struct {
	*xml.Decoder
	tok xml.Token
	err error
}

func (s *scanner) scan() bool {
	if s.err != nil {
		return false
	}
	s.tok, s.err = s.Token()
	return s.err == nil
}

// Parse builds a tree of Elements by reading an XML document.  The
// byte slice passed to Parse is expected to be a valid XML document
// with a single root element.
func Parse(doc []byte) (*Element, error) {
	d := xml.NewDecoder(bytes.NewReader(doc))
	scanner := scanner{Decoder: d}
	root := new(Element)

	for scanner.scan() {
		if start, ok := scanner.tok.(xml.StartElement); ok {
			root.StartElement = start
			break
		}
	}
	if scanner.err != nil {
		return nil, scanner.err
	}
	if err := root.parse(&scanner, doc, 0); err != nil {
		return nil, err
	}
	return root, nil
}

func (el *Element) parse(scanner *scanner, data []byte, depth int) error {
	if depth > recursionLimit {
		return errDeepXML
	}
	el.pushNS(el.StartElement)

	begin := scanner.InputOffset()
	end := begin
walk:
	for scanner.scan() {
		switch tok := scanner.tok.(type) {
		case xml.StartElement:
			child := Element{StartElement: tok.Copy(), Scope: el.Scope}
			if err := child.parse(scanner, data, depth+1); err != nil {
				return err
			}
			el.Children = append(el.Children, child)
		case xml.EndElement:
			if tok.Name != el.Name {
				return fmt.Errorf("Expecting </%s>, got </%s>", el.Prefix(el.Name), el.Prefix(tok.Name))
			}
			el.Content = data[int(begin):int(end)]
			break walk
		}
		end = scanner.InputOffset()
	}
	return scanner.err
}

// The walk method calls the walkFunc for each of the Element's children.
// If the WalkFunc returns a non-nil error, Walk will return it
// immediately.
func (el *Element) walk(fn walkFunc) error {
	for i := 0; i < len(el.Children); i++ {
		fn(&el.Children[i])
	}
	return nil
}

// SetAttr adds an XML attribute to an Element's existing Attributes.
// If the attribute already exists, it is replaced.
func (el *Element) SetAttr(space, local, value string) {
	for i, a := range el.StartElement.Attr {
		if a.Name.Local != local {
			continue
		}
		if space == "" || a.Name.Space == space {
			el.StartElement.Attr[i].Value = value
			return
		}
	}
	el.StartElement.Attr = append(el.StartElement.Attr, xml.Attr{
		Name:  xml.Name{space, local},
		Value: value,
	})
}

// walkFunc is the type of the function called for each of an Element's
// children.
type walkFunc func(*Element)

// SearchFunc traverses the Element tree in depth-first order and returns
// a slice of Elements for which the function fn returns true. Note that
// SearchFunc does not search the children of Elements that match the search;
// there is no parent-child relationship between the Elements returned in
// the result.
func (root *Element) SearchFunc(fn func(*Element) bool) []*Element {
	var results []*Element
	var search func(el *Element)

	search = func(el *Element) {
		if fn(el) {
			results = append(results, el)
		}
		el.walk(search)
	}
	root.walk(search)
	return results
}

// Search searches the Element tree for Elements with an xml tag
// matching the name and xml namespace. If space is the empty string,
// any namespace is matched.
func (root *Element) Search(space, local string) []*Element {
	return root.SearchFunc(func(el *Element) bool {
		if local != el.Name.Local {
			return false
		}
		return space == "" || space == el.Name.Space
	})
}
<|endoftext|>


<commit_before>/*
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 main

import (
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"strconv"
	"strings"
	"time"

	"github.com/porjo/go-flowrate/flowrate"
	"golang.org/x/net/context"
	"golang.org/x/oauth2"
	"google.golang.org/api/googleapi"
	"google.golang.org/api/youtube/v3"
)

var (
	filename     = flag.String("filename", "", "Filename to upload. Can be a URL")
	title        = flag.String("title", "Video Title", "Video title")
	description  = flag.String("description", "uploaded by youtubeuploader", "Video description")
	category     = flag.String("category", "", "Video category")
	keywords     = flag.String("keywords", "", "Comma separated list of video keywords")
	privacy      = flag.String("privacy", "private", "Video privacy status")
	showProgress = flag.Bool("progress", false, "Show progress indicator")
	rate         = flag.Int("ratelimit", 0, "Rate limit upload in KiB/s. No limit by default")
)

func main() {
	flag.Parse()

	if *filename == "" {
		fmt.Printf("You must provide a filename of a video file to upload\n")
		flag.PrintDefaults()
		os.Exit(1)
	}

	var reader io.Reader
	var filesize int64

	if strings.HasPrefix(*filename, "http") {
		resp, err := http.Head(*filename)
		if err != nil {
			log.Fatalf("Error opening %v: %v", *filename, err)
		}
		lenStr := resp.Header.Get("content-length")
		if lenStr != "" {
			filesize, err = strconv.ParseInt(lenStr, 10, 64)
			if err != nil {
				log.Fatal(err)
			}
		}

		resp, err = http.Get(*filename)
		if err != nil {
			log.Fatalf("Error opening %v: %v", *filename, err)
		}
		reader = resp.Body
		filesize = resp.ContentLength
		defer resp.Body.Close()
	} else {
		file, err := os.Open(*filename)
		if err != nil {
			log.Fatalf("Error opening %v: %v", *filename, err)
		}
		fileInfo, err := file.Stat()
		if err != nil {
			log.Fatalf("Error stating file %v: %v", *filename, err)
		}
		filesize = fileInfo.Size()
		reader = file
		defer file.Close()
	}

	ctx := context.Background()
	transport := &limitTransport{rt: http.DefaultTransport, filesize: filesize}
	ctx = context.WithValue(ctx, oauth2.HTTPClient, &http.Client{
		Transport: transport,
	})

	if *showProgress {
		ticker := time.NewTicker(time.Millisecond * 500).C
		quitChan := make(chan bool)
		defer func() {
			quitChan <- true
		}()
		go func() {
			for {
				select {
				case <-ticker:
					if transport.reader != nil {
						s := transport.reader.Monitor.Status()
						fmt.Printf("\rProgress: %.2f Mbps, %d / %d (%s) ETA %s", float32(s.CurRate*8)/(1000*1000), s.Bytes, filesize, s.Progress, s.TimeRem)
					}
				case <-quitChan:
					return
				}
			}
		}()
	}
	client, err := buildOAuthHTTPClient(ctx, youtube.YoutubeUploadScope)
	if err != nil {
		log.Fatalf("Error building OAuth client: %v", err)
	}

	service, err := youtube.New(client)
	if err != nil {
		log.Fatalf("Error creating YouTube client: %v", err)
	}

	upload := &youtube.Video{
		Snippet: &youtube.VideoSnippet{
			Title:       *title,
			Description: *description,
			CategoryId:  *category,
		},
		Status: &youtube.VideoStatus{PrivacyStatus: *privacy},
	}

	// The API returns a 400 Bad Request response if tags is an empty string.
	if strings.Trim(*keywords, "") != "" {
		upload.Snippet.Tags = strings.Split(*keywords, ",")
	}

	call := service.Videos.Insert("snippet,status", upload)

	var option googleapi.MediaOption
	var video *youtube.Video

	// our RoundTrip gets bypassed if the filesize < DefaultUploadChunkSize
	if googleapi.DefaultUploadChunkSize > filesize {
		option = googleapi.ChunkSize(int(filesize / 2))
	} else {
		option = googleapi.ChunkSize(googleapi.DefaultUploadChunkSize)
	}
	video, err = call.Media(reader, option).Do()
	if err != nil {
		if video != nil {
			log.Fatalf("Error making YouTube API call: %v, %v", err, video.HTTPStatusCode)
		} else {
			log.Fatalf("Error making YouTube API call: %v", err)
		}
	}
	fmt.Printf("\nUpload successful! Video ID: %v\n", video.Id)
}

type limitTransport struct {
	rt       http.RoundTripper
	reader   *flowrate.Reader
	filesize int64
}

func (t *limitTransport) RoundTrip(r *http.Request) (res *http.Response, err error) {

	// FIXME need a better way to detect which roundtrip is the media upload
	if r.ContentLength > 1000 {
		var monitor *flowrate.Monitor

		if t.reader != nil {
			monitor = t.reader.Monitor
		}
		t.reader = flowrate.NewReader(r.Body, int64(*rate*1000))

		if monitor != nil {
			// carry over stats to new limiter
			t.reader.Monitor = monitor
		} else {
			t.reader.Monitor.SetTransferSize(t.filesize)
		}
		r.Body = ioutil.NopCloser(t.reader)
	}

	return t.rt.RoundTrip(r)
}
<commit_msg>tweaks<commit_after>/*
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 main

import (
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"strconv"
	"strings"
	"time"

	"github.com/porjo/go-flowrate/flowrate"
	"golang.org/x/net/context"
	"golang.org/x/oauth2"
	"google.golang.org/api/googleapi"
	"google.golang.org/api/youtube/v3"
)

var (
	filename    = flag.String("filename", "", "Filename to upload. Can be a URL")
	title       = flag.String("title", "Video Title", "Video title")
	description = flag.String("description", "uploaded by youtubeuploader", "Video description")
	category    = flag.String("category", "", "Video category")
	keywords    = flag.String("keywords", "", "Comma separated list of video keywords")
	privacy     = flag.String("privacy", "private", "Video privacy status")
	quiet       = flag.Bool("quiet", false, "Suppress progress indicator")
	rate        = flag.Int("ratelimit", 0, "Rate limit upload in KiB/s. No limit by default")
)

func main() {
	flag.Parse()

	if *filename == "" {
		fmt.Printf("You must provide a filename of a video file to upload\n")
		flag.PrintDefaults()
		os.Exit(1)
	}

	var reader io.Reader
	var filesize int64

	if strings.HasPrefix(*filename, "http") {
		resp, err := http.Head(*filename)
		if err != nil {
			log.Fatalf("Error opening %v: %v", *filename, err)
		}
		lenStr := resp.Header.Get("content-length")
		if lenStr != "" {
			filesize, err = strconv.ParseInt(lenStr, 10, 64)
			if err != nil {
				log.Fatal(err)
			}
		}

		resp, err = http.Get(*filename)
		if err != nil {
			log.Fatalf("Error opening %v: %v", *filename, err)
		}
		reader = resp.Body
		filesize = resp.ContentLength
		defer resp.Body.Close()
	} else {
		file, err := os.Open(*filename)
		if err != nil {
			log.Fatalf("Error opening %v: %v", *filename, err)
		}
		fileInfo, err := file.Stat()
		if err != nil {
			log.Fatalf("Error stating file %v: %v", *filename, err)
		}
		filesize = fileInfo.Size()
		reader = file
		defer file.Close()
	}

	ctx := context.Background()
	transport := &limitTransport{rt: http.DefaultTransport, filesize: filesize}
	ctx = context.WithValue(ctx, oauth2.HTTPClient, &http.Client{
		Transport: transport,
	})

	if !*quiet {
		ticker := time.NewTicker(time.Millisecond * 500).C
		quitChan := make(chan bool)
		defer func() {
			quitChan <- true
		}()
		go func() {
			for {
				select {
				case <-ticker:
					if transport.reader != nil {
						s := transport.reader.Monitor.Status()
						fmt.Printf("\rProgress: %.2f Mbps, %d / %d (%s) ETA %s", float32(s.CurRate*8)/(1000*1000), s.Bytes, filesize, s.Progress, s.TimeRem)
					}
				case <-quitChan:
					return
				}
			}
		}()
	}
	client, err := buildOAuthHTTPClient(ctx, youtube.YoutubeUploadScope)
	if err != nil {
		log.Fatalf("Error building OAuth client: %v", err)
	}

	service, err := youtube.New(client)
	if err != nil {
		log.Fatalf("Error creating YouTube client: %v", err)
	}

	upload := &youtube.Video{
		Snippet: &youtube.VideoSnippet{
			Title:       *title,
			Description: *description,
			CategoryId:  *category,
		},
		Status: &youtube.VideoStatus{PrivacyStatus: *privacy},
	}

	// The API returns a 400 Bad Request response if tags is an empty string.
	if strings.Trim(*keywords, "") != "" {
		upload.Snippet.Tags = strings.Split(*keywords, ",")
	}

	call := service.Videos.Insert("snippet,status", upload)

	var option googleapi.MediaOption
	var video *youtube.Video

	// our RoundTrip gets bypassed if the filesize < DefaultUploadChunkSize
	if googleapi.DefaultUploadChunkSize > filesize {
		option = googleapi.ChunkSize(int(filesize / 2))
	} else {
		option = googleapi.ChunkSize(googleapi.DefaultUploadChunkSize)
	}

	fmt.Printf("Uploading file '%s'...\n", *filename)

	video, err = call.Media(reader, option).Do()
	if err != nil {
		if video != nil {
			log.Fatalf("Error making YouTube API call: %v, %v", err, video.HTTPStatusCode)
		} else {
			log.Fatalf("Error making YouTube API call: %v", err)
		}
	}
	fmt.Printf("\nUpload successful! Video ID: %v\n", video.Id)
}

type limitTransport struct {
	rt       http.RoundTripper
	reader   *flowrate.Reader
	filesize int64
}

func (t *limitTransport) RoundTrip(r *http.Request) (res *http.Response, err error) {

	// FIXME need a better way to detect which roundtrip is the media upload
	if r.ContentLength > 1000 {
		var monitor *flowrate.Monitor

		if t.reader != nil {
			monitor = t.reader.Monitor
		}
		t.reader = flowrate.NewReader(r.Body, int64(*rate*1000))

		if monitor != nil {
			// carry over stats to new limiter
			t.reader.Monitor = monitor
		} else {
			t.reader.Monitor.SetTransferSize(t.filesize)
		}
		r.Body = ioutil.NopCloser(t.reader)
	}

	return t.rt.RoundTrip(r)
}
<|endoftext|>


<commit_before>package docker

import (
	"archive/tar"
	"bytes"
	"errors"
	"github.com/dotcloud/docker/archive"
	"github.com/dotcloud/docker/graphdriver"
	"github.com/dotcloud/docker/utils"
	"io"
	"io/ioutil"
	"os"
	"testing"
	"time"
)

func TestInit(t *testing.T) {
	graph := tempGraph(t)
	defer os.RemoveAll(graph.Root)
	// Root should exist
	if _, err := os.Stat(graph.Root); err != nil {
		t.Fatal(err)
	}
	// Map() should be empty
	if l, err := graph.Map(); err != nil {
		t.Fatal(err)
	} else if len(l) != 0 {
		t.Fatalf("len(Map()) should return %d, not %d", 0, len(l))
	}
}

// Test that Register can be interrupted cleanly without side effects
func TestInterruptedRegister(t *testing.T) {
	graph := tempGraph(t)
	defer os.RemoveAll(graph.Root)
	badArchive, w := io.Pipe() // Use a pipe reader as a fake archive which never yields data
	image := &Image{
		ID:      GenerateID(),
		Comment: "testing",
		Created: time.Now(),
	}
	go graph.Register(nil, badArchive, image)
	time.Sleep(200 * time.Millisecond)
	w.CloseWithError(errors.New("But I'm not a tarball!")) // (Nobody's perfect, darling)
	if _, err := graph.Get(image.ID); err == nil {
		t.Fatal("Image should not exist after Register is interrupted")
	}
	// Registering the same image again should succeed if the first register was interrupted
	goodArchive, err := fakeTar()
	if err != nil {
		t.Fatal(err)
	}
	if err := graph.Register(nil, goodArchive, image); err != nil {
		t.Fatal(err)
	}
}

// FIXME: Do more extensive tests (ex: create multiple, delete, recreate;
//       create multiple, check the amount of images and paths, etc..)
func TestGraphCreate(t *testing.T) {
	graph := tempGraph(t)
	defer os.RemoveAll(graph.Root)
	archive, err := fakeTar()
	if err != nil {
		t.Fatal(err)
	}
	image, err := graph.Create(archive, nil, "Testing", "", nil)
	if err != nil {
		t.Fatal(err)
	}
	if err := ValidateID(image.ID); err != nil {
		t.Fatal(err)
	}
	if image.Comment != "Testing" {
		t.Fatalf("Wrong comment: should be '%s', not '%s'", "Testing", image.Comment)
	}
	if image.DockerVersion != VERSION {
		t.Fatalf("Wrong docker_version: should be '%s', not '%s'", VERSION, image.DockerVersion)
	}
	images, err := graph.Map()
	if err != nil {
		t.Fatal(err)
	} else if l := len(images); l != 1 {
		t.Fatalf("Wrong number of images. Should be %d, not %d", 1, l)
	}
	if images[image.ID] == nil {
		t.Fatalf("Could not find image with id %s", image.ID)
	}
}

func TestRegister(t *testing.T) {
	graph := tempGraph(t)
	defer os.RemoveAll(graph.Root)
	archive, err := fakeTar()
	if err != nil {
		t.Fatal(err)
	}
	image := &Image{
		ID:      GenerateID(),
		Comment: "testing",
		Created: time.Now(),
	}
	err = graph.Register(nil, archive, image)
	if err != nil {
		t.Fatal(err)
	}
	if images, err := graph.Map(); err != nil {
		t.Fatal(err)
	} else if l := len(images); l != 1 {
		t.Fatalf("Wrong number of images. Should be %d, not %d", 1, l)
	}
	if resultImg, err := graph.Get(image.ID); err != nil {
		t.Fatal(err)
	} else {
		if resultImg.ID != image.ID {
			t.Fatalf("Wrong image ID. Should be '%s', not '%s'", image.ID, resultImg.ID)
		}
		if resultImg.Comment != image.Comment {
			t.Fatalf("Wrong image comment. Should be '%s', not '%s'", image.Comment, resultImg.Comment)
		}
	}
}

// Test that an image can be deleted by its shorthand prefix
func TestDeletePrefix(t *testing.T) {
	graph := tempGraph(t)
	defer os.RemoveAll(graph.Root)
	img := createTestImage(graph, t)
	if err := graph.Delete(utils.TruncateID(img.ID)); err != nil {
		t.Fatal(err)
	}
	assertNImages(graph, t, 0)
}

func createTestImage(graph *Graph, t *testing.T) *Image {
	archive, err := fakeTar()
	if err != nil {
		t.Fatal(err)
	}
	img, err := graph.Create(archive, nil, "Test image", "", nil)
	if err != nil {
		t.Fatal(err)
	}
	return img
}

func TestDelete(t *testing.T) {
	graph := tempGraph(t)
	defer os.RemoveAll(graph.Root)
	archive, err := fakeTar()
	if err != nil {
		t.Fatal(err)
	}
	assertNImages(graph, t, 0)
	img, err := graph.Create(archive, nil, "Bla bla", "", nil)
	if err != nil {
		t.Fatal(err)
	}
	assertNImages(graph, t, 1)
	if err := graph.Delete(img.ID); err != nil {
		t.Fatal(err)
	}
	assertNImages(graph, t, 0)

	archive, err = fakeTar()
	if err != nil {
		t.Fatal(err)
	}
	// Test 2 create (same name) / 1 delete
	img1, err := graph.Create(archive, nil, "Testing", "", nil)
	if err != nil {
		t.Fatal(err)
	}
	archive, err = fakeTar()
	if err != nil {
		t.Fatal(err)
	}
	if _, err = graph.Create(archive, nil, "Testing", "", nil); err != nil {
		t.Fatal(err)
	}
	assertNImages(graph, t, 2)
	if err := graph.Delete(img1.ID); err != nil {
		t.Fatal(err)
	}
	assertNImages(graph, t, 1)

	// Test delete wrong name
	if err := graph.Delete("Not_foo"); err == nil {
		t.Fatalf("Deleting wrong ID should return an error")
	}
	assertNImages(graph, t, 1)

	archive, err = fakeTar()
	if err != nil {
		t.Fatal(err)
	}
	// Test delete twice (pull -> rm -> pull -> rm)
	if err := graph.Register(nil, archive, img1); err != nil {
		t.Fatal(err)
	}
	if err := graph.Delete(img1.ID); err != nil {
		t.Fatal(err)
	}
	assertNImages(graph, t, 1)
}

func TestByParent(t *testing.T) {
	archive1, _ := fakeTar()
	archive2, _ := fakeTar()
	archive3, _ := fakeTar()

	graph := tempGraph(t)
	defer os.RemoveAll(graph.Root)
	parentImage := &Image{
		ID:      GenerateID(),
		Comment: "parent",
		Created: time.Now(),
		Parent:  "",
	}
	childImage1 := &Image{
		ID:      GenerateID(),
		Comment: "child1",
		Created: time.Now(),
		Parent:  parentImage.ID,
	}
	childImage2 := &Image{
		ID:      GenerateID(),
		Comment: "child2",
		Created: time.Now(),
		Parent:  parentImage.ID,
	}
	_ = graph.Register(nil, archive1, parentImage)
	_ = graph.Register(nil, archive2, childImage1)
	_ = graph.Register(nil, archive3, childImage2)

	byParent, err := graph.ByParent()
	if err != nil {
		t.Fatal(err)
	}
	numChildren := len(byParent[parentImage.ID])
	if numChildren != 2 {
		t.Fatalf("Expected 2 children, found %d", numChildren)
	}
}

func assertNImages(graph *Graph, t *testing.T, n int) {
	if images, err := graph.Map(); err != nil {
		t.Fatal(err)
	} else if actualN := len(images); actualN != n {
		t.Fatalf("Expected %d images, found %d", n, actualN)
	}
}

/*
 * HELPER FUNCTIONS
 */

func tempGraph(t *testing.T) *Graph {
	tmp, err := ioutil.TempDir("", "docker-graph-")
	if err != nil {
		t.Fatal(err)
	}
	backend, err := graphdriver.New(tmp)
	if err != nil {
		t.Fatal(err)
	}
	graph, err := NewGraph(tmp, backend)
	if err != nil {
		t.Fatal(err)
	}
	return graph
}

func testArchive(t *testing.T) archive.Archive {
	archive, err := fakeTar()
	if err != nil {
		t.Fatal(err)
	}
	return archive
}

func fakeTar() (io.Reader, error) {
	content := []byte("Hello world!\n")
	buf := new(bytes.Buffer)
	tw := tar.NewWriter(buf)
	for _, name := range []string{"/etc/postgres/postgres.conf", "/etc/passwd", "/var/log/postgres/postgres.conf"} {
		hdr := new(tar.Header)
		hdr.Size = int64(len(content))
		hdr.Name = name
		if err := tw.WriteHeader(hdr); err != nil {
			return nil, err
		}
		tw.Write([]byte(content))
	}
	tw.Close()
	return buf, nil
}
<commit_msg>graph_test: Clean up drivers allocated in tempGraph()<commit_after>package docker

import (
	"archive/tar"
	"bytes"
	"errors"
	"github.com/dotcloud/docker/archive"
	"github.com/dotcloud/docker/graphdriver"
	"github.com/dotcloud/docker/utils"
	"io"
	"io/ioutil"
	"os"
	"testing"
	"time"
)

func TestInit(t *testing.T) {
	graph := tempGraph(t)
	defer nukeGraph(graph)
	// Root should exist
	if _, err := os.Stat(graph.Root); err != nil {
		t.Fatal(err)
	}
	// Map() should be empty
	if l, err := graph.Map(); err != nil {
		t.Fatal(err)
	} else if len(l) != 0 {
		t.Fatalf("len(Map()) should return %d, not %d", 0, len(l))
	}
}

// Test that Register can be interrupted cleanly without side effects
func TestInterruptedRegister(t *testing.T) {
	graph := tempGraph(t)
	defer nukeGraph(graph)
	badArchive, w := io.Pipe() // Use a pipe reader as a fake archive which never yields data
	image := &Image{
		ID:      GenerateID(),
		Comment: "testing",
		Created: time.Now(),
	}
	go graph.Register(nil, badArchive, image)
	time.Sleep(200 * time.Millisecond)
	w.CloseWithError(errors.New("But I'm not a tarball!")) // (Nobody's perfect, darling)
	if _, err := graph.Get(image.ID); err == nil {
		t.Fatal("Image should not exist after Register is interrupted")
	}
	// Registering the same image again should succeed if the first register was interrupted
	goodArchive, err := fakeTar()
	if err != nil {
		t.Fatal(err)
	}
	if err := graph.Register(nil, goodArchive, image); err != nil {
		t.Fatal(err)
	}
}

// FIXME: Do more extensive tests (ex: create multiple, delete, recreate;
//       create multiple, check the amount of images and paths, etc..)
func TestGraphCreate(t *testing.T) {
	graph := tempGraph(t)
	defer nukeGraph(graph)
	archive, err := fakeTar()
	if err != nil {
		t.Fatal(err)
	}
	image, err := graph.Create(archive, nil, "Testing", "", nil)
	if err != nil {
		t.Fatal(err)
	}
	if err := ValidateID(image.ID); err != nil {
		t.Fatal(err)
	}
	if image.Comment != "Testing" {
		t.Fatalf("Wrong comment: should be '%s', not '%s'", "Testing", image.Comment)
	}
	if image.DockerVersion != VERSION {
		t.Fatalf("Wrong docker_version: should be '%s', not '%s'", VERSION, image.DockerVersion)
	}
	images, err := graph.Map()
	if err != nil {
		t.Fatal(err)
	} else if l := len(images); l != 1 {
		t.Fatalf("Wrong number of images. Should be %d, not %d", 1, l)
	}
	if images[image.ID] == nil {
		t.Fatalf("Could not find image with id %s", image.ID)
	}
}

func TestRegister(t *testing.T) {
	graph := tempGraph(t)
	defer nukeGraph(graph)
	archive, err := fakeTar()
	if err != nil {
		t.Fatal(err)
	}
	image := &Image{
		ID:      GenerateID(),
		Comment: "testing",
		Created: time.Now(),
	}
	err = graph.Register(nil, archive, image)
	if err != nil {
		t.Fatal(err)
	}
	if images, err := graph.Map(); err != nil {
		t.Fatal(err)
	} else if l := len(images); l != 1 {
		t.Fatalf("Wrong number of images. Should be %d, not %d", 1, l)
	}
	if resultImg, err := graph.Get(image.ID); err != nil {
		t.Fatal(err)
	} else {
		if resultImg.ID != image.ID {
			t.Fatalf("Wrong image ID. Should be '%s', not '%s'", image.ID, resultImg.ID)
		}
		if resultImg.Comment != image.Comment {
			t.Fatalf("Wrong image comment. Should be '%s', not '%s'", image.Comment, resultImg.Comment)
		}
	}
}

// Test that an image can be deleted by its shorthand prefix
func TestDeletePrefix(t *testing.T) {
	graph := tempGraph(t)
	defer nukeGraph(graph)
	img := createTestImage(graph, t)
	if err := graph.Delete(utils.TruncateID(img.ID)); err != nil {
		t.Fatal(err)
	}
	assertNImages(graph, t, 0)
}

func createTestImage(graph *Graph, t *testing.T) *Image {
	archive, err := fakeTar()
	if err != nil {
		t.Fatal(err)
	}
	img, err := graph.Create(archive, nil, "Test image", "", nil)
	if err != nil {
		t.Fatal(err)
	}
	return img
}

func TestDelete(t *testing.T) {
	graph := tempGraph(t)
	defer nukeGraph(graph)
	archive, err := fakeTar()
	if err != nil {
		t.Fatal(err)
	}
	assertNImages(graph, t, 0)
	img, err := graph.Create(archive, nil, "Bla bla", "", nil)
	if err != nil {
		t.Fatal(err)
	}
	assertNImages(graph, t, 1)
	if err := graph.Delete(img.ID); err != nil {
		t.Fatal(err)
	}
	assertNImages(graph, t, 0)

	archive, err = fakeTar()
	if err != nil {
		t.Fatal(err)
	}
	// Test 2 create (same name) / 1 delete
	img1, err := graph.Create(archive, nil, "Testing", "", nil)
	if err != nil {
		t.Fatal(err)
	}
	archive, err = fakeTar()
	if err != nil {
		t.Fatal(err)
	}
	if _, err = graph.Create(archive, nil, "Testing", "", nil); err != nil {
		t.Fatal(err)
	}
	assertNImages(graph, t, 2)
	if err := graph.Delete(img1.ID); err != nil {
		t.Fatal(err)
	}
	assertNImages(graph, t, 1)

	// Test delete wrong name
	if err := graph.Delete("Not_foo"); err == nil {
		t.Fatalf("Deleting wrong ID should return an error")
	}
	assertNImages(graph, t, 1)

	archive, err = fakeTar()
	if err != nil {
		t.Fatal(err)
	}
	// Test delete twice (pull -> rm -> pull -> rm)
	if err := graph.Register(nil, archive, img1); err != nil {
		t.Fatal(err)
	}
	if err := graph.Delete(img1.ID); err != nil {
		t.Fatal(err)
	}
	assertNImages(graph, t, 1)
}

func TestByParent(t *testing.T) {
	archive1, _ := fakeTar()
	archive2, _ := fakeTar()
	archive3, _ := fakeTar()

	graph := tempGraph(t)
	defer nukeGraph(graph)
	parentImage := &Image{
		ID:      GenerateID(),
		Comment: "parent",
		Created: time.Now(),
		Parent:  "",
	}
	childImage1 := &Image{
		ID:      GenerateID(),
		Comment: "child1",
		Created: time.Now(),
		Parent:  parentImage.ID,
	}
	childImage2 := &Image{
		ID:      GenerateID(),
		Comment: "child2",
		Created: time.Now(),
		Parent:  parentImage.ID,
	}
	_ = graph.Register(nil, archive1, parentImage)
	_ = graph.Register(nil, archive2, childImage1)
	_ = graph.Register(nil, archive3, childImage2)

	byParent, err := graph.ByParent()
	if err != nil {
		t.Fatal(err)
	}
	numChildren := len(byParent[parentImage.ID])
	if numChildren != 2 {
		t.Fatalf("Expected 2 children, found %d", numChildren)
	}
}

func assertNImages(graph *Graph, t *testing.T, n int) {
	if images, err := graph.Map(); err != nil {
		t.Fatal(err)
	} else if actualN := len(images); actualN != n {
		t.Fatalf("Expected %d images, found %d", n, actualN)
	}
}

/*
 * HELPER FUNCTIONS
 */

func tempGraph(t *testing.T) *Graph {
	tmp, err := ioutil.TempDir("", "docker-graph-")
	if err != nil {
		t.Fatal(err)
	}
	backend, err := graphdriver.New(tmp)
	if err != nil {
		t.Fatal(err)
	}
	graph, err := NewGraph(tmp, backend)
	if err != nil {
		t.Fatal(err)
	}
	return graph
}

func nukeGraph(graph *Graph) {
	graph.driver.Cleanup()
	os.RemoveAll(graph.Root)
}

func testArchive(t *testing.T) archive.Archive {
	archive, err := fakeTar()
	if err != nil {
		t.Fatal(err)
	}
	return archive
}

func fakeTar() (io.Reader, error) {
	content := []byte("Hello world!\n")
	buf := new(bytes.Buffer)
	tw := tar.NewWriter(buf)
	for _, name := range []string{"/etc/postgres/postgres.conf", "/etc/passwd", "/var/log/postgres/postgres.conf"} {
		hdr := new(tar.Header)
		hdr.Size = int64(len(content))
		hdr.Name = name
		if err := tw.WriteHeader(hdr); err != nil {
			return nil, err
		}
		tw.Write([]byte(content))
	}
	tw.Close()
	return buf, nil
}
<|endoftext|>


<commit_before>//go:build windows
// +build windows

package enforcerui

import (
	"context"
	"errors"
	"fmt"
	"runtime"
	"sync"

	"github.com/lxn/walk"
	"github.com/scjalliance/resourceful/policy"
)

type trayData struct {
	Window     *walk.MainWindow
	Tray       *walk.NotifyIcon
	Completion <-chan int
	Err        error
}

// Tray is an enforcer system tray agent.
type Tray struct {
	icon    *walk.Icon
	name    string
	version string

	mutex   sync.RWMutex
	stop    context.CancelFunc
	stopped <-chan struct{}
	states  chan<- State
	notices chan<- Notice
}

// NewTray returns a new system tray instance.
func NewTray(icon *walk.Icon, name, version string) *Tray {
	return &Tray{
		icon:    icon,
		name:    name,
		version: version,
	}
}

// Start causes the tray to begin operation.
func (t *Tray) Start() error {
	t.mutex.Lock()
	defer t.mutex.Unlock()

	if t.stop != nil {
		return errors.New("the tray instance is already running")
	}

	// Create the tray objects and start the tray on a dedicated thread
	startup := make(chan trayData)
	go t.run(startup)

	// Collect information about the tray once it's been initialized
	data := <-startup
	if data.Err != nil {
		return data.Err
	}

	// Prepare a context to stop the tray on request
	ctx, stop := context.WithCancel(context.Background())
	stopped := make(chan struct{})

	// Prepare input channels
	states := make(chan State, 128)
	notices := make(chan Notice, 128)

	// Start a manager that updates the tray in response to input
	go t.manage(ctx, data.Window, data.Tray, data.Completion, stopped, states, notices)

	t.stop = stop
	t.stopped = stopped
	t.states = states
	t.notices = notices

	return nil
}

// Stop instructs the tray to cease operation and waits for it to close.
//
// The tray will stop automatically if its state channel is closed.
func (t *Tray) Stop() error {
	t.mutex.RLock()
	stop, stopped := t.stop, t.stopped
	t.mutex.RUnlock()

	if t.stop == nil {
		return errors.New("the tray instance is not running")
	}

	stop()

	//fmt.Printf("waiting for close\n")
	<-stopped

	return nil
}

// Update updates the state of the tray.
func (t *Tray) Update(state State) {
	t.mutex.RLock()
	defer t.mutex.RUnlock()
	if t.states != nil {
		select {
		case t.states <- state:
		default:
		}
	}
}

// Notify causes the tray to send a notification.
func (t *Tray) Notify(notice Notice) {
	t.mutex.RLock()
	defer t.mutex.RUnlock()
	if t.notices != nil {
		select {
		case t.notices <- notice:
		default:
		}
	}
}

func (t *Tray) run(startup chan<- trayData) {
	runtime.LockOSThread()
	defer runtime.UnlockOSThread()

	mw, ni, err := createSystemTray(t.icon)
	if err != nil {
		startup <- trayData{Err: err}
		close(startup)
	}
	defer ni.Dispose()
	defer mw.Dispose()

	completion := make(chan int)
	defer close(completion)

	startup <- trayData{
		Window:     mw,
		Tray:       ni,
		Completion: completion,
	}
	close(startup)

	result := mw.Run()

	t.mutex.Lock()
	t.stop = nil
	t.stopped = nil
	close(t.states)
	t.states = nil
	close(t.notices)
	t.notices = nil
	t.mutex.Unlock()

	completion <- result
}

func (t *Tray) manage(ctx context.Context, window *walk.MainWindow, ni *walk.NotifyIcon, completion <-chan int, stopped chan<- struct{}, states <-chan State, notices <-chan Notice) {
	defer close(stopped)
	for {
		select {
		case <-completion:
			return
		case state := <-states:
			summary := state.Summary()
			info := fmt.Sprintf("%s %s", t.name, t.version)
			window.Synchronize(func() {
				// Update tool tip
				//fmt.Println(summary)
				ni.SetToolTip(summary)

				// Update menu
				actions := ni.ContextMenu().Actions()
				actions.Clear()
				{
					action := walk.NewAction()
					action.SetText(info)
					action.SetEnabled(false)
					actions.Add(action)
				}
				actions.Add(walk.NewSeparatorAction())
				for _, pol := range state.Policies {
					action := walk.NewAction()
					desc := pol.Resource
					if name := pol.Properties["resource.name"]; name != "" {
						desc = name
					}
					if pol.Limit != policy.DefaultLimit {
						desc = fmt.Sprintf("%s: %d", desc, pol.Limit)
					}
					action.SetText(desc)
					actions.Add(action)
				}
			})
		case notice := <-notices:
			window.Synchronize(func() {
				ni.ShowCustom(
					notice.Title,
					notice.Message,
					t.icon,
				)
			})
		case <-ctx.Done():
			// Here we use the synchronize function to ensure that our call to Close
			// pushes the WM_CLOSE message onto the message queue of the correct
			// thread. If we call Close() directly it could fail silently and
			// deadlock.
			window.Synchronize(func() {
				window.Close()
			})
			<-completion
			return
		}
	}
}

func createSystemTray(icon *walk.Icon) (mw *walk.MainWindow, ni *walk.NotifyIcon, err error) {
	mw, err = walk.NewMainWindow()
	if err != nil {
		return nil, nil, fmt.Errorf("failed to create window: %v", err)
	}

	ni, err = createNotifyIcon(mw, icon)
	if err != nil {
		mw.Dispose()
		return nil, nil, fmt.Errorf("failed to create system tray: %v", err)
	}

	return mw, ni, nil
}

func createNotifyIcon(form walk.Form, icon *walk.Icon) (*walk.NotifyIcon, error) {
	ni, err := walk.NewNotifyIcon(form)
	if err != nil {
		return nil, fmt.Errorf("creation failed: %v", err)
	}

	if err := ni.SetIcon(icon); err != nil {
		ni.Dispose()
		return nil, fmt.Errorf("unable to set icon: %v", err)
	}

	if err := ni.SetVisible(true); err != nil {
		ni.Dispose()
		return nil, fmt.Errorf("unable to make the system tray visible: %v", err)
	}

	return ni, nil
}
<commit_msg>enforcerui: Include current leases holders in tray menu<commit_after>//go:build windows
// +build windows

package enforcerui

import (
	"context"
	"errors"
	"fmt"
	"runtime"
	"sync"

	"github.com/lxn/walk"
	"github.com/scjalliance/resourceful/policy"
)

type trayData struct {
	Window     *walk.MainWindow
	Tray       *walk.NotifyIcon
	Completion <-chan int
	Err        error
}

// Tray is an enforcer system tray agent.
type Tray struct {
	icon    *walk.Icon
	name    string
	version string

	mutex   sync.RWMutex
	stop    context.CancelFunc
	stopped <-chan struct{}
	states  chan<- State
	notices chan<- Notice
}

// NewTray returns a new system tray instance.
func NewTray(icon *walk.Icon, name, version string) *Tray {
	return &Tray{
		icon:    icon,
		name:    name,
		version: version,
	}
}

// Start causes the tray to begin operation.
func (t *Tray) Start() error {
	t.mutex.Lock()
	defer t.mutex.Unlock()

	if t.stop != nil {
		return errors.New("the tray instance is already running")
	}

	// Create the tray objects and start the tray on a dedicated thread
	startup := make(chan trayData)
	go t.run(startup)

	// Collect information about the tray once it's been initialized
	data := <-startup
	if data.Err != nil {
		return data.Err
	}

	// Prepare a context to stop the tray on request
	ctx, stop := context.WithCancel(context.Background())
	stopped := make(chan struct{})

	// Prepare input channels
	states := make(chan State, 128)
	notices := make(chan Notice, 128)

	// Start a manager that updates the tray in response to input
	go t.manage(ctx, data.Window, data.Tray, data.Completion, stopped, states, notices)

	t.stop = stop
	t.stopped = stopped
	t.states = states
	t.notices = notices

	return nil
}

// Stop instructs the tray to cease operation and waits for it to close.
//
// The tray will stop automatically if its state channel is closed.
func (t *Tray) Stop() error {
	t.mutex.RLock()
	stop, stopped := t.stop, t.stopped
	t.mutex.RUnlock()

	if t.stop == nil {
		return errors.New("the tray instance is not running")
	}

	stop()

	//fmt.Printf("waiting for close\n")
	<-stopped

	return nil
}

// Update updates the state of the tray.
func (t *Tray) Update(state State) {
	t.mutex.RLock()
	defer t.mutex.RUnlock()
	if t.states != nil {
		select {
		case t.states <- state:
		default:
		}
	}
}

// Notify causes the tray to send a notification.
func (t *Tray) Notify(notice Notice) {
	t.mutex.RLock()
	defer t.mutex.RUnlock()
	if t.notices != nil {
		select {
		case t.notices <- notice:
		default:
		}
	}
}

func (t *Tray) run(startup chan<- trayData) {
	runtime.LockOSThread()
	defer runtime.UnlockOSThread()

	mw, ni, err := createSystemTray(t.icon)
	if err != nil {
		startup <- trayData{Err: err}
		close(startup)
	}
	defer ni.Dispose()
	defer mw.Dispose()

	completion := make(chan int)
	defer close(completion)

	startup <- trayData{
		Window:     mw,
		Tray:       ni,
		Completion: completion,
	}
	close(startup)

	result := mw.Run()

	t.mutex.Lock()
	t.stop = nil
	t.stopped = nil
	close(t.states)
	t.states = nil
	close(t.notices)
	t.notices = nil
	t.mutex.Unlock()

	completion <- result
}

func (t *Tray) manage(ctx context.Context, window *walk.MainWindow, ni *walk.NotifyIcon, completion <-chan int, stopped chan<- struct{}, states <-chan State, notices <-chan Notice) {
	defer close(stopped)
	for {
		select {
		case <-completion:
			return
		case state := <-states:
			summary := state.Summary()
			info := fmt.Sprintf("%s %s", t.name, t.version)
			window.Synchronize(func() {
				// Update tool tip
				//fmt.Println(summary)
				ni.SetToolTip(summary)

				// Update menu
				actions := ni.ContextMenu().Actions()
				actions.Clear()
				{
					action := walk.NewAction()
					action.SetText(info)
					action.SetEnabled(false)
					actions.Add(action)
				}
				if len(state.Policies) > 0 {
					actions.Add(walk.NewSeparatorAction())
					for _, pol := range state.Policies {
						action := walk.NewAction()
						desc := pol.Resource
						if name := pol.Properties["resource.name"]; name != "" {
							desc = name
						}
						if pol.Limit != policy.DefaultLimit {
							desc = fmt.Sprintf("%s: %d", desc, pol.Limit)
						}
						action.SetText(desc)
						actions.Add(action)
						leases := state.Leases.Resource(pol.Resource)
						for _, lease := range leases {
							action := walk.NewAction()
							action.SetText(fmt.Sprintf("  %s", lease.HostUser()))
							actions.Add(action)
						}
					}
				}
			})
		case notice := <-notices:
			window.Synchronize(func() {
				ni.ShowCustom(
					notice.Title,
					notice.Message,
					t.icon,
				)
			})
		case <-ctx.Done():
			// Here we use the synchronize function to ensure that our call to Close
			// pushes the WM_CLOSE message onto the message queue of the correct
			// thread. If we call Close() directly it could fail silently and
			// deadlock.
			window.Synchronize(func() {
				window.Close()
			})
			<-completion
			return
		}
	}
}

func createSystemTray(icon *walk.Icon) (mw *walk.MainWindow, ni *walk.NotifyIcon, err error) {
	mw, err = walk.NewMainWindow()
	if err != nil {
		return nil, nil, fmt.Errorf("failed to create window: %v", err)
	}

	ni, err = createNotifyIcon(mw, icon)
	if err != nil {
		mw.Dispose()
		return nil, nil, fmt.Errorf("failed to create system tray: %v", err)
	}

	return mw, ni, nil
}

func createNotifyIcon(form walk.Form, icon *walk.Icon) (*walk.NotifyIcon, error) {
	ni, err := walk.NewNotifyIcon(form)
	if err != nil {
		return nil, fmt.Errorf("creation failed: %v", err)
	}

	if err := ni.SetIcon(icon); err != nil {
		ni.Dispose()
		return nil, fmt.Errorf("unable to set icon: %v", err)
	}

	if err := ni.SetVisible(true); err != nil {
		ni.Dispose()
		return nil, fmt.Errorf("unable to make the system tray visible: %v", err)
	}

	return ni, nil
}
<|endoftext|>


<commit_before>// Copyright 2015 Square 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 query

import (
	"errors"
	"fmt"

	"github.com/square/metrics/api"
	"github.com/square/metrics/query/aggregate"
)

// EvaluationContext is the central piece of logic, providing
// helper funcions & varaibles to evaluate a given piece of
// metrics query.
// * Contains Backend object, which can be used to fetch data
// from the backend system.s
// * Contains current timerange being queried for - this can be
// changed by say, application of time shift function.
type EvaluationContext struct {
	// Backend to fetch data from
	Backend api.Backend
	// Timerange to fetch data from
	Timerange api.Timerange
	// SampleMethod to use when up/downsampling to match the requested resolution
	SampleMethod api.SampleMethod
	Predicate    api.Predicate
}

// A value is the result of evaluating an expression.
// They can be floating point values, strings, or series lists.
type value interface {
	toSeriesList(api.Timerange) (api.SeriesList, error)
	toString() (string, error)
	toScalar() (float64, error)
}

type conversionError struct {
	from string
	to   string
}

func (e conversionError) Error() string {
	return fmt.Sprintf("cannot convert from type %s to type %s", e.from, e.to)
}

// A seriesListValue is a value which holds a SeriesList
type seriesListValue api.SeriesList

func (value seriesListValue) toSeriesList(time api.Timerange) (api.SeriesList, error) {
	return api.SeriesList(value), nil
}
func (value seriesListValue) toString() (string, error) {
	return "", conversionError{"SeriesList", "string"}
}
func (value seriesListValue) toScalar() (float64, error) {
	return 0, conversionError{"SeriesList", "scalar"}
}

// A stringValue holds a string
type stringValue string

func (value stringValue) toSeriesList(time api.Timerange) (api.SeriesList, error) {
	return api.SeriesList{}, conversionError{"string", "SeriesList"}
}
func (value stringValue) toString() (string, error) {
	return string(value), nil
}
func (value stringValue) toScalar() (float64, error) {
	return 0, conversionError{"string", "scalar"}
}

// A scalarValue holds a float and can be converted to a serieslist
type scalarValue float64

func (value scalarValue) toSeriesList(timerange api.Timerange) (api.SeriesList, error) {
	if !timerange.IsValid() {
		return api.SeriesList{}, errors.New("Invalid context.Timerange")
	}

	series := make([]float64, timerange.Slots())
	for i := range series {
		series[i] = float64(value)
	}

	return api.SeriesList{
		Series:    []api.Timeseries{api.Timeseries{series, api.NewTagSet()}},
		Timerange: timerange,
	}, nil
}

func (value scalarValue) toString() (string, error) {
	return "", conversionError{"scalar", "string"}
}

func (value scalarValue) toScalar() (float64, error) {
	return float64(value), nil
}

// Expression is a piece of code, which can be evaluated in a given
// EvaluationContext. EvaluationContext must never be changed in an Evalute().
//
// The contract of Expressions is that leaf nodes must sample a resulting
// timeseries according to the resolution specified in its EvaluationContext's
// Timerange. Internal nodes may assume that results from evaluating child
// Expressions correspond to the timerange in the current EvaluationContext.
type Expression interface {
	// Evaluate the given expression.
	Evaluate(context EvaluationContext) (value, error)
}

func evaluateToSeriesList(e Expression, context EvaluationContext) (api.SeriesList, error) {
	value, err := e.Evaluate(context)
	if err != nil {
		return api.SeriesList{}, err
	}
	return value.toSeriesList(context.Timerange)
}

// Implementations
// ===============

// Generates a Timeseries from the encapsulated scalar.
func (expr scalarExpression) Evaluate(context EvaluationContext) (value, error) {
	return scalarValue(expr.value), nil
}

func (expr *metricFetchExpression) Evaluate(context EvaluationContext) (value, error) {
	// Merge predicates appropriately
	var predicate api.Predicate
	if context.Predicate == nil {
		predicate = expr.predicate
	} else if expr.predicate == nil {
		predicate = context.Predicate
	} else {
		predicate = &andPredicate{[]api.Predicate{expr.predicate, context.Predicate}}
	}

	series, err := context.Backend.FetchSeries(api.TaggedMetric{api.MetricKey(expr.metricName), nil}, predicate, context.SampleMethod, context.Timerange)
	if err != nil {
		return seriesListValue{}, err
	}
	return seriesListValue(*series), err
}

func (expr *functionExpression) Evaluate(context EvaluationContext) (value, error) {

	name := expr.functionName

	operatorMap := map[string]func(float64, float64) float64{
		"+": func(x, y float64) float64 { return x + y },
		"-": func(x, y float64) float64 { return x - y },
		"*": func(x, y float64) float64 { return x * y },
		"/": func(x, y float64) float64 { return x / y },
	}

	if operator, ok := operatorMap[name]; ok {
		// Evaluation of a binary operator:
		// Verify that exactly 2 arguments are given.
		if len(expr.arguments) != 2 {
			return nil, errors.New(fmt.Sprintf("Function `%s` expects 2 operands but received %d (%+v)", name, len(expr.arguments), expr.arguments))
		}
		left, err := expr.arguments[0].Evaluate(context)
		if err != nil {
			return nil, err
		}
		right, err := expr.arguments[1].Evaluate(context)
		if err != nil {
			return nil, err
		}
		return evaluateBinaryOperation(context, name, left, right, operator)
	}

	if aggregator, ok := aggregate.GetAggregate(name); ok {
		// Verify that exactly 1 argument is given.
		if len(expr.arguments) != 1 {
			return nil, errors.New(fmt.Sprintf("Function `%s` expects 1 argument but received %d (%+v)", name, len(expr.arguments), expr.arguments))
		}
		argument := expr.arguments[0]
		value, err := argument.Evaluate(context)
		if err != nil {
			return nil, err
		}
		list, err := value.toSeriesList(context.Timerange)
		if err != nil {
			return nil, err
		}
		series := aggregate.AggregateBy(list, aggregator, expr.groupBy)
		return seriesListValue(series), nil
	}

	if transform, ok := GetTransformation(name); ok {
		//Verify that at least one argument is given.
		if len(expr.arguments) != 1 {
			return nil, errors.New(fmt.Sprintf("Function `%s` expects at least 1 argument but was given 0", name))
		}
		first, err := expr.arguments[0].Evaluate(context)
		if err != nil {
			return nil, err
		}
		list, err := first.toSeriesList(context.Timerange)
		if err != nil {
			return nil, err
		}
		// Evaluate all the other parameters:
		rest := expr.arguments[1:]
		parameters := make([]value, len(rest))
		for i := range parameters {
			parameters[i], err = rest[i].Evaluate(context)
			if err != nil {
				return nil, err
			}
		}
		series, err := ApplyTransform(list, transform, parameters)
		if err != nil {
			return nil, err
		}
		return seriesListValue(series), nil
	}

	return nil, errors.New(fmt.Sprintf("unknown function name `%s`", name))
}

//
// Auxiliary functions
//

// evaluateExpression wraps expr.Evaluate() to provide common messaging
// for errors. This can get pretty messy if the Expression we evaluate
// isn't a leaf node, but a leaf fails.
func evaluateExpression(context EvaluationContext, expr Expression) (value, error) {
	result, err := expr.Evaluate(context)
	if err != nil {
		return nil, errors.New(fmt.Sprintf("Evaluation of expression %+v failed:\n%s\n", expr, err.Error()))
	}
	return result, err
}

// evaluateBinaryOperation applies an arbirary binary operation to two
// Expressions.
func evaluateBinaryOperation(
	context EvaluationContext,
	functionName string,
	leftValue value,
	rightValue value,
	evaluate func(float64, float64) float64,
) (value, error) {

	leftList, err := leftValue.toSeriesList(context.Timerange)
	if err != nil {
		return nil, err
	}
	rightList, err := rightValue.toSeriesList(context.Timerange)
	if err != nil {
		return nil, err
	}

	joined := join([]api.SeriesList{leftList, rightList})

	result := make([]api.Timeseries, len(joined.Rows))

	for i, row := range joined.Rows {
		left := row.Row[0]
		right := row.Row[1]
		array := make([]float64, len(left.Values))
		for j := 0; j < len(left.Values); j++ {
			array[j] = evaluate(left.Values[j], right.Values[j])
		}
		result[i] = api.Timeseries{array, row.TagSet}
	}

	return seriesListValue(api.SeriesList{
		Series:    result,
		Timerange: context.Timerange,
	}), nil
}

// evaluateExpressions evaluates all provided Expressions in the
// EvaluationContext. If any evaluations error, evaluateExpressions will
// propagate that error. The resulting SeriesLists will be in an order
// corresponding to the provided Expresesions.
func evaluateExpressions(context EvaluationContext, expressions []Expression) ([]value, error) {
	if len(expressions) == 0 {
		return []value{}, nil
	}

	results := make([]value, len(expressions))
	for i, expr := range expressions {
		result, err := expr.Evaluate(context)
		if err != nil {
			return nil, err
		}

		results[i] = result
	}

	return results, nil
}
<commit_msg>bugfix<commit_after>// Copyright 2015 Square 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 query

import (
	"errors"
	"fmt"

	"github.com/square/metrics/api"
	"github.com/square/metrics/query/aggregate"
)

// EvaluationContext is the central piece of logic, providing
// helper funcions & varaibles to evaluate a given piece of
// metrics query.
// * Contains Backend object, which can be used to fetch data
// from the backend system.s
// * Contains current timerange being queried for - this can be
// changed by say, application of time shift function.
type EvaluationContext struct {
	// Backend to fetch data from
	Backend api.Backend
	// Timerange to fetch data from
	Timerange api.Timerange
	// SampleMethod to use when up/downsampling to match the requested resolution
	SampleMethod api.SampleMethod
	Predicate    api.Predicate
}

// A value is the result of evaluating an expression.
// They can be floating point values, strings, or series lists.
type value interface {
	toSeriesList(api.Timerange) (api.SeriesList, error)
	toString() (string, error)
	toScalar() (float64, error)
}

type conversionError struct {
	from string
	to   string
}

func (e conversionError) Error() string {
	return fmt.Sprintf("cannot convert from type %s to type %s", e.from, e.to)
}

// A seriesListValue is a value which holds a SeriesList
type seriesListValue api.SeriesList

func (value seriesListValue) toSeriesList(time api.Timerange) (api.SeriesList, error) {
	return api.SeriesList(value), nil
}
func (value seriesListValue) toString() (string, error) {
	return "", conversionError{"SeriesList", "string"}
}
func (value seriesListValue) toScalar() (float64, error) {
	return 0, conversionError{"SeriesList", "scalar"}
}

// A stringValue holds a string
type stringValue string

func (value stringValue) toSeriesList(time api.Timerange) (api.SeriesList, error) {
	return api.SeriesList{}, conversionError{"string", "SeriesList"}
}
func (value stringValue) toString() (string, error) {
	return string(value), nil
}
func (value stringValue) toScalar() (float64, error) {
	return 0, conversionError{"string", "scalar"}
}

// A scalarValue holds a float and can be converted to a serieslist
type scalarValue float64

func (value scalarValue) toSeriesList(timerange api.Timerange) (api.SeriesList, error) {
	if !timerange.IsValid() {
		return api.SeriesList{}, errors.New("Invalid context.Timerange")
	}

	series := make([]float64, timerange.Slots())
	for i := range series {
		series[i] = float64(value)
	}

	return api.SeriesList{
		Series:    []api.Timeseries{api.Timeseries{series, api.NewTagSet()}},
		Timerange: timerange,
	}, nil
}

func (value scalarValue) toString() (string, error) {
	return "", conversionError{"scalar", "string"}
}

func (value scalarValue) toScalar() (float64, error) {
	return float64(value), nil
}

// Expression is a piece of code, which can be evaluated in a given
// EvaluationContext. EvaluationContext must never be changed in an Evalute().
//
// The contract of Expressions is that leaf nodes must sample a resulting
// timeseries according to the resolution specified in its EvaluationContext's
// Timerange. Internal nodes may assume that results from evaluating child
// Expressions correspond to the timerange in the current EvaluationContext.
type Expression interface {
	// Evaluate the given expression.
	Evaluate(context EvaluationContext) (value, error)
}

func evaluateToSeriesList(e Expression, context EvaluationContext) (api.SeriesList, error) {
	value, err := e.Evaluate(context)
	if err != nil {
		return api.SeriesList{}, err
	}
	return value.toSeriesList(context.Timerange)
}

// Implementations
// ===============

// Generates a Timeseries from the encapsulated scalar.
func (expr scalarExpression) Evaluate(context EvaluationContext) (value, error) {
	return scalarValue(expr.value), nil
}

func (expr *metricFetchExpression) Evaluate(context EvaluationContext) (value, error) {
	// Merge predicates appropriately
	var predicate api.Predicate
	if context.Predicate == nil {
		predicate = expr.predicate
	} else if expr.predicate == nil {
		predicate = context.Predicate
	} else {
		predicate = &andPredicate{[]api.Predicate{expr.predicate, context.Predicate}}
	}

	series, err := context.Backend.FetchSeries(api.TaggedMetric{api.MetricKey(expr.metricName), nil}, predicate, context.SampleMethod, context.Timerange)
	if err != nil {
		return seriesListValue{}, err
	}
	return seriesListValue(*series), err
}

func (expr *functionExpression) Evaluate(context EvaluationContext) (value, error) {

	name := expr.functionName

	operatorMap := map[string]func(float64, float64) float64{
		"+": func(x, y float64) float64 { return x + y },
		"-": func(x, y float64) float64 { return x - y },
		"*": func(x, y float64) float64 { return x * y },
		"/": func(x, y float64) float64 { return x / y },
	}

	if operator, ok := operatorMap[name]; ok {
		// Evaluation of a binary operator:
		// Verify that exactly 2 arguments are given.
		if len(expr.arguments) != 2 {
			return nil, errors.New(fmt.Sprintf("Function `%s` expects 2 operands but received %d (%+v)", name, len(expr.arguments), expr.arguments))
		}
		left, err := expr.arguments[0].Evaluate(context)
		if err != nil {
			return nil, err
		}
		right, err := expr.arguments[1].Evaluate(context)
		if err != nil {
			return nil, err
		}
		return evaluateBinaryOperation(context, name, left, right, operator)
	}

	if aggregator, ok := aggregate.GetAggregate(name); ok {
		// Verify that exactly 1 argument is given.
		if len(expr.arguments) != 1 {
			return nil, errors.New(fmt.Sprintf("Function `%s` expects 1 argument but received %d (%+v)", name, len(expr.arguments), expr.arguments))
		}
		argument := expr.arguments[0]
		value, err := argument.Evaluate(context)
		if err != nil {
			return nil, err
		}
		list, err := value.toSeriesList(context.Timerange)
		if err != nil {
			return nil, err
		}
		series := aggregate.AggregateBy(list, aggregator, expr.groupBy)
		return seriesListValue(series), nil
	}

	if transform, ok := GetTransformation(name); ok {
		//Verify that at least one argument is given.
		if len(expr.arguments) == 0 {
			return nil, errors.New(fmt.Sprintf("Function `%s` expects at least 1 argument but was given 0", name))
		}
		first, err := expr.arguments[0].Evaluate(context)
		if err != nil {
			return nil, err
		}
		list, err := first.toSeriesList(context.Timerange)
		if err != nil {
			return nil, err
		}
		// Evaluate all the other parameters:
		rest := expr.arguments[1:]
		parameters := make([]value, len(rest))
		for i := range parameters {
			parameters[i], err = rest[i].Evaluate(context)
			if err != nil {
				return nil, err
			}
		}
		series, err := ApplyTransform(list, transform, parameters)
		if err != nil {
			return nil, err
		}
		return seriesListValue(series), nil
	}

	return nil, errors.New(fmt.Sprintf("unknown function name `%s`", name))
}

//
// Auxiliary functions
//

// evaluateExpression wraps expr.Evaluate() to provide common messaging
// for errors. This can get pretty messy if the Expression we evaluate
// isn't a leaf node, but a leaf fails.
func evaluateExpression(context EvaluationContext, expr Expression) (value, error) {
	result, err := expr.Evaluate(context)
	if err != nil {
		return nil, errors.New(fmt.Sprintf("Evaluation of expression %+v failed:\n%s\n", expr, err.Error()))
	}
	return result, err
}

// evaluateBinaryOperation applies an arbirary binary operation to two
// Expressions.
func evaluateBinaryOperation(
	context EvaluationContext,
	functionName string,
	leftValue value,
	rightValue value,
	evaluate func(float64, float64) float64,
) (value, error) {

	leftList, err := leftValue.toSeriesList(context.Timerange)
	if err != nil {
		return nil, err
	}
	rightList, err := rightValue.toSeriesList(context.Timerange)
	if err != nil {
		return nil, err
	}

	joined := join([]api.SeriesList{leftList, rightList})

	result := make([]api.Timeseries, len(joined.Rows))

	for i, row := range joined.Rows {
		left := row.Row[0]
		right := row.Row[1]
		array := make([]float64, len(left.Values))
		for j := 0; j < len(left.Values); j++ {
			array[j] = evaluate(left.Values[j], right.Values[j])
		}
		result[i] = api.Timeseries{array, row.TagSet}
	}

	return seriesListValue(api.SeriesList{
		Series:    result,
		Timerange: context.Timerange,
	}), nil
}

// evaluateExpressions evaluates all provided Expressions in the
// EvaluationContext. If any evaluations error, evaluateExpressions will
// propagate that error. The resulting SeriesLists will be in an order
// corresponding to the provided Expresesions.
func evaluateExpressions(context EvaluationContext, expressions []Expression) ([]value, error) {
	if len(expressions) == 0 {
		return []value{}, nil
	}

	results := make([]value, len(expressions))
	for i, expr := range expressions {
		result, err := expr.Evaluate(context)
		if err != nil {
			return nil, err
		}

		results[i] = result
	}

	return results, nil
}
<|endoftext|>