<commit_before>// A simple logging module that mimics the behavior of Python's logging module.
//
// All it does basically is wrap Go's logger with nice multi-level logging calls, and
// allows you to set the logging level of your app in runtime.
//
// Logging is done just like calling fmt.Sprintf:
// 		logging.Info("This object is %s and that is %s", obj, that)
//
// example output:
//	2013/05/07 01:20:26 INFO @ db.go:528: Registering plugin REPLICATION
//	2013/05/07 01:20:26 INFO @ db.go:562: Registered 6 plugins and 22 commands
//	2013/05/07 01:20:26 INFO @ slave.go:277: Running replication watchdog loop!
//	2013/05/07 01:20:26 INFO @ redis.go:49: Redis adapter listening on 0.0.0.0:2000
//	2013/05/07 01:20:26 WARN @ main.go:69: Starting adapter...
//	2013/05/07 01:20:26 INFO @ db.go:966: Finished dump load. Loaded 2 objects from dump
//	2013/05/07 01:22:26 INFO @ db.go:329: Checking persistence... 0 changes since 2m0.000297531s
//	2013/05/07 01:22:26 INFO @ db.go:337: No need to save the db. no changes...
//	2013/05/07 01:22:26 DEBUG @ db.go:341: Sleeping for 2m0s
//
package logging

import (
	"log"
	"fmt"
	"io"
	"runtime/debug"
	"path"
	"runtime"
)

const (
	DEBUG = 1
	INFO = 2
	WARNING = 4
	WARN = 4
	ERROR = 8
	CRITICAL  = 16
	QUIET = ERROR | CRITICAL  //setting for errors only
	NORMAL = INFO | WARN | ERROR | CRITICAL // default setting - all besides debug
	ALL = 255
	NOTHING = 0
)

//default logging level is ALL
var level int = ALL

// Set the logging level.
//
// Contrary to Python that specifies a minimal level, this logger is set with a bit mask
// of active levels.
//
// e.g. for INFO and ERROR use:
// 		SetLevel(logging.INFO | logging.ERROR)
//
// For everything but debug and info use:
// 		SetLevel(logging.ALL &^ (logging.INFO | logging.DEBUG))
//
func SetLevel(l int) {
	level = l
}

// Set the output writer. for now it just wraps log.SetOutput()
func SetOutPut(w io.Writer) {
	log.SetOutput(w)
}

//get the stack (line + file) context to return the caller to the log
func getContext() (file string, line int) {

	_, file, line, _ = runtime.Caller(3)
	file = path.Base(file)

	return
}

//Output debug logging messages
func Debug(msg string, args ...interface{}) {
	if level & DEBUG != 0 {
		writeMessage("DEBUG", msg, args...)
	}
}

//format the message
func writeMessage(level string, msg string, args ...interface {} ) {
	f, l := getContext()
	log.Printf(fmt.Sprintf("%s @ %s:%d: %s", level, f, l, msg), args...)
}

//output INFO level messages
func Info(msg string, args ...interface{}) {

	if level & INFO != 0 {

		writeMessage("INFO", msg, args...)

	}
}

//output WARNING level messages
func Warning(msg string, args ...interface{}) {
	if level & WARN != 0 {
		writeMessage("WARNING", msg, args...)
	}
}

//output ERROR level messages
func Error(msg string, args ...interface{}) {
	if level & ERROR != 0 {
		writeMessage("ERROR", msg, args...)
	}
}

//Output a CRITICAL level message while showing a stack trace
func Critical(msg string, args ...interface{}) {
	if level & CRITICAL != 0 {
		writeMessage("CRITICAL", msg, args...)
		log.Println(string(debug.Stack()))
	}
}

// Raise a PANIC while writing the stack trace to the log
func Panic(msg string, args ...interface{}) {
	log.Println(string(debug.Stack()))
	log.Panicf(msg, args...)

}

<commit_msg>Fixed capitazlization of SetOutput()<commit_after>// A simple logging module that mimics the behavior of Python's logging module.
//
// All it does basically is wrap Go's logger with nice multi-level logging calls, and
// allows you to set the logging level of your app in runtime.
//
// Logging is done just like calling fmt.Sprintf:
// 		logging.Info("This object is %s and that is %s", obj, that)
//
// example output:
//	2013/05/07 01:20:26 INFO @ db.go:528: Registering plugin REPLICATION
//	2013/05/07 01:20:26 INFO @ db.go:562: Registered 6 plugins and 22 commands
//	2013/05/07 01:20:26 INFO @ slave.go:277: Running replication watchdog loop!
//	2013/05/07 01:20:26 INFO @ redis.go:49: Redis adapter listening on 0.0.0.0:2000
//	2013/05/07 01:20:26 WARN @ main.go:69: Starting adapter...
//	2013/05/07 01:20:26 INFO @ db.go:966: Finished dump load. Loaded 2 objects from dump
//	2013/05/07 01:22:26 INFO @ db.go:329: Checking persistence... 0 changes since 2m0.000297531s
//	2013/05/07 01:22:26 INFO @ db.go:337: No need to save the db. no changes...
//	2013/05/07 01:22:26 DEBUG @ db.go:341: Sleeping for 2m0s
//
package logging

import (
	"log"
	"fmt"
	"io"
	"runtime/debug"
	"path"
	"runtime"
)

const (
	DEBUG = 1
	INFO = 2
	WARNING = 4
	WARN = 4
	ERROR = 8
	CRITICAL  = 16
	QUIET = ERROR | CRITICAL  //setting for errors only
	NORMAL = INFO | WARN | ERROR | CRITICAL // default setting - all besides debug
	ALL = 255
	NOTHING = 0
)

//default logging level is ALL
var level int = ALL

// Set the logging level.
//
// Contrary to Python that specifies a minimal level, this logger is set with a bit mask
// of active levels.
//
// e.g. for INFO and ERROR use:
// 		SetLevel(logging.INFO | logging.ERROR)
//
// For everything but debug and info use:
// 		SetLevel(logging.ALL &^ (logging.INFO | logging.DEBUG))
//
func SetLevel(l int) {
	level = l
}

// Set the output writer. for now it just wraps log.SetOutput()
func SetOutput(w io.Writer) {
	log.SetOutput(w)
}

//get the stack (line + file) context to return the caller to the log
func getContext() (file string, line int) {

	_, file, line, _ = runtime.Caller(3)
	file = path.Base(file)

	return
}

//Output debug logging messages
func Debug(msg string, args ...interface{}) {
	if level & DEBUG != 0 {
		writeMessage("DEBUG", msg, args...)
	}
}

//format the message
func writeMessage(level string, msg string, args ...interface {} ) {
	f, l := getContext()
	log.Printf(fmt.Sprintf("%s @ %s:%d: %s", level, f, l, msg), args...)
}

//output INFO level messages
func Info(msg string, args ...interface{}) {

	if level & INFO != 0 {

		writeMessage("INFO", msg, args...)

	}
}

//output WARNING level messages
func Warning(msg string, args ...interface{}) {
	if level & WARN != 0 {
		writeMessage("WARNING", msg, args...)
	}
}

//output ERROR level messages
func Error(msg string, args ...interface{}) {
	if level & ERROR != 0 {
		writeMessage("ERROR", msg, args...)
	}
}

//Output a CRITICAL level message while showing a stack trace
func Critical(msg string, args ...interface{}) {
	if level & CRITICAL != 0 {
		writeMessage("CRITICAL", msg, args...)
		log.Println(string(debug.Stack()))
	}
}

// Raise a PANIC while writing the stack trace to the log
func Panic(msg string, args ...interface{}) {
	log.Println(string(debug.Stack()))
	log.Panicf(msg, args...)

}

<|endoftext|>


<commit_before>package logging

import (
	"context"
	"log"
	"time"
)

type LogEntry struct {
	level        string
	dateOccurred time.Time
	application  string
	message      string
}

func Debug(application string, message string) LogEntry {
	return LogEntry{
		level:        "debug",
		dateOccurred: time.Now().UTC(),
		message:      message,
	}
}

func Error(application string, message string) LogEntry {
	return LogEntry{
		level:        "error",
		dateOccurred: time.Now().UTC(),
		message:      message,
	}
}

func Warn(application string, message string) LogEntry {
	return LogEntry{
		level:        "warn",
		dateOccurred: time.Now().UTC(),
		message:      message,
	}
}

func Info(application string, message string) LogEntry {
	return LogEntry{
		level:        "info",
		dateOccurred: time.Now().UTC(),
		message:      message,
	}
}

func NewLogEntry(context context.Context) {

	logEntry := Debug("TestApplication v0.1", "Some random message")
	log.Println(logEntry)
}
<commit_msg>added more example log entries<commit_after>package logging

import (
	"context"
	"log"
	"time"
)

type LogEntry struct {
	level        string
	dateOccurred time.Time
	application  string
	message      string
}

func Debug(application string, message string) LogEntry {
	return LogEntry{
		level:        "debug",
		dateOccurred: time.Now().UTC(),
		message:      message,
	}
}

func Error(application string, message string) LogEntry {
	return LogEntry{
		level:        "error",
		dateOccurred: time.Now().UTC(),
		message:      message,
	}
}

func Warn(application string, message string) LogEntry {
	return LogEntry{
		level:        "warn",
		dateOccurred: time.Now().UTC(),
		message:      message,
	}
}

func Info(application string, message string) LogEntry {
	return LogEntry{
		level:        "info",
		dateOccurred: time.Now().UTC(),
		message:      message,
	}
}

func NewLogEntry(context context.Context) {

	logEntry := Debug("TestApplication v0.1", "Some random message")
	log.Println(logEntry)

	logEntry = Warn("TestApplication v0.1", "Warning, getting hot!")
	log.Println(logEntry)

	logEntry = Error("TestApplication v0.1", "App has blown up!")
	log.Println(logEntry)
}
<|endoftext|>


<commit_before>package movavg_test

import (
	"fmt"

	. "github.com/mxmCherry/movavg"
)

func ExampleSMA() {

	sma := ThreadSafe(NewSMA(3))

	fmt.Println(sma.Avg())

	fmt.Println(sma.Add(1))
	fmt.Println(sma.Avg())

	fmt.Println(sma.Add(2))
	fmt.Println(sma.Avg())

	fmt.Println(sma.Add(3))
	fmt.Println(sma.Avg())

	fmt.Println(sma.Add(4))
	fmt.Println(sma.Avg())

	fmt.Println(sma.Add(5))
	fmt.Println(sma.Avg())

	// Output:
	// 0
	// 1
	// 1
	// 1.5
	// 1.5
	// 2
	// 2
	// 3
	// 3
	// 4
	// 4
}
<commit_msg>Minor comments improvs<commit_after>package movavg_test

import (
	"fmt"

	. "github.com/mxmCherry/movavg"
)

func ExampleSMA() {

	sma := ThreadSafe(NewSMA(3))

	fmt.Println(sma.Avg()) // 0

	fmt.Println(sma.Add(1)) // 1
	fmt.Println(sma.Avg())  // 1

	fmt.Println(sma.Add(2)) // 1.5
	fmt.Println(sma.Avg())  // 1.5

	fmt.Println(sma.Add(3)) // 2
	fmt.Println(sma.Avg())  // 2

	fmt.Println(sma.Add(4)) // 3
	fmt.Println(sma.Avg())  // 3

	fmt.Println(sma.Add(5)) // 4
	fmt.Println(sma.Avg())  // 4

	// Output:
	// 0
	// 1
	// 1
	// 1.5
	// 1.5
	// 2
	// 2
	// 3
	// 3
	// 4
	// 4
}
<|endoftext|>


<commit_before>package cmd

const Version = "3.13.1"
<commit_msg>version bump<commit_after>package cmd

const Version = "3.13.2"
<|endoftext|>


<commit_before>package cmd

// Version defines the current Pop version.
const Version = "v4.6.0"
<commit_msg>Reset version to development<commit_after>package cmd

// Version defines the current Pop version.
const Version = "development"
<|endoftext|>


<commit_before>// +build !appengine

package logutil

import (
	"bytes"
	"fmt"
	"io"
	"runtime"
	"unicode/utf8"
)

const ()

func sprintf(format string, args ...interface{}) string {
	_, file, line, _ := runtime.Caller(2)
	var fileLine = fmt.Sprintf("%s:%d]", file, line)
	return fmt.Sprintf("%s %s", fileLine, fmt.Sprintf(format, args...))
}

func Debugf(c appengine.Context, format string, args ...interface{}) {
	c.Debugf(sprintf(format, args...))
}

func Infof(c appengine.Context, format string, args ...interface{}) {
	c.Infof(sprintf(format, args...))
}

func Warningf(c appengine.Context, format string, args ...interface{}) {
	c.Warningf(sprintf(format, args...))
}

func Errorf(c appengine.Context, format string, args ...interface{}) {
	c.Errorf(sprintf(format, args...))
}

func Criticalf(c appengine.Context, format string, args ...interface{}) {
	c.Criticalf(sprintf(format, args...))
}

func ErrorStackTracef(c appengine.Context, format string, args ...interface{}) {
	c.Errorf(sprintf("%s", StackTrace(2, 2048)))
	c.Errorf(sprintf(format, args...))
}

//https://github.com/knightso/base/blob/master/errors/errors.go
func StackTrace(skip, maxBytes int) []byte {
	// this func is debug purpose and ignores errors

	buf := make([]byte, maxBytes)
	n := runtime.Stack(buf, false)
	var gotall bool
	if n < len(buf) {
		buf = buf[:n]
		gotall = true
	} else {
		for !utf8.Valid(buf) || len(buf) == 0 {
			buf = buf[:len(buf)-1]
		}
	}

	var w bytes.Buffer

	writeOrSkip := func(buf []byte, w io.Writer, line int) {
		if line == 1 || line > 1+skip*2 {
			w.Write(buf)
		}
	}

	line := 1
	for { //EOFまでloopする
		lf := bytes.IndexByte(buf, '\n')
		if lf < 0 {
			writeOrSkip(buf, &w, line)
			break
		}
		writeOrSkip(buf[:lf+1], &w, line)
		buf = buf[lf+1:]
		line++
	}

	if !gotall {
		w.WriteString("\n        ... (omitted)")
	}
	w.WriteString("\n")

	return w.Bytes()
}
<commit_msg>minimum log level<commit_after>// +build !appengine

package logutil

import (
	"bytes"
	"fmt"
	"io"
	"runtime"
	"unicode/utf8"
)

const ()

func sprintf(format string, args ...interface{}) string {
	_, file, line, _ := runtime.Caller(2)
	var fileLine = fmt.Sprintf("%s:%d]", file, line)
	return fmt.Sprintf("%s %s", fileLine, fmt.Sprintf(format, args...))
}

func Debugf(c appengine.Context, format string, args ...interface{}) {
	c.Debugf(sprintf(format, args...))
}

func Infof(c appengine.Context, format string, args ...interface{}) {
	c.Infof(sprintf(format, args...))
}

func ErrorStackTracef(c appengine.Context, format string, args ...interface{}) {
	c.Errorf(sprintf("%s", StackTrace(2, 2048)))
	c.Errorf(sprintf(format, args...))
}

//https://github.com/knightso/base/blob/master/errors/errors.go
func StackTrace(skip, maxBytes int) []byte {
	// this func is debug purpose and ignores errors

	buf := make([]byte, maxBytes)
	n := runtime.Stack(buf, false)
	var gotall bool
	if n < len(buf) {
		buf = buf[:n]
		gotall = true
	} else {
		for !utf8.Valid(buf) || len(buf) == 0 {
			buf = buf[:len(buf)-1]
		}
	}

	var w bytes.Buffer

	writeOrSkip := func(buf []byte, w io.Writer, line int) {
		if line == 1 || line > 1+skip*2 {
			w.Write(buf)
		}
	}

	line := 1
	for { //EOFまでloopする
		lf := bytes.IndexByte(buf, '\n')
		if lf < 0 {
			writeOrSkip(buf, &w, line)
			break
		}
		writeOrSkip(buf[:lf+1], &w, line)
		buf = buf[lf+1:]
		line++
	}

	if !gotall {
		w.WriteString("\n        ... (omitted)")
	}
	w.WriteString("\n")

	return w.Bytes()
}
<|endoftext|>


<commit_before><commit_msg>Solved the lower problem<commit_after>package main

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"os"
	"strings"
)

func readLines(path string) (lines []string, err error) {
	var (
		file   *os.File
		part   []byte
		prefix bool
	)
	if file, err = os.Open(path); err != nil {
		return
	}
	defer file.Close()

	reader := bufio.NewReader(file)
	buffer := bytes.NewBuffer(make([]byte, 0))
	for {
		if part, prefix, err = reader.ReadLine(); err != nil {
			break
		}
		buffer.Write(part)
		if !prefix {
			lines = append(lines, buffer.String())
			buffer.Reset()
		}
	}
	if err == io.EOF {
		err = nil
	}
	return
}

func main() {
	lines, err := readLines(os.Args[1])
	if err != nil {
		fmt.Println("Error: %s\n", err)
		return
	}
	for _, line := range lines {
		fmt.Print(strings.ToLower(line))
		fmt.Println()
	}

}
<|endoftext|>


<commit_before>package models

import (
	"errors"
	"time"

	"github.com/koding/bongo"
)

type MessageReply struct {
	// unique identifier of the MessageReply
	Id int64 `json:"id"`

	// Id of the interacted message
	MessageId int64 `json:"messageId"             sql:"NOT NULL"`

	// Id of the reply
	ReplyId int64 `json:"replyId"                 sql:"NOT NULL"`

	// Creation of the MessageReply
	CreatedAt time.Time `json:"createdAt"         sql:"NOT NULL"`
}

func (m MessageReply) GetId() int64 {
	return m.Id
}

func (m MessageReply) TableName() string {
	return "api.message_reply"
}

func NewMessageReply() *MessageReply {
	return &MessageReply{}
}

func (m *MessageReply) AfterCreate() {
	bongo.B.AfterCreate(m)
}

func (m *MessageReply) AfterUpdate() {
	bongo.B.AfterUpdate(m)
}

func (m MessageReply) AfterDelete() {
	bongo.B.AfterDelete(m)
}

func (m *MessageReply) ById(id int64) error {
	return bongo.B.ById(m, id)
}

func (m *MessageReply) Create() error {
	return bongo.B.Create(m)
}

func (m *MessageReply) Some(data interface{}, q *bongo.Query) error {
	return bongo.B.Some(m, data, q)
}

func (m *MessageReply) One(q *bongo.Query) error {
	return bongo.B.One(m, m, q)
}

func (m *MessageReply) Delete() error {
	selector := map[string]interface{}{
		"message_id": m.MessageId,
		"reply_id":   m.ReplyId,
	}

	if err := m.One(bongo.NewQS(selector)); err != nil {
		return err
	}

	err := bongo.B.Delete(m)
	if err != nil {
		return err
	}

	return nil
}

func (m *MessageReply) DeleteByOrQuery(messageId int64) error {
	var messageReplies []MessageReply
	query := bongo.B.DB.Table(m.TableName())
	query = query.Where("message_id = ? or reply_id = ?", messageId, messageId)

	if err := query.Find(&messageReplies).Error; err != nil {
		return err
	}

	if messageReplies == nil {
		return nil
	}

	if len(messageReplies) == 0 {
		return nil
	}

	for _, messageReply := range messageReplies {
		err := bongo.B.Delete(messageReply)
		if err != nil {
			return err
		}
	}

	return nil
}

func (m *MessageReply) List(query *Query) ([]ChannelMessage, error) {
	return m.fetchMessages(query)
}

func (m *MessageReply) ListAll() ([]ChannelMessage, error) {
	query := NewQuery()
	query.Limit = 0
	query.Skip = 0
	return m.fetchMessages(query)
}

func (m *MessageReply) fetchMessages(query *Query) ([]ChannelMessage, error) {
	var replies []int64

	if m.MessageId == 0 {
		return nil, errors.New("MessageId is not set")
	}

	q := &bongo.Query{
		Selector: map[string]interface{}{
			"message_id": m.MessageId,
		},
		Pluck:      "reply_id",
		Pagination: *bongo.NewPagination(query.Limit, query.Skip),
		Sort:       map[string]string{"created_at": "DESC"},
	}

	bongoQuery := bongo.B.BuildQuery(m, q)
	if !query.From.IsZero() {
		bongoQuery = bongoQuery.Where("created_at < ?", query.From)
	}

	bongoQuery = bongoQuery.Pluck(q.Pluck, &replies)
	if err := bongoQuery.Error; err != nil {
		return nil, err
	}

	parent := NewChannelMessage()
	channelMessageReplies, err := parent.FetchByIds(replies)
	if err != nil {
		return nil, err
	}

	return channelMessageReplies, nil
}

func (m *MessageReply) Count() (int, error) {
	if m.MessageId == 0 {
		return 0, errors.New("MessageId is not set")
	}

	return bongo.B.Count(m,
		"message_id = ?",
		m.MessageId,
	)
}
<commit_msg>Social: Parent ChannelMessage is fetched via reply message id<commit_after>package models

import (
	"errors"
	"time"

	"github.com/koding/bongo"
)

type MessageReply struct {
	// unique identifier of the MessageReply
	Id int64 `json:"id"`

	// Id of the interacted message
	MessageId int64 `json:"messageId"             sql:"NOT NULL"`

	// Id of the reply
	ReplyId int64 `json:"replyId"                 sql:"NOT NULL"`

	// Creation of the MessageReply
	CreatedAt time.Time `json:"createdAt"         sql:"NOT NULL"`
}

func (m MessageReply) GetId() int64 {
	return m.Id
}

func (m MessageReply) TableName() string {
	return "api.message_reply"
}

func NewMessageReply() *MessageReply {
	return &MessageReply{}
}

func (m *MessageReply) AfterCreate() {
	bongo.B.AfterCreate(m)
}

func (m *MessageReply) AfterUpdate() {
	bongo.B.AfterUpdate(m)
}

func (m MessageReply) AfterDelete() {
	bongo.B.AfterDelete(m)
}

func (m *MessageReply) ById(id int64) error {
	return bongo.B.ById(m, id)
}

func (m *MessageReply) Create() error {
	return bongo.B.Create(m)
}

func (m *MessageReply) Some(data interface{}, q *bongo.Query) error {
	return bongo.B.Some(m, data, q)
}

func (m *MessageReply) One(q *bongo.Query) error {
	return bongo.B.One(m, m, q)
}

func (m *MessageReply) Delete() error {
	selector := map[string]interface{}{
		"message_id": m.MessageId,
		"reply_id":   m.ReplyId,
	}

	if err := m.One(bongo.NewQS(selector)); err != nil {
		return err
	}

	err := bongo.B.Delete(m)
	if err != nil {
		return err
	}

	return nil
}

func (m *MessageReply) DeleteByOrQuery(messageId int64) error {
	var messageReplies []MessageReply
	query := bongo.B.DB.Table(m.TableName())
	query = query.Where("message_id = ? or reply_id = ?", messageId, messageId)

	if err := query.Find(&messageReplies).Error; err != nil {
		return err
	}

	if messageReplies == nil {
		return nil
	}

	if len(messageReplies) == 0 {
		return nil
	}

	for _, messageReply := range messageReplies {
		err := bongo.B.Delete(messageReply)
		if err != nil {
			return err
		}
	}

	return nil
}

func (m *MessageReply) List(query *Query) ([]ChannelMessage, error) {
	return m.fetchMessages(query)
}

func (m *MessageReply) ListAll() ([]ChannelMessage, error) {
	query := NewQuery()
	query.Limit = 0
	query.Skip = 0
	return m.fetchMessages(query)
}

func (m *MessageReply) fetchMessages(query *Query) ([]ChannelMessage, error) {
	var replies []int64

	if m.MessageId == 0 {
		return nil, errors.New("MessageId is not set")
	}

	q := &bongo.Query{
		Selector: map[string]interface{}{
			"message_id": m.MessageId,
		},
		Pluck:      "reply_id",
		Pagination: *bongo.NewPagination(query.Limit, query.Skip),
		Sort:       map[string]string{"created_at": "DESC"},
	}

	bongoQuery := bongo.B.BuildQuery(m, q)
	if !query.From.IsZero() {
		bongoQuery = bongoQuery.Where("created_at < ?", query.From)
	}

	bongoQuery = bongoQuery.Pluck(q.Pluck, &replies)
	if err := bongoQuery.Error; err != nil {
		return nil, err
	}

	parent := NewChannelMessage()
	channelMessageReplies, err := parent.FetchByIds(replies)
	if err != nil {
		return nil, err
	}

	return channelMessageReplies, nil
}

func (m *MessageReply) Count() (int, error) {
	if m.MessageId == 0 {
		return 0, errors.New("MessageId is not set")
	}

	return bongo.B.Count(m,
		"message_id = ?",
		m.MessageId,
	)
}

func (m *MessageReply) FetchRepliedMessage() (*ChannelMessage, error) {
	if m.ReplyId == 0 {
		return nil, errors.New("ReplyId is not set")
	}

	q := &bongo.Query{
		Selector: map[string]interface{}{
			"reply_id": m.ReplyId,
		},
	}

	if err := m.One(q); err != nil {
		return nil, err
	}

	parent := NewChannelMessage()
	if err := parent.ById(m.MessageId); err != nil {
		return nil, err
	}

	return parent, nil
}
<|endoftext|>


<commit_before>package finalize

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

export APP_ROOT=$HOME
export LD_LIBRARY_PATH=$APP_ROOT/nginx/lib:$LD_LIBRARY_PATH

mv $APP_ROOT/nginx/conf/nginx.conf $APP_ROOT/nginx/conf/nginx.conf.erb
erb $APP_ROOT/nginx/conf/nginx.conf.erb > $APP_ROOT/nginx/conf/nginx.conf

if [[ ! -f $APP_ROOT/nginx/logs/access.log ]]; then
    mkfifo $APP_ROOT/nginx/logs/access.log
fi

if [[ ! -f $APP_ROOT/nginx/logs/error.log ]]; then
    mkfifo $APP_ROOT/nginx/logs/error.log
fi
`

	startLoggingScript = `
cat < $APP_ROOT/nginx/logs/access.log &
(>&2 cat) < $APP_ROOT/nginx/logs/error.log &
`

	startCommand = `#!/bin/sh
set -ex
$APP_ROOT/start_logging.sh
nginx -p $APP_ROOT/nginx -c $APP_ROOT/nginx/conf/nginx.conf
`

	nginxConfTemplate = `
worker_processes 1;
daemon off;

error_log <%= ENV["APP_ROOT"] %>/nginx/logs/error.log;
events { worker_connections 1024; }

http {
  charset utf-8;
  log_format cloudfoundry '$http_x_forwarded_for - $http_referer - [$time_local] "$request" $status $body_bytes_sent';
  access_log <%= ENV["APP_ROOT"] %>/nginx/logs/access.log cloudfoundry;
  default_type application/octet-stream;
  include mime.types;
  sendfile on;

  gzip on;
  gzip_disable "msie6";
  gzip_comp_level 6;
  gzip_min_length 1100;
  gzip_buffers 16 8k;
  gzip_proxied any;
  gunzip on;
  gzip_static always;
  gzip_types text/plain text/css text/js text/xml text/javascript application/javascript application/x-javascript application/json application/xml application/xml+rss;
  gzip_vary on;

  tcp_nopush on;
  keepalive_timeout 30;
  port_in_redirect off; # Ensure that redirects don't include the internal container PORT - <%= ENV["PORT"] %>
  server_tokens off;

  server {
    listen <%= ENV["PORT"] %>;
    server_name localhost;

    location / {
      root <%= ENV["APP_ROOT"] %>/public;

      {{if .PushState}}
        if (!-e $request_filename) {
          rewrite ^(.*)$ / break;
        }
      {{end}}

      index index.html index.htm Default.htm;

      {{if .DirectoryIndex}}
        autoindex on;
      {{end}}

      {{if .BasicAuth}}
        auth_basic "Restricted";  #For Basic Auth
        auth_basic_user_file <%= ENV["APP_ROOT"] %>/nginx/conf/.htpasswd;
      {{end}}

      {{if .ForceHTTPS}}
        if ($http_x_forwarded_proto != "https") {
          return 301 https://$host$request_uri;
        }
      {{else}}
      <% if ENV["FORCE_HTTPS"] %>
        if ($http_x_forwarded_proto != "https") {
          return 301 https://$host$request_uri;
        }
      <% end %>
      {{end}}

      {{if .SSI}}
        ssi on;
      {{end}}

      {{if .HSTS}}
        add_header Strict-Transport-Security "max-age=31536000{{if .HSTSIncludeSubDomains}}; includeSubDomains{{end}}{{if .HSTSPreload}}; preload{{end}}";
      {{end}}

      {{if ne .LocationInclude ""}}
        include {{.LocationInclude}};
      {{end}}
    }

  {{if not .HostDotFiles}}
    location ~ /\. {
      deny all;
      return 404;
    }
  {{end}}
  }
}
`
	MimeTypes = `
types {
  text/html html htm shtml;
  text/css css;
  text/xml xml;
  image/gif gif;
  image/jpeg jpeg jpg;
  application/x-javascript js;
  application/atom+xml atom;
  application/rss+xml rss;
  font/ttf ttf;
  font/woff woff;
  font/woff2 woff2;
  text/mathml mml;
  text/plain txt;
  text/vnd.sun.j2me.app-descriptor jad;
  text/vnd.wap.wml wml;
  text/x-component htc;
  text/cache-manifest manifest;
  image/png png;
  image/tiff tif tiff;
  image/vnd.wap.wbmp wbmp;
  image/x-icon ico;
  image/x-jng jng;
  image/x-ms-bmp bmp;
  image/svg+xml svg svgz;
  image/webp webp;
  application/java-archive jar war ear;
  application/mac-binhex40 hqx;
  application/msword doc;
  application/pdf pdf;
  application/postscript ps eps ai;
  application/rtf rtf;
  application/vnd.ms-excel xls;
  application/vnd.ms-powerpoint ppt;
  application/vnd.wap.wmlc wmlc;
  application/vnd.google-earth.kml+xml  kml;
  application/vnd.google-earth.kmz kmz;
  application/x-7z-compressed 7z;
  application/x-cocoa cco;
  application/x-java-archive-diff jardiff;
  application/x-java-jnlp-file jnlp;
  application/x-makeself run;
  application/x-perl pl pm;
  application/x-pilot prc pdb;
  application/x-rar-compressed rar;
  application/x-redhat-package-manager  rpm;
  application/x-sea sea;
  application/x-shockwave-flash swf;
  application/x-stuffit sit;
  application/x-tcl tcl tk;
  application/x-x509-ca-cert der pem crt;
  application/x-xpinstall xpi;
  application/xhtml+xml xhtml;
  application/zip zip;
  application/octet-stream bin exe dll;
  application/octet-stream deb;
  application/octet-stream dmg;
  application/octet-stream eot;
  application/octet-stream iso img;
  application/octet-stream msi msp msm;
  application/json json;
  audio/midi mid midi kar;
  audio/mpeg mp3;
  audio/ogg ogg;
  audio/x-m4a m4a;
  audio/x-realaudio ra;
  video/3gpp 3gpp 3gp;
  video/mp4 mp4;
  video/mpeg mpeg mpg;
  video/quicktime mov;
  video/webm webm;
  video/x-flv flv;
  video/x-m4v m4v;
  video/x-mng mng;
  video/x-ms-asf asx asf;
  video/x-ms-wmv wmv;
  video/x-msvideo avi;
}
`
)
<commit_msg>Move nginx config directives from location to server block<commit_after>package finalize

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

export APP_ROOT=$HOME
export LD_LIBRARY_PATH=$APP_ROOT/nginx/lib:$LD_LIBRARY_PATH

mv $APP_ROOT/nginx/conf/nginx.conf $APP_ROOT/nginx/conf/nginx.conf.erb
erb $APP_ROOT/nginx/conf/nginx.conf.erb > $APP_ROOT/nginx/conf/nginx.conf

if [[ ! -f $APP_ROOT/nginx/logs/access.log ]]; then
    mkfifo $APP_ROOT/nginx/logs/access.log
fi

if [[ ! -f $APP_ROOT/nginx/logs/error.log ]]; then
    mkfifo $APP_ROOT/nginx/logs/error.log
fi
`

	startLoggingScript = `
cat < $APP_ROOT/nginx/logs/access.log &
(>&2 cat) < $APP_ROOT/nginx/logs/error.log &
`

	startCommand = `#!/bin/sh
set -ex
$APP_ROOT/start_logging.sh
nginx -p $APP_ROOT/nginx -c $APP_ROOT/nginx/conf/nginx.conf
`

	nginxConfTemplate = `
worker_processes 1;
daemon off;

error_log <%= ENV["APP_ROOT"] %>/nginx/logs/error.log;
events { worker_connections 1024; }

http {
  charset utf-8;
  log_format cloudfoundry '$http_x_forwarded_for - $http_referer - [$time_local] "$request" $status $body_bytes_sent';
  access_log <%= ENV["APP_ROOT"] %>/nginx/logs/access.log cloudfoundry;
  default_type application/octet-stream;
  include mime.types;
  sendfile on;

  gzip on;
  gzip_disable "msie6";
  gzip_comp_level 6;
  gzip_min_length 1100;
  gzip_buffers 16 8k;
  gzip_proxied any;
  gunzip on;
  gzip_static always;
  gzip_types text/plain text/css text/js text/xml text/javascript application/javascript application/x-javascript application/json application/xml application/xml+rss;
  gzip_vary on;

  tcp_nopush on;
  keepalive_timeout 30;
  port_in_redirect off; # Ensure that redirects don't include the internal container PORT - <%= ENV["PORT"] %>
  server_tokens off;

  server {
    listen <%= ENV["PORT"] %>;
    server_name localhost;

    root <%= ENV["APP_ROOT"] %>/public;

  {{if .PushState}}
    if (!-e $request_filename) {
      rewrite ^(.*)$ / break;
    }
  {{end}}

    index index.html index.htm Default.htm;

  {{if .DirectoryIndex}}
    autoindex on;
  {{end}}

  {{if .BasicAuth}}
    auth_basic "Restricted";  #For Basic Auth
    auth_basic_user_file <%= ENV["APP_ROOT"] %>/nginx/conf/.htpasswd;
  {{end}}

  {{if .ForceHTTPS}}
    if ($http_x_forwarded_proto != "https") {
      return 301 https://$host$request_uri;
    }
  {{else}}
  <% if ENV["FORCE_HTTPS"] %>
    if ($http_x_forwarded_proto != "https") {
      return 301 https://$host$request_uri;
    }
  <% end %>
  {{end}}

  {{if .SSI}}
    ssi on;
  {{end}}

  {{if .HSTS}}
    add_header Strict-Transport-Security "max-age=31536000{{if .HSTSIncludeSubDomains}}; includeSubDomains{{end}}{{if .HSTSPreload}}; preload{{end}}";
  {{end}}

  {{if not .HostDotFiles}}
    location ~ /\. {
      deny all;
      return 404;
    }
  {{end}}

  {{if ne .LocationInclude ""}}
    location / {
      include {{.LocationInclude}};
    }
  {{end}}
  }
}
`
	MimeTypes = `
types {
  text/html html htm shtml;
  text/css css;
  text/xml xml;
  image/gif gif;
  image/jpeg jpeg jpg;
  application/x-javascript js;
  application/atom+xml atom;
  application/rss+xml rss;
  font/ttf ttf;
  font/woff woff;
  font/woff2 woff2;
  text/mathml mml;
  text/plain txt;
  text/vnd.sun.j2me.app-descriptor jad;
  text/vnd.wap.wml wml;
  text/x-component htc;
  text/cache-manifest manifest;
  image/png png;
  image/tiff tif tiff;
  image/vnd.wap.wbmp wbmp;
  image/x-icon ico;
  image/x-jng jng;
  image/x-ms-bmp bmp;
  image/svg+xml svg svgz;
  image/webp webp;
  application/java-archive jar war ear;
  application/mac-binhex40 hqx;
  application/msword doc;
  application/pdf pdf;
  application/postscript ps eps ai;
  application/rtf rtf;
  application/vnd.ms-excel xls;
  application/vnd.ms-powerpoint ppt;
  application/vnd.wap.wmlc wmlc;
  application/vnd.google-earth.kml+xml  kml;
  application/vnd.google-earth.kmz kmz;
  application/x-7z-compressed 7z;
  application/x-cocoa cco;
  application/x-java-archive-diff jardiff;
  application/x-java-jnlp-file jnlp;
  application/x-makeself run;
  application/x-perl pl pm;
  application/x-pilot prc pdb;
  application/x-rar-compressed rar;
  application/x-redhat-package-manager  rpm;
  application/x-sea sea;
  application/x-shockwave-flash swf;
  application/x-stuffit sit;
  application/x-tcl tcl tk;
  application/x-x509-ca-cert der pem crt;
  application/x-xpinstall xpi;
  application/xhtml+xml xhtml;
  application/zip zip;
  application/octet-stream bin exe dll;
  application/octet-stream deb;
  application/octet-stream dmg;
  application/octet-stream eot;
  application/octet-stream iso img;
  application/octet-stream msi msp msm;
  application/json json;
  audio/midi mid midi kar;
  audio/mpeg mp3;
  audio/ogg ogg;
  audio/x-m4a m4a;
  audio/x-realaudio ra;
  video/3gpp 3gpp 3gp;
  video/mp4 mp4;
  video/mpeg mpeg mpg;
  video/quicktime mov;
  video/webm webm;
  video/x-flv flv;
  video/x-m4v m4v;
  video/x-mng mng;
  video/x-ms-asf asx asf;
  video/x-ms-wmv wmv;
  video/x-msvideo avi;
}
`
)
<|endoftext|>


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

package main

import (
	"fmt"
	"log"
	"os"
	"runtime/pprof"

	"github.com/ikravets/errs"
	"github.com/jessevdk/go-flags"

	"my/ev/cmd"
)

func main() {
	var err error
	log.SetFlags(log.Ldate | log.Ltime | log.Lshortfile)

	var opts struct {
		LogFileName string `short:"l" long:"log" value-name:"FILE" default:"/dev/stderr" default-mask:"stderr" description:"log file"`
		ProfileCpu  string `long:"profile-cpu" value-name:"FILE"`
		ProfileMem  string `long:"profile-mem" value-name:"FILE"`
	}

	parser := flags.NewParser(&opts, flags.PassDoubleDash|flags.HelpFlag)
	cmd.Registry.ConfigParser(parser)
	_, err = parser.Parse()
	if e, ok := err.(*flags.Error); ok && e.Type != flags.ErrUnknown {
		fmt.Printf("%s\n", e.Message)
		if e.Type == flags.ErrHelp {
			os.Exit(0)
		} else {
			os.Exit(1)
		}
	}
	errs.CheckE(err)

	logFile, err := os.OpenFile(opts.LogFileName, os.O_WRONLY|os.O_CREATE|os.O_APPEND, 0644)
	errs.CheckE(err)
	log.SetOutput(logFile)
	defer func() { errs.CheckE(logFile.Close()) }()
	log.Printf("args: %v\n", os.Args)

	if opts.ProfileCpu != "" {
		profFile, err := os.Create(opts.ProfileCpu)
		errs.CheckE(err)
		errs.CheckE(pprof.StartCPUProfile(profFile))
		defer pprof.StopCPUProfile()
	}

	cmd.Registry.ParsingFinished()

	if opts.ProfileMem != "" {
		profFile, err := os.Create(opts.ProfileMem)
		errs.CheckE(err)
		errs.CheckE(pprof.WriteHeapProfile(profFile))
		errs.CheckE(profFile.Close())
	}
	log.Println("finished")
}
<commit_msg>main: exit 1 on error<commit_after>// Copyright (c) Ilia Kravets, 2015. All rights reserved. PROVIDED "AS IS"
// WITHOUT ANY WARRANTY, EXPRESS OR IMPLIED. See LICENSE file for details.

package main

import (
	"fmt"
	"log"
	"os"
	"runtime/pprof"

	"github.com/ikravets/errs"
	"github.com/jessevdk/go-flags"

	"my/ev/cmd"
)

func main() {
	if err := do_main(); err != nil {
		log.Printf("exit error: %s\n", err)
		os.Exit(1)
	} else {
		log.Println("exit ok")
	}
}

func do_main() (err error) {
	defer errs.PassE(&err)
	log.SetFlags(log.Ldate | log.Ltime | log.Lshortfile)

	var opts struct {
		LogFileName string `short:"l" long:"log" value-name:"FILE" default:"/dev/stderr" default-mask:"stderr" description:"log file"`
		ProfileCpu  string `long:"profile-cpu" value-name:"FILE"`
		ProfileMem  string `long:"profile-mem" value-name:"FILE"`
	}

	parser := flags.NewParser(&opts, flags.PassDoubleDash|flags.HelpFlag)
	cmd.Registry.ConfigParser(parser)
	_, err = parser.Parse()
	if e, ok := err.(*flags.Error); ok && e.Type != flags.ErrUnknown {
		fmt.Printf("%s\n", e.Message)
		if e.Type == flags.ErrHelp {
			os.Exit(0)
		} else {
			os.Exit(1)
		}
	}
	errs.CheckE(err)

	logFile, err := os.OpenFile(opts.LogFileName, os.O_WRONLY|os.O_CREATE|os.O_APPEND, 0644)
	errs.CheckE(err)
	log.SetOutput(logFile)
	log.Printf("args: %v\n", os.Args)

	if opts.ProfileCpu != "" {
		profFile, err := os.Create(opts.ProfileCpu)
		errs.CheckE(err)
		errs.CheckE(pprof.StartCPUProfile(profFile))
		defer pprof.StopCPUProfile()
	}

	err = cmd.Registry.ParsingFinished()

	if opts.ProfileMem != "" {
		profFile, err := os.Create(opts.ProfileMem)
		errs.CheckE(err)
		errs.CheckE(pprof.WriteHeapProfile(profFile))
		errs.CheckE(profFile.Close())
	}
	return
}
<|endoftext|>


<commit_before>package cli

import (
	"bufio"
	"context"
	"fmt"
	"io"
	"log"
	"os"
	"os/exec"
	"os/signal"
	"path/filepath"
	"syscall"
	"time"

	"github.com/cheggaaa/pb"
	"github.com/dustin/go-humanize"
	"github.com/pkg/errors"
	"github.com/spf13/cobra"
	"github.com/taku-k/polymerase/pkg/base"
	"github.com/taku-k/polymerase/pkg/storage/storagepb"
	"github.com/taku-k/polymerase/pkg/utils/dirutil"
	pexec "github.com/taku-k/polymerase/pkg/utils/exec"
	"go.uber.org/ratelimit"
	"golang.org/x/sync/errgroup"
)

const (
	progressBarWidth = 80
)

type restoreContext struct {
	*base.Config

	from string

	applyPrepare bool

	maxBandWidth string

	latest bool

	decompressCmd string
}

var restoreCmd = &cobra.Command{
	Use:   "restore",
	Short: "Receives backup data to restore from a polymerase server",
	RunE:  runRestore,
}

func runRestore(cmd *cobra.Command, args []string) error {
	if len(args) > 0 {
		return usageAndError(cmd)
	}

	if restoreCtx.from == "" && !restoreCtx.latest {
		return errors.New("You must specify `from` option or `latest` flag.")
	}
	if db == "" {
		return errors.New("You must specify `db`")
	}
	var maxBandWidth uint64
	if restoreCtx.maxBandWidth != "" {
		if bw, err := humanize.ParseBytes(restoreCtx.maxBandWidth); err != nil {
			return errors.Wrap(err, "Cannot parse -max-bandwidth")
		} else {
			maxBandWidth = bw
		}
	}

	// If `from` is not specified and `latest` option is added,
	// restoreCtx.from is set as tomorrow.
	if restoreCtx.from == "" && restoreCtx.latest {
		restoreCtx.from = time.Now().AddDate(0, 0, 1).Format("2006-01-02")
	}

	// Signal
	signalCh := make(chan os.Signal, 1)
	signal.Notify(signalCh, syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT)

	errCh := make(chan error, 1)
	finishCh := make(chan struct{})

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	go func() {
		scli, err := getStorageClient(ctx, db)
		if err != nil {
			errCh <- err
			return
		}
		res, err := scli.GetKeysAtPoint(context.Background(), &storagepb.GetKeysAtPointRequest{
			Db:   db,
			From: restoreCtx.from,
		})

		restoreDir, err := filepath.Abs("polymerase-restore")
		if err != nil {
			errCh <- err
			return
		}
		if err := dirutil.MkdirAllWithLog(restoreDir); err != nil {
			errCh <- err
			return
		}
		log.Printf("Restore data directory: %v\n", restoreDir)

		pbs := make([]*pb.ProgressBar, len(res.Keys))
		pbs[len(res.Keys)-1] = pb.New64(int64(res.Keys[len(res.Keys)-1].Size)).Prefix("base | ")
		for inc, idx := len(res.Keys)-1, 0; inc > 0; inc -= 1 {
			info := res.Keys[idx]
			pbs[idx] = pb.New64(int64(info.Size)).Prefix(fmt.Sprintf("inc%d | ", inc))
		}
		for _, bar := range pbs {
			bar.SetWidth(progressBarWidth)
			bar.SetUnits(pb.U_BYTES)
		}

		pool, err := pb.StartPool(pbs...)
		if err != nil {
			errCh <- err
			return
		}

		g, _ := errgroup.WithContext(ctx)

		for inc, idx := len(res.Keys)-1, 0; inc > 0; inc -= 1 {
			info := res.Keys[idx]
			bar := pbs[idx]
			inc := inc
			g.Go(func() error {
				return getIncBackup(scli, info, restoreDir, inc, bar, maxBandWidth)
			})
			idx += 1
		}
		g.Go(func() error {
			return getFullBackup(scli, res.Keys[len(res.Keys)-1], restoreDir, pbs[len(res.Keys)-1], maxBandWidth)
		})
		if err := g.Wait(); err != nil {
			pool.Stop()
			errCh <- err
			return
		}
		pool.Stop()

		// Automatically preparing backups only when applyPrepare flag is true.
		if restoreCtx.applyPrepare {
			os.Chdir(restoreDir)
			c := pexec.PrepareBaseBackup(ctx, len(res.Keys) == 1, xtrabackupCfg)
			if err := c.Run(); err != nil {
				errCh <- errors.Wrap(err, fmt.Sprintf("failed preparing base: %v", c.Args))
			}
			for inc := 1; inc < len(res.Keys); inc += 1 {
				c := pexec.PrepareIncBackup(ctx, inc, inc == len(res.Keys)-1, xtrabackupCfg)
				if err := c.Run(); err != nil {
					errCh <- errors.Wrap(err, fmt.Sprintf("failed preparing inc%d: %v", inc, c.Args))
				}
			}
		}

		finishCh <- struct{}{}
	}()

	select {
	case err := <-errCh:
		return err
	case <-signalCh:
	case <-finishCh:
	}
	return nil
}

func getIncBackup(
	cli storagepb.StorageServiceClient,
	info *storagepb.BackupFileInfo,
	restoreDir string,
	inc int,
	bar *pb.ProgressBar,
	maxBandwidth uint64,
) error {
	fn := filepath.Join(restoreDir, fmt.Sprintf("inc%d.xb.gz", inc))
	if err := getBackup(cli, info, fn, bar, maxBandwidth); err != nil {
		return err
	}
	if err := unzipIncBackupCmd(context.TODO(), fn, restoreDir, inc); err != nil {
		return err
	}
	return nil
}

func getFullBackup(
	cli storagepb.StorageServiceClient,
	info *storagepb.BackupFileInfo,
	restoreDir string,
	bar *pb.ProgressBar,
	maxBandwidth uint64,
) error {
	fn := filepath.Join(restoreDir, "base.tar.gz")
	if err := getBackup(cli, info, fn, bar, maxBandwidth); err != nil {
		return err
	}
	if err := unzipFullBackupCmd(context.TODO(), fn, restoreDir); err != nil {
		return err
	}
	return nil
}

func getBackup(
	cli storagepb.StorageServiceClient,
	info *storagepb.BackupFileInfo,
	fn string,
	bar *pb.ProgressBar,
	maxBandWitdh uint64,
) error {
	stream, err := cli.GetFileByKey(context.Background(), &storagepb.GetFileByKeyRequest{
		Key:         info.Key,
		StorageType: info.StorageType,
	})
	if err != nil {
		return err
	}
	f, err := os.Create(fn)
	if err != nil {
		return err
	}
	bufw := bufio.NewWriter(f)
	multiw := io.MultiWriter(bufw, bar)

	if err := readFromStreamWithRateLimit(stream, multiw, maxBandWitdh); err != nil {
		return err
	}

	bar.Finish()
	return nil
}

func readFromStreamWithRateLimit(
	stream storagepb.StorageService_GetFileByKeyClient,
	writer io.Writer,
	maxBandwidth uint64,
) error {
	fs, err := stream.Recv()
	if err == io.EOF {
		return nil
	}
	if err != nil {
		return err
	}
	writer.Write(fs.Content)

	// Determine the number of operations to perform per second
	// based on maxBandwidth.
	var rl ratelimit.Limiter
	if maxBandwidth == 0 {
		rl = ratelimit.NewUnlimited()
	} else {
		log.Printf("Set max bandwidth %s/sec\n", humanize.Bytes(maxBandwidth))
		rl = ratelimit.New(int(maxBandwidth / uint64(len(fs.Content))))
	}

	for {
		rl.Take()
		fs, err := stream.Recv()
		if err == io.EOF {
			break
		}
		if err != nil {
			return err
		}
		writer.Write(fs.Content)
	}

	return nil
}

func unzipIncBackupCmd(ctx context.Context, name, dir string, inc int) error {
	if dir == "" {
		return errors.New("directory path cannot be empty")
	}

	odir := filepath.Join(dir, fmt.Sprintf("inc%d", inc))
	if err := os.MkdirAll(odir, 0755); err != nil {
		return errors.Wrap(err, odir+" dir cannot be created")
	}

	cmd := exec.CommandContext(
		ctx,
		"sh", "-c", fmt.Sprintf("%s -c %s | xbstream -x -C %s", restoreCtx.decompressCmd, name, odir))
	if err := cmd.Run(); err != nil {
		return errors.Wrap(err, fmt.Sprintf("failed `%s -c %s | xbstream -x -C %s", restoreCtx.decompressCmd, name, odir))
	}

	if err := os.Remove(name); err != nil {
		return errors.Wrap(err, "failed to remove "+name)
	}

	return nil
}

func unzipFullBackupCmd(ctx context.Context, name, dir string) error {
	odir := filepath.Join(dir, "base")
	if err := os.MkdirAll(odir, 0755); err != nil {
		return errors.Wrap(err, odir+" dir cannot be created")
	}

	cl := fmt.Sprintf("%s < %s | tar -xC %s", restoreCtx.decompressCmd, name, odir)
	if err := exec.CommandContext(ctx, "sh", "-c", cl).Run(); err != nil {
		return errors.Wrap(err, "Failed unzip")
	}

	if err := os.Remove(name); err != nil {
		return err
	}

	return nil
}
<commit_msg>Fix bug that index is not incremented<commit_after>package cli

import (
	"bufio"
	"context"
	"fmt"
	"io"
	"log"
	"os"
	"os/exec"
	"os/signal"
	"path/filepath"
	"syscall"
	"time"

	"github.com/cheggaaa/pb"
	"github.com/dustin/go-humanize"
	"github.com/pkg/errors"
	"github.com/spf13/cobra"
	"github.com/taku-k/polymerase/pkg/base"
	"github.com/taku-k/polymerase/pkg/storage/storagepb"
	"github.com/taku-k/polymerase/pkg/utils/dirutil"
	pexec "github.com/taku-k/polymerase/pkg/utils/exec"
	"go.uber.org/ratelimit"
	"golang.org/x/sync/errgroup"
)

const (
	progressBarWidth = 80
)

type restoreContext struct {
	*base.Config

	from string

	applyPrepare bool

	maxBandWidth string

	latest bool

	decompressCmd string
}

var restoreCmd = &cobra.Command{
	Use:   "restore",
	Short: "Receives backup data to restore from a polymerase server",
	RunE:  runRestore,
}

func runRestore(cmd *cobra.Command, args []string) error {
	if len(args) > 0 {
		return usageAndError(cmd)
	}

	if restoreCtx.from == "" && !restoreCtx.latest {
		return errors.New("You must specify `from` option or `latest` flag.")
	}
	if db == "" {
		return errors.New("You must specify `db`")
	}
	var maxBandWidth uint64
	if restoreCtx.maxBandWidth != "" {
		if bw, err := humanize.ParseBytes(restoreCtx.maxBandWidth); err != nil {
			return errors.Wrap(err, "Cannot parse -max-bandwidth")
		} else {
			maxBandWidth = bw
		}
	}

	// If `from` is not specified and `latest` option is added,
	// restoreCtx.from is set as tomorrow.
	if restoreCtx.from == "" && restoreCtx.latest {
		restoreCtx.from = time.Now().AddDate(0, 0, 1).Format("2006-01-02")
	}

	// Signal
	signalCh := make(chan os.Signal, 1)
	signal.Notify(signalCh, syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT)

	errCh := make(chan error, 1)
	finishCh := make(chan struct{})

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	go func() {
		scli, err := getStorageClient(ctx, db)
		if err != nil {
			errCh <- err
			return
		}
		res, err := scli.GetKeysAtPoint(context.Background(), &storagepb.GetKeysAtPointRequest{
			Db:   db,
			From: restoreCtx.from,
		})

		restoreDir, err := filepath.Abs("polymerase-restore")
		if err != nil {
			errCh <- err
			return
		}
		if err := dirutil.MkdirAllWithLog(restoreDir); err != nil {
			errCh <- err
			return
		}
		log.Printf("Restore data directory: %v\n", restoreDir)

		pbs := make([]*pb.ProgressBar, len(res.Keys))
		pbs[len(res.Keys)-1] = pb.New64(int64(res.Keys[len(res.Keys)-1].Size)).Prefix("base | ")
		for i := 0; i < len(res.Keys)-1; i += 1 {
			inc := len(res.Keys) - i - 1
			info := res.Keys[i]
			pbs[i] = pb.New64(int64(info.Size)).Prefix(fmt.Sprintf("inc%d | ", inc))
		}
		for _, bar := range pbs {
			bar.SetWidth(progressBarWidth)
			bar.SetUnits(pb.U_BYTES)
		}

		pool, err := pb.StartPool(pbs...)
		if err != nil {
			errCh <- err
			return
		}

		g, _ := errgroup.WithContext(ctx)

		for inc, idx := len(res.Keys)-1, 0; inc > 0; inc -= 1 {
			info := res.Keys[idx]
			bar := pbs[idx]
			inc := inc
			g.Go(func() error {
				return getIncBackup(scli, info, restoreDir, inc, bar, maxBandWidth)
			})
			idx += 1
		}
		g.Go(func() error {
			return getFullBackup(scli, res.Keys[len(res.Keys)-1], restoreDir, pbs[len(res.Keys)-1], maxBandWidth)
		})
		if err := g.Wait(); err != nil {
			pool.Stop()
			errCh <- err
			return
		}
		pool.Stop()

		// Automatically preparing backups only when applyPrepare flag is true.
		if restoreCtx.applyPrepare {
			os.Chdir(restoreDir)
			c := pexec.PrepareBaseBackup(ctx, len(res.Keys) == 1, xtrabackupCfg)
			if err := c.Run(); err != nil {
				errCh <- errors.Wrap(err, fmt.Sprintf("failed preparing base: %v", c.Args))
			}
			for inc := 1; inc < len(res.Keys); inc += 1 {
				c := pexec.PrepareIncBackup(ctx, inc, inc == len(res.Keys)-1, xtrabackupCfg)
				if err := c.Run(); err != nil {
					errCh <- errors.Wrap(err, fmt.Sprintf("failed preparing inc%d: %v", inc, c.Args))
				}
			}
		}

		finishCh <- struct{}{}
	}()

	select {
	case err := <-errCh:
		return err
	case <-signalCh:
	case <-finishCh:
	}
	return nil
}

func getIncBackup(
	cli storagepb.StorageServiceClient,
	info *storagepb.BackupFileInfo,
	restoreDir string,
	inc int,
	bar *pb.ProgressBar,
	maxBandwidth uint64,
) error {
	fn := filepath.Join(restoreDir, fmt.Sprintf("inc%d.xb.gz", inc))
	if err := getBackup(cli, info, fn, bar, maxBandwidth); err != nil {
		return err
	}
	if err := unzipIncBackupCmd(context.TODO(), fn, restoreDir, inc); err != nil {
		return err
	}
	return nil
}

func getFullBackup(
	cli storagepb.StorageServiceClient,
	info *storagepb.BackupFileInfo,
	restoreDir string,
	bar *pb.ProgressBar,
	maxBandwidth uint64,
) error {
	fn := filepath.Join(restoreDir, "base.tar.gz")
	if err := getBackup(cli, info, fn, bar, maxBandwidth); err != nil {
		return err
	}
	if err := unzipFullBackupCmd(context.TODO(), fn, restoreDir); err != nil {
		return err
	}
	return nil
}

func getBackup(
	cli storagepb.StorageServiceClient,
	info *storagepb.BackupFileInfo,
	fn string,
	bar *pb.ProgressBar,
	maxBandWitdh uint64,
) error {
	stream, err := cli.GetFileByKey(context.Background(), &storagepb.GetFileByKeyRequest{
		Key:         info.Key,
		StorageType: info.StorageType,
	})
	if err != nil {
		return err
	}
	f, err := os.Create(fn)
	if err != nil {
		return err
	}
	bufw := bufio.NewWriter(f)
	multiw := io.MultiWriter(bufw, bar)

	if err := readFromStreamWithRateLimit(stream, multiw, maxBandWitdh); err != nil {
		return err
	}

	bar.Finish()
	return nil
}

func readFromStreamWithRateLimit(
	stream storagepb.StorageService_GetFileByKeyClient,
	writer io.Writer,
	maxBandwidth uint64,
) error {
	fs, err := stream.Recv()
	if err == io.EOF {
		return nil
	}
	if err != nil {
		return err
	}
	writer.Write(fs.Content)

	// Determine the number of operations to perform per second
	// based on maxBandwidth.
	var rl ratelimit.Limiter
	if maxBandwidth == 0 {
		rl = ratelimit.NewUnlimited()
	} else {
		log.Printf("Set max bandwidth %s/sec\n", humanize.Bytes(maxBandwidth))
		rl = ratelimit.New(int(maxBandwidth / uint64(len(fs.Content))))
	}

	for {
		rl.Take()
		fs, err := stream.Recv()
		if err == io.EOF {
			break
		}
		if err != nil {
			return err
		}
		writer.Write(fs.Content)
	}

	return nil
}

func unzipIncBackupCmd(ctx context.Context, name, dir string, inc int) error {
	if dir == "" {
		return errors.New("directory path cannot be empty")
	}

	odir := filepath.Join(dir, fmt.Sprintf("inc%d", inc))
	if err := os.MkdirAll(odir, 0755); err != nil {
		return errors.Wrap(err, odir+" dir cannot be created")
	}

	cmd := exec.CommandContext(
		ctx,
		"sh", "-c", fmt.Sprintf("%s -c %s | xbstream -x -C %s", restoreCtx.decompressCmd, name, odir))
	if err := cmd.Run(); err != nil {
		return errors.Wrap(err, fmt.Sprintf("failed `%s -c %s | xbstream -x -C %s", restoreCtx.decompressCmd, name, odir))
	}

	if err := os.Remove(name); err != nil {
		return errors.Wrap(err, "failed to remove "+name)
	}

	return nil
}

func unzipFullBackupCmd(ctx context.Context, name, dir string) error {
	odir := filepath.Join(dir, "base")
	if err := os.MkdirAll(odir, 0755); err != nil {
		return errors.Wrap(err, odir+" dir cannot be created")
	}

	cl := fmt.Sprintf("%s < %s | tar -xC %s", restoreCtx.decompressCmd, name, odir)
	if err := exec.CommandContext(ctx, "sh", "-c", cl).Run(); err != nil {
		return errors.Wrap(err, "Failed unzip")
	}

	if err := os.Remove(name); err != nil {
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>package cli

import (
	"bufio"
	"context"
	"fmt"
	"io"
	"log"
	"math"
	"os"
	"os/exec"
	"os/signal"
	"path/filepath"
	"strconv"
	"strings"
	"syscall"
	"time"

	"github.com/cheggaaa/pb"
	"github.com/dustin/go-humanize"
	"github.com/elastic/gosigar"
	"github.com/pkg/errors"
	"github.com/spf13/cobra"
	"github.com/taku-k/polymerase/pkg/base"
	"github.com/taku-k/polymerase/pkg/storage/storagepb"
	"github.com/taku-k/polymerase/pkg/utils/dirutil"
	pexec "github.com/taku-k/polymerase/pkg/utils/exec"
	"go.uber.org/ratelimit"
	"golang.org/x/sync/errgroup"
)

const (
	progressBarWidth = 80

	defaultApplyPrepare = false
	defaultUseMemory    = "100MB"
)

type restoreContext struct {
	*base.Config

	from string

	applyPrepare bool

	maxBandWidth MaxBandWidthType

	latest bool

	decompressCmd string

	useMemory UseMemoryType
}

func MakeRestoreContext(cfg *base.Config) *restoreContext {
	// Use half the total memory
	mem := gosigar.Mem{}
	um := defaultUseMemory
	if err := mem.Get(); err == nil {
		ss := strings.Split(humanize.Bytes(mem.Total/2), " ")
		hm, err := strconv.ParseFloat(ss[0], 64)
		if err != nil {
			panic(err)
		}
		// Round off to the nearest whole number
		um = fmt.Sprintf("%d%s", int(math.Floor(hm+0.5)), ss[1])
	}

	return &restoreContext{
		Config:       cfg,
		applyPrepare: defaultApplyPrepare,
		useMemory:    UseMemoryType(um),
	}
}

var restoreCmd = &cobra.Command{
	Use:   "restore",
	Short: "Receives backup data to restore from a polymerase server",
	RunE:  runRestore,
}

func runRestore(cmd *cobra.Command, args []string) error {
	if len(args) > 0 {
		return usageAndError(cmd)
	}

	if restoreCtx.from == "" && !restoreCtx.latest {
		return errors.New("You must specify `from` option or `latest` flag.")
	}
	if db == "" {
		return errors.New("You must specify `db`")
	}
	maxBandWidth := uint64(restoreCtx.maxBandWidth)

	// If `from` is not specified and `latest` option is added,
	// restoreCtx.from is set as tomorrow.
	if restoreCtx.from == "" && restoreCtx.latest {
		restoreCtx.from = time.Now().AddDate(0, 0, 1).Format("2006-01-02")
	}

	// Signal
	signalCh := make(chan os.Signal, 1)
	signal.Notify(signalCh, syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT)

	errCh := make(chan error, 1)
	finishCh := make(chan struct{})

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	go func() {
		scli, err := getSuitableStorageClient(ctx, db)
		if err != nil {
			errCh <- err
			return
		}
		res, err := scli.GetKeysAtPoint(context.Background(), &storagepb.GetKeysAtPointRequest{
			Db:   db,
			From: restoreCtx.from,
		})

		restoreDir, err := filepath.Abs("polymerase-restore")
		if err != nil {
			errCh <- err
			return
		}
		if err := dirutil.MkdirAllWithLog(restoreDir); err != nil {
			errCh <- err
			return
		}
		log.Printf("Restore data directory: %v\n", restoreDir)

		pbs := make([]*pb.ProgressBar, len(res.Keys))
		pbs[len(res.Keys)-1] = pb.New64(int64(res.Keys[len(res.Keys)-1].Size)).Prefix("base | ")
		for i := 0; i < len(res.Keys)-1; i += 1 {
			inc := len(res.Keys) - i - 1
			info := res.Keys[i]
			pbs[i] = pb.New64(int64(info.Size)).Prefix(fmt.Sprintf("inc%d | ", inc))
		}
		for _, bar := range pbs {
			bar.SetWidth(progressBarWidth)
			bar.SetUnits(pb.U_BYTES)
		}

		pool, err := pb.StartPool(pbs...)
		if err != nil {
			errCh <- err
			return
		}

		g, _ := errgroup.WithContext(ctx)

		for inc, idx := len(res.Keys)-1, 0; inc > 0; inc -= 1 {
			info := res.Keys[idx]
			bar := pbs[idx]
			inc := inc
			g.Go(func() error {
				return getIncBackup(scli, info, restoreDir, inc, bar, maxBandWidth)
			})
			idx += 1
		}
		g.Go(func() error {
			return getFullBackup(scli, res.Keys[len(res.Keys)-1], restoreDir, pbs[len(res.Keys)-1], maxBandWidth)
		})
		if err := g.Wait(); err != nil {
			pool.Stop()
			errCh <- err
			return
		}
		pool.Stop()

		// Automatically preparing backups only when applyPrepare flag is true.
		if restoreCtx.applyPrepare {
			os.Chdir(restoreDir)
			c, err := pexec.PrepareBaseBackup(ctx, len(res.Keys) == 1, xtrabackupCfg)
			if err != nil {
				errCh <- err
				return
			}
			log.Println(pexec.StringWithMaskPassword(c))
			c.Stdout = os.Stderr
			c.Stderr = os.Stderr
			if err := c.Run(); err != nil {
				errCh <- errors.Wrap(err, fmt.Sprintf("failed preparing base: %v", c.Args))
				return
			}
			for inc := 1; inc < len(res.Keys); inc += 1 {
				c, err := pexec.PrepareIncBackup(ctx, inc, inc == len(res.Keys)-1, xtrabackupCfg)
				if err != nil {
					errCh <- err
					return
				}
				log.Println(pexec.StringWithMaskPassword(c))
				c.Stdout = os.Stderr
				c.Stderr = os.Stderr
				if err := c.Run(); err != nil {
					errCh <- errors.Wrap(err, fmt.Sprintf("failed preparing inc%d: %v", inc, c.Args))
					return
				}
			}
		}

		finishCh <- struct{}{}
	}()

	select {
	case err := <-errCh:
		return err
	case <-signalCh:
	case <-finishCh:
	}
	return nil
}

func getIncBackup(
	cli storagepb.StorageServiceClient,
	info *storagepb.BackupFileInfo,
	restoreDir string,
	inc int,
	bar *pb.ProgressBar,
	maxBandwidth uint64,
) error {
	odir := filepath.Join(restoreDir, fmt.Sprintf("inc%d", inc))
	if err := os.MkdirAll(odir, 0755); err != nil {
		return errors.Wrap(err, odir+" dir cannot be created")
	}
	extractCmd := exec.Command(
		"sh",
		"-c",
		fmt.Sprintf("%s -c -d | xbstream -x -C %s", restoreCtx.decompressCmd, odir))
	if err := getBackup(cli, info, bar, maxBandwidth, extractCmd); err != nil {
		return err
	}
	return nil
}

func getFullBackup(
	cli storagepb.StorageServiceClient,
	info *storagepb.BackupFileInfo,
	restoreDir string,
	bar *pb.ProgressBar,
	maxBandwidth uint64,
) error {
	odir := filepath.Join(restoreDir, "base")
	if err := os.MkdirAll(odir, 0755); err != nil {
		return errors.Wrap(err, odir+" dir cannot be created")
	}
	extractCmd := exec.Command(
		"sh",
		"-c",
		fmt.Sprintf("%s -c -d | tar -xC %s", restoreCtx.decompressCmd, odir))
	if err := getBackup(cli, info, bar, maxBandwidth, extractCmd); err != nil {
		return err
	}
	return nil
}

func getBackup(
	cli storagepb.StorageServiceClient,
	info *storagepb.BackupFileInfo,
	bar *pb.ProgressBar,
	maxBandWitdh uint64,
	extractCmd *exec.Cmd,
) error {
	stream, err := cli.GetFileByKey(context.Background(), &storagepb.GetFileByKeyRequest{
		Key:         info.Key,
		StorageType: info.StorageType,
	})
	if err != nil {
		return err
	}
	extractCmd.Stderr = os.Stderr
	extractCmdStdin, err := extractCmd.StdinPipe()
	if err != nil {
		return err
	}
	bufw := bufio.NewWriter(extractCmdStdin)
	multiw := io.MultiWriter(bufw, bar)

	finishCh := make(chan struct{})
	errCh := make(chan error)

	go func() {
		err := extractCmd.Start()
		if err != nil {
			errCh <- err
			return
		}
		extractCmd.Wait()
		finishCh <- struct{}{}
	}()

	go writeOutFromStreamWithRateLimit(stream, multiw, maxBandWitdh, errCh, extractCmdStdin)

	// Wait
	select {
	case err := <-errCh:
		return err
	case <-finishCh:
		break
	}

	bar.Finish()
	return nil
}

func writeOutFromStreamWithRateLimit(
	stream storagepb.StorageService_GetFileByKeyClient,
	writer io.Writer,
	maxBandwidth uint64,
	errCh chan error,
	toCloseW io.WriteCloser,
) {
	defer toCloseW.Close()

	fs, err := stream.Recv()
	if err == io.EOF {
		return
	}
	if err != nil {
		errCh <- err
		return
	}
	writer.Write(fs.Content)

	// Determine the number of operations to perform per second
	// based on maxBandwidth.
	var rl ratelimit.Limiter
	if maxBandwidth == 0 {
		rl = ratelimit.NewUnlimited()
	} else {
		log.Printf("Set max bandwidth %s/sec\n", humanize.Bytes(maxBandwidth))
		rl = ratelimit.New(int(maxBandwidth / uint64(len(fs.Content))))
	}

	for {
		rl.Take()
		fs, err := stream.Recv()
		if err == io.EOF {
			break
		}
		if err != nil {
			errCh <- err
			return
		}
		writer.Write(fs.Content)
	}
	return
}
<commit_msg>Make restore command download sequentialy<commit_after>package cli

import (
	"bufio"
	"context"
	"fmt"
	"io"
	"log"
	"math"
	"os"
	"os/exec"
	"os/signal"
	"path/filepath"
	"strconv"
	"strings"
	"syscall"
	"time"

	"github.com/cheggaaa/pb"
	"github.com/dustin/go-humanize"
	"github.com/elastic/gosigar"
	"github.com/pkg/errors"
	"github.com/spf13/cobra"
	"github.com/taku-k/polymerase/pkg/base"
	"github.com/taku-k/polymerase/pkg/storage/storagepb"
	"github.com/taku-k/polymerase/pkg/utils/dirutil"
	pexec "github.com/taku-k/polymerase/pkg/utils/exec"
	"go.uber.org/ratelimit"
)

const (
	progressBarWidth = 80

	defaultApplyPrepare = false
	defaultUseMemory    = "100MB"
)

type restoreContext struct {
	*base.Config

	from string

	applyPrepare bool

	maxBandWidth MaxBandWidthType

	latest bool

	decompressCmd string

	useMemory UseMemoryType
}

func MakeRestoreContext(cfg *base.Config) *restoreContext {
	// Use half the total memory
	mem := gosigar.Mem{}
	um := defaultUseMemory
	if err := mem.Get(); err == nil {
		ss := strings.Split(humanize.Bytes(mem.Total/2), " ")
		hm, err := strconv.ParseFloat(ss[0], 64)
		if err != nil {
			panic(err)
		}
		// Round off to the nearest whole number
		um = fmt.Sprintf("%d%s", int(math.Floor(hm+0.5)), ss[1])
	}

	return &restoreContext{
		Config:       cfg,
		applyPrepare: defaultApplyPrepare,
		useMemory:    UseMemoryType(um),
	}
}

var restoreCmd = &cobra.Command{
	Use:   "restore",
	Short: "Receives backup data to restore from a polymerase server",
	RunE:  runRestore,
}

func runRestore(cmd *cobra.Command, args []string) error {
	if len(args) > 0 {
		return usageAndError(cmd)
	}

	if restoreCtx.from == "" && !restoreCtx.latest {
		return errors.New("You must specify `from` option or `latest` flag.")
	}
	if db == "" {
		return errors.New("You must specify `db`")
	}

	// If `from` is not specified and `latest` option is added,
	// restoreCtx.from is set as tomorrow.
	if restoreCtx.from == "" && restoreCtx.latest {
		restoreCtx.from = time.Now().AddDate(0, 0, 1).Format("2006-01-02")
	}

	// Signal
	signalCh := make(chan os.Signal, 1)
	signal.Notify(signalCh, syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT)

	errCh := make(chan error, 1)
	finishCh := make(chan struct{})

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	go doRestore(ctx, errCh, finishCh)

	select {
	case err := <-errCh:
		return err
	case <-signalCh:
	case <-finishCh:
	}
	return nil
}

func doRestore(ctx context.Context, errCh chan error, finishCh chan struct{}) {
	scli, err := getSuitableStorageClient(ctx, db)
	if err != nil {
		errCh <- err
		return
	}
	res, err := scli.GetKeysAtPoint(context.Background(), &storagepb.GetKeysAtPointRequest{
		Db:   db,
		From: restoreCtx.from,
	})

	restoreDir, err := filepath.Abs("polymerase-restore")
	if err != nil {
		errCh <- err
		return
	}
	if err := dirutil.MkdirAllWithLog(restoreDir); err != nil {
		errCh <- err
		return
	}
	log.Printf("Restore data directory: %v\n", restoreDir)

	pbs := make([]*pb.ProgressBar, len(res.Keys))
	pbs[len(res.Keys)-1] = pb.New64(int64(res.Keys[len(res.Keys)-1].Size)).Prefix("base | ")
	for i := 0; i < len(res.Keys)-1; i += 1 {
		inc := len(res.Keys) - i - 1
		info := res.Keys[i]
		pbs[i] = pb.New64(int64(info.Size)).Prefix(fmt.Sprintf("inc%d | ", inc))
	}
	for _, bar := range pbs {
		bar.SetWidth(progressBarWidth)
		bar.SetUnits(pb.U_BYTES)
	}

	pool, err := pb.StartPool(pbs...)
	if err != nil {
		errCh <- err
		return
	}

	maxBandWidth := uint64(restoreCtx.maxBandWidth)
	if maxBandWidth != 0 {
		log.Printf("Set max bandwidth %s/sec\n", humanize.Bytes(maxBandWidth))
	}

	for inc, idx := len(res.Keys)-1, 0; inc > 0; inc -= 1 {
		info := res.Keys[idx]
		bar := pbs[idx]
		inc := inc
		if err := getIncBackup(scli, info, restoreDir, inc, bar, maxBandWidth); err != nil {
			pool.Stop()
			errCh <- err
			return
		}
		idx += 1
	}
	if err := getFullBackup(scli, res.Keys[len(res.Keys)-1], restoreDir, pbs[len(res.Keys)-1], maxBandWidth); err != nil {
		pool.Stop()
		errCh <- err
		return
	}
	pool.Stop()

	// Automatically preparing backups only when applyPrepare flag is true.
	if restoreCtx.applyPrepare {
		os.Chdir(restoreDir)
		c, err := pexec.PrepareBaseBackup(ctx, len(res.Keys) == 1, xtrabackupCfg)
		if err != nil {
			errCh <- err
			return
		}
		log.Println(pexec.StringWithMaskPassword(c))
		c.Stdout = os.Stderr
		c.Stderr = os.Stderr
		if err := c.Run(); err != nil {
			errCh <- errors.Wrap(err, fmt.Sprintf("failed preparing base: %v", c.Args))
			return
		}
		for inc := 1; inc < len(res.Keys); inc += 1 {
			c, err := pexec.PrepareIncBackup(ctx, inc, inc == len(res.Keys)-1, xtrabackupCfg)
			if err != nil {
				errCh <- err
				return
			}
			log.Println(pexec.StringWithMaskPassword(c))
			c.Stdout = os.Stderr
			c.Stderr = os.Stderr
			if err := c.Run(); err != nil {
				errCh <- errors.Wrap(err, fmt.Sprintf("failed preparing inc%d: %v", inc, c.Args))
				return
			}
		}
	}

	finishCh <- struct{}{}
}

func getIncBackup(
	cli storagepb.StorageServiceClient,
	info *storagepb.BackupFileInfo,
	restoreDir string,
	inc int,
	bar *pb.ProgressBar,
	maxBandwidth uint64,
) error {
	odir := filepath.Join(restoreDir, fmt.Sprintf("inc%d", inc))
	if err := os.MkdirAll(odir, 0755); err != nil {
		return errors.Wrap(err, odir+" dir cannot be created")
	}
	extractCmd := exec.Command(
		"sh",
		"-c",
		fmt.Sprintf("%s -c -d | xbstream -x -C %s", restoreCtx.decompressCmd, odir))
	if err := getBackup(cli, info, bar, maxBandwidth, extractCmd); err != nil {
		return err
	}
	return nil
}

func getFullBackup(
	cli storagepb.StorageServiceClient,
	info *storagepb.BackupFileInfo,
	restoreDir string,
	bar *pb.ProgressBar,
	maxBandwidth uint64,
) error {
	odir := filepath.Join(restoreDir, "base")
	if err := os.MkdirAll(odir, 0755); err != nil {
		return errors.Wrap(err, odir+" dir cannot be created")
	}
	extractCmd := exec.Command(
		"sh",
		"-c",
		fmt.Sprintf("%s -c -d | tar -xC %s", restoreCtx.decompressCmd, odir))
	if err := getBackup(cli, info, bar, maxBandwidth, extractCmd); err != nil {
		return err
	}
	return nil
}

func getBackup(
	cli storagepb.StorageServiceClient,
	info *storagepb.BackupFileInfo,
	bar *pb.ProgressBar,
	maxBandWitdh uint64,
	extractCmd *exec.Cmd,
) error {
	stream, err := cli.GetFileByKey(context.Background(), &storagepb.GetFileByKeyRequest{
		Key:         info.Key,
		StorageType: info.StorageType,
	})
	if err != nil {
		return err
	}
	extractCmd.Stderr = os.Stderr
	extractCmdStdin, err := extractCmd.StdinPipe()
	if err != nil {
		return err
	}
	bufw := bufio.NewWriter(extractCmdStdin)
	multiw := io.MultiWriter(bufw, bar)

	finishCh := make(chan struct{})
	errCh := make(chan error)

	go func() {
		err := extractCmd.Start()
		if err != nil {
			errCh <- err
			return
		}
		extractCmd.Wait()
		finishCh <- struct{}{}
	}()

	go writeOutFromStreamWithRateLimit(stream, multiw, maxBandWitdh, errCh, extractCmdStdin)

	// Wait
	select {
	case err := <-errCh:
		return err
	case <-finishCh:
		break
	}

	bar.Finish()
	return nil
}

func writeOutFromStreamWithRateLimit(
	stream storagepb.StorageService_GetFileByKeyClient,
	writer io.Writer,
	maxBandwidth uint64,
	errCh chan error,
	toCloseW io.WriteCloser,
) {
	defer toCloseW.Close()

	fs, err := stream.Recv()
	if err == io.EOF {
		return
	}
	if err != nil {
		errCh <- err
		return
	}
	writer.Write(fs.Content)

	// Determine the number of operations to perform per second
	// based on maxBandwidth.
	var rl ratelimit.Limiter
	if maxBandwidth == 0 {
		rl = ratelimit.NewUnlimited()
	} else {
		rl = ratelimit.New(int(maxBandwidth / uint64(len(fs.Content))))
	}

	for {
		rl.Take()
		fs, err := stream.Recv()
		if err == io.EOF {
			break
		}
		if err != nil {
			errCh <- err
			return
		}
		writer.Write(fs.Content)
	}
	return
}
<|endoftext|>


<commit_before>/*
   Copyright 2020 Docker Compose CLI authors

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
*/

package compose

import (
	"context"
	"fmt"
	"io"

	"github.com/docker/compose/v2/pkg/api"

	"github.com/compose-spec/compose-go/types"
	"github.com/docker/cli/cli/streams"
	moby "github.com/docker/docker/api/types"
	"github.com/docker/docker/pkg/ioutils"
	"github.com/docker/docker/pkg/stdcopy"
	"github.com/docker/docker/pkg/stringid"
	"github.com/moby/term"
)

func (s *composeService) RunOneOffContainer(ctx context.Context, project *types.Project, opts api.RunOptions) (int, error) {
	containerID, err := s.prepareRun(ctx, project, opts)
	if err != nil {
		return 0, err
	}

	if opts.Detach {
		err := s.apiClient.ContainerStart(ctx, containerID, moby.ContainerStartOptions{})
		if err != nil {
			return 0, err
		}
		fmt.Fprintln(opts.Stdout, containerID)
		return 0, nil
	}

	return s.runInteractive(ctx, containerID, opts)
}

func (s *composeService) runInteractive(ctx context.Context, containerID string, opts api.RunOptions) (int, error) {
	r, err := s.getEscapeKeyProxy(opts.Stdin)
	if err != nil {
		return 0, err
	}

	stdin, stdout, err := s.getContainerStreams(ctx, containerID)
	if err != nil {
		return 0, err
	}

	in := streams.NewIn(opts.Stdin)
	if in.IsTerminal() {
		state, err := term.SetRawTerminal(in.FD())
		if err != nil {
			return 0, err
		}
		defer term.RestoreTerminal(in.FD(), state) //nolint:errcheck
	}

	outputDone := make(chan error)
	inputDone := make(chan error)

	go func() {
		if opts.Tty {
			_, err := io.Copy(opts.Stdout, stdout) //nolint:errcheck
			outputDone <- err
		} else {
			_, err := stdcopy.StdCopy(opts.Stdout, opts.Stderr, stdout) //nolint:errcheck
			outputDone <- err
		}
		stdout.Close() //nolint:errcheck
	}()

	go func() {
		_, err := io.Copy(stdin, r)
		inputDone <- err
		stdin.Close() //nolint:errcheck
	}()

	err = s.apiClient.ContainerStart(ctx, containerID, moby.ContainerStartOptions{})
	if err != nil {
		return 0, err
	}

	s.monitorTTySize(ctx, containerID, s.apiClient.ContainerResize)

	for {
		select {
		case err := <-outputDone:
			return s.terminateRun(ctx, containerID, opts, err)
		case err := <-inputDone:
			if _, ok := err.(term.EscapeError); ok {
				return 0, nil
			}
			if err != nil {
				return 0, err
			}
			// Wait for output to complete streaming
		case <-ctx.Done():
			return 0, ctx.Err()
		}
	}
}

func (s *composeService) terminateRun(ctx context.Context, containerID string, opts api.RunOptions, err error) (int, error) {
	if err != nil {
		return 0, err
	}
	inspect, err := s.apiClient.ContainerInspect(ctx, containerID)
	if err != nil {
		return 0, err
	}
	exitCode := 0
	if inspect.State != nil {
		exitCode = inspect.State.ExitCode
	}
	if opts.AutoRemove {
		err = s.apiClient.ContainerRemove(ctx, containerID, moby.ContainerRemoveOptions{})
	}
	return exitCode, err
}

func (s *composeService) prepareRun(ctx context.Context, project *types.Project, opts api.RunOptions) (string, error) {
	if err := prepareVolumes(project); err != nil { // all dependencies already checked, but might miss service img
		return "", err
	}
	service, err := project.GetService(opts.Service)
	if err != nil {
		return "", err
	}

	applyRunOptions(project, &service, opts)

	slug := stringid.GenerateRandomID()
	if service.ContainerName == "" {
		service.ContainerName = fmt.Sprintf("%s_%s_run_%s", project.Name, service.Name, stringid.TruncateID(slug))
	}
	service.Scale = 1
	service.StdinOpen = true
	service.Restart = ""
	if service.Deploy != nil {
		service.Deploy.RestartPolicy = nil
	}
	service.Labels = service.Labels.Add(api.SlugLabel, slug)
	service.Labels = service.Labels.Add(api.OneoffLabel, "True")

	if err := s.ensureImagesExists(ctx, project, false); err != nil { // all dependencies already checked, but might miss service img
		return "", err
	}
	if !opts.NoDeps {
		if err := s.waitDependencies(ctx, project, service); err != nil {
			return "", err
		}
	}
	created, err := s.createContainer(ctx, project, service, service.ContainerName, 1, opts.Detach && opts.AutoRemove, opts.UseNetworkAliases)
	if err != nil {
		return "", err
	}
	containerID := created.ID
	return containerID, nil
}

func (s *composeService) getEscapeKeyProxy(r io.ReadCloser) (io.ReadCloser, error) {
	var escapeKeys = []byte{16, 17}
	if s.configFile.DetachKeys != "" {
		customEscapeKeys, err := term.ToBytes(s.configFile.DetachKeys)
		if err != nil {
			return nil, err
		}
		escapeKeys = customEscapeKeys
	}
	return ioutils.NewReadCloserWrapper(term.NewEscapeProxy(r, escapeKeys), r.Close), nil
}

func applyRunOptions(project *types.Project, service *types.ServiceConfig, opts api.RunOptions) {
	service.Tty = opts.Tty
	service.ContainerName = opts.Name

	if len(opts.Command) > 0 {
		service.Command = opts.Command
	}
	if len(opts.User) > 0 {
		service.User = opts.User
	}
	if len(opts.WorkingDir) > 0 {
		service.WorkingDir = opts.WorkingDir
	}
	if opts.Entrypoint != nil {
		service.Entrypoint = opts.Entrypoint
	}
	if len(opts.Environment) > 0 {
		env := types.NewMappingWithEquals(opts.Environment)
		projectEnv := env.Resolve(func(s string) (string, bool) {
			v, ok := project.Environment[s]
			return v, ok
		}).RemoveEmpty()
		service.Environment.OverrideBy(projectEnv)
	}
	for k, v := range opts.Labels {
		service.Labels.Add(k, v)
	}
}
<commit_msg>waitForContainer to ensure stopped before remove<commit_after>/*
   Copyright 2020 Docker Compose CLI authors

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
*/

package compose

import (
	"context"
	"fmt"
	"io"

	"github.com/compose-spec/compose-go/types"
	"github.com/docker/cli/cli/streams"
	"github.com/docker/compose/v2/pkg/api"
	moby "github.com/docker/docker/api/types"
	"github.com/docker/docker/api/types/container"
	"github.com/docker/docker/pkg/ioutils"
	"github.com/docker/docker/pkg/stdcopy"
	"github.com/docker/docker/pkg/stringid"
	"github.com/moby/term"
)

func (s *composeService) RunOneOffContainer(ctx context.Context, project *types.Project, opts api.RunOptions) (int, error) {
	containerID, err := s.prepareRun(ctx, project, opts)
	if err != nil {
		return 0, err
	}

	if opts.Detach {
		err := s.apiClient.ContainerStart(ctx, containerID, moby.ContainerStartOptions{})
		if err != nil {
			return 0, err
		}
		fmt.Fprintln(opts.Stdout, containerID)
		return 0, nil
	}

	return s.runInteractive(ctx, containerID, opts)
}

func (s *composeService) runInteractive(ctx context.Context, containerID string, opts api.RunOptions) (int, error) {
	r, err := s.getEscapeKeyProxy(opts.Stdin)
	if err != nil {
		return 0, err
	}

	stdin, stdout, err := s.getContainerStreams(ctx, containerID)
	if err != nil {
		return 0, err
	}

	in := streams.NewIn(opts.Stdin)
	if in.IsTerminal() {
		state, err := term.SetRawTerminal(in.FD())
		if err != nil {
			return 0, err
		}
		defer term.RestoreTerminal(in.FD(), state) //nolint:errcheck
	}

	outputDone := make(chan error)
	inputDone := make(chan error)

	go func() {
		if opts.Tty {
			_, err := io.Copy(opts.Stdout, stdout) //nolint:errcheck
			outputDone <- err
		} else {
			_, err := stdcopy.StdCopy(opts.Stdout, opts.Stderr, stdout) //nolint:errcheck
			outputDone <- err
		}
		stdout.Close() //nolint:errcheck
	}()

	go func() {
		_, err := io.Copy(stdin, r)
		inputDone <- err
		stdin.Close() //nolint:errcheck
	}()

	err = s.apiClient.ContainerStart(ctx, containerID, moby.ContainerStartOptions{})
	if err != nil {
		return 0, err
	}

	s.monitorTTySize(ctx, containerID, s.apiClient.ContainerResize)

	for {
		select {
		case err := <-outputDone:
			if err != nil {
				return 0, err
			}
			return s.terminateRun(ctx, containerID, opts)
		case err := <-inputDone:
			if _, ok := err.(term.EscapeError); ok {
				return 0, nil
			}
			if err != nil {
				return 0, err
			}
			// Wait for output to complete streaming
		case <-ctx.Done():
			return 0, ctx.Err()
		}
	}
}

func (s *composeService) terminateRun(ctx context.Context, containerID string, opts api.RunOptions) (exitCode int, err error) {
	exitCh, errCh := s.apiClient.ContainerWait(ctx, containerID, container.WaitConditionNotRunning)
	select {
	case exit := <-exitCh:
		exitCode = int(exit.StatusCode)
	case err = <-errCh:
		return
	}
	if opts.AutoRemove {
		err = s.apiClient.ContainerRemove(ctx, containerID, moby.ContainerRemoveOptions{})
	}
	return
}

func (s *composeService) prepareRun(ctx context.Context, project *types.Project, opts api.RunOptions) (string, error) {
	if err := prepareVolumes(project); err != nil { // all dependencies already checked, but might miss service img
		return "", err
	}
	service, err := project.GetService(opts.Service)
	if err != nil {
		return "", err
	}

	applyRunOptions(project, &service, opts)

	slug := stringid.GenerateRandomID()
	if service.ContainerName == "" {
		service.ContainerName = fmt.Sprintf("%s_%s_run_%s", project.Name, service.Name, stringid.TruncateID(slug))
	}
	service.Scale = 1
	service.StdinOpen = true
	service.Restart = ""
	if service.Deploy != nil {
		service.Deploy.RestartPolicy = nil
	}
	service.Labels = service.Labels.Add(api.SlugLabel, slug)
	service.Labels = service.Labels.Add(api.OneoffLabel, "True")

	if err := s.ensureImagesExists(ctx, project, false); err != nil { // all dependencies already checked, but might miss service img
		return "", err
	}
	if !opts.NoDeps {
		if err := s.waitDependencies(ctx, project, service); err != nil {
			return "", err
		}
	}
	created, err := s.createContainer(ctx, project, service, service.ContainerName, 1, opts.Detach && opts.AutoRemove, opts.UseNetworkAliases)
	if err != nil {
		return "", err
	}
	containerID := created.ID
	return containerID, nil
}

func (s *composeService) getEscapeKeyProxy(r io.ReadCloser) (io.ReadCloser, error) {
	var escapeKeys = []byte{16, 17}
	if s.configFile.DetachKeys != "" {
		customEscapeKeys, err := term.ToBytes(s.configFile.DetachKeys)
		if err != nil {
			return nil, err
		}
		escapeKeys = customEscapeKeys
	}
	return ioutils.NewReadCloserWrapper(term.NewEscapeProxy(r, escapeKeys), r.Close), nil
}

func applyRunOptions(project *types.Project, service *types.ServiceConfig, opts api.RunOptions) {
	service.Tty = opts.Tty
	service.ContainerName = opts.Name

	if len(opts.Command) > 0 {
		service.Command = opts.Command
	}
	if len(opts.User) > 0 {
		service.User = opts.User
	}
	if len(opts.WorkingDir) > 0 {
		service.WorkingDir = opts.WorkingDir
	}
	if opts.Entrypoint != nil {
		service.Entrypoint = opts.Entrypoint
	}
	if len(opts.Environment) > 0 {
		env := types.NewMappingWithEquals(opts.Environment)
		projectEnv := env.Resolve(func(s string) (string, bool) {
			v, ok := project.Environment[s]
			return v, ok
		}).RemoveEmpty()
		service.Environment.OverrideBy(projectEnv)
	}
	for k, v := range opts.Labels {
		service.Labels.Add(k, v)
	}
}
<|endoftext|>


<commit_before>package helm

import (
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"regexp"
	"strings"

	"github.com/Masterminds/semver"
	"sigs.k8s.io/yaml"
)

// LoadChartfile opens a Chartfile tree
func LoadChartfile(projectRoot string) (*Charts, error) {
	// make sure project root is valid
	abs, err := filepath.Abs(projectRoot)
	if err != nil {
		return nil, err
	}

	// open chartfile
	chartfile := filepath.Join(abs, Filename)
	data, err := ioutil.ReadFile(chartfile)
	if err != nil {
		return nil, err
	}

	// parse it
	c := Chartfile{
		Version:   Version,
		Directory: DefaultDir,
	}
	if err := yaml.UnmarshalStrict(data, &c); err != nil {
		return nil, err
	}

	for i, r := range c.Requires {
		if r.Chart == "" {
			return nil, fmt.Errorf("requirements[%v]: 'chart' must be set", i)
		}
	}

	// return Charts handle
	charts := &Charts{
		Manifest:    c,
		projectRoot: abs,

		// default to ExecHelm, but allow injecting from the outside
		Helm: ExecHelm{},
	}
	return charts, nil
}

// Charts exposes the central Chartfile management functions
type Charts struct {
	// Manifest are the chartfile.yaml contents. It holds data about the developers intentions
	Manifest Chartfile

	// projectRoot is the enclosing directory of chartfile.yaml
	projectRoot string

	// Helm is the helm implementation underneath. ExecHelm is the default, but
	// any implementation of the Helm interface may be used
	Helm Helm
}

// chartManifest represents a Helm chart's Chart.yaml
type chartManifest struct {
	Version string `yaml:"version"`
}

// ChartDir returns the directory pulled charts are saved in
func (c Charts) ChartDir() string {
	return filepath.Join(c.projectRoot, c.Manifest.Directory)
}

// ManifestFile returns the full path to the chartfile.yaml
func (c Charts) ManifestFile() string {
	return filepath.Join(c.projectRoot, Filename)
}

// Vendor pulls all Charts specified in the manifest into the local charts
// directory. It fetches the repository index before doing so.
func (c Charts) Vendor() error {
	dir := c.ChartDir()
	if err := os.MkdirAll(dir, os.ModePerm); err != nil {
		return err
	}

	log.Println("Syncing Repositories ...")
	if err := c.Helm.RepoUpdate(Opts{Repositories: c.Manifest.Repositories}); err != nil {
		return err
	}

	log.Println("Pulling Charts ...")
	for _, r := range c.Manifest.Requires {
		chartName := parseReqName(r.Chart)
		chartPath := filepath.Join(dir, chartName)

		_, err := os.Stat(chartPath)
		if err == nil {
			chartManifestPath := filepath.Join(chartPath, "Chart.yaml")
			chartManifestBytes, err := ioutil.ReadFile(chartManifestPath)
			if err != nil {
				return fmt.Errorf("reading chart manifest: %w", err)
			}
			var chartYAML chartManifest
			if err := yaml.Unmarshal(chartManifestBytes, &chartYAML); err != nil {
				return fmt.Errorf("unmarshalling chart manifest: %w", err)
			}

			if chartYAML.Version == r.Version.String() {
				log.Printf(" %s@%s exists", r.Chart, r.Version.String())
				continue
			} else {
				log.Printf("Removing %s@%s", r.Chart, r.Version.String())
				if err := os.RemoveAll(chartPath); err != nil {
					return err
				}
			}
		} else if !os.IsNotExist(err) {
			return err
		}

		err = c.Helm.Pull(r.Chart, r.Version.String(), PullOpts{
			Destination: dir,
			Opts:        Opts{Repositories: c.Manifest.Repositories},
		})
		if err != nil {
			return err
		}

		log.Printf(" %s@%s downloaded", r.Chart, r.Version.String())
	}

	return nil
}

// Add adds every Chart in reqs to the Manifest after validation, and runs
// Vendor afterwards
func (c *Charts) Add(reqs []string) error {
	log.Printf("Adding %v Charts ...", len(reqs))

	skip := func(s string, err error) {
		log.Printf(" Skipping %s: %s.", s, err)
	}

	// parse new charts, append in memory
	added := 0
	for _, s := range reqs {
		r, err := parseReq(s)
		if err != nil {
			skip(s, err)
			continue
		}

		if c.Manifest.Requires.Has(*r) {
			skip(s, fmt.Errorf("already exists"))
			continue
		}

		c.Manifest.Requires = append(c.Manifest.Requires, *r)
		added++
		log.Println(" OK:", s)
	}

	// write out
	if err := write(c.Manifest, c.ManifestFile()); err != nil {
		return err
	}

	// skipped some? fail then
	if added != len(reqs) {
		return fmt.Errorf("%v Charts were skipped. Please check above logs for details", len(reqs)-added)
	}

	// worked fine? vendor it
	log.Printf("Added %v Charts to helmfile.yaml. Vendoring ...", added)
	return c.Vendor()
}

func (c *Charts) AddRepos(repos ...Repo) error {
	added := 0
	for _, r := range repos {
		if c.Manifest.Repositories.Has(r) {
			log.Printf("Skipping %s: Already exists", r.Name)
			continue
		}

		c.Manifest.Repositories = append(c.Manifest.Repositories, r)
	}

	// write out
	if err := write(c.Manifest, c.ManifestFile()); err != nil {
		return err
	}

	if added != len(repos) {
		return fmt.Errorf("%v Chart(s) were skipped. Please check above logs for details", len(repos)-added)
	}

	return nil
}

func InitChartfile(path string) (*Charts, error) {
	c := Chartfile{
		Version: Version,
		Repositories: []Repo{{
			Name: "stable",
			URL:  "https://charts.helm.sh/stable",
		}},
		Requires: make(Requirements, 0),
	}

	if err := write(c, path); err != nil {
		return nil, err
	}

	return LoadChartfile(filepath.Dir(path))
}

// write saves a Chartfile to dest
func write(c Chartfile, dest string) error {
	data, err := yaml.Marshal(c)
	if err != nil {
		return err
	}

	return ioutil.WriteFile(dest, data, 0644)
}

var chartExp = regexp.MustCompile(`\w+\/.+@.+`)

// parseReq parses a requirement from a string of the format `repo/name@version`
func parseReq(s string) (*Requirement, error) {
	if !chartExp.MatchString(s) {
		return nil, fmt.Errorf("not of form 'repo/chart@version'")
	}

	elems := strings.Split(s, "@")
	chart := elems[0]
	ver, err := semver.NewVersion(elems[1])
	if errors.Is(err, semver.ErrInvalidSemVer) {
		return nil, fmt.Errorf("version is invalid")
	} else if err != nil {
		return nil, fmt.Errorf("version is invalid: %s", err)
	}

	return &Requirement{
		Chart:   chart,
		Version: *ver,
	}, nil
}

// parseReqName parses a name from a string of the format `repo/name`
func parseReqName(s string) string {
	elems := strings.Split(s, "/")
	name := elems[1]
	return name
}
<commit_msg>feat(tool/charts): Only update helm repositories when necessary (#535)<commit_after>package helm

import (
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"regexp"
	"strings"

	"github.com/Masterminds/semver"
	"sigs.k8s.io/yaml"
)

// LoadChartfile opens a Chartfile tree
func LoadChartfile(projectRoot string) (*Charts, error) {
	// make sure project root is valid
	abs, err := filepath.Abs(projectRoot)
	if err != nil {
		return nil, err
	}

	// open chartfile
	chartfile := filepath.Join(abs, Filename)
	data, err := ioutil.ReadFile(chartfile)
	if err != nil {
		return nil, err
	}

	// parse it
	c := Chartfile{
		Version:   Version,
		Directory: DefaultDir,
	}
	if err := yaml.UnmarshalStrict(data, &c); err != nil {
		return nil, err
	}

	for i, r := range c.Requires {
		if r.Chart == "" {
			return nil, fmt.Errorf("requirements[%v]: 'chart' must be set", i)
		}
	}

	// return Charts handle
	charts := &Charts{
		Manifest:    c,
		projectRoot: abs,

		// default to ExecHelm, but allow injecting from the outside
		Helm: ExecHelm{},
	}
	return charts, nil
}

// Charts exposes the central Chartfile management functions
type Charts struct {
	// Manifest are the chartfile.yaml contents. It holds data about the developers intentions
	Manifest Chartfile

	// projectRoot is the enclosing directory of chartfile.yaml
	projectRoot string

	// Helm is the helm implementation underneath. ExecHelm is the default, but
	// any implementation of the Helm interface may be used
	Helm Helm
}

// chartManifest represents a Helm chart's Chart.yaml
type chartManifest struct {
	Version string `yaml:"version"`
}

// ChartDir returns the directory pulled charts are saved in
func (c Charts) ChartDir() string {
	return filepath.Join(c.projectRoot, c.Manifest.Directory)
}

// ManifestFile returns the full path to the chartfile.yaml
func (c Charts) ManifestFile() string {
	return filepath.Join(c.projectRoot, Filename)
}

// Vendor pulls all Charts specified in the manifest into the local charts
// directory. It fetches the repository index before doing so.
func (c Charts) Vendor() error {
	dir := c.ChartDir()
	if err := os.MkdirAll(dir, os.ModePerm); err != nil {
		return err
	}

	repositoriesUpdated := false
	log.Println("Pulling Charts ...")
	for _, r := range c.Manifest.Requires {
		chartName := parseReqName(r.Chart)
		chartPath := filepath.Join(dir, chartName)

		_, err := os.Stat(chartPath)
		if err == nil {
			chartManifestPath := filepath.Join(chartPath, "Chart.yaml")
			chartManifestBytes, err := ioutil.ReadFile(chartManifestPath)
			if err != nil {
				return fmt.Errorf("reading chart manifest: %w", err)
			}
			var chartYAML chartManifest
			if err := yaml.Unmarshal(chartManifestBytes, &chartYAML); err != nil {
				return fmt.Errorf("unmarshalling chart manifest: %w", err)
			}

			if chartYAML.Version == r.Version.String() {
				log.Printf(" %s@%s exists", r.Chart, r.Version.String())
				continue
			} else {
				log.Printf("Removing %s@%s", r.Chart, r.Version.String())
				if err := os.RemoveAll(chartPath); err != nil {
					return err
				}
			}
		} else if !os.IsNotExist(err) {
			return err
		}

		if !repositoriesUpdated {
			log.Println("Syncing Repositories ...")
			if err := c.Helm.RepoUpdate(Opts{Repositories: c.Manifest.Repositories}); err != nil {
				return err
			}
			repositoriesUpdated = true
		}
		err = c.Helm.Pull(r.Chart, r.Version.String(), PullOpts{
			Destination: dir,
			Opts:        Opts{Repositories: c.Manifest.Repositories},
		})
		if err != nil {
			return err
		}

		log.Printf(" %s@%s downloaded", r.Chart, r.Version.String())
	}

	return nil
}

// Add adds every Chart in reqs to the Manifest after validation, and runs
// Vendor afterwards
func (c *Charts) Add(reqs []string) error {
	log.Printf("Adding %v Charts ...", len(reqs))

	skip := func(s string, err error) {
		log.Printf(" Skipping %s: %s.", s, err)
	}

	// parse new charts, append in memory
	added := 0
	for _, s := range reqs {
		r, err := parseReq(s)
		if err != nil {
			skip(s, err)
			continue
		}

		if c.Manifest.Requires.Has(*r) {
			skip(s, fmt.Errorf("already exists"))
			continue
		}

		c.Manifest.Requires = append(c.Manifest.Requires, *r)
		added++
		log.Println(" OK:", s)
	}

	// write out
	if err := write(c.Manifest, c.ManifestFile()); err != nil {
		return err
	}

	// skipped some? fail then
	if added != len(reqs) {
		return fmt.Errorf("%v Charts were skipped. Please check above logs for details", len(reqs)-added)
	}

	// worked fine? vendor it
	log.Printf("Added %v Charts to helmfile.yaml. Vendoring ...", added)
	return c.Vendor()
}

func (c *Charts) AddRepos(repos ...Repo) error {
	added := 0
	for _, r := range repos {
		if c.Manifest.Repositories.Has(r) {
			log.Printf("Skipping %s: Already exists", r.Name)
			continue
		}

		c.Manifest.Repositories = append(c.Manifest.Repositories, r)
	}

	// write out
	if err := write(c.Manifest, c.ManifestFile()); err != nil {
		return err
	}

	if added != len(repos) {
		return fmt.Errorf("%v Chart(s) were skipped. Please check above logs for details", len(repos)-added)
	}

	return nil
}

func InitChartfile(path string) (*Charts, error) {
	c := Chartfile{
		Version: Version,
		Repositories: []Repo{{
			Name: "stable",
			URL:  "https://charts.helm.sh/stable",
		}},
		Requires: make(Requirements, 0),
	}

	if err := write(c, path); err != nil {
		return nil, err
	}

	return LoadChartfile(filepath.Dir(path))
}

// write saves a Chartfile to dest
func write(c Chartfile, dest string) error {
	data, err := yaml.Marshal(c)
	if err != nil {
		return err
	}

	return ioutil.WriteFile(dest, data, 0644)
}

var chartExp = regexp.MustCompile(`\w+\/.+@.+`)

// parseReq parses a requirement from a string of the format `repo/name@version`
func parseReq(s string) (*Requirement, error) {
	if !chartExp.MatchString(s) {
		return nil, fmt.Errorf("not of form 'repo/chart@version'")
	}

	elems := strings.Split(s, "@")
	chart := elems[0]
	ver, err := semver.NewVersion(elems[1])
	if errors.Is(err, semver.ErrInvalidSemVer) {
		return nil, fmt.Errorf("version is invalid")
	} else if err != nil {
		return nil, fmt.Errorf("version is invalid: %s", err)
	}

	return &Requirement{
		Chart:   chart,
		Version: *ver,
	}, nil
}

// parseReqName parses a name from a string of the format `repo/name`
func parseReqName(s string) string {
	elems := strings.Split(s, "/")
	name := elems[1]
	return name
}
<|endoftext|>


<commit_before>// Copyright (C) 2017 File Maps Backend Authors.
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

package httpd

import (
	"encoding/json"
	log "github.com/Sirupsen/logrus"
	"github.com/julienschmidt/httprouter"
	"net"
	"net/http"
	"strings"

	"github.com/filemaps/filemaps/pkg/config"
	"github.com/filemaps/filemaps/pkg/model"
)

var (
	// Value of CORS header Access-Control-Allow-Origin
	CorsAllow string
)

// RunHTTP starts HTTP server
func RunHTTP(addr string, webUIPath string) {
	router := httprouter.New()
	route(router, webUIPath)
	handler := authMiddleware(router)
	log.WithFields(log.Fields{
		"transport": "HTTP",
		"addr":      addr,
	}).Info("Starting server")
	log.Fatal(http.ListenAndServe(addr, handler))
}

// WriteJSON writes JSON response
func WriteJSON(w http.ResponseWriter, v interface{}) error {
	b, err := json.Marshal(v)
	if err != nil {
		return err
	}

	w.Header().Set("Content-Type", "application/json; charset=utf-8")

	// CORS header Access-Control-Allow-Origin for development
	if CorsAllow != "" {
		w.Header().Set("Access-Control-Allow-Origin", CorsAllow)
	}

	w.Write(b)
	return nil
}

// WriteJSONError writes error JSON response
func WriteJSONError(w http.ResponseWriter, code int, err string) {
	w.WriteHeader(code)
	WriteJSON(w, map[string]string{
		"error": err,
	})
}

// authMiddleware authenticates the request.
// Request must come from trusted address or X-API-Key header must
// contain a valid API key.
func authMiddleware(handler http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		model.GetAPIKeyManager()
		if addrIsTrusted(r.RemoteAddr) {
			handler.ServeHTTP(w, r)
			return
		}

		if model.GetAPIKeyManager().IsValidAPIKey(r.Header.Get("X-API-Key")) {
			handler.ServeHTTP(w, r)
			return
		}

		log.WithFields(log.Fields{
			"requestURI": r.RequestURI,
			"remoteAddr": r.RemoteAddr,
		}).Error("Access denied")
		w.WriteHeader(403)
	})
}

// addrIsTrusted returns true if given address is trusted.
// addr is request.RemoteAddr which has format IP:port
func addrIsTrusted(addr string) bool {
	// strip port
	addr = addr[:strings.LastIndex(addr, ":")]
	// remove square brackets from ipv6 addr
	addr = strings.Replace(addr, "[", "", -1)
	addr = strings.Replace(addr, "]", "", -1)
	remoteIP := net.ParseIP(addr)
	if remoteIP == nil {
		log.WithFields(log.Fields{
			"ip": addr,
		}).Error("Could not parse remote IP")
		return false
	}

	// trust loopback addresses
	if remoteIP.IsLoopback() {
		return true
	}

	// check trusted addresses from config
	cfg := config.GetConfiguration()
	addrs := strings.Split(cfg.TrustedAddresses, ",")
	for _, a := range addrs {
		trustedIP := net.ParseIP(a)
		if trustedIP != nil && trustedIP.Equal(remoteIP) {
			return true
		}
	}

	return false
}
<commit_msg>CORS middleware for Access-Control-Allow-Origin header<commit_after>// Copyright (C) 2017 File Maps Backend Authors.
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

package httpd

import (
	"encoding/json"
	log "github.com/Sirupsen/logrus"
	"github.com/julienschmidt/httprouter"
	"net"
	"net/http"
	"strings"

	"github.com/filemaps/filemaps/pkg/config"
	"github.com/filemaps/filemaps/pkg/model"
)

var (
	// Value of CORS header Access-Control-Allow-Origin
	CorsAllow string
)

// RunHTTP starts HTTP server
func RunHTTP(addr string, webUIPath string) {
	router := httprouter.New()
	route(router, webUIPath)
	handler := corsMiddleware(router)
	handler = authMiddleware(handler)
	log.WithFields(log.Fields{
		"transport": "HTTP",
		"addr":      addr,
	}).Info("Starting server")
	log.Fatal(http.ListenAndServe(addr, handler))
}

// WriteJSON writes JSON response
func WriteJSON(w http.ResponseWriter, v interface{}) error {
	b, err := json.Marshal(v)
	if err != nil {
		return err
	}

	w.Header().Set("Content-Type", "application/json; charset=utf-8")
	w.Write(b)
	return nil
}

// WriteJSONError writes error JSON response
func WriteJSONError(w http.ResponseWriter, code int, err string) {
	w.WriteHeader(code)
	WriteJSON(w, map[string]string{
		"error": err,
	})
}

// corsMiddleware adds Access-Control-Allow-Origin header for CORS.
func corsMiddleware(handler http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if CorsAllow != "" {
			w.Header().Set("Access-Control-Allow-Origin", CorsAllow)
		}
		handler.ServeHTTP(w, r)
		return
	})
}

// authMiddleware authenticates the request.
// Request must come from trusted address or X-API-Key header must
// contain a valid API key.
func authMiddleware(handler http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		model.GetAPIKeyManager()
		if addrIsTrusted(r.RemoteAddr) {
			handler.ServeHTTP(w, r)
			return
		}

		if model.GetAPIKeyManager().IsValidAPIKey(r.Header.Get("X-API-Key")) {
			handler.ServeHTTP(w, r)
			return
		}

		log.WithFields(log.Fields{
			"requestURI": r.RequestURI,
			"remoteAddr": r.RemoteAddr,
		}).Error("Access denied")
		w.WriteHeader(403)
	})
}

// addrIsTrusted returns true if given address is trusted.
// addr is request.RemoteAddr which has format IP:port
func addrIsTrusted(addr string) bool {
	// strip port
	addr = addr[:strings.LastIndex(addr, ":")]
	// remove square brackets from ipv6 addr
	addr = strings.Replace(addr, "[", "", -1)
	addr = strings.Replace(addr, "]", "", -1)
	remoteIP := net.ParseIP(addr)
	if remoteIP == nil {
		log.WithFields(log.Fields{
			"ip": addr,
		}).Error("Could not parse remote IP")
		return false
	}

	// trust loopback addresses
	if remoteIP.IsLoopback() {
		return true
	}

	// check trusted addresses from config
	cfg := config.GetConfiguration()
	addrs := strings.Split(cfg.TrustedAddresses, ",")
	for _, a := range addrs {
		trustedIP := net.ParseIP(a)
		if trustedIP != nil && trustedIP.Equal(remoteIP) {
			return true
		}
	}

	return false
}

func writeCORSHeaders(w http.ResponseWriter) {
	// CORS header Access-Control-Allow-Origin for development
	if CorsAllow != "" {
		w.Header().Set("Access-Control-Allow-Origin", CorsAllow)
	}
}
<|endoftext|>