// loadKeyspaceAndBlacklistRules does what the name suggests:
// 1. load and build keyrange query rules
// 2. load and build blacklist query rules
func (agent *ActionAgent) loadKeyspaceAndBlacklistRules(tablet *pbt.Tablet, blacklistedTables []string) (err error) {
// Keyrange rules
keyrangeRules := tabletserver.NewQueryRules()
if key.KeyRangeIsPartial(tablet.KeyRange) {
log.Infof("Restricting to keyrange: %v", tablet.KeyRange)
dmlPlans := []struct {
planID planbuilder.PlanType
onAbsent bool
}{
{planbuilder.PlanInsertPK, true},
{planbuilder.PlanInsertSubquery, true},
{planbuilder.PlanPassDML, false},
{planbuilder.PlanDMLPK, false},
{planbuilder.PlanDMLSubquery, false},
{planbuilder.PlanUpsertPK, false},
}
for _, plan := range dmlPlans {
qr := tabletserver.NewQueryRule(
fmt.Sprintf("enforce keyspace_id range for %v", plan.planID),
fmt.Sprintf("keyspace_id_not_in_range_%v", plan.planID),
tabletserver.QRFail,
)
qr.AddPlanCond(plan.planID)
err := qr.AddBindVarCond("keyspace_id", plan.onAbsent, true, tabletserver.QRNotIn, tablet.KeyRange)
if err != nil {
return fmt.Errorf("Unable to add keyspace rule: %v", err)
}
keyrangeRules.Add(qr)
}
}
// Blacklisted tables
blacklistRules := tabletserver.NewQueryRules()
if len(blacklistedTables) > 0 {
// tables, first resolve wildcards
tables, err := mysqlctl.ResolveTables(agent.MysqlDaemon, topoproto.TabletDbName(tablet), blacklistedTables)
if err != nil {
return err
}
log.Infof("Blacklisting tables %v", strings.Join(tables, ", "))
qr := tabletserver.NewQueryRule("enforce blacklisted tables", "blacklisted_table", tabletserver.QRFailRetry)
for _, t := range tables {
qr.AddTableCond(t)
}
blacklistRules.Add(qr)
}
// Push all three sets of QueryRules to TabletServerRpcService
loadRuleErr := agent.QueryServiceControl.SetQueryRules(keyrangeQueryRules, keyrangeRules)
if loadRuleErr != nil {
log.Warningf("Fail to load query rule set %s: %s", keyrangeQueryRules, loadRuleErr)
}
loadRuleErr = agent.QueryServiceControl.SetQueryRules(blacklistQueryRules, blacklistRules)
if loadRuleErr != nil {
log.Warningf("Fail to load query rule set %s: %s", blacklistQueryRules, loadRuleErr)
}
return nil
}
// NewZkCustomRule Creates new ZkCustomRule structure
func NewZkCustomRule(zkconn zk.Conn) *ZkCustomRule {
return &ZkCustomRule{
zconn: zkconn,
currentRuleSet: tabletserver.NewQueryRules(),
currentRuleSetVersion: InvalidQueryRulesVersion,
finish: make(chan int, 1)}
}
// loadBlacklistRules loads and builds the blacklist query rules
func (agent *ActionAgent) loadBlacklistRules(tablet *topodatapb.Tablet, blacklistedTables []string) (err error) {
blacklistRules := tabletserver.NewQueryRules()
if len(blacklistedTables) > 0 {
// tables, first resolve wildcards
tables, err := mysqlctl.ResolveTables(agent.MysqlDaemon, topoproto.TabletDbName(tablet), blacklistedTables)
if err != nil {
return err
}
// Verify that at least one table matches the wildcards, so
// that we don't add a rule to blacklist all tables
if len(tables) > 0 {
log.Infof("Blacklisting tables %v", strings.Join(tables, ", "))
qr := tabletserver.NewQueryRule("enforce blacklisted tables", "blacklisted_table", tabletserver.QRFailRetry)
for _, t := range tables {
qr.AddTableCond(t)
}
blacklistRules.Add(qr)
}
}
loadRuleErr := agent.QueryServiceControl.SetQueryRules(blacklistQueryRules, blacklistRules)
if loadRuleErr != nil {
log.Warningf("Fail to load query rule set %s: %s", blacklistQueryRules, loadRuleErr)
}
return nil
}
// Open try to build query rules from local file and push the rules to vttablet
func (fcr *FileCustomRule) Open(qsc tabletserver.QueryServiceControl, rulePath string) error {
fcr.path = rulePath
if fcr.path == "" {
// Don't go further if path is empty
return nil
}
data, err := ioutil.ReadFile(fcr.path)
if err != nil {
log.Warningf("Error reading file %v: %v", fcr.path, err)
// Don't update any internal cache, just return error
return err
}
qrs := tabletserver.NewQueryRules()
err = qrs.UnmarshalJSON(data)
if err != nil {
log.Warningf("Error unmarshaling query rules %v", err)
return err
}
fcr.currentRuleSetTimestamp = time.Now().Unix()
fcr.currentRuleSet = qrs.Copy()
// Push query rules to vttablet
qsc.SetQueryRules(FileCustomRuleSource, qrs.Copy())
log.Infof("Custom rule loaded from file: %s", fcr.path)
return nil
}
func loadCustomRules(customrules string) *ts.QueryRules {
if customrules == "" {
return ts.NewQueryRules()
}
data, err := ioutil.ReadFile(customrules)
if err != nil {
relog.Fatal("Error reading file %v: %v", customrules, err)
}
qrs := ts.NewQueryRules()
err = qrs.UnmarshalJSON(data)
if err != nil {
relog.Fatal("Error unmarshaling query rules %v", err)
}
return qrs
}
// initializeKeyRangeRule will create and set the key range rules
func (agent *ActionAgent) initializeKeyRangeRule(ctx context.Context, keyspace string, keyRange *topodatapb.KeyRange) error {
// check we have a partial key range
if !key.KeyRangeIsPartial(keyRange) {
log.Infof("Tablet covers the full KeyRange, not adding KeyRange query rule")
return nil
}
// read the keyspace to get the sharding column name
keyspaceInfo, err := agent.TopoServer.GetKeyspace(ctx, keyspace)
if err != nil {
return fmt.Errorf("cannot read keyspace %v to get sharding key: %v", keyspace, err)
}
if keyspaceInfo.ShardingColumnName == "" {
log.Infof("Keyspace %v has an empty ShardingColumnName, not adding KeyRange query rule", keyspace)
return nil
}
// create the rules
log.Infof("Restricting to keyrange: %v", key.KeyRangeString(keyRange))
keyrangeRules := tabletserver.NewQueryRules()
dmlPlans := []struct {
planID planbuilder.PlanType
onAbsent bool
}{
{planbuilder.PlanInsertPK, true},
{planbuilder.PlanInsertSubquery, true},
{planbuilder.PlanPassDML, false},
{planbuilder.PlanDMLPK, false},
{planbuilder.PlanDMLSubquery, false},
{planbuilder.PlanUpsertPK, false},
}
for _, plan := range dmlPlans {
qr := tabletserver.NewQueryRule(
fmt.Sprintf("enforce %v range for %v", keyspaceInfo.ShardingColumnName, plan.planID),
fmt.Sprintf("%v_not_in_range_%v", keyspaceInfo.ShardingColumnName, plan.planID),
tabletserver.QRFail,
)
qr.AddPlanCond(plan.planID)
err := qr.AddBindVarCond(keyspaceInfo.ShardingColumnName, plan.onAbsent, true, tabletserver.QRNotIn, keyRange)
if err != nil {
return fmt.Errorf("Unable to add key range rule: %v", err)
}
keyrangeRules.Add(qr)
}
// and load them
agent.QueryServiceControl.RegisterQueryRuleSource(keyrangeQueryRules)
if err := agent.QueryServiceControl.SetQueryRules(keyrangeQueryRules, keyrangeRules); err != nil {
return fmt.Errorf("failed to load query rule set %s: %s", keyrangeQueryRules, err)
}
return nil
}
func TestQueryRules(t *testing.T) {
rules := tabletserver.NewQueryRules()
err := rules.UnmarshalJSON(rulesJSON)
if err != nil {
t.Error(err)
return
}
err = framework.Server.SetQueryRules("endtoend", rules)
want := "Rule source identifier endtoend is not valid"
if err == nil || err.Error() != want {
t.Errorf("Error: %v, want %s", err, want)
}
framework.Server.RegisterQueryRuleSource("endtoend")
defer framework.Server.UnRegisterQueryRuleSource("endtoend")
err = framework.Server.SetQueryRules("endtoend", rules)
if err != nil {
t.Error(err)
return
}
client := framework.NewClient()
query := "select * from vitess_test where intval=:asdfg"
bv := map[string]interface{}{"asdfg": 1}
_, err = client.Execute(query, bv)
want = "error: Query disallowed due to rule: disallow bindvar 'asdfg'"
if err == nil || err.Error() != want {
t.Errorf("Error: %v, want %s", err, want)
}
_, err = client.StreamExecute(query, bv)
want = "error: Query disallowed due to rule: disallow bindvar 'asdfg'"
if err == nil || err.Error() != want {
t.Errorf("Error: %v, want %s", err, want)
}
err = framework.Server.SetQueryRules("endtoend", nil)
if err != nil {
t.Error(err)
return
}
_, err = client.Execute(query, bv)
if err != nil {
t.Error(err)
return
}
}
// refreshData gets query rules from Zookeeper and refresh internal QueryRules cache
// this function will also call TabletServer.SetQueryRules to propagate rule changes to query service
func (zkcr *ZkCustomRule) refreshData(qsc tabletserver.QueryServiceControl, nodeRemoval bool) error {
data, stat, err := zkcr.zconn.Get(zkcr.path)
zkcr.mu.Lock()
defer zkcr.mu.Unlock()
if err == nil {
qrs := tabletserver.NewQueryRules()
if !nodeRemoval {
err = qrs.UnmarshalJSON([]byte(data))
if err != nil {
log.Warningf("Error unmarshaling query rules %v, original data '%s'", err, data)
return nil
}
}
zkcr.currentRuleSetVersion = stat.Mzxid()
if !reflect.DeepEqual(zkcr.currentRuleSet, qrs) {
zkcr.currentRuleSet = qrs.Copy()
qsc.SetQueryRules(ZkCustomRuleSource, qrs.Copy())
log.Infof("Custom rule version %v fetched from Zookeeper and applied to vttablet", zkcr.currentRuleSetVersion)
}
return nil
}
log.Warningf("Error encountered when trying to get data and watch from Zk: %v", err)
return err
}
func TestZkCustomRule(t *testing.T) {
tqsc := tabletservermock.NewController()
setUpFakeZk(t)
zkcr := NewZkCustomRule(conn)
err := zkcr.Open(tqsc, "/zk/fake/customrules/testrules")
if err != nil {
t.Fatalf("Cannot open zookeeper custom rule service, err=%v", err)
}
var qrs *tabletserver.QueryRules
// Test if we can successfully fetch the original rule (test GetRules)
qrs, _, err = zkcr.GetRules()
if err != nil {
t.Fatalf("GetRules of ZkCustomRule should always return nil error, but we receive %v", err)
}
qr := qrs.Find("r1")
if qr == nil {
t.Fatalf("Expect custom rule r1 to be found, but got nothing, qrs=%v", qrs)
}
// Test updating rules
conn.Set("/zk/fake/customrules/testrules", customRule2, -1)
<-time.After(time.Second) //Wait for the polling thread to respond
qrs, _, err = zkcr.GetRules()
if err != nil {
t.Fatalf("GetRules of ZkCustomRule should always return nil error, but we receive %v", err)
}
qr = qrs.Find("r2")
if qr == nil {
t.Fatalf("Expect custom rule r2 to be found, but got nothing, qrs=%v", qrs)
}
qr = qrs.Find("r1")
if qr != nil {
t.Fatalf("Custom rule r1 should not be found after r2 is set")
}
// Test rule path removal
conn.Delete("/zk/fake/customrules/testrules", -1)
<-time.After(time.Second)
qrs, _, err = zkcr.GetRules()
if err != nil {
t.Fatalf("GetRules of ZkCustomRule should always return nil error, but we receive %v", err)
}
if reflect.DeepEqual(qrs, tabletserver.NewQueryRules()) {
t.Fatalf("Expect empty rule at this point")
}
// Test rule path revival
conn.Create("/zk/fake/customrules/testrules", "customrule2", 0, zookeeper.WorldACL(zookeeper.PERM_ALL))
conn.Set("/zk/fake/customrules/testrules", customRule2, -1)
<-time.After(time.Second) //Wait for the polling thread to respond
qrs, _, err = zkcr.GetRules()
if err != nil {
t.Fatalf("GetRules of ZkCustomRule should always return nil error, but we receive %v", err)
}
qr = qrs.Find("r2")
if qr == nil {
t.Fatalf("Expect custom rule r2 to be found, but got nothing, qrs=%v", qrs)
}
zkcr.Close()
}
// NewFileCustomRule returns pointer to new FileCustomRule structure
func NewFileCustomRule() (fcr *FileCustomRule) {
fcr = new(FileCustomRule)
fcr.path = ""
fcr.currentRuleSet = tabletserver.NewQueryRules()
return fcr
}
func TestQueryRules(t *testing.T) {
rules := tabletserver.NewQueryRules()
err := rules.UnmarshalJSON(rulesJSON)
if err != nil {
t.Error(err)
return
}
err = framework.Server.SetQueryRules("endtoend", rules)
want := "Rule source identifier endtoend is not valid"
if err == nil || err.Error() != want {
t.Errorf("Error: %v, want %s", err, want)
}
framework.Server.RegisterQueryRuleSource("endtoend")
defer framework.Server.UnRegisterQueryRuleSource("endtoend")
err = framework.Server.SetQueryRules("endtoend", rules)
if err != nil {
t.Error(err)
return
}
rulesJSON := compacted(framework.FetchURL("/debug/query_rules"))
want = compacted(`{
"endtoend":[{
"Description": "disallow bindvar 'asdfg'",
"Name": "r1",
"BindVarConds":[{
"Name": "asdfg",
"OnAbsent": false,
"Operator": ""
}],
"Action": "FAIL"
}]
}`)
if rulesJSON != want {
t.Errorf("/debug/query_rules:\n%v, want\n%s", rulesJSON, want)
}
client := framework.NewClient()
query := "select * from vitess_test where intval=:asdfg"
bv := map[string]interface{}{"asdfg": 1}
_, err = client.Execute(query, bv)
want = "error: Query disallowed due to rule: disallow bindvar 'asdfg'"
if err == nil || err.Error() != want {
t.Errorf("Error: %v, want %s", err, want)
}
_, err = client.StreamExecute(query, bv)
want = "error: Query disallowed due to rule: disallow bindvar 'asdfg'"
if err == nil || err.Error() != want {
t.Errorf("Error: %v, want %s", err, want)
}
err = framework.Server.SetQueryRules("endtoend", nil)
if err != nil {
t.Error(err)
return
}
_, err = client.Execute(query, bv)
if err != nil {
t.Error(err)
return
}
}