Query killer

go.mod

@@ -23,6 +23,7 @@ require (
 	github.com/corpix/uarand v0.1.1 // indirect
 	github.com/cyberdelia/go-metrics-graphite v0.0.0-20161219230853-39f87cc3b432
 	github.com/evanphx/json-patch v4.5.0+incompatible
+	github.com/fsnotify/fsnotify v1.4.9
 	github.com/go-martini/martini v0.0.0-20170121215854-22fa46961aab
 	github.com/go-sql-driver/mysql v1.5.0
 	github.com/gogo/protobuf v1.3.1 // indirect

go.sum

@@ -173,14 +173,15 @@ github.com/emicklei/go-restful v2.9.5+incompatible/go.mod h1:otzb+WCGbkyDHkqmQmT
 github.com/evanphx/json-patch v4.2.0+incompatible/go.mod h1:50XU6AFN0ol/bzJsmQLiYLvXMP4fmwYFNcr97nuDLSk=
 github.com/evanphx/json-patch v4.5.0+incompatible h1:ouOWdg56aJriqS0huScTkVXPC5IcNrDCXZ6OoTAWu7M=
 github.com/evanphx/json-patch v4.5.0+incompatible/go.mod h1:50XU6AFN0ol/bzJsmQLiYLvXMP4fmwYFNcr97nuDLSk=
-github.com/fatih/color v1.7.0 h1:DkWD4oS2D8LGGgTQ6IvwJJXSL5Vp2ffcQg58nFV38Ys=
 github.com/fatih/color v1.7.0/go.mod h1:Zm6kSWBoL9eyXnKyktHP6abPY2pDugNf5KwzbycvMj4=
 github.com/fatih/color v1.9.0 h1:8xPHl4/q1VyqGIPif1F+1V3Y3lSmrq01EabUW3CoW5s=
 github.com/fatih/color v1.9.0/go.mod h1:eQcE1qtQxscV5RaZvpXrrb8Drkc3/DdQ+uUYCNjL+zU=
 github.com/felixge/httpsnoop v1.0.1 h1:lvB5Jl89CsZtGIWuTcDM1E/vkVs49/Ml7JJe07l8SPQ=
 github.com/felixge/httpsnoop v1.0.1/go.mod h1:m8KPJKqk1gH5J9DgRY2ASl2lWCfGKXixSwevea8zH2U=
 github.com/fsnotify/fsnotify v1.4.7 h1:IXs+QLmnXW2CcXuY+8Mzv/fWEsPGWxqefPtCP5CnV9I=
 github.com/fsnotify/fsnotify v1.4.7/go.mod h1:jwhsz4b93w/PPRr/qN1Yymfu8t87LnFCMoQvtojpjFo=
+github.com/fsnotify/fsnotify v1.4.9 h1:hsms1Qyu0jgnwNXIxa+/V/PDsU6CfLf6CNO8H7IWoS4=
+github.com/fsnotify/fsnotify v1.4.9/go.mod h1:znqG4EE+3YCdAaPaxE2ZRY/06pZUdp0tY4IgpuI1SZQ=
 github.com/ghodss/yaml v0.0.0-20150909031657-73d445a93680/go.mod h1:4dBDuWmgqj2HViK6kFavaiC9ZROes6MMH2rRYeMEF04=
 github.com/ghodss/yaml v1.0.0 h1:wQHKEahhL6wmXdzwWG11gIVCkOv05bNOh+Rxn0yngAk=
 github.com/ghodss/yaml v1.0.0/go.mod h1:4dBDuWmgqj2HViK6kFavaiC9ZROes6MMH2rRYeMEF04=
@@ -821,6 +822,7 @@ golang.org/x/sys v0.0.0-20190624142023-c5567b49c5d0/go.mod h1:h1NjWce9XRLGQEsW7w
 golang.org/x/sys v0.0.0-20190826190057-c7b8b68b1456/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20190922100055-0a153f010e69/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20190924154521-2837fb4f24fe/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
+golang.org/x/sys v0.0.0-20191005200804-aed5e4c7ecf9/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20191008105621-543471e840be/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20191026070338-33540a1f6037/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20200116001909-b77594299b42/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=

go/cmd/rulesctl/cmd/add.go

@@ -0,0 +1,174 @@
+package cmd
+
+import (
+	"log"
+	"os"
+	"strings"
+
+	"github.com/spf13/cobra"
+
+	"vitess.io/vitess/go/cmd/rulesctl/common"
+	"vitess.io/vitess/go/vt/vttablet/tabletserver/planbuilder"
+	vtrules "vitess.io/vitess/go/vt/vttablet/tabletserver/rules"
+)
+
+var (
+	addOptDryrun            bool
+	addOptName              string
+	addOptDescription       string
+	addOptAction            string
+	addOptPlans             []string
+	addOptTables            []string
+	addOptQueryRE           string
+	addOptLeadingCommentRE  string
+	addOptTrailingCommentRE string
+	// TODO: other stuff, bind vars etc
+)
+
+func runAdd(cmd *cobra.Command, args []string) {
+	rulePlans := mkPlanSlice()
+	ruleAction := mkAction()
+
+	rule := vtrules.NewQueryRule(addOptDescription, addOptName, ruleAction)
+	for _, pt := range rulePlans {
+		rule.AddPlanCond(pt)
+	}
+
+	for _, t := range addOptTables {
+		rule.AddTableCond(t)
+	}
+
+	if addOptQueryRE != "" {
+		if err := rule.SetQueryCond(addOptQueryRE); err != nil {
+			log.Fatalf("Query condition invalid '%v': %v", addOptQueryRE, err)
+		}
+	}
+	if addOptLeadingCommentRE != "" {
+		if err := rule.SetLeadingCommentCond(addOptLeadingCommentRE); err != nil {
+			log.Fatalf("Leading comment condition invalid '%v': %v", addOptLeadingCommentRE, err)
+		}
+	}
+	if addOptTrailingCommentRE != "" {
+		if err := rule.SetTrailingCommentCond(addOptTrailingCommentRE); err != nil {
+			log.Fatalf("Trailing comment condition invalid '%v': %v", addOptTrailingCommentRE, err)
+		}
+	}
+
+	var rules *vtrules.Rules
+	_, err := os.Stat(configFile)
+	if os.IsNotExist(err) {
+		rules = vtrules.New()
+	} else {
+		rules = common.GetRules(configFile)
+	}
+	existingRule := rules.Find(rule.Name)
+	if existingRule != nil {
+		log.Fatalf("Rule by name %q already exists", rule.Name)
+	}
+	rules.Add(rule)
+
+	if addOptDryrun {
+		common.MustPrintJSON(rules)
+	} else {
+		common.MustWriteJSON(rules, configFile)
+	}
+}
+
+func mkPlanSlice() []planbuilder.PlanType {
+	if len(addOptPlans) == 0 {
+		return nil
+	}
+
+	plans := []planbuilder.PlanType{}
+	badPlans := []string{}
+
+	for _, p := range addOptPlans {
+		if pbn, ok := planbuilder.PlanByNameIC(p); ok {
+			plans = append(plans, pbn)
+		} else {
+			badPlans = append(badPlans, p)
+		}
+	}
+
+	if len(badPlans) != 0 {
+		log.Fatalf("Unknown PlanType(s) %q", badPlans)
+	}
+
+	return plans
+}
+
+func mkAction() vtrules.Action {
+	switch strings.ToLower(addOptAction) {
+	case "fail":
+		return vtrules.QRFail
+	case "fail_retry":
+		return vtrules.QRFailRetry
+	case "continue":
+		return vtrules.QRContinue
+	default:
+		log.Fatalf("Unknown action '%v'", addOptAction)
+	}
+
+	panic("Nope")
+}
+
+func Add() *cobra.Command {
+	addCmd := &cobra.Command{
+		Use:   "add-rule",
+		Short: "Adds a rule to the config file",
+		Args:  cobra.NoArgs,
+		Run:   runAdd,
+	}
+
+	addCmd.Flags().BoolVarP(
+		&addOptDryrun,
+		"dry-run", "d",
+		false,
+		"Instead of writing the config file back print the result to stdout")
+	addCmd.Flags().StringVarP(
+		&addOptName,
+		"name", "n",
+		"",
+		"The name of the rule to add (required)")
+	addCmd.Flags().StringVarP(
+		&addOptDescription,
+		"description", "e",
+		"",
+		"The purpose/description of the rule being added")
+	addCmd.Flags().StringVarP(
+		&addOptAction,
+		"action", "a",
+		"",
+		"What action should be taken when this rule is matched {continue, fail, fail-retry} (required)")
+	addCmd.Flags().StringSliceVarP(
+		&addOptPlans,
+		"plan", "p",
+		nil,
+		"Which query plan types does this rule match; see \"explain query-plans\" for details; may be specified multiple times")
+	addCmd.Flags().StringSliceVarP(
+		&addOptTables,
+		"table", "t",
+		nil,
+		"Queries will only match if running against these tables; may be specified multiple times")
+	addCmd.Flags().StringVarP(
+		&addOptQueryRE,
+		"query", "q",
+		"",
+		"A regexp that will be applied to a query in order to determine if it matches")
+	addCmd.Flags().StringVarP(
+		&addOptLeadingCommentRE,
+		"leading-comment", "l",
+		"",
+		"A regexp that will be applied to comments prefacing a SQL statement")
+	addCmd.Flags().StringVarP(
+		&addOptTrailingCommentRE,
+		"trailing-comment", "r",
+		"",
+		"A regexp that will be applied to comments after a SQL statement")
+
+	for _, f := range []string{"name", "action"} {
+		addCmd.MarkFlagRequired(f)
+	}
+
+	return addCmd
+}

go/cmd/rulesctl/cmd/explain.go

@@ -0,0 +1,46 @@
+package cmd
+
+import (
+	"fmt"
+
+	"github.com/spf13/cobra"
+
+	"vitess.io/vitess/go/vt/vttablet/tabletserver/planbuilder"
+)
+
+func Explain() *cobra.Command {
+	explain := &cobra.Command{
+		Use:   "explain [concept]",
+		Short: "Explains a concept, valid options are: query-plans",
+		Args:  cobra.ExactArgs(1),
+		Run:   runExplain,
+	}
+	return explain
+}
+
+func runExplain(cmd *cobra.Command, args []string) {
+	lookup := map[string]func(){
+		"query-plans": helpQueryPlans,
+	}
+
+	if fn, ok := lookup[args[0]]; ok {
+		fn()
+	} else {
+		fmt.Printf("I don't know anything about %q, sorry!", args[0])
+	}
+}
+
+func helpQueryPlans() {
+	fmt.Printf(`Query Plans!
+
+A query plan is the type of work the Tablet is about to do. When used in a rule
+it can be used to limit the class of queries that a rule can impact. In other
+words it will allow you to say "this rule only fails inserts" or "this rule only
+fails selects."
+
+The list of valid plan types that can be used follows:
+`)
+	for i := 0; i < int(planbuilder.NumPlans); i++ {
+		fmt.Printf("  - %v\n", planbuilder.PlanType(i).String())
+	}
+}

go/cmd/rulesctl/cmd/list.go

@@ -0,0 +1,55 @@
+package cmd
+
+import (
+	"github.com/spf13/cobra"
+
+	"vitess.io/vitess/go/cmd/rulesctl/common"
+)
+
+func List() *cobra.Command {
+	var listOptName string
+	var listOptNamesOnly bool
+	listCmd := &cobra.Command{
+		Use:   "list",
+		Short: "Display the rules in the config file",
+		Args:  cobra.NoArgs,
+	}
+
+	listCmd.Flags().StringVarP(
+		&listOptName,
+		"name", "n",
+		"",
+		"Display a named rule (optional)")
+	listCmd.Flags().BoolVar(
+		&listOptNamesOnly,
+		"names-only",
+		false,
+		"Lists only the names of the rules in the config file")
+
+	listCmd.Run = func(cmd *cobra.Command, args []string) {
+		rules := common.GetRules(configFile)
+
+		var out interface{}
+		if listOptName == "" {
+			if listOptNamesOnly {
+				out = []string{}
+				for _, r := range rules.CopyUnderlying() {
+					out = append(out.([]string), r.Name)
+				}
+			} else {
+				out = rules
+			}
+		} else {
+			out = rules.Find(listOptName)
+			if listOptNamesOnly && out != nil {
+				out = listOptName
+			} else if listOptNamesOnly {
+				out = ""
+			}
+		}
+
+		common.MustPrintJSON(out)
+	}
+
+	return listCmd
+}

go/cmd/rulesctl/cmd/main.go

@@ -0,0 +1,38 @@
+package cmd
+
+import (
+	"flag"
+	"os"
+
+	"github.com/spf13/cobra"
+)
+
+var configFile string
+
+func Main() *cobra.Command {
+	rootCmd := &cobra.Command{
+		Use:  "rulesctl",
+		Args: cobra.NoArgs,
+		PreRun: func(cmd *cobra.Command, args []string) {
+			tmp := os.Args
+			os.Args = os.Args[0:1]
+			flag.Parse()
+			os.Args = tmp
+		},
+		Run: func(cmd *cobra.Command, _ []string) { cmd.Help() },
+	}
+
+	rootCmd.PersistentFlags().StringVarP(
+		&configFile,
+		"config-file", "f",
+		"rules.json",
+		"the config file we will be using to store query rules")
+	rootCmd.MarkPersistentFlagFilename("config-file")
+
+	rootCmd.AddCommand(List())
+	rootCmd.AddCommand(Remove())
+	rootCmd.AddCommand(Add())
+	rootCmd.AddCommand(Explain())
+
+	return rootCmd
+}

go/cmd/rulesctl/cmd/remove.go

@@ -0,0 +1,48 @@
+package cmd
+
+import (
+	"fmt"
+
+	"github.com/spf13/cobra"
+
+	"vitess.io/vitess/go/cmd/rulesctl/common"
+)
+
+func Remove() *cobra.Command {
+	var removeOptName string
+	var removeOptDryRun bool
+
+	removeCmd := &cobra.Command{
+		Use:   "remove-rule",
+		Short: "Removes a named rule from the config file",
+		Args:  cobra.NoArgs,
+	}
+
+	removeCmd.Flags().StringVarP(
+		&removeOptName,
+		"name", "n",
+		"",
+		"The named rule to remove (required)")
+	removeCmd.Flags().BoolVarP(
+		&removeOptDryRun,
+		"dry-run", "d",
+		false,
+		"Instead of writing the config file back print the result to stdout")
+	removeCmd.MarkFlagRequired("name")
+
+	removeCmd.Run = func(cmd *cobra.Command, args []string) {
+		rules := common.GetRules(configFile)
+		if deleted := rules.Delete(removeOptName); deleted == nil {
+			fmt.Printf("No rule found: '%v'", removeOptName)
+			return
+		}
+
+		if removeOptDryRun {
+			common.MustPrintJSON(rules)
+		} else {
+			common.MustWriteJSON(rules, configFile)
+		}
+	}
+
+	return removeCmd
+}