<commit_before>package main

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

	"github.com/levenlabs/go-llog"
	"github.com/levenlabs/go-srvclient"
	"github.com/levenlabs/golib/rpcutil"
	"github.com/mediocregopher/lever"
)

func main() {
	l := lever.New("rpc-cli", &lever.Opts{
		HelpHeader:         "Usage: rpc-cli [options] <url> <service.method> <parameters>\n",
		DisallowConfigFile: true,
	})
	l.Add(lever.Param{
		Name:        "--llog",
		Description: "Instead of outputting the return as json, output a llog-style message with the returned json object as the key/value params",
		Flag:        true,
	})
	l.Parse()

	argv := l.ParamRest()
	if len(argv) != 3 {
		fmt.Print(l.Help())
		exit(0)
	}

	u, method, body := argv[0], argv[1], argv[2]
	u = srvclient.MaybeSRVURL(u)

	var ret interface{}
	err := rpcutil.JSONRPC2RawCall(u, &ret, method, body)
	if err != nil {
		llog.Error("error calling rpc method", llog.KV{"err": err})
		exit(1)
	}

	if l.ParamFlag("--llog") {
		retm, ok := ret.(map[string]interface{})
		if !ok {
			llog.Error("return value not a json object, can't llog")
			exit(1)
		}
		llog.Info("output from "+method, llog.KV(retm))
		exit(0)
	}

	out, err := json.MarshalIndent(ret, "", "    ")
	if err != nil {
		llog.Error("error marshalling json", llog.KV{"err": err})
		exit(1)
	}

	fmt.Println(string(out))
}

func exit(i int) {
	time.Sleep(100 * time.Millisecond)
	os.Exit(i)
}
<commit_msg>make pretty-fying the rpc-cli output optional<commit_after>package main

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

	"github.com/levenlabs/go-llog"
	"github.com/levenlabs/go-srvclient"
	"github.com/levenlabs/golib/rpcutil"
	"github.com/mediocregopher/lever"
)

func main() {
	l := lever.New("rpc-cli", &lever.Opts{
		HelpHeader:         "Usage: rpc-cli [options] <url> <service.method> <parameters>\n",
		DisallowConfigFile: true,
	})
	l.Add(lever.Param{
		Name:        "--llog",
		Description: "Instead of outputting the return as json, output a llog-style message with the returned json object as the key/value params",
		Flag:        true,
	})
	l.Add(lever.Param{
		Name:        "--pretty",
		Description: "Prettify the returned JSON",
		Flag:        true,
	})
	l.Parse()

	argv := l.ParamRest()
	if len(argv) != 3 {
		fmt.Print(l.Help())
		exit(0)
	}

	u, method, body := argv[0], argv[1], argv[2]
	u = srvclient.MaybeSRVURL(u)

	var ret interface{}
	err := rpcutil.JSONRPC2RawCall(u, &ret, method, body)
	if err != nil {
		llog.Error("error calling rpc method", llog.KV{"err": err})
		exit(1)
	}

	if l.ParamFlag("--llog") {
		retm, ok := ret.(map[string]interface{})
		if !ok {
			llog.Error("return value not a json object, can't llog")
			exit(1)
		}
		llog.Info("output from "+method, llog.KV(retm))
		exit(0)
	}

	var out []byte

	if l.ParamFlag("--pretty") {
		out, err = json.MarshalIndent(ret, "", "    ")
	} else {
		out, err = json.Marshal(ret)
	}
	if err != nil {
		llog.Error("error marshalling json", llog.KV{"err": err})
		exit(1)
	}

	fmt.Println(string(out))
}

func exit(i int) {
	time.Sleep(100 * time.Millisecond)
	os.Exit(i)
}
<|endoftext|>


<commit_before>// Copyright 2017 HyperHQ Inc.
//
// SPDX-License-Identifier: Apache-2.0
//

package main

import (
	"errors"
	"io"
	"os"
	"os/signal"
	"sync"
	"syscall"

	"github.com/moby/moby/pkg/term"
	"github.com/sirupsen/logrus"
	context "golang.org/x/net/context"

	pb "github.com/kata-containers/agent/protocols/grpc"
)

const sigChanSize = 2048

var sigIgnored = map[syscall.Signal]bool{
	syscall.SIGCHLD: true,
	syscall.SIGPIPE: true,
}

type shim struct {
	containerID string
	execID      string

	ctx   context.Context
	agent *shimAgent
}

func newShim(addr, containerID, execID string) (*shim, error) {
	agent, err := newShimAgent(addr)
	if err != nil {
		return nil, err
	}

	return &shim{containerID: containerID,
		execID: execID,
		ctx:    context.Background(),
		agent:  agent}, nil
}

func (s *shim) proxyStdio(wg *sync.WaitGroup, terminal bool) {
	// don't wait the copying of the stdin, because `io.Copy(inPipe, os.Stdin)`
	// can't terminate when no input. todo: find a better way.
	wg.Add(1)
	if !terminal {
		// In case it's not a terminal, we also need to get the output
		// from stderr.
		wg.Add(1)
	}

	inPipe, outPipe, errPipe := shimStdioPipe(s.ctx, s.agent, s.containerID, s.execID)
	go func() {
		_, err1 := io.Copy(inPipe, os.Stdin)
		_, err2 := s.agent.CloseStdin(s.ctx, &pb.CloseStdinRequest{
			ContainerId: s.containerID,
			ExecId:      s.execID})
		if err1 != nil {
			logger().WithError(err1).Warn("copy stdin failed")
		}
		if err2 != nil {
			logger().WithError(err2).Warn("close stdin failed")
		}
	}()

	go func() {
		_, err := io.Copy(os.Stdout, outPipe)
		if err != nil {
			logger().WithError(err).Info("copy stdout failed")
		}

		wg.Done()
	}()

	if !terminal {
		go func() {
			_, err := io.Copy(os.Stderr, errPipe)
			if err != nil {
				logger().WithError(err).Info("copy stderr failed")
			}

			wg.Done()
		}()
	}
}

// handleSignals performs all signal handling.
//
// The tty parameter is specific to SIGWINCH handling.
func (s *shim) handleSignals(tty *os.File) chan os.Signal {
	sigc := make(chan os.Signal, sigChanSize)
	// handle all signals for the process.
	signal.Notify(sigc)
	signal.Ignore(syscall.SIGCHLD, syscall.SIGPIPE)

	go func() {
		for sig := range sigc {
			sysSig, ok := sig.(syscall.Signal)
			if !ok {
				err := errors.New("unknown signal")
				logger().WithError(err).WithField("signal", sig.String()).Error()
				continue
			}

			if sigIgnored[sysSig] {
				//ignore these
				continue
			}

			if sysSig == syscall.SIGWINCH {
				s.resizeTty(tty)

				// Don't actually send the signal to the agent
				// in the container since the resize call will
				// request (via the agent) that the kernel send
				// the signal to the real workload.
				continue
			} else if debug && nonFatalSignal(sysSig) {
				logger().WithField("signal", sig).Debug("handling signal")
				backtrace()
			}

			// forward this signal to container
			_, err := s.agent.SignalProcess(s.ctx, &pb.SignalProcessRequest{
				ContainerId: s.containerID,
				ExecId:      s.execID,
				Signal:      uint32(sysSig)})
			if err != nil {
				logger().WithError(err).WithField("signal", sig.String()).Error("forward signal failed")
			}

			if fatalSignal(sysSig) {
				logger().WithField("signal", sig).Error("received fatal signal")
				die()
			}
		}
	}()
	return sigc
}

func (s *shim) resizeTty(fromTty *os.File) error {
	fd := fromTty.Fd()

	ws, err := term.GetWinsize(fd)
	if err != nil {
		logger().WithError(err).WithField("fd", fd).Info("Error getting window size")
		return nil
	}

	_, err = s.agent.TtyWinResize(s.ctx, &pb.TtyWinResizeRequest{
		ContainerId: s.containerID,
		ExecId:      s.execID,
		Row:         uint32(ws.Height),
		Column:      uint32(ws.Width)})
	if err != nil {
		logger().WithError(err).WithFields(logrus.Fields{
			"window-height": ws.Height,
			"window-width":  ws.Width,
		}).Error("set window size failed")
	}

	return err
}

func (s *shim) wait() (int32, error) {
	resp, err := s.agent.WaitProcess(s.ctx, &pb.WaitProcessRequest{
		ContainerId: s.containerID,
		ExecId:      s.execID})
	if err != nil {
		return 0, err
	}

	return resp.Status, nil
}
<commit_msg>signals: Always show message when handling a signal<commit_after>// Copyright 2017 HyperHQ Inc.
//
// SPDX-License-Identifier: Apache-2.0
//

package main

import (
	"errors"
	"io"
	"os"
	"os/signal"
	"sync"
	"syscall"

	"github.com/moby/moby/pkg/term"
	"github.com/sirupsen/logrus"
	context "golang.org/x/net/context"

	pb "github.com/kata-containers/agent/protocols/grpc"
)

const sigChanSize = 2048

var sigIgnored = map[syscall.Signal]bool{
	syscall.SIGCHLD: true,
	syscall.SIGPIPE: true,
}

type shim struct {
	containerID string
	execID      string

	ctx   context.Context
	agent *shimAgent
}

func newShim(addr, containerID, execID string) (*shim, error) {
	agent, err := newShimAgent(addr)
	if err != nil {
		return nil, err
	}

	return &shim{containerID: containerID,
		execID: execID,
		ctx:    context.Background(),
		agent:  agent}, nil
}

func (s *shim) proxyStdio(wg *sync.WaitGroup, terminal bool) {
	// don't wait the copying of the stdin, because `io.Copy(inPipe, os.Stdin)`
	// can't terminate when no input. todo: find a better way.
	wg.Add(1)
	if !terminal {
		// In case it's not a terminal, we also need to get the output
		// from stderr.
		wg.Add(1)
	}

	inPipe, outPipe, errPipe := shimStdioPipe(s.ctx, s.agent, s.containerID, s.execID)
	go func() {
		_, err1 := io.Copy(inPipe, os.Stdin)
		_, err2 := s.agent.CloseStdin(s.ctx, &pb.CloseStdinRequest{
			ContainerId: s.containerID,
			ExecId:      s.execID})
		if err1 != nil {
			logger().WithError(err1).Warn("copy stdin failed")
		}
		if err2 != nil {
			logger().WithError(err2).Warn("close stdin failed")
		}
	}()

	go func() {
		_, err := io.Copy(os.Stdout, outPipe)
		if err != nil {
			logger().WithError(err).Info("copy stdout failed")
		}

		wg.Done()
	}()

	if !terminal {
		go func() {
			_, err := io.Copy(os.Stderr, errPipe)
			if err != nil {
				logger().WithError(err).Info("copy stderr failed")
			}

			wg.Done()
		}()
	}
}

// handleSignals performs all signal handling.
//
// The tty parameter is specific to SIGWINCH handling.
func (s *shim) handleSignals(tty *os.File) chan os.Signal {
	sigc := make(chan os.Signal, sigChanSize)
	// handle all signals for the process.
	signal.Notify(sigc)
	signal.Ignore(syscall.SIGCHLD, syscall.SIGPIPE)

	go func() {
		for sig := range sigc {
			sysSig, ok := sig.(syscall.Signal)
			if !ok {
				err := errors.New("unknown signal")
				logger().WithError(err).WithField("signal", sig.String()).Error()
				continue
			}

			if sigIgnored[sysSig] {
				//ignore these
				continue
			}

			logger().WithField("signal", sig).Debug("handling signal")

			if sysSig == syscall.SIGWINCH {
				s.resizeTty(tty)

				// Don't actually send the signal to the agent
				// in the container since the resize call will
				// request (via the agent) that the kernel send
				// the signal to the real workload.
				continue
			} else if debug && nonFatalSignal(sysSig) {
				// only backtrace in debug mode for security
				// reasons.
				backtrace()
			}

			// forward this signal to container
			_, err := s.agent.SignalProcess(s.ctx, &pb.SignalProcessRequest{
				ContainerId: s.containerID,
				ExecId:      s.execID,
				Signal:      uint32(sysSig)})
			if err != nil {
				logger().WithError(err).WithField("signal", sig.String()).Error("forward signal failed")
			}

			if fatalSignal(sysSig) {
				logger().WithField("signal", sig).Error("received fatal signal")
				die()
			}
		}
	}()
	return sigc
}

func (s *shim) resizeTty(fromTty *os.File) error {
	fd := fromTty.Fd()

	ws, err := term.GetWinsize(fd)
	if err != nil {
		logger().WithError(err).WithField("fd", fd).Info("Error getting window size")
		return nil
	}

	_, err = s.agent.TtyWinResize(s.ctx, &pb.TtyWinResizeRequest{
		ContainerId: s.containerID,
		ExecId:      s.execID,
		Row:         uint32(ws.Height),
		Column:      uint32(ws.Width)})
	if err != nil {
		logger().WithError(err).WithFields(logrus.Fields{
			"window-height": ws.Height,
			"window-width":  ws.Width,
		}).Error("set window size failed")
	}

	return err
}

func (s *shim) wait() (int32, error) {
	resp, err := s.agent.WaitProcess(s.ctx, &pb.WaitProcessRequest{
		ContainerId: s.containerID,
		ExecId:      s.execID})
	if err != nil {
		return 0, err
	}

	return resp.Status, nil
}
<|endoftext|>


<commit_before>package gorpc

import (
	"bytes"
	"sync"
	"time"

	"github.com/gsdocker/gserrors"
	"github.com/gsdocker/gslogger"
)

// Sink .
type Sink interface {
	Channel
	AddService(dispatcher Dispatcher)
	RemoveService(dispatcher Dispatcher)
	MessageReceived(message *Message) error
}

type _Sink struct {
	gslogger.Log                       // Mixin Log APIs
	sync.RWMutex                       // mutex
	name         string                // name
	timeout      time.Duration         // timeout
	seqID        uint32                // sequence id
	dispatchers  map[uint16]Dispatcher // register dispatchers
	promises     map[uint32]Promise    // rpc promise
	channel      MessageChannel        // channel
	eventLoop    EventLoop             // event loop
}

// NewSink .
func NewSink(name string, eventLoop EventLoop, channel MessageChannel, timeout time.Duration) Sink {
	return &_Sink{
		Log:         gslogger.Get("sink"),
		name:        name,
		timeout:     timeout,
		dispatchers: make(map[uint16]Dispatcher),
		promises:    make(map[uint32]Promise),
		channel:     channel,
		eventLoop:   eventLoop,
	}
}

// Register .
func (sink *_Sink) AddService(dispatcher Dispatcher) {
	sink.Lock()
	defer sink.Unlock()

	sink.dispatchers[dispatcher.ID()] = dispatcher
}

// Unreigster .
func (sink *_Sink) RemoveService(dispatcher Dispatcher) {
	sink.Lock()
	defer sink.Unlock()

	if dispatcher, ok := sink.dispatchers[dispatcher.ID()]; ok && dispatcher == dispatcher {
		delete(sink.dispatchers, dispatcher.ID())
	}
}

func (sink *_Sink) Promise() (Promise, uint32) {

	sink.Lock()
	defer sink.Unlock()

	for {
		seqID := sink.seqID

		sink.seqID++

		_, ok := sink.promises[seqID]

		if ok {
			continue
		}

		promise := NewPromise(sink.eventLoop, sink.timeout, func() {
			sink.Lock()
			defer sink.Unlock()

			delete(sink.promises, seqID)
		})

		sink.promises[seqID] = promise

		return promise, seqID
	}
}

// Post .
func (sink *_Sink) Post(call *Request) error {

	var buff bytes.Buffer
	err := WriteRequest(&buff, call)
	if err != nil {
		return err
	}

	message := NewMessage()
	message.Code = CodeRequest
	message.Content = buff.Bytes()

	err = sink.channel.SendMessage(message)
	if err != nil {
		return err
	}

	sink.V("%s post request(%d:%d:%d)", sink.name, call.ID, call.Service, call.Method)

	return nil
}

// Send .
func (sink *_Sink) Send(call *Request) (Future, error) {

	promise, id := sink.Promise()

	call.ID = id

	var buff bytes.Buffer

	err := WriteRequest(&buff, call)

	if err != nil {
		promise.Cancel()
		return nil, err
	}

	message := NewMessage()

	message.Code = CodeRequest

	message.Content = buff.Bytes()

	err = sink.channel.SendMessage(message)

	if err != nil {
		promise.Cancel()
		return nil, err
	}

	sink.V("%s send request(%d:%d:%d)", sink.name, call.ID, call.Service, call.Method)

	return promise, nil
}

func (sink *_Sink) dispatchResponse(response *Response) {
	sink.Lock()
	defer sink.Unlock()

	if promise, ok := sink.promises[response.ID]; ok {
		go promise.Notify(response, nil)
		delete(sink.promises, response.ID)
		return
	}

	sink.W("%s unhandle response(%d) %s", sink.name, response.ID, gserrors.Newf(nil, ""))
}

func (sink *_Sink) dispatch(id uint16) (dispatcher Dispatcher, ok bool) {
	sink.RLock()
	defer sink.RUnlock()

	dispatcher, ok = sink.dispatchers[id]

	return
}

func (sink *_Sink) dispatchRequest(message *Message) error {

	request, err := ReadRequest(bytes.NewBuffer(message.Content))

	if err != nil {
		sink.E("[%s] unmarshal request error\n%s", sink.name, err)
		return err
	}

	sink.V("%s received request(%d:%d:%d)", sink.name, request.ID, request.Service, request.Method)

	if dispatcher, ok := sink.dispatch(request.Service); ok {

		response, err := dispatcher.Dispatch(request)

		if err != nil {
			sink.E("dispatch request(%d)(%d:%d) error\n%s", request.ID, request.Service, request.Method, err)
			return nil
		}

		if response == nil {
			return nil
		}

		var buff bytes.Buffer

		err = WriteResponse(&buff, response)

		if err != nil {
			sink.E("marshal request(%d)(%d:%d)'s response error\n%s", request.ID, request.Service, request.Method, err)
			return nil
		}

		message.Code = CodeResponse

		message.Content = buff.Bytes()

		sink.channel.SendMessage(message)

		sink.V("[%s] send response(%d)(%d)", sink.name, response.ID, response.Exception)

		return nil
	}

	sink.W("[%s] unhandle request(%d)(%d:%d)", sink.name, request.ID, request.Service, request.Method)

	return nil
}

func (sink *_Sink) MessageReceived(message *Message) error {

	switch message.Code {
	case CodeRequest:

		sink.dispatchRequest(message)

		return nil

	case CodeResponse:

		response, err := ReadResponse(bytes.NewBuffer(message.Content))

		if err != nil {
			sink.E("[%s] unmarshal response error\n%s", sink.name, err)
			return err
		}

		sink.dispatchResponse(response)

		return nil

	default:

		sink.W("[%s] unsupport message(%s)", sink.name, message.Code)

		return nil
	}
}
<commit_msg>add clear services for sink<commit_after>package gorpc

import (
	"bytes"
	"sync"
	"time"

	"github.com/gsdocker/gserrors"
	"github.com/gsdocker/gslogger"
)

// Sink .
type Sink interface {
	Channel
	AddService(dispatcher Dispatcher)
	RemoveService(dispatcher Dispatcher)
	ClearServices()
	MessageReceived(message *Message) error
}

type _Sink struct {
	gslogger.Log                       // Mixin Log APIs
	sync.RWMutex                       // mutex
	name         string                // name
	timeout      time.Duration         // timeout
	seqID        uint32                // sequence id
	dispatchers  map[uint16]Dispatcher // register dispatchers
	promises     map[uint32]Promise    // rpc promise
	channel      MessageChannel        // channel
	eventLoop    EventLoop             // event loop
}

// NewSink .
func NewSink(name string, eventLoop EventLoop, channel MessageChannel, timeout time.Duration) Sink {
	return &_Sink{
		Log:         gslogger.Get("sink"),
		name:        name,
		timeout:     timeout,
		dispatchers: make(map[uint16]Dispatcher),
		promises:    make(map[uint32]Promise),
		channel:     channel,
		eventLoop:   eventLoop,
	}
}

// Register .
func (sink *_Sink) AddService(dispatcher Dispatcher) {
	sink.Lock()
	defer sink.Unlock()

	sink.dispatchers[dispatcher.ID()] = dispatcher
}

// Unreigster .
func (sink *_Sink) RemoveService(dispatcher Dispatcher) {
	sink.Lock()
	defer sink.Unlock()

	if dispatcher, ok := sink.dispatchers[dispatcher.ID()]; ok && dispatcher == dispatcher {
		delete(sink.dispatchers, dispatcher.ID())
	}
}

// ClearServices .
func (sink *_Sink) ClearServices() {
	sink.Lock()
	defer sink.Unlock()

	sink.dispatchers = make(map[uint16]Dispatcher)
}

func (sink *_Sink) Promise() (Promise, uint32) {

	sink.Lock()
	defer sink.Unlock()

	for {
		seqID := sink.seqID

		sink.seqID++

		_, ok := sink.promises[seqID]

		if ok {
			continue
		}

		promise := NewPromise(sink.eventLoop, sink.timeout, func() {
			sink.Lock()
			defer sink.Unlock()

			delete(sink.promises, seqID)
		})

		sink.promises[seqID] = promise

		return promise, seqID
	}
}

// Post .
func (sink *_Sink) Post(call *Request) error {

	var buff bytes.Buffer
	err := WriteRequest(&buff, call)
	if err != nil {
		return err
	}

	message := NewMessage()
	message.Code = CodeRequest
	message.Content = buff.Bytes()

	err = sink.channel.SendMessage(message)
	if err != nil {
		return err
	}

	sink.V("%s post request(%d:%d:%d)", sink.name, call.ID, call.Service, call.Method)

	return nil
}

// Send .
func (sink *_Sink) Send(call *Request) (Future, error) {

	promise, id := sink.Promise()

	call.ID = id

	var buff bytes.Buffer

	err := WriteRequest(&buff, call)

	if err != nil {
		promise.Cancel()
		return nil, err
	}

	message := NewMessage()

	message.Code = CodeRequest

	message.Content = buff.Bytes()

	err = sink.channel.SendMessage(message)

	if err != nil {
		promise.Cancel()
		return nil, err
	}

	sink.V("%s send request(%d:%d:%d)", sink.name, call.ID, call.Service, call.Method)

	return promise, nil
}

func (sink *_Sink) dispatchResponse(response *Response) {
	sink.Lock()
	defer sink.Unlock()

	if promise, ok := sink.promises[response.ID]; ok {
		go promise.Notify(response, nil)
		delete(sink.promises, response.ID)
		return
	}

	sink.W("%s unhandle response(%d) %s", sink.name, response.ID, gserrors.Newf(nil, ""))
}

func (sink *_Sink) dispatch(id uint16) (dispatcher Dispatcher, ok bool) {
	sink.RLock()
	defer sink.RUnlock()

	dispatcher, ok = sink.dispatchers[id]

	return
}

func (sink *_Sink) dispatchRequest(message *Message) error {

	request, err := ReadRequest(bytes.NewBuffer(message.Content))

	if err != nil {
		sink.E("[%s] unmarshal request error\n%s", sink.name, err)
		return err
	}

	sink.V("%s received request(%d:%d:%d)", sink.name, request.ID, request.Service, request.Method)

	if dispatcher, ok := sink.dispatch(request.Service); ok {

		response, err := dispatcher.Dispatch(request)

		if err != nil {
			sink.E("dispatch request(%d)(%d:%d) error\n%s", request.ID, request.Service, request.Method, err)
			return nil
		}

		if response == nil {
			return nil
		}

		var buff bytes.Buffer

		err = WriteResponse(&buff, response)

		if err != nil {
			sink.E("marshal request(%d)(%d:%d)'s response error\n%s", request.ID, request.Service, request.Method, err)
			return nil
		}

		message.Code = CodeResponse

		message.Content = buff.Bytes()

		sink.channel.SendMessage(message)

		sink.V("[%s] send response(%d)(%d)", sink.name, response.ID, response.Exception)

		return nil
	}

	sink.W("[%s] unhandle request(%d)(%d:%d)", sink.name, request.ID, request.Service, request.Method)

	return nil
}

func (sink *_Sink) MessageReceived(message *Message) error {

	switch message.Code {
	case CodeRequest:

		sink.dispatchRequest(message)

		return nil

	case CodeResponse:

		response, err := ReadResponse(bytes.NewBuffer(message.Content))

		if err != nil {
			sink.E("[%s] unmarshal response error\n%s", sink.name, err)
			return err
		}

		sink.dispatchResponse(response)

		return nil

	default:

		sink.W("[%s] unsupport message(%s)", sink.name, message.Code)

		return nil
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"mime"
	"os"
	"path/filepath"
	"text/template"
	"time"

	"launchpad.net/goamz/aws"
	"launchpad.net/goamz/s3"
)

var (
	MsgCopyingFile  = "Copying File: %s"
	MsgGenerateFile = "Generating Page: %s"
	MsgUploadFile   = "Uploading: %s"
	MsgUsingConfig  = "Loading Config: %s"
)

type Site struct {
	Src  string // Directory where Jekyll will look to transform files
	Dest string // Directory where Jekyll will write files to
	Conf Config // Configuration date from the _config.yml file

	posts []Page             // Posts thet need to be generated
	pages []Page             // Pages that need to be generated
	files []string           // Static files to get copied to the destination
	templ *template.Template // Compiled templates
}

func NewSite(src, dest string) (*Site, error) {

	// Parse the _config.yml file
	path := filepath.Join(src, "_config.yml")
	conf, err := ParseConfig(path)
	logf(MsgUsingConfig, path)
	if err != nil {
		return nil, err
	}

	site := Site{
		Src:  src,
		Dest: dest,
		Conf: conf,
	}

	// Recursively process all files in the source directory
	// and parse pages, posts, templates, etc
	if err := site.read(); err != nil {
		return nil, err
	}

	return &site, nil
}

// Reloads the site into memory
func (s *Site) Reload() error {
	s.posts = []Page{}
	s.pages = []Page{}
	s.files = []string{}
	s.templ = nil
	return s.read()
}

// Prepares the source directory for site generation
func (s *Site) Prep() error {
	return os.MkdirAll(s.Dest, 0755)
}

// Removes the existing site (typically in _site).
func (s *Site) Clear() error {
	return os.RemoveAll(s.Dest)
}

// Generates a static website based on Jekyll standard layout.
func (s *Site) Generate() error {

	// Remove previously generated site, and then (re)create the
	// destination directory
	if err := s.Clear(); err != nil {
		return err
	}
	if err := s.Prep(); err != nil {
		return err
	}

	// Generate all Pages and Posts and static files
	if err := s.writePages(); err != nil {
		return err
	}
	if err := s.writeStatic(); err != nil {
		return err
	}

	return nil
}

// Deploys a site to S3.
func (s *Site) Deploy(user, pass, url string) error {

	auth := aws.Auth{AccessKey: user, SecretKey: pass}
	b := s3.New(auth, aws.USEast).Bucket(url)

	// walks _site directory and uploads file to S3
	walker := func(fn string, fi os.FileInfo, err error) error {
		if fi.IsDir() {
			return nil
		}

		rel, _ := filepath.Rel(s.Dest, fn)
		typ := mime.TypeByExtension(filepath.Ext(rel))
		content, err := ioutil.ReadFile(fn)
		logf(MsgUploadFile, rel)
		if err != nil {
			return err
		}

		// try to upload the file ... sometimes this fails due to amazon
		// issues. If so, we'll re-try
		if err := b.Put(rel, content, typ, s3.PublicRead); err != nil {
			time.Sleep(100 * time.Millisecond) // sleep so that we don't immediately retry
			return b.Put(rel, content, typ, s3.PublicRead)
		}

		// file upload was a success, return nil
		return nil
	}

	return filepath.Walk(s.Dest, walker)
}

// Helper function to traverse the source directory and identify all posts,
// projects, templates, etc and parse.
func (s *Site) read() error {

	// Lists of templates (_layouts, _includes) that we find that
	// will need to be compiled
	layouts := []string{}

	// func to walk the jekyll directory structure
	walker := func(fn string, fi os.FileInfo, err error) error {
		rel, _ := filepath.Rel(s.Src, fn)
		switch {
		case err != nil:
			return nil

		case fi.IsDir() && isHiddenOrTemp(fn):
			return filepath.SkipDir

		// Ignore directories
		case fi.IsDir():
			return nil

		// Ignore Hidden or Temp files
		// (starting with . or ending with ~)
		case isHiddenOrTemp(rel):
			return nil

		// Parse Templates
		case isTemplate(rel):
			layouts = append(layouts, fn)

		// Parse Posts
		case isPost(rel):
			post, err := ParsePost(rel)
			if err != nil {
				return err
			}
			// TODO: this is a hack to get the posts in rev chronological order
			s.posts = append([]Page{post}, s.posts...) //s.posts, post)

		// Parse Pages
		case isPage(rel):
			page, err := ParsePage(rel)
			if err != nil {
				return err
			}
			s.pages = append(s.pages, page)

		// Move static files, no processing required
		case isStatic(rel):
			s.files = append(s.files, rel)
		}
		return nil
	}

	// Walk the diretory recursively to get a list of all posts,
	// pages, templates and static files.
	err := filepath.Walk(s.Src, walker)
	if err != nil {
		return err
	}

	// Compile all templates found, if any
	if len(layouts) > 0 {
		s.templ, err = template.New("layouts").Funcs(funcMap).ParseFiles(layouts...)
		if err != nil {
			return err
		}
	}

	// Add the posts, timestamp, etc to the Site Params
	s.Conf.Set("posts", s.posts)
	s.Conf.Set("time", time.Now())
	s.calculateTags()
	s.calculateCategories()

	return nil
}

// Helper function to write all pages and posts to the destination directory
// during site generation.
func (s *Site) writePages() error {

	// There is really no difference between a Page and a Post (other than
	// initial parsing) so we can combine the lists and use the same rendering
	// code for both.
	pages := []Page{}
	pages = append(pages, s.pages...)
	pages = append(pages, s.posts...)

	for _, page := range pages {
		url := page.GetUrl()
		layout := page.GetLayout()

		// is the layout provided? or is it nil /empty?
		//layoutNil := layout == "" || layout == "nil"

		// make sure the posts's parent dir exists
		d := filepath.Join(s.Dest, filepath.Dir(url))
		f := filepath.Join(s.Dest, url)
		if err := os.MkdirAll(d, 0755); err != nil {
			return err
		}

		// if markdown, need to convert to html
		// otherwise just convert raw html to a string
		//var content string
		//if isMarkdown(page.GetExt()) {
		//	content = string(blackfriday.MarkdownCommon(raw))
		//} else {
		//	content = string(raw)
		//}

		//data passed in to each template
		data := map[string]interface{}{
			"site": s.Conf,
			"page": page,
		}

		// treat all non-markdown pages as templates
		content := page.GetContent()
		if isMarkdown(page.GetExt()) == false {
			// this code will add the page to the list of templates,
			// will execute the template, and then set the content
			// to the rendered template

			if s.templ == nil {
				return fmt.Errorf("No templates defined for page: %s", url)
			}

			t, err := s.templ.New(url).Parse(content)
			if err != nil {
				return err
			}
			var buf bytes.Buffer
			err = t.ExecuteTemplate(&buf, url, data)
			if err != nil {
				return err
			}
			content = buf.String()
		}

		// add document body to the map
		data["content"] = content

		// write the template to a buffer
		// NOTE: if template is nil or empty, then we should parse the
		//       content as if it were a template
		var buf bytes.Buffer
		if layout == "" || layout == "nil" {
			//t, err := s.templ.New(url).Parse(content);
			//if err != nil { return err }
			//err = t.ExecuteTemplate(&buf, url, data);
			//if err != nil { return err }

			buf.WriteString(content)
		} else {
			layout = appendExt(layout, ".html")
			s.templ.ExecuteTemplate(&buf, layout, data)
		}

		logf(MsgGenerateFile, url)
		if err := ioutil.WriteFile(f, buf.Bytes(), 0644); err != nil {
			return err
		}
	}

	return nil
}

// Helper function to write all static files to the destination directory
// during site generation. This will also take care of creating any parent
// directories, if necessary.
func (s *Site) writeStatic() error {

	for _, file := range s.files {
		from := filepath.Join(s.Src, file)
		to := filepath.Join(s.Dest, file)
		logf(MsgCopyingFile, file)
		if err := copyTo(from, to); err != nil {
			return err
		}
	}

	return nil
}

// Helper function to aggregate a list of all categories and their
// related posts.
func (s *Site) calculateCategories() {

	categories := make(map[string][]Page)
	for _, post := range s.posts {
		for _, category := range post.GetCategories() {
			if posts, ok := categories[category]; ok == false {
				categories[category] = append(posts, post)
			} else {
				categories[category] = []Page{post}
			}
		}
	}

	s.Conf.Set("categories", categories)
}

// Helper function to aggregate a list of all tags and their
// related posts.
func (s *Site) calculateTags() {

	tags := make(map[string][]Page)
	for _, post := range s.posts {
		for _, tag := range post.GetTags() {
			if posts, ok := tags[tag]; ok == false {
				tags[tag] = append(posts, post)
			} else {
				tags[tag] = []Page{post}
			}
		}
	}

	s.Conf.Set("tags", tags)
}
<commit_msg>goamz moved<commit_after>package main

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"mime"
	"os"
	"path/filepath"
	"text/template"
	"time"

	"gopkg.in/amz.v1/aws"
	"gopkg.in/amz.v1/s3"
)

var (
	MsgCopyingFile  = "Copying File: %s"
	MsgGenerateFile = "Generating Page: %s"
	MsgUploadFile   = "Uploading: %s"
	MsgUsingConfig  = "Loading Config: %s"
)

type Site struct {
	Src  string // Directory where Jekyll will look to transform files
	Dest string // Directory where Jekyll will write files to
	Conf Config // Configuration date from the _config.yml file

	posts []Page             // Posts thet need to be generated
	pages []Page             // Pages that need to be generated
	files []string           // Static files to get copied to the destination
	templ *template.Template // Compiled templates
}

func NewSite(src, dest string) (*Site, error) {

	// Parse the _config.yml file
	path := filepath.Join(src, "_config.yml")
	conf, err := ParseConfig(path)
	logf(MsgUsingConfig, path)
	if err != nil {
		return nil, err
	}

	site := Site{
		Src:  src,
		Dest: dest,
		Conf: conf,
	}

	// Recursively process all files in the source directory
	// and parse pages, posts, templates, etc
	if err := site.read(); err != nil {
		return nil, err
	}

	return &site, nil
}

// Reloads the site into memory
func (s *Site) Reload() error {
	s.posts = []Page{}
	s.pages = []Page{}
	s.files = []string{}
	s.templ = nil
	return s.read()
}

// Prepares the source directory for site generation
func (s *Site) Prep() error {
	return os.MkdirAll(s.Dest, 0755)
}

// Removes the existing site (typically in _site).
func (s *Site) Clear() error {
	return os.RemoveAll(s.Dest)
}

// Generates a static website based on Jekyll standard layout.
func (s *Site) Generate() error {

	// Remove previously generated site, and then (re)create the
	// destination directory
	if err := s.Clear(); err != nil {
		return err
	}
	if err := s.Prep(); err != nil {
		return err
	}

	// Generate all Pages and Posts and static files
	if err := s.writePages(); err != nil {
		return err
	}
	if err := s.writeStatic(); err != nil {
		return err
	}

	return nil
}

// Deploys a site to S3.
func (s *Site) Deploy(user, pass, url string) error {

	auth := aws.Auth{AccessKey: user, SecretKey: pass}
	b := s3.New(auth, aws.USEast).Bucket(url)

	// walks _site directory and uploads file to S3
	walker := func(fn string, fi os.FileInfo, err error) error {
		if fi.IsDir() {
			return nil
		}

		rel, _ := filepath.Rel(s.Dest, fn)
		typ := mime.TypeByExtension(filepath.Ext(rel))
		content, err := ioutil.ReadFile(fn)
		logf(MsgUploadFile, rel)
		if err != nil {
			return err
		}

		// try to upload the file ... sometimes this fails due to amazon
		// issues. If so, we'll re-try
		if err := b.Put(rel, content, typ, s3.PublicRead); err != nil {
			time.Sleep(100 * time.Millisecond) // sleep so that we don't immediately retry
			return b.Put(rel, content, typ, s3.PublicRead)
		}

		// file upload was a success, return nil
		return nil
	}

	return filepath.Walk(s.Dest, walker)
}

// Helper function to traverse the source directory and identify all posts,
// projects, templates, etc and parse.
func (s *Site) read() error {

	// Lists of templates (_layouts, _includes) that we find that
	// will need to be compiled
	layouts := []string{}

	// func to walk the jekyll directory structure
	walker := func(fn string, fi os.FileInfo, err error) error {
		rel, _ := filepath.Rel(s.Src, fn)
		switch {
		case err != nil:
			return nil

		case fi.IsDir() && isHiddenOrTemp(fn):
			return filepath.SkipDir

		// Ignore directories
		case fi.IsDir():
			return nil

		// Ignore Hidden or Temp files
		// (starting with . or ending with ~)
		case isHiddenOrTemp(rel):
			return nil

		// Parse Templates
		case isTemplate(rel):
			layouts = append(layouts, fn)

		// Parse Posts
		case isPost(rel):
			post, err := ParsePost(rel)
			if err != nil {
				return err
			}
			// TODO: this is a hack to get the posts in rev chronological order
			s.posts = append([]Page{post}, s.posts...) //s.posts, post)

		// Parse Pages
		case isPage(rel):
			page, err := ParsePage(rel)
			if err != nil {
				return err
			}
			s.pages = append(s.pages, page)

		// Move static files, no processing required
		case isStatic(rel):
			s.files = append(s.files, rel)
		}
		return nil
	}

	// Walk the diretory recursively to get a list of all posts,
	// pages, templates and static files.
	err := filepath.Walk(s.Src, walker)
	if err != nil {
		return err
	}

	// Compile all templates found, if any
	if len(layouts) > 0 {
		s.templ, err = template.New("layouts").Funcs(funcMap).ParseFiles(layouts...)
		if err != nil {
			return err
		}
	}

	// Add the posts, timestamp, etc to the Site Params
	s.Conf.Set("posts", s.posts)
	s.Conf.Set("time", time.Now())
	s.calculateTags()
	s.calculateCategories()

	return nil
}

// Helper function to write all pages and posts to the destination directory
// during site generation.
func (s *Site) writePages() error {

	// There is really no difference between a Page and a Post (other than
	// initial parsing) so we can combine the lists and use the same rendering
	// code for both.
	pages := []Page{}
	pages = append(pages, s.pages...)
	pages = append(pages, s.posts...)

	for _, page := range pages {
		url := page.GetUrl()
		layout := page.GetLayout()

		// is the layout provided? or is it nil /empty?
		//layoutNil := layout == "" || layout == "nil"

		// make sure the posts's parent dir exists
		d := filepath.Join(s.Dest, filepath.Dir(url))
		f := filepath.Join(s.Dest, url)
		if err := os.MkdirAll(d, 0755); err != nil {
			return err
		}

		// if markdown, need to convert to html
		// otherwise just convert raw html to a string
		//var content string
		//if isMarkdown(page.GetExt()) {
		//	content = string(blackfriday.MarkdownCommon(raw))
		//} else {
		//	content = string(raw)
		//}

		//data passed in to each template
		data := map[string]interface{}{
			"site": s.Conf,
			"page": page,
		}

		// treat all non-markdown pages as templates
		content := page.GetContent()
		if isMarkdown(page.GetExt()) == false {
			// this code will add the page to the list of templates,
			// will execute the template, and then set the content
			// to the rendered template

			if s.templ == nil {
				return fmt.Errorf("No templates defined for page: %s", url)
			}

			t, err := s.templ.New(url).Parse(content)
			if err != nil {
				return err
			}
			var buf bytes.Buffer
			err = t.ExecuteTemplate(&buf, url, data)
			if err != nil {
				return err
			}
			content = buf.String()
		}

		// add document body to the map
		data["content"] = content

		// write the template to a buffer
		// NOTE: if template is nil or empty, then we should parse the
		//       content as if it were a template
		var buf bytes.Buffer
		if layout == "" || layout == "nil" {
			//t, err := s.templ.New(url).Parse(content);
			//if err != nil { return err }
			//err = t.ExecuteTemplate(&buf, url, data);
			//if err != nil { return err }

			buf.WriteString(content)
		} else {
			layout = appendExt(layout, ".html")
			s.templ.ExecuteTemplate(&buf, layout, data)
		}

		logf(MsgGenerateFile, url)
		if err := ioutil.WriteFile(f, buf.Bytes(), 0644); err != nil {
			return err
		}
	}

	return nil
}

// Helper function to write all static files to the destination directory
// during site generation. This will also take care of creating any parent
// directories, if necessary.
func (s *Site) writeStatic() error {

	for _, file := range s.files {
		from := filepath.Join(s.Src, file)
		to := filepath.Join(s.Dest, file)
		logf(MsgCopyingFile, file)
		if err := copyTo(from, to); err != nil {
			return err
		}
	}

	return nil
}

// Helper function to aggregate a list of all categories and their
// related posts.
func (s *Site) calculateCategories() {

	categories := make(map[string][]Page)
	for _, post := range s.posts {
		for _, category := range post.GetCategories() {
			if posts, ok := categories[category]; ok == false {
				categories[category] = append(posts, post)
			} else {
				categories[category] = []Page{post}
			}
		}
	}

	s.Conf.Set("categories", categories)
}

// Helper function to aggregate a list of all tags and their
// related posts.
func (s *Site) calculateTags() {

	tags := make(map[string][]Page)
	for _, post := range s.posts {
		for _, tag := range post.GetTags() {
			if posts, ok := tags[tag]; ok == false {
				tags[tag] = append(posts, post)
			} else {
				tags[tag] = []Page{post}
			}
		}
	}

	s.Conf.Set("tags", tags)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"net/http"
	"os"
	"strings"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	"github.com/urfave/cli"
	"k8s.io/klog"

	"github.com/gesellix/couchdb-prometheus-exporter/lib"
)

type exporterConfigType struct {
	listenAddress   string
	metricsEndpoint string
	couchdbURI      string
	couchdbUsername string
	couchdbPassword string
	couchdbInsecure bool
	databases       string
	databaseViews   bool
	schedulerJobs   bool
}

type loggingConfigType struct {
	toStderr        bool   // The -logtostderr flag.
	alsoToStderr    bool   // The -alsologtostderr flag.
	verbosity       int    // V logging level, the value of the -v flag/
	stderrThreshold int    // The -stderrthreshold flag.
	logDir          string // The -log_dir flag.
}

var exporterConfig exporterConfigType

// custom exposed logging config flags
var loggingConfig loggingConfigType

var appFlags []cli.Flag

// TODO graceful migration to new parameter names
// 1) Warn, that these parameters are deprecated and will be removed/renamed
// 2) Fail at startup, when deprecated parameters are used. Maybe allow override by explicit "i-know-what-i-am-doing"-parameter
// 3) Remove (ignore) deprecated parameters
func init() {
	// TODO replace mechanism with urfave/cli
	//flag.String(flag.DefaultConfigFlagname, "", "path to config file")

	appFlags = []cli.Flag{
		cli.StringFlag{
			Name:  "telemetry.address",
			Usage: "Address on which to expose metrics",
			//EnvVar:      "TELEMETRY_ADDRESS",
			Hidden:      false,
			Value:       "localhost:9984",
			Destination: &exporterConfig.listenAddress,
		},
		cli.StringFlag{
			Name:  "telemetry.endpoint",
			Usage: "Path under which to expose metrics",
			//EnvVar:      "TELEMETRY_ENDPOINT",
			Hidden:      false,
			Value:       "/metrics",
			Destination: &exporterConfig.metricsEndpoint,
		},
		cli.StringFlag{
			Name:  "couchdb.uri",
			Usage: "URI to the CouchDB instance",
			//EnvVar:      "COUCHDB_URI",
			Hidden:      false,
			Value:       "http://localhost:5984",
			Destination: &exporterConfig.couchdbURI,
		},
		cli.StringFlag{
			Name:  "couchdb.username",
			Usage: "Basic auth username for the CouchDB instance",
			//EnvVar:      "COUCHDB_USERNAME",
			Hidden:      false,
			Value:       "",
			Destination: &exporterConfig.couchdbUsername,
		},
		cli.StringFlag{
			Name:  "couchdb.password",
			Usage: "Basic auth password for the CouchDB instance",
			//EnvVar:      "COUCHDB_PASSWORD",
			Hidden:      false,
			Value:       "",
			Destination: &exporterConfig.couchdbPassword,
		},
		// TODO doesn't print the default when showing the command help
		cli.BoolTFlag{
			Name:  "couchdb.insecure",
			Usage: "Ignore server certificate if using https",
			//EnvVar:      "COUCHDB_INSECURE",
			Hidden:      false,
			Destination: &exporterConfig.couchdbInsecure,
		},
		// TODO use cli.StringSliceFlag?
		cli.StringFlag{
			Name:  "databases",
			Usage: fmt.Sprintf("Comma separated list of database names, or '%s'", lib.AllDbs),
			//EnvVar:      "DATABASES",
			Hidden:      false,
			Value:       "",
			Destination: &exporterConfig.databases,
		},
		// TODO doesn't print the default when showing the command help
		cli.BoolTFlag{
			Name:  "databases.views",
			Usage: "Collect view details of every observed database",
			//EnvVar:      "DATABASES_VIEWS",
			Hidden:      false,
			Destination: &exporterConfig.databaseViews,
		},
		// TODO doesn't print the default when showing the command help
		cli.BoolFlag{
			Name:  "scheduler.jobs",
			Usage: "Collect active replication jobs (CouchDB 2.x+ only)",
			//EnvVar:      "SCHEDULER_JOBS",
			Hidden:      false,
			Destination: &exporterConfig.schedulerJobs,
		},

		cli.BoolFlag{
			Name:        "logtostderr",
			Usage:       "log to standard error instead of files",
			Hidden:      false,
			Destination: &loggingConfig.toStderr,
		},
		cli.BoolFlag{
			Name:        "alsologtostderr",
			Usage:       "log to standard error as well as files",
			Hidden:      false,
			Destination: &loggingConfig.alsoToStderr,
		},
		// TODO clashes with urfave/cli `--version` shortcut `-v`.
		// TODO should be of type `Level` or at least match int32
		//cli.IntFlag{
		//	Name:        "v",
		//	Usage:       "log level for V logs",
		//	Hidden:      false,
		//	Destination: &loggingConfig.verbosity,
		//},
		// TODO should be of type `severity` or at least match int32
		cli.IntFlag{
			Name:        "stderrthreshold",
			Usage:       "logs at or above this threshold go to stderr",
			Hidden:      false,
			Destination: &loggingConfig.stderrThreshold,
		},
		cli.StringFlag{
			Name:        "log_dir",
			Usage:       "If non-empty, write log files in this directory",
			Hidden:      false,
			Destination: &loggingConfig.logDir,
		},
	}

	//klogFlags := flag.NewFlagSet("klog", flag.ContinueOnError)
	klog.InitFlags(nil)
}

func main() {
	// > Programs should call Flush before exiting to guarantee all log output is written.
	defer klog.Flush()

	var appAction = func(c *cli.Context) error {
		var databases []string
		if *&exporterConfig.databases != "" {
			databases = strings.Split(*&exporterConfig.databases, ",")
		}

		exporter := lib.NewExporter(
			*&exporterConfig.couchdbURI,
			lib.BasicAuth{
				Username: *&exporterConfig.couchdbUsername,
				Password: *&exporterConfig.couchdbPassword},
			lib.CollectorConfig{
				Databases:            databases,
				CollectViews:         *&exporterConfig.databaseViews,
				CollectSchedulerJobs: *&exporterConfig.schedulerJobs,
			},
			*&exporterConfig.couchdbInsecure)
		prometheus.MustRegister(exporter)

		http.Handle(*&exporterConfig.metricsEndpoint, promhttp.Handler())
		http.HandleFunc("/status", func(w http.ResponseWriter, r *http.Request) {
			_, err := fmt.Fprint(w, "OK")
			if err != nil {
				klog.Error(err)
			}
		})
		http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
			http.Error(w, fmt.Sprintf("Please GET %s", *&exporterConfig.metricsEndpoint), http.StatusNotFound)
		})

		klog.Infof("Starting exporter at '%s' to read from CouchDB at '%s'", *&exporterConfig.listenAddress, *&exporterConfig.couchdbURI)
		err := http.ListenAndServe(*&exporterConfig.listenAddress, nil)
		if err != nil {
			klog.Fatal(err)
		}
		return err
	}

	app := cli.NewApp()
	app.Name = "CouchDB Prometheus Exporter"
	app.Usage = ""
	app.Description = "CouchDB stats exporter for Prometheus"
	app.Version = ""
	app.Flags = appFlags
	app.Action = appAction

	err := app.Run(os.Args)
	if err != nil {
		klog.Fatal(err)
	}
}
<commit_msg>Replace the former logging verbosity shortcut "v" to "verbosity"<commit_after>package main

import (
	"fmt"
	"net/http"
	"os"
	"strings"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	"github.com/urfave/cli"
	"k8s.io/klog"

	"github.com/gesellix/couchdb-prometheus-exporter/lib"
)

type exporterConfigType struct {
	listenAddress   string
	metricsEndpoint string
	couchdbURI      string
	couchdbUsername string
	couchdbPassword string
	couchdbInsecure bool
	databases       string
	databaseViews   bool
	schedulerJobs   bool
}

type loggingConfigType struct {
	toStderr        bool   // The -logtostderr flag.
	alsoToStderr    bool   // The -alsologtostderr flag.
	verbosity       int    // V logging level, the value of the -v flag/
	stderrThreshold int    // The -stderrthreshold flag.
	logDir          string // The -log_dir flag.
}

var exporterConfig exporterConfigType

// custom exposed logging config flags
var loggingConfig loggingConfigType

var appFlags []cli.Flag

// TODO graceful migration to new parameter names
// 1) Warn, that these parameters are deprecated and will be removed/renamed
// 2) Fail at startup, when deprecated parameters are used. Maybe allow override by explicit "i-know-what-i-am-doing"-parameter
// 3) Remove (ignore) deprecated parameters
func init() {
	// TODO replace mechanism with urfave/cli
	//flag.String(flag.DefaultConfigFlagname, "", "path to config file")

	appFlags = []cli.Flag{
		cli.StringFlag{
			Name:  "telemetry.address",
			Usage: "Address on which to expose metrics",
			//EnvVar:      "TELEMETRY_ADDRESS",
			Hidden:      false,
			Value:       "localhost:9984",
			Destination: &exporterConfig.listenAddress,
		},
		cli.StringFlag{
			Name:  "telemetry.endpoint",
			Usage: "Path under which to expose metrics",
			//EnvVar:      "TELEMETRY_ENDPOINT",
			Hidden:      false,
			Value:       "/metrics",
			Destination: &exporterConfig.metricsEndpoint,
		},
		cli.StringFlag{
			Name:  "couchdb.uri",
			Usage: "URI to the CouchDB instance",
			//EnvVar:      "COUCHDB_URI",
			Hidden:      false,
			Value:       "http://localhost:5984",
			Destination: &exporterConfig.couchdbURI,
		},
		cli.StringFlag{
			Name:  "couchdb.username",
			Usage: "Basic auth username for the CouchDB instance",
			//EnvVar:      "COUCHDB_USERNAME",
			Hidden:      false,
			Value:       "",
			Destination: &exporterConfig.couchdbUsername,
		},
		cli.StringFlag{
			Name:  "couchdb.password",
			Usage: "Basic auth password for the CouchDB instance",
			//EnvVar:      "COUCHDB_PASSWORD",
			Hidden:      false,
			Value:       "",
			Destination: &exporterConfig.couchdbPassword,
		},
		// TODO doesn't print the default when showing the command help
		cli.BoolTFlag{
			Name:  "couchdb.insecure",
			Usage: "Ignore server certificate if using https",
			//EnvVar:      "COUCHDB_INSECURE",
			Hidden:      false,
			Destination: &exporterConfig.couchdbInsecure,
		},
		// TODO use cli.StringSliceFlag?
		cli.StringFlag{
			Name:  "databases",
			Usage: fmt.Sprintf("Comma separated list of database names, or '%s'", lib.AllDbs),
			//EnvVar:      "DATABASES",
			Hidden:      false,
			Value:       "",
			Destination: &exporterConfig.databases,
		},
		// TODO doesn't print the default when showing the command help
		cli.BoolTFlag{
			Name:  "databases.views",
			Usage: "Collect view details of every observed database",
			//EnvVar:      "DATABASES_VIEWS",
			Hidden:      false,
			Destination: &exporterConfig.databaseViews,
		},
		// TODO doesn't print the default when showing the command help
		cli.BoolFlag{
			Name:  "scheduler.jobs",
			Usage: "Collect active replication jobs (CouchDB 2.x+ only)",
			//EnvVar:      "SCHEDULER_JOBS",
			Hidden:      false,
			Destination: &exporterConfig.schedulerJobs,
		},

		cli.BoolFlag{
			Name:        "logtostderr",
			Usage:       "log to standard error instead of files",
			Hidden:      false,
			Destination: &loggingConfig.toStderr,
		},
		cli.BoolFlag{
			Name:        "alsologtostderr",
			Usage:       "log to standard error as well as files",
			Hidden:      false,
			Destination: &loggingConfig.alsoToStderr,
		},
		// TODO `v` clashed with urfave/cli's `--version` shortcut `-v`.
		// TODO should be of type `Level` or at least match int32
		cli.IntFlag{
			Name:        "verbosity",
			Usage:       "log level for V logs",
			Hidden:      false,
			Destination: &loggingConfig.verbosity,
		},
		// TODO should be of type `severity` or at least match int32
		cli.IntFlag{
			Name:        "stderrthreshold",
			Usage:       "logs at or above this threshold go to stderr",
			Hidden:      false,
			Destination: &loggingConfig.stderrThreshold,
		},
		cli.StringFlag{
			Name:        "log_dir",
			Usage:       "If non-empty, write log files in this directory",
			Hidden:      false,
			Destination: &loggingConfig.logDir,
		},
	}

	//klogFlags := flag.NewFlagSet("klog", flag.ContinueOnError)
	klog.InitFlags(nil)
}

func main() {
	defer klog.Flush()

	var appAction = func(c *cli.Context) error {
		var databases []string
		if *&exporterConfig.databases != "" {
			databases = strings.Split(*&exporterConfig.databases, ",")
		}

		exporter := lib.NewExporter(
			*&exporterConfig.couchdbURI,
			lib.BasicAuth{
				Username: *&exporterConfig.couchdbUsername,
				Password: *&exporterConfig.couchdbPassword},
			lib.CollectorConfig{
				Databases:            databases,
				CollectViews:         *&exporterConfig.databaseViews,
				CollectSchedulerJobs: *&exporterConfig.schedulerJobs,
			},
			*&exporterConfig.couchdbInsecure)
		prometheus.MustRegister(exporter)

		http.Handle(*&exporterConfig.metricsEndpoint, promhttp.Handler())
		http.HandleFunc("/status", func(w http.ResponseWriter, r *http.Request) {
			_, err := fmt.Fprint(w, "OK")
			if err != nil {
				klog.Error(err)
			}
		})
		http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
			http.Error(w, fmt.Sprintf("Please GET %s", *&exporterConfig.metricsEndpoint), http.StatusNotFound)
		})

		klog.Infof("Starting exporter at '%s' to read from CouchDB at '%s'", *&exporterConfig.listenAddress, *&exporterConfig.couchdbURI)
		err := http.ListenAndServe(*&exporterConfig.listenAddress, nil)
		if err != nil {
			klog.Fatal(err)
		}
		return err
	}

	app := cli.NewApp()
	app.Name = "CouchDB Prometheus Exporter"
	app.Usage = ""
	app.Description = "CouchDB stats exporter for Prometheus"
	app.Version = ""
	app.Flags = appFlags
	app.Action = appAction

	err := app.Run(os.Args)
	if err != nil {
		klog.Fatal(err)
	}
}
<|endoftext|>


<commit_before>// Package mount provides a Datastore that has other Datastores
// mounted at various key prefixes and is threadsafe
package syncmount

import (
	"errors"
	"fmt"
	"io"
	"sort"
	"strings"
	"sync"

	ds "github.com/ipfs/go-datastore"
	"github.com/ipfs/go-datastore/query"
)

var (
	ErrNoMount = errors.New("no datastore mounted for this key")
)

type Mount struct {
	Prefix    ds.Key
	Datastore ds.Datastore
}

func New(mounts []Mount) *Datastore {
	// make a copy so we're sure it doesn't mutate
	m := make([]Mount, len(mounts))
	for i, v := range mounts {
		m[i] = v
	}
	sort.Slice(m, func(i, j int) bool { return m[i].Prefix.String() > m[j].Prefix.String() })
	return &Datastore{mounts: m}
}

type Datastore struct {
	mounts []Mount
	lk     sync.Mutex
}

var _ ds.Datastore = (*Datastore)(nil)

func (d *Datastore) lookup(key ds.Key) (ds.Datastore, ds.Key, ds.Key) {
	d.lk.Lock()
	defer d.lk.Unlock()
	for _, m := range d.mounts {
		if m.Prefix.Equal(key) || m.Prefix.IsAncestorOf(key) {
			s := strings.TrimPrefix(key.String(), m.Prefix.String())
			k := ds.NewKey(s)
			return m.Datastore, m.Prefix, k
		}
	}
	return nil, ds.NewKey("/"), key
}

// lookupAll returns all mounts that might contain keys that are descendant of <key>
//
// Matching: /ao/e
//
// /          B /ao/e
// /a/        not matching
// /ao/       B /e
// /ao/e/     A /
// /ao/e/uh/  A /
// /aoe/      not matching
func (d *Datastore) lookupAll(key ds.Key) (dst []ds.Datastore, mountpoint, rest []ds.Key) {
	d.lk.Lock()
	defer d.lk.Unlock()

	for _, m := range d.mounts {
		p := m.Prefix.String()
		if len(p) > 1 {
			p = p + "/"
		}

		if strings.HasPrefix(p, key.String()) {
			dst = append(dst, m.Datastore)
			mountpoint = append(mountpoint, m.Prefix)
			rest = append(rest, ds.NewKey("/"))
		} else if strings.HasPrefix(key.String(), p) {
			r := strings.TrimPrefix(key.String(), m.Prefix.String())

			dst = append(dst, m.Datastore)
			mountpoint = append(mountpoint, m.Prefix)
			rest = append(rest, ds.NewKey(r))
		}
	}
	return dst, mountpoint, rest
}

func (d *Datastore) Put(key ds.Key, value interface{}) error {
	cds, _, k := d.lookup(key)
	if cds == nil {
		return ErrNoMount
	}
	return cds.Put(k, value)
}

func (d *Datastore) Get(key ds.Key) (value interface{}, err error) {
	cds, _, k := d.lookup(key)
	if cds == nil {
		return nil, ds.ErrNotFound
	}
	return cds.Get(k)
}

func (d *Datastore) Has(key ds.Key) (exists bool, err error) {
	cds, _, k := d.lookup(key)
	if cds == nil {
		return false, nil
	}
	return cds.Has(k)
}

func (d *Datastore) Delete(key ds.Key) error {
	cds, _, k := d.lookup(key)
	if cds == nil {
		return ds.ErrNotFound
	}
	return cds.Delete(k)
}

func (d *Datastore) Query(q query.Query) (query.Results, error) {
	if len(q.Filters) > 0 ||
		len(q.Orders) > 0 ||
		q.Limit > 0 ||
		q.Offset > 0 {
		// TODO this is still overly simplistic, but the only callers are
		// `ipfs refs local` and ipfs-ds-convert.
		return nil, errors.New("mount only supports listing all prefixed keys in random order")
	}
	prefix := ds.NewKey(q.Prefix)
	dses, mounts, rests := d.lookupAll(prefix)

	// current itorator state
	var res query.Results
	var mount ds.Key
	i := 0

	return query.ResultsFromIterator(q, query.Iterator{
		Next: func() (query.Result, bool) {
			var r query.Result
			var more bool

			for try := true; try; try = len(dses) > i {
				if res == nil {
					if len(dses) <= i {
						//This should not happen normally
						return query.Result{}, false
					}

					dst := dses[i]
					mount = mounts[i]
					rest := rests[i]

					q2 := q
					q2.Prefix = rest.String()
					r, err := dst.Query(q2)
					if err != nil {
						return query.Result{Error: err}, false
					}
					res = r
				}

				r, more = res.NextSync()
				if !more {
					err := res.Close()
					if err != nil {
						return query.Result{Error: err}, false
					}
					res = nil

					i++
					more = len(dses) > i
				} else {
					break
				}
			}

			r.Key = mount.Child(ds.RawKey(r.Key)).String()
			return r, more
		},
		Close: func() error {
			if len(mounts) > i && res != nil {
				return res.Close()
			}
			return nil
		},
	}), nil
}

func (d *Datastore) IsThreadSafe() {}

func (d *Datastore) Close() error {
	for _, d := range d.mounts {
		if c, ok := d.Datastore.(io.Closer); ok {
			err := c.Close()
			if err != nil {
				return err
			}
		}
	}
	return nil
}

type mountBatch struct {
	mounts map[string]ds.Batch
	lk     sync.Mutex

	d *Datastore
}

func (d *Datastore) Batch() (ds.Batch, error) {
	return &mountBatch{
		mounts: make(map[string]ds.Batch),
		d:      d,
	}, nil
}

func (mt *mountBatch) lookupBatch(key ds.Key) (ds.Batch, ds.Key, error) {
	mt.lk.Lock()
	defer mt.lk.Unlock()

	child, loc, rest := mt.d.lookup(key)
	t, ok := mt.mounts[loc.String()]
	if !ok {
		bds, ok := child.(ds.Batching)
		if !ok {
			return nil, ds.NewKey(""), ds.ErrBatchUnsupported
		}
		var err error
		t, err = bds.Batch()
		if err != nil {
			return nil, ds.NewKey(""), err
		}
		mt.mounts[loc.String()] = t
	}
	return t, rest, nil
}

func (mt *mountBatch) Put(key ds.Key, val interface{}) error {
	t, rest, err := mt.lookupBatch(key)
	if err != nil {
		return err
	}

	return t.Put(rest, val)
}

func (mt *mountBatch) Delete(key ds.Key) error {
	t, rest, err := mt.lookupBatch(key)
	if err != nil {
		return err
	}

	return t.Delete(rest)
}

func (mt *mountBatch) Commit() error {
	for _, t := range mt.mounts {
		err := t.Commit()
		if err != nil {
			return err
		}
	}
	return nil
}

func (d *Datastore) Check() error {
	for _, m := range d.mounts {
		if c, ok := m.Datastore.(ds.CheckedDatastore); ok {
			if err := c.Check(); err != nil {
				return fmt.Errorf("checking datastore at %s: %s", m.Prefix.String(), err.Error())
			}
		}
	}
	return nil
}

func (d *Datastore) Scrub() error {
	for _, m := range d.mounts {
		if c, ok := m.Datastore.(ds.ScrubbedDatastore); ok {
			if err := c.Scrub(); err != nil {
				return fmt.Errorf("scrubbing datastore at %s: %s", m.Prefix.String(), err.Error())
			}
		}
	}
	return nil
}

func (d *Datastore) CollectGarbage() error {
	for _, m := range d.mounts {
		if c, ok := m.Datastore.(ds.GCDatastore); ok {
			if err := c.CollectGarbage(); err != nil {
				return fmt.Errorf("gc on datastore at %s: %s", m.Prefix.String(), err.Error())
			}
		}
	}
	return nil
}
<commit_msg>Fix lock usage in syncmount<commit_after>// Package mount provides a Datastore that has other Datastores
// mounted at various key prefixes and is threadsafe
package syncmount

import (
	"errors"
	"fmt"
	"io"
	"sort"
	"strings"
	"sync"

	ds "github.com/ipfs/go-datastore"
	"github.com/ipfs/go-datastore/query"
)

var (
	ErrNoMount = errors.New("no datastore mounted for this key")
)

type Mount struct {
	Prefix    ds.Key
	Datastore ds.Datastore
}

func New(mounts []Mount) *Datastore {
	// make a copy so we're sure it doesn't mutate
	m := make([]Mount, len(mounts))
	for i, v := range mounts {
		m[i] = v
	}
	sort.Slice(m, func(i, j int) bool { return m[i].Prefix.String() > m[j].Prefix.String() })
	return &Datastore{mounts: m}
}

type Datastore struct {
	mounts []Mount
}

var _ ds.Datastore = (*Datastore)(nil)

func (d *Datastore) lookup(key ds.Key) (ds.Datastore, ds.Key, ds.Key) {
	for _, m := range d.mounts {
		if m.Prefix.Equal(key) || m.Prefix.IsAncestorOf(key) {
			s := strings.TrimPrefix(key.String(), m.Prefix.String())
			k := ds.NewKey(s)
			return m.Datastore, m.Prefix, k
		}
	}
	return nil, ds.NewKey("/"), key
}

// lookupAll returns all mounts that might contain keys that are descendant of <key>
//
// Matching: /ao/e
//
// /          B /ao/e
// /a/        not matching
// /ao/       B /e
// /ao/e/     A /
// /ao/e/uh/  A /
// /aoe/      not matching
func (d *Datastore) lookupAll(key ds.Key) (dst []ds.Datastore, mountpoint, rest []ds.Key) {
	for _, m := range d.mounts {
		p := m.Prefix.String()
		if len(p) > 1 {
			p = p + "/"
		}

		if strings.HasPrefix(p, key.String()) {
			dst = append(dst, m.Datastore)
			mountpoint = append(mountpoint, m.Prefix)
			rest = append(rest, ds.NewKey("/"))
		} else if strings.HasPrefix(key.String(), p) {
			r := strings.TrimPrefix(key.String(), m.Prefix.String())

			dst = append(dst, m.Datastore)
			mountpoint = append(mountpoint, m.Prefix)
			rest = append(rest, ds.NewKey(r))
		}
	}
	return dst, mountpoint, rest
}

func (d *Datastore) Put(key ds.Key, value interface{}) error {
	cds, _, k := d.lookup(key)
	if cds == nil {
		return ErrNoMount
	}
	return cds.Put(k, value)
}

func (d *Datastore) Get(key ds.Key) (value interface{}, err error) {
	cds, _, k := d.lookup(key)
	if cds == nil {
		return nil, ds.ErrNotFound
	}
	return cds.Get(k)
}

func (d *Datastore) Has(key ds.Key) (exists bool, err error) {
	cds, _, k := d.lookup(key)
	if cds == nil {
		return false, nil
	}
	return cds.Has(k)
}

func (d *Datastore) Delete(key ds.Key) error {
	cds, _, k := d.lookup(key)
	if cds == nil {
		return ds.ErrNotFound
	}
	return cds.Delete(k)
}

func (d *Datastore) Query(q query.Query) (query.Results, error) {
	if len(q.Filters) > 0 ||
		len(q.Orders) > 0 ||
		q.Limit > 0 ||
		q.Offset > 0 {
		// TODO this is still overly simplistic, but the only callers are
		// `ipfs refs local` and ipfs-ds-convert.
		return nil, errors.New("mount only supports listing all prefixed keys in random order")
	}
	prefix := ds.NewKey(q.Prefix)
	dses, mounts, rests := d.lookupAll(prefix)

	// current itorator state
	var res query.Results
	var mount ds.Key
	i := 0

	return query.ResultsFromIterator(q, query.Iterator{
		Next: func() (query.Result, bool) {
			var r query.Result
			var more bool

			for try := true; try; try = len(dses) > i {
				if res == nil {
					if len(dses) <= i {
						//This should not happen normally
						return query.Result{}, false
					}

					dst := dses[i]
					mount = mounts[i]
					rest := rests[i]

					q2 := q
					q2.Prefix = rest.String()
					r, err := dst.Query(q2)
					if err != nil {
						return query.Result{Error: err}, false
					}
					res = r
				}

				r, more = res.NextSync()
				if !more {
					err := res.Close()
					if err != nil {
						return query.Result{Error: err}, false
					}
					res = nil

					i++
					more = len(dses) > i
				} else {
					break
				}
			}

			r.Key = mount.Child(ds.RawKey(r.Key)).String()
			return r, more
		},
		Close: func() error {
			if len(mounts) > i && res != nil {
				return res.Close()
			}
			return nil
		},
	}), nil
}

func (d *Datastore) IsThreadSafe() {}

func (d *Datastore) Close() error {
	for _, d := range d.mounts {
		if c, ok := d.Datastore.(io.Closer); ok {
			err := c.Close()
			if err != nil {
				return err
			}
		}
	}
	return nil
}

type mountBatch struct {
	mounts map[string]ds.Batch
	lk     sync.Mutex

	d *Datastore
}

func (d *Datastore) Batch() (ds.Batch, error) {
	return &mountBatch{
		mounts: make(map[string]ds.Batch),
		d:      d,
	}, nil
}

func (mt *mountBatch) lookupBatch(key ds.Key) (ds.Batch, ds.Key, error) {
	mt.lk.Lock()
	defer mt.lk.Unlock()

	child, loc, rest := mt.d.lookup(key)
	t, ok := mt.mounts[loc.String()]
	if !ok {
		bds, ok := child.(ds.Batching)
		if !ok {
			return nil, ds.NewKey(""), ds.ErrBatchUnsupported
		}
		var err error
		t, err = bds.Batch()
		if err != nil {
			return nil, ds.NewKey(""), err
		}
		mt.mounts[loc.String()] = t
	}
	return t, rest, nil
}

func (mt *mountBatch) Put(key ds.Key, val interface{}) error {
	t, rest, err := mt.lookupBatch(key)
	if err != nil {
		return err
	}

	return t.Put(rest, val)
}

func (mt *mountBatch) Delete(key ds.Key) error {
	t, rest, err := mt.lookupBatch(key)
	if err != nil {
		return err
	}

	return t.Delete(rest)
}

func (mt *mountBatch) Commit() error {
	mt.lk.Lock()
	defer mt.lk.Unlock()

	for _, t := range mt.mounts {
		err := t.Commit()
		if err != nil {
			return err
		}
	}
	return nil
}

func (d *Datastore) Check() error {
	for _, m := range d.mounts {
		if c, ok := m.Datastore.(ds.CheckedDatastore); ok {
			if err := c.Check(); err != nil {
				return fmt.Errorf("checking datastore at %s: %s", m.Prefix.String(), err.Error())
			}
		}
	}
	return nil
}

func (d *Datastore) Scrub() error {
	for _, m := range d.mounts {
		if c, ok := m.Datastore.(ds.ScrubbedDatastore); ok {
			if err := c.Scrub(); err != nil {
				return fmt.Errorf("scrubbing datastore at %s: %s", m.Prefix.String(), err.Error())
			}
		}
	}
	return nil
}

func (d *Datastore) CollectGarbage() error {
	for _, m := range d.mounts {
		if c, ok := m.Datastore.(ds.GCDatastore); ok {
			if err := c.CollectGarbage(); err != nil {
				return fmt.Errorf("gc on datastore at %s: %s", m.Prefix.String(), err.Error())
			}
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"runtime/pprof"
	"sort"
	"text/tabwriter"
)

const VERSION = `0.1`

var languages = []Language{
	Language{"C", mExt(".c", ".h"), cComments},
	Language{"C++", mExt(".cc", ".cpp", ".cxx", ".hh", ".hpp", ".hxx"), cComments},
	Language{"Go", mExt(".go"), cComments},
	Language{"Haskell", mExt(".hs", ".lhs"), noComments},
	Language{"Perl", mExt(".pl", ".pm"), shComments},

	Language{"Shell", mExt(".sh"), shComments},
	Language{"Bash", mExt(".bash"), shComments},

	Language{"Python", mExt(".py"), noComments},
	Language{"Assembly", mExt(".asm", ".s"), semiComments},
	Language{"Lisp", mExt(".lsp", ".lisp"), semiComments},

	Language{"Make", mName("makefile", "Makefile", "MAKEFILE"), shComments},
	Language{"Jam", mName("Jamfile", "Jamrules"), shComments},

	Language{"Markdown", mExt(".md"), noComments},

	Language{"HTML", mExt(".htm", ".html", ".xhtml"), xmlComments},
	Language{"XML", mExt(".xml"), xmlComments},
	Language{"CSS", mExt(".css"), cssComments},
	Language{"JavaScript", mExt(".js"), cComments},
}

type Commenter struct {
	LineComment  string
	StartComment string
	EndComment   string
	Nesting      bool
}

var (
	noComments = Commenter{"\000", "\000", "\000", false}
	xmlComments = Commenter{"\000", `<!--`, `-->`, false}
	cComments  = Commenter{`//`, `/*`, `*/`, false}
	cssComments  = Commenter{"\000", `/*`, `*/`, false}
	shComments = Commenter{`#`, "\000", "\000", false}
	semiComments = Commenter{`;`, "\000", "\000", false}
)

type Language struct {
	Namer
	Matcher
	Commenter
}

// TODO work properly with unicode
func (l Language) Update(c []byte, s *Stats) {
	s.FileCount++

	inComment := 0 // this is an int for nesting
	inLComment := false
	blank := true
	lc := []byte(l.LineComment)
	sc := []byte(l.StartComment)
	ec := []byte(l.EndComment)
	lp, sp, ep := 0, 0, 0

	for _, b := range c {
		if inComment == 0 && b == lc[lp] {
			lp++
			if lp == len(lc) {
				inLComment = true
				lp = 0
			}
		} else { lp = 0 }
		if !inLComment && b == sc[sp] {
			sp++
			if sp == len(sc) {
				inComment++
				if inComment > 1 && !l.Nesting {
					inComment = 1
				}
				sp = 0
			}
		} else { sp = 0 }
		if !inLComment && inComment > 0 && b == ec[ep] {
			ep++
			if ep == len(ec) {
				if inComment > 0 { inComment-- }
				ep = 0
			}
		} else { ep = 0 }

		if b != byte(' ') && b != byte('\t') && b != byte('\n') {
			blank = false
		}

		// BUG(srl): lines with comment don't count towards code
		// Note that lines with both code and comment count towards
		// each, but are not counted twice in the total.
		if b == byte('\n') {
			s.TotalLines++
			if inComment > 0 || inLComment {
				inLComment = false
				s.CommentLines++
			} else if blank {
				s.BlankLines++
			} else { s.CodeLines++ }
			blank = true
			continue
		}
	}
}

type Namer string

func (l Namer) Name() string { return string(l) }

type Matcher func(string) bool

func (m Matcher) Match(fname string) bool { return m(fname) }

func mExt(exts ...string) Matcher {
	return func(fname string) bool {
		for _, ext := range exts {
			if ext == path.Ext(fname) {
				return true
			}
		}
		return false
	}
}

func mName(names ...string) Matcher {
	return func(fname string) bool {
		for _, name := range names {
			if name == path.Base(fname) {
				return true
			}
		}
		return false
	}
}

type Stats struct {
	FileCount    int
	TotalLines   int
	CodeLines    int
	BlankLines   int
	CommentLines int
}

var info = map[string]*Stats{}

func handleFile(fname string) {
	var l Language
	ok := false
	for _, lang := range languages {
		if lang.Match(fname) {
			ok = true
			l = lang
			break
		}
	}
	if !ok {
		return // ignore this file
	}
	i, ok := info[l.Name()]
	if !ok {
		i = &Stats{}
		info[l.Name()] = i
	}
	c, err := ioutil.ReadFile(fname)
	if err != nil {
		fmt.Fprintf(os.Stderr, "  ! %s\n", fname)
		return
	}
	l.Update(c, i)
}

var files []string

func add(n string) {
	fi, err := os.Stat(n)
	if err != nil {
		goto invalid
	}
	if fi.IsDir() {
		fs, err := ioutil.ReadDir(n)
		if err != nil {
			goto invalid
		}
		for _, f := range fs {
			if f.Name()[0] != '.' {
				add(path.Join(n, f.Name()))
			}
		}
		return
	}
	if fi.Mode()&os.ModeType == 0 {
		files = append(files, n)
		return
	}

	println(fi.Mode())

invalid:
	fmt.Fprintf(os.Stderr, "  ! %s\n", n)
}

type LData []LResult

func (d LData) Len() int { return len(d) }

func (d LData) Less(i, j int) bool {
	if d[i].CodeLines == d[j].CodeLines {
		return d[i].Name > d[j].Name
	}
	return d[i].CodeLines > d[j].CodeLines
}

func (d LData) Swap(i, j int) {
	d[i], d[j] = d[j], d[i]
}

type LResult struct {
	Name string
	FileCount int
	CodeLines int
	CommentLines int
	BlankLines int
	TotalLines int
}

func (r *LResult) Add(a LResult) {
	r.FileCount += a.FileCount
	r.CodeLines += a.CodeLines
	r.CommentLines += a.CommentLines
	r.BlankLines += a.BlankLines
	r.TotalLines += a.TotalLines
}

func printJSON() {
	bs, err := json.MarshalIndent(info, "", "  ")
	if err != nil { panic(err) }
	fmt.Println(string(bs))
}

func printInfo() {
	w := tabwriter.NewWriter(os.Stdout, 2, 8, 2, ' ', tabwriter.AlignRight)
	fmt.Fprintln(w, "Language\tFiles\tCode\tComment\tBlank\tTotal\t")
	d := LData([]LResult{})
	total := &LResult{}
	total.Name = "Total"
	for n, i := range info {
		r := LResult{
			n,
			i.FileCount,
			i.CodeLines,
			i.CommentLines,
			i.BlankLines,
			i.TotalLines,
		}
		d = append(d, r)
		total.Add(r)
	}
	d = append(d, *total)
	sort.Sort(d)
	//d[0].Name = "Total"
	for _, i := range d {
		fmt.Fprintf(
			w,
			"%s\t%d\t%d\t%d\t%d\t%d\t\n",
			i.Name,
			i.FileCount,
			i.CodeLines,
			i.CommentLines,
			i.BlankLines,
			i.TotalLines)
	}

	w.Flush()
}

var (
	cpuprofile = flag.String("cpuprofile", "", "write cpu profile to file")
	useJson = flag.Bool("json", false, "JSON-format output")
	version = flag.Bool("V", false, "display version info and exit")
)

func main() {
	flag.Parse()
	if *version {
		fmt.Printf("sloc %s\n", VERSION)
		return
	}
	if *cpuprofile != "" {
		f, err := os.Create(*cpuprofile)
		if err != nil {
			fmt.Fprintf(os.Stderr, "error: %s\n", err.Error())
			return
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	args := flag.Args()
	if len(args) == 0 {
		args = append(args, `.`)
	}

	for _, n := range args {
		add(n)
	}

	for _, f := range files {
		handleFile(f)
	}

	if *useJson {
		printJSON()
	} else {
		printInfo()
	}
}
<commit_msg>Paying attention to .gitignore files<commit_after>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"runtime/pprof"
	"sort"
	"strings"
	"text/tabwriter"
)

const VERSION = `0.1`

var languages = []Language{
	Language{"C", mExt(".c", ".h"), cComments},
	Language{"C++", mExt(".cc", ".cpp", ".cxx", ".hh", ".hpp", ".hxx"), cComments},
	Language{"Go", mExt(".go"), cComments},
	Language{"Haskell", mExt(".hs", ".lhs"), noComments},
	Language{"Perl", mExt(".pl", ".pm"), shComments},

	Language{"Shell", mExt(".sh"), shComments},
	Language{"Bash", mExt(".bash"), shComments},

	Language{"Python", mExt(".py"), noComments},
	Language{"Assembly", mExt(".asm", ".s"), semiComments},
	Language{"Lisp", mExt(".lsp", ".lisp"), semiComments},

	Language{"Make", mName("makefile", "Makefile", "MAKEFILE"), shComments},
	Language{"Jam", mName("Jamfile", "Jamrules"), shComments},

	Language{"Markdown", mExt(".md"), noComments},

	Language{"HTML", mExt(".htm", ".html", ".xhtml"), xmlComments},
	Language{"XML", mExt(".xml"), xmlComments},
	Language{"CSS", mExt(".css"), cssComments},
	Language{"JavaScript", mExt(".js"), cComments},
}

type Commenter struct {
	LineComment  string
	StartComment string
	EndComment   string
	Nesting      bool
}

var (
	noComments = Commenter{"\000", "\000", "\000", false}
	xmlComments = Commenter{"\000", `<!--`, `-->`, false}
	cComments  = Commenter{`//`, `/*`, `*/`, false}
	cssComments  = Commenter{"\000", `/*`, `*/`, false}
	shComments = Commenter{`#`, "\000", "\000", false}
	semiComments = Commenter{`;`, "\000", "\000", false}
)

type Language struct {
	Namer
	Matcher
	Commenter
}

// TODO work properly with unicode
func (l Language) Update(c []byte, s *Stats) {
	s.FileCount++

	inComment := 0 // this is an int for nesting
	inLComment := false
	blank := true
	lc := []byte(l.LineComment)
	sc := []byte(l.StartComment)
	ec := []byte(l.EndComment)
	lp, sp, ep := 0, 0, 0

	for _, b := range c {
		if inComment == 0 && b == lc[lp] {
			lp++
			if lp == len(lc) {
				inLComment = true
				lp = 0
			}
		} else { lp = 0 }
		if !inLComment && b == sc[sp] {
			sp++
			if sp == len(sc) {
				inComment++
				if inComment > 1 && !l.Nesting {
					inComment = 1
				}
				sp = 0
			}
		} else { sp = 0 }
		if !inLComment && inComment > 0 && b == ec[ep] {
			ep++
			if ep == len(ec) {
				if inComment > 0 { inComment-- }
				ep = 0
			}
		} else { ep = 0 }

		if b != byte(' ') && b != byte('\t') && b != byte('\n') {
			blank = false
		}

		// BUG(srl): lines with comment don't count towards code
		// Note that lines with both code and comment count towards
		// each, but are not counted twice in the total.
		if b == byte('\n') {
			s.TotalLines++
			if inComment > 0 || inLComment {
				inLComment = false
				s.CommentLines++
			} else if blank {
				s.BlankLines++
			} else { s.CodeLines++ }
			blank = true
			continue
		}
	}
}

type Namer string

func (l Namer) Name() string { return string(l) }

type Matcher func(string) bool

func (m Matcher) Match(fname string) bool { return m(fname) }

func mExt(exts ...string) Matcher {
	return func(fname string) bool {
		for _, ext := range exts {
			if ext == filepath.Ext(fname) {
				return true
			}
		}
		return false
	}
}

func mName(names ...string) Matcher {
	return func(fname string) bool {
		for _, name := range names {
			if name == filepath.Base(fname) {
				return true
			}
		}
		return false
	}
}

type Stats struct {
	FileCount    int
	TotalLines   int
	CodeLines    int
	BlankLines   int
	CommentLines int
}

var info = map[string]*Stats{}

func handleFile(fname string) {
	var l Language
	ok := false
	for _, lang := range languages {
		if lang.Match(fname) {
			ok = true
			l = lang
			break
		}
	}
	if !ok {
		return // ignore this file
	}
	i, ok := info[l.Name()]
	if !ok {
		i = &Stats{}
		info[l.Name()] = i
	}
	c, err := ioutil.ReadFile(fname)
	if err != nil {
		fmt.Fprintf(os.Stderr, "  ! %s\n", fname)
		return
	}
	l.Update(c, i)
}

var ignoreGlobs []string
var files []string

func add(n string) {
	fi, err := os.Stat(n)
	if err != nil {
		goto invalid
	}
	if fi.IsDir() {
		n += "/"
	}
	for _, p := range ignoreGlobs {
		if m, _ := filepath.Match(p, n); m {
			return
		}
	}
	if fi.IsDir() {
		fs, err := ioutil.ReadDir(n)
		if err != nil {
			goto invalid
		}
		for _, f := range fs {
			if f.Name()[0] != '.' {
				add(filepath.Join(n, f.Name()))
			}
		}
		return
	}
	if fi.Mode()&os.ModeType == 0 {
		files = append(files, n)
		return
	}

	println(fi.Mode())

invalid:
	fmt.Fprintf(os.Stderr, "  ! %s\n", n)
}

type LData []LResult

func (d LData) Len() int { return len(d) }

func (d LData) Less(i, j int) bool {
	if d[i].CodeLines == d[j].CodeLines {
		return d[i].Name > d[j].Name
	}
	return d[i].CodeLines > d[j].CodeLines
}

func (d LData) Swap(i, j int) {
	d[i], d[j] = d[j], d[i]
}

type LResult struct {
	Name string
	FileCount int
	CodeLines int
	CommentLines int
	BlankLines int
	TotalLines int
}

func (r *LResult) Add(a LResult) {
	r.FileCount += a.FileCount
	r.CodeLines += a.CodeLines
	r.CommentLines += a.CommentLines
	r.BlankLines += a.BlankLines
	r.TotalLines += a.TotalLines
}

func printJSON() {
	bs, err := json.MarshalIndent(info, "", "  ")
	if err != nil { panic(err) }
	fmt.Println(string(bs))
}

func printInfo() {
	w := tabwriter.NewWriter(os.Stdout, 2, 8, 2, ' ', tabwriter.AlignRight)
	fmt.Fprintln(w, "Language\tFiles\tCode\tComment\tBlank\tTotal\t")
	d := LData([]LResult{})
	total := &LResult{}
	total.Name = "Total"
	for n, i := range info {
		r := LResult{
			n,
			i.FileCount,
			i.CodeLines,
			i.CommentLines,
			i.BlankLines,
			i.TotalLines,
		}
		d = append(d, r)
		total.Add(r)
	}
	d = append(d, *total)
	sort.Sort(d)
	//d[0].Name = "Total"
	for _, i := range d {
		fmt.Fprintf(
			w,
			"%s\t%d\t%d\t%d\t%d\t%d\t\n",
			i.Name,
			i.FileCount,
			i.CodeLines,
			i.CommentLines,
			i.BlankLines,
			i.TotalLines)
	}

	w.Flush()
}

func tryReadGitIgnore() {
	ign, err := ioutil.ReadFile(".gitignore")
	if err != nil {
		return
	}
	gs := strings.Split(string(ign), "\n")
	ignoreGlobs = append(ignoreGlobs, gs...)
}

var (
	cpuprofile = flag.String("cpuprofile", "", "write cpu profile to file")
	useJson = flag.Bool("json", false, "JSON-format output")
	version = flag.Bool("V", false, "display version info and exit")
)

func main() {
	flag.Parse()
	if *version {
		fmt.Printf("sloc %s\n", VERSION)
		return
	}
	if *cpuprofile != "" {
		f, err := os.Create(*cpuprofile)
		if err != nil {
			fmt.Fprintf(os.Stderr, "error: %s\n", err.Error())
			return
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	args := flag.Args()
	if len(args) == 0 {
		args = append(args, `.`)
	}

	tryReadGitIgnore()

	for _, n := range args {
		add(n)
	}

	for _, f := range files {
		handleFile(f)
	}

	if *useJson {
		printJSON()
	} else {
		printInfo()
	}
}
<|endoftext|>


<commit_before>// Package slug provides utilities for generating URL-friendly slugs.
package slug

import (
	"strings"
	"unicode"

	"github.com/fiam/gounidecode/unidecode"
)

// Slug returns a slug generated from s, with the maximum length l and the separator sep.
func Slug(s string, l int, sep string) string {
	decoded := unidecode.Unidecode(s)
	slug := ""

	for _, d := range decoded {
		r := convert(d, sep)

		switch r {
		case "":
			continue
		case sep:
			if len(slug) > 0 && !strings.HasSuffix(slug, sep) {
				slug += sep
			}
		default:
			slug += string(r)
		}

		if len(slug) >= l {
			break
		}
	}

	return strings.TrimSuffix(slug, sep)
}

// convert returns the slug representation of r, using the separator sep.
func convert(r rune, sep string) string {
	switch {
	case 'a' <= r && r <= 'z' || '0' <= r && r <= '9':
		return string(r)
	case 'A' <= r && r <= 'Z':
		return string(unicode.ToLower(r))
	case r == ' ' || r == ',' || r == '.' || r == '/' || r == '\\' || r == '-' || r == '_' ||
		r == '=':
		return sep
	}
	return ""
}
<commit_msg>Add usage example<commit_after>// Package slug provides an utility for generating URL-friendly slugs.
//
// Example:
//
//	package main
//
//	import (
//		"github.com/peferron/go-slug"
//	)
//
//	func main() {
//		s := slug.Slug("Hello World!", 50, "-")
//		print(s) // Prints "hello-world"
//	}
package slug

import (
	"strings"
	"unicode"

	"github.com/fiam/gounidecode/unidecode"
)

// Slug returns a slug generated from s, with the maximum length l and the separator sep.
func Slug(s string, l int, sep string) string {
	decoded := unidecode.Unidecode(s)
	slug := ""

	for _, d := range decoded {
		r := convert(d, sep)

		switch r {
		case "":
			continue
		case sep:
			if len(slug) > 0 && !strings.HasSuffix(slug, sep) {
				slug += sep
			}
		default:
			slug += string(r)
		}

		if len(slug) >= l {
			break
		}
	}

	return strings.TrimSuffix(slug, sep)
}

// convert returns the slug representation of r, using the separator sep.
func convert(r rune, sep string) string {
	switch {
	case 'a' <= r && r <= 'z' || '0' <= r && r <= '9':
		return string(r)
	case 'A' <= r && r <= 'Z':
		return string(unicode.ToLower(r))
	case r == ' ' || r == ',' || r == '.' || r == '/' || r == '\\' || r == '-' || r == '_' ||
		r == '=':
		return sep
	}
	return ""
}
<|endoftext|>


<commit_before>// Copyright 2016 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.
package application

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"os"
	"strings"
	"unicode/utf8"

	"github.com/juju/cmd"
	"github.com/juju/errors"
	"github.com/juju/gnuflag"

	"github.com/juju/juju/api/application"
	"github.com/juju/juju/apiserver/params"
	"github.com/juju/juju/cmd/juju/block"
	"github.com/juju/juju/cmd/modelcmd"
	"github.com/juju/juju/cmd/output"
	"github.com/juju/utils/keyvalues"
)

const (
	configSummary = `Gets, sets, or resets configuration for a deployed application.`
	configDetails = `By default, all configuration (keys, values, metadata) for the application are
displayed if a key is not specified.

Output includes the name of the charm used to deploy the application and a
listing of the application-specific configuration settings.
See ` + "`juju status`" + ` for application names.

Examples:
    juju config <application-name>
	juju config --format=json <application-name>
    juju config <application-name> wait-timeout
	juju config <application-name> key=value key2=value2
	juju config <application-name> --reset key key2
	juju config <application-name> --file path/to/config.yaml

See also:
    deploy
    status
`
	maxValueSize = 5242880 // Max size for a config file.
)

// NewConfigCommand returns a command used to get application attributes.
func NewConfigCommand() cmd.Command {
	return modelcmd.Wrap(&configCommand{})
}

type attributes map[string]string

// configCommand get, sets, and resets configuration values of an application.
type configCommand struct {
	api configCommandAPI
	modelcmd.ModelCommandBase
	out cmd.Output

	action          func(*cmd.Context) error // get, set, or reset action set in  Init
	applicationName string
	configFile      cmd.FileVar
	useFile         bool
	keys            []string
	reset           bool
	values          attributes
}

// configCommandAPI is an interface to allow passing in a fake implementation under test.
type configCommandAPI interface {
	Close() error
	Update(args params.ApplicationUpdate) error
	Get(application string) (*params.ApplicationGetResults, error)
	Set(application string, options map[string]string) error
	Unset(application string, options []string) error
}

// Info is part of the cmd.Command interface.
func (c *configCommand) Info() *cmd.Info {
	return &cmd.Info{
		Name:    "config",
		Args:    "<application name> [<attribute-key>[<=attribute-value>] ...]",
		Purpose: configSummary,
		Doc:     configDetails,
	}
}

// SetFlags is part of the cmd.Command interface.
func (c *configCommand) SetFlags(f *gnuflag.FlagSet) {
	c.ModelCommandBase.SetFlags(f)
	c.out.AddFlags(f, "yaml", output.DefaultFormatters)
	f.Var(&c.configFile, "file", "path to yaml-formatted application config")
	f.BoolVar(&c.reset, "reset", false, "Reset the proveded keys to be empty")
}

// getAPI either uses the fake API set at test time or that is nil, gets a real
// API and sets that as the API.
func (c *configCommand) getAPI() (func(), error) {
	if c.api != nil {
		return func() { c.api.Close() }, nil
	}
	root, err := c.NewAPIRoot()
	if err != nil {
		return nil, errors.Trace(err)
	}
	c.api = application.NewClient(root)
	return func() { c.api.Close() }, nil
}

// Init is part of the cmd.Command interface.
func (c *configCommand) Init(args []string) error {
	if len(args) == 0 || len(strings.Split(args[0], "=")) > 1 {
		return errors.New("no application name specified")
	}
	c.applicationName = args[0]
	if c.reset {
		c.action = c.resetConfig
		return c.parseReset(args[1:])
	}
	if c.configFile.Path != "" {
		return c.parseSet(args[1:], true)
	}
	if len(args[1:]) > 0 && strings.Contains(args[1], "=") {
		return c.parseSet(args[1:], false)
	}
	return c.parseGet(args[1:])
}

// parseReset parses command line args when the --reset flag is supplied.
func (c *configCommand) parseReset(args []string) error {
	if len(args) == 0 {
		return errors.New("no configuration options specified")
	}
	c.action = c.resetConfig
	c.keys = args

	return nil
}

// parseSet parses the command line args when --file is set or if the
// positional args are key=value pairs.
func (c *configCommand) parseSet(args []string, file bool) error {
	if file && len(args) > 0 {
		return errors.New("cannot specify --file and key=value arguments simultaneously")
	}
	c.action = c.setConfig
	if file {
		c.useFile = true
		return nil
	}

	settings, err := keyvalues.Parse(args, true)
	if err != nil {
		return err
	}
	c.values = settings

	return nil
}

// parseGet parses the command line args if we aren't setting or resetting.
func (c *configCommand) parseGet(args []string) error {
	if len(args) > 1 {
		return errors.New("can only retrieve a single value, or all values")
	}
	c.action = c.getConfig
	c.keys = args
	return nil
}

// Run implements the cmd.Command interface.
func (c *configCommand) Run(ctx *cmd.Context) error {
	closer, err := c.getAPI()
	if err != nil {
		return errors.Trace(err)
	}
	defer closer()

	return c.action(ctx)
}

// resetConfig is the run action when we are resetting attributes.
func (c *configCommand) resetConfig(ctx *cmd.Context) error {
	return block.ProcessBlockedError(c.api.Unset(c.applicationName, c.keys), block.BlockChange)
}

// setConfig is the run action when we are setting new attribute values as args
// or as a file passed in.
func (c *configCommand) setConfig(ctx *cmd.Context) error {
	if c.useFile {
		b, err := c.configFile.Read(ctx)
		if err != nil {
			return err
		}
		return block.ProcessBlockedError(
			c.api.Update(
				params.ApplicationUpdate{
					ApplicationName: c.applicationName,
					SettingsYAML:    string(b)}), block.BlockChange)
	}
	settings := map[string]string{}
	for k, v := range c.values {
		//empty string is also valid as a setting value
		if v == "" {
			settings[k] = v
			continue
		}

		if v[0] != '@' {
			if !utf8.ValidString(v) {
				return errors.Errorf("value for option %q contains non-UTF-8 sequences", k)
			}
			settings[k] = v
			continue
		}
		nv, err := readValue(ctx, v[1:])
		if err != nil {
			return err
		}
		if !utf8.ValidString(nv) {
			return errors.Errorf("value for option %q contains non-UTF-8 sequences", k)
		}
		settings[k] = nv
	}

	result, err := c.api.Get(c.applicationName)
	if err != nil {
		return err
	}

	for k, v := range settings {
		configValue := result.Config[k]

		configValueMap, ok := configValue.(map[string]interface{})
		if ok {
			// convert the value to string and compare
			if fmt.Sprintf("%v", configValueMap["value"]) == v {
				logger.Warningf("the configuration setting %q already has the value %q", k, v)
			}
		}
	}

	return block.ProcessBlockedError(c.api.Set(c.applicationName, settings), block.BlockChange)
}

// getConfig is the run action to return one or all configuration values.
func (c *configCommand) getConfig(ctx *cmd.Context) error {
	results, err := c.api.Get(c.applicationName)
	if err != nil {
		return err
	}
	if len(c.keys) == 1 {
		key := c.keys[0]
		info, found := results.Config[key].(map[string]interface{})
		if !found {
			return errors.Errorf("key %q not found in %q application settings.", key, c.applicationName)
		}
		out := &bytes.Buffer{}
		err := cmd.FormatYaml(out, info["value"])
		if err != nil {
			return err
		}
		fmt.Fprint(ctx.Stdout, out.String())
		return nil
	}

	resultsMap := map[string]interface{}{
		"application": results.Application,
		"charm":       results.Charm,
		"settings":    results.Config,
	}
	return c.out.Write(ctx, resultsMap)
}

// readValue reads the value of an option out of the named file.
// An empty content is valid, like in parsing the options. The upper
// size is 5M.
func readValue(ctx *cmd.Context, filename string) (string, error) {
	absFilename := ctx.AbsPath(filename)
	fi, err := os.Stat(absFilename)
	if err != nil {
		return "", errors.Errorf("cannot read option from file %q: %v", filename, err)
	}
	if fi.Size() > maxValueSize {
		return "", errors.Errorf("size of option file is larger than 5M")
	}
	content, err := ioutil.ReadFile(ctx.AbsPath(filename))
	if err != nil {
		return "", errors.Errorf("cannot read option from file %q: %v", filename, err)
	}
	return string(content), nil
}
<commit_msg>Minor optimizations<commit_after>// Copyright 2016 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.
package application

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"os"
	"strings"
	"unicode/utf8"

	"github.com/juju/cmd"
	"github.com/juju/errors"
	"github.com/juju/gnuflag"

	"github.com/juju/juju/api/application"
	"github.com/juju/juju/apiserver/params"
	"github.com/juju/juju/cmd/juju/block"
	"github.com/juju/juju/cmd/modelcmd"
	"github.com/juju/juju/cmd/output"
	"github.com/juju/utils/keyvalues"
)

const (
	configSummary = `Gets, sets, or resets configuration for a deployed application.`
	configDetails = `By default, all configuration (keys, values, metadata) for the application are
displayed if a key is not specified.

Output includes the name of the charm used to deploy the application and a
listing of the application-specific configuration settings.
See ` + "`juju status`" + ` for application names.

Examples:
    juju config <application-name>
	juju config --format=json <application-name>
    juju config <application-name> wait-timeout
	juju config <application-name> key=value key2=value2
	juju config <application-name> --reset key key2
	juju config <application-name> --file path/to/config.yaml

See also:
    deploy
    status
`
	maxValueSize = 5242880 // Max size for a config file.
)

// NewConfigCommand returns a command used to get application attributes.
func NewConfigCommand() cmd.Command {
	return modelcmd.Wrap(&configCommand{})
}

type attributes map[string]string

// configCommand get, sets, and resets configuration values of an application.
type configCommand struct {
	api configCommandAPI
	modelcmd.ModelCommandBase
	out cmd.Output

	action          func(*cmd.Context) error // get, set, or reset action set in  Init
	applicationName string
	configFile      cmd.FileVar
	keys            []string
	reset           bool
	useFile         bool
	values          attributes
}

// configCommandAPI is an interface to allow passing in a fake implementation under test.
type configCommandAPI interface {
	Close() error
	Update(args params.ApplicationUpdate) error
	Get(application string) (*params.ApplicationGetResults, error)
	Set(application string, options map[string]string) error
	Unset(application string, options []string) error
}

// Info is part of the cmd.Command interface.
func (c *configCommand) Info() *cmd.Info {
	return &cmd.Info{
		Name:    "config",
		Args:    "<application name> [<attribute-key>[<=attribute-value>] ...]",
		Purpose: configSummary,
		Doc:     configDetails,
	}
}

// SetFlags is part of the cmd.Command interface.
func (c *configCommand) SetFlags(f *gnuflag.FlagSet) {
	c.ModelCommandBase.SetFlags(f)
	c.out.AddFlags(f, "yaml", output.DefaultFormatters)
	f.Var(&c.configFile, "file", "path to yaml-formatted application config")
	f.BoolVar(&c.reset, "reset", false, "Reset the proveded keys to be empty")
}

// getAPI either uses the fake API set at test time or that is nil, gets a real
// API and sets that as the API.
func (c *configCommand) getAPI() (func(), error) {
	if c.api != nil {
		return func() { c.api.Close() }, nil
	}
	root, err := c.NewAPIRoot()
	if err != nil {
		return nil, errors.Trace(err)
	}
	c.api = application.NewClient(root)
	return func() { c.api.Close() }, nil
}

// Init is part of the cmd.Command interface.
func (c *configCommand) Init(args []string) error {
	if len(args) == 0 || len(strings.Split(args[0], "=")) > 1 {
		return errors.New("no application name specified")
	}
	c.applicationName = args[0]
	if c.reset {
		c.action = c.resetConfig
		return c.parseReset(args[1:])
	}
	if c.configFile.Path != "" {
		return c.parseSet(args[1:], true)
	}
	if len(args[1:]) > 0 && strings.Contains(args[1], "=") {
		return c.parseSet(args[1:], false)
	}
	return c.parseGet(args[1:])
}

// parseReset parses command line args when the --reset flag is supplied.
func (c *configCommand) parseReset(args []string) error {
	if len(args) == 0 {
		return errors.New("no configuration options specified")
	}
	c.action = c.resetConfig
	c.keys = args

	return nil
}

// parseSet parses the command line args when --file is set or if the
// positional args are key=value pairs.
func (c *configCommand) parseSet(args []string, file bool) error {
	if file && len(args) > 0 {
		return errors.New("cannot specify --file and key=value arguments simultaneously")
	}
	c.action = c.setConfig
	if file {
		c.useFile = true
		return nil
	}

	settings, err := keyvalues.Parse(args, true)
	if err != nil {
		return err
	}
	c.values = settings

	return nil
}

// parseGet parses the command line args if we aren't setting or resetting.
func (c *configCommand) parseGet(args []string) error {
	if len(args) > 1 {
		return errors.New("can only retrieve a single value, or all values")
	}
	c.action = c.getConfig
	c.keys = args
	return nil
}

// Run implements the cmd.Command interface.
func (c *configCommand) Run(ctx *cmd.Context) error {
	closer, err := c.getAPI()
	if err != nil {
		return errors.Trace(err)
	}
	defer closer()

	return c.action(ctx)
}

// resetConfig is the run action when we are resetting attributes.
func (c *configCommand) resetConfig(ctx *cmd.Context) error {
	return block.ProcessBlockedError(c.api.Unset(c.applicationName, c.keys), block.BlockChange)
}

// setConfigFromFile sets the application configuration from settings passed
// in a YAML file.
func (c *configCommand) setConfigFromFile(ctx *cmd.Context) error {
	b, err := c.configFile.Read(ctx)
	if err != nil {
		return err
	}
	return block.ProcessBlockedError(
		c.api.Update(
			params.ApplicationUpdate{
				ApplicationName: c.applicationName,
				SettingsYAML:    string(b)}), block.BlockChange)

}

// validateValues reads the values provided as args and validates that they are
// valid UTF-8.
func (c *configCommand) validateValues(ctx *cmd.Context) (map[string]string, error) {
	settings := map[string]string{}
	for k, v := range c.values {
		//empty string is also valid as a setting value
		if v == "" {
			settings[k] = v
			continue
		}

		if v[0] != '@' {
			if !utf8.ValidString(v) {
				return nil, errors.Errorf("value for option %q contains non-UTF-8 sequences", k)
			}
			settings[k] = v
			continue
		}
		nv, err := readValue(ctx, v[1:])
		if err != nil {
			return nil, err
		}
		if !utf8.ValidString(nv) {
			return nil, errors.Errorf("value for option %q contains non-UTF-8 sequences", k)
		}
		settings[k] = nv
	}
	return settings, nil

}

// setConfig is the run action when we are setting new attribute values as args
// or as a file passed in.
func (c *configCommand) setConfig(ctx *cmd.Context) error {
	if c.useFile {
		return c.setConfigFromFile(ctx)
	}

	settings, err := c.validateValues(ctx)
	if err != nil {
		return errors.Trace(err)
	}

	result, err := c.api.Get(c.applicationName)
	if err != nil {
		return err
	}

	for k, v := range settings {
		configValue := result.Config[k]

		configValueMap, ok := configValue.(map[string]interface{})
		if ok {
			// convert the value to string and compare
			if fmt.Sprintf("%v", configValueMap["value"]) == v {
				logger.Warningf("the configuration setting %q already has the value %q", k, v)
			}
		}
	}

	return block.ProcessBlockedError(c.api.Set(c.applicationName, settings), block.BlockChange)
}

// getConfig is the run action to return one or all configuration values.
func (c *configCommand) getConfig(ctx *cmd.Context) error {
	results, err := c.api.Get(c.applicationName)
	if err != nil {
		return err
	}
	if len(c.keys) == 1 {
		key := c.keys[0]
		info, found := results.Config[key].(map[string]interface{})
		if !found {
			return errors.Errorf("key %q not found in %q application settings.", key, c.applicationName)
		}
		out := &bytes.Buffer{}
		err := cmd.FormatYaml(out, info["value"])
		if err != nil {
			return err
		}
		fmt.Fprint(ctx.Stdout, out.String())
		return nil
	}

	resultsMap := map[string]interface{}{
		"application": results.Application,
		"charm":       results.Charm,
		"settings":    results.Config,
	}
	return c.out.Write(ctx, resultsMap)
}

// readValue reads the value of an option out of the named file.
// An empty content is valid, like in parsing the options. The upper
// size is 5M.
func readValue(ctx *cmd.Context, filename string) (string, error) {
	absFilename := ctx.AbsPath(filename)
	fi, err := os.Stat(absFilename)
	if err != nil {
		return "", errors.Errorf("cannot read option from file %q: %v", filename, err)
	}
	if fi.Size() > maxValueSize {
		return "", errors.Errorf("size of option file is larger than 5M")
	}
	content, err := ioutil.ReadFile(ctx.AbsPath(filename))
	if err != nil {
		return "", errors.Errorf("cannot read option from file %q: %v", filename, err)
	}
	return string(content), nil
}
<|endoftext|>