<commit_before>package movie

import (
	"context"
	"time"

	"github.com/gilcrest/errs"
	"github.com/gilcrest/go-api-basic/domain/random"
	"github.com/google/uuid"
)

// Movie holds details of a movie
type Movie struct {
	ID              uuid.UUID
	ExtlID          string
	Title           string
	Year            int
	Rated           string
	Released        time.Time
	RunTime         int
	Director        string
	Writer          string
	CreateTimestamp time.Time
	UpdateTimestamp time.Time
}

// Validate does basic input validation and ensures the struct is
// properly constructed
func (m *Movie) Validate() error {
	const op errs.Op = "domain/Movie.validate"

	switch {
	case m.Title == "":
		return errs.E(op, errs.Validation, errs.Parameter("Title"), errs.MissingField("Title"))
	case m.Year < 1878:
		return errs.E(op, errs.Validation, errs.Parameter("Year"), "The first film was in 1878, Year must be >= 1878")
	case m.Rated == "":
		return errs.E(op, errs.Validation, errs.Parameter("Rated"), errs.MissingField("Rated"))
	case m.Released.IsZero() == true:
		return errs.E(op, errs.Validation, errs.Parameter("ReleaseDate"), "Released must have a value")
	case m.RunTime <= 0:
		return errs.E(op, errs.Validation, errs.Parameter("RunTime"), "Run time must be greater than zero")
	case m.Director == "":
		return errs.E(op, errs.Validation, errs.Parameter("Director"), errs.MissingField("Director"))
	case m.Writer == "":
		return errs.E(op, errs.Validation, errs.Parameter("Writer"), errs.MissingField("Writer"))
	}

	return nil
}

// Add performs business validations prior to writing to the db
func (m *Movie) Add(ctx context.Context) error {
	const op errs.Op = "movie/Movie.Add"

	m.ID = uuid.New()
	id, err := random.CryptoString(15)
	if err != nil {
		return errs.E(op, errs.Validation, errs.Internal, err)
	}
	m.ExtlID = id

	return nil
}

// Update performs business validations prior to writing to the db
func (m *Movie) Update(ctx context.Context, id string) error {
	const op errs.Op = "movie/Movie.Update"

	m.ExtlID = id
	m.UpdateTimestamp = time.Now().UTC()

	return nil
}
<commit_msg>Field name change for readability<commit_after>package movie

import (
	"context"
	"time"

	"github.com/gilcrest/errs"
	"github.com/gilcrest/go-api-basic/domain/random"
	"github.com/google/uuid"
)

// Movie holds details of a movie
type Movie struct {
	ID              uuid.UUID
	ExternalID      string
	Title           string
	Year            int
	Rated           string
	Released        time.Time
	RunTime         int
	Director        string
	Writer          string
	CreateTimestamp time.Time
	UpdateTimestamp time.Time
}

// Validate does basic input validation and ensures the struct is
// properly constructed
func (m *Movie) Validate() error {
	const op errs.Op = "domain/Movie.validate"

	switch {
	case m.Title == "":
		return errs.E(op, errs.Validation, errs.Parameter("Title"), errs.MissingField("Title"))
	case m.Year < 1878:
		return errs.E(op, errs.Validation, errs.Parameter("Year"), "The first film was in 1878, Year must be >= 1878")
	case m.Rated == "":
		return errs.E(op, errs.Validation, errs.Parameter("Rated"), errs.MissingField("Rated"))
	case m.Released.IsZero() == true:
		return errs.E(op, errs.Validation, errs.Parameter("ReleaseDate"), "Released must have a value")
	case m.RunTime <= 0:
		return errs.E(op, errs.Validation, errs.Parameter("RunTime"), "Run time must be greater than zero")
	case m.Director == "":
		return errs.E(op, errs.Validation, errs.Parameter("Director"), errs.MissingField("Director"))
	case m.Writer == "":
		return errs.E(op, errs.Validation, errs.Parameter("Writer"), errs.MissingField("Writer"))
	}

	return nil
}

// Add performs business validations prior to writing to the db
func (m *Movie) Add(ctx context.Context) error {
	const op errs.Op = "movie/Movie.Add"

	m.ID = uuid.New()
	id, err := random.CryptoString(15)
	if err != nil {
		return errs.E(op, errs.Validation, errs.Internal, err)
	}
	m.ExternalID = id

	return nil
}

// Update performs business validations prior to writing to the db
func (m *Movie) Update(ctx context.Context, id string) error {
	const op errs.Op = "movie/Movie.Update"

	m.ExternalID = id
	m.UpdateTimestamp = time.Now().UTC()

	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"

	"net"
	"reflect"

	log "github.com/sirupsen/logrus"
	"github.com/gonuts/cbor"
	"github.com/nttdots/go-dots/coap"
	"github.com/nttdots/go-dots/dots_common"
	"github.com/nttdots/go-dots/dots_common/messages"
	"github.com/nttdots/go-dots/dots_server/controllers"
	"github.com/nttdots/go-dots/dots_server/models"
	dtls "github.com/nttdots/go-dtls"
)

type ControllerInfo struct {
	Controller         controllers.ControllerInterface
	RequestMessageType reflect.Type
}

type ControllerInfoMap map[string]ControllerInfo
type DotsServiceMethod func(request interface{}, customer *models.Customer) (controllers.Response, error)

/*
 * Router struct invokes appropriate API controllers based on request-uris.
 */
type Router struct {
	ControllerMap map[string]ControllerInfo
}

func NewRouter() *Router {
	r := new(Router)
	r.ControllerMap = make(ControllerInfoMap)

	return r
}

/*
 * Register an API route based on the message code.
 */
func (r *Router) Register(code messages.Code, controller controllers.ControllerInterface) {
	messageType := messages.MessageTypes[code]
	r.ControllerMap[messageType.Path] = ControllerInfo{
		Controller:         controller,
		RequestMessageType: messageType.Type,
	}
}

/*
 * Obtain the corresponding API controller to the request.
 */
func (r *Router) getMethod(controller controllers.ControllerInterface, request *coap.Message) (DotsServiceMethod, error) {
	var code coap.COAPCode = request.Code
	var method DotsServiceMethod = nil

	switch code {
	case coap.GET:
		method = controller.Get
	case coap.POST:
		method = controller.Post
	case coap.PUT:
		method = controller.Put
	case coap.DELETE:
		method = controller.Delete
	}
	if method == nil {
		e := fmt.Sprintf("Unknonw COAPCode for Method %d:", code)
		return method, errors.New(e)
	}
	return method, nil
}

/*
 * Unmarshals request-body JSONs.
 *
 * parameter:
 *  request: CoAP request
 *  messageType: Type to unmarshal
 * return:
 *  1: Object unmarshaled
 *  2: error
 */
func (r *Router) loadJson(request *coap.Message, messageType reflect.Type) (interface{}, error) {

	m := reflect.New(messageType).Interface()
	err := json.Unmarshal(request.Payload, &m)

	return m, err
}

/*
 * Unmarshals request-body CBORs.
 *
 * parameter:
 *  request: CoAP request
 *  messageType: Type to unmarshal
 * return:
 *  1: Object unmarshaled
 *  2: error
 */
func (r *Router) UnmarshalCbor(request *coap.Message, messageType reflect.Type) (interface{}, error) {

	if len(request.Payload) == 0 {
		return nil, nil
	}

	m := reflect.New(messageType).Interface()
	cborReader := bytes.NewReader(request.Payload)

	d := cbor.NewDecoder(cborReader)
	err := d.Decode(m)

	return m, err
}

/*
 * Convert to CBOR format.
 *
 * parameter:
 *  message: Object to be encoded.
 * return:
 *  1: CBOR message
 *  2: error
*/
func (r *Router) MarshalCbor(message interface{}) ([]byte, error) {

	cborWriter := bytes.NewBuffer(nil)
	e := cbor.NewEncoder(cborWriter)

	err := e.Encode(message)

	return cborWriter.Bytes(), err

}

func (r *Router) createResponse(request *coap.Message, controllerResponse []byte,
	responseType dots_common.Type, responseCode dots_common.Code) *coap.Message {
	var result *coap.Message = nil

	result = &coap.Message{
		Type:      responseType.CoAPType(),
		Code:      responseCode.CoAPCode(),
		MessageID: request.MessageID,
		Token:     request.Token,
		Payload:   controllerResponse,
	}
	result.SetOption(coap.ContentFormat, coap.AppCbor)

	return result
}

func (r *Router) callController(request *coap.Message, customer *models.Customer) *coap.Message {

	controllerInfo, ok := r.ControllerMap[request.PathString()]
	if !ok {
		return r.createResponse(request, nil, dots_common.NonConfirmable, dots_common.MethodNotAllowed)
	}
	log.WithFields(log.Fields{
		"controller": controllerInfo,
	}).Debug("controller decided.")

	method, err := r.getMethod(controllerInfo.Controller, request)
	log.WithFields(log.Fields{
		"method": method,
	}).Debug("method decided.")

	requestStructure, err := r.UnmarshalCbor(request, controllerInfo.RequestMessageType)
	if err != nil {
		return r.createResponse(request, nil, dots_common.NonConfirmable, dots_common.InternalServerError)
	}

	log.WithFields(log.Fields{
		"params": fmt.Sprintf("%+v", requestStructure),
	}).Debug("call controller method with message.")
	result, err := method(requestStructure, customer)

	if err != nil {
		switch e := err.(type) {
		case controllers.Error:
			responseCbor, cborErr := r.MarshalCbor(e.Body)
			if cborErr != nil {
				responseCbor = nil
			}
			return r.createResponse(request, responseCbor, e.Type, e.Code)

		case error:
			return r.createResponse(request, nil, dots_common.NonConfirmable, dots_common.InternalServerError)
		}
	}

	responseCbor, cborErr := r.MarshalCbor(result.Body)
	if cborErr != nil {
		return r.createResponse(request, nil, dots_common.NonConfirmable, dots_common.InternalServerError)
	}
	return r.createResponse(request, responseCbor, result.Type, result.Code)

}

/*
 * Receive CoAP messages
 *  1. Identify the message source customer.
 *  2. Invoke the appropriate API controller.
 *
 * parameter:
 *  l: connection object to the dots client.
 *  a: client IP address
 *  request: CoAP request message
 * return:
 *  1: CoAP request message
*/
func (r *Router) Serve(l net.Conn, a net.Addr, request *coap.Message) *coap.Message {
	log.WithFields(log.Fields{
		"path":      request.PathString(),
		"from":      a,
		"messageId": request.MessageID,
	}).Info("Got message")

	conn, ok := l.(dtls.DTLSServerConn)
	if !ok {
		return r.createResponse(request, nil, dots_common.NonConfirmable, dots_common.InternalServerError)
	}

	commonName := conn.GetClientCN()
	if commonName == "" {
		log.Errorln("Not found CommonName")
		return r.createResponse(request, nil, dots_common.NonConfirmable, dots_common.Forbidden)
	}

	customer, err := models.GetCustomerByCommonName(commonName)
	if err != nil || customer.Id == 0 {
		log.WithFields(log.Fields{
			"common-name": commonName,
		}).Error("client does not exist.")
		return r.createResponse(request, nil, dots_common.NonConfirmable, dots_common.Forbidden)
	}

	log.WithFields(log.Fields{
		"customer.id":   customer.Id,
		"customer.name": customer.Name,
	}).Debug("find client.")
	log.Debug(CoapHeaderDisplay(request))

	return r.callController(request, customer)
}
<commit_msg>fixed cbor mapping on dots_server<commit_after>package main

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"

	"net"
	"reflect"

	log "github.com/sirupsen/logrus"
	"github.com/gonuts/cbor"
	"github.com/ugorji/go/codec"
	"github.com/nttdots/go-dots/coap"
	"github.com/nttdots/go-dots/dots_common"
	"github.com/nttdots/go-dots/dots_common/messages"
	"github.com/nttdots/go-dots/dots_server/controllers"
	"github.com/nttdots/go-dots/dots_server/models"
	dtls "github.com/nttdots/go-dtls"
)

type ControllerInfo struct {
	Controller         controllers.ControllerInterface
	RequestMessageType reflect.Type
}

type ControllerInfoMap map[string]ControllerInfo
type DotsServiceMethod func(request interface{}, customer *models.Customer) (controllers.Response, error)

/*
 * Router struct invokes appropriate API controllers based on request-uris.
 */
type Router struct {
	ControllerMap map[string]ControllerInfo
}

func NewRouter() *Router {
	r := new(Router)
	r.ControllerMap = make(ControllerInfoMap)

	return r
}

/*
 * Register an API route based on the message code.
 */
func (r *Router) Register(code messages.Code, controller controllers.ControllerInterface) {
	messageType := messages.MessageTypes[code]
	r.ControllerMap[messageType.Path] = ControllerInfo{
		Controller:         controller,
		RequestMessageType: messageType.Type,
	}
}

/*
 * Obtain the corresponding API controller to the request.
 */
func (r *Router) getMethod(controller controllers.ControllerInterface, request *coap.Message) (DotsServiceMethod, error) {
	var code coap.COAPCode = request.Code
	var method DotsServiceMethod = nil

	switch code {
	case coap.GET:
		method = controller.Get
	case coap.POST:
		method = controller.Post
	case coap.PUT:
		method = controller.Put
	case coap.DELETE:
		method = controller.Delete
	}
	if method == nil {
		e := fmt.Sprintf("Unknonw COAPCode for Method %d:", code)
		return method, errors.New(e)
	}
	return method, nil
}

/*
 * Unmarshals request-body JSONs.
 *
 * parameter:
 *  request: CoAP request
 *  messageType: Type to unmarshal
 * return:
 *  1: Object unmarshaled
 *  2: error
 */
func (r *Router) loadJson(request *coap.Message, messageType reflect.Type) (interface{}, error) {

	m := reflect.New(messageType).Interface()
	err := json.Unmarshal(request.Payload, &m)

	return m, err
}

/*
 * Unmarshals request-body CBORs.
 *
 * parameter:
 *  request: CoAP request
 *  messageType: Type to unmarshal
 * return:
 *  1: Object unmarshaled
 *  2: error
 */
func (r *Router) UnmarshalCbor(request *coap.Message, messageType reflect.Type) (interface{}, error) {

	if len(request.Payload) == 0 {
		return nil, nil
	}

	m := reflect.New(messageType).Interface()
	cborReader := bytes.NewReader(request.Payload)

	cborDecHandle := new(codec.CborHandle)
	cborDecHandle.SetUseIntElmOfStruct(true)
	d := codec.NewDecoder(cborReader, cborDecHandle)
	err := d.Decode(m)

	return m, err
}

/*
 * Convert to CBOR format.
 *
 * parameter:
 *  message: Object to be encoded.
 * return:
 *  1: CBOR message
 *  2: error
*/
func (r *Router) MarshalCbor(message interface{}) ([]byte, error) {

	cborWriter := bytes.NewBuffer(nil)
	e := cbor.NewEncoder(cborWriter)

	err := e.Encode(message)

	return cborWriter.Bytes(), err

}

func (r *Router) createResponse(request *coap.Message, controllerResponse []byte,
	responseType dots_common.Type, responseCode dots_common.Code) *coap.Message {
	var result *coap.Message = nil

	result = &coap.Message{
		Type:      responseType.CoAPType(),
		Code:      responseCode.CoAPCode(),
		MessageID: request.MessageID,
		Token:     request.Token,
		Payload:   controllerResponse,
	}
	result.SetOption(coap.ContentFormat, coap.AppCbor)

	return result
}

func (r *Router) callController(request *coap.Message, customer *models.Customer) *coap.Message {

	controllerInfo, ok := r.ControllerMap[request.PathString()]
	log.Debugf("callController -in. path=%s, ok=%v\n", request.PathString(), ok)
	if !ok {
		return r.createResponse(request, nil, dots_common.NonConfirmable, dots_common.MethodNotAllowed)
	}
	log.WithFields(log.Fields{
		"controller": controllerInfo,
	}).Debug("controller decided.")

	method, err := r.getMethod(controllerInfo.Controller, request)
	log.WithFields(log.Fields{
		"method": method,
	}).Debug("method decided.")

	requestStructure, err := r.UnmarshalCbor(request, controllerInfo.RequestMessageType)
	if err != nil {
		return r.createResponse(request, nil, dots_common.NonConfirmable, dots_common.InternalServerError)
	}

	log.WithFields(log.Fields{
		"params": fmt.Sprintf("%+v", requestStructure),
	}).Debug("call controller method with message.")
	result, err := method(requestStructure, customer)

	if err != nil {
		switch e := err.(type) {
		case controllers.Error:
			responseCbor, cborErr := r.MarshalCbor(e.Body)
			if cborErr != nil {
				responseCbor = nil
			}
			return r.createResponse(request, responseCbor, e.Type, e.Code)

		case error:
			return r.createResponse(request, nil, dots_common.NonConfirmable, dots_common.InternalServerError)
		}
	}

	responseCbor, cborErr := r.MarshalCbor(result.Body)
	if cborErr != nil {
		return r.createResponse(request, nil, dots_common.NonConfirmable, dots_common.InternalServerError)
	}
	return r.createResponse(request, responseCbor, result.Type, result.Code)

}

/*
 * Receive CoAP messages
 *  1. Identify the message source customer.
 *  2. Invoke the appropriate API controller.
 *
 * parameter:
 *  l: connection object to the dots client.
 *  a: client IP address
 *  request: CoAP request message
 * return:
 *  1: CoAP request message
*/
func (r *Router) Serve(l net.Conn, a net.Addr, request *coap.Message) *coap.Message {
	log.WithFields(log.Fields{
		"path":      request.PathString(),
		"from":      a,
		"messageId": request.MessageID,
	}).Info("Got message")

	conn, ok := l.(dtls.DTLSServerConn)
	if !ok {
		return r.createResponse(request, nil, dots_common.NonConfirmable, dots_common.InternalServerError)
	}

	commonName := conn.GetClientCN()
	if commonName == "" {
		log.Errorln("Not found CommonName")
		return r.createResponse(request, nil, dots_common.NonConfirmable, dots_common.Forbidden)
	}

	customer, err := models.GetCustomerByCommonName(commonName)
	if err != nil || customer.Id == 0 {
		log.WithFields(log.Fields{
			"common-name": commonName,
		}).Error("client does not exist.")
		return r.createResponse(request, nil, dots_common.NonConfirmable, dots_common.Forbidden)
	}

	log.WithFields(log.Fields{
		"customer.id":   customer.Id,
		"customer.name": customer.Name,
	}).Debug("find client.")
	log.Debug(CoapHeaderDisplay(request))

	return r.callController(request, customer)
}
<|endoftext|>


<commit_before>package editor

import (
	"bytes"
	"strings"
	"testing"

	"github.com/elpinal/coco3/config"
)

type testScreen struct {
}

func (ts *testScreen) Refresh(prompt string, s []rune, pos int) {
}

func (ts *testScreen) SetLastLine(msg string) {
}

func TestEditor(t *testing.T) {
	inBuf := strings.NewReader("aaa" + string(CharCtrlM))
	var outBuf, errBuf bytes.Buffer
	e := New(&testScreen{}, &config.Config{}, inBuf, &outBuf, &errBuf)
	s, err := e.Read()
	if err != nil {
		t.Error(err)
	}
	if want := "aaa"; string(s) != want {
		t.Errorf("got %q, want %q", string(s), want)
	}
	e.Clear()
}
<commit_msg>Check if output and error are empty<commit_after>package editor

import (
	"bytes"
	"strings"
	"testing"

	"github.com/elpinal/coco3/config"
)

type testScreen struct {
}

func (ts *testScreen) Refresh(prompt string, s []rune, pos int) {
}

func (ts *testScreen) SetLastLine(msg string) {
}

func TestEditor(t *testing.T) {
	inBuf := strings.NewReader("aaa" + string(CharCtrlM))
	var outBuf, errBuf bytes.Buffer
	e := New(&testScreen{}, &config.Config{}, inBuf, &outBuf, &errBuf)
	s, err := e.Read()
	if err != nil {
		t.Error(err)
	}
	if want := "aaa"; string(s) != want {
		t.Errorf("got %q, want %q", string(s), want)
	}
	if got := outBuf.String(); got != "" {
		t.Errorf("got %q, want %q", got, "")
	}
	if got := errBuf.String(); got != "" {
		t.Errorf("got %q, want %q", got, "")
	}
	e.Clear()
}
<|endoftext|>


<commit_before>package elastic

import (
	"context"

	"bitbucket.org/firstrow/logvoyage/models"
	"bitbucket.org/firstrow/logvoyage/shared/config"
	"gopkg.in/olivere/elastic.v5"
)

// LogRecord fetched from storage
type LogRecord struct {
	Source   string
	Datetime int64
}

// SearchLogsResult contains logs and total number of records in storage.
type SearchLogsResult struct {
	Logs  []string `json:"logs"`
	Total int64    `json:"total"`
}

// SearchLogs sends query to elastic search index
// types - ealstic types to search on.
// queryString - user provided data.
func SearchLogs(user *models.User, project *models.Project, types []string, queryString string, page int) (SearchLogsResult, error) {
	ctx := context.Background()
	es, _ := elastic.NewClient(elastic.SetURL(config.Get("elastic.url")))

	q := buildQuery(queryString)
	s := es.Search().
		Index(project.IndexName()).
		Type(types...).
		From(page*5).Size(5).
		Sort("_datetime", false).
		Query(q)

	searchResult, err := s.Do(ctx)

	if err != nil {
		// Index not found. That's ok, user didn't sent any data for now.
		// Otherwise error should be handled.
		if elastic.IsNotFound(err) {
			return SearchLogsResult{}, nil
		}
		return SearchLogsResult{}, err
	}

	if searchResult.Hits.TotalHits > 0 {
		var result = make([]string, len(searchResult.Hits.Hits))
		for i, hit := range searchResult.Hits.Hits {
			result[i] = string(*hit.Source)
		}

		return SearchLogsResult{
			Logs:  result,
			Total: searchResult.Hits.TotalHits,
		}, nil
	}

	return SearchLogsResult{}, nil
}

// If queryString is empty - return all records.
// Else, use query string dsl.
func buildQuery(queryString string) elastic.Query {
	if len(queryString) == 0 {
		return elastic.NewMatchAllQuery()
	}
	return elastic.NewQueryStringQuery(queryString).DefaultField("msg")
}
<commit_msg>Moved page size to constant<commit_after>package elastic

import (
	"context"

	"bitbucket.org/firstrow/logvoyage/models"
	"bitbucket.org/firstrow/logvoyage/shared/config"
	"gopkg.in/olivere/elastic.v5"
)

// How many log records display per-page.
const pageSize = 10

// LogRecord fetched from storage
type LogRecord struct {
	Source   string
	Datetime int64
}

// SearchLogsResult contains logs and total number of records in storage.
type SearchLogsResult struct {
	Logs  []string `json:"logs"`
	Total int64    `json:"total"`
}

// SearchLogs sends query to elastic search index
// types - ealstic types to search on.
// queryString - user provided data.
func SearchLogs(user *models.User, project *models.Project, types []string, queryString string, page int) (SearchLogsResult, error) {
	ctx := context.Background()
	es, _ := elastic.NewClient(elastic.SetURL(config.Get("elastic.url")))

	q := buildQuery(queryString)
	s := es.Search().
		Index(project.IndexName()).
		Type(types...).
		From(page*pageSize).Size(pageSize).
		Sort("_datetime", false).
		Query(q)

	searchResult, err := s.Do(ctx)

	if err != nil {
		// Index not found. That's ok, user didn't sent any data for now.
		// Otherwise error should be handled.
		if elastic.IsNotFound(err) {
			return SearchLogsResult{}, nil
		}
		return SearchLogsResult{}, err
	}

	if searchResult.Hits.TotalHits > 0 {
		var result = make([]string, len(searchResult.Hits.Hits))
		for i, hit := range searchResult.Hits.Hits {
			result[i] = string(*hit.Source)
		}

		return SearchLogsResult{
			Logs:  result,
			Total: searchResult.Hits.TotalHits,
		}, nil
	}

	return SearchLogsResult{}, nil
}

// If queryString is empty - return all records.
// Else, use query string dsl.
func buildQuery(queryString string) elastic.Query {
	if len(queryString) == 0 {
		return elastic.NewMatchAllQuery()
	}
	return elastic.NewQueryStringQuery(queryString).DefaultField("msg")
}
<|endoftext|>


<commit_before>package tarexport

import (
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"

	"github.com/Sirupsen/logrus"
	"github.com/docker/docker/image"
	"github.com/docker/docker/image/v1"
	"github.com/docker/docker/layer"
	"github.com/docker/docker/pkg/archive"
	"github.com/docker/docker/pkg/chrootarchive"
	"github.com/docker/docker/pkg/progress"
	"github.com/docker/docker/pkg/streamformatter"
	"github.com/docker/docker/pkg/stringid"
	"github.com/docker/docker/pkg/symlink"
	"github.com/docker/docker/reference"
)

func (l *tarexporter) Load(inTar io.ReadCloser, outStream io.Writer, quiet bool) error {
	var (
		sf             = streamformatter.NewJSONStreamFormatter()
		progressOutput progress.Output
	)
	if !quiet {
		progressOutput = sf.NewProgressOutput(outStream, false)
	} else {
		progressOutput = nil
	}

	tmpDir, err := ioutil.TempDir("", "docker-import-")
	if err != nil {
		return err
	}
	defer os.RemoveAll(tmpDir)

	if err := chrootarchive.Untar(inTar, tmpDir, nil); err != nil {
		return err
	}
	// read manifest, if no file then load in legacy mode
	manifestPath, err := safePath(tmpDir, manifestFileName)
	if err != nil {
		return err
	}
	manifestFile, err := os.Open(manifestPath)
	if err != nil {
		if os.IsNotExist(err) {
			return l.legacyLoad(tmpDir, outStream, progressOutput)
		}
		return manifestFile.Close()
	}
	defer manifestFile.Close()

	var manifest []manifestItem
	if err := json.NewDecoder(manifestFile).Decode(&manifest); err != nil {
		return err
	}

	for _, m := range manifest {
		configPath, err := safePath(tmpDir, m.Config)
		if err != nil {
			return err
		}
		config, err := ioutil.ReadFile(configPath)
		if err != nil {
			return err
		}
		img, err := image.NewFromJSON(config)
		if err != nil {
			return err
		}
		var rootFS image.RootFS
		rootFS = *img.RootFS
		rootFS.DiffIDs = nil

		if expected, actual := len(m.Layers), len(img.RootFS.DiffIDs); expected != actual {
			return fmt.Errorf("invalid manifest, layers length mismatch: expected %q, got %q", expected, actual)
		}

		for i, diffID := range img.RootFS.DiffIDs {
			layerPath, err := safePath(tmpDir, m.Layers[i])
			if err != nil {
				return err
			}
			r := rootFS
			r.Append(diffID)
			newLayer, err := l.ls.Get(r.ChainID())
			if err != nil {
				newLayer, err = l.loadLayer(layerPath, rootFS, diffID.String(), progressOutput)
				if err != nil {
					return err
				}
			}
			defer layer.ReleaseAndLog(l.ls, newLayer)
			if expected, actual := diffID, newLayer.DiffID(); expected != actual {
				return fmt.Errorf("invalid diffID for layer %d: expected %q, got %q", i, expected, actual)
			}
			rootFS.Append(diffID)
		}

		imgID, err := l.is.Create(config)
		if err != nil {
			return err
		}

		for _, repoTag := range m.RepoTags {
			named, err := reference.ParseNamed(repoTag)
			if err != nil {
				return err
			}
			ref, ok := named.(reference.NamedTagged)
			if !ok {
				return fmt.Errorf("invalid tag %q", repoTag)
			}
			l.setLoadedTag(ref, imgID, outStream)
		}

	}

	return nil
}

func (l *tarexporter) loadLayer(filename string, rootFS image.RootFS, id string, progressOutput progress.Output) (layer.Layer, error) {
	rawTar, err := os.Open(filename)
	if err != nil {
		logrus.Debugf("Error reading embedded tar: %v", err)
		return nil, err
	}
	defer rawTar.Close()

	inflatedLayerData, err := archive.DecompressStream(rawTar)
	if err != nil {
		return nil, err
	}
	defer inflatedLayerData.Close()

	if progressOutput != nil {
		fileInfo, err := os.Stat(filename)
		if err != nil {
			logrus.Debugf("Error statting file: %v", err)
			return nil, err
		}

		progressReader := progress.NewProgressReader(inflatedLayerData, progressOutput, fileInfo.Size(), stringid.TruncateID(id), "Loading layer")

		return l.ls.Register(progressReader, rootFS.ChainID())
	}
	return l.ls.Register(inflatedLayerData, rootFS.ChainID())
}

func (l *tarexporter) setLoadedTag(ref reference.NamedTagged, imgID image.ID, outStream io.Writer) error {
	if prevID, err := l.rs.Get(ref); err == nil && prevID != imgID {
		fmt.Fprintf(outStream, "The image %s already exists, renaming the old one with ID %s to empty string\n", ref.String(), string(prevID)) // todo: this message is wrong in case of multiple tags
	}

	if err := l.rs.AddTag(ref, imgID, true); err != nil {
		return err
	}
	return nil
}

func (l *tarexporter) legacyLoad(tmpDir string, outStream io.Writer, progressOutput progress.Output) error {
	legacyLoadedMap := make(map[string]image.ID)

	dirs, err := ioutil.ReadDir(tmpDir)
	if err != nil {
		return err
	}

	// every dir represents an image
	for _, d := range dirs {
		if d.IsDir() {
			if err := l.legacyLoadImage(d.Name(), tmpDir, legacyLoadedMap, progressOutput); err != nil {
				return err
			}
		}
	}

	// load tags from repositories file
	repositoriesPath, err := safePath(tmpDir, legacyRepositoriesFileName)
	if err != nil {
		return err
	}
	repositoriesFile, err := os.Open(repositoriesPath)
	if err != nil {
		if !os.IsNotExist(err) {
			return err
		}
		return repositoriesFile.Close()
	}
	defer repositoriesFile.Close()

	repositories := make(map[string]map[string]string)
	if err := json.NewDecoder(repositoriesFile).Decode(&repositories); err != nil {
		return err
	}

	for name, tagMap := range repositories {
		for tag, oldID := range tagMap {
			imgID, ok := legacyLoadedMap[oldID]
			if !ok {
				return fmt.Errorf("invalid target ID: %v", oldID)
			}
			named, err := reference.WithName(name)
			if err != nil {
				return err
			}
			ref, err := reference.WithTag(named, tag)
			if err != nil {
				return err
			}
			l.setLoadedTag(ref, imgID, outStream)
		}
	}

	return nil
}

func (l *tarexporter) legacyLoadImage(oldID, sourceDir string, loadedMap map[string]image.ID, progressOutput progress.Output) error {
	if _, loaded := loadedMap[oldID]; loaded {
		return nil
	}
	configPath, err := safePath(sourceDir, filepath.Join(oldID, legacyConfigFileName))
	if err != nil {
		return err
	}
	imageJSON, err := ioutil.ReadFile(configPath)
	if err != nil {
		logrus.Debugf("Error reading json: %v", err)
		return err
	}

	var img struct{ Parent string }
	if err := json.Unmarshal(imageJSON, &img); err != nil {
		return err
	}

	var parentID image.ID
	if img.Parent != "" {
		for {
			var loaded bool
			if parentID, loaded = loadedMap[img.Parent]; !loaded {
				if err := l.legacyLoadImage(img.Parent, sourceDir, loadedMap, progressOutput); err != nil {
					return err
				}
			} else {
				break
			}
		}
	}

	// todo: try to connect with migrate code
	rootFS := image.NewRootFS()
	var history []image.History

	if parentID != "" {
		parentImg, err := l.is.Get(parentID)
		if err != nil {
			return err
		}

		rootFS = parentImg.RootFS
		history = parentImg.History
	}

	layerPath, err := safePath(sourceDir, filepath.Join(oldID, legacyLayerFileName))
	if err != nil {
		return err
	}
	newLayer, err := l.loadLayer(layerPath, *rootFS, oldID, progressOutput)
	if err != nil {
		return err
	}
	rootFS.Append(newLayer.DiffID())

	h, err := v1.HistoryFromConfig(imageJSON, false)
	if err != nil {
		return err
	}
	history = append(history, h)

	config, err := v1.MakeConfigFromV1Config(imageJSON, rootFS, history)
	if err != nil {
		return err
	}
	imgID, err := l.is.Create(config)
	if err != nil {
		return err
	}

	metadata, err := l.ls.Release(newLayer)
	layer.LogReleaseMetadata(metadata)
	if err != nil {
		return err
	}

	if parentID != "" {
		if err := l.is.SetParent(imgID, parentID); err != nil {
			return err
		}
	}

	loadedMap[oldID] = imgID
	return nil
}

func safePath(base, path string) (string, error) {
	return symlink.FollowSymlinkInScope(filepath.Join(base, path), base)
}
<commit_msg>carry 17329<commit_after>package tarexport

import (
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"

	"github.com/Sirupsen/logrus"
	"github.com/docker/docker/image"
	"github.com/docker/docker/image/v1"
	"github.com/docker/docker/layer"
	"github.com/docker/docker/pkg/archive"
	"github.com/docker/docker/pkg/chrootarchive"
	"github.com/docker/docker/pkg/progress"
	"github.com/docker/docker/pkg/streamformatter"
	"github.com/docker/docker/pkg/stringid"
	"github.com/docker/docker/pkg/symlink"
	"github.com/docker/docker/reference"
)

func (l *tarexporter) Load(inTar io.ReadCloser, outStream io.Writer, quiet bool) error {
	var (
		sf             = streamformatter.NewJSONStreamFormatter()
		progressOutput progress.Output
	)
	if !quiet {
		progressOutput = sf.NewProgressOutput(outStream, false)
	}

	tmpDir, err := ioutil.TempDir("", "docker-import-")
	if err != nil {
		return err
	}
	defer os.RemoveAll(tmpDir)

	if err := chrootarchive.Untar(inTar, tmpDir, nil); err != nil {
		return err
	}
	// read manifest, if no file then load in legacy mode
	manifestPath, err := safePath(tmpDir, manifestFileName)
	if err != nil {
		return err
	}
	manifestFile, err := os.Open(manifestPath)
	if err != nil {
		if os.IsNotExist(err) {
			return l.legacyLoad(tmpDir, outStream, progressOutput)
		}
		return manifestFile.Close()
	}
	defer manifestFile.Close()

	var manifest []manifestItem
	if err := json.NewDecoder(manifestFile).Decode(&manifest); err != nil {
		return err
	}

	for _, m := range manifest {
		configPath, err := safePath(tmpDir, m.Config)
		if err != nil {
			return err
		}
		config, err := ioutil.ReadFile(configPath)
		if err != nil {
			return err
		}
		img, err := image.NewFromJSON(config)
		if err != nil {
			return err
		}
		var rootFS image.RootFS
		rootFS = *img.RootFS
		rootFS.DiffIDs = nil

		if expected, actual := len(m.Layers), len(img.RootFS.DiffIDs); expected != actual {
			return fmt.Errorf("invalid manifest, layers length mismatch: expected %q, got %q", expected, actual)
		}

		for i, diffID := range img.RootFS.DiffIDs {
			layerPath, err := safePath(tmpDir, m.Layers[i])
			if err != nil {
				return err
			}
			r := rootFS
			r.Append(diffID)
			newLayer, err := l.ls.Get(r.ChainID())
			if err != nil {
				newLayer, err = l.loadLayer(layerPath, rootFS, diffID.String(), progressOutput)
				if err != nil {
					return err
				}
			}
			defer layer.ReleaseAndLog(l.ls, newLayer)
			if expected, actual := diffID, newLayer.DiffID(); expected != actual {
				return fmt.Errorf("invalid diffID for layer %d: expected %q, got %q", i, expected, actual)
			}
			rootFS.Append(diffID)
		}

		imgID, err := l.is.Create(config)
		if err != nil {
			return err
		}

		for _, repoTag := range m.RepoTags {
			named, err := reference.ParseNamed(repoTag)
			if err != nil {
				return err
			}
			ref, ok := named.(reference.NamedTagged)
			if !ok {
				return fmt.Errorf("invalid tag %q", repoTag)
			}
			l.setLoadedTag(ref, imgID, outStream)
		}

	}

	return nil
}

func (l *tarexporter) loadLayer(filename string, rootFS image.RootFS, id string, progressOutput progress.Output) (layer.Layer, error) {
	rawTar, err := os.Open(filename)
	if err != nil {
		logrus.Debugf("Error reading embedded tar: %v", err)
		return nil, err
	}
	defer rawTar.Close()

	inflatedLayerData, err := archive.DecompressStream(rawTar)
	if err != nil {
		return nil, err
	}
	defer inflatedLayerData.Close()

	if progressOutput != nil {
		fileInfo, err := os.Stat(filename)
		if err != nil {
			logrus.Debugf("Error statting file: %v", err)
			return nil, err
		}

		progressReader := progress.NewProgressReader(inflatedLayerData, progressOutput, fileInfo.Size(), stringid.TruncateID(id), "Loading layer")

		return l.ls.Register(progressReader, rootFS.ChainID())
	}
	return l.ls.Register(inflatedLayerData, rootFS.ChainID())
}

func (l *tarexporter) setLoadedTag(ref reference.NamedTagged, imgID image.ID, outStream io.Writer) error {
	if prevID, err := l.rs.Get(ref); err == nil && prevID != imgID {
		fmt.Fprintf(outStream, "The image %s already exists, renaming the old one with ID %s to empty string\n", ref.String(), string(prevID)) // todo: this message is wrong in case of multiple tags
	}

	if err := l.rs.AddTag(ref, imgID, true); err != nil {
		return err
	}
	return nil
}

func (l *tarexporter) legacyLoad(tmpDir string, outStream io.Writer, progressOutput progress.Output) error {
	legacyLoadedMap := make(map[string]image.ID)

	dirs, err := ioutil.ReadDir(tmpDir)
	if err != nil {
		return err
	}

	// every dir represents an image
	for _, d := range dirs {
		if d.IsDir() {
			if err := l.legacyLoadImage(d.Name(), tmpDir, legacyLoadedMap, progressOutput); err != nil {
				return err
			}
		}
	}

	// load tags from repositories file
	repositoriesPath, err := safePath(tmpDir, legacyRepositoriesFileName)
	if err != nil {
		return err
	}
	repositoriesFile, err := os.Open(repositoriesPath)
	if err != nil {
		if !os.IsNotExist(err) {
			return err
		}
		return repositoriesFile.Close()
	}
	defer repositoriesFile.Close()

	repositories := make(map[string]map[string]string)
	if err := json.NewDecoder(repositoriesFile).Decode(&repositories); err != nil {
		return err
	}

	for name, tagMap := range repositories {
		for tag, oldID := range tagMap {
			imgID, ok := legacyLoadedMap[oldID]
			if !ok {
				return fmt.Errorf("invalid target ID: %v", oldID)
			}
			named, err := reference.WithName(name)
			if err != nil {
				return err
			}
			ref, err := reference.WithTag(named, tag)
			if err != nil {
				return err
			}
			l.setLoadedTag(ref, imgID, outStream)
		}
	}

	return nil
}

func (l *tarexporter) legacyLoadImage(oldID, sourceDir string, loadedMap map[string]image.ID, progressOutput progress.Output) error {
	if _, loaded := loadedMap[oldID]; loaded {
		return nil
	}
	configPath, err := safePath(sourceDir, filepath.Join(oldID, legacyConfigFileName))
	if err != nil {
		return err
	}
	imageJSON, err := ioutil.ReadFile(configPath)
	if err != nil {
		logrus.Debugf("Error reading json: %v", err)
		return err
	}

	var img struct{ Parent string }
	if err := json.Unmarshal(imageJSON, &img); err != nil {
		return err
	}

	var parentID image.ID
	if img.Parent != "" {
		for {
			var loaded bool
			if parentID, loaded = loadedMap[img.Parent]; !loaded {
				if err := l.legacyLoadImage(img.Parent, sourceDir, loadedMap, progressOutput); err != nil {
					return err
				}
			} else {
				break
			}
		}
	}

	// todo: try to connect with migrate code
	rootFS := image.NewRootFS()
	var history []image.History

	if parentID != "" {
		parentImg, err := l.is.Get(parentID)
		if err != nil {
			return err
		}

		rootFS = parentImg.RootFS
		history = parentImg.History
	}

	layerPath, err := safePath(sourceDir, filepath.Join(oldID, legacyLayerFileName))
	if err != nil {
		return err
	}
	newLayer, err := l.loadLayer(layerPath, *rootFS, oldID, progressOutput)
	if err != nil {
		return err
	}
	rootFS.Append(newLayer.DiffID())

	h, err := v1.HistoryFromConfig(imageJSON, false)
	if err != nil {
		return err
	}
	history = append(history, h)

	config, err := v1.MakeConfigFromV1Config(imageJSON, rootFS, history)
	if err != nil {
		return err
	}
	imgID, err := l.is.Create(config)
	if err != nil {
		return err
	}

	metadata, err := l.ls.Release(newLayer)
	layer.LogReleaseMetadata(metadata)
	if err != nil {
		return err
	}

	if parentID != "" {
		if err := l.is.SetParent(imgID, parentID); err != nil {
			return err
		}
	}

	loadedMap[oldID] = imgID
	return nil
}

func safePath(base, path string) (string, error) {
	return symlink.FollowSymlinkInScope(filepath.Join(base, path), base)
}
<|endoftext|>


<commit_before>// Package v8 contains code for importing data from 0.8 instances of InfluxDB.
package v8 // import "github.com/influxdata/influxdb/importer/v8"

import (
	"bufio"
	"compress/gzip"
	"fmt"
	"io"
	"log"
	"os"
	"strings"
	"time"

	"github.com/influxdata/influxdb/client"
)

const batchSize = 5000

// Config is the config used to initialize a Importer importer
type Config struct {
	Path       string // Path to import data.
	Version    string
	Compressed bool // Whether import data is gzipped.
	PPS        int  // points per second importer imports with.

	client.Config
}

// NewConfig returns an initialized *Config
func NewConfig() Config {
	return Config{Config: client.NewConfig()}
}

// Importer is the importer used for importing 0.8 data
type Importer struct {
	client                *client.Client
	database              string
	retentionPolicy       string
	config                Config
	batch                 []string
	totalInserts          int
	failedInserts         int
	totalCommands         int
	throttlePointsWritten int
	lastWrite             time.Time
	throttle              *time.Ticker

	stderrLogger *log.Logger
	stdoutLogger *log.Logger
}

// NewImporter will return an intialized Importer struct
func NewImporter(config Config) *Importer {
	config.UserAgent = fmt.Sprintf("influxDB importer/%s", config.Version)
	return &Importer{
		config:       config,
		batch:        make([]string, 0, batchSize),
		stdoutLogger: log.New(os.Stdout, "", log.LstdFlags),
		stderrLogger: log.New(os.Stderr, "", log.LstdFlags),
	}
}

// Import processes the specified file in the Config and writes the data to the databases in chunks specified by batchSize
func (i *Importer) Import() error {
	// Create a client and try to connect.
	cl, err := client.NewClient(i.config.Config)
	if err != nil {
		return fmt.Errorf("could not create client %s", err)
	}
	i.client = cl
	if _, _, e := i.client.Ping(); e != nil {
		return fmt.Errorf("failed to connect to %s\n", i.client.Addr())
	}

	// Validate args
	if i.config.Path == "" {
		return fmt.Errorf("file argument required")
	}

	defer func() {
		if i.totalInserts > 0 {
			i.stdoutLogger.Printf("Processed %d commands\n", i.totalCommands)
			i.stdoutLogger.Printf("Processed %d inserts\n", i.totalInserts)
			i.stdoutLogger.Printf("Failed %d inserts\n", i.failedInserts)
		}
	}()

	// Open the file
	f, err := os.Open(i.config.Path)
	if err != nil {
		return err
	}
	defer f.Close()

	var r io.Reader

	// If gzipped, wrap in a gzip reader
	if i.config.Compressed {
		gr, err := gzip.NewReader(f)
		if err != nil {
			return err
		}
		defer gr.Close()
		// Set the reader to the gzip reader
		r = gr
	} else {
		// Standard text file so our reader can just be the file
		r = f
	}

	// Get our reader
	scanner := bufio.NewReader(r)

	// Process the DDL
	if err := i.processDDL(scanner); err != nil {
		return fmt.Errorf("reading standard input: %s", err)
	}

	// Set up our throttle channel.  Since there is effectively no other activity at this point
	// the smaller resolution gets us much closer to the requested PPS
	i.throttle = time.NewTicker(time.Microsecond)
	defer i.throttle.Stop()

	// Prime the last write
	i.lastWrite = time.Now()

	// Process the DML
	if err := i.processDML(scanner); err != nil {
		return fmt.Errorf("reading standard input: %s", err)
	}

	// If there were any failed inserts then return an error so that a non-zero
	// exit code can be returned.
	if i.failedInserts > 0 {
		plural := " was"
		if i.failedInserts > 1 {
			plural = "s were"
		}

		return fmt.Errorf("%d point%s not inserted", i.failedInserts, plural)
	}

	return nil
}

func (i *Importer) processDDL(scanner *bufio.Reader) error {
	for {
		line, err := scanner.ReadString(byte('\n'))
		if err != nil && err != io.EOF {
			return err
		} else if err == io.EOF {
			return nil
		}
		// If we find the DML token, we are done with DDL
		if strings.HasPrefix(line, "# DML") {
			return nil
		}
		if strings.HasPrefix(line, "#") {
			continue
		}
		// Skip blank lines
		if strings.TrimSpace(line) == "" {
			continue
		}
		i.queryExecutor(line)
	}
}

func (i *Importer) processDML(scanner *bufio.Reader) error {
	start := time.Now()
	for {
		line, err := scanner.ReadString(byte('\n'))
		if err != nil && err != io.EOF {
			return err
		} else if err == io.EOF {
			// Call batchWrite one last time to flush anything out in the batch
			i.batchWrite()
			return nil
		}
		if strings.HasPrefix(line, "# CONTEXT-DATABASE:") {
			i.database = strings.TrimSpace(strings.Split(line, ":")[1])
		}
		if strings.HasPrefix(line, "# CONTEXT-RETENTION-POLICY:") {
			i.retentionPolicy = strings.TrimSpace(strings.Split(line, ":")[1])
		}
		if strings.HasPrefix(line, "#") {
			continue
		}
		// Skip blank lines
		if strings.TrimSpace(line) == "" {
			continue
		}
		i.batchAccumulator(line, start)
	}
}

func (i *Importer) execute(command string) {
	response, err := i.client.Query(client.Query{Command: command, Database: i.database})
	if err != nil {
		i.stderrLogger.Printf("error: %s\n", err)
		return
	}
	if err := response.Error(); err != nil {
		i.stderrLogger.Printf("error: %s\n", response.Error())
	}
}

func (i *Importer) queryExecutor(command string) {
	i.totalCommands++
	i.execute(command)
}

func (i *Importer) batchAccumulator(line string, start time.Time) {
	i.batch = append(i.batch, line)
	if len(i.batch) == batchSize {
		i.batchWrite()
		i.batch = i.batch[:0]
		// Give some status feedback every 100000 lines processed
		processed := i.totalInserts + i.failedInserts
		if processed%100000 == 0 {
			since := time.Since(start)
			pps := float64(processed) / since.Seconds()
			i.stdoutLogger.Printf("Processed %d lines.  Time elapsed: %s.  Points per second (PPS): %d", processed, since.String(), int64(pps))
		}
	}
}

func (i *Importer) batchWrite() {
	// Accumulate the batch size to see how many points we have written this second
	i.throttlePointsWritten += len(i.batch)

	// Find out when we last wrote data
	since := time.Since(i.lastWrite)

	// Check to see if we've exceeded our points per second for the current timeframe
	var currentPPS int
	if since.Seconds() > 0 {
		currentPPS = int(float64(i.throttlePointsWritten) / since.Seconds())
	} else {
		currentPPS = i.throttlePointsWritten
	}

	// If our currentPPS is greater than the PPS specified, then we wait and retry
	if int(currentPPS) > i.config.PPS && i.config.PPS != 0 {
		// Wait for the next tick
		<-i.throttle.C

		// Decrement the batch size back out as it is going to get called again
		i.throttlePointsWritten -= len(i.batch)
		i.batchWrite()
		return
	}

	_, e := i.client.WriteLineProtocol(strings.Join(i.batch, "\n"), i.database, i.retentionPolicy, i.config.Precision, i.config.WriteConsistency)
	if e != nil {
		i.stderrLogger.Println("error writing batch: ", e)
		i.stderrLogger.Println(strings.Join(i.batch, "\n"))
		i.failedInserts += len(i.batch)
	} else {
		i.totalInserts += len(i.batch)
	}
	i.throttlePointsWritten = 0
	i.lastWrite = time.Now()
}
<commit_msg>Fix imports of multiple databases in a single import file from `influx -import`<commit_after>// Package v8 contains code for importing data from 0.8 instances of InfluxDB.
package v8 // import "github.com/influxdata/influxdb/importer/v8"

import (
	"bufio"
	"compress/gzip"
	"fmt"
	"io"
	"log"
	"os"
	"strings"
	"time"

	"github.com/influxdata/influxdb/client"
)

const batchSize = 5000

// Config is the config used to initialize a Importer importer
type Config struct {
	Path       string // Path to import data.
	Version    string
	Compressed bool // Whether import data is gzipped.
	PPS        int  // points per second importer imports with.

	client.Config
}

// NewConfig returns an initialized *Config
func NewConfig() Config {
	return Config{Config: client.NewConfig()}
}

// Importer is the importer used for importing 0.8 data
type Importer struct {
	client                *client.Client
	database              string
	retentionPolicy       string
	config                Config
	batch                 []string
	totalInserts          int
	failedInserts         int
	totalCommands         int
	throttlePointsWritten int
	startTime             time.Time
	lastWrite             time.Time
	throttle              *time.Ticker

	stderrLogger *log.Logger
	stdoutLogger *log.Logger
}

// NewImporter will return an intialized Importer struct
func NewImporter(config Config) *Importer {
	config.UserAgent = fmt.Sprintf("influxDB importer/%s", config.Version)
	return &Importer{
		config:       config,
		batch:        make([]string, 0, batchSize),
		stdoutLogger: log.New(os.Stdout, "", log.LstdFlags),
		stderrLogger: log.New(os.Stderr, "", log.LstdFlags),
	}
}

// Import processes the specified file in the Config and writes the data to the databases in chunks specified by batchSize
func (i *Importer) Import() error {
	// Create a client and try to connect.
	cl, err := client.NewClient(i.config.Config)
	if err != nil {
		return fmt.Errorf("could not create client %s", err)
	}
	i.client = cl
	if _, _, e := i.client.Ping(); e != nil {
		return fmt.Errorf("failed to connect to %s\n", i.client.Addr())
	}

	// Validate args
	if i.config.Path == "" {
		return fmt.Errorf("file argument required")
	}

	defer func() {
		if i.totalInserts > 0 {
			i.stdoutLogger.Printf("Processed %d commands\n", i.totalCommands)
			i.stdoutLogger.Printf("Processed %d inserts\n", i.totalInserts)
			i.stdoutLogger.Printf("Failed %d inserts\n", i.failedInserts)
		}
	}()

	// Open the file
	f, err := os.Open(i.config.Path)
	if err != nil {
		return err
	}
	defer f.Close()

	var r io.Reader

	// If gzipped, wrap in a gzip reader
	if i.config.Compressed {
		gr, err := gzip.NewReader(f)
		if err != nil {
			return err
		}
		defer gr.Close()
		// Set the reader to the gzip reader
		r = gr
	} else {
		// Standard text file so our reader can just be the file
		r = f
	}

	// Get our reader
	scanner := bufio.NewReader(r)

	// Process the DDL
	if err := i.processDDL(scanner); err != nil {
		return fmt.Errorf("reading standard input: %s", err)
	}

	// Set up our throttle channel.  Since there is effectively no other activity at this point
	// the smaller resolution gets us much closer to the requested PPS
	i.throttle = time.NewTicker(time.Microsecond)
	defer i.throttle.Stop()

	// Prime the last write
	i.lastWrite = time.Now()

	// Process the DML
	if err := i.processDML(scanner); err != nil {
		return fmt.Errorf("reading standard input: %s", err)
	}

	// If there were any failed inserts then return an error so that a non-zero
	// exit code can be returned.
	if i.failedInserts > 0 {
		plural := " was"
		if i.failedInserts > 1 {
			plural = "s were"
		}

		return fmt.Errorf("%d point%s not inserted", i.failedInserts, plural)
	}

	return nil
}

func (i *Importer) processDDL(scanner *bufio.Reader) error {
	for {
		line, err := scanner.ReadString(byte('\n'))
		if err != nil && err != io.EOF {
			return err
		} else if err == io.EOF {
			return nil
		}
		// If we find the DML token, we are done with DDL
		if strings.HasPrefix(line, "# DML") {
			return nil
		}
		if strings.HasPrefix(line, "#") {
			continue
		}
		// Skip blank lines
		if strings.TrimSpace(line) == "" {
			continue
		}
		i.queryExecutor(line)
	}
}

func (i *Importer) processDML(scanner *bufio.Reader) error {
	i.startTime = time.Now()
	for {
		line, err := scanner.ReadString(byte('\n'))
		if err != nil && err != io.EOF {
			return err
		} else if err == io.EOF {
			// Call batchWrite one last time to flush anything out in the batch
			i.batchWrite()
			return nil
		}
		if strings.HasPrefix(line, "# CONTEXT-DATABASE:") {
			i.batchWrite()
			i.database = strings.TrimSpace(strings.Split(line, ":")[1])
		}
		if strings.HasPrefix(line, "# CONTEXT-RETENTION-POLICY:") {
			i.batchWrite()
			i.retentionPolicy = strings.TrimSpace(strings.Split(line, ":")[1])
		}
		if strings.HasPrefix(line, "#") {
			continue
		}
		// Skip blank lines
		if strings.TrimSpace(line) == "" {
			continue
		}
		i.batchAccumulator(line)
	}
}

func (i *Importer) execute(command string) {
	response, err := i.client.Query(client.Query{Command: command, Database: i.database})
	if err != nil {
		i.stderrLogger.Printf("error: %s\n", err)
		return
	}
	if err := response.Error(); err != nil {
		i.stderrLogger.Printf("error: %s\n", response.Error())
	}
}

func (i *Importer) queryExecutor(command string) {
	i.totalCommands++
	i.execute(command)
}

func (i *Importer) batchAccumulator(line string) {
	i.batch = append(i.batch, line)
	if len(i.batch) == batchSize {
		i.batchWrite()
	}
}

func (i *Importer) batchWrite() {
	// Exit early if there are no points in the batch.
	if len(i.batch) == 0 {
		return
	}

	// Accumulate the batch size to see how many points we have written this second
	i.throttlePointsWritten += len(i.batch)

	// Find out when we last wrote data
	since := time.Since(i.lastWrite)

	// Check to see if we've exceeded our points per second for the current timeframe
	var currentPPS int
	if since.Seconds() > 0 {
		currentPPS = int(float64(i.throttlePointsWritten) / since.Seconds())
	} else {
		currentPPS = i.throttlePointsWritten
	}

	// If our currentPPS is greater than the PPS specified, then we wait and retry
	if int(currentPPS) > i.config.PPS && i.config.PPS != 0 {
		// Wait for the next tick
		<-i.throttle.C

		// Decrement the batch size back out as it is going to get called again
		i.throttlePointsWritten -= len(i.batch)
		i.batchWrite()
		return
	}

	_, e := i.client.WriteLineProtocol(strings.Join(i.batch, "\n"), i.database, i.retentionPolicy, i.config.Precision, i.config.WriteConsistency)
	if e != nil {
		i.stderrLogger.Println("error writing batch: ", e)
		i.stderrLogger.Println(strings.Join(i.batch, "\n"))
		i.failedInserts += len(i.batch)
	} else {
		i.totalInserts += len(i.batch)
	}
	i.throttlePointsWritten = 0
	i.lastWrite = time.Now()

	// Clear the batch and record the number of processed points.
	i.batch = i.batch[:0]
	// Give some status feedback every 100000 lines processed
	processed := i.totalInserts + i.failedInserts
	if processed%100000 == 0 {
		since := time.Since(i.startTime)
		pps := float64(processed) / since.Seconds()
		i.stdoutLogger.Printf("Processed %d lines.  Time elapsed: %s.  Points per second (PPS): %d", processed, since.String(), int64(pps))
	}
}
<|endoftext|>


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

// Package io provides basic interfaces to I/O primitives.
// Its primary job is to wrap existing implementations of such primitives,
// such as those in package os, into shared public interfaces that
// abstract the functionality, plus some other related primitives.
package io

import "os"

// Error represents an unexpected I/O behavior.
type Error struct {
	os.ErrorString
}

// ErrShortWrite means that a write accepted fewer bytes than requested
// but failed to return an explicit error.
var ErrShortWrite os.Error = &Error{"short write"}

// ErrShortBuffer means that a read required a longer buffer than was provided.
var ErrShortBuffer os.Error = &Error{"short buffer"}

// ErrUnexpectedEOF means that os.EOF was encountered in the
// middle of reading a fixed-size block or data structure.
var ErrUnexpectedEOF os.Error = &Error{"unexpected EOF"}

// Reader is the interface that wraps the basic Read method.
//
// Read reads up to len(p) bytes into p.  It returns the number of bytes
// read (0 <= n <= len(p)) and any error encountered.
// Even if Read returns n < len(p),
// it may use all of p as scratch space during the call.
// If some data is available but not len(p) bytes, Read conventionally
// returns what is available rather than block waiting for more.
//
// At the end of the input stream, Read returns 0, os.EOF.
// Read may return a non-zero number of bytes with a non-nil err.
// In particular, a Read that exhausts the input may return n > 0, os.EOF.
type Reader interface {
	Read(p []byte) (n int, err os.Error)
}

// Writer is the interface that wraps the basic Write method.
//
// Write writes len(p) bytes from p to the underlying data stream.
// It returns the number of bytes written from p (0 <= n <= len(p))
// and any error encountered that caused the write to stop early.
// Write must return a non-nil error if it returns n < len(p).
type Writer interface {
	Write(p []byte) (n int, err os.Error)
}

// Closer is the interface that wraps the basic Close method.
type Closer interface {
	Close() os.Error
}

// Seeker is the interface that wraps the basic Seek method.
//
// Seek sets the offset for the next Read or Write to offset,
// interpreted according to whence: 0 means relative to the origin of
// the file, 1 means relative to the current offset, and 2 means
// relative to the end.  Seek returns the new offset and an Error, if
// any.
type Seeker interface {
	Seek(offset int64, whence int) (ret int64, err os.Error)
}

// ReadWriter is the interface that groups the basic Read and Write methods.
type ReadWriter interface {
	Reader
	Writer
}

// ReadCloser is the interface that groups the basic Read and Close methods.
type ReadCloser interface {
	Reader
	Closer
}

// WriteCloser is the interface that groups the basic Write and Close methods.
type WriteCloser interface {
	Writer
	Closer
}

// ReadWriteCloser is the interface that groups the basic Read, Write and Close methods.
type ReadWriteCloser interface {
	Reader
	Writer
	Closer
}

// ReadSeeker is the interface that groups the basic Read and Seek methods.
type ReadSeeker interface {
	Reader
	Seeker
}

// WriteSeeker is the interface that groups the basic Write and Seek methods.
type WriteSeeker interface {
	Writer
	Seeker
}

// ReadWriteSeeker is the interface that groups the basic Read, Write and Seek methods.
type ReadWriteSeeker interface {
	Reader
	Writer
	Seeker
}

// ReaderFrom is the interface that wraps the ReadFrom method.
type ReaderFrom interface {
	ReadFrom(r Reader) (n int64, err os.Error)
}

// WriterTo is the interface that wraps the WriteTo method.
type WriterTo interface {
	WriteTo(w Writer) (n int64, err os.Error)
}

// ReaderAt is the interface that wraps the basic ReadAt method.
//
// ReadAt reads len(p) bytes into p starting at offset off in the
// underlying data stream.  It returns the number of bytes
// read (0 <= n <= len(p)) and any error encountered.
//
// Even if ReadAt returns n < len(p),
// it may use all of p as scratch space during the call.
// If some data is available but not len(p) bytes, ReadAt blocks
// until either all the data is available or an error occurs.
//
// At the end of the input stream, ReadAt returns 0, os.EOF.
// ReadAt may return a non-zero number of bytes with a non-nil err.
// In particular, a ReadAt that exhausts the input may return n > 0, os.EOF.
//
// If ReadAt is reading from an data stream with a seek offset,
// ReadAt should not affect nor be affected by the underlying
// seek offset.
type ReaderAt interface {
	ReadAt(p []byte, off int64) (n int, err os.Error)
}

// WriterAt is the interface that wraps the basic WriteAt method.
//
// WriteAt writes len(p) bytes from p to the underlying data stream
// at offset off.  It returns the number of bytes written from p (0 <= n <= len(p))
// and any error encountered that caused the write to stop early.
// WriteAt must return a non-nil error if it returns n < len(p).
type WriterAt interface {
	WriteAt(p []byte, off int64) (n int, err os.Error)
}

// ByteReader is the interface that wraps the ReadByte method.
//
// ReadByte reads and returns the next byte from the input.
// If no byte is available, err will be set.
type ByteReader interface {
	ReadByte() (c byte, err os.Error)
}

// RuneReader is the interface that wraps the ReadRune method.
//
// ReadRune reads a single UTF-8 encoded Unicode character
// and returns the rune and its size in bytes. If no character is
// available, err will be set.
type RuneReader interface {
	ReadRune() (rune int, size int, err os.Error)
}

// WriteString writes the contents of the string s to w, which accepts an array of bytes.
func WriteString(w Writer, s string) (n int, err os.Error) {
	return w.Write([]byte(s))
}

// ReadAtLeast reads from r into buf until it has read at least min bytes.
// It returns the number of bytes copied and an error if fewer bytes were read.
// The error is os.EOF only if no bytes were read.
// If an EOF happens after reading fewer than min bytes,
// ReadAtLeast returns ErrUnexpectedEOF.
// If min is greater than the length of buf, ReadAtLeast returns ErrShortBuffer.
func ReadAtLeast(r Reader, buf []byte, min int) (n int, err os.Error) {
	if len(buf) < min {
		return 0, ErrShortBuffer
	}
	for n < min && err == nil {
		var nn int
		nn, err = r.Read(buf[n:])
		n += nn
	}
	if err == os.EOF {
		if n >= min {
			err = nil
		} else if n > 0 {
			err = ErrUnexpectedEOF
		}
	}
	return
}

// ReadFull reads exactly len(buf) bytes from r into buf.
// It returns the number of bytes copied and an error if fewer bytes were read.
// The error is os.EOF only if no bytes were read.
// If an EOF happens after reading some but not all the bytes,
// ReadFull returns ErrUnexpectedEOF.
func ReadFull(r Reader, buf []byte) (n int, err os.Error) {
	return ReadAtLeast(r, buf, len(buf))
}

// Copyn copies n bytes (or until an error) from src to dst.
// It returns the number of bytes copied and the error, if any.
//
// If dst implements the ReaderFrom interface,
// the copy is implemented by calling dst.ReadFrom(src).
func Copyn(dst Writer, src Reader, n int64) (written int64, err os.Error) {
	// If the writer has a ReadFrom method, use it to do the copy.
	// Avoids a buffer allocation and a copy.
	if rt, ok := dst.(ReaderFrom); ok {
		written, err = rt.ReadFrom(LimitReader(src, n))
		if written < n && err == nil {
			// rt stopped early; must have been EOF.
			err = os.EOF
		}
		return
	}
	buf := make([]byte, 32*1024)
	for written < n {
		l := len(buf)
		if d := n - written; d < int64(l) {
			l = int(d)
		}
		nr, er := src.Read(buf[0:l])
		if nr > 0 {
			nw, ew := dst.Write(buf[0:nr])
			if nw > 0 {
				written += int64(nw)
			}
			if ew != nil {
				err = ew
				break
			}
			if nr != nw {
				err = ErrShortWrite
				break
			}
		}
		if er != nil {
			err = er
			break
		}
	}
	return written, err
}

// Copy copies from src to dst until either EOF is reached
// on src or an error occurs.  It returns the number of bytes
// copied and the error, if any.
//
// If dst implements the ReaderFrom interface,
// the copy is implemented by calling dst.ReadFrom(src).
// Otherwise, if src implements the WriterTo interface,
// the copy is implemented by calling src.WriteTo(dst).
func Copy(dst Writer, src Reader) (written int64, err os.Error) {
	// If the writer has a ReadFrom method, use it to do the copy.
	// Avoids an allocation and a copy.
	if rt, ok := dst.(ReaderFrom); ok {
		return rt.ReadFrom(src)
	}
	// Similarly, if the reader has a WriteTo method, use it to do the copy.
	if wt, ok := src.(WriterTo); ok {
		return wt.WriteTo(dst)
	}
	buf := make([]byte, 32*1024)
	for {
		nr, er := src.Read(buf)
		if nr > 0 {
			nw, ew := dst.Write(buf[0:nr])
			if nw > 0 {
				written += int64(nw)
			}
			if ew != nil {
				err = ew
				break
			}
			if nr != nw {
				err = ErrShortWrite
				break
			}
		}
		if er == os.EOF {
			break
		}
		if er != nil {
			err = er
			break
		}
	}
	return written, err
}

// LimitReader returns a Reader that reads from r
// but stops with os.EOF after n bytes.
// The underlying implementation is a *LimitedReader.
func LimitReader(r Reader, n int64) Reader { return &LimitedReader{r, n} }

// A LimitedReader reads from R but limits the amount of
// data returned to just N bytes. Each call to Read
// updates N to reflect the new amount remaining.
type LimitedReader struct {
	R Reader // underlying reader
	N int64  // max bytes remaining
}

func (l *LimitedReader) Read(p []byte) (n int, err os.Error) {
	if l.N <= 0 {
		return 0, os.EOF
	}
	if int64(len(p)) > l.N {
		p = p[0:l.N]
	}
	n, err = l.R.Read(p)
	l.N -= int64(n)
	return
}

// NewSectionReader returns a SectionReader that reads from r
// starting at offset off and stops with os.EOF after n bytes.
func NewSectionReader(r ReaderAt, off int64, n int64) *SectionReader {
	return &SectionReader{r, off, off, off + n}
}

// SectionReader implements Read, Seek, and ReadAt on a section
// of an underlying ReaderAt.
type SectionReader struct {
	r     ReaderAt
	base  int64
	off   int64
	limit int64
}

func (s *SectionReader) Read(p []byte) (n int, err os.Error) {
	if s.off >= s.limit {
		return 0, os.EOF
	}
	if max := s.limit - s.off; int64(len(p)) > max {
		p = p[0:max]
	}
	n, err = s.r.ReadAt(p, s.off)
	s.off += int64(n)
	return
}

func (s *SectionReader) Seek(offset int64, whence int) (ret int64, err os.Error) {
	switch whence {
	default:
		return 0, os.EINVAL
	case 0:
		offset += s.base
	case 1:
		offset += s.off
	case 2:
		offset += s.limit
	}
	if offset < s.base || offset > s.limit {
		return 0, os.EINVAL
	}
	s.off = offset
	return offset - s.base, nil
}

func (s *SectionReader) ReadAt(p []byte, off int64) (n int, err os.Error) {
	if off < 0 || off >= s.limit-s.base {
		return 0, os.EOF
	}
	off += s.base
	if max := s.limit - off; int64(len(p)) > max {
		p = p[0:max]
	}
	return s.r.ReadAt(p, off)
}

// Size returns the size of the section in bytes.
func (s *SectionReader) Size() int64 { return s.limit - s.base }
<commit_msg>io: add ByteScanner, RuneScanner interfaces<commit_after>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package io provides basic interfaces to I/O primitives.
// Its primary job is to wrap existing implementations of such primitives,
// such as those in package os, into shared public interfaces that
// abstract the functionality, plus some other related primitives.
package io

import "os"

// Error represents an unexpected I/O behavior.
type Error struct {
	os.ErrorString
}

// ErrShortWrite means that a write accepted fewer bytes than requested
// but failed to return an explicit error.
var ErrShortWrite os.Error = &Error{"short write"}

// ErrShortBuffer means that a read required a longer buffer than was provided.
var ErrShortBuffer os.Error = &Error{"short buffer"}

// ErrUnexpectedEOF means that os.EOF was encountered in the
// middle of reading a fixed-size block or data structure.
var ErrUnexpectedEOF os.Error = &Error{"unexpected EOF"}

// Reader is the interface that wraps the basic Read method.
//
// Read reads up to len(p) bytes into p.  It returns the number of bytes
// read (0 <= n <= len(p)) and any error encountered.
// Even if Read returns n < len(p),
// it may use all of p as scratch space during the call.
// If some data is available but not len(p) bytes, Read conventionally
// returns what is available rather than block waiting for more.
//
// At the end of the input stream, Read returns 0, os.EOF.
// Read may return a non-zero number of bytes with a non-nil err.
// In particular, a Read that exhausts the input may return n > 0, os.EOF.
type Reader interface {
	Read(p []byte) (n int, err os.Error)
}

// Writer is the interface that wraps the basic Write method.
//
// Write writes len(p) bytes from p to the underlying data stream.
// It returns the number of bytes written from p (0 <= n <= len(p))
// and any error encountered that caused the write to stop early.
// Write must return a non-nil error if it returns n < len(p).
type Writer interface {
	Write(p []byte) (n int, err os.Error)
}

// Closer is the interface that wraps the basic Close method.
type Closer interface {
	Close() os.Error
}

// Seeker is the interface that wraps the basic Seek method.
//
// Seek sets the offset for the next Read or Write to offset,
// interpreted according to whence: 0 means relative to the origin of
// the file, 1 means relative to the current offset, and 2 means
// relative to the end.  Seek returns the new offset and an Error, if
// any.
type Seeker interface {
	Seek(offset int64, whence int) (ret int64, err os.Error)
}

// ReadWriter is the interface that groups the basic Read and Write methods.
type ReadWriter interface {
	Reader
	Writer
}

// ReadCloser is the interface that groups the basic Read and Close methods.
type ReadCloser interface {
	Reader
	Closer
}

// WriteCloser is the interface that groups the basic Write and Close methods.
type WriteCloser interface {
	Writer
	Closer
}

// ReadWriteCloser is the interface that groups the basic Read, Write and Close methods.
type ReadWriteCloser interface {
	Reader
	Writer
	Closer
}

// ReadSeeker is the interface that groups the basic Read and Seek methods.
type ReadSeeker interface {
	Reader
	Seeker
}

// WriteSeeker is the interface that groups the basic Write and Seek methods.
type WriteSeeker interface {
	Writer
	Seeker
}

// ReadWriteSeeker is the interface that groups the basic Read, Write and Seek methods.
type ReadWriteSeeker interface {
	Reader
	Writer
	Seeker
}

// ReaderFrom is the interface that wraps the ReadFrom method.
type ReaderFrom interface {
	ReadFrom(r Reader) (n int64, err os.Error)
}

// WriterTo is the interface that wraps the WriteTo method.
type WriterTo interface {
	WriteTo(w Writer) (n int64, err os.Error)
}

// ReaderAt is the interface that wraps the basic ReadAt method.
//
// ReadAt reads len(p) bytes into p starting at offset off in the
// underlying data stream.  It returns the number of bytes
// read (0 <= n <= len(p)) and any error encountered.
//
// Even if ReadAt returns n < len(p),
// it may use all of p as scratch space during the call.
// If some data is available but not len(p) bytes, ReadAt blocks
// until either all the data is available or an error occurs.
//
// At the end of the input stream, ReadAt returns 0, os.EOF.
// ReadAt may return a non-zero number of bytes with a non-nil err.
// In particular, a ReadAt that exhausts the input may return n > 0, os.EOF.
//
// If ReadAt is reading from an data stream with a seek offset,
// ReadAt should not affect nor be affected by the underlying
// seek offset.
type ReaderAt interface {
	ReadAt(p []byte, off int64) (n int, err os.Error)
}

// WriterAt is the interface that wraps the basic WriteAt method.
//
// WriteAt writes len(p) bytes from p to the underlying data stream
// at offset off.  It returns the number of bytes written from p (0 <= n <= len(p))
// and any error encountered that caused the write to stop early.
// WriteAt must return a non-nil error if it returns n < len(p).
type WriterAt interface {
	WriteAt(p []byte, off int64) (n int, err os.Error)
}

// ByteReader is the interface that wraps the ReadByte method.
//
// ReadByte reads and returns the next byte from the input.
// If no byte is available, err will be set.
type ByteReader interface {
	ReadByte() (c byte, err os.Error)
}

// ByteScanner is the interface that adds the UnreadByte method to the
// basic ReadByte method.
//
// UnreadByte causes the next call to ReadByte to return the same byte
// as the previous call to ReadByte.
// It may be an error to call UnreadByte twice without an intervening
// call to ReadByte.
type ByteScanner interface {
	ByteReader
	UnreadByte() os.Error
}

// RuneReader is the interface that wraps the ReadRune method.
//
// ReadRune reads a single UTF-8 encoded Unicode character
// and returns the rune and its size in bytes. If no character is
// available, err will be set.
type RuneReader interface {
	ReadRune() (rune int, size int, err os.Error)
}

// RuneScanner is the interface that adds the UnreadRune method to the
// basic ReadRune method.
//
// UnreadRune causes the next call to ReadRune to return the same rune
// as the previous call to ReadRune.
// It may be an error to call UnreadRune twice without an intervening
// call to ReadRune.
type RuneScanner interface {
	RuneReader
	UnreadRune() os.Error
}

// WriteString writes the contents of the string s to w, which accepts an array of bytes.
func WriteString(w Writer, s string) (n int, err os.Error) {
	return w.Write([]byte(s))
}

// ReadAtLeast reads from r into buf until it has read at least min bytes.
// It returns the number of bytes copied and an error if fewer bytes were read.
// The error is os.EOF only if no bytes were read.
// If an EOF happens after reading fewer than min bytes,
// ReadAtLeast returns ErrUnexpectedEOF.
// If min is greater than the length of buf, ReadAtLeast returns ErrShortBuffer.
func ReadAtLeast(r Reader, buf []byte, min int) (n int, err os.Error) {
	if len(buf) < min {
		return 0, ErrShortBuffer
	}
	for n < min && err == nil {
		var nn int
		nn, err = r.Read(buf[n:])
		n += nn
	}
	if err == os.EOF {
		if n >= min {
			err = nil
		} else if n > 0 {
			err = ErrUnexpectedEOF
		}
	}
	return
}

// ReadFull reads exactly len(buf) bytes from r into buf.
// It returns the number of bytes copied and an error if fewer bytes were read.
// The error is os.EOF only if no bytes were read.
// If an EOF happens after reading some but not all the bytes,
// ReadFull returns ErrUnexpectedEOF.
func ReadFull(r Reader, buf []byte) (n int, err os.Error) {
	return ReadAtLeast(r, buf, len(buf))
}

// Copyn copies n bytes (or until an error) from src to dst.
// It returns the number of bytes copied and the error, if any.
//
// If dst implements the ReaderFrom interface,
// the copy is implemented by calling dst.ReadFrom(src).
func Copyn(dst Writer, src Reader, n int64) (written int64, err os.Error) {
	// If the writer has a ReadFrom method, use it to do the copy.
	// Avoids a buffer allocation and a copy.
	if rt, ok := dst.(ReaderFrom); ok {
		written, err = rt.ReadFrom(LimitReader(src, n))
		if written < n && err == nil {
			// rt stopped early; must have been EOF.
			err = os.EOF
		}
		return
	}
	buf := make([]byte, 32*1024)
	for written < n {
		l := len(buf)
		if d := n - written; d < int64(l) {
			l = int(d)
		}
		nr, er := src.Read(buf[0:l])
		if nr > 0 {
			nw, ew := dst.Write(buf[0:nr])
			if nw > 0 {
				written += int64(nw)
			}
			if ew != nil {
				err = ew
				break
			}
			if nr != nw {
				err = ErrShortWrite
				break
			}
		}
		if er != nil {
			err = er
			break
		}
	}
	return written, err
}

// Copy copies from src to dst until either EOF is reached
// on src or an error occurs.  It returns the number of bytes
// copied and the error, if any.
//
// If dst implements the ReaderFrom interface,
// the copy is implemented by calling dst.ReadFrom(src).
// Otherwise, if src implements the WriterTo interface,
// the copy is implemented by calling src.WriteTo(dst).
func Copy(dst Writer, src Reader) (written int64, err os.Error) {
	// If the writer has a ReadFrom method, use it to do the copy.
	// Avoids an allocation and a copy.
	if rt, ok := dst.(ReaderFrom); ok {
		return rt.ReadFrom(src)
	}
	// Similarly, if the reader has a WriteTo method, use it to do the copy.
	if wt, ok := src.(WriterTo); ok {
		return wt.WriteTo(dst)
	}
	buf := make([]byte, 32*1024)
	for {
		nr, er := src.Read(buf)
		if nr > 0 {
			nw, ew := dst.Write(buf[0:nr])
			if nw > 0 {
				written += int64(nw)
			}
			if ew != nil {
				err = ew
				break
			}
			if nr != nw {
				err = ErrShortWrite
				break
			}
		}
		if er == os.EOF {
			break
		}
		if er != nil {
			err = er
			break
		}
	}
	return written, err
}

// LimitReader returns a Reader that reads from r
// but stops with os.EOF after n bytes.
// The underlying implementation is a *LimitedReader.
func LimitReader(r Reader, n int64) Reader { return &LimitedReader{r, n} }

// A LimitedReader reads from R but limits the amount of
// data returned to just N bytes. Each call to Read
// updates N to reflect the new amount remaining.
type LimitedReader struct {
	R Reader // underlying reader
	N int64  // max bytes remaining
}

func (l *LimitedReader) Read(p []byte) (n int, err os.Error) {
	if l.N <= 0 {
		return 0, os.EOF
	}
	if int64(len(p)) > l.N {
		p = p[0:l.N]
	}
	n, err = l.R.Read(p)
	l.N -= int64(n)
	return
}

// NewSectionReader returns a SectionReader that reads from r
// starting at offset off and stops with os.EOF after n bytes.
func NewSectionReader(r ReaderAt, off int64, n int64) *SectionReader {
	return &SectionReader{r, off, off, off + n}
}

// SectionReader implements Read, Seek, and ReadAt on a section
// of an underlying ReaderAt.
type SectionReader struct {
	r     ReaderAt
	base  int64
	off   int64
	limit int64
}

func (s *SectionReader) Read(p []byte) (n int, err os.Error) {
	if s.off >= s.limit {
		return 0, os.EOF
	}
	if max := s.limit - s.off; int64(len(p)) > max {
		p = p[0:max]
	}
	n, err = s.r.ReadAt(p, s.off)
	s.off += int64(n)
	return
}

func (s *SectionReader) Seek(offset int64, whence int) (ret int64, err os.Error) {
	switch whence {
	default:
		return 0, os.EINVAL
	case 0:
		offset += s.base
	case 1:
		offset += s.off
	case 2:
		offset += s.limit
	}
	if offset < s.base || offset > s.limit {
		return 0, os.EINVAL
	}
	s.off = offset
	return offset - s.base, nil
}

func (s *SectionReader) ReadAt(p []byte, off int64) (n int, err os.Error) {
	if off < 0 || off >= s.limit-s.base {
		return 0, os.EOF
	}
	off += s.base
	if max := s.limit - off; int64(len(p)) > max {
		p = p[0:max]
	}
	return s.r.ReadAt(p, off)
}

// Size returns the size of the section in bytes.
func (s *SectionReader) Size() int64 { return s.limit - s.base }
<|endoftext|>


<commit_before>package input

import (
	"errors"
	"fmt"
	"path/filepath"

	"github.com/cweill/gotests/internal/models"
)

var ErrNoFilesFound = errors.New("no files found")

// Returns all the Golang files for the given path. Ignores hidden files.
func Files(srcPath string) ([]models.Path, error) {
	var srcPaths []models.Path
	srcPath, err := filepath.Abs(srcPath)
	if err != nil {
		return nil, fmt.Errorf("filepath.Abs: %v\n", err)
	}
	if filepath.Ext(srcPath) == "" {
		ps, err := filepath.Glob(srcPath + "/*.go")
		if err != nil {
			return nil, fmt.Errorf("filepath.Glob: %v\n", err)
		}
		for _, p := range ps {
			if isHiddenFile(p) {
				continue
			}
			src := models.Path(p)
			if !src.IsTestPath() {
				srcPaths = append(srcPaths, src)
			}
		}
		return srcPaths, nil
	}
	if filepath.Ext(srcPath) == ".go" && !isHiddenFile(srcPath) {
		src := models.Path(srcPath)
		if !src.IsTestPath() {
			srcPaths = append(srcPaths, src)
		}
		return srcPaths, nil
	}
	return nil, fmt.Errorf("no files found at %v", srcPath)
}

func isHiddenFile(path string) bool {
	return []rune(filepath.Base(path))[0] == '.'
}
<commit_msg>Split input.Files into smaller functions.<commit_after>package input

import (
	"fmt"
	"path/filepath"

	"github.com/cweill/gotests/internal/models"
)

// Returns all the Golang files for the given path. Ignores hidden files.
func Files(srcPath string) ([]models.Path, error) {
	srcPath, err := filepath.Abs(srcPath)
	if err != nil {
		return nil, fmt.Errorf("filepath.Abs: %v\n", err)
	}
	if filepath.Ext(srcPath) == "" {
		return dirFiles(srcPath)
	}
	return file(srcPath)
}

func dirFiles(srcPath string) ([]models.Path, error) {
	ps, err := filepath.Glob(srcPath + "/*.go")
	if err != nil {
		return nil, fmt.Errorf("filepath.Glob: %v\n", err)
	}
	var srcPaths []models.Path
	for _, p := range ps {
		src := models.Path(p)
		if isHiddenFile(p) || src.IsTestPath() {
			continue
		}
		srcPaths = append(srcPaths, src)
	}
	return srcPaths, nil
}

func file(srcPath string) ([]models.Path, error) {
	src := models.Path(srcPath)
	if filepath.Ext(srcPath) != ".go" || isHiddenFile(srcPath) || src.IsTestPath() {
		return nil, fmt.Errorf("no Go source files found at %v", srcPath)
	}
	return []models.Path{src}, nil
}

func isHiddenFile(path string) bool {
	return []rune(filepath.Base(path))[0] == '.'
}
<|endoftext|>


<commit_before>package k8s

import (
	"os"
	"time"

	//"github.com/keel-hq/keel/constants"
	"github.com/keel-hq/keel/internal/workgroup"
	"github.com/sirupsen/logrus"

	apps_v1 "k8s.io/api/apps/v1"
	v1beta1 "k8s.io/api/batch/v1beta1"

	v1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/fields"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/tools/cache"
)

const EnvNamespaceInstall = "NAMESPACE_INSTALL"

// WatchDeployments creates a SharedInformer for apps/v1.Deployments and registers it with g.
func WatchDeployments(g *workgroup.Group, client *kubernetes.Clientset, log logrus.FieldLogger, rs ...cache.ResourceEventHandler) {
	watch(g, client.AppsV1().RESTClient(), log, "deployments", new(apps_v1.Deployment), rs...)
}

// WatchStatefulSets creates a SharedInformer for apps/v1.StatefulSet and registers it with g.
func WatchStatefulSets(g *workgroup.Group, client *kubernetes.Clientset, log logrus.FieldLogger, rs ...cache.ResourceEventHandler) {
	watch(g, client.AppsV1().RESTClient(), log, "statefulsets", new(apps_v1.StatefulSet), rs...)
}

// WatchDaemonSets creates a SharedInformer for apps/v1.DaemonSet and registers it with g.
func WatchDaemonSets(g *workgroup.Group, client *kubernetes.Clientset, log logrus.FieldLogger, rs ...cache.ResourceEventHandler) {
	watch(g, client.AppsV1().RESTClient(), log, "daemonsets", new(apps_v1.DaemonSet), rs...)
}

// WatchCronJobs creates a SharedInformer for v1beta1.CronJob and registers it with g.
func WatchCronJobs(g *workgroup.Group, client *kubernetes.Clientset, log logrus.FieldLogger, rs ...cache.ResourceEventHandler) {
	watch(g, client.BatchV1beta1().RESTClient(), log, "cronjobs", new(v1beta1.CronJob), rs...)
}

func watch(g *workgroup.Group, c cache.Getter, log logrus.FieldLogger, resource string, objType runtime.Object, rs ...cache.ResourceEventHandler) {
	namespace_scan := EnvNamespaceInstall
	if os.Getenv(EnvNamespaceInstall) == "keel" {
		namespace_scan = v1.NamespaceAll
	} else if os.Getenv(EnvNamespaceInstall) == "" {
		namespace_scan = v1.NamespaceAll
	} else {
		namespace_scan = os.Getenv(EnvNamespaceInstall)
	}

	//ethos_namespace := "ns-team-3di-services-stage"
	lw := cache.NewListWatchFromClient(c, resource, namespace_scan, fields.Everything())
	sw := cache.NewSharedInformer(lw, objType, 30*time.Minute)
	for _, r := range rs {
		sw.AddEventHandler(r)
	}
	g.Add(func(stop <-chan struct{}) {
		log := log.WithField("resource", resource)
		log.Println("started")
		defer log.Println("stopped")
		sw.Run(stop)
	})
}

type buffer struct {
	ev chan interface{}
	logrus.StdLogger
	rh cache.ResourceEventHandler
}

type addEvent struct {
	obj interface{}
}

type updateEvent struct {
	oldObj, newObj interface{}
}

type deleteEvent struct {
	obj interface{}
}

// NewBuffer returns a ResourceEventHandler which buffers and serialises ResourceEventHandler events.
func NewBuffer(g *workgroup.Group, rh cache.ResourceEventHandler, log logrus.FieldLogger, size int) cache.ResourceEventHandler {
	buf := &buffer{
		ev:        make(chan interface{}, size),
		StdLogger: log.WithField("context", "buffer"),
		rh:        rh,
	}
	g.Add(buf.loop)
	return buf
}

func (b *buffer) loop(stop <-chan struct{}) {
	b.Println("started")
	defer b.Println("stopped")

	for {
		select {
		case ev := <-b.ev:
			switch ev := ev.(type) {
			case *addEvent:
				b.rh.OnAdd(ev.obj)
			case *updateEvent:
				b.rh.OnUpdate(ev.oldObj, ev.newObj)
			case *deleteEvent:
				b.rh.OnDelete(ev.obj)
			default:
				b.Printf("unhandled event type: %T: %v", ev, ev)
			}
		case <-stop:
			return
		}
	}
}

func (b *buffer) OnAdd(obj interface{}) {
	b.send(&addEvent{obj})
}

func (b *buffer) OnUpdate(oldObj, newObj interface{}) {
	b.send(&updateEvent{oldObj, newObj})
}

func (b *buffer) OnDelete(obj interface{}) {
	b.send(&deleteEvent{obj})
}

func (b *buffer) send(ev interface{}) {
	select {
	case b.ev <- ev:
		// all good
	default:
		b.Printf("event channel is full, len: %v, cap: %v", len(b.ev), cap(b.ev))
		b.ev <- ev
	}
}
<commit_msg>Adding the good constant + env<commit_after>package k8s

import (
	"os"
	"time"

	"github.com/keel-hq/keel/constants"
	"github.com/keel-hq/keel/internal/workgroup"
	"github.com/sirupsen/logrus"

	apps_v1 "k8s.io/api/apps/v1"
	v1beta1 "k8s.io/api/batch/v1beta1"

	v1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/fields"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/tools/cache"
)

// WatchDeployments creates a SharedInformer for apps/v1.Deployments and registers it with g.
func WatchDeployments(g *workgroup.Group, client *kubernetes.Clientset, log logrus.FieldLogger, rs ...cache.ResourceEventHandler) {
	watch(g, client.AppsV1().RESTClient(), log, "deployments", new(apps_v1.Deployment), rs...)
}

// WatchStatefulSets creates a SharedInformer for apps/v1.StatefulSet and registers it with g.
func WatchStatefulSets(g *workgroup.Group, client *kubernetes.Clientset, log logrus.FieldLogger, rs ...cache.ResourceEventHandler) {
	watch(g, client.AppsV1().RESTClient(), log, "statefulsets", new(apps_v1.StatefulSet), rs...)
}

// WatchDaemonSets creates a SharedInformer for apps/v1.DaemonSet and registers it with g.
func WatchDaemonSets(g *workgroup.Group, client *kubernetes.Clientset, log logrus.FieldLogger, rs ...cache.ResourceEventHandler) {
	watch(g, client.AppsV1().RESTClient(), log, "daemonsets", new(apps_v1.DaemonSet), rs...)
}

// WatchCronJobs creates a SharedInformer for v1beta1.CronJob and registers it with g.
func WatchCronJobs(g *workgroup.Group, client *kubernetes.Clientset, log logrus.FieldLogger, rs ...cache.ResourceEventHandler) {
	watch(g, client.BatchV1beta1().RESTClient(), log, "cronjobs", new(v1beta1.CronJob), rs...)
}

func watch(g *workgroup.Group, c cache.Getter, log logrus.FieldLogger, resource string, objType runtime.Object, rs ...cache.ResourceEventHandler) {
	//Check if the env var RESTRICTED_NAMESPACE is empty or equal to keel
	// If equal to keel or empty, the scan will be over all the cluster
	// If RESTRICTED_NAMESPACE is different than keel or empty, keel will scan in the defined namespace
	namespace_scan := "keel"
	if os.Getenv(constants.EnvRestrictedNamespace) == "keel" {
		namespace_scan = v1.NamespaceAll
	} else if os.Getenv(constants.EnvRestrictedNamespace) == "" {
		namespace_scan = v1.NamespaceAll
	} else {
		namespace_scan = os.Getenv(constants.EnvRestrictedNamespace)
	}

	lw := cache.NewListWatchFromClient(c, resource, namespace_scan, fields.Everything())
	sw := cache.NewSharedInformer(lw, objType, 30*time.Minute)
	for _, r := range rs {
		sw.AddEventHandler(r)
	}
	g.Add(func(stop <-chan struct{}) {
		log := log.WithField("resource", resource)
		log.Println("started")
		defer log.Println("stopped")
		sw.Run(stop)
	})
}

type buffer struct {
	ev chan interface{}
	logrus.StdLogger
	rh cache.ResourceEventHandler
}

type addEvent struct {
	obj interface{}
}

type updateEvent struct {
	oldObj, newObj interface{}
}

type deleteEvent struct {
	obj interface{}
}

// NewBuffer returns a ResourceEventHandler which buffers and serialises ResourceEventHandler events.
func NewBuffer(g *workgroup.Group, rh cache.ResourceEventHandler, log logrus.FieldLogger, size int) cache.ResourceEventHandler {
	buf := &buffer{
		ev:        make(chan interface{}, size),
		StdLogger: log.WithField("context", "buffer"),
		rh:        rh,
	}
	g.Add(buf.loop)
	return buf
}

func (b *buffer) loop(stop <-chan struct{}) {
	b.Println("started")
	defer b.Println("stopped")

	for {
		select {
		case ev := <-b.ev:
			switch ev := ev.(type) {
			case *addEvent:
				b.rh.OnAdd(ev.obj)
			case *updateEvent:
				b.rh.OnUpdate(ev.oldObj, ev.newObj)
			case *deleteEvent:
				b.rh.OnDelete(ev.obj)
			default:
				b.Printf("unhandled event type: %T: %v", ev, ev)
			}
		case <-stop:
			return
		}
	}
}

func (b *buffer) OnAdd(obj interface{}) {
	b.send(&addEvent{obj})
}

func (b *buffer) OnUpdate(oldObj, newObj interface{}) {
	b.send(&updateEvent{oldObj, newObj})
}

func (b *buffer) OnDelete(obj interface{}) {
	b.send(&deleteEvent{obj})
}

func (b *buffer) send(ev interface{}) {
	select {
	case b.ev <- ev:
		// all good
	default:
		b.Printf("event channel is full, len: %v, cap: %v", len(b.ev), cap(b.ev))
		b.ev <- ev
	}
}
<|endoftext|>