<commit_before>package epoch

import (
	"database/sql/driver"
	"encoding/json"
	"fmt"
	"strconv"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/dynamodb"
	"github.com/pkg/errors"
)

type Millis int64

const (
	Scale = int64(time.Millisecond)
)

var (
	timeFormats = []string{
		time.RFC3339,
		time.RFC1123,
	}
)

func Now() Millis {
	return Time(time.Now())
}

func (em Millis) Add(d time.Duration) Millis {
	return em + Millis(d/time.Millisecond)
}

func (em Millis) Time() time.Time {
	v := int64(em) * Scale
	return time.Unix(v/int64(time.Second), v%int64(time.Second))
}

func (em Millis) Format(layout string) string {
	return em.Time().In(PT).Format(layout)
}

func (em Millis) AttributeValue() *dynamodb.AttributeValue {
	return &dynamodb.AttributeValue{
		N: aws.String(strconv.FormatInt(int64(em), 10)),
	}
}

func (em Millis) Value() (driver.Value, error) {
	if em == 0 {
		return nil, nil
	}
	return em.Time(), nil
}

func (em *Millis) Scan(src interface{}) error {
	if src == nil {
		return nil
	}

	switch v := src.(type) {
	case time.Time:
		if t := Time(v); t > 0 {
			*em = t
		}
		return nil

	default:
		return fmt.Errorf("unhandled type, %#v", v)
	}
}

func (em Millis) Int64() int64 {
	return int64(em)
}

func (em Millis) MarshalJSON() ([]byte, error) {
	if em == 0 {
		return []byte("null"), nil
	}

	t := em.Time().In(PT)
	v := map[string]interface{}{
		"Date":  t.Format("1/2/2006"),
		"Time":  t.Format(time.Kitchen),
		"Value": em.Int64(),
		"date":  t.Format("1/2/2006"),
		"time":  t.Format(time.Kitchen),
		"value": em.Int64(),
		"ago":   (Now() - em).Ago(),
	}

	return json.Marshal(v)
}

type millisModel struct {
	Value int64
}

func (em *Millis) UnmarshalJSON(data []byte) error {
	if data[0] == '"' {
		v := ""
		err := json.Unmarshal(data, &v)
		if err != nil {
			return err
		}

		for _, layout := range timeFormats {
			t, err := time.Parse(layout, v)
			if err == nil {
				*em = Time(t)
				return nil
			}
		}

		return errors.New("invalid time format")

	} else if data[0] == '{' {
		v := millisModel{}
		err := json.Unmarshal(data, &v)
		if err != nil {
			return err
		}

		*em = Millis(v.Value)

	} else {
		v := int64(0)
		err := json.Unmarshal(data, &v)
		if err != nil {
			return err
		}

		*em = Millis(v)
	}

	return nil
}

func UnixNano(v int64) Millis {
	return Millis(v / Scale)
}

func Time(t time.Time) Millis {
	return UnixNano(t.UnixNano())
}

func (em Millis) Ago() string {
	if em < 0 {
		return "-"
	}

	secs := em / 1000
	if secs < 60 {
		return "just now"
	}

	minutes := secs / 60
	if minutes == 1 {
		return "1 minute ago"
	} else if minutes < 60 {
		return strconv.Itoa(int(minutes)) + " minutes ago"
	}

	hours := minutes / 60
	if hours == 1 {
		return "1 hour ago"
	} else if hours < 24 {
		return strconv.Itoa(int(hours)) + " hours ago"
	}

	days := hours / 24
	if days == 1 {
		return "1 day ago"
	} else if days < 30 {
		return strconv.Itoa(int(days)) + " days ago"
	}

	months := days / 30
	if months == 1 {
		return "1 month ago"
	} else if months < 12 {
		return strconv.Itoa(int(months)) + " months ago"
	}

	years := months / 12
	if years == 1 {
		return "1 year ago"
	}
	return strconv.Itoa(int(years)) + " years ago"
}
<commit_msg>- removing lower case references from time<commit_after>package epoch

import (
	"database/sql/driver"
	"encoding/json"
	"fmt"
	"strconv"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/dynamodb"
	"github.com/pkg/errors"
)

type Millis int64

const (
	Scale = int64(time.Millisecond)
)

var (
	timeFormats = []string{
		time.RFC3339,
		time.RFC1123,
	}
)

func Now() Millis {
	return Time(time.Now())
}

func (em Millis) Add(d time.Duration) Millis {
	return em + Millis(d/time.Millisecond)
}

func (em Millis) Time() time.Time {
	v := int64(em) * Scale
	return time.Unix(v/int64(time.Second), v%int64(time.Second))
}

func (em Millis) Format(layout string) string {
	return em.Time().In(PT).Format(layout)
}

func (em Millis) AttributeValue() *dynamodb.AttributeValue {
	return &dynamodb.AttributeValue{
		N: aws.String(strconv.FormatInt(int64(em), 10)),
	}
}

func (em Millis) Value() (driver.Value, error) {
	if em == 0 {
		return nil, nil
	}
	return em.Time(), nil
}

func (em *Millis) Scan(src interface{}) error {
	if src == nil {
		return nil
	}

	switch v := src.(type) {
	case time.Time:
		if t := Time(v); t > 0 {
			*em = t
		}
		return nil

	default:
		return fmt.Errorf("unhandled type, %#v", v)
	}
}

func (em Millis) Int64() int64 {
	return int64(em)
}

func (em Millis) MarshalJSON() ([]byte, error) {
	if em == 0 {
		return []byte("null"), nil
	}

	t := em.Time().In(PT)
	v := map[string]interface{}{
		"Date":  t.Format("1/2/2006"),
		"Time":  t.Format(time.Kitchen),
		"Value": em.Int64(),
		"Ago":   (Now() - em).Ago(),
	}

	return json.Marshal(v)
}

type millisModel struct {
	Value int64
}

func (em *Millis) UnmarshalJSON(data []byte) error {
	if data[0] == '"' {
		v := ""
		err := json.Unmarshal(data, &v)
		if err != nil {
			return err
		}

		for _, layout := range timeFormats {
			t, err := time.Parse(layout, v)
			if err == nil {
				*em = Time(t)
				return nil
			}
		}

		return errors.New("invalid time format")

	} else if data[0] == '{' {
		v := millisModel{}
		err := json.Unmarshal(data, &v)
		if err != nil {
			return err
		}

		*em = Millis(v.Value)

	} else {
		v := int64(0)
		err := json.Unmarshal(data, &v)
		if err != nil {
			return err
		}

		*em = Millis(v)
	}

	return nil
}

func UnixNano(v int64) Millis {
	return Millis(v / Scale)
}

func Time(t time.Time) Millis {
	return UnixNano(t.UnixNano())
}

func (em Millis) Ago() string {
	if em < 0 {
		return "-"
	}

	secs := em / 1000
	if secs < 60 {
		return "just now"
	}

	minutes := secs / 60
	if minutes == 1 {
		return "1 minute ago"
	} else if minutes < 60 {
		return strconv.Itoa(int(minutes)) + " minutes ago"
	}

	hours := minutes / 60
	if hours == 1 {
		return "1 hour ago"
	} else if hours < 24 {
		return strconv.Itoa(int(hours)) + " hours ago"
	}

	days := hours / 24
	if days == 1 {
		return "1 day ago"
	} else if days < 30 {
		return strconv.Itoa(int(days)) + " days ago"
	}

	months := days / 30
	if months == 1 {
		return "1 month ago"
	} else if months < 12 {
		return strconv.Itoa(int(months)) + " months ago"
	}

	years := months / 12
	if years == 1 {
		return "1 year ago"
	}
	return strconv.Itoa(int(years)) + " years ago"
}
<|endoftext|>


<commit_before>// The erisdb package contains tendermint-specific services that goes with the
// server.
package erisdb

import (
	"github.com/eris-ltd/eris-db/Godeps/_workspace/src/github.com/tendermint/log15"
	. "github.com/eris-ltd/eris-db/Godeps/_workspace/src/github.com/tendermint/tendermint/common"
	cfg "github.com/eris-ltd/eris-db/Godeps/_workspace/src/github.com/tendermint/tendermint/config"
	tmcfg "github.com/eris-ltd/eris-db/Godeps/_workspace/src/github.com/tendermint/tendermint/config/tendermint"
	"github.com/eris-ltd/eris-db/Godeps/_workspace/src/github.com/tendermint/tendermint/node"
	"github.com/eris-ltd/eris-db/Godeps/_workspace/src/github.com/tendermint/tendermint/p2p"
	ep "github.com/eris-ltd/eris-db/erisdb/pipe"
	"github.com/eris-ltd/eris-db/server"
	"path"
)

const ERISDB_VERSION = "0.10.3"
const TENDERMINT_VERSION = "0.5.0"

var log = log15.New("module", "eris/erisdb_server")
var tmConfig cfg.Config

// This function returns a properly configured ErisDb server process with a running
// tendermint node attached to it. To start listening for incoming requests, call
// 'Start()' on the process. Make sure to register any start event listeners before
// that.
func ServeErisDB(workDir string) (*server.ServeProcess, error) {
	log.Info("ErisDB Serve initializing.")
	errEns := EnsureDir(workDir)

	if errEns != nil {
		return nil, errEns
	}

	var sConf *server.ServerConfig

	sConfPath := path.Join(workDir, "server_conf.toml")
	if !FileExists(sConfPath) {
		log.Info("No server configuration, using default.")
		log.Info("Writing to: " + sConfPath)
		sConf = server.DefaultServerConfig()
		errW := server.WriteServerConfig(sConfPath, sConf)
		if errW != nil {
			panic(errW)
		}
	} else {
		var errRSC error
		sConf, errRSC = server.ReadServerConfig(sConfPath)
		if errRSC != nil {
			log.Error("Server config file error.", "error", errRSC.Error())
		}
	}

	// Get tendermint configuration
	tmConfig = tmcfg.GetConfig(workDir)
	tmConfig.Set("version", TENDERMINT_VERSION)
	cfg.ApplyConfig(tmConfig) // Notify modules of new config

	// Set the node up.
	nodeRd := make(chan struct{})
	nd := node.NewNode()
	// Load the supporting objects.
	pipe := ep.NewPipe(nd)
	codec := &TCodec{}
	evtSubs := NewEventSubscriptions(pipe.Events())
	// The services.
	tmwss := NewErisDbWsService(codec, pipe)
	tmjs := NewErisDbJsonService(codec, pipe, evtSubs)
	// The servers.
	jsonServer := NewJsonRpcServer(tmjs)
	restServer := NewRestServer(codec, pipe, evtSubs)
	wsServer := server.NewWebSocketServer(sConf.WebSocket.MaxWebSocketSessions, tmwss)
	// Create a server process.
	proc := server.NewServeProcess(sConf, jsonServer, restServer, wsServer)

	stopChan := proc.StopEventChannel()
	go startNode(nd, nodeRd, stopChan)
	<-nodeRd
	return proc, nil
}

// Private. Create a new node
func startNode(nd *node.Node, ready chan struct{}, shutDown <-chan struct{}) {
	laddr := tmConfig.GetString("node_laddr")
	if laddr != "" {
		l := p2p.NewDefaultListener("tcp", laddr, false)
		nd.AddListener(l)
	}

	nd.Start()

	// If seedNode is provided by config, dial out.

	if len(tmConfig.GetString("seeds")) > 0 {
		nd.DialSeed()
	}

	ready <- struct{}{}
	// Block until everything is shut down.
	<-shutDown
	nd.Stop()
}<commit_msg>version bump<commit_after>// The erisdb package contains tendermint-specific services that goes with the
// server.
package erisdb

import (
	"github.com/eris-ltd/eris-db/Godeps/_workspace/src/github.com/tendermint/log15"
	. "github.com/eris-ltd/eris-db/Godeps/_workspace/src/github.com/tendermint/tendermint/common"
	cfg "github.com/eris-ltd/eris-db/Godeps/_workspace/src/github.com/tendermint/tendermint/config"
	tmcfg "github.com/eris-ltd/eris-db/Godeps/_workspace/src/github.com/tendermint/tendermint/config/tendermint"
	"github.com/eris-ltd/eris-db/Godeps/_workspace/src/github.com/tendermint/tendermint/node"
	"github.com/eris-ltd/eris-db/Godeps/_workspace/src/github.com/tendermint/tendermint/p2p"
	ep "github.com/eris-ltd/eris-db/erisdb/pipe"
	"github.com/eris-ltd/eris-db/server"
	"path"
)

const ERISDB_VERSION = "0.10.4"
const TENDERMINT_VERSION = "0.5.0"

var log = log15.New("module", "eris/erisdb_server")
var tmConfig cfg.Config

// This function returns a properly configured ErisDb server process with a running
// tendermint node attached to it. To start listening for incoming requests, call
// 'Start()' on the process. Make sure to register any start event listeners before
// that.
func ServeErisDB(workDir string) (*server.ServeProcess, error) {
	log.Info("ErisDB Serve initializing.")
	errEns := EnsureDir(workDir)

	if errEns != nil {
		return nil, errEns
	}

	var sConf *server.ServerConfig

	sConfPath := path.Join(workDir, "server_conf.toml")
	if !FileExists(sConfPath) {
		log.Info("No server configuration, using default.")
		log.Info("Writing to: " + sConfPath)
		sConf = server.DefaultServerConfig()
		errW := server.WriteServerConfig(sConfPath, sConf)
		if errW != nil {
			panic(errW)
		}
	} else {
		var errRSC error
		sConf, errRSC = server.ReadServerConfig(sConfPath)
		if errRSC != nil {
			log.Error("Server config file error.", "error", errRSC.Error())
		}
	}

	// Get tendermint configuration
	tmConfig = tmcfg.GetConfig(workDir)
	tmConfig.Set("version", TENDERMINT_VERSION)
	cfg.ApplyConfig(tmConfig) // Notify modules of new config

	// Set the node up.
	nodeRd := make(chan struct{})
	nd := node.NewNode()
	// Load the supporting objects.
	pipe := ep.NewPipe(nd)
	codec := &TCodec{}
	evtSubs := NewEventSubscriptions(pipe.Events())
	// The services.
	tmwss := NewErisDbWsService(codec, pipe)
	tmjs := NewErisDbJsonService(codec, pipe, evtSubs)
	// The servers.
	jsonServer := NewJsonRpcServer(tmjs)
	restServer := NewRestServer(codec, pipe, evtSubs)
	wsServer := server.NewWebSocketServer(sConf.WebSocket.MaxWebSocketSessions, tmwss)
	// Create a server process.
	proc := server.NewServeProcess(sConf, jsonServer, restServer, wsServer)

	stopChan := proc.StopEventChannel()
	go startNode(nd, nodeRd, stopChan)
	<-nodeRd
	return proc, nil
}

// Private. Create a new node
func startNode(nd *node.Node, ready chan struct{}, shutDown <-chan struct{}) {
	laddr := tmConfig.GetString("node_laddr")
	if laddr != "" {
		l := p2p.NewDefaultListener("tcp", laddr, false)
		nd.AddListener(l)
	}

	nd.Start()

	// If seedNode is provided by config, dial out.

	if len(tmConfig.GetString("seeds")) > 0 {
		nd.DialSeed()
	}

	ready <- struct{}{}
	// Block until everything is shut down.
	<-shutDown
	nd.Stop()
}<|endoftext|>


<commit_before>package errors

import "fmt"

// Code references to the kind of error that was returned
type Code string

// StatusMessage in order to make it simple
// to keep a list of the used codes and default messages
type StatusMessage struct {
	Message    string
	StatusCode int
}

// NewErrorStatusMessage build a message with a status code
func NewErrorStatusMessage(message string, status int) StatusMessage {
	return StatusMessage{Message: message, StatusCode: status}
}

const (
	// ErrorCodeResourceNotFound error when the requested resource was not found
	ErrorCodeResourceNotFound = "resource_not_exist"
	// ErrorCodeInvalidArgs error when the request data is incorrect or incomplete
	ErrorCodeInvalidArgs = "invalid_args"
	// ErrorCodeBadRequest error when the needed input is not provided
	ErrorCodeBadRequest = "bad_request"
	//ErrorCodeResourceAlreadyExists when a resource already exists
	ErrorCodeResourceAlreadyExists = "resource_already_exist"
	// ErrorCodeUnknownIssue when the issue is unknown
	ErrorCodeUnknownIssue = "unknown_issue"
	// ErrorCodePermissionDenied when user dont have necessary permissions
	ErrorCodePermissionDenied = "permission_denied"
	// ErrorCodeResourceStateConflict when the resource is in another state and generates a conflict
	ErrorCodeResourceStateConflict = "resource_state_conflict"
	// ErrorCodeNotImplemented when some method is not implemented
	ErrorCodeNotImplemented = "not_implemented"
	// ErrorCodeUnauthorized when user is not authorized
	ErrorCodeUnauthorized = "unauthorized"
	// ErrorCodeNotFound when is not found
	ErrorCodeNotFound = "not_found"
	// ErrorCodeElasticSearchError when using a elastic search and it returned an error
	ErrorCodeElasticSearchError = "elasticsearch_error"
)

var (
	// ErrorMessageList general default messages for all error codes
	ErrorMessageList = map[Code]StatusMessage{
		ErrorCodeResourceNotFound:      NewErrorStatusMessage("Resource not found", 404),
		ErrorCodeInvalidArgs:           NewErrorStatusMessage("Invalid parameters were passed.", 400),
		ErrorCodeBadRequest:            NewErrorStatusMessage("Required data not valid.", 400),
		ErrorCodeResourceAlreadyExists: NewErrorStatusMessage("The posted resource already existed.", 400),
		ErrorCodeUnknownIssue:          NewErrorStatusMessage("Unknown issue was caught and message was not specified.", 500),
		ErrorCodePermissionDenied:      NewErrorStatusMessage("Current user has no permission to perform the action.", 403),
		ErrorCodeResourceStateConflict: NewErrorStatusMessage("The posted resource already existed.", 409),
		ErrorCodeNotImplemented:        NewErrorStatusMessage("Method not implemented", 501),
		ErrorCodeElasticSearchError:    NewErrorStatusMessage("Elastic search error.", 400),
	}
)

// AddError add error
func AddError(code Code, message string, status int) {
	ErrorMessageList[code] = NewErrorStatusMessage(message, status)
}

// AlaudaError common error struct used in the whole application
type AlaudaError struct {
	Source     string                `json:"source"`
	Message    string                `json:"message"`
	Code       Code                  `json:"code"`
	Fields     []map[string][]string `json:"fields,omitempty"`
	StatusCode int                   `json:"-"`
}

// New Constructor function for the error structure
func New(source string, code Code) *AlaudaError {
	var (
		message = "Error not described"
		status  = 999
	)
	if val, ok := ErrorMessageList[code]; ok {
		message = val.Message
		status = val.StatusCode
	}
	return &AlaudaError{
		Source:     source,
		Code:       code,
		Message:    message,
		StatusCode: status,
	}
}

// NewCommon return an error from a common error
func NewCommon(source string, err error) *AlaudaError {
	if err == nil {
		return nil
	}
	if alErr, ok := err.(*AlaudaError); ok {
		return alErr
	}
	return &AlaudaError{
		Source:     source,
		Code:       ErrorCodeUnknownIssue,
		Message:    err.Error(),
		StatusCode: ErrorMessageList[ErrorCodeUnknownIssue].StatusCode,
	}
}

// Error satisfies the error interface
func (h *AlaudaError) Error() string {
	return fmt.Sprintf("%s: %s", h.Code, h.Message)
}

// SetMessage sets a message using the given format and parameters
// returns itself
func (h *AlaudaError) SetMessage(format string, args ...interface{}) *AlaudaError {
	h.Message = fmt.Sprintf(format, args...)
	return h
}

// SetCodeString sets a code as a string for custom codes
func (h *AlaudaError) SetCodeString(code string) *AlaudaError {
	return h.SetCode(Code(code))
}

// SetCode sets a code and returns itself for chaining calls
func (h *AlaudaError) SetCode(code Code) *AlaudaError {
	h.Code = code
	return h
}

// SetSource sets the source and returns itself for chaining calls
func (h *AlaudaError) SetSource(source string) *AlaudaError {
	h.Source = source
	return h
}

// AddFieldError adds a field error and return itself
func (h *AlaudaError) AddFieldError(field string, message ...string) *AlaudaError {
	h.initializeFields()
	if value, ok := h.Fields[0][field]; ok {
		value = append(value, message...)
		return h
	}
	h.Fields[0][field] = message
	return h
}

// AddFieldsFromError will add the field errors of the given error if any
// will also add in the current error using a formated string as base
func (h *AlaudaError) AddFieldsFromError(format string, index int, err *AlaudaError) *AlaudaError {
	if err == nil || err.Fields == nil || len(err.Fields) == 0 || len(err.Fields[0]) == 0 {
		return h
	}
	for k, v := range err.Fields[0] {
		h.AddFieldError(fmt.Sprintf(format, index, k), v...)
	}
	return h
}

func (h *AlaudaError) initializeFields() {
	if h.Fields == nil {
		h.Fields = []map[string][]string{
			map[string][]string{},
		}
	}
}

// GetError will create an error or return if already instatiated
func GetError(source string, err *AlaudaError, code Code) *AlaudaError {
	if err == nil {
		err = New(source, code)
	}
	return err
}
<commit_msg>feature: Adding database error standard error code<commit_after>package errors

import "fmt"

// Code references to the kind of error that was returned
type Code string

// StatusMessage in order to make it simple
// to keep a list of the used codes and default messages
type StatusMessage struct {
	Message    string
	StatusCode int
}

// NewErrorStatusMessage build a message with a status code
func NewErrorStatusMessage(message string, status int) StatusMessage {
	return StatusMessage{Message: message, StatusCode: status}
}

const (
	// ErrorCodeResourceNotFound error when the requested resource was not found
	ErrorCodeResourceNotFound = "resource_not_exist"
	// ErrorCodeInvalidArgs error when the request data is incorrect or incomplete
	ErrorCodeInvalidArgs = "invalid_args"
	// ErrorCodeBadRequest error when the needed input is not provided
	ErrorCodeBadRequest = "bad_request"
	//ErrorCodeResourceAlreadyExists when a resource already exists
	ErrorCodeResourceAlreadyExists = "resource_already_exist"
	// ErrorCodeUnknownIssue when the issue is unknown
	ErrorCodeUnknownIssue = "unknown_issue"
	// ErrorCodePermissionDenied when user dont have necessary permissions
	ErrorCodePermissionDenied = "permission_denied"
	// ErrorCodeResourceStateConflict when the resource is in another state and generates a conflict
	ErrorCodeResourceStateConflict = "resource_state_conflict"
	// ErrorCodeNotImplemented when some method is not implemented
	ErrorCodeNotImplemented = "not_implemented"
	// ErrorCodeUnauthorized when user is not authorized
	ErrorCodeUnauthorized = "unauthorized"
	// ErrorCodeNotFound when is not found
	ErrorCodeNotFound = "not_found"
	// ErrorCodeElasticSearchError when using a elastic search and it returned an error
	ErrorCodeElasticSearchError = "elasticsearch_error"
	// ErrorCodeDatabaseError when using database commands and it returned an error
	ErrorCodeDatabaseError = "database_error"
)

var (
	// ErrorMessageList general default messages for all error codes
	ErrorMessageList = map[Code]StatusMessage{
		ErrorCodeResourceNotFound:      NewErrorStatusMessage("Resource not found", 404),
		ErrorCodeInvalidArgs:           NewErrorStatusMessage("Invalid parameters were passed.", 400),
		ErrorCodeBadRequest:            NewErrorStatusMessage("Required data not valid.", 400),
		ErrorCodeResourceAlreadyExists: NewErrorStatusMessage("The posted resource already existed.", 400),
		ErrorCodeUnknownIssue:          NewErrorStatusMessage("Unknown issue was caught and message was not specified.", 500),
		ErrorCodePermissionDenied:      NewErrorStatusMessage("Current user has no permission to perform the action.", 403),
		ErrorCodeResourceStateConflict: NewErrorStatusMessage("The posted resource already existed.", 409),
		ErrorCodeNotImplemented:        NewErrorStatusMessage("Method not implemented", 501),
		ErrorCodeElasticSearchError:    NewErrorStatusMessage("Elastic search error.", 500),
		ErrorCodeDatabaseError:         NewErrorStatusMessage("Database error.", 500),
	}
)

// AddError add error
func AddError(code Code, message string, status int) {
	ErrorMessageList[code] = NewErrorStatusMessage(message, status)
}

// AlaudaError common error struct used in the whole application
type AlaudaError struct {
	Source     string                `json:"source"`
	Message    string                `json:"message"`
	Code       Code                  `json:"code"`
	Fields     []map[string][]string `json:"fields,omitempty"`
	StatusCode int                   `json:"-"`
}

// New Constructor function for the error structure
func New(source string, code Code) *AlaudaError {
	var (
		message = "Error not described"
		status  = 999
	)
	if val, ok := ErrorMessageList[code]; ok {
		message = val.Message
		status = val.StatusCode
	}
	return &AlaudaError{
		Source:     source,
		Code:       code,
		Message:    message,
		StatusCode: status,
	}
}

// NewCommon return an error from a common error
func NewCommon(source string, err error) *AlaudaError {
	if err == nil {
		return nil
	}
	if alErr, ok := err.(*AlaudaError); ok {
		return alErr
	}
	return &AlaudaError{
		Source:     source,
		Code:       ErrorCodeUnknownIssue,
		Message:    err.Error(),
		StatusCode: ErrorMessageList[ErrorCodeUnknownIssue].StatusCode,
	}
}

// Error satisfies the error interface
func (h *AlaudaError) Error() string {
	return fmt.Sprintf("%s: %s", h.Code, h.Message)
}

// SetMessage sets a message using the given format and parameters
// returns itself
func (h *AlaudaError) SetMessage(format string, args ...interface{}) *AlaudaError {
	h.Message = fmt.Sprintf(format, args...)
	return h
}

// SetCodeString sets a code as a string for custom codes
func (h *AlaudaError) SetCodeString(code string) *AlaudaError {
	return h.SetCode(Code(code))
}

// SetCode sets a code and returns itself for chaining calls
func (h *AlaudaError) SetCode(code Code) *AlaudaError {
	h.Code = code
	return h
}

// SetSource sets the source and returns itself for chaining calls
func (h *AlaudaError) SetSource(source string) *AlaudaError {
	h.Source = source
	return h
}

// AddFieldError adds a field error and return itself
func (h *AlaudaError) AddFieldError(field string, message ...string) *AlaudaError {
	h.initializeFields()
	if value, ok := h.Fields[0][field]; ok {
		value = append(value, message...)
		return h
	}
	h.Fields[0][field] = message
	return h
}

// AddFieldsFromError will add the field errors of the given error if any
// will also add in the current error using a formated string as base
func (h *AlaudaError) AddFieldsFromError(format string, index int, err *AlaudaError) *AlaudaError {
	if err == nil || err.Fields == nil || len(err.Fields) == 0 || len(err.Fields[0]) == 0 {
		return h
	}
	for k, v := range err.Fields[0] {
		h.AddFieldError(fmt.Sprintf(format, index, k), v...)
	}
	return h
}

func (h *AlaudaError) initializeFields() {
	if h.Fields == nil {
		h.Fields = []map[string][]string{
			map[string][]string{},
		}
	}
}

// GetError will create an error or return if already instatiated
func GetError(source string, err *AlaudaError, code Code) *AlaudaError {
	if err == nil {
		err = New(source, code)
	}
	return err
}
<|endoftext|>


<commit_before>package pipe

import (
	"crypto/sha1"
	"fmt"
	"io"
	"os"
	_ "testing"
)

func dump(filters ...Filter) {
	fmt.Println("-------")
	Print(filters...)
}

func ExampleEmpty() {
	Print()
	// Output:
}

func ExampleSingle() {
	Print(Echo("foo"))
	// Output:
	// foo
}

func ExampleMultiple() {
	Print(Echo("foo"), Echo("bar"))
	// Output:
	// foo
	// bar
}

func ExampleSequence() {
	Print(Sequence())
	fmt.Println("---")
	Print(Sequence(Echo("1 of 1")))
	Print(Sequence(Echo("1 of 2"), Echo("2 of 2")))
	// Output:
	// ---
	// 1 of 1
	// 1 of 2
	// 2 of 2
}

func ExampleEach() {
	for s := range Each(Numbers(2, 4)) {
		fmt.Print(s)
	}
	// Output:
	// 234
}

func ExampleIf() {
	Print(Numbers(1, 12), If(func(s string) bool { return len(s) > 1 }))
	// Output:
	// 10
	// 11
	// 12
}

func ExampleGrep() {
	Print(Numbers(1, 12), Grep(".."))
	// Output:
	// 10
	// 11
	// 12
}

func ExampleGrepNot() {
	Print(Numbers(1, 12), GrepNot("^.$"))
	// Output:
	// 10
	// 11
	// 12
}

func ExampleUniq() {
	Print(Echo("a", "b", "b", "c"), Uniq)
	// Output:
	// a
	// b
	// c
}

func ExampleUniqWithCount() {
	Print(Echo("a", "b", "b", "c"), UniqWithCount)
	// Output:
	// 1 a
	// 2 b
	// 1 c
}

func ExampleParallel() {
	Print(
		Numbers(1, 3),
		Parallel(4, func(s string, out chan<- string) { out <- s }),
		Sort(),
	)
	// Output:
	// 1
	// 2
	// 3
}

func ExampleReplaceMatch() {
	Print(Numbers(1, 5), ReplaceMatch("(3)", "$1$1"))
	// Output:
	// 1
	// 2
	// 33
	// 4
	// 5
}

func ExampleDeleteMatch() {
	Print(Numbers(1, 5), DeleteMatch("[24]"))
	// Output:
	// 1
	//
	// 3
	//
	// 5
}

func sortData() Filter {
	return Echo(
		"8 1",
		"8 3 x",
		"8 3 w",
		"8 2",
		"4 5",
		"9 3",
		"12 13",
		"12 5",
	)
}

func ExampleSort() {
	Print(sortData(), Sort())
	// Output:
	// 12 13
	// 12 5
	// 4 5
	// 8 1
	// 8 2
	// 8 3 w
	// 8 3 x
	// 9 3
}

func ExampleSort_TextCol() {
	Print(sortData(), Sort(Text(2)))
	// Output:
	// 8 1
	// 12 13
	// 8 2
	// 8 3 w
	// 8 3 x
	// 9 3
	// 12 5
	// 4 5
}

func ExampleSort_TwoText() {
	Print(sortData(), Sort(Text(1), Text(2)))
	// Output:
	// 12 13
	// 12 5
	// 4 5
	// 8 1
	// 8 2
	// 8 3 w
	// 8 3 x
	// 9 3
}

func ExampleSort_TwoNum() {
	Print(sortData(), Sort(Num(1), Num(2)))
	// Output:
	// 4 5
	// 8 1
	// 8 2
	// 8 3 w
	// 8 3 x
	// 9 3
	// 12 5
	// 12 13
}

func ExampleSort_Mix() {
	Print(sortData(), Sort(Text(1), Num(2)))
	// Output:
	// 12 5
	// 12 13
	// 4 5
	// 8 1
	// 8 2
	// 8 3 w
	// 8 3 x
	// 9 3
}

func ExampleSort_Rev() {
	Print(sortData(), Sort(Rev(Num(1)), Num(2)))
	// Output:
	// 12 5
	// 12 13
	// 9 3
	// 8 1
	// 8 2
	// 8 3 w
	// 8 3 x
	// 4 5
}

func ExampleReverse() {
	Print(Echo("a", "b"), Reverse)
	// Output:
	// b
	// a
}

func ExampleFirst() {
	Print(Numbers(1, 10), First(3))
	// Output:
	// 1
	// 2
	// 3
}

func ExampleLast() {
	Print(Numbers(1, 10), Last(2))
	// Output:
	// 9
	// 10
}

func ExampleDropFirst() {
	Print(Numbers(1, 10), DropFirst(8))
	// Output:
	// 9
	// 10
}

func ExampleDropLast() {
	Print(Numbers(1, 10), DropLast(8))
	// Output:
	// 1
	// 2
}

func ExampleNumberLines() {
	Print(Echo("a", "b"), NumberLines)
	// Output:
	//     1 a
	//     2 b
}

func ExampleCut() {
	Print(Echo("hello", "world."), Cut(2, 4))
	// Output:
	// llo
	// rld
}

func ExmapleSelect() {
	Print(Echo("hello world"), Select(2, 0, 1))
	// Output:
	// world hello world hello
}

func ExampleFind() {
	Print(Find(FILES, "."))
	Print(Find(DIRS, "."))
}

func ExampleCat() {
	Print(Cat("pipe_test.go"))
}

func ExampleSystem() {
	Print(
		System("find", ".", "-type", "f", "-print"),
		Grep(`^\./pipe`),
		Sort(),
	)
	// TODO: Remove output checking if it becomes fragile.
	// Output:
	// ./pipe
	// ./pipe.go
	// ./pipe_test.go
}

func ExampleMix() {
	dbl := func(arg Arg) {
		for s := range arg.in {
			arg.out <- s
			arg.out <- s
		}
	}

	Print(Numbers(1, 100),
		Grep("3"),
		GrepNot("7"),
		dbl,
		Uniq,
		ReplaceMatch("^(.)$", "x$1"),
		Sort(),
		ReplaceMatch("^(.)", "$1 "),
		dbl,
		DeleteMatch(" .$"),
		UniqWithCount,
		Sort(Num(1)),
		Reverse,
	)
	// Output:
	// 18 3
	// 2 x
	// 2 9
	// 2 8
	// 2 6
	// 2 5
	// 2 4
	// 2 2
	// 2 1
}

func ExampleHash() {
	hash := func(f string, out chan<- string) {
		file, err := os.Open(f)
		if err != nil {
			fmt.Fprintln(os.Stderr, err)
			return
		}
		hasher := sha1.New()
		_, err = io.Copy(hasher, file)
		file.Close()
		if err != nil {
			fmt.Fprintln(os.Stderr, err)
			return
		}
		out <- fmt.Sprintf("%x %s", hasher.Sum(nil), f)
	}

	Print(
		Find(FILES, "/home/sanjay/tmp"),
		Grep("/tmp/x"),
		GrepNot("/sub2/"),
		Parallel(4, hash),
		ReplaceMatch(" /home/sanjay/", " HOME/"))

	Print(
		Find(FILES, "/home/sanjay/tmp/y"),
		GrepNot(`/home/sanjay/(\.Trash|Library)/`),
		Parallel(4, hash),
		Sort(Text(2)),
	)

	Print(
		System("find", "/home/sanjay/tmp/y", "-type", "f", "-print"),
		Parallel(4, hash),
		Sort(Text(2)),
	)

}
<commit_msg>more smaller tests<commit_after>package pipe

import (
	"crypto/sha1"
	"fmt"
	"io"
	"os"
	_ "testing"
)

func dump(filters ...Filter) {
	fmt.Println("-------")
	Print(filters...)
}

func ExampleEmpty() {
	Print()
	// Output:
}

func ExampleSingle() {
	Print(Echo("foo"))
	// Output:
	// foo
}

func ExampleMultiple() {
	Print(Echo("foo"), Echo("bar"))
	// Output:
	// foo
	// bar
}

func ExampleSequence() {
	Print(Sequence())
	fmt.Println("---")
	Print(Sequence(Echo("1 of 1")))
	Print(Sequence(Echo("1 of 2"), Echo("2 of 2")))
	// Output:
	// ---
	// 1 of 1
	// 1 of 2
	// 2 of 2
}

func ExampleEach() {
	for s := range Each(Numbers(2, 4)) {
		fmt.Print(s)
	}
	// Output:
	// 234
}

func ExampleIf() {
	Print(Numbers(1, 12), If(func(s string) bool { return len(s) > 1 }))
	// Output:
	// 10
	// 11
	// 12
}

func ExampleGrep() {
	Print(Numbers(1, 12), Grep(".."))
	// Output:
	// 10
	// 11
	// 12
}

func ExampleGrepNot() {
	Print(Numbers(1, 12), GrepNot("^.$"))
	// Output:
	// 10
	// 11
	// 12
}

func ExampleUniq() {
	Print(Echo("a", "b", "b", "c"), Uniq)
	// Output:
	// a
	// b
	// c
}

func ExampleUniqWithCount() {
	Print(Echo("a", "b", "b", "c"), UniqWithCount)
	// Output:
	// 1 a
	// 2 b
	// 1 c
}

func ExampleParallel() {
	Print(
		Numbers(1, 3),
		Parallel(4, func(s string, out chan<- string) { out <- s }),
		Sort(),
	)
	// Output:
	// 1
	// 2
	// 3
}

func ExampleReplaceMatch() {
	Print(Numbers(1, 5), ReplaceMatch("(3)", "$1$1"))
	// Output:
	// 1
	// 2
	// 33
	// 4
	// 5
}

func ExampleDeleteMatch() {
	Print(Numbers(1, 5), DeleteMatch("[24]"))
	// Output:
	// 1
	//
	// 3
	//
	// 5
}

func sortData() Filter {
	return Echo(
		"8 1",
		"8 3 x",
		"8 3 w",
		"8 2",
		"4 5",
		"9 3",
		"12 13",
		"12 5",
	)
}

func ExampleSort() {
	Print(sortData(), Sort())
	// Output:
	// 12 13
	// 12 5
	// 4 5
	// 8 1
	// 8 2
	// 8 3 w
	// 8 3 x
	// 9 3
}

func ExampleSort_TextCol() {
	Print(sortData(), Sort(Text(2)))
	// Output:
	// 8 1
	// 12 13
	// 8 2
	// 8 3 w
	// 8 3 x
	// 9 3
	// 12 5
	// 4 5
}

func ExampleSort_TwoText() {
	Print(sortData(), Sort(Text(1), Text(2)))
	// Output:
	// 12 13
	// 12 5
	// 4 5
	// 8 1
	// 8 2
	// 8 3 w
	// 8 3 x
	// 9 3
}

func ExampleSort_TwoNum() {
	Print(sortData(), Sort(Num(1), Num(2)))
	// Output:
	// 4 5
	// 8 1
	// 8 2
	// 8 3 w
	// 8 3 x
	// 9 3
	// 12 5
	// 12 13
}

func ExampleSort_Mix() {
	Print(sortData(), Sort(Text(1), Num(2)))
	// Output:
	// 12 5
	// 12 13
	// 4 5
	// 8 1
	// 8 2
	// 8 3 w
	// 8 3 x
	// 9 3
}

func ExampleSort_Rev() {
	Print(sortData(), Sort(Rev(Num(1)), Num(2)))
	// Output:
	// 12 5
	// 12 13
	// 9 3
	// 8 1
	// 8 2
	// 8 3 w
	// 8 3 x
	// 4 5
}

func ExampleReverse() {
	Print(Echo("a", "b"), Reverse)
	// Output:
	// b
	// a
}

func ExampleFirst() {
	Print(Numbers(1, 10), First(3))
	// Output:
	// 1
	// 2
	// 3
}

func ExampleLast() {
	Print(Numbers(1, 10), Last(2))
	// Output:
	// 9
	// 10
}

func ExampleDropFirst() {
	Print(Numbers(1, 10), DropFirst(8))
	// Output:
	// 9
	// 10
}

func ExampleDropLast() {
	Print(Numbers(1, 10), DropLast(8))
	// Output:
	// 1
	// 2
}

func ExampleNumberLines() {
	Print(Echo("a", "b"), NumberLines)
	// Output:
	//     1 a
	//     2 b
}

func ExampleCut() {
	Print(Echo("hello", "world."), Cut(2, 4))
	// Output:
	// llo
	// rld
}

func ExmapleSelect() {
	Print(Echo("hello world"), Select(2, 0, 1))
	// Output:
	// world hello world hello
}

func ExampleFind() {
	Print(Find(DIRS, "."), GrepNot("git"), Echo("---"))
	Print(Find(FILES, "."), Grep("pipe"))
	// Output:
	// .
	// ---
	// pipe.go
	// pipe_test.go
}

func ExampleCat() {
	Print(Cat("pipe_test.go"))
}

func ExampleSystem() {
	Print(
		System("find", ".", "-type", "f", "-print"),
		Grep(`^\./pipe.*\.go$`),
		Sort(),
	)
	// TODO: Remove output checking if it becomes fragile.

	// Output:
	// ./pipe.go
	// ./pipe_test.go
}

func ExampleMix() {
	dbl := func(arg Arg) {
		for s := range arg.in {
			arg.out <- s
			arg.out <- s
		}
	}

	Print(Numbers(1, 100),
		Grep("3"),
		GrepNot("7"),
		dbl,
		Uniq,
		ReplaceMatch("^(.)$", "x$1"),
		Sort(),
		ReplaceMatch("^(.)", "$1 "),
		dbl,
		DeleteMatch(" .$"),
		UniqWithCount,
		Sort(Num(1)),
		Reverse,
	)
	// Output:
	// 18 3
	// 2 x
	// 2 9
	// 2 8
	// 2 6
	// 2 5
	// 2 4
	// 2 2
	// 2 1
}

func ExampleHash() {
	hash := func(f string, out chan<- string) {
		file, err := os.Open(f)
		if err != nil {
			fmt.Fprintln(os.Stderr, err)
			return
		}
		hasher := sha1.New()
		_, err = io.Copy(hasher, file)
		file.Close()
		if err != nil {
			fmt.Fprintln(os.Stderr, err)
			return
		}
		out <- fmt.Sprintf("%x %s", hasher.Sum(nil), f)
	}

	Print(
		Find(FILES, "/home/sanjay/tmp"),
		Grep("/tmp/x"),
		GrepNot("/sub2/"),
		Parallel(4, hash),
		ReplaceMatch(" /home/sanjay/", " HOME/"))

	Print(
		Find(FILES, "/home/sanjay/tmp/y"),
		GrepNot(`/home/sanjay/(\.Trash|Library)/`),
		Parallel(4, hash),
		Sort(Text(2)),
	)

	Print(
		System("find", "/home/sanjay/tmp/y", "-type", "f", "-print"),
		Parallel(4, hash),
		Sort(Text(2)),
	)

}
<|endoftext|>


<commit_before>//  Copyright (c) 2014 Couchbase, 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 collector

import (
	"time"

	"github.com/blevesearch/bleve/index"
	"github.com/blevesearch/bleve/search"
	"golang.org/x/net/context"
)

// PreAllocSizeSkipCap will cap preallocation to this amount when
// size+skip exceeds this value
var PreAllocSizeSkipCap = 1000

type collectorCompare func(i, j *search.DocumentMatch) int

type collectorFixup func(d *search.DocumentMatch) error

// TopNCollector collects the top N hits, optionally skipping some results
type TopNCollector struct {
	size          int
	skip          int
	total         uint64
	maxScore      float64
	took          time.Duration
	sort          search.SortOrder
	results       search.DocumentMatchCollection
	facetsBuilder *search.FacetsBuilder

	store *collectStoreSlice

	needDocIds    bool
	neededFields  []string
	cachedScoring []bool
	cachedDesc    []bool

	lowestMatchOutsideResults *search.DocumentMatch
}

// CheckDoneEvery controls how frequently we check the context deadline
const CheckDoneEvery = uint64(1024)

// NewTopNCollector builds a collector to find the top 'size' hits
// skipping over the first 'skip' hits
// ordering hits by the provided sort order
func NewTopNCollector(size int, skip int, sort search.SortOrder) *TopNCollector {
	hc := &TopNCollector{size: size, skip: skip, sort: sort}

	// pre-allocate space on the store to avoid reslicing
	// unless the size + skip is too large, then cap it
	// everything should still work, just reslices as necessary
	backingSize := size + skip + 1
	if size+skip > PreAllocSizeSkipCap {
		backingSize = PreAllocSizeSkipCap + 1
	}

	hc.store = newStoreSlice(backingSize, func(i, j *search.DocumentMatch) int {
		return hc.sort.Compare(hc.cachedScoring, hc.cachedDesc, i, j)
	})

	// these lookups traverse an interface, so do once up-front
	if sort.RequiresDocID() {
		hc.needDocIds = true
	}
	hc.neededFields = sort.RequiredFields()
	hc.cachedScoring = sort.CacheIsScore()
	hc.cachedDesc = sort.CacheDescending()

	return hc
}

// Collect goes to the index to find the matching documents
func (hc *TopNCollector) Collect(ctx context.Context, searcher search.Searcher, reader index.IndexReader) error {
	startTime := time.Now()
	var err error
	var next *search.DocumentMatch

	// pre-allocate enough space in the DocumentMatchPool
	// unless the size + skip is too large, then cap it
	// everything should still work, just allocates DocumentMatches on demand
	backingSize := hc.size + hc.skip + 1
	if hc.size+hc.skip > PreAllocSizeSkipCap {
		backingSize = PreAllocSizeSkipCap + 1
	}
	searchContext := &search.SearchContext{
		DocumentMatchPool: search.NewDocumentMatchPool(backingSize+searcher.DocumentMatchPoolSize(), len(hc.sort)),
	}

	select {
	case <-ctx.Done():
		return ctx.Err()
	default:
		next, err = searcher.Next(searchContext)
	}
	for err == nil && next != nil {
		if hc.total%CheckDoneEvery == 0 {
			select {
			case <-ctx.Done():
				return ctx.Err()
			default:
			}
		}

		err = hc.collectSingle(searchContext, reader, next)
		if err != nil {
			break
		}

		next, err = searcher.Next(searchContext)
	}
	// compute search duration
	hc.took = time.Since(startTime)
	if err != nil {
		return err
	}
	// finalize actual results
	err = hc.finalizeResults(reader)
	if err != nil {
		return err
	}
	return nil
}

var sortByScoreOpt = []string{"_score"}

func (hc *TopNCollector) collectSingle(ctx *search.SearchContext, reader index.IndexReader, d *search.DocumentMatch) error {

	// visit field terms for features that require it (sort, facets)
	err := hc.visitFieldTerms(reader, d)
	if err != nil {
		return err
	}

	// increment total hits
	hc.total++
	d.HitNumber = hc.total

	// update max score
	if d.Score > hc.maxScore {
		hc.maxScore = d.Score
	}

	// see if we need to load ID (at this early stage, for example to sort on it)
	if hc.needDocIds {
		d.ID, err = reader.ExternalID(d.IndexInternalID)
		if err != nil {
			return err
		}
	}

	// compute this hits sort value
	if len(hc.sort) == 1 && hc.cachedScoring[0] {
		d.Sort = sortByScoreOpt
	} else {
		hc.sort.Value(d)
	}

	// optimization, we track lowest sorting hit already removed from heap
	// with this one comparison, we can avoid all heap operations if
	// this hit would have been added and then immediately removed
	if hc.lowestMatchOutsideResults != nil {
		cmp := hc.sort.Compare(hc.cachedScoring, hc.cachedDesc, d, hc.lowestMatchOutsideResults)
		if cmp >= 0 {
			// this hit can't possibly be in the result set, so avoid heap ops
			ctx.DocumentMatchPool.Put(d)
			return nil
		}
	}

	hc.store.Add(d)
	if hc.store.Len() > hc.size+hc.skip {
		removed := hc.store.RemoveLast()
		if hc.lowestMatchOutsideResults == nil {
			hc.lowestMatchOutsideResults = removed
		} else {
			cmp := hc.sort.Compare(hc.cachedScoring, hc.cachedDesc, removed, hc.lowestMatchOutsideResults)
			if cmp < 0 {
				tmp := hc.lowestMatchOutsideResults
				hc.lowestMatchOutsideResults = removed
				ctx.DocumentMatchPool.Put(tmp)
			}
		}
	}

	return nil
}

// visitFieldTerms is responsible for visiting the field terms of the
// search hit, and passing visited terms to the sort and facet builder
func (hc *TopNCollector) visitFieldTerms(reader index.IndexReader, d *search.DocumentMatch) error {
	if hc.facetsBuilder != nil {
		hc.facetsBuilder.StartDoc()
	}

	err := reader.DocumentVisitFieldTerms(d.IndexInternalID, hc.neededFields, func(field string, term []byte) {
		if hc.facetsBuilder != nil {
			hc.facetsBuilder.UpdateVisitor(field, term)
		}
		hc.sort.UpdateVisitor(field, term)
	})

	if hc.facetsBuilder != nil {
		hc.facetsBuilder.EndDoc()
	}

	return err
}

// SetFacetsBuilder registers a facet builder for this collector
func (hc *TopNCollector) SetFacetsBuilder(facetsBuilder *search.FacetsBuilder) {
	hc.facetsBuilder = facetsBuilder
	hc.neededFields = append(hc.neededFields, hc.facetsBuilder.RequiredFields()...)
}

// finalizeResults starts with the heap containing the final top size+skip
// it now throws away the results to be skipped
// and does final doc id lookup (if necessary)
func (hc *TopNCollector) finalizeResults(r index.IndexReader) error {
	var err error
	hc.results, err = hc.store.Final(hc.skip, func(doc *search.DocumentMatch) error {
		if doc.ID == "" {
			// look up the id since we need it for lookup
			var err error
			doc.ID, err = r.ExternalID(doc.IndexInternalID)
			if err != nil {
				return err
			}
		}
		return nil
	})

	return err
}

// Results returns the collected hits
func (hc *TopNCollector) Results() search.DocumentMatchCollection {
	return hc.results
}

// Total returns the total number of hits
func (hc *TopNCollector) Total() uint64 {
	return hc.total
}

// MaxScore returns the maximum score seen across all the hits
func (hc *TopNCollector) MaxScore() float64 {
	return hc.maxScore
}

// Took returns the time spent collecting hits
func (hc *TopNCollector) Took() time.Duration {
	return hc.took
}

// FacetResults returns the computed facets results
func (hc *TopNCollector) FacetResults() search.FacetResults {
	if hc.facetsBuilder != nil {
		return hc.facetsBuilder.Results()
	}
	return search.FacetResults{}
}
<commit_msg>fix perf regression, unnecessarily loading backindex<commit_after>//  Copyright (c) 2014 Couchbase, 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 collector

import (
	"time"

	"github.com/blevesearch/bleve/index"
	"github.com/blevesearch/bleve/search"
	"golang.org/x/net/context"
)

// PreAllocSizeSkipCap will cap preallocation to this amount when
// size+skip exceeds this value
var PreAllocSizeSkipCap = 1000

type collectorCompare func(i, j *search.DocumentMatch) int

type collectorFixup func(d *search.DocumentMatch) error

// TopNCollector collects the top N hits, optionally skipping some results
type TopNCollector struct {
	size          int
	skip          int
	total         uint64
	maxScore      float64
	took          time.Duration
	sort          search.SortOrder
	results       search.DocumentMatchCollection
	facetsBuilder *search.FacetsBuilder

	store *collectStoreSlice

	needDocIds    bool
	neededFields  []string
	cachedScoring []bool
	cachedDesc    []bool

	lowestMatchOutsideResults *search.DocumentMatch
}

// CheckDoneEvery controls how frequently we check the context deadline
const CheckDoneEvery = uint64(1024)

// NewTopNCollector builds a collector to find the top 'size' hits
// skipping over the first 'skip' hits
// ordering hits by the provided sort order
func NewTopNCollector(size int, skip int, sort search.SortOrder) *TopNCollector {
	hc := &TopNCollector{size: size, skip: skip, sort: sort}

	// pre-allocate space on the store to avoid reslicing
	// unless the size + skip is too large, then cap it
	// everything should still work, just reslices as necessary
	backingSize := size + skip + 1
	if size+skip > PreAllocSizeSkipCap {
		backingSize = PreAllocSizeSkipCap + 1
	}

	hc.store = newStoreSlice(backingSize, func(i, j *search.DocumentMatch) int {
		return hc.sort.Compare(hc.cachedScoring, hc.cachedDesc, i, j)
	})

	// these lookups traverse an interface, so do once up-front
	if sort.RequiresDocID() {
		hc.needDocIds = true
	}
	hc.neededFields = sort.RequiredFields()
	hc.cachedScoring = sort.CacheIsScore()
	hc.cachedDesc = sort.CacheDescending()

	return hc
}

// Collect goes to the index to find the matching documents
func (hc *TopNCollector) Collect(ctx context.Context, searcher search.Searcher, reader index.IndexReader) error {
	startTime := time.Now()
	var err error
	var next *search.DocumentMatch

	// pre-allocate enough space in the DocumentMatchPool
	// unless the size + skip is too large, then cap it
	// everything should still work, just allocates DocumentMatches on demand
	backingSize := hc.size + hc.skip + 1
	if hc.size+hc.skip > PreAllocSizeSkipCap {
		backingSize = PreAllocSizeSkipCap + 1
	}
	searchContext := &search.SearchContext{
		DocumentMatchPool: search.NewDocumentMatchPool(backingSize+searcher.DocumentMatchPoolSize(), len(hc.sort)),
	}

	select {
	case <-ctx.Done():
		return ctx.Err()
	default:
		next, err = searcher.Next(searchContext)
	}
	for err == nil && next != nil {
		if hc.total%CheckDoneEvery == 0 {
			select {
			case <-ctx.Done():
				return ctx.Err()
			default:
			}
		}

		err = hc.collectSingle(searchContext, reader, next)
		if err != nil {
			break
		}

		next, err = searcher.Next(searchContext)
	}
	// compute search duration
	hc.took = time.Since(startTime)
	if err != nil {
		return err
	}
	// finalize actual results
	err = hc.finalizeResults(reader)
	if err != nil {
		return err
	}
	return nil
}

var sortByScoreOpt = []string{"_score"}

func (hc *TopNCollector) collectSingle(ctx *search.SearchContext, reader index.IndexReader, d *search.DocumentMatch) error {
	var err error

	// visit field terms for features that require it (sort, facets)
	if len(hc.neededFields) > 0 {
		err = hc.visitFieldTerms(reader, d)
		if err != nil {
			return err
		}
	}

	// increment total hits
	hc.total++
	d.HitNumber = hc.total

	// update max score
	if d.Score > hc.maxScore {
		hc.maxScore = d.Score
	}

	// see if we need to load ID (at this early stage, for example to sort on it)
	if hc.needDocIds {
		d.ID, err = reader.ExternalID(d.IndexInternalID)
		if err != nil {
			return err
		}
	}

	// compute this hits sort value
	if len(hc.sort) == 1 && hc.cachedScoring[0] {
		d.Sort = sortByScoreOpt
	} else {
		hc.sort.Value(d)
	}

	// optimization, we track lowest sorting hit already removed from heap
	// with this one comparison, we can avoid all heap operations if
	// this hit would have been added and then immediately removed
	if hc.lowestMatchOutsideResults != nil {
		cmp := hc.sort.Compare(hc.cachedScoring, hc.cachedDesc, d, hc.lowestMatchOutsideResults)
		if cmp >= 0 {
			// this hit can't possibly be in the result set, so avoid heap ops
			ctx.DocumentMatchPool.Put(d)
			return nil
		}
	}

	hc.store.Add(d)
	if hc.store.Len() > hc.size+hc.skip {
		removed := hc.store.RemoveLast()
		if hc.lowestMatchOutsideResults == nil {
			hc.lowestMatchOutsideResults = removed
		} else {
			cmp := hc.sort.Compare(hc.cachedScoring, hc.cachedDesc, removed, hc.lowestMatchOutsideResults)
			if cmp < 0 {
				tmp := hc.lowestMatchOutsideResults
				hc.lowestMatchOutsideResults = removed
				ctx.DocumentMatchPool.Put(tmp)
			}
		}
	}

	return nil
}

// visitFieldTerms is responsible for visiting the field terms of the
// search hit, and passing visited terms to the sort and facet builder
func (hc *TopNCollector) visitFieldTerms(reader index.IndexReader, d *search.DocumentMatch) error {
	if hc.facetsBuilder != nil {
		hc.facetsBuilder.StartDoc()
	}

	err := reader.DocumentVisitFieldTerms(d.IndexInternalID, hc.neededFields, func(field string, term []byte) {
		if hc.facetsBuilder != nil {
			hc.facetsBuilder.UpdateVisitor(field, term)
		}
		hc.sort.UpdateVisitor(field, term)
	})

	if hc.facetsBuilder != nil {
		hc.facetsBuilder.EndDoc()
	}

	return err
}

// SetFacetsBuilder registers a facet builder for this collector
func (hc *TopNCollector) SetFacetsBuilder(facetsBuilder *search.FacetsBuilder) {
	hc.facetsBuilder = facetsBuilder
	hc.neededFields = append(hc.neededFields, hc.facetsBuilder.RequiredFields()...)
}

// finalizeResults starts with the heap containing the final top size+skip
// it now throws away the results to be skipped
// and does final doc id lookup (if necessary)
func (hc *TopNCollector) finalizeResults(r index.IndexReader) error {
	var err error
	hc.results, err = hc.store.Final(hc.skip, func(doc *search.DocumentMatch) error {
		if doc.ID == "" {
			// look up the id since we need it for lookup
			var err error
			doc.ID, err = r.ExternalID(doc.IndexInternalID)
			if err != nil {
				return err
			}
		}
		return nil
	})

	return err
}

// Results returns the collected hits
func (hc *TopNCollector) Results() search.DocumentMatchCollection {
	return hc.results
}

// Total returns the total number of hits
func (hc *TopNCollector) Total() uint64 {
	return hc.total
}

// MaxScore returns the maximum score seen across all the hits
func (hc *TopNCollector) MaxScore() float64 {
	return hc.maxScore
}

// Took returns the time spent collecting hits
func (hc *TopNCollector) Took() time.Duration {
	return hc.took
}

// FacetResults returns the computed facets results
func (hc *TopNCollector) FacetResults() search.FacetResults {
	if hc.facetsBuilder != nil {
		return hc.facetsBuilder.Results()
	}
	return search.FacetResults{}
}
<|endoftext|>


<commit_before>/*
Package db contains functions related to database.
 * db.go includes conneting to db, query and insertion.
 * dbschema.go contains definitions of struct for db tables.
   Currently it only contains Module struct.
*/
package db

import (
	"database/sql"
	"fmt"
	"os"
	"strconv"

	// Go postgres driver for Go's database/sql package
	_ "github.com/lib/pq"
)

// Put all used SQL statments here for easy maintainess and tracking.
// To support more operations (query/update), just append this const list.
const (
	insertModule            = `INSERT INTO modules (orgName, name, version, data) VALUES($1, $2, $3, $4)`
	allModules              = `select * from modules`
	modulesByOrgName        = `select * from modules where orgName=$1`
	modulesByName           = `select * from modules where name=$1`
	modulesByVersion        = `select * from modules where version=$1`
	modulesByNameAndVersion = `select * from modules where name=$1 and version=$2`
)

var db *sql.DB

// Establish connection to database, *db* variable is assigned when opening database.
// This should only be called once before any other database function is called.
//
// Users need set environment variables for connection, including
//  * DB_HOST:          host address of target db instances, by default: localhost.
//  * DB_PORT:          port number of postgres db, by default: 5432.
//  * DB_USERNAME:      username of database, error would be returned if not set.
//  * DB_PWD:           password of target database, error would be returned if not set.
//  * DB_NAME:          name of database for connection, error would be returned if not set.
//  * DB_SOCKER_DIR:    directory of Unix socket in Cloud Run which serves as Cloud SQL
//                      Auth proxy to connect to postgres database.
//                      If service is deployed on Cloud Run, just use the default value.
//                      By default, it is set to `/cloudsql`.
func ConnectDB() error {
	var ok bool
	var err error
	var port int         // port number of target database
	var user string      // username of target database
	var password string  // password of target database
	var dbname string    // name of target database
	var host string      // host address of target database
	var socketDir string // (Cloud Run only) Directory of Unix socket
	var psqlconn string  // connection string used to connect to traget database

	// read db config from env

	if portStr, ok := os.LookupEnv("DB_PORT"); !ok {
		fmt.Println("DB_PORT not set, set port to 5432")
		port = 5432
	} else {
		if port, err = strconv.Atoi(portStr); err != nil {
			return fmt.Errorf("parse port failed: %v", err)
		}
	}

	if user, ok = os.LookupEnv("DB_USERNAME"); !ok {
		return fmt.Errorf("DB_USERNAME not set")
	}

	if password, ok = os.LookupEnv("DB_PWD"); !ok {
		return fmt.Errorf("DB_PWD not set")
	}

	if dbname, ok = os.LookupEnv("DB_NAME"); !ok {
		return fmt.Errorf("DB_NAME not set")
	}

	if socketDir, ok = os.LookupEnv("DB_SOCKET_DIR"); !ok {
		socketDir = "/cloudsql"
	}

	if host, ok = os.LookupEnv("DB_HOST"); !ok || host == "localhost" {
		if !ok {
			fmt.Println("DB_HOST not set, set host to localhost")
			host = "localhost"
		}
		// This connection string is used if service is not deployed on Cloud Run,
		// instead connection is made from localhost via Cloud SQL proxy.
		psqlconn = fmt.Sprintf("host=%s port=%d user=%s password=%s dbname=%s sslmode=disable", host, port, user, password, dbname)
	} else {
		psqlconn = fmt.Sprintf("host=%s/%s port=%d user=%s password=%s dbname=%s", socketDir, host, port, user, password, dbname)
	}

	// open database
	db, err = sql.Open("postgres", psqlconn)
	if err != nil {
		return fmt.Errorf("open database failed: %v", err)
	}

	// see if connection is established successfully
	if err = db.Ping(); err != nil {
		return fmt.Errorf("ping database failed: %v", err)
	}

	return nil
}

// Insert module into database given values of four field of MODULE schema.
// Error is returned when insertion failed.
func InsertModule(orgName string, name string, version string, data string) error {
	if _, err := db.Exec(insertModule, orgName, name, version, data); err != nil {
		return fmt.Errorf("insert module into db failed: %v", err)
	}
	return nil
}

// Scan from queried modules from rows one by one, rows are *not* closed inside.
// Return slice of db Module struct each field of which corresponds to one column in db.
// Error is returned when scan rows failed.
func ReadModluesByRow(rows *sql.Rows) ([]Module, error) {
	var modules []Module
	for rows.Next() {
		var module Module
		if err := rows.Scan(&module.OrgName, &module.Name, &module.Version, &module.Data); err != nil {
			return nil, fmt.Errorf("scan db rows failure, %v", err)
		}
		modules = append(modules, module)
	}
	return modules, nil
}

// Query modules of organization with *orgName* from database.
// If orgName is null then directly query all modules.
// Return slice of db Module struct each field of which corresponds to one column in db.
// Error is returned when query or reading data failed.
func QueryModulesByOrgName(orgName *string) ([]Module, error) {
	var rows *sql.Rows
	var err error

	if orgName != nil {
		rows, err = db.Query(modulesByOrgName, *orgName)
		if err != nil {
			return nil, fmt.Errorf("QueryModulesByOrgName failed: %v, query param orgName: %s", err, *orgName)
		}
	} else {
		rows, err = db.Query(allModules)
		if err != nil {
			return nil, fmt.Errorf("QueryModulesByOrgName failed: %v", err)
		}
	}

	defer rows.Close()

	return ReadModluesByRow(rows)
}

// Query modules by its key (name, version), it is possible that parameters are null.
// If both parameters are null, this equals query for all modules.
// Return slice of db Module struct each field of which corresponds to one column in db.
// Error is returned when query or reading data failed.
func QueryModulesByKey(name *string, version *string) ([]Module, error) {
	var rows *sql.Rows
	var err error

	if name != nil && version != nil {
		rows, err = db.Query(modulesByNameAndVersion, *name, *version)
		if err != nil {
			return nil, fmt.Errorf("QueryModulesByKey failed: %v, query param name: %s, version: %s", err, *name, *version)
		}
	} else if name != nil {
		rows, err = db.Query(modulesByName, *name)
		if err != nil {
			return nil, fmt.Errorf("QueryModulesByKey failed: %v, query param name: %s", err, *name)
		}
	} else if version != nil {
		rows, err = db.Query(modulesByVersion, *version)
		if err != nil {
			return nil, fmt.Errorf("QueryModulesByKey failed: %v, query param version: %s", err, *version)
		}
	} else {
		rows, err = db.Query(allModules)
		if err != nil {
			return nil, fmt.Errorf("QueryModulesByKey failed: %v", err)
		}
	}

	defer rows.Close()

	return ReadModluesByRow(rows)
}

// Close db connection
func Close() error {
	return db.Close()
}
<commit_msg>Dynamically generate query statements based its query parameters<commit_after>/*
Package db contains functions related to database.
 * db.go includes conneting to db, query and insertion.
 * dbschema.go contains definitions of struct for db tables.
   Currently it only contains Module struct.
*/
package db

import (
	"database/sql"
	"fmt"
	"os"
	"strconv"

	// Go postgres driver for Go's database/sql package
	_ "github.com/lib/pq"
)

// SQL stataments used in this package.
// Query statements can be appended based on its query parameters.
const (
	insertModule  = `INSERT INTO modules (orgName, name, version, data) VALUES($1, $2, $3, $4)`
	selectModules = `select * from modules`
)

var db *sql.DB

// Establish connection to database, *db* variable is assigned when opening database.
// This should only be called once before any other database function is called.
//
// Users need set environment variables for connection, including
//  * DB_HOST:          host address of target db instances, by default: localhost.
//  * DB_PORT:          port number of postgres db, by default: 5432.
//  * DB_USERNAME:      username of database, error would be returned if not set.
//  * DB_PWD:           password of target database, error would be returned if not set.
//  * DB_NAME:          name of database for connection, error would be returned if not set.
//  * DB_SOCKER_DIR:    directory of Unix socket in Cloud Run which serves as Cloud SQL
//                      Auth proxy to connect to postgres database.
//                      If service is deployed on Cloud Run, just use the default value.
//                      By default, it is set to `/cloudsql`.
func ConnectDB() error {
	var ok bool
	var err error
	var port int         // port number of target database
	var user string      // username of target database
	var password string  // password of target database
	var dbname string    // name of target database
	var host string      // host address of target database
	var socketDir string // (Cloud Run only) Directory of Unix socket
	var psqlconn string  // connection string used to connect to traget database

	// read db config from env

	if portStr, ok := os.LookupEnv("DB_PORT"); !ok {
		fmt.Println("DB_PORT not set, set port to 5432")
		port = 5432
	} else {
		if port, err = strconv.Atoi(portStr); err != nil {
			return fmt.Errorf("parse port failed: %v", err)
		}
	}

	if user, ok = os.LookupEnv("DB_USERNAME"); !ok {
		return fmt.Errorf("DB_USERNAME not set")
	}

	if password, ok = os.LookupEnv("DB_PWD"); !ok {
		return fmt.Errorf("DB_PWD not set")
	}

	if dbname, ok = os.LookupEnv("DB_NAME"); !ok {
		return fmt.Errorf("DB_NAME not set")
	}

	if socketDir, ok = os.LookupEnv("DB_SOCKET_DIR"); !ok {
		socketDir = "/cloudsql"
	}

	if host, ok = os.LookupEnv("DB_HOST"); !ok || host == "localhost" {
		if !ok {
			fmt.Println("DB_HOST not set, set host to localhost")
			host = "localhost"
		}
		// This connection string is used if service is not deployed on Cloud Run,
		// instead connection is made from localhost via Cloud SQL proxy.
		psqlconn = fmt.Sprintf("host=%s port=%d user=%s password=%s dbname=%s sslmode=disable", host, port, user, password, dbname)
	} else {
		psqlconn = fmt.Sprintf("host=%s/%s port=%d user=%s password=%s dbname=%s", socketDir, host, port, user, password, dbname)
	}

	// open database
	db, err = sql.Open("postgres", psqlconn)
	if err != nil {
		return fmt.Errorf("open database failed: %v", err)
	}

	// see if connection is established successfully
	if err = db.Ping(); err != nil {
		return fmt.Errorf("ping database failed: %v", err)
	}

	return nil
}

// Insert module into database given values of four field of MODULE schema.
// Error is returned when insertion failed.
func InsertModule(orgName string, name string, version string, data string) error {
	if _, err := db.Exec(insertModule, orgName, name, version, data); err != nil {
		return fmt.Errorf("insert module into db failed: %v", err)
	}
	return nil
}

// Scan from queried modules from rows one by one, rows are *not* closed inside.
// Return slice of db Module struct each field of which corresponds to one column in db.
// Error is returned when scan rows failed.
func ReadModluesByRow(rows *sql.Rows) ([]Module, error) {
	var modules []Module
	for rows.Next() {
		var module Module
		if err := rows.Scan(&module.OrgName, &module.Name, &module.Version, &module.Data); err != nil {
			return nil, fmt.Errorf("scan db rows failure, %v", err)
		}
		modules = append(modules, module)
	}
	return modules, nil
}

// This function is used to generate query statement string based on query parameters.
// * parmNames is a list of names of all non-nil query parameters.
// * baseQuery is query statement without any query parameters.
func FormatQueryStr(parmNames []string, baseQuery string) string {
	queryStmt := baseQuery
	for i := 0; i < len(parmNames); i++ {
		if i == 0 {
			queryStmt += " where"
		} else {
			queryStmt += " and"
		}
		queryStmt += fmt.Sprintf(" %s=$%d", parmNames[i], i+1)
	}
	return queryStmt
}

// Query modules of organization with *orgName* from database.
// If orgName is null then directly query all modules.
// Return slice of db Module struct each field of which corresponds to one column in db.
// Error is returned when query or reading data failed.
func QueryModulesByOrgName(orgName *string) ([]Module, error) {
	parms := []interface{}{} // parms is used to store value of non-nil query parameters
	parmNames := []string{}  // parmNames is used to store name of non-nil query parameters

	if orgName != nil {
		parms = append(parms, *orgName)
		parmNames = append(parmNames, "orgName")
	}

	// Format query statement string based on non-nil query parameters
	queryStmt := FormatQueryStr(parmNames, selectModules)

	rows, err := db.Query(queryStmt, parms...)
	if err != nil {
		return nil, fmt.Errorf("QueryModulesByOrgName failed: %v", err)
	}

	defer rows.Close()

	return ReadModluesByRow(rows)
}

// Query modules by its key (name, version), it is possible that parameters are null.
// If both parameters are null, this equals query for all modules.
// Return slice of db Module struct each field of which corresponds to one column in db.
// Error is returned when query or reading data failed.
func QueryModulesByKey(name *string, version *string) ([]Module, error) {
	parms := []interface{}{} // parms is used to store value of non-nil query parameters
	parmNames := []string{}  // parmNames is used to store name of non-nil query parameters

	if name != nil {
		parms = append(parms, *name)
		parmNames = append(parmNames, "name")
	}

	if version != nil {
		parms = append(parms, *version)
		parmNames = append(parmNames, "version")
	}

	// Format query statement string based on non-nil query parameters
	queryStmt := FormatQueryStr(parmNames, selectModules)

	rows, err := db.Query(queryStmt, parms...)
	if err != nil {
		return nil, fmt.Errorf("QueryModulesByOrgName failed: %v", err)
	}

	defer rows.Close()

	return ReadModluesByRow(rows)
}

// Close db connection
func Close() error {
	return db.Close()
}
<|endoftext|>


<commit_before>package main

import (
	"context"
	"strings"
	"testing"
	"time"

	"github.com/ejholmes/walk/internal/dag"
	"github.com/stretchr/testify/assert"
)

var ctx = context.Background()

func TestPlan(t *testing.T) {
	clean(t)

	err := Exec(ctx, "test/110-compile/all")
	assert.NoError(t, err)
}

func TestPlan_CyclicDependencies(t *testing.T) {
	clean(t)

	err := Exec(ctx, "test/000-cyclic/all").(*dag.MultiError)
	assert.Equal(t, 1, len(err.Errors))
	assert.True(t, strings.Contains(err.Errors[0].Error(), "Cycle"))
}

func TestPlan_Cancel(t *testing.T) {
	clean(t)

	ctx, cancel := context.WithTimeout(ctx, 1*time.Second)
	defer cancel()
	err := Exec(ctx, "test/000-cancel/all").(*dag.MultiError)
	assert.Equal(t, 1, len(err.Errors))
	assert.True(t, strings.Contains(err.Errors[0].Error(), "signal: killed"))
}

func clean(t testing.TB) {
	err := Exec(ctx, "test/clean")
	assert.NoError(t, err)
}
<commit_msg>Add test for specifying multiple targets<commit_after>package main

import (
	"context"
	"strings"
	"testing"
	"time"

	"github.com/ejholmes/walk/internal/dag"
	"github.com/stretchr/testify/assert"
)

var ctx = context.Background()

func TestPlan(t *testing.T) {
	clean(t)

	err := Exec(ctx, "test/110-compile/all")
	assert.NoError(t, err)
}

func TestPlan_Multi(t *testing.T) {
	clean(t)

	err := Exec(ctx, "test/110-compile/all", "test/111-compile/all")
	assert.NoError(t, err)
}

func TestPlan_CyclicDependencies(t *testing.T) {
	clean(t)

	err := Exec(ctx, "test/000-cyclic/all").(*dag.MultiError)
	assert.Equal(t, 1, len(err.Errors))
	assert.True(t, strings.Contains(err.Errors[0].Error(), "Cycle"))
}

func TestPlan_Cancel(t *testing.T) {
	clean(t)

	ctx, cancel := context.WithTimeout(ctx, 1*time.Second)
	defer cancel()
	err := Exec(ctx, "test/000-cancel/all").(*dag.MultiError)
	assert.Equal(t, 1, len(err.Errors))
	assert.True(t, strings.Contains(err.Errors[0].Error(), "signal: killed"))
}

func clean(t testing.TB) {
	err := Exec(ctx, "test/clean")
	assert.NoError(t, err)
}
<|endoftext|>


<commit_before>package records

import (
	"encoding/json"
	"io/ioutil"
	"reflect"
	"testing"

	"github.com/mesosphere/mesos-dns/logging"
	"github.com/mesosphere/mesos-dns/records/labels"
)

func init() {
	logging.VerboseFlag = false
	logging.VeryVerboseFlag = false
	logging.SetupLogs()
}

func TestMasterRecord(t *testing.T) {
	// masterRecord(domain string, masters []string, leader string)
	type expectedRR struct {
		name  string
		host  string
		rtype string
	}
	tt := []struct {
		domain  string
		masters []string
		leader  string
		expect  []expectedRR
	}{
		{"foo.com", nil, "", nil},
		{"foo.com", nil, "@", nil},
		{"foo.com", nil, "1@", nil},
		{"foo.com", nil, "@2", nil},
		{"foo.com", nil, "3@4", nil},
		{"foo.com", nil, "5@6:7",
			[]expectedRR{
				{"leader.foo.com.", "6", "A"},
				{"master.foo.com.", "6", "A"},
				{"master0.foo.com.", "6", "A"},
				{"_leader._tcp.foo.com.", "leader.foo.com.:7", "SRV"},
				{"_leader._udp.foo.com.", "leader.foo.com.:7", "SRV"},
			}},
		// single master: leader and fallback
		{"foo.com", []string{"6:7"}, "5@6:7",
			[]expectedRR{
				{"leader.foo.com.", "6", "A"},
				{"master.foo.com.", "6", "A"},
				{"master0.foo.com.", "6", "A"},
				{"_leader._tcp.foo.com.", "leader.foo.com.:7", "SRV"},
				{"_leader._udp.foo.com.", "leader.foo.com.:7", "SRV"},
			}},
		// leader not in fallback list
		{"foo.com", []string{"8:9"}, "5@6:7",
			[]expectedRR{
				{"leader.foo.com.", "6", "A"},
				{"master.foo.com.", "6", "A"},
				{"master.foo.com.", "8", "A"},
				{"master1.foo.com.", "6", "A"},
				{"master0.foo.com.", "8", "A"},
				{"_leader._tcp.foo.com.", "leader.foo.com.:7", "SRV"},
				{"_leader._udp.foo.com.", "leader.foo.com.:7", "SRV"},
			}},
		// duplicate fallback masters, leader not in fallback list
		{"foo.com", []string{"8:9", "8:9"}, "5@6:7",
			[]expectedRR{
				{"leader.foo.com.", "6", "A"},
				{"master.foo.com.", "6", "A"},
				{"master.foo.com.", "8", "A"},
				{"master1.foo.com.", "6", "A"},
				{"master0.foo.com.", "8", "A"},
				{"_leader._tcp.foo.com.", "leader.foo.com.:7", "SRV"},
				{"_leader._udp.foo.com.", "leader.foo.com.:7", "SRV"},
			}},
		// leader that's also listed in the fallback list (at the end)
		{"foo.com", []string{"8:9", "6:7"}, "5@6:7",
			[]expectedRR{
				{"leader.foo.com.", "6", "A"},
				{"master.foo.com.", "6", "A"},
				{"master.foo.com.", "8", "A"},
				{"master1.foo.com.", "6", "A"},
				{"master0.foo.com.", "8", "A"},
				{"_leader._tcp.foo.com.", "leader.foo.com.:7", "SRV"},
				{"_leader._udp.foo.com.", "leader.foo.com.:7", "SRV"},
			}},
		// duplicate leading masters in the fallback list
		{"foo.com", []string{"8:9", "6:7", "6:7"}, "5@6:7",
			[]expectedRR{
				{"leader.foo.com.", "6", "A"},
				{"master.foo.com.", "6", "A"},
				{"master.foo.com.", "8", "A"},
				{"master1.foo.com.", "6", "A"},
				{"master0.foo.com.", "8", "A"},
				{"_leader._tcp.foo.com.", "leader.foo.com.:7", "SRV"},
				{"_leader._udp.foo.com.", "leader.foo.com.:7", "SRV"},
			}},
		// leader that's also listed in the fallback list (in the middle)
		{"foo.com", []string{"8:9", "6:7", "bob:0"}, "5@6:7",
			[]expectedRR{
				{"leader.foo.com.", "6", "A"},
				{"master.foo.com.", "6", "A"},
				{"master.foo.com.", "8", "A"},
				{"master.foo.com.", "bob", "A"},
				{"master0.foo.com.", "8", "A"},
				{"master1.foo.com.", "6", "A"},
				{"master2.foo.com.", "bob", "A"},
				{"_leader._tcp.foo.com.", "leader.foo.com.:7", "SRV"},
				{"_leader._udp.foo.com.", "leader.foo.com.:7", "SRV"},
			}},
	}
	for i, tc := range tt {
		rg := &RecordGenerator{}
		rg.As = make(rrs)
		rg.SRVs = make(rrs)
		t.Logf("test case %d", i+1)
		rg.masterRecord(tc.domain, tc.masters, tc.leader)
		if tc.expect == nil {
			if len(rg.As) > 0 {
				t.Fatalf("test case %d: unexpected As: %v", i+1, rg.As)
			}
			if len(rg.SRVs) > 0 {
				t.Fatalf("test case %d: unexpected SRVs: %v", i+1, rg.SRVs)
			}
		}
		expectedA := make(rrs)
		expectedSRV := make(rrs)
		for _, e := range tc.expect {
			found := rg.exists(e.name, e.host, e.rtype)
			if !found {
				t.Fatalf("test case %d: missing expected record: name=%q host=%q rtype=%s, As=%v", i+1, e.name, e.host, e.rtype, rg.As)
			}
			if e.rtype == "A" {
				expectedA[e.name] = append(expectedA[e.name], e.host)
			} else {
				expectedSRV[e.name] = append(expectedSRV[e.name], e.host)
			}
		}
		if !reflect.DeepEqual(rg.As, expectedA) {
			t.Fatalf("test case %d: expected As of %v instead of %v", i+1, expectedA, rg.As)
		}
		if !reflect.DeepEqual(rg.SRVs, expectedSRV) {
			t.Fatalf("test case %d: expected SRVs of %v instead of %v", i+1, expectedSRV, rg.SRVs)
		}
	}
}

func TestSanitizedSlaveAddress(t *testing.T) {
	spec := labels.ForRFC952()
	x := sanitizedSlaveAddress("1.2.3.4", spec)
	if x != "1.2.3.4" {
		t.Fatalf("unexpected slave address %q", x)
	}

	x = sanitizedSlaveAddress("localhost", spec)
	if x != "127.0.0.1" {
		t.Fatalf("unexpected slave address %q", x)
	}

	x = sanitizedSlaveAddress("unbelievable.domain.acme", spec)
	if x != "unbelievable.domain.acme" {
		t.Fatalf("unexpected slave address %q", x)
	}

	x = sanitizedSlaveAddress("unbelievable<>.domain!@#...acme", spec)
	if x != "unbelievable.domain.acme" {
		t.Fatalf("unexpected slave address %q", x)
	}
}

func TestYankPorts(t *testing.T) {
	p := "[31328-31328]"

	ports := yankPorts(p)

	if ports[0] != "31328" {
		t.Error("not parsing port")
	}
}

func TestMultipleYankPorts(t *testing.T) {
	p := "[31111-31111, 31113-31113]"

	ports := yankPorts(p)

	if len(ports) != 2 {
		t.Error("not parsing ports")
	}

	if ports[0] != "31111" {
		t.Error("not parsing port")
	}

	if ports[1] != "31113" {
		t.Error("not parsing port")
	}
}

func TestRangePorts(t *testing.T) {
	p := "[31115-31117]"

	ports := yankPorts(p)

	if len(ports) != 3 {
		t.Error("not parsing ports")
	}

	if ports[0] != "31115" {
		t.Error("not parsing port")
	}

	if ports[1] != "31116" {
		t.Error("not parsing port")
	}

	if ports[2] != "31117" {
		t.Error("not parsing port")
	}

}

func TestLeaderIP(t *testing.T) {
	l := "master@144.76.157.37:5050"

	ip := leaderIP(l)

	if ip != "144.76.157.37" {
		t.Error("not parsing ip")
	}
}

// ensure we are parsing what we think we are
func TestInsertState(t *testing.T) {

	var sj StateJSON

	b, err := ioutil.ReadFile("../factories/fake.json")
	if err != nil {
		t.Error("missing test data")
	}

	err = json.Unmarshal(b, &sj)
	if err != nil {
		t.Error(err)
	}
	sj.Leader = "master@144.76.157.37:5050"

	masters := []string{"144.76.157.37:5050"}
	spec := labels.ForRFC952()
	rg := &RecordGenerator{}
	rg.InsertState(sj, "mesos", "mesos-dns.mesos.", "127.0.0.1", masters, spec)

	// ensure we are only collecting running tasks
	_, ok := rg.SRVs["_poseidon._tcp.marathon.mesos."]
	if ok {
		t.Error("should not find this not-running task - SRV record")
	}

	_, ok = rg.As["liquor-store.marathon.mesos."]
	if !ok {
		t.Error("should find this running task - A record")
	}

	_, ok = rg.As["poseidon.marathon.mesos."]
	if ok {
		t.Error("should not find this not-running task - A record")
	}

	_, ok = rg.As["master.mesos."]
	if !ok {
		t.Error("should find a running master - A record")
	}

	_, ok = rg.As["master0.mesos."]
	if !ok {
		t.Error("should find a running master0 - A record")
	}

	_, ok = rg.As["leader.mesos."]
	if !ok {
		t.Error("should find a leading master - A record")
	}

	_, ok = rg.SRVs["_leader._tcp.mesos."]
	if !ok {
		t.Error("should find a leading master - SRV record")
	}

	// test for 10 SRV names
	if len(rg.SRVs) != 10 {
		t.Error("not enough SRVs")
	}

	// test for 5 A names
	if len(rg.As) != 13 {
		t.Error("not enough As")
	}

	// ensure we translate the framework name as well
	_, ok = rg.As["some-box.chronoswithaspaceandmixe.mesos."]
	if !ok {
		t.Error("should find this task w/a space in the framework name - A record")
	}

	// ensure we find this SRV
	rrs := rg.SRVs["_liquor-store._tcp.marathon.mesos."]
	// ensure there are 3 RRDATA answers for this SRV name
	if len(rrs) != 3 {
		t.Error("not enough SRV records")
	}

	// ensure we don't find this as a SRV record
	rrs = rg.SRVs["_liquor-store.marathon.mesos."]
	if len(rrs) != 0 {
		t.Error("not a proper SRV record")
	}

}

// ensure we only generate one A record for each host
func TestNTasks(t *testing.T) {
	rg := &RecordGenerator{}
	rg.As = make(rrs)

	rg.insertRR("blah.mesos", "10.0.0.1", "A")
	rg.insertRR("blah.mesos", "10.0.0.1", "A")
	rg.insertRR("blah.mesos", "10.0.0.2", "A")

	k, _ := rg.As["blah.mesos"]

	if len(k) != 2 {
		t.Error("should only have 2 A records")
	}
}
<commit_msg>Use fatal erros in tests when fake.json is invalid<commit_after>package records

import (
	"encoding/json"
	"io/ioutil"
	"reflect"
	"testing"

	"github.com/mesosphere/mesos-dns/logging"
	"github.com/mesosphere/mesos-dns/records/labels"
)

func init() {
	logging.VerboseFlag = false
	logging.VeryVerboseFlag = false
	logging.SetupLogs()
}

func TestMasterRecord(t *testing.T) {
	// masterRecord(domain string, masters []string, leader string)
	type expectedRR struct {
		name  string
		host  string
		rtype string
	}
	tt := []struct {
		domain  string
		masters []string
		leader  string
		expect  []expectedRR
	}{
		{"foo.com", nil, "", nil},
		{"foo.com", nil, "@", nil},
		{"foo.com", nil, "1@", nil},
		{"foo.com", nil, "@2", nil},
		{"foo.com", nil, "3@4", nil},
		{"foo.com", nil, "5@6:7",
			[]expectedRR{
				{"leader.foo.com.", "6", "A"},
				{"master.foo.com.", "6", "A"},
				{"master0.foo.com.", "6", "A"},
				{"_leader._tcp.foo.com.", "leader.foo.com.:7", "SRV"},
				{"_leader._udp.foo.com.", "leader.foo.com.:7", "SRV"},
			}},
		// single master: leader and fallback
		{"foo.com", []string{"6:7"}, "5@6:7",
			[]expectedRR{
				{"leader.foo.com.", "6", "A"},
				{"master.foo.com.", "6", "A"},
				{"master0.foo.com.", "6", "A"},
				{"_leader._tcp.foo.com.", "leader.foo.com.:7", "SRV"},
				{"_leader._udp.foo.com.", "leader.foo.com.:7", "SRV"},
			}},
		// leader not in fallback list
		{"foo.com", []string{"8:9"}, "5@6:7",
			[]expectedRR{
				{"leader.foo.com.", "6", "A"},
				{"master.foo.com.", "6", "A"},
				{"master.foo.com.", "8", "A"},
				{"master1.foo.com.", "6", "A"},
				{"master0.foo.com.", "8", "A"},
				{"_leader._tcp.foo.com.", "leader.foo.com.:7", "SRV"},
				{"_leader._udp.foo.com.", "leader.foo.com.:7", "SRV"},
			}},
		// duplicate fallback masters, leader not in fallback list
		{"foo.com", []string{"8:9", "8:9"}, "5@6:7",
			[]expectedRR{
				{"leader.foo.com.", "6", "A"},
				{"master.foo.com.", "6", "A"},
				{"master.foo.com.", "8", "A"},
				{"master1.foo.com.", "6", "A"},
				{"master0.foo.com.", "8", "A"},
				{"_leader._tcp.foo.com.", "leader.foo.com.:7", "SRV"},
				{"_leader._udp.foo.com.", "leader.foo.com.:7", "SRV"},
			}},
		// leader that's also listed in the fallback list (at the end)
		{"foo.com", []string{"8:9", "6:7"}, "5@6:7",
			[]expectedRR{
				{"leader.foo.com.", "6", "A"},
				{"master.foo.com.", "6", "A"},
				{"master.foo.com.", "8", "A"},
				{"master1.foo.com.", "6", "A"},
				{"master0.foo.com.", "8", "A"},
				{"_leader._tcp.foo.com.", "leader.foo.com.:7", "SRV"},
				{"_leader._udp.foo.com.", "leader.foo.com.:7", "SRV"},
			}},
		// duplicate leading masters in the fallback list
		{"foo.com", []string{"8:9", "6:7", "6:7"}, "5@6:7",
			[]expectedRR{
				{"leader.foo.com.", "6", "A"},
				{"master.foo.com.", "6", "A"},
				{"master.foo.com.", "8", "A"},
				{"master1.foo.com.", "6", "A"},
				{"master0.foo.com.", "8", "A"},
				{"_leader._tcp.foo.com.", "leader.foo.com.:7", "SRV"},
				{"_leader._udp.foo.com.", "leader.foo.com.:7", "SRV"},
			}},
		// leader that's also listed in the fallback list (in the middle)
		{"foo.com", []string{"8:9", "6:7", "bob:0"}, "5@6:7",
			[]expectedRR{
				{"leader.foo.com.", "6", "A"},
				{"master.foo.com.", "6", "A"},
				{"master.foo.com.", "8", "A"},
				{"master.foo.com.", "bob", "A"},
				{"master0.foo.com.", "8", "A"},
				{"master1.foo.com.", "6", "A"},
				{"master2.foo.com.", "bob", "A"},
				{"_leader._tcp.foo.com.", "leader.foo.com.:7", "SRV"},
				{"_leader._udp.foo.com.", "leader.foo.com.:7", "SRV"},
			}},
	}
	for i, tc := range tt {
		rg := &RecordGenerator{}
		rg.As = make(rrs)
		rg.SRVs = make(rrs)
		t.Logf("test case %d", i+1)
		rg.masterRecord(tc.domain, tc.masters, tc.leader)
		if tc.expect == nil {
			if len(rg.As) > 0 {
				t.Fatalf("test case %d: unexpected As: %v", i+1, rg.As)
			}
			if len(rg.SRVs) > 0 {
				t.Fatalf("test case %d: unexpected SRVs: %v", i+1, rg.SRVs)
			}
		}
		expectedA := make(rrs)
		expectedSRV := make(rrs)
		for _, e := range tc.expect {
			found := rg.exists(e.name, e.host, e.rtype)
			if !found {
				t.Fatalf("test case %d: missing expected record: name=%q host=%q rtype=%s, As=%v", i+1, e.name, e.host, e.rtype, rg.As)
			}
			if e.rtype == "A" {
				expectedA[e.name] = append(expectedA[e.name], e.host)
			} else {
				expectedSRV[e.name] = append(expectedSRV[e.name], e.host)
			}
		}
		if !reflect.DeepEqual(rg.As, expectedA) {
			t.Fatalf("test case %d: expected As of %v instead of %v", i+1, expectedA, rg.As)
		}
		if !reflect.DeepEqual(rg.SRVs, expectedSRV) {
			t.Fatalf("test case %d: expected SRVs of %v instead of %v", i+1, expectedSRV, rg.SRVs)
		}
	}
}

func TestSanitizedSlaveAddress(t *testing.T) {
	spec := labels.ForRFC952()
	x := sanitizedSlaveAddress("1.2.3.4", spec)
	if x != "1.2.3.4" {
		t.Fatalf("unexpected slave address %q", x)
	}

	x = sanitizedSlaveAddress("localhost", spec)
	if x != "127.0.0.1" {
		t.Fatalf("unexpected slave address %q", x)
	}

	x = sanitizedSlaveAddress("unbelievable.domain.acme", spec)
	if x != "unbelievable.domain.acme" {
		t.Fatalf("unexpected slave address %q", x)
	}

	x = sanitizedSlaveAddress("unbelievable<>.domain!@#...acme", spec)
	if x != "unbelievable.domain.acme" {
		t.Fatalf("unexpected slave address %q", x)
	}
}

func TestYankPorts(t *testing.T) {
	p := "[31328-31328]"

	ports := yankPorts(p)

	if ports[0] != "31328" {
		t.Error("not parsing port")
	}
}

func TestMultipleYankPorts(t *testing.T) {
	p := "[31111-31111, 31113-31113]"

	ports := yankPorts(p)

	if len(ports) != 2 {
		t.Error("not parsing ports")
	}

	if ports[0] != "31111" {
		t.Error("not parsing port")
	}

	if ports[1] != "31113" {
		t.Error("not parsing port")
	}
}

func TestRangePorts(t *testing.T) {
	p := "[31115-31117]"

	ports := yankPorts(p)

	if len(ports) != 3 {
		t.Error("not parsing ports")
	}

	if ports[0] != "31115" {
		t.Error("not parsing port")
	}

	if ports[1] != "31116" {
		t.Error("not parsing port")
	}

	if ports[2] != "31117" {
		t.Error("not parsing port")
	}

}

func TestLeaderIP(t *testing.T) {
	l := "master@144.76.157.37:5050"

	ip := leaderIP(l)

	if ip != "144.76.157.37" {
		t.Error("not parsing ip")
	}
}

// ensure we are parsing what we think we are
func TestInsertState(t *testing.T) {
	var sj StateJSON

	b, err := ioutil.ReadFile("../factories/fake.json")
	if err != nil {
		t.Fatal("missing test data")
	}

	err = json.Unmarshal(b, &sj)
	if err != nil {
		t.Fatal(err)
	}
	sj.Leader = "master@144.76.157.37:5050"

	masters := []string{"144.76.157.37:5050"}
	spec := labels.ForRFC952()
	rg := &RecordGenerator{}
	rg.InsertState(sj, "mesos", "mesos-dns.mesos.", "127.0.0.1", masters, spec)

	// ensure we are only collecting running tasks
	_, ok := rg.SRVs["_poseidon._tcp.marathon.mesos."]
	if ok {
		t.Error("should not find this not-running task - SRV record")
	}

	_, ok = rg.As["liquor-store.marathon.mesos."]
	if !ok {
		t.Error("should find this running task - A record")
	}

	_, ok = rg.As["poseidon.marathon.mesos."]
	if ok {
		t.Error("should not find this not-running task - A record")
	}

	_, ok = rg.As["master.mesos."]
	if !ok {
		t.Error("should find a running master - A record")
	}

	_, ok = rg.As["master0.mesos."]
	if !ok {
		t.Error("should find a running master0 - A record")
	}

	_, ok = rg.As["leader.mesos."]
	if !ok {
		t.Error("should find a leading master - A record")
	}

	_, ok = rg.SRVs["_leader._tcp.mesos."]
	if !ok {
		t.Error("should find a leading master - SRV record")
	}

	// test for 10 SRV names
	if len(rg.SRVs) != 10 {
		t.Error("not enough SRVs")
	}

	// test for 5 A names
	if len(rg.As) != 13 {
		t.Error("not enough As")
	}

	// ensure we translate the framework name as well
	_, ok = rg.As["some-box.chronoswithaspaceandmixe.mesos."]
	if !ok {
		t.Error("should find this task w/a space in the framework name - A record")
	}

	// ensure we find this SRV
	rrs := rg.SRVs["_liquor-store._tcp.marathon.mesos."]
	// ensure there are 3 RRDATA answers for this SRV name
	if len(rrs) != 3 {
		t.Error("not enough SRV records")
	}

	// ensure we don't find this as a SRV record
	rrs = rg.SRVs["_liquor-store.marathon.mesos."]
	if len(rrs) != 0 {
		t.Error("not a proper SRV record")
	}

}

// ensure we only generate one A record for each host
func TestNTasks(t *testing.T) {
	rg := &RecordGenerator{}
	rg.As = make(rrs)

	rg.insertRR("blah.mesos", "10.0.0.1", "A")
	rg.insertRR("blah.mesos", "10.0.0.1", "A")
	rg.insertRR("blah.mesos", "10.0.0.2", "A")

	k, _ := rg.As["blah.mesos"]

	if len(k) != 2 {
		t.Error("should only have 2 A records")
	}
}
<|endoftext|>


<commit_before>package main

import (
	"archive/zip"
	"database/sql"
	"errors"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"github.com/briandowns/spinner"
	"github.com/mattn/go-zglob"

	"golang.org/x/tools/godoc/vfs/zipfs"

	"github.com/beevik/etree"
	_ "github.com/mattn/go-sqlite3"
)

var version = "dev"

func copyFile(src, dst string) error {
	in, err := os.Open(src)
	if err != nil {
		return err
	}
	defer in.Close()

	out, err := os.Create(dst)
	if err != nil {
		return err
	}
	defer out.Close()

	_, err = io.Copy(out, in)
	if err != nil {
		return err
	}
	return out.Close()
}

func pathToContentID(koboPath, path string) (imageID string, err error) {
	relPath, err := filepath.Rel(koboPath, path)
	if err != nil {
		return "", fmt.Errorf("could not get relative path to file: %v", err)
	}

	contentID := fmt.Sprintf("file:///mnt/onboard/%s", relPath)

	return contentID, nil
}

func contentIDToImageID(contentID string) string {
	imageID := contentID

	imageID = strings.Replace(imageID, " ", "_", -1)
	imageID = strings.Replace(imageID, "/", "_", -1)
	imageID = strings.Replace(imageID, ":", "_", -1)
	imageID = strings.Replace(imageID, ".", "_", -1)

	return imageID
}

func updateSeriesMeta(db *sql.DB, imageID, series string, seriesNumber float64) (int64, error) {
	res, err := db.Exec("UPDATE content SET Series=?, SeriesNumber=? WHERE ImageID=?", sql.NullString{
		String: series,
		Valid:  series != "",
	}, sql.NullString{
		String: fmt.Sprintf("%v", seriesNumber),
		Valid:  seriesNumber > 0,
	}, imageID)

	if err != nil {
		return 0, err
	}

	return res.RowsAffected()
}

func getEPUBMeta(path string) (string, float64, error) {
	zr, err := zip.OpenReader(path)
	if err != nil {
		return "", 0, err
	}

	zfs := zipfs.New(zr, "epub")
	rsk, err := zfs.Open("/META-INF/container.xml")
	if err != nil {
		return "", 0, err
	}
	defer rsk.Close()

	container := etree.NewDocument()
	_, err = container.ReadFrom(rsk)
	if err != nil {
		return "", 0, err
	}

	rootfile := ""
	for _, e := range container.FindElements("//rootfiles/rootfile[@full-path]") {
		rootfile = e.SelectAttrValue("full-path", "")
	}

	if rootfile == "" {
		return "", 0, errors.New("Cannot parse container")
	}

	rrsk, err := zfs.Open("/" + rootfile)
	if err != nil {
		return "", 0, err
	}
	defer rrsk.Close()

	opf := etree.NewDocument()
	_, err = opf.ReadFrom(rrsk)
	if err != nil {
		return "", 0, err
	}

	var series string
	for _, e := range opf.FindElements("//meta[@name='calibre:series']") {
		series = e.SelectAttrValue("content", "")
		break
	}

	var seriesNumber float64
	for _, e := range opf.FindElements("//meta[@name='calibre:series_index']") {
		i, err := strconv.ParseFloat(e.SelectAttrValue("content", "0"), 64)
		if err == nil {
			seriesNumber = i
			break
		}
	}

	return series, seriesNumber, nil
}

func updateSeriesMetaFromEPUB(s *spinner.Spinner, db *sql.DB, koboPath, epubPath string) (int64, error) {
	series, seriesNumber, err := getEPUBMeta(epubPath)
	if err != nil {
		return 0, err
	}

	cid, err := pathToContentID(koboPath, epubPath)
	if err != nil {
		return 0, err
	}

	iid := contentIDToImageID(cid)

	if s != nil {
		s.Suffix = fmt.Sprintf(" UPDATE %s => [%s %v]\n", iid, series, seriesNumber)
	} else {
		fmt.Printf("INFO: UPDATE %s => [%s %v]\n", iid, series, seriesNumber)
	}

	return updateSeriesMeta(db, iid, series, seriesNumber)
}

func loadKoboDB(koboPath string) (*sql.DB, error) {
	koboDBPath := filepath.Join(koboPath, ".kobo/KoboReader.sqlite")
	koboDBBackupPath := filepath.Join(koboPath, "KoboReader.sqlite.bak")

	if _, err := os.Stat(koboDBPath); os.IsNotExist(err) {
		return nil, fmt.Errorf("Kobo database %s does not exist", koboDBPath)
	}

	copyFile(koboDBPath, koboDBBackupPath)

	return sql.Open("sqlite3", koboDBPath)
}

func main() {
	if len(os.Args) < 2 || len(os.Args) > 3 {
		fmt.Printf("USAGE: %s KOBO_ROOT_PATH [EPUB_PATH]\n", filepath.Base(os.Args[0]))
		os.Exit(1)
	}

	koboPath, err := filepath.Abs(os.Args[1])
	if err != nil {
		fmt.Printf("FATAL: Could resolve Kobo path %s: %v\n", os.Args[1], err)
		os.Exit(1)
	}

	if _, err := os.Stat(filepath.Join(koboPath, ".kobo")); os.IsNotExist(err) {
		fmt.Printf("FATAL: %s is not a valid path to a Kobo eReader.\n", os.Args[1])
		fmt.Printf("USAGE: %s KOBO_ROOT_PATH [EPUB_PATH]\n", filepath.Base(os.Args[0]))
		os.Exit(1)
	}

	if len(os.Args) == 3 {
		epubPath, err := filepath.Abs(os.Args[2])
		if err != nil {
			fmt.Printf("FATAL: Could resolve ePub path %s: %v\n", os.Args[2], err)
			os.Exit(1)
		}

		if !strings.HasPrefix(epubPath, koboPath) {
			fmt.Printf("FATAL: ePub file not in the specified Kobo path.\n")
			os.Exit(1)
		}

		db, err := loadKoboDB(koboPath)
		if err != nil {
			fmt.Printf("FATAL: Could not open Kobo database: %v\n", err)
			os.Exit(1)
		}

		ra, err := updateSeriesMetaFromEPUB(nil, db, koboPath, epubPath)
		if err != nil {
			fmt.Printf("ERROR: Could not update series metadata: %v\n", err)
			os.Exit(1)
		} else if ra < 1 {
			fmt.Printf("ERROR: Could not update series metadata: no database entry for book. Please let the kobo import the book before using this tool.\n")
		} else if ra > 1 {
			fmt.Printf("WARN: More than 1 match for book in database.\n")
		}
	} else {
		s := spinner.New(spinner.CharSets[11], 100*time.Millisecond)
		s.Start()

		s.Suffix = " Opening Kobo database"
		db, err := loadKoboDB(koboPath)
		if err != nil {
			s.Stop()
			fmt.Printf("FATAL: Could not open Kobo database: %v\n", err)
			os.Exit(1)
		}

		s.Suffix = " Searching for epub files"
		matches, err := zglob.Glob(filepath.Join(koboPath, "**/*.epub"))
		if err != nil {
			s.Stop()
			fmt.Printf("FATAL: Error searching for epub files: %v\n", err)
			os.Exit(1)
		}

		s.Suffix = " Filtering epub files"
		epubs := []string{}
		for _, match := range matches {
			if strings.HasPrefix(filepath.Base(match), ".") {
				continue
			}
			epubs = append(epubs, match)
		}

		s.Stop()
		fmt.Printf("INFO: Found %v epub files\n", len(epubs))
		s.Start()

		errcount := 0
		for _, epub := range epubs {
			ra, err := updateSeriesMetaFromEPUB(s, db, koboPath, epub)

			b, err := filepath.Rel(koboPath, epub)
			if err != nil {
				b = filepath.Base(epub)
			}

			if err != nil {
				s.Stop()
				fmt.Printf("ERROR: Could not update series metadata for %s: %v\n", b, err)
				s.Start()
				errcount++
			} else if ra < 1 {
				s.Stop()
				fmt.Printf("ERROR: Could not update series metadata for %s: no entry in database for book. Please let the kobo import the book before using this tool.\n", b)
				s.Start()
				errcount++
			} else if ra > 1 {
				s.Stop()
				fmt.Printf("WARN: More than 1 match for book in database: %s.\n", b)
				s.Start()
			}
		}

		time.Sleep(time.Second)
		s.Stop()
		fmt.Println()
		fmt.Printf("INFO: Finished updating metadata. %v books processed. %v errors.\n", len(epubs), errcount)
	}
}
<commit_msg>Update<commit_after>package main

import (
	"archive/zip"
	"database/sql"
	"errors"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"github.com/briandowns/spinner"
	"github.com/mattn/go-zglob"

	"golang.org/x/tools/godoc/vfs/zipfs"

	"github.com/beevik/etree"
	_ "github.com/mattn/go-sqlite3"
)

var version = "dev"

func copyFile(src, dst string) error {
	in, err := os.Open(src)
	if err != nil {
		return err
	}
	defer in.Close()

	out, err := os.Create(dst)
	if err != nil {
		return err
	}
	defer out.Close()

	_, err = io.Copy(out, in)
	if err != nil {
		return err
	}
	return out.Close()
}

func pathToContentID(koboPath, path string) (imageID string, err error) {
	relPath, err := filepath.Rel(koboPath, path)
	if err != nil {
		return "", fmt.Errorf("could not get relative path to file: %v", err)
	}

	contentID := fmt.Sprintf("file:///mnt/onboard/%s", relPath)

	return contentID, nil
}

func contentIDToImageID(contentID string) string {
	imageID := contentID

	imageID = strings.Replace(imageID, " ", "_", -1)
	imageID = strings.Replace(imageID, "/", "_", -1)
	imageID = strings.Replace(imageID, ":", "_", -1)
	imageID = strings.Replace(imageID, ".", "_", -1)

	return imageID
}

func updateSeriesMeta(db *sql.DB, imageID, series string, seriesNumber float64) (int64, error) {
	res, err := db.Exec("UPDATE content SET Series=?, SeriesNumber=? WHERE ImageID=?", sql.NullString{
		String: series,
		Valid:  series != "",
	}, sql.NullString{
		String: fmt.Sprintf("%v", seriesNumber),
		Valid:  seriesNumber > 0,
	}, imageID)

	if err != nil {
		return 0, err
	}

	return res.RowsAffected()
}

func getEPUBMeta(path string) (string, float64, error) {
	zr, err := zip.OpenReader(path)
	if err != nil {
		return "", 0, err
	}

	zfs := zipfs.New(zr, "epub")
	rsk, err := zfs.Open("/META-INF/container.xml")
	if err != nil {
		return "", 0, err
	}
	defer rsk.Close()

	container := etree.NewDocument()
	_, err = container.ReadFrom(rsk)
	if err != nil {
		return "", 0, err
	}

	rootfile := ""
	for _, e := range container.FindElements("//rootfiles/rootfile[@full-path]") {
		rootfile = e.SelectAttrValue("full-path", "")
	}

	if rootfile == "" {
		return "", 0, errors.New("Cannot parse container")
	}

	rrsk, err := zfs.Open("/" + rootfile)
	if err != nil {
		return "", 0, err
	}
	defer rrsk.Close()

	opf := etree.NewDocument()
	_, err = opf.ReadFrom(rrsk)
	if err != nil {
		return "", 0, err
	}

	var series string
	for _, e := range opf.FindElements("//meta[@name='calibre:series']") {
		series = e.SelectAttrValue("content", "")
		break
	}

	var seriesNumber float64
	for _, e := range opf.FindElements("//meta[@name='calibre:series_index']") {
		i, err := strconv.ParseFloat(e.SelectAttrValue("content", "0"), 64)
		if err == nil {
			seriesNumber = i
			break
		}
	}

	return series, seriesNumber, nil
}

func updateSeriesMetaFromEPUB(s *spinner.Spinner, db *sql.DB, koboPath, epubPath string) (int64, error) {
	series, seriesNumber, err := getEPUBMeta(epubPath)
	if err != nil {
		return 0, err
	}

	cid, err := pathToContentID(koboPath, epubPath)
	if err != nil {
		return 0, err
	}

	iid := contentIDToImageID(cid)

	if s != nil {
		s.Suffix = fmt.Sprintf(" UPDATE %s => [%s %v]\n", iid, series, seriesNumber)
	} else {
		fmt.Printf("INFO: UPDATE %s => [%s %v]\n", iid, series, seriesNumber)
	}

	return updateSeriesMeta(db, iid, series, seriesNumber)
}

func loadKoboDB(koboPath string) (*sql.DB, error) {
	koboDBPath := filepath.Join(koboPath, ".kobo/KoboReader.sqlite")
	koboDBBackupPath := filepath.Join(koboPath, "KoboReader.sqlite.bak")

	if _, err := os.Stat(koboDBPath); os.IsNotExist(err) {
		return nil, fmt.Errorf("Kobo database %s does not exist", koboDBPath)
	}

	copyFile(koboDBPath, koboDBBackupPath)

	return sql.Open("sqlite3", koboDBPath)
}

func main() {
	if len(os.Args) < 2 || len(os.Args) > 3 {
		fmt.Printf("USAGE: %s KOBO_ROOT_PATH [EPUB_PATH]\n", filepath.Base(os.Args[0]))
		os.Exit(1)
	}

	koboPath, err := filepath.Abs(os.Args[1])
	if err != nil {
		fmt.Printf("FATAL: Could resolve Kobo path %s: %v\n", os.Args[1], err)
		os.Exit(1)
	}

	if _, err := os.Stat(filepath.Join(koboPath, ".kobo")); os.IsNotExist(err) {
		fmt.Printf("FATAL: %s is not a valid path to a Kobo eReader.\n", os.Args[1])
		fmt.Printf("USAGE: %s KOBO_ROOT_PATH [EPUB_PATH]\n", filepath.Base(os.Args[0]))
		os.Exit(1)
	}

	if len(os.Args) == 3 {
		epubPath, err := filepath.Abs(os.Args[2])
		if err != nil {
			fmt.Printf("FATAL: Could resolve ePub path %s: %v\n", os.Args[2], err)
			os.Exit(1)
		}

		if !strings.HasPrefix(epubPath, koboPath) {
			fmt.Printf("FATAL: ePub file not in the specified Kobo path.\n")
			os.Exit(1)
		}

		db, err := loadKoboDB(koboPath)
		if err != nil {
			fmt.Printf("FATAL: Could not open Kobo database: %v\n", err)
			os.Exit(1)
		}

		ra, err := updateSeriesMetaFromEPUB(nil, db, koboPath, epubPath)
		if err != nil {
			fmt.Printf("ERROR: Could not update series metadata: %v\n", err)
			os.Exit(1)
		} else if ra < 1 {
			fmt.Printf("ERROR: Could not update series metadata: no database entry for book. Please let the kobo import the book before using this tool.\n")
		} else if ra > 1 {
			fmt.Printf("WARN: More than 1 match for book in database.\n")
		}
	} else {
		s := spinner.New(spinner.CharSets[11], 100*time.Millisecond)
		s.Start()

		s.Suffix = " Opening Kobo database"
		db, err := loadKoboDB(koboPath)
		if err != nil {
			s.Stop()
			fmt.Printf("FATAL: Could not open Kobo database: %v\n", err)
			os.Exit(1)
		}

		s.Suffix = " Searching for epub files"
		matches, err := zglob.Glob(filepath.Join(koboPath, "**/*.epub"))
		if err != nil {
			s.Stop()
			fmt.Printf("FATAL: Error searching for epub files: %v\n", err)
			os.Exit(1)
		}

		s.Suffix = " Filtering epub files"
		epubs := []string{}
		for _, match := range matches {
			if strings.HasPrefix(filepath.Base(match), ".") {
				continue
			}
			epubs = append(epubs, match)
		}

		s.Stop()
		fmt.Printf("INFO: Found %v epub files\n\n", len(epubs))
		s.Start()

		errcount := 0
		for _, epub := range epubs {
			ra, err := updateSeriesMetaFromEPUB(s, db, koboPath, epub)

			b, err := filepath.Rel(koboPath, epub)
			if err != nil {
				b = filepath.Base(epub)
			}

			if err != nil {
				s.Stop()
				fmt.Printf("ERROR: Could not update series metadata for %s: %v\n", b, err)
				s.Start()
				errcount++
			} else if ra < 1 {
				s.Stop()
				fmt.Printf("ERROR: Could not update series metadata for %s: no entry in database for book. Please let the kobo import the book before using this tool.\n", b)
				s.Start()
				errcount++
			} else if ra > 1 {
				s.Stop()
				fmt.Printf("WARN: More than 1 match for book in database: %s.\n", b)
				s.Start()
			}
		}

		time.Sleep(time.Second)
		s.Stop()
		fmt.Println()
		fmt.Printf("INFO: Finished updating metadata. %v books processed. %v errors.\n", len(epubs), errcount)
	}
}
<|endoftext|>