<commit_before>package ffmpeg

import (
	"io/ioutil"
	"os"
	"testing"
)

func Test_SignDataCreate(t *testing.T) {
	_, dir := setupTest(t)
	defer os.RemoveAll(dir)

	in := &TranscodeOptionsIn{Fname: "../transcoder/test.ts"}
	out := []TranscodeOptions{{
		Oname:        dir + "/signtest1.ts",
		Profile:      P720p60fps16x9,
		AudioEncoder: ComponentOptions{Name: "copy"},
		CalcSign:     true,
	}, {
		Oname:        dir + "/signtest2.ts",
		Profile:      P360p30fps16x9,
		AudioEncoder: ComponentOptions{Name: "copy"},
		CalcSign:     true,
	}, {
		Oname:        dir + "/signtest3.ts",
		Profile:      P360p30fps16x9,
		AudioEncoder: ComponentOptions{Name: "copy"},
	}}
	_, err := Transcode3(in, out)
	if err != nil {
		t.Error(err)
	}
	_, err = os.Stat(dir + "/signtest1.ts.bin")
	if os.IsNotExist(err) {
		t.Error(err)
	}
	_, err = os.Stat(dir + "/signtest2.ts.bin")
	if os.IsNotExist(err) {
		t.Error(err)
	}
	_, err = os.Stat(dir + "/signtest3.ts.bin")
	if os.IsNotExist(err) == false {
		t.Error(err)
	}

}
func Test_SignDataCompare(t *testing.T) {

	res, err := CompareSignatureByPath("../data/sign_sw1.bin", "../data/sign_nv1.bin")
	if err != nil || res != true {
		t.Error(err)
	}
	res, err = CompareSignatureByPath("../data/sign_sw2.bin", "../data/sign_nv2.bin")
	if err != nil || res != true {
		t.Error(err)
	}
	res, err = CompareSignatureByPath("../data/sign_sw1.bin", "../data/sign_sw2.bin")
	if err != nil || res != false {
		t.Error(err)
	}
	res, err = CompareSignatureByPath("../data/sign_nv1.bin", "../data/sign_nv2.bin")
	if err != nil || res != false {
		t.Error(err)
	}
	res, err = CompareSignatureByPath("../data/sign_sw1.bin", "../data/nodata.bin")
	if err == nil || res != false {
		t.Error(err)
	}

	//test CompareSignatureByBuffer function
	data0, err := ioutil.ReadFile("../data/sign_sw1.bin")
	if err != nil {
		t.Error(err)
	}
	data1, err := ioutil.ReadFile("../data/sign_sw2.bin")
	if err != nil {
		t.Error(err)
	}
	data2, err := ioutil.ReadFile("../data/sign_nv1.bin")
	if err != nil {
		t.Error(err)
	}
	res, err = CompareSignatureByBuffer(data0, data2)
	if err != nil || res != true {
		t.Error(err)
	}
	res, err = CompareSignatureByBuffer(data0, data1)
	if err != nil || res != false {
		t.Error(err)
	}
}
<commit_msg>ffmpeg: reproduce filepath issue<commit_after>package ffmpeg

import (
	"io/ioutil"
	"os"
	"path/filepath"
	"testing"
)

func Test_SignDataCreate(t *testing.T) {
	_, dir := setupTest(t)
	defer os.RemoveAll(dir)

	in := &TranscodeOptionsIn{Fname: "../transcoder/test.ts"}
	out := []TranscodeOptions{{
		Oname:        dir + "/signtest1.ts",
		Profile:      P720p60fps16x9,
		AudioEncoder: ComponentOptions{Name: "copy"},
		CalcSign:     true,
	}, {
		Oname:        dir + "/signtest2.ts",
		Profile:      P360p30fps16x9,
		AudioEncoder: ComponentOptions{Name: "copy"},
		CalcSign:     true,
	}, {
		Oname:        dir + "/signtest3.ts",
		Profile:      P360p30fps16x9,
		AudioEncoder: ComponentOptions{Name: "copy"},
	}}
	_, err := Transcode3(in, out)
	if err != nil {
		t.Error(err)
	}
	_, err = os.Stat(dir + "/signtest1.ts.bin")
	if os.IsNotExist(err) {
		t.Error(err)
	}
	_, err = os.Stat(dir + "/signtest2.ts.bin")
	if os.IsNotExist(err) {
		t.Error(err)
	}
	_, err = os.Stat(dir + "/signtest3.ts.bin")
	if os.IsNotExist(err) == false {
		t.Error(err)
	}
}

func Test_SignUnescapedFilepath(t *testing.T) {
	_, dir := setupTest(t)
	defer os.RemoveAll(dir)

	// Test an output file name that contains special chars
	// like ":" and "\" that FFmpeg treats differently in AVOption
	oname, err := filepath.Abs(dir + "/out:720p\\test.ts")
	if err != nil {
		t.Error(err)
	}
	in := &TranscodeOptionsIn{Fname: "../transcoder/test.ts"}
	out := []TranscodeOptions{{
		Oname:        oname,
		Profile:      P720p60fps16x9,
		AudioEncoder: ComponentOptions{Name: "copy"},
		CalcSign:     true,
	}}
	_, err = Transcode3(in, out)

	// Our transcoder module should correctly escape those characters
	// before passing them onto the signature filter
	if err != nil {
		t.Error(err)
	}
	_, err = os.Stat(oname + ".bin")
	if os.IsNotExist(err) {
		t.Error(err)
	}

	// Test same output file, but with a windows style absolute path
	// need to prefix it like /tmp/<dir>/ on unix systems, because without it
	// ffmpeg thinks it's a protocol called "C"
	out = []TranscodeOptions{{
		Oname:        dir + "/" + "C:\\Users\\test\\.lpData\\out720ptest.ts",
		Profile:      P720p60fps16x9,
		AudioEncoder: ComponentOptions{Name: "copy"},
		CalcSign:     true,
	}}
	_, err = Transcode3(in, out)

	if err != nil {
		t.Error(err)
	}
	_, err = os.Stat(oname + ".bin")
	if os.IsNotExist(err) {
		t.Error(err)
	}
}

func Test_SignDataCompare(t *testing.T) {

	res, err := CompareSignatureByPath("../data/sign_sw1.bin", "../data/sign_nv1.bin")
	if err != nil || res != true {
		t.Error(err)
	}
	res, err = CompareSignatureByPath("../data/sign_sw2.bin", "../data/sign_nv2.bin")
	if err != nil || res != true {
		t.Error(err)
	}
	res, err = CompareSignatureByPath("../data/sign_sw1.bin", "../data/sign_sw2.bin")
	if err != nil || res != false {
		t.Error(err)
	}
	res, err = CompareSignatureByPath("../data/sign_nv1.bin", "../data/sign_nv2.bin")
	if err != nil || res != false {
		t.Error(err)
	}
	res, err = CompareSignatureByPath("../data/sign_sw1.bin", "../data/nodata.bin")
	if err == nil || res != false {
		t.Error(err)
	}

	//test CompareSignatureByBuffer function
	data0, err := ioutil.ReadFile("../data/sign_sw1.bin")
	if err != nil {
		t.Error(err)
	}
	data1, err := ioutil.ReadFile("../data/sign_sw2.bin")
	if err != nil {
		t.Error(err)
	}
	data2, err := ioutil.ReadFile("../data/sign_nv1.bin")
	if err != nil {
		t.Error(err)
	}
	res, err = CompareSignatureByBuffer(data0, data2)
	if err != nil || res != true {
		t.Error(err)
	}
	res, err = CompareSignatureByBuffer(data0, data1)
	if err != nil || res != false {
		t.Error(err)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"encoding/json"
	"flag"
	"fmt"
	"github.com/datastream/nsq/nsq"
	"io"
	"io/ioutil"
	"log"
	"os"
	"os/signal"
	"strconv"
	"syscall"
	"time"
)

var (
	conf_file   = flag.String("conf", "config.json", "config file")
	nsq_address = flag.String("nsq_address", "127.0.0.1:4150", "nsq")
)

func main() {
	flag.Parse()
	w := nsq.NewWriter()
	err := w.ConnectToNSQ(*nsq_address)
	if err != nil {
		log.Fatal("fail to connect nsq", err)
	}
	setting, err := ReadConfig(*conf_file)
	if err != nil {
		log.Fatal("fail to read config", err)
	}
	exitchan := make(chan int)
	offset := read_stat(setting)
	for k, v := range setting {
		go read_log(v, offset[k], k, w, exitchan)
	}
	termchan := make(chan os.Signal, 1)
	signal.Notify(termchan, syscall.SIGINT, syscall.SIGTERM)
	<-termchan
	close(exitchan)
}

func read_stat(setting map[string]string) map[string]int64 {
	stat := make(map[string]int64)
	for k, _ := range setting {
		stat_file, err := os.Open(k)
		if err != nil {
			stat[k] = 0
			continue
		}
		s, err := ioutil.ReadAll(stat_file)
		if err != nil {
			stat[k] = 0
			continue
		}
		i, err := strconv.ParseInt(string(s), 10, 64)
		stat[k] = i
	}
	return stat
}

func sync_stat(stat string, value int64) {
	fd, err := os.Create(stat)
	if err != nil {
		log.Println("fail to create ", stat, err)
	}
	defer fd.Close()
	fd.WriteString(fmt.Sprintf("%d", value))
}

func ReadConfig(file string) (map[string]string, error) {
	var setting map[string]string
	config_file, err := os.Open(file)
	config, err := ioutil.ReadAll(config_file)
	if err != nil {
		return nil, err
	}
	defer config_file.Close()
	if err := json.Unmarshal(config, &setting); err != nil {
		return nil, err
	}
	return setting, nil
}

func read_log(file string, offset int64, topic string, w *nsq.Writer, exitchan chan int) {
	fd, err := os.Open(file)
	if err != nil {
		log.Println(err)
		return
	}
	defer fd.Close()
	reader := bufio.NewReader(fd)
	tick := time.Tick(time.Second * 10)
	for {
		line, err := reader.ReadString('\n')
		if err != nil {
			log.Println("retry")
			line, err = reader.ReadString('\n')
		}
		if err == io.EOF {
			log.Println("READ EOF")
			size0, err := fd.Seek(0, 1)
			if err != nil {
				log.Println(er)
				return
			}
			fd, err = os.Open(file)
			if err != nil {
				log.Println("open failed", err)
				return
			}
			_, err = fd.Seek(size0, 0)
			if err != nil {
				log.Println("reopen", file)
				fd.Seek(0, 0)
			}
			reader = bufio.NewReader(fd)
			continue
		}
		if err != nil {
			log.Println(err)
			return
		}
		cmd := nsq.Publish(topic, []byte(line))
		_, _, err = w.Write(cmd)
		if err != nil {
			log.Println("NSQ writer", err)
			w.ConnectToNSQ(*nsq_address)
		}
		select {
		case <-tick:
			size, _ := fd.Seek(0, 1)
			sync_stat(topic, size)
		case <-exitchan:
			size, _ := fd.Seek(0, 1)
			sync_stat(topic, size)
			return
		default:
		}
	}
}
<commit_msg>seek to last offset<commit_after>package main

import (
	"bufio"
	"encoding/json"
	"flag"
	"fmt"
	"github.com/datastream/nsq/nsq"
	"io"
	"io/ioutil"
	"log"
	"os"
	"os/signal"
	"strconv"
	"syscall"
	"time"
)

var (
	conf_file   = flag.String("conf", "config.json", "config file")
	nsq_address = flag.String("nsq_address", "127.0.0.1:4150", "nsq")
)

func main() {
	flag.Parse()
	w := nsq.NewWriter()
	err := w.ConnectToNSQ(*nsq_address)
	if err != nil {
		log.Fatal("fail to connect nsq", err)
	}
	setting, err := ReadConfig(*conf_file)
	if err != nil {
		log.Fatal("fail to read config", err)
	}
	exitchan := make(chan int)
	offset := read_stat(setting)
	for k, v := range setting {
		go read_log(v, offset[k], k, w, exitchan)
	}
	termchan := make(chan os.Signal, 1)
	signal.Notify(termchan, syscall.SIGINT, syscall.SIGTERM)
	<-termchan
	close(exitchan)
}

func read_stat(setting map[string]string) map[string]int64 {
	stat := make(map[string]int64)
	for k, _ := range setting {
		stat_file, err := os.Open(k)
		if err != nil {
			stat[k] = 0
			continue
		}
		s, err := ioutil.ReadAll(stat_file)
		if err != nil {
			stat[k] = 0
			continue
		}
		i, err := strconv.ParseInt(string(s), 10, 64)
		stat[k] = i
	}
	return stat
}

func sync_stat(stat string, value int64) {
	fd, err := os.Create(stat)
	if err != nil {
		log.Println("fail to create ", stat, err)
	}
	defer fd.Close()
	fd.WriteString(fmt.Sprintf("%d", value))
}

func ReadConfig(file string) (map[string]string, error) {
	var setting map[string]string
	config_file, err := os.Open(file)
	config, err := ioutil.ReadAll(config_file)
	if err != nil {
		return nil, err
	}
	defer config_file.Close()
	if err := json.Unmarshal(config, &setting); err != nil {
		return nil, err
	}
	return setting, nil
}

func read_log(file string, offset int64, topic string, w *nsq.Writer, exitchan chan int) {
	fd, err := os.Open(file)
	if err != nil {
		log.Println(err)
		return
	}
	defer fd.Close()
	_, err = fd.Seek(offset, 0)
	if err != nil {
		fd.Seek(0, 0)
	}
	reader := bufio.NewReader(fd)
	tick := time.Tick(time.Second * 10)
	for {
		line, err := reader.ReadString('\n')
		if err != nil {
			log.Println("retry")
			line, err = reader.ReadString('\n')
		}
		if err == io.EOF {
			log.Println("READ EOF")
			size0, err := fd.Seek(0, 1)
			if err != nil {
				log.Println("fail to read current offset", err)
				return
			}
			fd, err = os.Open(file)
			if err != nil {
				log.Println("open failed", err)
				return
			}
			_, err = fd.Seek(size0, 0)
			if err != nil {
				log.Println("reopen", file)
				fd.Seek(0, 0)
			}
			reader = bufio.NewReader(fd)
			continue
		}
		if err != nil {
			log.Println(err)
			return
		}
		cmd := nsq.Publish(topic, []byte(line))
		_, _, err = w.Write(cmd)
		if err != nil {
			log.Println("NSQ writer", err)
			w.ConnectToNSQ(*nsq_address)
		}
		select {
		case <-tick:
			size, _ := fd.Seek(0, 1)
			sync_stat(topic, size)
		case <-exitchan:
			size, _ := fd.Seek(0, 1)
			sync_stat(topic, size)
			return
		default:
		}
	}
}
<|endoftext|>


<commit_before>package qbit

import (
	"errors"
	"fmt"
	"strings"
)

func NewBuilder() *Builder {
	return &Builder{
		prettyMode: true,
		query:      NewQuery(),
		errors:     []error{},
	}
}

// builder is a subset of dialect could be used for common sql queries
// it has all the common functions except multiple statements & table crudders
type Builder struct {
	prettyMode bool
	query      *Query
	errors     []error
}

// function builds the query clauses and returns the sql and bindings
func (b *Builder) Build() (string, []interface{}, error) {

	if b.HasError() {
		errs := []string{}
		for _, err := range b.errors {
			errs = append(errs, err.Error())
		}
		err := errors.New(strings.Join(errs, "\n"))
		b.errors = []error{}
		return "", nil, err
	}

	var delimiter string

	if b.prettyMode {
		delimiter = "\n"
	} else {
		delimiter = " "
	}

	bindings := b.query.Bindings()
	sql := fmt.Sprintf("%s;", strings.Join(b.query.Clauses(), delimiter))

	b.query = NewQuery()

	return sql, bindings, nil
}

func (b *Builder) AddError(err error) {
	b.errors = append(b.errors, err)
}

func (b *Builder) HasError() bool {
	return len(b.errors) > 0
}

func (b *Builder) questionMarks(values ...interface{}) []string {
	questionMarks := make([]string, len(values))
	for k, _ := range values {
		questionMarks[k] = "?"
	}
	return questionMarks
}

// function generates an "insert into %s(%s)" statement
func (b *Builder) Insert(table string, columns ...string) *Builder {
	clause := fmt.Sprintf("INSERT INTO %s(%s)", table, strings.Join(columns, ", "))
	b.query.AddClause(clause)
	return b
}

// function generates "values(%s)" statement and add bindings for each value
func (b *Builder) Values(values ...interface{}) *Builder {
	b.query.AddBinding(values...)
	clause := fmt.Sprintf("VALUES (%s)", strings.Join(b.questionMarks(values...), ", "))
	b.query.AddClause(clause)
	return b
}

// function generates "update %s" statement
func (b *Builder) Update(table string) *Builder {
	clause := fmt.Sprintf("UPDATE %s", table)
	b.query.AddClause(clause)
	return b
}

// function generates "set a = ?" statement for each key a and add bindings for map value
func (b *Builder) Set(m map[string]interface{}) *Builder {
	updates := []string{}
	for k, v := range m {
		updates = append(updates, fmt.Sprintf("%s = ?", k))
		b.query.AddBinding(v)
	}
	clause := fmt.Sprintf("SET %s", strings.Join(updates, ", "))
	b.query.AddClause(clause)
	return b
}

// function generates "select %s" statement
func (b *Builder) Select(columns ...string) *Builder {
	clause := fmt.Sprintf("SELECT %s", strings.Join(columns, ", "))
	b.query.AddClause(clause)
	return b
}

// function generates "from %s" statement for each table name
func (b *Builder) From(tables ...string) *Builder {
	b.query.AddClause(fmt.Sprintf("FROM %s", strings.Join(tables, ", ")))
	return b
}

// function generates "inner join %s on %s" statement for each expression
func (b *Builder) InnerJoin(table string, expressions ...string) *Builder {
	b.query.AddClause(fmt.Sprintf("INNER JOIN %s ON %s", table, strings.Join(expressions, " ")))
	return b
}

// function generates "cross join %s" statement for table
func (b *Builder) CrossJoin(table string) *Builder {
	b.query.AddClause(fmt.Sprintf("CROSS JOIN %s", table))
	return b
}

// function generates "left outer join %s on %s" statement for each expression
func (b *Builder) LeftOuterJoin(table string, expressions ...string) *Builder {
	b.query.AddClause(fmt.Sprintf("LEFT OUTER JOIN %s ON %s", table, strings.Join(expressions, " ")))
	return b
}

// function generates "right outer join %s on %s" statement for each expression
func (b *Builder) RightOuterJoin(table string, expressions ...string) *Builder {
	b.query.AddClause(fmt.Sprintf("RIGHT OUTER JOIN %s ON %s", table, strings.Join(expressions, " ")))
	return b
}

// function generates "full outer join %s on %s" for each expression
func (b *Builder) FullOuterJoin(table string, expressions ...string) *Builder {
	b.query.AddClause(fmt.Sprintf("FULL OUTER JOIN %s ON %s", table, strings.Join(expressions, " ")))
	return b
}

// function generates "where %s" for the expression and adds bindings for each value
func (b *Builder) Where(expression string, bindings ...interface{}) *Builder {
	b.query.AddClause(fmt.Sprintf("WHERE %s", expression))
	b.query.AddBinding(bindings...)
	return b
}

// function generates "order by %s" for each expression
func (b *Builder) OrderBy(expressions ...string) *Builder {
	b.query.AddClause(fmt.Sprintf("ORDER BY %s", strings.Join(expressions, ", ")))
	return b
}

// function generates "group by %s" for each column
func (b *Builder) GroupBy(columns ...string) *Builder {
	b.query.AddClause(fmt.Sprintf("GROUP BY %s", strings.Join(columns, ", ")))
	return b
}

// function generates "having %s" for each expression
func (b *Builder) Having(expressions ...string) *Builder {
	b.query.AddClause(fmt.Sprintf("HAVING %s", strings.Join(expressions, ", ")))
	return b
}

// function generates limit %d offset %d for offset and count
func (b *Builder) Limit(offset int, count int) *Builder {
	b.query.AddClause(fmt.Sprintf("LIMIT %d OFFSET %d", count, offset))
	return b
}

// aggregates

// function generates "avg(%s)" statement for column
func (b *Builder) Avg(column string) string {
	return fmt.Sprintf("AVG(%s)", column)
}

// function generates "count(%s)" statement for column
func (b *Builder) Count(column string) string {
	return fmt.Sprintf("COUNT(%s)", column)
}

// function generates "sum(%s)" statement for column
func (b *Builder) Sum(column string) string {
	return fmt.Sprintf("SUM(%s)", column)
}

// function generates "min(%s)" statement for column
func (b *Builder) Min(column string) string {
	return fmt.Sprintf("MIN(%s)", column)
}

// function generates "max(%s)" statement for column
func (b *Builder) Max(column string) string {
	return fmt.Sprintf("MAX(%s)", column)
}

// expressions

// function generates "%s not in (%s)" for key and adds bindings for each value
func (b *Builder) NotIn(key string, values ...interface{}) string {
	b.query.AddBinding(values...)
	return fmt.Sprintf("%s NOT IN (%s)", key, strings.Join(b.questionMarks(values...), ","))
}

// function generates "%s in (%s)" for key and adds bindings for each value
func (b *Builder) In(key string, values ...interface{}) string {
	b.query.AddBinding(values...)
	return fmt.Sprintf("%s IN (%s)", key, strings.Join(b.questionMarks(values...), ","))
}

// function generates "%s != ?" for key and adds binding for value
func (b *Builder) NotEq(key string, value interface{}) string {
	b.query.AddBinding(value)
	return fmt.Sprintf("%s != ?", key)
}

// function generates "%s = ?" for key and adds binding for value
func (b *Builder) Eq(key string, value interface{}) string {
	b.query.AddBinding(value)
	return fmt.Sprintf("%s = ?", key)
}

// function generates "%s > ?" for key and adds binding for value
func (b *Builder) Gt(key string, value interface{}) string {
	b.query.AddBinding(value)
	return fmt.Sprintf("%s > ?", key)
}

// function generates "%s >= ?" for key and adds binding for value
func (b *Builder) Gte(key string, value interface{}) string {
	b.query.AddBinding(value)
	return fmt.Sprintf("%s >= ?", key)
}

// function generates "%s < ?" for key and adds binding for value
func (b *Builder) St(key string, value interface{}) string {
	b.query.AddBinding(value)
	return fmt.Sprintf("%s < ?", key)
}

// function generates "%s <= ?" for key and adds binding for value
func (b *Builder) Ste(key string, value interface{}) string {
	b.query.AddBinding(value)
	return fmt.Sprintf("%s <= ?", key)
}

// function generates " AND " between any number of expressions
func (b *Builder) And(expressions ...string) string {
	return fmt.Sprintf("(%s)", strings.Join(expressions, " AND "))
}

// function generates " OR " between any number of expressions
func (b *Builder) Or(expressions ...string) string {
	return strings.Join(expressions, " OR ")
}

// function generates generic CREATE TABLE statement
func (b *Builder) CreateTable(table string, fields []string, constraints []string) *Builder {

	b.query.AddClause(fmt.Sprintf("CREATE TABLE %s(", table))

	for k, f := range fields {
		clause := fmt.Sprintf("\t%s", f)
		if len(fields)-1 > k {
			clause += ","
		}
		b.query.AddClause(clause)
	}

	for _, c := range constraints {
		b.query.AddClause(fmt.Sprintf("%s, ", c))
	}

	b.query.AddClause(")")
	return b
}

// function generates generic ALTER TABLE statement
func (b *Builder) AlterTable(table string) *Builder {

	b.query.AddClause(fmt.Sprintf("ALTER TABLE %s", table))
	return b
}

// function generates generic DROP TABLE statement
func (b *Builder) DropTable(table string) *Builder {

	b.query.AddClause(fmt.Sprintf("DROP TABLE %s", table))
	return b
}

// function generates generic ADD COLUMN statement
func (b *Builder) Add(colName string, colType string) *Builder {

	b.query.AddClause(fmt.Sprintf("ADD %s %s", colName, colType))
	return b
}

// function generates generic DROP COLUMN statement
func (b *Builder) Drop(colName string) *Builder {

	b.query.AddClause(fmt.Sprintf("DROP %s", colName))
	return b
}
<commit_msg>add Delete() function to query builder<commit_after>package qbit

import (
	"errors"
	"fmt"
	"strings"
)

func NewBuilder() *Builder {
	return &Builder{
		prettyMode: true,
		query:      NewQuery(),
		errors:     []error{},
	}
}

// builder is a subset of dialect could be used for common sql queries
// it has all the common functions except multiple statements & table crudders
type Builder struct {
	prettyMode bool
	query      *Query
	errors     []error
}

// function builds the query clauses and returns the sql and bindings
func (b *Builder) Build() (string, []interface{}, error) {

	if b.HasError() {
		errs := []string{}
		for _, err := range b.errors {
			errs = append(errs, err.Error())
		}
		err := errors.New(strings.Join(errs, "\n"))
		b.errors = []error{}
		return "", nil, err
	}

	var delimiter string

	if b.prettyMode {
		delimiter = "\n"
	} else {
		delimiter = " "
	}

	bindings := b.query.Bindings()
	sql := fmt.Sprintf("%s;", strings.Join(b.query.Clauses(), delimiter))

	b.query = NewQuery()

	return sql, bindings, nil
}

func (b *Builder) AddError(err error) {
	b.errors = append(b.errors, err)
}

func (b *Builder) HasError() bool {
	return len(b.errors) > 0
}

func (b *Builder) questionMarks(values ...interface{}) []string {
	questionMarks := make([]string, len(values))
	for k, _ := range values {
		questionMarks[k] = "?"
	}
	return questionMarks
}

// function generates an "insert into %s(%s)" statement
func (b *Builder) Insert(table string, columns ...string) *Builder {
	clause := fmt.Sprintf("INSERT INTO %s(%s)", table, strings.Join(columns, ", "))
	b.query.AddClause(clause)
	return b
}

// function generates "values(%s)" statement and add bindings for each value
func (b *Builder) Values(values ...interface{}) *Builder {
	b.query.AddBinding(values...)
	clause := fmt.Sprintf("VALUES (%s)", strings.Join(b.questionMarks(values...), ", "))
	b.query.AddClause(clause)
	return b
}

// function generates "update %s" statement
func (b *Builder) Update(table string) *Builder {
	clause := fmt.Sprintf("UPDATE %s", table)
	b.query.AddClause(clause)
	return b
}

// function generates "set a = ?" statement for each key a and add bindings for map value
func (b *Builder) Set(m map[string]interface{}) *Builder {
	updates := []string{}
	for k, v := range m {
		updates = append(updates, fmt.Sprintf("%s = ?", k))
		b.query.AddBinding(v)
	}
	clause := fmt.Sprintf("SET %s", strings.Join(updates, ", "))
	b.query.AddClause(clause)
	return b
}

// function generates "delete" statement
func (b *Builder) Delete(table string) *Builder {
	b.query.AddClause(fmt.Sprintf("DELETE FROM %s", table))
	return b
}

// function generates "select %s" statement
func (b *Builder) Select(columns ...string) *Builder {
	clause := fmt.Sprintf("SELECT %s", strings.Join(columns, ", "))
	b.query.AddClause(clause)
	return b
}

// function generates "from %s" statement for each table name
func (b *Builder) From(tables ...string) *Builder {
	b.query.AddClause(fmt.Sprintf("FROM %s", strings.Join(tables, ", ")))
	return b
}

// function generates "inner join %s on %s" statement for each expression
func (b *Builder) InnerJoin(table string, expressions ...string) *Builder {
	b.query.AddClause(fmt.Sprintf("INNER JOIN %s ON %s", table, strings.Join(expressions, " ")))
	return b
}

// function generates "cross join %s" statement for table
func (b *Builder) CrossJoin(table string) *Builder {
	b.query.AddClause(fmt.Sprintf("CROSS JOIN %s", table))
	return b
}

// function generates "left outer join %s on %s" statement for each expression
func (b *Builder) LeftOuterJoin(table string, expressions ...string) *Builder {
	b.query.AddClause(fmt.Sprintf("LEFT OUTER JOIN %s ON %s", table, strings.Join(expressions, " ")))
	return b
}

// function generates "right outer join %s on %s" statement for each expression
func (b *Builder) RightOuterJoin(table string, expressions ...string) *Builder {
	b.query.AddClause(fmt.Sprintf("RIGHT OUTER JOIN %s ON %s", table, strings.Join(expressions, " ")))
	return b
}

// function generates "full outer join %s on %s" for each expression
func (b *Builder) FullOuterJoin(table string, expressions ...string) *Builder {
	b.query.AddClause(fmt.Sprintf("FULL OUTER JOIN %s ON %s", table, strings.Join(expressions, " ")))
	return b
}

// function generates "where %s" for the expression and adds bindings for each value
func (b *Builder) Where(expression string, bindings ...interface{}) *Builder {
	b.query.AddClause(fmt.Sprintf("WHERE %s", expression))
	b.query.AddBinding(bindings...)
	return b
}

// function generates "order by %s" for each expression
func (b *Builder) OrderBy(expressions ...string) *Builder {
	b.query.AddClause(fmt.Sprintf("ORDER BY %s", strings.Join(expressions, ", ")))
	return b
}

// function generates "group by %s" for each column
func (b *Builder) GroupBy(columns ...string) *Builder {
	b.query.AddClause(fmt.Sprintf("GROUP BY %s", strings.Join(columns, ", ")))
	return b
}

// function generates "having %s" for each expression
func (b *Builder) Having(expressions ...string) *Builder {
	b.query.AddClause(fmt.Sprintf("HAVING %s", strings.Join(expressions, ", ")))
	return b
}

// function generates limit %d offset %d for offset and count
func (b *Builder) Limit(offset int, count int) *Builder {
	b.query.AddClause(fmt.Sprintf("LIMIT %d OFFSET %d", count, offset))
	return b
}

// aggregates

// function generates "avg(%s)" statement for column
func (b *Builder) Avg(column string) string {
	return fmt.Sprintf("AVG(%s)", column)
}

// function generates "count(%s)" statement for column
func (b *Builder) Count(column string) string {
	return fmt.Sprintf("COUNT(%s)", column)
}

// function generates "sum(%s)" statement for column
func (b *Builder) Sum(column string) string {
	return fmt.Sprintf("SUM(%s)", column)
}

// function generates "min(%s)" statement for column
func (b *Builder) Min(column string) string {
	return fmt.Sprintf("MIN(%s)", column)
}

// function generates "max(%s)" statement for column
func (b *Builder) Max(column string) string {
	return fmt.Sprintf("MAX(%s)", column)
}

// expressions

// function generates "%s not in (%s)" for key and adds bindings for each value
func (b *Builder) NotIn(key string, values ...interface{}) string {
	b.query.AddBinding(values...)
	return fmt.Sprintf("%s NOT IN (%s)", key, strings.Join(b.questionMarks(values...), ","))
}

// function generates "%s in (%s)" for key and adds bindings for each value
func (b *Builder) In(key string, values ...interface{}) string {
	b.query.AddBinding(values...)
	return fmt.Sprintf("%s IN (%s)", key, strings.Join(b.questionMarks(values...), ","))
}

// function generates "%s != ?" for key and adds binding for value
func (b *Builder) NotEq(key string, value interface{}) string {
	b.query.AddBinding(value)
	return fmt.Sprintf("%s != ?", key)
}

// function generates "%s = ?" for key and adds binding for value
func (b *Builder) Eq(key string, value interface{}) string {
	b.query.AddBinding(value)
	return fmt.Sprintf("%s = ?", key)
}

// function generates "%s > ?" for key and adds binding for value
func (b *Builder) Gt(key string, value interface{}) string {
	b.query.AddBinding(value)
	return fmt.Sprintf("%s > ?", key)
}

// function generates "%s >= ?" for key and adds binding for value
func (b *Builder) Gte(key string, value interface{}) string {
	b.query.AddBinding(value)
	return fmt.Sprintf("%s >= ?", key)
}

// function generates "%s < ?" for key and adds binding for value
func (b *Builder) St(key string, value interface{}) string {
	b.query.AddBinding(value)
	return fmt.Sprintf("%s < ?", key)
}

// function generates "%s <= ?" for key and adds binding for value
func (b *Builder) Ste(key string, value interface{}) string {
	b.query.AddBinding(value)
	return fmt.Sprintf("%s <= ?", key)
}

// function generates " AND " between any number of expressions
func (b *Builder) And(expressions ...string) string {
	return fmt.Sprintf("(%s)", strings.Join(expressions, " AND "))
}

// function generates " OR " between any number of expressions
func (b *Builder) Or(expressions ...string) string {
	return strings.Join(expressions, " OR ")
}

// function generates generic CREATE TABLE statement
func (b *Builder) CreateTable(table string, fields []string, constraints []string) *Builder {

	b.query.AddClause(fmt.Sprintf("CREATE TABLE %s(", table))

	for k, f := range fields {
		clause := fmt.Sprintf("\t%s", f)
		if len(fields)-1 > k {
			clause += ","
		}
		b.query.AddClause(clause)
	}

	for _, c := range constraints {
		b.query.AddClause(fmt.Sprintf("%s, ", c))
	}

	b.query.AddClause(")")
	return b
}

// function generates generic ALTER TABLE statement
func (b *Builder) AlterTable(table string) *Builder {

	b.query.AddClause(fmt.Sprintf("ALTER TABLE %s", table))
	return b
}

// function generates generic DROP TABLE statement
func (b *Builder) DropTable(table string) *Builder {

	b.query.AddClause(fmt.Sprintf("DROP TABLE %s", table))
	return b
}

// function generates generic ADD COLUMN statement
func (b *Builder) Add(colName string, colType string) *Builder {

	b.query.AddClause(fmt.Sprintf("ADD %s %s", colName, colType))
	return b
}

// function generates generic DROP COLUMN statement
func (b *Builder) Drop(colName string) *Builder {

	b.query.AddClause(fmt.Sprintf("DROP %s", colName))
	return b
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"io"
	"time"

	"github.com/cloudfoundry-incubator/veritas/config_finder"
	"github.com/cloudfoundry-incubator/veritas/fetch_store"
	"github.com/cloudfoundry-incubator/veritas/print_store"
	"github.com/cloudfoundry-incubator/veritas/say"
)

func DumpStoreCommand() Command {
	var (
		etcdClusterFlag string
		tasks           bool
		lrps            bool
		services        bool
		verbose         bool
		rate            time.Duration
	)

	flagSet := flag.NewFlagSet("dump-store", flag.ExitOnError)
	flagSet.StringVar(&etcdClusterFlag, "etcdCluster", "", "comma-separated etcd cluster urls")
	flagSet.BoolVar(&verbose, "v", false, "be verbose")
	flagSet.BoolVar(&tasks, "tasks", true, "print tasks")
	flagSet.BoolVar(&lrps, "lrps", true, "print lrps")
	flagSet.BoolVar(&services, "services", true, "print services")
	flagSet.DurationVar(&rate, "rate", time.Duration(0), "rate at which to poll the store")

	return Command{
		Name:        "dump-store",
		Description: "- Fetch and print contents of the BBS",
		FlagSet:     flagSet,
		Run: func(args []string) {
			etcdCluster, err := config_finder.FindETCDCluster(etcdClusterFlag)
			ExitIfError("Could not find etcd cluster", err)

			if rate == 0 {
				err = dump(etcdCluster, verbose, tasks, lrps, services)
				ExitIfError("Failed to dump", err)
				return
			}

			ticker := time.NewTicker(rate)
			for {
				<-ticker.C
				say.Clear()
				err = dump(etcdCluster, verbose, tasks, lrps, services)
				if err != nil {
					say.Println(0, say.Red("Failed to dump: %s", err.Error()))
				}
			}
		},
	}
}

func dump(etcdCluster []string, verbose bool, tasks bool, lrps bool, services bool) error {
	reader, writer := io.Pipe()

	errs := make(chan error)
	go func() {
		err := fetch_store.Fetch(etcdCluster, false, writer)
		errs <- err
	}()
	go func() {
		err := print_store.PrintStore(verbose, tasks, lrps, services, reader)
		errs <- err
	}()

	err1 := <-errs
	err2 := <-errs
	if err1 != nil {
		return err1
	}
	if err2 != nil {
		return err2
	}
	return nil
}
<commit_msg>fix dump<commit_after>package main

import (
	"flag"
	"io"
	"time"

	"github.com/cloudfoundry-incubator/veritas/config_finder"
	"github.com/cloudfoundry-incubator/veritas/fetch_store"
	"github.com/cloudfoundry-incubator/veritas/print_store"
	"github.com/cloudfoundry-incubator/veritas/say"
)

func DumpStoreCommand() Command {
	var (
		etcdClusterFlag string
		tasks           bool
		lrps            bool
		services        bool
		verbose         bool
		rate            time.Duration
	)

	flagSet := flag.NewFlagSet("dump-store", flag.ExitOnError)
	flagSet.StringVar(&etcdClusterFlag, "etcdCluster", "", "comma-separated etcd cluster urls")
	flagSet.BoolVar(&verbose, "v", false, "be verbose")
	flagSet.BoolVar(&tasks, "tasks", true, "print tasks")
	flagSet.BoolVar(&lrps, "lrps", true, "print lrps")
	flagSet.BoolVar(&services, "services", true, "print services")
	flagSet.DurationVar(&rate, "rate", time.Duration(0), "rate at which to poll the store")

	return Command{
		Name:        "dump-store",
		Description: "- Fetch and print contents of the BBS",
		FlagSet:     flagSet,
		Run: func(args []string) {
			etcdCluster, err := config_finder.FindETCDCluster(etcdClusterFlag)
			ExitIfError("Could not find etcd cluster", err)

			if rate == 0 {
				err = dump(etcdCluster, verbose, tasks, lrps, services)
				ExitIfError("Failed to dump", err)
				return
			}

			ticker := time.NewTicker(rate)
			for {
				<-ticker.C
				say.Clear()
				err = dump(etcdCluster, verbose, tasks, lrps, services)
				if err != nil {
					say.Println(0, say.Red("Failed to dump: %s", err.Error()))
				}
			}
		},
	}
}

func dump(etcdCluster []string, verbose bool, tasks bool, lrps bool, services bool) error {
	reader, writer := io.Pipe()

	errs := make(chan error)
	go func() {
		err := fetch_store.Fetch(etcdCluster, false, writer)
		errs <- err
	}()
	go func() {
		err := print_store.PrintStore(verbose, tasks, lrps, services, reader)
		errs <- err
	}()

	err1 := <-errs
	if err1 != nil {
		return err1
	}
	err2 := <-errs
	if err2 != nil {
		return err2
	}
	return nil
}
<|endoftext|>


<commit_before>package sshvault

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"testing"

	"github.com/ssh-vault/crypto"
	"github.com/ssh-vault/crypto/aead"
)

// These are done in one function to avoid declaring global variables in a test
// file.
func TestVaultFunctionsFingerprint(t *testing.T) {
	dir, err := ioutil.TempDir("", "vault")
	if err != nil {
		t.Error(err)
	}
	defer os.RemoveAll(dir) // clean up

	tmpfile := filepath.Join(dir, "vault")

	vault, err := New("55:cd:f2:7e:4c:0b:e5:a7:6e:6c:fc:6b:8e:58:9d:15", "test_data/id_rsa.pub", "", "create", tmpfile)
	if err != nil {
		t.Error(err)
	}

	PKCS8, err := vault.PKCS8()
	if err != nil {
		t.Error(err)
	}

	vault.PublicKey, err = vault.GetRSAPublicKey(PKCS8)
	if err != nil {
		t.Error(err)
	}

	vault.Fingerprint, err = vault.GenFingerprint(PKCS8)
	if err != nil {
		t.Error(err)
	}

	if vault.Password, err = crypto.GenerateNonce(32); err != nil {
		t.Error(err)
	}

	// Skip vault.Create because we don't need/want to interact with an editor
	// for tests.
	in := []byte("The quick brown fox jumps over the lazy dog")

	out, err := aead.Encrypt(vault.Password, in, []byte(vault.Fingerprint))
	if err != nil {
		t.Error(err)
	}

	if err = vault.Close(out); err != nil {
		t.Error(err)
	}

	enc1, err := ioutil.ReadFile(tmpfile)
	if err != nil {
		t.Error(err)
	}

	plaintext, err := vault.View()
	if err != nil {
		t.Error(err)
	}

	if !bytes.Equal(in, plaintext) {
		t.Error("in != out")
	}

	os.Setenv("EDITOR", "cat")
	edited, err := vault.Edit(plaintext)
	if err != nil {
		t.Error(err)
	}

	out, err = aead.Encrypt(vault.Password, edited, []byte(vault.Fingerprint))
	if err != nil {
		t.Error(err)
	}

	if err = vault.Close(out); err != nil {
		t.Error(err)
	}

	plaintext, err = vault.View()
	if err != nil {
		t.Error(err)
	}

	enc2, err := ioutil.ReadFile(tmpfile)
	if err != nil {
		t.Error(err)
	}

	if !bytes.Equal(edited, plaintext) {
		t.Error("edited != plaintext ")
	}

	if bytes.Equal(enc1, enc2) {
		t.Error("Expecting different encrypted outputs")
	}
}

func TestVaultFunctionsSTDOUTFingerprint(t *testing.T) {
	dir, err := ioutil.TempDir("", "vault")
	if err != nil {
		t.Error(err)
	}
	defer os.RemoveAll(dir) // clean up

	var testTable = []struct {
		Fingerprint string
		KeyPath     string
		Option      string
	}{
		{"55:cd:f2:7e:4c:0b:e5:a7:6e:6c:fc:6b:8e:58:9d:15", "test_data/id_rsa.pub", "create"},
		{"55:cd:f2:7e:4c:0b:e5:a7:6e:6c:fc:6b:8e:58:9d:15", "test_data/id_rsa_extra_linebreak", "create"},
	}

	for _, tt := range testTable {
		vault, err := New(tt.Fingerprint, tt.KeyPath, "", tt.Option, "")
		if err != nil {
			t.Error(err)
		}

		PKCS8, err := vault.PKCS8()
		if err != nil {
			t.Error(err)
		}

		vault.PublicKey, err = vault.GetRSAPublicKey(PKCS8)
		if err != nil {
			t.Error(err)
		}

		vault.Fingerprint, err = vault.GenFingerprint(PKCS8)
		if err != nil {
			t.Error(err)
		}

		if vault.Password, err = crypto.GenerateNonce(32); err != nil {
			t.Error(err)
		}

		// Skip vault.Create because we don't need/want to interact with an editor
		in := []byte("The quick brown fox jumps over the lazy dog")

		out, err := aead.Encrypt(vault.Password, in, []byte(vault.Fingerprint))
		if err != nil {
			t.Error(err)
		}

		rescueStdout := os.Stdout // keep backup of the real stdout
		r, w, _ := os.Pipe()
		os.Stdout = w

		if err = vault.Close(out); err != nil {
			t.Error(err)
		}

		w.Close()
		outStdout, _ := ioutil.ReadAll(r)
		os.Stdout = rescueStdout
		tmpfile, err := ioutil.TempFile("", "stdout")
		if err != nil {
			t.Error(err)
		}
		tmpfile.Write([]byte(outStdout))
		vault.vault = tmpfile.Name()

		plaintext, err := vault.View()
		if err != nil {
			t.Error(err)
		}

		if !bytes.Equal(in, plaintext) {
			t.Error("in != out")
		}
	}
}

func TestVaultNewFingerprint(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		expect(t, "ssh-vault", r.Header.Get("User-agent"))
		fmt.Fprintln(w, "ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAACAQDjjM4JEyg1T8j5YICtqslLNp2UGg80CppTM3ZYu73pEmDhMwbLfdhuI56AQZgWViFsF/7QHDJPcRY2Piu38b4kizTSM0QHEOC7CTo+vnzxptlKLGT1y2mcY1P9VXzCBMSWQN9/vGasgl/sUp1zcTvVT0CjjA6k1dJM6/+aDVtCsFa851VkwbeIsWl5BAHLyL+ur5BX93/BxYnRcYl7ooheuEWWokyWJ0IwEFToPMHAthTbDn1P17wYF43oscTORsFBfkP1JLBKHPDPJCGcBgQButL/srLJf6o44fScAYL99s1dQ/Qqv31aygDmwLdKEDldNnWEaJZ+iidEiIlPtAnLYGnVVA4u+NA2p3egrUrLWmpPjMX6XSb2VRHllzCcY4vZ4F2ud2TFaYG6N+9+vRCdxB+LFcHhm7ottI4vnC5P1bbMagjmFne0+TSKrAfMCw59eiQd8yZVMoE2yPXjFOQt6EOBvB4OHv1AaVt2q0PGqSkv5vIhgsKJWx/6IUj0Kz24hDiMipFb0jL3xstvizAllpC6yF26Ju/nwF03eJJGGxJjrxYd4P5/rY6SWY3yakiUN7pUBgUK2Ok3K3/+BTy5Aag8OXcvOZJumr2X2Wn9DweQeCRjC8UqFDKALqA/3vopZ2S59V4WOg3sV94hEig/KHLISNge1Uatn+qosK2sPw==")
	}))
	defer ts.Close()
	_, err := New("55:cd:f2:7e:4c:0b:e5:a7:6e:6c:fc:6b:8e:58:9d:15", "", ts.URL, "view", "")
	if err != nil {
		t.Error(err)
	}
}
func TestVaultNewFingerprintBadKey(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		expect(t, "ssh-vault", r.Header.Get("User-agent"))
		fmt.Fprintln(w, "ssh-rsa FOO")
	}))
	defer ts.Close()
	_, err := New("55:cd:f2:7e:4c:0b:e5:a7:6e:6c:fc:6b:8e:58:9d:15", "", ts.URL, "view", "")
	if err == nil {
		t.Error("Expecting error: Use a public ssh key: illegal base64 ...")
	}
}

func TestVaultNewNoKeyFingerprint(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		expect(t, "ssh-vault", r.Header.Get("User-agent"))
	}))
	defer ts.Close()
	_, err := New("55:cd:f2:7e:4c:0b:e5:a7:6e:6c:fc:6b:8e:58:9d:XX", "", ts.URL, "view", "")
	if err == nil {
		t.Error("Expecting error")
	}
}
<commit_msg>test travis<commit_after>package sshvault

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"testing"

	"github.com/ssh-vault/crypto"
	"github.com/ssh-vault/crypto/aead"
)

// These are done in one function to avoid declaring global variables in a test
// file.
func TestVaultFunctionsFingerprint(t *testing.T) {
	dir, err := ioutil.TempDir("", "vault")
	if err != nil {
		t.Error(err)
	}
	defer os.RemoveAll(dir) // clean up

	tmpfile := filepath.Join(dir, "vault")

	vault, err := New("55:cd:f2:7e:4c:0b:e5:a7:6e:6c:fc:6b:8e:58:9d:15", "test_data/id_rsa.pub", "", "create", tmpfile)
	if err != nil {
		t.Error(err)
	}

	PKCS8, err := vault.PKCS8()
	if err != nil {
		t.Error(err)
	}

	vault.PublicKey, err = vault.GetRSAPublicKey(PKCS8)
	if err != nil {
		t.Error(err)
	}

	vault.Fingerprint, err = vault.GenFingerprint(PKCS8)
	if err != nil {
		t.Error(err)
	}

	if vault.Password, err = crypto.GenerateNonce(32); err != nil {
		t.Error(err)
	}

	// Skip vault.Create because we don't need/want to interact with an editor
	// for tests.
	in := []byte("The quick brown fox jumps over the lazy dog")

	out, err := aead.Encrypt(vault.Password, in, []byte(vault.Fingerprint))
	if err != nil {
		t.Error(err)
	}

	if err = vault.Close(out); err != nil {
		t.Error(err)
	}

	enc1, err := ioutil.ReadFile(tmpfile)
	if err != nil {
		t.Error(err)
	}

	plaintext, err := vault.View()
	if err != nil {
		t.Error(err)
	}

	if !bytes.Equal(in, plaintext) {
		t.Error("in != out")
	}

	os.Setenv("EDITOR", "cat")
	edited, err := vault.Edit(plaintext)
	if err != nil {
		t.Error(err)
	}

	out, err = aead.Encrypt(vault.Password, edited, []byte(vault.Fingerprint))
	if err != nil {
		t.Error(err)
	}

	if err = vault.Close(out); err != nil {
		t.Error(err)
	}

	plaintext, err = vault.View()
	if err != nil {
		t.Error(err)
	}

	enc2, err := ioutil.ReadFile(tmpfile)
	if err != nil {
		t.Error(err)
	}

	if !bytes.Equal(edited, plaintext) {
		t.Error("edited != plaintext ")
	}

	if bytes.Equal(enc1, enc2) {
		t.Error("Expecting different encrypted outputs")
	}
}

func TestVaultFunctionsSTDOUTFingerprint(t *testing.T) {
	dir, err := ioutil.TempDir("", "vault")
	if err != nil {
		t.Error(err)
	}
	defer os.RemoveAll(dir) // clean up

	var testTable = []struct {
		Fingerprint string
		KeyPath     string
		Option      string
	}{
		{"55:cd:f2:7e:4c:0b:e5:a7:6e:6c:fc:6b:8e:58:9d:15", "test_data/id_rsa.pub", "create"},
		{"55:cd:f2:7e:4c:0b:e5:a7:6e:6c:fc:6b:8e:58:9d:15", "test_data/id_rsa_extra_linebreak", "create"},
	}

	for _, tt := range testTable {
		vault, err := New(tt.Fingerprint, tt.KeyPath, "", tt.Option, "")
		if err != nil {
			t.Error(err)
		}

		PKCS8, err := vault.PKCS8()
		if err != nil {
			t.Error(err)
		}
		t.Logf("%x", PKCS8)

		vault.PublicKey, err = vault.GetRSAPublicKey(PKCS8)
		if err != nil {
			t.Error(err)
		}

		vault.Fingerprint, err = vault.GenFingerprint(PKCS8)
		if err != nil {
			t.Error(err)
		}

		if vault.Password, err = crypto.GenerateNonce(32); err != nil {
			t.Error(err)
		}

		// Skip vault.Create because we don't need/want to interact with an editor
		in := []byte("The quick brown fox jumps over the lazy dog")

		out, err := aead.Encrypt(vault.Password, in, []byte(vault.Fingerprint))
		if err != nil {
			t.Error(err)
		}

		rescueStdout := os.Stdout // keep backup of the real stdout
		r, w, _ := os.Pipe()
		os.Stdout = w

		if err = vault.Close(out); err != nil {
			t.Error(err)
		}

		w.Close()
		outStdout, _ := ioutil.ReadAll(r)
		os.Stdout = rescueStdout
		tmpfile, err := ioutil.TempFile("", "stdout")
		if err != nil {
			t.Error(err)
		}
		tmpfile.Write([]byte(outStdout))
		vault.vault = tmpfile.Name()

		plaintext, err := vault.View()
		if err != nil {
			t.Error(err)
		}

		if !bytes.Equal(in, plaintext) {
			t.Error("in != out")
		}
	}
}

func TestVaultNewFingerprint(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		expect(t, "ssh-vault", r.Header.Get("User-agent"))
		fmt.Fprintln(w, "ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAACAQDjjM4JEyg1T8j5YICtqslLNp2UGg80CppTM3ZYu73pEmDhMwbLfdhuI56AQZgWViFsF/7QHDJPcRY2Piu38b4kizTSM0QHEOC7CTo+vnzxptlKLGT1y2mcY1P9VXzCBMSWQN9/vGasgl/sUp1zcTvVT0CjjA6k1dJM6/+aDVtCsFa851VkwbeIsWl5BAHLyL+ur5BX93/BxYnRcYl7ooheuEWWokyWJ0IwEFToPMHAthTbDn1P17wYF43oscTORsFBfkP1JLBKHPDPJCGcBgQButL/srLJf6o44fScAYL99s1dQ/Qqv31aygDmwLdKEDldNnWEaJZ+iidEiIlPtAnLYGnVVA4u+NA2p3egrUrLWmpPjMX6XSb2VRHllzCcY4vZ4F2ud2TFaYG6N+9+vRCdxB+LFcHhm7ottI4vnC5P1bbMagjmFne0+TSKrAfMCw59eiQd8yZVMoE2yPXjFOQt6EOBvB4OHv1AaVt2q0PGqSkv5vIhgsKJWx/6IUj0Kz24hDiMipFb0jL3xstvizAllpC6yF26Ju/nwF03eJJGGxJjrxYd4P5/rY6SWY3yakiUN7pUBgUK2Ok3K3/+BTy5Aag8OXcvOZJumr2X2Wn9DweQeCRjC8UqFDKALqA/3vopZ2S59V4WOg3sV94hEig/KHLISNge1Uatn+qosK2sPw==")
	}))
	defer ts.Close()
	_, err := New("55:cd:f2:7e:4c:0b:e5:a7:6e:6c:fc:6b:8e:58:9d:15", "", ts.URL, "view", "")
	if err != nil {
		t.Error(err)
	}
}
func TestVaultNewFingerprintBadKey(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		expect(t, "ssh-vault", r.Header.Get("User-agent"))
		fmt.Fprintln(w, "ssh-rsa FOO")
	}))
	defer ts.Close()
	_, err := New("55:cd:f2:7e:4c:0b:e5:a7:6e:6c:fc:6b:8e:58:9d:15", "", ts.URL, "view", "")
	if err == nil {
		t.Error("Expecting error: Use a public ssh key: illegal base64 ...")
	}
}

func TestVaultNewNoKeyFingerprint(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		expect(t, "ssh-vault", r.Header.Get("User-agent"))
	}))
	defer ts.Close()
	_, err := New("55:cd:f2:7e:4c:0b:e5:a7:6e:6c:fc:6b:8e:58:9d:XX", "", ts.URL, "view", "")
	if err == nil {
		t.Error("Expecting error")
	}
}
<|endoftext|>


<commit_before>package shellwords

import (
	"os"
	"reflect"
	"testing"
)

var testcases = []struct {
	line     string
	expected []string
}{
	{`var --bar=baz`, []string{`var`, `--bar=baz`}},
	{`var --bar="baz"`, []string{`var`, `--bar=baz`}},
	{`var "--bar=baz"`, []string{`var`, `--bar=baz`}},
	{`var "--bar='baz'"`, []string{`var`, `--bar='baz'`}},
	{"var --bar=`baz`", []string{`var`, "--bar=`baz`"}},
	{`var "--bar=\"baz'"`, []string{`var`, `--bar="baz'`}},
	{`var "--bar=\'baz\'"`, []string{`var`, `--bar='baz'`}},
	{`var --bar='\'`, []string{`var`, `--bar=\`}},
	{`var "--bar baz"`, []string{`var`, `--bar baz`}},
	{`var --"bar baz"`, []string{`var`, `--bar baz`}},
	{`var  --"bar baz"`, []string{`var`, `--bar baz`}},
}

func TestSimple(t *testing.T) {
	for _, testcase := range testcases {
		args, err := Parse(testcase.line)
		if err != nil {
			t.Fatal(err)
		}
		if !reflect.DeepEqual(args, testcase.expected) {
			t.Fatalf("Expected %#v, but %#v:", testcase.expected, args)
		}
	}
}

func TestError(t *testing.T) {
	_, err := Parse("foo '")
	if err == nil {
		t.Fatal("Should be an error")
	}
	_, err = Parse(`foo "`)
	if err == nil {
		t.Fatal("Should be an error")
	}

	_, err = Parse("foo `")
	if err == nil {
		t.Fatal("Should be an error")
	}
}

func TestLastSpace(t *testing.T) {
	args, err := Parse("foo bar\\  ")
	if err != nil {
		t.Fatal(err)
	}
	if len(args) != 2 {
		t.Fatal("Should have two elements")
	}
	if args[0] != "foo" {
		t.Fatal("1st element should be `foo`")
	}
	if args[1] != "bar " {
		t.Fatal("1st element should be `bar `")
	}
}

func TestBacktick(t *testing.T) {
	goversion, err := shellRun("go version")
	if err != nil {
		t.Fatal(err)
	}

	parser := NewParser()
	parser.ParseBacktick = true
	args, err := parser.Parse("echo `go version`")
	if err != nil {
		t.Fatal(err)
	}
	expected := []string{"echo", goversion}
	if !reflect.DeepEqual(args, expected) {
		t.Fatalf("Expected %#v, but %#v:", expected, args)
	}

	args, err = parser.Parse(`echo $(echo "foo")`)
	if err != nil {
		t.Fatal(err)
	}
	expected = []string{"echo", "foo"}
	if !reflect.DeepEqual(args, expected) {
		t.Fatalf("Expected %#v, but %#v:", expected, args)
	}
}

func TestBacktickError(t *testing.T) {
	parser := NewParser()
	parser.ParseBacktick = true
	_, err := parser.Parse("echo `go Version`")
	if err == nil {
		t.Fatal("Should be an error")
	}
	expected := "exit status 2:go: unknown subcommand \"Version\"\nRun 'go help' for usage.\n"
	if expected != err.Error() {
		t.Fatalf("Expected %q, but %q", expected, err.Error())
	}
}

func TestEnv(t *testing.T) {
	os.Setenv("FOO", "bar")

	parser := NewParser()
	parser.ParseEnv = true
	args, err := parser.Parse("echo $FOO")
	if err != nil {
		t.Fatal(err)
	}
	expected := []string{"echo", "bar"}
	if !reflect.DeepEqual(args, expected) {
		t.Fatalf("Expected %#v, but %#v:", expected, args)
	}
}

func TestNoEnv(t *testing.T) {
	parser := NewParser()
	parser.ParseEnv = true
	args, err := parser.Parse("echo $BAR")
	if err != nil {
		t.Fatal(err)
	}
	expected := []string{"echo", ""}
	if !reflect.DeepEqual(args, expected) {
		t.Fatalf("Expected %#v, but %#v:", expected, args)
	}
}

func TestDupEnv(t *testing.T) {
	os.Setenv("FOO", "bar")
	os.Setenv("FOO_BAR", "baz")

	parser := NewParser()
	parser.ParseEnv = true
	args, err := parser.Parse("echo $$FOO$")
	if err != nil {
		t.Fatal(err)
	}
	expected := []string{"echo", "$bar$"}
	if !reflect.DeepEqual(args, expected) {
		t.Fatalf("Expected %#v, but %#v:", expected, args)
	}

	args, err = parser.Parse("echo $${FOO_BAR}$")
	if err != nil {
		t.Fatal(err)
	}
	expected = []string{"echo", "$baz$"}
	if !reflect.DeepEqual(args, expected) {
		t.Fatalf("Expected %#v, but %#v:", expected, args)
	}
}

func TestHaveMore(t *testing.T) {
	parser := NewParser()
	parser.ParseEnv = true

	line := "echo foo; seq 1 10"
	args, err := parser.Parse(line)
	if err != nil {
		t.Fatalf(err.Error())
	}
	expected := []string{"echo", "foo"}
	if !reflect.DeepEqual(args, expected) {
		t.Fatalf("Expected %#v, but %#v:", expected, args)
	}

	if parser.Position == 0 {
		t.Fatalf("Commands should be remaining")
	}

	line = string([]rune(line)[parser.Position+1:])
	args, err = parser.Parse(line)
	if err != nil {
		t.Fatalf(err.Error())
	}
	expected = []string{"seq", "1", "10"}
	if !reflect.DeepEqual(args, expected) {
		t.Fatalf("Expected %#v, but %#v:", expected, args)
	}

	if parser.Position > 0 {
		t.Fatalf("Commands should not be remaining")
	}
}
<commit_msg>add test<commit_after>package shellwords

import (
	"os"
	"reflect"
	"testing"
)

var testcases = []struct {
	line     string
	expected []string
}{
	{`var --bar=baz`, []string{`var`, `--bar=baz`}},
	{`var --bar="baz"`, []string{`var`, `--bar=baz`}},
	{`var "--bar=baz"`, []string{`var`, `--bar=baz`}},
	{`var "--bar='baz'"`, []string{`var`, `--bar='baz'`}},
	{"var --bar=`baz`", []string{`var`, "--bar=`baz`"}},
	{`var "--bar=\"baz'"`, []string{`var`, `--bar="baz'`}},
	{`var "--bar=\'baz\'"`, []string{`var`, `--bar='baz'`}},
	{`var --bar='\'`, []string{`var`, `--bar=\`}},
	{`var "--bar baz"`, []string{`var`, `--bar baz`}},
	{`var --"bar baz"`, []string{`var`, `--bar baz`}},
	{`var  --"bar baz"`, []string{`var`, `--bar baz`}},
}

func TestSimple(t *testing.T) {
	for _, testcase := range testcases {
		args, err := Parse(testcase.line)
		if err != nil {
			t.Fatal(err)
		}
		if !reflect.DeepEqual(args, testcase.expected) {
			t.Fatalf("Expected %#v, but %#v:", testcase.expected, args)
		}
	}
}

func TestError(t *testing.T) {
	_, err := Parse("foo '")
	if err == nil {
		t.Fatal("Should be an error")
	}
	_, err = Parse(`foo "`)
	if err == nil {
		t.Fatal("Should be an error")
	}

	_, err = Parse("foo `")
	if err == nil {
		t.Fatal("Should be an error")
	}
}

func TestLastSpace(t *testing.T) {
	args, err := Parse("foo bar\\  ")
	if err != nil {
		t.Fatal(err)
	}
	if len(args) != 2 {
		t.Fatal("Should have two elements")
	}
	if args[0] != "foo" {
		t.Fatal("1st element should be `foo`")
	}
	if args[1] != "bar " {
		t.Fatal("1st element should be `bar `")
	}
}

func TestBacktick(t *testing.T) {
	goversion, err := shellRun("go version")
	if err != nil {
		t.Fatal(err)
	}

	parser := NewParser()
	parser.ParseBacktick = true
	args, err := parser.Parse("echo `go version`")
	if err != nil {
		t.Fatal(err)
	}
	expected := []string{"echo", goversion}
	if !reflect.DeepEqual(args, expected) {
		t.Fatalf("Expected %#v, but %#v:", expected, args)
	}

	args, err = parser.Parse(`echo $(echo "foo")`)
	if err != nil {
		t.Fatal(err)
	}
	expected = []string{"echo", "foo"}
	if !reflect.DeepEqual(args, expected) {
		t.Fatalf("Expected %#v, but %#v:", expected, args)
	}

	args, err = parser.Parse(`echo $(echo1)`)
	if err == nil {
		t.Fatal("Should be an error")
	}
}

func TestBacktickError(t *testing.T) {
	parser := NewParser()
	parser.ParseBacktick = true
	_, err := parser.Parse("echo `go Version`")
	if err == nil {
		t.Fatal("Should be an error")
	}
	expected := "exit status 2:go: unknown subcommand \"Version\"\nRun 'go help' for usage.\n"
	if expected != err.Error() {
		t.Fatalf("Expected %q, but %q", expected, err.Error())
	}
}

func TestEnv(t *testing.T) {
	os.Setenv("FOO", "bar")

	parser := NewParser()
	parser.ParseEnv = true
	args, err := parser.Parse("echo $FOO")
	if err != nil {
		t.Fatal(err)
	}
	expected := []string{"echo", "bar"}
	if !reflect.DeepEqual(args, expected) {
		t.Fatalf("Expected %#v, but %#v:", expected, args)
	}
}

func TestNoEnv(t *testing.T) {
	parser := NewParser()
	parser.ParseEnv = true
	args, err := parser.Parse("echo $BAR")
	if err != nil {
		t.Fatal(err)
	}
	expected := []string{"echo", ""}
	if !reflect.DeepEqual(args, expected) {
		t.Fatalf("Expected %#v, but %#v:", expected, args)
	}
}

func TestDupEnv(t *testing.T) {
	os.Setenv("FOO", "bar")
	os.Setenv("FOO_BAR", "baz")

	parser := NewParser()
	parser.ParseEnv = true
	args, err := parser.Parse("echo $$FOO$")
	if err != nil {
		t.Fatal(err)
	}
	expected := []string{"echo", "$bar$"}
	if !reflect.DeepEqual(args, expected) {
		t.Fatalf("Expected %#v, but %#v:", expected, args)
	}

	args, err = parser.Parse("echo $${FOO_BAR}$")
	if err != nil {
		t.Fatal(err)
	}
	expected = []string{"echo", "$baz$"}
	if !reflect.DeepEqual(args, expected) {
		t.Fatalf("Expected %#v, but %#v:", expected, args)
	}
}

func TestHaveMore(t *testing.T) {
	parser := NewParser()
	parser.ParseEnv = true

	line := "echo foo; seq 1 10"
	args, err := parser.Parse(line)
	if err != nil {
		t.Fatalf(err.Error())
	}
	expected := []string{"echo", "foo"}
	if !reflect.DeepEqual(args, expected) {
		t.Fatalf("Expected %#v, but %#v:", expected, args)
	}

	if parser.Position == 0 {
		t.Fatalf("Commands should be remaining")
	}

	line = string([]rune(line)[parser.Position+1:])
	args, err = parser.Parse(line)
	if err != nil {
		t.Fatalf(err.Error())
	}
	expected = []string{"seq", "1", "10"}
	if !reflect.DeepEqual(args, expected) {
		t.Fatalf("Expected %#v, but %#v:", expected, args)
	}

	if parser.Position > 0 {
		t.Fatalf("Commands should not be remaining")
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2020 The Vitess 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 sqlparser

import (
	querypb "vitess.io/vitess/go/vt/proto/query"
	vtrpcpb "vitess.io/vitess/go/vt/proto/vtrpc"
	"vitess.io/vitess/go/vt/vterrors"

	"strings"

	"vitess.io/vitess/go/vt/sysvars"
)

// RewriteASTResult contains the rewritten ast and meta information about it
type RewriteASTResult struct {
	*BindVarNeeds
	AST Statement // The rewritten AST
}

// PrepareAST will normalize the query
func PrepareAST(in Statement, bindVars map[string]*querypb.BindVariable, prefix string, parameterize bool, keyspace string) (*RewriteASTResult, error) {
	if parameterize {
		Normalize(in, bindVars, prefix)
	}
	return RewriteAST(in, keyspace)
}

// RewriteAST rewrites the whole AST, replacing function calls and adding column aliases to queries
func RewriteAST(in Statement, keyspace string) (*RewriteASTResult, error) {
	er := newExpressionRewriter(keyspace)
	er.shouldRewriteDatabaseFunc = shouldRewriteDatabaseFunc(in)
	setRewriter := &setNormalizer{}
	out, ok := Rewrite(in, er.rewrite, setRewriter.rewriteSetComingUp).(Statement)
	if !ok {
		return nil, vterrors.Errorf(vtrpcpb.Code_INTERNAL, "statement rewriting returned a non statement: %s", String(out))
	}
	if setRewriter.err != nil {
		return nil, setRewriter.err
	}

	r := &RewriteASTResult{
		AST:          out,
		BindVarNeeds: er.bindVars,
	}
	return r, nil
}

func shouldRewriteDatabaseFunc(in Statement) bool {
	selct, ok := in.(*Select)
	if !ok {
		return false
	}
	if len(selct.From) != 1 {
		return false
	}
	aliasedTable, ok := selct.From[0].(*AliasedTableExpr)
	if !ok {
		return false
	}
	tableName, ok := aliasedTable.Expr.(TableName)
	if !ok {
		return false
	}
	return tableName.Name.String() == "dual"
}

type expressionRewriter struct {
	bindVars                  *BindVarNeeds
	shouldRewriteDatabaseFunc bool
	err                       error

	// we need to know this to make a decision if we can safely rewrite JOIN USING => JOIN ON
	hasStarInSelect bool

	keyspace string
}

func newExpressionRewriter(keyspace string) *expressionRewriter {
	return &expressionRewriter{bindVars: &BindVarNeeds{}, keyspace: keyspace}
}

const (
	//LastInsertIDName is a reserved bind var name for last_insert_id()
	LastInsertIDName = "__lastInsertId"

	//DBVarName is a reserved bind var name for database()
	DBVarName = "__vtdbname"

	//FoundRowsName is a reserved bind var name for found_rows()
	FoundRowsName = "__vtfrows"

	//RowCountName is a reserved bind var name for row_count()
	RowCountName = "__vtrcount"

	//UserDefinedVariableName is what we prepend bind var names for user defined variables
	UserDefinedVariableName = "__vtudv"
)

func (er *expressionRewriter) rewriteAliasedExpr(node *AliasedExpr) (*BindVarNeeds, error) {
	inner := newExpressionRewriter("")
	inner.shouldRewriteDatabaseFunc = er.shouldRewriteDatabaseFunc
	tmp := Rewrite(node.Expr, inner.rewrite, nil)
	newExpr, ok := tmp.(Expr)
	if !ok {
		return nil, vterrors.Errorf(vtrpcpb.Code_INTERNAL, "failed to rewrite AST. function expected to return Expr returned a %s", String(tmp))
	}
	node.Expr = newExpr
	return inner.bindVars, nil
}

func (er *expressionRewriter) rewrite(cursor *Cursor) bool {
	switch node := cursor.Node().(type) {
	// select last_insert_id() -> select :__lastInsertId as `last_insert_id()`
	case *Select:
		for _, col := range node.SelectExprs {
			_, hasStar := col.(*StarExpr)
			if hasStar {
				er.hasStarInSelect = true
			}

			aliasedExpr, ok := col.(*AliasedExpr)
			if ok && aliasedExpr.As.IsEmpty() {
				buf := NewTrackedBuffer(nil)
				aliasedExpr.Expr.Format(buf)
				innerBindVarNeeds, err := er.rewriteAliasedExpr(aliasedExpr)
				if err != nil {
					er.err = err
					return false
				}
				if innerBindVarNeeds.HasRewrites() {
					aliasedExpr.As = NewColIdent(buf.String())
				}
				er.bindVars.MergeWith(innerBindVarNeeds)
			}
		}
	case *FuncExpr:
		er.funcRewrite(cursor, node)
	case *ColName:
		switch node.Name.at {
		case SingleAt:
			er.udvRewrite(cursor, node)
		case DoubleAt:
			er.sysVarRewrite(cursor, node)
		}
	case *Subquery:
		er.unnestSubQueries(cursor, node)
	case JoinCondition:
		er.rewriteJoinCondition(cursor, node)
	case *AliasedTableExpr:
		if !SystemSchema(er.keyspace) {
			break
		}
		aliasTableName, ok := node.Expr.(TableName)
		if !ok {
			return true
		}
		if er.keyspace != "" && aliasTableName.Qualifier.IsEmpty() {
			aliasTableName.Qualifier = NewTableIdent(er.keyspace)
			node.Expr = aliasTableName
			cursor.Replace(node)
		}
	}
	return true
}

func (er *expressionRewriter) rewriteJoinCondition(cursor *Cursor, node JoinCondition) {
	if node.Using != nil && !er.hasStarInSelect {
		joinTableExpr, ok := cursor.Parent().(*JoinTableExpr)
		if !ok {
			// this is not possible with the current AST
			return
		}
		leftTable, leftOk := joinTableExpr.LeftExpr.(*AliasedTableExpr)
		rightTable, rightOk := joinTableExpr.RightExpr.(*AliasedTableExpr)
		if !(leftOk && rightOk) {
			// we only deal with simple FROM A JOIN B USING queries at the moment
			return
		}
		lft, err := leftTable.TableName()
		if err != nil {
			er.err = err
			return
		}
		rgt, err := rightTable.TableName()
		if err != nil {
			er.err = err
			return
		}
		newCondition := JoinCondition{}
		for _, colIdent := range node.Using {
			lftCol := NewColNameWithQualifier(colIdent.String(), lft)
			rgtCol := NewColNameWithQualifier(colIdent.String(), rgt)
			cmp := &ComparisonExpr{
				Operator: EqualOp,
				Left:     lftCol,
				Right:    rgtCol,
			}
			if newCondition.On == nil {
				newCondition.On = cmp
			} else {
				newCondition.On = &AndExpr{Left: newCondition.On, Right: cmp}
			}
		}
		cursor.Replace(newCondition)
	}
}

func (er *expressionRewriter) sysVarRewrite(cursor *Cursor, node *ColName) {
	lowered := node.Name.Lowered()
	switch lowered {
	case sysvars.Autocommit.Name,
		sysvars.ClientFoundRows.Name,
		sysvars.SkipQueryPlanCache.Name,
		sysvars.SQLSelectLimit.Name,
		sysvars.TransactionMode.Name,
		sysvars.Workload.Name,
		sysvars.DDLStrategy.Name,
		sysvars.ReadAfterWriteGTID.Name,
		sysvars.ReadAfterWriteTimeOut.Name,
		sysvars.SessionTrackGTIDs.Name:
		cursor.Replace(bindVarExpression("__vt" + lowered))
		er.bindVars.AddSysVar(lowered)
	}
}

func (er *expressionRewriter) udvRewrite(cursor *Cursor, node *ColName) {
	udv := strings.ToLower(node.Name.CompliantName())
	cursor.Replace(bindVarExpression(UserDefinedVariableName + udv))
	er.bindVars.AddUserDefVar(udv)
}

var funcRewrites = map[string]string{
	"last_insert_id": LastInsertIDName,
	"database":       DBVarName,
	"schema":         DBVarName,
	"found_rows":     FoundRowsName,
	"row_count":      RowCountName,
}

func (er *expressionRewriter) funcRewrite(cursor *Cursor, node *FuncExpr) {
	bindVar, found := funcRewrites[node.Name.Lowered()]
	if found {
		if bindVar == DBVarName && !er.shouldRewriteDatabaseFunc {
			return
		}
		if len(node.Exprs) > 0 {
			er.err = vterrors.Errorf(vtrpcpb.Code_UNIMPLEMENTED, "Argument to %s() not supported", node.Name.Lowered())
			return
		}
		cursor.Replace(bindVarExpression(bindVar))
		er.bindVars.AddFuncResult(bindVar)
	}
}

func (er *expressionRewriter) unnestSubQueries(cursor *Cursor, subquery *Subquery) {
	sel, isSimpleSelect := subquery.Select.(*Select)
	if !isSimpleSelect {
		return
	}

	if !(len(sel.SelectExprs) != 1 ||
		len(sel.OrderBy) != 0 ||
		len(sel.GroupBy) != 0 ||
		len(sel.From) != 1 ||
		sel.Where == nil ||
		sel.Having == nil ||
		sel.Limit == nil) && sel.Lock == NoLock {
		return
	}
	aliasedTable, ok := sel.From[0].(*AliasedTableExpr)
	if !ok {
		return
	}
	table, ok := aliasedTable.Expr.(TableName)
	if !ok || table.Name.String() != "dual" {
		return
	}
	expr, ok := sel.SelectExprs[0].(*AliasedExpr)
	if !ok {
		return
	}
	er.bindVars.NoteRewrite()
	// we need to make sure that the inner expression also gets rewritten,
	// so we fire off another rewriter traversal here
	rewrittenExpr := Rewrite(expr.Expr, er.rewrite, nil)
	cursor.Replace(rewrittenExpr)
}

func bindVarExpression(name string) Expr {
	return NewArgument([]byte(":" + name))
}

// SystemSchema returns true if the schema passed is system schema
func SystemSchema(schema string) bool {
	return strings.EqualFold(schema, "information_schema") ||
		strings.EqualFold(schema, "performance_schema") ||
		strings.EqualFold(schema, "sys") ||
		strings.EqualFold(schema, "mysql")
}
<commit_msg>addressed review comments<commit_after>/*
Copyright 2020 The Vitess 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 sqlparser

import (
	querypb "vitess.io/vitess/go/vt/proto/query"
	vtrpcpb "vitess.io/vitess/go/vt/proto/vtrpc"
	"vitess.io/vitess/go/vt/vterrors"

	"strings"

	"vitess.io/vitess/go/vt/sysvars"
)

// RewriteASTResult contains the rewritten ast and meta information about it
type RewriteASTResult struct {
	*BindVarNeeds
	AST Statement // The rewritten AST
}

// PrepareAST will normalize the query
func PrepareAST(in Statement, bindVars map[string]*querypb.BindVariable, prefix string, parameterize bool, keyspace string) (*RewriteASTResult, error) {
	if parameterize {
		Normalize(in, bindVars, prefix)
	}
	return RewriteAST(in, keyspace)
}

// RewriteAST rewrites the whole AST, replacing function calls and adding column aliases to queries
func RewriteAST(in Statement, keyspace string) (*RewriteASTResult, error) {
	er := newExpressionRewriter(keyspace)
	er.shouldRewriteDatabaseFunc = shouldRewriteDatabaseFunc(in)
	setRewriter := &setNormalizer{}
	out, ok := Rewrite(in, er.rewrite, setRewriter.rewriteSetComingUp).(Statement)
	if !ok {
		return nil, vterrors.Errorf(vtrpcpb.Code_INTERNAL, "statement rewriting returned a non statement: %s", String(out))
	}
	if setRewriter.err != nil {
		return nil, setRewriter.err
	}

	r := &RewriteASTResult{
		AST:          out,
		BindVarNeeds: er.bindVars,
	}
	return r, nil
}

func shouldRewriteDatabaseFunc(in Statement) bool {
	selct, ok := in.(*Select)
	if !ok {
		return false
	}
	if len(selct.From) != 1 {
		return false
	}
	aliasedTable, ok := selct.From[0].(*AliasedTableExpr)
	if !ok {
		return false
	}
	tableName, ok := aliasedTable.Expr.(TableName)
	if !ok {
		return false
	}
	return tableName.Name.String() == "dual"
}

type expressionRewriter struct {
	bindVars                  *BindVarNeeds
	shouldRewriteDatabaseFunc bool
	err                       error

	// we need to know this to make a decision if we can safely rewrite JOIN USING => JOIN ON
	hasStarInSelect bool

	keyspace string
}

func newExpressionRewriter(keyspace string) *expressionRewriter {
	return &expressionRewriter{bindVars: &BindVarNeeds{}, keyspace: keyspace}
}

const (
	//LastInsertIDName is a reserved bind var name for last_insert_id()
	LastInsertIDName = "__lastInsertId"

	//DBVarName is a reserved bind var name for database()
	DBVarName = "__vtdbname"

	//FoundRowsName is a reserved bind var name for found_rows()
	FoundRowsName = "__vtfrows"

	//RowCountName is a reserved bind var name for row_count()
	RowCountName = "__vtrcount"

	//UserDefinedVariableName is what we prepend bind var names for user defined variables
	UserDefinedVariableName = "__vtudv"
)

func (er *expressionRewriter) rewriteAliasedExpr(node *AliasedExpr) (*BindVarNeeds, error) {
	inner := newExpressionRewriter(er.keyspace)
	inner.shouldRewriteDatabaseFunc = er.shouldRewriteDatabaseFunc
	tmp := Rewrite(node.Expr, inner.rewrite, nil)
	newExpr, ok := tmp.(Expr)
	if !ok {
		return nil, vterrors.Errorf(vtrpcpb.Code_INTERNAL, "failed to rewrite AST. function expected to return Expr returned a %s", String(tmp))
	}
	node.Expr = newExpr
	return inner.bindVars, nil
}

func (er *expressionRewriter) rewrite(cursor *Cursor) bool {
	switch node := cursor.Node().(type) {
	// select last_insert_id() -> select :__lastInsertId as `last_insert_id()`
	case *Select:
		for _, col := range node.SelectExprs {
			_, hasStar := col.(*StarExpr)
			if hasStar {
				er.hasStarInSelect = true
			}

			aliasedExpr, ok := col.(*AliasedExpr)
			if ok && aliasedExpr.As.IsEmpty() {
				buf := NewTrackedBuffer(nil)
				aliasedExpr.Expr.Format(buf)
				innerBindVarNeeds, err := er.rewriteAliasedExpr(aliasedExpr)
				if err != nil {
					er.err = err
					return false
				}
				if innerBindVarNeeds.HasRewrites() {
					aliasedExpr.As = NewColIdent(buf.String())
				}
				er.bindVars.MergeWith(innerBindVarNeeds)
			}
		}
	case *FuncExpr:
		er.funcRewrite(cursor, node)
	case *ColName:
		switch node.Name.at {
		case SingleAt:
			er.udvRewrite(cursor, node)
		case DoubleAt:
			er.sysVarRewrite(cursor, node)
		}
	case *Subquery:
		er.unnestSubQueries(cursor, node)
	case JoinCondition:
		er.rewriteJoinCondition(cursor, node)
	case *AliasedTableExpr:
		if !SystemSchema(er.keyspace) {
			break
		}
		aliasTableName, ok := node.Expr.(TableName)
		if !ok {
			return true
		}
		if er.keyspace != "" && aliasTableName.Qualifier.IsEmpty() {
			aliasTableName.Qualifier = NewTableIdent(er.keyspace)
			node.Expr = aliasTableName
			cursor.Replace(node)
		}
	}
	return true
}

func (er *expressionRewriter) rewriteJoinCondition(cursor *Cursor, node JoinCondition) {
	if node.Using != nil && !er.hasStarInSelect {
		joinTableExpr, ok := cursor.Parent().(*JoinTableExpr)
		if !ok {
			// this is not possible with the current AST
			return
		}
		leftTable, leftOk := joinTableExpr.LeftExpr.(*AliasedTableExpr)
		rightTable, rightOk := joinTableExpr.RightExpr.(*AliasedTableExpr)
		if !(leftOk && rightOk) {
			// we only deal with simple FROM A JOIN B USING queries at the moment
			return
		}
		lft, err := leftTable.TableName()
		if err != nil {
			er.err = err
			return
		}
		rgt, err := rightTable.TableName()
		if err != nil {
			er.err = err
			return
		}
		newCondition := JoinCondition{}
		for _, colIdent := range node.Using {
			lftCol := NewColNameWithQualifier(colIdent.String(), lft)
			rgtCol := NewColNameWithQualifier(colIdent.String(), rgt)
			cmp := &ComparisonExpr{
				Operator: EqualOp,
				Left:     lftCol,
				Right:    rgtCol,
			}
			if newCondition.On == nil {
				newCondition.On = cmp
			} else {
				newCondition.On = &AndExpr{Left: newCondition.On, Right: cmp}
			}
		}
		cursor.Replace(newCondition)
	}
}

func (er *expressionRewriter) sysVarRewrite(cursor *Cursor, node *ColName) {
	lowered := node.Name.Lowered()
	switch lowered {
	case sysvars.Autocommit.Name,
		sysvars.ClientFoundRows.Name,
		sysvars.SkipQueryPlanCache.Name,
		sysvars.SQLSelectLimit.Name,
		sysvars.TransactionMode.Name,
		sysvars.Workload.Name,
		sysvars.DDLStrategy.Name,
		sysvars.ReadAfterWriteGTID.Name,
		sysvars.ReadAfterWriteTimeOut.Name,
		sysvars.SessionTrackGTIDs.Name:
		cursor.Replace(bindVarExpression("__vt" + lowered))
		er.bindVars.AddSysVar(lowered)
	}
}

func (er *expressionRewriter) udvRewrite(cursor *Cursor, node *ColName) {
	udv := strings.ToLower(node.Name.CompliantName())
	cursor.Replace(bindVarExpression(UserDefinedVariableName + udv))
	er.bindVars.AddUserDefVar(udv)
}

var funcRewrites = map[string]string{
	"last_insert_id": LastInsertIDName,
	"database":       DBVarName,
	"schema":         DBVarName,
	"found_rows":     FoundRowsName,
	"row_count":      RowCountName,
}

func (er *expressionRewriter) funcRewrite(cursor *Cursor, node *FuncExpr) {
	bindVar, found := funcRewrites[node.Name.Lowered()]
	if found {
		if bindVar == DBVarName && !er.shouldRewriteDatabaseFunc {
			return
		}
		if len(node.Exprs) > 0 {
			er.err = vterrors.Errorf(vtrpcpb.Code_UNIMPLEMENTED, "Argument to %s() not supported", node.Name.Lowered())
			return
		}
		cursor.Replace(bindVarExpression(bindVar))
		er.bindVars.AddFuncResult(bindVar)
	}
}

func (er *expressionRewriter) unnestSubQueries(cursor *Cursor, subquery *Subquery) {
	sel, isSimpleSelect := subquery.Select.(*Select)
	if !isSimpleSelect {
		return
	}

	if !(len(sel.SelectExprs) != 1 ||
		len(sel.OrderBy) != 0 ||
		len(sel.GroupBy) != 0 ||
		len(sel.From) != 1 ||
		sel.Where == nil ||
		sel.Having == nil ||
		sel.Limit == nil) && sel.Lock == NoLock {
		return
	}
	aliasedTable, ok := sel.From[0].(*AliasedTableExpr)
	if !ok {
		return
	}
	table, ok := aliasedTable.Expr.(TableName)
	if !ok || table.Name.String() != "dual" {
		return
	}
	expr, ok := sel.SelectExprs[0].(*AliasedExpr)
	if !ok {
		return
	}
	er.bindVars.NoteRewrite()
	// we need to make sure that the inner expression also gets rewritten,
	// so we fire off another rewriter traversal here
	rewrittenExpr := Rewrite(expr.Expr, er.rewrite, nil)
	cursor.Replace(rewrittenExpr)
}

func bindVarExpression(name string) Expr {
	return NewArgument([]byte(":" + name))
}

// SystemSchema returns true if the schema passed is system schema
func SystemSchema(schema string) bool {
	return strings.EqualFold(schema, "information_schema") ||
		strings.EqualFold(schema, "performance_schema") ||
		strings.EqualFold(schema, "sys") ||
		strings.EqualFold(schema, "mysql")
}
<|endoftext|>


<commit_before>// Copyright 2015, David Howden
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package index defines functionality for creating and manipulating a Tchaik music index.
package index

import (
	"compress/gzip"
	"encoding/json"
	"fmt"
	"io"
	"time"
)

// Library is an interface which defines methods for listing tracks.
type Library interface {
	// Tracks returns a slice of all the tracks in the library.
	Tracks() []Track

	// Track returns the track from the given identifier and true if successful,
	// false otherwise.
	Track(identifier string) (Track, bool)
}

// Track is an interface which defines methods for retrieving track metadata.  Methods return
// zero values if attributes are unset and panic on undefined attributes and type mismatches.
type Track interface {
	// GetString returns a string value for the given attribute name.  Panics
	// if no such string attribute exists.
	GetString(string) string

	// GetStrings returns a list of strings for the given attribute name.
	// Panics if no such attribute exists.
	GetStrings(string) []string

	// GetInt returns an int value for the given attribute name.  Panics if no such
	// attribute exists.
	GetInt(string) int

	// GetTime returns a time.Time value for the given attribute name. Panics if no
	// such attribute exists.
	GetTime(string) time.Time
}

// Convert reads all the data exported by the Library and writes into the standard
// tchaik Library implementation.
// NB: The identifier field is set to be the value of ID on every track, regardless
// of whether this value has already been set in the input Library.
func Convert(l Library, id string) *library {
	allTracks := l.Tracks()
	tracks := make(map[string]*track, len(allTracks))

	for _, t := range allTracks {
		identifier := t.GetString(id)
		tracks[identifier] = &track{
			// string fields
			ID:          identifier,
			Name:        t.GetString("Name"),
			Album:       t.GetString("Album"),
			AlbumArtist: t.GetString("AlbumArtist"),
			Artist:      t.GetString("Artist"),
			Composer:    t.GetString("Composer"),
			Genre:       t.GetString("Genre"),
			Location:    t.GetString("Location"),

			// integer fields
			TotalTime:   t.GetInt("TotalTime"),
			Year:        t.GetInt("Year"),
			DiscNumber:  t.GetInt("DiscNumber"),
			TrackNumber: t.GetInt("TrackNumber"),
			TrackCount:  t.GetInt("TrackCount"),
			DiscCount:   t.GetInt("DiscCount"),
			BitRate:     t.GetInt("BitRate"),

			// date fields
			DateAdded:    t.GetTime("DateAdded"),
			DateModified: t.GetTime("DateModified"),
		}
	}
	return &library{
		tracks,
	}
}

// library is the default internal implementation Library which acts as the data
// source for all media tracks.
type library struct {
	trks map[string]*track
}

// Tracks implements Library.
func (l *library) Tracks() []Track {
	tracks := make([]Track, 0, len(l.trks))
	for _, t := range l.trks {
		tracks = append(tracks, t)
	}
	return tracks
}

// Track implements Library.
func (l *library) Track(id string) (Track, bool) {
	t, ok := l.trks[id]
	return t, ok
}

func (l *library) MarshalJSON() ([]byte, error) {
	return json.Marshal(l.trks)
}

func (l *library) UnmarshalJSON(b []byte) error {
	l.trks = make(map[string]*track)
	return json.Unmarshal(b, &l.trks)
}

// WriteTo writes the Library data to the writer, currently using gzipped-JSON.
func WriteTo(l Library, w io.Writer) error {
	gzw, err := gzip.NewWriterLevel(w, gzip.BestSpeed)
	if err != nil {
		return err
	}
	defer gzw.Close()
	enc := json.NewEncoder(gzw)
	return enc.Encode(l)
}

// ReadFrom reads the gzipped-JSON representation of a Library.
func ReadFrom(r io.Reader) (Library, error) {
	gzr, err := gzip.NewReader(r)
	if err != nil {
		return nil, err
	}
	defer gzr.Close()
	dec := json.NewDecoder(gzr)
	l := &library{}
	err = dec.Decode(l)
	return l, err
}

// track is the default implementation of the Track interface.
type track struct {
	ID          string `json:"id,omitempty"`
	Name        string `json:"name,omitempty"`
	Album       string `json:"album,omitempty"`
	AlbumArtist string `json:"albumArtist,omitempty"`
	Artist      string `json:"artist,omitempty"`
	Composer    string `json:"composer,omitempty"`
	Genre       string `json:"genre,omitempty"`
	Location    string `json:"location,omitempty"`

	TotalTime   int `json:"totalTime,omitempty"`
	Year        int `json:"year,omitempty"`
	DiscNumber  int `json:"discNumber,omitempty"`
	TrackNumber int `json:"trackNumber,omitempty"`
	TrackCount  int `json:"trackCount,omitempty"`
	DiscCount   int `json:"discCount,omitempty"`
	BitRate     int `json:"bitRate,omitempty"`

	DateAdded    time.Time `json:"dateAdded,omitempty"`
	DateModified time.Time `json:"dateModified,omitempty"`
}

// GetString implements Track.
func (t *track) GetString(name string) string {
	switch name {
	case "ID":
		return t.ID
	case "Name":
		return t.Name
	case "Album":
		return t.Album
	case "AlbumArtist":
		return t.AlbumArtist
	case "Artist":
		return t.Artist
	case "Composer":
		return t.Composer
	case "Genre":
		return t.Genre
	case "Location":
		return t.Location
	}
	panic(fmt.Sprintf("unknown string field '%v'", name))
}

// DefaultGetStrings is a function which returns the default value for a GetStrings
// attribute which is based on an existing GetString attribute.  In particular, we handle
// the case where an empty 'GetString' attribute would be "", whereas the corresponding
// 'GetStrings' method should return 'nil' and not '[]string{""}'.
func DefaultGetStrings(t Track, f string) []string {
	v := t.GetString(f)
	if v != "" {
		return []string{v}
	}
	return nil
}

// GetStrings implements Track.
func (t *track) GetStrings(name string) []string {
	switch name {
	case "Artist", "AlbumArtist", "Composer":
		return DefaultGetStrings(t, name)
	}
	panic(fmt.Sprintf("unknown strings field '%v", name))
}

// GetInt implements Track.
func (t *track) GetInt(name string) int {
	switch name {
	case "TotalTime":
		return t.TotalTime
	case "Year":
		return t.Year
	case "DiscNumber":
		return t.DiscNumber
	case "TrackNumber":
		return t.TrackNumber
	case "TrackCount":
		return t.TrackCount
	case "DiscCount":
		return t.DiscCount
	case "BitRate":
		return t.BitRate
	}
	panic(fmt.Sprintf("unknown int field '%v'", name))
}

// GetTime implements Track.
func (t *track) GetTime(name string) time.Time {
	switch name {
	case "DateAdded":
		return t.DateAdded
	case "DateModified":
		return t.DateModified
	}
	panic(fmt.Sprintf("unknown time field '%v'", name))
}
<commit_msg>Add Kind field to tchaik Track type.<commit_after>// Copyright 2015, David Howden
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package index defines functionality for creating and manipulating a Tchaik music index.
package index

import (
	"compress/gzip"
	"encoding/json"
	"fmt"
	"io"
	"time"
)

// Library is an interface which defines methods for listing tracks.
type Library interface {
	// Tracks returns a slice of all the tracks in the library.
	Tracks() []Track

	// Track returns the track from the given identifier and true if successful,
	// false otherwise.
	Track(identifier string) (Track, bool)
}

// Track is an interface which defines methods for retrieving track metadata.  Methods return
// zero values if attributes are unset and panic on undefined attributes and type mismatches.
type Track interface {
	// GetString returns a string value for the given attribute name.  Panics
	// if no such string attribute exists.
	GetString(string) string

	// GetStrings returns a list of strings for the given attribute name.
	// Panics if no such attribute exists.
	GetStrings(string) []string

	// GetInt returns an int value for the given attribute name.  Panics if no such
	// attribute exists.
	GetInt(string) int

	// GetTime returns a time.Time value for the given attribute name. Panics if no
	// such attribute exists.
	GetTime(string) time.Time
}

// Convert reads all the data exported by the Library and writes into the standard
// tchaik Library implementation.
// NB: The identifier field is set to be the value of ID on every track, regardless
// of whether this value has already been set in the input Library.
func Convert(l Library, id string) *library {
	allTracks := l.Tracks()
	tracks := make(map[string]*track, len(allTracks))

	for _, t := range allTracks {
		identifier := t.GetString(id)
		tracks[identifier] = &track{
			// string fields
			ID:          identifier,
			Name:        t.GetString("Name"),
			Album:       t.GetString("Album"),
			AlbumArtist: t.GetString("AlbumArtist"),
			Artist:      t.GetString("Artist"),
			Composer:    t.GetString("Composer"),
			Genre:       t.GetString("Genre"),
			Location:    t.GetString("Location"),

			// integer fields
			TotalTime:   t.GetInt("TotalTime"),
			Year:        t.GetInt("Year"),
			DiscNumber:  t.GetInt("DiscNumber"),
			TrackNumber: t.GetInt("TrackNumber"),
			TrackCount:  t.GetInt("TrackCount"),
			DiscCount:   t.GetInt("DiscCount"),
			BitRate:     t.GetInt("BitRate"),

			// date fields
			DateAdded:    t.GetTime("DateAdded"),
			DateModified: t.GetTime("DateModified"),
		}
	}
	return &library{
		tracks,
	}
}

// library is the default internal implementation Library which acts as the data
// source for all media tracks.
type library struct {
	trks map[string]*track
}

// Tracks implements Library.
func (l *library) Tracks() []Track {
	tracks := make([]Track, 0, len(l.trks))
	for _, t := range l.trks {
		tracks = append(tracks, t)
	}
	return tracks
}

// Track implements Library.
func (l *library) Track(id string) (Track, bool) {
	t, ok := l.trks[id]
	return t, ok
}

func (l *library) MarshalJSON() ([]byte, error) {
	return json.Marshal(l.trks)
}

func (l *library) UnmarshalJSON(b []byte) error {
	l.trks = make(map[string]*track)
	return json.Unmarshal(b, &l.trks)
}

// WriteTo writes the Library data to the writer, currently using gzipped-JSON.
func WriteTo(l Library, w io.Writer) error {
	gzw, err := gzip.NewWriterLevel(w, gzip.BestSpeed)
	if err != nil {
		return err
	}
	defer gzw.Close()
	enc := json.NewEncoder(gzw)
	return enc.Encode(l)
}

// ReadFrom reads the gzipped-JSON representation of a Library.
func ReadFrom(r io.Reader) (Library, error) {
	gzr, err := gzip.NewReader(r)
	if err != nil {
		return nil, err
	}
	defer gzr.Close()
	dec := json.NewDecoder(gzr)
	l := &library{}
	err = dec.Decode(l)
	return l, err
}

// track is the default implementation of the Track interface.
type track struct {
	ID          string `json:"id,omitempty"`
	Name        string `json:"name,omitempty"`
	Album       string `json:"album,omitempty"`
	AlbumArtist string `json:"albumArtist,omitempty"`
	Artist      string `json:"artist,omitempty"`
	Composer    string `json:"composer,omitempty"`
	Genre       string `json:"genre,omitempty"`
	Location    string `json:"location,omitempty"`
	Kind        string `json:"kind"`

	TotalTime   int `json:"totalTime,omitempty"`
	Year        int `json:"year,omitempty"`
	DiscNumber  int `json:"discNumber,omitempty"`
	TrackNumber int `json:"trackNumber,omitempty"`
	TrackCount  int `json:"trackCount,omitempty"`
	DiscCount   int `json:"discCount,omitempty"`
	BitRate     int `json:"bitRate,omitempty"`

	DateAdded    time.Time `json:"dateAdded,omitempty"`
	DateModified time.Time `json:"dateModified,omitempty"`
}

// GetString implements Track.
func (t *track) GetString(name string) string {
	switch name {
	case "ID":
		return t.ID
	case "Name":
		return t.Name
	case "Album":
		return t.Album
	case "AlbumArtist":
		return t.AlbumArtist
	case "Artist":
		return t.Artist
	case "Composer":
		return t.Composer
	case "Genre":
		return t.Genre
	case "Location":
		return t.Location
	case "Kind":
		return t.Kind
	}
	panic(fmt.Sprintf("unknown string field '%v'", name))
}

// DefaultGetStrings is a function which returns the default value for a GetStrings
// attribute which is based on an existing GetString attribute.  In particular, we handle
// the case where an empty 'GetString' attribute would be "", whereas the corresponding
// 'GetStrings' method should return 'nil' and not '[]string{""}'.
func DefaultGetStrings(t Track, f string) []string {
	v := t.GetString(f)
	if v != "" {
		return []string{v}
	}
	return nil
}

// GetStrings implements Track.
func (t *track) GetStrings(name string) []string {
	switch name {
	case "Artist", "AlbumArtist", "Composer":
		return DefaultGetStrings(t, name)
	}
	panic(fmt.Sprintf("unknown strings field '%v", name))
}

// GetInt implements Track.
func (t *track) GetInt(name string) int {
	switch name {
	case "TotalTime":
		return t.TotalTime
	case "Year":
		return t.Year
	case "DiscNumber":
		return t.DiscNumber
	case "TrackNumber":
		return t.TrackNumber
	case "TrackCount":
		return t.TrackCount
	case "DiscCount":
		return t.DiscCount
	case "BitRate":
		return t.BitRate
	}
	panic(fmt.Sprintf("unknown int field '%v'", name))
}

// GetTime implements Track.
func (t *track) GetTime(name string) time.Time {
	switch name {
	case "DateAdded":
		return t.DateAdded
	case "DateModified":
		return t.DateModified
	}
	panic(fmt.Sprintf("unknown time field '%v'", name))
}
<|endoftext|>