Merge pull request #7784 from tinyspeck/extend-queryrules

[tablet, queryrules] Extend query rules to check comments

go.mod

@@ -25,6 +25,7 @@ require (
  github.com/cyberdelia/go-metrics-graphite v0.0.0-20161219230853-39f87cc3b432
  github.com/dave/jennifer v1.4.1
  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.1-0.20210202043019-fe2230a8b20c
  github.com/gogo/protobuf v1.3.1

go.sum

@@ -182,6 +182,8 @@ github.com/felixge/httpsnoop v1.0.1 h1:lvB5Jl89CsZtGIWuTcDM1E/vkVs49/Ml7JJe07l8S
 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=
@@ -783,6 +785,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

@@ -13,13 +13,15 @@ import (
 )
 
 var (
- addOptDryrun bool
- addOptName string
- addOptDescription string
- addOptAction string
- addOptPlans []string
- addOptTables []string
- addOptQueryRE string
+ addOptDryrun bool
+ addOptName string
+ addOptDescription string
+ addOptAction string
+ addOptPlans []string
+ addOptTables []string
+ addOptQueryRE string
+ addOptLeadingCommentRE string
+ addOptTrailingCommentRE string
  // TODO: other stuff, bind vars etc
 )
 
@@ -36,8 +38,20 @@ func runAdd(cmd *cobra.Command, args []string) {
  rule.AddTableCond(t)
  }
 
- if err := rule.SetQueryCond(addOptQueryRE); err != nil {
- log.Fatalf("Query condition invalid '%v': %v", addOptQueryRE, err)
+ 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
@@ -141,6 +155,16 @@ func Add() *cobra.Command {
  "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)

go/vt/vttablet/customrule/filecustomrule/filecustomrule.go

@@ -1,5 +1,5 @@
 /*
-Copyright 2019 The Vitess Authors.
+Copyright 2021 The Vitess Authors.
 
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
@@ -20,9 +20,13 @@ package filecustomrule
 import (
  "flag"
  "io/ioutil"
+ "path"
  "time"
 
+ "github.com/fsnotify/fsnotify"
+
  "vitess.io/vitess/go/vt/log"
+ "vitess.io/vitess/go/vt/servenv"
  "vitess.io/vitess/go/vt/vttablet/tabletserver"
  "vitess.io/vitess/go/vt/vttablet/tabletserver/rules"
 )
@@ -32,6 +36,8 @@ var (
  fileCustomRule = NewFileCustomRule()
  // Commandline flag to specify rule path
  fileRulePath = flag.String("filecustomrules", "", "file based custom rule path")
+
+ fileRuleShouldWatch = flag.Bool("filecustomrules_watch", false, "set up a watch on the target file and reload query rules when it changes")
 )
 
 // FileCustomRule is an implementation of CustomRuleManager, it reads custom query
@@ -102,6 +108,45 @@ func ActivateFileCustomRules(qsc tabletserver.Controller) {
  if *fileRulePath != "" {
  qsc.RegisterQueryRuleSource(FileCustomRuleSource)
  fileCustomRule.Open(qsc, *fileRulePath)
+
+ if *fileRuleShouldWatch {
+ baseDir := path.Dir(*fileRulePath)
+ ruleFileName := path.Base(*fileRulePath)
+
+ watcher, err := fsnotify.NewWatcher()
+ if err != nil {
+ log.Fatalf("Unable create new fsnotify watcher: %v", err)
+ }
+ servenv.OnTerm(func() { watcher.Close() })
+
+ go func(tsc tabletserver.Controller) {
+ for {
+ select {
+ case evt, ok := <-watcher.Events:
+ if !ok {
+ return
+ }
+ if path.Base(evt.Name) != ruleFileName {
+ continue
+ }
+ if err := fileCustomRule.Open(tsc, *fileRulePath); err != nil {
+ log.Infof("Failed to load custom rules from %q: %v", *fileRulePath, err)
+ } else {
+ log.Infof("Loaded custom rules from %q", *fileRulePath)
+ }
+ case err, ok := <-watcher.Errors:
+ if !ok {
+ return
+ }
+ log.Errorf("Error watching %v: %v", *fileRulePath, err)
+ }
+ }
+ }(qsc)
+
+ if err = watcher.Add(baseDir); err != nil {
+ log.Fatalf("Unable to set up watcher for %v + %v: %v", baseDir, ruleFileName, err)
+ }
+ }
  }
 }
 

go/vt/vttablet/tabletserver/query_executor.go

@@ -382,7 +382,7 @@ func (qre *QueryExecutor) checkPermissions() error {
  remoteAddr = ci.RemoteAddr()
  username = ci.Username()
  }
- action, desc := qre.plan.Rules.GetAction(remoteAddr, username, qre.bindVars)
+ action, desc := qre.plan.Rules.GetAction(remoteAddr, username, qre.bindVars, qre.marginComments)
  switch action {
  case rules.QRFail:
  return vterrors.Errorf(vtrpcpb.Code_INVALID_ARGUMENT, "disallowed due to rule: %s", desc)

go/vt/vttablet/tabletserver/rules/rules.go

@@ -27,6 +27,7 @@ import (
  "vitess.io/vitess/go/vt/vtgate/evalengine"
 
  "vitess.io/vitess/go/sqltypes"
+ "vitess.io/vitess/go/vt/sqlparser"
  "vitess.io/vitess/go/vt/vterrors"
  "vitess.io/vitess/go/vt/vttablet/tabletserver/planbuilder"
 
@@ -166,9 +167,14 @@ func (qrs *Rules) FilterByPlan(query string, planid planbuilder.PlanType, tableN
 }
 
 // GetAction runs the input against the rules engine and returns the action to be performed.
-func (qrs *Rules) GetAction(ip, user string, bindVars map[string]*querypb.BindVariable) (action Action, desc string) {
+func (qrs *Rules) GetAction(
+ ip,
+ user string,
+ bindVars map[string]*querypb.BindVariable,
+ marginComments sqlparser.MarginComments,
+) (action Action, desc string) {
  for _, qr := range qrs.rules {
- if act := qr.GetAction(ip, user, bindVars); act != QRContinue {
+ if act := qr.GetAction(ip, user, bindVars, marginComments); act != QRContinue {
  return act, qr.Description
  }
  }
@@ -192,7 +198,7 @@ type Rule struct {
  // All defined conditions must match for the rule to fire (AND).
 
  // Regexp conditions. nil conditions are ignored (TRUE).
- requestIP, user, query namedRegexp
+ requestIP, user, query, leadingComment, trailingComment namedRegexp
 
  // Any matched plan will make this condition true (OR)
  plans []planbuilder.PlanType
@@ -241,6 +247,8 @@ func (qr *Rule) Equal(other *Rule) bool {
  qr.requestIP.Equal(other.requestIP) &&
  qr.user.Equal(other.user) &&
  qr.query.Equal(other.query) &&
+ qr.leadingComment.Equal(other.leadingComment) &&
+ qr.trailingComment.Equal(other.trailingComment) &&
  reflect.DeepEqual(qr.plans, other.plans) &&
  reflect.DeepEqual(qr.tableNames, other.tableNames) &&
  reflect.DeepEqual(qr.bindVarConds, other.bindVarConds) &&
@@ -250,12 +258,14 @@ func (qr *Rule) Equal(other *Rule) bool {
 // Copy performs a deep copy of a Rule.
 func (qr *Rule) Copy() (newqr *Rule) {
  newqr = &Rule{
- Description: qr.Description,
- Name: qr.Name,
- requestIP: qr.requestIP,
- user: qr.user,
- query: qr.query,
- act: qr.act,
+ Description: qr.Description,
+ Name: qr.Name,
+ requestIP: qr.requestIP,
+ user: qr.user,
+ query: qr.query,
+ leadingComment: qr.leadingComment,
+ trailingComment: qr.trailingComment,
+ act: qr.act,
  }
  if qr.plans != nil {
  newqr.plans = make([]planbuilder.PlanType, len(qr.plans))
@@ -286,6 +296,12 @@ func (qr *Rule) MarshalJSON() ([]byte, error) {
  if qr.query.Regexp != nil {
  safeEncode(b, `,"Query":`, qr.query)
  }
+ if qr.leadingComment.Regexp != nil {
+ safeEncode(b, `,"LeadingComment":`, qr.leadingComment)
+ }
+ if qr.trailingComment.Regexp != nil {
+ safeEncode(b, `,"TrailingComment":`, qr.trailingComment)
+ }
  if qr.plans != nil {
  safeEncode(b, `,"Plans":`, qr.plans)
  }
@@ -339,6 +355,20 @@ func (qr *Rule) SetQueryCond(pattern string) (err error) {
  return
 }
 
+// SetLeadingCommentCond adds a regular expression condition for a leading query comment.
+func (qr *Rule) SetLeadingCommentCond(pattern string) (err error) {
+ qr.leadingComment.name = pattern
+ qr.leadingComment.Regexp, err = regexp.Compile(makeExact(pattern))
+ return
+}
+
+// SetTrailingCommentCond adds a regular expression condition for a trailing query comment.
+func (qr *Rule) SetTrailingCommentCond(pattern string) (err error) {
+ qr.trailingComment.name = pattern
+ qr.trailingComment.Regexp, err = regexp.Compile(makeExact(pattern))
+ return
+}
+
 // makeExact forces a full string match for the regex instead of substring
 func makeExact(pattern string) string {
  return fmt.Sprintf("^%s$", pattern)
@@ -418,13 +448,26 @@ func (qr *Rule) FilterByPlan(query string, planid planbuilder.PlanType, tableNam
  }
  newqr = qr.Copy()
  newqr.query = namedRegexp{}
+ // Note we explicitly don't remove the leading/trailing comments as they
+ // must be evaluated at execution time.
  newqr.plans = nil
  newqr.tableNames = nil
  return newqr
 }
 
 // GetAction returns the action for a single rule.
-func (qr *Rule) GetAction(ip, user string, bindVars map[string]*querypb.BindVariable) Action {
+func (qr *Rule) GetAction(
+ ip,
+ user string,
+ bindVars map[string]*querypb.BindVariable,
+ marginComments sqlparser.MarginComments,
+) Action {
+ if !reMatch(qr.leadingComment.Regexp, marginComments.Leading) {
+ return QRContinue
+ }
+ if !reMatch(qr.trailingComment.Regexp, marginComments.Trailing) {
+ return QRContinue
+ }
  if !reMatch(qr.requestIP.Regexp, ip) {
  return QRContinue
  }
@@ -791,7 +834,7 @@ func BuildQueryRule(ruleInfo map[string]interface{}) (qr *Rule, err error) {
  var lv []interface{}
  var ok bool
  switch k {
- case "Name", "Description", "RequestIP", "User", "Query", "Action":
+ case "Name", "Description", "RequestIP", "User", "Query", "Action", "LeadingComment", "TrailingComment":
  sv, ok = v.(string)
  if !ok {
  return nil, vterrors.Errorf(vtrpcpb.Code_INVALID_ARGUMENT, "want string for %s", k)
@@ -824,6 +867,16 @@ func BuildQueryRule(ruleInfo map[string]interface{}) (qr *Rule, err error) {
  if err != nil {
  return nil, vterrors.Errorf(vtrpcpb.Code_INVALID_ARGUMENT, "could not set Query condition: %v", sv)
  }
+ case "LeadingComment":
+ err = qr.SetLeadingCommentCond(sv)
+ if err != nil {
+ return nil, vterrors.Errorf(vtrpcpb.Code_INVALID_ARGUMENT, "could not set LeadingComment condition: %v", sv)
+ }
+ case "TrailingComment":
+ err = qr.SetTrailingCommentCond(sv)
+ if err != nil {
+ return nil, vterrors.Errorf(vtrpcpb.Code_INVALID_ARGUMENT, "could not set TrailingComment condition: %v", sv)
+ }
  case "Plans":
  for _, p := range lv {
  pv, ok := p.(string)