<commit_before>package data

import (
	"fmt"
	"log"
	"sort"
	"strings"
)

type Asset struct {
	Currency string `json:"currency"`
	Issuer   string `json:"issuer,omitempty"`
}

func NewAsset(s string) (*Asset, error) {
	if s == "XRP" {
		return &Asset{
			Currency: s,
		}, nil
	}
	parts := strings.Split(s, "/")
	if len(parts) != 2 {
		return nil, fmt.Errorf("bad asset: %s", s)
	}
	return &Asset{
		Currency: parts[0],
		Issuer:   parts[1],
	}, nil
}

func (a *Asset) IsNative() bool {
	return a.Currency == "XRP"
}

func (a Asset) String() string {
	if a.IsNative() {
		return a.Currency
	}
	return fmt.Sprintf("%s/%-34s", a.Currency, a.Issuer)
}

type OrderBookOffer struct {
	Offer
	OwnerFunds      Value          `json:"owner_funds"`
	Quality         NonNativeValue `json:"quality"`
	TakerGetsFunded *Amount        `json:"taker_gets_funded"`
	TakerPaysFunded *Amount        `json:"taker_pays_funded"`
}

type AccountOffer struct {
	Flags      LedgerEntryFlag `json:"flags"`
	Quality    NonNativeValue  `json:"quality"`
	Sequence   uint32          `json:"seq"`
	TakerGets  Amount          `json:"taker_gets"`
	TakerPays  Amount          `json:"taker_pays"`
	Expiration *uint32         `json:"expiration"`
}

type AccountOfferSlice []AccountOffer

func (s AccountOfferSlice) Len() int           { return len(s) }
func (s AccountOfferSlice) Swap(i, j int)      { s[i], s[j] = s[j], s[i] }
func (s AccountOfferSlice) Less(i, j int) bool { return s[i].Sequence > s[j].Sequence }

func (s AccountOfferSlice) GetSequences(pays, gets *Asset) []uint32 {
	// TODO: improve performance
	var sequences []uint32
	for i := range s {
		if s[i].TakerPays.Currency.String() == pays.Currency &&
			s[i].TakerPays.Issuer.String() == pays.Issuer &&
			s[i].TakerGets.Currency.String() == gets.Currency &&
			s[i].TakerGets.Issuer.String() == gets.Issuer {
			sequences = append(sequences, s[i].Sequence)
		}
	}
	return sequences
}

func (s AccountOfferSlice) Find(sequence uint32) int {
	return sort.Search(len(s), func(i int) bool {
		return s[i].Sequence <= sequence
	})
}

func (s AccountOfferSlice) Get(sequence uint32) *AccountOffer {
	if i := s.Find(sequence); i < len(s) && s[i].Sequence == sequence {
		return &s[i]
	}
	return nil
}

func defaultUint32(v *uint32) uint32 {
	if v == nil {
		return 0
	}
	return *v
}

func (s *AccountOfferSlice) Add(offer *Offer) bool {
	quality, err := offer.TakerPays.Divide(offer.TakerGets)
	if err != nil {
		panic(fmt.Sprintf("impossible quality: %s %s", offer.TakerPays, offer.TakerGets))
	}
	o := AccountOffer{
		Flags:      LedgerEntryFlag(defaultUint32((*uint32)(offer.Flags))),
		Sequence:   *offer.Sequence,
		TakerGets:  *offer.TakerGets,
		TakerPays:  *offer.TakerPays,
		Quality:    NonNativeValue{*quality.Value},
		Expiration: offer.Expiration,
	}
	i := s.Find(*offer.Sequence)
	switch {
	case i == len(*s):
		*s = append(*s, o)
		return true
	case (*s)[i].Sequence != *offer.Sequence:
		*s = append(*s, AccountOffer{})
		copy((*s)[i+1:], (*s)[i:])
		(*s)[i] = o
		return true
	default:
		return false
	}
}

func (s AccountOfferSlice) Update(offer *Offer) bool {
	if existing := s.Get(*offer.Sequence); existing != nil {
		existing.TakerGets = *offer.TakerGets
		existing.TakerPays = *offer.TakerPays
		return true
	}
	return false
}

func (s *AccountOfferSlice) Delete(offer *Offer) bool {
	if i := s.Find(*offer.Sequence); i < len(*s) && (*s)[i].Sequence == *offer.Sequence {
		*s = append((*s)[:i], (*s)[i+1:]...)
		return true
	}
	return false
}

type AccountLine struct {
	Account      Account        `json:"account"`
	Balance      NonNativeValue `json:"balance"`
	Currency     Currency       `json:"currency"`
	Limit        NonNativeValue `json:"limit"`
	LimitPeer    NonNativeValue `json:"limit_peer"`
	NoRipple     bool           `json:"no_ripple"`
	NoRipplePeer bool           `json:"no_ripple_peer"`
	QualityIn    uint32         `json:"quality_in"`
	QualityOut   uint32         `json:"quality_out"`
}

func (l *AccountLine) Asset() *Asset {
	return &Asset{
		Currency: l.Currency.String(),
		Issuer:   l.Account.String(),
	}
}

func (l *AccountLine) CompareByCurrencyAccount(other *AccountLine) int {
	if cmp := l.Currency.Compare(other.Currency); cmp != 0 {
		return cmp
	}
	return l.Account.Compare(other.Account)
}

func (l *AccountLine) CompareByCurrencyAmount(other *AccountLine) int {
	if cmp := l.Currency.Compare(other.Currency); cmp != 0 {
		return cmp
	}
	return l.Balance.Abs().Compare(*other.Balance.Abs())
}

type AccountLineSlice []AccountLine

type byCurrencyAccount struct {
	AccountLineSlice
}

func (by *byCurrencyAccount) Less(i, j int) bool {
	return by.AccountLineSlice[i].CompareByCurrencyAccount(&by.AccountLineSlice[j]) < 0
}

type byCurrencyAmount struct {
	AccountLineSlice
}

func (by *byCurrencyAmount) Less(i, j int) bool {
	return by.AccountLineSlice[i].CompareByCurrencyAmount(&by.AccountLineSlice[j]) < 0
}

func (s AccountLineSlice) Len() int               { return len(s) }
func (s AccountLineSlice) Swap(i, j int)          { s[i], s[j] = s[j], s[i] }
func (s AccountLineSlice) SortbyCurrencyAccount() { sort.Sort(&byCurrencyAccount{s}) }
func (s AccountLineSlice) SortByCurrencyAmount()  { sort.Sort(&byCurrencyAmount{s}) }

func (s AccountLineSlice) Find(account Account, currency Currency) int {
	return sort.Search(len(s), func(i int) bool {
		switch s[i].Currency.Compare(currency) {
		case 1:
			return true
		case -1:
			return false
		default:
			return s[i].Account.Compare(account) >= 0
		}
	})
}

type highLowFunc func(balance, limit, limitPeer *Amount, noRipple, noRipplePeer bool, qualityIn, qualityOut uint32) bool

func highLow(account Account, rs *RippleState, f highLowFunc) bool {
	high, low := account.Equals(rs.HighLimit.Issuer), account.Equals(rs.LowLimit.Issuer)
	switch {
	case high == low:
		return false
	case high:
		noRipple, noRipplePeer := *rs.Flags&LsHighNoRipple > 0, *rs.Flags&LsLowNoRipple > 0
		return f(rs.Balance.Negate(), rs.HighLimit, rs.LowLimit, noRipple, noRipplePeer, defaultUint32(rs.HighQualityIn), defaultUint32(rs.HighQualityOut))
	default:
		noRipple, noRipplePeer := *rs.Flags&LsLowNoRipple > 0, *rs.Flags&LsHighNoRipple > 0
		return f(rs.Balance, rs.LowLimit, rs.HighLimit, noRipple, noRipplePeer, defaultUint32(rs.LowQualityIn), defaultUint32(rs.LowQualityOut))
	}
}

func (s *AccountLineSlice) Add(account Account, rs *RippleState) bool {
	var add = func(balance, limit, limitPeer *Amount, noRipple, noRipplePeer bool, qualityIn, qualityOut uint32) bool {
		line := AccountLine{
			Account:      limitPeer.Issuer,
			Balance:      NonNativeValue{*balance.Value},
			Currency:     limit.Currency,
			Limit:        NonNativeValue{*limit.Value},
			LimitPeer:    NonNativeValue{*limitPeer.Value},
			NoRipple:     noRipple,
			NoRipplePeer: noRipplePeer,
			QualityIn:    qualityIn,
			QualityOut:   qualityOut,
		}
		i := s.Find(limitPeer.Issuer, limitPeer.Currency)
		log.Println("Add", account, limitPeer.Issuer, limitPeer.Currency, len(*s), i)
		switch {
		case i == len(*s):
			*s = append(*s, line)
			return true
		case (*s)[i].Account.Equals(limitPeer.Issuer) && (*s)[i].Currency.Equals(limitPeer.Currency):
			return false
		default:
			*s = append((*s)[:i], append(AccountLineSlice{line}, (*s)[i:]...)...)
			return true
		}
	}
	return highLow(account, rs, add)
}

func (s AccountLineSlice) Update(account Account, rs *RippleState) bool {
	var update = func(balance, limit, limitPeer *Amount, noripple, noRipplePeer bool, qualityIn, qualityOut uint32) bool {
		i := s.Find(limitPeer.Issuer, limitPeer.Currency)
		if i == len(s) || !s[i].Account.Equals(limitPeer.Issuer) || !s[i].Currency.Equals(limitPeer.Currency) {
			return false
		}
		s[i].Balance = NonNativeValue{*balance.Value}
		s[i].Limit = NonNativeValue{*limit.Value}
		s[i].LimitPeer = NonNativeValue{*limitPeer.Value}
		s[i].NoRipple = noripple
		s[i].NoRipplePeer = noRipplePeer
		return true
	}
	return highLow(account, rs, update)
}

func (s *AccountLineSlice) Delete(account Account, rs *RippleState) bool {
	var del = func(balance, limit, limitPeer *Amount, noripple, noRipplePeer bool, qualityIn, qualityOut uint32) bool {
		i := s.Find(limitPeer.Issuer, limitPeer.Currency)
		log.Println("Delete:", account, limitPeer.Issuer, limitPeer.Currency, len(*s), i)
		if i < len(*s) && (*s)[i].Account.Equals(limitPeer.Issuer) && (*s)[i].Currency.Equals(limitPeer.Currency) {
			*s = append((*s)[:i], (*s)[i+1:]...)
			return true
		}
		return false
	}
	return highLow(account, rs, del)
}
<commit_msg>Fix Asset formatting<commit_after>package data

import (
	"fmt"
	"log"
	"sort"
	"strings"
)

type Asset struct {
	Currency string `json:"currency"`
	Issuer   string `json:"issuer,omitempty"`
}

func NewAsset(s string) (*Asset, error) {
	if s == "XRP" {
		return &Asset{
			Currency: s,
		}, nil
	}
	parts := strings.Split(s, "/")
	if len(parts) != 2 {
		return nil, fmt.Errorf("bad asset: %s", s)
	}
	return &Asset{
		Currency: parts[0],
		Issuer:   parts[1],
	}, nil
}

func (a *Asset) IsNative() bool {
	return a.Currency == "XRP"
}

func (a Asset) String() string {
	if a.IsNative() {
		return a.Currency
	}
	return fmt.Sprintf("%s/%s", a.Currency, a.Issuer)
}

type OrderBookOffer struct {
	Offer
	OwnerFunds      Value          `json:"owner_funds"`
	Quality         NonNativeValue `json:"quality"`
	TakerGetsFunded *Amount        `json:"taker_gets_funded"`
	TakerPaysFunded *Amount        `json:"taker_pays_funded"`
}

type AccountOffer struct {
	Flags      LedgerEntryFlag `json:"flags"`
	Quality    NonNativeValue  `json:"quality"`
	Sequence   uint32          `json:"seq"`
	TakerGets  Amount          `json:"taker_gets"`
	TakerPays  Amount          `json:"taker_pays"`
	Expiration *uint32         `json:"expiration"`
}

type AccountOfferSlice []AccountOffer

func (s AccountOfferSlice) Len() int           { return len(s) }
func (s AccountOfferSlice) Swap(i, j int)      { s[i], s[j] = s[j], s[i] }
func (s AccountOfferSlice) Less(i, j int) bool { return s[i].Sequence > s[j].Sequence }

func (s AccountOfferSlice) GetSequences(pays, gets *Asset) []uint32 {
	// TODO: improve performance
	var sequences []uint32
	for i := range s {
		if s[i].TakerPays.Currency.String() == pays.Currency &&
			s[i].TakerPays.Issuer.String() == pays.Issuer &&
			s[i].TakerGets.Currency.String() == gets.Currency &&
			s[i].TakerGets.Issuer.String() == gets.Issuer {
			sequences = append(sequences, s[i].Sequence)
		}
	}
	return sequences
}

func (s AccountOfferSlice) Find(sequence uint32) int {
	return sort.Search(len(s), func(i int) bool {
		return s[i].Sequence <= sequence
	})
}

func (s AccountOfferSlice) Get(sequence uint32) *AccountOffer {
	if i := s.Find(sequence); i < len(s) && s[i].Sequence == sequence {
		return &s[i]
	}
	return nil
}

func defaultUint32(v *uint32) uint32 {
	if v == nil {
		return 0
	}
	return *v
}

func (s *AccountOfferSlice) Add(offer *Offer) bool {
	quality, err := offer.TakerPays.Divide(offer.TakerGets)
	if err != nil {
		panic(fmt.Sprintf("impossible quality: %s %s", offer.TakerPays, offer.TakerGets))
	}
	o := AccountOffer{
		Flags:      LedgerEntryFlag(defaultUint32((*uint32)(offer.Flags))),
		Sequence:   *offer.Sequence,
		TakerGets:  *offer.TakerGets,
		TakerPays:  *offer.TakerPays,
		Quality:    NonNativeValue{*quality.Value},
		Expiration: offer.Expiration,
	}
	i := s.Find(*offer.Sequence)
	switch {
	case i == len(*s):
		*s = append(*s, o)
		return true
	case (*s)[i].Sequence != *offer.Sequence:
		*s = append(*s, AccountOffer{})
		copy((*s)[i+1:], (*s)[i:])
		(*s)[i] = o
		return true
	default:
		return false
	}
}

func (s AccountOfferSlice) Update(offer *Offer) bool {
	if existing := s.Get(*offer.Sequence); existing != nil {
		existing.TakerGets = *offer.TakerGets
		existing.TakerPays = *offer.TakerPays
		return true
	}
	return false
}

func (s *AccountOfferSlice) Delete(offer *Offer) bool {
	if i := s.Find(*offer.Sequence); i < len(*s) && (*s)[i].Sequence == *offer.Sequence {
		*s = append((*s)[:i], (*s)[i+1:]...)
		return true
	}
	return false
}

type AccountLine struct {
	Account      Account        `json:"account"`
	Balance      NonNativeValue `json:"balance"`
	Currency     Currency       `json:"currency"`
	Limit        NonNativeValue `json:"limit"`
	LimitPeer    NonNativeValue `json:"limit_peer"`
	NoRipple     bool           `json:"no_ripple"`
	NoRipplePeer bool           `json:"no_ripple_peer"`
	QualityIn    uint32         `json:"quality_in"`
	QualityOut   uint32         `json:"quality_out"`
}

func (l *AccountLine) Asset() *Asset {
	return &Asset{
		Currency: l.Currency.String(),
		Issuer:   l.Account.String(),
	}
}

func (l *AccountLine) CompareByCurrencyAccount(other *AccountLine) int {
	if cmp := l.Currency.Compare(other.Currency); cmp != 0 {
		return cmp
	}
	return l.Account.Compare(other.Account)
}

func (l *AccountLine) CompareByCurrencyAmount(other *AccountLine) int {
	if cmp := l.Currency.Compare(other.Currency); cmp != 0 {
		return cmp
	}
	return l.Balance.Abs().Compare(*other.Balance.Abs())
}

type AccountLineSlice []AccountLine

type byCurrencyAccount struct {
	AccountLineSlice
}

func (by *byCurrencyAccount) Less(i, j int) bool {
	return by.AccountLineSlice[i].CompareByCurrencyAccount(&by.AccountLineSlice[j]) < 0
}

type byCurrencyAmount struct {
	AccountLineSlice
}

func (by *byCurrencyAmount) Less(i, j int) bool {
	return by.AccountLineSlice[i].CompareByCurrencyAmount(&by.AccountLineSlice[j]) < 0
}

func (s AccountLineSlice) Len() int               { return len(s) }
func (s AccountLineSlice) Swap(i, j int)          { s[i], s[j] = s[j], s[i] }
func (s AccountLineSlice) SortbyCurrencyAccount() { sort.Sort(&byCurrencyAccount{s}) }
func (s AccountLineSlice) SortByCurrencyAmount()  { sort.Sort(&byCurrencyAmount{s}) }

func (s AccountLineSlice) Find(account Account, currency Currency) int {
	return sort.Search(len(s), func(i int) bool {
		switch s[i].Currency.Compare(currency) {
		case 1:
			return true
		case -1:
			return false
		default:
			return s[i].Account.Compare(account) >= 0
		}
	})
}

type highLowFunc func(balance, limit, limitPeer *Amount, noRipple, noRipplePeer bool, qualityIn, qualityOut uint32) bool

func highLow(account Account, rs *RippleState, f highLowFunc) bool {
	high, low := account.Equals(rs.HighLimit.Issuer), account.Equals(rs.LowLimit.Issuer)
	switch {
	case high == low:
		return false
	case high:
		noRipple, noRipplePeer := *rs.Flags&LsHighNoRipple > 0, *rs.Flags&LsLowNoRipple > 0
		return f(rs.Balance.Negate(), rs.HighLimit, rs.LowLimit, noRipple, noRipplePeer, defaultUint32(rs.HighQualityIn), defaultUint32(rs.HighQualityOut))
	default:
		noRipple, noRipplePeer := *rs.Flags&LsLowNoRipple > 0, *rs.Flags&LsHighNoRipple > 0
		return f(rs.Balance, rs.LowLimit, rs.HighLimit, noRipple, noRipplePeer, defaultUint32(rs.LowQualityIn), defaultUint32(rs.LowQualityOut))
	}
}

func (s *AccountLineSlice) Add(account Account, rs *RippleState) bool {
	var add = func(balance, limit, limitPeer *Amount, noRipple, noRipplePeer bool, qualityIn, qualityOut uint32) bool {
		line := AccountLine{
			Account:      limitPeer.Issuer,
			Balance:      NonNativeValue{*balance.Value},
			Currency:     limit.Currency,
			Limit:        NonNativeValue{*limit.Value},
			LimitPeer:    NonNativeValue{*limitPeer.Value},
			NoRipple:     noRipple,
			NoRipplePeer: noRipplePeer,
			QualityIn:    qualityIn,
			QualityOut:   qualityOut,
		}
		i := s.Find(limitPeer.Issuer, limitPeer.Currency)
		log.Println("Add", account, limitPeer.Issuer, limitPeer.Currency, len(*s), i)
		switch {
		case i == len(*s):
			*s = append(*s, line)
			return true
		case (*s)[i].Account.Equals(limitPeer.Issuer) && (*s)[i].Currency.Equals(limitPeer.Currency):
			return false
		default:
			*s = append((*s)[:i], append(AccountLineSlice{line}, (*s)[i:]...)...)
			return true
		}
	}
	return highLow(account, rs, add)
}

func (s AccountLineSlice) Update(account Account, rs *RippleState) bool {
	var update = func(balance, limit, limitPeer *Amount, noripple, noRipplePeer bool, qualityIn, qualityOut uint32) bool {
		i := s.Find(limitPeer.Issuer, limitPeer.Currency)
		if i == len(s) || !s[i].Account.Equals(limitPeer.Issuer) || !s[i].Currency.Equals(limitPeer.Currency) {
			return false
		}
		s[i].Balance = NonNativeValue{*balance.Value}
		s[i].Limit = NonNativeValue{*limit.Value}
		s[i].LimitPeer = NonNativeValue{*limitPeer.Value}
		s[i].NoRipple = noripple
		s[i].NoRipplePeer = noRipplePeer
		return true
	}
	return highLow(account, rs, update)
}

func (s *AccountLineSlice) Delete(account Account, rs *RippleState) bool {
	var del = func(balance, limit, limitPeer *Amount, noripple, noRipplePeer bool, qualityIn, qualityOut uint32) bool {
		i := s.Find(limitPeer.Issuer, limitPeer.Currency)
		log.Println("Delete:", account, limitPeer.Issuer, limitPeer.Currency, len(*s), i)
		if i < len(*s) && (*s)[i].Account.Equals(limitPeer.Issuer) && (*s)[i].Currency.Equals(limitPeer.Currency) {
			*s = append((*s)[:i], (*s)[i+1:]...)
			return true
		}
		return false
	}
	return highLow(account, rs, del)
}
<|endoftext|>


<commit_before>package db

import (
	"fmt"
	"strings"

	"github.com/greensnark/go-sequell/qyaml"
	"github.com/greensnark/go-sequell/schema"
)

type CrawlSchema struct {
	FieldParser             *FieldParser
	Tables                  []*CrawlTable
	TableVariantPrefixes    []string
	VariantNamePrefixMap    map[string]string
	LookupTables            []*LookupTable
	FieldNameLookupTableMap map[string]*LookupTable
}

type CrawlTable struct {
	Name             string
	Fields           []*Field
	PrimaryKeyField  *Field
	CompositeIndexes []*Index
}

type Index struct {
	Columns []string
}

type LookupTable struct {
	CrawlTable
	ReferencingFields []*Field
	DerivedFields     []*Field
}

func MustLoadSchema(schemaDef qyaml.Yaml) *CrawlSchema {
	sch, err := LoadSchema(schemaDef)
	if err != nil {
		panic(err)
	}
	return sch
}

func LoadSchema(schemaDef qyaml.Yaml) (*CrawlSchema, error) {
	schema := CrawlSchema{
		FieldParser:             NewFieldParser(schemaDef),
		FieldNameLookupTableMap: map[string]*LookupTable{},
	}
	schema.ParseVariants(schemaDef.StringMap("game-type-prefixes"))
	for iname, definition := range schemaDef.Map("lookup-tables") {
		name := iname.(string)
		err := schema.ParseLookupTable(name, definition)
		if err != nil {
			return nil, err
		}
	}
	for _, table := range schemaDef.StringSlice("query-tables") {
		if err := schema.ParseTable(table, schemaDef); err != nil {
			return nil, err
		}
	}
	return &schema, nil
}

func (s *CrawlSchema) LookupTable(name string) *LookupTable {
	for _, lt := range s.LookupTables {
		if lt.Name == name {
			return lt
		}
	}
	return nil
}

func (s *CrawlSchema) ParseVariants(variantMap map[string]string) {
	s.VariantNamePrefixMap = variantMap
	s.TableVariantPrefixes = make([]string, len(variantMap))
	i := 0
	for _, prefix := range variantMap {
		s.TableVariantPrefixes[i] = prefix
		i++
	}
	if len(s.TableVariantPrefixes) == 0 {
		s.TableVariantPrefixes = []string{""}
	}
}

func (s *CrawlSchema) ParseTable(name string, schemaDef qyaml.Yaml) (err error) {
	fields, err := s.ParseTableFields(name, schemaDef)
	if err != nil {
		return err
	}

	indexes, err := s.ParseCompositeIndexes(name, fields, schemaDef)
	if err != nil {
		return err
	}
	table := &CrawlTable{
		Name:             name,
		Fields:           fields,
		CompositeIndexes: indexes,
	}
	for _, field := range fields {
		if field.PrimaryKey {
			table.PrimaryKeyField = field
			break
		}
	}
	s.RegisterLookupTables(fields)
	s.Tables = append(s.Tables, table)
	return nil
}

func (s *CrawlSchema) ParseTableFields(name string, schemaDef qyaml.Yaml) (tableFields []*Field, err error) {
	annotatedFieldNames := schemaDef.StringSlice(name + "-fields-with-type")
	return s.ParseFields(annotatedFieldNames)
}

func (s *CrawlSchema) ParseFields(annotatedFieldNames []string) (tableFields []*Field, err error) {
	tableFields = make([]*Field, len(annotatedFieldNames))
	for i, field := range annotatedFieldNames {
		tableFields[i], err = s.FieldParser.ParseField(field)
		if err != nil {
			return nil, err
		}
	}
	return
}

func (s *CrawlSchema) ParseCompositeIndexes(name string, fields []*Field, schemaDef qyaml.Yaml) ([]*Index, error) {
	indexDefs := schemaDef.Slice(name + "-indexes")
	compositeIndexes := make([]*Index, len(indexDefs))

	findField := func(name string) *Field {
		for _, f := range fields {
			if f.Name == name {
				return f
			}
		}
		return nil
	}

	fieldSqlNames := func(names []string) ([]string, error) {
		res := make([]string, len(names))
		for i, name := range names {
			field := findField(name)
			if field == nil {
				return nil, fmt.Errorf("no field definition for '%s'", name)
			}
			res[i] = field.RefName()
		}
		return res, nil
	}

	for i, def := range indexDefs {
		fields := qyaml.IStringSlice(def)
		if len(fields) == 0 {
			return nil, fmt.Errorf("No fields defined for index on %s with spec %#v\n",
				name, def)
		}
		sqlNames, err := fieldSqlNames(fields)
		if err != nil {
			return nil, err
		}
		compositeIndexes[i] = &Index{
			Columns: sqlNames,
		}
	}
	return compositeIndexes, nil
}

func (s *CrawlSchema) ParseLookupTable(name string, defn interface{}) error {
	switch tdef := defn.(type) {
	case []interface{}:
		fields, err := s.ParseFields(qyaml.IStringSlice(tdef))
		if err != nil {
			return err
		}
		s.AddLookupTable(name, s.markReferenceFields(fields), nil)
	case map[interface{}]interface{}:
		fields, err := s.ParseFields(qyaml.IStringSlice(tdef["fields"]))
		if err != nil {
			return err
		}
		generatedFields, err := s.ParseFields(
			qyaml.IStringSlice(tdef["generated-fields"]))
		if err != nil {
			return err
		}
		s.AddLookupTable(name, s.markReferenceFields(fields), generatedFields)
	}
	return nil
}

func (s *CrawlSchema) markReferenceFields(fields []*Field) []*Field {
	for _, f := range fields {
		f.ForeignKeyLookup = true
	}
	return fields
}

func (s *CrawlSchema) RegisterLookupTables(fields []*Field) {
	for _, field := range fields {
		if field.ForeignKeyLookup {
			lookup := s.FindLookupTableForField(field.Name)
			if lookup == nil {
				lookup = s.AddLookupTable(field.Name, []*Field{field}, nil)
			}
			field.ForeignKeyTable = lookup.TableName()
		}
	}
}

func (s *CrawlSchema) FindLookupTableForField(fieldName string) *LookupTable {
	return s.FieldNameLookupTableMap[fieldName]
}

func (s *CrawlSchema) AddLookupTable(name string, fields []*Field, generatedFields []*Field) *LookupTable {
	idField, err := s.FieldParser.ParseField("idIB%*")
	if err != nil {
		panic(err)
	}
	tableFields := make([]*Field, 2+len(generatedFields))
	tableFields[0] = idField
	tableFields[1] = fields[0]
	for i, g := range generatedFields {
		tableFields[i+2] = g
	}
	lookupTable := &LookupTable{
		CrawlTable: CrawlTable{
			Name:   name,
			Fields: tableFields,
		},
		ReferencingFields: fields,
		DerivedFields:     generatedFields,
	}
	for _, field := range fields {
		s.FieldNameLookupTableMap[field.Name] = lookupTable
	}
	s.LookupTables = append(s.LookupTables, lookupTable)
	return lookupTable
}

func IndexName(table string, fields []string, unique bool) string {
	base := "ind_" + table
	if unique {
		base += "_uniq"
	}
	return base + "_" + strings.Join(fields, "_")
}

func (s *CrawlSchema) Schema() *schema.Schema {
	return &schema.Schema{
		Tables: s.SchemaTables(),
	}
}

func (s *CrawlSchema) SchemaTables() []*schema.Table {
	tables := make([]*schema.Table,
		len(s.LookupTables)+len(s.Tables)*len(s.TableVariantPrefixes))
	for i, lookup := range s.LookupTables {
		tables[i] = lookup.SchemaTable()
	}

	i := len(s.LookupTables)
	for _, prefix := range s.TableVariantPrefixes {
		for _, table := range s.Tables {
			tables[i] = table.SchemaTablePrefixed(prefix)
			i++
		}
	}

	return tables
}

func (t *CrawlTable) SchemaTablePrefixed(prefix string) *schema.Table {
	tableName := prefix + t.Name
	return &schema.Table{
		Name:        tableName,
		Columns:     t.SchemaColumns(),
		Indexes:     t.SchemaIndexes(prefix),
		Constraints: t.SchemaConstraints(tableName),
	}
}

func (t *CrawlTable) SchemaColumns() []*schema.Column {
	cols := make([]*schema.Column, len(t.Fields))
	for i, field := range t.Fields {
		cols[i] = field.SchemaColumn()
	}
	return cols
}

func (t *CrawlTable) SchemaIndexes(prefix string) []*schema.Index {
	indexes := []*schema.Index{}
	add := func(i *schema.Index) {
		if i != nil {
			indexes = append(indexes, i)
		}
	}

	tableName := prefix + t.Name
	for _, compIndex := range t.CompositeIndexes {
		add(&schema.Index{
			Name:      t.SchemaIndexName(prefix, compIndex.Columns, false),
			TableName: tableName,
			Columns:   compIndex.Columns,
		})
	}

	for _, field := range t.Fields {
		if field.NeedsIndex() {
			cols := []string{field.RefName()}
			add(&schema.Index{
				Name:      t.SchemaIndexName(prefix, cols, false),
				TableName: tableName,
				Columns:   cols,
			})
		}
	}

	return indexes
}

func (t *CrawlTable) SchemaIndexName(prefix string, columns []string, unique bool) string {
	return IndexName(prefix+t.Name, columns, unique)
}

func (t *CrawlTable) SchemaConstraints(table string) []schema.Constraint {
	constraints := []schema.Constraint{}
	add := func(c schema.Constraint) {
		if c != nil {
			constraints = append(constraints, c)
		}
	}

	if t.PrimaryKeyField != nil {
		add(schema.PrimaryKeyConstraint{
			ConstraintName: table + "_pk",
			Column:         t.PrimaryKeyField.SqlName,
		})
	}
	for _, field := range t.Fields {
		if field.ForeignKeyLookup {
			add(field.ForeignKeyConstraint(table))
		}
	}

	return constraints
}

func (l *LookupTable) CaseSensitive() bool {
	return l.LookupField().CaseSensitive
}

// ReferencingFieldCount returns the number of fields in a single row
// that may refer to the lookup table.
func (l *LookupTable) ReferencingFieldCount() int {
	return len(l.ReferencingFields)
}

func (l *LookupTable) LookupField() *Field {
	return l.Fields[1]
}

func (l *LookupTable) TableName() string {
	return "l_" + l.Name
}

func (l *LookupTable) SchemaTable() *schema.Table {
	return &schema.Table{
		Name:    l.TableName(),
		Columns: l.SchemaColumns(),
		Constraints: []schema.Constraint{
			schema.PrimaryKeyConstraint{
				ConstraintName: l.TableName() + "_pk",
				Column:         "id",
			},
		},
		Indexes: l.SchemaIndexes(),
	}
}

func (l *LookupTable) SchemaColumns() []*schema.Column {
	cols := make([]*schema.Column, len(l.Fields))
	for i, field := range l.Fields {
		cols[i] = field.LookupSchemaColumn()
	}
	return cols
}

func (l *LookupTable) SchemaIndexes() []*schema.Index {
	indexes := []*schema.Index{
		&schema.Index{
			Name:      IndexName(l.TableName(), []string{l.Fields[1].SqlName}, true),
			TableName: l.TableName(),
			Columns:   []string{l.Fields[1].SqlName},
			Unique:    true,
			Force:     l.Fields[1].ForceIndex,
		},
	}
	for _, field := range l.Fields {
		if field.SqlLookupExpr != "" {
			indexes = append(indexes, &schema.Index{
				Name:      IndexName(l.TableName(), []string{field.SqlName}, false),
				TableName: l.TableName(),
				Columns:   []string{field.SqlLookupExpr},
			})
		}
	}
	return indexes
}
<commit_msg>Force all lookup tables to be indexed for bulk load.<commit_after>package db

import (
	"fmt"
	"strings"

	"github.com/greensnark/go-sequell/qyaml"
	"github.com/greensnark/go-sequell/schema"
)

type CrawlSchema struct {
	FieldParser             *FieldParser
	Tables                  []*CrawlTable
	TableVariantPrefixes    []string
	VariantNamePrefixMap    map[string]string
	LookupTables            []*LookupTable
	FieldNameLookupTableMap map[string]*LookupTable
}

type CrawlTable struct {
	Name             string
	Fields           []*Field
	PrimaryKeyField  *Field
	CompositeIndexes []*Index
}

type Index struct {
	Columns []string
}

type LookupTable struct {
	CrawlTable
	ReferencingFields []*Field
	DerivedFields     []*Field
}

func MustLoadSchema(schemaDef qyaml.Yaml) *CrawlSchema {
	sch, err := LoadSchema(schemaDef)
	if err != nil {
		panic(err)
	}
	return sch
}

func LoadSchema(schemaDef qyaml.Yaml) (*CrawlSchema, error) {
	schema := CrawlSchema{
		FieldParser:             NewFieldParser(schemaDef),
		FieldNameLookupTableMap: map[string]*LookupTable{},
	}
	schema.ParseVariants(schemaDef.StringMap("game-type-prefixes"))
	for iname, definition := range schemaDef.Map("lookup-tables") {
		name := iname.(string)
		err := schema.ParseLookupTable(name, definition)
		if err != nil {
			return nil, err
		}
	}
	for _, table := range schemaDef.StringSlice("query-tables") {
		if err := schema.ParseTable(table, schemaDef); err != nil {
			return nil, err
		}
	}
	return &schema, nil
}

func (s *CrawlSchema) LookupTable(name string) *LookupTable {
	for _, lt := range s.LookupTables {
		if lt.Name == name {
			return lt
		}
	}
	return nil
}

func (s *CrawlSchema) ParseVariants(variantMap map[string]string) {
	s.VariantNamePrefixMap = variantMap
	s.TableVariantPrefixes = make([]string, len(variantMap))
	i := 0
	for _, prefix := range variantMap {
		s.TableVariantPrefixes[i] = prefix
		i++
	}
	if len(s.TableVariantPrefixes) == 0 {
		s.TableVariantPrefixes = []string{""}
	}
}

func (s *CrawlSchema) ParseTable(name string, schemaDef qyaml.Yaml) (err error) {
	fields, err := s.ParseTableFields(name, schemaDef)
	if err != nil {
		return err
	}

	indexes, err := s.ParseCompositeIndexes(name, fields, schemaDef)
	if err != nil {
		return err
	}
	table := &CrawlTable{
		Name:             name,
		Fields:           fields,
		CompositeIndexes: indexes,
	}
	for _, field := range fields {
		if field.PrimaryKey {
			table.PrimaryKeyField = field
			break
		}
	}
	s.RegisterLookupTables(fields)
	s.Tables = append(s.Tables, table)
	return nil
}

func (s *CrawlSchema) ParseTableFields(name string, schemaDef qyaml.Yaml) (tableFields []*Field, err error) {
	annotatedFieldNames := schemaDef.StringSlice(name + "-fields-with-type")
	return s.ParseFields(annotatedFieldNames)
}

func (s *CrawlSchema) ParseFields(annotatedFieldNames []string) (tableFields []*Field, err error) {
	tableFields = make([]*Field, len(annotatedFieldNames))
	for i, field := range annotatedFieldNames {
		tableFields[i], err = s.FieldParser.ParseField(field)
		if err != nil {
			return nil, err
		}
	}
	return
}

func (s *CrawlSchema) ParseCompositeIndexes(name string, fields []*Field, schemaDef qyaml.Yaml) ([]*Index, error) {
	indexDefs := schemaDef.Slice(name + "-indexes")
	compositeIndexes := make([]*Index, len(indexDefs))

	findField := func(name string) *Field {
		for _, f := range fields {
			if f.Name == name {
				return f
			}
		}
		return nil
	}

	fieldSqlNames := func(names []string) ([]string, error) {
		res := make([]string, len(names))
		for i, name := range names {
			field := findField(name)
			if field == nil {
				return nil, fmt.Errorf("no field definition for '%s'", name)
			}
			res[i] = field.RefName()
		}
		return res, nil
	}

	for i, def := range indexDefs {
		fields := qyaml.IStringSlice(def)
		if len(fields) == 0 {
			return nil, fmt.Errorf("No fields defined for index on %s with spec %#v\n",
				name, def)
		}
		sqlNames, err := fieldSqlNames(fields)
		if err != nil {
			return nil, err
		}
		compositeIndexes[i] = &Index{
			Columns: sqlNames,
		}
	}
	return compositeIndexes, nil
}

func (s *CrawlSchema) ParseLookupTable(name string, defn interface{}) error {
	switch tdef := defn.(type) {
	case []interface{}:
		fields, err := s.ParseFields(qyaml.IStringSlice(tdef))
		if err != nil {
			return err
		}
		s.AddLookupTable(name, s.markReferenceFields(fields), nil)
	case map[interface{}]interface{}:
		fields, err := s.ParseFields(qyaml.IStringSlice(tdef["fields"]))
		if err != nil {
			return err
		}
		generatedFields, err := s.ParseFields(
			qyaml.IStringSlice(tdef["generated-fields"]))
		if err != nil {
			return err
		}
		s.AddLookupTable(name, s.markReferenceFields(fields), generatedFields)
	}
	return nil
}

func (s *CrawlSchema) markReferenceFields(fields []*Field) []*Field {
	for _, f := range fields {
		f.ForeignKeyLookup = true
	}
	return fields
}

func (s *CrawlSchema) RegisterLookupTables(fields []*Field) {
	for _, field := range fields {
		if field.ForeignKeyLookup {
			lookup := s.FindLookupTableForField(field.Name)
			if lookup == nil {
				lookup = s.AddLookupTable(field.Name, []*Field{field}, nil)
			}
			field.ForeignKeyTable = lookup.TableName()
		}
	}
}

func (s *CrawlSchema) FindLookupTableForField(fieldName string) *LookupTable {
	return s.FieldNameLookupTableMap[fieldName]
}

func (s *CrawlSchema) AddLookupTable(name string, fields []*Field, generatedFields []*Field) *LookupTable {
	idField, err := s.FieldParser.ParseField("idIB%*")
	if err != nil {
		panic(err)
	}
	tableFields := make([]*Field, 2+len(generatedFields))
	tableFields[0] = idField
	tableFields[1] = fields[0]
	for i, g := range generatedFields {
		tableFields[i+2] = g
	}
	lookupTable := &LookupTable{
		CrawlTable: CrawlTable{
			Name:   name,
			Fields: tableFields,
		},
		ReferencingFields: fields,
		DerivedFields:     generatedFields,
	}
	for _, field := range fields {
		s.FieldNameLookupTableMap[field.Name] = lookupTable
	}
	s.LookupTables = append(s.LookupTables, lookupTable)
	return lookupTable
}

func IndexName(table string, fields []string, unique bool) string {
	base := "ind_" + table
	if unique {
		base += "_uniq"
	}
	return base + "_" + strings.Join(fields, "_")
}

func (s *CrawlSchema) Schema() *schema.Schema {
	return &schema.Schema{
		Tables: s.SchemaTables(),
	}
}

func (s *CrawlSchema) SchemaTables() []*schema.Table {
	tables := make([]*schema.Table,
		len(s.LookupTables)+len(s.Tables)*len(s.TableVariantPrefixes))
	for i, lookup := range s.LookupTables {
		tables[i] = lookup.SchemaTable()
	}

	i := len(s.LookupTables)
	for _, prefix := range s.TableVariantPrefixes {
		for _, table := range s.Tables {
			tables[i] = table.SchemaTablePrefixed(prefix)
			i++
		}
	}

	return tables
}

func (t *CrawlTable) SchemaTablePrefixed(prefix string) *schema.Table {
	tableName := prefix + t.Name
	return &schema.Table{
		Name:        tableName,
		Columns:     t.SchemaColumns(),
		Indexes:     t.SchemaIndexes(prefix),
		Constraints: t.SchemaConstraints(tableName),
	}
}

func (t *CrawlTable) SchemaColumns() []*schema.Column {
	cols := make([]*schema.Column, len(t.Fields))
	for i, field := range t.Fields {
		cols[i] = field.SchemaColumn()
	}
	return cols
}

func (t *CrawlTable) SchemaIndexes(prefix string) []*schema.Index {
	indexes := []*schema.Index{}
	add := func(i *schema.Index) {
		if i != nil {
			indexes = append(indexes, i)
		}
	}

	tableName := prefix + t.Name
	for _, compIndex := range t.CompositeIndexes {
		add(&schema.Index{
			Name:      t.SchemaIndexName(prefix, compIndex.Columns, false),
			TableName: tableName,
			Columns:   compIndex.Columns,
		})
	}

	for _, field := range t.Fields {
		if field.NeedsIndex() {
			cols := []string{field.RefName()}
			add(&schema.Index{
				Name:      t.SchemaIndexName(prefix, cols, false),
				TableName: tableName,
				Columns:   cols,
			})
		}
	}

	return indexes
}

func (t *CrawlTable) SchemaIndexName(prefix string, columns []string, unique bool) string {
	return IndexName(prefix+t.Name, columns, unique)
}

func (t *CrawlTable) SchemaConstraints(table string) []schema.Constraint {
	constraints := []schema.Constraint{}
	add := func(c schema.Constraint) {
		if c != nil {
			constraints = append(constraints, c)
		}
	}

	if t.PrimaryKeyField != nil {
		add(schema.PrimaryKeyConstraint{
			ConstraintName: table + "_pk",
			Column:         t.PrimaryKeyField.SqlName,
		})
	}
	for _, field := range t.Fields {
		if field.ForeignKeyLookup {
			add(field.ForeignKeyConstraint(table))
		}
	}

	return constraints
}

func (l *LookupTable) CaseSensitive() bool {
	return l.LookupField().CaseSensitive
}

// ReferencingFieldCount returns the number of fields in a single row
// that may refer to the lookup table.
func (l *LookupTable) ReferencingFieldCount() int {
	return len(l.ReferencingFields)
}

func (l *LookupTable) LookupField() *Field {
	return l.Fields[1]
}

func (l *LookupTable) TableName() string {
	return "l_" + l.Name
}

func (l *LookupTable) SchemaTable() *schema.Table {
	return &schema.Table{
		Name:    l.TableName(),
		Columns: l.SchemaColumns(),
		Constraints: []schema.Constraint{
			schema.PrimaryKeyConstraint{
				ConstraintName: l.TableName() + "_pk",
				Column:         "id",
			},
		},
		Indexes: l.SchemaIndexes(),
	}
}

func (l *LookupTable) SchemaColumns() []*schema.Column {
	cols := make([]*schema.Column, len(l.Fields))
	for i, field := range l.Fields {
		cols[i] = field.LookupSchemaColumn()
	}
	return cols
}

func (l *LookupTable) SchemaIndexes() []*schema.Index {
	indexes := []*schema.Index{
		&schema.Index{
			Name:      IndexName(l.TableName(), []string{l.Fields[1].SqlName}, true),
			TableName: l.TableName(),
			Columns:   []string{l.Fields[1].SqlName},
			Unique:    true,
			Force:     true,
		},
	}
	for _, field := range l.Fields {
		if field.SqlLookupExpr != "" {
			indexes = append(indexes, &schema.Index{
				Name:      IndexName(l.TableName(), []string{field.SqlName}, false),
				TableName: l.TableName(),
				Columns:   []string{field.SqlLookupExpr},
			})
		}
	}
	return indexes
}
<|endoftext|>


<commit_before>package query

import (
	pb "github.com/mediachain/concat/proto"
	"reflect"
	"testing"
)

var simpleq []string = []string{
	"SELECT * FROM foo.bar",
	"SELECT body FROM foo.bar",
	"SELECT id FROM foo.bar",
	"SELECT source FROM foo.bar",
	"SELECT timestamp FROM foo.bar",
	"SELECT (body, source) FROM foo.bar",
	"SELECT COUNT(*) FROM foo.bar",
	"SELECT namespace FROM *",
	"SELECT * FROM foo.bar.*",
	"SELECT * FROM foo.bar-baz-with-dashes",
	"SELECT * FROM foo.bar WHERE id = abc",
	"SELECT * FROM foo.bar WHERE publisher = abc",
	"SELECT * FROM foo.bar WHERE source = abc",
	"SELECT * FROM foo.bar WHERE timestamp < 1474000000",
	"SELECT * FROM foo.bar WHERE timestamp <= 1474000000",
	"SELECT * FROM foo.bar WHERE timestamp = 1474000000",
	"SELECT * FROM foo.bar WHERE timestamp >= 1474000000",
	"SELECT * FROM foo.bar WHERE timestamp > 1474000000",
	"SELECT * FROM foo.bar WHERE publisher = abc AND timestamp > 1474000000",
	"SELECT * FROM foo.bar WHERE publisher = abc OR timestamp > 1474000000",
	"SELECT * FROM foo.bar WHERE (publisher = abc AND timestamp > 1474000000) OR timestamp < 1474000000",
	"SELECT * FROM foo.bar WHERE publisher = abc LIMIT 10",
	"SELECT * FROM foo.bar LIMIT 10"}

func checkError(t *testing.T, where string, err error) {
	if err != nil {
		t.Logf("QUERY: %s", where)
		t.Error(err)
	}
}

func checkErrorNow(t *testing.T, where string, err error) {
	if err != nil {
		t.Logf("QUERY: %s", where)
		t.Log(err)
		t.FailNow()
	}
}

func checkResultLen(t *testing.T, where string, res []interface{}, xlen int) bool {
	if len(res) == xlen {
		return true
	}

	t.Logf("QUERY: %s", where)
	t.Errorf("Bad result length: expected %d elements, but got %v", xlen, res)
	return false
}

func checkContains(t *testing.T, where string, res []interface{}, val interface{}) {
	for _, v := range res {
		if reflect.DeepEqual(v, val) {
			return
		}
	}
	t.Logf("QUERY: %s", where)
	t.Errorf("%v is not in result set", val)
}

func TestQuerySyntax(t *testing.T) {
	for _, qs := range simpleq {
		_, err := ParseQuery(qs)
		checkError(t, qs, err)
	}
}

func TestQueryEval(t *testing.T) {
	a := &pb.Statement{
		Id:        "a",
		Publisher: "A",
		Namespace: "foo.a",
		Body:      &pb.Statement_Simple{&pb.SimpleStatement{Object: "QmAAA"}},
		Timestamp: 100}

	b := &pb.Statement{
		Id:        "b",
		Publisher: "B",
		Namespace: "foo.b",
		Body:      &pb.Statement_Simple{&pb.SimpleStatement{Object: "QmBBB"}},
		Timestamp: 200}

	c := &pb.Statement{
		Id:        "c",
		Publisher: "A",
		Namespace: "bar.c",
		Body:      &pb.Statement_Simple{&pb.SimpleStatement{Object: "QmCCC"}},
		Timestamp: 300}

	stmts := []*pb.Statement{a, b, c}

	// check namespace filter
	qs := "SELECT * FROM *"
	q, err := ParseQuery(qs)
	checkErrorNow(t, qs, err)

	res, err := EvalQuery(q, stmts)
	checkErrorNow(t, qs, err)

	if checkResultLen(t, qs, res, 3) {
		checkContains(t, qs, res, a)
		checkContains(t, qs, res, b)
		checkContains(t, qs, res, c)
	}

	qs = "SELECT * FROM foo.a"
	q, err = ParseQuery(qs)
	checkErrorNow(t, qs, err)

	res, err = EvalQuery(q, stmts)
	checkErrorNow(t, qs, err)

	if checkResultLen(t, qs, res, 1) {
		checkContains(t, qs, res, a)
	}

	qs = "SELECT * FROM foo.*"
	q, err = ParseQuery(qs)
	checkErrorNow(t, qs, err)

	res, err = EvalQuery(q, stmts)
	checkErrorNow(t, qs, err)

	if checkResultLen(t, qs, res, 2) {
		checkContains(t, qs, res, a)
		checkContains(t, qs, res, b)
	}

	// check all simple selectors
	qs = "SELECT body FROM foo.a"
	q, err = ParseQuery(qs)
	checkErrorNow(t, qs, err)

	res, err = EvalQuery(q, stmts)
	checkErrorNow(t, qs, err)

	if checkResultLen(t, qs, res, 1) {
		checkContains(t, qs, res, a.Body)
	}

	qs = "SELECT id FROM foo.a"
	q, err = ParseQuery(qs)
	checkErrorNow(t, qs, err)

	res, err = EvalQuery(q, stmts)
	checkErrorNow(t, qs, err)

	if checkResultLen(t, qs, res, 1) {
		checkContains(t, qs, res, "a")
	}

	qs = "SELECT publisher FROM foo.a"
	q, err = ParseQuery(qs)
	checkErrorNow(t, qs, err)

	res, err = EvalQuery(q, stmts)
	checkErrorNow(t, qs, err)

	if checkResultLen(t, qs, res, 1) {
		checkContains(t, qs, res, "A")
	}

	qs = "SELECT source FROM foo.a"
	q, err = ParseQuery(qs)
	checkErrorNow(t, qs, err)

	res, err = EvalQuery(q, stmts)
	checkErrorNow(t, qs, err)

	if checkResultLen(t, qs, res, 1) {
		checkContains(t, qs, res, "A")
	}

	qs = "SELECT namespace FROM foo.a"
	q, err = ParseQuery(qs)
	checkErrorNow(t, qs, err)

	res, err = EvalQuery(q, stmts)
	checkErrorNow(t, qs, err)

	if checkResultLen(t, qs, res, 1) {
		checkContains(t, qs, res, "foo.a")
	}

	qs = "SELECT timestamp FROM foo.a"
	q, err = ParseQuery(qs)
	checkErrorNow(t, qs, err)

	res, err = EvalQuery(q, stmts)
	checkErrorNow(t, qs, err)

	if checkResultLen(t, qs, res, 1) {
		checkContains(t, qs, res, int64(100))
	}

	// check compound selection
	qs = "SELECT (id, publisher) FROM foo.a"
	q, err = ParseQuery(qs)
	checkErrorNow(t, qs, err)

	res, err = EvalQuery(q, stmts)
	checkErrorNow(t, qs, err)

	if checkResultLen(t, qs, res, 1) {
		checkContains(t, qs, res, map[string]interface{}{"id": "a", "publisher": "A"})
	}

	// check the function
	qs = "SELECT COUNT(*) FROM foo.a"
	q, err = ParseQuery(qs)
	checkErrorNow(t, qs, err)

	res, err = EvalQuery(q, stmts)
	checkErrorNow(t, qs, err)

	if checkResultLen(t, qs, res, 1) {
		checkContains(t, qs, res, 1)
	}

	// check simple selection criteria
	qs = "SELECT * FROM * WHERE id = a"
	q, err = ParseQuery(qs)
	checkErrorNow(t, qs, err)

	res, err = EvalQuery(q, stmts)
	checkErrorNow(t, qs, err)

	if checkResultLen(t, qs, res, 1) {
		checkContains(t, qs, res, a)
	}

	qs = "SELECT * FROM * WHERE publisher = A"
	q, err = ParseQuery(qs)
	checkErrorNow(t, qs, err)

	res, err = EvalQuery(q, stmts)
	checkErrorNow(t, qs, err)

	if checkResultLen(t, qs, res, 2) {
		checkContains(t, qs, res, a)
		checkContains(t, qs, res, c)
	}

	qs = "SELECT * FROM * WHERE source = A"
	q, err = ParseQuery(qs)
	checkErrorNow(t, qs, err)

	res, err = EvalQuery(q, stmts)
	checkErrorNow(t, qs, err)

	if checkResultLen(t, qs, res, 2) {
		checkContains(t, qs, res, a)
		checkContains(t, qs, res, c)
	}

	qs = "SELECT * FROM * WHERE timestamp < 200"
	q, err = ParseQuery(qs)
	checkErrorNow(t, qs, err)

	res, err = EvalQuery(q, stmts)
	checkErrorNow(t, qs, err)

	if checkResultLen(t, qs, res, 1) {
		checkContains(t, qs, res, a)
	}

	qs = "SELECT * FROM * WHERE timestamp <= 200"
	q, err = ParseQuery(qs)
	checkErrorNow(t, qs, err)

	res, err = EvalQuery(q, stmts)
	checkErrorNow(t, qs, err)

	if checkResultLen(t, qs, res, 2) {
		checkContains(t, qs, res, a)
		checkContains(t, qs, res, b)
	}

	qs = "SELECT * FROM * WHERE timestamp = 200"
	q, err = ParseQuery(qs)
	checkErrorNow(t, qs, err)

	res, err = EvalQuery(q, stmts)
	checkErrorNow(t, qs, err)

	if checkResultLen(t, qs, res, 1) {
		checkContains(t, qs, res, b)
	}

	qs = "SELECT * FROM * WHERE timestamp >= 200"
	q, err = ParseQuery(qs)
	checkErrorNow(t, qs, err)

	res, err = EvalQuery(q, stmts)
	checkErrorNow(t, qs, err)

	if checkResultLen(t, qs, res, 2) {
		checkContains(t, qs, res, b)
		checkContains(t, qs, res, c)
	}

	qs = "SELECT * FROM * WHERE timestamp > 200"
	q, err = ParseQuery(qs)
	checkErrorNow(t, qs, err)

	res, err = EvalQuery(q, stmts)
	checkErrorNow(t, qs, err)

	if checkResultLen(t, qs, res, 1) {
		checkContains(t, qs, res, c)
	}

	// check compound selection criteria
	qs = "SELECT * FROM * WHERE timestamp >= 200 AND publisher = A"
	q, err = ParseQuery(qs)
	checkErrorNow(t, qs, err)

	res, err = EvalQuery(q, stmts)
	checkErrorNow(t, qs, err)

	if checkResultLen(t, qs, res, 1) {
		checkContains(t, qs, res, c)
	}

	qs = "SELECT * FROM * WHERE timestamp < 200 OR publisher = B"
	q, err = ParseQuery(qs)
	checkErrorNow(t, qs, err)

	res, err = EvalQuery(q, stmts)
	checkErrorNow(t, qs, err)

	if checkResultLen(t, qs, res, 2) {
		checkContains(t, qs, res, a)
		checkContains(t, qs, res, b)
	}

	qs = "SELECT * FROM * WHERE publisher = B OR (publisher = A AND timestamp < 200)"
	q, err = ParseQuery(qs)
	checkErrorNow(t, qs, err)

	res, err = EvalQuery(q, stmts)
	checkErrorNow(t, qs, err)

	if checkResultLen(t, qs, res, 2) {
		checkContains(t, qs, res, a)
		checkContains(t, qs, res, b)
	}

}
<commit_msg>query: test LIMIT<commit_after>package query

import (
	pb "github.com/mediachain/concat/proto"
	"reflect"
	"testing"
)

var simpleq []string = []string{
	"SELECT * FROM foo.bar",
	"SELECT body FROM foo.bar",
	"SELECT id FROM foo.bar",
	"SELECT source FROM foo.bar",
	"SELECT timestamp FROM foo.bar",
	"SELECT (body, source) FROM foo.bar",
	"SELECT COUNT(*) FROM foo.bar",
	"SELECT namespace FROM *",
	"SELECT * FROM foo.bar.*",
	"SELECT * FROM foo.bar-baz-with-dashes",
	"SELECT * FROM foo.bar WHERE id = abc",
	"SELECT * FROM foo.bar WHERE publisher = abc",
	"SELECT * FROM foo.bar WHERE source = abc",
	"SELECT * FROM foo.bar WHERE timestamp < 1474000000",
	"SELECT * FROM foo.bar WHERE timestamp <= 1474000000",
	"SELECT * FROM foo.bar WHERE timestamp = 1474000000",
	"SELECT * FROM foo.bar WHERE timestamp >= 1474000000",
	"SELECT * FROM foo.bar WHERE timestamp > 1474000000",
	"SELECT * FROM foo.bar WHERE publisher = abc AND timestamp > 1474000000",
	"SELECT * FROM foo.bar WHERE publisher = abc OR timestamp > 1474000000",
	"SELECT * FROM foo.bar WHERE (publisher = abc AND timestamp > 1474000000) OR timestamp < 1474000000",
	"SELECT * FROM foo.bar WHERE publisher = abc LIMIT 10",
	"SELECT * FROM foo.bar LIMIT 10"}

func checkError(t *testing.T, where string, err error) {
	if err != nil {
		t.Logf("QUERY: %s", where)
		t.Error(err)
	}
}

func checkErrorNow(t *testing.T, where string, err error) {
	if err != nil {
		t.Logf("QUERY: %s", where)
		t.Log(err)
		t.FailNow()
	}
}

func checkResultLen(t *testing.T, where string, res []interface{}, xlen int) bool {
	if len(res) == xlen {
		return true
	}

	t.Logf("QUERY: %s", where)
	t.Errorf("Bad result length: expected %d elements, but got %v", xlen, res)
	return false
}

func checkContains(t *testing.T, where string, res []interface{}, val interface{}) {
	for _, v := range res {
		if reflect.DeepEqual(v, val) {
			return
		}
	}
	t.Logf("QUERY: %s", where)
	t.Errorf("%v is not in result set", val)
}

func TestQuerySyntax(t *testing.T) {
	for _, qs := range simpleq {
		_, err := ParseQuery(qs)
		checkError(t, qs, err)
	}
}

func TestQueryEval(t *testing.T) {
	a := &pb.Statement{
		Id:        "a",
		Publisher: "A",
		Namespace: "foo.a",
		Body:      &pb.Statement_Simple{&pb.SimpleStatement{Object: "QmAAA"}},
		Timestamp: 100}

	b := &pb.Statement{
		Id:        "b",
		Publisher: "B",
		Namespace: "foo.b",
		Body:      &pb.Statement_Simple{&pb.SimpleStatement{Object: "QmBBB"}},
		Timestamp: 200}

	c := &pb.Statement{
		Id:        "c",
		Publisher: "A",
		Namespace: "bar.c",
		Body:      &pb.Statement_Simple{&pb.SimpleStatement{Object: "QmCCC"}},
		Timestamp: 300}

	stmts := []*pb.Statement{a, b, c}

	// check namespace filter
	qs := "SELECT * FROM *"
	q, err := ParseQuery(qs)
	checkErrorNow(t, qs, err)

	res, err := EvalQuery(q, stmts)
	checkErrorNow(t, qs, err)

	if checkResultLen(t, qs, res, 3) {
		checkContains(t, qs, res, a)
		checkContains(t, qs, res, b)
		checkContains(t, qs, res, c)
	}

	qs = "SELECT * FROM foo.a"
	q, err = ParseQuery(qs)
	checkErrorNow(t, qs, err)

	res, err = EvalQuery(q, stmts)
	checkErrorNow(t, qs, err)

	if checkResultLen(t, qs, res, 1) {
		checkContains(t, qs, res, a)
	}

	qs = "SELECT * FROM foo.*"
	q, err = ParseQuery(qs)
	checkErrorNow(t, qs, err)

	res, err = EvalQuery(q, stmts)
	checkErrorNow(t, qs, err)

	if checkResultLen(t, qs, res, 2) {
		checkContains(t, qs, res, a)
		checkContains(t, qs, res, b)
	}

	// check all simple selectors
	qs = "SELECT body FROM foo.a"
	q, err = ParseQuery(qs)
	checkErrorNow(t, qs, err)

	res, err = EvalQuery(q, stmts)
	checkErrorNow(t, qs, err)

	if checkResultLen(t, qs, res, 1) {
		checkContains(t, qs, res, a.Body)
	}

	qs = "SELECT id FROM foo.a"
	q, err = ParseQuery(qs)
	checkErrorNow(t, qs, err)

	res, err = EvalQuery(q, stmts)
	checkErrorNow(t, qs, err)

	if checkResultLen(t, qs, res, 1) {
		checkContains(t, qs, res, "a")
	}

	qs = "SELECT publisher FROM foo.a"
	q, err = ParseQuery(qs)
	checkErrorNow(t, qs, err)

	res, err = EvalQuery(q, stmts)
	checkErrorNow(t, qs, err)

	if checkResultLen(t, qs, res, 1) {
		checkContains(t, qs, res, "A")
	}

	qs = "SELECT source FROM foo.a"
	q, err = ParseQuery(qs)
	checkErrorNow(t, qs, err)

	res, err = EvalQuery(q, stmts)
	checkErrorNow(t, qs, err)

	if checkResultLen(t, qs, res, 1) {
		checkContains(t, qs, res, "A")
	}

	qs = "SELECT namespace FROM foo.a"
	q, err = ParseQuery(qs)
	checkErrorNow(t, qs, err)

	res, err = EvalQuery(q, stmts)
	checkErrorNow(t, qs, err)

	if checkResultLen(t, qs, res, 1) {
		checkContains(t, qs, res, "foo.a")
	}

	qs = "SELECT timestamp FROM foo.a"
	q, err = ParseQuery(qs)
	checkErrorNow(t, qs, err)

	res, err = EvalQuery(q, stmts)
	checkErrorNow(t, qs, err)

	if checkResultLen(t, qs, res, 1) {
		checkContains(t, qs, res, int64(100))
	}

	// check compound selection
	qs = "SELECT (id, publisher) FROM foo.a"
	q, err = ParseQuery(qs)
	checkErrorNow(t, qs, err)

	res, err = EvalQuery(q, stmts)
	checkErrorNow(t, qs, err)

	if checkResultLen(t, qs, res, 1) {
		checkContains(t, qs, res, map[string]interface{}{"id": "a", "publisher": "A"})
	}

	// check the function
	qs = "SELECT COUNT(*) FROM foo.a"
	q, err = ParseQuery(qs)
	checkErrorNow(t, qs, err)

	res, err = EvalQuery(q, stmts)
	checkErrorNow(t, qs, err)

	if checkResultLen(t, qs, res, 1) {
		checkContains(t, qs, res, 1)
	}

	// check the limits -- order is unpredictable, so just check the count
	qs = "SELECT COUNT(*) FROM * LIMIT 1"
	q, err = ParseQuery(qs)
	checkErrorNow(t, qs, err)

	res, err = EvalQuery(q, stmts)
	checkErrorNow(t, qs, err)

	if checkResultLen(t, qs, res, 1) {
		checkContains(t, qs, res, 1)
	}

	// check simple selection criteria
	qs = "SELECT * FROM * WHERE id = a"
	q, err = ParseQuery(qs)
	checkErrorNow(t, qs, err)

	res, err = EvalQuery(q, stmts)
	checkErrorNow(t, qs, err)

	if checkResultLen(t, qs, res, 1) {
		checkContains(t, qs, res, a)
	}

	qs = "SELECT * FROM * WHERE publisher = A"
	q, err = ParseQuery(qs)
	checkErrorNow(t, qs, err)

	res, err = EvalQuery(q, stmts)
	checkErrorNow(t, qs, err)

	if checkResultLen(t, qs, res, 2) {
		checkContains(t, qs, res, a)
		checkContains(t, qs, res, c)
	}

	qs = "SELECT * FROM * WHERE source = A"
	q, err = ParseQuery(qs)
	checkErrorNow(t, qs, err)

	res, err = EvalQuery(q, stmts)
	checkErrorNow(t, qs, err)

	if checkResultLen(t, qs, res, 2) {
		checkContains(t, qs, res, a)
		checkContains(t, qs, res, c)
	}

	qs = "SELECT * FROM * WHERE timestamp < 200"
	q, err = ParseQuery(qs)
	checkErrorNow(t, qs, err)

	res, err = EvalQuery(q, stmts)
	checkErrorNow(t, qs, err)

	if checkResultLen(t, qs, res, 1) {
		checkContains(t, qs, res, a)
	}

	qs = "SELECT * FROM * WHERE timestamp <= 200"
	q, err = ParseQuery(qs)
	checkErrorNow(t, qs, err)

	res, err = EvalQuery(q, stmts)
	checkErrorNow(t, qs, err)

	if checkResultLen(t, qs, res, 2) {
		checkContains(t, qs, res, a)
		checkContains(t, qs, res, b)
	}

	qs = "SELECT * FROM * WHERE timestamp = 200"
	q, err = ParseQuery(qs)
	checkErrorNow(t, qs, err)

	res, err = EvalQuery(q, stmts)
	checkErrorNow(t, qs, err)

	if checkResultLen(t, qs, res, 1) {
		checkContains(t, qs, res, b)
	}

	qs = "SELECT * FROM * WHERE timestamp >= 200"
	q, err = ParseQuery(qs)
	checkErrorNow(t, qs, err)

	res, err = EvalQuery(q, stmts)
	checkErrorNow(t, qs, err)

	if checkResultLen(t, qs, res, 2) {
		checkContains(t, qs, res, b)
		checkContains(t, qs, res, c)
	}

	qs = "SELECT * FROM * WHERE timestamp > 200"
	q, err = ParseQuery(qs)
	checkErrorNow(t, qs, err)

	res, err = EvalQuery(q, stmts)
	checkErrorNow(t, qs, err)

	if checkResultLen(t, qs, res, 1) {
		checkContains(t, qs, res, c)
	}

	// check compound selection criteria
	qs = "SELECT * FROM * WHERE timestamp >= 200 AND publisher = A"
	q, err = ParseQuery(qs)
	checkErrorNow(t, qs, err)

	res, err = EvalQuery(q, stmts)
	checkErrorNow(t, qs, err)

	if checkResultLen(t, qs, res, 1) {
		checkContains(t, qs, res, c)
	}

	qs = "SELECT * FROM * WHERE timestamp < 200 OR publisher = B"
	q, err = ParseQuery(qs)
	checkErrorNow(t, qs, err)

	res, err = EvalQuery(q, stmts)
	checkErrorNow(t, qs, err)

	if checkResultLen(t, qs, res, 2) {
		checkContains(t, qs, res, a)
		checkContains(t, qs, res, b)
	}

	qs = "SELECT * FROM * WHERE publisher = B OR (publisher = A AND timestamp < 200)"
	q, err = ParseQuery(qs)
	checkErrorNow(t, qs, err)

	res, err = EvalQuery(q, stmts)
	checkErrorNow(t, qs, err)

	if checkResultLen(t, qs, res, 2) {
		checkContains(t, qs, res, a)
		checkContains(t, qs, res, b)
	}

}
<|endoftext|>


<commit_before>/*
Copyright [2017] [Elre Oldewage]

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing,
software distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and limitations under the License.
*/

//Package main registers an HTTP handler at "/" that handles incoming requests for stubborn-manatee.
package main

import (
	"flag"
	"fmt"
	_ "fmt"
	"log"
	"net/http"

	"github.com/gorilla/mux"
	ctrls "github.com/stubborn-manatee/server/controllers"
)

//Every handler should be wrapped by this error handler
func errorHandler(f func(http.ResponseWriter, *http.Request) error) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		err := f(w, r)
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			log.Printf("Handling %q: %v", r.RequestURI, err)
		}
	}
}

func redirectToHttps(w http.ResponseWriter, r *http.Request) {
	http.Redirect(w, r, "https://"+r.URL.Host+r.RequestURI, http.StatusMovedPermanently)
}

func main() {
	contentPath := flag.String("c", "../../content/", "Path to the content directory")
	serverPort := flag.Int("p", 8080, "Port on which server will be hosted")
	flag.Parse()
	r := mux.NewRouter()
	postController := ctrls.NewPostController()
	r.HandleFunc("/post", postController.GetPost).Methods("GET")
	r.HandleFunc("/post", postController.CreatePost).Methods("POST")
	r.PathPrefix("/").Handler(http.FileServer(http.Dir(*contentPath)))

	http.HandleFunc( "/", redirectToHttps)

	go http.ListenAndServe(":80", nil)
	fmt.Println("Hosting stubborn-manatee on port", fmt.Sprintf("%v", *serverPort))
	err := http.ListenAndServeTLS(fmt.Sprintf(":%v", *serverPort), "/etc/letsencrypt/live/stubborn-manatee.co.za/cert.pem", "/etc/letsencrypt/live/stubborn-manatee.co.za/privkey.pem", r)
	if err != nil {
		panic(err);
	}
}
<commit_msg>Hopefully fixed the http to https redirect<commit_after>/*
Copyright [2017] [Elre Oldewage]

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing,
software distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and limitations under the License.
*/

//Package main registers an HTTP handler at "/" that handles incoming requests for stubborn-manatee.
package main

import (
	"flag"
	"fmt"
	_ "fmt"
	"log"
	"net/http"
	"net/url"

	"github.com/gorilla/mux"
	ctrls "github.com/stubborn-manatee/server/controllers"
)

//Every handler should be wrapped by this error handler
func errorHandler(f func(http.ResponseWriter, *http.Request) error) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		err := f(w, r)
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			log.Printf("Handling %q: %v", r.RequestURI, err)
		}
	}
}

func redirectToHttps(w http.ResponseWriter, r *http.Request) {
	targetUrl := url.URL{Scheme: "https", Host: r.Host, Path: r.URL.Path, RawQuery: r.URL.RawQuery}
	http.Redirect(w, r, targetUrl.String(), http.StatusMovedPermanently)
}

func main() {
	contentPath := flag.String("c", "../../content/", "Path to the content directory")
	serverPort := flag.Int("p", 8080, "Port on which server will be hosted")
	flag.Parse()
	r := mux.NewRouter()
	postController := ctrls.NewPostController()
	r.HandleFunc("/post", postController.GetPost).Methods("GET")
	r.HandleFunc("/post", postController.CreatePost).Methods("POST")
	r.PathPrefix("/").Handler(http.FileServer(http.Dir(*contentPath)))

	http.HandleFunc("/", redirectToHttps)

	go http.ListenAndServe(":80", nil)
	fmt.Println("Hosting stubborn-manatee on port", fmt.Sprintf("%v", *serverPort))
	err := http.ListenAndServeTLS(fmt.Sprintf(":%v", *serverPort), "/etc/letsencrypt/live/stubborn-manatee.co.za/cert.pem", "/etc/letsencrypt/live/stubborn-manatee.co.za/privkey.pem", r)
	if err != nil {
		panic(err)
	}
}
<|endoftext|>


<commit_before>// Package agent runs on servers with computing resources, and executes
// tasks sent by driver.
package agent

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"net/url"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"sync"

	"github.com/chrislusf/glow/driver/cmd"
	"github.com/chrislusf/glow/resource"
	"github.com/chrislusf/glow/resource/service_discovery/client"
	"github.com/chrislusf/glow/util"
	"github.com/golang/protobuf/proto"
)

type AgentServerOption struct {
	Master       *string
	Port         *int
	Dir          *string
	DataCenter   *string
	Rack         *string
	MaxExecutor  *int
	MemoryMB     *int64
	CPULevel     *int
	CleanRestart *bool
}

type AgentServer struct {
	Option                *AgentServerOption
	Master                string
	Port                  int
	wg                    sync.WaitGroup
	listener              net.Listener
	computeResource       *resource.ComputeResource
	allocatedResource     *resource.ComputeResource
	allocatedResourceLock sync.Mutex
	storageBackend        *ManagedDatasetShards
}

func NewAgentServer(option *AgentServerOption) *AgentServer {
	absoluteDir, err := filepath.Abs(util.CleanPath(*option.Dir))
	if err != nil {
		panic(err)
	}
	println("starting in", absoluteDir)
	option.Dir = &absoluteDir

	as := &AgentServer{
		Option:         option,
		Master:         *option.Master,
		Port:           *option.Port,
		storageBackend: NewManagedDatasetShards(*option.Dir, *option.Port),
		computeResource: &resource.ComputeResource{
			CPUCount: *option.MaxExecutor,
			CPULevel: *option.CPULevel,
			MemoryMB: *option.MemoryMB,
		},
		allocatedResource: &resource.ComputeResource{},
	}

	err = as.Init()
	if err != nil {
		panic(err)
	}

	return as
}

// Start starts to listen on a port, returning the listening port
// r.Port can be pre-set or leave it as zero
// The actual port set to r.Port
func (r *AgentServer) Init() (err error) {
	r.listener, err = net.Listen("tcp", ":"+strconv.Itoa(r.Port))
	if err != nil {
		log.Fatal(err)
	}

	r.Port = r.listener.Addr().(*net.TCPAddr).Port
	fmt.Println("AgentServer starts on:", r.Port)

	if *r.Option.CleanRestart {
		if fileInfos, err := ioutil.ReadDir(r.storageBackend.dir); err == nil {
			for _, fi := range fileInfos {
				name := fi.Name()
				if !fi.IsDir() && strings.HasSuffix(name, ".dat") {
					// println("removing old dat file:", name)
					os.Remove(filepath.Join(r.storageBackend.dir, name))
				}
			}
		}
	}

	return
}

func (as *AgentServer) Run() {
	//register agent
	killHeartBeaterChan := make(chan bool, 1)
	go client.NewHeartBeater(as.Port, as.Master).StartAgentHeartBeat(killHeartBeaterChan, func(values url.Values) {
		resource.AddToValues(values, as.computeResource, as.allocatedResource)
		values.Add("dataCenter", *as.Option.DataCenter)
		values.Add("rack", *as.Option.Rack)
	})

	for {
		// Listen for an incoming connection.
		conn, err := as.listener.Accept()
		if err != nil {
			fmt.Println("Error accepting: ", err.Error())
			os.Exit(1)
		}
		// Handle connections in a new goroutine.
		as.wg.Add(1)
		go func() {
			defer as.wg.Done()
			defer conn.Close()
			as.handleRequest(conn)
		}()
	}
}

func (r *AgentServer) Stop() {
	r.listener.Close()
	r.wg.Wait()
}

// Handles incoming requests.
func (r *AgentServer) handleRequest(conn net.Conn) {

	buf := make([]byte, 4)

	f, message, err := util.ReadBytes(conn, buf)
	if f != util.Data {
		//strange if this happens
		println("read", len(message.Bytes()), "request flag:", f, "data", string(message.Data()))
		return
	}

	if err != nil {
		log.Printf("Failed to read command %s:%v", string(message.Data()), err)
	}
	if bytes.HasPrefix(message.Data(), []byte("PUT ")) {
		name := string(message.Data()[4:])
		r.handleWriteConnection(conn, name)
	} else if bytes.HasPrefix(message.Data(), []byte("GET ")) {
		name := string(message.Data()[4:])
		offset := util.ReadUint64(conn)
		r.handleLocalReadConnection(conn, name, int64(offset))
	} else if bytes.HasPrefix(message.Data(), []byte("CMD ")) {
		newCmd := &cmd.ControlMessage{}
		err := proto.Unmarshal(message.Data()[4:], newCmd)
		if err != nil {
			log.Fatal("unmarshaling error: ", err)
		}
		reply := r.handleCommandConnection(conn, newCmd)
		if reply != nil {
			data, err := proto.Marshal(reply)
			if err != nil {
				log.Fatal("marshaling error: ", err)
			}
			conn.Write(data)
		}
	}

}
<commit_msg>accurate data file suffix name<commit_after>// Package agent runs on servers with computing resources, and executes
// tasks sent by driver.
package agent

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"net/url"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"sync"

	"github.com/chrislusf/glow/driver/cmd"
	"github.com/chrislusf/glow/resource"
	"github.com/chrislusf/glow/resource/service_discovery/client"
	"github.com/chrislusf/glow/util"
	"github.com/golang/protobuf/proto"
)

type AgentServerOption struct {
	Master       *string
	Port         *int
	Dir          *string
	DataCenter   *string
	Rack         *string
	MaxExecutor  *int
	MemoryMB     *int64
	CPULevel     *int
	CleanRestart *bool
}

type AgentServer struct {
	Option                *AgentServerOption
	Master                string
	Port                  int
	wg                    sync.WaitGroup
	listener              net.Listener
	computeResource       *resource.ComputeResource
	allocatedResource     *resource.ComputeResource
	allocatedResourceLock sync.Mutex
	storageBackend        *ManagedDatasetShards
}

func NewAgentServer(option *AgentServerOption) *AgentServer {
	absoluteDir, err := filepath.Abs(util.CleanPath(*option.Dir))
	if err != nil {
		panic(err)
	}
	println("starting in", absoluteDir)
	option.Dir = &absoluteDir

	as := &AgentServer{
		Option:         option,
		Master:         *option.Master,
		Port:           *option.Port,
		storageBackend: NewManagedDatasetShards(*option.Dir, *option.Port),
		computeResource: &resource.ComputeResource{
			CPUCount: *option.MaxExecutor,
			CPULevel: *option.CPULevel,
			MemoryMB: *option.MemoryMB,
		},
		allocatedResource: &resource.ComputeResource{},
	}

	err = as.Init()
	if err != nil {
		panic(err)
	}

	return as
}

// Start starts to listen on a port, returning the listening port
// r.Port can be pre-set or leave it as zero
// The actual port set to r.Port
func (r *AgentServer) Init() (err error) {
	r.listener, err = net.Listen("tcp", ":"+strconv.Itoa(r.Port))
	if err != nil {
		log.Fatal(err)
	}

	r.Port = r.listener.Addr().(*net.TCPAddr).Port
	fmt.Println("AgentServer starts on:", r.Port)

	if *r.Option.CleanRestart {
		if fileInfos, err := ioutil.ReadDir(r.storageBackend.dir); err == nil {
			suffix := fmt.Sprintf("-%d.dat", r.Port)
			for _, fi := range fileInfos {
				name := fi.Name()
				if !fi.IsDir() && strings.HasSuffix(name, suffix) {
					// println("removing old dat file:", name)
					os.Remove(filepath.Join(r.storageBackend.dir, name))
				}
			}
		}
	}

	return
}

func (as *AgentServer) Run() {
	//register agent
	killHeartBeaterChan := make(chan bool, 1)
	go client.NewHeartBeater(as.Port, as.Master).StartAgentHeartBeat(killHeartBeaterChan, func(values url.Values) {
		resource.AddToValues(values, as.computeResource, as.allocatedResource)
		values.Add("dataCenter", *as.Option.DataCenter)
		values.Add("rack", *as.Option.Rack)
	})

	for {
		// Listen for an incoming connection.
		conn, err := as.listener.Accept()
		if err != nil {
			fmt.Println("Error accepting: ", err.Error())
			os.Exit(1)
		}
		// Handle connections in a new goroutine.
		as.wg.Add(1)
		go func() {
			defer as.wg.Done()
			defer conn.Close()
			as.handleRequest(conn)
		}()
	}
}

func (r *AgentServer) Stop() {
	r.listener.Close()
	r.wg.Wait()
}

// Handles incoming requests.
func (r *AgentServer) handleRequest(conn net.Conn) {

	buf := make([]byte, 4)

	f, message, err := util.ReadBytes(conn, buf)
	if f != util.Data {
		//strange if this happens
		println("read", len(message.Bytes()), "request flag:", f, "data", string(message.Data()))
		return
	}

	if err != nil {
		log.Printf("Failed to read command %s:%v", string(message.Data()), err)
	}
	if bytes.HasPrefix(message.Data(), []byte("PUT ")) {
		name := string(message.Data()[4:])
		r.handleWriteConnection(conn, name)
	} else if bytes.HasPrefix(message.Data(), []byte("GET ")) {
		name := string(message.Data()[4:])
		offset := util.ReadUint64(conn)
		r.handleLocalReadConnection(conn, name, int64(offset))
	} else if bytes.HasPrefix(message.Data(), []byte("CMD ")) {
		newCmd := &cmd.ControlMessage{}
		err := proto.Unmarshal(message.Data()[4:], newCmd)
		if err != nil {
			log.Fatal("unmarshaling error: ", err)
		}
		reply := r.handleCommandConnection(conn, newCmd)
		if reply != nil {
			data, err := proto.Marshal(reply)
			if err != nil {
				log.Fatal("marshaling error: ", err)
			}
			conn.Write(data)
		}
	}

}
<|endoftext|>


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

package agent

import (
	"crypto/tls"
	"crypto/x509"
	"encoding/json"
	"encoding/pem"
	"github.com/gocd-contrib/gocd-golang-agent/protocol"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"
	"runtime"
)

func ReadGoServerCACert() error {
	_, err := os.Stat(config.GoServerCAFile)
	if err == nil {
		return nil
	}

	LogInfo("fetching Go server[%v] CA certificate", config.ServerHostAndPort)
	conn, err := tls.Dial("tcp", config.ServerHostAndPort, &tls.Config{
		InsecureSkipVerify: true,
	})
	if err != nil {
		logger.Error.Printf("failed to connect: " + err.Error())
		return err
	}
	defer conn.Close()
	state := conn.ConnectionState()
	certOut, err := os.Create(config.GoServerCAFile)
	if err != nil {
		logger.Error.Printf("failed to open %v for writing: %s", config.GoServerCAFile, err)
		return err
	}
	defer certOut.Close()
	for i:=0; i < len(state.PeerCertificates); i++ {
		pem.Encode(certOut, &pem.Block{Type: "CERTIFICATE", Bytes: state.PeerCertificates[i].Raw})
	}
	return nil
}

func GoServerRootCAs() (*x509.CertPool, error) {
	caCert, err := ioutil.ReadFile(config.GoServerCAFile)
	if err != nil {
		return nil, err
	}
	roots := x509.NewCertPool()
	ok := roots.AppendCertsFromPEM([]byte(caCert))
	if !ok {
		return nil, Err("failed to parse root certificate")
	}
	return roots, nil
}

func GoServerTlsConfig(withClientCert bool) (*tls.Config, error) {
	certs := make([]tls.Certificate, 0)
	if withClientCert {
		cert, err := tls.LoadX509KeyPair(config.AgentCertFile, config.AgentPrivateKeyFile)
		if err != nil {
			return nil, err
		}
		certs = append(certs, cert)
	}
	roots, err := GoServerRootCAs()
	if err != nil {
		return nil, err
	}
	serverName, err := extractServerDN(config.GoServerCAFile)
	if err != nil {
		return nil, err
	}
	return &tls.Config{
		Certificates: certs,
		RootCAs:      roots,
		ServerName:   serverName,
	}, nil
}

func GoServerRemoteClient(withClientCert bool) (*http.Client, error) {
	config, err := GoServerTlsConfig(withClientCert)
	if err != nil {
		return nil, err
	}
	tr := &http.Transport{
		TLSClientConfig: config,
	}
	return &http.Client{Transport: tr}, nil
}

func Register() error {
	if err := ReadGoServerCACert(); err != nil {
		return err
	}
	if err := readAgentKeyAndCerts(registerData()); err != nil {
		return err
	}
	return nil
}

func CleanRegistration() error {
	files := []string{config.GoServerCAFile,
		config.AgentPrivateKeyFile,
		config.AgentCertFile}
	for _, f := range files {
		_, err := os.Stat(f)
		if err == nil {
			err := os.Remove(f)
			if err != nil {
				return err
			}
		}
	}
	return nil
}

func registerData() map[string]string {
	return map[string]string{
		"hostname":                      config.Hostname,
		"uuid":                          AgentId,
		"location":                      config.WorkingDir,
		"operatingSystem":               runtime.GOOS,
		"usablespace":                   UsableSpaceString(),
		"agentAutoRegisterKey":          config.AgentAutoRegisterKey,
		"agentAutoRegisterResources":    config.AgentAutoRegisterResources,
		"agentAutoRegisterEnvironments": config.AgentAutoRegisterEnvironments,
		"agentAutoRegisterHostname":     config.Hostname,
		"elasticAgentId":                config.AgentAutoRegisterElasticAgentId,
		"elasticPluginId":               config.AgentAutoRegisterElasticPluginId,
	}
}

func readAgentKeyAndCerts(params map[string]string) error {
	_, agentPrivateKeyFileErr := os.Stat(config.AgentPrivateKeyFile)
	_, agentCertFileErr := os.Stat(config.AgentCertFile)
	if agentPrivateKeyFileErr == nil && agentCertFileErr == nil {
		return nil
	}

	form := url.Values{}
	for k, v := range params {
		form.Add(k, v)
	}

	client, err := GoServerRemoteClient(false)
	if err != nil {
		return err
	}

	url, err := config.RegistrationURL()
	LogInfo("fetching agent key and certificates from: %v", url)
	if err != nil {
		return err
	}
	resp, err := client.PostForm(url.String(), form)
	if err != nil {
		return err
	}

	defer resp.Body.Close()
	var registration protocol.Registration

	dec := json.NewDecoder(resp.Body)

	if err := dec.Decode(&registration); err != nil {
		return err
	}
	if registration.AgentCertificate == "" {
		return Err("Register failed, probably need approve agent registration on Server side")
	}

	ioutil.WriteFile(config.AgentPrivateKeyFile, []byte(registration.AgentPrivateKey), 0600)
	ioutil.WriteFile(config.AgentCertFile, []byte(registration.AgentCertificate), 0600)
	return nil
}

func extractServerDN(certFileName string) (string, error) {
	pemBlock, err := ioutil.ReadFile(certFileName)
	if err != nil {
		return "", err
	}

	der, _ := pem.Decode(pemBlock)
	cert, err := x509.ParseCertificate(der.Bytes)
	if err != nil {
		return "", err
	}
	return cert.Subject.CommonName, nil
}
<commit_msg>add supportsBuildCommandProtocol to registration<commit_after>/*
 * Copyright 2016 ThoughtWorks, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *  http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package agent

import (
	"crypto/tls"
	"crypto/x509"
	"encoding/json"
	"encoding/pem"
	"github.com/gocd-contrib/gocd-golang-agent/protocol"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"
	"runtime"
)

func ReadGoServerCACert() error {
	_, err := os.Stat(config.GoServerCAFile)
	if err == nil {
		return nil
	}

	LogInfo("fetching Go server[%v] CA certificate", config.ServerHostAndPort)
	conn, err := tls.Dial("tcp", config.ServerHostAndPort, &tls.Config{
		InsecureSkipVerify: true,
	})
	if err != nil {
		logger.Error.Printf("failed to connect: " + err.Error())
		return err
	}
	defer conn.Close()
	state := conn.ConnectionState()
	certOut, err := os.Create(config.GoServerCAFile)
	if err != nil {
		logger.Error.Printf("failed to open %v for writing: %s", config.GoServerCAFile, err)
		return err
	}
	defer certOut.Close()
	for i:=0; i < len(state.PeerCertificates); i++ {
		pem.Encode(certOut, &pem.Block{Type: "CERTIFICATE", Bytes: state.PeerCertificates[i].Raw})
	}
	return nil
}

func GoServerRootCAs() (*x509.CertPool, error) {
	caCert, err := ioutil.ReadFile(config.GoServerCAFile)
	if err != nil {
		return nil, err
	}
	roots := x509.NewCertPool()
	ok := roots.AppendCertsFromPEM([]byte(caCert))
	if !ok {
		return nil, Err("failed to parse root certificate")
	}
	return roots, nil
}

func GoServerTlsConfig(withClientCert bool) (*tls.Config, error) {
	certs := make([]tls.Certificate, 0)
	if withClientCert {
		cert, err := tls.LoadX509KeyPair(config.AgentCertFile, config.AgentPrivateKeyFile)
		if err != nil {
			return nil, err
		}
		certs = append(certs, cert)
	}
	roots, err := GoServerRootCAs()
	if err != nil {
		return nil, err
	}
	serverName, err := extractServerDN(config.GoServerCAFile)
	if err != nil {
		return nil, err
	}
	return &tls.Config{
		Certificates: certs,
		RootCAs:      roots,
		ServerName:   serverName,
	}, nil
}

func GoServerRemoteClient(withClientCert bool) (*http.Client, error) {
	config, err := GoServerTlsConfig(withClientCert)
	if err != nil {
		return nil, err
	}
	tr := &http.Transport{
		TLSClientConfig: config,
	}
	return &http.Client{Transport: tr}, nil
}

func Register() error {
	if err := ReadGoServerCACert(); err != nil {
		return err
	}
	if err := readAgentKeyAndCerts(registerData()); err != nil {
		return err
	}
	return nil
}

func CleanRegistration() error {
	files := []string{config.GoServerCAFile,
		config.AgentPrivateKeyFile,
		config.AgentCertFile}
	for _, f := range files {
		_, err := os.Stat(f)
		if err == nil {
			err := os.Remove(f)
			if err != nil {
				return err
			}
		}
	}
	return nil
}

func registerData() map[string]string {
	return map[string]string{
		"hostname":                      config.Hostname,
		"uuid":                          AgentId,
		"location":                      config.WorkingDir,
		"operatingSystem":               runtime.GOOS,
		"usablespace":                   UsableSpaceString(),
		"agentAutoRegisterKey":          config.AgentAutoRegisterKey,
		"agentAutoRegisterResources":    config.AgentAutoRegisterResources,
		"agentAutoRegisterEnvironments": config.AgentAutoRegisterEnvironments,
		"agentAutoRegisterHostname":     config.Hostname,
		"elasticAgentId":                config.AgentAutoRegisterElasticAgentId,
		"elasticPluginId":               config.AgentAutoRegisterElasticPluginId,
		"supportsBuildCommandProtocol":  "true",
	}
}

func readAgentKeyAndCerts(params map[string]string) error {
	_, agentPrivateKeyFileErr := os.Stat(config.AgentPrivateKeyFile)
	_, agentCertFileErr := os.Stat(config.AgentCertFile)
	if agentPrivateKeyFileErr == nil && agentCertFileErr == nil {
		return nil
	}

	form := url.Values{}
	for k, v := range params {
		form.Add(k, v)
	}

	client, err := GoServerRemoteClient(false)
	if err != nil {
		return err
	}

	url, err := config.RegistrationURL()
	LogInfo("fetching agent key and certificates from: %v", url)
	if err != nil {
		return err
	}
	resp, err := client.PostForm(url.String(), form)
	if err != nil {
		return err
	}

	defer resp.Body.Close()
	var registration protocol.Registration

	dec := json.NewDecoder(resp.Body)

	if err := dec.Decode(&registration); err != nil {
		return err
	}
	if registration.AgentCertificate == "" {
		return Err("Register failed, probably need approve agent registration on Server side")
	}

	ioutil.WriteFile(config.AgentPrivateKeyFile, []byte(registration.AgentPrivateKey), 0600)
	ioutil.WriteFile(config.AgentCertFile, []byte(registration.AgentCertificate), 0600)
	return nil
}

func extractServerDN(certFileName string) (string, error) {
	pemBlock, err := ioutil.ReadFile(certFileName)
	if err != nil {
		return "", err
	}

	der, _ := pem.Decode(pemBlock)
	cert, err := x509.ParseCertificate(der.Bytes)
	if err != nil {
		return "", err
	}
	return cert.Subject.CommonName, nil
}
<|endoftext|>


<commit_before>// Copyright 2015 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package systemd

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"strings"

	"github.com/coreos/go-systemd/unit"
	"github.com/juju/errors"

	"github.com/juju/juju/service/common"
)

var limitMap = map[string]string{
	"as":         "LimitAS",
	"core":       "LimitCORE",
	"cpu":        "LimitCPU",
	"data":       "LimitDATA",
	"fsize":      "LimitFSIZE",
	"memlock":    "LimitMEMLOCK",
	"msgqueue":   "LimitMSGQUEUE",
	"nice":       "LimitNICE",
	"nofile":     "LimitNOFILE",
	"nproc":      "LimitNPROC",
	"rss":        "LimitRSS",
	"rtprio":     "LimitRTPRIO",
	"sigpending": "LimitSIGPENDING",
	"stack":      "LimitSTACK",
}

// normalize adjusts the conf to more standardized content and
// returns a new Conf with that updated content. It also returns the
// content of any script file that should accompany the conf.
func normalize(conf common.Conf, scriptPath string) (common.Conf, []byte) {
	var data []byte

	if conf.ExtraScript != "" {
		conf.ExecStart = conf.ExtraScript + "\n" + conf.ExecStart
		conf.ExtraScript = ""
	}
	if !isSimpleCommand(conf.ExecStart) {
		data = []byte(conf.ExecStart)
		conf.ExecStart = scriptPath
	}

	if len(conf.Env) == 0 {
		conf.Env = nil
	}

	if len(conf.Limit) == 0 {
		conf.Limit = nil
	}

	output := common.Unquote(conf.Output)
	if strings.HasPrefix(output, "/") {
		// Due to isSimpleCommand we can always do this safely.
		conf.ExecStart += " &> " + conf.Output
		conf.Output = ""
	}

	return conf, data
}

func isSimpleCommand(cmd string) bool {
	if strings.ContainsAny(cmd, "\n;|><&") {
		return false
	}

	return true
}

func validate(name string, conf common.Conf) error {
	if name == "" {
		return errors.NotValidf("missing service name")
	}

	if err := conf.Validate(); err != nil {
		return errors.Trace(err)
	}

	if conf.ExtraScript != "" {
		return errors.NotValidf("unexpected ExtraScript")
	}

	if conf.Output != "" && conf.Output != "syslog" {
		return errors.NotValidf("conf.Output value %q (Options are syslog)", conf.Output)
	}
	// We ignore Desc and InitDir.

	for k := range conf.Limit {
		if _, ok := limitMap[k]; !ok {
			return errors.NotValidf("conf.Limit key %q", k)
		}
	}

	if conf.Transient {
		// TODO(ericsnow) This needs to be sorted out.
		return errors.NotSupportedf("Conf.Transient")
	}

	if conf.AfterStopped != "" {
		// TODO(ericsnow) This needs to be sorted out.
		return errors.NotSupportedf("Conf.AfterStopped")
	}

	return nil
}

// serialize returns the data that should be written to disk for the
// provided Conf, rendered in the systemd unit file format.
func serialize(name string, conf common.Conf) ([]byte, error) {
	if err := validate(name, conf); err != nil {
		return nil, errors.Trace(err)
	}

	var unitOptions []*unit.UnitOption
	unitOptions = append(unitOptions, serializeUnit(conf)...)
	unitOptions = append(unitOptions, serializeService(conf)...)
	unitOptions = append(unitOptions, serializeInstall(conf)...)

	data, err := ioutil.ReadAll(unit.Serialize(unitOptions))
	return data, errors.Trace(err)
}

func serializeUnit(conf common.Conf) []*unit.UnitOption {
	var unitOptions []*unit.UnitOption

	unitOptions = append(unitOptions, &unit.UnitOption{
		Section: "Unit",
		Name:    "Description",
		Value:   conf.Desc,
	})

	after := []string{
		"syslog.target",
		"network.target",
		"systemd-user-sessions.service",
	}
	for _, name := range after {
		unitOptions = append(unitOptions, &unit.UnitOption{
			Section: "Unit",
			Name:    "After",
			Value:   name,
		})
	}

	return unitOptions
}

func serializeService(conf common.Conf) []*unit.UnitOption {
	var unitOptions []*unit.UnitOption

	unitOptions = append(unitOptions, &unit.UnitOption{
		Section: "Service",
		Name:    "Type",
		Value:   "forking",
	})

	if conf.Output != "" {
		unitOptions = append(unitOptions, &unit.UnitOption{
			Section: "Service",
			Name:    "StandardOutput",
			Value:   conf.Output,
		})
		unitOptions = append(unitOptions, &unit.UnitOption{
			Section: "Service",
			Name:    "StandardError",
			Value:   conf.Output,
		})
	}

	for k, v := range conf.Env {
		unitOptions = append(unitOptions, &unit.UnitOption{
			Section: "Service",
			Name:    "Environment",
			Value:   fmt.Sprintf(`"%q=%q"`, k, v),
		})
	}

	for k, v := range conf.Limit {
		unitOptions = append(unitOptions, &unit.UnitOption{
			Section: "Service",
			Name:    limitMap[k],
			Value:   v,
		})
	}

	unitOptions = append(unitOptions, &unit.UnitOption{
		Section: "Service",
		Name:    "ExecStart",
		Value:   conf.ExecStart,
	})

	unitOptions = append(unitOptions, &unit.UnitOption{
		Section: "Service",
		Name:    "RemainAfterExit",
		Value:   "yes",
	})

	unitOptions = append(unitOptions, &unit.UnitOption{
		Section: "Service",
		Name:    "Restart",
		Value:   "always",
	})

	if conf.ExecStopPost != "" {
		unitOptions = append(unitOptions, &unit.UnitOption{
			Section: "Service",
			Name:    "ExecStopPost",
			Value:   conf.ExecStopPost,
		})
	}

	return unitOptions
}

func serializeInstall(conf common.Conf) []*unit.UnitOption {
	var unitOptions []*unit.UnitOption

	unitOptions = append(unitOptions, &unit.UnitOption{
		Section: "Install",
		Name:    "WantedBy",
		Value:   "multi-user.target",
	})

	return unitOptions
}

// deserialize parses the provided data (in the systemd unit file
// format) and populates a new Conf with the result.
func deserialize(data []byte) (common.Conf, error) {
	opts, err := unit.Deserialize(bytes.NewBuffer(data))
	if err != nil {
		return common.Conf{}, errors.Trace(err)
	}
	return deserializeOptions(opts)
}

func deserializeOptions(opts []*unit.UnitOption) (common.Conf, error) {
	var conf common.Conf

	for _, uo := range opts {
		switch uo.Section {
		case "Unit":
			switch uo.Name {
			case "Description":
				conf.Desc = uo.Value
			case "After":
				// Do nothing until we support it in common.Conf.
			default:
				return conf, errors.NotSupportedf("Unit directive %q", uo.Name)
			}
		case "Service":
			switch {
			case uo.Name == "ExecStart":
				conf.ExecStart = uo.Value
			case uo.Name == "StandardError", uo.Name == "StandardOutput":
				// TODO(wwitzel3) We serialize Standard(Error|Output)
				// to the same thing, but we should probably make sure they match
				conf.Output = uo.Value
			case uo.Name == "Environment":
				if conf.Env == nil {
					conf.Env = make(map[string]string)
				}
				var value = uo.Value
				if strings.HasPrefix(value, `"`) && strings.HasSuffix(value, `"`) {
					value = value[1 : len(value)-1]
				}
				parts := strings.SplitN(value, "=", 2)
				if len(parts) != 2 {
					return conf, errors.NotValidf("service environment value %q", uo.Value)
				}
				conf.Env[parts[0]] = parts[1]
			case strings.HasPrefix(uo.Name, "Limit"):
				if conf.Limit == nil {
					conf.Limit = make(map[string]string)
				}
				for k, v := range limitMap {
					if v == uo.Name {
						conf.Limit[k] = v
						break
					}
				}
			case uo.Name == "Type":
				// Do nothing until we support it in common.Conf.
			case uo.Name == "RemainAfterExit":
				// Do nothing until we support it in common.Conf.
			case uo.Name == "Restart":
				// Do nothing until we support it in common.Conf.
			default:
				return conf, errors.NotSupportedf("Service directive %q", uo.Name)
			}
		case "Install":
			switch uo.Name {
			case "WantedBy":
				if uo.Value != "multi-user.target" {
					return conf, errors.NotValidf("unit target %q", uo.Value)
				}
			default:
				return conf, errors.NotSupportedf("Install directive %q", uo.Name)
			}
		default:
			return conf, errors.NotSupportedf("section %q", uo.Name)
		}
	}

	err := validate("<>", conf)
	return conf, errors.Trace(err)
}
<commit_msg>Support Conf.AfterStopped.<commit_after>// Copyright 2015 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package systemd

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"strings"

	"github.com/coreos/go-systemd/unit"
	"github.com/juju/errors"

	"github.com/juju/juju/service/common"
)

var limitMap = map[string]string{
	"as":         "LimitAS",
	"core":       "LimitCORE",
	"cpu":        "LimitCPU",
	"data":       "LimitDATA",
	"fsize":      "LimitFSIZE",
	"memlock":    "LimitMEMLOCK",
	"msgqueue":   "LimitMSGQUEUE",
	"nice":       "LimitNICE",
	"nofile":     "LimitNOFILE",
	"nproc":      "LimitNPROC",
	"rss":        "LimitRSS",
	"rtprio":     "LimitRTPRIO",
	"sigpending": "LimitSIGPENDING",
	"stack":      "LimitSTACK",
}

// normalize adjusts the conf to more standardized content and
// returns a new Conf with that updated content. It also returns the
// content of any script file that should accompany the conf.
func normalize(conf common.Conf, scriptPath string) (common.Conf, []byte) {
	var data []byte

	if conf.ExtraScript != "" {
		conf.ExecStart = conf.ExtraScript + "\n" + conf.ExecStart
		conf.ExtraScript = ""
	}
	if !isSimpleCommand(conf.ExecStart) {
		data = []byte(conf.ExecStart)
		conf.ExecStart = scriptPath
	}

	if len(conf.Env) == 0 {
		conf.Env = nil
	}

	if len(conf.Limit) == 0 {
		conf.Limit = nil
	}

	output := common.Unquote(conf.Output)
	if strings.HasPrefix(output, "/") {
		// Due to isSimpleCommand we can always do this safely.
		conf.ExecStart += " &> " + conf.Output
		conf.Output = ""
	}

	return conf, data
}

func isSimpleCommand(cmd string) bool {
	if strings.ContainsAny(cmd, "\n;|><&") {
		return false
	}

	return true
}

func validate(name string, conf common.Conf) error {
	if name == "" {
		return errors.NotValidf("missing service name")
	}

	if err := conf.Validate(); err != nil {
		return errors.Trace(err)
	}

	if conf.ExtraScript != "" {
		return errors.NotValidf("unexpected ExtraScript")
	}

	if conf.Output != "" && conf.Output != "syslog" {
		return errors.NotValidf("conf.Output value %q (Options are syslog)", conf.Output)
	}
	// We ignore Desc and InitDir.

	for k := range conf.Limit {
		if _, ok := limitMap[k]; !ok {
			return errors.NotValidf("conf.Limit key %q", k)
		}
	}

	if conf.Transient {
		// TODO(ericsnow) This needs to be sorted out.
		return errors.NotSupportedf("Conf.Transient")
	}

	return nil
}

// serialize returns the data that should be written to disk for the
// provided Conf, rendered in the systemd unit file format.
func serialize(name string, conf common.Conf) ([]byte, error) {
	if err := validate(name, conf); err != nil {
		return nil, errors.Trace(err)
	}

	var unitOptions []*unit.UnitOption
	unitOptions = append(unitOptions, serializeUnit(conf)...)
	unitOptions = append(unitOptions, serializeService(conf)...)
	unitOptions = append(unitOptions, serializeInstall(conf)...)

	data, err := ioutil.ReadAll(unit.Serialize(unitOptions))
	return data, errors.Trace(err)
}

func serializeUnit(conf common.Conf) []*unit.UnitOption {
	var unitOptions []*unit.UnitOption

	unitOptions = append(unitOptions, &unit.UnitOption{
		Section: "Unit",
		Name:    "Description",
		Value:   conf.Desc,
	})

	after := []string{
		"syslog.target",
		"network.target",
		"systemd-user-sessions.service",
	}
	for _, name := range after {
		unitOptions = append(unitOptions, &unit.UnitOption{
			Section: "Unit",
			Name:    "After",
			Value:   name,
		})
	}

	if conf.AfterStopped != "" {
		unitOptions = append(unitOptions, &unit.UnitOption{
			Section: "Unit",
			Name:    "After",
			Value:   conf.AfterStopped,
		})
		unitOptions = append(unitOptions, &unit.UnitOption{
			Section: "Unit",
			Name:    "Conflicts",
			Value:   conf.AfterStopped,
		})
	}

	return unitOptions
}

func serializeService(conf common.Conf) []*unit.UnitOption {
	var unitOptions []*unit.UnitOption

	unitOptions = append(unitOptions, &unit.UnitOption{
		Section: "Service",
		Name:    "Type",
		Value:   "forking",
	})

	if conf.Output != "" {
		unitOptions = append(unitOptions, &unit.UnitOption{
			Section: "Service",
			Name:    "StandardOutput",
			Value:   conf.Output,
		})
		unitOptions = append(unitOptions, &unit.UnitOption{
			Section: "Service",
			Name:    "StandardError",
			Value:   conf.Output,
		})
	}

	for k, v := range conf.Env {
		unitOptions = append(unitOptions, &unit.UnitOption{
			Section: "Service",
			Name:    "Environment",
			Value:   fmt.Sprintf(`"%q=%q"`, k, v),
		})
	}

	for k, v := range conf.Limit {
		unitOptions = append(unitOptions, &unit.UnitOption{
			Section: "Service",
			Name:    limitMap[k],
			Value:   v,
		})
	}

	unitOptions = append(unitOptions, &unit.UnitOption{
		Section: "Service",
		Name:    "ExecStart",
		Value:   conf.ExecStart,
	})

	unitOptions = append(unitOptions, &unit.UnitOption{
		Section: "Service",
		Name:    "RemainAfterExit",
		Value:   "yes",
	})

	unitOptions = append(unitOptions, &unit.UnitOption{
		Section: "Service",
		Name:    "Restart",
		Value:   "always",
	})

	if conf.ExecStopPost != "" {
		unitOptions = append(unitOptions, &unit.UnitOption{
			Section: "Service",
			Name:    "ExecStopPost",
			Value:   conf.ExecStopPost,
		})
	}

	return unitOptions
}

func serializeInstall(conf common.Conf) []*unit.UnitOption {
	var unitOptions []*unit.UnitOption

	unitOptions = append(unitOptions, &unit.UnitOption{
		Section: "Install",
		Name:    "WantedBy",
		Value:   "multi-user.target",
	})

	return unitOptions
}

// deserialize parses the provided data (in the systemd unit file
// format) and populates a new Conf with the result.
func deserialize(data []byte) (common.Conf, error) {
	opts, err := unit.Deserialize(bytes.NewBuffer(data))
	if err != nil {
		return common.Conf{}, errors.Trace(err)
	}
	return deserializeOptions(opts)
}

func deserializeOptions(opts []*unit.UnitOption) (common.Conf, error) {
	var conf common.Conf

	for _, uo := range opts {
		switch uo.Section {
		case "Unit":
			switch uo.Name {
			case "Description":
				conf.Desc = uo.Value
			case "After":
				// Do nothing until we support it in common.Conf.
			default:
				return conf, errors.NotSupportedf("Unit directive %q", uo.Name)
			}
		case "Service":
			switch {
			case uo.Name == "ExecStart":
				conf.ExecStart = uo.Value
			case uo.Name == "StandardError", uo.Name == "StandardOutput":
				// TODO(wwitzel3) We serialize Standard(Error|Output)
				// to the same thing, but we should probably make sure they match
				conf.Output = uo.Value
			case uo.Name == "Environment":
				if conf.Env == nil {
					conf.Env = make(map[string]string)
				}
				var value = uo.Value
				if strings.HasPrefix(value, `"`) && strings.HasSuffix(value, `"`) {
					value = value[1 : len(value)-1]
				}
				parts := strings.SplitN(value, "=", 2)
				if len(parts) != 2 {
					return conf, errors.NotValidf("service environment value %q", uo.Value)
				}
				conf.Env[parts[0]] = parts[1]
			case strings.HasPrefix(uo.Name, "Limit"):
				if conf.Limit == nil {
					conf.Limit = make(map[string]string)
				}
				for k, v := range limitMap {
					if v == uo.Name {
						conf.Limit[k] = v
						break
					}
				}
			case uo.Name == "Type":
				// Do nothing until we support it in common.Conf.
			case uo.Name == "RemainAfterExit":
				// Do nothing until we support it in common.Conf.
			case uo.Name == "Restart":
				// Do nothing until we support it in common.Conf.
			default:
				return conf, errors.NotSupportedf("Service directive %q", uo.Name)
			}
		case "Install":
			switch uo.Name {
			case "WantedBy":
				if uo.Value != "multi-user.target" {
					return conf, errors.NotValidf("unit target %q", uo.Value)
				}
			default:
				return conf, errors.NotSupportedf("Install directive %q", uo.Name)
			}
		default:
			return conf, errors.NotSupportedf("section %q", uo.Name)
		}
	}

	err := validate("<>", conf)
	return conf, errors.Trace(err)
}
<|endoftext|>