vitess.io/vitess@v0.16.2/go/test/endtoend/reparent/prssettingspool/main_test.go (about)

     1  /*
     2  Copyright 2023 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 misc
    18  
    19  import (
    20  	"context"
    21  	_ "embed"
    22  	"flag"
    23  	"os"
    24  	"testing"
    25  	"time"
    26  
    27  	"github.com/stretchr/testify/require"
    28  
    29  	"vitess.io/vitess/go/mysql"
    30  	"vitess.io/vitess/go/test/endtoend/cluster"
    31  	rutils "vitess.io/vitess/go/test/endtoend/reparent/utils"
    32  	"vitess.io/vitess/go/test/endtoend/utils"
    33  )
    34  
    35  var (
    36  	clusterInstance *cluster.LocalProcessCluster
    37  	vtParams        mysql.ConnParams
    38  	keyspaceName    = "ks"
    39  	cell            = "test"
    40  
    41  	//go:embed schema.sql
    42  	schemaSQL string
    43  )
    44  
    45  func TestMain(m *testing.M) {
    46  	defer cluster.PanicHandler(nil)
    47  	flag.Parse()
    48  
    49  	exitCode := func() int {
    50  		clusterInstance = cluster.NewCluster(cell, "localhost")
    51  		defer clusterInstance.Teardown()
    52  
    53  		// Start topo server
    54  		err := clusterInstance.StartTopo()
    55  		if err != nil {
    56  			return 1
    57  		}
    58  
    59  		// Start keyspace
    60  		keyspace := &cluster.Keyspace{
    61  			Name:      keyspaceName,
    62  			SchemaSQL: schemaSQL,
    63  		}
    64  		clusterInstance.VtTabletExtraArgs = append(clusterInstance.VtTabletExtraArgs,
    65  			"--queryserver-enable-settings-pool")
    66  		err = clusterInstance.StartUnshardedKeyspace(*keyspace, 2, false)
    67  		if err != nil {
    68  			return 1
    69  		}
    70  
    71  		// Start vtgate
    72  		clusterInstance.VtGateExtraArgs = append(clusterInstance.VtGateExtraArgs,
    73  			"--planner-version", "gen4")
    74  		err = clusterInstance.StartVtgate()
    75  		if err != nil {
    76  			return 1
    77  		}
    78  
    79  		vtParams = mysql.ConnParams{
    80  			Host: clusterInstance.Hostname,
    81  			Port: clusterInstance.VtgateMySQLPort,
    82  		}
    83  		return m.Run()
    84  	}()
    85  	os.Exit(exitCode)
    86  }
    87  
    88  func TestSettingsPoolWithTXAndPRS(t *testing.T) {
    89  	ctx := context.Background()
    90  	conn, err := mysql.Connect(ctx, &vtParams)
    91  	require.NoError(t, err)
    92  	defer conn.Close()
    93  
    94  	// set a system settings that will trigger reserved connection usage.
    95  	utils.Exec(t, conn, "set default_week_format = 5")
    96  
    97  	// have transaction on the session
    98  	utils.Exec(t, conn, "begin")
    99  	utils.Exec(t, conn, "select id1, id2 from t1")
   100  	utils.Exec(t, conn, "commit")
   101  
   102  	tablets := clusterInstance.Keyspaces[0].Shards[0].Vttablets
   103  
   104  	// prs should happen without any error.
   105  	text, err := rutils.Prs(t, clusterInstance, tablets[1])
   106  	require.NoError(t, err, text)
   107  	rutils.WaitForTabletToBeServing(t, clusterInstance, tablets[0], 1*time.Minute)
   108  
   109  	defer func() {
   110  		// reset state
   111  		text, err = rutils.Prs(t, clusterInstance, tablets[0])
   112  		require.NoError(t, err, text)
   113  		rutils.WaitForTabletToBeServing(t, clusterInstance, tablets[1], 1*time.Minute)
   114  	}()
   115  
   116  	// no error should occur and it should go to the right tablet.
   117  	utils.Exec(t, conn, "select id1, id2 from t1")
   118  }
   119  
   120  func TestSettingsPoolWithoutTXAndPRS(t *testing.T) {
   121  	ctx := context.Background()
   122  	conn, err := mysql.Connect(ctx, &vtParams)
   123  	require.NoError(t, err)
   124  	defer conn.Close()
   125  
   126  	// set a system settings that will trigger reserved connection usage.
   127  	utils.Exec(t, conn, "set default_week_format = 5")
   128  
   129  	// execute non-tx query
   130  	utils.Exec(t, conn, "select id1, id2 from t1")
   131  
   132  	tablets := clusterInstance.Keyspaces[0].Shards[0].Vttablets
   133  
   134  	// prs should happen without any error.
   135  	text, err := rutils.Prs(t, clusterInstance, tablets[1])
   136  	require.NoError(t, err, text)
   137  	rutils.WaitForTabletToBeServing(t, clusterInstance, tablets[0], 1*time.Minute)
   138  	defer func() {
   139  		// reset state
   140  		text, err = rutils.Prs(t, clusterInstance, tablets[0])
   141  		require.NoError(t, err, text)
   142  		rutils.WaitForTabletToBeServing(t, clusterInstance, tablets[1], 1*time.Minute)
   143  	}()
   144  
   145  	// no error should occur and it should go to the right tablet.
   146  	utils.Exec(t, conn, "select id1, id2 from t1")
   147  
   148  }