<commit_before>/*
RpcClient for Go RPC Servers
Copyright (C) ITsysCOM GmbH

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program.  If not, see <http://www.gnu.org/licenses/>
*/

package rpcclient

import (
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"log/syslog"
	"math"
	"math/rand"
	"net"
	"net/http"
	"net/rpc"
	"net/rpc/jsonrpc"
	"reflect"
	"strings"
	"sync"
	"syscall"
	"time"
)

const (
	JSON_RPC       = "json"
	JSON_HTTP      = "http_jsonrpc"
	GOB_RPC        = "gob"
	INTERNAL_RPC   = "*internal"
	POOL_FIRST     = "first"
	POOL_RANDOM    = "random"
	POOL_NEXT      = "next"
	POOL_BROADCAST = "broadcast"
)

var (
	ErrReqUnsynchronized       = errors.New("REQ_UNSYNCHRONIZED")
	ErrUnsupporteServiceMethod = errors.New("UNSUPPORTED_SERVICE_METHOD")
	ErrWrongArgsType           = errors.New("WRONG_ARGS_TYPE")
	ErrWrongReplyType          = errors.New("WRONG_REPLY_TYPE")
	ErrDisconnected            = errors.New("DISCONNECTED")
	ErrReplyTimeout            = errors.New("REPLY_TIMEOUT")
	ErrFailedReconnect         = errors.New("FAILED_RECONNECT")
	ErrInternallyDisconnected  = errors.New("INTERNALLY_DISCONNECTED")
	logger                     *syslog.Writer
)

func init() {
	logger, _ = syslog.New(syslog.LOG_INFO, "RPCClient") // If we need to report anything to syslog
}

// successive Fibonacci numbers.
func Fib() func() time.Duration {
	a, b := 0, 1
	return func() time.Duration {
		a, b = b, a+b
		return time.Duration(a*10) * time.Millisecond
	}
}

func NewRpcClient(transport, addr string, connectAttempts, reconnects int, connTimeout, replyTimeout time.Duration, codec string, internalConn RpcClientConnection) (*RpcClient, error) {
	if codec == INTERNAL_RPC && reflect.ValueOf(internalConn).IsNil() {
		return nil, ErrInternallyDisconnected
	}
	var err error
	rpcClient := &RpcClient{transport: transport, address: addr, reconnects: reconnects,
		connTimeout: connTimeout, replyTimeout: replyTimeout, codec: codec, connection: internalConn}
	delay := Fib()
	for i := 0; i < connectAttempts; i++ {
		err = rpcClient.connect()
		if err == nil { //Connected so no need to reiterate
			break
		}
		time.Sleep(delay())
	}
	return rpcClient, err
}

type RpcClient struct {
	transport    string
	address      string
	reconnects   int
	connTimeout  time.Duration
	replyTimeout time.Duration
	codec        string // JSON_RPC or GOB_RPC
	connection   RpcClientConnection
	connMux      sync.RWMutex // protects connection
}

func (self *RpcClient) connect() (err error) {
	self.connMux.Lock()
	defer self.connMux.Unlock()
	if self.codec == INTERNAL_RPC {
		return
	} else if self.codec == JSON_HTTP {
		self.connection = &HttpJsonRpcClient{httpClient: new(http.Client), url: self.address}
		return
	}
	// RPC compliant connections here, manually create connection to timeout
	netconn, err := net.DialTimeout(self.transport, self.address, self.connTimeout)
	if err != nil {
		self.connection = nil // So we don't wrap nil into the interface
		return err
	}
	if self.codec == JSON_RPC {
		self.connection = jsonrpc.NewClient(netconn)
	} else {
		self.connection = rpc.NewClient(netconn)
	}
	return
}

func (self *RpcClient) isConnected() bool {
	self.connMux.RLock()
	defer self.connMux.RUnlock()
	return self.connection != nil
}

func (self *RpcClient) disconnect() (err error) {
	switch self.codec {
	case INTERNAL_RPC, JSON_HTTP:
	default:
		self.connMux.Lock()
		if self.connection != nil {
			self.connection.(*rpc.Client).Close()
			self.connection = nil
		}
		self.connMux.Unlock()
	}
	return nil
}

func (self *RpcClient) reconnect() (err error) {
	if self.reconnects == 0 {
		return nil // no reconnects
	}
	self.disconnect()            // make sure we have cleared the connection so it can be garbage collected
	if self.codec == JSON_HTTP { // http client has automatic reconnects in place
		return self.connect()
	}
	i := 0
	delay := Fib()
	for {
		i++
		if self.reconnects != -1 && i > self.reconnects { // Maximum reconnects reached, -1 for infinite reconnects
			break
		}
		if err = self.connect(); err == nil { // No error on connect, succcess
			return nil
		}
		time.Sleep(delay()) // Cound not reconnect, retry
	}
	return ErrFailedReconnect
}

func (self *RpcClient) Call(serviceMethod string, args interface{}, reply interface{}) (err error) {
	if args == nil {
		return fmt.Errorf("nil rpc in argument method: %s in: %v out: %v", serviceMethod, args, reply)
	}
	errChan := make(chan error, 1)
	go func(serviceMethod string, args interface{}, reply interface{}) {
		self.connMux.RLock()
		if self.connection == nil {
			errChan <- ErrDisconnected
		} else {
			errChan <- self.connection.Call(serviceMethod, args, reply)
		}
		self.connMux.RUnlock()
	}(serviceMethod, args, reply)
	select {
	case err = <-errChan:
	case <-time.After(self.replyTimeout):
		err = ErrReplyTimeout
	}

	if isNetworkError(err) && err != ErrReplyTimeout && self.reconnects != 0 { // ReplyTimeout should not reconnect since it creates loop
		if errReconnect := self.reconnect(); errReconnect != nil {
			return err
		} else { // Run command after reconnect
			self.connMux.RLock()
			defer self.connMux.RUnlock()
			return self.connection.Call(serviceMethod, args, reply)
		}
	}
	return err
}

// Connection used in RpcClient, as interface so we can combine the rpc.RpcClient with http one or websocket
type RpcClientConnection interface {
	Call(string, interface{}, interface{}) error
}

// Response received for
type JsonRpcResponse struct {
	Id     uint64
	Result *json.RawMessage
	Error  interface{}
}

type HttpJsonRpcClient struct {
	httpClient *http.Client
	id         uint64
	url        string
}

func (self *HttpJsonRpcClient) Call(serviceMethod string, args interface{}, reply interface{}) error {
	self.id += 1
	id := self.id
	data, err := json.Marshal(map[string]interface{}{
		"method": serviceMethod,
		"id":     self.id,
		"params": [1]interface{}{args},
	})
	if err != nil {
		return err
	}
	resp, err := self.httpClient.Post(self.url, "application/json", ioutil.NopCloser(strings.NewReader(string(data)))) // Closer so we automatically have close after response
	if err != nil {
		return err
	}
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return err
	}
	var jsonRsp JsonRpcResponse
	err = json.Unmarshal(body, &jsonRsp)
	if err != nil {
		return err
	}
	if jsonRsp.Id != id {
		return ErrReqUnsynchronized
	}
	if jsonRsp.Error != nil || jsonRsp.Result == nil {
		x, ok := jsonRsp.Error.(string)
		if !ok {
			return fmt.Errorf("invalid error %v", jsonRsp.Error)
		}
		if x == "" {
			x = "unspecified error"
		}
		return errors.New(x)
	}
	return json.Unmarshal(*jsonRsp.Result, reply)
}

type RpcClientPool struct {
	transmissionType string
	connections      []RpcClientConnection
	counter          int
	replyTimeout     time.Duration
}

func NewRpcClientPool(transmissionType string, replyTimeout time.Duration) *RpcClientPool {
	return &RpcClientPool{transmissionType: transmissionType, replyTimeout: replyTimeout}
}

func (pool *RpcClientPool) AddClient(rcc RpcClientConnection) {
	if rcc != nil && !reflect.ValueOf(rcc).IsNil() {
		pool.connections = append(pool.connections, rcc)
	}
}

func (pool *RpcClientPool) Call(serviceMethod string, args interface{}, reply interface{}) (err error) {
	switch pool.transmissionType {
	case POOL_BROADCAST:
		replyChan := make(chan *rpcReplyError, len(pool.connections))
		for _, rc := range pool.connections {
			go func(conn RpcClientConnection) {
				// make a new pointer of the same type
				rpl := reflect.New(reflect.TypeOf(reflect.ValueOf(reply).Elem().Interface()))
				err := conn.Call(serviceMethod, args, rpl.Interface())
				if !isNetworkError(err) {
					replyChan <- &rpcReplyError{reply: rpl.Interface(), err: err}
				}
			}(rc)
		}
		//get first response with timeout
		var re *rpcReplyError
		select {
		case re = <-replyChan:
		case <-time.After(pool.replyTimeout):
			return ErrReplyTimeout
		}
		// put received value in the orig reply
		reflect.ValueOf(reply).Elem().Set(reflect.ValueOf(re.reply).Elem())
		return re.err
	case POOL_FIRST:
		for _, rc := range pool.connections {
			err = rc.Call(serviceMethod, args, reply)
			if isNetworkError(err) {
				continue
			}
			return
		}
	case POOL_NEXT:
		ln := len(pool.connections)
		rrIndexes := roundIndex(int(math.Mod(float64(pool.counter), float64(ln))), ln)
		pool.counter++
		for _, index := range rrIndexes {
			err = pool.connections[index].Call(serviceMethod, args, reply)
			if isNetworkError(err) {
				continue
			}
			return
		}
	case POOL_RANDOM:
		rand.Seed(time.Now().UnixNano())
		randomIndex := rand.Perm(len(pool.connections))
		for _, index := range randomIndex {
			err = pool.connections[index].Call(serviceMethod, args, reply)
			if isNetworkError(err) {
				continue
			}
			return
		}
	}
	return
}

type rpcReplyError struct {
	reply interface{}
	err   error
}

// generates round robin indexes for a slice of length max
// starting from index start
func roundIndex(start, max int) []int {
	if start < 0 {
		start = 0
	}
	result := make([]int, max)
	for i := 0; i < max; i++ {
		if start+i < max {
			result[i] = start + i
		} else {
			result[i] = int(math.Abs(float64(max - (start + i))))
		}
	}
	return result
}

func isNetworkError(err error) bool {
	if err == nil {
		return false
	}
	if operr, ok := err.(*net.OpError); ok && strings.HasSuffix(operr.Err.Error(), syscall.ECONNRESET.Error()) { // connection reset
		return true
	}
	return err == rpc.ErrShutdown ||
		err == ErrReqUnsynchronized ||
		err == ErrDisconnected ||
		err == ErrReplyTimeout ||
		strings.HasPrefix(err.Error(), "rpc: can't find service")
}
<commit_msg>Error for connect attempts on nil connection for internal codec<commit_after>/*
RpcClient for Go RPC Servers
Copyright (C) ITsysCOM GmbH

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program.  If not, see <http://www.gnu.org/licenses/>
*/

package rpcclient

import (
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"log/syslog"
	"math"
	"math/rand"
	"net"
	"net/http"
	"net/rpc"
	"net/rpc/jsonrpc"
	"reflect"
	"strings"
	"sync"
	"syscall"
	"time"
)

const (
	JSON_RPC       = "json"
	JSON_HTTP      = "http_jsonrpc"
	GOB_RPC        = "gob"
	INTERNAL_RPC   = "*internal"
	POOL_FIRST     = "first"
	POOL_RANDOM    = "random"
	POOL_NEXT      = "next"
	POOL_BROADCAST = "broadcast"
)

var (
	ErrReqUnsynchronized       = errors.New("REQ_UNSYNCHRONIZED")
	ErrUnsupporteServiceMethod = errors.New("UNSUPPORTED_SERVICE_METHOD")
	ErrWrongArgsType           = errors.New("WRONG_ARGS_TYPE")
	ErrWrongReplyType          = errors.New("WRONG_REPLY_TYPE")
	ErrDisconnected            = errors.New("DISCONNECTED")
	ErrReplyTimeout            = errors.New("REPLY_TIMEOUT")
	ErrFailedReconnect         = errors.New("FAILED_RECONNECT")
	ErrInternallyDisconnected  = errors.New("INTERNALLY_DISCONNECTED")
	ErrUnsupportedCodec        = errors.New("UNSUPPORTED_CODEC")
	logger                     *syslog.Writer
)

func init() {
	logger, _ = syslog.New(syslog.LOG_INFO, "RPCClient") // If we need to report anything to syslog
}

// successive Fibonacci numbers.
func Fib() func() time.Duration {
	a, b := 0, 1
	return func() time.Duration {
		a, b = b, a+b
		return time.Duration(a*10) * time.Millisecond
	}
}

func NewRpcClient(transport, addr string, connectAttempts, reconnects int, connTimeout, replyTimeout time.Duration, codec string, internalConn RpcClientConnection) (*RpcClient, error) {
	if codec != INTERNAL_RPC && codec != JSON_RPC && codec != JSON_HTTP && codec != GOB_RPC {
		return nil, ErrUnsupportedCodec
	}
	if codec == INTERNAL_RPC && reflect.ValueOf(internalConn).IsNil() {
		return nil, ErrInternallyDisconnected
	}
	var err error
	rpcClient := &RpcClient{transport: transport, address: addr, reconnects: reconnects,
		connTimeout: connTimeout, replyTimeout: replyTimeout, codec: codec, connection: internalConn}
	delay := Fib()
	for i := 0; i < connectAttempts; i++ {
		err = rpcClient.connect()
		if err == nil { //Connected so no need to reiterate
			break
		}
		time.Sleep(delay())
	}
	return rpcClient, err
}

type RpcClient struct {
	transport    string
	address      string
	reconnects   int
	connTimeout  time.Duration
	replyTimeout time.Duration
	codec        string // JSON_RPC or GOB_RPC
	connection   RpcClientConnection
	connMux      sync.RWMutex // protects connection
}

func (self *RpcClient) connect() (err error) {
	self.connMux.Lock()
	defer self.connMux.Unlock()
	if self.codec == INTERNAL_RPC {
		if self.connection == nil {
			return ErrDisconnected
		}
		return
	} else if self.codec == JSON_HTTP {
		self.connection = &HttpJsonRpcClient{httpClient: new(http.Client), url: self.address}
		return
	}
	// RPC compliant connections here, manually create connection to timeout
	netconn, err := net.DialTimeout(self.transport, self.address, self.connTimeout)
	if err != nil {
		self.connection = nil // So we don't wrap nil into the interface
		return err
	}
	if self.codec == JSON_RPC {
		self.connection = jsonrpc.NewClient(netconn)
	} else {
		self.connection = rpc.NewClient(netconn)
	}
	return
}

func (self *RpcClient) isConnected() bool {
	self.connMux.RLock()
	defer self.connMux.RUnlock()
	return self.connection != nil
}

func (self *RpcClient) disconnect() (err error) {
	switch self.codec {
	case INTERNAL_RPC, JSON_HTTP:
	default:
		self.connMux.Lock()
		if self.connection != nil {
			self.connection.(*rpc.Client).Close()
			self.connection = nil
		}
		self.connMux.Unlock()
	}
	return nil
}

func (self *RpcClient) reconnect() (err error) {
	if self.reconnects == 0 {
		return nil // no reconnects
	}
	self.disconnect()            // make sure we have cleared the connection so it can be garbage collected
	if self.codec == JSON_HTTP { // http client has automatic reconnects in place
		return self.connect()
	}
	i := 0
	delay := Fib()
	for {
		i++
		if self.reconnects != -1 && i > self.reconnects { // Maximum reconnects reached, -1 for infinite reconnects
			break
		}
		if err = self.connect(); err == nil { // No error on connect, succcess
			return nil
		}
		time.Sleep(delay()) // Cound not reconnect, retry
	}
	return ErrFailedReconnect
}

func (self *RpcClient) Call(serviceMethod string, args interface{}, reply interface{}) (err error) {
	if args == nil {
		return fmt.Errorf("nil rpc in argument method: %s in: %v out: %v", serviceMethod, args, reply)
	}
	errChan := make(chan error, 1)
	go func(serviceMethod string, args interface{}, reply interface{}) {
		self.connMux.RLock()
		if self.connection == nil {
			errChan <- ErrDisconnected
		} else {
			errChan <- self.connection.Call(serviceMethod, args, reply)
		}
		self.connMux.RUnlock()
	}(serviceMethod, args, reply)
	select {
	case err = <-errChan:
	case <-time.After(self.replyTimeout):
		err = ErrReplyTimeout
	}

	if isNetworkError(err) && err != ErrReplyTimeout && self.reconnects != 0 { // ReplyTimeout should not reconnect since it creates loop
		if errReconnect := self.reconnect(); errReconnect != nil {
			return err
		} else { // Run command after reconnect
			self.connMux.RLock()
			defer self.connMux.RUnlock()
			return self.connection.Call(serviceMethod, args, reply)
		}
	}
	return err
}

// Connection used in RpcClient, as interface so we can combine the rpc.RpcClient with http one or websocket
type RpcClientConnection interface {
	Call(string, interface{}, interface{}) error
}

// Response received for
type JsonRpcResponse struct {
	Id     uint64
	Result *json.RawMessage
	Error  interface{}
}

type HttpJsonRpcClient struct {
	httpClient *http.Client
	id         uint64
	url        string
}

func (self *HttpJsonRpcClient) Call(serviceMethod string, args interface{}, reply interface{}) error {
	self.id += 1
	id := self.id
	data, err := json.Marshal(map[string]interface{}{
		"method": serviceMethod,
		"id":     self.id,
		"params": [1]interface{}{args},
	})
	if err != nil {
		return err
	}
	resp, err := self.httpClient.Post(self.url, "application/json", ioutil.NopCloser(strings.NewReader(string(data)))) // Closer so we automatically have close after response
	if err != nil {
		return err
	}
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return err
	}
	var jsonRsp JsonRpcResponse
	err = json.Unmarshal(body, &jsonRsp)
	if err != nil {
		return err
	}
	if jsonRsp.Id != id {
		return ErrReqUnsynchronized
	}
	if jsonRsp.Error != nil || jsonRsp.Result == nil {
		x, ok := jsonRsp.Error.(string)
		if !ok {
			return fmt.Errorf("invalid error %v", jsonRsp.Error)
		}
		if x == "" {
			x = "unspecified error"
		}
		return errors.New(x)
	}
	return json.Unmarshal(*jsonRsp.Result, reply)
}

type RpcClientPool struct {
	transmissionType string
	connections      []RpcClientConnection
	counter          int
	replyTimeout     time.Duration
}

func NewRpcClientPool(transmissionType string, replyTimeout time.Duration) *RpcClientPool {
	return &RpcClientPool{transmissionType: transmissionType, replyTimeout: replyTimeout}
}

func (pool *RpcClientPool) AddClient(rcc RpcClientConnection) {
	if rcc != nil && !reflect.ValueOf(rcc).IsNil() {
		pool.connections = append(pool.connections, rcc)
	}
}

func (pool *RpcClientPool) Call(serviceMethod string, args interface{}, reply interface{}) (err error) {
	switch pool.transmissionType {
	case POOL_BROADCAST:
		replyChan := make(chan *rpcReplyError, len(pool.connections))
		for _, rc := range pool.connections {
			go func(conn RpcClientConnection) {
				// make a new pointer of the same type
				rpl := reflect.New(reflect.TypeOf(reflect.ValueOf(reply).Elem().Interface()))
				err := conn.Call(serviceMethod, args, rpl.Interface())
				if !isNetworkError(err) {
					replyChan <- &rpcReplyError{reply: rpl.Interface(), err: err}
				}
			}(rc)
		}
		//get first response with timeout
		var re *rpcReplyError
		select {
		case re = <-replyChan:
		case <-time.After(pool.replyTimeout):
			return ErrReplyTimeout
		}
		// put received value in the orig reply
		reflect.ValueOf(reply).Elem().Set(reflect.ValueOf(re.reply).Elem())
		return re.err
	case POOL_FIRST:
		for _, rc := range pool.connections {
			err = rc.Call(serviceMethod, args, reply)
			if isNetworkError(err) {
				continue
			}
			return
		}
	case POOL_NEXT:
		ln := len(pool.connections)
		rrIndexes := roundIndex(int(math.Mod(float64(pool.counter), float64(ln))), ln)
		pool.counter++
		for _, index := range rrIndexes {
			err = pool.connections[index].Call(serviceMethod, args, reply)
			if isNetworkError(err) {
				continue
			}
			return
		}
	case POOL_RANDOM:
		rand.Seed(time.Now().UnixNano())
		randomIndex := rand.Perm(len(pool.connections))
		for _, index := range randomIndex {
			err = pool.connections[index].Call(serviceMethod, args, reply)
			if isNetworkError(err) {
				continue
			}
			return
		}
	}
	return
}

type rpcReplyError struct {
	reply interface{}
	err   error
}

// generates round robin indexes for a slice of length max
// starting from index start
func roundIndex(start, max int) []int {
	if start < 0 {
		start = 0
	}
	result := make([]int, max)
	for i := 0; i < max; i++ {
		if start+i < max {
			result[i] = start + i
		} else {
			result[i] = int(math.Abs(float64(max - (start + i))))
		}
	}
	return result
}

func isNetworkError(err error) bool {
	if err == nil {
		return false
	}
	if operr, ok := err.(*net.OpError); ok && strings.HasSuffix(operr.Err.Error(), syscall.ECONNRESET.Error()) { // connection reset
		return true
	}
	return err == rpc.ErrShutdown ||
		err == ErrReqUnsynchronized ||
		err == ErrDisconnected ||
		err == ErrReplyTimeout ||
		strings.HasPrefix(err.Error(), "rpc: can't find service")
}
<|endoftext|>


<commit_before>package agent

import (
	"bytes"
	"encoding/json"
	"fmt"
	"os"
	"strings"
	"testing"

	"github.com/buildkite/agent/v3/env"
	"github.com/stretchr/testify/assert"
)

func TestPipelineParserParsesYaml(t *testing.T) {
	environ := env.FromSlice([]string{`ENV_VAR_FRIEND="friend"`})

	result, err := PipelineParser{
		Filename: "awesome.yml",
		Pipeline: []byte("steps:\n  - label: \"hello ${ENV_VAR_FRIEND}\""),
		Env:      environ}.Parse()

	assert.NoError(t, err)
	j, err := json.Marshal(result)
	assert.Equal(t, `{"steps":[{"label":"hello \"friend\""}]}`, string(j))
}

func TestPipelineParserParsesYamlWithNoInterpolation(t *testing.T) {
	result, err := PipelineParser{
		Filename:        "awesome.yml",
		Pipeline:        []byte("steps:\n  - label: \"hello ${ENV_VAR_FRIEND}\""),
		NoInterpolation: true,
	}.Parse()

	assert.NoError(t, err)
	j, err := json.Marshal(result)
	assert.Equal(t, `{"steps":[{"label":"hello ${ENV_VAR_FRIEND}"}]}`, string(j))
}

func TestPipelineParserSupportsYamlMergesAndAnchors(t *testing.T) {
	complexYAML := `---
base_step: &base_step
  type: script
  agent_query_rules:
    - queue=default

steps:
  - <<: *base_step
    name: ':docker: building image'
    command: docker build .
    agents:
      queue: default`

	result, err := PipelineParser{
		Filename: "awesome.yml",
		Pipeline: []byte(complexYAML)}.Parse()

	assert.NoError(t, err)
	j, err := json.Marshal(result)
	assert.Equal(t, `{"base_step":{"type":"script","agent_query_rules":["queue=default"]},"steps":[{"type":"script","agent_query_rules":["queue=default"],"name":":docker: building image","command":"docker build .","agents":{"queue":"default"}}]}`, string(j))
}

func TestPipelineParserReturnsYamlParsingErrors(t *testing.T) {
	_, err := PipelineParser{Filename: "awesome.yml", Pipeline: []byte("steps: %blah%")}.Parse()
	assert.Error(t, err, `Failed to parse awesome.yml: found character that cannot start any token`, fmt.Sprintf("%s", err))
}

func TestPipelineParserReturnsJsonParsingErrors(t *testing.T) {
	_, err := PipelineParser{Filename: "awesome.json", Pipeline: []byte("{")}.Parse()
	assert.Error(t, err, `Failed to parse awesome.json: line 1: did not find expected node content`, fmt.Sprintf("%s", err))
}

func TestPipelineParserParsesJson(t *testing.T) {
	environ := env.FromSlice([]string{`ENV_VAR_FRIEND="friend"`})

	result, err := PipelineParser{
		Filename: "thing.json",
		Pipeline: []byte("\n\n     \n  { \"foo\": \"bye ${ENV_VAR_FRIEND}\" }\n"),
		Env:      environ}.Parse()

	assert.NoError(t, err)
	j, err := json.Marshal(result)
	assert.Equal(t, `{"foo":"bye \"friend\""}`, string(j))
}

func TestPipelineParserParsesJsonObjects(t *testing.T) {
	environ := env.FromSlice([]string{`ENV_VAR_FRIEND="friend"`})

	result, err := PipelineParser{Pipeline: []byte("\n\n     \n  { \"foo\": \"bye ${ENV_VAR_FRIEND}\" }\n"), Env: environ}.Parse()
	assert.NoError(t, err)
	j, err := json.Marshal(result)
	assert.Equal(t, `{"foo":"bye \"friend\""}`, string(j))
}

func TestPipelineParserParsesJsonArrays(t *testing.T) {
	environ := env.FromSlice([]string{`ENV_VAR_FRIEND="friend"`})

	result, err := PipelineParser{Pipeline: []byte("\n\n     \n  [ { \"foo\": \"bye ${ENV_VAR_FRIEND}\" } ]\n"), Env: environ}.Parse()
	assert.NoError(t, err)
	j, err := json.Marshal(result)
	assert.Equal(t, `{"steps":[{"foo":"bye \"friend\""}]}`, string(j))
}

func TestPipelineParserParsesTopLevelSteps(t *testing.T) {
	result, err := PipelineParser{Pipeline: []byte("---\n- name: Build\n  command: echo hello world\n- wait\n"), Env: nil}.Parse()
	assert.NoError(t, err)
	j, err := json.Marshal(result)
	assert.Equal(t, `{"steps":[{"name":"Build","command":"echo hello world"},"wait"]}`, string(j))
}

func TestPipelineParserPreservesBools(t *testing.T) {
	result, err := PipelineParser{Pipeline: []byte("steps:\n  - trigger: hello\n    async: true")}.Parse()
	assert.Nil(t, err)
	j, err := json.Marshal(result)
	assert.Equal(t, `{"steps":[{"trigger":"hello","async":true}]}`, string(j))
}

func TestPipelineParserPreservesInts(t *testing.T) {
	result, err := PipelineParser{Pipeline: []byte("steps:\n  - label: hello\n    parallelism: 10")}.Parse()
	assert.Nil(t, err)
	j, err := json.Marshal(result)
	assert.Equal(t, `{"steps":[{"label":"hello","parallelism":10}]}`, string(j))
}

func TestPipelineParserPreservesNull(t *testing.T) {
	result, err := PipelineParser{Pipeline: []byte("steps:\n  - wait: ~")}.Parse()
	assert.Nil(t, err)
	j, err := json.Marshal(result)
	assert.Equal(t, `{"steps":[{"wait":null}]}`, string(j))
}

func TestPipelineParserPreservesFloats(t *testing.T) {
	result, err := PipelineParser{Pipeline: []byte("steps:\n  - trigger: hello\n    llamas: 3.142")}.Parse()
	assert.Nil(t, err)
	j, err := json.Marshal(result)
	assert.Equal(t, `{"steps":[{"trigger":"hello","llamas":3.142}]}`, string(j))
}

func TestPipelineParserHandlesDates(t *testing.T) {
	result, err := PipelineParser{Pipeline: []byte("steps:\n  - trigger: hello\n    llamas: 2002-08-15T17:18:23.18-06:00")}.Parse()
	assert.Nil(t, err)
	j, err := json.Marshal(result)
	assert.Equal(t, `{"steps":[{"trigger":"hello","llamas":"2002-08-15T17:18:23.18-06:00"}]}`, string(j))
}

func TestPipelineParserInterpolatesKeysAsWellAsValues(t *testing.T) {
	var pipeline = `{
		"env": {
			"${FROM_ENV}TEST1": "MyTest",
			"TEST2": "${FROM_ENV}"
		}
	}`

	var decoded struct {
		Env map[string]string `json:"env"`
	}

	environ := env.FromSlice([]string{`FROM_ENV=llamas`})

	result, err := PipelineParser{Pipeline: []byte(pipeline), Env: environ}.Parse()
	if err != nil {
		t.Fatal(err)
	}

	err = decodeIntoStruct(&decoded, result)
	if err != nil {
		t.Fatal(err)
	}

	assert.Equal(t, `MyTest`, decoded.Env["llamasTEST1"])
	assert.Equal(t, `llamas`, decoded.Env["TEST2"])
}

func TestPipelineParserLoadsGlobalEnvBlockFirst(t *testing.T) {
	var pipeline = `{
		"env": {
			"TEAM1": "England",
			"TEAM2": "Australia",
			"HEADLINE": "${TEAM1} smashes ${TEAM2} to win the ashes in ${YEAR_FROM_SHELL}!!"
		},
		"steps": [{
			"command": "echo ${HEADLINE}"
		}]
	}`

	var decoded struct {
		Env   map[string]string `json:"env"`
		Steps []struct {
			Command string `json:"command"`
		} `json:"steps"`
	}

	environ := env.FromSlice([]string{`YEAR_FROM_SHELL=1912`})

	result, err := PipelineParser{Pipeline: []byte(pipeline), Env: environ}.Parse()
	if err != nil {
		t.Fatal(err)
	}

	err = decodeIntoStruct(&decoded, result)
	if err != nil {
		t.Fatal(err)
	}

	assert.Equal(t, `England`, decoded.Env["TEAM1"])
	assert.Equal(t, `England smashes Australia to win the ashes in 1912!!`, decoded.Env["HEADLINE"])
	assert.Equal(t, `echo England smashes Australia to win the ashes in 1912!!`, decoded.Steps[0].Command)
}

func decodeIntoStruct(into interface{}, from interface{}) error {
	b, err := json.Marshal(from)
	if err != nil {
		return err
	}
	return json.Unmarshal(b, into)
}

func TestPipelineParserLoadsSystemEnvironment(t *testing.T) {
	var pipeline = `{
		"steps": [{
			"command": "echo ${LLAMAS_ROCK?}"
		}]
	}`

	var decoded struct {
		Steps []struct {
			Command string `json:"command"`
		} `json:"steps"`
	}

	_, err := PipelineParser{Pipeline: []byte(pipeline)}.Parse()
	if err == nil {
		t.Fatalf("Expected $LLAMAS_ROCK: not set")
	}

	os.Setenv("LLAMAS_ROCK", "absolutely")
	defer os.Unsetenv("LLAMAS_ROCK")

	result2, err := PipelineParser{Pipeline: []byte(pipeline)}.Parse()
	if err != nil {
		t.Fatal(err)
	}

	err = decodeIntoStruct(&decoded, result2)
	if err != nil {
		t.Fatal(err)
	}

	if decoded.Steps[0].Command != "echo absolutely" {
		t.Fatalf("Unexpected: %q", decoded.Steps[0].Command)
	}
}

func TestPipelineParserPreservesOrderOfPlugins(t *testing.T) {
	var pipeline = `---
steps:
  - name: ":s3: xxx"
    command: "script/buildkite/xxx.sh"
    plugins:
      xxx/aws-assume-role#v0.1.0:
        role: arn:aws:iam::xxx:role/xxx
      ecr#v1.1.4:
        login: true
        account_ids: xxx
        registry_region: us-east-1
      docker-compose#v2.5.1:
        run: xxx
        config: .buildkite/docker/docker-compose.yml
        env:
          - AWS_ACCESS_KEY_ID
          - AWS_SECRET_ACCESS_KEY
          - AWS_SESSION_TOKEN
    agents:
      queue: xxx`

	result, err := PipelineParser{Pipeline: []byte(pipeline), Env: nil}.Parse()
	if err != nil {
		t.Fatal(err)
	}

	buf := &bytes.Buffer{}
	err = json.NewEncoder(buf).Encode(result)
	if err != nil {
		t.Fatal(err)
	}

	expected := `{"steps":[{"name":":s3: xxx","command":"script/buildkite/xxx.sh","plugins":{"xxx/aws-assume-role#v0.1.0":{"role":"arn:aws:iam::xxx:role/xxx"},"ecr#v1.1.4":{"login":true,"account_ids":"xxx","registry_region":"us-east-1"},"docker-compose#v2.5.1":{"run":"xxx","config":".buildkite/docker/docker-compose.yml","env":["AWS_ACCESS_KEY_ID","AWS_SECRET_ACCESS_KEY","AWS_SESSION_TOKEN"]}},"agents":{"queue":"xxx"}}]}`
	assert.Equal(t, expected, strings.TrimSpace(buf.String()))
}
<commit_msg>pipeline_parser_test: demonstrate escaping dollar sign<commit_after>package agent

import (
	"bytes"
	"encoding/json"
	"fmt"
	"os"
	"strings"
	"testing"

	"github.com/buildkite/agent/v3/env"
	"github.com/stretchr/testify/assert"
)

func TestPipelineParserParsesYaml(t *testing.T) {
	environ := env.FromSlice([]string{`ENV_VAR_FRIEND="friend"`})

	result, err := PipelineParser{
		Filename: "awesome.yml",
		Pipeline: []byte("steps:\n  - label: \"hello ${ENV_VAR_FRIEND}\""),
		Env:      environ}.Parse()

	assert.NoError(t, err)
	j, err := json.Marshal(result)
	assert.Equal(t, `{"steps":[{"label":"hello \"friend\""}]}`, string(j))
}

func TestPipelineParserParsesYamlWithNoInterpolation(t *testing.T) {
	result, err := PipelineParser{
		Filename:        "awesome.yml",
		Pipeline:        []byte("steps:\n  - label: \"hello ${ENV_VAR_FRIEND}\""),
		NoInterpolation: true,
	}.Parse()

	assert.NoError(t, err)
	j, err := json.Marshal(result)
	assert.Equal(t, `{"steps":[{"label":"hello ${ENV_VAR_FRIEND}"}]}`, string(j))
}

func TestPipelineParserSupportsYamlMergesAndAnchors(t *testing.T) {
	complexYAML := `---
base_step: &base_step
  type: script
  agent_query_rules:
    - queue=default

steps:
  - <<: *base_step
    name: ':docker: building image'
    command: docker build .
    agents:
      queue: default`

	result, err := PipelineParser{
		Filename: "awesome.yml",
		Pipeline: []byte(complexYAML)}.Parse()

	assert.NoError(t, err)
	j, err := json.Marshal(result)
	assert.Equal(t, `{"base_step":{"type":"script","agent_query_rules":["queue=default"]},"steps":[{"type":"script","agent_query_rules":["queue=default"],"name":":docker: building image","command":"docker build .","agents":{"queue":"default"}}]}`, string(j))
}

func TestPipelineParserReturnsYamlParsingErrors(t *testing.T) {
	_, err := PipelineParser{Filename: "awesome.yml", Pipeline: []byte("steps: %blah%")}.Parse()
	assert.Error(t, err, `Failed to parse awesome.yml: found character that cannot start any token`, fmt.Sprintf("%s", err))
}

func TestPipelineParserReturnsJsonParsingErrors(t *testing.T) {
	_, err := PipelineParser{Filename: "awesome.json", Pipeline: []byte("{")}.Parse()
	assert.Error(t, err, `Failed to parse awesome.json: line 1: did not find expected node content`, fmt.Sprintf("%s", err))
}

func TestPipelineParserParsesJson(t *testing.T) {
	environ := env.FromSlice([]string{`ENV_VAR_FRIEND="friend"`})

	result, err := PipelineParser{
		Filename: "thing.json",
		Pipeline: []byte("\n\n     \n  { \"foo\": \"bye ${ENV_VAR_FRIEND}\" }\n"),
		Env:      environ}.Parse()

	assert.NoError(t, err)
	j, err := json.Marshal(result)
	assert.Equal(t, `{"foo":"bye \"friend\""}`, string(j))
}

func TestPipelineParserParsesJsonObjects(t *testing.T) {
	environ := env.FromSlice([]string{`ENV_VAR_FRIEND="friend"`})

	result, err := PipelineParser{Pipeline: []byte("\n\n     \n  { \"foo\": \"bye ${ENV_VAR_FRIEND}\" }\n"), Env: environ}.Parse()
	assert.NoError(t, err)
	j, err := json.Marshal(result)
	assert.Equal(t, `{"foo":"bye \"friend\""}`, string(j))
}

func TestPipelineParserParsesJsonArrays(t *testing.T) {
	environ := env.FromSlice([]string{`ENV_VAR_FRIEND="friend"`})

	result, err := PipelineParser{Pipeline: []byte("\n\n     \n  [ { \"foo\": \"bye ${ENV_VAR_FRIEND}\" } ]\n"), Env: environ}.Parse()
	assert.NoError(t, err)
	j, err := json.Marshal(result)
	assert.Equal(t, `{"steps":[{"foo":"bye \"friend\""}]}`, string(j))
}

func TestPipelineParserParsesTopLevelSteps(t *testing.T) {
	result, err := PipelineParser{Pipeline: []byte("---\n- name: Build\n  command: echo hello world\n- wait\n"), Env: nil}.Parse()
	assert.NoError(t, err)
	j, err := json.Marshal(result)
	assert.Equal(t, `{"steps":[{"name":"Build","command":"echo hello world"},"wait"]}`, string(j))
}

func TestPipelineParserPreservesBools(t *testing.T) {
	result, err := PipelineParser{Pipeline: []byte("steps:\n  - trigger: hello\n    async: true")}.Parse()
	assert.Nil(t, err)
	j, err := json.Marshal(result)
	assert.Equal(t, `{"steps":[{"trigger":"hello","async":true}]}`, string(j))
}

func TestPipelineParserPreservesInts(t *testing.T) {
	result, err := PipelineParser{Pipeline: []byte("steps:\n  - label: hello\n    parallelism: 10")}.Parse()
	assert.Nil(t, err)
	j, err := json.Marshal(result)
	assert.Equal(t, `{"steps":[{"label":"hello","parallelism":10}]}`, string(j))
}

func TestPipelineParserPreservesNull(t *testing.T) {
	result, err := PipelineParser{Pipeline: []byte("steps:\n  - wait: ~")}.Parse()
	assert.Nil(t, err)
	j, err := json.Marshal(result)
	assert.Equal(t, `{"steps":[{"wait":null}]}`, string(j))
}

func TestPipelineParserPreservesFloats(t *testing.T) {
	result, err := PipelineParser{Pipeline: []byte("steps:\n  - trigger: hello\n    llamas: 3.142")}.Parse()
	assert.Nil(t, err)
	j, err := json.Marshal(result)
	assert.Equal(t, `{"steps":[{"trigger":"hello","llamas":3.142}]}`, string(j))
}

func TestPipelineParserHandlesDates(t *testing.T) {
	result, err := PipelineParser{Pipeline: []byte("steps:\n  - trigger: hello\n    llamas: 2002-08-15T17:18:23.18-06:00")}.Parse()
	assert.Nil(t, err)
	j, err := json.Marshal(result)
	assert.Equal(t, `{"steps":[{"trigger":"hello","llamas":"2002-08-15T17:18:23.18-06:00"}]}`, string(j))
}

func TestPipelineParserInterpolatesKeysAsWellAsValues(t *testing.T) {
	var pipeline = `{
		"env": {
			"${FROM_ENV}TEST1": "MyTest",
			"TEST2": "${FROM_ENV}"
		}
	}`

	var decoded struct {
		Env map[string]string `json:"env"`
	}

	environ := env.FromSlice([]string{`FROM_ENV=llamas`})

	result, err := PipelineParser{Pipeline: []byte(pipeline), Env: environ}.Parse()
	if err != nil {
		t.Fatal(err)
	}

	err = decodeIntoStruct(&decoded, result)
	if err != nil {
		t.Fatal(err)
	}

	assert.Equal(t, `MyTest`, decoded.Env["llamasTEST1"])
	assert.Equal(t, `llamas`, decoded.Env["TEST2"])
}

func TestPipelineParserLoadsGlobalEnvBlockFirst(t *testing.T) {
	var pipeline = `{
		"env": {
			"TEAM1": "England",
			"TEAM2": "Australia",
			"HEADLINE": "${TEAM1} smashes ${TEAM2} to win the ashes in ${YEAR_FROM_SHELL}!!"
		},
		"steps": [{
			"command": "echo ${HEADLINE}"
		}]
	}`

	var decoded struct {
		Env   map[string]string `json:"env"`
		Steps []struct {
			Command string `json:"command"`
		} `json:"steps"`
	}

	environ := env.FromSlice([]string{`YEAR_FROM_SHELL=1912`})

	result, err := PipelineParser{Pipeline: []byte(pipeline), Env: environ}.Parse()
	if err != nil {
		t.Fatal(err)
	}

	err = decodeIntoStruct(&decoded, result)
	if err != nil {
		t.Fatal(err)
	}

	assert.Equal(t, `England`, decoded.Env["TEAM1"])
	assert.Equal(t, `England smashes Australia to win the ashes in 1912!!`, decoded.Env["HEADLINE"])
	assert.Equal(t, `echo England smashes Australia to win the ashes in 1912!!`, decoded.Steps[0].Command)
}

func decodeIntoStruct(into interface{}, from interface{}) error {
	b, err := json.Marshal(from)
	if err != nil {
		return err
	}
	return json.Unmarshal(b, into)
}

func TestPipelineParserLoadsSystemEnvironment(t *testing.T) {
	var pipeline = `{
		"steps": [{
			"command": "echo ${LLAMAS_ROCK?}"
		}]
	}`

	var decoded struct {
		Steps []struct {
			Command string `json:"command"`
		} `json:"steps"`
	}

	_, err := PipelineParser{Pipeline: []byte(pipeline)}.Parse()
	if err == nil {
		t.Fatalf("Expected $LLAMAS_ROCK: not set")
	}

	os.Setenv("LLAMAS_ROCK", "absolutely")
	defer os.Unsetenv("LLAMAS_ROCK")

	result2, err := PipelineParser{Pipeline: []byte(pipeline)}.Parse()
	if err != nil {
		t.Fatal(err)
	}

	err = decodeIntoStruct(&decoded, result2)
	if err != nil {
		t.Fatal(err)
	}

	if decoded.Steps[0].Command != "echo absolutely" {
		t.Fatalf("Unexpected: %q", decoded.Steps[0].Command)
	}
}

func TestPipelineParserPreservesOrderOfPlugins(t *testing.T) {
	var pipeline = `---
steps:
  - name: ":s3: xxx"
    command: "script/buildkite/xxx.sh"
    plugins:
      xxx/aws-assume-role#v0.1.0:
        role: arn:aws:iam::xxx:role/xxx
      ecr#v1.1.4:
        login: true
        account_ids: xxx
        registry_region: us-east-1
      docker-compose#v2.5.1:
        run: xxx
        config: .buildkite/docker/docker-compose.yml
        env:
          - AWS_ACCESS_KEY_ID
          - AWS_SECRET_ACCESS_KEY
          - AWS_SESSION_TOKEN
    agents:
      queue: xxx`

	result, err := PipelineParser{Pipeline: []byte(pipeline), Env: nil}.Parse()
	if err != nil {
		t.Fatal(err)
	}

	buf := &bytes.Buffer{}
	err = json.NewEncoder(buf).Encode(result)
	if err != nil {
		t.Fatal(err)
	}

	expected := `{"steps":[{"name":":s3: xxx","command":"script/buildkite/xxx.sh","plugins":{"xxx/aws-assume-role#v0.1.0":{"role":"arn:aws:iam::xxx:role/xxx"},"ecr#v1.1.4":{"login":true,"account_ids":"xxx","registry_region":"us-east-1"},"docker-compose#v2.5.1":{"run":"xxx","config":".buildkite/docker/docker-compose.yml","env":["AWS_ACCESS_KEY_ID","AWS_SECRET_ACCESS_KEY","AWS_SESSION_TOKEN"]}},"agents":{"queue":"xxx"}}]}`
	assert.Equal(t, expected, strings.TrimSpace(buf.String()))
}

func TestPipelineParserParsesConditionalWithEndOfLineAnchorDollarSign(t *testing.T) {
	for _, row := range []struct {
		noInterpolation bool
		pipeline        string
	}{
		// dollar sign must be escaped when interpolation is in effect
		{false, "steps:\n  - if: build.env(\"ACCOUNT\") =~ /^(foo|bar)\\$/"},
		{true, "steps:\n  - if: build.env(\"ACCOUNT\") =~ /^(foo|bar)$/"},
	} {
		result, err := PipelineParser{
			Pipeline:        []byte(row.pipeline),
			NoInterpolation: row.noInterpolation,
			Env:             env.New(),
		}.Parse()
		assert.NoError(t, err)
		j, _ := json.Marshal(result)
		assert.Equal(t, `{"steps":[{"if":"build.env(\"ACCOUNT\") =~ /^(foo|bar)$/"}]}`, string(j))
	}

}
<|endoftext|>


<commit_before>package indicators

import (
	"errors"
	"github.com/thetruetrade/gotrade"
)

// An Exponential Moving Average Indicator (Ema), no storage, for use in other indicators
type EmaWithoutStorage struct {
	*baseIndicator
	*baseFloatBounds

	// private variables
	periodTotal          float64
	periodCounter        int
	multiplier           float64
	previousEma          float64
	valueAvailableAction ValueAvailableActionFloat
	timePeriod           int
}

// NewEmaWithoutStorage creates an Exponential Moving Average Indicator (Ema) without storage
func NewEmaWithoutStorage(timePeriod int, valueAvailableAction ValueAvailableActionFloat) (indicator *EmaWithoutStorage, err error) {

	// an indicator without storage MUST have a value available action
	if valueAvailableAction == nil {
		return nil, ErrValueAvailableActionIsNil
	}

	// the minimum timeperiod for this indicator is 2
	if timePeriod < 2 {
		return nil, errors.New("timePeriod is less than the minimum (2)")
	}

	// check the maximum timeperiod
	if timePeriod > MaximumLookbackPeriod {
		return nil, errors.New("timePeriod is greater than the maximum (100000)")
	}

	lookback := timePeriod - 1
	ind := EmaWithoutStorage{
		baseIndicator:        newBaseIndicator(lookback),
		baseFloatBounds:      newBaseFloatBounds(),
		periodCounter:        timePeriod * -1,
		multiplier:           float64(2.0 / float64(timePeriod+1.0)),
		valueAvailableAction: valueAvailableAction,
		timePeriod:           timePeriod,
	}

	return &ind, err
}

// An Exponential Moving Average Indicator (Ema)
type Ema struct {
	*EmaWithoutStorage
	selectData gotrade.DataSelectionFunc

	// public variables
	Data []float64
}

// NewEma creates an Exponential Moving Average Indicator (Ema) for online usage
func NewEma(timePeriod int, selectData gotrade.DataSelectionFunc) (indicator *Ema, err error) {
	newEma := Ema{selectData: selectData}
	newEma.EmaWithoutStorage, err = NewEmaWithoutStorage(timePeriod,
		func(dataItem float64, streamBarIndex int) {
			newEma.Data = append(newEma.Data, dataItem)
		})

	return &newEma, err
}

// NewDefaultEma creates an Exponential Moving Average (Ema) for online usage with default parameters
//	- timePeriod: 25
func NewDefaultEma() (indicator *Ema, err error) {
	timePeriod := 25
	return NewEma(timePeriod, gotrade.UseClosePrice)
}

// NewEmaWithSrcLen creates an Exponential Moving Average (Ema) for offline usage
func NewEmaWithSrcLen(sourceLength uint, timePeriod int, selectData gotrade.DataSelectionFunc) (indicator *Ema, err error) {
	ind, err := NewEma(timePeriod, selectData)

	// only initialise the storage if there is enough source data to require it
	if sourceLength-uint(ind.GetLookbackPeriod()) > 1 {
		ind.Data = make([]float64, 0, sourceLength-uint(ind.GetLookbackPeriod()))
	}

	return ind, err
}

// NewDefaultEmaWithSrcLen creates an Exponential Moving Average (Ema) for offline usage with default parameters
func NewDefaultEmaWithSrcLen(sourceLength uint) (indicator *Ema, err error) {
	ind, err := NewDefaultEma()

	// only initialise the storage if there is enough source data to require it
	if sourceLength-uint(ind.GetLookbackPeriod()) > 1 {
		ind.Data = make([]float64, 0, sourceLength-uint(ind.GetLookbackPeriod()))
	}

	return ind, err
}

// NewEmaForStream creates an Exponential Moving Average (Ema) for online usage with a source data stream
func NewEmaForStream(priceStream gotrade.DOHLCVStreamSubscriber, timePeriod int, selectData gotrade.DataSelectionFunc) (indicator *Ema, err error) {
	ind, err := NewEma(timePeriod, selectData)
	priceStream.AddTickSubscription(ind)
	return ind, err
}

// NewDefaultEmaForStream creates an Exponential Moving Average (Ema) for online usage with a source data stream
func NewDefaultEmaForStream(priceStream gotrade.DOHLCVStreamSubscriber) (indicator *Ema, err error) {
	ind, err := NewDefaultEma()
	priceStream.AddTickSubscription(ind)
	return ind, err
}

// NewEmaForStreamWithSrcLen creates an Exponential Moving Average (Ema) for offline usage with a source data stream
func NewEmaForStreamWithSrcLen(sourceLength uint, priceStream gotrade.DOHLCVStreamSubscriber, timePeriod int, selectData gotrade.DataSelectionFunc) (indicator *Ema, err error) {
	ind, err := NewEmaWithSrcLen(sourceLength, timePeriod, selectData)
	priceStream.AddTickSubscription(ind)
	return ind, err
}

// NewDefaultEmaForStreamWithSrcLen creates an Exponential Moving Average (Ema) for offline usage with a source data stream
func NewDefaultEmaForStreamWithSrcLen(sourceLength uint, priceStream gotrade.DOHLCVStreamSubscriber) (indicator *Ema, err error) {
	ind, err := NewDefaultEmaWithSrcLen(sourceLength)
	priceStream.AddTickSubscription(ind)
	return ind, err
}

// ReceiveDOHLCVTick consumes a source data DOHLCV price tick
func (ema *Ema) ReceiveDOHLCVTick(tickData gotrade.DOHLCV, streamBarIndex int) {
	var selectedData = ema.selectData(tickData)
	ema.ReceiveTick(selectedData, streamBarIndex)
}

func (ema *EmaWithoutStorage) ReceiveTick(tickData float64, streamBarIndex int) {
	ema.periodCounter += 1
	if ema.periodCounter < 0 {
		ema.periodTotal += tickData
	} else if ema.periodCounter == 0 {

		// increment the number of results this indicator can be expected to return
		ema.dataLength += 1

		if ema.validFromBar == -1 {
			// set the streamBarIndex from which this indicator returns valid results
			ema.validFromBar = streamBarIndex
		}

		ema.periodTotal += tickData
		result := ema.periodTotal / float64(ema.timePeriod)
		ema.previousEma = result

		// update the maximum result value
		if result > ema.maxValue {
			ema.maxValue = result
		}

		// update the minimum result value
		if result < ema.minValue {
			ema.minValue = result
		}

		// notify of a new result value though the value available action
		ema.valueAvailableAction(ema.previousEma, streamBarIndex)

	} else if ema.periodCounter > 0 {
		// increment the number of results this indicator can be expected to return
		ema.dataLength += 1

		result := (tickData-ema.previousEma)*ema.multiplier + ema.previousEma
		ema.previousEma = result

		// update the maximum result value
		if result > ema.maxValue {
			ema.maxValue = result
		}

		// update the minimum result value
		if result < ema.minValue {
			ema.minValue = result
		}

		// notify of a new result value though the value available action
		ema.valueAvailableAction(ema.previousEma, streamBarIndex)
	}
}
<commit_msg>#76 Remove duplication - ema<commit_after>package indicators

import (
	"errors"
	"github.com/thetruetrade/gotrade"
)

// An Exponential Moving Average Indicator (Ema), no storage, for use in other indicators
type EmaWithoutStorage struct {
	*baseIndicatorWithFloatBounds

	// private variables
	periodTotal   float64
	periodCounter int
	multiplier    float64
	previousEma   float64
	timePeriod    int
}

// NewEmaWithoutStorage creates an Exponential Moving Average Indicator (Ema) without storage
func NewEmaWithoutStorage(timePeriod int, valueAvailableAction ValueAvailableActionFloat) (indicator *EmaWithoutStorage, err error) {

	// an indicator without storage MUST have a value available action
	if valueAvailableAction == nil {
		return nil, ErrValueAvailableActionIsNil
	}

	// the minimum timeperiod for this indicator is 2
	if timePeriod < 2 {
		return nil, errors.New("timePeriod is less than the minimum (2)")
	}

	// check the maximum timeperiod
	if timePeriod > MaximumLookbackPeriod {
		return nil, errors.New("timePeriod is greater than the maximum (100000)")
	}

	lookback := timePeriod - 1
	ind := EmaWithoutStorage{
		baseIndicatorWithFloatBounds: newBaseIndicatorWithFloatBounds(lookback, valueAvailableAction),
		periodCounter:                timePeriod * -1,
		multiplier:                   float64(2.0 / float64(timePeriod+1.0)),
		timePeriod:                   timePeriod,
	}

	return &ind, err
}

// An Exponential Moving Average Indicator (Ema)
type Ema struct {
	*EmaWithoutStorage
	selectData gotrade.DataSelectionFunc

	// public variables
	Data []float64
}

// NewEma creates an Exponential Moving Average Indicator (Ema) for online usage
func NewEma(timePeriod int, selectData gotrade.DataSelectionFunc) (indicator *Ema, err error) {
	newEma := Ema{selectData: selectData}
	newEma.EmaWithoutStorage, err = NewEmaWithoutStorage(timePeriod,
		func(dataItem float64, streamBarIndex int) {
			newEma.Data = append(newEma.Data, dataItem)
		})

	return &newEma, err
}

// NewDefaultEma creates an Exponential Moving Average (Ema) for online usage with default parameters
//	- timePeriod: 25
func NewDefaultEma() (indicator *Ema, err error) {
	timePeriod := 25
	return NewEma(timePeriod, gotrade.UseClosePrice)
}

// NewEmaWithSrcLen creates an Exponential Moving Average (Ema) for offline usage
func NewEmaWithSrcLen(sourceLength uint, timePeriod int, selectData gotrade.DataSelectionFunc) (indicator *Ema, err error) {
	ind, err := NewEma(timePeriod, selectData)

	// only initialise the storage if there is enough source data to require it
	if sourceLength-uint(ind.GetLookbackPeriod()) > 1 {
		ind.Data = make([]float64, 0, sourceLength-uint(ind.GetLookbackPeriod()))
	}

	return ind, err
}

// NewDefaultEmaWithSrcLen creates an Exponential Moving Average (Ema) for offline usage with default parameters
func NewDefaultEmaWithSrcLen(sourceLength uint) (indicator *Ema, err error) {
	ind, err := NewDefaultEma()

	// only initialise the storage if there is enough source data to require it
	if sourceLength-uint(ind.GetLookbackPeriod()) > 1 {
		ind.Data = make([]float64, 0, sourceLength-uint(ind.GetLookbackPeriod()))
	}

	return ind, err
}

// NewEmaForStream creates an Exponential Moving Average (Ema) for online usage with a source data stream
func NewEmaForStream(priceStream gotrade.DOHLCVStreamSubscriber, timePeriod int, selectData gotrade.DataSelectionFunc) (indicator *Ema, err error) {
	ind, err := NewEma(timePeriod, selectData)
	priceStream.AddTickSubscription(ind)
	return ind, err
}

// NewDefaultEmaForStream creates an Exponential Moving Average (Ema) for online usage with a source data stream
func NewDefaultEmaForStream(priceStream gotrade.DOHLCVStreamSubscriber) (indicator *Ema, err error) {
	ind, err := NewDefaultEma()
	priceStream.AddTickSubscription(ind)
	return ind, err
}

// NewEmaForStreamWithSrcLen creates an Exponential Moving Average (Ema) for offline usage with a source data stream
func NewEmaForStreamWithSrcLen(sourceLength uint, priceStream gotrade.DOHLCVStreamSubscriber, timePeriod int, selectData gotrade.DataSelectionFunc) (indicator *Ema, err error) {
	ind, err := NewEmaWithSrcLen(sourceLength, timePeriod, selectData)
	priceStream.AddTickSubscription(ind)
	return ind, err
}

// NewDefaultEmaForStreamWithSrcLen creates an Exponential Moving Average (Ema) for offline usage with a source data stream
func NewDefaultEmaForStreamWithSrcLen(sourceLength uint, priceStream gotrade.DOHLCVStreamSubscriber) (indicator *Ema, err error) {
	ind, err := NewDefaultEmaWithSrcLen(sourceLength)
	priceStream.AddTickSubscription(ind)
	return ind, err
}

// ReceiveDOHLCVTick consumes a source data DOHLCV price tick
func (ind *Ema) ReceiveDOHLCVTick(tickData gotrade.DOHLCV, streamBarIndex int) {
	var selectedData = ind.selectData(tickData)
	ind.ReceiveTick(selectedData, streamBarIndex)
}

func (ind *EmaWithoutStorage) ReceiveTick(tickData float64, streamBarIndex int) {
	ind.periodCounter += 1
	if ind.periodCounter < 0 {
		ind.periodTotal += tickData
	} else if ind.periodCounter == 0 {

		ind.periodTotal += tickData
		result := ind.periodTotal / float64(ind.timePeriod)
		ind.previousEma = result

		ind.UpdateIndicatorWithNewValue(result, streamBarIndex)

	} else if ind.periodCounter > 0 {

		result := (tickData-ind.previousEma)*ind.multiplier + ind.previousEma
		ind.previousEma = result

		ind.UpdateIndicatorWithNewValue(result, streamBarIndex)
	}
}
<|endoftext|>


<commit_before>package ops

import (
	"fmt"
	"regexp"
	"strings"
)

// negate returns the negative version of the input.  We avoid any
// conversion to a number her and just treat it as a raw string.
func negate(s string) string {
	if strings.HasPrefix(s, "-") {
		return s[1:]
	}
	return "-" + s
}

// Mutator is a function that knows how to mutate a Line
type Mutator func(l *Line) error

// Line represents a line we read from a csv file
type Line struct {
	LineNo int
	Record []string
}

// NewLine creates a Line structure
func NewLine(lineNo int, record []string) *Line {
	return &Line{lineNo, record}
}

func ReplaceHeader(header []string) Mutator {
	return func(l *Line) error {
		if l.LineNo == 1 {
			l.Record = header
		}
		return nil
	}
}

func MoveAndNegateIfPresent(from int, to int) Mutator {
	return func(l *Line) error {
		if l.LineNo != 1 {
			value := l.Record[from]
			if value != "" {
				l.Record[to] = negate(value)
				l.Record[from] = ""
			}
		}
		return nil
	}
}

func Negate(i int) Mutator {
	return func(l *Line) error {
		if l.LineNo != 1 {
			value := l.Record[i]
			l.Record[i] = negate(value)
		}
		return nil
	}
}

// CheckWithdrawal converts an entry like 'Check Withdrawal: #999999'
// to '(#999999)'
func CheckWithdrawal(i int) Mutator {
	r := regexp.MustCompile(`^Check Withdrawal: (#\d+).*$`)
	return func(l *Line) error {
		if l.LineNo != 1 {
			value := l.Record[i]
			if matches := r.FindStringSubmatch(value); matches != nil {
				l.Record[i] = fmt.Sprintf("(%s)", matches[1])
			}
		}
		return nil
	}
}

// RemoveText removes the substring rem, if it is present, up to one time.
func RemoveText(i int, rem string) Mutator {
	return func(l *Line) error {
		if l.LineNo != 1 {
			value := l.Record[i]
			l.Record[i] = strings.Replace(value, rem, "", 1)
		}
		return nil
	}
}

func DeparenNegatives(i int) Mutator {
	return func(l *Line) error {
		if l.LineNo != 1 {
			value := l.Record[i]
			if strings.HasPrefix(value, "(") && strings.HasSuffix(value, ")") {
				value = strings.TrimPrefix(value, "(")
				value = strings.TrimSuffix(value, ")")
				value = negate(value)
				l.Record[i] = value
			}
		}
		return nil
	}
}

func EnsureDollars(i int) Mutator {
	return func(l *Line) error {
		if l.LineNo != 1 {
			value := l.Record[i]
			if value != "" && (!strings.HasPrefix(value, "$")) {
				l.Record[i] = "$" + value
			}
		}
		return nil
	}
}

func StripCommas(i int) Mutator {
	return func(l *Line) error {
		if l.LineNo != 1 {
			value := l.Record[i]
			if strings.Contains(value, ",") {
				l.Record[i] = strings.Replace(value, ",", "", -1)
			}
		}
		return nil
	}
}

func StripSuffix(i int, suffix string) Mutator {
	return func(l *Line) error {
		if l.LineNo != 1 {
			l.Record[i] = strings.TrimSuffix(l.Record[i], suffix)
		}
		return nil
	}
}

func StripNewlines(l *Line) error {
	for i, c := range l.Record {
		if strings.Contains(c, "\n") {
			l.Record[i] = strings.Replace(c, "\n", "", -1)
		}
	}
	return nil
}
<commit_msg>add func comment<commit_after>package ops

import (
	"fmt"
	"regexp"
	"strings"
)

// negate returns the negative version of the input.  We avoid any
// conversion to a number her and just treat it as a raw string.
func negate(s string) string {
	if strings.HasPrefix(s, "-") {
		return s[1:]
	}
	return "-" + s
}

// Mutator is a function that knows how to mutate a Line
type Mutator func(l *Line) error

// Line represents a line we read from a csv file
type Line struct {
	LineNo int
	Record []string
}

// NewLine creates a Line structure
func NewLine(lineNo int, record []string) *Line {
	return &Line{lineNo, record}
}

// ReplaceHeader removes the first line
func ReplaceHeader(header []string) Mutator {
	return func(l *Line) error {
		if l.LineNo == 1 {
			l.Record = header
		}
		return nil
	}
}

func MoveAndNegateIfPresent(from int, to int) Mutator {
	return func(l *Line) error {
		if l.LineNo != 1 {
			value := l.Record[from]
			if value != "" {
				l.Record[to] = negate(value)
				l.Record[from] = ""
			}
		}
		return nil
	}
}

func Negate(i int) Mutator {
	return func(l *Line) error {
		if l.LineNo != 1 {
			value := l.Record[i]
			l.Record[i] = negate(value)
		}
		return nil
	}
}

// CheckWithdrawal converts an entry like 'Check Withdrawal: #999999'
// to '(#999999)'
func CheckWithdrawal(i int) Mutator {
	r := regexp.MustCompile(`^Check Withdrawal: (#\d+).*$`)
	return func(l *Line) error {
		if l.LineNo != 1 {
			value := l.Record[i]
			if matches := r.FindStringSubmatch(value); matches != nil {
				l.Record[i] = fmt.Sprintf("(%s)", matches[1])
			}
		}
		return nil
	}
}

// RemoveText removes the substring rem, if it is present, up to one time.
func RemoveText(i int, rem string) Mutator {
	return func(l *Line) error {
		if l.LineNo != 1 {
			value := l.Record[i]
			l.Record[i] = strings.Replace(value, rem, "", 1)
		}
		return nil
	}
}

func DeparenNegatives(i int) Mutator {
	return func(l *Line) error {
		if l.LineNo != 1 {
			value := l.Record[i]
			if strings.HasPrefix(value, "(") && strings.HasSuffix(value, ")") {
				value = strings.TrimPrefix(value, "(")
				value = strings.TrimSuffix(value, ")")
				value = negate(value)
				l.Record[i] = value
			}
		}
		return nil
	}
}

func EnsureDollars(i int) Mutator {
	return func(l *Line) error {
		if l.LineNo != 1 {
			value := l.Record[i]
			if value != "" && (!strings.HasPrefix(value, "$")) {
				l.Record[i] = "$" + value
			}
		}
		return nil
	}
}

func StripCommas(i int) Mutator {
	return func(l *Line) error {
		if l.LineNo != 1 {
			value := l.Record[i]
			if strings.Contains(value, ",") {
				l.Record[i] = strings.Replace(value, ",", "", -1)
			}
		}
		return nil
	}
}

func StripSuffix(i int, suffix string) Mutator {
	return func(l *Line) error {
		if l.LineNo != 1 {
			l.Record[i] = strings.TrimSuffix(l.Record[i], suffix)
		}
		return nil
	}
}

func StripNewlines(l *Line) error {
	for i, c := range l.Record {
		if strings.Contains(c, "\n") {
			l.Record[i] = strings.Replace(c, "\n", "", -1)
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package appdash

import (
	"errors"
	"fmt"
	"net/url"
	"reflect"
	"time"

	influxDBClient "github.com/influxdb/influxdb/client"
	influxDBServer "github.com/influxdb/influxdb/cmd/influxd/run"
	influxDBModels "github.com/influxdb/influxdb/models"
)

const (
	dbName               string = "appdash" // InfluxDB db name.
	spanMeasurementName  string = "spans"   // InfluxDB container name for trace spans.
	defaultTracesPerPage int    = 10        // Default number of traces per page.
)

// Compile-time "implements" check.
var _ interface {
	Store
	Queryer
} = (*InfluxDBStore)(nil)

// TODO: should be a constant.
var zeroID = fmt.Sprintf("%016x", uint64(0))

type InfluxDBStore struct {
	con           *influxDBClient.Client // InfluxDB client connection.
	server        *influxDBServer.Server // InfluxDB API server.
	tracesPerPage int                    // Number of traces per page.
}

func (in *InfluxDBStore) Collect(id SpanID, anns ...Annotation) error {
	p, err := in.findSpanPoint(id)
	if err != nil {
		return err
	}

	// trace_id, span_id & parent_id are set as tags
	// because InfluxDB tags are indexed & those values
	// are used later on queries.
	tags := map[string]string{
		"trace_id":  id.Trace.String(),
		"span_id":   id.Span.String(),
		"parent_id": id.Parent.String(),
	}

	// Saving annotations as InfluxDB measurement spans fields
	// which are not indexed.
	fields := make(map[string]interface{}, len(anns))
	for _, ann := range anns {
		fields[ann.Key] = string(ann.Value)
	}

	if p != nil { // span exists on DB.
		p.Measurement = spanMeasurementName
		p.Tags = tags
		p.Fields = fields
	} else { // new span to be saved on DB.
		p = &influxDBClient.Point{
			Measurement: spanMeasurementName,
			Tags:        tags,   // indexed metadata.
			Fields:      fields, // non-indexed metadata.
			Time:        time.Now().UTC(),
		}
	}

	// InfluxDB point represents a single span.
	pts := []influxDBClient.Point{*p}
	bps := influxDBClient.BatchPoints{
		Points:          pts,
		Database:        dbName,
		RetentionPolicy: "default",
	}
	_, writeErr := in.con.Write(bps)
	if writeErr != nil {
		return writeErr
	}
	return nil
}

func (in *InfluxDBStore) Trace(id ID) (*Trace, error) {
	trace := &Trace{}

	// GROUP BY * -> meaning group by all tags(trace_id, span_id & parent_id)
	// grouping by all tags includes those and it's values on the query response.
	q := fmt.Sprintf("SELECT * FROM spans WHERE trace_id='%s' GROUP BY *", id)
	result, err := in.executeOneQuery(q)
	if err != nil {
		return nil, err
	}

	// result.Series -> A slice containing all the spans.
	if len(result.Series) == 0 {
		return nil, errors.New("trace not found")
	}

	// Iterate over series(spans) to create trace children's & set trace fields.
	var rootSpanSet bool
	for _, s := range result.Series {
		var isRootSpan bool
		span, err := newSpanFromRow(&s)
		if err != nil {
			return nil, err
		}
		annotations, err := annotationsFromRow(&s)
		if err != nil {
			return trace, nil
		}
		span.Annotations = *annotations
		if span.ID.IsRoot() && rootSpanSet {
			return nil, errors.New("unexpected multiple root spans")
		}
		if span.ID.IsRoot() && !rootSpanSet {
			isRootSpan = true
		}
		if isRootSpan { // root span.
			trace.Span = *span
			rootSpanSet = true
		} else { // children span.
			trace.Sub = append(trace.Sub, &Trace{Span: *span})
		}
	}
	return trace, nil
}

func (in *InfluxDBStore) Traces() ([]*Trace, error) {
	traces := make([]*Trace, 0)

	// GROUP BY * -> meaning group by all tags(trace_id, span_id & parent_id)
	// grouping by all tags includes those and it's values on the query response.
	rootSpansQuery := fmt.Sprintf("SELECT * FROM spans WHERE parent_id='%s' GROUP BY * LIMIT %d", zeroID, in.tracesPerPage)
	rootSpansResult, err := in.executeOneQuery(rootSpansQuery)
	if err != nil {
		return nil, err
	}

	// result.Series -> A slice containing all the spans.
	if len(rootSpansResult.Series) == 0 {
		return traces, nil
	}

	// Cache to keep track of traces to be returned.
	tracesCache := make(map[ID]*Trace, 0)

	// Iterate over series(spans) to create traces.
	for _, s := range rootSpansResult.Series {
		span, err := newSpanFromRow(&s)
		if err != nil {
			return nil, err
		}
		annotations, err := annotationsFromRow(&s)
		if err != nil {
			return nil, err
		}
		span.Annotations = *annotations
		_, present := tracesCache[span.ID.Trace]
		if !present {
			tracesCache[span.ID.Trace] = &Trace{Span: *span}
		} else {
			return nil, errors.New("duplicated root span")
		}
	}

	// Using 'OR' since 'IN' not supported yet.
	where := `WHERE `
	var i int = 1
	for _, trace := range tracesCache {
		where += fmt.Sprintf("(trace_id='%s' AND parent_id!='%s')", trace.Span.ID.Trace, zeroID)

		// Adds 'OR' except for last iteration.
		if i != len(tracesCache) && len(tracesCache) > 1 {
			where += " OR "
		}
		i += 1
	}

	// Queries for all children spans of the traces to be returned.
	childrenSpansQuery := fmt.Sprintf("SELECT * FROM spans %s GROUP BY *", where)
	childrenSpansResult, err := in.executeOneQuery(childrenSpansQuery)
	if err != nil {
		return nil, err
	}

	// Iterate over series(children spans) to create sub-traces
	// and associates sub-traces with it's parent trace.
	for _, s := range childrenSpansResult.Series {
		span, err := newSpanFromRow(&s)
		if err != nil {
			return nil, err
		}
		annotations, err := annotationsFromRow(&s)
		if err != nil {
			return nil, err
		}
		span.Annotations = *annotations
		trace, present := tracesCache[span.ID.Trace]
		if !present { // Root trace not added.
			return nil, errors.New("parent not found")
		} else { // Root trace already added so append a sub-trace.
			trace.Sub = append(trace.Sub, &Trace{Span: *span})
		}
	}
	for _, trace := range tracesCache {
		traces = append(traces, trace)
	}
	return traces, nil
}

func (in *InfluxDBStore) Close() {
	in.server.Close()
}

func (in *InfluxDBStore) createDBIfNotExists() error {
	// If no errors query execution was successfully - either DB was created or already exists.
	response, err := in.con.Query(influxDBClient.Query{
		Command: fmt.Sprintf("%s %s", "CREATE DATABASE IF NOT EXISTS", dbName),
	})
	if err != nil {
		return err
	}
	if response.Error() != nil {
		return response.Error()
	}
	return nil
}

func (in *InfluxDBStore) executeOneQuery(command string) (*influxDBClient.Result, error) {
	response, err := in.con.Query(influxDBClient.Query{
		Command:  command,
		Database: dbName,
	})
	if err != nil {
		return nil, err
	}
	if response.Error() != nil {
		return nil, response.Error()
	}

	// Expecting one result, since a single query is executed.
	if len(response.Results) != 1 {
		return nil, errors.New("unexpected number of results for an influxdb single query")
	}
	return &response.Results[0], nil
}

func (in *InfluxDBStore) findSpanPoint(ID SpanID) (*influxDBClient.Point, error) {
	q := fmt.Sprintf(`
		SELECT * FROM spans WHERE trace_id='%s' AND span_id='%s' AND parent_id='%s' GROUP BY *
	`, ID.Trace, ID.Span, ID.Parent)
	result, err := in.executeOneQuery(q)
	if err != nil {
		return nil, err
	}
	if len(result.Series) == 0 {
		return nil, nil
	}
	if len(result.Series) > 1 {
		return nil, errors.New("unexpected multiple series")
	}
	r := result.Series[0]
	if len(r.Values) == 0 {
		return nil, errors.New("unexpected empty series")
	}
	p := influxDBClient.Point{}
	fields := r.Values[0]
	for i, field := range fields {
		key := r.Columns[i]
		if key == "time" {
			switch field.(type) {
			case string:
				t, err := time.Parse(time.RFC3339Nano, field.(string))
				if err != nil {
					return nil, err
				}
				p.Time = t
			default:
				return nil, fmt.Errorf("unexpected time type: %v", reflect.TypeOf(field))
			}
		}
	}
	return &p, err
}

func (in *InfluxDBStore) init(server *influxDBServer.Server) error {
	in.server = server
	url, err := url.Parse(fmt.Sprintf("http://%s:%d", influxDBClient.DefaultHost, influxDBClient.DefaultPort))
	if err != nil {
		return err
	}
	con, err := influxDBClient.NewClient(influxDBClient.Config{URL: *url})
	if err != nil {
		return err
	}
	in.con = con
	if err := in.createDBIfNotExists(); err != nil {
		return err
	}
	in.tracesPerPage = defaultTracesPerPage
	return nil
}

func annotationsFromRow(r *influxDBModels.Row) (*Annotations, error) {
	// Actually an influxDBModels.Row represents a single InfluxDB serie.
	// r.Values[n] is a slice containing span's annotation values.
	var fields []interface{}
	if len(r.Values) == 1 {
		fields = r.Values[0]
	}

	// len(r.Values) might be greater than one - meaning there are
	// some spans to drop, see: InfluxDBStore.Collect(...).
	// If so last one is picked.
	if len(r.Values) > 1 {
		fields = r.Values[len(r.Values)-1]
	}
	annotations := make(Annotations, 0)

	// Iterates over fields which represent span's annotation values.
	for i, field := range fields {
		// It is safe to do column[0] (eg. 'Server.Request.Method')
		// matches fields[0] (eg. 'GET')
		key := r.Columns[i]
		var value []byte
		switch field.(type) {
		case string:
			value = []byte(field.(string))
		case nil:
		default:
			return nil, fmt.Errorf("unexpected field type: %v", reflect.TypeOf(field))
		}
		a := Annotation{
			Key:   key,
			Value: value,
		}
		annotations = append(annotations, a)
	}

	return &annotations, nil
}

func newSpanFromRow(r *influxDBModels.Row) (*Span, error) {
	span := &Span{}
	traceID, err := ParseID(r.Tags["trace_id"])
	if err != nil {
		return nil, err
	}
	spanID, err := ParseID(r.Tags["span_id"])
	if err != nil {
		return nil, err
	}
	parentID, err := ParseID(r.Tags["parent_id"])
	if err != nil {
		return nil, err
	}
	span.ID = SpanID{
		Trace:  ID(traceID),
		Span:   ID(spanID),
		Parent: ID(parentID),
	}
	return span, nil
}

type InfluxDBStoreConfig struct {
	Server    *influxDBServer.Config
	BuildInfo *influxDBServer.BuildInfo
}

func NewInfluxDBStore(config InfluxDBStoreConfig) (*InfluxDBStore, error) {
	//TODO: add Authentication.
	s, err := influxDBServer.NewServer(config.Server, config.BuildInfo)
	if err != nil {
		return nil, err
	}
	if err := s.Open(); err != nil {
		return nil, err
	}
	var in InfluxDBStore
	if err := in.init(s); err != nil {
		return nil, err
	}
	return &in, nil
}
<commit_msg>updates to preserve existing span fields<commit_after>package appdash

import (
	"errors"
	"fmt"
	"net/url"
	"reflect"
	"time"

	influxDBClient "github.com/influxdb/influxdb/client"
	influxDBServer "github.com/influxdb/influxdb/cmd/influxd/run"
	influxDBModels "github.com/influxdb/influxdb/models"
)

const (
	dbName               string = "appdash" // InfluxDB db name.
	spanMeasurementName  string = "spans"   // InfluxDB container name for trace spans.
	defaultTracesPerPage int    = 10        // Default number of traces per page.
)

// Compile-time "implements" check.
var _ interface {
	Store
	Queryer
} = (*InfluxDBStore)(nil)

// TODO: should be a constant.
var zeroID = fmt.Sprintf("%016x", uint64(0))

// pointFields -> influxDBClient.Point.Fields
type pointFields map[string]interface{}

type InfluxDBStore struct {
	con           *influxDBClient.Client // InfluxDB client connection.
	server        *influxDBServer.Server // InfluxDB API server.
	tracesPerPage int                    // Number of traces per page.
}

func (in *InfluxDBStore) Collect(id SpanID, anns ...Annotation) error {
	p, err := in.findSpanPoint(id)
	if err != nil {
		return err
	}

	// trace_id, span_id & parent_id are set as tags
	// because InfluxDB tags are indexed & those values
	// are used later on queries.
	tags := map[string]string{
		"trace_id":  id.Trace.String(),
		"span_id":   id.Span.String(),
		"parent_id": id.Parent.String(),
	}

	// Saving annotations as InfluxDB measurement spans fields
	// which are not indexed.
	fields := make(map[string]interface{}, len(anns))
	for _, ann := range anns {
		fields[ann.Key] = string(ann.Value)
	}

	if p != nil { // span exists on DB.
		p.Measurement = spanMeasurementName
		p.Tags = tags
		// Extends fields in order to preserve the p.Fields
		// already saved on DB.
		p.Fields = extendFields(fields, p.Fields)
	} else { // new span to be saved on DB.
		p = &influxDBClient.Point{
			Measurement: spanMeasurementName,
			Tags:        tags,   // indexed metadata.
			Fields:      fields, // non-indexed metadata.
			Time:        time.Now().UTC(),
		}
	}

	// InfluxDB point represents a single span.
	pts := []influxDBClient.Point{*p}
	bps := influxDBClient.BatchPoints{
		Points:          pts,
		Database:        dbName,
		RetentionPolicy: "default",
	}
	_, writeErr := in.con.Write(bps)
	if writeErr != nil {
		return writeErr
	}
	return nil
}

func (in *InfluxDBStore) Trace(id ID) (*Trace, error) {
	trace := &Trace{}

	// GROUP BY * -> meaning group by all tags(trace_id, span_id & parent_id)
	// grouping by all tags includes those and it's values on the query response.
	q := fmt.Sprintf("SELECT * FROM spans WHERE trace_id='%s' GROUP BY *", id)
	result, err := in.executeOneQuery(q)
	if err != nil {
		return nil, err
	}

	// result.Series -> A slice containing all the spans.
	if len(result.Series) == 0 {
		return nil, errors.New("trace not found")
	}

	// Iterate over series(spans) to create trace children's & set trace fields.
	var rootSpanSet bool
	for _, s := range result.Series {
		var isRootSpan bool
		span, err := newSpanFromRow(&s)
		if err != nil {
			return nil, err
		}
		annotations, err := annotationsFromRow(&s)
		if err != nil {
			return trace, nil
		}
		span.Annotations = *annotations
		if span.ID.IsRoot() && rootSpanSet {
			return nil, errors.New("unexpected multiple root spans")
		}
		if span.ID.IsRoot() && !rootSpanSet {
			isRootSpan = true
		}
		if isRootSpan { // root span.
			trace.Span = *span
			rootSpanSet = true
		} else { // children span.
			trace.Sub = append(trace.Sub, &Trace{Span: *span})
		}
	}
	return trace, nil
}

func (in *InfluxDBStore) Traces() ([]*Trace, error) {
	traces := make([]*Trace, 0)

	// GROUP BY * -> meaning group by all tags(trace_id, span_id & parent_id)
	// grouping by all tags includes those and it's values on the query response.
	rootSpansQuery := fmt.Sprintf("SELECT * FROM spans WHERE parent_id='%s' GROUP BY * LIMIT %d", zeroID, in.tracesPerPage)
	rootSpansResult, err := in.executeOneQuery(rootSpansQuery)
	if err != nil {
		return nil, err
	}

	// result.Series -> A slice containing all the spans.
	if len(rootSpansResult.Series) == 0 {
		return traces, nil
	}

	// Cache to keep track of traces to be returned.
	tracesCache := make(map[ID]*Trace, 0)

	// Iterate over series(spans) to create traces.
	for _, s := range rootSpansResult.Series {
		span, err := newSpanFromRow(&s)
		if err != nil {
			return nil, err
		}
		annotations, err := annotationsFromRow(&s)
		if err != nil {
			return nil, err
		}
		span.Annotations = *annotations
		_, present := tracesCache[span.ID.Trace]
		if !present {
			tracesCache[span.ID.Trace] = &Trace{Span: *span}
		} else {
			return nil, errors.New("duplicated root span")
		}
	}

	// Using 'OR' since 'IN' not supported yet.
	where := `WHERE `
	var i int = 1
	for _, trace := range tracesCache {
		where += fmt.Sprintf("(trace_id='%s' AND parent_id!='%s')", trace.Span.ID.Trace, zeroID)

		// Adds 'OR' except for last iteration.
		if i != len(tracesCache) && len(tracesCache) > 1 {
			where += " OR "
		}
		i += 1
	}

	// Queries for all children spans of the traces to be returned.
	childrenSpansQuery := fmt.Sprintf("SELECT * FROM spans %s GROUP BY *", where)
	childrenSpansResult, err := in.executeOneQuery(childrenSpansQuery)
	if err != nil {
		return nil, err
	}

	// Iterate over series(children spans) to create sub-traces
	// and associates sub-traces with it's parent trace.
	for _, s := range childrenSpansResult.Series {
		span, err := newSpanFromRow(&s)
		if err != nil {
			return nil, err
		}
		annotations, err := annotationsFromRow(&s)
		if err != nil {
			return nil, err
		}
		span.Annotations = *annotations
		trace, present := tracesCache[span.ID.Trace]
		if !present { // Root trace not added.
			return nil, errors.New("parent not found")
		} else { // Root trace already added so append a sub-trace.
			trace.Sub = append(trace.Sub, &Trace{Span: *span})
		}
	}
	for _, trace := range tracesCache {
		traces = append(traces, trace)
	}
	return traces, nil
}

func (in *InfluxDBStore) Close() {
	in.server.Close()
}

func (in *InfluxDBStore) createDBIfNotExists() error {
	// If no errors query execution was successfully - either DB was created or already exists.
	response, err := in.con.Query(influxDBClient.Query{
		Command: fmt.Sprintf("%s %s", "CREATE DATABASE IF NOT EXISTS", dbName),
	})
	if err != nil {
		return err
	}
	if response.Error() != nil {
		return response.Error()
	}
	return nil
}

func (in *InfluxDBStore) executeOneQuery(command string) (*influxDBClient.Result, error) {
	response, err := in.con.Query(influxDBClient.Query{
		Command:  command,
		Database: dbName,
	})
	if err != nil {
		return nil, err
	}
	if response.Error() != nil {
		return nil, response.Error()
	}

	// Expecting one result, since a single query is executed.
	if len(response.Results) != 1 {
		return nil, errors.New("unexpected number of results for an influxdb single query")
	}
	return &response.Results[0], nil
}

func (in *InfluxDBStore) findSpanPoint(ID SpanID) (*influxDBClient.Point, error) {
	q := fmt.Sprintf(`
		SELECT * FROM spans WHERE trace_id='%s' AND span_id='%s' AND parent_id='%s' GROUP BY *
	`, ID.Trace, ID.Span, ID.Parent)
	result, err := in.executeOneQuery(q)
	if err != nil {
		return nil, err
	}
	if len(result.Series) == 0 {
		return nil, nil
	}
	if len(result.Series) > 1 {
		return nil, errors.New("unexpected multiple series")
	}
	r := result.Series[0]
	if len(r.Values) == 0 {
		return nil, errors.New("unexpected empty series")
	}
	p := influxDBClient.Point{
		Fields: make(pointFields, 0),
	}
	fields := r.Values[0]
	for i, field := range fields {
		key := r.Columns[i]
		switch field.(type) {
		case string:
			// time field is set by InfluxDB not related to annotations.
			if key == "time" {
				t, err := time.Parse(time.RFC3339Nano, field.(string))
				if err != nil {
					return nil, err
				}
				p.Time = t
			}
			if key == "Name" {
				p.Fields[key] = field.(string)
			}
		case nil:
			continue
		default:
			return nil, fmt.Errorf("unexpected time type: %v", reflect.TypeOf(field))
		}
	}
	return &p, err
}

func (in *InfluxDBStore) init(server *influxDBServer.Server) error {
	in.server = server
	url, err := url.Parse(fmt.Sprintf("http://%s:%d", influxDBClient.DefaultHost, influxDBClient.DefaultPort))
	if err != nil {
		return err
	}
	con, err := influxDBClient.NewClient(influxDBClient.Config{URL: *url})
	if err != nil {
		return err
	}
	in.con = con
	if err := in.createDBIfNotExists(); err != nil {
		return err
	}
	in.tracesPerPage = defaultTracesPerPage
	return nil
}

func annotationsFromRow(r *influxDBModels.Row) (*Annotations, error) {
	// Actually an influxDBModels.Row represents a single InfluxDB serie.
	// r.Values[n] is a slice containing span's annotation values.
	var fields []interface{}
	if len(r.Values) == 1 {
		fields = r.Values[0]
	}

	// len(r.Values) might be greater than one - meaning there are
	// some spans to drop, see: InfluxDBStore.Collect(...).
	// If so last one is picked.
	if len(r.Values) > 1 {
		fields = r.Values[len(r.Values)-1]
	}
	annotations := make(Annotations, 0)

	// Iterates over fields which represent span's annotation values.
	for i, field := range fields {
		// It is safe to do column[0] (eg. 'Server.Request.Method')
		// matches fields[0] (eg. 'GET')
		key := r.Columns[i]
		var value []byte
		switch field.(type) {
		case string:
			value = []byte(field.(string))
		case nil:
		default:
			return nil, fmt.Errorf("unexpected field type: %v", reflect.TypeOf(field))
		}
		a := Annotation{
			Key:   key,
			Value: value,
		}
		annotations = append(annotations, a)
	}

	return &annotations, nil
}

// extendFields replaces existing items on dst from src.
func extendFields(dst, src pointFields) pointFields {
	for k, v := range src {
		if _, present := dst[k]; present {
			dst[k] = v
		}
	}
	return dst
}

func newSpanFromRow(r *influxDBModels.Row) (*Span, error) {
	span := &Span{}
	traceID, err := ParseID(r.Tags["trace_id"])
	if err != nil {
		return nil, err
	}
	spanID, err := ParseID(r.Tags["span_id"])
	if err != nil {
		return nil, err
	}
	parentID, err := ParseID(r.Tags["parent_id"])
	if err != nil {
		return nil, err
	}
	span.ID = SpanID{
		Trace:  ID(traceID),
		Span:   ID(spanID),
		Parent: ID(parentID),
	}
	return span, nil
}

type InfluxDBStoreConfig struct {
	Server    *influxDBServer.Config
	BuildInfo *influxDBServer.BuildInfo
}

func NewInfluxDBStore(config InfluxDBStoreConfig) (*InfluxDBStore, error) {
	//TODO: add Authentication.
	s, err := influxDBServer.NewServer(config.Server, config.BuildInfo)
	if err != nil {
		return nil, err
	}
	if err := s.Open(); err != nil {
		return nil, err
	}
	var in InfluxDBStore
	if err := in.init(s); err != nil {
		return nil, err
	}
	return &in, nil
}
<|endoftext|>


<commit_before>package main

import "github.com/dvrkps/dojo/initfuncs/aaa"

func init() {
	println("main/main: top init")
}

func main() {
	println("main")
	aaa.Aaa()
}

func init() {
	println("main/main: bottom init")
}
<commit_msg>initfuncs: call zzz<commit_after>package main

import (
	"github.com/dvrkps/dojo/initfuncs/aaa"
	"github.com/dvrkps/dojo/initfuncs/zzz"
)

func init() {
	println("main/main: top init")
}

func main() {
	println("main")
	aaa.Aaa()
	zzz.Zzz()
}

func init() {
	println("main/main: bottom init")
}
<|endoftext|>


<commit_before>// Package auth implements http request authentication
package auth

import (
	"errors"
	"fmt"
	//"github.com/MG-RAST/Shock/shock-server/auth/basic"
	"github.com/MG-RAST/Shock/shock-server/auth/globus"
	"github.com/MG-RAST/Shock/shock-server/auth/mgrast"
	"github.com/MG-RAST/Shock/shock-server/conf"
	e "github.com/MG-RAST/Shock/shock-server/errors"
	"github.com/MG-RAST/Shock/shock-server/logger"
	"github.com/MG-RAST/Shock/shock-server/user"
	"os"
)

// authCache is a
var authCache cache
var authMethods []func(string) (*user.User, error)

func Initialize() {
	authCache = cache{m: make(map[string]cacheValue)}
	authMethods = []func(string) (*user.User, error){}
	//if conf.Conf["basic_auth"] != "" {
	//	authMethods = append(authMethods, basic.Auth)
	//}
	if conf.Conf["globus_token_url"] != "" && conf.Conf["globus_profile_url"] != "" {
		authMethods = append(authMethods, globus.Auth)
	}
	if conf.Conf["mgrast_oauth_url"] != "" {
		authMethods = append(authMethods, mgrast.Auth)
	}
}

func Authenticate(header string) (u *user.User, err error) {
	if u = authCache.lookup(header); u != nil {
		return u, nil
	} else {
		for _, auth := range authMethods {
			if u, err := auth(header); u != nil && err == nil {
				if &u == nil {
					err_msg := fmt.Sprintf("ERROR: received pointer to empty user object from auth method for header = %v\n", header)
					logger.Error(err_msg)
					fmt.Fprintln(os.Stderr, err_msg)
					return nil, errors.New(e.InvalidAuth)
				} else {
					authCache.add(header, u)
					return u, nil
				}
			}
		}
	}
	return nil, errors.New(e.InvalidAuth)
}
<commit_msg>Removed unnecessary error checking.<commit_after>// Package auth implements http request authentication
package auth

import (
	"errors"
	//"github.com/MG-RAST/Shock/shock-server/auth/basic"
	"github.com/MG-RAST/Shock/shock-server/auth/globus"
	"github.com/MG-RAST/Shock/shock-server/auth/mgrast"
	"github.com/MG-RAST/Shock/shock-server/conf"
	e "github.com/MG-RAST/Shock/shock-server/errors"
	"github.com/MG-RAST/Shock/shock-server/user"
)

// authCache is a
var authCache cache
var authMethods []func(string) (*user.User, error)

func Initialize() {
	authCache = cache{m: make(map[string]cacheValue)}
	authMethods = []func(string) (*user.User, error){}
	//if conf.Conf["basic_auth"] != "" {
	//	authMethods = append(authMethods, basic.Auth)
	//}
	if conf.Conf["globus_token_url"] != "" && conf.Conf["globus_profile_url"] != "" {
		authMethods = append(authMethods, globus.Auth)
	}
	if conf.Conf["mgrast_oauth_url"] != "" {
		authMethods = append(authMethods, mgrast.Auth)
	}
}

func Authenticate(header string) (u *user.User, err error) {
	if u = authCache.lookup(header); u != nil {
		return u, nil
	} else {
		for _, auth := range authMethods {
			if u, err := auth(header); u != nil && err == nil {
				authCache.add(header, u)
				return u, nil
			}
		}
	}
	return nil, errors.New(e.InvalidAuth)
}
<|endoftext|>


<commit_before>package goline
/*
 *  Filename:    question.go
 *  Package:     goline
 *  Author:      Bryan Matsuo <bmatsuo@soe.ucsc.edu>
 *  Created:     Sat Aug 13 02:30:29 PDT 2011
 *  Description:
 */
import (
    "reflect"
    "strings"
    "bufio"
    "fmt"
    "os"
)

type ErrorPrecision struct {
    Wide, Thin interface{}
}

func (e ErrorPrecision) String() string {
    return fmt.Sprintf("Input out of destination range (%v -> %v)", e.Wide, e.Thin)
}

//  Prompt the user for text input. The result is stored in dest, which must
//  be a pointer to a native Go type (int, uint16, string, float32, ...).
//  Slice types are not currently supported. List input must be done with a
//  *string destination and post-processing.
func Ask(dest interface{}, msg string, config func(*Answer)) os.Error {
    if k := reflect.TypeOf(dest).Kind(); k != reflect.Ptr && k != reflect.Slice {
        return fmt.Errorf("Ask(...) requires a Ptr type, not %s", k.String())
    } else if k == reflect.Slice {
        return fmt.Errorf("Ask(...) can not currently assign to slices.")
    }

    var t Type
    switch dest.(type) {
    case *uint:
        t = Uint
    case *uint8:
        t = Uint
    case *uint16:
        t = Uint
    case *uint32:
        t = Uint
    case *uint64:
        t = Uint
    case *int:
        t = Int
    case *int8:
        t = Int
    case *int16:
        t = Int
    case *int32:
        t = Int
    case *int64:
        t = Int
    case *float32:
        t = Float
    case *float64:
        t = Float
    case *string:
        t = String
    default:
        fmt.Errorf("Unusable destination")
    }
    a := newAnswer(t)
    a.Question = msg
    config(a)

    if err := a.tryFirstAnswer(); err != nil {
        return err
    } else if a.val != nil {
        return nil
    }

    prompt := msg
    r := bufio.NewReader(os.Stdin)
    for {
        fmt.Print(strings.Trim(prompt, "\n") + a.defaultString())
        var resp []byte
        for cont := true; cont; {
            s, isPrefix, err := r.ReadLine()
            cont = isPrefix
            if err != nil {
                return err
            }
            resp = append(resp, s...)
        }
        if err := a.parse(string(resp)); err != nil {
            switch err.(type) {
            case RecoverableError:
                fmt.Printf("Error: %s\n", err.String())
                prompt = a.Responses[AskOnError]
                continue
            default:
                return err
            }
        }

        // Cast the result from a wide (e.g. 64bit) type to the desired type.
        // This should not fail under any normal circumstances, so failure
        // should break the loop.
        var errCast os.Error
        switch dest.(type) {
        case *uint:
            d := dest.(*uint)
            *(d) = uint(a.val.(uint64))
            if x := uint64(*(d)); x != a.val.(uint64) {
                errCast = ErrorPrecision{a.val.(uint64), x}
            }
        case *uint8:
            d := dest.(*uint8)
            *(d) = uint8(a.val.(uint64))
            if x := uint64(*(d)); x != a.val.(uint64) {
                errCast = ErrorPrecision{a.val.(uint64), x}
            }
        case *uint16:
            d := dest.(*uint16)
            *(d) = uint16(a.val.(uint64))
            if x := uint64(*(d)); x != a.val.(uint64) {
                errCast = ErrorPrecision{a.val.(uint64), x}
            }
        case *uint32:
            d := dest.(*uint32)
            *(d) = uint32(a.val.(uint64))
            if x := uint64(*(d)); x != a.val.(uint64) {
                errCast = ErrorPrecision{a.val.(uint64), x}
            }
        case *uint64:
            *(dest.(*uint64)) = a.val.(uint64)
        case *int:
            d := dest.(*int)
            *(d) = int(a.val.(int64))
            if x := int64(*(d)); x != a.val.(int64) {
                errCast = ErrorPrecision{a.val.(int64), x}
            }
        case *int8:
            d := dest.(*int8)
            *(d) = int8(a.val.(int64))
            if x := int64(*(d)); x != a.val.(int64) {
                errCast = ErrorPrecision{a.val.(int64), x}
            }
        case *int16:
            d := dest.(*int16)
            *(d) = int16(a.val.(int64))
            if x := int64(*(d)); x != a.val.(int64) {
                errCast = ErrorPrecision{a.val.(int64), x}
            }
        case *int32:
            d := dest.(*int32)
            *(d) = int32(a.val.(int64))
            if x := int64(*(d)); x != a.val.(int64) {
                errCast = ErrorPrecision{a.val.(int64), x}
            }
        case *int64:
            *(dest.(*int64)) = a.val.(int64)
        case *float32:
            d := dest.(*float32)
            *(d) = float32(a.val.(float64))
            if x := float64(*(d)); x != a.val.(float64) {
                errCast = ErrorPrecision{a.val.(float64), x}
            }
        case *float64:
            *(dest.(*float64)) = a.val.(float64)
        case *string:
            *(dest.(*string)) = a.val.(string)
        default:
            errCast = fmt.Errorf("Unexpected cast type")
        }
        if errCast != nil {
            return errCast
        }
        break
    }
    return nil
}
<commit_msg>Fix an error with FirstAnswer short-circuiting.<commit_after>package goline
/*
 *  Filename:    question.go
 *  Package:     goline
 *  Author:      Bryan Matsuo <bmatsuo@soe.ucsc.edu>
 *  Created:     Sat Aug 13 02:30:29 PDT 2011
 *  Description:
 */
import (
    "reflect"
    "strings"
    "bufio"
    "fmt"
    "os"
)

type ErrorPrecision struct {
    Wide, Thin interface{}
}

func (e ErrorPrecision) String() string {
    return fmt.Sprintf("Input out of destination range (%v -> %v)", e.Wide, e.Thin)
}

//  Prompt the user for text input. The result is stored in dest, which must
//  be a pointer to a native Go type (int, uint16, string, float32, ...).
//  Slice types are not currently supported. List input must be done with a
//  *string destination and post-processing.
func Ask(dest interface{}, msg string, config func(*Answer)) os.Error {
    if k := reflect.TypeOf(dest).Kind(); k != reflect.Ptr && k != reflect.Slice {
        return fmt.Errorf("Ask(...) requires a Ptr type, not %s", k.String())
    } else if k == reflect.Slice {
        return fmt.Errorf("Ask(...) can not currently assign to slices.")
    }

    var t Type
    switch dest.(type) {
    case *uint:
        t = Uint
    case *uint8:
        t = Uint
    case *uint16:
        t = Uint
    case *uint32:
        t = Uint
    case *uint64:
        t = Uint
    case *int:
        t = Int
    case *int8:
        t = Int
    case *int16:
        t = Int
    case *int32:
        t = Int
    case *int64:
        t = Int
    case *float32:
        t = Float
    case *float64:
        t = Float
    case *string:
        t = String
    default:
        fmt.Errorf("Unusable destination")
    }
    a := newAnswer(t)
    a.Question = msg
    config(a)

    if err := a.tryFirstAnswer(); err == nil && a.val != nil {
        return nil
    }

    prompt := msg
    r := bufio.NewReader(os.Stdin)
    for {
        fmt.Print(strings.Trim(prompt, "\n") + a.defaultString())
        var resp []byte
        for cont := true; cont; {
            s, isPrefix, err := r.ReadLine()
            cont = isPrefix
            if err != nil {
                return err
            }
            resp = append(resp, s...)
        }
        if err := a.parse(string(resp)); err != nil {
            switch err.(type) {
            case RecoverableError:
                fmt.Printf("Error: %s\n", err.String())
                prompt = a.Responses[AskOnError]
                continue
            default:
                return err
            }
        }

        // Cast the result from a wide (e.g. 64bit) type to the desired type.
        // This should not fail under any normal circumstances, so failure
        // should break the loop.
        var errCast os.Error
        switch dest.(type) {
        case *uint:
            d := dest.(*uint)
            *(d) = uint(a.val.(uint64))
            if x := uint64(*(d)); x != a.val.(uint64) {
                errCast = ErrorPrecision{a.val.(uint64), x}
            }
        case *uint8:
            d := dest.(*uint8)
            *(d) = uint8(a.val.(uint64))
            if x := uint64(*(d)); x != a.val.(uint64) {
                errCast = ErrorPrecision{a.val.(uint64), x}
            }
        case *uint16:
            d := dest.(*uint16)
            *(d) = uint16(a.val.(uint64))
            if x := uint64(*(d)); x != a.val.(uint64) {
                errCast = ErrorPrecision{a.val.(uint64), x}
            }
        case *uint32:
            d := dest.(*uint32)
            *(d) = uint32(a.val.(uint64))
            if x := uint64(*(d)); x != a.val.(uint64) {
                errCast = ErrorPrecision{a.val.(uint64), x}
            }
        case *uint64:
            *(dest.(*uint64)) = a.val.(uint64)
        case *int:
            d := dest.(*int)
            *(d) = int(a.val.(int64))
            if x := int64(*(d)); x != a.val.(int64) {
                errCast = ErrorPrecision{a.val.(int64), x}
            }
        case *int8:
            d := dest.(*int8)
            *(d) = int8(a.val.(int64))
            if x := int64(*(d)); x != a.val.(int64) {
                errCast = ErrorPrecision{a.val.(int64), x}
            }
        case *int16:
            d := dest.(*int16)
            *(d) = int16(a.val.(int64))
            if x := int64(*(d)); x != a.val.(int64) {
                errCast = ErrorPrecision{a.val.(int64), x}
            }
        case *int32:
            d := dest.(*int32)
            *(d) = int32(a.val.(int64))
            if x := int64(*(d)); x != a.val.(int64) {
                errCast = ErrorPrecision{a.val.(int64), x}
            }
        case *int64:
            *(dest.(*int64)) = a.val.(int64)
        case *float32:
            d := dest.(*float32)
            *(d) = float32(a.val.(float64))
            if x := float64(*(d)); x != a.val.(float64) {
                errCast = ErrorPrecision{a.val.(float64), x}
            }
        case *float64:
            *(dest.(*float64)) = a.val.(float64)
        case *string:
            *(dest.(*string)) = a.val.(string)
        default:
            errCast = fmt.Errorf("Unexpected cast type")
        }
        if errCast != nil {
            return errCast
        }
        break
    }
    return nil
}
<|endoftext|>