Golang tabletserver.NewQueryRules函数代码示例

// 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
	}
}