<commit_before>package lib

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

	log "github.com/sirupsen/logrus"

	cvmfs "github.com/cvmfs/ducc/cvmfs"
	da "github.com/cvmfs/ducc/docker-api"
	exec "github.com/cvmfs/ducc/exec"
	l "github.com/cvmfs/ducc/log"
)

func FindAllUsedFlatImages(CVMFSRepo string) ([]string, error) {
	root := filepath.Join("/", "cvmfs", CVMFSRepo)
	root_components := strings.Split(root, string(os.PathSeparator))
	result := make([]string, 0)
	walker := func(path string, info os.FileInfo, err error) error {
		// some kind of error, we don't really care and we just move on.
		if err != nil {
			l.LogE(err).WithFields(log.Fields{"path": path}).Warning("Error in opening the path, moving on.")
			return nil
		}
		components := strings.Split(path, string(os.PathSeparator))
		// first root directory, not sure if this ever happen
		if len(components) == len(root_components) {
			return nil
		}
		// checking if we are in a hidden directory
		// if we are, we skip it all
		first_dir := components[len(root_components)]
		if strings.HasPrefix(first_dir, ".") {
			return filepath.SkipDir
		}
		// let's check if we have reach a symlink
		// if we are in a symlink, we should capture
		// the image digest
		// we don't need to break the walk since Walk
		// does not follow symlinks
		if info.Mode()&os.ModeSymlink != 0 {
			realName, _ := filepath.EvalSymlinks(path)
			if err != nil {
				return nil
			}
			result = append(result, realName)
		}
		// general case we keep iterating
		return nil
	}
	filepath.Walk(root, walker)
	return result, nil
}

func FindAllFlatImages(CVMFSRepo string) ([]string, error) {
	root := filepath.Join("/", "cvmfs", CVMFSRepo, ".flat")
	root_components := strings.Split(root, string(os.PathSeparator))
	result := make([]string, 0)
	walker := func(path string, info os.FileInfo, err error) error {
		if err != nil {
			l.LogE(err).WithFields(log.Fields{"path": path}).Warning("Error in opening the path, moving on.")
			return nil
		}
		components := strings.Split(path, string(os.PathSeparator))
		if len(components) == len(root_components)+2 && info.IsDir() {
			result = append(result, path)
			return filepath.SkipDir
		}
		if len(components) > len(root_components)+2 {
			return filepath.SkipDir
		}
		if len(components) < len(root_components)+2 {
			return nil
		}
		// general case we keep iterating
		return nil
	}
	filepath.Walk(root, walker)
	return result, nil
}

func FindAllLayers(CVMFSRepo string) ([]string, error) {
	root := filepath.Join("/", "cvmfs", CVMFSRepo, ".layers")
	root_components := strings.Split(root, string(os.PathSeparator))
	result := make([]string, 0)
	walker := func(path string, info os.FileInfo, err error) error {
		if err != nil {
			l.LogE(err).WithFields(log.Fields{"path": path}).Warning("Error in opening the path, moving on.")
			return nil
		}
		components := strings.Split(path, string(os.PathSeparator))
		if len(components) == len(root_components)+2 && info.IsDir() {
			result = append(result, path)
			return filepath.SkipDir
		}
		if len(components) > len(root_components)+2 {
			return filepath.SkipDir
		}
		if len(components) < len(root_components)+2 {
			return nil
		}
		// general case we keep iterating
		return nil
	}
	filepath.Walk(root, walker)
	return result, nil
}

func FindAllUsedLayers(CVMFSRepo string) ([]string, error) {
	root := filepath.Join("/", "cvmfs", CVMFSRepo, ".metadata")
	result := make([]string, 0)
	walker := func(path string, info os.FileInfo, err error) error {
		if err != nil {
			l.LogE(err).WithFields(log.Fields{"path": path}).Warning("Error in opening the path, moving on.")
			return nil
		}
		if info.Name() == "manifest.json" {
			bytes, err := ioutil.ReadFile(path)
			if err != nil {
				return filepath.SkipDir
			}
			var manifest da.Manifest
			err = json.Unmarshal(bytes, &manifest)
			if err != nil {
				return filepath.SkipDir
			}
			for _, layerStruct := range manifest.Layers {
				layer := strings.Split(layerStruct.Digest, ":")[1]
				layerPath := filepath.Join("/", "cvmfs", CVMFSRepo, ".layers", layer[0:2], layer)
				result = append(result, layerPath)
			}
			return filepath.SkipDir
		}
		return nil
	}
	filepath.Walk(root, walker)
	return result, nil
}

func FindPodmanPathsToDelete(CVMFSRepo string, layersToDelete []string) ([]string, error) {
	podmanPathsToDelete := make([]string, 0)
	layerInfoPath := filepath.Join("/cvmfs", CVMFSRepo, "podmanStore", "overlay-layers", "layers.json")

	layersToDeleteMap := make(map[string]bool)
	for _, layerpath := range layersToDelete {
		components := strings.Split(layerpath, string(os.PathSeparator))
		layerid := components[len(components)-1]
		layersToDeleteMap[layerid] = true
	}

	_, err := os.Stat(layerInfoPath)
	if err == nil {
		layersdata := []LayerInfo{}
		file, err := ioutil.ReadFile(layerInfoPath)
		if err != nil {
			l.LogE(err).Error("Error in reading layers.json file")
			return podmanPathsToDelete, err
		}
		json.Unmarshal(file, &layersdata)

		newlayersdata := []LayerInfo{}
		for _, info := range layersdata {
			layerid := strings.Split(info.CompressedDiffDigest, ":")[1]
			if layersToDeleteMap[layerid] {
				podmanLayerPath := filepath.Join("/cvmfs", CVMFSRepo, "podmanStore", "overlay", info.ID)

				linkFilePath := filepath.Join(podmanLayerPath, "link")
				data, err := ioutil.ReadFile(linkFilePath)
				if err != nil {
					l.LogE(err).Error("Error in reading link file")
					return podmanPathsToDelete, err
				}
				id := string(data)

				linkDirPath := filepath.Join("/cvmfs", CVMFSRepo, "overlay", "l", id)
				podmanPathsToDelete = append(podmanPathsToDelete, podmanLayerPath)
				podmanPathsToDelete = append(podmanPathsToDelete, linkDirPath)
				continue
			}
			newlayersdata = append(newlayersdata, info)
		}

		layerInfo, err := json.MarshalIndent(newlayersdata, "", " ")
		if err != nil {
			l.LogE(err).Error("Error in marshaling json data for layers.json")
			return podmanPathsToDelete, err
		}

		err = cvmfs.WriteDataToCvmfs(CVMFSRepo, cvmfs.TrimCVMFSRepoPrefix(layerInfoPath), layerInfo)
		if err != nil {
			l.LogE(err).Error("Error in writing layers.json")
			return podmanPathsToDelete, err
		}
	}
	return podmanPathsToDelete, nil
}

func FindImageToGarbageCollect(CVMFSRepo string) ([]da.Manifest, error) {
	removeSchedulePath := cvmfs.RemoveScheduleLocation(CVMFSRepo)
	llog := func(l *log.Entry) *log.Entry {
		return l.WithFields(log.Fields{
			"action": "find image to garbage collect in schedule file",
			"file":   removeSchedulePath})
	}

	var schedule []da.Manifest

	_, err := os.Stat(removeSchedulePath)
	if os.IsNotExist(err) {
		return schedule, nil
	}
	if err != nil {
		llog(l.LogE(err)).Error("Error in stating the schedule file")
		return schedule, err
	}
	scheduleFileRO, err := os.Open(removeSchedulePath)
	if err != nil {
		llog(l.LogE(err)).Error("Error in opening the schedule file")
		return schedule, err
	}

	scheduleBytes, err := ioutil.ReadAll(scheduleFileRO)
	if err != nil {
		llog(l.LogE(err)).Error("Impossible to read the schedule file")
		return schedule, err
	}

	err = scheduleFileRO.Close()
	if err != nil {
		llog(l.LogE(err)).Error("Impossible to close the schedule file")
		return schedule, err
	}

	err = json.Unmarshal(scheduleBytes, &schedule)
	if err != nil {
		llog(l.LogE(err)).Error("Impossible to unmarshal the schedule file")
		return schedule, err
	}

	return schedule, nil
}

// with image and layer we pass the digest of the layer and the digest of the image,
// both without the sha256: prefix
func GarbageCollectSingleLayer(CVMFSRepo, image, layer string) error {
	backlink, err := cvmfs.GetBacklinkFromLayer(CVMFSRepo, layer)
	llog := func(l *log.Entry) *log.Entry {
		return l.WithFields(log.Fields{"action": "garbage collect layer",
			"repo":  CVMFSRepo,
			"image": image,
			"layer": layer})
	}
	if err != nil {
		llog(l.LogE(err)).Error("Impossible to retrieve the backlink information")
		return err
	}
	var newOrigin []string
	for _, origin := range backlink.Origin {
		withoutPrefix := strings.Split(origin, ":")[1]
		if withoutPrefix != image {
			newOrigin = append(newOrigin, origin)
		}
	}
	if len(newOrigin) > 0 {
		backlink.Origin = newOrigin
		backLinkMarshall, err := json.Marshal(backlink)
		if err != nil {
			llog(l.LogE(err)).Error("Error in marshaling the new backlink")
			return err
		}

		backlinkPath := cvmfs.GetBacklinkPath(CVMFSRepo, layer)

		err = exec.ExecCommand("cvmfs_server", "transaction", CVMFSRepo).Start()
		if err != nil {
			llog(l.LogE(err)).Error("Error in opening the transaction")
			return err
		}

		dir := filepath.Dir(backlinkPath)
		if _, err := os.Stat(dir); os.IsNotExist(err) {
			err = os.MkdirAll(dir, 0666)
			if err != nil {
				llog(l.LogE(err)).WithFields(log.Fields{"directory": dir}).Error(
					"Error in creating the directory for the backlinks file, skipping...")
				return err
			}
		}

		err = ioutil.WriteFile(backlinkPath, backLinkMarshall, 0666)
		if err != nil {
			llog(l.LogE(err)).WithFields(log.Fields{"file": backlinkPath}).Error(
				"Error in writing the backlink file, skipping...")
			return err
		}

		err = exec.ExecCommand("cvmfs_server", "publish", CVMFSRepo).Start()
		if err != nil {
			llog(l.LogE(err)).Error("Error in publishing after adding the backlinks")
			return err
		}
		// write it to file
		return nil
	} else {
		err = cvmfs.RemoveLayer(CVMFSRepo, layer)
		if err != nil {
			llog(l.LogE(err)).Error("Error in deleting the layer")
		}
		return err
	}
}
<commit_msg>remove raw use of exec and replace with cvmfs.WithinTransaction<commit_after>package lib

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

	log "github.com/sirupsen/logrus"

	cvmfs "github.com/cvmfs/ducc/cvmfs"
	da "github.com/cvmfs/ducc/docker-api"
	l "github.com/cvmfs/ducc/log"
)

func FindAllUsedFlatImages(CVMFSRepo string) ([]string, error) {
	root := filepath.Join("/", "cvmfs", CVMFSRepo)
	root_components := strings.Split(root, string(os.PathSeparator))
	result := make([]string, 0)
	walker := func(path string, info os.FileInfo, err error) error {
		// some kind of error, we don't really care and we just move on.
		if err != nil {
			l.LogE(err).WithFields(log.Fields{"path": path}).Warning("Error in opening the path, moving on.")
			return nil
		}
		components := strings.Split(path, string(os.PathSeparator))
		// first root directory, not sure if this ever happen
		if len(components) == len(root_components) {
			return nil
		}
		// checking if we are in a hidden directory
		// if we are, we skip it all
		first_dir := components[len(root_components)]
		if strings.HasPrefix(first_dir, ".") {
			return filepath.SkipDir
		}
		// let's check if we have reach a symlink
		// if we are in a symlink, we should capture
		// the image digest
		// we don't need to break the walk since Walk
		// does not follow symlinks
		if info.Mode()&os.ModeSymlink != 0 {
			realName, _ := filepath.EvalSymlinks(path)
			if err != nil {
				return nil
			}
			result = append(result, realName)
		}
		// general case we keep iterating
		return nil
	}
	filepath.Walk(root, walker)
	return result, nil
}

func FindAllFlatImages(CVMFSRepo string) ([]string, error) {
	root := filepath.Join("/", "cvmfs", CVMFSRepo, ".flat")
	root_components := strings.Split(root, string(os.PathSeparator))
	result := make([]string, 0)
	walker := func(path string, info os.FileInfo, err error) error {
		if err != nil {
			l.LogE(err).WithFields(log.Fields{"path": path}).Warning("Error in opening the path, moving on.")
			return nil
		}
		components := strings.Split(path, string(os.PathSeparator))
		if len(components) == len(root_components)+2 && info.IsDir() {
			result = append(result, path)
			return filepath.SkipDir
		}
		if len(components) > len(root_components)+2 {
			return filepath.SkipDir
		}
		if len(components) < len(root_components)+2 {
			return nil
		}
		// general case we keep iterating
		return nil
	}
	filepath.Walk(root, walker)
	return result, nil
}

func FindAllLayers(CVMFSRepo string) ([]string, error) {
	root := filepath.Join("/", "cvmfs", CVMFSRepo, ".layers")
	root_components := strings.Split(root, string(os.PathSeparator))
	result := make([]string, 0)
	walker := func(path string, info os.FileInfo, err error) error {
		if err != nil {
			l.LogE(err).WithFields(log.Fields{"path": path}).Warning("Error in opening the path, moving on.")
			return nil
		}
		components := strings.Split(path, string(os.PathSeparator))
		if len(components) == len(root_components)+2 && info.IsDir() {
			result = append(result, path)
			return filepath.SkipDir
		}
		if len(components) > len(root_components)+2 {
			return filepath.SkipDir
		}
		if len(components) < len(root_components)+2 {
			return nil
		}
		// general case we keep iterating
		return nil
	}
	filepath.Walk(root, walker)
	return result, nil
}

func FindAllUsedLayers(CVMFSRepo string) ([]string, error) {
	root := filepath.Join("/", "cvmfs", CVMFSRepo, ".metadata")
	result := make([]string, 0)
	walker := func(path string, info os.FileInfo, err error) error {
		if err != nil {
			l.LogE(err).WithFields(log.Fields{"path": path}).Warning("Error in opening the path, moving on.")
			return nil
		}
		if info.Name() == "manifest.json" {
			bytes, err := ioutil.ReadFile(path)
			if err != nil {
				return filepath.SkipDir
			}
			var manifest da.Manifest
			err = json.Unmarshal(bytes, &manifest)
			if err != nil {
				return filepath.SkipDir
			}
			for _, layerStruct := range manifest.Layers {
				layer := strings.Split(layerStruct.Digest, ":")[1]
				layerPath := filepath.Join("/", "cvmfs", CVMFSRepo, ".layers", layer[0:2], layer)
				result = append(result, layerPath)
			}
			return filepath.SkipDir
		}
		return nil
	}
	filepath.Walk(root, walker)
	return result, nil
}

func FindPodmanPathsToDelete(CVMFSRepo string, layersToDelete []string) ([]string, error) {
	podmanPathsToDelete := make([]string, 0)
	layerInfoPath := filepath.Join("/cvmfs", CVMFSRepo, "podmanStore", "overlay-layers", "layers.json")

	layersToDeleteMap := make(map[string]bool)
	for _, layerpath := range layersToDelete {
		components := strings.Split(layerpath, string(os.PathSeparator))
		layerid := components[len(components)-1]
		layersToDeleteMap[layerid] = true
	}

	_, err := os.Stat(layerInfoPath)
	if err == nil {
		layersdata := []LayerInfo{}
		file, err := ioutil.ReadFile(layerInfoPath)
		if err != nil {
			l.LogE(err).Error("Error in reading layers.json file")
			return podmanPathsToDelete, err
		}
		json.Unmarshal(file, &layersdata)

		newlayersdata := []LayerInfo{}
		for _, info := range layersdata {
			layerid := strings.Split(info.CompressedDiffDigest, ":")[1]
			if layersToDeleteMap[layerid] {
				podmanLayerPath := filepath.Join("/cvmfs", CVMFSRepo, "podmanStore", "overlay", info.ID)

				linkFilePath := filepath.Join(podmanLayerPath, "link")
				data, err := ioutil.ReadFile(linkFilePath)
				if err != nil {
					l.LogE(err).Error("Error in reading link file")
					return podmanPathsToDelete, err
				}
				id := string(data)

				linkDirPath := filepath.Join("/cvmfs", CVMFSRepo, "overlay", "l", id)
				podmanPathsToDelete = append(podmanPathsToDelete, podmanLayerPath)
				podmanPathsToDelete = append(podmanPathsToDelete, linkDirPath)
				continue
			}
			newlayersdata = append(newlayersdata, info)
		}

		layerInfo, err := json.MarshalIndent(newlayersdata, "", " ")
		if err != nil {
			l.LogE(err).Error("Error in marshaling json data for layers.json")
			return podmanPathsToDelete, err
		}

		err = cvmfs.WriteDataToCvmfs(CVMFSRepo, cvmfs.TrimCVMFSRepoPrefix(layerInfoPath), layerInfo)
		if err != nil {
			l.LogE(err).Error("Error in writing layers.json")
			return podmanPathsToDelete, err
		}
	}
	return podmanPathsToDelete, nil
}

func FindImageToGarbageCollect(CVMFSRepo string) ([]da.Manifest, error) {
	removeSchedulePath := cvmfs.RemoveScheduleLocation(CVMFSRepo)
	llog := func(l *log.Entry) *log.Entry {
		return l.WithFields(log.Fields{
			"action": "find image to garbage collect in schedule file",
			"file":   removeSchedulePath})
	}

	var schedule []da.Manifest

	_, err := os.Stat(removeSchedulePath)
	if os.IsNotExist(err) {
		return schedule, nil
	}
	if err != nil {
		llog(l.LogE(err)).Error("Error in stating the schedule file")
		return schedule, err
	}
	scheduleFileRO, err := os.Open(removeSchedulePath)
	if err != nil {
		llog(l.LogE(err)).Error("Error in opening the schedule file")
		return schedule, err
	}

	scheduleBytes, err := ioutil.ReadAll(scheduleFileRO)
	if err != nil {
		llog(l.LogE(err)).Error("Impossible to read the schedule file")
		return schedule, err
	}

	err = scheduleFileRO.Close()
	if err != nil {
		llog(l.LogE(err)).Error("Impossible to close the schedule file")
		return schedule, err
	}

	err = json.Unmarshal(scheduleBytes, &schedule)
	if err != nil {
		llog(l.LogE(err)).Error("Impossible to unmarshal the schedule file")
		return schedule, err
	}

	return schedule, nil
}

// with image and layer we pass the digest of the layer and the digest of the image,
// both without the sha256: prefix
func GarbageCollectSingleLayer(CVMFSRepo, image, layer string) error {
	backlink, err := cvmfs.GetBacklinkFromLayer(CVMFSRepo, layer)
	llog := func(l *log.Entry) *log.Entry {
		return l.WithFields(log.Fields{"action": "garbage collect layer",
			"repo":  CVMFSRepo,
			"image": image,
			"layer": layer})
	}
	if err != nil {
		llog(l.LogE(err)).Error("Impossible to retrieve the backlink information")
		return err
	}
	var newOrigin []string
	for _, origin := range backlink.Origin {
		withoutPrefix := strings.Split(origin, ":")[1]
		if withoutPrefix != image {
			newOrigin = append(newOrigin, origin)
		}
	}
	if len(newOrigin) > 0 {
		backlink.Origin = newOrigin
		backLinkMarshall, err := json.Marshal(backlink)
		if err != nil {
			llog(l.LogE(err)).Error("Error in marshaling the new backlink")
			return err
		}

		backlinkPath := cvmfs.GetBacklinkPath(CVMFSRepo, layer)

		err = cvmfs.WithinTransaction(CVMFSRepo, func() error {
			dir := filepath.Dir(backlinkPath)
			if _, err := os.Stat(dir); os.IsNotExist(err) {
				err = os.MkdirAll(dir, 0666)
				if err != nil {
					llog(l.LogE(err)).WithFields(log.Fields{"directory": dir}).Error(
						"Error in creating the directory for the backlinks file, skipping...")
					return err
				}
			}

			err = ioutil.WriteFile(backlinkPath, backLinkMarshall, 0666)
			if err != nil {
				llog(l.LogE(err)).WithFields(log.Fields{"file": backlinkPath}).Error(
					"Error in writing the backlink file, skipping...")
				return err
			}
			return nil
		})

		if err != nil {
			llog(l.LogE(err)).Error("Error in publishing after adding the backlinks")
			return err
		}
		// write it to file
		return nil
	} else {
		err = cvmfs.RemoveLayer(CVMFSRepo, layer)
		if err != nil {
			llog(l.LogE(err)).Error("Error in deleting the layer")
		}
		return err
	}
}
<|endoftext|>


<commit_before>package bongo

import (
	"encoding/json"
	"fmt"
	"strconv"
	"strings"

	"github.com/jinzhu/gorm"
)

// Fetch fetches the data from db by given parameters(fields of the struct)
func (b *Bongo) Fetch(i Modellable) error {
	if i.GetId() == 0 {
		return IdIsNotSet
	}

	if err := b.DB.Table(i.TableName()).
		Find(i).
		Error; err != nil {
		return err
	}

	return nil
}

// ById Fetches data from db by it's id
func (b *Bongo) ById(i Modellable, id int64) error {
	if err := b.DB.
		Table(i.TableName()).
		Where("id = ?", id).
		Find(i).
		Error; err != nil {
		return err
	}

	return nil
}

// Creates a new record with the given struct and its fields
func (b *Bongo) Create(i Modellable) error {
	if err := b.DB.Save(i).Error; err != nil {
		return err
	}

	return nil
}

// Update updates all fields of a struct with assigned data
func (b *Bongo) Update(i Modellable) error {
	if i.GetId() == 0 {
		return IdIsNotSet
	}

	// Update and Create is using the Save method, so they are
	// same functions but GORM handles, AfterCreate and AfterUpdate
	// in correct manner
	return b.Create(i)
}

// Delete deletes the data by it's id, it doesnt take any other fields
// into consideration
func (b *Bongo) Delete(i Modellable) error {
	if i.GetId() == 0 {
		return IdIsNotSet
	}

	if err := b.DB.Delete(i).Error; err != nil {
		return err
	}

	return nil
}

// FetchByIds fetches records by their ids and returns results in the same order
// as the ids; if no records in db we don't return error
func (b *Bongo) FetchByIds(i Modellable, data interface{}, ids []int64) error {
	if len(ids) == 0 {
		return nil
	}

	orderByQuery := ""
	comma := ""
	for _, id := range ids {
		orderByQuery = orderByQuery + comma + " id = " + strconv.FormatInt(id, 10) + " desc"
		comma = ","
	}

	// init query
	query := b.DB.Model(i)

	// add table name
	query = query.Table(i.TableName())

	query = query.Order(orderByQuery)

	query = query.Where(ids)

	query = query.Find(data)

	// supress not found errors
	return CheckErr(query)

}

func (b *Bongo) UpdatePartial(i Modellable, set map[string]interface{}) error {
	if i.GetId() == 0 {
		return IdIsNotSet
	}

	// init query
	query := b.DB

	query = query.Table(i.TableName())

	query = query.Where(i.GetId())

	if err := query.Update(set).Error; err != nil {
		return err
	}

	if err := b.Fetch(i); err != nil {
		return err
	}

	b.AfterUpdate(i)
	return nil
}

// selector, set
func (b *Bongo) UpdateMulti(i Modellable, rest ...map[string]interface{}) error {
	var set, selector map[string]interface{}

	switch len(rest) {
	case 1:
		set = rest[0]
		selector = nil
	case 2:
		selector = rest[0]
		set = rest[1]
	default:
		return WrongParameter
	}

	query := b.DB.Table(i.TableName())

	//add selector
	query = addWhere(query, selector)

	if err := query.Updates(set).Error; err != nil {
		return err
	}

	return nil
}

func (b *Bongo) Count(i Modellable, where ...interface{}) (int, error) {
	var count int

	// init query
	query := b.DB.Model(i)

	// add table name
	query = query.Table(i.TableName())

	// add query
	query = query.Where(where[0], where[1:len(where)]...)

	return count, query.Count(&count).Error
}

func (b *Bongo) CountWithQuery(i Modellable, q *Query) (int, error) {
	query := b.BuildQuery(i, q)
	var count int
	return count, query.Count(&count).Error
}

type Scope func(d *gorm.DB) *gorm.DB

type Query struct {
	Selector   map[string]interface{}
	Sort       map[string]string
	Pluck      string
	Pagination Pagination
	Scopes     []Scope
}

func (q *Query) AddScope(scope Scope) {
	if q.Scopes == nil {
		q.Scopes = make([]Scope, 0)
	}

	q.Scopes = append(q.Scopes, scope)
}

type Pagination struct {
	Limit int
	Skip  int
}

func NewPagination(limit int, skip int) *Pagination {
	return &Pagination{
		Limit: limit,
		Skip:  skip,
	}
}

func NewQS(selector map[string]interface{}) *Query {
	return &Query{
		Selector: selector,
	}
}

// selector, sort, limit, pluck,
func (b *Bongo) Some(i Modellable, data interface{}, q *Query) error {
	err := b.executeQuery(i, data, q)
	if err == gorm.RecordNotFound {
		return nil
	}
	return err
}

func (b *Bongo) One(i Modellable, data interface{}, q *Query) error {
	q.Pagination.Limit = 1
	return b.executeQuery(i, data, q)
}

func (b *Bongo) BuildQuery(i Modellable, q *Query) *gorm.DB {
	// init query
	query := b.DB.Model(i)

	// add table name
	query = query.Table(i.TableName())

	// add sort options
	query = addSort(query, q.Sort)

	query = addSkip(query, q.Pagination.Skip)

	query = addLimit(query, q.Pagination.Limit)

	// add selector
	query = addWhere(query, q.Selector)

	// put scopes
	if q.Scopes != nil && len(q.Scopes) > 0 {
		for _, scope := range q.Scopes {
			query = query.Scopes(scope)
		}
	}

	return query
}

func (b *Bongo) executeQuery(i Modellable, data interface{}, q *Query) error {
	// init query
	query := b.BuildQuery(i, q)

	var err error
	// TODO refactor this part
	if q.Pluck != "" {
		if strings.Contains(q.Pluck, ",") {
			// add pluck data
			query = addPluck(query, q.Pluck)

			err = query.Find(data).Error
		} else {
			err = query.Pluck(q.Pluck, data).Error
		}
	} else {
		err = query.Find(data).Error
	}

	return err
}

func (b *Bongo) PublishEvent(eventName string, i Modellable) error {
	data, err := json.Marshal(i)
	if err != nil {
		b.log.Error("Error while marshalling for publish %s", err)
		return err
	}

	err = b.Broker.Publish(i.TableName()+"_"+eventName, data)
	if err != nil {
		b.log.Error("Error while publishing %s", err)
		return err
	}

	return nil
}

func (b *Bongo) AfterCreate(i Modellable) {
	b.PublishEvent("created", i)
}

func (b *Bongo) AfterUpdate(i Modellable) {
	b.PublishEvent("updated", i)
}

func (b *Bongo) AfterDelete(i Modellable) {
	b.PublishEvent("deleted", i)
}

// addSort injects sort parameters into query
func addSort(query *gorm.DB, options map[string]string) *gorm.DB {

	if options == nil {
		return query
	}

	if len(options) == 0 {
		return query
	}

	var opts []string
	for key, val := range options {
		opts = append(opts, fmt.Sprintf("%s %v", key, val))
	}
	return query.Order(strings.Join(opts, ","))
}

// addPluck basically adds select statement for
// only required fields
func addPluck(query *gorm.DB, plucked string) *gorm.DB {
	if plucked == "" {
		return query
	}

	return query.Select(plucked)
}

// addWhere adds where query
func addWhere(query *gorm.DB, selector map[string]interface{}) *gorm.DB {
	if selector == nil {
		return query
	}

	// instead sending one selector, do chaining here
	return query.Where(selector)
}

// addSkip adds skip parameter into sql query
func addSkip(query *gorm.DB, skip int) *gorm.DB {
	if skip > 0 {
		return query.Offset(skip)
	}

	return query
}

// addLimit adds limit into query if set
func addLimit(query *gorm.DB, limit int) *gorm.DB {
	// if limit is minus or 0 ignore
	if limit > 0 {
		return query.Limit(limit)
	}

	return query
}

// CheckErr checks error exitence and returns
// if found, but this function suppress RecordNotFound errors
func CheckErr(res *gorm.DB) error {
	if res == nil {
		return nil
	}

	if res.Error == nil {
		return nil
	}

	if res.Error == gorm.RecordNotFound {
		return nil
	}

	return res.Error
}
<commit_msg>Bongo: inject caching into bongo<commit_after>package bongo

import (
	"encoding/json"
	"fmt"
	"strconv"
	"strings"

	"github.com/jinzhu/gorm"
)

// Fetch fetches the data from db by given parameters(fields of the struct)
func (b *Bongo) Fetch(i Modellable, id int64) error {
	if id == 0 {
		return IdIsNotSet
	}

	if data, ok := i.(Cacher); ok {
		err := b.GetFromCache(i, data, id)
		if err == nil {
			return nil
		}
		if err == RecordNotFound {
			return err
		}
	}

	if err := b.DB.Table(i.TableName()).
		Find(i).
		Error; err != nil {
		return err
	}

	return nil
}

// ById Fetches data from db by it's id
func (b *Bongo) ById(i Modellable, id int64) error {
	if err := b.DB.
		Table(i.TableName()).
		Where("id = ?", id).
		Find(i).
		Error; err != nil {
		return err
	}

	return nil
}

// Creates a new record with the given struct and its fields
func (b *Bongo) Create(i Modellable) error {
	if err := b.DB.Save(i).Error; err != nil {
		return err
	}

	return nil
}

// Update updates all fields of a struct with assigned data
func (b *Bongo) Update(i Modellable) error {
	if i.GetId() == 0 {
		return IdIsNotSet
	}

	// Update and Create is using the Save method, so they are
	// same functions but GORM handles, AfterCreate and AfterUpdate
	// in correct manner
	return b.Create(i)
}

// Delete deletes the data by it's id, it doesnt take any other fields
// into consideration
func (b *Bongo) Delete(i Modellable) error {
	if i.GetId() == 0 {
		return IdIsNotSet
	}

	if err := b.DB.Delete(i).Error; err != nil {
		return err
	}

	return nil
}

// FetchByIds fetches records by their ids and returns results in the same order
// as the ids; if no records in db we don't return error
func (b *Bongo) FetchByIds(i Modellable, data interface{}, ids []int64) error {
	if len(ids) == 0 {
		return nil
	}

	orderByQuery := ""
	comma := ""
	for _, id := range ids {
		orderByQuery = orderByQuery + comma + " id = " + strconv.FormatInt(id, 10) + " desc"
		comma = ","
	}

	// init query
	query := b.DB.Model(i)

	// add table name
	query = query.Table(i.TableName())

	query = query.Order(orderByQuery)

	query = query.Where(ids)

	query = query.Find(data)

	// supress not found errors
	return CheckErr(query)

}

func (b *Bongo) UpdatePartial(i Modellable, set map[string]interface{}) error {
	if i.GetId() == 0 {
		return IdIsNotSet
	}

	// init query
	query := b.DB

	query = query.Table(i.TableName())

	query = query.Where(i.GetId())

	if err := query.Update(set).Error; err != nil {
		return err
	}

	if err := b.ById(i, i.GetId()); err != nil {
		return err
	}

	b.AfterUpdate(i)
	return nil
}

// selector, set
func (b *Bongo) UpdateMulti(i Modellable, rest ...map[string]interface{}) error {
	var set, selector map[string]interface{}

	switch len(rest) {
	case 1:
		set = rest[0]
		selector = nil
	case 2:
		selector = rest[0]
		set = rest[1]
	default:
		return WrongParameter
	}

	query := b.DB.Table(i.TableName())

	//add selector
	query = addWhere(query, selector)

	if err := query.Updates(set).Error; err != nil {
		return err
	}

	return nil
}

func (b *Bongo) Count(i Modellable, where ...interface{}) (int, error) {
	var count int

	// init query
	query := b.DB.Model(i)

	// add table name
	query = query.Table(i.TableName())

	// add query
	query = query.Where(where[0], where[1:len(where)]...)

	return count, query.Count(&count).Error
}

func (b *Bongo) CountWithQuery(i Modellable, q *Query) (int, error) {
	query := b.BuildQuery(i, q)
	var count int
	return count, query.Count(&count).Error
}

type Scope func(d *gorm.DB) *gorm.DB

type Query struct {
	Selector   map[string]interface{}
	Sort       map[string]string
	Pluck      string
	Pagination Pagination
	Scopes     []Scope
}

func (q *Query) AddScope(scope Scope) {
	if q.Scopes == nil {
		q.Scopes = make([]Scope, 0)
	}

	q.Scopes = append(q.Scopes, scope)
}

type Pagination struct {
	Limit int
	Skip  int
}

func NewPagination(limit int, skip int) *Pagination {
	return &Pagination{
		Limit: limit,
		Skip:  skip,
	}
}

func NewQS(selector map[string]interface{}) *Query {
	return &Query{
		Selector: selector,
	}
}

// selector, sort, limit, pluck,
func (b *Bongo) Some(i Modellable, data interface{}, q *Query) error {
	err := b.executeQuery(i, data, q)
	if err == gorm.RecordNotFound {
		return nil
	}
	return err
}

func (b *Bongo) One(i Modellable, data interface{}, q *Query) error {
	q.Pagination.Limit = 1
	return b.executeQuery(i, data, q)
}

func (b *Bongo) BuildQuery(i Modellable, q *Query) *gorm.DB {
	// init query
	query := b.DB.Model(i)

	// add table name
	query = query.Table(i.TableName())

	// add sort options
	query = addSort(query, q.Sort)

	query = addSkip(query, q.Pagination.Skip)

	query = addLimit(query, q.Pagination.Limit)

	// add selector
	query = addWhere(query, q.Selector)

	// put scopes
	if q.Scopes != nil && len(q.Scopes) > 0 {
		for _, scope := range q.Scopes {
			query = query.Scopes(scope)
		}
	}

	return query
}

func (b *Bongo) executeQuery(i Modellable, data interface{}, q *Query) error {
	// init query
	query := b.BuildQuery(i, q)

	var err error
	// TODO refactor this part
	if q.Pluck != "" {
		if strings.Contains(q.Pluck, ",") {
			// add pluck data
			query = addPluck(query, q.Pluck)

			err = query.Find(data).Error
		} else {
			err = query.Pluck(q.Pluck, data).Error
		}
	} else {
		err = query.Find(data).Error
	}

	return err
}

func (b *Bongo) PublishEvent(eventName string, i Modellable) error {
	data, err := json.Marshal(i)
	if err != nil {
		b.log.Error("Error while marshalling for publish %s", err)
		return err
	}

	err = b.Broker.Publish(i.TableName()+"_"+eventName, data)
	if err != nil {
		b.log.Error("Error while publishing %s", err)
		return err
	}

	return nil
}

func (b *Bongo) AfterCreate(i Modellable) error {
	b.PublishEvent("created", i)
	return b.AddToCache(i)
}

func (b *Bongo) AfterUpdate(i Modellable) error {
	b.PublishEvent("updated", i)
	return b.AddToCache(i)
}

func (b *Bongo) AfterDelete(i Modellable) error {
	b.PublishEvent("deleted", i)
	return b.AddToCache(i)
}

func (b *Bongo) AddToCache(i Modellable) error {
	data, ok := i.(Cacher)
	if !ok {
		return nil
	}

	err := b.SetToCache(data)
	if err != ErrCacheIsNotEnabled {
		return err
	}

	return nil
}

// addSort injects sort parameters into query
func addSort(query *gorm.DB, options map[string]string) *gorm.DB {

	if options == nil {
		return query
	}

	if len(options) == 0 {
		return query
	}

	var opts []string
	for key, val := range options {
		opts = append(opts, fmt.Sprintf("%s %v", key, val))
	}
	return query.Order(strings.Join(opts, ","))
}

// addPluck basically adds select statement for
// only required fields
func addPluck(query *gorm.DB, plucked string) *gorm.DB {
	if plucked == "" {
		return query
	}

	return query.Select(plucked)
}

// addWhere adds where query
func addWhere(query *gorm.DB, selector map[string]interface{}) *gorm.DB {
	if selector == nil {
		return query
	}

	// instead sending one selector, do chaining here
	return query.Where(selector)
}

// addSkip adds skip parameter into sql query
func addSkip(query *gorm.DB, skip int) *gorm.DB {
	if skip > 0 {
		return query.Offset(skip)
	}

	return query
}

// addLimit adds limit into query if set
func addLimit(query *gorm.DB, limit int) *gorm.DB {
	// if limit is minus or 0 ignore
	if limit > 0 {
		return query.Limit(limit)
	}

	return query
}

// CheckErr checks error exitence and returns
// if found, but this function suppress RecordNotFound errors
func CheckErr(res *gorm.DB) error {
	if res == nil {
		return nil
	}

	if res.Error == nil {
		return nil
	}

	if res.Error == gorm.RecordNotFound {
		return nil
	}

	return res.Error
}
<|endoftext|>


<commit_before>package etcd_test

import (
	"flag"
	"os"
	"testing"
	"time"

	updater "github.com/mwitkow-io/go-flagz/etcd"
	"github.com/mwitkow-io/go-flagz/test_etcd"

	"github.com/Sirupsen/logrus"
	etcd "github.com/coreos/etcd/client"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/suite"
	"golang.org/x/net/context"
)

const (
	prefix = "/updater_test/"
)

var (
	logger = logrus.StandardLogger()
	ctxNil = context.Background()
)

// Define the suite, and absorb the built-in basic suite
// functionality from testify - including assertion methods.
type UpdaterTestSuite struct {
	suite.Suite
	keys etcd.KeysAPI

	flagSet *flag.FlagSet
	updater *updater.Updater
}

// Clean up the etcd state before each test.
func (s *UpdaterTestSuite) SetupTest() {
	s.keys.Delete(ctxNil, prefix, &etcd.DeleteOptions{Dir: true, Recursive: true})
	_, err := s.keys.Set(ctxNil, prefix, "", &etcd.SetOptions{Dir: true})
	if err != nil {
		s.T().Fatalf("cannot create empty dir %v: %v", prefix, err)
	}
	s.flagSet = flag.NewFlagSet("updater_test", flag.ContinueOnError)
	s.updater, err = updater.New(s.flagSet, s.keys, prefix, &testingLog{T: s.T()})
	if err != nil {
		s.T().Fatalf("cannot create updater: %v", err)
	}
}

func (s *UpdaterTestSuite) setFlagzValue(flagzName string, value string) {
	_, err := s.keys.Set(ctxNil, prefix+flagzName, value, &etcd.SetOptions{})
	if err != nil {
		s.T().Fatalf("failed setting flagz value: %v")
	}
}

func (s *UpdaterTestSuite) getFlagzValue(flagzName string) string {
	resp, err := s.keys.Get(ctxNil, prefix+flagzName, &etcd.GetOptions{})
	if err != nil {
		s.T().Fatalf("failed getting flagz value: %v")
	}
	return resp.Node.Value
}

// Tear down the updater
func (s *UpdaterTestSuite) TearDownTest() {
	s.updater.Stop()
	time.Sleep(100 * time.Millisecond)
}

func (s *UpdaterTestSuite) Test_ErrorsOnInitialUnknownFlag() {
	s.flagSet.Int("someint", 1337, "some int usage")
	s.setFlagzValue("anotherint", "999")
	s.Require().Error(s.updater.Initialize(), "initialize should complain about unknown flag")
}

func (s *UpdaterTestSuite) Test_SetsInitialValues() {
	someInt := s.flagSet.Int("someint", 1337, "some int usage")
	someString := s.flagSet.String("somestring", "initial_value", "some int usage")
	anotherString := s.flagSet.String("anotherstring", "default_value", "some int usage")
	s.setFlagzValue("someint", "2015")
	s.setFlagzValue("somestring", "changed_value")
	s.Require().NoError(s.updater.Initialize())

	s.Require().Equal(2015, *someInt, "int flag should change value")
	s.Require().Equal("changed_value", *someString, "string flag should change value")
	s.Require().Equal("default_value", *anotherString, "anotherstring should be unchanged")
}

func (s *UpdaterTestSuite) Test_DynamicUpdate() {
	someInt := s.flagSet.Int("someint", 1337, "some int usage")
	s.Require().NoError(s.updater.Initialize())
	s.Require().NoError(s.updater.Start())
	s.Require().Equal(1337, *someInt, "int flag should not change value")
	s.setFlagzValue("someint", "2014")
	eventually(s.T(), 1*time.Second,
		assert.Equal, 2014,
		func() interface{} { return *someInt },
		"someint value should change")
	s.setFlagzValue("someint", "2015")
	eventually(s.T(), 1*time.Second,
		assert.Equal, 2015,
		func() interface{} { return *someInt },
		"someint value should change")
	s.setFlagzValue("someint", "2016")
	eventually(s.T(), 1*time.Second,
		assert.Equal, 2016,
		func() interface{} { return *someInt },
		"someint value should change")
}

func (s *UpdaterTestSuite) Test_DynamicUpdateRestoresGoodState() {
	someInt := s.flagSet.Int("someint", 1337, "some int usage")
	someFloat := s.flagSet.Float64("somefloat", 1.337, "some int usage")
	s.setFlagzValue("someint", "2015")
	s.Require().NoError(s.updater.Initialize())
	s.Require().NoError(s.updater.Start())
	s.Require().Equal(2015, *someInt, "int flag should change value")
	s.Require().Equal(1.337, *someFloat, "float flag should not change value")

	// Bad update causing a rollback.
	s.setFlagzValue("someint", "randombleh")
	eventually(s.T(), 1*time.Second,
		assert.Equal,
		"2015",
		func() interface{} {
			return s.getFlagzValue("someint")
		},
		"someint failure should revert etcd value to 2015")

	// Make sure we can continue updating.
	s.setFlagzValue("someint", "2016")
	s.setFlagzValue("somefloat", "3.14")
	eventually(s.T(), 1*time.Second,
		assert.Equal, 2016,
		func() interface{} { return *someInt },
		"someint value should change, after rolled back")
	eventually(s.T(), 1*time.Second,
		assert.Equal, 3.14,
		func() interface{} { return *someFloat },
		"somefloat value should change")

}

func TestUpdaterSuite(t *testing.T) {
	server, err := test_etcd.New(os.Stderr)
	if err != nil {
		t.Fatalf("failed starting test server: %v", err)
	}
	t.Logf("will use etcd test endpoint: %v", server.Endpoint)
	defer func() {
		server.Stop()
		t.Logf("cleaned up etcd test server")
	}()
	suite.Run(t, &UpdaterTestSuite{keys: etcd.NewKeysAPI(server.Client)})
}

type assertFunc func(T assert.TestingT, expected, actual interface{}, msgAndArgs ...interface{}) bool
type getter func() interface{}

// eventually tries a given Assert function 5 times over the period of time.
func eventually(T *testing.T, duration time.Duration,
	af assertFunc, expected interface{}, actual getter, msgAndArgs ...interface{}) {
	increment := duration / 5
	for i := 0; i < 5; i++ {
		time.Sleep(increment)
		if af(T, expected, actual(), msgAndArgs...) {
			return
		}
	}
	T.FailNow()
}

// Abstraction that allows us to pass the *testing.T as a logger to the updater.
type testingLog struct {
	T *testing.T
}

func (tl *testingLog) Printf(format string, v ...interface{}) {
	tl.T.Logf(format+"\n", v...)
}
<commit_msg>Fix mwitkow-io to mwitkow path.<commit_after>package etcd_test

import (
	"flag"
	"os"
	"testing"
	"time"

	updater "github.com/mwitkow/go-flagz/etcd"
	"github.com/mwitkow/go-flagz/test_etcd"

	"github.com/Sirupsen/logrus"
	etcd "github.com/coreos/etcd/client"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/suite"
	"golang.org/x/net/context"
)

const (
	prefix = "/updater_test/"
)

var (
	logger = logrus.StandardLogger()
	ctxNil = context.Background()
)

// Define the suite, and absorb the built-in basic suite
// functionality from testify - including assertion methods.
type UpdaterTestSuite struct {
	suite.Suite
	keys etcd.KeysAPI

	flagSet *flag.FlagSet
	updater *updater.Updater
}

// Clean up the etcd state before each test.
func (s *UpdaterTestSuite) SetupTest() {
	s.keys.Delete(ctxNil, prefix, &etcd.DeleteOptions{Dir: true, Recursive: true})
	_, err := s.keys.Set(ctxNil, prefix, "", &etcd.SetOptions{Dir: true})
	if err != nil {
		s.T().Fatalf("cannot create empty dir %v: %v", prefix, err)
	}
	s.flagSet = flag.NewFlagSet("updater_test", flag.ContinueOnError)
	s.updater, err = updater.New(s.flagSet, s.keys, prefix, &testingLog{T: s.T()})
	if err != nil {
		s.T().Fatalf("cannot create updater: %v", err)
	}
}

func (s *UpdaterTestSuite) setFlagzValue(flagzName string, value string) {
	_, err := s.keys.Set(ctxNil, prefix+flagzName, value, &etcd.SetOptions{})
	if err != nil {
		s.T().Fatalf("failed setting flagz value: %v")
	}
}

func (s *UpdaterTestSuite) getFlagzValue(flagzName string) string {
	resp, err := s.keys.Get(ctxNil, prefix+flagzName, &etcd.GetOptions{})
	if err != nil {
		s.T().Fatalf("failed getting flagz value: %v")
	}
	return resp.Node.Value
}

// Tear down the updater
func (s *UpdaterTestSuite) TearDownTest() {
	s.updater.Stop()
	time.Sleep(100 * time.Millisecond)
}

func (s *UpdaterTestSuite) Test_ErrorsOnInitialUnknownFlag() {
	s.flagSet.Int("someint", 1337, "some int usage")
	s.setFlagzValue("anotherint", "999")
	s.Require().Error(s.updater.Initialize(), "initialize should complain about unknown flag")
}

func (s *UpdaterTestSuite) Test_SetsInitialValues() {
	someInt := s.flagSet.Int("someint", 1337, "some int usage")
	someString := s.flagSet.String("somestring", "initial_value", "some int usage")
	anotherString := s.flagSet.String("anotherstring", "default_value", "some int usage")
	s.setFlagzValue("someint", "2015")
	s.setFlagzValue("somestring", "changed_value")
	s.Require().NoError(s.updater.Initialize())

	s.Require().Equal(2015, *someInt, "int flag should change value")
	s.Require().Equal("changed_value", *someString, "string flag should change value")
	s.Require().Equal("default_value", *anotherString, "anotherstring should be unchanged")
}

func (s *UpdaterTestSuite) Test_DynamicUpdate() {
	someInt := s.flagSet.Int("someint", 1337, "some int usage")
	s.Require().NoError(s.updater.Initialize())
	s.Require().NoError(s.updater.Start())
	s.Require().Equal(1337, *someInt, "int flag should not change value")
	s.setFlagzValue("someint", "2014")
	eventually(s.T(), 1*time.Second,
		assert.Equal, 2014,
		func() interface{} { return *someInt },
		"someint value should change")
	s.setFlagzValue("someint", "2015")
	eventually(s.T(), 1*time.Second,
		assert.Equal, 2015,
		func() interface{} { return *someInt },
		"someint value should change")
	s.setFlagzValue("someint", "2016")
	eventually(s.T(), 1*time.Second,
		assert.Equal, 2016,
		func() interface{} { return *someInt },
		"someint value should change")
}

func (s *UpdaterTestSuite) Test_DynamicUpdateRestoresGoodState() {
	someInt := s.flagSet.Int("someint", 1337, "some int usage")
	someFloat := s.flagSet.Float64("somefloat", 1.337, "some int usage")
	s.setFlagzValue("someint", "2015")
	s.Require().NoError(s.updater.Initialize())
	s.Require().NoError(s.updater.Start())
	s.Require().Equal(2015, *someInt, "int flag should change value")
	s.Require().Equal(1.337, *someFloat, "float flag should not change value")

	// Bad update causing a rollback.
	s.setFlagzValue("someint", "randombleh")
	eventually(s.T(), 1*time.Second,
		assert.Equal,
		"2015",
		func() interface{} {
			return s.getFlagzValue("someint")
		},
		"someint failure should revert etcd value to 2015")

	// Make sure we can continue updating.
	s.setFlagzValue("someint", "2016")
	s.setFlagzValue("somefloat", "3.14")
	eventually(s.T(), 1*time.Second,
		assert.Equal, 2016,
		func() interface{} { return *someInt },
		"someint value should change, after rolled back")
	eventually(s.T(), 1*time.Second,
		assert.Equal, 3.14,
		func() interface{} { return *someFloat },
		"somefloat value should change")

}

func TestUpdaterSuite(t *testing.T) {
	server, err := test_etcd.New(os.Stderr)
	if err != nil {
		t.Fatalf("failed starting test server: %v", err)
	}
	t.Logf("will use etcd test endpoint: %v", server.Endpoint)
	defer func() {
		server.Stop()
		t.Logf("cleaned up etcd test server")
	}()
	suite.Run(t, &UpdaterTestSuite{keys: etcd.NewKeysAPI(server.Client)})
}

type assertFunc func(T assert.TestingT, expected, actual interface{}, msgAndArgs ...interface{}) bool
type getter func() interface{}

// eventually tries a given Assert function 5 times over the period of time.
func eventually(T *testing.T, duration time.Duration,
	af assertFunc, expected interface{}, actual getter, msgAndArgs ...interface{}) {
	increment := duration / 5
	for i := 0; i < 5; i++ {
		time.Sleep(increment)
		if af(T, expected, actual(), msgAndArgs...) {
			return
		}
	}
	T.FailNow()
}

// Abstraction that allows us to pass the *testing.T as a logger to the updater.
type testingLog struct {
	T *testing.T
}

func (tl *testingLog) Printf(format string, v ...interface{}) {
	tl.T.Logf(format+"\n", v...)
}
<|endoftext|>


<commit_before>/*
Copyright 2017 Google Inc.

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 vtexplain

import (
	"fmt"
	"io/ioutil"
	"os/exec"
	"strings"
	"testing"

	"github.com/youtube/vitess/go/testfiles"
)

var testOutputTempDir string

func defaultTestOpts() *Options {
	return &Options{
		ReplicationMode: "ROW",
		NumShards:       4,
		Normalize:       true,
	}
}

func initTest(opts *Options, t *testing.T) {
	schema, err := ioutil.ReadFile(testfiles.Locate("vtexplain/test-schema.sql"))
	if err != nil {
		t.Fatalf("error: %v", err)
	}

	vSchema, err := ioutil.ReadFile(testfiles.Locate("vtexplain/test-vschema.json"))
	if err != nil {
		t.Fatalf("error: %v", err)
	}

	err = Init(string(vSchema), string(schema), opts)
	if err != nil {
		t.Fatalf("vtexplain Init error: %v", err)
	}

}

func testExplain(testcase string, opts *Options, t *testing.T) {
	// TODO(sougou): remove this skip for manual testing.
	// We need to find a better testing strategy.
	// These tests are almost always failing on travis.
	t.Skip()

	initTest(opts, t)

	sqlFile := testfiles.Locate(fmt.Sprintf("vtexplain/%s-queries.sql", testcase))
	sql, err := ioutil.ReadFile(sqlFile)

	jsonOutFile := testfiles.Locate(fmt.Sprintf("vtexplain/%s-output.json", testcase))
	jsonOut, err := ioutil.ReadFile(jsonOutFile)

	textOutFile := testfiles.Locate(fmt.Sprintf("vtexplain/%s-output.txt", testcase))
	textOut, err := ioutil.ReadFile(textOutFile)

	explains, err := Run(string(sql))
	if err != nil {
		t.Fatalf("vtexplain error: %v", err)
	}
	if explains == nil {
		t.Fatalf("vtexplain error running %s: no explain", string(sql))
	}

	explainJSON := ExplainsAsJSON(explains)
	if strings.TrimSpace(string(explainJSON)) != strings.TrimSpace(string(jsonOut)) {
		// Print the json that was actually returned and also dump to a
		// temp file to be able to diff the results.
		t.Errorf("json output did not match")
		t.Logf("got:\n%s\n", string(explainJSON))

		if testOutputTempDir == "" {
			testOutputTempDir, err = ioutil.TempDir("", "vtexplain_output")
			if err != nil {
				t.Fatalf("error getting tempdir: %v", err)
			}
		}
		gotFile := fmt.Sprintf("%s/%s-output.json", testOutputTempDir, testcase)
		ioutil.WriteFile(gotFile, []byte(explainJSON), 0644)

		command := exec.Command("diff", "-u", jsonOutFile, gotFile)
		out, _ := command.CombinedOutput()
		t.Logf("diff:\n%s\n", out)
	}

	explainText := ExplainsAsText(explains)
	if strings.TrimSpace(string(explainText)) != strings.TrimSpace(string(textOut)) {
		// Print the Text that was actually returned and also dump to a
		// temp file to be able to diff the results.
		t.Errorf("Text output did not match")
		t.Logf("got:\n%s\n", string(explainText))

		if testOutputTempDir == "" {
			testOutputTempDir, err = ioutil.TempDir("", "vtexplain_output")
			if err != nil {
				t.Fatalf("error getting tempdir: %v", err)
			}
		}
		gotFile := fmt.Sprintf("%s/%s-output.txt", testOutputTempDir, testcase)
		ioutil.WriteFile(gotFile, []byte(explainText), 0644)

		command := exec.Command("diff", "-u", textOutFile, gotFile)
		out, _ := command.CombinedOutput()
		t.Logf("diff:\n%s\n", out)
	}
}

func TestUnsharded(t *testing.T) {
	testExplain("unsharded", defaultTestOpts(), t)
}

func TestSelectSharded(t *testing.T) {
	testExplain("selectsharded", defaultTestOpts(), t)
}

func TestInsertSharded(t *testing.T) {
	testExplain("insertsharded", defaultTestOpts(), t)
}

func TestUpdateSharded(t *testing.T) {
	testExplain("updatesharded", defaultTestOpts(), t)
}

func TestDeleteSharded(t *testing.T) {
	testExplain("deletesharded", defaultTestOpts(), t)
}

func TestOptions(t *testing.T) {
	opts := &Options{
		ReplicationMode: "STATEMENT",
		NumShards:       4,
		Normalize:       false,
	}

	testExplain("options", opts, t)
}

func TestComments(t *testing.T) {
	testExplain("comments", defaultTestOpts(), t)
}

func TestErrors(t *testing.T) {
	initTest(defaultTestOpts(), t)

	tests := []struct {
		SQL string
		Err string
	}{
		{
			SQL: "INVALID SQL",
			Err: "vtexplain execute error: unrecognized statement: INVALID SQL in INVALID SQL",
		},

		{
			SQL: "SELECT * FROM THIS IS NOT SQL",
			Err: "vtexplain execute error: syntax error at position 22 near 'is' in SELECT * FROM THIS IS NOT SQL",
		},

		{
			SQL: "SELECT * FROM table_not_in_vschema",
			Err: "vtexplain execute error: table table_not_in_vschema not found in SELECT * FROM table_not_in_vschema",
		},

		{
			SQL: "SELECT * FROM table_not_in_schema",
			Err: "vtexplain execute error: target: ks_unsharded.-.master, used tablet: explainCell-0 (ks_unsharded/-), table table_not_in_schema not found in schema in SELECT * FROM table_not_in_schema",
		},
	}

	for _, test := range tests {
		_, err := Run(test.SQL)
		if err == nil || err.Error() != test.Err {
			t.Errorf("Run(%s): %v, want %s", test.SQL, err, test.Err)
		}
	}
}
<commit_msg>re-enable the vtexplain tests<commit_after>/*
Copyright 2017 Google Inc.

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 vtexplain

import (
	"fmt"
	"io/ioutil"
	"os/exec"
	"strings"
	"testing"

	"github.com/youtube/vitess/go/testfiles"
)

var testOutputTempDir string

func defaultTestOpts() *Options {
	return &Options{
		ReplicationMode: "ROW",
		NumShards:       4,
		Normalize:       true,
	}
}

func initTest(opts *Options, t *testing.T) {
	schema, err := ioutil.ReadFile(testfiles.Locate("vtexplain/test-schema.sql"))
	if err != nil {
		t.Fatalf("error: %v", err)
	}

	vSchema, err := ioutil.ReadFile(testfiles.Locate("vtexplain/test-vschema.json"))
	if err != nil {
		t.Fatalf("error: %v", err)
	}

	err = Init(string(vSchema), string(schema), opts)
	if err != nil {
		t.Fatalf("vtexplain Init error: %v", err)
	}

}

func testExplain(testcase string, opts *Options, t *testing.T) {
	initTest(opts, t)

	sqlFile := testfiles.Locate(fmt.Sprintf("vtexplain/%s-queries.sql", testcase))
	sql, err := ioutil.ReadFile(sqlFile)

	jsonOutFile := testfiles.Locate(fmt.Sprintf("vtexplain/%s-output.json", testcase))
	jsonOut, err := ioutil.ReadFile(jsonOutFile)

	textOutFile := testfiles.Locate(fmt.Sprintf("vtexplain/%s-output.txt", testcase))
	textOut, err := ioutil.ReadFile(textOutFile)

	explains, err := Run(string(sql))
	if err != nil {
		t.Fatalf("vtexplain error: %v", err)
	}
	if explains == nil {
		t.Fatalf("vtexplain error running %s: no explain", string(sql))
	}

	explainJSON := ExplainsAsJSON(explains)
	if strings.TrimSpace(string(explainJSON)) != strings.TrimSpace(string(jsonOut)) {
		// Print the json that was actually returned and also dump to a
		// temp file to be able to diff the results.
		t.Errorf("json output did not match")
		t.Logf("got:\n%s\n", string(explainJSON))

		if testOutputTempDir == "" {
			testOutputTempDir, err = ioutil.TempDir("", "vtexplain_output")
			if err != nil {
				t.Fatalf("error getting tempdir: %v", err)
			}
		}
		gotFile := fmt.Sprintf("%s/%s-output.json", testOutputTempDir, testcase)
		ioutil.WriteFile(gotFile, []byte(explainJSON), 0644)

		command := exec.Command("diff", "-u", jsonOutFile, gotFile)
		out, _ := command.CombinedOutput()
		t.Logf("diff:\n%s\n", out)
	}

	explainText := ExplainsAsText(explains)
	if strings.TrimSpace(string(explainText)) != strings.TrimSpace(string(textOut)) {
		// Print the Text that was actually returned and also dump to a
		// temp file to be able to diff the results.
		t.Errorf("Text output did not match")
		t.Logf("got:\n%s\n", string(explainText))

		if testOutputTempDir == "" {
			testOutputTempDir, err = ioutil.TempDir("", "vtexplain_output")
			if err != nil {
				t.Fatalf("error getting tempdir: %v", err)
			}
		}
		gotFile := fmt.Sprintf("%s/%s-output.txt", testOutputTempDir, testcase)
		ioutil.WriteFile(gotFile, []byte(explainText), 0644)

		command := exec.Command("diff", "-u", textOutFile, gotFile)
		out, _ := command.CombinedOutput()
		t.Logf("diff:\n%s\n", out)
	}
}

func TestUnsharded(t *testing.T) {
	testExplain("unsharded", defaultTestOpts(), t)
}

func TestSelectSharded(t *testing.T) {
	testExplain("selectsharded", defaultTestOpts(), t)
}

func TestInsertSharded(t *testing.T) {
	testExplain("insertsharded", defaultTestOpts(), t)
}

func TestUpdateSharded(t *testing.T) {
	testExplain("updatesharded", defaultTestOpts(), t)
}

func TestDeleteSharded(t *testing.T) {
	testExplain("deletesharded", defaultTestOpts(), t)
}

func TestOptions(t *testing.T) {
	opts := &Options{
		ReplicationMode: "STATEMENT",
		NumShards:       4,
		Normalize:       false,
	}

	testExplain("options", opts, t)
}

func TestComments(t *testing.T) {
	testExplain("comments", defaultTestOpts(), t)
}

func TestErrors(t *testing.T) {
	initTest(defaultTestOpts(), t)

	tests := []struct {
		SQL string
		Err string
	}{
		{
			SQL: "INVALID SQL",
			Err: "vtexplain execute error: unrecognized statement: INVALID SQL in INVALID SQL",
		},

		{
			SQL: "SELECT * FROM THIS IS NOT SQL",
			Err: "vtexplain execute error: syntax error at position 22 near 'is' in SELECT * FROM THIS IS NOT SQL",
		},

		{
			SQL: "SELECT * FROM table_not_in_vschema",
			Err: "vtexplain execute error: table table_not_in_vschema not found in SELECT * FROM table_not_in_vschema",
		},

		{
			SQL: "SELECT * FROM table_not_in_schema",
			Err: "vtexplain execute error: target: ks_unsharded.-.master, used tablet: explainCell-0 (ks_unsharded/-), table table_not_in_schema not found in schema in SELECT * FROM table_not_in_schema",
		},
	}

	for _, test := range tests {
		_, err := Run(test.SQL)
		if err == nil || err.Error() != test.Err {
			t.Errorf("Run(%s): %v, want %s", test.SQL, err, test.Err)
		}
	}
}
<|endoftext|>


<commit_before>// ----------------------------------------------------------------------------
//
//     ***     AUTO GENERATED CODE    ***    AUTO GENERATED CODE     ***
//
// ----------------------------------------------------------------------------
//
//     This file is automatically generated by Magic Modules and manual
//     changes will be clobbered when the file is regenerated.
//
//     Please read more about how to change this file in
//     .github/CONTRIBUTING.md.
//
// ----------------------------------------------------------------------------

package google

import "reflect"

func GetStorageObjectAccessControlCaiObject(d TerraformResourceData, config *Config) ([]Asset, error) {
	name, err := assetName(d, config, "//storage.googleapis.com/b/{{bucket}}/o/{{object}}/acl/{{entity}}")
	if err != nil {
		return []Asset{}, err
	}
	if obj, err := GetStorageObjectAccessControlApiObject(d, config); err == nil {
		return []Asset{{
			Name: name,
			Type: "storage.googleapis.com/ObjectAccessControl",
			Resource: &AssetResource{
				Version:              "v1",
				DiscoveryDocumentURI: "https://www.googleapis.com/discovery/v1/apis/storage/v1/rest",
				DiscoveryName:        "ObjectAccessControl",
				Data:                 obj,
			},
		}}, nil
	} else {
		return []Asset{}, err
	}
}

func GetStorageObjectAccessControlApiObject(d TerraformResourceData, config *Config) (map[string]interface{}, error) {
	obj := make(map[string]interface{})
	bucketProp, err := expandStorageObjectAccessControlBucket(d.Get("bucket"), d, config)
	if err != nil {
		return nil, err
	} else if v, ok := d.GetOkExists("bucket"); !isEmptyValue(reflect.ValueOf(bucketProp)) && (ok || !reflect.DeepEqual(v, bucketProp)) {
		obj["bucket"] = bucketProp
	}
	entityProp, err := expandStorageObjectAccessControlEntity(d.Get("entity"), d, config)
	if err != nil {
		return nil, err
	} else if v, ok := d.GetOkExists("entity"); !isEmptyValue(reflect.ValueOf(entityProp)) && (ok || !reflect.DeepEqual(v, entityProp)) {
		obj["entity"] = entityProp
	}
	objectProp, err := expandStorageObjectAccessControlObject(d.Get("object"), d, config)
	if err != nil {
		return nil, err
	} else if v, ok := d.GetOkExists("object"); !isEmptyValue(reflect.ValueOf(objectProp)) && (ok || !reflect.DeepEqual(v, objectProp)) {
		obj["object"] = objectProp
	}
	roleProp, err := expandStorageObjectAccessControlRole(d.Get("role"), d, config)
	if err != nil {
		return nil, err
	} else if v, ok := d.GetOkExists("role"); !isEmptyValue(reflect.ValueOf(roleProp)) && (ok || !reflect.DeepEqual(v, roleProp)) {
		obj["role"] = roleProp
	}

	return obj, nil
}

func expandStorageObjectAccessControlBucket(v interface{}, d TerraformResourceData, config *Config) (interface{}, error) {
	return v, nil
}

func expandStorageObjectAccessControlEntity(v interface{}, d TerraformResourceData, config *Config) (interface{}, error) {
	return v, nil
}

func expandStorageObjectAccessControlObject(v interface{}, d TerraformResourceData, config *Config) (interface{}, error) {
	return v, nil
}

func expandStorageObjectAccessControlRole(v interface{}, d TerraformResourceData, config *Config) (interface{}, error) {
	return v, nil
}
<commit_msg>Properly encode paths for google_storage_object_access_control (#4946) (#738)<commit_after>// ----------------------------------------------------------------------------
//
//     ***     AUTO GENERATED CODE    ***    AUTO GENERATED CODE     ***
//
// ----------------------------------------------------------------------------
//
//     This file is automatically generated by Magic Modules and manual
//     changes will be clobbered when the file is regenerated.
//
//     Please read more about how to change this file in
//     .github/CONTRIBUTING.md.
//
// ----------------------------------------------------------------------------

package google

import "reflect"

func GetStorageObjectAccessControlCaiObject(d TerraformResourceData, config *Config) ([]Asset, error) {
	name, err := assetName(d, config, "//storage.googleapis.com/b/{{bucket}}/o/{{%object}}/acl/{{entity}}")
	if err != nil {
		return []Asset{}, err
	}
	if obj, err := GetStorageObjectAccessControlApiObject(d, config); err == nil {
		return []Asset{{
			Name: name,
			Type: "storage.googleapis.com/ObjectAccessControl",
			Resource: &AssetResource{
				Version:              "v1",
				DiscoveryDocumentURI: "https://www.googleapis.com/discovery/v1/apis/storage/v1/rest",
				DiscoveryName:        "ObjectAccessControl",
				Data:                 obj,
			},
		}}, nil
	} else {
		return []Asset{}, err
	}
}

func GetStorageObjectAccessControlApiObject(d TerraformResourceData, config *Config) (map[string]interface{}, error) {
	obj := make(map[string]interface{})
	bucketProp, err := expandStorageObjectAccessControlBucket(d.Get("bucket"), d, config)
	if err != nil {
		return nil, err
	} else if v, ok := d.GetOkExists("bucket"); !isEmptyValue(reflect.ValueOf(bucketProp)) && (ok || !reflect.DeepEqual(v, bucketProp)) {
		obj["bucket"] = bucketProp
	}
	entityProp, err := expandStorageObjectAccessControlEntity(d.Get("entity"), d, config)
	if err != nil {
		return nil, err
	} else if v, ok := d.GetOkExists("entity"); !isEmptyValue(reflect.ValueOf(entityProp)) && (ok || !reflect.DeepEqual(v, entityProp)) {
		obj["entity"] = entityProp
	}
	objectProp, err := expandStorageObjectAccessControlObject(d.Get("object"), d, config)
	if err != nil {
		return nil, err
	} else if v, ok := d.GetOkExists("object"); !isEmptyValue(reflect.ValueOf(objectProp)) && (ok || !reflect.DeepEqual(v, objectProp)) {
		obj["object"] = objectProp
	}
	roleProp, err := expandStorageObjectAccessControlRole(d.Get("role"), d, config)
	if err != nil {
		return nil, err
	} else if v, ok := d.GetOkExists("role"); !isEmptyValue(reflect.ValueOf(roleProp)) && (ok || !reflect.DeepEqual(v, roleProp)) {
		obj["role"] = roleProp
	}

	return obj, nil
}

func expandStorageObjectAccessControlBucket(v interface{}, d TerraformResourceData, config *Config) (interface{}, error) {
	return v, nil
}

func expandStorageObjectAccessControlEntity(v interface{}, d TerraformResourceData, config *Config) (interface{}, error) {
	return v, nil
}

func expandStorageObjectAccessControlObject(v interface{}, d TerraformResourceData, config *Config) (interface{}, error) {
	return v, nil
}

func expandStorageObjectAccessControlRole(v interface{}, d TerraformResourceData, config *Config) (interface{}, error) {
	return v, nil
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"

	"github.com/gin-gonic/gin"
	"github.com/opentracing/opentracing-go"
	"github.com/opentracing/opentracing-go/ext"
	"github.com/rs/zerolog/log"
)

// OpenTracingMiddleware creates a span for each request
func OpenTracingMiddleware() gin.HandlerFunc {
	return func(c *gin.Context) {

		path := c.Request.URL.Path
		if path == "/liveness" || path == "/readiness" {
			// don't log these requests, only execute them
			c.Next()
			return
		}

		// retrieve span context from upstream caller if available
		tracingCtx, err := opentracing.GlobalTracer().Extract(opentracing.HTTPHeaders, opentracing.HTTPHeadersCarrier(c.Request.Header))
		if err != nil {
			log.Warn().Err(err).Msgf("Failed extracting trace context from http headers for %v %v", c.Request.Method, c.Request.URL.Path)
		}

		// create a span for the http request
		span := opentracing.StartSpan(fmt.Sprintf("%v %v", c.Request.Method, c.Request.URL.Path), ext.RPCServerOption(tracingCtx))
		defer span.Finish()

		ext.SpanKindRPCServer.Set(span)
		ext.HTTPMethod.Set(span, c.Request.Method)
		ext.HTTPUrl.Set(span, c.Request.URL.String())

		// store the span in the request context
		c.Request = c.Request.WithContext(opentracing.ContextWithSpan(c.Request.Context(), span))

		// process request
		c.Next()
	}
}
<commit_msg>set http status after request<commit_after>package main

import (
	"fmt"

	"github.com/gin-gonic/gin"
	"github.com/opentracing/opentracing-go"
	"github.com/opentracing/opentracing-go/ext"
	"github.com/rs/zerolog/log"
)

// OpenTracingMiddleware creates a span for each request
func OpenTracingMiddleware() gin.HandlerFunc {
	return func(c *gin.Context) {

		path := c.Request.URL.Path
		if path == "/liveness" || path == "/readiness" {
			// don't log these requests, only execute them
			c.Next()
			return
		}

		// retrieve span context from upstream caller if available
		tracingCtx, err := opentracing.GlobalTracer().Extract(opentracing.HTTPHeaders, opentracing.HTTPHeadersCarrier(c.Request.Header))
		if err != nil {
			log.Warn().Err(err).Msgf("Failed extracting trace context from http headers for %v %v", c.Request.Method, c.Request.URL.Path)
		}

		// create a span for the http request
		span := opentracing.StartSpan(fmt.Sprintf("%v %v", c.Request.Method, c.Request.URL.Path), ext.RPCServerOption(tracingCtx))
		defer span.Finish()

		ext.SpanKindRPCServer.Set(span)
		ext.HTTPMethod.Set(span, c.Request.Method)
		ext.HTTPUrl.Set(span, c.Request.URL.String())

		// store the span in the request context
		c.Request = c.Request.WithContext(opentracing.ContextWithSpan(c.Request.Context(), span))

		// process request
		c.Next()

		ext.HTTPStatusCode.Set(span, uint16(c.Writer.Status()))
	}
}
<|endoftext|>


<commit_before>package main

import (
	"context"
	"flag"
	"fmt"
	"io"
	"log"
	"time"

	flatbuffers "github.com/google/flatbuffers/go"
	models "github.com/google/flatbuffers/grpc/examples/go/greeter/models"
	"google.golang.org/grpc"
)

var (
	addr = "3000"
	name = flag.String("name", "Flatbuffers", "name to be sent to server :D")
)

func printSayHello(client models.GreeterClient, name string) {
	log.Printf("Name to be sent (%s)", name)
	b := flatbuffers.NewBuilder(0)
	i := b.CreateString(name)
	models.HelloRequestStart(b)
	models.HelloRequestAddName(b, i)
	b.Finish(models.HelloRequestEnd(b))

	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	request, err := client.SayHello(ctx, b, grpc.CallContentSubtype("flatbuffers"))
	if err != nil {
		log.Fatalf("%v.SayHello(_) = _, %v: ", client, err)
	}
	log.Printf("server said %q", request.Message())
}

func printSayManyHello(client models.GreeterClient, name string) {
	log.Printf("Name to be sent (%s)", name)
	b := flatbuffers.NewBuilder(0)
	i := b.CreateString(name)
	models.HelloRequestStart(b)
	models.HelloRequestAddName(b, i)
	b.Finish(models.HelloRequestEnd(b))

	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	stream, err := client.SayManyHellos(ctx, b, grpc.CallContentSubtype("flatbuffers"))
	if err != nil {
		log.Fatalf("%v.SayManyHellos(_) = _, %v", client, err)
	}
	for {
		request, err := stream.Recv()
		if err == io.EOF {
			break
		}
		if err != nil {
			log.Fatalf("%v.SayManyHellos(_) = _, %v", client, err)
		}
		log.Printf("server said %q", request.Message())
	}
}

func main() {
	flag.Parse()
	conn, err := grpc.Dial(fmt.Sprintf("localhost:%d", 3000), grpc.WithInsecure(), grpc.WithCodec(flatbuffers.FlatbuffersCodec{}))
	if err != nil {
		log.Fatalf("fail to dial: %v", err)
	}
	defer conn.Close()
	client := models.NewGreeterClient(conn)
	printSayHello(client, *name)
	printSayManyHello(client, *name)
}
<commit_msg>fix(#7360): grpc used deprecated functions (#7361)<commit_after>package main

import (
	"context"
	"flag"
	"fmt"
	"io"
	"log"
	"time"

	flatbuffers "github.com/google/flatbuffers/go"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials/insecure"

	models "github.com/google/flatbuffers/grpc/examples/go/greeter/models"
)

var (
	addr = "3000"
	name = flag.String("name", "Flatbuffers", "name to be sent to server :D")
)

func printSayHello(client models.GreeterClient, name string) {
	log.Printf("Name to be sent (%s)", name)
	b := flatbuffers.NewBuilder(0)
	i := b.CreateString(name)
	models.HelloRequestStart(b)
	models.HelloRequestAddName(b, i)
	b.Finish(models.HelloRequestEnd(b))

	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	request, err := client.SayHello(ctx, b, grpc.CallContentSubtype("flatbuffers"))
	if err != nil {
		log.Fatalf("%v.SayHello(_) = _, %v: ", client, err)
	}
	log.Printf("server said %q", request.Message())
}

func printSayManyHello(client models.GreeterClient, name string) {
	log.Printf("Name to be sent (%s)", name)
	b := flatbuffers.NewBuilder(0)
	i := b.CreateString(name)
	models.HelloRequestStart(b)
	models.HelloRequestAddName(b, i)
	b.Finish(models.HelloRequestEnd(b))

	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	stream, err := client.SayManyHellos(ctx, b, grpc.CallContentSubtype("flatbuffers"))
	if err != nil {
		log.Fatalf("%v.SayManyHellos(_) = _, %v", client, err)
	}
	for {
		request, err := stream.Recv()
		if err == io.EOF {
			break
		}
		if err != nil {
			log.Fatalf("%v.SayManyHellos(_) = _, %v", client, err)
		}
		log.Printf("server said %q", request.Message())
	}
}

func main() {
	flag.Parse()
	conn, err := grpc.Dial(fmt.Sprintf("localhost:%d", 3000),
		grpc.WithTransportCredentials(insecure.NewCredentials()),
		grpc.WithDefaultCallOptions(grpc.ForceCodec(flatbuffers.FlatbuffersCodec{})))
	if err != nil {
		log.Fatalf("fail to dial: %v", err)
	}
	defer conn.Close()
	client := models.NewGreeterClient(conn)
	printSayHello(client, *name)
	printSayManyHello(client, *name)
}
<|endoftext|>


<commit_before>package tls

import (
	"crypto"

	"github.com/xenolf/lego/acme"
)

type acmeClient interface {
	SetChallengeProvider(acme.Challenge, acme.ChallengeProvider) error
	Register() (*acme.RegistrationResource, error)
	AgreeToTOS() error
	ObtainCertificate([]string, bool, crypto.PrivateKey, bool) (acme.CertificateResource, map[string]error)
	RenewCertificate(cert acme.CertificateResource, bundle, mustStaple bool) (acme.CertificateResource, error)
}

type legoAcme struct {
	*acme.Client
}

var newAcmeClient = newLegoAcmeClient

func newLegoAcmeClient(user acme.User, cp acme.ChallengeProvider) (acmeClient, error) {
	c, err := acme.NewClient(defaultCAURL, user, "")
	if err != nil {
		return nil, err
	}
	err = c.SetChallengeProvider(acme.HTTP01, cp)
	if err != nil {
		return nil, err
	}
	return legoAcme{c}, nil
}
<commit_msg>tls: Only support http-01 challenge.<commit_after>package tls

import (
	"crypto"

	"github.com/xenolf/lego/acme"
)

type acmeClient interface {
	SetChallengeProvider(acme.Challenge, acme.ChallengeProvider) error
	Register() (*acme.RegistrationResource, error)
	AgreeToTOS() error
	ObtainCertificate([]string, bool, crypto.PrivateKey, bool) (acme.CertificateResource, map[string]error)
	RenewCertificate(cert acme.CertificateResource, bundle, mustStaple bool) (acme.CertificateResource, error)
}

type legoAcme struct {
	*acme.Client
}

var newAcmeClient = newLegoAcmeClient

func newLegoAcmeClient(user acme.User, cp acme.ChallengeProvider) (acmeClient, error) {
	c, err := acme.NewClient(defaultCAURL, user, "")
	if err != nil {
		return nil, err
	}
	err = c.SetChallengeProvider(acme.HTTP01, cp)
	if err != nil {
		return nil, err
	}
	// only support http-01
	c.ExcludeChallenges([]acme.Challenge{acme.TLSSNI01, acme.DNS01})
	return legoAcme{c}, nil
}
<|endoftext|>


<commit_before>//+build windows

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

package oswrap

import (
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"syscall"
	"unsafe"

	"golang.org/x/sys/windows"
)

var (
	kernel32       = windows.NewLazySystemDLL("Kernel32.dll")
	procMoveFileEx = kernel32.NewProc("MoveFileExW")
)

type (
	DWORD   uint32
	LPCTSTR *uint16
)

const MOVEFILE_DELAY_UNTIL_REBOOT = 4

// normPath transforms a windows path into an extended-length path as described in
// https://msdn.microsoft.com/en-us/library/windows/desktop/aa365247(v=vs.85).aspx#maxpath
func normPath(path string) (string, error) {
	if strings.HasPrefix(path, "\\\\?\\") {
		return path, nil
	}

	path, err := filepath.Abs(path)
	if err != nil {
		return "", err
	}
	path = filepath.Clean(path)
	return "\\\\?\\" + path, nil
}

// RemoveOnReboot schedules a file for removal on next reboot.
func RemoveOnReboot(name string) error {
	name, err := normPath(name)
	if err != nil {
		return err
	}

	nPtr, err := syscall.UTF16PtrFromString(name)
	if err != nil {
		return fmt.Errorf("error encoding path to UTF16: %v", err)
	}
	ret, _, _ := procMoveFileEx.Call(
		uintptr(unsafe.Pointer(nPtr)),
		uintptr(0),
		MOVEFILE_DELAY_UNTIL_REBOOT)
	if ret == 0 {
		return errors.New("return code of '0' (failure) from MoveFileEx")
	}
	return nil
}

// Open calls os.Open with name normalized
func Open(name string) (*os.File, error) {
	name, err := normPath(name)
	if err != nil {
		return nil, err
	}
	return os.Open(name)
}

// Create calls os.Create with name normalized
func Create(name string) (*os.File, error) {
	name, err := normPath(name)
	if err != nil {
		return nil, err
	}
	return os.Create(name)
}

// OpenFile calls os.OpenFile with name normalized
func OpenFile(name string, flag int, perm os.FileMode) (*os.File, error) {
	name, err := normPath(name)
	if err != nil {
		return nil, err
	}
	return os.OpenFile(name, flag, perm)
}

// Remove calls os.Remove with name normalized
func Remove(name string) error {
	name, err := normPath(name)
	if err != nil {
		return nil
	}
	return os.Remove(name)
}

// RemoveAll calls os.RemoveAll with name normalized
func RemoveAll(name string) error {
	name, err := normPath(name)
	if err != nil {
		return nil
	}
	return os.RemoveAll(name)
}

// Mkdir calls os.Mkdir with name normalized
func Mkdir(name string, mode os.FileMode) error {
	name, err := normPath(name)
	if err != nil {
		return err
	}
	return os.Mkdir(name, mode)
}

// MkdirAll calls os.MkdirAll with name normalized
func MkdirAll(name string, mode os.FileMode) error {
	name, err := normPath(name)
	if err != nil {
		return err
	}

	// os.MkdirAll does not work with extended-length paths if
	// nothing in the path exists.
	if err := mkRootDir(name, mode); err != nil {
		return err
	}

	return os.MkdirAll(name, mode)
}

// Rename calls os.Rename with name normalized
func Rename(oldpath, newpath string) error {
	oldpath, err := normPath(oldpath)
	if err != nil {
		return err
	}
	newpath, err = normPath(newpath)
	if err != nil {
		return err
	}
	return os.Rename(oldpath, newpath)
}

// Lstat calls os.Lstat with name normalized
func Lstat(name string) (os.FileInfo, error) {
	name, err := normPath(name)
	if err != nil {
		return nil, err
	}
	return os.Lstat(name)
}

// Stat calls os.Stat with name normalized
func Stat(name string) (os.FileInfo, error) {
	name, err := normPath(name)
	if err != nil {
		return nil, err
	}
	return os.Stat(name)
}

// Walk calls filepath.Walk with name normalized, and un-normalizes name before
// calling walkFn
func Walk(root string, walkFn filepath.WalkFunc) error {
	newroot, err := normPath(root)
	if err != nil {
		return err
	}
	return filepath.Walk(newroot, func(path string, info os.FileInfo, err error) error {
		oldpath := root + strings.TrimPrefix(path, newroot)
		return walkFn(oldpath, info, err)
	})
}
<commit_msg>Test remove operations on Windows for paths that look sensitive. (#53)<commit_after>//+build windows

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

package oswrap

import (
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"syscall"
	"unsafe"

	"golang.org/x/sys/windows"
)

var (
	kernel32       = windows.NewLazySystemDLL("Kernel32.dll")
	procMoveFileEx = kernel32.NewProc("MoveFileExW")
)

type (
	DWORD   uint32
	LPCTSTR *uint16
)

const MOVEFILE_DELAY_UNTIL_REBOOT = 4

// normPath transforms a windows path into an extended-length path as described in
// https://msdn.microsoft.com/en-us/library/windows/desktop/aa365247(v=vs.85).aspx#maxpath
func normPath(path string) (string, error) {
	if strings.HasPrefix(path, "\\\\?\\") {
		return path, nil
	}

	path, err := filepath.Abs(path)
	if err != nil {
		return "", err
	}
	path = filepath.Clean(path)
	return "\\\\?\\" + path, nil
}

// RemoveOnReboot schedules a file for removal on next reboot.
func RemoveOnReboot(name string) error {
	name, err := normPath(name)
	if err != nil {
		return err
	}

	nPtr, err := syscall.UTF16PtrFromString(name)
	if err != nil {
		return fmt.Errorf("error encoding path to UTF16: %v", err)
	}
	ret, _, _ := procMoveFileEx.Call(
		uintptr(unsafe.Pointer(nPtr)),
		uintptr(0),
		MOVEFILE_DELAY_UNTIL_REBOOT)
	if ret == 0 {
		return errors.New("return code of '0' (failure) from MoveFileEx")
	}
	return nil
}

// Open calls os.Open with name normalized
func Open(name string) (*os.File, error) {
	name, err := normPath(name)
	if err != nil {
		return nil, err
	}
	return os.Open(name)
}

// Create calls os.Create with name normalized
func Create(name string) (*os.File, error) {
	name, err := normPath(name)
	if err != nil {
		return nil, err
	}
	return os.Create(name)
}

// OpenFile calls os.OpenFile with name normalized
func OpenFile(name string, flag int, perm os.FileMode) (*os.File, error) {
	name, err := normPath(name)
	if err != nil {
		return nil, err
	}
	return os.OpenFile(name, flag, perm)
}

func isProtectedPath(path string) (bool, error) {
	if path == "" {
		return false, errors.New("cannot use empty path")
	}
	paths := []string{}
	for _, v := range []string{"SystemDrive", "SystemRoot", "ProgramFiles", "ProgramFiles(x86)"} {
		paths = append(paths, os.Getenv(v))
	}
	paths = append(paths, filepath.Join(os.Getenv("SystemRoot"), "system32"))

	for _, p := range paths {
		if strings.ToLower(filepath.Clean(p)) == strings.ToLower(filepath.Clean(path)) {
			return true, nil
		}
	}
	return false, nil
}

// Remove calls os.Remove with name normalized
func Remove(name string) error {
	if bad, err := isProtectedPath(name); bad || err != nil {
		return fmt.Errorf("path %s appears to be protected", name)
	}
	name, err := normPath(name)
	if err != nil {
		return nil
	}
	return os.Remove(name)
}

// RemoveAll calls os.RemoveAll with name normalized
func RemoveAll(name string) error {
	if bad, err := isProtectedPath(name); bad || err != nil {
		return fmt.Errorf("path %s appears to be protected", name)
	}
	name, err := normPath(name)
	if err != nil {
		return nil
	}
	return os.RemoveAll(name)
}

// Mkdir calls os.Mkdir with name normalized
func Mkdir(name string, mode os.FileMode) error {
	name, err := normPath(name)
	if err != nil {
		return err
	}
	return os.Mkdir(name, mode)
}

// MkdirAll calls os.MkdirAll with name normalized
func MkdirAll(name string, mode os.FileMode) error {
	name, err := normPath(name)
	if err != nil {
		return err
	}

	// os.MkdirAll does not work with extended-length paths if
	// nothing in the path exists.
	if err := mkRootDir(name, mode); err != nil {
		return err
	}

	return os.MkdirAll(name, mode)
}

// Rename calls os.Rename with name normalized
func Rename(oldpath, newpath string) error {
	oldpath, err := normPath(oldpath)
	if err != nil {
		return err
	}
	newpath, err = normPath(newpath)
	if err != nil {
		return err
	}
	return os.Rename(oldpath, newpath)
}

// Lstat calls os.Lstat with name normalized
func Lstat(name string) (os.FileInfo, error) {
	name, err := normPath(name)
	if err != nil {
		return nil, err
	}
	return os.Lstat(name)
}

// Stat calls os.Stat with name normalized
func Stat(name string) (os.FileInfo, error) {
	name, err := normPath(name)
	if err != nil {
		return nil, err
	}
	return os.Stat(name)
}

// Walk calls filepath.Walk with name normalized, and un-normalizes name before
// calling walkFn
func Walk(root string, walkFn filepath.WalkFunc) error {
	newroot, err := normPath(root)
	if err != nil {
		return err
	}
	return filepath.Walk(newroot, func(path string, info os.FileInfo, err error) error {
		oldpath := root + strings.TrimPrefix(path, newroot)
		return walkFn(oldpath, info, err)
	})
}
<|endoftext|>


<commit_before>package main

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

	"database/sql"

	postgres "github.com/lib/pq"

	"github.com/Shopify/sarama"

	log "github.com/Sirupsen/logrus"
	cli "gopkg.in/urfave/cli.v2"
)

var consumerTblName = postgres.QuoteIdentifier("kafka_consumer")

type WAL struct {
	Type       string          `json:"type"`
	InstanceId string          `json:"instanceId"`
	Table      string          `json:"table"`
	Host       string          `json:"host"`
	Key        string          `json:"key"`
	CreatedAt  time.Time       `json:"created_at"`
	Data       json.RawMessage `json:"data"`
}

func main() {
	app := &cli.App{
		Name:    "kafka2psql",
		Usage:   `Store Kafka Topic To PostgreSQL Table`,
		Version: "0.1",
		Flags: []cli.Flag{
			&cli.StringSliceFlag{
				Name:  "brokers, b",
				Value: cli.NewStringSlice("localhost:9092"),
				Usage: "kafka brokers address",
			},
			&cli.StringFlag{
				Name:  "table-topic",
				Value: "WAL",
				Usage: "topic name that contains the table",
			},
			&cli.StringFlag{
				Name:  "table",
				Value: "user_updates",
				Usage: "table name in WAL to archive",
			},
			&cli.StringFlag{
				Name:  "pq",
				Value: "postgres://127.0.0.1:5432/pipeline?sslmode=disable",
				Usage: "psql url",
			},
			&cli.StringFlag{
				Name:  "pq-tblname",
				Value: "log20060102",
				Usage: "psql table name, aware of timeformat in golang",
			},
			&cli.DurationFlag{
				Name:  "commit-interval",
				Value: time.Second,
				Usage: "interval for committing pending data to psql",
			},
		},
		Action: processor,
	}
	app.Run(os.Args)
}

func processor(c *cli.Context) error {
	brokers := c.StringSlice("brokers")
	table_topic := c.String("table-topic")
	table := c.String("table")
	pq := c.String("pq")
	pq_tblname := c.String("pq-tblname")
	commit_interval := c.Duration("commit-interval")

	log.Println("brokers:", brokers)
	log.Println("table-topic:", table_topic)
	log.Println("table:", table)
	log.Println("pq:", pq)
	log.Println("pq-tblname:", pq_tblname)
	log.Println("commit-interval:", commit_interval)

	// unique consumer name to store in psql
	consumerId := fmt.Sprintf("%v-%v-%v", table_topic, table, pq_tblname)
	log.Println("consumerId:", consumerId)

	// connect to postgres
	db, err := sql.Open("postgres", pq)
	if err != nil {
		log.Fatal(err)
	}
	defer db.Close()

	consumer, err := sarama.NewConsumer(brokers, nil)
	if err != nil {
		log.Fatalln(err)
	}

	defer func() {
		if err := consumer.Close(); err != nil {
			log.Fatalln(err)
		}
	}()

	// table creation
	db.Exec(fmt.Sprintf("CREATE TABLE %s (id TEXT PRIMARY KEY, value BIGINT)", consumerTblName))
	lastTblName := postgres.QuoteIdentifier(time.Now().Format(pq_tblname))
	db.Exec("CREATE TABLE " + lastTblName + "(id TEXT PRIMARY KEY, data JSONB)")

	// read offset
	offset := sarama.OffsetOldest
	err = db.QueryRow("SELECT value FROM kafka_consumer WHERE id = $1 LIMIT 1", consumerId).Scan(&offset)
	if err != nil {
		log.Println(err)
	}
	log.Println("consuming from offset:", offset)

	tableConsumer, err := consumer.ConsumePartition(table_topic, 0, offset)
	if err != nil {
		log.Fatalln(err)
	}

	defer func() {
		if err := tableConsumer.Close(); err != nil {
			log.Fatalln(err)
		}
	}()

	log.Println("started")
	commitTicker := time.NewTicker(commit_interval)
	pending := make(map[string][]byte)

	for {
		select {
		case msg := <-tableConsumer.Messages():
			wal := &WAL{}
			if err := json.Unmarshal(msg.Value, wal); err == nil {
				if wal.Table == table { // table filter
					// create new table if necessary
					tblName := postgres.QuoteIdentifier(time.Now().Format(pq_tblname))
					if tblName != lastTblName {
						commit(lastTblName, consumerId, db, pending, offset)
						pending = make(map[string][]byte)
						// CREATE TABLE
						db.Exec("CREATE TABLE " + tblName + "(id TEXT PRIMARY KEY, data JSONB)")
						lastTblName = tblName
					}

					// pending
					pending[wal.Key] = msg.Value
					offset = msg.Offset
				}
			} else {
				log.Println(err)
			}
		case <-commitTicker.C:
			commit(lastTblName, consumerId, db, pending, offset)
			pending = make(map[string][]byte)
		}
	}
}

func commit(tblname, consumerId string, db *sql.DB, pending map[string][]byte, offset int64) {
	if len(pending) == 0 {
		return
	}

	for key, value := range pending {
		if r, err := db.Exec(fmt.Sprintf("INSERT INTO %s (id, data) VALUES ($1,$2) ON CONFLICT(id) DO UPDATE SET data = EXCLUDED.data",
			tblname), key, string(value)); err == nil {
		} else {
			log.Println(r, err)
		}
	}

	// write offset
	if r, err := db.Exec(fmt.Sprintf("INSERT INTO %s (id, value) VALUES ($1,$2) ON CONFLICT(id) DO UPDATE SET value=EXCLUDED.value",
		consumerTblName), consumerId, offset); err != nil {
		log.Println(r, err)
	}
	log.Println("written:", len(pending), "offset:", offset)
}
<commit_msg>fix the mis-guiding pq-tblname<commit_after>package main

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

	"database/sql"

	postgres "github.com/lib/pq"

	"github.com/Shopify/sarama"

	log "github.com/Sirupsen/logrus"
	cli "gopkg.in/urfave/cli.v2"
)

var consumerTblName = postgres.QuoteIdentifier("kafka_consumer")

type WAL struct {
	Type       string          `json:"type"`
	InstanceId string          `json:"instanceId"`
	Table      string          `json:"table"`
	Host       string          `json:"host"`
	Key        string          `json:"key"`
	CreatedAt  time.Time       `json:"created_at"`
	Data       json.RawMessage `json:"data"`
}

func main() {
	app := &cli.App{
		Name:    "kafka2psql",
		Usage:   `Store Kafka Topic To PostgreSQL Table`,
		Version: "0.1",
		Flags: []cli.Flag{
			&cli.StringSliceFlag{
				Name:  "brokers, b",
				Value: cli.NewStringSlice("localhost:9092"),
				Usage: "kafka brokers address",
			},
			&cli.StringFlag{
				Name:  "table-topic",
				Value: "WAL",
				Usage: "topic name that contains the table",
			},
			&cli.StringFlag{
				Name:  "table",
				Value: "user_updates",
				Usage: "table name in WAL to archive",
			},
			&cli.StringFlag{
				Name:  "pq",
				Value: "postgres://127.0.0.1:5432/pipeline?sslmode=disable",
				Usage: "psql url",
			},
			&cli.StringFlag{
				Name:  "pq-tblname",
				Value: "users",
				Usage: "psql table name, aware of timeformat in golang",
			},
			&cli.DurationFlag{
				Name:  "commit-interval",
				Value: time.Second,
				Usage: "interval for committing pending data to psql",
			},
		},
		Action: processor,
	}
	app.Run(os.Args)
}

func processor(c *cli.Context) error {
	brokers := c.StringSlice("brokers")
	table_topic := c.String("table-topic")
	table := c.String("table")
	pq := c.String("pq")
	pq_tblname := c.String("pq-tblname")
	commit_interval := c.Duration("commit-interval")

	log.Println("brokers:", brokers)
	log.Println("table-topic:", table_topic)
	log.Println("table:", table)
	log.Println("pq:", pq)
	log.Println("pq-tblname:", pq_tblname)
	log.Println("commit-interval:", commit_interval)

	// unique consumer name to store in psql
	consumerId := fmt.Sprintf("%v-%v-%v", table_topic, table, pq_tblname)
	log.Println("consumerId:", consumerId)

	// connect to postgres
	db, err := sql.Open("postgres", pq)
	if err != nil {
		log.Fatal(err)
	}
	defer db.Close()

	consumer, err := sarama.NewConsumer(brokers, nil)
	if err != nil {
		log.Fatalln(err)
	}

	defer func() {
		if err := consumer.Close(); err != nil {
			log.Fatalln(err)
		}
	}()

	// table creation
	db.Exec(fmt.Sprintf("CREATE TABLE %s (id TEXT PRIMARY KEY, value BIGINT)", consumerTblName))
	lastTblName := postgres.QuoteIdentifier(time.Now().Format(pq_tblname))
	db.Exec("CREATE TABLE " + lastTblName + "(id TEXT PRIMARY KEY, data JSONB)")

	// read offset
	offset := sarama.OffsetOldest
	err = db.QueryRow("SELECT value FROM kafka_consumer WHERE id = $1 LIMIT 1", consumerId).Scan(&offset)
	if err != nil {
		log.Println(err)
	}
	log.Println("consuming from offset:", offset)

	tableConsumer, err := consumer.ConsumePartition(table_topic, 0, offset)
	if err != nil {
		log.Fatalln(err)
	}

	defer func() {
		if err := tableConsumer.Close(); err != nil {
			log.Fatalln(err)
		}
	}()

	log.Println("started")
	commitTicker := time.NewTicker(commit_interval)
	pending := make(map[string][]byte)

	for {
		select {
		case msg := <-tableConsumer.Messages():
			wal := &WAL{}
			if err := json.Unmarshal(msg.Value, wal); err == nil {
				if wal.Table == table { // table filter
					// create new table if necessary
					tblName := postgres.QuoteIdentifier(time.Now().Format(pq_tblname))
					if tblName != lastTblName {
						commit(lastTblName, consumerId, db, pending, offset)
						pending = make(map[string][]byte)
						// CREATE TABLE
						db.Exec("CREATE TABLE " + tblName + "(id TEXT PRIMARY KEY, data JSONB)")
						lastTblName = tblName
					}

					// pending
					pending[wal.Key] = msg.Value
					offset = msg.Offset
				}
			} else {
				log.Println(err)
			}
		case <-commitTicker.C:
			commit(lastTblName, consumerId, db, pending, offset)
			pending = make(map[string][]byte)
		}
	}
}

func commit(tblname, consumerId string, db *sql.DB, pending map[string][]byte, offset int64) {
	if len(pending) == 0 {
		return
	}

	for key, value := range pending {
		if r, err := db.Exec(fmt.Sprintf("INSERT INTO %s (id, data) VALUES ($1,$2) ON CONFLICT(id) DO UPDATE SET data = EXCLUDED.data",
			tblname), key, string(value)); err == nil {
		} else {
			log.Println(r, err)
		}
	}

	// write offset
	if r, err := db.Exec(fmt.Sprintf("INSERT INTO %s (id, value) VALUES ($1,$2) ON CONFLICT(id) DO UPDATE SET value=EXCLUDED.value",
		consumerTblName), consumerId, offset); err != nil {
		log.Println(r, err)
	}
	log.Println("written:", len(pending), "offset:", offset)
}
<|endoftext|>


<commit_before>package service

import "fmt"

//Builder is a simple implementation of ContainerBuilder
type Builder struct {
	definitions []Definition
}

//Insert a new definition into the Builder
func (b *Builder) Insert(def Definition) {
	b.definitions = append(b.definitions, def)
}

//Build builds the container once all definitions have been place in it
//dependencies need to be inserted in order for now
func (b *Builder) Build() (Container, error) {
	numDefs := len(b.definitions)
	servs := make(map[string]interface{}, numDefs)
	for _, def := range b.definitions {
		numDeps := len(def.Dependencies)
		if numDeps == 0 {
			service, err := def.Initializer.Init(map[string]interface{}{}, def.Configuration)
			if err != nil {
				return nil, err
			}
			servs[def.Name] = service
		} else {
			deps := make(map[string]interface{}, numDeps)
			for _, name := range def.Dependencies {
				dep, ok := servs[name]
				if !ok {
					return nil, fmt.Errorf("service: Could not find dependency %q for service %q. Please make sure to insert them in order", name, def.Name)
				}
				deps[name] = dep
			}
			service, err := def.Initializer.Init(deps, def.Configuration)
			if err != nil {
				return nil, err
			}
			servs[def.Name] = service
		}
	}

	return SimpleContainer{
		services: servs,
	}, nil
}
<commit_msg>builder: Simplify Build method<commit_after>package service

import "fmt"

//Builder is a simple implementation of ContainerBuilder
type Builder struct {
	definitions []Definition
}

//Insert a new definition into the Builder
func (b *Builder) Insert(def Definition) {
	b.definitions = append(b.definitions, def)
}

//Build builds the container once all definitions have been place in it
//dependencies need to be inserted in order for now
func (b *Builder) Build() (Container, error) {
	numDefs := len(b.definitions)
	servs := make(map[string]interface{}, numDefs)
	for _, def := range b.definitions {
		numDeps := len(def.Dependencies)
		deps := make(map[string]interface{}, numDeps)
		for _, name := range def.Dependencies {
			dep, ok := servs[name]
			if !ok {
				return nil, fmt.Errorf("service: Could not find dependency %q for service %q. Please make sure to insert them in order", name, def.Name)
			}
			deps[name] = dep
		}
		service, err := def.Initializer.Init(deps, def.Configuration)
		if err != nil {
			return nil, err
		}
		servs[def.Name] = service
	}

	return SimpleContainer{
		services: servs,
	}, nil
}
<|endoftext|>