Merge pull request #35 from vitessio/tester-demo

Vitess tester demo

Makefile

@@ -20,8 +20,8 @@ default: check_version build
 build:
 	$(GO) build -o vt ./go/vt
 
-test: build
-	$(GO) test -cover ./go/...
+test:
+	$(GO) test ./go/...
 
 tidy:
 	$(GO) mod tidy

go/cmd/tester.go

@@ -44,6 +44,7 @@ func testerCmd() *cobra.Command {
 	cmd.Flags().BoolVar(&cfg.OLAP, "olap", false, "Use OLAP to run the queries.")
 	cmd.Flags().BoolVar(&cfg.XUnit, "xunit", false, "Get output in an xml file instead of errors directory")
 	cmd.Flags().BoolVar(&cfg.Sharded, "sharded", false, "Run all tests on a sharded keyspace and using auto-vschema. This cannot be used with either -vschema or -vtexplain-vschema.")
+	cmd.Flags().IntVar(&cfg.NumberOfShards, "number-of-shards", 0, "Number of shards to use for the sharded keyspace.")
 
 	return cmd
 }

go/tester/comparing_query_runner.go

@@ -20,10 +20,12 @@ import (
 	"fmt"
 
 	log "github.com/sirupsen/logrus"
+	"vitess.io/vitess/go/test/endtoend/cluster"
 	"vitess.io/vitess/go/test/endtoend/utils"
 	"vitess.io/vitess/go/vt/sqlparser"
 
 	"github.com/vitessio/vt/go/data"
+	"github.com/vitessio/vt/go/tester/state"
 )
 
 type (
@@ -32,56 +34,62 @@ type (
 		reporter          Reporter
 		handleCreateTable CreateTableHandler
 		comparer          utils.MySQLCompare
+		cluster           *cluster.LocalProcessCluster
 	}
 	CreateTableHandler          func(create *sqlparser.CreateTable) func()
 	ComparingQueryRunnerFactory struct{}
 )
 
 func (f ComparingQueryRunnerFactory) Close() {}
 
-func (f ComparingQueryRunnerFactory) NewQueryRunner(reporter Reporter, handleCreateTable CreateTableHandler, comparer utils.MySQLCompare) QueryRunner {
-	return newComparingQueryRunner(reporter, handleCreateTable, comparer)
+func (f ComparingQueryRunnerFactory) NewQueryRunner(
+	reporter Reporter,
+	handleCreateTable CreateTableHandler,
+	comparer utils.MySQLCompare,
+	cluster *cluster.LocalProcessCluster,
+) QueryRunner {
+	return newComparingQueryRunner(reporter, handleCreateTable, comparer, cluster)
 }
 
 func newComparingQueryRunner(
 	reporter Reporter,
 	handleCreateTable CreateTableHandler,
 	comparer utils.MySQLCompare,
+	cluster *cluster.LocalProcessCluster,
 ) *ComparingQueryRunner {
 	return &ComparingQueryRunner{
 		reporter:          reporter,
 		handleCreateTable: handleCreateTable,
 		comparer:          comparer,
+		cluster:           cluster,
 	}
 }
 
-func (nqr ComparingQueryRunner) runQuery(q data.Query, expectedErrs bool, ast sqlparser.Statement) error {
-	return nqr.execute(q, expectedErrs, ast)
+func (nqr ComparingQueryRunner) runQuery(q data.Query, ast sqlparser.Statement, state *state.State) error {
+	return nqr.execute(q, ast, state)
 }
 
-func (nqr *ComparingQueryRunner) execute(query data.Query, expectedErrs bool, ast sqlparser.Statement) error {
+func (nqr *ComparingQueryRunner) execute(query data.Query, ast sqlparser.Statement, state *state.State) error {
 	if len(query.Query) == 0 {
 		return nil
 	}
 
-	if err := nqr.executeStmt(query.Query, ast, expectedErrs); err != nil {
+	if err := nqr.executeStmt(query.Query, ast, state); err != nil {
 		return fmt.Errorf("run \"%v\" at line %d err %v", query.Query, query.Line, err)
 	}
-	// clear expected errors after we execute
-	expectedErrs = false
 
 	return nil
 }
 
-func (nqr *ComparingQueryRunner) executeStmt(query string, ast sqlparser.Statement, expectedErrs bool) (err error) {
+func (nqr *ComparingQueryRunner) executeStmt(query string, ast sqlparser.Statement, state *state.State) (err error) {
 	_, commentOnly := ast.(*sqlparser.CommentOnly)
 	if commentOnly {
 		return nil
 	}
 
 	log.Debugf("executeStmt: %s", query)
 	create, isCreateStatement := ast.(*sqlparser.CreateTable)
-	if isCreateStatement && !expectedErrs {
+	if isCreateStatement && !state.IsErrorExpectedSet() && state.RunOnVitess() {
 		closer := nqr.handleCreateTable(create)
 		defer func() {
 			if err == nil {
@@ -91,14 +99,76 @@ func (nqr *ComparingQueryRunner) executeStmt(query string, ast sqlparser.Stateme
 	}
 
 	switch {
-	case expectedErrs:
-		_, err := nqr.comparer.ExecAllowAndCompareError(query, utils.CompareOptions{CompareColumnNames: true})
-		if err == nil {
-			// If we expected an error, but didn't get one, return an error
-			return fmt.Errorf("expected error, but got none")
+	case state.CheckAndClearErrorExpected():
+		err := nqr.execAndExpectErr(query)
+		if err != nil {
+			nqr.reporter.AddFailure(err)
 		}
 	default:
-		_ = nqr.comparer.Exec(query)
+		var err error
+		switch {
+		case state.CheckAndClearReference():
+			return nqr.executeReference(query, ast)
+		case state.NormalExecution():
+			nqr.comparer.Exec(query)
+		case state.IsVitessOnlySet():
+			_, err = nqr.comparer.VtConn.ExecuteFetch(query, 1000, true)
+		case state.IsMySQLOnlySet():
+			_, err = nqr.comparer.MySQLConn.ExecuteFetch(query, 1000, true)
+		}
+		if err != nil {
+			nqr.reporter.AddFailure(err)
+		}
+	}
+	return nil
+}
+
+func (nqr *ComparingQueryRunner) execAndExpectErr(query string) error {
+	_, err := nqr.comparer.ExecAllowAndCompareError(query, utils.CompareOptions{CompareColumnNames: true})
+	if err == nil {
+		// If we expected an error, but didn't get one, return an error
+		return fmt.Errorf("expected error, but got none")
+	}
+	return nil
+}
+
+func (nqr *ComparingQueryRunner) executeReference(query string, ast sqlparser.Statement) error {
+	_, err := nqr.comparer.MySQLConn.ExecuteFetch(query, 1000, true)
+	if err != nil {
+		return err
+	}
+
+	tables := sqlparser.ExtractAllTables(ast)
+	if len(tables) != 1 {
+		return fmt.Errorf("expected exactly one table in the query, got %d", len(tables))
+	}
+
+	tableName := tables[0]
+
+	tbl, err := vschema.FindTable("" /*empty means global search*/, tableName)
+	if err != nil {
+		return err
 	}
+
+	for _, ks := range nqr.cluster.Keyspaces {
+		if ks.Name == tbl.Keyspace.Name {
+			for _, shard := range ks.Shards {
+				_, err := nqr.comparer.VtConn.ExecuteFetch(fmt.Sprintf("use `%s/%s`", ks.Name, shard.Name), 1000, true)
+				if err != nil {
+					return fmt.Errorf("error setting keyspace/shard: %w", err)
+				}
+				_, err = nqr.comparer.VtConn.ExecuteFetch(query, 1000, true)
+				if err != nil {
+					return fmt.Errorf("error executing query on vtgate: %w", err)
+				}
+			}
+			q := fmt.Sprintf("use %s", ks.Name)
+			_, err = nqr.comparer.VtConn.ExecuteFetch(q, 1000, true)
+			if err != nil {
+				return fmt.Errorf("error setting keyspace: %s %w", q, err)
+			}
+		}
+	}
+
 	return nil
 }

go/tester/execute.go

@@ -78,6 +78,7 @@ func ExecuteTests(
 func SetupCluster(
 	vschemaFile, vtexplainVschemaFile string,
 	sharded bool,
+	numberOfShards int,
 ) (clusterInstance *cluster.LocalProcessCluster, vtParams, mysqlParams mysql.ConnParams, ksNames []string, close func()) {
 	clusterInstance = cluster.NewCluster(defaultCellName, "localhost")
 
@@ -103,8 +104,9 @@ func SetupCluster(
 		}
 
 		if vschemaKs.Keyspace.Sharded {
-			fmt.Printf("starting sharded keyspace: '%s'\n", keyspace.Name)
-			err = clusterInstance.StartKeyspace(*keyspace, []string{"-80", "80-"}, 0, false)
+			shardRanges := generateShardRanges(numberOfShards)
+			fmt.Printf("starting sharded keyspace: '%s' with shards %v\n", keyspace.Name, shardRanges)
+			err = clusterInstance.StartKeyspace(*keyspace, shardRanges, 0, false)
 			errCheck(err)
 		} else {
 			fmt.Printf("starting unsharded keyspace: '%s'\n", keyspace.Name)
@@ -162,6 +164,34 @@ func SetupCluster(
 	}
 }
 
+func generateShardRanges(numberOfShards int) []string {
+	if numberOfShards <= 0 {
+		return []string{}
+	}
+
+	if numberOfShards == 1 {
+		return []string{"-"}
+	}
+
+	ranges := make([]string, numberOfShards)
+	step := 0x100 / numberOfShards
+
+	for i := 0; i < numberOfShards; i++ {
+		start := i * step
+		end := (i + 1) * step
+
+		if i == 0 {
+			ranges[i] = fmt.Sprintf("-%02x", end)
+		} else if i == numberOfShards-1 {
+			ranges[i] = fmt.Sprintf("%02x-", start)
+		} else {
+			ranges[i] = fmt.Sprintf("%02x-%02x", start, end)
+		}
+	}
+
+	return ranges
+}
+
 func defaultVschema(defaultKeyspaceName string) vindexes.VSchema {
 	return vindexes.VSchema{
 		Keyspaces: map[string]*vindexes.KeyspaceSchema{

go/tester/execute_test.go

@@ -0,0 +1,43 @@
+/*
+Copyright 2024 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 tester
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestGenerateShardRanges(t *testing.T) {
+	tests := []struct {
+		numberOfShards int
+		expected       []string
+	}{
+		{1, []string{"-"}},
+		{2, []string{"-80", "80-"}},
+		{4, []string{"-40", "40-80", "80-c0", "c0-"}},
+		{20, []string{"-0c", "0c-18", "18-24", "24-30", "30-3c", "3c-48", "48-54", "54-60", "60-6c", "6c-78", "78-84", "84-90", "90-9c", "9c-a8", "a8-b4", "b4-c0", "c0-cc", "cc-d8", "d8-e4", "e4-"}},
+	}
+
+	for _, tt := range tests {
+		t.Run(fmt.Sprintf("%d shards", tt.numberOfShards), func(t *testing.T) {
+			result := generateShardRanges(tt.numberOfShards)
+			assert.Equal(t, tt.expected, result)
+		})
+	}
+}

go/tester/run.go

@@ -17,12 +17,12 @@ limitations under the License.
 package tester
 
 import (
-	"github.com/fatih/color"
 	"io"
 	"net/http"
 	"os"
 	"time"
 
+	"github.com/fatih/color"
 	log "github.com/sirupsen/logrus"
 	"vitess.io/vitess/go/test/endtoend/cluster"
 )
@@ -36,6 +36,14 @@ type Config struct {
 	VtExplainVschemaFile string
 	TraceFile            string
 	Tests                []string
+	NumberOfShards       int
+}
+
+func (cfg Config) GetNumberOfShards() int {
+	if cfg.NumberOfShards == 0 {
+		return 2
+	}
+	return cfg.NumberOfShards
 }
 
 func Run(cfg Config) {
@@ -51,6 +59,11 @@ func Run(cfg Config) {
 		os.Exit(1)
 	}
 
+	if cfg.NumberOfShards > 0 && !(cfg.Sharded || cfg.VschemaFile != "" || cfg.VtExplainVschemaFile != "") {
+		log.Errorf("number-of-shards can only be used with -sharded, -vschema or -vtexplain-vschema")
+		os.Exit(1)
+	}
+
 	if ll := os.Getenv("LOG_LEVEL"); ll != "" {
 		cfg.LogLevel = ll
 	}
@@ -69,7 +82,7 @@ func Run(cfg Config) {
 
 	log.Infof("running tests: %v", cfg.Tests)
 
-	clusterInstance, vtParams, mysqlParams, ksNames, closer := SetupCluster(cfg.VschemaFile, cfg.VtExplainVschemaFile, cfg.Sharded)
+	clusterInstance, vtParams, mysqlParams, ksNames, closer := SetupCluster(cfg.VschemaFile, cfg.VtExplainVschemaFile, cfg.Sharded, cfg.GetNumberOfShards())
 	defer closer()
 
 	// remove errors folder if exists