<commit_before>package unilog

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"os"
	"os/exec"
	"os/signal"
	"strings"
	"syscall"
	"text/template"
	"time"

	"github.com/DataDog/datadog-go/statsd"
	"github.com/getsentry/raven-go"

	"github.com/stripe/unilog/clevels"
	flag "launchpad.net/gnuflag"
)

// hold the argument passed in with "-statstags"
var statstags string

// A filter to be applied to log lines prior to prefixing them
// with a timestamp and logging them.
type Filter func(string) string

// Unilog represents a unilog process. unilog is intended to be used
// as a standalone application, but is exported as a package to allow
// users to perform compile-time configuration to simplify deployment.
type Unilog struct {
	// Sentry DSN for reporting Unilog errors
	// If this is unset, unilog will not report errors to Sentry
	SentryDSN string
	// StatsdAddress for sending metrics
	// If this is unset, it wlil default to "127.0.0.1:8200" -> TODO: is this what we want?
	StatsdAddress string
	// The email address from which unilog will send mail on
	// errors
	MailTo string
	// The email address to which unilog will email breakages. If
	// either MailTo or MailFrom is unset, unilog will not
	// generate email.
	MailFrom string

	// A series of filters which will be applied to each log line
	// in order
	Filters []Filter

	// The version that unilog will report on the command-line and
	// in error emails. Defaults to the toplevel Version constant.
	Version string
	// The number of log lines to buffer in-memory, in case
	// unilog's disk writer falls behind. Note that when talking
	// to unilog over a pipe, the kernel also maintains an
	// in-kernel pipe buffer, sized 64kb on Linux.
	BufferLines int

	Name    string
	Verbose bool
	Debug   bool

	lines     <-chan string
	errs      <-chan error
	sigReopen <-chan os.Signal
	sigTerm   <-chan os.Signal
	sigQuit   <-chan os.Signal
	shutdown  chan struct{}
	file      io.WriteCloser
	target    string

	b struct {
		broken bool
		at     time.Time
		count  int
	}

	exit           func(int)
	shouldShutdown bool
}

func stringFlag(val *string, longname, shortname, init, help string) {
	flag.StringVar(val, longname, init, help)
	flag.StringVar(val, shortname, init, help)
}

func boolFlag(val *bool, longname, shortname string, init bool, help string) {
	flag.BoolVar(val, longname, init, help)
	flag.BoolVar(val, shortname, init, help)
}

func (u *Unilog) fillDefaults() {
	u.exit = os.Exit
	if u.Version == "" {
		u.Version = Version
	}
	if u.BufferLines == 0 {
		u.BufferLines = DefaultBuffer
	}
}

func (u *Unilog) addFlags() {
	stringFlag(&u.Name, "name", "a", "", "Name of logged program")
	boolFlag(&u.Verbose, "verbose", "v", false, "Echo lines to stdout")
	boolFlag(&u.Debug, "debug", "d", false, "Print debug messages")
	flag.StringVar(&u.MailFrom, "mailfrom", u.MailFrom, "Address to send error emails from")
	flag.StringVar(&u.MailTo, "mailto", u.MailTo, "Address to send error emails to")
	flag.StringVar(&u.SentryDSN, "sentrydsn", u.SentryDSN, "Sentry DSN to send errors to")
	flag.StringVar(&u.StatsdAddress, "statsdaddress", "127.0.0.1:8200", "Address to send statsd metrics to")
	flag.StringVar(&clevels.AusterityFile, "austerityfile", clevels.AusterityFile, "(optional) Location of file to read austerity level from")
	stringFlag(&statstags, "statstags", "s", "", `(optional) tags to include with all statsd metrics (e.g. "foo:bar,baz:quz")`)
}

var emailTemplate = template.Must(template.New("email").Parse(`From: {{.From}}
To: {{.To}}
Subject: [unilog] {{.Name}} could not {{.Action}}

Hi there,

This is unilog reporting from {{.Hostname}}. I'm sad to report that
{{.Name}} is having some troubles writing to its log. I got caught up
trying to log a line to {{.Target}}.

To avoid spamming you, I'm going to shut up for an hour. Please fix me.

{{.Error}}
--
Sent from unilog {{.Version}}
`))

const (
	// Version is the Unilog version. Reported in emails and in
	// response to --version on the command line. Can be overriden
	// by the Version field in a Unilog object.
	Version = "0.3"
	// DefaultBuffer is the default size (in lines) of the
	// in-process line buffer
	DefaultBuffer = 1 << 12
)

var Stats *statsd.Client

func readlines(in io.Reader, bufsize int, shutdown chan struct{}) (<-chan string, <-chan error) {
	linec := make(chan string, bufsize)
	errc := make(chan error, 1)

	u := NewUnilogReader(in, shutdown)
	r := bufio.NewReader(u)

	go func() {
		var err error
		var s string

		for err == nil {
			s, err = r.ReadString('\n')
			if s != "" {
				s = strings.TrimRight(s, "\n")
				linec <- s
				if Stats != nil {
					Stats.Count("unilog.bytes", int64(len(s)), nil, .1)
				}
			}
		}

		if err != io.EOF {
			errc <- err
		}
		close(linec)
	}()

	return linec, errc
}

func (u *Unilog) reopen() error {
	if u.target == "-" {
		u.file = os.Stdout
		return nil
	}

	if u.file != nil {
		u.file.Close()
		u.file = nil
	}

	var e error
	if u.file, e = os.OpenFile(u.target, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0644); e != nil {
		u.file = nil
		return e
	}
	return nil
}

func (u *Unilog) format(line string) string {
	for _, filter := range u.Filters {
		if filter != nil {
			line = filter(line)
		}
	}
	return fmt.Sprintf("[%s] %s\n", time.Now().Format("2006-01-02 15:04:05.000000"), line)
}

func (u *Unilog) logLine(line string) {
	formatted := u.format(line)
	if u.Verbose {
		defer io.WriteString(os.Stdout, formatted)
	}

	var e error
	if u.file == nil {
		e = u.reopen()
	}
	if e != nil {
		u.handleError("reopen_file", e)
		return
	}
	_, e = io.WriteString(u.file, formatted)
	if e != nil {
		u.handleError("write_to_log", e)
	} else {
		u.b.broken = false
	}
}

func (u *Unilog) run() {
	for {
		u.selector()
	}
}

func (u *Unilog) selector() {
	select {
	case e := <-u.errs:
		if e != nil && e != io.EOF {
			panic(e)
		} else {
			return
		}
	case <-u.sigReopen:
		u.reopen()
	case <-u.sigTerm:
		select {
		case u.shutdown <- struct{}{}:
			u.shouldShutdown = true
		default:
		}
	case <-u.sigQuit:
		if u.shouldShutdown {
			u.exit(1)
			return
		}
	case line, ok := <-u.lines:
		if !ok {
			return
		}
		u.logLine(line)
	}
}

func (u *Unilog) handleError(action string, e error) {
	if !u.b.broken {
		u.b.broken = true
		u.b.at = time.Now()
		u.b.count = 0
	} else if time.Since(u.b.at) > time.Hour {
		u.b.at = time.Now()
		u.b.count = 0
	}

	if u.Debug {
		fmt.Fprintf(os.Stderr, "Could not %s: %s\n", action, e.Error())
	}

	if Stats != nil {
		emsg := fmt.Sprintf("err_action:%s", action)
		Stats.Count("unilog.errors_total", 1, []string{emsg}, 1)
	}

	if u.b.count == 0 && u.SentryDSN != "" {
		hostname, _ := os.Hostname()
		keys := map[string]string{
			"Hostname": hostname,
			"Action":   action,
			"Name":     u.Name,
			"Target":   u.target,
			"Error":    e.Error(),
			"Version":  Version,
		}
		raven.CaptureError(e, keys)
	}

	if u.b.count == 0 && u.MailFrom != "" && u.MailTo != "" {
		message := new(bytes.Buffer)
		hostname, _ := os.Hostname()
		emailTemplate.Execute(message, map[string]string{
			"Hostname": hostname,
			"From":     u.MailFrom,
			"To":       u.MailTo,
			"Action":   action,
			"Name":     u.Name,
			"Target":   u.target,
			"Error":    e.Error(),
			"Version":  Version,
		})
		cmd := exec.Command("sendmail", "-t")
		cmd.Stdin = message
		cmd.Run()
	}

	u.b.count++
}

func (u *Unilog) Main() {
	u.fillDefaults()

	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage: %s [options] dstfile\n", os.Args[0])
		flag.PrintDefaults()
	}

	u.addFlags()
	var flagVersion bool
	boolFlag(&flagVersion, "version", "V", false, "Print the version number and exit")

	flag.Parse(true)

	if flagVersion {
		fmt.Printf("This is unilog v%s\n", Version)
		return
	}
	args := flag.Args()
	if len(args) != 1 {
		flag.Usage()
		os.Exit(1)
	}

	reopen := make(chan os.Signal, 2)
	signal.Notify(reopen, syscall.SIGALRM, syscall.SIGHUP)
	u.sigReopen = reopen

	term := make(chan os.Signal, 2)
	signal.Notify(term, syscall.SIGTERM, syscall.SIGINT)
	u.sigTerm = term

	quit := make(chan os.Signal, 2)
	signal.Notify(quit, syscall.SIGQUIT)
	u.sigQuit = quit

	u.shutdown = make(chan struct{})
	u.target = flag.Arg(0)
	u.reopen()

	u.lines, u.errs = readlines(os.Stdin, u.BufferLines, u.shutdown)

	fileName := u.target

	Stats, _ = statsd.New(u.StatsdAddress)

	Stats.Tags = append(Stats.Tags, fmt.Sprintf("file_name:%s", fileName))
	if statstags != "" {
		Stats.Tags = append(Stats.Tags, strings.Split(statstags, ",")...)
	}

	clevels.Stats = Stats

	_ = raven.SetDSN(u.SentryDSN)

	u.run()
}
<commit_msg>Selector => tick<commit_after>package unilog

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"os"
	"os/exec"
	"os/signal"
	"strings"
	"syscall"
	"text/template"
	"time"

	"github.com/DataDog/datadog-go/statsd"
	"github.com/getsentry/raven-go"

	"github.com/stripe/unilog/clevels"
	flag "launchpad.net/gnuflag"
)

// hold the argument passed in with "-statstags"
var statstags string

// A filter to be applied to log lines prior to prefixing them
// with a timestamp and logging them.
type Filter func(string) string

// Unilog represents a unilog process. unilog is intended to be used
// as a standalone application, but is exported as a package to allow
// users to perform compile-time configuration to simplify deployment.
type Unilog struct {
	// Sentry DSN for reporting Unilog errors
	// If this is unset, unilog will not report errors to Sentry
	SentryDSN string
	// StatsdAddress for sending metrics
	// If this is unset, it wlil default to "127.0.0.1:8200" -> TODO: is this what we want?
	StatsdAddress string
	// The email address from which unilog will send mail on
	// errors
	MailTo string
	// The email address to which unilog will email breakages. If
	// either MailTo or MailFrom is unset, unilog will not
	// generate email.
	MailFrom string

	// A series of filters which will be applied to each log line
	// in order
	Filters []Filter

	// The version that unilog will report on the command-line and
	// in error emails. Defaults to the toplevel Version constant.
	Version string
	// The number of log lines to buffer in-memory, in case
	// unilog's disk writer falls behind. Note that when talking
	// to unilog over a pipe, the kernel also maintains an
	// in-kernel pipe buffer, sized 64kb on Linux.
	BufferLines int

	Name    string
	Verbose bool
	Debug   bool

	lines     <-chan string
	errs      <-chan error
	sigReopen <-chan os.Signal
	sigTerm   <-chan os.Signal
	sigQuit   <-chan os.Signal
	shutdown  chan struct{}
	file      io.WriteCloser
	target    string

	b struct {
		broken bool
		at     time.Time
		count  int
	}

	exit           func(int)
	shouldShutdown bool
}

func stringFlag(val *string, longname, shortname, init, help string) {
	flag.StringVar(val, longname, init, help)
	flag.StringVar(val, shortname, init, help)
}

func boolFlag(val *bool, longname, shortname string, init bool, help string) {
	flag.BoolVar(val, longname, init, help)
	flag.BoolVar(val, shortname, init, help)
}

func (u *Unilog) fillDefaults() {
	u.exit = os.Exit
	if u.Version == "" {
		u.Version = Version
	}
	if u.BufferLines == 0 {
		u.BufferLines = DefaultBuffer
	}
}

func (u *Unilog) addFlags() {
	stringFlag(&u.Name, "name", "a", "", "Name of logged program")
	boolFlag(&u.Verbose, "verbose", "v", false, "Echo lines to stdout")
	boolFlag(&u.Debug, "debug", "d", false, "Print debug messages")
	flag.StringVar(&u.MailFrom, "mailfrom", u.MailFrom, "Address to send error emails from")
	flag.StringVar(&u.MailTo, "mailto", u.MailTo, "Address to send error emails to")
	flag.StringVar(&u.SentryDSN, "sentrydsn", u.SentryDSN, "Sentry DSN to send errors to")
	flag.StringVar(&u.StatsdAddress, "statsdaddress", "127.0.0.1:8200", "Address to send statsd metrics to")
	flag.StringVar(&clevels.AusterityFile, "austerityfile", clevels.AusterityFile, "(optional) Location of file to read austerity level from")
	stringFlag(&statstags, "statstags", "s", "", `(optional) tags to include with all statsd metrics (e.g. "foo:bar,baz:quz")`)
}

var emailTemplate = template.Must(template.New("email").Parse(`From: {{.From}}
To: {{.To}}
Subject: [unilog] {{.Name}} could not {{.Action}}

Hi there,

This is unilog reporting from {{.Hostname}}. I'm sad to report that
{{.Name}} is having some troubles writing to its log. I got caught up
trying to log a line to {{.Target}}.

To avoid spamming you, I'm going to shut up for an hour. Please fix me.

{{.Error}}
--
Sent from unilog {{.Version}}
`))

const (
	// Version is the Unilog version. Reported in emails and in
	// response to --version on the command line. Can be overriden
	// by the Version field in a Unilog object.
	Version = "0.3"
	// DefaultBuffer is the default size (in lines) of the
	// in-process line buffer
	DefaultBuffer = 1 << 12
)

var Stats *statsd.Client

func readlines(in io.Reader, bufsize int, shutdown chan struct{}) (<-chan string, <-chan error) {
	linec := make(chan string, bufsize)
	errc := make(chan error, 1)

	u := NewUnilogReader(in, shutdown)
	r := bufio.NewReader(u)

	go func() {
		var err error
		var s string

		for err == nil {
			s, err = r.ReadString('\n')
			if s != "" {
				s = strings.TrimRight(s, "\n")
				linec <- s
				if Stats != nil {
					Stats.Count("unilog.bytes", int64(len(s)), nil, .1)
				}
			}
		}

		if err != io.EOF {
			errc <- err
		}
		close(linec)
	}()

	return linec, errc
}

func (u *Unilog) reopen() error {
	if u.target == "-" {
		u.file = os.Stdout
		return nil
	}

	if u.file != nil {
		u.file.Close()
		u.file = nil
	}

	var e error
	if u.file, e = os.OpenFile(u.target, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0644); e != nil {
		u.file = nil
		return e
	}
	return nil
}

func (u *Unilog) format(line string) string {
	for _, filter := range u.Filters {
		if filter != nil {
			line = filter(line)
		}
	}
	return fmt.Sprintf("[%s] %s\n", time.Now().Format("2006-01-02 15:04:05.000000"), line)
}

func (u *Unilog) logLine(line string) {
	formatted := u.format(line)
	if u.Verbose {
		defer io.WriteString(os.Stdout, formatted)
	}

	var e error
	if u.file == nil {
		e = u.reopen()
	}
	if e != nil {
		u.handleError("reopen_file", e)
		return
	}
	_, e = io.WriteString(u.file, formatted)
	if e != nil {
		u.handleError("write_to_log", e)
	} else {
		u.b.broken = false
	}
}

func (u *Unilog) run() {
	for {
		if !u.tick() {
			return
		}
	}
}

// "tick" is Unilog's event loop
// returns true if Unilog should keep running,
// and false if it should stop.
func (u *Unilog) tick() bool {
	select {
	case e := <-u.errs:
		if e != nil && e != io.EOF {
			panic(e)
		} else {
			return false
		}
	case <-u.sigReopen:
		u.reopen()
	case <-u.sigTerm:
		select {
		case u.shutdown <- struct{}{}:
			u.shouldShutdown = true
		default:
		}
	case <-u.sigQuit:
		if u.shouldShutdown {
			u.exit(1)
			return false
		}
	case line, ok := <-u.lines:
		if !ok {
			return false
		}
		u.logLine(line)
	}
	return true
}

func (u *Unilog) handleError(action string, e error) {
	if !u.b.broken {
		u.b.broken = true
		u.b.at = time.Now()
		u.b.count = 0
	} else if time.Since(u.b.at) > time.Hour {
		u.b.at = time.Now()
		u.b.count = 0
	}

	if u.Debug {
		fmt.Fprintf(os.Stderr, "Could not %s: %s\n", action, e.Error())
	}

	if Stats != nil {
		emsg := fmt.Sprintf("err_action:%s", action)
		Stats.Count("unilog.errors_total", 1, []string{emsg}, 1)
	}

	if u.b.count == 0 && u.SentryDSN != "" {
		hostname, _ := os.Hostname()
		keys := map[string]string{
			"Hostname": hostname,
			"Action":   action,
			"Name":     u.Name,
			"Target":   u.target,
			"Error":    e.Error(),
			"Version":  Version,
		}
		raven.CaptureError(e, keys)
	}

	if u.b.count == 0 && u.MailFrom != "" && u.MailTo != "" {
		message := new(bytes.Buffer)
		hostname, _ := os.Hostname()
		emailTemplate.Execute(message, map[string]string{
			"Hostname": hostname,
			"From":     u.MailFrom,
			"To":       u.MailTo,
			"Action":   action,
			"Name":     u.Name,
			"Target":   u.target,
			"Error":    e.Error(),
			"Version":  Version,
		})
		cmd := exec.Command("sendmail", "-t")
		cmd.Stdin = message
		cmd.Run()
	}

	u.b.count++
}

func (u *Unilog) Main() {
	u.fillDefaults()

	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage: %s [options] dstfile\n", os.Args[0])
		flag.PrintDefaults()
	}

	u.addFlags()
	var flagVersion bool
	boolFlag(&flagVersion, "version", "V", false, "Print the version number and exit")

	flag.Parse(true)

	if flagVersion {
		fmt.Printf("This is unilog v%s\n", Version)
		return
	}
	args := flag.Args()
	if len(args) != 1 {
		flag.Usage()
		os.Exit(1)
	}

	reopen := make(chan os.Signal, 2)
	signal.Notify(reopen, syscall.SIGALRM, syscall.SIGHUP)
	u.sigReopen = reopen

	term := make(chan os.Signal, 2)
	signal.Notify(term, syscall.SIGTERM, syscall.SIGINT)
	u.sigTerm = term

	quit := make(chan os.Signal, 2)
	signal.Notify(quit, syscall.SIGQUIT)
	u.sigQuit = quit

	u.shutdown = make(chan struct{})
	u.target = flag.Arg(0)
	u.reopen()

	u.lines, u.errs = readlines(os.Stdin, u.BufferLines, u.shutdown)

	fileName := u.target

	Stats, _ = statsd.New(u.StatsdAddress)

	Stats.Tags = append(Stats.Tags, fmt.Sprintf("file_name:%s", fileName))
	if statstags != "" {
		Stats.Tags = append(Stats.Tags, strings.Split(statstags, ",")...)
	}

	clevels.Stats = Stats

	_ = raven.SetDSN(u.SentryDSN)

	u.run()
}
<|endoftext|>


<commit_before>// +build linux

package main

import (
	"encoding/json"
	"fmt"
	"os"
	"strconv"

	"github.com/docker/go-units"
	"github.com/opencontainers/runtime-spec/specs-go"
	"github.com/urfave/cli"
)

func u64Ptr(i uint64) *uint64 { return &i }
func u16Ptr(i uint16) *uint16 { return &i }

var updateCommand = cli.Command{
	Name:      "update",
	Usage:     "update container resource constraints",
	ArgsUsage: `<container-id>`,
	Flags: []cli.Flag{
		cli.StringFlag{
			Name:  "resources, r",
			Value: "",
			Usage: `path to the file containing the resources to update or '-' to read from the standard input

The accepted format is as follow (unchanged values can be omitted):

{
  "memory": {
    "limit": 0,
    "reservation": 0,
    "swap": 0,
    "kernel": 0,
    "kernelTCP": 0
  },
  "cpu": {
    "shares": 0,
    "quota": 0,
    "period": 0,
    "cpus": "",
    "mems": ""
  },
  "blockIO": {
    "blkioWeight": 0
  },
}

Note: if data is to be read from a file or the standard input, all
other options are ignored.
`,
		},

		cli.IntFlag{
			Name:  "blkio-weight",
			Usage: "Specifies per cgroup weight, range is from 10 to 1000",
		},
		cli.StringFlag{
			Name:  "cpu-period",
			Usage: "CPU period to be used for hardcapping (in usecs). 0 to use system default",
		},
		cli.StringFlag{
			Name:  "cpu-quota",
			Usage: "CPU hardcap limit (in usecs). Allowed cpu time in a given period",
		},
		cli.StringFlag{
			Name:  "cpu-share",
			Usage: "CPU shares (relative weight vs. other containers)",
		},
		cli.StringFlag{
			Name:  "cpuset-cpus",
			Usage: "CPU(s) to use",
		},
		cli.StringFlag{
			Name:  "cpuset-mems",
			Usage: "Memory node(s) to use",
		},
		cli.StringFlag{
			Name:  "kernel-memory",
			Usage: "Kernel memory limit (in bytes)",
		},
		cli.StringFlag{
			Name:  "kernel-memory-tcp",
			Usage: "Kernel memory limit (in bytes) for tcp buffer",
		},
		cli.StringFlag{
			Name:  "memory",
			Usage: "Memory limit (in bytes)",
		},
		cli.StringFlag{
			Name:  "memory-reservation",
			Usage: "Memory reservation or soft_limit (in bytes)",
		},
		cli.StringFlag{
			Name:  "memory-swap",
			Usage: "Total memory usage (memory + swap); set `-1` to enable unlimited swap",
		},
	},
	Action: func(context *cli.Context) error {
		container, err := getContainer(context)
		if err != nil {
			return err
		}

		r := specs.Resources{
			Memory: &specs.Memory{
				Limit:       u64Ptr(0),
				Reservation: u64Ptr(0),
				Swap:        u64Ptr(0),
				Kernel:      u64Ptr(0),
				KernelTCP:   u64Ptr(0),
			},
			CPU: &specs.CPU{
				Shares: u64Ptr(0),
				Quota:  u64Ptr(0),
				Period: u64Ptr(0),
				Cpus:   sPtr(""),
				Mems:   sPtr(""),
			},
			BlockIO: &specs.BlockIO{
				Weight: u16Ptr(0),
			},
		}

		config := container.Config()

		if in := context.String("resources"); in != "" {
			var (
				f   *os.File
				err error
			)
			switch in {
			case "-":
				f = os.Stdin
			default:
				f, err = os.Open(in)
				if err != nil {
					return err
				}
			}
			err = json.NewDecoder(f).Decode(&r)
			if err != nil {
				return err
			}
		} else {
			if val := context.Int("blkio-weight"); val != 0 {
				r.BlockIO.Weight = u16Ptr(uint16(val))
			}
			if val := context.String("cpuset-cpus"); val != "" {
				r.CPU.Cpus = &val
			}
			if val := context.String("cpuset-mems"); val != "" {
				r.CPU.Mems = &val
			}

			for opt, dest := range map[string]*uint64{
				"cpu-period": r.CPU.Period,
				"cpu-quota":  r.CPU.Quota,
				"cpu-share":  r.CPU.Shares,
			} {
				if val := context.String(opt); val != "" {
					var err error
					*dest, err = strconv.ParseUint(val, 10, 64)
					if err != nil {
						return fmt.Errorf("invalid value for %s: %s", opt, err)
					}
				}
			}

			for opt, dest := range map[string]*uint64{
				"kernel-memory":      r.Memory.Kernel,
				"kernel-memory-tcp":  r.Memory.KernelTCP,
				"memory":             r.Memory.Limit,
				"memory-reservation": r.Memory.Reservation,
				"memory-swap":        r.Memory.Swap,
			} {
				if val := context.String(opt); val != "" {
					v, err := units.RAMInBytes(val)
					if err != nil {
						return fmt.Errorf("invalid value for %s: %s", opt, err)
					}
					*dest = uint64(v)
				}
			}
		}

		// Update the value
		config.Cgroups.Resources.BlkioWeight = *r.BlockIO.Weight
		config.Cgroups.Resources.CpuPeriod = int64(*r.CPU.Period)
		config.Cgroups.Resources.CpuQuota = int64(*r.CPU.Quota)
		config.Cgroups.Resources.CpuShares = int64(*r.CPU.Shares)
		config.Cgroups.Resources.CpusetCpus = *r.CPU.Cpus
		config.Cgroups.Resources.CpusetMems = *r.CPU.Mems
		config.Cgroups.Resources.KernelMemory = int64(*r.Memory.Kernel)
		config.Cgroups.Resources.KernelMemoryTCP = int64(*r.Memory.KernelTCP)
		config.Cgroups.Resources.Memory = int64(*r.Memory.Limit)
		config.Cgroups.Resources.MemoryReservation = int64(*r.Memory.Reservation)
		config.Cgroups.Resources.MemorySwap = int64(*r.Memory.Swap)

		if err := container.Set(config); err != nil {
			return err
		}
		return nil
	},
}
<commit_msg>Fix help message for memory-swap<commit_after>// +build linux

package main

import (
	"encoding/json"
	"fmt"
	"os"
	"strconv"

	"github.com/docker/go-units"
	"github.com/opencontainers/runtime-spec/specs-go"
	"github.com/urfave/cli"
)

func u64Ptr(i uint64) *uint64 { return &i }
func u16Ptr(i uint16) *uint16 { return &i }

var updateCommand = cli.Command{
	Name:      "update",
	Usage:     "update container resource constraints",
	ArgsUsage: `<container-id>`,
	Flags: []cli.Flag{
		cli.StringFlag{
			Name:  "resources, r",
			Value: "",
			Usage: `path to the file containing the resources to update or '-' to read from the standard input

The accepted format is as follow (unchanged values can be omitted):

{
  "memory": {
    "limit": 0,
    "reservation": 0,
    "swap": 0,
    "kernel": 0,
    "kernelTCP": 0
  },
  "cpu": {
    "shares": 0,
    "quota": 0,
    "period": 0,
    "cpus": "",
    "mems": ""
  },
  "blockIO": {
    "blkioWeight": 0
  },
}

Note: if data is to be read from a file or the standard input, all
other options are ignored.
`,
		},

		cli.IntFlag{
			Name:  "blkio-weight",
			Usage: "Specifies per cgroup weight, range is from 10 to 1000",
		},
		cli.StringFlag{
			Name:  "cpu-period",
			Usage: "CPU period to be used for hardcapping (in usecs). 0 to use system default",
		},
		cli.StringFlag{
			Name:  "cpu-quota",
			Usage: "CPU hardcap limit (in usecs). Allowed cpu time in a given period",
		},
		cli.StringFlag{
			Name:  "cpu-share",
			Usage: "CPU shares (relative weight vs. other containers)",
		},
		cli.StringFlag{
			Name:  "cpuset-cpus",
			Usage: "CPU(s) to use",
		},
		cli.StringFlag{
			Name:  "cpuset-mems",
			Usage: "Memory node(s) to use",
		},
		cli.StringFlag{
			Name:  "kernel-memory",
			Usage: "Kernel memory limit (in bytes)",
		},
		cli.StringFlag{
			Name:  "kernel-memory-tcp",
			Usage: "Kernel memory limit (in bytes) for tcp buffer",
		},
		cli.StringFlag{
			Name:  "memory",
			Usage: "Memory limit (in bytes)",
		},
		cli.StringFlag{
			Name:  "memory-reservation",
			Usage: "Memory reservation or soft_limit (in bytes)",
		},
		cli.StringFlag{
			Name:  "memory-swap",
			Usage: "Total memory usage (memory + swap); set '-1' to enable unlimited swap",
		},
	},
	Action: func(context *cli.Context) error {
		container, err := getContainer(context)
		if err != nil {
			return err
		}

		r := specs.Resources{
			Memory: &specs.Memory{
				Limit:       u64Ptr(0),
				Reservation: u64Ptr(0),
				Swap:        u64Ptr(0),
				Kernel:      u64Ptr(0),
				KernelTCP:   u64Ptr(0),
			},
			CPU: &specs.CPU{
				Shares: u64Ptr(0),
				Quota:  u64Ptr(0),
				Period: u64Ptr(0),
				Cpus:   sPtr(""),
				Mems:   sPtr(""),
			},
			BlockIO: &specs.BlockIO{
				Weight: u16Ptr(0),
			},
		}

		config := container.Config()

		if in := context.String("resources"); in != "" {
			var (
				f   *os.File
				err error
			)
			switch in {
			case "-":
				f = os.Stdin
			default:
				f, err = os.Open(in)
				if err != nil {
					return err
				}
			}
			err = json.NewDecoder(f).Decode(&r)
			if err != nil {
				return err
			}
		} else {
			if val := context.Int("blkio-weight"); val != 0 {
				r.BlockIO.Weight = u16Ptr(uint16(val))
			}
			if val := context.String("cpuset-cpus"); val != "" {
				r.CPU.Cpus = &val
			}
			if val := context.String("cpuset-mems"); val != "" {
				r.CPU.Mems = &val
			}

			for opt, dest := range map[string]*uint64{
				"cpu-period": r.CPU.Period,
				"cpu-quota":  r.CPU.Quota,
				"cpu-share":  r.CPU.Shares,
			} {
				if val := context.String(opt); val != "" {
					var err error
					*dest, err = strconv.ParseUint(val, 10, 64)
					if err != nil {
						return fmt.Errorf("invalid value for %s: %s", opt, err)
					}
				}
			}

			for opt, dest := range map[string]*uint64{
				"kernel-memory":      r.Memory.Kernel,
				"kernel-memory-tcp":  r.Memory.KernelTCP,
				"memory":             r.Memory.Limit,
				"memory-reservation": r.Memory.Reservation,
				"memory-swap":        r.Memory.Swap,
			} {
				if val := context.String(opt); val != "" {
					v, err := units.RAMInBytes(val)
					if err != nil {
						return fmt.Errorf("invalid value for %s: %s", opt, err)
					}
					*dest = uint64(v)
				}
			}
		}

		// Update the value
		config.Cgroups.Resources.BlkioWeight = *r.BlockIO.Weight
		config.Cgroups.Resources.CpuPeriod = int64(*r.CPU.Period)
		config.Cgroups.Resources.CpuQuota = int64(*r.CPU.Quota)
		config.Cgroups.Resources.CpuShares = int64(*r.CPU.Shares)
		config.Cgroups.Resources.CpusetCpus = *r.CPU.Cpus
		config.Cgroups.Resources.CpusetMems = *r.CPU.Mems
		config.Cgroups.Resources.KernelMemory = int64(*r.Memory.Kernel)
		config.Cgroups.Resources.KernelMemoryTCP = int64(*r.Memory.KernelTCP)
		config.Cgroups.Resources.Memory = int64(*r.Memory.Limit)
		config.Cgroups.Resources.MemoryReservation = int64(*r.Memory.Reservation)
		config.Cgroups.Resources.MemorySwap = int64(*r.Memory.Swap)

		if err := container.Set(config); err != nil {
			return err
		}
		return nil
	},
}
<|endoftext|>


<commit_before>package main

import (
	"compress/gzip"
	"crypto/sha1"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"syscall"
	"time"

	"github.com/dickeyxxx/golock"
)

var updateTopic = &Topic{
	Name:        "update",
	Description: "update heroku-cli",
}

var updateCmd = &Command{
	Topic:       "update",
	Description: "updates heroku-cli",
	Args:        []Arg{{Name: "channel", Optional: true}},
	Run: func(ctx *Context) {
		channel := ctx.Args.(map[string]string)["channel"]
		if channel == "" {
			channel = "master"
		}
		Update(channel)
	},
}

var binPath = filepath.Join(AppDir, "heroku-cli")
var updateLockPath = filepath.Join(AppDir, "updating.lock")
var autoupdateFile = filepath.Join(AppDir, "autoupdate")

func init() {
	if runtime.GOOS == "windows" {
		binPath = binPath + ".exe"
	}
}

// Update updates the CLI and plugins
func Update(channel string) {
	Logln("updating")
	if err := golock.Lock(updateLockPath); err != nil {
		panic(err)
	}
	defer golock.Unlock(updateLockPath)
	done := make(chan bool)
	go func() {
		updateCLI(channel)
		updatePlugins()
		done <- true
	}()
	select {
	case <-time.After(time.Second * 120):
		Errln("Timed out while updating")
		Logln("Timed out while updating")
	case <-done:
		touchAutoupdateFile()
	}
	Logln("done updating")
}

func updatePlugins() {
	Logln("updating plugins")
	b, _ := node.UpdatePackages()
	if len(b) > 0 {
		Logln("clearing plugins cache")
		ClearPluginCache()
		WritePluginCache(GetPlugins())
	}
}

func updateCLI(channel string) {
	manifest := getUpdateManifest(channel)
	if manifest.Version == Version && manifest.Channel == Channel {
		return
	}
	if !updatable() {
		Errf("Out of date: You are running %s but %s is out.\n", Version, manifest.Version)
		return
	}
	Logln("updating from %s to %s (%s)", Version, manifest.Version, manifest.Channel)
	build := manifest.Builds[runtime.GOOS][runtime.GOARCH]
	// on windows we can't remove an existing file or remove the running binary
	// so we download the file to binName.new
	// move the running binary to binName.old (deleting any existing file first)
	// rename the downloaded file to binName
	if err := downloadBin(binPath+".new", build.URL); err != nil {
		panic(err)
	}
	if fileSha1(binPath+".new") != build.Sha1 {
		panic("SHA mismatch")
	}
	os.Remove(binPath + ".old")
	if err := os.Rename(binPath, binPath+".old"); err != nil {
		panic(err)
	}
	if err := os.Rename(binPath+".new", binPath); err != nil {
		panic(err)
	}
	reexecBin()
}

// IsUpdateNeeded checks if an update is available
func IsUpdateNeeded(t string) bool {
	f, err := os.Stat(autoupdateFile)
	if err != nil {
		return true
	}
	if t == "soft" {
		return time.Since(f.ModTime()) > 1*time.Hour
	}
	return time.Since(f.ModTime()) > 168*time.Hour
}

func touchAutoupdateFile() {
	out, err := os.OpenFile(autoupdateFile, os.O_WRONLY|os.O_CREATE, 0644)
	if err != nil {
		panic(err)
	}
	out.WriteString(time.Now().String())
}

type manifest struct {
	Channel, Version string
	Builds           map[string]map[string]struct {
		URL, Sha1 string
	}
}

func getUpdateManifest(channel string) manifest {
	res, err := http.Get("https://d1gvo455cekpjp.cloudfront.net/" + channel + "/manifest.json")
	if err != nil {
		panic(err)
	}
	var m manifest
	json.NewDecoder(res.Body).Decode(&m)
	return m
}

func updatable() bool {
	path, err := filepath.Abs(os.Args[0])
	if err != nil {
		Errln(err)
	}
	return path == binPath
}

func downloadBin(path, url string) error {
	out, err := os.OpenFile(path, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0755)
	if err != nil {
		return err
	}
	client := &http.Client{}
	req, err := http.NewRequest("GET", url+".gz", nil)
	if err != nil {
		return err
	}
	req.Header.Add("Accept-Encoding", "gzip")
	resp, err := client.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	uncompressed, err := gzip.NewReader(resp.Body)
	if err != nil {
		return err
	}
	_, err = io.Copy(out, uncompressed)
	if err != nil {
		return err
	}
	return out.Close()
}

func fileSha1(path string) string {
	data, err := ioutil.ReadFile(path)
	if err != nil {
		panic(err)
	}
	return fmt.Sprintf("%x", sha1.Sum(data))
}

func reexecBin() {
	if runtime.GOOS == "windows" {
		cmd := exec.Command(binPath, os.Args[1:]...)
		cmd.Stdin = os.Stdin
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr
		cmd.Run()
		os.Exit(0)
	} else {
		if err := syscall.Exec(binPath, os.Args, os.Environ()); err != nil {
			panic(err)
		}
		os.Exit(99) // should never happen
	}
}

// TriggerBackgroundUpdate will trigger an update to the client in the background
func TriggerBackgroundUpdate() {
	if err := exec.Command("./heroku-cli", "update").Start(); err != nil {
		panic(err)
	}
}
<commit_msg>added note for updates<commit_after>package main

import (
	"compress/gzip"
	"crypto/sha1"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"syscall"
	"time"

	"github.com/dickeyxxx/golock"
)

var updateTopic = &Topic{
	Name:        "update",
	Description: "update heroku-cli",
}

var updateCmd = &Command{
	Topic:       "update",
	Description: "updates heroku-cli",
	Args:        []Arg{{Name: "channel", Optional: true}},
	Run: func(ctx *Context) {
		channel := ctx.Args.(map[string]string)["channel"]
		if channel == "" {
			channel = "master"
		}
		Update(channel)
	},
}

var binPath = filepath.Join(AppDir, "heroku-cli")
var updateLockPath = filepath.Join(AppDir, "updating.lock")
var autoupdateFile = filepath.Join(AppDir, "autoupdate")

func init() {
	if runtime.GOOS == "windows" {
		binPath = binPath + ".exe"
	}
}

// Update updates the CLI and plugins
func Update(channel string) {
	Errln("updating cli...")
	Logln("updating")
	if err := golock.Lock(updateLockPath); err != nil {
		panic(err)
	}
	defer golock.Unlock(updateLockPath)
	done := make(chan bool)
	go func() {
		updateCLI(channel)
		updatePlugins()
		done <- true
	}()
	select {
	case <-time.After(time.Second * 120):
		Errln("Timed out while updating")
		Logln("Timed out while updating")
	case <-done:
		touchAutoupdateFile()
	}
	Logln("done updating")
}

func updatePlugins() {
	Logln("updating plugins")
	b, _ := node.UpdatePackages()
	if len(b) > 0 {
		Logln("clearing plugins cache")
		ClearPluginCache()
		WritePluginCache(GetPlugins())
	}
}

func updateCLI(channel string) {
	manifest := getUpdateManifest(channel)
	if manifest.Version == Version && manifest.Channel == Channel {
		return
	}
	if !updatable() {
		Errf("Out of date: You are running %s but %s is out.\n", Version, manifest.Version)
		return
	}
	Logln("updating from %s to %s (%s)", Version, manifest.Version, manifest.Channel)
	build := manifest.Builds[runtime.GOOS][runtime.GOARCH]
	// on windows we can't remove an existing file or remove the running binary
	// so we download the file to binName.new
	// move the running binary to binName.old (deleting any existing file first)
	// rename the downloaded file to binName
	if err := downloadBin(binPath+".new", build.URL); err != nil {
		panic(err)
	}
	if fileSha1(binPath+".new") != build.Sha1 {
		panic("SHA mismatch")
	}
	os.Remove(binPath + ".old")
	if err := os.Rename(binPath, binPath+".old"); err != nil {
		panic(err)
	}
	if err := os.Rename(binPath+".new", binPath); err != nil {
		panic(err)
	}
	reexecBin()
}

// IsUpdateNeeded checks if an update is available
func IsUpdateNeeded(t string) bool {
	f, err := os.Stat(autoupdateFile)
	if err != nil {
		return true
	}
	if t == "soft" {
		return time.Since(f.ModTime()) > 1*time.Hour
	}
	return time.Since(f.ModTime()) > 168*time.Hour
}

func touchAutoupdateFile() {
	out, err := os.OpenFile(autoupdateFile, os.O_WRONLY|os.O_CREATE, 0644)
	if err != nil {
		panic(err)
	}
	out.WriteString(time.Now().String())
}

type manifest struct {
	Channel, Version string
	Builds           map[string]map[string]struct {
		URL, Sha1 string
	}
}

func getUpdateManifest(channel string) manifest {
	res, err := http.Get("https://d1gvo455cekpjp.cloudfront.net/" + channel + "/manifest.json")
	if err != nil {
		panic(err)
	}
	var m manifest
	json.NewDecoder(res.Body).Decode(&m)
	return m
}

func updatable() bool {
	path, err := filepath.Abs(os.Args[0])
	if err != nil {
		Errln(err)
	}
	return path == binPath
}

func downloadBin(path, url string) error {
	out, err := os.OpenFile(path, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0755)
	if err != nil {
		return err
	}
	client := &http.Client{}
	req, err := http.NewRequest("GET", url+".gz", nil)
	if err != nil {
		return err
	}
	req.Header.Add("Accept-Encoding", "gzip")
	resp, err := client.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	uncompressed, err := gzip.NewReader(resp.Body)
	if err != nil {
		return err
	}
	_, err = io.Copy(out, uncompressed)
	if err != nil {
		return err
	}
	return out.Close()
}

func fileSha1(path string) string {
	data, err := ioutil.ReadFile(path)
	if err != nil {
		panic(err)
	}
	return fmt.Sprintf("%x", sha1.Sum(data))
}

func reexecBin() {
	if runtime.GOOS == "windows" {
		cmd := exec.Command(binPath, os.Args[1:]...)
		cmd.Stdin = os.Stdin
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr
		cmd.Run()
		os.Exit(0)
	} else {
		if err := syscall.Exec(binPath, os.Args, os.Environ()); err != nil {
			panic(err)
		}
		os.Exit(99) // should never happen
	}
}

// TriggerBackgroundUpdate will trigger an update to the client in the background
func TriggerBackgroundUpdate() {
	if err := exec.Command("./heroku-cli", "update").Start(); err != nil {
		panic(err)
	}
}
<|endoftext|>


<commit_before>// Copyright (c) 2017 VMware, Inc. All Rights Reserved.
//
// 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 dao

import (
	"fmt"
	"strconv"
	"testing"

	"github.com/stretchr/testify/require"
	"github.com/vmware/harbor/src/common/models"
)

var (
	project    = "library"
	name       = "library/repository-test"
	repository = &models.RepoRecord{
		Name:      name,
		ProjectID: 1,
	}
)

func TestGetRepositoryByProjectName(t *testing.T) {
	if err := addRepository(repository); err != nil {
		t.Fatalf("failed to add repository %s: %v", name, err)
	}
	defer func() {
		if err := deleteRepository(name); err != nil {
			t.Fatalf("failed to delete repository %s: %v", name, err)
		}
	}()

	repositories, err := GetRepositoryByProjectName(project)
	if err != nil {
		t.Fatalf("failed to get repositories of project %s: %v",
			project, err)
	}

	if len(repositories) == 0 {
		t.Fatal("unexpected length of repositories: 0, at least 1")
	}

	exist := false
	for _, repo := range repositories {
		if repo.Name == name {
			exist = true
			break
		}
	}
	if !exist {
		t.Errorf("there is no repository whose name is %s", name)
	}
}

func TestGetTotalOfRepositories(t *testing.T) {
	total, err := GetTotalOfRepositories("")
	if err != nil {
		t.Fatalf("failed to get total of repositoreis: %v", err)
	}

	if err := addRepository(repository); err != nil {
		t.Fatalf("failed to add repository %s: %v", name, err)
	}
	defer func() {
		if err := deleteRepository(name); err != nil {
			t.Fatalf("failed to delete repository %s: %v", name, err)
		}
	}()

	n, err := GetTotalOfRepositories("")
	if err != nil {
		t.Fatalf("failed to get total of repositoreis: %v", err)
	}

	if n != total+1 {
		t.Errorf("unexpected total: %d != %d", n, total+1)
	}
}

func TestGetTopRepos(t *testing.T) {
	var err error
	require := require.New(t)

	require.NoError(GetOrmer().Begin())
	defer func() {
		require.NoError(GetOrmer().Rollback())
	}()

	projectIDs := []int64{}

	project1 := models.Project{
		OwnerID: 1,
		Name:    "project1",
		Public:  0,
	}
	project1.ProjectID, err = AddProject(project1)
	require.NoError(err)
	projectIDs = append(projectIDs, project1.ProjectID)

	project2 := models.Project{
		OwnerID: 1,
		Name:    "project2",
		Public:  0,
	}
	project2.ProjectID, err = AddProject(project2)
	require.NoError(err)
	projectIDs = append(projectIDs, project2.ProjectID)

	repository1 := &models.RepoRecord{
		Name:      fmt.Sprintf("%v/repository1", project1.Name),
		ProjectID: project1.ProjectID,
	}
	err = AddRepository(*repository1)
	require.NoError(err)
	require.NoError(IncreasePullCount(repository1.Name))

	repository2 := &models.RepoRecord{
		Name:      fmt.Sprintf("%v/repository2", project1.Name),
		ProjectID: project1.ProjectID,
	}
	err = AddRepository(*repository2)
	require.NoError(err)
	require.NoError(IncreasePullCount(repository2.Name))
	require.NoError(IncreasePullCount(repository2.Name))

	repository3 := &models.RepoRecord{
		Name:      fmt.Sprintf("%v/repository3", project2.Name),
		ProjectID: project2.ProjectID,
	}
	err = AddRepository(*repository3)
	require.NoError(err)
	require.NoError(IncreasePullCount(repository3.Name))
	require.NoError(IncreasePullCount(repository3.Name))
	require.NoError(IncreasePullCount(repository3.Name))

	topRepos, err := GetTopRepos(projectIDs, 100)
	require.NoError(err)
	require.Len(topRepos, 3)
	require.Equal(topRepos[0].Name, repository3.Name)
}

func TestGetTotalOfRepositoriesByProject(t *testing.T) {
	var projectID int64 = 1
	repoName := "library/total_count"

	total, err := GetTotalOfRepositoriesByProject([]int64{projectID}, repoName)
	if err != nil {
		t.Errorf("failed to get total of repositoreis of project %d: %v", projectID, err)
		return
	}

	if err := addRepository(&models.RepoRecord{
		Name:      repoName,
		ProjectID: projectID,
	}); err != nil {
		t.Errorf("failed to add repository %s: %v", repoName, err)
		return
	}
	defer func() {
		if err := deleteRepository(repoName); err != nil {
			t.Errorf("failed to delete repository %s: %v", name, err)
			return
		}
	}()

	n, err := GetTotalOfRepositoriesByProject([]int64{projectID}, repoName)
	if err != nil {
		t.Errorf("failed to get total of repositoreis of project %d: %v", projectID, err)
		return
	}

	if n != total+1 {
		t.Errorf("unexpected total: %d != %d", n, total+1)
	}
}

func TestGetRepositoriesByProject(t *testing.T) {
	var projectID int64 = 1
	repoName := "library/repository"

	if err := addRepository(&models.RepoRecord{
		Name:      repoName,
		ProjectID: projectID,
	}); err != nil {
		t.Errorf("failed to add repository %s: %v", repoName, err)
		return
	}
	defer func() {
		if err := deleteRepository(repoName); err != nil {
			t.Errorf("failed to delete repository %s: %v", name, err)
			return
		}
	}()

	repositories, err := GetRepositoriesByProject(projectID, repoName, 10, 0)
	if err != nil {
		t.Errorf("failed to get repositoreis of project %d: %v", projectID, err)
		return
	}

	t.Log(repositories)

	for _, repository := range repositories {
		if repository.Name == repoName {
			return
		}
	}

	t.Errorf("repository %s not found", repoName)
}

func TestGetAllRepositories(t *testing.T) {
	var reposBefore []*models.RepoRecord
	reposBefore, err := GetAllRepositories()
	if err != nil {
		t.Fatalf("error occurred while getting all registories. %v", err)
		return
	}
	allBefore := len(reposBefore)

	for i := 0; i < 1200; i++ {
		t := strconv.Itoa(i)
		repoRecord := models.RepoRecord{
			Name:      "test" + t,
			ProjectID: 1,
		}
		if err := AddRepository(repoRecord); err != nil {
			t.Fatalf("Error happens when adding the repository: %v", err)
			return
		}
	}

	var reposAfter []*models.RepoRecord
	reposAfter, err2 := GetAllRepositories()
	if err2 != nil {
		t.Fatalf("error occurred while getting all registories. %v", err2)
		return
	}
	allAfter := len(reposAfter)

	if allAfter != allBefore+1200 {
		t.Errorf("unexpected total: %d != %d", allAfter, allBefore+1200)
	}
}

func addRepository(repository *models.RepoRecord) error {
	return AddRepository(*repository)
}

func deleteRepository(name string) error {
	return DeleteRepository(name)
}
<commit_msg>update<commit_after>// Copyright (c) 2017 VMware, Inc. All Rights Reserved.
//
// 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 dao

import (
	"fmt"
	"strconv"
	"testing"

	"github.com/stretchr/testify/require"
	"github.com/vmware/harbor/src/common/models"
)

var (
	project    = "library"
	name       = "library/repository-test"
	repository = &models.RepoRecord{
		Name:      name,
		ProjectID: 1,
	}
)

func TestGetRepositoryByProjectName(t *testing.T) {
	if err := addRepository(repository); err != nil {
		t.Fatalf("failed to add repository %s: %v", name, err)
	}
	defer func() {
		if err := deleteRepository(name); err != nil {
			t.Fatalf("failed to delete repository %s: %v", name, err)
		}
	}()

	repositories, err := GetRepositoryByProjectName(project)
	if err != nil {
		t.Fatalf("failed to get repositories of project %s: %v",
			project, err)
	}

	if len(repositories) == 0 {
		t.Fatal("unexpected length of repositories: 0, at least 1")
	}

	exist := false
	for _, repo := range repositories {
		if repo.Name == name {
			exist = true
			break
		}
	}
	if !exist {
		t.Errorf("there is no repository whose name is %s", name)
	}
}

func TestGetTotalOfRepositories(t *testing.T) {
	total, err := GetTotalOfRepositories("")
	if err != nil {
		t.Fatalf("failed to get total of repositoreis: %v", err)
	}

	if err := addRepository(repository); err != nil {
		t.Fatalf("failed to add repository %s: %v", name, err)
	}
	defer func() {
		if err := deleteRepository(name); err != nil {
			t.Fatalf("failed to delete repository %s: %v", name, err)
		}
	}()

	n, err := GetTotalOfRepositories("")
	if err != nil {
		t.Fatalf("failed to get total of repositoreis: %v", err)
	}

	if n != total+1 {
		t.Errorf("unexpected total: %d != %d", n, total+1)
	}
}

func TestGetTopRepos(t *testing.T) {
	var err error
	require := require.New(t)

	require.NoError(GetOrmer().Begin())
	defer func() {
		require.NoError(GetOrmer().Rollback())
	}()

	projectIDs := []int64{}

	project1 := models.Project{
		OwnerID: 1,
		Name:    "project1",
		Public:  0,
	}
	project1.ProjectID, err = AddProject(project1)
	require.NoError(err)
	projectIDs = append(projectIDs, project1.ProjectID)

	project2 := models.Project{
		OwnerID: 1,
		Name:    "project2",
		Public:  0,
	}
	project2.ProjectID, err = AddProject(project2)
	require.NoError(err)
	projectIDs = append(projectIDs, project2.ProjectID)

	repository1 := &models.RepoRecord{
		Name:      fmt.Sprintf("%v/repository1", project1.Name),
		ProjectID: project1.ProjectID,
	}
	err = AddRepository(*repository1)
	require.NoError(err)
	require.NoError(IncreasePullCount(repository1.Name))

	repository2 := &models.RepoRecord{
		Name:      fmt.Sprintf("%v/repository2", project1.Name),
		ProjectID: project1.ProjectID,
	}
	err = AddRepository(*repository2)
	require.NoError(err)
	require.NoError(IncreasePullCount(repository2.Name))
	require.NoError(IncreasePullCount(repository2.Name))

	repository3 := &models.RepoRecord{
		Name:      fmt.Sprintf("%v/repository3", project2.Name),
		ProjectID: project2.ProjectID,
	}
	err = AddRepository(*repository3)
	require.NoError(err)
	require.NoError(IncreasePullCount(repository3.Name))
	require.NoError(IncreasePullCount(repository3.Name))
	require.NoError(IncreasePullCount(repository3.Name))

	topRepos, err := GetTopRepos(projectIDs, 100)
	require.NoError(err)
	require.Len(topRepos, 3)
	require.Equal(topRepos[0].Name, repository3.Name)
}

func TestGetTotalOfRepositoriesByProject(t *testing.T) {
	var projectID int64 = 1
	repoName := "library/total_count"

	total, err := GetTotalOfRepositoriesByProject([]int64{projectID}, repoName)
	if err != nil {
		t.Errorf("failed to get total of repositoreis of project %d: %v", projectID, err)
		return
	}

	if err := addRepository(&models.RepoRecord{
		Name:      repoName,
		ProjectID: projectID,
	}); err != nil {
		t.Errorf("failed to add repository %s: %v", repoName, err)
		return
	}
	defer func() {
		if err := deleteRepository(repoName); err != nil {
			t.Errorf("failed to delete repository %s: %v", name, err)
			return
		}
	}()

	n, err := GetTotalOfRepositoriesByProject([]int64{projectID}, repoName)
	if err != nil {
		t.Errorf("failed to get total of repositoreis of project %d: %v", projectID, err)
		return
	}

	if n != total+1 {
		t.Errorf("unexpected total: %d != %d", n, total+1)
	}
}

func TestGetRepositoriesByProject(t *testing.T) {
	var projectID int64 = 1
	repoName := "library/repository"

	if err := addRepository(&models.RepoRecord{
		Name:      repoName,
		ProjectID: projectID,
	}); err != nil {
		t.Errorf("failed to add repository %s: %v", repoName, err)
		return
	}
	defer func() {
		if err := deleteRepository(repoName); err != nil {
			t.Errorf("failed to delete repository %s: %v", name, err)
			return
		}
	}()

	repositories, err := GetRepositoriesByProject(projectID, repoName, 10, 0)
	if err != nil {
		t.Errorf("failed to get repositoreis of project %d: %v", projectID, err)
		return
	}

	t.Log(repositories)

	for _, repository := range repositories {
		if repository.Name == repoName {
			return
		}
	}

	t.Errorf("repository %s not found", repoName)
}

func TestGetAllRepositories(t *testing.T) {
	var reposBefore []*models.RepoRecord
	reposBefore, err := GetAllRepositories()
	if err != nil {
		t.Fatalf("error occurred while getting all registories. %v", err)
		return
	}
	allBefore := len(reposBefore)

	for i := 0; i < 1200; i++ {
		end := strconv.Itoa(i)
		repoRecord := models.RepoRecord{
			Name:      "test" + end,
			ProjectID: 1,
		}
		if err := AddRepository(repoRecord); err != nil {
			t.Fatalf("Error happens when adding the repository: %v", err)
			return
		}
	}

	var reposAfter []*models.RepoRecord
	reposAfter, err2 := GetAllRepositories()
	if err2 != nil {
		t.Fatalf("error occurred while getting all registories. %v", err2)
		return
	}
	allAfter := len(reposAfter)

	if allAfter != allBefore+1200 {
		t.Errorf("unexpected total: %d != %d", allAfter, allBefore+1200)
	}
}

func addRepository(repository *models.RepoRecord) error {
	return AddRepository(*repository)
}

func deleteRepository(name string) error {
	return DeleteRepository(name)
}
<|endoftext|>


<commit_before>// Automatically generated by MockGen. DO NOT EDIT!
// Source: tracker.go

package state

import (
	gomock "gomock.googlecode.com/hg/gomock"
)

// Mock of StateTracker interface
type MockStateTracker struct {
	ctrl     *gomock.Controller
	recorder *_MockStateTrackerRecorder
}

// Recorder for MockStateTracker (not exported)
type _MockStateTrackerRecorder struct {
	mock *MockStateTracker
}

func NewMockStateTracker(ctrl *gomock.Controller) *MockStateTracker {
	mock := &MockStateTracker{ctrl: ctrl}
	mock.recorder = &_MockStateTrackerRecorder{mock}
	return mock
}

func (m *MockStateTracker) EXPECT() *_MockStateTrackerRecorder {
	return m.recorder
}

func (m *MockStateTracker) NewNick(nick string) *Nick {
	ret := m.ctrl.Call(m, "NewNick", nick)
	ret0, _ := ret[0].(*Nick)
	return ret0
}

func (mr *_MockStateTrackerRecorder) NewNick(arg0 interface{}) *gomock.Call {
	return mr.mock.ctrl.RecordCall(mr.mock, "NewNick", arg0)
}

func (m *MockStateTracker) GetNick(nick string) *Nick {
	ret := m.ctrl.Call(m, "GetNick", nick)
	ret0, _ := ret[0].(*Nick)
	return ret0
}

func (mr *_MockStateTrackerRecorder) GetNick(arg0 interface{}) *gomock.Call {
	return mr.mock.ctrl.RecordCall(mr.mock, "GetNick", arg0)
}

func (m *MockStateTracker) ReNick(old string, neu string) {
	m.ctrl.Call(m, "ReNick", old, neu)
}

func (mr *_MockStateTrackerRecorder) ReNick(arg0, arg1 interface{}) *gomock.Call {
	return mr.mock.ctrl.RecordCall(mr.mock, "ReNick", arg0, arg1)
}

func (m *MockStateTracker) DelNick(nick string) {
	m.ctrl.Call(m, "DelNick", nick)
}

func (mr *_MockStateTrackerRecorder) DelNick(arg0 interface{}) *gomock.Call {
	return mr.mock.ctrl.RecordCall(mr.mock, "DelNick", arg0)
}

func (m *MockStateTracker) NewChannel(channel string) *Channel {
	ret := m.ctrl.Call(m, "NewChannel", channel)
	ret0, _ := ret[0].(*Channel)
	return ret0
}

func (mr *_MockStateTrackerRecorder) NewChannel(arg0 interface{}) *gomock.Call {
	return mr.mock.ctrl.RecordCall(mr.mock, "NewChannel", arg0)
}

func (m *MockStateTracker) GetChannel(channel string) *Channel {
	ret := m.ctrl.Call(m, "GetChannel", channel)
	ret0, _ := ret[0].(*Channel)
	return ret0
}

func (mr *_MockStateTrackerRecorder) GetChannel(arg0 interface{}) *gomock.Call {
	return mr.mock.ctrl.RecordCall(mr.mock, "GetChannel", arg0)
}

func (m *MockStateTracker) DelChannel(channel string) {
	m.ctrl.Call(m, "DelChannel", channel)
}

func (mr *_MockStateTrackerRecorder) DelChannel(arg0 interface{}) *gomock.Call {
	return mr.mock.ctrl.RecordCall(mr.mock, "DelChannel", arg0)
}

func (m *MockStateTracker) Me() *Nick {
	ret := m.ctrl.Call(m, "Me")
	ret0, _ := ret[0].(*Nick)
	return ret0
}

func (mr *_MockStateTrackerRecorder) Me() *gomock.Call {
	return mr.mock.ctrl.RecordCall(mr.mock, "Me")
}

func (m *MockStateTracker) IsOn(channel string, nick string) (*ChanPrivs, bool) {
	ret := m.ctrl.Call(m, "IsOn", channel, nick)
	ret0, _ := ret[0].(*ChanPrivs)
	ret1, _ := ret[1].(bool)
	return ret0, ret1
}

func (mr *_MockStateTrackerRecorder) IsOn(arg0, arg1 interface{}) *gomock.Call {
	return mr.mock.ctrl.RecordCall(mr.mock, "IsOn", arg0, arg1)
}

func (m *MockStateTracker) Associate(channel *Channel, nick *Nick) *ChanPrivs {
	ret := m.ctrl.Call(m, "Associate", channel, nick)
	ret0, _ := ret[0].(*ChanPrivs)
	return ret0
}

func (mr *_MockStateTrackerRecorder) Associate(arg0, arg1 interface{}) *gomock.Call {
	return mr.mock.ctrl.RecordCall(mr.mock, "Associate", arg0, arg1)
}

func (m *MockStateTracker) Dissociate(channel *Channel, nick *Nick) {
	m.ctrl.Call(m, "Dissociate", channel, nick)
}

func (mr *_MockStateTrackerRecorder) Dissociate(arg0, arg1 interface{}) *gomock.Call {
	return mr.mock.ctrl.RecordCall(mr.mock, "Dissociate", arg0, arg1)
}

func (m *MockStateTracker) Wipe() {
	m.ctrl.Call(m, "Wipe")
}

func (mr *_MockStateTrackerRecorder) Wipe() *gomock.Call {
	return mr.mock.ctrl.RecordCall(mr.mock, "Wipe")
}

func (m *MockStateTracker) String() string {
	ret := m.ctrl.Call(m, "String")
	ret0, _ := ret[0].(string)
	return ret0
}

func (mr *_MockStateTrackerRecorder) String() *gomock.Call {
	return mr.mock.ctrl.RecordCall(mr.mock, "String")
}
<commit_msg>Update weekly to reflect the new location of gomock.<commit_after>// Automatically generated by MockGen. DO NOT EDIT!
// Source: tracker.go

package state

import (
	gomock "github.com/dsymonds/gomock/gomock"
)

// Mock of StateTracker interface
type MockStateTracker struct {
	ctrl     *gomock.Controller
	recorder *_MockStateTrackerRecorder
}

// Recorder for MockStateTracker (not exported)
type _MockStateTrackerRecorder struct {
	mock *MockStateTracker
}

func NewMockStateTracker(ctrl *gomock.Controller) *MockStateTracker {
	mock := &MockStateTracker{ctrl: ctrl}
	mock.recorder = &_MockStateTrackerRecorder{mock}
	return mock
}

func (_m *MockStateTracker) EXPECT() *_MockStateTrackerRecorder {
	return _m.recorder
}

func (_m *MockStateTracker) NewNick(nick string) *Nick {
	ret := _m.ctrl.Call(_m, "NewNick", nick)
	ret0, _ := ret[0].(*Nick)
	return ret0
}

func (_mr *_MockStateTrackerRecorder) NewNick(arg0 interface{}) *gomock.Call {
	return _mr.mock.ctrl.RecordCall(_mr.mock, "NewNick", arg0)
}

func (_m *MockStateTracker) GetNick(nick string) *Nick {
	ret := _m.ctrl.Call(_m, "GetNick", nick)
	ret0, _ := ret[0].(*Nick)
	return ret0
}

func (_mr *_MockStateTrackerRecorder) GetNick(arg0 interface{}) *gomock.Call {
	return _mr.mock.ctrl.RecordCall(_mr.mock, "GetNick", arg0)
}

func (_m *MockStateTracker) ReNick(old string, neu string) {
	_m.ctrl.Call(_m, "ReNick", old, neu)
}

func (_mr *_MockStateTrackerRecorder) ReNick(arg0, arg1 interface{}) *gomock.Call {
	return _mr.mock.ctrl.RecordCall(_mr.mock, "ReNick", arg0, arg1)
}

func (_m *MockStateTracker) DelNick(nick string) {
	_m.ctrl.Call(_m, "DelNick", nick)
}

func (_mr *_MockStateTrackerRecorder) DelNick(arg0 interface{}) *gomock.Call {
	return _mr.mock.ctrl.RecordCall(_mr.mock, "DelNick", arg0)
}

func (_m *MockStateTracker) NewChannel(channel string) *Channel {
	ret := _m.ctrl.Call(_m, "NewChannel", channel)
	ret0, _ := ret[0].(*Channel)
	return ret0
}

func (_mr *_MockStateTrackerRecorder) NewChannel(arg0 interface{}) *gomock.Call {
	return _mr.mock.ctrl.RecordCall(_mr.mock, "NewChannel", arg0)
}

func (_m *MockStateTracker) GetChannel(channel string) *Channel {
	ret := _m.ctrl.Call(_m, "GetChannel", channel)
	ret0, _ := ret[0].(*Channel)
	return ret0
}

func (_mr *_MockStateTrackerRecorder) GetChannel(arg0 interface{}) *gomock.Call {
	return _mr.mock.ctrl.RecordCall(_mr.mock, "GetChannel", arg0)
}

func (_m *MockStateTracker) DelChannel(channel string) {
	_m.ctrl.Call(_m, "DelChannel", channel)
}

func (_mr *_MockStateTrackerRecorder) DelChannel(arg0 interface{}) *gomock.Call {
	return _mr.mock.ctrl.RecordCall(_mr.mock, "DelChannel", arg0)
}

func (_m *MockStateTracker) Me() *Nick {
	ret := _m.ctrl.Call(_m, "Me")
	ret0, _ := ret[0].(*Nick)
	return ret0
}

func (_mr *_MockStateTrackerRecorder) Me() *gomock.Call {
	return _mr.mock.ctrl.RecordCall(_mr.mock, "Me")
}

func (_m *MockStateTracker) IsOn(channel string, nick string) (*ChanPrivs, bool) {
	ret := _m.ctrl.Call(_m, "IsOn", channel, nick)
	ret0, _ := ret[0].(*ChanPrivs)
	ret1, _ := ret[1].(bool)
	return ret0, ret1
}

func (_mr *_MockStateTrackerRecorder) IsOn(arg0, arg1 interface{}) *gomock.Call {
	return _mr.mock.ctrl.RecordCall(_mr.mock, "IsOn", arg0, arg1)
}

func (_m *MockStateTracker) Associate(channel *Channel, nick *Nick) *ChanPrivs {
	ret := _m.ctrl.Call(_m, "Associate", channel, nick)
	ret0, _ := ret[0].(*ChanPrivs)
	return ret0
}

func (_mr *_MockStateTrackerRecorder) Associate(arg0, arg1 interface{}) *gomock.Call {
	return _mr.mock.ctrl.RecordCall(_mr.mock, "Associate", arg0, arg1)
}

func (_m *MockStateTracker) Dissociate(channel *Channel, nick *Nick) {
	_m.ctrl.Call(_m, "Dissociate", channel, nick)
}

func (_mr *_MockStateTrackerRecorder) Dissociate(arg0, arg1 interface{}) *gomock.Call {
	return _mr.mock.ctrl.RecordCall(_mr.mock, "Dissociate", arg0, arg1)
}

func (_m *MockStateTracker) Wipe() {
	_m.ctrl.Call(_m, "Wipe")
}

func (_mr *_MockStateTrackerRecorder) Wipe() *gomock.Call {
	return _mr.mock.ctrl.RecordCall(_mr.mock, "Wipe")
}

func (_m *MockStateTracker) String() string {
	ret := _m.ctrl.Call(_m, "String")
	ret0, _ := ret[0].(string)
	return ret0
}

func (_mr *_MockStateTrackerRecorder) String() *gomock.Call {
	return _mr.mock.ctrl.RecordCall(_mr.mock, "String")
}
<|endoftext|>


<commit_before>/*

 */
package spaassets

import (
	"fmt"
	"net/http"
	"path/filepath"
	"strings"
)

// DeepLink middleware handles deep-linking into a single page application (SPA).
//
// This middleware should wrap your asset handler, and expects all requests are
// requesting a file with an extention. (ie: .js, .html, .png, .jpg etc)
//
// If it recieves a request targeting a directory (with or without a trailing '/'),
// it will rewrite the request to target appURL
//
// appURL should be the location of the SPA
//
func DeepLink(next http.Handler, appURL string) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if ext := filepath.Ext(r.URL.Path); ext == "" {
			r.URL.Path = appURL
		}
		next.ServeHTTP(w, r)
	})
}

// NoStoreCacheHandler applies Cache-control: no-store header according to options.
//
// Options
// include: indicates if matches should be included/excluded.
// prefix:  prefix is matched to the beginning of url.
//
func NoStoreCacheHandler(next http.Handler, include bool, prefix string) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if strings.HasPrefix(r.URL.Path, prefix) == include {
			w.Header().Set("Cache-control", "no-store")
		}
		next.ServeHTTP(w, r)
	})
}

// MaxAgeCacheHandler applies Cache-control: max-age=<age> header according to options.
//
// Options
// age: cache experation in seconds
// include: indicates if matches should be included/excluded.
// prefix:  prefix is matched to the beginning of url.
//
func MaxAgeCacheHandler(next http.Handler, age int, include bool, prefix string) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if strings.HasPrefix(r.URL.Path, prefix) == include {
			w.Header().Set("Cache-control", fmt.Sprintf("max-age=%d", age))
		}
		next.ServeHTTP(w, r)
	})
}
<commit_msg>Handle deep linking for Angular 2 router.  The new router has optional parameters that are separated by a semicolon (;) instead of question mark (?).  Handle this by removing everything after the first semicolon.  Don't think this will affect other valid routes. (#3)<commit_after>/*

 */
package spaassets

import (
	"fmt"
	"net/http"
	"path/filepath"
	"strings"
)

// DeepLink middleware handles deep-linking into a single page application (SPA).
//
// This middleware should wrap your asset handler, and expects all requests are
// requesting a file with an extention. (ie: .js, .html, .png, .jpg etc)
//
// If it recieves a request targeting a directory (with or without a trailing '/'),
// it will rewrite the request to target appURL
//
// appURL should be the location of the SPA
//
func DeepLink(next http.Handler, appURL string) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		path := r.URL.Path

		// Handle Angular 2 optional route parameters by removing them.
		if i := strings.Index(path, ";"); i > -1 {
			path = path[:i]
		}

		if ext := filepath.Ext(path); ext == "" {
			r.URL.Path = appURL
		}
		next.ServeHTTP(w, r)
	})
}

// NoStoreCacheHandler applies Cache-control: no-store header according to options.
//
// Options
// include: indicates if matches should be included/excluded.
// prefix:  prefix is matched to the beginning of url.
//
func NoStoreCacheHandler(next http.Handler, include bool, prefix string) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if strings.HasPrefix(r.URL.Path, prefix) == include {
			w.Header().Set("Cache-control", "no-store")
		}
		next.ServeHTTP(w, r)
	})
}

// MaxAgeCacheHandler applies Cache-control: max-age=<age> header according to options.
//
// Options
// age: cache experation in seconds
// include: indicates if matches should be included/excluded.
// prefix:  prefix is matched to the beginning of url.
//
func MaxAgeCacheHandler(next http.Handler, age int, include bool, prefix string) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if strings.HasPrefix(r.URL.Path, prefix) == include {
			w.Header().Set("Cache-control", fmt.Sprintf("max-age=%d", age))
		}
		next.ServeHTTP(w, r)
	})
}
<|endoftext|>


<commit_before>package cloud

import (
	"errors"
	"fmt"
	"net/http"

	"golang.org/x/net/context"
	"golang.org/x/oauth2"
	"golang.org/x/oauth2/google"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"
)

type dialOpt struct {
	endpoint string
	scopes   []string

	tokenSource oauth2.TokenSource

	httpClient *http.Client
	grpcClient *grpc.ClientConn
}

// ClientOption is used when construct clients for each cloud service.
type ClientOption interface {
	resolve(*dialOpt)
}

// WithTokenSource returns a ClientOption that specifies an OAuth2 token
// source to be used as the basis for authentication.
func WithTokenSource(s oauth2.TokenSource) ClientOption {
	return withTokenSource{s}
}

type withTokenSource struct{ ts oauth2.TokenSource }

func (w withTokenSource) resolve(o *dialOpt) {
	o.tokenSource = w.ts
}

// WithEndpoint returns a ClientOption that overrides the default endpoint
// to be used for a service.
func WithEndpoint(url string) ClientOption {
	return withEndpoint(url)
}

type withEndpoint string

func (w withEndpoint) resolve(o *dialOpt) {
	o.endpoint = string(w)
}

// WithScopes returns a ClientOption that overrides the default OAuth2 scopes
// to be used for a service.
func WithScopes(scope ...string) ClientOption {
	return withScopes(scope)
}

type withScopes []string

func (w withScopes) resolve(o *dialOpt) {
	o.scopes = []string(w)
}

// WithBaseHTTP returns a ClientOption that specifies the HTTP client to
// use as the basis of communications. This option may only be used with
// services that support HTTP as their communication transport.
func WithBaseHTTP(client *http.Client) ClientOption {
	return withBaseHTTP{client}
}

type withBaseHTTP struct{ client *http.Client }

func (w withBaseHTTP) resolve(o *dialOpt) {
	o.httpClient = w.client
}

// WithBaseGRPC returns a ClientOption that specifies the GRPC client
// connection to use as the basis of communications. This option many only be
// used with services that support HRPC as their communication transport.
func WithBaseGRPC(client *grpc.ClientConn) ClientOption {
	return withBaseGRPC{client}
}

type withBaseGRPC struct{ client *grpc.ClientConn }

func (w withBaseGRPC) resolve(o *dialOpt) {
	o.grpcClient = w.client
}

// DialHTTP returns an HTTP client for use communicating with a Google cloud
// service, configured with the given ClientOptions. Most developers should
// call the relevant NewClient method for the target service rather than
// invoking DialHTTP directly.
func DialHTTP(ctx context.Context, opt ...ClientOption) (*http.Client, error) {
	var o dialOpt
	for _, opt := range opt {
		opt.resolve(&o)
	}
	if o.grpcClient != nil {
		return nil, errors.New("unsupported GRPC base transport specified")
	}
	// TODO(djd): Wrap all http.Client's with appropriate internal version to add
	// UserAgent header and prepend correct endpoint.
	if o.httpClient != nil {
		return o.httpClient, nil
	}
	if o.tokenSource == nil {
		var err error
		o.tokenSource, err = google.DefaultTokenSource(ctx, o.scopes...)
		if err != nil {
			return nil, fmt.Errorf("google.DefaultTokenSource: %v", err)
		}
	}
	return oauth2.NewClient(ctx, o.tokenSource), nil
}

// DialGRPC returns a GRPC connection for use communicating with a Google cloud
// service, configured with the given ClientOptions. Most developers should
// call the relevant NewClient method for the target service rather than
// invoking DialGRPC directly.
func DialGRPC(ctx context.Context, opt ...ClientOption) (*grpc.ClientConn, error) {
	var o dialOpt
	for _, opt := range opt {
		opt.resolve(&o)
	}
	if o.httpClient != nil {
		return nil, errors.New("unsupported HTTP base transport specified")
	}
	if o.grpcClient != nil {
		return o.grpcClient, nil
	}
	if o.tokenSource == nil {
		var err error
		o.tokenSource, err = google.DefaultTokenSource(ctx, o.scopes...)
		if err != nil {
			return nil, fmt.Errorf("google.DefaultTokenSource: %v", err)
		}
	}
	grpcOpts := []grpc.DialOption{
		grpc.WithPerRPCCredentials(credentials.TokenSource{o.tokenSource}),
		grpc.WithTransportCredentials(credentials.NewClientTLSFromCert(nil, "")),
	}
	return grpc.Dial(o.endpoint, grpcOpts...)
}
<commit_msg>cloud: Add licence information to option.go<commit_after>/*
Copyright 2015 Google Inc. All Rights Reserved.

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 cloud

import (
	"errors"
	"fmt"
	"net/http"

	"golang.org/x/net/context"
	"golang.org/x/oauth2"
	"golang.org/x/oauth2/google"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"
)

type dialOpt struct {
	endpoint string
	scopes   []string

	tokenSource oauth2.TokenSource

	httpClient *http.Client
	grpcClient *grpc.ClientConn
}

// ClientOption is used when construct clients for each cloud service.
type ClientOption interface {
	resolve(*dialOpt)
}

// WithTokenSource returns a ClientOption that specifies an OAuth2 token
// source to be used as the basis for authentication.
func WithTokenSource(s oauth2.TokenSource) ClientOption {
	return withTokenSource{s}
}

type withTokenSource struct{ ts oauth2.TokenSource }

func (w withTokenSource) resolve(o *dialOpt) {
	o.tokenSource = w.ts
}

// WithEndpoint returns a ClientOption that overrides the default endpoint
// to be used for a service.
func WithEndpoint(url string) ClientOption {
	return withEndpoint(url)
}

type withEndpoint string

func (w withEndpoint) resolve(o *dialOpt) {
	o.endpoint = string(w)
}

// WithScopes returns a ClientOption that overrides the default OAuth2 scopes
// to be used for a service.
func WithScopes(scope ...string) ClientOption {
	return withScopes(scope)
}

type withScopes []string

func (w withScopes) resolve(o *dialOpt) {
	o.scopes = []string(w)
}

// WithBaseHTTP returns a ClientOption that specifies the HTTP client to
// use as the basis of communications. This option may only be used with
// services that support HTTP as their communication transport.
func WithBaseHTTP(client *http.Client) ClientOption {
	return withBaseHTTP{client}
}

type withBaseHTTP struct{ client *http.Client }

func (w withBaseHTTP) resolve(o *dialOpt) {
	o.httpClient = w.client
}

// WithBaseGRPC returns a ClientOption that specifies the GRPC client
// connection to use as the basis of communications. This option many only be
// used with services that support HRPC as their communication transport.
func WithBaseGRPC(client *grpc.ClientConn) ClientOption {
	return withBaseGRPC{client}
}

type withBaseGRPC struct{ client *grpc.ClientConn }

func (w withBaseGRPC) resolve(o *dialOpt) {
	o.grpcClient = w.client
}

// DialHTTP returns an HTTP client for use communicating with a Google cloud
// service, configured with the given ClientOptions. Most developers should
// call the relevant NewClient method for the target service rather than
// invoking DialHTTP directly.
func DialHTTP(ctx context.Context, opt ...ClientOption) (*http.Client, error) {
	var o dialOpt
	for _, opt := range opt {
		opt.resolve(&o)
	}
	if o.grpcClient != nil {
		return nil, errors.New("unsupported GRPC base transport specified")
	}
	// TODO(djd): Wrap all http.Client's with appropriate internal version to add
	// UserAgent header and prepend correct endpoint.
	if o.httpClient != nil {
		return o.httpClient, nil
	}
	if o.tokenSource == nil {
		var err error
		o.tokenSource, err = google.DefaultTokenSource(ctx, o.scopes...)
		if err != nil {
			return nil, fmt.Errorf("google.DefaultTokenSource: %v", err)
		}
	}
	return oauth2.NewClient(ctx, o.tokenSource), nil
}

// DialGRPC returns a GRPC connection for use communicating with a Google cloud
// service, configured with the given ClientOptions. Most developers should
// call the relevant NewClient method for the target service rather than
// invoking DialGRPC directly.
func DialGRPC(ctx context.Context, opt ...ClientOption) (*grpc.ClientConn, error) {
	var o dialOpt
	for _, opt := range opt {
		opt.resolve(&o)
	}
	if o.httpClient != nil {
		return nil, errors.New("unsupported HTTP base transport specified")
	}
	if o.grpcClient != nil {
		return o.grpcClient, nil
	}
	if o.tokenSource == nil {
		var err error
		o.tokenSource, err = google.DefaultTokenSource(ctx, o.scopes...)
		if err != nil {
			return nil, fmt.Errorf("google.DefaultTokenSource: %v", err)
		}
	}
	grpcOpts := []grpc.DialOption{
		grpc.WithPerRPCCredentials(credentials.TokenSource{o.tokenSource}),
		grpc.WithTransportCredentials(credentials.NewClientTLSFromCert(nil, "")),
	}
	return grpc.Dial(o.endpoint, grpcOpts...)
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/sha1"
	"database/sql"
	"encoding/base64"
	"encoding/json"
	"encoding/xml"
	"fmt"
	"io"
	"net/http"
	"os"
	"path/filepath"
	"strconv"
	"time"

	"github.com/dchest/uniuri"
	_ "github.com/go-sql-driver/mysql"
)

const (
	LENGTH     = 6
	PORT       = ":8080"
	DIRECTORY  = "/tmp/"
	UPADDRESS  = "http://localhost"
	dbUSERNAME = ""
	dbNAME     = ""
	dbPASSWORD = ""
	DATABASE   = dbUSERNAME + ":" + dbPASSWORD + "@/" + dbNAME + "?charset=utf8"
	MAXSIZE    = 10 * 1024 * 1024
)

type Result struct {
	URL  string `json:"url"`
	Name string `json:"name"`
	Hash string `json:"hash"`
	Size int64  `json:"size"`
}

type Response struct {
	Success     bool     `json:"success"`
	ErrorCode   int      `json:"errorcode,omitempty"`
	Description string   `json:"description,omitempty"`
	Files       []Result `json:"files,omitempty"`
}

func check(err error) {
	if err != nil {
		fmt.Println(err)
		return
	}
}

func generateName() string {
	name := uniuri.NewLen(LENGTH)
	db, err := sql.Open("mysql", DATABASE)
	check(err)
	query, err := db.Query("select id from files where id=?", name)
	if err != sql.ErrNoRows {
		for query.Next() {
			generateName()
		}
	}
	db.Close()

	return name
}
func respond(w http.ResponseWriter, output string, resp Response) {
	if resp.ErrorCode != 0 {
		resp.Files = []Result{}
		resp.Success = false
	} else {
		resp.Success = true
	}

	switch output {
	case "json":

		w.Header().Set("Content-Type", "application/json")
		err := json.NewEncoder(w).Encode(resp)

		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
	case "xml":
		x, err := xml.MarshalIndent(resp, "", "  ")
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}

		w.Header().Set("Content-Type", "application/xml")
		w.Write(x)

	case "html":
		w.Header().Set("Content-Type", "text/html")
		for _, file := range resp.Files {
			io.WriteString(w, ""+file.URL+"
")
		}

	case "gyazo", "text":
		w.Header().Set("Content-Type", "plain/text")
		for _, file := range resp.Files {
			io.WriteString(w, file.URL+"\n")
		}

	case "csv":
		w.Header().Set("Content-Type", "text/csv")
		io.WriteString(w, "name, url, hash, size\n")
		for _, file := range resp.Files {
			io.WriteString(w, file.Name+","+file.URL+","+file.Hash+","+strconv.FormatInt(file.Size, 10)+"\n")
		}

	default:
		w.Header().Set("Content-Type", "application/json")
		err := json.NewEncoder(w).Encode(resp)

		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}

	}

}
func grillHandler(w http.ResponseWriter, r *http.Request) {
}
func uploadHandler(w http.ResponseWriter, r *http.Request) {
	output := r.FormValue("output")
	resp := Response{Files: []Result{}}
	reader, err := r.MultipartReader()
	if err != nil {
		resp.ErrorCode = http.StatusInternalServerError
		resp.Description = err.Error()
		resp.Success = false
		respond(w, output, resp)
		return
	}
	db, err := sql.Open("mysql", DATABASE)
	check(err)
	for {
		part, err := reader.NextPart()
		if err == io.EOF {
			break
		}

		if part.FileName() == "" {
			continue
		}

		s := generateName()
		extName := filepath.Ext(part.FileName())
		filename := s + extName
		dst, err := os.Create(DIRECTORY + filename)
		defer dst.Close()

		if err != nil {
			resp.ErrorCode = http.StatusInternalServerError
			resp.Description = err.Error()
			respond(w, output, resp)
			return
		}

		h := sha1.New()
		t := io.TeeReader(part, h)
		_, err = io.Copy(dst, t)

		if err != nil {
			resp.ErrorCode = http.StatusInternalServerError
			resp.Description = err.Error()
			respond(w, output, resp)
			return
		}
		hash := h.Sum(nil)
		sha1 := base64.URLEncoding.EncodeToString(hash)
		stat, _ := dst.Stat()
		size := stat.Size()
		if size > MAXSIZE {
			resp.ErrorCode = http.StatusRequestEntityTooLarge
			resp.Description = err.Error()
			respond(w, output, resp)
			return
		}

		originalname := part.FileName()
		err = db.QueryRow("select originalname, filename, size from files where hash=?", sha1).Scan(&originalname, &filename, &size)
		res := Result{
			URL:  UPADDRESS + "/" + filename,
			Name: originalname,
			Hash: sha1,
			Size: size,
		}
		if err == sql.ErrNoRows {
			query, err := db.Prepare("INSERT into files(hash, originalname, filename, size, date) values(?, ?, ?, ?, ?)")
			check(err)
			_, err = query.Exec(res.Hash, res.Name, filename, res.Size, time.Now().Format("2016-01-02"))
			check(err)
		}
		resp.Files = append(resp.Files, res)
	}
	respond(w, output, resp)
}

func main() {
	http.HandleFunc("/upload.php", uploadHandler)
	http.HandleFunc("/grill.php", grillHandler)
	err := http.ListenAndServe(PORT, nil)
	if err != nil {
		panic(err)
	}

}
<commit_msg>Remove respond<commit_after>package main

import (
	"crypto/sha1"
	"database/sql"
	"encoding/base64"
	"encoding/json"
	"encoding/xml"
	"fmt"
	"io"
	"net/http"
	"os"
	"path/filepath"
	"strconv"
	"time"

	"github.com/dchest/uniuri"
	_ "github.com/go-sql-driver/mysql"
)

const (
	LENGTH     = 6
	PORT       = ":8080"
	DIRECTORY  = "/tmp/"
	UPADDRESS  = "http://localhost"
	dbUSERNAME = ""
	dbNAME     = ""
	dbPASSWORD = ""
	DATABASE   = dbUSERNAME + ":" + dbPASSWORD + "@/" + dbNAME + "?charset=utf8"
	MAXSIZE    = 10 * 1024 * 1024
)

type Result struct {
	URL  string `json:"url"`
	Name string `json:"name"`
	Hash string `json:"hash"`
	Size int64  `json:"size"`
}

type Response struct {
	Success     bool     `json:"success"`
	ErrorCode   int      `json:"errorcode,omitempty"`
	Description string   `json:"description,omitempty"`
	Files       []Result `json:"files,omitempty"`
}

func check(err error) {
	if err != nil {
		fmt.Println(err)
		return
	}
}

func generateName() string {
	name := uniuri.NewLen(LENGTH)
	db, err := sql.Open("mysql", DATABASE)
	check(err)
	query, err := db.Query("select id from files where id=?", name)
	if err != sql.ErrNoRows {
		for query.Next() {
			generateName()
		}
	}
	db.Close()

	return name
}
func respond(w http.ResponseWriter, output string, resp Response) {
	if resp.ErrorCode != 0 {
		resp.Files = []Result{}
		resp.Success = false
	} else {
		resp.Success = true
	}

	switch output {
	case "json":

		w.Header().Set("Content-Type", "application/json")
		err := json.NewEncoder(w).Encode(resp)

		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
	case "xml":
		x, err := xml.MarshalIndent(resp, "", "  ")
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}

		w.Header().Set("Content-Type", "application/xml")
		w.Write(x)

	case "html":
		w.Header().Set("Content-Type", "text/html")
		for _, file := range resp.Files {
			io.WriteString(w, ""+file.URL+"
")
		}

	case "gyazo", "text":
		w.Header().Set("Content-Type", "plain/text")
		for _, file := range resp.Files {
			io.WriteString(w, file.URL+"\n")
		}

	case "csv":
		w.Header().Set("Content-Type", "text/csv")
		io.WriteString(w, "name, url, hash, size\n")
		for _, file := range resp.Files {
			io.WriteString(w, file.Name+","+file.URL+","+file.Hash+","+strconv.FormatInt(file.Size, 10)+"\n")
		}

	default:
		w.Header().Set("Content-Type", "application/json")
		err := json.NewEncoder(w).Encode(resp)

		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}

	}

}
func grillHandler(w http.ResponseWriter, r *http.Request) {
}
func uploadHandler(w http.ResponseWriter, r *http.Request) {
	output := r.FormValue("output")
	resp := Response{Files: []Result{}}
	reader, err := r.MultipartReader()
	if err != nil {
		resp.ErrorCode = http.StatusInternalServerError
		resp.Description = err.Error()
		return
	}
	db, err := sql.Open("mysql", DATABASE)
	check(err)
	for {
		part, err := reader.NextPart()
		if err == io.EOF {
			break
		}

		if part.FileName() == "" {
			continue
		}

		s := generateName()
		extName := filepath.Ext(part.FileName())
		filename := s + extName
		dst, err := os.Create(DIRECTORY + filename)
		defer dst.Close()

		if err != nil {
			resp.ErrorCode = http.StatusInternalServerError
			resp.Description = err.Error()
			respond(w, output, resp)
			return
		}

		h := sha1.New()
		t := io.TeeReader(part, h)
		_, err = io.Copy(dst, t)

		if err != nil {
			resp.ErrorCode = http.StatusInternalServerError
			resp.Description = err.Error()
			respond(w, output, resp)
			return
		}
		hash := h.Sum(nil)
		sha1 := base64.URLEncoding.EncodeToString(hash)
		stat, _ := dst.Stat()
		size := stat.Size()
		if size > MAXSIZE {
			resp.ErrorCode = http.StatusRequestEntityTooLarge
			resp.Description = err.Error()
			respond(w, output, resp)
			return
		}

		originalname := part.FileName()
		err = db.QueryRow("select originalname, filename, size from files where hash=?", sha1).Scan(&originalname, &filename, &size)
		res := Result{
			URL:  UPADDRESS + "/" + filename,
			Name: originalname,
			Hash: sha1,
			Size: size,
		}
		if err == sql.ErrNoRows {
			query, err := db.Prepare("INSERT into files(hash, originalname, filename, size, date) values(?, ?, ?, ?, ?)")
			check(err)
			_, err = query.Exec(res.Hash, res.Name, filename, res.Size, time.Now().Format("2016-01-02"))
			check(err)
		}
		resp.Files = append(resp.Files, res)
	}
	respond(w, output, resp)
}

func main() {
	http.HandleFunc("/upload.php", uploadHandler)
	http.HandleFunc("/grill.php", grillHandler)
	err := http.ListenAndServe(PORT, nil)
	if err != nil {
		panic(err)
	}

}
<|endoftext|>


<commit_before>/*
Copyright 2021 The Vitess Authors.

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

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

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

package integration

import (
	"encoding/json"
	"fmt"
	"math/rand"
	"os"
	"strings"
	"testing"
	"time"

	"vitess.io/vitess/go/mysql/collations"
	"vitess.io/vitess/go/vt/sqlparser"
	"vitess.io/vitess/go/vt/vtgate/evalengine"
)

type (
	gencase struct {
		rand         *rand.Rand
		ratioTuple   int
		ratioSubexpr int
		tupleLen     int

		operators  []string
		primitives []string
	}

	dummyCollation collations.ID
)

func (g *gencase) arg(tuple bool) string {
	if tuple || g.rand.Intn(g.ratioTuple) == 0 {
		var exprs []string
		for i := 0; i < g.tupleLen; i++ {
			exprs = append(exprs, g.arg(false))
		}
		return fmt.Sprintf("(%s)", strings.Join(exprs, ", "))
	}
	if g.rand.Intn(g.ratioSubexpr) == 0 {
		return fmt.Sprintf("(%s)", g.expr())
	}
	return g.primitives[g.rand.Intn(len(g.primitives))]
}

func (g *gencase) expr() string {
	op := g.operators[g.rand.Intn(len(g.operators))]
	return fmt.Sprintf("%s %s %s", g.arg(false), op, g.arg(op == "IN" || op == "NOT IN"))
}

func (d dummyCollation) ColumnLookup(_ *sqlparser.ColName) (int, error) {
	panic("not supported")
}

func (d dummyCollation) CollationIDLookup(_ sqlparser.Expr) collations.ID {
	return collations.ID(d)
}

func TestTypes(t *testing.T) {
	var conn = mysqlconn(t)
	defer conn.Close()

	var queries = []string{
		"1 > 3",
		"3 > 1",
		"-1 > -1",
		"1 = 1",
		"-1 = 1",
		"1 IN (1, -2, 3)",
		"1 LIKE 1",
		"-1 LIKE -1",
		"-1 LIKE 1",
		`"foo" IN ("bar", "FOO", "baz")`,
		`'pokemon' LIKE 'poke%'`,
		`(1, 2) = (1, 2)`,
		`1 = 'sad'`,
		`(1, 2) = (1, 3)`,
	}

	for _, query := range queries {
		remote, err := conn.ExecuteFetch("SELECT "+query, 1, false)
		if err != nil {
			t.Fatal(err)
		}
		t.Logf("%s => %s", query, remote.Rows[0][0])
	}
}

func TestGenerateFuzzCases(t *testing.T) {
	const Total = 100
	const Seed = 1234

	type evaltest struct {
		Query string
		Value string `json:",omitempty"`
		Error string `json:",omitempty"`
	}

	var golden []evaltest
	var gen = gencase{
		rand:         rand.New(rand.NewSource(Seed)),
		ratioTuple:   8,
		ratioSubexpr: 8,
		tupleLen:     4,
		operators: []string{
			"+", "-", "/", "*", "=", "!=", "<=>", "<", "<=", ">", ">=", "IN", "NOT IN", "LIKE", "NOT LIKE",
		},
		primitives: []string{
			"1", "0", "-1", `"foo"`, `"FOO"`, `"fOo"`, "NULL",
		},
	}

	var conn = mysqlconn(t)
	defer conn.Close()

	for len(golden) < Total {
		query := "SELECT " + gen.expr()

		stmt, err := sqlparser.Parse(query)
		if err != nil {
			t.Fatalf("bad codegen: %v", err)
		}

		var eval evalengine.EvalResult
		var evaluated bool

		astExpr := stmt.(*sqlparser.Select).SelectExprs[0].(*sqlparser.AliasedExpr).Expr
		local, localErr := func() (expr evalengine.Expr, err error) {
			defer func() {
				if r := recover(); r != nil {
					err = fmt.Errorf("PANIC: %v", r)
				}
			}()
			expr, err = evalengine.ConvertEx(astExpr, dummyCollation(45), true)
			return
		}()
		if localErr == nil {
			evaluated = true
			eval, localErr = func() (eval evalengine.EvalResult, err error) {
				defer func() {
					if r := recover(); r != nil {
						err = fmt.Errorf("PANIC: %v", r)
					}
				}()
				eval, err = local.Evaluate(nil)
				return
			}()
		}

		remote, remoteErr := conn.ExecuteFetch(query, 1, false)

		if localErr != nil {
			if remoteErr == nil {
				t.Errorf("local query %q failed: %v (eval=%v); mysql response: %s", query, localErr, evaluated, remote.Rows[0][0].String())
				goto failed
			}
			if !strings.Contains(remoteErr.Error(), localErr.Error()) {
				t.Errorf("mismatch in errors for %q: local=%q (eval=%v), remote=%q", query, localErr.Error(), evaluated, remoteErr.Error())
				goto failed
			}
			continue
		}

		if remoteErr != nil {
			t.Errorf("remote query %q failed: %v, local=%s", query, remoteErr.Error(), eval.Value().String())
			goto failed
		}

		if eval.Value().String() != remote.Rows[0][0].String() {
			t.Errorf("mismatch for query %q: local=%v, remote=%v", query, eval.Value().String(), remote.Rows[0][0].String())
			goto failed
		}

		// OK -- all match
		continue

	failed:
		if remoteErr != nil {
			golden = append(golden, evaltest{
				Query: query,
				Error: remoteErr.Error(),
			})
		} else {
			golden = append(golden, evaltest{
				Query: query,
				Value: remote.Rows[0][0].String(),
			})
		}
	}

	out, err := os.Create(fmt.Sprintf("testdata/mysql_golden_%d.json", time.Now().Unix()))
	if err != nil {
		t.Fatal(err)
	}
	defer out.Close()

	enc := json.NewEncoder(out)
	enc.SetIndent("", "    ")
	enc.Encode(golden)
}
<commit_msg>evalengine: do not run fuzz in CI by default<commit_after>/*
Copyright 2021 The Vitess Authors.

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

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

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

package integration

import (
	"encoding/json"
	"flag"
	"fmt"
	"math/rand"
	"os"
	"strings"
	"testing"
	"time"

	"vitess.io/vitess/go/mysql/collations"
	"vitess.io/vitess/go/vt/sqlparser"
	"vitess.io/vitess/go/vt/vtgate/evalengine"
)

type (
	gencase struct {
		rand         *rand.Rand
		ratioTuple   int
		ratioSubexpr int
		tupleLen     int

		operators  []string
		primitives []string
	}

	dummyCollation collations.ID
)

func (g *gencase) arg(tuple bool) string {
	if tuple || g.rand.Intn(g.ratioTuple) == 0 {
		var exprs []string
		for i := 0; i < g.tupleLen; i++ {
			exprs = append(exprs, g.arg(false))
		}
		return fmt.Sprintf("(%s)", strings.Join(exprs, ", "))
	}
	if g.rand.Intn(g.ratioSubexpr) == 0 {
		return fmt.Sprintf("(%s)", g.expr())
	}
	return g.primitives[g.rand.Intn(len(g.primitives))]
}

func (g *gencase) expr() string {
	op := g.operators[g.rand.Intn(len(g.operators))]
	return fmt.Sprintf("%s %s %s", g.arg(false), op, g.arg(op == "IN" || op == "NOT IN"))
}

func (d dummyCollation) ColumnLookup(_ *sqlparser.ColName) (int, error) {
	panic("not supported")
}

func (d dummyCollation) CollationIDLookup(_ sqlparser.Expr) collations.ID {
	return collations.ID(d)
}

func TestTypes(t *testing.T) {
	var conn = mysqlconn(t)
	defer conn.Close()

	var queries = []string{
		"1 > 3",
		"3 > 1",
		"-1 > -1",
		"1 = 1",
		"-1 = 1",
		"1 IN (1, -2, 3)",
		"1 LIKE 1",
		"-1 LIKE -1",
		"-1 LIKE 1",
		`"foo" IN ("bar", "FOO", "baz")`,
		`'pokemon' LIKE 'poke%'`,
		`(1, 2) = (1, 2)`,
		`1 = 'sad'`,
		`(1, 2) = (1, 3)`,
	}

	for _, query := range queries {
		remote, err := conn.ExecuteFetch("SELECT "+query, 1, false)
		if err != nil {
			t.Fatal(err)
		}
		t.Logf("%s => %s", query, remote.Rows[0][0])
	}
}

var fuzzMaxFailures = flag.Int("fuzz-total", 0, "maximum number of failures to fuzz for")
var fuzzSeed = flag.Int64("fuzz-seed", 1234, "RNG seed when generating fuzz expressions")

func TestGenerateFuzzCases(t *testing.T) {
	if *fuzzMaxFailures <= 0 {
		t.Skipf("skipping fuzz test generation")
	}

	type evaltest struct {
		Query string
		Value string `json:",omitempty"`
		Error string `json:",omitempty"`
	}

	var golden []evaltest
	var gen = gencase{
		rand:         rand.New(rand.NewSource(*fuzzSeed)),
		ratioTuple:   8,
		ratioSubexpr: 8,
		tupleLen:     4,
		operators: []string{
			"+", "-", "/", "*", "=", "!=", "<=>", "<", "<=", ">", ">=", "IN", "NOT IN", "LIKE", "NOT LIKE",
		},
		primitives: []string{
			"1", "0", "-1", `"foo"`, `"FOO"`, `"fOo"`, "NULL",
		},
	}

	var conn = mysqlconn(t)
	defer conn.Close()

	for len(golden) < *fuzzMaxFailures {
		query := "SELECT " + gen.expr()

		stmt, err := sqlparser.Parse(query)
		if err != nil {
			t.Fatalf("bad codegen: %v", err)
		}

		var eval evalengine.EvalResult
		var evaluated bool

		astExpr := stmt.(*sqlparser.Select).SelectExprs[0].(*sqlparser.AliasedExpr).Expr
		local, localErr := func() (expr evalengine.Expr, err error) {
			defer func() {
				if r := recover(); r != nil {
					err = fmt.Errorf("PANIC: %v", r)
				}
			}()
			expr, err = evalengine.ConvertEx(astExpr, dummyCollation(45), true)
			return
		}()
		if localErr == nil {
			evaluated = true
			eval, localErr = func() (eval evalengine.EvalResult, err error) {
				defer func() {
					if r := recover(); r != nil {
						err = fmt.Errorf("PANIC: %v", r)
					}
				}()
				eval, err = local.Evaluate(nil)
				return
			}()
		}

		remote, remoteErr := conn.ExecuteFetch(query, 1, false)

		if localErr != nil {
			if remoteErr == nil {
				t.Errorf("local query %q failed: %v (eval=%v); mysql response: %s", query, localErr, evaluated, remote.Rows[0][0].String())
				goto failed
			}
			if !strings.Contains(remoteErr.Error(), localErr.Error()) {
				t.Errorf("mismatch in errors for %q: local=%q (eval=%v), remote=%q", query, localErr.Error(), evaluated, remoteErr.Error())
				goto failed
			}
			continue
		}

		if remoteErr != nil {
			t.Errorf("remote query %q failed: %v, local=%s", query, remoteErr.Error(), eval.Value().String())
			goto failed
		}

		if eval.Value().String() != remote.Rows[0][0].String() {
			t.Errorf("mismatch for query %q: local=%v, remote=%v", query, eval.Value().String(), remote.Rows[0][0].String())
			goto failed
		}

		// OK -- all match
		continue

	failed:
		if remoteErr != nil {
			golden = append(golden, evaltest{
				Query: query,
				Error: remoteErr.Error(),
			})
		} else {
			golden = append(golden, evaltest{
				Query: query,
				Value: remote.Rows[0][0].String(),
			})
		}
	}

	out, err := os.Create(fmt.Sprintf("testdata/mysql_golden_%d.json", time.Now().Unix()))
	if err != nil {
		t.Fatal(err)
	}
	defer out.Close()

	enc := json.NewEncoder(out)
	enc.SetIndent("", "    ")
	enc.Encode(golden)
}
<|endoftext|>