<commit_before>// Package store automatically configures a database to store structured information in an sql database
package store

import (
	"database/sql"
	"errors"
	"reflect"
	"strings"
	"sync"
)

const (
	add = iota
	get
	update
	remove
	getPage
	count
)

type field struct {
	name      string
	pos       int
	isPointer bool
	isStruct  bool
}

type typeInfo struct {
	primary    int
	fields     []field
	statements []*sql.Stmt
}

type Store struct {
	db    *sql.DB
	types map[string]typeInfo
	mutex sync.Mutex
}

func New(driverName, dataSourceName string) (*Store, error) {
	db, err := sql.Open(driverName, dataSourceName)
	if err != nil {
		return nil, err
	}
	return &Store{
		db:    db,
		types: make(map[string]typeInfo),
	}, nil
}

func (s *Store) Close() error {
	err := s.db.Close()
	s.db = nil
	return err
}

func isPointerStruct(i interface{}) bool {
	t := reflect.TypeOf(i)
	return t.Kind() == reflect.Ptr && t.Elem().Kind() == reflect.Struct
}

func (s *Store) Register(i interface{}) error {
	if s.db == nil {
		return DBClosed
	} else if !isPointerStruct(i) {
		return NoPointerStruct
	}
	s.mutex.Lock()
	defer s.mutex.Unlock()
	return s.defineType(i)
}

func (s *Store) defineType(i interface{}) error {
	name := typeName(i)
	if _, ok := s.types[name]; ok {
		return nil
	}

	s.types[name] = typeInfo{}

	v := reflect.ValueOf(i).Elem()
	numFields := v.Type().NumField()
	fields := make([]field, 0, numFields)
	id := 0
	idType := 0

	for n := 0; n < numFields; n++ {
		f := v.Type().Field(n)
		if f.PkgPath != "" { // not exported
			continue
		}
		fieldName := f.Name
		if fn := f.Tag.Get("store"); fn != "" {
			fieldName = fn
		}
		if fieldName == "-" { // Skip field
			continue
		}
		tmp := strings.ToLower(fieldName)
		for _, tf := range fields {
			if strings.ToLower(tf.name) == tmp {
				return DuplicateColumn
			}
		}
		isPointer := f.Type.Kind() == reflect.Ptr
		var iface interface{}
		if isPointer {
			iface = v.Field(n).Interface()
		} else {
			iface = v.Field(n).Addr().Interface()
		}
		isStruct := false
		if isPointerStruct(iface) {
			s.defineType(iface)
			isStruct = true
		} else if !isValidType(iface) {
			continue
		}
		if isValidKeyType(iface) {
			if idType < 3 && f.Tag.Get("key") == "1" {
				idType = 3
				id = len(fields)
			} else if idType < 2 && strings.ToLower(fieldName) == "id" {
				idType = 2
				id = len(fields)
			} else if idType < 1 {
				idType = 1
				id = len(fields)
			}
		}
		fields = append(fields, field{
			fieldName,
			n,
			isPointer,
			isStruct,
		})
	}
	if idType == 0 {
		return NoKey
	}
	s.types[name] = typeInfo{
		primary: id,
	}

	// create statements
	var (
		sqlVars, sqlParams, setSQLParams, tableVars string
		doneFirst, doneFirstNonKey                  bool
	)

	for pos, f := range fields {
		if doneFirst {
			tableVars += ", "
		} else {
			doneFirst = true
		}
		if pos != id {
			if doneFirstNonKey {
				sqlVars += ", "
				setSQLParams += ", "
				sqlParams += ", "
			} else {
				doneFirstNonKey = true
			}
		}
		var varType string
		if f.isStruct {
			varType = "INTEGER"
		} else {
			varType = getType(i, f.pos)
		}
		tableVars += "[" + f.name + "] " + varType
		if pos == id {
			tableVars += " PRIMARY KEY AUTOINCREMENT"
		} else {
			sqlVars += "[" + f.name + "]"
			setSQLParams += "[" + f.name + "] = ?"
			sqlParams += "?"
		}
	}

	statements := make([]*sql.Stmt, 6)

	sql := "CREATE TABLE IF NOT EXISTS [" + name + "](" + tableVars + ");"
	_, err := s.db.Exec(sql)
	if err != nil {
		return err
	}

	sql = "INSERT INTO [" + name + "] (" + sqlVars + ") VALUES (" + sqlParams + ");"
	stmt, err := s.db.Prepare(sql)
	if err != nil {
		return err
	}
	statements[add] = stmt

	sql = "SELECT " + sqlVars + " FROM [" + name + "] WHERE [" + fields[id].name + "] = ? LIMIT 1;"
	stmt, err = s.db.Prepare(sql)
	if err != nil {
		return err
	}
	statements[get] = stmt

	sql = "UPDATE [" + name + "] SET " + setSQLParams + " WHERE [" + fields[id].name + "] = ?;"
	stmt, err = s.db.Prepare(sql)
	if err != nil {
		return err
	}
	statements[update] = stmt

	sql = "DELETE FROM [" + name + "] WHERE [" + fields[id].name + "] = ?;"
	stmt, err = s.db.Prepare(sql)
	if err != nil {
		return err
	}
	statements[remove] = stmt

	sql = "SELECT " + sqlVars + ", [" + fields[id].name + "] FROM [" + name + "] ORDER BY [" + fields[id].name + "] LIMIT ? OFFSET ?;"
	stmt, err = s.db.Prepare(sql)
	if err != nil {
		return err
	}
	statements[getPage] = stmt

	sql = "SELECT COUNT(1) FROM [" + name + "];"
	stmt, err = s.db.Prepare(sql)
	if err != nil {
		return err
	}
	statements[count] = stmt

	s.types[name] = typeInfo{
		primary:    id,
		fields:     fields,
		statements: statements,
	}
	return nil
}

func (s *Store) Set(is ...interface{}) error {
	var toSet []interface{}
	for _, i := range is {
		t, ok := s.types[typeName(i)]
		if !ok {
			return UnregisteredType
		}
		toSet = toSet[:0]
		err := s.Set(i, &t, &toSet)
		if err != nil {
			return err
		}
	}
	return nil
}

func (s *Store) set(i interface{}, t *typeInfo, toSet *[]interface{}) error {
	for _, oi := range *toSet {
		if oi == i {
			return nil
		}
	}
	(*toSet) = append(*toSet, i)
	id := t.GetID(i)
	isUpdate := id != 0
	vars := make([]interface{}, 0, len(t.fields))
	for pos, f := range t.fields {
		if pos == t.primary {
			continue
		}
		if f.isStruct {
			ni := getFieldPointer(i, f.pos)
			nt := s.types[typeName(ni)]
			err := s.Set(ni, &nt, toSet)
			if err != nil {
				return err
			}
			v := getFieldPointer(i, f.pos)
			vt := s.types[typeName(v)]
			vars = append(vars, getField(v, vt.fields[vt.primary].pos))
		} else {
			vars = append(vars, getField(i, f.pos))
		}
	}
	if isUpdate {
		vars = append(vars, id)
		_, err := t.statements[update].Exec(vars...)
		return err
	}
	r, err := t.statements[add].Exec(vars...)
	if err != nil {
		return err
	}
	lid, err := r.LastInsertId()
	if err != nil {
		return err
	}
	t.SetID(i, lid)
	return nil
}

func (s *Store) Count(i interface{}) (int, error) {
	s.mutex.Lock()
	defer s.mutex.Unlock()
	if !isPointerStruct(i) {
		return 0, NoPointerStruct
	}
	name := typeName(i)
	stmt := s.types[name].statements[count]
	res, err := stmt.Query()
	if err != nil {
		return err
	}
	num := 0
	err = res.Scan(&num)
	return num, err
}

// Errors

var (
	DBClosed         = errors.New("database already closed")
	NoPointerStruct  = errors.New("given variable is not a pointer to a struct")
	NoKey            = errors.New("could not determine key")
	DuplicateColumn  = errors.New("duplicate column name found")
	UnregisteredType = errors.New("type not registered")
)
<commit_msg>Initial Get implementation<commit_after>// Package store automatically configures a database to store structured information in an sql database
package store

import (
	"database/sql"
	"errors"
	"reflect"
	"strings"
	"sync"
)

const (
	add = iota
	get
	update
	remove
	getPage
	count
)

type field struct {
	name      string
	pos       int
	isPointer bool
	isStruct  bool
}

type typeInfo struct {
	primary    int
	fields     []field
	statements []*sql.Stmt
}

type Store struct {
	db    *sql.DB
	types map[string]typeInfo
	mutex sync.Mutex
}

func New(driverName, dataSourceName string) (*Store, error) {
	db, err := sql.Open(driverName, dataSourceName)
	if err != nil {
		return nil, err
	}
	return &Store{
		db:    db,
		types: make(map[string]typeInfo),
	}, nil
}

func (s *Store) Close() error {
	err := s.db.Close()
	s.db = nil
	return err
}

func isPointerStruct(i interface{}) bool {
	t := reflect.TypeOf(i)
	return t.Kind() == reflect.Ptr && t.Elem().Kind() == reflect.Struct
}

func (s *Store) Register(i interface{}) error {
	if s.db == nil {
		return DBClosed
	} else if !isPointerStruct(i) {
		return NoPointerStruct
	}
	s.mutex.Lock()
	defer s.mutex.Unlock()
	return s.defineType(i)
}

func (s *Store) defineType(i interface{}) error {
	name := typeName(i)
	if _, ok := s.types[name]; ok {
		return nil
	}

	s.types[name] = typeInfo{}

	v := reflect.ValueOf(i).Elem()
	numFields := v.Type().NumField()
	fields := make([]field, 0, numFields)
	id := 0
	idType := 0

	for n := 0; n < numFields; n++ {
		f := v.Type().Field(n)
		if f.PkgPath != "" { // not exported
			continue
		}
		fieldName := f.Name
		if fn := f.Tag.Get("store"); fn != "" {
			fieldName = fn
		}
		if fieldName == "-" { // Skip field
			continue
		}
		tmp := strings.ToLower(fieldName)
		for _, tf := range fields {
			if strings.ToLower(tf.name) == tmp {
				return DuplicateColumn
			}
		}
		isPointer := f.Type.Kind() == reflect.Ptr
		var iface interface{}
		if isPointer {
			iface = v.Field(n).Interface()
		} else {
			iface = v.Field(n).Addr().Interface()
		}
		isStruct := false
		if isPointerStruct(iface) {
			s.defineType(iface)
			isStruct = true
		} else if !isValidType(iface) {
			continue
		}
		if isValidKeyType(iface) {
			if idType < 3 && f.Tag.Get("key") == "1" {
				idType = 3
				id = len(fields)
			} else if idType < 2 && strings.ToLower(fieldName) == "id" {
				idType = 2
				id = len(fields)
			} else if idType < 1 {
				idType = 1
				id = len(fields)
			}
		}
		fields = append(fields, field{
			fieldName,
			n,
			isPointer,
			isStruct,
		})
	}
	if idType == 0 {
		return NoKey
	}
	s.types[name] = typeInfo{
		primary: id,
	}

	// create statements
	var (
		sqlVars, sqlParams, setSQLParams, tableVars string
		doneFirst, doneFirstNonKey                  bool
	)

	for pos, f := range fields {
		if doneFirst {
			tableVars += ", "
		} else {
			doneFirst = true
		}
		if pos != id {
			if doneFirstNonKey {
				sqlVars += ", "
				setSQLParams += ", "
				sqlParams += ", "
			} else {
				doneFirstNonKey = true
			}
		}
		var varType string
		if f.isStruct {
			varType = "INTEGER"
		} else {
			varType = getType(i, f.pos)
		}
		tableVars += "[" + f.name + "] " + varType
		if pos == id {
			tableVars += " PRIMARY KEY AUTOINCREMENT"
		} else {
			sqlVars += "[" + f.name + "]"
			setSQLParams += "[" + f.name + "] = ?"
			sqlParams += "?"
		}
	}

	statements := make([]*sql.Stmt, 6)

	sql := "CREATE TABLE IF NOT EXISTS [" + name + "](" + tableVars + ");"
	_, err := s.db.Exec(sql)
	if err != nil {
		return err
	}

	sql = "INSERT INTO [" + name + "] (" + sqlVars + ") VALUES (" + sqlParams + ");"
	stmt, err := s.db.Prepare(sql)
	if err != nil {
		return err
	}
	statements[add] = stmt

	sql = "SELECT " + sqlVars + " FROM [" + name + "] WHERE [" + fields[id].name + "] = ? LIMIT 1;"
	stmt, err = s.db.Prepare(sql)
	if err != nil {
		return err
	}
	statements[get] = stmt

	sql = "UPDATE [" + name + "] SET " + setSQLParams + " WHERE [" + fields[id].name + "] = ?;"
	stmt, err = s.db.Prepare(sql)
	if err != nil {
		return err
	}
	statements[update] = stmt

	sql = "DELETE FROM [" + name + "] WHERE [" + fields[id].name + "] = ?;"
	stmt, err = s.db.Prepare(sql)
	if err != nil {
		return err
	}
	statements[remove] = stmt

	sql = "SELECT " + sqlVars + ", [" + fields[id].name + "] FROM [" + name + "] ORDER BY [" + fields[id].name + "] LIMIT ? OFFSET ?;"
	stmt, err = s.db.Prepare(sql)
	if err != nil {
		return err
	}
	statements[getPage] = stmt

	sql = "SELECT COUNT(1) FROM [" + name + "];"
	stmt, err = s.db.Prepare(sql)
	if err != nil {
		return err
	}
	statements[count] = stmt

	s.types[name] = typeInfo{
		primary:    id,
		fields:     fields,
		statements: statements,
	}
	return nil
}

func (s *Store) Set(is ...interface{}) error {
	var toSet []interface{}
	for _, i := range is {
		t, ok := s.types[typeName(i)]
		if !ok {
			return UnregisteredType
		}
		toSet = toSet[:0]
		err := s.Set(i, &t, &toSet)
		if err != nil {
			return err
		}
	}
	return nil
}

func (s *Store) set(i interface{}, t *typeInfo, toSet *[]interface{}) error {
	for _, oi := range *toSet {
		if oi == i {
			return nil
		}
	}
	(*toSet) = append(*toSet, i)
	id := t.GetID(i)
	isUpdate := id != 0
	vars := make([]interface{}, 0, len(t.fields))
	for pos, f := range t.fields {
		if pos == t.primary {
			continue
		}
		if f.isStruct {
			ni := getFieldPointer(i, f.pos)
			nt := s.types[typeName(ni)]
			err := s.Set(ni, &nt, toSet)
			if err != nil {
				return err
			}
			vars = append(vars, getField(ni, nt.fields[nt.primary].pos))
		} else {
			vars = append(vars, getField(i, f.pos))
		}
	}
	if isUpdate {
		vars = append(vars, id)
		_, err := t.statements[update].Exec(vars...)
		return err
	}
	r, err := t.statements[add].Exec(vars...)
	if err != nil {
		return err
	}
	lid, err := r.LastInsertId()
	if err != nil {
		return err
	}
	t.SetID(i, lid)
	return nil
}

func (s *Store) Get(is ...interface{}) error {
	for _, i := range is {
		t, ok := s.types[typeName(i)]
		if !ok {
			return UnregisteredType
		}
		id := t.GetID(i)
		if id == 0 {
			continue
		}
		vars := make([]interface{}, 0, len(t.fields))
		var toGet []interface{}
		for pos, f := range t.fields {
			if pos == t.primary {
				continue
			}
			if f.isStruct {
				ni := getFieldPointer(i, f.pos)
				nt := s.types[typeName(ni)]
				toGet = append(toGet, ni)
				vars = append(vars, getFieldPointer(ni, nt.fields[nt.primary].pos))
			} else {
				vars = append(vars, getFieldPointer(i, f.pos))
			}
		}
		row := t.statements[get].QueryRow(id)
		err := row.Scan(vars...)
		if err != nil {
			return err
		}
		if len(toGet) > 0 {
			if err = s.Get(toGet...); err != nil {
				return err
			}
		}
	}
	return nil
}

func (s *Store) Count(i interface{}) (int, error) {
	s.mutex.Lock()
	defer s.mutex.Unlock()
	if !isPointerStruct(i) {
		return 0, NoPointerStruct
	}
	name := typeName(i)
	stmt := s.types[name].statements[count]
	res, err := stmt.Query()
	if err != nil {
		return err
	}
	num := 0
	err = res.Scan(&num)
	return num, err
}

// Errors

var (
	DBClosed         = errors.New("database already closed")
	NoPointerStruct  = errors.New("given variable is not a pointer to a struct")
	NoKey            = errors.New("could not determine key")
	DuplicateColumn  = errors.New("duplicate column name found")
	UnregisteredType = errors.New("type not registered")
)
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"github.com/HearthSim/stove/bnet"
	"github.com/HearthSim/stove/pegasus"
)

const (
	CONN_HOST = "localhost"
	CONN_PORT = 1119
)

func main() {
	serv := bnet.NewServer()
	serv.RegisterGameServer("WTCG", pegasus.NewServer(serv))

	addr := fmt.Sprintf("%s:%d", CONN_HOST, CONN_PORT)
	fmt.Printf("Listening on %s ...\n", addr)
	serv.ListenAndServe(addr)
}
<commit_msg>Optionally get the bind address from the command line<commit_after>package main

import (
	"flag"
	"strings"
	"fmt"
	"github.com/HearthSim/stove/bnet"
	"github.com/HearthSim/stove/pegasus"
)

const (
	CONN_DEFAULT_HOST = "localhost"
	CONN_DEFAULT_PORT = 1119
)

func main() {
	addr := fmt.Sprintf("%s:%d", CONN_DEFAULT_HOST, CONN_DEFAULT_PORT)
	flag.StringVar(&addr, "bind", addr, "The address to run on")
	flag.Parse()

	if !strings.Contains(addr, ":") {
		addr = fmt.Sprintf("%s:%d", addr, CONN_DEFAULT_PORT)
	}

	serv := bnet.NewServer()
	serv.RegisterGameServer("WTCG", pegasus.NewServer(serv))

	fmt.Printf("Listening on %s ...\n", addr)
	serv.ListenAndServe(addr)
}
<|endoftext|>


<commit_before>// Copyright 2016 The Cockroach Authors.
//
// 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 apd

import "math/big"

// digitsLookupTable is used to map binary digit counts to their corresponding
// decimal border values. The map relies on the proof that (without leading zeros)
// for any given number of binary digits r, such that the number represented is
// between 2^r and 2^(r+1)-1, there are only two possible decimal digit counts
// k and k+1 that the binary r digits could be representing.
//
// Using this proof, for a given digit count, the map will return the lower number
// of decimal digits (k) the binary digit count could represent, along with the
// value of the border between the two decimal digit counts (10^k).
const digitsTableSize = 128

var digitsLookupTable [digitsTableSize + 1]tableVal

type tableVal struct {
	digits int64
	border big.Int
}

func init() {
	curVal := big.NewInt(1)
	curExp := new(big.Int)
	for i := 1; i <= digitsTableSize; i++ {
		if i > 1 {
			curVal.Lsh(curVal, 1)
		}

		elem := &digitsLookupTable[i]
		elem.digits = int64(len(curVal.String()))

		elem.border.SetInt64(10)
		curExp.SetInt64(elem.digits)
		elem.border.Exp(&elem.border, curExp, nil)
	}
}

func lookupBits(bitLen int) (tableVal, bool) {
	if bitLen > 0 && bitLen < len(digitsLookupTable) {
		return digitsLookupTable[bitLen], true
	}
	return tableVal{}, false
}

// numDigits returns the number of decimal digits that make up
// big.Int value. The function first attempts to look this digit
// count up in the digitsLookupTable. If the value is not there,
// it defaults to constructing a string value for the big.Int and
// using this to determine the number of digits.
func (d *Decimal) numDigits() int64 {
	return numDigits(&d.Coeff)
}

func numDigits(b *big.Int) int64 {
	bl := b.BitLen()
	if val, ok := lookupBits(bl); ok {
		ab := new(big.Int).Abs(b)
		if ab.Cmp(&val.border) < 0 {
			return val.digits
		}
		return val.digits + 1
	}

	n := int64(float64(bl) / digitsToBitsRatio)
	a := new(big.Int).Abs(b)
	e := new(big.Int).Exp(bigTen, big.NewInt(n), nil)
	if a.Cmp(e) >= 0 {
		n++
	}
	return n
}
<commit_msg>Return correct numDigits for 0<commit_after>// Copyright 2016 The Cockroach Authors.
//
// 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 apd

import "math/big"

// digitsLookupTable is used to map binary digit counts to their corresponding
// decimal border values. The map relies on the proof that (without leading zeros)
// for any given number of binary digits r, such that the number represented is
// between 2^r and 2^(r+1)-1, there are only two possible decimal digit counts
// k and k+1 that the binary r digits could be representing.
//
// Using this proof, for a given digit count, the map will return the lower number
// of decimal digits (k) the binary digit count could represent, along with the
// value of the border between the two decimal digit counts (10^k).
const digitsTableSize = 128

var digitsLookupTable [digitsTableSize + 1]tableVal

type tableVal struct {
	digits int64
	border big.Int
}

func init() {
	curVal := big.NewInt(1)
	curExp := new(big.Int)
	for i := 1; i <= digitsTableSize; i++ {
		if i > 1 {
			curVal.Lsh(curVal, 1)
		}

		elem := &digitsLookupTable[i]
		elem.digits = int64(len(curVal.String()))

		elem.border.SetInt64(10)
		curExp.SetInt64(elem.digits)
		elem.border.Exp(&elem.border, curExp, nil)
	}
}

func lookupBits(bitLen int) (tableVal, bool) {
	if bitLen > 0 && bitLen < len(digitsLookupTable) {
		return digitsLookupTable[bitLen], true
	}
	return tableVal{}, false
}

// numDigits returns the number of decimal digits that make up
// big.Int value. The function first attempts to look this digit
// count up in the digitsLookupTable. If the value is not there,
// it defaults to constructing a string value for the big.Int and
// using this to determine the number of digits.
func (d *Decimal) numDigits() int64 {
	return numDigits(&d.Coeff)
}

func numDigits(b *big.Int) int64 {
	bl := b.BitLen()
	if bl == 0 {
		return 1
	}
	if val, ok := lookupBits(bl); ok {
		ab := new(big.Int).Abs(b)
		if ab.Cmp(&val.border) < 0 {
			return val.digits
		}
		return val.digits + 1
	}

	n := int64(float64(bl) / digitsToBitsRatio)
	a := new(big.Int).Abs(b)
	e := new(big.Int).Exp(bigTen, big.NewInt(n), nil)
	if a.Cmp(e) >= 0 {
		n++
	}
	return n
}
<|endoftext|>


<commit_before>package gocli

import (
	"bufio"
	"fmt"
	"math"
	"os"
	"reflect"
	"regexp"
	"strconv"
	"strings"
	"unicode/utf8"
)

func NewTable() *Table {
	return &Table{
		Columns:   [][]string{},
		Lengths:   map[int]int{},
		Separator: "\t",
	}
}

type Table struct {
	Columns   [][]string
	Lengths   map[int]int
	Separator string

	SortBy    int
	HasHeader bool
}

func (t *Table) Select(message string) int {
	for {
		fmt.Fprintf(os.Stdout, t.StringWithIndex()+"\n"+message+": ")
		scanner := bufio.NewScanner(os.Stdin)
		scanner.Scan()
		i, e := strconv.Atoi(scanner.Text())
		if e == nil {
			if i > 0 && i <= len(t.Columns) {
				return i - 1
			}
		}
	}
}

func (t *Table) Len() int { return len(t.Columns) }

func (t *Table) Swap(a, b int) { t.Columns[a], t.Columns[b] = t.Columns[b], t.Columns[a] }

func (t *Table) Less(a, b int) bool {
	if t.HasHeader {
		if a == 0 {
			return true
		} else if b == 0 {
			return false
		}
	}
	if len(t.Columns[a]) <= t.SortBy {
		return false
	} else if len(t.Columns[b]) <= t.SortBy {
		return true
	} else {
		return fmt.Sprint(t.Columns[a][t.SortBy]) <= fmt.Sprintf(t.Columns[b][t.SortBy])
	}
	return true
}

func (t *Table) String() string {
	return strings.Join(t.Lines(false), "\n")
}

func (t *Table) StringWithIndex() string {
	return strings.Join(t.Lines(true), "\n")
}

var uncolorRegexp = regexp.MustCompile("\033\\[38;5;\\d+m([^\033]+)\033\\[0m")

func stringLength(s string) int {
	return utf8.RuneCountInString((uncolorRegexp.ReplaceAllString(s, "$1")))
}

func (t *Table) Lines(printIndex bool) (lines []string) {
	for row, col := range t.Columns {
		cl := []string{}
		if printIndex {
			col = append([]string{strconv.Itoa(row + 1)}, col...)
		}
		for i, v := range col {
			theLen := t.Lengths[i]
			if printIndex {
				if i == 0 {
					theLen = intLength(len(t.Columns))
				} else {
					theLen = t.Lengths[i-1]
				}
			}
			pad := theLen - stringLength(v)
			cl = append(cl, v+strings.Repeat(" ", pad))
		}
		lines = append(lines, strings.Join(cl, t.Separator))
	}
	return
}

func intLength(i int) int {
	if i == 0 {
		return 1
	} else if i < 0 {
		return intLength(int(math.Abs(float64(i)))) + 1
	}
	return int(math.Ceil(math.Log10(float64(i + 1))))
}

func (t *Table) AddStrings(list []string) {
	for i, s := range list {
		length := stringLength(s)
		if width := t.Lengths[i]; width < length {
			t.Lengths[i] = length
		}
	}
	t.Columns = append(t.Columns, list)
}

// Add adds a column to the table
func (t *Table) Add(cols ...interface{}) {
	converted := make([]string, 0, len(cols))
	for _, v := range cols {
		converted = append(converted, vToS(v))
	}
	t.AddStrings(converted)
}

func vToS(in interface{}) string {
	if in == nil {
		return "<nil>"
	}
	switch c := in.(type) {
	case *string:
		return *c
	case *int:
		return fmt.Sprint(*c)
	case *int64:
		return fmt.Sprint(*c)
	case *float64:
		return fmt.Sprint(*c)
	default:
		return fmt.Sprint(in)
	}
}

func (t *Table) Header(cols ...interface{}) {
	t.HasHeader = true
	t.Add(cols...)
}

// Dereferencing pointers if not nil
// TODO: Please someone tell me, how to do this right!
func (t *Table) AddP(cols ...interface{}) {
	converted := make([]string, 0, len(cols))
	var str string
	for _, v := range cols {
		if value := reflect.ValueOf(v); value.Kind() == reflect.Ptr {
			indirect := reflect.Indirect(value)
			switch {
			case indirect != reflect.Zero(value.Type()) && indirect.IsValid() == true:
				switch {
				case indirect.Kind() == reflect.String:
					str = fmt.Sprint(indirect.String())
				case indirect.Kind() == reflect.Int:
					str = fmt.Sprint(indirect.Int())
				case indirect.Kind() == reflect.Float32:
					str = fmt.Sprint(indirect.Float())
				case indirect.Kind() == reflect.Bool:
					str = fmt.Sprint(indirect.Bool())
				case indirect.Kind() == reflect.Slice:
					str = fmt.Sprint(v)
				default:
					str = fmt.Sprint(v)
				}
			default:
				str = ""
			}
		} else {
			str = fmt.Sprint(v)
		}
		converted = append(converted, str)
	}
	t.AddStrings(converted)
}
<commit_msg>support nil timers<commit_after>package gocli

import (
	"bufio"
	"fmt"
	"math"
	"os"
	"reflect"
	"regexp"
	"strconv"
	"strings"
	"time"
	"unicode/utf8"
)

func NewTable() *Table {
	return &Table{
		Columns:   [][]string{},
		Lengths:   map[int]int{},
		Separator: "\t",
	}
}

type Table struct {
	Columns   [][]string
	Lengths   map[int]int
	Separator string

	SortBy    int
	HasHeader bool
}

func (t *Table) Select(message string) int {
	for {
		fmt.Fprintf(os.Stdout, t.StringWithIndex()+"\n"+message+": ")
		scanner := bufio.NewScanner(os.Stdin)
		scanner.Scan()
		i, e := strconv.Atoi(scanner.Text())
		if e == nil {
			if i > 0 && i <= len(t.Columns) {
				return i - 1
			}
		}
	}
}

func (t *Table) Len() int { return len(t.Columns) }

func (t *Table) Swap(a, b int) { t.Columns[a], t.Columns[b] = t.Columns[b], t.Columns[a] }

func (t *Table) Less(a, b int) bool {
	if t.HasHeader {
		if a == 0 {
			return true
		} else if b == 0 {
			return false
		}
	}
	if len(t.Columns[a]) <= t.SortBy {
		return false
	} else if len(t.Columns[b]) <= t.SortBy {
		return true
	} else {
		return fmt.Sprint(t.Columns[a][t.SortBy]) <= fmt.Sprintf(t.Columns[b][t.SortBy])
	}
	return true
}

func (t *Table) String() string {
	return strings.Join(t.Lines(false), "\n")
}

func (t *Table) StringWithIndex() string {
	return strings.Join(t.Lines(true), "\n")
}

var uncolorRegexp = regexp.MustCompile("\033\\[38;5;\\d+m([^\033]+)\033\\[0m")

func stringLength(s string) int {
	return utf8.RuneCountInString((uncolorRegexp.ReplaceAllString(s, "$1")))
}

func (t *Table) Lines(printIndex bool) (lines []string) {
	for row, col := range t.Columns {
		cl := []string{}
		if printIndex {
			col = append([]string{strconv.Itoa(row + 1)}, col...)
		}
		for i, v := range col {
			theLen := t.Lengths[i]
			if printIndex {
				if i == 0 {
					theLen = intLength(len(t.Columns))
				} else {
					theLen = t.Lengths[i-1]
				}
			}
			pad := theLen - stringLength(v)
			cl = append(cl, v+strings.Repeat(" ", pad))
		}
		lines = append(lines, strings.Join(cl, t.Separator))
	}
	return
}

func intLength(i int) int {
	if i == 0 {
		return 1
	} else if i < 0 {
		return intLength(int(math.Abs(float64(i)))) + 1
	}
	return int(math.Ceil(math.Log10(float64(i + 1))))
}

func (t *Table) AddStrings(list []string) {
	for i, s := range list {
		length := stringLength(s)
		if width := t.Lengths[i]; width < length {
			t.Lengths[i] = length
		}
	}
	t.Columns = append(t.Columns, list)
}

// Add adds a column to the table
func (t *Table) Add(cols ...interface{}) {
	converted := make([]string, 0, len(cols))
	for _, v := range cols {
		converted = append(converted, vToS(v))
	}
	t.AddStrings(converted)
}

func vToS(in interface{}) string {
	v := reflect.ValueOf(in)
	if v.Kind() == reflect.Ptr && v.IsNil() {
		return "<nil>"
	}
	switch c := in.(type) {
	case *string:
		return *c
	case *int:
		return fmt.Sprint(*c)
	case *int64:
		return fmt.Sprint(*c)
	case *float64:
		return fmt.Sprint(*c)
	case *time.Time:
		return fmt.Sprint(*c)
	default:
		return fmt.Sprint(in)
	}
}

func (t *Table) Header(cols ...interface{}) {
	t.HasHeader = true
	t.Add(cols...)
}

// Dereferencing pointers if not nil
// TODO: Please someone tell me, how to do this right!
func (t *Table) AddP(cols ...interface{}) {
	converted := make([]string, 0, len(cols))
	var str string
	for _, v := range cols {
		if value := reflect.ValueOf(v); value.Kind() == reflect.Ptr {
			indirect := reflect.Indirect(value)
			switch {
			case indirect != reflect.Zero(value.Type()) && indirect.IsValid() == true:
				switch {
				case indirect.Kind() == reflect.String:
					str = fmt.Sprint(indirect.String())
				case indirect.Kind() == reflect.Int:
					str = fmt.Sprint(indirect.Int())
				case indirect.Kind() == reflect.Float32:
					str = fmt.Sprint(indirect.Float())
				case indirect.Kind() == reflect.Bool:
					str = fmt.Sprint(indirect.Bool())
				case indirect.Kind() == reflect.Slice:
					str = fmt.Sprint(v)
				default:
					str = fmt.Sprint(v)
				}
			default:
				str = ""
			}
		} else {
			str = fmt.Sprint(v)
		}
		converted = append(converted, str)
	}
	t.AddStrings(converted)
}
<|endoftext|>


<commit_before>package sqldb

import (
	"bytes"
	"database/sql"
	"fmt"
	"reflect"
	"strings"
	"unicode"
	"unicode/utf8"
)

type Column struct {
	Name         string
	Type         string
	Precision    string
	DBType       string
	Primary      bool
	AutoIncr     bool
	Notnull      bool
	Default      bool
	DefaultVal   string
	Unique       bool
	UniqueName   string
	ForeignTable string
	ForeignCol   string

	Field reflect.StructField
}

type Table struct {
	Name string
	Cols []Column
	Type reflect.Type
}

type Parser struct {
	DBDialect       DBDialect
	FieldTag        string
	Default         bool
	Notnull         bool
	TablenamePrefix string
	NameMapper      NameMapper
}

func (p *Parser) initDefault() {
	if p.DBDialect == nil {
		p.DBDialect = Postgres{}
	}
	if p.FieldTag == "" {
		p.FieldTag = "sqldb"
	}
	if p.NameMapper == nil {
		p.NameMapper = SnakeCase
	}
}

func (p *Parser) CreateTables(db *sql.DB, models ...interface{}) error {
	for _, mod := range models {
		table, err := p.StructTable(mod)
		if err != nil {
			return err
		}
		s, err := p.SQLCreate(table)
		if err != nil {
			return fmt.Errorf("%s: %s", table.Name, err.Error())
		}
		_, err = db.Exec(s)
		if err != nil {
			return fmt.Errorf("%s: %s", table.Name, err.Error())
		}
	}
	return nil
}

func (p *Parser) EscapeName(name string) string {
	return `"` + name + `"`
}

func (p *Parser) SQLCreate(table Table) (string, error) {
	var buf bytes.Buffer
	fmt.Fprintf(&buf, "CREATE TABLE IF NOT EXISTS %s (\n", p.EscapeName(table.Name))
	var (
		uniques   map[string][]string
		primaries []string
		foreigns  []int
		lastQuite string
	)
	for i, col := range table.Cols {
		dbTyp, defaultVal, err := p.DBDialect.Type(col.Type, col.Precision, col.DefaultVal)
		if err != nil {
			return "", err
		}
		if col.Primary {
			primaries = append(primaries, p.EscapeName(col.Name))
		}
		if col.ForeignTable != "" {
			foreigns = append(foreigns, i)
		}
		if col.DBType != "" {
			dbTyp = col.DBType
		}
		var constraints string
		if col.Unique {
			if col.UniqueName == "" {
				constraints += " UNIQUE"
			} else {
				if uniques == nil {
					uniques = make(map[string][]string)
				}
				uniques[col.UniqueName] = append(uniques[col.UniqueName], p.EscapeName(col.Name))
			}
		}
		if col.AutoIncr {
			constraints += " AUTO INCREAMENT"
		}
		if !col.Notnull {
			constraints += " NOT NULL"
		}
		if col.Default {
			constraints += " DEFAULT " + defaultVal
		}
		lastQuite = ""
		if i != len(table.Cols)-1 || len(primaries) != 0 || len(uniques) != 0 || len(foreigns) != 0 {
			lastQuite = ","
		}
		fmt.Fprintf(&buf, "    %s %s %s%s\n", p.EscapeName(col.Name), dbTyp, constraints, lastQuite)
	}
	if len(primaries) > 0 {
		lastQuite = ""
		if len(uniques) != 0 || len(foreigns) != 0 {
			lastQuite = ","
		}
		fmt.Fprintf(&buf, "    PRIMARY KEY (%s)%s\n", strings.Join(primaries, ","), lastQuite)
	}
	for name, keys := range uniques {
		lastQuite = ""
		if len(foreigns) != 0 || len(uniques) != 1 {
			lastQuite = ","
		}
		fmt.Fprintf(&buf, "    CONSTRAINT %s UNIQUE (%s)%s\n", name, strings.Join(keys, ","), lastQuite)
		delete(uniques, name)
	}
	for i, index := range foreigns {
		col := table.Cols[index]
		lastQuite = ""
		if i != len(foreigns)-1 {
			lastQuite = ","
		}
		fmt.Fprintf(&buf, "    FOREIGN KEY(%s) REFERENCES %s(%s)%s\n", p.EscapeName(col.Name), col.ForeignTable, col.ForeignCol, lastQuite)
	}
	fmt.Fprintf(&buf, ");\n")
	return buf.String(), nil
}

func (p *Parser) parseColumn(t *Table, f reflect.StructField) (Column, error) {
	col := Column{
		Name:    p.NameMapper(f.Name),
		Type:    f.Type.Kind().String(),
		Default: p.Default,
		Notnull: !p.Notnull,
		Field:   f,
	}

	var conds []string
	tag := strings.TrimSpace(f.Tag.Get(p.FieldTag))
	if tag != "" {
		conds = strings.Split(tag, " ")
	}
	for _, sec := range conds {
		sec = strings.TrimSpace(sec)
		var (
			keyCond  = strings.SplitN(sec, ":", 2)
			condName = keyCond[0]
			condVal  string
		)
		if len(keyCond) > 1 {
			condVal = keyCond[1]
		}

		switch condName {
		case "table":
			t.Name = condVal
		case "col":
			col.Name = condVal
			if condVal == "" {
				return col, fmt.Errorf("invalid column name")
			}
			if condVal == "-" {
				col.Name = ""
				return col, nil
			}
			col.Name = condVal
		case "type":
			if condVal == "" {
				return col, fmt.Errorf("invalid column type: %s", col.Name)
			}
			col.Type = condVal
		case "precision":
			if condVal == "" {
				return col, fmt.Errorf("invalid column precision: %s", col.Name)
			}
			col.Precision = condVal
		case "dbtype":
			if condVal == "" {
				return col, fmt.Errorf("invalid column db type: %s", col.Name)
			}
			col.DBType = condVal
		case "pk":
			col.Primary = condVal == "" || condVal == "true"
		case "autoincr":
			col.AutoIncr = condVal == "" || condVal == "true"
		case "notnull":
			col.Notnull = condVal == "" || condVal == "true"
		case "default":
			col.Default = condVal != "-"
			if p.Default {
				col.DefaultVal = condVal
			}
		case "unique":
			col.Unique = true
			col.UniqueName = condVal
		case "fk":
			fkConds := strings.SplitN(condVal, ".", 2)
			if len(fkConds) != 2 || fkConds[0] == "" || fkConds[1] == "" {
				return col, fmt.Errorf("invalid foreign key: %s", condVal)
			}
			col.ForeignTable = fkConds[0]
			col.ForeignCol = fkConds[1]
		default:
			return col, fmt.Errorf("unsupported tag: %s", condName)
		}
	}
	return col, nil
}

func (p *Parser) isPrimary(t reflect.Kind) bool {
	switch t {
	case reflect.Bool,
		reflect.Int,
		reflect.Int8,
		reflect.Int16,
		reflect.Int32,
		reflect.Int64,
		reflect.Uint,
		reflect.Uint8,
		reflect.Uint16,
		reflect.Uint32,
		reflect.Uint64,
		reflect.Float32,
		reflect.Float64,
		reflect.String:
		return true
	}
	return false
}

func (p *Parser) shouldIgnore(f *reflect.StructField) bool {
	if f.Tag.Get(p.FieldTag) == "-" {
		return true
	}
	if f.Type.Kind() == reflect.Struct {
		return !f.Anonymous
	}
	if !p.isPrimary(f.Type.Kind()) {
		return true
	}
	return unicode.IsLower([]rune(f.Name)[0])
}

func (p *Parser) structFields(fields []reflect.StructField, t reflect.Type) []reflect.StructField {
	n := t.NumField()

	var anonymousStructs []reflect.StructField
	for i := 0; i < n; i++ {
		f := t.Field(i)
		if p.shouldIgnore(&f) {
			continue
		}
		if f.Anonymous && f.Type.Kind() == reflect.Struct {
			anonymousStructs = append(anonymousStructs, f)
			p.structFields(fields, f.Type)
		} else {
			var override bool
			for i := range fields {
				if fields[i].Name == f.Name {
					override = true
					break
				}
			}
			if !override {
				fields = append(fields, f)
			}
		}
	}
	for _, f := range anonymousStructs {
		fields = p.structFields(fields, f.Type)
	}
	return fields
}

func (p *Parser) StructTable(v interface{}) (Table, error) {
	p.initDefault()

	refv := reflect.ValueOf(v)
	if refv.Kind() == reflect.Ptr {
		refv = refv.Elem()
	}
	if refv.Kind() != reflect.Struct {
		return Table{}, fmt.Errorf("invalid artument type, expect (pointer of) structure")
	}
	reft := refv.Type()

	t := Table{
		Name: p.TablenamePrefix + p.NameMapper(reft.Name()),
		Type: reft,
	}
	fields := p.structFields(nil, reft)
	for _, f := range fields {
		col, err := p.parseColumn(&t, f)
		if err != nil {
			return t, err
		}
		if col.Name != "" {
			t.Cols = append(t.Cols, col)
		}
	}
	return t, nil
}

func SnakeCase(s string) string {
	var (
		hasUpper  bool
		size      = utf8.RuneCountInString(s)
		prevUpper bool
	)
	for i, r := range s {
		if unicode.IsUpper(r) {
			hasUpper = true
			if i != 0 && !prevUpper {
				size++
			}
			prevUpper = true
		} else {
			prevUpper = false
		}
	}
	if !hasUpper {
		return s
	}
	var (
		buf = make([]rune, 0, size)
	)
	prevUpper = false
	for i, r := range s {
		isUpper := unicode.IsUpper(r)
		if isUpper && i != 0 && !prevUpper {
			buf = append(buf, '_')
		}
		if prevUpper = isUpper; isUpper {
			buf = append(buf, unicode.ToLower(r))
		} else {
			buf = append(buf, r)
		}
	}
	return string(buf)
}
<commit_msg>fixed embed field index<commit_after>package sqldb

import (
	"bytes"
	"database/sql"
	"fmt"
	"reflect"
	"strings"
	"unicode"
	"unicode/utf8"
)

type Column struct {
	Name         string
	Type         string
	Precision    string
	DBType       string
	Primary      bool
	AutoIncr     bool
	Notnull      bool
	Default      bool
	DefaultVal   string
	Unique       bool
	UniqueName   string
	ForeignTable string
	ForeignCol   string

	Field reflect.StructField
}

type Table struct {
	Name string
	Cols []Column
	Type reflect.Type
}

type Parser struct {
	DBDialect       DBDialect
	FieldTag        string
	Default         bool
	Notnull         bool
	TablenamePrefix string
	NameMapper      NameMapper
}

func (p *Parser) initDefault() {
	if p.DBDialect == nil {
		p.DBDialect = Postgres{}
	}
	if p.FieldTag == "" {
		p.FieldTag = "sqldb"
	}
	if p.NameMapper == nil {
		p.NameMapper = SnakeCase
	}
}

func (p *Parser) CreateTables(db *sql.DB, models ...interface{}) error {
	for _, mod := range models {
		table, err := p.StructTable(mod)
		if err != nil {
			return err
		}
		s, err := p.SQLCreate(table)
		if err != nil {
			return fmt.Errorf("%s: %s", table.Name, err.Error())
		}
		_, err = db.Exec(s)
		if err != nil {
			return fmt.Errorf("%s: %s", table.Name, err.Error())
		}
	}
	return nil
}

func (p *Parser) EscapeName(name string) string {
	return `"` + name + `"`
}

func (p *Parser) SQLCreate(table Table) (string, error) {
	var buf bytes.Buffer
	fmt.Fprintf(&buf, "CREATE TABLE IF NOT EXISTS %s (\n", p.EscapeName(table.Name))
	var (
		uniques   map[string][]string
		primaries []string
		foreigns  []int
		lastQuite string
	)
	for i, col := range table.Cols {
		dbTyp, defaultVal, err := p.DBDialect.Type(col.Type, col.Precision, col.DefaultVal)
		if err != nil {
			return "", err
		}
		if col.Primary {
			primaries = append(primaries, p.EscapeName(col.Name))
		}
		if col.ForeignTable != "" {
			foreigns = append(foreigns, i)
		}
		if col.DBType != "" {
			dbTyp = col.DBType
		}
		var constraints string
		if col.Unique {
			if col.UniqueName == "" {
				constraints += " UNIQUE"
			} else {
				if uniques == nil {
					uniques = make(map[string][]string)
				}
				uniques[col.UniqueName] = append(uniques[col.UniqueName], p.EscapeName(col.Name))
			}
		}
		if col.AutoIncr {
			constraints += " AUTO INCREAMENT"
		}
		if !col.Notnull {
			constraints += " NOT NULL"
		}
		if col.Default {
			constraints += " DEFAULT " + defaultVal
		}
		lastQuite = ""
		if i != len(table.Cols)-1 || len(primaries) != 0 || len(uniques) != 0 || len(foreigns) != 0 {
			lastQuite = ","
		}
		fmt.Fprintf(&buf, "    %s %s %s%s\n", p.EscapeName(col.Name), dbTyp, constraints, lastQuite)
	}
	if len(primaries) > 0 {
		lastQuite = ""
		if len(uniques) != 0 || len(foreigns) != 0 {
			lastQuite = ","
		}
		fmt.Fprintf(&buf, "    PRIMARY KEY (%s)%s\n", strings.Join(primaries, ","), lastQuite)
	}
	for name, keys := range uniques {
		lastQuite = ""
		if len(foreigns) != 0 || len(uniques) != 1 {
			lastQuite = ","
		}
		fmt.Fprintf(&buf, "    CONSTRAINT %s UNIQUE (%s)%s\n", name, strings.Join(keys, ","), lastQuite)
		delete(uniques, name)
	}
	for i, index := range foreigns {
		col := table.Cols[index]
		lastQuite = ""
		if i != len(foreigns)-1 {
			lastQuite = ","
		}
		fmt.Fprintf(&buf, "    FOREIGN KEY(%s) REFERENCES %s(%s)%s\n", p.EscapeName(col.Name), col.ForeignTable, col.ForeignCol, lastQuite)
	}
	fmt.Fprintf(&buf, ");\n")
	return buf.String(), nil
}

func (p *Parser) parseColumn(t *Table, f reflect.StructField) (Column, error) {
	col := Column{
		Name:    p.NameMapper(f.Name),
		Type:    f.Type.Kind().String(),
		Default: p.Default,
		Notnull: !p.Notnull,
		Field:   f,
	}

	var conds []string
	tag := strings.TrimSpace(f.Tag.Get(p.FieldTag))
	if tag != "" {
		conds = strings.Split(tag, " ")
	}
	for _, sec := range conds {
		sec = strings.TrimSpace(sec)
		var (
			keyCond  = strings.SplitN(sec, ":", 2)
			condName = keyCond[0]
			condVal  string
		)
		if len(keyCond) > 1 {
			condVal = keyCond[1]
		}

		switch condName {
		case "table":
			t.Name = condVal
		case "col":
			col.Name = condVal
			if condVal == "" {
				return col, fmt.Errorf("invalid column name")
			}
			if condVal == "-" {
				col.Name = ""
				return col, nil
			}
			col.Name = condVal
		case "type":
			if condVal == "" {
				return col, fmt.Errorf("invalid column type: %s", col.Name)
			}
			col.Type = condVal
		case "precision":
			if condVal == "" {
				return col, fmt.Errorf("invalid column precision: %s", col.Name)
			}
			col.Precision = condVal
		case "dbtype":
			if condVal == "" {
				return col, fmt.Errorf("invalid column db type: %s", col.Name)
			}
			col.DBType = condVal
		case "pk":
			col.Primary = condVal == "" || condVal == "true"
		case "autoincr":
			col.AutoIncr = condVal == "" || condVal == "true"
		case "notnull":
			col.Notnull = condVal == "" || condVal == "true"
		case "default":
			col.Default = condVal != "-"
			if p.Default {
				col.DefaultVal = condVal
			}
		case "unique":
			col.Unique = true
			col.UniqueName = condVal
		case "fk":
			fkConds := strings.SplitN(condVal, ".", 2)
			if len(fkConds) != 2 || fkConds[0] == "" || fkConds[1] == "" {
				return col, fmt.Errorf("invalid foreign key: %s", condVal)
			}
			col.ForeignTable = fkConds[0]
			col.ForeignCol = fkConds[1]
		default:
			return col, fmt.Errorf("unsupported tag: %s", condName)
		}
	}
	return col, nil
}

func (p *Parser) isPrimary(t reflect.Kind) bool {
	switch t {
	case reflect.Bool,
		reflect.Int,
		reflect.Int8,
		reflect.Int16,
		reflect.Int32,
		reflect.Int64,
		reflect.Uint,
		reflect.Uint8,
		reflect.Uint16,
		reflect.Uint32,
		reflect.Uint64,
		reflect.Float32,
		reflect.Float64,
		reflect.String:
		return true
	}
	return false
}

func (p *Parser) shouldIgnore(f *reflect.StructField) bool {
	if f.Tag.Get(p.FieldTag) == "-" {
		return true
	}
	if f.Type.Kind() == reflect.Struct {
		return !f.Anonymous
	}
	if !p.isPrimary(f.Type.Kind()) {
		return true
	}
	return unicode.IsLower([]rune(f.Name)[0])
}

func (p *Parser) concatIndexes(parent, child []int) []int {
	if len(parent) == 0 {
		return child
	}
	indexes := make([]int, 0, len(parent)+len(child))
	indexes = append(indexes, parent...)
	indexes = append(indexes, child...)
	return indexes
}

func (p *Parser) structFields(fields []reflect.StructField, parentFieldIndexes []int, t reflect.Type) []reflect.StructField {
	n := t.NumField()

	var anonymousStructs []reflect.StructField
	for i := 0; i < n; i++ {
		f := t.Field(i)
		if p.shouldIgnore(&f) {
			continue
		}
		if f.Anonymous && f.Type.Kind() == reflect.Struct {
			anonymousStructs = append(anonymousStructs, f)
		} else {
			var override bool
			for i := range fields {
				if fields[i].Name == f.Name {
					override = true
					break
				}
			}
			if !override {
				f.Index = p.concatIndexes(parentFieldIndexes, f.Index)
				fields = append(fields, f)
			}
		}
	}
	for _, f := range anonymousStructs {
		fields = p.structFields(fields, p.concatIndexes(parentFieldIndexes, f.Index), f.Type)
	}
	return fields
}

func (p *Parser) StructTable(v interface{}) (Table, error) {
	p.initDefault()

	refv := reflect.ValueOf(v)
	if refv.Kind() == reflect.Ptr {
		refv = refv.Elem()
	}
	if refv.Kind() != reflect.Struct {
		return Table{}, fmt.Errorf("invalid artument type, expect (pointer of) structure")
	}
	reft := refv.Type()

	t := Table{
		Name: p.TablenamePrefix + p.NameMapper(reft.Name()),
		Type: reft,
	}
	fields := p.structFields(nil, nil, reft)
	for _, f := range fields {
		col, err := p.parseColumn(&t, f)
		if err != nil {
			return t, err
		}
		if col.Name != "" {
			t.Cols = append(t.Cols, col)
		}
	}
	return t, nil
}

func SnakeCase(s string) string {
	var (
		hasUpper  bool
		size      = utf8.RuneCountInString(s)
		prevUpper bool
	)
	for i, r := range s {
		if unicode.IsUpper(r) {
			hasUpper = true
			if i != 0 && !prevUpper {
				size++
			}
			prevUpper = true
		} else {
			prevUpper = false
		}
	}
	if !hasUpper {
		return s
	}
	var (
		buf = make([]rune, 0, size)
	)
	prevUpper = false
	for i, r := range s {
		isUpper := unicode.IsUpper(r)
		if isUpper && i != 0 && !prevUpper {
			buf = append(buf, '_')
		}
		if prevUpper = isUpper; isUpper {
			buf = append(buf, unicode.ToLower(r))
		} else {
			buf = append(buf, r)
		}
	}
	return string(buf)
}
<|endoftext|>


<commit_before>package somaproto

import (
	"github.com/satori/go.uuid"
)

type ProtoRequestTeam struct {
	Team   ProtoTeam       `json:"team,omitempty"`
	Filter ProtoTeamFilter `json:"filter,omitempty"`
}

type ProtoResultTeam struct {
	Code   uint16      `json:"code,omitempty"`
	Status string      `json:"status,omitempty"`
	Text   []string    `json:"text,omitempty"`
	Teams  []ProtoTeam `json:"teams,omitempty"`
}

type ProtoTeam struct {
	TeamId   uuid.UUID         `json:"teamid,omitempty"`
	TeamName string            `json:"teamname,omitempty"`
	LdapId   string            `json:"ldapid,omitempty"`
	System   bool              `json:"system,omitempty"`
	Details  *ProtoTeamDetails `json:"details,omitempty"`
}

type ProtoTeamDetails struct {
	CreatedAt string   `json:"createdat,omitempty"`
	CreatedBy string   `json:"createdby,omitempty"`
	Members   []string `json:"members,omitempty"`
}

type ProtoTeamFilter struct {
	TeamName string `json:"teamname,omitempty"`
	LdapId   string `json:"ldapid,omitempty"`
	System   bool   `json:"system,omitempty"`
}

// vim: ts=4 sw=4 sts=4 noet fenc=utf-8 ffs=unix
<commit_msg>Update structs for team<commit_after>package somaproto

type ProtoRequestTeam struct {
	Team   *ProtoTeam       `json:"team,omitempty"`
	Filter *ProtoTeamFilter `json:"filter,omitempty"`
}

type ProtoResultTeam struct {
	Code   uint16      `json:"code,omitempty"`
	Status string      `json:"status,omitempty"`
	Text   []string    `json:"text,omitempty"`
	Teams  []ProtoTeam `json:"teams,omitempty"`
}

type ProtoTeam struct {
	Id      string            `json:"id,omitempty"`
	Name    string            `json:"name,omitempty"`
	Ldap    string            `json:"ldap,omitempty"`
	System  bool              `json:"system,omitempty"`
	Details *ProtoTeamDetails `json:"details,omitempty"`
}

type ProtoTeamDetails struct {
	CreatedAt string   `json:"createdat,omitempty"`
	CreatedBy string   `json:"createdby,omitempty"`
	Members   []string `json:"members,omitempty"`
}

type ProtoTeamFilter struct {
	Name   string `json:"name,omitempty"`
	Ldap   string `json:"ldap,omitempty"`
	System bool   `json:"system,omitempty"`
}

// vim: ts=4 sw=4 sts=4 noet fenc=utf-8 ffs=unix
<|endoftext|>


<commit_before>package main

import (
	"crypto/tls"
	"encoding/json"
	"fmt"
	"net/rpc"
	"net/rpc/jsonrpc"
	"os"

	"github.com/bitmark-inc/exitwithstatus"
	"github.com/bitmark-inc/getoptions"
)

type RPCEmptyArguments struct{}

type ConnClient struct {
	Clients []string `json:"clients"`
}

type RPCClient struct {
	Client *rpc.Client
}

// GetNodeInfo will get the node info of a node from bitmark rpc
func (r *RPCClient) GetNodeInfo() (json.RawMessage, error) {
	args := RPCEmptyArguments{}
	var reply json.RawMessage
	err := r.Client.Call("Node.Info", &args, &reply)
	return reply, err
}

// GetSubscribers will get all its subscribers of a node from bitmark rpc
func (r *RPCClient) GetSubscribers() (ConnClient, error) {
	args := RPCEmptyArguments{}
	var reply ConnClient
	err := r.Client.Call("Node.Subscribers", &args, &reply)
	return reply, err
}

// GetConnectors will get all its connectors of a node from bitmark rpc
func (r *RPCClient) GetConnectors() (ConnClient, error) {
	args := RPCEmptyArguments{}
	var reply ConnClient
	err := r.Client.Call("Node.Connectors", &args, &reply)
	return reply, err
}

func (r *RPCClient) GetAllInfo() (reply map[string]interface{}, err error) {
	node, err := r.GetNodeInfo()
	if err != nil {
		return
	}
	sbsc, err := r.GetSubscribers()
	if err != nil {
		return
	}
	conn, err := r.GetConnectors()
	if err != nil {
		return
	}
	reply = map[string]interface{}{
		"node": node,
		"sbsc": sbsc,
		"conn": conn,
	}
	return
}

func main() {
	defer exitwithstatus.Handler()

	flags := []getoptions.Option{
		{Long: "help", HasArg: getoptions.NO_ARGUMENT, Short: 'h'},
		{Long: "info-type", HasArg: getoptions.OPTIONAL_ARGUMENT, Short: 'i'},
	}

	program, options, arguments, err := getoptions.GetOS(flags)
	if err != nil {
		exitwithstatus.Message("option parse error: %s", err)
	}

	if len(options["help"]) > 0 {
		exitwithstatus.Message("usage: %s [--help] [--info-type=TYPE] [host:port]", program)
	}

	// set the default info type
	infoType := []string{"node"}

	if len(options["info-type"]) != 0 {
		infoType = options["info-type"]
	}

	var hostPort string
	if len(arguments) != 0 {
		hostPort = arguments[0]
	}

	// establish rpc connection over tls
	conn, err := tls.Dial("tcp", hostPort, &tls.Config{
		InsecureSkipVerify: true,
	})
	if err != nil {
		exitwithstatus.Message("dial error: %s", err)
	}
	defer conn.Close()
	client := jsonrpc.NewClient(conn)

	r := RPCClient{client}
	reply := map[string]interface{}{}

	for _, t := range infoType {
		var v interface{}
		switch t {
		case "all":
			reply, err = r.GetAllInfo()
			break
		case "node":
			v, err = r.GetNodeInfo()
			reply["node"] = v
		case "sbsc":
			v, err = r.GetSubscribers()
			reply["sbsc"] = v
		case "conn":
			v, err = r.GetConnectors()
			reply["conn"] = v
		default:
			err = fmt.Errorf("incorrect info type provided: %s", infoType)
		}
	}

	if err != nil {
		exitwithstatus.Message("rpc error: %s", err)
	}

	b, err := json.Marshal(reply)
	if err != nil {
		exitwithstatus.Message("incorrect json marshal: %s", err)
	}

	fmt.Printf("%s", b)
	os.Exit(0)
}
<commit_msg>[command] add rpc host for each info return<commit_after>package main

import (
	"crypto/tls"
	"encoding/json"
	"fmt"
	"net/rpc"
	"net/rpc/jsonrpc"
	"os"

	"github.com/bitmark-inc/exitwithstatus"
	"github.com/bitmark-inc/getoptions"
)

type RPCEmptyArguments struct{}

type ConnClient struct {
	Clients []string `json:"clients"`
}

type RPCClient struct {
	Client *rpc.Client
}

// GetNodeInfo will get the node info of a node from bitmark rpc
func (r *RPCClient) GetNodeInfo() (json.RawMessage, error) {
	args := RPCEmptyArguments{}
	var reply json.RawMessage
	err := r.Client.Call("Node.Info", &args, &reply)
	return reply, err
}

// GetSubscribers will get all its subscribers of a node from bitmark rpc
func (r *RPCClient) GetSubscribers() (ConnClient, error) {
	args := RPCEmptyArguments{}
	var reply ConnClient
	err := r.Client.Call("Node.Subscribers", &args, &reply)
	return reply, err
}

// GetConnectors will get all its connectors of a node from bitmark rpc
func (r *RPCClient) GetConnectors() (ConnClient, error) {
	args := RPCEmptyArguments{}
	var reply ConnClient
	err := r.Client.Call("Node.Connectors", &args, &reply)
	return reply, err
}

func (r *RPCClient) GetAllInfo() (reply map[string]interface{}, err error) {
	node, err := r.GetNodeInfo()
	if err != nil {
		return
	}
	sbsc, err := r.GetSubscribers()
	if err != nil {
		return
	}
	conn, err := r.GetConnectors()
	if err != nil {
		return
	}
	reply = map[string]interface{}{
		"node": node,
		"sbsc": sbsc,
		"conn": conn,
	}
	return
}

func main() {
	defer exitwithstatus.Handler()

	flags := []getoptions.Option{
		{Long: "help", HasArg: getoptions.NO_ARGUMENT, Short: 'h'},
		{Long: "info-type", HasArg: getoptions.OPTIONAL_ARGUMENT, Short: 'i'},
	}

	program, options, arguments, err := getoptions.GetOS(flags)
	if err != nil {
		exitwithstatus.Message("option parse error: %s", err)
	}

	if len(options["help"]) > 0 {
		exitwithstatus.Message("usage: %s [--help] [--info-type=TYPE] [host:port]", program)
	}

	// set the default info type
	infoType := []string{"node"}

	if len(options["info-type"]) != 0 {
		infoType = options["info-type"]
	}

	var hostPort string
	if len(arguments) != 0 {
		hostPort = arguments[0]
	}

	// establish rpc connection over tls
	conn, err := tls.Dial("tcp", hostPort, &tls.Config{
		InsecureSkipVerify: true,
	})
	if err != nil {
		exitwithstatus.Message("dial error: %s", err)
	}
	defer conn.Close()
	client := jsonrpc.NewClient(conn)

	r := RPCClient{client}

	reply := map[string]interface{}{
		"host": fmt.Sprintf("tcp://%s", hostPort),
	}

	for _, t := range infoType {
		var v interface{}
		switch t {
		case "all":
			reply, err = r.GetAllInfo()
			break
		case "node":
			v, err = r.GetNodeInfo()
			reply["node"] = v
		case "sbsc":
			v, err = r.GetSubscribers()
			reply["sbsc"] = v
		case "conn":
			v, err = r.GetConnectors()
			reply["conn"] = v
		default:
			err = fmt.Errorf("incorrect info type provided: %s", infoType)
		}
	}

	if err != nil {
		exitwithstatus.Message("rpc error: %s", err)
	}

	b, err := json.Marshal(reply)
	if err != nil {
		exitwithstatus.Message("incorrect json marshal: %s", err)
	}

	fmt.Printf("%s", b)
	os.Exit(0)
}
<|endoftext|>


<commit_before>package compressor

import (
	"compress/gzip"
	"io"
	"os"
)

type Compressor interface {
	Compress(src string, dst string) error
}

func NewTgz() Compressor {
	return &tgzCompressor{}
}

type tgzCompressor struct{}

func (compressor *tgzCompressor) Compress(src string, dest string) error {
	reader := NewTarReader(src)

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

	gw := gzip.NewWriter(fw)
	defer gw.Close()

	_, err = io.Copy(gw, reader)
	return err
}
<commit_msg>use WriteTar in compressor<commit_after>package compressor

import (
	"compress/gzip"
	"os"
)

type Compressor interface {
	Compress(src string, dst string) error
}

func NewTgz() Compressor {
	return &tgzCompressor{}
}

type tgzCompressor struct{}

func (compressor *tgzCompressor) Compress(src string, dest string) error {
	fw, err := os.Create(dest)
	if err != nil {
		return err
	}
	defer fw.Close()

	gw := gzip.NewWriter(fw)
	defer gw.Close()

	return WriteTar(src, gw)
}
<|endoftext|>


<commit_before>package datastore

import (
	"database/sql"
	"errors"

	"github.com/anthonynsimon/parrot/datastore/postgres"
	"github.com/anthonynsimon/parrot/model"
)

var (
	ErrNoDB           = errors.New("couldn't get DB")
	ErrNotImplemented = errors.New("database not implemented")
)

type Datastore struct {
	Store
}

type Store interface {
	model.DocStorer
	Ping() error
	Close() error
}

func NewDatastore(name string, url string) (*Datastore, error) {
	var ds *Datastore

	switch name {
	case "postgres":
		conn, err := sql.Open("postgres", url)
		if err != nil {
			return nil, err
		}
		p := &postgres.PostgresDB{DB: conn}
		ds = &Datastore{p}
	default:
		return nil, ErrNotImplemented
	}

	return ds, nil
}
<commit_msg>Fmt code<commit_after>package datastore

import (
	"database/sql"
	"errors"

	"github.com/anthonynsimon/parrot/datastore/postgres"
	"github.com/anthonynsimon/parrot/model"
)

type Store interface {
	model.DocStorer
	Ping() error
	Close() error
}

var (
	ErrNoDB           = errors.New("couldn't get DB")
	ErrNotImplemented = errors.New("database not implemented")
)

type Datastore struct {
	Store
}

func NewDatastore(name string, url string) (*Datastore, error) {
	var ds *Datastore

	switch name {
	case "postgres":
		conn, err := sql.Open("postgres", url)
		if err != nil {
			return nil, err
		}
		p := &postgres.PostgresDB{DB: conn}
		ds = &Datastore{p}
	default:
		return nil, ErrNotImplemented
	}

	return ds, nil
}
<|endoftext|>


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

package main

import (
	"bytes"
	"os"
	"path/filepath"
	"text/template"
	"time"

	"github.com/Masterminds/sprig"
	"github.com/banzaicloud/bank-vaults/pkg/vault"
	"github.com/fsnotify/fsnotify"
	"github.com/sirupsen/logrus"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"
)

const (
	cfgVaultConfigFile = "vault-config-file"
	cfgFatal           = "fatal"
)

var configureCmd = &cobra.Command{
	Use:   "configure",
	Short: "Configures a Vault based on a YAML/JSON configuration file",
	Long: `This configuration is an extension to what is available through the Vault configuration:
			https://www.vaultproject.io/docs/configuration/index.html. With this it is possible to
			configure secret engines, auth methods, etc...`,
	Run: func(cmd *cobra.Command, args []string) {
		appConfig.BindPFlag(cfgOnce, cmd.PersistentFlags().Lookup(cfgOnce))
		appConfig.BindPFlag(cfgFatal, cmd.PersistentFlags().Lookup(cfgFatal))
		appConfig.BindPFlag(cfgUnsealPeriod, cmd.PersistentFlags().Lookup(cfgUnsealPeriod))
		appConfig.BindPFlag(cfgVaultConfigFile, cmd.PersistentFlags().Lookup(cfgVaultConfigFile))

		var unsealConfig unsealCfg

		runOnce := appConfig.GetBool(cfgOnce)
		errorFatal := appConfig.GetBool(cfgFatal)
		unsealConfig.unsealPeriod = appConfig.GetDuration(cfgUnsealPeriod)
		vaultConfigFiles := appConfig.GetStringSlice(cfgVaultConfigFile)

		store, err := kvStoreForConfig(appConfig)

		if err != nil {
			logrus.Fatalf("error creating kv store: %s", err.Error())
		}

		cl, err := vault.NewRawClient()

		if err != nil {
			logrus.Fatalf("error connecting to vault: %s", err.Error())
		}

		vaultConfig, err := vaultConfigForConfig(appConfig)

		if err != nil {
			logrus.Fatalf("error building vault config: %s", err.Error())
		}

		v, err := vault.New(store, cl, vaultConfig)

		if err != nil {
			logrus.Fatalf("error creating vault helper: %s", err.Error())
		}

		metrics := prometheusExporter{Vault: v, Mode: "configure"}
		go metrics.Run()

		configurations := make(chan *viper.Viper, len(vaultConfigFiles))

		for _, vaultConfigFile := range vaultConfigFiles {
			configurations <- parseConfiguration(vaultConfigFile)
		}

		if !runOnce {
			go watchConfigurations(vaultConfigFiles, configurations)
		} else {
			close(configurations)
		}

		for config := range configurations {

			logrus.Infoln("config file has changed:", config.ConfigFileUsed())

			func() {
				for {
					logrus.Infof("checking if vault is sealed...")
					sealed, err := v.Sealed()
					if err != nil {
						logrus.Errorf("error checking if vault is sealed: %s, waiting %s before trying again...", err.Error(), unsealConfig.unsealPeriod)
						time.Sleep(unsealConfig.unsealPeriod)
						continue
					}

					// If vault is sealed, we stop here and wait another unsealPeriod
					if sealed {
						logrus.Infof("vault is sealed, waiting %s before trying again...", unsealConfig.unsealPeriod)
						time.Sleep(unsealConfig.unsealPeriod)
						continue
					}

					logrus.Info("vault is unsealed, configuring...")

					if err = v.Configure(config); err != nil {
						logrus.Errorf("error configuring vault: %s", err.Error())
						if errorFatal {
							os.Exit(1)
						}
						failedConfigurationsCount++
						return
					}

					successfulConfigurationsCount++
					logrus.Info("successfully configured vault")
					return
				}
			}()
		}
	},
}

func watchConfigurations(vaultConfigFiles []string, configurations chan *viper.Viper) {
	watcher, err := fsnotify.NewWatcher()
	if err != nil {
		logrus.Fatal(err)
	}
	defer watcher.Close()

	for _, vaultConfigFile := range vaultConfigFiles {
		// we have to watch the entire directory to pick up renames/atomic saves in a cross-platform way
		configFile := filepath.Clean(vaultConfigFile)
		configDir, _ := filepath.Split(configFile)

		done := make(chan bool)
		go func() {
			for {
				select {
				case event := <-watcher.Events:
					// we only care about the config file or the ConfigMap directory (if in Kubernetes)
					if filepath.Clean(event.Name) == configFile || filepath.Base(event.Name) == "..data" {
						if event.Op&fsnotify.Write == fsnotify.Write || event.Op&fsnotify.Create == fsnotify.Create {
							configurations <- parseConfiguration(configFile)
						}
					}
				case err := <-watcher.Errors:
					logrus.Error(err)
				}
			}
		}()

		watcher.Add(configDir)
		<-done
	}
}

func parseConfiguration(vaultConfigFile string) *viper.Viper {

	config := viper.New()

	templateName := filepath.Base(vaultConfigFile)

	configTemplate, err := template.New(templateName).
		Funcs(sprig.TxtFuncMap()).
		Delims("${", "}").
		ParseFiles(vaultConfigFile)

	if err != nil {
		logrus.Fatalf("error parsing vault config template: %s", err.Error())
	}

	buffer := bytes.NewBuffer(nil)

	err = configTemplate.ExecuteTemplate(buffer, templateName, nil)
	if err != nil {
		logrus.Fatalf("error executing vault config template: %s", err.Error())
	}

	config.SetConfigFile(vaultConfigFile)

	err = config.ReadConfig(buffer)
	if err != nil {
		logrus.Fatalf("error reading vault config file: %s", err.Error())
	}

	return config
}

func init() {
	configureCmd.PersistentFlags().Bool(cfgOnce, false, "Run configure only once")
	configureCmd.PersistentFlags().Bool(cfgFatal, false, "Make configuration errors fatal to the configurator")
	configureCmd.PersistentFlags().Duration(cfgUnsealPeriod, time.Second*5, "How often to attempt to unseal the Vault instance")
	configureCmd.PersistentFlags().StringSlice(cfgVaultConfigFile, []string{vault.DefaultConfigFile}, "The filename of the YAML/JSON Vault configuration")

	rootCmd.AddCommand(configureCmd)
}
<commit_msg>Make Watch function work on multiple Directories<commit_after>// Copyright © 2019 Banzai Cloud
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"bytes"
	"os"
	"path/filepath"
	"strings"
	"text/template"
	"time"

	"github.com/Masterminds/sprig"
	"github.com/banzaicloud/bank-vaults/pkg/vault"
	"github.com/fsnotify/fsnotify"
	"github.com/sirupsen/logrus"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"
)

const (
	cfgVaultConfigFile = "vault-config-file"
	cfgFatal           = "fatal"
)

var configureCmd = &cobra.Command{
	Use:   "configure",
	Short: "Configures a Vault based on a YAML/JSON configuration file",
	Long: `This configuration is an extension to what is available through the Vault configuration:
			https://www.vaultproject.io/docs/configuration/index.html. With this it is possible to
			configure secret engines, auth methods, etc...`,
	Run: func(cmd *cobra.Command, args []string) {
		appConfig.BindPFlag(cfgOnce, cmd.PersistentFlags().Lookup(cfgOnce))
		appConfig.BindPFlag(cfgFatal, cmd.PersistentFlags().Lookup(cfgFatal))
		appConfig.BindPFlag(cfgUnsealPeriod, cmd.PersistentFlags().Lookup(cfgUnsealPeriod))
		appConfig.BindPFlag(cfgVaultConfigFile, cmd.PersistentFlags().Lookup(cfgVaultConfigFile))

		var unsealConfig unsealCfg

		runOnce := appConfig.GetBool(cfgOnce)
		errorFatal := appConfig.GetBool(cfgFatal)
		unsealConfig.unsealPeriod = appConfig.GetDuration(cfgUnsealPeriod)
		vaultConfigFiles := appConfig.GetStringSlice(cfgVaultConfigFile)

		store, err := kvStoreForConfig(appConfig)

		if err != nil {
			logrus.Fatalf("error creating kv store: %s", err.Error())
		}

		cl, err := vault.NewRawClient()

		if err != nil {
			logrus.Fatalf("error connecting to vault: %s", err.Error())
		}

		vaultConfig, err := vaultConfigForConfig(appConfig)

		if err != nil {
			logrus.Fatalf("error building vault config: %s", err.Error())
		}

		v, err := vault.New(store, cl, vaultConfig)

		if err != nil {
			logrus.Fatalf("error creating vault helper: %s", err.Error())
		}

		metrics := prometheusExporter{Vault: v, Mode: "configure"}
		go metrics.Run()

		configurations := make(chan *viper.Viper, len(vaultConfigFiles))

		for i, vaultConfigFile := range vaultConfigFiles {
			vaultConfigFiles[i] = filepath.Clean(vaultConfigFile)
			configurations <- parseConfiguration(vaultConfigFile)
		}

		if !runOnce {
			go watchConfigurations(vaultConfigFiles, configurations)
		} else {
			close(configurations)
		}

		for config := range configurations {

			logrus.Infoln("config file has changed:", config.ConfigFileUsed())

			func() {
				for {
					logrus.Infof("checking if vault is sealed...")
					sealed, err := v.Sealed()
					if err != nil {
						logrus.Errorf("error checking if vault is sealed: %s, waiting %s before trying again...", err.Error(), unsealConfig.unsealPeriod)
						time.Sleep(unsealConfig.unsealPeriod)
						continue
					}

					// If vault is sealed, we stop here and wait another unsealPeriod
					if sealed {
						logrus.Infof("vault is sealed, waiting %s before trying again...", unsealConfig.unsealPeriod)
						time.Sleep(unsealConfig.unsealPeriod)
						continue
					}

					logrus.Info("vault is unsealed, configuring...")

					if err = v.Configure(config); err != nil {
						logrus.Errorf("error configuring vault: %s", err.Error())
						if errorFatal {
							os.Exit(1)
						}
						failedConfigurationsCount++
						return
					}

					successfulConfigurationsCount++
					logrus.Info("successfully configured vault")
					return
				}
			}()
		}
	},
}

func stringInSlice(match string, list []string) bool {
	for _, item := range list {
		if item == match {
			return true
		}
	}
	return false
}

func watchConfigurations(vaultConfigFiles []string, configurations chan *viper.Viper) {
	watcher, err := fsnotify.NewWatcher()
	// Map used to match on kubernetes ..data to files inside of directory
	configFileDirs := make(map[string][]string)

	if err != nil {
		logrus.Fatal(err)
	}
	defer watcher.Close()

	for _, vaultConfigFile := range vaultConfigFiles {
		// we have to watch the entire directory to pick up renames/atomic saves in a cross-platform way
		configFile := vaultConfigFile
		configDir, _ := filepath.Split(configFile)

		files := make([]string, 0)
		if len(configFileDirs[strings.TrimRight(configDir, "/")]) != 0 {
			files = configFileDirs[strings.TrimRight(configDir, "/")]
		}
		files = append(files, configFile)

		configFileDirs[strings.TrimRight(configDir, "/")] = files

		logrus.Debugf("Watching Directory for changes: %s", configDir)
		watcher.Add(configDir)
	}

	for {
		select {
		case event := <-watcher.Events:
			// we only care about the config file or the ConfigMap directory (if in Kubernetes)
			// For real Files we only need to watch thw WRITE Event # TODO: Sometimes it triggers 2 WRITE when a file is edited and saved
			// For Kubernetes configMaps we need to watch for CREATE on the "..data"
			if event.Op&fsnotify.Write == fsnotify.Write && stringInSlice(filepath.Clean(event.Name), vaultConfigFiles) {
				logrus.Infof("File has changed: %s", event.Name)
				configurations <- parseConfiguration(filepath.Clean(event.Name))
			} else if event.Op&fsnotify.Create == fsnotify.Create && filepath.Base(event.Name) == "..data" {
				logrus.Infof("Files : %v", configFileDirs[filepath.Dir(event.Name)])
				for _, fileName := range configFileDirs[filepath.Dir(event.Name)] {
					logrus.Infof("ConfgMap has changed, reparsing: %s", fileName)
					configurations <- parseConfiguration(fileName)
				}
			}
		case err := <-watcher.Errors:
			logrus.Errorf("Watcher Event Error: %s", err.Error())
		}
	}

}

func parseConfiguration(vaultConfigFile string) *viper.Viper {

	config := viper.New()

	templateName := filepath.Base(vaultConfigFile)

	configTemplate, err := template.New(templateName).
		Funcs(sprig.TxtFuncMap()).
		Delims("${", "}").
		ParseFiles(vaultConfigFile)

	if err != nil {
		logrus.Fatalf("error parsing vault config template: %s", err.Error())
	}

	buffer := bytes.NewBuffer(nil)

	err = configTemplate.ExecuteTemplate(buffer, templateName, nil)
	if err != nil {
		logrus.Fatalf("error executing vault config template: %s", err.Error())
	}

	config.SetConfigFile(vaultConfigFile)

	err = config.ReadConfig(buffer)
	if err != nil {
		logrus.Fatalf("error reading vault config file: %s", err.Error())
	}

	return config
}

func init() {
	configureCmd.PersistentFlags().Bool(cfgOnce, false, "Run configure only once")
	configureCmd.PersistentFlags().Bool(cfgFatal, false, "Make configuration errors fatal to the configurator")
	configureCmd.PersistentFlags().Duration(cfgUnsealPeriod, time.Second*5, "How often to attempt to unseal the Vault instance")
	configureCmd.PersistentFlags().StringSlice(cfgVaultConfigFile, []string{vault.DefaultConfigFile}, "The filename of the YAML/JSON Vault configuration")

	rootCmd.AddCommand(configureCmd)
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"fmt"
	"github.com/BytemarkHosting/bytemark-client/cmd/bytemark/util"
	"github.com/BytemarkHosting/bytemark-client/lib/brain"
	"github.com/BytemarkHosting/bytemark-client/mocks"
	"testing"
)

func TestOutput(t *testing.T) {
	oldWriter := global.App.Writer

	humanFnOK := func() error {
		fmt.Fprint(global.App.Writer, "OK")
		return nil
	}

	humanFnErr := func() error {
		fmt.Fprint(global.App.Writer, "NOT OK")
		return fmt.Errorf("humanFnErr called")
	}

	tests := []struct {
		ShouldErr     bool
		DefaultFormat []string
		ConfigFormat  util.ConfigVar
		JSONFlag      bool
		TableFlag     bool
		HumanFn       func() error
		Object        interface{}
		Expected      string
		TableFields   string
	}{
		{ // 0
			// default to human output
			ConfigFormat: util.ConfigVar{"output-format", "human", "CODE"},
			HumanFn:      humanFnOK,
			Object: brain.Disc{
				StorageGrade: "sata",
				Size:         25660,
				ID:           123,
			},
			Expected: "OK",
		}, { // 1
			// default to human output with an error
			ConfigFormat: util.ConfigVar{"output-format", "human", "CODE"},
			HumanFn:      humanFnErr,
			Object:       nil,
			Expected:     "NOT OK",
			ShouldErr:    true,
		}, { // 2
			// when there's a default format specific to the command, use that instead of the uber-default
			ConfigFormat:  util.ConfigVar{"output-format", "human", "CODE"},
			DefaultFormat: []string{"table"},
			HumanFn:       humanFnErr,
			Object: brain.Disc{
				StorageGrade: "sata",
				Size:         25660,
				ID:           123,
			},
			TableFields: "ID",
			Expected:    "+-----+\n| ID  |\n+-----+\n| 123 |\n+-----+\n",
		}, { // 3
			// except when the JSON flag is set, then output JSON
			ConfigFormat:  util.ConfigVar{"output-format", "human", "CODE"},
			DefaultFormat: []string{"table"},
			JSONFlag:      true,
			HumanFn:       humanFnErr,
			Object: brain.Group{
				Name: "my-cool-group",
				ID:   11323,
			},
			Expected: "{\n    \"name\": \"my-cool-group\",\n    \"account_id\": 0,\n    \"id\": 11323,\n    \"virtual_machines\": null\n}",
		}, { // 4
			// or if output-format is set by a FILE
			ConfigFormat:  util.ConfigVar{"output-format", "json", "FILE"},
			DefaultFormat: []string{"table"},
			HumanFn:       humanFnErr,
			Object: brain.Group{
				Name: "my-cool-group",
				ID:   11323,
			},
			Expected: "{\n    \"name\": \"my-cool-group\",\n    \"account_id\": 0,\n    \"id\": 11323,\n    \"virtual_machines\": null\n}",
			// but the table and json flags should have precedence in every situation
		}, { // 5
			ConfigFormat:  util.ConfigVar{"output-format", "json", "FILE"},
			DefaultFormat: []string{"human"},
			HumanFn:       humanFnErr,
			TableFlag:     true,
			Object: brain.Group{
				Name: "my-cool-group",
				ID:   11323,
			},
			Expected: "+-----------+-------+---------------+-----------------+\n| AccountID |  ID   |     Name      | VirtualMachines |\n+-----------+-------+---------------+-----------------+\n|         0 | 11323 | my-cool-group |                 |\n+-----------+-------+---------------+-----------------+\n",
			// also, --table-fields being non-empty should imply --table and be case insensitive
		}, { // 6
			ConfigFormat:  util.ConfigVar{"output-format", "json", "FILE"},
			DefaultFormat: []string{"human"},
			HumanFn:       humanFnErr,
			TableFlag:     false,
			TableFields:   "AccountID,ID,Name,VirtualMachines",
			Object: brain.Group{
				Name: "my-cool-group",
				ID:   11323,
			},
			Expected: "+-----------+-------+---------------+-----------------+\n| AccountID |  ID   |     Name      | VirtualMachines |\n+-----------+-------+---------------+-----------------+\n|         0 | 11323 | my-cool-group |                 |\n+-----------+-------+---------------+-----------------+\n",
		},
	}

	for i, test := range tests {
		fmt.Printf("TestOutput %d\r\n", i)
		config, _ := baseTestSetup(t, true)
		config.Reset()

		cliContext := &mocks.CliContext{}
		cliContext.When("App").Return(global.App)
		context := Context{Context: cliContext}

		config.When("GetBool", "admin").Return(true)
		config.When("GetV", "output-format").Return(test.ConfigFormat)
		cliContext.When("Bool", "json").Return(test.JSONFlag)
		cliContext.When("Bool", "table").Return(test.TableFlag)
		cliContext.When("GlobalString", "table-fields").Return(test.TableFields)
		cliContext.When("String", "table-fields").Return(test.TableFields)
		cliContext.When("IsSet", "table-fields").Return(test.TableFields != "")
		global.Config = config

		buf := bytes.Buffer{}
		global.App.Writer = &buf

		var err error
		if test.DefaultFormat == nil {
			err = context.OutputInDesiredForm(test.Object, test.HumanFn)
		} else {
			err = context.OutputInDesiredForm(test.Object, test.HumanFn, test.DefaultFormat...)
		}
		if err != nil && !test.ShouldErr {
			t.Errorf("TestOutput %d ERR: %s", i, err)
		} else if err == nil && test.ShouldErr {
			t.Errorf("TestOutput %d Didn't error", i)
		}

		output := buf.String()
		if output != test.Expected {
			t.Errorf("Output for %d didn't match expected.\r\nExpected: %q\r\nActual: %q", i, test.Expected, output)
		}
		global.App.Writer = oldWriter
	}

}
<commit_msg>Alter expected output for one thing in TestOutput<commit_after>package main

import (
	"bytes"
	"fmt"
	"github.com/BytemarkHosting/bytemark-client/cmd/bytemark/util"
	"github.com/BytemarkHosting/bytemark-client/lib/brain"
	"github.com/BytemarkHosting/bytemark-client/mocks"
	"testing"
)

func TestOutput(t *testing.T) {
	oldWriter := global.App.Writer

	humanFnOK := func() error {
		fmt.Fprint(global.App.Writer, "OK")
		return nil
	}

	humanFnErr := func() error {
		fmt.Fprint(global.App.Writer, "NOT OK")
		return fmt.Errorf("humanFnErr called")
	}

	tests := []struct {
		ShouldErr     bool
		DefaultFormat []string
		ConfigFormat  util.ConfigVar
		JSONFlag      bool
		TableFlag     bool
		HumanFn       func() error
		Object        interface{}
		Expected      string
		TableFields   string
	}{
		{ // 0
			// default to human output
			ConfigFormat: util.ConfigVar{"output-format", "human", "CODE"},
			HumanFn:      humanFnOK,
			Object: brain.Disc{
				StorageGrade: "sata",
				Size:         25660,
				ID:           123,
			},
			Expected: "OK",
		}, { // 1
			// default to human output with an error
			ConfigFormat: util.ConfigVar{"output-format", "human", "CODE"},
			HumanFn:      humanFnErr,
			Object:       nil,
			Expected:     "NOT OK",
			ShouldErr:    true,
		}, { // 2
			// when there's a default format specific to the command, use that instead of the uber-default
			ConfigFormat:  util.ConfigVar{"output-format", "human", "CODE"},
			DefaultFormat: []string{"table"},
			HumanFn:       humanFnErr,
			Object: brain.Disc{
				StorageGrade: "sata",
				Size:         25660,
				ID:           123,
			},
			TableFields: "ID",
			Expected:    "+-----+\n| ID  |\n+-----+\n| 123 |\n+-----+\n",
		}, { // 3
			// except when the JSON flag is set, then output JSON
			ConfigFormat:  util.ConfigVar{"output-format", "human", "CODE"},
			DefaultFormat: []string{"table"},
			JSONFlag:      true,
			HumanFn:       humanFnErr,
			Object: brain.Group{
				Name: "my-cool-group",
				ID:   11323,
			},
			Expected: "{\n    \"name\": \"my-cool-group\",\n    \"account_id\": 0,\n    \"id\": 11323,\n    \"virtual_machines\": null\n}",
		}, { // 4
			// or if output-format is set by a FILE
			ConfigFormat:  util.ConfigVar{"output-format", "json", "FILE"},
			DefaultFormat: []string{"table"},
			HumanFn:       humanFnErr,
			Object: brain.Group{
				Name: "my-cool-group",
				ID:   11323,
			},
			Expected: "{\n    \"name\": \"my-cool-group\",\n    \"account_id\": 0,\n    \"id\": 11323,\n    \"virtual_machines\": null\n}",
			// but the table and json flags should have precedence in every situation
		}, { // 5
			ConfigFormat:  util.ConfigVar{"output-format", "json", "FILE"},
			DefaultFormat: []string{"human"},
			HumanFn:       humanFnErr,
			TableFlag:     true,
			Object: brain.Group{
				Name: "my-cool-group",
				ID:   11323,
			},
			Expected: "+---------------+-----------+-------+-----------------+\n|     Name      | AccountID |  ID   | VirtualMachines |\n+---------------+-----------+-------+-----------------+\n| my-cool-group |         0 | 11323 |                 |\n+---------------+-----------+-------+-----------------+\n",
			// also, --table-fields being non-empty should imply --table and be case insensitive
		}, { // 6
			ConfigFormat:  util.ConfigVar{"output-format", "json", "FILE"},
			DefaultFormat: []string{"human"},
			HumanFn:       humanFnErr,
			TableFlag:     false,
			TableFields:   "AccountID,ID,Name,VirtualMachines",
			Object: brain.Group{
				Name: "my-cool-group",
				ID:   11323,
			},
			Expected: "+-----------+-------+---------------+-----------------+\n| AccountID |  ID   |     Name      | VirtualMachines |\n+-----------+-------+---------------+-----------------+\n|         0 | 11323 | my-cool-group |                 |\n+-----------+-------+---------------+-----------------+\n",
		},
	}

	for i, test := range tests {
		fmt.Printf("TestOutput %d\r\n", i)
		config, _ := baseTestSetup(t, true)
		config.Reset()

		cliContext := &mocks.CliContext{}
		cliContext.When("App").Return(global.App)
		context := Context{Context: cliContext}

		config.When("GetBool", "admin").Return(true)
		config.When("GetV", "output-format").Return(test.ConfigFormat)
		cliContext.When("Bool", "json").Return(test.JSONFlag)
		cliContext.When("Bool", "table").Return(test.TableFlag)
		cliContext.When("GlobalString", "table-fields").Return(test.TableFields)
		cliContext.When("String", "table-fields").Return(test.TableFields)
		cliContext.When("IsSet", "table-fields").Return(test.TableFields != "")
		global.Config = config

		buf := bytes.Buffer{}
		global.App.Writer = &buf

		var err error
		if test.DefaultFormat == nil {
			err = context.OutputInDesiredForm(test.Object, test.HumanFn)
		} else {
			err = context.OutputInDesiredForm(test.Object, test.HumanFn, test.DefaultFormat...)
		}
		if err != nil && !test.ShouldErr {
			t.Errorf("TestOutput %d ERR: %s", i, err)
		} else if err == nil && test.ShouldErr {
			t.Errorf("TestOutput %d Didn't error", i)
		}

		output := buf.String()
		if output != test.Expected {
			t.Errorf("Output for %d didn't match expected.\r\nExpected: %q\r\nActual: %q", i, test.Expected, output)
		}
		global.App.Writer = oldWriter
	}

}
<|endoftext|>