Implements CLI tool for rule management

go/cmd/rulesctl/add.go

@@ -0,0 +1,149 @@
+package main
+
+import (
+	"log"
+	"os"
+	"strings"
+
+	"github.com/spf13/cobra"
+
+	"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
+	// 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 err := rule.SetQueryCond(addOptQueryRE); err != nil {
+		log.Fatalf("Query condition invalid '%v': %v", addOptQueryRE, err)
+	}
+
+	var rules *vtrules.Rules
+	_, err := os.Stat(configFile)
+	if os.IsNotExist(err) {
+		rules = vtrules.New()
+	} else {
+		rules = getRules()
+	}
+	existingRule := rules.Find(rule.Name)
+	if existingRule != nil {
+		log.Fatalf("Rule by name %q already exists", rule.Name)
+	}
+	rules.Add(rule)
+
+	if addOptDryrun {
+		mustPrintJSON(rules)
+	} else {
+		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 getAddCmd() *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")
+
+	for _, f := range []string{"name", "action"} {
+		addCmd.MarkFlagRequired(f)
+	}
+
+	return addCmd
+}

go/cmd/rulesctl/explain.go

@@ -0,0 +1,46 @@
+package main
+
+import (
+	"fmt"
+
+	"github.com/spf13/cobra"
+
+	"vitess.io/vitess/go/vt/vttablet/tabletserver/planbuilder"
+)
+
+func getExplainCmd() *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/main.go

@@ -0,0 +1,162 @@
+package main
+
+import (
+	"encoding/json"
+	"flag"
+	"fmt"
+	"io/ioutil"
+	"log"
+	"os"
+
+	"github.com/spf13/cobra"
+
+	vtfcr "vitess.io/vitess/go/vt/vttablet/customrule/filecustomrule"
+	"vitess.io/vitess/go/vt/vttablet/tabletserver/rules"
+)
+
+var configFile string
+
+func main() {
+	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
+		},
+	}
+	rootCmd.Run = func(_ *cobra.Command, _ []string) { rootCmd.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(getListCmd())
+	rootCmd.AddCommand(getRemoveCmd())
+	rootCmd.AddCommand(getAddCmd())
+	rootCmd.AddCommand(getExplainCmd())
+
+	if err := rootCmd.Execute(); err != nil {
+		log.Printf("%v", err)
+	}
+}
+
+func getRemoveCmd() *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 := getRules()
+		if deleted := rules.Delete(removeOptName); deleted == nil {
+			fmt.Printf("No rule found: '%v'", removeOptName)
+			return
+		}
+
+		if removeOptDryRun {
+			mustPrintJSON(rules)
+		} else {
+			mustWriteJSON(rules, configFile)
+		}
+	}
+
+	return removeCmd
+}
+
+func getListCmd() *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 := getRules()
+
+		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 = ""
+			}
+		}
+
+		mustPrintJSON(out)
+	}
+
+	return listCmd
+}
+
+func getRules() *rules.Rules {
+	rules, err := vtfcr.ParseRules(configFile)
+	if err != nil {
+		log.Fatalf("Failure attempting to parse rules: %v", err)
+	}
+	return rules
+}
+
+func mustPrintJSON(obj interface{}) {
+	enc, err := json.MarshalIndent(obj, "", "  ")
+	if err != nil {
+		log.Fatalf("Unable to marshal object: %v", err)
+	}
+	fmt.Printf("%v\n", string(enc))
+}
+
+func mustWriteJSON(obj interface{}, path string) {
+	enc, err := json.MarshalIndent(obj, "", "  ")
+	if err != nil {
+		log.Fatalf("Unable to marshal object: %v", err)
+	}
+
+	err = ioutil.WriteFile(path, enc, 0400)
+	if err != nil {
+		log.Fatalf("Unable to save new JSON: %v", err)
+	}
+}