<commit_before>package helm

import (
	"log"
	"net/http"
	"net/url"
	"os"
	"strings"
	"time"
)

const (
	get     = "GET"
	head    = "HEAD"
	post    = "POST"
	put     = "PUT"
	patch   = "PATCH"
	deleteh = "DELETE"
)

// Handle is just like "net/http" Handlers, only takes params.
type Handle func(http.ResponseWriter, *http.Request, url.Values)

// Middleware is just like the Handle type, but has a boolean return. True
// means to keep processing the rest of the middleware chain, false means end.
// If you return false to end the request-response cycle you MUST
// write something back to the client, otherwise it will be left hanging.
type Middleware func(http.ResponseWriter, *http.Request, url.Values) bool

// Router name says it all.
type Router struct {
	tree        *node
	rootHandler Handle
	middleware  []Middleware
	l           *log.Logger
}

// New creates a new router. Take the root/fall through route
// like how the default mux works. Only difference is in this case,
// you have to specific one.
func New(rootHandler Handle) *Router {
	node := node{component: "/", isNamedParam: false, methods: make(map[string]*route)}
	return &Router{tree: &node, rootHandler: rootHandler, l: log.New(os.Stdout, "[helm] ", 0)}
}

// Handle takes an http handler, method and pattern for a route.
func (r *Router) Handle(method, path string, handler Handle, middleware ...Middleware) {
	if path[0] != '/' {
		panic("Path has to start with a /.")
	}
	r.tree.addNode(method, path, handler, middleware...)
}

// GET same as Handle only the method is already implied.
func (r *Router) GET(path string, handler Handle, middleware ...Middleware) {
	r.Handle(get, path, handler, middleware...)
}

// HEAD same as Handle only the method is already implied.
func (r *Router) HEAD(path string, handler Handle, middleware ...Middleware) {
	r.Handle(head, path, handler, middleware...)
}

// POST same as Handle only the method is already implied.
func (r *Router) POST(path string, handler Handle, middleware ...Middleware) {
	r.Handle(post, path, handler, middleware...)
}

// PUT same as Handle only the method is already implied.
func (r *Router) PUT(path string, handler Handle, middleware ...Middleware) {
	r.Handle(put, path, handler, middleware...)
}

// PATCH same as Handle only the method is already implied.
func (r *Router) PATCH(path string, handler Handle, middleware ...Middleware) { // might make this and put one.
	r.Handle(patch, path, handler, middleware...)
}

// DELETE same as Handle only the method is already implied.
func (r *Router) DELETE(path string, handler Handle, middleware ...Middleware) {
	r.Handle(deleteh, path, handler, middleware...)
}

// Add Middleware adds middleware to all of the routes.
func (r *Router) AddMiddleware(middleware ...Middleware) {
	r.middleware = append(r.middleware, middleware...)
}

// Run is a simple wrapper around http.ListenAndServe.
func (r *Router) Run(address string) {
	r.l.Println("Running on", address)
	http.ListenAndServe(address, r)
}

// runMiddleware loops over the slice of middleware and call to each of the middleware handlers.
func runMiddleware(w http.ResponseWriter, req *http.Request, params url.Values, middleware ...Middleware) bool {
	for _, m := range middleware {
		if !m(w, req, params) {
			return false // the middleware returned false, so end processing the chain.
		}
	}
	return true
}

// Needed by "net/http" to handle http requests and be a mux to http.ListenAndServe.
func (r *Router) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	start := time.Now()
	r.l.Printf("Started %s %s", req.Method, req.URL.Path)
	cw := responseWriter{w, 200}

	req.ParseForm()
	params := req.Form
	if !runMiddleware(&cw, req, params, r.middleware...) {
		r.l.Printf("Completed %d %s in %v", cw.status, http.StatusText(cw.status), time.Since(start))
		return // end the chain.
	}
	node, _ := r.tree.traverse(strings.Split(req.URL.Path, "/")[1:], params)
	if handler := node.methods[req.Method]; handler != nil {
		if !runMiddleware(&cw, req, params, handler.middleware...) {
			r.l.Printf("Completed %d %s in %v", cw.status, http.StatusText(cw.status), time.Since(start))
			return
		}
		handler.handler(&cw, req, params)
	} else {
		r.rootHandler(&cw, req, params)
	}

	r.l.Printf("Completed %d %s in %v", cw.status, http.StatusText(cw.status), time.Since(start))
}
<commit_msg>comment fix<commit_after>package helm

import (
	"log"
	"net/http"
	"net/url"
	"os"
	"strings"
	"time"
)

const (
	get     = "GET"
	head    = "HEAD"
	post    = "POST"
	put     = "PUT"
	patch   = "PATCH"
	deleteh = "DELETE"
)

// Handle is just like "net/http" Handlers, only takes params.
type Handle func(http.ResponseWriter, *http.Request, url.Values)

// Middleware is just like the Handle type, but has a boolean return. True
// means to keep processing the rest of the middleware chain, false means end.
// If you return false to end the request-response cycle you MUST
// write something back to the client, otherwise it will be left hanging.
type Middleware func(http.ResponseWriter, *http.Request, url.Values) bool

// Router name says it all.
type Router struct {
	tree        *node
	rootHandler Handle
	middleware  []Middleware
	l           *log.Logger
}

// New creates a new router. Take the root/fall through route
// like how the default mux works. Only difference is in this case,
// you have to specific one.
func New(rootHandler Handle) *Router {
	node := node{component: "/", isNamedParam: false, methods: make(map[string]*route)}
	return &Router{tree: &node, rootHandler: rootHandler, l: log.New(os.Stdout, "[helm] ", 0)}
}

// Handle takes an http handler, method and pattern for a route.
func (r *Router) Handle(method, path string, handler Handle, middleware ...Middleware) {
	if path[0] != '/' {
		panic("Path has to start with a /.")
	}
	r.tree.addNode(method, path, handler, middleware...)
}

// GET same as Handle only the method is already implied.
func (r *Router) GET(path string, handler Handle, middleware ...Middleware) {
	r.Handle(get, path, handler, middleware...)
}

// HEAD same as Handle only the method is already implied.
func (r *Router) HEAD(path string, handler Handle, middleware ...Middleware) {
	r.Handle(head, path, handler, middleware...)
}

// POST same as Handle only the method is already implied.
func (r *Router) POST(path string, handler Handle, middleware ...Middleware) {
	r.Handle(post, path, handler, middleware...)
}

// PUT same as Handle only the method is already implied.
func (r *Router) PUT(path string, handler Handle, middleware ...Middleware) {
	r.Handle(put, path, handler, middleware...)
}

// PATCH same as Handle only the method is already implied.
func (r *Router) PATCH(path string, handler Handle, middleware ...Middleware) { // might make this and put one.
	r.Handle(patch, path, handler, middleware...)
}

// DELETE same as Handle only the method is already implied.
func (r *Router) DELETE(path string, handler Handle, middleware ...Middleware) {
	r.Handle(deleteh, path, handler, middleware...)
}

// AddMiddleware adds middleware to all of the routes.
func (r *Router) AddMiddleware(middleware ...Middleware) {
	r.middleware = append(r.middleware, middleware...)
}

// Run is a simple wrapper around http.ListenAndServe.
func (r *Router) Run(address string) {
	r.l.Println("Running on", address)
	http.ListenAndServe(address, r)
}

// runMiddleware loops over the slice of middleware and call to each of the middleware handlers.
func runMiddleware(w http.ResponseWriter, req *http.Request, params url.Values, middleware ...Middleware) bool {
	for _, m := range middleware {
		if !m(w, req, params) {
			return false // the middleware returned false, so end processing the chain.
		}
	}
	return true
}

// Needed by "net/http" to handle http requests and be a mux to http.ListenAndServe.
func (r *Router) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	start := time.Now()
	r.l.Printf("Started %s %s", req.Method, req.URL.Path)
	cw := responseWriter{w, 200}

	req.ParseForm()
	params := req.Form
	if !runMiddleware(&cw, req, params, r.middleware...) {
		r.l.Printf("Completed %d %s in %v", cw.status, http.StatusText(cw.status), time.Since(start))
		return // end the chain.
	}
	node, _ := r.tree.traverse(strings.Split(req.URL.Path, "/")[1:], params)
	if handler := node.methods[req.Method]; handler != nil {
		if !runMiddleware(&cw, req, params, handler.middleware...) {
			r.l.Printf("Completed %d %s in %v", cw.status, http.StatusText(cw.status), time.Since(start))
			return
		}
		handler.handler(&cw, req, params)
	} else {
		r.rootHandler(&cw, req, params)
	}

	r.l.Printf("Completed %d %s in %v", cw.status, http.StatusText(cw.status), time.Since(start))
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"net/http"
	"net/url"
	"sync"
	"time"

	"github.com/prometheus/client_golang/prometheus"

	"github.com/alphagov/router/handlers"
	"github.com/alphagov/router/logger"
	"github.com/alphagov/router/triemux"
	"github.com/globalsign/mgo"
)

// Router is a wrapper around an HTTP multiplexer (trie.Mux) which retrieves its
// routes from a passed mongo database.
type Router struct {
	mux                   *triemux.Mux
	lock                  sync.RWMutex
	mongoURL              string
	mongoDbName           string
	backendConnectTimeout time.Duration
	backendHeaderTimeout  time.Duration
	logger                logger.Logger
}

type Backend struct {
	BackendID  string `bson:"backend_id"`
	BackendURL string `bson:"backend_url"`
}

type Route struct {
	IncomingPath string `bson:"incoming_path"`
	RouteType    string `bson:"route_type"`
	Handler      string `bson:"handler"`
	BackendID    string `bson:"backend_id"`
	RedirectTo   string `bson:"redirect_to"`
	RedirectType string `bson:"redirect_type"`
	SegmentsMode string `bson:"segments_mode"`
	Disabled     bool   `bson:"disabled"`
}

// NewRouter returns a new empty router instance. You will still need to call
// ReloadRoutes() to do the initial route load.
func NewRouter(mongoURL, mongoDbName, backendConnectTimeout, backendHeaderTimeout, logFileName string) (rt *Router, err error) {
	beConnTimeout, err := time.ParseDuration(backendConnectTimeout)
	if err != nil {
		return nil, err
	}
	beHeaderTimeout, err := time.ParseDuration(backendHeaderTimeout)
	if err != nil {
		return nil, err
	}
	logInfo("router: using backend connect timeout:", beConnTimeout)
	logInfo("router: using backend header timeout:", beHeaderTimeout)

	l, err := logger.New(logFileName)
	if err != nil {
		return nil, err
	}
	logInfo("router: logging errors as JSON to", logFileName)

	rt = &Router{
		mux:                   triemux.NewMux(),
		mongoURL:              mongoURL,
		mongoDbName:           mongoDbName,
		backendConnectTimeout: beConnTimeout,
		backendHeaderTimeout:  beHeaderTimeout,
		logger:                l,
	}
	return rt, nil
}

// ServeHTTP delegates responsibility for serving requests to the proxy mux
// instance for this router.
func (rt *Router) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	defer func() {
		if r := recover(); r != nil {
			logWarn("router: recovered from panic in ServeHTTP:", r)

			errorMessage := fmt.Sprintf("panic: %v", r)
			err := logger.RecoveredError{ErrorMessage: errorMessage}

			logger.NotifySentry(logger.ReportableError{Error: err, Request: req})
			rt.logger.LogFromClientRequest(map[string]interface{}{
				"error":  errorMessage,
				"status": http.StatusInternalServerError,
			}, req)

			w.WriteHeader(http.StatusInternalServerError)

			internalServerErrorCountMetric.With(prometheus.Labels{"host": req.Host}).Inc()
		}
	}()
	rt.lock.RLock()
	mux := rt.mux
	rt.lock.RUnlock()

	mux.ServeHTTP(w, req)
}

// ReloadRoutes reloads the routes for this Router instance on the fly. It will
// create a new proxy mux, load applications (backends) and routes into it, and
// then flip the "mux" pointer in the Router.
func (rt *Router) ReloadRoutes() {
	routeReloadCountMetric.Inc()

	defer func() {
		if r := recover(); r != nil {
			logWarn("router: recovered from panic in ReloadRoutes:", r)
			logInfo("router: original routes have not been modified")
			errorMessage := fmt.Sprintf("panic: %v", r)
			err := logger.RecoveredError{ErrorMessage: errorMessage}
			logger.NotifySentry(logger.ReportableError{Error: err})

			routeReloadErrorCountMetric.Inc()
		}
	}()

	logDebug("mgo: connecting to", rt.mongoURL)
	sess, err := mgo.Dial(rt.mongoURL)
	if err != nil {
		panic(fmt.Sprintln("mgo:", err))
	}
	defer sess.Close()
	sess.SetMode(mgo.Strong, true)

	db := sess.DB(rt.mongoDbName)

	logInfo("router: reloading routes")
	newmux := triemux.NewMux()

	backends := rt.loadBackends(db.C("backends"))
	loadRoutes(db.C("routes"), newmux, backends)

	rt.lock.Lock()
	rt.mux = newmux
	rt.lock.Unlock()

	logInfo(fmt.Sprintf("router: reloaded %d routes (checksum: %x)", rt.mux.RouteCount(), rt.mux.RouteChecksum()))

	routesCountMetric.Set(float64(rt.mux.RouteCount()))
}

// loadBackends is a helper function which loads backends from the
// passed mongo collection, constructs a Handler for each one, and returns
// them in map keyed on the backend_id
func (rt *Router) loadBackends(c *mgo.Collection) (backends map[string]http.Handler) {
	backend := &Backend{}
	backends = make(map[string]http.Handler)

	iter := c.Find(nil).Iter()

	for iter.Next(&backend) {
		backendURL, err := url.Parse(backend.BackendURL)
		if err != nil {
			logWarn(fmt.Sprintf("router: couldn't parse URL %s for backend %s "+
				"(error: %v), skipping!", backend.BackendURL, backend.BackendID, err))
			continue
		}

		backends[backend.BackendID] = handlers.NewBackendHandler(
			backend.BackendID,
			backendURL,
			rt.backendConnectTimeout, rt.backendHeaderTimeout,
			rt.logger,
		)
	}

	if err := iter.Err(); err != nil {
		panic(err)
	}

	return
}

// loadRoutes is a helper function which loads routes from the passed mongo
// collection and registers them with the passed proxy mux.
func loadRoutes(c *mgo.Collection, mux *triemux.Mux, backends map[string]http.Handler) {
	route := &Route{}

	iter := c.Find(nil).Sort("incoming_path", "route_type").Iter()

	goneHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		http.Error(w, "410 Gone", http.StatusGone)
	})
	unavailableHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		http.Error(w, "503 Service Unavailable", http.StatusServiceUnavailable)
	})

	for iter.Next(&route) {
		prefix := (route.RouteType == "prefix")

		// the database contains paths with % encoded routes.
		// Unescape them here because the http.Request objects we match against contain the unescaped variants.
		incomingURL, err := url.Parse(route.IncomingPath)
		if err != nil {
			logWarn(fmt.Sprintf("router: found route %+v with invalid incoming path '%s', skipping!", route, route.IncomingPath))
			continue
		}

		if route.Disabled {
			mux.Handle(incomingURL.Path, prefix, unavailableHandler)
			logDebug(fmt.Sprintf("router: registered %s (prefix: %v)(disabled) -> Unavailable", incomingURL.Path, prefix))
			continue
		}

		switch route.Handler {
		case "backend":
			handler, ok := backends[route.BackendID]
			if !ok {
				logWarn(fmt.Sprintf("router: found route %+v which references unknown backend "+
					"%s, skipping!", route, route.BackendID))
				continue
			}
			mux.Handle(incomingURL.Path, prefix, handler)
			logDebug(fmt.Sprintf("router: registered %s (prefix: %v) for %s",
				incomingURL.Path, prefix, route.BackendID))
		case "redirect":
			redirectTemporarily := (route.RedirectType == "temporary")
			handler := handlers.NewRedirectHandler(incomingURL.Path, route.RedirectTo, shouldPreserveSegments(route), redirectTemporarily)
			mux.Handle(incomingURL.Path, prefix, handler)
			logDebug(fmt.Sprintf("router: registered %s (prefix: %v) -> %s",
				incomingURL.Path, prefix, route.RedirectTo))
		case "gone":
			mux.Handle(incomingURL.Path, prefix, goneHandler)
			logDebug(fmt.Sprintf("router: registered %s (prefix: %v) -> Gone", incomingURL.Path, prefix))
		case "boom":
			// Special handler so that we can test failure behaviour.
			mux.Handle(incomingURL.Path, prefix, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				panic("Boom!!!")
			}))
			logDebug(fmt.Sprintf("router: registered %s (prefix: %v) -> Boom!!!", incomingURL.Path, prefix))
		default:
			logWarn(fmt.Sprintf("router: found route %+v with unknown handler type "+
				"%s, skipping!", route, route.Handler))
			continue
		}
	}

	if err := iter.Err(); err != nil {
		panic(err)
	}
}

func (rt *Router) RouteStats() (stats map[string]interface{}) {
	rt.lock.RLock()
	mux := rt.mux
	rt.lock.RUnlock()

	stats = make(map[string]interface{})
	stats["count"] = mux.RouteCount()
	stats["checksum"] = fmt.Sprintf("%x", mux.RouteChecksum())
	return
}

func shouldPreserveSegments(route *Route) bool {
	switch {
	case route.RouteType == "exact" && route.SegmentsMode == "preserve":
		return true
	case route.RouteType == "exact":
		return false
	case route.RouteType == "prefix" && route.SegmentsMode == "ignore":
		return false
	case route.RouteType == "prefix":
		return true
	}
	return false
}
<commit_msg>Move the reload route count metric inside the defer<commit_after>package main

import (
	"fmt"
	"net/http"
	"net/url"
	"sync"
	"time"

	"github.com/prometheus/client_golang/prometheus"

	"github.com/alphagov/router/handlers"
	"github.com/alphagov/router/logger"
	"github.com/alphagov/router/triemux"
	"github.com/globalsign/mgo"
)

// Router is a wrapper around an HTTP multiplexer (trie.Mux) which retrieves its
// routes from a passed mongo database.
type Router struct {
	mux                   *triemux.Mux
	lock                  sync.RWMutex
	mongoURL              string
	mongoDbName           string
	backendConnectTimeout time.Duration
	backendHeaderTimeout  time.Duration
	logger                logger.Logger
}

type Backend struct {
	BackendID  string `bson:"backend_id"`
	BackendURL string `bson:"backend_url"`
}

type Route struct {
	IncomingPath string `bson:"incoming_path"`
	RouteType    string `bson:"route_type"`
	Handler      string `bson:"handler"`
	BackendID    string `bson:"backend_id"`
	RedirectTo   string `bson:"redirect_to"`
	RedirectType string `bson:"redirect_type"`
	SegmentsMode string `bson:"segments_mode"`
	Disabled     bool   `bson:"disabled"`
}

// NewRouter returns a new empty router instance. You will still need to call
// ReloadRoutes() to do the initial route load.
func NewRouter(mongoURL, mongoDbName, backendConnectTimeout, backendHeaderTimeout, logFileName string) (rt *Router, err error) {
	beConnTimeout, err := time.ParseDuration(backendConnectTimeout)
	if err != nil {
		return nil, err
	}
	beHeaderTimeout, err := time.ParseDuration(backendHeaderTimeout)
	if err != nil {
		return nil, err
	}
	logInfo("router: using backend connect timeout:", beConnTimeout)
	logInfo("router: using backend header timeout:", beHeaderTimeout)

	l, err := logger.New(logFileName)
	if err != nil {
		return nil, err
	}
	logInfo("router: logging errors as JSON to", logFileName)

	rt = &Router{
		mux:                   triemux.NewMux(),
		mongoURL:              mongoURL,
		mongoDbName:           mongoDbName,
		backendConnectTimeout: beConnTimeout,
		backendHeaderTimeout:  beHeaderTimeout,
		logger:                l,
	}
	return rt, nil
}

// ServeHTTP delegates responsibility for serving requests to the proxy mux
// instance for this router.
func (rt *Router) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	defer func() {
		if r := recover(); r != nil {
			logWarn("router: recovered from panic in ServeHTTP:", r)

			errorMessage := fmt.Sprintf("panic: %v", r)
			err := logger.RecoveredError{ErrorMessage: errorMessage}

			logger.NotifySentry(logger.ReportableError{Error: err, Request: req})
			rt.logger.LogFromClientRequest(map[string]interface{}{
				"error":  errorMessage,
				"status": http.StatusInternalServerError,
			}, req)

			w.WriteHeader(http.StatusInternalServerError)

			internalServerErrorCountMetric.With(prometheus.Labels{"host": req.Host}).Inc()
		}
	}()
	rt.lock.RLock()
	mux := rt.mux
	rt.lock.RUnlock()

	mux.ServeHTTP(w, req)
}

// ReloadRoutes reloads the routes for this Router instance on the fly. It will
// create a new proxy mux, load applications (backends) and routes into it, and
// then flip the "mux" pointer in the Router.
func (rt *Router) ReloadRoutes() {
	defer func() {
		// increment this metric regardless of whether the route reload succeeded
		routeReloadCountMetric.Inc()

		if r := recover(); r != nil {
			logWarn("router: recovered from panic in ReloadRoutes:", r)
			logInfo("router: original routes have not been modified")
			errorMessage := fmt.Sprintf("panic: %v", r)
			err := logger.RecoveredError{ErrorMessage: errorMessage}
			logger.NotifySentry(logger.ReportableError{Error: err})

			routeReloadErrorCountMetric.Inc()
		}
	}()

	logDebug("mgo: connecting to", rt.mongoURL)
	sess, err := mgo.Dial(rt.mongoURL)
	if err != nil {
		panic(fmt.Sprintln("mgo:", err))
	}
	defer sess.Close()
	sess.SetMode(mgo.Strong, true)

	db := sess.DB(rt.mongoDbName)

	logInfo("router: reloading routes")
	newmux := triemux.NewMux()

	backends := rt.loadBackends(db.C("backends"))
	loadRoutes(db.C("routes"), newmux, backends)

	rt.lock.Lock()
	rt.mux = newmux
	rt.lock.Unlock()

	logInfo(fmt.Sprintf("router: reloaded %d routes (checksum: %x)", rt.mux.RouteCount(), rt.mux.RouteChecksum()))

	routesCountMetric.Set(float64(rt.mux.RouteCount()))
}

// loadBackends is a helper function which loads backends from the
// passed mongo collection, constructs a Handler for each one, and returns
// them in map keyed on the backend_id
func (rt *Router) loadBackends(c *mgo.Collection) (backends map[string]http.Handler) {
	backend := &Backend{}
	backends = make(map[string]http.Handler)

	iter := c.Find(nil).Iter()

	for iter.Next(&backend) {
		backendURL, err := url.Parse(backend.BackendURL)
		if err != nil {
			logWarn(fmt.Sprintf("router: couldn't parse URL %s for backend %s "+
				"(error: %v), skipping!", backend.BackendURL, backend.BackendID, err))
			continue
		}

		backends[backend.BackendID] = handlers.NewBackendHandler(
			backend.BackendID,
			backendURL,
			rt.backendConnectTimeout, rt.backendHeaderTimeout,
			rt.logger,
		)
	}

	if err := iter.Err(); err != nil {
		panic(err)
	}

	return
}

// loadRoutes is a helper function which loads routes from the passed mongo
// collection and registers them with the passed proxy mux.
func loadRoutes(c *mgo.Collection, mux *triemux.Mux, backends map[string]http.Handler) {
	route := &Route{}

	iter := c.Find(nil).Sort("incoming_path", "route_type").Iter()

	goneHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		http.Error(w, "410 Gone", http.StatusGone)
	})
	unavailableHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		http.Error(w, "503 Service Unavailable", http.StatusServiceUnavailable)
	})

	for iter.Next(&route) {
		prefix := (route.RouteType == "prefix")

		// the database contains paths with % encoded routes.
		// Unescape them here because the http.Request objects we match against contain the unescaped variants.
		incomingURL, err := url.Parse(route.IncomingPath)
		if err != nil {
			logWarn(fmt.Sprintf("router: found route %+v with invalid incoming path '%s', skipping!", route, route.IncomingPath))
			continue
		}

		if route.Disabled {
			mux.Handle(incomingURL.Path, prefix, unavailableHandler)
			logDebug(fmt.Sprintf("router: registered %s (prefix: %v)(disabled) -> Unavailable", incomingURL.Path, prefix))
			continue
		}

		switch route.Handler {
		case "backend":
			handler, ok := backends[route.BackendID]
			if !ok {
				logWarn(fmt.Sprintf("router: found route %+v which references unknown backend "+
					"%s, skipping!", route, route.BackendID))
				continue
			}
			mux.Handle(incomingURL.Path, prefix, handler)
			logDebug(fmt.Sprintf("router: registered %s (prefix: %v) for %s",
				incomingURL.Path, prefix, route.BackendID))
		case "redirect":
			redirectTemporarily := (route.RedirectType == "temporary")
			handler := handlers.NewRedirectHandler(incomingURL.Path, route.RedirectTo, shouldPreserveSegments(route), redirectTemporarily)
			mux.Handle(incomingURL.Path, prefix, handler)
			logDebug(fmt.Sprintf("router: registered %s (prefix: %v) -> %s",
				incomingURL.Path, prefix, route.RedirectTo))
		case "gone":
			mux.Handle(incomingURL.Path, prefix, goneHandler)
			logDebug(fmt.Sprintf("router: registered %s (prefix: %v) -> Gone", incomingURL.Path, prefix))
		case "boom":
			// Special handler so that we can test failure behaviour.
			mux.Handle(incomingURL.Path, prefix, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				panic("Boom!!!")
			}))
			logDebug(fmt.Sprintf("router: registered %s (prefix: %v) -> Boom!!!", incomingURL.Path, prefix))
		default:
			logWarn(fmt.Sprintf("router: found route %+v with unknown handler type "+
				"%s, skipping!", route, route.Handler))
			continue
		}
	}

	if err := iter.Err(); err != nil {
		panic(err)
	}
}

func (rt *Router) RouteStats() (stats map[string]interface{}) {
	rt.lock.RLock()
	mux := rt.mux
	rt.lock.RUnlock()

	stats = make(map[string]interface{})
	stats["count"] = mux.RouteCount()
	stats["checksum"] = fmt.Sprintf("%x", mux.RouteChecksum())
	return
}

func shouldPreserveSegments(route *Route) bool {
	switch {
	case route.RouteType == "exact" && route.SegmentsMode == "preserve":
		return true
	case route.RouteType == "exact":
		return false
	case route.RouteType == "prefix" && route.SegmentsMode == "ignore":
		return false
	case route.RouteType == "prefix":
		return true
	}
	return false
}
<|endoftext|>


<commit_before>package helm

import (
	"io"
	"log"
	"net/http"
	"net/url"
	"strings"
	"time"
)

const (
	get     = "GET"
	head    = "HEAD"
	post    = "POST"
	put     = "PUT"
	patch   = "PATCH"
	deleteh = "DELETE"
)

// Handle is just like "net/http" Handlers, only takes params.
type Handle func(http.ResponseWriter, *http.Request, url.Values)

// Middleware is just like the Handle type, but has a boolean return. True
// means to keep processing the rest of the middleware chain, false means end.
// If you return false to end the request-response cycle you MUST
// write something back to the client, otherwise it will be left hanging.
type Middleware func(http.ResponseWriter, *http.Request, url.Values) bool

// Router name says it all.
type Router struct {
	tree           *node
	rootHandler    Handle
	middleware     []Middleware
	l              *log.Logger
	LoggingEnabled bool
}

// New creates a new router. Take the root/fall through route
// like how the default mux works. Only difference is in this case,
// you have to specific one.
func New(rootHandler Handle) *Router {
	node := node{component: "/", isNamedParam: false, methods: make(map[string]*route)}
	return &Router{tree: &node, rootHandler: rootHandler}
}

// EnableLogging sets logging to supplied writer.
func (r *Router) EnableLogging(w io.Writer) {
	r.l = log.New(w, "[helm] ", 0)
	r.LoggingEnabled = true
}

// Handle takes an http handler, method and pattern for a route.
func (r *Router) Handle(method, path string, handler Handle, middleware ...Middleware) {
	if path[0] != '/' {
		panic("Path has to start with a /.")
	}
	r.tree.addNode(method, path, handler, middleware...)
}

// GET same as Handle only the method is already implied.
func (r *Router) GET(path string, handler Handle, middleware ...Middleware) {
	r.Handle(get, path, handler, middleware...)
}

// HEAD same as Handle only the method is already implied.
func (r *Router) HEAD(path string, handler Handle, middleware ...Middleware) {
	r.Handle(head, path, handler, middleware...)
}

// POST same as Handle only the method is already implied.
func (r *Router) POST(path string, handler Handle, middleware ...Middleware) {
	r.Handle(post, path, handler, middleware...)
}

// PUT same as Handle only the method is already implied.
func (r *Router) PUT(path string, handler Handle, middleware ...Middleware) {
	r.Handle(put, path, handler, middleware...)
}

// PATCH same as Handle only the method is already implied.
func (r *Router) PATCH(path string, handler Handle, middleware ...Middleware) { // might make this and put one.
	r.Handle(patch, path, handler, middleware...)
}

// DELETE same as Handle only the method is already implied.
func (r *Router) DELETE(path string, handler Handle, middleware ...Middleware) {
	r.Handle(deleteh, path, handler, middleware...)
}

// Use adds middleware to all of the routes.
func (r *Router) Use(middleware ...Middleware) {
	r.middleware = append(r.middleware, middleware...)
}

// Run is a simple wrapper around http.ListenAndServe.
func (r *Router) Run(address string) {
	r.l.Println("Running on", address)
	http.ListenAndServe(address, r)
}

// runMiddleware loops over the slice of middleware and call to each of the middleware handlers.
func runMiddleware(w http.ResponseWriter, req *http.Request, params url.Values, middleware ...Middleware) bool {
	for _, m := range middleware {
		if !m(w, req, params) {
			return false // the middleware returned false, so end processing the chain.
		}
	}
	return true
}

// Needed by "net/http" to handle http requests and be a mux to http.ListenAndServe.
func (r *Router) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	defer clear(req)
	cw := w
	if r.LoggingEnabled {
		r.l.Printf("Started %s %s", req.Method, req.URL.Path)
		cw = &responseWriter{w, 200}
		start := time.Now()
		defer func(time.Time) {
			status := cw.(*responseWriter).status
			r.l.Printf("Completed %d %s in %v", status, http.StatusText(status), time.Since(start))
		}(start)
	}

	req.ParseMultipartForm(10 * 1024 * 1024) // 10MB. Should probably make this configurable...
	params := req.Form
	if !runMiddleware(cw, req, params, r.middleware...) {
		return // end the chain.
	}
	node, _ := r.tree.traverse(strings.Split(req.URL.Path, "/")[1:], params)
	if handler := node.methods[req.Method]; handler != nil {
		if !runMiddleware(cw, req, params, handler.middleware...) {
			return
		}
		handler.handler(cw, req, params)
	} else {
		r.rootHandler(cw, req, params)
	}
}
<commit_msg>fixed conflicts<commit_after>package helm

import (
	"errors"
	"fmt"
	"io"
	"log"
	"net/http"
	"net/url"
	"strings"
	"time"
)

const (
	get     = "GET"
	head    = "HEAD"
	post    = "POST"
	put     = "PUT"
	patch   = "PATCH"
	deleteh = "DELETE"
)

// Handle is just like "net/http" Handlers, only takes params.
type Handle func(http.ResponseWriter, *http.Request, url.Values)

// Middleware is just like the Handle type, but has a boolean return. True
// means to keep processing the rest of the middleware chain, false means end.
// If you return false to end the request-response cycle you MUST
// write something back to the client, otherwise it will be left hanging.
type Middleware func(http.ResponseWriter, *http.Request, url.Values) bool

// Router name says it all.
type Router struct {
	tree           *node
	rootHandler    Handle
	middleware     []Middleware
	l              *log.Logger
	LoggingEnabled bool
}

type Param struct {
	Name     string
	Required bool
}

// New creates a new router. Take the root/fall through route
// like how the default mux works. Only difference is in this case,
// you have to specific one.
func New(rootHandler Handle) *Router {
	node := node{component: "/", isNamedParam: false, methods: make(map[string]*route)}
	return &Router{tree: &node, rootHandler: rootHandler}
}

// EnableLogging sets logging to supplied writer.
func (r *Router) EnableLogging(w io.Writer) {
	r.l = log.New(w, "[helm] ", 0)
	r.LoggingEnabled = true
}

// Handle takes an http handler, method and pattern for a route.
func (r *Router) Handle(method, path string, handler Handle, middleware ...Middleware) {
	if path[0] != '/' {
		panic("Path has to start with a /.")
	}
	r.tree.addNode(method, path, handler, middleware...)
}

// GET same as Handle only the method is already implied.
func (r *Router) GET(path string, handler Handle, middleware ...Middleware) {
	r.Handle(get, path, handler, middleware...)
}

// HEAD same as Handle only the method is already implied.
func (r *Router) HEAD(path string, handler Handle, middleware ...Middleware) {
	r.Handle(head, path, handler, middleware...)
}

// POST same as Handle only the method is already implied.
func (r *Router) POST(path string, handler Handle, middleware ...Middleware) {
	r.Handle(post, path, handler, middleware...)
}

// PUT same as Handle only the method is already implied.
func (r *Router) PUT(path string, handler Handle, middleware ...Middleware) {
	r.Handle(put, path, handler, middleware...)
}

// PATCH same as Handle only the method is already implied.
func (r *Router) PATCH(path string, handler Handle, middleware ...Middleware) { // might make this and put one.
	r.Handle(patch, path, handler, middleware...)
}

// DELETE same as Handle only the method is already implied.
func (r *Router) DELETE(path string, handler Handle, middleware ...Middleware) {
	r.Handle(deleteh, path, handler, middleware...)
}

// Use adds middleware to all of the routes.
func (r *Router) Use(middleware ...Middleware) {
	r.middleware = append(r.middleware, middleware...)
}

// Run is a simple wrapper around http.ListenAndServe.
func (r *Router) Run(address string) {
	r.l.Println("Running on", address)
	http.ListenAndServe(address, r)
}

// runMiddleware loops over the slice of middleware and call to each of the middleware handlers.
func runMiddleware(w http.ResponseWriter, req *http.Request, params url.Values, middleware ...Middleware) bool {
	for _, m := range middleware {
		if !m(w, req, params) {
			return false // the middleware returned false, so end processing the chain.
		}
	}
	return true
}

// Needed by "net/http" to handle http requests and be a mux to http.ListenAndServe.
func (r *Router) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	defer clear(req)
	cw := w
	if r.LoggingEnabled {
		r.l.Printf("Started %s %s", req.Method, req.URL.Path)
		cw = &responseWriter{w, 200}
		start := time.Now()
		defer func(time.Time) {
			status := cw.(*responseWriter).status
			r.l.Printf("Completed %d %s in %v", status, http.StatusText(status), time.Since(start))
		}(start)
	}

	req.ParseMultipartForm(10 * 1024 * 1024) // 10MB. Should probably make this configurable...
	params := req.Form
	if !runMiddleware(cw, req, params, r.middleware...) {
		return // end the chain.
	}
	node, _ := r.tree.traverse(strings.Split(req.URL.Path, "/")[1:], params)
	if handler := node.methods[req.Method]; handler != nil {
		if !runMiddleware(cw, req, params, handler.middleware...) {
			return
		}
		handler.handler(cw, req, params)
	} else {
		r.rootHandler(cw, req, params)
	}
}

// ValidateParams is used for validating and sanizating params. Since HTTP params can have
// same name for multiple params, if this happens it will just use the first one.
func ValidateParams(params url.Values, desiredParams []Param) (map[string]string, error) {
	paramValues := make(map[string]string)
	for _, param := range desiredParams {
		p, ok := params[param.Name]
		if !ok && param.Required && p[0] != "" {
			return nil, errors.New(fmt.Sprintf("Required parameter (%s) not valid", param.Name))
		} else if ok {
			paramValues[param.Name] = p[0]
		}
	}
	return paramValues, nil
}
<|endoftext|>


<commit_before>package main

import ()

func BeARemote(config MinionConfig) {
}

// vim: foldmethod=marker
<commit_msg>add rpc hookup<commit_after>package main

import (
	"fmt"
	"log"

	"pault.ag/go/service"
)

func BeARemote(config MinionConfig) {
	node := MinionNode{}
	log.Printf("Bringing remote online\n")
	node.Register()
	log.Printf("Diling coordinator\n")
	conn, err := service.DialFromKeys(
		fmt.Sprintf("%s:%d", config.Host, config.Port),
		config.Cert, config.Key, config.CaCert,
	)
	if err != nil {
		log.Fatalf("Error! %s\n", err)
	}
	log.Printf("Bringing RPC online")
	client := service.Client(conn)
	log.Printf("%s\n", client)
}

// vim: foldmethod=marker
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"net"
	"sync"
	"time"

	"github.com/alienth/fastlyctl/util"
	"github.com/alienth/go-fastly"
	"github.com/juju/ratelimit"
)

type rateBucket struct {
	ratelimit.Bucket
	lastUsed time.Time
}

type ipRate struct {
	ip          *net.IP
	buckets     map[Dimension]*rateBucket
	entries     []*util.ACLEntry
	limited     bool
	shouldLimit bool
	list        *IPList

	FirstHit    int64     `json:"first_hit,omitempty"`
	LastHit     epochTime `json:"last_hit,omitempty"`
	LastLimit   int64     `json:"last_limit,omitempty"`
	Hits        int       `json:"hits,omitempty"`
	Strikes     int       `json:"strikes,omitempty"`
	Expire      int64     `json:"-"`
	LimitExpire int64     `json:"limit_expire,omitempty"`

	sync.RWMutex
}

// Records a hit and returns true if it is over limit.
func (ipr *ipRate) Hit(ts time.Time, dimension *Dimension) bool {
	ipr.Lock()
	defer ipr.Unlock()

	rate := float64(ipr.list.Requests) / ipr.list.Time.Duration.Seconds()
	var found bool
	// If DimensionValues were specified in our IPList, check to see if the
	// dimension passed matches that value. If it doesn't, zero out the
	// Dimension so that we just track by IP address.
	if len(ipr.list.DimensionValues) > 0 {
		for _, value := range ipr.list.DimensionValues {
			if value == dimension.Value {
				found = true
				break
			}
		}
		if found != true {
			dimension = &Dimension{}
		}
	}
	var bucket *rateBucket
	if dimension != nil && ipr.list.DimensionShared {
		sharedBuckets := ipr.list.sharedBuckets
		sharedBuckets.Lock()
		if bucket, found = sharedBuckets.m[*dimension]; !found {
			bucket = &rateBucket{Bucket: *ratelimit.NewBucketWithRate(rate, ipr.list.Requests)}
			sharedBuckets.m[*dimension] = bucket
		}
		sharedBuckets.Unlock()
		ipr.buckets[*dimension] = bucket
	}
	if bucket, found = ipr.buckets[*dimension]; !found {
		bucket = &rateBucket{Bucket: *ratelimit.NewBucketWithRate(rate, ipr.list.Requests)}
		ipr.buckets[*dimension] = bucket
	}
	var overlimit bool
	waitTime := bucket.Take(1)
	bucket.lastUsed = ts
	if waitTime != 0 {
		overlimit = true
	}
	if ipr.FirstHit == 0 {
		ipr.FirstHit = time.Now().Unix()
	}
	ipr.LastHit.Time = ts
	ipr.Hits++
	ipr.Expire = time.Now().Add(ipr.list.Expire.Duration).Unix()
	return overlimit
}

// Limit adds an IP to a fastly edge ACL
func (ipr *ipRate) Limit(service *fastly.Service) error {
	ipr.Lock()
	defer ipr.Unlock()

	if !ipr.shouldLimit {
		return nil
	}

	// Return if this IP is already limited on this service.
	for _, e := range ipr.entries {
		if e.ServiceID == service.ID {
			return nil
		}
	}

	ipr.LastLimit = time.Now().Unix()
	if !ipr.limited {
		// Only increase the duration time if we're not already
		// limited.  This is because we might just be applying a limit
		// to a new service that we saw a hit on.
		ipr.Strikes++
	}
	limitDuration := ipr.list.LimitDuration.multiply(float64(ipr.Strikes))
	ipr.LimitExpire = time.Now().Add(limitDuration.Duration).Unix()
	ipr.Expire = time.Now().Add(time.Duration(24) * time.Hour).Unix()
	comment, err := json.Marshal(ipr)
	if err != nil {
		return err
	}
	entry, err := util.NewACLEntry(client, service.Name, aclName, ipr.ip.String(), 0, string(comment), false)
	if err != nil {
		return err
	}
	fmt.Printf("Limiting IP %s for %d minutes on service %s\n", ipr.ip.String(), int(limitDuration.Minutes()), service.Name)
	if !noop {
		if err = entry.Add(); err != nil {
			return err
		}
	}

	ipr.limited = true
	ipr.entries = append(ipr.entries, entry)
	return nil
}

// Removes an IP from ratelimits
func (ipr *ipRate) RemoveLimit() error {
	ipr.Lock()
	defer ipr.Unlock()
	if len(ipr.entries) > 0 {
		fmt.Printf("Unlimiting IP %s\n", ipr.ip.String())
		// defer the filtration in case we get an error during the removal loop
		defer func(ipr *ipRate) {
			newEntries := ipr.entries[:0]
			for _, e := range ipr.entries {
				if e != nil {
					newEntries = append(newEntries, e)
				}
			}
			ipr.entries = newEntries
		}(ipr)
		for i, entry := range ipr.entries {
			if !noop {
				if err := entry.Remove(); err != nil {
					return fmt.Errorf("Error removing limit for IP %s: %s", ipr.ip.String(), err)
				}
			}
			ipr.entries[i] = nil
		}
		ipr.limited = false
	}
	return nil
}

// clean will free any shared bucket from our IPList if this ipRate was the last
// to utilize that shared bucket.
func (ipr *ipRate) cleanSharedBuckets() {
	ipr.Lock()
	defer ipr.Unlock()
	sharedBuckets := ipr.list.sharedBuckets
	sharedBuckets.Lock()
	defer sharedBuckets.Unlock()
	for dimension, bucket := range ipr.buckets {
		if bucket.lastUsed == ipr.LastHit.Time {
			delete(sharedBuckets.m, dimension)
		}
	}
}
<commit_msg>Add comment.<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"net"
	"sync"
	"time"

	"github.com/alienth/fastlyctl/util"
	"github.com/alienth/go-fastly"
	"github.com/juju/ratelimit"
)

type rateBucket struct {
	ratelimit.Bucket
	lastUsed time.Time
}

type ipRate struct {
	ip          *net.IP
	buckets     map[Dimension]*rateBucket
	entries     []*util.ACLEntry
	limited     bool
	shouldLimit bool
	list        *IPList

	FirstHit    int64     `json:"first_hit,omitempty"`
	LastHit     epochTime `json:"last_hit,omitempty"`
	LastLimit   int64     `json:"last_limit,omitempty"`
	Hits        int       `json:"hits,omitempty"`
	Strikes     int       `json:"strikes,omitempty"`
	Expire      int64     `json:"-"`
	LimitExpire int64     `json:"limit_expire,omitempty"`

	sync.RWMutex
}

// Records a hit and returns true if it is over limit.
func (ipr *ipRate) Hit(ts time.Time, dimension *Dimension) bool {
	ipr.Lock()
	defer ipr.Unlock()

	rate := float64(ipr.list.Requests) / ipr.list.Time.Duration.Seconds()
	var found bool
	// If DimensionValues were specified in our IPList, check to see if the
	// dimension passed matches that value. If it doesn't, zero out the
	// Dimension so that we just track by IP address.
	if len(ipr.list.DimensionValues) > 0 {
		for _, value := range ipr.list.DimensionValues {
			if value == dimension.Value {
				found = true
				break
			}
		}
		if found != true {
			dimension = &Dimension{}
		}
	}
	var bucket *rateBucket
	if dimension != nil && ipr.list.DimensionShared {
		sharedBuckets := ipr.list.sharedBuckets
		sharedBuckets.Lock()
		if bucket, found = sharedBuckets.m[*dimension]; !found {
			bucket = &rateBucket{Bucket: *ratelimit.NewBucketWithRate(rate, ipr.list.Requests)}
			sharedBuckets.m[*dimension] = bucket
		}
		sharedBuckets.Unlock()
		ipr.buckets[*dimension] = bucket
	}
	if bucket, found = ipr.buckets[*dimension]; !found {
		bucket = &rateBucket{Bucket: *ratelimit.NewBucketWithRate(rate, ipr.list.Requests)}
		ipr.buckets[*dimension] = bucket
	}
	var overlimit bool
	waitTime := bucket.Take(1)
	bucket.lastUsed = ts
	if waitTime != 0 {
		overlimit = true
	}
	if ipr.FirstHit == 0 {
		ipr.FirstHit = time.Now().Unix()
	}
	ipr.LastHit.Time = ts
	ipr.Hits++
	ipr.Expire = time.Now().Add(ipr.list.Expire.Duration).Unix()
	return overlimit
}

// Limit adds an IP to a fastly edge ACL
func (ipr *ipRate) Limit(service *fastly.Service) error {
	ipr.Lock()
	defer ipr.Unlock()

	if !ipr.shouldLimit {
		return nil
	}

	// Return if this IP is already limited on this service.
	for _, e := range ipr.entries {
		if e.ServiceID == service.ID {
			return nil
		}
	}

	ipr.LastLimit = time.Now().Unix()
	if !ipr.limited {
		// Only increase the duration time if we're not already
		// limited.  This is because we might just be applying a limit
		// to a new service that we saw a hit on.
		ipr.Strikes++
	}
	limitDuration := ipr.list.LimitDuration.multiply(float64(ipr.Strikes))
	ipr.LimitExpire = time.Now().Add(limitDuration.Duration).Unix()
	ipr.Expire = time.Now().Add(time.Duration(24) * time.Hour).Unix()
	comment, err := json.Marshal(ipr)
	if err != nil {
		return err
	}
	entry, err := util.NewACLEntry(client, service.Name, aclName, ipr.ip.String(), 0, string(comment), false)
	if err != nil {
		return err
	}
	fmt.Printf("Limiting IP %s for %d minutes on service %s\n", ipr.ip.String(), int(limitDuration.Minutes()), service.Name)
	if !noop {
		if err = entry.Add(); err != nil {
			return err
		}
	}

	ipr.limited = true
	ipr.entries = append(ipr.entries, entry)
	return nil
}

// Removes an IP from ratelimits
func (ipr *ipRate) RemoveLimit() error {
	ipr.Lock()
	defer ipr.Unlock()
	if len(ipr.entries) > 0 {
		fmt.Printf("Unlimiting IP %s\n", ipr.ip.String())
		// defer the filtration in case we get an error during the removal loop
		defer func(ipr *ipRate) {
			newEntries := ipr.entries[:0]
			for _, e := range ipr.entries {
				if e != nil {
					newEntries = append(newEntries, e)
				}
			}
			// Unless we had a removal error, this should be an empty list.
			ipr.entries = newEntries
		}(ipr)
		for i, entry := range ipr.entries {
			if !noop {
				if err := entry.Remove(); err != nil {
					return fmt.Errorf("Error removing limit for IP %s: %s", ipr.ip.String(), err)
				}
			}
			ipr.entries[i] = nil
		}
		ipr.limited = false
	}
	return nil
}

// clean will free any shared bucket from our IPList if this ipRate was the last
// to utilize that shared bucket.
func (ipr *ipRate) cleanSharedBuckets() {
	ipr.Lock()
	defer ipr.Unlock()
	sharedBuckets := ipr.list.sharedBuckets
	sharedBuckets.Lock()
	defer sharedBuckets.Unlock()
	for dimension, bucket := range ipr.buckets {
		if bucket.lastUsed == ipr.LastHit.Time {
			delete(sharedBuckets.m, dimension)
		}
	}
}
<|endoftext|>


<commit_before>package redis

import "time"

// NewCmdResult returns a Cmd initalised with val and err for testing
func NewCmdResult(val interface{}, err error) *Cmd {
	var cmd Cmd
	cmd.val = val
	cmd.setErr(err)
	return &cmd
}

// NewSliceResult returns a SliceCmd initalised with val and err for testing
func NewSliceResult(val []interface{}, err error) *SliceCmd {
	var cmd SliceCmd
	cmd.val = val
	cmd.setErr(err)
	return &cmd
}

// NewStatusResult returns a StatusCmd initalised with val and err for testing
func NewStatusResult(val string, err error) *StatusCmd {
	var cmd StatusCmd
	cmd.val = val
	cmd.setErr(err)
	return &cmd
}

// NewIntResult returns an IntCmd initalised with val and err for testing
func NewIntResult(val int64, err error) *IntCmd {
	var cmd IntCmd
	cmd.val = val
	cmd.setErr(err)
	return &cmd
}

// NewDurationResult returns a DurationCmd initalised with val and err for testing
func NewDurationResult(val time.Duration, err error) *DurationCmd {
	var cmd DurationCmd
	cmd.val = val
	cmd.setErr(err)
	return &cmd
}

// NewBoolResult returns a BoolCmd initalised with val and err for testing
func NewBoolResult(val bool, err error) *BoolCmd {
	var cmd BoolCmd
	cmd.val = val
	cmd.setErr(err)
	return &cmd
}

// NewStringResult returns a StringCmd initalised with val and err for testing
func NewStringResult(val string, err error) *StringCmd {
	var cmd StringCmd
	cmd.val = []byte(val)
	cmd.setErr(err)
	return &cmd
}

// NewFloatResult returns a FloatCmd initalised with val and err for testing
func NewFloatResult(val float64, err error) *FloatCmd {
	var cmd FloatCmd
	cmd.val = val
	cmd.setErr(err)
	return &cmd
}

// NewStringSliceResult returns a StringSliceCmd initalised with val and err for testing
func NewStringSliceResult(val []string, err error) *StringSliceCmd {
	var cmd StringSliceCmd
	cmd.val = val
	cmd.setErr(err)
	return &cmd
}

// NewBoolSliceResult returns a BoolSliceCmd initalised with val and err for testing
func NewBoolSliceResult(val []bool, err error) *BoolSliceCmd {
	var cmd BoolSliceCmd
	cmd.val = val
	cmd.setErr(err)
	return &cmd
}

// NewStringStringMapResult returns a StringStringMapCmd initalised with val and err for testing
func NewStringStringMapResult(val map[string]string, err error) *StringStringMapCmd {
	var cmd StringStringMapCmd
	cmd.val = val
	cmd.setErr(err)
	return &cmd
}

// NewStringIntMapCmdResult returns a StringIntMapCmd initalised with val and err for testing
func NewStringIntMapCmdResult(val map[string]int64, err error) *StringIntMapCmd {
	var cmd StringIntMapCmd
	cmd.val = val
	cmd.setErr(err)
	return &cmd
}

// NewZSliceCmdResult returns a ZSliceCmd initalised with val and err for testing
func NewZSliceCmdResult(val []Z, err error) *ZSliceCmd {
	var cmd ZSliceCmd
	cmd.val = val
	cmd.setErr(err)
	return &cmd
}

// NewScanCmdResult returns a ScanCmd initalised with val and err for testing
func NewScanCmdResult(keys []string, cursor uint64, err error) *ScanCmd {
	var cmd ScanCmd
	cmd.page = keys
	cmd.cursor = cursor
	cmd.setErr(err)
	return &cmd
}

// NewClusterSlotsCmdResult returns a ClusterSlotsCmd initalised with val and err for testing
func NewClusterSlotsCmdResult(val []ClusterSlot, err error) *ClusterSlotsCmd {
	var cmd ClusterSlotsCmd
	cmd.val = val
	cmd.setErr(err)
	return &cmd
}

// NewGeoLocationCmdResult returns a GeoLocationCmd initalised with val and err for testing
func NewGeoLocationCmdResult(val []GeoLocation, err error) *GeoLocationCmd {
	var cmd GeoLocationCmd
	cmd.locations = val
	cmd.setErr(err)
	return &cmd
}

// NewCommandsInfoCmdResult returns a CommandsInfoCmd initalised with val and err for testing
func NewCommandsInfoCmdResult(val map[string]*CommandInfo, err error) *CommandsInfoCmd {
	var cmd CommandsInfoCmd
	cmd.val = val
	cmd.setErr(err)
	return &cmd
}
<commit_msg>Update result.go (#754)<commit_after>package redis

import "time"

// NewCmdResult returns a Cmd initialised with val and err for testing
func NewCmdResult(val interface{}, err error) *Cmd {
	var cmd Cmd
	cmd.val = val
	cmd.setErr(err)
	return &cmd
}

// NewSliceResult returns a SliceCmd initialised with val and err for testing
func NewSliceResult(val []interface{}, err error) *SliceCmd {
	var cmd SliceCmd
	cmd.val = val
	cmd.setErr(err)
	return &cmd
}

// NewStatusResult returns a StatusCmd initialised with val and err for testing
func NewStatusResult(val string, err error) *StatusCmd {
	var cmd StatusCmd
	cmd.val = val
	cmd.setErr(err)
	return &cmd
}

// NewIntResult returns an IntCmd initialised with val and err for testing
func NewIntResult(val int64, err error) *IntCmd {
	var cmd IntCmd
	cmd.val = val
	cmd.setErr(err)
	return &cmd
}

// NewDurationResult returns a DurationCmd initialised with val and err for testing
func NewDurationResult(val time.Duration, err error) *DurationCmd {
	var cmd DurationCmd
	cmd.val = val
	cmd.setErr(err)
	return &cmd
}

// NewBoolResult returns a BoolCmd initialised with val and err for testing
func NewBoolResult(val bool, err error) *BoolCmd {
	var cmd BoolCmd
	cmd.val = val
	cmd.setErr(err)
	return &cmd
}

// NewStringResult returns a StringCmd initialised with val and err for testing
func NewStringResult(val string, err error) *StringCmd {
	var cmd StringCmd
	cmd.val = []byte(val)
	cmd.setErr(err)
	return &cmd
}

// NewFloatResult returns a FloatCmd initialised with val and err for testing
func NewFloatResult(val float64, err error) *FloatCmd {
	var cmd FloatCmd
	cmd.val = val
	cmd.setErr(err)
	return &cmd
}

// NewStringSliceResult returns a StringSliceCmd initialised with val and err for testing
func NewStringSliceResult(val []string, err error) *StringSliceCmd {
	var cmd StringSliceCmd
	cmd.val = val
	cmd.setErr(err)
	return &cmd
}

// NewBoolSliceResult returns a BoolSliceCmd initialised with val and err for testing
func NewBoolSliceResult(val []bool, err error) *BoolSliceCmd {
	var cmd BoolSliceCmd
	cmd.val = val
	cmd.setErr(err)
	return &cmd
}

// NewStringStringMapResult returns a StringStringMapCmd initialised with val and err for testing
func NewStringStringMapResult(val map[string]string, err error) *StringStringMapCmd {
	var cmd StringStringMapCmd
	cmd.val = val
	cmd.setErr(err)
	return &cmd
}

// NewStringIntMapCmdResult returns a StringIntMapCmd initialised with val and err for testing
func NewStringIntMapCmdResult(val map[string]int64, err error) *StringIntMapCmd {
	var cmd StringIntMapCmd
	cmd.val = val
	cmd.setErr(err)
	return &cmd
}

// NewZSliceCmdResult returns a ZSliceCmd initialised with val and err for testing
func NewZSliceCmdResult(val []Z, err error) *ZSliceCmd {
	var cmd ZSliceCmd
	cmd.val = val
	cmd.setErr(err)
	return &cmd
}

// NewScanCmdResult returns a ScanCmd initialised with val and err for testing
func NewScanCmdResult(keys []string, cursor uint64, err error) *ScanCmd {
	var cmd ScanCmd
	cmd.page = keys
	cmd.cursor = cursor
	cmd.setErr(err)
	return &cmd
}

// NewClusterSlotsCmdResult returns a ClusterSlotsCmd initialised with val and err for testing
func NewClusterSlotsCmdResult(val []ClusterSlot, err error) *ClusterSlotsCmd {
	var cmd ClusterSlotsCmd
	cmd.val = val
	cmd.setErr(err)
	return &cmd
}

// NewGeoLocationCmdResult returns a GeoLocationCmd initialised with val and err for testing
func NewGeoLocationCmdResult(val []GeoLocation, err error) *GeoLocationCmd {
	var cmd GeoLocationCmd
	cmd.locations = val
	cmd.setErr(err)
	return &cmd
}

// NewCommandsInfoCmdResult returns a CommandsInfoCmd initialised with val and err for testing
func NewCommandsInfoCmdResult(val map[string]*CommandInfo, err error) *CommandsInfoCmd {
	var cmd CommandsInfoCmd
	cmd.val = val
	cmd.setErr(err)
	return &cmd
}
<|endoftext|>


<commit_before>package espsdk

import (
	"io/ioutil"
	"net/http"
	"time"

	log "github.com/Sirupsen/logrus"
)

// A Result contains information relative to a completed request, including
// the time elapsed to fulfill the request and any errors.
type result struct {
	Response *response     `json:"response"`
	Payload  []byte        `json:"-"`
	Duration time.Duration `json:"response_ms"`
	Err      error         `json:"-"`
}

func getResult(c *http.Client, req *http.Request) *result {
	httpCommand := req.Method + " " + string(req.URL.Path)
	start := start(httpCommand)
	resp, err := c.Do(req)
	duration := elapsed(httpCommand, start) / time.Millisecond
	if err != nil {
		log.Fatal(err)
		return buildResult(resp, nil, duration)
	}
	defer resp.Body.Close()

	payload, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Fatal(err)
		return buildResult(resp, payload, duration)
	}
	if resp.StatusCode < 200 || resp.StatusCode > 299 {
		log.Warnf("HTTP %s", resp.Status)
	}
	return buildResult(resp, payload, duration)
}

func buildResult(resp *http.Response, payload []byte, duration time.Duration) *result {
	return &result{
		&response{
			resp.StatusCode,
			resp.Status,
		},
		payload, duration, nil}
}
<commit_msg>isolate fields for easier parsing<commit_after>package espsdk

import (
	"io/ioutil"
	"net/http"
	"time"

	log "github.com/Sirupsen/logrus"
)

// A Result contains information relative to a completed request, including
// the time elapsed to fulfill the request and any errors.
type result struct {
	Response *response     `json:"response"`
	Payload  []byte        `json:"-"`
	Duration time.Duration `json:"response_ms"`
	Err      error         `json:"-"`
}

func getResult(c *http.Client, req *http.Request) *result {
	httpCommand := req.Method + " " + string(req.URL.Path)
	start := start(httpCommand)
	resp, err := c.Do(req)
	duration := elapsed(httpCommand, start) / time.Millisecond
	if err != nil {
		log.Fatal(err)
		return buildResult(resp, nil, duration)
	}
	defer resp.Body.Close()

	payload, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Fatal(err)
		return buildResult(resp, payload, duration)
	}
	if resp.StatusCode < 200 || resp.StatusCode > 299 {
		log.WithFields(log.Fields{
			"object":      "response",
			"status_code": resp.StatusCode,
			"status":      resp.Status,
		}).Warn()
	}
	return buildResult(resp, payload, duration)
}

func buildResult(resp *http.Response, payload []byte, duration time.Duration) *result {
	return &result{
		&response{
			resp.StatusCode,
			resp.Status,
		},
		payload, duration, nil}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"os"
	"strings"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/cfpb/rhobot/config"
	"github.com/cfpb/rhobot/database"
	"github.com/cfpb/rhobot/gocd"
	"github.com/cfpb/rhobot/healthcheck"
	"github.com/cfpb/rhobot/report"
	"github.com/codegangsta/cli"
)

func main() {

	app := cli.NewApp()
	app.Name = "Rhobot"
	app.Usage = "Rhobot is a database development tool that uses DevOps best practices."
	app.EnableBashCompletion = true

	config := config.NewConfig()

	logLevelFlag := cli.StringFlag{
		Name:  "loglevel, lvl",
		Value: "",
		Usage: "sets the log level for Rhobot",
	}
	gocdHostFlag := cli.StringFlag{
		Name:  "host",
		Value: "",
		Usage: "host of the GoCD server",
	}
	reportFileFlag := cli.StringFlag{
		Name:  "report",
		Value: "",
		Usage: "path to the healthcheck report",
	}
	dburiFlag := cli.StringFlag{
		Name:  "dburi",
		Value: "",
		Usage: "database uri postgres://user:password@host:port/database",
	}
	emailListFlag := cli.StringFlag{
		Name:  "email",
		Value: "",
		Usage: "yaml file containing email distribution list",
	}

	app.Flags = []cli.Flag{logLevelFlag}
	app.Commands = []cli.Command{
		{
			Name:    "run",
			Aliases: []string{},
			Usage:   "healthchecks|pipeline|tbd",
			Subcommands: []cli.Command{
				{
					Name:  "healthchecks",
					Usage: "HEALTHCHECK_FILE [--dburi DATABASE_URI] [--report REPORT_FILE] [--email DISTRIBUTION_FILE]",
					Flags: []cli.Flag{
						reportFileFlag,
						dburiFlag,
						emailListFlag,
					},
					Action: func(c *cli.Context) {
						if c.String("loglevel") != "" {
							config.SetLogLevel(c.String("loglevel"))
						}

						// variables to be populated by cli args
						var healthcheckPath string
						var reportPath string
						var emailListPath string

						if c.Args().Get(0) != "" {
							healthcheckPath = c.Args().Get(0)
						} else {
							log.Error("You must provide the path to the healthcheck file.")
							return
						}
						log.Info("Running health checks from ", healthcheckPath)

						if c.String("dburi") != "" {
							config.SetDBURI(c.String("dburi"))
						}
						log.Debug("DB_URI: ", config.DBURI())

						if c.String("report") != "" {
							reportPath = c.String("report")
							log.Debugf("Generating report at %v", reportPath)
						}

						if c.String("email") != "" {
							emailListPath = c.String("email")
							log.Debugf("Emailing report to %v", emailListPath)
						}

						healthcheckRunner(config, healthcheckPath, reportPath, emailListPath)
						log.Info("Success!")
					},
				},
				{
					Name:    "pipeline",
					Aliases: []string{},
					Usage:   "Interact with GoCD pipeline",
					Subcommands: []cli.Command{
						{
							Name:  "push",
							Usage: "PATH [PIPELINE_GROUP]",
							Flags: []cli.Flag{
								gocdHostFlag,
							},
							Action: func(c *cli.Context) {
								if c.String("loglevel") != "" {
									config.SetLogLevel(c.String("loglevel"))
								}

								if c.String("host") != "" {
									log.Debug("Setting GoCD host: ", c.String("host"))
									config.SetGoCDHost(c.String("host"))
								}

								if len(c.Args()) > 0 {
									path := c.Args()[0]
									group := c.Args().Get(1)
									log.Infof("Pushing config from %v to pipeline group %v...", path, group)
									if err := gocd.Push(config.GoCDURL(), path, group); err != nil {
										log.Fatal("Failed to push pipeline config: ", err)
									}
									log.Info("Success!")
								} else {
									log.Fatal("A path to the pipeline config to push is required.")
								}
							},
						},
						{
							Name:  "pull",
							Usage: "PATH",
							Flags: []cli.Flag{
								gocdHostFlag,
							},
							Action: func(c *cli.Context) {
								if c.String("loglevel") != "" {
									config.SetLogLevel(c.String("loglevel"))
								}

								if c.String("host") != "" {
									config.SetGoCDHost(c.String("host"))
								}

								if len(c.Args()) > 0 {
									path := c.Args()[0]
									log.Infof("Pulling config from %v to %v...", config.GoCDURL(), path)
									if err := gocd.Pull(config.GoCDURL(), path); err != nil {
										log.Fatal("Failed to pull pipeline config: ", err)
									}
									log.Info("Success!")
								} else {
									log.Fatal("A path to pull the pipeline config to is required.")
								}
							},
						},
						{
							Name:  "clone",
							Usage: "PIPELINE_NAME PATH",
							Flags: []cli.Flag{
								gocdHostFlag,
							},
							Action: func(c *cli.Context) {
								if c.String("loglevel") != "" {
									config.SetLogLevel(c.String("loglevel"))
								}

								if c.String("host") != "" {
									config.SetGoCDHost(c.String("host"))
								}

								if len(c.Args()) > 1 {
									name := c.Args()[0]
									path := c.Args()[1]
									log.Infof("Cloning pipeline %v to %v...", name, path)
									if err := gocd.Clone(config.GoCDURL(), path, name); err != nil {
										log.Fatal("Failed to clone pipeline config: ", err)
									}
									log.Info("Success!")
								} else {
									log.Fatal("A pipeline name and a path to clone to are required.")
								}
							},
						},
					},
				},
			},
		},
	}

	app.Run(os.Args)
}

func healthcheckRunner(config *config.Config, healthcheckPath string, reportPath string, emailListPath string) {
	healthChecks, err := healthcheck.ReadHealthCheckYAMLFromFile(healthcheckPath)
	if err != nil {
		log.Fatal("Failed to read healthchecks: ", err)
	}
	cxn := database.GetPGConnection(config.DBURI())

	// TODO the error returned from PreformHealthChecks determis a bad exit
	results, _ := healthcheck.PreformHealthChecks(healthChecks, cxn)
	var elements []report.Element
	for _, val := range results {
		elements = append(elements, val)
	}

	// Make Templated report
	metadata := map[string]interface{}{
		"name":      healthChecks.Name,
		"db_name":   config.PgDatabase,
		"footer":    healthcheck.FooterHealthcheck,
		"timestamp": time.Now().UTC().String(),
	}

	prr := report.NewPongo2ReportRunnerFromString(healthcheck.TemplateHealthcheck)
	rs := report.Set{Elements: elements, Metadata: metadata}
	reader, _ := prr.ReportReader(rs)

	// Write report to file
	if reportPath != "" {
		fhr := report.FileHandler{Filename: reportPath}
		_ = fhr.HandleReport(reader)
	}

	// Email report
	if emailListPath != "" {

		df, err := report.ReadDistributionFormatYAMLFromFile(emailListPath)
		if err != nil {
			log.Fatal("Failed to read distribution format: ", err)
		}

		for _, level := range report.LogLevelArray {

			// TODO calculate a subject line to include hostname, DB, # of failed HCs
			hcName := metadata["name"]
			if hcName == "" {
				hcName = "healthchecks"
			}
			subjectStr := fmt.Sprintf("%s for %s at %s level",
				hcName, metadata["db_name"], strings.ToUpper(level))

			logFilteredSet := report.FilterReportSet(rs, level)
			reader, _ := prr.ReportReader(logFilteredSet)
			recipients := df.GetEmails(level)

			if recipients != nil && len(recipients) != 0 && len(logFilteredSet.Elements) != 0 {
				log.Infof("Send %s to: %v", subjectStr, recipients)
				ehr := report.EmailHandler{
					SMTPHost:    config.SMTPHost,
					SMTPPort:    config.SMTPPort,
					SenderEmail: config.SMTPEmail,
					SenderName:  config.SMTPName,
					Subject:     subjectStr,
					Recipients:  recipients,
					HTML:        true,
				}
				err = ehr.HandleReport(reader)
				if err != nil {
					log.Warn("Failed to email report: ", err)
				}
			}
		}

	}
}
<commit_msg>remove run command<commit_after>package main

import (
	"fmt"
	"os"
	"strings"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/cfpb/rhobot/config"
	"github.com/cfpb/rhobot/database"
	"github.com/cfpb/rhobot/gocd"
	"github.com/cfpb/rhobot/healthcheck"
	"github.com/cfpb/rhobot/report"
	"github.com/codegangsta/cli"
)

func main() {

	app := cli.NewApp()
	app.Name = "Rhobot"
	app.Usage = "Rhobot is a database development tool that uses DevOps best practices."
	app.EnableBashCompletion = true

	config := config.NewConfig()

	logLevelFlag := cli.StringFlag{
		Name:  "loglevel, lvl",
		Value: "",
		Usage: "sets the log level for Rhobot",
	}
	gocdHostFlag := cli.StringFlag{
		Name:  "host",
		Value: "",
		Usage: "host of the GoCD server",
	}
	reportFileFlag := cli.StringFlag{
		Name:  "report",
		Value: "",
		Usage: "path to the healthcheck report",
	}
	dburiFlag := cli.StringFlag{
		Name:  "dburi",
		Value: "",
		Usage: "database uri postgres://user:password@host:port/database",
	}
	emailListFlag := cli.StringFlag{
		Name:  "email",
		Value: "",
		Usage: "yaml file containing email distribution list",
	}

	app.Flags = []cli.Flag{logLevelFlag}
	app.Commands = []cli.Command{
		{
			Name:  "healthchecks",
			Usage: "HEALTHCHECK_FILE [--dburi DATABASE_URI] [--report REPORT_FILE] [--email DISTRIBUTION_FILE]",
			Flags: []cli.Flag{
				reportFileFlag,
				dburiFlag,
				emailListFlag,
			},
			Action: func(c *cli.Context) {
				if c.String("loglevel") != "" {
					config.SetLogLevel(c.String("loglevel"))
				}

				// variables to be populated by cli args
				var healthcheckPath string
				var reportPath string
				var emailListPath string

				if c.Args().Get(0) != "" {
					healthcheckPath = c.Args().Get(0)
				} else {
					log.Error("You must provide the path to the healthcheck file.")
					return
				}
				log.Info("Running health checks from ", healthcheckPath)

				if c.String("dburi") != "" {
					config.SetDBURI(c.String("dburi"))
				}
				log.Debug("DB_URI: ", config.DBURI())

				if c.String("report") != "" {
					reportPath = c.String("report")
					log.Debugf("Generating report at %v", reportPath)
				}

				if c.String("email") != "" {
					emailListPath = c.String("email")
					log.Debugf("Emailing report to %v", emailListPath)
				}

				healthcheckRunner(config, healthcheckPath, reportPath, emailListPath)
				log.Info("Success!")
			},
		},
		{
			Name:    "pipeline",
			Aliases: []string{},
			Usage:   "Interact with GoCD pipeline",
			Subcommands: []cli.Command{
				{
					Name:  "push",
					Usage: "PATH [PIPELINE_GROUP]",
					Flags: []cli.Flag{
						gocdHostFlag,
					},
					Action: func(c *cli.Context) {
						if c.String("loglevel") != "" {
							config.SetLogLevel(c.String("loglevel"))
						}

						if c.String("host") != "" {
							log.Debug("Setting GoCD host: ", c.String("host"))
							config.SetGoCDHost(c.String("host"))
						}

						if len(c.Args()) > 0 {
							path := c.Args()[0]
							group := c.Args().Get(1)
							log.Infof("Pushing config from %v to pipeline group %v...", path, group)
							if err := gocd.Push(config.GoCDURL(), path, group); err != nil {
								log.Fatal("Failed to push pipeline config: ", err)
							}
							log.Info("Success!")
						} else {
							log.Fatal("A path to the pipeline config to push is required.")
						}
					},
				},
				{
					Name:  "pull",
					Usage: "PATH",
					Flags: []cli.Flag{
						gocdHostFlag,
					},
					Action: func(c *cli.Context) {
						if c.String("loglevel") != "" {
							config.SetLogLevel(c.String("loglevel"))
						}

						if c.String("host") != "" {
							config.SetGoCDHost(c.String("host"))
						}

						if len(c.Args()) > 0 {
							path := c.Args()[0]
							log.Infof("Pulling config from %v to %v...", config.GoCDURL(), path)
							if err := gocd.Pull(config.GoCDURL(), path); err != nil {
								log.Fatal("Failed to pull pipeline config: ", err)
							}
							log.Info("Success!")
						} else {
							log.Fatal("A path to pull the pipeline config to is required.")
						}
					},
				},
				{
					Name:  "clone",
					Usage: "PIPELINE_NAME PATH",
					Flags: []cli.Flag{
						gocdHostFlag,
					},
					Action: func(c *cli.Context) {
						if c.String("loglevel") != "" {
							config.SetLogLevel(c.String("loglevel"))
						}

						if c.String("host") != "" {
							config.SetGoCDHost(c.String("host"))
						}

						if len(c.Args()) > 1 {
							name := c.Args()[0]
							path := c.Args()[1]
							log.Infof("Cloning pipeline %v to %v...", name, path)
							if err := gocd.Clone(config.GoCDURL(), path, name); err != nil {
								log.Fatal("Failed to clone pipeline config: ", err)
							}
							log.Info("Success!")
						} else {
							log.Fatal("A pipeline name and a path to clone to are required.")
						}
					},
				},
			},
		},
	}

	app.Run(os.Args)
}

func healthcheckRunner(config *config.Config, healthcheckPath string, reportPath string, emailListPath string) {
	healthChecks, err := healthcheck.ReadHealthCheckYAMLFromFile(healthcheckPath)
	if err != nil {
		log.Fatal("Failed to read healthchecks: ", err)
	}
	cxn := database.GetPGConnection(config.DBURI())

	// TODO the error returned from PreformHealthChecks determis a bad exit
	results, _ := healthcheck.PreformHealthChecks(healthChecks, cxn)
	var elements []report.Element
	for _, val := range results {
		elements = append(elements, val)
	}

	// Make Templated report
	metadata := map[string]interface{}{
		"name":      healthChecks.Name,
		"db_name":   config.PgDatabase,
		"footer":    healthcheck.FooterHealthcheck,
		"timestamp": time.Now().UTC().String(),
	}

	prr := report.NewPongo2ReportRunnerFromString(healthcheck.TemplateHealthcheck)
	rs := report.Set{Elements: elements, Metadata: metadata}
	reader, _ := prr.ReportReader(rs)

	// Write report to file
	if reportPath != "" {
		fhr := report.FileHandler{Filename: reportPath}
		_ = fhr.HandleReport(reader)
	}

	// Email report
	if emailListPath != "" {

		df, err := report.ReadDistributionFormatYAMLFromFile(emailListPath)
		if err != nil {
			log.Fatal("Failed to read distribution format: ", err)
		}

		for _, level := range report.LogLevelArray {

			// TODO calculate a subject line to include hostname, DB, # of failed HCs
			hcName := metadata["name"]
			if hcName == "" {
				hcName = "healthchecks"
			}
			subjectStr := fmt.Sprintf("%s for %s at %s level",
				hcName, metadata["db_name"], strings.ToUpper(level))

			logFilteredSet := report.FilterReportSet(rs, level)
			reader, _ := prr.ReportReader(logFilteredSet)
			recipients := df.GetEmails(level)

			if recipients != nil && len(recipients) != 0 && len(logFilteredSet.Elements) != 0 {
				log.Infof("Send %s to: %v", subjectStr, recipients)
				ehr := report.EmailHandler{
					SMTPHost:    config.SMTPHost,
					SMTPPort:    config.SMTPPort,
					SenderEmail: config.SMTPEmail,
					SenderName:  config.SMTPName,
					Subject:     subjectStr,
					Recipients:  recipients,
					HTML:        true,
				}
				err = ehr.HandleReport(reader)
				if err != nil {
					log.Warn("Failed to email report: ", err)
				}
			}
		}

	}
}
<|endoftext|>


<commit_before>package peco

import (
	"encoding/json"
	"fmt"
	"os"

	"github.com/nsf/termbox-go"
)

type KeymapHandler func(*Input, termbox.Event)
type Keymap map[termbox.Key]KeymapHandler
type KeymapStringKey string
type KeymapStringHandler string

// This map is populated using some magic numbers, which must match
// the values defined in termbox-go. Verification against the actual
// termbox constants are done in the test
var stringToKey = map[string]termbox.Key{}

func init() {
	fidx := 12
	for k := termbox.KeyF1; k > termbox.KeyF12; k-- {
		sk := fmt.Sprintf("F%d", fidx)
		stringToKey[sk] = k
		fidx--
	}

	names := []string{
		"Insert",
		"Delete",
		"Home",
		"End",
		"Pgup",
		"Pgdn",
		"ArrowUp",
		"ArrowDown",
		"ArrowLeft",
		"ArrowRight",
	}
	for i, n := range names {
		stringToKey[n] = termbox.Key(int(termbox.KeyF12) - (i + 1))
	}

	names = []string{
		"Left",
		"Middle",
		"Right",
	}
	for i, n := range names {
		sk := fmt.Sprintf("Mouse%s", n)
		stringToKey[sk] = termbox.Key(int(termbox.KeyArrowRight) - (i + 2))
	}

	whacky := [][]string{
		{"~", "2", "Space"},
		{"a"},
		{"b"},
		{"c"},
		{"d"},
		{"e"},
		{"f"},
		{"g"},
		{"h"},
		{"i"},
		{"j"},
		{"k"},
		{"l"},
		{"m"},
		{"n"},
		{"o"},
		{"p"},
		{"q"},
		{"r"},
		{"s"},
		{"t"},
		{"u"},
		{"v"},
		{"w"},
		{"x"},
		{"y"},
		{"z"},
		{"[", "3"},
		{"4", "\\"},
		{"5", "]"},
		{"6"},
		{"7", "/", "_"},
	}
	for i, list := range whacky {
		for _, n := range list {
			sk := fmt.Sprintf("C-%s", n)
			stringToKey[sk] = termbox.Key(int(termbox.KeyCtrlTilde) + i)
		}
	}

	stringToKey["BS"] = termbox.KeyBackspace
	stringToKey["Tab"] = termbox.KeyTab
	stringToKey["Enter"] = termbox.KeyEnter
	stringToKey["Esc"] = termbox.KeyEsc
	stringToKey["Space"] = termbox.KeySpace
	stringToKey["BS2"] = termbox.KeyBackspace2
	stringToKey["C-8"] = termbox.KeyCtrl8

	//	panic(fmt.Sprintf("%#q", stringToKey))
}

func handleAcceptChar(i *Input, ev termbox.Event) {
	if ev.Key == termbox.KeySpace {
		ev.Ch = ' '
	}

	if ev.Ch > 0 {
		if len(i.query) == i.caretPos {
			i.query = append(i.query, ev.Ch)
		} else {
			buf := make([]rune, len(i.query)+1)
			copy(buf, i.query[:i.caretPos])
			buf[i.caretPos] = ev.Ch
			copy(buf[i.caretPos+1:], i.query[i.caretPos:])
			i.query = buf
		}
		i.caretPos++
		i.ExecQuery(string(i.query))
	}
}

// peco.Finish -> end program, exit with success
func handleFinish(i *Input, _ termbox.Event) {
	if len(i.current) == 1 {
		i.result = i.current[0].line
	} else if i.selectedLine > 0 && i.selectedLine < len(i.current) {
		i.result = i.current[i.selectedLine-1].line
	}
	i.Finish()
}

// peco.Cancel -> end program, exit with failure
func handleCancel(i *Input, ev termbox.Event) {
	i.ExitStatus = 1
	i.Finish()
}

func handleSelectPrevious(i *Input, ev termbox.Event) {
	i.PagingCh() <- ToPrevLine
	i.DrawMatches(nil)
}

func handleSelectNext(i *Input, ev termbox.Event) {
	i.PagingCh() <- ToNextLine
	i.DrawMatches(nil)
}

func handleSelectPreviousPage(i *Input, ev termbox.Event) {
	i.PagingCh() <- ToPrevPage
	i.DrawMatches(nil)
}

func handleSelectNextPage(i *Input, ev termbox.Event) {
	i.PagingCh() <- ToNextPage
	i.DrawMatches(nil)
}

func handleForwardChar(i *Input, _ termbox.Event) {
	if i.caretPos >= len(i.query) {
		return
	}
	i.caretPos++
	i.DrawMatches(nil)
}

func handleBackwardChar(i *Input, _ termbox.Event) {
	if i.caretPos <= 0 {
		return
	}
	i.caretPos--
	i.DrawMatches(nil)
}

func handleDeleteBackwardChar(i *Input, ev termbox.Event) {
	if len(i.query) <= 0 {
		return
	}

	if i.caretPos == len(i.query) {
		i.query = i.query[:len(i.query)-1]
	} else {
		buf := make([]rune, len(i.query)-1)
		copy(buf, i.query[:i.caretPos])
		copy(buf[i.caretPos-1:], i.query[i.caretPos:])
		i.query = buf
	}
	i.caretPos--
	if len(i.query) > 0 {
		i.ExecQuery(string(i.query))
		return
	}

	i.current = nil
	i.DrawMatches(nil)
}

func (ksk KeymapStringKey) ToKey() (k termbox.Key, err error) {
	k, ok := stringToKey[string(ksk)]
	if !ok {
		err = fmt.Errorf("No such key %s", ksk)
	}
	return
}

func (ksh KeymapStringHandler) ToHandler() (h KeymapHandler, err error) {
	switch ksh {
	case "peco.ForwardChar":
		h = handleForwardChar
	case "peco.BackwardChar":
		h = handleBackwardChar
	case "peco.DeleteBackwardChar":
		h = handleDeleteBackwardChar
	case "peco.SelectPreviousPage":
		h = handleSelectPreviousPage
	case "peco.SelectNextPage":
		h = handleSelectNextPage
	case "peco.SelectPrevious":
		h = handleSelectPrevious
	case "peco.SelectNext":
		h = handleSelectNext
	case "peco.Finish":
		h = handleFinish
	case "peco.Cancel":
		h = handleCancel
	default:
		err = fmt.Errorf("No such handler %s", ksh)
	}
	return
}

func NewKeymap() Keymap {
	return Keymap{
		termbox.KeyEsc:        handleCancel,
		termbox.KeyEnter:      handleFinish,
		termbox.KeyArrowUp:    handleSelectPrevious,
		termbox.KeyCtrlK:      handleSelectPrevious,
		termbox.KeyArrowDown:  handleSelectNext,
		termbox.KeyCtrlJ:      handleSelectNext,
		termbox.KeyArrowLeft:  handleSelectPreviousPage,
		termbox.KeyArrowRight: handleSelectNextPage,
		termbox.KeyBackspace:  handleDeleteBackwardChar,
		termbox.KeyBackspace2: handleDeleteBackwardChar,
	}
}

func (km Keymap) Handler(k termbox.Key) KeymapHandler {
	h, ok := km[k]
	if ok {
		return h
	}
	return handleAcceptChar
}

func (km Keymap) UnmarshalJSON(buf []byte) error {
	raw := map[string]string{}
	if err := json.Unmarshal(buf, &raw); err != nil {
		return err
	}

	for ks, vs := range raw {
		k, err := KeymapStringKey(ks).ToKey()
		if err != nil {
			fmt.Fprintf(os.Stderr, "Unknown key %s", ks)
			continue
		}

		v, err := KeymapStringHandler(vs).ToHandler()
		if err != nil {
			fmt.Fprintf(os.Stderr, "Unknown handler %s", vs)
			continue
		}

		km[k] = v
	}

	return nil
}
<commit_msg>Handle BeginningOfLine + EndOfLine<commit_after>package peco

import (
	"encoding/json"
	"fmt"
	"os"

	"github.com/nsf/termbox-go"
)

type KeymapHandler func(*Input, termbox.Event)
type Keymap map[termbox.Key]KeymapHandler
type KeymapStringKey string
type KeymapStringHandler string

// This map is populated using some magic numbers, which must match
// the values defined in termbox-go. Verification against the actual
// termbox constants are done in the test
var stringToKey = map[string]termbox.Key{}

func init() {
	fidx := 12
	for k := termbox.KeyF1; k > termbox.KeyF12; k-- {
		sk := fmt.Sprintf("F%d", fidx)
		stringToKey[sk] = k
		fidx--
	}

	names := []string{
		"Insert",
		"Delete",
		"Home",
		"End",
		"Pgup",
		"Pgdn",
		"ArrowUp",
		"ArrowDown",
		"ArrowLeft",
		"ArrowRight",
	}
	for i, n := range names {
		stringToKey[n] = termbox.Key(int(termbox.KeyF12) - (i + 1))
	}

	names = []string{
		"Left",
		"Middle",
		"Right",
	}
	for i, n := range names {
		sk := fmt.Sprintf("Mouse%s", n)
		stringToKey[sk] = termbox.Key(int(termbox.KeyArrowRight) - (i + 2))
	}

	whacky := [][]string{
		{"~", "2", "Space"},
		{"a"},
		{"b"},
		{"c"},
		{"d"},
		{"e"},
		{"f"},
		{"g"},
		{"h"},
		{"i"},
		{"j"},
		{"k"},
		{"l"},
		{"m"},
		{"n"},
		{"o"},
		{"p"},
		{"q"},
		{"r"},
		{"s"},
		{"t"},
		{"u"},
		{"v"},
		{"w"},
		{"x"},
		{"y"},
		{"z"},
		{"[", "3"},
		{"4", "\\"},
		{"5", "]"},
		{"6"},
		{"7", "/", "_"},
	}
	for i, list := range whacky {
		for _, n := range list {
			sk := fmt.Sprintf("C-%s", n)
			stringToKey[sk] = termbox.Key(int(termbox.KeyCtrlTilde) + i)
		}
	}

	stringToKey["BS"] = termbox.KeyBackspace
	stringToKey["Tab"] = termbox.KeyTab
	stringToKey["Enter"] = termbox.KeyEnter
	stringToKey["Esc"] = termbox.KeyEsc
	stringToKey["Space"] = termbox.KeySpace
	stringToKey["BS2"] = termbox.KeyBackspace2
	stringToKey["C-8"] = termbox.KeyCtrl8

	//	panic(fmt.Sprintf("%#q", stringToKey))
}

func handleAcceptChar(i *Input, ev termbox.Event) {
	if ev.Key == termbox.KeySpace {
		ev.Ch = ' '
	}

	if ev.Ch > 0 {
		if len(i.query) == i.caretPos {
			i.query = append(i.query, ev.Ch)
		} else {
			buf := make([]rune, len(i.query)+1)
			copy(buf, i.query[:i.caretPos])
			buf[i.caretPos] = ev.Ch
			copy(buf[i.caretPos+1:], i.query[i.caretPos:])
			i.query = buf
		}
		i.caretPos++
		i.ExecQuery(string(i.query))
	}
}

// peco.Finish -> end program, exit with success
func handleFinish(i *Input, _ termbox.Event) {
	if len(i.current) == 1 {
		i.result = i.current[0].line
	} else if i.selectedLine > 0 && i.selectedLine < len(i.current) {
		i.result = i.current[i.selectedLine-1].line
	}
	i.Finish()
}

// peco.Cancel -> end program, exit with failure
func handleCancel(i *Input, ev termbox.Event) {
	i.ExitStatus = 1
	i.Finish()
}

func handleSelectPrevious(i *Input, ev termbox.Event) {
	i.PagingCh() <- ToPrevLine
	i.DrawMatches(nil)
}

func handleSelectNext(i *Input, ev termbox.Event) {
	i.PagingCh() <- ToNextLine
	i.DrawMatches(nil)
}

func handleSelectPreviousPage(i *Input, ev termbox.Event) {
	i.PagingCh() <- ToPrevPage
	i.DrawMatches(nil)
}

func handleSelectNextPage(i *Input, ev termbox.Event) {
	i.PagingCh() <- ToNextPage
	i.DrawMatches(nil)
}

func handleForwardChar(i *Input, _ termbox.Event) {
	if i.caretPos >= len(i.query) {
		return
	}
	i.caretPos++
	i.DrawMatches(nil)
}

func handleBackwardChar(i *Input, _ termbox.Event) {
	if i.caretPos <= 0 {
		return
	}
	i.caretPos--
	i.DrawMatches(nil)
}

func handleBeginningOfLine(i *Input, _ termbox.Event) {
	i.caretPos = 0
	i.DrawMatches(nil)
}

func handleEndOfLine(i *Input, _ termbox.Event) {
	i.caretPos = len(i.query)
	i.DrawMatches(nil)
}

func handleDeleteBackwardChar(i *Input, ev termbox.Event) {
	if len(i.query) <= 0 {
		return
	}

	switch i.caretPos {
	case 0:
		// No op
		return
	case len(i.query):
		i.query = i.query[:len(i.query)-1]
	default:
		buf := make([]rune, len(i.query)-1)
		copy(buf, i.query[:i.caretPos])
		copy(buf[i.caretPos-1:], i.query[i.caretPos:])
		i.query = buf
	}
	i.caretPos--
	if len(i.query) > 0 {
		i.ExecQuery(string(i.query))
		return
	}

	i.current = nil
	i.DrawMatches(nil)
}

func (ksk KeymapStringKey) ToKey() (k termbox.Key, err error) {
	k, ok := stringToKey[string(ksk)]
	if !ok {
		err = fmt.Errorf("No such key %s", ksk)
	}
	return
}

func (ksh KeymapStringHandler) ToHandler() (h KeymapHandler, err error) {
	switch ksh {
	case "peco.BeginningOfLine":
		h = handleBeginningOfLine
	case "peco.EndOfLine":
		h = handleEndOfLine
	case "peco.ForwardChar":
		h = handleForwardChar
	case "peco.BackwardChar":
		h = handleBackwardChar
	case "peco.DeleteBackwardChar":
		h = handleDeleteBackwardChar
	case "peco.SelectPreviousPage":
		h = handleSelectPreviousPage
	case "peco.SelectNextPage":
		h = handleSelectNextPage
	case "peco.SelectPrevious":
		h = handleSelectPrevious
	case "peco.SelectNext":
		h = handleSelectNext
	case "peco.Finish":
		h = handleFinish
	case "peco.Cancel":
		h = handleCancel
	default:
		err = fmt.Errorf("No such handler %s", ksh)
	}
	return
}

func NewKeymap() Keymap {
	return Keymap{
		termbox.KeyEsc:        handleCancel,
		termbox.KeyEnter:      handleFinish,
		termbox.KeyArrowUp:    handleSelectPrevious,
		termbox.KeyCtrlK:      handleSelectPrevious,
		termbox.KeyArrowDown:  handleSelectNext,
		termbox.KeyCtrlJ:      handleSelectNext,
		termbox.KeyArrowLeft:  handleSelectPreviousPage,
		termbox.KeyArrowRight: handleSelectNextPage,
		termbox.KeyBackspace:  handleDeleteBackwardChar,
		termbox.KeyBackspace2: handleDeleteBackwardChar,
	}
}

func (km Keymap) Handler(k termbox.Key) KeymapHandler {
	h, ok := km[k]
	if ok {
		return h
	}
	return handleAcceptChar
}

func (km Keymap) UnmarshalJSON(buf []byte) error {
	raw := map[string]string{}
	if err := json.Unmarshal(buf, &raw); err != nil {
		return err
	}

	for ks, vs := range raw {
		k, err := KeymapStringKey(ks).ToKey()
		if err != nil {
			fmt.Fprintf(os.Stderr, "Unknown key %s", ks)
			continue
		}

		v, err := KeymapStringHandler(vs).ToHandler()
		if err != nil {
			fmt.Fprintf(os.Stderr, "Unknown handler %s", vs)
			continue
		}

		km[k] = v
	}

	return nil
}
<|endoftext|>


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

package main

import (
	"fmt"
	"io"
	"os"
	"strings"

	"github.com/spf13/cobra"
	"github.com/spf13/pflag"
	"github.com/chzyer/readline"
	"github.com/vivekn/autocomplete"

	"k8s.io/kubernetes/pkg/kubectl/cmd"
	cmdutil "k8s.io/kubernetes/pkg/kubectl/cmd/util"
)

func main() {
	cmdutil.BehaviorOnFatal(func(msg string, code int) {
		fmt.Println(msg)
	})

	cmd := cmd.NewKubectlCommand(cmdutil.NewFactory(nil), os.Stdin, os.Stdout, os.Stderr)
	l, err := readline.NewEx(&readline.Config{
		Prompt:       ">>> ",
		AutoComplete: &CommandCompleter{cmd},
	})
	if err != nil {
		panic(err)
	}
	defer l.Close()

	for {
		line, err := l.Readline()
		if err == readline.ErrInterrupt {
			if len(line) == 0 {
				break
			} else {
				continue
			}
		} else if err == io.EOF {
			break
		}
		line = strings.TrimSpace(line)
		cmd.ResetFlags()
		cmd.SetArgs(strings.Split(line, " "))
		cmd.Execute()
	}
}

type CommandCompleter struct {
	Root *cobra.Command
}

func (cc *CommandCompleter) Do(line []rune, pos int) (newLine [][]rune, offset int) {
	cmd := cc.Root
	index := strings.LastIndex(string(line[:pos]), " ") + 1
	word := string(line[:pos])
	if index > 0 {
		word = word[index:pos]
		var err error
		cmd, _, err = cc.Root.Find(strings.Split(string(line), " "))
		if err != nil {
			return
		}
	}
	for _, completion := range completions(word, cmd) {
		if len(word) >= len(completion) {
			if len(word) == len(completion) {
				newLine = append(newLine, []rune{' '})
			} else {
				newLine = append(newLine, []rune(completion))
			}
			offset = len(completion)
		} else {
			newLine = append(newLine, []rune(completion)[len(word):])
			offset = len(word)
		}

	}
	return
}

func completions(prefix string, cmd *cobra.Command) (completions []string) {
	if strings.HasPrefix(prefix, "-") {
		completions = flags(cmd)
	} else {
		completions = subCommands(cmd)
		if len(completions) == 0 {
			completions = resourceTypes(cmd)
		}
	}
	trie := trie.NewTrie()
	for _, c := range completions {
		trie.Insert(c)
	}
	completions, _ = trie.AutoComplete(prefix)
	return
}

func subCommands(cmd *cobra.Command) []string {
	prefixes := make([]string, len(cmd.Commands()))
	for i, c := range cmd.Commands() {
		prefixes[i] = c.Name()
	}
	return prefixes
}

func resourceTypes(cmd *cobra.Command) []string {
	return cmd.ValidArgs
}

func flags(cmd *cobra.Command) []string {
	flags := []string{}
	fn := func(f *pflag.Flag) {
		flag := "--" + f.Name
		if len(f.NoOptDefVal) == 0 {
			flag += "="
		}
		flags = append(flags, flag)
	}
	cmd.NonInheritedFlags().VisitAll(fn)
	cmd.InheritedFlags().VisitAll(fn)
	return flags
}
<commit_msg>Add an `sc` command for setting the context<commit_after>// Copyright 2016 Red Hat, Inc, and individual contributors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"bufio"
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"strings"

	"github.com/chzyer/readline"
	"github.com/spf13/cobra"
	"github.com/spf13/pflag"
	"github.com/vivekn/autocomplete"

	"k8s.io/kubernetes/pkg/kubectl/cmd"
	cmdutil "k8s.io/kubernetes/pkg/kubectl/cmd/util"
)

type InternalCommand func(*kubesh, []string) error

type kubesh struct {
	factory          *cmdutil.Factory
	cmd              *cobra.Command
	context          []string
	rl               *readline.Instance
	internalCommands map[string]InternalCommand
}

func main() {
	cmdutil.BehaviorOnFatal(func(msg string, code int) {
		fmt.Println(msg)
	})

	factory := cmdutil.NewFactory(nil)
	cmd := cmd.NewKubectlCommand(factory, os.Stdin, os.Stdout, os.Stderr)
	rl, err := readline.NewEx(&readline.Config{
		Prompt:       "> ",
		AutoComplete: &CommandCompleter{cmd},
	})
	if err != nil {
		panic(err)
	}

	defer rl.Close()

	sh := kubesh{
		factory: factory,
		cmd:     cmd,
		rl:      rl,
		internalCommands: map[string]InternalCommand{
			"exit": func(_ *kubesh, _ []string) error {
				fmt.Println("Bye!")
				os.Exit(0)

				return nil
			},

			"sc": setContextCommand,
		},
	}

	for {
		line, err := sh.rl.Readline()
		if err == readline.ErrInterrupt {
			if len(line) == 0 {
				break
			} else {
				continue
			}
		} else if err == io.EOF {
			break
		}
		line = strings.TrimSpace(line)
		args := strings.Split(line, " ")
		internal, err := sh.runInternalCommand(args)
		if err != nil {
			panic(err)
		}

		if !internal {
			sh.cmd.ResetFlags()
			sh.cmd.SetArgs(args)
			sh.cmd.Execute()
		}
	}
}

type CommandCompleter struct {
	Root *cobra.Command
}

func (cc *CommandCompleter) Do(line []rune, pos int) (newLine [][]rune, offset int) {
	cmd := cc.Root
	index := strings.LastIndex(string(line[:pos]), " ") + 1
	word := string(line[:pos])
	if index > 0 {
		word = word[index:pos]
		var err error
		cmd, _, err = cc.Root.Find(strings.Split(string(line), " "))
		if err != nil {
			return
		}
	}
	for _, completion := range completions(word, cmd) {
		if len(word) >= len(completion) {
			if len(word) == len(completion) {
				newLine = append(newLine, []rune{' '})
			} else {
				newLine = append(newLine, []rune(completion))
			}
			offset = len(completion)
		} else {
			newLine = append(newLine, []rune(completion)[len(word):])
			offset = len(word)
		}

	}
	return
}

func completions(prefix string, cmd *cobra.Command) (completions []string) {
	if strings.HasPrefix(prefix, "-") {
		completions = flags(cmd)
	} else {
		completions = subCommands(cmd)
		if len(completions) == 0 {
			completions = resourceTypes(cmd)
		}
	}
	trie := trie.NewTrie()
	for _, c := range completions {
		trie.Insert(c)
	}
	completions, _ = trie.AutoComplete(prefix)
	return
}

func subCommands(cmd *cobra.Command) []string {
	prefixes := make([]string, len(cmd.Commands()))
	for i, c := range cmd.Commands() {
		prefixes[i] = c.Name()
	}
	return prefixes
}

func resourceTypes(cmd *cobra.Command) []string {
	return cmd.ValidArgs
}

func flags(cmd *cobra.Command) []string {
	flags := []string{}
	fn := func(f *pflag.Flag) {
		flag := "--" + f.Name
		if len(f.NoOptDefVal) == 0 {
			flag += "="
		}
		flags = append(flags, flag)
	}
	cmd.NonInheritedFlags().VisitAll(fn)
	cmd.InheritedFlags().VisitAll(fn)
	return flags
}

func (sh *kubesh) runInternalCommand(args []string) (bool, error) {
	if len(args) > 0 {
		if f := sh.internalCommands[args[0]]; f != nil {

			return true, f(sh, args)
		}
	}

	return false, nil
}

func setContextCommand(sh *kubesh, args []string) error {
	if len(args) == 1 || len(args) > 3 {
		fmt.Println("Usage: " + args[0] + " TYPE [NAME]")

		//TODO: return an error?
		return nil
	}

	typeOnly := len(args) == 2

	var buff bytes.Buffer
	writer := bufio.NewWriter(&buff)
	cmd := cmd.NewKubectlCommand(sh.factory, os.Stdin, writer, os.Stderr)
	callArgs := []string{"get", "--output=json"}
	callArgs = append(callArgs, args[1:]...)

	cmd.SetArgs(callArgs)
	cmd.Execute()

	writer.Flush()
	content, err := ioutil.ReadAll(bufio.NewReader(&buff))
	if err != nil {
		fmt.Println(err)

		return nil
	}

	var result map[string]interface{}
	json.Unmarshal(content, &result)

	if typeOnly {
		// this is fucking disgusting
		// reads {"items": [{"kind": x}]}
		typeName := result["items"].([]interface{})[0].(map[string]interface{})["kind"].(string)
		sh.context = []string{strings.ToLower(typeName)}
	} else {
		// equally fucking disgusting
		// reads {"kind": x, "metadata": {"name": y}}
		typeName := result["kind"].(string)
		resourceName := result["metadata"].(map[string]interface{})["name"].(string)
		sh.context = []string{strings.ToLower(typeName), resourceName}
	}

	sh.rl.SetPrompt(prompt(sh.context))

	return nil
}

func prompt(context []string) string {
	return strings.Join(context, ":") + "> "
}
<|endoftext|>


<commit_before>/*
Copyright 2015 The Kubernetes Authors All rights reserved.

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 to keep track of API Versions that should be registered in api.Scheme.
package registered

import (
	"os"
	"strings"

	"github.com/golang/glog"
)

// List of registered API versions.
// The list is in the order of most preferred to the least.
var RegisteredVersions []string

func init() {
	validAPIVersions := map[string]bool{
		"v1":      true,
		"v1beta3": true,
	}

	// The default list of supported api versions, in order of most preferred to the least.
	defaultSupportedVersions := "v1"
	// Env var KUBE_API_VERSIONS is a comma separated list of API versions that should be registered in the scheme.
	// The versions should be in the order of most preferred to the least.
	supportedVersions := os.Getenv("KUBE_API_VERSIONS")
	if supportedVersions == "" {
		supportedVersions = defaultSupportedVersions
	}
	versions := strings.Split(supportedVersions, ",")
	for _, version := range versions {
		// Verify that the version is valid.
		valid, ok := validAPIVersions[version]
		if !ok || !valid {
			// Not a valid API version.
			glog.Fatalf("invalid api version: %s in KUBE_API_VERSIONS: %s. List of valid API versions: %v",
				version, os.Getenv("KUBE_API_VERSIONS"), validAPIVersions)
		}
		RegisteredVersions = append(RegisteredVersions, version)
	}
}

// Returns true if the given api version is one of the registered api versions.
func IsRegisteredAPIVersion(version string) bool {
	for _, apiVersion := range RegisteredVersions {
		if apiVersion == version {
			return true
		}
	}
	return false
}
<commit_msg>UPSTREAM: <carry>: Leave v1beta3 enabled for now<commit_after>/*
Copyright 2015 The Kubernetes Authors All rights reserved.

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 to keep track of API Versions that should be registered in api.Scheme.
package registered

import (
	"os"
	"strings"

	"github.com/golang/glog"
)

// List of registered API versions.
// The list is in the order of most preferred to the least.
var RegisteredVersions []string

func init() {
	validAPIVersions := map[string]bool{
		"v1":      true,
		"v1beta3": true,
	}

	// The default list of supported api versions, in order of most preferred to the least.
	defaultSupportedVersions := "v1,v1beta3"
	// Env var KUBE_API_VERSIONS is a comma separated list of API versions that should be registered in the scheme.
	// The versions should be in the order of most preferred to the least.
	supportedVersions := os.Getenv("KUBE_API_VERSIONS")
	if supportedVersions == "" {
		supportedVersions = defaultSupportedVersions
	}
	versions := strings.Split(supportedVersions, ",")
	for _, version := range versions {
		// Verify that the version is valid.
		valid, ok := validAPIVersions[version]
		if !ok || !valid {
			// Not a valid API version.
			glog.Fatalf("invalid api version: %s in KUBE_API_VERSIONS: %s. List of valid API versions: %v",
				version, os.Getenv("KUBE_API_VERSIONS"), validAPIVersions)
		}
		RegisteredVersions = append(RegisteredVersions, version)
	}
}

// Returns true if the given api version is one of the registered api versions.
func IsRegisteredAPIVersion(version string) bool {
	for _, apiVersion := range RegisteredVersions {
		if apiVersion == version {
			return true
		}
	}
	return false
}
<|endoftext|>