diff --git a/go.mod b/go.mod index f1b2d8186d2..cd9a7eadab0 100644 --- a/go.mod +++ b/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 diff --git a/go.sum b/go.sum index c700b4a7578..da382725e13 100644 --- a/go.sum +++ b/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= diff --git a/go/vt/vttablet/customrule/filecustomrule/filecustomrule.go b/go/vt/vttablet/customrule/filecustomrule/filecustomrule.go index caa35a9a998..306d72f44d6 100644 --- a/go/vt/vttablet/customrule/filecustomrule/filecustomrule.go +++ b/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) + ch, cancelFn := w.OnInterval(*fileRulePollInterval) + servenv.OnTerm(cancelFn) + 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) + } } } diff --git a/go/vt/vttablet/tabletserver/query_executor.go b/go/vt/vttablet/tabletserver/query_executor.go index 70c363fe3b4..74b608727ba 100644 --- a/go/vt/vttablet/tabletserver/query_executor.go +++ b/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) diff --git a/go/vt/vttablet/tabletserver/rules/rules.go b/go/vt/vttablet/tabletserver/rules/rules.go index 8dd25ef7450..3fe7bc636af 100644 --- a/go/vt/vttablet/tabletserver/rules/rules.go +++ b/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) diff --git a/go/vt/vttablet/tabletserver/rules/rules_test.go b/go/vt/vttablet/tabletserver/rules/rules_test.go index 3a25ce74fca..4d7bc62c331 100644 --- a/go/vt/vttablet/tabletserver/rules/rules_test.go +++ b/go/vt/vttablet/tabletserver/rules/rules_test.go @@ -25,6 +25,7 @@ import ( "testing" "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" @@ -496,31 +497,68 @@ func TestAction(t *testing.T) { bv := make(map[string]*querypb.BindVariable) bv["a"] = sqltypes.Uint64BindVariable(0) - action, desc := qrs.GetAction("123", "user1", bv) + + mc := sqlparser.MarginComments{ + Leading: "some comments leading the query", + Trailing: "other trailing comments", + } + + action, desc := qrs.GetAction("123", "user1", bv, mc) if action != QRFail { t.Errorf("want fail") } if desc != "rule 1" { t.Errorf("want rule 1, got %s", desc) } - action, desc = qrs.GetAction("1234", "user", bv) + action, desc = qrs.GetAction("1234", "user", bv, mc) if action != QRFailRetry { t.Errorf("want fail_retry") } if desc != "rule 2" { t.Errorf("want rule 2, got %s", desc) } - action, _ = qrs.GetAction("1234", "user1", bv) + action, _ = qrs.GetAction("1234", "user1", bv, mc) if action != QRContinue { t.Errorf("want continue") } + bv["a"] = sqltypes.Uint64BindVariable(1) - action, desc = qrs.GetAction("1234", "user1", bv) + action, desc = qrs.GetAction("1234", "user1", bv, mc) if action != QRFail { t.Errorf("want fail") } if desc != "rule 3" { - t.Errorf("want rule 2, got %s", desc) + t.Errorf("want rule 3, got %s", desc) + } + + // reset bound variable 'a' to 0 so it doesn't match rule 3 + bv["a"] = sqltypes.Uint64BindVariable(0) + + qr4 := NewQueryRule("rule 4", "r4", QRFail) + qr4.SetTrailingCommentCond(".*trailing.*") + + newQrs := qrs.Copy() + newQrs.Add(qr4) + + action, desc = newQrs.GetAction("1234", "user1", bv, mc) + if action != QRFail { + t.Errorf("want fail") + } + if desc != "rule 4" { + t.Errorf("want rule 4, got %s", desc) + } + + qr5 := NewQueryRule("rule 5", "r4", QRFail) + qr5.SetTrailingCommentCond(".*trailing.*") + + newQrs = qrs.Copy() + newQrs.Add(qr5) + action, desc = newQrs.GetAction("1234", "user1", bv, mc) + if action != QRFail { + t.Errorf("want fail") + } + if desc != "rule 5" { + t.Errorf("want rule 5, got %s", desc) } }