<commit_before>// Copyright 2014, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package tabletserver

import (
	"errors"
	"sync"

	"github.com/youtube/vitess/go/vt/tabletserver/planbuilder"
)

// QueryRuleInfo is the maintainer of QueryRules from multiple sources
type QueryRuleInfo struct {
	mu            sync.Mutex
	queryRulesMap map[string]*QueryRules
}

// Names for QueryRules coming from different sources
// QueryRules from keyrange
const KeyrangeQueryRules string = "KEYRANGE_QUERY_RULES"

// QueryRules from blacklist
const BlacklistQueryRules string = "BLACKLIST_QUERY_RULES"

// QueryRules from custom rules
const CustomQueryRules string = "CUSTOM_QUERY_RULES"

func NewQueryRuleInfo() *QueryRuleInfo {
	qri := &QueryRuleInfo{
		queryRulesMap: map[string]*QueryRules{
			KeyrangeQueryRules:  NewQueryRules(),
			BlacklistQueryRules: NewQueryRules(),
			CustomQueryRules:    NewQueryRules(),
		},
	}
	return qri
}

// SetRules takes a external QueryRules structure and overwrite one of the
// internal QueryRules as designated by queryRuleSet parameter
func (qri *QueryRuleInfo) SetRules(queryRuleSet string, newRules *QueryRules) error {
	qri.mu.Lock()
	defer qri.mu.Unlock()
	if _, ok := qri.queryRulesMap[queryRuleSet]; ok {
		qri.queryRulesMap[queryRuleSet] = newRules.Copy()
		return nil
	} else {
		return errors.New("QueryRules identifier " + queryRuleSet + " is not valid")
	}
}

// GetRules returns the corresponding QueryRules as designated by queryRuleSet parameter
func (qri *QueryRuleInfo) GetRules(queryRuleSet string) (error, *QueryRules) {
	qri.mu.Lock()
	defer qri.mu.Unlock()
	if ruleset, ok := qri.queryRulesMap[queryRuleSet]; ok {
		return nil, ruleset.Copy()
	} else {
		return errors.New("QueryRules identifier " + queryRuleSet + " is not valid"), nil
	}
}

// filterByPlan creates a new QueryRules by prefiltering on all query rules that are contained in internal
// QueryRules structures, in other words, query rules from all predefined sources will be applied
func (qri *QueryRuleInfo) filterByPlan(query string, planid planbuilder.PlanType, tableName string) (newqrs *QueryRules) {
	newqrs = NewQueryRules()
	newqrs.Append(qri.queryRulesMap[KeyrangeQueryRules].filterByPlan(query, planid, tableName))
	newqrs.Append(qri.queryRulesMap[BlacklistQueryRules].filterByPlan(query, planid, tableName))
	newqrs.Append(qri.queryRulesMap[CustomQueryRules].filterByPlan(query, planid, tableName))
	return newqrs
}<commit_msg>jeffjiang: fix code styles and add sanity check for QueryRuleInfo.SetRules<commit_after>// Copyright 2014, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package tabletserver

import (
	"errors"
	"sync"

	"github.com/youtube/vitess/go/vt/tabletserver/planbuilder"
)

// QueryRuleInfo is the maintainer of QueryRules from multiple sources
type QueryRuleInfo struct {
	mu            sync.Mutex
	queryRulesMap map[string]*QueryRules
}

// Names for QueryRules coming from different sources
// QueryRules from keyrange
const KeyrangeQueryRules string = "KEYRANGE_QUERY_RULES"

// QueryRules from blacklist
const BlacklistQueryRules string = "BLACKLIST_QUERY_RULES"

// QueryRules from custom rules
const CustomQueryRules string = "CUSTOM_QUERY_RULES"

func NewQueryRuleInfo() *QueryRuleInfo {
	qri := &QueryRuleInfo{
		queryRulesMap: map[string]*QueryRules{
			KeyrangeQueryRules:  NewQueryRules(),
			BlacklistQueryRules: NewQueryRules(),
			CustomQueryRules:    NewQueryRules(),
		},
	}
	return qri
}

// SetRules takes a external QueryRules structure and overwrite one of the
// internal QueryRules as designated by queryRuleSet parameter
func (qri *QueryRuleInfo) SetRules(queryRuleSet string, newRules *QueryRules) error {
	if newRules == nil {
		newRules = NewQueryRules()
	}
	qri.mu.Lock()
	defer qri.mu.Unlock()
	if _, ok := qri.queryRulesMap[queryRuleSet]; ok {
		qri.queryRulesMap[queryRuleSet] = newRules.Copy()
		return nil
	}
	return errors.New("QueryRules identifier " + queryRuleSet + " is not valid")
}

// GetRules returns the corresponding QueryRules as designated by queryRuleSet parameter
func (qri *QueryRuleInfo) GetRules(queryRuleSet string) (error, *QueryRules) {
	qri.mu.Lock()
	defer qri.mu.Unlock()
	if ruleset, ok := qri.queryRulesMap[queryRuleSet]; ok {
		return nil, ruleset.Copy()
	}
	return errors.New("QueryRules identifier " + queryRuleSet + " is not valid"), nil
}

// filterByPlan creates a new QueryRules by prefiltering on all query rules that are contained in internal
// QueryRules structures, in other words, query rules from all predefined sources will be applied
func (qri *QueryRuleInfo) filterByPlan(query string, planid planbuilder.PlanType, tableName string) (newqrs *QueryRules) {
	newqrs = NewQueryRules()
	newqrs.Append(qri.queryRulesMap[KeyrangeQueryRules].filterByPlan(query, planid, tableName))
	newqrs.Append(qri.queryRulesMap[BlacklistQueryRules].filterByPlan(query, planid, tableName))
	newqrs.Append(qri.queryRulesMap[CustomQueryRules].filterByPlan(query, planid, tableName))
	return newqrs
}
{'fraction_non_alphanumeric': 0.0905580227843213, 'fraction_numerical': 0.0015446997489862908, 'mean_word_length': 5.1520190023752965, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '653784', 'n_tokens_mistral': 1615, 'n_tokens_neox': 1574, 'n_words': 470}