vitess.io/vitess@v0.16.2/go/test/endtoend/vtgate/gen4/main_test.go (about)

     1  /*
     2  Copyright 2021 The Vitess Authors.
     3  
     4  Licensed under the Apache License, Version 2.0 (the "License");
     5  you may not use this file except in compliance with the License.
     6  You may obtain a copy of the License at
     7  
     8      http://www.apache.org/licenses/LICENSE-2.0
     9  
    10  Unless required by applicable law or agreed to in writing, software
    11  distributed under the License is distributed on an "AS IS" BASIS,
    12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  See the License for the specific language governing permissions and
    14  limitations under the License.
    15  */
    16  
    17  package vtgate
    18  
    19  import (
    20  	_ "embed"
    21  	"flag"
    22  	"fmt"
    23  	"os"
    24  	"testing"
    25  
    26  	"github.com/stretchr/testify/require"
    27  
    28  	"vitess.io/vitess/go/test/endtoend/utils"
    29  
    30  	"vitess.io/vitess/go/vt/vtgate/planbuilder"
    31  
    32  	"vitess.io/vitess/go/mysql"
    33  	"vitess.io/vitess/go/test/endtoend/cluster"
    34  )
    35  
    36  var (
    37  	clusterInstance *cluster.LocalProcessCluster
    38  	vtParams        mysql.ConnParams
    39  	mysqlParams     mysql.ConnParams
    40  	shardedKs       = "ks"
    41  	unshardedKs     = "uks"
    42  	shardedKsShards = []string{"-19a0", "19a0-20", "20-20c0", "20c0-"}
    43  	Cell            = "test"
    44  	//go:embed sharded_schema.sql
    45  	shardedSchemaSQL string
    46  
    47  	//go:embed unsharded_schema.sql
    48  	unshardedSchemaSQL string
    49  
    50  	//go:embed sharded_vschema.json
    51  	shardedVSchema string
    52  
    53  	//go:embed unsharded_vschema.json
    54  	unshardedVSchema string
    55  
    56  	routingRules = `
    57  {"rules": [
    58    {
    59      "from_table": "ks.t1000",
    60  	"to_tables": ["ks.t1"]
    61    }
    62  ]}
    63  `
    64  )
    65  
    66  func TestMain(m *testing.M) {
    67  	defer cluster.PanicHandler(nil)
    68  	flag.Parse()
    69  
    70  	exitCode := func() int {
    71  		clusterInstance = cluster.NewCluster(Cell, "localhost")
    72  		defer clusterInstance.Teardown()
    73  
    74  		// Start topo server
    75  		err := clusterInstance.StartTopo()
    76  		if err != nil {
    77  			return 1
    78  		}
    79  
    80  		// Start keyspace
    81  		sKs := &cluster.Keyspace{
    82  			Name:      shardedKs,
    83  			SchemaSQL: shardedSchemaSQL,
    84  			VSchema:   shardedVSchema,
    85  		}
    86  
    87  		clusterInstance.VtGateExtraArgs = []string{"--schema_change_signal"}
    88  		clusterInstance.VtTabletExtraArgs = []string{"--queryserver-config-schema-change-signal", "--queryserver-config-schema-change-signal-interval", "0.1"}
    89  		err = clusterInstance.StartKeyspace(*sKs, shardedKsShards, 0, false)
    90  		if err != nil {
    91  			return 1
    92  		}
    93  
    94  		uKs := &cluster.Keyspace{
    95  			Name:      unshardedKs,
    96  			SchemaSQL: unshardedSchemaSQL,
    97  			VSchema:   unshardedVSchema,
    98  		}
    99  		err = clusterInstance.StartUnshardedKeyspace(*uKs, 0, false)
   100  		if err != nil {
   101  			return 1
   102  		}
   103  
   104  		// apply routing rules
   105  		err = clusterInstance.VtctlclientProcess.ApplyRoutingRules(routingRules)
   106  		if err != nil {
   107  			return 1
   108  		}
   109  
   110  		err = clusterInstance.VtctlclientProcess.ExecuteCommand("RebuildVSchemaGraph")
   111  		if err != nil {
   112  			return 1
   113  		}
   114  
   115  		// Start vtgate
   116  		clusterInstance.VtGatePlannerVersion = planbuilder.Gen4 // enable Gen4 planner.
   117  		err = clusterInstance.StartVtgate()
   118  		if err != nil {
   119  			return 1
   120  		}
   121  		vtParams = mysql.ConnParams{
   122  			Host: clusterInstance.Hostname,
   123  			Port: clusterInstance.VtgateMySQLPort,
   124  		}
   125  
   126  		conn, closer, err := utils.NewMySQL(clusterInstance, shardedKs, shardedSchemaSQL)
   127  		if err != nil {
   128  			fmt.Println(err)
   129  			return 1
   130  		}
   131  		defer closer()
   132  		mysqlParams = conn
   133  		return m.Run()
   134  	}()
   135  	os.Exit(exitCode)
   136  }
   137  
   138  func start(t *testing.T) (utils.MySQLCompare, func()) {
   139  	mcmp, err := utils.NewMySQLCompare(t, vtParams, mysqlParams)
   140  	require.NoError(t, err)
   141  	deleteAll := func() {
   142  		_, _ = utils.ExecAllowError(t, mcmp.VtConn, "set workload = oltp")
   143  
   144  		tables := []string{"t1", "t2", "t3", "user_region", "region_tbl", "multicol_tbl", "t1_id2_idx", "t2_id4_idx", "u_a", "u_b"}
   145  		for _, table := range tables {
   146  			_, _ = mcmp.ExecAndIgnore("delete from " + table)
   147  		}
   148  	}
   149  
   150  	deleteAll()
   151  
   152  	return mcmp, func() {
   153  		deleteAll()
   154  		mcmp.Close()
   155  		cluster.PanicHandler(t)
   156  	}
   157  }