<commit_before>package graph

import (
	"github.com/emirpasic/gods/lists/arraylist"

	"fmt"
	"net/http"
	"regexp"
	"strconv"
	"strings"

	h "github.com/Cepave/open-falcon-backend/modules/f2e-api/app/helper"
	"github.com/Cepave/open-falcon-backend/modules/f2e-api/app/helper/filter"
	m "github.com/Cepave/open-falcon-backend/modules/f2e-api/app/model/graph"

	"github.com/Cepave/open-falcon-backend/modules/f2e-api/app/model/boss"
	"github.com/gin-gonic/gin"
)

func HostsSearching(c *gin.Context) {
	AccpectTypes := arraylist.New()
	AccpectTypes.Add("platform", "idc", "isp", "province", "hostname", "hostgroup")
	q := c.DefaultQuery("q", "--")
	limitStr := c.DefaultQuery("limit", "100")
	limit := 100
	if ok, _ := regexp.MatchString(`\d+`, limitStr); ok {
		limit, _ = strconv.Atoi(limitStr)
	}
	ftype := c.DefaultQuery("filter_type", "all")
	if q == "--" {
		h.JSONR(c, badstatus, "q is empty, please check it!")
		return
	}
	if !(ftype == "all" || AccpectTypes.Contains(ftype)) {
		h.JSONR(c, badstatus, fmt.Sprintf("filter_type got error type, please check it!, only support: %v", AccpectTypes))
		return
	}
	bossList := boss.GetBossObjs()
	res := map[string]interface{}{}
	switch ftype {
	case "platform":
		res = map[string]interface{}{
			"platform": filter.PlatformFilter(bossList, q, limit),
		}
	case "idc":
		res = map[string]interface{}{
			"idc": filter.IdcFilter(bossList, q, limit),
		}
	case "isp":
		res = map[string]interface{}{
			"isp": filter.IspFilter(bossList, q, limit),
		}
	case "province":
		res = map[string]interface{}{
			"province": filter.ProvinceFilter(bossList, q, limit),
		}
	case "hostname":
		res = map[string]interface{}{
			"hostname": filter.HostNameFilter(bossList, q, limit),
		}
	case "hostgroup":
		res = map[string]interface{}{
			"hostgroup": filter.HostGroupFilter(q, limit),
		}
	case "all":
		res = map[string]interface{}{
			"platform":  filter.PlatformFilter(bossList, q, limit),
			"idc":       filter.IdcFilter(bossList, q, limit),
			"isp":       filter.IspFilter(bossList, q, limit),
			"province":  filter.ProvinceFilter(bossList, q, limit),
			"hostname":  filter.HostNameFilter(bossList, q, limit),
			"hostgroup": filter.HostGroupFilter(q, limit),
		}
	}
	h.JSONR(c, res)
	return
}

type APIEndpointsQuerySubMetricInputs struct {
	Endpoints   []string `json:"endpoints" form:"endpoints" binding:"required"`
	MetricQuery string   `json:"metric_query" form:"metric_query"`
	RawKeys     bool     `json:"raw_keys" form:"raw_keys"`
}

func EndpointsQuerySubMetric(c *gin.Context) {
	inputs := APIEndpointsQuerySubMetricInputs{}
	inputs.MetricQuery = ".+"
	inputs.RawKeys = true
	if err := c.Bind(&inputs); err != nil {
		h.JSONR(c, http.StatusBadRequest, err.Error())
		return
	}
	if inputs.MetricQuery == "" {
		inputs.MetricQuery = ".+"
	}
	enps := inputs.Endpoints
	enpids := []int64{}
	db.Graph.Table("endpoint").Select("id").Where("endpoint IN (?)", enps).Pluck("id", &enpids)
	metricSlist := map[string]int{}
	if len(enpids) != 0 {
		counters := []string{}
		db.Graph.Table("endpoint_counter").Where("endpoint_id IN (?) AND counter regexp ? ", enpids, inputs.MetricQuery).Pluck("DISTINCT(counter)", &counters)
		if inputs.RawKeys {
			h.JSONR(c, counters)
		} else {
			for _, ct := range counters {
				c2 := getFirstCutStringOfMetric(ct)
				if v, ok := metricSlist[c2]; ok {
					metricSlist[c2] = v + 1
				} else {
					metricSlist[c2] = 1
				}
			}
			h.JSONR(c, metricSlist)
		}
		return
	}
	return
}

type APIEndpointsGetMetricBySubStarInput struct {
	Endpoints      []string `json:"endpoints" form:"endpoints" binding:"required"`
	MetricQuery    string   `json:"metric_query" form:"metric_query" binding:"required"`
	StarStr        []string `json:"star_strings" form:"star_strings" binding:"required"`
	SelectedMetric []string `json:"selected_metric" form:"selected_metric"`
}

func EndpointsGetMetricBySubStar(c *gin.Context) {
	//the same with other one , so reuse it
	inputs := APIEndpointsGetMetricBySubStarInput{}
	if err := c.Bind(&inputs); err != nil {
		h.JSONR(c, http.StatusBadRequest, err.Error())
		return
	}
	enps := inputs.Endpoints
	enpids := []int{}
	db.Graph.Table("endpoint").Select("id").Where("endpoint IN (?)", enps).Pluck("id", &enpids)
	var regexFilterTmp []string
	for _, n := range inputs.StarStr {
		n2 := fmt.Sprintf("counter regexp '^%s\\.'", n)
		regexFilterTmp = append(regexFilterTmp, n2)
	}
	regexFilter := strings.Join(regexFilterTmp, " OR ")
	metricSlist := inputs.SelectedMetric
	if len(enpids) != 0 {
		var counters []m.EndpointCounter
		db.Graph.Table("endpoint_counter").
			Select("DISTINCT(counter) as counter").
			Where(fmt.Sprintf(
				"endpoint_id IN (?) AND (%s) AND counter regexp ? ", regexFilter), enpids, inputs.MetricQuery).
			Scan(&counters)

		for _, ct := range counters {
			metricSlist = append(metricSlist, ct.Counter)
		}
	}
	h.JSONR(c, metricSlist)
	return
}

func getFirstCutStringOfMetric(c string) string {
	patt := regexp.MustCompile("^([^\\.]+)\\.")
	matchPatt := patt.FindStringSubmatch(c)
	if len(matchPatt) == 0 {
		return c
	} else {
		return matchPatt[1]
	}
}
<commit_msg>1. change query host of boss list to graph's endpoint table 2. support hostgroup / host query<commit_after>package graph

import (
	"github.com/emirpasic/gods/lists/arraylist"

	"fmt"
	"net/http"
	"regexp"
	"strconv"
	"strings"

	h "github.com/Cepave/open-falcon-backend/modules/f2e-api/app/helper"
	"github.com/Cepave/open-falcon-backend/modules/f2e-api/app/helper/filter"
	m "github.com/Cepave/open-falcon-backend/modules/f2e-api/app/model/graph"

	"github.com/Cepave/open-falcon-backend/modules/f2e-api/app/model/boss"
	"github.com/gin-gonic/gin"
)

func HostsSearching(c *gin.Context) {
	AccpectTypes := arraylist.New()
	AccpectTypes.Add("platform", "idc", "isp", "province", "hostname", "hostgroup")
	q := c.DefaultQuery("q", "--")
	limitStr := c.DefaultQuery("limit", "100")
	limit := 100
	if ok, _ := regexp.MatchString(`\d+`, limitStr); ok {
		limit, _ = strconv.Atoi(limitStr)
	}
	ftype := c.DefaultQuery("filter_type", "all")
	if q == "--" {
		h.JSONR(c, badstatus, "q is empty, please check it!")
		return
	}
	if !(ftype == "all" || AccpectTypes.Contains(ftype)) {
		h.JSONR(c, badstatus, fmt.Sprintf("filter_type got error type, please check it!, only support: %v", AccpectTypes))
		return
	}
	bossList := []boss.BossHost{}
	if ftype != "hostname" && ftype != "hostgroup" {
		bossList = boss.GetBossObjs()
	}
	res := map[string]interface{}{}
	switch ftype {
	case "platform":
		res = map[string]interface{}{
			"platform": filter.PlatformFilter(bossList, q, limit),
		}
	case "idc":
		res = map[string]interface{}{
			"idc": filter.IdcFilter(bossList, q, limit),
		}
	case "isp":
		res = map[string]interface{}{
			"isp": filter.IspFilter(bossList, q, limit),
		}
	case "province":
		res = map[string]interface{}{
			"province": filter.ProvinceFilter(bossList, q, limit),
		}
	case "hostname":
		res = map[string]interface{}{
			"hostname": filter.HostFilter(q, limit),
		}
	case "hostgroup":
		res = map[string]interface{}{
			"hostgroup": filter.HostGroupFilter(q, limit),
		}
	case "all":
		res = map[string]interface{}{
			"platform":  filter.PlatformFilter(bossList, q, limit),
			"idc":       filter.IdcFilter(bossList, q, limit),
			"isp":       filter.IspFilter(bossList, q, limit),
			"province":  filter.ProvinceFilter(bossList, q, limit),
			"hostname":  filter.HostFilter(q, limit),
			"hostgroup": filter.HostGroupFilter(q, limit),
		}
	}
	h.JSONR(c, res)
	return
}

type APIEndpointsQuerySubMetricInputs struct {
	Endpoints   []string `json:"endpoints" form:"endpoints" binding:"required"`
	MetricQuery string   `json:"metric_query" form:"metric_query"`
	RawKeys     bool     `json:"raw_keys" form:"raw_keys"`
}

func EndpointsQuerySubMetric(c *gin.Context) {
	inputs := APIEndpointsQuerySubMetricInputs{
		MetricQuery: ".+",
		RawKeys:     true,
	}
	if err := c.Bind(&inputs); err != nil {
		h.JSONR(c, http.StatusBadRequest, err.Error())
		return
	}
	enps := inputs.Endpoints
	enpids := []int64{}
	db.Graph.Table("endpoint").Select("id").Where("endpoint IN (?)", enps).Pluck("id", &enpids)
	metricSlist := map[string]int{}
	if len(enpids) != 0 {
		counters := []string{}
		db.Graph.Table("endpoint_counter").Where("endpoint_id IN (?) AND counter regexp ? ", enpids, inputs.MetricQuery).Pluck("DISTINCT(counter)", &counters)
		if inputs.RawKeys {
			h.JSONR(c, counters)
		} else {
			for _, ct := range counters {
				c2 := getFirstCutStringOfMetric(ct)
				if v, ok := metricSlist[c2]; ok {
					metricSlist[c2] = v + 1
				} else {
					metricSlist[c2] = 1
				}
			}
			h.JSONR(c, metricSlist)
		}
	} else {
		h.JSONR(c, badstatus, "no metrice found")
	}
	return
}

type APIEndpointsGetMetricBySubStarInput struct {
	Endpoints      []string `json:"endpoints" form:"endpoints" binding:"required"`
	MetricQuery    string   `json:"metric_query" form:"metric_query" binding:"required"`
	StarStr        []string `json:"star_strings" form:"star_strings" binding:"required"`
	SelectedMetric []string `json:"selected_metric" form:"selected_metric"`
}

func EndpointsGetMetricBySubStar(c *gin.Context) {
	//the same with other one , so reuse it
	inputs := APIEndpointsGetMetricBySubStarInput{}
	if err := c.Bind(&inputs); err != nil {
		h.JSONR(c, http.StatusBadRequest, err.Error())
		return
	}
	enps := inputs.Endpoints
	enpids := []int{}
	db.Graph.Table("endpoint").Select("id").Where("endpoint IN (?)", enps).Pluck("id", &enpids)
	var regexFilterTmp []string
	for _, n := range inputs.StarStr {
		n2 := fmt.Sprintf("counter regexp '^%s\\.'", n)
		regexFilterTmp = append(regexFilterTmp, n2)
	}
	regexFilter := strings.Join(regexFilterTmp, " OR ")
	metricSlist := inputs.SelectedMetric
	if len(enpids) != 0 {
		var counters []m.EndpointCounter
		db.Graph.Table("endpoint_counter").
			Select("DISTINCT(counter) as counter").
			Where(fmt.Sprintf(
				"endpoint_id IN (?) AND (%s) AND counter regexp ? ", regexFilter), enpids, inputs.MetricQuery).
			Scan(&counters)

		for _, ct := range counters {
			metricSlist = append(metricSlist, ct.Counter)
		}
	}
	h.JSONR(c, metricSlist)
	return
}

func getFirstCutStringOfMetric(c string) string {
	patt := regexp.MustCompile("^([^\\.]+)\\.")
	matchPatt := patt.FindStringSubmatch(c)
	if len(matchPatt) == 0 {
		return c
	} else {
		return matchPatt[1]
	}
}
<|endoftext|>


<commit_before>package main

import (
	//"appengine"
	"appengine/datastore"
	"github.com/GoogleCloudPlatform/go-endpoints/endpoints"
	//"log"
	//"net/url"
	"time"
)

// API for Match
type APIMatch struct{}

// Type of Match
type Match struct {
	UID          *datastore.Key `json:"uid" datastore:"-"`
	UIDCountryA  *datastore.Key `json:"uidCountryA" `
	UIDCountryB  *datastore.Key `json:"uidCountryB" `
	UIDWinner    *datastore.Key `json:"uidWinner" `
	UIDUser      *datastore.Key `json:"uidUser" `
	Round        int            `json:"round"`
	DateCreation time.Time      `json:"dateCreation"`
}

// List of Match
type Matchs struct {
	Matchs []Match `json:"matchs"`
}

// Type that is use for creating a Match
type MatchToCreate struct {
	UIDCountryA *datastore.Key
	UIDCountryB *datastore.Key
	UIDUser     *datastore.Key
	UIDWinner   *datastore.Key
	Round       int
}

// Create allow you to create a match
// waiting for a context and all data that are editable by the client side
// give back the current created Match or an error
func (APIMatch) Create(c endpoints.Context, r *MatchToCreate) (*Match, error) {

	k := datastore.NewIncompleteKey(c, "Match", nil)

	m := &Match{
		UIDCountryA:  r.UIDCountryA,
		UIDCountryB:  r.UIDCountryB,
		UIDUser:      r.UIDUser,
		UIDWinner:    r.UIDWinner,
		Round:        r.Round,
		DateCreation: time.Now(),
	}

	k, err := datastore.Put(c, k, m)

	if err != nil {
		return nil, err
	}

	m.UID = k

	return m, nil
}

// type for upgrading a Match
type MatchToUpdate struct {
	UID         *datastore.Key
	UIDCountryA *datastore.Key
	UIDCountryB *datastore.Key
	UIDWinner   *datastore.Key
	UIDUser     *datastore.Key
	Round       int
}

// Update allow you to update a Match
// waiting for a context and all data that are editable for the client side
// give back the updated Match or an error
func (APIMatch) Update(c endpoints.Context, r *MatchToUpdate) (*Match, error) {

	var m Match

	if err := datastore.Get(c, r.UID, &m); err == datastore.ErrNoSuchEntity {
		return nil, err
	} else if err != nil {
		return nil, err
	}

	m.UIDCountryA = r.UIDCountryA
	m.UIDCountryB = r.UIDCountryB
	m.UIDWinner = r.UIDWinner
	m.UIDUser = r.UIDUser
	m.Round = r.Round

	_, err := datastore.Put(c, r.UID, &m)

	if err != nil {
		return nil, err
	}

	return &m, nil

}

// type for getting a Match by UID
type MatchToGet struct {
	UID *datastore.Key
}

// Get give all data from the Match
// waiting for a context and a key
// give back a Match or an error
func (APIMatch) Get(c endpoints.Context, r *MatchToGet) (*Match, error) {

	var m Match

	if err := datastore.Get(c, r.UID, &m); err == datastore.ErrNoSuchEntity {
		return nil, endpoints.NewNotFoundError("User not found")
	} else if err != nil {
		return nil, err
	}

	m.UID = r.UID

	return &m, nil

}

// List give all Matchs that are stored into the datastore
// waiting for a context
// give back a list of Matchs or an error
func (APIMatch) List(c endpoints.Context) (*Matchs, error) {

	matchs := []Match{}

	keys, err := datastore.NewQuery("Match").GetAll(c, &matchs)

	if err != nil {
		return nil, err
	}
	for i, k := range keys {
		matchs[i].UID = k
	}

	return &Matchs{matchs}, nil
}

// type for getting a Match by UID
type ListByRoundRequest struct {
	Round int
}

// List give all Matchs that are stored into the datastore
// waiting for a context
// give back a list of Matchs or an error
func (APIMatch) ListByRound(c endpoints.Context, r *ListByRoundRequest) (*Matchs, error) {

	matchs := []Match{}

	q := datastore.NewQuery("Match").Filter("Round =", r.Round)
	keys, err := q.GetAll(c, &matchs)

	if err != nil {
		return nil, err
	}

	for i, k := range keys {
		matchs[i].UID = k
	}

	return &Matchs{matchs}, nil
}

// type for deleting a Match by UID
type MatchToDelete struct {
	UID *datastore.Key
}

// Delete allow you to delete a User
// waiging for a context and a user key
// give back an error if something went wrong
func (APIMatch) Delete(c endpoints.Context, r *MatchToDelete) error {

	err := datastore.Delete(c, r.UID)

	if err != nil {
		return nil
	}

	return err

}

// Delete allow you to delete a Match
// waiging for a context and a Match key
// give back an error if something went wrong

/*
func (APIMatch) RandomCreate(c endpoints) Matchs {

}
*/
<commit_msg>modify comments on matchs<commit_after>package main

import (
	//"appengine"
	"appengine/datastore"
	"github.com/GoogleCloudPlatform/go-endpoints/endpoints"
	//"log"
	//"net/url"
	"time"
)

// API for Match
type APIMatch struct{}

// Type of Match
type Match struct {
	UID          *datastore.Key `json:"uid" datastore:"-"`
	UIDCountryA  *datastore.Key `json:"uidCountryA" `
	UIDCountryB  *datastore.Key `json:"uidCountryB" `
	UIDWinner    *datastore.Key `json:"uidWinner" `
	UIDUser      *datastore.Key `json:"uidUser" `
	Round        int            `json:"round"`
	DateCreation time.Time      `json:"dateCreation"`
}

// List of Match
type Matchs struct {
	Matchs []Match `json:"matchs"`
}

// Type that is use for creating a Match
type MatchToCreate struct {
	UIDCountryA *datastore.Key
	UIDCountryB *datastore.Key
	UIDUser     *datastore.Key
	UIDWinner   *datastore.Key
	Round       int
}

// Create allow you to create a match
// waiting for a context and all data that are editable by the client side
// give back the current created Match or an error
func (APIMatch) Create(c endpoints.Context, r *MatchToCreate) (*Match, error) {

	k := datastore.NewIncompleteKey(c, "Match", nil)

	m := &Match{
		UIDCountryA:  r.UIDCountryA,
		UIDCountryB:  r.UIDCountryB,
		UIDUser:      r.UIDUser,
		UIDWinner:    r.UIDWinner,
		Round:        r.Round,
		DateCreation: time.Now(),
	}

	k, err := datastore.Put(c, k, m)

	if err != nil {
		return nil, err
	}

	m.UID = k

	return m, nil
}

// type for upgrading a Match
type MatchToUpdate struct {
	UID         *datastore.Key
	UIDCountryA *datastore.Key
	UIDCountryB *datastore.Key
	UIDWinner   *datastore.Key
	UIDUser     *datastore.Key
	Round       int
}

// Update allow you to update a Match
// waiting for a context and all data that are editable for the client side
// give back the updated Match or an error
func (APIMatch) Update(c endpoints.Context, r *MatchToUpdate) (*Match, error) {

	var m Match

	if err := datastore.Get(c, r.UID, &m); err == datastore.ErrNoSuchEntity {
		return nil, err
	} else if err != nil {
		return nil, err
	}

	m.UIDCountryA = r.UIDCountryA
	m.UIDCountryB = r.UIDCountryB
	m.UIDWinner = r.UIDWinner
	m.UIDUser = r.UIDUser
	m.Round = r.Round

	_, err := datastore.Put(c, r.UID, &m)

	if err != nil {
		return nil, err
	}

	return &m, nil

}

// type for getting a Match by UID
type MatchToGet struct {
	UID *datastore.Key
}

// Get give all data from the Match
// waiting for a context and a key
// give back a Match or an error
func (APIMatch) Get(c endpoints.Context, r *MatchToGet) (*Match, error) {

	var m Match

	if err := datastore.Get(c, r.UID, &m); err == datastore.ErrNoSuchEntity {
		return nil, endpoints.NewNotFoundError("User not found")
	} else if err != nil {
		return nil, err
	}

	m.UID = r.UID

	return &m, nil

}

// List give all Matchs that are stored into the datastore
// waiting for a context
// give back a list of Matchs or an error
func (APIMatch) List(c endpoints.Context) (*Matchs, error) {

	matchs := []Match{}

	keys, err := datastore.NewQuery("Match").GetAll(c, &matchs)

	if err != nil {
		return nil, err
	}
	for i, k := range keys {
		matchs[i].UID = k
	}

	return &Matchs{matchs}, nil
}

// type for getting a Match by UID
type ListByRoundRequest struct {
	Round int
}

// List give all Matchs that are stored into the datastore
// waiting for a context
// give back a list of Matchs or an error
func (APIMatch) ListByRound(c endpoints.Context, r *ListByRoundRequest) (*Matchs, error) {

	matchs := []Match{}

	q := datastore.NewQuery("Match").Filter("Round =", r.Round)
	keys, err := q.GetAll(c, &matchs)

	if err != nil {
		return nil, err
	}

	for i, k := range keys {
		matchs[i].UID = k
	}

	return &Matchs{matchs}, nil
}

// type for deleting a Match by UID
type MatchToDelete struct {
	UID *datastore.Key
}

// Delete allow you to delete a Match
// waiging for a context and a match key
// give back an error if something went wrong
func (APIMatch) Delete(c endpoints.Context, r *MatchToDelete) error {

	err := datastore.Delete(c, r.UID)

	if err != nil {
		return nil
	}

	return err
}

// Delete allow you to delete a Match
// waiging for a context and a Match key
// give back an error if something went wrong

/*
func (APIMatch) RandomCreate(c endpoints) Matchs {

}
*/
<|endoftext|>


<commit_before>// Copyright 2013 Matthew Baird
// 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 elastigo

import (
	"encoding/json"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"net/url"
	"strings"
	"testing"
)

var (
	mux    *http.ServeMux
	server *httptest.Server
)

func setup(t *testing.T) *Conn {
	mux = http.NewServeMux()
	server = httptest.NewServer(mux)
	c := NewConn()

	serverURL, err := url.Parse(server.URL)
	if err != nil {
		t.Fatalf("Error: %v", err)
	}

	c.Domain = strings.Split(serverURL.Host, ":")[0]
	c.Port = strings.Split(serverURL.Host, ":")[1]

	return c
}

func teardown() {
	server.Close()
}

type TestStruct struct {
	Id            string `json:"id" elastic:"index:not_analyzed"`
	DontIndex     string `json:"dontIndex" elastic:"index:no"`
	Number        int    `json:"number" elastic:"type:integer,index:analyzed"`
	Omitted       string `json:"-"`
	NoJson        string `elastic:"type:string"`
	unexported    string
	JsonOmitEmpty string `json:"jsonOmitEmpty,omitempty" elastic:"type:string"`
	Embedded
	Nested       NestedStruct   `json:"nested"`
	NestedP      *NestedStruct  `json:"pointer_to_nested"`
	NestedS      []NestedStruct `json:"slice_of_nested"`
	MultiAnalyze string         `json:"multi_analyze"`
}

type Embedded struct {
	EmbeddedField string `json:"embeddedField" elastic:"type:string"`
}

type NestedStruct struct {
	NestedField string `json:"nestedField" elastic:"type:date"`
}

func TestPutMapping(t *testing.T) {
	c := setup(t)
	defer teardown()

	options := MappingOptions{
		Timestamp: TimestampOptions{Enabled: true},
		Id:        IdOptions{Index: "analyzed", Path: "id"},
		Properties: map[string]interface{}{
			// special properties that can't be expressed as tags
			"multi_analyze": map[string]interface{}{
				"type": "multi_field",
				"fields": map[string]map[string]string{
					"ma_analyzed":   {"type": "string", "index": "analyzed"},
					"ma_unanalyzed": {"type": "string", "index": "un_analyzed"},
				},
			},
		},
	}
	expValue := MappingForType("myType", MappingOptions{
		Timestamp: TimestampOptions{Enabled: true},
		Id:        IdOptions{Index: "analyzed", Path: "id"},
		Properties: map[string]interface{}{
			"NoJson":        map[string]string{"type": "string"},
			"dontIndex":     map[string]string{"index": "no"},
			"embeddedField": map[string]string{"type": "string"},
			"id":            map[string]string{"index": "not_analyzed"},
			"jsonOmitEmpty": map[string]string{"type": "string"},
			"number":        map[string]string{"index": "analyzed", "type": "integer"},
			"multi_analyze": map[string]interface{}{
				"type": "multi_field",
				"fields": map[string]map[string]string{
					"ma_analyzed":   {"type": "string", "index": "analyzed"},
					"ma_unanalyzed": {"type": "string", "index": "un_analyzed"},
				},
			},
			"nested": map[string]map[string]map[string]string{
				"properties": {
					"nestedField": {"type": "date"},
				},
			},
			"pointer_to_nested": map[string]map[string]map[string]string{
				"properties": {
					"nestedField": {"type": "date"},
				},
			},
			"slice_of_nested": map[string]map[string]map[string]string{
				"properties": {
					"nestedField": {"type": "date"},
				},
			},
		},
	})

	mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		var value map[string]interface{}
		bd, err := ioutil.ReadAll(r.Body)
		json.NewDecoder(strings.NewReader(string(bd))).Decode(&value)
		expValJson, err := json.MarshalIndent(expValue, "", "  ")
		if err != nil {
			t.Errorf("Got error: %v", err)
		}
		valJson, err := json.MarshalIndent(value, "", "  ")
		if err != nil {
			t.Errorf("Got error: %v", err)
		}

		if string(expValJson) != string(valJson) {
			t.Errorf("Expected %s but got %s", string(expValJson), string(valJson))
		}
	})

	err := c.PutMapping("myIndex", "myType", TestStruct{}, options)
	if err != nil {
		t.Errorf("Error: %v", err)
	}
}
<commit_msg>Fix misleading test case<commit_after>// Copyright 2013 Matthew Baird
// 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 elastigo

import (
	"encoding/json"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"net/url"
	"strings"
	"testing"
)

var (
	mux    *http.ServeMux
	server *httptest.Server
)

func setup(t *testing.T) *Conn {
	mux = http.NewServeMux()
	server = httptest.NewServer(mux)
	c := NewConn()

	serverURL, err := url.Parse(server.URL)
	if err != nil {
		t.Fatalf("Error: %v", err)
	}

	c.Domain = strings.Split(serverURL.Host, ":")[0]
	c.Port = strings.Split(serverURL.Host, ":")[1]

	return c
}

func teardown() {
	server.Close()
}

type TestStruct struct {
	Id            string `json:"id" elastic:"index:not_analyzed"`
	DontIndex     string `json:"dontIndex" elastic:"index:no"`
	Number        int    `json:"number" elastic:"type:integer,index:analyzed"`
	Omitted       string `json:"-"`
	NoJson        string `elastic:"type:string"`
	unexported    string
	JsonOmitEmpty string `json:"jsonOmitEmpty,omitempty" elastic:"type:string"`
	Embedded
	Nested       NestedStruct   `json:"nested"`
	NestedP      *NestedStruct  `json:"pointer_to_nested"`
	NestedS      []NestedStruct `json:"slice_of_nested"`
	MultiAnalyze string         `json:"multi_analyze"`
}

type Embedded struct {
	EmbeddedField string `json:"embeddedField" elastic:"type:string"`
}

type NestedStruct struct {
	NestedField string `json:"nestedField" elastic:"type:date"`
}

func TestPutMapping(t *testing.T) {
	c := setup(t)
	defer teardown()

	options := MappingOptions{
		Timestamp: TimestampOptions{Enabled: true},
		Id:        IdOptions{Index: "analyzed", Path: "id"},
		Properties: map[string]interface{}{
			// special properties that can't be expressed as tags
			"multi_analyze": map[string]interface{}{
				"type": "multi_field",
				"fields": map[string]map[string]string{
					"ma_analyzed":    {"type": "string", "index": "analyzed"},
					"ma_notanalyzed": {"type": "string", "index": "not_analyzed"},
				},
			},
		},
	}
	expValue := MappingForType("myType", MappingOptions{
		Timestamp: TimestampOptions{Enabled: true},
		Id:        IdOptions{Index: "analyzed", Path: "id"},
		Properties: map[string]interface{}{
			"NoJson":        map[string]string{"type": "string"},
			"dontIndex":     map[string]string{"index": "no"},
			"embeddedField": map[string]string{"type": "string"},
			"id":            map[string]string{"index": "not_analyzed"},
			"jsonOmitEmpty": map[string]string{"type": "string"},
			"number":        map[string]string{"index": "analyzed", "type": "integer"},
			"multi_analyze": map[string]interface{}{
				"type": "multi_field",
				"fields": map[string]map[string]string{
					"ma_analyzed":    {"type": "string", "index": "analyzed"},
					"ma_notanalyzed": {"type": "string", "index": "not_analyzed"},
				},
			},
			"nested": map[string]map[string]map[string]string{
				"properties": {
					"nestedField": {"type": "date"},
				},
			},
			"pointer_to_nested": map[string]map[string]map[string]string{
				"properties": {
					"nestedField": {"type": "date"},
				},
			},
			"slice_of_nested": map[string]map[string]map[string]string{
				"properties": {
					"nestedField": {"type": "date"},
				},
			},
		},
	})

	mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		var value map[string]interface{}
		bd, err := ioutil.ReadAll(r.Body)
		json.NewDecoder(strings.NewReader(string(bd))).Decode(&value)
		expValJson, err := json.MarshalIndent(expValue, "", "  ")
		if err != nil {
			t.Errorf("Got error: %v", err)
		}
		valJson, err := json.MarshalIndent(value, "", "  ")
		if err != nil {
			t.Errorf("Got error: %v", err)
		}

		if string(expValJson) != string(valJson) {
			t.Errorf("Expected %s but got %s", string(expValJson), string(valJson))
		}
	})

	err := c.PutMapping("myIndex", "myType", TestStruct{}, options)
	if err != nil {
		t.Errorf("Error: %v", err)
	}
}
<|endoftext|>


<commit_before>// Package oura implements the OAuth protocol for authenticating users through Oura API (for OuraRing).
package oura

import (
	"encoding/json"
	"io"
	"net/http"

	"fmt"

	"github.com/markbates/goth"
	"golang.org/x/oauth2"
)

const (
	authURL         string = "https://cloud.ouraring.com/oauth/authorize"
	tokenURL        string = "https://api.ouraring.com/oauth/token"
	endpointProfile string = "https://api.ouraring.com/v1/userinfo"
)

const (
	// ScopeEmail includes email address of the user
	ScopeEmail = "email"
	// ScopePersonal includes personal information (gender, age, height, weight)
	ScopePersonal = "personal"
	// ScopeDaily includes daily summaries of sleep, activity and readiness
	ScopeDaily = "daily"
)

// New creates a new Oura provider (for OuraRing), and sets up important connection details.
// You should always call `oura.New` to get a new Provider. Never try to create
// one manually.
func New(clientKey, secret, callbackURL string, scopes ...string) *Provider {
	p := &Provider{
		ClientKey:    clientKey,
		Secret:       secret,
		CallbackURL:  callbackURL,
		providerName: "oura",
	}
	p.config = newConfig(p, scopes)
	return p
}

// Provider is the implementation of `goth.Provider` for accessing Oura API.
type Provider struct {
	ClientKey    string
	Secret       string
	CallbackURL  string
	HTTPClient   *http.Client
	config       *oauth2.Config
	providerName string
}

// Name is the name used to retrieve this provider later.
func (p *Provider) Name() string {
	return p.providerName
}

// SetName is to update the name of the provider (needed in case of multiple providers of 1 type)
func (p *Provider) SetName(name string) {
	p.providerName = name
}

func (p *Provider) Client() *http.Client {
	return goth.HTTPClientWithFallBack(p.HTTPClient)
}

// Debug is a no-op for the oura package.
func (p *Provider) Debug(debug bool) {}

// BeginAuth asks Oura for an authentication end-point.
func (p *Provider) BeginAuth(state string) (goth.Session, error) {
	url := p.config.AuthCodeURL(state)
	session := &Session{
		AuthURL: url,
	}
	return session, nil
}

// FetchUser will go to Oura and access basic information about the user.
func (p *Provider) FetchUser(session goth.Session) (goth.User, error) {
	s := session.(*Session)
	user := goth.User{
		AccessToken:  s.AccessToken,
		Provider:     p.Name(),
		RefreshToken: s.RefreshToken,
		ExpiresAt:    s.ExpiresAt,
		UserID:       s.UserID,
	}

	if user.AccessToken == "" {
		// data is not yet retrieved since accessToken is still empty
		return user, fmt.Errorf("%s cannot get user information without accessToken", p.providerName)
	}

	req, err := http.NewRequest("GET", endpointProfile, nil)
	if err != nil {
		return user, err
	}
	req.Header.Set("Authorization", "Bearer "+s.AccessToken)
	resp, err := p.Client().Do(req)
	if err != nil {
		if resp != nil {
			resp.Body.Close()
		}
		return user, err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return user, fmt.Errorf("%s responded with a %d trying to fetch user information", p.providerName, resp.StatusCode)
	}

	//err = userFromReader(io.TeeReader(resp.Body, os.Stdout), &user)
	err = userFromReader(resp.Body, &user)
	return user, err
}

func userFromReader(reader io.Reader, user *goth.User) error {
	u := struct {
		Age    int     `json:"age"`
		Weight float32 `json:"weight"` // kg
		Height int     `json:"height"` // cm
		Gender string  `json:"gender"`
		Email  string  `json:"email"`
		UserID string  `json:"user_id"`
	}{}

	err := json.NewDecoder(reader).Decode(&u)
	if err != nil {
		return err
	}

	user.UserID = u.UserID
	user.Email = u.Email
	user.RawData = make(map[string]interface{})
	user.RawData["age"] = u.Age
	user.RawData["weight"] = u.Weight
	user.RawData["height"] = u.Height
	user.RawData["gender"] = u.Gender

	return err
}

func newConfig(provider *Provider, scopes []string) *oauth2.Config {
	c := &oauth2.Config{
		ClientID:     provider.ClientKey,
		ClientSecret: provider.Secret,
		RedirectURL:  provider.CallbackURL,
		Endpoint: oauth2.Endpoint{
			AuthURL:  authURL,
			TokenURL: tokenURL,
		},
		Scopes: []string{},
	}

	for _, scope := range scopes {
		c.Scopes = append(c.Scopes, scope)
	}

	return c
}

//RefreshToken get new access token based on the refresh token
func (p *Provider) RefreshToken(refreshToken string) (*oauth2.Token, error) {
	token := &oauth2.Token{RefreshToken: refreshToken}
	ts := p.config.TokenSource(oauth2.NoContext, token)
	newToken, err := ts.Token()
	if err != nil {
		return nil, err
	}
	return newToken, err
}

//RefreshTokenAvailable refresh token is not provided by oura
func (p *Provider) RefreshTokenAvailable() bool {
	return true
}
<commit_msg>Update FetchUser to only fill populated fields in RawData<commit_after>// Package oura implements the OAuth protocol for authenticating users through Oura API (for OuraRing).
package oura

import (
	"encoding/json"
	"io"
	"net/http"

	"fmt"

	"github.com/markbates/goth"
	"golang.org/x/oauth2"
)

const (
	authURL         string = "https://cloud.ouraring.com/oauth/authorize"
	tokenURL        string = "https://api.ouraring.com/oauth/token"
	endpointProfile string = "https://api.ouraring.com/v1/userinfo"
)

const (
	// ScopeEmail includes email address of the user
	ScopeEmail = "email"
	// ScopePersonal includes personal information (gender, age, height, weight)
	ScopePersonal = "personal"
	// ScopeDaily includes daily summaries of sleep, activity and readiness
	ScopeDaily = "daily"
)

// New creates a new Oura provider (for OuraRing), and sets up important connection details.
// You should always call `oura.New` to get a new Provider. Never try to create
// one manually.
func New(clientKey, secret, callbackURL string, scopes ...string) *Provider {
	p := &Provider{
		ClientKey:    clientKey,
		Secret:       secret,
		CallbackURL:  callbackURL,
		providerName: "oura",
	}
	p.config = newConfig(p, scopes)
	return p
}

// Provider is the implementation of `goth.Provider` for accessing Oura API.
type Provider struct {
	ClientKey    string
	Secret       string
	CallbackURL  string
	HTTPClient   *http.Client
	config       *oauth2.Config
	providerName string
}

// Name is the name used to retrieve this provider later.
func (p *Provider) Name() string {
	return p.providerName
}

// SetName is to update the name of the provider (needed in case of multiple providers of 1 type)
func (p *Provider) SetName(name string) {
	p.providerName = name
}

func (p *Provider) Client() *http.Client {
	return goth.HTTPClientWithFallBack(p.HTTPClient)
}

// Debug is a no-op for the oura package.
func (p *Provider) Debug(debug bool) {}

// BeginAuth asks Oura for an authentication end-point.
func (p *Provider) BeginAuth(state string) (goth.Session, error) {
	url := p.config.AuthCodeURL(state)
	session := &Session{
		AuthURL: url,
	}
	return session, nil
}

// FetchUser will go to Oura and access basic information about the user.
func (p *Provider) FetchUser(session goth.Session) (goth.User, error) {
	s := session.(*Session)
	user := goth.User{
		AccessToken:  s.AccessToken,
		Provider:     p.Name(),
		RefreshToken: s.RefreshToken,
		ExpiresAt:    s.ExpiresAt,
		UserID:       s.UserID,
	}

	if user.AccessToken == "" {
		// data is not yet retrieved since accessToken is still empty
		return user, fmt.Errorf("%s cannot get user information without accessToken", p.providerName)
	}

	req, err := http.NewRequest("GET", endpointProfile, nil)
	if err != nil {
		return user, err
	}
	req.Header.Set("Authorization", "Bearer "+s.AccessToken)
	resp, err := p.Client().Do(req)
	if err != nil {
		if resp != nil {
			resp.Body.Close()
		}
		return user, err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return user, fmt.Errorf("%s responded with a %d trying to fetch user information", p.providerName, resp.StatusCode)
	}

	//err = userFromReader(io.TeeReader(resp.Body, os.Stdout), &user)
	err = userFromReader(resp.Body, &user)
	return user, err
}

func userFromReader(reader io.Reader, user *goth.User) error {
	u := struct {
		Age    int     `json:"age"`
		Weight float32 `json:"weight"` // kg
		Height int     `json:"height"` // cm
		Gender string  `json:"gender"`
		Email  string  `json:"email"`
		UserID string  `json:"user_id"`
	}{}

	err := json.NewDecoder(reader).Decode(&u)
	if err != nil {
		return err
	}

	rawData := make(map[string]interface{})

	if u.Age != 0 {
		rawData["age"] = u.Age
	}
	if u.Weight != 0 {
		rawData["weight"] = u.Weight
	}
	if u.Height != 0 {
		rawData["height"] = u.Height
	}
	if u.Gender != "" {
		rawData["gender"] = u.Gender
	}

	user.UserID = u.UserID
	user.Email = u.Email
	if len(rawData) > 0 {
		user.RawData = rawData
	}

	return err
}

func newConfig(provider *Provider, scopes []string) *oauth2.Config {
	c := &oauth2.Config{
		ClientID:     provider.ClientKey,
		ClientSecret: provider.Secret,
		RedirectURL:  provider.CallbackURL,
		Endpoint: oauth2.Endpoint{
			AuthURL:  authURL,
			TokenURL: tokenURL,
		},
		Scopes: []string{},
	}

	for _, scope := range scopes {
		c.Scopes = append(c.Scopes, scope)
	}

	return c
}

//RefreshToken get new access token based on the refresh token
func (p *Provider) RefreshToken(refreshToken string) (*oauth2.Token, error) {
	token := &oauth2.Token{RefreshToken: refreshToken}
	ts := p.config.TokenSource(oauth2.NoContext, token)
	newToken, err := ts.Token()
	if err != nil {
		return nil, err
	}
	return newToken, err
}

//RefreshTokenAvailable refresh token is not provided by oura
func (p *Provider) RefreshTokenAvailable() bool {
	return true
}
<|endoftext|>


<commit_before>package providers

import (
	"context"
	"encoding/binary"
	"fmt"
	"strings"
	"time"

	lru "github.com/hashicorp/golang-lru"
	cid "github.com/ipfs/go-cid"
	ds "github.com/ipfs/go-datastore"
	autobatch "github.com/ipfs/go-datastore/autobatch"
	dsq "github.com/ipfs/go-datastore/query"
	logging "github.com/ipfs/go-log"
	goprocess "github.com/jbenet/goprocess"
	goprocessctx "github.com/jbenet/goprocess/context"
	peer "github.com/libp2p/go-libp2p-peer"
	base32 "github.com/whyrusleeping/base32"
)

var batchBufferSize = 256

var log = logging.Logger("providers")

var lruCacheSize = 256
var ProvideValidity = time.Hour * 24
var defaultCleanupInterval = time.Hour

type ProviderManager struct {
	// all non channel fields are meant to be accessed only within
	// the run method
	providers *lru.Cache
	dstore    ds.Datastore

	newprovs chan *addProv
	getprovs chan *getProv
	proc     goprocess.Process

	cleanupInterval time.Duration
}

type providerSet struct {
	providers []peer.ID
	set       map[peer.ID]time.Time
}

type addProv struct {
	k   cid.Cid
	val peer.ID
}

type getProv struct {
	k    cid.Cid
	resp chan []peer.ID
}

func NewProviderManager(ctx context.Context, local peer.ID, dstore ds.Batching) *ProviderManager {
	pm := new(ProviderManager)
	pm.getprovs = make(chan *getProv)
	pm.newprovs = make(chan *addProv)
	pm.dstore = autobatch.NewAutoBatching(dstore, batchBufferSize)
	cache, err := lru.New(lruCacheSize)
	if err != nil {
		panic(err) //only happens if negative value is passed to lru constructor
	}
	pm.providers = cache

	pm.proc = goprocessctx.WithContext(ctx)
	pm.cleanupInterval = defaultCleanupInterval
	pm.proc.Go(func(p goprocess.Process) { pm.run() })

	return pm
}

const providersKeyPrefix = "/providers/"

func mkProvKey(k cid.Cid) string {
	return providersKeyPrefix + base32.RawStdEncoding.EncodeToString(k.Bytes())
}

func (pm *ProviderManager) Process() goprocess.Process {
	return pm.proc
}

func (pm *ProviderManager) providersForKey(k cid.Cid) ([]peer.ID, error) {
	pset, err := pm.getProvSet(k)
	if err != nil {
		return nil, err
	}
	return pset.providers, nil
}

func (pm *ProviderManager) getProvSet(k cid.Cid) (*providerSet, error) {
	cached, ok := pm.providers.Get(k)
	if ok {
		return cached.(*providerSet), nil
	}

	pset, err := loadProvSet(pm.dstore, k)
	if err != nil {
		return nil, err
	}

	if len(pset.providers) > 0 {
		pm.providers.Add(k, pset)
	}

	return pset, nil
}

func loadProvSet(dstore ds.Datastore, k cid.Cid) (*providerSet, error) {
	res, err := dstore.Query(dsq.Query{Prefix: mkProvKey(k)})
	if err != nil {
		return nil, err
	}
	defer res.Close()

	now := time.Now()
	out := newProviderSet()
	for {
		e, ok := res.NextSync()
		if !ok {
			break
		}
		if e.Error != nil {
			log.Error("got an error: ", e.Error)
			continue
		}

		// check expiration time
		t, err := readTimeValue(e.Value)
		switch {
		case err != nil:
			// couldn't parse the time
			log.Warning("parsing providers record from disk: ", err)
			fallthrough
		case now.Sub(t) > ProvideValidity:
			// or just expired
			err = dstore.Delete(ds.RawKey(e.Key))
			if err != nil && err != ds.ErrNotFound {
				log.Warning("failed to remove provider record from disk: ", err)
			}
			continue
		}

		lix := strings.LastIndex(e.Key, "/")

		decstr, err := base32.RawStdEncoding.DecodeString(e.Key[lix+1:])
		if err != nil {
			log.Error("base32 decoding error: ", err)
			err = dstore.Delete(ds.RawKey(e.Key))
			if err != nil && err != ds.ErrNotFound {
				log.Warning("failed to remove provider record from disk: ", err)
			}
			continue
		}

		pid := peer.ID(decstr)

		out.setVal(pid, t)
	}

	return out, nil
}

func readTimeValue(data []byte) (time.Time, error) {
	nsec, n := binary.Varint(data)
	if n <= 0 {
		return time.Time{}, fmt.Errorf("failed to parse time")
	}

	return time.Unix(0, nsec), nil
}

func (pm *ProviderManager) addProv(k cid.Cid, p peer.ID) error {
	now := time.Now()
	if provs, ok := pm.providers.Get(k); ok {
		provs.(*providerSet).setVal(p, now)
	} // else not cached, just write through

	return writeProviderEntry(pm.dstore, k, p, now)
}

func writeProviderEntry(dstore ds.Datastore, k cid.Cid, p peer.ID, t time.Time) error {
	dsk := mkProvKey(k) + "/" + base32.RawStdEncoding.EncodeToString([]byte(p))

	buf := make([]byte, 16)
	n := binary.PutVarint(buf, t.UnixNano())

	return dstore.Put(ds.NewKey(dsk), buf[:n])
}

func (pm *ProviderManager) deleteProvSet(k cid.Cid) error {
	pm.providers.Remove(k)

	res, err := pm.dstore.Query(dsq.Query{
		KeysOnly: true,
		Prefix:   mkProvKey(k),
	})
	if err != nil {
		return err
	}

	entries, err := res.Rest()
	if err != nil {
		return err
	}

	for _, e := range entries {
		err := pm.dstore.Delete(ds.NewKey(e.Key))
		if err != nil {
			log.Error("deleting provider set: ", err)
		}
	}
	return nil
}

func (pm *ProviderManager) getProvKeys() (func() (cid.Cid, bool), error) {
	res, err := pm.dstore.Query(dsq.Query{
		KeysOnly: true,
		Prefix:   providersKeyPrefix,
	})
	if err != nil {
		return nil, err
	}

	iter := func() (cid.Cid, bool) {
		for e := range res.Next() {
			parts := strings.Split(e.Key, "/")
			if len(parts) != 4 {
				log.Warningf("incorrectly formatted provider entry in datastore: %s", e.Key)
				continue
			}
			decoded, err := base32.RawStdEncoding.DecodeString(parts[2])
			if err != nil {
				log.Warning("error decoding base32 provider key: %s: %s", parts[2], err)
				continue
			}

			c, err := cid.Cast(decoded)
			if err != nil {
				log.Warning("error casting key to cid from datastore key: %s", err)
				continue
			}

			return c, true
		}
		return cid.Cid{}, false
	}

	return iter, nil
}

func (pm *ProviderManager) run() {
	tick := time.NewTicker(pm.cleanupInterval)
	for {
		select {
		case np := <-pm.newprovs:
			err := pm.addProv(np.k, np.val)
			if err != nil {
				log.Error("error adding new providers: ", err)
			}
		case gp := <-pm.getprovs:
			provs, err := pm.providersForKey(gp.k)
			if err != nil && err != ds.ErrNotFound {
				log.Error("error reading providers: ", err)
			}

			gp.resp <- provs
		case <-tick.C:
			keys, err := pm.getProvKeys()
			if err != nil {
				log.Error("Error loading provider keys: ", err)
				continue
			}
			now := time.Now()
			for {
				k, ok := keys()
				if !ok {
					break
				}

				provs, err := pm.getProvSet(k)
				if err != nil {
					log.Error("error loading known provset: ", err)
					continue
				}
				for p, t := range provs.set {
					if now.Sub(t) > ProvideValidity {
						delete(provs.set, p)
					}
				}
				// have we run out of providers?
				if len(provs.set) == 0 {
					provs.providers = nil
					err := pm.deleteProvSet(k)
					if err != nil {
						log.Error("error deleting provider set: ", err)
					}
				} else if len(provs.set) < len(provs.providers) {
					// We must have modified the providers set, recompute.
					provs.providers = make([]peer.ID, 0, len(provs.set))
					for p := range provs.set {
						provs.providers = append(provs.providers, p)
					}
				}
			}
		case <-pm.proc.Closing():
			tick.Stop()
			return
		}
	}
}

// AddProvider adds a provider.
func (pm *ProviderManager) AddProvider(ctx context.Context, k cid.Cid, val peer.ID) {
	prov := &addProv{
		k:   k,
		val: val,
	}
	select {
	case pm.newprovs <- prov:
	case <-ctx.Done():
	}
}

// GetProviders returns the set of providers for the given key.
// This method _does not_ copy the set. Do not modify it.
func (pm *ProviderManager) GetProviders(ctx context.Context, k cid.Cid) []peer.ID {
	gp := &getProv{
		k:    k,
		resp: make(chan []peer.ID, 1), // buffered to prevent sender from blocking
	}
	select {
	case <-ctx.Done():
		return nil
	case pm.getprovs <- gp:
	}
	select {
	case <-ctx.Done():
		return nil
	case peers := <-gp.resp:
		return peers
	}
}

func newProviderSet() *providerSet {
	return &providerSet{
		set: make(map[peer.ID]time.Time),
	}
}

func (ps *providerSet) Add(p peer.ID) {
	ps.setVal(p, time.Now())
}

func (ps *providerSet) setVal(p peer.ID, t time.Time) {
	_, found := ps.set[p]
	if !found {
		ps.providers = append(ps.providers, p)
	}

	ps.set[p] = t
}
<commit_msg>providers: use the non-locking LRU<commit_after>package providers

import (
	"context"
	"encoding/binary"
	"fmt"
	"strings"
	"time"

	lru "github.com/hashicorp/golang-lru/simplelru"
	cid "github.com/ipfs/go-cid"
	ds "github.com/ipfs/go-datastore"
	autobatch "github.com/ipfs/go-datastore/autobatch"
	dsq "github.com/ipfs/go-datastore/query"
	logging "github.com/ipfs/go-log"
	goprocess "github.com/jbenet/goprocess"
	goprocessctx "github.com/jbenet/goprocess/context"
	peer "github.com/libp2p/go-libp2p-peer"
	base32 "github.com/whyrusleeping/base32"
)

var batchBufferSize = 256

var log = logging.Logger("providers")

var lruCacheSize = 256
var ProvideValidity = time.Hour * 24
var defaultCleanupInterval = time.Hour

type ProviderManager struct {
	// all non channel fields are meant to be accessed only within
	// the run method
	providers *lru.LRU
	dstore    ds.Datastore

	newprovs chan *addProv
	getprovs chan *getProv
	proc     goprocess.Process

	cleanupInterval time.Duration
}

type providerSet struct {
	providers []peer.ID
	set       map[peer.ID]time.Time
}

type addProv struct {
	k   cid.Cid
	val peer.ID
}

type getProv struct {
	k    cid.Cid
	resp chan []peer.ID
}

func NewProviderManager(ctx context.Context, local peer.ID, dstore ds.Batching) *ProviderManager {
	pm := new(ProviderManager)
	pm.getprovs = make(chan *getProv)
	pm.newprovs = make(chan *addProv)
	pm.dstore = autobatch.NewAutoBatching(dstore, batchBufferSize)
	cache, err := lru.NewLRU(lruCacheSize, nil)
	if err != nil {
		panic(err) //only happens if negative value is passed to lru constructor
	}
	pm.providers = cache

	pm.proc = goprocessctx.WithContext(ctx)
	pm.cleanupInterval = defaultCleanupInterval
	pm.proc.Go(func(p goprocess.Process) { pm.run() })

	return pm
}

const providersKeyPrefix = "/providers/"

func mkProvKey(k cid.Cid) string {
	return providersKeyPrefix + base32.RawStdEncoding.EncodeToString(k.Bytes())
}

func (pm *ProviderManager) Process() goprocess.Process {
	return pm.proc
}

func (pm *ProviderManager) providersForKey(k cid.Cid) ([]peer.ID, error) {
	pset, err := pm.getProvSet(k)
	if err != nil {
		return nil, err
	}
	return pset.providers, nil
}

func (pm *ProviderManager) getProvSet(k cid.Cid) (*providerSet, error) {
	cached, ok := pm.providers.Get(k)
	if ok {
		return cached.(*providerSet), nil
	}

	pset, err := loadProvSet(pm.dstore, k)
	if err != nil {
		return nil, err
	}

	if len(pset.providers) > 0 {
		pm.providers.Add(k, pset)
	}

	return pset, nil
}

func loadProvSet(dstore ds.Datastore, k cid.Cid) (*providerSet, error) {
	res, err := dstore.Query(dsq.Query{Prefix: mkProvKey(k)})
	if err != nil {
		return nil, err
	}
	defer res.Close()

	now := time.Now()
	out := newProviderSet()
	for {
		e, ok := res.NextSync()
		if !ok {
			break
		}
		if e.Error != nil {
			log.Error("got an error: ", e.Error)
			continue
		}

		// check expiration time
		t, err := readTimeValue(e.Value)
		switch {
		case err != nil:
			// couldn't parse the time
			log.Warning("parsing providers record from disk: ", err)
			fallthrough
		case now.Sub(t) > ProvideValidity:
			// or just expired
			err = dstore.Delete(ds.RawKey(e.Key))
			if err != nil && err != ds.ErrNotFound {
				log.Warning("failed to remove provider record from disk: ", err)
			}
			continue
		}

		lix := strings.LastIndex(e.Key, "/")

		decstr, err := base32.RawStdEncoding.DecodeString(e.Key[lix+1:])
		if err != nil {
			log.Error("base32 decoding error: ", err)
			err = dstore.Delete(ds.RawKey(e.Key))
			if err != nil && err != ds.ErrNotFound {
				log.Warning("failed to remove provider record from disk: ", err)
			}
			continue
		}

		pid := peer.ID(decstr)

		out.setVal(pid, t)
	}

	return out, nil
}

func readTimeValue(data []byte) (time.Time, error) {
	nsec, n := binary.Varint(data)
	if n <= 0 {
		return time.Time{}, fmt.Errorf("failed to parse time")
	}

	return time.Unix(0, nsec), nil
}

func (pm *ProviderManager) addProv(k cid.Cid, p peer.ID) error {
	now := time.Now()
	if provs, ok := pm.providers.Get(k); ok {
		provs.(*providerSet).setVal(p, now)
	} // else not cached, just write through

	return writeProviderEntry(pm.dstore, k, p, now)
}

func writeProviderEntry(dstore ds.Datastore, k cid.Cid, p peer.ID, t time.Time) error {
	dsk := mkProvKey(k) + "/" + base32.RawStdEncoding.EncodeToString([]byte(p))

	buf := make([]byte, 16)
	n := binary.PutVarint(buf, t.UnixNano())

	return dstore.Put(ds.NewKey(dsk), buf[:n])
}

func (pm *ProviderManager) deleteProvSet(k cid.Cid) error {
	pm.providers.Remove(k)

	res, err := pm.dstore.Query(dsq.Query{
		KeysOnly: true,
		Prefix:   mkProvKey(k),
	})
	if err != nil {
		return err
	}

	entries, err := res.Rest()
	if err != nil {
		return err
	}

	for _, e := range entries {
		err := pm.dstore.Delete(ds.NewKey(e.Key))
		if err != nil {
			log.Error("deleting provider set: ", err)
		}
	}
	return nil
}

func (pm *ProviderManager) getProvKeys() (func() (cid.Cid, bool), error) {
	res, err := pm.dstore.Query(dsq.Query{
		KeysOnly: true,
		Prefix:   providersKeyPrefix,
	})
	if err != nil {
		return nil, err
	}

	iter := func() (cid.Cid, bool) {
		for e := range res.Next() {
			parts := strings.Split(e.Key, "/")
			if len(parts) != 4 {
				log.Warningf("incorrectly formatted provider entry in datastore: %s", e.Key)
				continue
			}
			decoded, err := base32.RawStdEncoding.DecodeString(parts[2])
			if err != nil {
				log.Warning("error decoding base32 provider key: %s: %s", parts[2], err)
				continue
			}

			c, err := cid.Cast(decoded)
			if err != nil {
				log.Warning("error casting key to cid from datastore key: %s", err)
				continue
			}

			return c, true
		}
		return cid.Cid{}, false
	}

	return iter, nil
}

func (pm *ProviderManager) run() {
	tick := time.NewTicker(pm.cleanupInterval)
	for {
		select {
		case np := <-pm.newprovs:
			err := pm.addProv(np.k, np.val)
			if err != nil {
				log.Error("error adding new providers: ", err)
			}
		case gp := <-pm.getprovs:
			provs, err := pm.providersForKey(gp.k)
			if err != nil && err != ds.ErrNotFound {
				log.Error("error reading providers: ", err)
			}

			gp.resp <- provs
		case <-tick.C:
			keys, err := pm.getProvKeys()
			if err != nil {
				log.Error("Error loading provider keys: ", err)
				continue
			}
			now := time.Now()
			for {
				k, ok := keys()
				if !ok {
					break
				}

				provs, err := pm.getProvSet(k)
				if err != nil {
					log.Error("error loading known provset: ", err)
					continue
				}
				for p, t := range provs.set {
					if now.Sub(t) > ProvideValidity {
						delete(provs.set, p)
					}
				}
				// have we run out of providers?
				if len(provs.set) == 0 {
					provs.providers = nil
					err := pm.deleteProvSet(k)
					if err != nil {
						log.Error("error deleting provider set: ", err)
					}
				} else if len(provs.set) < len(provs.providers) {
					// We must have modified the providers set, recompute.
					provs.providers = make([]peer.ID, 0, len(provs.set))
					for p := range provs.set {
						provs.providers = append(provs.providers, p)
					}
				}
			}
		case <-pm.proc.Closing():
			tick.Stop()
			return
		}
	}
}

// AddProvider adds a provider.
func (pm *ProviderManager) AddProvider(ctx context.Context, k cid.Cid, val peer.ID) {
	prov := &addProv{
		k:   k,
		val: val,
	}
	select {
	case pm.newprovs <- prov:
	case <-ctx.Done():
	}
}

// GetProviders returns the set of providers for the given key.
// This method _does not_ copy the set. Do not modify it.
func (pm *ProviderManager) GetProviders(ctx context.Context, k cid.Cid) []peer.ID {
	gp := &getProv{
		k:    k,
		resp: make(chan []peer.ID, 1), // buffered to prevent sender from blocking
	}
	select {
	case <-ctx.Done():
		return nil
	case pm.getprovs <- gp:
	}
	select {
	case <-ctx.Done():
		return nil
	case peers := <-gp.resp:
		return peers
	}
}

func newProviderSet() *providerSet {
	return &providerSet{
		set: make(map[peer.ID]time.Time),
	}
}

func (ps *providerSet) Add(p peer.ID) {
	ps.setVal(p, time.Now())
}

func (ps *providerSet) setVal(p peer.ID, t time.Time) {
	_, found := ps.set[p]
	if !found {
		ps.providers = append(ps.providers, p)
	}

	ps.set[p] = t
}
<|endoftext|>


<commit_before>package server

import (
	"bufio"
	"bytes"
	"crypto/md5"
	"encoding/hex"
	"fmt"
	"io"
	"math/rand"
	"os"
	"path"
	"reflect"
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/pachyderm/pachyderm/src/client"
	"github.com/pachyderm/pachyderm/src/client/pfs"
	"github.com/pachyderm/pachyderm/src/client/pkg/require"
	"github.com/pachyderm/pachyderm/src/client/pps"
	versionlib "github.com/pachyderm/pachyderm/src/client/version"
	"github.com/pachyderm/pachyderm/src/server/pkg/backoff"
	tu "github.com/pachyderm/pachyderm/src/server/pkg/testutil"
	"github.com/pachyderm/pachyderm/src/server/pkg/workload"

	"github.com/golang/snappy"
)

const (
	KB = 1024
	MB = 1024 * KB
)

var pachClient *client.APIClient
var getPachClientOnce sync.Once

func getPachClient(t testing.TB) *client.APIClient {
	getPachClientOnce.Do(func() {
		var err error
		if addr := os.Getenv("PACHD_PORT_650_TCP_ADDR"); addr != "" {
			pachClient, err = client.NewInCluster()
		} else {
			pachClient, err = client.NewForTest()
		}
		require.NoError(t, err)
	})
	return pachClient
}

func collectCommitInfos(t testing.TB, commitInfoIter client.CommitInfoIterator) []*pfs.CommitInfo {
	var commitInfos []*pfs.CommitInfo
	for {
		commitInfo, err := commitInfoIter.Next()
		if err == io.EOF {
			return commitInfos
		}
		require.NoError(t, err)
		commitInfos = append(commitInfos, commitInfo)
	}
}

// TODO(msteffen) equivalent to funciton in src/server/auth/server/admin_test.go.
// These should be unified.
func RepoInfoToName(repoInfo interface{}) interface{} {
	return repoInfo.(*pfs.RepoInfo).Repo.Name
}

// testExtractRestored effectively implements both TestExtractRestoreObjects
// TestExtractRestoreNoObjects, as their logic is mostly the same
func testExtractRestore(t *testing.T, testObjects bool) {
	if testing.Short() {
		t.Skip("Skipping integration tests in short mode")
	}

	c := getPachClient(t)
	require.NoError(t, c.DeleteAll())

	dataRepo := tu.UniqueString("TestExtractRestoreObjects-in-")
	require.NoError(t, c.CreateRepo(dataRepo))

	// Create input data
	nCommits := 2
	r := rand.New(rand.NewSource(45))
	fileHashes := make([]string, 0, nCommits)
	for i := 0; i < nCommits; i++ {
		hash := md5.New()
		fileContent := workload.RandString(r, 40*MB)
		_, err := c.PutFile(dataRepo, "master", fmt.Sprintf("file-%d", i),
			io.TeeReader(strings.NewReader(fileContent), hash))
		require.NoError(t, err)
		fileHashes = append(fileHashes, hex.EncodeToString(hash.Sum(nil)))
	}

	// Create test pipelines
	numPipelines := 3
	var input, pipeline string
	input = dataRepo
	for i := 0; i < numPipelines; i++ {
		pipeline = tu.UniqueString(fmt.Sprintf("TestExtractRestoreObjects-P%d-", i))
		require.NoError(t, c.CreatePipeline(
			pipeline,
			"",
			[]string{"bash"},
			[]string{
				fmt.Sprintf("cp /pfs/%s/* /pfs/out/", input),
			},
			&pps.ParallelismSpec{
				Constant: 1,
			},
			client.NewPFSInput(input, "/*"),
			"",
			false,
		))
		input = pipeline
	}

	// Wait for pipelines to process input data
	commitIter, err := c.FlushCommit([]*pfs.Commit{client.NewCommit(dataRepo, "master")}, nil)
	require.NoError(t, err)
	commitInfos := collectCommitInfos(t, commitIter)
	require.Equal(t, numPipelines, len(commitInfos))

	// Extract existing cluster state
	ops, err := c.ExtractAll(testObjects)
	require.NoError(t, err)

	require.NoError(t, c.DeleteAll())

	if testObjects {
		// Delete existing objects
		require.NoError(t, c.GarbageCollect(10000))
	}

	// Restore metadata and possibly objects
	require.NoError(t, c.Restore(ops))

	// Wait for re-created pipelines to process recreated input data
	commitIter, err = c.FlushCommit([]*pfs.Commit{client.NewCommit(dataRepo, "master")}, nil)
	require.NoError(t, err)
	commitInfos = collectCommitInfos(t, commitIter)
	require.Equal(t, numPipelines, len(commitInfos))

	// Confirm all the recreated jobs passed
	require.NoErrorWithinTRetry(t, 30*time.Second, func() error {
		jobInfos, err := c.ListJob("", nil, nil, -1, true) // make sure jobs all succeeded
		if err != nil {
			return err
		}
		// Extract places commits before pipelines, so that all commits are processed
		// after Restore() (thus |jobInfos| = nCommits * numPipelines)
		if len(jobInfos) != nCommits*numPipelines {
			return fmt.Errorf("expected %d commits, but only encountered %d",
				nCommits*numPipelines, len(jobInfos))
		}
		for _, ji := range jobInfos {
			// race--we may call listJob between when a job's output commit is closed
			// and when its state is updated
			if ji.State != pps.JobState_JOB_SUCCESS {
				return fmt.Errorf("expected job %q to be in state JOB_SUCCESS but was %q",
					ji.Job.ID, ji.State.String())
			}
			// Job must ultimately succeed
			require.Equal(t, "JOB_SUCCESS", ji.State.String())
		}
		return nil
	})

	// Make sure all branches were recreated
	bis, err := c.ListBranch(dataRepo)
	require.NoError(t, err)
	require.Equal(t, 1, len(bis))

	// Check input data
	// This check uses a backoff because sometimes GetFile causes pachd to OOM
	var restoredFileHashes []string
	require.NoError(t, backoff.Retry(func() error {
		restoredFileHashes = make([]string, 0, nCommits)
		for i := 0; i < nCommits; i++ {
			hash := md5.New()
			err := c.GetFile(dataRepo, "master", fmt.Sprintf("file-%d", i), 0, 0, hash)
			if err != nil {
				return err
			}
			restoredFileHashes = append(restoredFileHashes, hex.EncodeToString(hash.Sum(nil)))
		}
		return nil
	}, backoff.NewTestingBackOff()))
	require.ElementsEqual(t, fileHashes, restoredFileHashes)

	// Check output data
	// This check uses a backoff because sometimes GetFile causes pachd to OOM
	require.NoError(t, backoff.Retry(func() error {
		restoredFileHashes = make([]string, 0, nCommits)
		for i := 0; i < nCommits; i++ {
			hash := md5.New()
			err := c.GetFile(pipeline, "master", fmt.Sprintf("file-%d", i), 0, 0, hash)
			if err != nil {
				return err
			}
			restoredFileHashes = append(restoredFileHashes, hex.EncodeToString(hash.Sum(nil)))
		}
		return nil
	}, backoff.NewTestingBackOff()))
	require.ElementsEqual(t, fileHashes, restoredFileHashes)
}

// TestExtractRestoreNoObjects tests extraction and restoration in the case
// where existing objects are re-used (common for cloud deployments, as objects
// are stored outside of kubernetes, in object store)
func TestExtractRestoreNoObjects(t *testing.T) {
	testExtractRestore(t, false)
}

// TestExtractRestoreObjects tests extraction and restoration of objects. Note
// that since 1.8, only data in input repos is referenced by objects, so this
// tests extracting/restoring an input repo.
func TestExtractRestoreObjects(t *testing.T) {
	testExtractRestore(t, true)
}

func TestExtractRestoreHeadlessBranches(t *testing.T) {
	if testing.Short() {
		t.Skip("Skipping integration tests in short mode")
	}

	c := getPachClient(t)
	require.NoError(t, c.DeleteAll())

	dataRepo := tu.UniqueString("TestExtractRestore_data")
	require.NoError(t, c.CreateRepo(dataRepo))

	// create a headless branch
	require.NoError(t, c.CreateBranch(dataRepo, "headless", "", nil))

	ops, err := c.ExtractAll(false)
	require.NoError(t, err)
	require.NoError(t, c.DeleteAll())
	require.NoError(t, c.Restore(ops))

	bis, err := c.ListBranch(dataRepo)
	require.NoError(t, err)
	require.Equal(t, 1, len(bis))
	require.Equal(t, "headless", bis[0].Branch.Name)
}

func TestExtractVersion(t *testing.T) {
	if testing.Short() {
		t.Skip("Skipping integration tests in short mode")
	}

	c := getPachClient(t)
	require.NoError(t, c.DeleteAll())

	dataRepo := tu.UniqueString("TestExtractRestore_data")
	require.NoError(t, c.CreateRepo(dataRepo))

	r := rand.New(rand.NewSource(45))
	_, err := c.PutFile(dataRepo, "master", "file", strings.NewReader(workload.RandString(r, 40*MB)))
	require.NoError(t, err)

	pipeline := tu.UniqueString("TestExtractRestore")
	require.NoError(t, c.CreatePipeline(
		pipeline,
		"",
		[]string{"bash"},
		[]string{
			fmt.Sprintf("cp /pfs/%s/* /pfs/out/", dataRepo),
		},
		&pps.ParallelismSpec{
			Constant: 1,
		},
		client.NewPFSInput(dataRepo, "/*"),
		"",
		false,
	))

	ops, err := c.ExtractAll(false)
	require.NoError(t, err)
	require.True(t, len(ops) > 0)

	// Check that every Op looks right; the version set matches pachd's version
	for _, op := range ops {
		opV := reflect.ValueOf(op).Elem()
		var versions, nonemptyVersions int
		for i := 0; i < opV.NumField(); i++ {
			fDesc := opV.Type().Field(i)
			if !strings.HasPrefix(fDesc.Name, "Op") {
				continue
			}
			versions++
			if strings.HasSuffix(fDesc.Name,
				fmt.Sprintf("%d_%d", versionlib.MajorVersion, versionlib.MinorVersion)) {
				require.False(t, opV.Field(i).IsNil())
				nonemptyVersions++
			} else {
				require.True(t, opV.Field(i).IsNil())
			}
		}
		require.Equal(t, 1, nonemptyVersions)
		require.True(t, versions > 1)
	}
}

func TestMigrateFrom1_7(t *testing.T) {
	if testing.Short() {
		t.Skip("Skipping integration tests in short mode")
	}

	// Clear pachyderm cluster (so that next cluster starts up in a clean environment)
	c := getPachClient(t)
	require.NoError(t, c.DeleteAll())

	// Restore dumped metadata (now that objects are present)
	md, err := os.Open(path.Join(os.Getenv("GOPATH"),
		"src/github.com/pachyderm/pachyderm/etc/testing/migration/1_7/sort.metadata"))
	require.NoError(t, err)
	require.NoError(t, c.RestoreReader(snappy.NewReader(md)))
	require.NoError(t, md.Close())

	// Wait for final imported commit to be processed
	commitIter, err := c.FlushCommit([]*pfs.Commit{client.NewCommit("left", "master")}, nil)
	require.NoError(t, err)
	commitInfos := collectCommitInfos(t, commitIter)
	// filter-left and filter-right both compute a join of left and
	// right--depending on when the final commit to 'left' was added, it may have
	// been processed multiple times (should be n * 3, as there are 3 pipelines)
	require.True(t, len(commitInfos) >= 3)

	// Inspect input
	commits, err := c.ListCommit("left", "master", "", 0)
	require.NoError(t, err)
	require.Equal(t, 3, len(commits))
	commits, err = c.ListCommit("right", "master", "", 0)
	require.NoError(t, err)
	require.Equal(t, 3, len(commits))

	// Inspect output
	repos, err := c.ListRepo()
	require.NoError(t, err)
	require.ElementsEqualUnderFn(t,
		[]string{"left", "right", "copy", "sort"},
		repos, RepoInfoToName)

	// make sure all numbers 0-99 are in /nums
	var buf bytes.Buffer
	require.NoError(t, c.GetFile("sort", "master", "/nums", 0, 0, &buf))
	s := bufio.NewScanner(&buf)
	numbers := make(map[string]struct{})
	for s.Scan() {
		numbers[s.Text()] = struct{}{}
	}
	require.Equal(t, 100, len(numbers)) // job processed all inputs

	// Confirm stats commits are present
	commits, err = c.ListCommit("sort", "stats", "", 0)
	require.NoError(t, err)
	require.Equal(t, 6, len(commits))
}

func int64p(i int64) *int64 {
	return &i
}
<commit_msg>Makes TestExtractRestore pass, and adds a few checks.<commit_after>package server

import (
	"bufio"
	"bytes"
	"crypto/md5"
	"encoding/hex"
	"fmt"
	"io"
	"math/rand"
	"os"
	"path"
	"reflect"
	"strings"
	"sync"
	"testing"

	"github.com/pachyderm/pachyderm/src/client"
	"github.com/pachyderm/pachyderm/src/client/pfs"
	"github.com/pachyderm/pachyderm/src/client/pkg/require"
	"github.com/pachyderm/pachyderm/src/client/pps"
	versionlib "github.com/pachyderm/pachyderm/src/client/version"
	"github.com/pachyderm/pachyderm/src/server/pkg/backoff"
	tu "github.com/pachyderm/pachyderm/src/server/pkg/testutil"
	"github.com/pachyderm/pachyderm/src/server/pkg/workload"

	"github.com/golang/snappy"
)

const (
	KB = 1024
	MB = 1024 * KB
)

var pachClient *client.APIClient
var getPachClientOnce sync.Once

func getPachClient(t testing.TB) *client.APIClient {
	getPachClientOnce.Do(func() {
		var err error
		if addr := os.Getenv("PACHD_PORT_650_TCP_ADDR"); addr != "" {
			pachClient, err = client.NewInCluster()
		} else {
			pachClient, err = client.NewForTest()
		}
		require.NoError(t, err)
	})
	return pachClient
}

func collectCommitInfos(t testing.TB, commitInfoIter client.CommitInfoIterator) []*pfs.CommitInfo {
	var commitInfos []*pfs.CommitInfo
	for {
		commitInfo, err := commitInfoIter.Next()
		if err == io.EOF {
			return commitInfos
		}
		require.NoError(t, err)
		commitInfos = append(commitInfos, commitInfo)
	}
}

// TODO(msteffen) equivalent to funciton in src/server/auth/server/admin_test.go.
// These should be unified.
func RepoInfoToName(repoInfo interface{}) interface{} {
	return repoInfo.(*pfs.RepoInfo).Repo.Name
}

// testExtractRestored effectively implements both TestExtractRestoreObjects
// TestExtractRestoreNoObjects, as their logic is mostly the same
func testExtractRestore(t *testing.T, testObjects bool) {
	if testing.Short() {
		t.Skip("Skipping integration tests in short mode")
	}

	c := getPachClient(t)
	require.NoError(t, c.DeleteAll())

	dataRepo := tu.UniqueString("TestExtractRestoreObjects-in-")
	require.NoError(t, c.CreateRepo(dataRepo))

	// Create input data
	nCommits := 2
	r := rand.New(rand.NewSource(45))
	fileHashes := make([]string, 0, nCommits)
	for i := 0; i < nCommits; i++ {
		hash := md5.New()
		fileContent := workload.RandString(r, 40*MB)
		_, err := c.PutFile(dataRepo, "master", fmt.Sprintf("file-%d", i),
			io.TeeReader(strings.NewReader(fileContent), hash))
		require.NoError(t, err)
		fileHashes = append(fileHashes, hex.EncodeToString(hash.Sum(nil)))
	}

	// Create test pipelines
	numPipelines := 3
	var input, pipeline string
	input = dataRepo
	for i := 0; i < numPipelines; i++ {
		pipeline = tu.UniqueString(fmt.Sprintf("TestExtractRestoreObjects-P%d-", i))
		require.NoError(t, c.CreatePipeline(
			pipeline,
			"",
			[]string{"bash"},
			[]string{
				fmt.Sprintf("cp /pfs/%s/* /pfs/out/", input),
			},
			&pps.ParallelismSpec{
				Constant: 1,
			},
			client.NewPFSInput(input, "/*"),
			"",
			false,
		))
		input = pipeline
	}

	// Wait for pipelines to process input data
	commitIter, err := c.FlushCommit([]*pfs.Commit{client.NewCommit(dataRepo, "master")}, nil)
	require.NoError(t, err)
	commitInfos := collectCommitInfos(t, commitIter)
	require.Equal(t, numPipelines, len(commitInfos))

	// Extract existing cluster state
	ops, err := c.ExtractAll(testObjects)
	require.NoError(t, err)

	require.NoError(t, c.DeleteAll())

	if testObjects {
		// Delete existing objects
		require.NoError(t, c.GarbageCollect(10000))
	}

	// Restore metadata and possibly objects
	require.NoError(t, c.Restore(ops))

	// Wait for re-created pipelines to process recreated input data
	commitIter, err = c.FlushCommit([]*pfs.Commit{client.NewCommit(dataRepo, "master")}, nil)
	require.NoError(t, err)
	commitInfos = collectCommitInfos(t, commitIter)
	require.Equal(t, numPipelines, len(commitInfos))

	// Confirm all the recreated jobs passed
	jis, err := pachClient.FlushJobAll([]*pfs.Commit{client.NewCommit(dataRepo, "master")}, nil)
	// One job per pipeline, as above
	require.Equal(t, numPipelines, len(jis))
	for _, ji := range jis {
		// Job must ultimately succeed
		require.Equal(t, "JOB_SUCCESS", ji.State.String())
		require.Equal(t, int64(2), ji.DataTotal)
		require.NotNil(t, ji.Started)
		require.NotNil(t, ji.Finished)
	}

	// Make sure all branches were recreated
	bis, err := c.ListBranch(dataRepo)
	require.NoError(t, err)
	require.Equal(t, 1, len(bis))

	// Check input data
	// This check uses a backoff because sometimes GetFile causes pachd to OOM
	var restoredFileHashes []string
	require.NoError(t, backoff.Retry(func() error {
		restoredFileHashes = make([]string, 0, nCommits)
		for i := 0; i < nCommits; i++ {
			hash := md5.New()
			err := c.GetFile(dataRepo, "master", fmt.Sprintf("file-%d", i), 0, 0, hash)
			if err != nil {
				return err
			}
			restoredFileHashes = append(restoredFileHashes, hex.EncodeToString(hash.Sum(nil)))
		}
		return nil
	}, backoff.NewTestingBackOff()))
	require.ElementsEqual(t, fileHashes, restoredFileHashes)

	// Check output data
	// This check uses a backoff because sometimes GetFile causes pachd to OOM
	require.NoError(t, backoff.Retry(func() error {
		restoredFileHashes = make([]string, 0, nCommits)
		for i := 0; i < nCommits; i++ {
			hash := md5.New()
			err := c.GetFile(pipeline, "master", fmt.Sprintf("file-%d", i), 0, 0, hash)
			if err != nil {
				return err
			}
			restoredFileHashes = append(restoredFileHashes, hex.EncodeToString(hash.Sum(nil)))
		}
		return nil
	}, backoff.NewTestingBackOff()))
	require.ElementsEqual(t, fileHashes, restoredFileHashes)
}

// TestExtractRestoreNoObjects tests extraction and restoration in the case
// where existing objects are re-used (common for cloud deployments, as objects
// are stored outside of kubernetes, in object store)
func TestExtractRestoreNoObjects(t *testing.T) {
	testExtractRestore(t, false)
}

// TestExtractRestoreObjects tests extraction and restoration of objects. Note
// that since 1.8, only data in input repos is referenced by objects, so this
// tests extracting/restoring an input repo.
func TestExtractRestoreObjects(t *testing.T) {
	testExtractRestore(t, true)
}

func TestExtractRestoreHeadlessBranches(t *testing.T) {
	if testing.Short() {
		t.Skip("Skipping integration tests in short mode")
	}

	c := getPachClient(t)
	require.NoError(t, c.DeleteAll())

	dataRepo := tu.UniqueString("TestExtractRestore_data")
	require.NoError(t, c.CreateRepo(dataRepo))

	// create a headless branch
	require.NoError(t, c.CreateBranch(dataRepo, "headless", "", nil))

	ops, err := c.ExtractAll(false)
	require.NoError(t, err)
	require.NoError(t, c.DeleteAll())
	require.NoError(t, c.Restore(ops))

	bis, err := c.ListBranch(dataRepo)
	require.NoError(t, err)
	require.Equal(t, 1, len(bis))
	require.Equal(t, "headless", bis[0].Branch.Name)
}

func TestExtractVersion(t *testing.T) {
	if testing.Short() {
		t.Skip("Skipping integration tests in short mode")
	}

	c := getPachClient(t)
	require.NoError(t, c.DeleteAll())

	dataRepo := tu.UniqueString("TestExtractRestore_data")
	require.NoError(t, c.CreateRepo(dataRepo))

	r := rand.New(rand.NewSource(45))
	_, err := c.PutFile(dataRepo, "master", "file", strings.NewReader(workload.RandString(r, 40*MB)))
	require.NoError(t, err)

	pipeline := tu.UniqueString("TestExtractRestore")
	require.NoError(t, c.CreatePipeline(
		pipeline,
		"",
		[]string{"bash"},
		[]string{
			fmt.Sprintf("cp /pfs/%s/* /pfs/out/", dataRepo),
		},
		&pps.ParallelismSpec{
			Constant: 1,
		},
		client.NewPFSInput(dataRepo, "/*"),
		"",
		false,
	))

	ops, err := c.ExtractAll(false)
	require.NoError(t, err)
	require.True(t, len(ops) > 0)

	// Check that every Op looks right; the version set matches pachd's version
	for _, op := range ops {
		opV := reflect.ValueOf(op).Elem()
		var versions, nonemptyVersions int
		for i := 0; i < opV.NumField(); i++ {
			fDesc := opV.Type().Field(i)
			if !strings.HasPrefix(fDesc.Name, "Op") {
				continue
			}
			versions++
			if strings.HasSuffix(fDesc.Name,
				fmt.Sprintf("%d_%d", versionlib.MajorVersion, versionlib.MinorVersion)) {
				require.False(t, opV.Field(i).IsNil())
				nonemptyVersions++
			} else {
				require.True(t, opV.Field(i).IsNil())
			}
		}
		require.Equal(t, 1, nonemptyVersions)
		require.True(t, versions > 1)
	}
}

func TestMigrateFrom1_7(t *testing.T) {
	if testing.Short() {
		t.Skip("Skipping integration tests in short mode")
	}

	// Clear pachyderm cluster (so that next cluster starts up in a clean environment)
	c := getPachClient(t)
	require.NoError(t, c.DeleteAll())

	// Restore dumped metadata (now that objects are present)
	md, err := os.Open(path.Join(os.Getenv("GOPATH"),
		"src/github.com/pachyderm/pachyderm/etc/testing/migration/1_7/sort.metadata"))
	require.NoError(t, err)
	require.NoError(t, c.RestoreReader(snappy.NewReader(md)))
	require.NoError(t, md.Close())

	// Wait for final imported commit to be processed
	commitIter, err := c.FlushCommit([]*pfs.Commit{client.NewCommit("left", "master")}, nil)
	require.NoError(t, err)
	commitInfos := collectCommitInfos(t, commitIter)
	// filter-left and filter-right both compute a join of left and
	// right--depending on when the final commit to 'left' was added, it may have
	// been processed multiple times (should be n * 3, as there are 3 pipelines)
	require.True(t, len(commitInfos) >= 3)

	// Inspect input
	commits, err := c.ListCommit("left", "master", "", 0)
	require.NoError(t, err)
	require.Equal(t, 3, len(commits))
	commits, err = c.ListCommit("right", "master", "", 0)
	require.NoError(t, err)
	require.Equal(t, 3, len(commits))

	// Inspect output
	repos, err := c.ListRepo()
	require.NoError(t, err)
	require.ElementsEqualUnderFn(t,
		[]string{"left", "right", "copy", "sort"},
		repos, RepoInfoToName)

	// make sure all numbers 0-99 are in /nums
	var buf bytes.Buffer
	require.NoError(t, c.GetFile("sort", "master", "/nums", 0, 0, &buf))
	s := bufio.NewScanner(&buf)
	numbers := make(map[string]struct{})
	for s.Scan() {
		numbers[s.Text()] = struct{}{}
	}
	require.Equal(t, 100, len(numbers)) // job processed all inputs

	// Confirm stats commits are present
	commits, err = c.ListCommit("sort", "stats", "", 0)
	require.NoError(t, err)
	require.Equal(t, 6, len(commits))
}

func int64p(i int64) *int64 {
	return &i
}
<|endoftext|>


<commit_before>package server

import (
	"bytes"
	"fmt"
	"io"
	"sync"

	"github.com/golang/snappy"

	"github.com/gogo/protobuf/types"
	"github.com/pachyderm/pachyderm/src/client"
	"github.com/pachyderm/pachyderm/src/client/admin"
	"github.com/pachyderm/pachyderm/src/client/pfs"
	"github.com/pachyderm/pachyderm/src/client/pkg/grpcutil"
	"github.com/pachyderm/pachyderm/src/client/pkg/pbutil"
	"github.com/pachyderm/pachyderm/src/client/pps"
	"github.com/pachyderm/pachyderm/src/server/pkg/log"
	"github.com/pachyderm/pachyderm/src/server/pkg/obj"
)

type apiServer struct {
	log.Logger
	address        string
	pachClient     *client.APIClient
	pachClientOnce sync.Once
}

func (a *apiServer) Extract(request *admin.ExtractRequest, extractServer admin.API_ExtractServer) (retErr error) {
	pachClient := a.getPachClient()
	pachClient = pachClient.WithCtx(extractServer.Context())
	handleOp := extractServer.Send
	if request.URL != "" {
		url, err := obj.ParseURL(request.URL)
		if err != nil {
			return fmt.Errorf("error parsing url %v: %v", request.URL, err)
		}
		objClient, err := obj.NewClientFromURLAndSecret(extractServer.Context(), url)
		if err != nil {
			return err
		}
		objW, err := objClient.Writer(url.Object)
		if err != nil {
			return err
		}
		snappyW := snappy.NewBufferedWriter(objW)
		defer func() {
			if err := snappyW.Close(); err != nil && retErr == nil {
				retErr = err
			}
		}()
		w := pbutil.NewWriter(snappyW)
		handleOp = func(op *admin.Op) error { return w.Write(op) }
	}
	if !request.NoObjects {
		w := &extractObjectWriter{extractServer}
		if err := pachClient.ListObject(func(object *pfs.Object) error {
			if err := pachClient.GetObject(object.Hash, w); err != nil {
				return err
			}
			// empty PutObjectRequest to indicate EOF
			return handleOp(&admin.Op{Op1_7: &admin.Op1_7{Object: &pfs.PutObjectRequest{}}})
		}); err != nil {
			return err
		}
		if err := pachClient.ListTag(func(resp *pfs.ListTagsResponse) error {
			return handleOp(&admin.Op{Op1_7: &admin.Op1_7{
				Tag: &pfs.TagObjectRequest{
					Object: resp.Object,
					Tags:   []*pfs.Tag{resp.Tag},
				},
			}})
		}); err != nil {
			return err
		}
	}
	if !request.NoRepos {
		ris, err := pachClient.ListRepo(nil)
		if err != nil {
			return err
		}
		for _, ri := range ris {
			if len(ri.Provenance) > 0 {
				continue
			}
			if err := handleOp(&admin.Op{Op1_7: &admin.Op1_7{
				Repo: &pfs.CreateRepoRequest{
					Repo:        ri.Repo,
					Provenance:  ri.Provenance,
					Description: ri.Description,
				}},
			}); err != nil {
				return err
			}
			cis, err := pachClient.ListCommit(ri.Repo.Name, "", "", 0)
			if err != nil {
				return err
			}
			for _, ci := range sortCommitInfos(cis) {
				// Even without a parent, ParentCommit is used to indicate which
				// repo to make the commit in.
				if ci.ParentCommit == nil {
					ci.ParentCommit = client.NewCommit(ri.Repo.Name, "")
				}
				if err := handleOp(&admin.Op{Op1_7: &admin.Op1_7{
					Commit: &pfs.BuildCommitRequest{
						Parent: ci.ParentCommit,
						Tree:   ci.Tree,
						ID:     ci.Commit.ID,
					},
				}}); err != nil {
					return err
				}
			}
			bis, err := pachClient.ListBranch(ri.Repo.Name)
			if err != nil {
				return err
			}
			for _, bi := range bis {
				if err := handleOp(&admin.Op{Op1_7: &admin.Op1_7{
					Branch: &pfs.SetBranchRequest{
						Commit: bi.Head,
						Branch: bi.Name,
					},
				}}); err != nil {
					return err
				}
			}
		}
	}
	if !request.NoPipelines {
		pis, err := pachClient.ListPipeline()
		if err != nil {
			return err
		}
		for _, pi := range pis {
			if err := handleOp(&admin.Op{Op1_7: &admin.Op1_7{
				Pipeline: &pps.CreatePipelineRequest{
					Pipeline:           pi.Pipeline,
					Transform:          pi.Transform,
					ParallelismSpec:    pi.ParallelismSpec,
					Egress:             pi.Egress,
					OutputBranch:       pi.OutputBranch,
					ScaleDownThreshold: pi.ScaleDownThreshold,
					ResourceRequests:   pi.ResourceRequests,
					ResourceLimits:     pi.ResourceLimits,
					Input:              pi.Input,
					Description:        pi.Description,
					Incremental:        pi.Incremental,
					CacheSize:          pi.CacheSize,
					EnableStats:        pi.EnableStats,
					Batch:              pi.Batch,
					MaxQueueSize:       pi.MaxQueueSize,
					Service:            pi.Service,
					ChunkSpec:          pi.ChunkSpec,
					DatumTimeout:       pi.DatumTimeout,
					JobTimeout:         pi.JobTimeout,
					Salt:               pi.Salt,
				}}}); err != nil {
				return err
			}
		}
	}
	return nil
}

func sortCommitInfos(cis []*pfs.CommitInfo) []*pfs.CommitInfo {
	commitMap := make(map[string]*pfs.CommitInfo)
	for _, ci := range cis {
		commitMap[ci.Commit.ID] = ci
	}
	var result []*pfs.CommitInfo
	for _, ci := range cis {
		if commitMap[ci.Commit.ID] == nil {
			continue
		}
		var localResult []*pfs.CommitInfo
		for ci != nil {
			localResult = append(localResult, ci)
			delete(commitMap, ci.Commit.ID)
			if ci.ParentCommit != nil {
				ci = commitMap[ci.ParentCommit.ID]
			} else {
				ci = nil
			}
		}
		for i := range localResult {
			result = append(result, localResult[len(localResult)-i-1])
		}
	}
	return result
}

func (a *apiServer) Restore(restoreServer admin.API_RestoreServer) (retErr error) {
	ctx := restoreServer.Context()
	pachClient := a.getPachClient()
	defer func() {
		for {
			_, err := restoreServer.Recv()
			if err != nil {
				break
			}
		}
		if err := restoreServer.SendAndClose(&types.Empty{}); err != nil && retErr == nil {
			retErr = err
		}
	}()
	var r pbutil.Reader
	for {
		var op *admin.Op
		if r == nil {
			req, err := restoreServer.Recv()
			if err != nil {
				if err == io.EOF {
					return nil
				}
				return err
			}
			if req.URL != "" {
				url, err := obj.ParseURL(req.URL)
				if err != nil {
					return fmt.Errorf("error parsing url %v: %v", req.URL, err)
				}
				objClient, err := obj.NewClientFromURLAndSecret(restoreServer.Context(), url)
				if err != nil {
					return err
				}
				objR, err := objClient.Reader(url.Object, 0, 0)
				if err != nil {
					return err
				}
				snappyR := snappy.NewReader(objR)
				r = pbutil.NewReader(snappyR)
				continue
			} else {
				op = req.Op
			}
		} else {
			op = &admin.Op{}
			if err := r.Read(op); err != nil {
				if err == io.EOF {
					return nil
				}
				return err
			}
		}
		switch {
		case op.Op1_7 != nil && op.Op1_7.Object != nil:
			r := &extractObjectReader{adminAPIRestoreServer: restoreServer}
			r.buf.Write(op.Op1_7.Object.Value)
			if _, _, err := pachClient.PutObject(r); err != nil {
				return fmt.Errorf("error putting object: %v", err)
			}
		case op.Op1_7 != nil && op.Op1_7.Tag != nil:
			if _, err := pachClient.ObjectAPIClient.TagObject(ctx, op.Op1_7.Tag); err != nil {
				return fmt.Errorf("error tagging object: %v", grpcutil.ScrubGRPC(err))
			}
		case op.Op1_7 != nil && op.Op1_7.Repo != nil:
			if _, err := pachClient.PfsAPIClient.CreateRepo(ctx, op.Op1_7.Repo); err != nil {
				return fmt.Errorf("error creating repo: %v", grpcutil.ScrubGRPC(err))
			}
		case op.Op1_7 != nil && op.Op1_7.Commit != nil:
			if _, err := pachClient.PfsAPIClient.BuildCommit(ctx, op.Op1_7.Commit); err != nil {
				return fmt.Errorf("error creating commit: %v", grpcutil.ScrubGRPC(err))
			}
		case op.Op1_7 != nil && op.Op1_7.Branch != nil:
			if _, err := pachClient.PfsAPIClient.SetBranch(ctx, op.Op1_7.Branch); err != nil {
				return fmt.Errorf("error creating branch: %v", grpcutil.ScrubGRPC(err))
			}
		case op.Op1_7 != nil && op.Op1_7.Pipeline != nil:
			if _, err := pachClient.PpsAPIClient.CreatePipeline(ctx, op.Op1_7.Pipeline); err != nil {
				return fmt.Errorf("error creating pipeline: %v", grpcutil.ScrubGRPC(err))
			}
		}
	}
}

func (a *apiServer) getPachClient() *client.APIClient {
	a.pachClientOnce.Do(func() {
		var err error
		a.pachClient, err = client.NewFromAddress(a.address)
		if err != nil {
			panic(fmt.Sprintf("pps failed to initialize pach client: %v", err))
		}
	})
	return a.pachClient
}

type extractObjectWriter struct {
	admin.API_ExtractServer
}

func (w *extractObjectWriter) Write(p []byte) (int, error) {
	chunkSize := grpcutil.MaxMsgSize / 2
	var n int
	for i := 0; i*(chunkSize) < len(p); i++ {
		value := p[i*chunkSize:]
		if len(value) > chunkSize {
			value = value[:chunkSize]
		}
		if err := w.Send(&admin.Op{Op1_7: &admin.Op1_7{Object: &pfs.PutObjectRequest{Value: value}}}); err != nil {
			return n, err
		}
		n += len(value)
	}
	return n, nil
}

type adminAPIRestoreServer admin.API_RestoreServer

type extractObjectReader struct {
	adminAPIRestoreServer
	buf bytes.Buffer
	eof bool
}

func (r *extractObjectReader) Read(p []byte) (int, error) {
	for len(p) > r.buf.Len() && !r.eof {
		request, err := r.Recv()
		if err != nil {
			return 0, grpcutil.ScrubGRPC(err)
		}
		op := request.Op
		if op.Op1_7.Object == nil {
			return 0, fmt.Errorf("expected an object, but got: %v", op)
		}
		r.buf.Write(op.Op1_7.Object.Value)
		if len(op.Op1_7.Object.Value) == 0 {
			r.eof = true
		}
	}
	return r.buf.Read(p)
}
<commit_msg>Extract commits / branches last.<commit_after>package server

import (
	"bytes"
	"fmt"
	"io"
	"sync"

	"github.com/golang/snappy"

	"github.com/gogo/protobuf/types"
	"github.com/pachyderm/pachyderm/src/client"
	"github.com/pachyderm/pachyderm/src/client/admin"
	"github.com/pachyderm/pachyderm/src/client/pfs"
	"github.com/pachyderm/pachyderm/src/client/pkg/grpcutil"
	"github.com/pachyderm/pachyderm/src/client/pkg/pbutil"
	"github.com/pachyderm/pachyderm/src/client/pps"
	"github.com/pachyderm/pachyderm/src/server/pkg/log"
	"github.com/pachyderm/pachyderm/src/server/pkg/obj"
)

type apiServer struct {
	log.Logger
	address        string
	pachClient     *client.APIClient
	pachClientOnce sync.Once
}

func (a *apiServer) Extract(request *admin.ExtractRequest, extractServer admin.API_ExtractServer) (retErr error) {
	pachClient := a.getPachClient()
	pachClient = pachClient.WithCtx(extractServer.Context())
	handleOp := extractServer.Send
	if request.URL != "" {
		url, err := obj.ParseURL(request.URL)
		if err != nil {
			return fmt.Errorf("error parsing url %v: %v", request.URL, err)
		}
		objClient, err := obj.NewClientFromURLAndSecret(extractServer.Context(), url)
		if err != nil {
			return err
		}
		objW, err := objClient.Writer(url.Object)
		if err != nil {
			return err
		}
		snappyW := snappy.NewBufferedWriter(objW)
		defer func() {
			if err := snappyW.Close(); err != nil && retErr == nil {
				retErr = err
			}
		}()
		w := pbutil.NewWriter(snappyW)
		handleOp = func(op *admin.Op) error { return w.Write(op) }
	}
	if !request.NoObjects {
		w := &extractObjectWriter{extractServer}
		if err := pachClient.ListObject(func(object *pfs.Object) error {
			if err := pachClient.GetObject(object.Hash, w); err != nil {
				return err
			}
			// empty PutObjectRequest to indicate EOF
			return handleOp(&admin.Op{Op1_7: &admin.Op1_7{Object: &pfs.PutObjectRequest{}}})
		}); err != nil {
			return err
		}
		if err := pachClient.ListTag(func(resp *pfs.ListTagsResponse) error {
			return handleOp(&admin.Op{Op1_7: &admin.Op1_7{
				Tag: &pfs.TagObjectRequest{
					Object: resp.Object,
					Tags:   []*pfs.Tag{resp.Tag},
				},
			}})
		}); err != nil {
			return err
		}
	}
	var repos []*pfs.Repo
	if !request.NoRepos {
		ris, err := pachClient.ListRepo(nil)
		if err != nil {
			return err
		}
		for _, ri := range ris {
			if len(ri.Provenance) > 0 {
				continue
			}
			if err := handleOp(&admin.Op{Op1_7: &admin.Op1_7{
				Repo: &pfs.CreateRepoRequest{
					Repo:        ri.Repo,
					Provenance:  ri.Provenance,
					Description: ri.Description,
				}},
			}); err != nil {
				return err
			}
			repos = append(repos, ri.Repo)
		}
	}
	if !request.NoPipelines {
		pis, err := pachClient.ListPipeline()
		if err != nil {
			return err
		}
		for _, pi := range pis {
			if err := handleOp(&admin.Op{Op1_7: &admin.Op1_7{
				Pipeline: &pps.CreatePipelineRequest{
					Pipeline:           pi.Pipeline,
					Transform:          pi.Transform,
					ParallelismSpec:    pi.ParallelismSpec,
					Egress:             pi.Egress,
					OutputBranch:       pi.OutputBranch,
					ScaleDownThreshold: pi.ScaleDownThreshold,
					ResourceRequests:   pi.ResourceRequests,
					ResourceLimits:     pi.ResourceLimits,
					Input:              pi.Input,
					Description:        pi.Description,
					Incremental:        pi.Incremental,
					CacheSize:          pi.CacheSize,
					EnableStats:        pi.EnableStats,
					Batch:              pi.Batch,
					MaxQueueSize:       pi.MaxQueueSize,
					Service:            pi.Service,
					ChunkSpec:          pi.ChunkSpec,
					DatumTimeout:       pi.DatumTimeout,
					JobTimeout:         pi.JobTimeout,
					Salt:               pi.Salt,
				}}}); err != nil {
				return err
			}
		}
	}
	// We send the actual commits last, that way pipelines will have already
	// been created and will recreate output commits for historical outputs.
	for _, repo := range repos {
		cis, err := pachClient.ListCommit(repo.Name, "", "", 0)
		if err != nil {
			return err
		}
		bis, err := pachClient.ListBranch(repo.Name)
		if err != nil {
			return err
		}
		for _, bcr := range buildCommitRequests(cis, bis) {
			if err := handleOp(&admin.Op{Op1_7: &admin.Op1_7{Commit: bcr}}); err != nil {
				return err
			}
		}
		for _, bi := range bis {
			if err := handleOp(&admin.Op{Op1_7: &admin.Op1_7{
				Branch: &pfs.SetBranchRequest{
					Commit: bi.Head,
					Branch: bi.Name,
				},
			}}); err != nil {
				return err
			}
		}
	}
	return nil
}

func buildCommitRequests(cis []*pfs.CommitInfo, bis []*pfs.BranchInfo) []*pfs.BuildCommitRequest {
	cis = sortCommitInfos(cis)
	result := make([]*pfs.BuildCommitRequest, len(cis))
	commitToBranch := make(map[string]string)
	for _, bi := range bis {
		if _, ok := commitToBranch[bi.Head.ID]; !ok || bi.Name == "master" {
			commitToBranch[bi.Head.ID] = bi.Name
		}
	}
	for i := range cis {
		ci := cis[len(cis)-i-1]
		branch := commitToBranch[ci.Commit.ID]
		// Even without a parent, ParentCommit is used to indicate which
		// repo to make the commit in.
		if ci.ParentCommit == nil {
			ci.ParentCommit = client.NewCommit(ci.Commit.Repo.Name, "")
		}
		result[len(cis)-i-1] = &pfs.BuildCommitRequest{
			Parent: ci.ParentCommit,
			Tree:   ci.Tree,
			ID:     ci.Commit.ID,
			Branch: branch,
		}
		if _, ok := commitToBranch[ci.ParentCommit.ID]; !ok || branch == "master" {
			commitToBranch[ci.ParentCommit.ID] = branch
		}
	}
	return result
}

func sortCommitInfos(cis []*pfs.CommitInfo) []*pfs.CommitInfo {
	commitMap := make(map[string]*pfs.CommitInfo)
	for _, ci := range cis {
		commitMap[ci.Commit.ID] = ci
	}
	var result []*pfs.CommitInfo
	for _, ci := range cis {
		if commitMap[ci.Commit.ID] == nil {
			continue
		}
		var localResult []*pfs.CommitInfo
		for ci != nil {
			localResult = append(localResult, ci)
			delete(commitMap, ci.Commit.ID)
			if ci.ParentCommit != nil {
				ci = commitMap[ci.ParentCommit.ID]
			} else {
				ci = nil
			}
		}
		for i := range localResult {
			result = append(result, localResult[len(localResult)-i-1])
		}
	}
	return result
}

func (a *apiServer) Restore(restoreServer admin.API_RestoreServer) (retErr error) {
	ctx := restoreServer.Context()
	pachClient := a.getPachClient()
	defer func() {
		for {
			_, err := restoreServer.Recv()
			if err != nil {
				break
			}
		}
		if err := restoreServer.SendAndClose(&types.Empty{}); err != nil && retErr == nil {
			retErr = err
		}
	}()
	var r pbutil.Reader
	for {
		var op *admin.Op
		if r == nil {
			req, err := restoreServer.Recv()
			if err != nil {
				if err == io.EOF {
					return nil
				}
				return err
			}
			if req.URL != "" {
				url, err := obj.ParseURL(req.URL)
				if err != nil {
					return fmt.Errorf("error parsing url %v: %v", req.URL, err)
				}
				objClient, err := obj.NewClientFromURLAndSecret(restoreServer.Context(), url)
				if err != nil {
					return err
				}
				objR, err := objClient.Reader(url.Object, 0, 0)
				if err != nil {
					return err
				}
				snappyR := snappy.NewReader(objR)
				r = pbutil.NewReader(snappyR)
				continue
			} else {
				op = req.Op
			}
		} else {
			op = &admin.Op{}
			if err := r.Read(op); err != nil {
				if err == io.EOF {
					return nil
				}
				return err
			}
		}
		switch {
		case op.Op1_7 != nil && op.Op1_7.Object != nil:
			r := &extractObjectReader{adminAPIRestoreServer: restoreServer}
			r.buf.Write(op.Op1_7.Object.Value)
			if _, _, err := pachClient.PutObject(r); err != nil {
				return fmt.Errorf("error putting object: %v", err)
			}
		case op.Op1_7 != nil && op.Op1_7.Tag != nil:
			if _, err := pachClient.ObjectAPIClient.TagObject(ctx, op.Op1_7.Tag); err != nil {
				return fmt.Errorf("error tagging object: %v", grpcutil.ScrubGRPC(err))
			}
		case op.Op1_7 != nil && op.Op1_7.Repo != nil:
			if _, err := pachClient.PfsAPIClient.CreateRepo(ctx, op.Op1_7.Repo); err != nil {
				return fmt.Errorf("error creating repo: %v", grpcutil.ScrubGRPC(err))
			}
		case op.Op1_7 != nil && op.Op1_7.Commit != nil:
			if _, err := pachClient.PfsAPIClient.BuildCommit(ctx, op.Op1_7.Commit); err != nil {
				return fmt.Errorf("error creating commit: %v", grpcutil.ScrubGRPC(err))
			}
		case op.Op1_7 != nil && op.Op1_7.Branch != nil:
			if _, err := pachClient.PfsAPIClient.SetBranch(ctx, op.Op1_7.Branch); err != nil {
				return fmt.Errorf("error creating branch: %v", grpcutil.ScrubGRPC(err))
			}
		case op.Op1_7 != nil && op.Op1_7.Pipeline != nil:
			if _, err := pachClient.PpsAPIClient.CreatePipeline(ctx, op.Op1_7.Pipeline); err != nil {
				return fmt.Errorf("error creating pipeline: %v", grpcutil.ScrubGRPC(err))
			}
		}
	}
}

func (a *apiServer) getPachClient() *client.APIClient {
	a.pachClientOnce.Do(func() {
		var err error
		a.pachClient, err = client.NewFromAddress(a.address)
		if err != nil {
			panic(fmt.Sprintf("pps failed to initialize pach client: %v", err))
		}
	})
	return a.pachClient
}

type extractObjectWriter struct {
	admin.API_ExtractServer
}

func (w *extractObjectWriter) Write(p []byte) (int, error) {
	chunkSize := grpcutil.MaxMsgSize / 2
	var n int
	for i := 0; i*(chunkSize) < len(p); i++ {
		value := p[i*chunkSize:]
		if len(value) > chunkSize {
			value = value[:chunkSize]
		}
		if err := w.Send(&admin.Op{Op1_7: &admin.Op1_7{Object: &pfs.PutObjectRequest{Value: value}}}); err != nil {
			return n, err
		}
		n += len(value)
	}
	return n, nil
}

type adminAPIRestoreServer admin.API_RestoreServer

type extractObjectReader struct {
	adminAPIRestoreServer
	buf bytes.Buffer
	eof bool
}

func (r *extractObjectReader) Read(p []byte) (int, error) {
	for len(p) > r.buf.Len() && !r.eof {
		request, err := r.Recv()
		if err != nil {
			return 0, grpcutil.ScrubGRPC(err)
		}
		op := request.Op
		if op.Op1_7.Object == nil {
			return 0, fmt.Errorf("expected an object, but got: %v", op)
		}
		r.buf.Write(op.Op1_7.Object.Value)
		if len(op.Op1_7.Object.Value) == 0 {
			r.eof = true
		}
	}
	return r.buf.Read(p)
}
<|endoftext|>