Fix stmt query does not follow rules

proxy/server/conn_query.go

@@ -113,6 +113,21 @@ func (c *ClientConn) handleQuery(sql string) (err error) {
 	return nil
 }
 
+func (c *ClientConn) getExecuteNode() (*backend.Node, error) {
+	if c.db != "" {
+		rule := c.schema.rule.Rules[c.db][""]
+		if rule != nil && len(rule.Nodes) > 0 {
+			return c.proxy.GetNode(rule.Nodes[0]), nil
+		}
+	}
+	// fallback to default
+	defaultRule := c.schema.rule.DefaultRule
+	if len(defaultRule.Nodes) == 0 {
+		return nil, errors.ErrNoDefaultNode
+	}
+	return c.proxy.GetNode(defaultRule.Nodes[0]), nil
+}
+
 func (c *ClientConn) getBackendConn(n *backend.Node, fromSlave bool) (co *backend.BackendConn, err error) {
 	if !c.isInTransaction() {
 		if fromSlave {

proxy/server/conn_stmt.go

@@ -21,7 +21,6 @@ import (
 	"strconv"
 	"strings"
 
-	"github.com/flike/kingshard/core/errors"
 	"github.com/flike/kingshard/core/golog"
 	"github.com/flike/kingshard/mysql"
 	"github.com/flike/kingshard/sqlparser"
@@ -72,9 +71,10 @@ func (c *ClientConn) handleStmtPrepare(sql string) error {
 
 	s.sql = sql
 
-	defaultRule := c.schema.rule.DefaultRule
-
-	n := c.proxy.GetNode(defaultRule.Nodes[0])
+	n, err := c.getExecuteNode()
+	if err != nil {
+		return err
+	}
 
 	co, err := c.getBackendConn(n, false)
 	defer c.closeConn(co, false)
@@ -258,14 +258,13 @@ func (c *ClientConn) handleStmtExecute(data []byte) error {
 }
 
 func (c *ClientConn) handlePrepareSelect(stmt *sqlparser.Select, sql string, args []interface{}) error {
-	defaultRule := c.schema.rule.DefaultRule
-	if len(defaultRule.Nodes) == 0 {
-		return errors.ErrNoDefaultNode
+	n, err := c.getExecuteNode()
+	if err != nil {
+		return err
 	}
-	defaultNode := c.proxy.GetNode(defaultRule.Nodes[0])
 
 	//choose connection in slave DB first
-	conn, err := c.getBackendConn(defaultNode, true)
+	conn, err := c.getBackendConn(n, true)
 	defer c.closeConn(conn, false)
 	if err != nil {
 		return err
@@ -295,14 +294,13 @@ func (c *ClientConn) handlePrepareSelect(stmt *sqlparser.Select, sql string, arg
 }
 
 func (c *ClientConn) handlePrepareExec(stmt sqlparser.Statement, sql string, args []interface{}) error {
-	defaultRule := c.schema.rule.DefaultRule
-	if len(defaultRule.Nodes) == 0 {
-		return errors.ErrNoDefaultNode
+	n, err := c.getExecuteNode()
+	if err != nil {
+		return err
 	}
-	defaultNode := c.proxy.GetNode(defaultRule.Nodes[0])
 
 	//execute in Master DB
-	conn, err := c.getBackendConn(defaultNode, false)
+	conn, err := c.getBackendConn(n, false)
 	defer c.closeConn(conn, false)
 	if err != nil {
 		return err