<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"net/http"
	"os"
	"os/exec"
	"strings"
	"time"

	"github.com/manifoldco/promptui"
)

func main() {
	fmt.Println("RPNow Admin Console")

	for {
		up, pid := isServerUp()

		var options []string
		if up {
			options = []string{"manage rooms", "stop server", "exit"}
		} else {
			options = []string{"start server", "exit"}
		}

		prompt := promptui.Select{
			Label: "RP Admin",
			Items: options,
		}
		_, action, err := prompt.Run()

		if err != nil {
			return
		}
		switch action {
		case "start server":
			server := exec.Command("rpnow")
			out, err := server.CombinedOutput()
			if err != nil {
				fmt.Printf("Error starting server: %s\n", err)
			}
			fmt.Printf("%s\n", out)
			fmt.Println("Server started")
		case "stop server":
			process, err := os.FindProcess(pid)
			if err != nil {
				fmt.Printf("Error finding process with pid %d: %s\n", pid, err)
				continue
			}
			err = process.Kill()
			if err != nil {
				fmt.Printf("Error killing server process: %s\n", err)
				continue
			}
			fmt.Println("Server stopped")
		case "manage rooms":
			editRps()
		case "exit":
			return
		}
	}
}

func editRps() {
	for {
		// Select an RP from the list
		rps, err := getRpList()
		if err != nil {
			panic(err)
		}

		rp := pickRp(rps)
		if rp == nil {
			return
		}
		for {
			// Expand & edit the selected RP
			urls, err := getRpUrls(rp.RPID)
			if err != nil {
				panic(err)
			}

			fmt.Println()
			fmt.Println(rp.Title)
			for _, url := range urls {
				fmt.Printf("*  %s\n", url.String())
			}

			prompt := promptui.Select{
				Label: fmt.Sprintf("Modify %q", rp.Title),
				Items: []string{"go back", "edit urls", "destroy rp"},
			}
			_, action, err := prompt.Run()
			if err != nil {
				panic(err)
			}
			if action == "go back" {
				break
			} else if action == "edit urls" {

			} else if action == "destroy rp" {
				killswitch := strings.ToUpper(fmt.Sprintf("destroy %s", rp.Title))
				prompt := promptui.Prompt{Label: fmt.Sprintf("Type %q", killswitch)}
				result, _ := prompt.Run()
				if result == killswitch {
					err := destroyRp(rp.RPID)
					if err != nil {
						panic(err)
					}
					fmt.Printf("BOOM! %q is no more.\n", rp.Title)
					break
				} else {
					fmt.Println("Incorrect. Will not delete.")
				}
			}
		}
	}
}

func isServerUp() (bool, int) {
	fmt.Print("Getting server status... ")

	res, err := http.Get("http://127.0.0.1:12789/status")
	if err != nil {
		fmt.Println("(server not running)")
		return false, 0
	}
	defer res.Body.Close()
	var status struct {
		RPNowLine string `json:"rpnow"`
		PID       int    `json:"pid"`
	}
	err = json.NewDecoder(res.Body).Decode(&status)
	if err != nil {
		fmt.Println("(bad response)")
		return false, 0
	}
	fmt.Println(status.RPNowLine)
	return status.RPNowLine == "ok", status.PID
}

func pickRp(rps []rpInfo) *rpInfo {
	fmt.Println()

	prompt := promptui.Select{
		Label: "Choose an RP",
		Items: rps,
		Searcher: func(input string, index int) bool {
			return strings.Contains(strings.ToLower(rps[index].Title), strings.ToLower(input))
		},
	}
	idx, _, err := prompt.Run()
	if err != nil {
		return nil
	}

	return &rps[idx]
}

type rpInfo struct {
	Title     string    `json:"title"`
	RPID      string    `json:"rpid"`
	Timestamp time.Time `json:"timestamp"`
}

func (r rpInfo) String() string {
	return fmt.Sprintf("%-30s (%s)", r.Title, r.Timestamp.Format("02 Jan 2006"))
}

func getRpList() ([]rpInfo, error) {
	res, err := http.Get("http://127.0.0.1:12789/rps")
	if err != nil {
		return nil, err
	}
	defer res.Body.Close()
	var rps []rpInfo
	err = json.NewDecoder(res.Body).Decode(&rps)
	if err != nil {
		return nil, err
	}
	return rps, nil
}

type rpURL struct {
	URL    string `json:"url"`
	Access string `json:"access"`
}

func (u *rpURL) String() string {
	if u.Access == "normal" {
		return "/rp/" + u.URL
	} else if u.Access == "read" {
		return "/read/" + u.URL
	} else {
		return "???" + u.Access + "???"
	}
}

func getRpUrls(rpid string) ([]rpURL, error) {
	res, err := http.Get("http://127.0.0.1:12789/rps/" + rpid)
	if err != nil {
		return nil, err
	}
	defer res.Body.Close()
	var urls []rpURL
	err = json.NewDecoder(res.Body).Decode(&urls)
	if err != nil {
		return nil, err
	}
	return urls, nil
}

func destroyRp(rpid string) error {
	req, err := http.NewRequest("DELETE", "http://127.0.0.1:12789/rps/"+rpid, nil)
	if err != nil {
		return err
	}
	var client http.Client
	res, err := client.Do(req)
	if err != nil {
		return err
	}
	defer res.Body.Close()
	return nil
}
<commit_msg>finish implementing "start server" and "test server" options<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"net/http"
	"os"
	"os/exec"
	"strings"
	"time"

	"github.com/manifoldco/promptui"
)

func main() {
	fmt.Println("RPNow Admin Console")

	for {
		fmt.Println()
		up, pid := isServerUp()

		var options []string
		if up {
			options = []string{"manage rooms", "stop server", "exit"}
		} else {
			options = []string{"start server", "test server", "exit"}
		}

		prompt := promptui.Select{
			Label: "RP Admin",
			Items: options,
		}
		_, action, err := prompt.Run()

		if err != nil {
			return
		}
		switch action {
		case "start server":
			cmd := exec.Command("../src/index.js")

			if err := cmd.Start(); err != nil {
				fmt.Printf("Error starting server: %s\n", err)
			}

			isUp := make(chan bool)

			go func() {
				cmd.Wait()
				isUp <- false
			}()

			go func() {
				for i := 0; i < 10; i++ {
					time.Sleep(time.Duration(250) * time.Millisecond)
					if rpnowStatus, _ := isServerUp(); rpnowStatus {
						isUp <- true
						return
					}
				}
				isUp <- false
			}()

			if <-isUp {
				fmt.Println("Server launched")
			} else {
				fmt.Println("Server failed")
			}

			close(isUp)
		case "test server":
			cmd := exec.Command("../src/index.js")

			cmd.Stdout = os.Stdout
			cmd.Stderr = os.Stdout

			if err := cmd.Start(); err != nil {
				fmt.Printf("Error starting server: %s\n", err)
			}

			if err := cmd.Wait(); err != nil {
				fmt.Printf("Server exited on error: %s\n", err)
			} else {
				fmt.Println("Server exited cleanly")
			}
		case "stop server":
			process, err := os.FindProcess(pid)
			if err != nil {
				fmt.Printf("Error finding process with pid %d: %s\n", pid, err)
				continue
			}
			err = process.Kill()
			if err != nil {
				fmt.Printf("Error killing server process: %s\n", err)
				continue
			}
			fmt.Println("Server stopped")
		case "manage rooms":
			editRps()
		case "exit":
			return
		}
	}
}

func editRps() {
	for {
		// Select an RP from the list
		rps, err := getRpList()
		if err != nil {
			panic(err)
		}

		rp := pickRp(rps)
		if rp == nil {
			return
		}
		for {
			// Expand & edit the selected RP
			urls, err := getRpUrls(rp.RPID)
			if err != nil {
				panic(err)
			}

			fmt.Println()
			fmt.Println(rp.Title)
			for _, url := range urls {
				fmt.Printf("*  %s\n", url.String())
			}

			prompt := promptui.Select{
				Label: fmt.Sprintf("Modify %q", rp.Title),
				Items: []string{"go back", "edit urls", "destroy rp"},
			}
			_, action, err := prompt.Run()
			if err != nil {
				panic(err)
			}
			if action == "go back" {
				break
			} else if action == "edit urls" {

			} else if action == "destroy rp" {
				killswitch := strings.ToUpper(fmt.Sprintf("destroy %s", rp.Title))
				prompt := promptui.Prompt{Label: fmt.Sprintf("Type %q", killswitch)}
				result, _ := prompt.Run()
				if result == killswitch {
					err := destroyRp(rp.RPID)
					if err != nil {
						panic(err)
					}
					fmt.Printf("BOOM! %q is no more.\n", rp.Title)
					break
				} else {
					fmt.Println("Incorrect. Will not delete.")
				}
			}
		}
	}
}

func isServerUp() (bool, int) {
	fmt.Print("server status... ")

	res, err := http.Get("http://127.0.0.1:12789/status")
	if err != nil {
		fmt.Println("(server not running)")
		return false, 0
	}
	defer res.Body.Close()
	var status struct {
		RPNowLine string `json:"rpnow"`
		PID       int    `json:"pid"`
	}
	err = json.NewDecoder(res.Body).Decode(&status)
	if err != nil {
		fmt.Println("(bad response)")
		return false, 0
	}
	fmt.Println(status.RPNowLine)
	return status.RPNowLine == "ok", status.PID
}

func pickRp(rps []rpInfo) *rpInfo {
	fmt.Println()

	prompt := promptui.Select{
		Label: "Choose an RP",
		Items: rps,
		Searcher: func(input string, index int) bool {
			return strings.Contains(strings.ToLower(rps[index].Title), strings.ToLower(input))
		},
	}
	idx, _, err := prompt.Run()
	if err != nil {
		return nil
	}

	return &rps[idx]
}

type rpInfo struct {
	Title     string    `json:"title"`
	RPID      string    `json:"rpid"`
	Timestamp time.Time `json:"timestamp"`
}

func (r rpInfo) String() string {
	return fmt.Sprintf("%-30s (%s)", r.Title, r.Timestamp.Format("02 Jan 2006"))
}

func getRpList() ([]rpInfo, error) {
	res, err := http.Get("http://127.0.0.1:12789/rps")
	if err != nil {
		return nil, err
	}
	defer res.Body.Close()
	var rps []rpInfo
	err = json.NewDecoder(res.Body).Decode(&rps)
	if err != nil {
		return nil, err
	}
	return rps, nil
}

type rpURL struct {
	URL    string `json:"url"`
	Access string `json:"access"`
}

func (u *rpURL) String() string {
	if u.Access == "normal" {
		return "/rp/" + u.URL
	} else if u.Access == "read" {
		return "/read/" + u.URL
	} else {
		return "???" + u.Access + "???"
	}
}

func getRpUrls(rpid string) ([]rpURL, error) {
	res, err := http.Get("http://127.0.0.1:12789/rps/" + rpid)
	if err != nil {
		return nil, err
	}
	defer res.Body.Close()
	var urls []rpURL
	err = json.NewDecoder(res.Body).Decode(&urls)
	if err != nil {
		return nil, err
	}
	return urls, nil
}

func destroyRp(rpid string) error {
	req, err := http.NewRequest("DELETE", "http://127.0.0.1:12789/rps/"+rpid, nil)
	if err != nil {
		return err
	}
	var client http.Client
	res, err := client.Do(req)
	if err != nil {
		return err
	}
	defer res.Body.Close()
	return nil
}
<|endoftext|>


<commit_before>// Package aphjsapi provides additional interfaces, wrapper and helper functions for original
// jsapi package("github.com/manyminds/api2go/jsapi")
package aphjsonapi

import (
	"fmt"
	"math"
	"reflect"
	"strings"

	"github.com/dictyBase/go-middlewares/middlewares/pagination"
	jsapi "github.com/manyminds/api2go/jsonapi"
)

// RelationshipLink is a container type for having information about
// relationship links
type RelationShipLink struct {
	Name string `validate:"required"`
	// To override the default links, it will be appended to
	// the base url.
	SuffixFragment string
	// The type(type key in JSAONAPI specs) of relationship resource
	Type string `validate:"required"`
}

// MarshalSelfRelations is an interface  for creating self relationship links
type MarshalSelfRelations interface {
	// Validates all instances of RelationshipLink structures, using
	// https://gopkg.in/go-playground/validator.v9 package is recommended
	ValidateSelfLinks() error
	GetSelfLinksInfo() []RelationShipLink
}

// MarshalRelatedRelations is an interface  for creating related relationship
// links
type MarshalRelatedRelations interface {
	// Validates all instances of RelationshipLink structures, using
	// https://gopkg.in/go-playground/validator.v9 package is recommended
	ValidateRelatedLinks() error
	GetRelatedLinksInfo() []RelationShipLink
}

// AttributeToDbRowMapper is an interface to provide mapping between jsapi
// attribute and database row names. This is useful for implementing filter
// query parameter
type AttributeToDbRowMapper interface {
	GetMap() map[string]string
}

// RelationshipAttributes is an interface to provide attribute fields of
// relationship resources. This is mandatory for supporting sparse fieldset
// query parameter.
type RelationshipAttribute interface {
	GetAttributeFields(string) []string
}

// MarshalWithPagination adds pagination information for collection resource
func MarshalWithPagination(data interface{}, ep jsapi.ServerInformation, opt *pagination.Props) (*jsapi.Document, error) {
	var jst *jsapi.Document
	if reflect.TypeOf(data).Kind() != reflect.Slice {
		return jst, fmt.Errorf("%s\n", "Only slice type is allowed for pagination")
	}
	jst, err := MarshalToStructWrapper(data, ep)
	if err != nil {
		return jst, err
	}
	baseLink := jst.Links.Self
	pageLink := &jsapi.Links{}
	pageLink.Self = generatePaginatedResourceLink(baseLink, opt.Current, opt.Entries)
	pageLink.First = generatePaginatedResourceLink(baseLink, 1, opt.Entries)
	if opt.Current != 1 {
		pageLink.Previous = generatePaginatedResourceLink(baseLink, opt.Current-1, opt.Entries)
	}
	lastPage := int(math.Floor(float64(opt.Records) / float64(opt.Entries)))
	pageLink.Last = generatePaginatedResourceLink(baseLink, lastPage, opt.Entries)
	if opt.Current != lastPage {
		pageLink.Next = generatePaginatedResourceLink(baseLink, opt.Current+1, opt.Entries)
	}
	jst.Links = pageLink
	jst.Meta = map[string]interface{}{
		"pagination": map[string]int{
			"records": opt.Records,
			"total":   lastPage,
			"size":    opt.Entries,
			"number":  opt.Current,
		},
	}
	return jst, nil
}

// MarshalToStructWrapper adds relationship information and returns a
// jsapi.Document structure for further json encoding
func MarshalToStructWrapper(data interface{}, ep jsapi.ServerInformation) (*jsapi.Document, error) {
	jst, err := jsapi.MarshalToStruct(data, ep)
	if err != nil {
		return jst, err
	}
	if len(jst.Data.DataArray) > 0 { //array resource objects
		// picking first element both from the generated and given typed structures
		elem := jst.Data.DataArray[0]
		value := reflect.ValueOf(data).Index(0).Interface()
		// link for the array resource itself
		jst.Links = &jsapi.Links{Self: generateMultiResourceLink(&elem, ep)}
		for i, d := range jst.Data.DataArray {
			// link for individual resource
			jst.Data.DataArray[i].Links = &jsapi.Links{Self: generateSingleResourceLink(&d, ep)}
			// Add relationships to every member
			r := generateRelationshipLinks(value, &d, ep)
			jst.Data.DataArray[i].Relationships = r
		}
	} else {
		jst.Links = &jsapi.Links{Self: generateSingleResourceLink(jst.Data.DataObject, ep)}
		relationships := generateRelationshipLinks(data, jst.Data.DataObject, ep)
		if len(relationships) > 0 {
			jst.Data.DataObject.Relationships = relationships
		}
	}
	return jst, nil
}

func generateBaseLink(ep jsapi.ServerInformation) string {
	return fmt.Sprintf(
		"%s/%s",
		strings.Trim(ep.GetBaseURL(), "/"),
		strings.Trim(ep.GetPrefix(), "/"),
	)
}

func generatePaginatedResourceLink(baseurl string, pagenum, pagesize int) string {
	return fmt.Sprintf(
		"%s?page[number]=%d&page[size]=%d",
		baseurl,
		pagenum,
		pagesize,
	)
}

func generateSingleResourceLink(jdata *jsapi.Data, ep jsapi.ServerInformation) string {
	return fmt.Sprintf(
		"%s/%s/%s",
		generateBaseLink(ep),
		jdata.Type,
		jdata.ID,
	)
}

func generateMultiResourceLink(jdata *jsapi.Data, ep jsapi.ServerInformation) string {
	return fmt.Sprintf(
		"%s/%s",
		generateBaseLink(ep),
		jdata.Type,
	)
}

func generateRelationshipLinks(data interface{}, jdata *jsapi.Data, ep jsapi.ServerInformation) map[string]jsapi.Relationship {
	relationships := make(map[string]jsapi.Relationship)
	baselink := generateBaseLink(ep)
	self, ok := data.(MarshalSelfRelations)
	if ok {
		for _, rel := range self.GetSelfLinksInfo() {
			links := &jsapi.Links{}
			if len(rel.SuffixFragment) > 0 {
				links.Self = fmt.Sprintf("%s/%s", baselink, strings.Trim(rel.SuffixFragment, "/"))
			} else {
				links.Self = fmt.Sprintf("%s/%s/%s/relationships/%s",
					baselink,
					jdata.Type,
					jdata.ID,
					rel.Name,
				)
			}
			relationships[rel.Name] = jsapi.Relationship{Links: links}
		}
	}
	related, ok := data.(MarshalRelatedRelations)
	if ok {
		for _, rel := range related.GetRelatedLinksInfo() {
			var rlink string
			if len(rel.SuffixFragment) > 0 {
				rlink = fmt.Sprintf("%s/%s", baselink, strings.Trim(rel.SuffixFragment, "/"))
			} else {
				rlink = fmt.Sprintf("%s/%s/%s/%s",
					baselink,
					jdata.Type,
					jdata.ID, rel.Name,
				)
			}
			if _, ok := relationships[rel.Name]; ok {
				relationships[rel.Name].Links.Related = rlink
			} else {
				relationships[rel.Name] = jsapi.Relationship{Links: &jsapi.Links{Related: rlink}}
			}
		}
	}
	return relationships
}

// MapFieldsToDbRow maps jsapi attributes to database row names
func MapFieldsToDbRow(data interface{}) map[string]string {
	m, ok := data.(AttributeToDbRowMapper)
	if ok {
		return m.GetMap()
	}
	frow := make(map[string]string)
	t := reflect.TypeOf(data)
	for i := 0; i < t.NumField(); i++ {
		tag := t.Field(i).Tag
		v, ok := tag.Lookup("json")
		if ok && v != "-" {
			r, dbok := tag.Lookup("db")
			if dbok && r != "-" {
				frow[v] = r
			}
		}
	}
	return frow
}

// GetTypeName gets the type name(type field) from a jsapi implementing
// interface. It is recommended to implement jsapi.EntityNamer interface to
// reduce the use of reflection
func GetTypeName(data interface{}) string {
	entity, ok := data.(jsapi.EntityNamer)
	if ok {
		return entity.GetName()
	}
	rType := reflect.TypeOf(data)
	if rType.Kind() == reflect.Ptr {
		return jsapi.Pluralize(jsapi.Jsonify(rType.Elem().Name()))
	}
	return jsapi.Pluralize(jsapi.Jsonify(rType.Name()))
}

// AttributeNames returns all JSAONAPI attribute names of data interface
func GetAttributeFields(data interface{}) ([]string, error) {
	var attr []string
	t := reflect.TypeOf(data)
	if t == nil {
		return attr, fmt.Errorf("received an nil interface")
	}
	var st reflect.Type
	if t.Kind() == reflect.Ptr {
		st = t.Elem()
	} else {
		st = t
	}
	for i := 0; i < st.NumField(); i++ {
		v, ok := st.Field(i).Tag.Lookup("json")
		if ok && v != "-" {
			attr = append(attr, v)
		}
	}
	return attr, nil
}

// GetAllRelationships returns all relationships of data interface
func GetAllRelationships(data interface{}) []RelationShipLink {
	var r []RelationShipLink
	self, ok := data.(MarshalSelfRelations)
	if ok {
		r = append(r, self.GetSelfLinksInfo()...)
	}
	related, ok := data.(MarshalRelatedRelations)
	if ok {
		r = append(r, related.GetRelatedLinksInfo()...)
	}
	return r
}

//GetRelatedTypes returns a map jsapi types of the related resources using
//reflection
func getRelatedTypeNames(data interface{}) []string {
	var names []string
	mtype := reflect.TypeOf((*jsapi.MarshalIdentifier)(nil)).Elem()
	t := reflect.TypeOf(data)
	for i := 0; i < t.NumField(); i++ {
		ftype := t.Field(i).Type
		if ftype.Kind() == reflect.Slice {
			if ftype.Elem().Implements(mtype) {
				names = append(
					names,
					jsapi.Pluralize(
						jsapi.Jsonify(
							ftype.Elem().Name(),
						),
					),
				)
				continue
			}
		}
		if ftype.Implements(mtype) {
			names = append(
				names,
				jsapi.Pluralize(
					jsapi.Jsonify(
						ftype.Name(),
					),
				),
			)
		}
	}
	return names
}
<commit_msg>changed the return type<commit_after>// Package aphjsapi provides additional interfaces, wrapper and helper functions for original
// jsapi package("github.com/manyminds/api2go/jsapi")
package aphjsonapi

import (
	"fmt"
	"math"
	"reflect"
	"strings"

	"github.com/dictyBase/go-middlewares/middlewares/pagination"
	jsapi "github.com/manyminds/api2go/jsonapi"
)

// RelationshipLink is a container type for having information about
// relationship links
type RelationShipLink struct {
	Name string `validate:"required"`
	// To override the default links, it will be appended to
	// the base url.
	SuffixFragment string
	// The type(type key in JSAONAPI specs) of relationship resource
	Type string `validate:"required"`
}

// MarshalSelfRelations is an interface  for creating self relationship links
type MarshalSelfRelations interface {
	// Validates all instances of RelationshipLink structures, using
	// https://gopkg.in/go-playground/validator.v9 package is recommended
	ValidateSelfLinks() error
	GetSelfLinksInfo() []RelationShipLink
}

// MarshalRelatedRelations is an interface  for creating related relationship
// links
type MarshalRelatedRelations interface {
	// Validates all instances of RelationshipLink structures, using
	// https://gopkg.in/go-playground/validator.v9 package is recommended
	ValidateRelatedLinks() error
	GetRelatedLinksInfo() []RelationShipLink
}

// AttributeToDbRowMapper is an interface to provide mapping between jsapi
// attribute and database row names. This is useful for implementing filter
// query parameter
type AttributeToDbRowMapper interface {
	GetMap() map[string]string
}

// RelationshipAttributes is an interface to provide attribute fields of
// relationship resources. This is mandatory for supporting sparse fieldset
// query parameter.
type RelationshipAttribute interface {
	GetAttributeFields(string) []string
}

// MarshalWithPagination adds pagination information for collection resource
func MarshalWithPagination(data interface{}, ep jsapi.ServerInformation, opt *pagination.Props) (*jsapi.Document, error) {
	var jst *jsapi.Document
	if reflect.TypeOf(data).Kind() != reflect.Slice {
		return jst, fmt.Errorf("%s\n", "Only slice type is allowed for pagination")
	}
	jst, err := MarshalToStructWrapper(data, ep)
	if err != nil {
		return jst, err
	}
	baseLink := jst.Links.Self
	pageLink := &jsapi.Links{}
	pageLink.Self = generatePaginatedResourceLink(baseLink, opt.Current, opt.Entries)
	pageLink.First = generatePaginatedResourceLink(baseLink, 1, opt.Entries)
	if opt.Current != 1 {
		pageLink.Previous = generatePaginatedResourceLink(baseLink, opt.Current-1, opt.Entries)
	}
	lastPage := int(math.Floor(float64(opt.Records) / float64(opt.Entries)))
	pageLink.Last = generatePaginatedResourceLink(baseLink, lastPage, opt.Entries)
	if opt.Current != lastPage {
		pageLink.Next = generatePaginatedResourceLink(baseLink, opt.Current+1, opt.Entries)
	}
	jst.Links = pageLink
	jst.Meta = map[string]interface{}{
		"pagination": map[string]int{
			"records": opt.Records,
			"total":   lastPage,
			"size":    opt.Entries,
			"number":  opt.Current,
		},
	}
	return jst, nil
}

// MarshalToStructWrapper adds relationship information and returns a
// jsapi.Document structure for further json encoding
func MarshalToStructWrapper(data interface{}, ep jsapi.ServerInformation) (*jsapi.Document, error) {
	jst, err := jsapi.MarshalToStruct(data, ep)
	if err != nil {
		return jst, err
	}
	if len(jst.Data.DataArray) > 0 { //array resource objects
		// picking first element both from the generated and given typed structures
		elem := jst.Data.DataArray[0]
		value := reflect.ValueOf(data).Index(0).Interface()
		// link for the array resource itself
		jst.Links = &jsapi.Links{Self: generateMultiResourceLink(&elem, ep)}
		for i, d := range jst.Data.DataArray {
			// link for individual resource
			jst.Data.DataArray[i].Links = &jsapi.Links{Self: generateSingleResourceLink(&d, ep)}
			// Add relationships to every member
			r := generateRelationshipLinks(value, &d, ep)
			jst.Data.DataArray[i].Relationships = r
		}
	} else {
		jst.Links = &jsapi.Links{Self: generateSingleResourceLink(jst.Data.DataObject, ep)}
		relationships := generateRelationshipLinks(data, jst.Data.DataObject, ep)
		if len(relationships) > 0 {
			jst.Data.DataObject.Relationships = relationships
		}
	}
	return jst, nil
}

func generateBaseLink(ep jsapi.ServerInformation) string {
	return fmt.Sprintf(
		"%s/%s",
		strings.Trim(ep.GetBaseURL(), "/"),
		strings.Trim(ep.GetPrefix(), "/"),
	)
}

func generatePaginatedResourceLink(baseurl string, pagenum, pagesize int) string {
	return fmt.Sprintf(
		"%s?page[number]=%d&page[size]=%d",
		baseurl,
		pagenum,
		pagesize,
	)
}

func generateSingleResourceLink(jdata *jsapi.Data, ep jsapi.ServerInformation) string {
	return fmt.Sprintf(
		"%s/%s/%s",
		generateBaseLink(ep),
		jdata.Type,
		jdata.ID,
	)
}

func generateMultiResourceLink(jdata *jsapi.Data, ep jsapi.ServerInformation) string {
	return fmt.Sprintf(
		"%s/%s",
		generateBaseLink(ep),
		jdata.Type,
	)
}

func generateRelationshipLinks(data interface{}, jdata *jsapi.Data, ep jsapi.ServerInformation) map[string]jsapi.Relationship {
	relationships := make(map[string]jsapi.Relationship)
	baselink := generateBaseLink(ep)
	self, ok := data.(MarshalSelfRelations)
	if ok {
		for _, rel := range self.GetSelfLinksInfo() {
			links := &jsapi.Links{}
			if len(rel.SuffixFragment) > 0 {
				links.Self = fmt.Sprintf("%s/%s", baselink, strings.Trim(rel.SuffixFragment, "/"))
			} else {
				links.Self = fmt.Sprintf("%s/%s/%s/relationships/%s",
					baselink,
					jdata.Type,
					jdata.ID,
					rel.Name,
				)
			}
			relationships[rel.Name] = jsapi.Relationship{Links: links}
		}
	}
	related, ok := data.(MarshalRelatedRelations)
	if ok {
		for _, rel := range related.GetRelatedLinksInfo() {
			var rlink string
			if len(rel.SuffixFragment) > 0 {
				rlink = fmt.Sprintf("%s/%s", baselink, strings.Trim(rel.SuffixFragment, "/"))
			} else {
				rlink = fmt.Sprintf("%s/%s/%s/%s",
					baselink,
					jdata.Type,
					jdata.ID, rel.Name,
				)
			}
			if _, ok := relationships[rel.Name]; ok {
				relationships[rel.Name].Links.Related = rlink
			} else {
				relationships[rel.Name] = jsapi.Relationship{Links: &jsapi.Links{Related: rlink}}
			}
		}
	}
	return relationships
}

// MapFieldsToDbRow maps jsapi attributes to database row names
func MapFieldsToDbRow(data interface{}) map[string]string {
	m, ok := data.(AttributeToDbRowMapper)
	if ok {
		return m.GetMap()
	}
	frow := make(map[string]string)
	t := reflect.TypeOf(data)
	for i := 0; i < t.NumField(); i++ {
		tag := t.Field(i).Tag
		v, ok := tag.Lookup("json")
		if ok && v != "-" {
			r, dbok := tag.Lookup("db")
			if dbok && r != "-" {
				frow[v] = r
			}
		}
	}
	return frow
}

// GetTypeName gets the type name(type field) from a jsapi implementing
// interface. It is recommended to implement jsapi.EntityNamer interface to
// reduce the use of reflection
func GetTypeName(data interface{}) string {
	entity, ok := data.(jsapi.EntityNamer)
	if ok {
		return entity.GetName()
	}
	rType := reflect.TypeOf(data)
	if rType.Kind() == reflect.Ptr {
		return jsapi.Pluralize(jsapi.Jsonify(rType.Elem().Name()))
	}
	return jsapi.Pluralize(jsapi.Jsonify(rType.Name()))
}

// AttributeNames returns all JSAONAPI attribute names of data interface
func GetAttributeFields(data interface{}) []string {
	var attr []string
	t := reflect.TypeOf(data)
	if t == nil {
		return attr
	}
	var st reflect.Type
	if t.Kind() == reflect.Ptr {
		st = t.Elem()
	} else {
		st = t
	}
	for i := 0; i < st.NumField(); i++ {
		v, ok := st.Field(i).Tag.Lookup("json")
		if ok && v != "-" {
			attr = append(attr, v)
		}
	}
	return attr, nil
}

// GetAllRelationships returns all relationships of data interface
func GetAllRelationships(data interface{}) []RelationShipLink {
	var r []RelationShipLink
	self, ok := data.(MarshalSelfRelations)
	if ok {
		r = append(r, self.GetSelfLinksInfo()...)
	}
	related, ok := data.(MarshalRelatedRelations)
	if ok {
		r = append(r, related.GetRelatedLinksInfo()...)
	}
	return r
}

//GetRelatedTypes returns a map jsapi types of the related resources using
//reflection
func getRelatedTypeNames(data interface{}) []string {
	var names []string
	mtype := reflect.TypeOf((*jsapi.MarshalIdentifier)(nil)).Elem()
	t := reflect.TypeOf(data)
	for i := 0; i < t.NumField(); i++ {
		ftype := t.Field(i).Type
		if ftype.Kind() == reflect.Slice {
			if ftype.Elem().Implements(mtype) {
				names = append(
					names,
					jsapi.Pluralize(
						jsapi.Jsonify(
							ftype.Elem().Name(),
						),
					),
				)
				continue
			}
		}
		if ftype.Implements(mtype) {
			names = append(
				names,
				jsapi.Pluralize(
					jsapi.Jsonify(
						ftype.Name(),
					),
				),
			)
		}
	}
	return names
}
<|endoftext|>


<commit_before>package pq

import (
	"fmt"
	"reflect"
	"sync"
	"time"

	"github.com/armon/relay"
	"github.com/armon/relay/broker"
)

// The pq package provides a simple priority queue. All that is required is a
// a known number of available priorities. Queues labeled from 0..N will be
// automatically created, where the higher numbers are higher priority queues.
// High-level methods are exposed to make dealing with the priority queues
// very simple and abstract.

var (
	// The duration of a read attempt from an individual AMQP queue. The
	// consumer methods use this value to run short Consume() calls so that
	// we can check for interrupts. The higher this number is, the longer it
	// it may take to interrupt a read and exit.
	consumeResetInterval = 500 * time.Millisecond

	// The minimum amount of time to wait after receiving a message for a higher
	// priority message to arrive.
	MinQuietPeriod = 10 * time.Millisecond
)

// PriorityQueue is a simple wrapper around a relay.Broker to manage
// a set of queues with varying priority.
type PriorityQueue struct {
	max          int
	source       broker.Broker
	prefix       string
	quietPeriod  time.Duration
	publishers   []broker.Publisher
	shutdownLock sync.Mutex
	shutdown     bool
}

// priorityResp is used as a container for response data during a Consume().
// Since the response must be fed down a channel and contains multiple values,
// we stuff them all into this struct and thread it through.
type priorityResp struct {
	value    interface{}
	priority int
	consumer broker.Consumer
}

// NewPriorityQueue returns a new priority queue from which a consumer or
// producer at a given priority can be easily retrieved.
func NewPriorityQueue(
	b broker.Broker,
	pri int,
	prefix string,
	quietPeriod time.Duration) (*PriorityQueue, error) {

	if b == nil {
		return nil, fmt.Errorf("Broker must not be nil")
	}
	if pri < 1 {
		return nil, fmt.Errorf("Must be 1 or more priorities")
	}

	// Guard against a quiet period which is too small
	if quietPeriod < MinQuietPeriod {
		quietPeriod = MinQuietPeriod
	}

	q := PriorityQueue{
		source:      b,
		prefix:      prefix,
		max:         pri - 1,
		quietPeriod: quietPeriod,
	}

	// Initialize the publisher cache
	q.publishers = make([]broker.Publisher, pri)

	return &q, nil
}

// queueName formats the name of a priority queue by appending its priority to
// the user-provided queue prefix.
func queueName(prefix string, pri int) string {
	return fmt.Sprintf("%s-%d", prefix, pri)
}

// Max returns the highest priority number.
func (q *PriorityQueue) Max() int {
	return q.max
}

// Min returns the lowest priority number. This is always 0.
func (q *PriorityQueue) Min() int {
	return 0
}

// publisher returns a new publisher from the priority indicated by pri.
func (q *PriorityQueue) publisher(pri int) (broker.Publisher, error) {
	if pri > q.Max() || pri < q.Min() {
		return nil, fmt.Errorf("Priority out of range: %d", pri)
	}

	if q.publishers[pri] == nil {
		pub, err := q.source.Publisher(queueName(q.prefix, pri))
		if err != nil {
			return nil, err
		}
		q.publishers[pri] = pub
	}

	return q.publishers[pri], nil
}

// consumer returns a new consumer with the indicated priority.
func (q *PriorityQueue) consumer(pri int) (broker.Consumer, error) {
	if pri > q.Max() || pri < q.Min() {
		return nil, fmt.Errorf("Priority out of range: %d", pri)
	}

	cons, err := q.source.Consumer(queueName(q.prefix, pri))
	if err != nil {
		return nil, err
	}

	return cons, nil
}

// Publish will publish a message at a given priority. The publisher is
// automatically closed afterward.
func (q *PriorityQueue) Publish(payload interface{}, pri int) error {
	pub, err := q.publisher(pri)
	if err != nil {
		return err
	}

	if err := pub.Publish(payload); err != nil {
		return err
	}

	return nil
}

// consume consumes a message from the priority queue. This is a blocking call
// which will watch every queue at every priority until a message is received on
// at least one of them. Once a message is received, we continue blocking for a
// configurable amount of time for any other higher priority message to arrive.
// After the quiet period, if no higher priority message has been received, all
// lower priority messages are marked for re-delivery, and the most urgent
// message is returned.
//
// The quiet period should almost always carry a value greater than zero. If
// there is no quiet period, then messages are returned as soon as they are
// received, which makes the quickness of the server and queue the decider of
// message priority. Therefore, we enforce a minimum quiet period.
//
// The consumer is also returned as part of the result, as it will contain
// the session open to the queue with an un-Ack()'ed message. It is the
// responsibility of the caller to Ack() and Close() the consumer.
func (q *PriorityQueue) consume(
	out interface{}, cancelCh chan struct{}) (broker.Consumer, int, error) {

	shutdownCh := make(chan struct{})

	// Populate the cancelCh if none was provided
	if cancelCh == nil {
		cancelCh = make(chan struct{})
	}

	// Initialize the data channels
	errCh := make(chan error, q.Max()+1)
	respCh := make(chan priorityResp, q.Max()+1)

	var consumers []broker.Consumer
	for i := q.Min(); i <= q.Max(); i++ {
		cons, err := q.consumer(i)
		if err != nil {
			return cons, 0, err
		}
		consumers = append(consumers, cons)
	}

	for i := q.Min(); i <= q.Max(); i++ {
		go func(cons broker.Consumer, pri int) {
			val := reflect.New(reflect.TypeOf(out)).Interface()
			for {
				select {
				case <-cancelCh:
					cons.Close()
					return
				case <-shutdownCh:
					cons.Close()
					return
				default:
				}

				// Keep consuming with timeout so we can bail if we need to
				if err := cons.ConsumeTimeout(&val, consumeResetInterval); err != nil {
					if err == relay.TimedOut {
						continue
					}
					errCh <- err
					return
				}

				// Check for cancellation
				select {
				case <-cancelCh:
					if cons != nil {
						cons.Nack()
						cons.Close()
					}
				default:
				}

				// Create the response object to send down the channel
				r := priorityResp{
					value:    val,
					priority: pri,
					consumer: cons,
				}
				respCh <- r
				return
			}
		}(consumers[i], i)
	}

	var wait <-chan time.Time
	var cons broker.Consumer
	highest := q.Min() - 1 // Allows Min() messages to be accepted

OUTER:
	for {
		select {
		case err := <-errCh:
			close(shutdownCh)
			return cons, 0, err

		case r := <-respCh:
			if r.priority <= highest {
				// Send a negative acknowledgement and close the consumer
				r.consumer.Nack()
				r.consumer.Close()
				continue
			}

			// Close the outstanding consumer. This is only set if some other
			// value had previously been considered the highest value.
			if cons != nil {
				cons.Nack()
				cons.Close()
			}

			// Received message was higher priority, so re-assign the results
			// and enter the quiet period for any higher-priority messages
			// to arrive.
			dst := reflect.Indirect(reflect.ValueOf(out))
			src := reflect.Indirect(reflect.ValueOf(r.value))
			dst.Set(reflect.Indirect(src))

			highest = r.priority
			cons = r.consumer
			wait = time.After(q.quietPeriod)

		case <-wait:
			// Signal all higher-priority readers to close, they didn't receive
			// any messages we can use.
			close(shutdownCh)
			return cons, highest, nil

		case <-cancelCh:
			break OUTER

		case <-shutdownCh:
			break OUTER
		}
	}

	return cons, 0, nil
}

// Consume is the public method for consuming data out of a priority queue. It
// will block until data is received, and returns the priority level of the
// consumed message along with any errors. The consumer is also returned, which
// should be Ack'ed and Closed by the caller.
func (q *PriorityQueue) Consume(out interface{}) (broker.Consumer, int, error) {
	return q.consume(out, nil)
}

// ConsumeCancel allows passing in a channel to signal that we should
// stop trying to consume a message. Internally this channel will be checked on
// a short interval, and will shut down the job if the channel has been closed.
func (q *PriorityQueue) ConsumeCancel(
	out interface{}, cancelCh chan struct{}) (broker.Consumer, int, error) {

	if cancelCh == nil {
		return nil, 0, fmt.Errorf("Cancellation channel cannot be nil")
	}
	return q.consume(out, cancelCh)
}

// Close will call a shutdown on all publishers we have used. By default, all
// publishers are kept open so that multiple calls to establish the sessions are
// not always required. This method shuts them all down and resets the pool.
func (q *PriorityQueue) Close() error {
	q.shutdownLock.Lock()
	defer q.shutdownLock.Unlock()

	if q.shutdown {
		return nil
	}
	q.shutdown = true

	for i, pub := range q.publishers {
		if pub != nil {
			if err := pub.Close(); err != nil {
				return err
			}
		}
		q.publishers[i] = nil
	}

	return nil
}
<commit_msg>pq: seal up possible consumer leaks<commit_after>package pq

import (
	"fmt"
	"reflect"
	"sync"
	"time"

	"github.com/armon/relay"
	"github.com/armon/relay/broker"
)

// The pq package provides a simple priority queue. All that is required is a
// a known number of available priorities. Queues labeled from 0..N will be
// automatically created, where the higher numbers are higher priority queues.
// High-level methods are exposed to make dealing with the priority queues
// very simple and abstract.

var (
	// The duration of a read attempt from an individual AMQP queue. The
	// consumer methods use this value to run short Consume() calls so that
	// we can check for interrupts. The higher this number is, the longer it
	// it may take to interrupt a read and exit.
	consumeResetInterval = 500 * time.Millisecond

	// The minimum amount of time to wait after receiving a message for a higher
	// priority message to arrive.
	MinQuietPeriod = 10 * time.Millisecond
)

// PriorityQueue is a simple wrapper around a relay.Broker to manage
// a set of queues with varying priority.
type PriorityQueue struct {
	max          int
	source       broker.Broker
	prefix       string
	quietPeriod  time.Duration
	publishers   []broker.Publisher
	shutdownLock sync.Mutex
	shutdown     bool
}

// priorityResp is used as a container for response data during a Consume().
// Since the response must be fed down a channel and contains multiple values,
// we stuff them all into this struct and thread it through.
type priorityResp struct {
	value    interface{}
	priority int
	consumer broker.Consumer
}

// NewPriorityQueue returns a new priority queue from which a consumer or
// producer at a given priority can be easily retrieved.
func NewPriorityQueue(
	b broker.Broker,
	pri int,
	prefix string,
	quietPeriod time.Duration) (*PriorityQueue, error) {

	if b == nil {
		return nil, fmt.Errorf("Broker must not be nil")
	}
	if pri < 1 {
		return nil, fmt.Errorf("Must be 1 or more priorities")
	}

	// Guard against a quiet period which is too small
	if quietPeriod < MinQuietPeriod {
		quietPeriod = MinQuietPeriod
	}

	q := PriorityQueue{
		source:      b,
		prefix:      prefix,
		max:         pri - 1,
		quietPeriod: quietPeriod,
	}

	// Initialize the publisher cache
	q.publishers = make([]broker.Publisher, pri)

	return &q, nil
}

// queueName formats the name of a priority queue by appending its priority to
// the user-provided queue prefix.
func queueName(prefix string, pri int) string {
	return fmt.Sprintf("%s-%d", prefix, pri)
}

// Max returns the highest priority number.
func (q *PriorityQueue) Max() int {
	return q.max
}

// Min returns the lowest priority number. This is always 0.
func (q *PriorityQueue) Min() int {
	return 0
}

// publisher returns a new publisher from the priority indicated by pri.
func (q *PriorityQueue) publisher(pri int) (broker.Publisher, error) {
	if pri > q.Max() || pri < q.Min() {
		return nil, fmt.Errorf("Priority out of range: %d", pri)
	}

	if q.publishers[pri] == nil {
		pub, err := q.source.Publisher(queueName(q.prefix, pri))
		if err != nil {
			return nil, err
		}
		q.publishers[pri] = pub
	}

	return q.publishers[pri], nil
}

// consumer returns a new consumer with the indicated priority.
func (q *PriorityQueue) consumer(pri int) (broker.Consumer, error) {
	if pri > q.Max() || pri < q.Min() {
		return nil, fmt.Errorf("Priority out of range: %d", pri)
	}

	cons, err := q.source.Consumer(queueName(q.prefix, pri))
	if err != nil {
		if cons != nil {
			cons.Close()
		}
		return nil, err
	}

	return cons, nil
}

// Publish will publish a message at a given priority. The publisher is
// automatically closed afterward.
func (q *PriorityQueue) Publish(payload interface{}, pri int) error {
	pub, err := q.publisher(pri)
	if err != nil {
		return err
	}

	if err := pub.Publish(payload); err != nil {
		return err
	}

	return nil
}

// consume consumes a message from the priority queue. This is a blocking call
// which will watch every queue at every priority until a message is received on
// at least one of them. Once a message is received, we continue blocking for a
// configurable amount of time for any other higher priority message to arrive.
// After the quiet period, if no higher priority message has been received, all
// lower priority messages are marked for re-delivery, and the most urgent
// message is returned.
//
// The quiet period should almost always carry a value greater than zero. If
// there is no quiet period, then messages are returned as soon as they are
// received, which makes the quickness of the server and queue the decider of
// message priority. Therefore, we enforce a minimum quiet period.
//
// The consumer is also returned as part of the result, as it will contain
// the session open to the queue with an un-Ack()'ed message. It is the
// responsibility of the caller to Ack() and Close() the consumer.
func (q *PriorityQueue) consume(
	out interface{}, cancelCh chan struct{}) (broker.Consumer, int, error) {

	shutdownCh := make(chan struct{})

	// Populate the cancelCh if none was provided
	if cancelCh == nil {
		cancelCh = make(chan struct{})
	}

	// Initialize the data channels
	errCh := make(chan error, q.Max()+1)
	respCh := make(chan priorityResp, q.Max()+1)

	var consumers []broker.Consumer
	for i := q.Min(); i <= q.Max(); i++ {
		cons, err := q.consumer(i)
		if err != nil {
			return nil, 0, err
		}
		consumers = append(consumers, cons)
	}

	for i := q.Min(); i <= q.Max(); i++ {
		go func(cons broker.Consumer, pri int) {
			val := reflect.New(reflect.TypeOf(out)).Interface()
			for {
				select {
				case <-cancelCh:
					cons.Close()
					return
				case <-shutdownCh:
					cons.Close()
					return
				default:
				}

				// Keep consuming with timeout so we can bail if we need to
				if err := cons.ConsumeTimeout(&val, consumeResetInterval); err != nil {
					if err == relay.TimedOut {
						continue
					}
					cons.Close()
					errCh <- err
					return
				}

				// Check for cancellation
				select {
				case <-cancelCh:
					cons.Close()
				default:
				}

				// Create the response object to send down the channel
				r := priorityResp{
					value:    val,
					priority: pri,
					consumer: cons,
				}
				respCh <- r
				return
			}
		}(consumers[i], i)
	}

	var wait <-chan time.Time
	var cons broker.Consumer
	highest := q.Min() - 1 // Allows Min() messages to be accepted

OUTER:
	for {
		select {
		case err := <-errCh:
			close(shutdownCh)
			return cons, 0, err

		case r := <-respCh:
			if r.priority <= highest {
				// Send a negative acknowledgement and close the consumer
				r.consumer.Nack()
				r.consumer.Close()
				continue
			}

			// Close the outstanding consumer. This is only set if some other
			// value had previously been considered the highest value.
			if cons != nil {
				cons.Nack()
				cons.Close()
			}

			// Received message was higher priority, so re-assign the results
			// and enter the quiet period for any higher-priority messages
			// to arrive.
			dst := reflect.Indirect(reflect.ValueOf(out))
			src := reflect.Indirect(reflect.ValueOf(r.value))
			dst.Set(reflect.Indirect(src))

			highest = r.priority
			cons = r.consumer
			wait = time.After(q.quietPeriod)

		case <-wait:
			// Signal all higher-priority readers to close, they didn't receive
			// any messages we can use.
			close(shutdownCh)
			return cons, highest, nil

		case <-cancelCh:
			break OUTER

		case <-shutdownCh:
			break OUTER
		}
	}

	return cons, 0, nil
}

// Consume is the public method for consuming data out of a priority queue. It
// will block until data is received, and returns the priority level of the
// consumed message along with any errors. The consumer is also returned, which
// should be Ack'ed and Closed by the caller.
func (q *PriorityQueue) Consume(out interface{}) (broker.Consumer, int, error) {
	return q.consume(out, nil)
}

// ConsumeCancel allows passing in a channel to signal that we should
// stop trying to consume a message. Internally this channel will be checked on
// a short interval, and will shut down the job if the channel has been closed.
func (q *PriorityQueue) ConsumeCancel(
	out interface{}, cancelCh chan struct{}) (broker.Consumer, int, error) {

	if cancelCh == nil {
		return nil, 0, fmt.Errorf("Cancellation channel cannot be nil")
	}
	return q.consume(out, cancelCh)
}

// Close will call a shutdown on all publishers we have used. By default, all
// publishers are kept open so that multiple calls to establish the sessions are
// not always required. This method shuts them all down and resets the pool.
func (q *PriorityQueue) Close() error {
	q.shutdownLock.Lock()
	defer q.shutdownLock.Unlock()

	if q.shutdown {
		return nil
	}
	q.shutdown = true

	for i, pub := range q.publishers {
		if pub != nil {
			if err := pub.Close(); err != nil {
				return err
			}
		}
		q.publishers[i] = nil
	}

	return nil
}
<|endoftext|>


<commit_before>/*
Copyright 2021 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 common

import (
	"fmt"
	"time"

	"github.com/prometheus/common/model"

	"k8s.io/klog"
	"k8s.io/perf-tests/clusterloader2/pkg/errors"
	"k8s.io/perf-tests/clusterloader2/pkg/measurement"
	measurementutil "k8s.io/perf-tests/clusterloader2/pkg/measurement/util"
	"k8s.io/perf-tests/clusterloader2/pkg/util"
)

const (
	name = "GenericPrometheusQuery"
)

func init() {
	create := func() measurement.Measurement {
		return CreatePrometheusMeasurement(&genericQueryGatherer{})
	}
	if err := measurement.Register(name, create); err != nil {
		klog.Fatalf("Cannot register %s: %v", name, err)
	}
}

type genericQueryGatherer struct {
	metricName    string
	metricVersion string
	queries       []genericQuery
	unit          string
}

type genericQuery struct {
	name         string
	query        string
	hasThreshold bool
	threshold    float64
}

func (g *genericQueryGatherer) Configure(config *measurement.Config) error {
	metricName, err := util.GetString(config.Params, "metricName")
	if err != nil {
		return err
	}

	metricVersion, err := util.GetString(config.Params, "metricVersion")
	if err != nil {
		return err
	}

	unit, err := util.GetString(config.Params, "unit")
	if err != nil {
		return err
	}

	queries, err := util.GetMapArray(config.Params, "queries")
	if err != nil {
		return err
	}
	var genericQueries []genericQuery
	for _, q := range queries {
		query, err := util.GetString(q, "query")
		if err != nil {
			return err
		}
		name, err := util.GetString(q, "name")
		if err != nil {
			return err
		}
		hasThreshold := true
		threshold, err := util.GetFloat64(q, "threshold")
		if util.IsErrKeyNotFound(err) {
			klog.V(2).Infof("No threshold set for %v: %v", metricName, name)
			hasThreshold = false
		} else if err != nil {
			return err
		}
		genericQueries = append(genericQueries, genericQuery{name, query, hasThreshold, threshold})
	}

	g.metricName = metricName
	g.metricVersion = metricVersion
	g.unit = unit
	g.queries = genericQueries
	return nil
}

func (g *genericQueryGatherer) IsEnabled(config *measurement.Config) bool {
	return true
}

func (g *genericQueryGatherer) Gather(executor QueryExecutor, startTime, endTime time.Time, config *measurement.Config) ([]measurement.Summary, error) {
	var errs []error
	data := map[string]float64{}
	for _, q := range g.queries {
		samples, err := g.query(q, executor, startTime, endTime)
		if err != nil {
			return nil, err
		}

		if len(samples) > 1 {
			errs = append(errs, errors.NewMetricViolationError(q.name, fmt.Sprintf("too many samples: query returned %v streams, expected 1", len(samples))))
		}
		if len(samples) == 0 {
			klog.Warningf("query returned no samples for %v: %v", g.metricName, q.name)
			continue
		}

		val := float64(samples[0].Value)
		if q.hasThreshold && val > q.threshold {
			errs = append(errs, errors.NewMetricViolationError(q.name, fmt.Sprintf("sample above threshold: want: less or equal than %v, got: %v", q.threshold, val)))
		}

		data[q.name] = val
	}
	summary, err := g.createSummary(g.metricName, data)
	if err != nil {
		return nil, err
	}
	if len(errs) > 0 {
		err = errors.NewMetricViolationError(g.metricName, fmt.Sprintf("%v", errs))
	}
	return []measurement.Summary{summary}, err
}

func (g *genericQueryGatherer) String() string {
	return name
}

func (g *genericQueryGatherer) query(q genericQuery, executor QueryExecutor, startTime, endTime time.Time) ([]*model.Sample, error) {
	duration := endTime.Sub(startTime)
	boundedQuery := fmt.Sprintf(q.query, measurementutil.ToPrometheusTime(duration))
	klog.V(2).Infof("bounded query: %s, duration: %v", boundedQuery, duration)
	return executor.Query(boundedQuery, endTime)
}

func (g *genericQueryGatherer) createSummary(name string, data map[string]float64) (measurement.Summary, error) {
	content, err := util.PrettyPrintJSON(&measurementutil.PerfData{
		Version: g.metricVersion,
		DataItems: []measurementutil.DataItem{
			{
				Data: data,
				Unit: g.unit,
			},
		},
	})
	if err != nil {
		return nil, err
	}
	return measurement.CreateSummary(name, "json", content), nil
}
<commit_msg>Improve logging in generic query measurement<commit_after>/*
Copyright 2021 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 common

import (
	"fmt"
	"time"

	"github.com/prometheus/common/model"

	"k8s.io/klog"
	"k8s.io/perf-tests/clusterloader2/pkg/errors"
	"k8s.io/perf-tests/clusterloader2/pkg/measurement"
	measurementutil "k8s.io/perf-tests/clusterloader2/pkg/measurement/util"
	"k8s.io/perf-tests/clusterloader2/pkg/util"
)

const (
	name = "GenericPrometheusQuery"
)

func init() {
	create := func() measurement.Measurement {
		return CreatePrometheusMeasurement(&genericQueryGatherer{})
	}
	if err := measurement.Register(name, create); err != nil {
		klog.Fatalf("Cannot register %s: %v", name, err)
	}
}

type genericQueryGatherer struct {
	metricName    string
	metricVersion string
	queries       []genericQuery
	unit          string
}

type genericQuery struct {
	name         string
	query        string
	hasThreshold bool
	threshold    float64
}

func (g *genericQueryGatherer) Configure(config *measurement.Config) error {
	metricName, err := util.GetString(config.Params, "metricName")
	if err != nil {
		return err
	}

	metricVersion, err := util.GetString(config.Params, "metricVersion")
	if err != nil {
		return err
	}

	unit, err := util.GetString(config.Params, "unit")
	if err != nil {
		return err
	}

	queries, err := util.GetMapArray(config.Params, "queries")
	if err != nil {
		return err
	}
	var genericQueries []genericQuery
	for _, q := range queries {
		query, err := util.GetString(q, "query")
		if err != nil {
			return err
		}
		name, err := util.GetString(q, "name")
		if err != nil {
			return err
		}
		hasThreshold := true
		threshold, err := util.GetFloat64(q, "threshold")
		if util.IsErrKeyNotFound(err) {
			klog.V(2).Infof("No threshold set for %v: %v", metricName, name)
			hasThreshold = false
		} else if err != nil {
			return err
		}
		genericQueries = append(genericQueries, genericQuery{name, query, hasThreshold, threshold})
	}

	g.metricName = metricName
	g.metricVersion = metricVersion
	g.unit = unit
	g.queries = genericQueries
	return nil
}

func (g *genericQueryGatherer) IsEnabled(config *measurement.Config) bool {
	return true
}

func (g *genericQueryGatherer) Gather(executor QueryExecutor, startTime, endTime time.Time, config *measurement.Config) ([]measurement.Summary, error) {
	var errs []error
	data := map[string]float64{}
	for _, q := range g.queries {
		samples, err := g.query(q, executor, startTime, endTime)
		if err != nil {
			return nil, err
		}

		if len(samples) > 1 {
			errs = append(errs, errors.NewMetricViolationError(q.name, fmt.Sprintf("too many samples: query returned %v streams, expected 1", len(samples))))
		}
		if len(samples) == 0 {
			klog.Warningf("query returned no samples for %v: %v", g.metricName, q.name)
			continue
		}

		val := float64(samples[0].Value)

		thresholdMsg := "none"
		if q.hasThreshold {
			thresholdMsg = fmt.Sprintf("%v", q.threshold)
		}
		klog.V(2).Infof("metric: %v: %v, value: %v, threshold: %v", g.metricName, q.name, val, thresholdMsg)

		if q.hasThreshold && val > q.threshold {
			errs = append(errs, errors.NewMetricViolationError(q.name, fmt.Sprintf("sample above threshold: want: less or equal than %v, got: %v", q.threshold, val)))
		}

		data[q.name] = val
	}
	summary, err := g.createSummary(g.metricName, data)
	if err != nil {
		return nil, err
	}
	if len(errs) > 0 {
		err = errors.NewMetricViolationError(g.metricName, fmt.Sprintf("%v", errs))
	}
	return []measurement.Summary{summary}, err
}

func (g *genericQueryGatherer) String() string {
	return name
}

func (g *genericQueryGatherer) query(q genericQuery, executor QueryExecutor, startTime, endTime time.Time) ([]*model.Sample, error) {
	duration := endTime.Sub(startTime)
	boundedQuery := fmt.Sprintf(q.query, measurementutil.ToPrometheusTime(duration))
	klog.V(2).Infof("bounded query: %s, duration: %v", boundedQuery, duration)
	return executor.Query(boundedQuery, endTime)
}

func (g *genericQueryGatherer) createSummary(name string, data map[string]float64) (measurement.Summary, error) {
	content, err := util.PrettyPrintJSON(&measurementutil.PerfData{
		Version: g.metricVersion,
		DataItems: []measurementutil.DataItem{
			{
				Data: data,
				Unit: g.unit,
			},
		},
	})
	if err != nil {
		return nil, err
	}
	return measurement.CreateSummary(name, "json", content), nil
}
<|endoftext|>


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

package github

import (
	"context"
	"fmt"
	"net/http"
	"net/url"
	"reflect"
	"testing"
	"time"
)

func TestActionsService_ListWorkflowJobs(t *testing.T) {
	client, mux, _, teardown := setup()
	defer teardown()

	mux.HandleFunc("/repos/o/r/actions/runs/29679449/jobs", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, "GET")
		testFormValues(t, r, values{"per_page": "2", "page": "2"})
		fmt.Fprint(w, `{"total_count":4,"jobs":[{"id":399444496,"run_id":29679449,"started_at":"2019-01-02T15:04:05Z","completed_at":"2020-01-02T15:04:05Z"},{"id":399444497,"run_id":29679449,"started_at":"2019-01-02T15:04:05Z","completed_at":"2020-01-02T15:04:05Z"}]}`)
	})

	opts := &ListWorkflowJobsOptions{ListOptions: ListOptions{Page: 2, PerPage: 2}}
	ctx := context.Background()
	jobs, _, err := client.Actions.ListWorkflowJobs(ctx, "o", "r", 29679449, opts)
	if err != nil {
		t.Errorf("Actions.ListWorkflowJobs returned error: %v", err)
	}

	want := &Jobs{
		TotalCount: Int(4),
		Jobs: []*WorkflowJob{
			{ID: Int64(399444496), RunID: Int64(29679449), StartedAt: &Timestamp{time.Date(2019, time.January, 02, 15, 04, 05, 0, time.UTC)}, CompletedAt: &Timestamp{time.Date(2020, time.January, 02, 15, 04, 05, 0, time.UTC)}},
			{ID: Int64(399444497), RunID: Int64(29679449), StartedAt: &Timestamp{time.Date(2019, time.January, 02, 15, 04, 05, 0, time.UTC)}, CompletedAt: &Timestamp{time.Date(2020, time.January, 02, 15, 04, 05, 0, time.UTC)}},
		},
	}
	if !reflect.DeepEqual(jobs, want) {
		t.Errorf("Actions.ListWorkflowJobs returned %+v, want %+v", jobs, want)
	}
}

func TestActionsService_ListWorkflowJobs_Filter(t *testing.T) {
	client, mux, _, teardown := setup()
	defer teardown()

	mux.HandleFunc("/repos/o/r/actions/runs/29679449/jobs", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, "GET")
		testFormValues(t, r, values{"filter": "all", "per_page": "2", "page": "2"})
		fmt.Fprint(w, `{"total_count":4,"jobs":[{"id":399444496,"run_id":29679449,"started_at":"2019-01-02T15:04:05Z","completed_at":"2020-01-02T15:04:05Z"},{"id":399444497,"run_id":29679449,"started_at":"2019-01-02T15:04:05Z","completed_at":"2020-01-02T15:04:05Z"}]}`)
	})

	opts := &ListWorkflowJobsOptions{Filter: "all", ListOptions: ListOptions{Page: 2, PerPage: 2}}
	ctx := context.Background()
	jobs, _, err := client.Actions.ListWorkflowJobs(ctx, "o", "r", 29679449, opts)
	if err != nil {
		t.Errorf("Actions.ListWorkflowJobs returned error: %v", err)
	}

	want := &Jobs{
		TotalCount: Int(4),
		Jobs: []*WorkflowJob{
			{ID: Int64(399444496), RunID: Int64(29679449), StartedAt: &Timestamp{time.Date(2019, time.January, 02, 15, 04, 05, 0, time.UTC)}, CompletedAt: &Timestamp{time.Date(2020, time.January, 02, 15, 04, 05, 0, time.UTC)}},
			{ID: Int64(399444497), RunID: Int64(29679449), StartedAt: &Timestamp{time.Date(2019, time.January, 02, 15, 04, 05, 0, time.UTC)}, CompletedAt: &Timestamp{time.Date(2020, time.January, 02, 15, 04, 05, 0, time.UTC)}},
		},
	}
	if !reflect.DeepEqual(jobs, want) {
		t.Errorf("Actions.ListWorkflowJobs returned %+v, want %+v", jobs, want)
	}
}

func TestActionsService_GetWorkflowJobByID(t *testing.T) {
	client, mux, _, teardown := setup()
	defer teardown()

	mux.HandleFunc("/repos/o/r/actions/jobs/399444496", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, "GET")
		fmt.Fprint(w, `{"id":399444496,"started_at":"2019-01-02T15:04:05Z","completed_at":"2020-01-02T15:04:05Z"}`)
	})

	ctx := context.Background()
	job, _, err := client.Actions.GetWorkflowJobByID(ctx, "o", "r", 399444496)
	if err != nil {
		t.Errorf("Actions.GetWorkflowJobByID returned error: %v", err)
	}

	want := &WorkflowJob{
		ID:          Int64(399444496),
		StartedAt:   &Timestamp{time.Date(2019, time.January, 02, 15, 04, 05, 0, time.UTC)},
		CompletedAt: &Timestamp{time.Date(2020, time.January, 02, 15, 04, 05, 0, time.UTC)},
	}
	if !reflect.DeepEqual(job, want) {
		t.Errorf("Actions.GetWorkflowJobByID returned %+v, want %+v", job, want)
	}
}

func TestActionsService_GetWorkflowJobLogs(t *testing.T) {
	client, mux, _, teardown := setup()
	defer teardown()

	mux.HandleFunc("/repos/o/r/actions/jobs/399444496/logs", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, "GET")
		http.Redirect(w, r, "http://github.com/a", http.StatusFound)
	})

	ctx := context.Background()
	url, resp, err := client.Actions.GetWorkflowJobLogs(ctx, "o", "r", 399444496, true)
	if err != nil {
		t.Errorf("Actions.GetWorkflowJobLogs returned error: %v", err)
	}
	if resp.StatusCode != http.StatusFound {
		t.Errorf("Actions.GetWorkflowJobLogs returned status: %d, want %d", resp.StatusCode, http.StatusFound)
	}
	want := "http://github.com/a"
	if url.String() != want {
		t.Errorf("Actions.GetWorkflowJobLogs returned %+v, want %+v", url.String(), want)
	}
}

func TestActionsService_GetWorkflowJobLogs_StatusMovedPermanently_dontFollowRedirects(t *testing.T) {
	client, mux, _, teardown := setup()
	defer teardown()

	mux.HandleFunc("/repos/o/r/actions/jobs/399444496/logs", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, "GET")
		http.Redirect(w, r, "http://github.com/a", http.StatusMovedPermanently)
	})

	ctx := context.Background()
	_, resp, _ := client.Actions.GetWorkflowJobLogs(ctx, "o", "r", 399444496, false)
	if resp.StatusCode != http.StatusMovedPermanently {
		t.Errorf("Actions.GetWorkflowJobLogs returned status: %d, want %d", resp.StatusCode, http.StatusMovedPermanently)
	}
}

func TestActionsService_GetWorkflowJobLogs_StatusMovedPermanently_followRedirects(t *testing.T) {
	client, mux, serverURL, teardown := setup()
	defer teardown()

	// Mock a redirect link, which leads to an archive link
	mux.HandleFunc("/repos/o/r/actions/jobs/399444496/logs", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, "GET")
		redirectURL, _ := url.Parse(serverURL + baseURLPath + "/redirect")
		http.Redirect(w, r, redirectURL.String(), http.StatusMovedPermanently)
	})

	mux.HandleFunc("/redirect", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, "GET")
		http.Redirect(w, r, "http://github.com/a", http.StatusFound)
	})

	ctx := context.Background()
	url, resp, err := client.Actions.GetWorkflowJobLogs(ctx, "o", "r", 399444496, true)
	if err != nil {
		t.Errorf("Actions.GetWorkflowJobLogs returned error: %v", err)
	}

	if resp.StatusCode != http.StatusFound {
		t.Errorf("Actions.GetWorkflowJobLogs returned status: %d, want %d", resp.StatusCode, http.StatusFound)
	}

	want := "http://github.com/a"
	if url.String() != want {
		t.Errorf("Actions.GetWorkflowJobLogs returned %+v, want %+v", url.String(), want)
	}
}
<commit_msg>Improve actions_workflow_jobs.go coverage (#1701)<commit_after>// Copyright 2020 The go-github AUTHORS. All rights reserved.
//
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package github

import (
	"context"
	"fmt"
	"net/http"
	"net/url"
	"reflect"
	"testing"
	"time"
)

func TestActionsService_ListWorkflowJobs(t *testing.T) {
	client, mux, _, teardown := setup()
	defer teardown()

	mux.HandleFunc("/repos/o/r/actions/runs/29679449/jobs", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, "GET")
		testFormValues(t, r, values{"per_page": "2", "page": "2"})
		fmt.Fprint(w, `{"total_count":4,"jobs":[{"id":399444496,"run_id":29679449,"started_at":"2019-01-02T15:04:05Z","completed_at":"2020-01-02T15:04:05Z"},{"id":399444497,"run_id":29679449,"started_at":"2019-01-02T15:04:05Z","completed_at":"2020-01-02T15:04:05Z"}]}`)
	})

	opts := &ListWorkflowJobsOptions{ListOptions: ListOptions{Page: 2, PerPage: 2}}
	ctx := context.Background()
	jobs, _, err := client.Actions.ListWorkflowJobs(ctx, "o", "r", 29679449, opts)
	if err != nil {
		t.Errorf("Actions.ListWorkflowJobs returned error: %v", err)
	}

	want := &Jobs{
		TotalCount: Int(4),
		Jobs: []*WorkflowJob{
			{ID: Int64(399444496), RunID: Int64(29679449), StartedAt: &Timestamp{time.Date(2019, time.January, 02, 15, 04, 05, 0, time.UTC)}, CompletedAt: &Timestamp{time.Date(2020, time.January, 02, 15, 04, 05, 0, time.UTC)}},
			{ID: Int64(399444497), RunID: Int64(29679449), StartedAt: &Timestamp{time.Date(2019, time.January, 02, 15, 04, 05, 0, time.UTC)}, CompletedAt: &Timestamp{time.Date(2020, time.January, 02, 15, 04, 05, 0, time.UTC)}},
		},
	}
	if !reflect.DeepEqual(jobs, want) {
		t.Errorf("Actions.ListWorkflowJobs returned %+v, want %+v", jobs, want)
	}

	const methodName = "ListWorkflowJobs"
	testBadOptions(t, methodName, func() (err error) {
		_, _, err = client.Actions.ListWorkflowJobs(ctx, "\n", "\n", 29679449, opts)
		return err
	})

	testNewRequestAndDoFailure(t, methodName, client, func() (*Response, error) {
		got, resp, err := client.Actions.ListWorkflowJobs(ctx, "o", "r", 29679449, opts)
		if got != nil {
			t.Errorf("testNewRequestAndDoFailure %v = %#v, want nil", methodName, got)
		}
		return resp, err
	})
}

func TestActionsService_ListWorkflowJobs_Filter(t *testing.T) {
	client, mux, _, teardown := setup()
	defer teardown()

	mux.HandleFunc("/repos/o/r/actions/runs/29679449/jobs", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, "GET")
		testFormValues(t, r, values{"filter": "all", "per_page": "2", "page": "2"})
		fmt.Fprint(w, `{"total_count":4,"jobs":[{"id":399444496,"run_id":29679449,"started_at":"2019-01-02T15:04:05Z","completed_at":"2020-01-02T15:04:05Z"},{"id":399444497,"run_id":29679449,"started_at":"2019-01-02T15:04:05Z","completed_at":"2020-01-02T15:04:05Z"}]}`)
	})

	opts := &ListWorkflowJobsOptions{Filter: "all", ListOptions: ListOptions{Page: 2, PerPage: 2}}
	ctx := context.Background()
	jobs, _, err := client.Actions.ListWorkflowJobs(ctx, "o", "r", 29679449, opts)
	if err != nil {
		t.Errorf("Actions.ListWorkflowJobs returned error: %v", err)
	}

	want := &Jobs{
		TotalCount: Int(4),
		Jobs: []*WorkflowJob{
			{ID: Int64(399444496), RunID: Int64(29679449), StartedAt: &Timestamp{time.Date(2019, time.January, 02, 15, 04, 05, 0, time.UTC)}, CompletedAt: &Timestamp{time.Date(2020, time.January, 02, 15, 04, 05, 0, time.UTC)}},
			{ID: Int64(399444497), RunID: Int64(29679449), StartedAt: &Timestamp{time.Date(2019, time.January, 02, 15, 04, 05, 0, time.UTC)}, CompletedAt: &Timestamp{time.Date(2020, time.January, 02, 15, 04, 05, 0, time.UTC)}},
		},
	}
	if !reflect.DeepEqual(jobs, want) {
		t.Errorf("Actions.ListWorkflowJobs returned %+v, want %+v", jobs, want)
	}
}

func TestActionsService_GetWorkflowJobByID(t *testing.T) {
	client, mux, _, teardown := setup()
	defer teardown()

	mux.HandleFunc("/repos/o/r/actions/jobs/399444496", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, "GET")
		fmt.Fprint(w, `{"id":399444496,"started_at":"2019-01-02T15:04:05Z","completed_at":"2020-01-02T15:04:05Z"}`)
	})

	ctx := context.Background()
	job, _, err := client.Actions.GetWorkflowJobByID(ctx, "o", "r", 399444496)
	if err != nil {
		t.Errorf("Actions.GetWorkflowJobByID returned error: %v", err)
	}

	want := &WorkflowJob{
		ID:          Int64(399444496),
		StartedAt:   &Timestamp{time.Date(2019, time.January, 02, 15, 04, 05, 0, time.UTC)},
		CompletedAt: &Timestamp{time.Date(2020, time.January, 02, 15, 04, 05, 0, time.UTC)},
	}
	if !reflect.DeepEqual(job, want) {
		t.Errorf("Actions.GetWorkflowJobByID returned %+v, want %+v", job, want)
	}

	const methodName = "GetWorkflowJobByID"
	testBadOptions(t, methodName, func() (err error) {
		_, _, err = client.Actions.GetWorkflowJobByID(ctx, "\n", "\n", 399444496)
		return err
	})

	testNewRequestAndDoFailure(t, methodName, client, func() (*Response, error) {
		got, resp, err := client.Actions.GetWorkflowJobByID(ctx, "o", "r", 399444496)
		if got != nil {
			t.Errorf("testNewRequestAndDoFailure %v = %#v, want nil", methodName, got)
		}
		return resp, err
	})
}

func TestActionsService_GetWorkflowJobLogs(t *testing.T) {
	client, mux, _, teardown := setup()
	defer teardown()

	mux.HandleFunc("/repos/o/r/actions/jobs/399444496/logs", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, "GET")
		http.Redirect(w, r, "http://github.com/a", http.StatusFound)
	})

	ctx := context.Background()
	url, resp, err := client.Actions.GetWorkflowJobLogs(ctx, "o", "r", 399444496, true)
	if err != nil {
		t.Errorf("Actions.GetWorkflowJobLogs returned error: %v", err)
	}
	if resp.StatusCode != http.StatusFound {
		t.Errorf("Actions.GetWorkflowJobLogs returned status: %d, want %d", resp.StatusCode, http.StatusFound)
	}
	want := "http://github.com/a"
	if url.String() != want {
		t.Errorf("Actions.GetWorkflowJobLogs returned %+v, want %+v", url.String(), want)
	}

	const methodName = "GetWorkflowJobLogs"
	testBadOptions(t, methodName, func() (err error) {
		_, _, err = client.Actions.GetWorkflowJobLogs(ctx, "\n", "\n", 399444496, true)
		return err
	})
}

func TestActionsService_GetWorkflowJobLogs_StatusMovedPermanently_dontFollowRedirects(t *testing.T) {
	client, mux, _, teardown := setup()
	defer teardown()

	mux.HandleFunc("/repos/o/r/actions/jobs/399444496/logs", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, "GET")
		http.Redirect(w, r, "http://github.com/a", http.StatusMovedPermanently)
	})

	ctx := context.Background()
	_, resp, _ := client.Actions.GetWorkflowJobLogs(ctx, "o", "r", 399444496, false)
	if resp.StatusCode != http.StatusMovedPermanently {
		t.Errorf("Actions.GetWorkflowJobLogs returned status: %d, want %d", resp.StatusCode, http.StatusMovedPermanently)
	}
}

func TestActionsService_GetWorkflowJobLogs_StatusMovedPermanently_followRedirects(t *testing.T) {
	client, mux, serverURL, teardown := setup()
	defer teardown()

	// Mock a redirect link, which leads to an archive link
	mux.HandleFunc("/repos/o/r/actions/jobs/399444496/logs", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, "GET")
		redirectURL, _ := url.Parse(serverURL + baseURLPath + "/redirect")
		http.Redirect(w, r, redirectURL.String(), http.StatusMovedPermanently)
	})

	mux.HandleFunc("/redirect", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, "GET")
		http.Redirect(w, r, "http://github.com/a", http.StatusFound)
	})

	ctx := context.Background()
	url, resp, err := client.Actions.GetWorkflowJobLogs(ctx, "o", "r", 399444496, true)
	if err != nil {
		t.Errorf("Actions.GetWorkflowJobLogs returned error: %v", err)
	}

	if resp.StatusCode != http.StatusFound {
		t.Errorf("Actions.GetWorkflowJobLogs returned status: %d, want %d", resp.StatusCode, http.StatusFound)
	}

	want := "http://github.com/a"
	if url.String() != want {
		t.Errorf("Actions.GetWorkflowJobLogs returned %+v, want %+v", url.String(), want)
	}
}
<|endoftext|>


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

package types2

// sanitizeInfo walks the types contained in info to ensure that all instances
// are expanded.
//
// This includes some objects that may be shared across concurrent
// type-checking passes (such as those in the universe scope), so we are
// careful here not to write types that are already sanitized. This avoids a
// data race as any shared types should already be sanitized.
func sanitizeInfo(info *Info) {
	var s sanitizer = make(map[Type]Type)

	// Note: Some map entries are not references.
	// If modified, they must be assigned back.

	for e, tv := range info.Types {
		if typ := s.typ(tv.Type); typ != tv.Type {
			tv.Type = typ
			info.Types[e] = tv
		}
	}

	for e, inf := range info.Inferred {
		changed := false
		for i, targ := range inf.Targs {
			if typ := s.typ(targ); typ != targ {
				inf.Targs[i] = typ
				changed = true
			}
		}
		if typ := s.typ(inf.Sig); typ != inf.Sig {
			inf.Sig = typ.(*Signature)
			changed = true
		}
		if changed {
			info.Inferred[e] = inf
		}
	}

	for _, obj := range info.Defs {
		if obj != nil {
			if typ := s.typ(obj.Type()); typ != obj.Type() {
				obj.setType(typ)
			}
		}
	}

	for _, obj := range info.Uses {
		if obj != nil {
			if typ := s.typ(obj.Type()); typ != obj.Type() {
				obj.setType(typ)
			}
		}
	}

	// TODO(gri) sanitize as needed
	// - info.Implicits
	// - info.Selections
	// - info.Scopes
	// - info.InitOrder
}

type sanitizer map[Type]Type

func (s sanitizer) typ(typ Type) Type {
	if t, found := s[typ]; found {
		return t
	}
	s[typ] = typ

	switch t := typ.(type) {
	case nil, *Basic, *bottom, *top:
		// nothing to do

	case *Array:
		if elem := s.typ(t.elem); elem != t.elem {
			t.elem = elem
		}

	case *Slice:
		if elem := s.typ(t.elem); elem != t.elem {
			t.elem = elem
		}

	case *Struct:
		s.varList(t.fields)

	case *Pointer:
		if base := s.typ(t.base); base != t.base {
			t.base = base
		}

	case *Tuple:
		s.tuple(t)

	case *Signature:
		s.var_(t.recv)
		s.tuple(t.params)
		s.tuple(t.results)

	case *Sum:
		s.typeList(t.types)

	case *Interface:
		s.funcList(t.methods)
		s.typ(t.types)
		s.typeList(t.embeddeds)
		s.funcList(t.allMethods)
		s.typ(t.allTypes)

	case *Map:
		if key := s.typ(t.key); key != t.key {
			t.key = key
		}
		if elem := s.typ(t.elem); elem != t.elem {
			t.elem = elem
		}

	case *Chan:
		if elem := s.typ(t.elem); elem != t.elem {
			t.elem = elem
		}

	case *Named:
		if orig := s.typ(t.fromRHS); orig != t.fromRHS {
			t.fromRHS = orig
		}
		if under := s.typ(t.underlying); under != t.underlying {
			t.underlying = under
		}
		s.typeList(t.targs)
		s.funcList(t.methods)

	case *TypeParam:
		if bound := s.typ(t.bound); bound != t.bound {
			t.bound = bound
		}

	case *instance:
		typ = t.expand()
		s[t] = typ

	default:
		panic("unimplemented")
	}

	return typ
}

func (s sanitizer) var_(v *Var) {
	if v != nil {
		if typ := s.typ(v.typ); typ != v.typ {
			v.typ = typ
		}
	}
}

func (s sanitizer) varList(list []*Var) {
	for _, v := range list {
		s.var_(v)
	}
}

func (s sanitizer) tuple(t *Tuple) {
	if t != nil {
		s.varList(t.vars)
	}
}

func (s sanitizer) func_(f *Func) {
	if f != nil {
		if typ := s.typ(f.typ); typ != f.typ {
			f.typ = typ
		}
	}
}

func (s sanitizer) funcList(list []*Func) {
	for _, f := range list {
		s.func_(f)
	}
}

func (s sanitizer) typeList(list []Type) {
	for i, t := range list {
		if typ := s.typ(t); typ != t {
			list[i] = typ
		}
	}
}
<commit_msg>cmd/compile/internal/types2: fix potential bugs in santitizer pass<commit_after>// Copyright 2020 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package types2

// sanitizeInfo walks the types contained in info to ensure that all instances
// are expanded.
//
// This includes some objects that may be shared across concurrent
// type-checking passes (such as those in the universe scope), so we are
// careful here not to write types that are already sanitized. This avoids a
// data race as any shared types should already be sanitized.
func sanitizeInfo(info *Info) {
	var s sanitizer = make(map[Type]Type)

	// Note: Some map entries are not references.
	// If modified, they must be assigned back.

	for e, tv := range info.Types {
		if typ := s.typ(tv.Type); typ != tv.Type {
			tv.Type = typ
			info.Types[e] = tv
		}
	}

	for e, inf := range info.Inferred {
		changed := false
		for i, targ := range inf.Targs {
			if typ := s.typ(targ); typ != targ {
				inf.Targs[i] = typ
				changed = true
			}
		}
		if typ := s.typ(inf.Sig); typ != inf.Sig {
			inf.Sig = typ.(*Signature)
			changed = true
		}
		if changed {
			info.Inferred[e] = inf
		}
	}

	for _, obj := range info.Defs {
		if obj != nil {
			if typ := s.typ(obj.Type()); typ != obj.Type() {
				obj.setType(typ)
			}
		}
	}

	for _, obj := range info.Uses {
		if obj != nil {
			if typ := s.typ(obj.Type()); typ != obj.Type() {
				obj.setType(typ)
			}
		}
	}

	// TODO(gri) sanitize as needed
	// - info.Implicits
	// - info.Selections
	// - info.Scopes
	// - info.InitOrder
}

type sanitizer map[Type]Type

func (s sanitizer) typ(typ Type) Type {
	if typ == nil {
		return nil
	}

	if t, found := s[typ]; found {
		return t
	}
	s[typ] = typ

	switch t := typ.(type) {
	case *Basic, *bottom, *top:
		// nothing to do

	case *Array:
		if elem := s.typ(t.elem); elem != t.elem {
			t.elem = elem
		}

	case *Slice:
		if elem := s.typ(t.elem); elem != t.elem {
			t.elem = elem
		}

	case *Struct:
		s.varList(t.fields)

	case *Pointer:
		if base := s.typ(t.base); base != t.base {
			t.base = base
		}

	case *Tuple:
		s.tuple(t)

	case *Signature:
		s.var_(t.recv)
		s.tuple(t.params)
		s.tuple(t.results)

	case *Sum:
		s.typeList(t.types)

	case *Interface:
		s.funcList(t.methods)
		if types := s.typ(t.types); types != t.types {
			t.types = types
		}
		s.typeList(t.embeddeds)
		s.funcList(t.allMethods)
		if allTypes := s.typ(t.allTypes); allTypes != t.allTypes {
			t.allTypes = allTypes
		}

	case *Map:
		if key := s.typ(t.key); key != t.key {
			t.key = key
		}
		if elem := s.typ(t.elem); elem != t.elem {
			t.elem = elem
		}

	case *Chan:
		if elem := s.typ(t.elem); elem != t.elem {
			t.elem = elem
		}

	case *Named:
		if orig := s.typ(t.fromRHS); orig != t.fromRHS {
			t.fromRHS = orig
		}
		if under := s.typ(t.underlying); under != t.underlying {
			t.underlying = under
		}
		s.typeList(t.targs)
		s.funcList(t.methods)

	case *TypeParam:
		if bound := s.typ(t.bound); bound != t.bound {
			t.bound = bound
		}

	case *instance:
		typ = t.expand()
		s[t] = typ

	default:
		panic("unimplemented")
	}

	return typ
}

func (s sanitizer) var_(v *Var) {
	if v != nil {
		if typ := s.typ(v.typ); typ != v.typ {
			v.typ = typ
		}
	}
}

func (s sanitizer) varList(list []*Var) {
	for _, v := range list {
		s.var_(v)
	}
}

func (s sanitizer) tuple(t *Tuple) {
	if t != nil {
		s.varList(t.vars)
	}
}

func (s sanitizer) func_(f *Func) {
	if f != nil {
		if typ := s.typ(f.typ); typ != f.typ {
			f.typ = typ
		}
	}
}

func (s sanitizer) funcList(list []*Func) {
	for _, f := range list {
		s.func_(f)
	}
}

func (s sanitizer) typeList(list []Type) {
	for i, t := range list {
		if typ := s.typ(t); typ != t {
			list[i] = typ
		}
	}
}
<|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 vtexplain

import (
	"encoding/json"
	"fmt"
	"os"
	"path"
	"strings"
	"testing"

	"github.com/google/go-cmp/cmp"
	"github.com/stretchr/testify/require"

	"vitess.io/vitess/go/vt/key"
	"vitess.io/vitess/go/vt/proto/topodata"
	"vitess.io/vitess/go/vt/topo"
)

var testOutputTempDir string

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

type testopts struct {
	shardmap map[string]map[string]*topo.ShardInfo
}

func initTest(mode string, opts *Options, topts *testopts, t *testing.T) {
	schema, err := os.ReadFile("testdata/test-schema.sql")
	require.NoError(t, err)

	vSchema, err := os.ReadFile("testdata/test-vschema.json")
	require.NoError(t, err)

	shardmap := ""
	if topts.shardmap != nil {
		shardmapBytes, err := json.Marshal(topts.shardmap)
		require.NoError(t, err)

		shardmap = string(shardmapBytes)
	}

	opts.ExecutionMode = mode
	err = Init(string(vSchema), string(schema), shardmap, opts)
	require.NoError(t, err, "vtexplain Init error\n%s", string(schema))
}

func testExplain(testcase string, opts *Options, t *testing.T) {
	modes := []string{
		ModeMulti,

		// TwoPC mode is functional, but the output isn't stable for
		// tests since there are timestamps in the value rows
		// ModeTwoPC,
	}

	for _, mode := range modes {
		runTestCase(testcase, mode, opts, &testopts{}, t)
	}
}

func runTestCase(testcase, mode string, opts *Options, topts *testopts, t *testing.T) {
	t.Run(testcase, func(t *testing.T) {
		initTest(mode, opts, topts, t)

		sqlFile := fmt.Sprintf("testdata/%s-queries.sql", testcase)
		sql, err := os.ReadFile(sqlFile)
		require.NoError(t, err, "vtexplain error")

		textOutFile := fmt.Sprintf("testdata/%s-output/%s-output.txt", mode, testcase)
		expected, _ := os.ReadFile(textOutFile)

		explains, err := Run(string(sql))
		require.NoError(t, err, "vtexplain error")
		require.NotNil(t, explains, "vtexplain error running %s: no explain", string(sql))

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

			if testOutputTempDir == "" {
				testOutputTempDir, err = os.MkdirTemp("", "vtexplain_output")
				require.NoError(t, err, "error getting tempdir")
			}
			gotFile := fmt.Sprintf("%s/%s-output.txt", testOutputTempDir, testcase)
			os.WriteFile(gotFile, []byte(explainText), 0644)

			t.Logf("run the following command to update the expected output:")
			t.Logf("cp %s/* %s", testOutputTempDir, path.Dir(textOutFile))
		}
	})
}

func TestExplain(t *testing.T) {
	type test struct {
		name string
		opts *Options
	}
	tests := []test{
		{"unsharded", defaultTestOpts()},
		{"selectsharded", defaultTestOpts()},
		{"insertsharded", defaultTestOpts()},
		{"updatesharded", defaultTestOpts()},
		{"deletesharded", defaultTestOpts()},
		{"comments", defaultTestOpts()},
		{"options", &Options{
			ReplicationMode: "STATEMENT",
			NumShards:       4,
			Normalize:       false,
		}},
		{"target", &Options{
			ReplicationMode: "ROW",
			NumShards:       4,
			Normalize:       false,
			Target:          "ks_sharded/40-80",
		}},
	}

	for _, tst := range tests {
		testExplain(tst.name, tst.opts, t)
	}
}

func TestErrors(t *testing.T) {
	initTest(ModeMulti, defaultTestOpts(), &testopts{}, t)

	tests := []struct {
		SQL string
		Err string
	}{
		{
			SQL: "INVALID SQL",
			Err: "vtexplain execute error in 'INVALID SQL': syntax error at position 8 near 'INVALID'",
		},

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

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

		{
			SQL: "SELECT * FROM table_not_in_schema",
			Err: "unknown error: unable to resolve table name table_not_in_schema",
		},
	}

	for _, test := range tests {
		t.Run(test.SQL, func(t *testing.T) {
			_, err := Run(test.SQL)
			require.Error(t, err)
			require.Contains(t, err.Error(), test.Err)
		})
	}
}

func TestJSONOutput(t *testing.T) {
	sql := "select 1 from user where id = 1"
	explains, err := Run(sql)
	require.NoError(t, err, "vtexplain error")
	require.NotNil(t, explains, "vtexplain error running %s: no explain", string(sql))

	explainJSON := ExplainsAsJSON(explains)

	var data interface{}
	err = json.Unmarshal([]byte(explainJSON), &data)
	require.NoError(t, err, "error unmarshaling json")

	array, ok := data.([]interface{})
	if !ok || len(array) != 1 {
		t.Errorf("expected single-element top-level array, got:\n%s", explainJSON)
	}

	explain, ok := array[0].(map[string]interface{})
	if !ok {
		t.Errorf("expected explain map, got:\n%s", explainJSON)
	}

	if explain["SQL"] != sql {
		t.Errorf("expected SQL, got:\n%s", explainJSON)
	}

	plans, ok := explain["Plans"].([]interface{})
	if !ok || len(plans) != 1 {
		t.Errorf("expected single-element plans array, got:\n%s", explainJSON)
	}

	actions, ok := explain["TabletActions"].(map[string]interface{})
	if !ok {
		t.Errorf("expected TabletActions map, got:\n%s", explainJSON)
	}

	actionsJSON, err := json.MarshalIndent(actions, "", "    ")
	require.NoError(t, err, "error in json marshal")
	wantJSON := `{
    "ks_sharded/-40": {
        "MysqlQueries": [
            {
                "SQL": "select 1 from ` + "`user`" + ` where id = 1 limit 10001",
                "Time": 1
            }
        ],
        "TabletQueries": [
            {
                "BindVars": {
                    "#maxLimit": "10001",
                    "vtg1": "1"
                },
                "SQL": "select :vtg1 from ` + "`user`" + ` where id = :vtg1",
                "Time": 1
            }
        ]
    }
}`
	diff := cmp.Diff(wantJSON, string(actionsJSON))
	if diff != "" {
		t.Errorf(diff)
	}
}

func testShardInfo(ks, start, end string, t *testing.T) *topo.ShardInfo {
	kr, err := key.ParseKeyRangeParts(start, end)
	require.NoError(t, err)

	return topo.NewShardInfo(
		ks,
		fmt.Sprintf("%s-%s", start, end),
		&topodata.Shard{KeyRange: kr},
		&vtexplainTestTopoVersion{},
	)
}

func TestUsingKeyspaceShardMap(t *testing.T) {
	tests := []struct {
		testcase      string
		ShardRangeMap map[string]map[string]*topo.ShardInfo
	}{
		{
			testcase: "select-sharded-8",
			ShardRangeMap: map[string]map[string]*topo.ShardInfo{
				"ks_sharded": {
					"-20":   testShardInfo("ks_sharded", "", "20", t),
					"20-40": testShardInfo("ks_sharded", "20", "40", t),
					"40-60": testShardInfo("ks_sharded", "40", "60", t),
					"60-80": testShardInfo("ks_sharded", "60", "80", t),
					"80-a0": testShardInfo("ks_sharded", "80", "a0", t),
					"a0-c0": testShardInfo("ks_sharded", "a0", "c0", t),
					"c0-e0": testShardInfo("ks_sharded", "c0", "e0", t),
					"e0-":   testShardInfo("ks_sharded", "e0", "", t),
				},
			},
		},
		{
			testcase: "uneven-keyspace",
			ShardRangeMap: map[string]map[string]*topo.ShardInfo{
				// Have mercy on the poor soul that has this keyspace sharding.
				// But, hey, vtexplain still works so they have that going for them.
				"ks_sharded": {
					"-80":   testShardInfo("ks_sharded", "", "80", t),
					"80-90": testShardInfo("ks_sharded", "80", "90", t),
					"90-a0": testShardInfo("ks_sharded", "90", "a0", t),
					"a0-e8": testShardInfo("ks_sharded", "a0", "e8", t),
					"e8-":   testShardInfo("ks_sharded", "e8", "", t),
				},
			},
		},
	}

	for _, test := range tests {
		runTestCase(test.testcase, ModeMulti, defaultTestOpts(), &testopts{test.ShardRangeMap}, t)
	}
}

type vtexplainTestTopoVersion struct{}

func (vtexplain *vtexplainTestTopoVersion) String() string { return "vtexplain-test-topo" }
<commit_msg>removed additional set collation_connection queries in vtexplain unit tests<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 vtexplain

import (
	"encoding/json"
	"fmt"
	"os"
	"path"
	"strings"
	"testing"

	"github.com/google/go-cmp/cmp"
	"github.com/stretchr/testify/require"

	"vitess.io/vitess/go/vt/key"
	"vitess.io/vitess/go/vt/proto/topodata"
	"vitess.io/vitess/go/vt/topo"
)

var testOutputTempDir string

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

type testopts struct {
	shardmap map[string]map[string]*topo.ShardInfo
}

func initTest(mode string, opts *Options, topts *testopts, t *testing.T) {
	schema, err := os.ReadFile("testdata/test-schema.sql")
	require.NoError(t, err)

	vSchema, err := os.ReadFile("testdata/test-vschema.json")
	require.NoError(t, err)

	shardmap := ""
	if topts.shardmap != nil {
		shardmapBytes, err := json.Marshal(topts.shardmap)
		require.NoError(t, err)

		shardmap = string(shardmapBytes)
	}

	opts.ExecutionMode = mode
	err = Init(string(vSchema), string(schema), shardmap, opts)
	require.NoError(t, err, "vtexplain Init error\n%s", string(schema))
}

func testExplain(testcase string, opts *Options, t *testing.T) {
	modes := []string{
		ModeMulti,

		// TwoPC mode is functional, but the output isn't stable for
		// tests since there are timestamps in the value rows
		// ModeTwoPC,
	}

	for _, mode := range modes {
		runTestCase(testcase, mode, opts, &testopts{}, t)
	}
}

func runTestCase(testcase, mode string, opts *Options, topts *testopts, t *testing.T) {
	t.Run(testcase, func(t *testing.T) {
		initTest(mode, opts, topts, t)

		sqlFile := fmt.Sprintf("testdata/%s-queries.sql", testcase)
		sql, err := os.ReadFile(sqlFile)
		require.NoError(t, err, "vtexplain error")

		textOutFile := fmt.Sprintf("testdata/%s-output/%s-output.txt", mode, testcase)
		expected, _ := os.ReadFile(textOutFile)

		explains, err := Run(string(sql))
		require.NoError(t, err, "vtexplain error")
		require.NotNil(t, explains, "vtexplain error running %s: no explain", string(sql))

		// We want to remove the additional `set collation_connection` queries that happen
		// when the tablet connects to MySQL to set the default collation.
		// Removing them lets us keep simpler expected output files.
		for _, e := range explains {
			for i, action := range e.TabletActions {
				var mysqlQueries []*MysqlQuery
				for _, query := range action.MysqlQueries {
					if !strings.Contains(strings.ToLower(query.SQL), "set collation_connection") {
						mysqlQueries = append(mysqlQueries, query)
					}
				}
				e.TabletActions[i].MysqlQueries = mysqlQueries
			}
		}

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

			if testOutputTempDir == "" {
				testOutputTempDir, err = os.MkdirTemp("", "vtexplain_output")
				require.NoError(t, err, "error getting tempdir")
			}
			gotFile := fmt.Sprintf("%s/%s-output.txt", testOutputTempDir, testcase)
			os.WriteFile(gotFile, []byte(explainText), 0644)

			t.Logf("run the following command to update the expected output:")
			t.Logf("cp %s/* %s", testOutputTempDir, path.Dir(textOutFile))
		}
	})
}

func TestExplain(t *testing.T) {
	type test struct {
		name string
		opts *Options
	}
	tests := []test{
		{"unsharded", defaultTestOpts()},
		{"selectsharded", defaultTestOpts()},
		{"insertsharded", defaultTestOpts()},
		{"updatesharded", defaultTestOpts()},
		{"deletesharded", defaultTestOpts()},
		{"comments", defaultTestOpts()},
		{"options", &Options{
			ReplicationMode: "STATEMENT",
			NumShards:       4,
			Normalize:       false,
		}},
		{"target", &Options{
			ReplicationMode: "ROW",
			NumShards:       4,
			Normalize:       false,
			Target:          "ks_sharded/40-80",
		}},
	}

	for _, tst := range tests {
		testExplain(tst.name, tst.opts, t)
	}
}

func TestErrors(t *testing.T) {
	initTest(ModeMulti, defaultTestOpts(), &testopts{}, t)

	tests := []struct {
		SQL string
		Err string
	}{
		{
			SQL: "INVALID SQL",
			Err: "vtexplain execute error in 'INVALID SQL': syntax error at position 8 near 'INVALID'",
		},

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

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

		{
			SQL: "SELECT * FROM table_not_in_schema",
			Err: "unknown error: unable to resolve table name table_not_in_schema",
		},
	}

	for _, test := range tests {
		t.Run(test.SQL, func(t *testing.T) {
			_, err := Run(test.SQL)
			require.Error(t, err)
			require.Contains(t, err.Error(), test.Err)
		})
	}
}

func TestJSONOutput(t *testing.T) {
	sql := "select 1 from user where id = 1"
	explains, err := Run(sql)
	require.NoError(t, err, "vtexplain error")
	require.NotNil(t, explains, "vtexplain error running %s: no explain", string(sql))

	for _, e := range explains {
		for i, action := range e.TabletActions {
			var mysqlQueries []*MysqlQuery
			for _, query := range action.MysqlQueries {
				if !strings.Contains(strings.ToLower(query.SQL), "set collation_connection") {
					mysqlQueries = append(mysqlQueries, query)
				}
			}
			e.TabletActions[i].MysqlQueries = mysqlQueries
		}
	}
	explainJSON := ExplainsAsJSON(explains)

	var data interface{}
	err = json.Unmarshal([]byte(explainJSON), &data)
	require.NoError(t, err, "error unmarshaling json")

	array, ok := data.([]interface{})
	if !ok || len(array) != 1 {
		t.Errorf("expected single-element top-level array, got:\n%s", explainJSON)
	}

	explain, ok := array[0].(map[string]interface{})
	if !ok {
		t.Errorf("expected explain map, got:\n%s", explainJSON)
	}

	if explain["SQL"] != sql {
		t.Errorf("expected SQL, got:\n%s", explainJSON)
	}

	plans, ok := explain["Plans"].([]interface{})
	if !ok || len(plans) != 1 {
		t.Errorf("expected single-element plans array, got:\n%s", explainJSON)
	}

	actions, ok := explain["TabletActions"].(map[string]interface{})
	if !ok {
		t.Errorf("expected TabletActions map, got:\n%s", explainJSON)
	}

	actionsJSON, err := json.MarshalIndent(actions, "", "    ")
	require.NoError(t, err, "error in json marshal")
	wantJSON := `{
    "ks_sharded/-40": {
        "MysqlQueries": [
            {
                "SQL": "select 1 from ` + "`user`" + ` where id = 1 limit 10001",
                "Time": 1
            }
        ],
        "TabletQueries": [
            {
                "BindVars": {
                    "#maxLimit": "10001",
                    "vtg1": "1"
                },
                "SQL": "select :vtg1 from ` + "`user`" + ` where id = :vtg1",
                "Time": 1
            }
        ]
    }
}`
	diff := cmp.Diff(wantJSON, string(actionsJSON))
	if diff != "" {
		t.Errorf(diff)
	}
}

func testShardInfo(ks, start, end string, t *testing.T) *topo.ShardInfo {
	kr, err := key.ParseKeyRangeParts(start, end)
	require.NoError(t, err)

	return topo.NewShardInfo(
		ks,
		fmt.Sprintf("%s-%s", start, end),
		&topodata.Shard{KeyRange: kr},
		&vtexplainTestTopoVersion{},
	)
}

func TestUsingKeyspaceShardMap(t *testing.T) {
	tests := []struct {
		testcase      string
		ShardRangeMap map[string]map[string]*topo.ShardInfo
	}{
		{
			testcase: "select-sharded-8",
			ShardRangeMap: map[string]map[string]*topo.ShardInfo{
				"ks_sharded": {
					"-20":   testShardInfo("ks_sharded", "", "20", t),
					"20-40": testShardInfo("ks_sharded", "20", "40", t),
					"40-60": testShardInfo("ks_sharded", "40", "60", t),
					"60-80": testShardInfo("ks_sharded", "60", "80", t),
					"80-a0": testShardInfo("ks_sharded", "80", "a0", t),
					"a0-c0": testShardInfo("ks_sharded", "a0", "c0", t),
					"c0-e0": testShardInfo("ks_sharded", "c0", "e0", t),
					"e0-":   testShardInfo("ks_sharded", "e0", "", t),
				},
			},
		},
		{
			testcase: "uneven-keyspace",
			ShardRangeMap: map[string]map[string]*topo.ShardInfo{
				// Have mercy on the poor soul that has this keyspace sharding.
				// But, hey, vtexplain still works so they have that going for them.
				"ks_sharded": {
					"-80":   testShardInfo("ks_sharded", "", "80", t),
					"80-90": testShardInfo("ks_sharded", "80", "90", t),
					"90-a0": testShardInfo("ks_sharded", "90", "a0", t),
					"a0-e8": testShardInfo("ks_sharded", "a0", "e8", t),
					"e8-":   testShardInfo("ks_sharded", "e8", "", t),
				},
			},
		},
	}

	for _, test := range tests {
		runTestCase(test.testcase, ModeMulti, defaultTestOpts(), &testopts{test.ShardRangeMap}, t)
	}
}

type vtexplainTestTopoVersion struct{}

func (vtexplain *vtexplainTestTopoVersion) String() string { return "vtexplain-test-topo" }
<|endoftext|>


<commit_before>package broker

import (
	"fmt"
	"sync"

	"github.com/hashicorp/raft"
	"github.com/travisjeffery/jocko/broker/metadata"
)

type brokerLookup struct {
	lock            sync.RWMutex
	addressToBroker map[raft.ServerAddress]*metadata.Broker
	idToBroker      map[raft.ServerID]*metadata.Broker
}

func NewBrokerLookup() *brokerLookup {
	return &brokerLookup{
		addressToBroker: make(map[raft.ServerAddress]*metadata.Broker),
		idToBroker:      make(map[raft.ServerID]*metadata.Broker),
	}
}

func (sl *brokerLookup) AddBroker(broker *metadata.Broker) {
	sl.lock.Lock()
	defer sl.lock.Unlock()
	sl.addressToBroker[raft.ServerAddress(broker.RaftAddr)] = broker
	sl.idToBroker[raft.ServerID(broker.ID)] = broker
}

func (sl *brokerLookup) BrokerByAddr(addr raft.ServerAddress) *metadata.Broker {
	sl.lock.RLock()
	defer sl.lock.RUnlock()
	svr, _ := sl.addressToBroker[addr]
	return svr
}

func (sl *brokerLookup) BrokerByID(id raft.ServerID) *metadata.Broker {
	sl.lock.RLock()
	defer sl.lock.RUnlock()
	svr, _ := sl.idToBroker[id]
	return svr
}

func (sl *brokerLookup) BrokerAddr(id raft.ServerID) (raft.ServerAddress, error) {
	sl.lock.RLock()
	defer sl.lock.RUnlock()
	svr, ok := sl.idToBroker[id]
	if !ok {
		return "", fmt.Errorf("no broker for id %v", id)
	}
	return raft.ServerAddress(svr.RaftAddr), nil
}

func (sl *brokerLookup) RemoveBroker(broker *metadata.Broker) {
	sl.lock.Lock()
	defer sl.lock.Unlock()
	delete(sl.addressToBroker, raft.ServerAddress(broker.RaftAddr))
	delete(sl.idToBroker, raft.ServerID(broker.ID))
}

func (sl *brokerLookup) Brokers() []*metadata.Broker {
	sl.lock.RLock()
	sl.lock.RUnlock()
	var ret []*metadata.Broker
	for _, svr := range sl.addressToBroker {
		ret = append(ret, svr)
	}
	return ret
}
<commit_msg>broker: fix unlock call<commit_after>package broker

import (
	"fmt"
	"sync"

	"github.com/hashicorp/raft"
	"github.com/travisjeffery/jocko/broker/metadata"
)

type brokerLookup struct {
	lock            sync.RWMutex
	addressToBroker map[raft.ServerAddress]*metadata.Broker
	idToBroker      map[raft.ServerID]*metadata.Broker
}

func NewBrokerLookup() *brokerLookup {
	return &brokerLookup{
		addressToBroker: make(map[raft.ServerAddress]*metadata.Broker),
		idToBroker:      make(map[raft.ServerID]*metadata.Broker),
	}
}

func (sl *brokerLookup) AddBroker(broker *metadata.Broker) {
	sl.lock.Lock()
	defer sl.lock.Unlock()
	sl.addressToBroker[raft.ServerAddress(broker.RaftAddr)] = broker
	sl.idToBroker[raft.ServerID(broker.ID)] = broker
}

func (sl *brokerLookup) BrokerByAddr(addr raft.ServerAddress) *metadata.Broker {
	sl.lock.RLock()
	defer sl.lock.RUnlock()
	svr, _ := sl.addressToBroker[addr]
	return svr
}

func (sl *brokerLookup) BrokerByID(id raft.ServerID) *metadata.Broker {
	sl.lock.RLock()
	defer sl.lock.RUnlock()
	svr, _ := sl.idToBroker[id]
	return svr
}

func (sl *brokerLookup) BrokerAddr(id raft.ServerID) (raft.ServerAddress, error) {
	sl.lock.RLock()
	defer sl.lock.RUnlock()
	svr, ok := sl.idToBroker[id]
	if !ok {
		return "", fmt.Errorf("no broker for id %v", id)
	}
	return raft.ServerAddress(svr.RaftAddr), nil
}

func (sl *brokerLookup) RemoveBroker(broker *metadata.Broker) {
	sl.lock.Lock()
	defer sl.lock.Unlock()
	delete(sl.addressToBroker, raft.ServerAddress(broker.RaftAddr))
	delete(sl.idToBroker, raft.ServerID(broker.ID))
}

func (sl *brokerLookup) Brokers() []*metadata.Broker {
	sl.lock.RLock()
	defer sl.lock.RUnlock()
	var ret []*metadata.Broker
	for _, svr := range sl.addressToBroker {
		ret = append(ret, svr)
	}
	return ret
}
<|endoftext|>