<commit_before>package testrunner

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"syscall"
	"time"

	"github.com/onsi/ginkgo/config"
	"github.com/onsi/ginkgo/ginkgo/testsuite"
	"github.com/onsi/ginkgo/internal/remote"
	"github.com/onsi/ginkgo/reporters/stenographer"
	"github.com/onsi/ginkgo/types"
)

type TestRunner struct {
	Suite testsuite.TestSuite

	compiled              bool
	compilationTargetPath string

	numCPU         int
	parallelStream bool
	race           bool
	cover          bool
	coverPkg       string
	tags           string
	additionalArgs []string
}

func New(suite testsuite.TestSuite, numCPU int, parallelStream bool, race bool, cover bool, coverPkg string, tags string, additionalArgs []string) *TestRunner {
	runner := &TestRunner{
		Suite:          suite,
		numCPU:         numCPU,
		parallelStream: parallelStream,
		race:           race,
		cover:          cover,
		coverPkg:       coverPkg,
		tags:           tags,
		additionalArgs: additionalArgs,
	}

	if !suite.Precompiled {
		dir, err := ioutil.TempDir("", "ginkgo")
		if err != nil {
			panic(fmt.Sprintf("coulnd't create temporary directory... might be time to rm -rf:\n%s", err.Error()))
		}
		runner.compilationTargetPath = filepath.Join(dir, suite.PackageName+".test")
	}

	return runner
}

func (t *TestRunner) Compile() error {
	return t.CompileTo(t.compilationTargetPath)
}

func (t *TestRunner) CompileTo(path string) error {
	if t.compiled {
		return nil
	}

	if t.Suite.Precompiled {
		return nil
	}

	args := []string{"test", "-c", "-i", "-o", path}
	if t.race {
		args = append(args, "-race")
	}
	if t.cover || t.coverPkg != "" {
		args = append(args, "-cover", "-covermode=atomic")
	}
	if t.coverPkg != "" {
		args = append(args, fmt.Sprintf("-coverpkg=%s", t.coverPkg))
	}
	if t.tags != "" {
		args = append(args, fmt.Sprintf("-tags=%s", t.tags))
	}

	cmd := exec.Command("go", args...)

	cmd.Dir = t.Suite.Path

	output, err := cmd.CombinedOutput()

	if err != nil {
		fixedOutput := fixCompilationOutput(string(output), t.Suite.Path)
		if len(output) > 0 {
			return fmt.Errorf("Failed to compile %s:\n\n%s", t.Suite.PackageName, fixedOutput)
		}
		return fmt.Errorf("Failed to compile %s", t.Suite.PackageName)
	}

	if fileExists(path) == false {
		compiledFile := filepath.Join(t.Suite.Path, t.Suite.PackageName+".test")
		if fileExists(compiledFile) {
			// seems like we are on an old go version that does not support the -o flag on go test
			// move the compiled test file to the desired location by hand
			err = os.Rename(compiledFile, path)
			if err != nil {
				// We cannot move the file, perhaps because the source and destination
				// are on different partitions. We can copy the file, however.
				err = copyFile(compiledFile, path)
				if err != nil {
					return fmt.Errorf("Failed to copy compiled file: %s", err)
				}
				err = os.Chmod(path, os.ModePerm)
				if err != nil {
					return fmt.Errorf("Failed to set permissions on compiled file: %s", err)
				}
			}
		} else {
			return fmt.Errorf("Failed to compile %s: output file %q could not be found", t.Suite.PackageName, path)
		}
	}

	t.compiled = true

	return nil
}

func fileExists(path string) bool {
	_, err := os.Stat(path)
	return err == nil || os.IsNotExist(err) == false
}

// copyFile copies the contents of the file named src to the file named
// by dst. The file will be created if it does not already exist. If the
// destination file exists, all it's contents will be replaced by the contents
// of the source file.
func copyFile(src, dst string) error {
	in, err := os.Open(src)
	if err != nil {
		return err
	}

	defer in.Close()

	out, err := os.Create(dst)
	if err != nil {
		return err
	}

	defer func() {
		closeErr := out.Close()
		if err == nil {
			err = closeErr
		}
	}()

	_, err = io.Copy(out, in)
	if err != nil {
		return err
	}

	err = out.Sync()
	return err
}

/*
go test -c -i spits package.test out into the cwd. there's no way to change this.

to make sure it doesn't generate conflicting .test files in the cwd, Compile() must switch the cwd to the test package.

unfortunately, this causes go test's compile output to be expressed *relative to the test package* instead of the cwd.

this makes it hard to reason about what failed, and also prevents iterm's Cmd+click from working.

fixCompilationOutput..... rewrites the output to fix the paths.

yeah......
*/
func fixCompilationOutput(output string, relToPath string) string {
	re := regexp.MustCompile(`^(\S.*\.go)\:\d+\:`)
	lines := strings.Split(output, "\n")
	for i, line := range lines {
		indices := re.FindStringSubmatchIndex(line)
		if len(indices) == 0 {
			continue
		}

		path := line[indices[2]:indices[3]]
		path = filepath.Join(relToPath, path)
		lines[i] = path + line[indices[3]:]
	}
	return strings.Join(lines, "\n")
}

func (t *TestRunner) Run() RunResult {
	if t.Suite.IsGinkgo {
		if t.numCPU > 1 {
			if t.parallelStream {
				return t.runAndStreamParallelGinkgoSuite()
			} else {
				return t.runParallelGinkgoSuite()
			}
		} else {
			return t.runSerialGinkgoSuite()
		}
	} else {
		return t.runGoTestSuite()
	}
}

func (t *TestRunner) CleanUp() {
	if t.Suite.Precompiled {
		return
	}
	os.RemoveAll(filepath.Dir(t.compilationTargetPath))
}

func (t *TestRunner) runSerialGinkgoSuite() RunResult {
	ginkgoArgs := config.BuildFlagArgs("ginkgo", config.GinkgoConfig, config.DefaultReporterConfig)
	return t.run(t.cmd(ginkgoArgs, os.Stdout, 1), nil)
}

func (t *TestRunner) runGoTestSuite() RunResult {
	return t.run(t.cmd([]string{"-test.v"}, os.Stdout, 1), nil)
}

func (t *TestRunner) runAndStreamParallelGinkgoSuite() RunResult {
	completions := make(chan RunResult)
	writers := make([]*logWriter, t.numCPU)

	server, err := remote.NewServer(t.numCPU)
	if err != nil {
		panic("Failed to start parallel spec server")
	}

	server.Start()
	defer server.Close()

	for cpu := 0; cpu < t.numCPU; cpu++ {
		config.GinkgoConfig.ParallelNode = cpu + 1
		config.GinkgoConfig.ParallelTotal = t.numCPU
		config.GinkgoConfig.SyncHost = server.Address()

		ginkgoArgs := config.BuildFlagArgs("ginkgo", config.GinkgoConfig, config.DefaultReporterConfig)

		writers[cpu] = newLogWriter(os.Stdout, cpu+1)

		cmd := t.cmd(ginkgoArgs, writers[cpu], cpu+1)

		server.RegisterAlive(cpu+1, func() bool {
			if cmd.ProcessState == nil {
				return true
			}
			return !cmd.ProcessState.Exited()
		})

		go t.run(cmd, completions)
	}

	res := PassingRunResult()

	for cpu := 0; cpu < t.numCPU; cpu++ {
		res = res.Merge(<-completions)
	}

	for _, writer := range writers {
		writer.Close()
	}

	os.Stdout.Sync()

	if t.cover || t.coverPkg != "" {
		t.combineCoverprofiles()
	}

	return res
}

func (t *TestRunner) runParallelGinkgoSuite() RunResult {
	result := make(chan bool)
	completions := make(chan RunResult)
	writers := make([]*logWriter, t.numCPU)
	reports := make([]*bytes.Buffer, t.numCPU)

	stenographer := stenographer.New(!config.DefaultReporterConfig.NoColor)
	aggregator := remote.NewAggregator(t.numCPU, result, config.DefaultReporterConfig, stenographer)

	server, err := remote.NewServer(t.numCPU)
	if err != nil {
		panic("Failed to start parallel spec server")
	}
	server.RegisterReporters(aggregator)
	server.Start()
	defer server.Close()

	for cpu := 0; cpu < t.numCPU; cpu++ {
		config.GinkgoConfig.ParallelNode = cpu + 1
		config.GinkgoConfig.ParallelTotal = t.numCPU
		config.GinkgoConfig.SyncHost = server.Address()
		config.GinkgoConfig.StreamHost = server.Address()

		ginkgoArgs := config.BuildFlagArgs("ginkgo", config.GinkgoConfig, config.DefaultReporterConfig)

		reports[cpu] = &bytes.Buffer{}
		writers[cpu] = newLogWriter(reports[cpu], cpu+1)

		cmd := t.cmd(ginkgoArgs, writers[cpu], cpu+1)

		server.RegisterAlive(cpu+1, func() bool {
			if cmd.ProcessState == nil {
				return true
			}
			return !cmd.ProcessState.Exited()
		})

		go t.run(cmd, completions)
	}

	res := PassingRunResult()

	for cpu := 0; cpu < t.numCPU; cpu++ {
		res = res.Merge(<-completions)
	}

	//all test processes are done, at this point
	//we should be able to wait for the aggregator to tell us that it's done

	select {
	case <-result:
		fmt.Println("")
	case <-time.After(time.Second):
		//the aggregator never got back to us!  something must have gone wrong
		fmt.Println(`
	 -------------------------------------------------------------------
	|                                                                   |
	|  Ginkgo timed out waiting for all parallel nodes to report back!  |
	|                                                                   |
	 -------------------------------------------------------------------
`)

		os.Stdout.Sync()

		for _, writer := range writers {
			writer.Close()
		}

		for _, report := range reports {
			fmt.Print(report.String())
		}

		os.Stdout.Sync()
	}

	if t.cover || t.coverPkg != "" {
		t.combineCoverprofiles()
	}

	return res
}

func (t *TestRunner) cmd(ginkgoArgs []string, stream io.Writer, node int) *exec.Cmd {
	args := []string{"--test.timeout=24h"}
	if t.cover || t.coverPkg != "" {
		coverprofile := "--test.coverprofile=" + t.Suite.PackageName + ".coverprofile"
		if t.numCPU > 1 {
			coverprofile = fmt.Sprintf("%s.%d", coverprofile, node)
		}
		args = append(args, coverprofile)
	}

	args = append(args, ginkgoArgs...)
	args = append(args, t.additionalArgs...)

	path := t.compilationTargetPath
	if t.Suite.Precompiled {
		path, _ = filepath.Abs(filepath.Join(t.Suite.Path, fmt.Sprintf("%s.test", t.Suite.PackageName)))
	}

	cmd := exec.Command(path, args...)

	cmd.Dir = t.Suite.Path
	cmd.Stderr = stream
	cmd.Stdout = stream

	return cmd
}

func (t *TestRunner) run(cmd *exec.Cmd, completions chan RunResult) RunResult {
	var res RunResult

	defer func() {
		if completions != nil {
			completions <- res
		}
	}()

	err := cmd.Start()
	if err != nil {
		fmt.Printf("Failed to run test suite!\n\t%s", err.Error())
		return res
	}

	cmd.Wait()
	exitStatus := cmd.ProcessState.Sys().(syscall.WaitStatus).ExitStatus()
	res.Passed = (exitStatus == 0) || (exitStatus == types.GINKGO_FOCUS_EXIT_CODE)
	res.HasProgrammaticFocus = (exitStatus == types.GINKGO_FOCUS_EXIT_CODE)

	return res
}

func (t *TestRunner) combineCoverprofiles() {
	profiles := []string{}
	for cpu := 1; cpu <= t.numCPU; cpu++ {
		coverFile := fmt.Sprintf("%s.coverprofile.%d", t.Suite.PackageName, cpu)
		coverFile = filepath.Join(t.Suite.Path, coverFile)
		coverProfile, err := ioutil.ReadFile(coverFile)
		os.Remove(coverFile)

		if err == nil {
			profiles = append(profiles, string(coverProfile))
		}
	}

	if len(profiles) != t.numCPU {
		return
	}

	lines := map[string]int{}
	lineOrder := []string{}
	for i, coverProfile := range profiles {
		for _, line := range strings.Split(string(coverProfile), "\n")[1:] {
			if len(line) == 0 {
				continue
			}
			components := strings.Split(line, " ")
			count, _ := strconv.Atoi(components[len(components)-1])
			prefix := strings.Join(components[0:len(components)-1], " ")
			lines[prefix] += count
			if i == 0 {
				lineOrder = append(lineOrder, prefix)
			}
		}
	}

	output := []string{"mode: atomic"}
	for _, line := range lineOrder {
		output = append(output, fmt.Sprintf("%s %d", line, lines[line]))
	}
	finalOutput := strings.Join(output, "\n")
	ioutil.WriteFile(filepath.Join(t.Suite.Path, fmt.Sprintf("%s.coverprofile", t.Suite.PackageName)), []byte(finalOutput), 0666)
}
<commit_msg>Preserve permissions when copying compiled test file.<commit_after>package testrunner

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"syscall"
	"time"

	"github.com/onsi/ginkgo/config"
	"github.com/onsi/ginkgo/ginkgo/testsuite"
	"github.com/onsi/ginkgo/internal/remote"
	"github.com/onsi/ginkgo/reporters/stenographer"
	"github.com/onsi/ginkgo/types"
)

type TestRunner struct {
	Suite testsuite.TestSuite

	compiled              bool
	compilationTargetPath string

	numCPU         int
	parallelStream bool
	race           bool
	cover          bool
	coverPkg       string
	tags           string
	additionalArgs []string
}

func New(suite testsuite.TestSuite, numCPU int, parallelStream bool, race bool, cover bool, coverPkg string, tags string, additionalArgs []string) *TestRunner {
	runner := &TestRunner{
		Suite:          suite,
		numCPU:         numCPU,
		parallelStream: parallelStream,
		race:           race,
		cover:          cover,
		coverPkg:       coverPkg,
		tags:           tags,
		additionalArgs: additionalArgs,
	}

	if !suite.Precompiled {
		dir, err := ioutil.TempDir("", "ginkgo")
		if err != nil {
			panic(fmt.Sprintf("coulnd't create temporary directory... might be time to rm -rf:\n%s", err.Error()))
		}
		runner.compilationTargetPath = filepath.Join(dir, suite.PackageName+".test")
	}

	return runner
}

func (t *TestRunner) Compile() error {
	return t.CompileTo(t.compilationTargetPath)
}

func (t *TestRunner) CompileTo(path string) error {
	if t.compiled {
		return nil
	}

	if t.Suite.Precompiled {
		return nil
	}

	args := []string{"test", "-c", "-i", "-o", path}
	if t.race {
		args = append(args, "-race")
	}
	if t.cover || t.coverPkg != "" {
		args = append(args, "-cover", "-covermode=atomic")
	}
	if t.coverPkg != "" {
		args = append(args, fmt.Sprintf("-coverpkg=%s", t.coverPkg))
	}
	if t.tags != "" {
		args = append(args, fmt.Sprintf("-tags=%s", t.tags))
	}

	cmd := exec.Command("go", args...)

	cmd.Dir = t.Suite.Path

	output, err := cmd.CombinedOutput()

	if err != nil {
		fixedOutput := fixCompilationOutput(string(output), t.Suite.Path)
		if len(output) > 0 {
			return fmt.Errorf("Failed to compile %s:\n\n%s", t.Suite.PackageName, fixedOutput)
		}
		return fmt.Errorf("Failed to compile %s", t.Suite.PackageName)
	}

	if fileExists(path) == false {
		compiledFile := filepath.Join(t.Suite.Path, t.Suite.PackageName+".test")
		if fileExists(compiledFile) {
			// seems like we are on an old go version that does not support the -o flag on go test
			// move the compiled test file to the desired location by hand
			err = os.Rename(compiledFile, path)
			if err != nil {
				// We cannot move the file, perhaps because the source and destination
				// are on different partitions. We can copy the file, however.
				err = copyFile(compiledFile, path)
				if err != nil {
					return fmt.Errorf("Failed to copy compiled file: %s", err)
				}
			}
		} else {
			return fmt.Errorf("Failed to compile %s: output file %q could not be found", t.Suite.PackageName, path)
		}
	}

	t.compiled = true

	return nil
}

func fileExists(path string) bool {
	_, err := os.Stat(path)
	return err == nil || os.IsNotExist(err) == false
}

// copyFile copies the contents of the file named src to the file named
// by dst. The file will be created if it does not already exist. If the
// destination file exists, all it's contents will be replaced by the contents
// of the source file.
func copyFile(src, dst string) error {
	srcInfo, err := os.Stat(src)
	if err != nil {
		return err
	}
	mode := srcInfo.Mode()

	in, err := os.Open(src)
	if err != nil {
		return err
	}

	defer in.Close()

	out, err := os.Create(dst)
	if err != nil {
		return err
	}

	defer func() {
		closeErr := out.Close()
		if err == nil {
			err = closeErr
		}
	}()

	_, err = io.Copy(out, in)
	if err != nil {
		return err
	}

	err = out.Sync()
	if err != nil {
		return err
	}

	return out.Chmod(mode)
}

/*
go test -c -i spits package.test out into the cwd. there's no way to change this.

to make sure it doesn't generate conflicting .test files in the cwd, Compile() must switch the cwd to the test package.

unfortunately, this causes go test's compile output to be expressed *relative to the test package* instead of the cwd.

this makes it hard to reason about what failed, and also prevents iterm's Cmd+click from working.

fixCompilationOutput..... rewrites the output to fix the paths.

yeah......
*/
func fixCompilationOutput(output string, relToPath string) string {
	re := regexp.MustCompile(`^(\S.*\.go)\:\d+\:`)
	lines := strings.Split(output, "\n")
	for i, line := range lines {
		indices := re.FindStringSubmatchIndex(line)
		if len(indices) == 0 {
			continue
		}

		path := line[indices[2]:indices[3]]
		path = filepath.Join(relToPath, path)
		lines[i] = path + line[indices[3]:]
	}
	return strings.Join(lines, "\n")
}

func (t *TestRunner) Run() RunResult {
	if t.Suite.IsGinkgo {
		if t.numCPU > 1 {
			if t.parallelStream {
				return t.runAndStreamParallelGinkgoSuite()
			} else {
				return t.runParallelGinkgoSuite()
			}
		} else {
			return t.runSerialGinkgoSuite()
		}
	} else {
		return t.runGoTestSuite()
	}
}

func (t *TestRunner) CleanUp() {
	if t.Suite.Precompiled {
		return
	}
	os.RemoveAll(filepath.Dir(t.compilationTargetPath))
}

func (t *TestRunner) runSerialGinkgoSuite() RunResult {
	ginkgoArgs := config.BuildFlagArgs("ginkgo", config.GinkgoConfig, config.DefaultReporterConfig)
	return t.run(t.cmd(ginkgoArgs, os.Stdout, 1), nil)
}

func (t *TestRunner) runGoTestSuite() RunResult {
	return t.run(t.cmd([]string{"-test.v"}, os.Stdout, 1), nil)
}

func (t *TestRunner) runAndStreamParallelGinkgoSuite() RunResult {
	completions := make(chan RunResult)
	writers := make([]*logWriter, t.numCPU)

	server, err := remote.NewServer(t.numCPU)
	if err != nil {
		panic("Failed to start parallel spec server")
	}

	server.Start()
	defer server.Close()

	for cpu := 0; cpu < t.numCPU; cpu++ {
		config.GinkgoConfig.ParallelNode = cpu + 1
		config.GinkgoConfig.ParallelTotal = t.numCPU
		config.GinkgoConfig.SyncHost = server.Address()

		ginkgoArgs := config.BuildFlagArgs("ginkgo", config.GinkgoConfig, config.DefaultReporterConfig)

		writers[cpu] = newLogWriter(os.Stdout, cpu+1)

		cmd := t.cmd(ginkgoArgs, writers[cpu], cpu+1)

		server.RegisterAlive(cpu+1, func() bool {
			if cmd.ProcessState == nil {
				return true
			}
			return !cmd.ProcessState.Exited()
		})

		go t.run(cmd, completions)
	}

	res := PassingRunResult()

	for cpu := 0; cpu < t.numCPU; cpu++ {
		res = res.Merge(<-completions)
	}

	for _, writer := range writers {
		writer.Close()
	}

	os.Stdout.Sync()

	if t.cover || t.coverPkg != "" {
		t.combineCoverprofiles()
	}

	return res
}

func (t *TestRunner) runParallelGinkgoSuite() RunResult {
	result := make(chan bool)
	completions := make(chan RunResult)
	writers := make([]*logWriter, t.numCPU)
	reports := make([]*bytes.Buffer, t.numCPU)

	stenographer := stenographer.New(!config.DefaultReporterConfig.NoColor)
	aggregator := remote.NewAggregator(t.numCPU, result, config.DefaultReporterConfig, stenographer)

	server, err := remote.NewServer(t.numCPU)
	if err != nil {
		panic("Failed to start parallel spec server")
	}
	server.RegisterReporters(aggregator)
	server.Start()
	defer server.Close()

	for cpu := 0; cpu < t.numCPU; cpu++ {
		config.GinkgoConfig.ParallelNode = cpu + 1
		config.GinkgoConfig.ParallelTotal = t.numCPU
		config.GinkgoConfig.SyncHost = server.Address()
		config.GinkgoConfig.StreamHost = server.Address()

		ginkgoArgs := config.BuildFlagArgs("ginkgo", config.GinkgoConfig, config.DefaultReporterConfig)

		reports[cpu] = &bytes.Buffer{}
		writers[cpu] = newLogWriter(reports[cpu], cpu+1)

		cmd := t.cmd(ginkgoArgs, writers[cpu], cpu+1)

		server.RegisterAlive(cpu+1, func() bool {
			if cmd.ProcessState == nil {
				return true
			}
			return !cmd.ProcessState.Exited()
		})

		go t.run(cmd, completions)
	}

	res := PassingRunResult()

	for cpu := 0; cpu < t.numCPU; cpu++ {
		res = res.Merge(<-completions)
	}

	//all test processes are done, at this point
	//we should be able to wait for the aggregator to tell us that it's done

	select {
	case <-result:
		fmt.Println("")
	case <-time.After(time.Second):
		//the aggregator never got back to us!  something must have gone wrong
		fmt.Println(`
	 -------------------------------------------------------------------
	|                                                                   |
	|  Ginkgo timed out waiting for all parallel nodes to report back!  |
	|                                                                   |
	 -------------------------------------------------------------------
`)

		os.Stdout.Sync()

		for _, writer := range writers {
			writer.Close()
		}

		for _, report := range reports {
			fmt.Print(report.String())
		}

		os.Stdout.Sync()
	}

	if t.cover || t.coverPkg != "" {
		t.combineCoverprofiles()
	}

	return res
}

func (t *TestRunner) cmd(ginkgoArgs []string, stream io.Writer, node int) *exec.Cmd {
	args := []string{"--test.timeout=24h"}
	if t.cover || t.coverPkg != "" {
		coverprofile := "--test.coverprofile=" + t.Suite.PackageName + ".coverprofile"
		if t.numCPU > 1 {
			coverprofile = fmt.Sprintf("%s.%d", coverprofile, node)
		}
		args = append(args, coverprofile)
	}

	args = append(args, ginkgoArgs...)
	args = append(args, t.additionalArgs...)

	path := t.compilationTargetPath
	if t.Suite.Precompiled {
		path, _ = filepath.Abs(filepath.Join(t.Suite.Path, fmt.Sprintf("%s.test", t.Suite.PackageName)))
	}

	cmd := exec.Command(path, args...)

	cmd.Dir = t.Suite.Path
	cmd.Stderr = stream
	cmd.Stdout = stream

	return cmd
}

func (t *TestRunner) run(cmd *exec.Cmd, completions chan RunResult) RunResult {
	var res RunResult

	defer func() {
		if completions != nil {
			completions <- res
		}
	}()

	err := cmd.Start()
	if err != nil {
		fmt.Printf("Failed to run test suite!\n\t%s", err.Error())
		return res
	}

	cmd.Wait()
	exitStatus := cmd.ProcessState.Sys().(syscall.WaitStatus).ExitStatus()
	res.Passed = (exitStatus == 0) || (exitStatus == types.GINKGO_FOCUS_EXIT_CODE)
	res.HasProgrammaticFocus = (exitStatus == types.GINKGO_FOCUS_EXIT_CODE)

	return res
}

func (t *TestRunner) combineCoverprofiles() {
	profiles := []string{}
	for cpu := 1; cpu <= t.numCPU; cpu++ {
		coverFile := fmt.Sprintf("%s.coverprofile.%d", t.Suite.PackageName, cpu)
		coverFile = filepath.Join(t.Suite.Path, coverFile)
		coverProfile, err := ioutil.ReadFile(coverFile)
		os.Remove(coverFile)

		if err == nil {
			profiles = append(profiles, string(coverProfile))
		}
	}

	if len(profiles) != t.numCPU {
		return
	}

	lines := map[string]int{}
	lineOrder := []string{}
	for i, coverProfile := range profiles {
		for _, line := range strings.Split(string(coverProfile), "\n")[1:] {
			if len(line) == 0 {
				continue
			}
			components := strings.Split(line, " ")
			count, _ := strconv.Atoi(components[len(components)-1])
			prefix := strings.Join(components[0:len(components)-1], " ")
			lines[prefix] += count
			if i == 0 {
				lineOrder = append(lineOrder, prefix)
			}
		}
	}

	output := []string{"mode: atomic"}
	for _, line := range lineOrder {
		output = append(output, fmt.Sprintf("%s %d", line, lines[line]))
	}
	finalOutput := strings.Join(output, "\n")
	ioutil.WriteFile(filepath.Join(t.Suite.Path, fmt.Sprintf("%s.coverprofile", t.Suite.PackageName)), []byte(finalOutput), 0666)
}
<|endoftext|>


<commit_before>// Copyright 2014 The go-github AUTHORS. All rights reserved.
//
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package github

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"reflect"
	"testing"
)

func TestIssuesService_ListMilestones(t *testing.T) {
	client, mux, _, teardown := setup()
	defer teardown()

	mux.HandleFunc("/repos/o/r/milestones", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, "GET")
		testFormValues(t, r, values{
			"state":     "closed",
			"sort":      "due_date",
			"direction": "asc",
			"page":      "2",
		})
		fmt.Fprint(w, `[{"number":1}]`)
	})

	opt := &MilestoneListOptions{"closed", "due_date", "asc", ListOptions{Page: 2}}
	ctx := context.Background()
	milestones, _, err := client.Issues.ListMilestones(ctx, "o", "r", opt)
	if err != nil {
		t.Errorf("IssuesService.ListMilestones returned error: %v", err)
	}

	want := []*Milestone{{Number: Int(1)}}
	if !reflect.DeepEqual(milestones, want) {
		t.Errorf("IssuesService.ListMilestones returned %+v, want %+v", milestones, want)
	}
}

func TestIssuesService_ListMilestones_invalidOwner(t *testing.T) {
	client, _, _, teardown := setup()
	defer teardown()

	ctx := context.Background()
	_, _, err := client.Issues.ListMilestones(ctx, "%", "r", nil)
	testURLParseError(t, err)
}

func TestIssuesService_GetMilestone(t *testing.T) {
	client, mux, _, teardown := setup()
	defer teardown()

	mux.HandleFunc("/repos/o/r/milestones/1", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, "GET")
		fmt.Fprint(w, `{"number":1}`)
	})

	ctx := context.Background()
	milestone, _, err := client.Issues.GetMilestone(ctx, "o", "r", 1)
	if err != nil {
		t.Errorf("IssuesService.GetMilestone returned error: %v", err)
	}

	want := &Milestone{Number: Int(1)}
	if !reflect.DeepEqual(milestone, want) {
		t.Errorf("IssuesService.GetMilestone returned %+v, want %+v", milestone, want)
	}
}

func TestIssuesService_GetMilestone_invalidOwner(t *testing.T) {
	client, _, _, teardown := setup()
	defer teardown()

	ctx := context.Background()
	_, _, err := client.Issues.GetMilestone(ctx, "%", "r", 1)
	testURLParseError(t, err)
}

func TestIssuesService_CreateMilestone(t *testing.T) {
	client, mux, _, teardown := setup()
	defer teardown()

	input := &Milestone{Title: String("t")}

	mux.HandleFunc("/repos/o/r/milestones", func(w http.ResponseWriter, r *http.Request) {
		v := new(Milestone)
		json.NewDecoder(r.Body).Decode(v)

		testMethod(t, r, "POST")
		if !reflect.DeepEqual(v, input) {
			t.Errorf("Request body = %+v, want %+v", v, input)
		}

		fmt.Fprint(w, `{"number":1}`)
	})

	ctx := context.Background()
	milestone, _, err := client.Issues.CreateMilestone(ctx, "o", "r", input)
	if err != nil {
		t.Errorf("IssuesService.CreateMilestone returned error: %v", err)
	}

	want := &Milestone{Number: Int(1)}
	if !reflect.DeepEqual(milestone, want) {
		t.Errorf("IssuesService.CreateMilestone returned %+v, want %+v", milestone, want)
	}
}

func TestIssuesService_CreateMilestone_invalidOwner(t *testing.T) {
	client, _, _, teardown := setup()
	defer teardown()

	ctx := context.Background()
	_, _, err := client.Issues.CreateMilestone(ctx, "%", "r", nil)
	testURLParseError(t, err)
}

func TestIssuesService_EditMilestone(t *testing.T) {
	client, mux, _, teardown := setup()
	defer teardown()

	input := &Milestone{Title: String("t")}

	mux.HandleFunc("/repos/o/r/milestones/1", func(w http.ResponseWriter, r *http.Request) {
		v := new(Milestone)
		json.NewDecoder(r.Body).Decode(v)

		testMethod(t, r, "PATCH")
		if !reflect.DeepEqual(v, input) {
			t.Errorf("Request body = %+v, want %+v", v, input)
		}

		fmt.Fprint(w, `{"number":1}`)
	})

	ctx := context.Background()
	milestone, _, err := client.Issues.EditMilestone(ctx, "o", "r", 1, input)
	if err != nil {
		t.Errorf("IssuesService.EditMilestone returned error: %v", err)
	}

	want := &Milestone{Number: Int(1)}
	if !reflect.DeepEqual(milestone, want) {
		t.Errorf("IssuesService.EditMilestone returned %+v, want %+v", milestone, want)
	}
}

func TestIssuesService_EditMilestone_invalidOwner(t *testing.T) {
	client, _, _, teardown := setup()
	defer teardown()

	ctx := context.Background()
	_, _, err := client.Issues.EditMilestone(ctx, "%", "r", 1, nil)
	testURLParseError(t, err)
}

func TestIssuesService_DeleteMilestone(t *testing.T) {
	client, mux, _, teardown := setup()
	defer teardown()

	mux.HandleFunc("/repos/o/r/milestones/1", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, "DELETE")
	})

	ctx := context.Background()
	_, err := client.Issues.DeleteMilestone(ctx, "o", "r", 1)
	if err != nil {
		t.Errorf("IssuesService.DeleteMilestone returned error: %v", err)
	}
}

func TestIssuesService_DeleteMilestone_invalidOwner(t *testing.T) {
	client, _, _, teardown := setup()
	defer teardown()

	ctx := context.Background()
	_, err := client.Issues.DeleteMilestone(ctx, "%", "r", 1)
	testURLParseError(t, err)
}
<commit_msg>Improve issues_milestones.go coverage (#1749)<commit_after>// Copyright 2014 The go-github AUTHORS. All rights reserved.
//
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package github

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"reflect"
	"testing"
)

func TestIssuesService_ListMilestones(t *testing.T) {
	client, mux, _, teardown := setup()
	defer teardown()

	mux.HandleFunc("/repos/o/r/milestones", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, "GET")
		testFormValues(t, r, values{
			"state":     "closed",
			"sort":      "due_date",
			"direction": "asc",
			"page":      "2",
		})
		fmt.Fprint(w, `[{"number":1}]`)
	})

	opt := &MilestoneListOptions{"closed", "due_date", "asc", ListOptions{Page: 2}}
	ctx := context.Background()
	milestones, _, err := client.Issues.ListMilestones(ctx, "o", "r", opt)
	if err != nil {
		t.Errorf("IssuesService.ListMilestones returned error: %v", err)
	}

	want := []*Milestone{{Number: Int(1)}}
	if !reflect.DeepEqual(milestones, want) {
		t.Errorf("IssuesService.ListMilestones returned %+v, want %+v", milestones, want)
	}

	const methodName = "ListMilestones"
	testBadOptions(t, methodName, func() (err error) {
		_, _, err = client.Issues.ListMilestones(ctx, "\n", "\n", opt)
		return err
	})

	testNewRequestAndDoFailure(t, methodName, client, func() (*Response, error) {
		got, resp, err := client.Issues.ListMilestones(ctx, "o", "r", opt)
		if got != nil {
			t.Errorf("testNewRequestAndDoFailure %v = %#v, want nil", methodName, got)
		}
		return resp, err
	})
}

func TestIssuesService_ListMilestones_invalidOwner(t *testing.T) {
	client, _, _, teardown := setup()
	defer teardown()

	ctx := context.Background()
	_, _, err := client.Issues.ListMilestones(ctx, "%", "r", nil)
	testURLParseError(t, err)
}

func TestIssuesService_GetMilestone(t *testing.T) {
	client, mux, _, teardown := setup()
	defer teardown()

	mux.HandleFunc("/repos/o/r/milestones/1", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, "GET")
		fmt.Fprint(w, `{"number":1}`)
	})

	ctx := context.Background()
	milestone, _, err := client.Issues.GetMilestone(ctx, "o", "r", 1)
	if err != nil {
		t.Errorf("IssuesService.GetMilestone returned error: %v", err)
	}

	want := &Milestone{Number: Int(1)}
	if !reflect.DeepEqual(milestone, want) {
		t.Errorf("IssuesService.GetMilestone returned %+v, want %+v", milestone, want)
	}

	const methodName = "GetMilestone"
	testBadOptions(t, methodName, func() (err error) {
		_, _, err = client.Issues.GetMilestone(ctx, "\n", "\n", -1)
		return err
	})

	testNewRequestAndDoFailure(t, methodName, client, func() (*Response, error) {
		got, resp, err := client.Issues.GetMilestone(ctx, "o", "r", 1)
		if got != nil {
			t.Errorf("testNewRequestAndDoFailure %v = %#v, want nil", methodName, got)
		}
		return resp, err
	})
}

func TestIssuesService_GetMilestone_invalidOwner(t *testing.T) {
	client, _, _, teardown := setup()
	defer teardown()

	ctx := context.Background()
	_, _, err := client.Issues.GetMilestone(ctx, "%", "r", 1)
	testURLParseError(t, err)
}

func TestIssuesService_CreateMilestone(t *testing.T) {
	client, mux, _, teardown := setup()
	defer teardown()

	input := &Milestone{Title: String("t")}

	mux.HandleFunc("/repos/o/r/milestones", func(w http.ResponseWriter, r *http.Request) {
		v := new(Milestone)
		json.NewDecoder(r.Body).Decode(v)

		testMethod(t, r, "POST")
		if !reflect.DeepEqual(v, input) {
			t.Errorf("Request body = %+v, want %+v", v, input)
		}

		fmt.Fprint(w, `{"number":1}`)
	})

	ctx := context.Background()
	milestone, _, err := client.Issues.CreateMilestone(ctx, "o", "r", input)
	if err != nil {
		t.Errorf("IssuesService.CreateMilestone returned error: %v", err)
	}

	want := &Milestone{Number: Int(1)}
	if !reflect.DeepEqual(milestone, want) {
		t.Errorf("IssuesService.CreateMilestone returned %+v, want %+v", milestone, want)
	}

	const methodName = "CreateMilestone"
	testBadOptions(t, methodName, func() (err error) {
		_, _, err = client.Issues.CreateMilestone(ctx, "\n", "\n", input)
		return err
	})

	testNewRequestAndDoFailure(t, methodName, client, func() (*Response, error) {
		got, resp, err := client.Issues.CreateMilestone(ctx, "o", "r", input)
		if got != nil {
			t.Errorf("testNewRequestAndDoFailure %v = %#v, want nil", methodName, got)
		}
		return resp, err
	})
}

func TestIssuesService_CreateMilestone_invalidOwner(t *testing.T) {
	client, _, _, teardown := setup()
	defer teardown()

	ctx := context.Background()
	_, _, err := client.Issues.CreateMilestone(ctx, "%", "r", nil)
	testURLParseError(t, err)
}

func TestIssuesService_EditMilestone(t *testing.T) {
	client, mux, _, teardown := setup()
	defer teardown()

	input := &Milestone{Title: String("t")}

	mux.HandleFunc("/repos/o/r/milestones/1", func(w http.ResponseWriter, r *http.Request) {
		v := new(Milestone)
		json.NewDecoder(r.Body).Decode(v)

		testMethod(t, r, "PATCH")
		if !reflect.DeepEqual(v, input) {
			t.Errorf("Request body = %+v, want %+v", v, input)
		}

		fmt.Fprint(w, `{"number":1}`)
	})

	ctx := context.Background()
	milestone, _, err := client.Issues.EditMilestone(ctx, "o", "r", 1, input)
	if err != nil {
		t.Errorf("IssuesService.EditMilestone returned error: %v", err)
	}

	want := &Milestone{Number: Int(1)}
	if !reflect.DeepEqual(milestone, want) {
		t.Errorf("IssuesService.EditMilestone returned %+v, want %+v", milestone, want)
	}

	const methodName = "EditMilestone"
	testBadOptions(t, methodName, func() (err error) {
		_, _, err = client.Issues.EditMilestone(ctx, "\n", "\n", -1, input)
		return err
	})

	testNewRequestAndDoFailure(t, methodName, client, func() (*Response, error) {
		got, resp, err := client.Issues.EditMilestone(ctx, "o", "r", 1, input)
		if got != nil {
			t.Errorf("testNewRequestAndDoFailure %v = %#v, want nil", methodName, got)
		}
		return resp, err
	})
}

func TestIssuesService_EditMilestone_invalidOwner(t *testing.T) {
	client, _, _, teardown := setup()
	defer teardown()

	ctx := context.Background()
	_, _, err := client.Issues.EditMilestone(ctx, "%", "r", 1, nil)
	testURLParseError(t, err)
}

func TestIssuesService_DeleteMilestone(t *testing.T) {
	client, mux, _, teardown := setup()
	defer teardown()

	mux.HandleFunc("/repos/o/r/milestones/1", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, "DELETE")
	})

	ctx := context.Background()
	_, err := client.Issues.DeleteMilestone(ctx, "o", "r", 1)
	if err != nil {
		t.Errorf("IssuesService.DeleteMilestone returned error: %v", err)
	}

	const methodName = "DeleteMilestone"
	testBadOptions(t, methodName, func() (err error) {
		_, err = client.Issues.DeleteMilestone(ctx, "\n", "\n", -1)
		return err
	})

	testNewRequestAndDoFailure(t, methodName, client, func() (*Response, error) {
		return client.Issues.DeleteMilestone(ctx, "o", "r", 1)
	})
}

func TestIssuesService_DeleteMilestone_invalidOwner(t *testing.T) {
	client, _, _, teardown := setup()
	defer teardown()

	ctx := context.Background()
	_, err := client.Issues.DeleteMilestone(ctx, "%", "r", 1)
	testURLParseError(t, err)
}
<|endoftext|>


<commit_before>package rest

import (
	"encoding/json"
	"fmt"
	"math/rand"
	"socialapi/models"
	"socialapi/request"
	"strconv"
	"time"

	"github.com/google/go-querystring/query"
)

func GetHistory(channelId int64, q *request.Query) (*models.HistoryResponse, error) {
	v, err := query.Values(q)
	if err != nil {
		return nil, err
	}

	url := fmt.Sprintf("/channel/%d/history?%s", channelId, v.Encode())
	res, err := sendRequest("GET", url, nil)
	if err != nil {
		return nil, err
	}

	var history models.HistoryResponse
	err = json.Unmarshal(res, &history)
	if err != nil {
		return nil, err
	}

	return &history, nil
}

func CountHistory(channelId int64) (*models.CountResponse, error) {
	c := models.NewChannelMessageList()
	c.ChannelId = channelId
	res, err := marshallAndSendRequest("POST", "/channel/count", c)
	if err != nil {
		return nil, err
	}

	var count models.CountResponse
	err = json.Unmarshal(res, &count)
	if err != nil {
		return nil, err
	}

	return &count, nil
}

func FetchChannels(accountId int64) ([]*models.Channel, error) {
	url := fmt.Sprintf("/account/%d/channels", accountId)
	res, err := sendRequest("GET", url, nil)
	if err != nil {
		return nil, err
	}

	var channels []*models.Channel
	err = json.Unmarshal(res, &channels)
	if err != nil {
		return nil, err
	}

	return channels, nil
}

func DeleteChannel(creatorId, channelId int64) error {
	c := models.NewChannel()
	c.CreatorId = creatorId
	c.Id = channelId

	url := fmt.Sprintf("/channel/%d/delete", channelId)
	_, err := sendModel("POST", url, c)
	if err != nil {
		return err
	}
	return nil
}

func CreateChannel(creatorId int64) (*models.Channel, error) {
	c := models.NewChannel()
	rand.Seed(time.Now().UnixNano())
	groupName := c.GroupName + strconv.Itoa(rand.Intn(100000000))

	return CreateChannelByGroupNameAndType(creatorId, groupName, c.TypeConstant)
}

func CreateChannelByGroupNameAndType(creatorId int64, groupName, typeConstant string) (*models.Channel, error) {
	c := models.NewChannel()
	c.GroupName = groupName
	c.CreatorId = creatorId
	c.TypeConstant = typeConstant
	c.Name = c.Name + strconv.Itoa(rand.Intn(100000000))
	cm, err := sendModel("POST", "/channel", c)
	if err != nil {
		return nil, err
	}
	return cm.(*models.Channel), nil
}

func UpdateChannel(cm *models.Channel) (*models.Channel, error) {
	url := fmt.Sprintf("/channel/%d", cm.Id)
	cmI, err := sendModel("POST", url, cm)
	if err != nil {
		return nil, err
	}

	return cmI.(*models.Channel), nil
}

func GetChannel(id int64) (*models.Channel, error) {

	url := fmt.Sprintf("/channel/%d", id)
	cm := models.NewChannel()
	cmI, err := sendModel("GET", url, cm)
	if err != nil {
		return nil, err
	}
	return cmI.(*models.Channel), nil
}

func CreateGroupActivityChannel(creatorId int64, groupName string) (*models.Channel, error) {
	c := models.NewChannel()
	c.GroupName = groupName
	c.CreatorId = creatorId
	c.TypeConstant = models.Channel_TYPE_GROUP
	c.Name = groupName

	cm, err := sendModel("POST", "/channel", c)
	if err != nil {
		return nil, err
	}

	return cm.(*models.Channel), nil
}
<commit_msg>social: fix wrong endpoint<commit_after>package rest

import (
	"encoding/json"
	"fmt"
	"math/rand"
	"socialapi/models"
	"socialapi/request"
	"strconv"
	"time"

	"github.com/google/go-querystring/query"
)

func GetHistory(channelId int64, q *request.Query) (*models.HistoryResponse, error) {
	v, err := query.Values(q)
	if err != nil {
		return nil, err
	}

	url := fmt.Sprintf("/channel/%d/history?%s", channelId, v.Encode())
	res, err := sendRequest("GET", url, nil)
	if err != nil {
		return nil, err
	}

	var history models.HistoryResponse
	err = json.Unmarshal(res, &history)
	if err != nil {
		return nil, err
	}

	return &history, nil
}

func CountHistory(channelId int64) (*models.CountResponse, error) {
	c := models.NewChannelMessageList()
	c.ChannelId = channelId

	url := fmt.Sprintf("/channel/%d/history/count", channelId)
	res, err := marshallAndSendRequest("GET", url, c)
	if err != nil {
		return nil, err
	}

	var count models.CountResponse
	err = json.Unmarshal(res, &count)
	if err != nil {
		return nil, err
	}

	return &count, nil
}

func FetchChannels(accountId int64) ([]*models.Channel, error) {
	url := fmt.Sprintf("/account/%d/channels", accountId)
	res, err := sendRequest("GET", url, nil)
	if err != nil {
		return nil, err
	}

	var channels []*models.Channel
	err = json.Unmarshal(res, &channels)
	if err != nil {
		return nil, err
	}

	return channels, nil
}

func DeleteChannel(creatorId, channelId int64) error {
	c := models.NewChannel()
	c.CreatorId = creatorId
	c.Id = channelId

	url := fmt.Sprintf("/channel/%d/delete", channelId)
	_, err := sendModel("POST", url, c)
	if err != nil {
		return err
	}
	return nil
}

func CreateChannel(creatorId int64) (*models.Channel, error) {
	c := models.NewChannel()
	rand.Seed(time.Now().UnixNano())
	groupName := c.GroupName + strconv.Itoa(rand.Intn(100000000))

	return CreateChannelByGroupNameAndType(creatorId, groupName, c.TypeConstant)
}

func CreateChannelByGroupNameAndType(creatorId int64, groupName, typeConstant string) (*models.Channel, error) {
	c := models.NewChannel()
	c.GroupName = groupName
	c.CreatorId = creatorId
	c.TypeConstant = typeConstant
	c.Name = c.Name + strconv.Itoa(rand.Intn(100000000))
	cm, err := sendModel("POST", "/channel", c)
	if err != nil {
		return nil, err
	}
	return cm.(*models.Channel), nil
}

func UpdateChannel(cm *models.Channel) (*models.Channel, error) {
	url := fmt.Sprintf("/channel/%d", cm.Id)
	cmI, err := sendModel("POST", url, cm)
	if err != nil {
		return nil, err
	}

	return cmI.(*models.Channel), nil
}

func GetChannel(id int64) (*models.Channel, error) {

	url := fmt.Sprintf("/channel/%d", id)
	cm := models.NewChannel()
	cmI, err := sendModel("GET", url, cm)
	if err != nil {
		return nil, err
	}
	return cmI.(*models.Channel), nil
}

func CreateGroupActivityChannel(creatorId int64, groupName string) (*models.Channel, error) {
	c := models.NewChannel()
	c.GroupName = groupName
	c.CreatorId = creatorId
	c.TypeConstant = models.Channel_TYPE_GROUP
	c.Name = groupName

	cm, err := sendModel("POST", "/channel", c)
	if err != nil {
		return nil, err
	}

	return cm.(*models.Channel), nil
}
<|endoftext|>


<commit_before>/*
Copyright 2019 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 tmutils

import (
	"crypto/md5"
	"encoding/hex"
	"fmt"
	"regexp"
	"strings"

	"github.com/golang/protobuf/proto"
	"vitess.io/vitess/go/vt/concurrency"

	tabletmanagerdatapb "vitess.io/vitess/go/vt/proto/tabletmanagerdata"
)

// This file contains helper methods to deal with Schema information.

const (
	// TableBaseTable indicates the table type is a base table.
	TableBaseTable = "BASE TABLE"
	// TableView indicates the table type is a view.
	TableView = "VIEW"
)

// TableDefinitionGetColumn returns the index of a column inside a
// TableDefinition.
func TableDefinitionGetColumn(td *tabletmanagerdatapb.TableDefinition, name string) (index int, ok bool) {
	lowered := strings.ToLower(name)
	for i, n := range td.Columns {
		if lowered == strings.ToLower(n) {
			return i, true
		}
	}
	return -1, false
}

// TableDefinitions is a list of TableDefinition, for sorting
type TableDefinitions []*tabletmanagerdatapb.TableDefinition

// Len returns TableDefinitions length.
func (tds TableDefinitions) Len() int {
	return len(tds)
}

// Swap used for sorting TableDefinitions.
func (tds TableDefinitions) Swap(i, j int) {
	tds[i], tds[j] = tds[j], tds[i]
}

// TableFilter is a filter for table names and types.
type TableFilter struct {
	includeViews bool

	filterTables bool
	tableNames   []string
	tableREs     []*regexp.Regexp

	filterExcludeTables bool
	excludeTableNames   []string
	excludeTableREs     []*regexp.Regexp
}

// NewTableFilter creates a TableFilter for whitelisted tables
// (tables), no blacklisted tables (excludeTables) and optionally
// views (includeViews).
func NewTableFilter(tables, excludeTables []string, includeViews bool) (*TableFilter, error) {
	f := &TableFilter{
		includeViews: includeViews,
	}

	// Build a list of regexp to match table names against.
	// We only use regexps if the name starts and ends with '/'.
	// Otherwise the entry in the arrays is nil, and we use the original
	// table name.
	if len(tables) > 0 {
		f.filterTables = true
		for _, table := range tables {
			if strings.HasPrefix(table, "/") && strings.HasSuffix(table, "/") {
				table = strings.Trim(table, "/")
				re, err := regexp.Compile(table)
				if err != nil {
					return nil, fmt.Errorf("cannot compile regexp %v for table: %v", table, err)
				}

				f.tableREs = append(f.tableREs, re)
			} else {
				f.tableNames = append(f.tableNames, table)
			}
		}
	}

	if len(excludeTables) > 0 {
		f.filterExcludeTables = true
		for _, table := range excludeTables {
			if strings.HasPrefix(table, "/") && strings.HasSuffix(table, "/") {
				table = strings.Trim(table, "/")
				re, err := regexp.Compile(table)
				if err != nil {
					return nil, fmt.Errorf("cannot compile regexp %v for excludeTable: %v", table, err)
				}

				f.excludeTableREs = append(f.tableREs, re)
			} else {
				f.excludeTableNames = append(f.excludeTableNames, table)
			}
		}
	}

	return f, nil
}

// Includes returns whether a tableName/tableType should be included in this TableFilter.
func (f *TableFilter) Includes(tableName string, tableType string) bool {
	if f.filterTables {
		matches := false
		for _, name := range f.tableNames {
			if strings.EqualFold(name, tableName) {
				matches = true
				break
			}
		}

		if !matches {
			for _, re := range f.tableREs {
				if re.MatchString(tableName) {
					matches = true
					break
				}
			}
		}

		if !matches {
			return false
		}
	}

	if f.filterExcludeTables {
		for _, name := range f.excludeTableNames {
			if strings.EqualFold(name, tableName) {
				return false
			}
		}

		for _, re := range f.excludeTableREs {
			if re.MatchString(tableName) {
				return false
			}
		}
	}

	if !f.includeViews && tableType == TableView {
		return false
	}

	return true
}

// FilterTables returns a copy which includes only whitelisted tables
// (tables), no blacklisted tables (excludeTables) and optionally
// views (includeViews).
func FilterTables(sd *tabletmanagerdatapb.SchemaDefinition, tables, excludeTables []string, includeViews bool) (*tabletmanagerdatapb.SchemaDefinition, error) {
	copy := *sd
	copy.TableDefinitions = make([]*tabletmanagerdatapb.TableDefinition, 0, len(sd.TableDefinitions))

	f, err := NewTableFilter(tables, excludeTables, includeViews)
	if err != nil {
		return nil, err
	}

	for _, table := range sd.TableDefinitions {
		if f.Includes(table.Name, table.Type) {
			copy.TableDefinitions = append(copy.TableDefinitions, table)
		}
	}

	// Regenerate hash over tables because it may have changed.
	if copy.Version != "" {
		GenerateSchemaVersion(&copy)
	}

	return &copy, nil
}

// GenerateSchemaVersion return a unique schema version string based on
// its TableDefinitions.
func GenerateSchemaVersion(sd *tabletmanagerdatapb.SchemaDefinition) {
	hasher := md5.New()
	for _, td := range sd.TableDefinitions {
		if _, err := hasher.Write([]byte(td.Schema)); err != nil {
			panic(err) // extremely unlikely
		}
	}
	sd.Version = hex.EncodeToString(hasher.Sum(nil))
}

// SchemaDefinitionGetTable returns TableDefinition for a given table name.
func SchemaDefinitionGetTable(sd *tabletmanagerdatapb.SchemaDefinition, table string) (td *tabletmanagerdatapb.TableDefinition, ok bool) {
	for _, td := range sd.TableDefinitions {
		if td.Name == table {
			return td, true
		}
	}
	return nil, false
}

// SchemaDefinitionToSQLStrings converts a SchemaDefinition to an array of SQL strings. The array contains all
// the SQL statements needed for creating the database, tables, and views - in that order.
// All SQL statements will have {{.DatabaseName}} in place of the actual db name.
func SchemaDefinitionToSQLStrings(sd *tabletmanagerdatapb.SchemaDefinition) []string {
	sqlStrings := make([]string, 0, len(sd.TableDefinitions)+1)
	createViewSQL := make([]string, 0, len(sd.TableDefinitions))

	sqlStrings = append(sqlStrings, sd.DatabaseSchema)

	for _, td := range sd.TableDefinitions {
		if td.Type == TableView {
			createViewSQL = append(createViewSQL, td.Schema)
		} else {
			lines := strings.Split(td.Schema, "\n")
			for i, line := range lines {
				if strings.HasPrefix(line, "CREATE TABLE `") {
					lines[i] = strings.Replace(line, "CREATE TABLE `", "CREATE TABLE `{{.DatabaseName}}`.`", 1)
				}
			}
			sqlStrings = append(sqlStrings, strings.Join(lines, "\n"))
		}
	}

	return append(sqlStrings, createViewSQL...)
}

// DiffSchema generates a report on what's different between two SchemaDefinitions
// including views.
func DiffSchema(leftName string, left *tabletmanagerdatapb.SchemaDefinition, rightName string, right *tabletmanagerdatapb.SchemaDefinition, er concurrency.ErrorRecorder) {
	if left == nil && right == nil {
		return
	}
	if left == nil || right == nil {
		er.RecordError(fmt.Errorf("schemas are different:\n%s: %v, %s: %v", leftName, left, rightName, right))
		return
	}
	if left.DatabaseSchema != right.DatabaseSchema {
		er.RecordError(fmt.Errorf("schemas are different:\n%s: %v\n differs from:\n%s: %v", leftName, left.DatabaseSchema, rightName, right.DatabaseSchema))
	}

	leftIndex := 0
	rightIndex := 0
	for leftIndex < len(left.TableDefinitions) && rightIndex < len(right.TableDefinitions) {
		// extra table on the left side
		if left.TableDefinitions[leftIndex].Name < right.TableDefinitions[rightIndex].Name {
			er.RecordError(fmt.Errorf("%v has an extra table named %v", leftName, left.TableDefinitions[leftIndex].Name))
			leftIndex++
			continue
		}

		// extra table on the right side
		if left.TableDefinitions[leftIndex].Name > right.TableDefinitions[rightIndex].Name {
			er.RecordError(fmt.Errorf("%v has an extra table named %v", rightName, right.TableDefinitions[rightIndex].Name))
			rightIndex++
			continue
		}

		// same name, let's see content
		if left.TableDefinitions[leftIndex].Schema != right.TableDefinitions[rightIndex].Schema {
			er.RecordError(fmt.Errorf("schemas differ on table %v:\n%s: %v\n differs from:\n%s: %v", left.TableDefinitions[leftIndex].Name, leftName, left.TableDefinitions[leftIndex].Schema, rightName, right.TableDefinitions[rightIndex].Schema))
		}

		if left.TableDefinitions[leftIndex].Type != right.TableDefinitions[rightIndex].Type {
			er.RecordError(fmt.Errorf("schemas differ on table type for table %v:\n%s: %v\n differs from:\n%s: %v", left.TableDefinitions[leftIndex].Name, leftName, left.TableDefinitions[leftIndex].Type, rightName, right.TableDefinitions[rightIndex].Type))
		}

		leftIndex++
		rightIndex++
	}

	for leftIndex < len(left.TableDefinitions) {
		if left.TableDefinitions[leftIndex].Type == TableBaseTable {
			er.RecordError(fmt.Errorf("%v has an extra table named %v", leftName, left.TableDefinitions[leftIndex].Name))
		}
		if left.TableDefinitions[leftIndex].Type == TableView {
			er.RecordError(fmt.Errorf("%v has an extra view named %v", leftName, left.TableDefinitions[leftIndex].Name))
		}
		leftIndex++
	}
	for rightIndex < len(right.TableDefinitions) {
		if right.TableDefinitions[rightIndex].Type == TableBaseTable {
			er.RecordError(fmt.Errorf("%v has an extra table named %v", rightName, right.TableDefinitions[rightIndex].Name))
		}
		if right.TableDefinitions[rightIndex].Type == TableView {
			er.RecordError(fmt.Errorf("%v has an extra view named %v", rightName, right.TableDefinitions[rightIndex].Name))
		}
		rightIndex++
	}
}

// DiffSchemaToArray diffs two schemas and return the schema diffs if there is any.
func DiffSchemaToArray(leftName string, left *tabletmanagerdatapb.SchemaDefinition, rightName string, right *tabletmanagerdatapb.SchemaDefinition) (result []string) {
	er := concurrency.AllErrorRecorder{}
	DiffSchema(leftName, left, rightName, right, &er)
	if er.HasErrors() {
		return er.ErrorStrings()
	}
	return nil
}

// SchemaChange contains all necessary information to apply a schema change.
// It should not be sent over the wire, it's just a set of parameters.
type SchemaChange struct {
	SQL              string
	Force            bool
	AllowReplication bool
	BeforeSchema     *tabletmanagerdatapb.SchemaDefinition
	AfterSchema      *tabletmanagerdatapb.SchemaDefinition
}

// Equal compares two SchemaChange objects.
func (s *SchemaChange) Equal(s2 *SchemaChange) bool {
	return s.SQL == s2.SQL &&
		s.Force == s2.Force &&
		s.AllowReplication == s2.AllowReplication &&
		proto.Equal(s.BeforeSchema, s2.BeforeSchema) &&
		proto.Equal(s.AfterSchema, s2.AfterSchema)
}
<commit_msg>Reverting breaking change in #6242: local example was failing<commit_after>/*
Copyright 2019 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 tmutils

import (
	"crypto/md5"
	"encoding/hex"
	"fmt"
	"regexp"
	"strings"

	"github.com/golang/protobuf/proto"
	"vitess.io/vitess/go/vt/concurrency"

	tabletmanagerdatapb "vitess.io/vitess/go/vt/proto/tabletmanagerdata"
)

// This file contains helper methods to deal with Schema information.

const (
	// TableBaseTable indicates the table type is a base table.
	TableBaseTable = "BASE TABLE"
	// TableView indicates the table type is a view.
	TableView = "VIEW"
)

// TableDefinitionGetColumn returns the index of a column inside a
// TableDefinition.
func TableDefinitionGetColumn(td *tabletmanagerdatapb.TableDefinition, name string) (index int, ok bool) {
	lowered := strings.ToLower(name)
	for i, n := range td.Columns {
		if lowered == strings.ToLower(n) {
			return i, true
		}
	}
	return -1, false
}

// TableDefinitions is a list of TableDefinition, for sorting
type TableDefinitions []*tabletmanagerdatapb.TableDefinition

// Len returns TableDefinitions length.
func (tds TableDefinitions) Len() int {
	return len(tds)
}

// Swap used for sorting TableDefinitions.
func (tds TableDefinitions) Swap(i, j int) {
	tds[i], tds[j] = tds[j], tds[i]
}

// TableFilter is a filter for table names and types.
type TableFilter struct {
	includeViews bool

	filterTables bool
	tableNames   []string
	tableREs     []*regexp.Regexp

	filterExcludeTables bool
	excludeTableNames   []string
	excludeTableREs     []*regexp.Regexp
}

// NewTableFilter creates a TableFilter for whitelisted tables
// (tables), no blacklisted tables (excludeTables) and optionally
// views (includeViews).
func NewTableFilter(tables, excludeTables []string, includeViews bool) (*TableFilter, error) {
	f := &TableFilter{
		includeViews: includeViews,
	}

	// Build a list of regexp to match table names against.
	// We only use regexps if the name starts and ends with '/'.
	// Otherwise the entry in the arrays is nil, and we use the original
	// table name.
	if len(tables) > 0 {
		f.filterTables = true
		for _, table := range tables {
			if strings.HasPrefix(table, "/") {
				table = strings.Trim(table, "/")
				re, err := regexp.Compile(table)
				if err != nil {
					return nil, fmt.Errorf("cannot compile regexp %v for table: %v", table, err)
				}

				f.tableREs = append(f.tableREs, re)
			} else {
				f.tableNames = append(f.tableNames, table)
			}
		}
	}

	if len(excludeTables) > 0 {
		f.filterExcludeTables = true
		for _, table := range excludeTables {
			if strings.HasPrefix(table, "/") {
				table = strings.Trim(table, "/")
				re, err := regexp.Compile(table)
				if err != nil {
					return nil, fmt.Errorf("cannot compile regexp %v for excludeTable: %v", table, err)
				}

				f.excludeTableREs = append(f.tableREs, re)
			} else {
				f.excludeTableNames = append(f.excludeTableNames, table)
			}
		}
	}

	return f, nil
}

// Includes returns whether a tableName/tableType should be included in this TableFilter.
func (f *TableFilter) Includes(tableName string, tableType string) bool {
	if f.filterTables {
		matches := false
		for _, name := range f.tableNames {
			if strings.EqualFold(name, tableName) {
				matches = true
				break
			}
		}

		if !matches {
			for _, re := range f.tableREs {
				if re.MatchString(tableName) {
					matches = true
					break
				}
			}
		}

		if !matches {
			return false
		}
	}

	if f.filterExcludeTables {
		for _, name := range f.excludeTableNames {
			if strings.EqualFold(name, tableName) {
				return false
			}
		}

		for _, re := range f.excludeTableREs {
			if re.MatchString(tableName) {
				return false
			}
		}
	}

	if !f.includeViews && tableType == TableView {
		return false
	}

	return true
}

// FilterTables returns a copy which includes only whitelisted tables
// (tables), no blacklisted tables (excludeTables) and optionally
// views (includeViews).
func FilterTables(sd *tabletmanagerdatapb.SchemaDefinition, tables, excludeTables []string, includeViews bool) (*tabletmanagerdatapb.SchemaDefinition, error) {
	copy := *sd
	copy.TableDefinitions = make([]*tabletmanagerdatapb.TableDefinition, 0, len(sd.TableDefinitions))

	f, err := NewTableFilter(tables, excludeTables, includeViews)
	if err != nil {
		return nil, err
	}

	for _, table := range sd.TableDefinitions {
		if f.Includes(table.Name, table.Type) {
			copy.TableDefinitions = append(copy.TableDefinitions, table)
		}
	}

	// Regenerate hash over tables because it may have changed.
	if copy.Version != "" {
		GenerateSchemaVersion(&copy)
	}

	return &copy, nil
}

// GenerateSchemaVersion return a unique schema version string based on
// its TableDefinitions.
func GenerateSchemaVersion(sd *tabletmanagerdatapb.SchemaDefinition) {
	hasher := md5.New()
	for _, td := range sd.TableDefinitions {
		if _, err := hasher.Write([]byte(td.Schema)); err != nil {
			panic(err) // extremely unlikely
		}
	}
	sd.Version = hex.EncodeToString(hasher.Sum(nil))
}

// SchemaDefinitionGetTable returns TableDefinition for a given table name.
func SchemaDefinitionGetTable(sd *tabletmanagerdatapb.SchemaDefinition, table string) (td *tabletmanagerdatapb.TableDefinition, ok bool) {
	for _, td := range sd.TableDefinitions {
		if td.Name == table {
			return td, true
		}
	}
	return nil, false
}

// SchemaDefinitionToSQLStrings converts a SchemaDefinition to an array of SQL strings. The array contains all
// the SQL statements needed for creating the database, tables, and views - in that order.
// All SQL statements will have {{.DatabaseName}} in place of the actual db name.
func SchemaDefinitionToSQLStrings(sd *tabletmanagerdatapb.SchemaDefinition) []string {
	sqlStrings := make([]string, 0, len(sd.TableDefinitions)+1)
	createViewSQL := make([]string, 0, len(sd.TableDefinitions))

	sqlStrings = append(sqlStrings, sd.DatabaseSchema)

	for _, td := range sd.TableDefinitions {
		if td.Type == TableView {
			createViewSQL = append(createViewSQL, td.Schema)
		} else {
			lines := strings.Split(td.Schema, "\n")
			for i, line := range lines {
				if strings.HasPrefix(line, "CREATE TABLE `") {
					lines[i] = strings.Replace(line, "CREATE TABLE `", "CREATE TABLE `{{.DatabaseName}}`.`", 1)
				}
			}
			sqlStrings = append(sqlStrings, strings.Join(lines, "\n"))
		}
	}

	return append(sqlStrings, createViewSQL...)
}

// DiffSchema generates a report on what's different between two SchemaDefinitions
// including views.
func DiffSchema(leftName string, left *tabletmanagerdatapb.SchemaDefinition, rightName string, right *tabletmanagerdatapb.SchemaDefinition, er concurrency.ErrorRecorder) {
	if left == nil && right == nil {
		return
	}
	if left == nil || right == nil {
		er.RecordError(fmt.Errorf("schemas are different:\n%s: %v, %s: %v", leftName, left, rightName, right))
		return
	}
	if left.DatabaseSchema != right.DatabaseSchema {
		er.RecordError(fmt.Errorf("schemas are different:\n%s: %v\n differs from:\n%s: %v", leftName, left.DatabaseSchema, rightName, right.DatabaseSchema))
	}

	leftIndex := 0
	rightIndex := 0
	for leftIndex < len(left.TableDefinitions) && rightIndex < len(right.TableDefinitions) {
		// extra table on the left side
		if left.TableDefinitions[leftIndex].Name < right.TableDefinitions[rightIndex].Name {
			er.RecordError(fmt.Errorf("%v has an extra table named %v", leftName, left.TableDefinitions[leftIndex].Name))
			leftIndex++
			continue
		}

		// extra table on the right side
		if left.TableDefinitions[leftIndex].Name > right.TableDefinitions[rightIndex].Name {
			er.RecordError(fmt.Errorf("%v has an extra table named %v", rightName, right.TableDefinitions[rightIndex].Name))
			rightIndex++
			continue
		}

		// same name, let's see content
		if left.TableDefinitions[leftIndex].Schema != right.TableDefinitions[rightIndex].Schema {
			er.RecordError(fmt.Errorf("schemas differ on table %v:\n%s: %v\n differs from:\n%s: %v", left.TableDefinitions[leftIndex].Name, leftName, left.TableDefinitions[leftIndex].Schema, rightName, right.TableDefinitions[rightIndex].Schema))
		}

		if left.TableDefinitions[leftIndex].Type != right.TableDefinitions[rightIndex].Type {
			er.RecordError(fmt.Errorf("schemas differ on table type for table %v:\n%s: %v\n differs from:\n%s: %v", left.TableDefinitions[leftIndex].Name, leftName, left.TableDefinitions[leftIndex].Type, rightName, right.TableDefinitions[rightIndex].Type))
		}

		leftIndex++
		rightIndex++
	}

	for leftIndex < len(left.TableDefinitions) {
		if left.TableDefinitions[leftIndex].Type == TableBaseTable {
			er.RecordError(fmt.Errorf("%v has an extra table named %v", leftName, left.TableDefinitions[leftIndex].Name))
		}
		if left.TableDefinitions[leftIndex].Type == TableView {
			er.RecordError(fmt.Errorf("%v has an extra view named %v", leftName, left.TableDefinitions[leftIndex].Name))
		}
		leftIndex++
	}
	for rightIndex < len(right.TableDefinitions) {
		if right.TableDefinitions[rightIndex].Type == TableBaseTable {
			er.RecordError(fmt.Errorf("%v has an extra table named %v", rightName, right.TableDefinitions[rightIndex].Name))
		}
		if right.TableDefinitions[rightIndex].Type == TableView {
			er.RecordError(fmt.Errorf("%v has an extra view named %v", rightName, right.TableDefinitions[rightIndex].Name))
		}
		rightIndex++
	}
}

// DiffSchemaToArray diffs two schemas and return the schema diffs if there is any.
func DiffSchemaToArray(leftName string, left *tabletmanagerdatapb.SchemaDefinition, rightName string, right *tabletmanagerdatapb.SchemaDefinition) (result []string) {
	er := concurrency.AllErrorRecorder{}
	DiffSchema(leftName, left, rightName, right, &er)
	if er.HasErrors() {
		return er.ErrorStrings()
	}
	return nil
}

// SchemaChange contains all necessary information to apply a schema change.
// It should not be sent over the wire, it's just a set of parameters.
type SchemaChange struct {
	SQL              string
	Force            bool
	AllowReplication bool
	BeforeSchema     *tabletmanagerdatapb.SchemaDefinition
	AfterSchema      *tabletmanagerdatapb.SchemaDefinition
}

// Equal compares two SchemaChange objects.
func (s *SchemaChange) Equal(s2 *SchemaChange) bool {
	return s.SQL == s2.SQL &&
		s.Force == s2.Force &&
		s.AllowReplication == s2.AllowReplication &&
		proto.Equal(s.BeforeSchema, s2.BeforeSchema) &&
		proto.Equal(s.AfterSchema, s2.AfterSchema)
}
<|endoftext|>


<commit_before>// Copyright 2012, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package tabletserver

import (
	"fmt"
	"strconv"
	"strings"
	"sync"

	log "github.com/golang/glog"
	querypb "github.com/youtube/vitess/go/vt/proto/query"
	"github.com/youtube/vitess/go/vt/schema"
	"github.com/youtube/vitess/go/vt/sqlparser"
	"golang.org/x/net/context"
)

// TableInfo contains the tabletserver related info for a table.
// It's a superset of schema.Table.
type TableInfo struct {
	*schema.Table

	// Seq must be locked before accessing the sequence vars.
	// If CurVal==LastVal, we have to cache new values.
	Seq     sync.Mutex
	NextVal int64
	LastVal int64
}

// NewTableInfo creates a new TableInfo.
func NewTableInfo(conn *DBConn, tableName string, tableType string, comment string) (ti *TableInfo, err error) {
	ti, err = loadTableInfo(conn, tableName)
	if err != nil {
		return nil, err
	}
	if strings.Contains(comment, "vitess_sequence") {
		ti.Type = schema.Sequence
	}
	return ti, nil
}

func loadTableInfo(conn *DBConn, tableName string) (ti *TableInfo, err error) {
	ti = &TableInfo{Table: schema.NewTable(tableName)}
	sqlTableName := sqlparser.String(ti.Name)
	if err = ti.fetchColumns(conn, sqlTableName); err != nil {
		return nil, err
	}
	if err = ti.fetchIndexes(conn, sqlTableName); err != nil {
		return nil, err
	}
	return ti, nil
}

func (ti *TableInfo) fetchColumns(conn *DBConn, sqlTableName string) error {
	qr, err := conn.Exec(context.Background(), fmt.Sprintf("select * from %s where 1 != 1", sqlTableName), 10000, true)
	if err != nil {
		return err
	}
	fieldTypes := make(map[string]querypb.Type, len(qr.Fields))
	for _, field := range qr.Fields {
		fieldTypes[field.Name] = field.Type
	}
	columns, err := conn.Exec(context.Background(), fmt.Sprintf("describe %s", sqlTableName), 10000, false)
	if err != nil {
		return err
	}
	for _, row := range columns.Rows {
		name := row[0].String()
		columnType, ok := fieldTypes[name]
		if !ok {
			log.Warningf("Table: %s, column %s not found in select list, skipping.", ti.Name, name)
			continue
		}
		ti.AddColumn(name, columnType, row[4], row[5].String())
	}
	return nil
}

// SetPK sets the pk columns for a TableInfo.
func (ti *TableInfo) SetPK(colnames []string) error {
	pkIndex := schema.NewIndex("PRIMARY")
	colnums := make([]int, len(colnames))
	for i, colname := range colnames {
		colnums[i] = ti.FindColumn(sqlparser.NewColIdent(colname))
		if colnums[i] == -1 {
			return fmt.Errorf("column %s not found", colname)
		}
		pkIndex.AddColumn(colname, 1)
	}
	for _, col := range ti.Columns {
		pkIndex.DataColumns = append(pkIndex.DataColumns, col.Name)
	}
	ti.Indexes = append(ti.Indexes, nil)
	copy(ti.Indexes[1:], ti.Indexes[:len(ti.Indexes)-1])
	ti.Indexes[0] = pkIndex
	ti.PKColumns = colnums
	return nil
}

func (ti *TableInfo) fetchIndexes(conn *DBConn, sqlTableName string) error {
	indexes, err := conn.Exec(context.Background(), fmt.Sprintf("show index from %s", sqlTableName), 10000, false)
	if err != nil {
		return err
	}
	var currentIndex *schema.Index
	currentName := ""
	for _, row := range indexes.Rows {
		indexName := row[2].String()
		if currentName != indexName {
			currentIndex = ti.AddIndex(indexName)
			currentName = indexName
		}
		var cardinality uint64
		if !row[6].IsNull() {
			cardinality, err = strconv.ParseUint(row[6].String(), 0, 64)
			if err != nil {
				log.Warningf("%s", err)
			}
		}
		currentIndex.AddColumn(row[4].String(), cardinality)
	}
	if !ti.HasPrimary() {
		return nil
	}
	pkIndex := ti.Indexes[0]
	ti.PKColumns = make([]int, len(pkIndex.Columns))
	for i, pkCol := range pkIndex.Columns {
		ti.PKColumns[i] = ti.FindColumn(pkCol)
	}
	// Primary key contains all table columns
	for _, col := range ti.Columns {
		pkIndex.DataColumns = append(pkIndex.DataColumns, col.Name)
	}
	// Secondary indices contain all primary key columns
	for i := 1; i < len(ti.Indexes); i++ {
		for _, c := range ti.Indexes[i].Columns {
			ti.Indexes[i].DataColumns = append(ti.Indexes[i].DataColumns, c)
		}
		for _, c := range pkIndex.Columns {
			// pk columns may already be part of the index. So,
			// check before adding.
			if ti.Indexes[i].FindDataColumn(c) != -1 {
				continue
			}
			ti.Indexes[i].DataColumns = append(ti.Indexes[i].DataColumns, c)
		}
	}
	return nil
}
<commit_msg>Making maxrows consistent with what we expect.<commit_after>// Copyright 2012, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package tabletserver

import (
	"fmt"
	"strconv"
	"strings"
	"sync"

	log "github.com/golang/glog"
	querypb "github.com/youtube/vitess/go/vt/proto/query"
	"github.com/youtube/vitess/go/vt/schema"
	"github.com/youtube/vitess/go/vt/sqlparser"
	"golang.org/x/net/context"
)

// TableInfo contains the tabletserver related info for a table.
// It's a superset of schema.Table.
type TableInfo struct {
	*schema.Table

	// Seq must be locked before accessing the sequence vars.
	// If CurVal==LastVal, we have to cache new values.
	Seq     sync.Mutex
	NextVal int64
	LastVal int64
}

// NewTableInfo creates a new TableInfo.
func NewTableInfo(conn *DBConn, tableName string, tableType string, comment string) (ti *TableInfo, err error) {
	ti, err = loadTableInfo(conn, tableName)
	if err != nil {
		return nil, err
	}
	if strings.Contains(comment, "vitess_sequence") {
		ti.Type = schema.Sequence
	}
	return ti, nil
}

func loadTableInfo(conn *DBConn, tableName string) (ti *TableInfo, err error) {
	ti = &TableInfo{Table: schema.NewTable(tableName)}
	sqlTableName := sqlparser.String(ti.Name)
	if err = ti.fetchColumns(conn, sqlTableName); err != nil {
		return nil, err
	}
	if err = ti.fetchIndexes(conn, sqlTableName); err != nil {
		return nil, err
	}
	return ti, nil
}

func (ti *TableInfo) fetchColumns(conn *DBConn, sqlTableName string) error {
	qr, err := conn.Exec(context.Background(), fmt.Sprintf("select * from %s where 1 != 1", sqlTableName), 0, true)
	if err != nil {
		return err
	}
	fieldTypes := make(map[string]querypb.Type, len(qr.Fields))
	for _, field := range qr.Fields {
		fieldTypes[field.Name] = field.Type
	}
	columns, err := conn.Exec(context.Background(), fmt.Sprintf("describe %s", sqlTableName), 10000, false)
	if err != nil {
		return err
	}
	for _, row := range columns.Rows {
		name := row[0].String()
		columnType, ok := fieldTypes[name]
		if !ok {
			log.Warningf("Table: %s, column %s not found in select list, skipping.", ti.Name, name)
			continue
		}
		ti.AddColumn(name, columnType, row[4], row[5].String())
	}
	return nil
}

// SetPK sets the pk columns for a TableInfo.
func (ti *TableInfo) SetPK(colnames []string) error {
	pkIndex := schema.NewIndex("PRIMARY")
	colnums := make([]int, len(colnames))
	for i, colname := range colnames {
		colnums[i] = ti.FindColumn(sqlparser.NewColIdent(colname))
		if colnums[i] == -1 {
			return fmt.Errorf("column %s not found", colname)
		}
		pkIndex.AddColumn(colname, 1)
	}
	for _, col := range ti.Columns {
		pkIndex.DataColumns = append(pkIndex.DataColumns, col.Name)
	}
	ti.Indexes = append(ti.Indexes, nil)
	copy(ti.Indexes[1:], ti.Indexes[:len(ti.Indexes)-1])
	ti.Indexes[0] = pkIndex
	ti.PKColumns = colnums
	return nil
}

func (ti *TableInfo) fetchIndexes(conn *DBConn, sqlTableName string) error {
	indexes, err := conn.Exec(context.Background(), fmt.Sprintf("show index from %s", sqlTableName), 10000, false)
	if err != nil {
		return err
	}
	var currentIndex *schema.Index
	currentName := ""
	for _, row := range indexes.Rows {
		indexName := row[2].String()
		if currentName != indexName {
			currentIndex = ti.AddIndex(indexName)
			currentName = indexName
		}
		var cardinality uint64
		if !row[6].IsNull() {
			cardinality, err = strconv.ParseUint(row[6].String(), 0, 64)
			if err != nil {
				log.Warningf("%s", err)
			}
		}
		currentIndex.AddColumn(row[4].String(), cardinality)
	}
	if !ti.HasPrimary() {
		return nil
	}
	pkIndex := ti.Indexes[0]
	ti.PKColumns = make([]int, len(pkIndex.Columns))
	for i, pkCol := range pkIndex.Columns {
		ti.PKColumns[i] = ti.FindColumn(pkCol)
	}
	// Primary key contains all table columns
	for _, col := range ti.Columns {
		pkIndex.DataColumns = append(pkIndex.DataColumns, col.Name)
	}
	// Secondary indices contain all primary key columns
	for i := 1; i < len(ti.Indexes); i++ {
		for _, c := range ti.Indexes[i].Columns {
			ti.Indexes[i].DataColumns = append(ti.Indexes[i].DataColumns, c)
		}
		for _, c := range pkIndex.Columns {
			// pk columns may already be part of the index. So,
			// check before adding.
			if ti.Indexes[i].FindDataColumn(c) != -1 {
				continue
			}
			ti.Indexes[i].DataColumns = append(ti.Indexes[i].DataColumns, c)
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package goreleaserlib

import (
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"testing"

	yaml "gopkg.in/yaml.v1"

	"github.com/goreleaser/goreleaser/config"
	"github.com/stretchr/testify/assert"
)

func init() {
	_ = os.Unsetenv("GITHUB_TOKEN")
}

func TestRelease(t *testing.T) {
	var assert = assert.New(t)
	_, back := setup(t)
	defer back()
	var flags = fakeFlags{
		flags: map[string]string{
			"skip-publish":  "true",
			"skip-validate": "true",
		},
	}
	assert.NoError(Release(flags))
}

func TestConfigFileIsSetAndDontExist(t *testing.T) {
	var assert = assert.New(t)
	var flags = fakeFlags{
		flags: map[string]string{
			"config": "/this/wont/exist",
		},
	}
	assert.Error(Release(flags))
}

func TestReleaseNotesFileDontExist(t *testing.T) {
	var assert = assert.New(t)
	var flags = fakeFlags{
		flags: map[string]string{
			"release-notes": "/this/also/wont/exist",
		},
	}
	assert.Error(Release(flags))
}

func TestCustomReleaseNotesFile(t *testing.T) {
	var assert = assert.New(t)
	folder, back := setup(t)
	defer back()
	var releaseNotes = filepath.Join(folder, "notes.md")
	createFile(t, releaseNotes, "nothing important at all")
	var flags = fakeFlags{
		flags: map[string]string{
			"release-notes": releaseNotes,
			"skip-publish":  "true",
			"skip-validate": "true",
		},
	}
	assert.NoError(Release(flags))
}

func TestBrokenPipe(t *testing.T) {
	var assert = assert.New(t)
	_, back := setup(t)
	defer back()
	createFile(t, "main.go", "not a valid go file")
	var flags = fakeFlags{
		flags: map[string]string{
			"skip-publish":  "true",
			"skip-validate": "true",
		},
	}
	assert.Error(Release(flags))
}

// fakeFlags is a mock of the cli flags
type fakeFlags struct {
	flags map[string]string
}

func (f fakeFlags) IsSet(s string) bool {
	return f.flags[s] != ""
}
func (f fakeFlags) String(s string) string {
	return f.flags[s]
}
func (f fakeFlags) Bool(s string) bool {
	return f.flags[s] == "true"
}

func setup(t *testing.T) (current string, back func()) {
	var assert = assert.New(t)
	folder, err := ioutil.TempDir("", "goreleaser")
	assert.NoError(err)
	log.Println("Folder:", folder)
	previous, err := os.Getwd()
	assert.NoError(err)
	assert.NoError(os.Chdir(folder))
	var gitCmds = [][]string{
		{"init"},
		{"add", "-A"},
		{"commit", "--allow-empty", "-m", "asdf"},
		{"tag", "v0.0.1"},
		{"commit", "--allow-empty", "-m", "asas89d"},
		{"commit", "--allow-empty", "-m", "assssf"},
		{"commit", "--allow-empty", "-m", "assd"},
		{"tag", "v0.0.2"},
		{"remote", "add", "origin", "git@github.com:goreleaser/fake.git"},
	}
	createGoreleaserYaml(t)
	createMainGo(t)
	for _, cmd := range gitCmds {
		var args = []string{
			"-c",
			"user.name='GoReleaser'",
			"-c",
			"user.email='test@goreleaser.github.com'",
		}
		args = append(args, cmd...)
		assert.NoError(exec.Command("git", args...).Run())
	}
	return folder, func() {
		assert.NoError(os.Chdir(previous))
	}
}

func createFile(t *testing.T, filename, contents string) {
	var assert = assert.New(t)
	assert.NoError(ioutil.WriteFile(filename, []byte(contents), 0644))
}

func createMainGo(t *testing.T) {
	createFile(t, "main.go", "package main\nfunc main() {println(0)}")
}

func createGoreleaserYaml(t *testing.T) {
	var yaml = `build:
  binary: fake
  goos:
    - linux
  goarch:
    - amd64
release:
  github:
    owner: goreleaser
    name: fake
`
	createFile(t, "goreleaser.yml", yaml)
}

func TestInitProject(t *testing.T) {
	var filename = "test_goreleaser.yml"

	defer func() {
		if _, err := os.Stat(filename); !os.IsNotExist(err) {
			if err != nil {
				t.Fatal(err.Error())
			}

			if err := os.Remove(filename); err != nil {
				t.Fatal(err.Error())
			}
		}
	}()

	if err := InitProject(filename); err != nil {
		t.Fatalf("exepcted InitProject() to run, but got %v", err.Error())
	}

	file, err := os.Open(filename)
	if err != nil {
		t.Fatal(err.Error())
	}

	out, err := ioutil.ReadAll(file)
	if err != nil {
		t.Fatal(err.Error())
	}

	config := config.Project{}
	assert.NoError(t, yaml.Unmarshal(out, &config))
}

func TestInitProjectFileExist(t *testing.T) {
	var filename = "test_goreleaser.yml"

	createFile(t, filename, "")

	defer func() {
		if _, err := os.Stat(filename); !os.IsNotExist(err) {
			if err != nil {
				t.Fatal(err.Error())
			}

			if err := os.Remove(filename); err != nil {
				t.Fatal(err.Error())
			}
		}
	}()

	assert.Error(t, InitProject(filename))
}
<commit_msg>moving tests<commit_after>package goreleaserlib

import (
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"testing"

	yaml "gopkg.in/yaml.v1"

	"github.com/goreleaser/goreleaser/config"
	"github.com/stretchr/testify/assert"
)

func init() {
	_ = os.Unsetenv("GITHUB_TOKEN")
}

func TestRelease(t *testing.T) {
	var assert = assert.New(t)
	_, back := setup(t)
	defer back()
	var flags = fakeFlags{
		flags: map[string]string{
			"skip-publish":  "true",
			"skip-validate": "true",
		},
	}
	assert.NoError(Release(flags))
}

func TestConfigFileIsSetAndDontExist(t *testing.T) {
	var assert = assert.New(t)
	var flags = fakeFlags{
		flags: map[string]string{
			"config": "/this/wont/exist",
		},
	}
	assert.Error(Release(flags))
}

func TestReleaseNotesFileDontExist(t *testing.T) {
	var assert = assert.New(t)
	var flags = fakeFlags{
		flags: map[string]string{
			"release-notes": "/this/also/wont/exist",
		},
	}
	assert.Error(Release(flags))
}

func TestCustomReleaseNotesFile(t *testing.T) {
	var assert = assert.New(t)
	folder, back := setup(t)
	defer back()
	var releaseNotes = filepath.Join(folder, "notes.md")
	createFile(t, releaseNotes, "nothing important at all")
	var flags = fakeFlags{
		flags: map[string]string{
			"release-notes": releaseNotes,
			"skip-publish":  "true",
			"skip-validate": "true",
		},
	}
	assert.NoError(Release(flags))
}

func TestBrokenPipe(t *testing.T) {
	var assert = assert.New(t)
	_, back := setup(t)
	defer back()
	createFile(t, "main.go", "not a valid go file")
	var flags = fakeFlags{
		flags: map[string]string{
			"skip-publish":  "true",
			"skip-validate": "true",
		},
	}
	assert.Error(Release(flags))
}

func TestInitProject(t *testing.T) {
	var filename = "test_goreleaser.yml"

	defer func() {
		if _, err := os.Stat(filename); !os.IsNotExist(err) {
			if err != nil {
				t.Fatal(err.Error())
			}

			if err := os.Remove(filename); err != nil {
				t.Fatal(err.Error())
			}
		}
	}()

	if err := InitProject(filename); err != nil {
		t.Fatalf("exepcted InitProject() to run, but got %v", err.Error())
	}

	file, err := os.Open(filename)
	if err != nil {
		t.Fatal(err.Error())
	}

	out, err := ioutil.ReadAll(file)
	if err != nil {
		t.Fatal(err.Error())
	}

	config := config.Project{}
	assert.NoError(t, yaml.Unmarshal(out, &config))
}

func TestInitProjectFileExist(t *testing.T) {
	var filename = "test_goreleaser.yml"

	createFile(t, filename, "")

	defer func() {
		if _, err := os.Stat(filename); !os.IsNotExist(err) {
			if err != nil {
				t.Fatal(err.Error())
			}

			if err := os.Remove(filename); err != nil {
				t.Fatal(err.Error())
			}
		}
	}()

	assert.Error(t, InitProject(filename))
}

// fakeFlags is a mock of the cli flags
type fakeFlags struct {
	flags map[string]string
}

func (f fakeFlags) IsSet(s string) bool {
	return f.flags[s] != ""
}
func (f fakeFlags) String(s string) string {
	return f.flags[s]
}
func (f fakeFlags) Bool(s string) bool {
	return f.flags[s] == "true"
}

func setup(t *testing.T) (current string, back func()) {
	var assert = assert.New(t)
	folder, err := ioutil.TempDir("", "goreleaser")
	assert.NoError(err)
	log.Println("Folder:", folder)
	previous, err := os.Getwd()
	assert.NoError(err)
	assert.NoError(os.Chdir(folder))
	var gitCmds = [][]string{
		{"init"},
		{"add", "-A"},
		{"commit", "--allow-empty", "-m", "asdf"},
		{"tag", "v0.0.1"},
		{"commit", "--allow-empty", "-m", "asas89d"},
		{"commit", "--allow-empty", "-m", "assssf"},
		{"commit", "--allow-empty", "-m", "assd"},
		{"tag", "v0.0.2"},
		{"remote", "add", "origin", "git@github.com:goreleaser/fake.git"},
	}
	createGoreleaserYaml(t)
	createMainGo(t)
	for _, cmd := range gitCmds {
		var args = []string{
			"-c",
			"user.name='GoReleaser'",
			"-c",
			"user.email='test@goreleaser.github.com'",
		}
		args = append(args, cmd...)
		assert.NoError(exec.Command("git", args...).Run())
	}
	return folder, func() {
		assert.NoError(os.Chdir(previous))
	}
}

func createFile(t *testing.T, filename, contents string) {
	var assert = assert.New(t)
	assert.NoError(ioutil.WriteFile(filename, []byte(contents), 0644))
}

func createMainGo(t *testing.T) {
	createFile(t, "main.go", "package main\nfunc main() {println(0)}")
}

func createGoreleaserYaml(t *testing.T) {
	var yaml = `build:
  binary: fake
  goos:
    - linux
  goarch:
    - amd64
release:
  github:
    owner: goreleaser
    name: fake
`
	createFile(t, "goreleaser.yml", yaml)
}
<|endoftext|>


<commit_before><commit_msg>Add: TypeToString<commit_after><|endoftext|>


<commit_before><commit_msg>update makeRaw<commit_after><|endoftext|>


<commit_before>package main

import (
	"fmt"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"sort"
	"strings"
	"syscall"

	"github.com/mitchellh/osext"
)

const envPrefix = "VAGRANT_OLD_ENV"

func main() {
	debug := os.Getenv("VAGRANT_DEBUG_LAUNCHER") != ""

	// Get the path to the executable. This path doesn't resolve symlinks
	// so we have to do that afterwards to find the real binary.
	path, err := osext.Executable()
	if err != nil {
		fmt.Fprintf(os.Stderr, "Failed to load Vagrant: %s\n", err)
		os.Exit(1)
	}
	if debug {
		log.Printf("launcher: path = %s", path)
	}
	for {
		fi, err := os.Lstat(path)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Failed to stat executable: %s\n", err)
			os.Exit(1)
		}
		if fi.Mode()&os.ModeSymlink == 0 {
			break
		}

		// The executable is a symlink, so resolve it
		path, err = os.Readlink(path)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Failed to load Vagrant: %s\n", err)
			os.Exit(1)
		}
		if debug {
			log.Printf("launcher: resolved symlink = %s", path)
		}
	}

	// Determine some basic directories that we use throughout
	path = filepath.Dir(filepath.Clean(path))
	installerDir := filepath.Dir(path)
	embeddedDir := filepath.Join(installerDir, "embedded")
	if debug {
		log.Printf("launcher: installerDir = %s", installerDir)
		log.Printf("launcher: embeddedDir = %s", embeddedDir)
	}

	// Find the Vagrant gem
	gemPaths, err := filepath.Glob(
		filepath.Join(embeddedDir, "gems", "gems", "vagrant-*"))
	if err != nil {
		fmt.Fprintf(os.Stderr, "Failed to find Vagrant: %s\n", err)
		os.Exit(1)
	}
	if debug {
		log.Printf("launcher: gemPaths (initial) = %#v", gemPaths)
	}
	for i := 0; i < len(gemPaths); i++ {
		fullPath := filepath.Join(gemPaths[i], "lib", "vagrant", "pre-rubygems.rb")
		if _, err := os.Stat(fullPath); err != nil {
			if debug {
				log.Printf("launcher: bad gemPath += %s", fullPath)
			}

			gemPaths = append(gemPaths[:i], gemPaths[i+1:]...)
			i--
		}
	}
	if len(gemPaths) == 0 {
		fmt.Fprintf(os.Stderr, "Failed to find Vagrant!\n")
		os.Exit(1)
	}
	gemPath := gemPaths[len(gemPaths)-1]
	vagrantExecutable := filepath.Join(gemPath, "bin", "vagrant")
	if debug {
		log.Printf("launcher: gemPaths (final) = %#v", gemPaths)
		log.Printf("launcher: gemPath = %s", gemPath)
	}

	// Setup the CPP/LDFLAGS so that native extensions can be
	// properly compiled into the Vagrant environment.
	cppflags := "-I" + filepath.Join(embeddedDir, "include")
	ldflags := "-L" + filepath.Join(embeddedDir, "lib")
	if original := os.Getenv("CPPFLAGS"); original != "" {
		cppflags = original + " " + cppflags
	}
	if original := os.Getenv("LDFLAGS"); original != "" {
		ldflags = original + " " + ldflags
	}

	// Set the PATH to include the proper paths into our embedded dir
	path = os.Getenv("PATH")
	if runtime.GOOS == "windows" {
		path = fmt.Sprintf(
			"%s;%s;%s",
			filepath.Join(embeddedDir, "bin"),
			filepath.Join(embeddedDir, "gnuwin32", "bin"),
			path)
	} else {
		path = fmt.Sprintf("%s:%s",
			filepath.Join(embeddedDir, "bin"), path)
	}

	newEnv := map[string]string{
		// Setup the environment to prefer our embedded dir over
		// anything the user might have setup on his/her system.
		"CPPFLAGS":      cppflags,
		"GEM_HOME":      filepath.Join(embeddedDir, "gems"),
		"GEM_PATH":      filepath.Join(embeddedDir, "gems"),
		"GEMRC":         filepath.Join(embeddedDir, "etc", "gemrc"),
		"LDFLAGS":       ldflags,
		"PATH":          path,
		"SSL_CERT_FILE": filepath.Join(embeddedDir, "cacert.pem"),

		// Environmental variables used by Vagrant itself
		"VAGRANT_EXECUTABLE":             vagrantExecutable,
		"VAGRANT_INSTALLER_ENV":          "1",
		"VAGRANT_INSTALLER_EMBEDDED_DIR": embeddedDir,
		"VAGRANT_INSTALLER_VERSION":      "2",
	}

	// Unset any RUBYOPT, we don't want this bleeding into our runtime
	newEnv["RUBYOPT"] = ""

	// Store the "current" environment so Vagrant can restore it when shelling
	// out.
	for _, value := range os.Environ() {
		parts := strings.SplitN(value, "=", 2)
		k := parts[0]
		v := parts[1]

		key := fmt.Sprintf("%s_%s", envPrefix, k)
		newEnv[key] = v
	}
	if debug {
		keys := make([]string, 0, len(newEnv))
		for k, _ := range newEnv {
			keys = append(keys, k)
		}
		sort.Strings(keys)

		for _, k := range keys {
			log.Printf("launcher: env %q = %q", k, newEnv[k])
		}
	}

	// Set all the environmental variables
	for k, v := range newEnv {
		if err := os.Setenv(k, v); err != nil {
			fmt.Fprintf(os.Stderr, "Error setting env var %s: %s\n", k, err)
			os.Exit(1)
		}
	}

	// Determine the path to Ruby and then start the Vagrant process
	rubyPath := filepath.Join(embeddedDir, "bin", "ruby")
	if runtime.GOOS == "windows" {
		rubyPath += ".exe"
	}

	cmd := exec.Command(rubyPath)
	cmd.Args = make([]string, len(os.Args)+1)
	cmd.Args[0] = "ruby"
	cmd.Args[1] = filepath.Join(gemPath, "lib", "vagrant", "pre-rubygems.rb")
	copy(cmd.Args[2:], os.Args[1:])
	if debug {
		log.Printf("launcher: rubyPath = %s", rubyPath)
		log.Printf("launcher: args = %#v", cmd.Args)
	}

	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	if err := cmd.Start(); err != nil {
		fmt.Fprintf(os.Stderr, "Exec error: %s\n", err)
		os.Exit(1)
	}

	exitCode := 0
	if err := cmd.Wait(); err != nil {
		if exiterr, ok := err.(*exec.ExitError); ok {
			// The program has exited with an exit code != 0

			// This works on both Unix and Windows. Although package
			// syscall is generally platform dependent, WaitStatus is
			// defined for both Unix and Windows and in both cases has
			// an ExitStatus() method with the same signature.
			if status, ok := exiterr.Sys().(syscall.WaitStatus); ok {
				exitCode = status.ExitStatus()
			}
		}
	}

	os.Exit(exitCode)
}
<commit_msg>launcher: less allocations<commit_after>package main

import (
	"fmt"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"sort"
	"strings"
	"syscall"

	"github.com/mitchellh/osext"
)

const envPrefix = "VAGRANT_OLD_ENV"

func main() {
	debug := os.Getenv("VAGRANT_DEBUG_LAUNCHER") != ""

	// Get the path to the executable. This path doesn't resolve symlinks
	// so we have to do that afterwards to find the real binary.
	path, err := osext.Executable()
	if err != nil {
		fmt.Fprintf(os.Stderr, "Failed to load Vagrant: %s\n", err)
		os.Exit(1)
	}
	if debug {
		log.Printf("launcher: path = %s", path)
	}
	for {
		fi, err := os.Lstat(path)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Failed to stat executable: %s\n", err)
			os.Exit(1)
		}
		if fi.Mode()&os.ModeSymlink == 0 {
			break
		}

		// The executable is a symlink, so resolve it
		path, err = os.Readlink(path)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Failed to load Vagrant: %s\n", err)
			os.Exit(1)
		}
		if debug {
			log.Printf("launcher: resolved symlink = %s", path)
		}
	}

	// Determine some basic directories that we use throughout
	path = filepath.Dir(filepath.Clean(path))
	installerDir := filepath.Dir(path)
	embeddedDir := filepath.Join(installerDir, "embedded")
	if debug {
		log.Printf("launcher: installerDir = %s", installerDir)
		log.Printf("launcher: embeddedDir = %s", embeddedDir)
	}

	// Find the Vagrant gem
	gemPaths, err := filepath.Glob(
		filepath.Join(embeddedDir, "gems", "gems", "vagrant-*"))
	if err != nil {
		fmt.Fprintf(os.Stderr, "Failed to find Vagrant: %s\n", err)
		os.Exit(1)
	}
	if debug {
		log.Printf("launcher: gemPaths (initial) = %#v", gemPaths)
	}
	for i := 0; i < len(gemPaths); i++ {
		fullPath := filepath.Join(gemPaths[i], "lib", "vagrant", "pre-rubygems.rb")
		if _, err := os.Stat(fullPath); err != nil {
			if debug {
				log.Printf("launcher: bad gemPath += %s", fullPath)
			}

			gemPaths = append(gemPaths[:i], gemPaths[i+1:]...)
			i--
		}
	}
	if len(gemPaths) == 0 {
		fmt.Fprintf(os.Stderr, "Failed to find Vagrant!\n")
		os.Exit(1)
	}
	gemPath := gemPaths[len(gemPaths)-1]
	vagrantExecutable := filepath.Join(gemPath, "bin", "vagrant")
	if debug {
		log.Printf("launcher: gemPaths (final) = %#v", gemPaths)
		log.Printf("launcher: gemPath = %s", gemPath)
	}

	// Setup the CPP/LDFLAGS so that native extensions can be
	// properly compiled into the Vagrant environment.
	cppflags := "-I" + filepath.Join(embeddedDir, "include")
	ldflags := "-L" + filepath.Join(embeddedDir, "lib")
	if original := os.Getenv("CPPFLAGS"); original != "" {
		cppflags = original + " " + cppflags
	}
	if original := os.Getenv("LDFLAGS"); original != "" {
		ldflags = original + " " + ldflags
	}

	// Set the PATH to include the proper paths into our embedded dir
	path = os.Getenv("PATH")
	if runtime.GOOS == "windows" {
		path = fmt.Sprintf(
			"%s;%s;%s",
			filepath.Join(embeddedDir, "bin"),
			filepath.Join(embeddedDir, "gnuwin32", "bin"),
			path)
	} else {
		path = fmt.Sprintf("%s:%s",
			filepath.Join(embeddedDir, "bin"), path)
	}

	newEnv := map[string]string{
		// Setup the environment to prefer our embedded dir over
		// anything the user might have setup on his/her system.
		"CPPFLAGS":      cppflags,
		"GEM_HOME":      filepath.Join(embeddedDir, "gems"),
		"GEM_PATH":      filepath.Join(embeddedDir, "gems"),
		"GEMRC":         filepath.Join(embeddedDir, "etc", "gemrc"),
		"LDFLAGS":       ldflags,
		"PATH":          path,
		"SSL_CERT_FILE": filepath.Join(embeddedDir, "cacert.pem"),

		// Environmental variables used by Vagrant itself
		"VAGRANT_EXECUTABLE":             vagrantExecutable,
		"VAGRANT_INSTALLER_ENV":          "1",
		"VAGRANT_INSTALLER_EMBEDDED_DIR": embeddedDir,
		"VAGRANT_INSTALLER_VERSION":      "2",
	}

	// Unset any RUBYOPT, we don't want this bleeding into our runtime
	newEnv["RUBYOPT"] = ""

	// Store the "current" environment so Vagrant can restore it when shelling
	// out.
	for _, value := range os.Environ() {
		idx := strings.IndexRune(value, '=')
		key := fmt.Sprintf("%s_%s", envPrefix, value[:idx])
		newEnv[key] = value[idx+1:]
	}
	if debug {
		keys := make([]string, 0, len(newEnv))
		for k, _ := range newEnv {
			keys = append(keys, k)
		}
		sort.Strings(keys)

		for _, k := range keys {
			log.Printf("launcher: env %q = %q", k, newEnv[k])
		}
	}

	// Set all the environmental variables
	for k, v := range newEnv {
		if err := os.Setenv(k, v); err != nil {
			fmt.Fprintf(os.Stderr, "Error setting env var %s: %s\n", k, err)
			os.Exit(1)
		}
	}

	// Determine the path to Ruby and then start the Vagrant process
	rubyPath := filepath.Join(embeddedDir, "bin", "ruby")
	if runtime.GOOS == "windows" {
		rubyPath += ".exe"
	}

	cmd := exec.Command(rubyPath)
	cmd.Args = make([]string, len(os.Args)+1)
	cmd.Args[0] = "ruby"
	cmd.Args[1] = filepath.Join(gemPath, "lib", "vagrant", "pre-rubygems.rb")
	copy(cmd.Args[2:], os.Args[1:])
	if debug {
		log.Printf("launcher: rubyPath = %s", rubyPath)
		log.Printf("launcher: args = %#v", cmd.Args)
	}

	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	if err := cmd.Start(); err != nil {
		fmt.Fprintf(os.Stderr, "Exec error: %s\n", err)
		os.Exit(1)
	}

	exitCode := 0
	if err := cmd.Wait(); err != nil {
		if exiterr, ok := err.(*exec.ExitError); ok {
			// The program has exited with an exit code != 0

			// This works on both Unix and Windows. Although package
			// syscall is generally platform dependent, WaitStatus is
			// defined for both Unix and Windows and in both cases has
			// an ExitStatus() method with the same signature.
			if status, ok := exiterr.Sys().(syscall.WaitStatus); ok {
				exitCode = status.ExitStatus()
			}
		}
	}

	os.Exit(exitCode)
}
<|endoftext|>