<commit_before>package main

import (
	"fmt"
	"github.com/robertkrimen/otto"
	"gopkg.in/alecthomas/kingpin.v2"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"sort"
	"strings"
)

var (
	codeListPtr  = ArgsStrList(kingpin.Arg("code", "the js code you want to run").Required())
	lineSeqPtr   = kingpin.Flag("line-seq", "the char used for split line").Short('l').Default("\n").String()
	columnSeqPtr = kingpin.Flag("column-seq", "the char used for split column").Short('c').Default(" +").String()
	funcListPtr  = ArgsStrList(kingpin.Flag("funcion", "function").Short('f'))
	pathListPtr  = ArgsStrList(kingpin.Flag("path", "command search path").Short('p'))
	jsListPtr    = ArgsStrList(kingpin.Flag("js", "Javascript file").Short('j'))
	reversePtr   = kingpin.Flag("reverse", "execute code with reverse order").Short('r').Bool()
)

type argsStrList []string

func (i *argsStrList) Set(value string) error {
	*i = append(*i, value)
	return nil
}

func (i *argsStrList) String() string {
	return ""
}

func (i *argsStrList) IsCumulative() bool {
	return true
}

func ArgsStrList(s kingpin.Settings) (target *[]string) {
	target = new([]string)
	s.SetValue((*argsStrList)(target))
	return
}

var builtinFuncCode = []string{
	`
printArray = function(arr) {
	for(i=0;i<arr.length;i+=1){
		console.log(arr[i])
	}
}
printA=printArray;
printL=printArray;
`,
	`
function isFunc(val) {
	if(typeof val === 'function') {
		return true;
	}else{
		return false;
	}
}
`,
	`
mapIf = function(arr, func, preFunc, postFunc) {
	ret = []
	for(i=0;i<arr.length;i+=1) {
		val = arr[i];
		if(isFunc(preFunc) && !preFunc(val)){
			continue;
		}
		out = func(val);
		if(isFunc(postFunc) && !postFunc(out)){
			continue;
		}
		ret.push(val);
	}
	return ret;
}

`,
	`
printColumn = function(arr, column, func) {
	if(isFunc(func)) printFunc=func;
	else printFunc=println;
	arr.map(function(line, index, x){
		printFunc(line[column]);
	});
}
printC=printColumn;
`,
}

func readAll() string {
	bytes, err := ioutil.ReadAll(os.Stdin)
	if err == nil {
		return string(bytes)
	} else {
		panic("Failed to read stdin")
	}
}

func readJSFile(vm *otto.Otto, path string) {
	file, err := os.Open(path)
	defer file.Close()
	if err != nil {
		panic(err)
	}
	bytes, err := ioutil.ReadAll(file)
	if err != nil {
		panic(err)
	}
	_, err = vm.Run(string(bytes))
	if err != nil {
		panic(err)
	}
}

type Matrix [][]string

func setBuiltinFunc(vm *otto.Otto) {
	for _, ele := range builtinFuncCode {
		_, err := vm.Run(ele)
		if err != nil {
			panic(err)
		}
	}

	vm.Set("exec", func(call otto.FunctionCall) otto.Value {
		args := make([]string, len(call.ArgumentList))
		cmdPath, err := call.Argument(0).ToString()
		if err != nil {
			panic(err)
		}
		for i := 1; i < len(call.ArgumentList); i += 1 {
			ret, err := call.Argument(i).ToString()
			if err == nil {
				args[i] = ret
			} else {
				panic(err)
			}

		}
		output := callExternalFunc(cmdPath, args)
		result, err := vm.ToValue(output)
		if err != nil {
			panic(err)
		}

		return result
	})

	vm.Set("execStdin", func(call otto.FunctionCall) otto.Value {
		args := make([]string, len(call.ArgumentList))
		cmdPath, err := call.Argument(0).ToString()
		if err != nil {
			panic(err)
		}
		stdin, err := call.Argument(1).ToString()
		if err != nil {
			panic(err)
		}
		for i := 2; i < len(call.ArgumentList); i += 1 {
			ret, err := call.Argument(i).ToString()
			if err == nil {
				args[i] = ret
			} else {
				panic(err)
			}

		}
		output := callExternalFuncWithStdin(cmdPath, stdin, args)
		result, err := vm.ToValue(output)
		if err != nil {
			panic(err)
		}

		return result
	})
	vm.Set("print", func(call otto.FunctionCall) otto.Value {
		for i := 0; i < len(call.ArgumentList); i += 1 {
			str, err := call.Argument(i).ToString()
			if err == nil {
				fmt.Print(str)
			} else {
				panic(err)
			}
		}
		val, _ := vm.ToValue("")
		return val
	})

	sprintf := func(call otto.FunctionCall) string {
		args := make([]interface{}, len(call.ArgumentList)-1)
		format, err := call.Argument(0).ToString()
		if err != nil {
			panic(err)
		}
		for i := 1; i < len(call.ArgumentList); i += 1 {
			str, err := call.Argument(i).ToString()
			if err == nil {
				args[i-1] = str
			} else {
				panic(err)
			}
		}
		return fmt.Sprintf(format, args...)
	}

	vm.Set("sprint", func(call otto.FunctionCall) otto.Value {
		val := sprintf(call)
		ret, _ := vm.ToValue(val)
		return ret
	})
	vm.Set("printf", func(call otto.FunctionCall) otto.Value {
		val := sprintf(call)
		fmt.Print(val)
		ret, _ := vm.ToValue(val)
		return ret
	})
}

func callExternalFunc(cmd string, args []string) string {
	out, err := exec.Command(cmd, args...).Output()
	if err != nil {
		panic(err)
	}
	return string(out)
}

func callExternalFuncWithStdin(cmd string, stdin string, args []string) string {
	p := exec.Command(cmd, args...)
	p.Stdin = strings.NewReader(stdin)
	out, err := p.Output()
	if err != nil {
		panic(err)
	}
	return string(out)
}

func addPATHEnv(path string) {
	oldEnv := os.Getenv("PATH")
	newEnv := path + ":" + oldEnv
	os.Setenv("PATH", newEnv)
}

func getWd() string {
	cwd, err := os.Getwd()
	if err != nil {
		panic(err)
	}
	return cwd
}

func initExternelFunc(vm *otto.Otto, cmdPath string) {
	funcName := cmdPath[:len(cmdPath)-len(filepath.Ext(cmdPath))]
	vm.Set(funcName, func(call otto.FunctionCall) otto.Value {
		args := make([]string, len(call.ArgumentList))
		for i := 0; i < len(call.ArgumentList); i += 1 {
			ret, err := call.Argument(i).ToString()
			if err == nil {
				args[i] = ret
			} else {
				panic(err)
			}

		}
		output := callExternalFunc(cmdPath, args)
		result, err := vm.ToValue(output)
		if err != nil {
			panic(err)
		}
		return result
	})

}

func main() {
	addPATHEnv(getWd())
	fmt.Printf("")
	kingpin.Parse()

	stdin := readAll()

	lineSpliter := regexp.MustCompile(*lineSeqPtr)
	columnSpilter := regexp.MustCompile(*columnSeqPtr)

	lines := lineSpliter.Split(stdin, -1)
	matrixPtr := new(Matrix)

	for i := 0; i < len(lines); i += 1 {
		*matrixPtr = append(*matrixPtr, columnSpilter.Split(lines[i], -1))
	}

	vm := otto.New()

	vm.Set("stdin", stdin)
	vm.Set("stdout", "")
	vm.Set("lines", lines)
	vm.Set("matrix", *matrixPtr)
	vm.Run("println = console.log")
	setBuiltinFunc(vm)

	for i := 0; i < len(*funcListPtr); i += 1 {
		initExternelFunc(vm, (*funcListPtr)[i])
	}

	for i := 0; i < len(*pathListPtr); i += 1 {
		addPATHEnv((*pathListPtr)[i])
	}

	for i := 0; i < len(*jsListPtr); i += 1 {
		readJSFile(vm, (*jsListPtr)[i])
	}

	codeList := sort.StringSlice(*codeListPtr)

	if *reversePtr {
		for i, j := 0, len(codeList)-1; i < j; i, j = i+1, j-1 {
			codeList[i], codeList[j] = codeList[j], codeList[i]
		}
	}

	for i := 0; i < len(codeList); i += 1 {
		_, err := vm.Run(codeList[i])
		if err != nil {
			panic(err)
		}
	}

}
<commit_msg>add nostdin flag<commit_after>package main

import (
	"fmt"
	"github.com/robertkrimen/otto"
	"gopkg.in/alecthomas/kingpin.v2"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"sort"
	"strings"
)

var (
	codeListPtr  = ArgsStrList(kingpin.Arg("code", "the js code you want to run").Required())
	lineSeqPtr   = kingpin.Flag("line-seq", "the char used for split line").Short('l').Default("\n").String()
	columnSeqPtr = kingpin.Flag("column-seq", "the char used for split column").Short('c').Default(" +").String()
	funcListPtr  = ArgsStrList(kingpin.Flag("funcion", "function").Short('f'))
	pathListPtr  = ArgsStrList(kingpin.Flag("path", "command search path").Short('p'))
	jsListPtr    = ArgsStrList(kingpin.Flag("js", "Javascript file").Short('j'))
	reversePtr   = kingpin.Flag("reverse", "execute code with reverse order").Short('r').Bool()
	nostdinPtr   = kingpin.Flag("nostdin", "Dont't read from stdin").Bool()
)

type argsStrList []string

func (i *argsStrList) Set(value string) error {
	*i = append(*i, value)
	return nil
}

func (i *argsStrList) String() string {
	return ""
}

func (i *argsStrList) IsCumulative() bool {
	return true
}

func ArgsStrList(s kingpin.Settings) (target *[]string) {
	target = new([]string)
	s.SetValue((*argsStrList)(target))
	return
}

var builtinFuncCode = []string{
	`
printArray = function(arr) {
	for(i=0;i<arr.length;i+=1){
		console.log(arr[i])
	}
}
printA=printArray;
printL=printArray;
`,
	`
function isFunc(val) {
	if(typeof val === 'function') {
		return true;
	}else{
		return false;
	}
}
`,
	`
mapIf = function(arr, func, preFunc, postFunc) {
	ret = []
	for(i=0;i<arr.length;i+=1) {
		val = arr[i];
		if(isFunc(preFunc) && !preFunc(val)){
			continue;
		}
		out = func(val);
		if(isFunc(postFunc) && !postFunc(out)){
			continue;
		}
		ret.push(val);
	}
	return ret;
}

`,
	`
printColumn = function(arr, column, func) {
	if(isFunc(func)) printFunc=func;
	else printFunc=println;
	arr.map(function(line, index, x){
		printFunc(line[column]);
	});
}
printC=printColumn;
`,
}

func readAll() string {
	bytes, err := ioutil.ReadAll(os.Stdin)
	if err == nil {
		return string(bytes)
	} else {
		panic("Failed to read stdin")
	}
}

func readJSFile(vm *otto.Otto, path string) {
	file, err := os.Open(path)
	defer file.Close()
	if err != nil {
		panic(err)
	}
	bytes, err := ioutil.ReadAll(file)
	if err != nil {
		panic(err)
	}
	_, err = vm.Run(string(bytes))
	if err != nil {
		panic(err)
	}
}

type Matrix [][]string

func setBuiltinFunc(vm *otto.Otto) {
	for _, ele := range builtinFuncCode {
		_, err := vm.Run(ele)
		if err != nil {
			panic(err)
		}
	}

	vm.Set("exec", func(call otto.FunctionCall) otto.Value {
		args := make([]string, len(call.ArgumentList))
		cmdPath, err := call.Argument(0).ToString()
		if err != nil {
			panic(err)
		}
		for i := 1; i < len(call.ArgumentList); i += 1 {
			ret, err := call.Argument(i).ToString()
			if err == nil {
				args[i] = ret
			} else {
				panic(err)
			}

		}
		output := callExternalFunc(cmdPath, args)
		result, err := vm.ToValue(output)
		if err != nil {
			panic(err)
		}

		return result
	})

	vm.Set("execStdin", func(call otto.FunctionCall) otto.Value {
		args := make([]string, len(call.ArgumentList))
		cmdPath, err := call.Argument(0).ToString()
		if err != nil {
			panic(err)
		}
		stdin, err := call.Argument(1).ToString()
		if err != nil {
			panic(err)
		}
		for i := 2; i < len(call.ArgumentList); i += 1 {
			ret, err := call.Argument(i).ToString()
			if err == nil {
				args[i] = ret
			} else {
				panic(err)
			}

		}
		output := callExternalFuncWithStdin(cmdPath, stdin, args)
		result, err := vm.ToValue(output)
		if err != nil {
			panic(err)
		}

		return result
	})
	vm.Set("print", func(call otto.FunctionCall) otto.Value {
		for i := 0; i < len(call.ArgumentList); i += 1 {
			str, err := call.Argument(i).ToString()
			if err == nil {
				fmt.Print(str)
			} else {
				panic(err)
			}
		}
		val, _ := vm.ToValue("")
		return val
	})

	sprintf := func(call otto.FunctionCall) string {
		args := make([]interface{}, len(call.ArgumentList)-1)
		format, err := call.Argument(0).ToString()
		if err != nil {
			panic(err)
		}
		for i := 1; i < len(call.ArgumentList); i += 1 {
			str, err := call.Argument(i).ToString()
			if err == nil {
				args[i-1] = str
			} else {
				panic(err)
			}
		}
		return fmt.Sprintf(format, args...)
	}

	vm.Set("sprint", func(call otto.FunctionCall) otto.Value {
		val := sprintf(call)
		ret, _ := vm.ToValue(val)
		return ret
	})
	vm.Set("printf", func(call otto.FunctionCall) otto.Value {
		val := sprintf(call)
		fmt.Print(val)
		ret, _ := vm.ToValue(val)
		return ret
	})
}

func callExternalFunc(cmd string, args []string) string {
	out, err := exec.Command(cmd, args...).Output()
	if err != nil {
		panic(err)
	}
	return string(out)
}

func callExternalFuncWithStdin(cmd string, stdin string, args []string) string {
	p := exec.Command(cmd, args...)
	p.Stdin = strings.NewReader(stdin)
	out, err := p.Output()
	if err != nil {
		panic(err)
	}
	return string(out)
}

func addPATHEnv(path string) {
	oldEnv := os.Getenv("PATH")
	newEnv := path + ":" + oldEnv
	os.Setenv("PATH", newEnv)
}

func getWd() string {
	cwd, err := os.Getwd()
	if err != nil {
		panic(err)
	}
	return cwd
}

func initExternelFunc(vm *otto.Otto, cmdPath string) {
	funcName := cmdPath[:len(cmdPath)-len(filepath.Ext(cmdPath))]
	vm.Set(funcName, func(call otto.FunctionCall) otto.Value {
		args := make([]string, len(call.ArgumentList))
		for i := 0; i < len(call.ArgumentList); i += 1 {
			ret, err := call.Argument(i).ToString()
			if err == nil {
				args[i] = ret
			} else {
				panic(err)
			}

		}
		output := callExternalFunc(cmdPath, args)
		result, err := vm.ToValue(output)
		if err != nil {
			panic(err)
		}
		return result
	})

}

func main() {
	addPATHEnv(getWd())
	fmt.Printf("")
	kingpin.Parse()

	var stdin string

	if *nostdinPtr {
		stdin = ""
	}else{
		stdin = readAll()
	}

	lineSpliter := regexp.MustCompile(*lineSeqPtr)
	columnSpilter := regexp.MustCompile(*columnSeqPtr)

	lines := lineSpliter.Split(stdin, -1)
	matrixPtr := new(Matrix)

	for i := 0; i < len(lines); i += 1 {
		*matrixPtr = append(*matrixPtr, columnSpilter.Split(lines[i], -1))
	}

	vm := otto.New()

	vm.Set("stdin", stdin)
	vm.Set("stdout", "")
	vm.Set("lines", lines)
	vm.Set("matrix", *matrixPtr)
	vm.Run("println = console.log")
	setBuiltinFunc(vm)

	for i := 0; i < len(*funcListPtr); i += 1 {
		initExternelFunc(vm, (*funcListPtr)[i])
	}

	for i := 0; i < len(*pathListPtr); i += 1 {
		addPATHEnv((*pathListPtr)[i])
	}

	for i := 0; i < len(*jsListPtr); i += 1 {
		readJSFile(vm, (*jsListPtr)[i])
	}

	codeList := sort.StringSlice(*codeListPtr)

	if *reversePtr {
		for i, j := 0, len(codeList)-1; i < j; i, j = i+1, j-1 {
			codeList[i], codeList[j] = codeList[j], codeList[i]
		}
	}

	for i := 0; i < len(codeList); i += 1 {
		_, err := vm.Run(codeList[i])
		if err != nil {
			panic(err)
		}
	}

}
<|endoftext|>


<commit_before>package junos

import (
	"encoding/xml"
	"errors"
	"fmt"
	"github.com/Juniper/go-netconf/netconf"
	"log"
)

// Session holds the connection information to our Junos device.
type Session struct {
	Conn *netconf.Session
}

// rollbackXML parses our rollback configuration.
type rollbackXML struct {
	XMLName xml.Name `xml:"rollback-information"`
	Config  string   `xml:"configuration-information>configuration-output"`
}

// RescueXML parses our rescue configuration.
type rescueXML struct {
	XMLName xml.Name `xml:"rescue-information"`
	Config  string   `xml:"configuration-information>configuration-output"`
}

// CommandXML parses our operational command responses.
type commandXML struct {
	Config string `xml:",innerxml"`
}

// NewSession establishes a new connection to a Junos device that we will use
// to run our commands against.
func NewSession(host, user, password string) *Session {
	s, err := netconf.DialSSH(host, netconf.SSHConfigPassword(user, password))
	if err != nil {
		log.Fatal(err)
	}

	return &Session{
		Conn: s,
	}
}

// Lock locks the candidate configuration.
func (s *Session) Lock() error {
	resp, err := s.Conn.Exec(rpcCommand["lock"])
	if err != nil {
		log.Fatal(err)
	}

	if resp.Ok == false {
		for _, m := range resp.Errors {
			return errors.New(m.Message)
		}
	}

	return nil
}

// Unlock unlocks the candidate configuration.
func (s *Session) Unlock() error {
	resp, err := s.Conn.Exec(rpcCommand["unlock"])
	if err != nil {
		log.Fatal(err)
	}

	if resp.Ok == false {
		for _, m := range resp.Errors {
			return errors.New(m.Message)
		}
	}

	return nil
}

// GetRollbackConfig returns the configuration of the given rollback state.
func (s *Session) GetRollbackConfig(number int) (string, error) {
	rb := &rollbackXML{}
	command := fmt.Sprintf(rpcCommand["get-rollback-information"], number)
	reply, err := s.Conn.Exec(command)

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

	if reply.Ok == false {
		for _, m := range reply.Errors {
			return "", errors.New(m.Message)
		}
	}

	err = xml.Unmarshal([]byte(reply.Data), rb)
	if err != nil {
		log.Fatal(err)
	}

	return rb.Config, nil
}

// RollbackDiff compares the current active configuration to a given rollback configuration.
func (s *Session) RollbackDiff(compare int) (string, error) {
	rb := &rollbackXML{}
	command := fmt.Sprintf(rpcCommand["get-rollback-information-compare"], compare)
	reply, err := s.Conn.Exec(command)

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

	if reply.Ok == false {
		for _, m := range reply.Errors {
			return "", errors.New(m.Message)
		}
	}

	err = xml.Unmarshal([]byte(reply.Data), rb)
	if err != nil {
		log.Fatal(err)
	}

	return rb.Config, nil
}

// GetRescueConfig returns the rescue configuration.
func (s *Session) GetRescueConfig() (string, error) {
	rescue := &rescueXML{}
	reply, err := s.Conn.Exec(rpcCommand["get-rescue-information"])

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

	if reply.Ok == false {
		for _, m := range reply.Errors {
			return "", errors.New(m.Message)
		}
	}

	err = xml.Unmarshal([]byte(reply.Data), rescue)
	if err != nil {
		log.Fatal(err)
	}

	if rescue.Config == "" {
		return "No rescue configuration set.", nil
	}

	return rescue.Config, nil
}

// Command runs any operational mode command, such as "show", "request", etc..
// format is either "text" or "xml".
func (s *Session) Command(cmd, format string) (string, error) {
	c := &commandXML{}
	var command string

	switch format {
	case "xml":
		command = fmt.Sprintf(rpcCommand["command-xml"], cmd)
	default:
		command = fmt.Sprintf(rpcCommand["command"], cmd)
	}
	reply, err := s.Conn.Exec(command)
	if err != nil {
		log.Fatal(err)
	}

	if reply.Ok == false {
		for _, m := range reply.Errors {
			return "", errors.New(m.Message)
		}
	}

	err = xml.Unmarshal([]byte(reply.Data), &c)
	if err != nil {
		log.Fatal(err)
	}

	if c.Config == "" {
		return "No output available.", nil
	}

	return c.Config, nil
}

// Close disconnects and closes the session to our Junos device.
func (s *Session) Close() {
	s.Conn.Close()
}
<commit_msg>Updated documentation<commit_after>package junos

import (
	"encoding/xml"
	"errors"
	"fmt"
	"github.com/Juniper/go-netconf/netconf"
	"log"
)

// Session holds the connection information to our Junos device.
type Session struct {
	Conn *netconf.Session
}

// rollbackXML parses our rollback configuration.
type rollbackXML struct {
	XMLName xml.Name `xml:"rollback-information"`
	Config  string   `xml:"configuration-information>configuration-output"`
}

// RescueXML parses our rescue configuration.
type rescueXML struct {
	XMLName xml.Name `xml:"rescue-information"`
	Config  string   `xml:"configuration-information>configuration-output"`
}

// CommandXML parses our operational command responses.
type commandXML struct {
	Config string `xml:",innerxml"`
}

// NewSession establishes a new connection to a Junos device that we will use
// to run our commands against.
func NewSession(host, user, password string) *Session {
	s, err := netconf.DialSSH(host, netconf.SSHConfigPassword(user, password))
	if err != nil {
		log.Fatal(err)
	}

	return &Session{
		Conn: s,
	}
}

// Lock locks the candidate configuration.
func (s *Session) Lock() error {
	resp, err := s.Conn.Exec(rpcCommand["lock"])
	if err != nil {
		log.Fatal(err)
	}

	if resp.Ok == false {
		for _, m := range resp.Errors {
			return errors.New(m.Message)
		}
	}

	return nil
}

// Unlock unlocks the candidate configuration.
func (s *Session) Unlock() error {
	resp, err := s.Conn.Exec(rpcCommand["unlock"])
	if err != nil {
		log.Fatal(err)
	}

	if resp.Ok == false {
		for _, m := range resp.Errors {
			return errors.New(m.Message)
		}
	}

	return nil
}

// GetRollbackConfig returns the configuration of the given rollback state.
func (s *Session) GetRollbackConfig(number int) (string, error) {
	rb := &rollbackXML{}
	command := fmt.Sprintf(rpcCommand["get-rollback-information"], number)
	reply, err := s.Conn.Exec(command)

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

	if reply.Ok == false {
		for _, m := range reply.Errors {
			return "", errors.New(m.Message)
		}
	}

	err = xml.Unmarshal([]byte(reply.Data), rb)
	if err != nil {
		log.Fatal(err)
	}

	return rb.Config, nil
}

// RollbackDiff compares the current active configuration to a given rollback configuration.
func (s *Session) RollbackDiff(compare int) (string, error) {
	rb := &rollbackXML{}
	command := fmt.Sprintf(rpcCommand["get-rollback-information-compare"], compare)
	reply, err := s.Conn.Exec(command)

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

	if reply.Ok == false {
		for _, m := range reply.Errors {
			return "", errors.New(m.Message)
		}
	}

	err = xml.Unmarshal([]byte(reply.Data), rb)
	if err != nil {
		log.Fatal(err)
	}

	return rb.Config, nil
}

// GetRescueConfig returns the rescue configuration.
func (s *Session) GetRescueConfig() (string, error) {
	rescue := &rescueXML{}
	reply, err := s.Conn.Exec(rpcCommand["get-rescue-information"])

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

	if reply.Ok == false {
		for _, m := range reply.Errors {
			return "", errors.New(m.Message)
		}
	}

	err = xml.Unmarshal([]byte(reply.Data), rescue)
	if err != nil {
		log.Fatal(err)
	}

	if rescue.Config == "" {
		return "No rescue configuration set.", nil
	}

	return rescue.Config, nil
}

// Command runs any operational mode command, such as "show" or "request."
// Format is either "text" or "xml".
func (s *Session) Command(cmd, format string) (string, error) {
	c := &commandXML{}
	var command string

	switch format {
	case "xml":
		command = fmt.Sprintf(rpcCommand["command-xml"], cmd)
	default:
		command = fmt.Sprintf(rpcCommand["command"], cmd)
	}
	reply, err := s.Conn.Exec(command)
	if err != nil {
		log.Fatal(err)
	}

	if reply.Ok == false {
		for _, m := range reply.Errors {
			return "", errors.New(m.Message)
		}
	}

	err = xml.Unmarshal([]byte(reply.Data), &c)
	if err != nil {
		log.Fatal(err)
	}

	if c.Config == "" {
		return "No output available.", nil
	}

	return c.Config, nil
}

// Close disconnects and closes the session to our Junos device.
func (s *Session) Close() {
	s.Conn.Close()
}
<|endoftext|>


<commit_before>package pg_test

import (
	"bytes"
	"context"
	"fmt"
	"strings"
	"sync"
	"time"

	"github.com/go-pg/pg/v10"
	"github.com/go-pg/pg/v10/orm"
)

var (
	ctx  = context.Background()
	pgdb *pg.DB
)

func init() {
	pgdb = connect()

	err := pgdb.Model((*Flight)(nil)).DropTable(&orm.DropTableOptions{
		IfExists: true,
		Cascade:  true,
	})
	panicIf(err)

	err = pgdb.Model((*Flight)(nil)).CreateTable(nil)
	panicIf(err)
}

func connect() *pg.DB {
	return pg.Connect(pgOptions())
}

func panicIf(err error) {
	if err != nil {
		panic(err)
	}
}

func ExampleConnect() {
	db := pg.Connect(&pg.Options{
		User:     "postgres",
		Password: "",
		Database: "postgres",
	})
	defer db.Close()

	var n int
	_, err := db.QueryOne(pg.Scan(&n), "SELECT 1")
	panicIf(err)
	fmt.Println(n)
	// Output: 1
}

func ExampleDB_QueryOne() {
	var user struct {
		Name string
	}

	res, err := pgdb.QueryOne(&user, `
        WITH users (name) AS (VALUES (?))
        SELECT * FROM users
    `, "admin")
	panicIf(err)
	fmt.Println(res.RowsAffected())
	fmt.Println(user)
	// Output: 1
	// {admin}
}

func ExampleDB_QueryOne_returning_id() {
	_, err := pgdb.Exec(`CREATE TEMP TABLE users(id serial, name varchar(500))`)
	panicIf(err)

	var user struct {
		Id   int32
		Name string
	}
	user.Name = "admin"

	_, err = pgdb.QueryOne(&user, `
        INSERT INTO users (name) VALUES (?name) RETURNING id
    `, &user)
	panicIf(err)
	fmt.Println(user)
	// Output: {1 admin}
}

func ExampleDB_Exec() {
	res, err := pgdb.Exec(`CREATE TEMP TABLE test()`)
	panicIf(err)
	fmt.Println(res.RowsAffected())
	// Output: -1
}

func ExampleListener() {
	ln := pgdb.Listen(ctx, "mychan")
	defer ln.Close()

	ch := ln.Channel()

	go func() {
		time.Sleep(time.Millisecond)
		_, err := pgdb.Exec("NOTIFY mychan, ?", "hello world")
		panicIf(err)
	}()

	notif := <-ch
	fmt.Println(notif)
	// Output: &{mychan hello world}
}

func txExample() *pg.DB {
	db := pg.Connect(&pg.Options{
		User: "postgres",
	})

	queries := []string{
		`DROP TABLE IF EXISTS tx_test`,
		`CREATE TABLE tx_test(counter int)`,
		`INSERT INTO tx_test (counter) VALUES (0)`,
	}
	for _, q := range queries {
		_, err := db.Exec(q)
		panicIf(err)
	}

	return db
}

func ExampleDB_Begin() {
	db := txExample()

	incrInTx := func(db *pg.DB) error {
		tx, err := db.Begin()
		if err != nil {
			return err
		}
		// Rollback tx on error.
		defer tx.Rollback()

		var counter int
		_, err = tx.QueryOne(
			pg.Scan(&counter), `SELECT counter FROM tx_test FOR UPDATE`)
		if err != nil {
			return err
		}

		counter++

		_, err = tx.Exec(`UPDATE tx_test SET counter = ?`, counter)
		if err != nil {
			return err
		}

		return tx.Commit()
	}

	var wg sync.WaitGroup
	for i := 0; i < 10; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			err := incrInTx(db)
			panicIf(err)
		}()
	}
	wg.Wait()

	var counter int
	_, err := db.QueryOne(pg.Scan(&counter), `SELECT counter FROM tx_test`)
	panicIf(err)
	fmt.Println(counter)
	// Output: 10
}

func ExampleDB_RunInTransaction() {
	db := txExample()

	incrInTx := func(db *pg.DB) error {
		// Transaction is automatically rolled back on error.
		return db.RunInTransaction(ctx, func(tx *pg.Tx) error {
			var counter int
			_, err := tx.QueryOne(
				pg.Scan(&counter), `SELECT counter FROM tx_test FOR UPDATE`)
			if err != nil {
				return err
			}

			counter++

			_, err = tx.Exec(`UPDATE tx_test SET counter = ?`, counter)
			return err
		})
	}

	var wg sync.WaitGroup
	for i := 0; i < 10; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			err := incrInTx(db)
			panicIf(err)
		}()
	}
	wg.Wait()

	var counter int
	_, err := db.QueryOne(pg.Scan(&counter), `SELECT counter FROM tx_test`)
	panicIf(err)
	fmt.Println(counter)
	// Output: 10
}

func ExampleDB_Prepare() {
	stmt, err := pgdb.Prepare(`SELECT $1::text, $2::text`)
	panicIf(err)

	var s1, s2 string
	_, err = stmt.QueryOne(pg.Scan(&s1, &s2), "foo", "bar")
	panicIf(err)
	fmt.Println(s1, s2)
	// Output: foo bar
}

func ExampleDB_CreateTable() {
	type Model1 struct {
		Id int
	}

	type Model2 struct {
		Id   int
		Name string

		Model1Id int `pg:"on_delete:RESTRICT, on_update: CASCADE"`
		Model1   *Model1
	}

	for _, model := range []interface{}{&Model1{}, &Model2{}} {
		err := pgdb.Model(model).CreateTable(&orm.CreateTableOptions{
			Temp:          true, // create temp table
			FKConstraints: true,
		})
		panicIf(err)
	}

	var info []struct {
		ColumnName string
		DataType   string
	}
	_, err := pgdb.Query(&info, `
		SELECT column_name, data_type
		FROM information_schema.columns
		WHERE table_name = 'model2'
	`)
	panicIf(err)
	fmt.Println(info)
	// Output: [{id bigint} {name text} {model1_id bigint}]
}

func ExampleInts() {
	var nums pg.Ints
	_, err := pgdb.Query(&nums, `SELECT generate_series(0, 10)`)
	panicIf(err)
	fmt.Println(nums)
	// Output: [0 1 2 3 4 5 6 7 8 9 10]
}

func ExampleStrings() {
	var strs pg.Strings
	_, err := pgdb.Query(&strs, `
		WITH users AS (VALUES ('foo'), ('bar')) SELECT * FROM users
	`)
	panicIf(err)
	fmt.Println(strs)
	// Output: [foo bar]
}

func ExampleDB_CopyFrom() {
	_, err := pgdb.Exec(`CREATE TEMP TABLE words(word text, len int)`)
	panicIf(err)

	r := strings.NewReader("hello,5\nfoo,3\n")
	_, err = pgdb.CopyFrom(r, `COPY words FROM STDIN WITH CSV`)
	panicIf(err)

	var buf bytes.Buffer
	_, err = pgdb.CopyTo(&buf, `COPY words TO STDOUT WITH CSV`)
	panicIf(err)
	fmt.Println(buf.String())
	// Output: hello,5
	// foo,3
}

func ExampleDB_WithTimeout() {
	var count int
	// Use bigger timeout since this query is known to be slow.
	_, err := pgdb.WithTimeout(time.Minute).QueryOne(pg.Scan(&count), `
		SELECT count(*) FROM big_table
	`)
	panicIf(err)
}

func ExampleScan() {
	var s1, s2 string
	_, err := pgdb.QueryOne(pg.Scan(&s1, &s2), `SELECT ?, ?`, "foo", "bar")
	panicIf(err)
	fmt.Println(s1, s2)
	// Output: foo bar
}

func ExampleError() {
	flight := &Flight{
		Id: 123,
	}
	_, err := pgdb.Model(flight).Insert()
	panicIf(err)

	_, err = pgdb.Model(flight).Insert()
	if err != nil {
		pgErr, ok := err.(pg.Error)
		if ok && pgErr.IntegrityViolation() {
			fmt.Println("flight already exists:", err)
		} else {
			panic(err)
		}
	}
	// Output: flight already exists: ERROR #23505 duplicate key value violates unique constraint "flights_pkey"
}
<commit_msg>Fix test<commit_after>package pg_test

import (
	"bytes"
	"context"
	"fmt"
	"strings"
	"sync"
	"time"

	"github.com/go-pg/pg/v10"
	"github.com/go-pg/pg/v10/orm"
)

var (
	ctx  = context.Background()
	pgdb *pg.DB
)

func init() {
	pgdb = connect()

	err := pgdb.Model((*Flight)(nil)).DropTable(&orm.DropTableOptions{
		IfExists: true,
		Cascade:  true,
	})
	panicIf(err)

	err = pgdb.Model((*Flight)(nil)).CreateTable(nil)
	panicIf(err)
}

func connect() *pg.DB {
	return pg.Connect(pgOptions())
}

func panicIf(err error) {
	if err != nil {
		panic(err)
	}
}

func ExampleConnect() {
	db := pg.Connect(&pg.Options{
		User:     "postgres",
		Password: "",
		Database: "postgres",
	})
	defer db.Close()

	var n int
	_, err := db.QueryOne(pg.Scan(&n), "SELECT 1")
	panicIf(err)
	fmt.Println(n)
	// Output: 1
}

func ExampleDB_QueryOne() {
	var user struct {
		Name string
	}

	res, err := pgdb.QueryOne(&user, `
        WITH users (name) AS (VALUES (?))
        SELECT * FROM users
    `, "admin")
	panicIf(err)
	fmt.Println(res.RowsAffected())
	fmt.Println(user)
	// Output: 1
	// {admin}
}

func ExampleDB_QueryOne_returning_id() {
	_, err := pgdb.Exec(`CREATE TEMP TABLE users(id serial, name varchar(500))`)
	panicIf(err)

	var user struct {
		Id   int32
		Name string
	}
	user.Name = "admin"

	_, err = pgdb.QueryOne(&user, `
        INSERT INTO users (name) VALUES (?name) RETURNING id
    `, &user)
	panicIf(err)
	fmt.Println(user)
	// Output: {1 admin}
}

func ExampleDB_Exec() {
	res, err := pgdb.Exec(`CREATE TEMP TABLE test()`)
	panicIf(err)
	fmt.Println(res.RowsAffected())
	// Output: -1
}

func ExampleListener() {
	ln := pgdb.Listen(ctx, "mychan")
	defer ln.Close()

	ch := ln.Channel()

	go func() {
		time.Sleep(time.Millisecond)
		_, err := pgdb.Exec("NOTIFY mychan, ?", "hello world")
		panicIf(err)
	}()

	notif := <-ch
	fmt.Println(notif)
	// Output: {mychan hello world}
}

func txExample() *pg.DB {
	db := pg.Connect(&pg.Options{
		User: "postgres",
	})

	queries := []string{
		`DROP TABLE IF EXISTS tx_test`,
		`CREATE TABLE tx_test(counter int)`,
		`INSERT INTO tx_test (counter) VALUES (0)`,
	}
	for _, q := range queries {
		_, err := db.Exec(q)
		panicIf(err)
	}

	return db
}

func ExampleDB_Begin() {
	db := txExample()

	incrInTx := func(db *pg.DB) error {
		tx, err := db.Begin()
		if err != nil {
			return err
		}
		// Rollback tx on error.
		defer tx.Rollback()

		var counter int
		_, err = tx.QueryOne(
			pg.Scan(&counter), `SELECT counter FROM tx_test FOR UPDATE`)
		if err != nil {
			return err
		}

		counter++

		_, err = tx.Exec(`UPDATE tx_test SET counter = ?`, counter)
		if err != nil {
			return err
		}

		return tx.Commit()
	}

	var wg sync.WaitGroup
	for i := 0; i < 10; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			err := incrInTx(db)
			panicIf(err)
		}()
	}
	wg.Wait()

	var counter int
	_, err := db.QueryOne(pg.Scan(&counter), `SELECT counter FROM tx_test`)
	panicIf(err)
	fmt.Println(counter)
	// Output: 10
}

func ExampleDB_RunInTransaction() {
	db := txExample()

	incrInTx := func(db *pg.DB) error {
		// Transaction is automatically rolled back on error.
		return db.RunInTransaction(ctx, func(tx *pg.Tx) error {
			var counter int
			_, err := tx.QueryOne(
				pg.Scan(&counter), `SELECT counter FROM tx_test FOR UPDATE`)
			if err != nil {
				return err
			}

			counter++

			_, err = tx.Exec(`UPDATE tx_test SET counter = ?`, counter)
			return err
		})
	}

	var wg sync.WaitGroup
	for i := 0; i < 10; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			err := incrInTx(db)
			panicIf(err)
		}()
	}
	wg.Wait()

	var counter int
	_, err := db.QueryOne(pg.Scan(&counter), `SELECT counter FROM tx_test`)
	panicIf(err)
	fmt.Println(counter)
	// Output: 10
}

func ExampleDB_Prepare() {
	stmt, err := pgdb.Prepare(`SELECT $1::text, $2::text`)
	panicIf(err)

	var s1, s2 string
	_, err = stmt.QueryOne(pg.Scan(&s1, &s2), "foo", "bar")
	panicIf(err)
	fmt.Println(s1, s2)
	// Output: foo bar
}

func ExampleDB_CreateTable() {
	type Model1 struct {
		Id int
	}

	type Model2 struct {
		Id   int
		Name string

		Model1Id int `pg:"on_delete:RESTRICT, on_update: CASCADE"`
		Model1   *Model1
	}

	for _, model := range []interface{}{&Model1{}, &Model2{}} {
		err := pgdb.Model(model).CreateTable(&orm.CreateTableOptions{
			Temp:          true, // create temp table
			FKConstraints: true,
		})
		panicIf(err)
	}

	var info []struct {
		ColumnName string
		DataType   string
	}
	_, err := pgdb.Query(&info, `
		SELECT column_name, data_type
		FROM information_schema.columns
		WHERE table_name = 'model2'
	`)
	panicIf(err)
	fmt.Println(info)
	// Output: [{id bigint} {name text} {model1_id bigint}]
}

func ExampleInts() {
	var nums pg.Ints
	_, err := pgdb.Query(&nums, `SELECT generate_series(0, 10)`)
	panicIf(err)
	fmt.Println(nums)
	// Output: [0 1 2 3 4 5 6 7 8 9 10]
}

func ExampleStrings() {
	var strs pg.Strings
	_, err := pgdb.Query(&strs, `
		WITH users AS (VALUES ('foo'), ('bar')) SELECT * FROM users
	`)
	panicIf(err)
	fmt.Println(strs)
	// Output: [foo bar]
}

func ExampleDB_CopyFrom() {
	_, err := pgdb.Exec(`CREATE TEMP TABLE words(word text, len int)`)
	panicIf(err)

	r := strings.NewReader("hello,5\nfoo,3\n")
	_, err = pgdb.CopyFrom(r, `COPY words FROM STDIN WITH CSV`)
	panicIf(err)

	var buf bytes.Buffer
	_, err = pgdb.CopyTo(&buf, `COPY words TO STDOUT WITH CSV`)
	panicIf(err)
	fmt.Println(buf.String())
	// Output: hello,5
	// foo,3
}

func ExampleDB_WithTimeout() {
	var count int
	// Use bigger timeout since this query is known to be slow.
	_, err := pgdb.WithTimeout(time.Minute).QueryOne(pg.Scan(&count), `
		SELECT count(*) FROM big_table
	`)
	panicIf(err)
}

func ExampleScan() {
	var s1, s2 string
	_, err := pgdb.QueryOne(pg.Scan(&s1, &s2), `SELECT ?, ?`, "foo", "bar")
	panicIf(err)
	fmt.Println(s1, s2)
	// Output: foo bar
}

func ExampleError() {
	flight := &Flight{
		Id: 123,
	}
	_, err := pgdb.Model(flight).Insert()
	panicIf(err)

	_, err = pgdb.Model(flight).Insert()
	if err != nil {
		pgErr, ok := err.(pg.Error)
		if ok && pgErr.IntegrityViolation() {
			fmt.Println("flight already exists:", err)
		} else {
			panic(err)
		}
	}
	// Output: flight already exists: ERROR #23505 duplicate key value violates unique constraint "flights_pkey"
}
<|endoftext|>


<commit_before>// Copyright 2016 <chaishushan{AT}gmail.com>. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package ptyy_test

import (
	"fmt"

	"github.com/chai2010/ptyy"
)

func ExampleSearch_wuhan() {
	for _, v := range ptyy.Search("武汉", 0) {
		fmt.Printf("%s (%s)\n", v.Name, v.City)
	}

	// Output:
	// 武汉华仁医院 (武汉)
	// 武汉华夏医院 (武汉)
	// 武汉华美医院 (武汉)
	// 武汉当代佳丽医院 (武汉)
	// 武汉现代女子妇科医院 (武汉)
	// 武汉阳光女子医院 (武汉)
	// 湖北荣军医院 (武汉)
}

func ExampleSearch_wuhan_mimit3() {
	for _, v := range ptyy.Search("武汉", 3) {
		fmt.Printf("%s (%s)\n", v.Name, v.City)
	}

	// Output:
	// 武汉华仁医院 (武汉)
	// 武汉华夏医院 (武汉)
	// 武汉华美医院 (武汉)
}
<commit_msg>fix example code<commit_after>// Copyright 2016 <chaishushan{AT}gmail.com>. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package ptyy_test

import (
	"fmt"

	"github.com/chai2010/ptyy"
)

func ExampleSearch_wuhan() {
	for _, v := range ptyy.Search("武汉", 0) {
		fmt.Printf("%s (%s)\n", v.Name, v.City)
	}

	// Output:
	// 武汉华仁医院 (武汉)
	// 武汉华夏医院 (武汉)
	// 武汉华美医院 (武汉)
	// 武汉当代佳丽医院 (武汉)
	// 武汉现代女子妇科医院 (武汉)
	// 武汉阳光女子医院 (武汉)
	// 湖北荣军医院 (武汉)
}

func ExampleSearch_wuhanLimit3() {
	for _, v := range ptyy.Search("武汉", 3) {
		fmt.Printf("%s (%s)\n", v.Name, v.City)
	}

	// Output:
	// 武汉华仁医院 (武汉)
	// 武汉华夏医院 (武汉)
	// 武汉华美医院 (武汉)
}
<|endoftext|>


<commit_before>package devastator_test

import (
	"log"

	"github.com/nbusy/devastator"
)

const debug = false

// Example demonstrating the use of Devastator server.
func Example() {
	s, err := devastator.NewServer(nil, nil, "", debug)
	if err == nil && s != nil {
		log.Println("Connected")
	}
}
<commit_msg>execute and verify example code<commit_after>package devastator_test

import (
	"log"

	"github.com/nbusy/devastator"
)

const debug = false

// Example demonstrating the use of Devastator server.
func Example() {
	s, err := devastator.NewServer(nil, nil, "", debug)
	if err == nil && s != nil {
		log.Println("Connected")
		// Output: Connected
	}
}
<|endoftext|>


<commit_before>package plugin_test

import (
	"log"
	"os"
	"strings"

	"net/rpc/jsonrpc"

	"github.com/natefinch/plugin"
)

// This function should be called from the master program that wants to run
// plugins to extend its functionality.
//
// StartWithCodec starts a plugin at path "foo", using the JSON-RPC codec, and
// writing its output to this application's Stderr.  The application can
// then call methods on the rpc client returned using the standard rpc
// pattern.
func ExampleStartWithCodec() {
	foo, err := plugin.StartWithCodec(jsonrpc.NewClient, "/var/lib/foo", os.Stderr)
	if err != nil {
		log.Fatalf("failed to load foo plugin: %s", err)
	}
	var reply string
	foo.Call("Foo.ToUpper", "something", &reply)
}

// This function should be called from the plugin program that wants to provide
// functionality for the master program.
//
// ProvideCodec starts an RPC server that reads from stdin and writes to stdout.
// It provides functions attached to the API value passed in.  This function
// will block forever, so it is common to simply put this at theend of the
// plugin's main function.
func ExampleNewServerWithCodec() {
	p := plugin.NewServerWithCodec(jsonrpc.NewServerCodec)
	if err := p.RegisterName("Foo", API{}); err != nil {
		log.Fatalf("can't register api: %s", err)
	}
	p.Serve()
}

// API is an example type to show how to serve methods over RPC.
type API struct{}

// ToUpper is an example function that gets served over RPC.  See net/rpc for
// details on how to server functionality over RPC.
func (API) ToUpper(input string, output *string) error {
	*output = strings.ToUpper(input)
	return nil
}
<commit_msg>typos<commit_after>package plugin_test

import (
	"log"
	"os"
	"strings"

	"net/rpc/jsonrpc"

	"github.com/natefinch/plugin"
)

// This function should be called from the master program that wants to run
// plugins to extend its functionality.
//
// StartWithCodec starts a plugin at path "foo", using the JSON-RPC codec, and
// writing its output to this application's Stderr.  The application can
// then call methods on the rpc client returned using the standard rpc
// pattern.
func ExampleStartWithCodec() {
	foo, err := plugin.StartWithCodec(jsonrpc.NewClient, "/var/lib/foo", os.Stderr)
	if err != nil {
		log.Fatalf("failed to load foo plugin: %s", err)
	}
	var reply string
	foo.Call("Foo.ToUpper", "something", &reply)
}

// This function should be called from the plugin program that wants to provide
// functionality for the master program.
//
// NewServerWithCodec starts an RPC server that reads from stdin and writes to
// stdout. It provides functions attached to the API value passed in.
// Server.Serve() will block forever, so it is common to simply put this at the
// end of the plugin's main function.
func ExampleNewServerWithCodec() {
	p := plugin.NewServerWithCodec(jsonrpc.NewServerCodec)
	if err := p.RegisterName("Foo", API{}); err != nil {
		log.Fatalf("can't register api: %s", err)
	}
	p.Serve()
}

// API is an example type to show how to serve methods over RPC.
type API struct{}

// ToUpper is an example function that gets served over RPC.  See net/rpc for
// details on how to server functionality over RPC.
func (API) ToUpper(input string, output *string) error {
	*output = strings.ToUpper(input)
	return nil
}
<|endoftext|>


<commit_before>package gomail_test

import (
	"fmt"
	"html/template"
	"io"
	"log"
	"time"

	"gopkg.in/gomail.v2"
)

func Example() {
	m := gomail.NewMessage()
	m.SetHeader("From", "alex@example.com")
	m.SetHeader("To", "bob@example.com", "cora@example.com")
	m.SetAddressHeader("Cc", "dan@example.com", "Dan")
	m.SetHeader("Subject", "Hello!")
	m.SetBody("text/html", "Hello Bob and Cora!")
	m.Attach("/home/Alex/lolcat.jpg")

	d := gomail.NewPlainDialer("smtp.example.com", 587, "user", "123456")

	// Send the email to Bob, Cora and Dan.
	if err := d.DialAndSend(m); err != nil {
		panic(err)
	}
}

// A daemon that listens to a channel and sends all incoming messages.
func Example_daemon() {
	ch := make(chan *gomail.Message)

	go func() {
		d := gomail.NewPlainDialer("smtp.example.com", 587, "user", "123456")

		var s gomail.SendCloser
		var err error
		open := false
		for {
			select {
			case m, ok := <-ch:
				if !ok {
					return
				}
				if !open {
					if s, err = d.Dial(); err != nil {
						panic(err)
					}
					open = true
				}
				if err := gomail.Send(s, m); err != nil {
					log.Print(err)
				}
			// Close the connection to the SMTP server if no email was sent in
			// the last 30 seconds.
			case <-time.After(30 * time.Second):
				if open {
					if err := s.Close(); err != nil {
						panic(err)
					}
					open = false
				}
			}
		}
	}()

	// Use the channel in your program to send emails.

	// Close the channel to stop the mail daemon.
	close(ch)
}

// Efficiently send a customized newsletter to a list of recipients.
func Example_newsletter() {
	// The list of recipients.
	var list []struct {
		Name    string
		Address string
	}

	d := gomail.NewPlainDialer("smtp.example.com", 587, "user", "123456")
	s, err := d.Dial()
	if err != nil {
		panic(err)
	}

	m := gomail.NewMessage()
	for _, r := range list {
		m.SetHeader("From", "no-reply@example.com")
		m.SetAddressHeader("To", r.Address, r.Name)
		m.SetHeader("Subject", "Newsletter #1")
		m.SetBody("text/html", fmt.Sprintf("Hello %s!", r.Name))

		if err := gomail.Send(s, m); err != nil {
			log.Printf("Could not send email to %q: %v", r.Address, err)
		}
		m.Reset()
	}
}

// Send an email using a local SMTP server.
func Example_noAuth() {
	m := gomail.NewMessage()
	m.SetHeader("From", "from@example.com")
	m.SetHeader("To", "to@example.com")
	m.SetHeader("Subject", "Hello!")
	m.SetBody("text/plain", "Hello!")

	d := gomail.Dialer{Host: "localhost", Port: 587}
	if err := d.DialAndSend(m); err != nil {
		panic(err)
	}
}

// Send an email using an API or postfix.
func Example_noSMTP() {
	m := gomail.NewMessage()
	m.SetHeader("From", "from@example.com")
	m.SetHeader("To", "to@example.com")
	m.SetHeader("Subject", "Hello!")
	m.SetBody("text/plain", "Hello!")

	s := gomail.SendFunc(func(from string, to []string, msg io.WriterTo) error {
		// Implements you email-sending function, for example by calling
		// an API, or running postfix, etc.
		fmt.Println("From:", from)
		fmt.Println("To:", to)
		return nil
	})

	if err := gomail.Send(s, m); err != nil {
		panic(err)
	}
	// Output:
	// From: from@example.com
	// To: [to@example.com]
}

var m *gomail.Message

func ExampleSetCopyFunc() {
	m.Attach("foo.txt", gomail.SetCopyFunc(func(w io.Writer) error {
		_, err := w.Write([]byte("Content of foo.txt"))
		return err
	}))
}

func ExampleSetHeader() {
	h := map[string][]string{"Content-ID": {"<foo@bar.mail>"}}
	m.Attach("foo.jpg", gomail.SetHeader(h))
}

func ExampleMessage_AddAlternative() {
	m.SetBody("text/plain", "Hello!")
	m.AddAlternative("text/html", "
Hello!
")
}

func ExampleMessage_AddAlternativeWriter() {
	t := template.Must(template.New("example").Parse("Hello {{.}}!"))
	m.AddAlternativeWriter("text/plain", func(w io.Writer) error {
		return t.Execute(w, "Bob")
	})
}

func ExampleMessage_Attach() {
	m.Attach("/tmp/image.jpg")
}

func ExampleMessage_Embed() {
	m.Embed("/tmp/image.jpg")
	m.SetBody("text/html", `<img src="cid:image.jpg" alt="My image" />`)
}

func ExampleMessage_FormatAddress() {
	m.SetHeader("To", m.FormatAddress("bob@example.com", "Bob"), m.FormatAddress("cora@example.com", "Cora"))
}

func ExampleMessage_FormatDate() {
	m.SetHeaders(map[string][]string{
		"X-Date": {m.FormatDate(time.Now())},
	})
}

func ExampleMessage_SetAddressHeader() {
	m.SetAddressHeader("To", "bob@example.com", "Bob")
}

func ExampleMessage_SetBody() {
	m.SetBody("text/plain", "Hello!")
}

func ExampleMessage_SetDateHeader() {
	m.SetDateHeader("X-Date", time.Now())
}

func ExampleMessage_SetHeader() {
	m.SetHeader("Subject", "Hello!")
}

func ExampleMessage_SetHeaders() {
	m.SetHeaders(map[string][]string{
		"From":    {m.FormatAddress("alex@example.com", "Alex")},
		"To":      {"bob@example.com", "cora@example.com"},
		"Subject": {"Hello"},
	})
}

func ExampleSetCharset() {
	m = gomail.NewMessage(gomail.SetCharset("ISO-8859-1"))
}

func ExampleSetEncoding() {
	m = gomail.NewMessage(gomail.SetEncoding(gomail.Base64))
}

func ExampleSetPartEncoding() {
	m.SetBody("text/plain", "Hello!", gomail.SetPartEncoding(gomail.Unencoded))
}
<commit_msg>Added an example using the CRAM-MD5 authentication mechanism<commit_after>package gomail_test

import (
	"fmt"
	"html/template"
	"io"
	"log"
	"net/smtp"
	"time"

	"gopkg.in/gomail.v2"
)

func Example() {
	m := gomail.NewMessage()
	m.SetHeader("From", "alex@example.com")
	m.SetHeader("To", "bob@example.com", "cora@example.com")
	m.SetAddressHeader("Cc", "dan@example.com", "Dan")
	m.SetHeader("Subject", "Hello!")
	m.SetBody("text/html", "Hello Bob and Cora!")
	m.Attach("/home/Alex/lolcat.jpg")

	d := gomail.NewPlainDialer("smtp.example.com", 587, "user", "123456")

	// Send the email to Bob, Cora and Dan.
	if err := d.DialAndSend(m); err != nil {
		panic(err)
	}
}

// A daemon that listens to a channel and sends all incoming messages.
func Example_daemon() {
	ch := make(chan *gomail.Message)

	go func() {
		d := gomail.NewPlainDialer("smtp.example.com", 587, "user", "123456")

		var s gomail.SendCloser
		var err error
		open := false
		for {
			select {
			case m, ok := <-ch:
				if !ok {
					return
				}
				if !open {
					if s, err = d.Dial(); err != nil {
						panic(err)
					}
					open = true
				}
				if err := gomail.Send(s, m); err != nil {
					log.Print(err)
				}
			// Close the connection to the SMTP server if no email was sent in
			// the last 30 seconds.
			case <-time.After(30 * time.Second):
				if open {
					if err := s.Close(); err != nil {
						panic(err)
					}
					open = false
				}
			}
		}
	}()

	// Use the channel in your program to send emails.

	// Close the channel to stop the mail daemon.
	close(ch)
}

// Efficiently send a customized newsletter to a list of recipients.
func Example_newsletter() {
	// The list of recipients.
	var list []struct {
		Name    string
		Address string
	}

	d := gomail.NewPlainDialer("smtp.example.com", 587, "user", "123456")
	s, err := d.Dial()
	if err != nil {
		panic(err)
	}

	m := gomail.NewMessage()
	for _, r := range list {
		m.SetHeader("From", "no-reply@example.com")
		m.SetAddressHeader("To", r.Address, r.Name)
		m.SetHeader("Subject", "Newsletter #1")
		m.SetBody("text/html", fmt.Sprintf("Hello %s!", r.Name))

		if err := gomail.Send(s, m); err != nil {
			log.Printf("Could not send email to %q: %v", r.Address, err)
		}
		m.Reset()
	}
}

// Send an email using a local SMTP server.
func Example_noAuth() {
	m := gomail.NewMessage()
	m.SetHeader("From", "from@example.com")
	m.SetHeader("To", "to@example.com")
	m.SetHeader("Subject", "Hello!")
	m.SetBody("text/plain", "Hello!")

	d := gomail.Dialer{Host: "localhost", Port: 587}
	if err := d.DialAndSend(m); err != nil {
		panic(err)
	}
}

// Send an email using the CRAM-MD5 authentication mechanism.
func Example_CRAM_MD5() {
	m := gomail.NewMessage()
	m.SetHeader("From", "from@example.com")
	m.SetHeader("To", "to@example.com")
	m.SetHeader("Subject", "Hello!")
	m.SetBody("text/plain", "Hello!")

	d := gomail.Dialer{
		Host: "localhost",
		Port: 587,
		Auth: smtp.CRAMMD5Auth("username", "secret"),
	}
	if err := d.DialAndSend(m); err != nil {
		panic(err)
	}
}

// Send an email using an API or postfix.
func Example_noSMTP() {
	m := gomail.NewMessage()
	m.SetHeader("From", "from@example.com")
	m.SetHeader("To", "to@example.com")
	m.SetHeader("Subject", "Hello!")
	m.SetBody("text/plain", "Hello!")

	s := gomail.SendFunc(func(from string, to []string, msg io.WriterTo) error {
		// Implements you email-sending function, for example by calling
		// an API, or running postfix, etc.
		fmt.Println("From:", from)
		fmt.Println("To:", to)
		return nil
	})

	if err := gomail.Send(s, m); err != nil {
		panic(err)
	}
	// Output:
	// From: from@example.com
	// To: [to@example.com]
}

var m *gomail.Message

func ExampleSetCopyFunc() {
	m.Attach("foo.txt", gomail.SetCopyFunc(func(w io.Writer) error {
		_, err := w.Write([]byte("Content of foo.txt"))
		return err
	}))
}

func ExampleSetHeader() {
	h := map[string][]string{"Content-ID": {"<foo@bar.mail>"}}
	m.Attach("foo.jpg", gomail.SetHeader(h))
}

func ExampleMessage_AddAlternative() {
	m.SetBody("text/plain", "Hello!")
	m.AddAlternative("text/html", "
Hello!
")
}

func ExampleMessage_AddAlternativeWriter() {
	t := template.Must(template.New("example").Parse("Hello {{.}}!"))
	m.AddAlternativeWriter("text/plain", func(w io.Writer) error {
		return t.Execute(w, "Bob")
	})
}

func ExampleMessage_Attach() {
	m.Attach("/tmp/image.jpg")
}

func ExampleMessage_Embed() {
	m.Embed("/tmp/image.jpg")
	m.SetBody("text/html", `<img src="cid:image.jpg" alt="My image" />`)
}

func ExampleMessage_FormatAddress() {
	m.SetHeader("To", m.FormatAddress("bob@example.com", "Bob"), m.FormatAddress("cora@example.com", "Cora"))
}

func ExampleMessage_FormatDate() {
	m.SetHeaders(map[string][]string{
		"X-Date": {m.FormatDate(time.Now())},
	})
}

func ExampleMessage_SetAddressHeader() {
	m.SetAddressHeader("To", "bob@example.com", "Bob")
}

func ExampleMessage_SetBody() {
	m.SetBody("text/plain", "Hello!")
}

func ExampleMessage_SetDateHeader() {
	m.SetDateHeader("X-Date", time.Now())
}

func ExampleMessage_SetHeader() {
	m.SetHeader("Subject", "Hello!")
}

func ExampleMessage_SetHeaders() {
	m.SetHeaders(map[string][]string{
		"From":    {m.FormatAddress("alex@example.com", "Alex")},
		"To":      {"bob@example.com", "cora@example.com"},
		"Subject": {"Hello"},
	})
}

func ExampleSetCharset() {
	m = gomail.NewMessage(gomail.SetCharset("ISO-8859-1"))
}

func ExampleSetEncoding() {
	m = gomail.NewMessage(gomail.SetEncoding(gomail.Base64))
}

func ExampleSetPartEncoding() {
	m.SetBody("text/plain", "Hello!", gomail.SetPartEncoding(gomail.Unencoded))
}
<|endoftext|>


<commit_before>// Copyright 2020 Google LLC
//
// 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
//
// 	https://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 cors

import (
	"net/textproto"
	"strconv"
	"strings"

	"github.com/google/go-safeweb/safehttp"
)

const customHeader string = "X-Cors"

var disallowedMethods = map[string]bool{
	safehttp.MethodHead: true,
	safehttp.MethodGet:  true,
	safehttp.MethodPost: true,
}

var disallowedContentTypes = map[string]bool{
	"application/x-www-form-urlencoded": true,
	"multipart/form-data":               true,
	"text/plain":                        true,
}

// Interceptor handles CORS requests based on its settings. The request methods
// GET, HEAD and POST are banned and will be responded to with 405 Method Not Allowed.
// The content types "application/x-www-form-urlencoded", "multipart/form-data"
// and "text/plain" are also banned and will be responded to with 415 Unsupported
// Media Type. Each CORS request must contain the header "X-Cors: 1" or
// "Sec-Fetch-Mode: cors". All of this is to prevent XSRF.
//
// For more info about CORS, see: https://developer.mozilla.org/en-US/docs/Web/HTTP/CORS
type Interceptor struct {
	// AllowedOrigins determines which origins should be allowed in the
	// Access-Control-Allow-Origin header.
	AllowedOrigins map[string]bool
	allowedHeaders map[string]bool
	// ExposedHeaders determines which headers should be set in the
	// Access-Control-Expose-Headers header. If ExposedHeaders is nil, then the
	// header is not set.
	ExposedHeaders []string
	// AllowCredentials determines if Access-Control-Allow-Credentials should be
	// set to true when requests containing Cookie headers arrive.
	AllowCredentials bool
	// MaxAge sets the Access-Control-Max-Age header.
	// MaxAge=0 results in MaxAge: 5.
	MaxAge int
}

// Default creates a CORS Interceptor with default settings.
// Those defaults are:
//   - No Exposed Headers
//   - No Allowed Headers
//   - AllowCredentials: false
//   - MaxAge: 5 seconds
func Default(allowedOrigins ...string) *Interceptor {
	ao := map[string]bool{}
	for _, o := range allowedOrigins {
		ao[o] = true
	}
	return &Interceptor{
		AllowedOrigins: ao,
		allowedHeaders: map[string]bool{},
	}
}

// SetAllowedHeaders sets which headers should be allowed in the Access-Control-Allow-Headers
// header. The headers are first canonicalized using textproto.CanonicalMIMEHeaderKey.
func (it *Interceptor) SetAllowedHeaders(headers ...string) {
	it.allowedHeaders = map[string]bool{}
	for _, h := range headers {
		it.allowedHeaders[textproto.CanonicalMIMEHeaderKey(h)] = true
	}
}

// Before handle the incoming request and sets headers or responds with an error
// accordingly.
func (it *Interceptor) Before(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
	origin := r.Header.Get("Origin")
	if !it.AllowedOrigins[origin] {
		return w.ClientError(safehttp.StatusForbidden)
	}
	h := w.Header()
	allowOrigin, err := h.Claim("Access-Control-Allow-Origin")
	if err != nil {
		return w.ServerError(safehttp.StatusInternalServerError)
	}
	vary, err := h.Claim("Vary")
	if err != nil {
		return w.ServerError(safehttp.StatusInternalServerError)
	}

	allowCredentials, err := h.Claim("Access-Control-Allow-Credentials")
	if err != nil {
		return w.ServerError(safehttp.StatusInternalServerError)
	}
	if r.Header.Get("Cookie") != "" && !it.AllowCredentials {
		return w.ClientError(safehttp.StatusForbidden)
	}

	// Used to set headers in preflight() or actualRequest() if no error occurs.
	setHeaders := func() {
		allowOrigin([]string{origin})
		vary([]string{"Origin"})
		if r.Header.Get("Cookie") != "" && it.AllowCredentials {
			// TODO: handle other credentials than cookies:
			// https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Access-Control-Allow-Credentials
			allowCredentials([]string{"true"})
		}
	}

	if r.Method() == safehttp.MethodOptions {
		return it.preflight(setHeaders, w, r)
	} else {
		return it.actualRequest(setHeaders, w, r)
	}
}

// preflight handles requests that have the method OPTIONS.
func (it *Interceptor) preflight(setHeaders func(), w *safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
	rh := r.Header
	method := rh.Get("Access-Control-Request-Method")
	if disallowedMethods[method] {
		return w.ClientError(safehttp.StatusForbidden)
	}
	wh := w.Header()
	allowMethods, err := wh.Claim("Access-Control-Allow-Methods")
	if err != nil {
		return w.ServerError(safehttp.StatusInternalServerError)
	}

	headers := rh.Get("Access-Control-Request-Headers")
	if headers != "" {
		for _, h := range strings.Split(headers, ", ") {
			h = textproto.CanonicalMIMEHeaderKey(h)
			if !it.allowedHeaders[h] && h != customHeader {
				return w.ClientError(safehttp.StatusForbidden)
			}
		}
	}
	allowHeaders, err := wh.Claim("Access-Control-Allow-Headers")
	if err != nil {
		return w.ServerError(safehttp.StatusInternalServerError)
	}

	maxAge, err := wh.Claim("Access-Control-Max-Age")
	if err != nil {
		return w.ServerError(safehttp.StatusInternalServerError)
	}

	allowMethods([]string{method})
	if headers != "" {
		allowHeaders([]string{headers})
	}
	n := it.MaxAge
	if n == 0 {
		n = 5
	}
	maxAge([]string{strconv.Itoa(n)})
	setHeaders()
	return w.NoContent()
}

// actualRequest handles all requests that are not preflight requests.
func (it *Interceptor) actualRequest(setHeaders func(), w *safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
	h := r.Header
	if h.Get("Sec-Fetch-Mode") != "cors" && h.Get(customHeader) != "1" {
		return w.ClientError(safehttp.StatusPreconditionFailed)
	}

	if disallowedMethods[r.Method()] {
		return w.ClientError(safehttp.StatusMethodNotAllowed)
	}

	if disallowedContentTypes[h.Get("Content-Type")] {
		return w.ClientError(safehttp.StatusUnsupportedMediaType)
	}

	exposeHeaders, err := w.Header().Claim("Access-Control-Expose-Headers")
	if err != nil {
		return w.ServerError(safehttp.StatusInternalServerError)
	}

	if it.ExposedHeaders != nil {
		exposeHeaders([]string{strings.Join(it.ExposedHeaders, ", ")})
	}
	setHeaders()
	return safehttp.Result{}
}
<commit_msg>Rephrased comments according to code review comments.<commit_after>// Copyright 2020 Google LLC
//
// 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
//
// 	https://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 cors

import (
	"net/textproto"
	"strconv"
	"strings"

	"github.com/google/go-safeweb/safehttp"
)

const customHeader string = "X-Cors"

var disallowedMethods = map[string]bool{
	safehttp.MethodHead: true,
	safehttp.MethodGet:  true,
	safehttp.MethodPost: true,
}

var disallowedContentTypes = map[string]bool{
	"application/x-www-form-urlencoded": true,
	"multipart/form-data":               true,
	"text/plain":                        true,
}

// Interceptor handles CORS requests based on its settings.
//
// For more info about CORS, see: https://developer.mozilla.org/en-US/docs/Web/HTTP/CORS
//
// Constraints
//
// - The request methods GET, HEAD and POST are banned and will result in a 405
//   Method Not Allowed response.
// - The content types "application/x-www-form-urlencoded", "multipart/form-data"
//   and "text/plain" are also banned and will result in a  415 Unsupported Media
//   Type response.
// - Each CORS request must contain the header "X-Cors: 1" or "Sec-Fetch-Mode: cors".
// All of this is to prevent XSRF.
type Interceptor struct {
	// AllowedOrigins determines which origins should be allowed in the
	// Access-Control-Allow-Origin header.
	AllowedOrigins map[string]bool
	allowedHeaders map[string]bool
	// ExposedHeaders determines which headers should be set in the
	// Access-Control-Expose-Headers header. If ExposedHeaders is nil, then the
	// header is not set.
	ExposedHeaders []string
	// AllowCredentials determines if Access-Control-Allow-Credentials should be
	// set to true when requests containing Cookie headers arrive.
	AllowCredentials bool
	// MaxAge sets the Access-Control-Max-Age header.
	// MaxAge=0 results in MaxAge: 5.
	MaxAge int
}

// Default creates a CORS Interceptor with default settings.
// Those defaults are:
//   - No Exposed Headers
//   - No Allowed Headers
//   - AllowCredentials: false
//   - MaxAge: 5 seconds
func Default(allowedOrigins ...string) *Interceptor {
	ao := map[string]bool{}
	for _, o := range allowedOrigins {
		ao[o] = true
	}
	return &Interceptor{
		AllowedOrigins: ao,
		allowedHeaders: map[string]bool{},
	}
}

// SetAllowedHeaders sets which headers should be allowed in the Access-Control-Allow-Headers
// header. The headers are first canonicalized using textproto.CanonicalMIMEHeaderKey.
func (it *Interceptor) SetAllowedHeaders(headers ...string) {
	it.allowedHeaders = map[string]bool{}
	for _, h := range headers {
		it.allowedHeaders[textproto.CanonicalMIMEHeaderKey(h)] = true
	}
}

// Before handles the incoming request and sets headers or responds with an error
// accordingly.
func (it *Interceptor) Before(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
	origin := r.Header.Get("Origin")
	if !it.AllowedOrigins[origin] {
		return w.ClientError(safehttp.StatusForbidden)
	}
	h := w.Header()
	allowOrigin, err := h.Claim("Access-Control-Allow-Origin")
	if err != nil {
		return w.ServerError(safehttp.StatusInternalServerError)
	}
	vary, err := h.Claim("Vary")
	if err != nil {
		return w.ServerError(safehttp.StatusInternalServerError)
	}

	allowCredentials, err := h.Claim("Access-Control-Allow-Credentials")
	if err != nil {
		return w.ServerError(safehttp.StatusInternalServerError)
	}
	if r.Header.Get("Cookie") != "" && !it.AllowCredentials {
		return w.ClientError(safehttp.StatusForbidden)
	}

	// Used to set headers in preflight() or actualRequest() if no error occurs.
	setHeaders := func() {
		allowOrigin([]string{origin})
		vary([]string{"Origin"})
		if r.Header.Get("Cookie") != "" && it.AllowCredentials {
			// TODO: handle other credentials than cookies:
			// https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Access-Control-Allow-Credentials
			allowCredentials([]string{"true"})
		}
	}

	if r.Method() == safehttp.MethodOptions {
		return it.preflight(setHeaders, w, r)
	} else {
		return it.actualRequest(setHeaders, w, r)
	}
}

// preflight handles requests that have the method OPTIONS.
func (it *Interceptor) preflight(setHeaders func(), w *safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
	rh := r.Header
	method := rh.Get("Access-Control-Request-Method")
	if disallowedMethods[method] {
		return w.ClientError(safehttp.StatusForbidden)
	}
	wh := w.Header()
	allowMethods, err := wh.Claim("Access-Control-Allow-Methods")
	if err != nil {
		return w.ServerError(safehttp.StatusInternalServerError)
	}

	headers := rh.Get("Access-Control-Request-Headers")
	if headers != "" {
		for _, h := range strings.Split(headers, ", ") {
			h = textproto.CanonicalMIMEHeaderKey(h)
			if !it.allowedHeaders[h] && h != customHeader {
				return w.ClientError(safehttp.StatusForbidden)
			}
		}
	}
	allowHeaders, err := wh.Claim("Access-Control-Allow-Headers")
	if err != nil {
		return w.ServerError(safehttp.StatusInternalServerError)
	}

	maxAge, err := wh.Claim("Access-Control-Max-Age")
	if err != nil {
		return w.ServerError(safehttp.StatusInternalServerError)
	}

	allowMethods([]string{method})
	if headers != "" {
		allowHeaders([]string{headers})
	}
	n := it.MaxAge
	if n == 0 {
		n = 5
	}
	maxAge([]string{strconv.Itoa(n)})
	setHeaders()
	return w.NoContent()
}

// actualRequest handles all requests that are not preflight requests.
func (it *Interceptor) actualRequest(setHeaders func(), w *safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
	h := r.Header
	if h.Get("Sec-Fetch-Mode") != "cors" && h.Get(customHeader) != "1" {
		return w.ClientError(safehttp.StatusPreconditionFailed)
	}

	if disallowedMethods[r.Method()] {
		return w.ClientError(safehttp.StatusMethodNotAllowed)
	}

	if disallowedContentTypes[h.Get("Content-Type")] {
		return w.ClientError(safehttp.StatusUnsupportedMediaType)
	}

	exposeHeaders, err := w.Header().Claim("Access-Control-Expose-Headers")
	if err != nil {
		return w.ServerError(safehttp.StatusInternalServerError)
	}

	if it.ExposedHeaders != nil {
		exposeHeaders([]string{strings.Join(it.ExposedHeaders, ", ")})
	}
	setHeaders()
	return safehttp.Result{}
}
<|endoftext|>


<commit_before>package exec

import (
	"database/sql/driver"
	"encoding/gob"
	"fmt"
	"strings"
	"time"

	u "github.com/araddon/gou"

	"github.com/araddon/qlbridge/datasource"
	"github.com/araddon/qlbridge/expr"
	"github.com/araddon/qlbridge/plan"
	"github.com/araddon/qlbridge/rel"
	"github.com/araddon/qlbridge/value"
	"github.com/araddon/qlbridge/vm"
)

var (
	_ = u.EMPTY

	// Ensure that we implement the Task Runner interface
	_ TaskRunner = (*GroupBy)(nil)
)

func init() {
	gob.Register(AggPartial{})
}

// Group by:   Sql Group By Operator
//   creates a hashable key commposed of key = {each,value,of,column,in,groupby}
//
// A very stupid naive parallel groupby holds values in memory
//
//   task   ->  groupby  -->
//
type GroupBy struct {
	*TaskBase
	closed bool
	p      *plan.GroupBy
}

func NewGroupBy(ctx *plan.Context, p *plan.GroupBy) *GroupBy {
	m := &GroupBy{
		TaskBase: NewTaskBase(ctx),
		p:        p,
	}
	return m
}

// Group by:   Sql Group By Operator finalizer for partials
//
type GroupByFinal struct {
	*TaskBase
	p          *plan.GroupBy
	complete   chan bool
	closed     bool
	isComplete bool
}

func NewGroupByFinal(ctx *plan.Context, p *plan.GroupBy) *GroupByFinal {
	m := &GroupByFinal{
		TaskBase: NewTaskBase(ctx),
		p:        p,
		complete: make(chan bool),
	}
	return m
}

func (m *GroupBy) Run() error {
	defer m.Ctx.Recover()
	defer close(m.msgOutCh)

	outCh := m.MessageOut()
	inCh := m.MessageIn()

	columns := m.p.Stmt.Columns
	colIndex := m.p.Stmt.ColIndexes()

	aggs, err := buildAggs(m.p)
	if err != nil {
		return err
	}

	// are are going to hold entire row in memory while we are calculating
	//  so obviously not scalable.
	gb := make(map[string][]*datasource.SqlDriverMessageMap)

msgReadLoop:
	for {

		select {
		case <-m.SigChan():
			u.Warnf("got signal quit")
			return nil
		case msg, ok := <-inCh:
			if !ok {
				//u.Debugf("NICE, got closed channel shutdown")
				break msgReadLoop
			} else {
				var sdm *datasource.SqlDriverMessageMap

				switch mt := msg.(type) {
				case *datasource.SqlDriverMessageMap:
					sdm = mt
				default:

					msgReader, isContextReader := msg.(expr.ContextReader)
					if !isContextReader {
						err := fmt.Errorf("To use Join must use SqlDriverMessageMap but got %T", msg)
						u.Errorf("unrecognized msg %T", msg)
						close(m.TaskBase.sigCh)
						return err
					}

					sdm = datasource.NewSqlDriverMessageMapCtx(msg.Id(), msgReader, colIndex)
				}
				// We are going to use VM Engine to create a value for each statement in group by
				//  then join each value together to create a unique key.
				keys := make([]string, len(m.p.Stmt.GroupBy))
				for i, col := range m.p.Stmt.GroupBy {
					if col.Expr != nil {
						if key, ok := vm.Eval(sdm, col.Expr); ok {
							//u.Debugf("msgtype:%T  key:%q for-expr:%s", sdm, key, col.Expr)
							keys[i] = key.ToString()
						} else {
							// Is this an error?
							//u.Warnf("no key?  %s for %+v", col.Expr, sdm)
						}
					} else {
						u.Warnf("no col.expr? %#v", col)
					}
				}
				key := strings.Join(keys, ",")
				//u.Infof("found key:%s for %+v", key, sdm)
				gb[key] = append(gb[key], sdm)
			}
		}
	}

	i := uint64(0)
	for key, v := range gb {
		//u.Debugf("got %s:%v msgs", k, len(v))

		for _, mm := range v {
			for i, col := range columns {
				//u.Debugf("col: idx:%v sidx: %v pidx:%v key:%v   %s", col.Index, col.SourceIndex, col.ParentIndex, col.Key(), col.Expr)

				if col.Expr == nil {
					u.Warnf("wat?   nil col expr? %#v", col)
				} else {
					v, ok := vm.Eval(mm, col.Expr)
					//u.Infof("mt: %T  mm %#v", mm, mm)
					if !ok || v == nil {
						//u.Debugf("evaled nil? key=%v  val=%v expr:%s", col.Key(), v, col.Expr.String())
						//u.Infof("mt: %T  mm %#v", mm, mm)
						aggs[i].Do(value.NewNilValue())
					} else {
						//u.Debugf("evaled: key=%v  val=%v", col.Key(), v.Value())
						aggs[i].Do(v)
					}
				}
			}
		}

		row := make([]driver.Value, len(columns))
		for i, agg := range aggs {
			row[i] = driver.Value(agg.Result())
			agg.Reset()
			//u.Debugf("agg result: %#v  %v", row[i], row[i])
		}

		if m.p.Partial {
			// Partial results, append key at end?  shouldn't be able to be fit in message itself?
			row = append(row, key)
			//u.Debugf("GroupBy output row? key:%s %#v", key, row)
		}
		//u.Debugf("row: %v  cols:%v", row, colIndex)
		outCh <- datasource.NewSqlDriverMessageMap(i, row, colIndex)
		i++
	}

	return nil
}

func (m *GroupByFinal) Run() error {
	defer m.Ctx.Recover()
	defer close(m.msgOutCh)

	outCh := m.MessageOut()
	inCh := m.MessageIn()

	columns := m.p.Stmt.Columns
	colIndex := m.p.Stmt.ColIndexes()

	m.p.Partial = false
	aggs, err := buildAggs(m.p)
	if err != nil {
		return err
	}

	gb := make(map[string][][]driver.Value)

msgReadLoop:
	for {

		select {
		case <-m.SigChan():
			u.Warnf("got signal quit")
			return nil
		case msg, ok := <-inCh:
			if !ok {
				//u.Debugf("GroupByFinal, got closed channel shutdown")
				break msgReadLoop
			} else {
				//u.Infof("got gbfinal message %#v", msg)
				switch mt := msg.(type) {
				case *datasource.SqlDriverMessageMap:
					if len(mt.Vals) != len(columns)+1 {
						u.Warnf("Wrong number of values? %#v", mt)
					}
					key, ok := mt.Vals[len(mt.Vals)-1].(string)
					if !ok {
						u.Warnf("expected key?  %#v", mt.Vals)
					}
					vals := mt.Vals[0 : len(mt.Vals)-1]
					//u.Infof("found key:%s for %#v", key, mt.Vals)
					gb[key] = append(gb[key], vals)
				default:
					err := fmt.Errorf("To use Join must use SqlDriverMessageMap but got %T", msg)
					u.Errorf("unrecognized msg %T", msg)
					close(m.TaskBase.sigCh)
					return err
				}
			}
		}
	}

	i := uint64(0)
	for _, vals := range gb {
		//u.Debugf("got %s:%v msgs", key, vals)

		for _, dv := range vals {
			for i, col := range columns {
				//u.Debugf("col: idx:%v sidx: %v pidx:%v key:%v   %s", col.Index, col.SourceIndex, col.ParentIndex, col.Key(), col.Expr)

				if col.Expr == nil {
					u.Warnf("wat?   nil col expr? %#v", col)
				} else {
					v := dv[i]
					switch vt := v.(type) {
					case *AggPartial:
						//u.Debugf("evaled: key=%v  val=%v", col.Key(), v.Value())
						aggs[i].Merge(vt)
					case AggPartial:
						aggs[i].Merge(&vt)
					default:
						u.Warnf("unhandled type: %#v", v)
					}
				}
			}
		}

		row := make([]driver.Value, len(columns))
		for i, agg := range aggs {
			row[i] = driver.Value(agg.Result())
			agg.Reset()
			//u.Debugf("agg result: %#v  %v", row[i], row[i])
		}
		//u.Debugf("GroupBy output row? %v", row)
		outCh <- datasource.NewSqlDriverMessageMap(i, row, colIndex)
		i++
	}

	m.isComplete = true
	close(m.complete)
	return nil
}

func (m *GroupBy) Close() error {
	if m.closed {
		return nil
	}
	m.closed = true
	return m.TaskBase.Close()
}
func (m *GroupByFinal) Close() error {
	if m.closed {
		return nil
	}
	m.closed = true

	ticker := time.NewTicker(30 * time.Second)
	defer ticker.Stop()

	//u.Infof("%p group by final Close() waiting for complete", m)
	select {
	case <-ticker.C:
		u.Warnf("timeout???? ")
	case <-m.complete:
		//u.Warnf("%p got groupbyfinal complete", m)
	}

	return m.TaskBase.Close()
}

type AggPartial struct {
	Ct int64
	N  float64
}

type AggFunc func(v value.Value)
type resultFunc func() interface{}
type Aggregator interface {
	Do(v value.Value)
	Result() interface{}
	Reset()
	Merge(*AggPartial)
}
type agg struct {
	do     AggFunc
	result resultFunc
}
type groupByFunc struct {
	last interface{}
}

func (m *groupByFunc) Do(v value.Value)    { m.last = v.Value() }
func (m *groupByFunc) Result() interface{} { return m.last }
func (m *groupByFunc) Reset()              { m.last = nil }
func (m *groupByFunc) Merge(a *AggPartial) {}
func NewGroupByValue(col *rel.Column) Aggregator {
	return &groupByFunc{}
}

type sum struct {
	partial bool
	ct      int64
	n       float64
}

func (m *sum) Do(v value.Value) {
	m.ct++
	switch vt := v.(type) {
	case value.IntValue:
		m.n += vt.Float()
	case value.NumberValue:
		m.n += vt.Val()
	}
}
func (m *sum) Result() interface{} {
	if !m.partial {
		return m.n
	}
	return &AggPartial{
		m.ct,
		m.n,
	}
}
func (m *sum) Reset() { m.n = 0 }
func (m *sum) Merge(a *AggPartial) {
	m.ct += a.Ct
	m.n += a.N
}
func NewSum(col *rel.Column, partial bool) Aggregator {
	return &sum{partial: partial}
}

type avg struct {
	partial bool
	ct      int64
	n       float64
}

func (m *avg) Do(v value.Value) {
	m.ct++
	switch vt := v.(type) {
	case value.IntValue:
		m.n += vt.Float()
	case value.NumberValue:
		m.n += vt.Val()
	}
}
func (m *avg) Result() interface{} {
	if !m.partial {
		return m.n / float64(m.ct)
	}
	return &AggPartial{
		m.ct,
		m.n,
	}
}
func (m *avg) Reset() { m.n = 0; m.ct = 0 }
func (m *avg) Merge(a *AggPartial) {
	m.ct += a.Ct
	m.n += a.N
}
func NewAvg(col *rel.Column, partial bool) Aggregator {
	return &avg{partial: partial}
}

type count struct {
	n int64
}

func (m *count) Do(v value.Value)    { m.n++ }
func (m *count) Result() interface{} { return m.n }
func (m *count) Reset()              { m.n = 0 }
func (m *count) Merge(a *AggPartial) {
	m.n += a.Ct
}
func NewCount(col *rel.Column) Aggregator {
	return &count{}
}

func buildAggs(p *plan.GroupBy) ([]Aggregator, error) {
	//u.Debugf("build aggs: partial:%v  sql:%s", p.Partial, p.Stmt)
	aggs := make([]Aggregator, len(p.Stmt.Columns))
colLoop:
	for colIdx, col := range p.Stmt.Columns {
		for _, gb := range p.Stmt.GroupBy {
			if gb.As == col.As {
				// simple Non Aggregate Value
				aggs[colIdx] = NewGroupByValue(col)
				continue colLoop
			}
		}
		// Since we made it here, ann aggregate func
		//  move to a registry of some kind to allow extension
		switch n := col.Expr.(type) {
		case *expr.FuncNode:

			// TODO:  extract to a UDF Registry Similar to builtins
			switch strings.ToLower(n.Name) {
			case "avg":
				aggs[colIdx] = NewAvg(col, p.Partial)
			case "count":
				aggs[colIdx] = NewCount(col)
			case "sum":
				aggs[colIdx] = NewSum(col, p.Partial)
			default:
				return nil, fmt.Errorf("Not impelemneted groupby for column: %s", col.Expr)
			}
		case *expr.BinaryNode:
			// binary logic?
			return nil, fmt.Errorf("Not impelemneted groupby for column: %s", col.Expr)
		}
	}
	return aggs, nil
}
<commit_msg>better support for distributed aggregations<commit_after>package exec

import (
	"database/sql/driver"
	"encoding/gob"
	"fmt"
	"strings"
	"time"

	u "github.com/araddon/gou"

	"github.com/araddon/qlbridge/datasource"
	"github.com/araddon/qlbridge/expr"
	"github.com/araddon/qlbridge/plan"
	"github.com/araddon/qlbridge/rel"
	"github.com/araddon/qlbridge/value"
	"github.com/araddon/qlbridge/vm"
)

var (
	_ = u.EMPTY

	// Ensure that we implement the Task Runner interface
	_ TaskRunner = (*GroupBy)(nil)
)

func init() {
	gob.Register(AggPartial{})
}

// Group by:   Sql Group By Operator
//   creates a hashable key commposed of key = {each,value,of,column,in,groupby}
//
// A very stupid naive parallel groupby holds values in memory
//
//   task   ->  groupby  -->
//
type GroupBy struct {
	*TaskBase
	closed bool
	p      *plan.GroupBy
}

func NewGroupBy(ctx *plan.Context, p *plan.GroupBy) *GroupBy {
	m := &GroupBy{
		TaskBase: NewTaskBase(ctx),
		p:        p,
	}
	return m
}

// Group by:   Sql Group By Operator finalizer for partials
//
type GroupByFinal struct {
	*TaskBase
	p          *plan.GroupBy
	complete   chan bool
	closed     bool
	isComplete bool
}

func NewGroupByFinal(ctx *plan.Context, p *plan.GroupBy) *GroupByFinal {
	m := &GroupByFinal{
		TaskBase: NewTaskBase(ctx),
		p:        p,
		complete: make(chan bool),
	}
	return m
}

func (m *GroupBy) Run() error {
	defer m.Ctx.Recover()
	defer close(m.msgOutCh)

	outCh := m.MessageOut()
	inCh := m.MessageIn()

	columns := m.p.Stmt.Columns
	colIndex := m.p.Stmt.ColIndexes()

	aggs, err := buildAggs(m.p)
	if err != nil {
		return err
	}

	// are are going to hold entire row in memory while we are calculating
	//  so obviously not scalable.
	gb := make(map[string][]*datasource.SqlDriverMessageMap)

msgReadLoop:
	for {

		select {
		case <-m.SigChan():
			u.Warnf("got signal quit")
			return nil
		case msg, ok := <-inCh:
			if !ok {
				//u.Debugf("NICE, got closed channel shutdown")
				break msgReadLoop
			} else {
				var sdm *datasource.SqlDriverMessageMap

				switch mt := msg.(type) {
				case *datasource.SqlDriverMessageMap:
					sdm = mt
				default:

					msgReader, isContextReader := msg.(expr.ContextReader)
					if !isContextReader {
						err := fmt.Errorf("To use Join must use SqlDriverMessageMap but got %T", msg)
						u.Errorf("unrecognized msg %T", msg)
						close(m.TaskBase.sigCh)
						return err
					}

					sdm = datasource.NewSqlDriverMessageMapCtx(msg.Id(), msgReader, colIndex)
				}

				// We are going to use VM Engine to create a value for each statement in group by
				//  then join each value together to create a unique key.
				keys := make([]string, len(m.p.Stmt.GroupBy))
				for i, col := range m.p.Stmt.GroupBy {
					if col.Expr != nil {
						if key, ok := vm.Eval(sdm, col.Expr); ok {
							//u.Debugf("msgtype:%T  key:%q for-expr:%s", sdm, key, col.Expr)
							keys[i] = key.ToString()
						} else {
							// Is this an error?
							//u.Warnf("no key?  %s for %+v", col.Expr, sdm)
						}
					} else {
						u.Warnf("no col.expr? %#v", col)
					}
				}
				key := strings.Join(keys, ",")
				//u.Infof("found key:%s for %+v", key, sdm)
				gb[key] = append(gb[key], sdm)
			}
		}
	}

	i := uint64(0)
	for key, v := range gb {
		//u.Debugf("got %s:%v msgs", k, len(v))

		for _, mm := range v {
			for i, col := range columns {
				//u.Debugf("col: idx:%v sidx: %v pidx:%v key:%v   %s", col.Index, col.SourceIndex, col.ParentIndex, col.Key(), col.Expr)

				if col.Expr == nil {
					u.Warnf("wat?   nil col expr? %#v", col)
				} else {
					v, ok := vm.Eval(mm, col.Expr)
					//u.Infof("mt: %T  mm %#v", mm, mm)
					if !ok || v == nil {
						//u.Debugf("evaled nil? key=%v  val=%v expr:%s", col.Key(), v, col.Expr.String())
						//u.Infof("mt: %T  mm %#v", mm, mm)
						aggs[i].Do(value.NewNilValue())
					} else {
						//u.Debugf("evaled: key=%v  val=%v", col.Key(), v.Value())
						aggs[i].Do(v)
					}
				}
			}
		}

		row := make([]driver.Value, len(columns))
		for i, agg := range aggs {
			row[i] = driver.Value(agg.Result())
			agg.Reset()
			//u.Debugf("agg result: %#v  %v", row[i], row[i])
		}

		if m.p.Partial {
			// Partial results, append key at end?  shouldn't be able to be fit in message itself?
			row = append(row, key)
			//u.Debugf("GroupBy output row? key:%s %#v", key, row)
		}
		//u.Debugf("row: %v  cols:%v", row, colIndex)
		outCh <- datasource.NewSqlDriverMessageMap(i, row, colIndex)
		i++
	}

	return nil
}

func (m *GroupByFinal) Run() error {
	defer m.Ctx.Recover()
	defer close(m.msgOutCh)

	outCh := m.MessageOut()
	inCh := m.MessageIn()

	columns := m.p.Stmt.Columns
	colIndex := m.p.Stmt.ColIndexes()

	m.p.Partial = false
	aggs, err := buildAggs(m.p)
	if err != nil {
		return err
	}

	gb := make(map[string][][]driver.Value)

msgReadLoop:
	for {

		select {
		case <-m.SigChan():
			u.Warnf("got signal quit")
			return nil
		case msg, ok := <-inCh:
			if !ok {
				//u.Debugf("GroupByFinal, got closed channel shutdown")
				break msgReadLoop
			} else {
				//u.Infof("got gbfinal message %#v", msg)
				switch mt := msg.(type) {
				case *datasource.SqlDriverMessageMap:
					if len(mt.Vals) != len(columns)+1 {
						u.Warnf("Wrong number of values? %#v", mt)
					}
					key, ok := mt.Vals[len(mt.Vals)-1].(string)
					if !ok {
						u.Warnf("expected key?  %#v", mt.Vals)
					}
					vals := mt.Vals[0 : len(mt.Vals)-1]
					//u.Infof("found key:%s for %#v", key, mt.Vals)
					gb[key] = append(gb[key], vals)
				default:
					err := fmt.Errorf("To use Join must use SqlDriverMessageMap but got %T", msg)
					u.Errorf("unrecognized msg %T", msg)
					close(m.TaskBase.sigCh)
					return err
				}
			}
		}
	}

	i := uint64(0)
	for _, vals := range gb {
		//u.Debugf("got %s:%v msgs", key, vals)

		for _, dv := range vals {
			for i, col := range columns {
				//u.Debugf("col: idx:%v sidx: %v pidx:%v key:%v   %s", col.Index, col.SourceIndex, col.ParentIndex, col.Key(), col.Expr)
				if i-1 >= len(dv) {
					u.Errorf("what??? %v  dv: %d   %#v", i, len(dv), dv)
				}
				if col.Expr == nil {
					u.Warnf("wat?   nil col expr? %#v", col)
				} else {
					v := dv[i]
					switch vt := v.(type) {
					case *AggPartial:
						//u.Debugf("evaled: key=%v  val=%v", col.Key(), v.Value())
						aggs[i].Merge(vt)
					case AggPartial:
						aggs[i].Merge(&vt)
					case int64:
						aggs[i].Merge(&AggPartial{Ct: vt})
					case string:
						aggs[i] = &groupByFunc{vt}
					default:
						u.Warnf("unhandled type: %#v", v)
					}
				}
			}
		}

		row := make([]driver.Value, len(columns))
		for i, agg := range aggs {
			row[i] = driver.Value(agg.Result())
			agg.Reset()
			//u.Debugf("agg result: %#v  %v", row[i], row[i])
		}
		//u.Debugf("GroupBy output row? %v", row)
		outCh <- datasource.NewSqlDriverMessageMap(i, row, colIndex)
		i++
	}

	m.isComplete = true
	close(m.complete)
	return nil
}

func (m *GroupBy) Close() error {
	if m.closed {
		return nil
	}
	m.closed = true
	return m.TaskBase.Close()
}
func (m *GroupByFinal) Close() error {
	if m.closed {
		return nil
	}
	m.closed = true

	ticker := time.NewTicker(30 * time.Second)
	defer ticker.Stop()

	//u.Infof("%p group by final Close() waiting for complete", m)
	select {
	case <-ticker.C:
		u.Warnf("timeout???? ")
	case <-m.complete:
		//u.Warnf("%p got groupbyfinal complete", m)
	}

	return m.TaskBase.Close()
}

type AggPartial struct {
	Ct int64
	N  float64
}

type AggFunc func(v value.Value)
type resultFunc func() interface{}
type Aggregator interface {
	Do(v value.Value)
	Result() interface{}
	Reset()
	Merge(*AggPartial)
}
type agg struct {
	do     AggFunc
	result resultFunc
}
type groupByFunc struct {
	last interface{}
}

func (m *groupByFunc) Do(v value.Value)    { m.last = v.Value() }
func (m *groupByFunc) Result() interface{} { return m.last }
func (m *groupByFunc) Reset()              { m.last = nil }
func (m *groupByFunc) Merge(a *AggPartial) {}
func NewGroupByValue(col *rel.Column) Aggregator {
	return &groupByFunc{}
}

type sum struct {
	partial bool
	ct      int64
	n       float64
}

func (m *sum) Do(v value.Value) {
	m.ct++
	switch vt := v.(type) {
	case value.IntValue:
		m.n += vt.Float()
	case value.NumberValue:
		m.n += vt.Val()
	}
}
func (m *sum) Result() interface{} {
	if !m.partial {
		return m.n
	}
	return &AggPartial{
		m.ct,
		m.n,
	}
}
func (m *sum) Reset() { m.n = 0 }
func (m *sum) Merge(a *AggPartial) {
	m.ct += a.Ct
	m.n += a.N
}
func NewSum(col *rel.Column, partial bool) Aggregator {
	return &sum{partial: partial}
}

type avg struct {
	partial bool
	ct      int64
	n       float64
}

func (m *avg) Do(v value.Value) {
	m.ct++
	switch vt := v.(type) {
	case value.IntValue:
		m.n += vt.Float()
	case value.NumberValue:
		m.n += vt.Val()
	}
}
func (m *avg) Result() interface{} {
	if !m.partial {
		return m.n / float64(m.ct)
	}
	return &AggPartial{
		m.ct,
		m.n,
	}
}
func (m *avg) Reset() { m.n = 0; m.ct = 0 }
func (m *avg) Merge(a *AggPartial) {
	m.ct += a.Ct
	m.n += a.N
}
func NewAvg(col *rel.Column, partial bool) Aggregator {
	return &avg{partial: partial}
}

type count struct {
	n int64
}

func (m *count) Do(v value.Value) { m.n++ }
func (m *count) Result() interface{} {
	return m.n
}
func (m *count) Reset() { m.n = 0 }
func (m *count) Merge(a *AggPartial) {
	m.n += a.Ct
}
func NewCount(col *rel.Column) Aggregator {
	return &count{}
}

func buildAggs(p *plan.GroupBy) ([]Aggregator, error) {
	//u.Debugf("build aggs: partial:%v  sql:%s", p.Partial, p.Stmt)
	aggs := make([]Aggregator, len(p.Stmt.Columns))
colLoop:
	for colIdx, col := range p.Stmt.Columns {
		for _, gb := range p.Stmt.GroupBy {
			if gb.As == col.As {
				// simple Non Aggregate Value
				aggs[colIdx] = NewGroupByValue(col)
				continue colLoop
			}
		}
		// Since we made it here, ann aggregate func
		//  move to a registry of some kind to allow extension
		switch n := col.Expr.(type) {
		case *expr.FuncNode:

			// TODO:  extract to a UDF Registry Similar to builtins
			switch strings.ToLower(n.Name) {
			case "avg":
				aggs[colIdx] = NewAvg(col, p.Partial)
			case "count":
				aggs[colIdx] = NewCount(col)
			case "sum":
				aggs[colIdx] = NewSum(col, p.Partial)
			default:
				return nil, fmt.Errorf("Not impelemneted groupby for column: %s", col.Expr)
			}
		case *expr.BinaryNode:
			// binary logic?
			return nil, fmt.Errorf("Not impelemneted groupby for column: %s", col.Expr)
		}
	}
	return aggs, nil
}
<|endoftext|>


<commit_before>/*
 * Copyright 2017 the original author or 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 main

import (
	"context"
	"fmt"
	"gopkg.in/Shopify/sarama.v1"
	"github.com/bsm/sarama-cluster"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"os/signal"
	"time"
	"github.com/satori/go.uuid"
	"github.com/sk8sio/http-gateway/pkg/message"
)

// Function messageHandler creates an http handler that posts the http body as a message to Kafka, replying
// immediately with a successful http response
func messageHandler(producer sarama.AsyncProducer) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		topic := r.URL.Path[len("/messages/"):]
		b, err := ioutil.ReadAll(r.Body)
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
		scsMessage := message.Message{Payload: b, Headers: nil}

		bytesOut, err := message.EncodeMessage(scsMessage)
		if err != nil {
			log.Printf("Error encoding message: %v", err)
			return
		}
		kafkaMsg := &sarama.ProducerMessage{Topic: topic, Value: sarama.ByteEncoder(bytesOut)}

		select {
		case producer.Input() <- kafkaMsg:
			w.Write([]byte("message published to topic: " + topic + "\n"))
		}
	}
}

// Function replyHandler creates an http handler that posts the http body as a message to Kafka, then waits
// for a message on a go channel it creates for a reply (this is expected to be set by the main thread) and sends
// that as an http response.
func replyHandler(producer sarama.AsyncProducer, replies map[string]chan message.Message) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		topic := r.URL.Path[len("/requests/"):]
		b, err := ioutil.ReadAll(r.Body)
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
		correlationId := uuid.NewV4().String()
		replyChan := make(chan message.Message)
		replies[correlationId] = replyChan
		scsMessage := message.Message{Payload: b, Headers: map[string]interface{}{"correlationId": correlationId}}

		bytesOut, err := message.EncodeMessage(scsMessage)
		if err != nil {
			log.Printf("Error encoding message: %v", err)
			return
		}
		kafkaMsg := &sarama.ProducerMessage{Topic: topic, Value: sarama.ByteEncoder(bytesOut)}

		select {
		case producer.Input() <- kafkaMsg:
			select {
			case reply := <-replyChan:
				delete(replies, correlationId)
				p := reply.Payload
				w.Write(p.([]byte)) // TODO equivalent of Spring's HttpMessageConverter handling
			case <- time.After(time.Second * 60):
				delete(replies, correlationId)
				w.WriteHeader(404)
			}
		}
	}
}

func healthHandler() func(http.ResponseWriter, *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte(`{"status":"UP"}`))
	}
}

func startHttpServer(producer sarama.AsyncProducer, replies map[string]chan message.Message) *http.Server {
	srv := &http.Server{Addr: ":8080"}

	http.HandleFunc("/messages/", messageHandler(producer))
	http.HandleFunc("/requests/", replyHandler(producer, replies))
	http.HandleFunc("/application/status", healthHandler())

	go func() {
		if err := srv.ListenAndServe(); err != nil {
			log.Printf("Httpserver: ListenAndServe() error: %s", err)
		}
	}()

	log.Printf("Listening on %v", srv.Addr)
	return srv
}

func main() {
	// Trap SIGINT to trigger a shutdown.
	signals := make(chan os.Signal, 1)
	signal.Notify(signals, os.Interrupt)

	// Key is correlationId, value is channel used to pass message received from main Kafka consumer loop
	replies := make(map[string]chan message.Message)

	brokers := []string{os.Getenv("SPRING_CLOUD_STREAM_KAFKA_BINDER_BROKERS")}
	producer, err := sarama.NewAsyncProducer(brokers, nil)
	if err != nil {
		panic(err)
	}
	defer func() {
		if err := producer.Close(); err != nil {
			log.Fatalln(err)
		}
	}()

	consumerConfig := makeConsumerConfig()
	consumer, err := cluster.NewConsumer(brokers, "gateway", []string{"replies"}, consumerConfig)
	if err != nil {
		panic(err)
	}
	defer consumer.Close()
	if consumerConfig.Consumer.Return.Errors {
		go consumeErrors(consumer)
	}
	if consumerConfig.Group.Return.Notifications {
		go consumeNotifications(consumer)
	}

	srv := startHttpServer(producer, replies)

MainLoop:
	for {
		select {
		case <-signals:
			fmt.Println("Shutting Down...")
			timeout, c := context.WithTimeout(context.Background(), 1*time.Second)
			defer c()
			if err := srv.Shutdown(timeout); err != nil {
				panic(err) // failure/timeout shutting down the server gracefully
			}
			break MainLoop
		case msg, ok := <-consumer.Messages():
			if ok {
				messageWithHeaders, err := message.ExtractMessage(msg.Value)
				if err != nil {
					log.Println("Failed to extract message ", err)
					break
				}
				correlationId, ok := messageWithHeaders.Headers["correlationId"].(string)
				if ok {
					c := replies[correlationId]
					if c != nil {
						fmt.Printf("Sending %v\n", messageWithHeaders)
						c <- messageWithHeaders
						consumer.MarkOffset(msg, "") // mark message as processed
					} else {
						log.Printf("Did not find communication channel for correlationId %v. Timed out?", correlationId)
						consumer.MarkOffset(msg, "") // mark message as processed
					}
				}
			}

		case err := <-producer.Errors():
			log.Println("Failed to produce kafka message ", err)
		}
	}
}

func makeConsumerConfig() *cluster.Config {
	consumerConfig := cluster.NewConfig()
	consumerConfig.Consumer.Return.Errors = true
	consumerConfig.Group.Return.Notifications = true
	return consumerConfig
}

func consumeNotifications(consumer *cluster.Consumer) {
	for ntf := range consumer.Notifications() {
		log.Printf("Rebalanced: %+v\n", ntf)
	}
}
func consumeErrors(consumer *cluster.Consumer) {
	for err := range consumer.Errors() {
		log.Printf("Error: %s\n", err.Error())
	}
}
<commit_msg>Trap other signals<commit_after>/*
 * Copyright 2017 the original author or 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 main

import (
	"context"
	"fmt"
	"gopkg.in/Shopify/sarama.v1"
	"github.com/bsm/sarama-cluster"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"os/signal"
	"time"
	"github.com/satori/go.uuid"
	"github.com/sk8sio/http-gateway/pkg/message"
	"syscall"
)

// Function messageHandler creates an http handler that posts the http body as a message to Kafka, replying
// immediately with a successful http response
func messageHandler(producer sarama.AsyncProducer) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		topic := r.URL.Path[len("/messages/"):]
		b, err := ioutil.ReadAll(r.Body)
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
		scsMessage := message.Message{Payload: b, Headers: nil}

		bytesOut, err := message.EncodeMessage(scsMessage)
		if err != nil {
			log.Printf("Error encoding message: %v", err)
			return
		}
		kafkaMsg := &sarama.ProducerMessage{Topic: topic, Value: sarama.ByteEncoder(bytesOut)}

		select {
		case producer.Input() <- kafkaMsg:
			w.Write([]byte("message published to topic: " + topic + "\n"))
		}
	}
}

// Function replyHandler creates an http handler that posts the http body as a message to Kafka, then waits
// for a message on a go channel it creates for a reply (this is expected to be set by the main thread) and sends
// that as an http response.
func replyHandler(producer sarama.AsyncProducer, replies map[string]chan message.Message) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		topic := r.URL.Path[len("/requests/"):]
		b, err := ioutil.ReadAll(r.Body)
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
		correlationId := uuid.NewV4().String()
		replyChan := make(chan message.Message)
		replies[correlationId] = replyChan
		scsMessage := message.Message{Payload: b, Headers: map[string]interface{}{"correlationId": correlationId}}

		bytesOut, err := message.EncodeMessage(scsMessage)
		if err != nil {
			log.Printf("Error encoding message: %v", err)
			return
		}
		kafkaMsg := &sarama.ProducerMessage{Topic: topic, Value: sarama.ByteEncoder(bytesOut)}

		select {
		case producer.Input() <- kafkaMsg:
			select {
			case reply := <-replyChan:
				delete(replies, correlationId)
				p := reply.Payload
				w.Write(p.([]byte)) // TODO equivalent of Spring's HttpMessageConverter handling
			case <- time.After(time.Second * 60):
				delete(replies, correlationId)
				w.WriteHeader(404)
			}
		}
	}
}

func healthHandler() func(http.ResponseWriter, *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte(`{"status":"UP"}`))
	}
}

func startHttpServer(producer sarama.AsyncProducer, replies map[string]chan message.Message) *http.Server {
	srv := &http.Server{Addr: ":8080"}

	http.HandleFunc("/messages/", messageHandler(producer))
	http.HandleFunc("/requests/", replyHandler(producer, replies))
	http.HandleFunc("/application/status", healthHandler())

	go func() {
		if err := srv.ListenAndServe(); err != nil {
			log.Printf("Httpserver: ListenAndServe() error: %s", err)
		}
	}()

	log.Printf("Listening on %v", srv.Addr)
	return srv
}

func main() {
	// Trap signals to trigger a proper shutdown.
	signals := make(chan os.Signal, 1)
	signal.Notify(signals, os.Interrupt, syscall.SIGTERM, os.Kill)

	// Key is correlationId, value is channel used to pass message received from main Kafka consumer loop
	replies := make(map[string]chan message.Message)

	brokers := []string{os.Getenv("SPRING_CLOUD_STREAM_KAFKA_BINDER_BROKERS")}
	producer, err := sarama.NewAsyncProducer(brokers, nil)
	if err != nil {
		panic(err)
	}
	defer func() {
		if err := producer.Close(); err != nil {
			log.Fatalln(err)
		}
	}()

	consumerConfig := makeConsumerConfig()
	consumer, err := cluster.NewConsumer(brokers, "gateway", []string{"replies"}, consumerConfig)
	if err != nil {
		panic(err)
	}
	defer consumer.Close()
	if consumerConfig.Consumer.Return.Errors {
		go consumeErrors(consumer)
	}
	if consumerConfig.Group.Return.Notifications {
		go consumeNotifications(consumer)
	}

	srv := startHttpServer(producer, replies)

MainLoop:
	for {
		select {
		case <-signals:
			fmt.Println("Shutting Down...")
			timeout, c := context.WithTimeout(context.Background(), 1*time.Second)
			defer c()
			if err := srv.Shutdown(timeout); err != nil {
				panic(err) // failure/timeout shutting down the server gracefully
			}
			break MainLoop
		case msg, ok := <-consumer.Messages():
			if ok {
				messageWithHeaders, err := message.ExtractMessage(msg.Value)
				if err != nil {
					log.Println("Failed to extract message ", err)
					break
				}
				correlationId, ok := messageWithHeaders.Headers["correlationId"].(string)
				if ok {
					c := replies[correlationId]
					if c != nil {
						fmt.Printf("Sending %v\n", messageWithHeaders)
						c <- messageWithHeaders
						consumer.MarkOffset(msg, "") // mark message as processed
					} else {
						log.Printf("Did not find communication channel for correlationId %v. Timed out?", correlationId)
						consumer.MarkOffset(msg, "") // mark message as processed
					}
				}
			}

		case err := <-producer.Errors():
			log.Println("Failed to produce kafka message ", err)
		}
	}
}

func makeConsumerConfig() *cluster.Config {
	consumerConfig := cluster.NewConfig()
	consumerConfig.Consumer.Return.Errors = true
	consumerConfig.Group.Return.Notifications = true
	return consumerConfig
}

func consumeNotifications(consumer *cluster.Consumer) {
	for ntf := range consumer.Notifications() {
		log.Printf("Rebalanced: %+v\n", ntf)
	}
}
func consumeErrors(consumer *cluster.Consumer) {
	for err := range consumer.Errors() {
		log.Printf("Error: %s\n", err.Error())
	}
}
<|endoftext|>


<commit_before>package db

import "fmt"

//ErrExists is the error that should be returned from a db function if an item
// could not be inserted because it already exists in the database
type ErrExists struct {
	message string
}

//NewErrExists makes a new ErrExists object... works a hell of a lot like fmt.Errorf
func NewErrExists(format string, args ...interface{}) error {
	return ErrExists{
		message: fmt.Sprintf(format, args...),
	}
}

func (e ErrExists) Error() string {
	return e.message
}
<commit_msg>Remove unused db/errors.go<commit_after><|endoftext|>