<commit_before>package brats_test

import (
	"fmt"
	"os/exec"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"time"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/gbytes"
	"github.com/onsi/gomega/gexec"
)

func extractAzIpsMap(regex *regexp.Regexp, contents string) map[string][]string {
	out := map[string][]string{
		"z1": {},
		"z2": {},
	}

	instances := regex.FindAllStringSubmatch(contents, -1)
	Expect(instances).ToNot(BeNil())
	for _, q := range instances {
		out[q[1]] = append(out[q[1]], q[2])
	}

	return out
}

func mustGetLatestDnsVersions() []int {
	session, err := gexec.Start(exec.Command(
		boshBinaryPath, "-n",
		"-d", deploymentName,
		"ssh",
		"-c", "sudo cat /var/vcap/instance/dns/records.json",
	), GinkgoWriter, GinkgoWriter)
	Expect(err).ToNot(HaveOccurred())
	Eventually(session, time.Minute).Should(gexec.Exit(0))

	trimmedOutput := strings.TrimSpace(string(session.Out.Contents()))

	results := extractDnsVersionsList(trimmedOutput)
	Expect(len(results)).To(BeNumerically(">", 0))

	return results
}

var versionSegmentsPattern = regexp.MustCompile(`"version":(\d+)`)

func extractDnsVersionsList(sshContents string) []int {
	matches := versionSegmentsPattern.FindAllStringSubmatch(sshContents, -1)
	Expect(matches).ToNot(BeNil())
	results := make([]int, len(matches))

	for i, match := range matches {
		Expect(len(match)).To(Equal(2))
		value, err := strconv.Atoi(match[1])
		Expect(err).ToNot(HaveOccurred())

		results[i] = value
	}
	return results
}

var _ = Describe("BoshDns", func() {
	var (
		manifestPath              string
		linkedTemplateReleasePath string
	)

	BeforeEach(func() {
		startInnerBosh()

		session, err := gexec.Start(exec.Command(boshBinaryPath, "-n", "upload-stemcell", candidateWardenLinuxStemcellPath), GinkgoWriter, GinkgoWriter)
		Expect(err).ToNot(HaveOccurred())
		Eventually(session, 5*time.Minute).Should(gexec.Exit(0))

		manifestPath, err = filepath.Abs("../assets/dns-with-templates-manifest.yml")
		Expect(err).NotTo(HaveOccurred())

		linkedTemplateReleasePath, err = filepath.Abs("../assets/linked-templates-release")
		Expect(err).NotTo(HaveOccurred())
	})

	AfterEach(stopInnerBosh)

	Context("having enabled short dns addresses", func() {
		BeforeEach(func() {
			opFilePath, err := filepath.Abs("../assets/op-enable-short-dns-addresses.yml")
			Expect(err).NotTo(HaveOccurred())

			session, err := gexec.Start(exec.Command(
				boshBinaryPath, "deploy",
				"-n",
				"-d", deploymentName,
				manifestPath,
				"-o", opFilePath,
				"-v", fmt.Sprintf("dns-release-path=%s", dnsReleasePath),
				"-v", fmt.Sprintf("linked-template-release-path=%s", linkedTemplateReleasePath),
			), GinkgoWriter, GinkgoWriter)
			Expect(err).ToNot(HaveOccurred())
			Eventually(session, 15*time.Minute).Should(gexec.Exit(0))
		})

		It("can find instances using the address helper with short names", func() {
			session, err := gexec.Start(exec.Command(
				boshBinaryPath, "-n",
				"-d", deploymentName,
				"instances",
				"--column", "instance",
				"--column", "az",
				"--column", "ips",
			), GinkgoWriter, GinkgoWriter)
			Expect(err).ToNot(HaveOccurred())
			Eventually(session, time.Minute).Should(gexec.Exit(0))

			instanceList := session.Out.Contents()

			matchExpression := regexp.MustCompile(`provider\S+\s+(z1|z2)\s+(\S+)`)
			knownProviders := extractAzIpsMap(matchExpression, string(instanceList))

			session, err = gexec.Start(exec.Command(boshBinaryPath,
				"-d", deploymentName,
				"run-errand", "query-all",
			), GinkgoWriter, GinkgoWriter)
			Expect(err).ToNot(HaveOccurred())
			Eventually(session, time.Minute).Should(gexec.Exit(0))

			Expect(session.Out).To(gbytes.Say("ANSWER: 3"))

			output := string(session.Out.Contents())

			for _, ips := range knownProviders {
				for _, ip := range ips {
					Expect(output).To(MatchRegexp(`q-s0\.g-\d+\.bosh\.\s+\d+\s+IN\s+A\s+%s`, ip))
				}
			}
		})

		FIt("can find instances using the address helper with short names by network and instance ID", func() {
			session, err := gexec.Start(exec.Command(
				boshBinaryPath, "-n",
				"-d", deploymentName,
				"instances",
				"--column", "instance",
				"--column", "az",
				"--column", "ips",
			), GinkgoWriter, GinkgoWriter)
			Expect(err).ToNot(HaveOccurred())
			Eventually(session, time.Minute).Should(gexec.Exit(0))

			instanceList := session.Out.Contents()

			matchExpression := regexp.MustCompile(`provider\S+\s+(z1)\s+(\S+)`)
			knownProviders := extractAzIpsMap(matchExpression, string(instanceList))

			session, err = gexec.Start(exec.Command(boshBinaryPath,
				"-d", deploymentName,
				"run-errand", "query-individual-instance",
			), GinkgoWriter, GinkgoWriter)
			Expect(err).ToNot(HaveOccurred())
			Eventually(session, time.Minute).Should(gexec.Exit(0))

			Expect(session.Out).To(gbytes.Say("ANSWER: 1"))

			output := string(session.Out.Contents())

			ip := knownProviders["z1"][0]
			Expect(output).To(MatchRegexp(`q-n\d+m\d+\.g-\d\.bosh\.\s+\d+\s+IN\s+A\s+%s`, ip))
		})
	})

	Context("When deploying vms across different azs", func() {
		BeforeEach(func() {
			session, err := gexec.Start(exec.Command(
				boshBinaryPath, "deploy",
				"-n",
				"-d", deploymentName,
				manifestPath,
				"-v", fmt.Sprintf("dns-release-path=%s", dnsReleasePath),
				"-v", fmt.Sprintf("linked-template-release-path=%s", linkedTemplateReleasePath),
			), GinkgoWriter, GinkgoWriter)
			Expect(err).ToNot(HaveOccurred())
			Eventually(session, 15*time.Minute).Should(gexec.Exit(0))
		})

		It("can find instances using the address helper", func() {
			session, err := gexec.Start(exec.Command(
				boshBinaryPath, "-n",
				"-d", deploymentName,
				"instances",
				"--column", "instance",
				"--column", "az",
				"--column", "ips",
			), GinkgoWriter, GinkgoWriter)
			Expect(err).ToNot(HaveOccurred())
			Eventually(session, time.Minute).Should(gexec.Exit(0))

			instanceList := session.Out.Contents()

			By("finding instances in all AZs", func() {
				matchExpression := regexp.MustCompile(`provider\S+\s+(z1|z2)\s+(\S+)`)
				knownProviders := extractAzIpsMap(matchExpression, string(instanceList))

				session, err = gexec.Start(exec.Command(boshBinaryPath,
					"-d", deploymentName,
					"run-errand", "query-all",
				), GinkgoWriter, GinkgoWriter)
				Expect(err).ToNot(HaveOccurred())
				Eventually(session, time.Minute).Should(gexec.Exit(0))

				Expect(session.Out).To(gbytes.Say("ANSWER: 3"))
				output := string(session.Out.Contents())

				for _, ips := range knownProviders {
					for _, ip := range ips {
						Expect(output).To(ContainSubstring(ip))
					}
				}
			})

			By("finding instances filtering by AZ", func() {
				matchExpression := regexp.MustCompile(`provider\S+\s+(z1)\s+(\S+)`)
				knownProviders := extractAzIpsMap(matchExpression, string(instanceList))

				session, err = gexec.Start(exec.Command(boshBinaryPath,
					"-d", deploymentName,
					"run-errand", "query-with-az-filter",
				), GinkgoWriter, GinkgoWriter)
				Expect(err).ToNot(HaveOccurred())
				Eventually(session, time.Minute).Should(gexec.Exit(0))

				Expect(session.Out).To(gbytes.Say("ANSWER: 2"))
				output := string(session.Out.Contents())

				for _, ips := range knownProviders {
					for _, ip := range ips {
						Expect(output).To(ContainSubstring(ip))
					}
				}
			})
		})

		It("can force a new DNS blob to propagate to ALL vms", func() {
			versionPerInstance := mustGetLatestDnsVersions()
			previousMax := -1
			for _, version := range versionPerInstance {
				if previousMax < version {
					previousMax = version
				}
			}

			session, err := gexec.Start(exec.Command("ssh",
				fmt.Sprintf("%s@%s", innerDirectorUser, innerDirectorIP),
				"-i", innerBoshJumpboxPrivateKeyPath,
				"-oStrictHostKeyChecking=no",
				"sudo /var/vcap/jobs/director/bin/sync_dns_ctl force"),
				GinkgoWriter, GinkgoWriter)
			Expect(err).ToNot(HaveOccurred())
			Eventually(session, 2*time.Minute).Should(gexec.Exit(0))

			newVersionPerInstance := mustGetLatestDnsVersions()
			firstNewVersion := newVersionPerInstance[0]
			Expect(firstNewVersion).To(BeNumerically(">", previousMax))
			for _, version := range newVersionPerInstance {
				Expect(version).To(Equal(firstNewVersion))
			}
		})
	})
})
<commit_msg>Temporarily pend n/m short dns test<commit_after>package brats_test

import (
	"fmt"
	"os/exec"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"time"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/gbytes"
	"github.com/onsi/gomega/gexec"
)

func extractAzIpsMap(regex *regexp.Regexp, contents string) map[string][]string {
	out := map[string][]string{
		"z1": {},
		"z2": {},
	}

	instances := regex.FindAllStringSubmatch(contents, -1)
	Expect(instances).ToNot(BeNil())
	for _, q := range instances {
		out[q[1]] = append(out[q[1]], q[2])
	}

	return out
}

func mustGetLatestDnsVersions() []int {
	session, err := gexec.Start(exec.Command(
		boshBinaryPath, "-n",
		"-d", deploymentName,
		"ssh",
		"-c", "sudo cat /var/vcap/instance/dns/records.json",
	), GinkgoWriter, GinkgoWriter)
	Expect(err).ToNot(HaveOccurred())
	Eventually(session, time.Minute).Should(gexec.Exit(0))

	trimmedOutput := strings.TrimSpace(string(session.Out.Contents()))

	results := extractDnsVersionsList(trimmedOutput)
	Expect(len(results)).To(BeNumerically(">", 0))

	return results
}

var versionSegmentsPattern = regexp.MustCompile(`"version":(\d+)`)

func extractDnsVersionsList(sshContents string) []int {
	matches := versionSegmentsPattern.FindAllStringSubmatch(sshContents, -1)
	Expect(matches).ToNot(BeNil())
	results := make([]int, len(matches))

	for i, match := range matches {
		Expect(len(match)).To(Equal(2))
		value, err := strconv.Atoi(match[1])
		Expect(err).ToNot(HaveOccurred())

		results[i] = value
	}
	return results
}

var _ = Describe("BoshDns", func() {
	var (
		manifestPath              string
		linkedTemplateReleasePath string
	)

	BeforeEach(func() {
		startInnerBosh()

		session, err := gexec.Start(exec.Command(boshBinaryPath, "-n", "upload-stemcell", candidateWardenLinuxStemcellPath), GinkgoWriter, GinkgoWriter)
		Expect(err).ToNot(HaveOccurred())
		Eventually(session, 5*time.Minute).Should(gexec.Exit(0))

		manifestPath, err = filepath.Abs("../assets/dns-with-templates-manifest.yml")
		Expect(err).NotTo(HaveOccurred())

		linkedTemplateReleasePath, err = filepath.Abs("../assets/linked-templates-release")
		Expect(err).NotTo(HaveOccurred())
	})

	AfterEach(stopInnerBosh)

	Context("having enabled short dns addresses", func() {
		BeforeEach(func() {
			opFilePath, err := filepath.Abs("../assets/op-enable-short-dns-addresses.yml")
			Expect(err).NotTo(HaveOccurred())

			session, err := gexec.Start(exec.Command(
				boshBinaryPath, "deploy",
				"-n",
				"-d", deploymentName,
				manifestPath,
				"-o", opFilePath,
				"-v", fmt.Sprintf("dns-release-path=%s", dnsReleasePath),
				"-v", fmt.Sprintf("linked-template-release-path=%s", linkedTemplateReleasePath),
			), GinkgoWriter, GinkgoWriter)
			Expect(err).ToNot(HaveOccurred())
			Eventually(session, 15*time.Minute).Should(gexec.Exit(0))
		})

		It("can find instances using the address helper with short names", func() {
			session, err := gexec.Start(exec.Command(
				boshBinaryPath, "-n",
				"-d", deploymentName,
				"instances",
				"--column", "instance",
				"--column", "az",
				"--column", "ips",
			), GinkgoWriter, GinkgoWriter)
			Expect(err).ToNot(HaveOccurred())
			Eventually(session, time.Minute).Should(gexec.Exit(0))

			instanceList := session.Out.Contents()

			matchExpression := regexp.MustCompile(`provider\S+\s+(z1|z2)\s+(\S+)`)
			knownProviders := extractAzIpsMap(matchExpression, string(instanceList))

			session, err = gexec.Start(exec.Command(boshBinaryPath,
				"-d", deploymentName,
				"run-errand", "query-all",
			), GinkgoWriter, GinkgoWriter)
			Expect(err).ToNot(HaveOccurred())
			Eventually(session, time.Minute).Should(gexec.Exit(0))

			Expect(session.Out).To(gbytes.Say("ANSWER: 3"))

			output := string(session.Out.Contents())

			for _, ips := range knownProviders {
				for _, ip := range ips {
					Expect(output).To(MatchRegexp(`q-s0\.g-\d+\.bosh\.\s+\d+\s+IN\s+A\s+%s`, ip))
				}
			}
		})

		PIt("can find instances using the address helper with short names by network and instance ID", func() {
			session, err := gexec.Start(exec.Command(
				boshBinaryPath, "-n",
				"-d", deploymentName,
				"instances",
				"--column", "instance",
				"--column", "az",
				"--column", "ips",
			), GinkgoWriter, GinkgoWriter)
			Expect(err).ToNot(HaveOccurred())
			Eventually(session, time.Minute).Should(gexec.Exit(0))

			instanceList := session.Out.Contents()

			matchExpression := regexp.MustCompile(`provider\S+\s+(z1)\s+(\S+)`)
			knownProviders := extractAzIpsMap(matchExpression, string(instanceList))

			session, err = gexec.Start(exec.Command(boshBinaryPath,
				"-d", deploymentName,
				"run-errand", "query-individual-instance",
			), GinkgoWriter, GinkgoWriter)
			Expect(err).ToNot(HaveOccurred())
			Eventually(session, time.Minute).Should(gexec.Exit(0))

			Expect(session.Out).To(gbytes.Say("ANSWER: 1"))

			output := string(session.Out.Contents())

			ip := knownProviders["z1"][0]
			Expect(output).To(MatchRegexp(`q-n\d+m\d+\.g-\d\.bosh\.\s+\d+\s+IN\s+A\s+%s`, ip))
		})
	})

	Context("When deploying vms across different azs", func() {
		BeforeEach(func() {
			session, err := gexec.Start(exec.Command(
				boshBinaryPath, "deploy",
				"-n",
				"-d", deploymentName,
				manifestPath,
				"-v", fmt.Sprintf("dns-release-path=%s", dnsReleasePath),
				"-v", fmt.Sprintf("linked-template-release-path=%s", linkedTemplateReleasePath),
			), GinkgoWriter, GinkgoWriter)
			Expect(err).ToNot(HaveOccurred())
			Eventually(session, 15*time.Minute).Should(gexec.Exit(0))
		})

		It("can find instances using the address helper", func() {
			session, err := gexec.Start(exec.Command(
				boshBinaryPath, "-n",
				"-d", deploymentName,
				"instances",
				"--column", "instance",
				"--column", "az",
				"--column", "ips",
			), GinkgoWriter, GinkgoWriter)
			Expect(err).ToNot(HaveOccurred())
			Eventually(session, time.Minute).Should(gexec.Exit(0))

			instanceList := session.Out.Contents()

			By("finding instances in all AZs", func() {
				matchExpression := regexp.MustCompile(`provider\S+\s+(z1|z2)\s+(\S+)`)
				knownProviders := extractAzIpsMap(matchExpression, string(instanceList))

				session, err = gexec.Start(exec.Command(boshBinaryPath,
					"-d", deploymentName,
					"run-errand", "query-all",
				), GinkgoWriter, GinkgoWriter)
				Expect(err).ToNot(HaveOccurred())
				Eventually(session, time.Minute).Should(gexec.Exit(0))

				Expect(session.Out).To(gbytes.Say("ANSWER: 3"))
				output := string(session.Out.Contents())

				for _, ips := range knownProviders {
					for _, ip := range ips {
						Expect(output).To(ContainSubstring(ip))
					}
				}
			})

			By("finding instances filtering by AZ", func() {
				matchExpression := regexp.MustCompile(`provider\S+\s+(z1)\s+(\S+)`)
				knownProviders := extractAzIpsMap(matchExpression, string(instanceList))

				session, err = gexec.Start(exec.Command(boshBinaryPath,
					"-d", deploymentName,
					"run-errand", "query-with-az-filter",
				), GinkgoWriter, GinkgoWriter)
				Expect(err).ToNot(HaveOccurred())
				Eventually(session, time.Minute).Should(gexec.Exit(0))

				Expect(session.Out).To(gbytes.Say("ANSWER: 2"))
				output := string(session.Out.Contents())

				for _, ips := range knownProviders {
					for _, ip := range ips {
						Expect(output).To(ContainSubstring(ip))
					}
				}
			})
		})

		It("can force a new DNS blob to propagate to ALL vms", func() {
			versionPerInstance := mustGetLatestDnsVersions()
			previousMax := -1
			for _, version := range versionPerInstance {
				if previousMax < version {
					previousMax = version
				}
			}

			session, err := gexec.Start(exec.Command("ssh",
				fmt.Sprintf("%s@%s", innerDirectorUser, innerDirectorIP),
				"-i", innerBoshJumpboxPrivateKeyPath,
				"-oStrictHostKeyChecking=no",
				"sudo /var/vcap/jobs/director/bin/sync_dns_ctl force"),
				GinkgoWriter, GinkgoWriter)
			Expect(err).ToNot(HaveOccurred())
			Eventually(session, 2*time.Minute).Should(gexec.Exit(0))

			newVersionPerInstance := mustGetLatestDnsVersions()
			firstNewVersion := newVersionPerInstance[0]
			Expect(firstNewVersion).To(BeNumerically(">", previousMax))
			for _, version := range newVersionPerInstance {
				Expect(version).To(Equal(firstNewVersion))
			}
		})
	})
})
<|endoftext|>


<commit_before>package gitosis

import (
	"github.com/timeredbull/tsuru/config"
	"io/ioutil"
	. "launchpad.net/gocheck"
	"os"
	"os/exec"
	"path"
	"strings"
	"testing"
)

func Test(t *testing.T) { TestingT(t) }

type S struct {
	gitRoot     string
	gitosisBare string
	gitosisRepo string
}

var _ = Suite(&S{})

type isInGitosisChecker struct{}

func (c *isInGitosisChecker) Info() *CheckerInfo {
	return &CheckerInfo{Name: "IsInGitosis", Params: []string{"str"}}
}

func (c *isInGitosisChecker) Check(params []interface{}, names []string) (bool, string) {
	if len(params) != 1 {
		return false, "you should provide one string parameter"
	}
	str, ok := params[0].(string)
	if !ok {
		return false, "the parameter should be a string"
	}
	gitosisRepo, err := config.GetString("git:gitosis-repo")
	if err != nil {
		return false, "failed to get config"
	}
	path := path.Join(gitosisRepo, "gitosis.conf")
	f, err := os.Open(path)
	if err != nil {
		return false, err.Error()
	}
	defer f.Close()
	content, err := ioutil.ReadAll(f)
	if err != nil {
		return false, err.Error()
	}
	return strings.Contains(string(content), str), ""
}

var IsInGitosis, NotInGitosis Checker = &isInGitosisChecker{}, Not(IsInGitosis)

func (s *S) SetUpSuite(c *C) {
	err := config.ReadConfigFile("../etc/tsuru.conf")
	c.Assert(err, IsNil)
	s.gitRoot, err = config.GetString("git:root")
	c.Assert(err, IsNil)
	s.gitosisBare, err = config.GetString("git:gitosis-bare")
	c.Assert(err, IsNil)
	s.gitosisRepo, err = config.GetString("git:gitosis-repo")
	err = os.RemoveAll(s.gitRoot)
	c.Assert(err, IsNil)
	err = os.MkdirAll(s.gitRoot, 0777)
	c.Assert(err, IsNil)
	err = exec.Command("git", "init", "--bare", s.gitosisBare).Run()
	c.Assert(err, IsNil)
	err = exec.Command("git", "clone", s.gitosisBare, s.gitosisRepo).Run()
	c.Assert(err, IsNil)
}

func (s *S) SetUpTest(c *C) {
	fpath := path.Join(s.gitosisRepo, "gitosis.conf")
	f, err := os.Create(fpath)
	c.Assert(err, IsNil)
	f.Close()
}

func (s *S) TearDownSuite(c *C) {
	err := os.RemoveAll(s.gitRoot)
	c.Assert(err, IsNil)
}

func (s *S) TearDownTest(c *C) {
	_, err := runGit("rm", "gitosis.conf")
	if err == nil {
		err = pushToGitosis("removing test file")
		c.Assert(err, IsNil)
	}
}

func (s *S) lastBareCommit(c *C) string {
	bareOutput, err := exec.Command("git", "--git-dir="+s.gitosisBare, "log", "-1", "--pretty=format:%s").CombinedOutput()
	c.Assert(err, IsNil)
	return string(bareOutput)
}
<commit_msg>gitosis: set log file for use in tests<commit_after>package gitosis

import (
	"github.com/timeredbull/tsuru/config"
	"github.com/timeredbull/tsuru/log"
	"io/ioutil"
	. "launchpad.net/gocheck"
	stdlog "log"
	"os"
	"os/exec"
	"path"
	"strings"
	"testing"
)

func Test(t *testing.T) { TestingT(t) }

type S struct {
	gitRoot     string
	gitosisBare string
	gitosisRepo string
	logFile     *os.File
}

var _ = Suite(&S{})

type isInGitosisChecker struct{}

func (c *isInGitosisChecker) Info() *CheckerInfo {
	return &CheckerInfo{Name: "IsInGitosis", Params: []string{"str"}}
}

func (c *isInGitosisChecker) Check(params []interface{}, names []string) (bool, string) {
	if len(params) != 1 {
		return false, "you should provide one string parameter"
	}
	str, ok := params[0].(string)
	if !ok {
		return false, "the parameter should be a string"
	}
	gitosisRepo, err := config.GetString("git:gitosis-repo")
	if err != nil {
		return false, "failed to get config"
	}
	path := path.Join(gitosisRepo, "gitosis.conf")
	f, err := os.Open(path)
	if err != nil {
		return false, err.Error()
	}
	defer f.Close()
	content, err := ioutil.ReadAll(f)
	if err != nil {
		return false, err.Error()
	}
	return strings.Contains(string(content), str), ""
}

var IsInGitosis, NotInGitosis Checker = &isInGitosisChecker{}, Not(IsInGitosis)

func (s *S) SetUpSuite(c *C) {
	err := config.ReadConfigFile("../etc/tsuru.conf")
	c.Assert(err, IsNil)
	s.gitRoot, err = config.GetString("git:root")
	c.Assert(err, IsNil)
	s.gitosisBare, err = config.GetString("git:gitosis-bare")
	c.Assert(err, IsNil)
	s.gitosisRepo, err = config.GetString("git:gitosis-repo")
	err = os.RemoveAll(s.gitRoot)
	c.Assert(err, IsNil)
	err = os.MkdirAll(s.gitRoot, 0777)
	c.Assert(err, IsNil)
	err = exec.Command("git", "init", "--bare", s.gitosisBare).Run()
	c.Assert(err, IsNil)
	err = exec.Command("git", "clone", s.gitosisBare, s.gitosisRepo).Run()
	c.Assert(err, IsNil)
	s.logFile, err = os.Create("/tmp/tsuru-tests.log")
	c.Assert(err, IsNil)
	log.Target = stdlog.New(s.logFile, "[tsuru-tests]", stdlog.LstdFlags|stdlog.Llongfile)
}

func (s *S) SetUpTest(c *C) {
	fpath := path.Join(s.gitosisRepo, "gitosis.conf")
	f, err := os.Create(fpath)
	c.Assert(err, IsNil)
	f.Close()
}

func (s *S) TearDownSuite(c *C) {
	defer s.logFile.Close()
	err := os.RemoveAll(s.gitRoot)
	c.Assert(err, IsNil)
}

func (s *S) TearDownTest(c *C) {
	_, err := runGit("rm", "gitosis.conf")
	if err == nil {
		err = pushToGitosis("removing test file")
		c.Assert(err, IsNil)
	}
}

func (s *S) lastBareCommit(c *C) string {
	bareOutput, err := exec.Command("git", "--git-dir="+s.gitosisBare, "log", "-1", "--pretty=format:%s").CombinedOutput()
	c.Assert(err, IsNil)
	return string(bareOutput)
}
<|endoftext|>


<commit_before>package aws

import (
	"bytes"
	"fmt"
	"log"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/codebuild"
	"github.com/hashicorp/terraform/helper/hashcode"
	"github.com/hashicorp/terraform/helper/schema"
	"github.com/hashicorp/terraform/helper/validation"
)

func resourceAwsCodeBuildWebhook() *schema.Resource {
	return &schema.Resource{
		Create: resourceAwsCodeBuildWebhookCreate,
		Read:   resourceAwsCodeBuildWebhookRead,
		Delete: resourceAwsCodeBuildWebhookDelete,
		Update: resourceAwsCodeBuildWebhookUpdate,

		Importer: &schema.ResourceImporter{
			State: schema.ImportStatePassthrough,
		},

		Schema: map[string]*schema.Schema{
			"project_name": {
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},
			"branch_filter": {
				Type:     schema.TypeString,
				Optional: true,
			},
			"filter_group": {
				Type:     schema.TypeSet,
				Optional: true,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"type": {
							Type:     schema.TypeString,
							Required: true,
							ValidateFunc: validation.StringInSlice([]string{
								codebuild.WebhookFilterTypeEvent,
								codebuild.WebhookFilterTypeActorAccountId,
								codebuild.WebhookFilterTypeBaseRef,
								codebuild.WebhookFilterTypeFilePath,
								codebuild.WebhookFilterTypeHeadRef,
							}, false),
						},
						"exclude_matched_pattern": {
							Type:     schema.TypeBool,
							Optional: true,
							Default:  false,
						},
						"pattern": {
							Type:     schema.TypeString,
							Required: true,
						},
					},
				},
				Set: resourceAwsCodeBuildWebhookFilterHash,
			},
			"payload_url": {
				Type:     schema.TypeString,
				Computed: true,
			},
			"secret": {
				Type:      schema.TypeString,
				Computed:  true,
				Sensitive: true,
			},
			"url": {
				Type:     schema.TypeString,
				Computed: true,
			},
		},
	}
}

func resourceAwsCodeBuildWebhookCreate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).codebuildconn

	webhookFilterGroups := expandWebhookFilterGroup(d)

	input := &codebuild.CreateWebhookInput{
		ProjectName:  aws.String(d.Get("project_name").(string)),
		FilterGroups: webhookFilterGroups,
	}

	// The CodeBuild API requires this to be non-empty if defined
	if v, ok := d.GetOk("branch_filter"); ok {
		input.BranchFilter = aws.String(v.(string))
	}

	log.Printf("[DEBUG] Creating CodeBuild Webhook: %s", input)
	resp, err := conn.CreateWebhook(input)

	if err != nil {
		return fmt.Errorf("error creating CodeBuild Webhook: %s", err)
	}

	// Secret is only returned on create, so capture it at the start
	d.Set("secret", resp.Webhook.Secret)
	d.SetId(d.Get("project_name").(string))

	return resourceAwsCodeBuildWebhookRead(d, meta)
}

func expandWebhookFilterGroup(d *schema.ResourceData) [][]*codebuild.WebhookFilter {
	webhookFilters := make([]*codebuild.WebhookFilter, 0)

	configsList := d.Get("filter_group").(*schema.Set).List()

	if len(configsList) == 0 {
		return nil
	}

	for _, config := range configsList {
		filter := expandWebhookFilterData(config.(map[string]interface{}))
		webhookFilters = append(webhookFilters, &filter)
	}

	return [][]*codebuild.WebhookFilter{webhookFilters}
}

func expandWebhookFilterData(data map[string]interface{}) codebuild.WebhookFilter {
	filter := codebuild.WebhookFilter{
		Type:                  aws.String(data["type"].(string)),
		ExcludeMatchedPattern: aws.Bool(data["exclude_matched_pattern"].(bool)),
	}

	if v := data["pattern"]; v != nil {
		filter.Pattern = aws.String(v.(string))
	}

	return filter
}

func resourceAwsCodeBuildWebhookRead(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).codebuildconn

	resp, err := conn.BatchGetProjects(&codebuild.BatchGetProjectsInput{
		Names: []*string{
			aws.String(d.Id()),
		},
	})

	if err != nil {
		return err
	}

	if len(resp.Projects) == 0 {
		log.Printf("[WARN] CodeBuild Project %q not found, removing from state", d.Id())
		d.SetId("")
		return nil
	}

	project := resp.Projects[0]

	if project.Webhook == nil {
		log.Printf("[WARN] CodeBuild Project %q webhook not found, removing from state", d.Id())
		d.SetId("")
		return nil
	}

	d.Set("branch_filter", project.Webhook.BranchFilter)
	d.Set("filter_group", flattenAwsCodeBuildWebhookFilterGroups(project.Webhook.FilterGroups[0]))
	d.Set("payload_url", project.Webhook.PayloadUrl)
	d.Set("project_name", project.Name)
	d.Set("url", project.Webhook.Url)
	// The secret is never returned after creation, so don't set it here

	return nil
}

func resourceAwsCodeBuildWebhookUpdate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).codebuildconn

	var err error
	filterGroups := expandWebhookFilterGroup(d)

	if len(filterGroups) >= 1 {
		_, err = conn.UpdateWebhook(&codebuild.UpdateWebhookInput{
			ProjectName:  aws.String(d.Id()),
			FilterGroups: filterGroups,
			RotateSecret: aws.Bool(false),
		})
	} else {
		_, err = conn.UpdateWebhook(&codebuild.UpdateWebhookInput{
			ProjectName:  aws.String(d.Id()),
			BranchFilter: aws.String(d.Get("branch_filter").(string)),
			RotateSecret: aws.Bool(false),
		})
	}

	if err != nil {
		return err
	}

	return resourceAwsCodeBuildWebhookRead(d, meta)
}

func resourceAwsCodeBuildWebhookDelete(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).codebuildconn

	_, err := conn.DeleteWebhook(&codebuild.DeleteWebhookInput{
		ProjectName: aws.String(d.Id()),
	})

	if err != nil {
		if isAWSErr(err, codebuild.ErrCodeResourceNotFoundException, "") {
			return nil
		}
		return err
	}

	return nil
}

func flattenAwsCodeBuildWebhookFilterGroups(filterList []*codebuild.WebhookFilter) *schema.Set {
	filterSet := schema.Set{
		F: resourceAwsCodeBuildWebhookFilterHash,
	}

	for _, filters := range filterList {
		filterSet.Add(flattenAwsCodeBuildWebhookFilterData(*filters))
	}
	return &filterSet
}

func resourceAwsCodeBuildWebhookFilterHash(v interface{}) int {
	var buf bytes.Buffer
	m := v.(map[string]interface{})

	buf.WriteString(fmt.Sprintf("%s-", m["type"].(string)))
	buf.WriteString(fmt.Sprintf("%s-", m["pattern"].(string)))
	buf.WriteString(fmt.Sprintf("%q", m["exclude_matched_pattern"]))

	return hashcode.String(buf.String())
}

func flattenAwsCodeBuildWebhookFilterData(filter codebuild.WebhookFilter) map[string]interface{} {
	values := map[string]interface{}{}

	values["type"] = *filter.Type
	values["pattern"] = *filter.Pattern
	values["exclude_matched_pattern"] = *filter.ExcludeMatchedPattern

	return values
}
<commit_msg>Add multiple filter_groups<commit_after>package aws

import (
	"bytes"
	"fmt"
	"github.com/hashicorp/terraform/helper/validation"
	"log"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/codebuild"
	"github.com/hashicorp/terraform/helper/hashcode"
	"github.com/hashicorp/terraform/helper/schema"
)

func resourceAwsCodeBuildWebhook() *schema.Resource {
	return &schema.Resource{
		Create: resourceAwsCodeBuildWebhookCreate,
		Read:   resourceAwsCodeBuildWebhookRead,
		Delete: resourceAwsCodeBuildWebhookDelete,
		Update: resourceAwsCodeBuildWebhookUpdate,

		Importer: &schema.ResourceImporter{
			State: schema.ImportStatePassthrough,
		},

		Schema: map[string]*schema.Schema{
			"project_name": {
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},
			"branch_filter": {
				Type:          schema.TypeString,
				Optional:      true,
				ConflictsWith: []string{"filter_group"},
			},
			"filter_group": {
				Type:     schema.TypeSet,
				Optional: true,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"filter": {
							Type:     schema.TypeList,
							Optional: true,
							Elem: &schema.Resource{
								Schema: map[string]*schema.Schema{
									"type": {
										Type:     schema.TypeString,
										Required: true,
										ValidateFunc: validation.StringInSlice([]string{
											codebuild.WebhookFilterTypeEvent,
											codebuild.WebhookFilterTypeActorAccountId,
											codebuild.WebhookFilterTypeBaseRef,
											codebuild.WebhookFilterTypeFilePath,
											codebuild.WebhookFilterTypeHeadRef,
										}, false),
									},
									"exclude_matched_pattern": {
										Type:     schema.TypeBool,
										Optional: true,
										Default:  false,
									},
									"pattern": {
										Type:     schema.TypeString,
										Required: true,
									},
								},
							},
						},
					},
				},
				ConflictsWith: []string{"branch_filter"},
			},
			"payload_url": {
				Type:     schema.TypeString,
				Computed: true,
			},
			"secret": {
				Type:      schema.TypeString,
				Computed:  true,
				Sensitive: true,
			},
			"url": {
				Type:     schema.TypeString,
				Computed: true,
			},
		},
	}
}

func resourceAwsCodeBuildWebhookCreate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).codebuildconn

	input := &codebuild.CreateWebhookInput{
		ProjectName:  aws.String(d.Get("project_name").(string)),
		FilterGroups: expandWebhookFilterGroups(d),
	}

	// The CodeBuild API requires this to be non-empty if defined
	if v, ok := d.GetOk("branch_filter"); ok {
		input.BranchFilter = aws.String(v.(string))
	}

	log.Printf("[DEBUG] Creating CodeBuild Webhook: %s", input)
	resp, err := conn.CreateWebhook(input)

	if err != nil {
		return fmt.Errorf("error creating CodeBuild Webhook: %s", err)
	}

	// Secret is only returned on create, so capture it at the start
	d.Set("secret", resp.Webhook.Secret)
	d.SetId(d.Get("project_name").(string))

	return resourceAwsCodeBuildWebhookRead(d, meta)
}

func expandWebhookFilterGroups(d *schema.ResourceData) [][]*codebuild.WebhookFilter {
	configs := d.Get("filter_group").(*schema.Set).List()

	webhookFilters := make([][]*codebuild.WebhookFilter, 0)

	if len(configs) == 0 {
		return nil
	}

	for _, config := range configs {
		filters := expandWebhookFilterData(config.(map[string]interface{}))
		webhookFilters = append(webhookFilters, filters)
	}

	return webhookFilters
}

func expandWebhookFilterData(data map[string]interface{}) []*codebuild.WebhookFilter {
	filters := make([]*codebuild.WebhookFilter, 0)

	filterConfigs := data["filter"].([]interface{})

	for i, filterConfig := range filterConfigs {
		filter := filterConfig.(map[string]interface{})
		filters = append(filters, &codebuild.WebhookFilter{
			Type:                  aws.String(filter["type"].(string)),
			ExcludeMatchedPattern: aws.Bool(filter["exclude_matched_pattern"].(bool)),
		})
		if v := filter["pattern"]; v != nil {
			filters[i].Pattern = aws.String(v.(string))
		}
	}

	return filters
}

func resourceAwsCodeBuildWebhookRead(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).codebuildconn

	resp, err := conn.BatchGetProjects(&codebuild.BatchGetProjectsInput{
		Names: []*string{
			aws.String(d.Id()),
		},
	})

	if err != nil {
		return err
	}

	if len(resp.Projects) == 0 {
		log.Printf("[WARN] CodeBuild Project %q not found, removing from state", d.Id())
		d.SetId("")
		return nil
	}

	project := resp.Projects[0]

	if project.Webhook == nil {
		log.Printf("[WARN] CodeBuild Project %q webhook not found, removing from state", d.Id())
		d.SetId("")
		return nil
	}

	d.Set("branch_filter", project.Webhook.BranchFilter)
	d.Set("filter_group", flattenAwsCodeBuildWebhookFilterGroups(project.Webhook.FilterGroups))
	d.Set("payload_url", project.Webhook.PayloadUrl)
	d.Set("project_name", project.Name)
	d.Set("url", project.Webhook.Url)
	// The secret is never returned after creation, so don't set it here

	return nil
}

func resourceAwsCodeBuildWebhookUpdate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).codebuildconn

	var err error
	filterGroups := expandWebhookFilterGroups(d)

	if len(filterGroups) >= 1 {
		_, err = conn.UpdateWebhook(&codebuild.UpdateWebhookInput{
			ProjectName:  aws.String(d.Id()),
			FilterGroups: filterGroups,
			RotateSecret: aws.Bool(false),
		})
	} else {
		_, err = conn.UpdateWebhook(&codebuild.UpdateWebhookInput{
			ProjectName:  aws.String(d.Id()),
			BranchFilter: aws.String(d.Get("branch_filter").(string)),
			RotateSecret: aws.Bool(false),
		})
	}

	if err != nil {
		return err
	}

	return resourceAwsCodeBuildWebhookRead(d, meta)
}

func resourceAwsCodeBuildWebhookDelete(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).codebuildconn

	_, err := conn.DeleteWebhook(&codebuild.DeleteWebhookInput{
		ProjectName: aws.String(d.Id()),
	})

	if err != nil {
		if isAWSErr(err, codebuild.ErrCodeResourceNotFoundException, "") {
			return nil
		}
		return err
	}

	return nil
}

func flattenAwsCodeBuildWebhookFilterGroups(filterList [][]*codebuild.WebhookFilter) *schema.Set {
	filterSet := schema.Set{
		F: resourceAwsCodeBuildWebhookFilterHash,
	}

	for _, filters := range filterList {
		filterSet.Add(flattenAwsCodeBuildWebhookFilterData(filters))
	}
	return &filterSet
}

func resourceAwsCodeBuildWebhookFilterHash(v interface{}) int {
	var buf bytes.Buffer
	m := v.([]map[string]interface{})

	for _, f := range m {
		buf.WriteString(fmt.Sprintf("%s-", f["type"].(*string)))
		buf.WriteString(fmt.Sprintf("%s-", f["pattern"].(*string)))
		buf.WriteString(fmt.Sprintf("%q", f["exclude_matched_pattern"]))
	}

	return hashcode.String(buf.String())
}

func flattenAwsCodeBuildWebhookFilterData(filter []*codebuild.WebhookFilter) []map[string]interface{} {
	values := make([]map[string]interface{}, 0)

	for _, f := range filter {
		values = append(values, map[string]interface{}{
			"type":                    f.Type,
			"pattern":                 f.Pattern,
			"exclude_matched_pattern": f.ExcludeMatchedPattern,
		})
	}

	return values
}
<|endoftext|>


<commit_before>package engine

import (
	"fmt"
	"github.com/keybase/client/go/libkb"
	"io"
)

type PGPSignEngine struct {
	arg *PGPSignArg
	libkb.Contextified
}

type PGPSignArg struct {
	Sink     io.WriteCloser
	Source   io.ReadCloser
	Binary   bool
	KeyQuery string
}

func (p *PGPSignEngine) GetPrereqs() EnginePrereqs {
	return EnginePrereqs{
		Session: true,
	}
}

func (p *PGPSignEngine) Name() string {
	return "PGPSignEngine"
}

func (p *PGPSignEngine) RequiredUIs() []libkb.UIKind {
	return []libkb.UIKind{
		libkb.SecretUIKind,
	}
}

func (s *PGPSignEngine) SubConsumers() []libkb.UIConsumer {
	return nil
}

func NewPGPSignEngine(arg *PGPSignArg) *PGPSignEngine {
	return &PGPSignEngine{arg: arg}
}

func (p *PGPSignEngine) Run(ctx *Context, args interface{}, reply interface{}) (err error) {
	var key libkb.GenericKey
	var pgp *libkb.PgpKeyBundle
	var ok bool
	var dumpTo io.WriteCloser
	var written int64

	defer func() {
		if dumpTo != nil {
			dumpTo.Close()
		}
	}()

	ska := libkb.SecretKeyArg{
		Reason:   "command-line signature",
		PGPOnly:  true,
		KeyQuery: p.arg.KeyQuery,
		Ui:       ctx.SecretUI,
	}

	key, err = p.G().Keyrings.GetSecretKey(ska)

	if err != nil {
		return
	} else if pgp, ok = key.(*libkb.PgpKeyBundle); !ok {
		err = fmt.Errorf("Can only sign with PGP keys (for now)")
		return
	} else if key == nil {
		err = fmt.Errorf("No secret key available")
		return
	}

	dumpTo, err = libkb.AttachedSignWrapper(p.arg.Sink, *pgp, !p.arg.Binary)
	if err != nil {
		return
	}

	written, err = io.Copy(dumpTo, p.arg.Source)

	if err == nil && written == 0 {
		err = fmt.Errorf("Empty source file, nothing to sign")
	}

	return
}
<commit_msg>Close files in exported stream table for #259.<commit_after>package engine

import (
	"fmt"
	"github.com/keybase/client/go/libkb"
	"io"
)

type PGPSignEngine struct {
	arg *PGPSignArg
	libkb.Contextified
}

type PGPSignArg struct {
	Sink     io.WriteCloser
	Source   io.ReadCloser
	Binary   bool
	KeyQuery string
}

func (p *PGPSignEngine) GetPrereqs() EnginePrereqs {
	return EnginePrereqs{
		Session: true,
	}
}

func (p *PGPSignEngine) Name() string {
	return "PGPSignEngine"
}

func (p *PGPSignEngine) RequiredUIs() []libkb.UIKind {
	return []libkb.UIKind{
		libkb.SecretUIKind,
	}
}

func (s *PGPSignEngine) SubConsumers() []libkb.UIConsumer {
	return nil
}

func NewPGPSignEngine(arg *PGPSignArg) *PGPSignEngine {
	return &PGPSignEngine{arg: arg}
}

func (p *PGPSignEngine) Run(ctx *Context, args interface{}, reply interface{}) (err error) {
	var key libkb.GenericKey
	var pgp *libkb.PgpKeyBundle
	var ok bool
	var dumpTo io.WriteCloser
	var written int64

	defer func() {
		if dumpTo != nil {
			dumpTo.Close()
		}
		p.arg.Sink.Close()
		p.arg.Source.Close()
	}()

	ska := libkb.SecretKeyArg{
		Reason:   "command-line signature",
		PGPOnly:  true,
		KeyQuery: p.arg.KeyQuery,
		Ui:       ctx.SecretUI,
	}

	key, err = p.G().Keyrings.GetSecretKey(ska)

	if err != nil {
		return
	} else if pgp, ok = key.(*libkb.PgpKeyBundle); !ok {
		err = fmt.Errorf("Can only sign with PGP keys (for now)")
		return
	} else if key == nil {
		err = fmt.Errorf("No secret key available")
		return
	}

	dumpTo, err = libkb.AttachedSignWrapper(p.arg.Sink, *pgp, !p.arg.Binary)
	if err != nil {
		return
	}

	written, err = io.Copy(dumpTo, p.arg.Source)

	if err == nil && written == 0 {
		err = fmt.Errorf("Empty source file, nothing to sign")
	}

	return
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	//"github.com/davecgh/go-spew/spew"
	"fmt"
	"os"
	"path/filepath"
	"reflect"
	"strings"

	"github.com/elazarl/hadoophelpers/go/lib/readline"
	"github.com/elazarl/hadoophelpers/go/lib/hadoopconf"
	"github.com/foize/go.sgr"
	"github.com/jessevdk/go-flags"
	"github.com/elazarl/hadoophelpers/go/lib/table"
	//"github.com/wsxiaoys/terminal"
)

type getOpts struct {}

type setOpts struct {}

type envAddOpts struct {
	Append bool `long:"append" default:"false" description:"append value to environment variable"`
}

type envDelOpts struct {}

type envSetOpts struct {}

type envOpts struct {}

func (o getOpts) Execute(args []string) error {
	opt.executed = true
	if opt.completeOpts != nil {
		groups := getmygroups(o, &opt)
		options := getGroupOptions(groups)
		opt.completeOpts = append(options, opt.getConf().Keys()...)
		return nil
	}
	if len(args) == 0 {
		return errors.New("get must have nonzero number arguments")
	}
	t := table.New(4)
	c := opt.getConf()
	keys := []string{}
	for _, key := range c.Keys() {
		for _, arg := range args {
			if ok, _ := filepath.Match(arg, key); ok {
				keys = append(keys, key)
				break
			}
		}
	}
	if opt.UseColors() {
		t.CellConf[0].PadLeft = []byte(sgr.FgGrey)
		t.CellConf[1].PadLeft = []byte(sgr.FgCyan)
		t.CellConf[2].PadLeft = []byte(sgr.FgGrey)
		t.CellConf[3].PadLeft = []byte(sgr.ResetForegroundColor + sgr.Bold)
		t.CellConf[3].PadRight = []byte(sgr.Reset)
	}
	for _, arg := range keys {
		v, src := c.SourceGet(arg)
		if v == "" && src == hadoopconf.NoSource {
			t.Add("", arg, "", "no property")
		} else {
			t.Add(filepath.Base(src.Source), arg, "=", v)
		}
	}
	fmt.Print(t.String())
	return nil
}

func (o setOpts) Execute(args []string) error {
	opt.executed = true
	if opt.completeOpts != nil {
		options := getGroupOptions(getmygroups(o, &opt))
		for _, v := range opt.getConf().Keys() {
			opt.completeOpts = append(opt.completeOpts, v + "=")
		}
		for _, v := range options {
			opt.completeOpts = append(opt.completeOpts, v + " ")
		}
		readline.SuppressAppend()
		return nil
	}
	if len(args) == 0 {
		return errors.New("get must have nonzero number arguments")
	}
	for _, arg := range args {
		parts := strings.SplitN(arg, "=", 2)
		if len(parts) != 2 {
			return errors.New("set accepts arguments of the form x=y, no '=' in " + arg)
		}
		opt.getConf().SetIfExist(parts[0], parts[1])
	}
	opt.getConf().Save()
	return nil
}

func assignmentTable() *table.Table {
	t := table.New(4)
	if opt.UseColors() {
		t.CellConf[0].PadLeft = []byte(sgr.FgGrey)
		t.CellConf[1].PadLeft = []byte(sgr.FgCyan)
		t.CellConf[2].PadLeft = []byte(sgr.FgGrey)
		t.CellConf[3].PadLeft = []byte(sgr.ResetForegroundColor + sgr.Bold)
		t.CellConf[3].PadRight = []byte(sgr.Reset)
	}
	return t
}

func (o envSetOpts) Execute(args []string) error {
	opt.executed = true
	if opt.completeOpts != nil {
		options := getGroupOptions(getmygroups(o, &opt))
		if len(args) == 0 {
			opt.completeOpts = append(options, opt.getEnv().Keys()...)
		} else {
			if v := opt.getEnv().Get(args[0]); v != nil {
				opt.completeOpts = append(opt.completeOpts, v.Val)
			}
		}
		return nil
	}
	if len(args) == 0 {
		return errors.New("get must have nonzero number arguments")
	}
	v := opt.getEnv().Get(args[0])
	if v == nil {
		fmt.Println("No such variable", v)
	}
	t := assignmentTable()
	t.Add(filepath.Base(v.Source), v.Name, "was", v.Val)
	v.Val = strings.Join(args[1:], " ")
	t.Add("", "", "now", v.Val)
	if err := opt.getEnv().Save(); err != nil {
		return err
	}
	fmt.Print(t.String())
	return nil
}

func (o envAddOpts) Execute(args []string) error {
	opt.executed = true
	if opt.completeOpts != nil {
		options := getGroupOptions(getmygroups(o, &opt))
		if len(args) == 0 {
			opt.completeOpts = append(options, opt.getEnv().Keys()...)
		}
		return nil
	}
	if len(args) == 0 {
		return errors.New("get must have nonzero number arguments")
	}
	v := opt.getEnv().Get(args[0])
	if v == nil {
		fmt.Println("No such variable", v)
	}
	t := assignmentTable()
	t.Add(filepath.Base(v.Source), v.Name, "was", v.Val)
	v.Prepend(strings.Join(args[1:], " "))
	t.Add("", "", "now", v.Val)
	if err := opt.getEnv().Save(); err != nil {
		return err
	}
	fmt.Print(t.String())
	return nil
}

func (o envDelOpts) Execute(args []string) error {
	opt.executed = true
	if opt.completeOpts != nil {
		options := getGroupOptions(getmygroups(o, &opt))
		if len(args) == 0 {
			opt.completeOpts = append(options, opt.getEnv().Keys()...)
		} else {
			if v := opt.getEnv().Get(args[0]); v != nil {
				opt.completeOpts = append(opt.completeOpts, parseCommandLine(v.Val)...)
			}
		}
		return nil
	}
	if len(args) == 0 {
		return errors.New("get must have nonzero number arguments")
	}
	v := opt.getEnv().Get(args[0])
	if v == nil {
		fmt.Println("No such variable", v)
	}
	t := assignmentTable()
	t.Add(filepath.Base(v.Source), v.Name, "was", v.Val)
	v.Del(strings.Join(args[1:], " "))
	t.Add("", "", "now", v.Val)
	if err := opt.getEnv().Save(); err != nil {
		return err
	}
	fmt.Print(t.String())
	return nil
}

func (o envOpts) Execute(args []string) error {
	opt.executed = true
	if opt.completeOpts != nil {
		options := getGroupOptions(getmygroups(o, &opt))
		opt.completeOpts = append(options, opt.getEnv().Keys()...)
		return nil
	}
	if len(args) == 0 {
		return errors.New("get must have nonzero number arguments")
	}
	t := assignmentTable()
	c := opt.getEnv()
	keys := []string{}
	for _, key := range c.Keys() {
		for _, arg := range args {
			if ok, _ := filepath.Match(arg, key); ok {
				keys = append(keys, key)
				break
			}
		}
	}
	for _, arg := range keys {
		v := c.Get(arg)
		if v == nil {
			t.Add("", arg, "", "no property")
		} else {
			t.Add(filepath.Base(v.Source), arg, "=", v.Val)
		}
	}
	fmt.Print(t.String())
	return nil
}

func (o *gOpts) UseColors() bool {
	if o.Color == "auto" {
		return IsTerminal(os.Stdout.Fd())
	}
	return o.Color == "true" || o.Color == "t" || o.Color == "1"
}

type gOpts struct {
	Get getOpts `command:"get"`
	Set setOpts `command:"set"`
	SetEnv envSetOpts `command:"envset"`
	AddEnv envAddOpts `command:"envadd"`
	DelEnv envDelOpts `command:"envdel"`
	Env envOpts `command:"env"`
	Verbose bool `short:"v" long:"verbose" default:"true" description:"Show verbose debug information"`
	Color string `long:"color" description:"use colors on output" default:"auto"`
	ConfPath string `short:"c" long:"conf" description:"Set hadoop configuration dir"`
	conf *hadoopconf.HadoopConf
	env hadoopconf.Envs
	executed bool
	// set this to []string{} if you want command line options to autocomplete instead of executing themselves
	completeOpts []string
	parser *flags.Parser
}

func (opt *gOpts) setConfPath() {
	var p = "."
	if opt.ConfPath != "" {
		p = opt.ConfPath
	} else if os.Getenv("HADOOP_CONF") != "" {
		p = os.Getenv("HADOOP_CONF")
	}
	opt.ConfPath = p
}

func (opt *gOpts) getEnv() hadoopconf.Envs {
	if opt.env != nil {
		return opt.env
	}
	var err error
	opt.setConfPath()
	opt.env, err = hadoopconf.NewEnv(opt.ConfPath)
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
	return opt.env
}

func (opt *gOpts) getConf() *hadoopconf.HadoopConf {
	if opt.conf != nil {
		return opt.conf
	}
	var err error
	var p = "."
	if opt.ConfPath != "" {
		p = opt.ConfPath
	} else if os.Getenv("HADOOP_CONF") != "" {
		p = os.Getenv("HADOOP_CONF")
	}
	opt.conf, err = hadoopconf.New(p)
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
	return opt.conf
}

var opt gOpts
var conf *hadoopconf.HadoopConf

// bash-like command line parser, splits string to arguments
func parseCommandLine(line string) []string {
	args := []string{}
	type State int
	const (
		REGULAR State = iota
		IN_DQUOTE
		IN_QUOTE
		IN_ESCAPE
	)
	state := REGULAR
	escapePrevState := REGULAR
	next := []byte{}
	seenQuote := false
	for _, r := range []byte(line) {
		switch state {
		case REGULAR:
			if r == ' ' || r == '\t' {
				if len(next) > 0 || seenQuote {
					args = append(args, string(next))
				}
				next = nil
				seenQuote = false
			} else if r == '\'' {
				seenQuote = true
				state = IN_QUOTE
			} else if r == '"' {
				seenQuote = true
				state = IN_DQUOTE
			} else if r == '\\' {
				state = IN_ESCAPE
				escapePrevState = REGULAR
			} else {
				next = append(next, r)
			}
		case IN_QUOTE:
			if r == '\'' {
				state = REGULAR
			} else if r == '\\' {
				state = IN_ESCAPE
				escapePrevState = IN_QUOTE
			} else {
				next = append(next, r)
			}
		case IN_DQUOTE:
			if r == '"' {
				state = REGULAR
			} else if r == '\\' {
				state = IN_ESCAPE
				escapePrevState = IN_DQUOTE
			} else {
				next = append(next, r)
			}
		case IN_ESCAPE:
			next = append(next, r)
			state = escapePrevState
		}
	}
	if len(next) > 0 || seenQuote {
		args = append(args, string(next))
	}
	return args
}

func main() {
	parser := flags.NewParser(&opt, flags.HelpFlag | flags.PassDoubleDash | flags.IgnoreUnknown)
	if _, err := parser.ParseArgs(os.Args[1:]); err != nil && opt.executed {
		fmt.Println("dead:", err)
		os.Exit(1)
	}
	opt.getConf() // make sure we have correct conf
	if !opt.executed {
		if !IsTerminal(os.Stdout.Fd()) {
			fmt.Println("terminal not recognized or not supported (windows)")
			return
		}
		readline.Completer = func (line string, start, end int) (string, []string) {
			completionparser := flags.NewParser(&opt, flags.HelpFlag | flags.PassDoubleDash | flags.IgnoreUnknown)
			opt.executed = false
			args := parseCommandLine(line[:end])
			if line[end-1] == ' ' || line[end-1] == '\t' {
				return "", Complete(completionparser, args, "")
			}
			return "", Complete(completionparser, args[:len(args)-1], args[len(args)-1])
		}
		for {
			str, ok := readline.Readline("hadoopconf> ")
			if !ok {
				break
			}
			opt.completeOpts = nil
			args := parseCommandLine(str)
			if args, err := parser.ParseArgs(args); err != nil {
				fmt.Println("error:", err)
			} else if len(args) > 0 {
				fmt.Println("excessive arguments:", args)
			}
		}
	}
}

// given
// type T struct {
//     foo Foo `command:"moo"`
//     bar Bar `command:"maa"`
// }
// getField(Foo{}, T{}) == "moo"
// getField(Bar{}, T{}) == "maa"
// getField(Baz{}, T{}) panics
func getField(typ interface{}, strct interface{}) string {
	v := reflect.TypeOf(strct)
	if v.Kind() == reflect.Ptr {
		v = v.Elem()
	}
	for i := 0; i < v.NumField(); i++ {
		f := v.Field(i)
		if f.Type.AssignableTo(reflect.TypeOf(typ)) {
			return reflect.StructTag(f.Tag).Get("command")
		}
	}
	panic("cannot find type in struct")
}

func getmygroups(o, strct interface{}) *flags.Group {
	field := getField(o, strct)
	for _, group := range opt.parser.Groups {
		if mygroup, ok := group.Commands[field]; ok {
			return mygroup
		}
	}
	panic("my field not avail")
}
<commit_msg>bug: autocomplete empty line<commit_after>package main

import (
	"errors"
	//"github.com/davecgh/go-spew/spew"
	"fmt"
	"os"
	"path/filepath"
	"reflect"
	"strings"

	"github.com/elazarl/hadoophelpers/go/lib/readline"
	"github.com/elazarl/hadoophelpers/go/lib/hadoopconf"
	"github.com/foize/go.sgr"
	"github.com/jessevdk/go-flags"
	"github.com/elazarl/hadoophelpers/go/lib/table"
	//"github.com/wsxiaoys/terminal"
)

type getOpts struct {}

type setOpts struct {}

type envAddOpts struct {
	Append bool `long:"append" default:"false" description:"append value to environment variable"`
}

type envDelOpts struct {}

type envSetOpts struct {}

type envOpts struct {}

func (o getOpts) Execute(args []string) error {
	opt.executed = true
	if opt.completeOpts != nil {
		groups := getmygroups(o, &opt)
		options := getGroupOptions(groups)
		opt.completeOpts = append(options, opt.getConf().Keys()...)
		return nil
	}
	if len(args) == 0 {
		return errors.New("get must have nonzero number arguments")
	}
	t := table.New(4)
	c := opt.getConf()
	keys := []string{}
	for _, key := range c.Keys() {
		for _, arg := range args {
			if ok, _ := filepath.Match(arg, key); ok {
				keys = append(keys, key)
				break
			}
		}
	}
	if opt.UseColors() {
		t.CellConf[0].PadLeft = []byte(sgr.FgGrey)
		t.CellConf[1].PadLeft = []byte(sgr.FgCyan)
		t.CellConf[2].PadLeft = []byte(sgr.FgGrey)
		t.CellConf[3].PadLeft = []byte(sgr.ResetForegroundColor + sgr.Bold)
		t.CellConf[3].PadRight = []byte(sgr.Reset)
	}
	for _, arg := range keys {
		v, src := c.SourceGet(arg)
		if v == "" && src == hadoopconf.NoSource {
			t.Add("", arg, "", "no property")
		} else {
			t.Add(filepath.Base(src.Source), arg, "=", v)
		}
	}
	fmt.Print(t.String())
	return nil
}

func (o setOpts) Execute(args []string) error {
	opt.executed = true
	if opt.completeOpts != nil {
		options := getGroupOptions(getmygroups(o, &opt))
		for _, v := range opt.getConf().Keys() {
			opt.completeOpts = append(opt.completeOpts, v + "=")
		}
		for _, v := range options {
			opt.completeOpts = append(opt.completeOpts, v + " ")
		}
		readline.SuppressAppend()
		return nil
	}
	if len(args) == 0 {
		return errors.New("get must have nonzero number arguments")
	}
	for _, arg := range args {
		parts := strings.SplitN(arg, "=", 2)
		if len(parts) != 2 {
			return errors.New("set accepts arguments of the form x=y, no '=' in " + arg)
		}
		opt.getConf().SetIfExist(parts[0], parts[1])
	}
	opt.getConf().Save()
	return nil
}

func assignmentTable() *table.Table {
	t := table.New(4)
	if opt.UseColors() {
		t.CellConf[0].PadLeft = []byte(sgr.FgGrey)
		t.CellConf[1].PadLeft = []byte(sgr.FgCyan)
		t.CellConf[2].PadLeft = []byte(sgr.FgGrey)
		t.CellConf[3].PadLeft = []byte(sgr.ResetForegroundColor + sgr.Bold)
		t.CellConf[3].PadRight = []byte(sgr.Reset)
	}
	return t
}

func (o envSetOpts) Execute(args []string) error {
	opt.executed = true
	if opt.completeOpts != nil {
		options := getGroupOptions(getmygroups(o, &opt))
		if len(args) == 0 {
			opt.completeOpts = append(options, opt.getEnv().Keys()...)
		} else {
			if v := opt.getEnv().Get(args[0]); v != nil {
				opt.completeOpts = append(opt.completeOpts, v.Val)
			}
		}
		return nil
	}
	if len(args) == 0 {
		return errors.New("get must have nonzero number arguments")
	}
	v := opt.getEnv().Get(args[0])
	if v == nil {
		fmt.Println("No such variable", v)
	}
	t := assignmentTable()
	t.Add(filepath.Base(v.Source), v.Name, "was", v.Val)
	v.Val = strings.Join(args[1:], " ")
	t.Add("", "", "now", v.Val)
	if err := opt.getEnv().Save(); err != nil {
		return err
	}
	fmt.Print(t.String())
	return nil
}

func (o envAddOpts) Execute(args []string) error {
	opt.executed = true
	if opt.completeOpts != nil {
		options := getGroupOptions(getmygroups(o, &opt))
		if len(args) == 0 {
			opt.completeOpts = append(options, opt.getEnv().Keys()...)
		}
		return nil
	}
	if len(args) == 0 {
		return errors.New("get must have nonzero number arguments")
	}
	v := opt.getEnv().Get(args[0])
	if v == nil {
		fmt.Println("No such variable", v)
	}
	t := assignmentTable()
	t.Add(filepath.Base(v.Source), v.Name, "was", v.Val)
	v.Prepend(strings.Join(args[1:], " "))
	t.Add("", "", "now", v.Val)
	if err := opt.getEnv().Save(); err != nil {
		return err
	}
	fmt.Print(t.String())
	return nil
}

func (o envDelOpts) Execute(args []string) error {
	opt.executed = true
	if opt.completeOpts != nil {
		options := getGroupOptions(getmygroups(o, &opt))
		if len(args) == 0 {
			opt.completeOpts = append(options, opt.getEnv().Keys()...)
		} else {
			if v := opt.getEnv().Get(args[0]); v != nil {
				opt.completeOpts = append(opt.completeOpts, parseCommandLine(v.Val)...)
			}
		}
		return nil
	}
	if len(args) == 0 {
		return errors.New("get must have nonzero number arguments")
	}
	v := opt.getEnv().Get(args[0])
	if v == nil {
		fmt.Println("No such variable", v)
	}
	t := assignmentTable()
	t.Add(filepath.Base(v.Source), v.Name, "was", v.Val)
	v.Del(strings.Join(args[1:], " "))
	t.Add("", "", "now", v.Val)
	if err := opt.getEnv().Save(); err != nil {
		return err
	}
	fmt.Print(t.String())
	return nil
}

func (o envOpts) Execute(args []string) error {
	opt.executed = true
	if opt.completeOpts != nil {
		options := getGroupOptions(getmygroups(o, &opt))
		opt.completeOpts = append(options, opt.getEnv().Keys()...)
		return nil
	}
	if len(args) == 0 {
		return errors.New("get must have nonzero number arguments")
	}
	t := assignmentTable()
	c := opt.getEnv()
	keys := []string{}
	for _, key := range c.Keys() {
		for _, arg := range args {
			if ok, _ := filepath.Match(arg, key); ok {
				keys = append(keys, key)
				break
			}
		}
	}
	for _, arg := range keys {
		v := c.Get(arg)
		if v == nil {
			t.Add("", arg, "", "no property")
		} else {
			t.Add(filepath.Base(v.Source), arg, "=", v.Val)
		}
	}
	fmt.Print(t.String())
	return nil
}

func (o *gOpts) UseColors() bool {
	if o.Color == "auto" {
		return IsTerminal(os.Stdout.Fd())
	}
	return o.Color == "true" || o.Color == "t" || o.Color == "1"
}

type gOpts struct {
	Get getOpts `command:"get"`
	Set setOpts `command:"set"`
	SetEnv envSetOpts `command:"envset"`
	AddEnv envAddOpts `command:"envadd"`
	DelEnv envDelOpts `command:"envdel"`
	Env envOpts `command:"env"`
	Verbose bool `short:"v" long:"verbose" default:"true" description:"Show verbose debug information"`
	Color string `long:"color" description:"use colors on output" default:"auto"`
	ConfPath string `short:"c" long:"conf" description:"Set hadoop configuration dir"`
	conf *hadoopconf.HadoopConf
	env hadoopconf.Envs
	executed bool
	// set this to []string{} if you want command line options to autocomplete instead of executing themselves
	completeOpts []string
	parser *flags.Parser
}

func (opt *gOpts) setConfPath() {
	var p = "."
	if opt.ConfPath != "" {
		p = opt.ConfPath
	} else if os.Getenv("HADOOP_CONF") != "" {
		p = os.Getenv("HADOOP_CONF")
	}
	opt.ConfPath = p
}

func (opt *gOpts) getEnv() hadoopconf.Envs {
	if opt.env != nil {
		return opt.env
	}
	var err error
	opt.setConfPath()
	opt.env, err = hadoopconf.NewEnv(opt.ConfPath)
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
	return opt.env
}

func (opt *gOpts) getConf() *hadoopconf.HadoopConf {
	if opt.conf != nil {
		return opt.conf
	}
	var err error
	var p = "."
	if opt.ConfPath != "" {
		p = opt.ConfPath
	} else if os.Getenv("HADOOP_CONF") != "" {
		p = os.Getenv("HADOOP_CONF")
	}
	opt.conf, err = hadoopconf.New(p)
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
	return opt.conf
}

var opt gOpts
var conf *hadoopconf.HadoopConf

// bash-like command line parser, splits string to arguments
func parseCommandLine(line string) []string {
	args := []string{}
	type State int
	const (
		REGULAR State = iota
		IN_DQUOTE
		IN_QUOTE
		IN_ESCAPE
	)
	state := REGULAR
	escapePrevState := REGULAR
	next := []byte{}
	seenQuote := false
	for _, r := range []byte(line) {
		switch state {
		case REGULAR:
			if r == ' ' || r == '\t' {
				if len(next) > 0 || seenQuote {
					args = append(args, string(next))
				}
				next = nil
				seenQuote = false
			} else if r == '\'' {
				seenQuote = true
				state = IN_QUOTE
			} else if r == '"' {
				seenQuote = true
				state = IN_DQUOTE
			} else if r == '\\' {
				state = IN_ESCAPE
				escapePrevState = REGULAR
			} else {
				next = append(next, r)
			}
		case IN_QUOTE:
			if r == '\'' {
				state = REGULAR
			} else if r == '\\' {
				state = IN_ESCAPE
				escapePrevState = IN_QUOTE
			} else {
				next = append(next, r)
			}
		case IN_DQUOTE:
			if r == '"' {
				state = REGULAR
			} else if r == '\\' {
				state = IN_ESCAPE
				escapePrevState = IN_DQUOTE
			} else {
				next = append(next, r)
			}
		case IN_ESCAPE:
			next = append(next, r)
			state = escapePrevState
		}
	}
	if len(next) > 0 || seenQuote {
		args = append(args, string(next))
	}
	return args
}

func main() {
	parser := flags.NewParser(&opt, flags.HelpFlag | flags.PassDoubleDash | flags.IgnoreUnknown)
	if _, err := parser.ParseArgs(os.Args[1:]); err != nil && opt.executed {
		fmt.Println("dead:", err)
		os.Exit(1)
	}
	opt.getConf() // make sure we have correct conf
	if !opt.executed {
		if !IsTerminal(os.Stdout.Fd()) {
			fmt.Println("terminal not recognized or not supported (windows)")
			return
		}
		readline.Completer = func (line string, start, end int) (string, []string) {
			completionparser := flags.NewParser(&opt, flags.HelpFlag | flags.PassDoubleDash | flags.IgnoreUnknown)
			opt.executed = false
			args := parseCommandLine(line[:end])
			if len(line) == 0 || line[end-1] == ' ' || line[end-1] == '\t' {
				return "", Complete(completionparser, args, "")
			}
			return "", Complete(completionparser, args[:len(args)-1], args[len(args)-1])
		}
		for {
			str, ok := readline.Readline("hadoopconf> ")
			if !ok {
				break
			}
			opt.completeOpts = nil
			args := parseCommandLine(str)
			if args, err := parser.ParseArgs(args); err != nil {
				fmt.Println("error:", err)
			} else if len(args) > 0 {
				fmt.Println("excessive arguments:", args)
			}
		}
	}
}

// given
// type T struct {
//     foo Foo `command:"moo"`
//     bar Bar `command:"maa"`
// }
// getField(Foo{}, T{}) == "moo"
// getField(Bar{}, T{}) == "maa"
// getField(Baz{}, T{}) panics
func getField(typ interface{}, strct interface{}) string {
	v := reflect.TypeOf(strct)
	if v.Kind() == reflect.Ptr {
		v = v.Elem()
	}
	for i := 0; i < v.NumField(); i++ {
		f := v.Field(i)
		if f.Type.AssignableTo(reflect.TypeOf(typ)) {
			return reflect.StructTag(f.Tag).Get("command")
		}
	}
	panic("cannot find type in struct")
}

func getmygroups(o, strct interface{}) *flags.Group {
	field := getField(o, strct)
	for _, group := range opt.parser.Groups {
		if mygroup, ok := group.Commands[field]; ok {
			return mygroup
		}
	}
	panic("my field not avail")
}
<|endoftext|>


<commit_before>/*
Copyright 2017 Google Inc.

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 agreedto 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 worker

import (
	"fmt"

	"vitess.io/vitess/go/vt/vterrors"

	"golang.org/x/net/context"

	"vitess.io/vitess/go/sqlescape"
	"vitess.io/vitess/go/sqltypes"
	"vitess.io/vitess/go/vt/topo/topoproto"
	"vitess.io/vitess/go/vt/wrangler"

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

var (
	completeChunk       = chunk{sqltypes.NULL, sqltypes.NULL, 1, 1}
	singleCompleteChunk = []chunk{completeChunk}
)

// chunk holds the information which subset of the table should be worked on.
// The subset is the range of rows in the range [start, end) where start and end
// both refer to the first column of the primary key.
// If the column is not numeric, both start and end will be sqltypes.NULL.
type chunk struct {
	start sqltypes.Value
	end   sqltypes.Value
	// number records the position of this chunk among all "total" chunks.
	// The lowest value is 1.
	number int
	// total is the total number of chunks this chunk belongs to.
	total int
}

// String returns a human-readable presentation of the chunk range.
func (c chunk) String() string {
	// Pad the chunk number such that all log messages align nicely.
	digits := digits(c.total)
	return fmt.Sprintf("%*d/%d", digits, c.number, c.total)
}

func digits(i int) int {
	digits := 1
	for {
		i /= 10
		if i == 0 {
			break
		}
		digits++
	}
	return digits
}

// generateChunks returns an array of chunks to use for splitting up a table
// into multiple data chunks. It only works for tables with a primary key
// whose first column is a numeric type.
func generateChunks(ctx context.Context, wr *wrangler.Wrangler, tablet *topodatapb.Tablet, td *tabletmanagerdatapb.TableDefinition, chunkCount, minRowsPerChunk int) ([]chunk, error) {
	if len(td.PrimaryKeyColumns) == 0 {
		// No explicit primary key. Cannot chunk the rows then.
		wr.Logger().Infof("table=%v: Not splitting the table into multiple chunks because it has no primary key columns. This will reduce the performance of the clone.", td.Name)
		return singleCompleteChunk, nil
	}
	if td.RowCount < 2*uint64(minRowsPerChunk) {
		// The automatic adjustment of "chunkCount" based on "minRowsPerChunk"
		// below would set "chunkCount" to less than 2 i.e. 1 or 0 chunks.
		// In practice in this case there should be exactly one chunk.
		// Return early in this case and notice the user about this.
		wr.Logger().Infof("table=%v: Not splitting the table into multiple chunks because it has only %d rows.", td.Name, td.RowCount)
		return singleCompleteChunk, nil
	}
	if chunkCount == 1 {
		return singleCompleteChunk, nil
	}

	// Get the MIN and MAX of the leading column of the primary key.
	query := fmt.Sprintf("SELECT MIN(%v), MAX(%v) FROM %v.%v", sqlescape.EscapeID(td.PrimaryKeyColumns[0]), sqlescape.EscapeID(td.PrimaryKeyColumns[0]), sqlescape.EscapeID(topoproto.TabletDbName(tablet)), sqlescape.EscapeID(td.Name))
	shortCtx, cancel := context.WithTimeout(ctx, *remoteActionsTimeout)
	qr, err := wr.TabletManagerClient().ExecuteFetchAsApp(shortCtx, tablet, true, []byte(query), 1)
	cancel()
	if err != nil {
		return nil, vterrors.Wrapf(err, "tablet: %v, table: %v: cannot determine MIN and MAX of the first primary key column. ExecuteFetchAsApp", topoproto.TabletAliasString(tablet.Alias), td.Name, err)
	}
	if len(qr.Rows) != 1 {
		return nil, fmt.Errorf("tablet: %v, table: %v: cannot determine MIN and MAX of the first primary key column. Zero rows were returned", topoproto.TabletAliasString(tablet.Alias), td.Name)
	}

	result := sqltypes.Proto3ToResult(qr)
	min, _ := sqltypes.ToNative(result.Rows[0][0])
	max, _ := sqltypes.ToNative(result.Rows[0][1])

	if min == nil || max == nil {
		wr.Logger().Infof("table=%v: Not splitting the table into multiple chunks, min or max is NULL: %v", td.Name, qr.Rows[0])
		return singleCompleteChunk, nil
	}

	// Determine the average number of rows per chunk for the given chunkCount.
	avgRowsPerChunk := td.RowCount / uint64(chunkCount)
	if avgRowsPerChunk < uint64(minRowsPerChunk) {
		// Reduce the chunkCount to fulfill minRowsPerChunk.
		newChunkCount := td.RowCount / uint64(minRowsPerChunk)
		wr.Logger().Infof("table=%v: Reducing the number of chunks from the default %d to %d to make sure that each chunk has at least %d rows.", td.Name, chunkCount, newChunkCount, minRowsPerChunk)
		chunkCount = int(newChunkCount)
	}

	// TODO(mberlin): Write a unit test for this part of the function.
	var interval interface{}
	chunks := make([]chunk, chunkCount)
	switch min := min.(type) {
	case int64:
		max := max.(int64)
		interval = (max - min) / int64(chunkCount)
		if interval == 0 {
			wr.Logger().Infof("table=%v: Not splitting the table into multiple chunks, interval=0: %v to %v", td.Name, min, max)
			return singleCompleteChunk, nil
		}
	case uint64:
		max := max.(uint64)
		interval = (max - min) / uint64(chunkCount)
		if interval == 0 {
			wr.Logger().Infof("table=%v: Not splitting the table into multiple chunks, interval=0: %v to %v", td.Name, min, max)
			return singleCompleteChunk, nil
		}
	case float64:
		max := max.(float64)
		interval = (max - min) / float64(chunkCount)
		if interval == 0 {
			wr.Logger().Infof("table=%v: Not splitting the table into multiple chunks, interval=0: %v to %v", td.Name, min, max)
			return singleCompleteChunk, nil
		}
	default:
		wr.Logger().Infof("table=%v: Not splitting the table into multiple chunks, primary key not numeric.", td.Name)
		return singleCompleteChunk, nil
	}

	// Create chunks.
	start := min
	for i := 0; i < chunkCount; i++ {
		end := add(start, interval)
		chunk, err := toChunk(start, end, i+1, chunkCount)
		if err != nil {
			return nil, vterrors.Wrapf(err, "tablet: %v, table: %v", topoproto.TabletAliasString(tablet.Alias), td.Name)
		}
		chunks[i] = chunk
		start = end
	}

	// Clear out the MIN and MAX on the first and last chunk respectively
	// because other shards might have smaller or higher values than the one we
	// looked at.
	chunks[0].start = sqltypes.NULL
	chunks[chunkCount-1].end = sqltypes.NULL
	return chunks, nil
}

func add(start, interval interface{}) interface{} {
	switch start := start.(type) {
	case int64:
		return start + interval.(int64)
	case uint64:
		return start + interval.(uint64)
	case float64:
		return start + interval.(float64)
	default:
		panic(fmt.Sprintf("unsupported type %T for interval start: %v", start, start))
	}
}

func toChunk(start, end interface{}, number, total int) (chunk, error) {
	startValue, err := sqltypes.InterfaceToValue(start)
	if err != nil {
		return chunk{}, vterrors.Wrapf(err, "failed to convert calculated start value (%v) into internal sqltypes.Value", start)
	}
	endValue, err := sqltypes.InterfaceToValue(end)
	if err != nil {
		return chunk{}, vterrors.Wrapf(err, "failed to convert calculated end value (%v) into internal sqltypes.Value", end)
	}
	return chunk{startValue, endValue, number, total}, nil
}
<commit_msg>Fix govet error<commit_after>/*
Copyright 2017 Google Inc.

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 agreedto 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 worker

import (
	"fmt"

	"vitess.io/vitess/go/vt/vterrors"

	"golang.org/x/net/context"

	"vitess.io/vitess/go/sqlescape"
	"vitess.io/vitess/go/sqltypes"
	"vitess.io/vitess/go/vt/topo/topoproto"
	"vitess.io/vitess/go/vt/wrangler"

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

var (
	completeChunk       = chunk{sqltypes.NULL, sqltypes.NULL, 1, 1}
	singleCompleteChunk = []chunk{completeChunk}
)

// chunk holds the information which subset of the table should be worked on.
// The subset is the range of rows in the range [start, end) where start and end
// both refer to the first column of the primary key.
// If the column is not numeric, both start and end will be sqltypes.NULL.
type chunk struct {
	start sqltypes.Value
	end   sqltypes.Value
	// number records the position of this chunk among all "total" chunks.
	// The lowest value is 1.
	number int
	// total is the total number of chunks this chunk belongs to.
	total int
}

// String returns a human-readable presentation of the chunk range.
func (c chunk) String() string {
	// Pad the chunk number such that all log messages align nicely.
	digits := digits(c.total)
	return fmt.Sprintf("%*d/%d", digits, c.number, c.total)
}

func digits(i int) int {
	digits := 1
	for {
		i /= 10
		if i == 0 {
			break
		}
		digits++
	}
	return digits
}

// generateChunks returns an array of chunks to use for splitting up a table
// into multiple data chunks. It only works for tables with a primary key
// whose first column is a numeric type.
func generateChunks(ctx context.Context, wr *wrangler.Wrangler, tablet *topodatapb.Tablet, td *tabletmanagerdatapb.TableDefinition, chunkCount, minRowsPerChunk int) ([]chunk, error) {
	if len(td.PrimaryKeyColumns) == 0 {
		// No explicit primary key. Cannot chunk the rows then.
		wr.Logger().Infof("table=%v: Not splitting the table into multiple chunks because it has no primary key columns. This will reduce the performance of the clone.", td.Name)
		return singleCompleteChunk, nil
	}
	if td.RowCount < 2*uint64(minRowsPerChunk) {
		// The automatic adjustment of "chunkCount" based on "minRowsPerChunk"
		// below would set "chunkCount" to less than 2 i.e. 1 or 0 chunks.
		// In practice in this case there should be exactly one chunk.
		// Return early in this case and notice the user about this.
		wr.Logger().Infof("table=%v: Not splitting the table into multiple chunks because it has only %d rows.", td.Name, td.RowCount)
		return singleCompleteChunk, nil
	}
	if chunkCount == 1 {
		return singleCompleteChunk, nil
	}

	// Get the MIN and MAX of the leading column of the primary key.
	query := fmt.Sprintf("SELECT MIN(%v), MAX(%v) FROM %v.%v", sqlescape.EscapeID(td.PrimaryKeyColumns[0]), sqlescape.EscapeID(td.PrimaryKeyColumns[0]), sqlescape.EscapeID(topoproto.TabletDbName(tablet)), sqlescape.EscapeID(td.Name))
	shortCtx, cancel := context.WithTimeout(ctx, *remoteActionsTimeout)
	qr, err := wr.TabletManagerClient().ExecuteFetchAsApp(shortCtx, tablet, true, []byte(query), 1)
	cancel()
	if err != nil {
		return nil, vterrors.Wrapf(err, "tablet: %v, table: %v: cannot determine MIN and MAX of the first primary key column. ExecuteFetchAsApp", topoproto.TabletAliasString(tablet.Alias), td.Name)
	}
	if len(qr.Rows) != 1 {
		return nil, fmt.Errorf("tablet: %v, table: %v: cannot determine MIN and MAX of the first primary key column. Zero rows were returned", topoproto.TabletAliasString(tablet.Alias), td.Name)
	}

	result := sqltypes.Proto3ToResult(qr)
	min, _ := sqltypes.ToNative(result.Rows[0][0])
	max, _ := sqltypes.ToNative(result.Rows[0][1])

	if min == nil || max == nil {
		wr.Logger().Infof("table=%v: Not splitting the table into multiple chunks, min or max is NULL: %v", td.Name, qr.Rows[0])
		return singleCompleteChunk, nil
	}

	// Determine the average number of rows per chunk for the given chunkCount.
	avgRowsPerChunk := td.RowCount / uint64(chunkCount)
	if avgRowsPerChunk < uint64(minRowsPerChunk) {
		// Reduce the chunkCount to fulfill minRowsPerChunk.
		newChunkCount := td.RowCount / uint64(minRowsPerChunk)
		wr.Logger().Infof("table=%v: Reducing the number of chunks from the default %d to %d to make sure that each chunk has at least %d rows.", td.Name, chunkCount, newChunkCount, minRowsPerChunk)
		chunkCount = int(newChunkCount)
	}

	// TODO(mberlin): Write a unit test for this part of the function.
	var interval interface{}
	chunks := make([]chunk, chunkCount)
	switch min := min.(type) {
	case int64:
		max := max.(int64)
		interval = (max - min) / int64(chunkCount)
		if interval == 0 {
			wr.Logger().Infof("table=%v: Not splitting the table into multiple chunks, interval=0: %v to %v", td.Name, min, max)
			return singleCompleteChunk, nil
		}
	case uint64:
		max := max.(uint64)
		interval = (max - min) / uint64(chunkCount)
		if interval == 0 {
			wr.Logger().Infof("table=%v: Not splitting the table into multiple chunks, interval=0: %v to %v", td.Name, min, max)
			return singleCompleteChunk, nil
		}
	case float64:
		max := max.(float64)
		interval = (max - min) / float64(chunkCount)
		if interval == 0 {
			wr.Logger().Infof("table=%v: Not splitting the table into multiple chunks, interval=0: %v to %v", td.Name, min, max)
			return singleCompleteChunk, nil
		}
	default:
		wr.Logger().Infof("table=%v: Not splitting the table into multiple chunks, primary key not numeric.", td.Name)
		return singleCompleteChunk, nil
	}

	// Create chunks.
	start := min
	for i := 0; i < chunkCount; i++ {
		end := add(start, interval)
		chunk, err := toChunk(start, end, i+1, chunkCount)
		if err != nil {
			return nil, vterrors.Wrapf(err, "tablet: %v, table: %v", topoproto.TabletAliasString(tablet.Alias), td.Name)
		}
		chunks[i] = chunk
		start = end
	}

	// Clear out the MIN and MAX on the first and last chunk respectively
	// because other shards might have smaller or higher values than the one we
	// looked at.
	chunks[0].start = sqltypes.NULL
	chunks[chunkCount-1].end = sqltypes.NULL
	return chunks, nil
}

func add(start, interval interface{}) interface{} {
	switch start := start.(type) {
	case int64:
		return start + interval.(int64)
	case uint64:
		return start + interval.(uint64)
	case float64:
		return start + interval.(float64)
	default:
		panic(fmt.Sprintf("unsupported type %T for interval start: %v", start, start))
	}
}

func toChunk(start, end interface{}, number, total int) (chunk, error) {
	startValue, err := sqltypes.InterfaceToValue(start)
	if err != nil {
		return chunk{}, vterrors.Wrapf(err, "failed to convert calculated start value (%v) into internal sqltypes.Value", start)
	}
	endValue, err := sqltypes.InterfaceToValue(end)
	if err != nil {
		return chunk{}, vterrors.Wrapf(err, "failed to convert calculated end value (%v) into internal sqltypes.Value", end)
	}
	return chunk{startValue, endValue, number, total}, nil
}
<|endoftext|>


<commit_before>package utils

import (
	"io/ioutil"
	"encoding/json"
	"github.com/Sirupsen/logrus"
	dockerclient "github.com/fsouza/go-dockerclient"
	"errors"
	"strconv"
	"bytes"
)

const (
	ENDPOINT = "unix:///var/run/docker.sock"
	SHELL = "/bin/sh"
	COMMAND_OPTION = "-c"
)

type OptionalContainerArgs struct {
	Network   string
	Cmd       string
	HostBinds []string
}

func readInContainerTemplate(name string) (dockerclient.CreateContainerOptions, error) {
	var container dockerclient.CreateContainerOptions
	file, err := ioutil.ReadFile("testdata/" + name + ".json")
	if err != nil {
		return container, err
	}

	if err := json.Unmarshal(file, &container); err != nil {
		return container, err
	}
	return container, nil
}

func readInNetworkTemplate(name string) (dockerclient.CreateNetworkOptions, error) {
	var network dockerclient.CreateNetworkOptions
	file, err := ioutil.ReadFile("testdata/" + name + ".json")
	if err != nil {
		return network, err
	}
	if err := json.Unmarshal(file, &network); err != nil {
		return network, err
	}
	return network, nil
}

func DeleteContainer(id string) error {
	containerClient, err := dockerclient.NewClient(ENDPOINT)
	if err != nil {
		return err
	}
	options := dockerclient.RemoveContainerOptions{
		ID: id,
		RemoveVolumes: true,
		Force: true,
	}
	logrus.Debugf("Removing container %s", id)
	if err := containerClient.RemoveContainer(options); err != nil {
		return err
	}
	return nil
}

func StartContainer(id string) error {
	containerClient, err := dockerclient.NewClient(ENDPOINT)
	if err != nil {
		return err
	}
	logrus.Debugf("Starting container %s", id)
	if err := containerClient.StartContainer(id, nil); err != nil {
		return err
	}
	return nil
}

func StopContainer(id string) error {
	containerClient, err := dockerclient.NewClient(ENDPOINT)
	if err != nil {
		return err
	}
	logrus.Debugf("Stopping container %s", id)
	if err := containerClient.StopContainer(id, 0); err != nil {
		return err
	}
	return nil
}

func CreateContainer(description string, optionalArgs *OptionalContainerArgs) (string, error) {
	containerOptions, err := readInContainerTemplate(description)
	if err != nil {
		return "", err
	}
	if optionalArgs.Network != "" {
		containerOptions.HostConfig.NetworkMode = optionalArgs.Network
	}
	if optionalArgs.Cmd != "" {
		containerOptions.Config.Cmd = []string{
			SHELL,
			COMMAND_OPTION,
			optionalArgs.Cmd,
		}
	}
	if optionalArgs.HostBinds != nil {
		containerOptions.HostConfig.Binds = optionalArgs.HostBinds
	}
	logrus.Infof("%s, %s", description, optionalArgs)
	containerClient, err := dockerclient.NewClient(ENDPOINT)
	if err != nil {
		return "", err
	}
	container, err := containerClient.CreateContainer(containerOptions)
	if err != nil {
		return "", err
	}
	return container.ID, nil
}

func ExecuteContainer(containerID string, commands []string) (error) {
	client, err := dockerclient.NewClient(ENDPOINT)
	if err != nil {
		return err
	}
	execInstance, err := client.CreateExec(dockerclient.CreateExecOptions{
		Container: containerID,
		AttachStdin: false,
		AttachStdout: true,
		AttachStderr: true,
		Tty: false,
		Cmd: commands,
	})
	logrus.Debugf("container %s, exec instance %s", containerID, execInstance.ID)
	if err != nil {
		return err
	}
	var stdout bytes.Buffer
	err = client.StartExec(execInstance.ID, dockerclient.StartExecOptions{
		OutputStream: &stdout,
		Detach: false,
		Tty: false,
		RawTerminal: true,
	})
	if err := client.StartExec(execInstance.ID, dockerclient.StartExecOptions{
		OutputStream: &stdout,
		Detach: false,
		Tty: false,
		RawTerminal: true,
	}); err != nil {
		return err
	}
	execResult, err := client.InspectExec(execInstance.ID)
	if err != nil {
		return err
	}
	if execResult.ExitCode != 0 {
		logrus.Errorf("docker exec failed with exit code %d, %s", execResult.ExitCode, stdout.String())
		return errors.New(strconv.Itoa(execResult.ExitCode))
	}
	return nil
}

func CreateNetwork(description string, cidr string) (string, error) {
	options, err := readInNetworkTemplate(description)
	if err != nil {
		return "", err
	}
	if cidr != "" {
		options.IPAM.Config = []dockerclient.IPAMConfig{
			{Subnet: cidr + "/16" },
		}
	}
	client, err := dockerclient.NewClient(ENDPOINT)
	if err != nil {
		return "", err
	}
	logrus.Debugf("Creating network %s", options.Name)
	network, err := client.CreateNetwork(options)
	if err != nil {
		return "", err
	}
	return network.ID, nil
}

func DeleteNetwork(id string) error {
	client, err := dockerclient.NewClient(ENDPOINT)
	if err != nil {
		return err
	}
	logrus.Debugf("Removing network %s", id)
	err = client.RemoveNetwork(id)
	if err != nil {
		return err
	}
	return nil
}<commit_msg>Cleaning up unused functions.<commit_after>package utils

import (
	"io/ioutil"
	"encoding/json"
	"github.com/Sirupsen/logrus"
	dockerclient "github.com/fsouza/go-dockerclient"
	"errors"
	"strconv"
	"bytes"
)

const (
	ENDPOINT = "unix:///var/run/docker.sock"
	SHELL = "/bin/sh"
	COMMAND_OPTION = "-c"
)

type OptionalContainerArgs struct {
	Network   string
	Cmd       string
	HostBinds []string
}

func readInContainerTemplate(name string) (dockerclient.CreateContainerOptions, error) {
	var container dockerclient.CreateContainerOptions
	file, err := ioutil.ReadFile("testdata/" + name + ".json")
	if err != nil {
		return container, err
	}

	if err := json.Unmarshal(file, &container); err != nil {
		return container, err
	}
	return container, nil
}

func readInNetworkTemplate(name string) (dockerclient.CreateNetworkOptions, error) {
	var network dockerclient.CreateNetworkOptions
	file, err := ioutil.ReadFile("testdata/" + name + ".json")
	if err != nil {
		return network, err
	}
	if err := json.Unmarshal(file, &network); err != nil {
		return network, err
	}
	return network, nil
}

func DeleteContainer(id string) error {
	containerClient, err := dockerclient.NewClient(ENDPOINT)
	if err != nil {
		return err
	}
	options := dockerclient.RemoveContainerOptions{
		ID: id,
		RemoveVolumes: true,
		Force: true,
	}
	logrus.Debugf("Removing container %s", id)
	if err := containerClient.RemoveContainer(options); err != nil {
		return err
	}
	return nil
}

func StartContainer(id string) error {
	containerClient, err := dockerclient.NewClient(ENDPOINT)
	if err != nil {
		return err
	}
	logrus.Debugf("Starting container %s", id)
	if err := containerClient.StartContainer(id, nil); err != nil {
		return err
	}
	return nil
}

func StopContainer(id string) error {
	containerClient, err := dockerclient.NewClient(ENDPOINT)
	if err != nil {
		return err
	}
	logrus.Debugf("Stopping container %s", id)
	if err := containerClient.StopContainer(id, 0); err != nil {
		return err
	}
	return nil
}

func CreateContainer(description string, optionalArgs *OptionalContainerArgs) (string, error) {
	containerOptions, err := readInContainerTemplate(description)
	if err != nil {
		return "", err
	}
	if optionalArgs.Network != "" {
		containerOptions.HostConfig.NetworkMode = optionalArgs.Network
	}
	if optionalArgs.Cmd != "" {
		containerOptions.Config.Cmd = []string{
			SHELL,
			COMMAND_OPTION,
			optionalArgs.Cmd,
		}
	}
	if optionalArgs.HostBinds != nil {
		containerOptions.HostConfig.Binds = optionalArgs.HostBinds
	}
	containerClient, err := dockerclient.NewClient(ENDPOINT)
	if err != nil {
		return "", err
	}
	container, err := containerClient.CreateContainer(containerOptions)
	if err != nil {
		return "", err
	}
	return container.ID, nil
}

func ExecuteContainer(containerID string, commands []string) (error) {
	client, err := dockerclient.NewClient(ENDPOINT)
	if err != nil {
		return err
	}
	execInstance, err := client.CreateExec(dockerclient.CreateExecOptions{
		Container: containerID,
		AttachStdin: false,
		AttachStdout: true,
		AttachStderr: true,
		Tty: false,
		Cmd: commands,
	})
	logrus.Debugf("container %s, exec instance %s", containerID, execInstance.ID)
	if err != nil {
		return err
	}
	var stdout bytes.Buffer
	err = client.StartExec(execInstance.ID, dockerclient.StartExecOptions{
		OutputStream: &stdout,
		Detach: false,
		Tty: false,
		RawTerminal: true,
	})
	if err := client.StartExec(execInstance.ID, dockerclient.StartExecOptions{
		OutputStream: &stdout,
		Detach: false,
		Tty: false,
		RawTerminal: true,
	}); err != nil {
		return err
	}
	execResult, err := client.InspectExec(execInstance.ID)
	if err != nil {
		return err
	}
	if execResult.ExitCode != 0 {
		logrus.Errorf("docker exec failed with exit code %d, %s", execResult.ExitCode, stdout.String())
		return errors.New(strconv.Itoa(execResult.ExitCode))
	}
	return nil
}

func CreateNetwork(description string, cidr string) (string, error) {
	options, err := readInNetworkTemplate(description)
	if err != nil {
		return "", err
	}
	if cidr != "" {
		options.IPAM.Config = []dockerclient.IPAMConfig{
			{Subnet: cidr + "/16" },
		}
	}
	client, err := dockerclient.NewClient(ENDPOINT)
	if err != nil {
		return "", err
	}
	logrus.Debugf("Creating network %s", options.Name)
	network, err := client.CreateNetwork(options)
	if err != nil {
		return "", err
	}
	return network.ID, nil
}

func DeleteNetwork(id string) error {
	client, err := dockerclient.NewClient(ENDPOINT)
	if err != nil {
		return err
	}
	logrus.Debugf("Removing network %s", id)
	err = client.RemoveNetwork(id)
	if err != nil {
		return err
	}
	return nil
}<|endoftext|>