Add bench command

cmd/genji/commands/app.go

@@ -35,6 +35,7 @@ func NewApp() *cli.App {
 		NewVersionCommand(),
 		NewDumpCommand(),
 		NewRestoreCommand(),
+		NewBenchCommand(),
 	}
 
 	// Root command

cmd/genji/commands/bench.go

@@ -0,0 +1,121 @@
+package commands
+
+import (
+	"errors"
+
+	"github.com/genjidb/genji/cmd/genji/dbutil"
+	"github.com/urfave/cli/v2"
+)
+
+// NewBenchCommand returns a cli.Command for "genji bench".
+func NewBenchCommand() *cli.Command {
+	cmd := cli.Command{
+		Name:      "bench",
+		Usage:     "Simple load testing command",
+		UsageText: `genji bench query`,
+		Description: `The bench command runs a query repeatedly (100 times by default, -n option) and outputs a series of results.
+Each result represent the average time for a given sample of queries (10 by default, -s/--sample option).
+
+$ genji bench -n 200 -s 5 "SELECT 1"
+{
+	"totalQueries": 5,
+	"sampleSpeed": "2.191µs"
+}
+{
+	"totalQueries": 10,
+	"sampleSpeed": "1.941µs"
+}
+{
+	"totalQueries": 15,
+	"sampleSpeed": "2.237µs"
+}
+...
+
+
+By default, queries are run in-memory. To choose a different engine, use the -e/--engine and -p/--path options.
+The database will be created if it doesn't exist.
+
+$ genji bench -e bolt -p my.db "SELECT 1"
+$ genji bench -e badger -p mydb/ "SELECT 1"
+
+To prepare the database before running a query, use the -i/--init option
+
+$ genji bench -p "CREATE TABLE foo; INSERT INTO foo(a) VALUES (1), (2), (3)" "SELECT * FROM foo"
+
+By default, each query is run in a separate transaction. To run everything, including the setup,
+in the same transaction, use -t`,
+		Flags: []cli.Flag{
+			&cli.StringFlag{
+				Name:    "engine",
+				Aliases: []string{"e"},
+				Usage:   "name of the engine to use, options are 'bolt', 'badger' or 'memory'. Default to 'memory'",
+				Value:   "memory",
+			},
+			&cli.StringFlag{
+				Name:    "path",
+				Aliases: []string{"p"},
+				Usage:   "Path of the database to open or create. Only valid if for bolt or badger engines",
+			},
+			&cli.StringFlag{
+				Name:    "init",
+				Aliases: []string{"i"},
+				Usage:   "Queries to run to initialize the database before running the benchmark.",
+			},
+			&cli.BoolFlag{
+				Name:    "tx",
+				Aliases: []string{"x"},
+				Usage:   "Run everything in the same transaction.",
+			},
+			&cli.IntFlag{
+				Name:    "number",
+				Aliases: []string{"n"},
+				Value:   100,
+				Usage:   "Total number of queries to run.",
+			},
+			&cli.IntFlag{
+				Name:    "sample",
+				Aliases: []string{"s"},
+				Value:   10,
+				Usage:   "Number of queries to use to determine the average speed of the query.",
+			},
+			&cli.BoolFlag{
+				Name:  "prepare",
+				Usage: "Prepare the query before running the benchmark",
+			},
+			&cli.BoolFlag{
+				Name:  "csv",
+				Usage: "Output the results in csv",
+			},
+		},
+	}
+
+	cmd.Action = func(c *cli.Context) error {
+		query := c.Args().First()
+		if query == "" {
+			return errors.New(cmd.UsageText)
+		}
+
+		engine := c.String("engine")
+		path := c.String("path")
+		if engine == "" {
+			return errors.New(cmd.UsageText)
+		}
+
+		db, err := dbutil.OpenDB(c.Context, path, engine, dbutil.DBOptions{})
+		if err != nil {
+			return err
+		}
+		defer db.Close()
+
+		return dbutil.Bench(c.Context, db, query, dbutil.BenchOptions{
+			Init:       c.String("init"),
+			N:          c.Int("number"),
+			SampleSize: c.Int("sample"),
+			SameTx:     c.Bool("tx"),
+			Prepare:    c.Bool("prepare"),
+			CSV:        c.Bool("csv"),
+		})
+	}
+
+	return &cmd
+}

cmd/genji/dbutil/bench.go

@@ -0,0 +1,150 @@
+package dbutil
+
+import (
+	"context"
+	"encoding/csv"
+	"encoding/json"
+	"io"
+	"os"
+	"strconv"
+	"strings"
+	"time"
+
+	"github.com/genjidb/genji"
+)
+
+type BenchOptions struct {
+	Init       string
+	N          int
+	SampleSize int
+	SameTx     bool
+	Prepare    bool
+	CSV        bool
+}
+
+type preparer interface {
+	Prepare(q string) (*genji.Statement, error)
+}
+
+type execer func(q string, args ...interface{}) error
+
+// Bench takes a database and dumps its content as SQL queries in the given writer.
+// If tables is provided, only selected tables will be outputted.
+func Bench(ctx context.Context, db *genji.DB, query string, opt BenchOptions) error {
+	var p preparer = db
+	var e execer = db.Exec
+
+	if opt.SameTx {
+		tx, err := db.Begin(true)
+		if err != nil {
+			return err
+		}
+		defer tx.Rollback()
+		p = tx
+		e = tx.Exec
+	}
+
+	if opt.Init != "" {
+		err := e(opt.Init)
+		if err != nil {
+			return err
+		}
+	}
+
+	if opt.Prepare {
+		stmt, err := p.Prepare(query)
+		if err != nil {
+			return err
+		}
+		e = func(q string, args ...interface{}) error {
+			return stmt.Exec()
+		}
+	}
+
+	var enc encoder
+	if opt.CSV {
+		enc = newCSVWriter(os.Stdout)
+	} else {
+		enc = newJSONWriter(os.Stdout)
+	}
+
+	var totalDuration time.Duration
+	for i := 0; i < opt.N; i += opt.SampleSize {
+		var total time.Duration
+
+		for j := 0; j < opt.SampleSize; j++ {
+			start := time.Now()
+
+			err := e(query)
+			total += time.Since(start)
+			if err != nil {
+				return err
+			}
+		}
+
+		totalDuration += total
+		avg := total / time.Duration(opt.SampleSize)
+		qps := int(time.Second / avg)
+
+		err := enc(map[string]interface{}{
+			"totalQueries":     i + opt.SampleSize,
+			"averageDuration":  avg,
+			"queriesPerSecond": qps,
+			"totalDuration":    totalDuration,
+		})
+		if err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+type encoder func(map[string]interface{}) error
+
+func newJSONWriter(w io.Writer) func(map[string]interface{}) error {
+	enc := json.NewEncoder(w)
+	enc.SetIndent("", "  ")
+	return func(m map[string]interface{}) error {
+		return enc.Encode(m)
+	}
+}
+
+func newCSVWriter(w io.Writer) func(map[string]interface{}) error {
+	enc := csv.NewWriter(w)
+	enc.Comma = ';'
+	header := []string{"totalQueries", "averageDuration", "queriesPerSecond", "totalDuration"}
+	var headerWritten bool
+
+	return func(m map[string]interface{}) error {
+		if !headerWritten {
+			err := enc.Write(header)
+			if err != nil {
+				return err
+			}
+			headerWritten = true
+		}
+		err := enc.Write([]string{
+			strconv.Itoa(m["totalQueries"].(int)),
+			durationToString(m["averageDuration"].(time.Duration)),
+			strconv.Itoa(m["queriesPerSecond"].(int)),
+			durationToString(m["totalDuration"].(time.Duration)),
+		})
+		if err != nil {
+			return err
+		}
+		enc.Flush()
+		return enc.Error()
+	}
+}
+
+func durationToMilliseconds(d time.Duration) float64 {
+	m := d / time.Millisecond
+	nsec := d % time.Millisecond
+	return float64(m) + float64(nsec)/1e6
+}
+
+func durationToString(d time.Duration) string {
+	ms := durationToMilliseconds(d)
+	return strings.Replace(strconv.FormatFloat(ms, 'f', -1, 64), ".", ",", 1)
+}

internal/sql/parser/parser_fuzz_test.go

@@ -1,5 +1,5 @@
-//go:build go1.17
-// +build go1.17
+//go:build go1.18
+// +build go1.18
 
 package parser