Backport of USE and DATABASE() fix

go/mysql/conn.go

@@ -778,10 +778,7 @@ func (c *Conn) handleNextCommand(handler Handler) error {
 	case ComInitDB:
 		db := c.parseComInitDB(data)
 		c.recycleReadPacket()
-		c.schemaName = db
-		handler.ComInitDB(c, db)
-		if err := c.writeOKPacket(0, 0, c.StatusFlags, 0); err != nil {
-			log.Errorf("Error writing ComInitDB result to %s: %v", c, err)
+		if err := c.execQuery(fmt.Sprintf("use `%s`", db), handler, false); err != nil {
 			return err
 		}
 	case ComQuery:

go/mysql/fakesqldb/server.go

@@ -312,10 +312,6 @@ func (db *DB) ConnectionClosed(c *mysql.Conn) {
 	delete(db.connections, c.ConnectionID)
 }
 
-// ComInitDB is part of the mysql.Handler interface.
-func (db *DB) ComInitDB(c *mysql.Conn, schemaName string) {
-}
-
 // ComQuery is part of the mysql.Handler interface.
 func (db *DB) ComQuery(c *mysql.Conn, query string, callback func(*sqltypes.Result) error) error {
 	return db.Handler.HandleQuery(c, query, callback)

go/mysql/server.go

@@ -18,6 +18,7 @@ package mysql
 
 import (
 	"crypto/tls"
+	"fmt"
 	"io"
 	"net"
 	"strings"
@@ -90,10 +91,6 @@ type Handler interface {
 	// ConnectionClosed is called when a connection is closed.
 	ConnectionClosed(c *Conn)
 
-	// InitDB is called once at the beginning to set db name,
-	// and subsequently for every ComInitDB event.
-	ComInitDB(c *Conn, schemaName string)
-
 	// ComQuery is called when a connection receives a query.
 	// Note the contents of the query slice may change after
 	// the first call to callback. So the Handler should not
@@ -441,6 +438,17 @@ func (l *Listener) handle(conn net.Conn, connectionID uint32, acceptTime time.Ti
 		defer connCountPerUser.Add(c.User, -1)
 	}
 
+	// Set initial db name.
+	if c.schemaName != "" {
+		err = l.handler.ComQuery(c, fmt.Sprintf("use `%s`", c.schemaName), func(result *sqltypes.Result) error {
+			return nil
+		})
+		if err != nil {
+			c.writeErrorPacketFromError(err)
+			return
+		}
+	}
+
 	// Negotiation worked, send OK packet.
 	if err := c.writeOKPacket(0, 0, c.StatusFlags, 0); err != nil {
 		log.Errorf("Cannot write OK packet to %s: %v", c, err)
@@ -457,9 +465,6 @@ func (l *Listener) handle(conn net.Conn, connectionID uint32, acceptTime time.Ti
 		log.Warningf("Slow connection from %s: %v", c, connectTime)
 	}
 
-	// Set initial db name.
-	l.handler.ComInitDB(c, c.schemaName)
-
 	for {
 		err := c.handleNextCommand(l.handler)
 		if err != nil {

go/mysql/server_test.go

@@ -113,9 +113,6 @@ func (th *testHandler) NewConnection(c *Conn) {
 func (th *testHandler) ConnectionClosed(c *Conn) {
 }
 
-func (th *testHandler) ComInitDB(c *Conn, schemaName string) {
-}
-
 func (th *testHandler) ComQuery(c *Conn, query string, callback func(*sqltypes.Result) error) error {
 	if result := th.Result(); result != nil {
 		callback(result)

go/vt/vtgate/executor.go

@@ -230,12 +230,7 @@ func (e *Executor) execute(ctx context.Context, safeSession *SafeSession, sql st
 // addNeededBindVars adds bind vars that are needed by the plan
 func (e *Executor) addNeededBindVars(bindVarNeeds sqlparser.BindVarNeeds, bindVars map[string]*querypb.BindVariable, session *SafeSession) error {
 	if bindVarNeeds.NeedDatabase {
-		keyspace, _, _, _ := e.ParseDestinationTarget(session.TargetString)
-		if keyspace == "" {
-			bindVars[sqlparser.DBVarName] = sqltypes.NullBindVariable
-		} else {
-			bindVars[sqlparser.DBVarName] = sqltypes.StringBindVariable(keyspace)
-		}
+		bindVars[sqlparser.DBVarName] = sqltypes.StringBindVariable(session.TargetString)
 	}
 
 	if bindVarNeeds.NeedLastInsertID {

go/vt/vtgate/executor_select_test.go

@@ -389,7 +389,7 @@ func TestSelectDatabase(t *testing.T) {
 			{Name: "database()", Type: sqltypes.VarBinary},
 		},
 		Rows: [][]sqltypes.Value{{
-			sqltypes.NewVarBinary("TestExecutor"),
+			sqltypes.NewVarBinary("TestExecutor@master"),
 		}},
 	}
 	require.NoError(t, err)

go/vt/vtgate/executor_test.go

@@ -207,7 +207,7 @@ func TestDirectTargetRewrites(t *testing.T) {
 	require.NoError(t, err)
 	testQueries(t, "sbclookup", sbclookup, []*querypb.BoundQuery{{
 		Sql:           "select :__vtdbname as `database()` from dual",
-		BindVariables: map[string]*querypb.BindVariable{"__vtdbname": sqltypes.StringBindVariable("TestUnsharded")},
+		BindVariables: map[string]*querypb.BindVariable{"__vtdbname": sqltypes.StringBindVariable("TestUnsharded/0@master")},
 	}})
 }
 
@@ -1038,7 +1038,7 @@ func TestExecutorUse(t *testing.T) {
 	}
 
 	_, err = executor.Execute(context.Background(), "TestExecute", NewSafeSession(&vtgatepb.Session{}), "use UnexistentKeyspace", nil)
-	wantErr = "invalid keyspace provided: UnexistentKeyspace"
+	wantErr = "Unknown database 'UnexistentKeyspace' (errno 1049) (sqlstate 42000)"
 	if err == nil || err.Error() != wantErr {
 		t.Errorf("got: %v, want %v", err, wantErr)
 	}

go/vt/vtgate/mysql_protocol_test.go

@@ -118,7 +118,7 @@ func TestMySQLProtocolExecuteUseStatement(t *testing.T) {
 	// No such keyspace this will fail
 	_, err = c.ExecuteFetch("use InvalidKeyspace", 0, false)
 	require.Error(t, err)
-	assert.Contains(t, err.Error(), "invalid keyspace provided: InvalidKeyspace")
+	assert.Contains(t, err.Error(), "Unknown database 'InvalidKeyspace' (errno 1049) (sqlstate 42000)")
 
 	// That doesn't reset the vitess_target
 	qr, err = c.ExecuteFetch("show vitess_target", 1, false)
@@ -135,18 +135,8 @@ func TestMySQLProtocolExecuteUseStatement(t *testing.T) {
 }
 
 func TestMysqlProtocolInvalidDB(t *testing.T) {
-	c, err := mysqlConnect(&mysql.ConnParams{DbName: "invalidDB"})
-	if err != nil {
-		t.Fatal(err)
-	}
-	defer c.Close()
-
-	_, err = c.ExecuteFetch("select id from t1", 10, true /* wantfields */)
-	c.Close()
-	want := "vtgate: : keyspace invalidDB not found in vschema (errno 1105) (sqlstate HY000) during query: select id from t1"
-	if err == nil || err.Error() != want {
-		t.Errorf("exec with db:\n%v, want\n%s", err, want)
-	}
+	_, err := mysqlConnect(&mysql.ConnParams{DbName: "invalidDB"})
+	require.EqualError(t, err, "vtgate: : Unknown database 'invalidDB' (errno 1049) (sqlstate 42000) (errno 1049) (sqlstate 42000)")
 }
 
 func TestMySQLProtocolClientFoundRows(t *testing.T) {

go/vt/vtgate/plan_executor_select_test.go

@@ -391,7 +391,7 @@ func TestPlanSelectDatabase(t *testing.T) {
 			{Name: "database()", Type: sqltypes.VarBinary},
 		},
 		Rows: [][]sqltypes.Value{{
-			sqltypes.NewVarBinary("TestExecutor"),
+			sqltypes.NewVarBinary("TestExecutor@master"),
 		}},
 	}
 	require.NoError(t, err)

go/vt/vtgate/plan_executor_test.go

@@ -205,7 +205,7 @@ func TestPlanDirectTargetRewrites(t *testing.T) {
 	require.NoError(t, err)
 	testQueries(t, "sbclookup", sbclookup, []*querypb.BoundQuery{{
 		Sql:           "select :__vtdbname as `database()` from dual",
-		BindVariables: map[string]*querypb.BindVariable{"__vtdbname": sqltypes.StringBindVariable("TestUnsharded")},
+		BindVariables: map[string]*querypb.BindVariable{"__vtdbname": sqltypes.StringBindVariable("TestUnsharded/0@master")},
 	}})
 }
 
@@ -996,7 +996,7 @@ func TestPlanExecutorUse(t *testing.T) {
 	}
 
 	_, err = executor.Execute(context.Background(), "TestExecute", NewSafeSession(&vtgatepb.Session{}), "use UnexistentKeyspace", nil)
-	wantErr = "invalid keyspace provided: UnexistentKeyspace"
+	wantErr = "Unknown database 'UnexistentKeyspace' (errno 1049) (sqlstate 42000)"
 	if err == nil || err.Error() != wantErr {
 		t.Errorf("got: %v, want %v", err, wantErr)
 	}

go/vt/vtgate/plugin_mysql_server.go

@@ -166,10 +166,6 @@ func startSpan(ctx context.Context, query, label string) (trace.Span, context.Co
 	return startSpanTestable(ctx, query, label, trace.NewSpan, trace.NewFromString)
 }
 
-func (vh *vtgateHandler) ComInitDB(c *mysql.Conn, schemaName string) {
-	vh.session(c).TargetString = schemaName
-}
-
 func (vh *vtgateHandler) ComQuery(c *mysql.Conn, query string, callback func(*sqltypes.Result) error) error {
 	ctx := context.Background()
 	var cancel context.CancelFunc

go/vt/vtgate/plugin_mysql_server_test.go

@@ -43,9 +43,6 @@ func (th *testHandler) NewConnection(c *mysql.Conn) {
 func (th *testHandler) ConnectionClosed(c *mysql.Conn) {
 }
 
-func (th *testHandler) ComInitDB(c *mysql.Conn, schemaName string) {
-}
-
 func (th *testHandler) ComQuery(c *mysql.Conn, q string, callback func(*sqltypes.Result) error) error {
 	return nil
 }

go/vt/vtgate/vcursor_impl.go

@@ -21,6 +21,8 @@ import (
 	"sync/atomic"
 	"time"
 
+	"vitess.io/vitess/go/mysql"
+
 	"vitess.io/vitess/go/vt/callerid"
 	vschemapb "vitess.io/vitess/go/vt/proto/vschema"
 	"vitess.io/vitess/go/vt/topotools"
@@ -322,7 +324,7 @@ func (vc *vcursorImpl) SetTarget(target string) error {
 		return err
 	}
 	if _, ok := vc.vschema.Keyspaces[keyspace]; keyspace != "" && !ok {
-		return vterrors.Errorf(vtrpcpb.Code_INVALID_ARGUMENT, "invalid keyspace provided: %s", keyspace)
+		return mysql.NewSQLError(mysql.ERBadDb, "42000", "Unknown database '%s'", keyspace)
 	}
 
 	if vc.safeSession.InTransaction() && tabletType != topodatapb.TabletType_MASTER {

go/vt/vtgate/vcursor_impl_test.go

@@ -197,7 +197,7 @@ func TestSetTarget(t *testing.T) {
 	}, {
 		vschema:       vschemaWith2KS,
 		targetString:  "ks3",
-		expectedError: "invalid keyspace provided: ks3",
+		expectedError: "Unknown database 'ks3' (errno 1049) (sqlstate 42000)",
 	}, {
 		vschema:       vschemaWith2KS,
 		targetString:  "ks2@replica",