<commit_before>// pushk pushes a new version of an app.
//
// See flag.Usage for details.
package main

import (
	"context"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"regexp"
	"sort"
	"strings"

	"github.com/spf13/viper"
	"go.skia.org/infra/go/auth"
	"go.skia.org/infra/go/common"
	"go.skia.org/infra/go/exec"
	"go.skia.org/infra/go/gcr"
	"go.skia.org/infra/go/git"
	"go.skia.org/infra/go/kube/clusterconfig"
	"go.skia.org/infra/go/sklog"
	"go.skia.org/infra/go/util"
)

const (
	// containerRegistryProject is the GCP project in which we store our Docker
	// images via Google Cloud Container Registry.
	containerRegistryProject = "skia-public"

	// Max number of revisions of an image to print when using --list.
	maxListSize = 10
)

func init() {
	flag.Usage = func() {
		fmt.Printf("Usage: pushk <flags> [one or more image names]\n\n")
		fmt.Printf(`pushk pushes a new version of an app.

The command:
  1. Modifies the kubernetes yaml files with the new image.
  2. Commits the changes to the config repo.
  3. Applies the changes with kubectl.

The config is stored in a separate repo that will automaticaly be checked out
under /tmp by default, or the value of the PUSHK_GITDIR environment variable if set.

The command applies the changes by default, or just changes the local yaml files
if --dry-run is supplied.

Examples:
  # Push an exact tag.
  pushk gcr.io/skia-public/fiddler:694900e3ca9468784a5794dc53382d1c8411ab07

  # Push the latest version of docserver.
  pushk docserver --message="Fix bug #1234"

  # Push the latest version of docserver to the skia-corp cluster.
  pushk docserver --only-cluster=skia-corp --message="Fix bug #1234"

  # Push the latest version of docserver and iap-proxy
  pushk docserver iap-proxy

  # Rollback docserver.
  pushk --rollback docserver

  # List the last few versions of the docserver image. Doesn't apply anything.
  pushk --list docserver

  # Compute any changes a push to docserver will make, but do not apply them.
  # Note that the YAML file(s) will be updated, but not committed or pushed.
  pushk --dry-run docserver

ENV:

  The config repo is checked out by default into '/tmp'. This can be
  changed by setting the environment variable PUSKH_GITDIR.
`)
		flag.PrintDefaults()
	}
}

// flags
var (
	onlyCluster  = flag.String("only-cluster", "", "If set then only push to the specified cluster.")
	configFile   = flag.String("config-file", "", "Absolute filename of the config.json file.")
	dryRun       = flag.Bool("dry-run", false, "If true then do not run the kubectl command to apply the changes, and do not commit the changes to the config repo.")
	ignoreDirty  = flag.Bool("ignore-dirty", false, "If true, then do not fail out if the git repo is dirty.")
	list         = flag.Bool("list", false, "List the last few versions of the given image.")
	message      = flag.String("message", "Push", "Message to go along with the change.")
	rollback     = flag.Bool("rollback", false, "If true go back to the second most recent image, otherwise use most recent image.")
	runningInK8s = flag.Bool("running-in-k8s", false, "If true, then does not use flags that do not work in the k8s environment. Eg: '--cluster' when doing 'kubectl apply'.")
)

var (
	validTag = regexp.MustCompile(`^\d\d\d\d-\d\d-\d\dT\d\d_\d\d_\d\dZ-.+$`)

	config *viper.Viper
)

// filter strips the list of tags down to only the ones that conform to our
// constraints and also checks that there are enough tags. The results
// are sorted in ascending order, so oldest tags are first, newest tags
// are last.
func filter(tags []string) ([]string, error) {
	validTags := []string{}
	for _, t := range tags {
		if validTag.MatchString(t) {
			validTags = append(validTags, t)
		}
	}
	sort.Strings(validTags)
	if len(validTags) == 0 {
		return nil, fmt.Errorf("Not enough tags returned.")
	}
	return validTags, nil
}

// tagProvider is a type that returns the correct tag to push for the given imageName.
type tagProvider func(imageName string) ([]string, error)

// imageFromCmdLineImage handles image names, which can be either short, ala 'fiddler', or exact,
// such as gcr.io/skia-public/fiddler:694900e3ca9468784a5794dc53382d1c8411ab07, both of which can
// appear on the command-line.
func imageFromCmdLineImage(imageName string, tp tagProvider) (string, error) {
	if strings.HasPrefix(imageName, "gcr.io/") {
		if *rollback {
			return "", fmt.Errorf("Supplying a fully qualified image name and the --rollback flag are mutually exclusive.")
		}
		if *list {
			return "", fmt.Errorf("Supplying a fully qualified image name and the --list flag are mutually exclusive.")
		}
		return imageName, nil
	}
	// Get all the tags for the selected image.
	tags, err := tp(imageName)
	if err != nil {
		return "", fmt.Errorf("Tag provider failed: %s", err)
	}

	// Filter the tags
	tags, err = filter(tags)
	if err != nil {
		return "", fmt.Errorf("Failed to filter: %s", err)
	}

	if *list {
		if len(tags) > maxListSize {
			tags = tags[len(tags)-maxListSize:]
		}
		for _, tag := range tags {
			fmt.Println(tag)
		}
	}

	// Pick the target tag we want to move to.
	tag := tags[len(tags)-1]
	if *rollback {
		if len(tags) < 2 {
			return "", fmt.Errorf("No version to rollback to.")
		}
		tag = tags[len(tags)-2]
	}

	// The full docker image name and tag of the image we want to deploy.
	return fmt.Sprintf("%s/%s/%s:%s", gcr.SERVER, containerRegistryProject, imageName, tag), nil
}

// byClusterFromChanged returns a map from cluster name to the list of modified
// files in that cluster.
func byClusterFromChanged(gitDir string, changed util.StringSet) (map[string][]string, error) {
	// Find all the directory names, which are really cluster names.
	// filenames will be absolute directory names, e.g.
	// /tmp/k8s-config/skia-public/task-scheduler-be-staging.yaml
	byCluster := map[string][]string{}

	// The first part of that is
	for _, filename := range changed.Keys() {
		// /tmp/k8s-config/skia-public/task-scheduler-be-staging.yaml => skia-public/task-scheduler-be-staging.yaml
		rel, err := filepath.Rel(gitDir, filename)
		if err != nil {
			return nil, err
		}
		// skia-public/task-scheduler-be-staging.yaml => skia-public
		cluster := filepath.Dir(rel)
		arr, ok := byCluster[cluster]
		if !ok {
			arr = []string{filename}
		} else {
			arr = append(arr, filename)
		}
		byCluster[cluster] = arr
	}
	return byCluster, nil
}

func main() {
	common.Init()

	var err error
	var checkout *git.Checkout
	ctx := context.Background()
	config, checkout, err = clusterconfig.NewWithCheckout(ctx, *configFile)
	if err != nil {
		sklog.Fatal(err)
	}

	output, err := checkout.Git(ctx, "status", "-s")
	if err != nil {
		sklog.Fatal(err)
	}
	if strings.TrimSpace(output) != "" {
		if !*ignoreDirty {
			sklog.Fatalf("Found dirty checkout in %s:\n%s", checkout.Dir(), output)
		}
	} else {
		if err := checkout.Update(ctx); err != nil {
			sklog.Fatal(err)
		}
	}

	dirMatch := "*"
	if *onlyCluster != "" {
		dirMatch = *onlyCluster
	}
	glob := fmt.Sprintf("/%s/*.yaml", dirMatch)
	// Get all the yaml files.
	filenames, err := filepath.Glob(filepath.Join(checkout.Dir(), glob))
	if err != nil {
		sklog.Fatal(err)
	}

	tokenSource := auth.NewGCloudTokenSource(containerRegistryProject)
	imageNames := flag.Args()
	if len(imageNames) == 0 {
		fmt.Printf("At least one image name needs to be supplied.")
		flag.Usage()
		os.Exit(1)
	}
	sklog.Infof("Pushing the following images: %q", imageNames)

	gcrTagProvider := func(imageName string) ([]string, error) {
		return gcr.NewClient(tokenSource, containerRegistryProject, imageName).Tags()
	}

	changed := util.StringSet{}
	for _, imageName := range imageNames {
		image, err := imageFromCmdLineImage(imageName, gcrTagProvider)
		if err != nil {
			sklog.Fatal(err)
		}
		if *list {
			// imageFromCmdLineImage printed out the tags, so nothing more to do.
			continue
		}

		// imageRegex has the following groups returned on match:
		// 0 - the entire line
		// 1 - the prefix, i.e. image:, with correct spacing.
		// 2 - full image name
		// 3 - just the tag
		//
		// We pull out the 'prefix' so we can use it when
		// we rewrite the image: line so the indent level is
		// unchanged.
		parts := strings.SplitN(image, ":", 2)
		if len(parts) != 2 {
			sklog.Fatalf("Failed to split imageName: %v", parts)
		}
		imageNoTag := parts[0]
		imageRegex := regexp.MustCompile(fmt.Sprintf(`^(\s+image:\s+)(%s):.*$`, imageNoTag))

		// Loop over all the yaml files and update tags for the given imageName.
		for _, filename := range filenames {
			b, err := ioutil.ReadFile(filename)
			if err != nil {
				sklog.Errorf("Failed to read %q (skipping): %s", filename, err)
				continue
			}
			lines := strings.Split(string(b), "\n")
			for i, line := range lines {
				matches := imageRegex.FindStringSubmatch(line)
				if len(matches) != 3 {
					continue
				}
				changed[filename] = true
				lines[i] = matches[1] + image
			}
			if changed[filename] {
				err := util.WithWriteFile(filename, func(w io.Writer) error {
					_, err := w.Write([]byte(strings.Join(lines, "\n")))
					return err
				})
				if err != nil {
					sklog.Fatalf("Failed to write update config file %q: %s", filename, err)
				}
			}
		}
	}

	// Were any files updated?
	if len(changed) != 0 {
		byCluster, err := byClusterFromChanged(checkout.Dir(), changed)
		if err != nil {
			sklog.Fatal(err)
		}

		// Then loop over cluster names and apply all changed files for that
		// cluster.
		for cluster, files := range byCluster {
			clusterConfig := config.GetStringMapString(fmt.Sprintf("clusters.%s", cluster))

			filenameFlag := fmt.Sprintf("--filename=%s\n", strings.Join(files, ","))
			if !*dryRun {
				for filename := range changed {
					// /tmp/k8s-config/skia-public/task-scheduler-be-staging.yaml => skia-public/task-scheduler-be-staging.yaml
					rel, err := filepath.Rel(checkout.Dir(), filename)
					if err != nil {
						sklog.Fatal(err)
					}
					msg, err := checkout.Git(ctx, "add", rel)
					if err != nil {
						sklog.Fatalf("Failed to stage changes to the config repo: %s: %q", err, msg)
					}
				}

				kubectlArgs := []string{"apply", filenameFlag}
				if !*runningInK8s {
					kubectlArgs = append(kubectlArgs, "--cluster", clusterConfig["context_name"])
				}
				if err := exec.Run(context.Background(), &exec.Command{
					Name:      "kubectl",
					Args:      kubectlArgs,
					LogStderr: true,
					LogStdout: true,
				}); err != nil {
					sklog.Errorf("Failed to run: %s", err)
				}
			}
			fmt.Printf("\nkubectl apply %s --cluster %s\n", filenameFlag, clusterConfig["context_name"])
		}

		// Once everything is pushed, then commit and push the changes.
		msg, err := checkout.Git(ctx, "diff", "--cached", "--name-only")
		if err != nil {
			sklog.Fatalf("Failed to diff :%s: %q", err, msg)
		}
		if msg == "" {
			sklog.Infof("Not pushing since no files changed.")
			return
		}
		msg, err = checkout.Git(ctx, "commit", "-m", *message)
		if err != nil {
			sklog.Fatalf("Failed to commit to the config repo: %s: %q", err, msg)
		}
		msg, err = checkout.Git(ctx, "push", "origin", "master")
		if err != nil {
			sklog.Fatalf("Failed to push the config repo: %s: %q", err, msg)
		}
	} else {
		fmt.Println("Nothing to do.")
	}
}
<commit_msg>[pushk] Add ability to not override dirty images<commit_after>// pushk pushes a new version of an app.
//
// See flag.Usage for details.
package main

import (
	"context"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"regexp"
	"sort"
	"strings"

	"github.com/spf13/viper"
	"go.skia.org/infra/go/auth"
	"go.skia.org/infra/go/common"
	"go.skia.org/infra/go/exec"
	"go.skia.org/infra/go/gcr"
	"go.skia.org/infra/go/git"
	"go.skia.org/infra/go/kube/clusterconfig"
	"go.skia.org/infra/go/sklog"
	"go.skia.org/infra/go/util"
)

const (
	// containerRegistryProject is the GCP project in which we store our Docker
	// images via Google Cloud Container Registry.
	containerRegistryProject = "skia-public"

	// Max number of revisions of an image to print when using --list.
	maxListSize = 10

	// All dirty images are tagged with this suffix.
	dirtyImageTagSuffix = "-dirty"
)

func init() {
	flag.Usage = func() {
		fmt.Printf("Usage: pushk <flags> [one or more image names]\n\n")
		fmt.Printf(`pushk pushes a new version of an app.

The command:
  1. Modifies the kubernetes yaml files with the new image.
  2. Commits the changes to the config repo.
  3. Applies the changes with kubectl.

The config is stored in a separate repo that will automaticaly be checked out
under /tmp by default, or the value of the PUSHK_GITDIR environment variable if set.

The command applies the changes by default, or just changes the local yaml files
if --dry-run is supplied.

Examples:
  # Push an exact tag.
  pushk gcr.io/skia-public/fiddler:694900e3ca9468784a5794dc53382d1c8411ab07

  # Push the latest version of docserver.
  pushk docserver --message="Fix bug #1234"

  # Push the latest version of docserver to the skia-corp cluster.
  pushk docserver --only-cluster=skia-corp --message="Fix bug #1234"

  # Push the latest version of docserver and iap-proxy
  pushk docserver iap-proxy

  # Rollback docserver.
  pushk --rollback docserver

  # List the last few versions of the docserver image. Doesn't apply anything.
  pushk --list docserver

  # Compute any changes a push to docserver will make, but do not apply them.
  # Note that the YAML file(s) will be updated, but not committed or pushed.
  pushk --dry-run docserver

ENV:

  The config repo is checked out by default into '/tmp'. This can be
  changed by setting the environment variable PUSKH_GITDIR.
`)
		flag.PrintDefaults()
	}
}

// flags
var (
	onlyCluster             = flag.String("only-cluster", "", "If set then only push to the specified cluster.")
	configFile              = flag.String("config-file", "", "Absolute filename of the config.json file.")
	dryRun                  = flag.Bool("dry-run", false, "If true then do not run the kubectl command to apply the changes, and do not commit the changes to the config repo.")
	ignoreDirty             = flag.Bool("ignore-dirty", false, "If true, then do not fail out if the git repo is dirty.")
	list                    = flag.Bool("list", false, "List the last few versions of the given image.")
	message                 = flag.String("message", "Push", "Message to go along with the change.")
	rollback                = flag.Bool("rollback", false, "If true go back to the second most recent image, otherwise use most recent image.")
	runningInK8s            = flag.Bool("running-in-k8s", false, "If true, then does not use flags that do not work in the k8s environment. Eg: '--cluster' when doing 'kubectl apply'.")
	doNotOverrideDirtyImage = flag.Bool("do-not-override-dirty-image", false, "If true, then do not push if the latest checkedin image is dirty. Caveat: This only checks the k8s-config repository to determine if image is dirty, it does not check the live running k8s containers.")
)

var (
	validTag = regexp.MustCompile(`^\d\d\d\d-\d\d-\d\dT\d\d_\d\d_\d\dZ-.+$`)

	config *viper.Viper
)

// filter strips the list of tags down to only the ones that conform to our
// constraints and also checks that there are enough tags. The results
// are sorted in ascending order, so oldest tags are first, newest tags
// are last.
func filter(tags []string) ([]string, error) {
	validTags := []string{}
	for _, t := range tags {
		if validTag.MatchString(t) {
			validTags = append(validTags, t)
		}
	}
	sort.Strings(validTags)
	if len(validTags) == 0 {
		return nil, fmt.Errorf("Not enough tags returned.")
	}
	return validTags, nil
}

// tagProvider is a type that returns the correct tag to push for the given imageName.
type tagProvider func(imageName string) ([]string, error)

// imageFromCmdLineImage handles image names, which can be either short, ala 'fiddler', or exact,
// such as gcr.io/skia-public/fiddler:694900e3ca9468784a5794dc53382d1c8411ab07, both of which can
// appear on the command-line.
func imageFromCmdLineImage(imageName string, tp tagProvider) (string, error) {
	if strings.HasPrefix(imageName, "gcr.io/") {
		if *rollback {
			return "", fmt.Errorf("Supplying a fully qualified image name and the --rollback flag are mutually exclusive.")
		}
		if *list {
			return "", fmt.Errorf("Supplying a fully qualified image name and the --list flag are mutually exclusive.")
		}
		return imageName, nil
	}
	// Get all the tags for the selected image.
	tags, err := tp(imageName)
	if err != nil {
		return "", fmt.Errorf("Tag provider failed: %s", err)
	}

	// Filter the tags
	tags, err = filter(tags)
	if err != nil {
		return "", fmt.Errorf("Failed to filter: %s", err)
	}

	if *list {
		if len(tags) > maxListSize {
			tags = tags[len(tags)-maxListSize:]
		}
		for _, tag := range tags {
			fmt.Println(tag)
		}
	}

	// Pick the target tag we want to move to.
	tag := tags[len(tags)-1]
	if *rollback {
		if len(tags) < 2 {
			return "", fmt.Errorf("No version to rollback to.")
		}
		tag = tags[len(tags)-2]
	}

	// The full docker image name and tag of the image we want to deploy.
	return fmt.Sprintf("%s/%s/%s:%s", gcr.SERVER, containerRegistryProject, imageName, tag), nil
}

// byClusterFromChanged returns a map from cluster name to the list of modified
// files in that cluster.
func byClusterFromChanged(gitDir string, changed util.StringSet) (map[string][]string, error) {
	// Find all the directory names, which are really cluster names.
	// filenames will be absolute directory names, e.g.
	// /tmp/k8s-config/skia-public/task-scheduler-be-staging.yaml
	byCluster := map[string][]string{}

	// The first part of that is
	for _, filename := range changed.Keys() {
		// /tmp/k8s-config/skia-public/task-scheduler-be-staging.yaml => skia-public/task-scheduler-be-staging.yaml
		rel, err := filepath.Rel(gitDir, filename)
		if err != nil {
			return nil, err
		}
		// skia-public/task-scheduler-be-staging.yaml => skia-public
		cluster := filepath.Dir(rel)
		arr, ok := byCluster[cluster]
		if !ok {
			arr = []string{filename}
		} else {
			arr = append(arr, filename)
		}
		byCluster[cluster] = arr
	}
	return byCluster, nil
}

func main() {
	common.Init()

	var err error
	var checkout *git.Checkout
	ctx := context.Background()
	config, checkout, err = clusterconfig.NewWithCheckout(ctx, *configFile)
	if err != nil {
		sklog.Fatal(err)
	}

	output, err := checkout.Git(ctx, "status", "-s")
	if err != nil {
		sklog.Fatal(err)
	}
	if strings.TrimSpace(output) != "" {
		if !*ignoreDirty {
			sklog.Fatalf("Found dirty checkout in %s:\n%s", checkout.Dir(), output)
		}
	} else {
		if err := checkout.Update(ctx); err != nil {
			sklog.Fatal(err)
		}
	}

	dirMatch := "*"
	if *onlyCluster != "" {
		dirMatch = *onlyCluster
	}
	glob := fmt.Sprintf("/%s/*.yaml", dirMatch)
	// Get all the yaml files.
	filenames, err := filepath.Glob(filepath.Join(checkout.Dir(), glob))
	if err != nil {
		sklog.Fatal(err)
	}

	tokenSource := auth.NewGCloudTokenSource(containerRegistryProject)
	imageNames := flag.Args()
	if len(imageNames) == 0 {
		fmt.Printf("At least one image name needs to be supplied.")
		flag.Usage()
		os.Exit(1)
	}
	sklog.Infof("Pushing the following images: %q", imageNames)

	gcrTagProvider := func(imageName string) ([]string, error) {
		return gcr.NewClient(tokenSource, containerRegistryProject, imageName).Tags()
	}

	changed := util.StringSet{}
	for _, imageName := range imageNames {
		image, err := imageFromCmdLineImage(imageName, gcrTagProvider)
		if err != nil {
			sklog.Fatal(err)
		}
		if *list {
			// imageFromCmdLineImage printed out the tags, so nothing more to do.
			continue
		}

		// imageRegex has the following groups returned on match:
		// 0 - the entire line
		// 1 - the prefix, i.e. image:, with correct spacing.
		// 2 - full image name
		// 3 - just the tag
		//
		// We pull out the 'prefix' so we can use it when
		// we rewrite the image: line so the indent level is
		// unchanged.
		parts := strings.SplitN(image, ":", 2)
		if len(parts) != 2 {
			sklog.Fatalf("Failed to split imageName: %v", parts)
		}
		imageNoTag := parts[0]
		imageRegex := regexp.MustCompile(fmt.Sprintf(`^(\s+image:\s+)(%s):(.*)$`, imageNoTag))

		// Loop over all the yaml files and update tags for the given imageName.
		for _, filename := range filenames {
			b, err := ioutil.ReadFile(filename)
			if err != nil {
				sklog.Errorf("Failed to read %q (skipping): %s", filename, err)
				continue
			}
			lines := strings.Split(string(b), "\n")
			for i, line := range lines {
				matches := imageRegex.FindStringSubmatch(line)
				if len(matches) != 4 {
					continue
				}
				if *doNotOverrideDirtyImage && strings.HasSuffix(matches[3], dirtyImageTagSuffix) {
					sklog.Infof("%s is dirty. Not pushing to it since --do-not-override-dirty-image is set.", image)
					continue
				}

				changed[filename] = true
				lines[i] = matches[1] + image
			}
			if changed[filename] {
				err := util.WithWriteFile(filename, func(w io.Writer) error {
					_, err := w.Write([]byte(strings.Join(lines, "\n")))
					return err
				})
				if err != nil {
					sklog.Fatalf("Failed to write update config file %q: %s", filename, err)
				}
			}
		}
	}

	// Were any files updated?
	if len(changed) != 0 {
		byCluster, err := byClusterFromChanged(checkout.Dir(), changed)
		if err != nil {
			sklog.Fatal(err)
		}

		// Then loop over cluster names and apply all changed files for that
		// cluster.
		for cluster, files := range byCluster {
			clusterConfig := config.GetStringMapString(fmt.Sprintf("clusters.%s", cluster))

			filenameFlag := fmt.Sprintf("--filename=%s\n", strings.Join(files, ","))
			if !*dryRun {
				for filename := range changed {
					// /tmp/k8s-config/skia-public/task-scheduler-be-staging.yaml => skia-public/task-scheduler-be-staging.yaml
					rel, err := filepath.Rel(checkout.Dir(), filename)
					if err != nil {
						sklog.Fatal(err)
					}
					msg, err := checkout.Git(ctx, "add", rel)
					if err != nil {
						sklog.Fatalf("Failed to stage changes to the config repo: %s: %q", err, msg)
					}
				}

				kubectlArgs := []string{"apply", filenameFlag}
				if !*runningInK8s {
					kubectlArgs = append(kubectlArgs, "--cluster", clusterConfig["context_name"])
				}
				if err := exec.Run(context.Background(), &exec.Command{
					Name:      "kubectl",
					Args:      kubectlArgs,
					LogStderr: true,
					LogStdout: true,
				}); err != nil {
					sklog.Errorf("Failed to run: %s", err)
				}
			}
			fmt.Printf("\nkubectl apply %s --cluster %s\n", filenameFlag, clusterConfig["context_name"])
		}

		// Once everything is pushed, then commit and push the changes.
		msg, err := checkout.Git(ctx, "diff", "--cached", "--name-only")
		if err != nil {
			sklog.Fatalf("Failed to diff :%s: %q", err, msg)
		}
		if msg == "" {
			sklog.Infof("Not pushing since no files changed.")
			return
		}
		msg, err = checkout.Git(ctx, "commit", "-m", *message)
		if err != nil {
			sklog.Fatalf("Failed to commit to the config repo: %s: %q", err, msg)
		}
		msg, err = checkout.Git(ctx, "push", "origin", "master")
		if err != nil {
			sklog.Fatalf("Failed to push the config repo: %s: %q", err, msg)
		}
	} else {
		fmt.Println("Nothing to do.")
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"log"
	"os"
	"os/signal"
	"syscall"
	"time"
)

// Server represents server data.
type Server struct {
	counter int
	quit    chan os.Signal
}

// NewServer init and starts server.
func NewServer() {

	// Handle SIGINT and SIGTERM.
	ch := make(chan os.Signal, 1)
	signal.Notify(ch, syscall.SIGINT, syscall.SIGTERM)

	s := &Server{
		quit: ch,
	}

	s.run()
}

const counterDelay = 1 * time.Second

const maxCount = 42

func (s *Server) run() {

	log.Print("server start")

	for {
		select {

		case <-s.quit:
			s.stop()
			log.Print("server stop, counter:", s.counter)

			return

		default:
			if s.counter > maxCount {
				return
			}
			s.counter++
			log.Print(s.counter)
			time.Sleep(counterDelay)
		}
	}
}

func (s *Server) stop() {

	for s.counter > 0 {
		s.counter--
		log.Print("stop:", s.counter)
		time.Sleep(counterDelay)
	}
}

func main() {
	NewServer()
	fmt.Println("the end")
}
<commit_msg>remove graceful_stop<commit_after><|endoftext|>


<commit_before>package kvexpress

import (
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path"
	"sort"
	"strings"
)

func ReadFile(filepath string) string {
	dat, err := ioutil.ReadFile(filepath)
	check(err)
	return string(dat)
}

func SortFile(file string) string {
	lines := strings.Split(file, "\n")
	sort.Strings(lines)
	return strings.Join(lines, "\n")
}

func WriteFile(data string, filepath string, perms int, direction string) {
	err := ioutil.WriteFile(filepath, []byte(data), os.FileMode(perms))
	check(err)
	log.Print(direction, ": file_wrote='true' location='", filepath, "' permissions='", perms, "'")
}

func CompareFilename(file string) string {
	compare := fmt.Sprintf("%s.compare", path.Base(file))
	full_path := path.Join(path.Dir(file), compare)
	log.Print("in: file='compare' full_path='", full_path, "'")
	return full_path
}

func LastFilename(file string) string {
	last := fmt.Sprintf("%s.last", path.Base(file))
	full_path := path.Join(path.Dir(file), last)
	log.Print("in: file='last' full_path='", full_path, "'")
	return full_path
}

func CheckLastFile(file string, perms int) {
	if _, err := os.Stat(file); err != nil {
		log.Print("in: Last File: ", file, " does not exist.")
		WriteFile("This is a blank file.\n", file, perms, "in")
	}
}
<commit_msg>When sorting - remove blank lines. Closes #28.<commit_after>package kvexpress

import (
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path"
	"sort"
	"strings"
)

func ReadFile(filepath string) string {
	dat, err := ioutil.ReadFile(filepath)
	check(err)
	return string(dat)
}

func SortFile(file string) string {
	lines := strings.Split(file, "\n")
	lines = BlankLineStrip(lines)
	sort.Strings(lines)
	return strings.Join(lines, "\n")
}

func BlankLineStrip(data []string) []string {
	var stripped []string
	for _, str := range data {
		if str != "" {
			stripped = append(stripped, str)
		}
	}
	return stripped
}

func WriteFile(data string, filepath string, perms int, direction string) {
	err := ioutil.WriteFile(filepath, []byte(data), os.FileMode(perms))
	check(err)
	log.Print(direction, ": file_wrote='true' location='", filepath, "' permissions='", perms, "'")
}

func CompareFilename(file string) string {
	compare := fmt.Sprintf("%s.compare", path.Base(file))
	full_path := path.Join(path.Dir(file), compare)
	log.Print("in: file='compare' full_path='", full_path, "'")
	return full_path
}

func LastFilename(file string) string {
	last := fmt.Sprintf("%s.last", path.Base(file))
	full_path := path.Join(path.Dir(file), last)
	log.Print("in: file='last' full_path='", full_path, "'")
	return full_path
}

func CheckLastFile(file string, perms int) {
	if _, err := os.Stat(file); err != nil {
		log.Print("in: Last File: ", file, " does not exist.")
		WriteFile("This is a blank file.\n", file, perms, "in")
	}
}
<|endoftext|>


<commit_before>package store

import (
	"database/sql"
	"errors"
	"fmt"
	"log"
	"strings"
	"time"

	gologme "github.com/erasche/gologme/types"
	_ "github.com/lib/pq"
)

//The first implementation.
type PostgreSQLDataStore struct {
	DSN string
	DB  *sql.DB
}

func (ds *PostgreSQLDataStore) SetupDb() {
	_, err := ds.DB.Exec(
		strings.Replace(DB_SCHEMA, "id integer not null primary key autoincrement", "id serial not null primary key", -1),
	)
	if err != nil {
		log.Fatal(err)
	}
}

func (ds *PostgreSQLDataStore) LogToDb(uid int, windowlogs []*gologme.WindowLogs, keylogs []*gologme.KeyLogs) {
	tx, err := ds.DB.Begin()
	if err != nil {
		log.Fatal(err)
	}
	wl_stmt, err := tx.Prepare("insert into windowLogs (uid, time, name) values (?, ?, ?)")
	if err != nil {
		log.Fatal(err)
	}
	kl_stmt, err := tx.Prepare("insert into keyLogs (uid, time, count) values (?, ?, ?)")
	if err != nil {
		log.Fatal(err)
	}
	defer wl_stmt.Close()
	defer kl_stmt.Close()

	wll := len(windowlogs)
	log.Printf("%d window logs %d key logs from [%d]\n", wll, len(keylogs), uid)

	for _, w := range keylogs {
		_, err = kl_stmt.Exec(uid, w.Time.Unix(), w.Count)
		if err != nil {
			log.Fatal(err)
		}
	}

	for i, w := range windowlogs {
		_, err = wl_stmt.Exec(uid, w.Time.Unix(), w.Name)
		if err != nil {
			log.Fatal(err)
		}
		// Dunno why windowLogs comes through two too big, so whatever.
		if i >= wll-1 {
			break
		}
	}

	tx.Commit()
}

func (ds *PostgreSQLDataStore) CheckAuth(user string, key string) (int, error) {
	// Pretty assuredly not safe from timing attacks.
	var id int
	err := ds.DB.QueryRow("SELECT id FROM users WHERE username = ? AND api_key = ?", user, key).Scan(&id)

	if err != nil {
		if err == sql.ErrNoRows {
			return -1, UserNotFoundError
		}
		return -1, err
	}
	return id, nil
}

func (ds *PostgreSQLDataStore) Name() string {
	return "PostgreSQLDataStore"
}

func (ds *PostgreSQLDataStore) MaxDate() int {
	var mtime int
	err := ds.DB.QueryRow("SELECT time FROM windowLogs ORDER BY time DESC LIMIT 1").Scan(&mtime)

	switch {
	case err == sql.ErrNoRows:
		log.Printf("No data available")
	case err != nil:
		log.Fatal(err)
	}
	return mtime
}

func (ds *PostgreSQLDataStore) MinDate() int {
	var mtime int
	err := ds.DB.QueryRow("SELECT time FROM windowLogs ORDER BY time ASC LIMIT 1").Scan(&mtime)

	switch {
	case err == sql.ErrNoRows:
		log.Printf("No data available")
	case err != nil:
		log.Fatal(err)
	}
	return mtime
}

func (ds *PostgreSQLDataStore) FindUserNameById(id int) (string, error) {
	var username string
	err := ds.DB.QueryRow("SELECT username FROM users WHERE id = ?", id).Scan(&username)
	if err != nil {
		if err == sql.ErrNoRows {
			return "", UserNotFoundError
		}
		return "", err
	}
	return username, nil
}

func (ds *PostgreSQLDataStore) exportWindowLogsByRange(t0 int64, t1 int64) []*gologme.SEvent {
	stmt, err := ds.DB.Prepare("select time, name from windowLogs where time >= ? and time < ? order by id")
	if err != nil {
		log.Fatal(err)
	}
	defer stmt.Close()

	rows, err := stmt.Query(t0, t1)
	if err != nil {
		log.Fatal(err)
	}
	defer rows.Close()

	logs := make([]*gologme.SEvent, 0)
	for rows.Next() {
		var (
			t int
			s string
		)
		rows.Scan(&t, &s)

		logs = append(
			logs,
			&gologme.SEvent{
				T: t,
				S: s,
			},
		)
	}
	return logs
}

func (ds *PostgreSQLDataStore) exportKeyLogsByRange(t0 int64, t1 int64) []*gologme.IEvent {
	stmt, err := ds.DB.Prepare("select time, count from keyLogs where time >= ? and time < ? order by id")
	if err != nil {
		log.Fatal(err)
	}
	defer stmt.Close()

	rows, err := stmt.Query(t0, t1)
	if err != nil {
		log.Fatal(err)
	}
	defer rows.Close()

	logs := make([]*gologme.IEvent, 0)
	for rows.Next() {
		var (
			t int
			s int
		)
		rows.Scan(&t, &s)

		logs = append(
			logs,
			&gologme.IEvent{
				T: t,
				S: s,
			},
		)
	}
	return logs
}

func (ds *PostgreSQLDataStore) exportBlog(t0 int64, t1 int64) []*gologme.SEvent {
	stmt, err := ds.DB.Prepare("select time, type, contents from notes where time >= ? and time < ? and type = ? order by id")
	if err != nil {
		log.Fatal(err)
	}
	defer stmt.Close()

	rows, err := stmt.Query(t0, t1, gologme.BLOG_TYPE)
	if err != nil {
		log.Fatal(err)
	}
	defer rows.Close()

	logs := make([]*gologme.SEvent, 0)
	for rows.Next() {
		var (
			Time     int
			Type     int
			Contents string
		)
		rows.Scan(&Time, &Type, &Contents)
		logs = append(
			logs,
			&gologme.SEvent{
				T: Time,
				S: Contents,
			},
		)
	}
	return logs
}

func (ds *PostgreSQLDataStore) exportNotes(t0 int64, t1 int64) []*gologme.SEvent {
	stmt, err := ds.DB.Prepare("select time, type, contents from notes where time >= ? and time < ? and type = ? order by id")
	if err != nil {
		log.Fatal(err)
	}
	defer stmt.Close()

	rows, err := stmt.Query(t0, t1, gologme.NOTE_TYPE)
	if err != nil {
		log.Fatal(err)
	}
	defer rows.Close()

	logs := make([]*gologme.SEvent, 0)
	for rows.Next() {
		var (
			Time     int
			Type     int
			Contents string
		)
		rows.Scan(&Time, &Type, &Contents)
		logs = append(
			logs,
			&gologme.SEvent{
				T: Time,
				S: Contents,
			},
		)
	}
	return logs
}

func (ds *PostgreSQLDataStore) ExportEventsByDate(tm time.Time) *gologme.EventLog {
	t0 := Ulogme7amTime(tm)
	t1 := Ulogme7amTime(Tomorrow(tm))

	blog := ds.exportBlog(t0, t1)
	var blogstr string
	if len(blog) > 0 {
		blogstr = blog[0].S
	} else {
		blogstr = ""
	}

	return &gologme.EventLog{
		Window_events:  ds.exportWindowLogsByRange(t0, t1),
		Keyfreq_events: ds.exportKeyLogsByRange(t0, t1),
		Note_events:    ds.exportNotes(t0, t1),
		Blog:           blogstr,
	}
}

func NewPostgreSQLDataStore(conf map[string]string) (DataStore, error) {
	var dsn string
	if val, ok := conf["DATASTORE_URL"]; ok {
		dsn = val
	} else {
		return nil, errors.New(fmt.Sprintf("%s is required for the postgres datastore", "DATASTORE_URL"))
	}

	db, err := sql.Open("postgres", dsn)
	if err != nil {
		log.Panicf("Failed to connect to datastore: %s", err.Error())
		return nil, FailedToConnect
	}

	return &PostgreSQLDataStore{
		DSN: dsn,
		DB:  db,
	}, nil
}
<commit_msg>Fix for postgres<commit_after>package store

import (
	"database/sql"
	"errors"
	"fmt"
	"log"
	"strings"
	"time"

	gologme "github.com/erasche/gologme/types"
	_ "github.com/lib/pq"
)

//The first implementation.
type PostgreSQLDataStore struct {
	DSN string
	DB  *sql.DB
}

func (ds *PostgreSQLDataStore) SetupDb() {
	_, err := ds.DB.Exec(
		strings.Replace(DB_SCHEMA, "id integer not null primary key autoincrement", "id serial not null primary key", -1),
	)
	if err != nil {
		log.Fatal(err)
	}
}

func (ds *PostgreSQLDataStore) LogToDb(uid int, windowlogs []*gologme.WindowLogs, keylogs []*gologme.KeyLogs) {
	tx, err := ds.DB.Begin()
	if err != nil {
		log.Fatal(err)
	}
	wl_stmt, err := tx.Prepare("insert into windowLogs (uid, time, name) values ($1, $2, $3)")
	if err != nil {
		log.Fatal(err)
	}
	kl_stmt, err := tx.Prepare("insert into keyLogs (uid, time, count) values ($1, $2, $3)")
	if err != nil {
		log.Fatal(err)
	}
	defer wl_stmt.Close()
	defer kl_stmt.Close()

	wll := len(windowlogs)
	log.Printf("%d window logs %d key logs from [%d]\n", wll, len(keylogs), uid)

	for _, w := range keylogs {
		_, err = kl_stmt.Exec(uid, w.Time.Unix(), w.Count)
		if err != nil {
			log.Fatal(err)
		}
	}

	for i, w := range windowlogs {
		_, err = wl_stmt.Exec(uid, w.Time.Unix(), w.Name)
		if err != nil {
			log.Fatal(err)
		}
		// Dunno why windowLogs comes through two too big, so whatever.
		if i >= wll-1 {
			break
		}
	}

	tx.Commit()
}

func (ds *PostgreSQLDataStore) CheckAuth(user string, key string) (int, error) {
	// Pretty assuredly not safe from timing attacks.
	var id int
	err := ds.DB.QueryRow("SELECT id FROM users WHERE username = $1 AND api_key = $2", user, key).Scan(&id)

	if err != nil {
		if err == sql.ErrNoRows {
			return -1, UserNotFoundError
		}
		return -1, err
	}
	return id, nil
}

func (ds *PostgreSQLDataStore) Name() string {
	return "PostgreSQLDataStore"
}

func (ds *PostgreSQLDataStore) MaxDate() int {
	var mtime int
	err := ds.DB.QueryRow("SELECT time FROM windowLogs ORDER BY time DESC LIMIT 1").Scan(&mtime)

	switch {
	case err == sql.ErrNoRows:
		log.Printf("No data available")
	case err != nil:
		log.Fatal(err)
	}
	return mtime
}

func (ds *PostgreSQLDataStore) MinDate() int {
	var mtime int
	err := ds.DB.QueryRow("SELECT time FROM windowLogs ORDER BY time ASC LIMIT 1").Scan(&mtime)

	switch {
	case err == sql.ErrNoRows:
		log.Printf("No data available")
	case err != nil:
		log.Fatal(err)
	}
	return mtime
}

func (ds *PostgreSQLDataStore) FindUserNameById(id int) (string, error) {
	var username string
	err := ds.DB.QueryRow("SELECT username FROM users WHERE id = $1", id).Scan(&username)
	if err != nil {
		if err == sql.ErrNoRows {
			return "", UserNotFoundError
		}
		return "", err
	}
	return username, nil
}

func (ds *PostgreSQLDataStore) exportWindowLogsByRange(t0 int64, t1 int64) []*gologme.SEvent {
	stmt, err := ds.DB.Prepare("select time, name from windowLogs where time >= $1 and time < $2 order by id")
	if err != nil {
		log.Fatal(err)
	}
	defer stmt.Close()

	rows, err := stmt.Query(t0, t1)
	if err != nil {
		log.Fatal(err)
	}
	defer rows.Close()

	logs := make([]*gologme.SEvent, 0)
	for rows.Next() {
		var (
			t int
			s string
		)
		rows.Scan(&t, &s)

		logs = append(
			logs,
			&gologme.SEvent{
				T: t,
				S: s,
			},
		)
	}
	return logs
}

func (ds *PostgreSQLDataStore) exportKeyLogsByRange(t0 int64, t1 int64) []*gologme.IEvent {
	stmt, err := ds.DB.Prepare("select time, count from keyLogs where time >= $1 and time < $2 order by id")
	if err != nil {
		log.Fatal(err)
	}
	defer stmt.Close()

	rows, err := stmt.Query(t0, t1)
	if err != nil {
		log.Fatal(err)
	}
	defer rows.Close()

	logs := make([]*gologme.IEvent, 0)
	for rows.Next() {
		var (
			t int
			s int
		)
		rows.Scan(&t, &s)

		logs = append(
			logs,
			&gologme.IEvent{
				T: t,
				S: s,
			},
		)
	}
	return logs
}

func (ds *PostgreSQLDataStore) exportBlog(t0 int64, t1 int64) []*gologme.SEvent {
	stmt, err := ds.DB.Prepare("select time, type, contents from notes where time >= $1 and time < $2 and type = $3 order by id")
	if err != nil {
		log.Fatal(err)
	}
	defer stmt.Close()

	rows, err := stmt.Query(t0, t1, gologme.BLOG_TYPE)
	if err != nil {
		log.Fatal(err)
	}
	defer rows.Close()

	logs := make([]*gologme.SEvent, 0)
	for rows.Next() {
		var (
			Time     int
			Type     int
			Contents string
		)
		rows.Scan(&Time, &Type, &Contents)
		logs = append(
			logs,
			&gologme.SEvent{
				T: Time,
				S: Contents,
			},
		)
	}
	return logs
}

func (ds *PostgreSQLDataStore) exportNotes(t0 int64, t1 int64) []*gologme.SEvent {
	stmt, err := ds.DB.Prepare("select time, type, contents from notes where time >= $1 and time < $2 and type = $3 order by id")
	if err != nil {
		log.Fatal(err)
	}
	defer stmt.Close()

	rows, err := stmt.Query(t0, t1, gologme.NOTE_TYPE)
	if err != nil {
		log.Fatal(err)
	}
	defer rows.Close()

	logs := make([]*gologme.SEvent, 0)
	for rows.Next() {
		var (
			Time     int
			Type     int
			Contents string
		)
		rows.Scan(&Time, &Type, &Contents)
		logs = append(
			logs,
			&gologme.SEvent{
				T: Time,
				S: Contents,
			},
		)
	}
	return logs
}

func (ds *PostgreSQLDataStore) ExportEventsByDate(tm time.Time) *gologme.EventLog {
	t0 := Ulogme7amTime(tm)
	t1 := Ulogme7amTime(Tomorrow(tm))

	blog := ds.exportBlog(t0, t1)
	var blogstr string
	if len(blog) > 0 {
		blogstr = blog[0].S
	} else {
		blogstr = ""
	}

	return &gologme.EventLog{
		Window_events:  ds.exportWindowLogsByRange(t0, t1),
		Keyfreq_events: ds.exportKeyLogsByRange(t0, t1),
		Note_events:    ds.exportNotes(t0, t1),
		Blog:           blogstr,
	}
}

func NewPostgreSQLDataStore(conf map[string]string) (DataStore, error) {
	var dsn string
	if val, ok := conf["DATASTORE_URL"]; ok {
		dsn = val
	} else {
		return nil, errors.New(fmt.Sprintf("%s is required for the postgres datastore", "DATASTORE_URL"))
	}

	db, err := sql.Open("postgres", dsn)
	if err != nil {
		log.Panicf("Failed to connect to datastore: %s", err.Error())
		return nil, FailedToConnect
	}

	return &PostgreSQLDataStore{
		DSN: dsn,
		DB:  db,
	}, nil
}
<|endoftext|>


<commit_before>package version

// Maj is the major release
const Maj = "0"

// Min is the minor release
const Min = "5"

// Fix is the patch fix number
const Fix = "0"

var (
	// Version is The full version string
	Version = "0.5.0"

	// GitCommit is set with --ldflags "-X main.gitCommit=$(git rev-parse HEAD)"
	GitCommit string
)

func init() {
	if GitCommit != "" {
		Version += "-" + GitCommit[:8]
	}
}
<commit_msg>Add flag in version.go<commit_after>package version

const flag = "develop"

var (
	// Version is The full version string
	Version = "0.5.0"

	// GitCommit is set with --ldflags "-X main.gitCommit=$(git rev-parse HEAD)"
	GitCommit string
)

func init() {
	Version += "-" + flag

	if GitCommit != "" {
		Version += "-" + GitCommit[:8]
	}
}
<|endoftext|>


<commit_before>package db

import (
	"encoding/json"
	"io/ioutil"
	"os"
	"path"
	"testing"
)

/*
 * Lowest-layer database tests
 */

func Test_DbFileCreation(t *testing.T) {
	dir, err := ioutil.TempDir("", "rqlite-test-")
	defer os.RemoveAll(dir)

	db, err := Open(path.Join(dir, "test_db"))
	if err != nil {
		t.Fatalf("failed to open new database: %s", err.Error())
	}
	if db == nil {
		t.Fatal("database is nil")
	}
	err = db.Close()
	if err != nil {
		t.Fatalf("failed to close database: %s", err.Error())
	}
}

func Test_TableCreation(t *testing.T) {
	db, path := mustOpenDatabase()
	defer db.Close()
	defer os.Remove(path)

	err := db.Execute("create table foo (id integer not null primary key, name text)")
	if err != nil {
		t.Fatalf("failed to create table: %s", err.Error())
	}

	r, err := db.Query("SELECT * FROM foo")
	if err != nil {
		t.Fatalf("failed to query empty table: %s", err.Error())
	}
	if len(r) != 0 {
		t.Fatalf("expected 0 results, got %d", len(r))
	}
}

func Test_SimpleStatements(t *testing.T) {
	db, path := mustOpenDatabase()
	defer db.Close()
	defer os.Remove(path)

	if err := db.Execute("create table foo (id integer not null primary key, name text)"); err != nil {
		t.Fatalf("failed to create table: %s", err.Error())
	}

	if err := db.Execute(`INSERT INTO foo(name) VALUES("fiona")`); err != nil {
		t.Fatalf("failed to insert record: %s", err.Error())
	}
	r, err := db.Query("SELECT * FROM foo")
	if err != nil {
		t.Fatalf("failed to query empty table: %s", err.Error())
	}
	if exp, got := `[{"id":"1","name":"fiona"}]`, asJson(r); exp != got {
		t.Fatalf("unexpected results for query, expected %s, got %s", exp, got)
	}

	if err := db.Execute(`INSERT INTO foo(name) VALUES("aoife")`); err != nil {
		t.Fatalf("failed to insert record: %s", err.Error())
	}
	r, err = db.Query("SELECT * FROM foo")
	if err != nil {
		t.Fatalf("failed to query table: %s", err.Error())
	}
	if exp, got := `[{"id":"1","name":"fiona"},{"id":"2","name":"aoife"}]`, asJson(r); exp != got {
		t.Fatalf("unexpected results for query, expected %s, got %s", exp, got)
	}

	if err := db.Execute("UPDATE foo SET Name='Who knows?' WHERE Id=1"); err != nil {
		t.Fatalf("failed to update record: %s", err.Error())
	}
	r, err = db.Query("SELECT * FROM foo")
	if err != nil {
		t.Fatalf("failed to query table: %s", err.Error())
	}
	if exp, got := `[{"id":"1","name":"Who knows?"},{"id":"2","name":"aoife"}]`, asJson(r); exp != got {
		t.Fatalf("unexpected results for query, expected %s, got %s", exp, got)
	}

	if err := db.Execute("DELETE FROM foo WHERE Id=2"); err != nil {
		t.Fatalf("failed to delete record: %s", err.Error())
	}
	r, err = db.Query("SELECT * FROM foo")
	if err != nil {
		t.Fatalf("failed to query table: %s", err.Error())
	}
	if exp, got := `[{"id":"1","name":"Who knows?"}]`, asJson(r); exp != got {
		t.Fatalf("unexpected results for query, expected %s, got %s", exp, got)
	}
}

func Test_FailingSimpleStatements(t *testing.T) {
	db, path := mustOpenDatabase()
	defer db.Close()
	defer os.Remove(path)

	err := db.Execute(`INSERT INTO foo(name) VALUES("fiona")`)
	if err == nil {
		t.Fatal("inserted record into empty table OK")
	}
	if err.Error() != "no such table: foo" {
		t.Fatal("unexpected error returned")
	}

	if err = db.Execute("create table foo (id integer not null primary key, name text)"); err != nil {
		t.Fatalf("failed to create table: %s", err.Error())
	}
	if err = db.Execute("create table foo (id integer not null primary key, name text)"); err == nil {
		t.Fatal("duplicate table created OK")
	}
	if err.Error() != "table foo already exists" {
		t.Fatalf("unexpected error returned: %s", err.Error())
	}

	if err = db.Execute(`INSERT INTO foo(id, name) VALUES(11, "fiona")`); err != nil {
		t.Fatalf("failed to insert record: %s", err.Error())
	}
	if err = db.Execute(`INSERT INTO foo(id, name) VALUES(11, "fiona")`); err == nil {
		t.Fatal("duplicated record inserted OK")
	}
	if err.Error() != "UNIQUE constraint failed: foo.id" {
		t.Fatal("unexpected error returned")
	}

	if err = db.Execute("SELECT * FROM bar"); err == nil {
		t.Fatal("no error when querying non-existant table")
	}
	if err.Error() != "no such table: bar" {
		t.Fatal("unexpected error returned")
	}

	if err = db.Execute("utter nonsense"); err == nil {
		t.Fatal("no error when issuing nonsense query")
	}
	if err.Error() != `near "utter": syntax error` {
		t.Fatal("unexpected error returned")
	}
}

func Test_SimpleTransactions(t *testing.T) {
	db, path := mustOpenDatabase()
	defer db.Close()
	defer os.Remove(path)

	if err := db.Execute("create table foo (id integer not null primary key, name text)"); err != nil {
		t.Fatalf("failed to create table: %s", err.Error())
	}

	if err := db.StartTransaction(); err != nil {
		t.Fatalf("failed to start transaction: %s", err.Error())
	}
	for i := 0; i < 10; i++ {
		_ = db.Execute(`INSERT INTO foo(name) VALUES("philip")`)
	}
	if err := db.CommitTransaction(); err != nil {
		t.Fatalf("failed to commit transaction: %s", err.Error())
	}

	r, err := db.Query("SELECT name FROM foo")
	if err != nil {
		t.Fatalf("failed to query after commited transaction: %s", err.Error())
	}
	if len(r) != 10 {
		t.Fatalf("incorrect number of results returned: %d", len(r))
	}

	if err := db.StartTransaction(); err != nil {
		t.Fatalf("failed to start transaction: %s", err.Error())
	}
	for i := 0; i < 10; i++ {
		_ = db.Execute(`INSERT INTO foo(name) VALUES("philip")`)
	}
	if err := db.RollbackTransaction(); err != nil {
		t.Fatalf("failed to rollback transaction: %s", err.Error())
	}

	r, err = db.Query("SELECT name FROM foo")
	if err != nil {
		t.Fatalf("failed to query after commited transaction: %s", err.Error())
	}
	if len(r) != 10 {
		t.Fatalf("incorrect number of results returned: %d", len(r))
	}
}

func Test_TransactionsConstraintViolation(t *testing.T) {
	db, path := mustOpenDatabase()
	defer db.Close()
	defer os.Remove(path)
	var err error

	if err = db.Execute("create table foo (id integer not null primary key, name text)"); err != nil {
		t.Fatalf("failed to create table: %s", err.Error())
	}

	if err = db.StartTransaction(); err != nil {
		t.Fatalf("failed to start transaction: %s", err.Error())
	}
	if err = db.Execute(`INSERT INTO foo(id, name) VALUES(11, "fiona")`); err != nil {
		t.Fatalf("failed to insert record: %s", err.Error())
	}
	if err = db.Execute(`INSERT INTO foo(id, name) VALUES(11, "fiona")`); err == nil {
		t.Fatal("duplicated record inserted OK")
	}
	if err.Error() != "UNIQUE constraint failed: foo.id" {
		t.Fatal("unexpected error returned")
	}
	if err := db.RollbackTransaction(); err != nil {
		t.Fatalf("failed to rollback transaction: %s", err.Error())
	}
	r, err := db.Query("SELECT name FROM foo")
	if err != nil {
		t.Fatalf("failed to query after commited transaction: %s", err.Error())
	}
	if len(r) != 0 {
		t.Fatalf("incorrect number of results returned: %d", len(r))
	}
}

// func Test_TransactionsHardFail(t *testing.T) {
// 	dir, err := ioutil.TempDir("", "rqlite-test-")
// 	defer os.RemoveAll(dir)
// 	db := New(path.Join(dir, "test_db"))

// 	err = db.Execute("create table foo (id integer not null primary key, name text)")
// 	c.Assert(err, IsNil)
// 	err = db.Execute("INSERT INTO foo(id, name) VALUES(1, \"fiona\")")
// 	c.Assert(err, IsNil)

// 	err = db.StartTransaction()
// 	c.Assert(err, IsNil)
// 	err = db.Execute("INSERT INTO foo(id, name) VALUES(2, \"dana\")")
// 	c.Assert(err, IsNil)
// 	db.Close()

// 	db = Open(path.Join(dir, "test_db"))
// 	c.Assert(db, NotNil)
// 	r, err := db.Query("SELECT * FROM foo")
// 	c.Assert(len(r), Equals, 1)
// 	c.Assert(r[0]["name"], Equals, "fiona")
// 	db.Close()
// }

func mustOpenDatabase() (*DB, string) {
	var err error
	f, err := ioutil.TempFile("", "rqlilte-test-")
	if err != nil {
		panic("failed to create temp file")
	}
	f.Close()

	db, err := Open(f.Name())
	if err != nil {
		panic("failed to open database")
	}

	return db, f.Name()
}

func asJson(v interface{}) string {
	b, err := json.Marshal(v)
	if err != nil {
		panic("failed to JSON marshal value")
	}
	return string(b)
}
<commit_msg>Test_TransactionsHardFail now passes<commit_after>package db

import (
	"encoding/json"
	"io/ioutil"
	"os"
	"path"
	"testing"
)

/*
 * Lowest-layer database tests
 */

func Test_DbFileCreation(t *testing.T) {
	dir, err := ioutil.TempDir("", "rqlite-test-")
	defer os.RemoveAll(dir)

	db, err := Open(path.Join(dir, "test_db"))
	if err != nil {
		t.Fatalf("failed to open new database: %s", err.Error())
	}
	if db == nil {
		t.Fatal("database is nil")
	}
	err = db.Close()
	if err != nil {
		t.Fatalf("failed to close database: %s", err.Error())
	}
}

func Test_TableCreation(t *testing.T) {
	db, path := mustOpenDatabase()
	defer db.Close()
	defer os.Remove(path)

	err := db.Execute("create table foo (id integer not null primary key, name text)")
	if err != nil {
		t.Fatalf("failed to create table: %s", err.Error())
	}

	r, err := db.Query("SELECT * FROM foo")
	if err != nil {
		t.Fatalf("failed to query empty table: %s", err.Error())
	}
	if len(r) != 0 {
		t.Fatalf("expected 0 results, got %d", len(r))
	}
}

func Test_SimpleStatements(t *testing.T) {
	db, path := mustOpenDatabase()
	defer db.Close()
	defer os.Remove(path)

	if err := db.Execute("create table foo (id integer not null primary key, name text)"); err != nil {
		t.Fatalf("failed to create table: %s", err.Error())
	}

	if err := db.Execute(`INSERT INTO foo(name) VALUES("fiona")`); err != nil {
		t.Fatalf("failed to insert record: %s", err.Error())
	}
	r, err := db.Query("SELECT * FROM foo")
	if err != nil {
		t.Fatalf("failed to query empty table: %s", err.Error())
	}
	if exp, got := `[{"id":"1","name":"fiona"}]`, asJson(r); exp != got {
		t.Fatalf("unexpected results for query, expected %s, got %s", exp, got)
	}

	if err := db.Execute(`INSERT INTO foo(name) VALUES("aoife")`); err != nil {
		t.Fatalf("failed to insert record: %s", err.Error())
	}
	r, err = db.Query("SELECT * FROM foo")
	if err != nil {
		t.Fatalf("failed to query table: %s", err.Error())
	}
	if exp, got := `[{"id":"1","name":"fiona"},{"id":"2","name":"aoife"}]`, asJson(r); exp != got {
		t.Fatalf("unexpected results for query, expected %s, got %s", exp, got)
	}

	if err := db.Execute("UPDATE foo SET Name='Who knows?' WHERE Id=1"); err != nil {
		t.Fatalf("failed to update record: %s", err.Error())
	}
	r, err = db.Query("SELECT * FROM foo")
	if err != nil {
		t.Fatalf("failed to query table: %s", err.Error())
	}
	if exp, got := `[{"id":"1","name":"Who knows?"},{"id":"2","name":"aoife"}]`, asJson(r); exp != got {
		t.Fatalf("unexpected results for query, expected %s, got %s", exp, got)
	}

	if err := db.Execute("DELETE FROM foo WHERE Id=2"); err != nil {
		t.Fatalf("failed to delete record: %s", err.Error())
	}
	r, err = db.Query("SELECT * FROM foo")
	if err != nil {
		t.Fatalf("failed to query table: %s", err.Error())
	}
	if exp, got := `[{"id":"1","name":"Who knows?"}]`, asJson(r); exp != got {
		t.Fatalf("unexpected results for query, expected %s, got %s", exp, got)
	}
}

func Test_FailingSimpleStatements(t *testing.T) {
	db, path := mustOpenDatabase()
	defer db.Close()
	defer os.Remove(path)

	err := db.Execute(`INSERT INTO foo(name) VALUES("fiona")`)
	if err == nil {
		t.Fatal("inserted record into empty table OK")
	}
	if err.Error() != "no such table: foo" {
		t.Fatal("unexpected error returned")
	}

	if err = db.Execute("create table foo (id integer not null primary key, name text)"); err != nil {
		t.Fatalf("failed to create table: %s", err.Error())
	}
	if err = db.Execute("create table foo (id integer not null primary key, name text)"); err == nil {
		t.Fatal("duplicate table created OK")
	}
	if err.Error() != "table foo already exists" {
		t.Fatalf("unexpected error returned: %s", err.Error())
	}

	if err = db.Execute(`INSERT INTO foo(id, name) VALUES(11, "fiona")`); err != nil {
		t.Fatalf("failed to insert record: %s", err.Error())
	}
	if err = db.Execute(`INSERT INTO foo(id, name) VALUES(11, "fiona")`); err == nil {
		t.Fatal("duplicated record inserted OK")
	}
	if err.Error() != "UNIQUE constraint failed: foo.id" {
		t.Fatal("unexpected error returned")
	}

	if err = db.Execute("SELECT * FROM bar"); err == nil {
		t.Fatal("no error when querying non-existant table")
	}
	if err.Error() != "no such table: bar" {
		t.Fatal("unexpected error returned")
	}

	if err = db.Execute("utter nonsense"); err == nil {
		t.Fatal("no error when issuing nonsense query")
	}
	if err.Error() != `near "utter": syntax error` {
		t.Fatal("unexpected error returned")
	}
}

func Test_SimpleTransactions(t *testing.T) {
	db, path := mustOpenDatabase()
	defer db.Close()
	defer os.Remove(path)

	if err := db.Execute("create table foo (id integer not null primary key, name text)"); err != nil {
		t.Fatalf("failed to create table: %s", err.Error())
	}

	if err := db.StartTransaction(); err != nil {
		t.Fatalf("failed to start transaction: %s", err.Error())
	}
	for i := 0; i < 10; i++ {
		_ = db.Execute(`INSERT INTO foo(name) VALUES("philip")`)
	}
	if err := db.CommitTransaction(); err != nil {
		t.Fatalf("failed to commit transaction: %s", err.Error())
	}

	r, err := db.Query("SELECT name FROM foo")
	if err != nil {
		t.Fatalf("failed to query after commited transaction: %s", err.Error())
	}
	if len(r) != 10 {
		t.Fatalf("incorrect number of results returned: %d", len(r))
	}

	if err := db.StartTransaction(); err != nil {
		t.Fatalf("failed to start transaction: %s", err.Error())
	}
	for i := 0; i < 10; i++ {
		_ = db.Execute(`INSERT INTO foo(name) VALUES("philip")`)
	}
	if err := db.RollbackTransaction(); err != nil {
		t.Fatalf("failed to rollback transaction: %s", err.Error())
	}

	r, err = db.Query("SELECT name FROM foo")
	if err != nil {
		t.Fatalf("failed to query after commited transaction: %s", err.Error())
	}
	if len(r) != 10 {
		t.Fatalf("incorrect number of results returned: %d", len(r))
	}
}

func Test_TransactionsConstraintViolation(t *testing.T) {
	db, path := mustOpenDatabase()
	defer db.Close()
	defer os.Remove(path)
	var err error

	if err = db.Execute("create table foo (id integer not null primary key, name text)"); err != nil {
		t.Fatalf("failed to create table: %s", err.Error())
	}

	if err = db.StartTransaction(); err != nil {
		t.Fatalf("failed to start transaction: %s", err.Error())
	}
	if err = db.Execute(`INSERT INTO foo(id, name) VALUES(11, "fiona")`); err != nil {
		t.Fatalf("failed to insert record: %s", err.Error())
	}
	if err = db.Execute(`INSERT INTO foo(id, name) VALUES(11, "fiona")`); err == nil {
		t.Fatal("duplicated record inserted OK")
	}
	if err.Error() != "UNIQUE constraint failed: foo.id" {
		t.Fatal("unexpected error returned")
	}
	if err := db.RollbackTransaction(); err != nil {
		t.Fatalf("failed to rollback transaction: %s", err.Error())
	}
	r, err := db.Query("SELECT name FROM foo")
	if err != nil {
		t.Fatalf("failed to query after commited transaction: %s", err.Error())
	}
	if len(r) != 0 {
		t.Fatalf("incorrect number of results returned: %d", len(r))
	}
}

func Test_TransactionsHardFail(t *testing.T) {
	db, path := mustOpenDatabase()
	defer db.Close()
	defer os.Remove(path)
	var err error

	if err = db.Execute("create table foo (id integer not null primary key, name text)"); err != nil {
		t.Fatalf("failed to create table: %s", err.Error())
	}
	if err = db.Execute(`INSERT INTO foo(id, name) VALUES(1, "fiona")`); err != nil {
		t.Fatalf("failed to insert record: %s", err.Error())
	}

	if err = db.StartTransaction(); err != nil {
		t.Fatalf("failed to start transaction: %s", err.Error())
	}
	if err = db.Execute(`INSERT INTO foo(id, name) VALUES(2, "fiona")`); err != nil {
		t.Fatalf("failed to insert record: %s", err.Error())
	}
	if err = db.Close(); err != nil {
		t.Fatalf("failed to close database: %s", err.Error())
	}

	db, err = Open(path)
	if err != nil {
		t.Fatalf("failed to re-open database: %s", err.Error())
	}
	r, err := db.Query("SELECT * FROM foo")
	if err != nil {
		t.Fatalf("failed to query after hard close: %s", err.Error())
	}
	if len(r) != 1 {
		t.Fatalf("incorrect number of results returned: %d", len(r))
	}
}

func mustOpenDatabase() (*DB, string) {
	var err error
	f, err := ioutil.TempFile("", "rqlilte-test-")
	if err != nil {
		panic("failed to create temp file")
	}
	f.Close()

	db, err := Open(f.Name())
	if err != nil {
		panic("failed to open database")
	}

	return db, f.Name()
}

func asJson(v interface{}) string {
	b, err := json.Marshal(v)
	if err != nil {
		panic("failed to JSON marshal value")
	}
	return string(b)
}
<|endoftext|>


<commit_before>package gb

import (
	"os"
	"path/filepath"
	"runtime"
	"time"

	"github.com/constabulary/gb/debug"
)

// pkgpath returns the destination for object cached for this Package.
func pkgpath(pkg *Package) string {
	importpath := filepath.FromSlash(pkg.ImportPath) + ".a"
	switch {
	case pkg.isCrossCompile():
		return filepath.Join(pkg.Pkgdir(), importpath)
	case pkg.Standard && pkg.race:
		// race enabled standard lib
		return filepath.Join(runtime.GOROOT(), "pkg", pkg.gotargetos+"_"+pkg.gotargetarch+"_race", importpath)
	case pkg.Standard:
		// standard lib
		return filepath.Join(runtime.GOROOT(), "pkg", pkg.gotargetos+"_"+pkg.gotargetarch, importpath)
	default:
		return filepath.Join(pkg.Pkgdir(), importpath)
	}
}

// installpath returns the distination to cache this package's compiled .a file.
// pkgpath and installpath differ in that the former returns the location where you will find
// a previously cached .a file, the latter returns the location where an installed file
// will be placed.
//
// The difference is subtle. pkgpath must deal with the possibility that the file is from the
// standard library and is previously compiled. installpath will always return a path for the
// project's pkg/ directory in the case that the stdlib is out of date, or not compiled for
// a specific architecture.
func installpath(pkg *Package) string {
	if pkg.Scope == "test" {
		panic("installpath called with test scope")
	}
	return filepath.Join(pkg.Pkgdir(), filepath.FromSlash(pkg.ImportPath)+".a")
}

// isStale returns true if the source pkg is considered to be stale with
// respect to its installed version.
func isStale(pkg *Package) bool {
	switch pkg.ImportPath {
	case "C", "unsafe":
		// synthetic packages are never stale
		return false
	}

	if !pkg.Standard && pkg.Force {
		return true
	}

	// tests are always stale, they are never installed
	if pkg.Scope == "test" {
		return true
	}

	// Package is stale if completely unbuilt.
	var built time.Time
	if fi, err := os.Stat(pkgpath(pkg)); err == nil {
		built = fi.ModTime()
	}

	if built.IsZero() {
		debug.Debugf("%s is missing", pkgpath(pkg))
		return true
	}

	olderThan := func(file string) bool {
		fi, err := os.Stat(file)
		return err != nil || fi.ModTime().After(built)
	}

	newerThan := func(file string) bool {
		fi, err := os.Stat(file)
		return err != nil || fi.ModTime().Before(built)
	}

	// As a courtesy to developers installing new versions of the compiler
	// frequently, define that packages are stale if they are
	// older than the compiler, and commands if they are older than
	// the linker.  This heuristic will not work if the binaries are
	// back-dated, as some binary distributions may do, but it does handle
	// a very common case.
	if !pkg.Standard {
		if olderThan(pkg.tc.compiler()) {
			debug.Debugf("%s is older than %s", pkgpath(pkg), pkg.tc.compiler())
			return true
		}
		if pkg.IsCommand() && olderThan(pkg.tc.linker()) {
			debug.Debugf("%s is older than %s", pkgpath(pkg), pkg.tc.compiler())
			return true
		}
	}

	// Package is stale if a dependency is newer.
	for _, p := range pkg.Imports() {
		if p.ImportPath == "C" || p.ImportPath == "unsafe" {
			continue // ignore stale imports of synthetic packages
		}
		if olderThan(pkgpath(p)) {
			debug.Debugf("%s is older than %s", pkgpath(pkg), pkgpath(p))
			return true
		}
	}

	// if the main package is up to date but _newer_ than the binary (which
	// could have been removed), then consider it stale.
	if pkg.isMain() && newerThan(pkg.Binfile()) {
		debug.Debugf("%s is newer than %s", pkgpath(pkg), pkg.Binfile())
		return true
	}

	srcs := stringList(pkg.GoFiles, pkg.CFiles, pkg.CXXFiles, pkg.MFiles, pkg.HFiles, pkg.SFiles, pkg.CgoFiles, pkg.SysoFiles, pkg.SwigFiles, pkg.SwigCXXFiles)

	for _, src := range srcs {
		if olderThan(filepath.Join(pkg.Dir, src)) {
			debug.Debugf("%s is older than %s", pkgpath(pkg), filepath.Join(pkg.Dir, src))
			return true
		}
	}

	return false
}

func stringList(args ...[]string) []string {
	var l []string
	for _, arg := range args {
		l = append(l, arg...)
	}
	return l
}
<commit_msg>Always assume the stdlib is up to date if present<commit_after>package gb

import (
	"os"
	"path/filepath"
	"runtime"
	"time"

	"github.com/constabulary/gb/debug"
)

// pkgpath returns the destination for object cached for this Package.
func pkgpath(pkg *Package) string {
	importpath := filepath.FromSlash(pkg.ImportPath) + ".a"
	switch {
	case pkg.isCrossCompile():
		return filepath.Join(pkg.Pkgdir(), importpath)
	case pkg.Standard && pkg.race:
		// race enabled standard lib
		return filepath.Join(runtime.GOROOT(), "pkg", pkg.gotargetos+"_"+pkg.gotargetarch+"_race", importpath)
	case pkg.Standard:
		// standard lib
		return filepath.Join(runtime.GOROOT(), "pkg", pkg.gotargetos+"_"+pkg.gotargetarch, importpath)
	default:
		return filepath.Join(pkg.Pkgdir(), importpath)
	}
}

// installpath returns the distination to cache this package's compiled .a file.
// pkgpath and installpath differ in that the former returns the location where you will find
// a previously cached .a file, the latter returns the location where an installed file
// will be placed.
//
// The difference is subtle. pkgpath must deal with the possibility that the file is from the
// standard library and is previously compiled. installpath will always return a path for the
// project's pkg/ directory in the case that the stdlib is out of date, or not compiled for
// a specific architecture.
func installpath(pkg *Package) string {
	if pkg.Scope == "test" {
		panic("installpath called with test scope")
	}
	return filepath.Join(pkg.Pkgdir(), filepath.FromSlash(pkg.ImportPath)+".a")
}

// isStale returns true if the source pkg is considered to be stale with
// respect to its installed version.
func isStale(pkg *Package) bool {
	switch pkg.ImportPath {
	case "C", "unsafe":
		// synthetic packages are never stale
		return false
	}

	if !pkg.Standard && pkg.Force {
		return true
	}

	// tests are always stale, they are never installed
	if pkg.Scope == "test" {
		return true
	}

	// Package is stale if completely unbuilt.
	var built time.Time
	if fi, err := os.Stat(pkgpath(pkg)); err == nil {
		built = fi.ModTime()
	}

	if built.IsZero() {
		debug.Debugf("%s is missing", pkgpath(pkg))
		return true
	}

	olderThan := func(file string) bool {
		fi, err := os.Stat(file)
		return err != nil || fi.ModTime().After(built)
	}

	newerThan := func(file string) bool {
		fi, err := os.Stat(file)
		return err != nil || fi.ModTime().Before(built)
	}

	// As a courtesy to developers installing new versions of the compiler
	// frequently, define that packages are stale if they are
	// older than the compiler, and commands if they are older than
	// the linker.  This heuristic will not work if the binaries are
	// back-dated, as some binary distributions may do, but it does handle
	// a very common case.
	if !pkg.Standard {
		if olderThan(pkg.tc.compiler()) {
			debug.Debugf("%s is older than %s", pkgpath(pkg), pkg.tc.compiler())
			return true
		}
		if pkg.IsCommand() && olderThan(pkg.tc.linker()) {
			debug.Debugf("%s is older than %s", pkgpath(pkg), pkg.tc.compiler())
			return true
		}
	}

	// Package is stale if a dependency is newer.
	for _, p := range pkg.Imports() {
		if p.ImportPath == "C" || p.ImportPath == "unsafe" {
			continue // ignore stale imports of synthetic packages
		}
		if olderThan(pkgpath(p)) {
			debug.Debugf("%s is older than %s", pkgpath(pkg), pkgpath(p))
			return true
		}
	}

	// if the main package is up to date but _newer_ than the binary (which
	// could have been removed), then consider it stale.
	if pkg.isMain() && newerThan(pkg.Binfile()) {
		debug.Debugf("%s is newer than %s", pkgpath(pkg), pkg.Binfile())
		return true
	}

	if pkg.Standard && !pkg.isCrossCompile() {
		// if this is a standard lib package, and we are not cross compiling
		// then assume the package is up to date. This also works around
		// golang/go#13769.
		return false
	}

	srcs := stringList(pkg.GoFiles, pkg.CFiles, pkg.CXXFiles, pkg.MFiles, pkg.HFiles, pkg.SFiles, pkg.CgoFiles, pkg.SysoFiles, pkg.SwigFiles, pkg.SwigCXXFiles)

	for _, src := range srcs {
		if olderThan(filepath.Join(pkg.Dir, src)) {
			debug.Debugf("%s is older than %s", pkgpath(pkg), filepath.Join(pkg.Dir, src))
			return true
		}
	}

	return false
}

func stringList(args ...[]string) []string {
	var l []string
	for _, arg := range args {
		l = append(l, arg...)
	}
	return l
}
<|endoftext|>


<commit_before>// Copyright 2019 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package main

import (
	"context"
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"time"

	"flag"
	"github.com/google/cabbie/notification"
	"github.com/google/cabbie/cablib"
	"github.com/google/cabbie/download"
	"github.com/google/cabbie/install"
	"github.com/google/cabbie/search"
	"github.com/google/cabbie/session"
	"github.com/google/cabbie/updatecollection"
	"github.com/google/subcommands"
)

// Available flags
type installCmd struct {
	drivers, deadlineOnly, virusDef bool
	kbs                             string
}

type installRsp struct {
	hResult        string
	resultCode     int
	rebootRequired bool
}

func (installCmd) Name() string     { return "install" }
func (installCmd) Synopsis() string { return "Install selected available updates." }
func (installCmd) Usage() string {
	return fmt.Sprintf("%s install [--drivers | --virusDef | --kbs=\"<KBNumber>\"]\n", filepath.Base(os.Args[0]))
}

func (i *installCmd) SetFlags(f *flag.FlagSet) {
	f.BoolVar(&i.drivers, "drivers", false, "Install available drivers.")
	f.BoolVar(&i.virusDef, "virus_def", false, "Update virus definitions.")
	f.BoolVar(&i.deadlineOnly, "deadlineOnly", false, fmt.Sprintf("Install available updates older than %d days", config.Deadline))
	f.StringVar(&i.kbs, "kbs", "", "Comma separated string of KB numbers in the form of 1234567.")
}

func (i installCmd) Execute(_ context.Context, flags *flag.FlagSet, _ ...interface{}) subcommands.ExitStatus {
	// TODO: Fix logic to allow only 0 to 1 flags at a time.
	if i.drivers && i.virusDef && i.kbs != "" {
		fmt.Println("drivers and virus_def flags can not be passed at the same time.")
		fmt.Printf("%s\nUsage: %s\n", i.Synopsis(), i.Usage())
		return subcommands.ExitUsageError
	}

	if err := i.installUpdates(); err != nil {
		fmt.Printf("Failed to install updates: %v", err)
		elog.Error(113, fmt.Sprintf("Failed to install updates: %v", err))
		return subcommands.ExitFailure
	}

	select {
	case <-rebootEvent:
		fmt.Println("Please reboot to finalize the update installation.")
		return 6
	default:
		fmt.Println("No reboot needed.")
	}

	return subcommands.ExitSuccess
}

func (i *installCmd) criteria() (string, []string) {
	// Set search criteria and required categories.
	var c string
	var rc []string
	switch {
	case i.drivers:
		c = "Type='Driver'"
		rc = append(rc, "Drivers")
		elog.Info(0021, fmt.Sprintf("Starting search for updated drivers: %s", c))
	case i.virusDef:
		c = fmt.Sprintf("%s AND CategoryIDs contains '%s'", search.BasicSearch, search.DefinitionUpdates)
		rc = append(rc, "Definition Updates")
		elog.Info(0022, fmt.Sprintf("Starting search for virus definitions:\n%s", c))
	case i.kbs != "":
		c = search.BasicSearch
		elog.Info(0023, fmt.Sprintf("Starting search for KB's %q:\n%s", i.kbs, c))
	default:
		c = search.BasicSearch
		rc = config.RequiredCategories
		elog.Info(0024, fmt.Sprintf("Starting search for general updates: %s", c))
	}
	return c, rc
}

func installingMessage() {
	elog.Info(2, "Cabbie is installing new updates.")

	if err := notification.NewNotification(cablib.SvcName, notification.NewInstallingMessage(), "installingUpdates"); err != nil {
		elog.Error(6, fmt.Sprintf("Failed to create notification:\n%v", err))
	}
}

func rebootMessage(seconds int) {
	elog.Info(2, "Updates have been installed, please reboot to complete the installation...")

	if err := notification.NewNotification(cablib.SvcName, notification.NewRebootMessage(seconds), "rebootPending"); err != nil {
		elog.Error(6, fmt.Sprintf("Failed to create notification:\n%v", err))
	}
}

func downloadCollection(s *session.UpdateSession, c *updatecollection.Collection) (int, error) {
	d, err := download.NewDownloader(s, c)
	if err != nil {
		return 0, fmt.Errorf("error creating downloader:\n %v", err)
	}
	defer d.Close()

	if err := d.Download(); err != nil {
		return 0, fmt.Errorf("error downloading updates:\n %v", err)
	}

	return d.ResultCode()
}

func installCollection(s *session.UpdateSession, c *updatecollection.Collection) (*installRsp, error) {
	inst, err := install.NewInstaller(s, c)
	if err != nil {
		return nil, fmt.Errorf("error creating installer: \n %v", err)
	}
	defer inst.Close()

	if err := inst.Install(); err != nil {
		return nil, fmt.Errorf("error installing updates:\n %v", err)
	}

	rc, err := inst.ResultCode()
	if err != nil {
		return nil, fmt.Errorf("error getting install ResultCode:\n %v", err)
	}

	hr, err := inst.HResult()
	if err != nil {
		return nil, fmt.Errorf("error getting install ReturnCode:\n %v", err)
	}

	rb, err := inst.RebootRequired()

	return &installRsp{
		hResult:        hr,
		resultCode:     rc,
		rebootRequired: rb,
	}, err
}

func (i *installCmd) installUpdates() error {
	var rebootRequired bool
	// Check for reboot status when not installing virus definitions.
	if !(i.virusDef) {
		rebootRequired, err := cablib.RebootRequired()
		if err != nil {
			return fmt.Errorf("failed to determine reboot status: %v", err)
		}

		if rebootRequired {
			rebootEvent <- rebootRequired
			return nil
		}
	}

	// Start Windows update session
	s, err := session.New()
	if err != nil {
		return fmt.Errorf("failed to create new Windows Update session: %v", err)
	}
	defer s.Close()

	criteria, rc := i.criteria()

	q, err := search.NewSearcher(s, criteria, config.WSUSServers, config.EnableThirdParty)
	if err != nil {
		return fmt.Errorf("failed to create a new searcher object: %v", err)
	}
	defer q.Close()

	uc, err := q.QueryUpdates()
	if er := searchHResult.Set(q.SearchHResult); er != nil {
		elog.Error(206, fmt.Sprintf("Error posting metric:\n%v", er))
	}
	if err != nil {
		return fmt.Errorf("error encountered when attempting to query for updates: %v", err)
	}
	defer uc.Close()

	if len(uc.Updates) == 0 {
		elog.Info(002, "No updates found to install.")
		return nil
	}
	elog.Info(4, fmt.Sprintf("Updates Found:\n%s", strings.Join(uc.Titles(), "\n\n")))

	installMsgPopped := i.virusDef

	kbs := NewKBSet(i.kbs)
	for _, u := range uc.Updates {
		if !(u.InCategories(rc)) {
			elog.Info(1, fmt.Sprintf("Skipping update %s.\nRequiredClassifications:\n%v\nUpdate classifications:\n%v",
				u.Title,
				rc,
				u.Categories))
			continue
		}

		if !(u.EulaAccepted) {
			elog.Info(002, fmt.Sprintf("Accepting EULA for update: %s", u.Title))
			if err := u.AcceptEula(); err != nil {
				elog.Error(202, fmt.Sprintf("Failed to accept EULA for update %s:\n%s", u.Title, err))
			}
		}

		if kbs.Size() > 0 {
			if !kbs.Search(u.KBArticleIDs) {
				elog.Info(1, fmt.Sprintf("Skipping update %s.\nRequired KBs:\n%s\nUpdate KBs:\n%v",
					u.Title,
					kbs,
					u.KBArticleIDs))
				continue
			}
		}
		if i.deadlineOnly {
			deadline := time.Duration(config.Deadline) * 24 * time.Hour
			pastDeadline := time.Now().After(u.LastDeploymentChangeTime.Add(deadline))
			if !pastDeadline {
				elog.Info(002,
					fmt.Sprintf("Skipping update %s.\nUpdate deployed on %v has not reached the %d day threshold.",
						u.Title,
						u.LastDeploymentChangeTime,
						config.Deadline))
				continue
			}
		}

		c, err := updatecollection.New()
		if err != nil {
			elog.Error(202, fmt.Sprintf("Failed to create collection: %v", err))
			continue
		}
		c.Add(u.Item)

		if !installMsgPopped && !u.InCategories([]string{"Definition Updates"}) {
			installingMessage()
			installMsgPopped = true
		}
		elog.Info(002, fmt.Sprintf("Downloading Update:\n%v", u))

		rc, err := downloadCollection(s, c)
		if err != nil {
			elog.Error(203, fmt.Sprintf("%v", err))
			c.Close()
			continue
		}
		if rc == 2 {
			elog.Info(002, fmt.Sprintf("Successfully downloaded update:\n %s", u.Title))
		} else {

			elog.Error(204, fmt.Sprintf("Failed to download update:\n %s\n ReturnCode: %d", u.Title, rc))
			c.Close()
			continue
		}

		elog.Info(002, fmt.Sprintf("Installing Update:\n%v", u))

		rsp, err := installCollection(s, c)
		if err != nil {
			elog.Error(205, fmt.Sprintf("%v", err))
			c.Close()
			continue
		}

		if err := installHResult.Set(rsp.hResult); err != nil {
			elog.Error(206, fmt.Sprintf("Error posting metric:\n%v", err))
		}
		if rsp.resultCode == 2 {
			elog.Info(002, fmt.Sprintf("Successfully installed update:\n%s\nHResult Code: %s", u.Title, rsp.hResult))
		} else {
			elog.Error(206, fmt.Sprintf("Failed to install update:\n%s\nReturnCode: %d\nHResult Code: %s", u.Title, rsp.resultCode, rsp.hResult))
			c.Close()
			continue
		}

		elog.Info(002, fmt.Sprintf("Install Reboot Required: %t", rsp.rebootRequired))
		if !rebootRequired {
			rebootRequired = rsp.rebootRequired
		}
		c.Close()
	}

	if rebootRequired {
		rebootMessage(int(config.RebootDelay))
		if err := cablib.SetRebootTime(config.RebootDelay); err != nil {
			elog.Error(306, fmt.Sprintf("Failed to run reboot command:\n%v", err))
		}
		rebootEvent <- rebootRequired
	}

	return nil
}
<commit_msg>Set forced reboot time if updates were installed outside of Cabbie.<commit_after>// Copyright 2019 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package main

import (
	"context"
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"time"

	"flag"
	"github.com/google/cabbie/notification"
	"github.com/google/cabbie/cablib"
	"github.com/google/cabbie/download"
	"github.com/google/cabbie/install"
	"github.com/google/cabbie/search"
	"github.com/google/cabbie/session"
	"github.com/google/cabbie/updatecollection"
	"github.com/google/subcommands"
)

// Available flags
type installCmd struct {
	drivers, deadlineOnly, virusDef bool
	kbs                             string
}

type installRsp struct {
	hResult        string
	resultCode     int
	rebootRequired bool
}

func (installCmd) Name() string     { return "install" }
func (installCmd) Synopsis() string { return "Install selected available updates." }
func (installCmd) Usage() string {
	return fmt.Sprintf("%s install [--drivers | --virusDef | --kbs=\"<KBNumber>\"]\n", filepath.Base(os.Args[0]))
}

func (i *installCmd) SetFlags(f *flag.FlagSet) {
	f.BoolVar(&i.drivers, "drivers", false, "Install available drivers.")
	f.BoolVar(&i.virusDef, "virus_def", false, "Update virus definitions.")
	f.BoolVar(&i.deadlineOnly, "deadlineOnly", false, fmt.Sprintf("Install available updates older than %d days", config.Deadline))
	f.StringVar(&i.kbs, "kbs", "", "Comma separated string of KB numbers in the form of 1234567.")
}

func (i installCmd) Execute(_ context.Context, flags *flag.FlagSet, _ ...interface{}) subcommands.ExitStatus {
	// TODO: Fix logic to allow only 0 to 1 flags at a time.
	if i.drivers && i.virusDef && i.kbs != "" {
		fmt.Println("drivers and virus_def flags can not be passed at the same time.")
		fmt.Printf("%s\nUsage: %s\n", i.Synopsis(), i.Usage())
		return subcommands.ExitUsageError
	}

	if err := i.installUpdates(); err != nil {
		fmt.Printf("Failed to install updates: %v", err)
		elog.Error(113, fmt.Sprintf("Failed to install updates: %v", err))
		return subcommands.ExitFailure
	}

	select {
	case <-rebootEvent:
		fmt.Println("Please reboot to finalize the update installation.")
		return 6
	default:
		fmt.Println("No reboot needed.")
	}

	return subcommands.ExitSuccess
}

func (i *installCmd) criteria() (string, []string) {
	// Set search criteria and required categories.
	var c string
	var rc []string
	switch {
	case i.drivers:
		c = "Type='Driver'"
		rc = append(rc, "Drivers")
		elog.Info(0021, fmt.Sprintf("Starting search for updated drivers: %s", c))
	case i.virusDef:
		c = fmt.Sprintf("%s AND CategoryIDs contains '%s'", search.BasicSearch, search.DefinitionUpdates)
		rc = append(rc, "Definition Updates")
		elog.Info(0022, fmt.Sprintf("Starting search for virus definitions:\n%s", c))
	case i.kbs != "":
		c = search.BasicSearch
		elog.Info(0023, fmt.Sprintf("Starting search for KB's %q:\n%s", i.kbs, c))
	default:
		c = search.BasicSearch
		rc = config.RequiredCategories
		elog.Info(0024, fmt.Sprintf("Starting search for general updates: %s", c))
	}
	return c, rc
}

func installingMessage() {
	elog.Info(2, "Cabbie is installing new updates.")

	if err := notification.NewNotification(cablib.SvcName, notification.NewInstallingMessage(), "installingUpdates"); err != nil {
		elog.Error(6, fmt.Sprintf("Failed to create notification:\n%v", err))
	}
}

func rebootMessage(seconds int) {
	elog.Info(2, "Updates have been installed, please reboot to complete the installation...")

	if err := notification.NewNotification(cablib.SvcName, notification.NewRebootMessage(seconds), "rebootPending"); err != nil {
		elog.Error(6, fmt.Sprintf("Failed to create notification:\n%v", err))
	}
}

func downloadCollection(s *session.UpdateSession, c *updatecollection.Collection) (int, error) {
	d, err := download.NewDownloader(s, c)
	if err != nil {
		return 0, fmt.Errorf("error creating downloader:\n %v", err)
	}
	defer d.Close()

	if err := d.Download(); err != nil {
		return 0, fmt.Errorf("error downloading updates:\n %v", err)
	}

	return d.ResultCode()
}

func installCollection(s *session.UpdateSession, c *updatecollection.Collection) (*installRsp, error) {
	inst, err := install.NewInstaller(s, c)
	if err != nil {
		return nil, fmt.Errorf("error creating installer: \n %v", err)
	}
	defer inst.Close()

	if err := inst.Install(); err != nil {
		return nil, fmt.Errorf("error installing updates:\n %v", err)
	}

	rc, err := inst.ResultCode()
	if err != nil {
		return nil, fmt.Errorf("error getting install ResultCode:\n %v", err)
	}

	hr, err := inst.HResult()
	if err != nil {
		return nil, fmt.Errorf("error getting install ReturnCode:\n %v", err)
	}

	rb, err := inst.RebootRequired()

	return &installRsp{
		hResult:        hr,
		resultCode:     rc,
		rebootRequired: rb,
	}, err
}

func (i *installCmd) installUpdates() error {
	var rebootRequired bool
	// Check for reboot status when not installing virus definitions.
	if !(i.virusDef) {
		rebootRequired, err := cablib.RebootRequired()
		if err != nil {
			return fmt.Errorf("failed to determine reboot status: %v", err)
		}

		if rebootRequired {
			t, err := cablib.RebootTime()
			if err != nil {
				return fmt.Errorf("Error getting reboot time: %v", err)
			}
			if t.IsZero() {
				// Set reboot time if a reboot is pending but no time has been set.
				// This can happen when a user installs updates outside of Cabbie.
				rebootMessage(int(config.RebootDelay))
				if err := cablib.SetRebootTime(config.RebootDelay); err != nil {
					return fmt.Errorf("Failed to set reboot time:\n%v", err)
				}
			}
			rebootEvent <- rebootRequired
			return nil
		}
	}

	// Start Windows update session
	s, err := session.New()
	if err != nil {
		return fmt.Errorf("failed to create new Windows Update session: %v", err)
	}
	defer s.Close()

	criteria, rc := i.criteria()

	q, err := search.NewSearcher(s, criteria, config.WSUSServers, config.EnableThirdParty)
	if err != nil {
		return fmt.Errorf("failed to create a new searcher object: %v", err)
	}
	defer q.Close()

	uc, err := q.QueryUpdates()
	if er := searchHResult.Set(q.SearchHResult); er != nil {
		elog.Error(206, fmt.Sprintf("Error posting metric:\n%v", er))
	}
	if err != nil {
		return fmt.Errorf("error encountered when attempting to query for updates: %v", err)
	}
	defer uc.Close()

	if len(uc.Updates) == 0 {
		elog.Info(002, "No updates found to install.")
		return nil
	}
	elog.Info(4, fmt.Sprintf("Updates Found:\n%s", strings.Join(uc.Titles(), "\n\n")))

	installMsgPopped := i.virusDef

	kbs := NewKBSet(i.kbs)
	for _, u := range uc.Updates {
		if !(u.InCategories(rc)) {
			elog.Info(1, fmt.Sprintf("Skipping update %s.\nRequiredClassifications:\n%v\nUpdate classifications:\n%v",
				u.Title,
				rc,
				u.Categories))
			continue
		}

		if !(u.EulaAccepted) {
			elog.Info(002, fmt.Sprintf("Accepting EULA for update: %s", u.Title))
			if err := u.AcceptEula(); err != nil {
				elog.Error(202, fmt.Sprintf("Failed to accept EULA for update %s:\n%s", u.Title, err))
			}
		}

		if kbs.Size() > 0 {
			if !kbs.Search(u.KBArticleIDs) {
				elog.Info(1, fmt.Sprintf("Skipping update %s.\nRequired KBs:\n%s\nUpdate KBs:\n%v",
					u.Title,
					kbs,
					u.KBArticleIDs))
				continue
			}
		}
		if i.deadlineOnly {
			deadline := time.Duration(config.Deadline) * 24 * time.Hour
			pastDeadline := time.Now().After(u.LastDeploymentChangeTime.Add(deadline))
			if !pastDeadline {
				elog.Info(002,
					fmt.Sprintf("Skipping update %s.\nUpdate deployed on %v has not reached the %d day threshold.",
						u.Title,
						u.LastDeploymentChangeTime,
						config.Deadline))
				continue
			}
		}

		c, err := updatecollection.New()
		if err != nil {
			elog.Error(202, fmt.Sprintf("Failed to create collection: %v", err))
			continue
		}
		c.Add(u.Item)

		if !installMsgPopped && !u.InCategories([]string{"Definition Updates"}) {
			installingMessage()
			installMsgPopped = true
		}
		elog.Info(002, fmt.Sprintf("Downloading Update:\n%v", u))

		rc, err := downloadCollection(s, c)
		if err != nil {
			elog.Error(203, fmt.Sprintf("%v", err))
			c.Close()
			continue
		}
		if rc == 2 {
			elog.Info(002, fmt.Sprintf("Successfully downloaded update:\n %s", u.Title))
		} else {

			elog.Error(204, fmt.Sprintf("Failed to download update:\n %s\n ReturnCode: %d", u.Title, rc))
			c.Close()
			continue
		}

		elog.Info(002, fmt.Sprintf("Installing Update:\n%v", u))

		rsp, err := installCollection(s, c)
		if err != nil {
			elog.Error(205, fmt.Sprintf("%v", err))
			c.Close()
			continue
		}

		if err := installHResult.Set(rsp.hResult); err != nil {
			elog.Error(206, fmt.Sprintf("Error posting metric:\n%v", err))
		}
		if rsp.resultCode == 2 {
			elog.Info(002, fmt.Sprintf("Successfully installed update:\n%s\nHResult Code: %s", u.Title, rsp.hResult))
		} else {
			elog.Error(206, fmt.Sprintf("Failed to install update:\n%s\nReturnCode: %d\nHResult Code: %s", u.Title, rsp.resultCode, rsp.hResult))
			c.Close()
			continue
		}

		elog.Info(002, fmt.Sprintf("Install Reboot Required: %t", rsp.rebootRequired))
		if !rebootRequired {
			rebootRequired = rsp.rebootRequired
		}
		c.Close()
	}

	if rebootRequired {
		rebootMessage(int(config.RebootDelay))
		if err := cablib.SetRebootTime(config.RebootDelay); err != nil {
			elog.Error(306, fmt.Sprintf("Failed to run reboot command:\n%v", err))
		}
		rebootEvent <- rebootRequired
	}

	return nil
}
<|endoftext|>


<commit_before>/*
Copyright 2017 The Kubernetes 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 plugins

import (
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"

	"github.com/ghodss/yaml"
	"github.com/golang/glog"

	"k8s.io/client-go/tools/clientcmd"
)

const PluginDescriptorFilename = "plugin.yaml"

// PluginLoader is capable of loading a list of plugin descriptions.
type PluginLoader interface {
	Load() (Plugins, error)
}

// DirectoryPluginLoader is a PluginLoader that loads plugin descriptions
// from a given directory in the filesystem. Plugins are located in subdirs
// under the loader "root", where each subdir must contain, at least, a plugin
// descriptor file called "plugin.yaml" that translates into a PluginDescription.
type DirectoryPluginLoader struct {
	Directory string
}

// Load reads the directory the loader holds and loads plugin descriptions.
func (l *DirectoryPluginLoader) Load() (Plugins, error) {
	if len(l.Directory) == 0 {
		return nil, fmt.Errorf("directory not specified")
	}

	list := Plugins{}

	stat, err := os.Stat(l.Directory)
	if err != nil {
		return nil, err
	}
	if !stat.IsDir() {
		return nil, fmt.Errorf("not a directory: %s", l.Directory)
	}

	base, err := filepath.Abs(l.Directory)
	if err != nil {
		return nil, err
	}

	// read the base directory tree searching for plugin descriptors
	// fails silently (descriptors unable to be read or unmarshalled are logged but skipped)
	err = filepath.Walk(base, func(path string, fileInfo os.FileInfo, walkErr error) error {
		if walkErr != nil || fileInfo.IsDir() || fileInfo.Name() != PluginDescriptorFilename {
			return nil
		}

		file, err := ioutil.ReadFile(path)
		if err != nil {
			glog.V(1).Infof("Unable to read plugin descriptor %s: %v", path, err)
			return nil
		}

		plugin := &Plugin{}
		if err := yaml.Unmarshal(file, plugin); err != nil {
			glog.V(1).Infof("Unable to unmarshal plugin descriptor %s: %v", path, err)
			return nil
		}

		if err := plugin.Validate(); err != nil {
			glog.V(1).Infof("%v", err)
			return nil
		}

		plugin.Dir = filepath.Dir(path)
		plugin.DescriptorName = fileInfo.Name()

		glog.V(6).Infof("Plugin loaded: %s", plugin.Name)
		list = append(list, plugin)

		return nil
	})

	return list, err
}

// UserDirPluginLoader is a PluginLoader that loads plugins from the
// "plugins" directory under the user's kubeconfig dir (usually "~/.kube/plugins/").
func UserDirPluginLoader() PluginLoader {
	dir := filepath.Join(clientcmd.RecommendedConfigDir, "plugins")
	return &DirectoryPluginLoader{
		Directory: dir,
	}
}

// PathFromEnvVarPluginLoader is a PluginLoader that loads plugins from one or more
// directories specified by the provided env var name. In case the env var is not
// set, the PluginLoader just loads nothing. A list of subdirectories can be provided,
// which will be appended to each path specified by the env var.
func PathFromEnvVarPluginLoader(envVarName string, subdirs ...string) PluginLoader {
	env := os.Getenv(envVarName)
	if len(env) == 0 {
		return &DummyPluginLoader{}
	}
	loader := MultiPluginLoader{}
	for _, path := range filepath.SplitList(env) {
		dir := append([]string{path}, subdirs...)
		loader = append(loader, &DirectoryPluginLoader{
			Directory: filepath.Join(dir...),
		})
	}
	return loader
}

// PluginsEnvVarPluginLoader is a PluginLoader that loads plugins from one or more
// directories specified by the KUBECTL_PLUGINS_PATH env var.
func PluginsEnvVarPluginLoader() PluginLoader {
	return PathFromEnvVarPluginLoader("KUBECTL_PLUGINS_PATH")
}

// XDGDataPluginLoader is a PluginLoader that loads plugins from one or more
// directories specified by the XDG system directory structure spec in the
// XDG_DATA_DIRS env var, plus the "kubectl/plugins/" suffix. According to the
// spec, if XDG_DATA_DIRS is not set it defaults to "/usr/local/share:/usr/share".
func XDGDataPluginLoader() PluginLoader {
	envVarName := "XDG_DATA_DIRS"
	if len(os.Getenv(envVarName)) > 0 {
		return PathFromEnvVarPluginLoader(envVarName, "kubectl", "plugins")
	}
	return TolerantMultiPluginLoader{
		&DirectoryPluginLoader{
			Directory: "/usr/local/share",
		},
		&DirectoryPluginLoader{
			Directory: "/usr/share",
		},
	}
}

// MultiPluginLoader is a PluginLoader that can encapsulate multiple plugin loaders,
// a successful loading means every encapsulated loader was able to load without errors.
type MultiPluginLoader []PluginLoader

func (l MultiPluginLoader) Load() (Plugins, error) {
	plugins := Plugins{}
	for _, loader := range l {
		loaded, err := loader.Load()
		if err != nil {
			return nil, err
		}
		plugins = append(plugins, loaded...)
	}
	return plugins, nil
}

// TolerantMultiPluginLoader is a PluginLoader than encapsulates multiple plugins loaders,
// but is tolerant to errors while loading from them.
type TolerantMultiPluginLoader []PluginLoader

func (l TolerantMultiPluginLoader) Load() (Plugins, error) {
	plugins := Plugins{}
	for _, loader := range l {
		loaded, _ := loader.Load()
		if loaded != nil {
			plugins = append(plugins, loaded...)
		}
	}
	return plugins, nil
}

// DummyPluginLoader loads nothing.
type DummyPluginLoader struct{}

func (l *DummyPluginLoader) Load() (Plugins, error) {
	return Plugins{}, nil
}
<commit_msg>Fix XDG-based kubectl plugin dirs<commit_after>/*
Copyright 2017 The Kubernetes 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 plugins

import (
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"

	"github.com/ghodss/yaml"
	"github.com/golang/glog"

	"k8s.io/client-go/tools/clientcmd"
)

const PluginDescriptorFilename = "plugin.yaml"

// PluginLoader is capable of loading a list of plugin descriptions.
type PluginLoader interface {
	Load() (Plugins, error)
}

// DirectoryPluginLoader is a PluginLoader that loads plugin descriptions
// from a given directory in the filesystem. Plugins are located in subdirs
// under the loader "root", where each subdir must contain, at least, a plugin
// descriptor file called "plugin.yaml" that translates into a PluginDescription.
type DirectoryPluginLoader struct {
	Directory string
}

// Load reads the directory the loader holds and loads plugin descriptions.
func (l *DirectoryPluginLoader) Load() (Plugins, error) {
	if len(l.Directory) == 0 {
		return nil, fmt.Errorf("directory not specified")
	}

	list := Plugins{}

	stat, err := os.Stat(l.Directory)
	if err != nil {
		return nil, err
	}
	if !stat.IsDir() {
		return nil, fmt.Errorf("not a directory: %s", l.Directory)
	}

	base, err := filepath.Abs(l.Directory)
	if err != nil {
		return nil, err
	}

	// read the base directory tree searching for plugin descriptors
	// fails silently (descriptors unable to be read or unmarshalled are logged but skipped)
	err = filepath.Walk(base, func(path string, fileInfo os.FileInfo, walkErr error) error {
		if walkErr != nil || fileInfo.IsDir() || fileInfo.Name() != PluginDescriptorFilename {
			return nil
		}

		file, err := ioutil.ReadFile(path)
		if err != nil {
			glog.V(1).Infof("Unable to read plugin descriptor %s: %v", path, err)
			return nil
		}

		plugin := &Plugin{}
		if err := yaml.Unmarshal(file, plugin); err != nil {
			glog.V(1).Infof("Unable to unmarshal plugin descriptor %s: %v", path, err)
			return nil
		}

		if err := plugin.Validate(); err != nil {
			glog.V(1).Infof("%v", err)
			return nil
		}

		plugin.Dir = filepath.Dir(path)
		plugin.DescriptorName = fileInfo.Name()

		glog.V(6).Infof("Plugin loaded: %s", plugin.Name)
		list = append(list, plugin)

		return nil
	})

	return list, err
}

// UserDirPluginLoader is a PluginLoader that loads plugins from the
// "plugins" directory under the user's kubeconfig dir (usually "~/.kube/plugins/").
func UserDirPluginLoader() PluginLoader {
	dir := filepath.Join(clientcmd.RecommendedConfigDir, "plugins")
	return &DirectoryPluginLoader{
		Directory: dir,
	}
}

// PathFromEnvVarPluginLoader is a PluginLoader that loads plugins from one or more
// directories specified by the provided env var name. In case the env var is not
// set, the PluginLoader just loads nothing. A list of subdirectories can be provided,
// which will be appended to each path specified by the env var.
func PathFromEnvVarPluginLoader(envVarName string, subdirs ...string) PluginLoader {
	env := os.Getenv(envVarName)
	if len(env) == 0 {
		return &DummyPluginLoader{}
	}
	loader := MultiPluginLoader{}
	for _, path := range filepath.SplitList(env) {
		dir := append([]string{path}, subdirs...)
		loader = append(loader, &DirectoryPluginLoader{
			Directory: filepath.Join(dir...),
		})
	}
	return loader
}

// PluginsEnvVarPluginLoader is a PluginLoader that loads plugins from one or more
// directories specified by the KUBECTL_PLUGINS_PATH env var.
func PluginsEnvVarPluginLoader() PluginLoader {
	return PathFromEnvVarPluginLoader("KUBECTL_PLUGINS_PATH")
}

// XDGDataPluginLoader is a PluginLoader that loads plugins from one or more
// directories specified by the XDG system directory structure spec in the
// XDG_DATA_DIRS env var, plus the "kubectl/plugins/" suffix. According to the
// spec, if XDG_DATA_DIRS is not set it defaults to "/usr/local/share:/usr/share".
func XDGDataPluginLoader() PluginLoader {
	envVarName := "XDG_DATA_DIRS"
	if len(os.Getenv(envVarName)) > 0 {
		return PathFromEnvVarPluginLoader(envVarName, "kubectl", "plugins")
	}
	return TolerantMultiPluginLoader{
		&DirectoryPluginLoader{
			Directory: "/usr/local/share/kubectl/plugins",
		},
		&DirectoryPluginLoader{
			Directory: "/usr/share/kubectl/plugins",
		},
	}
}

// MultiPluginLoader is a PluginLoader that can encapsulate multiple plugin loaders,
// a successful loading means every encapsulated loader was able to load without errors.
type MultiPluginLoader []PluginLoader

func (l MultiPluginLoader) Load() (Plugins, error) {
	plugins := Plugins{}
	for _, loader := range l {
		loaded, err := loader.Load()
		if err != nil {
			return nil, err
		}
		plugins = append(plugins, loaded...)
	}
	return plugins, nil
}

// TolerantMultiPluginLoader is a PluginLoader than encapsulates multiple plugins loaders,
// but is tolerant to errors while loading from them.
type TolerantMultiPluginLoader []PluginLoader

func (l TolerantMultiPluginLoader) Load() (Plugins, error) {
	plugins := Plugins{}
	for _, loader := range l {
		loaded, _ := loader.Load()
		if loaded != nil {
			plugins = append(plugins, loaded...)
		}
	}
	return plugins, nil
}

// DummyPluginLoader loads nothing.
type DummyPluginLoader struct{}

func (l *DummyPluginLoader) Load() (Plugins, error) {
	return Plugins{}, nil
}
<|endoftext|>