dev sql blacklist and whitelist (#59)

conf/namespace/test_namespace.yaml

@@ -4,11 +4,12 @@ frontend:
   allowed_dbs:
     - "test_weir_db"
   slow_sql_time: 50
-  denied_sqls:
-    - sql: "select * from test"
-      ttl: 0
-    - sql: "select count(*) from test"
-      ttl: 60
+  sql_blacklist:
+    - sql: "select * from tbl0"
+    - sql: "select * from tbl1"
+  sql_whitelist:
+    - sql: "select * from tbl2"
+    - sql: "select * from tbl3"
   denied_ips:
   idle_timeout: 3600
   users:
@@ -28,9 +29,9 @@ breaker:
   scope: "sql"
   strategies:
     - min_qps: 3
-      failure_ratethreshold: 0
+      failure_rate_threshold: 0
       failure_num: 5
       sql_timeout_ms: 2000
       open_status_duration_ms: 5000
       size: 10
-      cellInterval_ms: 1000
+      cell_interval_ms: 1000

pkg/config/marshaller_test.go

@@ -12,7 +12,6 @@ var testNamespaceConfig = Namespace{
 	Frontend: FrontendNamespace{
 		AllowedDBs:  []string{"db0", "db1"},
 		SlowSQLTime: 10,
-		DeniedSQLs:  []string{"known", "unknown"},
 		DeniedIPs:   []string{"127.0.0.0", "128.0.0.0"},
 		IdleTimeout: 10,
 		Users: []FrontendUserInfo{

pkg/config/namespace.go

@@ -8,26 +8,26 @@ type Namespace struct {
 	Breaker   BreakerInfo       `yaml:"breaker"`
 }
 
-type DeniedSqlInfo struct {
-	Sql string `yaml:"sql"`
-	Ttl int64  `yaml:"ttl"`
-}
-
 type FrontendNamespace struct {
 	AllowedDBs      []string           `yaml:"allowed_dbs"`
 	SlowSQLTime     int                `yaml:"slow_sql_time"`
-	DeniedSQLs      []DeniedSqlInfo    `yaml:"denied_sqls"`
 	DeniedIPs       []string           `yaml:"denied_ips"`
 	IdleTimeout     int                `yaml:"idle_timeout"`
 	IsGlobalBreaker bool               `yaml:"global_breaker_switch"`
 	Users           []FrontendUserInfo `yaml:"users"`
+	SQLBlackList    []SQLInfo          `yaml:"sql_blacklist"`
+	SQLWhiteList    []SQLInfo          `yaml:"sql_whitelist"`
 }
 
 type FrontendUserInfo struct {
 	Username string `yaml:"username"`
 	Password string `yaml:"password"`
 }
 
+type SQLInfo struct {
+	SQL string `yaml:"sql"`
+}
+
 type BackendNamespace struct {
 	Username     string   `yaml:"username"`
 	Password     string   `yaml:"password"`
@@ -40,11 +40,11 @@ type BackendNamespace struct {
 type StrategyInfo struct {
 	MinQps               int64 `yaml:"min_qps"`
 	SqlTimeoutMs         int64 `yaml:"sql_timeout_ms"`
-	FailureRatethreshold int64 `yaml:"failure_ratethreshold"`
+	FailureRatethreshold int64 `yaml:"failure_rate_threshold"`
 	FailureNum           int64 `yaml:"failure_num"`
-	OpenStatusDurationMs int64 `yaml:"open_status_duratio_ms"`
+	OpenStatusDurationMs int64 `yaml:"open_status_duration_ms"`
 	Size                 int64 `yaml:"size"`
-	CellIntervalMs       int64 `yaml:"cellInterval_ms"`
+	CellIntervalMs       int64 `yaml:"cell_interval_ms"`
 }
 
 type BreakerInfo struct {

pkg/config/namespace_example.yaml

@@ -24,9 +24,9 @@ breaker:
   scope: "sql"
   strategies:
     - min_qps: 3
-      failure_ratethreshold: 0
+      failure_rate_threshold: 0
       failure_num: 5
       sql_timeoutMs: 2000
-      open_status_durationMs: 5000
+      open_status_duration_ms: 5000
       size: 10
-      cellIntervalMs: 1000
+      cell_interval_ms: 1000

pkg/proxy/driver/domain.go

@@ -14,6 +14,7 @@ type Namespace interface {
 	IsDatabaseAllowed(db string) bool
 	ListDatabases() []string
 	IsDeniedSQL(sqlFeature uint32) bool
+	IsAllowedSQL(sqlFeature uint32) bool
 	GetPooledConn(context.Context) (PooledBackendConn, error)
 	IncrConnCount()
 	DescConnCount()

pkg/proxy/driver/queryctx.go

@@ -125,6 +125,10 @@ func (q *QueryCtxImpl) Execute(ctx context.Context, sql string) (*gomysql.Result
 		return nil, mysql.NewErrf(mysql.ErrUnknown, "statement is denied")
 	}
 
+	if q.isStmtAllowed(ctx, sqlDigest) {
+		return q.execute(ctx, sql, stmt)
+	}
+
 	if !q.isStmtNeedToCheckCircuitBreaking(stmt) {
 		return q.execute(ctx, sql, stmt)
 	}

pkg/proxy/driver/queryctx_exec.go

@@ -16,11 +16,14 @@ import (
 	"go.uber.org/zap"
 )
 
-// TODO(eastfisher): implement this function
 func (q *QueryCtxImpl) isStmtDenied(ctx context.Context, sqlDigest uint32) bool {
 	return q.ns.IsDeniedSQL(sqlDigest)
 }
 
+func (q *QueryCtxImpl) isStmtAllowed(ctx context.Context, sqlDigest uint32) bool {
+	return q.ns.IsAllowedSQL(sqlDigest)
+}
+
 func (q *QueryCtxImpl) getBreakerName(ctx context.Context, sql string, breaker Breaker) (string, bool) {
 	switch breaker.GetBreakerScope() {
 	case "namesapce":

pkg/proxy/namespace/builder.go

@@ -1,13 +1,16 @@
 package namespace
 
 import (
+	"hash/crc32"
 	"time"
 
 	"github.com/pingcap-incubator/weir/pkg/config"
 	"github.com/pingcap-incubator/weir/pkg/proxy/backend"
 	"github.com/pingcap-incubator/weir/pkg/proxy/driver"
+	wast "github.com/pingcap-incubator/weir/pkg/util/ast"
 	"github.com/pingcap-incubator/weir/pkg/util/datastructure"
 	"github.com/pingcap/errors"
+	"github.com/pingcap/parser"
 )
 
 type NamespaceImpl struct {
@@ -77,6 +80,42 @@ func BuildFrontend(cfg *config.FrontendNamespace) (Frontend, error) {
 	}
 	fns.userPasswd = userPasswds
 
+	sqlBlacklist := make(map[uint32]SQLInfo)
+	fns.sqlBlacklist = sqlBlacklist
+
+	p := parser.New()
+	for _, deniedSQL := range cfg.SQLBlackList {
+		stmtNodes, _, err := p.Parse(deniedSQL.SQL, "", "")
+		if err != nil {
+			return nil, err
+		}
+		if len(stmtNodes) != 1 {
+			return nil, nil
+		}
+		v, err := wast.ExtractAstVisit(stmtNodes[0])
+		if err != nil {
+			return nil, err
+		}
+		fns.sqlBlacklist[crc32.ChecksumIEEE([]byte(v.SqlFeature()))] = SQLInfo{SQL: deniedSQL.SQL}
+	}
+
+	sqlWhitelist := make(map[uint32]SQLInfo)
+	fns.sqlWhitelist = sqlWhitelist
+	for _, allowedSQL := range cfg.SQLWhiteList {
+		stmtNodes, _, err := p.Parse(allowedSQL.SQL, "", "")
+		if err != nil {
+			return nil, err
+		}
+		if len(stmtNodes) != 1 {
+			return nil, nil
+		}
+		v, err := wast.ExtractAstVisit(stmtNodes[0])
+		if err != nil {
+			return nil, err
+		}
+		fns.sqlWhitelist[crc32.ChecksumIEEE([]byte(v.SqlFeature()))] = SQLInfo{SQL: allowedSQL.SQL}
+	}
+
 	return fns, nil
 }
 

pkg/proxy/namespace/domain.go

@@ -11,6 +11,7 @@ type Namespace interface {
 	IsDatabaseAllowed(db string) bool
 	ListDatabases() []string
 	IsDeniedSQL(sqlFeature uint32) bool
+	IsAllowedSQL(sqlFeature uint32) bool
 	GetPooledConn(context.Context) (driver.PooledBackendConn, error)
 	Close()
 	GetBreaker() (driver.Breaker, error)
@@ -21,6 +22,7 @@ type Frontend interface {
 	IsDatabaseAllowed(db string) bool
 	ListDatabases() []string
 	IsDeniedSQL(sqlFeature uint32) bool
+	IsAllowedSQL(sqlFeature uint32) bool
 }
 
 type Backend interface {

pkg/proxy/namespace/frontend.go

@@ -3,62 +3,19 @@ package namespace
 import (
 	"bytes"
 
-	"github.com/pingcap-incubator/weir/pkg/config"
-	wast "github.com/pingcap-incubator/weir/pkg/util/ast"
-	"github.com/pingcap-incubator/weir/pkg/util/datastructure"
 	"github.com/pingcap-incubator/weir/pkg/util/passwd"
-	"github.com/pingcap/parser"
-	"hash/crc32"
-	"time"
 )
 
-type DeniedSQLInfo struct {
-	Sql string
-	Ttl int64
+type SQLInfo struct {
+	SQL string
 }
 
 type FrontendNamespace struct {
 	allowedDBs   []string
 	allowedDBSet map[string]struct{}
 	userPasswd   map[string]string
-	deniedSQLs   map[uint32]DeniedSQLInfo
-}
-
-func CreateFrontendNamespace(namespace string, cfg *config.FrontendNamespace) (*FrontendNamespace, error) {
-	fns := &FrontendNamespace{
-		allowedDBs: cfg.AllowedDBs,
-	}
-	fns.allowedDBSet = datastructure.StringSliceToSet(cfg.AllowedDBs)
-
-	userPasswd := make(map[string]string)
-	for _, userInfo := range cfg.Users {
-		userPasswd[userInfo.Username] = userPasswd[userInfo.Password]
-	}
-	fns.userPasswd = userPasswd
-
-	deniedSQLs := make(map[uint32]DeniedSQLInfo)
-	fns.deniedSQLs = deniedSQLs
-
-	p := parser.New()
-	for _, deniedSQL := range cfg.DeniedSQLs {
-		stmtNodes, _, err := p.Parse(deniedSQL.Sql, "", "")
-		if err != nil {
-			return nil, err
-		}
-		if len(stmtNodes) != 1 {
-			return nil, nil
-		}
-		v, err := wast.ExtractAstVisit(stmtNodes[0])
-		if err != nil {
-			return nil, err
-		}
-		if deniedSQL.Ttl != 0 {
-			deniedSQL.Ttl = time.Now().Unix() + deniedSQL.Ttl
-		}
-		fns.deniedSQLs[crc32.ChecksumIEEE([]byte(v.SqlFeature()))] = DeniedSQLInfo{Sql: deniedSQL.Sql, Ttl: deniedSQL.Ttl}
-	}
-
-	return fns, nil
+	sqlBlacklist map[uint32]SQLInfo
+	sqlWhitelist map[uint32]SQLInfo
 }
 
 func (n *FrontendNamespace) Auth(username string, passwdBytes []byte, salt []byte) bool {
@@ -82,14 +39,11 @@ func (n *FrontendNamespace) ListDatabases() []string {
 }
 
 func (n *FrontendNamespace) IsDeniedSQL(sqlFeature uint32) bool {
-	if val, ok := n.deniedSQLs[sqlFeature]; ok {
-		if val.Ttl == 0 {
-			return true
-		}
-		if time.Now().Unix() > val.Ttl {
-			return false
-		}
-		return true
-	}
-	return false
+	_, ok := n.sqlBlacklist[sqlFeature]
+	return ok
+}
+
+func (n *FrontendNamespace) IsAllowedSQL(sqlFeature uint32) bool {
+	_, ok := n.sqlWhitelist[sqlFeature]
+	return ok
 }

pkg/proxy/namespace/namespace.go

@@ -77,6 +77,10 @@ func (n *NamespaceWrapper) IsDeniedSQL(sqlFeature uint32) bool {
 	return n.mustGetCurrentNamespace().IsDeniedSQL(sqlFeature)
 }
 
+func (n *NamespaceWrapper) IsAllowedSQL(sqlFeature uint32) bool {
+	return n.mustGetCurrentNamespace().IsAllowedSQL(sqlFeature)
+}
+
 func (n *NamespaceWrapper) GetPooledConn(ctx context.Context) (driver.PooledBackendConn, error) {
 	return n.mustGetCurrentNamespace().GetPooledConn(ctx)
 }