From 9803c86b9d7b4ca5fb11a4513824fdc66b63cf83 Mon Sep 17 00:00:00 2001 From: Richard Bailey Date: Wed, 24 Mar 2021 16:02:45 -0700 Subject: [PATCH 1/3] leading+trailing comment regex + file watch --- go.mod | 1 + go.sum | 2 + .../filecustomrule/filecustomrule.go | 21 ++++++++- go/vt/vttablet/tabletserver/rules/rules.go | 45 ++++++++++++++++++- 4 files changed, 65 insertions(+), 4 deletions(-) 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..a29b686f599 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 @@ -88,8 +92,21 @@ func (fcr *FileCustomRule) GetRules() (qrs *rules.Rules, version int64, err erro // ActivateFileCustomRules activates this static file based custom rule mechanism func ActivateFileCustomRules(qsc tabletserver.Controller) { if *fileRulePath != "" { - qsc.RegisterQueryRuleSource(FileCustomRuleSource) - fileCustomRule.Open(qsc, *fileRulePath) + if *fileRulePollInterval <= time.Duration(0) { + qsc.RegisterQueryRuleSource(FileCustomRuleSource) + fileCustomRule.Open(qsc, *fileRulePath) + } else { + 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) + } + } + }(qsc) + } } } diff --git a/go/vt/vttablet/tabletserver/rules/rules.go b/go/vt/vttablet/tabletserver/rules/rules.go index 8dd25ef7450..763350a7adc 100644 --- a/go/vt/vttablet/tabletserver/rules/rules.go +++ b/go/vt/vttablet/tabletserver/rules/rules.go @@ -24,6 +24,7 @@ import ( "regexp" "strconv" + "vitess.io/vitess/go/vt/sqlparser" "vitess.io/vitess/go/vt/vtgate/evalengine" "vitess.io/vitess/go/sqltypes" @@ -182,7 +183,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 @@ -276,6 +277,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 +336,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) @@ -406,6 +427,16 @@ func (qr *Rule) FilterByPlan(query string, planid planbuilder.PlanType, tableNam if !tableMatch(qr.tableNames, tableName) { return nil } + // check to see if we need to extract margin comments + if qr.leadingComment.Regexp != nil || qr.trailingComment.Regexp != nil { + _, comments := sqlparser.SplitMarginComments(query) + if !reMatch(qr.leadingComment.Regexp, comments.Leading) { + return nil + } + if !reMatch(qr.trailingComment.Regexp, comments.Trailing) { + return nil + } + } newqr = qr.Copy() newqr.query = namedRegexp{} newqr.plans = nil @@ -781,7 +812,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 +845,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) From 54e44d29ecbb9d76ff85d5428fd9a7260aef9138 Mon Sep 17 00:00:00 2001 From: Richard Bailey Date: Thu, 25 Mar 2021 18:53:49 -0700 Subject: [PATCH 2/3] route query comments into GetAction call; let filterbyplan continue to work as expected Signed-off-by: Richard Bailey --- .../filecustomrule/filecustomrule.go | 9 +-- go/vt/vttablet/tabletserver/query_executor.go | 2 +- go/vt/vttablet/tabletserver/rules/rules.go | 62 +++++++++++++------ 3 files changed, 48 insertions(+), 25 deletions(-) diff --git a/go/vt/vttablet/customrule/filecustomrule/filecustomrule.go b/go/vt/vttablet/customrule/filecustomrule/filecustomrule.go index a29b686f599..7135fa21743 100644 --- a/go/vt/vttablet/customrule/filecustomrule/filecustomrule.go +++ b/go/vt/vttablet/customrule/filecustomrule/filecustomrule.go @@ -92,10 +92,9 @@ func (fcr *FileCustomRule) GetRules() (qrs *rules.Rules, version int64, err erro // ActivateFileCustomRules activates this static file based custom rule mechanism func ActivateFileCustomRules(qsc tabletserver.Controller) { if *fileRulePath != "" { - if *fileRulePollInterval <= time.Duration(0) { - qsc.RegisterQueryRuleSource(FileCustomRuleSource) - fileCustomRule.Open(qsc, *fileRulePath) - } else { + 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) @@ -103,6 +102,8 @@ func ActivateFileCustomRules(qsc 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 763350a7adc..3fe7bc636af 100644 --- a/go/vt/vttablet/tabletserver/rules/rules.go +++ b/go/vt/vttablet/tabletserver/rules/rules.go @@ -24,10 +24,11 @@ import ( "regexp" "strconv" - "vitess.io/vitess/go/vt/sqlparser" "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" @@ -157,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 } } @@ -232,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) && @@ -241,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)) @@ -418,34 +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 } - // check to see if we need to extract margin comments - if qr.leadingComment.Regexp != nil || qr.trailingComment.Regexp != nil { - _, comments := sqlparser.SplitMarginComments(query) - if !reMatch(qr.leadingComment.Regexp, comments.Leading) { - return nil - } - if !reMatch(qr.trailingComment.Regexp, comments.Trailing) { - 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 } From 4bb0568cf52c4028af2c2661c9c693b158c610bc Mon Sep 17 00:00:00 2001 From: Richard Bailey Date: Thu, 1 Apr 2021 01:47:22 -0700 Subject: [PATCH 3/3] fix tests, register to cancel the watch Signed-off-by: Richard Bailey --- .../filecustomrule/filecustomrule.go | 4 +- .../vttablet/tabletserver/rules/rules_test.go | 48 +++++++++++++++++-- 2 files changed, 45 insertions(+), 7 deletions(-) diff --git a/go/vt/vttablet/customrule/filecustomrule/filecustomrule.go b/go/vt/vttablet/customrule/filecustomrule/filecustomrule.go index 7135fa21743..306d72f44d6 100644 --- a/go/vt/vttablet/customrule/filecustomrule/filecustomrule.go +++ b/go/vt/vttablet/customrule/filecustomrule/filecustomrule.go @@ -96,8 +96,8 @@ func ActivateFileCustomRules(qsc tabletserver.Controller) { 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) + ch, cancelFn := w.OnInterval(*fileRulePollInterval) + servenv.OnTerm(cancelFn) go func(tsc tabletserver.Controller) { for range ch { if err := fileCustomRule.Open(tsc, *fileRulePath); err != nil { 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) } }