<commit_before>package models

import (
	"encoding/json"
	"strings"

	"testing"

	"github.com/cihangir/gene/generators/folders"
	"github.com/cihangir/gene/generators/validators"
	"github.com/cihangir/gene/testdata"
	"github.com/cihangir/schema"

	"github.com/cihangir/gene/writers"
)

func TestGenerateModel(t *testing.T) {
	var s schema.Schema
	if err := json.Unmarshal([]byte(testdata.JSON1), &s); err != nil {
		t.Fatal(err.Error())
	}

	model, err := GenerateModel(&s)
	if err != nil {
		t.Fatal(err.Error())
	}

	folders.EnsureFolders("/tmp/", folders.FolderStucture)
	fileName := "/tmp/models/" + s.Title + ".go"

	err = writers.WriteFormattedFile(fileName, model)
	if err != nil {
		t.Fatal(err.Error())
	}

}

func TestGenerateSchema(t *testing.T) {
	s := &schema.Schema{}
	if err := json.Unmarshal([]byte(testdata.JSON1), s); err != nil {
		t.Fatal(err.Error())
	}

	// replace "~" with "`"
	result := strings.Replace(`
// Account represents a registered User
type Account struct {
	CreatedAt              time.Time ~json:"createdAt"~              // Profile's creation time
	EmailAddress           string    ~json:"emailAddress"~           // Email Address of the Account
	EmailStatusConstant    string    ~json:"emailStatusConstant"~    // Status of the Account's Email
	ID                     int64     ~json:"id,string"~              // The unique identifier for a Account's Profile
	Password               string    ~json:"password"~               // Salted Password of the Account
	PasswordStatusConstant string    ~json:"passwordStatusConstant"~ // Status of the Account's Password
	ProfileID              int64     ~json:"profileId"~              // The unique identifier for a Account's Profile
	Salt                   string    ~json:"salt"~                   // Salt used to hash Password of the Account
	StatusConstant         string    ~json:"statusConstant"~         // Status of the Account
	URL                    string    ~json:"url"~                    // Salted Password of the Account
	URLName                string    ~json:"urlName"~                // Salted Password of the Account
}`, "~", "`", -1)

	code, err := GenerateSchema(s.Definitions["Account"])
	if err != nil {
		t.Fatal(err.Error())
	}

	if result != string(code) {
		// fmt.Printf("foo %# v", pretty.Formatter(difflib.Diff([]string{result}, []string{string(code)})))
		t.Fatalf("Schema is not same. Wanted: %s, Get: %s", result, string(code))
	}
}

func TestGenerateValidators(t *testing.T) {
	s := &schema.Schema{}
	if err := json.Unmarshal([]byte(testdata.JSON1), s); err != nil {
		t.Fatal(err.Error())
	}
	result := `
// Validate validates the Account struct
func (a *Account) Validate() error {
	return govalidator.NewMulti(govalidator.Date(a.CreatedAt),
		govalidator.MaxLength(a.Salt, 255),
		govalidator.Min(float64(a.ID), 1.000000),
		govalidator.Min(float64(a.ProfileID), 1.000000),
		govalidator.MinLength(a.Password, 6),
		govalidator.MinLength(a.URL, 6),
		govalidator.MinLength(a.URLName, 6),
		govalidator.OneOf(a.EmailStatusConstant, []string{
			AccountEmailStatusConstant.Verified,
			AccountEmailStatusConstant.NotVerified,
		}),
		govalidator.OneOf(a.PasswordStatusConstant, []string{
			AccountPasswordStatusConstant.Valid,
			AccountPasswordStatusConstant.NeedsReset,
			AccountPasswordStatusConstant.Generated,
		}),
		govalidator.OneOf(a.StatusConstant, []string{
			AccountStatusConstant.Registered,
			AccountStatusConstant.Unregistered,
			AccountStatusConstant.NeedsManualVerification,
		})).Validate()
}`

	code, err := validators.Generate(s.Definitions["Account"])
	if err != nil {
		t.Fatal(err.Error())
	}

	if result != string(code) {
		t.Fatalf("Schema is not same. Wanted: %s, Get: %s", result, string(code))
	}
}

func TestGenerateFunctions(t *testing.T) {
	var s schema.Schema
	if err := json.Unmarshal([]byte(testdata.JSON1), &s); err != nil {
		t.Fatal(err.Error())
	}

	_, err := GenerateFunctions(&s)
	if err != nil {
		t.Fatal(err.Error())
	}
}
<commit_msg>Generator: fix tab/space problem<commit_after>package models

import (
	"encoding/json"
	"strings"

	"testing"

	"github.com/cihangir/gene/generators/folders"
	"github.com/cihangir/gene/generators/validators"
	"github.com/cihangir/gene/testdata"
	"github.com/cihangir/schema"

	"github.com/cihangir/gene/writers"
)

func TestGenerateModel(t *testing.T) {
	var s schema.Schema
	if err := json.Unmarshal([]byte(testdata.JSON1), &s); err != nil {
		t.Fatal(err.Error())
	}

	model, err := GenerateModel(&s)
	if err != nil {
		t.Fatal(err.Error())
	}

	folders.EnsureFolders("/tmp/", folders.FolderStucture)
	fileName := "/tmp/models/" + s.Title + ".go"

	err = writers.WriteFormattedFile(fileName, model)
	if err != nil {
		t.Fatal(err.Error())
	}

}

func TestGenerateSchema(t *testing.T) {
	s := &schema.Schema{}
	if err := json.Unmarshal([]byte(testdata.JSON1), s); err != nil {
		t.Fatal(err.Error())
	}

	// replace "~" with "`"
	result := strings.Replace(`
// Account represents a registered User
type Account struct {
	CreatedAt              time.Time ~json:"createdAt,omitempty"~              // Profile's creation time
	EmailAddress           string    ~json:"emailAddress"~                     // Email Address of the Account
	EmailStatusConstant    string    ~json:"emailStatusConstant,omitempty"~    // Status of the Account's Email
	ID                     int64     ~json:"id,omitempty,string"~              // The unique identifier for a Account's Profile
	Password               string    ~json:"password"~                         // Salted Password of the Account
	PasswordStatusConstant string    ~json:"passwordStatusConstant,omitempty"~ // Status of the Account's Password
	ProfileID              int64     ~json:"profileId,omitempty,string"~       // The unique identifier for a Account's Profile
	Salt                   string    ~json:"salt,omitempty"~                   // Salt used to hash Password of the Account
	StatusConstant         string    ~json:"statusConstant,omitempty"~         // Status of the Account
	URL                    string    ~json:"url,omitempty"~                    // Salted Password of the Account
	URLName                string    ~json:"urlName,omitempty"~                // Salted Password of the Account
}`, "~", "`", -1)

	code, err := GenerateSchema(s.Definitions["Account"])
	if err != nil {
		t.Fatal(err.Error())
	}

	if result != string(code) {
		// fmt.Printf("foo %# v", pretty.Formatter(difflib.Diff([]string{result}, []string{string(code)})))
		t.Fatalf("Schema is not same. Wanted: %s, Get: %s", result, string(code))
	}
}

func TestGenerateValidators(t *testing.T) {
	s := &schema.Schema{}
	if err := json.Unmarshal([]byte(testdata.JSON1), s); err != nil {
		t.Fatal(err.Error())
	}
	result := `
// Validate validates the Account struct
func (a *Account) Validate() error {
	return govalidator.NewMulti(govalidator.Date(a.CreatedAt),
		govalidator.MaxLength(a.Salt, 255),
		govalidator.Min(float64(a.ID), 1.000000),
		govalidator.Min(float64(a.ProfileID), 1.000000),
		govalidator.MinLength(a.Password, 6),
		govalidator.MinLength(a.URL, 6),
		govalidator.MinLength(a.URLName, 6),
		govalidator.OneOf(a.EmailStatusConstant, []string{
			AccountEmailStatusConstant.Verified,
			AccountEmailStatusConstant.NotVerified,
		}),
		govalidator.OneOf(a.PasswordStatusConstant, []string{
			AccountPasswordStatusConstant.Valid,
			AccountPasswordStatusConstant.NeedsReset,
			AccountPasswordStatusConstant.Generated,
		}),
		govalidator.OneOf(a.StatusConstant, []string{
			AccountStatusConstant.Registered,
			AccountStatusConstant.Unregistered,
			AccountStatusConstant.NeedsManualVerification,
		})).Validate()
}`

	code, err := validators.Generate(s.Definitions["Account"])
	if err != nil {
		t.Fatal(err.Error())
	}

	if result != string(code) {
		t.Fatalf("Schema is not same. Wanted: %s, Get: %s", result, string(code))
	}
}

func TestGenerateFunctions(t *testing.T) {
	var s schema.Schema
	if err := json.Unmarshal([]byte(testdata.JSON1), &s); err != nil {
		t.Fatal(err.Error())
	}

	_, err := GenerateFunctions(&s)
	if err != nil {
		t.Fatal(err.Error())
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Vitess Authors.

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

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

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

package mysql

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

	"vitess.io/vitess/go/vt/proto/vtrpc"
	"vitess.io/vitess/go/vt/vterrors"
)

const (
	// MaximumPositionSize is the maximum size of a
	// replication position. It is used as the maximum column size in the _vt.reparent_journal and
	// other related tables. A row has a maximum size of 65535 bytes. So
	// we want to stay under that. We use VARBINARY so the
	// character set doesn't matter, we only store ascii
	// characters anyway.
	MaximumPositionSize = 64000
)

// Position represents the information necessary to describe which
// transactions a server has seen, so that it can request a replication stream
// from a new source that picks up where it left off.
//
// This must be a concrete struct because custom Unmarshalers can't be
// registered on an interface.
//
// The == operator should not be used with Position, because the
// underlying GTIDSet might use slices, which are not comparable. Using == in
// those cases will result in a run-time panic.
type Position struct {
	// This is a zero byte compile-time check that no one is trying to
	// use == or != with Position. Without this, we won't know there's
	// a problem until the runtime panic. Note that this must not be
	// the last field of the struct, or else the Go compiler will add
	// padding to prevent pointers to this field from becoming invalid.
	_ [0]struct{ _ []byte }

	// GTIDSet is the underlying GTID set. It must not be anonymous,
	// or else Position would itself also implement the GTIDSet interface.
	GTIDSet GTIDSet
}

// Equal returns true if this position is equal to another.
func (rp Position) Equal(other Position) bool {
	if rp.GTIDSet == nil {
		return other.GTIDSet == nil
	}
	return rp.GTIDSet.Equal(other.GTIDSet)
}

// AtLeast returns true if this position is equal to or after another.
func (rp Position) AtLeast(other Position) bool {
	if other.GTIDSet == nil {
		return true
	}
	if rp.GTIDSet == nil {
		return false
	}
	return rp.GTIDSet.Contains(other.GTIDSet)
}

// String returns a string representation of the underlying GTIDSet.
// If the set is nil, it returns "<nil>" in the style of Sprintf("%v", nil).
func (rp Position) String() string {
	if rp.GTIDSet == nil {
		return "<nil>"
	}
	return rp.GTIDSet.String()
}

// IsZero returns true if this is the zero value, Position{}.
func (rp Position) IsZero() bool {
	return rp.GTIDSet == nil
}

// AppendGTID returns a new Position that represents the position
// after the given GTID is replicated.
func AppendGTID(rp Position, gtid GTID) Position {
	if gtid == nil {
		return rp
	}
	if rp.GTIDSet == nil {
		return Position{GTIDSet: gtid.GTIDSet()}
	}
	return Position{GTIDSet: rp.GTIDSet.AddGTID(gtid)}
}

// MustParsePosition calls ParsePosition and panics
// on error.
func MustParsePosition(flavor, value string) Position {
	rp, err := ParsePosition(flavor, value)
	if err != nil {
		panic(err)
	}
	return rp
}

// EncodePosition returns a string that contains both the flavor
// and value of the Position, so that the correct parser can be
// selected when that string is passed to DecodePosition.
func EncodePosition(rp Position) string {
	if rp.GTIDSet == nil {
		return ""
	}
	return fmt.Sprintf("%s/%s", rp.GTIDSet.Flavor(), rp.GTIDSet.String())
}

// DecodePosition converts a string in the format returned by
// EncodePosition back into a Position value with the
// correct underlying flavor.
func DecodePosition(s string) (rp Position, err error) {
	if s == "" {
		return rp, nil
	}

	parts := strings.SplitN(s, "/", 2)
	if len(parts) != 2 {
		// There is no flavor. Try looking for a default parser.
		return ParsePosition("", s)
	}
	return ParsePosition(parts[0], parts[1])
}

// ParsePosition calls the parser for the specified flavor.
func ParsePosition(flavor, value string) (rp Position, err error) {
	parser := gtidSetParsers[flavor]
	if parser == nil {
		return rp, vterrors.Errorf(vtrpc.Code_INTERNAL, "parse error: unknown GTIDSet flavor %#v", flavor)
	}
	gtidSet, err := parser(value)
	if err != nil {
		return rp, err
	}
	rp.GTIDSet = gtidSet
	return rp, err
}

// MarshalJSON implements encoding/json.Marshaler.
func (rp Position) MarshalJSON() ([]byte, error) {
	return json.Marshal(EncodePosition(rp))
}

// UnmarshalJSON implements encoding/json.Unmarshaler.
func (rp *Position) UnmarshalJSON(buf []byte) error {
	var s string
	err := json.Unmarshal(buf, &s)
	if err != nil {
		return err
	}

	*rp, err = DecodePosition(s)
	if err != nil {
		return err
	}
	return nil
}

// MatchesFlavor will take a flavor string, and return whether the positions GTIDSet matches the supplied flavor.
// The caller should use the constants Mysql56FlavorID, MariadbFlavorID, or FilePosFlavorID when supplying the flavor string.
func (rp *Position) MatchesFlavor(flavor string) bool {
	switch flavor {
	case Mysql56FlavorID:
		_, matches := rp.GTIDSet.(Mysql56GTIDSet)
		return matches
	case MariadbFlavorID:
		_, matches := rp.GTIDSet.(MariadbGTIDSet)
		return matches
	case FilePosFlavorID:
		_, matches := rp.GTIDSet.(filePosGTID)
		return matches
	}
	return false
}

// Comparable returns whether the receiver is comparable to the supplied position, based on whether one
// of the two positions contains the other.
func (rp *Position) Comparable(other Position) bool {
	return rp.GTIDSet.Contains(other.GTIDSet) || other.GTIDSet.Contains(rp.GTIDSet)
}

// AllPositionsComparable returns true if all positions in the supplied list are comparable with one another, and false
// if any are non-comparable.
func AllPositionsComparable(positions []Position) bool {
	for i := 0; i < len(positions); i++ {
		for j := i + 1; j < len(positions); j++ {
			if !positions[i].Comparable(positions[j]) {
				return false
			}
		}
	}

	return true
}
<commit_msg>docs: added comments explaining the functioning of the AtLeast function<commit_after>/*
Copyright 2019 The Vitess Authors.

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

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

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

package mysql

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

	"vitess.io/vitess/go/vt/proto/vtrpc"
	"vitess.io/vitess/go/vt/vterrors"
)

const (
	// MaximumPositionSize is the maximum size of a
	// replication position. It is used as the maximum column size in the _vt.reparent_journal and
	// other related tables. A row has a maximum size of 65535 bytes. So
	// we want to stay under that. We use VARBINARY so the
	// character set doesn't matter, we only store ascii
	// characters anyway.
	MaximumPositionSize = 64000
)

// Position represents the information necessary to describe which
// transactions a server has seen, so that it can request a replication stream
// from a new source that picks up where it left off.
//
// This must be a concrete struct because custom Unmarshalers can't be
// registered on an interface.
//
// The == operator should not be used with Position, because the
// underlying GTIDSet might use slices, which are not comparable. Using == in
// those cases will result in a run-time panic.
type Position struct {
	// This is a zero byte compile-time check that no one is trying to
	// use == or != with Position. Without this, we won't know there's
	// a problem until the runtime panic. Note that this must not be
	// the last field of the struct, or else the Go compiler will add
	// padding to prevent pointers to this field from becoming invalid.
	_ [0]struct{ _ []byte }

	// GTIDSet is the underlying GTID set. It must not be anonymous,
	// or else Position would itself also implement the GTIDSet interface.
	GTIDSet GTIDSet
}

// Equal returns true if this position is equal to another.
func (rp Position) Equal(other Position) bool {
	if rp.GTIDSet == nil {
		return other.GTIDSet == nil
	}
	return rp.GTIDSet.Equal(other.GTIDSet)
}

// AtLeast returns true if this position is equal to or after another.
func (rp Position) AtLeast(other Position) bool {
	// The underlying call to the Contains method of the GTIDSet interface
	// does not guarantee handling the nil cases correctly.
	// So, we have to do nil case handling here
	if other.GTIDSet == nil {
		// If the other GTIDSet is nil, then it is contained
		// in all possible GTIDSets, even nil ones.
		return true
	}
	if rp.GTIDSet == nil {
		// Here rp GTIDSet is nil but the other GTIDSet isn't.
		// So it is not contained in the rp GTIDSet.
		return false
	}
	return rp.GTIDSet.Contains(other.GTIDSet)
}

// String returns a string representation of the underlying GTIDSet.
// If the set is nil, it returns "<nil>" in the style of Sprintf("%v", nil).
func (rp Position) String() string {
	if rp.GTIDSet == nil {
		return "<nil>"
	}
	return rp.GTIDSet.String()
}

// IsZero returns true if this is the zero value, Position{}.
func (rp Position) IsZero() bool {
	return rp.GTIDSet == nil
}

// AppendGTID returns a new Position that represents the position
// after the given GTID is replicated.
func AppendGTID(rp Position, gtid GTID) Position {
	if gtid == nil {
		return rp
	}
	if rp.GTIDSet == nil {
		return Position{GTIDSet: gtid.GTIDSet()}
	}
	return Position{GTIDSet: rp.GTIDSet.AddGTID(gtid)}
}

// MustParsePosition calls ParsePosition and panics
// on error.
func MustParsePosition(flavor, value string) Position {
	rp, err := ParsePosition(flavor, value)
	if err != nil {
		panic(err)
	}
	return rp
}

// EncodePosition returns a string that contains both the flavor
// and value of the Position, so that the correct parser can be
// selected when that string is passed to DecodePosition.
func EncodePosition(rp Position) string {
	if rp.GTIDSet == nil {
		return ""
	}
	return fmt.Sprintf("%s/%s", rp.GTIDSet.Flavor(), rp.GTIDSet.String())
}

// DecodePosition converts a string in the format returned by
// EncodePosition back into a Position value with the
// correct underlying flavor.
func DecodePosition(s string) (rp Position, err error) {
	if s == "" {
		return rp, nil
	}

	parts := strings.SplitN(s, "/", 2)
	if len(parts) != 2 {
		// There is no flavor. Try looking for a default parser.
		return ParsePosition("", s)
	}
	return ParsePosition(parts[0], parts[1])
}

// ParsePosition calls the parser for the specified flavor.
func ParsePosition(flavor, value string) (rp Position, err error) {
	parser := gtidSetParsers[flavor]
	if parser == nil {
		return rp, vterrors.Errorf(vtrpc.Code_INTERNAL, "parse error: unknown GTIDSet flavor %#v", flavor)
	}
	gtidSet, err := parser(value)
	if err != nil {
		return rp, err
	}
	rp.GTIDSet = gtidSet
	return rp, err
}

// MarshalJSON implements encoding/json.Marshaler.
func (rp Position) MarshalJSON() ([]byte, error) {
	return json.Marshal(EncodePosition(rp))
}

// UnmarshalJSON implements encoding/json.Unmarshaler.
func (rp *Position) UnmarshalJSON(buf []byte) error {
	var s string
	err := json.Unmarshal(buf, &s)
	if err != nil {
		return err
	}

	*rp, err = DecodePosition(s)
	if err != nil {
		return err
	}
	return nil
}

// MatchesFlavor will take a flavor string, and return whether the positions GTIDSet matches the supplied flavor.
// The caller should use the constants Mysql56FlavorID, MariadbFlavorID, or FilePosFlavorID when supplying the flavor string.
func (rp *Position) MatchesFlavor(flavor string) bool {
	switch flavor {
	case Mysql56FlavorID:
		_, matches := rp.GTIDSet.(Mysql56GTIDSet)
		return matches
	case MariadbFlavorID:
		_, matches := rp.GTIDSet.(MariadbGTIDSet)
		return matches
	case FilePosFlavorID:
		_, matches := rp.GTIDSet.(filePosGTID)
		return matches
	}
	return false
}

// Comparable returns whether the receiver is comparable to the supplied position, based on whether one
// of the two positions contains the other.
func (rp *Position) Comparable(other Position) bool {
	return rp.GTIDSet.Contains(other.GTIDSet) || other.GTIDSet.Contains(rp.GTIDSet)
}

// AllPositionsComparable returns true if all positions in the supplied list are comparable with one another, and false
// if any are non-comparable.
func AllPositionsComparable(positions []Position) bool {
	for i := 0; i < len(positions); i++ {
		for j := i + 1; j < len(positions); j++ {
			if !positions[i].Comparable(positions[j]) {
				return false
			}
		}
	}

	return true
}
<|endoftext|>


<commit_before>package stellarsvc

import (
	"context"
	"errors"

	"github.com/keybase/client/go/libkb"
	"github.com/keybase/client/go/protocol/stellar1"
	"github.com/keybase/client/go/stellar"
	"github.com/keybase/client/go/stellar/remote"
	"github.com/keybase/stellarnet"
)

// Service handlers

// CreateWallet creates and posts an initial stellar bundle for a user.
// Only succeeds if they do not already have one.
// Safe to call even if the user has a bundle already.
func CreateWallet(ctx context.Context, g *libkb.GlobalContext) (created bool, err error) {
	return stellar.CreateWallet(ctx, g)
}

type balancesResult struct {
	Status   libkb.AppStatus    `json:"status"`
	Balances []stellar1.Balance `json:"balances"`
}

func (b *balancesResult) GetAppStatus() *libkb.AppStatus {
	return &b.Status
}

func Balances(ctx context.Context, g *libkb.GlobalContext, accountID stellar1.AccountID) ([]stellar1.Balance, error) {
	apiArg := libkb.APIArg{
		Endpoint:    "stellar/balances",
		SessionType: libkb.APISessionTypeREQUIRED,
		Args:        libkb.HTTPArgs{"account_id": libkb.S{Val: string(accountID)}},
		NetContext:  ctx,
	}

	var res balancesResult
	if err := g.API.GetDecode(apiArg, &res); err != nil {
		return nil, err
	}

	return res.Balances, nil
}

func balanceXLM(ctx context.Context, g *libkb.GlobalContext, accountID stellar1.AccountID) (stellar1.Balance, error) {
	balances, err := Balances(ctx, g, accountID)
	if err != nil {
		return stellar1.Balance{}, err
	}

	for _, b := range balances {
		if b.Asset.Type == "native" {
			return b, nil
		}
	}

	return stellar1.Balance{}, errors.New("no native balance")
}

func lookupSenderPrimary(ctx context.Context, g *libkb.GlobalContext) (stellar1.BundleEntry, error) {
	bundle, _, err := remote.Fetch(ctx, g)
	if err != nil {
		return stellar1.BundleEntry{}, err
	}

	primary, err := bundle.PrimaryAccount()
	if err != nil {
		return stellar1.BundleEntry{}, err
	}
	if len(primary.Signers) == 0 {
		return stellar1.BundleEntry{}, errors.New("no signer for primary bundle")
	}
	if len(primary.Signers) > 1 {
		return stellar1.BundleEntry{}, errors.New("only single signer supported")
	}

	return primary, nil
}

type Recipient struct {
	Input     string
	User      *libkb.User
	AccountID stellarnet.AddressStr
}

// TODO: actually lookup the recipient and handle
// 1. stellar addresses GXXXXX
// 2. stellar federation address rebecca*keybase.io
// 3. keybase username
// 4. keybase assertion
func lookupRecipient(ctx context.Context, g *libkb.GlobalContext, to string) (*Recipient, error) {
	r := Recipient{
		Input: to,
	}

	// temporary: only handle stellar addresses
	var err error
	r.AccountID, err = stellarnet.NewAddressStr(to)
	if err != nil {
		return nil, err
	}

	return &r, nil
}

type submitResult struct {
	Status        libkb.AppStatus        `json:"status"`
	PaymentResult stellar1.PaymentResult `json:"payment_result"`
}

func (s *submitResult) GetAppStatus() *libkb.AppStatus {
	return &s.Status
}

func postFromCurrentUser(ctx context.Context, g *libkb.GlobalContext, acctID stellarnet.AddressStr) stellar1.PaymentPost {
	uid, deviceID, _, _, _ := g.ActiveDevice.AllFields()
	return stellar1.PaymentPost{
		Members: stellar1.Members{
			FromStellar:  stellar1.AccountID(acctID.String()),
			FromKeybase:  g.Env.GetUsername().String(),
			FromUID:      uid,
			FromDeviceID: deviceID,
		},
	}
}

func Send(ctx context.Context, g *libkb.GlobalContext, arg stellar1.SendLocalArg) (stellar1.PaymentResult, error) {
	// look up sender wallet
	primary, err := lookupSenderPrimary(ctx, g)
	if err != nil {
		return stellar1.PaymentResult{}, err
	}
	primaryAccountID, err := stellarnet.NewAddressStr(primary.AccountID.String())
	if err != nil {
		return stellar1.PaymentResult{}, err
	}
	primarySeed, err := stellarnet.NewSeedStr(primary.Signers[0].SecureNoLogString())
	if err != nil {
		return stellar1.PaymentResult{}, err
	}
	senderAcct := stellarnet.NewAccount(primaryAccountID)

	// look up recipient
	recipient, err := lookupRecipient(ctx, g, arg.Recipient)
	if err != nil {
		return stellar1.PaymentResult{}, err
	}

	post := postFromCurrentUser(ctx, g, primaryAccountID)

	// Note:
	// CreateAccountXLMTransaction and PaymentXLMTransaction use horizon
	// to get the sequence number.  In the future we could provide an RPC
	// to get it instead of having the clients go direct to horizon.

	// check if recipient account exists
	_, err = balanceXLM(ctx, g, stellar1.AccountID(recipient.AccountID.String()))
	if err != nil {
		// if no balance, create_account operation
		// we could check here to make sure that amount is at least 1XLM
		// but for now, just let stellar-core tell us there was an error
		post.StellarAccountSeqno, post.SignedTransaction, err = senderAcct.CreateAccountXLMTransaction(primarySeed, recipient.AccountID, arg.Amount)
		if err != nil {
			return stellar1.PaymentResult{}, err
		}
	} else {
		// if balance, payment operation
		post.StellarAccountSeqno, post.SignedTransaction, err = senderAcct.PaymentXLMTransaction(primarySeed, recipient.AccountID, arg.Amount)
		if err != nil {
			return stellar1.PaymentResult{}, err
		}
	}

	// submit the transaction
	payload := make(libkb.JSONPayload)
	payload["payment"] = post
	apiArg := libkb.APIArg{
		Endpoint:    "stellar/submitpayment",
		SessionType: libkb.APISessionTypeREQUIRED,
		JSONPayload: payload,
		NetContext:  ctx,
	}

	var res submitResult
	if err := g.API.PostDecode(apiArg, &res); err != nil {
		return stellar1.PaymentResult{}, err
	}

	return res.PaymentResult, nil
}

func OwnAccount(ctx context.Context, g *libkb.GlobalContext, accountID stellar1.AccountID) (isOwn bool, err error) {
	bundle, _, err := remote.Fetch(ctx, g)
	if err != nil {
		return false, err
	}
	for _, account := range bundle.Accounts {
		if account.AccountID.Eq(accountID) {
			return true, nil
		}
	}
	return false, nil
}

func ImportSecretKey(ctx context.Context, g *libkb.GlobalContext, secretKey stellar1.SecretKey, makePrimary bool) (err error) {
	return stellar.ImportSecretKey(ctx, g, secretKey, makePrimary)
}

func Dump(ctx context.Context, g *libkb.GlobalContext) (res stellar1.Bundle, err error) {
	if g.Env.GetRunMode() != libkb.DevelRunMode {
		return res, errors.New("WalletDump only supported in devel run mode")
	}
	res, _, err = remote.Fetch(ctx, g)
	return res, err
}
<commit_msg>Lookup keybase recipient users (#11144)<commit_after>package stellarsvc

import (
	"context"
	"errors"
	"fmt"

	"github.com/keybase/client/go/libkb"
	"github.com/keybase/client/go/protocol/stellar1"
	"github.com/keybase/client/go/stellar"
	"github.com/keybase/client/go/stellar/remote"
	"github.com/keybase/stellarnet"
)

// Service handlers

// CreateWallet creates and posts an initial stellar bundle for a user.
// Only succeeds if they do not already have one.
// Safe to call even if the user has a bundle already.
func CreateWallet(ctx context.Context, g *libkb.GlobalContext) (created bool, err error) {
	return stellar.CreateWallet(ctx, g)
}

type balancesResult struct {
	Status   libkb.AppStatus    `json:"status"`
	Balances []stellar1.Balance `json:"balances"`
}

func (b *balancesResult) GetAppStatus() *libkb.AppStatus {
	return &b.Status
}

func Balances(ctx context.Context, g *libkb.GlobalContext, accountID stellar1.AccountID) ([]stellar1.Balance, error) {
	apiArg := libkb.APIArg{
		Endpoint:    "stellar/balances",
		SessionType: libkb.APISessionTypeREQUIRED,
		Args:        libkb.HTTPArgs{"account_id": libkb.S{Val: string(accountID)}},
		NetContext:  ctx,
	}

	var res balancesResult
	if err := g.API.GetDecode(apiArg, &res); err != nil {
		return nil, err
	}

	return res.Balances, nil
}

func balanceXLM(ctx context.Context, g *libkb.GlobalContext, accountID stellar1.AccountID) (stellar1.Balance, error) {
	balances, err := Balances(ctx, g, accountID)
	if err != nil {
		return stellar1.Balance{}, err
	}

	for _, b := range balances {
		if b.Asset.Type == "native" {
			return b, nil
		}
	}

	return stellar1.Balance{}, errors.New("no native balance")
}

func lookupSenderPrimary(ctx context.Context, g *libkb.GlobalContext) (stellar1.BundleEntry, error) {
	bundle, _, err := remote.Fetch(ctx, g)
	if err != nil {
		return stellar1.BundleEntry{}, err
	}

	primary, err := bundle.PrimaryAccount()
	if err != nil {
		return stellar1.BundleEntry{}, err
	}
	if len(primary.Signers) == 0 {
		return stellar1.BundleEntry{}, errors.New("no signer for primary bundle")
	}
	if len(primary.Signers) > 1 {
		return stellar1.BundleEntry{}, errors.New("only single signer supported")
	}

	return primary, nil
}

type Recipient struct {
	Input     string
	User      *libkb.User
	AccountID stellarnet.AddressStr
}

// TODO: actually lookup the recipient and handle
// 1. stellar addresses GXXXXX
// 2. stellar federation address rebecca*keybase.io
// 3. keybase username
// 4. keybase assertion
func lookupRecipient(ctx context.Context, g *libkb.GlobalContext, to string) (*Recipient, error) {
	r := Recipient{
		Input: to,
	}

	if to[0] == 'G' && len(to) > 16 {
		var err error
		r.AccountID, err = stellarnet.NewAddressStr(to)
		if err != nil {
			return nil, err
		}

		return &r, nil
	}

	user, err := libkb.LoadUser(libkb.NewLoadUserByNameArg(g, to))
	if err != nil {
		return nil, err
	}
	accountID := user.StellarWalletAddress()
	if accountID == nil {
		return nil, fmt.Errorf("keybase user %s does not have a stellar wallet", to)
	}
	r.AccountID, err = stellarnet.NewAddressStr(accountID.String())
	if err != nil {
		return nil, err
	}
	r.User = user

	return &r, nil
}

type submitResult struct {
	Status        libkb.AppStatus        `json:"status"`
	PaymentResult stellar1.PaymentResult `json:"payment_result"`
}

func (s *submitResult) GetAppStatus() *libkb.AppStatus {
	return &s.Status
}

func postFromCurrentUser(ctx context.Context, g *libkb.GlobalContext, acctID stellarnet.AddressStr, recipient *Recipient) stellar1.PaymentPost {
	uid, deviceID, _, _, _ := g.ActiveDevice.AllFields()
	post := stellar1.PaymentPost{
		Members: stellar1.Members{
			FromStellar:  stellar1.AccountID(acctID.String()),
			FromKeybase:  g.Env.GetUsername().String(),
			FromUID:      uid,
			FromDeviceID: deviceID,
		},
	}

	if recipient != nil {
		post.Members.ToStellar = stellar1.AccountID(recipient.AccountID.String())
		if recipient.User != nil {
			post.Members.ToUID = recipient.User.GetUID()
			post.Members.ToKeybase = recipient.User.GetName()
		}
	}

	return post
}

func Send(ctx context.Context, g *libkb.GlobalContext, arg stellar1.SendLocalArg) (stellar1.PaymentResult, error) {
	// look up sender wallet
	primary, err := lookupSenderPrimary(ctx, g)
	if err != nil {
		return stellar1.PaymentResult{}, err
	}
	primaryAccountID, err := stellarnet.NewAddressStr(primary.AccountID.String())
	if err != nil {
		return stellar1.PaymentResult{}, err
	}
	primarySeed, err := stellarnet.NewSeedStr(primary.Signers[0].SecureNoLogString())
	if err != nil {
		return stellar1.PaymentResult{}, err
	}
	senderAcct := stellarnet.NewAccount(primaryAccountID)

	// look up recipient
	recipient, err := lookupRecipient(ctx, g, arg.Recipient)
	if err != nil {
		return stellar1.PaymentResult{}, err
	}

	post := postFromCurrentUser(ctx, g, primaryAccountID, recipient)

	// Note:
	// CreateAccountXLMTransaction and PaymentXLMTransaction use horizon
	// to get the sequence number.  In the future we could provide an RPC
	// to get it instead of having the clients go direct to horizon.

	// check if recipient account exists
	_, err = balanceXLM(ctx, g, stellar1.AccountID(recipient.AccountID.String()))
	if err != nil {
		// if no balance, create_account operation
		// we could check here to make sure that amount is at least 1XLM
		// but for now, just let stellar-core tell us there was an error
		post.StellarAccountSeqno, post.SignedTransaction, err = senderAcct.CreateAccountXLMTransaction(primarySeed, recipient.AccountID, arg.Amount)
		if err != nil {
			return stellar1.PaymentResult{}, err
		}
	} else {
		// if balance, payment operation
		post.StellarAccountSeqno, post.SignedTransaction, err = senderAcct.PaymentXLMTransaction(primarySeed, recipient.AccountID, arg.Amount)
		if err != nil {
			return stellar1.PaymentResult{}, err
		}
	}

	// submit the transaction
	payload := make(libkb.JSONPayload)
	payload["payment"] = post
	apiArg := libkb.APIArg{
		Endpoint:    "stellar/submitpayment",
		SessionType: libkb.APISessionTypeREQUIRED,
		JSONPayload: payload,
		NetContext:  ctx,
	}

	var res submitResult
	if err := g.API.PostDecode(apiArg, &res); err != nil {
		return stellar1.PaymentResult{}, err
	}

	return res.PaymentResult, nil
}

func OwnAccount(ctx context.Context, g *libkb.GlobalContext, accountID stellar1.AccountID) (isOwn bool, err error) {
	bundle, _, err := remote.Fetch(ctx, g)
	if err != nil {
		return false, err
	}
	for _, account := range bundle.Accounts {
		if account.AccountID.Eq(accountID) {
			return true, nil
		}
	}
	return false, nil
}

func ImportSecretKey(ctx context.Context, g *libkb.GlobalContext, secretKey stellar1.SecretKey, makePrimary bool) (err error) {
	return stellar.ImportSecretKey(ctx, g, secretKey, makePrimary)
}

func Dump(ctx context.Context, g *libkb.GlobalContext) (res stellar1.Bundle, err error) {
	if g.Env.GetRunMode() != libkb.DevelRunMode {
		return res, errors.New("WalletDump only supported in devel run mode")
	}
	res, _, err = remote.Fetch(ctx, g)
	return res, err
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Vitess Authors.

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

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

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

package discovery

import (
	"fmt"
	"math/rand"
	"strings"
	"sync"
	"time"

	"vitess.io/vitess/go/vt/topo/topoproto"

	vtrpcpb "vitess.io/vitess/go/vt/proto/vtrpc"

	"vitess.io/vitess/go/vt/vttablet/tabletconn"

	"vitess.io/vitess/go/vt/log"

	"golang.org/x/net/context"
	topodatapb "vitess.io/vitess/go/vt/proto/topodata"
	"vitess.io/vitess/go/vt/topo"
	"vitess.io/vitess/go/vt/vterrors"
)

var (
	tabletPickerRetryDelay   = 30 * time.Second
	muTabletPickerRetryDelay sync.Mutex
)

// GetTabletPickerRetryDelay synchronizes changes to tabletPickerRetryDelay. Used in tests only at the moment
func GetTabletPickerRetryDelay() time.Duration {
	muTabletPickerRetryDelay.Lock()
	defer muTabletPickerRetryDelay.Unlock()
	return tabletPickerRetryDelay
}

// SetTabletPickerRetryDelay synchronizes reads for tabletPickerRetryDelay. Used in tests only at the moment
func SetTabletPickerRetryDelay(delay time.Duration) {
	muTabletPickerRetryDelay.Lock()
	defer muTabletPickerRetryDelay.Unlock()
	tabletPickerRetryDelay = delay
}

// TabletPicker gives a simplified API for picking tablets.
type TabletPicker struct {
	ts          *topo.Server
	cells       []string
	keyspace    string
	shard       string
	tabletTypes []topodatapb.TabletType
}

// NewTabletPicker returns a TabletPicker.
func NewTabletPicker(ts *topo.Server, cells []string, keyspace, shard, tabletTypesStr string) (*TabletPicker, error) {
	tabletTypes, err := topoproto.ParseTabletTypes(tabletTypesStr)
	if err != nil {
		return nil, vterrors.Errorf(vtrpcpb.Code_FAILED_PRECONDITION, "failed to parse list of tablet types: %v", tabletTypesStr)
	}
	var missingFields []string
	if keyspace == "" {
		missingFields = append(missingFields, "Keyspace")
	}
	if shard == "" {
		missingFields = append(missingFields, "Shard")
	}
	if len(cells) == 0 {
		missingFields = append(missingFields, "Cells")
	}
	if len(missingFields) > 0 {
		//log.Errorf("missing picker fields %s", debug.Stack()) //FIXME: remove after all tests run
		return nil, vterrors.Errorf(vtrpcpb.Code_FAILED_PRECONDITION,
			fmt.Sprintf("Missing required field(s) for tablet picker: %s", strings.Join(missingFields, ", ")))
	}
	return &TabletPicker{
		ts:          ts,
		cells:       cells,
		keyspace:    keyspace,
		shard:       shard,
		tabletTypes: tabletTypes,
	}, nil
}

// PickForStreaming picks an available tablet
// All tablets that belong to tp.cells are evaluated and one is
// chosen at random
func (tp *TabletPicker) PickForStreaming(ctx context.Context) (*topodatapb.Tablet, error) {
	// keep trying at intervals (tabletPickerRetryDelay) until a tablet is found
	// or the context is canceled
	for {
		select {
		case <-ctx.Done():
			return nil, vterrors.Errorf(vtrpcpb.Code_CANCELED, "context has expired")
		default:
		}
		candidates := tp.getMatchingTablets(ctx)
		if len(candidates) == 0 {
			// if no candidates were found, sleep and try again
			log.Infof("No tablet found for streaming, sleeping for %d seconds", int(GetTabletPickerRetryDelay()/1e9))
			time.Sleep(GetTabletPickerRetryDelay())
			continue
		}
		// try at most len(candidate) times to find a healthy tablet
		for i := 0; i < len(candidates); i++ {
			idx := rand.Intn(len(candidates))
			ti := candidates[idx]
			// get tablet
			// try to connect to tablet
			conn, err := tabletconn.GetDialer()(ti.Tablet, true)
			if err != nil {
				log.Warningf("unable to connect to tablet for alias %v", ti.Alias)
				candidates = append(candidates[:idx], candidates[idx+1:]...)
				if len(candidates) == 0 {
					break
				}
				continue
			}
			// OK to use ctx here because it is not actually used by the underlying Close implementation
			_ = conn.Close(ctx)
			return ti.Tablet, nil
		}
	}
}

// getMatchingTablets returns a list of TabletInfo for tablets
// that match the cells, keyspace, shard and tabletTypes for this TabletPicker
func (tp *TabletPicker) getMatchingTablets(ctx context.Context) []*topo.TabletInfo {
	// Special handling for MASTER tablet type
	// Since there is only one master, we ignore cell and find the master
	aliases := make([]*topodatapb.TabletAlias, 0)
	if len(tp.tabletTypes) == 1 && tp.tabletTypes[0] == topodatapb.TabletType_MASTER {
		shortCtx, cancel := context.WithTimeout(ctx, *topo.RemoteOperationTimeout)
		defer cancel()
		si, err := tp.ts.GetShard(shortCtx, tp.keyspace, tp.shard)
		if err != nil {
			return nil
		}
		aliases = append(aliases, si.MasterAlias)
	} else {
		actualCells := make([]string, 0)
		for _, cell := range tp.cells {
			// check if cell is actually an alias
			// non-blocking read so that this is fast
			shortCtx, cancel := context.WithTimeout(ctx, *topo.RemoteOperationTimeout)
			defer cancel()
			alias, err := tp.ts.GetCellsAlias(shortCtx, cell, false)
			if err != nil {
				// either cellAlias doesn't exist or it isn't a cell alias at all. In that case assume it is a cell
				actualCells = append(actualCells, cell)
			} else {
				actualCells = append(actualCells, alias.Cells...)
			}
		}
		for _, cell := range actualCells {
			shortCtx, cancel := context.WithTimeout(ctx, *topo.RemoteOperationTimeout)
			defer cancel()
			// match cell, keyspace and shard
			sri, err := tp.ts.GetShardReplication(shortCtx, cell, tp.keyspace, tp.shard)
			if err != nil {
				//log.Errorf("missing shard in topo %s", debug.Stack()) //FIXME: remove after all tests run

				//log.Warningf("error %v from GetShardReplication for %v %v %v", err, cell, tp.keyspace, tp.shard)
				continue
			}

			for _, node := range sri.Nodes {
				aliases = append(aliases, node.TabletAlias)
			}
		}
	}

	if len(aliases) == 0 {
		return nil
	}
	shortCtx, cancel := context.WithTimeout(ctx, *topo.RemoteOperationTimeout)
	defer cancel()
	tabletMap, err := tp.ts.GetTabletMap(shortCtx, aliases)
	if err != nil {
		log.Warningf("error fetching tablets from topo: %v", err)
		return nil
	}
	tablets := make([]*topo.TabletInfo, 0, len(aliases))
	for _, tabletAlias := range aliases {
		tabletInfo, ok := tabletMap[topoproto.TabletAliasString(tabletAlias)]
		if !ok {
			// tablet disappeared on us (GetTabletMap ignores
			// topo.ErrNoNode), just echo a warning
			log.Warningf("failed to load tablet %v", tabletAlias)
		} else if topoproto.IsTypeInList(tabletInfo.Type, tp.tabletTypes) {
			tablets = append(tablets, tabletInfo)
		}
	}
	return tablets
}

func init() {
	// TODO(sougou): consolidate this call to be once per process.
	rand.Seed(time.Now().UnixNano())
}
<commit_msg>tablet picker: use partial result from GetTabletMap if it is available. check for cell before cellsAlias to reduce number of topo errors.<commit_after>/*
Copyright 2019 The Vitess Authors.

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

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

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

package discovery

import (
	"fmt"
	"math/rand"
	"strings"
	"sync"
	"time"

	"vitess.io/vitess/go/vt/topo/topoproto"

	vtrpcpb "vitess.io/vitess/go/vt/proto/vtrpc"

	"vitess.io/vitess/go/vt/vttablet/tabletconn"

	"vitess.io/vitess/go/vt/log"

	"golang.org/x/net/context"
	topodatapb "vitess.io/vitess/go/vt/proto/topodata"
	"vitess.io/vitess/go/vt/topo"
	"vitess.io/vitess/go/vt/vterrors"
)

var (
	tabletPickerRetryDelay   = 30 * time.Second
	muTabletPickerRetryDelay sync.Mutex
)

// GetTabletPickerRetryDelay synchronizes changes to tabletPickerRetryDelay. Used in tests only at the moment
func GetTabletPickerRetryDelay() time.Duration {
	muTabletPickerRetryDelay.Lock()
	defer muTabletPickerRetryDelay.Unlock()
	return tabletPickerRetryDelay
}

// SetTabletPickerRetryDelay synchronizes reads for tabletPickerRetryDelay. Used in tests only at the moment
func SetTabletPickerRetryDelay(delay time.Duration) {
	muTabletPickerRetryDelay.Lock()
	defer muTabletPickerRetryDelay.Unlock()
	tabletPickerRetryDelay = delay
}

// TabletPicker gives a simplified API for picking tablets.
type TabletPicker struct {
	ts          *topo.Server
	cells       []string
	keyspace    string
	shard       string
	tabletTypes []topodatapb.TabletType
}

// NewTabletPicker returns a TabletPicker.
func NewTabletPicker(ts *topo.Server, cells []string, keyspace, shard, tabletTypesStr string) (*TabletPicker, error) {
	tabletTypes, err := topoproto.ParseTabletTypes(tabletTypesStr)
	if err != nil {
		return nil, vterrors.Errorf(vtrpcpb.Code_FAILED_PRECONDITION, "failed to parse list of tablet types: %v", tabletTypesStr)
	}
	var missingFields []string
	if keyspace == "" {
		missingFields = append(missingFields, "Keyspace")
	}
	if shard == "" {
		missingFields = append(missingFields, "Shard")
	}
	if len(cells) == 0 {
		missingFields = append(missingFields, "Cells")
	}
	if len(missingFields) > 0 {
		//log.Errorf("missing picker fields %s", debug.Stack()) //FIXME: remove after all tests run
		return nil, vterrors.Errorf(vtrpcpb.Code_FAILED_PRECONDITION,
			fmt.Sprintf("Missing required field(s) for tablet picker: %s", strings.Join(missingFields, ", ")))
	}
	return &TabletPicker{
		ts:          ts,
		cells:       cells,
		keyspace:    keyspace,
		shard:       shard,
		tabletTypes: tabletTypes,
	}, nil
}

// PickForStreaming picks an available tablet
// All tablets that belong to tp.cells are evaluated and one is
// chosen at random
func (tp *TabletPicker) PickForStreaming(ctx context.Context) (*topodatapb.Tablet, error) {
	// keep trying at intervals (tabletPickerRetryDelay) until a tablet is found
	// or the context is canceled
	for {
		select {
		case <-ctx.Done():
			return nil, vterrors.Errorf(vtrpcpb.Code_CANCELED, "context has expired")
		default:
		}
		candidates := tp.getMatchingTablets(ctx)
		if len(candidates) == 0 {
			// if no candidates were found, sleep and try again
			log.Infof("No tablet found for streaming, sleeping for %d seconds", int(GetTabletPickerRetryDelay()/1e9))
			time.Sleep(GetTabletPickerRetryDelay())
			continue
		}
		// try at most len(candidate) times to find a healthy tablet
		for i := 0; i < len(candidates); i++ {
			idx := rand.Intn(len(candidates))
			ti := candidates[idx]
			// get tablet
			// try to connect to tablet
			conn, err := tabletconn.GetDialer()(ti.Tablet, true)
			if err != nil {
				log.Warningf("unable to connect to tablet for alias %v", ti.Alias)
				candidates = append(candidates[:idx], candidates[idx+1:]...)
				if len(candidates) == 0 {
					break
				}
				continue
			}
			// OK to use ctx here because it is not actually used by the underlying Close implementation
			_ = conn.Close(ctx)
			return ti.Tablet, nil
		}
	}
}

// getMatchingTablets returns a list of TabletInfo for tablets
// that match the cells, keyspace, shard and tabletTypes for this TabletPicker
func (tp *TabletPicker) getMatchingTablets(ctx context.Context) []*topo.TabletInfo {
	// Special handling for MASTER tablet type
	// Since there is only one master, we ignore cell and find the master
	aliases := make([]*topodatapb.TabletAlias, 0)
	if len(tp.tabletTypes) == 1 && tp.tabletTypes[0] == topodatapb.TabletType_MASTER {
		shortCtx, cancel := context.WithTimeout(ctx, *topo.RemoteOperationTimeout)
		defer cancel()
		si, err := tp.ts.GetShard(shortCtx, tp.keyspace, tp.shard)
		if err != nil {
			return nil
		}
		aliases = append(aliases, si.MasterAlias)
	} else {
		actualCells := make([]string, 0)
		for _, cell := range tp.cells {
			// check if cell is actually an alias
			// non-blocking read so that this is fast
			shortCtx, cancel := context.WithTimeout(ctx, *topo.RemoteOperationTimeout)
			defer cancel()
			_, err := tp.ts.GetCellInfo(ctx, cell, false)
			if err != nil {
				// not a valid cell, check whether it is a cell alias
				alias, err := tp.ts.GetCellsAlias(shortCtx, cell, false)
				// if we get an error, either cellAlias doesn't exist or it isn't a cell alias at all. Ignore and continue
				if err == nil {
					actualCells = append(actualCells, alias.Cells...)
				}
			} else {
				// valid cell, add it to our list
				actualCells = append(actualCells, cell)
			}
		}
		for _, cell := range actualCells {
			shortCtx, cancel := context.WithTimeout(ctx, *topo.RemoteOperationTimeout)
			defer cancel()
			// match cell, keyspace and shard
			sri, err := tp.ts.GetShardReplication(shortCtx, cell, tp.keyspace, tp.shard)
			if err != nil {
				continue
			}

			for _, node := range sri.Nodes {
				aliases = append(aliases, node.TabletAlias)
			}
		}
	}

	if len(aliases) == 0 {
		return nil
	}
	shortCtx, cancel := context.WithTimeout(ctx, *topo.RemoteOperationTimeout)
	defer cancel()
	tabletMap, err := tp.ts.GetTabletMap(shortCtx, aliases)
	if err != nil {
		log.Warningf("error fetching tablets from topo: %v", err)
		// If we get a partial result we can still use it, otherwise return
		if len(tabletMap) == 0 {
			return nil
		}
	}
	tablets := make([]*topo.TabletInfo, 0, len(aliases))
	for _, tabletAlias := range aliases {
		tabletInfo, ok := tabletMap[topoproto.TabletAliasString(tabletAlias)]
		if !ok {
			// Either tablet disappeared on us, or we got a partial result (GetTabletMap ignores
			// topo.ErrNoNode). Just log a warning
			log.Warningf("failed to load tablet %v", tabletAlias)
		} else if topoproto.IsTypeInList(tabletInfo.Type, tp.tabletTypes) {
			tablets = append(tablets, tabletInfo)
		}
	}
	return tablets
}

func init() {
	// TODO(sougou): consolidate this call to be once per process.
	rand.Seed(time.Now().UnixNano())
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"os"

	"golang.org/x/net/html"
)

func main() {
	noLinks := flag.Bool("nolinks", false, "Do not print links.")
	noCount := flag.Bool("nocount", false, "Do not print element count.")
	flag.Parse()

	doc, err := html.Parse(os.Stdin)

	if err != nil {
		fmt.Fprintf(os.Stderr, "findlinks: %v\n", err)
		os.Exit(1)
	}
	if !*noLinks {
		for _, link := range collectLinks(nil, doc) {
			fmt.Println(link)
		}
	}
	if !*noCount {
		m := make(map[string]int)
		countElements(m, doc)
		for k, v := range m {
			fmt.Printf("%-10s: % 4d\n", k, v)
		}
	}
}

func collectLinks(links []string, n *html.Node) []string {
	if n == nil {
		return links
	}

	if n.Type == html.ElementNode && n.Data == "a" {
		for _, a := range n.Attr {
			if a.Key == "href" {
				links = append(links, a.Val)
			}
		}
	}
	links = collectLinks(links, n.FirstChild)
	return collectLinks(links, n.NextSibling)
}

func countElements(m map[string]int, n *html.Node) {
	if n.Type == html.ElementNode {
		m[n.Data]++
	}
	for c := n.FirstChild; c != nil; c = c.NextSibling {
		countElements(m, c)
	}
}
<commit_msg>[5.3] Solution added for showing text for nodes.<commit_after>package main

import (
	"flag"
	"fmt"
	"os"
	"strings"

	"golang.org/x/net/html"
)

func main() {
	noLinks := flag.Bool("nolinks", false, "Do not print links.")
	noCount := flag.Bool("nocount", false, "Do not print element count.")
	noText := flag.Bool("notext", false, "Do not print text of nodes.")
	flag.Parse()

	doc, err := html.Parse(os.Stdin)

	if err != nil {
		fmt.Fprintf(os.Stderr, "findlinks: %v\n", err)
		os.Exit(1)
	}
	if !*noLinks {
		for _, link := range collectLinks(nil, doc) {
			fmt.Println(link)
		}
	}
	if !*noCount {
		m := make(map[string]int)
		countElements(m, doc)
		for k, v := range m {
			fmt.Printf("%-10s: % 4d\n", k, v)
		}
	}
	if !*noText {
		for _, text := range collectText(nil, doc) {
			fmt.Println(text)
		}
	}
}

func collectLinks(links []string, n *html.Node) []string {
	if n == nil {
		return links
	}

	if n.Type == html.ElementNode && n.Data == "a" {
		for _, a := range n.Attr {
			if a.Key == "href" {
				links = append(links, a.Val)
			}
		}
	}
	links = collectLinks(links, n.FirstChild)
	return collectLinks(links, n.NextSibling)
}

func countElements(m map[string]int, n *html.Node) {
	if n.Type == html.ElementNode {
		m[n.Data]++
	}
	for c := n.FirstChild; c != nil; c = c.NextSibling {
		countElements(m, c)
	}
}

func collectText(txt []string, n *html.Node) []string {
	if n == nil {
		return txt
	}

	if n.Type == html.TextNode {
		if n.Parent.Data != "script" && n.Parent.Data != "style" {
			for _, line := range strings.Split(n.Data, "\n") {
				line = strings.TrimLeft(line, " \t")
				if len(line) == 0 {
					continue
				}
				txt = append(txt, line)
			}
		}
	}

	for c := n.FirstChild; c != nil; c = c.NextSibling {
		txt = collectText(txt, c)
	}
	return txt
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"github.com/bobtfish/sensuplugingo/handler"
)

func main() {
	event, err := handler.GetEventFromStdin()
	if err != nil {
		panic("Error getting Event")
	}
	fmt.Println(event)
}
<commit_msg>Basic structure for a handler which may do something<commit_after>package main

import (
	"fmt"
	"github.com/bobtfish/sensuplugingo/handler"
	"github.com/bobtfish/sensuplugingo"
)

func FormatMessage(e sensuplugingo.Event) string {
	return("foo")
}

func GetChannels() (c []string) {
	return c
}

func Send(channel string, message string) error {
	// system('/nail/sys/bin/nodebot', channel, message)
	return(nil)
}

func main() {
	event, err := handler.GetEventFromStdin()
	if err != nil {
		panic("Error getting Event")
	}
	message := FormatMessage(event)
	channels := GetChannels()
	for _, channel := range channels {
		Send(channel, message)
	}
	fmt.Println(event)
}
<|endoftext|>


<commit_before>package testhelper

import "testing"

// This function compares two arbitrary values and performs a comparison. If the
// comparison fails, a fatal error is raised that will fail the test
func Compare(t *testing.T, expected interface{}, actual interface{}) {
	if expected != actual {
		t.Fatalf("Expected [%#v] but got [%#v]", expected, actual)
	}
}

// A convenience function for checking whether an error value is an actual error
func CheckErr(t *testing.T, e error) {
	if e != nil {
		t.Fatalf("Unexpected error: %#v", e)
	}
}
<commit_msg>Adding deep compare<commit_after>package testhelper

import (
	"reflect"
	"testing"
)

// This function compares two arbitrary values and performs a comparison. If the
// comparison fails, a fatal error is raised that will fail the test
func Equals(t *testing.T, expected, actual interface{}) {
	if expected != actual {
		t.Fatalf("Expected [%#v] but got [%#v]", expected, actual)
	}
}

// This function, like Equals, performs a comparison - but on more complex
// structures that requires deeper inspection
func DeepEquals(t *testing.T, actual, expected interface{}) {
	if !reflect.DeepEqual(actual, expected) {
		t.Fatalf("Expected %#v but got %#v", expected, actual)
	}
}

// A convenience function for checking whether an error value is an actual error
func CheckErr(t *testing.T, e error) {
	if e != nil {
		t.Fatalf("Unexpected error: %#v", e)
	}
}
<|endoftext|>


<commit_before>// Requirements: smem (https://www.selenic.com/smem/).
//
// Config file: SmemStats.conf
// Example: {
//   "user": "some-user", <-- This user should be able to access the /proc/<pid>/smaps files listed in the whitelist below
//   "procsWhitelist": "apache2|tmux",
//   "smemPath": "/usr/bin/smem" <-- Path to the smem executable
// }

package collector

import (
	"fullerite/config"
	"fullerite/metric"
	"os/exec"
	"strconv"
	"strings"

	l "github.com/Sirupsen/logrus"
)

type smemStatLine struct {
	proc string
	pss  float64
	rss  float64
	vss  float64
}

// SmemStats Collector to record smem stats
type SmemStats struct {
	baseCollector
	user             string
	whitelistedProcs string
	smemPath         string
}

var (
	requiredConfigs = []string{"user", "procsWhitelist"}

	execCommand   = exec.Command
	commandOutput = (*exec.Cmd).Output
	getSmemStats  = (*SmemStats).getSmemStats
)

func init() {
	RegisterCollector("SmemStats", newSmemStats)
}

func newSmemStats(channel chan metric.Metric, initialInterval int, log *l.Entry) Collector {
	s := new(SmemStats)
	s.log = log
	s.channel = channel
	s.interval = initialInterval
	s.name = "SmemStats"

	return s
}

// Configure Override *baseCollector.Configure(); will fetch the whitelisted processes
func (s *SmemStats) Configure(configMap map[string]interface{}) {
	s.configureCommonParams(configMap)
	cfg := config.GetAsMap(configMap)

	if user, exists := cfg["user"]; exists {
		s.user = user
	} else {
		s.log.Warn("Required config does not exist for SmemStats: user")
	}

	if whitelist, exists := cfg["procsWhitelist"]; exists {
		s.whitelistedProcs = whitelist
	} else {
		s.log.Warn("Required config does not exist for SmemStats: procsWhitelist")
	}

	if smemPath, exists := cfg["smemPath"]; exists {
		s.smemPath = smemPath
	} else {
		s.log.Warn("Required config does not exist for SmemStats: smemPath")
	}
}

// Collect periodically call smem periodically
func (s *SmemStats) Collect() {
	if s.whitelistedProcs == "" || s.user == "" || s.smemPath == "" {
		return
	}

	for _, stat := range getSmemStats(s) {
		s.Channel() <- metric.WithValue(stat.proc+".smem.pss", stat.pss)
		s.Channel() <- metric.WithValue(stat.proc+".smem.vss", stat.vss)
		s.Channel() <- metric.WithValue(stat.proc+".smem.rss", stat.rss)
	}
}

func (s *SmemStats) getSmemStats() []smemStatLine {
	var out []byte
	var err error

	cmd := execCommand(
		"/usr/bin/sudo",
		"-u", s.user,
		s.smemPath,
		"-P", s.whitelistedProcs,
		"-c", "pss rss vss name")
	if out, err = commandOutput(cmd); err != nil {
		s.log.Error(err.Error())
		return nil
	}

	return s.parseSmemLines(string(out))
}

func (s *SmemStats) parseSmemLines(out string) []smemStatLine {
	raw := strings.Trim(out, "\n")
	lines := strings.Split(raw, "\n")
	stats := []smemStatLine{}

	for _, line := range lines {
		parts := strings.Fields(line)
		stats = append(stats, smemStatLine{
			proc: parts[3],
			pss:  strToFloat(parts[0]),
			rss:  strToFloat(parts[1]),
			vss:  strToFloat(parts[2]),
		})
	}

	return stats
}

func strToFloat(val string) float64 {
	if i, err := strconv.ParseFloat(val, 64); err == nil {
		return i
	}

	return 0
}
<commit_msg>Choose sending metrics from configuration<commit_after>// Requirements: smem (https://www.selenic.com/smem/).
//
// Config file: SmemStats.conf
// Example: {
//   "user": "some-user", <-- This user should be able to access the /proc/<pid>/smaps files listed in the whitelist below
//   "procsWhitelist": "apache2|tmux",
//   "smemPath": "/usr/bin/smem" <-- Path to the smem executable
// }

package collector

import (
	"fullerite/config"
	"fullerite/metric"
	"os/exec"
	"strconv"
	"strings"

	l "github.com/Sirupsen/logrus"
)

type smemStatLine struct {
	proc string
	pss  float64
	rss  float64
	vss  float64
}

// SmemStats Collector to record smem stats
type SmemStats struct {
	baseCollector
	user               string
	whitelistedProcs   string
	smemPath           string
	whitelistedMetrics []string
}

var (
	requiredConfigs = []string{"user", "procsWhitelist"}

	execCommand   = exec.Command
	commandOutput = (*exec.Cmd).Output
	getSmemStats  = (*SmemStats).getSmemStats
)

func init() {
	RegisterCollector("SmemStats", newSmemStats)
}

func newSmemStats(channel chan metric.Metric, initialInterval int, log *l.Entry) Collector {
	s := new(SmemStats)
	s.log = log
	s.channel = channel
	s.interval = initialInterval
	s.name = "SmemStats"

	return s
}

// Configure Override *baseCollector.Configure(); will fetch the whitelisted processes
func (s *SmemStats) Configure(configMap map[string]interface{}) {
	s.configureCommonParams(configMap)

	if user, exists := configMap["user"]; exists {
		s.user = user.(string)
	} else {
		s.log.Warn("Required config does not exist for SmemStats: user")
	}

	if whitelist, exists := configMap["procsWhitelist"]; exists {
		s.whitelistedProcs = whitelist.(string)
	} else {
		s.log.Warn("Required config does not exist for SmemStats: procsWhitelist")
	}

	if smemPath, exists := configMap["smemPath"]; exists {
		s.smemPath = smemPath.(string)
	} else {
		s.log.Warn("Required config does not exist for SmemStats: smemPath")
	}

	if whitelist, exists := configMap["metricsWhitelist"]; exists {
		s.whitelistedMetrics = config.GetAsSlice(whitelist)
	} else {
		s.log.Warn("Required config does not exist for SmemStats: metricsWhitelist")
	}
}

// Collect periodically call smem periodically
func (s *SmemStats) Collect() {
	if s.whitelistedProcs == "" || s.user == "" || s.smemPath == "" {
		return
	}

	for _, stat := range getSmemStats(s) {
		for _, element := range s.whitelistedMetrics {
			switch element {
			case "pss":
				s.Channel() <- metric.WithValue(stat.proc+".smem.pss", stat.pss)
			case "vss":
				s.Channel() <- metric.WithValue(stat.proc+".smem.vss", stat.vss)
			case "rss":
				s.Channel() <- metric.WithValue(stat.proc+".smem.rss", stat.rss)
			}
		}
	}
}

func (s *SmemStats) getSmemStats() []smemStatLine {
	var out []byte
	var err error

	cmd := execCommand(
		"/usr/bin/sudo",
		"-u", s.user,
		s.smemPath,
		"-P", s.whitelistedProcs,
		"-c", "pss rss vss name")
	if out, err = commandOutput(cmd); err != nil {
		s.log.Error(err.Error())
		return nil
	}

	return s.parseSmemLines(string(out))
}

func (s *SmemStats) parseSmemLines(out string) []smemStatLine {
	raw := strings.Trim(out, "\n")
	lines := strings.Split(raw, "\n")
	stats := []smemStatLine{}

	for _, line := range lines {
		parts := strings.Fields(line)
		stats = append(stats, smemStatLine{
			proc: parts[3],
			pss:  strToFloat(parts[0]),
			rss:  strToFloat(parts[1]),
			vss:  strToFloat(parts[2]),
		})
	}

	return stats
}

func strToFloat(val string) float64 {
	if i, err := strconv.ParseFloat(val, 64); err == nil {
		return i
	}

	return 0
}
<|endoftext|>


<commit_before>/*
Copyright 2022 The Kubernetes Authors.

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

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

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

package kubernetes

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"strings"
	"testing"
	"time"

	"github.com/stretchr/testify/require"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
	"k8s.io/apimachinery/pkg/types"
	"k8s.io/apimachinery/pkg/util/yaml"
	"sigs.k8s.io/controller-runtime/pkg/client"
)

// MustApplyWithCleanup creates or updates Kubernetes resources defined with the
// provided YAML file and registers a cleanup function for resources it created.
// Note that this does not remove resources that already existed in the cluster.
func MustApplyWithCleanup(t *testing.T, c client.Client, location string, gcName string, cleanup bool) {
	data, err := getContentsFromPathOrURL(location)
	require.NoError(t, err)

	decoder := yaml.NewYAMLOrJSONDecoder(data, 4096)
	for {
		uObj := unstructured.Unstructured{}
		if decodeErr := decoder.Decode(&uObj); decodeErr != nil {
			if errors.Is(decodeErr, io.EOF) {
				break
			}
			t.Logf("manifest: %s", data.String())
			require.NoErrorf(t, decodeErr, "error parsing manifest")
		}
		if len(uObj.Object) == 0 {
			continue
		}

		if uObj.GetKind() == "Gateway" {
			err = unstructured.SetNestedField(uObj.Object, gcName, "spec", "gatewayClassName")
			require.NoErrorf(t, err, "error setting `spec.gatewayClassName` on %s Gateway resource", uObj.GetName())
		}

		ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
		defer cancel()

		namespacedName := types.NamespacedName{Namespace: uObj.GetNamespace(), Name: uObj.GetName()}
		fetchedObj := uObj.DeepCopy()
		err := c.Get(ctx, namespacedName, fetchedObj)
		if err != nil {
			if !apierrors.IsNotFound(err) {
				require.NoErrorf(t, err, "error getting resource")
			}
			t.Logf("Creating %s %s", uObj.GetName(), uObj.GetKind())
			err = c.Create(ctx, &uObj)
			require.NoErrorf(t, err, "error creating resource")

			if cleanup {
				t.Cleanup(func() {
					err = c.Delete(ctx, &uObj)
					require.NoErrorf(t, err, "error deleting resource")
				})
			}
			continue
		}

		uObj.SetResourceVersion(fetchedObj.GetResourceVersion())
		t.Logf("Updating %s %s", uObj.GetName(), uObj.GetKind())
		err = c.Update(ctx, &uObj)

		if cleanup {
			t.Cleanup(func() {
				err = c.Delete(ctx, &uObj)
				require.NoErrorf(t, err, "error deleting resource")
			})
		}
		require.NoErrorf(t, err, "error updating resource")
	}
}

// getContentsFromPathOrURL takes a string that can either be a local file
// path or an https:// URL to YAML manifests and provides the contents.
func getContentsFromPathOrURL(location string) (*bytes.Buffer, error) {
	manifests := new(bytes.Buffer)
	if strings.HasPrefix(location, "http://") {
		return nil, fmt.Errorf("data can't be retrieved from %s: http is not supported, use https", location)
	} else if strings.HasPrefix(location, "https://") {
		ctx, cancel := context.WithTimeout(context.Background(), time.Second*10)
		defer cancel()

		req, err := http.NewRequestWithContext(ctx, http.MethodGet, location, nil)
		if err != nil {
			return nil, err
		}

		resp, err := http.DefaultClient.Do(req)
		if err != nil {
			return nil, err
		}
		defer resp.Body.Close()

		count, err := manifests.ReadFrom(resp.Body)
		if err != nil {
			return nil, err
		}

		if resp.ContentLength != -1 && count != resp.ContentLength {
			return nil, fmt.Errorf("received %d bytes from %s, expected %d", count, location, resp.ContentLength)
		}
	} else {
		b, err := ioutil.ReadFile(location)
		if err != nil {
			return nil, err
		}

		count, err := manifests.Read(b)
		if err != nil {
			return nil, err
		}

		if count != len(b) {
			return nil, fmt.Errorf("expected to write %d bytes to memory, only wrote %d", len(b), count)
		}
	}

	return manifests, nil
}
<commit_msg>conformance: fix reading local manifest files<commit_after>/*
Copyright 2022 The Kubernetes Authors.

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

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

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

package kubernetes

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"strings"
	"testing"
	"time"

	"github.com/stretchr/testify/require"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
	"k8s.io/apimachinery/pkg/types"
	"k8s.io/apimachinery/pkg/util/yaml"
	"sigs.k8s.io/controller-runtime/pkg/client"
)

// MustApplyWithCleanup creates or updates Kubernetes resources defined with the
// provided YAML file and registers a cleanup function for resources it created.
// Note that this does not remove resources that already existed in the cluster.
func MustApplyWithCleanup(t *testing.T, c client.Client, location string, gcName string, cleanup bool) {
	data, err := getContentsFromPathOrURL(location)
	require.NoError(t, err)

	decoder := yaml.NewYAMLOrJSONDecoder(data, 4096)
	for {
		uObj := unstructured.Unstructured{}
		if decodeErr := decoder.Decode(&uObj); decodeErr != nil {
			if errors.Is(decodeErr, io.EOF) {
				break
			}
			t.Logf("manifest: %s", data.String())
			require.NoErrorf(t, decodeErr, "error parsing manifest")
		}
		if len(uObj.Object) == 0 {
			continue
		}

		if uObj.GetKind() == "Gateway" {
			err = unstructured.SetNestedField(uObj.Object, gcName, "spec", "gatewayClassName")
			require.NoErrorf(t, err, "error setting `spec.gatewayClassName` on %s Gateway resource", uObj.GetName())
		}

		ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
		defer cancel()

		namespacedName := types.NamespacedName{Namespace: uObj.GetNamespace(), Name: uObj.GetName()}
		fetchedObj := uObj.DeepCopy()
		err := c.Get(ctx, namespacedName, fetchedObj)
		if err != nil {
			if !apierrors.IsNotFound(err) {
				require.NoErrorf(t, err, "error getting resource")
			}
			t.Logf("Creating %s %s", uObj.GetName(), uObj.GetKind())
			err = c.Create(ctx, &uObj)
			require.NoErrorf(t, err, "error creating resource")

			if cleanup {
				t.Cleanup(func() {
					err = c.Delete(ctx, &uObj)
					require.NoErrorf(t, err, "error deleting resource")
				})
			}
			continue
		}

		uObj.SetResourceVersion(fetchedObj.GetResourceVersion())
		t.Logf("Updating %s %s", uObj.GetName(), uObj.GetKind())
		err = c.Update(ctx, &uObj)

		if cleanup {
			t.Cleanup(func() {
				err = c.Delete(ctx, &uObj)
				require.NoErrorf(t, err, "error deleting resource")
			})
		}
		require.NoErrorf(t, err, "error updating resource")
	}
}

// getContentsFromPathOrURL takes a string that can either be a local file
// path or an https:// URL to YAML manifests and provides the contents.
func getContentsFromPathOrURL(location string) (*bytes.Buffer, error) {
	manifests := new(bytes.Buffer)
	if strings.HasPrefix(location, "http://") {
		return nil, fmt.Errorf("data can't be retrieved from %s: http is not supported, use https", location)
	} else if strings.HasPrefix(location, "https://") {
		ctx, cancel := context.WithTimeout(context.Background(), time.Second*10)
		defer cancel()

		req, err := http.NewRequestWithContext(ctx, http.MethodGet, location, nil)
		if err != nil {
			return nil, err
		}

		resp, err := http.DefaultClient.Do(req)
		if err != nil {
			return nil, err
		}
		defer resp.Body.Close()

		count, err := manifests.ReadFrom(resp.Body)
		if err != nil {
			return nil, err
		}

		if resp.ContentLength != -1 && count != resp.ContentLength {
			return nil, fmt.Errorf("received %d bytes from %s, expected %d", count, location, resp.ContentLength)
		}
	} else {
		b, err := ioutil.ReadFile(location)
		if err != nil {
			return nil, err
		}

		count, err := manifests.Write(b)
		if err != nil {
			return nil, err
		}

		if count != len(b) {
			return nil, fmt.Errorf("expected to write %d bytes to memory, only wrote %d", len(b), count)
		}
	}

	return manifests, nil
}
<|endoftext|>


<commit_before>package bubbles

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"sync"
	"time"
)

const (
	// DefaultMaxDocumentsPerBatch is the number of documents a batch needs to
	// have before it is send. This is per connection.
	DefaultMaxDocumentsPerBatch = 1000

	// DefaultFlushTimeout is the maximum time we batch something before we try
	// to send it to a server.
	DefaultFlushTimeout = 10 * time.Second

	// DefaultServerTimeout is the time we give ElasticSearch to respond. This
	// is also the maximum time Stop() will take.
	DefaultServerTimeout = 10 * time.Second

	// DefaultConnCount is the number of connections per hosts.
	DefaultConnCount = 2

	serverErrorWait = 500 * time.Millisecond
	serverErrorWaitMax = 16 * time.Second

	defaultElasticSearchPort = "9200"
)

var (
	errInvalidResponse = errors.New("invalid response")
)

// Bubbles is the main struct to control a queue of Actions going to the
// ElasticSearch servers.
type Bubbles struct {
	q                chan Action
	retryQ           chan Action
	quit             chan struct{}
	wg               sync.WaitGroup
	maxDocumentCount int
	connCount        int
	flushTimeout     time.Duration
	serverTimeout    time.Duration
	e                Errer
}

// Opt is any option to New().
type Opt func(*Bubbles)

// OptConnCount is an option to New() to specify the number of connections per
// host. The default is DefaultConnCount.
func OptConnCount(n int) Opt {
	return func(b *Bubbles) {
		b.connCount = n
	}
}

// OptFlush is an option to New() to specify the flush timeout of a batch. The
// default is DefaultFlushTimeout.
func OptFlush(d time.Duration) Opt {
	return func(b *Bubbles) {
		b.flushTimeout = d
	}
}

// OptServerTimeout is an option to New() to specify the timeout of a single
// batch POST to ElasticSearch. This value is also the maximum time Stop() will
// take. All actions in a bulk which is timed out will be retried. The default
// is DefaultServerTimeout.
func OptServerTimeout(d time.Duration) Opt {
	return func(b *Bubbles) {
		b.serverTimeout = d
	}
}

// OptMaxDocs is an option to New() to specify maximum number of documents in a
// single batch. The default is DefaultMaxDocumentsPerBatch.
func OptMaxDocs(n int) Opt {
	return func(b *Bubbles) {
		b.maxDocumentCount = n
	}
}

// OptErrer is an option to New() to specify an error handler. The default
// handler uses the log module.
func OptErrer(e Errer) Opt {
	return func(b *Bubbles) {
		b.e = e
	}
}

// New makes a new ElasticSearch bulk inserter. It needs a list with 'ip' or
// 'ip:port' addresses, options are added via the Opt* functions.
func New(addrs []string, opts ...Opt) *Bubbles {
	b := Bubbles{
		q:                make(chan Action),
		quit:             make(chan struct{}),
		maxDocumentCount: DefaultMaxDocumentsPerBatch,
		connCount:        DefaultConnCount,
		flushTimeout:     DefaultFlushTimeout,
		serverTimeout:    DefaultServerTimeout,
		e:                DefaultErrer{},
	}
	for _, o := range opts {
		o(&b)
	}
	b.retryQ = make(chan Action, len(addrs)*b.connCount*b.maxDocumentCount)

	// Start a go routine per connection per host
	for _, a := range addrs {
		addr := withPort(a, defaultElasticSearchPort)
		for i := 0; i < b.connCount; i++ {
			b.wg.Add(1)
			go func(a string) {
				client(&b, a)
				b.wg.Done()
			}(addr)
		}
	}
	return &b
}

// Enqueue returns the queue to add Actions in a routine. It will block if all bulk
// processors are busy.
func (b *Bubbles) Enqueue() chan<- Action {
	return b.q
}

// Stop shuts down all ElasticSearch clients. It'll return all Action entries
// which were not yet processed, or were up for a retry. It can take
// OptServerTimeout to complete.
func (b *Bubbles) Stop() []Action {
	close(b.quit)
	// There is no explicit timeout, we rely on b.serverTimeout to shut down
	// everything.
	b.wg.Wait()

	// Collect and return elements which are in flight.
	close(b.retryQ)
	close(b.q)
	pending := make([]Action, 0, len(b.q)+len(b.retryQ))
	for a := range b.q {
		pending = append(pending, a)
	}
	for a := range b.retryQ {
		pending = append(pending, a)
	}
	return pending
}

// client talks to ElasticSearch. This runs in a go routine in a loop and deals
// with a single ElasticSearch address.
func client(b *Bubbles, addr string) {
	url := fmt.Sprintf("http://%s/_bulk", addr)

	cl := http.Client{
		Timeout: b.serverTimeout,
		CheckRedirect: func(req *http.Request, via []*http.Request) error {
			return errors.New("no redirect")
		},
	}

	wait := serverErrorWait
	for {
		select {
		case <-b.quit:
			return
		default:
		}
		if runBatch(b, cl, url) {
			// Some error, back off.
			select {
			case <-b.quit:
				return
			case <-time.After(wait):
				wait *= 2
				if wait >= serverErrorWaitMax {
					wait = serverErrorWaitMax
				}
			}
			continue
		}
		wait = serverErrorWait
	}
}

// runBatch gathers and deals with a batch of actions. It returns
// whether there was an error.
func runBatch(b *Bubbles, cl http.Client, url string) bool {
	actions := make([]Action, 0, b.maxDocumentCount)
	// First use all retry actions.
retry:
	for len(actions) < b.maxDocumentCount {
		select {
		case a := <-b.retryQ:
			actions = append(actions, a)
		default:
			// no more retry actions queued
			break retry
		}
	}

	var t <-chan time.Time
gather:
	for len(actions) < b.maxDocumentCount {
		if t == nil && len(actions) > 0 {
			// Set timeout on the first element we read
			t = time.After(b.flushTimeout)
		}
		select {
		case <-b.quit:
			for _, a := range actions {
				b.retryQ <- a
			}
			return false
		case <-t:
			// this case is not enabled until we've got an action
			break gather
		case a := <-b.retryQ:
			actions = append(actions, a)
		case a := <-b.q:
			actions = append(actions, a)
		}
	}

	res, err := postActions(cl, url, actions)
	if err != nil {
		// A server error. Retry these actions later.
		b.e.Error(err)
		for _, a := range actions {
			b.retryQ <- a
		}
		return true
	}

	// Server has accepted the request an sich, but there can be errors in the
	// individual actions.
	if !res.Errors {
		// Simple case, no errors present.
		return false
	}

	// Invalid response from ElasticSearch.
	if len(actions) != len(res.Items) {
		b.e.Error(errInvalidResponse)
		for _, a := range actions {
			b.retryQ <- a
		}
		return true
	}
	// Figure out which actions have errors.
	for i, e := range res.Items {
		a := actions[i]
		el, ok := e[string(a.Type)]
		if !ok {
			// Unexpected reply from ElasticSearch.
			b.e.Error(errInvalidResponse)
			b.retryQ <- a
			continue
		}

		c := el.Status
		switch {
		case c >= 200 && c < 300:
			// Document accepted by ElasticSearch.
		case c == 429 || (c >= 500 && c < 600):
			// Server error. Retry it.
			// We get a 429 when the bulk queue is full, which we just retry as
			// well.
			b.e.Warning(ActionError{
				Action: a,
				Msg:    fmt.Sprintf("transient error %d: %s", c, el.Error),
				Server: url,
			})
			b.retryQ <- a
		case c >= 400 && c < 500:
			// Some error. Nothing we can do with it.
			b.e.Error(ActionError{
				Action: a,
				Msg:    fmt.Sprintf("error %d: %s", c, el.Error),
				Server: url,
			})
		default:
			// No idea.
			b.e.Error(fmt.Errorf("unwelcome response %d: %s", c, el.Error))
		}
	}
	return true
}

type bulkRes struct {
	Took   int  `json:"took"`
	Errors bool `json:"errors"`
	Items  []map[string]struct {
		Index   string `json:"_index"`
		Type    string `json:"_type"`
		ID      string `json:"_id"`
		Version int    `json:"_version"`
		Status  int    `json:"status"`
		Error   string `json:"error"`
	} `json:"items"`
}

func postActions(cl http.Client, url string, actions []Action) (*bulkRes, error) {
	buf := bytes.Buffer{}
	for _, a := range actions {
		buf.Write(a.Buf())
	}

	// This doesn't Chunk.
	resp, err := cl.Post(url, "application/json", &buf)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	if resp.StatusCode != 200 {
		return nil, fmt.Errorf("status %d: %s", resp.StatusCode, string(body))
	}

	var bulk bulkRes
	if err := json.Unmarshal(body, &bulk); err != nil {
		return nil, err
	}
	return &bulk, nil
}

// ActionError wraps an Action we won't retry. It implements the error interface.
type ActionError struct {
	Action Action
	Msg    string
	Server string
}

func (e ActionError) Error() string {
	return fmt.Sprintf("%s: %s %s", e.Server, e.Action.Type, e.Msg)
}

// withPort adds a default port to an address string.
func withPort(a, port string) string {
	if _, _, err := net.SplitHostPort(a); err != nil {
		// no port found.
		return net.JoinHostPort(a, port)
	}
	return a
}
<commit_msg>Add status code to ActionError<commit_after>package bubbles

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"sync"
	"time"
)

const (
	// DefaultMaxDocumentsPerBatch is the number of documents a batch needs to
	// have before it is send. This is per connection.
	DefaultMaxDocumentsPerBatch = 1000

	// DefaultFlushTimeout is the maximum time we batch something before we try
	// to send it to a server.
	DefaultFlushTimeout = 10 * time.Second

	// DefaultServerTimeout is the time we give ElasticSearch to respond. This
	// is also the maximum time Stop() will take.
	DefaultServerTimeout = 10 * time.Second

	// DefaultConnCount is the number of connections per hosts.
	DefaultConnCount = 2

	serverErrorWait    = 500 * time.Millisecond
	serverErrorWaitMax = 16 * time.Second

	defaultElasticSearchPort = "9200"
)

var (
	errInvalidResponse = errors.New("invalid response")
)

// Bubbles is the main struct to control a queue of Actions going to the
// ElasticSearch servers.
type Bubbles struct {
	q                chan Action
	retryQ           chan Action
	quit             chan struct{}
	wg               sync.WaitGroup
	maxDocumentCount int
	connCount        int
	flushTimeout     time.Duration
	serverTimeout    time.Duration
	e                Errer
}

// Opt is any option to New().
type Opt func(*Bubbles)

// OptConnCount is an option to New() to specify the number of connections per
// host. The default is DefaultConnCount.
func OptConnCount(n int) Opt {
	return func(b *Bubbles) {
		b.connCount = n
	}
}

// OptFlush is an option to New() to specify the flush timeout of a batch. The
// default is DefaultFlushTimeout.
func OptFlush(d time.Duration) Opt {
	return func(b *Bubbles) {
		b.flushTimeout = d
	}
}

// OptServerTimeout is an option to New() to specify the timeout of a single
// batch POST to ElasticSearch. This value is also the maximum time Stop() will
// take. All actions in a bulk which is timed out will be retried. The default
// is DefaultServerTimeout.
func OptServerTimeout(d time.Duration) Opt {
	return func(b *Bubbles) {
		b.serverTimeout = d
	}
}

// OptMaxDocs is an option to New() to specify maximum number of documents in a
// single batch. The default is DefaultMaxDocumentsPerBatch.
func OptMaxDocs(n int) Opt {
	return func(b *Bubbles) {
		b.maxDocumentCount = n
	}
}

// OptErrer is an option to New() to specify an error handler. The default
// handler uses the log module.
func OptErrer(e Errer) Opt {
	return func(b *Bubbles) {
		b.e = e
	}
}

// New makes a new ElasticSearch bulk inserter. It needs a list with 'ip' or
// 'ip:port' addresses, options are added via the Opt* functions.
func New(addrs []string, opts ...Opt) *Bubbles {
	b := Bubbles{
		q:                make(chan Action),
		quit:             make(chan struct{}),
		maxDocumentCount: DefaultMaxDocumentsPerBatch,
		connCount:        DefaultConnCount,
		flushTimeout:     DefaultFlushTimeout,
		serverTimeout:    DefaultServerTimeout,
		e:                DefaultErrer{},
	}
	for _, o := range opts {
		o(&b)
	}
	b.retryQ = make(chan Action, len(addrs)*b.connCount*b.maxDocumentCount)

	// Start a go routine per connection per host
	for _, a := range addrs {
		addr := withPort(a, defaultElasticSearchPort)
		for i := 0; i < b.connCount; i++ {
			b.wg.Add(1)
			go func(a string) {
				client(&b, a)
				b.wg.Done()
			}(addr)
		}
	}
	return &b
}

// Enqueue returns the queue to add Actions in a routine. It will block if all bulk
// processors are busy.
func (b *Bubbles) Enqueue() chan<- Action {
	return b.q
}

// Stop shuts down all ElasticSearch clients. It'll return all Action entries
// which were not yet processed, or were up for a retry. It can take
// OptServerTimeout to complete.
func (b *Bubbles) Stop() []Action {
	close(b.quit)
	// There is no explicit timeout, we rely on b.serverTimeout to shut down
	// everything.
	b.wg.Wait()

	// Collect and return elements which are in flight.
	close(b.retryQ)
	close(b.q)
	pending := make([]Action, 0, len(b.q)+len(b.retryQ))
	for a := range b.q {
		pending = append(pending, a)
	}
	for a := range b.retryQ {
		pending = append(pending, a)
	}
	return pending
}

// client talks to ElasticSearch. This runs in a go routine in a loop and deals
// with a single ElasticSearch address.
func client(b *Bubbles, addr string) {
	url := fmt.Sprintf("http://%s/_bulk", addr)

	cl := http.Client{
		Timeout: b.serverTimeout,
		CheckRedirect: func(req *http.Request, via []*http.Request) error {
			return errors.New("no redirect")
		},
	}

	wait := serverErrorWait
	for {
		select {
		case <-b.quit:
			return
		default:
		}
		if runBatch(b, cl, url) {
			// Some error, back off.
			select {
			case <-b.quit:
				return
			case <-time.After(wait):
				wait *= 2
				if wait >= serverErrorWaitMax {
					wait = serverErrorWaitMax
				}
			}
			continue
		}
		wait = serverErrorWait
	}
}

// runBatch gathers and deals with a batch of actions. It returns
// whether there was an error.
func runBatch(b *Bubbles, cl http.Client, url string) bool {
	actions := make([]Action, 0, b.maxDocumentCount)
	// First use all retry actions.
retry:
	for len(actions) < b.maxDocumentCount {
		select {
		case a := <-b.retryQ:
			actions = append(actions, a)
		default:
			// no more retry actions queued
			break retry
		}
	}

	var t <-chan time.Time
gather:
	for len(actions) < b.maxDocumentCount {
		if t == nil && len(actions) > 0 {
			// Set timeout on the first element we read
			t = time.After(b.flushTimeout)
		}
		select {
		case <-b.quit:
			for _, a := range actions {
				b.retryQ <- a
			}
			return false
		case <-t:
			// this case is not enabled until we've got an action
			break gather
		case a := <-b.retryQ:
			actions = append(actions, a)
		case a := <-b.q:
			actions = append(actions, a)
		}
	}

	res, err := postActions(cl, url, actions)
	if err != nil {
		// A server error. Retry these actions later.
		b.e.Error(err)
		for _, a := range actions {
			b.retryQ <- a
		}
		return true
	}

	// Server has accepted the request an sich, but there can be errors in the
	// individual actions.
	if !res.Errors {
		// Simple case, no errors present.
		return false
	}

	// Invalid response from ElasticSearch.
	if len(actions) != len(res.Items) {
		b.e.Error(errInvalidResponse)
		for _, a := range actions {
			b.retryQ <- a
		}
		return true
	}
	// Figure out which actions have errors.
	for i, e := range res.Items {
		a := actions[i]
		el, ok := e[string(a.Type)]
		if !ok {
			// Unexpected reply from ElasticSearch.
			b.e.Error(errInvalidResponse)
			b.retryQ <- a
			continue
		}

		c := el.Status
		switch {
		case c >= 200 && c < 300:
			// Document accepted by ElasticSearch.
		case c == 429 || (c >= 500 && c < 600):
			// Server error. Retry it.
			// We get a 429 when the bulk queue is full, which we just retry as
			// well.
			b.e.Warning(ActionError{
				Action:     a,
				StatusCode: c,
				Msg:        fmt.Sprintf("transient error %d: %s", c, el.Error),
				Server:     url,
			})
			b.retryQ <- a
		case c >= 400 && c < 500:
			// Some error. Nothing we can do with it.
			b.e.Error(ActionError{
				Action:     a,
				StatusCode: c,
				Msg:        fmt.Sprintf("error %d: %s", c, el.Error),
				Server:     url,
			})
		default:
			// No idea.
			b.e.Error(fmt.Errorf("unwelcome response %d: %s", c, el.Error))
		}
	}
	return true
}

type bulkRes struct {
	Took   int  `json:"took"`
	Errors bool `json:"errors"`
	Items  []map[string]struct {
		Index   string `json:"_index"`
		Type    string `json:"_type"`
		ID      string `json:"_id"`
		Version int    `json:"_version"`
		Status  int    `json:"status"`
		Error   string `json:"error"`
	} `json:"items"`
}

func postActions(cl http.Client, url string, actions []Action) (*bulkRes, error) {
	buf := bytes.Buffer{}
	for _, a := range actions {
		buf.Write(a.Buf())
	}

	// This doesn't Chunk.
	resp, err := cl.Post(url, "application/json", &buf)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	if resp.StatusCode != 200 {
		return nil, fmt.Errorf("status %d: %s", resp.StatusCode, string(body))
	}

	var bulk bulkRes
	if err := json.Unmarshal(body, &bulk); err != nil {
		return nil, err
	}
	return &bulk, nil
}

// ActionError wraps an Action we won't retry. It implements the error interface.
type ActionError struct {
	Action     Action
	StatusCode int
	Msg        string
	Server     string
}

func (e ActionError) Error() string {
	return fmt.Sprintf("%s: %s %s", e.Server, e.Action.Type, e.Msg)
}

// withPort adds a default port to an address string.
func withPort(a, port string) string {
	if _, _, err := net.SplitHostPort(a); err != nil {
		// no port found.
		return net.JoinHostPort(a, port)
	}
	return a
}
<|endoftext|>


<commit_before>package cobra

import (
	"fmt"
	"io"
	"os"
	"strings"
	"text/template"

	"github.com/spf13/pflag"
)

var (
	funcMap = template.FuncMap{
		"constructPath": constructPath,
		"subCmdList":    subCmdList,
		"extractFlags":  extractFlags,
		"cmdName":       cmdName,
		"simpleFlag":    simpleFlag,
	}
	zshCompletionText = `
{{/* for pflag.Flag (specifically annotations) */}}
{{define "flagAnnotations" -}}
{{with index .Annotations "cobra_annotation_bash_completion_filename_extensions"}}:filename:_files{{end}}
{{- end}}

{{/* for pflag.Flag with short and long options */}}
{{define "complexFlag" -}}
"(-{{.Shorthand}} --{{.Name}})"{-{{.Shorthand}},--{{.Name}}}"[{{.Usage}}]{{template "flagAnnotations" .}}"
{{- end}}

{{/* for pflag.Flag with either short or long options */}}
{{define "simpleFlag" -}}
"{{with .Name}}--{{.}}{{else}}-{{.Shorthand}}{{end}}[{{.Usage}}]{{template "flagAnnotations" .}}"
{{- end}}

{{/* should accept Command (that contains subcommands) as parameter */}}
{{define "argumentsC" -}}
function {{constructPath .}} {
  local -a commands

  _arguments -C \{{- range extractFlags .}}
    {{if simpleFlag .}}{{template "simpleFlag" .}}{{else}}{{template "complexFlag" .}}{{end}} \{{- end}}
    "1: :->cmnds" \
    "*::arg:->args"

  case $state in
  cmnds)
    commands=({{range .Commands}}{{if not .Hidden}}
      "{{cmdName .}}:{{.Short}}"{{end}}{{end}}
    )
    _describe "command" commands
    ;;
  esac

  case "$words[1]" in {{- range .Commands}}{{if not .Hidden}}
  {{cmdName .}})
    {{constructPath .}}
    ;;{{end}}{{end}}
  esac
}
{{range .Commands}}{{if not .Hidden}}
{{template "selectCmdTemplate" .}}
{{- end}}{{end}}
{{- end}}

{{/* should accept Command without subcommands as parameter */}}
{{define "arguments" -}}
function {{constructPath .}} {
{{"  _arguments"}}{{range extractFlags .}} \
    {{if simpleFlag .}}{{template "simpleFlag" .}}{{else}}{{template "complexFlag" .}}{{end -}}
{{end}}
}
{{end}}

{{/* dispatcher for commands with or without subcommands */}}
{{define "selectCmdTemplate" -}}
{{if .Hidden}}{{/* ignore hidden*/}}{{else -}}
{{if .Commands}}{{template "argumentsC" .}}{{else}}{{template "arguments" .}}{{end}}
{{- end}}
{{- end}}

{{/* template entry point */}}
{{define "Main" -}}
#compdef _{{cmdName .}} {{cmdName .}}

{{template "selectCmdTemplate" .}}
{{end}}
`
)

// GenZshCompletionFile generates zsh completion file.
func (c *Command) GenZshCompletionFile(filename string) error {
	outFile, err := os.Create(filename)
	if err != nil {
		return err
	}
	defer outFile.Close()

	return c.GenZshCompletion(outFile)
}

// GenZshCompletion generates a zsh completion file and writes to the passed writer.
func (c *Command) GenZshCompletion(w io.Writer) error {
	tmpl, err := template.New("Main").Funcs(funcMap).Parse(zshCompletionText)
	if err != nil {
		return fmt.Errorf("error creating zsh completion template: %v", err)
	}
	return tmpl.Execute(w, c)
}

func constructPath(c *Command) string {
	var path []string
	tmpCmd := c
	path = append(path, tmpCmd.Name())

	for {
		if !tmpCmd.HasParent() {
			break
		}
		tmpCmd = tmpCmd.Parent()
		path = append(path, tmpCmd.Name())
	}

	// reverse path
	for left, right := 0, len(path)-1; left < right; left, right = left+1, right-1 {
		path[left], path[right] = path[right], path[left]
	}

	return "_" + strings.Join(path, "_")
}

// subCmdList returns a space separated list of subcommands names
func subCmdList(c *Command) string {
	var subCmds []string

	for _, cmd := range c.Commands() {
		if cmd.Hidden {
			continue
		}
		subCmds = append(subCmds, cmd.Name())
	}

	return strings.Join(subCmds, " ")
}

func extractFlags(c *Command) []*pflag.Flag {
	var flags []*pflag.Flag
	c.LocalFlags().VisitAll(func(f *pflag.Flag) {
		if !f.Hidden {
			flags = append(flags, f)
		}
	})
	c.InheritedFlags().VisitAll(func(f *pflag.Flag) {
		if !f.Hidden {
			flags = append(flags, f)
		}
	})
	return flags
}

// cmdName returns the command's innvocation
func cmdName(c *Command) string {
	return c.Name()
}

func simpleFlag(p *pflag.Flag) bool {
	return p.Name == "" || p.Shorthand == ""
}
<commit_msg>zsh-completion template refactoring:<commit_after>package cobra

import (
	"fmt"
	"io"
	"os"
	"text/template"

	"github.com/spf13/pflag"
)

var (
	funcMap = template.FuncMap{
		"genZshFuncName": generateZshCompletionFuncName,
		"extractFlags":   extractFlags,
		"simpleFlag":     simpleFlag,
	}
	zshCompletionText = `
{{/* for pflag.Flag (specifically annotations) */}}
{{define "flagAnnotations" -}}
{{with index .Annotations "cobra_annotation_bash_completion_filename_extensions"}}:filename:_files{{end}}
{{- end}}

{{/* for pflag.Flag with short and long options */}}
{{define "complexFlag" -}}
"(-{{.Shorthand}} --{{.Name}})"{-{{.Shorthand}},--{{.Name}}}"[{{.Usage}}]{{template "flagAnnotations" .}}"
{{- end}}

{{/* for pflag.Flag with either short or long options */}}
{{define "simpleFlag" -}}
"{{with .Name}}--{{.}}{{else}}-{{.Shorthand}}{{end}}[{{.Usage}}]{{template "flagAnnotations" .}}"
{{- end}}

{{/* should accept Command (that contains subcommands) as parameter */}}
{{define "argumentsC" -}}
{{ $cmdPath := genZshFuncName .}}
function {{$cmdPath}} {
  local -a commands

  _arguments -C \{{- range extractFlags .}}
    {{if simpleFlag .}}{{template "simpleFlag" .}}{{else}}{{template "complexFlag" .}}{{end}} \{{- end}}
    "1: :->cmnds" \
    "*::arg:->args"

  case $state in
  cmnds)
    commands=({{range .Commands}}{{if not .Hidden}}
      "{{.Name}}:{{.Short}}"{{end}}{{end}}
    )
    _describe "command" commands
    ;;
  esac

  case "$words[1]" in {{- range .Commands}}{{if not .Hidden}}
  {{.Name}})
    {{$cmdPath}}_{{.Name}}
    ;;{{end}}{{end}}
  esac
}
{{range .Commands}}{{if not .Hidden}}
{{template "selectCmdTemplate" .}}
{{- end}}{{end}}
{{- end}}

{{/* should accept Command without subcommands as parameter */}}
{{define "arguments" -}}
function {{genZshFuncName .}} {
{{"  _arguments"}}{{range extractFlags .}} \
    {{if simpleFlag .}}{{template "simpleFlag" .}}{{else}}{{template "complexFlag" .}}{{end -}}
{{end}}
}
{{end}}

{{/* dispatcher for commands with or without subcommands */}}
{{define "selectCmdTemplate" -}}
{{if .Hidden}}{{/* ignore hidden*/}}{{else -}}
{{if .Commands}}{{template "argumentsC" .}}{{else}}{{template "arguments" .}}{{end}}
{{- end}}
{{- end}}

{{/* template entry point */}}
{{define "Main" -}}
#compdef _{{.Name}} {{.Name}}

{{template "selectCmdTemplate" .}}
{{end}}
`
)

// GenZshCompletionFile generates zsh completion file.
func (c *Command) GenZshCompletionFile(filename string) error {
	outFile, err := os.Create(filename)
	if err != nil {
		return err
	}
	defer outFile.Close()

	return c.GenZshCompletion(outFile)
}

// GenZshCompletion generates a zsh completion file and writes to the passed writer.
func (c *Command) GenZshCompletion(w io.Writer) error {
	tmpl, err := template.New("Main").Funcs(funcMap).Parse(zshCompletionText)
	if err != nil {
		return fmt.Errorf("error creating zsh completion template: %v", err)
	}
	return tmpl.Execute(w, c)
}

func generateZshCompletionFuncName(c *Command) string {
	if c.HasParent() {
		return generateZshCompletionFuncName(c.Parent()) + "_" + c.Name()
	}
	return "_" + c.Name()
}

func extractFlags(c *Command) []*pflag.Flag {
	var flags []*pflag.Flag
	c.LocalFlags().VisitAll(func(f *pflag.Flag) {
		if !f.Hidden {
			flags = append(flags, f)
		}
	})
	c.InheritedFlags().VisitAll(func(f *pflag.Flag) {
		if !f.Hidden {
			flags = append(flags, f)
		}
	})
	return flags
}

func simpleFlag(p *pflag.Flag) bool {
	return p.Name == "" || p.Shorthand == ""
}
<|endoftext|>