<commit_before>package main

import (
	"context"
	"errors"
	"fmt"
	log "github.com/Sirupsen/logrus"
	"github.com/ghodss/yaml"
	"github.com/loadimpact/speedboat/api"
	"github.com/loadimpact/speedboat/js"
	"github.com/loadimpact/speedboat/lib"
	"github.com/loadimpact/speedboat/simple"
	"github.com/loadimpact/speedboat/stats"
	"github.com/loadimpact/speedboat/stats/influxdb"
	"github.com/loadimpact/speedboat/ui"
	"gopkg.in/urfave/cli.v1"
	"io/ioutil"
	"net/url"
	"os"
	"os/signal"
	"sort"
	"strings"
	"sync"
	"syscall"
	"time"
)

const (
	TypeAuto = "auto"
	TypeURL  = "url"
	TypeJS   = "js"
)

var (
	ErrUnknownType = errors.New("Unable to infer type from argument; specify with -t/--type")
	ErrInvalidType = errors.New("Invalid type specified, see --help")
)

var commandRun = cli.Command{
	Name:      "run",
	Usage:     "Starts running a load test",
	ArgsUsage: "url|filename",
	Flags: []cli.Flag{
		cli.Int64Flag{
			Name:  "vus, u",
			Usage: "virtual users to simulate",
			Value: 10,
		},
		cli.Int64Flag{
			Name:  "max, m",
			Usage: "max number of virtual users, if more than --vus",
		},
		cli.DurationFlag{
			Name:  "duration, d",
			Usage: "test duration, 0 to run until cancelled",
			Value: 10 * time.Second,
		},
		cli.BoolFlag{
			Name:  "run, r",
			Usage: "start test immediately",
		},
		cli.StringFlag{
			Name:  "type, t",
			Usage: "input type, one of: auto, url, js",
			Value: "auto",
		},
		cli.BoolFlag{
			Name:  "quit, q",
			Usage: "quit immediately on test completion",
		},
		cli.BoolFlag{
			Name:  "quit-on-taint",
			Usage: "quit immediately if the test gets tainted",
		},
		cli.StringFlag{
			Name:  "out, o",
			Usage: "output metrics to an external data store",
		},
		cli.StringSliceFlag{
			Name:  "config, c",
			Usage: "read additional config files",
		},
	},
	Action: actionRun,
	Description: `Run starts a load test.

   This is the main entry point to Speedboat, and will do two things:

   - Construct an Engine and provide it with a Runner, depending on the first
     argument and the --type flag, which is used to execute the test.

   - Start an a web server on the address specified by the global --address
     flag, which serves a web interface and a REST API for remote control.

   For ease of use, you may also pass initial status parameters (vus, max,
   duration) to 'run', which will be applied through a normal API call.`,
}

var commandInspect = cli.Command{
	Name:      "inspect",
	Aliases:   []string{"i"},
	Usage:     "Merges and prints test configuration",
	ArgsUsage: "url|filename",
	Flags: []cli.Flag{
		cli.StringFlag{
			Name:  "type, t",
			Usage: "input type, one of: auto, url, js",
			Value: "auto",
		},
		cli.StringSliceFlag{
			Name:  "config, c",
			Usage: "read additional config files",
		},
	},
	Action: actionInspect,
}

func guessType(filename string) string {
	switch {
	case strings.Contains(filename, "://"):
		return TypeURL
	case strings.HasSuffix(filename, ".js"):
		return TypeJS
	default:
		return ""
	}
}

func makeRunner(filename, t string, opts *lib.Options) (lib.Runner, error) {
	if t == TypeAuto {
		t = guessType(filename)
	}

	switch t {
	case "":
		return nil, ErrUnknownType
	case TypeURL:
		return simple.New(filename)
	case TypeJS:
		rt, err := js.New()
		if err != nil {
			return nil, err
		}

		exports, err := rt.Load(filename)
		if err != nil {
			return nil, err
		}

		if err := rt.ExtractOptions(exports, opts); err != nil {
			return nil, err
		}
		return js.NewRunner(rt, exports)
	default:
		return nil, ErrInvalidType
	}
}

func parseCollectorString(s string) (t string, u *url.URL, err error) {
	parts := strings.SplitN(s, "=", 2)
	if len(parts) != 2 {
		return "", nil, errors.New("Malformed output; must be in the form 'type=url'")
	}

	u, err = url.Parse(parts[1])
	if err != nil {
		return "", nil, err
	}

	return parts[0], u, nil
}

func makeCollector(s string) (stats.Collector, error) {
	t, u, err := parseCollectorString(s)
	if err != nil {
		return nil, err
	}

	switch t {
	case "influxdb":
		return influxdb.New(u)
	default:
		return nil, errors.New("Unknown output type: " + t)
	}
}

func actionRun(cc *cli.Context) error {
	wg := sync.WaitGroup{}

	args := cc.Args()
	if len(args) != 1 {
		return cli.NewExitError("Wrong number of arguments!", 1)
	}

	// Collect CLI arguments, most (not all) relating to options.
	addr := cc.GlobalString("address")
	out := cc.String("out")
	opts := lib.Options{
		Run:         cliBool(cc, "run"),
		VUs:         cliInt64(cc, "vus"),
		VUsMax:      cliInt64(cc, "vus-max"),
		Duration:    cliDuration(cc, "duration"),
		Quit:        cliBool(cc, "quit"),
		QuitOnTaint: cliBool(cc, "quit-on-taint"),
	}

	// Make the Runner, extract script-defined options.
	filename := args[0]
	runnerType := cc.String("type")
	runnerOpts := lib.Options{}
	runner, err := makeRunner(filename, runnerType, &runnerOpts)
	if err != nil {
		log.WithError(err).Error("Couldn't create a runner")
		return err
	}
	opts = opts.Apply(runnerOpts)

	// Read config files.
	for _, filename := range cc.StringSlice("config") {
		data, err := ioutil.ReadFile(filename)
		if err != nil {
			return cli.NewExitError(err.Error(), 1)
		}

		var configOpts lib.Options
		if err := yaml.Unmarshal(data, &configOpts); err != nil {
			return cli.NewExitError(err.Error(), 1)
		}
		opts = opts.Apply(configOpts)
	}

	// CLI options have defaults, which are set as invalid, but have potentially nonzero values.
	// Flipping the Valid flag for all invalid options thus applies all defaults.
	opts = opts.SetAllValid(true)

	// Make the metric collector, if requested.
	var collector stats.Collector
	collectorString := "-"
	if out != "" {
		c, err := makeCollector(out)
		if err != nil {
			log.WithError(err).Error("Couldn't create output")
			return err
		}
		collector = c
		collectorString = fmt.Sprint(collector)
	}

	// Make the Engine
	engine, err := lib.NewEngine(runner)
	if err != nil {
		log.WithError(err).Error("Couldn't create the engine")
		return err
	}
	engineC, engineCancel := context.WithCancel(context.Background())
	engine.Collector = collector

	// Make the API Server
	srv := &api.Server{
		Engine: engine,
		Info:   lib.Info{Version: cc.App.Version},
	}
	srvC, srvCancel := context.WithCancel(context.Background())

	// Run the engine and API server in the background
	wg.Add(2)
	go func() {
		defer func() {
			log.Debug("Engine terminated")
			wg.Done()
		}()
		log.Debug("Starting engine...")
		if err := engine.Run(engineC, opts); err != nil {
			log.WithError(err).Error("Engine Error")
		}
		engineCancel()
	}()
	go func() {
		defer func() {
			log.Debug("API Server terminated")
			wg.Done()
		}()
		log.WithField("addr", addr).Debug("API Server starting...")
		srv.Run(srvC, addr)
		srvCancel()
	}()

	// Print the banner!
	fmt.Printf("Welcome to Speedboat v%s!\n", cc.App.Version)
	fmt.Printf("\n")
	fmt.Printf("  execution: local\n")
	fmt.Printf("     output: %s\n", collectorString)
	fmt.Printf("     script: %s\n", filename)
	fmt.Printf("             ↳ duration: %s\n", opts.Duration.String)
	fmt.Printf("             ↳ vus: %d, max: %d\n", opts.VUs.Int64, opts.VUsMax.Int64)
	fmt.Printf("\n")
	fmt.Printf("  web ui: http://%s/\n", addr)
	fmt.Printf("\n")

	progressBar := ui.ProgressBar{Width: 60}
	fmt.Printf(" starting %s -- / --\r", progressBar.String())

	// Wait for a signal or timeout before shutting down
	signals := make(chan os.Signal)
	signal.Notify(signals, os.Interrupt, syscall.SIGTERM)
	ticker := time.NewTicker(10 * time.Millisecond)

loop:
	for {
		select {
		case <-ticker.C:
			statusString := "running"
			if !engine.Status.Running.Bool {
				if engine.IsRunning() {
					statusString = "paused"
				} else {
					statusString = "stopping"
				}
			}

			atTime := time.Duration(engine.Status.AtTime.Int64)
			totalTime, finite := engine.TotalTime()
			progress := 0.0
			if finite {
				progress = float64(atTime) / float64(totalTime)
			}

			progressBar.Progress = progress
			fmt.Printf("%10s %s %10s / %s\r",
				statusString,
				progressBar.String(),
				atTime-(atTime%(100*time.Millisecond)),
				totalTime-(totalTime%(100*time.Millisecond)),
			)
		case <-srvC.Done():
			log.Debug("API server terminated; shutting down...")
			break loop
		case <-engineC.Done():
			log.Debug("Engine terminated; shutting down...")
			break loop
		case sig := <-signals:
			log.WithField("signal", sig).Debug("Signal received; shutting down...")
			break loop
		}
	}

	// Shut down the API server and engine.
	srvCancel()
	engineCancel()
	wg.Wait()

	// Test done, leave that status as the final progress bar!
	atTime := time.Duration(engine.Status.AtTime.Int64)
	progressBar.Progress = 1.0
	fmt.Printf("      done %s %10s / %s\n",
		progressBar.String(),
		atTime-(atTime%(100*time.Millisecond)),
		atTime-(atTime%(100*time.Millisecond)),
	)
	fmt.Printf("\n")

	// Print groups.
	var printGroup func(g *lib.Group, level int)
	printGroup = func(g *lib.Group, level int) {
		indent := strings.Repeat("  ", level)

		if g.Name != "" && g.Parent != nil {
			fmt.Printf("%s█ %s\n", indent, g.Name)
		}

		if len(g.Checks) > 0 {
			if g.Name != "" && g.Parent != nil {
				fmt.Printf("\n")
			}
			for _, check := range g.Checks {
				icon := "✓"
				if check.Fails > 0 {
					icon = "✗"
				}
				fmt.Printf("%s  %s %2.2f%% - %s\n",
					indent,
					icon,
					100*(float64(check.Passes)/float64(check.Passes+check.Fails)),
					check.Name,
				)
			}
			fmt.Printf("\n")
		}
		if len(g.Groups) > 0 {
			if g.Name != "" && g.Parent != nil && len(g.Checks) > 0 {
				fmt.Printf("\n")
			}
			for _, g := range g.Groups {
				printGroup(g, level+1)
			}
		}
	}

	groups := engine.Runner.GetGroups()
	for _, g := range groups {
		if g.Parent != nil {
			continue
		}
		printGroup(g, 1)
	}

	// Sort and print metrics.
	metrics := make(map[string]*stats.Metric, len(engine.Metrics))
	metricNames := make([]string, 0, len(engine.Metrics))
	for m, _ := range engine.Metrics {
		metrics[m.Name] = m
		metricNames = append(metricNames, m.Name)
	}
	sort.Strings(metricNames)

	for _, name := range metricNames {
		m := metrics[name]
		m.Sample = engine.Metrics[m].Format()
		val := metrics[name].Humanize()
		if val == "0" {
			continue
		}
		icon := "•"
		for _, threshold := range engine.Thresholds[name] {
			icon = "✓"
			if threshold.Failed {
				icon = "✗"
				break
			}
		}
		fmt.Printf("  %s %s: %s\n", icon, name, val)
	}

	if engine.Status.Tainted.Bool {
		return cli.NewExitError("", 99)
	}
	return nil
}

func actionInspect(cc *cli.Context) error {
	args := cc.Args()
	if len(args) != 1 {
		return cli.NewExitError("Wrong number of arguments!", 1)
	}
	filename := args[0]

	t := cc.String("type")
	if t == TypeAuto {
		t = guessType(filename)
	}

	var opts lib.Options
	switch t {
	case TypeJS:
		r, err := js.New()
		if err != nil {
			return cli.NewExitError(err.Error(), 1)
		}

		exports, err := r.Load(filename)
		if err != nil {
			return cli.NewExitError(err.Error(), 1)
		}
		if err := r.ExtractOptions(exports, &opts); err != nil {
			return cli.NewExitError(err.Error(), 1)
		}
	}

	for _, filename := range cc.StringSlice("config") {
		data, err := ioutil.ReadFile(filename)
		if err != nil {
			return cli.NewExitError(err.Error(), 1)
		}

		var configOpts lib.Options
		if err := yaml.Unmarshal(data, &configOpts); err != nil {
			return cli.NewExitError(err.Error(), 1)
		}
		opts = opts.Apply(configOpts)
	}

	return dumpYAML(opts)
}
<commit_msg>[tweak] Emphasize threshold-less metrics less<commit_after>package main

import (
	"context"
	"errors"
	"fmt"
	log "github.com/Sirupsen/logrus"
	"github.com/ghodss/yaml"
	"github.com/loadimpact/speedboat/api"
	"github.com/loadimpact/speedboat/js"
	"github.com/loadimpact/speedboat/lib"
	"github.com/loadimpact/speedboat/simple"
	"github.com/loadimpact/speedboat/stats"
	"github.com/loadimpact/speedboat/stats/influxdb"
	"github.com/loadimpact/speedboat/ui"
	"gopkg.in/urfave/cli.v1"
	"io/ioutil"
	"net/url"
	"os"
	"os/signal"
	"sort"
	"strings"
	"sync"
	"syscall"
	"time"
)

const (
	TypeAuto = "auto"
	TypeURL  = "url"
	TypeJS   = "js"
)

var (
	ErrUnknownType = errors.New("Unable to infer type from argument; specify with -t/--type")
	ErrInvalidType = errors.New("Invalid type specified, see --help")
)

var commandRun = cli.Command{
	Name:      "run",
	Usage:     "Starts running a load test",
	ArgsUsage: "url|filename",
	Flags: []cli.Flag{
		cli.Int64Flag{
			Name:  "vus, u",
			Usage: "virtual users to simulate",
			Value: 10,
		},
		cli.Int64Flag{
			Name:  "max, m",
			Usage: "max number of virtual users, if more than --vus",
		},
		cli.DurationFlag{
			Name:  "duration, d",
			Usage: "test duration, 0 to run until cancelled",
			Value: 10 * time.Second,
		},
		cli.BoolFlag{
			Name:  "run, r",
			Usage: "start test immediately",
		},
		cli.StringFlag{
			Name:  "type, t",
			Usage: "input type, one of: auto, url, js",
			Value: "auto",
		},
		cli.BoolFlag{
			Name:  "quit, q",
			Usage: "quit immediately on test completion",
		},
		cli.BoolFlag{
			Name:  "quit-on-taint",
			Usage: "quit immediately if the test gets tainted",
		},
		cli.StringFlag{
			Name:  "out, o",
			Usage: "output metrics to an external data store",
		},
		cli.StringSliceFlag{
			Name:  "config, c",
			Usage: "read additional config files",
		},
	},
	Action: actionRun,
	Description: `Run starts a load test.

   This is the main entry point to Speedboat, and will do two things:

   - Construct an Engine and provide it with a Runner, depending on the first
     argument and the --type flag, which is used to execute the test.

   - Start an a web server on the address specified by the global --address
     flag, which serves a web interface and a REST API for remote control.

   For ease of use, you may also pass initial status parameters (vus, max,
   duration) to 'run', which will be applied through a normal API call.`,
}

var commandInspect = cli.Command{
	Name:      "inspect",
	Aliases:   []string{"i"},
	Usage:     "Merges and prints test configuration",
	ArgsUsage: "url|filename",
	Flags: []cli.Flag{
		cli.StringFlag{
			Name:  "type, t",
			Usage: "input type, one of: auto, url, js",
			Value: "auto",
		},
		cli.StringSliceFlag{
			Name:  "config, c",
			Usage: "read additional config files",
		},
	},
	Action: actionInspect,
}

func guessType(filename string) string {
	switch {
	case strings.Contains(filename, "://"):
		return TypeURL
	case strings.HasSuffix(filename, ".js"):
		return TypeJS
	default:
		return ""
	}
}

func makeRunner(filename, t string, opts *lib.Options) (lib.Runner, error) {
	if t == TypeAuto {
		t = guessType(filename)
	}

	switch t {
	case "":
		return nil, ErrUnknownType
	case TypeURL:
		return simple.New(filename)
	case TypeJS:
		rt, err := js.New()
		if err != nil {
			return nil, err
		}

		exports, err := rt.Load(filename)
		if err != nil {
			return nil, err
		}

		if err := rt.ExtractOptions(exports, opts); err != nil {
			return nil, err
		}
		return js.NewRunner(rt, exports)
	default:
		return nil, ErrInvalidType
	}
}

func parseCollectorString(s string) (t string, u *url.URL, err error) {
	parts := strings.SplitN(s, "=", 2)
	if len(parts) != 2 {
		return "", nil, errors.New("Malformed output; must be in the form 'type=url'")
	}

	u, err = url.Parse(parts[1])
	if err != nil {
		return "", nil, err
	}

	return parts[0], u, nil
}

func makeCollector(s string) (stats.Collector, error) {
	t, u, err := parseCollectorString(s)
	if err != nil {
		return nil, err
	}

	switch t {
	case "influxdb":
		return influxdb.New(u)
	default:
		return nil, errors.New("Unknown output type: " + t)
	}
}

func actionRun(cc *cli.Context) error {
	wg := sync.WaitGroup{}

	args := cc.Args()
	if len(args) != 1 {
		return cli.NewExitError("Wrong number of arguments!", 1)
	}

	// Collect CLI arguments, most (not all) relating to options.
	addr := cc.GlobalString("address")
	out := cc.String("out")
	opts := lib.Options{
		Run:         cliBool(cc, "run"),
		VUs:         cliInt64(cc, "vus"),
		VUsMax:      cliInt64(cc, "vus-max"),
		Duration:    cliDuration(cc, "duration"),
		Quit:        cliBool(cc, "quit"),
		QuitOnTaint: cliBool(cc, "quit-on-taint"),
	}

	// Make the Runner, extract script-defined options.
	filename := args[0]
	runnerType := cc.String("type")
	runnerOpts := lib.Options{}
	runner, err := makeRunner(filename, runnerType, &runnerOpts)
	if err != nil {
		log.WithError(err).Error("Couldn't create a runner")
		return err
	}
	opts = opts.Apply(runnerOpts)

	// Read config files.
	for _, filename := range cc.StringSlice("config") {
		data, err := ioutil.ReadFile(filename)
		if err != nil {
			return cli.NewExitError(err.Error(), 1)
		}

		var configOpts lib.Options
		if err := yaml.Unmarshal(data, &configOpts); err != nil {
			return cli.NewExitError(err.Error(), 1)
		}
		opts = opts.Apply(configOpts)
	}

	// CLI options have defaults, which are set as invalid, but have potentially nonzero values.
	// Flipping the Valid flag for all invalid options thus applies all defaults.
	opts = opts.SetAllValid(true)

	// Make the metric collector, if requested.
	var collector stats.Collector
	collectorString := "-"
	if out != "" {
		c, err := makeCollector(out)
		if err != nil {
			log.WithError(err).Error("Couldn't create output")
			return err
		}
		collector = c
		collectorString = fmt.Sprint(collector)
	}

	// Make the Engine
	engine, err := lib.NewEngine(runner)
	if err != nil {
		log.WithError(err).Error("Couldn't create the engine")
		return err
	}
	engineC, engineCancel := context.WithCancel(context.Background())
	engine.Collector = collector

	// Make the API Server
	srv := &api.Server{
		Engine: engine,
		Info:   lib.Info{Version: cc.App.Version},
	}
	srvC, srvCancel := context.WithCancel(context.Background())

	// Run the engine and API server in the background
	wg.Add(2)
	go func() {
		defer func() {
			log.Debug("Engine terminated")
			wg.Done()
		}()
		log.Debug("Starting engine...")
		if err := engine.Run(engineC, opts); err != nil {
			log.WithError(err).Error("Engine Error")
		}
		engineCancel()
	}()
	go func() {
		defer func() {
			log.Debug("API Server terminated")
			wg.Done()
		}()
		log.WithField("addr", addr).Debug("API Server starting...")
		srv.Run(srvC, addr)
		srvCancel()
	}()

	// Print the banner!
	fmt.Printf("Welcome to Speedboat v%s!\n", cc.App.Version)
	fmt.Printf("\n")
	fmt.Printf("  execution: local\n")
	fmt.Printf("     output: %s\n", collectorString)
	fmt.Printf("     script: %s\n", filename)
	fmt.Printf("             ↳ duration: %s\n", opts.Duration.String)
	fmt.Printf("             ↳ vus: %d, max: %d\n", opts.VUs.Int64, opts.VUsMax.Int64)
	fmt.Printf("\n")
	fmt.Printf("  web ui: http://%s/\n", addr)
	fmt.Printf("\n")

	progressBar := ui.ProgressBar{Width: 60}
	fmt.Printf(" starting %s -- / --\r", progressBar.String())

	// Wait for a signal or timeout before shutting down
	signals := make(chan os.Signal)
	signal.Notify(signals, os.Interrupt, syscall.SIGTERM)
	ticker := time.NewTicker(10 * time.Millisecond)

loop:
	for {
		select {
		case <-ticker.C:
			statusString := "running"
			if !engine.Status.Running.Bool {
				if engine.IsRunning() {
					statusString = "paused"
				} else {
					statusString = "stopping"
				}
			}

			atTime := time.Duration(engine.Status.AtTime.Int64)
			totalTime, finite := engine.TotalTime()
			progress := 0.0
			if finite {
				progress = float64(atTime) / float64(totalTime)
			}

			progressBar.Progress = progress
			fmt.Printf("%10s %s %10s / %s\r",
				statusString,
				progressBar.String(),
				atTime-(atTime%(100*time.Millisecond)),
				totalTime-(totalTime%(100*time.Millisecond)),
			)
		case <-srvC.Done():
			log.Debug("API server terminated; shutting down...")
			break loop
		case <-engineC.Done():
			log.Debug("Engine terminated; shutting down...")
			break loop
		case sig := <-signals:
			log.WithField("signal", sig).Debug("Signal received; shutting down...")
			break loop
		}
	}

	// Shut down the API server and engine.
	srvCancel()
	engineCancel()
	wg.Wait()

	// Test done, leave that status as the final progress bar!
	atTime := time.Duration(engine.Status.AtTime.Int64)
	progressBar.Progress = 1.0
	fmt.Printf("      done %s %10s / %s\n",
		progressBar.String(),
		atTime-(atTime%(100*time.Millisecond)),
		atTime-(atTime%(100*time.Millisecond)),
	)
	fmt.Printf("\n")

	// Print groups.
	var printGroup func(g *lib.Group, level int)
	printGroup = func(g *lib.Group, level int) {
		indent := strings.Repeat("  ", level)

		if g.Name != "" && g.Parent != nil {
			fmt.Printf("%s█ %s\n", indent, g.Name)
		}

		if len(g.Checks) > 0 {
			if g.Name != "" && g.Parent != nil {
				fmt.Printf("\n")
			}
			for _, check := range g.Checks {
				icon := "✓"
				if check.Fails > 0 {
					icon = "✗"
				}
				fmt.Printf("%s  %s %2.2f%% - %s\n",
					indent,
					icon,
					100*(float64(check.Passes)/float64(check.Passes+check.Fails)),
					check.Name,
				)
			}
			fmt.Printf("\n")
		}
		if len(g.Groups) > 0 {
			if g.Name != "" && g.Parent != nil && len(g.Checks) > 0 {
				fmt.Printf("\n")
			}
			for _, g := range g.Groups {
				printGroup(g, level+1)
			}
		}
	}

	groups := engine.Runner.GetGroups()
	for _, g := range groups {
		if g.Parent != nil {
			continue
		}
		printGroup(g, 1)
	}

	// Sort and print metrics.
	metrics := make(map[string]*stats.Metric, len(engine.Metrics))
	metricNames := make([]string, 0, len(engine.Metrics))
	for m, _ := range engine.Metrics {
		metrics[m.Name] = m
		metricNames = append(metricNames, m.Name)
	}
	sort.Strings(metricNames)

	for _, name := range metricNames {
		m := metrics[name]
		m.Sample = engine.Metrics[m].Format()
		val := metrics[name].Humanize()
		if val == "0" {
			continue
		}
		icon := " "
		for _, threshold := range engine.Thresholds[name] {
			icon = "✓"
			if threshold.Failed {
				icon = "✗"
				break
			}
		}
		fmt.Printf("  %s %s: %s\n", icon, name, val)
	}

	if engine.Status.Tainted.Bool {
		return cli.NewExitError("", 99)
	}
	return nil
}

func actionInspect(cc *cli.Context) error {
	args := cc.Args()
	if len(args) != 1 {
		return cli.NewExitError("Wrong number of arguments!", 1)
	}
	filename := args[0]

	t := cc.String("type")
	if t == TypeAuto {
		t = guessType(filename)
	}

	var opts lib.Options
	switch t {
	case TypeJS:
		r, err := js.New()
		if err != nil {
			return cli.NewExitError(err.Error(), 1)
		}

		exports, err := r.Load(filename)
		if err != nil {
			return cli.NewExitError(err.Error(), 1)
		}
		if err := r.ExtractOptions(exports, &opts); err != nil {
			return cli.NewExitError(err.Error(), 1)
		}
	}

	for _, filename := range cc.StringSlice("config") {
		data, err := ioutil.ReadFile(filename)
		if err != nil {
			return cli.NewExitError(err.Error(), 1)
		}

		var configOpts lib.Options
		if err := yaml.Unmarshal(data, &configOpts); err != nil {
			return cli.NewExitError(err.Error(), 1)
		}
		opts = opts.Apply(configOpts)
	}

	return dumpYAML(opts)
}
<|endoftext|>


<commit_before>// Package exphttp implements HTTP request/response timing collection for the
// standard `net/http` package.
//
// Although it's more efficient with a little tweak.
package exphttp

import (
	"expvar"
	"fmt"
	"log"
	"net/http"
	"time"
)

// DefaultGranularity is the default level of granularity for recording
// counters over time. For example, a RateCounter with an interval of 1 minute
// and a granulartiy of 30 will be accurate to within 2 seconds.
//
// Thus you want to make sure your polling interval is greater than the
// measurement interval divided by granularity.
const DefaultGranularity = 32

var expHandlers = expvar.NewMap("exphttp")

// ExpHandlerFunc is a http.HandlerFunc that returns it's own HTTP StatusCode.
type ExpHandlerFunc func(w http.ResponseWriter, r *http.Request) int

type getStatusCode struct {
	http.ResponseWriter
	code int
}

func (w *getStatusCode) WriteHeader(c int) {
	w.code = c
	w.ResponseWriter.WriteHeader(c)
}

// MakeExpHandlerFunc wraps a http.HandlerFunc so that the response status code
// is accessible. It is more efficient to update your code to implement
// ExpHandlerFunc and return the code directly.
func MakeExpHandlerFunc(h http.HandlerFunc) ExpHandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) int {
		w2 := &getStatusCode{w, 500}
		h(w2, r)
		return w2.code
	}
}

// ExpHandler is an http.Handler that exposes request/response timing
// information via the `expvar` stdlib package.
type ExpHandler struct {
	// Name of the handler/endpoint.
	Name string

	// Stats contains the request/response stats that are exposed.
	Stats *expvar.Map

	// Durations are the time spans for the rate counters. Only parsed once in
	// the first incoming request to populate counters.
	Durations map[string]time.Duration

	// HandlerFunc is the ExpHandlerFunc that is tracked.
	HandlerFunc ExpHandlerFunc

	didInit      bool
	reqCounters  []*RateCounter
	respCounters []*RateCounter
}

// NewExpHandler creates a new ExpHandler, publishes a new expvar.Map to track
// it, sets a default Durations={"min": time.Minute}, and adds name to the
// exposed "exphttp" map so that polling code can auto-discover.
func NewExpHandler(name string, h ExpHandlerFunc) *ExpHandler {
	e := &ExpHandler{
		Name:        name,
		Stats:       expvar.NewMap(name),
		Durations:   map[string]time.Duration{"min": time.Minute},
		HandlerFunc: h,
	}

	expHandlers.Add(name, 1)
	return e
}

func (e *ExpHandler) init() {
	e.reqCounters = make([]*RateCounter, 0, len(e.Durations))
	e.respCounters = make([]*RateCounter, 0, len(e.Durations))

	for key, dur := range e.Durations {
		r1 := NewRateCounter(dur)
		r2 := NewRateCounter(dur)
		e.Stats.Set("requests_per_"+key, r1)
		e.Stats.Set("responses_per_"+key, r2)
		e.reqCounters = append(e.reqCounters, r1)
		e.respCounters = append(e.respCounters, r2)
	}
	e.didInit = true
}

// ServeHTTP implements the http.Handler interface.
func (e *ExpHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if !e.didInit {
		e.init()
	}

	e.Stats.Add("requests", 1)
	for _, rc := range e.reqCounters {
		rc.Add(1)
	}

	startTime := time.Now()
	defer func() {
		if p := recover(); p != nil {
			elap := time.Now().Sub(startTime).Nanoseconds()

			log.Println("caught panic: ", p)
			e.Stats.Add("panics", 1)
			e.Stats.Add("responses", 1)
			for _, rc := range e.respCounters {
				rc.Add(1)
			}
			e.Stats.Add("responses.500", 1)
			e.Stats.Add("responses.500.total_ns", elap)

			http.Error(w, "server error", http.StatusInternalServerError)
		}
	}()
	////////

	code := e.HandlerFunc(w, r)

	////////
	elapsed := time.Now().Sub(startTime).Nanoseconds()
	log.Println(float64(elapsed)/1000000.0, "ms -- ", code, "--", r.URL)

	e.Stats.Add("responses", 1)
	for _, rc := range e.respCounters {
		rc.Add(1)
	}

	switch code {
	case http.StatusOK:
		e.Stats.Add("responses.200", 1)
		e.Stats.Add("responses.200.total_ns", elapsed)
	case http.StatusBadRequest:
		e.Stats.Add("responses.400", 1)
		e.Stats.Add("responses.400.total_ns", elapsed)
	case http.StatusUnauthorized:
		e.Stats.Add("responses.401", 1)
		e.Stats.Add("responses.401.total_ns", elapsed)
	case http.StatusInternalServerError:
		e.Stats.Add("responses.500", 1)
		e.Stats.Add("responses.500.total_ns", elapsed)
	default:
		e.Stats.Add(fmt.Sprintf("responses.%d", code), 1)
		e.Stats.Add(fmt.Sprintf("responses.%d.total_ns", code), elapsed)
	}
}
<commit_msg>make logging optional<commit_after>// Package exphttp implements HTTP request/response timing collection for the
// standard `net/http` package.
//
// Example usage:
//     func getThePage(w http.ResponseWriter, r *http.Request) {
//        fmt.Fprint(w, "neat page")
//     }
//
//     http.HandleFunc("/thepage/", getThePage)
//
//     // same page, but now with stats tracked!
//     http.Handle("/thepage2/", exphttp.NewExpHandler("thepage2", MakeExpHandlerFunc(getThePage)))
//
// Although it's more efficient with a slightly tweaked http.HandlerFunc that
// returns the response status code. Here the handler is modified to return its
// status code work with exphttp:
//     func getThePage2(w http.ResponseWriter, r *http.Request) int {
//        fmt.Fprint(w, "neat page")
//        return http.StatusOK
//     }
//
//     http.Handle("/thepage2/", exphttp.NewExpHandler("thepage2", getThePage2))
//
package exphttp

import (
	"expvar"
	"fmt"
	"log"
	"net/http"
	"os"
	"time"
)

// DefaultGranularity is the default level of granularity for recording
// counters over time. For example, a RateCounter with an interval of 1 minute
// and a granulartiy of 30 will be accurate to within 2 seconds.
//
// Thus you want to make sure your polling interval is greater than the
// measurement interval divided by granularity.
const DefaultGranularity = 32

var expHandlers = expvar.NewMap("exphttp")
// DefaultLogger is used when creating new ExpHandlers, and used to log requests
// and timing info to Stderr.
//
// Set to nil to disable before calling NewExpHandler()
var DefaultLogger = log.New(os.Stderr, "", log.LstdFlags)

// ExpHandlerFunc is a http.HandlerFunc that returns it's own HTTP StatusCode.
type ExpHandlerFunc func(w http.ResponseWriter, r *http.Request) int

type getStatusCode struct {
	http.ResponseWriter
	code int
}

func (w *getStatusCode) WriteHeader(c int) {
	w.code = c
	w.ResponseWriter.WriteHeader(c)
}

// MakeExpHandlerFunc wraps a http.HandlerFunc so that the response status code
// is accessible. It is more efficient to update your code to implement
// ExpHandlerFunc and return the code directly.
func MakeExpHandlerFunc(h http.HandlerFunc) ExpHandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) int {
		w2 := &getStatusCode{w, 500}
		h(w2, r)
		return w2.code
	}
}

// ExpHandler is an http.Handler that exposes request/response timing
// information via the `expvar` stdlib package.
type ExpHandler struct {
	// Name of the handler/endpoint.
	Name string

	// Stats contains the request/response stats that are exposed.
	Stats *expvar.Map

	// Durations are the time spans for the rate counters. Only parsed once in
	// the first incoming request to populate counters.
	Durations map[string]time.Duration

	// HandlerFunc is the ExpHandlerFunc that is tracked.
	HandlerFunc ExpHandlerFunc

	// Log requests to this logger if non-nil.
	Log *log.Logger

	didInit      bool
	reqCounters  []*RateCounter
	respCounters []*RateCounter
}

// NewExpHandler creates a new ExpHandler, publishes a new expvar.Map to track
// it, sets a default Durations={"min": time.Minute}, set Log=DefaultLogger,
// and adds name to the exposed "exphttp" map so that stats polling code
// can auto-discover.
func NewExpHandler(name string, h ExpHandlerFunc) *ExpHandler {
	e := &ExpHandler{
		Name:        name,
		Stats:       expvar.NewMap(name),
		Durations:   map[string]time.Duration{"min": time.Minute},
		HandlerFunc: h,
		Log:         DefaultLogger,
	}

	expHandlers.Add(name, 1)
	return e
}

func (e *ExpHandler) init() {
	e.reqCounters = make([]*RateCounter, 0, len(e.Durations))
	e.respCounters = make([]*RateCounter, 0, len(e.Durations))

	for key, dur := range e.Durations {
		r1 := NewRateCounter(dur)
		r2 := NewRateCounter(dur)
		e.Stats.Set("requests_per_"+key, r1)
		e.Stats.Set("responses_per_"+key, r2)
		e.reqCounters = append(e.reqCounters, r1)
		e.respCounters = append(e.respCounters, r2)
	}
	e.didInit = true
}

// ServeHTTP implements the http.Handler interface.
func (e *ExpHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if !e.didInit {
		e.init()
	}

	e.Stats.Add("requests", 1)
	for _, rc := range e.reqCounters {
		rc.Add(1)
	}

	startTime := time.Now()
	defer func() {
		if p := recover(); p != nil {
			elap := time.Now().Sub(startTime).Nanoseconds()

			if e.Log != nil {
				e.Log.Println("caught panic: ", p)
			}
			e.Stats.Add("panics", 1)
			e.Stats.Add("responses", 1)
			for _, rc := range e.respCounters {
				rc.Add(1)
			}
			e.Stats.Add("responses.500", 1)
			e.Stats.Add("responses.500.total_ns", elap)

			http.Error(w, "server error", http.StatusInternalServerError)
		}
	}()
	////////

	code := e.HandlerFunc(w, r)

	////////
	elapsed := time.Now().Sub(startTime).Nanoseconds()
	if e.Log != nil {
		e.Log.Println(float64(elapsed)/1000000.0, "ms --", code, "--", r.Method, r.URL)
	}

	e.Stats.Add("responses", 1)
	for _, rc := range e.respCounters {
		rc.Add(1)
	}

	switch code {
	case http.StatusOK:
		e.Stats.Add("responses.200", 1)
		e.Stats.Add("responses.200.total_ns", elapsed)
	case http.StatusBadRequest:
		e.Stats.Add("responses.400", 1)
		e.Stats.Add("responses.400.total_ns", elapsed)
	case http.StatusUnauthorized:
		e.Stats.Add("responses.401", 1)
		e.Stats.Add("responses.401.total_ns", elapsed)
	case http.StatusInternalServerError:
		e.Stats.Add("responses.500", 1)
		e.Stats.Add("responses.500.total_ns", elapsed)
	default:
		e.Stats.Add(fmt.Sprintf("responses.%d", code), 1)
		e.Stats.Add(fmt.Sprintf("responses.%d.total_ns", code), elapsed)
	}
}
<|endoftext|>


<commit_before>package filters

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

	"github.com/solher/arangolite/filters"
)

// Filter defines a way of filtering AQL queries.
type Filter struct {
	Offset  int                      `json:"offset"`
	Limit   int                      `json:"limit"`
	Sort    []string                 `json:"sort"`
	Where   []map[string]interface{} `json:"where"`
	Options map[string]interface{}   `json:"options"`
}

type processedFilter struct {
	OffsetLimit string
	Sort        string
	Where       string
}

// FromRequest returns a filter object from a http request.
func FromRequest(r *http.Request) (*filters.Filter, error) {
	param := r.URL.Query().Get("filter")

	if len(param) == 0 {
		param = r.URL.Query().Get("Filter")
	}

	if len(param) == 0 {
		return nil, nil
	}

	filter, err := FromJSON(param)
	if err != nil {
		return nil, err
	}

	return filter, nil
}

// FromJSON converts a JSON filter to a Filter object.
func FromJSON(jsonFilter string) (*Filter, error) {
	filter := &Filter{}

	if err := json.Unmarshal([]byte(jsonFilter), filter); err != nil {
		return nil, err
	}

	return filter, nil
}

// ToAQL converts a Filter object to its translation in AQL.
// "tmpVar" is the AQL var name to apply the filter on.
func ToAQL(tmpVar string, f *Filter) (string, error) {
	fp := newFilterProcessor(tmpVar)
	filter, err := fp.Process(f)
	if err != nil {
		return "", err
	}

	aqlFilter := bytes.NewBuffer(nil)

	if len(filter.Where) != 0 {
		aqlFilter.WriteString("FILTER ")
		aqlFilter.WriteString(filter.Where)
	}

	if len(filter.Sort) != 0 {
		aqlFilter.WriteString(" SORT ")
		aqlFilter.WriteString(filter.Sort)
	}

	if len(filter.OffsetLimit) != 0 {
		aqlFilter.WriteString(" LIMIT ")
		aqlFilter.WriteString(filter.OffsetLimit)
	}

	return aqlFilter.String(), nil
}
<commit_msg>Method added to get a filter from a http request fixed.<commit_after>package filters

import (
	"bytes"
	"encoding/json"
	"net/http"
)

// Filter defines a way of filtering AQL queries.
type Filter struct {
	Offset  int                      `json:"offset"`
	Limit   int                      `json:"limit"`
	Sort    []string                 `json:"sort"`
	Where   []map[string]interface{} `json:"where"`
	Options map[string]interface{}   `json:"options"`
}

type processedFilter struct {
	OffsetLimit string
	Sort        string
	Where       string
}

// FromRequest returns a filter object from a http request.
func FromRequest(r *http.Request) (*Filter, error) {
	param := r.URL.Query().Get("filter")

	if len(param) == 0 {
		param = r.URL.Query().Get("Filter")
	}

	if len(param) == 0 {
		return nil, nil
	}

	filter, err := FromJSON(param)
	if err != nil {
		return nil, err
	}

	return filter, nil
}

// FromJSON converts a JSON filter to a Filter object.
func FromJSON(jsonFilter string) (*Filter, error) {
	filter := &Filter{}

	if err := json.Unmarshal([]byte(jsonFilter), filter); err != nil {
		return nil, err
	}

	return filter, nil
}

// ToAQL converts a Filter object to its translation in AQL.
// "tmpVar" is the AQL var name to apply the filter on.
func ToAQL(tmpVar string, f *Filter) (string, error) {
	fp := newFilterProcessor(tmpVar)
	filter, err := fp.Process(f)
	if err != nil {
		return "", err
	}

	aqlFilter := bytes.NewBuffer(nil)

	if len(filter.Where) != 0 {
		aqlFilter.WriteString("FILTER ")
		aqlFilter.WriteString(filter.Where)
	}

	if len(filter.Sort) != 0 {
		aqlFilter.WriteString(" SORT ")
		aqlFilter.WriteString(filter.Sort)
	}

	if len(filter.OffsetLimit) != 0 {
		aqlFilter.WriteString(" LIMIT ")
		aqlFilter.WriteString(filter.OffsetLimit)
	}

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


<commit_before>package clone

import (
	"crypto/md5"
	"encoding/hex"
	"github.com/jeffail/tunny"
	csv "github.com/whosonfirst/go-whosonfirst-csv"
	log "github.com/whosonfirst/go-whosonfirst-log"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"path"
	"path/filepath"
	"runtime"
	"strings"
	"sync"
	"sync/atomic"
)

type WOFClone struct {
	Source    string
	Dest      string
	Count     int64
	Success   int64
	Error     int64
	Skipped   int64
	Scheduled int64
	Completed int64
	Logger    *log.WOFLogger
	client    *http.Client
	pool      *tunny.WorkPool
}

func NewWOFClone(source string, dest string, procs int, logger *log.WOFLogger) *WOFClone {

	// cd := &sync.Cond{L: &sync.Mutex{}}

	cl := &http.Client{}

	runtime.GOMAXPROCS(procs)

	pool, _ := tunny.CreatePoolGeneric(procs).Open()

	c := WOFClone{
		Count:   0,
		Success: 0,
		Error:   0,
		Skipped: 0,
		Source:  source,
		Dest:    dest,
		Logger:  logger,
		client:  cl,
		pool:    pool,
	}

	return &c
}

func (c *WOFClone) CloneMetaFile(file string) error {

	abs_path, _ := filepath.Abs(file)
	// c.Logger.Debug("Parse meta file %s", abs_path)

	reader, read_err := csv.NewDictReader(abs_path)

	if read_err != nil {
		c.Logger.Error("Failed to read %s, because %v", abs_path, read_err)
		return read_err
	}

	wg := new(sync.WaitGroup)

	for {
		row, err := reader.Read()

		if err == io.EOF {
			break
		}

		if err != nil {
			return err
		}

		rel_path, ok := row["path"]

		if !ok {
			continue
		}

		ensure_changes := true

		/*
			skip_existing := true

			remote := c.Source + rel_path
			local := path.Join(c.Dest, rel_path)

			if !os.IsNotExist(err) {

			   if skip_existing {
			      c.Logger.Debug("%s already exists and we are skipping things that exist", local)
			      continue
			   }

			   // hash file here...

			   change, _ := c.HasHashChanged(file_hash, remote)

			   if !change{
			      c.logger.Info("no changes to %s", file_hash)
			      ensure_changes = false
			}
		*/

		wg.Add(1)
		atomic.AddInt64(&c.Scheduled, 1)

		go func() {

			defer wg.Done()

			_, err = c.pool.SendWork(func() {

				cl_err := c.ClonePath(rel_path, ensure_changes)

				if cl_err != nil {
					atomic.AddInt64(&c.Error, 1)
				} else {
					atomic.AddInt64(&c.Success, 1)
				}

				atomic.AddInt64(&c.Completed, 1)
				c.Status()
			})

		}()
	}

	wg.Wait()
	return nil
}

func (c *WOFClone) ClonePath(rel_path string, ensure_changes bool) error {

	atomic.AddInt64(&c.Count, 1)

	remote := c.Source + rel_path
	local := path.Join(c.Dest, rel_path)

	_, err := os.Stat(local)

	if !os.IsNotExist(err) && ensure_changes {

		change, _ := c.HasChanged(local, remote)

		if !change {

			c.Logger.Debug("%s has not changed so skipping", local)
			atomic.AddInt64(&c.Skipped, 1)
			return nil
		}

	}

	process_err := c.Process(remote, local)

	if process_err != nil {
		atomic.AddInt64(&c.Error, 1)
		return process_err
	}

	return nil
}

// don't return true if there's a problem - move that logic up above

func (c *WOFClone) HasChanged(local string, remote string) (bool, error) {

	change := true

	body, err := ioutil.ReadFile(local)

	if err != nil {
		c.Logger.Error("Failed to read %s, becase %v", local, err)
		return change, err
	}

	hash := md5.Sum(body)
	local_hash := hex.EncodeToString(hash[:])

	return c.HasHashChanged(local_hash, remote)
}

func (c *WOFClone) HasHashChanged(local_hash string, remote string) (bool, error) {

	change := true

	rsp, err := c.Fetch("HEAD", remote)

	if err != nil {
		return change, err
	}

	defer rsp.Body.Close()

	etag := rsp.Header.Get("Etag")
	remote_hash := strings.Replace(etag, "\"", "", -1)

	if local_hash == remote_hash {
		change = false
	}

	return change, nil
}

func (c *WOFClone) Process(remote string, local string) error {

	c.Logger.Debug("fetch %s and store in %s", remote, local)

	local_root := path.Dir(local)

	_, err := os.Stat(local_root)

	if os.IsNotExist(err) {
		c.Logger.Info("create %s", local_root)
		os.MkdirAll(local_root, 0755)
	}

	rsp, fetch_err := c.Fetch("GET", remote)

	if fetch_err != nil {
		return fetch_err
	}

	defer rsp.Body.Close()

	contents, read_err := ioutil.ReadAll(rsp.Body)

	if read_err != nil {
		c.Logger.Error("failed to read body for %s, because %v", remote, read_err)
		return read_err
	}

	go func() error {

		write_err := ioutil.WriteFile(local, contents, 0644)

		if write_err != nil {
			c.Logger.Error("Failed to write %s, because %v", local, write_err)

			atomic.AddInt64(&c.Success, -1)
			atomic.AddInt64(&c.Error, 1)

			return write_err
		}

		c.Logger.Debug("Wrote %s to disk", local)
		return nil
	}()

	return nil
}

func (c *WOFClone) Fetch(method string, url string) (*http.Response, error) {

	c.Logger.Debug("%s %s", method, url)

	req, _ := http.NewRequest(method, url, nil)
	req.Close = true

	rsp, err := c.client.Do(req)

	if err != nil {
		c.Logger.Error("Failed to %s %s, because %v", method, url, err)
		// golog.Fatal(err)
		return nil, err
	}

	return rsp, err
}

func (c *WOFClone) Status() {
	c.Logger.Info("scheduled: %d completed: %d success: %d error: %d goroutines: %d", c.Scheduled, c.Completed, c.Success, c.Error, runtime.NumGoroutine())
}
<commit_msg>more stuff inre issue #5<commit_after>package clone

import (
	"crypto/md5"
	"encoding/hex"
	"github.com/jeffail/tunny"
	csv "github.com/whosonfirst/go-whosonfirst-csv"
	log "github.com/whosonfirst/go-whosonfirst-log"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"path"
	"path/filepath"
	"runtime"
	"strings"
	"sync"
	"sync/atomic"
)

type WOFClone struct {
	Source    string
	Dest      string
	Count     int64
	Success   int64
	Error     int64
	Skipped   int64
	Scheduled int64
	Completed int64
	Logger    *log.WOFLogger
	client    *http.Client
	pool      *tunny.WorkPool
}

func NewWOFClone(source string, dest string, procs int, logger *log.WOFLogger) *WOFClone {

	// cd := &sync.Cond{L: &sync.Mutex{}}

	cl := &http.Client{}

	runtime.GOMAXPROCS(procs)

	pool, _ := tunny.CreatePoolGeneric(procs).Open()

	c := WOFClone{
		Count:   0,
		Success: 0,
		Error:   0,
		Skipped: 0,
		Source:  source,
		Dest:    dest,
		Logger:  logger,
		client:  cl,
		pool:    pool,
	}

	return &c
}

func (c *WOFClone) CloneMetaFile(file string) error {

	abs_path, _ := filepath.Abs(file)
	// c.Logger.Debug("Parse meta file %s", abs_path)

	reader, read_err := csv.NewDictReader(abs_path)

	if read_err != nil {
		c.Logger.Error("Failed to read %s, because %v", abs_path, read_err)
		return read_err
	}

	wg := new(sync.WaitGroup)

	for {
		row, err := reader.Read()

		if err == io.EOF {
			break
		}

		if err != nil {
			return err
		}

		rel_path, ok := row["path"]

		if !ok {
			continue
		}

		// this does not account for counts - need to think about who calls what, when
		// probably needs to moved in to ClonePath... (20151105/thisisaaronland)

		ensure_changes := true
		skip_existing := false

		remote := c.Source + rel_path
		local := path.Join(c.Dest, rel_path)

		if !os.IsNotExist(err) {

			if skip_existing {
				c.Logger.Debug("%s already exists and we are skipping things that exist", local)
				continue
			}

			change, _ := c.HasChanged(local, remote)

			if !change {
				c.Logger.Info("no changes to %s", local)
				continue
			}
		}

		wg.Add(1)
		atomic.AddInt64(&c.Scheduled, 1)

		go func() {

			defer wg.Done()

			_, err = c.pool.SendWork(func() {

				cl_err := c.ClonePath(rel_path, ensure_changes)

				if cl_err != nil {
					atomic.AddInt64(&c.Error, 1)
				} else {
					atomic.AddInt64(&c.Success, 1)
				}

				atomic.AddInt64(&c.Completed, 1)
				c.Status()
			})

		}()
	}

	wg.Wait()
	return nil
}

func (c *WOFClone) ClonePath(rel_path string, ensure_changes bool) error {

	atomic.AddInt64(&c.Count, 1)

	remote := c.Source + rel_path
	local := path.Join(c.Dest, rel_path)

	_, err := os.Stat(local)

	if !os.IsNotExist(err) && ensure_changes {

		change, _ := c.HasChanged(local, remote)

		if !change {

			c.Logger.Debug("%s has not changed so skipping", local)
			atomic.AddInt64(&c.Skipped, 1)
			return nil
		}

	}

	process_err := c.Process(remote, local)

	if process_err != nil {
		atomic.AddInt64(&c.Error, 1)
		return process_err
	}

	return nil
}

// don't return true if there's a problem - move that logic up above

func (c *WOFClone) HasChanged(local string, remote string) (bool, error) {

	change := true

	body, err := ioutil.ReadFile(local)

	if err != nil {
		c.Logger.Error("Failed to read %s, becase %v", local, err)
		return change, err
	}

	hash := md5.Sum(body)
	local_hash := hex.EncodeToString(hash[:])

	return c.HasHashChanged(local_hash, remote)
}

func (c *WOFClone) HasHashChanged(local_hash string, remote string) (bool, error) {

	change := true

	rsp, err := c.Fetch("HEAD", remote)

	if err != nil {
		return change, err
	}

	defer rsp.Body.Close()

	etag := rsp.Header.Get("Etag")
	remote_hash := strings.Replace(etag, "\"", "", -1)

	if local_hash == remote_hash {
		change = false
	}

	return change, nil
}

func (c *WOFClone) Process(remote string, local string) error {

	c.Logger.Debug("fetch %s and store in %s", remote, local)

	local_root := path.Dir(local)

	_, err := os.Stat(local_root)

	if os.IsNotExist(err) {
		c.Logger.Info("create %s", local_root)
		os.MkdirAll(local_root, 0755)
	}

	rsp, fetch_err := c.Fetch("GET", remote)

	if fetch_err != nil {
		return fetch_err
	}

	defer rsp.Body.Close()

	contents, read_err := ioutil.ReadAll(rsp.Body)

	if read_err != nil {
		c.Logger.Error("failed to read body for %s, because %v", remote, read_err)
		return read_err
	}

	go func() error {

		write_err := ioutil.WriteFile(local, contents, 0644)

		if write_err != nil {
			c.Logger.Error("Failed to write %s, because %v", local, write_err)

			atomic.AddInt64(&c.Success, -1)
			atomic.AddInt64(&c.Error, 1)

			return write_err
		}

		c.Logger.Debug("Wrote %s to disk", local)
		return nil
	}()

	return nil
}

func (c *WOFClone) Fetch(method string, url string) (*http.Response, error) {

	c.Logger.Debug("%s %s", method, url)

	req, _ := http.NewRequest(method, url, nil)
	req.Close = true

	rsp, err := c.client.Do(req)

	if err != nil {
		c.Logger.Error("Failed to %s %s, because %v", method, url, err)
		// golog.Fatal(err)
		return nil, err
	}

	return rsp, err
}

func (c *WOFClone) Status() {
	c.Logger.Info("scheduled: %d completed: %d success: %d error: %d goroutines: %d", c.Scheduled, c.Completed, c.Success, c.Error, runtime.NumGoroutine())
}
<|endoftext|>


<commit_before>package gorocksdb

import (
	"fmt"
	. "github.com/smartystreets/goconvey/convey"
	"os"
	"testing"
)

type testSliceTransformHandler struct {
	initiated bool
}

func (self *testSliceTransformHandler) Transform(src []byte) []byte {
	fmt.Println("Transform", string(src))
	return src
}

func (self *testSliceTransformHandler) InDomain(src []byte) bool {
	fmt.Println("InDomain", string(src))
	return true
}

func (self *testSliceTransformHandler) InRange(src []byte) bool {
	fmt.Println(src)
	return true
}

func (self *testSliceTransformHandler) Name() string {
	self.initiated = true
	return "gorocksdb.test"
}

func TestNewSliceTransform(t *testing.T) {
	dbName := os.TempDir() + "/TestNewSliceTransform"

	Convey("Subject: Custom slice transform", t, func() {
		Convey("When create a custom slice transform it should not panic", func() {
			handler := &testSliceTransformHandler{}
			sliceTransform := NewSliceTransform(handler)

			Convey("When passed to the db as prefix extractor it should not panic", func() {
				options := NewDefaultOptions()
				DestroyDb(dbName, options)
				options.SetCreateIfMissing(true)
				options.SetPrefixExtractor(sliceTransform)
				options.SetWholeKeyFiltering(true)

				_, err := OpenDb(options, dbName)
				So(err, ShouldBeNil)
				So(handler.initiated, ShouldBeTrue)
			})
		})
	})
}

func TestNewFixedPrefixTransform(t *testing.T) {
	dbName := os.TempDir() + "/TestNewFixedPrefixTransform"

	Convey("Subject: Create fixed prefix transform", t, func() {
		Convey("When create a fixed prefix transform then it should not panic", func() {
			sliceTransform := NewFixedPrefixTransform(3)

			Convey("When passed to the db as prefix extractor then it should not panic", func() {
				options := NewDefaultOptions()
				DestroyDb(dbName, options)
				options.SetCreateIfMissing(true)
				options.SetPrefixExtractor(sliceTransform)
				options.SetWholeKeyFiltering(true)
				options.SetFilterPolicy(NewBloomFilter(8))

				db, err := OpenDb(options, dbName)
				So(err, ShouldBeNil)

				Convey("When add 3 values with key prefix 'foo' and 3 values with key prefix 'bar' then it should not panic", func() {
					wo := NewDefaultWriteOptions()
					So(db.Put(wo, []byte("foo1"), []byte("foo")), ShouldBeNil)
					So(db.Put(wo, []byte("foo2"), []byte("foo")), ShouldBeNil)
					So(db.Put(wo, []byte("foo3"), []byte("foo")), ShouldBeNil)
					So(db.Put(wo, []byte("bar1"), []byte("bar")), ShouldBeNil)
					So(db.Put(wo, []byte("bar2"), []byte("bar")), ShouldBeNil)
					So(db.Put(wo, []byte("bar3"), []byte("bar")), ShouldBeNil)

					Convey("When create an interator and seek to prefix 'bar' then it should return only 3 keys", func() {
						ro := NewDefaultReadOptions()
						ro.SetPrefixSeek(true)
						ro.SetPrefix([]byte("bar"))

						it := db.NewIterator(ro)
						it.Seek([]byte("bar"))
						numFound := 0
						for ; it.Valid(); it.Next() {
							numFound++

						}
						So(it.Err(), ShouldBeNil)
						So(numFound, ShouldEqual, 3)
						it.Close()
					})
				})
			})
		})
	})
}
<commit_msg>[SliceTransform] fixed prefix seek tests<commit_after>package gorocksdb

import (
	. "github.com/smartystreets/goconvey/convey"
	"os"
	"testing"
)

type testSliceTransformHandler struct {
	initiated bool
}

func (self *testSliceTransformHandler) Transform(src []byte) []byte {
	return src[0:3]
}

func (self *testSliceTransformHandler) InDomain(src []byte) bool {
	return len(src) >= 3
}

func (self *testSliceTransformHandler) InRange(src []byte) bool {
	return len(src) == 3
}

func (self *testSliceTransformHandler) Name() string {
	self.initiated = true
	return "gorocksdb.test"
}

func TestCustomSliceTransform(t *testing.T) {
	dbName := os.TempDir() + "/TestNewSliceTransform"

	Convey("Subject: Prefix filtering with custom slice transform", t, func() {
		handler := &testSliceTransformHandler{}
		sliceTransform := NewSliceTransform(handler)

		options := NewDefaultOptions()
		DestroyDb(dbName, options)

		options.SetFilterPolicy(NewBloomFilter(10))
		options.SetPrefixExtractor(sliceTransform)
		options.SetHashSkipListRep(50000, 4, 4)
		options.SetCreateIfMissing(true)

		db, err := OpenDb(options, dbName)
		defer db.Close()

		So(err, ShouldBeNil)

		wo := NewDefaultWriteOptions()
		So(db.Put(wo, []byte("foo1"), []byte("foo")), ShouldBeNil)
		So(db.Put(wo, []byte("foo2"), []byte("foo")), ShouldBeNil)
		So(db.Put(wo, []byte("foo3"), []byte("foo")), ShouldBeNil)
		So(db.Put(wo, []byte("bar1"), []byte("bar")), ShouldBeNil)
		So(db.Put(wo, []byte("bar2"), []byte("bar")), ShouldBeNil)
		So(db.Put(wo, []byte("bar3"), []byte("bar")), ShouldBeNil)

		ro := NewDefaultReadOptions()
		ro.SetPrefixSeek(true)

		it := db.NewIterator(ro)
		defer it.Close()
		numFound := 0
		for it.Seek([]byte("bar")); it.Valid(); it.Next() {
			numFound++
		}

		So(it.Err(), ShouldBeNil)
		So(numFound, ShouldEqual, 3)
	})
}

func TestFixedPrefixTransform(t *testing.T) {
	dbName := os.TempDir() + "/TestNewFixedPrefixTransform"

	Convey("Subject: Prefix filtering with native fixed prefix transform", t, func() {
		options := NewDefaultOptions()
		DestroyDb(dbName, options)

		options.SetFilterPolicy(NewBloomFilter(10))
		options.SetPrefixExtractor(NewFixedPrefixTransform(3))
		options.SetHashSkipListRep(50000, 4, 4)
		options.SetCreateIfMissing(true)

		db, err := OpenDb(options, dbName)
		defer db.Close()

		So(err, ShouldBeNil)

		wo := NewDefaultWriteOptions()
		So(db.Put(wo, []byte("foo1"), []byte("foo")), ShouldBeNil)
		So(db.Put(wo, []byte("foo2"), []byte("foo")), ShouldBeNil)
		So(db.Put(wo, []byte("foo3"), []byte("foo")), ShouldBeNil)
		So(db.Put(wo, []byte("bar1"), []byte("bar")), ShouldBeNil)
		So(db.Put(wo, []byte("bar2"), []byte("bar")), ShouldBeNil)
		So(db.Put(wo, []byte("bar3"), []byte("bar")), ShouldBeNil)

		ro := NewDefaultReadOptions()
		ro.SetPrefixSeek(true)

		it := db.NewIterator(ro)
		defer it.Close()
		numFound := 0
		for it.Seek([]byte("bar")); it.Valid(); it.Next() {
			numFound++
		}

		So(it.Err(), ShouldBeNil)
		So(numFound, ShouldEqual, 3)
	})
}
<|endoftext|>


<commit_before>package rpcd

import (
	"encoding/gob"
	"errors"
	"github.com/Symantec/Dominator/lib/srpc"
	"github.com/Symantec/Dominator/proto/imageserver"
	"github.com/proxypoke/group.go"
)

func (t *srpcType) ChownDirectory(conn *srpc.Conn) error {
	var request imageserver.ChangeOwnerRequest
	var response imageserver.ChangeOwnerResponse
	decoder := gob.NewDecoder(conn)
	if err := decoder.Decode(&request); err != nil {
		_, err = conn.WriteString(err.Error() + "\n")
		return err
	}
	if err := t.changeOwner(request, &response, conn.Username()); err != nil {
		_, err = conn.WriteString(err.Error() + "\n")
		return err
	}
	if _, err := conn.WriteString("\n"); err != nil {
		return err
	}
	return gob.NewEncoder(conn).Encode(response)
}

func (t *srpcType) changeOwner(request imageserver.ChangeOwnerRequest,
	reply *imageserver.ChangeOwnerResponse, username string) error {
	if username == "" {
		return errors.New("no username: unauthenticated connection")
	}
	if _, err := group.Lookup(request.OwnerGroup); err != nil {
		return err
	}
	t.logger.Printf("ChownDirectory(%s) to: %s by %s\n",
		request.DirectoryName, request.OwnerGroup, username)
	return t.imageDataBase.ChownDirectory(request.DirectoryName,
		request.OwnerGroup)
}
<commit_msg>Support empty OwnerGroup in ImageServer.ChownDirectory() SRPC method.<commit_after>package rpcd

import (
	"encoding/gob"
	"errors"
	"github.com/Symantec/Dominator/lib/srpc"
	"github.com/Symantec/Dominator/proto/imageserver"
	"github.com/proxypoke/group.go"
)

func (t *srpcType) ChownDirectory(conn *srpc.Conn) error {
	var request imageserver.ChangeOwnerRequest
	var response imageserver.ChangeOwnerResponse
	decoder := gob.NewDecoder(conn)
	if err := decoder.Decode(&request); err != nil {
		_, err = conn.WriteString(err.Error() + "\n")
		return err
	}
	if err := t.changeOwner(request, &response, conn.Username()); err != nil {
		_, err = conn.WriteString(err.Error() + "\n")
		return err
	}
	if _, err := conn.WriteString("\n"); err != nil {
		return err
	}
	return gob.NewEncoder(conn).Encode(response)
}

func (t *srpcType) changeOwner(request imageserver.ChangeOwnerRequest,
	reply *imageserver.ChangeOwnerResponse, username string) error {
	if username == "" {
		return errors.New("no username: unauthenticated connection")
	}
	if request.OwnerGroup != "" {
		if _, err := group.Lookup(request.OwnerGroup); err != nil {
			return err
		}
	}
	t.logger.Printf("ChownDirectory(%s) to: \"%s\" by %s\n",
		request.DirectoryName, request.OwnerGroup, username)
	return t.imageDataBase.ChownDirectory(request.DirectoryName,
		request.OwnerGroup)
}
<|endoftext|>


<commit_before>// Tideland Go CouchDB Client - Find - Unit Tests
//
// Copyright (C) 2016-2017 Frank Mueller / Tideland / Oldenburg / Germany
//
// All rights reserved. Use of this source code is governed
// by the new BSD license.

package find_test

import (
	"strings"
	"testing"
	"time"

	"github.com/tideland/golib/audit"
	"github.com/tideland/golib/etc"
	"github.com/tideland/golib/identifier"
	"github.com/tideland/golib/logger"

	"github.com/tideland/gocouch/couchdb"
	"github.com/tideland/gocouch/find"
)

//--------------------
// IMPORTS
//--------------------

//--------------------
// CONSTANTS
//--------------------

const (
	Cfg = "{etc {hostname localhost}{port 5984}{database tgocouch-testing-<<DATABASE>>}{debug-logging true}}"
)

//--------------------
// TESTS
//--------------------

// TestSimpleFind tests calling find with a simple selector.
func TestSimpleFind(t *testing.T) {
	assert := audit.NewTestingAssertion(t, true)
	cdb, cleanup := prepareFilledDatabase("find-simple", 1000, assert)
	defer cleanup()

	// Try to find some documents a simple way.
	selector := find.Select(find.Or(
		find.And(
			find.LowerThan("age", 30),
			find.Equal("active", false),
		),
		find.And(
			find.GreaterThan("age", 60),
			find.Equal("active", "true"),
		),
	))
	frs := find.Find(cdb, selector, find.Fields("name", "age", "active"))
	assert.Nil(frs.Error())
	assert.True(frs.IsOK())

	err := frs.Do(func(document couchdb.Unmarshable) error {
		fields := struct {
			Name   string `json:"name"`
			Age    int    `json:"age"`
			Active bool   `json:"active"`
		}{}
		if err := document.Unmarshal(&fields); err != nil {
			return err
		}
		assert.True((fields.Age < 30 && !fields.Active) || (fields.Age > 60 && fields.Active))
		assert.Logf("person with name %s has age %d and activity status is %v", fields.Name, fields.Age, fields.Active)
		return nil
	})
	assert.Nil(err)
}

// TestLimitedFind tests retrieving a larger number but set the limit.
func TestLimitedFind(t *testing.T) {
	assert := audit.NewTestingAssertion(t, true)
	cdb, cleanup := prepareFilledDatabase("find-limited", 1000, assert)
	defer cleanup()

	// Limit found documents.
	selector := find.Select(find.Equal("active", true))
	frs := find.Find(cdb, selector, find.Fields("name", "active"), find.Limit(5))
	assert.NotNil(frs)
	assert.True(frs.IsOK())
	assert.Length(frs, 5)

	// Greater limit.
	frs = find.Find(cdb, selector, find.Fields("name", "active"), find.Limit(100))
	assert.NotNil(frs)
	assert.True(frs.IsOK())
	assert.Length(frs, 100)
}

// TestSortedFind tests retrieving a larger number in a sorted way.
func TestSortedFind(t *testing.T) {
	assert := audit.NewTestingAssertion(t, true)
	cdb, cleanup := prepareFilledDatabase("find-sorted", 1000, assert)
	defer cleanup()

	// Sorting field has to be part of selector.
	selector := find.Select(
		find.GreaterThan("name", nil),
		find.Equal("active", true),
	)
	frs := find.Find(cdb, selector, find.Fields("name", "age"), find.Sort(find.Ascending("name")), find.Limit(1000))
	assert.NotNil(frs)
	assert.Equal(frs.Error(), nil)
	assert.True(frs.IsOK())

	name := ""
	err := frs.Do(func(document couchdb.Unmarshable) error {
		fields := struct {
			Name string `json:"name"`
			Age  int    `json:"age"`
		}{}
		if err := document.Unmarshal(&fields); err != nil {
			return err
		}
		assert.Logf("person with name %s has age %d", fields.Name, fields.Age)
		assert.True(fields.Name >= name)
		name = fields.Name
		return nil
	})
	assert.Nil(err)
}

// TestFindExists tests calling find with an exists selector.
func TestFindExists(t *testing.T) {
	assert := audit.NewTestingAssertion(t, true)
	cdb, cleanup := prepareFilledDatabase("find-exists", 1000, assert)
	defer cleanup()

	// Try to find some documents having an existing "last_active".
	selector := find.Select(
		find.Exists("last_active"),
		find.LowerEqualThan("age", 25),
	)
	frs := find.Find(cdb, selector, find.Fields("name", "age", "active", "last_active"))
	assert.NotNil(frs)
	assert.True(frs.IsOK())

	err := frs.Do(func(document couchdb.Unmarshable) error {
		fields := struct {
			Name       string `json:"name"`
			Age        int    `json:"age"`
			Active     bool   `json:"active"`
			LastActive int64  `json:"last_active"`
		}{}
		if err := document.Unmarshal(&fields); err != nil {
			return err
		}
		assert.True(fields.Age <= 25 && fields.LastActive > 0 && fields.Active)
		lastActive := time.Unix(fields.LastActive, 0).Format(time.RFC1123)
		assert.Logf("person with name %s (age %d) has been last active at %v", fields.Name, fields.Age, lastActive)
		return nil
	})
	assert.Nil(err)

	// Now look for existing "last_active" but "active" is false. So
	// no results.
	selector = find.Select(
		find.Exists("last_active"),
		find.NotEqual("active", true),
	)
	frs = find.Find(cdb, selector, find.Fields("name", "age", "active", "last_active"))
	assert.NotNil(frs)
	assert.True(frs.IsOK())
	assert.Equal(frs.Len(), 0)
}

// TestOneCriterion tests using only one criterion, here a regular expression.
func TestOneCriterion(t *testing.T) {
	assert := audit.NewTestingAssertion(t, true)
	cdb, cleanup := prepareFilledDatabase("find-one-criterion", 1000, assert)
	defer cleanup()

	// Try to find some documents having an existing "last_active".
	selector := find.Select(find.RegExp("name", ".*Adam.*"))
	frs := find.Find(cdb, selector, find.Fields("name", "age", "active"))
	assert.NotNil(frs)
	assert.Nil(frs.Error())
	assert.True(frs.IsOK())

	err := frs.Do(func(document couchdb.Unmarshable) error {
		fields := struct {
			Name   string `json:"name"`
			Age    int    `json:"age"`
			Active bool   `json:"active"`
		}{}
		if err := document.Unmarshal(&fields); err != nil {
			return err
		}
		assert.Match(fields.Name, ".*Adam.*")
		return nil
	})
	assert.Nil(err)
}

// TestMatches tests using element and all match operators.
func TestMatch(t *testing.T) {
	assert := audit.NewTestingAssertion(t, true)
	cdb, cleanup := prepareFilledDatabase("find-match", 1000, assert)
	defer cleanup()

	assert.NotNil(cdb)
}

//--------------------
// HELPERS
//--------------------

// Note is used for the tests.
type Note struct {
	Title string `json:"title"`
	Text  string `json:"text"`
}

// Person is used for the tests.
type Person struct {
	DocumentID       string `json:"_id,omitempty"`
	DocumentRevision string `json:"_rev,omitempty"`

	Name       string `json:"name"`
	Age        int    `json:"age"`
	Shifts     []int  `json:"shifts"`
	Active     bool   `json:"active"`
	LastActive int64  `json:"last_active,omitempty"`
	Notes      []Note `json:"notes"`
}

// prepareFilledDatabase opens the database, deletes a possible test
// database, creates it newly and adds some data.
func prepareFilledDatabase(database string, count int, assert audit.Assertion) (couchdb.CouchDB, func()) {
	logger.SetLevel(logger.LevelDebug)
	cfgstr := strings.Replace(Cfg, "<<DATABASE>>", database, 1)
	cfg, err := etc.ReadString(cfgstr)
	assert.Nil(err)
	cdb, err := couchdb.Open(cfg)
	assert.Nil(err)
	rs := cdb.DeleteDatabase()
	rs = cdb.CreateDatabase()
	assert.True(rs.IsOK())
	err = find.CreateIndex(cdb, find.NewIndex("name"))
	assert.Nil(err)

	gen := audit.NewGenerator(audit.FixedRand())
	runs := count / 1000
	for outer := 0; outer < runs; outer++ {
		assert.Logf("filling database run %d of %d", outer+1, runs)
		docs := []interface{}{}
		for inner := 0; inner < 1000; inner++ {
			first, middle, last := gen.Name()
			person := Person{
				DocumentID: identifier.Identifier(last, first, outer, inner),
				Name:       first + " " + middle + " " + last,
				Age:        gen.Int(18, 65),
				Shifts:     []int{gen.Int(1, 3), gen.Int(1, 3), gen.Int(1, 3)},
				Active:     gen.FlipCoin(75),
			}
			if person.Active {
				person.LastActive = gen.Time(time.UTC, time.Now().Add(-24*time.Hour), 24*time.Hour).Unix()
			}
			for j := 0; j < gen.Int(3, 9); j++ {
				note := Note{
					Title: gen.Sentence(),
					Text:  gen.Paragraph(),
				}
				person.Notes = append(person.Notes, note)
			}
			docs = append(docs, person)
		}
		results, err := cdb.BulkWriteDocuments(docs)
		assert.Nil(err)
		for _, result := range results {
			assert.True(result.OK)
		}
	}

	return cdb, func() { cdb.DeleteDatabase() }
}

// EOF
<commit_msg>Added test for matching find operators<commit_after>// Tideland Go CouchDB Client - Find - Unit Tests
//
// Copyright (C) 2016-2017 Frank Mueller / Tideland / Oldenburg / Germany
//
// All rights reserved. Use of this source code is governed
// by the new BSD license.

package find_test

import (
	"strings"
	"testing"
	"time"

	"github.com/tideland/golib/audit"
	"github.com/tideland/golib/etc"
	"github.com/tideland/golib/identifier"
	"github.com/tideland/golib/logger"

	"github.com/tideland/gocouch/couchdb"
	"github.com/tideland/gocouch/find"
)

//--------------------
// IMPORTS
//--------------------

//--------------------
// CONSTANTS
//--------------------

const (
	Cfg = "{etc {hostname localhost}{port 5984}{database tgocouch-testing-<<DATABASE>>}{debug-logging true}}"
)

//--------------------
// TESTS
//--------------------

// TestSimpleFind tests calling find with a simple selector.
func TestSimpleFind(t *testing.T) {
	assert := audit.NewTestingAssertion(t, true)
	cdb, cleanup := prepareFilledDatabase("find-simple", 1000, assert)
	defer cleanup()

	// Try to find some documents a simple way.
	selector := find.Select(find.Or(
		find.And(
			find.LowerThan("age", 30),
			find.Equal("active", false),
		),
		find.And(
			find.GreaterThan("age", 60),
			find.Equal("active", "true"),
		),
	))
	frs := find.Find(cdb, selector, find.Fields("name", "age", "active"))
	assert.Nil(frs.Error())
	assert.True(frs.IsOK())

	err := frs.Do(func(document couchdb.Unmarshable) error {
		fields := struct {
			Name   string `json:"name"`
			Age    int    `json:"age"`
			Active bool   `json:"active"`
		}{}
		if err := document.Unmarshal(&fields); err != nil {
			return err
		}
		assert.True((fields.Age < 30 && !fields.Active) || (fields.Age > 60 && fields.Active))
		assert.Logf("person with name %s has age %d and activity status is %v", fields.Name, fields.Age, fields.Active)
		return nil
	})
	assert.Nil(err)
}

// TestLimitedFind tests retrieving a larger number but set the limit.
func TestLimitedFind(t *testing.T) {
	assert := audit.NewTestingAssertion(t, true)
	cdb, cleanup := prepareFilledDatabase("find-limited", 1000, assert)
	defer cleanup()

	// Limit found documents.
	selector := find.Select(find.Equal("active", true))
	frs := find.Find(cdb, selector, find.Fields("name", "active"), find.Limit(5))
	assert.NotNil(frs)
	assert.True(frs.IsOK())
	assert.Length(frs, 5)

	// Greater limit.
	frs = find.Find(cdb, selector, find.Fields("name", "active"), find.Limit(100))
	assert.NotNil(frs)
	assert.True(frs.IsOK())
	assert.Length(frs, 100)
}

// TestSortedFind tests retrieving a larger number in a sorted way.
func TestSortedFind(t *testing.T) {
	assert := audit.NewTestingAssertion(t, true)
	cdb, cleanup := prepareFilledDatabase("find-sorted", 1000, assert)
	defer cleanup()

	// Sorting field has to be part of selector.
	selector := find.Select(
		find.GreaterThan("name", nil),
		find.Equal("active", true),
	)
	frs := find.Find(cdb, selector, find.Fields("name", "age"), find.Sort(find.Ascending("name")), find.Limit(1000))
	assert.NotNil(frs)
	assert.Equal(frs.Error(), nil)
	assert.True(frs.IsOK())

	name := ""
	err := frs.Do(func(document couchdb.Unmarshable) error {
		fields := struct {
			Name string `json:"name"`
			Age  int    `json:"age"`
		}{}
		if err := document.Unmarshal(&fields); err != nil {
			return err
		}
		assert.Logf("person with name %s has age %d", fields.Name, fields.Age)
		assert.True(fields.Name >= name)
		name = fields.Name
		return nil
	})
	assert.Nil(err)
}

// TestFindExists tests calling find with an exists selector.
func TestFindExists(t *testing.T) {
	assert := audit.NewTestingAssertion(t, true)
	cdb, cleanup := prepareFilledDatabase("find-exists", 1000, assert)
	defer cleanup()

	// Try to find some documents having an existing "last_active".
	selector := find.Select(
		find.Exists("last_active"),
		find.LowerEqualThan("age", 25),
	)
	frs := find.Find(cdb, selector, find.Fields("name", "age", "active", "last_active"))
	assert.NotNil(frs)
	assert.True(frs.IsOK())

	err := frs.Do(func(document couchdb.Unmarshable) error {
		fields := struct {
			Name       string `json:"name"`
			Age        int    `json:"age"`
			Active     bool   `json:"active"`
			LastActive int64  `json:"last_active"`
		}{}
		if err := document.Unmarshal(&fields); err != nil {
			return err
		}
		assert.True(fields.Age <= 25 && fields.LastActive > 0 && fields.Active)
		lastActive := time.Unix(fields.LastActive, 0).Format(time.RFC1123)
		assert.Logf("person with name %s (age %d) has been last active at %v", fields.Name, fields.Age, lastActive)
		return nil
	})
	assert.Nil(err)

	// Now look for existing "last_active" but "active" is false. So
	// no results.
	selector = find.Select(
		find.Exists("last_active"),
		find.NotEqual("active", true),
	)
	frs = find.Find(cdb, selector, find.Fields("name", "age", "active", "last_active"))
	assert.NotNil(frs)
	assert.True(frs.IsOK())
	assert.Equal(frs.Len(), 0)
}

// TestOneCriterion tests using only one criterion, here a regular expression.
func TestOneCriterion(t *testing.T) {
	assert := audit.NewTestingAssertion(t, true)
	cdb, cleanup := prepareFilledDatabase("find-one-criterion", 1000, assert)
	defer cleanup()

	// Try to find some documents having an existing "last_active".
	selector := find.Select(find.RegExp("name", ".*Adam.*"))
	frs := find.Find(cdb, selector, find.Fields("name", "age", "active"))
	assert.NotNil(frs)
	assert.True(frs.IsOK())

	err := frs.Do(func(document couchdb.Unmarshable) error {
		fields := struct {
			Name   string `json:"name"`
			Age    int    `json:"age"`
			Active bool   `json:"active"`
		}{}
		if err := document.Unmarshal(&fields); err != nil {
			return err
		}
		assert.Match(fields.Name, ".*Adam.*")
		return nil
	})
	assert.Nil(err)
}

// TestMatches tests using element and all match operators.
func TestMatches(t *testing.T) {
	assert := audit.NewTestingAssertion(t, true)
	cdb, cleanup := prepareFilledDatabase("find-match", 1000, assert)
	defer cleanup()

	// Find with at least one matching element.
	selector := find.Select(find.MatchElement("shifts", find.Equal("", 3)))
	frs := find.Find(cdb, selector, find.Fields("name", "shifts"))
	assert.NotNil(frs)
	assert.True(frs.IsOK())

	err := frs.Do(func(document couchdb.Unmarshable) error {
		fields := struct {
			Name   string `json:"name"`
			Shifts []int  `json:"shifts"`
		}{}
		if err := document.Unmarshal(&fields); err != nil {
			return err
		}
		assert.Logf("NAME %v SHIFTS %v", fields.Name, fields.Shifts)
		assert.Contents(3, fields.Shifts)
		return nil
	})
	assert.Nil(err)

	// Find with all matching elements (dumb query,
	// but checking combination).
	selector = find.Select(
		find.MatchAll("shifts",
			find.GreaterThan("", 1),
			find.LowerThan("", 3),
		),
	)
	frs = find.Find(cdb, selector, find.Fields("name", "shifts"))
	assert.NotNil(frs)
	assert.Nil(frs.Error())
	assert.True(frs.IsOK())

	err = frs.Do(func(document couchdb.Unmarshable) error {
		fields := struct {
			Name   string `json:"name"`
			Shifts []int  `json:"shifts"`
		}{}
		if err := document.Unmarshal(&fields); err != nil {
			return err
		}
		assert.Logf("NAME %v SHIFTS %v", fields.Name, fields.Shifts)
		assert.Equal(fields.Shifts, []int{2, 2, 2})
		return nil
	})
	assert.Nil(err)
}

//--------------------
// HELPERS
//--------------------

// Note is used for the tests.
type Note struct {
	Title string `json:"title"`
	Text  string `json:"text"`
}

// Person is used for the tests.
type Person struct {
	DocumentID       string `json:"_id,omitempty"`
	DocumentRevision string `json:"_rev,omitempty"`

	Name       string `json:"name"`
	Age        int    `json:"age"`
	Shifts     []int  `json:"shifts"`
	Active     bool   `json:"active"`
	LastActive int64  `json:"last_active,omitempty"`
	Notes      []Note `json:"notes"`
}

// prepareFilledDatabase opens the database, deletes a possible test
// database, creates it newly and adds some data.
func prepareFilledDatabase(database string, count int, assert audit.Assertion) (couchdb.CouchDB, func()) {
	logger.SetLevel(logger.LevelDebug)
	cfgstr := strings.Replace(Cfg, "<<DATABASE>>", database, 1)
	cfg, err := etc.ReadString(cfgstr)
	assert.Nil(err)
	cdb, err := couchdb.Open(cfg)
	assert.Nil(err)
	rs := cdb.DeleteDatabase()
	rs = cdb.CreateDatabase()
	assert.True(rs.IsOK())
	err = find.CreateIndex(cdb, find.NewIndex("name"))
	assert.Nil(err)

	gen := audit.NewGenerator(audit.FixedRand())
	runs := count / 1000
	for outer := 0; outer < runs; outer++ {
		assert.Logf("filling database run %d of %d", outer+1, runs)
		docs := []interface{}{}
		for inner := 0; inner < 1000; inner++ {
			first, middle, last := gen.Name()
			person := Person{
				DocumentID: identifier.Identifier(last, first, outer, inner),
				Name:       first + " " + middle + " " + last,
				Age:        gen.Int(18, 65),
				Shifts:     []int{gen.Int(1, 3), gen.Int(1, 3), gen.Int(1, 3)},
				Active:     gen.FlipCoin(75),
			}
			if person.Active {
				person.LastActive = gen.Time(time.UTC, time.Now().Add(-24*time.Hour), 24*time.Hour).Unix()
			}
			for j := 0; j < gen.Int(3, 9); j++ {
				note := Note{
					Title: gen.Sentence(),
					Text:  gen.Paragraph(),
				}
				person.Notes = append(person.Notes, note)
			}
			docs = append(docs, person)
		}
		results, err := cdb.BulkWriteDocuments(docs)
		assert.Nil(err)
		for _, result := range results {
			assert.True(result.OK)
		}
	}

	return cdb, func() { cdb.DeleteDatabase() }
}

// EOF
<|endoftext|>


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

package concrete_test

import (
	"testing"

	"github.com/gonum/graph"
	"github.com/gonum/graph/concrete"
)

var _ graph.Graph = &concrete.Graph{}
var _ graph.Graph = &concrete.Graph{}

func TestAssertMutableNotDirected(t *testing.T) {
	graph.VetMutableGraph(concrete.NewMutableGraph())
}

// var _ gr.EdgeListGraph = &concrete.Graph{}
<commit_msg>Fixed mutability assertion<commit_after>// Copyright ©2014 The gonum Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package concrete_test

import (
	"testing"

	"github.com/gonum/graph"
	"github.com/gonum/graph/concrete"
)

var _ graph.Graph = &concrete.Graph{}
var _ graph.Graph = &concrete.Graph{}

func TestAssertMutableNotDirected(t *testing.T) {
	var g graph.MutableGraph = concrete.NewGraph()
	if _, ok := g.(graph.DirectedGraph); ok {
		t.Fatal("concrete.Graph is directed, but a MutableGraph cannot safely be directed!")
	}
}

// var _ gr.EdgeListGraph = &concrete.Graph{}
<|endoftext|>


<commit_before>package helpers

import (
	"github.com/globalsign/mgo"
	"github.com/globalsign/mgo/bson"

	"reflect"

	"strings"

	"crypto/tls"

	"net"

	"fmt"
	"time"

	"github.com/Seklfreak/Robyul2/cache"
	"github.com/Seklfreak/Robyul2/models"
	"github.com/pkg/errors"
)

var (
	mDbSession  *mgo.Session
	mDbDatabase string
)

type mgoLogger struct {
}

func (mgol mgoLogger) Output(calldepth int, s string) error {
	cache.GetLogger().WithField("module", "mdb").Info(s)
	return nil
}

// ConnectDB connects to mongodb and stores the session
func ConnectMDB(url string, database string) {
	var err error

	log := cache.GetLogger()
	log.WithField("module", "mdb").Info("Connecting to " + url)

	mgoL := new(mgoLogger)
	mgo.SetLogger(mgoL)

	cache.GetLogger()

	newUrl := strings.TrimSuffix(url, "?ssl=true")
	newUrl = strings.Replace(newUrl, "ssl=true&", "", -1)

	dialInfo, err := mgo.ParseURL(newUrl)
	if err != nil {
		log.WithField("module", "mdb").Error(err.Error())
		panic(err)
	}

	// setup TLS if we use SSL
	if newUrl != url {
		tlsConfig := &tls.Config{}
		tlsConfig.InsecureSkipVerify = true

		dialInfo.DialServer = func(addr *mgo.ServerAddr) (net.Conn, error) {
			conn, err := tls.Dial("tcp", addr.String(), tlsConfig)
			return conn, err
		}
	}

	mDbSession, err = mgo.DialWithInfo(dialInfo)
	if err != nil {
		log.WithField("module", "mdb").Error(err.Error())
		panic(err)
	}

	mDbSession.SetMode(mgo.Monotonic, true)
	mDbSession.SetSafe(&mgo.Safe{WMode: "majority"})

	mDbDatabase = database

	log.WithField("module", "mdb").Info("Connected!")
}

// GetDB is a simple getter for the mongodb database.
func GetMDb() *mgo.Database {
	return mDbSession.DB(mDbDatabase)
}

// GetDB is a simple getter for the mongodb session.
func GetMDbSession() *mgo.Session {
	return mDbSession
}

func MDbInsert(collection models.MongoDbCollection, data interface{}) (rid bson.ObjectId, err error) {
	ptr := reflect.New(reflect.TypeOf(data))
	temp := ptr.Elem()
	temp.Set(reflect.ValueOf(data))

	v := temp.FieldByName("ID")

	if !v.IsValid() {
		return bson.ObjectId(""), errors.New("invalid data")
	}

	newID := v.String()
	if newID == "" {
		newID = string(bson.NewObjectId())
		v.SetString(newID)
	}

	start := time.Now()
	err = GetMDb().C(collection.String()).Insert(temp.Interface())
	took := time.Since(start)

	if cache.HasKeen() {
		go func() {
			defer Recover()

			err := cache.GetKeen().AddEvent("Robyul_MongoDB", &KeenMongoDbEvent{
				Seconds:    took.Seconds(),
				Type:       "insert",
				Method:     "MDbInsert()",
				Collection: collection.String(),
				Data:       fmt.Sprintf("%+v", data),
			})
			if err != nil {
				cache.GetLogger().WithField("module", "mdb").Error("Error logging MongoDB request to keen: ", err.Error())
			}
		}()
	}

	if err != nil {
		return bson.ObjectId(""), err
	}

	return bson.ObjectId(newID), nil
}

func MDbInsertWithoutLogging(collection models.MongoDbCollection, data interface{}) (rid bson.ObjectId, err error) {
	ptr := reflect.New(reflect.TypeOf(data))
	temp := ptr.Elem()
	temp.Set(reflect.ValueOf(data))

	v := temp.FieldByName("ID")

	if !v.IsValid() {
		return bson.ObjectId(""), errors.New("invalid data")
	}

	newID := v.String()
	if newID == "" {
		newID = string(bson.NewObjectId())
		v.SetString(newID)
	}

	err = GetMDb().C(collection.String()).Insert(temp.Interface())

	if err != nil {
		return bson.ObjectId(""), err
	}

	return bson.ObjectId(newID), nil
}

func MDbUpdate(collection models.MongoDbCollection, id bson.ObjectId, data interface{}) (rid bson.ObjectId, err error) {
	if !id.Valid() {
		return bson.ObjectId(""), errors.New("invalid id")
	}

	start := time.Now()
	err = GetMDb().C(collection.String()).UpdateId(id, data)
	took := time.Since(start)

	if cache.HasKeen() {
		go func() {
			defer Recover()

			err := cache.GetKeen().AddEvent("Robyul_MongoDB", &KeenMongoDbEvent{
				Seconds:    took.Seconds(),
				Type:       "update",
				Method:     "MDbUpdate()",
				Collection: collection.String(),
				Id:         id.String(),
				Data:       fmt.Sprintf("%+v", data),
			})
			if err != nil {
				cache.GetLogger().WithField("module", "mdb").Error("Error logging MongoDB request to keen: ", err.Error())
			}
		}()
	}

	if err != nil {
		return bson.ObjectId(""), err
	}

	return id, nil
}

func MDbUpdateWithoutLogging(collection models.MongoDbCollection, id bson.ObjectId, data interface{}) (rid bson.ObjectId, err error) {
	if !id.Valid() {
		return bson.ObjectId(""), errors.New("invalid id")
	}

	err = GetMDb().C(collection.String()).UpdateId(id, data)

	if err != nil {
		return bson.ObjectId(""), err
	}

	return id, nil
}

func MDbUpsertID(collection models.MongoDbCollection, id bson.ObjectId, data interface{}) (rid bson.ObjectId, err error) {
	if !id.Valid() {
		id = bson.NewObjectId()
	}

	start := time.Now()
	_, err = GetMDb().C(collection.String()).UpsertId(id, data)
	took := time.Since(start)

	if cache.HasKeen() {
		go func() {
			defer Recover()

			err := cache.GetKeen().AddEvent("Robyul_MongoDB", &KeenMongoDbEvent{
				Seconds:    took.Seconds(),
				Type:       "upsert",
				Method:     "MDbUpsertID()",
				Collection: collection.String(),
				Id:         id.String(),
				Data:       fmt.Sprintf("%+v", data),
			})
			if err != nil {
				cache.GetLogger().WithField("module", "mdb").Error("Error logging MongoDB request to keen: ", err.Error())
			}
		}()
	}

	if err != nil {
		return bson.ObjectId(""), err
	}

	return id, nil
}

func MDbUpsert(collection models.MongoDbCollection, selector interface{}, data interface{}) (err error) {
	start := time.Now()
	_, err = GetMDb().C(collection.String()).Upsert(selector, data)
	took := time.Since(start)

	if cache.HasKeen() {
		go func() {
			defer Recover()

			err := cache.GetKeen().AddEvent("Robyul_MongoDB", &KeenMongoDbEvent{
				Seconds:    took.Seconds(),
				Type:       "upsert",
				Method:     "MDbUpsert()",
				Collection: collection.String(),
				Query:      fmt.Sprintf("%+v", selector),
				Data:       fmt.Sprintf("%+v", data),
			})
			if err != nil {
				cache.GetLogger().WithField("module", "mdb").Error("Error logging MongoDB request to keen: ", err.Error())
			}
		}()
	}

	return err
}

func MDbDelete(collection models.MongoDbCollection, id bson.ObjectId) (err error) {
	if !id.Valid() {
		return errors.New("invalid id")
	}

	start := time.Now()
	err = GetMDb().C(collection.String()).RemoveId(id)
	took := time.Since(start)

	if cache.HasKeen() {
		go func() {
			defer Recover()

			err := cache.GetKeen().AddEvent("Robyul_MongoDB", &KeenMongoDbEvent{
				Seconds:    took.Seconds(),
				Type:       "remove",
				Method:     "MDbDelete()",
				Collection: collection.String(),
				Id:         id.String(),
			})
			if err != nil {
				cache.GetLogger().WithField("module", "mdb").Error("Error logging MongoDB request to keen: ", err.Error())
			}
		}()
	}

	if err != nil {
		return err
	}

	return nil
}

func MdbCollection(collection models.MongoDbCollection) (query *mgo.Collection) {
	return GetMDb().C(collection.String())
}

func MDbIter(query *mgo.Query) (iter *mgo.Iter) {
	start := time.Now()
	iter = query.Iter()
	took := time.Since(start)
	if cache.HasKeen() {
		go func() {
			defer Recover()

			queryValue := reflect.ValueOf(*query)
			queryOp := queryValue.FieldByName("query").FieldByName("op")

			err := cache.GetKeen().AddEvent("Robyul_MongoDB", &KeenMongoDbEvent{
				Seconds:    took.Seconds(),
				Type:       "query",
				Method:     "MdbIter()",
				Collection: queryOp.FieldByName("collection").String(),
				Query:      fmt.Sprintf("%+v", reflect.ValueOf(queryOp.FieldByName("query")).Interface()),
				Skip:       queryOp.FieldByName("skip").Int(),
				Limit:      queryOp.FieldByName("limit").Int(),
			})
			if err != nil {
				cache.GetLogger().WithField("module", "mdb").Error("Error logging MongoDB request to keen: ", err.Error())
			}
		}()
	}
	return
}

func MdbOne(query *mgo.Query, object interface{}) (err error) {
	start := time.Now()
	err = query.One(object)
	took := time.Since(start)
	if cache.HasKeen() {
		go func() {
			defer Recover()

			queryValue := reflect.ValueOf(*query)
			queryOp := queryValue.FieldByName("query").FieldByName("op")

			err := cache.GetKeen().AddEvent("Robyul_MongoDB", &KeenMongoDbEvent{
				Seconds:    took.Seconds(),
				Type:       "query",
				Method:     "MdbOne()",
				Collection: queryOp.FieldByName("collection").String(),
				Query:      fmt.Sprintf("%+v", reflect.ValueOf(queryOp.FieldByName("query")).Interface()),
				Skip:       queryOp.FieldByName("skip").Int(),
				Limit:      queryOp.FieldByName("limit").Int(),
			})
			if err != nil {
				cache.GetLogger().WithField("module", "mdb").Error("Error logging MongoDB request to keen: ", err.Error())
			}
		}()
	}
	return
}

func MdbOneWithoutLogging(query *mgo.Query, object interface{}) (err error) {
	err = query.One(object)
	return
}

type KeenMongoDbEvent struct {
	Seconds    float64
	Collection string
	Type       string
	Method     string
	Query      string `json:",omitempty"`
	Skip       int64  `json:",omitempty"`
	Limit      int64  `json:",omitempty"`
	Id         string `json:",omitempty"`
	Data       string `json:",omitempty"`
}
<commit_msg>[mongodb] normalizes collection name before sending it to keen.io<commit_after>package helpers

import (
	"github.com/globalsign/mgo"
	"github.com/globalsign/mgo/bson"

	"reflect"

	"strings"

	"crypto/tls"

	"net"

	"fmt"
	"time"

	"github.com/Seklfreak/Robyul2/cache"
	"github.com/Seklfreak/Robyul2/models"
	"github.com/pkg/errors"
)

var (
	mDbSession  *mgo.Session
	mDbDatabase string
)

type mgoLogger struct {
}

func (mgol mgoLogger) Output(calldepth int, s string) error {
	cache.GetLogger().WithField("module", "mdb").Info(s)
	return nil
}

// ConnectDB connects to mongodb and stores the session
func ConnectMDB(url string, database string) {
	var err error

	log := cache.GetLogger()
	log.WithField("module", "mdb").Info("Connecting to " + url)

	mgoL := new(mgoLogger)
	mgo.SetLogger(mgoL)

	cache.GetLogger()

	newUrl := strings.TrimSuffix(url, "?ssl=true")
	newUrl = strings.Replace(newUrl, "ssl=true&", "", -1)

	dialInfo, err := mgo.ParseURL(newUrl)
	if err != nil {
		log.WithField("module", "mdb").Error(err.Error())
		panic(err)
	}

	// setup TLS if we use SSL
	if newUrl != url {
		tlsConfig := &tls.Config{}
		tlsConfig.InsecureSkipVerify = true

		dialInfo.DialServer = func(addr *mgo.ServerAddr) (net.Conn, error) {
			conn, err := tls.Dial("tcp", addr.String(), tlsConfig)
			return conn, err
		}
	}

	mDbSession, err = mgo.DialWithInfo(dialInfo)
	if err != nil {
		log.WithField("module", "mdb").Error(err.Error())
		panic(err)
	}

	mDbSession.SetMode(mgo.Monotonic, true)
	mDbSession.SetSafe(&mgo.Safe{WMode: "majority"})

	mDbDatabase = database

	log.WithField("module", "mdb").Info("Connected!")
}

// GetDB is a simple getter for the mongodb database.
func GetMDb() *mgo.Database {
	return mDbSession.DB(mDbDatabase)
}

// GetDB is a simple getter for the mongodb session.
func GetMDbSession() *mgo.Session {
	return mDbSession
}

func MDbInsert(collection models.MongoDbCollection, data interface{}) (rid bson.ObjectId, err error) {
	ptr := reflect.New(reflect.TypeOf(data))
	temp := ptr.Elem()
	temp.Set(reflect.ValueOf(data))

	v := temp.FieldByName("ID")

	if !v.IsValid() {
		return bson.ObjectId(""), errors.New("invalid data")
	}

	newID := v.String()
	if newID == "" {
		newID = string(bson.NewObjectId())
		v.SetString(newID)
	}

	start := time.Now()
	err = GetMDb().C(collection.String()).Insert(temp.Interface())
	took := time.Since(start)

	if cache.HasKeen() {
		go func() {
			defer Recover()

			err := cache.GetKeen().AddEvent("Robyul_MongoDB", &KeenMongoDbEvent{
				Seconds:    took.Seconds(),
				Type:       "insert",
				Method:     "MDbInsert()",
				Collection: stripRobyulDatabaseFromCollection(collection.String()),
				Data:       fmt.Sprintf("%+v", data),
			})
			if err != nil {
				cache.GetLogger().WithField("module", "mdb").Error("Error logging MongoDB request to keen: ", err.Error())
			}
		}()
	}

	if err != nil {
		return bson.ObjectId(""), err
	}

	return bson.ObjectId(newID), nil
}

func MDbInsertWithoutLogging(collection models.MongoDbCollection, data interface{}) (rid bson.ObjectId, err error) {
	ptr := reflect.New(reflect.TypeOf(data))
	temp := ptr.Elem()
	temp.Set(reflect.ValueOf(data))

	v := temp.FieldByName("ID")

	if !v.IsValid() {
		return bson.ObjectId(""), errors.New("invalid data")
	}

	newID := v.String()
	if newID == "" {
		newID = string(bson.NewObjectId())
		v.SetString(newID)
	}

	err = GetMDb().C(collection.String()).Insert(temp.Interface())

	if err != nil {
		return bson.ObjectId(""), err
	}

	return bson.ObjectId(newID), nil
}

func MDbUpdate(collection models.MongoDbCollection, id bson.ObjectId, data interface{}) (rid bson.ObjectId, err error) {
	if !id.Valid() {
		return bson.ObjectId(""), errors.New("invalid id")
	}

	start := time.Now()
	err = GetMDb().C(collection.String()).UpdateId(id, data)
	took := time.Since(start)

	if cache.HasKeen() {
		go func() {
			defer Recover()

			err := cache.GetKeen().AddEvent("Robyul_MongoDB", &KeenMongoDbEvent{
				Seconds:    took.Seconds(),
				Type:       "update",
				Method:     "MDbUpdate()",
				Collection: stripRobyulDatabaseFromCollection(collection.String()),
				Id:         id.String(),
				Data:       fmt.Sprintf("%+v", data),
			})
			if err != nil {
				cache.GetLogger().WithField("module", "mdb").Error("Error logging MongoDB request to keen: ", err.Error())
			}
		}()
	}

	if err != nil {
		return bson.ObjectId(""), err
	}

	return id, nil
}

func MDbUpdateWithoutLogging(collection models.MongoDbCollection, id bson.ObjectId, data interface{}) (rid bson.ObjectId, err error) {
	if !id.Valid() {
		return bson.ObjectId(""), errors.New("invalid id")
	}

	err = GetMDb().C(collection.String()).UpdateId(id, data)

	if err != nil {
		return bson.ObjectId(""), err
	}

	return id, nil
}

func MDbUpsertID(collection models.MongoDbCollection, id bson.ObjectId, data interface{}) (rid bson.ObjectId, err error) {
	if !id.Valid() {
		id = bson.NewObjectId()
	}

	start := time.Now()
	_, err = GetMDb().C(collection.String()).UpsertId(id, data)
	took := time.Since(start)

	if cache.HasKeen() {
		go func() {
			defer Recover()

			err := cache.GetKeen().AddEvent("Robyul_MongoDB", &KeenMongoDbEvent{
				Seconds:    took.Seconds(),
				Type:       "upsert",
				Method:     "MDbUpsertID()",
				Collection: stripRobyulDatabaseFromCollection(collection.String()),
				Id:         id.String(),
				Data:       fmt.Sprintf("%+v", data),
			})
			if err != nil {
				cache.GetLogger().WithField("module", "mdb").Error("Error logging MongoDB request to keen: ", err.Error())
			}
		}()
	}

	if err != nil {
		return bson.ObjectId(""), err
	}

	return id, nil
}

func MDbUpsert(collection models.MongoDbCollection, selector interface{}, data interface{}) (err error) {
	start := time.Now()
	_, err = GetMDb().C(collection.String()).Upsert(selector, data)
	took := time.Since(start)

	if cache.HasKeen() {
		go func() {
			defer Recover()

			err := cache.GetKeen().AddEvent("Robyul_MongoDB", &KeenMongoDbEvent{
				Seconds:    took.Seconds(),
				Type:       "upsert",
				Method:     "MDbUpsert()",
				Collection: stripRobyulDatabaseFromCollection(collection.String()),
				Query:      fmt.Sprintf("%+v", selector),
				Data:       fmt.Sprintf("%+v", data),
			})
			if err != nil {
				cache.GetLogger().WithField("module", "mdb").Error("Error logging MongoDB request to keen: ", err.Error())
			}
		}()
	}

	return err
}

func MDbDelete(collection models.MongoDbCollection, id bson.ObjectId) (err error) {
	if !id.Valid() {
		return errors.New("invalid id")
	}

	start := time.Now()
	err = GetMDb().C(collection.String()).RemoveId(id)
	took := time.Since(start)

	if cache.HasKeen() {
		go func() {
			defer Recover()

			err := cache.GetKeen().AddEvent("Robyul_MongoDB", &KeenMongoDbEvent{
				Seconds:    took.Seconds(),
				Type:       "remove",
				Method:     "MDbDelete()",
				Collection: stripRobyulDatabaseFromCollection(collection.String()),
				Id:         id.String(),
			})
			if err != nil {
				cache.GetLogger().WithField("module", "mdb").Error("Error logging MongoDB request to keen: ", err.Error())
			}
		}()
	}

	if err != nil {
		return err
	}

	return nil
}

func MdbCollection(collection models.MongoDbCollection) (query *mgo.Collection) {
	return GetMDb().C(collection.String())
}

func MDbIter(query *mgo.Query) (iter *mgo.Iter) {
	start := time.Now()
	iter = query.Iter()
	took := time.Since(start)
	if cache.HasKeen() {
		go func() {
			defer Recover()

			queryValue := reflect.ValueOf(*query)
			queryOp := queryValue.FieldByName("query").FieldByName("op")

			err := cache.GetKeen().AddEvent("Robyul_MongoDB", &KeenMongoDbEvent{
				Seconds:    took.Seconds(),
				Type:       "query",
				Method:     "MdbIter()",
				Collection: stripRobyulDatabaseFromCollection(queryOp.FieldByName("collection").String()),
				Query:      fmt.Sprintf("%+v", reflect.ValueOf(queryOp.FieldByName("query")).Interface()),
				Skip:       queryOp.FieldByName("skip").Int(),
				Limit:      queryOp.FieldByName("limit").Int(),
			})
			if err != nil {
				cache.GetLogger().WithField("module", "mdb").Error("Error logging MongoDB request to keen: ", err.Error())
			}
		}()
	}
	return
}

func MdbOne(query *mgo.Query, object interface{}) (err error) {
	start := time.Now()
	err = query.One(object)
	took := time.Since(start)
	if cache.HasKeen() {
		go func() {
			defer Recover()

			queryValue := reflect.ValueOf(*query)
			queryOp := queryValue.FieldByName("query").FieldByName("op")

			err := cache.GetKeen().AddEvent("Robyul_MongoDB", &KeenMongoDbEvent{
				Seconds:    took.Seconds(),
				Type:       "query",
				Method:     "MdbOne()",
				Collection: stripRobyulDatabaseFromCollection(queryOp.FieldByName("collection").String()),
				Query:      fmt.Sprintf("%+v", reflect.ValueOf(queryOp.FieldByName("query")).Interface()),
				Skip:       queryOp.FieldByName("skip").Int(),
				Limit:      queryOp.FieldByName("limit").Int(),
			})
			if err != nil {
				cache.GetLogger().WithField("module", "mdb").Error("Error logging MongoDB request to keen: ", err.Error())
			}
		}()
	}
	return
}

func MdbOneWithoutLogging(query *mgo.Query, object interface{}) (err error) {
	err = query.One(object)
	return
}

func stripRobyulDatabaseFromCollection(input string) (output string) {
	return strings.TrimPrefix(input, mDbDatabase+".")
}

type KeenMongoDbEvent struct {
	Seconds    float64
	Collection string
	Type       string
	Method     string
	Query      string `json:",omitempty"`
	Skip       int64  `json:",omitempty"`
	Limit      int64  `json:",omitempty"`
	Id         string `json:",omitempty"`
	Data       string `json:",omitempty"`
}
<|endoftext|>