[Prototype] More functional query killing 🎁

go.mod

@@ -22,6 +22,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/falun/watch v0.0.0-20200621174039-43c45ed8922d
 	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

@@ -151,6 +151,8 @@ 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/falun/watch v0.0.0-20200621174039-43c45ed8922d h1:PhAfeZuGr4z11ushRkiy4JD+gskpF2NihVumlAvSNr4=
+github.com/falun/watch v0.0.0-20200621174039-43c45ed8922d/go.mod h1:PXtPcVJjFbxuU5Yxv5F1q2ZmsGHnbMaML6TuCAX9TOo=
 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=

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

@@ -22,6 +22,8 @@ import (
 	"io/ioutil"
 	"time"
 
+	"github.com/falun/watch"
+
 	"vitess.io/vitess/go/vt/log"
 	"vitess.io/vitess/go/vt/vttablet/tabletserver"
 	"vitess.io/vitess/go/vt/vttablet/tabletserver/rules"
@@ -32,6 +34,8 @@ var (
 	fileCustomRule = NewFileCustomRule()
 	// Commandline flag to specify rule path
 	fileRulePath = flag.String("filecustomrules", "", "file based custom rule path")
+
+	fileRulePollInterval = flag.Duration("filecustomrules_interval", 0, "How often should we poll the rule file <= 0 means we will not poll")
 )
 
 // FileCustomRule is an implementation of CustomRuleManager, it reads custom query
@@ -90,6 +94,20 @@ func ActivateFileCustomRules(qsc tabletserver.Controller) {
 	if *fileRulePath != "" {
 		qsc.RegisterQueryRuleSource(FileCustomRuleSource)
 		fileCustomRule.Open(qsc, *fileRulePath)
+		if *fileRulePollInterval > time.Duration(0) {
+			w := watch.File(*fileRulePath)
+			// this returns a cancelFn that for cleanliness we would register to run on server shut down
+			ch, _ := w.OnInterval(*fileRulePollInterval)
+			go func(tsc tabletserver.Controller) {
+				for range ch {
+					if err := fileCustomRule.Open(tsc, *fileRulePath); err != nil {
+						log.Infof("Failed to load fileCustomRule %q: %v", *fileRulePath, err)
+					} else {
+						log.Infof("Opened %q", *fileRulePath)
+					}
+				}
+			}(qsc)
+		}
 	}
 }
 

go/vt/vttablet/tabletserver/query_executor.go

@@ -309,7 +309,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,8 @@ import (
 	"vitess.io/vitess/go/vt/vtgate/evalengine"
 
 	"vitess.io/vitess/go/sqltypes"
+	"vitess.io/vitess/go/vt/log"
+	"vitess.io/vitess/go/vt/sqlparser"
 	"vitess.io/vitess/go/vt/vterrors"
 	"vitess.io/vitess/go/vt/vttablet/tabletserver/planbuilder"
 
@@ -156,9 +158,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
 		}
 	}
@@ -182,7 +189,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
@@ -231,6 +238,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) &&
@@ -240,12 +249,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))
@@ -276,6 +287,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.query)
+	}
+	if qr.trailingComment.Regexp != nil {
+		safeEncode(b, `,"TrailingComment":`, qr.query)
+	}
 	if qr.plans != nil {
 		safeEncode(b, `,"Plans":`, qr.plans)
 	}
@@ -329,6 +346,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)
@@ -397,24 +428,46 @@ Error:
 // than the plan and query. If the plan and query don't match the Rule,
 // then it returns nil.
 func (qr *Rule) FilterByPlan(query string, planid planbuilder.PlanType, tableName string) (newqr *Rule) {
+	log.Infof("Running FilterByPlan on %v / %v", qr.Name, qr.Description)
+	log.Infof("  query: %q", query)
 	if !reMatch(qr.query.Regexp, query) {
+		log.Infof("bail: query")
 		return nil
 	}
 	if !planMatch(qr.plans, planid) {
+		log.Infof("bail: plan")
 		return nil
 	}
 	if !tableMatch(qr.tableNames, tableName) {
+		log.Infof("bail: plan")
 		return nil
 	}
 	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
+	log.Infof("Passed FilterByPlan")
 	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 {
+	// check to see if we need to extract margin comments
+	if qr.leadingComment.Regexp != nil || qr.trailingComment.Regexp != nil {
+		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
 	}
@@ -781,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)
@@ -814,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)