<commit_before>package stepman

import (
	"archive/zip"
	"encoding/json"
	"errors"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"path"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/bitrise-io/go-pathutil/pathutil"
	"github.com/bitrise-io/stepman/models"
	"gopkg.in/yaml.v2"
)

// DebugMode ...
var DebugMode bool

func parseStepYml(pth, id, version string) (models.StepModel, error) {
	bytes, err := ioutil.ReadFile(pth)
	if err != nil {
		return models.StepModel{}, err
	}

	var stepJSON models.StepModel
	if err := yaml.Unmarshal(bytes, &stepJSON); err != nil {
		return models.StepModel{}, err
	}

	stepJSON.ID = id
	stepJSON.VersionTag = version
	stepJSON.StepLibSource = CollectionURI

	return stepJSON, nil
}

// ParseStepCollection ...
func ParseStepCollection(pth string) (models.StepCollectionModel, error) {
	bytes, err := ioutil.ReadFile(pth)
	if err != nil {
		return models.StepCollectionModel{}, err
	}

	var stepCollection models.StepCollectionModel
	if err := yaml.Unmarshal(bytes, &stepCollection); err != nil {
		return models.StepCollectionModel{}, err
	}
	return stepCollection, nil
}

// DownloadStep ...
func DownloadStep(collection models.StepCollectionModel, step models.StepModel) error {
	downloadLocations := collection.GetDownloadLocations(step)

	success := false
	for _, downloadLocationMap := range downloadLocations {
		for key, value := range downloadLocationMap {
			switch key {
			case "zip":
				log.Info("[STEPMAN] - Downloading step from:", value)
				if err := DownloadAndUnZIP(value, GetStepPath(step)); err != nil {
					log.Error("[STEPMAN] - Failed to download step.zip:", err)
				} else {
					success = true
					return nil
				}
			case "git":
				log.Info("[STEPMAN] - Git clone step from:", value)
				if err := DoGitClone(value, GetStepPath(step)); err != nil {
					log.Error("[STEPMAN] - Failed to clone step:", err)
				} else {
					success = true
					return nil
				}
			default:
				log.Error("[STEPMAN] - Invalid download location")
				break
			}
		}
	}

	if !success {
		return errors.New("Failed to download step")
	}
	return nil
}

// GetStepPath ...
// cach is cool
func GetStepPath(step models.StepModel) string {
	return GetCurrentStepCacheDir() + step.ID + "/" + step.VersionTag + "/"
}

// semantic version (X.Y.Z)
// true if version 2 is greater then version 1
func isVersionGrater(version1, version2 string) bool {
	version1Slice := strings.Split(version1, ".")
	version2Slice := strings.Split(version2, ".")

	for i, num := range version1Slice {
		num1, err1 := strconv.ParseInt(num, 0, 64)
		if err1 != nil {
			log.Error("[STEPMAN] - Failed to parse int:", err1)
			return false
		}

		num2, err2 := strconv.ParseInt(version2Slice[i], 0, 64)
		if err2 != nil {
			log.Error("[STEPMAN] - Failed to parse int:", err2)
			return false
		}

		if num2 > num1 {
			return true
		}
	}
	return false
}

func addStepToStepGroup(step models.StepModel, stepGroup models.StepGroupModel) models.StepGroupModel {
	var newStepGroup models.StepGroupModel
	if len(stepGroup.Versions) > 0 {
		// Step Group already created -> new version of step
		newStepGroup = stepGroup

		if isVersionGrater(newStepGroup.Latest.VersionTag, step.VersionTag) {
			newStepGroup.Latest = step
		}
	} else {
		// Create Step Group
		newStepGroup = models.StepGroupModel{}
		newStepGroup.Latest = step
	}

	versions := make([]models.StepModel, len(newStepGroup.Versions))
	for idx, step := range newStepGroup.Versions {
		versions[idx] = step
	}
	versions = append(versions, step)

	newStepGroup.Versions = versions
	newStepGroup.ID = step.ID
	return newStepGroup
}

func generateFormattedJSONForStepsSpec(templateCollection models.StepCollectionModel) ([]byte, error) {
	collection := models.StepCollectionModel{
		FormatVersion:        templateCollection.FormatVersion,
		GeneratedAtTimeStamp: time.Now().Unix(),
		SteplibSource:        CollectionURI,
		DownloadLocations:    templateCollection.DownloadLocations,
	}

	stepHash := models.StepHash{}

	stepsSpecDir := GetCurrentStepCollectionPath()
	err := filepath.Walk(stepsSpecDir, func(path string, f os.FileInfo, err error) error {
		truncatedPath := strings.Replace(path, stepsSpecDir, "", -1)
		match, matchErr := regexp.MatchString("([a-z]+).yml", truncatedPath)
		if matchErr != nil {
			return matchErr
		}

		if match {
			components := strings.Split(truncatedPath, "/")
			if len(components) == 4 {
				name := components[1]
				version := components[2]

				step, parseErr := parseStepYml(path, name, version)
				if parseErr != nil {
					return parseErr
				}
				stepGroup := addStepToStepGroup(step, stepHash[name])

				stepHash[name] = stepGroup
			} else {
				log.Debug("[STEPMAN] - Path:", truncatedPath)
				log.Debug("[STEPMAN] - Legth:", len(components))
			}
		}

		return err
	})
	if err != nil {
		log.Error("[STEPMAN] - Failed to walk through path:", err)
	}

	collection.Steps = stepHash

	var bytes []byte
	// if DebugMode == true {
	bytes, err = json.MarshalIndent(collection, "", "\t")
	// } else {
	// 	bytes, err = json.Marshal(collection)
	// }
	if err != nil {
		log.Error("[STEPMAN] - Failed to parse json:", err)
		return []byte{}, err
	}

	return bytes, nil
}

// WriteStepSpecToFile ...
func WriteStepSpecToFile(templateCollection models.StepCollectionModel) error {
	pth := GetCurrentStepSpecPath()

	if exist, err := pathutil.IsPathExists(pth); err != nil {
		log.Error("[STEPMAN] - Failed to check path:", err)
		return err
	} else if exist == false {
		dir, _ := path.Split(pth)
		err := os.MkdirAll(dir, 0777)
		if err != nil {
			return err
		}
	} else {
		err := os.Remove(pth)
		if err != nil {
			return err
		}
	}

	file, err := os.OpenFile(pth, os.O_RDWR|os.O_CREATE, 0666)
	if err != nil {
		return err
	}
	defer func() {
		err := file.Close()
		if err != nil {
			log.Error("[STEPMAN] - Failed to close file:", err)
		}
	}()

	jsonContBytes, err := generateFormattedJSONForStepsSpec(templateCollection)
	if err != nil {
		return err
	}

	_, err = file.Write(jsonContBytes)
	if err != nil {
		return err
	}

	return nil
}

// ReadStepSpec ...
func ReadStepSpec() (models.StepCollectionModel, error) {
	pth := GetCurrentStepSpecPath()
	file, err := os.Open(pth)
	if err != nil {
		return models.StepCollectionModel{}, err
	}

	var stepCollection models.StepCollectionModel
	parser := json.NewDecoder(file)
	if err = parser.Decode(&stepCollection); err != nil {
		return models.StepCollectionModel{}, err
	}
	return stepCollection, err
}

// DownloadAndUnZIP ...
func DownloadAndUnZIP(url, pth string) error {
	filePath := os.TempDir() + "step.zip"
	file, err := os.Create(filePath)
	if err != nil {
		return err
	}
	defer func() {
		if err := file.Close(); err != nil {
			log.Fatal("Failed to close file:", err)
		}
		if err := os.Remove(filePath); err != nil {
			log.Fatal("Failed to remove file:", err)
		}
	}()

	response, err := http.Get(url)
	if err != nil {
		return err
	}
	defer func() {
		if err := response.Body.Close(); err != nil {
			log.Fatal("Failed to close response body:", err)
		}
	}()

	if response.StatusCode != http.StatusOK {
		errorMsg := "Failed to download step.zip from: " + url
		return errors.New(errorMsg)
	}

	log.Info("Successfully downloaded step.zip")
	if _, err := io.Copy(file, response.Body); err != nil {
		return err
	}

	return unzip(filePath, pth)
}

func unzip(src, dest string) error {
	r, err := zip.OpenReader(src)
	if err != nil {
		return err
	}
	defer func() {
		if err := r.Close(); err != nil {
			log.Fatal(err)
		}
	}()

	if err := os.MkdirAll(dest, 0755); err != nil {
		return err
	}

	// Closure to address file descriptors issue with all the deferred .Close() methods
	extractAndWriteFile := func(f *zip.File) error {
		rc, err := f.Open()
		if err != nil {
			return err
		}
		defer func() {
			if err := rc.Close(); err != nil {
				log.Fatal(err)
			}
		}()

		path := filepath.Join(dest, f.Name)

		if f.FileInfo().IsDir() {
			if err := os.MkdirAll(path, f.Mode()); err != nil {
				return err
			}
		} else {
			f, err := os.OpenFile(path, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, f.Mode())
			if err != nil {
				return err
			}
			defer func() {
				if err := f.Close(); err != nil {
					log.Fatal(err)
				}
			}()

			if _, err = io.Copy(f, rc); err != nil {
				return err
			}
		}
		return nil
	}

	for _, f := range r.File {
		if err := extractAndWriteFile(f); err != nil {
			return err
		}
	}
	return nil
}
<commit_msg>code style<commit_after>package stepman

import (
	"archive/zip"
	"encoding/json"
	"errors"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"path"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/bitrise-io/go-pathutil/pathutil"
	"github.com/bitrise-io/stepman/models"
	"gopkg.in/yaml.v2"
)

// DebugMode ...
var DebugMode bool

func parseStepYml(pth, id, version string) (models.StepModel, error) {
	bytes, err := ioutil.ReadFile(pth)
	if err != nil {
		return models.StepModel{}, err
	}

	var stepJSON models.StepModel
	if err := yaml.Unmarshal(bytes, &stepJSON); err != nil {
		return models.StepModel{}, err
	}

	stepJSON.ID = id
	stepJSON.VersionTag = version
	stepJSON.StepLibSource = CollectionURI

	return stepJSON, nil
}

// ParseStepCollection ...
func ParseStepCollection(pth string) (models.StepCollectionModel, error) {
	bytes, err := ioutil.ReadFile(pth)
	if err != nil {
		return models.StepCollectionModel{}, err
	}

	var stepCollection models.StepCollectionModel
	if err := yaml.Unmarshal(bytes, &stepCollection); err != nil {
		return models.StepCollectionModel{}, err
	}
	return stepCollection, nil
}

// DownloadStep ...
func DownloadStep(collection models.StepCollectionModel, step models.StepModel) error {
	downloadLocations := collection.GetDownloadLocations(step)

	success := false
	for _, downloadLocationMap := range downloadLocations {
		for key, value := range downloadLocationMap {
			switch key {
			case "zip":
				log.Info("[STEPMAN] - Downloading step from:", value)
				if err := DownloadAndUnZIP(value, GetStepPath(step)); err != nil {
					log.Error("[STEPMAN] - Failed to download step.zip:", err)
				} else {
					success = true
					return nil
				}
			case "git":
				log.Info("[STEPMAN] - Git clone step from:", value)
				if err := DoGitClone(value, GetStepPath(step)); err != nil {
					log.Error("[STEPMAN] - Failed to clone step:", err)
				} else {
					success = true
					return nil
				}
			default:
				log.Error("[STEPMAN] - Invalid download location")
			}
		}
	}

	if !success {
		return errors.New("Failed to download step")
	}
	return nil
}

// GetStepPath ...
// cach is cool
func GetStepPath(step models.StepModel) string {
	return GetCurrentStepCacheDir() + step.ID + "/" + step.VersionTag + "/"
}

// semantic version (X.Y.Z)
// true if version 2 is greater then version 1
func isVersionGrater(version1, version2 string) bool {
	version1Slice := strings.Split(version1, ".")
	version2Slice := strings.Split(version2, ".")

	for i, num := range version1Slice {
		num1, err1 := strconv.ParseInt(num, 0, 64)
		if err1 != nil {
			log.Error("[STEPMAN] - Failed to parse int:", err1)
			return false
		}

		num2, err2 := strconv.ParseInt(version2Slice[i], 0, 64)
		if err2 != nil {
			log.Error("[STEPMAN] - Failed to parse int:", err2)
			return false
		}

		if num2 > num1 {
			return true
		}
	}
	return false
}

func addStepToStepGroup(step models.StepModel, stepGroup models.StepGroupModel) models.StepGroupModel {
	var newStepGroup models.StepGroupModel
	if len(stepGroup.Versions) > 0 {
		// Step Group already created -> new version of step
		newStepGroup = stepGroup

		if isVersionGrater(newStepGroup.Latest.VersionTag, step.VersionTag) {
			newStepGroup.Latest = step
		}
	} else {
		// Create Step Group
		newStepGroup = models.StepGroupModel{}
		newStepGroup.Latest = step
	}

	versions := make([]models.StepModel, len(newStepGroup.Versions))
	for idx, step := range newStepGroup.Versions {
		versions[idx] = step
	}
	versions = append(versions, step)

	newStepGroup.Versions = versions
	newStepGroup.ID = step.ID
	return newStepGroup
}

func generateFormattedJSONForStepsSpec(templateCollection models.StepCollectionModel) ([]byte, error) {
	collection := models.StepCollectionModel{
		FormatVersion:        templateCollection.FormatVersion,
		GeneratedAtTimeStamp: time.Now().Unix(),
		SteplibSource:        CollectionURI,
		DownloadLocations:    templateCollection.DownloadLocations,
	}

	stepHash := models.StepHash{}

	stepsSpecDir := GetCurrentStepCollectionPath()
	err := filepath.Walk(stepsSpecDir, func(path string, f os.FileInfo, err error) error {
		truncatedPath := strings.Replace(path, stepsSpecDir, "", -1)
		match, matchErr := regexp.MatchString("([a-z]+).yml", truncatedPath)
		if matchErr != nil {
			return matchErr
		}

		if match {
			components := strings.Split(truncatedPath, "/")
			if len(components) == 4 {
				name := components[1]
				version := components[2]

				step, parseErr := parseStepYml(path, name, version)
				if parseErr != nil {
					return parseErr
				}
				stepGroup := addStepToStepGroup(step, stepHash[name])

				stepHash[name] = stepGroup
			} else {
				log.Debug("[STEPMAN] - Path:", truncatedPath)
				log.Debug("[STEPMAN] - Legth:", len(components))
			}
		}

		return err
	})
	if err != nil {
		log.Error("[STEPMAN] - Failed to walk through path:", err)
	}

	collection.Steps = stepHash

	var bytes []byte
	// if DebugMode == true {
	bytes, err = json.MarshalIndent(collection, "", "\t")
	// } else {
	// 	bytes, err = json.Marshal(collection)
	// }
	if err != nil {
		log.Error("[STEPMAN] - Failed to parse json:", err)
		return []byte{}, err
	}

	return bytes, nil
}

// WriteStepSpecToFile ...
func WriteStepSpecToFile(templateCollection models.StepCollectionModel) error {
	pth := GetCurrentStepSpecPath()

	if exist, err := pathutil.IsPathExists(pth); err != nil {
		log.Error("[STEPMAN] - Failed to check path:", err)
		return err
	} else if exist == false {
		dir, _ := path.Split(pth)
		err := os.MkdirAll(dir, 0777)
		if err != nil {
			return err
		}
	} else {
		err := os.Remove(pth)
		if err != nil {
			return err
		}
	}

	file, err := os.OpenFile(pth, os.O_RDWR|os.O_CREATE, 0666)
	if err != nil {
		return err
	}
	defer func() {
		err := file.Close()
		if err != nil {
			log.Error("[STEPMAN] - Failed to close file:", err)
		}
	}()

	jsonContBytes, err := generateFormattedJSONForStepsSpec(templateCollection)
	if err != nil {
		return err
	}

	_, err = file.Write(jsonContBytes)
	if err != nil {
		return err
	}

	return nil
}

// ReadStepSpec ...
func ReadStepSpec() (models.StepCollectionModel, error) {
	pth := GetCurrentStepSpecPath()
	file, err := os.Open(pth)
	if err != nil {
		return models.StepCollectionModel{}, err
	}

	var stepCollection models.StepCollectionModel
	parser := json.NewDecoder(file)
	if err = parser.Decode(&stepCollection); err != nil {
		return models.StepCollectionModel{}, err
	}
	return stepCollection, err
}

// DownloadAndUnZIP ...
func DownloadAndUnZIP(url, pth string) error {
	filePath := os.TempDir() + "step.zip"
	file, err := os.Create(filePath)
	if err != nil {
		return err
	}
	defer func() {
		if err := file.Close(); err != nil {
			log.Fatal("Failed to close file:", err)
		}
		if err := os.Remove(filePath); err != nil {
			log.Fatal("Failed to remove file:", err)
		}
	}()

	response, err := http.Get(url)
	if err != nil {
		return err
	}
	defer func() {
		if err := response.Body.Close(); err != nil {
			log.Fatal("Failed to close response body:", err)
		}
	}()

	if response.StatusCode != http.StatusOK {
		errorMsg := "Failed to download step.zip from: " + url
		return errors.New(errorMsg)
	}

	log.Info("Successfully downloaded step.zip")
	if _, err := io.Copy(file, response.Body); err != nil {
		return err
	}

	return unzip(filePath, pth)
}

func unzip(src, dest string) error {
	r, err := zip.OpenReader(src)
	if err != nil {
		return err
	}
	defer func() {
		if err := r.Close(); err != nil {
			log.Fatal(err)
		}
	}()

	if err := os.MkdirAll(dest, 0755); err != nil {
		return err
	}

	// Closure to address file descriptors issue with all the deferred .Close() methods
	extractAndWriteFile := func(f *zip.File) error {
		rc, err := f.Open()
		if err != nil {
			return err
		}
		defer func() {
			if err := rc.Close(); err != nil {
				log.Fatal(err)
			}
		}()

		path := filepath.Join(dest, f.Name)

		if f.FileInfo().IsDir() {
			if err := os.MkdirAll(path, f.Mode()); err != nil {
				return err
			}
		} else {
			f, err := os.OpenFile(path, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, f.Mode())
			if err != nil {
				return err
			}
			defer func() {
				if err := f.Close(); err != nil {
					log.Fatal(err)
				}
			}()

			if _, err = io.Copy(f, rc); err != nil {
				return err
			}
		}
		return nil
	}

	for _, f := range r.File {
		if err := extractAndWriteFile(f); err != nil {
			return err
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package store

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

	"github.com/Sirupsen/logrus"
	"github.com/aslanides/gocrud/req"
	"github.com/aslanides/gocrud/x"
)

// Update stores the create and update instructions, acting as the modifier
// to the entity Update relates to.
type Update struct {
	kind     string
	id       string
	source   string
	children []*Update
	parent   *Update
	edges    map[string]interface{}
	NanoTs   int64
}

// NewUpdate is the main entrypoint to updates. Returns back a Update
// object pointer, to run create and update operations on.
func NewUpdate(kind, id string) *Update {
	log.WithFields(logrus.Fields{
		"func": "NewUpdate",
		"kind": kind,
		"id":   id,
	}).Debug("Called")
	n := new(Update)
	n.kind = kind
	n.id = id
	n.NanoTs = time.Now().UnixNano()
	return n
}

// SetSource sets the author of the update. Generally, the userid of the
// modifier.
func (n *Update) SetSource(source string) *Update {
	n.source = source
	return n
}

// AddChild creates a new entity with the given kind, and creates a
// directed relationship from current entity to this new child entity.
// This is useful to generate arbitrarily deep and complex data structures,
// for e.g. Posts by Users, Comments on Posts, Likes on Posts etc.
//
// Retuns the Update pointer for the child entity, so any update operations
// done on this pointer would be reflected in the child entity. Note that
// a child can only have one parent by design.
func (n *Update) AddChild(kind string) *Update {
	log.WithField("childkind", kind).Debug("AddChild")
	child := new(Update)
	child.parent = n
	child.kind = kind
	child.NanoTs = n.NanoTs
	child.source = n.source
	n.children = append(n.children, child)
	return child
}

// Set allows you to set the property and value on the current entity.
// This would effectively replace any other value this property had,
// on this entity node pointer represents.
// But, no actual data would be overwritten though, as per the retention
// principle.
func (n *Update) Set(property string, value interface{}) *Update {
	log.WithField(property, value).Debug("Set")
	if n.edges == nil {
		n.edges = make(map[string]interface{})
	}
	n.edges[property] = value
	return n
}

// Marks the current entity for deletion. This is equivalent to doing a
// Set("delete_me", true), and then running q.FilterOut("delete_me") during
// query phase. Nothing is actually deleted though, as per the retention
// principle.
func (n *Update) MarkDeleted() *Update {
	return n.Set("_delete_", true)
}

func (n *Update) recPrint(l int) {
	log.Printf("Update[%d]: %+v", l, n)
	for _, child := range n.children {
		child.recPrint(l + 1)
	}
}

func (n *Update) root() *Update {
	for n.parent != nil {
		n = n.parent
	}
	return n
}

// Print finds the root from the given Update pointer, and does a recursive
// print on the tree for debugging purposes.
func (n *Update) Print() *Update {
	n = n.root()
	n.recPrint(0)
	return n
}

func (n *Update) Id() string {
	return n.id
}

func (n *Update) setCommitTs(tsNano int64) {
	n.NanoTs = tsNano
	for _, child := range n.children {
		child.setCommitTs(tsNano)
	}
}

// SetCommitTs allows you to set the commit timestamp of the update.
// The timestamp provided should be in nano-seconds. SetCommitTs
// can only be called on the root update node.
func (n *Update) SetCommitTs(tsNano int64) *Update {
	if n.parent != nil {
		log.Error("SetCommitTs called on a child node. Ignoring...")
		return n
	}
	n.setCommitTs(tsNano)
	return n
}

func (n *Update) doExecute(c *req.Context, its *[]*x.Instruction) error {
	for pred, val := range n.edges {
		if len(n.source) == 0 {
			return errors.New(fmt.Sprintf(
				"No source specified for id: %v kind: %v", n.id, n.kind))
		}

		i := new(x.Instruction)
		i.SubjectId = n.id
		i.SubjectType = n.kind
		i.Predicate = pred

		if b, err := json.Marshal(val); err != nil {
			return err
		} else {
			i.Object = b
		}
		i.Source = n.source
		i.NanoTs = n.NanoTs
		log.WithField("instruction", i).Debug("Pushing to list")
		*its = append(*its, i)
	}

	if len(n.children) == 0 {
		return nil
	}
	if len(n.source) == 0 {
		return errors.New(fmt.Sprintf(
			"No source specified for id: %v kind: %v", n.id, n.kind))
	}

	// Children can only be added, not deleted via API. But they can be stopped
	// from being retrieved.
	// Scenario: How do I stop childA from being retrieved?
	// Answer:
	// Modify child by adding a 'deleted' edge
	// Get(ChildKind, ChildId).Set("deleted", true).Execute(c)
	//
	// Then for retrieval from parent:
	// NewQuery(ParentKind, ParentId).Collect(ChildKind).FilterOut("deleted")
	// This would remove all children with a 'deleted' edge.
	// Better still
	// Get(ChildKind, ChildId).MarkDeleted().Execute(c)
	// would automatically filter out that child and it's children
	// from being retrieved.

	for _, child := range n.children {
		if len(child.id) > 0 {
			log.WithField("child_id", child.id).Fatal(
				"Child id should be empty for all current use cases")
			return errors.New("Non empty child id")
		}

		for idx := 0; ; idx++ { // Retry loop.
			child.id = x.UniqueString(c.NumCharsUnique)
			log.WithField("id", child.id).Debug("Checking availability of new id")
			if isnew := Get().IsNew(child.id); isnew {
				log.WithField("id", child.id).Debug("New id available")
				break
			}
			if idx >= 30 {
				return errors.New("Unable to find new id")
			}
		}
		// Create edge from parent to child
		i := new(x.Instruction)
		i.SubjectId = n.id
		i.SubjectType = n.kind
		i.Predicate = child.kind
		i.ObjectId = child.id
		i.Source = n.source
		i.NanoTs = n.NanoTs
		log.WithField("instruction", i).Debug("Pushing to list")
		*its = append(*its, i)

		// Create edge from child to parent
		i = new(x.Instruction)
		i.SubjectId = child.id
		i.SubjectType = child.kind
		i.Predicate = "_parent_"
		i.ObjectId = n.id
		i.Source = n.source
		i.NanoTs = n.NanoTs
		log.WithField("instruction", i).Debug("Pushing to list")
		*its = append(*its, i)

		if err := child.doExecute(c, its); err != nil {
			return err
		}
	}
	return nil
}

// Execute finds the root from the given Update pointer, recursively generates
// the set of instructions to store, and commits them. Returns any errors
// encountered during these steps.
func (n *Update) Execute(c *req.Context) error {
	if c.NumCharsUnique <= 0 {
		log.Fatal("Invalid number of chars for generating unique ids. Set req.Context.NumCharsUnique")
		return errors.New("Invalid req.Context.NumCharsUnique")
	}

	n = n.root()

	var its []*x.Instruction
	err := n.doExecute(c, &its)
	if err != nil {
		return err
	}
	if len(its) == 0 {
		return errors.New("No instructions generated")
	}

	if rerr := Get().Commit(its); rerr != nil {
		return rerr
	}

	if c.HasIndexer {
		// This block of code figures out which entities have been modified,
		// and sends them off to be updated by indexer.
		updates := make(map[x.Entity]bool)
		for _, it := range its {
			e := x.Entity{Kind: it.SubjectType, Id: it.SubjectId}
			if _, present := updates[e]; present { // find distinct entities
				continue
			}
			c.Updates <- e
			updates[e] = true
		}
	}
	return nil
}
<commit_msg>Add MarkRestored()<commit_after>package store

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

	"github.com/Sirupsen/logrus"
	"github.com/aslanides/gocrud/req"
	"github.com/aslanides/gocrud/x"
)

// Update stores the create and update instructions, acting as the modifier
// to the entity Update relates to.
type Update struct {
	kind     string
	id       string
	source   string
	children []*Update
	parent   *Update
	edges    map[string]interface{}
	NanoTs   int64
}

// NewUpdate is the main entrypoint to updates. Returns back a Update
// object pointer, to run create and update operations on.
func NewUpdate(kind, id string) *Update {
	log.WithFields(logrus.Fields{
		"func": "NewUpdate",
		"kind": kind,
		"id":   id,
	}).Debug("Called")
	n := new(Update)
	n.kind = kind
	n.id = id
	n.NanoTs = time.Now().UnixNano()
	return n
}

// SetSource sets the author of the update. Generally, the userid of the
// modifier.
func (n *Update) SetSource(source string) *Update {
	n.source = source
	return n
}

// AddChild creates a new entity with the given kind, and creates a
// directed relationship from current entity to this new child entity.
// This is useful to generate arbitrarily deep and complex data structures,
// for e.g. Posts by Users, Comments on Posts, Likes on Posts etc.
//
// Retuns the Update pointer for the child entity, so any update operations
// done on this pointer would be reflected in the child entity. Note that
// a child can only have one parent by design.
func (n *Update) AddChild(kind string) *Update {
	log.WithField("childkind", kind).Debug("AddChild")
	child := new(Update)
	child.parent = n
	child.kind = kind
	child.NanoTs = n.NanoTs
	child.source = n.source
	n.children = append(n.children, child)
	return child
}

// Set allows you to set the property and value on the current entity.
// This would effectively replace any other value this property had,
// on this entity node pointer represents.
// But, no actual data would be overwritten though, as per the retention
// principle.
func (n *Update) Set(property string, value interface{}) *Update {
	log.WithField(property, value).Debug("Set")
	if n.edges == nil {
		n.edges = make(map[string]interface{})
	}
	n.edges[property] = value
	return n
}

// Marks the current entity for deletion. This is equivalent to doing a
// Set("delete_me", true), and then running q.FilterOut("delete_me") during
// query phase. Nothing is actually deleted though, as per the retention
// principle.
func (n *Update) MarkDeleted() *Update {
	return n.Set("_delete_", true)
}

func (n *Update) MarkRestored() *Update {
	return n.Set("_delete_", false)
}

func (n *Update) recPrint(l int) {
	log.Printf("Update[%d]: %+v", l, n)
	for _, child := range n.children {
		child.recPrint(l + 1)
	}
}

func (n *Update) root() *Update {
	for n.parent != nil {
		n = n.parent
	}
	return n
}

// Print finds the root from the given Update pointer, and does a recursive
// print on the tree for debugging purposes.
func (n *Update) Print() *Update {
	n = n.root()
	n.recPrint(0)
	return n
}

func (n *Update) Id() string {
	return n.id
}

func (n *Update) setCommitTs(tsNano int64) {
	n.NanoTs = tsNano
	for _, child := range n.children {
		child.setCommitTs(tsNano)
	}
}

// SetCommitTs allows you to set the commit timestamp of the update.
// The timestamp provided should be in nano-seconds. SetCommitTs
// can only be called on the root update node.
func (n *Update) SetCommitTs(tsNano int64) *Update {
	if n.parent != nil {
		log.Error("SetCommitTs called on a child node. Ignoring...")
		return n
	}
	n.setCommitTs(tsNano)
	return n
}

func (n *Update) doExecute(c *req.Context, its *[]*x.Instruction) error {
	for pred, val := range n.edges {
		if len(n.source) == 0 {
			return errors.New(fmt.Sprintf(
				"No source specified for id: %v kind: %v", n.id, n.kind))
		}

		i := new(x.Instruction)
		i.SubjectId = n.id
		i.SubjectType = n.kind
		i.Predicate = pred

		if b, err := json.Marshal(val); err != nil {
			return err
		} else {
			i.Object = b
		}
		i.Source = n.source
		i.NanoTs = n.NanoTs
		log.WithField("instruction", i).Debug("Pushing to list")
		*its = append(*its, i)
	}

	if len(n.children) == 0 {
		return nil
	}
	if len(n.source) == 0 {
		return errors.New(fmt.Sprintf(
			"No source specified for id: %v kind: %v", n.id, n.kind))
	}

	// Children can only be added, not deleted via API. But they can be stopped
	// from being retrieved.
	// Scenario: How do I stop childA from being retrieved?
	// Answer:
	// Modify child by adding a 'deleted' edge
	// Get(ChildKind, ChildId).Set("deleted", true).Execute(c)
	//
	// Then for retrieval from parent:
	// NewQuery(ParentKind, ParentId).Collect(ChildKind).FilterOut("deleted")
	// This would remove all children with a 'deleted' edge.
	// Better still
	// Get(ChildKind, ChildId).MarkDeleted().Execute(c)
	// would automatically filter out that child and it's children
	// from being retrieved.

	for _, child := range n.children {
		if len(child.id) > 0 {
			log.WithField("child_id", child.id).Fatal(
				"Child id should be empty for all current use cases")
			return errors.New("Non empty child id")
		}

		for idx := 0; ; idx++ { // Retry loop.
			child.id = x.UniqueString(c.NumCharsUnique)
			log.WithField("id", child.id).Debug("Checking availability of new id")
			if isnew := Get().IsNew(child.id); isnew {
				log.WithField("id", child.id).Debug("New id available")
				break
			}
			if idx >= 30 {
				return errors.New("Unable to find new id")
			}
		}
		// Create edge from parent to child
		i := new(x.Instruction)
		i.SubjectId = n.id
		i.SubjectType = n.kind
		i.Predicate = child.kind
		i.ObjectId = child.id
		i.Source = n.source
		i.NanoTs = n.NanoTs
		log.WithField("instruction", i).Debug("Pushing to list")
		*its = append(*its, i)

		// Create edge from child to parent
		i = new(x.Instruction)
		i.SubjectId = child.id
		i.SubjectType = child.kind
		i.Predicate = "_parent_"
		i.ObjectId = n.id
		i.Source = n.source
		i.NanoTs = n.NanoTs
		log.WithField("instruction", i).Debug("Pushing to list")
		*its = append(*its, i)

		if err := child.doExecute(c, its); err != nil {
			return err
		}
	}
	return nil
}

// Execute finds the root from the given Update pointer, recursively generates
// the set of instructions to store, and commits them. Returns any errors
// encountered during these steps.
func (n *Update) Execute(c *req.Context) error {
	if c.NumCharsUnique <= 0 {
		log.Fatal("Invalid number of chars for generating unique ids. Set req.Context.NumCharsUnique")
		return errors.New("Invalid req.Context.NumCharsUnique")
	}

	n = n.root()

	var its []*x.Instruction
	err := n.doExecute(c, &its)
	if err != nil {
		return err
	}
	if len(its) == 0 {
		return errors.New("No instructions generated")
	}

	if rerr := Get().Commit(its); rerr != nil {
		return rerr
	}

	if c.HasIndexer {
		// This block of code figures out which entities have been modified,
		// and sends them off to be updated by indexer.
		updates := make(map[x.Entity]bool)
		for _, it := range its {
			e := x.Entity{Kind: it.SubjectType, Id: it.SubjectId}
			if _, present := updates[e]; present { // find distinct entities
				continue
			}
			c.Updates <- e
			updates[e] = true
		}
	}
	return nil
}
<|endoftext|>


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

// Package publisher provides a type to publish names to a mounttable.
package publisher

// TODO(toddw): Add unittests.

import (
	"fmt"
	"sort"
	"strings"
	"time"

	"v.io/v23/context"
	"v.io/v23/namespace"
	"v.io/v23/naming"
	"v.io/v23/rpc"
	"v.io/v23/verror"
	"v.io/x/lib/vlog"
)

// Publisher manages the publishing of servers in mounttable.
type Publisher interface {
	// AddServer adds a new server to be mounted.
	AddServer(server string)
	// RemoveServer removes a server from the list of mounts.
	RemoveServer(server string)
	// AddName adds a new name for all servers to be mounted as.
	AddName(name string, ServesMountTable bool, IsLeaf bool)
	// RemoveName removes a name.
	RemoveName(name string)
	// Status returns a snapshot of the publisher's current state.
	Status() rpc.MountState
	// DebugString returns a string representation of the publisher
	// meant solely for debugging.
	DebugString() string
	// Stop causes the publishing to stop and initiates unmounting of the
	// mounted names.  Stop performs the unmounting asynchronously, and
	// WaitForStop should be used to wait until it is done.
	// Once Stop is called Add/RemoveServer and AddName become noops.
	Stop()
	// WaitForStop waits until all unmounting initiated by Stop is finished.
	WaitForStop()
}

// The publisher adds this much slack to each TTL.
const mountTTLSlack = 20 * time.Second

// publisher maintains the name->server associations in the mounttable.  It
// spawns its own goroutine that does the actual work; the publisher itself
// simply coordinates concurrent access by sending and receiving on the
// appropriate channels.
type publisher struct {
	cmdchan  chan interface{} // value is one of {server,name,debug}Cmd
	stopchan chan struct{}    // closed when no longer accepting commands.
	donechan chan struct{}    // closed when the publisher is done
}

type addServerCmd struct {
	server string // server to add
}

type removeServerCmd struct {
	server string // server to remove
}

type addNameCmd struct {
	name string // name to add
	mt   bool   // true if server serves a mount table
	leaf bool   // true if server is a leaf
}

type removeNameCmd struct {
	name string // name to remove
}

type debugCmd chan string // debug string is sent when the cmd is done

type statusCmd chan rpc.MountState // status info is sent when cmd is done

type stopCmd struct{} // sent to the runloop when we want it to exit.

// New returns a new publisher that updates mounts on ns every period.
func New(ctx *context.T, ns namespace.T, period time.Duration) Publisher {
	p := &publisher{
		cmdchan:  make(chan interface{}),
		stopchan: make(chan struct{}),
		donechan: make(chan struct{}),
	}
	go runLoop(ctx, p.cmdchan, p.donechan, ns, period)
	return p
}

func (p *publisher) sendCmd(cmd interface{}) bool {
	select {
	case p.cmdchan <- cmd:
		return true
	case <-p.stopchan:
		return false
	case <-p.donechan:
		return false
	}
}

func (p *publisher) AddServer(server string) {
	p.sendCmd(addServerCmd{server})
}

func (p *publisher) RemoveServer(server string) {
	p.sendCmd(removeServerCmd{server})
}

func (p *publisher) AddName(name string, mt bool, leaf bool) {
	p.sendCmd(addNameCmd{name, mt, leaf})
}

func (p *publisher) RemoveName(name string) {
	p.sendCmd(removeNameCmd{name})
}

func (p *publisher) Status() rpc.MountState {
	status := make(statusCmd)
	if p.sendCmd(status) {
		return <-status
	}
	return rpc.MountState{}
}

func (p *publisher) DebugString() (dbg string) {
	debug := make(debugCmd)
	if p.sendCmd(debug) {
		dbg = <-debug
	} else {
		dbg = "stopped"
	}
	return
}

// Stop stops the publisher, which in practical terms means un-mounting
// everything and preventing any further publish operations.  The caller can
// be confident that no new names or servers will get published once Stop
// returns.  To wait for existing mounts to be cleaned up, use WaitForStop.
//
// Stopping the publisher is irreversible.
//
// Once the publisher is stopped, any further calls on its public methods
// (including Stop) are no-ops.
func (p *publisher) Stop() {
	p.sendCmd(stopCmd{})
	close(p.stopchan) // stop accepting new commands now.
}

func (p *publisher) WaitForStop() {
	<-p.donechan
}

func runLoop(ctx *context.T, cmdchan chan interface{}, donechan chan struct{}, ns namespace.T, period time.Duration) {
	vlog.VI(2).Info("rpc pub: start runLoop")
	state := newPubState(ctx, ns, period)
	for {
		select {
		case cmd := <-cmdchan:
			switch tcmd := cmd.(type) {
			case stopCmd:
				state.unmountAll()
				close(donechan)
				vlog.VI(2).Info("rpc pub: exit runLoop")
				return
			case addServerCmd:
				state.addServer(tcmd.server)
			case removeServerCmd:
				state.removeServer(tcmd.server)
			case addNameCmd:
				state.addName(tcmd.name, tcmd.mt, tcmd.leaf)
			case removeNameCmd:
				state.removeName(tcmd.name)
			case statusCmd:
				tcmd <- state.getStatus()
				close(tcmd)
			case debugCmd:
				tcmd <- state.debugString()
				close(tcmd)
			}
		case <-state.timeout():
			// Sync everything once every period, to refresh the ttls.
			state.sync()
		}
	}
}

type mountKey struct {
	name, server string
}

// pubState maintains the state for our periodic mounts.  It is not thread-safe;
// it's only used in the sequential publisher runLoop.
type pubState struct {
	ctx      *context.T
	ns       namespace.T
	period   time.Duration
	deadline time.Time           // deadline for the next sync call
	names    map[string]nameAttr // names that have been added
	servers  map[string]bool     // servers that have been added, true
	// map each (name,server) to its status.
	mounts map[mountKey]*rpc.MountStatus
}

type nameAttr struct {
	servesMT bool
	isLeaf   bool
}

func newPubState(ctx *context.T, ns namespace.T, period time.Duration) *pubState {
	return &pubState{
		ctx:      ctx,
		ns:       ns,
		period:   period,
		deadline: time.Now().Add(period),
		names:    make(map[string]nameAttr),
		servers:  make(map[string]bool),
		mounts:   make(map[mountKey]*rpc.MountStatus),
	}
}

func (ps *pubState) timeout() <-chan time.Time {
	return time.After(ps.deadline.Sub(time.Now()))
}

func (ps *pubState) addName(name string, mt bool, leaf bool) {
	// Each non-dup name that is added causes new mounts to be created for all
	// existing servers.
	if _, exists := ps.names[name]; exists {
		return
	}
	attr := nameAttr{mt, leaf}
	ps.names[name] = attr
	for server, _ := range ps.servers {
		status := new(rpc.MountStatus)
		ps.mounts[mountKey{name, server}] = status
		ps.mount(name, server, status, attr)
	}
}

func (ps *pubState) removeName(name string) {
	if _, exists := ps.names[name]; !exists {
		return
	}
	for server, _ := range ps.servers {
		if status, exists := ps.mounts[mountKey{name, server}]; exists {
			ps.unmount(name, server, status)
		}
	}
	delete(ps.names, name)
}

func (ps *pubState) addServer(server string) {
	// Each non-dup server that is added causes new mounts to be created for all
	// existing names.
	if _, exists := ps.servers[server]; !exists {
		ps.servers[server] = true
		for name, attr := range ps.names {
			status := new(rpc.MountStatus)
			ps.mounts[mountKey{name, server}] = status
			ps.mount(name, server, status, attr)
		}
	}
}

func (ps *pubState) removeServer(server string) {
	if _, exists := ps.servers[server]; !exists {
		return
	}
	delete(ps.servers, server)
	for name, _ := range ps.names {
		if status, exists := ps.mounts[mountKey{name, server}]; exists {
			ps.unmount(name, server, status)
		}
	}
}

func (ps *pubState) mount(name, server string, status *rpc.MountStatus, attr nameAttr) {
	// Always mount with ttl = period + slack, regardless of whether this is
	// triggered by a newly added server or name, or by sync.  The next call
	// to sync will occur within the next period, and refresh all mounts.
	ttl := ps.period + mountTTLSlack
	last := status
	status.LastMount = time.Now()
	status.LastMountErr = ps.ns.Mount(ps.ctx, name, server, ttl, naming.ServesMountTable(attr.servesMT), naming.IsLeaf(attr.isLeaf))
	status.TTL = ttl
	// If the mount status changed, log it.
	if status.LastMountErr != nil {
		if verror.ErrorID(last.LastMountErr) != verror.ErrorID(status.LastMountErr) || vlog.V(2) {
			vlog.Errorf("rpc pub: couldn't mount(%v, %v, %v): %v", name, server, ttl, status.LastMountErr)
		}
	} else {
		if last.LastMount.IsZero() || last.LastMountErr != nil || vlog.V(2) {
			vlog.Infof("rpc pub: mount(%v, %v, %v)", name, server, ttl)
		}
	}
}

func (ps *pubState) sync() {
	ps.deadline = time.Now().Add(ps.period) // set deadline for the next sync
	for key, status := range ps.mounts {
		if status.LastUnmountErr != nil {
			// Desired state is "unmounted", failed at previous attempt. Retry.
			ps.unmount(key.name, key.server, status)
		} else {
			ps.mount(key.name, key.server, status, ps.names[key.name])
		}
	}
}

func (ps *pubState) unmount(name, server string, status *rpc.MountStatus) {
	status.LastUnmount = time.Now()
	status.LastUnmountErr = ps.ns.Unmount(ps.ctx, name, server)
	if status.LastUnmountErr != nil {
		vlog.Errorf("rpc pub: couldn't unmount(%v, %v): %v", name, server, status.LastUnmountErr)
	} else {
		vlog.VI(1).Infof("rpc pub: unmount(%v, %v)", name, server)
		delete(ps.mounts, mountKey{name, server})
	}
}

func (ps *pubState) unmountAll() {
	for key, status := range ps.mounts {
		ps.unmount(key.name, key.server, status)
	}
}

func copyNamesToSlice(sl map[string]nameAttr) []string {
	var ret []string
	for s, _ := range sl {
		if len(s) == 0 {
			continue
		}
		ret = append(ret, s)
	}
	return ret
}

func copyServersToSlice(sl map[string]bool) []string {
	var ret []string
	for s, _ := range sl {
		if len(s) == 0 {
			continue
		}
		ret = append(ret, s)
	}
	return ret
}

func (ps *pubState) getStatus() rpc.MountState {
	st := make([]rpc.MountStatus, 0, len(ps.mounts))
	names := copyNamesToSlice(ps.names)
	servers := copyServersToSlice(ps.servers)
	sort.Strings(names)
	sort.Strings(servers)
	for _, name := range names {
		for _, server := range servers {
			if v := ps.mounts[mountKey{name, server}]; v != nil {
				mst := *v
				mst.Name = name
				mst.Server = server
				st = append(st, mst)
			}
		}
	}
	return st
}

// TODO(toddw): sort the names/servers so that the output order is stable.
func (ps *pubState) debugString() string {
	l := make([]string, 2+len(ps.mounts))
	l = append(l, fmt.Sprintf("Publisher period:%v deadline:%v", ps.period, ps.deadline))
	l = append(l, "==============================Mounts============================================")
	for key, status := range ps.mounts {
		l = append(l, fmt.Sprintf("[%s,%s] mount(%v, %v, %s) unmount(%v, %v)", key.name, key.server, status.LastMount, status.LastMountErr, status.TTL, status.LastUnmount, status.LastUnmountErr))
	}
	return strings.Join(l, "\n")
}
<commit_msg>ref/runtime/internal/lib/publisher: When you stop the publisher, don't retry unmount.<commit_after>// Copyright 2015 The Vanadium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package publisher provides a type to publish names to a mounttable.
package publisher

// TODO(toddw): Add unittests.

import (
	"fmt"
	"sort"
	"strings"
	"time"

	"v.io/v23/context"
	"v.io/v23/namespace"
	"v.io/v23/naming"
	"v.io/v23/options"
	"v.io/v23/rpc"
	"v.io/v23/verror"
	"v.io/x/lib/vlog"
)

// Publisher manages the publishing of servers in mounttable.
type Publisher interface {
	// AddServer adds a new server to be mounted.
	AddServer(server string)
	// RemoveServer removes a server from the list of mounts.
	RemoveServer(server string)
	// AddName adds a new name for all servers to be mounted as.
	AddName(name string, ServesMountTable bool, IsLeaf bool)
	// RemoveName removes a name.
	RemoveName(name string)
	// Status returns a snapshot of the publisher's current state.
	Status() rpc.MountState
	// DebugString returns a string representation of the publisher
	// meant solely for debugging.
	DebugString() string
	// Stop causes the publishing to stop and initiates unmounting of the
	// mounted names.  Stop performs the unmounting asynchronously, and
	// WaitForStop should be used to wait until it is done.
	// Once Stop is called Add/RemoveServer and AddName become noops.
	Stop()
	// WaitForStop waits until all unmounting initiated by Stop is finished.
	WaitForStop()
}

// The publisher adds this much slack to each TTL.
const mountTTLSlack = 20 * time.Second

// publisher maintains the name->server associations in the mounttable.  It
// spawns its own goroutine that does the actual work; the publisher itself
// simply coordinates concurrent access by sending and receiving on the
// appropriate channels.
type publisher struct {
	cmdchan  chan interface{} // value is one of {server,name,debug}Cmd
	stopchan chan struct{}    // closed when no longer accepting commands.
	donechan chan struct{}    // closed when the publisher is done
}

type addServerCmd struct {
	server string // server to add
}

type removeServerCmd struct {
	server string // server to remove
}

type addNameCmd struct {
	name string // name to add
	mt   bool   // true if server serves a mount table
	leaf bool   // true if server is a leaf
}

type removeNameCmd struct {
	name string // name to remove
}

type debugCmd chan string // debug string is sent when the cmd is done

type statusCmd chan rpc.MountState // status info is sent when cmd is done

type stopCmd struct{} // sent to the runloop when we want it to exit.

// New returns a new publisher that updates mounts on ns every period.
func New(ctx *context.T, ns namespace.T, period time.Duration) Publisher {
	p := &publisher{
		cmdchan:  make(chan interface{}),
		stopchan: make(chan struct{}),
		donechan: make(chan struct{}),
	}
	go runLoop(ctx, p.cmdchan, p.donechan, ns, period)
	return p
}

func (p *publisher) sendCmd(cmd interface{}) bool {
	select {
	case p.cmdchan <- cmd:
		return true
	case <-p.stopchan:
		return false
	case <-p.donechan:
		return false
	}
}

func (p *publisher) AddServer(server string) {
	p.sendCmd(addServerCmd{server})
}

func (p *publisher) RemoveServer(server string) {
	p.sendCmd(removeServerCmd{server})
}

func (p *publisher) AddName(name string, mt bool, leaf bool) {
	p.sendCmd(addNameCmd{name, mt, leaf})
}

func (p *publisher) RemoveName(name string) {
	p.sendCmd(removeNameCmd{name})
}

func (p *publisher) Status() rpc.MountState {
	status := make(statusCmd)
	if p.sendCmd(status) {
		return <-status
	}
	return rpc.MountState{}
}

func (p *publisher) DebugString() (dbg string) {
	debug := make(debugCmd)
	if p.sendCmd(debug) {
		dbg = <-debug
	} else {
		dbg = "stopped"
	}
	return
}

// Stop stops the publisher, which in practical terms means un-mounting
// everything and preventing any further publish operations.  The caller can
// be confident that no new names or servers will get published once Stop
// returns.  To wait for existing mounts to be cleaned up, use WaitForStop.
//
// Stopping the publisher is irreversible.
//
// Once the publisher is stopped, any further calls on its public methods
// (including Stop) are no-ops.
func (p *publisher) Stop() {
	p.sendCmd(stopCmd{})
	close(p.stopchan) // stop accepting new commands now.
}

func (p *publisher) WaitForStop() {
	<-p.donechan
}

func runLoop(ctx *context.T, cmdchan chan interface{}, donechan chan struct{}, ns namespace.T, period time.Duration) {
	vlog.VI(2).Info("rpc pub: start runLoop")
	state := newPubState(ctx, ns, period)
	for {
		select {
		case cmd := <-cmdchan:
			switch tcmd := cmd.(type) {
			case stopCmd:
				state.unmountAll()
				close(donechan)
				vlog.VI(2).Info("rpc pub: exit runLoop")
				return
			case addServerCmd:
				state.addServer(tcmd.server)
			case removeServerCmd:
				state.removeServer(tcmd.server)
			case addNameCmd:
				state.addName(tcmd.name, tcmd.mt, tcmd.leaf)
			case removeNameCmd:
				state.removeName(tcmd.name)
			case statusCmd:
				tcmd <- state.getStatus()
				close(tcmd)
			case debugCmd:
				tcmd <- state.debugString()
				close(tcmd)
			}
		case <-state.timeout():
			// Sync everything once every period, to refresh the ttls.
			state.sync()
		}
	}
}

type mountKey struct {
	name, server string
}

// pubState maintains the state for our periodic mounts.  It is not thread-safe;
// it's only used in the sequential publisher runLoop.
type pubState struct {
	ctx      *context.T
	ns       namespace.T
	period   time.Duration
	deadline time.Time           // deadline for the next sync call
	names    map[string]nameAttr // names that have been added
	servers  map[string]bool     // servers that have been added, true
	// map each (name,server) to its status.
	mounts map[mountKey]*rpc.MountStatus
}

type nameAttr struct {
	servesMT bool
	isLeaf   bool
}

func newPubState(ctx *context.T, ns namespace.T, period time.Duration) *pubState {
	return &pubState{
		ctx:      ctx,
		ns:       ns,
		period:   period,
		deadline: time.Now().Add(period),
		names:    make(map[string]nameAttr),
		servers:  make(map[string]bool),
		mounts:   make(map[mountKey]*rpc.MountStatus),
	}
}

func (ps *pubState) timeout() <-chan time.Time {
	return time.After(ps.deadline.Sub(time.Now()))
}

func (ps *pubState) addName(name string, mt bool, leaf bool) {
	// Each non-dup name that is added causes new mounts to be created for all
	// existing servers.
	if _, exists := ps.names[name]; exists {
		return
	}
	attr := nameAttr{mt, leaf}
	ps.names[name] = attr
	for server, _ := range ps.servers {
		status := new(rpc.MountStatus)
		ps.mounts[mountKey{name, server}] = status
		ps.mount(name, server, status, attr)
	}
}

func (ps *pubState) removeName(name string) {
	if _, exists := ps.names[name]; !exists {
		return
	}
	for server, _ := range ps.servers {
		if status, exists := ps.mounts[mountKey{name, server}]; exists {
			ps.unmount(name, server, status, true)
		}
	}
	delete(ps.names, name)
}

func (ps *pubState) addServer(server string) {
	// Each non-dup server that is added causes new mounts to be created for all
	// existing names.
	if _, exists := ps.servers[server]; !exists {
		ps.servers[server] = true
		for name, attr := range ps.names {
			status := new(rpc.MountStatus)
			ps.mounts[mountKey{name, server}] = status
			ps.mount(name, server, status, attr)
		}
	}
}

func (ps *pubState) removeServer(server string) {
	if _, exists := ps.servers[server]; !exists {
		return
	}
	delete(ps.servers, server)
	for name, _ := range ps.names {
		if status, exists := ps.mounts[mountKey{name, server}]; exists {
			ps.unmount(name, server, status, true)
		}
	}
}

func (ps *pubState) mount(name, server string, status *rpc.MountStatus, attr nameAttr) {
	// Always mount with ttl = period + slack, regardless of whether this is
	// triggered by a newly added server or name, or by sync.  The next call
	// to sync will occur within the next period, and refresh all mounts.
	ttl := ps.period + mountTTLSlack
	last := status
	status.LastMount = time.Now()
	status.LastMountErr = ps.ns.Mount(ps.ctx, name, server, ttl, naming.ServesMountTable(attr.servesMT), naming.IsLeaf(attr.isLeaf))
	status.TTL = ttl
	// If the mount status changed, log it.
	if status.LastMountErr != nil {
		if verror.ErrorID(last.LastMountErr) != verror.ErrorID(status.LastMountErr) || vlog.V(2) {
			vlog.Errorf("rpc pub: couldn't mount(%v, %v, %v): %v", name, server, ttl, status.LastMountErr)
		}
	} else {
		if last.LastMount.IsZero() || last.LastMountErr != nil || vlog.V(2) {
			vlog.Infof("rpc pub: mount(%v, %v, %v)", name, server, ttl)
		}
	}
}

func (ps *pubState) sync() {
	ps.deadline = time.Now().Add(ps.period) // set deadline for the next sync
	for key, status := range ps.mounts {
		if status.LastUnmountErr != nil {
			// Desired state is "unmounted", failed at previous attempt. Retry.
			ps.unmount(key.name, key.server, status, true)
		} else {
			ps.mount(key.name, key.server, status, ps.names[key.name])
		}
	}
}

func (ps *pubState) unmount(name, server string, status *rpc.MountStatus, retry bool) {
	status.LastUnmount = time.Now()
	var opts []naming.NamespaceOpt
	if !retry {
		opts = []naming.NamespaceOpt{options.NoRetry{}}
	}
	status.LastUnmountErr = ps.ns.Unmount(ps.ctx, name, server, opts...)
	if status.LastUnmountErr != nil {
		vlog.Errorf("rpc pub: couldn't unmount(%v, %v): %v", name, server, status.LastUnmountErr)
	} else {
		vlog.VI(1).Infof("rpc pub: unmount(%v, %v)", name, server)
		delete(ps.mounts, mountKey{name, server})
	}
}

func (ps *pubState) unmountAll() {
	for key, status := range ps.mounts {
		ps.unmount(key.name, key.server, status, false)
	}
}

func copyNamesToSlice(sl map[string]nameAttr) []string {
	var ret []string
	for s, _ := range sl {
		if len(s) == 0 {
			continue
		}
		ret = append(ret, s)
	}
	return ret
}

func copyServersToSlice(sl map[string]bool) []string {
	var ret []string
	for s, _ := range sl {
		if len(s) == 0 {
			continue
		}
		ret = append(ret, s)
	}
	return ret
}

func (ps *pubState) getStatus() rpc.MountState {
	st := make([]rpc.MountStatus, 0, len(ps.mounts))
	names := copyNamesToSlice(ps.names)
	servers := copyServersToSlice(ps.servers)
	sort.Strings(names)
	sort.Strings(servers)
	for _, name := range names {
		for _, server := range servers {
			if v := ps.mounts[mountKey{name, server}]; v != nil {
				mst := *v
				mst.Name = name
				mst.Server = server
				st = append(st, mst)
			}
		}
	}
	return st
}

// TODO(toddw): sort the names/servers so that the output order is stable.
func (ps *pubState) debugString() string {
	l := make([]string, 2+len(ps.mounts))
	l = append(l, fmt.Sprintf("Publisher period:%v deadline:%v", ps.period, ps.deadline))
	l = append(l, "==============================Mounts============================================")
	for key, status := range ps.mounts {
		l = append(l, fmt.Sprintf("[%s,%s] mount(%v, %v, %s) unmount(%v, %v)", key.name, key.server, status.LastMount, status.LastMountErr, status.TTL, status.LastUnmount, status.LastUnmountErr))
	}
	return strings.Join(l, "\n")
}
<|endoftext|>


<commit_before>package leetcode

import (
	"github.com/stretchr/testify/assert"
	"testing"
)

func Test_groupAnagrams(t *testing.T) {
	assert := assert.New(t)

	assert.Equal(groupAnagrams([]string{"eat", "tea", "tan", "ate", "nat", "bat"}), [][]string{
		{"eat", "tea", "ate"}, {"tan", "nat"}, {"bat"},
	})
}
<commit_msg>enhance unordered output test<commit_after>package leetcode

import (
	"github.com/stretchr/testify/assert"
	"testing"
)

func Test_groupAnagrams(t *testing.T) {
	assert := assert.New(t)

	output := groupAnagrams([]string{"eat", "tea", "tan", "ate", "nat", "bat"})
	assert.Equal(len(output), 3)
	for _, sub := range output {
		switch len(sub) {
		case 1:
			assert.Equal(sub, []string{"bat"})
		case 2:
			assert.Equal(sub, []string{"tan", "nat"})
		case 3:
			assert.Equal(sub, []string{"eat", "tea", "ate"})
		}
	}
}
<|endoftext|>


<commit_before>// Copyright © 2016 The Things Network
// Use of this source code is governed by the MIT license that can be found in the LICENSE file.

package cmd

import (
	"encoding/json"
	"fmt"
	"net/http"
	"net/url"
	"strings"

	"github.com/TheThingsNetwork/ttn/ttnctl/util"
	"github.com/apex/log"
	"github.com/gosuri/uitable"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"
)

type app struct {
	EUI   string `json:"eui"`
	Name  string `json:"name"`
	Owner string `json:"owner"`
}

// applicationsCmd represents the applications command
var applicationsCmd = &cobra.Command{
	Use:   "applications",
	Short: "Show applications",
	Long:  `ttnctl applications retrieves the applications of the logged on user.`,
	Run: func(cmd *cobra.Command, args []string) {
		server := viper.GetString("ttn-account-server")
		req, err := util.NewRequestWithAuth(server, "GET", fmt.Sprintf("%s/applications", server), nil)
		if err != nil {
			ctx.WithError(err).Fatal("Failed to create authenticated request")
		}

		client := &http.Client{}
		resp, err := client.Do(req)
		if err != nil {
			ctx.WithError(err).Fatal("Failed to get applications")
		}
		if resp.StatusCode != http.StatusOK {
			ctx.Fatalf("Failed to get applications: %s", resp.Status)
		}

		defer resp.Body.Close()
		decoder := json.NewDecoder(resp.Body)
		var apps []*app
		err = decoder.Decode(&apps)
		if err != nil {
			ctx.WithError(err).Fatal("Failed to read applications")
		}

		ctx.Infof("Found %d application(s)", len(apps))
		table := uitable.New()
		table.MaxColWidth = 70
		table.AddRow("EUI", "Name", "Owner")
		for _, app := range apps {
			table.AddRow(app.EUI, app.Name, app.Owner)
		}
		fmt.Println(table)
	},
}

var applicationsCreateCmd = &cobra.Command{
	Use:   "create [eui] [name]",
	Short: "Create a new application",
	Long:  `ttnctl applications create creates a new application.`,
	Run: func(cmd *cobra.Command, args []string) {
		if len(args) < 2 {
			cmd.Help()
			return
		}

		appEUI, err := util.Parse64(args[0])
		if err != nil {
			ctx.Fatalf("Invalid AppEUI: %s", err)
		}

		server := viper.GetString("ttn-account-server")
		values := url.Values{
			"eui":  {fmt.Sprintf("%X", appEUI)},
			"name": {args[1]},
		}
		req, err := util.NewRequestWithAuth(server, "POST", fmt.Sprintf("%s/applications", server), strings.NewReader(values.Encode()))
		if err != nil {
			ctx.WithError(err).Fatal("Failed to create authenticated request")
		}

		req.Header.Set("Content-Type", "application/x-www-form-urlencoded")

		client := &http.Client{}
		resp, err := client.Do(req)
		if err != nil {
			ctx.WithError(err).Fatal("Failed to create application")
		}
		if resp.StatusCode != http.StatusCreated {
			ctx.Fatalf("Failed to create application: %s", resp.Status)
		}

		ctx.Info("Application created successfully")

		// We need to refresh the token to add the new application to the set of
		// claims
		_, err = util.RefreshToken(server)
		if err != nil {
			log.WithError(err).Warn("Failed to refresh token. Please login")
		}
	},
}

func init() {
	RootCmd.AddCommand(applicationsCmd)
	applicationsCmd.AddCommand(applicationsCreateCmd)
}
<commit_msg>Show access keys and delete application<commit_after>// Copyright © 2016 The Things Network
// Use of this source code is governed by the MIT license that can be found in the LICENSE file.

package cmd

import (
	"encoding/json"
	"fmt"
	"net/http"
	"net/url"
	"strings"

	"github.com/TheThingsNetwork/ttn/ttnctl/util"
	"github.com/apex/log"
	"github.com/gosuri/uitable"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"
)

type app struct {
	EUI        string   `json:"eui"`
	Name       string   `json:"name"`
	Owner      string   `json:"owner"`
	AccessKeys []string `json:"accessKeys"`
	Valid      bool     `json:"valid"`
}

// applicationsCmd represents the applications command
var applicationsCmd = &cobra.Command{
	Use:   "applications",
	Short: "Show applications",
	Long:  `ttnctl applications retrieves the applications of the logged on user.`,
	Run: func(cmd *cobra.Command, args []string) {
		server := viper.GetString("ttn-account-server")
		req, err := util.NewRequestWithAuth(server, "GET", fmt.Sprintf("%s/applications", server), nil)
		if err != nil {
			ctx.WithError(err).Fatal("Failed to create authenticated request")
		}

		client := &http.Client{}
		resp, err := client.Do(req)
		if err != nil {
			ctx.WithError(err).Fatal("Failed to get applications")
		}
		if resp.StatusCode != http.StatusOK {
			ctx.Fatalf("Failed to get applications: %s", resp.Status)
		}

		defer resp.Body.Close()
		decoder := json.NewDecoder(resp.Body)
		var apps []*app
		err = decoder.Decode(&apps)
		if err != nil {
			ctx.WithError(err).Fatal("Failed to read applications")
		}

		ctx.Infof("Found %d application(s)", len(apps))
		table := uitable.New()
		table.MaxColWidth = 70
		table.AddRow("EUI", "Name", "Owner", "Access Keys", "Valid")
		for _, app := range apps {
			table.AddRow(app.EUI, app.Name, app.Owner, strings.Join(app.AccessKeys, ", "), app.Valid)
		}

		fmt.Println(table)
	},
}

var applicationsCreateCmd = &cobra.Command{
	Use:   "create [eui] [name]",
	Short: "Create a new application",
	Long:  `ttnctl applications create creates a new application.`,
	Run: func(cmd *cobra.Command, args []string) {
		if len(args) < 2 {
			cmd.Help()
			return
		}

		appEUI, err := util.Parse64(args[0])
		if err != nil {
			ctx.Fatalf("Invalid AppEUI: %s", err)
		}

		server := viper.GetString("ttn-account-server")
		values := url.Values{
			"eui":  {fmt.Sprintf("%X", appEUI)},
			"name": {args[1]},
		}
		req, err := util.NewRequestWithAuth(server, "POST", fmt.Sprintf("%s/applications", server), strings.NewReader(values.Encode()))
		if err != nil {
			ctx.WithError(err).Fatal("Failed to create authenticated request")
		}

		req.Header.Set("Content-Type", "application/x-www-form-urlencoded")

		client := &http.Client{}
		resp, err := client.Do(req)
		if err != nil {
			ctx.WithError(err).Fatal("Failed to create application")
		}
		if resp.StatusCode != http.StatusCreated {
			ctx.Fatalf("Failed to create application: %s", resp.Status)
		}

		ctx.Info("Application created successfully")

		// We need to refresh the token to add the new application to the set of
		// claims
		_, err = util.RefreshToken(server)
		if err != nil {
			log.WithError(err).Warn("Failed to refresh token. Please login")
		}
	},
}

var applicationsDeleteCmd = &cobra.Command{
	Use:   "delete [eui]",
	Short: "Delete an application",
	Long:  `ttnctl application delete deletes an existing application.`,
	Run: func(cmd *cobra.Command, args []string) {
		if len(args) < 1 {
			cmd.Help()
			return
		}

		appEUI, err := util.Parse64(args[0])
		if err != nil {
			ctx.Fatalf("Invalid AppEUI: %s", err)
		}

		server := viper.GetString("ttn-account-server")
		req, err := util.NewRequestWithAuth(server, "DELETE", fmt.Sprintf("%s/applications/%s", server, fmt.Sprintf("%X", appEUI)), nil)
		if err != nil {
			ctx.WithError(err).Fatal("Failed to create authenticated request")
		}

		client := &http.Client{}
		resp, err := client.Do(req)
		if err != nil {
			ctx.WithError(err).Fatal("Failed to delete application")
		}
		if resp.StatusCode != http.StatusOK {
			ctx.Fatalf("Failed to delete application: %s", resp.Status)
		}

		ctx.Info("Application deleted successfully")

		// We need to refresh the token to remove the application from the set of
		// claims
		_, err = util.RefreshToken(server)
		if err != nil {
			log.WithError(err).Warn("Failed to refresh token. Please login")
		}
	},
}

func init() {
	RootCmd.AddCommand(applicationsCmd)
	applicationsCmd.AddCommand(applicationsCreateCmd)
	applicationsCmd.AddCommand(applicationsDeleteCmd)
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"

	"github.com/gpahal/shlex"
)

var goPath = ""
var execCmd = ""
var envPath = ".gobu"
var globalVersion = "1.7"
var onlinePath = "https://storage.googleapis.com/golang/go%s.%s-%s.%s"

func userHomeDir() string {
	if runtime.GOOS == "windows" {
		home := os.Getenv("HOMEDRIVE") + os.Getenv("HOMEPATH")
		if home == "" {
			home = os.Getenv("USERPROFILE")
		}
		return home
	}
	return os.Getenv("HOME")
}

func init() {
	p, err := os.Getwd()
	if err != nil {
		log.Fatal(err)
	}
	goPath = p
	envPath = filepath.Join(userHomeDir(), envPath)

	flag.StringVar(&envPath, "env_path", envPath, "Location of GO instalation")
	flag.StringVar(&globalVersion, "version", globalVersion, "Version of Golang you wish to use")
	flag.StringVar(&goPath, "GOPATH", goPath, "Overide GOPATH")
	flag.StringVar(&execCmd, "exec", "", "Run command instead of default shell")
}

func createStore(version string) {
	os.MkdirAll(filepath.Join(envPath, version), 0755)
	os.MkdirAll(filepath.Join(envPath, version, "local"), 0755)
}

func runShell(version string) {
	if os.Getenv("GOBU") != "" {
		log.Fatalln("Already in boostraped env!")
	}
	os.Setenv("GOBU", "1")

	log.Println(">> You are now in a new GOBU shell. To exit, type 'exit'")
	defaultShell := resolveBinary(os.Getenv("SHELL"))

	run(version, defaultShell, []string{defaultShell})
	log.Println("Exited gobu shell")
}

func resolveBinary(bin string) string {
	if !filepath.IsAbs(bin) {
		rbin, err := exec.LookPath(bin)
		if err != nil {
			log.Panic(err)
		}
		return rbin
	}
	return bin
}

func run(version, cmd string, cmdArgs []string) *os.ProcessState {

	path := os.Getenv("PATH")
	cwd, err := os.Getwd()
	if err != nil {
		log.Panic(err)
	}
	goroot := filepath.Join(envPath, version, "go")

	os.Setenv("GO15VENDOREXPERIMENT", "1")
	os.Setenv("GOPATH", goPath)
	log.Println("GOPATH", goPath)
	os.Setenv("GOROOT", goroot)
	log.Println("GOROOT", goroot)
	// Sometimes we want to use tools from local install
	os.Setenv("PATH", goPath+"/bin:"+goroot+"/bin:"+path)

	pa := os.ProcAttr{
		Files: []*os.File{os.Stdin, os.Stdout, os.Stderr},
		Dir:   cwd,
	}

	proc, err := os.StartProcess(cmd, cmdArgs, &pa)
	if err != nil {
		log.Panic(err)
	}

	state, err := proc.Wait()
	if err != nil {
		log.Panic(err)
	}
	return state
}

func main() {
	flag.Parse()
	arch := runtime.GOARCH
	// Fix for special case of arm version naming
	if arch == "arm" {
		arch = "armv6l"
	}

	extension := ".tar.gz"

	if runtime.GOOS == "windows" {
		extension = ".zip"
	}

	onlinePath = fmt.Sprintf(onlinePath, globalVersion, runtime.GOOS, arch, extension)

	createStore(globalVersion)
	download(globalVersion)
	if execCmd != "" {
		execCmdParsed, err := shlex.Split(execCmd)
		if err != nil {
			log.Panic(err)
		}
		bin := resolveBinary(execCmdParsed[0])
		run(globalVersion, bin, execCmdParsed[:1])
	} else {
		runShell(globalVersion)
	}

}
<commit_msg>fixed extensions<commit_after>package main

import (
	"flag"
	"fmt"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"

	"github.com/gpahal/shlex"
)

var goPath = ""
var execCmd = ""
var envPath = ".gobu"
var globalVersion = "1.7"
var onlinePath = "https://storage.googleapis.com/golang/go%s.%s-%s.%s"

func userHomeDir() string {
	if runtime.GOOS == "windows" {
		home := os.Getenv("HOMEDRIVE") + os.Getenv("HOMEPATH")
		if home == "" {
			home = os.Getenv("USERPROFILE")
		}
		return home
	}
	return os.Getenv("HOME")
}

func init() {
	p, err := os.Getwd()
	if err != nil {
		log.Fatal(err)
	}
	goPath = p
	envPath = filepath.Join(userHomeDir(), envPath)

	flag.StringVar(&envPath, "env_path", envPath, "Location of GO instalation")
	flag.StringVar(&globalVersion, "version", globalVersion, "Version of Golang you wish to use")
	flag.StringVar(&goPath, "GOPATH", goPath, "Overide GOPATH")
	flag.StringVar(&execCmd, "exec", "", "Run command instead of default shell")
}

func createStore(version string) {
	os.MkdirAll(filepath.Join(envPath, version), 0755)
	os.MkdirAll(filepath.Join(envPath, version, "local"), 0755)
}

func runShell(version string) {
	if os.Getenv("GOBU") != "" {
		log.Fatalln("Already in boostraped env!")
	}
	os.Setenv("GOBU", "1")

	log.Println(">> You are now in a new GOBU shell. To exit, type 'exit'")
	defaultShell := resolveBinary(os.Getenv("SHELL"))

	run(version, defaultShell, []string{defaultShell})
	log.Println("Exited gobu shell")
}

func resolveBinary(bin string) string {
	if !filepath.IsAbs(bin) {
		rbin, err := exec.LookPath(bin)
		if err != nil {
			log.Panic(err)
		}
		return rbin
	}
	return bin
}

func run(version, cmd string, cmdArgs []string) *os.ProcessState {

	path := os.Getenv("PATH")
	cwd, err := os.Getwd()
	if err != nil {
		log.Panic(err)
	}
	goroot := filepath.Join(envPath, version, "go")

	os.Setenv("GO15VENDOREXPERIMENT", "1")
	os.Setenv("GOPATH", goPath)
	log.Println("GOPATH", goPath)
	os.Setenv("GOROOT", goroot)
	log.Println("GOROOT", goroot)
	// Sometimes we want to use tools from local install
	os.Setenv("PATH", goPath+"/bin:"+goroot+"/bin:"+path)

	pa := os.ProcAttr{
		Files: []*os.File{os.Stdin, os.Stdout, os.Stderr},
		Dir:   cwd,
	}

	proc, err := os.StartProcess(cmd, cmdArgs, &pa)
	if err != nil {
		log.Panic(err)
	}

	state, err := proc.Wait()
	if err != nil {
		log.Panic(err)
	}
	return state
}

func main() {
	flag.Parse()
	arch := runtime.GOARCH
	// Fix for special case of arm version naming
	if arch == "arm" {
		arch = "armv6l"
	}

	extension := "tar.gz"

	if runtime.GOOS == "windows" {
		extension = "zip"
	}

	onlinePath = fmt.Sprintf(onlinePath, globalVersion, runtime.GOOS, arch, extension)

	createStore(globalVersion)
	download(globalVersion)
	if execCmd != "" {
		execCmdParsed, err := shlex.Split(execCmd)
		if err != nil {
			log.Panic(err)
		}
		bin := resolveBinary(execCmdParsed[0])
		run(globalVersion, bin, execCmdParsed[:1])
	} else {
		runShell(globalVersion)
	}

}
<|endoftext|>


<commit_before>package deploy

import (
	"io/ioutil"
	"os"
	"path"
	"sort"
	"strings"
	"testing"

	"encoding/json"

	"github.com/coreos/alm/pkg/controller/registry"
	"github.com/coreos/go-semver/semver"
	"github.com/ghodss/yaml"
	"github.com/stretchr/testify/require"
	"k8s.io/api/core/v1"
)

const manifestDir = "./tectonic-alm-operator/manifests"

// BySemverDir lets us sort os.FileInfo by interpreting the filename as a semver version,
// which is how manifest directories are stored
type BySemverDir []os.FileInfo

func (s BySemverDir) Len() int      { return len(s) }
func (s BySemverDir) Swap(i, j int) { s[i], s[j] = s[j], s[i] }
func (s BySemverDir) Less(i, j int) bool {
	semverA := semver.New(s[i].Name())
	semverB := semver.New(s[j].Name())
	return semverA.LessThan(*semverB)
}

// loadCatalogFromFile loads an in memory catalog from a file path. Only used for testing.
func loadCatalogFromFile(path string) (*registry.InMem, error) {
	loader := registry.ConfigMapCatalogResourceLoader{
		Catalog: registry.NewInMem(),
	}
	currentBytes, err := ioutil.ReadFile(path)
	if err != nil {
		return nil, err
	}
	currentJsonBytes, err := yaml.YAMLToJSON(currentBytes)
	if err != nil {
		return nil, err
	}
	var currentConfigMap v1.ConfigMap
	err = json.Unmarshal(currentJsonBytes, &currentConfigMap)
	if err != nil {
		return nil, err
	}
	err = loader.LoadCatalogResourcesFromConfigMap(&currentConfigMap)
	if err != nil {
		return nil, err
	}
	return loader.Catalog, nil
}

func TestCatalogVersions(t *testing.T) {
	// for each version of the catalog, load (version-1) and verify that each OCS that has a replaces field
	// points to an OCS in the previous version of the catalog
	files, err := ioutil.ReadDir(manifestDir)
	require.NoError(t, err)

	versionDirs := []os.FileInfo{}
	for _, f := range files {
		if f.IsDir() {
			versionDirs = append(versionDirs, f)
		}
	}

	// sort manifest directories by semver
	sort.Sort(BySemverDir(versionDirs))

	// versions before this don't contain the catalog
	oldestVersion := semver.New("0.2.1")

	// collect catalog configmaps
	catalogConfigMaps := []string{}
	for _, versioned := range versionDirs {

		// ignore old versions that don't have the catalog
		semverDirName := semver.New(versioned.Name())
		if semverDirName.LessThan(*oldestVersion) {
			continue
		}

		// get the path of each version of the catalog
		manifestFiles, err := ioutil.ReadDir(path.Join(manifestDir, versioned.Name()))
		require.NoError(t, err)
		for _, f := range manifestFiles {
			if strings.HasSuffix(f.Name(), "configmap.yaml") {
				catalogConfigMaps = append(catalogConfigMaps, path.Join(manifestDir, versioned.Name(), f.Name()))
			}
		}
	}

	// ensure services in <version> that have a `replaces` field replace something in <version - 1>
	for i := 0; i < len(catalogConfigMaps)-1; i++ {
		t.Logf("comparing %s and %s", catalogConfigMaps[i], catalogConfigMaps[i+1])

		currentCatalog, err := loadCatalogFromFile(catalogConfigMaps[i])
		require.NoError(t, err)

		nextCatalog, err := loadCatalogFromFile(catalogConfigMaps[i+1])
		require.NoError(t, err)

		nextServices, err := nextCatalog.ListServices()
		require.NoError(t, err)
		for _, csv := range nextServices {
			if csv.Spec.Replaces != "" {
				oldCSV, err := currentCatalog.FindCSVByName(csv.Spec.Replaces)
				require.NoError(t, err)
				require.NotNil(t, oldCSV)
			}
		}
	}
}
<commit_msg>update catalog version test to handle multiple catalogs<commit_after>package deploy

import (
	"io/ioutil"
	"os"
	"path"
	"sort"
	"strings"
	"testing"

	"encoding/json"

	"github.com/coreos/alm/pkg/controller/registry"
	"github.com/coreos/go-semver/semver"
	"github.com/ghodss/yaml"
	"github.com/stretchr/testify/require"
	"k8s.io/api/core/v1"
)

const manifestDir = "./tectonic-alm-operator/manifests"

// BySemverDir lets us sort os.FileInfo by interpreting the filename as a semver version,
// which is how manifest directories are stored
type BySemverDir []os.FileInfo

func (s BySemverDir) Len() int      { return len(s) }
func (s BySemverDir) Swap(i, j int) { s[i], s[j] = s[j], s[i] }
func (s BySemverDir) Less(i, j int) bool {
	semverA := semver.New(s[i].Name())
	semverB := semver.New(s[j].Name())
	return semverA.LessThan(*semverB)
}

// LoadedCatalog wraps an in mem catalog with version metadata
type LoadedCatalog struct {
	Registry *registry.InMem
	Name     string
	Version  string
}

// loadCatalogFromFile loads an in memory catalog from a file path. Only used for testing.
func loadCatalogFromFile(path string) (*LoadedCatalog, error) {
	loader := registry.ConfigMapCatalogResourceLoader{
		Catalog: registry.NewInMem(),
	}
	currentBytes, err := ioutil.ReadFile(path)
	if err != nil {
		return nil, err
	}
	currentJsonBytes, err := yaml.YAMLToJSON(currentBytes)
	if err != nil {
		return nil, err
	}
	var currentConfigMap v1.ConfigMap
	err = json.Unmarshal(currentJsonBytes, &currentConfigMap)
	if err != nil {
		return nil, err
	}
	err = loader.LoadCatalogResourcesFromConfigMap(&currentConfigMap)
	if err != nil {
		return nil, err
	}
	return &LoadedCatalog{
		Registry: loader.Catalog,
		Name:     currentConfigMap.Name,
	}, nil
}

func TestCatalogVersions(t *testing.T) {
	// for each version of the catalog, load (version-1) and verify that each OCS that has a replaces field
	// points to an OCS in the previous version of the catalog
	files, err := ioutil.ReadDir(manifestDir)
	require.NoError(t, err)

	versionDirs := []os.FileInfo{}
	for _, f := range files {
		if f.IsDir() {
			versionDirs = append(versionDirs, f)
		}
	}

	// sort manifest directories by semver
	sort.Sort(BySemverDir(versionDirs))

	// versions before this don't contain the catalog
	oldestVersion := semver.New("0.2.1")

	// load all available catalogs
	catalogNameVersions := map[string][]*LoadedCatalog{}
	for _, versioned := range versionDirs {

		// ignore old versions that don't have the catalog
		semverDirName := semver.New(versioned.Name())
		if semverDirName.LessThan(*oldestVersion) {
			continue
		}

		// get the path of each version of the catalog
		manifestFiles, err := ioutil.ReadDir(path.Join(manifestDir, versioned.Name()))
		require.NoError(t, err)
		for _, f := range manifestFiles {
			if strings.HasSuffix(f.Name(), "configmap.yaml") {
				loadedCatalog, err := loadCatalogFromFile(path.Join(manifestDir, versioned.Name(), f.Name()))
				require.NoError(t, err)
				loadedCatalog.Version = semverDirName.String()
				if _, ok := catalogNameVersions[loadedCatalog.Name]; !ok {
					catalogNameVersions[loadedCatalog.Name] = []*LoadedCatalog{}
				}
				catalogNameVersions[loadedCatalog.Name] = append(catalogNameVersions[loadedCatalog.Name], loadedCatalog)
			}
		}
	}

	// ensure services in <version> that have a `replaces` field replace something in <version - 1>
	for catalogName, catalogVersions := range catalogNameVersions {
		for i := 0; i < len(catalogVersions)-1; i++ {
			currentCatalog := catalogVersions[i]
			nextCatalog := catalogVersions[i+1]

			t.Logf("comparing %s version %s to %s", catalogName, currentCatalog.Version, nextCatalog.Version)

			nextServices, err := nextCatalog.Registry.ListServices()
			require.NoError(t, err)
			for _, csv := range nextServices {
				if csv.Spec.Replaces != "" {
					oldCSV, err := currentCatalog.Registry.FindCSVByName(csv.Spec.Replaces)
					require.NoError(t, err)
					require.NotNil(t, oldCSV)
				}
			}
		}
	}
}
<|endoftext|>


<commit_before>// Copyright 2016 tsuru-client authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package admin

import (
	"encoding/json"
	"fmt"
	"net/http"
	"net/url"
	"strconv"
	"strings"

	"github.com/ajg/form"
	"github.com/tsuru/gnuflag"
	"github.com/tsuru/tsuru/cmd"
	"github.com/tsuru/tsuru/errors"
	"github.com/tsuru/tsuru/provision"
)

type AddPoolToSchedulerCmd struct {
	public       bool
	defaultPool  bool
	forceDefault bool
	provisioner  string
	fs           *gnuflag.FlagSet
}

func (AddPoolToSchedulerCmd) Info() *cmd.Info {
	return &cmd.Info{
		Name:  "pool-add",
		Usage: "pool-add <pool> [-p/--public] [-d/--default] [--provisioner <name>] [-f/--force]",
		Desc: `Adds a new pool.

Each docker node added using [[node-add]] command belongs to one pool.
Also, when creating a new application a pool must be chosen and this means
that all units of the created application will be spawned in nodes belonging
to the chosen pool.`,
		MinArgs: 1,
	}
}

func (c *AddPoolToSchedulerCmd) Flags() *gnuflag.FlagSet {
	if c.fs == nil {
		c.fs = gnuflag.NewFlagSet("", gnuflag.ExitOnError)
		msg := "Make pool public (all teams can use it)"
		c.fs.BoolVar(&c.public, "public", false, msg)
		c.fs.BoolVar(&c.public, "p", false, msg)
		msg = "Make pool default (when none is specified during [[app-create]] this pool will be used)"
		c.fs.BoolVar(&c.defaultPool, "default", false, msg)
		c.fs.BoolVar(&c.defaultPool, "d", false, msg)
		msg = "Force overwrite default pool"
		c.fs.BoolVar(&c.forceDefault, "force", false, msg)
		c.fs.BoolVar(&c.forceDefault, "f", false, msg)
		msg = "Provisioner associated to the pool (empty for default docker provisioner)"
		c.fs.StringVar(&c.provisioner, "provisioner", "", msg)
	}
	return c.fs
}

func (c *AddPoolToSchedulerCmd) Run(ctx *cmd.Context, client *cmd.Client) error {
	v := url.Values{}
	v.Set("name", ctx.Args[0])
	v.Set("public", strconv.FormatBool(c.public))
	v.Set("default", strconv.FormatBool(c.defaultPool))
	v.Set("force", strconv.FormatBool(c.forceDefault))
	v.Set("provisioner", c.provisioner)
	u, err := cmd.GetURL("/pools")
	err = doRequest(client, u, "POST", v.Encode())
	if err != nil {
		if e, ok := err.(*errors.HTTP); ok && e.Code == http.StatusPreconditionFailed {
			retryMessage := "WARNING: Default pool already exist. Do you want change to %s pool? (y/n) "
			v.Set("force", "true")
			url, _ := cmd.GetURL("/pools")
			successMessage := "Pool successfully registered.\n"
			failMessage := "Pool add aborted.\n"
			return confirmAction(ctx, client, url, "POST", v.Encode(), retryMessage, failMessage, successMessage)
		}
		return err
	}
	ctx.Stdout.Write([]byte("Pool successfully registered.\n"))
	return nil
}

type UpdatePoolToSchedulerCmd struct {
	public       pointerBoolFlag
	defaultPool  pointerBoolFlag
	forceDefault bool
	fs           *gnuflag.FlagSet
}

func (UpdatePoolToSchedulerCmd) Info() *cmd.Info {
	return &cmd.Info{
		Name:    "pool-update",
		Usage:   "pool-update <pool> [--public=true/false] [--default=true/false] [-f/--force]",
		Desc:    `Updates attributes for a pool.`,
		MinArgs: 1,
	}
}

func (c *UpdatePoolToSchedulerCmd) Flags() *gnuflag.FlagSet {
	if c.fs == nil {
		c.fs = gnuflag.NewFlagSet("", gnuflag.ExitOnError)
		msg := "Make pool public (all teams can use it)"
		c.fs.Var(&c.public, "public", msg)
		msg = "Make pool default (when none is specified during [[app-create]] this pool will be used)"
		c.fs.Var(&c.defaultPool, "default", msg)
		c.fs.BoolVar(&c.forceDefault, "force", false, "Force pool to be default.")
		c.fs.BoolVar(&c.forceDefault, "f", false, "Force pool to be default.")
	}
	return c.fs
}

func (c *UpdatePoolToSchedulerCmd) Run(ctx *cmd.Context, client *cmd.Client) error {
	v := url.Values{}
	if c.public.value == nil {
		v.Set("public", "")
	} else {
		v.Set("public", strconv.FormatBool(*c.public.value))
	}
	if c.defaultPool.value == nil {
		v.Set("default", "")
	} else {
		v.Set("default", strconv.FormatBool(*c.defaultPool.value))
	}
	v.Set("force", strconv.FormatBool(c.forceDefault))
	u, err := cmd.GetURL(fmt.Sprintf("/pools/%s", ctx.Args[0]))
	err = doRequest(client, u, "PUT", v.Encode())
	if err != nil {
		if e, ok := err.(*errors.HTTP); ok && e.Code == http.StatusPreconditionFailed {
			retryMessage := "WARNING: Default pool already exist. Do you want change to %s pool? (y/n) "
			failMessage := "Pool update aborted.\n"
			successMessage := "Pool successfully updated.\n"
			v.Set("force", "true")
			u, err = cmd.GetURL(fmt.Sprintf("/pools/%s", ctx.Args[0]))
			return confirmAction(ctx, client, u, "PUT", v.Encode(), retryMessage, failMessage, successMessage)
		}
		return err
	}
	ctx.Stdout.Write([]byte("Pool successfully updated.\n"))
	return nil
}

type RemovePoolFromSchedulerCmd struct {
	cmd.ConfirmationCommand
}

func (c *RemovePoolFromSchedulerCmd) Info() *cmd.Info {
	return &cmd.Info{
		Name:    "pool-remove",
		Usage:   "pool-remove <pool> [-y]",
		Desc:    "Remove an existing pool.",
		MinArgs: 1,
	}
}

func (c *RemovePoolFromSchedulerCmd) Run(ctx *cmd.Context, client *cmd.Client) error {
	if !c.Confirm(ctx, fmt.Sprintf("Are you sure you want to remove \"%s\" pool?", ctx.Args[0])) {
		return nil
	}
	url, err := cmd.GetURL(fmt.Sprintf("/pools/%s", ctx.Args[0]))
	if err != nil {
		return err
	}
	req, err := http.NewRequest(http.MethodDelete, url, nil)
	if err != nil {
		return err
	}
	_, err = client.Do(req)
	if err != nil {
		return err
	}
	ctx.Stdout.Write([]byte("Pool successfully removed.\n"))
	return nil
}

type AddTeamsToPoolCmd struct{}

func (AddTeamsToPoolCmd) Info() *cmd.Info {
	return &cmd.Info{
		Name:  "pool-teams-add",
		Usage: "pool-teams-add <pool> <teams>...",
		Desc: `Adds teams to a pool. This will make the specified pool available when
creating a new application for one of the added teams.`,
		MinArgs: 2,
	}
}

func (AddTeamsToPoolCmd) Run(ctx *cmd.Context, client *cmd.Client) error {
	v := url.Values{}
	for _, team := range ctx.Args[1:] {
		v.Add("team", team)
	}
	u, err := cmd.GetURL(fmt.Sprintf("/pools/%s/team", ctx.Args[0]))
	if err != nil {
		return err
	}
	req, err := http.NewRequest("POST", u, strings.NewReader(v.Encode()))
	if err != nil {
		return err
	}
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	_, err = client.Do(req)
	if err != nil {
		return err
	}
	ctx.Stdout.Write([]byte("Teams successfully registered.\n"))
	return nil
}

type RemoveTeamsFromPoolCmd struct{}

func (RemoveTeamsFromPoolCmd) Info() *cmd.Info {
	return &cmd.Info{
		Name:  "pool-teams-remove",
		Usage: "pool-teams-remove <pool> <teams>...",
		Desc: `Removes teams from a pool. Listed teams will be no longer able to use this
pool when creating a new application.`,
		MinArgs: 2,
	}
}

func (RemoveTeamsFromPoolCmd) Run(ctx *cmd.Context, client *cmd.Client) error {
	v := url.Values{}
	for _, team := range ctx.Args[1:] {
		v.Add("team", team)
	}
	u, err := cmd.GetURL(fmt.Sprintf("/pools/%s/team?%s", ctx.Args[0], v.Encode()))
	if err != nil {
		return err
	}
	req, err := http.NewRequest(http.MethodDelete, u, nil)
	if err != nil {
		return err
	}
	_, err = client.Do(req)
	if err != nil {
		return err
	}
	ctx.Stdout.Write([]byte("Teams successfully removed.\n"))
	return nil
}

type pointerBoolFlag struct {
	value *bool
}

func (p *pointerBoolFlag) String() string {
	if p.value == nil {
		return "not set"
	}
	return fmt.Sprintf("%v", *p.value)
}

func (p *pointerBoolFlag) Set(value string) error {
	if value == "" {
		return nil
	}
	v, err := strconv.ParseBool(value)
	if err != nil {
		return err
	}
	p.value = &v
	return nil
}

func doRequest(client *cmd.Client, url, method, body string) error {
	req, err := http.NewRequest(method, url, strings.NewReader(body))
	if err != nil {
		return err
	}
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	_, err = client.Do(req)
	return err
}

func confirmAction(ctx *cmd.Context, client *cmd.Client, url, method, body string, retryMessage, failMessage, successMessage string) error {
	var answer string
	fmt.Fprintf(ctx.Stdout, retryMessage, ctx.Args[0])
	fmt.Fscanf(ctx.Stdin, "%s", &answer)
	if answer == "y" || answer == "yes" {
		err := doRequest(client, url, method, body)
		if err != nil {
			return err
		}
		ctx.Stdout.Write([]byte(successMessage))
		return nil

	}
	ctx.Stdout.Write([]byte(failMessage))
	return nil
}

type PoolConstraintList struct{}

func (c *PoolConstraintList) Info() *cmd.Info {
	return &cmd.Info{
		Name:  "pool-constraint-list",
		Usage: "pool-constraint-list",
		Desc:  "List pool constraints.",
	}
}

func (c *PoolConstraintList) Run(ctx *cmd.Context, client *cmd.Client) error {
	url, err := cmd.GetURLVersion("1.3", "/constraints")
	if err != nil {
		return err
	}
	req, err := http.NewRequest(http.MethodGet, url, nil)
	if err != nil {
		return err
	}
	resp, err := client.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	var constraints []provision.PoolConstraint
	err = json.NewDecoder(resp.Body).Decode(&constraints)
	if err != nil {
		return err
	}
	tbl := cmd.NewTable()
	tbl.Headers = cmd.Row{"Pool Expression", "Field", "Values", "Blacklist"}
	for _, c := range constraints {
		tbl.AddRow(cmd.Row{c.PoolExpr, c.Field, strings.Join(c.Values, ","), strconv.FormatBool(c.Blacklist)})
	}
	tbl.SortByColumn(0)
	ctx.Stdout.Write([]byte(tbl.String()))
	return nil
}

type PoolConstraintSet struct {
	append    bool
	blacklist bool
	fs        *gnuflag.FlagSet
}

func (c *PoolConstraintSet) Flags() *gnuflag.FlagSet {
	if c.fs == nil {
		c.fs = gnuflag.NewFlagSet("", gnuflag.ExitOnError)
		c.fs.BoolVar(&c.append, "append", false, "Append to existing constraint.")
		c.fs.BoolVar(&c.append, "a", false, "Append to existing constraint.")
		c.fs.BoolVar(&c.blacklist, "b", false, "Blacklist constraint.")
		c.fs.BoolVar(&c.blacklist, "blacklist", false, "Blacklist constraint.")
	}
	return c.fs
}
func (c *PoolConstraintSet) Info() *cmd.Info {
	return &cmd.Info{
		Name:  "pool-constraint-set",
		Usage: "pool-constraint-set <poolExpression> <field> [<values>]... [-b/--blacklist] [-a/--append]",
		Desc: `Set a constraint on a pool expression.

Examples:

[[tsuru pool-constraint-set dev_pool team "*" # allows every team to use the pool "dev_pool" ]]
[[tsuru pool-constraint-set "dev_*" router prod_router --blacklist # disallows "prod_router" to be used on every pool with "dev_" prefix ]]
[[tsuru pool-constraint-set prod_pool team team2 team3 --append # adds "team2" and "team3" to the list of teams allowed to use pool "prod_pool"]]`,
		MinArgs: 2,
	}
}

func (c *PoolConstraintSet) Run(ctx *cmd.Context, client *cmd.Client) error {
	u, err := cmd.GetURLVersion("1.3", "/constraints")
	if err != nil {
		return err
	}
	values := ctx.Args[2:]
	var allValues []string
	for _, v := range values {
		allValues = append(allValues, strings.Split(v, ",")...)
	}
	constraint := provision.PoolConstraint{
		PoolExpr:  ctx.Args[0],
		Field:     ctx.Args[1],
		Blacklist: c.blacklist,
		Values:    allValues,
	}
	v, err := form.EncodeToValues(constraint)
	if err != nil {
		return err
	}
	v.Set("append", strconv.FormatBool(c.append))
	err = doRequest(client, u, http.MethodPut, v.Encode())
	if err != nil {
		return err
	}
	ctx.Stdout.Write([]byte("Constraint successfully set.\n"))
	return nil
}
<commit_msg>pool: docs to allow restricting services to certain pools<commit_after>// Copyright 2016 tsuru-client authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package admin

import (
	"encoding/json"
	"fmt"
	"net/http"
	"net/url"
	"strconv"
	"strings"

	"github.com/ajg/form"
	"github.com/tsuru/gnuflag"
	"github.com/tsuru/tsuru/cmd"
	"github.com/tsuru/tsuru/errors"
	"github.com/tsuru/tsuru/provision"
)

type AddPoolToSchedulerCmd struct {
	public       bool
	defaultPool  bool
	forceDefault bool
	provisioner  string
	fs           *gnuflag.FlagSet
}

func (AddPoolToSchedulerCmd) Info() *cmd.Info {
	return &cmd.Info{
		Name:  "pool-add",
		Usage: "pool-add <pool> [-p/--public] [-d/--default] [--provisioner <name>] [-f/--force]",
		Desc: `Adds a new pool.

Each docker node added using [[node-add]] command belongs to one pool.
Also, when creating a new application a pool must be chosen and this means
that all units of the created application will be spawned in nodes belonging
to the chosen pool.`,
		MinArgs: 1,
	}
}

func (c *AddPoolToSchedulerCmd) Flags() *gnuflag.FlagSet {
	if c.fs == nil {
		c.fs = gnuflag.NewFlagSet("", gnuflag.ExitOnError)
		msg := "Make pool public (all teams can use it)"
		c.fs.BoolVar(&c.public, "public", false, msg)
		c.fs.BoolVar(&c.public, "p", false, msg)
		msg = "Make pool default (when none is specified during [[app-create]] this pool will be used)"
		c.fs.BoolVar(&c.defaultPool, "default", false, msg)
		c.fs.BoolVar(&c.defaultPool, "d", false, msg)
		msg = "Force overwrite default pool"
		c.fs.BoolVar(&c.forceDefault, "force", false, msg)
		c.fs.BoolVar(&c.forceDefault, "f", false, msg)
		msg = "Provisioner associated to the pool (empty for default docker provisioner)"
		c.fs.StringVar(&c.provisioner, "provisioner", "", msg)
	}
	return c.fs
}

func (c *AddPoolToSchedulerCmd) Run(ctx *cmd.Context, client *cmd.Client) error {
	v := url.Values{}
	v.Set("name", ctx.Args[0])
	v.Set("public", strconv.FormatBool(c.public))
	v.Set("default", strconv.FormatBool(c.defaultPool))
	v.Set("force", strconv.FormatBool(c.forceDefault))
	v.Set("provisioner", c.provisioner)
	u, err := cmd.GetURL("/pools")
	err = doRequest(client, u, "POST", v.Encode())
	if err != nil {
		if e, ok := err.(*errors.HTTP); ok && e.Code == http.StatusPreconditionFailed {
			retryMessage := "WARNING: Default pool already exist. Do you want change to %s pool? (y/n) "
			v.Set("force", "true")
			url, _ := cmd.GetURL("/pools")
			successMessage := "Pool successfully registered.\n"
			failMessage := "Pool add aborted.\n"
			return confirmAction(ctx, client, url, "POST", v.Encode(), retryMessage, failMessage, successMessage)
		}
		return err
	}
	ctx.Stdout.Write([]byte("Pool successfully registered.\n"))
	return nil
}

type UpdatePoolToSchedulerCmd struct {
	public       pointerBoolFlag
	defaultPool  pointerBoolFlag
	forceDefault bool
	fs           *gnuflag.FlagSet
}

func (UpdatePoolToSchedulerCmd) Info() *cmd.Info {
	return &cmd.Info{
		Name:    "pool-update",
		Usage:   "pool-update <pool> [--public=true/false] [--default=true/false] [-f/--force]",
		Desc:    `Updates attributes for a pool.`,
		MinArgs: 1,
	}
}

func (c *UpdatePoolToSchedulerCmd) Flags() *gnuflag.FlagSet {
	if c.fs == nil {
		c.fs = gnuflag.NewFlagSet("", gnuflag.ExitOnError)
		msg := "Make pool public (all teams can use it)"
		c.fs.Var(&c.public, "public", msg)
		msg = "Make pool default (when none is specified during [[app-create]] this pool will be used)"
		c.fs.Var(&c.defaultPool, "default", msg)
		c.fs.BoolVar(&c.forceDefault, "force", false, "Force pool to be default.")
		c.fs.BoolVar(&c.forceDefault, "f", false, "Force pool to be default.")
	}
	return c.fs
}

func (c *UpdatePoolToSchedulerCmd) Run(ctx *cmd.Context, client *cmd.Client) error {
	v := url.Values{}
	if c.public.value == nil {
		v.Set("public", "")
	} else {
		v.Set("public", strconv.FormatBool(*c.public.value))
	}
	if c.defaultPool.value == nil {
		v.Set("default", "")
	} else {
		v.Set("default", strconv.FormatBool(*c.defaultPool.value))
	}
	v.Set("force", strconv.FormatBool(c.forceDefault))
	u, err := cmd.GetURL(fmt.Sprintf("/pools/%s", ctx.Args[0]))
	err = doRequest(client, u, "PUT", v.Encode())
	if err != nil {
		if e, ok := err.(*errors.HTTP); ok && e.Code == http.StatusPreconditionFailed {
			retryMessage := "WARNING: Default pool already exist. Do you want change to %s pool? (y/n) "
			failMessage := "Pool update aborted.\n"
			successMessage := "Pool successfully updated.\n"
			v.Set("force", "true")
			u, err = cmd.GetURL(fmt.Sprintf("/pools/%s", ctx.Args[0]))
			return confirmAction(ctx, client, u, "PUT", v.Encode(), retryMessage, failMessage, successMessage)
		}
		return err
	}
	ctx.Stdout.Write([]byte("Pool successfully updated.\n"))
	return nil
}

type RemovePoolFromSchedulerCmd struct {
	cmd.ConfirmationCommand
}

func (c *RemovePoolFromSchedulerCmd) Info() *cmd.Info {
	return &cmd.Info{
		Name:    "pool-remove",
		Usage:   "pool-remove <pool> [-y]",
		Desc:    "Remove an existing pool.",
		MinArgs: 1,
	}
}

func (c *RemovePoolFromSchedulerCmd) Run(ctx *cmd.Context, client *cmd.Client) error {
	if !c.Confirm(ctx, fmt.Sprintf("Are you sure you want to remove \"%s\" pool?", ctx.Args[0])) {
		return nil
	}
	url, err := cmd.GetURL(fmt.Sprintf("/pools/%s", ctx.Args[0]))
	if err != nil {
		return err
	}
	req, err := http.NewRequest(http.MethodDelete, url, nil)
	if err != nil {
		return err
	}
	_, err = client.Do(req)
	if err != nil {
		return err
	}
	ctx.Stdout.Write([]byte("Pool successfully removed.\n"))
	return nil
}

type AddTeamsToPoolCmd struct{}

func (AddTeamsToPoolCmd) Info() *cmd.Info {
	return &cmd.Info{
		Name:  "pool-teams-add",
		Usage: "pool-teams-add <pool> <teams>...",
		Desc: `Adds teams to a pool. This will make the specified pool available when
creating a new application for one of the added teams.`,
		MinArgs: 2,
	}
}

func (AddTeamsToPoolCmd) Run(ctx *cmd.Context, client *cmd.Client) error {
	v := url.Values{}
	for _, team := range ctx.Args[1:] {
		v.Add("team", team)
	}
	u, err := cmd.GetURL(fmt.Sprintf("/pools/%s/team", ctx.Args[0]))
	if err != nil {
		return err
	}
	req, err := http.NewRequest("POST", u, strings.NewReader(v.Encode()))
	if err != nil {
		return err
	}
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	_, err = client.Do(req)
	if err != nil {
		return err
	}
	ctx.Stdout.Write([]byte("Teams successfully registered.\n"))
	return nil
}

type RemoveTeamsFromPoolCmd struct{}

func (RemoveTeamsFromPoolCmd) Info() *cmd.Info {
	return &cmd.Info{
		Name:  "pool-teams-remove",
		Usage: "pool-teams-remove <pool> <teams>...",
		Desc: `Removes teams from a pool. Listed teams will be no longer able to use this
pool when creating a new application.`,
		MinArgs: 2,
	}
}

func (RemoveTeamsFromPoolCmd) Run(ctx *cmd.Context, client *cmd.Client) error {
	v := url.Values{}
	for _, team := range ctx.Args[1:] {
		v.Add("team", team)
	}
	u, err := cmd.GetURL(fmt.Sprintf("/pools/%s/team?%s", ctx.Args[0], v.Encode()))
	if err != nil {
		return err
	}
	req, err := http.NewRequest(http.MethodDelete, u, nil)
	if err != nil {
		return err
	}
	_, err = client.Do(req)
	if err != nil {
		return err
	}
	ctx.Stdout.Write([]byte("Teams successfully removed.\n"))
	return nil
}

type pointerBoolFlag struct {
	value *bool
}

func (p *pointerBoolFlag) String() string {
	if p.value == nil {
		return "not set"
	}
	return fmt.Sprintf("%v", *p.value)
}

func (p *pointerBoolFlag) Set(value string) error {
	if value == "" {
		return nil
	}
	v, err := strconv.ParseBool(value)
	if err != nil {
		return err
	}
	p.value = &v
	return nil
}

func doRequest(client *cmd.Client, url, method, body string) error {
	req, err := http.NewRequest(method, url, strings.NewReader(body))
	if err != nil {
		return err
	}
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	_, err = client.Do(req)
	return err
}

func confirmAction(ctx *cmd.Context, client *cmd.Client, url, method, body string, retryMessage, failMessage, successMessage string) error {
	var answer string
	fmt.Fprintf(ctx.Stdout, retryMessage, ctx.Args[0])
	fmt.Fscanf(ctx.Stdin, "%s", &answer)
	if answer == "y" || answer == "yes" {
		err := doRequest(client, url, method, body)
		if err != nil {
			return err
		}
		ctx.Stdout.Write([]byte(successMessage))
		return nil

	}
	ctx.Stdout.Write([]byte(failMessage))
	return nil
}

type PoolConstraintList struct{}

func (c *PoolConstraintList) Info() *cmd.Info {
	return &cmd.Info{
		Name:  "pool-constraint-list",
		Usage: "pool-constraint-list",
		Desc:  "List pool constraints.",
	}
}

func (c *PoolConstraintList) Run(ctx *cmd.Context, client *cmd.Client) error {
	url, err := cmd.GetURLVersion("1.3", "/constraints")
	if err != nil {
		return err
	}
	req, err := http.NewRequest(http.MethodGet, url, nil)
	if err != nil {
		return err
	}
	resp, err := client.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	var constraints []provision.PoolConstraint
	err = json.NewDecoder(resp.Body).Decode(&constraints)
	if err != nil {
		return err
	}
	tbl := cmd.NewTable()
	tbl.Headers = cmd.Row{"Pool Expression", "Field", "Values", "Blacklist"}
	for _, c := range constraints {
		tbl.AddRow(cmd.Row{c.PoolExpr, c.Field, strings.Join(c.Values, ","), strconv.FormatBool(c.Blacklist)})
	}
	tbl.SortByColumn(0)
	ctx.Stdout.Write([]byte(tbl.String()))
	return nil
}

type PoolConstraintSet struct {
	append    bool
	blacklist bool
	fs        *gnuflag.FlagSet
}

func (c *PoolConstraintSet) Flags() *gnuflag.FlagSet {
	if c.fs == nil {
		c.fs = gnuflag.NewFlagSet("", gnuflag.ExitOnError)
		c.fs.BoolVar(&c.append, "append", false, "Append to existing constraint.")
		c.fs.BoolVar(&c.append, "a", false, "Append to existing constraint.")
		c.fs.BoolVar(&c.blacklist, "b", false, "Blacklist constraint.")
		c.fs.BoolVar(&c.blacklist, "blacklist", false, "Blacklist constraint.")
	}
	return c.fs
}
func (c *PoolConstraintSet) Info() *cmd.Info {
	return &cmd.Info{
		Name:  "pool-constraint-set",
		Usage: "pool-constraint-set <poolExpression> <field> [<values>]... [-b/--blacklist] [-a/--append]",
		Desc: `Set a constraint on a pool expression.

Examples:

[[tsuru pool-constraint-set dev_pool team "*" # allows every team to use the pool "dev_pool" ]]
[[tsuru pool-constraint-set "dev_*" router prod_router --blacklist # disallows "prod_router" to be used on every pool with "dev_" prefix ]]
[[tsuru pool-constraint-set prod_pool team team2 team3 --append # adds "team2" and "team3" to the list of teams allowed to use pool "prod_pool"]]
[[tsuru pool-constraint-set prod_pool service service1 service2 --append # adds "service1" and "service2" to the list of services allowed to be used on pool "prod_pool"]]
[[tsuru pool-constraint-set prod_pool service service1 --blacklist # disallows "service1" to be used on pool "prod_pool"]]`,
		MinArgs: 2,
	}
}

func (c *PoolConstraintSet) Run(ctx *cmd.Context, client *cmd.Client) error {
	u, err := cmd.GetURLVersion("1.3", "/constraints")
	if err != nil {
		return err
	}
	values := ctx.Args[2:]
	var allValues []string
	for _, v := range values {
		allValues = append(allValues, strings.Split(v, ",")...)
	}
	constraint := provision.PoolConstraint{
		PoolExpr:  ctx.Args[0],
		Field:     ctx.Args[1],
		Blacklist: c.blacklist,
		Values:    allValues,
	}
	v, err := form.EncodeToValues(constraint)
	if err != nil {
		return err
	}
	v.Set("append", strconv.FormatBool(c.append))
	err = doRequest(client, u, http.MethodPut, v.Encode())
	if err != nil {
		return err
	}
	ctx.Stdout.Write([]byte("Constraint successfully set.\n"))
	return nil
}
<|endoftext|>