<commit_before>package elasticorm

import (
	"context"
	"encoding/json"
	"fmt"
	"reflect"
	"strings"

	"github.com/pkg/errors"
	elastic "gopkg.in/olivere/elastic.v5"
)

// NewDatastore returns a fresh instance of an elasticorm datastore
func NewDatastore(esc *elastic.Client, options ...DatastoreOptFunc) (*Datastore, error) {
	ds := &Datastore{
		elasticClient: esc,
		Ctx:           context.Background(),
	}
	var err error
	for _, opt := range options {
		err = opt(ds)
	}
	return ds, err
}

// DatastoreOptFunc is used as a parameter to NewDatastore and provides a way of configuration
type DatastoreOptFunc func(*Datastore) error

// QueryOptFunc is used as a parameter
type QueryOptFunc func(*elastic.SearchService) error

// Datstore is an instance to communicate easy with elasticsearch
// It is meant to be used for one struct and helps storing and retrieving it in/from elasticsearch
// It leverages the great elastic package from olivere
type Datastore struct {
	elasticClient   *elastic.Client
	Ctx             context.Context
	indexName       string
	goType          reflect.Type
	idFieldName     string          // the name of the structs field to store the ID
	typeName        string          // in elasticsearch
	indexDefinition IndexDefinition // in elasticsearch
}

// EnsureIndexExists checks wether the needed index for this datastore exists. It it doesn't it gets created
// the name of the datastore is determined by the structs name (+ plural s)
func (ds *Datastore) EnsureIndexExists() error {
	if ds.indexName == `` {
		return errors.New(`EnsureIndexExists failed, because no index name is defined`)
	}
	exists, err := ds.elasticClient.
		IndexExists(ds.indexName).
		Do(ds.Ctx)

	if exists || err != nil {
		return err
	}

	err = ds.createIndex()
	if err != nil {
		return err
	}
	ds.Refresh()
	return nil
}

func (ds *Datastore) createIndex() error {
	JSON, err := json.Marshal(ds.indexDefinition)
	if err != nil {
		return err
	}
	ack, err := ds.elasticClient.
		CreateIndex(ds.indexName).
		BodyString(string(JSON)).
		Do(ds.Ctx)
	if err != nil || !ack.Acknowledged {
		return errors.Wrap(
			err, fmt.Sprintf("creating elasticsearch index %s failed - %s", ds.indexName, string(JSON)),
		)
	}
	return nil
}

func (ds *Datastore) EnsureIndexDoesntExist() error {
	if ds.indexName == `` {
		return errors.New(`EnsureIndexDoesntExists failed, because no index name is defined`)
	}
	exists, err := ds.elasticClient.
		IndexExists(ds.indexName).
		Do(ds.Ctx)

	if !exists || err != nil {
		return err
	}
	res, err := ds.elasticClient.
		DeleteIndex(ds.indexName).
		Do(context.Background())
	if err != nil || !res.Acknowledged {
		return errors.Wrap(
			err, fmt.Sprintf("deleting elasticsearch index %s failed", ds.indexName),
		)
	}
	return nil
}

func (ds *Datastore) CleanUp() error {
	err := ds.EnsureIndexDoesntExist()
	if err != nil {
		return err
	}
	err = ds.EnsureIndexExists()
	if err != nil {
		return err
	}
	return ds.Refresh()
}

func (ds *Datastore) Refresh() error {
	_, err := ds.elasticClient.Refresh().Do(ds.Ctx)
	return err
}

// ForStruct generates a DatastoreOptFunc
// It is used to generate a default mapping, type name and index name by analysing a provided struct
func ForStruct(i interface{}) DatastoreOptFunc {
	return func(ds *Datastore) error {
		ds.goType = reflect.TypeOf(i)
		typeName, err := typeNameFromStruct(i)
		if err != nil {
			return err
		}
		ds.typeName = typeName
		ds.indexName = typeName + `s`
		indexDefinition, err := NewIndexDefinition(
			AddMappingFromStruct(ds.typeName, i),
		)
		if err != nil {
			return err
		}
		ds.indexDefinition = indexDefinition
		// TODO configure ID
		ds.idFieldName = `ID`
		return nil
	}
}

func typeNameFromStruct(i interface{}) (string, error) {
	typeName := getType(i)
	if typeName == `` {
		return ``, errors.New(`Could not determine type name from struct`)
	}

	typeName = strings.ToLower(typeName)
	typeName = strings.TrimLeft(typeName, `*`)
	return typeName, nil
}

func getType(myvar interface{}) string {
	return nameOfType(reflect.TypeOf(myvar))
}

func nameOfType(t reflect.Type) string {
	if t.Kind() == reflect.Ptr {
		return "*" + t.Elem().Name()
	} else {
		return t.Name()
	}
}

func (ds *Datastore) Create(o interface{}) error {
	err := ds.isSaveableType(o)
	if err != nil {
		return err
	}
	put, err := ds.elasticClient.Index().
		Index(ds.indexName).
		Type(ds.typeName).
		BodyJson(o).
		Do(ds.Ctx)

	if err != nil {
		return err
	}
	if !put.Created {
		return ErrCreationFailed
	}
	return ds.setID(o, put.Id)
}

func (ds *Datastore) isSaveableType(i interface{}) error {
	rt := reflect.TypeOf(i)
	if nameOfType(rt) != nameOfType(ds.goType) {
		return errors.Wrapf(
			ErrInvalidType,
			`create failed for %s, expected %s`, nameOfType(rt), nameOfType(ds.goType),
		)
	}
	rv := reflect.ValueOf(i)
	if rv.Kind() != reflect.Ptr || rv.IsNil() {
		return errors.Wrap(ErrInvalidType, `no pointer given`)
	}
	return nil
}

func (ds *Datastore) Find(ID string, result interface{}) error {
	err := ds.isSaveableType(result)
	if err != nil {
		return err
	}
	res, err := ds.elasticClient.Get().
		Index(ds.indexName).
		Type(ds.typeName).
		Id(ID).
		Do(ds.Ctx)

	if elastic.IsNotFound(err) || (res != nil && !res.Found) {
		return ErrNotFound
	}
	if err != nil {
		return err
	}

	return ds.decodeElasticResponse(res.Source, res.Id, result)
}

func (ds *Datastore) Update(o interface{}) error {
	err := ds.isSaveableType(o)
	if err != nil {
		return err
	}
	ID, err := ds.getID(o)
	if err != nil {
		return err
	}
	if ID == `` {
		return errors.New(`can't save struct with empty ID`)
	}
	_, err = ds.elasticClient.Update().
		Index(ds.indexName).
		Type(ds.typeName).
		Id(ID).
		Doc(o).
		Do(ds.Ctx)

	return err
}

func (ds *Datastore) FindOneBy(fieldName string, value interface{}, result interface{}, opts ...QueryOptFunc) error {
	elasticFieldName, err := ds.indexDefinition.elasticFieldName(ds.typeName, fieldName)
	if err != nil {
		return err
	}
	q := ds.elasticClient.Search().
		Index(ds.indexName).
		Query(elastic.NewBoolQuery().Filter(elastic.NewTermQuery(elasticFieldName, value))).
		From(0).Size(1)

	for _, opt := range opts {
		err := opt(q)
		if err != nil {
			return err
		}
	}

	res, err := q.Do(ds.Ctx)
	if err != nil {
		return err
	}

	if res.TotalHits() < 1 {
		return ErrNotFound
	}
	return ds.decodeElasticResponse(res.Hits.Hits[0].Source, res.Hits.Hits[0].Id, result)
}

func (ds *Datastore) FindAll(results interface{}, opts ...QueryOptFunc) error {
	q := ds.elasticClient.Search().
		Index(ds.indexName).
		Query(elastic.NewMatchAllQuery())

	for _, opt := range opts {
		err := opt(q)
		if err != nil {
			return err
		}
	}

	res, err := q.Do(ds.Ctx)
	if err != nil {
		return err
	}
	return ds.decodeElasticResponses(res.Hits.Hits, results)
}

func (ds *Datastore) SetSorting(fieldName string, order string) QueryOptFunc {
	return func(srv *elastic.SearchService) error {
		if order != `asc` && order != `desc` {
			return errors.New(`sorting order must be asc or desc`)
		}
		elasticFieldName, err := ds.indexDefinition.elasticFieldName(ds.typeName, fieldName)
		// TODO loosen coupling with indexDefinition
		if err != nil {
			return err
		}
		srv.Sort(elasticFieldName+`.raw`, order == `asc`)
		return nil
	}
}

func (ds *Datastore) FilterByField(fieldName string, value interface{}) QueryOptFunc {
	return func(srv *elastic.SearchService) error {
		// TODO loosen coupling with indexDefinition
		elasticFieldName, err := ds.indexDefinition.elasticFieldName(ds.typeName, fieldName)
		if err != nil {
			return err
		}
		q := elastic.NewBoolQuery()
		q.Filter(elastic.NewTermQuery(elasticFieldName, value))
		srv.Query(q)
		return nil
	}
}

func (ds *Datastore) Offset(offset int) QueryOptFunc {
	return func(srv *elastic.SearchService) error {
		srv.From(offset)
		return nil
	}
}

func (ds *Datastore) Limit(limit int) QueryOptFunc {
	return func(srv *elastic.SearchService) error {
		srv.Size(limit)
		return nil
	}
}

type BoundingBox struct {
	Top    float64
	Left   float64
	Bottom float64
	Right  float64
}

func NewBoundingBox(top, right, bottom, left float64) BoundingBox {
	return BoundingBox{
		Top:    top,
		Right:  right,
		Bottom: bottom,
		Left:   left,
	}
}

func (ds *Datastore) FindByGeoBoundingBox(fieldName string, box BoundingBox, results interface{}) error {
	elasticFieldName, err := ds.indexDefinition.elasticFieldName(ds.typeName, fieldName)
	if err != nil {
		return err
	}
	query := elastic.NewGeoBoundingBoxQuery(elasticFieldName).
		TopLeft(box.Top, box.Left).
		BottomRight(box.Bottom, box.Right)

	res, err := ds.elasticClient.Search().
		Index(ds.indexName).
		// TODO query type
		Query(elastic.NewBoolQuery().Filter(query)).
		Do(ds.Ctx)

	if err != nil {
		return err
	}

	return ds.decodeElasticResponses(res.Hits.Hits, results)
}

func (ds *Datastore) FindByGeoDistance(fieldName string, lat float64, lon float64, distance string, results interface{}) error {
	elasticFieldName, err := ds.indexDefinition.elasticFieldName(ds.typeName, fieldName)
	if err != nil {
		return err
	}
	query := elastic.NewGeoDistanceQuery(elasticFieldName).
		Lat(lat).
		Lon(lon).
		Distance(distance)

	res, err := ds.elasticClient.Search().
		Index(ds.indexName).
		// TODO query type
		Query(query).
		SortBy(elastic.NewGeoDistanceSort(elasticFieldName).Point(lat, lon)).
		Do(ds.Ctx)

	if err != nil {
		return err
	}

	return ds.decodeElasticResponses(res.Hits.Hits, results)
}

func (ds *Datastore) decodeElasticResponses(hits []*elastic.SearchHit, results interface{}) error {
	resultsv := reflect.ValueOf(results)
	if resultsv.Kind() != reflect.Ptr || resultsv.Elem().Kind() != reflect.Slice {
		return errors.New("result argument must be a slice address")
	}

	slicev := resultsv.Elem()
	elemt := slicev.Type().Elem()

	for _, hit := range hits {
		elemp := reflect.New(elemt)
		err := ds.decodeElasticResponse(hit.Source, hit.Id, elemp.Interface())
		if err != nil {
			return err
		}
		slicev = reflect.Append(slicev, elemp.Elem())
	}

	resultsv.Elem().Set(slicev.Slice(0, len(hits)))
	return nil
}

func (ds *Datastore) setID(o interface{}, ID string) error {
	eo := reflect.ValueOf(o).Elem()
	if eo.Kind() != reflect.Struct {
		return errors.Wrap(ErrInvalidType, `setID failed`)
	}
	idField := eo.FieldByName(ds.idFieldName)
	if !idField.IsValid() || !idField.CanSet() || idField.Kind() != reflect.String {
		return ErrInvalidIDField
	}
	idField.SetString(ID)
	return nil
}

func (ds *Datastore) getID(o interface{}) (string, error) {
	eo := reflect.ValueOf(o).Elem()
	if eo.Kind() != reflect.Struct {
		return ``, errors.Wrap(ErrInvalidType, `getID failed`)
	}
	idField := eo.FieldByName(ds.idFieldName)
	if !idField.IsValid() || idField.Kind() != reflect.String {
		return ``, ErrInvalidIDField
	}
	return idField.String(), nil
}

func (ds *Datastore) decodeElasticResponse(source *json.RawMessage, ID string, o interface{}) error {
	err := json.Unmarshal(*source, o)
	if err != nil {
		return err
	}
	ds.setID(o, ID)
	return nil
}
<commit_msg>[FEATURE] configure ID field<commit_after>package elasticorm

import (
	"context"
	"encoding/json"
	"fmt"
	"reflect"
	"strings"

	"github.com/pkg/errors"
	elastic "gopkg.in/olivere/elastic.v5"
)

// NewDatastore returns a fresh instance of an elasticorm datastore
func NewDatastore(esc *elastic.Client, options ...DatastoreOptFunc) (*Datastore, error) {
	ds := &Datastore{
		elasticClient: esc,
		Ctx:           context.Background(),
	}
	var err error
	for _, opt := range options {
		err = opt(ds)
	}
	return ds, err
}

// DatastoreOptFunc is used as a parameter to NewDatastore and provides a way of configuration
type DatastoreOptFunc func(*Datastore) error

// QueryOptFunc is used as a parameter
type QueryOptFunc func(*elastic.SearchService) error

// Datstore is an instance to communicate easy with elasticsearch
// It is meant to be used for one struct and helps storing and retrieving it in/from elasticsearch
// It leverages the great elastic package from olivere
type Datastore struct {
	elasticClient   *elastic.Client
	Ctx             context.Context
	indexName       string
	goType          reflect.Type
	idFieldName     string          // the name of the structs field to store the ID
	typeName        string          // in elasticsearch
	indexDefinition IndexDefinition // in elasticsearch
}

// EnsureIndexExists checks wether the needed index for this datastore exists. It it doesn't it gets created
// the name of the datastore is determined by the structs name (+ plural s)
func (ds *Datastore) EnsureIndexExists() error {
	if ds.indexName == `` {
		return errors.New(`EnsureIndexExists failed, because no index name is defined`)
	}
	exists, err := ds.elasticClient.
		IndexExists(ds.indexName).
		Do(ds.Ctx)

	if exists || err != nil {
		return err
	}

	err = ds.createIndex()
	if err != nil {
		return err
	}
	ds.Refresh()
	return nil
}

func (ds *Datastore) createIndex() error {
	JSON, err := json.Marshal(ds.indexDefinition)
	if err != nil {
		return err
	}
	ack, err := ds.elasticClient.
		CreateIndex(ds.indexName).
		BodyString(string(JSON)).
		Do(ds.Ctx)
	if err != nil || !ack.Acknowledged {
		return errors.Wrap(
			err, fmt.Sprintf("creating elasticsearch index %s failed - %s", ds.indexName, string(JSON)),
		)
	}
	return nil
}

func (ds *Datastore) EnsureIndexDoesntExist() error {
	if ds.indexName == `` {
		return errors.New(`EnsureIndexDoesntExists failed, because no index name is defined`)
	}
	exists, err := ds.elasticClient.
		IndexExists(ds.indexName).
		Do(ds.Ctx)

	if !exists || err != nil {
		return err
	}
	res, err := ds.elasticClient.
		DeleteIndex(ds.indexName).
		Do(context.Background())
	if err != nil || !res.Acknowledged {
		return errors.Wrap(
			err, fmt.Sprintf("deleting elasticsearch index %s failed", ds.indexName),
		)
	}
	return nil
}

func (ds *Datastore) CleanUp() error {
	err := ds.EnsureIndexDoesntExist()
	if err != nil {
		return err
	}
	err = ds.EnsureIndexExists()
	if err != nil {
		return err
	}
	return ds.Refresh()
}

func (ds *Datastore) Refresh() error {
	_, err := ds.elasticClient.Refresh().Do(ds.Ctx)
	return err
}

// ForStruct generates a DatastoreOptFunc
// It is used to generate a default mapping, type name and index name by analysing a provided struct
func ForStruct(i interface{}) DatastoreOptFunc {
	return func(ds *Datastore) error {
		ds.goType = reflect.TypeOf(i)
		typeName, err := typeNameFromStruct(i)
		if err != nil {
			return err
		}
		ds.typeName = typeName
		ds.indexName = typeName + `s`
		indexDefinition, err := NewIndexDefinition(
			AddMappingFromStruct(ds.typeName, i),
		)
		if err != nil {
			return err
		}
		ds.indexDefinition = indexDefinition
		if ds.idFieldName == "" {
			ds.idFieldName = "ID"
		}
		return nil
	}
}

func WithIDField(name string) DatastoreOptFunc {
	return func(ds *Datastore) error {
		ds.idFieldName = name
		return nil
	}
}

func typeNameFromStruct(i interface{}) (string, error) {
	typeName := getType(i)
	if typeName == `` {
		return ``, errors.New(`Could not determine type name from struct`)
	}

	typeName = strings.ToLower(typeName)
	typeName = strings.TrimLeft(typeName, `*`)
	return typeName, nil
}

func getType(myvar interface{}) string {
	return nameOfType(reflect.TypeOf(myvar))
}

func nameOfType(t reflect.Type) string {
	if t.Kind() == reflect.Ptr {
		return "*" + t.Elem().Name()
	} else {
		return t.Name()
	}
}

func (ds *Datastore) Create(o interface{}) error {
	err := ds.isSaveableType(o)
	if err != nil {
		return err
	}
	put, err := ds.elasticClient.Index().
		Index(ds.indexName).
		Type(ds.typeName).
		BodyJson(o).
		Do(ds.Ctx)

	if err != nil {
		return err
	}
	if !put.Created {
		return ErrCreationFailed
	}
	return ds.setID(o, put.Id)
}

func (ds *Datastore) isSaveableType(i interface{}) error {
	rt := reflect.TypeOf(i)
	if nameOfType(rt) != nameOfType(ds.goType) {
		return errors.Wrapf(
			ErrInvalidType,
			`create failed for %s, expected %s`, nameOfType(rt), nameOfType(ds.goType),
		)
	}
	rv := reflect.ValueOf(i)
	if rv.Kind() != reflect.Ptr || rv.IsNil() {
		return errors.Wrap(ErrInvalidType, `no pointer given`)
	}
	return nil
}

func (ds *Datastore) Find(ID string, result interface{}) error {
	err := ds.isSaveableType(result)
	if err != nil {
		return err
	}
	res, err := ds.elasticClient.Get().
		Index(ds.indexName).
		Type(ds.typeName).
		Id(ID).
		Do(ds.Ctx)

	if elastic.IsNotFound(err) || (res != nil && !res.Found) {
		return ErrNotFound
	}
	if err != nil {
		return err
	}

	return ds.decodeElasticResponse(res.Source, res.Id, result)
}

func (ds *Datastore) Update(o interface{}) error {
	err := ds.isSaveableType(o)
	if err != nil {
		return err
	}
	ID, err := ds.getID(o)
	if err != nil {
		return err
	}
	if ID == `` {
		return errors.New(`can't save struct with empty ID`)
	}
	_, err = ds.elasticClient.Update().
		Index(ds.indexName).
		Type(ds.typeName).
		Id(ID).
		Doc(o).
		Do(ds.Ctx)

	return err
}

func (ds *Datastore) FindOneBy(fieldName string, value interface{}, result interface{}, opts ...QueryOptFunc) error {
	elasticFieldName, err := ds.indexDefinition.elasticFieldName(ds.typeName, fieldName)
	if err != nil {
		return err
	}
	q := ds.elasticClient.Search().
		Index(ds.indexName).
		Query(elastic.NewBoolQuery().Filter(elastic.NewTermQuery(elasticFieldName, value))).
		From(0).Size(1)

	for _, opt := range opts {
		err := opt(q)
		if err != nil {
			return err
		}
	}

	res, err := q.Do(ds.Ctx)
	if err != nil {
		return err
	}

	if res.TotalHits() < 1 {
		return ErrNotFound
	}
	return ds.decodeElasticResponse(res.Hits.Hits[0].Source, res.Hits.Hits[0].Id, result)
}

func (ds *Datastore) FindAll(results interface{}, opts ...QueryOptFunc) error {
	q := ds.elasticClient.Search().
		Index(ds.indexName).
		Query(elastic.NewMatchAllQuery())

	for _, opt := range opts {
		err := opt(q)
		if err != nil {
			return err
		}
	}

	res, err := q.Do(ds.Ctx)
	if err != nil {
		return err
	}
	return ds.decodeElasticResponses(res.Hits.Hits, results)
}

func (ds *Datastore) SetSorting(fieldName string, order string) QueryOptFunc {
	return func(srv *elastic.SearchService) error {
		if order != `asc` && order != `desc` {
			return errors.New(`sorting order must be asc or desc`)
		}
		elasticFieldName, err := ds.indexDefinition.elasticFieldName(ds.typeName, fieldName)
		// TODO loosen coupling with indexDefinition
		if err != nil {
			return err
		}
		srv.Sort(elasticFieldName+`.raw`, order == `asc`)
		return nil
	}
}

func (ds *Datastore) FilterByField(fieldName string, value interface{}) QueryOptFunc {
	return func(srv *elastic.SearchService) error {
		// TODO loosen coupling with indexDefinition
		elasticFieldName, err := ds.indexDefinition.elasticFieldName(ds.typeName, fieldName)
		if err != nil {
			return err
		}
		q := elastic.NewBoolQuery()
		q.Filter(elastic.NewTermQuery(elasticFieldName, value))
		srv.Query(q)
		return nil
	}
}

func (ds *Datastore) Offset(offset int) QueryOptFunc {
	return func(srv *elastic.SearchService) error {
		srv.From(offset)
		return nil
	}
}

func (ds *Datastore) Limit(limit int) QueryOptFunc {
	return func(srv *elastic.SearchService) error {
		srv.Size(limit)
		return nil
	}
}

type BoundingBox struct {
	Top    float64
	Left   float64
	Bottom float64
	Right  float64
}

func NewBoundingBox(top, right, bottom, left float64) BoundingBox {
	return BoundingBox{
		Top:    top,
		Right:  right,
		Bottom: bottom,
		Left:   left,
	}
}

func (ds *Datastore) FindByGeoBoundingBox(fieldName string, box BoundingBox, results interface{}) error {
	elasticFieldName, err := ds.indexDefinition.elasticFieldName(ds.typeName, fieldName)
	if err != nil {
		return err
	}
	query := elastic.NewGeoBoundingBoxQuery(elasticFieldName).
		TopLeft(box.Top, box.Left).
		BottomRight(box.Bottom, box.Right)

	res, err := ds.elasticClient.Search().
		Index(ds.indexName).
		// TODO query type
		Query(elastic.NewBoolQuery().Filter(query)).
		Do(ds.Ctx)

	if err != nil {
		return err
	}

	return ds.decodeElasticResponses(res.Hits.Hits, results)
}

func (ds *Datastore) FindByGeoDistance(fieldName string, lat float64, lon float64, distance string, results interface{}) error {
	elasticFieldName, err := ds.indexDefinition.elasticFieldName(ds.typeName, fieldName)
	if err != nil {
		return err
	}
	query := elastic.NewGeoDistanceQuery(elasticFieldName).
		Lat(lat).
		Lon(lon).
		Distance(distance)

	res, err := ds.elasticClient.Search().
		Index(ds.indexName).
		// TODO query type
		Query(query).
		SortBy(elastic.NewGeoDistanceSort(elasticFieldName).Point(lat, lon)).
		Do(ds.Ctx)

	if err != nil {
		return err
	}

	return ds.decodeElasticResponses(res.Hits.Hits, results)
}

func (ds *Datastore) decodeElasticResponses(hits []*elastic.SearchHit, results interface{}) error {
	resultsv := reflect.ValueOf(results)
	if resultsv.Kind() != reflect.Ptr || resultsv.Elem().Kind() != reflect.Slice {
		return errors.New("result argument must be a slice address")
	}

	slicev := resultsv.Elem()
	elemt := slicev.Type().Elem()

	for _, hit := range hits {
		elemp := reflect.New(elemt)
		err := ds.decodeElasticResponse(hit.Source, hit.Id, elemp.Interface())
		if err != nil {
			return err
		}
		slicev = reflect.Append(slicev, elemp.Elem())
	}

	resultsv.Elem().Set(slicev.Slice(0, len(hits)))
	return nil
}

func (ds *Datastore) setID(o interface{}, ID string) error {
	eo := reflect.ValueOf(o).Elem()
	if eo.Kind() != reflect.Struct {
		return errors.Wrap(ErrInvalidType, `setID failed`)
	}
	idField := eo.FieldByName(ds.idFieldName)
	if !idField.IsValid() || !idField.CanSet() || idField.Kind() != reflect.String {
		return ErrInvalidIDField
	}
	idField.SetString(ID)
	return nil
}

func (ds *Datastore) getID(o interface{}) (string, error) {
	eo := reflect.ValueOf(o).Elem()
	if eo.Kind() != reflect.Struct {
		return ``, errors.Wrap(ErrInvalidType, `getID failed`)
	}
	idField := eo.FieldByName(ds.idFieldName)
	if !idField.IsValid() || idField.Kind() != reflect.String {
		return ``, ErrInvalidIDField
	}
	return idField.String(), nil
}

func (ds *Datastore) decodeElasticResponse(source *json.RawMessage, ID string, o interface{}) error {
	err := json.Unmarshal(*source, o)
	if err != nil {
		return err
	}
	ds.setID(o, ID)
	return nil
}
<|endoftext|>


<commit_before>// Package datatypes is a set of packages that provides many tools which helps you treating some complex structures and primitives types.
//
// The following packages are distributed under datatypes:
//  @datatypes/collection : Provides a new single-value list struct simplifying the iteration over it
//  @datatypes/dictionary : Provides a new object type which allows to treat them like <Key,Value> list
//  @datatypes/generic : Packages which includes some functionalities to treat 'generic' objects
//  @datatypes/string : Simple package which encapsulates some basic operations over strings
package datatypes

// Unused imports help docs.
import (
	// dictionary package
	_ "github.com/jaimelopez/datatypes/dictionary"
	// generic package
	_ "github.com/jaimelopez/datatypes/generic"
	// string package
	_ "github.com/jaimelopez/datatypes/string"
	// collection package
	_ "github.com/jaimelopez/datatypes/collection"
)
<commit_msg>Test for godoc...<commit_after>// Package datatypes is a set of packages that provides many tools which helps you treating some complex structures and primitives types.
//
// The following packages are distributed under datatypes:
//  @datatypes/collection : Provides a new single-value list struct simplifying the iteration over it
//  @datatypes/dictionary : Provides a new object type which allows to treat them like <Key,Value> list
//  @datatypes/generic : Packages which includes some functionalities to treat 'generic' objects
//  @datatypes/string : Simple package which encapsulates some basic operations over strings
package datatypes

// Unused imports help docs.
import (
	// collection package
	_ "github.com/jaimelopez/datatypes/collection"
	// dictionary package
	_ "github.com/jaimelopez/datatypes/dictionary"
	// generic package
	_ "github.com/jaimelopez/datatypes/generic"
	// string package
	_ "github.com/jaimelopez/datatypes/string"
)
<|endoftext|>


<commit_before>package date

import (
	"database/sql/driver"
	"errors"
	"time"

	"github.com/reflexionhealth/vanilla/clock"
)

var Clock *clock.Source = &clock.Default

const (
	MillisecondsInSecond = 1000
	MicrosecondsInSecond = 1000 * MillisecondsInSecond
	NanosecondsInSecond  = 1000 * MicrosecondsInSecond
)

const RFC3339 = "2006-01-02"

// Parse string into desired date format i.e RFC3339
func Parse(format string, source string) (Date, error) {
	t, err := time.Parse(format, source)
	return From(t), err
}

// Date is a plain date, without time or timezone info (use time.Time for those!)
type Date struct {
	Year  int
	Month time.Month
	Day   int

	// NOTE: time.Time does not preserve timezone when gob encoded
	// location must be private to support Gob encoding
	location *time.Location
}

func At(y int, m time.Month, d int, l *time.Location) Date {
	return Date{y, m, d, l}
}

func TodayIn(loc *time.Location) Date {
	t := Clock.In(loc)
	y, m, d := t.Date()
	return Date{y, m, d, loc}
}

func TodayUtc() Date {
	return TodayIn(time.UTC)
}

func YesterdayIn(loc *time.Location) Date {
	return TodayIn(loc).PrevDay()
}

// Create a Date from a time.Time object
func From(t time.Time) Date {
	y, m, d := t.Date()
	return Date{y, m, d, t.Location()}
}

func (d Date) DaysAfter(other Date) int {
	return int(d.BeginningOfDayIn(time.UTC).Sub(other.BeginningOfDayIn(time.UTC)).Hours() / 24)
}

func (d Date) AddDays(num int) Date {
	t := time.Date(d.Year, d.Month, d.Day, 0, 0, 0, 0, d.location).AddDate(0, 0, num)
	year, month, day := t.Date()
	return Date{year, month, day, d.location}
}

func (d Date) PrevDay() Date {
	if d.Day > 1 {
		return Date{d.Year, d.Month, d.Day - 1, d.location}
	} else if d.Month > 1 {
		return Date{d.Year, d.Month - 1, DaysInMonth(d.Month-1, d.Year), d.location}
	} else {
		return Date{d.Year - 1, time.December, 31, d.location}
	}
}

func (d Date) NextDay() Date {
	if d.Day < DaysInMonth(d.Month, d.Year) {
		return Date{d.Year, d.Month, d.Day + 1, d.location}
	} else if d.Month < time.December {
		return Date{d.Year, d.Month + 1, 1, d.location}
	} else {
		return Date{d.Year + 1, time.January, 1, d.location}
	}
}

// MostRecent returns today if today is the given weekday, otherwise it returns
// the date of the last day for that weekday.
//
// It is useful to get the beginning of the week:
//
//    today.MostRecent(time.Sunday) // in the US
//    today.MostRecent(time.Monday) // officially
//
func (d Date) MostRecent(weekday time.Weekday) Date {
	t := d.BeginningOfDay()
	today := t.Weekday()
	if today == weekday {
		return d
	} else if today > weekday {
		return From(t.AddDate(0, 0, int(weekday-today)))
	} else {
		return From(t.AddDate(0, 0, int(weekday-today)-7))
	}
}

// MostSoon returns today if today is the given weekday, otherwise it returns
// the date of the upcoming day for that weekday.
func (d Date) MostSoon(weekday time.Weekday) Date {
	if d.Weekday() == weekday {
		return d
	} else if d.Weekday() < weekday {
		return d.AddDays(int(weekday - d.Weekday()))
	} else {
		return d.AddDays(7 - int(d.Weekday()-weekday))
	}
}

func (d Date) Before(other Date) bool {
	return (d.Year < other.Year ||
		(d.Year == other.Year &&
			(d.Month < other.Month ||
				(d.Month == other.Month && d.Day < other.Day))))
}

func (d Date) Equal(other Date) bool {
	return d == other
}

func (d Date) After(other Date) bool {
	return (d.Year > other.Year ||
		(d.Year == other.Year &&
			(d.Month > other.Month ||
				(d.Month == other.Month && d.Day > other.Day))))
}

func (d Date) AtLeast(other Date) bool {
	return !other.After(d)
}

func (d Date) AtMost(other Date) bool {
	return !other.Before(d)
}

func (d Date) Weekday() time.Weekday {
	return d.BeginningOfDay().Weekday()
}

func (d Date) BeginningOfDay() time.Time {
	return time.Date(d.Year, d.Month, d.Day, 0, 0, 0, 0, d.location)
}

func (d Date) BeginningOfDayIn(timezone *time.Location) time.Time {
	return time.Date(d.Year, d.Month, d.Day, 0, 0, 0, 0, timezone)
}

func (d Date) EndOfDay() time.Time {
	return time.Date(d.Year, d.Month, d.Day, 23, 59, 59, NanosecondsInSecond-1, d.location)
}

func (d Date) EndOfDayIn(timezone *time.Location) time.Time {
	return time.Date(d.Year, d.Month, d.Day, 23, 59, 59, NanosecondsInSecond-1, timezone)
}

func (d Date) String() string {
	return d.Format("2006-01-02")
}

func (d Date) Format(format string) string {
	return d.BeginningOfDayIn(time.UTC).Format(format)
}

// Implements sql.Scanner interface
func (d *Date) Scan(src interface{}) error {
	t, ok := src.(time.Time)
	if !ok {
		return errors.New("date: scan value was not a Time")
	}

	*d = From(t)
	return nil
}

// Implements sql.driver.Valuer interface
func (d Date) Value() (driver.Value, error) {
	return d.BeginningOfDayIn(time.UTC), nil
}

// Implements json.Marshaler interface
func (d Date) MarshalJSON() ([]byte, error) {
	return []byte(d.Format(`"2006-01-02"`)), nil
}

// Implements json.Unmarshaler interface
func (d *Date) UnmarshalJSON(bytes []byte) error {
	t, err := time.Parse(`"2006-01-02"`, string(bytes))
	if err != nil {
		return err
	}

	*d = From(t)
	return nil
}

// TODO: Implment gob.GobEncoder and gob.GobDecoder to preserve timezone
// func (d Date) GobEncode() ([]byte, error) {}
// func (d *Date) GobDecode(bytes []byte) error {}

func IsLeapYear(year int) bool {
	return year%4 == 0 && (year%100 != 0 || year%400 == 0)
}

var DaysInNonLeapMonth = [12]int32{31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31}

func DaysInMonth(m time.Month, year int) int {
	if m == time.February && IsLeapYear(year) {
		return 29
	} else {
		return int(DaysInNonLeapMonth[m-1])
	}
}
<commit_msg>date: fix date equal comparison to ignore location<commit_after>package date

import (
	"database/sql/driver"
	"errors"
	"time"

	"github.com/reflexionhealth/vanilla/clock"
)

var Clock *clock.Source = &clock.Default

const (
	MillisecondsInSecond = 1000
	MicrosecondsInSecond = 1000 * MillisecondsInSecond
	NanosecondsInSecond  = 1000 * MicrosecondsInSecond
)

const RFC3339 = "2006-01-02"

// Parse string into desired date format i.e RFC3339
func Parse(format string, source string) (Date, error) {
	t, err := time.Parse(format, source)
	return From(t), err
}

// Date is a plain date, without time or timezone info (use time.Time for those!)
type Date struct {
	Year  int
	Month time.Month
	Day   int

	// NOTE: time.Time does not preserve timezone when gob encoded
	// location must be private to support Gob encoding
	location *time.Location
}

func At(y int, m time.Month, d int, l *time.Location) Date {
	return Date{y, m, d, l}
}

func TodayIn(loc *time.Location) Date {
	t := Clock.In(loc)
	y, m, d := t.Date()
	return Date{y, m, d, loc}
}

func TodayUtc() Date {
	return TodayIn(time.UTC)
}

func YesterdayIn(loc *time.Location) Date {
	return TodayIn(loc).PrevDay()
}

// Create a Date from a time.Time object
func From(t time.Time) Date {
	y, m, d := t.Date()
	return Date{y, m, d, t.Location()}
}

func (d Date) DaysAfter(other Date) int {
	return int(d.BeginningOfDayIn(time.UTC).Sub(other.BeginningOfDayIn(time.UTC)).Hours() / 24)
}

func (d Date) AddDays(num int) Date {
	t := time.Date(d.Year, d.Month, d.Day, 0, 0, 0, 0, d.location).AddDate(0, 0, num)
	year, month, day := t.Date()
	return Date{year, month, day, d.location}
}

func (d Date) PrevDay() Date {
	if d.Day > 1 {
		return Date{d.Year, d.Month, d.Day - 1, d.location}
	} else if d.Month > 1 {
		return Date{d.Year, d.Month - 1, DaysInMonth(d.Month-1, d.Year), d.location}
	} else {
		return Date{d.Year - 1, time.December, 31, d.location}
	}
}

func (d Date) NextDay() Date {
	if d.Day < DaysInMonth(d.Month, d.Year) {
		return Date{d.Year, d.Month, d.Day + 1, d.location}
	} else if d.Month < time.December {
		return Date{d.Year, d.Month + 1, 1, d.location}
	} else {
		return Date{d.Year + 1, time.January, 1, d.location}
	}
}

// MostRecent returns today if today is the given weekday, otherwise it returns
// the date of the last day for that weekday.
//
// It is useful to get the beginning of the week:
//
//    today.MostRecent(time.Sunday) // in the US
//    today.MostRecent(time.Monday) // officially
//
func (d Date) MostRecent(weekday time.Weekday) Date {
	t := d.BeginningOfDay()
	today := t.Weekday()
	if today == weekday {
		return d
	} else if today > weekday {
		return From(t.AddDate(0, 0, int(weekday-today)))
	} else {
		return From(t.AddDate(0, 0, int(weekday-today)-7))
	}
}

// MostSoon returns today if today is the given weekday, otherwise it returns
// the date of the upcoming day for that weekday.
func (d Date) MostSoon(weekday time.Weekday) Date {
	if d.Weekday() == weekday {
		return d
	} else if d.Weekday() < weekday {
		return d.AddDays(int(weekday - d.Weekday()))
	} else {
		return d.AddDays(7 - int(d.Weekday()-weekday))
	}
}

func (d Date) Before(other Date) bool {
	return (d.Year < other.Year ||
		(d.Year == other.Year &&
			(d.Month < other.Month ||
				(d.Month == other.Month && d.Day < other.Day))))
}

func (d Date) Equal(other Date) bool {
	return (d.Month == other.Month) && (d.Day == other.Day) && (d.Year == other.Year)
}

func (d Date) After(other Date) bool {
	return (d.Year > other.Year ||
		(d.Year == other.Year &&
			(d.Month > other.Month ||
				(d.Month == other.Month && d.Day > other.Day))))
}

func (d Date) AtLeast(other Date) bool {
	return !other.After(d)
}

func (d Date) AtMost(other Date) bool {
	return !other.Before(d)
}

func (d Date) Weekday() time.Weekday {
	return d.BeginningOfDay().Weekday()
}

func (d Date) BeginningOfDay() time.Time {
	return time.Date(d.Year, d.Month, d.Day, 0, 0, 0, 0, d.location)
}

func (d Date) BeginningOfDayIn(timezone *time.Location) time.Time {
	return time.Date(d.Year, d.Month, d.Day, 0, 0, 0, 0, timezone)
}

func (d Date) EndOfDay() time.Time {
	return time.Date(d.Year, d.Month, d.Day, 23, 59, 59, NanosecondsInSecond-1, d.location)
}

func (d Date) EndOfDayIn(timezone *time.Location) time.Time {
	return time.Date(d.Year, d.Month, d.Day, 23, 59, 59, NanosecondsInSecond-1, timezone)
}

func (d Date) String() string {
	return d.Format("2006-01-02")
}

func (d Date) Format(format string) string {
	return d.BeginningOfDayIn(time.UTC).Format(format)
}

// Implements sql.Scanner interface
func (d *Date) Scan(src interface{}) error {
	t, ok := src.(time.Time)
	if !ok {
		return errors.New("date: scan value was not a Time")
	}

	*d = From(t)
	return nil
}

// Implements sql.driver.Valuer interface
func (d Date) Value() (driver.Value, error) {
	return d.BeginningOfDayIn(time.UTC), nil
}

// Implements json.Marshaler interface
func (d Date) MarshalJSON() ([]byte, error) {
	return []byte(d.Format(`"2006-01-02"`)), nil
}

// Implements json.Unmarshaler interface
func (d *Date) UnmarshalJSON(bytes []byte) error {
	t, err := time.Parse(`"2006-01-02"`, string(bytes))
	if err != nil {
		return err
	}

	*d = From(t)
	return nil
}

// TODO: Implment gob.GobEncoder and gob.GobDecoder to preserve timezone
// func (d Date) GobEncode() ([]byte, error) {}
// func (d *Date) GobDecode(bytes []byte) error {}

func IsLeapYear(year int) bool {
	return year%4 == 0 && (year%100 != 0 || year%400 == 0)
}

var DaysInNonLeapMonth = [12]int32{31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31}

func DaysInMonth(m time.Month, year int) int {
	if m == time.February && IsLeapYear(year) {
		return 29
	} else {
		return int(DaysInNonLeapMonth[m-1])
	}
}
<|endoftext|>


<commit_before>package messenger

type requestPayload struct {
	Recipient *recipient `json:"recipient,omitempty"`
	Message   *message   `json:"message,omitempty"`
}

type recipient struct {
	RecipientID string `json:"id,omitempty"`
}

type message struct {
	Text       string      `json:"text,omitempty"`
	Attachment *attachment `json:"attachment,omitempty"`
}

type attachment struct {
	Type    string   `json:"type,omitempty"`
	Payload *payload `json:"payload,omitempty"`
}

type payload struct {
	URL string `json:"url,omitempty"`
}
<commit_msg>Add facebook request payload structures.<commit_after>package messenger

type requestPayload struct {
	Recipient *recipient `json:"recipient,omitempty"`
	Message   *message   `json:"message,omitempty"`
}

type recipient struct {
	RecipientID string `json:"id,omitempty"`
}

type message struct {
	Text         string       `json:"text,omitempty"`
	Attachment   *attachment  `json:"attachment,omitempty"`
	QuickReplies []quickReply `json:"quick_replies,omitempty"`
}

type quickReply struct {
	ContentType string `json:"content_type,omitempty"`
	Title       string `json:"title,omitempty"`
	Payload     string `json:"payload,omitempty"`
	ImageURL    string `json:"image_url,omitempty"`
}

type attachment struct {
	Type    string   `json:"type,omitempty"`
	Payload *payload `json:"payload,omitempty"`
}

type payload struct {
	URL           string       `json:"url,omitempty"`
	TemplateType  string       `json:"template_type,omitempty"`
	Text          string       `json:"text,omitempty"`
	RecipientName string       `json:"recipient_name,omitempty"`
	OrderNumber   string       `json:"order_number,omitempty"`
	Currency      string       `json:"currency,omitempty"`
	PaymentMethod string       `json:"payment_method,omitempty"`
	OrderURL      string       `json:"order_url,omitempty"`
	Timestamp     string       `json:"timestamp,omitempty"`
	Elements      []element    `json:"elements,omitempty"`
	Address       *address     `json:"address,omitempty"`
	Summary       *summary     `json:"summary,omitempty"`
	Adjustments   []adjustment `json:"adjustments,omitempty"`
}

type element struct {
	Title    string   `json:"title,omitempty"`
	ItemURL  string   `json:"item_url,omitempty"`
	ImageURL string   `json:"image_url,omitempty"`
	Subtitle string   `json:"subtitle,omitempty"`
	Quantity int      `json:"quantity,omitempty"`
	Price    float64  `json:"price,omitempty"`
	Currency string   `json:"currency,omitempty"`
	Buttons  []button `json:"buttons,omitempty"`
}

type button struct {
	Type                string          `json:"type,omitempy"`
	Title               string          `json:"title,omitempty"`
	URL                 string          `json:"url,omitempy"`
	Payload             string          `json:"payload,omitempy"`
	WebviewHeightRatio  string          `json:"webview_height_ratio,omitempty"`
	MessengerExtensions bool            `json:"messenger_extensions,omitempty"`
	FallbackURL         string          `json:"fallback_url"`
	PaymentSummary      *paymentSummary `json:"payment_summary,omitempty"`
}

type address struct {
	Street1    string `json:"street_1,omitempty"`
	Street2    string `json:"street_2,omitempty"`
	City       string `json:"city,omitempty"`
	PostalCode string `json:"postal_code,omitempty"`
	State      string `json:"state,omitempty"`
	Country    string `json:"country,omitempty"`
}

type summary struct {
	Subtotal     float64 `json:"subtotal,omitempty"`
	ShippingCost float64 `json:"shipping_cost,omitempty"`
	TotalTax     float64 `json:"total_tax,omitempty"`
	TotalCost    float64 `json:"total_cost,omitempty"`
}

type adjustment struct {
	Name   string  `json:"name,omitempty"`
	Amount float64 `json:"amount,omitempty"`
}

type paymentSummary struct {
	Currency          string   `json:"currency,omitempty"`
	PaymentType       string   `json:"payment_type,omitempty"`
	MerchantName      string   `json:"merchant_name,omitempty"`
	RequestedUserInfo []string `json:"requested_user_info,omitempty"`
	PriceList         []price  `json:"price_list,omitempty"`
}

type price struct {
	Label  string `json:"label,omitempty"`
	Amount string `json:"amount,omitempty"`
}
<|endoftext|>


<commit_before>package couchbase

import (
	"log"
	"time"

	"fmt"
	"github.com/couchbase/gomemcached/client"
)

// A UprFeed streams mutation events from a bucket.
//
// Events from the bucket can be read from the channel 'C'.  Remember
// to call Close() on it when you're done, unless its channel has
// closed itself already.
type UprFeed struct {
	C <-chan *memcached.UprEvent

	bucket     *Bucket
	nodeFeeds  map[string]*FeedInfo     // The UPR feeds of the individual nodes
	output     chan *memcached.UprEvent // Same as C but writeably-typed
	quit       chan bool
	name       string // name of this UPR feed
	sequence   uint32 // sequence number for this feed
	connected  bool
	killSwitch chan bool
}

// UprFeed from a single connection
type FeedInfo struct {
	uprFeed   *memcached.UprFeed // UPR feed handle
	host      string             // hostname
	connected bool               // connected
}

type FailoverLog map[uint16]memcached.FailoverLog

// GetFailoverLogs, get the failover logs for a set of vbucket ids
func (b *Bucket) GetFailoverLogs(vBuckets []uint16) (FailoverLog, error) {

	// map vbids to their corresponding hosts
	vbHostList := make(map[string][]uint16)
	vbm := b.VBServerMap()
	if len(vbm.VBucketMap) < len(vBuckets) {
		return nil, fmt.Errorf("vbmap smaller than vbucket list: %v vs. %v",
			vbm.VBucketMap, vBuckets)
	}

	for _, vb := range vBuckets {
		masterID := vbm.VBucketMap[vb][0]
		master := b.getMasterNode(masterID)
		if master == "" {
			return nil, fmt.Errorf("No master found for vb %d", vb)
		}

		vbList := vbHostList[master]
		if vbList == nil {
			vbList = make([]uint16, 0)
		}
		vbList = append(vbList, vb)
		vbHostList[master] = vbList
	}

	failoverLogMap := make(FailoverLog)
	for _, serverConn := range b.getConnPools() {

		vbList := vbHostList[serverConn.host]
		if vbList == nil {
			continue
		}

		mc, err := serverConn.Get()
		if err != nil {
			log.Printf("No Free connections for vblist %v", vbList)
			return nil, fmt.Errorf("No Free connections for host %s",
				serverConn.host)

		}
		defer serverConn.Return(mc)
		failoverlogs, err := mc.UprGetFailoverLog(vbList)
		if err != nil {
			return nil, fmt.Errorf("Error getting failover log %s host %s",
				err.Error(), serverConn.host)

		}

		for vb, log := range failoverlogs {
			failoverLogMap[vb] = *log
		}
	}

	return failoverLogMap, nil
}

// StartUprFeed creates and starts a new Upr feed
// No data will be sent on the channel unless vbuckets streams are requested
func (b *Bucket) StartUprFeed(name string, sequence uint32) (*UprFeed, error) {

	feed := &UprFeed{
		bucket:     b,
		output:     make(chan *memcached.UprEvent, 10),
		quit:       make(chan bool),
		nodeFeeds:  make(map[string]*FeedInfo, 0),
		name:       name,
		sequence:   sequence,
		killSwitch: make(chan bool),
	}

	err := feed.connectToNodes()
	if err != nil {
		return nil, fmt.Errorf("Cannot connect to bucket %s", err.Error())
	}
	feed.connected = true
	go feed.run()

	feed.C = feed.output
	return feed, nil
}

// AddUprStream starts a stream for a vb on a feed
func (feed *UprFeed) UprRequestStream(vb uint16, flags uint32,
	vuuid, startSequence, endSequence, snapStart, snapEnd uint64) error {

	vbm := feed.bucket.VBServerMap()
	if len(vbm.VBucketMap) < int(vb) {
		return fmt.Errorf("vbmap smaller than vbucket list: %v vs. %v",
			vb, vbm.VBucketMap)
	}

	if int(vb) >= len(vbm.VBucketMap) {
		return fmt.Errorf("Invalid vbucket id %d", vb)
	}

	masterID := vbm.VBucketMap[vb][0]
	master := feed.bucket.getMasterNode(masterID)
	if master == "" {
		return fmt.Errorf("Master node not found for vbucket %d", vb)
	}
	singleFeed := feed.nodeFeeds[master]
	if singleFeed == nil {
		return fmt.Errorf("UprFeed for this host not found")
	}

	if err := singleFeed.uprFeed.UprRequestStream(vb, flags,
		vuuid, startSequence, endSequence, snapStart, snapEnd); err != nil {
		return err
	}

	return nil
}

// Goroutine that runs the feed
func (feed *UprFeed) run() {
	retryInterval := initialRetryInterval
	bucketOK := true
	for {
		// Connect to the UPR feed of each server node:
		if bucketOK {
			// Run until one of the sub-feeds fails:
			select {
			case <-feed.killSwitch:
			case <-feed.quit:
				return
			}
			//feed.closeNodeFeeds()
			retryInterval = initialRetryInterval
		}

		// On error, try to refresh the bucket in case the list of nodes changed:
		log.Printf("go-couchbase: UPR connection lost; reconnecting to bucket %q in %v",
			feed.bucket.Name, retryInterval)

		if err := feed.bucket.Refresh(); err != nil {
			log.Printf("Unable to refresh bucket %s ", err.Error())
			feed.closeNodeFeeds()
		}
		// this will only connect to nodes that are not connected or changed
		// user will have to reconnect the stream
		err := feed.connectToNodes()
		bucketOK = err == nil

		select {
		case <-time.After(retryInterval):
		case <-feed.quit:
			return
		}
		if retryInterval *= 2; retryInterval > maximumRetryInterval {
			retryInterval = maximumRetryInterval
		}
	}
}

func (feed *UprFeed) connectToNodes() (err error) {
	for _, serverConn := range feed.bucket.getConnPools() {

		// this maybe a reconnection, so check if the connection to the node
		// already exists. Connect only if the node is not found in the list
		// or connected == false
		nodeFeed := feed.nodeFeeds[serverConn.host]

		if nodeFeed != nil && nodeFeed.connected == true {
			continue
		}

		var singleFeed *memcached.UprFeed
		var name string
		if feed.name == "" {
			name = "DefaultUprClient"
		} else {
			name = feed.name
		}
		singleFeed, err = serverConn.StartUprFeed(name, feed.sequence)
		if err != nil {
			log.Printf("go-couchbase: Error connecting to upr feed of %s: %v", serverConn.host, err)
			feed.closeNodeFeeds()
			return
		}
		// add the node to the connection map
		feedInfo := &FeedInfo{uprFeed: singleFeed, connected: true, host: serverConn.host}
		feed.nodeFeeds[serverConn.host] = feedInfo
		go feed.forwardUprEvents(feedInfo, feed.killSwitch, serverConn.host)
	}
	return
}

// Goroutine that forwards Upr events from a single node's feed to the aggregate feed.
func (feed *UprFeed) forwardUprEvents(nodeFeed *FeedInfo, killSwitch chan bool, host string) {
	singleFeed := nodeFeed.uprFeed
	for {
		select {
		case event, ok := <-singleFeed.C:
			if !ok {
				if singleFeed.Error != nil {
					log.Printf("go-couchbase: Upr feed from %s failed: %v", host, singleFeed.Error)
				}
				killSwitch <- true
				return
			}
			feed.output <- event
		case <-feed.quit:
			nodeFeed.connected = false
			return
		}
	}
}

func (feed *UprFeed) closeNodeFeeds() {
	for _, f := range feed.nodeFeeds {
		f.uprFeed.Close()
	}
	feed.nodeFeeds = nil
}

// Close a Upr feed.
func (feed *UprFeed) Close() error {
	select {
	case <-feed.quit:
		return nil
	default:
	}

	feed.closeNodeFeeds()
	close(feed.quit)
	close(feed.output)
	return nil
}
<commit_msg>Refresh vbucket on NOT_MY_VBUCKET error<commit_after>package couchbase

import (
	"log"
	"time"

	"fmt"
	"github.com/couchbase/gomemcached"
	"github.com/couchbase/gomemcached/client"
)

// A UprFeed streams mutation events from a bucket.
//
// Events from the bucket can be read from the channel 'C'.  Remember
// to call Close() on it when you're done, unless its channel has
// closed itself already.
type UprFeed struct {
	C <-chan *memcached.UprEvent

	bucket     *Bucket
	nodeFeeds  map[string]*FeedInfo     // The UPR feeds of the individual nodes
	output     chan *memcached.UprEvent // Same as C but writeably-typed
	quit       chan bool
	name       string // name of this UPR feed
	sequence   uint32 // sequence number for this feed
	connected  bool
	killSwitch chan bool
}

// UprFeed from a single connection
type FeedInfo struct {
	uprFeed   *memcached.UprFeed // UPR feed handle
	host      string             // hostname
	connected bool               // connected
}

type FailoverLog map[uint16]memcached.FailoverLog

// GetFailoverLogs, get the failover logs for a set of vbucket ids
func (b *Bucket) GetFailoverLogs(vBuckets []uint16) (FailoverLog, error) {

	// map vbids to their corresponding hosts
	vbHostList := make(map[string][]uint16)
	vbm := b.VBServerMap()
	if len(vbm.VBucketMap) < len(vBuckets) {
		return nil, fmt.Errorf("vbmap smaller than vbucket list: %v vs. %v",
			vbm.VBucketMap, vBuckets)
	}

	for _, vb := range vBuckets {
		masterID := vbm.VBucketMap[vb][0]
		master := b.getMasterNode(masterID)
		if master == "" {
			return nil, fmt.Errorf("No master found for vb %d", vb)
		}

		vbList := vbHostList[master]
		if vbList == nil {
			vbList = make([]uint16, 0)
		}
		vbList = append(vbList, vb)
		vbHostList[master] = vbList
	}

	failoverLogMap := make(FailoverLog)
	for _, serverConn := range b.getConnPools() {

		vbList := vbHostList[serverConn.host]
		if vbList == nil {
			continue
		}

		mc, err := serverConn.Get()
		if err != nil {
			log.Printf("No Free connections for vblist %v", vbList)
			return nil, fmt.Errorf("No Free connections for host %s",
				serverConn.host)

		}
		defer serverConn.Return(mc)
		failoverlogs, err := mc.UprGetFailoverLog(vbList)
		if err != nil {
			return nil, fmt.Errorf("Error getting failover log %s host %s",
				err.Error(), serverConn.host)

		}

		for vb, log := range failoverlogs {
			failoverLogMap[vb] = *log
		}
	}

	return failoverLogMap, nil
}

// StartUprFeed creates and starts a new Upr feed
// No data will be sent on the channel unless vbuckets streams are requested
func (b *Bucket) StartUprFeed(name string, sequence uint32) (*UprFeed, error) {

	feed := &UprFeed{
		bucket:     b,
		output:     make(chan *memcached.UprEvent, 10),
		quit:       make(chan bool),
		nodeFeeds:  make(map[string]*FeedInfo, 0),
		name:       name,
		sequence:   sequence,
		killSwitch: make(chan bool),
	}

	err := feed.connectToNodes()
	if err != nil {
		return nil, fmt.Errorf("Cannot connect to bucket %s", err.Error())
	}
	feed.connected = true
	go feed.run()

	feed.C = feed.output
	return feed, nil
}

// AddUprStream starts a stream for a vb on a feed
func (feed *UprFeed) UprRequestStream(vb uint16, flags uint32,
	vuuid, startSequence, endSequence, snapStart, snapEnd uint64) error {

	vbm := feed.bucket.VBServerMap()
	if len(vbm.VBucketMap) < int(vb) {
		return fmt.Errorf("vbmap smaller than vbucket list: %v vs. %v",
			vb, vbm.VBucketMap)
	}

	if int(vb) >= len(vbm.VBucketMap) {
		return fmt.Errorf("Invalid vbucket id %d", vb)
	}

	masterID := vbm.VBucketMap[vb][0]
	master := feed.bucket.getMasterNode(masterID)
	if master == "" {
		return fmt.Errorf("Master node not found for vbucket %d", vb)
	}
	singleFeed := feed.nodeFeeds[master]
	if singleFeed == nil {
		return fmt.Errorf("UprFeed for this host not found")
	}

	if err := singleFeed.uprFeed.UprRequestStream(vb, flags,
		vuuid, startSequence, endSequence, snapStart, snapEnd); err != nil {
		return err
	}

	return nil
}

// Goroutine that runs the feed
func (feed *UprFeed) run() {
	retryInterval := initialRetryInterval
	bucketOK := true
	for {
		// Connect to the UPR feed of each server node:
		if bucketOK {
			// Run until one of the sub-feeds fails:
			select {
			case <-feed.killSwitch:
			case <-feed.quit:
				return
			}
			//feed.closeNodeFeeds()
			retryInterval = initialRetryInterval
		}

		// On error, try to refresh the bucket in case the list of nodes changed:
		log.Printf("go-couchbase: UPR connection lost; reconnecting to bucket %q in %v",
			feed.bucket.Name, retryInterval)

		if err := feed.bucket.Refresh(); err != nil {
			log.Printf("Unable to refresh bucket %s ", err.Error())
			feed.closeNodeFeeds()
		}
		// this will only connect to nodes that are not connected or changed
		// user will have to reconnect the stream
		err := feed.connectToNodes()
		bucketOK = err == nil

		select {
		case <-time.After(retryInterval):
		case <-feed.quit:
			return
		}
		if retryInterval *= 2; retryInterval > maximumRetryInterval {
			retryInterval = maximumRetryInterval
		}
	}
}

func (feed *UprFeed) connectToNodes() (err error) {
	for _, serverConn := range feed.bucket.getConnPools() {

		// this maybe a reconnection, so check if the connection to the node
		// already exists. Connect only if the node is not found in the list
		// or connected == false
		nodeFeed := feed.nodeFeeds[serverConn.host]

		if nodeFeed != nil && nodeFeed.connected == true {
			continue
		}

		var singleFeed *memcached.UprFeed
		var name string
		if feed.name == "" {
			name = "DefaultUprClient"
		} else {
			name = feed.name
		}
		singleFeed, err = serverConn.StartUprFeed(name, feed.sequence)
		if err != nil {
			log.Printf("go-couchbase: Error connecting to upr feed of %s: %v", serverConn.host, err)
			feed.closeNodeFeeds()
			return
		}
		// add the node to the connection map
		feedInfo := &FeedInfo{uprFeed: singleFeed, connected: true, host: serverConn.host}
		feed.nodeFeeds[serverConn.host] = feedInfo
		go feed.forwardUprEvents(feedInfo, feed.killSwitch, serverConn.host)
	}
	return
}

// Goroutine that forwards Upr events from a single node's feed to the aggregate feed.
func (feed *UprFeed) forwardUprEvents(nodeFeed *FeedInfo, killSwitch chan bool, host string) {
	singleFeed := nodeFeed.uprFeed
	for {
		select {
		case event, ok := <-singleFeed.C:
			if !ok {
				if singleFeed.Error != nil {
					log.Printf("go-couchbase: Upr feed from %s failed: %v", host, singleFeed.Error)
				}
				killSwitch <- true
				return
			}
			feed.output <- event
			if event.Status == gomemcached.NOT_MY_VBUCKET {
				log.Printf(" Got a not my vbucket error !! ")
				if err := feed.bucket.Refresh(); err != nil {
					log.Printf("Unable to refresh bucket %s ", err.Error())
					feed.closeNodeFeeds()
					return
				}
				// this will only connect to nodes that are not connected or changed
				// user will have to reconnect the stream
				if err := feed.connectToNodes(); err != nil {
					log.Printf("Unable to connect to nodes %s", err.Error())
					return
				}

			}
		case <-feed.quit:
			nodeFeed.connected = false
			return
		}
	}
}

func (feed *UprFeed) closeNodeFeeds() {
	for _, f := range feed.nodeFeeds {
		f.uprFeed.Close()
	}
	feed.nodeFeeds = nil
}

// Close a Upr feed.
func (feed *UprFeed) Close() error {
	select {
	case <-feed.quit:
		return nil
	default:
	}

	feed.closeNodeFeeds()
	close(feed.quit)
	close(feed.output)
	return nil
}
<|endoftext|>


<commit_before>// goupr package provides client sdk for UPR connections and streaming.
package couchbase

import (
	"encoding/binary"
	"fmt"
	mcd "github.com/dustin/gomemcached"
	mc "github.com/dustin/gomemcached/client"
	"log"
	"time"
)

const (
	UPR_OPEN         = mcd.CommandCode(0x50) // Open a upr connection with `name`
	UPR_ADD_STREAM   = mcd.CommandCode(0x51) // Sent by ebucketmigrator to upr consumer
	UPR_CLOSE_STREAM = mcd.CommandCode(0x52) // Sent by ebucketmigrator to upr consumer
	UPR_FAILOVER_LOG = mcd.CommandCode(0x54) // Request all known failover ids for restart
	UPR_STREAM_REQ   = mcd.CommandCode(0x53) // Stream request from consumer to producer
	UPR_STREAM_END   = mcd.CommandCode(0x55) // Sent by producer when it is going to end stream
	UPR_SNAPSHOTM    = mcd.CommandCode(0x56) // Sent by producer for a new snapshot
	UPR_MUTATION     = mcd.CommandCode(0x57) // Notifies SET/ADD/REPLACE/etc.  on the server
	UPR_DELETION     = mcd.CommandCode(0x58) // Notifies DELETE on the server
	UPR_EXPIRATION   = mcd.CommandCode(0x59) // Notifies key expiration
	UPR_FLUSH        = mcd.CommandCode(0x5a) // Notifies vbucket flush
)

const (
	ROLLBACK = mcd.Status(0x23)
)

// UprStream will maintain stream information per vbucket
type UprStream struct {
	Vbucket  uint16 // vbucket id
	Vuuid    uint64 // vbucket uuid
	Opaque   uint32 // messages from producer to this stream have same value
	Highseq  uint64 // to be supplied by the application
	Startseq uint64 // to be supplied by the application
	Endseq   uint64 // to be supplied by the application
	Flog     FailoverLog
}

// UprEvent objects will be created for stream mutations and deletions and
// published on UprFeed:C channel.
type UprEvent struct {
	Bucket  string // bucket name for this event
	Opstr   string // TODO: Make this consistent with TAP
	Vbucket uint16 // vbucket number
	Seqno   uint64 // sequence number
	Key     []byte
	Value   []byte
}

// UprFeed is per bucket structure managing connections to all nodes and
// vbucket streams.
type UprFeed struct {
	// Exported channel where an aggregate of all UprEvent are sent to app.
	C <-chan UprEvent
	c chan UprEvent

	bucket  *Bucket                   // upr client for bucket
	name    string                    // name of the connection used in UPR_OPEN
	vbmap   map[uint16]*uprConnection // vbucket-number->connection mapping
	streams map[uint16]*UprStream     // vbucket-number->stream mapping
	// `quit` channel is used to signal that a Close() is called on UprFeed
	quit chan bool
}

// uprConnection structure maintains an active memcached connection.
type uprConnection struct {
	host string // host to which `mc` is connected.
	conn *mc.Client
}

type msgT struct {
	uprconn *uprConnection
	pkt     mcd.MCRequest
	err     error
}

var eventTypes = map[mcd.CommandCode]string{ // Refer UprEvent
	UPR_MUTATION: "UPR_MUTATION",
	UPR_DELETION: "UPR_DELETION",
}

// StartUprFeed creates a feed that aggregates all mutations for the bucket
// and publishes them as UprEvent on UprFeed:C channel.
func StartUprFeed(b *Bucket, name string,
	streams map[uint16]*UprStream) (feed *UprFeed, err error) {

	uprconns, err := connectToNodes(b, name)
	if err == nil {
		vbmap := vbConns(b, uprconns)
		if streams == nil {
			streams = freshStreams(vbmap)
		}
		streams, err = startStreams(streams, vbmap)
		if err != nil {
			closeConnections(uprconns)
			return nil, err
		}
		feed = &UprFeed{
			bucket:  b,
			name:    name,
			vbmap:   vbmap,
			streams: streams,
			quit:    make(chan bool),
			c:       make(chan UprEvent, 16),
		}
		feed.C = feed.c
		go feed.doSession(uprconns)
	}
	return feed, err
}

// GetFailoverLogs return a list of vuuid and sequence number for all
// vbuckets.
func GetFailoverLogs(b *Bucket, name string) ([]FailoverLog, error) {

	var flog FailoverLog
	var err error
	uprconns, err := connectToNodes(b, name)
	if err == nil {
		flogs := make([]FailoverLog, 0)
		vbmap := vbConns(b, uprconns)
		for vb, uprconn := range vbmap {
			if flog, err = RequestFailoverLog(uprconn.conn, vb); err != nil {
				return nil, err
			}
			flogs = append(flogs, flog)
		}
		closeConnections(uprconns)
		return flogs, nil
	}
	return nil, err
}

// Close UprFeed. Does the opposite of StartUprFeed()
func (feed *UprFeed) Close() {
	if feed.hasQuit() {
		return
	}
	log.Println("Closing feed")
	close(feed.quit)
	feed.vbmap = nil
	feed.streams = nil
}

// Get stream will return the UprStream structure for vbucket `vbucket`.
// Caller can use the UprStream information to restart the stream later.
func (feed *UprFeed) GetStream(vbno uint16) *UprStream {
	return feed.streams[vbno]
}

func freshStreams(vbmaps map[uint16]*uprConnection) map[uint16]*UprStream {
	streams := make(map[uint16]*UprStream)
	for vbno := range vbmaps {
		streams[vbno] = &UprStream{
			Vbucket:  vbno,
			Vuuid:    0,
			Opaque:   uint32(vbno),
			Highseq:  0,
			Startseq: 0,
			Endseq:   0xFFFFFFFFFFFFFFFF,
		}
	}
	return streams
}

// connect with all servers holding data for `feed.bucket` and try reconnecting
// until `feed` is closed.
func (feed *UprFeed) doSession(uprconns []*uprConnection) {
	msgch := make(chan msgT)
	killSwitch := make(chan bool)

	for _, uprconn := range uprconns {
		go doReceive(uprconn, uprconn.host, msgch, killSwitch)
	}
	feed.doEvents(msgch)

	close(killSwitch)
	closeConnections(uprconns)
	close(feed.c)
}

// TODO: This function is not used at present.
func (feed *UprFeed) retryConnections(
	uprconns []*uprConnection) ([]*uprConnection, bool) {

	var err error

	log.Println("Retrying connections ...")
	retryInterval := initialRetryInterval
	for {
		closeConnections(uprconns)
		uprconns, err = connectToNodes(feed.bucket, feed.name)
		if err == nil {
			feed.vbmap = vbConns(feed.bucket, uprconns)
			feed.streams, err = startStreams(feed.streams, feed.vbmap)
			if err == nil {
				return uprconns, false
			}
		} else {
			log.Println("retry:", err)
		}
		log.Printf("Retrying after %v seconds ...", retryInterval)
		select {
		case <-time.After(retryInterval):
		case <-feed.quit:
			return nil, true
		}
		if retryInterval *= 2; retryInterval > maximumRetryInterval {
			retryInterval = maximumRetryInterval
		}
	}
	return uprconns, false
}

func (feed *UprFeed) doEvents(msgch <-chan msgT) bool {
	for {
		select {
		case msg, ok := <-msgch:
			if !ok {
				return false
			}
			if msg.err != nil {
				log.Printf(
					"Received error from %v: %v\n", msg.uprconn.host, msg.err)
				return false
			} else if err := handleUprMessage(feed, &msg.pkt); err != nil {
				log.Println(err)
				return false
			}
		case <-feed.quit:
			return true
		}
	}
}

func handleUprMessage(feed *UprFeed, req *mcd.MCRequest) (err error) {
	vb := uint16(req.Opaque)
	stream := feed.streams[uint16(req.Opaque)]
	switch req.Opcode {
	case UPR_STREAM_REQ:
		rollb, flog, err := handleStreamResponse(Request2Response(req))
		if err == nil {
			if flog != nil {
				stream.Flog = flog
			} else if rollb > 0 {
				uprconn := feed.vbmap[vb]
				log.Println("Requesting a rollback for %v to sequence %v",
					vb, rollb)
				flags := uint32(0)
				err = RequestStream(
					uprconn.conn, flags, req.Opaque, vb, stream.Vuuid,
					rollb, stream.Endseq, stream.Highseq)
			}
		}
	case UPR_MUTATION, UPR_DELETION:
		e := feed.makeUprEvent(req)
		stream.Startseq = e.Seqno
		feed.c <- e
	case UPR_STREAM_END:
		res := Request2Response(req)
		err = fmt.Errorf("Stream %v is ending", uint16(res.Opaque))
	case UPR_SNAPSHOTM:
	case UPR_CLOSE_STREAM, UPR_EXPIRATION, UPR_FLUSH, UPR_ADD_STREAM:
		err = fmt.Errorf("Opcode %v not implemented", req.Opcode)
	default:
		err = fmt.Errorf("ERROR: un-known opcode received %v", req)
	}
	return
}

func handleStreamResponse(res *mcd.MCResponse) (uint64, FailoverLog, error) {
	var rollback uint64
	var err error
	var flog FailoverLog

	switch {
	case res.Status == ROLLBACK && len(res.Extras) != 8:
		err = fmt.Errorf("invalid rollback %v\n", res.Extras)
	case res.Status == ROLLBACK:
		rollback = binary.BigEndian.Uint64(res.Extras)
	case res.Status != mcd.SUCCESS:
		err = fmt.Errorf("Unexpected status %v", res.Status)
	}
	if err == nil {
		if rollback > 0 {
			return rollback, flog, err
		} else {
			flog, err = ParseFailoverLog(res.Body[:])
		}
	}
	return rollback, flog, err
}

func connectToNodes(b *Bucket, name string) ([]*uprConnection, error) {
	var conn *mc.Client
	var err error

	uprconns := make([]*uprConnection, 0)
	for _, cp := range b.getConnPools() {
		if cp == nil {
			return nil, fmt.Errorf("go-couchbase: no connection pool")
		}
		if conn, err = cp.Get(); err != nil {
			return nil, err
		}
		// A connection can't be used after UPR; Dont' count it against the
		// connection pool capacity
		<-cp.createsem
		if uprconn, err := connectToNode(b, conn, name, cp.host); err != nil {
			return nil, err
		} else {
			uprconns = append(uprconns, uprconn)
		}
	}
	return uprconns, nil
}

func connectToNode(b *Bucket, conn *mc.Client,
	name, host string) (uprconn *uprConnection, err error) {
	if err = UprOpen(conn, name, uint32(0x1) /*flags*/); err != nil {
		return
	}
	uprconn = &uprConnection{
		host: host,
		conn: conn,
	}
	log.Printf("Connected to host %v (%p)\n", host, conn)
	return
}

func vbConns(b *Bucket, uprconns []*uprConnection) map[uint16]*uprConnection {

	servers, vbmaps := b.VBSMJson.ServerList, b.VBSMJson.VBucketMap
	vbconns := make(map[uint16]*uprConnection)
	for _, uprconn := range uprconns {
		// Collect vbuckets under this connection
		for vbno := range vbmaps {
			host := servers[vbmaps[int(vbno)][0]]
			if uprconn.host == host {
				vbconns[uint16(vbno)] = uprconn
			}
		}
	}
	return vbconns
}

func startStreams(streams map[uint16]*UprStream,
	vbmap map[uint16]*uprConnection) (map[uint16]*UprStream, error) {

	for vb, uprconn := range vbmap {
		stream := streams[vb]
		flags := uint32(0)
		err := RequestStream(
			uprconn.conn, flags, uint32(vb), vb, stream.Vuuid,
			stream.Startseq, stream.Endseq, stream.Highseq)
		if err != nil {
			return nil, err
		}
		log.Printf(
			"Posted stream request for vbucket %v from %v (vuuid:%v)\n",
			vb, stream.Startseq, stream.Vuuid)
	}
	return streams, nil
}

func doReceive(uprconn *uprConnection, host string,
	msgch chan msgT, killSwitch chan bool) {

	var hdr [mcd.HDR_LEN]byte
	var msg msgT
	var pkt mcd.MCRequest
	var err error

	mcconn := uprconn.conn.GetConnection()

loop:
	for {
		if _, err = pkt.Receive(mcconn, hdr[:]); err != nil {
			msg = msgT{uprconn: uprconn, err: err}
		} else {
			msg = msgT{uprconn: uprconn, pkt: pkt}
		}
		select {
		case msgch <- msg:
		case <-killSwitch:
			break loop
		}
	}
	return
}

func (feed *UprFeed) makeUprEvent(req *mcd.MCRequest) UprEvent {
	bySeqno := binary.BigEndian.Uint64(req.Extras[:8])
	e := UprEvent{
		Bucket:  feed.bucket.Name,
		Opstr:   eventTypes[req.Opcode],
		Vbucket: req.VBucket,
		Seqno:   bySeqno,
		Key:     req.Key,
		Value:   req.Body,
	}
	return e
}

func (feed *UprFeed) hasQuit() bool {
	select {
	case <-feed.quit:
		return true
	default:
		return false
	}
}

func closeConnections(uprconns []*uprConnection) {
	for _, uprconn := range uprconns {
		log.Printf("Closing connection for %v: %p\n", uprconn.host, uprconn.conn)
		if uprconn.conn != nil {
			uprconn.conn.Close()
		}
		uprconn.conn = nil
	}
}
<commit_msg>Renamed GetConnection() to Hijack()<commit_after>// goupr package provides client sdk for UPR connections and streaming.
package couchbase

import (
	"encoding/binary"
	"fmt"
	mcd "github.com/dustin/gomemcached"
	mc "github.com/dustin/gomemcached/client"
	"log"
	"time"
)

const (
	UPR_OPEN         = mcd.CommandCode(0x50) // Open a upr connection with `name`
	UPR_ADD_STREAM   = mcd.CommandCode(0x51) // Sent by ebucketmigrator to upr consumer
	UPR_CLOSE_STREAM = mcd.CommandCode(0x52) // Sent by ebucketmigrator to upr consumer
	UPR_FAILOVER_LOG = mcd.CommandCode(0x54) // Request all known failover ids for restart
	UPR_STREAM_REQ   = mcd.CommandCode(0x53) // Stream request from consumer to producer
	UPR_STREAM_END   = mcd.CommandCode(0x55) // Sent by producer when it is going to end stream
	UPR_SNAPSHOTM    = mcd.CommandCode(0x56) // Sent by producer for a new snapshot
	UPR_MUTATION     = mcd.CommandCode(0x57) // Notifies SET/ADD/REPLACE/etc.  on the server
	UPR_DELETION     = mcd.CommandCode(0x58) // Notifies DELETE on the server
	UPR_EXPIRATION   = mcd.CommandCode(0x59) // Notifies key expiration
	UPR_FLUSH        = mcd.CommandCode(0x5a) // Notifies vbucket flush
)

const (
	ROLLBACK = mcd.Status(0x23)
)

// UprStream will maintain stream information per vbucket
type UprStream struct {
	Vbucket  uint16 // vbucket id
	Vuuid    uint64 // vbucket uuid
	Opaque   uint32 // messages from producer to this stream have same value
	Highseq  uint64 // to be supplied by the application
	Startseq uint64 // to be supplied by the application
	Endseq   uint64 // to be supplied by the application
	Flog     FailoverLog
}

// UprEvent objects will be created for stream mutations and deletions and
// published on UprFeed:C channel.
type UprEvent struct {
	Bucket  string // bucket name for this event
	Opstr   string // TODO: Make this consistent with TAP
	Vbucket uint16 // vbucket number
	Seqno   uint64 // sequence number
	Key     []byte
	Value   []byte
}

// UprFeed is per bucket structure managing connections to all nodes and
// vbucket streams.
type UprFeed struct {
	// Exported channel where an aggregate of all UprEvent are sent to app.
	C <-chan UprEvent
	c chan UprEvent

	bucket  *Bucket                   // upr client for bucket
	name    string                    // name of the connection used in UPR_OPEN
	vbmap   map[uint16]*uprConnection // vbucket-number->connection mapping
	streams map[uint16]*UprStream     // vbucket-number->stream mapping
	// `quit` channel is used to signal that a Close() is called on UprFeed
	quit chan bool
}

// uprConnection structure maintains an active memcached connection.
type uprConnection struct {
	host string // host to which `mc` is connected.
	conn *mc.Client
}

type msgT struct {
	uprconn *uprConnection
	pkt     mcd.MCRequest
	err     error
}

var eventTypes = map[mcd.CommandCode]string{ // Refer UprEvent
	UPR_MUTATION: "UPR_MUTATION",
	UPR_DELETION: "UPR_DELETION",
}

// StartUprFeed creates a feed that aggregates all mutations for the bucket
// and publishes them as UprEvent on UprFeed:C channel.
func StartUprFeed(b *Bucket, name string,
	streams map[uint16]*UprStream) (feed *UprFeed, err error) {

	uprconns, err := connectToNodes(b, name)
	if err == nil {
		vbmap := vbConns(b, uprconns)
		if streams == nil {
			streams = freshStreams(vbmap)
		}
		streams, err = startStreams(streams, vbmap)
		if err != nil {
			closeConnections(uprconns)
			return nil, err
		}
		feed = &UprFeed{
			bucket:  b,
			name:    name,
			vbmap:   vbmap,
			streams: streams,
			quit:    make(chan bool),
			c:       make(chan UprEvent, 16),
		}
		feed.C = feed.c
		go feed.doSession(uprconns)
	}
	return feed, err
}

// GetFailoverLogs return a list of vuuid and sequence number for all
// vbuckets.
func GetFailoverLogs(b *Bucket, name string) ([]FailoverLog, error) {

	var flog FailoverLog
	var err error
	uprconns, err := connectToNodes(b, name)
	if err == nil {
		flogs := make([]FailoverLog, 0)
		vbmap := vbConns(b, uprconns)
		for vb, uprconn := range vbmap {
			if flog, err = RequestFailoverLog(uprconn.conn, vb); err != nil {
				return nil, err
			}
			flogs = append(flogs, flog)
		}
		closeConnections(uprconns)
		return flogs, nil
	}
	return nil, err
}

// Close UprFeed. Does the opposite of StartUprFeed()
func (feed *UprFeed) Close() {
	if feed.hasQuit() {
		return
	}
	log.Println("Closing feed")
	close(feed.quit)
	feed.vbmap = nil
	feed.streams = nil
}

// Get stream will return the UprStream structure for vbucket `vbucket`.
// Caller can use the UprStream information to restart the stream later.
func (feed *UprFeed) GetStream(vbno uint16) *UprStream {
	return feed.streams[vbno]
}

func freshStreams(vbmaps map[uint16]*uprConnection) map[uint16]*UprStream {
	streams := make(map[uint16]*UprStream)
	for vbno := range vbmaps {
		streams[vbno] = &UprStream{
			Vbucket:  vbno,
			Vuuid:    0,
			Opaque:   uint32(vbno),
			Highseq:  0,
			Startseq: 0,
			Endseq:   0xFFFFFFFFFFFFFFFF,
		}
	}
	return streams
}

// connect with all servers holding data for `feed.bucket` and try reconnecting
// until `feed` is closed.
func (feed *UprFeed) doSession(uprconns []*uprConnection) {
	msgch := make(chan msgT)
	killSwitch := make(chan bool)

	for _, uprconn := range uprconns {
		go doReceive(uprconn, uprconn.host, msgch, killSwitch)
	}
	feed.doEvents(msgch)

	close(killSwitch)
	closeConnections(uprconns)
	close(feed.c)
}

// TODO: This function is not used at present.
func (feed *UprFeed) retryConnections(
	uprconns []*uprConnection) ([]*uprConnection, bool) {

	var err error

	log.Println("Retrying connections ...")
	retryInterval := initialRetryInterval
	for {
		closeConnections(uprconns)
		uprconns, err = connectToNodes(feed.bucket, feed.name)
		if err == nil {
			feed.vbmap = vbConns(feed.bucket, uprconns)
			feed.streams, err = startStreams(feed.streams, feed.vbmap)
			if err == nil {
				return uprconns, false
			}
		} else {
			log.Println("retry:", err)
		}
		log.Printf("Retrying after %v seconds ...", retryInterval)
		select {
		case <-time.After(retryInterval):
		case <-feed.quit:
			return nil, true
		}
		if retryInterval *= 2; retryInterval > maximumRetryInterval {
			retryInterval = maximumRetryInterval
		}
	}
	return uprconns, false
}

func (feed *UprFeed) doEvents(msgch <-chan msgT) bool {
	for {
		select {
		case msg, ok := <-msgch:
			if !ok {
				return false
			}
			if msg.err != nil {
				log.Printf(
					"Received error from %v: %v\n", msg.uprconn.host, msg.err)
				return false
			} else if err := handleUprMessage(feed, &msg.pkt); err != nil {
				log.Println(err)
				return false
			}
		case <-feed.quit:
			return true
		}
	}
}

func handleUprMessage(feed *UprFeed, req *mcd.MCRequest) (err error) {
	vb := uint16(req.Opaque)
	stream := feed.streams[uint16(req.Opaque)]
	switch req.Opcode {
	case UPR_STREAM_REQ:
		rollb, flog, err := handleStreamResponse(Request2Response(req))
		if err == nil {
			if flog != nil {
				stream.Flog = flog
			} else if rollb > 0 {
				uprconn := feed.vbmap[vb]
				log.Println("Requesting a rollback for %v to sequence %v",
					vb, rollb)
				flags := uint32(0)
				err = RequestStream(
					uprconn.conn, flags, req.Opaque, vb, stream.Vuuid,
					rollb, stream.Endseq, stream.Highseq)
			}
		}
	case UPR_MUTATION, UPR_DELETION:
		e := feed.makeUprEvent(req)
		stream.Startseq = e.Seqno
		feed.c <- e
	case UPR_STREAM_END:
		res := Request2Response(req)
		err = fmt.Errorf("Stream %v is ending", uint16(res.Opaque))
	case UPR_SNAPSHOTM:
	case UPR_CLOSE_STREAM, UPR_EXPIRATION, UPR_FLUSH, UPR_ADD_STREAM:
		err = fmt.Errorf("Opcode %v not implemented", req.Opcode)
	default:
		err = fmt.Errorf("ERROR: un-known opcode received %v", req)
	}
	return
}

func handleStreamResponse(res *mcd.MCResponse) (uint64, FailoverLog, error) {
	var rollback uint64
	var err error
	var flog FailoverLog

	switch {
	case res.Status == ROLLBACK && len(res.Extras) != 8:
		err = fmt.Errorf("invalid rollback %v\n", res.Extras)
	case res.Status == ROLLBACK:
		rollback = binary.BigEndian.Uint64(res.Extras)
	case res.Status != mcd.SUCCESS:
		err = fmt.Errorf("Unexpected status %v", res.Status)
	}
	if err == nil {
		if rollback > 0 {
			return rollback, flog, err
		} else {
			flog, err = ParseFailoverLog(res.Body[:])
		}
	}
	return rollback, flog, err
}

func connectToNodes(b *Bucket, name string) ([]*uprConnection, error) {
	var conn *mc.Client
	var err error

	uprconns := make([]*uprConnection, 0)
	for _, cp := range b.getConnPools() {
		if cp == nil {
			return nil, fmt.Errorf("go-couchbase: no connection pool")
		}
		if conn, err = cp.Get(); err != nil {
			return nil, err
		}
		// A connection can't be used after UPR; Dont' count it against the
		// connection pool capacity
		<-cp.createsem
		if uprconn, err := connectToNode(b, conn, name, cp.host); err != nil {
			return nil, err
		} else {
			uprconns = append(uprconns, uprconn)
		}
	}
	return uprconns, nil
}

func connectToNode(b *Bucket, conn *mc.Client,
	name, host string) (uprconn *uprConnection, err error) {
	if err = UprOpen(conn, name, uint32(0x1) /*flags*/); err != nil {
		return
	}
	uprconn = &uprConnection{
		host: host,
		conn: conn,
	}
	log.Printf("Connected to host %v (%p)\n", host, conn)
	return
}

func vbConns(b *Bucket, uprconns []*uprConnection) map[uint16]*uprConnection {

	servers, vbmaps := b.VBSMJson.ServerList, b.VBSMJson.VBucketMap
	vbconns := make(map[uint16]*uprConnection)
	for _, uprconn := range uprconns {
		// Collect vbuckets under this connection
		for vbno := range vbmaps {
			host := servers[vbmaps[int(vbno)][0]]
			if uprconn.host == host {
				vbconns[uint16(vbno)] = uprconn
			}
		}
	}
	return vbconns
}

func startStreams(streams map[uint16]*UprStream,
	vbmap map[uint16]*uprConnection) (map[uint16]*UprStream, error) {

	for vb, uprconn := range vbmap {
		stream := streams[vb]
		flags := uint32(0)
		err := RequestStream(
			uprconn.conn, flags, uint32(vb), vb, stream.Vuuid,
			stream.Startseq, stream.Endseq, stream.Highseq)
		if err != nil {
			return nil, err
		}
		log.Printf(
			"Posted stream request for vbucket %v from %v (vuuid:%v)\n",
			vb, stream.Startseq, stream.Vuuid)
	}
	return streams, nil
}

func doReceive(uprconn *uprConnection, host string,
	msgch chan msgT, killSwitch chan bool) {

	var hdr [mcd.HDR_LEN]byte
	var msg msgT
	var pkt mcd.MCRequest
	var err error

	mcconn := uprconn.conn.Hijack()

loop:
	for {
		if _, err = pkt.Receive(mcconn, hdr[:]); err != nil {
			msg = msgT{uprconn: uprconn, err: err}
		} else {
			msg = msgT{uprconn: uprconn, pkt: pkt}
		}
		select {
		case msgch <- msg:
		case <-killSwitch:
			break loop
		}
	}
	return
}

func (feed *UprFeed) makeUprEvent(req *mcd.MCRequest) UprEvent {
	bySeqno := binary.BigEndian.Uint64(req.Extras[:8])
	e := UprEvent{
		Bucket:  feed.bucket.Name,
		Opstr:   eventTypes[req.Opcode],
		Vbucket: req.VBucket,
		Seqno:   bySeqno,
		Key:     req.Key,
		Value:   req.Body,
	}
	return e
}

func (feed *UprFeed) hasQuit() bool {
	select {
	case <-feed.quit:
		return true
	default:
		return false
	}
}

func closeConnections(uprconns []*uprConnection) {
	for _, uprconn := range uprconns {
		log.Printf("Closing connection for %v: %p\n", uprconn.host, uprconn.conn)
		if uprconn.conn != nil {
			uprconn.conn.Close()
		}
		uprconn.conn = nil
	}
}
<|endoftext|>


<commit_before>package regsrc

import (
	"regexp"
	"strings"

	"golang.org/x/net/idna"
)

var (
	// InvalidHostString is a placeholder returned when a raw host can't be
	// converted by IDNA spec. It will never be returned for any host for which
	// Valid() is true.
	InvalidHostString = "<invalid host>"

	// urlLabelEndSubRe is a sub-expression that matches any character that's
	// allowed at the start or end of a URL label according to RFC1123.
	urlLabelEndSubRe = "[0-9A-Za-z]"

	// urlLabelEndSubRe is a sub-expression that matches any character that's
	// allowed at in a non-start or end of a URL label according to RFC1123.
	urlLabelMidSubRe = "[0-9A-Za-z-]"

	// urlLabelUnicodeSubRe is a sub-expression that matches any non-ascii char
	// in an IDN (Unicode) display URL. It's not strict - there are only ~15k
	// valid Unicode points in IDN RFC (some with conditions). We are just going
	// with being liberal with matching and then erroring if we fail to convert
	// to punycode later (which validates chars fully). This at least ensures
	// ascii chars dissalowed by the RC1123 parts above don't become legal
	// again.
	urlLabelUnicodeSubRe = "[^[:ascii:]]"

	// hostLabelSubRe is the sub-expression that matches a valid hostname label.
	// It does not anchor the start or end so it can be composed into more
	// complex RegExps below. Note that for sanity we don't handle disallowing
	// raw punycode in this regexp (esp. since re2 doesn't support negative
	// lookbehind, but we can capture it's presence here to check later).
	hostLabelSubRe = "" +
		// Match valid initial char, or unicode char
		"(?:" + urlLabelEndSubRe + "|" + urlLabelUnicodeSubRe + ")" +
		// Optionally, match 0 to 61 valid URL or Unicode chars,
		// followed by one valid end char or unicode char
		"(?:" +
		"(?:" + urlLabelMidSubRe + "|" + urlLabelUnicodeSubRe + "){0,61}" +
		"(?:" + urlLabelEndSubRe + "|" + urlLabelUnicodeSubRe + ")" +
		")?"

	// hostSubRe is the sub-expression that matches a valid host prefix.
	// Allows custom port.
	hostSubRe = hostLabelSubRe + "(?:\\." + hostLabelSubRe + ")+(?::\\d+)?"

	// hostRe is a regexp that matches a valid host prefix. Additional
	// validation of unicode strings is needed for matches.
	hostRe = regexp.MustCompile("^" + hostSubRe + "$")
)

// FriendlyHost describes a registry instance identified in source strings by a
// simple bare hostname like registry.terraform.io.
type FriendlyHost struct {
	Raw string
}

func NewFriendlyHost(host string) *FriendlyHost {
	return &FriendlyHost{Raw: host}
}

// ParseFriendlyHost attempts to parse a valid "friendly host" prefix from the
// given string. If no valid prefix is found, host will be nil and rest will
// contain the full source string. The host prefix must terminate at the end of
// the input or at the first / character. If one or more characters exist after
// the first /, they will be returned as rest (without the / delimiter).
// Hostnames containing punycode WILL be parsed successfully since they may have
// come from an internal normalized source string, however should be considered
// invalid if the string came from a user directly. This must be checked
// explicitly for user-input strings by calling Valid() on the
// returned host.
func ParseFriendlyHost(source string) (host *FriendlyHost, rest string) {
	parts := strings.SplitN(source, "/", 2)

	if hostRe.MatchString(parts[0]) {
		host = &FriendlyHost{Raw: parts[0]}
		if len(parts) == 2 {
			rest = parts[1]
		}
		return
	}

	// No match, return whole string as rest along with nil host
	rest = source
	return
}

// Valid returns whether the host prefix is considered valid in any case.
// Example of invalid prefixes might include ones that don't conform to the host
// name specifications. Not that IDN prefixes containing punycode are not valid
// input which we expect to always be in user-input or normalised display form.
func (h *FriendlyHost) Valid() bool {
	if h.Display() == InvalidHostString {
		return false
	}
	if h.Normalized() == InvalidHostString {
		return false
	}
	if containsPuny(h.Raw) {
		return false
	}
	return true
}

// Display returns the host formatted for display to the user in CLI or web
// output.
func (h *FriendlyHost) Display() string {
	parts := strings.SplitN(h.Raw, ":", 2)
	var err error
	parts[0], err = idna.Display.ToUnicode(parts[0])
	if err != nil {
		return InvalidHostString
	}
	return strings.Join(parts, ":")
}

// Normalized returns the host formatted for internal reference or comparison.
func (h *FriendlyHost) Normalized() string {
	// For now IDNA does all the normalisation we need including case-folding
	// pure ASCII to lower. But breaks if a custom port is included while we
	// want to allow that and normalize comparison including it,
	parts := strings.SplitN(h.Raw, ":", 2)
	var err error
	parts[0], err = idna.Lookup.ToASCII(parts[0])
	if err != nil {
		return InvalidHostString
	}
	return strings.Join(parts, ":")
}

// String returns the host formatted as the user originally typed it assuming it
// was parsed from user input.
func (h *FriendlyHost) String() string {
	return h.Raw
}

// Equal compares the FriendlyHost against another instance taking normalization
// into account.
func (h *FriendlyHost) Equal(other *FriendlyHost) bool {
	if other == nil {
		return false
	}
	return h.Normalized() == other.Normalized()
}

func containsPuny(host string) bool {
	for _, lbl := range strings.Split(host, ".") {
		if strings.HasPrefix(strings.ToLower(lbl), "xn--") {
			return true
		}
	}
	return false
}
<commit_msg>use svchost.Hostname for FriendlyHost validation<commit_after>package regsrc

import (
	"regexp"
	"strings"

	"github.com/hashicorp/terraform/svchost"
)

var (
	// InvalidHostString is a placeholder returned when a raw host can't be
	// converted by IDNA spec. It will never be returned for any host for which
	// Valid() is true.
	InvalidHostString = "<invalid host>"

	// urlLabelEndSubRe is a sub-expression that matches any character that's
	// allowed at the start or end of a URL label according to RFC1123.
	urlLabelEndSubRe = "[0-9A-Za-z]"

	// urlLabelEndSubRe is a sub-expression that matches any character that's
	// allowed at in a non-start or end of a URL label according to RFC1123.
	urlLabelMidSubRe = "[0-9A-Za-z-]"

	// urlLabelUnicodeSubRe is a sub-expression that matches any non-ascii char
	// in an IDN (Unicode) display URL. It's not strict - there are only ~15k
	// valid Unicode points in IDN RFC (some with conditions). We are just going
	// with being liberal with matching and then erroring if we fail to convert
	// to punycode later (which validates chars fully). This at least ensures
	// ascii chars dissalowed by the RC1123 parts above don't become legal
	// again.
	urlLabelUnicodeSubRe = "[^[:ascii:]]"

	// hostLabelSubRe is the sub-expression that matches a valid hostname label.
	// It does not anchor the start or end so it can be composed into more
	// complex RegExps below. Note that for sanity we don't handle disallowing
	// raw punycode in this regexp (esp. since re2 doesn't support negative
	// lookbehind, but we can capture it's presence here to check later).
	hostLabelSubRe = "" +
		// Match valid initial char, or unicode char
		"(?:" + urlLabelEndSubRe + "|" + urlLabelUnicodeSubRe + ")" +
		// Optionally, match 0 to 61 valid URL or Unicode chars,
		// followed by one valid end char or unicode char
		"(?:" +
		"(?:" + urlLabelMidSubRe + "|" + urlLabelUnicodeSubRe + "){0,61}" +
		"(?:" + urlLabelEndSubRe + "|" + urlLabelUnicodeSubRe + ")" +
		")?"

	// hostSubRe is the sub-expression that matches a valid host prefix.
	// Allows custom port.
	hostSubRe = hostLabelSubRe + "(?:\\." + hostLabelSubRe + ")+(?::\\d+)?"

	// hostRe is a regexp that matches a valid host prefix. Additional
	// validation of unicode strings is needed for matches.
	hostRe = regexp.MustCompile("^" + hostSubRe + "$")
)

// FriendlyHost describes a registry instance identified in source strings by a
// simple bare hostname like registry.terraform.io.
type FriendlyHost struct {
	Raw string
}

func NewFriendlyHost(host string) *FriendlyHost {
	return &FriendlyHost{Raw: host}
}

// ParseFriendlyHost attempts to parse a valid "friendly host" prefix from the
// given string. If no valid prefix is found, host will be nil and rest will
// contain the full source string. The host prefix must terminate at the end of
// the input or at the first / character. If one or more characters exist after
// the first /, they will be returned as rest (without the / delimiter).
// Hostnames containing punycode WILL be parsed successfully since they may have
// come from an internal normalized source string, however should be considered
// invalid if the string came from a user directly. This must be checked
// explicitly for user-input strings by calling Valid() on the
// returned host.
func ParseFriendlyHost(source string) (host *FriendlyHost, rest string) {
	parts := strings.SplitN(source, "/", 2)

	if hostRe.MatchString(parts[0]) {
		host = &FriendlyHost{Raw: parts[0]}
		if len(parts) == 2 {
			rest = parts[1]
		}
		return
	}

	// No match, return whole string as rest along with nil host
	rest = source
	return
}

// Valid returns whether the host prefix is considered valid in any case.
// Example of invalid prefixes might include ones that don't conform to the host
// name specifications. Not that IDN prefixes containing punycode are not valid
// input which we expect to always be in user-input or normalised display form.
func (h *FriendlyHost) Valid() bool {
	return svchost.IsValid(h.Raw)
}

// Display returns the host formatted for display to the user in CLI or web
// output.
func (h *FriendlyHost) Display() string {
	return svchost.ForDisplay(h.Raw)
}

// Normalized returns the host formatted for internal reference or comparison.
func (h *FriendlyHost) Normalized() string {
	host, err := svchost.ForComparison(h.Raw)
	if err != nil {
		return InvalidHostString
	}
	return string(host)
}

// String returns the host formatted as the user originally typed it assuming it
// was parsed from user input.
func (h *FriendlyHost) String() string {
	return h.Raw
}

// Equal compares the FriendlyHost against another instance taking normalization
// into account. Invalid hosts cannot be compared and will always return false.
func (h *FriendlyHost) Equal(other *FriendlyHost) bool {
	if other == nil {
		return false
	}

	return h.Normalized() == other.Normalized()
}

func (h *FriendlyHost) SvcHost() (svchost.Hostname, error) {
	return svchost.ForComparison(h.Raw)
}
<|endoftext|>


<commit_before>package db

import (
	"bytes"
	"encoding/gob"
	"fmt"

	"github.com/boltdb/bolt"
)

type DB interface {
	Load(string, interface{}) (bool, error)
	Save(string, interface{}) error
}

type BoltDB struct {
	db *bolt.DB
}

var (
	versionIdent       = []byte("version")
	persistenceVersion = []byte{1, 0} // major.minor
	topBucket          = []byte("top")
)

const (
	NameIdent = "peername"
	FileName  = "data.db"
)

func NewBoltDB(dbPathname string) (*BoltDB, error) {
	db, err := bolt.Open(dbPathname, 0660, nil)
	if err != nil {
		return nil, fmt.Errorf("[boltDB] Unable to open %s: %s", dbPathname, err)
	}
	// check the persistence version
	err = db.Update(func(tx *bolt.Tx) error {
		if top := tx.Bucket(topBucket); top == nil {
			top, err := tx.CreateBucket(topBucket)
			if err != nil {
				return err
			}
			if err := top.Put(versionIdent, persistenceVersion); err != nil {
				return err
			}
		} else {
			if checkVersion := top.Get(versionIdent); checkVersion != nil {
				if checkVersion[0] != persistenceVersion[0] {
					return fmt.Errorf("[boltDB] Cannot use persistence file %s - version %x", dbPathname, checkVersion)
				}
			}
		}
		return nil
	})
	return &BoltDB{db: db}, err
}

func (d *BoltDB) Load(ident string, data interface{}) (bool, error) {
	found := true
	err := d.db.View(func(tx *bolt.Tx) error {
		b := tx.Bucket(topBucket)
		v := b.Get([]byte(ident))
		if v == nil {
			found = false
			return nil
		}
		reader := bytes.NewReader(v)
		decoder := gob.NewDecoder(reader)
		err := decoder.Decode(data)
		return err
	})
	return found, err
}

func (d *BoltDB) Save(ident string, data interface{}) error {
	return d.db.Update(func(tx *bolt.Tx) error {
		b := tx.Bucket(topBucket)
		buf := new(bytes.Buffer)
		enc := gob.NewEncoder(buf)
		if err := enc.Encode(data); err != nil {
			return err
		}
		return b.Put([]byte(ident), buf.Bytes())
	})
}

func (d *BoltDB) Close() error {
	return d.db.Close()
}
<commit_msg>Allow opening persistence file read-only<commit_after>package db

import (
	"bytes"
	"encoding/gob"
	"fmt"
	"time"

	"github.com/boltdb/bolt"
)

type DB interface {
	Load(string, interface{}) (bool, error)
	Save(string, interface{}) error
}

type BoltDB struct {
	db *bolt.DB
}

var (
	versionIdent       = []byte("version")
	persistenceVersion = []byte{1, 0} // major.minor
	topBucket          = []byte("top")
)

const (
	NameIdent = "peername"
	FileName  = "data.db"
)

func NewBoltDB(dbPathname string) (*BoltDB, error) {
	db, err := bolt.Open(dbPathname, 0660, nil)
	if err != nil {
		return nil, fmt.Errorf("[boltDB] Unable to open %s: %s", dbPathname, err)
	}
	err = db.Update(checkVersion(false))
	return &BoltDB{db: db}, err
}

func NewBoltDBReadOnly(dbPathname string) (*BoltDB, error) {
	options := bolt.Options{Timeout: time.Millisecond * 50, ReadOnly: true}
	db, err := bolt.Open(dbPathname, 0660, &options)
	if err != nil {
		return nil, fmt.Errorf("[boltDB] Unable to open %s: %s", dbPathname, err)
	}
	err = db.View(checkVersion(true))
	if err != nil {
		return nil, fmt.Errorf("[boltDB] Cannot use persistence file %s: %s", dbPathname, err)
	}
	return &BoltDB{db: db}, nil
}

func checkVersion(readOnly bool) func(tx *bolt.Tx) error {
	return func(tx *bolt.Tx) error {
		if top := tx.Bucket(topBucket); top == nil {
			if readOnly {
				return fmt.Errorf("no top bucket")
			}
			top, err := tx.CreateBucket(topBucket)
			if err != nil {
				return err
			}
			if err := top.Put(versionIdent, persistenceVersion); err != nil {
				return err
			}
		} else {
			if checkVersion := top.Get(versionIdent); checkVersion != nil {
				if checkVersion[0] != persistenceVersion[0] {
					return fmt.Errorf("mismatched version %x", checkVersion)
				}
			}
		}
		return nil
	}
}

func (d *BoltDB) Load(ident string, data interface{}) (bool, error) {
	found := true
	err := d.db.View(func(tx *bolt.Tx) error {
		b := tx.Bucket(topBucket)
		v := b.Get([]byte(ident))
		if v == nil {
			found = false
			return nil
		}
		reader := bytes.NewReader(v)
		decoder := gob.NewDecoder(reader)
		err := decoder.Decode(data)
		return err
	})
	return found, err
}

func (d *BoltDB) Save(ident string, data interface{}) error {
	return d.db.Update(func(tx *bolt.Tx) error {
		b := tx.Bucket(topBucket)
		buf := new(bytes.Buffer)
		enc := gob.NewEncoder(buf)
		if err := enc.Encode(data); err != nil {
			return err
		}
		return b.Put([]byte(ident), buf.Bytes())
	})
}

func (d *BoltDB) Close() error {
	return d.db.Close()
}
<|endoftext|>


<commit_before>package vcr

import (
	"net/http"
)

// Client request record type
type RequestRecord {
	// Body of request
	Body string `yaml:"body"`

	// Request headers
	Headers http.Header `yaml:"headers"`

	// Request URL
	URL string `yaml:"url"`

	// Request method
	Method string `yaml:"method"`
}

// Server response record type
type ResponseRecord struct {
	// Body of response
	Body string `yaml:"body"`

	// Response headers
	Headers http.Header `yaml:"headers"`

	// Response status message
	Status string `yaml:"status"`

	// Response status code
	Code int `yaml:"code"`
}

// CassetteRecord type represents an HTTP interaction between
// client and server. It contains the client request and the
// returned response from the server for this request.
type CassetteRecord struct {
	Request  RequestRecord  `yaml:"request"`
	Response ResponseRecord `yaml:"response"`
}
<commit_msg>Commit vcr.Play helper function<commit_after>package vcr

import (
	"log"

	"github.com/dnaeon/vcr/recorder"
)

// Helper wrapper around recorder.Recorder type
// The returned Recorder must be stopped by the
// caller when no longer needed
func Play(cassette string) *recorder.Recorder {
	r = recorder.NewRecorder(cassette)

	err := r.Start()
	if err != nil {
		log.Fatal(err)
	}

	return r
}
<|endoftext|>