<commit_before>/* Search functions */

/*
 * Copyright (c) 2013-2016, Jeremy Bingham (<jeremy@goiardi.gl>)
 *
 * 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 main

import (
	"encoding/json"
	"fmt"
	"github.com/ctdk/goiardi/actor"
	"github.com/ctdk/goiardi/client"
	"github.com/ctdk/goiardi/config"
	"github.com/ctdk/goiardi/databag"
	"github.com/ctdk/goiardi/environment"
	"github.com/ctdk/goiardi/indexer"
	"github.com/ctdk/goiardi/node"
	"github.com/ctdk/goiardi/role"
	"github.com/ctdk/goiardi/search"
	"github.com/ctdk/goiardi/util"
	"github.com/tideland/golib/logger"
	"net/http"
	"net/url"
	"regexp"
	"strconv"
	"sync"
)

var riM *sync.Mutex

func init() {
	riM = new(sync.Mutex)
}

func searchHandler(w http.ResponseWriter, r *http.Request) {
	/* ... and we need search to run the environment tests, so here we
	 * go. */
	w.Header().Set("Content-Type", "application/json")
	searchResponse := make(map[string]interface{})
	pathArray := splitPath(r.URL.Path)
	pathArrayLen := len(pathArray)

	opUser, oerr := actor.GetReqUser(r.Header.Get("X-OPS-USERID"))
	if oerr != nil {
		jsonErrorReport(w, r, oerr.Error(), oerr.Status())
		return
	}

	/* set up query params for searching */
	var (
		paramQuery string
		paramsRows int
		sortOrder  string
		start      int
	)
	r.ParseForm()
	if q, found := r.Form["q"]; found {
		if len(q) < 0 {
			jsonErrorReport(w, r, "No query string specified for search", http.StatusBadRequest)
			return
		}
		paramQuery = q[0]
	} else if pathArrayLen != 1 {
		/* default to "*:*" for a search term */
		paramQuery = "*:*"
	}
	if pr, found := r.Form["rows"]; found {
		if len(pr) > 0 {
			paramsRows, _ = strconv.Atoi(pr[0])
		}
	} else {
		paramsRows = 1000
	}
	sortOrder = "id ASC"
	if s, found := r.Form["sort"]; found {
		if len(s) > 0 {
			if s[0] != "" {
				sortOrder = s[0]
			}
		} else {
			sortOrder = "id ASC"
		}
	}
	if st, found := r.Form["start"]; found {
		if len(st) > 0 {
			start, _ = strconv.Atoi(st[0])
		}
	} else {
		start = 0
	}

	var searcher search.Searcher
	if config.Config.PgSearch {
		searcher = &search.PostgresSearch{}
	} else {
		searcher = &search.TrieSearch{}
	}

	if pathArrayLen == 1 {
		/* base end points */
		switch r.Method {
		case "GET":
			if opUser.IsValidator() {
				jsonErrorReport(w, r, "You are not allowed to perform this action", http.StatusForbidden)
				return
			}
			searchEndpoints := searcher.GetEndpoints()
			for _, s := range searchEndpoints {
				searchResponse[s] = util.CustomURL(fmt.Sprintf("/search/%s", s))
			}
		default:
			jsonErrorReport(w, r, "Method not allowed", http.StatusMethodNotAllowed)
			return
		}
	} else if pathArrayLen == 2 {
		switch r.Method {
		case "GET", "POST":
			if opUser.IsValidator() {
				jsonErrorReport(w, r, "You are not allowed to perform this action", http.StatusForbidden)
				return
			}
			var qerr error
			paramQuery, qerr = url.QueryUnescape(paramQuery)
			if qerr != nil {
				jsonErrorReport(w, r, qerr.Error(), http.StatusBadRequest)
				return
			}
			/* start figuring out what comes in POSTS now,
			 * so the partial search tests don't complain
			 * anymore. */
			var partialData map[string]interface{}
			if r.Method == "POST" {
				var perr error
				partialData, perr = parseObjJSON(r.Body)
				if perr != nil {
					jsonErrorReport(w, r, perr.Error(), http.StatusBadRequest)
					return
				}
			}

			idx := pathArray[1]
			res, err := searcher.Search(idx, paramQuery, paramsRows, sortOrder, start, partialData)

			if err != nil {
				statusCode := http.StatusBadRequest
				re := regexp.MustCompile(`^I don't know how to search for .*? data objects.`)
				if re.MatchString(err.Error()) {
					statusCode = http.StatusNotFound
				}
				jsonErrorReport(w, r, err.Error(), statusCode)
				return
			}

			searchResponse["total"] = len(res)
			searchResponse["start"] = start
			searchResponse["rows"] = res
		default:
			jsonErrorReport(w, r, "Method not allowed", http.StatusMethodNotAllowed)
			return
		}
	} else {
		/* Say what? Bad request. */
		jsonErrorReport(w, r, "Bad request", http.StatusBadRequest)
		return
	}

	enc := json.NewEncoder(w)
	if err := enc.Encode(&searchResponse); err != nil {
		jsonErrorReport(w, r, err.Error(), http.StatusInternalServerError)
	}
}

func reindexHandler(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/json")
	reindexResponse := make(map[string]interface{})
	opUser, oerr := actor.GetReqUser(r.Header.Get("X-OPS-USERID"))
	if oerr != nil {
		jsonErrorReport(w, r, oerr.Error(), oerr.Status())
		return
	}
	switch r.Method {
	case "POST":
		if !opUser.IsAdmin() {
			jsonErrorReport(w, r, "You are not allowed to perform that action.", http.StatusForbidden)
			return
		}
		reindexAll()
		reindexResponse["reindex"] = "OK"
	default:
		jsonErrorReport(w, r, "Method not allowed. If you're trying to do something with a data bag named 'reindex', it's not going to work I'm afraid.", http.StatusMethodNotAllowed)
		return
	}
	enc := json.NewEncoder(w)
	if err := enc.Encode(&reindexResponse); err != nil {
		jsonErrorReport(w, r, err.Error(), http.StatusInternalServerError)
	}
}

func reindexAll() {
	// Take the mutex before starting to reindex everything. This way at
	// least reindexing jobs won't pile up on top of each other all trying
	// to execute simultaneously.
	riM.Lock()
	defer riM.Unlock()

	reindexObjs := make([]indexer.Indexable, 0, 100)
	// We clear the index, *then* do the fetch because if
	// something comes in between the time we fetch the
	// objects to reindex and when it gets done, they'll
	// just be added naturally
	indexer.ClearIndex()

	for _, v := range client.AllClients() {
		reindexObjs = append(reindexObjs, v)
	}
	for _, v := range node.AllNodes() {
		reindexObjs = append(reindexObjs, v)
	}
	for _, v := range role.AllRoles() {
		reindexObjs = append(reindexObjs, v)
	}
	for _, v := range environment.AllEnvironments() {
		reindexObjs = append(reindexObjs, v)
	}
	defaultEnv, _ := environment.Get("_default")
	reindexObjs = append(reindexObjs, defaultEnv)
	// data bags have to be done separately
	dbags := databag.GetList()
	for _, db := range dbags {
		dbag, err := databag.Get(db)
		if err != nil {
			continue
		}
		dbis := make([]indexer.Indexable, dbag.NumDBItems())
		i := 0
		allDBItems, derr := dbag.AllDBItems()
		if derr != nil {
			logger.Errorf(derr.Error())
			continue
		}
		for _, k := range allDBItems {
			n := k
			dbis[i] = n
			i++
		}
		reindexObjs = append(reindexObjs, dbis...)
	}
	indexer.ReIndex(reindexObjs)
	return
}
<commit_msg>break up reindexing into smaller chunks to reindex at a time<commit_after>/* Search functions */

/*
 * Copyright (c) 2013-2016, Jeremy Bingham (<jeremy@goiardi.gl>)
 *
 * 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 main

import (
	"encoding/json"
	"fmt"
	"github.com/ctdk/goiardi/actor"
	"github.com/ctdk/goiardi/client"
	"github.com/ctdk/goiardi/config"
	"github.com/ctdk/goiardi/databag"
	"github.com/ctdk/goiardi/environment"
	"github.com/ctdk/goiardi/indexer"
	"github.com/ctdk/goiardi/node"
	"github.com/ctdk/goiardi/role"
	"github.com/ctdk/goiardi/search"
	"github.com/ctdk/goiardi/util"
	"github.com/tideland/golib/logger"
	"net/http"
	"net/url"
	"regexp"
	"strconv"
	"sync"
)

var riM *sync.Mutex

func init() {
	riM = new(sync.Mutex)
}

func searchHandler(w http.ResponseWriter, r *http.Request) {
	/* ... and we need search to run the environment tests, so here we
	 * go. */
	w.Header().Set("Content-Type", "application/json")
	searchResponse := make(map[string]interface{})
	pathArray := splitPath(r.URL.Path)
	pathArrayLen := len(pathArray)

	opUser, oerr := actor.GetReqUser(r.Header.Get("X-OPS-USERID"))
	if oerr != nil {
		jsonErrorReport(w, r, oerr.Error(), oerr.Status())
		return
	}

	/* set up query params for searching */
	var (
		paramQuery string
		paramsRows int
		sortOrder  string
		start      int
	)
	r.ParseForm()
	if q, found := r.Form["q"]; found {
		if len(q) < 0 {
			jsonErrorReport(w, r, "No query string specified for search", http.StatusBadRequest)
			return
		}
		paramQuery = q[0]
	} else if pathArrayLen != 1 {
		/* default to "*:*" for a search term */
		paramQuery = "*:*"
	}
	if pr, found := r.Form["rows"]; found {
		if len(pr) > 0 {
			paramsRows, _ = strconv.Atoi(pr[0])
		}
	} else {
		paramsRows = 1000
	}
	sortOrder = "id ASC"
	if s, found := r.Form["sort"]; found {
		if len(s) > 0 {
			if s[0] != "" {
				sortOrder = s[0]
			}
		} else {
			sortOrder = "id ASC"
		}
	}
	if st, found := r.Form["start"]; found {
		if len(st) > 0 {
			start, _ = strconv.Atoi(st[0])
		}
	} else {
		start = 0
	}

	var searcher search.Searcher
	if config.Config.PgSearch {
		searcher = &search.PostgresSearch{}
	} else {
		searcher = &search.TrieSearch{}
	}

	if pathArrayLen == 1 {
		/* base end points */
		switch r.Method {
		case "GET":
			if opUser.IsValidator() {
				jsonErrorReport(w, r, "You are not allowed to perform this action", http.StatusForbidden)
				return
			}
			searchEndpoints := searcher.GetEndpoints()
			for _, s := range searchEndpoints {
				searchResponse[s] = util.CustomURL(fmt.Sprintf("/search/%s", s))
			}
		default:
			jsonErrorReport(w, r, "Method not allowed", http.StatusMethodNotAllowed)
			return
		}
	} else if pathArrayLen == 2 {
		switch r.Method {
		case "GET", "POST":
			if opUser.IsValidator() {
				jsonErrorReport(w, r, "You are not allowed to perform this action", http.StatusForbidden)
				return
			}
			var qerr error
			paramQuery, qerr = url.QueryUnescape(paramQuery)
			if qerr != nil {
				jsonErrorReport(w, r, qerr.Error(), http.StatusBadRequest)
				return
			}
			/* start figuring out what comes in POSTS now,
			 * so the partial search tests don't complain
			 * anymore. */
			var partialData map[string]interface{}
			if r.Method == "POST" {
				var perr error
				partialData, perr = parseObjJSON(r.Body)
				if perr != nil {
					jsonErrorReport(w, r, perr.Error(), http.StatusBadRequest)
					return
				}
			}

			idx := pathArray[1]
			res, err := searcher.Search(idx, paramQuery, paramsRows, sortOrder, start, partialData)

			if err != nil {
				statusCode := http.StatusBadRequest
				re := regexp.MustCompile(`^I don't know how to search for .*? data objects.`)
				if re.MatchString(err.Error()) {
					statusCode = http.StatusNotFound
				}
				jsonErrorReport(w, r, err.Error(), statusCode)
				return
			}

			searchResponse["total"] = len(res)
			searchResponse["start"] = start
			searchResponse["rows"] = res
		default:
			jsonErrorReport(w, r, "Method not allowed", http.StatusMethodNotAllowed)
			return
		}
	} else {
		/* Say what? Bad request. */
		jsonErrorReport(w, r, "Bad request", http.StatusBadRequest)
		return
	}

	enc := json.NewEncoder(w)
	if err := enc.Encode(&searchResponse); err != nil {
		jsonErrorReport(w, r, err.Error(), http.StatusInternalServerError)
	}
}

func reindexHandler(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/json")
	reindexResponse := make(map[string]interface{})
	opUser, oerr := actor.GetReqUser(r.Header.Get("X-OPS-USERID"))
	if oerr != nil {
		jsonErrorReport(w, r, oerr.Error(), oerr.Status())
		return
	}
	switch r.Method {
	case "POST":
		if !opUser.IsAdmin() {
			jsonErrorReport(w, r, "You are not allowed to perform that action.", http.StatusForbidden)
			return
		}
		reindexAll()
		reindexResponse["reindex"] = "OK"
	default:
		jsonErrorReport(w, r, "Method not allowed. If you're trying to do something with a data bag named 'reindex', it's not going to work I'm afraid.", http.StatusMethodNotAllowed)
		return
	}
	enc := json.NewEncoder(w)
	if err := enc.Encode(&reindexResponse); err != nil {
		jsonErrorReport(w, r, err.Error(), http.StatusInternalServerError)
	}
}

func reindexAll() {
	// Take the mutex before starting to reindex everything. This way at
	// least reindexing jobs won't pile up on top of each other all trying
	// to execute simultaneously.
	riM.Lock()
	defer riM.Unlock()

	reindexObjs := make([]indexer.Indexable, 0, 100)
	// We clear the index, *then* do the fetch because if
	// something comes in between the time we fetch the
	// objects to reindex and when it gets done, they'll
	// just be added naturally
	indexer.ClearIndex()

	// Send the objects to be reindexed in somewhat more manageable chunks
	clientObjs := make([]indexer.Indexable, 0, 100)
	for _, v := range client.AllClients() {
		clientObjs = append(clientObjs, v)
	}
	indexer.ReIndex(clientObjs)

	nodeObjs := make([]indexer.Indexable, 0, 100)
	for _, v := range node.AllNodes() {
		nodeObjs = append(nodeObjs, v)
	}
	indexer.ReIndex(nodeObjs)

	roleObjs := make([]indexer.Indexable, 0, 100)
	for _, v := range role.AllRoles() {
		roleObjs = append(roleObjs, v)
	}
	indexer.ReIndex(roleObjs)

	environmentObjs := make([]indexer.Indexable, 0, 100)
	for _, v := range environment.AllEnvironments() {
		environmentObjs = append(environmentObjs, v)
	}
	defaultEnv, _ := environment.Get("_default")
	environmentObjs = append(environmentObjs, defaultEnv)
	indexer.ReIndex(environmentObjs)

	dbagObjs := make([]indexer.Indexable, 0, 100)
	// data bags have to be done separately
	dbags := databag.GetList()
	for _, db := range dbags {
		dbag, err := databag.Get(db)
		if err != nil {
			continue
		}
		dbis := make([]indexer.Indexable, dbag.NumDBItems())
		i := 0
		allDBItems, derr := dbag.AllDBItems()
		if derr != nil {
			logger.Errorf(derr.Error())
			continue
		}
		for _, k := range allDBItems {
			n := k
			dbis[i] = n
			i++
		}
		dbagObjs = append(dbagObjs, dbis...)
	}
	indexer.ReIndex(dbagObjs)
	return
}
<|endoftext|>


<commit_before>package elastic

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"log"
)

type Dict map[string]interface{}

// fields of a Search API call
const (
	// APIs
	EXPLAIN  = "explain"
	VALIDATE = "validate"
	SEARCH   = "search"
	// Query elements
	ALL     = "_all"
	INCLUDE = "include_in_all"
	SOURCE  = "_source"
	// url params
	SEARCH_TYPE = "search_type"
	SCROLL      = "scroll"
)

/*
 * a request representing a search
 */
type Search struct {
	url    string
	params map[string]string
	query  Dict
}

type Query interface {
	Name() string
	KV() Dict
}

/*
 * General purpose query
 */
type Object struct {
	name string
	kv   Dict
}

func (this *Object) Name() string {
	return this.name
}

func (this *Object) KV() Dict {
	return this.kv
}

/*
 * Create a new query object
 */
func NewQuery(name string) *Object {
	return &Object{name: name, kv: make(Dict)}
}

/*
 * Get a string representation of this object
 */
func (this *Object) String() string {
	return String(this.KV())
}

/*
 * Create an Explain request, that will return explanation for why a document is returned by the query
 */
func (this *Elasticsearch) Explain(index, class string, id int64) *Search {
	var url string = this.request(index, class, id, EXPLAIN)
	return &Search{url: url, query: make(Dict)}
}

/*
 * Create a Validate request
 */
func (this *Elasticsearch) Validate(index, class string, explain bool) *Search {
	var url string = this.request(index, class, -1, VALIDATE) + "/query"
	if explain {
		url += "?" + EXPLAIN
	}
	return &Search{url: url, query: make(Dict)}
}

/*
 * Create a Search request
 */
func (this *Elasticsearch) Search(index, class string) *Search {
	var url string = this.request(index, class, -1, SEARCH)
	return &Search{url: url, params: make(map[string]string), query: make(Dict)}
}

/*
 * Create a new Search API call
 */
func newSearch() *Search {
	return &Search{url: "", params: make(map[string]string), query: make(Dict)}
}

/*
 * Add a url parameter/value, e.g. search_type (count, query_and_fetch, dfs_query_then_fetch/dfs_query_and_fetch, scan)
 */
func (this *Search) AddParam(name, value string) *Search {
	this.params[name] = value
	return this
}

/*
* Add a query to this search request
 */
func (this *Search) AddQuery(query Query) *Search {
	this.query[query.Name()] = query.KV()
	return this
}

/*
 * Add to _source (i.e. specify another field that should be extracted)
 */
func (this *Search) AddSource(source string) *Search {
	var sources []string
	if this.query[SOURCE] == nil {
		sources = []string{}
	} else {
		sources = this.query[SOURCE].([]string)
	}
	sources = append(sources, source)
	this.query[SOURCE] = sources
	return this
}

/*
 * Add a query argument/value, e.g. size, from, etc.
 */
func (this *Search) Add(argument string, value interface{}) *Search {
	this.query[argument] = value
	return this
}

/*
 * Get a string representation of this Search API call
 */
func (this *Search) String() string {
	body := ""
	if len(this.query) > 0 {
		body = String(this.query)
	}
	return body
}

/*
 * Construct the url of this Search API call
 */
func (this *Search) urlString() string {
	return urlString(this.url, this.params)
}

/*
 * request mappings between the json fields and how Elasticsearch store them
 * GET /:index/:type/_search
 */
func (this *Search) Get() {
	// construct the url
	url := this.urlString()
	// construct the body
	body := this.String()
	var data io.Reader = nil
	if body != "" {
		data = bytes.NewReader([]byte(body))
	}
	// submit the request
	log.Println("GET", url)
	log.Println("query", body)
	reader, err := exec("GET", url, data)
	if err != nil {
		log.Println(err)
		return
	}
	if data, err := ioutil.ReadAll(reader); err == nil {
		fmt.Println(string(data))
	}
}

/*
 * Add a query argument/value
 */
func (this *Object) Add(argument string, value interface{}) *Object {
	this.kv[argument] = value
	return this
}

/*
 * specify multiple values to match
 */
func (this *Object) AddMultiple(argument string, values ...interface{}) *Object {
	this.kv[argument] = values
	return this
}

/*
 * Add a sub query (e.g. a field query)
 */
func (this *Object) AddQuery(query Query) *Object {
	this.kv[query.Name()] = query.KV()
	return this
}

/*
 * Boolean clause, it is a complex clause that allows to combine other clauses as 'must' match, 'must_not' match, 'should' match.
 */
type Bool struct {
	name string
	kv   Dict
}

func (this *Bool) Name() string {
	return this.name
}
func (this *Bool) KV() Dict {
	return this.kv
}

/*
 * Create a new 'bool' clause
 */
func NewBool() *Bool {
	kv := make(Dict)
	return &Bool{name: "bool", kv: kv}
}

/*
 * Add a 'must' clause to this 'bool' clause
 */
func (this *Bool) AddMust(query Query) *Bool {
	this.add("must", query)
	return this
}

/*
 * Add a 'must_not' clause to this 'bool' clause
 */
func (this *Bool) AddMustNot(query Query) *Bool {
	this.add("must_not", query)
	return this
}

/*
 * Add a 'should' clause to this 'bool' clause
 */
func (this *Bool) AddShould(query Query) *Bool {
	this.add("should", query)
	return this
}

/*
 * add a clause
 */
func (this *Bool) add(key string, query Query) {
	collection := this.kv[key]
	//TODO check if query.Name exists, otherwise transform the map to array
	if collection == nil {
		// at first the collection is a map
		collection = make(Dict)
		collection.(Dict)[query.Name()] = query.KV()
	} else {
		// when more items are added, then it becomes an array
		dict := make(Dict)
		dict[query.Name()] = query.KV()
		// check if it is a map
		if _, ok := collection.(Dict); ok {
			array := []Dict{} // transform previous map into array
			for k, v := range collection.(Dict) {
				d := make(Dict)
				d[k] = v
				array = append(array, d)
			}
			collection = array
		}
		collection = append(collection.([]Dict), dict)
	}
	this.kv[key] = collection
}

/*
 * Create a new 'terms' filter, it is like 'term' but can match multiple values
 */
func NewTerms() *Object {
	return NewQuery("terms")
}

/*
 * Create a new 'term' filter
 */
func NewTerm() *Object {
	return NewQuery("term")
}

/*
 * Create a new `exists` filter.
 */
func NewExists() *Object {
	return NewQuery("exists")
}

/*
 * Create a new `missing` filter (the inverse of `exists`)
 */
func NewMissing() *Object {
	return NewQuery("missing")
}
<commit_msg>can add attributes to bool queries<commit_after>package elastic

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"log"
)

type Dict map[string]interface{}

// fields of a Search API call
const (
	// APIs
	EXPLAIN  = "explain"
	VALIDATE = "validate"
	SEARCH   = "search"
	// Query elements
	ALL     = "_all"
	INCLUDE = "include_in_all"
	SOURCE  = "_source"
	// url params
	SEARCH_TYPE = "search_type"
	SCROLL      = "scroll"
)

/*
 * a request representing a search
 */
type Search struct {
	url    string
	params map[string]string
	query  Dict
}

type Query interface {
	Name() string
	KV() Dict
}

/*
 * General purpose query
 */
type Object struct {
	name string
	kv   Dict
}

func (this *Object) Name() string {
	return this.name
}

func (this *Object) KV() Dict {
	return this.kv
}

/*
 * Create a new query object
 */
func NewQuery(name string) *Object {
	return &Object{name: name, kv: make(Dict)}
}

/*
 * Get a string representation of this object
 */
func (this *Object) String() string {
	return String(this.KV())
}

/*
 * Create an Explain request, that will return explanation for why a document is returned by the query
 */
func (this *Elasticsearch) Explain(index, class string, id int64) *Search {
	var url string = this.request(index, class, id, EXPLAIN)
	return &Search{url: url, query: make(Dict)}
}

/*
 * Create a Validate request
 */
func (this *Elasticsearch) Validate(index, class string, explain bool) *Search {
	var url string = this.request(index, class, -1, VALIDATE) + "/query"
	if explain {
		url += "?" + EXPLAIN
	}
	return &Search{url: url, query: make(Dict)}
}

/*
 * Create a Search request
 */
func (this *Elasticsearch) Search(index, class string) *Search {
	var url string = this.request(index, class, -1, SEARCH)
	return &Search{url: url, params: make(map[string]string), query: make(Dict)}
}

/*
 * Create a new Search API call
 */
func newSearch() *Search {
	return &Search{url: "", params: make(map[string]string), query: make(Dict)}
}

/*
 * Add a url parameter/value, e.g. search_type (count, query_and_fetch, dfs_query_then_fetch/dfs_query_and_fetch, scan)
 */
func (this *Search) AddParam(name, value string) *Search {
	this.params[name] = value
	return this
}

/*
 * Pretiffy the response result
 */
func (this *Search) Pretty() *Search {
	this.AddParam("pretty", "")
	return this
}

/*
* Add a query to this search request
 */
func (this *Search) AddQuery(query Query) *Search {
	this.query[query.Name()] = query.KV()
	return this
}

/*
 * Add to _source (i.e. specify another field that should be extracted)
 */
func (this *Search) AddSource(source string) *Search {
	var sources []string
	if this.query[SOURCE] == nil {
		sources = []string{}
	} else {
		sources = this.query[SOURCE].([]string)
	}
	sources = append(sources, source)
	this.query[SOURCE] = sources
	return this
}

/*
 * Add a query argument/value, e.g. size, from, etc.
 */
func (this *Search) Add(argument string, value interface{}) *Search {
	this.query[argument] = value
	return this
}

/*
 * Get a string representation of this Search API call
 */
func (this *Search) String() string {
	body := ""
	if len(this.query) > 0 {
		body = String(this.query)
	}
	return body
}

/*
 * Construct the url of this Search API call
 */
func (this *Search) urlString() string {
	return urlString(this.url, this.params)
}

/*
 * request mappings between the json fields and how Elasticsearch store them
 * GET /:index/:type/_search
 */
func (this *Search) Get() {
	// construct the url
	url := this.urlString()
	// construct the body
	query := this.String()
	var body io.Reader
	if query != "" {
		body = bytes.NewReader([]byte(query))
	}
	// submit the request
	log.Println("GET", url, query)
	reader, err := exec("GET", url, body)
	if err != nil {
		log.Println(err)
		return
	}
	if data, err := ioutil.ReadAll(reader); err == nil {
		fmt.Println(string(data))
	}
}

/*
 * Add a query argument/value
 */
func (this *Object) Add(argument string, value interface{}) *Object {
	this.kv[argument] = value
	return this
}

/*
 * specify multiple values to match
 */
func (this *Object) AddMultiple(argument string, values ...interface{}) *Object {
	this.kv[argument] = values
	return this
}

/*
 * Add a sub query (e.g. a field query)
 */
func (this *Object) AddQuery(query Query) *Object {
	this.kv[query.Name()] = query.KV()
	return this
}

/*
 * Boolean clause, it is a complex clause that allows to combine other clauses as 'must' match, 'must_not' match, 'should' match.
 */
type Bool struct {
	name string
	kv   Dict
}

func (this *Bool) Name() string {
	return this.name
}
func (this *Bool) KV() Dict {
	return this.kv
}

/*
 * Create a new 'bool' clause
 */
func NewBool() *Bool {
	kv := make(Dict)
	return &Bool{name: "bool", kv: kv}
}

/*
 * Add a 'must' clause to this 'bool' clause
 */
func (this *Bool) AddMust(query Query) *Bool {
	this.add("must", query)
	return this
}

/*
 * Add a 'must_not' clause to this 'bool' clause
 */
func (this *Bool) AddMustNot(query Query) *Bool {
	this.add("must_not", query)
	return this
}

/*
 * Add a 'should' clause to this 'bool' clause
 */
func (this *Bool) AddShould(query Query) *Bool {
	this.add("should", query)
	return this
}

/*
 * Add a parameter to this `bool` query
 */
func (this *Bool) Add(name string, value interface{}) *Bool {
	this.kv[name] = value
	return this
}

/*
 * add a clause
 */
func (this *Bool) add(key string, query Query) {
	collection := this.kv[key]
	// check if query.Name exists, otherwise transform the map to array
	if collection == nil {
		// at first the collection is a map
		collection = make(Dict)
		collection.(Dict)[query.Name()] = query.KV()
	} else {
		// when more items are added, then it becomes an array
		dict := make(Dict)
		dict[query.Name()] = query.KV()
		// check if it is a map
		if _, ok := collection.(Dict); ok {
			array := []Dict{} // transform previous map into array
			for k, v := range collection.(Dict) {
				d := make(Dict)
				d[k] = v
				array = append(array, d)
			}
			collection = array
		}
		collection = append(collection.([]Dict), dict)
	}
	this.kv[key] = collection
}

/*
 * Create a new 'terms' filter, it is like 'term' but can match multiple values
 */
func NewTerms() *Object {
	return NewQuery("terms")
}

/*
 * Create a new 'term' filter
 */
func NewTerm() *Object {
	return NewQuery("term")
}

/*
 * Create a new `exists` filter.
 */
func NewExists() *Object {
	return NewQuery("exists")
}

/*
 * Create a new `missing` filter (the inverse of `exists`)
 */
func NewMissing() *Object {
	return NewQuery("missing")
}
<|endoftext|>


<commit_before>package sheet

import (
	"bufio"
	"errors"
	"fmt"
	"io"
	"os"
	"strconv"
	"strings"

	"scim/sheet/align"
)

var columnArr []string

func init() {
	columnArr = make([]string, 26)
	columnArr[0] = "A"
	columnArr[1] = "B"
	columnArr[2] = "C"
	columnArr[3] = "D"
	columnArr[4] = "E"
	columnArr[5] = "F"
	columnArr[6] = "G"
	columnArr[7] = "H"
	columnArr[8] = "I"
	columnArr[9] = "J"
	columnArr[10] = "K"
	columnArr[11] = "L"
	columnArr[12] = "M"
	columnArr[13] = "N"
	columnArr[14] = "O"
	columnArr[15] = "P"
	columnArr[16] = "Q"
	columnArr[17] = "R"
	columnArr[18] = "S"
	columnArr[19] = "T"
}

const (
	STARTING_COLUMN_WIDTH = 10
)

type ColumnFormat struct {
	width     int
	precision int
	ctype     int
}

type Sheet struct {
	Filename string

	SelectedCell  string
	columnFormats map[string]ColumnFormat
	data          map[string]*Cell

	// display window
	startRow, startCol       int
	displayRows, displayCols int
}

func NewSheet(filename string) Sheet {
	s := Sheet{Filename: filename, SelectedCell: "A0", columnFormats: make(map[string]ColumnFormat), data: make(map[string]*Cell)}

	// Load file
	if file, err := os.Open(filename); err == nil {
		scanner := bufio.NewScanner(file)
		for scanner.Scan() {
			line := scanner.Text()
			if strings.HasPrefix(line, "#") || len(line) < 1 {
				continue
			}
			words := strings.Split(line, " ")
			cmd := ""
			adrs := ""
			val := ""
			if len(words) >= 2 {
				cmd = words[0]
				adrs = words[1]
			}
			if len(words) >= 4 {
				val = strings.Join(words[3:], " ")
			}
			if len(val) > 1 && val[0] == '"' {
				val = val[1 : len(val)-1]
			}
			switch cmd {
			case "leftstring":
				s.SetCell(adrs, &Cell{stringType: true, alignment: align.AlignLeft, value: val})
			case "rightstrng":
				s.SetCell(adrs, &Cell{stringType: true, alignment: align.AlignRight, value: val})
			case "label":
				s.SetCell(adrs, &Cell{stringType: true, alignment: align.AlignCenter, value: val})
			case "let":
				s.SetCell(adrs, &Cell{stringType: false, alignment: align.AlignRight, value: val})
			case "goto":
				s.SelectedCell = adrs
			case "format":
				width, _ := strconv.ParseInt(words[2], 10, 64)
				precision, _ := strconv.ParseInt(words[3], 10, 64)
				ctype, _ := strconv.ParseInt(words[4], 10, 64)
				s.columnFormats[adrs] = ColumnFormat{width: int(width), precision: int(precision), ctype: int(ctype)}
			}
		}
	}

	return s
}

func (s *Sheet) writeFormats(w io.Writer) {
	for k, cFormat := range s.columnFormats {
		fmt.Fprintf(w, "format %s %d %d %d\n", k, cFormat.width, cFormat.precision, cFormat.ctype)
	}
}

func (s *Sheet) getPrecision(address string) int {
	column := address[1:]
	if cFormat, found := s.columnFormats[column]; found {
		return cFormat.precision
	} else {
		return 2
	}
}

func (s *Sheet) DisplayFormat(address string) string {
	column := address[1:]
	if cFormat, found := s.columnFormats[column]; found {
		return fmt.Sprintf("%d %d %d", cFormat.width, cFormat.precision, cFormat.ctype)
	} else {
		return fmt.Sprintf("%d %d %d", STARTING_COLUMN_WIDTH, 2, 0)
	}
}

func (s *Sheet) getColumnWidth(column string) int {
	if cFormat, found := s.columnFormats[column]; found {
		return cFormat.width
	} else {
		s.columnFormats[column] = ColumnFormat{width: STARTING_COLUMN_WIDTH, precision: 2, ctype: 0}
		return STARTING_COLUMN_WIDTH
	}
}

func (s *Sheet) increaseColumnWidth(column string) {
	if cFormat, found := s.columnFormats[column]; found {
		cFormat.width += 1
		s.columnFormats[column] = cFormat
	} else {
		s.columnFormats[column] = ColumnFormat{width: STARTING_COLUMN_WIDTH + 1, precision: 2, ctype: 0}
	}
}

func (s *Sheet) decreaseColumnWidth(column string) {
	if cFormat, found := s.columnFormats[column]; found {
		if cFormat.width > 1 {
			cFormat.width--
			s.columnFormats[column] = cFormat
		}
	} else {
		s.columnFormats[column] = ColumnFormat{width: STARTING_COLUMN_WIDTH - 1, precision: 2, ctype: 0}
	}
}

func (s *Sheet) ClearCell(address string) {
	delete(s.data, address)
}

func (s *Sheet) GetCell(address string) (*Cell, error) {
	if cell, found := s.data[address]; found {
		return cell, nil
	} else if address == s.SelectedCell {
		return &Cell{}, nil
	}
	return nil, errors.New("Cell does not exist in spreadsheet.")
}

func (s *Sheet) SetCell(address string, cell *Cell) {
	// TODO: more work here to set refs and calc disp value
	s.data[address] = cell
	s.display()
}

func (s *Sheet) Save() error {
	if outfile, err := os.Create(s.Filename); err == nil {
		fmt.Fprintln(outfile, "# This data file was generated by Spreadsheet Calculator IMproved.")
		fmt.Fprintln(outfile, "# You almost certainly shouldn't edit it.")
		fmt.Fprintln(outfile, "")

		s.writeFormats(outfile)

		for addr, cell := range s.data {
			cell.write(outfile, addr)
		}
		fmt.Fprintf(outfile, "goto %s A0", s.SelectedCell)
		outfile.Close()
		return nil
	} else {
		return err
	}
}
<commit_msg>Add forward and back references for each cell used in expressions.<commit_after>package sheet

import (
	"bufio"
	"errors"
	"fmt"
	"io"
	"os"
	"strconv"
	"strings"

	"scim/evaler"
	"scim/sheet/align"
)

var columnArr []string

func init() {
	columnArr = make([]string, 26)
	columnArr[0] = "A"
	columnArr[1] = "B"
	columnArr[2] = "C"
	columnArr[3] = "D"
	columnArr[4] = "E"
	columnArr[5] = "F"
	columnArr[6] = "G"
	columnArr[7] = "H"
	columnArr[8] = "I"
	columnArr[9] = "J"
	columnArr[10] = "K"
	columnArr[11] = "L"
	columnArr[12] = "M"
	columnArr[13] = "N"
	columnArr[14] = "O"
	columnArr[15] = "P"
	columnArr[16] = "Q"
	columnArr[17] = "R"
	columnArr[18] = "S"
	columnArr[19] = "T"
}

const (
	STARTING_COLUMN_WIDTH = 10
)

type ColumnFormat struct {
	width     int
	precision int
	ctype     int
}

type Sheet struct {
	Filename string

	SelectedCell  string
	columnFormats map[string]ColumnFormat
	data          map[string]*Cell

	// display window
	startRow, startCol       int
	displayRows, displayCols int
}

func NewSheet(filename string) Sheet {
	s := Sheet{Filename: filename, SelectedCell: "A0", columnFormats: make(map[string]ColumnFormat), data: make(map[string]*Cell)}

	// Load file
	if file, err := os.Open(filename); err == nil {
		scanner := bufio.NewScanner(file)
		for scanner.Scan() {
			line := scanner.Text()
			if strings.HasPrefix(line, "#") || len(line) < 1 {
				continue
			}
			words := strings.Split(line, " ")
			cmd := ""
			adrs := ""
			val := ""
			if len(words) >= 2 {
				cmd = words[0]
				adrs = words[1]
			}
			if len(words) >= 4 {
				val = strings.Join(words[3:], " ")
			}
			if len(val) > 1 && val[0] == '"' {
				val = val[1 : len(val)-1]
			}
			switch cmd {
			case "leftstring":
				s.SetCell(adrs, &Cell{stringType: true, alignment: align.AlignLeft, value: val})
			case "rightstrng":
				s.SetCell(adrs, &Cell{stringType: true, alignment: align.AlignRight, value: val})
			case "label":
				s.SetCell(adrs, &Cell{stringType: true, alignment: align.AlignCenter, value: val})
			case "let":
				s.SetCell(adrs, &Cell{stringType: false, alignment: align.AlignRight, value: val})
			case "goto":
				s.SelectedCell = adrs
			case "format":
				width, _ := strconv.ParseInt(words[2], 10, 64)
				precision, _ := strconv.ParseInt(words[3], 10, 64)
				ctype, _ := strconv.ParseInt(words[4], 10, 64)
				s.columnFormats[adrs] = ColumnFormat{width: int(width), precision: int(precision), ctype: int(ctype)}
			}
		}
	}

	return s
}

func (s *Sheet) writeFormats(w io.Writer) {
	for k, cFormat := range s.columnFormats {
		fmt.Fprintf(w, "format %s %d %d %d\n", k, cFormat.width, cFormat.precision, cFormat.ctype)
	}
}

func (s *Sheet) getPrecision(address string) int {
	column := address[1:]
	if cFormat, found := s.columnFormats[column]; found {
		return cFormat.precision
	} else {
		return 2
	}
}

func (s *Sheet) DisplayFormat(address string) string {
	column := address[1:]
	if cFormat, found := s.columnFormats[column]; found {
		return fmt.Sprintf("%d %d %d", cFormat.width, cFormat.precision, cFormat.ctype)
	} else {
		return fmt.Sprintf("%d %d %d", STARTING_COLUMN_WIDTH, 2, 0)
	}
}

func (s *Sheet) getColumnWidth(column string) int {
	if cFormat, found := s.columnFormats[column]; found {
		return cFormat.width
	} else {
		s.columnFormats[column] = ColumnFormat{width: STARTING_COLUMN_WIDTH, precision: 2, ctype: 0}
		return STARTING_COLUMN_WIDTH
	}
}

func (s *Sheet) increaseColumnWidth(column string) {
	if cFormat, found := s.columnFormats[column]; found {
		cFormat.width += 1
		s.columnFormats[column] = cFormat
	} else {
		s.columnFormats[column] = ColumnFormat{width: STARTING_COLUMN_WIDTH + 1, precision: 2, ctype: 0}
	}
}

func (s *Sheet) decreaseColumnWidth(column string) {
	if cFormat, found := s.columnFormats[column]; found {
		if cFormat.width > 1 {
			cFormat.width--
			s.columnFormats[column] = cFormat
		}
	} else {
		s.columnFormats[column] = ColumnFormat{width: STARTING_COLUMN_WIDTH - 1, precision: 2, ctype: 0}
	}
}

func (s *Sheet) ClearCell(address string) {
	if cell, err := s.GetCell(address); err == nil {
		for forRef, _ := range cell.forwardRefs {
			if forCell, forErr := s.GetCell(forRef); forErr == nil {
				delete(forCell.backRefs, forRef)
			}
		}
	}
	delete(s.data, address)
}

func (s *Sheet) GetCell(address string) (*Cell, error) {
	if cell, found := s.data[address]; found {
		return cell, nil
	} else if address == s.SelectedCell {
		return &Cell{}, nil
	}
	return nil, errors.New("Cell does not exist in spreadsheet.")
}

func (s *Sheet) SetCell(address string, cell *Cell) {
	if currentCell, found := s.data[address]; found {
		cell.backRefs = currentCell.backRefs
	}
	if !cell.stringType {
		postfix := evaler.GetPostfix(cell.value)
		for _, token := range postfix {
			if evaler.IsCellAddr(token) {
				cell.forwardRefs[token] = struct{}{}
				if tokCell, tokErr := s.GetCell(token); tokErr == nil {
					tokCell.backRefs[address] = struct{}{}
				}
			}
		}
	}
	s.data[address] = cell

	// TODO: change to display current cell and all back references
	s.display()
}

func (s *Sheet) Save() error {
	if outfile, err := os.Create(s.Filename); err == nil {
		fmt.Fprintln(outfile, "# This data file was generated by Spreadsheet Calculator IMproved.")
		fmt.Fprintln(outfile, "# You almost certainly shouldn't edit it.")
		fmt.Fprintln(outfile, "")

		s.writeFormats(outfile)

		for addr, cell := range s.data {
			cell.write(outfile, addr)
		}
		fmt.Fprintf(outfile, "goto %s A0", s.SelectedCell)
		outfile.Close()
		return nil
	} else {
		return err
	}
}
<|endoftext|>


<commit_before><commit_msg>Add renaming of generic oauth2 provider<commit_after><|endoftext|>


<commit_before>// Copyright 2022 The Matrix.org Foundation C.I.C.
//
// 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 test

import (
	"crypto/sha256"
	"database/sql"
	"encoding/hex"
	"fmt"
	"os"
	"os/exec"
	"os/user"
	"testing"

	"github.com/lib/pq"
)

type DBType int

var DBTypeSQLite DBType = 1
var DBTypePostgres DBType = 2

var Quiet = false
var Required = os.Getenv("DENDRITE_TEST_SKIP_NODB") == ""

func fatalError(t *testing.T, format string, args ...interface{}) {
	if Required {
		t.Fatalf(format, args...)
	} else {
		t.Skipf(format, args...)
	}
}

func createLocalDB(t *testing.T, dbName string) {
	if !Quiet {
		t.Log("Note: tests require a postgres install accessible to the current user")
	}
	createDB := exec.Command("createdb", dbName)
	if !Quiet {
		createDB.Stdout = os.Stdout
		createDB.Stderr = os.Stderr
	}
	err := createDB.Run()
	if err != nil && !Quiet {
		fmt.Println("createLocalDB returned error:", err)
	}
}

func createRemoteDB(t *testing.T, dbName, user, connStr string) {
	db, err := sql.Open("postgres", connStr+" dbname=postgres")
	if err != nil {
		fatalError(t, "failed to open postgres conn with connstr=%s : %s", connStr, err)
	}
	_, err = db.Exec(fmt.Sprintf(`CREATE DATABASE %s;`, dbName))
	if err != nil {
		pqErr, ok := err.(*pq.Error)
		if !ok {
			t.Fatalf("failed to CREATE DATABASE: %s", err)
		}
		// we ignore duplicate database error as we expect this
		if pqErr.Code != "42P04" {
			t.Fatalf("failed to CREATE DATABASE with code=%s msg=%s", pqErr.Code, pqErr.Message)
		}
	}
	_, err = db.Exec(fmt.Sprintf(`GRANT ALL PRIVILEGES ON DATABASE %s TO %s`, dbName, user))
	if err != nil {
		t.Fatalf("failed to GRANT: %s", err)
	}
	_ = db.Close()
}

func currentUser() string {
	user, err := user.Current()
	if err != nil {
		if !Quiet {
			fmt.Println("cannot get current user: ", err)
		}
		os.Exit(2)
	}
	return user.Username
}

// Prepare a sqlite or postgres connection string for testing.
// Returns the connection string to use and a close function which must be called when the test finishes.
// Calling this function twice will return the same database, which will have data from previous tests
// unless close() is called.
// TODO: namespace for concurrent package tests
func PrepareDBConnectionString(t *testing.T, dbType DBType) (connStr string, close func()) {
	if dbType == DBTypeSQLite {
		// this will be made in the current working directory which namespaces concurrent package runs correctly
		dbname := "dendrite_test.db"
		return fmt.Sprintf("file:%s", dbname), func() {
			err := os.Remove(dbname)
			if err != nil {
				t.Fatalf("failed to cleanup sqlite db '%s': %s", dbname, err)
			}
		}
	}

	// Required vars: user and db
	// We'll try to infer from the local env if they are missing
	user := os.Getenv("POSTGRES_USER")
	if user == "" {
		user = currentUser()
	}
	connStr = fmt.Sprintf(
		"user=%s sslmode=disable",
		user,
	)
	// optional vars, used in CI
	password := os.Getenv("POSTGRES_PASSWORD")
	if password != "" {
		connStr += fmt.Sprintf(" password=%s", password)
	}
	host := os.Getenv("POSTGRES_HOST")
	if host != "" {
		connStr += fmt.Sprintf(" host=%s", host)
	}

	// superuser database
	postgresDB := os.Getenv("POSTGRES_DB")
	// we cannot use 'dendrite_test' here else 2x concurrently running packages will try to use the same db.
	// instead, hash the current working directory, snaffle the first 16 bytes and append that to dendrite_test
	// and use that as the unique db name. We do this because packages are per-directory hence by hashing the
	// working (test) directory we ensure we get a consistent hash and don't hash against concurrent packages.
	wd, err := os.Getwd()
	if err != nil {
		t.Fatalf("cannot get working directory: %s", err)
	}
	hash := sha256.Sum256([]byte(wd))
	dbName := fmt.Sprintf("dendrite_test_%s", hex.EncodeToString(hash[:16]))
	if postgresDB == "" { // local server, use createdb
		createLocalDB(t, dbName)
	} else { // remote server, shell into the postgres user and CREATE DATABASE
		createRemoteDB(t, dbName, user, connStr)
	}
	connStr += fmt.Sprintf(" dbname=%s", dbName)

	return connStr, func() {
		// Drop all tables on the database to get a fresh instance
		db, err := sql.Open("postgres", connStr)
		if err != nil {
			t.Fatalf("failed to connect to postgres db '%s': %s", connStr, err)
		}
		_, err = db.Exec(`DROP SCHEMA public CASCADE;
		CREATE SCHEMA public;`)
		if err != nil {
			t.Fatalf("failed to cleanup postgres db '%s': %s", connStr, err)
		}
		_ = db.Close()
	}
}

// Creates subtests with each known DBType
func WithAllDatabases(t *testing.T, testFn func(t *testing.T, db DBType)) {
	dbs := map[string]DBType{
		"postgres": DBTypePostgres,
		"sqlite":   DBTypeSQLite,
	}
	for dbName, dbType := range dbs {
		dbt := dbType
		t.Run(dbName, func(tt *testing.T) {
			tt.Parallel()
			testFn(tt, dbt)
		})
	}
}
<commit_msg>Really SKIP_NODB (#2472)<commit_after>// Copyright 2022 The Matrix.org Foundation C.I.C.
//
// 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 test

import (
	"crypto/sha256"
	"database/sql"
	"encoding/hex"
	"fmt"
	"os"
	"os/exec"
	"os/user"
	"testing"

	"github.com/lib/pq"
)

type DBType int

var DBTypeSQLite DBType = 1
var DBTypePostgres DBType = 2

var Quiet = false
var Required = os.Getenv("DENDRITE_TEST_SKIP_NODB") == ""

func fatalError(t *testing.T, format string, args ...interface{}) {
	if Required {
		t.Fatalf(format, args...)
	} else {
		t.Skipf(format, args...)
	}
}

func createLocalDB(t *testing.T, dbName string) {
	if _, err := exec.LookPath("createdb"); err != nil {
		fatalError(t, "Note: tests require a postgres install accessible to the current user")
		return
	}
	createDB := exec.Command("createdb", dbName)
	if !Quiet {
		createDB.Stdout = os.Stdout
		createDB.Stderr = os.Stderr
	}
	err := createDB.Run()
	if err != nil && !Quiet {
		fmt.Println("createLocalDB returned error:", err)
	}
}

func createRemoteDB(t *testing.T, dbName, user, connStr string) {
	db, err := sql.Open("postgres", connStr+" dbname=postgres")
	if err != nil {
		fatalError(t, "failed to open postgres conn with connstr=%s : %s", connStr, err)
	}
	if err = db.Ping(); err != nil {
		fatalError(t, "failed to open postgres conn with connstr=%s : %s", connStr, err)
	}
	_, err = db.Exec(fmt.Sprintf(`CREATE DATABASE %s;`, dbName))
	if err != nil {
		pqErr, ok := err.(*pq.Error)
		if !ok {
			t.Fatalf("failed to CREATE DATABASE: %s", err)
		}
		// we ignore duplicate database error as we expect this
		if pqErr.Code != "42P04" {
			t.Fatalf("failed to CREATE DATABASE with code=%s msg=%s", pqErr.Code, pqErr.Message)
		}
	}
	_, err = db.Exec(fmt.Sprintf(`GRANT ALL PRIVILEGES ON DATABASE %s TO %s`, dbName, user))
	if err != nil {
		t.Fatalf("failed to GRANT: %s", err)
	}
	_ = db.Close()
}

func currentUser() string {
	user, err := user.Current()
	if err != nil {
		if !Quiet {
			fmt.Println("cannot get current user: ", err)
		}
		os.Exit(2)
	}
	return user.Username
}

// Prepare a sqlite or postgres connection string for testing.
// Returns the connection string to use and a close function which must be called when the test finishes.
// Calling this function twice will return the same database, which will have data from previous tests
// unless close() is called.
// TODO: namespace for concurrent package tests
func PrepareDBConnectionString(t *testing.T, dbType DBType) (connStr string, close func()) {
	if dbType == DBTypeSQLite {
		// this will be made in the current working directory which namespaces concurrent package runs correctly
		dbname := "dendrite_test.db"
		return fmt.Sprintf("file:%s", dbname), func() {
			err := os.Remove(dbname)
			if err != nil {
				t.Fatalf("failed to cleanup sqlite db '%s': %s", dbname, err)
			}
		}
	}

	// Required vars: user and db
	// We'll try to infer from the local env if they are missing
	user := os.Getenv("POSTGRES_USER")
	if user == "" {
		user = currentUser()
	}
	connStr = fmt.Sprintf(
		"user=%s sslmode=disable",
		user,
	)
	// optional vars, used in CI
	password := os.Getenv("POSTGRES_PASSWORD")
	if password != "" {
		connStr += fmt.Sprintf(" password=%s", password)
	}
	host := os.Getenv("POSTGRES_HOST")
	if host != "" {
		connStr += fmt.Sprintf(" host=%s", host)
	}

	// superuser database
	postgresDB := os.Getenv("POSTGRES_DB")
	// we cannot use 'dendrite_test' here else 2x concurrently running packages will try to use the same db.
	// instead, hash the current working directory, snaffle the first 16 bytes and append that to dendrite_test
	// and use that as the unique db name. We do this because packages are per-directory hence by hashing the
	// working (test) directory we ensure we get a consistent hash and don't hash against concurrent packages.
	wd, err := os.Getwd()
	if err != nil {
		t.Fatalf("cannot get working directory: %s", err)
	}
	hash := sha256.Sum256([]byte(wd))
	dbName := fmt.Sprintf("dendrite_test_%s", hex.EncodeToString(hash[:16]))
	if postgresDB == "" { // local server, use createdb
		createLocalDB(t, dbName)
	} else { // remote server, shell into the postgres user and CREATE DATABASE
		createRemoteDB(t, dbName, user, connStr)
	}
	connStr += fmt.Sprintf(" dbname=%s", dbName)

	return connStr, func() {
		// Drop all tables on the database to get a fresh instance
		db, err := sql.Open("postgres", connStr)
		if err != nil {
			t.Fatalf("failed to connect to postgres db '%s': %s", connStr, err)
		}
		_, err = db.Exec(`DROP SCHEMA public CASCADE;
		CREATE SCHEMA public;`)
		if err != nil {
			t.Fatalf("failed to cleanup postgres db '%s': %s", connStr, err)
		}
		_ = db.Close()
	}
}

// Creates subtests with each known DBType
func WithAllDatabases(t *testing.T, testFn func(t *testing.T, db DBType)) {
	dbs := map[string]DBType{
		"postgres": DBTypePostgres,
		"sqlite":   DBTypeSQLite,
	}
	for dbName, dbType := range dbs {
		dbt := dbType
		t.Run(dbName, func(tt *testing.T) {
			tt.Parallel()
			testFn(tt, dbt)
		})
	}
}
<|endoftext|>


<commit_before>package errors

import "sync"

// ErrorList is used to chain a list of potential errors and is thread-safe.
type ErrorList struct {
	mux  sync.RWMutex
	errs []error
}

// Error will return the string-form of the errors.
// Implements the error interface.
func (e *ErrorList) Error() string {
	if e == nil {
		return ""
	}

	e.mux.RLock()
	defer e.mux.RUnlock()

	if len(e.errs) == 0 {
		return ""
	}

	if len(e.errs) == 1 {
		return e.errs[0].Error()
	}

	b := []byte("the following errors occured:\n")
	for _, err := range e.errs {
		b = append(b, err.Error()...)
		b = append(b, '\n')
	}

	return string(b)
}

// Err will return an error if the errorlist is not empty.
// If there's only 1 error, it will be directly returned.
// If the errorlist is empty - nil is returned.
func (e *ErrorList) Err() (err error) {
	if e == nil {
		return
	}
	e.mux.RLock()
	switch len(e.errs) {
	case 0: // do nothing
	case 1:
		err = e.errs[0]
	default:
		err = e
	}
	e.mux.RLock()
	return
}

// Push will push an error to the errorlist
// If err is a errorlist, it will be merged.
// If the errorlist is nil, it will be created.
func (e *ErrorList) Push(err error) {
	if err == nil {
		return
	}

	if e == nil {
		*e = ErrorList{}
	}

	e.mux.Lock()
	switch v := err.(type) {
	case *ErrorList:
		v.ForEach(func(err error) {
			e.errs = append(e.errs, err)
		})
	default:
		e.errs = append(e.errs, err)
	}
	e.mux.Unlock()
}

// ForEach will iterate through all of the errors within the error list.
func (e *ErrorList) ForEach(fn func(error)) {
	if e == nil {
		return
	}

	e.mux.RLock()
	for _, err := range e.errs {
		fn(err)
	}
	e.mux.RUnlock()
}

// Copy will copy the items from the inbound error list to the source
func (e *ErrorList) Copy(in *ErrorList) {
	e.errs = append(e.errs, in.errs...)
}

// Len will return the length of the inner errors list.
func (e *ErrorList) Len() (n int) {
	if e == nil {
		return
	}

	e.mux.RLock()
	n = len(e.errs)
	e.mux.RUnlock()
	return
}
<commit_msg>Fix race condition<commit_after>package errors

import "sync"

// ErrorList is used to chain a list of potential errors and is thread-safe.
type ErrorList struct {
	mux  sync.RWMutex
	errs []error
}

// Error will return the string-form of the errors.
// Implements the error interface.
func (e *ErrorList) Error() string {
	if e == nil {
		return ""
	}

	e.mux.RLock()
	defer e.mux.RUnlock()

	if len(e.errs) == 0 {
		return ""
	}

	if len(e.errs) == 1 {
		return e.errs[0].Error()
	}

	b := []byte("the following errors occured:\n")
	for _, err := range e.errs {
		b = append(b, err.Error()...)
		b = append(b, '\n')
	}

	return string(b)
}

// Err will return an error if the errorlist is not empty.
// If there's only 1 error, it will be directly returned.
// If the errorlist is empty - nil is returned.
func (e *ErrorList) Err() (err error) {
	if e == nil {
		return
	}
	e.mux.RLock()
	switch len(e.errs) {
	case 0: // do nothing
	case 1:
		err = e.errs[0]
	default:
		err = e
	}
	e.mux.RLock()
	return
}

// Push will push an error to the errorlist
// If err is a errorlist, it will be merged.
// If the errorlist is nil, it will be created.
func (e *ErrorList) Push(err error) {
	if err == nil {
		return
	}

	if e == nil {
		*e = ErrorList{}
	}

	e.mux.Lock()
	switch v := err.(type) {
	case *ErrorList:
		v.ForEach(func(err error) {
			e.errs = append(e.errs, err)
		})
	default:
		e.errs = append(e.errs, err)
	}
	e.mux.Unlock()
}

// ForEach will iterate through all of the errors within the error list.
func (e *ErrorList) ForEach(fn func(error)) {
	if e == nil {
		return
	}

	e.mux.RLock()
	for _, err := range e.errs {
		fn(err)
	}
	e.mux.RUnlock()
}

// Copy will copy the items from the inbound error list to the source
func (e *ErrorList) Copy(in *ErrorList) {
	if in == nil {
		return
	}

	e.mux.Lock()
	defer e.mux.Unlock()

	in.mux.RLock()
	defer in.mux.RUnlock()

	e.errs = append(e.errs, in.errs...)
}

// Len will return the length of the inner errors list.
func (e *ErrorList) Len() (n int) {
	if e == nil {
		return
	}

	e.mux.RLock()
	n = len(e.errs)
	e.mux.RUnlock()
	return
}
<|endoftext|>


<commit_before>package simpleredis

import (
	"errors"
	"strconv"
	"strings"

	"github.com/garyburd/redigo/redis"
)

// Common for each of the redis datastructures used here
type redisDatastructure struct {
	pool    *ConnectionPool
	id      string
	dbindex int
}

type (
	// A pool of readily available Redis connections
	ConnectionPool redis.Pool

	List     redisDatastructure
	Set      redisDatastructure
	HashMap  redisDatastructure
	KeyValue redisDatastructure
)

const (
	// Version number. Stable API within major version numbers.
	Version = 1.0
	// The default [url]:port that Redis is running at
	defaultRedisServer = ":6379"
)

var (
	// How many connections should stay ready for requests, at a maximum?
	// When an idle connection is used, new idle connections are created.
	maxIdleConnections = 3
)

/* --- Helper functions --- */

// Connect to the local instance of Redis at port 6379
func newRedisConnection() (redis.Conn, error) {
	return newRedisConnectionTo(defaultRedisServer)
}

// Connect to host:port, host may be omitted, so ":6379" is valid.
// Will not try to AUTH with any given password (password@host:port).
func newRedisConnectionTo(hostColonPort string) (redis.Conn, error) {
	// Discard the password, if provided
	if _, theRest, ok := twoFields(hostColonPort, "@"); ok {
		hostColonPort = theRest
	}
	return redis.Dial("tcp", hostColonPort)
}

// Get a string from a list of results at a given position
func getString(bi []interface{}, i int) string {
	return string(bi[i].([]uint8))
}

// Test if the local Redis server is up and running
func TestConnection() (err error) {
	return TestConnectionHost(defaultRedisServer)
}

// Test if a given Redis server at host:port is up and running.
// Does not try to PING or AUTH.
func TestConnectionHost(hostColonPort string) (err error) {
	// Connect to the given host:port
	conn, err := newRedisConnectionTo(hostColonPort)
	defer conn.Close()
	defer func() {
		if r := recover(); r != nil {
			err = errors.New("Could not connect to redis server: " + hostColonPort)
		}
	}()
	return err
}

/* --- ConnectionPool functions --- */

// Create a new connection pool
func NewConnectionPool() *ConnectionPool {
	// The second argument is the maximum number of idle connections
	redisPool := redis.NewPool(newRedisConnection, maxIdleConnections)
	pool := ConnectionPool(*redisPool)
	return &pool
}

// Split a string into two parts, given a delimiter.
// Returns the two parts and true if it works out.
func twoFields(s, delim string) (string, string, bool) {
	if strings.Count(s, delim) != 1 {
		return s, "", false
	}
	fields := strings.Split(s, delim)
	return fields[0], fields[1], true
}

// Create a new connection pool given a host:port string.
// A password may be supplied as well, on the form "password@host:port".
func NewConnectionPoolHost(hostColonPort string) *ConnectionPool {
	// Create a redis Pool
	redisPool := redis.NewPool(
		// Anonymous function for calling new RedisConnectionTo with the host:port
		func() (redis.Conn, error) {
			conn, err := newRedisConnectionTo(hostColonPort)
			if err != nil {
				return nil, err
			}
			// If a password is given, use it to authenticate
			if password, _, ok := twoFields(hostColonPort, "@"); ok {
				if _, err := conn.Do("AUTH", password); err != nil {
					conn.Close()
					return nil, err
				}
			}
			return conn, err
		},
		// Maximum number of idle connections to the redis database
		maxIdleConnections)
	pool := ConnectionPool(*redisPool)
	return &pool
}

// Set the number of maximum *idle* connections standing ready when
// creating new connection pools. When an idle connection is used,
// a new idle connection is created. The default is 3 and should be fine
// for most cases.
func SetMaxIdleConnections(maximum int) {
	maxIdleConnections = maximum
}

// Get one of the available connections from the connection pool, given a database index
func (pool *ConnectionPool) Get(dbindex int) redis.Conn {
	redisPool := redis.Pool(*pool)
	conn := redisPool.Get()
	// The default database index is 0
	if dbindex != 0 {
		// SELECT is not critical, ignore the return values
		conn.Do("SELECT", strconv.Itoa(dbindex))
	}
	return conn
}

// Ping the server by sending a PING command
func (pool *ConnectionPool) Ping() (pong bool) {
	redisPool := redis.Pool(*pool)
	conn := redisPool.Get()
	_, err := conn.Do("PING")
	return err == nil
}

// Close down the connection pool
func (pool *ConnectionPool) Close() {
	redisPool := redis.Pool(*pool)
	redisPool.Close()
}

/* --- List functions --- */

// Create a new list
func NewList(pool *ConnectionPool, id string) *List {
	return &List{pool, id, 0}
}

// Select a different database
func (rl *List) SelectDatabase(dbindex int) {
	rl.dbindex = dbindex
}

// Add an element to the list
func (rl *List) Add(value string) error {
	conn := rl.pool.Get(rl.dbindex)
	_, err := conn.Do("RPUSH", rl.id, value)
	return err
}

// Get all elements of a list
func (rl *List) GetAll() ([]string, error) {
	conn := rl.pool.Get(rl.dbindex)
	result, err := redis.Values(conn.Do("LRANGE", rl.id, "0", "-1"))
	strs := make([]string, len(result))
	for i := 0; i < len(result); i++ {
		strs[i] = getString(result, i)
	}
	return strs, err
}

// Get the last element of a list
func (rl *List) GetLast() (string, error) {
	conn := rl.pool.Get(rl.dbindex)
	result, err := redis.Values(conn.Do("LRANGE", rl.id, "-1", "-1"))
	if len(result) == 1 {
		return getString(result, 0), err
	}
	return "", err
}

// Get the last N elements of a list
func (rl *List) GetLastN(n int) ([]string, error) {
	conn := rl.pool.Get(rl.dbindex)
	result, err := redis.Values(conn.Do("LRANGE", rl.id, "-"+strconv.Itoa(n), "-1"))
	strs := make([]string, len(result))
	for i := 0; i < len(result); i++ {
		strs[i] = getString(result, i)
	}
	return strs, err
}

// Remove this list
func (rl *List) Remove() error {
	conn := rl.pool.Get(rl.dbindex)
	_, err := conn.Do("DEL", rl.id)
	return err
}

/* --- Set functions --- */

// Create a new set
func NewSet(pool *ConnectionPool, id string) *Set {
	return &Set{pool, id, 0}
}

// Select a different database
func (rs *Set) SelectDatabase(dbindex int) {
	rs.dbindex = dbindex
}

// Add an element to the set
func (rs *Set) Add(value string) error {
	conn := rs.pool.Get(rs.dbindex)
	_, err := conn.Do("SADD", rs.id, value)
	return err
}

// Check if a given value is in the set
func (rs *Set) Has(value string) (bool, error) {
	conn := rs.pool.Get(rs.dbindex)
	retval, err := conn.Do("SISMEMBER", rs.id, value)
	if err != nil {
		panic(err)
	}
	return redis.Bool(retval, err)
}

// Get all elements of the set
func (rs *Set) GetAll() ([]string, error) {
	conn := rs.pool.Get(rs.dbindex)
	result, err := redis.Values(conn.Do("SMEMBERS", rs.id))
	strs := make([]string, len(result))
	for i := 0; i < len(result); i++ {
		strs[i] = getString(result, i)
	}
	return strs, err
}

// Remove an element from the set
func (rs *Set) Del(value string) error {
	conn := rs.pool.Get(rs.dbindex)
	_, err := conn.Do("SREM", rs.id, value)
	return err
}

// Remove this set
func (rs *Set) Remove() error {
	conn := rs.pool.Get(rs.dbindex)
	_, err := conn.Do("DEL", rs.id)
	return err
}

/* --- HashMap functions --- */

// Create a new hashmap
func NewHashMap(pool *ConnectionPool, id string) *HashMap {
	return &HashMap{pool, id, 0}
}

// Select a different database
func (rh *HashMap) SelectDatabase(dbindex int) {
	rh.dbindex = dbindex
}

// Set a value in a hashmap given the element id (for instance a user id) and the key (for instance "password")
func (rh *HashMap) Set(elementid, key, value string) error {
	conn := rh.pool.Get(rh.dbindex)
	_, err := conn.Do("HSET", rh.id+":"+elementid, key, value)
	return err
}

// Get a value from a hashmap given the element id (for instance a user id) and the key (for instance "password")
func (rh *HashMap) Get(elementid, key string) (string, error) {
	conn := rh.pool.Get(rh.dbindex)
	result, err := redis.String(conn.Do("HGET", rh.id+":"+elementid, key))
	if err != nil {
		return "", err
	}
	return result, nil
}

// Check if a given elementid + key is in the hash map
func (rh *HashMap) Has(elementid, key string) (bool, error) {
	conn := rh.pool.Get(rh.dbindex)
	retval, err := conn.Do("HEXISTS", rh.id+":"+elementid, key)
	if err != nil {
		panic(err)
	}
	return redis.Bool(retval, err)
}

// Check if a given elementid exists as a hash map at all
func (rh *HashMap) Exists(elementid string) (bool, error) {
	// TODO: key is not meant to be a wildcard, check for "*"
	return hasKey(rh.pool, rh.id+":"+elementid, rh.dbindex)
}

// Get all elementid's for all hash elements
func (rh *HashMap) GetAll() ([]string, error) {
	conn := rh.pool.Get(rh.dbindex)
	result, err := redis.Values(conn.Do("KEYS", rh.id+":*"))
	strs := make([]string, len(result))
	idlen := len(rh.id)
	for i := 0; i < len(result); i++ {
		strs[i] = getString(result, i)[idlen+1:]
	}
	return strs, err
}

// Remove a key for an entry in a hashmap (for instance the email field for a user)
func (rh *HashMap) DelKey(elementid, key string) error {
	conn := rh.pool.Get(rh.dbindex)
	_, err := conn.Do("HDEL", rh.id+":"+elementid, key)
	return err
}

// Remove an element (for instance a user)
func (rh *HashMap) Del(elementid string) error {
	conn := rh.pool.Get(rh.dbindex)
	_, err := conn.Do("DEL", rh.id+":"+elementid)
	return err
}

// Remove this hashmap
func (rh *HashMap) Remove() error {
	conn := rh.pool.Get(rh.dbindex)
	_, err := conn.Do("DEL", rh.id)
	return err
}

/* --- KeyValue functions --- */

// Create a new key/value
func NewKeyValue(pool *ConnectionPool, id string) *KeyValue {
	return &KeyValue{pool, id, 0}
}

// Select a different database
func (rkv *KeyValue) SelectDatabase(dbindex int) {
	rkv.dbindex = dbindex
}

// Set a key and value
func (rkv *KeyValue) Set(key, value string) error {
	conn := rkv.pool.Get(rkv.dbindex)
	_, err := conn.Do("SET", rkv.id+":"+key, value)
	return err
}

// Get a value given a key
func (rkv *KeyValue) Get(key string) (string, error) {
	conn := rkv.pool.Get(rkv.dbindex)
	result, err := redis.String(conn.Do("GET", rkv.id+":"+key))
	if err != nil {
		return "", err
	}
	return result, nil
}

// Remove a key
func (rkv *KeyValue) Del(key string) error {
	conn := rkv.pool.Get(rkv.dbindex)
	_, err := conn.Do("DEL", rkv.id+":"+key)
	return err
}

// Remove this key/value
func (rkv *KeyValue) Remove() error {
	conn := rkv.pool.Get(rkv.dbindex)
	_, err := conn.Do("DEL", rkv.id)
	return err
}

// --- Generic redis functions ---

// Check if a key exists. The key can be a wildcard (ie. "user*").
func hasKey(pool *ConnectionPool, wildcard string, dbindex int) (bool, error) {
	conn := pool.Get(dbindex)
	result, err := redis.Values(conn.Do("KEYS", wildcard))
	if err != nil {
		return false, err
	}
	return len(result) > 0, nil
}
<commit_msg>No empty passwords<commit_after>package simpleredis

import (
	"errors"
	"strconv"
	"strings"

	"github.com/garyburd/redigo/redis"
)

// Common for each of the redis datastructures used here
type redisDatastructure struct {
	pool    *ConnectionPool
	id      string
	dbindex int
}

type (
	// A pool of readily available Redis connections
	ConnectionPool redis.Pool

	List     redisDatastructure
	Set      redisDatastructure
	HashMap  redisDatastructure
	KeyValue redisDatastructure
)

const (
	// Version number. Stable API within major version numbers.
	Version = 1.0
	// The default [url]:port that Redis is running at
	defaultRedisServer = ":6379"
)

var (
	// How many connections should stay ready for requests, at a maximum?
	// When an idle connection is used, new idle connections are created.
	maxIdleConnections = 3
)

/* --- Helper functions --- */

// Connect to the local instance of Redis at port 6379
func newRedisConnection() (redis.Conn, error) {
	return newRedisConnectionTo(defaultRedisServer)
}

// Connect to host:port, host may be omitted, so ":6379" is valid.
// Will not try to AUTH with any given password (password@host:port).
func newRedisConnectionTo(hostColonPort string) (redis.Conn, error) {
	// Discard the password, if provided
	if _, theRest, ok := twoFields(hostColonPort, "@"); ok {
		hostColonPort = theRest
	}
	return redis.Dial("tcp", hostColonPort)
}

// Get a string from a list of results at a given position
func getString(bi []interface{}, i int) string {
	return string(bi[i].([]uint8))
}

// Test if the local Redis server is up and running
func TestConnection() (err error) {
	return TestConnectionHost(defaultRedisServer)
}

// Test if a given Redis server at host:port is up and running.
// Does not try to PING or AUTH.
func TestConnectionHost(hostColonPort string) (err error) {
	// Connect to the given host:port
	conn, err := newRedisConnectionTo(hostColonPort)
	defer conn.Close()
	defer func() {
		if r := recover(); r != nil {
			err = errors.New("Could not connect to redis server: " + hostColonPort)
		}
	}()
	return err
}

/* --- ConnectionPool functions --- */

// Create a new connection pool
func NewConnectionPool() *ConnectionPool {
	// The second argument is the maximum number of idle connections
	redisPool := redis.NewPool(newRedisConnection, maxIdleConnections)
	pool := ConnectionPool(*redisPool)
	return &pool
}

// Split a string into two parts, given a delimiter.
// Returns the two parts and true if it works out.
func twoFields(s, delim string) (string, string, bool) {
	if strings.Count(s, delim) != 1 {
		return s, "", false
	}
	fields := strings.Split(s, delim)
	return fields[0], fields[1], true
}

// Create a new connection pool given a host:port string.
// A password may be supplied as well, on the form "password@host:port".
func NewConnectionPoolHost(hostColonPort string) *ConnectionPool {
	// Create a redis Pool
	redisPool := redis.NewPool(
		// Anonymous function for calling new RedisConnectionTo with the host:port
		func() (redis.Conn, error) {
			conn, err := newRedisConnectionTo(hostColonPort)
			if err != nil {
				return nil, err
			}
			// If a password is given, use it to authenticate
			if password, _, ok := twoFields(hostColonPort, "@"); ok {
				if password != "" {
					if _, err := conn.Do("AUTH", password); err != nil {
						conn.Close()
						return nil, err
					}
				}
			}
			return conn, err
		},
		// Maximum number of idle connections to the redis database
		maxIdleConnections)
	pool := ConnectionPool(*redisPool)
	return &pool
}

// Set the number of maximum *idle* connections standing ready when
// creating new connection pools. When an idle connection is used,
// a new idle connection is created. The default is 3 and should be fine
// for most cases.
func SetMaxIdleConnections(maximum int) {
	maxIdleConnections = maximum
}

// Get one of the available connections from the connection pool, given a database index
func (pool *ConnectionPool) Get(dbindex int) redis.Conn {
	redisPool := redis.Pool(*pool)
	conn := redisPool.Get()
	// The default database index is 0
	if dbindex != 0 {
		// SELECT is not critical, ignore the return values
		conn.Do("SELECT", strconv.Itoa(dbindex))
	}
	return conn
}

// Ping the server by sending a PING command
func (pool *ConnectionPool) Ping() (pong bool) {
	redisPool := redis.Pool(*pool)
	conn := redisPool.Get()
	_, err := conn.Do("PING")
	return err == nil
}

// Close down the connection pool
func (pool *ConnectionPool) Close() {
	redisPool := redis.Pool(*pool)
	redisPool.Close()
}

/* --- List functions --- */

// Create a new list
func NewList(pool *ConnectionPool, id string) *List {
	return &List{pool, id, 0}
}

// Select a different database
func (rl *List) SelectDatabase(dbindex int) {
	rl.dbindex = dbindex
}

// Add an element to the list
func (rl *List) Add(value string) error {
	conn := rl.pool.Get(rl.dbindex)
	_, err := conn.Do("RPUSH", rl.id, value)
	return err
}

// Get all elements of a list
func (rl *List) GetAll() ([]string, error) {
	conn := rl.pool.Get(rl.dbindex)
	result, err := redis.Values(conn.Do("LRANGE", rl.id, "0", "-1"))
	strs := make([]string, len(result))
	for i := 0; i < len(result); i++ {
		strs[i] = getString(result, i)
	}
	return strs, err
}

// Get the last element of a list
func (rl *List) GetLast() (string, error) {
	conn := rl.pool.Get(rl.dbindex)
	result, err := redis.Values(conn.Do("LRANGE", rl.id, "-1", "-1"))
	if len(result) == 1 {
		return getString(result, 0), err
	}
	return "", err
}

// Get the last N elements of a list
func (rl *List) GetLastN(n int) ([]string, error) {
	conn := rl.pool.Get(rl.dbindex)
	result, err := redis.Values(conn.Do("LRANGE", rl.id, "-"+strconv.Itoa(n), "-1"))
	strs := make([]string, len(result))
	for i := 0; i < len(result); i++ {
		strs[i] = getString(result, i)
	}
	return strs, err
}

// Remove this list
func (rl *List) Remove() error {
	conn := rl.pool.Get(rl.dbindex)
	_, err := conn.Do("DEL", rl.id)
	return err
}

/* --- Set functions --- */

// Create a new set
func NewSet(pool *ConnectionPool, id string) *Set {
	return &Set{pool, id, 0}
}

// Select a different database
func (rs *Set) SelectDatabase(dbindex int) {
	rs.dbindex = dbindex
}

// Add an element to the set
func (rs *Set) Add(value string) error {
	conn := rs.pool.Get(rs.dbindex)
	_, err := conn.Do("SADD", rs.id, value)
	return err
}

// Check if a given value is in the set
func (rs *Set) Has(value string) (bool, error) {
	conn := rs.pool.Get(rs.dbindex)
	retval, err := conn.Do("SISMEMBER", rs.id, value)
	if err != nil {
		panic(err)
	}
	return redis.Bool(retval, err)
}

// Get all elements of the set
func (rs *Set) GetAll() ([]string, error) {
	conn := rs.pool.Get(rs.dbindex)
	result, err := redis.Values(conn.Do("SMEMBERS", rs.id))
	strs := make([]string, len(result))
	for i := 0; i < len(result); i++ {
		strs[i] = getString(result, i)
	}
	return strs, err
}

// Remove an element from the set
func (rs *Set) Del(value string) error {
	conn := rs.pool.Get(rs.dbindex)
	_, err := conn.Do("SREM", rs.id, value)
	return err
}

// Remove this set
func (rs *Set) Remove() error {
	conn := rs.pool.Get(rs.dbindex)
	_, err := conn.Do("DEL", rs.id)
	return err
}

/* --- HashMap functions --- */

// Create a new hashmap
func NewHashMap(pool *ConnectionPool, id string) *HashMap {
	return &HashMap{pool, id, 0}
}

// Select a different database
func (rh *HashMap) SelectDatabase(dbindex int) {
	rh.dbindex = dbindex
}

// Set a value in a hashmap given the element id (for instance a user id) and the key (for instance "password")
func (rh *HashMap) Set(elementid, key, value string) error {
	conn := rh.pool.Get(rh.dbindex)
	_, err := conn.Do("HSET", rh.id+":"+elementid, key, value)
	return err
}

// Get a value from a hashmap given the element id (for instance a user id) and the key (for instance "password")
func (rh *HashMap) Get(elementid, key string) (string, error) {
	conn := rh.pool.Get(rh.dbindex)
	result, err := redis.String(conn.Do("HGET", rh.id+":"+elementid, key))
	if err != nil {
		return "", err
	}
	return result, nil
}

// Check if a given elementid + key is in the hash map
func (rh *HashMap) Has(elementid, key string) (bool, error) {
	conn := rh.pool.Get(rh.dbindex)
	retval, err := conn.Do("HEXISTS", rh.id+":"+elementid, key)
	if err != nil {
		panic(err)
	}
	return redis.Bool(retval, err)
}

// Check if a given elementid exists as a hash map at all
func (rh *HashMap) Exists(elementid string) (bool, error) {
	// TODO: key is not meant to be a wildcard, check for "*"
	return hasKey(rh.pool, rh.id+":"+elementid, rh.dbindex)
}

// Get all elementid's for all hash elements
func (rh *HashMap) GetAll() ([]string, error) {
	conn := rh.pool.Get(rh.dbindex)
	result, err := redis.Values(conn.Do("KEYS", rh.id+":*"))
	strs := make([]string, len(result))
	idlen := len(rh.id)
	for i := 0; i < len(result); i++ {
		strs[i] = getString(result, i)[idlen+1:]
	}
	return strs, err
}

// Remove a key for an entry in a hashmap (for instance the email field for a user)
func (rh *HashMap) DelKey(elementid, key string) error {
	conn := rh.pool.Get(rh.dbindex)
	_, err := conn.Do("HDEL", rh.id+":"+elementid, key)
	return err
}

// Remove an element (for instance a user)
func (rh *HashMap) Del(elementid string) error {
	conn := rh.pool.Get(rh.dbindex)
	_, err := conn.Do("DEL", rh.id+":"+elementid)
	return err
}

// Remove this hashmap
func (rh *HashMap) Remove() error {
	conn := rh.pool.Get(rh.dbindex)
	_, err := conn.Do("DEL", rh.id)
	return err
}

/* --- KeyValue functions --- */

// Create a new key/value
func NewKeyValue(pool *ConnectionPool, id string) *KeyValue {
	return &KeyValue{pool, id, 0}
}

// Select a different database
func (rkv *KeyValue) SelectDatabase(dbindex int) {
	rkv.dbindex = dbindex
}

// Set a key and value
func (rkv *KeyValue) Set(key, value string) error {
	conn := rkv.pool.Get(rkv.dbindex)
	_, err := conn.Do("SET", rkv.id+":"+key, value)
	return err
}

// Get a value given a key
func (rkv *KeyValue) Get(key string) (string, error) {
	conn := rkv.pool.Get(rkv.dbindex)
	result, err := redis.String(conn.Do("GET", rkv.id+":"+key))
	if err != nil {
		return "", err
	}
	return result, nil
}

// Remove a key
func (rkv *KeyValue) Del(key string) error {
	conn := rkv.pool.Get(rkv.dbindex)
	_, err := conn.Do("DEL", rkv.id+":"+key)
	return err
}

// Remove this key/value
func (rkv *KeyValue) Remove() error {
	conn := rkv.pool.Get(rkv.dbindex)
	_, err := conn.Do("DEL", rkv.id)
	return err
}

// --- Generic redis functions ---

// Check if a key exists. The key can be a wildcard (ie. "user*").
func hasKey(pool *ConnectionPool, wildcard string, dbindex int) (bool, error) {
	conn := pool.Get(dbindex)
	result, err := redis.Values(conn.Do("KEYS", wildcard))
	if err != nil {
		return false, err
	}
	return len(result) > 0, nil
}
<|endoftext|>


<commit_before>package contacts

import (
	"context"
	"fmt"

	"github.com/keybase/client/go/encrypteddb"

	"github.com/keybase/client/go/libkb"
	"github.com/keybase/client/go/protocol/keybase1"
)

// Saving contact list into encrypted db.

// Cache resolutions of a lookup ran on entire contact list provided by the
// frontend. Assume every time SaveContacts is called, entire contact list is
// passed as an argument. Always cache the result of last resolution, do not do
// any result merging.

type SavedContactsStore struct {
	encryptedDB *encrypteddb.EncryptedDB
}

var _ libkb.SyncedContactListProvider = (*SavedContactsStore)(nil)

// NewSavedContactsStore creates a new SavedContactsStore for global context.
// The store is used to securely store list of resolved contacts.
func NewSavedContactsStore(g *libkb.GlobalContext) *SavedContactsStore {
	keyFn := func(ctx context.Context) ([32]byte, error) {
		return encrypteddb.GetSecretBoxKey(ctx, g, encrypteddb.DefaultSecretUI,
			libkb.EncryptionReasonContactsLocalStorage, "encrypting local contact list")
	}
	dbFn := func(g *libkb.GlobalContext) *libkb.JSONLocalDb {
		return g.LocalDb
	}
	return &SavedContactsStore{
		encryptedDB: encrypteddb.New(g, dbFn, keyFn),
	}
}

func ServiceInit(g *libkb.GlobalContext) {
	g.SyncedContactList = NewSavedContactsStore(g)
}

func savedContactsDbKey(uid keybase1.UID) libkb.DbKey {
	return libkb.DbKey{
		Typ: libkb.DBSavedContacts,
		Key: fmt.Sprintf("%v", uid),
	}
}

type savedContactsCache struct {
	Contacts []keybase1.ProcessedContact
	Version  int
}

const savedContactsCurrentVer = 1

func assertionToNameDbKey(uid keybase1.UID) libkb.DbKey {
	return libkb.DbKey{
		Typ: libkb.DBSavedContacts,
		Key: fmt.Sprintf("lookup:%v", uid),
	}
}

type assertionToNameCache struct {
	AssertionToName map[string]string
	Version         int
}

const assertionToNameCurrentVer = 1

func ResolveAndSaveContacts(mctx libkb.MetaContext, provider ContactsProvider, contacts []keybase1.Contact) (newlyResolved []keybase1.ProcessedContact, err error) {
	resolveResults, err := ResolveContacts(mctx, provider, contacts, keybase1.RegionCode(""))
	if err != nil {
		return nil, err
	}

	// find newly resolved
	s := mctx.G().SyncedContactList
	currentContacts, err := s.RetrieveContacts(mctx)
	if err == nil {
		unres := make(map[string]struct{})
		for _, contact := range currentContacts {
			if !contact.Resolved {
				unres[contact.Assertion] = struct{}{}
			}
		}

		for _, result := range resolveResults {
			if _, ok := unres[result.Assertion]; ok && result.Resolved {
				newlyResolved = append(newlyResolved, result)
			}
		}
	} else {
		mctx.Warning("error retrieving synced contacts; continuing: %s", err)
	}

	if len(newlyResolved) > 0 {
		resolutionsForPeoplePage := make([]ContactResolution, len(newlyResolved))
		for i, contact := range newlyResolved {
			resolutionsForPeoplePage[i] = ContactResolution{
				Description: fmt.Sprintf("%s (%s)", contact.ContactName,
					contact.Component.ValueString()),
				ResolvedUser: keybase1.User{
					Uid:      contact.Uid,
					Username: contact.Username,
				},
			}
		}
		err = SendEncryptedContactResolutionToServer(mctx, resolutionsForPeoplePage)
		if err != nil {
			mctx.Warning("Could not add resolved contacts to people page: %v; returning contacts anyway", err)
		}
	}

	return newlyResolved, s.SaveProcessedContacts(mctx, resolveResults)
}

func makeAssertionToName(contacts []keybase1.ProcessedContact) (res map[string]string) {
	res = make(map[string]string)
	toRemove := make(map[string]struct{})
	for _, contact := range contacts {
		if _, ok := res[contact.Assertion]; ok {
			// multiple contacts match this assertion, remove once we're done
			toRemove[contact.Assertion] = struct{}{}
			continue
		}
		res[contact.Assertion] = contact.ContactName
	}
	for remove := range toRemove {
		delete(res, remove)
	}
	return res
}

func (s *SavedContactsStore) SaveProcessedContacts(mctx libkb.MetaContext, contacts []keybase1.ProcessedContact) (err error) {
	val := savedContactsCache{
		Contacts: contacts,
		Version:  savedContactsCurrentVer,
	}

	cacheKey := savedContactsDbKey(mctx.CurrentUID())
	err = s.encryptedDB.Put(mctx.Ctx(), cacheKey, val)
	if err != nil {
		return err
	}

	assertionToName := makeAssertionToName(contacts)
	lookupVal := assertionToNameCache{
		AssertionToName: assertionToName,
		Version:         assertionToNameCurrentVer,
	}

	cacheKey = assertionToNameDbKey(mctx.CurrentUID())
	err = s.encryptedDB.Put(mctx.Ctx(), cacheKey, lookupVal)
	return err
}

func (s *SavedContactsStore) RetrieveContacts(mctx libkb.MetaContext) (ret []keybase1.ProcessedContact, err error) {
	cacheKey := savedContactsDbKey(mctx.CurrentUID())
	var cache savedContactsCache
	found, err := s.encryptedDB.Get(mctx.Ctx(), cacheKey, &cache)
	if err != nil {
		return nil, err
	}
	if !found {
		return ret, nil
	}
	if cache.Version != savedContactsCurrentVer {
		mctx.Warning("synced contact list found but had an old version (found: %d, need: %d), returning empty list",
			cache.Version, savedContactsCurrentVer)
		return ret, nil
	}
	return cache.Contacts, nil
}

func (s *SavedContactsStore) RetrieveAssertionToName(mctx libkb.MetaContext) (ret map[string]string, err error) {
	cacheKey := assertionToNameDbKey(mctx.CurrentUID())
	var cache assertionToNameCache
	found, err := s.encryptedDB.Get(mctx.Ctx(), cacheKey, &cache)
	if err != nil {
		return nil, err
	}
	if !found {
		return ret, nil
	}
	if cache.Version != assertionToNameCurrentVer {
		mctx.Warning("assertion to name found but had an old version (found: %d, need: %d), returning empty map",
			cache.Version, assertionToNameCurrentVer)
		return ret, nil
	}
	return cache.AssertionToName, nil
}

func (s *SavedContactsStore) UnresolveContactsWithComponent(mctx libkb.MetaContext,
	phoneNumber *keybase1.PhoneNumber, email *keybase1.EmailAddress) {
	// TODO: Use a phoneNumber | email variant instead of two pointers.
	contactList, err := s.RetrieveContacts(mctx)
	if err != nil {
		mctx.Warning("Failed to get cached contact list: %x", err)
		return
	}
	for i, con := range contactList {
		var unresolve bool
		switch {
		case phoneNumber != nil && con.Component.PhoneNumber != nil:
			unresolve = *con.Component.PhoneNumber == keybase1.RawPhoneNumber(*phoneNumber)
		case email != nil && con.Component.Email != nil:
			unresolve = *con.Component.Email == *email
		}

		if unresolve {
			// Unresolve contact.
			con.Resolved = false
			con.Username = ""
			con.Uid = ""
			con.Following = false
			con.FullName = ""
			// TODO: DisplayName/DisplayLabel logic infects yet another file /
			// module. But it will sort itself out once we get rid of both.
			con.DisplayName = con.ContactName
			con.DisplayLabel = con.Component.FormatDisplayLabel(false /* addLabel */)
			contactList[i] = con
		}
	}
	err = s.SaveProcessedContacts(mctx, contactList)
	if err != nil {
		mctx.Warning("Failed to put cached contact list: %x", err)
	}
}
<commit_msg>get rid of double parens in contact resolution notifications (#19514)<commit_after>package contacts

import (
	"context"
	"fmt"

	"github.com/keybase/client/go/encrypteddb"

	"github.com/keybase/client/go/libkb"
	"github.com/keybase/client/go/protocol/keybase1"
)

// Saving contact list into encrypted db.

// Cache resolutions of a lookup ran on entire contact list provided by the
// frontend. Assume every time SaveContacts is called, entire contact list is
// passed as an argument. Always cache the result of last resolution, do not do
// any result merging.

type SavedContactsStore struct {
	encryptedDB *encrypteddb.EncryptedDB
}

var _ libkb.SyncedContactListProvider = (*SavedContactsStore)(nil)

// NewSavedContactsStore creates a new SavedContactsStore for global context.
// The store is used to securely store list of resolved contacts.
func NewSavedContactsStore(g *libkb.GlobalContext) *SavedContactsStore {
	keyFn := func(ctx context.Context) ([32]byte, error) {
		return encrypteddb.GetSecretBoxKey(ctx, g, encrypteddb.DefaultSecretUI,
			libkb.EncryptionReasonContactsLocalStorage, "encrypting local contact list")
	}
	dbFn := func(g *libkb.GlobalContext) *libkb.JSONLocalDb {
		return g.LocalDb
	}
	return &SavedContactsStore{
		encryptedDB: encrypteddb.New(g, dbFn, keyFn),
	}
}

func ServiceInit(g *libkb.GlobalContext) {
	g.SyncedContactList = NewSavedContactsStore(g)
}

func savedContactsDbKey(uid keybase1.UID) libkb.DbKey {
	return libkb.DbKey{
		Typ: libkb.DBSavedContacts,
		Key: fmt.Sprintf("%v", uid),
	}
}

type savedContactsCache struct {
	Contacts []keybase1.ProcessedContact
	Version  int
}

const savedContactsCurrentVer = 1

func assertionToNameDbKey(uid keybase1.UID) libkb.DbKey {
	return libkb.DbKey{
		Typ: libkb.DBSavedContacts,
		Key: fmt.Sprintf("lookup:%v", uid),
	}
}

type assertionToNameCache struct {
	AssertionToName map[string]string
	Version         int
}

const assertionToNameCurrentVer = 1

func ResolveAndSaveContacts(mctx libkb.MetaContext, provider ContactsProvider, contacts []keybase1.Contact) (newlyResolved []keybase1.ProcessedContact, err error) {
	resolveResults, err := ResolveContacts(mctx, provider, contacts, keybase1.RegionCode(""))
	if err != nil {
		return nil, err
	}

	// find newly resolved
	s := mctx.G().SyncedContactList
	currentContacts, err := s.RetrieveContacts(mctx)
	if err == nil {
		unres := make(map[string]struct{})
		for _, contact := range currentContacts {
			if !contact.Resolved {
				unres[contact.Assertion] = struct{}{}
			}
		}

		for _, result := range resolveResults {
			if _, ok := unres[result.Assertion]; ok && result.Resolved {
				newlyResolved = append(newlyResolved, result)
			}
		}
	} else {
		mctx.Warning("error retrieving synced contacts; continuing: %s", err)
	}

	if len(newlyResolved) > 0 {
		resolutionsForPeoplePage := make([]ContactResolution, len(newlyResolved))
		for i, contact := range newlyResolved {
			resolutionsForPeoplePage[i] = ContactResolution{
				Description: fmt.Sprintf("%s — %s", contact.ContactName,
					contact.Component.ValueString()),
				ResolvedUser: keybase1.User{
					Uid:      contact.Uid,
					Username: contact.Username,
				},
			}
		}
		err = SendEncryptedContactResolutionToServer(mctx, resolutionsForPeoplePage)
		if err != nil {
			mctx.Warning("Could not add resolved contacts to people page: %v; returning contacts anyway", err)
		}
	}

	return newlyResolved, s.SaveProcessedContacts(mctx, resolveResults)
}

func makeAssertionToName(contacts []keybase1.ProcessedContact) (res map[string]string) {
	res = make(map[string]string)
	toRemove := make(map[string]struct{})
	for _, contact := range contacts {
		if _, ok := res[contact.Assertion]; ok {
			// multiple contacts match this assertion, remove once we're done
			toRemove[contact.Assertion] = struct{}{}
			continue
		}
		res[contact.Assertion] = contact.ContactName
	}
	for remove := range toRemove {
		delete(res, remove)
	}
	return res
}

func (s *SavedContactsStore) SaveProcessedContacts(mctx libkb.MetaContext, contacts []keybase1.ProcessedContact) (err error) {
	val := savedContactsCache{
		Contacts: contacts,
		Version:  savedContactsCurrentVer,
	}

	cacheKey := savedContactsDbKey(mctx.CurrentUID())
	err = s.encryptedDB.Put(mctx.Ctx(), cacheKey, val)
	if err != nil {
		return err
	}

	assertionToName := makeAssertionToName(contacts)
	lookupVal := assertionToNameCache{
		AssertionToName: assertionToName,
		Version:         assertionToNameCurrentVer,
	}

	cacheKey = assertionToNameDbKey(mctx.CurrentUID())
	err = s.encryptedDB.Put(mctx.Ctx(), cacheKey, lookupVal)
	return err
}

func (s *SavedContactsStore) RetrieveContacts(mctx libkb.MetaContext) (ret []keybase1.ProcessedContact, err error) {
	cacheKey := savedContactsDbKey(mctx.CurrentUID())
	var cache savedContactsCache
	found, err := s.encryptedDB.Get(mctx.Ctx(), cacheKey, &cache)
	if err != nil {
		return nil, err
	}
	if !found {
		return ret, nil
	}
	if cache.Version != savedContactsCurrentVer {
		mctx.Warning("synced contact list found but had an old version (found: %d, need: %d), returning empty list",
			cache.Version, savedContactsCurrentVer)
		return ret, nil
	}
	return cache.Contacts, nil
}

func (s *SavedContactsStore) RetrieveAssertionToName(mctx libkb.MetaContext) (ret map[string]string, err error) {
	cacheKey := assertionToNameDbKey(mctx.CurrentUID())
	var cache assertionToNameCache
	found, err := s.encryptedDB.Get(mctx.Ctx(), cacheKey, &cache)
	if err != nil {
		return nil, err
	}
	if !found {
		return ret, nil
	}
	if cache.Version != assertionToNameCurrentVer {
		mctx.Warning("assertion to name found but had an old version (found: %d, need: %d), returning empty map",
			cache.Version, assertionToNameCurrentVer)
		return ret, nil
	}
	return cache.AssertionToName, nil
}

func (s *SavedContactsStore) UnresolveContactsWithComponent(mctx libkb.MetaContext,
	phoneNumber *keybase1.PhoneNumber, email *keybase1.EmailAddress) {
	// TODO: Use a phoneNumber | email variant instead of two pointers.
	contactList, err := s.RetrieveContacts(mctx)
	if err != nil {
		mctx.Warning("Failed to get cached contact list: %x", err)
		return
	}
	for i, con := range contactList {
		var unresolve bool
		switch {
		case phoneNumber != nil && con.Component.PhoneNumber != nil:
			unresolve = *con.Component.PhoneNumber == keybase1.RawPhoneNumber(*phoneNumber)
		case email != nil && con.Component.Email != nil:
			unresolve = *con.Component.Email == *email
		}

		if unresolve {
			// Unresolve contact.
			con.Resolved = false
			con.Username = ""
			con.Uid = ""
			con.Following = false
			con.FullName = ""
			// TODO: DisplayName/DisplayLabel logic infects yet another file /
			// module. But it will sort itself out once we get rid of both.
			con.DisplayName = con.ContactName
			con.DisplayLabel = con.Component.FormatDisplayLabel(false /* addLabel */)
			contactList[i] = con
		}
	}
	err = s.SaveProcessedContacts(mctx, contactList)
	if err != nil {
		mctx.Warning("Failed to put cached contact list: %x", err)
	}
}
<|endoftext|>