Repository: vitessio/vitess
Pull request #10047 (2022-04-06 13:15:31)

13.0 fix: allow multiple distinct columns on single shards
## Description

This is a backport of #9940.

## Checklist
- [x] Tests were added or are not required
- [x] Documentation was added or is not required



Patch:
diff --git a/go/vt/vterrors/constants.go b/go/vt/vterrors/constants.go
index 2fcbfda13d..c9530daa36 100644
--- a/go/vt/vterrors/constants.go
+++ b/go/vt/vterrors/constants.go
@@ -27,6 +27,12 @@ const (
 // RxOp regex for operation not allowed error
 var RxOp = regexp.MustCompile("operation not allowed in state (NOT_SERVING|SHUTTING_DOWN)")
 
+// TxEngineClosed for transaction engine closed error
+const TxEngineClosed = "tx engine can't accept new connections in state %v"
+
+// RxTxEngineClosed regex for operation not allowed error
+var RxTxEngineClosed = regexp.MustCompile("tx engine can't accept new connections in state (NotServing|Transitioning)")
+
 // WrongTablet for invalid tablet type error
 const WrongTablet = "wrong tablet type"
 
diff --git a/go/vt/vtgate/scatter_conn.go b/go/vt/vtgate/scatter_conn.go
index 43ed3f0446..713f4faece 100644
--- a/go/vt/vtgate/scatter_conn.go
+++ b/go/vt/vtgate/scatter_conn.go
@@ -204,7 +204,7 @@ func (stc *ScatterConn) ExecuteMultiShard(
 				}
 			}
 
-			qs, err = getQueryService(rs, info)
+			qs, err = getQueryService(rs, info, session, false)
 			if err != nil {
 				return nil, err
 			}
@@ -297,7 +297,7 @@ func checkAndResetShardSession(info *shardActionInfo, err error, session *SafeSe
 	return retry
 }
 
-func getQueryService(rs *srvtopo.ResolvedShard, info *shardActionInfo) (queryservice.QueryService, error) {
+func getQueryService(rs *srvtopo.ResolvedShard, info *shardActionInfo, session *SafeSession, skipReset bool) (queryservice.QueryService, error) {
 	_, usingLegacyGw := rs.Gateway.(*DiscoveryGateway)
 	if usingLegacyGw &&
 		(info.actionNeeded == reserve || info.actionNeeded == reserveBegin) {
@@ -306,7 +306,21 @@ func getQueryService(rs *srvtopo.ResolvedShard, info *shardActionInfo) (queryser
 	if usingLegacyGw || info.alias == nil {
 		return rs.Gateway, nil
 	}
-	return rs.Gateway.QueryServiceByAlias(info.alias, rs.Target)
+	qs, err := rs.Gateway.QueryServiceByAlias(info.alias, rs.Target)
+	if err == nil || skipReset {
+		return qs, err
+	}
+	// If the session info has only reserved connection and no transaction then we will route it through gateway
+	// Otherwise, we will fail.
+	if info.reservedID == 0 || info.transactionID != 0 {
+		return nil, err
+	}
+	err = session.ResetShard(info.alias)
+	if err != nil {
+		return nil, err
+	}
+	// Returning rs.Gateway will make the gateway to choose new healthy tablet for the targeted tablet type.
+	return rs.Gateway, nil
 }
 
 func (stc *ScatterConn) processOneStreamingResult(mu *sync.Mutex, fieldSent *bool, qr *sqltypes.Result, callback func(*sqltypes.Result) error) error {
@@ -373,7 +387,7 @@ func (stc *ScatterConn) StreamExecuteMulti(
 				}
 			}
 
-			qs, err = getQueryService(rs, info)
+			qs, err = getQueryService(rs, info, session, false)
 			if err != nil {
 				return nil, err
 			}
@@ -695,7 +709,7 @@ func (stc *ScatterConn) ExecuteLock(
 		_ = stc.txConn.ReleaseLock(ctx, session)
 		return nil, vterrors.Wrap(err, "Any previous held locks are released")
 	}
-	qs, err := getQueryService(rs, info)
+	qs, err := getQueryService(rs, info, nil, true)
 	if err != nil {
 		return nil, err
 	}
@@ -741,16 +755,33 @@ func wasConnectionClosed(err error) bool {
 	sqlErr := mysql.NewSQLErrorFromError(err).(*mysql.SQLError)
 	message := sqlErr.Error()
 
-	return sqlErr.Number() == mysql.CRServerGone ||
-		sqlErr.Number() == mysql.CRServerLost ||
-		(sqlErr.Number() == mysql.ERQueryInterrupted && txClosed.MatchString(message))
+	switch sqlErr.Number() {
+	case mysql.CRServerGone, mysql.CRServerLost:
+		return true
+	case mysql.ERQueryInterrupted:
+		return txClosed.MatchString(message)
+	default:
+		return false
+	}
 }
 
 func requireNewQS(err error, target *querypb.Target) bool {
 	code := vterrors.Code(err)
 	msg := err.Error()
-	return (code == vtrpcpb.Code_FAILED_PRECONDITION && vterrors.RxWrongTablet.MatchString(msg)) ||
-		(code == vtrpcpb.Code_CLUSTER_EVENT && ((target != nil && target.TabletType == topodatapb.TabletType_PRIMARY) || vterrors.RxOp.MatchString(msg)))
+	if (code == vtrpcpb.Code_FAILED_PRECONDITION && vterrors.RxWrongTablet.MatchString(msg)) ||
+		(code == vtrpcpb.Code_CLUSTER_EVENT && ((target != nil && target.TabletType == topodatapb.TabletType_PRIMARY) || vterrors.RxOp.MatchString(msg))) ||
+		(code == vtrpcpb.Code_UNAVAILABLE && vterrors.RxTxEngineClosed.MatchString(msg)) {
+		return true
+	}
+
+	sqlErr := mysql.NewSQLErrorFromError(err).(*mysql.SQLError)
+
+	switch sqlErr.Number() {
+	case mysql.CRConnectionError, mysql.CRConnHostError:
+		return true
+	default:
+		return false
+	}
 }
 
 // actionInfo looks at the current session, and returns information about what needs to be done for this tablet
diff --git a/go/vt/vtgate/semantics/analyzer.go b/go/vt/vtgate/semantics/analyzer.go
index d5b1fd16c0..39d4832610 100644
--- a/go/vt/vtgate/semantics/analyzer.go
+++ b/go/vt/vtgate/semantics/analyzer.go
@@ -289,7 +289,7 @@ func (a *analyzer) checkForInvalidConstructs(cursor *sqlparser.Cursor) error {
 
 		if node.Distinct {
 			err := vterrors.Errorf(vtrpcpb.Code_INVALID_ARGUMENT, "syntax error: %s", sqlparser.String(node))
-			if len(node.Exprs) != 1 {
+			if len(node.Exprs) < 1 {
 				return err
 			} else if _, ok := node.Exprs[0].(*sqlparser.AliasedExpr); !ok {
 				return err
diff --git a/go/vt/vttablet/tabletserver/tabletserver.go b/go/vt/vttablet/tabletserver/tabletserver.go
index 69454efc42..c33bc53119 100644
--- a/go/vt/vttablet/tabletserver/tabletserver.go
+++ b/go/vt/vttablet/tabletserver/tabletserver.go
@@ -1444,7 +1444,7 @@ func convertErrorCode(err error) vtrpcpb.Code {
 		errCode = vtrpcpb.Code_RESOURCE_EXHAUSTED
 	case mysql.ERLockWaitTimeout:
 		errCode = vtrpcpb.Code_DEADLINE_EXCEEDED
-	case mysql.CRServerGone, mysql.ERServerShutdown, mysql.ERServerIsntAvailable:
+	case mysql.CRServerGone, mysql.ERServerShutdown, mysql.ERServerIsntAvailable, mysql.CRConnectionError, mysql.CRConnHostError:
 		errCode = vtrpcpb.Code_UNAVAILABLE
 	case mysql.ERFormNotFound, mysql.ERKeyNotFound, mysql.ERBadFieldError, mysql.ERNoSuchThread, mysql.ERUnknownTable, mysql.ERCantFindUDF, mysql.ERNonExistingGrant,
 		mysql.ERNoSuchTable, mysql.ERNonExistingTableGrant, mysql.ERKeyDoesNotExist:
diff --git a/go/vt/vttablet/tabletserver/tx_engine.go b/go/vt/vttablet/tabletserver/tx_engine.go
index 1c1753a4e1..c16c8b3ae2 100644
--- a/go/vt/vttablet/tabletserver/tx_engine.go
+++ b/go/vt/vttablet/tabletserver/tx_engine.go
@@ -210,7 +210,7 @@ func (te *TxEngine) isTxPoolAvailable(addToWaitGroup func(int)) error {
 
 	canOpenTransactions := te.state == AcceptingReadOnly || te.state == AcceptingReadAndWrite
 	if !canOpenTransactions {
-		return vterrors.Errorf(vtrpc.Code_UNAVAILABLE, "tx engine can't accept new connections in state %v", te.state)
+		return vterrors.Errorf(vtrpc.Code_UNAVAILABLE, vterrors.TxEngineClosed, te.state)
 	}
 	addToWaitGroup(1)
 	return nil


Test patch:
diff --git a/go/test/endtoend/vtgate/reservedconn/reconnect3/main_test.go b/go/test/endtoend/vtgate/reservedconn/reconnect3/main_test.go
new file mode 100644
index 0000000000..d20e1e14c6
--- /dev/null
+++ b/go/test/endtoend/vtgate/reservedconn/reconnect3/main_test.go
@@ -0,0 +1,96 @@
+/*
+Copyright 2022 The Vitess Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package reservedconn
+
+import (
+	"context"
+	"flag"
+	"os"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"vitess.io/vitess/go/mysql"
+	"vitess.io/vitess/go/test/endtoend/cluster"
+)
+
+var (
+	clusterInstance *cluster.LocalProcessCluster
+	vtParams        mysql.ConnParams
+	keyspaceName    = "ks"
+	cell            = "zone1"
+	hostname        = "localhost"
+	sqlSchema       = `create table test(id bigint primary key)Engine=InnoDB;`
+)
+
+func TestMain(m *testing.M) {
+	defer cluster.PanicHandler(nil)
+	flag.Parse()
+
+	exitCode := func() int {
+		clusterInstance = cluster.NewCluster(cell, hostname)
+		defer clusterInstance.Teardown()
+
+		// Start topo server
+		if err := clusterInstance.StartTopo(); err != nil {
+			return 1
+		}
+
+		// Start keyspace
+		keyspace := &cluster.Keyspace{
+			Name:      keyspaceName,
+			SchemaSQL: sqlSchema,
+		}
+		if err := clusterInstance.StartUnshardedKeyspace(*keyspace, 2, false); err != nil {
+			return 1
+		}
+
+		// Start vtgate
+		clusterInstance.VtGateExtraArgs = []string{"--enable_system_settings=true"}
+		if err := clusterInstance.StartVtgate(); err != nil {
+			return 1
+		}
+
+		vtParams = mysql.ConnParams{
+			Host: clusterInstance.Hostname,
+			Port: clusterInstance.VtgateMySQLPort,
+		}
+		return m.Run()
+	}()
+	os.Exit(exitCode)
+}
+
+func TestMysqlDownServingChange(t *testing.T) {
+	conn, err := mysql.Connect(context.Background(), &vtParams)
+	require.NoError(t, err)
+	defer conn.Close()
+
+	_, err = conn.ExecuteFetch("set default_week_format = 1", 5, false)
+	require.NoError(t, err)
+	_, err = conn.ExecuteFetch("select /*vt+ PLANNER=gen4 */ * from test", 5, false)
+	require.NoError(t, err)
+
+	primaryTablet := clusterInstance.Keyspaces[0].Shards[0].PrimaryTablet()
+	require.NoError(t,
+		primaryTablet.MysqlctlProcess.Stop())
+	require.NoError(t,
+		clusterInstance.VtctlclientProcess.ExecuteCommand("EmergencyReparentShard", "-keyspace_shard", "ks/0"))
+
+	// This should work without any error.
+	_, err = conn.ExecuteFetch("select /*vt+ PLANNER=gen4 */ * from test", 5, false)
+	require.NoError(t, err)
+}
diff --git a/go/test/endtoend/vtgate/reservedconn/reconnect4/main_test.go b/go/test/endtoend/vtgate/reservedconn/reconnect4/main_test.go
new file mode 100644
index 0000000000..22e2119b45
--- /dev/null
+++ b/go/test/endtoend/vtgate/reservedconn/reconnect4/main_test.go
@@ -0,0 +1,98 @@
+/*
+Copyright 2022 The Vitess Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package reservedconn
+
+import (
+	"context"
+	"flag"
+	"os"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"vitess.io/vitess/go/mysql"
+	"vitess.io/vitess/go/test/endtoend/cluster"
+)
+
+var (
+	clusterInstance *cluster.LocalProcessCluster
+	vtParams        mysql.ConnParams
+	keyspaceName    = "ks"
+	cell            = "zone1"
+	hostname        = "localhost"
+	sqlSchema       = `create table test(id bigint primary key)Engine=InnoDB;`
+)
+
+func TestMain(m *testing.M) {
+	defer cluster.PanicHandler(nil)
+	flag.Parse()
+
+	exitCode := func() int {
+		clusterInstance = cluster.NewCluster(cell, hostname)
+		defer clusterInstance.Teardown()
+
+		// Start topo server
+		if err := clusterInstance.StartTopo(); err != nil {
+			return 1
+		}
+
+		// Start keyspace
+		keyspace := &cluster.Keyspace{
+			Name:      keyspaceName,
+			SchemaSQL: sqlSchema,
+		}
+		if err := clusterInstance.StartUnshardedKeyspace(*keyspace, 2, false); err != nil {
+			return 1
+		}
+
+		// Start vtgate
+		clusterInstance.VtGateExtraArgs = []string{"--enable_system_settings=true"}
+		if err := clusterInstance.StartVtgate(); err != nil {
+			return 1
+		}
+
+		vtParams = mysql.ConnParams{
+			Host: clusterInstance.Hostname,
+			Port: clusterInstance.VtgateMySQLPort,
+		}
+		return m.Run()
+	}()
+	os.Exit(exitCode)
+}
+
+func TestVttabletDownServingChange(t *testing.T) {
+	conn, err := mysql.Connect(context.Background(), &vtParams)
+	require.NoError(t, err)
+	defer conn.Close()
+
+	_, err = conn.ExecuteFetch("set default_week_format = 1", 5, false)
+	require.NoError(t, err)
+	_, err = conn.ExecuteFetch("select /*vt+ PLANNER=gen4 */ * from test", 5, false)
+	require.NoError(t, err)
+
+	primaryTablet := clusterInstance.Keyspaces[0].Shards[0].PrimaryTablet()
+	require.NoError(t,
+		primaryTablet.MysqlctlProcess.Stop())
+	// kill vttablet process
+	_ = primaryTablet.VttabletProcess.TearDown()
+	require.NoError(t,
+		clusterInstance.VtctlclientProcess.ExecuteCommand("EmergencyReparentShard", "-keyspace_shard", "ks/0"))
+
+	// This should work without any error.
+	_, err = conn.ExecuteFetch("select /*vt+ PLANNER=gen4 */ * from test", 5, false)
+	require.NoError(t, err)
+}
diff --git a/go/vt/vtgate/planbuilder/testdata/aggr_cases.txt b/go/vt/vtgate/planbuilder/testdata/aggr_cases.txt
index 2677992171..e6e6c5d962 100644
--- a/go/vt/vtgate/planbuilder/testdata/aggr_cases.txt
+++ b/go/vt/vtgate/planbuilder/testdata/aggr_cases.txt
@@ -3116,3 +3116,25 @@ Gen4 plan same as above
     ]
   }
 }
+
+"select count(distinct user_id, name) from unsharded"
+{
+  "QueryType": "SELECT",
+  "Original": "select count(distinct user_id, name) from unsharded",
+  "Instructions": {
+    "OperatorType": "Route",
+    "Variant": "Unsharded",
+    "Keyspace": {
+      "Name": "main",
+      "Sharded": false
+    },
+    "FieldQuery": "select count(distinct user_id, `name`) from unsharded where 1 != 1",
+    "Query": "select count(distinct user_id, `name`) from unsharded",
+    "Table": "unsharded"
+  }
+}
+Gen4 plan same as above
+
+"select count(distinct user_id, name) from user"
+"unsupported: only one expression allowed inside aggregates: count(distinct user_id, `name`)"
+Gen4 error: aggregate functions take a single argument 'count(distinct user_id, `name`)'
diff --git a/test/config.json b/test/config.json
index 24bcf08c18..44b35e3ed1 100644
--- a/test/config.json
+++ b/test/config.json
@@ -822,6 +822,24 @@
 			"RetryMax": 1,
 			"Tags": []
 		},
+		"vtgate_reserved_conn3": {
+			"File": "unused.go",
+			"Args": ["vitess.io/vitess/go/test/endtoend/vtgate/reservedconn/reconnect3"],
+			"Command": [],
+			"Manual": false,
+			"Shard": "vtgate_reservedconn",
+			"RetryMax": 1,
+			"Tags": []
+		},
+		"vtgate_reserved_conn4": {
+			"File": "unused.go",
+			"Args": ["vitess.io/vitess/go/test/endtoend/vtgate/reservedconn/reconnect4"],
+			"Command": [],
+			"Manual": false,
+			"Shard": "vtgate_reservedconn",
+			"RetryMax": 1,
+			"Tags": []
+		},
 		"vtgate_tablet_healthcheck_cache": {
 			"File": "unused.go",
 			"Args": ["vitess.io/vitess/go/test/endtoend/vtgate/tablet_healthcheck_cache", "-timeout", "45m"],