<commit_before>/*
rqlite -- a replicated SQLite database.

rqlite is a distributed system that provides a replicated SQLite database.
rqlite is written in Go and uses Raft to achieve consensus across all the
instances of the SQLite databases. rqlite ensures that every change made to
the database is made to a majority of underlying SQLite files, or none-at-all.
*/

package main

import (
	"bytes"
	"crypto/tls"
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"net"
	"net/http"
	"os"
	"os/signal"
	"path/filepath"
	"runtime/pprof"
	"time"

	"github.com/rqlite/rqlite/auth"
	"github.com/rqlite/rqlite/cluster"
	httpd "github.com/rqlite/rqlite/http"
	"github.com/rqlite/rqlite/store"
	"github.com/rqlite/rqlite/tcp"
)

const sqliteDSN = "db.sqlite"

const logo = `
            _ _ _           _
           | (_) |         | |
  _ __ __ _| |_| |_ ___  __| |
 | '__/ _  | | | __/ _ \/ _  |
 | | | (_| | | | ||  __/ (_| |
 |_|  \__, |_|_|\__\___|\__,_|
         | |
         |_|
`

// These variables are populated via the Go linker.
var (
	version   = "2"
	commit    string
	branch    string
	buildtime string
)

const (
	muxRaftHeader = 1 // Raft consensus communications
	muxMetaHeader = 2 // Cluster meta communications
)

const (
	publishPeerDelay   = 1 * time.Second
	publishPeerTimeout = 30 * time.Second
)

var httpAddr string
var authFile string
var x509Cert string
var x509Key string
var raftAddr string
var joinAddr string
var noVerify bool
var expvar bool
var dsn string
var inMem bool
var showVersion bool
var cpuprofile string

const desc = `rqlite is a distributed system that provides a replicated SQLite database.`

func init() {
	flag.StringVar(&httpAddr, "http", "localhost:4001", "HTTP query server address. Set X.509 cert and key for HTTPS.")
	flag.StringVar(&x509Cert, "x509cert", "", "Path to X.509 certificate")
	flag.StringVar(&x509Key, "x509key", "", "Path to X.509 private key for certificate")
	flag.StringVar(&authFile, "auth", "", "Path to authentication and authorization file. If not set, not enabled.")
	flag.StringVar(&raftAddr, "raft", "localhost:4002", "Raft communication bind address")
	flag.StringVar(&joinAddr, "join", "", "protocol://host:port of leader to join")
	flag.BoolVar(&noVerify, "noverify", false, "Skip verification of any HTTPS cert when joining")
	flag.BoolVar(&expvar, "expvar", true, "Serve expvar data on HTTP server")
	flag.StringVar(&dsn, "dsn", "", `SQLite DSN parameters. E.g. "cache=shared&mode=memory"`)
	flag.BoolVar(&inMem, "mem", false, "Use an in-memory database")
	flag.BoolVar(&showVersion, "version", false, "Show version information and exit")
	flag.StringVar(&cpuprofile, "cpuprofile", "", "Write CPU profile to file")
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "\n%s\n\n", desc)
		fmt.Fprintf(os.Stderr, "Usage: %s [arguments] <data directory>\n", os.Args[0])
		flag.PrintDefaults()
	}

	// If commit, branch, or build time are not set, make that clear.
	if commit == "" {
		commit = "unknown"
	}
	if branch == "" {
		branch = "unknown"
	}
	if buildtime == "" {
		buildtime = "unknown"
	}
}

func main() {
	flag.Parse()

	if showVersion {
		fmt.Printf("rqlited version %s (commit %s)\n", version, commit)
		os.Exit(0)
	}

	// Ensure the data path is set.
	if flag.NArg() == 0 {
		flag.Usage()
		os.Exit(1)
	}

	dataPath := flag.Arg(0)

	// Display logo.
	fmt.Println(logo)

	// Configure logging and pump out initial message.
	log.SetFlags(log.LstdFlags)
	log.SetPrefix("[rqlited] ")
	log.Printf("rqlited starting, version %s, commit %s, branch %s", version, commit, branch)

	// Set up profiling, if requested.
	if cpuprofile != "" {
		log.Println("profiling enabled")
		f, err := os.Create(cpuprofile)
		if err != nil {
			log.Printf("unable to create path: %s", err.Error())
		}
		defer f.Close()

		err = pprof.StartCPUProfile(f)
		if err != nil {
			log.Fatalf("unable to start CPU Profile: %s", err.Error())
		}

		defer pprof.StopCPUProfile()
	}

	// Set up TCP communication between nodes.
	ln, err := net.Listen("tcp", raftAddr)
	if err != nil {
		log.Fatalf("failed to listen on %s: %s", raftAddr, err.Error())
	}
	mux := tcp.NewMux(ln)
	go mux.Serve()

	// Start up mux and get transports for cluster.
	raftTn := mux.Listen(muxRaftHeader)

	// Create and open the store.
	dataPath, err = filepath.Abs(dataPath)
	if err != nil {
		log.Fatalf("failed to determine absolute data path: %s", err.Error())
	}
	dbConf := store.NewDBConfig(dsn, inMem)
	store := store.New(dbConf, dataPath, raftTn)
	if err := store.Open(joinAddr == ""); err != nil {
		log.Fatalf("failed to open store: %s", err.Error())
	}

	// Create and configure cluster service.
	tn := mux.Listen(muxMetaHeader)
	cs := cluster.NewService(tn, store)
	if err := cs.Open(); err != nil {
		log.Fatalf("failed to open cluster service: %s", err.Error())
	}

	// If join was specified, make the join request.
	if joinAddr != "" {
		if !store.JoinRequired() {
			log.Println("node is already member of cluster, ignoring join request")
		} else {
			if err := join(joinAddr, noVerify, raftAddr); err != nil {
				log.Fatalf("failed to join node at %s: %s", joinAddr, err.Error())
			}
			log.Println("successfully joined node at", joinAddr)
		}
	}

	// Create HTTP server and load authentication information, if supplied.
	var s *httpd.Service
	if authFile != "" {
		f, err := os.Open(authFile)
		if err != nil {
			log.Fatalf("failed to open authentication file %s: %s", authFile, err.Error())
		}
		credentialStore := auth.NewCredentialsStore()
		if err := credentialStore.Load(f); err != nil {
			log.Fatalf("failed to load authentication file: %s", err.Error())
		}
		s = httpd.New(httpAddr, store, credentialStore)
	} else {
		s = httpd.New(httpAddr, store, nil)
	}

	// Publish to the cluster the mapping between this Raft address and API address.
	if err := publishAPIAddr(cs, raftAddr, httpAddr, publishPeerTimeout); err != nil {
		log.Fatalf("failed to set peer for %s to %s: %s", raftAddr, httpAddr, err.Error())
	}
	log.Printf("set peer for %s to %s", raftAddr, httpAddr)

	s.CertFile = x509Cert
	s.KeyFile = x509Key
	s.Expvar = expvar
	s.BuildInfo = map[string]interface{}{
		"commit":     commit,
		"branch":     branch,
		"version":    version,
		"build_time": buildtime,
	}
	if err := s.Start(); err != nil {
		log.Fatalf("failed to start HTTP server: %s", err.Error())
	}

	terminate := make(chan os.Signal, 1)
	signal.Notify(terminate, os.Interrupt)
	<-terminate
	if err := store.Close(true); err != nil {
		log.Printf("failed to close store: %s", err.Error())
	}
	log.Println("rqlite server stopped")
}

func join(joinAddr string, skipVerify bool, raftAddr string) error {
	b, err := json.Marshal(map[string]string{"addr": raftAddr})
	if err != nil {
		return err
	}

	// Check for protocol scheme, and insert default if necessary.
	fullAddr := httpd.NormalizeAddr(fmt.Sprintf("%s/join", joinAddr))

	// Enable skipVerify as requested.
	tr := &http.Transport{
		TLSClientConfig: &tls.Config{InsecureSkipVerify: skipVerify},
	}
	client := &http.Client{Transport: tr}

	// Attempt to join.
	resp, err := client.Post(fullAddr, "application-type/json", bytes.NewReader(b))
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	if resp.StatusCode != 200 {
		return fmt.Errorf("failed to join, node returned: %s", resp.Status)
	}

	return nil
}

func publishAPIAddr(c *cluster.Service, raftAddr, apiAddr string, t time.Duration) error {
	tck := time.NewTicker(publishPeerDelay)
	defer tck.Stop()
	tmr := time.NewTimer(t)
	defer tmr.Stop()

	for {
		select {
		case <-tck.C:
			if err := c.SetPeer(raftAddr, apiAddr); err != nil {
				log.Printf("failed to set peer for %s to %s: %s (retrying)",
					raftAddr, apiAddr, err.Error())
				continue
			}
			return nil
		case <-tmr.C:
			return fmt.Errorf("set peer timeout expired")
		}
	}
}
<commit_msg>Publish peer before starting http service<commit_after>/*
rqlite -- a replicated SQLite database.

rqlite is a distributed system that provides a replicated SQLite database.
rqlite is written in Go and uses Raft to achieve consensus across all the
instances of the SQLite databases. rqlite ensures that every change made to
the database is made to a majority of underlying SQLite files, or none-at-all.
*/

package main

import (
	"bytes"
	"crypto/tls"
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"net"
	"net/http"
	"os"
	"os/signal"
	"path/filepath"
	"runtime/pprof"
	"time"

	"github.com/rqlite/rqlite/auth"
	"github.com/rqlite/rqlite/cluster"
	httpd "github.com/rqlite/rqlite/http"
	"github.com/rqlite/rqlite/store"
	"github.com/rqlite/rqlite/tcp"
)

const sqliteDSN = "db.sqlite"

const logo = `
            _ _ _           _
           | (_) |         | |
  _ __ __ _| |_| |_ ___  __| |
 | '__/ _  | | | __/ _ \/ _  |
 | | | (_| | | | ||  __/ (_| |
 |_|  \__, |_|_|\__\___|\__,_|
         | |
         |_|
`

// These variables are populated via the Go linker.
var (
	version   = "2"
	commit    string
	branch    string
	buildtime string
)

const (
	muxRaftHeader = 1 // Raft consensus communications
	muxMetaHeader = 2 // Cluster meta communications
)

const (
	publishPeerDelay   = 1 * time.Second
	publishPeerTimeout = 30 * time.Second
)

var httpAddr string
var authFile string
var x509Cert string
var x509Key string
var raftAddr string
var joinAddr string
var noVerify bool
var expvar bool
var dsn string
var inMem bool
var showVersion bool
var cpuprofile string

const desc = `rqlite is a distributed system that provides a replicated SQLite database.`

func init() {
	flag.StringVar(&httpAddr, "http", "localhost:4001", "HTTP query server address. Set X.509 cert and key for HTTPS.")
	flag.StringVar(&x509Cert, "x509cert", "", "Path to X.509 certificate")
	flag.StringVar(&x509Key, "x509key", "", "Path to X.509 private key for certificate")
	flag.StringVar(&authFile, "auth", "", "Path to authentication and authorization file. If not set, not enabled.")
	flag.StringVar(&raftAddr, "raft", "localhost:4002", "Raft communication bind address")
	flag.StringVar(&joinAddr, "join", "", "protocol://host:port of leader to join")
	flag.BoolVar(&noVerify, "noverify", false, "Skip verification of any HTTPS cert when joining")
	flag.BoolVar(&expvar, "expvar", true, "Serve expvar data on HTTP server")
	flag.StringVar(&dsn, "dsn", "", `SQLite DSN parameters. E.g. "cache=shared&mode=memory"`)
	flag.BoolVar(&inMem, "mem", false, "Use an in-memory database")
	flag.BoolVar(&showVersion, "version", false, "Show version information and exit")
	flag.StringVar(&cpuprofile, "cpuprofile", "", "Write CPU profile to file")
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "\n%s\n\n", desc)
		fmt.Fprintf(os.Stderr, "Usage: %s [arguments] <data directory>\n", os.Args[0])
		flag.PrintDefaults()
	}

	// If commit, branch, or build time are not set, make that clear.
	if commit == "" {
		commit = "unknown"
	}
	if branch == "" {
		branch = "unknown"
	}
	if buildtime == "" {
		buildtime = "unknown"
	}
}

func main() {
	flag.Parse()

	if showVersion {
		fmt.Printf("rqlited version %s (commit %s)\n", version, commit)
		os.Exit(0)
	}

	// Ensure the data path is set.
	if flag.NArg() == 0 {
		flag.Usage()
		os.Exit(1)
	}

	dataPath := flag.Arg(0)

	// Display logo.
	fmt.Println(logo)

	// Configure logging and pump out initial message.
	log.SetFlags(log.LstdFlags)
	log.SetPrefix("[rqlited] ")
	log.Printf("rqlited starting, version %s, commit %s, branch %s", version, commit, branch)

	// Set up profiling, if requested.
	if cpuprofile != "" {
		log.Println("profiling enabled")
		f, err := os.Create(cpuprofile)
		if err != nil {
			log.Printf("unable to create path: %s", err.Error())
		}
		defer f.Close()

		err = pprof.StartCPUProfile(f)
		if err != nil {
			log.Fatalf("unable to start CPU Profile: %s", err.Error())
		}

		defer pprof.StopCPUProfile()
	}

	// Set up TCP communication between nodes.
	ln, err := net.Listen("tcp", raftAddr)
	if err != nil {
		log.Fatalf("failed to listen on %s: %s", raftAddr, err.Error())
	}
	mux := tcp.NewMux(ln)
	go mux.Serve()

	// Start up mux and get transports for cluster.
	raftTn := mux.Listen(muxRaftHeader)

	// Create and open the store.
	dataPath, err = filepath.Abs(dataPath)
	if err != nil {
		log.Fatalf("failed to determine absolute data path: %s", err.Error())
	}
	dbConf := store.NewDBConfig(dsn, inMem)
	store := store.New(dbConf, dataPath, raftTn)
	if err := store.Open(joinAddr == ""); err != nil {
		log.Fatalf("failed to open store: %s", err.Error())
	}

	// Create and configure cluster service.
	tn := mux.Listen(muxMetaHeader)
	cs := cluster.NewService(tn, store)
	if err := cs.Open(); err != nil {
		log.Fatalf("failed to open cluster service: %s", err.Error())
	}

	// If join was specified, make the join request.
	if joinAddr != "" {
		if !store.JoinRequired() {
			log.Println("node is already member of cluster, ignoring join request")
		} else {
			if err := join(joinAddr, noVerify, raftAddr); err != nil {
				log.Fatalf("failed to join node at %s: %s", joinAddr, err.Error())
			}
			log.Println("successfully joined node at", joinAddr)
		}
	}

	// Publish to the cluster the mapping between this Raft address and API address.
	if err := publishAPIAddr(cs, raftAddr, httpAddr, publishPeerTimeout); err != nil {
		log.Fatalf("failed to set peer for %s to %s: %s", raftAddr, httpAddr, err.Error())
	}
	log.Printf("set peer for %s to %s", raftAddr, httpAddr)

	// Create HTTP server and load authentication information, if supplied.
	var s *httpd.Service
	if authFile != "" {
		f, err := os.Open(authFile)
		if err != nil {
			log.Fatalf("failed to open authentication file %s: %s", authFile, err.Error())
		}
		credentialStore := auth.NewCredentialsStore()
		if err := credentialStore.Load(f); err != nil {
			log.Fatalf("failed to load authentication file: %s", err.Error())
		}
		s = httpd.New(httpAddr, store, credentialStore)
	} else {
		s = httpd.New(httpAddr, store, nil)
	}

	s.CertFile = x509Cert
	s.KeyFile = x509Key
	s.Expvar = expvar
	s.BuildInfo = map[string]interface{}{
		"commit":     commit,
		"branch":     branch,
		"version":    version,
		"build_time": buildtime,
	}
	if err := s.Start(); err != nil {
		log.Fatalf("failed to start HTTP server: %s", err.Error())
	}

	terminate := make(chan os.Signal, 1)
	signal.Notify(terminate, os.Interrupt)
	<-terminate
	if err := store.Close(true); err != nil {
		log.Printf("failed to close store: %s", err.Error())
	}
	log.Println("rqlite server stopped")
}

func join(joinAddr string, skipVerify bool, raftAddr string) error {
	b, err := json.Marshal(map[string]string{"addr": raftAddr})
	if err != nil {
		return err
	}

	// Check for protocol scheme, and insert default if necessary.
	fullAddr := httpd.NormalizeAddr(fmt.Sprintf("%s/join", joinAddr))

	// Enable skipVerify as requested.
	tr := &http.Transport{
		TLSClientConfig: &tls.Config{InsecureSkipVerify: skipVerify},
	}
	client := &http.Client{Transport: tr}

	// Attempt to join.
	resp, err := client.Post(fullAddr, "application-type/json", bytes.NewReader(b))
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	if resp.StatusCode != 200 {
		return fmt.Errorf("failed to join, node returned: %s", resp.Status)
	}

	return nil
}

func publishAPIAddr(c *cluster.Service, raftAddr, apiAddr string, t time.Duration) error {
	tck := time.NewTicker(publishPeerDelay)
	defer tck.Stop()
	tmr := time.NewTimer(t)
	defer tmr.Stop()

	for {
		select {
		case <-tck.C:
			if err := c.SetPeer(raftAddr, apiAddr); err != nil {
				log.Printf("failed to set peer for %s to %s: %s (retrying)",
					raftAddr, apiAddr, err.Error())
				continue
			}
			return nil
		case <-tmr.C:
			return fmt.Errorf("set peer timeout expired")
		}
	}
}
<|endoftext|>


<commit_before>package packer

import (
	"fmt"
	"sort"

	"github.com/hashicorp/go-multierror"
	"github.com/hashicorp/go-version"
	"github.com/hashicorp/packer/template"
	"github.com/hashicorp/packer/template/interpolate"
)

// Core is the main executor of Packer. If Packer is being used as a
// library, this is the struct you'll want to instantiate to get anything done.
type Core struct {
	Template *template.Template

	components ComponentFinder
	variables  map[string]string
	builds     map[string]*template.Builder
	version    string
}

// CoreConfig is the structure for initializing a new Core. Once a CoreConfig
// is used to initialize a Core, it shouldn't be re-used or modified again.
type CoreConfig struct {
	Components ComponentFinder
	Template   *template.Template
	Variables  map[string]string
	Version    string
}

// The function type used to lookup Builder implementations.
type BuilderFunc func(name string) (Builder, error)

// The function type used to lookup Hook implementations.
type HookFunc func(name string) (Hook, error)

// The function type used to lookup PostProcessor implementations.
type PostProcessorFunc func(name string) (PostProcessor, error)

// The function type used to lookup Provisioner implementations.
type ProvisionerFunc func(name string) (Provisioner, error)

// ComponentFinder is a struct that contains the various function
// pointers necessary to look up components of Packer such as builders,
// commands, etc.
type ComponentFinder struct {
	Builder       BuilderFunc
	Hook          HookFunc
	PostProcessor PostProcessorFunc
	Provisioner   ProvisionerFunc
}

// NewCore creates a new Core.
func NewCore(c *CoreConfig) (*Core, error) {
	result := &Core{
		Template:   c.Template,
		components: c.Components,
		variables:  c.Variables,
		version:    c.Version,
	}
	if err := result.validate(); err != nil {
		return nil, err
	}
	if err := result.init(); err != nil {
		return nil, err
	}

	// Go through and interpolate all the build names. We shuld be able
	// to do this at this point with the variables.
	result.builds = make(map[string]*template.Builder)
	for _, b := range c.Template.Builders {
		v, err := interpolate.Render(b.Name, result.Context())
		if err != nil {
			return nil, fmt.Errorf(
				"Error interpolating builder '%s': %s",
				b.Name, err)
		}

		result.builds[v] = b
	}

	return result, nil
}

// BuildNames returns the builds that are available in this configured core.
func (c *Core) BuildNames() []string {
	r := make([]string, 0, len(c.builds))
	for n := range c.builds {
		r = append(r, n)
	}
	sort.Strings(r)

	return r
}

// Build returns the Build object for the given name.
func (c *Core) Build(n string) (Build, error) {
	// Setup the builder
	configBuilder, ok := c.builds[n]
	if !ok {
		return nil, fmt.Errorf("no such build found: %s", n)
	}
	builder, err := c.components.Builder(configBuilder.Type)
	if err != nil {
		return nil, fmt.Errorf(
			"error initializing builder '%s': %s",
			configBuilder.Type, err)
	}
	if builder == nil {
		return nil, fmt.Errorf(
			"builder type not found: %s", configBuilder.Type)
	}

	// rawName is the uninterpolated name that we use for various lookups
	rawName := configBuilder.Name

	// Setup the provisioners for this build
	provisioners := make([]coreBuildProvisioner, 0, len(c.Template.Provisioners))
	for _, rawP := range c.Template.Provisioners {
		// If we're skipping this, then ignore it
		if rawP.Skip(rawName) {
			continue
		}

		// Get the provisioner
		provisioner, err := c.components.Provisioner(rawP.Type)
		if err != nil {
			return nil, fmt.Errorf(
				"error initializing provisioner '%s': %s",
				rawP.Type, err)
		}
		if provisioner == nil {
			return nil, fmt.Errorf(
				"provisioner type not found: %s", rawP.Type)
		}

		// Get the configuration
		config := make([]interface{}, 1, 2)
		config[0] = rawP.Config
		if rawP.Override != nil {
			if override, ok := rawP.Override[rawName]; ok {
				config = append(config, override)
			}
		}

		// If we're pausing, we wrap the provisioner in a special pauser.
		if rawP.PauseBefore > 0 {
			provisioner = &PausedProvisioner{
				PauseBefore: rawP.PauseBefore,
				Provisioner: provisioner,
			}
		}

		provisioners = append(provisioners, coreBuildProvisioner{
			pType:       rawP.Type,
			provisioner: provisioner,
			config:      config,
		})
	}

	// Setup the post-processors
	postProcessors := make([][]coreBuildPostProcessor, 0, len(c.Template.PostProcessors))
	for _, rawPs := range c.Template.PostProcessors {
		current := make([]coreBuildPostProcessor, 0, len(rawPs))
		for _, rawP := range rawPs {
			// If we skip, ignore
			if rawP.Skip(rawName) {
				continue
			}

			// Get the post-processor
			postProcessor, err := c.components.PostProcessor(rawP.Type)
			if err != nil {
				return nil, fmt.Errorf(
					"error initializing post-processor '%s': %s",
					rawP.Type, err)
			}
			if postProcessor == nil {
				return nil, fmt.Errorf(
					"post-processor type not found: %s", rawP.Type)
			}

			current = append(current, coreBuildPostProcessor{
				processor:         postProcessor,
				processorType:     rawP.Type,
				config:            rawP.Config,
				keepInputArtifact: rawP.KeepInputArtifact,
			})
		}

		// If we have no post-processors in this chain, just continue.
		if len(current) == 0 {
			continue
		}

		postProcessors = append(postProcessors, current)
	}

	// TODO hooks one day

	return &coreBuild{
		name:           n,
		builder:        builder,
		builderConfig:  configBuilder.Config,
		builderType:    configBuilder.Type,
		postProcessors: postProcessors,
		provisioners:   provisioners,
		templatePath:   c.Template.Path,
		variables:      c.variables,
	}, nil
}

// Context returns an interpolation context.
func (c *Core) Context() *interpolate.Context {
	return &interpolate.Context{
		TemplatePath:  c.Template.Path,
		UserVariables: c.variables,
	}
}

// validate does a full validation of the template.
//
// This will automatically call template.validate() in addition to doing
// richer semantic checks around variables and so on.
func (c *Core) validate() error {
	// First validate the template in general, we can't do anything else
	// unless the template itself is valid.
	if err := c.Template.Validate(); err != nil {
		return err
	}

	// Validate the minimum version is satisfied
	if c.Template.MinVersion != "" {
		versionActual, err := version.NewVersion(c.version)
		if err != nil {
			// This shouldn't happen since we set it via the compiler
			panic(err)
		}

		versionMin, err := version.NewVersion(c.Template.MinVersion)
		if err != nil {
			return fmt.Errorf(
				"min_version is invalid: %s", err)
		}

		if versionActual.LessThan(versionMin) {
			return fmt.Errorf(
				"This template requires Packer version %s or higher; using %s",
				versionMin,
				versionActual)
		}
	}

	// Validate variables are set
	var err error
	for n, v := range c.Template.Variables {
		if v.Required {
			if _, ok := c.variables[n]; !ok {
				err = multierror.Append(err, fmt.Errorf(
					"required variable not set: %s", n))
			}
		}
	}

	// TODO: validate all builders exist
	// TODO: ^^ provisioner
	// TODO: ^^ post-processor

	return err
}

func (c *Core) init() error {
	if c.variables == nil {
		c.variables = make(map[string]string)
	}

	// Go through the variables and interpolate the environment variables
	ctx := c.Context()
	ctx.EnableEnv = true
	ctx.UserVariables = nil
	for k, v := range c.Template.Variables {
		// Ignore variables that are required
		if v.Required {
			continue
		}

		// Ignore variables that have a value
		if _, ok := c.variables[k]; ok {
			continue
		}

		// Interpolate the default
		def, err := interpolate.Render(v.Default, ctx)
		if err != nil {
			return fmt.Errorf(
				"error interpolating default value for '%s': %s",
				k, err)
		}

		c.variables[k] = def
	}

	// Interpolate the push configuration
	if _, err := interpolate.RenderInterface(&c.Template.Push, c.Context()); err != nil {
		return fmt.Errorf("Error interpolating 'push': %s", err)
	}

	return nil
}
<commit_msg>spelling: should<commit_after>package packer

import (
	"fmt"
	"sort"

	"github.com/hashicorp/go-multierror"
	"github.com/hashicorp/go-version"
	"github.com/hashicorp/packer/template"
	"github.com/hashicorp/packer/template/interpolate"
)

// Core is the main executor of Packer. If Packer is being used as a
// library, this is the struct you'll want to instantiate to get anything done.
type Core struct {
	Template *template.Template

	components ComponentFinder
	variables  map[string]string
	builds     map[string]*template.Builder
	version    string
}

// CoreConfig is the structure for initializing a new Core. Once a CoreConfig
// is used to initialize a Core, it shouldn't be re-used or modified again.
type CoreConfig struct {
	Components ComponentFinder
	Template   *template.Template
	Variables  map[string]string
	Version    string
}

// The function type used to lookup Builder implementations.
type BuilderFunc func(name string) (Builder, error)

// The function type used to lookup Hook implementations.
type HookFunc func(name string) (Hook, error)

// The function type used to lookup PostProcessor implementations.
type PostProcessorFunc func(name string) (PostProcessor, error)

// The function type used to lookup Provisioner implementations.
type ProvisionerFunc func(name string) (Provisioner, error)

// ComponentFinder is a struct that contains the various function
// pointers necessary to look up components of Packer such as builders,
// commands, etc.
type ComponentFinder struct {
	Builder       BuilderFunc
	Hook          HookFunc
	PostProcessor PostProcessorFunc
	Provisioner   ProvisionerFunc
}

// NewCore creates a new Core.
func NewCore(c *CoreConfig) (*Core, error) {
	result := &Core{
		Template:   c.Template,
		components: c.Components,
		variables:  c.Variables,
		version:    c.Version,
	}
	if err := result.validate(); err != nil {
		return nil, err
	}
	if err := result.init(); err != nil {
		return nil, err
	}

	// Go through and interpolate all the build names. We should be able
	// to do this at this point with the variables.
	result.builds = make(map[string]*template.Builder)
	for _, b := range c.Template.Builders {
		v, err := interpolate.Render(b.Name, result.Context())
		if err != nil {
			return nil, fmt.Errorf(
				"Error interpolating builder '%s': %s",
				b.Name, err)
		}

		result.builds[v] = b
	}

	return result, nil
}

// BuildNames returns the builds that are available in this configured core.
func (c *Core) BuildNames() []string {
	r := make([]string, 0, len(c.builds))
	for n := range c.builds {
		r = append(r, n)
	}
	sort.Strings(r)

	return r
}

// Build returns the Build object for the given name.
func (c *Core) Build(n string) (Build, error) {
	// Setup the builder
	configBuilder, ok := c.builds[n]
	if !ok {
		return nil, fmt.Errorf("no such build found: %s", n)
	}
	builder, err := c.components.Builder(configBuilder.Type)
	if err != nil {
		return nil, fmt.Errorf(
			"error initializing builder '%s': %s",
			configBuilder.Type, err)
	}
	if builder == nil {
		return nil, fmt.Errorf(
			"builder type not found: %s", configBuilder.Type)
	}

	// rawName is the uninterpolated name that we use for various lookups
	rawName := configBuilder.Name

	// Setup the provisioners for this build
	provisioners := make([]coreBuildProvisioner, 0, len(c.Template.Provisioners))
	for _, rawP := range c.Template.Provisioners {
		// If we're skipping this, then ignore it
		if rawP.Skip(rawName) {
			continue
		}

		// Get the provisioner
		provisioner, err := c.components.Provisioner(rawP.Type)
		if err != nil {
			return nil, fmt.Errorf(
				"error initializing provisioner '%s': %s",
				rawP.Type, err)
		}
		if provisioner == nil {
			return nil, fmt.Errorf(
				"provisioner type not found: %s", rawP.Type)
		}

		// Get the configuration
		config := make([]interface{}, 1, 2)
		config[0] = rawP.Config
		if rawP.Override != nil {
			if override, ok := rawP.Override[rawName]; ok {
				config = append(config, override)
			}
		}

		// If we're pausing, we wrap the provisioner in a special pauser.
		if rawP.PauseBefore > 0 {
			provisioner = &PausedProvisioner{
				PauseBefore: rawP.PauseBefore,
				Provisioner: provisioner,
			}
		}

		provisioners = append(provisioners, coreBuildProvisioner{
			pType:       rawP.Type,
			provisioner: provisioner,
			config:      config,
		})
	}

	// Setup the post-processors
	postProcessors := make([][]coreBuildPostProcessor, 0, len(c.Template.PostProcessors))
	for _, rawPs := range c.Template.PostProcessors {
		current := make([]coreBuildPostProcessor, 0, len(rawPs))
		for _, rawP := range rawPs {
			// If we skip, ignore
			if rawP.Skip(rawName) {
				continue
			}

			// Get the post-processor
			postProcessor, err := c.components.PostProcessor(rawP.Type)
			if err != nil {
				return nil, fmt.Errorf(
					"error initializing post-processor '%s': %s",
					rawP.Type, err)
			}
			if postProcessor == nil {
				return nil, fmt.Errorf(
					"post-processor type not found: %s", rawP.Type)
			}

			current = append(current, coreBuildPostProcessor{
				processor:         postProcessor,
				processorType:     rawP.Type,
				config:            rawP.Config,
				keepInputArtifact: rawP.KeepInputArtifact,
			})
		}

		// If we have no post-processors in this chain, just continue.
		if len(current) == 0 {
			continue
		}

		postProcessors = append(postProcessors, current)
	}

	// TODO hooks one day

	return &coreBuild{
		name:           n,
		builder:        builder,
		builderConfig:  configBuilder.Config,
		builderType:    configBuilder.Type,
		postProcessors: postProcessors,
		provisioners:   provisioners,
		templatePath:   c.Template.Path,
		variables:      c.variables,
	}, nil
}

// Context returns an interpolation context.
func (c *Core) Context() *interpolate.Context {
	return &interpolate.Context{
		TemplatePath:  c.Template.Path,
		UserVariables: c.variables,
	}
}

// validate does a full validation of the template.
//
// This will automatically call template.validate() in addition to doing
// richer semantic checks around variables and so on.
func (c *Core) validate() error {
	// First validate the template in general, we can't do anything else
	// unless the template itself is valid.
	if err := c.Template.Validate(); err != nil {
		return err
	}

	// Validate the minimum version is satisfied
	if c.Template.MinVersion != "" {
		versionActual, err := version.NewVersion(c.version)
		if err != nil {
			// This shouldn't happen since we set it via the compiler
			panic(err)
		}

		versionMin, err := version.NewVersion(c.Template.MinVersion)
		if err != nil {
			return fmt.Errorf(
				"min_version is invalid: %s", err)
		}

		if versionActual.LessThan(versionMin) {
			return fmt.Errorf(
				"This template requires Packer version %s or higher; using %s",
				versionMin,
				versionActual)
		}
	}

	// Validate variables are set
	var err error
	for n, v := range c.Template.Variables {
		if v.Required {
			if _, ok := c.variables[n]; !ok {
				err = multierror.Append(err, fmt.Errorf(
					"required variable not set: %s", n))
			}
		}
	}

	// TODO: validate all builders exist
	// TODO: ^^ provisioner
	// TODO: ^^ post-processor

	return err
}

func (c *Core) init() error {
	if c.variables == nil {
		c.variables = make(map[string]string)
	}

	// Go through the variables and interpolate the environment variables
	ctx := c.Context()
	ctx.EnableEnv = true
	ctx.UserVariables = nil
	for k, v := range c.Template.Variables {
		// Ignore variables that are required
		if v.Required {
			continue
		}

		// Ignore variables that have a value
		if _, ok := c.variables[k]; ok {
			continue
		}

		// Interpolate the default
		def, err := interpolate.Render(v.Default, ctx)
		if err != nil {
			return fmt.Errorf(
				"error interpolating default value for '%s': %s",
				k, err)
		}

		c.variables[k] = def
	}

	// Interpolate the push configuration
	if _, err := interpolate.RenderInterface(&c.Template.Push, c.Context()); err != nil {
		return fmt.Errorf("Error interpolating 'push': %s", err)
	}

	return nil
}
<|endoftext|>


<commit_before>// Copyright (c) Ilia Kravets, 2015. All rights reserved. PROVIDED "AS IS"
// WITHOUT ANY WARRANTY, EXPRESS OR IMPLIED. See LICENSE file for details.

package anal

import (
	"sort"

	"my/errs"

	"my/ev/packet"
	"my/ev/sim"
)

type bookStat struct {
	maxLevels int
}

type HashFunc func(uint64) uint64

type orderHashStat struct {
	f             HashFunc
	bucketSize    map[uint64]int
	maxBucketSize map[uint64]int
}
type Analyzer struct {
	observer      observer
	bookStats     map[OptionSide]*bookStat
	orderHashStat []orderHashStat
}

func NewAnalyzer() *Analyzer {
	a := &Analyzer{
		bookStats: make(map[OptionSide]*bookStat),
	}
	a.observer.analyzer = a
	return a
}

type observer struct {
	analyzer *Analyzer
}

var _ sim.Observer = &observer{}

func (*observer) MessageArrived(*sim.SimMessage) {}
func (o *observer) OperationAppliedToOrders(op sim.SimOperation) {
	sess := op.GetMessage().Session.Index()
	//errs.Check(sess < 4, sess)
	errs.Check(sess < 32, sess)
	var key uint64
	var delta int
	switch o := op.(type) {
	case *sim.OperationAdd:
		key = uint64(sess)<<32 | uint64(o.OrderId.ToUint32())
		delta = 1
	case *sim.OperationRemove:
		key = uint64(sess)<<32 | uint64(o.GetOrigOrderId().ToUint32())
		delta = -1
	}
	for _, ohs := range o.analyzer.orderHashStat {
		keyHash := ohs.f(key)
		ohs.bucketSize[keyHash] += delta
		if ohs.bucketSize[keyHash] > ohs.maxBucketSize[keyHash] {
			ohs.maxBucketSize[keyHash] = ohs.bucketSize[keyHash]
		} else if ohs.bucketSize[keyHash] == 0 {
			delete(ohs.bucketSize, keyHash)
		}
	}
}
func (*observer) BeforeBookUpdate(sim.Book, sim.SimOperation) {}
func (o *observer) AfterBookUpdate(book sim.Book, op sim.SimOperation) {
	oid := op.GetOptionId()
	if oid.Invalid() {
		return
	}
	bs := o.analyzer.book(oid, op.GetSide())
	//bookSize := len(book.GetTop(oid, op.GetSide(), 0))
	b := book.GetTop(oid, op.GetSide(), 0)
	bookSize := len(b)
	if bs.maxLevels < bookSize {
		//log.Printf("%d %s %d: %v\n", oid, op.GetSide(), bookSize, b)
		bs.maxLevels = bookSize
	}
}

func (a *Analyzer) Observer() sim.Observer {
	return &a.observer
}
func (a *Analyzer) AddOrderHashFunction(f HashFunc) {
	ohs := orderHashStat{
		f:             f,
		bucketSize:    make(map[uint64]int),
		maxBucketSize: make(map[uint64]int),
	}
	a.orderHashStat = append(a.orderHashStat, ohs)
}

func (a *Analyzer) book(oid packet.OptionId, side packet.MarketSide) (bs *bookStat) {
	errs.Check(side == packet.MarketSideBid || side == packet.MarketSideAsk)
	key := OptionSide{
		Oid:  oid,
		Side: side,
	}
	var ok bool
	if bs, ok = a.bookStats[key]; !ok {
		bs = &bookStat{}
		a.bookStats[key] = bs
	}
	return
}

type OptionSide struct {
	Oid  packet.OptionId
	Side packet.MarketSide
}
type BSVal struct {
	Levels       int
	OptionNumber int
	Sample       []OptionSide
}
type BSHist []BSVal

func (a *Analyzer) BookSizeHist() BSHist {
	bsv := make(map[int]BSVal)
	var levels []int
	for k, bs := range a.bookStats {
		v := bsv[bs.maxLevels]
		v.OptionNumber++
		if len(v.Sample) < 10 {
			v.Sample = append(v.Sample, OptionSide{
				Oid:  k.Oid,
				Side: k.Side,
			})
		}
		if v.OptionNumber == 1 {
			v.Levels = bs.maxLevels
			levels = append(levels, v.Levels)
		}
		bsv[v.Levels] = v
	}
	sort.Ints(levels)

	var bsh BSHist
	for _, l := range levels {
		bsh = append(bsh, bsv[l])
	}
	return bsh
}

type HistVal struct {
	Bin   int
	Count int
}
type Hist []HistVal

func (a *Analyzer) OrdersHashCollisionHist() []Hist {
	var hists []Hist
	for _, ohs := range a.orderHashStat {
		collisionHist := make(map[int]int)
		for _, c := range ohs.maxBucketSize {
			collisionHist[c]++
		}
		var chKeys []int
		for k, _ := range collisionHist {
			chKeys = append(chKeys, k)
		}
		sort.Ints(chKeys)
		var hist Hist
		for _, k := range chKeys {
			hist = append(hist, HistVal{k, collisionHist[k]})
		}
		hists = append(hists, hist)
	}
	return hists
}
<commit_msg>anal: do not limit session number<commit_after>// Copyright (c) Ilia Kravets, 2015. All rights reserved. PROVIDED "AS IS"
// WITHOUT ANY WARRANTY, EXPRESS OR IMPLIED. See LICENSE file for details.

package anal

import (
	"sort"

	"my/errs"

	"my/ev/packet"
	"my/ev/sim"
)

type bookStat struct {
	maxLevels int
}

type HashFunc func(uint64) uint64

type orderHashStat struct {
	f             HashFunc
	bucketSize    map[uint64]int
	maxBucketSize map[uint64]int
}
type Analyzer struct {
	observer      observer
	bookStats     map[OptionSide]*bookStat
	orderHashStat []orderHashStat
}

func NewAnalyzer() *Analyzer {
	a := &Analyzer{
		bookStats: make(map[OptionSide]*bookStat),
	}
	a.observer.analyzer = a
	return a
}

type observer struct {
	analyzer *Analyzer
}

var _ sim.Observer = &observer{}

func (*observer) MessageArrived(*sim.SimMessage) {}
func (o *observer) OperationAppliedToOrders(op sim.SimOperation) {
	sess := op.GetMessage().Session.Index()
	//errs.Check(sess < 4, sess)
	//errs.Check(sess < 32, sess)
	var key uint64
	var delta int
	switch o := op.(type) {
	case *sim.OperationAdd:
		key = uint64(sess)<<32 | uint64(o.OrderId.ToUint32())
		delta = 1
	case *sim.OperationRemove:
		key = uint64(sess)<<32 | uint64(o.GetOrigOrderId().ToUint32())
		delta = -1
	}
	for _, ohs := range o.analyzer.orderHashStat {
		keyHash := ohs.f(key)
		ohs.bucketSize[keyHash] += delta
		if ohs.bucketSize[keyHash] > ohs.maxBucketSize[keyHash] {
			ohs.maxBucketSize[keyHash] = ohs.bucketSize[keyHash]
		} else if ohs.bucketSize[keyHash] == 0 {
			delete(ohs.bucketSize, keyHash)
		}
	}
}
func (*observer) BeforeBookUpdate(sim.Book, sim.SimOperation) {}
func (o *observer) AfterBookUpdate(book sim.Book, op sim.SimOperation) {
	oid := op.GetOptionId()
	if oid.Invalid() {
		return
	}
	bs := o.analyzer.book(oid, op.GetSide())
	//bookSize := len(book.GetTop(oid, op.GetSide(), 0))
	b := book.GetTop(oid, op.GetSide(), 0)
	bookSize := len(b)
	if bs.maxLevels < bookSize {
		//log.Printf("%d %s %d: %v\n", oid, op.GetSide(), bookSize, b)
		bs.maxLevels = bookSize
	}
}

func (a *Analyzer) Observer() sim.Observer {
	return &a.observer
}
func (a *Analyzer) AddOrderHashFunction(f HashFunc) {
	ohs := orderHashStat{
		f:             f,
		bucketSize:    make(map[uint64]int),
		maxBucketSize: make(map[uint64]int),
	}
	a.orderHashStat = append(a.orderHashStat, ohs)
}

func (a *Analyzer) book(oid packet.OptionId, side packet.MarketSide) (bs *bookStat) {
	errs.Check(side == packet.MarketSideBid || side == packet.MarketSideAsk)
	key := OptionSide{
		Oid:  oid,
		Side: side,
	}
	var ok bool
	if bs, ok = a.bookStats[key]; !ok {
		bs = &bookStat{}
		a.bookStats[key] = bs
	}
	return
}

type OptionSide struct {
	Oid  packet.OptionId
	Side packet.MarketSide
}
type BSVal struct {
	Levels       int
	OptionNumber int
	Sample       []OptionSide
}
type BSHist []BSVal

func (a *Analyzer) BookSizeHist() BSHist {
	bsv := make(map[int]BSVal)
	var levels []int
	for k, bs := range a.bookStats {
		v := bsv[bs.maxLevels]
		v.OptionNumber++
		if len(v.Sample) < 10 {
			v.Sample = append(v.Sample, OptionSide{
				Oid:  k.Oid,
				Side: k.Side,
			})
		}
		if v.OptionNumber == 1 {
			v.Levels = bs.maxLevels
			levels = append(levels, v.Levels)
		}
		bsv[v.Levels] = v
	}
	sort.Ints(levels)

	var bsh BSHist
	for _, l := range levels {
		bsh = append(bsh, bsv[l])
	}
	return bsh
}

type HistVal struct {
	Bin   int
	Count int
}
type Hist []HistVal

func (a *Analyzer) OrdersHashCollisionHist() []Hist {
	var hists []Hist
	for _, ohs := range a.orderHashStat {
		collisionHist := make(map[int]int)
		for _, c := range ohs.maxBucketSize {
			collisionHist[c]++
		}
		var chKeys []int
		for k, _ := range collisionHist {
			chKeys = append(chKeys, k)
		}
		sort.Ints(chKeys)
		var hist Hist
		for _, k := range chKeys {
			hist = append(hist, HistVal{k, collisionHist[k]})
		}
		hists = append(hists, hist)
	}
	return hists
}
<|endoftext|>


<commit_before>package main

import (
	"time"
	"flag"

	"github.com/mitchellh/cli"
	"github.com/aws/aws-sdk-go/service/ec2"
	//"os"
	"github.com/aws/aws-sdk-go/aws"
)

type CreateCommand struct {
	Ui 			cli.Ui
	InstanceId 	string
	Name 		string
	DryRun		bool
	NoReboot	bool
}

// descriptions for args
var createDscrInstanceId = "The instance from which to create the AMI"
var createDscrName = "The name of the AMI; the current timestamp will be automatically appended"
var createDscrDryRun = "Execute a simulated run"
var createDscrNoReboot = "If true, do not reboot the instance before creating the AMI. It is preferable to reboot the instance to guarantee a consistent filesystem when taking the snapshot, but the likelihood of an inconsistent snapshot is very low."

func (c *CreateCommand) Help() string {
	return `ec2-snapper create <args> [--help]

Create an AMI of the given EC2 instance.

Available args are:
--instance      ` + createDscrInstanceId + `
--name          ` + createDscrName + `
--dry-run       ` + createDscrDryRun + `
--no-reboot     ` + createDscrNoReboot
}

func (c *CreateCommand) Synopsis() string {
	return "Create an AMI of the given EC2 instance"
}

func (c *CreateCommand) Run(args []string) int {

	// Handle the command-line args
	cmdFlags := flag.NewFlagSet("create", flag.ExitOnError)
	cmdFlags.Usage = func() { c.Ui.Output(c.Help()) }

	cmdFlags.StringVar(&c.InstanceId, "instance", "", createDscrInstanceId)
	cmdFlags.StringVar(&c.Name, "name", "", createDscrName)
	cmdFlags.BoolVar(&c.DryRun, "dry-run", false, createDscrDryRun)
	cmdFlags.BoolVar(&c.NoReboot, "no-reboot", true, createDscrNoReboot)

	if err := cmdFlags.Parse(args); err != nil {
		return 1
	}

	// Check for required command-line args
	if c.InstanceId == "" {
		c.Ui.Error("ERROR: The argument '--instance' is required.")
		return 1
	}

	if c.Name == "" {
		c.Ui.Error("ERROR: The argument '--name' is required.")
		return 1
	}

	// Create an EC2 service object; AWS region is picked up from the "AWS_REGION" env var.
	svc := ec2.New(nil)

	// Generate a nicely formatted timestamp for right now
	const dateLayoutForAmiName = "2006-01-02 at 15_04_05 (MST)"
	t := time.Now()

	// Create the AMI Snapshot
	name := c.Name + " " + t.Format(dateLayoutForAmiName)

	c.Ui.Output("==> Creating AMI for " + c.InstanceId + "...")

	resp, err := svc.CreateImage(&ec2.CreateImageInput{
		Name: &name,
		InstanceID: &c.InstanceId,
		DryRun: &c.DryRun,
		NoReboot: &c.NoReboot })
	if err != nil {
		panic(err)
	}

	// Assign tags to this AMI.  We'll use these when it comes time to delete the AMI
	c.Ui.Output("==> Adding tags to AMI " + *resp.ImageID + "...")

	dateLayoutForTagValue := time.Now().Format(time.RFC3339)

	svc.CreateTags(&ec2.CreateTagsInput{
		Resources: []*string{resp.ImageID},
		Tags: []*ec2.Tag{
			&ec2.Tag{ Key: aws.String("ec2-snapper-instance-id"), Value: &c.InstanceId },
			&ec2.Tag{ Key: aws.String("ec2-snapper-snapshot-date"), Value: &dateLayoutForTagValue },
		},
	})

	c.Ui.Info("==> Success! Created " + *resp.ImageID + " named \"" + name + "\"")
	return 0
}

<commit_msg>Deleted extraneous "ec2-snapper-snapshot-date" tag<commit_after>package main

import (
	"time"
	"flag"

	"github.com/mitchellh/cli"
	"github.com/aws/aws-sdk-go/service/ec2"
	//"os"
	"github.com/aws/aws-sdk-go/aws"
)

type CreateCommand struct {
	Ui 			cli.Ui
	InstanceId 	string
	Name 		string
	DryRun		bool
	NoReboot	bool
}

// descriptions for args
var createDscrInstanceId = "The instance from which to create the AMI"
var createDscrName = "The name of the AMI; the current timestamp will be automatically appended"
var createDscrDryRun = "Execute a simulated run"
var createDscrNoReboot = "If true, do not reboot the instance before creating the AMI. It is preferable to reboot the instance to guarantee a consistent filesystem when taking the snapshot, but the likelihood of an inconsistent snapshot is very low."

func (c *CreateCommand) Help() string {
	return `ec2-snapper create <args> [--help]

Create an AMI of the given EC2 instance.

Available args are:
--instance      ` + createDscrInstanceId + `
--name          ` + createDscrName + `
--dry-run       ` + createDscrDryRun + `
--no-reboot     ` + createDscrNoReboot
}

func (c *CreateCommand) Synopsis() string {
	return "Create an AMI of the given EC2 instance"
}

func (c *CreateCommand) Run(args []string) int {

	// Handle the command-line args
	cmdFlags := flag.NewFlagSet("create", flag.ExitOnError)
	cmdFlags.Usage = func() { c.Ui.Output(c.Help()) }

	cmdFlags.StringVar(&c.InstanceId, "instance", "", createDscrInstanceId)
	cmdFlags.StringVar(&c.Name, "name", "", createDscrName)
	cmdFlags.BoolVar(&c.DryRun, "dry-run", false, createDscrDryRun)
	cmdFlags.BoolVar(&c.NoReboot, "no-reboot", true, createDscrNoReboot)

	if err := cmdFlags.Parse(args); err != nil {
		return 1
	}

	// Check for required command-line args
	if c.InstanceId == "" {
		c.Ui.Error("ERROR: The argument '--instance' is required.")
		return 1
	}

	if c.Name == "" {
		c.Ui.Error("ERROR: The argument '--name' is required.")
		return 1
	}

	// Create an EC2 service object; AWS region is picked up from the "AWS_REGION" env var.
	svc := ec2.New(nil)

	// Generate a nicely formatted timestamp for right now
	const dateLayoutForAmiName = "2006-01-02 at 15_04_05 (MST)"
	t := time.Now()

	// Create the AMI Snapshot
	name := c.Name + " " + t.Format(dateLayoutForAmiName)

	c.Ui.Output("==> Creating AMI for " + c.InstanceId + "...")

	resp, err := svc.CreateImage(&ec2.CreateImageInput{
		Name: &name,
		InstanceID: &c.InstanceId,
		DryRun: &c.DryRun,
		NoReboot: &c.NoReboot })
	if err != nil {
		panic(err)
	}

	// Assign tags to this AMI.  We'll use these when it comes time to delete the AMI
	c.Ui.Output("==> Adding tags to AMI " + *resp.ImageID + "...")

	dateLayoutForTagValue := time.Now().Format(time.RFC3339)

	svc.CreateTags(&ec2.CreateTagsInput{
		Resources: []*string{resp.ImageID},
		Tags: []*ec2.Tag{
			&ec2.Tag{ Key: aws.String("ec2-snapper-instance-id"), Value: &c.InstanceId },
		},
	})

	c.Ui.Info("==> Success! Created " + *resp.ImageID + " named \"" + name + "\"")
	return 0
}

<|endoftext|>


<commit_before>package commands

import (
	"fmt"
	"io"
	"runtime"

	version "github.com/ipfs/go-ipfs"
	fsrepo "github.com/ipfs/go-ipfs/repo/fsrepo"

	cmds "gx/ipfs/QmSXUokcP4TJpFfqozT69AVAYRtzXVMUjzQVkYX41R9Svs/go-ipfs-cmds"
	"gx/ipfs/Qmde5VP1qUkyQXKCfmEUA7bP64V2HAptbJ7phuPp7jXWwg/go-ipfs-cmdkit"
)

type VersionOutput struct {
	Version string
	Commit  string
	Repo    string
	System  string
	Golang  string
}

const (
	versionNumberOptionName = "number"
	versionCommitOptionName = "commit"
	versionRepoOptionName   = "repo"
	versionAllOptionName    = "all"
)

var VersionCmd = &cmds.Command{
	Helptext: cmdkit.HelpText{
		Tagline:          "Show ipfs version information.",
		ShortDescription: "Returns the current version of ipfs and exits.",
	},

	Options: []cmdkit.Option{
		cmdkit.BoolOption(versionNumberOptionName, "n", "Only show the version number."),
		cmdkit.BoolOption(versionCommitOptionName, "Show the commit hash."),
		cmdkit.BoolOption(versionRepoOptionName, "Show repo version."),
		cmdkit.BoolOption(versionAllOptionName, "Show all version information"),
	},
	Run: func(req *cmds.Request, res cmds.ResponseEmitter, env cmds.Environment) error {
		return cmds.EmitOnce(res, &VersionOutput{
			Version: version.CurrentVersionNumber,
			Commit:  version.CurrentCommit,
			Repo:    fmt.Sprint(fsrepo.RepoVersion),
			System:  runtime.GOARCH + "/" + runtime.GOOS, //TODO: Precise version here
			Golang:  runtime.Version(),
		})
	},
	Encoders: cmds.EncoderMap{
		cmds.Text: cmds.MakeTypedEncoder(func(req *cmds.Request, w io.Writer, version *VersionOutput) error {
			repo, _ := req.Options[versionRepoOptionName].(bool)
			if repo {
				fmt.Fprintln(w, version.Repo)
				return nil
			}

			commit, _ := req.Options[versionCommitOptionName].(bool)
			commitTxt := ""
			if commit {
				commitTxt = "-" + version.Commit
			}

			number, _ := req.Options[versionNumberOptionName].(bool)
			if number {
				fmt.Fprintln(w, version.Version+commitTxt)
				return nil
			}

			all, _ := req.Options[versionAllOptionName].(bool)
			if all {
				out := fmt.Sprintf("go-ipfs version: %s-%s\n"+
					"Repo version: %s\nSystem version: %s\nGolang version: %s\n",
					version.Version, version.Commit, version.Repo, version.System, version.Golang)
				fmt.Fprint(w, out)
				return nil
			}

			fmt.Fprint(w, fmt.Sprintf("ipfs version %s%s\n", version.Version, commitTxt))
			return nil
		}),
	},
	Type: VersionOutput{},
}
<commit_msg>fix(flag/version): flag `all` should have a higher priority<commit_after>package commands

import (
	"fmt"
	"io"
	"runtime"

	version "github.com/ipfs/go-ipfs"
	fsrepo "github.com/ipfs/go-ipfs/repo/fsrepo"

	cmds "gx/ipfs/QmSXUokcP4TJpFfqozT69AVAYRtzXVMUjzQVkYX41R9Svs/go-ipfs-cmds"
	"gx/ipfs/Qmde5VP1qUkyQXKCfmEUA7bP64V2HAptbJ7phuPp7jXWwg/go-ipfs-cmdkit"
)

type VersionOutput struct {
	Version string
	Commit  string
	Repo    string
	System  string
	Golang  string
}

const (
	versionNumberOptionName = "number"
	versionCommitOptionName = "commit"
	versionRepoOptionName   = "repo"
	versionAllOptionName    = "all"
)

var VersionCmd = &cmds.Command{
	Helptext: cmdkit.HelpText{
		Tagline:          "Show ipfs version information.",
		ShortDescription: "Returns the current version of ipfs and exits.",
	},

	Options: []cmdkit.Option{
		cmdkit.BoolOption(versionNumberOptionName, "n", "Only show the version number."),
		cmdkit.BoolOption(versionCommitOptionName, "Show the commit hash."),
		cmdkit.BoolOption(versionRepoOptionName, "Show repo version."),
		cmdkit.BoolOption(versionAllOptionName, "Show all version information"),
	},
	Run: func(req *cmds.Request, res cmds.ResponseEmitter, env cmds.Environment) error {
		return cmds.EmitOnce(res, &VersionOutput{
			Version: version.CurrentVersionNumber,
			Commit:  version.CurrentCommit,
			Repo:    fmt.Sprint(fsrepo.RepoVersion),
			System:  runtime.GOARCH + "/" + runtime.GOOS, //TODO: Precise version here
			Golang:  runtime.Version(),
		})
	},
	Encoders: cmds.EncoderMap{
		cmds.Text: cmds.MakeTypedEncoder(func(req *cmds.Request, w io.Writer, version *VersionOutput) error {
			commit, _ := req.Options[versionCommitOptionName].(bool)
			commitTxt := ""
			if commit {
				commitTxt = "-" + version.Commit
			}

			all, _ := req.Options[versionAllOptionName].(bool)
			if all {
				out := fmt.Sprintf("go-ipfs version: %s-%s\n"+
					"Repo version: %s\nSystem version: %s\nGolang version: %s\n",
					version.Version, version.Commit, version.Repo, version.System, version.Golang)
				fmt.Fprint(w, out)
				return nil
			}

			repo, _ := req.Options[versionRepoOptionName].(bool)
			if repo {
				fmt.Fprintln(w, version.Repo)
				return nil
			}

			number, _ := req.Options[versionNumberOptionName].(bool)
			if number {
				fmt.Fprintln(w, version.Version+commitTxt)
				return nil
			}

			fmt.Fprint(w, fmt.Sprintf("ipfs version %s%s\n", version.Version, commitTxt))
			return nil
		}),
	},
	Type: VersionOutput{},
}
<|endoftext|>


<commit_before>package commands

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/exec"

	cmds "github.com/jbenet/go-ipfs/commands"
	config "github.com/jbenet/go-ipfs/config"
	u "github.com/jbenet/go-ipfs/util"
)

type ConfigField struct {
	Key   string
	Value interface{}
}

var configCmd = &cmds.Command{
	Helptext: cmds.HelpText{
		Tagline: "get and set IPFS config values",
		Synopsis: `
ipfs config <key>          - Get value of <key>
ipfs config <key> <value>  - Set value of <key> to <value>
ipfs config --show         - Show config file
ipfs config --edit         - Edit config file in $EDITOR
`,
		ShortDescription: `
ipfs config controls configuration variables. It works like 'git config'.
The configuration values are stored in a config file inside your IPFS
repository.`,
		LongDescription: `
ipfs config controls configuration variables. It works
much like 'git config'. The configuration values are stored in a config
file inside your IPFS repository.

EXAMPLES:

Get the value of the 'datastore.path' key:

  ipfs config datastore.path

Set the value of the 'datastore.path' key:

  ipfs config datastore.path ~/.go-ipfs/datastore
`,
	},

	Arguments: []cmds.Argument{
		cmds.StringArg("key", true, false, "The key of the config entry (e.g. \"Addresses.API\")"),
		cmds.StringArg("value", false, false, "The value to set the config entry to"),
	},
	Options: []cmds.Option{
		cmds.StringOption("show", "s", "Show config file"),
		cmds.StringOption("edit", "e", "Edit config file in $EDITOR"),
	},
	Run: func(req cmds.Request) (interface{}, error) {
		args := req.Arguments()

		key, ok := args[0].(string)
		if !ok {
			return nil, u.ErrCast()
		}

		filename, err := config.Filename(req.Context().ConfigRoot)
		if err != nil {
			return nil, err
		}

		var value string
		if len(args) == 2 {
			var ok bool
			value, ok = args[1].(string)
			if !ok {
				return nil, u.ErrCast()
			}

			return setConfig(filename, key, value)

		} else {
			return getConfig(filename, key)
		}
	},
	Marshallers: map[cmds.EncodingType]cmds.Marshaller{
		cmds.Text: func(res cmds.Response) ([]byte, error) {
			if len(res.Request().Arguments()) == 2 {
				return nil, nil // dont output anything
			}

			v := res.Output().(*ConfigField)
			buf, err := config.HumanOutput(v.Value)
			if err != nil {
				return nil, err
			}
			buf = append(buf, byte('\n'))
			return buf, nil
		},
	},
	Type: &ConfigField{},
	Subcommands: map[string]*cmds.Command{
		"show": configShowCmd,
		"edit": configEditCmd,
	},
}

var configShowCmd = &cmds.Command{
	Description: "Outputs the content of the config file",
	Help: `WARNING: Your private key is stored in the config file, and it will be
included in the output of this command.
`,

	Run: func(req cmds.Request) (interface{}, error) {
		filename, err := config.Filename(req.Context().ConfigRoot)
		if err != nil {
			return nil, err
		}

		return showConfig(filename)
	},
}

var configEditCmd = &cmds.Command{
	Description: "Opens the config file for editing in $EDITOR",
	Help: `To use 'ipfs config edit', you must have the $EDITOR environment
variable set to your preferred text editor.
`,

	Run: func(req cmds.Request) (interface{}, error) {
		filename, err := config.Filename(req.Context().ConfigRoot)
		if err != nil {
			return nil, err
		}

		return nil, editConfig(filename)
	},
}

func getConfig(filename string, key string) (*ConfigField, error) {
	value, err := config.ReadConfigKey(filename, key)
	if err != nil {
		return nil, fmt.Errorf("Failed to get config value: %s", err)
	}

	return &ConfigField{
		Key:   key,
		Value: value,
	}, nil
}

func setConfig(filename string, key, value string) (*ConfigField, error) {
	err := config.WriteConfigKey(filename, key, value)
	if err != nil {
		return nil, fmt.Errorf("Failed to set config value: %s", err)
	}

	return getConfig(filename, key)
}

func showConfig(filename string) (io.Reader, error) {
	// TODO maybe we should omit privkey so we don't accidentally leak it?

	data, err := ioutil.ReadFile(filename)
	if err != nil {
		return nil, err
	}

	return bytes.NewReader(data), nil
}

func editConfig(filename string) error {
	editor := os.Getenv("EDITOR")
	if editor == "" {
		return errors.New("ENV variable $EDITOR not set")
	}

	cmd := exec.Command("sh", "-c", editor+" "+filename)
	cmd.Stdin, cmd.Stdout, cmd.Stderr = os.Stdin, os.Stdout, os.Stderr
	return cmd.Run()
}
<commit_msg>cmd2 fixed config panic<commit_after>package commands

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/exec"

	cmds "github.com/jbenet/go-ipfs/commands"
	config "github.com/jbenet/go-ipfs/config"
	u "github.com/jbenet/go-ipfs/util"
)

type ConfigField struct {
	Key   string
	Value interface{}
}

var configCmd = &cmds.Command{
	Helptext: cmds.HelpText{
		Tagline: "get and set IPFS config values",
		Synopsis: `
ipfs config <key>          - Get value of <key>
ipfs config <key> <value>  - Set value of <key> to <value>
ipfs config --show         - Show config file
ipfs config --edit         - Edit config file in $EDITOR
`,
		ShortDescription: `
ipfs config controls configuration variables. It works like 'git config'.
The configuration values are stored in a config file inside your IPFS
repository.`,
		LongDescription: `
ipfs config controls configuration variables. It works
much like 'git config'. The configuration values are stored in a config
file inside your IPFS repository.

EXAMPLES:

Get the value of the 'datastore.path' key:

  ipfs config datastore.path

Set the value of the 'datastore.path' key:

  ipfs config datastore.path ~/.go-ipfs/datastore
`,
	},

	Arguments: []cmds.Argument{
		cmds.StringArg("key", true, false, "The key of the config entry (e.g. \"Addresses.API\")"),
		cmds.StringArg("value", false, false, "The value to set the config entry to"),
	},
	Options: []cmds.Option{
		cmds.StringOption("show", "s", "Show config file"),
		cmds.StringOption("edit", "e", "Edit config file in $EDITOR"),
	},
	Run: func(req cmds.Request) (interface{}, error) {
		args := req.Arguments()

		key, ok := args[0].(string)
		if !ok {
			return nil, u.ErrCast()
		}

		filename, err := config.Filename(req.Context().ConfigRoot)
		if err != nil {
			return nil, err
		}

		var value string
		if len(args) == 2 {
			var ok bool
			value, ok = args[1].(string)
			if !ok {
				return nil, u.ErrCast()
			}

			return setConfig(filename, key, value)

		} else {
			return getConfig(filename, key)
		}
	},
	Marshallers: map[cmds.EncodingType]cmds.Marshaller{
		cmds.Text: func(res cmds.Response) ([]byte, error) {
			if len(res.Request().Arguments()) == 2 {
				return nil, nil // dont output anything
			}

			v := res.Output()
			if v == nil {
				k := res.Request().Arguments()[0]
				return nil, fmt.Errorf("config does not contain key: %s", k)
			}
			vf, ok := v.(*ConfigField)
			if !ok {
				return nil, u.ErrCast()
			}

			buf, err := config.HumanOutput(vf.Value)
			if err != nil {
				return nil, err
			}
			buf = append(buf, byte('\n'))
			return buf, nil
		},
	},
	Type: &ConfigField{},
	Subcommands: map[string]*cmds.Command{
		"show": configShowCmd,
		"edit": configEditCmd,
	},
}

var configShowCmd = &cmds.Command{
	Description: "Outputs the content of the config file",
	Help: `WARNING: Your private key is stored in the config file, and it will be
included in the output of this command.
`,

	Run: func(req cmds.Request) (interface{}, error) {
		filename, err := config.Filename(req.Context().ConfigRoot)
		if err != nil {
			return nil, err
		}

		return showConfig(filename)
	},
}

var configEditCmd = &cmds.Command{
	Description: "Opens the config file for editing in $EDITOR",
	Help: `To use 'ipfs config edit', you must have the $EDITOR environment
variable set to your preferred text editor.
`,

	Run: func(req cmds.Request) (interface{}, error) {
		filename, err := config.Filename(req.Context().ConfigRoot)
		if err != nil {
			return nil, err
		}

		return nil, editConfig(filename)
	},
}

func getConfig(filename string, key string) (*ConfigField, error) {
	value, err := config.ReadConfigKey(filename, key)
	if err != nil {
		return nil, fmt.Errorf("Failed to get config value: %s", err)
	}

	return &ConfigField{
		Key:   key,
		Value: value,
	}, nil
}

func setConfig(filename string, key, value string) (*ConfigField, error) {
	err := config.WriteConfigKey(filename, key, value)
	if err != nil {
		return nil, fmt.Errorf("Failed to set config value: %s", err)
	}

	return getConfig(filename, key)
}

func showConfig(filename string) (io.Reader, error) {
	// TODO maybe we should omit privkey so we don't accidentally leak it?

	data, err := ioutil.ReadFile(filename)
	if err != nil {
		return nil, err
	}

	return bytes.NewReader(data), nil
}

func editConfig(filename string) error {
	editor := os.Getenv("EDITOR")
	if editor == "" {
		return errors.New("ENV variable $EDITOR not set")
	}

	cmd := exec.Command("sh", "-c", editor+" "+filename)
	cmd.Stdin, cmd.Stdout, cmd.Stderr = os.Stdin, os.Stdout, os.Stderr
	return cmd.Run()
}
<|endoftext|>


<commit_before>package transport_test

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"strings"
	"time"

	"code.cloudfoundry.org/garden"
	"github.com/concourse/concourse/atc/db/dbfakes"
	gconn "github.com/concourse/concourse/atc/worker/gclient/connection"
	"github.com/concourse/concourse/atc/worker/transport"
	"github.com/concourse/concourse/atc/worker/transport/transportfakes"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/tedsuo/rata"

	"github.com/concourse/concourse/atc/db"
	"github.com/concourse/retryhttp/retryhttpfakes"
)

var _ = Describe("hijackStreamer", func() {
	var (
		savedWorker          *dbfakes.FakeWorker
		savedWorkerAddress   string
		fakeDB               *transportfakes.FakeTransportDB
		fakeRoundTripper     *transportfakes.FakeRoundTripper
		fakeHijackableClient *retryhttpfakes.FakeHijackableClient
		hijackStreamer       gconn.HijackStreamer
		fakeRequestGenerator *transportfakes.FakeRequestGenerator
		handler              string
		params               rata.Params
		query                url.Values
		contentType          string
	)
	BeforeEach(func() {
		savedWorkerAddress = "some-garden-addr"

		savedWorker = new(dbfakes.FakeWorker)

		savedWorker.GardenAddrReturns(&savedWorkerAddress)
		savedWorker.ExpiresAtReturns(time.Now().Add(123 * time.Minute))
		savedWorker.StateReturns(db.WorkerStateRunning)

		fakeDB = new(transportfakes.FakeTransportDB)
		fakeDB.GetWorkerReturns(savedWorker, true, nil)

		fakeRequestGenerator = new(transportfakes.FakeRequestGenerator)

		fakeRoundTripper = new(transportfakes.FakeRoundTripper)
		fakeHijackableClient = new(retryhttpfakes.FakeHijackableClient)

		hijackStreamer = &transport.WorkerHijackStreamer{
			HttpClient:       &http.Client{Transport: fakeRoundTripper},
			HijackableClient: fakeHijackableClient,
			Req:              fakeRequestGenerator,
		}

		handler = "Ping"
		params = map[string]string{"param1": "value1"}
		contentType = "application/json"
		query = map[string][]string{"key": []string{"some", "values"}}

		request, err := http.NewRequest("POST", "http://example.url", strings.NewReader("some-request-body"))
		Expect(err).NotTo(HaveOccurred())
		fakeRequestGenerator.CreateRequestReturns(request, nil)
	})

	Describe("hijackStreamer #Stream", func() {
		var (
			body             io.Reader
			actualReadCloser io.ReadCloser
			streamErr        error
			httpResp         http.Response
			expectedString   string
		)

		BeforeEach(func() {
			expectedString = "some-example-string"
			body = strings.NewReader(expectedString)

			fakeRoundTripper.RoundTripReturns(&httpResp, nil)
		})

		JustBeforeEach(func() {
			actualReadCloser, streamErr = hijackStreamer.Stream(handler, body, params, query, contentType)
		})

		Context("when httpResponse is success", func() {
			BeforeEach(func() {
				httpResp = http.Response{StatusCode: http.StatusOK, Body: ioutil.NopCloser(body)}
			})

			It("returns response body", func() {
				actualBodyBytes, err := ioutil.ReadAll(actualReadCloser)
				Expect(err).ToNot(HaveOccurred())
				Expect(expectedString).To(Equal(string(actualBodyBytes)))
			})
		})

		Context("when httpResponse is not success", func() {
			var fakeBody *transportfakes.FakeReadCloser

			BeforeEach(func() {
				fakeBody = new(transportfakes.FakeReadCloser)
				httpResp = http.Response{StatusCode: http.StatusTeapot, Body: fakeBody}

				bodyBuf, _ := json.Marshal(garden.Error{Err: errors.New("some-error")})
				realBody := strings.NewReader(string(bodyBuf))
				fakeBody.ReadStub = func(buf []byte) (int, error) {
					Expect(fakeBody.CloseCallCount()).To(BeZero())
					return realBody.Read(buf)
				}
			})

			It("closes httpResp.Body and returns error", func() {
				Expect(actualReadCloser).To(BeNil())
				Expect(streamErr).To(HaveOccurred())
				Expect(fakeBody.CloseCallCount()).To(Equal(1))
				Expect(streamErr).To(MatchError("some-error"))
			})
		})

		Context("when httpResponse is not success with bad response", func() {
			var fakeBody *transportfakes.FakeReadCloser

			BeforeEach(func() {
				fakeBody = new(transportfakes.FakeReadCloser)
				httpResp = http.Response{StatusCode: http.StatusTeapot, Body: fakeBody}

				realBody := strings.NewReader("some-error")
				fakeBody.ReadStub = func(buf []byte) (int, error) {
					Expect(fakeBody.CloseCallCount()).To(BeZero())
					return realBody.Read(buf)
				}
			})

			It("closes httpResp.Body and returns bad response", func() {
				Expect(actualReadCloser).To(BeNil())
				Expect(streamErr).To(HaveOccurred())
				Expect(fakeBody.CloseCallCount()).To(Equal(1))
				Expect(streamErr).To(MatchError(fmt.Errorf("bad response: %s", errors.New("invalid character 's' looking for beginning of value"))))
			})
		})

		Context("when httpResponse fails", func() {
			BeforeEach(func() {
				httpResp = http.Response{StatusCode: http.StatusTeapot, Body: ioutil.NopCloser(body)}
			})

			It("returns error", func() {
				Expect(actualReadCloser).To(BeNil())
				Expect(streamErr).To(HaveOccurred())
			})

			It("creates request with the right arguments", func() {
				Expect(fakeRequestGenerator.CreateRequestCallCount()).To(Equal(1))
				actualHandler, actualParams, actualBody := fakeRequestGenerator.CreateRequestArgsForCall(0)
				Expect(actualHandler).To(Equal(handler))
				Expect(actualParams).To(Equal(params))
				Expect(actualBody).To(Equal(body))
			})

			It("httpClient makes the right request", func() {
				expectedRequest, err := http.NewRequest("POST", "http://example.url", strings.NewReader("some-request-body"))
				expectedRequest.Header.Add("Content-Type", "application/json")
				expectedRequest.URL.RawQuery = "key=some&key=values"
				Expect(err).NotTo(HaveOccurred())
				Expect(fakeRoundTripper.RoundTripCallCount()).To(Equal(1))
				actualRequest := fakeRoundTripper.RoundTripArgsForCall(0)
				Expect(actualRequest.Method).To(Equal(expectedRequest.Method))
				Expect(actualRequest.URL).To(Equal(expectedRequest.URL))
				Expect(actualRequest.Header).To(Equal(expectedRequest.Header))

				s, err := ioutil.ReadAll(actualRequest.Body)
				Expect(err).NotTo(HaveOccurred())
				Expect(string(s)).To(Equal("some-request-body"))
			})

		})
	})

	Describe("hijackStreamer #Hijack", func() {
		var (
			body                   io.Reader
			hijackError            error
			httpResp               http.Response
			expectedString         string
			actualHijackedConn     net.Conn
			actualResponseReader   *bufio.Reader
			fakeHijackCloser       *retryhttpfakes.FakeHijackCloser
			expectedResponseReader *bufio.Reader
			fakeHijackedConn       *retryhttpfakes.FakeConn
		)

		BeforeEach(func() {
			expectedResponseReader = new(bufio.Reader)
			fakeHijackedConn = new(retryhttpfakes.FakeConn)

			expectedString = "some-example-string"
			body = strings.NewReader(expectedString)
			fakeHijackCloser = new(retryhttpfakes.FakeHijackCloser)
		})

		JustBeforeEach(func() {
			actualHijackedConn, actualResponseReader, hijackError = hijackStreamer.Hijack(context.TODO(), handler, body, params, query, contentType)
		})

		Context("when request is successful", func() {
			BeforeEach(func() {
				fakeHijackableClient.DoReturns(&httpResp, fakeHijackCloser, nil)
				httpResp = http.Response{StatusCode: http.StatusOK}
				fakeHijackCloser.HijackReturns(fakeHijackedConn, expectedResponseReader)
			})

			It("returns success response and hijackCloser", func() {
				Expect(hijackError).ToNot(HaveOccurred())
				Expect(fakeHijackCloser.HijackCallCount()).To(Equal(1))
				Expect(actualHijackedConn).To(Equal(fakeHijackedConn))
				Expect(actualResponseReader).To(Equal(expectedResponseReader))
			})
		})

		Context("when httpResponse is not success", func() {
			var fakeBody *transportfakes.FakeReadCloser

			BeforeEach(func() {
				fakeHijackableClient.DoReturns(&httpResp, fakeHijackCloser, nil)
				fakeBody = new(transportfakes.FakeReadCloser)
				httpResp = http.Response{StatusCode: http.StatusTeapot, Body: fakeBody}

				realBody := strings.NewReader("some-error")
				fakeBody.ReadStub = func(buf []byte) (int, error) {
					Expect(fakeBody.CloseCallCount()).To(BeZero())
					return realBody.Read(buf)
				}
			})

			It("closes httpResp.Body, hijackCloser and returns error", func() {
				Expect(fakeHijackCloser).NotTo(BeNil())
				Expect(hijackError).To(HaveOccurred())
				Expect(fakeBody.CloseCallCount()).To(Equal(1))
				Expect(hijackError).To(MatchError(fmt.Errorf("Backend error: Exit status: %d, message: %s", httpResp.StatusCode, "some-error")))
				Expect(fakeHijackCloser.CloseCallCount()).To(Equal(1))
			})
		})

		Context("when httpResponse fails with ExecutableNotFoundError", func() {
			var fakeBody *transportfakes.FakeReadCloser
			gerr := garden.ExecutableNotFoundError{Message: "sorry, couldn't find it"}

			BeforeEach(func() {
				fakeHijackableClient.DoReturns(&httpResp, fakeHijackCloser, nil)
				fakeBody = new(transportfakes.FakeReadCloser)
				httpResp = http.Response{StatusCode: http.StatusTeapot, Body: fakeBody}

				jsonBody, _ := garden.Error{Err: gerr}.MarshalJSON()
				realBody := bytes.NewReader(jsonBody)
				fakeBody.ReadStub = func(buf []byte) (int, error) {
					Expect(fakeBody.CloseCallCount()).To(BeZero())
					return realBody.Read(buf)
				}
			})

			It("closes httpResp.Body, hijackCloser and returns ExecutableNotFoundError", func() {
				Expect(fakeHijackCloser).NotTo(BeNil())
				Expect(hijackError).To(HaveOccurred())
				Expect(fakeBody.CloseCallCount()).To(Equal(1))
				Expect(hijackError).To(MatchError(gerr))
				Expect(fakeHijackCloser.CloseCallCount()).To(Equal(1))
			})
		})

		Context("when httpResponse is not success with bad response", func() {
			var fakeBody *transportfakes.FakeReadCloser

			BeforeEach(func() {
				fakeHijackableClient.DoReturns(&httpResp, fakeHijackCloser, nil)
				fakeBody = new(transportfakes.FakeReadCloser)
				httpResp = http.Response{StatusCode: http.StatusTeapot, Body: fakeBody}

				fakeBody.ReadStub = func(buf []byte) (int, error) {
					Expect(fakeBody.CloseCallCount()).To(BeZero())
					return 0, errors.New("error reading")
				}
			})

			It("closes httpResp.Body and returns bad response", func() {
				Expect(fakeHijackCloser).NotTo(BeNil())
				Expect(hijackError).To(HaveOccurred())
				Expect(fakeBody.CloseCallCount()).To(Equal(1))
				Expect(hijackError).To(MatchError(fmt.Errorf("Backend error: Exit status: %d, error reading response body: %s", httpResp.StatusCode, "error reading")))
				Expect(fakeHijackCloser.CloseCallCount()).To(Equal(1))
			})
		})

		Context("when httpResponse fails", func() {
			BeforeEach(func() {
				fakeHijackableClient.DoReturns(nil, nil, errors.New("Request failed"))
				httpResp = http.Response{StatusCode: http.StatusTeapot, Body: ioutil.NopCloser(body)}
			})

			It("returns error", func() {
				Expect(hijackError).To(HaveOccurred())
				Expect(actualHijackedConn).To(BeNil())
				Expect(actualResponseReader).To(BeNil())
			})

			It("makes the right request", func() {
				expectedRequest, err := http.NewRequest("POST", "http://example.url", strings.NewReader("some-request-body"))
				expectedRequest.Header.Add("Content-Type", "application/json")
				expectedRequest.URL.RawQuery = "key=some&key=values"
				Expect(err).NotTo(HaveOccurred())
				Expect(fakeHijackableClient.DoCallCount()).To(Equal(1))
				actualRequest := fakeHijackableClient.DoArgsForCall(0)

				Expect(actualRequest.Method).To(Equal(expectedRequest.Method))
				Expect(actualRequest.URL).To(Equal(expectedRequest.URL))
				Expect(actualRequest.Header).To(Equal(expectedRequest.Header))

				s, err := ioutil.ReadAll(actualRequest.Body)
				Expect(err).NotTo(HaveOccurred())
				Expect(string(s)).To(Equal("some-request-body"))
			})
		})
	})
})
<commit_msg>Fix hijack unit tests<commit_after>package transport_test

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"strings"
	"time"

	"code.cloudfoundry.org/garden"
	"github.com/concourse/concourse/atc/db/dbfakes"
	gconn "github.com/concourse/concourse/atc/worker/gclient/connection"
	"github.com/concourse/concourse/atc/worker/transport"
	"github.com/concourse/concourse/atc/worker/transport/transportfakes"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/tedsuo/rata"

	"github.com/concourse/concourse/atc/db"
	"github.com/concourse/retryhttp/retryhttpfakes"
)

var _ = Describe("hijackStreamer", func() {
	var (
		savedWorker          *dbfakes.FakeWorker
		savedWorkerAddress   string
		fakeDB               *transportfakes.FakeTransportDB
		fakeRoundTripper     *transportfakes.FakeRoundTripper
		fakeHijackableClient *retryhttpfakes.FakeHijackableClient
		hijackStreamer       gconn.HijackStreamer
		fakeRequestGenerator *transportfakes.FakeRequestGenerator
		handler              string
		params               rata.Params
		query                url.Values
		contentType          string
	)
	BeforeEach(func() {
		savedWorkerAddress = "some-garden-addr"

		savedWorker = new(dbfakes.FakeWorker)

		savedWorker.GardenAddrReturns(&savedWorkerAddress)
		savedWorker.ExpiresAtReturns(time.Now().Add(123 * time.Minute))
		savedWorker.StateReturns(db.WorkerStateRunning)

		fakeDB = new(transportfakes.FakeTransportDB)
		fakeDB.GetWorkerReturns(savedWorker, true, nil)

		fakeRequestGenerator = new(transportfakes.FakeRequestGenerator)

		fakeRoundTripper = new(transportfakes.FakeRoundTripper)
		fakeHijackableClient = new(retryhttpfakes.FakeHijackableClient)

		hijackStreamer = &transport.WorkerHijackStreamer{
			HttpClient:       &http.Client{Transport: fakeRoundTripper},
			HijackableClient: fakeHijackableClient,
			Req:              fakeRequestGenerator,
		}

		handler = "Ping"
		params = map[string]string{"param1": "value1"}
		contentType = "application/json"
		query = map[string][]string{"key": []string{"some", "values"}}

		request, err := http.NewRequest("POST", "http://example.url", strings.NewReader("some-request-body"))
		Expect(err).NotTo(HaveOccurred())
		fakeRequestGenerator.CreateRequestReturns(request, nil)
	})

	Describe("hijackStreamer #Stream", func() {
		var (
			body             io.Reader
			actualReadCloser io.ReadCloser
			streamErr        error
			httpResp         http.Response
			expectedString   string
		)

		BeforeEach(func() {
			expectedString = "some-example-string"
			body = strings.NewReader(expectedString)

			fakeRoundTripper.RoundTripReturns(&httpResp, nil)
		})

		JustBeforeEach(func() {
			actualReadCloser, streamErr = hijackStreamer.Stream(handler, body, params, query, contentType)
		})

		Context("when httpResponse is success", func() {
			BeforeEach(func() {
				httpResp = http.Response{StatusCode: http.StatusOK, Body: ioutil.NopCloser(body)}
			})

			It("returns response body", func() {
				actualBodyBytes, err := ioutil.ReadAll(actualReadCloser)
				Expect(err).ToNot(HaveOccurred())
				Expect(expectedString).To(Equal(string(actualBodyBytes)))
			})
		})

		Context("when httpResponse is not success", func() {
			var fakeBody *transportfakes.FakeReadCloser

			BeforeEach(func() {
				fakeBody = new(transportfakes.FakeReadCloser)
				httpResp = http.Response{StatusCode: http.StatusTeapot, Body: fakeBody}

				bodyBuf, _ := json.Marshal(garden.Error{Err: errors.New("some-error")})
				realBody := strings.NewReader(string(bodyBuf))
				fakeBody.ReadStub = func(buf []byte) (int, error) {
					Expect(fakeBody.CloseCallCount()).To(BeZero())
					return realBody.Read(buf)
				}
			})

			It("closes httpResp.Body and returns error", func() {
				Expect(actualReadCloser).To(BeNil())
				Expect(streamErr).To(HaveOccurred())
				Expect(fakeBody.CloseCallCount()).To(Equal(1))
				Expect(streamErr).To(MatchError("some-error"))
			})
		})

		Context("when httpResponse is not success with bad response", func() {
			var fakeBody *transportfakes.FakeReadCloser

			BeforeEach(func() {
				fakeBody = new(transportfakes.FakeReadCloser)
				httpResp = http.Response{StatusCode: http.StatusTeapot, Body: fakeBody}

				realBody := strings.NewReader("some-error")
				fakeBody.ReadStub = func(buf []byte) (int, error) {
					Expect(fakeBody.CloseCallCount()).To(BeZero())
					return realBody.Read(buf)
				}
			})

			It("closes httpResp.Body and returns bad response", func() {
				Expect(actualReadCloser).To(BeNil())
				Expect(streamErr).To(HaveOccurred())
				Expect(fakeBody.CloseCallCount()).To(Equal(1))
				Expect(streamErr).To(MatchError(fmt.Errorf("bad response: %s", errors.New("invalid character 's' looking for beginning of value"))))
			})
		})

		Context("when httpResponse fails", func() {
			BeforeEach(func() {
				httpResp = http.Response{StatusCode: http.StatusTeapot, Body: ioutil.NopCloser(body)}
			})

			It("returns error", func() {
				Expect(actualReadCloser).To(BeNil())
				Expect(streamErr).To(HaveOccurred())
			})

			It("creates request with the right arguments", func() {
				Expect(fakeRequestGenerator.CreateRequestCallCount()).To(Equal(1))
				actualHandler, actualParams, actualBody := fakeRequestGenerator.CreateRequestArgsForCall(0)
				Expect(actualHandler).To(Equal(handler))
				Expect(actualParams).To(Equal(params))
				Expect(actualBody).To(Equal(body))
			})

			It("httpClient makes the right request", func() {
				expectedRequest, err := http.NewRequest("POST", "http://example.url", strings.NewReader("some-request-body"))
				expectedRequest.Header.Add("Content-Type", "application/json")
				expectedRequest.URL.RawQuery = "key=some&key=values"
				Expect(err).NotTo(HaveOccurred())
				Expect(fakeRoundTripper.RoundTripCallCount()).To(Equal(1))
				actualRequest := fakeRoundTripper.RoundTripArgsForCall(0)
				Expect(actualRequest.Method).To(Equal(expectedRequest.Method))
				Expect(actualRequest.URL).To(Equal(expectedRequest.URL))
				Expect(actualRequest.Header).To(Equal(expectedRequest.Header))

				s, err := ioutil.ReadAll(actualRequest.Body)
				Expect(err).NotTo(HaveOccurred())
				Expect(string(s)).To(Equal("some-request-body"))
			})

		})
	})

	Describe("hijackStreamer #Hijack", func() {
		var (
			body                   io.Reader
			hijackError            error
			httpResp               http.Response
			expectedString         string
			actualHijackedConn     net.Conn
			actualResponseReader   *bufio.Reader
			fakeHijackCloser       *retryhttpfakes.FakeHijackCloser
			expectedResponseReader *bufio.Reader
			fakeHijackedConn       *retryhttpfakes.FakeConn
		)

		BeforeEach(func() {
			expectedResponseReader = new(bufio.Reader)
			fakeHijackedConn = new(retryhttpfakes.FakeConn)

			expectedString = "some-example-string"
			body = strings.NewReader(expectedString)
			fakeHijackCloser = new(retryhttpfakes.FakeHijackCloser)
		})

		JustBeforeEach(func() {
			actualHijackedConn, actualResponseReader, hijackError = hijackStreamer.Hijack(context.TODO(), handler, body, params, query, contentType)
		})

		Context("when request is successful", func() {
			BeforeEach(func() {
				fakeHijackableClient.DoReturns(&httpResp, fakeHijackCloser, nil)
				httpResp = http.Response{StatusCode: http.StatusOK}
				fakeHijackCloser.HijackReturns(fakeHijackedConn, expectedResponseReader)
			})

			It("returns success response and hijackCloser", func() {
				Expect(hijackError).ToNot(HaveOccurred())
				Expect(fakeHijackCloser.HijackCallCount()).To(Equal(1))
				Expect(actualHijackedConn).To(Equal(fakeHijackedConn))
				Expect(actualResponseReader).To(Equal(expectedResponseReader))
			})
		})

		Context("when httpResponse is not success", func() {
			var fakeBody *transportfakes.FakeReadCloser

			BeforeEach(func() {
				fakeHijackableClient.DoReturns(&httpResp, fakeHijackCloser, nil)
				fakeBody = new(transportfakes.FakeReadCloser)
				httpResp = http.Response{StatusCode: http.StatusTeapot, Body: fakeBody}

				realBody := strings.NewReader("some-error")
				fakeBody.ReadStub = func(buf []byte) (int, error) {
					Expect(fakeBody.CloseCallCount()).To(BeZero())
					return realBody.Read(buf)
				}
			})

			It("closes httpResp.Body, hijackCloser and returns error", func() {
				Expect(fakeHijackCloser).NotTo(BeNil())
				Expect(hijackError).To(HaveOccurred())
				Expect(fakeBody.CloseCallCount()).To(Equal(1))
				Expect(hijackError).To(MatchError(fmt.Errorf("backend error: Exit status: %d, message: %s", httpResp.StatusCode, "some-error")))
				Expect(fakeHijackCloser.CloseCallCount()).To(Equal(1))
			})
		})

		Context("when httpResponse fails with ExecutableNotFoundError", func() {
			var fakeBody *transportfakes.FakeReadCloser
			gerr := garden.ExecutableNotFoundError{Message: "sorry, couldn't find it"}

			BeforeEach(func() {
				fakeHijackableClient.DoReturns(&httpResp, fakeHijackCloser, nil)
				fakeBody = new(transportfakes.FakeReadCloser)
				httpResp = http.Response{StatusCode: http.StatusTeapot, Body: fakeBody}

				jsonBody, _ := garden.Error{Err: gerr}.MarshalJSON()
				realBody := bytes.NewReader(jsonBody)
				fakeBody.ReadStub = func(buf []byte) (int, error) {
					Expect(fakeBody.CloseCallCount()).To(BeZero())
					return realBody.Read(buf)
				}
			})

			It("closes httpResp.Body, hijackCloser and returns ExecutableNotFoundError", func() {
				Expect(fakeHijackCloser).NotTo(BeNil())
				Expect(hijackError).To(HaveOccurred())
				Expect(fakeBody.CloseCallCount()).To(Equal(1))
				Expect(hijackError).To(MatchError(gerr))
				Expect(fakeHijackCloser.CloseCallCount()).To(Equal(1))
			})
		})

		Context("when httpResponse is not success with bad response", func() {
			var fakeBody *transportfakes.FakeReadCloser

			BeforeEach(func() {
				fakeHijackableClient.DoReturns(&httpResp, fakeHijackCloser, nil)
				fakeBody = new(transportfakes.FakeReadCloser)
				httpResp = http.Response{StatusCode: http.StatusTeapot, Body: fakeBody}

				fakeBody.ReadStub = func(buf []byte) (int, error) {
					Expect(fakeBody.CloseCallCount()).To(BeZero())
					return 0, errors.New("error reading")
				}
			})

			It("closes httpResp.Body and returns bad response", func() {
				Expect(fakeHijackCloser).NotTo(BeNil())
				Expect(hijackError).To(HaveOccurred())
				Expect(fakeBody.CloseCallCount()).To(Equal(1))
				Expect(hijackError).To(MatchError(fmt.Errorf("backend error: Exit status: %d, error reading response body: %s", httpResp.StatusCode, "error reading")))
				Expect(fakeHijackCloser.CloseCallCount()).To(Equal(1))
			})
		})

		Context("when httpResponse fails", func() {
			BeforeEach(func() {
				fakeHijackableClient.DoReturns(nil, nil, errors.New("Request failed"))
				httpResp = http.Response{StatusCode: http.StatusTeapot, Body: ioutil.NopCloser(body)}
			})

			It("returns error", func() {
				Expect(hijackError).To(HaveOccurred())
				Expect(actualHijackedConn).To(BeNil())
				Expect(actualResponseReader).To(BeNil())
			})

			It("makes the right request", func() {
				expectedRequest, err := http.NewRequest("POST", "http://example.url", strings.NewReader("some-request-body"))
				expectedRequest.Header.Add("Content-Type", "application/json")
				expectedRequest.URL.RawQuery = "key=some&key=values"
				Expect(err).NotTo(HaveOccurred())
				Expect(fakeHijackableClient.DoCallCount()).To(Equal(1))
				actualRequest := fakeHijackableClient.DoArgsForCall(0)

				Expect(actualRequest.Method).To(Equal(expectedRequest.Method))
				Expect(actualRequest.URL).To(Equal(expectedRequest.URL))
				Expect(actualRequest.Header).To(Equal(expectedRequest.Header))

				s, err := ioutil.ReadAll(actualRequest.Body)
				Expect(err).NotTo(HaveOccurred())
				Expect(string(s)).To(Equal("some-request-body"))
			})
		})
	})
})
<|endoftext|>


<commit_before>package pg

import (
	"context"
	"database/sql"
	"fmt"

	"github.com/hashicorp/terraform/internal/backend"
	"github.com/hashicorp/terraform/internal/legacy/helper/schema"
	"github.com/lib/pq"
)

const (
	statesTableName = "states"
	statesIndexName = "states_by_name"
)

// New creates a new backend for Postgres remote state.
func New() backend.Backend {
	s := &schema.Backend{
		Schema: map[string]*schema.Schema{
			"conn_str": {
				Type:        schema.TypeString,
				Required:    true,
				Description: "Postgres connection string; a `postgres://` URL",
			},

			"schema_name": {
				Type:        schema.TypeString,
				Optional:    true,
				Description: "Name of the automatically managed Postgres schema to store state",
				Default:     "terraform_remote_state",
			},

			"skip_schema_creation": {
				Type:        schema.TypeBool,
				Optional:    true,
				Description: "If set to `true`, Terraform won't try to create the Postgres schema",
				Default:     false,
			},

			"skip_table_creation": {
				Type:        schema.TypeBool,
				Optional:    true,
				Description: "If set to `true`, Terraform won't try to create the Postgres table",
			},

			"skip_index_creation": {
				Type:        schema.TypeBool,
				Optional:    true,
				Description: "If set to `true`, Terraform won't try to create the Postgres index",
			},
		},
	}

	result := &Backend{Backend: s}
	result.Backend.ConfigureFunc = result.configure
	return result
}

type Backend struct {
	*schema.Backend

	// The fields below are set from configure
	db         *sql.DB
	configData *schema.ResourceData
	connStr    string
	schemaName string
}

func (b *Backend) configure(ctx context.Context) error {
	// Grab the resource data
	b.configData = schema.FromContextBackendConfig(ctx)
	data := b.configData

	b.connStr = data.Get("conn_str").(string)
	b.schemaName = pq.QuoteIdentifier(data.Get("schema_name").(string))

	db, err := sql.Open("postgres", b.connStr)
	if err != nil {
		return err
	}

	// Prepare database schema, tables, & indexes.
	var query string

	if !data.Get("skip_schema_creation").(bool) {
		// list all schemas to see if it exists
		var count int
		query = `select count(1) from information_schema.schemata where schema_name = $1`
		if err := db.QueryRow(query, data.Get("schema_name").(string)).Scan(&count); err != nil {
			return err
		}

		// skip schema creation if schema already exists
		// `CREATE SCHEMA IF NOT EXISTS` is to be avoided if ever
		// a user hasn't been granted the `CREATE SCHEMA` privilege
		if count < 1 {
			// tries to create the schema
			query = `CREATE SCHEMA IF NOT EXISTS %s`
			if _, err := db.Exec(fmt.Sprintf(query, b.schemaName)); err != nil {
				return err
			}
		}
	}

	if !data.Get("skip_table_creation").(bool) {
		if _, err := db.Exec("CREATE SEQUENCE IF NOT EXISTS public.global_states_id_seq AS bigint"); err != nil {
			return err
		}

		query = `CREATE TABLE IF NOT EXISTS %s.%s (
			id bigint NOT NULL DEFAULT nextval('public.global_states_id_seq') PRIMARY KEY,
			name text,
			data text
			)`
		if _, err := db.Exec(fmt.Sprintf(query, b.schemaName, statesTableName)); err != nil {
			return err
		}
	}

	if !data.Get("skip_index_creation").(bool) {
		query = `CREATE UNIQUE INDEX IF NOT EXISTS %s ON %s.%s (name)`
		if _, err := db.Exec(fmt.Sprintf(query, statesIndexName, b.schemaName, statesTableName)); err != nil {
			return err
		}
	}

	// Assign db after its schema is prepared.
	b.db = db

	return nil
}
<commit_msg>Add UNIQUE constraint in the states table for the pg backend<commit_after>package pg

import (
	"context"
	"database/sql"
	"fmt"

	"github.com/hashicorp/terraform/internal/backend"
	"github.com/hashicorp/terraform/internal/legacy/helper/schema"
	"github.com/lib/pq"
)

const (
	statesTableName = "states"
	statesIndexName = "states_by_name"
)

// New creates a new backend for Postgres remote state.
func New() backend.Backend {
	s := &schema.Backend{
		Schema: map[string]*schema.Schema{
			"conn_str": {
				Type:        schema.TypeString,
				Required:    true,
				Description: "Postgres connection string; a `postgres://` URL",
			},

			"schema_name": {
				Type:        schema.TypeString,
				Optional:    true,
				Description: "Name of the automatically managed Postgres schema to store state",
				Default:     "terraform_remote_state",
			},

			"skip_schema_creation": {
				Type:        schema.TypeBool,
				Optional:    true,
				Description: "If set to `true`, Terraform won't try to create the Postgres schema",
				Default:     false,
			},

			"skip_table_creation": {
				Type:        schema.TypeBool,
				Optional:    true,
				Description: "If set to `true`, Terraform won't try to create the Postgres table",
			},

			"skip_index_creation": {
				Type:        schema.TypeBool,
				Optional:    true,
				Description: "If set to `true`, Terraform won't try to create the Postgres index",
			},
		},
	}

	result := &Backend{Backend: s}
	result.Backend.ConfigureFunc = result.configure
	return result
}

type Backend struct {
	*schema.Backend

	// The fields below are set from configure
	db         *sql.DB
	configData *schema.ResourceData
	connStr    string
	schemaName string
}

func (b *Backend) configure(ctx context.Context) error {
	// Grab the resource data
	b.configData = schema.FromContextBackendConfig(ctx)
	data := b.configData

	b.connStr = data.Get("conn_str").(string)
	b.schemaName = pq.QuoteIdentifier(data.Get("schema_name").(string))

	db, err := sql.Open("postgres", b.connStr)
	if err != nil {
		return err
	}

	// Prepare database schema, tables, & indexes.
	var query string

	if !data.Get("skip_schema_creation").(bool) {
		// list all schemas to see if it exists
		var count int
		query = `select count(1) from information_schema.schemata where schema_name = $1`
		if err := db.QueryRow(query, data.Get("schema_name").(string)).Scan(&count); err != nil {
			return err
		}

		// skip schema creation if schema already exists
		// `CREATE SCHEMA IF NOT EXISTS` is to be avoided if ever
		// a user hasn't been granted the `CREATE SCHEMA` privilege
		if count < 1 {
			// tries to create the schema
			query = `CREATE SCHEMA IF NOT EXISTS %s`
			if _, err := db.Exec(fmt.Sprintf(query, b.schemaName)); err != nil {
				return err
			}
		}
	}

	if !data.Get("skip_table_creation").(bool) {
		if _, err := db.Exec("CREATE SEQUENCE IF NOT EXISTS public.global_states_id_seq AS bigint"); err != nil {
			return err
		}

		query = `CREATE TABLE IF NOT EXISTS %s.%s (
			id bigint NOT NULL DEFAULT nextval('public.global_states_id_seq') PRIMARY KEY,
			name text UNIQUE,
			data text
			)`
		if _, err := db.Exec(fmt.Sprintf(query, b.schemaName, statesTableName)); err != nil {
			return err
		}
	}

	if !data.Get("skip_index_creation").(bool) {
		query = `CREATE UNIQUE INDEX IF NOT EXISTS %s ON %s.%s (name)`
		if _, err := db.Exec(fmt.Sprintf(query, statesIndexName, b.schemaName, statesTableName)); err != nil {
			return err
		}
	}

	// Assign db after its schema is prepared.
	b.db = db

	return nil
}
<|endoftext|>