<commit_before>// Copyright 2015 ikawaha
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// 	You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package server

import (
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"html/template"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"os/exec"
	"strings"
	"time"

	"github.com/ikawaha/kagome/tokenizer"
)

// subcommand property
var (
	CommandName  = "server"
	Description  = `run tokenize server`
	usageMessage = "%s [-http=:6060] [-udic userdic_file] [-mode (normal|search|extended)]\n"
	ErrorWriter  = os.Stderr
)

// options
type option struct {
	http    string
	udic    string
	mode    string
	flagSet *flag.FlagSet
}

// ContinueOnError ErrorHandling // Return a descriptive error.
// ExitOnError                   // Call os.Exit(2).
// PanicOnError                  // Call panic with a descriptive error.flag.ContinueOnError
func newOption(w io.Writer, eh flag.ErrorHandling) (o *option) {
	o = &option{
		flagSet: flag.NewFlagSet(CommandName, eh),
	}
	// option settings
	o.flagSet.StringVar(&o.http, "http", ":6060", "HTTP service address (e.g., ':6060')")
	o.flagSet.StringVar(&o.udic, "udic", "", "user dictionary")
	o.flagSet.StringVar(&o.mode, "mode", "normal", "tokenize mode (normal|search|extended)")

	return
}

func (o *option) parse(args []string) (err error) {
	if err = o.flagSet.Parse(args); err != nil {
		return
	}
	// validations
	if nonFlag := o.flagSet.Args(); len(nonFlag) != 0 {
		return fmt.Errorf("invalid argument: %v", nonFlag)
	}
	if o.mode != "normal" && o.mode != "search" && o.mode != "extended" {
		return fmt.Errorf("unknown mode: %v", o.mode)
	}
	return
}

func OptionCheck(args []string) (err error) {
	opt := newOption(ioutil.Discard, flag.ContinueOnError)
	if e := opt.parse(args); e != nil {
		return fmt.Errorf("%v, %v", CommandName, e)
	}
	return nil
}

// command main
func command(opt *option) error {
	var udic tokenizer.UserDic
	if opt.udic != "" {
		var err error
		if udic, err = tokenizer.NewUserDic(opt.udic); err != nil {
			return err
		}
	}
	t := tokenizer.New()
	t.SetUserDic(udic)

	mux := http.NewServeMux()
	mux.Handle("/", &TokenizeDemoHandler{tokenizer: t})
	mux.Handle("/a", &TokenizeHandler{tokenizer: t})
	log.Fatal(http.ListenAndServe(opt.http, mux))

	return nil
}

type TokenizeHandler struct {
	tokenizer tokenizer.Tokenizer
}

func (h *TokenizeHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	type record struct {
		ID       int      `json:"id"`
		Start    int      `json:"start"`
		End      int      `json:"end"`
		Surface  string   `json:"surface"`
		Class    string   `json:"class"`
		Features []string `json:"features"`
	}

	var body struct {
		Input string `json:"sentence"`
	}
	e := json.NewDecoder(r.Body).Decode(&body)
	if e != nil {
		fmt.Fprintf(w, "{\"status\":false,\"error\":\"%v\"}", e)
		return
	}
	if body.Input == "" {
		fmt.Fprint(w, "{\"status\":true,\"tokens\":[]}")
		return
	}
	tokens := h.tokenizer.Analyze(body.Input, tokenizer.Normal)
	var rsp []record
	for _, tok := range tokens {
		if tok.ID == tokenizer.BosEosID {
			continue
		}
		fs := tok.Features()
		m := record{
			ID:       tok.ID,
			Class:    fmt.Sprintf("%v", tok.Class),
			Start:    tok.Start,
			End:      tok.End,
			Surface:  tok.Surface,
			Features: fs,
		}
		rsp = append(rsp, m)
	}
	j, e := json.Marshal(struct {
		Status bool     `json:"status"`
		Tokens []record `json:"tokens"`
	}{Status: true, Tokens: rsp})
	if e != nil {
		fmt.Fprintf(w, "{\"status\":false,\"error\":\"%v\"}", e)
		return
	}
	if _, e := w.Write(j); e != nil {
		fmt.Fprintf(w, "{\"status\":false,\"error\":\"%v\"}", e)
		return
	}
}

type TokenizeDemoHandler struct {
	tokenizer tokenizer.Tokenizer
}

func (h *TokenizeDemoHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	type record struct {
		Surface        string
		Pos            string
		Baseform       string
		Reading        string
		Pronounciation string
	}
	sen := r.FormValue("s")
	opt := r.FormValue("r")

	const cmdTimeout = 15 * time.Second
	var (
		records []record
		tokens  []tokenizer.Token
		svg     string
		cmdErr  string
	)

	switch opt {
	case "normal":
		tokens = h.tokenizer.Analyze(sen, tokenizer.Normal)
	case "search":
		tokens = h.tokenizer.Analyze(sen, tokenizer.Search)
	case "extended":
		tokens = h.tokenizer.Analyze(sen, tokenizer.Extended)
	case "lattice":
		if _, e := exec.LookPath("dot"); e != nil {
			cmdErr = "Error: graphviz is not in your furure"
			log.Print("graphviz is not in your future\n")
			break
		}
		var buf bytes.Buffer
		cmd := exec.Command("dot", "-Tsvg")
		r, w := io.Pipe()
		cmd.Stdin = r
		cmd.Stdout = &buf
		cmd.Stderr = ErrorWriter
		if err := cmd.Start(); err != nil {
			cmdErr = "Error"
			log.Printf("process done with error = %v", err)
		}
		h.tokenizer.Dot(sen, w)
		w.Close()

		done := make(chan error, 1)
		go func() {
			done <- cmd.Wait()
		}()
		select {
		case <-time.After(cmdTimeout):
			if err := cmd.Process.Kill(); err != nil {
				log.Fatal("failed to kill: ", err)
			}
			cmdErr = "Time out"
			<-done
		case err := <-done:
			if err != nil {
				cmdErr = "Error"
				log.Printf("process done with error = %v", err)
			}
		}
		svg = buf.String()
		if pos := strings.Index(svg, "<svg"); pos > 0 {
			svg = svg[pos:]
		}
	}
	for _, tok := range tokens {
		if tok.ID == tokenizer.BosEosID {
			continue
		}
		m := record{Surface: tok.Surface}
		fs := tok.Features()
		switch len(fs) {
		case 9:
			m.Pos = strings.Join(fs[0:5], ",")
			m.Baseform = fs[6]
			m.Reading = fs[7]
			m.Pronounciation = fs[8]
		case 7:
			m.Pos = strings.Join(fs[0:5], ",")
			m.Baseform = fs[6]
			m.Reading = "*"
			m.Pronounciation = "*"
		case 3:
			m.Pos = fs[0]
			m.Baseform = fs[1]
			m.Reading = fs[2]
			m.Pronounciation = "*"
		}
		records = append(records, m)
	}
	d := struct {
		Sentence string
		Tokens   []record
		CmdErr   string
		GraphSvg template.HTML
		RadioOpt string
	}{Sentence: sen, Tokens: records, CmdErr: cmdErr, GraphSvg: template.HTML(svg), RadioOpt: opt}
	t := template.Must(template.New("top").Parse(demoHTML))
	if e := t.Execute(w, d); e != nil {
		http.Error(w, e.Error(), http.StatusInternalServerError)
	}
}

func Run(args []string) error {
	opt := newOption(ErrorWriter, flag.ExitOnError)
	if e := opt.parse(args); e != nil {
		Usage()
		PrintDefaults(flag.ExitOnError)
		return fmt.Errorf("%v, %v", CommandName, e)
	}
	return command(opt)
}

func Usage() {
	fmt.Fprintf(os.Stderr, usageMessage, CommandName)
}

func PrintDefaults(eh flag.ErrorHandling) {
	o := newOption(ErrorWriter, eh)
	o.flagSet.PrintDefaults()
}

var demoHTML = `
<!DOCTYPE html>

  <meta charset="UTF-8">
  <title>Kagome demo - japanese morphological analyzer</title>

Kagome

    Kagome is an open source Japanese morphological analyzer written in Golang

Feature summary

Word segmentation. Segmenting text into words (or morphemes)

Part-of-speech tagging. Assign word-categories (nouns, verbs, particles, adjectives, etc.)

Lemmatization. Get dictionary forms for inflected verbs and adjectives

Readings. Extract readings for kanji.

  <form class="frm" action="/_demo" method="POST">

    <textarea class="txar" rows="3" name="s" placeholder="Enter Japanese text below in UTF-8 and click tokenize.">{{.Sentence}}</textarea>

<input type="radio" name="r" value="normal" checked>Normal

<input type="radio" name="r" value="search" {{if eq .RadioOpt "search"}}checked{{end}}>Search

<input type="radio" name="r" value="extended" {{if eq .RadioOpt "extended"}}checked{{end}}>Extended

<input type="radio" name="r" value="lattice" {{if eq .RadioOpt "lattice"}}checked{{end}}>Lattice

<input class="btn" type="submit" value="Tokenize"/>

  </form>
  {{if .CmdErr}}
    {{.CmdErr}}
  {{end}}
  {{if .GraphSvg}}
    {{.GraphSvg}}
  {{end}}
  {{if .Tokens}}

    <thread>

      Surface
      Part-of-Speech
      Base Form
      Reading
      Pronounciation

</thread>

    {{range .Tokens}}

      {{.Surface}}
      {{.Pos}}
      {{.Baseform}}
      {{.Reading}}
      {{.Pronounciation}}

    {{end}}

  {{end}}

`
<commit_msg>Fix an error message of graphviz time out<commit_after>// Copyright 2015 ikawaha
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// 	You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package server

import (
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"html/template"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"os/exec"
	"strings"
	"time"

	"github.com/ikawaha/kagome/tokenizer"
)

// subcommand property
var (
	CommandName  = "server"
	Description  = `run tokenize server`
	usageMessage = "%s [-http=:6060] [-udic userdic_file] [-mode (normal|search|extended)]\n"
	ErrorWriter  = os.Stderr
)

// options
type option struct {
	http    string
	udic    string
	mode    string
	flagSet *flag.FlagSet
}

// ContinueOnError ErrorHandling // Return a descriptive error.
// ExitOnError                   // Call os.Exit(2).
// PanicOnError                  // Call panic with a descriptive error.flag.ContinueOnError
func newOption(w io.Writer, eh flag.ErrorHandling) (o *option) {
	o = &option{
		flagSet: flag.NewFlagSet(CommandName, eh),
	}
	// option settings
	o.flagSet.StringVar(&o.http, "http", ":6060", "HTTP service address (e.g., ':6060')")
	o.flagSet.StringVar(&o.udic, "udic", "", "user dictionary")
	o.flagSet.StringVar(&o.mode, "mode", "normal", "tokenize mode (normal|search|extended)")

	return
}

func (o *option) parse(args []string) (err error) {
	if err = o.flagSet.Parse(args); err != nil {
		return
	}
	// validations
	if nonFlag := o.flagSet.Args(); len(nonFlag) != 0 {
		return fmt.Errorf("invalid argument: %v", nonFlag)
	}
	if o.mode != "normal" && o.mode != "search" && o.mode != "extended" {
		return fmt.Errorf("unknown mode: %v", o.mode)
	}
	return
}

func OptionCheck(args []string) (err error) {
	opt := newOption(ioutil.Discard, flag.ContinueOnError)
	if e := opt.parse(args); e != nil {
		return fmt.Errorf("%v, %v", CommandName, e)
	}
	return nil
}

// command main
func command(opt *option) error {
	var udic tokenizer.UserDic
	if opt.udic != "" {
		var err error
		if udic, err = tokenizer.NewUserDic(opt.udic); err != nil {
			return err
		}
	}
	t := tokenizer.New()
	t.SetUserDic(udic)

	mux := http.NewServeMux()
	mux.Handle("/", &TokenizeDemoHandler{tokenizer: t})
	mux.Handle("/a", &TokenizeHandler{tokenizer: t})
	log.Fatal(http.ListenAndServe(opt.http, mux))

	return nil
}

type TokenizeHandler struct {
	tokenizer tokenizer.Tokenizer
}

func (h *TokenizeHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	type record struct {
		ID       int      `json:"id"`
		Start    int      `json:"start"`
		End      int      `json:"end"`
		Surface  string   `json:"surface"`
		Class    string   `json:"class"`
		Features []string `json:"features"`
	}

	var body struct {
		Input string `json:"sentence"`
	}
	e := json.NewDecoder(r.Body).Decode(&body)
	if e != nil {
		fmt.Fprintf(w, "{\"status\":false,\"error\":\"%v\"}", e)
		return
	}
	if body.Input == "" {
		fmt.Fprint(w, "{\"status\":true,\"tokens\":[]}")
		return
	}
	tokens := h.tokenizer.Analyze(body.Input, tokenizer.Normal)
	var rsp []record
	for _, tok := range tokens {
		if tok.ID == tokenizer.BosEosID {
			continue
		}
		fs := tok.Features()
		m := record{
			ID:       tok.ID,
			Class:    fmt.Sprintf("%v", tok.Class),
			Start:    tok.Start,
			End:      tok.End,
			Surface:  tok.Surface,
			Features: fs,
		}
		rsp = append(rsp, m)
	}
	j, e := json.Marshal(struct {
		Status bool     `json:"status"`
		Tokens []record `json:"tokens"`
	}{Status: true, Tokens: rsp})
	if e != nil {
		fmt.Fprintf(w, "{\"status\":false,\"error\":\"%v\"}", e)
		return
	}
	if _, e := w.Write(j); e != nil {
		fmt.Fprintf(w, "{\"status\":false,\"error\":\"%v\"}", e)
		return
	}
}

type TokenizeDemoHandler struct {
	tokenizer tokenizer.Tokenizer
}

func (h *TokenizeDemoHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	type record struct {
		Surface        string
		Pos            string
		Baseform       string
		Reading        string
		Pronounciation string
	}
	sen := r.FormValue("s")
	opt := r.FormValue("r")

	const cmdTimeout = 15 * time.Second
	var (
		records []record
		tokens  []tokenizer.Token
		svg     string
		cmdErr  string
	)

	switch opt {
	case "normal":
		tokens = h.tokenizer.Analyze(sen, tokenizer.Normal)
	case "search":
		tokens = h.tokenizer.Analyze(sen, tokenizer.Search)
	case "extended":
		tokens = h.tokenizer.Analyze(sen, tokenizer.Extended)
	case "lattice":
		if _, e := exec.LookPath("dot"); e != nil {
			cmdErr = "Error: graphviz is not in your furure"
			log.Print("graphviz is not in your future\n")
			break
		}
		var buf bytes.Buffer
		cmd := exec.Command("dot", "-Tsvg")
		r, w := io.Pipe()
		cmd.Stdin = r
		cmd.Stdout = &buf
		cmd.Stderr = ErrorWriter
		if err := cmd.Start(); err != nil {
			cmdErr = "Error"
			log.Printf("process done with error = %v", err)
		}
		h.tokenizer.Dot(sen, w)
		w.Close()

		done := make(chan error, 1)
		go func() {
			done <- cmd.Wait()
		}()
		select {
		case <-time.After(cmdTimeout):
			if err := cmd.Process.Kill(); err != nil {
				log.Fatal("failed to kill: ", err)
			}
			cmdErr = "Error: Graphviz time out"
			<-done
		case err := <-done:
			if err != nil {
				cmdErr = "Error"
				log.Printf("process done with error = %v", err)
			}
		}
		svg = buf.String()
		if pos := strings.Index(svg, "<svg"); pos > 0 {
			svg = svg[pos:]
		}
	}
	for _, tok := range tokens {
		if tok.ID == tokenizer.BosEosID {
			continue
		}
		m := record{Surface: tok.Surface}
		fs := tok.Features()
		switch len(fs) {
		case 9:
			m.Pos = strings.Join(fs[0:5], ",")
			m.Baseform = fs[6]
			m.Reading = fs[7]
			m.Pronounciation = fs[8]
		case 7:
			m.Pos = strings.Join(fs[0:5], ",")
			m.Baseform = fs[6]
			m.Reading = "*"
			m.Pronounciation = "*"
		case 3:
			m.Pos = fs[0]
			m.Baseform = fs[1]
			m.Reading = fs[2]
			m.Pronounciation = "*"
		}
		records = append(records, m)
	}
	d := struct {
		Sentence string
		Tokens   []record
		CmdErr   string
		GraphSvg template.HTML
		RadioOpt string
	}{Sentence: sen, Tokens: records, CmdErr: cmdErr, GraphSvg: template.HTML(svg), RadioOpt: opt}
	t := template.Must(template.New("top").Parse(demoHTML))
	if e := t.Execute(w, d); e != nil {
		http.Error(w, e.Error(), http.StatusInternalServerError)
	}
}

func Run(args []string) error {
	opt := newOption(ErrorWriter, flag.ExitOnError)
	if e := opt.parse(args); e != nil {
		Usage()
		PrintDefaults(flag.ExitOnError)
		return fmt.Errorf("%v, %v", CommandName, e)
	}
	return command(opt)
}

func Usage() {
	fmt.Fprintf(os.Stderr, usageMessage, CommandName)
}

func PrintDefaults(eh flag.ErrorHandling) {
	o := newOption(ErrorWriter, eh)
	o.flagSet.PrintDefaults()
}

var demoHTML = `
<!DOCTYPE html>

  <meta charset="UTF-8">
  <title>Kagome demo - japanese morphological analyzer</title>

Kagome

    Kagome is an open source Japanese morphological analyzer written in Golang

Feature summary

Word segmentation. Segmenting text into words (or morphemes)

Part-of-speech tagging. Assign word-categories (nouns, verbs, particles, adjectives, etc.)

Lemmatization. Get dictionary forms for inflected verbs and adjectives

Readings. Extract readings for kanji.

  <form class="frm" action="/_demo" method="POST">

    <textarea class="txar" rows="3" name="s" placeholder="Enter Japanese text below in UTF-8 and click tokenize.">{{.Sentence}}</textarea>

<input type="radio" name="r" value="normal" checked>Normal

<input type="radio" name="r" value="search" {{if eq .RadioOpt "search"}}checked{{end}}>Search

<input type="radio" name="r" value="extended" {{if eq .RadioOpt "extended"}}checked{{end}}>Extended

<input type="radio" name="r" value="lattice" {{if eq .RadioOpt "lattice"}}checked{{end}}>Lattice

<input class="btn" type="submit" value="Tokenize"/>

  </form>
  {{if .CmdErr}}
    {{.CmdErr}}
  {{end}}
  {{if .GraphSvg}}
    {{.GraphSvg}}
  {{end}}
  {{if .Tokens}}

    <thread>

      Surface
      Part-of-Speech
      Base Form
      Reading
      Pronounciation

</thread>

    {{range .Tokens}}

      {{.Surface}}
      {{.Pos}}
      {{.Baseform}}
      {{.Reading}}
      {{.Pronounciation}}

    {{end}}

  {{end}}

`
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"os"
	"strconv"
	"strings"
	"time"

	"github.com/grafana/metrictank/cmd/mt-index-cat/out"
	"github.com/grafana/metrictank/idx/cassandra"
	"github.com/grafana/metrictank/logger"
	"github.com/raintank/dur"
	"github.com/raintank/schema"
	log "github.com/sirupsen/logrus"
)

func init() {
	formatter := &logger.TextFormatter{}
	formatter.TimestampFormat = "2006-01-02 15:04:05.000"
	log.SetFormatter(formatter)
	log.SetLevel(log.InfoLevel)
}

func perror(err error) {
	if err != nil {
		log.Fatal(err.Error())
	}
}

func main() {

	var addr string
	var prefix string
	var substr string
	var suffix string
	var tags string
	var from string
	var maxStale string
	var minStale string
	var verbose bool
	var limit int
	var partitionStr string

	globalFlags := flag.NewFlagSet("global config flags", flag.ExitOnError)
	globalFlags.StringVar(&addr, "addr", "http://localhost:6060", "graphite/metrictank address")
	globalFlags.StringVar(&prefix, "prefix", "", "only show metrics that have this prefix")
	globalFlags.StringVar(&substr, "substr", "", "only show metrics that have this substring")
	globalFlags.StringVar(&suffix, "suffix", "", "only show metrics that have this suffix")
	globalFlags.StringVar(&partitionStr, "partitions", "*", "only show metrics from the comma separated list of partitions or * for all")
	globalFlags.StringVar(&tags, "tags", "", "tag filter. empty (default), 'some', 'none', 'valid', or 'invalid'")
	globalFlags.StringVar(&from, "from", "30min", "for vegeta outputs, will generate requests for data starting from now minus... eg '30min', '5h', '14d', etc. or a unix timestamp")
	globalFlags.StringVar(&maxStale, "max-stale", "6h30min", "exclude series that have not been seen for this much time.  use 0 to disable")
	globalFlags.StringVar(&minStale, "min-stale", "0", "exclude series that have been seen in this much time.  use 0 to disable")
	globalFlags.IntVar(&limit, "limit", 0, "only show this many metrics.  use 0 to disable")
	globalFlags.BoolVar(&verbose, "verbose", false, "print stats to stderr")

	cassFlags := cassandra.ConfigSetup()

	outputs := []string{"dump", "list", "vegeta-render", "vegeta-render-patterns"}

	flag.Usage = func() {
		fmt.Println("mt-index-cat")
		fmt.Println()
		fmt.Println("Retrieves a metrictank index and dumps it in the requested format")
		fmt.Println("In particular, the vegeta outputs are handy to pipe requests for given series into the vegeta http benchmark tool")
		fmt.Println()
		fmt.Printf("Usage:\n\n")
		fmt.Printf("  mt-index-cat [global config flags] <idxtype> [idx config flags] output \n\n")
		fmt.Printf("global config flags:\n\n")
		globalFlags.PrintDefaults()
		fmt.Println()
		fmt.Println("tags filter:")
		fmt.Println("     ''        no filtering based on tags")
		fmt.Println("     'none'    only show metrics that have no tags")
		fmt.Println("     'some'    only show metrics that have one or more tags")
		fmt.Println("     'valid'   only show metrics whose tags (if any) are valid")
		fmt.Println("     'invalid' only show metrics that have one or more invalid tags")
		fmt.Println()
		fmt.Printf("idxtype: only 'cass' supported for now\n\n")
		fmt.Printf("cass config flags:\n\n")
		cassFlags.PrintDefaults()
		fmt.Println()
		fmt.Printf("output: either presets like %v\n", strings.Join(outputs, "|"))
		fmt.Printf("output: or custom templates like '{{.Id}} {{.OrgId}} {{.Name}} {{.Metric}} {{.Interval}} {{.Unit}} {{.Mtype}} {{.Tags}} {{.LastUpdate}} {{.Partition}}'\n\n\n")
		fmt.Println("You may also use processing functions in templates:")
		fmt.Println("pattern: transforms a graphite.style.metric.name into a pattern with wildcards inserted")
		fmt.Println("age: subtracts the passed integer (typically .LastUpdate) from the query time")
		fmt.Println("roundDuration: formats an integer-seconds duration using aggressive rounding. for the purpose of getting an idea of overal metrics age")
		fmt.Println("EXAMPLES:")
		fmt.Println("mt-index-cat -from 60min cass -hosts cassandra:9042 list")
		fmt.Println("mt-index-cat -from 60min cass -hosts cassandra:9042 'sumSeries({{.Name | pattern}})'")
		fmt.Println("mt-index-cat -from 60min cass -hosts cassandra:9042 'GET http://localhost:6060/render?target=sumSeries({{.Name | pattern}})&from=-6h\\nX-Org-Id: 1\\n\\n'")
		fmt.Println("mt-index-cat cass -hosts cassandra:9042 -timeout 60s '{{.LastUpdate | age | roundDuration}}\\n' | sort | uniq -c")
	}

	if len(os.Args) == 2 && (os.Args[1] == "-h" || os.Args[1] == "--help") {
		flag.Usage()
		os.Exit(0)
	}

	if len(os.Args) < 3 {
		flag.Usage()
		os.Exit(-1)
	}

	var partitions []int32
	if partitionStr != "*" {
		for _, p := range strings.Split(partitionStr, ",") {
			p = strings.TrimSpace(p)

			// handle trailing "," on the list of partitions.
			if p == "" {
				continue
			}

			id, err := strconv.ParseInt(p, 10, 32)
			if err != nil {
				log.Printf("invalid partition id %q. must be a int32", p)
				flag.Usage()
				os.Exit(-1)
			}
			partitions = append(partitions, int32(id))
		}
	}

	format := os.Args[len(os.Args)-1]
	var found bool
	if strings.Contains(format, "{{") {
		found = true
	} else {
		for _, output := range outputs {
			if format == output {
				found = true
				break
			}
		}
	}
	if !found {
		log.Printf("invalid output %q", format)
		flag.Usage()
		os.Exit(-1)
	}
	var cassI int
	for i, v := range os.Args {
		if v == "cass" {
			cassI = i
		}
	}
	if cassI == 0 {
		log.Println("only indextype 'cass' supported")
		flag.Usage()
		os.Exit(1)
	}

	if tags != "" && tags != "valid" && tags != "invalid" && tags != "some" && tags != "none" {
		log.Println("invalid tags filter")
		flag.Usage()
		os.Exit(1)
	}

	globalFlags.Parse(os.Args[1:cassI])
	cassFlags.Parse(os.Args[cassI+1 : len(os.Args)-1])
	cassandra.Enabled = true

	var show func(d schema.MetricDefinition)

	switch format {
	case "dump":
		show = out.Dump
	case "list":
		show = out.List
	case "vegeta-render":
		show = out.GetVegetaRender(addr, from)
	case "vegeta-render-patterns":
		show = out.GetVegetaRenderPattern(addr, from)
	default:
		show = out.Template(format)
	}

	idx := cassandra.New()
	err := idx.InitBare()
	perror(err)

	// from should either be a unix timestamp, or a specification that graphite/metrictank will recognize.
	_, err = strconv.Atoi(from)
	if err != nil {
		_, err = dur.ParseNDuration(from)
		perror(err)
	}

	var cutoff, cutoffMin int64
	now := time.Now().Unix()
	if maxStale != "0" {
		maxStaleInt, err := dur.ParseNDuration(maxStale)
		perror(err)
		cutoff = now - int64(maxStaleInt)
	}
	if minStale != "0" {
		minStaleInt, err := dur.ParseNDuration(minStale)
		perror(err)
		cutoffMin = now - int64(minStaleInt)
	}

	var defs []schema.MetricDefinition
	if len(partitions) == 0 {
		defs = idx.Load(nil, uint32(cutoff))
	} else {
		defs = idx.LoadPartitions(partitions, nil, uint32(cutoff))
	}
	// set this after doing the query, to assure age can't possibly be negative
	out.QueryTime = time.Now().Unix()
	total := len(defs)
	shown := 0

	for _, d := range defs {
		// note that prefix and substr can be "", meaning filter disabled.
		// the conditions handle this fine as well.
		if !strings.HasPrefix(d.Name, prefix) {
			continue
		}
		if !strings.HasSuffix(d.Name, suffix) {
			continue
		}
		if !strings.Contains(d.Name, substr) {
			continue
		}
		if tags == "none" && len(d.Tags) != 0 {
			continue
		}
		if tags == "some" && len(d.Tags) == 0 {
			continue
		}
		if tags == "valid" || tags == "invalid" {
			valid := schema.ValidateTags(d.Tags)

			// skip the metric if the validation result is not what we want
			if valid != (tags == "valid") {
				continue
			}
		}
		if cutoffMin != 0 && d.LastUpdate >= cutoffMin {
			continue
		}
		show(d)
		shown += 1
		if shown == limit {
			break
		}
	}

	if verbose {
		fmt.Fprintf(os.Stderr, "total: %d\n", total)
		fmt.Fprintf(os.Stderr, "shown: %d\n", shown)
	}
}
<commit_msg>process partitions after flag.Parse<commit_after>package main

import (
	"flag"
	"fmt"
	"os"
	"strconv"
	"strings"
	"time"

	"github.com/grafana/metrictank/cmd/mt-index-cat/out"
	"github.com/grafana/metrictank/idx/cassandra"
	"github.com/grafana/metrictank/logger"
	"github.com/raintank/dur"
	"github.com/raintank/schema"
	log "github.com/sirupsen/logrus"
)

func init() {
	formatter := &logger.TextFormatter{}
	formatter.TimestampFormat = "2006-01-02 15:04:05.000"
	log.SetFormatter(formatter)
	log.SetLevel(log.InfoLevel)
}

func perror(err error) {
	if err != nil {
		log.Fatal(err.Error())
	}
}

func main() {

	var addr string
	var prefix string
	var substr string
	var suffix string
	var tags string
	var from string
	var maxStale string
	var minStale string
	var verbose bool
	var limit int
	var partitionStr string

	globalFlags := flag.NewFlagSet("global config flags", flag.ExitOnError)
	globalFlags.StringVar(&addr, "addr", "http://localhost:6060", "graphite/metrictank address")
	globalFlags.StringVar(&prefix, "prefix", "", "only show metrics that have this prefix")
	globalFlags.StringVar(&substr, "substr", "", "only show metrics that have this substring")
	globalFlags.StringVar(&suffix, "suffix", "", "only show metrics that have this suffix")
	globalFlags.StringVar(&partitionStr, "partitions", "*", "only show metrics from the comma separated list of partitions or * for all")
	globalFlags.StringVar(&tags, "tags", "", "tag filter. empty (default), 'some', 'none', 'valid', or 'invalid'")
	globalFlags.StringVar(&from, "from", "30min", "for vegeta outputs, will generate requests for data starting from now minus... eg '30min', '5h', '14d', etc. or a unix timestamp")
	globalFlags.StringVar(&maxStale, "max-stale", "6h30min", "exclude series that have not been seen for this much time.  use 0 to disable")
	globalFlags.StringVar(&minStale, "min-stale", "0", "exclude series that have been seen in this much time.  use 0 to disable")
	globalFlags.IntVar(&limit, "limit", 0, "only show this many metrics.  use 0 to disable")
	globalFlags.BoolVar(&verbose, "verbose", false, "print stats to stderr")

	cassFlags := cassandra.ConfigSetup()

	outputs := []string{"dump", "list", "vegeta-render", "vegeta-render-patterns"}

	flag.Usage = func() {
		fmt.Println("mt-index-cat")
		fmt.Println()
		fmt.Println("Retrieves a metrictank index and dumps it in the requested format")
		fmt.Println("In particular, the vegeta outputs are handy to pipe requests for given series into the vegeta http benchmark tool")
		fmt.Println()
		fmt.Printf("Usage:\n\n")
		fmt.Printf("  mt-index-cat [global config flags] <idxtype> [idx config flags] output \n\n")
		fmt.Printf("global config flags:\n\n")
		globalFlags.PrintDefaults()
		fmt.Println()
		fmt.Println("tags filter:")
		fmt.Println("     ''        no filtering based on tags")
		fmt.Println("     'none'    only show metrics that have no tags")
		fmt.Println("     'some'    only show metrics that have one or more tags")
		fmt.Println("     'valid'   only show metrics whose tags (if any) are valid")
		fmt.Println("     'invalid' only show metrics that have one or more invalid tags")
		fmt.Println()
		fmt.Printf("idxtype: only 'cass' supported for now\n\n")
		fmt.Printf("cass config flags:\n\n")
		cassFlags.PrintDefaults()
		fmt.Println()
		fmt.Printf("output: either presets like %v\n", strings.Join(outputs, "|"))
		fmt.Printf("output: or custom templates like '{{.Id}} {{.OrgId}} {{.Name}} {{.Metric}} {{.Interval}} {{.Unit}} {{.Mtype}} {{.Tags}} {{.LastUpdate}} {{.Partition}}'\n\n\n")
		fmt.Println("You may also use processing functions in templates:")
		fmt.Println("pattern: transforms a graphite.style.metric.name into a pattern with wildcards inserted")
		fmt.Println("age: subtracts the passed integer (typically .LastUpdate) from the query time")
		fmt.Println("roundDuration: formats an integer-seconds duration using aggressive rounding. for the purpose of getting an idea of overal metrics age")
		fmt.Println("EXAMPLES:")
		fmt.Println("mt-index-cat -from 60min cass -hosts cassandra:9042 list")
		fmt.Println("mt-index-cat -from 60min cass -hosts cassandra:9042 'sumSeries({{.Name | pattern}})'")
		fmt.Println("mt-index-cat -from 60min cass -hosts cassandra:9042 'GET http://localhost:6060/render?target=sumSeries({{.Name | pattern}})&from=-6h\\nX-Org-Id: 1\\n\\n'")
		fmt.Println("mt-index-cat cass -hosts cassandra:9042 -timeout 60s '{{.LastUpdate | age | roundDuration}}\\n' | sort | uniq -c")
	}

	if len(os.Args) == 2 && (os.Args[1] == "-h" || os.Args[1] == "--help") {
		flag.Usage()
		os.Exit(0)
	}

	if len(os.Args) < 3 {
		flag.Usage()
		os.Exit(-1)
	}

	format := os.Args[len(os.Args)-1]
	var found bool
	if strings.Contains(format, "{{") {
		found = true
	} else {
		for _, output := range outputs {
			if format == output {
				found = true
				break
			}
		}
	}
	if !found {
		log.Printf("invalid output %q", format)
		flag.Usage()
		os.Exit(-1)
	}
	var cassI int
	for i, v := range os.Args {
		if v == "cass" {
			cassI = i
		}
	}
	if cassI == 0 {
		log.Println("only indextype 'cass' supported")
		flag.Usage()
		os.Exit(1)
	}

	if tags != "" && tags != "valid" && tags != "invalid" && tags != "some" && tags != "none" {
		log.Println("invalid tags filter")
		flag.Usage()
		os.Exit(1)
	}

	globalFlags.Parse(os.Args[1:cassI])
	cassFlags.Parse(os.Args[cassI+1 : len(os.Args)-1])
	cassandra.Enabled = true

	var show func(d schema.MetricDefinition)

	switch format {
	case "dump":
		show = out.Dump
	case "list":
		show = out.List
	case "vegeta-render":
		show = out.GetVegetaRender(addr, from)
	case "vegeta-render-patterns":
		show = out.GetVegetaRenderPattern(addr, from)
	default:
		show = out.Template(format)
	}

	idx := cassandra.New()
	err := idx.InitBare()
	perror(err)

	// from should either be a unix timestamp, or a specification that graphite/metrictank will recognize.
	_, err = strconv.Atoi(from)
	if err != nil {
		_, err = dur.ParseNDuration(from)
		perror(err)
	}

	var cutoff, cutoffMin int64
	now := time.Now().Unix()
	if maxStale != "0" {
		maxStaleInt, err := dur.ParseNDuration(maxStale)
		perror(err)
		cutoff = now - int64(maxStaleInt)
	}
	if minStale != "0" {
		minStaleInt, err := dur.ParseNDuration(minStale)
		perror(err)
		cutoffMin = now - int64(minStaleInt)
	}

	var partitions []int32
	log.Infof("parsing partitionStr %s", partitionStr)
	if partitionStr != "*" {
		for _, p := range strings.Split(partitionStr, ",") {
			p = strings.TrimSpace(p)

			// handle trailing "," on the list of partitions.
			if p == "" {
				continue
			}

			id, err := strconv.ParseInt(p, 10, 32)
			if err != nil {
				log.Printf("invalid partition id %q. must be a int32", p)
				flag.Usage()
				os.Exit(-1)
			}
			partitions = append(partitions, int32(id))
		}
	}

	var defs []schema.MetricDefinition
	if len(partitions) == 0 {
		log.Info("loading series from all partitions.")
		defs = idx.Load(nil, uint32(cutoff))
	} else {
		log.Infof("loading series from partitiosn %v", partitions)
		defs = idx.LoadPartitions(partitions, nil, uint32(cutoff))
	}
	// set this after doing the query, to assure age can't possibly be negative
	out.QueryTime = time.Now().Unix()
	total := len(defs)
	shown := 0

	for _, d := range defs {
		// note that prefix and substr can be "", meaning filter disabled.
		// the conditions handle this fine as well.
		if !strings.HasPrefix(d.Name, prefix) {
			continue
		}
		if !strings.HasSuffix(d.Name, suffix) {
			continue
		}
		if !strings.Contains(d.Name, substr) {
			continue
		}
		if tags == "none" && len(d.Tags) != 0 {
			continue
		}
		if tags == "some" && len(d.Tags) == 0 {
			continue
		}
		if tags == "valid" || tags == "invalid" {
			valid := schema.ValidateTags(d.Tags)

			// skip the metric if the validation result is not what we want
			if valid != (tags == "valid") {
				continue
			}
		}
		if cutoffMin != 0 && d.LastUpdate >= cutoffMin {
			continue
		}
		show(d)
		shown += 1
		if shown == limit {
			break
		}
	}

	if verbose {
		fmt.Fprintf(os.Stderr, "total: %d\n", total)
		fmt.Fprintf(os.Stderr, "shown: %d\n", shown)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"path/filepath"
	"runtime/pprof"
	"strings"
	"time"

	"github.com/nitrogen-lang/nitrogen/src/ast"
	"github.com/nitrogen-lang/nitrogen/src/compiler"
	"github.com/nitrogen-lang/nitrogen/src/compiler/marshal"
	"github.com/nitrogen-lang/nitrogen/src/lexer"
	"github.com/nitrogen-lang/nitrogen/src/moduleutils"
	"github.com/nitrogen-lang/nitrogen/src/object"
	"github.com/nitrogen-lang/nitrogen/src/parser"
	"github.com/nitrogen-lang/nitrogen/src/vm"

	_ "github.com/nitrogen-lang/nitrogen/src/builtins"
)

const (
	interactivePrompt = ">> "
	version           = "0.1.0"
)

type strSliceFlag []string

func (s *strSliceFlag) String() string {
	return strings.Join(*s, ":")
}

func (s *strSliceFlag) Set(st string) error {
	*s = append(*s, st)
	return nil
}

var (
	interactive  bool
	printAst     bool
	startSCGI    bool
	printVersion bool
	fullDebug    bool
	cpuprofile   string
	memprofile   string
	outputFile   string

	modulePaths     strSliceFlag
	autoloadModules strSliceFlag
)

func init() {
	pwd, _ := os.Getwd()
	modulePaths = append(modulePaths, pwd)

	envModPath := os.Getenv("NITROGEN_MODULES")
	if envModPath != "" {
		modulePaths = append(modulePaths, envModPath)
	}

	flag.BoolVar(&interactive, "i", false, "Interactive mode")
	flag.BoolVar(&printAst, "ast", false, "Print AST and exit")
	flag.BoolVar(&startSCGI, "scgi", false, "Start as an SCGI server")
	flag.BoolVar(&printVersion, "version", false, "Print version information")
	flag.BoolVar(&fullDebug, "debug", false, "Enable debug mode")
	flag.StringVar(&cpuprofile, "cpuprofile", "", "File to write CPU profile data")
	flag.StringVar(&memprofile, "memprofile", "", "File to write memory profile data")
	flag.StringVar(&outputFile, "o", "", "Output file of compiled bytecode")

	flag.Var(&modulePaths, "M", "Module search paths")
	flag.Var(&autoloadModules, "al", "Autoload modules")
}

func main() {
	flag.Parse()

	if printVersion {
		fmt.Printf("Nitrogen Interpreter Version %s\n", version)
		return
	}

	if len(autoloadModules) > 0 {
		if err := loadModules(modulePaths, autoloadModules); err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
	}

	if startSCGI {
		startSCGIServer()
		return
	}

	moduleutils.ParserSettings.Debug = fullDebug
	if interactive {
		fmt.Println("Nitrogen Programming Language")
		fmt.Println("Type in commands at the prompt")
		startRepl(os.Stdin, os.Stdout)
		return
	}

	if flag.NArg() == 0 {
		fmt.Println("No script given")
		os.Exit(1)
	}

	if cpuprofile != "" {
		f, err := os.Create(cpuprofile)
		if err != nil {
			log.Fatal(err)
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	env := makeEnv()

	var code *compiler.CodeBlock
	var program *ast.Program
	var err error
	if filepath.Ext(flag.Arg(0)) == ".nib" {
		code, err = marshal.ReadFile(flag.Arg(0))
		if err != nil {
			fmt.Println(err.Error())
			return
		}
	} else {
		program, err = moduleutils.ASTCache.GetTree(flag.Arg(0))
		if err != nil {
			fmt.Println(err.Error())
			return
		}

		if printAst {
			fmt.Println(program.String())
			return
		}
	}

	var result object.Object
	var start time.Time

	if code == nil {
		code = compiler.Compile(program, "__main")
	}

	if outputFile != "" {
		marshal.WriteFile(outputFile, code)
		return
	}

	start = time.Now()
	result = runCompiledCode(code, env)

	if fullDebug {
		fmt.Printf("Execution took %s\n", time.Now().Sub(start))
	}

	if result != nil && result != object.NullConst {
		if e, ok := result.(*object.Exception); ok {
			os.Stdout.WriteString("Uncaught Exception: ")
			os.Stdout.WriteString(e.Message)
			os.Stdout.Write([]byte{'\n'})
			os.Exit(1)
		}
		os.Stdout.WriteString(result.Inspect())
		os.Stdout.Write([]byte{'\n'})
	}

	if memprofile != "" {
		f, err := os.Create(memprofile)
		if err != nil {
			log.Fatal("could not create memory profile: ", err)
		}
		if err := pprof.WriteHeapProfile(f); err != nil {
			log.Fatal("could not write memory profile: ", err)
		}
		f.Close()
	}
}

func runCompiledCode(code *compiler.CodeBlock, env *object.Environment) object.Object {
	if fullDebug {
		code.Print("")
	}

	env.CreateConst("_FILE", object.MakeStringObj(code.Filename))

	vmsettings := vm.NewSettings()
	vmsettings.Debug = fullDebug
	machine := vm.NewVM(vmsettings)

	return machine.Execute(code, env)
}

func makeEnv() *object.Environment {
	env := object.NewEnvironment()
	env.CreateConst("_ENV", getExternalEnv())
	env.CreateConst("_ARGV", getScriptArgs(flag.Arg(0)))
	env.Create("_SEARCH_PATHS", object.MakeStringArray(modulePaths))
	return env
}

func getExternalEnv() *object.Hash {
	return stringMapToHash(getExtEnvMap())
}

func getExtEnvMap() map[string]string {
	env := os.Environ()
	m := make(map[string]string, len(env))
	for _, v := range env {
		val := strings.SplitN(v, "=", 2)
		m[val[0]] = val[1]
	}
	return m
}

func stringMapToHash(env map[string]string) *object.Hash {
	m := &object.Hash{Pairs: make(map[object.HashKey]object.HashPair)}
	for k, v := range env {
		key := &object.String{Value: k}
		m.Pairs[key.HashKey()] = object.HashPair{
			Key:   key,
			Value: &object.String{Value: v},
		}
	}
	return m
}

func getScriptArgs(filepath string) *object.Array {
	s := flag.Args()[1:]
	length := len(s) + 1
	newElements := make([]object.Object, length, length)
	newElements[0] = &object.String{Value: filepath}
	for i, v := range s {
		newElements[i+1] = &object.String{Value: v}
	}
	return &object.Array{Elements: newElements}
}

func startRepl(in io.Reader, out io.Writer) {
	scanner := bufio.NewScanner(in)
	env := makeEnv()

	var code *compiler.CodeBlock
	for {
		fmt.Fprint(out, interactivePrompt)
		scanned := scanner.Scan()
		if !scanned {
			return
		}

		line := strings.TrimSpace(scanner.Text())
		if line == "" {
			continue
		}
		if line == ".quit" || line == ".exit" {
			return
		}

		l := lexer.NewString(line)
		p := parser.New(l, moduleutils.ParserSettings)

		program := p.ParseProgram()
		if len(p.Errors()) != 0 {
			printParserErrors(out, p.Errors())
			continue
		}

		code = compiler.Compile(program, "__main")

		result := runCompiledCode(code, env)
		if result != nil && result != object.NullConst {
			io.WriteString(out, result.Inspect())
		}
		io.WriteString(out, "\n")
	}
}

func printParserErrors(out io.Writer, errors []string) {
	for _, msg := range errors {
		fmt.Fprintf(out, "ERROR: %s\n", msg)
	}
}
<commit_msg>Remove exec bit from go file<commit_after>package main

import (
	"bufio"
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"path/filepath"
	"runtime/pprof"
	"strings"
	"time"

	"github.com/nitrogen-lang/nitrogen/src/ast"
	"github.com/nitrogen-lang/nitrogen/src/compiler"
	"github.com/nitrogen-lang/nitrogen/src/compiler/marshal"
	"github.com/nitrogen-lang/nitrogen/src/lexer"
	"github.com/nitrogen-lang/nitrogen/src/moduleutils"
	"github.com/nitrogen-lang/nitrogen/src/object"
	"github.com/nitrogen-lang/nitrogen/src/parser"
	"github.com/nitrogen-lang/nitrogen/src/vm"

	_ "github.com/nitrogen-lang/nitrogen/src/builtins"
)

const (
	interactivePrompt = ">> "
	version           = "0.1.0"
)

type strSliceFlag []string

func (s *strSliceFlag) String() string {
	return strings.Join(*s, ":")
}

func (s *strSliceFlag) Set(st string) error {
	*s = append(*s, st)
	return nil
}

var (
	interactive  bool
	printAst     bool
	startSCGI    bool
	printVersion bool
	fullDebug    bool
	cpuprofile   string
	memprofile   string
	outputFile   string

	modulePaths     strSliceFlag
	autoloadModules strSliceFlag
)

func init() {
	pwd, _ := os.Getwd()
	modulePaths = append(modulePaths, pwd)

	envModPath := os.Getenv("NITROGEN_MODULES")
	if envModPath != "" {
		modulePaths = append(modulePaths, envModPath)
	}

	flag.BoolVar(&interactive, "i", false, "Interactive mode")
	flag.BoolVar(&printAst, "ast", false, "Print AST and exit")
	flag.BoolVar(&startSCGI, "scgi", false, "Start as an SCGI server")
	flag.BoolVar(&printVersion, "version", false, "Print version information")
	flag.BoolVar(&fullDebug, "debug", false, "Enable debug mode")
	flag.StringVar(&cpuprofile, "cpuprofile", "", "File to write CPU profile data")
	flag.StringVar(&memprofile, "memprofile", "", "File to write memory profile data")
	flag.StringVar(&outputFile, "o", "", "Output file of compiled bytecode")

	flag.Var(&modulePaths, "M", "Module search paths")
	flag.Var(&autoloadModules, "al", "Autoload modules")
}

func main() {
	flag.Parse()

	if printVersion {
		fmt.Printf("Nitrogen Interpreter Version %s\n", version)
		return
	}

	if len(autoloadModules) > 0 {
		if err := loadModules(modulePaths, autoloadModules); err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
	}

	if startSCGI {
		startSCGIServer()
		return
	}

	moduleutils.ParserSettings.Debug = fullDebug
	if interactive {
		fmt.Println("Nitrogen Programming Language")
		fmt.Println("Type in commands at the prompt")
		startRepl(os.Stdin, os.Stdout)
		return
	}

	if flag.NArg() == 0 {
		fmt.Println("No script given")
		os.Exit(1)
	}

	if cpuprofile != "" {
		f, err := os.Create(cpuprofile)
		if err != nil {
			log.Fatal(err)
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	env := makeEnv()

	var code *compiler.CodeBlock
	var program *ast.Program
	var err error
	if filepath.Ext(flag.Arg(0)) == ".nib" {
		code, err = marshal.ReadFile(flag.Arg(0))
		if err != nil {
			fmt.Println(err.Error())
			return
		}
	} else {
		program, err = moduleutils.ASTCache.GetTree(flag.Arg(0))
		if err != nil {
			fmt.Println(err.Error())
			return
		}

		if printAst {
			fmt.Println(program.String())
			return
		}
	}

	var result object.Object
	var start time.Time

	if code == nil {
		code = compiler.Compile(program, "__main")
	}

	if outputFile != "" {
		marshal.WriteFile(outputFile, code)
		return
	}

	start = time.Now()
	result = runCompiledCode(code, env)

	if fullDebug {
		fmt.Printf("Execution took %s\n", time.Now().Sub(start))
	}

	if result != nil && result != object.NullConst {
		if e, ok := result.(*object.Exception); ok {
			os.Stdout.WriteString("Uncaught Exception: ")
			os.Stdout.WriteString(e.Message)
			os.Stdout.Write([]byte{'\n'})
			os.Exit(1)
		}
		os.Stdout.WriteString(result.Inspect())
		os.Stdout.Write([]byte{'\n'})
	}

	if memprofile != "" {
		f, err := os.Create(memprofile)
		if err != nil {
			log.Fatal("could not create memory profile: ", err)
		}
		if err := pprof.WriteHeapProfile(f); err != nil {
			log.Fatal("could not write memory profile: ", err)
		}
		f.Close()
	}
}

func runCompiledCode(code *compiler.CodeBlock, env *object.Environment) object.Object {
	if fullDebug {
		code.Print("")
	}

	env.CreateConst("_FILE", object.MakeStringObj(code.Filename))

	vmsettings := vm.NewSettings()
	vmsettings.Debug = fullDebug
	machine := vm.NewVM(vmsettings)

	return machine.Execute(code, env)
}

func makeEnv() *object.Environment {
	env := object.NewEnvironment()
	env.CreateConst("_ENV", getExternalEnv())
	env.CreateConst("_ARGV", getScriptArgs(flag.Arg(0)))
	env.Create("_SEARCH_PATHS", object.MakeStringArray(modulePaths))
	return env
}

func getExternalEnv() *object.Hash {
	return stringMapToHash(getExtEnvMap())
}

func getExtEnvMap() map[string]string {
	env := os.Environ()
	m := make(map[string]string, len(env))
	for _, v := range env {
		val := strings.SplitN(v, "=", 2)
		m[val[0]] = val[1]
	}
	return m
}

func stringMapToHash(env map[string]string) *object.Hash {
	m := &object.Hash{Pairs: make(map[object.HashKey]object.HashPair)}
	for k, v := range env {
		key := &object.String{Value: k}
		m.Pairs[key.HashKey()] = object.HashPair{
			Key:   key,
			Value: &object.String{Value: v},
		}
	}
	return m
}

func getScriptArgs(filepath string) *object.Array {
	s := flag.Args()[1:]
	length := len(s) + 1
	newElements := make([]object.Object, length, length)
	newElements[0] = &object.String{Value: filepath}
	for i, v := range s {
		newElements[i+1] = &object.String{Value: v}
	}
	return &object.Array{Elements: newElements}
}

func startRepl(in io.Reader, out io.Writer) {
	scanner := bufio.NewScanner(in)
	env := makeEnv()

	var code *compiler.CodeBlock
	for {
		fmt.Fprint(out, interactivePrompt)
		scanned := scanner.Scan()
		if !scanned {
			return
		}

		line := strings.TrimSpace(scanner.Text())
		if line == "" {
			continue
		}
		if line == ".quit" || line == ".exit" {
			return
		}

		l := lexer.NewString(line)
		p := parser.New(l, moduleutils.ParserSettings)

		program := p.ParseProgram()
		if len(p.Errors()) != 0 {
			printParserErrors(out, p.Errors())
			continue
		}

		code = compiler.Compile(program, "__main")

		result := runCompiledCode(code, env)
		if result != nil && result != object.NullConst {
			io.WriteString(out, result.Inspect())
		}
		io.WriteString(out, "\n")
	}
}

func printParserErrors(out io.Writer, errors []string) {
	for _, msg := range errors {
		fmt.Fprintf(out, "ERROR: %s\n", msg)
	}
}
<|endoftext|>


<commit_before>package log

import (
	"io"
	"log"
	"os"

	"github.com/v2ray/v2ray-core/common/platform"
)

func createLogger(writer io.Writer) *log.Logger {
	return log.New(writer, "", log.Ldate|log.Ltime)
}

type logWriter interface {
	Log(LogEntry)
}

type noOpLogWriter struct {
}

func (this *noOpLogWriter) Log(entry LogEntry) {
	entry.Release()
}

type stdOutLogWriter struct {
	logger *log.Logger
}

func newStdOutLogWriter() logWriter {
	return &stdOutLogWriter{
		logger: createLogger(os.Stdout),
	}
}

func (this *stdOutLogWriter) Log(log LogEntry) {
	this.logger.Print(log.String() + platform.LineSeparator())
	log.Release()
}

type fileLogWriter struct {
	queue  chan LogEntry
	logger *log.Logger
	file   *os.File
}

func (this *fileLogWriter) Log(log LogEntry) {
	select {
	case this.queue <- log:
	default:
		log.Release()
		// We don't expect this to happen, but don't want to block main thread as well.
	}
}

func (this *fileLogWriter) run() {
	for entry := range this.queue {
		this.logger.Print(entry.String() + platform.LineSeparator())
		entry.Release()
	}
}

func (this *fileLogWriter) close() {
	this.file.Close()
}

func newFileLogWriter(path string) (*fileLogWriter, error) {
	file, err := os.OpenFile(path, os.O_APPEND|os.O_WRONLY|os.O_CREATE, 0600)
	if err != nil {
		return nil, err
	}
	logger := &fileLogWriter{
		queue:  make(chan LogEntry, 16),
		logger: log.New(file, "", log.Ldate|log.Ltime),
		file:   file,
	}
	go logger.run()
	return logger, nil
}
<commit_msg>safer for loop<commit_after>package log

import (
	"io"
	"log"
	"os"

	"github.com/v2ray/v2ray-core/common/platform"
)

func createLogger(writer io.Writer) *log.Logger {
	return log.New(writer, "", log.Ldate|log.Ltime)
}

type logWriter interface {
	Log(LogEntry)
}

type noOpLogWriter struct {
}

func (this *noOpLogWriter) Log(entry LogEntry) {
	entry.Release()
}

type stdOutLogWriter struct {
	logger *log.Logger
}

func newStdOutLogWriter() logWriter {
	return &stdOutLogWriter{
		logger: createLogger(os.Stdout),
	}
}

func (this *stdOutLogWriter) Log(log LogEntry) {
	this.logger.Print(log.String() + platform.LineSeparator())
	log.Release()
}

type fileLogWriter struct {
	queue  chan LogEntry
	logger *log.Logger
	file   *os.File
}

func (this *fileLogWriter) Log(log LogEntry) {
	select {
	case this.queue <- log:
	default:
		log.Release()
		// We don't expect this to happen, but don't want to block main thread as well.
	}
}

func (this *fileLogWriter) run() {
	for {
		entry := <-this.queue
		this.logger.Print(entry.String() + platform.LineSeparator())
		entry.Release()
		entry = nil
	}
}

func (this *fileLogWriter) close() {
	this.file.Close()
}

func newFileLogWriter(path string) (*fileLogWriter, error) {
	file, err := os.OpenFile(path, os.O_APPEND|os.O_WRONLY|os.O_CREATE, 0600)
	if err != nil {
		return nil, err
	}
	logger := &fileLogWriter{
		queue:  make(chan LogEntry, 16),
		logger: log.New(file, "", log.Ldate|log.Ltime),
		file:   file,
	}
	go logger.run()
	return logger, nil
}
<|endoftext|>


<commit_before>/*
Copyright 2020 Gravitational, 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,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package app

import (
	"net/http"
	"net/url"
	"path"

	"github.com/gravitational/teleport/lib/utils"
	"github.com/gravitational/trace"

	"github.com/julienschmidt/httprouter"
)

// withRouterAuth authenticates requests then hands the request to a
// httprouter.Handler handler.
func (h *Handler) withRouterAuth(handler routerAuthFunc) httprouter.Handle {
	return makeRouterHandler(func(w http.ResponseWriter, r *http.Request, p httprouter.Params) error {
		session, err := h.authenticate(r.Context(), r)
		if err != nil {
			return trace.Wrap(err)
		}
		if err := handler(w, r, p, session); err != nil {
			return trace.Wrap(err)
		}
		return nil
	})
}

// withAuth authenticates requests then hands the request to a http.HandlerFunc
// handler.
func (h *Handler) withAuth(handler handlerAuthFunc) http.HandlerFunc {
	return makeHandler(func(w http.ResponseWriter, r *http.Request) error {
		// If the caller fails to authenticate, redirect the caller to Teleport.
		session, err := h.authenticate(r.Context(), r)
		if err != nil {
			if redirectErr := h.redirectToLauncher(w, r, launcherURLParams{}); redirectErr == nil {
				return nil
			}
			return trace.Wrap(err)
		}
		if err := handler(w, r, session); err != nil {
			return trace.Wrap(err)
		}
		return nil
	})
}

// redirectToLauncher redirects to the proxy web's app launcher if the public
// address of the proxy is set.
func (h *Handler) redirectToLauncher(w http.ResponseWriter, r *http.Request, p launcherURLParams) error {
	if h.c.WebPublicAddr == "" {
		return trace.BadParameter("public address of the proxy is not set")
	}
	addr, err := utils.ParseAddr(r.Host)
	if err != nil {
		return trace.Wrap(err)
	}

	urlPath := []string{"web", "launch", addr.Host()}
	if p.clusterName != "" && p.publicAddr != "" {
		urlPath = append(urlPath, p.clusterName, p.publicAddr)
	}

	urlQuery := url.Values{}
	if p.stateToken != "" {
		urlQuery.Add("state", p.stateToken)
	}
	if p.awsRole != "" {
		urlQuery.Add("awsrole", p.awsRole)
	}

	u := url.URL{
		Scheme:   "https",
		Host:     h.c.WebPublicAddr,
		Path:     path.Join(urlPath...),
		RawQuery: urlQuery.Encode(),
	}
	http.Redirect(w, r, u.String(), http.StatusFound)
	return nil
}

type launcherURLParams struct {
	clusterName string
	publicAddr  string
	stateToken  string
	awsRole     string
}

// makeRouterHandler creates a httprouter.Handle.
func makeRouterHandler(handler routerFunc) httprouter.Handle {
	return func(w http.ResponseWriter, r *http.Request, p httprouter.Params) {
		if err := handler(w, r, p); err != nil {
			writeError(w, err)
			return
		}
	}
}

// makeHandler creates a http.HandlerFunc.
func makeHandler(handler handlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		if err := handler(w, r); err != nil {
			writeError(w, err)
			return
		}
	}
}

// writeError gets the HTTP status code from trace and writes the error to the
// response writer.
func writeError(w http.ResponseWriter, err error) {
	code := trace.ErrorToCode(err)
	http.Error(w, http.StatusText(code), code)
}

type routerFunc func(http.ResponseWriter, *http.Request, httprouter.Params) error
type routerAuthFunc func(http.ResponseWriter, *http.Request, httprouter.Params, *session) error

type handlerAuthFunc func(http.ResponseWriter, *http.Request, *session) error
type handlerFunc func(http.ResponseWriter, *http.Request) error
<commit_msg>Log when App Service fails due to empty `proxy_service.public_addr` (#10056)<commit_after>/*
Copyright 2020 Gravitational, 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,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package app

import (
	"net/http"
	"net/url"
	"path"

	"github.com/gravitational/teleport/lib/utils"
	"github.com/gravitational/trace"

	"github.com/julienschmidt/httprouter"
)

// withRouterAuth authenticates requests then hands the request to a
// httprouter.Handler handler.
func (h *Handler) withRouterAuth(handler routerAuthFunc) httprouter.Handle {
	return makeRouterHandler(func(w http.ResponseWriter, r *http.Request, p httprouter.Params) error {
		session, err := h.authenticate(r.Context(), r)
		if err != nil {
			return trace.Wrap(err)
		}
		if err := handler(w, r, p, session); err != nil {
			return trace.Wrap(err)
		}
		return nil
	})
}

// withAuth authenticates requests then hands the request to a http.HandlerFunc
// handler.
func (h *Handler) withAuth(handler handlerAuthFunc) http.HandlerFunc {
	return makeHandler(func(w http.ResponseWriter, r *http.Request) error {
		// If the caller fails to authenticate, redirect the caller to Teleport.
		session, err := h.authenticate(r.Context(), r)
		if err != nil {
			if redirectErr := h.redirectToLauncher(w, r, launcherURLParams{}); redirectErr == nil {
				return nil
			}
			return trace.Wrap(err)
		}
		if err := handler(w, r, session); err != nil {
			return trace.Wrap(err)
		}
		return nil
	})
}

// redirectToLauncher redirects to the proxy web's app launcher if the public
// address of the proxy is set.
func (h *Handler) redirectToLauncher(w http.ResponseWriter, r *http.Request, p launcherURLParams) error {
	if h.c.WebPublicAddr == "" {
		// The error below tends to be swallowed by the Web UI, so log a warning for
		// admins as well.
		h.log.Error("" +
			"Application Service requires public_addr to be set in the Teleport Proxy Service configuration. " +
			"Please contact your Teleport cluster administrator or refer to " +
			"https://goteleport.com/docs/application-access/guides/connecting-apps/#start-authproxy-service.")
		return trace.BadParameter("public address of the proxy is not set")
	}
	addr, err := utils.ParseAddr(r.Host)
	if err != nil {
		return trace.Wrap(err)
	}

	urlPath := []string{"web", "launch", addr.Host()}
	if p.clusterName != "" && p.publicAddr != "" {
		urlPath = append(urlPath, p.clusterName, p.publicAddr)
	}

	urlQuery := url.Values{}
	if p.stateToken != "" {
		urlQuery.Add("state", p.stateToken)
	}
	if p.awsRole != "" {
		urlQuery.Add("awsrole", p.awsRole)
	}

	u := url.URL{
		Scheme:   "https",
		Host:     h.c.WebPublicAddr,
		Path:     path.Join(urlPath...),
		RawQuery: urlQuery.Encode(),
	}
	http.Redirect(w, r, u.String(), http.StatusFound)
	return nil
}

type launcherURLParams struct {
	clusterName string
	publicAddr  string
	stateToken  string
	awsRole     string
}

// makeRouterHandler creates a httprouter.Handle.
func makeRouterHandler(handler routerFunc) httprouter.Handle {
	return func(w http.ResponseWriter, r *http.Request, p httprouter.Params) {
		if err := handler(w, r, p); err != nil {
			writeError(w, err)
			return
		}
	}
}

// makeHandler creates a http.HandlerFunc.
func makeHandler(handler handlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		if err := handler(w, r); err != nil {
			writeError(w, err)
			return
		}
	}
}

// writeError gets the HTTP status code from trace and writes the error to the
// response writer.
func writeError(w http.ResponseWriter, err error) {
	code := trace.ErrorToCode(err)
	http.Error(w, http.StatusText(code), code)
}

type routerFunc func(http.ResponseWriter, *http.Request, httprouter.Params) error
type routerAuthFunc func(http.ResponseWriter, *http.Request, httprouter.Params, *session) error

type handlerAuthFunc func(http.ResponseWriter, *http.Request, *session) error
type handlerFunc func(http.ResponseWriter, *http.Request) error
<|endoftext|>


<commit_before>/*                          _       _
 *__      _____  __ ___   ___  __ _| |_ ___
 *\ \ /\ / / _ \/ _` \ \ / / |/ _` | __/ _ \
 * \ V  V /  __/ (_| |\ V /| | (_| | ||  __/
 *  \_/\_/ \___|\__,_| \_/ |_|\__,_|\__\___|
 *
 * Copyright © 2016 - 2018 Weaviate. All rights reserved.
 * LICENSE: https://github.com/creativesoftwarefdn/weaviate/blob/develop/LICENSE.md
 * AUTHOR: Bob van Luijt (bob@kub.design)
 * See www.creativesoftwarefdn.org for details
 * Contact: @CreativeSofwFdn / bob@kub.design
 */

package config

import (
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"

	"github.com/creativesoftwarefdn/weaviate/models"

	"github.com/go-openapi/swag"

	"github.com/creativesoftwarefdn/weaviate/messages"
)

// DefaultConfigFile is the default file when no config file is provided
const DefaultConfigFile string = "./weaviate.conf.json"

// DefaultEnvironment is the default env when no env is provided
const DefaultEnvironment string = "development"

// Flags are input options
type Flags struct {
	ConfigSection string `long:"config" description:"the section inside the config file that has to be used"`
	ConfigFile    string `long:"config-file" description:"path to config file (default: ./weaviate.conf.json)"`
}

// File gives the outline of the config file
type File struct {
	Environments []Environment `json:"environments"`
}

// Environment outline of the environment inside the config file
type Environment struct {
	Name        string      `json:"name"`
	Database    Database    `json:"database"`
	Schemas     Schemas     `json:"schemas"`
	Broker      Broker      `json:"broker"`
	P2P         P2P         `json:"P2P"`
	Cache       Cache       `json:"cache"`
	Limit       int64       `json:"limit"`
	Debug       bool        `json:"debug"`
	Development Development `json:"development"`
}

// P2P sets the Peer to Peer values of the Weaviate network
type P2P struct {
	FriendlyPeer string `json:"friendlyPeer"`
	Vectors      struct {
		Cbor string `json:"cbor"`
		MD5  string `json:"MD5"`
	}
}

// Broker checks if broker details are set
type Broker struct {
	Host string `json:"host"`
	Port int32  `json:"port"`
}

// Database is the outline of the database
type Database struct {
	Name           string      `json:"name"`
	DatabaseConfig interface{} `json:"database_config"`
}

// Schemas contains the schema for 'things' and for 'actions'
type Schemas struct {
	Thing  string `json:"thing"`
	Action string `json:"action"`
}

// Cache is the outline of the cache-system
type Cache struct {
	Name string `json:"name"`
}

// Development is the outline of (temporary) config variables
// Note: the purpose is that these variables will be moved somewhere else in time
type Development struct {
	ExternalInstances []Instance `json:"external_instances"`
}

// Instance is the outline for an external instance whereto crossreferences can be resolved
type Instance struct {
	URL      string `json:"url"`
	APIKey   string `json:"api_key"`
	APIToken string `json:"api_token"`
}

// GetConfigOptionGroup creates a option group for swagger
func GetConfigOptionGroup() *swag.CommandLineOptionsGroup {
	commandLineOptionsGroup := swag.CommandLineOptionsGroup{
		ShortDescription: "Connector config & MQTT config",
		LongDescription:  "",
		Options:          &Flags{},
	}

	return &commandLineOptionsGroup
}

// WeaviateConfig represents the used schema's
type WeaviateConfig struct {
	Environment Environment
	Hostname    string
	Scheme      string
}

// GetHostAddress from config locations
func (f *WeaviateConfig) GetHostAddress() string {
	return fmt.Sprintf("%s://%s", f.Scheme, f.Hostname)
}

// LoadConfig from config locations
func (f *WeaviateConfig) LoadConfig(flags *swag.CommandLineOptionsGroup, m *messages.Messaging) error {
	// Get command line flags
	configEnvironment := flags.Options.(*Flags).ConfigSection
	configFileName := flags.Options.(*Flags).ConfigFile

	// Set default if not given
	if configFileName == "" {
		configFileName = DefaultConfigFile
		m.InfoMessage("Using default config file location '" + DefaultConfigFile + "'.")
	}

	// Read config file
	file, err := ioutil.ReadFile(configFileName)
	if err != nil {
		return errors.New("config file '" + configFileName + "' not found.")
	}

	// Set default env if not given
	if err != nil || configEnvironment == "" {
		configEnvironment = DefaultEnvironment
		m.InfoMessage("Using default environment '" + DefaultEnvironment + "'.")
	}

	// Read from the config file and add it to an object
	var configFile File
	err = json.Unmarshal(file, &configFile)

	// Return error if config file is incorrect
	if err != nil {
		return errors.New("error unmarshalling the config file. ")
	}

	// Loop through all values in object to see whether the given connection-name exists
	foundEnvironment := false
	for _, env := range configFile.Environments {
		if env.Name == configEnvironment {
			foundEnvironment = true

			// Get config interface data
			f.Environment = env
		}
	}

	// Return default database because no good config is found
	if !foundEnvironment {
		return errors.New("no environment found with name '" + configEnvironment + "'")
	}

	m.InfoMessage("Config file found, loading environment..")

	// Check the debug mode
	m.Debug = f.Environment.Debug
	if f.Environment.Debug {
		m.InfoMessage("Running in DEBUG-mode")
	}

	return nil
}

// GetInstance from config
func (f *WeaviateConfig) GetInstance(hostname string, keyToken *models.KeyTokenGetResponse) (instance Instance, err error) {
	err = nil

	found := false

	// For each instance, check if hostname is the same
	for _, v := range f.Environment.Development.ExternalInstances {
		if hostname == v.URL {
			instance = v
			found = true
			break
		}
	}

	if !found {
		// Didn't find item in list
		err = errors.New("can't find key for given instance")
		return
	}

	// If no key and token are given in the config, use the key and token used for the current request
	if instance.APIKey == "" {
		instance.APIKey = string(keyToken.KeyID)
	}

	if instance.APIToken == "" {
		instance.APIToken = string(keyToken.Token)
	}

	return
}
<commit_msg>gh-380: add contextionary field to environment<commit_after>/*                          _       _
 *__      _____  __ ___   ___  __ _| |_ ___
 *\ \ /\ / / _ \/ _` \ \ / / |/ _` | __/ _ \
 * \ V  V /  __/ (_| |\ V /| | (_| | ||  __/
 *  \_/\_/ \___|\__,_| \_/ |_|\__,_|\__\___|
 *
 * Copyright © 2016 - 2018 Weaviate. All rights reserved.
 * LICENSE: https://github.com/creativesoftwarefdn/weaviate/blob/develop/LICENSE.md
 * AUTHOR: Bob van Luijt (bob@kub.design)
 * See www.creativesoftwarefdn.org for details
 * Contact: @CreativeSofwFdn / bob@kub.design
 */

package config

import (
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"

	"github.com/creativesoftwarefdn/weaviate/models"

	"github.com/go-openapi/swag"

	"github.com/creativesoftwarefdn/weaviate/messages"
)

// DefaultConfigFile is the default file when no config file is provided
const DefaultConfigFile string = "./weaviate.conf.json"

// DefaultEnvironment is the default env when no env is provided
const DefaultEnvironment string = "development"

// Flags are input options
type Flags struct {
	ConfigSection string `long:"config" description:"the section inside the config file that has to be used"`
	ConfigFile    string `long:"config-file" description:"path to config file (default: ./weaviate.conf.json)"`
}

// File gives the outline of the config file
type File struct {
	Environments []Environment `json:"environments"`
}

// Environment outline of the environment inside the config file
type Environment struct {
	Name          string        `json:"name"`
	Database      Database      `json:"database"`
	Schemas       Schemas       `json:"schemas"`
	Broker        Broker        `json:"broker"`
	P2P           P2P           `json:"P2P"`
	Cache         Cache         `json:"cache"`
	Limit         int64         `json:"limit"`
	Debug         bool          `json:"debug"`
	Development   Development   `json:"development"`
	Contextionary Contextionary `json:"contextionary"`
}

type Contextionary struct {
	KNNFile      string `json:'knn_file'`
	IDXFile      string `json:'idx_file'`
	failOnGerund bool   `json:'fail_ongerund'` // is true by default.
}

// P2P sets the Peer to Peer values of the Weaviate network
type P2P struct {
	FriendlyPeer string `json:"friendlyPeer"`
	Vectors      struct {
		Cbor string `json:"cbor"`
		MD5  string `json:"MD5"`
	}
}

// Broker checks if broker details are set
type Broker struct {
	Host string `json:"host"`
	Port int32  `json:"port"`
}

// Database is the outline of the database
type Database struct {
	Name           string      `json:"name"`
	DatabaseConfig interface{} `json:"database_config"`
}

// Schemas contains the schema for 'things' and for 'actions'
type Schemas struct {
	Thing  string `json:"thing"`
	Action string `json:"action"`
}

// Cache is the outline of the cache-system
type Cache struct {
	Name string `json:"name"`
}

// Development is the outline of (temporary) config variables
// Note: the purpose is that these variables will be moved somewhere else in time
type Development struct {
	ExternalInstances []Instance `json:"external_instances"`
}

// Instance is the outline for an external instance whereto crossreferences can be resolved
type Instance struct {
	URL      string `json:"url"`
	APIKey   string `json:"api_key"`
	APIToken string `json:"api_token"`
}

// GetConfigOptionGroup creates a option group for swagger
func GetConfigOptionGroup() *swag.CommandLineOptionsGroup {
	commandLineOptionsGroup := swag.CommandLineOptionsGroup{
		ShortDescription: "Connector config & MQTT config",
		LongDescription:  "",
		Options:          &Flags{},
	}

	return &commandLineOptionsGroup
}

// WeaviateConfig represents the used schema's
type WeaviateConfig struct {
	Environment Environment
	Hostname    string
	Scheme      string
}

// GetHostAddress from config locations
func (f *WeaviateConfig) GetHostAddress() string {
	return fmt.Sprintf("%s://%s", f.Scheme, f.Hostname)
}

// LoadConfig from config locations
func (f *WeaviateConfig) LoadConfig(flags *swag.CommandLineOptionsGroup, m *messages.Messaging) error {
	// Get command line flags
	configEnvironment := flags.Options.(*Flags).ConfigSection
	configFileName := flags.Options.(*Flags).ConfigFile

	// Set default if not given
	if configFileName == "" {
		configFileName = DefaultConfigFile
		m.InfoMessage("Using default config file location '" + DefaultConfigFile + "'.")
	}

	// Read config file
	file, err := ioutil.ReadFile(configFileName)
	if err != nil {
		return errors.New("config file '" + configFileName + "' not found.")
	}

	// Set default env if not given
	if err != nil || configEnvironment == "" {
		configEnvironment = DefaultEnvironment
		m.InfoMessage("Using default environment '" + DefaultEnvironment + "'.")
	}

	// Read from the config file and add it to an object
	var configFile File
	err = json.Unmarshal(file, &configFile)

	// Return error if config file is incorrect
	if err != nil {
		return errors.New("error unmarshalling the config file. ")
	}

	// Loop through all values in object to see whether the given connection-name exists
	foundEnvironment := false
	for _, env := range configFile.Environments {
		if env.Name == configEnvironment {
			foundEnvironment = true

			// Get config interface data
			f.Environment = env
		}
	}

	// Return default database because no good config is found
	if !foundEnvironment {
		return errors.New("no environment found with name '" + configEnvironment + "'")
	}

	m.InfoMessage("Config file found, loading environment..")

	// Check the debug mode
	m.Debug = f.Environment.Debug
	if f.Environment.Debug {
		m.InfoMessage("Running in DEBUG-mode")
	}

	return nil
}

// GetInstance from config
func (f *WeaviateConfig) GetInstance(hostname string, keyToken *models.KeyTokenGetResponse) (instance Instance, err error) {
	err = nil

	found := false

	// For each instance, check if hostname is the same
	for _, v := range f.Environment.Development.ExternalInstances {
		if hostname == v.URL {
			instance = v
			found = true
			break
		}
	}

	if !found {
		// Didn't find item in list
		err = errors.New("can't find key for given instance")
		return
	}

	// If no key and token are given in the config, use the key and token used for the current request
	if instance.APIKey == "" {
		instance.APIKey = string(keyToken.KeyID)
	}

	if instance.APIToken == "" {
		instance.APIToken = string(keyToken.Token)
	}

	return
}
<|endoftext|>


<commit_before>package library

import (
	"github.com/mattbaird/elastigo/api"
	"github.com/mattbaird/elastigo/core"
	"github.com/nytlabs/streamtools/st/blocks" // blocks
	"github.com/nytlabs/streamtools/st/util"
)

// specify those channels we're going to use to communicate with streamtools
type ToElasticsearch struct {
	blocks.Block
	queryrule chan chan interface{}
	inrule    chan interface{}
	in        chan interface{}
	out       chan interface{}
	quit      chan interface{}
	host      string
	port      string
	index     string
	indextype string
}

// we need to build a simple factory so that streamtools can make new blocks of this kind
func NewToElasticsearch() blocks.BlockInterface {
	return &ToElasticsearch{}
}

// Setup is called once before running the block. We build up the channels and specify what kind of block this is.
func (b *ToElasticsearch) Setup() {
	b.Kind = "ToElasticsearch"
	b.in = b.InRoute("in")
	b.inrule = b.InRoute("rule")
	b.queryrule = b.QueryRoute("rule")
	b.quit = b.Quit()
	b.out = b.Broadcast()
}

// Run is the block's main loop. Here we listen on the different channels we set up.
// This block posts a message to a specified Elasticsearch index with the given type.
func (b *ToElasticsearch) Run() {
	for {
		select {
		case msgI := <-b.inrule:
			host, _ := util.ParseString(msgI, "Host")
			port, _ := util.ParseString(msgI, "Port")
			index, _ := util.ParseString(msgI, "Index")
			indextype, _ := util.ParseString(msgI, "IndexType")

			// Set the Elasticsearch Host/Port to Connect to
			api.Domain = host
			api.Port = port

			b.host = host
			b.port = port
			b.index = index
			b.indextype = indextype

		case respChan := <-b.queryrule:
			// deal with a query request
			respChan <- map[string]interface{}{
				"Host":      b.host,
				"Port":      b.port,
				"Index":     b.index,
				"IndexType": b.indextype,
			}
		case <-b.quit:
			// quit the block
			return
		case msg := <-b.in:
			_, err := core.Index(true, b.index, b.indextype, "", msg)
			if err != nil {
				b.Error(err)
			}
		}
	}
}
<commit_msg>Updated toElasticsearch for new core.Index syntax<commit_after>package library

import (
	"github.com/mattbaird/elastigo/api"
	"github.com/mattbaird/elastigo/core"
	"github.com/nytlabs/streamtools/st/blocks" // blocks
	"github.com/nytlabs/streamtools/st/util"
)

// specify those channels we're going to use to communicate with streamtools
type ToElasticsearch struct {
	blocks.Block
	queryrule chan chan interface{}
	inrule    chan interface{}
	in        chan interface{}
	out       chan interface{}
	quit      chan interface{}
	host      string
	port      string
	index     string
	indextype string
}

// we need to build a simple factory so that streamtools can make new blocks of this kind
func NewToElasticsearch() blocks.BlockInterface {
	return &ToElasticsearch{}
}

// Setup is called once before running the block. We build up the channels and specify what kind of block this is.
func (b *ToElasticsearch) Setup() {
	b.Kind = "ToElasticsearch"
	b.in = b.InRoute("in")
	b.inrule = b.InRoute("rule")
	b.queryrule = b.QueryRoute("rule")
	b.quit = b.Quit()
	b.out = b.Broadcast()
}

// Run is the block's main loop. Here we listen on the different channels we set up.
// This block posts a message to a specified Elasticsearch index with the given type.
func (b *ToElasticsearch) Run() {
	for {
		select {
		case msgI := <-b.inrule:
			host, _ := util.ParseString(msgI, "Host")
			port, _ := util.ParseString(msgI, "Port")
			index, _ := util.ParseString(msgI, "Index")
			indextype, _ := util.ParseString(msgI, "IndexType")

			// Set the Elasticsearch Host/Port to Connect to
			api.Domain = host
			api.Port = port

			b.host = host
			b.port = port
			b.index = index
			b.indextype = indextype

		case respChan := <-b.queryrule:
			// deal with a query request
			respChan <- map[string]interface{}{
				"Host":      b.host,
				"Port":      b.port,
				"Index":     b.index,
				"IndexType": b.indextype,
			}
		case <-b.quit:
			// quit the block
			return
		case msg := <-b.in:
			var args map[string]interface{}
			_, err := core.Index(b.index, b.indextype, "", args, msg)
			if err != nil {
				b.Error(err)
			}
		}
	}
}
<|endoftext|>


<commit_before>// Copyright 2015 The Cockroach Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
// implied. See the License for the specific language governing
// permissions and limitations under the License. See the AUTHORS file
// for names of contributors.
//
// Author: Bram Gruneir (bram+code@cockroachlabs.com)

package storage_test

import (
	"reflect"
	"testing"
	"time"

	"golang.org/x/net/context"

	"github.com/cockroachdb/cockroach/keys"
	"github.com/cockroachdb/cockroach/proto"
	"github.com/cockroachdb/cockroach/storage"
	"github.com/cockroachdb/cockroach/storage/engine"
	"github.com/cockroachdb/cockroach/util/hlc"
	"github.com/cockroachdb/cockroach/util/leaktest"
	gogoproto "github.com/gogo/protobuf/proto"
)

// compareStoreStatus ensures that the actual store status for the passed in
// store is updated correctly. It checks that the Desc.StoreID, Desc.Attrs,
// Desc.Node, Desc.Capacity.Capacity, NodeID, RangeCount, ReplicatedRangeCount
// are exactly correct and that the bytes and counts for Live, Key and Val are
// at least the expected value.
// The latest actual stats are returned.
func compareStoreStatus(t *testing.T, store *storage.Store, expectedStoreStatus *storage.StoreStatus, testNumber int) *storage.StoreStatus {
	storeStatusKey := keys.StoreStatusKey(int32(store.Ident.StoreID))
	gArgs, gReply := getArgs(storeStatusKey, 1, store.Ident.StoreID)
	if err := store.ExecuteCmd(context.Background(), proto.Call{Args: gArgs, Reply: gReply}); err != nil {
		t.Fatalf("%v: failure getting store status: %s", testNumber, err)
	}
	if gReply.Value == nil {
		t.Errorf("%v: could not find store status at: %s", testNumber, storeStatusKey)
	}
	storeStatus := &storage.StoreStatus{}
	if err := gogoproto.Unmarshal(gReply.Value.GetBytes(), storeStatus); err != nil {
		t.Fatalf("%v: could not unmarshal store status: %+v", testNumber, gReply)
	}

	// Values much match exactly.
	if expectedStoreStatus.Desc.StoreID != storeStatus.Desc.StoreID {
		t.Errorf("%v: actual Desc.StoreID does not match expected\nexpected: %+v\nactual: %v\n", testNumber, expectedStoreStatus, storeStatus)
	}
	if !reflect.DeepEqual(expectedStoreStatus.Desc.Attrs, storeStatus.Desc.Attrs) {
		t.Errorf("%v: actual Desc.Attrs does not match expected\nexpected: %+v\nactual: %v\n", testNumber, expectedStoreStatus, storeStatus)
	}
	if !reflect.DeepEqual(expectedStoreStatus.Desc.Node, storeStatus.Desc.Node) {
		t.Errorf("%v: actual Desc.Attrs does not match expected\nexpected: %+v\nactual: %v\n", testNumber, expectedStoreStatus, storeStatus)
	}
	if storeStatus.Desc.Capacity.Capacity != expectedStoreStatus.Desc.Capacity.Capacity {
		t.Errorf("%v: actual Desc.Capacity.Capacity does not match expected\nexpected: %+v\nactual: %v\n", testNumber, expectedStoreStatus, storeStatus)
	}
	if expectedStoreStatus.NodeID != storeStatus.NodeID {
		t.Errorf("%v: actual node ID does not match expected\nexpected: %+v\nactual: %v\n", testNumber, expectedStoreStatus, storeStatus)
	}
	if expectedStoreStatus.RangeCount != storeStatus.RangeCount {
		t.Errorf("%v: actual RangeCount does not match expected\nexpected: %+v\nactual: %v\n", testNumber, expectedStoreStatus, storeStatus)
	}
	if expectedStoreStatus.ReplicatedRangeCount != storeStatus.ReplicatedRangeCount {
		t.Errorf("%v: actual ReplicatedRangeCount does not match expected\nexpected: %+v\nactual: %v\n", testNumber, expectedStoreStatus, storeStatus)
	}

	// Values should be >= to expected values.
	if storeStatus.Stats.LiveBytes < expectedStoreStatus.Stats.LiveBytes {
		t.Errorf("%v: actual Live Bytes is not greater or equal to expected\nexpected: %+v\nactual: %v\n", testNumber, expectedStoreStatus, storeStatus)
	}
	if storeStatus.Stats.KeyBytes < expectedStoreStatus.Stats.KeyBytes {
		t.Errorf("%v: actual Key Bytes is not greater or equal to expected\nexpected: %+v\nactual: %v\n", testNumber, expectedStoreStatus, storeStatus)
	}
	if storeStatus.Stats.ValBytes < expectedStoreStatus.Stats.ValBytes {
		t.Errorf("%v: actual Val Bytes is not greater or equal to expected\nexpected: %+v\nactual: %v\n", testNumber, expectedStoreStatus, storeStatus)
	}
	if storeStatus.Stats.LiveCount < expectedStoreStatus.Stats.LiveCount {
		t.Errorf("%v: actual Live Count is not greater or equal to expected\nexpected: %+v\nactual: %v\n", testNumber, expectedStoreStatus, storeStatus)
	}
	if storeStatus.Stats.KeyCount < expectedStoreStatus.Stats.KeyCount {
		t.Errorf("%v: actual Key Count is not greater or equal to expected\nexpected: %+v\nactual: %v\n", testNumber, expectedStoreStatus, storeStatus)
	}
	if storeStatus.Stats.ValCount < expectedStoreStatus.Stats.ValCount {
		t.Errorf("%v: actual Val Count is not greater or equal to expected\nexpected: %+v\nactual: %v\n", testNumber, expectedStoreStatus, storeStatus)
	}
	return storeStatus
}

// TestStoreStatus checks the store status after each range scan to ensure that
// it is being updated correctly.
func TestStoreStatus(t *testing.T) {
	defer leaktest.AfterTest(t)
	ctx := &storage.TestStoreContext
	ctx.ScanInterval = time.Duration(10 * time.Millisecond)
	store, stopper := createTestStoreWithEngine(t, engine.NewInMem(proto.Attributes{}, 10<<20), hlc.NewClock(hlc.NewManualClock(0).UnixNano), true, ctx)
	defer stopper.Stop()
	splitKey := proto.Key("b")
	content := proto.Key("test content")

	storeDesc, err := store.Descriptor()
	if err != nil {
		t.Fatal(err)
	}

	store.WaitForInit()
	rng := store.LookupRange([]byte("a"), nil)
	// Perform a read from the range to ensure that the raft election has
	// completed.  We do not expect a value to be present.
	gArgs, gReply := getArgs([]byte("a"), rng.Desc().RaftID, store.StoreID())
	if err := store.ExecuteCmd(context.Background(), proto.Call{Args: gArgs, Reply: gReply}); err != nil {
		t.Fatal(err)
	}

	expectedStoreStatus := &storage.StoreStatus{
		Desc:                 *storeDesc,
		NodeID:               1,
		RangeCount:           1,
		LeaderRangeCount:     1,
		AvailableRangeCount:  1,
		ReplicatedRangeCount: 0,
		Stats: engine.MVCCStats{
			LiveBytes: 1,
			KeyBytes:  1,
			ValBytes:  1,
			LiveCount: 1,
			KeyCount:  1,
			ValCount:  1,
		},
	}
	// Always wait twice, to ensure a full scan has occurred.
	store.WaitForRangeScanCompletion()
	store.WaitForRangeScanCompletion()
	oldstats := compareStoreStatus(t, store, expectedStoreStatus, 0)

	// Write some values left and right of the proposed split key.
	pArgs, pReply := putArgs([]byte("a"), content, rng.Desc().RaftID, store.StoreID())
	if err := store.ExecuteCmd(context.Background(), proto.Call{Args: pArgs, Reply: pReply}); err != nil {
		t.Fatal(err)
	}
	pArgs, pReply = putArgs([]byte("c"), content, rng.Desc().RaftID, store.StoreID())
	if err := store.ExecuteCmd(context.Background(), proto.Call{Args: pArgs, Reply: pReply}); err != nil {
		t.Fatal(err)
	}

	expectedStoreStatus = &storage.StoreStatus{
		Desc:                 oldstats.Desc,
		NodeID:               1,
		RangeCount:           1,
		LeaderRangeCount:     1,
		AvailableRangeCount:  1,
		ReplicatedRangeCount: 0,
		Stats: engine.MVCCStats{
			LiveBytes: 1,
			KeyBytes:  1,
			ValBytes:  1,
			LiveCount: oldstats.Stats.LiveCount + 1,
			KeyCount:  oldstats.Stats.KeyCount + 1,
			ValCount:  oldstats.Stats.ValCount + 1,
		},
	}
	store.WaitForRangeScanCompletion()
	store.WaitForRangeScanCompletion()
	oldstats = compareStoreStatus(t, store, expectedStoreStatus, 1)

	// Split the range.
	args, reply := adminSplitArgs(proto.KeyMin, splitKey, rng.Desc().RaftID, store.StoreID())
	if err := store.ExecuteCmd(context.Background(), proto.Call{Args: args, Reply: reply}); err != nil {
		t.Fatal(err)
	}

	expectedStoreStatus = &storage.StoreStatus{
		Desc:                 oldstats.Desc,
		NodeID:               1,
		RangeCount:           2,
		LeaderRangeCount:     2,
		AvailableRangeCount:  2,
		ReplicatedRangeCount: 0,
		Stats: engine.MVCCStats{
			LiveBytes: 1,
			KeyBytes:  1,
			ValBytes:  1,
			LiveCount: oldstats.Stats.LiveCount,
			KeyCount:  oldstats.Stats.KeyCount,
			ValCount:  oldstats.Stats.ValCount,
		},
	}
	store.WaitForRangeScanCompletion()
	store.WaitForRangeScanCompletion()
	oldstats = compareStoreStatus(t, store, expectedStoreStatus, 2)

	// Write some values left and right of the split key.
	rng = store.LookupRange([]byte("aa"), nil)
	pArgs, pReply = putArgs([]byte("aa"), content, rng.Desc().RaftID, store.StoreID())
	if err := store.ExecuteCmd(context.Background(), proto.Call{Args: pArgs, Reply: pReply}); err != nil {
		t.Fatal(err)
	}
	rng2 := store.LookupRange([]byte("cc"), nil)
	pArgs, pReply = putArgs([]byte("cc"), content, rng2.Desc().RaftID, store.StoreID())
	if err := store.ExecuteCmd(context.Background(), proto.Call{Args: pArgs, Reply: pReply}); err != nil {
		t.Fatal(err)
	}

	expectedStoreStatus = &storage.StoreStatus{
		Desc:                 oldstats.Desc,
		NodeID:               1,
		RangeCount:           2,
		LeaderRangeCount:     2,
		AvailableRangeCount:  2,
		ReplicatedRangeCount: 0,
		Stats: engine.MVCCStats{
			LiveBytes: 1,
			KeyBytes:  1,
			ValBytes:  1,
			LiveCount: oldstats.Stats.LiveCount + 1,
			KeyCount:  oldstats.Stats.KeyCount + 1,
			ValCount:  oldstats.Stats.ValCount + 1,
		},
	}
	store.WaitForRangeScanCompletion()
	store.WaitForRangeScanCompletion()
	compareStoreStatus(t, store, expectedStoreStatus, 3)
}
<commit_msg>Make a copy of TestStoreContext in TestStoreStatus.<commit_after>// Copyright 2015 The Cockroach Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
// implied. See the License for the specific language governing
// permissions and limitations under the License. See the AUTHORS file
// for names of contributors.
//
// Author: Bram Gruneir (bram+code@cockroachlabs.com)

package storage_test

import (
	"reflect"
	"testing"
	"time"

	"golang.org/x/net/context"

	"github.com/cockroachdb/cockroach/keys"
	"github.com/cockroachdb/cockroach/proto"
	"github.com/cockroachdb/cockroach/storage"
	"github.com/cockroachdb/cockroach/storage/engine"
	"github.com/cockroachdb/cockroach/util/hlc"
	"github.com/cockroachdb/cockroach/util/leaktest"
	gogoproto "github.com/gogo/protobuf/proto"
)

// compareStoreStatus ensures that the actual store status for the passed in
// store is updated correctly. It checks that the Desc.StoreID, Desc.Attrs,
// Desc.Node, Desc.Capacity.Capacity, NodeID, RangeCount, ReplicatedRangeCount
// are exactly correct and that the bytes and counts for Live, Key and Val are
// at least the expected value.
// The latest actual stats are returned.
func compareStoreStatus(t *testing.T, store *storage.Store, expectedStoreStatus *storage.StoreStatus, testNumber int) *storage.StoreStatus {
	storeStatusKey := keys.StoreStatusKey(int32(store.Ident.StoreID))
	gArgs, gReply := getArgs(storeStatusKey, 1, store.Ident.StoreID)
	if err := store.ExecuteCmd(context.Background(), proto.Call{Args: gArgs, Reply: gReply}); err != nil {
		t.Fatalf("%v: failure getting store status: %s", testNumber, err)
	}
	if gReply.Value == nil {
		t.Errorf("%v: could not find store status at: %s", testNumber, storeStatusKey)
	}
	storeStatus := &storage.StoreStatus{}
	if err := gogoproto.Unmarshal(gReply.Value.GetBytes(), storeStatus); err != nil {
		t.Fatalf("%v: could not unmarshal store status: %+v", testNumber, gReply)
	}

	// Values much match exactly.
	if expectedStoreStatus.Desc.StoreID != storeStatus.Desc.StoreID {
		t.Errorf("%v: actual Desc.StoreID does not match expected\nexpected: %+v\nactual: %v\n", testNumber, expectedStoreStatus, storeStatus)
	}
	if !reflect.DeepEqual(expectedStoreStatus.Desc.Attrs, storeStatus.Desc.Attrs) {
		t.Errorf("%v: actual Desc.Attrs does not match expected\nexpected: %+v\nactual: %v\n", testNumber, expectedStoreStatus, storeStatus)
	}
	if !reflect.DeepEqual(expectedStoreStatus.Desc.Node, storeStatus.Desc.Node) {
		t.Errorf("%v: actual Desc.Attrs does not match expected\nexpected: %+v\nactual: %v\n", testNumber, expectedStoreStatus, storeStatus)
	}
	if storeStatus.Desc.Capacity.Capacity != expectedStoreStatus.Desc.Capacity.Capacity {
		t.Errorf("%v: actual Desc.Capacity.Capacity does not match expected\nexpected: %+v\nactual: %v\n", testNumber, expectedStoreStatus, storeStatus)
	}
	if expectedStoreStatus.NodeID != storeStatus.NodeID {
		t.Errorf("%v: actual node ID does not match expected\nexpected: %+v\nactual: %v\n", testNumber, expectedStoreStatus, storeStatus)
	}
	if expectedStoreStatus.RangeCount != storeStatus.RangeCount {
		t.Errorf("%v: actual RangeCount does not match expected\nexpected: %+v\nactual: %v\n", testNumber, expectedStoreStatus, storeStatus)
	}
	if expectedStoreStatus.ReplicatedRangeCount != storeStatus.ReplicatedRangeCount {
		t.Errorf("%v: actual ReplicatedRangeCount does not match expected\nexpected: %+v\nactual: %v\n", testNumber, expectedStoreStatus, storeStatus)
	}

	// Values should be >= to expected values.
	if storeStatus.Stats.LiveBytes < expectedStoreStatus.Stats.LiveBytes {
		t.Errorf("%v: actual Live Bytes is not greater or equal to expected\nexpected: %+v\nactual: %v\n", testNumber, expectedStoreStatus, storeStatus)
	}
	if storeStatus.Stats.KeyBytes < expectedStoreStatus.Stats.KeyBytes {
		t.Errorf("%v: actual Key Bytes is not greater or equal to expected\nexpected: %+v\nactual: %v\n", testNumber, expectedStoreStatus, storeStatus)
	}
	if storeStatus.Stats.ValBytes < expectedStoreStatus.Stats.ValBytes {
		t.Errorf("%v: actual Val Bytes is not greater or equal to expected\nexpected: %+v\nactual: %v\n", testNumber, expectedStoreStatus, storeStatus)
	}
	if storeStatus.Stats.LiveCount < expectedStoreStatus.Stats.LiveCount {
		t.Errorf("%v: actual Live Count is not greater or equal to expected\nexpected: %+v\nactual: %v\n", testNumber, expectedStoreStatus, storeStatus)
	}
	if storeStatus.Stats.KeyCount < expectedStoreStatus.Stats.KeyCount {
		t.Errorf("%v: actual Key Count is not greater or equal to expected\nexpected: %+v\nactual: %v\n", testNumber, expectedStoreStatus, storeStatus)
	}
	if storeStatus.Stats.ValCount < expectedStoreStatus.Stats.ValCount {
		t.Errorf("%v: actual Val Count is not greater or equal to expected\nexpected: %+v\nactual: %v\n", testNumber, expectedStoreStatus, storeStatus)
	}
	return storeStatus
}

// TestStoreStatus checks the store status after each range scan to ensure that
// it is being updated correctly.
func TestStoreStatus(t *testing.T) {
	defer leaktest.AfterTest(t)
	ctx := storage.TestStoreContext
	ctx.ScanInterval = time.Duration(10 * time.Millisecond)
	store, stopper := createTestStoreWithEngine(t, engine.NewInMem(proto.Attributes{}, 10<<20), hlc.NewClock(hlc.NewManualClock(0).UnixNano), true, &ctx)
	defer stopper.Stop()
	splitKey := proto.Key("b")
	content := proto.Key("test content")

	storeDesc, err := store.Descriptor()
	if err != nil {
		t.Fatal(err)
	}

	store.WaitForInit()
	rng := store.LookupRange([]byte("a"), nil)
	// Perform a read from the range to ensure that the raft election has
	// completed.  We do not expect a value to be present.
	gArgs, gReply := getArgs([]byte("a"), rng.Desc().RaftID, store.StoreID())
	if err := store.ExecuteCmd(context.Background(), proto.Call{Args: gArgs, Reply: gReply}); err != nil {
		t.Fatal(err)
	}

	expectedStoreStatus := &storage.StoreStatus{
		Desc:                 *storeDesc,
		NodeID:               1,
		RangeCount:           1,
		LeaderRangeCount:     1,
		AvailableRangeCount:  1,
		ReplicatedRangeCount: 0,
		Stats: engine.MVCCStats{
			LiveBytes: 1,
			KeyBytes:  1,
			ValBytes:  1,
			LiveCount: 1,
			KeyCount:  1,
			ValCount:  1,
		},
	}
	// Always wait twice, to ensure a full scan has occurred.
	store.WaitForRangeScanCompletion()
	store.WaitForRangeScanCompletion()
	oldstats := compareStoreStatus(t, store, expectedStoreStatus, 0)

	// Write some values left and right of the proposed split key.
	pArgs, pReply := putArgs([]byte("a"), content, rng.Desc().RaftID, store.StoreID())
	if err := store.ExecuteCmd(context.Background(), proto.Call{Args: pArgs, Reply: pReply}); err != nil {
		t.Fatal(err)
	}
	pArgs, pReply = putArgs([]byte("c"), content, rng.Desc().RaftID, store.StoreID())
	if err := store.ExecuteCmd(context.Background(), proto.Call{Args: pArgs, Reply: pReply}); err != nil {
		t.Fatal(err)
	}

	expectedStoreStatus = &storage.StoreStatus{
		Desc:                 oldstats.Desc,
		NodeID:               1,
		RangeCount:           1,
		LeaderRangeCount:     1,
		AvailableRangeCount:  1,
		ReplicatedRangeCount: 0,
		Stats: engine.MVCCStats{
			LiveBytes: 1,
			KeyBytes:  1,
			ValBytes:  1,
			LiveCount: oldstats.Stats.LiveCount + 1,
			KeyCount:  oldstats.Stats.KeyCount + 1,
			ValCount:  oldstats.Stats.ValCount + 1,
		},
	}
	store.WaitForRangeScanCompletion()
	store.WaitForRangeScanCompletion()
	oldstats = compareStoreStatus(t, store, expectedStoreStatus, 1)

	// Split the range.
	args, reply := adminSplitArgs(proto.KeyMin, splitKey, rng.Desc().RaftID, store.StoreID())
	if err := store.ExecuteCmd(context.Background(), proto.Call{Args: args, Reply: reply}); err != nil {
		t.Fatal(err)
	}

	expectedStoreStatus = &storage.StoreStatus{
		Desc:                 oldstats.Desc,
		NodeID:               1,
		RangeCount:           2,
		LeaderRangeCount:     2,
		AvailableRangeCount:  2,
		ReplicatedRangeCount: 0,
		Stats: engine.MVCCStats{
			LiveBytes: 1,
			KeyBytes:  1,
			ValBytes:  1,
			LiveCount: oldstats.Stats.LiveCount,
			KeyCount:  oldstats.Stats.KeyCount,
			ValCount:  oldstats.Stats.ValCount,
		},
	}
	store.WaitForRangeScanCompletion()
	store.WaitForRangeScanCompletion()
	oldstats = compareStoreStatus(t, store, expectedStoreStatus, 2)

	// Write some values left and right of the split key.
	rng = store.LookupRange([]byte("aa"), nil)
	pArgs, pReply = putArgs([]byte("aa"), content, rng.Desc().RaftID, store.StoreID())
	if err := store.ExecuteCmd(context.Background(), proto.Call{Args: pArgs, Reply: pReply}); err != nil {
		t.Fatal(err)
	}
	rng2 := store.LookupRange([]byte("cc"), nil)
	pArgs, pReply = putArgs([]byte("cc"), content, rng2.Desc().RaftID, store.StoreID())
	if err := store.ExecuteCmd(context.Background(), proto.Call{Args: pArgs, Reply: pReply}); err != nil {
		t.Fatal(err)
	}

	expectedStoreStatus = &storage.StoreStatus{
		Desc:                 oldstats.Desc,
		NodeID:               1,
		RangeCount:           2,
		LeaderRangeCount:     2,
		AvailableRangeCount:  2,
		ReplicatedRangeCount: 0,
		Stats: engine.MVCCStats{
			LiveBytes: 1,
			KeyBytes:  1,
			ValBytes:  1,
			LiveCount: oldstats.Stats.LiveCount + 1,
			KeyCount:  oldstats.Stats.KeyCount + 1,
			ValCount:  oldstats.Stats.ValCount + 1,
		},
	}
	store.WaitForRangeScanCompletion()
	store.WaitForRangeScanCompletion()
	compareStoreStatus(t, store, expectedStoreStatus, 3)
}
<|endoftext|>