<commit_before>package romulus

import (
	"fmt"
	"net/url"
	"strings"

	"code.google.com/p/go-uuid/uuid"

	"github.com/GoogleCloudPlatform/kubernetes/pkg/api"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/client"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/fields"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/labels"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/runtime"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/watch"
)

// EtcdPeerList is just a slice of etcd peers
type EtcdPeerList []string

// KubeClientConfig is an alias for kubernetes/pkg/client.Config
type KubeClientConfig client.Config

// ServiceSelector is a map of labels for selecting services
type ServiceSelector map[string]string

func (s ServiceSelector) fixNamespace() ServiceSelector {
	ss := make(map[string]string, len(s))
	for k := range s {
		key := k
		if !strings.HasPrefix(k, "romulus/") {
			key = fmt.Sprintf("romulus/%s", key)
		}
		ss[key] = s[k]
	}
	return ServiceSelector(ss)
}

func formatVulcanNamespace(v string) string {
	return fmt.Sprintf("/%s", strings.Trim((string)(v), "/"))
}

// Config is used to configure the Registrar
type Config struct {
	PeerList            EtcdPeerList
	KubeConfig          KubeClientConfig
	APIVersion          string
	Selector            ServiceSelector
	VulcanEtcdNamespace string
}

func (c *Config) kc() client.Config { return (client.Config)(c.KubeConfig) }
func (c *Config) ps() []string      { return ([]string)(c.PeerList) }

func (sl ServiceSelector) String() string {
	s := []string{}
	for k, v := range sl {
		s = append(s, strings.Join([]string{k, v}, "="))
	}
	return strings.Join(s, ", ")
}

// Registrar holds the kubernetes/pkg/client.Client and etcd.Client
type Registrar struct {
	k  *client.Client
	e  EtcdClient
	vk string
	v  string
	s  ServiceSelector
}

// NewRegistrar returns a ptr to a new Registrar from a Config
func NewRegistrar(c *Config) (*Registrar, error) {
	cf := c.kc()
	cl, err := client.New(&cf)
	if err != nil {
		return nil, err
	}
	return &Registrar{
		e:  NewEtcdClient(c.ps()),
		k:  cl,
		v:  c.APIVersion,
		s:  c.Selector,
		vk: formatVulcanNamespace(c.VulcanEtcdNamespace),
	}, nil
}

func (c *Registrar) endpointsEventChannel() (watch.Interface, error) {
	return c.k.Endpoints(api.NamespaceAll).Watch(labels.Everything(), fields.Everything(), "")
}

func (c *Registrar) serviceEventsChannel() (watch.Interface, error) {
	return c.k.Services(api.NamespaceAll).Watch(labels.Everything(), fields.Everything(), "")
}

func (c *Registrar) getService(name, ns string) (*api.Service, error) {
	s, e := c.k.Services(ns).Get(name)
	if e != nil || s == nil {
		return nil, Error{fmt.Sprintf("Unable to get service %q", name), e}
	}
	return s, nil
}

func (c *Registrar) getEndpoint(name, ns string) (*api.Endpoints, error) {
	en, e := c.k.Endpoints(ns).Get(name)
	if e != nil || en == nil {
		return nil, Error{fmt.Sprintf("Unable to get endpoint %q", name), e}
	}
	return en, nil
}

func (c *Registrar) pruneServers(bid uuid.UUID, sm ServerMap) error {
	k := fmt.Sprintf(srvrDirFmt, bid.String())
	ips, e := c.e.Keys(k)
	if e != nil {
		if isKeyNotFound(e) {
			return nil
		}
		return NewErr(e, "etcd error")
	}

	logf(fi{"servers": ips, "bcknd-id": bid.String()}).Debug("Gathered servers from etcd")
	for _, ip := range ips {
		if _, ok := sm[ip]; !ok {
			log().Debugf("Removing %s from etcd", ip)
			key := fmt.Sprintf("%s/%s", k, ip)
			if e := c.e.Del(key); e != nil {
				return Error{"etcd error", e}
			}
		}
	}
	return nil
}

func (reg *Registrar) handleDelete(r runtime.Object) error {
	switch o := r.(type) {
	case *api.Endpoints:
		return deregister(reg, o.ObjectMeta, false)
	case *api.Service:
		return deregister(reg, o.ObjectMeta, true)
	default:
		return NewErr(nil, "Unsupported api object: %v", r)
	}
}

func (reg *Registrar) handleUpdate(r runtime.Object) error {
	switch o := r.(type) {
	case *api.Service:
		return nil
	case *api.Endpoints:
		return register(reg, o)
	default:
		return NewErr(nil, "Unsupported api object: %v", r)
	}
}

func do(r *Registrar, e watch.Event) error {
	logf(fi{"event": e.Type}).Debug("Got a Kubernetes API event")
	switch e.Type {
	default:
		log().Debugf("Unsupported event type %q", e.Type)
		return nil
	case watch.Error:
		if a, ok := e.Object.(*api.Status); ok {
			e := fmt.Errorf("[%d] %v", a.Code, a.Reason)
			return Error{fmt.Sprintf("Kubernetes API failure: %s", a.Message), e}
		}
		return Error{"Unknown kubernetes api error", nil}
	case watch.Deleted:
		return r.handleDelete(e.Object)
	case watch.Added, watch.Modified:
		return r.handleUpdate(e.Object)
	}
}

func expandEndpoints(bid uuid.UUID, e *api.Endpoints) ServerMap {
	sm := ServerMap{}
	for _, es := range e.Subsets {
		for _, port := range es.Ports {
			if port.Protocol != api.ProtocolTCP {
				logf(fi{"bcknd-id": bid.String()}).Warnf("Unsupported protocol: %s", port.Protocol)
				continue
			}

			// TODO: Do we want to force ports to have a name?
			// if port.Name != "vulcan" {
			// 	log().Debugf("Not registering port %d", port.Port)
			// 	continue
			// }

			for _, ip := range es.Addresses {
				ur := fmt.Sprintf("http://%s:%d", ip.IP, port.Port)
				u, err := url.Parse(ur)
				if err != nil {
					logf(fi{"bcknd-id": bid.String()}).Warnf("Bad URL: %s", ur)
					continue
				}
				uu := (*URL)(u)
				sm[uu.GetHost()] = Server{
					Backend: bid,
					URL:     uu,
				}
			}
		}
	}
	return sm
}

func getUUID(o api.ObjectMeta) uuid.UUID {
	return uuid.Parse((string)(o.UID))
}

func registerable(s *api.Service, sl ServiceSelector) bool {
	for k, v := range sl.fixNamespace() {
		if sv, ok := s.Labels[k]; !ok || sv != v {
			if sv, ok := s.Annotations[k]; !ok || sv != v {
				return false
			}
		}
	}
	return api.IsServiceIPSet(s)
}
<commit_msg>Registrar updates<commit_after>package romulus

import (
	"fmt"
	"net/url"
	"strings"

	"code.google.com/p/go-uuid/uuid"

	"github.com/GoogleCloudPlatform/kubernetes/pkg/api"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/client"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/fields"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/labels"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/runtime"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/watch"
)

// EtcdPeerList is just a slice of etcd peers
type EtcdPeerList []string

// KubeClientConfig is an alias for kubernetes/pkg/client.Config
type KubeClientConfig client.Config

// ServiceSelector is a map of labels for selecting services
type ServiceSelector map[string]string

func (s ServiceSelector) fixNamespace() ServiceSelector {
	ss := make(map[string]string, len(s))
	for k := range s {
		key := k
		if !strings.HasPrefix(k, "romulus/") {
			key = fmt.Sprintf("romulus/%s", key)
		}
		ss[key] = s[k]
	}
	return ServiceSelector(ss)
}

func formatEtcdNamespace(v string) string {
	return fmt.Sprintf("/%s", strings.Trim(v, "/"))
}

// Config is used to configure the Registrar
type Config struct {
	PeerList             EtcdPeerList
	KubeConfig           KubeClientConfig
	APIVersion           string
	Selector             ServiceSelector
	VulcanEtcdNamespace  string
	RomulusEtcdNamespace string
}

func (c *Config) kc() client.Config { return (client.Config)(c.KubeConfig) }
func (c *Config) ps() []string      { return ([]string)(c.PeerList) }

func (sl ServiceSelector) String() string {
	s := []string{}
	for k, v := range sl {
		s = append(s, strings.Join([]string{k, v}, "="))
	}
	return strings.Join(s, ", ")
}

// Registrar holds the kubernetes/pkg/client.Client and etcd.Client
type Registrar struct {
	k  *client.Client
	e  EtcdClient
	vk string
	rk string
	v  string
	s  ServiceSelector
}

// NewRegistrar returns a ptr to a new Registrar from a Config
func NewRegistrar(c *Config) (*Registrar, error) {
	cf := c.kc()
	cl, err := client.New(&cf)
	if err != nil {
		return nil, err
	}
	return &Registrar{
		e:  NewEtcdClient(c.ps()),
		k:  cl,
		v:  c.APIVersion,
		s:  c.Selector,
		vk: formatEtcdNamespace(c.VulcanEtcdNamespace),
		rk: formatEtcdNamespace(c.RomulusEtcdNamespace),
	}, nil
}

func (c *Registrar) initEndpoints() (watch.Interface, error) {
	kf := "%s/backends"
	ids, err := c.e.Keys(fmt.Sprintf(kf, c.rk))
	if err != nil {
		return nil, NewErr(err, "etcd error")
	}
	for _, id := range ids {
		c.k.Endpoints(api.NamespaceAll)
	}

	return c.k.Endpoints(api.NamespaceAll).Watch(labels.Everything(), fields.Everything(), "")
}

func (c *Registrar) serviceEventsChannel() (watch.Interface, error) {
	return c.k.Services(api.NamespaceAll).Watch(labels.Everything(), fields.Everything(), "")
}

func (c *Registrar) getService(name, ns string) (*api.Service, error) {
	s, e := c.k.Services(ns).Get(name)
	if e != nil || s == nil {
		return nil, Error{fmt.Sprintf("Unable to get service %q", name), e}
	}
	return s, nil
}

func (c *Registrar) getEndpoint(name, ns string) (*api.Endpoints, error) {
	en, e := c.k.Endpoints(ns).Get(name)
	if e != nil || en == nil {
		return nil, Error{fmt.Sprintf("Unable to get endpoint %q", name), e}
	}
	return en, nil
}

func (c *Registrar) pruneServers(bid uuid.UUID, sm ServerMap) error {
	k := fmt.Sprintf(srvrDirFmt, bid.String())
	ips, e := c.e.Keys(k)
	if e != nil {
		if isKeyNotFound(e) {
			return nil
		}
		return NewErr(e, "etcd error")
	}

	logf(fi{"servers": ips, "bcknd-id": bid.String()}).Debug("Gathered servers from etcd")
	for _, ip := range ips {
		if _, ok := sm[ip]; !ok {
			log().Debugf("Removing %s from etcd", ip)
			key := fmt.Sprintf("%s/%s", k, ip)
			if e := c.e.Del(key); e != nil {
				return Error{"etcd error", e}
			}
		}
	}
	return nil
}

func (reg *Registrar) handleDelete(r runtime.Object) error {
	switch o := r.(type) {
	case *api.Endpoints:
		return deregister(reg, o.ObjectMeta, false)
	case *api.Service:
		return deregister(reg, o.ObjectMeta, true)
	default:
		return NewErr(nil, "Unsupported api object: %v", r)
	}
}

func (reg *Registrar) handleUpdate(r runtime.Object) error {
	switch o := r.(type) {
	case *api.Service:
		return nil
	case *api.Endpoints:
		return register(reg, o)
	default:
		return NewErr(nil, "Unsupported api object: %v", r)
	}
}

func do(r *Registrar, e watch.Event) error {
	logf(fi{"event": e.Type}).Debug("Got a Kubernetes API event")
	switch e.Type {
	default:
		log().Debugf("Unsupported event type %q", e.Type)
		return nil
	case watch.Error:
		if a, ok := e.Object.(*api.Status); ok {
			e := fmt.Errorf("[%d] %v", a.Code, a.Reason)
			return Error{fmt.Sprintf("Kubernetes API failure: %s", a.Message), e}
		}
		return Error{"Unknown kubernetes api error", nil}
	case watch.Deleted:
		return r.handleDelete(e.Object)
	case watch.Added, watch.Modified:
		return r.handleUpdate(e.Object)
	}
}

func expandEndpoints(bid uuid.UUID, e *api.Endpoints) ServerMap {
	sm := ServerMap{}
	for _, es := range e.Subsets {
		for _, port := range es.Ports {
			if port.Protocol != api.ProtocolTCP {
				logf(fi{"bcknd-id": bid.String()}).Warnf("Unsupported protocol: %s", port.Protocol)
				continue
			}

			// TODO: Do we want to force ports to have a name?
			// if port.Name != "vulcan" {
			// 	log().Debugf("Not registering port %d", port.Port)
			// 	continue
			// }

			for _, ip := range es.Addresses {
				ur := fmt.Sprintf("http://%s:%d", ip.IP, port.Port)
				u, err := url.Parse(ur)
				if err != nil {
					logf(fi{"bcknd-id": bid.String()}).Warnf("Bad URL: %s", ur)
					continue
				}
				uu := (*URL)(u)
				sm[uu.GetHost()] = Server{
					Backend: bid,
					URL:     uu,
				}
			}
		}
	}
	return sm
}

func getUUID(o api.ObjectMeta) uuid.UUID {
	return uuid.Parse((string)(o.UID))
}

func registerable(s *api.Service, sl ServiceSelector) bool {
	for k, v := range sl.fixNamespace() {
		if sv, ok := s.Labels[k]; !ok || sv != v {
			if sv, ok := s.Annotations[k]; !ok || sv != v {
				return false
			}
		}
	}
	return api.IsServiceIPSet(s)
}
<|endoftext|>


<commit_before>// roshi-server provides a REST-y HTTP service to interact with a farm.
package main

import (
	"encoding/json"
	_ "expvar"
	"flag"
	"fmt"
	logpkg "log"
	"net/http"
	_ "net/http/pprof"
	"net/url"
	"os"
	"runtime"
	"sort"
	"strconv"
	"strings"
	"time"

	"github.com/soundcloud/roshi/cluster"
	"github.com/soundcloud/roshi/common"
	"github.com/soundcloud/roshi/farm"
	"github.com/soundcloud/roshi/instrumentation/statsd"
	"github.com/soundcloud/roshi/shard"
	"github.com/soundcloud/roshi/vendor/g2s"
	"github.com/soundcloud/roshi/vendor/handy/breaker"
	"github.com/soundcloud/roshi/vendor/pat"
)

var (
	stats = g2s.Noop()
	log   = logpkg.New(os.Stdout, "", logpkg.Lmicroseconds)
)

func main() {
	var (
		redisInstances            = flag.String("redis.instances", "", "Semicolon-separated list of comma-separated lists of Redis instances")
		redisConnectTimeout       = flag.Duration("redis.connect.timeout", 3*time.Second, "Redis connect timeout")
		redisReadTimeout          = flag.Duration("redis.read.timeout", 3*time.Second, "Redis read timeout")
		redisWriteTimeout         = flag.Duration("redis.write.timeout", 3*time.Second, "Redis write timeout")
		redisMCPI                 = flag.Int("redis.mcpi", 10, "Max connections per Redis instance")
		redisHash                 = flag.String("redis.hash", "murmur3", "Redis hash function: murmur3, fnv, fnva")
		redisReadStrategy         = flag.String("redis.read.strategy", "SendAllReadAll", "Redis read strategy: SendAllReadAll, SendOneReadOne, SendAllReadFirstLinger, SendVarReadFirstLinger")
		redisReadThresholdRate    = flag.Int("redis.read.threshold.rate", 10, "Baseline SendAll reads per sec, additional reads are SendOne (SendVarReadFirstLinger strategy only)")
		redisReadThresholdLatency = flag.Duration("redis.read.threshold.latency", 50*time.Millisecond, "If a SendOne read has not returned anything after this latency, it's promoted to SendAll (SendVarReadFirstLinger strategy only)")
		redisRepairer             = flag.String("redis.repairer", "RateLimited", "Redis repairer: RateLimited, Nop")
		redisMaxRepairRate        = flag.Int("redis.repair.maxrate", 10, "Max repairs per second (RateLimited repairer only)")
		redisMaxRepairBacklog     = flag.Int("redis.repair.maxbacklog", 100000, "Max number of queued repairs (RateLimited repairer only)")
		maxSize                   = flag.Int("max.size", 10000, "Maximum number of events per key")
		statsdAddress             = flag.String("statsd.address", "", "Statsd address (blank to disable)")
		statsdSampleRate          = flag.Float64("statsd.sample.rate", 0.1, "Statsd sample rate for normal metrics")
		statsdBucketPrefix        = flag.String("statsd.bucket.prefix", "myservice.", "Statsd bucket key prefix, including trailing period")
		httpCircuitBreaker        = flag.Bool("http.circuit.breaker", true, "Enable HTTP server circuit breaker")
		httpAddress               = flag.String("http.address", ":6302", "HTTP listen address")
	)
	flag.Parse()
	log.Printf("GOMAXPROCS %d", runtime.GOMAXPROCS(-1))

	// Set up statsd instrumentation, if it's specified.
	if *statsdAddress != "" {
		var err error
		stats, err = g2s.Dial("udp", *statsdAddress)
		if err != nil {
			log.Fatal(err)
		}
	}

	// Parse read strategy.
	var readStrategy farm.ReadStrategy
	switch strings.ToLower(*redisReadStrategy) {
	case "sendallreadall":
		readStrategy = farm.SendAllReadAll
	case "sendonereadone":
		readStrategy = farm.SendOneReadOne
	case "sendallreadfirstlinger":
		readStrategy = farm.SendAllReadFirstLinger
	case "sendvarreadfirstlinger":
		readStrategy = farm.SendVarReadFirstLinger(*redisReadThresholdRate, *redisReadThresholdLatency)
	default:
		log.Fatalf("unknown read strategy '%s'", *redisReadStrategy)
	}
	log.Printf("using %s read strategy", *redisReadStrategy)

	// Parse repairer.
	var repairer farm.Repairer
	switch strings.ToLower(*redisRepairer) {
	case "nop":
		repairer = farm.NopRepairer
	case "ratelimited":
		repairer = farm.RateLimitedRepairer(*redisMaxRepairRate, *redisMaxRepairBacklog)
	default:
		log.Fatalf("unknown repairer '%s'", *redisRepairer)
	}

	// Parse hash function.
	var hashFunc func(string) uint32
	switch strings.ToLower(*redisHash) {
	case "murmur3":
		hashFunc = shard.Murmur3
	case "fnv":
		hashFunc = shard.FNV
	case "fnva":
		hashFunc = shard.FNVa
	default:
		log.Fatalf("unknown hash '%s'", *redisHash)
	}

	// Build the farm.
	farm, err := newFarm(
		*redisInstances,
		*redisConnectTimeout, *redisReadTimeout, *redisWriteTimeout,
		*redisMCPI,
		hashFunc,
		readStrategy,
		repairer,
		*maxSize,
		*statsdSampleRate,
		*statsdBucketPrefix,
	)
	if err != nil {
		log.Fatal(err)
	}

	// Build the HTTP server.
	r := pat.New()
	r.Add("GET", "/debug", http.DefaultServeMux)
	r.Get("/", handleSelect(farm))
	r.Post("/", handleInsert(farm))
	r.Delete("/", handleDelete(farm))
	h := http.Handler(r)
	if *httpCircuitBreaker {
		log.Printf("using HTTP circuit breaker")
		h = breaker.DefaultBreaker(h)
	}

	// Go for it.
	log.Printf("listening on %s", *httpAddress)
	log.Fatal(http.ListenAndServe(*httpAddress, h))
}

func newFarm(
	redisInstances string,
	connectTimeout, readTimeout, writeTimeout time.Duration,
	redisMCPI int,
	hash func(string) uint32,
	readStrategy farm.ReadStrategy,
	repairer farm.Repairer,
	maxSize int,
	statsdSampleRate float64,
	bucketPrefix string,
) (*farm.Farm, error) {
	instr := statsd.New(stats, float32(statsdSampleRate), bucketPrefix)

	clusters := []cluster.Cluster{}
	for i, clusterInstances := range strings.Split(redisInstances, ";") {
		addresses := stripBlank(strings.Split(clusterInstances, ","))
		if len(addresses) <= 0 {
			continue
		}
		clusters = append(clusters, cluster.New(
			shard.New(
				addresses,
				connectTimeout, readTimeout, writeTimeout,
				redisMCPI,
				hash,
			),
			maxSize,
			instr,
		))
		log.Printf("Redis cluster %d: %d instance(s)", i+1, len(addresses))
	}
	if len(clusters) <= 0 {
		return nil, fmt.Errorf("no cluster(s)")
	}

	return farm.New(
		clusters,
		readStrategy,
		repairer,
		instr,
	), nil
}

func handleSelect(selecter farm.Selecter) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		began := time.Now()

		if err := r.ParseForm(); err != nil {
			respondError(w, r.Method, r.URL.String(), http.StatusInternalServerError, err)
			return
		}
		offset := parseInt(r.Form, "offset", 0)
		limit := parseInt(r.Form, "limit", 10)
		coalesce := parseBool(r.Form, "coalesce", false)

		var keys [][]byte
		defer r.Body.Close()
		if err := json.NewDecoder(r.Body).Decode(&keys); err != nil {
			respondError(w, r.Method, r.URL.String(), http.StatusBadRequest, err)
			return
		}

		keyStrings := make([]string, len(keys))
		for i := range keys {
			keyStrings[i] = string(keys[i])
		}

		var records interface{}
		if coalesce {
			// We need to Select from 0 to offset+limit, flatten the map to a
			// single ordered slice, and then cut off the last limit elements.
			m, err := selecter.Select(keyStrings, 0, offset+limit)
			if err != nil {
				respondError(w, r.Method, r.URL.String(), http.StatusInternalServerError, err)
				return
			}
			records = flatten(m, offset, limit)
		} else {
			// We can directly Select using the given offset and limit.
			m, err := selecter.Select(keyStrings, offset, limit)
			if err != nil {
				respondError(w, r.Method, r.URL.String(), http.StatusInternalServerError, err)
				return
			}
			records = m
		}

		respondSelected(w, keys, offset, limit, records, time.Since(began))
	}
}

func handleInsert(inserter cluster.Inserter) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		began := time.Now()

		var tuples []common.KeyScoreMember
		if err := json.NewDecoder(r.Body).Decode(&tuples); err != nil {
			respondError(w, r.Method, r.URL.String(), http.StatusBadRequest, err)
			return
		}

		if err := inserter.Insert(tuples); err != nil {
			respondError(w, r.Method, r.URL.String(), http.StatusInternalServerError, err)
			return
		}

		respondInserted(w, len(tuples), time.Since(began))
	}
}

func handleDelete(deleter cluster.Deleter) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		began := time.Now()

		var tuples []common.KeyScoreMember
		if err := json.NewDecoder(r.Body).Decode(&tuples); err != nil {
			respondError(w, r.Method, r.URL.String(), http.StatusBadRequest, err)
			return
		}

		if err := deleter.Delete(tuples); err != nil {
			respondError(w, r.Method, r.URL.String(), http.StatusInternalServerError, err)
			return
		}

		respondDeleted(w, len(tuples), time.Since(began))
	}
}

func flatten(m map[string][]common.KeyScoreMember, offset, limit int) []common.KeyScoreMember {
	a := common.KeyScoreMembers{}
	for _, tuples := range m {
		a = append(a, tuples...)
	}

	sort.Sort(a)

	if len(a) < offset {
		return []common.KeyScoreMember{}
	}
	a = a[offset:]

	if len(a) > limit {
		a = a[:limit]
	}

	return a
}

func parseInt(values url.Values, key string, defaultValue int) int {
	value, err := strconv.ParseInt(values.Get(key), 10, 64)
	if err != nil {
		return defaultValue
	}
	return int(value)
}

func parseBool(values url.Values, key string, defaultValue bool) bool {
	value, err := strconv.ParseBool(values.Get(key))
	if err != nil {
		return defaultValue
	}
	return value
}

func respondInserted(w http.ResponseWriter, n int, duration time.Duration) {
	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(map[string]interface{}{
		"inserted": n,
		"duration": duration.String(),
	})
}

func respondSelected(w http.ResponseWriter, keys [][]byte, offset, limit int, records interface{}, duration time.Duration) {
	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(map[string]interface{}{
		"keys":     keys,
		"offset":   offset,
		"limit":    limit,
		"records":  records,
		"duration": duration.String(),
	})
}

func respondDeleted(w http.ResponseWriter, n int, duration time.Duration) {
	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(map[string]interface{}{
		"deleted":  n,
		"duration": duration.String(),
	})
}

func respondError(w http.ResponseWriter, method, url string, code int, err error) {
	log.Printf("%s %s: HTTP %d: %s", method, url, code, err)
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(code)
	json.NewEncoder(w).Encode(map[string]interface{}{
		"error":       err.Error(),
		"code":        code,
		"description": http.StatusText(code),
	})
}

func stripBlank(src []string) []string {
	dst := []string{}
	for _, s := range src {
		if s == "" {
			continue
		}
		dst = append(dst, s)
	}
	return dst
}

type writer struct{ *logpkg.Logger }

func (w writer) Write(p []byte) (int, error) { w.Print(string(p)); return len(p), nil }
<commit_msg>roshi-server: remove dead code<commit_after>// roshi-server provides a REST-y HTTP service to interact with a farm.
package main

import (
	"encoding/json"
	_ "expvar"
	"flag"
	"fmt"
	logpkg "log"
	"net/http"
	_ "net/http/pprof"
	"net/url"
	"os"
	"runtime"
	"sort"
	"strconv"
	"strings"
	"time"

	"github.com/soundcloud/roshi/cluster"
	"github.com/soundcloud/roshi/common"
	"github.com/soundcloud/roshi/farm"
	"github.com/soundcloud/roshi/instrumentation/statsd"
	"github.com/soundcloud/roshi/shard"
	"github.com/soundcloud/roshi/vendor/g2s"
	"github.com/soundcloud/roshi/vendor/handy/breaker"
	"github.com/soundcloud/roshi/vendor/pat"
)

var (
	stats = g2s.Noop()
	log   = logpkg.New(os.Stdout, "", logpkg.Lmicroseconds)
)

func main() {
	var (
		redisInstances            = flag.String("redis.instances", "", "Semicolon-separated list of comma-separated lists of Redis instances")
		redisConnectTimeout       = flag.Duration("redis.connect.timeout", 3*time.Second, "Redis connect timeout")
		redisReadTimeout          = flag.Duration("redis.read.timeout", 3*time.Second, "Redis read timeout")
		redisWriteTimeout         = flag.Duration("redis.write.timeout", 3*time.Second, "Redis write timeout")
		redisMCPI                 = flag.Int("redis.mcpi", 10, "Max connections per Redis instance")
		redisHash                 = flag.String("redis.hash", "murmur3", "Redis hash function: murmur3, fnv, fnva")
		redisReadStrategy         = flag.String("redis.read.strategy", "SendAllReadAll", "Redis read strategy: SendAllReadAll, SendOneReadOne, SendAllReadFirstLinger, SendVarReadFirstLinger")
		redisReadThresholdRate    = flag.Int("redis.read.threshold.rate", 10, "Baseline SendAll reads per sec, additional reads are SendOne (SendVarReadFirstLinger strategy only)")
		redisReadThresholdLatency = flag.Duration("redis.read.threshold.latency", 50*time.Millisecond, "If a SendOne read has not returned anything after this latency, it's promoted to SendAll (SendVarReadFirstLinger strategy only)")
		redisRepairer             = flag.String("redis.repairer", "RateLimited", "Redis repairer: RateLimited, Nop")
		redisMaxRepairRate        = flag.Int("redis.repair.maxrate", 10, "Max repairs per second (RateLimited repairer only)")
		redisMaxRepairBacklog     = flag.Int("redis.repair.maxbacklog", 100000, "Max number of queued repairs (RateLimited repairer only)")
		maxSize                   = flag.Int("max.size", 10000, "Maximum number of events per key")
		statsdAddress             = flag.String("statsd.address", "", "Statsd address (blank to disable)")
		statsdSampleRate          = flag.Float64("statsd.sample.rate", 0.1, "Statsd sample rate for normal metrics")
		statsdBucketPrefix        = flag.String("statsd.bucket.prefix", "myservice.", "Statsd bucket key prefix, including trailing period")
		httpCircuitBreaker        = flag.Bool("http.circuit.breaker", true, "Enable HTTP server circuit breaker")
		httpAddress               = flag.String("http.address", ":6302", "HTTP listen address")
	)
	flag.Parse()
	log.Printf("GOMAXPROCS %d", runtime.GOMAXPROCS(-1))

	// Set up statsd instrumentation, if it's specified.
	if *statsdAddress != "" {
		var err error
		stats, err = g2s.Dial("udp", *statsdAddress)
		if err != nil {
			log.Fatal(err)
		}
	}

	// Parse read strategy.
	var readStrategy farm.ReadStrategy
	switch strings.ToLower(*redisReadStrategy) {
	case "sendallreadall":
		readStrategy = farm.SendAllReadAll
	case "sendonereadone":
		readStrategy = farm.SendOneReadOne
	case "sendallreadfirstlinger":
		readStrategy = farm.SendAllReadFirstLinger
	case "sendvarreadfirstlinger":
		readStrategy = farm.SendVarReadFirstLinger(*redisReadThresholdRate, *redisReadThresholdLatency)
	default:
		log.Fatalf("unknown read strategy '%s'", *redisReadStrategy)
	}
	log.Printf("using %s read strategy", *redisReadStrategy)

	// Parse repairer.
	var repairer farm.Repairer
	switch strings.ToLower(*redisRepairer) {
	case "nop":
		repairer = farm.NopRepairer
	case "ratelimited":
		repairer = farm.RateLimitedRepairer(*redisMaxRepairRate, *redisMaxRepairBacklog)
	default:
		log.Fatalf("unknown repairer '%s'", *redisRepairer)
	}

	// Parse hash function.
	var hashFunc func(string) uint32
	switch strings.ToLower(*redisHash) {
	case "murmur3":
		hashFunc = shard.Murmur3
	case "fnv":
		hashFunc = shard.FNV
	case "fnva":
		hashFunc = shard.FNVa
	default:
		log.Fatalf("unknown hash '%s'", *redisHash)
	}

	// Build the farm.
	farm, err := newFarm(
		*redisInstances,
		*redisConnectTimeout, *redisReadTimeout, *redisWriteTimeout,
		*redisMCPI,
		hashFunc,
		readStrategy,
		repairer,
		*maxSize,
		*statsdSampleRate,
		*statsdBucketPrefix,
	)
	if err != nil {
		log.Fatal(err)
	}

	// Build the HTTP server.
	r := pat.New()
	r.Add("GET", "/debug", http.DefaultServeMux)
	r.Get("/", handleSelect(farm))
	r.Post("/", handleInsert(farm))
	r.Delete("/", handleDelete(farm))
	h := http.Handler(r)
	if *httpCircuitBreaker {
		log.Printf("using HTTP circuit breaker")
		h = breaker.DefaultBreaker(h)
	}

	// Go for it.
	log.Printf("listening on %s", *httpAddress)
	log.Fatal(http.ListenAndServe(*httpAddress, h))
}

func newFarm(
	redisInstances string,
	connectTimeout, readTimeout, writeTimeout time.Duration,
	redisMCPI int,
	hash func(string) uint32,
	readStrategy farm.ReadStrategy,
	repairer farm.Repairer,
	maxSize int,
	statsdSampleRate float64,
	bucketPrefix string,
) (*farm.Farm, error) {
	instr := statsd.New(stats, float32(statsdSampleRate), bucketPrefix)

	clusters := []cluster.Cluster{}
	for i, clusterInstances := range strings.Split(redisInstances, ";") {
		addresses := stripBlank(strings.Split(clusterInstances, ","))
		if len(addresses) <= 0 {
			continue
		}
		clusters = append(clusters, cluster.New(
			shard.New(
				addresses,
				connectTimeout, readTimeout, writeTimeout,
				redisMCPI,
				hash,
			),
			maxSize,
			instr,
		))
		log.Printf("Redis cluster %d: %d instance(s)", i+1, len(addresses))
	}
	if len(clusters) <= 0 {
		return nil, fmt.Errorf("no cluster(s)")
	}

	return farm.New(
		clusters,
		readStrategy,
		repairer,
		instr,
	), nil
}

func handleSelect(selecter farm.Selecter) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		began := time.Now()

		if err := r.ParseForm(); err != nil {
			respondError(w, r.Method, r.URL.String(), http.StatusInternalServerError, err)
			return
		}
		offset := parseInt(r.Form, "offset", 0)
		limit := parseInt(r.Form, "limit", 10)
		coalesce := parseBool(r.Form, "coalesce", false)

		var keys [][]byte
		defer r.Body.Close()
		if err := json.NewDecoder(r.Body).Decode(&keys); err != nil {
			respondError(w, r.Method, r.URL.String(), http.StatusBadRequest, err)
			return
		}

		keyStrings := make([]string, len(keys))
		for i := range keys {
			keyStrings[i] = string(keys[i])
		}

		var records interface{}
		if coalesce {
			// We need to Select from 0 to offset+limit, flatten the map to a
			// single ordered slice, and then cut off the last limit elements.
			m, err := selecter.Select(keyStrings, 0, offset+limit)
			if err != nil {
				respondError(w, r.Method, r.URL.String(), http.StatusInternalServerError, err)
				return
			}
			records = flatten(m, offset, limit)
		} else {
			// We can directly Select using the given offset and limit.
			m, err := selecter.Select(keyStrings, offset, limit)
			if err != nil {
				respondError(w, r.Method, r.URL.String(), http.StatusInternalServerError, err)
				return
			}
			records = m
		}

		respondSelected(w, keys, offset, limit, records, time.Since(began))
	}
}

func handleInsert(inserter cluster.Inserter) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		began := time.Now()

		var tuples []common.KeyScoreMember
		if err := json.NewDecoder(r.Body).Decode(&tuples); err != nil {
			respondError(w, r.Method, r.URL.String(), http.StatusBadRequest, err)
			return
		}

		if err := inserter.Insert(tuples); err != nil {
			respondError(w, r.Method, r.URL.String(), http.StatusInternalServerError, err)
			return
		}

		respondInserted(w, len(tuples), time.Since(began))
	}
}

func handleDelete(deleter cluster.Deleter) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		began := time.Now()

		var tuples []common.KeyScoreMember
		if err := json.NewDecoder(r.Body).Decode(&tuples); err != nil {
			respondError(w, r.Method, r.URL.String(), http.StatusBadRequest, err)
			return
		}

		if err := deleter.Delete(tuples); err != nil {
			respondError(w, r.Method, r.URL.String(), http.StatusInternalServerError, err)
			return
		}

		respondDeleted(w, len(tuples), time.Since(began))
	}
}

func flatten(m map[string][]common.KeyScoreMember, offset, limit int) []common.KeyScoreMember {
	a := common.KeyScoreMembers{}
	for _, tuples := range m {
		a = append(a, tuples...)
	}

	sort.Sort(a)

	if len(a) < offset {
		return []common.KeyScoreMember{}
	}
	a = a[offset:]

	if len(a) > limit {
		a = a[:limit]
	}

	return a
}

func parseInt(values url.Values, key string, defaultValue int) int {
	value, err := strconv.ParseInt(values.Get(key), 10, 64)
	if err != nil {
		return defaultValue
	}
	return int(value)
}

func parseBool(values url.Values, key string, defaultValue bool) bool {
	value, err := strconv.ParseBool(values.Get(key))
	if err != nil {
		return defaultValue
	}
	return value
}

func respondInserted(w http.ResponseWriter, n int, duration time.Duration) {
	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(map[string]interface{}{
		"inserted": n,
		"duration": duration.String(),
	})
}

func respondSelected(w http.ResponseWriter, keys [][]byte, offset, limit int, records interface{}, duration time.Duration) {
	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(map[string]interface{}{
		"keys":     keys,
		"offset":   offset,
		"limit":    limit,
		"records":  records,
		"duration": duration.String(),
	})
}

func respondDeleted(w http.ResponseWriter, n int, duration time.Duration) {
	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(map[string]interface{}{
		"deleted":  n,
		"duration": duration.String(),
	})
}

func respondError(w http.ResponseWriter, method, url string, code int, err error) {
	log.Printf("%s %s: HTTP %d: %s", method, url, code, err)
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(code)
	json.NewEncoder(w).Encode(map[string]interface{}{
		"error":       err.Error(),
		"code":        code,
		"description": http.StatusText(code),
	})
}

func stripBlank(src []string) []string {
	dst := []string{}
	for _, s := range src {
		if s == "" {
			continue
		}
		dst = append(dst, s)
	}
	return dst
}
<|endoftext|>


<commit_before>// Copyright 2015 PingCAP, Inc.
//
// 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,
// See the License for the specific language governing permissions and
// limitations under the License.

package rsets

import (
	"github.com/juju/errors"
	"github.com/pingcap/tidb/context"
	"github.com/pingcap/tidb/expression"
	"github.com/pingcap/tidb/field"
	"github.com/pingcap/tidb/plan"
	"github.com/pingcap/tidb/plan/plans"
)

var (
	_ plan.Planner = (*HavingRset)(nil)
)

// HavingRset is record set for having fields.
type HavingRset struct {
	Src        plan.Plan
	Expr       expression.Expression
	SelectList *plans.SelectList
	// Group by clause
	GroupBy []expression.Expression
}

// CheckAggregate will check whether order by has aggregate function or not,
// if has, we will add it to select list hidden field.
func (r *HavingRset) CheckAggregate(selectList *plans.SelectList) error {
	if expression.ContainAggregateFunc(r.Expr) {
		expr, err := selectList.UpdateAggFields(r.Expr)
		if err != nil {
			return errors.Errorf("%s in 'having clause'", err.Error())
		}

		r.Expr = expr
	}

	return nil
}

type havingVisitor struct {
	expression.BaseVisitor
	selectList *plans.SelectList

	// for group by
	groupBy []expression.Expression

	// true means we are visiting aggregate function arguments now.
	inAggregate bool
}

func (v *havingVisitor) visitIdentInAggregate(i *expression.Ident) (expression.Expression, error) {
	// if we are visiting aggregate function arguments, the identifier first checks in from table,
	// then in select list, and outer query finally.

	// find this identifier in FROM.
	idx := field.GetResultFieldIndex(i.L, v.selectList.FromFields, field.DefaultFieldFlag)
	if len(idx) > 0 {
		i.ReferScope = expression.IdentReferFromTable
		i.ReferIndex = idx[0]
		return i, nil
	}

	// check in select list.
	index, err := checkIdent(i, v.selectList, havingClause)
	if err != nil {
		return i, errors.Trace(err)
	}

	if index >= 0 {
		// find in select list
		i.ReferScope = expression.IdentReferSelectList
		i.ReferIndex = index
		return i, nil
	}

	// TODO: check in out query
	// TODO: return unknown field error, but now just return directly.
	// Because this may reference outer query.
	return i, nil
}

func (v *havingVisitor) checkIdentInGroupBy(i *expression.Ident) (*expression.Ident, bool, error) {
	for _, by := range v.groupBy {
		e := castIdent(by)
		if e == nil {
			// group by must be a identifier too
			continue
		}

		if !field.CheckFieldsEqual(e.L, i.L) {
			// not same, continue
			continue
		}

		// if group by references select list or having is qualified identifier,
		// no other check.
		if e.ReferScope == expression.IdentReferSelectList || field.IsQualifiedName(i.L) {
			i.ReferScope = e.ReferScope
			i.ReferIndex = e.ReferIndex
			return i, true, nil
		}

		// having is unqualified name, e.g, select * from t1, t2 group by t1.c having c.
		// both t1 and t2 have column c, we must check ambiguous here.

		idx := field.GetResultFieldIndex(i.L, v.selectList.FromFields, field.DefaultFieldFlag)
		if len(idx) > 1 {
			return i, false, errors.Errorf("Column '%s' in having clause is ambiguous", i)
		}

		i.ReferScope = e.ReferScope
		i.ReferIndex = e.ReferIndex
		return i, true, nil
	}

	return i, false, nil
}

func (v *havingVisitor) checkIdentInSelectList(i *expression.Ident) (*expression.Ident, bool, error) {
	index, err := checkIdent(i, v.selectList, havingClause)
	if err != nil {
		return i, false, errors.Trace(err)
	}

	if index >= 0 {
		// identifier references a select field. use it directly.
		// e,g. select c1 as c2 from t having c2, here c2 references c1.
		i.ReferScope = expression.IdentReferSelectList
		i.ReferIndex = index
		return i, true, nil
	}

	lastIndex := -1
	var lastFieldName string
	// we may meet this select c1 as c2 from t having c1, so we must check origin field name too.
	for index := 0; index < v.selectList.HiddenFieldOffset; index++ {
		e := castIdent(v.selectList.Fields[index].Expr)
		if e == nil {
			// not identifier
			continue
		}

		if !field.CheckFieldsEqual(e.L, i.L) {
			// not same, continue
			continue
		}

		if field.IsQualifiedName(i.L) {
			// qualified name, no need check
			i.ReferScope = expression.IdentReferSelectList
			i.ReferIndex = index
			return i, true, nil
		}

		if lastIndex == -1 {
			lastIndex = index
			lastFieldName = e.L
			continue
		}

		// we may meet select t1.c as a, t2.c as b from t1, t2 having c, must check ambiguous here
		if !field.CheckFieldsEqual(lastFieldName, e.L) {
			return i, false, errors.Errorf("Column '%s' in having clause is ambiguous", i)
		}

		i.ReferScope = expression.IdentReferSelectList
		i.ReferIndex = index
		return i, true, nil
	}

	if lastIndex != -1 {
		i.ReferScope = expression.IdentReferSelectList
		i.ReferIndex = lastIndex
		return i, true, nil
	}

	return i, false, nil
}

func (v *havingVisitor) VisitIdent(i *expression.Ident) (expression.Expression, error) {
	// Having has the most complex rule for ambiguous and identifer reference check.

	// Having ambiguous rule:
	//	select c1 as a, c2 as a from t having a is ambiguous
	//	select c1 as a, c2 as a from t having a + 1 is ambiguous
	//	select c1 as c2, c2 from t having c2 is ambiguous
	//	select c1 as c2, c2 from t having c2 + 1 is ambiguous

	// The identifier in having must exist in group by, select list and outer query, so select c1 from t having c2 is wrong,
	// the idenfitier will first try to find the reference in group by, then in select list and finally in outer query.

	// Having identifier reference check
	//	select c1 from t group by c2 having c2, c2 in having references group by c2.
	//	select c1 as c2 from t having c2, c2 in having references c1 in select field.
	//	select c1 as c2 from t group by c2 having c2, c2 in having references group by c2.
	// 	select c1 as c2 from t having sum(c2), c2 in having sum(c2) references t.c2 in from table.
	//	select c1 as c2 from t having sum(c2) + c2, c2 in having sum(c2) references t.c2 in from table, but another c2 references c1 in select list.
	//	select c1 as c2 from t group by c2 having sum(c2) + c2, c2 in having sum(c2) references t.c2 in from table, another c2 references c2 in group by.
	//	select c1 as a from t having a, a references c1 in select field.
	//	select c1 as a from t having sum(a), a in order by sum(a) references c1 in select field.
	//	select c1 as c2 from t having c1, c1 in having references c1 in select list.

	// TODO: unify VisitIdent for group by and order by visitor, they are little different.

	if v.inAggregate {
		return v.visitIdentInAggregate(i)
	}

	var (
		err error
		ok  bool
	)

	// first, check in group by
	i, ok, err = v.checkIdentInGroupBy(i)
	if err != nil || ok {
		return i, errors.Trace(err)
	}

	// then in select list
	i, ok, err = v.checkIdentInSelectList(i)
	if err != nil || ok {
		return i, errors.Trace(err)
	}

	// TODO: check in out query

	return i, errors.Errorf("Unknown column '%s' in 'having clause'", i)
}

func (v *havingVisitor) VisitPosition(p *expression.Position) (expression.Expression, error) {
	n := p.N

	// we have added order by to hidden field before, now visit it here.
	expr := v.selectList.Fields[n-1].Expr
	e, err := expr.Accept(v)
	if err != nil {
		return nil, errors.Trace(err)
	}
	v.selectList.Fields[n-1].Expr = e

	return p, nil
}

func (v *havingVisitor) VisitCall(c *expression.Call) (expression.Expression, error) {
	isAggregate, err := expression.IsAggregateFunc(c.F)
	if err != nil {
		return nil, errors.Trace(err)
	}

	v.inAggregate = isAggregate

	for i, e := range c.Args {
		c.Args[i], err = e.Accept(v)
		if err != nil {
			return nil, errors.Trace(err)
		}
	}

	v.inAggregate = false

	return c, nil
}

// Plan gets HavingPlan.
func (r *HavingRset) Plan(ctx context.Context) (plan.Plan, error) {
	visitor := &havingVisitor{
		selectList:  r.SelectList,
		groupBy:     r.GroupBy,
		inAggregate: false,
	}
	visitor.BaseVisitor.V = visitor

	e, err := r.Expr.Accept(visitor)
	if err != nil {
		return nil, errors.Trace(err)
	}

	r.Expr = e

	return &plans.HavingPlan{Src: r.Src, Expr: r.Expr, SelectList: r.SelectList}, nil
}

//  String implements fmt.Stringer interface.
func (r *HavingRset) String() string {
	return r.Expr.String()
}
<commit_msg>rsets: Address comment.<commit_after>// Copyright 2015 PingCAP, Inc.
//
// 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,
// See the License for the specific language governing permissions and
// limitations under the License.

package rsets

import (
	"github.com/juju/errors"
	"github.com/pingcap/tidb/context"
	"github.com/pingcap/tidb/expression"
	"github.com/pingcap/tidb/field"
	"github.com/pingcap/tidb/plan"
	"github.com/pingcap/tidb/plan/plans"
)

var (
	_ plan.Planner = (*HavingRset)(nil)
)

// HavingRset is record set for having fields.
type HavingRset struct {
	Src        plan.Plan
	Expr       expression.Expression
	SelectList *plans.SelectList
	// Group by clause
	GroupBy []expression.Expression
}

// CheckAggregate will check whether order by has aggregate function or not,
// if has, we will add it to select list hidden field.
func (r *HavingRset) CheckAggregate(selectList *plans.SelectList) error {
	if expression.ContainAggregateFunc(r.Expr) {
		expr, err := selectList.UpdateAggFields(r.Expr)
		if err != nil {
			return errors.Errorf("%s in 'having clause'", err.Error())
		}

		r.Expr = expr
	}

	return nil
}

type havingVisitor struct {
	expression.BaseVisitor
	selectList *plans.SelectList

	// for group by
	groupBy []expression.Expression

	// true means we are visiting aggregate function arguments now.
	inAggregate bool
}

func (v *havingVisitor) visitIdentInAggregate(i *expression.Ident) (expression.Expression, error) {
	// if we are visiting aggregate function arguments, the identifier first checks in from table,
	// then in select list, and outer query finally.

	// find this identifier in FROM.
	idx := field.GetResultFieldIndex(i.L, v.selectList.FromFields, field.DefaultFieldFlag)
	if len(idx) > 0 {
		i.ReferScope = expression.IdentReferFromTable
		i.ReferIndex = idx[0]
		return i, nil
	}

	// check in select list.
	index, err := checkIdent(i, v.selectList, havingClause)
	if err != nil {
		return i, errors.Trace(err)
	}

	if index >= 0 {
		// find in select list
		i.ReferScope = expression.IdentReferSelectList
		i.ReferIndex = index
		return i, nil
	}

	// TODO: check in out query
	// TODO: return unknown field error, but now just return directly.
	// Because this may reference outer query.
	return i, nil
}

func (v *havingVisitor) checkIdentInGroupBy(i *expression.Ident) (*expression.Ident, bool, error) {
	for _, by := range v.groupBy {
		e := castIdent(by)
		if e == nil {
			// group by must be a identifier too
			continue
		}

		if !field.CheckFieldsEqual(e.L, i.L) {
			// not same, continue
			continue
		}

		// if group by references select list or having is qualified identifier,
		// no other check.
		if e.ReferScope == expression.IdentReferSelectList || field.IsQualifiedName(i.L) {
			i.ReferScope = e.ReferScope
			i.ReferIndex = e.ReferIndex
			return i, true, nil
		}

		// having is unqualified name, e.g, select * from t1, t2 group by t1.c having c.
		// both t1 and t2 have column c, we must check ambiguous here.

		idx := field.GetResultFieldIndex(i.L, v.selectList.FromFields, field.DefaultFieldFlag)
		if len(idx) > 1 {
			return i, false, errors.Errorf("Column '%s' in having clause is ambiguous", i)
		}

		i.ReferScope = e.ReferScope
		i.ReferIndex = e.ReferIndex
		return i, true, nil
	}

	return i, false, nil
}

func (v *havingVisitor) checkIdentInSelectList(i *expression.Ident) (*expression.Ident, bool, error) {
	index, err := checkIdent(i, v.selectList, havingClause)
	if err != nil {
		return i, false, errors.Trace(err)
	}

	if index >= 0 {
		// identifier references a select field. use it directly.
		// e,g. select c1 as c2 from t having c2, here c2 references c1.
		i.ReferScope = expression.IdentReferSelectList
		i.ReferIndex = index
		return i, true, nil
	}

	lastIndex := -1
	var lastFieldName string
	// we may meet this select c1 as c2 from t having c1, so we must check origin field name too.
	for index := 0; index < v.selectList.HiddenFieldOffset; index++ {
		e := castIdent(v.selectList.Fields[index].Expr)
		if e == nil {
			// not identifier
			continue
		}

		if !field.CheckFieldsEqual(e.L, i.L) {
			// not same, continue
			continue
		}

		if field.IsQualifiedName(i.L) {
			// qualified name, no need check
			i.ReferScope = expression.IdentReferSelectList
			i.ReferIndex = index
			return i, true, nil
		}

		if lastIndex == -1 {
			lastIndex = index
			lastFieldName = e.L
			continue
		}

		// we may meet select t1.c as a, t2.c as b from t1, t2 having c, must check ambiguous here
		if !field.CheckFieldsEqual(lastFieldName, e.L) {
			return i, false, errors.Errorf("Column '%s' in having clause is ambiguous", i)
		}

		i.ReferScope = expression.IdentReferSelectList
		i.ReferIndex = index
		return i, true, nil
	}

	if lastIndex != -1 {
		i.ReferScope = expression.IdentReferSelectList
		i.ReferIndex = lastIndex
		return i, true, nil
	}

	return i, false, nil
}

func (v *havingVisitor) VisitIdent(i *expression.Ident) (expression.Expression, error) {
	// Having has the most complex rule for ambiguous and identifer reference check.

	// Having ambiguous rule:
	//	select c1 as a, c2 as a from t having a is ambiguous
	//	select c1 as a, c2 as a from t having a + 1 is ambiguous
	//	select c1 as c2, c2 from t having c2 is ambiguous
	//	select c1 as c2, c2 from t having c2 + 1 is ambiguous

	// The identifier in having must exist in group by, select list or outer query, so select c1 from t having c2 is wrong,
	// the identifier will first try to find the reference in group by, then in select list and finally in outer query.

	// Having identifier reference check
	//	select c1 from t group by c2 having c2, c2 in having references group by c2.
	//	select c1 as c2 from t having c2, c2 in having references c1 in select field.
	//	select c1 as c2 from t group by c2 having c2, c2 in having references group by c2.
	// 	select c1 as c2 from t having sum(c2), c2 in having sum(c2) references t.c2 in from table.
	//	select c1 as c2 from t having sum(c2) + c2, c2 in having sum(c2) references t.c2 in from table, but another c2 references c1 in select list.
	//	select c1 as c2 from t group by c2 having sum(c2) + c2, c2 in having sum(c2) references t.c2 in from table, another c2 references c2 in group by.
	//	select c1 as a from t having a, a references c1 in select field.
	//	select c1 as a from t having sum(a), a in order by sum(a) references c1 in select field.
	//	select c1 as c2 from t having c1, c1 in having references c1 in select list.

	// TODO: unify VisitIdent for group by and order by visitor, they are little different.

	if v.inAggregate {
		return v.visitIdentInAggregate(i)
	}

	var (
		err error
		ok  bool
	)

	// first, check in group by
	i, ok, err = v.checkIdentInGroupBy(i)
	if err != nil || ok {
		return i, errors.Trace(err)
	}

	// then in select list
	i, ok, err = v.checkIdentInSelectList(i)
	if err != nil || ok {
		return i, errors.Trace(err)
	}

	// TODO: check in out query

	return i, errors.Errorf("Unknown column '%s' in 'having clause'", i)
}

func (v *havingVisitor) VisitPosition(p *expression.Position) (expression.Expression, error) {
	n := p.N

	// we have added order by to hidden field before, now visit it here.
	expr := v.selectList.Fields[n-1].Expr
	e, err := expr.Accept(v)
	if err != nil {
		return nil, errors.Trace(err)
	}
	v.selectList.Fields[n-1].Expr = e

	return p, nil
}

func (v *havingVisitor) VisitCall(c *expression.Call) (expression.Expression, error) {
	isAggregate, err := expression.IsAggregateFunc(c.F)
	if err != nil {
		return nil, errors.Trace(err)
	}

	v.inAggregate = isAggregate

	for i, e := range c.Args {
		c.Args[i], err = e.Accept(v)
		if err != nil {
			return nil, errors.Trace(err)
		}
	}

	v.inAggregate = false

	return c, nil
}

// Plan gets HavingPlan.
func (r *HavingRset) Plan(ctx context.Context) (plan.Plan, error) {
	visitor := &havingVisitor{
		selectList:  r.SelectList,
		groupBy:     r.GroupBy,
		inAggregate: false,
	}
	visitor.BaseVisitor.V = visitor

	e, err := r.Expr.Accept(visitor)
	if err != nil {
		return nil, errors.Trace(err)
	}

	r.Expr = e

	return &plans.HavingPlan{Src: r.Src, Expr: r.Expr, SelectList: r.SelectList}, nil
}

//  String implements fmt.Stringer interface.
func (r *HavingRset) String() string {
	return r.Expr.String()
}
<|endoftext|>


<commit_before>package plugins

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"regexp"

	"github.com/InnovaCo/serve/manifest"
	"github.com/InnovaCo/serve/utils/gabs"
)

func init() {
	manifest.PluginRegestry.Add("gocd.pipeline.create", goCdPipelineCreate{})
}

/**
 * plugin for manifest section "goCd.pipeline.create"
 * section structure:
 *
 * goCd.pipeline.create:
 *   api-url: goCd_URL
 *   environment: ENV
 *   branch: BRANCH
 *   allowed-branches: [BRANCH, ...]
 *   pipeline:
 *     group: GROUP
 *     pipeline:
 *       according to the description: https://api.go.cd/current/#the-pipeline-config-object
 */

type goCdCredents struct {
	Login    string `json:"login"`
	Password string `json:"password"`
}

type goCdPipelineCreate struct{}

func (p goCdPipelineCreate) Run(data manifest.Manifest) error {
	name := data.GetString("pipeline.pipeline.name")
	url := data.GetString("api-url")
	body := data.GetTree("pipeline").String()
	branch := data.GetString("branch")

	m := false
	for _, b := range data.GetArray("allowed-branches") {
		re := b.Unwrap().(string)
		if re == branch {
			m = true
			break
		} else if m, _ = regexp.MatchString(re, branch); m {
			break
		}
	}

	if !m {
		log.Println("branch ", branch, " not in ", data.GetString("allowed-branches"))
		return nil
	}

	resp, err := goCdRequest("GET", url+"/pipelines/"+name, "", nil)
	if err != nil {
		return err
	}

	if resp.StatusCode == http.StatusOK {
		err = goCdUpdate(name, data.GetString("environment"), url, body, map[string]string{"If-Match": resp.Header.Get("ETag")})
	} else if resp.StatusCode == http.StatusNotFound {
		err = goCdCreate(name, data.GetString("environment"), url, body, nil)
	} else {
		log.Println("Operation error: " + resp.Status)
		return errors.New("Operation error: " + resp.Status)
	}

	if err != nil {
		log.Println(err)
		return err
	}

	return nil
}

func goCdCreate(name string, env string, resource string, body string, headers map[string]string) error {
	if resp, err := goCdRequest("POST", resource+"/pipelines", body, nil); err != nil {
		log.Println(err)
		return err
	} else if resp.StatusCode != http.StatusOK {
		return errors.New("Operation error: " + resp.Status)
	}
	data, tag, err := goCdChangeEnv(resource, env, name, "")
	if err != nil {
		log.Println(err)
		return err
	}

	log.Println(data)

	if resp, err := goCdRequest("PUT", resource+"/environments/"+env, data, map[string]string{"If-Match": tag}); err != nil {
		log.Println(err)
		return err
	} else if resp.StatusCode != http.StatusOK {
		log.Println("Operation error: " + resp.Status)
		return errors.New("Operation error: " + resp.Status)
	}

	return nil
}

func goCdUpdate(name string, env string, resource string, body string, headers map[string]string) error {
	if resp, err := goCdRequest("PUT", resource+"/pipelines/"+name, body, headers); err != nil {
		log.Println(err)
		return err
	} else if resp.StatusCode != http.StatusOK {
		log.Println("Operation error: " + resp.Status)
		return errors.New("Operation error: " + resp.Status)
	}

	cEnv, err := goCdFindEnv(resource, name)
	if err != nil {
		return err
	}

	if env != cEnv {
		data, tag, err := goCdChangeEnv(resource, cEnv, "", name)
		if err != nil {
			log.Println(err)
			return err
		}

		log.Println(data)

		if resp, err := goCdRequest("PUT", resource+"/environments/"+cEnv, data, map[string]string{"If-Match": tag}); err != nil {
			log.Println(err)
			return err
		} else if resp.StatusCode != http.StatusOK {
			log.Println("Operation error: " + resp.Status)
			return errors.New("Operation error: " + resp.Status)
		}
		//
		data, tag, err = goCdChangeEnv(resource, env, name, "")
		if err != nil {
			log.Println(err)
			return err
		}

		log.Println(data)

		if resp, err := goCdRequest("PUT", resource+"/environments/"+env, data, map[string]string{"If-Match": tag}); err != nil {
			log.Println(err)
			return err
		} else if resp.StatusCode != http.StatusOK {
			log.Println("Operation error: " + resp.Status)
			return errors.New("Operation error: " + resp.Status)
		}
	}

	return nil
}

func goCdDelete(name string, env string, resource string) error {
	data, tag, err := goCdChangeEnv(resource, env, "", name)
	if err != nil {
		log.Println(err)
		return err
	}

	log.Println(data)

	if resp, err := goCdRequest("PUT", resource+"/environments/"+env, data, map[string]string{"If-Match": tag}); err != nil {
		log.Println(err)
		return err
	} else if resp.StatusCode != http.StatusOK {
		log.Println("Operation error: " + resp.Status)
		return errors.New("Operation error: " + resp.Status)
	}

	if resp, err := goCdRequest("DELETE", resource+"/pipelines/"+name, "", nil); err != nil {
		log.Println(err)
		return err
	} else if resp.StatusCode != http.StatusOK {
		log.Println("Operation error: " + resp.Status)
		return errors.New("Operation error: " + resp.Status)
	}

	return nil
}

func goCdChangeEnv(resource string, env string, addPipeline string, delPipeline string) (string, string, error) {
	resp, err := goCdRequest("GET", resource+"/environments/"+env, "", nil)
	if err != nil {
		log.Println(err)
		return "", "", err
	}

	data, err := ChangeJSON(resp, addPipeline, delPipeline)
	if err != nil {
		log.Println(err)
		return "", "", err
	}

	return data, resp.Header.Get("ETag"), nil
}

func goCdFindEnv(resource string, pipeline string) (string, error) {
	resp, err := goCdRequest("GET", resource+"/environments", "", nil)
	if err != nil {
		log.Println(err)
		return "", err
	}

	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return "", errors.New("parse error")
	}

	tree, err := gabs.ParseJSON(body)
	if err != nil {
		return "", errors.New("parse error")
	}

	envs, _ := tree.Path("_embedded.environments").Children()
	for _, env := range envs {
		envName := env.Path("name").Data().(string)
		pipelines, _ := env.Path("pipelines").Children()
		for _, pline := range pipelines {
			if pline.Path("name").Data().(string) == pipeline {
				return envName, nil
			}
		}
	}

	return "", nil
}

func goCdRequest(method string, resource string, body string, headers map[string]string) (*http.Response, error) {
	req, _ := http.NewRequest(method, resource, bytes.NewReader([]byte(body)))

	for k, v := range headers {
		req.Header.Set(k, v)
	}

	req.Header.Set("Accept", "application/vnd.go.cd.v1+json")
	req.Header.Set("Content-Type", "application/json")

	data, err := ioutil.ReadFile("/etc/serve/gocd_credentials")
	if err != nil {
		return nil, errors.New(fmt.Sprintf("Credentias file error: %v", err))
	}

	creds := &goCdCredents{}
	json.Unmarshal(data, creds)

	req.SetBasicAuth(creds.Login, creds.Password)

	log.Printf(" --> %s %s:\n%s\n%s\n", method, resource, req.Header, body)

	return http.DefaultClient.Do(req)
}

func ChangeJSON(resp *http.Response, addPipeline string, delPipeline string) (string, error) {
	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)

	if err != nil {
		return "", errors.New("read body error")
	}

	tree, err := gabs.ParseJSON(body)
	if err != nil {
		return "", errors.New("parse error")
	}
	result := gabs.New()

	result.Set(tree.Path("name").Data(), "name")

	children, _ := tree.S("pipelines").Children()
	vals := []map[string]string{}
	for _, m := range children {
		name := m.Path("name").Data().(string)
		if (delPipeline != "") && (name == delPipeline) {
			continue
		}
		if (addPipeline != "") && (name == addPipeline) {
			addPipeline = ""
		}
		vals = append(vals, map[string]string{"name": name})
	}
	if addPipeline != "" {
		vals = append(vals, map[string]string{"name": addPipeline})
	}
	result.Set(vals, "pipelines")

	children, _ = tree.S("agents").Children()
	vals = []map[string]string{}
	for _, m := range children {
		vals = append(vals, map[string]string{"uuid": m.Path("uuid").Data().(string)})
	}
	result.Set(vals, "agents")
	result.Set(tree.Path("environment_variables").Data(), "environment_variables")

	return result.String(), nil
}
<commit_msg>change API version<commit_after>package plugins

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"regexp"

	"github.com/InnovaCo/serve/manifest"
	"github.com/InnovaCo/serve/utils/gabs"
)

func init() {
	manifest.PluginRegestry.Add("gocd.pipeline.create", goCdPipelineCreate{})
}

/**
 * plugin for manifest section "goCd.pipeline.create"
 * section structure:
 *
 * goCd.pipeline.create:
 *   api-url: goCd_URL
 *   environment: ENV
 *   branch: BRANCH
 *   allowed-branches: [BRANCH, ...]
 *   pipeline:
 *     group: GROUP
 *     pipeline:
 *       according to the description: https://api.go.cd/current/#the-pipeline-config-object
 */

type goCdCredents struct {
	Login    string `json:"login"`
	Password string `json:"password"`
}

type goCdPipelineCreate struct{}

func (p goCdPipelineCreate) Run(data manifest.Manifest) error {
	name := data.GetString("pipeline.pipeline.name")
	url := data.GetString("api-url")
	body := data.GetTree("pipeline").String()
	branch := data.GetString("branch")

	m := false
	for _, b := range data.GetArray("allowed-branches") {
		re := b.Unwrap().(string)
		if re == branch {
			m = true
			break
		} else if m, _ = regexp.MatchString(re, branch); m {
			break
		}
	}

	if !m {
		log.Println("branch ", branch, " not in ", data.GetString("allowed-branches"))
		return nil
	}

	resp, err := goCdRequest("GET", url+"/pipelines/"+name, "", nil)
	if err != nil {
		return err
	}

	if resp.StatusCode == http.StatusOK {
		err = goCdUpdate(name, data.GetString("environment"), url, body, map[string]string{"If-Match": resp.Header.Get("ETag")})
	} else if resp.StatusCode == http.StatusNotFound {
		err = goCdCreate(name, data.GetString("environment"), url, body, nil)
	} else {
		log.Println("Operation error: " + resp.Status)
		return errors.New("Operation error: " + resp.Status)
	}

	if err != nil {
		log.Println(err)
		return err
	}

	return nil
}

func goCdCreate(name string, env string, resource string, body string, headers map[string]string) error {
	if resp, err := goCdRequest("POST", resource+"/pipelines", body, nil); err != nil {
		log.Println(err)
		return err
	} else if resp.StatusCode != http.StatusOK {
		return errors.New("Operation error: " + resp.Status)
	}
	data, tag, err := goCdChangeEnv(resource, env, name, "")
	if err != nil {
		log.Println(err)
		return err
	}

	log.Println(data)

	if resp, err := goCdRequest("PUT", resource+"/environments/"+env, data, map[string]string{"If-Match": tag}); err != nil {
		log.Println(err)
		return err
	} else if resp.StatusCode != http.StatusOK {
		log.Println("Operation error: " + resp.Status)
		return errors.New("Operation error: " + resp.Status)
	}

	return nil
}

func goCdUpdate(name string, env string, resource string, body string, headers map[string]string) error {
	if resp, err := goCdRequest("PUT", resource+"/pipelines/"+name, body, headers); err != nil {
		log.Println(err)
		return err
	} else if resp.StatusCode != http.StatusOK {
		log.Println("Operation error: " + resp.Status)
		return errors.New("Operation error: " + resp.Status)
	}

	cEnv, err := goCdFindEnv(resource, name)
	if err != nil {
		return err
	}

	if env != cEnv {
		data, tag, err := goCdChangeEnv(resource, cEnv, "", name)
		if err != nil {
			log.Println(err)
			return err
		}

		log.Println(data)

		if resp, err := goCdRequest("PUT", resource+"/environments/"+cEnv, data, map[string]string{"If-Match": tag}); err != nil {
			log.Println(err)
			return err
		} else if resp.StatusCode != http.StatusOK {
			log.Println("Operation error: " + resp.Status)
			return errors.New("Operation error: " + resp.Status)
		}
		//
		data, tag, err = goCdChangeEnv(resource, env, name, "")
		if err != nil {
			log.Println(err)
			return err
		}

		log.Println(data)

		if resp, err := goCdRequest("PUT", resource+"/environments/"+env, data, map[string]string{"If-Match": tag}); err != nil {
			log.Println(err)
			return err
		} else if resp.StatusCode != http.StatusOK {
			log.Println("Operation error: " + resp.Status)
			return errors.New("Operation error: " + resp.Status)
		}
	}

	return nil
}

func goCdDelete(name string, env string, resource string) error {
	data, tag, err := goCdChangeEnv(resource, env, "", name)
	if err != nil {
		log.Println(err)
		return err
	}

	log.Println(data)

	if resp, err := goCdRequest("PUT", resource+"/environments/"+env, data, map[string]string{"If-Match": tag}); err != nil {
		log.Println(err)
		return err
	} else if resp.StatusCode != http.StatusOK {
		log.Println("Operation error: " + resp.Status)
		return errors.New("Operation error: " + resp.Status)
	}

	if resp, err := goCdRequest("DELETE", resource+"/pipelines/"+name, "", nil); err != nil {
		log.Println(err)
		return err
	} else if resp.StatusCode != http.StatusOK {
		log.Println("Operation error: " + resp.Status)
		return errors.New("Operation error: " + resp.Status)
	}

	return nil
}

func goCdChangeEnv(resource string, env string, addPipeline string, delPipeline string) (string, string, error) {
	resp, err := goCdRequest("GET", resource+"/environments/"+env, "", nil)
	if err != nil {
		log.Println(err)
		return "", "", err
	}

	data, err := ChangeJSON(resp, addPipeline, delPipeline)
	if err != nil {
		log.Println(err)
		return "", "", err
	}

	return data, resp.Header.Get("ETag"), nil
}

func goCdFindEnv(resource string, pipeline string) (string, error) {
	resp, err := goCdRequest("GET", resource+"/environments", "", nil)
	if err != nil {
		log.Println(err)
		return "", err
	}

	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return "", errors.New("parse error")
	}

	tree, err := gabs.ParseJSON(body)
	if err != nil {
		return "", errors.New("parse error")
	}

	envs, _ := tree.Path("_embedded.environments").Children()
	for _, env := range envs {
		envName := env.Path("name").Data().(string)
		pipelines, _ := env.Path("pipelines").Children()
		for _, pline := range pipelines {
			if pline.Path("name").Data().(string) == pipeline {
				return envName, nil
			}
		}
	}

	return "", nil
}

func goCdRequest(method string, resource string, body string, headers map[string]string) (*http.Response, error) {
	req, _ := http.NewRequest(method, resource, bytes.NewReader([]byte(body)))

	for k, v := range headers {
		req.Header.Set(k, v)
	}

	req.Header.Set("Accept", "application/vnd.go.cd.v2+json")
	req.Header.Set("Content-Type", "application/json")

	data, err := ioutil.ReadFile("/etc/serve/gocd_credentials")
	if err != nil {
		return nil, errors.New(fmt.Sprintf("Credentias file error: %v", err))
	}

	creds := &goCdCredents{}
	json.Unmarshal(data, creds)

	req.SetBasicAuth(creds.Login, creds.Password)

	log.Printf(" --> %s %s:\n%s\n%s\n", method, resource, req.Header, body)

	return http.DefaultClient.Do(req)
}

func ChangeJSON(resp *http.Response, addPipeline string, delPipeline string) (string, error) {
	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)

	if err != nil {
		return "", errors.New("read body error")
	}

	tree, err := gabs.ParseJSON(body)
	if err != nil {
		return "", errors.New("parse error")
	}
	result := gabs.New()

	result.Set(tree.Path("name").Data(), "name")

	children, _ := tree.S("pipelines").Children()
	vals := []map[string]string{}
	for _, m := range children {
		name := m.Path("name").Data().(string)
		if (delPipeline != "") && (name == delPipeline) {
			continue
		}
		if (addPipeline != "") && (name == addPipeline) {
			addPipeline = ""
		}
		vals = append(vals, map[string]string{"name": name})
	}
	if addPipeline != "" {
		vals = append(vals, map[string]string{"name": addPipeline})
	}
	result.Set(vals, "pipelines")

	children, _ = tree.S("agents").Children()
	vals = []map[string]string{}
	for _, m := range children {
		vals = append(vals, map[string]string{"uuid": m.Path("uuid").Data().(string)})
	}
	result.Set(vals, "agents")
	result.Set(tree.Path("environment_variables").Data(), "environment_variables")

	return result.String(), nil
}
<|endoftext|>


<commit_before>package cmp

import (
	"fmt"
	"reflect"
	"testing"
	"time"

	"github.com/influxdata/telegraf/testutil"

	"github.com/influxdata/telegraf"
)

func TestBuildPoint(t *testing.T) {
	var tagtests = []struct {
		ptIn  telegraf.Metric
		outPt Point
		err   error
	}{
		{
			testutil.TestMetric(float64(0.0), "testpt"),
			Point{
				float64(time.Date(2009, time.November, 10, 23, 0, 0, 0, time.UTC).Unix()),
				0.0,
			},
			nil,
		},
		{
			testutil.TestMetric(float64(1.0), "testpt"),
			Point{
				float64(time.Date(2009, time.November, 10, 23, 0, 0, 0, time.UTC).Unix()),
				1.0,
			},
			nil,
		},
		{
			testutil.TestMetric(int(10), "testpt"),
			Point{
				float64(time.Date(2009, time.November, 10, 23, 0, 0, 0, time.UTC).Unix()),
				10.0,
			},
			nil,
		},
		{
			testutil.TestMetric(int32(112345), "testpt"),
			Point{
				float64(time.Date(2009, time.November, 10, 23, 0, 0, 0, time.UTC).Unix()),
				112345.0,
			},
			nil,
		},
		{
			testutil.TestMetric(int64(112345), "testpt"),
			Point{
				float64(time.Date(2009, time.November, 10, 23, 0, 0, 0, time.UTC).Unix()),
				112345.0,
			},
			nil,
		},
		{
			testutil.TestMetric(float32(11234.5), "testpt"),
			Point{
				float64(time.Date(2009, time.November, 10, 23, 0, 0, 0, time.UTC).Unix()),
				11234.5,
			},
			nil,
		},
		{
			testutil.TestMetric("11234.5", "testpt"),
			Point{
				float64(time.Date(2009, time.November, 10, 23, 0, 0, 0, time.UTC).Unix()),
				11234.5,
			},
			fmt.Errorf("unable to extract value from Fields, undeterminable type"),
		},
	}
	for _, tt := range tagtests {
		pt, err := buildMetrics(tt.ptIn)
		if err != nil && tt.err == nil {
			t.Errorf("%s: unexpected error, %+v\n", tt.ptIn.Name(), err)
		}
		if tt.err != nil && err == nil {
			t.Errorf("%s: expected an error (%s) but none returned", tt.ptIn.Name(), tt.err.Error())
		}
		if !reflect.DeepEqual(pt["value"], tt.outPt) && tt.err == nil {
			t.Errorf("%s: \nexpected %+v\ngot %+v\n",
				tt.ptIn.Name(), tt.outPt, pt["value"])
		}
	}
}
<commit_msg>Remove the old tests as they are not working<commit_after>package cmp
<|endoftext|>


<commit_before>// Copyright 2015 Eleme Inc. All rights reserved.

// Package alerter implements an alerter to send sms/email messages
// on anomalies found.
package alerter

import (
	"encoding/json"
	"os/exec"
	"sync/atomic"
	"time"

	"github.com/eleme/banshee/config"
	"github.com/eleme/banshee/health"
	"github.com/eleme/banshee/models"
	"github.com/eleme/banshee/storage"
	"github.com/eleme/banshee/util/log"
	"github.com/eleme/banshee/util/safemap"
)

// Limit for buffered detected metric results, further results will be dropped
// if this limit is reached.
const bufferedMetricResultsLimit = 10 * 1024

// Alerter alerts on anomalies detected.
type Alerter struct {
	// Storage
	db *storage.DB
	// Config
	cfg *config.Config
	// Input
	In chan *models.Metric
	// Alertings stamps
	m *safemap.SafeMap
	// Alertings counters
	c *safemap.SafeMap
}

// Alerting message.
type msg struct {
	Project *models.Project `json:"project"`
	Metric  *models.Metric  `json:"metric"`
	User    *models.User    `json:"user"`
}

// New creates a alerter.
func New(cfg *config.Config, db *storage.DB) *Alerter {
	al := new(Alerter)
	al.cfg = cfg
	al.db = db
	al.In = make(chan *models.Metric, bufferedMetricResultsLimit)
	al.m = safemap.New()
	al.c = safemap.New()
	return al
}

// Start several goroutines to wait for detected metrics, then check each
// metric with all the rules, the configured shell command will be executed
// once a rule is hit.
func (al *Alerter) Start() {
	log.Info("start %d alerter workers..", al.cfg.Alerter.Workers)
	for i := 0; i < al.cfg.Alerter.Workers; i++ {
		go al.work()
	}
	go func() {
		ticker := time.NewTicker(time.Hour * 24)
		for _ = range ticker.C {
			al.c.Clear()
		}
	}()
}

// work waits for detected metrics, then check each metric with all the
// rules, the configured shell command will be executed once a rule is hit.
func (al *Alerter) work() {
	for {
		metric := <-al.In
		// Check interval.
		v, ok := al.m.Get(metric.Name)
		if ok && metric.Stamp-v.(uint32) < al.cfg.Alerter.Interval {
			continue
		}
		// Check alert times in one day
		v, ok = al.c.Get(metric.Name)
		if ok && atomic.LoadUint32(v.(*uint32)) > al.cfg.Alerter.OneDayLimit {
			log.Warn("%s hit alerting one day limit, skipping..", metric.Name)
			continue
		}
		if !ok {
			var newCounter uint32
			newCounter = 1
			al.c.Set(metric.Name, &newCounter)
		} else {
			atomic.AddUint32(v.(*uint32), 1)
		}
		// Universals
		var univs []models.User
		if err := al.db.Admin.DB().Where("universal = ?", true).Find(&univs).Error; err != nil {
			log.Error("get universal users: %v, skiping..", err)
			continue
		}
		for _, rule := range metric.TestedRules {
			// Project
			proj := &models.Project{}
			if err := al.db.Admin.DB().Model(rule).Related(proj).Error; err != nil {
				log.Error("project, %v, skiping..", err)
				continue
			}
			// Users
			var users []models.User
			if err := al.db.Admin.DB().Model(proj).Related(&users, "Users").Error; err != nil {
				log.Error("get users: %v, skiping..", err)
				continue
			}
			users = append(users, univs...)
			// Send
			for _, user := range users {
				if !user.EnableEmail && !user.EnablePhone {
					continue
				}
				d := &msg{
					Project: proj,
					Metric:  metric,
					User:    &user,
				}
				// Exec
				if len(al.cfg.Alerter.Command) == 0 {
					log.Warn("alert command not configured")
					continue
				}
				b, _ := json.Marshal(d)
				cmd := exec.Command(al.cfg.Alerter.Command, string(b))
				if err := cmd.Run(); err != nil {
					log.Error("exec %s: %v", al.cfg.Alerter.Command, err)
					continue
				}
				log.Info("send message to %s with %s ok", user.Name, metric.Name)
			}
			if len(users) != 0 {
				al.m.Set(metric.Name, metric.Stamp)
				health.IncrNumAlertingEvents(1)
			}
		}
	}
}
<commit_msg>Enable the sender command execution if user email and phone are both not enabled.<commit_after>// Copyright 2015 Eleme Inc. All rights reserved.

// Package alerter implements an alerter to send sms/email messages
// on anomalies found.
package alerter

import (
	"encoding/json"
	"os/exec"
	"sync/atomic"
	"time"

	"github.com/eleme/banshee/config"
	"github.com/eleme/banshee/health"
	"github.com/eleme/banshee/models"
	"github.com/eleme/banshee/storage"
	"github.com/eleme/banshee/util/log"
	"github.com/eleme/banshee/util/safemap"
)

// Limit for buffered detected metric results, further results will be dropped
// if this limit is reached.
const bufferedMetricResultsLimit = 10 * 1024

// Alerter alerts on anomalies detected.
type Alerter struct {
	// Storage
	db *storage.DB
	// Config
	cfg *config.Config
	// Input
	In chan *models.Metric
	// Alertings stamps
	m *safemap.SafeMap
	// Alertings counters
	c *safemap.SafeMap
}

// Alerting message.
type msg struct {
	Project *models.Project `json:"project"`
	Metric  *models.Metric  `json:"metric"`
	User    *models.User    `json:"user"`
}

// New creates a alerter.
func New(cfg *config.Config, db *storage.DB) *Alerter {
	al := new(Alerter)
	al.cfg = cfg
	al.db = db
	al.In = make(chan *models.Metric, bufferedMetricResultsLimit)
	al.m = safemap.New()
	al.c = safemap.New()
	return al
}

// Start several goroutines to wait for detected metrics, then check each
// metric with all the rules, the configured shell command will be executed
// once a rule is hit.
func (al *Alerter) Start() {
	log.Info("start %d alerter workers..", al.cfg.Alerter.Workers)
	for i := 0; i < al.cfg.Alerter.Workers; i++ {
		go al.work()
	}
	go func() {
		ticker := time.NewTicker(time.Hour * 24)
		for _ = range ticker.C {
			al.c.Clear()
		}
	}()
}

// work waits for detected metrics, then check each metric with all the
// rules, the configured shell command will be executed once a rule is hit.
func (al *Alerter) work() {
	for {
		metric := <-al.In
		// Check interval.
		v, ok := al.m.Get(metric.Name)
		if ok && metric.Stamp-v.(uint32) < al.cfg.Alerter.Interval {
			continue
		}
		// Check alert times in one day
		v, ok = al.c.Get(metric.Name)
		if ok && atomic.LoadUint32(v.(*uint32)) > al.cfg.Alerter.OneDayLimit {
			log.Warn("%s hit alerting one day limit, skipping..", metric.Name)
			continue
		}
		if !ok {
			var newCounter uint32
			newCounter = 1
			al.c.Set(metric.Name, &newCounter)
		} else {
			atomic.AddUint32(v.(*uint32), 1)
		}
		// Universals
		var univs []models.User
		if err := al.db.Admin.DB().Where("universal = ?", true).Find(&univs).Error; err != nil {
			log.Error("get universal users: %v, skiping..", err)
			continue
		}
		for _, rule := range metric.TestedRules {
			// Project
			proj := &models.Project{}
			if err := al.db.Admin.DB().Model(rule).Related(proj).Error; err != nil {
				log.Error("project, %v, skiping..", err)
				continue
			}
			// Users
			var users []models.User
			if err := al.db.Admin.DB().Model(proj).Related(&users, "Users").Error; err != nil {
				log.Error("get users: %v, skiping..", err)
				continue
			}
			users = append(users, univs...)
			// Send
			for _, user := range users {
				d := &msg{
					Project: proj,
					Metric:  metric,
					User:    &user,
				}
				// Exec
				if len(al.cfg.Alerter.Command) == 0 {
					log.Warn("alert command not configured")
					continue
				}
				b, _ := json.Marshal(d)
				cmd := exec.Command(al.cfg.Alerter.Command, string(b))
				if err := cmd.Run(); err != nil {
					log.Error("exec %s: %v", al.cfg.Alerter.Command, err)
					continue
				}
				log.Info("send message to %s with %s ok", user.Name, metric.Name)
			}
			if len(users) != 0 {
				al.m.Set(metric.Name, metric.Stamp)
				health.IncrNumAlertingEvents(1)
			}
		}
	}
}
<|endoftext|>


<commit_before>package plugin_status

import (
    "fmt"
    "github.com/thoj/go-ircevent"

    "github.com/coredump-ch/moss/conf"
    "github.com/coredump-ch/moss/plugin"
)

func printStatus(args string, e *irc.Event, con *irc.Connection) error {
    fmt.Println(e)
    con.Privmsg(conf.Channel, "Status unknown.")
    return nil
}

func init() {
    plugin.RegisterCommand("status", "PRIVMSG", printStatus)
}
<commit_msg>Removed unneeded print statement<commit_after>package plugin_status

import (
    "github.com/thoj/go-ircevent"

    "github.com/coredump-ch/moss/conf"
    "github.com/coredump-ch/moss/plugin"
)

func printStatus(args string, e *irc.Event, con *irc.Connection) error {
    con.Privmsg(conf.Channel, "Status unknown.")
    return nil
}

func init() {
    plugin.RegisterCommand("status", "PRIVMSG", printStatus)
}
<|endoftext|>


<commit_before>package discordbot

import (
	"log"
	"math/rand"
	"strings"
	"time"
	"unicode"

	"github.com/bwmarrin/discordgo"
	"github.com/heetch/sqalx"
	uuid "github.com/satori/go.uuid"
	"github.com/underlx/disturbancesmlx/dataobjects"
	"golang.org/x/text/runes"
	"golang.org/x/text/transform"
	"golang.org/x/text/unicode/norm"
)

var wordMap map[string]wordType
var lightTriggersMap map[string]lightTrigger
var lightTriggersLastUsage map[lastUsageKey]time.Time // maps lightTrigger IDs to the last time they were used
var stopMute map[string]time.Time                     // maps channel IDs to the time when the bot can talk again
var channelMute map[string]bool

var node sqalx.Node
var websiteURL string
var botLog *log.Logger
var session *discordgo.Session
var schedToLines func(schedules []*dataobjects.LobbySchedule) []string
var statusCallback func(status *dataobjects.Status)

type lightTrigger struct {
	wordType wordType
	id       string
}

type lastUsageKey struct {
	id        string
	channelID string
}

var botOwnerUserID string

var footerMessages = []string{
	"$mute para me mandar ir dar uma volta de Metro",
	"$mute para me calar por 15 minutos",
	"$mute e fico caladinho",
	"Estou a ser chato? Simimimimimim? Então $mute",
	"$mute e também faço greve",
	"$mute e vou fazer queixinhas ao sindicato",
	"Inoportuno? Então $mute",
	"Pareço uma idiotice artificial? $mute nisso",
	"Chato para caraças? Diga $mute",
	"A tentar ter uma conversa séria? $mute e calo-me",
	"Estou demasiado extrovertido? $mute",
	"$mute para me pôr no silêncio",
	"$mute para me mandar para o castigo",
	"$mute para me mandar ver se está a chover",
}

// wordType corresponds to a type of bot trigger word
type wordType int

const (
	wordTypeNetwork = iota
	wordTypeLine    = iota
	wordTypeStation = iota
	wordTypeLobby   = iota
	wordTypePOI     = iota
)

// Start starts the Discord bot
func Start(snode sqalx.Node, swebsiteURL, discordToken string, log *log.Logger,
	schedulesToLines func(schedules []*dataobjects.LobbySchedule) []string,
	statusCb func(status *dataobjects.Status)) error {
	node = snode
	websiteURL = swebsiteURL
	botLog = log
	schedToLines = schedulesToLines
	statusCallback = statusCb
	channelMute = make(map[string]bool)
	rand.Seed(time.Now().Unix())
	dg, err := discordgo.New("Bot " + discordToken)
	if err != nil {
		return err
	}
	session = dg

	selfApp, err := dg.Application("@me")
	if err != nil {
		return err
	}
	botOwnerUserID = selfApp.Owner.ID

	err = buildWordMap()
	if err != nil {
		return err
	}

	stopMute = make(map[string]time.Time)

	user, err := dg.User("@me")
	if err != nil {
		return err
	}
	if user.Username != "UnderLX" {
		_, err := dg.UserUpdate("", "", "UnderLX", "", "")
		if err != nil {
			return err
		}
	}
	dg.AddHandler(messageCreate)
	// Open a websocket connection to Discord and begin listening.
	return dg.Open()
}

// Stop stops the Discord bot
func Stop() {
	// Cleanly close down the Discord session.
	if session != nil {
		session.Close()
	}
}

func buildWordMap() error {
	wordMap = make(map[string]wordType)
	lightTriggersMap = make(map[string]lightTrigger)
	lightTriggersLastUsage = make(map[lastUsageKey]time.Time)

	tx, err := node.Beginx()
	if err != nil {
		return err
	}
	defer tx.Commit() // read-only tx

	networks, err := dataobjects.GetNetworks(tx)
	if err != nil {
		return err
	}
	for _, network := range networks {
		wordMap[network.ID] = wordTypeNetwork
	}

	lines, err := dataobjects.GetLines(tx)
	if err != nil {
		return err
	}
	for _, line := range lines {
		wordMap[line.ID] = wordTypeLine
		lightTriggersMap["linha "+line.Name] = lightTrigger{
			wordType: wordTypeLine,
			id:       line.ID}
	}

	stations, err := dataobjects.GetStations(tx)
	if err != nil {
		return err
	}
	for _, station := range stations {
		wordMap[station.ID] = wordTypeStation
		triggers := []string{
			"estação do " + station.Name,
			"estação da " + station.Name,
			"estação de " + station.Name,
			"estação " + station.Name,
		}
		for _, trigger := range triggers {
			lightTriggersMap[trigger] = lightTrigger{
				wordType: wordTypeStation,
				id:       station.ID}
		}
	}

	lobbies, err := dataobjects.GetLobbies(tx)
	if err != nil {
		return err
	}
	for _, lobby := range lobbies {
		wordMap[lobby.ID] = wordTypeLobby
	}

	pois, err := dataobjects.GetPOIs(tx)
	if err != nil {
		return err
	}
	for _, poi := range pois {
		wordMap[poi.ID] = wordTypePOI
	}

	return nil
}

// This function will be called (due to AddHandler above) every time a new
// message is created on any channel that the autenticated bot has access to.
func messageCreate(s *discordgo.Session, m *discordgo.MessageCreate) {
	// Ignore all messages created by the bot itself
	if m.Author.ID == s.State.User.ID {
		return
	}

	words := strings.Split(m.Content, " ")

	if parseCommands(s, m, words) {
		return
	}

	if !time.Now().After(stopMute[m.ChannelID]) || channelMute[m.ChannelID] {
		return
	}
	for _, word := range words {
		if wordType, ok := wordMap[word]; ok {
			sendReply(s, m, word, word, wordType)
		}
	}

	for lightTrigger, triggerInfo := range lightTriggersMap {
		if !strings.Contains(lightTrigger, " ") && len(m.Content) > len(lightTrigger) {
			lightTrigger = " " + lightTrigger + " "
		}
		t := transform.Chain(norm.NFD, runes.Remove(runes.In(unicode.Mn)), norm.NFC)
		noDiacriticsResult, _, _ := transform.String(t, lightTrigger)
		noDiacriticsMessage, _, _ := transform.String(t, strings.ToLower(m.Content))
		triggerWord := ""
		if strings.Contains(m.Content, lightTrigger) {
			triggerWord = strings.TrimSpace(lightTrigger)
		} else if needle := strings.ToLower(lightTrigger); strings.Contains(m.Content, needle) {
			triggerWord = strings.TrimSpace(needle)
		} else if needle := strings.ToLower(noDiacriticsResult); strings.Contains(m.Content, needle) {
			triggerWord = strings.TrimSpace(needle)
		} else if needle := strings.ToLower(noDiacriticsResult); strings.Contains(noDiacriticsMessage, needle) {
			triggerWord = strings.TrimSpace(needle)
		} else if needle := strings.ToLower(strings.TrimRight(noDiacriticsResult, " ")); strings.HasSuffix(noDiacriticsMessage, needle) {
			triggerWord = strings.TrimSpace(needle)
		} else if needle := strings.ToLower(strings.TrimLeft(noDiacriticsResult, " ")); strings.HasPrefix(noDiacriticsMessage, needle) {
			triggerWord = strings.TrimSpace(needle)
		}
		if triggerWord != "" {
			key := lastUsageKey{
				channelID: m.ChannelID,
				id:        triggerInfo.id}
			if t, ok := lightTriggersLastUsage[key]; ok && time.Since(t) < 10*time.Minute {
				continue
			}
			lightTriggersLastUsage[key] = time.Now()
			sendReply(s, m, triggerInfo.id, triggerWord, triggerInfo.wordType)
		}
	}
}

func parseCommands(s *discordgo.Session, m *discordgo.MessageCreate, words []string) bool {
	// whole-message commands
	switch m.Content {
	case "$mute":
		stopMute[m.ChannelID] = time.Now().Add(15 * time.Minute)
		s.ChannelMessageSend(m.ChannelID, "🤐 por 15 minutos")
		return true
	case "$unmute":
		stopMute[m.ChannelID] = time.Time{}
		s.ChannelMessageSend(m.ChannelID, "🤗")
		return true
	}

	if m.Author.ID == botOwnerUserID {
		switch words[0] {
		case "$setstatus":
			handleStatus(s, m, words[1:])
			return true
		case "$addlinestatus":
			handleLineStatus(s, m, words[1:])
			return true
		case "$permamute":
			channelMute[m.ChannelID] = true
			return true
		case "$permaunmute":
			channelMute[m.ChannelID] = false
			return true
		}
	}
	return false
}

func handleLineStatus(s *discordgo.Session, m *discordgo.MessageCreate, words []string) {
	if len(words) < 3 {
		s.ChannelMessageSend(m.ChannelID, "🆖 missing arguments")
		return
	}
	id, err := uuid.NewV4()
	if err != nil {
		s.ChannelMessageSend(m.ChannelID, "❌ "+err.Error())
		return
	}

	tx, err := node.Beginx()
	if err != nil {
		s.ChannelMessageSend(m.ChannelID, "❌ "+err.Error())
		return
	}
	defer tx.Commit() // read-only tx

	status := &dataobjects.Status{
		ID:   id.String(),
		Time: time.Now().UTC(),
		Source: &dataobjects.Source{
			ID:        "underlx-bot",
			Name:      "UnderLX Discord bot",
			Automatic: false,
			Official:  false,
		},
	}

	switch words[0] {
	case "up":
		status.IsDowntime = false
	case "down":
		status.IsDowntime = true
	default:
		s.ChannelMessageSend(m.ChannelID, "🆖 first argument must be `up` or `down`")
		return
	}

	line, err := dataobjects.GetLine(tx, words[1])
	if err != nil {
		lines, err := dataobjects.GetLines(tx)
		if err != nil {
			s.ChannelMessageSend(m.ChannelID, "❌ "+err.Error())
			return
		}
		lineIDs := make([]string, len(lines))
		for i := range lines {
			lineIDs[i] = "`" + lines[i].ID + "`"
		}
		s.ChannelMessageSend(m.ChannelID, "🆖 line ID must be one of ["+strings.Join(lineIDs, ",")+"]")
		return
	}

	status.Line = line
	status.Status = strings.Join(words[2:], " ")

	statusCallback(status)
	s.ChannelMessageSend(m.ChannelID, "✅")
}

func handleStatus(s *discordgo.Session, m *discordgo.MessageCreate, words []string) {
	var err error
	if len(words) == 0 {
		err = s.UpdateStatus(0, "")
	} else if len(words) > 1 {
		usd := &discordgo.UpdateStatusData{
			Status: "online",
		}

		switch words[1] {
		case "playing":
			usd.Game = &discordgo.Game{
				Name: strings.Join(words[2:], " "),
				Type: discordgo.GameTypeGame,
			}
		case "streaming":
			usd.Game = &discordgo.Game{
				Type: discordgo.GameTypeGame,
				URL:  strings.Join(words[2:], " "),
			}
		case "listening":
			usd.Game = &discordgo.Game{
				Name: strings.Join(words[2:], " "),
				Type: discordgo.GameTypeListening,
			}
		case "watching":
			usd.Game = &discordgo.Game{
				Name: strings.Join(words[2:], " "),
				Type: discordgo.GameTypeWatching,
			}
		default:
			usd.Game = &discordgo.Game{
				Name: strings.Join(words[1:], " "),
				Type: discordgo.GameTypeGame,
			}
		}

		err = s.UpdateStatusComplex(*usd)
	}
	if err != nil {
		s.ChannelMessageSend(m.ChannelID, "❌ "+err.Error())
	} else {
		s.ChannelMessageSend(m.ChannelID, "✅")
	}
}

func sendReply(s *discordgo.Session, m *discordgo.MessageCreate, trigger, origTrigger string, triggerType wordType) {
	var embed *Embed
	var err error
	switch triggerType {
	case wordTypeNetwork:
		embed, err = buildNetworkMessage(trigger)
	case wordTypeLine:
		embed, err = buildLineMessage(trigger)
	case wordTypeStation:
		embed, err = buildStationMessage(trigger)
	case wordTypeLobby:
		embed, err = buildLobbyMesage(trigger)
	}

	if err == nil && embed != nil {
		embed.SetFooter(origTrigger + " | " + footerMessages[rand.Intn(len(footerMessages))])
		embed.Timestamp = time.Now().Format(time.RFC3339Nano)
		s.ChannelMessageSendEmbed(m.ChannelID, embed.MessageEmbed)
	} else {
		botLog.Println(err)
	}
}

func getEmojiForLine(id string) string {
	switch id {
	case "pt-ml-azul":
		return "<:ml_azul:410577265420795904>"
	case "pt-ml-amarela":
		return "<:ml_amarela:410566925114933250>"
	case "pt-ml-verde":
		return "<:ml_verde:410577778862325764>"
	case "pt-ml-vermelha":
		return "<:ml_vermelha:410579362773991424>"
	}
	return ""
}
<commit_msg>Fix bug in discord bot<commit_after>package discordbot

import (
	"log"
	"math/rand"
	"strings"
	"time"
	"unicode"

	"github.com/bwmarrin/discordgo"
	"github.com/heetch/sqalx"
	uuid "github.com/satori/go.uuid"
	"github.com/underlx/disturbancesmlx/dataobjects"
	"golang.org/x/text/runes"
	"golang.org/x/text/transform"
	"golang.org/x/text/unicode/norm"
)

var wordMap map[string]wordType
var lightTriggersMap map[string]lightTrigger
var lightTriggersLastUsage map[lastUsageKey]time.Time // maps lightTrigger IDs to the last time they were used
var stopMute map[string]time.Time                     // maps channel IDs to the time when the bot can talk again
var channelMute map[string]bool

var node sqalx.Node
var websiteURL string
var botLog *log.Logger
var session *discordgo.Session
var schedToLines func(schedules []*dataobjects.LobbySchedule) []string
var statusCallback func(status *dataobjects.Status)

type lightTrigger struct {
	wordType wordType
	id       string
}

type lastUsageKey struct {
	id        string
	channelID string
}

var botOwnerUserID string

var footerMessages = []string{
	"$mute para me mandar ir dar uma volta de Metro",
	"$mute para me calar por 15 minutos",
	"$mute e fico caladinho",
	"Estou a ser chato? Simimimimimim? Então $mute",
	"$mute e também faço greve",
	"$mute e vou fazer queixinhas ao sindicato",
	"Inoportuno? Então $mute",
	"Pareço uma idiotice artificial? $mute nisso",
	"Chato para caraças? Diga $mute",
	"A tentar ter uma conversa séria? $mute e calo-me",
	"Estou demasiado extrovertido? $mute",
	"$mute para me pôr no silêncio",
	"$mute para me mandar para o castigo",
	"$mute para me mandar ver se está a chover",
}

// wordType corresponds to a type of bot trigger word
type wordType int

const (
	wordTypeNetwork = iota
	wordTypeLine    = iota
	wordTypeStation = iota
	wordTypeLobby   = iota
	wordTypePOI     = iota
)

// Start starts the Discord bot
func Start(snode sqalx.Node, swebsiteURL, discordToken string, log *log.Logger,
	schedulesToLines func(schedules []*dataobjects.LobbySchedule) []string,
	statusCb func(status *dataobjects.Status)) error {
	node = snode
	websiteURL = swebsiteURL
	botLog = log
	schedToLines = schedulesToLines
	statusCallback = statusCb
	channelMute = make(map[string]bool)
	rand.Seed(time.Now().Unix())
	dg, err := discordgo.New("Bot " + discordToken)
	if err != nil {
		return err
	}
	session = dg

	selfApp, err := dg.Application("@me")
	if err != nil {
		return err
	}
	botOwnerUserID = selfApp.Owner.ID

	err = buildWordMap()
	if err != nil {
		return err
	}

	stopMute = make(map[string]time.Time)

	user, err := dg.User("@me")
	if err != nil {
		return err
	}
	if user.Username != "UnderLX" {
		_, err := dg.UserUpdate("", "", "UnderLX", "", "")
		if err != nil {
			return err
		}
	}
	dg.AddHandler(messageCreate)
	// Open a websocket connection to Discord and begin listening.
	return dg.Open()
}

// Stop stops the Discord bot
func Stop() {
	// Cleanly close down the Discord session.
	if session != nil {
		session.Close()
	}
}

func buildWordMap() error {
	wordMap = make(map[string]wordType)
	lightTriggersMap = make(map[string]lightTrigger)
	lightTriggersLastUsage = make(map[lastUsageKey]time.Time)

	tx, err := node.Beginx()
	if err != nil {
		return err
	}
	defer tx.Commit() // read-only tx

	networks, err := dataobjects.GetNetworks(tx)
	if err != nil {
		return err
	}
	for _, network := range networks {
		wordMap[network.ID] = wordTypeNetwork
	}

	lines, err := dataobjects.GetLines(tx)
	if err != nil {
		return err
	}
	for _, line := range lines {
		wordMap[line.ID] = wordTypeLine
		lightTriggersMap["linha "+line.Name] = lightTrigger{
			wordType: wordTypeLine,
			id:       line.ID}
	}

	stations, err := dataobjects.GetStations(tx)
	if err != nil {
		return err
	}
	for _, station := range stations {
		wordMap[station.ID] = wordTypeStation
		triggers := []string{
			"estação do " + station.Name,
			"estação da " + station.Name,
			"estação de " + station.Name,
			"estação " + station.Name,
		}
		for _, trigger := range triggers {
			lightTriggersMap[trigger] = lightTrigger{
				wordType: wordTypeStation,
				id:       station.ID}
		}
	}

	lobbies, err := dataobjects.GetLobbies(tx)
	if err != nil {
		return err
	}
	for _, lobby := range lobbies {
		wordMap[lobby.ID] = wordTypeLobby
	}

	pois, err := dataobjects.GetPOIs(tx)
	if err != nil {
		return err
	}
	for _, poi := range pois {
		wordMap[poi.ID] = wordTypePOI
	}

	return nil
}

// This function will be called (due to AddHandler above) every time a new
// message is created on any channel that the autenticated bot has access to.
func messageCreate(s *discordgo.Session, m *discordgo.MessageCreate) {
	// Ignore all messages created by the bot itself
	if m.Author.ID == s.State.User.ID {
		return
	}

	words := strings.Split(m.Content, " ")

	if parseCommands(s, m, words) {
		return
	}

	if !time.Now().After(stopMute[m.ChannelID]) || channelMute[m.ChannelID] {
		return
	}
	for _, word := range words {
		if wordType, ok := wordMap[word]; ok {
			sendReply(s, m, word, word, wordType)
		}
	}

	for lightTrigger, triggerInfo := range lightTriggersMap {
		if !strings.Contains(lightTrigger, " ") && len(m.Content) > len(lightTrigger) {
			lightTrigger = " " + lightTrigger + " "
		}
		t := transform.Chain(norm.NFD, runes.Remove(runes.In(unicode.Mn)), norm.NFC)
		noDiacriticsResult, _, _ := transform.String(t, lightTrigger)
		noDiacriticsMessage, _, _ := transform.String(t, strings.ToLower(m.Content))
		triggerWord := ""
		if strings.Contains(m.Content, lightTrigger) {
			triggerWord = strings.TrimSpace(lightTrigger)
		} else if needle := strings.ToLower(lightTrigger); strings.Contains(m.Content, needle) {
			triggerWord = strings.TrimSpace(needle)
		} else if needle := strings.ToLower(noDiacriticsResult); strings.Contains(m.Content, needle) {
			triggerWord = strings.TrimSpace(needle)
		} else if needle := strings.ToLower(noDiacriticsResult); strings.Contains(noDiacriticsMessage, needle) {
			triggerWord = strings.TrimSpace(needle)
		} else if needle := strings.ToLower(strings.TrimRight(noDiacriticsResult, " ")); strings.HasSuffix(noDiacriticsMessage, needle) {
			triggerWord = strings.TrimSpace(needle)
		} else if needle := strings.ToLower(strings.TrimLeft(noDiacriticsResult, " ")); strings.HasPrefix(noDiacriticsMessage, needle) {
			triggerWord = strings.TrimSpace(needle)
		}
		if triggerWord != "" {
			key := lastUsageKey{
				channelID: m.ChannelID,
				id:        triggerInfo.id}
			if t, ok := lightTriggersLastUsage[key]; ok && time.Since(t) < 10*time.Minute {
				continue
			}
			lightTriggersLastUsage[key] = time.Now()
			sendReply(s, m, triggerInfo.id, triggerWord, triggerInfo.wordType)
		}
	}
}

func parseCommands(s *discordgo.Session, m *discordgo.MessageCreate, words []string) bool {
	// whole-message commands
	switch m.Content {
	case "$mute":
		stopMute[m.ChannelID] = time.Now().Add(15 * time.Minute)
		s.ChannelMessageSend(m.ChannelID, "🤐 por 15 minutos")
		return true
	case "$unmute":
		stopMute[m.ChannelID] = time.Time{}
		s.ChannelMessageSend(m.ChannelID, "🤗")
		return true
	}

	if m.Author.ID == botOwnerUserID {
		switch words[0] {
		case "$setstatus":
			handleStatus(s, m, words[1:])
			return true
		case "$addlinestatus":
			handleLineStatus(s, m, words[1:])
			return true
		case "$permamute":
			channelMute[m.ChannelID] = true
			return true
		case "$permaunmute":
			channelMute[m.ChannelID] = false
			return true
		}
	}
	return false
}

func handleLineStatus(s *discordgo.Session, m *discordgo.MessageCreate, words []string) {
	if len(words) < 3 {
		s.ChannelMessageSend(m.ChannelID, "🆖 missing arguments")
		return
	}
	id, err := uuid.NewV4()
	if err != nil {
		s.ChannelMessageSend(m.ChannelID, "❌ "+err.Error())
		return
	}

	tx, err := node.Beginx()
	if err != nil {
		s.ChannelMessageSend(m.ChannelID, "❌ "+err.Error())
		return
	}
	defer tx.Commit() // read-only tx

	status := &dataobjects.Status{
		ID:   id.String(),
		Time: time.Now().UTC(),
		Source: &dataobjects.Source{
			ID:        "underlx-bot",
			Name:      "UnderLX Discord bot",
			Automatic: false,
			Official:  false,
		},
	}

	switch words[0] {
	case "up":
		status.IsDowntime = false
	case "down":
		status.IsDowntime = true
	default:
		s.ChannelMessageSend(m.ChannelID, "🆖 first argument must be `up` or `down`")
		return
	}

	line, err := dataobjects.GetLine(tx, words[1])
	if err != nil {
		lines, err := dataobjects.GetLines(tx)
		if err != nil {
			s.ChannelMessageSend(m.ChannelID, "❌ "+err.Error())
			return
		}
		lineIDs := make([]string, len(lines))
		for i := range lines {
			lineIDs[i] = "`" + lines[i].ID + "`"
		}
		s.ChannelMessageSend(m.ChannelID, "🆖 line ID must be one of ["+strings.Join(lineIDs, ",")+"]")
		return
	}

	status.Line = line
	status.Status = strings.Join(words[2:], " ")

	statusCallback(status)
	s.ChannelMessageSend(m.ChannelID, "✅")
}

func handleStatus(s *discordgo.Session, m *discordgo.MessageCreate, words []string) {
	var err error
	if len(words) == 0 {
		err = s.UpdateStatus(0, "")
	} else if len(words) > 0 {
		usd := &discordgo.UpdateStatusData{
			Status: "online",
		}

		switch words[0] {
		case "playing":
			usd.Game = &discordgo.Game{
				Name: strings.Join(words[1:], " "),
				Type: discordgo.GameTypeGame,
			}
		case "streaming":
			usd.Game = &discordgo.Game{
				Type: discordgo.GameTypeGame,
				URL:  strings.Join(words[1:], " "),
			}
		case "listening":
			usd.Game = &discordgo.Game{
				Name: strings.Join(words[1:], " "),
				Type: discordgo.GameTypeListening,
			}
		case "watching":
			usd.Game = &discordgo.Game{
				Name: strings.Join(words[1:], " "),
				Type: discordgo.GameTypeWatching,
			}
		default:
			usd.Game = &discordgo.Game{
				Name: strings.Join(words, " "),
				Type: discordgo.GameTypeGame,
			}
		}

		err = s.UpdateStatusComplex(*usd)
	}
	if err != nil {
		s.ChannelMessageSend(m.ChannelID, "❌ "+err.Error())
	} else {
		s.ChannelMessageSend(m.ChannelID, "✅")
	}
}

func sendReply(s *discordgo.Session, m *discordgo.MessageCreate, trigger, origTrigger string, triggerType wordType) {
	var embed *Embed
	var err error
	switch triggerType {
	case wordTypeNetwork:
		embed, err = buildNetworkMessage(trigger)
	case wordTypeLine:
		embed, err = buildLineMessage(trigger)
	case wordTypeStation:
		embed, err = buildStationMessage(trigger)
	case wordTypeLobby:
		embed, err = buildLobbyMesage(trigger)
	}

	if err == nil && embed != nil {
		embed.SetFooter(origTrigger + " | " + footerMessages[rand.Intn(len(footerMessages))])
		embed.Timestamp = time.Now().Format(time.RFC3339Nano)
		s.ChannelMessageSendEmbed(m.ChannelID, embed.MessageEmbed)
	} else {
		botLog.Println(err)
	}
}

func getEmojiForLine(id string) string {
	switch id {
	case "pt-ml-azul":
		return "<:ml_azul:410577265420795904>"
	case "pt-ml-amarela":
		return "<:ml_amarela:410566925114933250>"
	case "pt-ml-verde":
		return "<:ml_verde:410577778862325764>"
	case "pt-ml-vermelha":
		return "<:ml_vermelha:410579362773991424>"
	}
	return ""
}
<|endoftext|>