<commit_before>// Copyright 2018 The Terraformer 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 cmd

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

	"github.com/GoogleCloudPlatform/terraformer/terraformutils/terraformerstring"

	"github.com/GoogleCloudPlatform/terraformer/terraformutils/providerwrapper"

	"github.com/spf13/pflag"

	"github.com/GoogleCloudPlatform/terraformer/terraformutils"
	"github.com/GoogleCloudPlatform/terraformer/terraformutils/terraformoutput"

	"github.com/spf13/cobra"
)

type ImportOptions struct {
	Resources   []string
	PathPattern string
	PathOutput  string
	State       string
	Bucket      string
	Profile     string
	Verbose     bool
	Zone        string
	Regions     []string
	Projects    []string
	Connect     bool
	Compact     bool
	Filter      []string
	Plan        bool `json:"-"`
	Output      string
}

const DefaultPathPattern = "{output}/{provider}/{service}/"
const DefaultPathOutput = "generated"
const DefaultState = "local"

func newImportCmd() *cobra.Command {
	options := ImportOptions{}
	cmd := &cobra.Command{
		Use:           "import",
		Short:         "Import current state to Terraform configuration",
		Long:          "Import current state to Terraform configuration",
		SilenceUsage:  true,
		SilenceErrors: false,
		//Version:       version.String(),
	}

	cmd.AddCommand(newCmdPlanImporter(options))
	for _, subcommand := range providerImporterSubcommands() {
		providerCommand := subcommand(options)
		_ = providerCommand.MarkPersistentFlagRequired("resources")
		cmd.AddCommand(providerCommand)
	}
	return cmd
}

func Import(provider terraformutils.ProviderGenerator, options ImportOptions, args []string) error {
	err := provider.Init(args)
	if err != nil {
		return err
	}
	plan := &ImportPlan{
		Provider:         provider.GetName(),
		Options:          options,
		Args:             args,
		ImportedResource: map[string][]terraformutils.Resource{},
	}

	if terraformerstring.ContainsString(options.Resources, "*") {
		log.Println("Attempting an import of ALL resources in " + provider.GetName())
		options.Resources = providerServices(provider)
	}

	for _, service := range options.Resources {
		log.Println(provider.GetName() + " importing... " + service)
		err = provider.InitService(service, options.Verbose)
		if err != nil {
			return err
		}
		provider.GetService().ParseFilters(options.Filter)
		err = provider.GetService().InitResources()
		if err != nil {
			return err
		}

		providerWrapper, err := providerwrapper.NewProviderWrapper(provider.GetName(), provider.GetConfig(), options.Verbose)
		if err != nil {
			return err
		}

		provider.GetService().PopulateIgnoreKeys(providerWrapper)
		provider.GetService().InitialCleanup()

		refreshedResources, err := terraformutils.RefreshResources(provider.GetService().GetResources(), providerWrapper)
		if err != nil {
			return err
		}
		provider.GetService().SetResources(refreshedResources)

		for i := range provider.GetService().GetResources() {
			err = provider.GetService().GetResources()[i].ConvertTFstate(providerWrapper)
			if err != nil {
				return err
			}
		}

		providerWrapper.Kill()

		provider.GetService().PostRefreshCleanup()

		// change structs with additional data for each resource
		err = provider.GetService().PostConvertHook()
		if err != nil {
			return err
		}
		plan.ImportedResource[service] = append(plan.ImportedResource[service], provider.GetService().GetResources()...)
	}
	if options.Plan {
		path := Path(options.PathPattern, provider.GetName(), "terraformer", options.PathOutput)
		return ExportPlanFile(plan, path, "plan.json")
	}
	return ImportFromPlan(provider, plan)
}

func ImportFromPlan(provider terraformutils.ProviderGenerator, plan *ImportPlan) error {
	options := plan.Options
	importedResource := plan.ImportedResource
	isServicePath := strings.Contains(options.PathPattern, "{service}")

	if options.Connect {
		log.Println(provider.GetName() + " Connecting.... ")
		importedResource = terraformutils.ConnectServices(importedResource, isServicePath, provider.GetResourceConnections())
	}

	if !isServicePath {
		var compactedResources []terraformutils.Resource
		for _, resources := range importedResource {
			compactedResources = append(compactedResources, resources...)
		}
		e := printService(provider, "", options, compactedResources, importedResource)
		if e != nil {
			return e
		}
	} else {
		for serviceName, resources := range importedResource {
			e := printService(provider, serviceName, options, resources, importedResource)
			if e != nil {
				return e
			}
		}
	}
	return nil
}

func printService(provider terraformutils.ProviderGenerator, serviceName string, options ImportOptions, resources []terraformutils.Resource, importedResource map[string][]terraformutils.Resource) error {
	log.Println(provider.GetName() + " save " + serviceName)
	// Print HCL files for Resources
	path := Path(options.PathPattern, provider.GetName(), serviceName, options.PathOutput)
	err := terraformoutput.OutputHclFiles(resources, provider, path, serviceName, options.Compact, options.Output)
	if err != nil {
		return err
	}
	tfStateFile, err := terraformutils.PrintTfState(resources)
	if err != nil {
		return err
	}
	// print or upload State file
	if options.State == "bucket" {
		log.Println(provider.GetName() + " upload tfstate to  bucket " + options.Bucket)
		bucket := terraformoutput.BucketState{
			Name: options.Bucket,
		}
		if err := bucket.BucketUpload(path, tfStateFile); err != nil {
			return err
		}
		// create Bucket file
		if bucketStateDataFile, err := terraformutils.Print(bucket.BucketGetTfData(path), map[string]struct{}{}, options.Output); err == nil {
			terraformoutput.PrintFile(path+"/bucket.tf", bucketStateDataFile)
		}
	} else {
		if serviceName == "" {
			log.Println(provider.GetName() + " save tfstate")
		} else {
			log.Println(provider.GetName() + " save tfstate for " + serviceName)
		}
		if err := ioutil.WriteFile(path+"/terraform.tfstate", tfStateFile, os.ModePerm); err != nil {
			return err
		}
	}
	// Print hcl variables.tf
	if serviceName != "" {
		if options.Connect && len(provider.GetResourceConnections()[serviceName]) > 0 {
			variables := map[string]map[string]map[string]interface{}{}
			variables["data"] = map[string]map[string]interface{}{}
			variables["data"]["terraform_remote_state"] = map[string]interface{}{}
			if options.State == "bucket" {
				bucket := terraformoutput.BucketState{
					Name: options.Bucket,
				}
				for k := range provider.GetResourceConnections()[serviceName] {
					if _, exist := importedResource[k]; !exist {
						continue
					}
					variables["data"]["terraform_remote_state"][k] = map[string]interface{}{
						"backend": "gcs",
						"config":  bucket.BucketGetTfData(strings.Replace(path, serviceName, k, -1)),
					}
				}
			} else {
				for k := range provider.GetResourceConnections()[serviceName] {
					if _, exist := importedResource[k]; !exist {
						continue
					}
					variables["data"]["terraform_remote_state"][k] = map[string]interface{}{
						"backend": "local",
						"config": [1]interface{}{map[string]interface{}{
							"path": strings.Repeat("../", strings.Count(path, "/")) + strings.Replace(path, serviceName, k, -1) + "terraform.tfstate",
						}},
					}
				}
			}
			// create variables file
			if len(provider.GetResourceConnections()[serviceName]) > 0 && options.Connect && len(variables["data"]["terraform_remote_state"]) > 0 {
				variablesFile, err := terraformutils.Print(variables, map[string]struct{}{"config": {}}, options.Output)
				if err != nil {
					return err
				}
				terraformoutput.PrintFile(path+"/variables."+terraformoutput.GetFileExtension(options.Output), variablesFile)
			}
		}
	} else {
		if options.Connect {
			variables := map[string]map[string]map[string]interface{}{}
			variables["data"] = map[string]map[string]interface{}{}
			variables["data"]["terraform_remote_state"] = map[string]interface{}{}
			if options.State == "bucket" {
				bucket := terraformoutput.BucketState{
					Name: options.Bucket,
				}
				variables["data"]["terraform_remote_state"]["local"] = map[string]interface{}{
					"backend": "gcs",
					"config":  bucket.BucketGetTfData(path),
				}
			} else {
				variables["data"]["terraform_remote_state"]["local"] = map[string]interface{}{
					"backend": "local",
					"config": map[string]interface{}{
						"path": "terraform.tfstate",
					},
				}
			}
			// create variables file
			if options.Connect {
				variablesFile, err := terraformutils.Print(variables, map[string]struct{}{"config": {}}, options.Output)
				if err != nil {
					return err
				}
				terraformoutput.PrintFile(path+"/variables."+terraformoutput.GetFileExtension(options.Output), variablesFile)
			}
		}
	}
	return nil
}

func Path(pathPattern, providerName, serviceName, output string) string {
	return strings.NewReplacer(
		"{provider}", providerName,
		"{service}", serviceName,
		"{output}", output,
	).Replace(pathPattern)
}

func listCmd(provider terraformutils.ProviderGenerator) *cobra.Command {
	cmd := &cobra.Command{
		Use:   "list",
		Short: "List supported resources for " + provider.GetName() + " provider",
		Long:  "List supported resources for " + provider.GetName() + " provider",
		RunE: func(cmd *cobra.Command, args []string) error {
			services := providerServices(provider)
			for _, k := range services {
				fmt.Println(k)
			}
			return nil
		},
	}
	cmd.Flags().AddFlag(&pflag.Flag{Name: "resources"})
	return cmd
}

func providerServices(provider terraformutils.ProviderGenerator) []string {
	var services []string
	for k := range provider.GetSupportedService() {
		services = append(services, k)
	}
	sort.Strings(services)
	return services
}

func baseProviderFlags(flag *pflag.FlagSet, options *ImportOptions, sampleRes, sampleFilters string) {
	flag.BoolVarP(&options.Connect, "connect", "c", true, "")
	flag.BoolVarP(&options.Compact, "compact", "C", false, "")
	flag.StringSliceVarP(&options.Resources, "resources", "r", []string{}, sampleRes)
	flag.StringVarP(&options.PathPattern, "path-pattern", "p", DefaultPathPattern, "{output}/{provider}/")
	flag.StringVarP(&options.PathOutput, "path-output", "o", DefaultPathOutput, "")
	flag.StringVarP(&options.State, "state", "s", DefaultState, "local or bucket")
	flag.StringVarP(&options.Bucket, "bucket", "b", "", "gs://terraform-state")
	flag.StringSliceVarP(&options.Filter, "filter", "f", []string{}, sampleFilters)
	flag.BoolVarP(&options.Verbose, "verbose", "v", false, "")
	flag.StringVarP(&options.Output, "output", "O", "hcl", "output format hcl or json")
}
<commit_msg>add -x / --excludes flag<commit_after>// Copyright 2018 The Terraformer 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 cmd

import (
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"sort"
	"strings"
	"github.com/GoogleCloudPlatform/terraformer/terraformutils/terraformerstring"

	"github.com/GoogleCloudPlatform/terraformer/terraformutils/providerwrapper"

	"github.com/spf13/pflag"

	"github.com/GoogleCloudPlatform/terraformer/terraformutils"
	"github.com/GoogleCloudPlatform/terraformer/terraformutils/terraformoutput"

	"github.com/spf13/cobra"
)

type ImportOptions struct {
	Resources   []string
	Excludes    []string
	PathPattern string
	PathOutput  string
	State       string
	Bucket      string
	Profile     string
	Verbose     bool
	Zone        string
	Regions     []string
	Projects    []string
	Connect     bool
	Compact     bool
	Filter      []string
	Plan        bool `json:"-"`
	Output      string
}

const DefaultPathPattern = "{output}/{provider}/{service}/"
const DefaultPathOutput = "generated"
const DefaultState = "local"

func newImportCmd() *cobra.Command {
	options := ImportOptions{}
	cmd := &cobra.Command{
		Use:           "import",
		Short:         "Import current state to Terraform configuration",
		Long:          "Import current state to Terraform configuration",
		SilenceUsage:  true,
		SilenceErrors: false,
		//Version:       version.String(),
	}

	cmd.AddCommand(newCmdPlanImporter(options))
	for _, subcommand := range providerImporterSubcommands() {
		providerCommand := subcommand(options)
		_ = providerCommand.MarkPersistentFlagRequired("resources")
		cmd.AddCommand(providerCommand)
	}
	return cmd
}

func Import(provider terraformutils.ProviderGenerator, options ImportOptions, args []string) error {
	err := provider.Init(args)
	if err != nil {
		return err
	}
	plan := &ImportPlan{
		Provider:         provider.GetName(),
		Options:          options,
		Args:             args,
		ImportedResource: map[string][]terraformutils.Resource{},
	}

	if terraformerstring.ContainsString(options.Resources, "*") {
		log.Println("Attempting an import of ALL resources in " + provider.GetName())
		options.Resources = providerServices(provider)
	}

	if (options.Excludes != nil) {
		localSlice := []string{}
		for _, r := range options.Resources {
			remove := false
			for _, e := range options.Excludes {
				if (r == e) {
					remove = true
					log.Println("Excluding resource " + e)
				}
			}
			if remove != true {
				localSlice = append(localSlice, r)
			}
		}
		options.Resources = localSlice
	}

	for _, service := range options.Resources {
		log.Println(provider.GetName() + " importing... " + service)
		err = provider.InitService(service, options.Verbose)
		if err != nil {
			return err
		}
		provider.GetService().ParseFilters(options.Filter)
		err = provider.GetService().InitResources()
		if err != nil {
			return err
		}

		providerWrapper, err := providerwrapper.NewProviderWrapper(provider.GetName(), provider.GetConfig(), options.Verbose)
		if err != nil {
			return err
		}

		provider.GetService().PopulateIgnoreKeys(providerWrapper)
		provider.GetService().InitialCleanup()

		refreshedResources, err := terraformutils.RefreshResources(provider.GetService().GetResources(), providerWrapper)
		if err != nil {
			return err
		}
		provider.GetService().SetResources(refreshedResources)

		for i := range provider.GetService().GetResources() {
			err = provider.GetService().GetResources()[i].ConvertTFstate(providerWrapper)
			if err != nil {
				return err
			}
		}

		providerWrapper.Kill()

		provider.GetService().PostRefreshCleanup()

		// change structs with additional data for each resource
		err = provider.GetService().PostConvertHook()
		if err != nil {
			return err
		}
		plan.ImportedResource[service] = append(plan.ImportedResource[service], provider.GetService().GetResources()...)
	}
	if options.Plan {
		path := Path(options.PathPattern, provider.GetName(), "terraformer", options.PathOutput)
		return ExportPlanFile(plan, path, "plan.json")
	}
	return ImportFromPlan(provider, plan)
}

func ImportFromPlan(provider terraformutils.ProviderGenerator, plan *ImportPlan) error {
	options := plan.Options
	importedResource := plan.ImportedResource
	isServicePath := strings.Contains(options.PathPattern, "{service}")

	if options.Connect {
		log.Println(provider.GetName() + " Connecting.... ")
		importedResource = terraformutils.ConnectServices(importedResource, isServicePath, provider.GetResourceConnections())
	}

	if !isServicePath {
		var compactedResources []terraformutils.Resource
		for _, resources := range importedResource {
			compactedResources = append(compactedResources, resources...)
		}
		e := printService(provider, "", options, compactedResources, importedResource)
		if e != nil {
			return e
		}
	} else {
		for serviceName, resources := range importedResource {
			e := printService(provider, serviceName, options, resources, importedResource)
			if e != nil {
				return e
			}
		}
	}
	return nil
}

func printService(provider terraformutils.ProviderGenerator, serviceName string, options ImportOptions, resources []terraformutils.Resource, importedResource map[string][]terraformutils.Resource) error {
	log.Println(provider.GetName() + " save " + serviceName)
	// Print HCL files for Resources
	path := Path(options.PathPattern, provider.GetName(), serviceName, options.PathOutput)
	err := terraformoutput.OutputHclFiles(resources, provider, path, serviceName, options.Compact, options.Output)
	if err != nil {
		return err
	}
	tfStateFile, err := terraformutils.PrintTfState(resources)
	if err != nil {
		return err
	}
	// print or upload State file
	if options.State == "bucket" {
		log.Println(provider.GetName() + " upload tfstate to  bucket " + options.Bucket)
		bucket := terraformoutput.BucketState{
			Name: options.Bucket,
		}
		if err := bucket.BucketUpload(path, tfStateFile); err != nil {
			return err
		}
		// create Bucket file
		if bucketStateDataFile, err := terraformutils.Print(bucket.BucketGetTfData(path), map[string]struct{}{}, options.Output); err == nil {
			terraformoutput.PrintFile(path+"/bucket.tf", bucketStateDataFile)
		}
	} else {
		if serviceName == "" {
			log.Println(provider.GetName() + " save tfstate")
		} else {
			log.Println(provider.GetName() + " save tfstate for " + serviceName)
		}
		if err := ioutil.WriteFile(path+"/terraform.tfstate", tfStateFile, os.ModePerm); err != nil {
			return err
		}
	}
	// Print hcl variables.tf
	if serviceName != "" {
		if options.Connect && len(provider.GetResourceConnections()[serviceName]) > 0 {
			variables := map[string]map[string]map[string]interface{}{}
			variables["data"] = map[string]map[string]interface{}{}
			variables["data"]["terraform_remote_state"] = map[string]interface{}{}
			if options.State == "bucket" {
				bucket := terraformoutput.BucketState{
					Name: options.Bucket,
				}
				for k := range provider.GetResourceConnections()[serviceName] {
					if _, exist := importedResource[k]; !exist {
						continue
					}
					variables["data"]["terraform_remote_state"][k] = map[string]interface{}{
						"backend": "gcs",
						"config":  bucket.BucketGetTfData(strings.Replace(path, serviceName, k, -1)),
					}
				}
			} else {
				for k := range provider.GetResourceConnections()[serviceName] {
					if _, exist := importedResource[k]; !exist {
						continue
					}
					variables["data"]["terraform_remote_state"][k] = map[string]interface{}{
						"backend": "local",
						"config": [1]interface{}{map[string]interface{}{
							"path": strings.Repeat("../", strings.Count(path, "/")) + strings.Replace(path, serviceName, k, -1) + "terraform.tfstate",
						}},
					}
				}
			}
			// create variables file
			if len(provider.GetResourceConnections()[serviceName]) > 0 && options.Connect && len(variables["data"]["terraform_remote_state"]) > 0 {
				variablesFile, err := terraformutils.Print(variables, map[string]struct{}{"config": {}}, options.Output)
				if err != nil {
					return err
				}
				terraformoutput.PrintFile(path+"/variables."+terraformoutput.GetFileExtension(options.Output), variablesFile)
			}
		}
	} else {
		if options.Connect {
			variables := map[string]map[string]map[string]interface{}{}
			variables["data"] = map[string]map[string]interface{}{}
			variables["data"]["terraform_remote_state"] = map[string]interface{}{}
			if options.State == "bucket" {
				bucket := terraformoutput.BucketState{
					Name: options.Bucket,
				}
				variables["data"]["terraform_remote_state"]["local"] = map[string]interface{}{
					"backend": "gcs",
					"config":  bucket.BucketGetTfData(path),
				}
			} else {
				variables["data"]["terraform_remote_state"]["local"] = map[string]interface{}{
					"backend": "local",
					"config": map[string]interface{}{
						"path": "terraform.tfstate",
					},
				}
			}
			// create variables file
			if options.Connect {
				variablesFile, err := terraformutils.Print(variables, map[string]struct{}{"config": {}}, options.Output)
				if err != nil {
					return err
				}
				terraformoutput.PrintFile(path+"/variables."+terraformoutput.GetFileExtension(options.Output), variablesFile)
			}
		}
	}
	return nil
}

func Path(pathPattern, providerName, serviceName, output string) string {
	return strings.NewReplacer(
		"{provider}", providerName,
		"{service}", serviceName,
		"{output}", output,
	).Replace(pathPattern)
}

func listCmd(provider terraformutils.ProviderGenerator) *cobra.Command {
	cmd := &cobra.Command{
		Use:   "list",
		Short: "List supported resources for " + provider.GetName() + " provider",
		Long:  "List supported resources for " + provider.GetName() + " provider",
		RunE: func(cmd *cobra.Command, args []string) error {
			services := providerServices(provider)
			for _, k := range services {
				fmt.Println(k)
			}
			return nil
		},
	}
	cmd.Flags().AddFlag(&pflag.Flag{Name: "resources"})
	return cmd
}

func providerServices(provider terraformutils.ProviderGenerator) []string {
	var services []string
	for k := range provider.GetSupportedService() {
		services = append(services, k)
	}
	sort.Strings(services)
	return services
}

func baseProviderFlags(flag *pflag.FlagSet, options *ImportOptions, sampleRes, sampleFilters string) {
	flag.BoolVarP(&options.Connect, "connect", "c", true, "")
	flag.BoolVarP(&options.Compact, "compact", "C", false, "")
	flag.StringSliceVarP(&options.Resources, "resources", "r", []string{}, sampleRes)
	flag.StringSliceVarP(&options.Excludes, "excludes", "x", []string{}, sampleRes)
	flag.StringVarP(&options.PathPattern, "path-pattern", "p", DefaultPathPattern, "{output}/{provider}/")
	flag.StringVarP(&options.PathOutput, "path-output", "o", DefaultPathOutput, "")
	flag.StringVarP(&options.State, "state", "s", DefaultState, "local or bucket")
	flag.StringVarP(&options.Bucket, "bucket", "b", "", "gs://terraform-state")
	flag.StringSliceVarP(&options.Filter, "filter", "f", []string{}, sampleFilters)
	flag.BoolVarP(&options.Verbose, "verbose", "v", false, "")
	flag.StringVarP(&options.Output, "output", "O", "hcl", "output format hcl or json")
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"os"
	"strings"
	"uctricks"
)

func main() {
	fontname := flag.String("f", "blackletter", "Specify which 'font' to use. Use -l to see all.")
	list := flag.Bool("l", false, "List fonts and exit.")
	noNewline := flag.Bool("n", false, "Don't print trailing newline.")

	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "usage: unifun [options] text\n")
		flag.PrintDefaults()
	}

	flag.Parse()

	// print a list of fonts if asked
	if *list {
		for _, name := range uctricks.FontList() {
			fmt.Println(name)
		}
		return
	}

	// concatentate all arguments
	s := strings.Join(flag.Args(), " ")

	// output the result
	font, err := uctricks.FontNamed(*fontname)
	if err != nil {
		fmt.Fprintf(os.Stderr, "No such font: %v", *fontname)
		os.Exit(1)
	}
	fmt.Print(font.Apply(s))

	// print a neline if required
	if !*noNewline {
		fmt.Println("")
	}
}
<commit_msg>be a nice citizen and append a newline to the error message<commit_after>package main

import (
	"flag"
	"fmt"
	"os"
	"strings"
	"uctricks"
)

func main() {
	fontname := flag.String("f", "blackletter", "Specify which 'font' to use. Use -l to see all.")
	list := flag.Bool("l", false, "List fonts and exit.")
	noNewline := flag.Bool("n", false, "Don't print trailing newline.")

	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "usage: unifun [options] text\n")
		flag.PrintDefaults()
	}

	flag.Parse()

	// print a list of fonts if asked
	if *list {
		for _, name := range uctricks.FontList() {
			fmt.Println(name)
		}
		return
	}

	// concatentate all arguments
	s := strings.Join(flag.Args(), " ")

	// output the result
	font, err := uctricks.FontNamed(*fontname)
	if err != nil {
		fmt.Fprintf(os.Stderr, "No such font: %v\n", *fontname)
		os.Exit(1)
	}
	fmt.Print(font.Apply(s))

	// print a neline if required
	if !*noNewline {
		fmt.Println("")
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"sort"
	"strings"
)

// `direnv export $0`
var CmdExport = &Cmd{
	Name:    "export",
	Desc:    "loads an .envrc and prints the diff in terms of exports",
	Args:    []string{"SHELL"},
	Private: true,
	Fn: func(env Env, args []string) (err error) {
		var oldEnv Env = env.Filtered()
		var newEnv Env
		var loadedRC *RC
		var foundRC *RC
		var config *Config
		var target string

		if len(args) > 1 {
			target = args[1]
		}

		shell := DetectShell(target)
		if shell == nil {
			return fmt.Errorf("Unknown target shell '%s'", target)
		}

		if config, err = LoadConfig(env); err != nil {
			return
		}

		loadedRC = config.LoadedRC()
		foundRC = config.FindRC()

		//fmt.Fprintf(os.Stderr, "%v %v\n", loadedRC, foundRC)

		if loadedRC == nil {
			if foundRC == nil {
				// We're done here.
				return nil
			}

			newEnv, err = loadRC(foundRC, config, oldEnv)
		} else {
			var backupEnv Env
			if backupEnv, err = config.EnvBackup(); err != nil {
				err = fmt.Errorf("EnvBackup() failed: %q", err)
				goto error
			}
			oldEnv = backupEnv.Filtered()
			if foundRC == nil {
				fmt.Fprintf(os.Stderr, "direnv: unloading\n")
				newEnv = oldEnv
			} else if loadedRC.path != foundRC.path {
				fmt.Fprintf(os.Stderr, "direnv: switching\n")
				newEnv, err = loadRC(foundRC, config, oldEnv)
			} else if loadedRC.mtime != foundRC.mtime {
				fmt.Fprintf(os.Stderr, "direnv: reloading\n")
				newEnv, err = loadRC(foundRC, config, oldEnv)
			} else {
				// Nothing to do. Env is loaded and hasn't changed
				return nil
			}
		}

	error:
		if err != nil {
			newEnv = oldEnv
			if foundRC != nil {
				// This should be nearby rc.Load()'s similar statement
				newEnv["DIRENV_DIR"] = "-" + filepath.Dir(foundRC.path)
				newEnv["DIRENV_MTIME"] = fmt.Sprintf("%d", foundRC.mtime)
				newEnv["DIRENV_BACKUP"] = oldEnv.Serialize()
			}
		}

		diff := EnvDiff(env, newEnv)

		diff2 := diff.Filtered()
		if len(diff2) > 0 {
			out := make([]string, len(diff2))
			i := 0
			for key, value := range diff2 {
				if value == "" {
					out[i] = "-" + key
				} else if oldEnv[key] == "" {
					out[i] = "+" + key
				} else {
					out[i] = "~" + key
				}
				i += 1
			}
			sort.Strings(out)
			fmt.Fprintf(os.Stderr, "direnv export: %s\n", strings.Join(out, " "))
		}

		str := EnvToShell(diff, shell)

		fmt.Print(str)
		return

	},
}

const NOT_ALLOWED = "%s is blocked from loading. Run `direnv allow` to approve it's content for loading.\n"

func loadRC(rc *RC, config *Config, env Env) (newEnv Env, err error) {
	if !rc.Allowed() {
		return nil, fmt.Errorf(NOT_ALLOWED, rc.RelTo(config.WorkDir))
	}

	argtmpl := `eval "$("%s" stdlib)" >&2 && source_env "%s" >&2 && "%s" dump`
	arg := fmt.Sprintf(argtmpl, config.SelfPath, rc.RelTo(config.WorkDir), config.SelfPath)
	cmd := exec.Command(config.BashPath, "--noprofile", "--norc", "-c", arg)

	cmd.Stderr = os.Stderr
	cmd.Env = env.ToGoEnv()
	cmd.Dir = config.WorkDir

	out, err := cmd.Output()
	if err != nil {
		err = fmt.Errorf("loadRC() failed to run bash: %q", err)
		return
	}

	newEnv, err = ParseEnv(string(out))
	if err != nil {
		err = fmt.Errorf("loadRC() ParseEnv failed: %q", err)
		return
	}

	newEnv["DIRENV_DIR"] = "-" + filepath.Dir(rc.path)
	newEnv["DIRENV_MTIME"] = fmt.Sprintf("%d", rc.mtime)
	newEnv["DIRENV_BACKUP"] = env.Serialize()

	return
}
<commit_msg>Better error message when direnv fails loading the .envrc<commit_after>package main

import (
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"sort"
	"strings"
)

// `direnv export $0`
var CmdExport = &Cmd{
	Name:    "export",
	Desc:    "loads an .envrc and prints the diff in terms of exports",
	Args:    []string{"SHELL"},
	Private: true,
	Fn: func(env Env, args []string) (err error) {
		var oldEnv Env = env.Filtered()
		var newEnv Env
		var loadedRC *RC
		var foundRC *RC
		var config *Config
		var target string

		if len(args) > 1 {
			target = args[1]
		}

		shell := DetectShell(target)
		if shell == nil {
			return fmt.Errorf("Unknown target shell '%s'", target)
		}

		if config, err = LoadConfig(env); err != nil {
			return
		}

		loadedRC = config.LoadedRC()
		foundRC = config.FindRC()

		//fmt.Fprintf(os.Stderr, "%v %v\n", loadedRC, foundRC)

		if loadedRC == nil {
			if foundRC == nil {
				// We're done here.
				return nil
			}

			newEnv, err = loadRC(foundRC, config, oldEnv)
		} else {
			var backupEnv Env
			if backupEnv, err = config.EnvBackup(); err != nil {
				err = fmt.Errorf("EnvBackup() failed: %q", err)
				goto error
			}
			oldEnv = backupEnv.Filtered()
			if foundRC == nil {
				fmt.Fprintf(os.Stderr, "direnv: unloading\n")
				newEnv = oldEnv
			} else if loadedRC.path != foundRC.path {
				fmt.Fprintf(os.Stderr, "direnv: switching\n")
				newEnv, err = loadRC(foundRC, config, oldEnv)
			} else if loadedRC.mtime != foundRC.mtime {
				fmt.Fprintf(os.Stderr, "direnv: reloading\n")
				newEnv, err = loadRC(foundRC, config, oldEnv)
			} else {
				// Nothing to do. Env is loaded and hasn't changed
				return nil
			}
		}

	error:
		if err != nil {
			newEnv = oldEnv
			if foundRC != nil {
				// This should be nearby rc.Load()'s similar statement
				newEnv["DIRENV_DIR"] = "-" + filepath.Dir(foundRC.path)
				newEnv["DIRENV_MTIME"] = fmt.Sprintf("%d", foundRC.mtime)
				newEnv["DIRENV_BACKUP"] = oldEnv.Serialize()
			}
		}

		diff := EnvDiff(env, newEnv)

		diff2 := diff.Filtered()
		if len(diff2) > 0 {
			out := make([]string, len(diff2))
			i := 0
			for key, value := range diff2 {
				if value == "" {
					out[i] = "-" + key
				} else if oldEnv[key] == "" {
					out[i] = "+" + key
				} else {
					out[i] = "~" + key
				}
				i += 1
			}
			sort.Strings(out)
			fmt.Fprintf(os.Stderr, "direnv export: %s\n", strings.Join(out, " "))
		}

		str := EnvToShell(diff, shell)

		fmt.Print(str)
		return

	},
}

const NOT_ALLOWED = "%s is blocked from loading. Run `direnv allow` to approve it's content for loading.\n"

func loadRC(rc *RC, config *Config, env Env) (newEnv Env, err error) {
	if !rc.Allowed() {
		return nil, fmt.Errorf(NOT_ALLOWED, rc.RelTo(config.WorkDir))
	}

	argtmpl := `eval "$("%s" stdlib)" >&2 && source_env "%s" >&2 && "%s" dump`
	arg := fmt.Sprintf(argtmpl, config.SelfPath, rc.RelTo(config.WorkDir), config.SelfPath)
	cmd := exec.Command(config.BashPath, "--noprofile", "--norc", "-c", arg)

	cmd.Stderr = os.Stderr
	cmd.Env = env.ToGoEnv()
	cmd.Dir = config.WorkDir

	out, err := cmd.Output()
	if err != nil {
		err = fmt.Errorf("direnv: loading error, %q", err)
		return
	}

	newEnv, err = ParseEnv(string(out))
	if err != nil {
		err = fmt.Errorf("loadRC() ParseEnv failed: %q", err)
		return
	}

	newEnv["DIRENV_DIR"] = "-" + filepath.Dir(rc.path)
	newEnv["DIRENV_MTIME"] = fmt.Sprintf("%d", rc.mtime)
	newEnv["DIRENV_BACKUP"] = env.Serialize()

	return
}
<|endoftext|>


<commit_before>// Copyright 2013-2017 the u-root Authors. All rights reserved
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Delete files.
//
// Synopsis:
//     rm [-Rrvi] FILE...
//
// Options:
//     -i: interactive mode
//     -v: verbose mode
//     -R: remove file hierarchies
//     -r: equivalent to -R
package main

import (
	"bufio"
	"flag"
	"fmt"
	"os"
	"path"
	"path/filepath"
)

var (
	flags struct {
		r bool
		v bool
		i bool
	}
	cmd = "rm [-Rrvi] file..."
)

func init() {
	defUsage := flag.Usage
	flag.Usage = func() {
		os.Args[0] = cmd
		defUsage()
	}
	flag.BoolVar(&flags.i, "i", false, "Interactive mode.")
	flag.BoolVar(&flags.v, "v", false, "Verbose mode.")
	flag.BoolVar(&flags.r, "R", false, "Remove file hierarchies")
	flag.BoolVar(&flags.r, "r", false, "Equivalent to -R.")
	flag.Parse()
}

func rm(files []string) error {
	f := os.Remove
	if flags.r {
		f = os.RemoveAll
	}

	workingPath, err := os.Getwd()
	if err != nil {
		return err
	}

	input := bufio.NewScanner(os.Stdin)
	for _, file := range files {
		if flags.i {
			fmt.Printf("rm: remove '%v'? ", file)
			input.Scan()
			if input.Text()[0] != 'y' {
				continue
			}
		}

		if err := f(file); err != nil {
			return err
		}

		if flags.v {
			toRemove := file
			if !path.IsAbs(file) {
				toRemove = filepath.Join(workingPath, file)
			}
			fmt.Printf("removed '%v'\n", toRemove)
		}
	}
	return nil
}

func main() {
	if flag.NArg() < 1 {
		flag.Usage()
		os.Exit(1)
	}

	if err := rm(flag.Args()); err != nil {
		fmt.Fprintf(os.Stderr, "%v\n", err)
		os.Exit(1)
	}
}
<commit_msg>rm: Fix out of bounds access; accept 'Y'<commit_after>// Copyright 2013-2017 the u-root Authors. All rights reserved
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Delete files.
//
// Synopsis:
//     rm [-Rrvi] FILE...
//
// Options:
//     -i: interactive mode
//     -v: verbose mode
//     -R: remove file hierarchies
//     -r: equivalent to -R
package main

import (
	"bufio"
	"flag"
	"fmt"
	"os"
	"path"
	"path/filepath"
	"strings"
)

var (
	flags struct {
		r bool
		v bool
		i bool
	}
	cmd = "rm [-Rrvi] file..."
)

func init() {
	defUsage := flag.Usage
	flag.Usage = func() {
		os.Args[0] = cmd
		defUsage()
	}
	flag.BoolVar(&flags.i, "i", false, "Interactive mode.")
	flag.BoolVar(&flags.v, "v", false, "Verbose mode.")
	flag.BoolVar(&flags.r, "R", false, "Remove file hierarchies")
	flag.BoolVar(&flags.r, "r", false, "Equivalent to -R.")
	flag.Parse()
}

func rm(files []string) error {
	f := os.Remove
	if flags.r {
		f = os.RemoveAll
	}

	workingPath, err := os.Getwd()
	if err != nil {
		return err
	}

	input := bufio.NewReader(os.Stdin)
	for _, file := range files {
		if flags.i {
			fmt.Printf("rm: remove '%v'? ", file)
			answer, err := input.ReadString('\n')
			if err != nil || strings.ToLower(answer)[0] != 'y' {
				continue
			}
		}

		if err := f(file); err != nil {
			return err
		}

		if flags.v {
			toRemove := file
			if !path.IsAbs(file) {
				toRemove = filepath.Join(workingPath, file)
			}
			fmt.Printf("removed '%v'\n", toRemove)
		}
	}
	return nil
}

func main() {
	if flag.NArg() < 1 {
		flag.Usage()
		os.Exit(1)
	}

	if err := rm(flag.Args()); err != nil {
		fmt.Fprintf(os.Stderr, "%v\n", err)
		os.Exit(1)
	}
}
<|endoftext|>


<commit_before>package server

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"math/rand"
	"net/http"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/gorilla/mux"
	"github.com/otoolep/raft"
	"github.com/otoolep/rqlite/command"
	"github.com/otoolep/rqlite/db"
	"github.com/otoolep/rqlite/interfaces"
	"github.com/rcrowley/go-metrics"

	log "code.google.com/p/log4go"
)

type FailedSqlStmt struct {
	Sql   string `json:"sql"`
	Error string `json:"error"`
}

type StmtResponse struct {
	Time     string          `json:"time"`
	Failures []FailedSqlStmt `json:"failures"`
}

type QueryResponse struct {
	Time     string          `json:"time"`
	Failures []FailedSqlStmt `json:"failures"`
	Rows     db.RowResults   `json:"rows"`
}

type ServerMetrics struct {
	registry          metrics.Registry
	queryReceived     metrics.Counter
	querySuccess      metrics.Counter
	queryFail         metrics.Counter
	executeReceived   metrics.Counter
	executeTxReceived metrics.Counter
	executeSuccess    metrics.Counter
	executeFail       metrics.Counter
}

// The raftd server is a combination of the Raft server and an HTTP
// server which acts as the transport.
type Server struct {
	name       string
	host       string
	port       int
	path       string
	router     *mux.Router
	raftServer raft.Server
	httpServer *http.Server
	database   *db.DB
	metrics    *ServerMetrics
	mutex      sync.RWMutex
}

// queryParam returns whether the given query param is set to true.
func queryParam(req *http.Request, param string) (bool, error) {
	err := req.ParseForm()
	if err != nil {
		return false, err
	}
	if _, ok := req.Form[param]; ok {
		return true, nil
	}
	return false, nil
}

// isPretty returns whether the HTTP response body should be pretty-printed.
func isPretty(req *http.Request) (bool, error) {
	return queryParam(req, "pretty")
}

// isTransaction returns whether the client requested an explicit
// transaction for the request.
func isTransaction(req *http.Request) (bool, error) {
	return queryParam(req, "transaction")
}

// Creates a new ServerMetrics object.
func NewServerMetrics() *ServerMetrics {
	m := &ServerMetrics{
		registry:          metrics.NewRegistry(),
		queryReceived:     metrics.NewCounter(),
		querySuccess:      metrics.NewCounter(),
		queryFail:         metrics.NewCounter(),
		executeReceived:   metrics.NewCounter(),
		executeTxReceived: metrics.NewCounter(),
		executeSuccess:    metrics.NewCounter(),
		executeFail:       metrics.NewCounter(),
	}

	m.registry.Register("query.Received", m.queryReceived)
	m.registry.Register("query.success", m.querySuccess)
	m.registry.Register("query.fail", m.queryFail)
	m.registry.Register("execute.Received", m.executeReceived)
	m.registry.Register("execute.tx.received", m.executeTxReceived)
	m.registry.Register("execute.success", m.executeSuccess)
	m.registry.Register("execute.fail", m.executeFail)
	return m
}

// Creates a new server.
func New(dataDir string, database *db.DB, host string, port int) *Server {
	s := &Server{
		host:     host,
		port:     port,
		path:     dataDir,
		database: database,
		metrics:  NewServerMetrics(),
		router:   mux.NewRouter(),
	}

	// Read existing name or generate a new one.
	if b, err := ioutil.ReadFile(filepath.Join(dataDir, "name")); err == nil {
		s.name = string(b)
	} else {
		s.name = fmt.Sprintf("%07x", rand.Int())[0:7]
		if err = ioutil.WriteFile(filepath.Join(dataDir, "name"), []byte(s.name), 0644); err != nil {
			panic(err)
		}
	}

	return s
}

// GetStatistics returns an object storing statistics, which supports JSON
// marshalling.
func (s *Server) GetStatistics() (metrics.Registry, error) {
	return s.metrics.registry, nil
}

// Returns the connection string.
func (s *Server) connectionString() string {
	return fmt.Sprintf("http://%s:%d", s.host, s.port)
}

// Starts the server.
func (s *Server) ListenAndServe(leader string) error {
	var err error

	log.Info("Initializing Raft Server: %s", s.path)

	// Initialize and start Raft server.
	transporter := raft.NewHTTPTransporter("/raft", 200*time.Millisecond)
	s.raftServer, err = raft.NewServer(s.name, s.path, transporter, nil, s.database, "")
	if err != nil {
		log.Error("Failed to create new Raft server", err.Error())
		return err
	}
	transporter.Install(s.raftServer, s)
	s.raftServer.Start()

	if leader != "" {
		// Join to leader if specified.

		log.Info("Attempting to join leader at %s", leader)

		if !s.raftServer.IsLogEmpty() {
			log.Error("Cannot join with an existing log")
			return errors.New("Cannot join with an existing log")
		}
		if err := s.Join(leader); err != nil {
			log.Error("Failed to join leader", err.Error())
			return err
		}

	} else if s.raftServer.IsLogEmpty() {
		// Initialize the server by joining itself.

		log.Info("Initializing new cluster")

		_, err := s.raftServer.Do(&raft.DefaultJoinCommand{
			Name:             s.raftServer.Name(),
			ConnectionString: s.connectionString(),
		})
		if err != nil {
			log.Error("Failed to join to self", err.Error())
		}

	} else {
		log.Info("Recovered from log")
	}

	log.Info("Initializing HTTP server")

	// Initialize and start HTTP server.
	s.httpServer = &http.Server{
		Addr:    fmt.Sprintf(":%d", s.port),
		Handler: s.router,
	}

	s.router.HandleFunc("/statistics", s.serveStatistics).Methods("GET")
	s.router.HandleFunc("/db", s.readHandler).Methods("GET")
	s.router.HandleFunc("/db", s.writeHandler).Methods("POST")
	s.router.HandleFunc("/join", s.joinHandler).Methods("POST")

	log.Info("Listening at %s", s.connectionString())

	return s.httpServer.ListenAndServe()
}

// This is a hack around Gorilla mux not providing the correct net/http
// HandleFunc() interface.
func (s *Server) HandleFunc(pattern string, handler func(http.ResponseWriter, *http.Request)) {
	s.router.HandleFunc(pattern, handler)
}

// Joins to the leader of an existing cluster.
func (s *Server) Join(leader string) error {
	command := &raft.DefaultJoinCommand{
		Name:             s.raftServer.Name(),
		ConnectionString: s.connectionString(),
	}

	var b bytes.Buffer
	json.NewEncoder(&b).Encode(command)
	resp, err := http.Post(fmt.Sprintf("http://%s/join", leader), "application/json", &b)
	if err != nil {
		return err
	}
	resp.Body.Close()

	return nil
}

func (s *Server) joinHandler(w http.ResponseWriter, req *http.Request) {
	command := &raft.DefaultJoinCommand{}

	if err := json.NewDecoder(req.Body).Decode(&command); err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	if _, err := s.raftServer.Do(command); err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
}

func (s *Server) readHandler(w http.ResponseWriter, req *http.Request) {
	log.Trace("readHandler for URL: %s", req.URL)
	s.metrics.queryReceived.Inc(1)

	var failures = make([]FailedSqlStmt, 0)

	b, err := ioutil.ReadAll(req.Body)
	if err != nil {
		log.Trace("Bad HTTP request", err.Error())
		w.WriteHeader(http.StatusBadRequest)
		s.metrics.queryFail.Inc(1)
		return
	}

	stmt := string(b)
	startTime := time.Now()
	r, err := s.database.Query(stmt)
	if err != nil {
		log.Trace("Bad SQL statement", err.Error())
		s.metrics.queryFail.Inc(1)
		failures = append(failures, FailedSqlStmt{stmt, err.Error()})
	} else {
		s.metrics.querySuccess.Inc(1)
	}
	duration := time.Since(startTime)

	rr := QueryResponse{Time: duration.String(), Failures: failures, Rows: r}
	pretty, _ := isPretty(req)
	if pretty {
		b, err = json.MarshalIndent(rr, "", "    ")
	} else {
		b, err = json.Marshal(rr)
	}
	if err != nil {
		log.Trace("Failed to marshal JSON data", err.Error())
		http.Error(w, err.Error(), http.StatusBadRequest) // Internal error actually
	} else {
		w.Write([]byte(b))
	}
}

func (s *Server) writeHandler(w http.ResponseWriter, req *http.Request) {
	log.Trace("writeHandler for URL: %s", req.URL)
	s.metrics.executeReceived.Inc(1)

	var failures = make([]FailedSqlStmt, 0)
	var startTime time.Time

	// Read the value from the POST body.
	b, err := ioutil.ReadAll(req.Body)
	if err != nil {
		log.Trace("Bad HTTP request", err.Error())
		s.metrics.executeFail.Inc(1)
		w.WriteHeader(http.StatusBadRequest)
		return
	}
	stmts := strings.Split(string(b), "\n")
	if stmts[len(stmts)-1] == "" {
		stmts = stmts[:len(stmts)-1]
	}

	log.Trace("Execute statement contains %d commands", len(stmts))
	if len(stmts) == 0 {
		log.Trace("No database execute commands supplied")
		s.metrics.executeFail.Inc(1)
		w.WriteHeader(http.StatusBadRequest)
		return
	}

	startTime = time.Now()
	transaction, _ := isTransaction(req)
	if transaction {
		log.Trace("Transaction requested")
		s.metrics.executeTxReceived.Inc(1)

		_, err = s.raftServer.Do(command.NewTransactionExecuteCommandSet(stmts))
		if err != nil {
			log.Trace("Transaction failed: %s", err.Error())
			s.metrics.executeFail.Inc(1)
			failures = append(failures, FailedSqlStmt{stmts[0], err.Error()})
		} else {
			s.metrics.executeSuccess.Inc(1)
		}
	} else {
		log.Trace("No transaction requested")
		for i := range stmts {
			_, err = s.raftServer.Do(command.NewExecuteCommand(stmts[i]))
			if err != nil {
				log.Trace("Execute statement %s failed: %s", stmts[i], err.Error())
				s.metrics.executeFail.Inc(1)
				failures = append(failures, FailedSqlStmt{stmts[i], err.Error()})
			} else {
				s.metrics.executeSuccess.Inc(1)
			}

		}
	}
	duration := time.Since(startTime)

	wr := StmtResponse{Time: duration.String(), Failures: failures}
	pretty, _ := isPretty(req)
	if pretty {
		b, err = json.MarshalIndent(wr, "", "    ")
	} else {
		b, err = json.Marshal(wr)
	}
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
	} else {
		w.Write([]byte(b))
	}
}

// serveStatistics returns the statistics for the program
func (s *Server) serveStatistics(w http.ResponseWriter, req *http.Request) {
	statistics := make(map[string]interface{})
	resources := map[string]interfaces.Statistics{"server": s}
	for k, v := range resources {
		s, err := v.GetStatistics()
		if err != nil {
			log.Error("failed to get " + k + " stats")
			http.Error(w, "failed to get "+k+" stats", http.StatusInternalServerError)
			return
		}
		statistics[k] = s
	}

	var b []byte
	var err error
	pretty, _ := isPretty(req)
	if pretty {
		b, err = json.MarshalIndent(statistics, "", "    ")
	} else {
		b, err = json.Marshal(statistics)
	}
	if err != nil {
		log.Error("failed to JSON marshal statistics map")
		http.Error(w, "failed to JSON marshal statistics map", http.StatusInternalServerError)
		return
	}
	w.Write(b)
}
<commit_msg>Count joins via metrics<commit_after>package server

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"math/rand"
	"net/http"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/gorilla/mux"
	"github.com/otoolep/raft"
	"github.com/otoolep/rqlite/command"
	"github.com/otoolep/rqlite/db"
	"github.com/otoolep/rqlite/interfaces"
	"github.com/rcrowley/go-metrics"

	log "code.google.com/p/log4go"
)

type FailedSqlStmt struct {
	Sql   string `json:"sql"`
	Error string `json:"error"`
}

type StmtResponse struct {
	Time     string          `json:"time"`
	Failures []FailedSqlStmt `json:"failures"`
}

type QueryResponse struct {
	Time     string          `json:"time"`
	Failures []FailedSqlStmt `json:"failures"`
	Rows     db.RowResults   `json:"rows"`
}

type ServerMetrics struct {
	registry          metrics.Registry
	joinSuccess       metrics.Counter
	joinFail          metrics.Counter
	queryReceived     metrics.Counter
	querySuccess      metrics.Counter
	queryFail         metrics.Counter
	executeReceived   metrics.Counter
	executeTxReceived metrics.Counter
	executeSuccess    metrics.Counter
	executeFail       metrics.Counter
}

// The raftd server is a combination of the Raft server and an HTTP
// server which acts as the transport.
type Server struct {
	name       string
	host       string
	port       int
	path       string
	router     *mux.Router
	raftServer raft.Server
	httpServer *http.Server
	database   *db.DB
	metrics    *ServerMetrics
	mutex      sync.RWMutex
}

// queryParam returns whether the given query param is set to true.
func queryParam(req *http.Request, param string) (bool, error) {
	err := req.ParseForm()
	if err != nil {
		return false, err
	}
	if _, ok := req.Form[param]; ok {
		return true, nil
	}
	return false, nil
}

// isPretty returns whether the HTTP response body should be pretty-printed.
func isPretty(req *http.Request) (bool, error) {
	return queryParam(req, "pretty")
}

// isTransaction returns whether the client requested an explicit
// transaction for the request.
func isTransaction(req *http.Request) (bool, error) {
	return queryParam(req, "transaction")
}

// Creates a new ServerMetrics object.
func NewServerMetrics() *ServerMetrics {
	m := &ServerMetrics{
		registry:          metrics.NewRegistry(),
		joinSuccess:       metrics.NewCounter(),
		joinFail:          metrics.NewCounter(),
		queryReceived:     metrics.NewCounter(),
		querySuccess:      metrics.NewCounter(),
		queryFail:         metrics.NewCounter(),
		executeReceived:   metrics.NewCounter(),
		executeTxReceived: metrics.NewCounter(),
		executeSuccess:    metrics.NewCounter(),
		executeFail:       metrics.NewCounter(),
	}

	m.registry.Register("join.succes", m.joinSuccess)
	m.registry.Register("join.fail", m.joinFail)
	m.registry.Register("query.Received", m.queryReceived)
	m.registry.Register("query.success", m.querySuccess)
	m.registry.Register("query.fail", m.queryFail)
	m.registry.Register("execute.Received", m.executeReceived)
	m.registry.Register("execute.tx.received", m.executeTxReceived)
	m.registry.Register("execute.success", m.executeSuccess)
	m.registry.Register("execute.fail", m.executeFail)
	return m
}

// Creates a new server.
func New(dataDir string, database *db.DB, host string, port int) *Server {
	s := &Server{
		host:     host,
		port:     port,
		path:     dataDir,
		database: database,
		metrics:  NewServerMetrics(),
		router:   mux.NewRouter(),
	}

	// Read existing name or generate a new one.
	if b, err := ioutil.ReadFile(filepath.Join(dataDir, "name")); err == nil {
		s.name = string(b)
	} else {
		s.name = fmt.Sprintf("%07x", rand.Int())[0:7]
		if err = ioutil.WriteFile(filepath.Join(dataDir, "name"), []byte(s.name), 0644); err != nil {
			panic(err)
		}
	}

	return s
}

// GetStatistics returns an object storing statistics, which supports JSON
// marshalling.
func (s *Server) GetStatistics() (metrics.Registry, error) {
	return s.metrics.registry, nil
}

// Returns the connection string.
func (s *Server) connectionString() string {
	return fmt.Sprintf("http://%s:%d", s.host, s.port)
}

// Starts the server.
func (s *Server) ListenAndServe(leader string) error {
	var err error

	log.Info("Initializing Raft Server: %s", s.path)

	// Initialize and start Raft server.
	transporter := raft.NewHTTPTransporter("/raft", 200*time.Millisecond)
	s.raftServer, err = raft.NewServer(s.name, s.path, transporter, nil, s.database, "")
	if err != nil {
		log.Error("Failed to create new Raft server", err.Error())
		return err
	}
	transporter.Install(s.raftServer, s)
	s.raftServer.Start()

	if leader != "" {
		// Join to leader if specified.

		log.Info("Attempting to join leader at %s", leader)

		if !s.raftServer.IsLogEmpty() {
			log.Error("Cannot join with an existing log")
			return errors.New("Cannot join with an existing log")
		}
		if err := s.Join(leader); err != nil {
			log.Error("Failed to join leader", err.Error())
			return err
		}

	} else if s.raftServer.IsLogEmpty() {
		// Initialize the server by joining itself.

		log.Info("Initializing new cluster")

		_, err := s.raftServer.Do(&raft.DefaultJoinCommand{
			Name:             s.raftServer.Name(),
			ConnectionString: s.connectionString(),
		})
		if err != nil {
			log.Error("Failed to join to self", err.Error())
		}

	} else {
		log.Info("Recovered from log")
	}

	log.Info("Initializing HTTP server")

	// Initialize and start HTTP server.
	s.httpServer = &http.Server{
		Addr:    fmt.Sprintf(":%d", s.port),
		Handler: s.router,
	}

	s.router.HandleFunc("/statistics", s.serveStatistics).Methods("GET")
	s.router.HandleFunc("/db", s.readHandler).Methods("GET")
	s.router.HandleFunc("/db", s.writeHandler).Methods("POST")
	s.router.HandleFunc("/join", s.joinHandler).Methods("POST")

	log.Info("Listening at %s", s.connectionString())

	return s.httpServer.ListenAndServe()
}

// This is a hack around Gorilla mux not providing the correct net/http
// HandleFunc() interface.
func (s *Server) HandleFunc(pattern string, handler func(http.ResponseWriter, *http.Request)) {
	s.router.HandleFunc(pattern, handler)
}

// Joins to the leader of an existing cluster.
func (s *Server) Join(leader string) error {
	command := &raft.DefaultJoinCommand{
		Name:             s.raftServer.Name(),
		ConnectionString: s.connectionString(),
	}

	var b bytes.Buffer
	json.NewEncoder(&b).Encode(command)
	resp, err := http.Post(fmt.Sprintf("http://%s/join", leader), "application/json", &b)
	if err != nil {
		return err
	}
	resp.Body.Close()

	return nil
}

func (s *Server) joinHandler(w http.ResponseWriter, req *http.Request) {
	command := &raft.DefaultJoinCommand{}

	if err := json.NewDecoder(req.Body).Decode(&command); err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		s.metrics.joinFail.Inc(1)
		return
	}
	if _, err := s.raftServer.Do(command); err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		s.metrics.joinFail.Inc(1)
		return
	}
	s.metrics.joinSuccess.Inc(1)
}

func (s *Server) readHandler(w http.ResponseWriter, req *http.Request) {
	log.Trace("readHandler for URL: %s", req.URL)
	s.metrics.queryReceived.Inc(1)

	var failures = make([]FailedSqlStmt, 0)

	b, err := ioutil.ReadAll(req.Body)
	if err != nil {
		log.Trace("Bad HTTP request", err.Error())
		w.WriteHeader(http.StatusBadRequest)
		s.metrics.queryFail.Inc(1)
		return
	}

	stmt := string(b)
	startTime := time.Now()
	r, err := s.database.Query(stmt)
	if err != nil {
		log.Trace("Bad SQL statement", err.Error())
		s.metrics.queryFail.Inc(1)
		failures = append(failures, FailedSqlStmt{stmt, err.Error()})
	} else {
		s.metrics.querySuccess.Inc(1)
	}
	duration := time.Since(startTime)

	rr := QueryResponse{Time: duration.String(), Failures: failures, Rows: r}
	pretty, _ := isPretty(req)
	if pretty {
		b, err = json.MarshalIndent(rr, "", "    ")
	} else {
		b, err = json.Marshal(rr)
	}
	if err != nil {
		log.Trace("Failed to marshal JSON data", err.Error())
		http.Error(w, err.Error(), http.StatusBadRequest) // Internal error actually
	} else {
		w.Write([]byte(b))
	}
}

func (s *Server) writeHandler(w http.ResponseWriter, req *http.Request) {
	log.Trace("writeHandler for URL: %s", req.URL)
	s.metrics.executeReceived.Inc(1)

	var failures = make([]FailedSqlStmt, 0)
	var startTime time.Time

	// Read the value from the POST body.
	b, err := ioutil.ReadAll(req.Body)
	if err != nil {
		log.Trace("Bad HTTP request", err.Error())
		s.metrics.executeFail.Inc(1)
		w.WriteHeader(http.StatusBadRequest)
		return
	}
	stmts := strings.Split(string(b), "\n")
	if stmts[len(stmts)-1] == "" {
		stmts = stmts[:len(stmts)-1]
	}

	log.Trace("Execute statement contains %d commands", len(stmts))
	if len(stmts) == 0 {
		log.Trace("No database execute commands supplied")
		s.metrics.executeFail.Inc(1)
		w.WriteHeader(http.StatusBadRequest)
		return
	}

	startTime = time.Now()
	transaction, _ := isTransaction(req)
	if transaction {
		log.Trace("Transaction requested")
		s.metrics.executeTxReceived.Inc(1)

		_, err = s.raftServer.Do(command.NewTransactionExecuteCommandSet(stmts))
		if err != nil {
			log.Trace("Transaction failed: %s", err.Error())
			s.metrics.executeFail.Inc(1)
			failures = append(failures, FailedSqlStmt{stmts[0], err.Error()})
		} else {
			s.metrics.executeSuccess.Inc(1)
		}
	} else {
		log.Trace("No transaction requested")
		for i := range stmts {
			_, err = s.raftServer.Do(command.NewExecuteCommand(stmts[i]))
			if err != nil {
				log.Trace("Execute statement %s failed: %s", stmts[i], err.Error())
				s.metrics.executeFail.Inc(1)
				failures = append(failures, FailedSqlStmt{stmts[i], err.Error()})
			} else {
				s.metrics.executeSuccess.Inc(1)
			}

		}
	}
	duration := time.Since(startTime)

	wr := StmtResponse{Time: duration.String(), Failures: failures}
	pretty, _ := isPretty(req)
	if pretty {
		b, err = json.MarshalIndent(wr, "", "    ")
	} else {
		b, err = json.Marshal(wr)
	}
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
	} else {
		w.Write([]byte(b))
	}
}

// serveStatistics returns the statistics for the program
func (s *Server) serveStatistics(w http.ResponseWriter, req *http.Request) {
	statistics := make(map[string]interface{})
	resources := map[string]interfaces.Statistics{"server": s}
	for k, v := range resources {
		s, err := v.GetStatistics()
		if err != nil {
			log.Error("failed to get " + k + " stats")
			http.Error(w, "failed to get "+k+" stats", http.StatusInternalServerError)
			return
		}
		statistics[k] = s
	}

	var b []byte
	var err error
	pretty, _ := isPretty(req)
	if pretty {
		b, err = json.MarshalIndent(statistics, "", "    ")
	} else {
		b, err = json.Marshal(statistics)
	}
	if err != nil {
		log.Error("failed to JSON marshal statistics map")
		http.Error(w, "failed to JSON marshal statistics map", http.StatusInternalServerError)
		return
	}
	w.Write(b)
}
<|endoftext|>


<commit_before>// Package log provides flexible logging redirection.
package log

import (
	"log"
	"sync"
)

func init() {
	log.SetFlags(log.LstdFlags | log.Lmicroseconds)
}

// Logger is a general-purpose interface for displaying messages.
type Logger interface {
	Log(items ...interface{})
}

// ConsoleLogger sends its input directly to the default Go logger tied to
// os.Stderr.
type ConsoleLogger struct{}

// Log sends a message to the default Go logger.
func (l ConsoleLogger) Log(items ...interface{}) {
	log.Println(items...)
}

// BufferLogger stores all log messages.  It is primarily useful during startup
// when the eventual Logger implementation to use is not yet known.
type BufferLogger struct {
	sync.Mutex
	buf [][]interface{}
}

// Log records a formatted log message for future retrieval.
func (b *BufferLogger) Log(items ...interface{}) {
	b.Lock()
	defer b.Unlock()
	b.buf = append(b.buf, items)
}

// WriteTo sends recorded log messages to another Logger in order.
func (b *BufferLogger) WriteTo(l Logger) {
	b.Lock()
	defer b.Unlock()
	for _, items := range b.buf {
		l.Log(items...)
	}
}

// ProxyLogger is a stable recipient of log messages that can be configured
// at runtime to forward the messages to different Logger implementations.
type ProxyLogger struct {
	sync.RWMutex
	l Logger
}

// Log forwards its message to the underlying Logger implementation.  Panics if
// called before the first call to SetLogger.
func (p *ProxyLogger) Log(items ...interface{}) {
	p.RLock()
	defer p.RUnlock()
	p.l.Log(items...)
}

// SetLogger assigns an underlying Logger implementation to this ProxyLogger.
func (p *ProxyLogger) SetLogger(l Logger) {
	p.Lock()
	defer p.Unlock()
	p.l = l
}

// ContextLogger is a logger that prepends a string to each log message.
type ContextLogger struct {
	context string
	l       Logger
}

// NewContext creates a new ContextLogger using l as the underlying destination.
func NewContext(l Logger, context string) Logger {
	return &ContextLogger{context, l}
}

// Log prepends the present context from NewContext and passes the resulting
// log message to the underlying Logger.
func (c *ContextLogger) Log(items ...interface{}) {
	args := append([]interface{}{c.context}, items...)
	c.l.Log(args...)
}
<commit_msg>adding context around a nil Logger should remain a nil Logger<commit_after>// Package log provides flexible logging redirection.
package log

import (
	"log"
	"sync"
)

func init() {
	log.SetFlags(log.LstdFlags | log.Lmicroseconds)
}

// Logger is a general-purpose interface for displaying messages.
type Logger interface {
	Log(items ...interface{})
}

// ConsoleLogger sends its input directly to the default Go logger tied to
// os.Stderr.
type ConsoleLogger struct{}

// Log sends a message to the default Go logger.
func (l ConsoleLogger) Log(items ...interface{}) {
	log.Println(items...)
}

// BufferLogger stores all log messages.  It is primarily useful during startup
// when the eventual Logger implementation to use is not yet known.
type BufferLogger struct {
	sync.Mutex
	buf [][]interface{}
}

// Log records a formatted log message for future retrieval.
func (b *BufferLogger) Log(items ...interface{}) {
	b.Lock()
	defer b.Unlock()
	b.buf = append(b.buf, items)
}

// WriteTo sends recorded log messages to another Logger in order.
func (b *BufferLogger) WriteTo(l Logger) {
	b.Lock()
	defer b.Unlock()
	for _, items := range b.buf {
		l.Log(items...)
	}
}

// ProxyLogger is a stable recipient of log messages that can be configured
// at runtime to forward the messages to different Logger implementations.
type ProxyLogger struct {
	sync.RWMutex
	l Logger
}

// Log forwards its message to the underlying Logger implementation.  Panics if
// called before the first call to SetLogger.
func (p *ProxyLogger) Log(items ...interface{}) {
	p.RLock()
	defer p.RUnlock()
	p.l.Log(items...)
}

// SetLogger assigns an underlying Logger implementation to this ProxyLogger.
func (p *ProxyLogger) SetLogger(l Logger) {
	p.Lock()
	defer p.Unlock()
	p.l = l
}

// ContextLogger is a logger that prepends a string to each log message.
type ContextLogger struct {
	context string
	l       Logger
}

// NewContext creates a new ContextLogger using l as the underlying destination.
func NewContext(l Logger, context string) Logger {
	if l == nil {
		return nil
	}
	return &ContextLogger{context, l}
}

// Log prepends the present context from NewContext and passes the resulting
// log message to the underlying Logger.
func (c *ContextLogger) Log(items ...interface{}) {
	args := append([]interface{}{c.context}, items...)
	c.l.Log(args...)
}
<|endoftext|>


<commit_before>package filters

import "github.com/chakrit/go-bunyan"
import z "github.com/zaiuz/zaiuz"

const LogFilterPrefix = "z.log"

type logFilterData struct {
	logger bunyan.Log
}

func LogFilter(name string, parent bunyan.Log) z.Filter {
	if parent == nil {
		parent = bunyan.NewStdLogger(name, bunyan.StdoutSink())
	}
	if name != "" {
		parent = parent.Child()
	}

	data := &logFilterData{parent}

	return func(action z.Action) z.Action {
		return func(c *z.Context) z.Result {
			// TODO: Adds request id to logger scope.
			var sublogger bunyan.Log
			rid := GetRequestId(c)
			if rid != "" {
				sublogger = parent.
					Record("request_id", rid).
					Child()
			} else {
				sublogger = parent.Child()
			}

			sublogger.Infof("%s %s", c.Request.Method, c.Request.URL.Path)
			c.Set(LogFilterPrefix, data)

			result := action(c)

			duration := GetDuration(c)
			if duration != 0 {
				sublogger.Infof("finish %s", duration.String())
			} else {
				sublogger.Infof("finish")
			}

			return result
		}
	}
}

func GetLogger(c *z.Context) bunyan.Log {
	logger_, ok := c.GetOk(LogFilterPrefix)
	if !ok {
		return nil
	}

	logger, ok := logger_.(bunyan.Log)
	if !ok {
		return nil
	}

	return logger
}
<commit_msg>Log some interesting header on incoming request.<commit_after>package filters

import "net/http"
import "github.com/chakrit/go-bunyan"
import z "github.com/zaiuz/zaiuz"

const LogFilterPrefix = "z.log"

type logFilterData struct {
	logger bunyan.Log
}

func LogFilter(name string, parent bunyan.Log) z.Filter {
	if parent == nil {
		parent = bunyan.NewStdLogger(name, bunyan.StdoutSink())
	}
	if name != "" {
		parent = parent.Child()
	}

	data := &logFilterData{parent}

	return func(action z.Action) z.Action {
		return func(c *z.Context) z.Result {
			var sublogger bunyan.Log
			rid := GetRequestId(c)
			if rid != "" {
				sublogger = parent.Record("request_id", rid).Child()
			} else {
				sublogger = parent.Child()
			}

			sublogger.Record("headers", headerDigest(c.Request.Header)).
				Infof("%s %s", c.Request.Method, c.Request.URL.Path)
			c.Set(LogFilterPrefix, data)

			result := action(c)

			duration := GetDuration(c)
			if duration != 0 {
				sublogger.Infof("finish %s", duration.String())
			} else {
				sublogger.Infof("finish")
			}

			return result
		}
	}
}

func GetLogger(c *z.Context) bunyan.Log {
	logger_, ok := c.GetOk(LogFilterPrefix)
	if !ok {
		return nil
	}

	logger, ok := logger_.(bunyan.Log)
	if !ok {
		return nil
	}

	return logger
}

func headerDigest(headers http.Header) map[string]string {
	interests := []string{"Accept", "User-Agent", "Referer"}

	result := map[string]string{}
	for _, header := range interests {
		values, ok := headers[header]
		if ok && len(values) > 0 {
			result[header] = values[0]
		}
	}

	return result
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"io"
	"testing"

	"github.com/flynn/go-dockerclient"
	"github.com/flynn/lorne/types"
	"github.com/flynn/sampi/types"
	"github.com/technoweenie/grohl"
)

type nullLogger struct{}

func (nullLogger) Log(grohl.Data) error { return nil }

func init() { grohl.SetLogger(nullLogger{}) }

type dockerClient struct {
	createErr error
	startErr  error
	pullErr   error
	created   *docker.Config
	pulled    string
	started   bool
	hostConf  *docker.HostConfig
}

func (c *dockerClient) CreateContainer(config *docker.Config) (*docker.Container, error) {
	if c.createErr != nil {
		err := c.createErr
		c.createErr = nil
		return nil, err
	}
	c.created = config
	return &docker.Container{ID: "asdf"}, nil
}

func (c *dockerClient) StartContainer(id string, config *docker.HostConfig) error {
	if id != "asdf" {
		return errors.New("Invalid ID")
	}
	if c.startErr != nil {
		return c.startErr
	}
	c.started = true
	c.hostConf = config
	return nil
}

func (c *dockerClient) PullImage(opts docker.PullImageOptions, w io.Writer) error {
	if c.pullErr != nil {
		return c.pullErr
	}
	c.pulled = opts.Repository
	return nil
}

func testProcess(job *sampi.Job, t *testing.T) (*State, *dockerClient) {
	client := &dockerClient{}
	return testProcessWithOpts(job, "", client, t), client
}

func testProcessWithOpts(job *sampi.Job, extAddr string, client *dockerClient, t *testing.T) *State {
	if client == nil {
		client = &dockerClient{}
	}
	state := processWithOpts(job, extAddr, client)

	if client.created != job.Config {
		t.Error("job not created")
	}
	if !client.started {
		t.Error("job not started")
	}
	sjob := state.GetJob("a")
	if sjob == nil || sjob.StartedAt.IsZero() || sjob.Status != lorne.StatusRunning || sjob.ContainerID != "asdf" {
		t.Error("incorrect state")
	}

	return state
}

func testProcessWithError(job *sampi.Job, client *dockerClient, err error, t *testing.T) *State {
	state := processWithOpts(job, "", client)

	sjob := state.GetJob(job.ID)
	if sjob.Status != lorne.StatusFailed {
		t.Errorf("expected status failed, got %s", sjob.Status)
	}
	if sjob.Error != err {
		t.Error("error not saved")
	}
	return state
}

func processWithOpts(job *sampi.Job, extAddr string, client *dockerClient) *State {
	ports := make(chan int)
	state := NewState()
	go allocatePorts(ports, 500, 501)
	(&jobProcessor{
		externalAddr: extAddr,
		docker:       client,
		state:        state,
		discoverd:    extAddr + ":1111",
	}).processJob(ports, job)
	return state
}

func TestProcessJob(t *testing.T) {
	testProcess(&sampi.Job{ID: "a", Config: &docker.Config{}}, t)
}

func TestProcessJobWithPort(t *testing.T) {
	job := &sampi.Job{TCPPorts: 1, ID: "a", Config: &docker.Config{}}
	_, client := testProcess(job, t)

	if len(job.Config.Env) == 0 || job.Config.Env[len(job.Config.Env)-1] != "PORT=500" {
		t.Error("port env not set")
	}
	if _, ok := job.Config.ExposedPorts["500/tcp"]; !ok {
		t.Error("exposed port not set")
	}
	if b := client.hostConf.PortBindings["500/tcp"]; len(b) == 0 || b[0].HostPort != "500" {
		t.Error("port binding not set")
	}
}

func TestProcessWithExtAddr(t *testing.T) {
	job := &sampi.Job{ID: "a", Config: &docker.Config{}}
	testProcessWithOpts(job, "10.10.10.1", nil, t)

	if !sliceHasString(job.Config.Env, "EXTERNAL_IP=10.10.10.1") {
		t.Error("EXTERNAL_IP not set")
	}
	if !sliceHasString(job.Config.Env, "DISCOVERD=10.10.10.1:1111") {
		t.Error("DISCOVERD not set")
	}
}

func sliceHasString(slice []string, str string) bool {
	for _, s := range slice {
		if s == str {
			return true
		}
	}
	return false
}

func TestProcessWithPull(t *testing.T) {
	job := &sampi.Job{ID: "a", Config: &docker.Config{Image: "test/foo"}}
	client := &dockerClient{createErr: docker.ErrNoSuchImage}
	testProcessWithOpts(job, "", client, t)

	if client.pulled != "test/foo" {
		t.Error("image not pulled")
	}
}

func TestProcessWithCreateFailure(t *testing.T) {
	job := &sampi.Job{ID: "a", Config: &docker.Config{}}
	err := errors.New("undefined failure")
	client := &dockerClient{createErr: err}
	testProcessWithError(job, client, err, t)
}

func TestProcessWithPullFailure(t *testing.T) {
	job := &sampi.Job{ID: "a", Config: &docker.Config{}}
	err := errors.New("undefined failure")
	client := &dockerClient{createErr: docker.ErrNoSuchImage, pullErr: err}
	testProcessWithError(job, client, err, t)
}

func TestProcessWithStartFailure(t *testing.T) {
	job := &sampi.Job{ID: "a", Config: &docker.Config{}}
	err := errors.New("undefined failure")
	client := &dockerClient{startErr: err}
	testProcessWithError(job, client, err, t)
}

type schedulerSyncClient struct {
	removeErr error
	removed   []string
}

func (s *schedulerSyncClient) RemoveJobs(jobs []string) error {
	if s.removeErr != nil {
		return s.removeErr
	}
	s.removed = append(s.removed, jobs...)
	return nil
}

func TestSyncScheduler(t *testing.T) {
	events := make(chan lorne.Event)
	client := &schedulerSyncClient{}
	done := make(chan struct{})
	go func() {
		syncScheduler(client, events)
		close(done)
	}()

	events <- lorne.Event{Event: "stop", JobID: "a"}
	close(events)
	<-done

	if len(client.removed) != 1 && client.removed[0] != "a" {
		t.Error("job not removed")
	}
}

type streamClient struct {
	events chan *docker.Event
}

func (s *streamClient) Events() (*docker.EventStream, error) {
	return &docker.EventStream{Events: s.events}, nil
}

func (s *streamClient) InspectContainer(id string) (*docker.Container, error) {
	if id != "1" {
		return nil, errors.New("incorrect id")
	}
	return &docker.Container{State: docker.State{ExitCode: 1}}, nil
}

func TestStreamEvents(t *testing.T) {
	events := make(chan *docker.Event)
	state := NewState()
	state.AddJob(&sampi.Job{ID: "a"})
	state.SetContainerID("a", "1")
	state.SetStatusRunning("a")

	done := make(chan struct{})
	go func() {
		streamEvents(&streamClient{events}, state)
		close(done)
	}()

	events <- &docker.Event{Status: "die", ID: "1"}
	close(events)
	<-done

	job := state.GetJob("a")
	if job.Status != lorne.StatusCrashed {
		t.Error("incorrect status")
	}
	if job.ExitCode != 1 {
		t.Error("incorrect exit code")
	}
}
<commit_msg>host: Fix port tests<commit_after>package main

import (
	"errors"
	"io"
	"testing"

	"github.com/flynn/go-dockerclient"
	"github.com/flynn/lorne/types"
	"github.com/flynn/sampi/types"
	"github.com/technoweenie/grohl"
)

type nullLogger struct{}

func (nullLogger) Log(grohl.Data) error { return nil }

func init() { grohl.SetLogger(nullLogger{}) }

type dockerClient struct {
	createErr error
	startErr  error
	pullErr   error
	created   *docker.Config
	pulled    string
	started   bool
	hostConf  *docker.HostConfig
}

func (c *dockerClient) CreateContainer(config *docker.Config) (*docker.Container, error) {
	if c.createErr != nil {
		err := c.createErr
		c.createErr = nil
		return nil, err
	}
	c.created = config
	return &docker.Container{ID: "asdf"}, nil
}

func (c *dockerClient) StartContainer(id string, config *docker.HostConfig) error {
	if id != "asdf" {
		return errors.New("Invalid ID")
	}
	if c.startErr != nil {
		return c.startErr
	}
	c.started = true
	c.hostConf = config
	return nil
}

func (c *dockerClient) PullImage(opts docker.PullImageOptions, w io.Writer) error {
	if c.pullErr != nil {
		return c.pullErr
	}
	c.pulled = opts.Repository
	return nil
}

func testProcess(job *sampi.Job, t *testing.T) (*State, *dockerClient) {
	client := &dockerClient{}
	return testProcessWithOpts(job, "", client, t), client
}

func testProcessWithOpts(job *sampi.Job, extAddr string, client *dockerClient, t *testing.T) *State {
	if client == nil {
		client = &dockerClient{}
	}
	state := processWithOpts(job, extAddr, client)

	if client.created != job.Config {
		t.Error("job not created")
	}
	if !client.started {
		t.Error("job not started")
	}
	sjob := state.GetJob("a")
	if sjob == nil || sjob.StartedAt.IsZero() || sjob.Status != lorne.StatusRunning || sjob.ContainerID != "asdf" {
		t.Error("incorrect state")
	}

	return state
}

func testProcessWithError(job *sampi.Job, client *dockerClient, err error, t *testing.T) *State {
	state := processWithOpts(job, "", client)

	sjob := state.GetJob(job.ID)
	if sjob.Status != lorne.StatusFailed {
		t.Errorf("expected status failed, got %s", sjob.Status)
	}
	if sjob.Error != err {
		t.Error("error not saved")
	}
	return state
}

func processWithOpts(job *sampi.Job, extAddr string, client *dockerClient) *State {
	ports := make(chan int)
	state := NewState()
	go allocatePorts(ports, 500, 505)
	(&jobProcessor{
		externalAddr: extAddr,
		docker:       client,
		state:        state,
		discoverd:    extAddr + ":1111",
	}).processJob(ports, job)
	return state
}

func TestProcessJob(t *testing.T) {
	testProcess(&sampi.Job{ID: "a", Config: &docker.Config{}}, t)
}

func TestProcessJobWithPorts(t *testing.T) {
	job := &sampi.Job{TCPPorts: 2, ID: "a", Config: &docker.Config{}}
	_, client := testProcess(job, t)

	if len(job.Config.Env) == 0 || !sliceHasString(job.Config.Env, "PORT=500") {
		t.Fatal("PORT env not set")
	}
	if !sliceHasString(job.Config.Env, "PORT_0=500") {
		t.Error("PORT_0 env not set")
	}
	if !sliceHasString(job.Config.Env, "PORT_1=501") {
		t.Error("PORT_1 env not set")
	}
	if _, ok := job.Config.ExposedPorts["500/tcp"]; !ok {
		t.Error("exposed port 500 not set")
	}
	if _, ok := job.Config.ExposedPorts["501/tcp"]; !ok {
		t.Error("exposed port 501 not set")
	}
	if b := client.hostConf.PortBindings["500/tcp"]; len(b) == 0 || b[0].HostPort != "500" {
		t.Error("port 500 binding not set")
	}
	if b := client.hostConf.PortBindings["501/tcp"]; len(b) == 0 || b[0].HostPort != "501" {
		t.Error("port 501 binding not set")
	}
}

func TestProcessWithExtAddr(t *testing.T) {
	job := &sampi.Job{ID: "a", Config: &docker.Config{}}
	testProcessWithOpts(job, "10.10.10.1", nil, t)

	if !sliceHasString(job.Config.Env, "EXTERNAL_IP=10.10.10.1") {
		t.Error("EXTERNAL_IP not set")
	}
	if !sliceHasString(job.Config.Env, "DISCOVERD=10.10.10.1:1111") {
		t.Error("DISCOVERD not set")
	}
}

func sliceHasString(slice []string, str string) bool {
	for _, s := range slice {
		if s == str {
			return true
		}
	}
	return false
}

func TestProcessWithPull(t *testing.T) {
	job := &sampi.Job{ID: "a", Config: &docker.Config{Image: "test/foo"}}
	client := &dockerClient{createErr: docker.ErrNoSuchImage}
	testProcessWithOpts(job, "", client, t)

	if client.pulled != "test/foo" {
		t.Error("image not pulled")
	}
}

func TestProcessWithCreateFailure(t *testing.T) {
	job := &sampi.Job{ID: "a", Config: &docker.Config{}}
	err := errors.New("undefined failure")
	client := &dockerClient{createErr: err}
	testProcessWithError(job, client, err, t)
}

func TestProcessWithPullFailure(t *testing.T) {
	job := &sampi.Job{ID: "a", Config: &docker.Config{}}
	err := errors.New("undefined failure")
	client := &dockerClient{createErr: docker.ErrNoSuchImage, pullErr: err}
	testProcessWithError(job, client, err, t)
}

func TestProcessWithStartFailure(t *testing.T) {
	job := &sampi.Job{ID: "a", Config: &docker.Config{}}
	err := errors.New("undefined failure")
	client := &dockerClient{startErr: err}
	testProcessWithError(job, client, err, t)
}

type schedulerSyncClient struct {
	removeErr error
	removed   []string
}

func (s *schedulerSyncClient) RemoveJobs(jobs []string) error {
	if s.removeErr != nil {
		return s.removeErr
	}
	s.removed = append(s.removed, jobs...)
	return nil
}

func TestSyncScheduler(t *testing.T) {
	events := make(chan lorne.Event)
	client := &schedulerSyncClient{}
	done := make(chan struct{})
	go func() {
		syncScheduler(client, events)
		close(done)
	}()

	events <- lorne.Event{Event: "stop", JobID: "a"}
	close(events)
	<-done

	if len(client.removed) != 1 && client.removed[0] != "a" {
		t.Error("job not removed")
	}
}

type streamClient struct {
	events chan *docker.Event
}

func (s *streamClient) Events() (*docker.EventStream, error) {
	return &docker.EventStream{Events: s.events}, nil
}

func (s *streamClient) InspectContainer(id string) (*docker.Container, error) {
	if id != "1" {
		return nil, errors.New("incorrect id")
	}
	return &docker.Container{State: docker.State{ExitCode: 1}}, nil
}

func TestStreamEvents(t *testing.T) {
	events := make(chan *docker.Event)
	state := NewState()
	state.AddJob(&sampi.Job{ID: "a"})
	state.SetContainerID("a", "1")
	state.SetStatusRunning("a")

	done := make(chan struct{})
	go func() {
		streamEvents(&streamClient{events}, state)
		close(done)
	}()

	events <- &docker.Event{Status: "die", ID: "1"}
	close(events)
	<-done

	job := state.GetJob("a")
	if job.Status != lorne.StatusCrashed {
		t.Error("incorrect status")
	}
	if job.ExitCode != 1 {
		t.Error("incorrect exit code")
	}
}
<|endoftext|>


<commit_before>package WebService

import (
	"fmt"
	"log"
	"net/http"

	"github.com/gaeshi-gaeshi/thermo-smart/TemperatureSensorController"
)

// Start is used to setup and start the web service
func Start() {
	http.HandleFunc("/temperature", temperatureHandler)
	http.ListenAndServe(":9999", nil)
}

func temperatureHandler(w http.ResponseWriter, r *http.Request) {
	if r.Method == "GET" {
		currentTemperature, error := TemperatureSensorController.ReadTemperature()
		if error != nil {
			log.Fatal(error)
		}

		fmt.Fprint(w, currentTemperature)
	} else if r.Method == "POST" {
		// Change target temperature
	}
}
<commit_msg>Created endpoint for controlling and getting status of heaters<commit_after>package WebService

import (
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"strconv"

	"github.com/gaeshi-gaeshi/thermo-smart/HeatersController"
	"github.com/gaeshi-gaeshi/thermo-smart/TemperatureSensorController"
)

// Start is used to setup and start the web service
func Start() {
	http.HandleFunc("/temperature", temperatureHandler)
	http.HandleFunc("/heaters", heatersHandler)
	http.ListenAndServe(":9999", nil)
}

func temperatureHandler(w http.ResponseWriter, r *http.Request) {
	if r.Method == "GET" {
		currentTemperature, error := TemperatureSensorController.ReadTemperature()
		if error != nil {
			log.Fatal(error)
		}

		fmt.Fprint(w, currentTemperature)
	} else if r.Method == "POST" {
		// Change target temperature
	} else {
		w.WriteHeader(http.StatusNotFound)
	}
}

func heatersHandler(w http.ResponseWriter, r *http.Request) {
	if r.Method == "GET" {
		fmt.Fprint(w, HeatersController.GetNumberOfWorkingHeaters())
	} else if r.Method == "POST" {
		bodyBytes, error := ioutil.ReadAll(r.Body)
		if error != nil {
			log.Fatal(error)
		}

		bodyString := string(bodyBytes)
		numberOfHeaters, error := strconv.Atoi(bodyString)
		if error != nil {
			log.Fatal(error)
		}

		HeatersController.SetNumberOfWorkingHeaters(numberOfHeaters)
	} else {
		w.WriteHeader(http.StatusNotFound)
	}
}
<|endoftext|>


<commit_before>package main

import (
  "log"
  "os"
  "time"
  "flag"
  "runtime/pprof"
)

var cpuprofile = flag.String("cpuprofile", "", "write cpu profile to file")
var spool_size = flag.Uint64("spool-size", 1024, "Maximum number of events to spool before a flush is forced.")
var idle_timeout = flag.Duration("idle-flush-time", 5 * time.Second, "Maximum time to wait for a full spool before flushing anyway")
var config_file = flag.String("config", "", "The config file to load")

func main() {
  flag.Parse()

  if *cpuprofile != "" {
    f, err := os.Create(*cpuprofile)
    if err != nil {
        log.Fatal(err)
    }
    pprof.StartCPUProfile(f)
    go func() {
      time.Sleep(60 * time.Second)
      pprof.StopCPUProfile()
      panic("done")
    }()
  }

  log.SetFlags(log.Ldate | log.Ltime | log.Lmicroseconds)

  config, err := LoadConfig(*config_file)
  if err != nil {
    return
  }

  event_chan := make(chan *FileEvent, 16)
  publisher_chan := make(chan []*FileEvent, 1)
  registrar_chan := make(chan []*FileEvent, 1)

  if len(config.Files) == 0 {
    log.Fatalf("No paths given. What files do you want me to watch?\n")
  }

  // The basic model of execution:
  // - prospector: finds files in paths/globs to harvest, starts harvesters
  // - harvester: reads a file, sends events to the spooler
  // - spooler: buffers events until ready to flush to the publisher
  // - publisher: writes to the network, notifies registrar
  // - registrar: records positions of files read
  // Finally, prospector uses the registrar information, on restart, to
  // determine where in each file to resume a harvester.

  // Prospect the globs/paths given on the command line and launch harvesters
  for _, fileconfig := range config.Files {
    go Prospect(fileconfig, event_chan)
  }

  // Harvesters dump events into the spooler.
  go Spool(event_chan, publisher_chan, *spool_size, *idle_timeout)

  go Publishv1(publisher_chan, registrar_chan, &config.Network)

  // registrar records last acknowledged positions in all files.
  Registrar(registrar_chan)
} /* main */
<commit_msg>- add syslog logging support<commit_after>package main

import (
  "log"
  "os"
  "time"
  "flag"
  "log/syslog"
  "runtime/pprof"
)

var cpuprofile = flag.String("cpuprofile", "", "write cpu profile to file")
var spool_size = flag.Uint64("spool-size", 1024, "Maximum number of events to spool before a flush is forced.")
var idle_timeout = flag.Duration("idle-flush-time", 5 * time.Second, "Maximum time to wait for a full spool before flushing anyway")
var config_file = flag.String("config", "", "The config file to load")
var use_syslog = flag.Bool("log-to-syslog", false, "Log to syslog instead of stdout")

func main() {
  flag.Parse()

  if *cpuprofile != "" {
    f, err := os.Create(*cpuprofile)
    if err != nil {
        log.Fatal(err)
    }
    pprof.StartCPUProfile(f)
    go func() {
      time.Sleep(60 * time.Second)
      pprof.StopCPUProfile()
      panic("done")
    }()
  }

  config, err := LoadConfig(*config_file)
  if err != nil {
    return
  }

  event_chan := make(chan *FileEvent, 16)
  publisher_chan := make(chan []*FileEvent, 1)
  registrar_chan := make(chan []*FileEvent, 1)

  if len(config.Files) == 0 {
    log.Fatalf("No paths given. What files do you want me to watch?\n")
  }

  // The basic model of execution:
  // - prospector: finds files in paths/globs to harvest, starts harvesters
  // - harvester: reads a file, sends events to the spooler
  // - spooler: buffers events until ready to flush to the publisher
  // - publisher: writes to the network, notifies registrar
  // - registrar: records positions of files read
  // Finally, prospector uses the registrar information, on restart, to
  // determine where in each file to resume a harvester.

  log.SetFlags(log.Ldate | log.Ltime | log.Lmicroseconds)
  if *use_syslog {
    writer, err := syslog.New(syslog.LOG_INFO | syslog.LOG_DAEMON, "lumberjack")
    if err != nil {
      log.Fatalf("Failed to open syslog: %s\n", err)
    }
    log.SetOutput(writer)
  }

  // Prospect the globs/paths given on the command line and launch harvesters
  for _, fileconfig := range config.Files {
    go Prospect(fileconfig, event_chan)
  }

  // Harvesters dump events into the spooler.
  go Spool(event_chan, publisher_chan, *spool_size, *idle_timeout)

  go Publishv1(publisher_chan, registrar_chan, &config.Network)

  // registrar records last acknowledged positions in all files.
  Registrar(registrar_chan)
} /* main */
<|endoftext|>