vitess.io/vitess@v0.16.2/go/vt/vttablet/tabletserver/query_engine_test.go (about)

     1  /*
     2  Copyright 2019 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 tabletserver
    18  
    19  import (
    20  	"context"
    21  	"expvar"
    22  	"fmt"
    23  	"math/rand"
    24  	"net/http"
    25  	"net/http/httptest"
    26  	"os"
    27  	"path"
    28  	"reflect"
    29  	"strings"
    30  	"sync"
    31  	"sync/atomic"
    32  	"testing"
    33  	"time"
    34  
    35  	"vitess.io/vitess/go/vt/sqlparser"
    36  
    37  	"vitess.io/vitess/go/mysql"
    38  
    39  	"github.com/stretchr/testify/assert"
    40  	"github.com/stretchr/testify/require"
    41  
    42  	"vitess.io/vitess/go/cache"
    43  	"vitess.io/vitess/go/mysql/fakesqldb"
    44  	"vitess.io/vitess/go/sqltypes"
    45  	"vitess.io/vitess/go/streamlog"
    46  	"vitess.io/vitess/go/vt/dbconfigs"
    47  	"vitess.io/vitess/go/vt/tableacl"
    48  	"vitess.io/vitess/go/vt/vttablet/tabletserver/planbuilder"
    49  	"vitess.io/vitess/go/vt/vttablet/tabletserver/schema"
    50  	"vitess.io/vitess/go/vt/vttablet/tabletserver/schema/schematest"
    51  	"vitess.io/vitess/go/vt/vttablet/tabletserver/tabletenv"
    52  
    53  	querypb "vitess.io/vitess/go/vt/proto/query"
    54  )
    55  
    56  func TestStrictMode(t *testing.T) {
    57  	db := fakesqldb.New(t)
    58  	defer db.Close()
    59  	schematest.AddDefaultQueries(db)
    60  
    61  	// Test default behavior.
    62  	config := tabletenv.NewDefaultConfig()
    63  	config.DB = newDBConfigs(db)
    64  	env := tabletenv.NewEnv(config, "TabletServerTest")
    65  	se := schema.NewEngine(env)
    66  	qe := NewQueryEngine(env, se)
    67  	qe.se.InitDBConfig(newDBConfigs(db).DbaWithDB())
    68  	qe.se.Open()
    69  	if err := qe.Open(); err != nil {
    70  		t.Error(err)
    71  	}
    72  	qe.Close()
    73  
    74  	// Check that we fail if STRICT_TRANS_TABLES or STRICT_ALL_TABLES is not set.
    75  	db.AddQuery(
    76  		"select @@global.sql_mode",
    77  		&sqltypes.Result{
    78  			Fields: []*querypb.Field{{Type: sqltypes.VarChar}},
    79  			Rows:   [][]sqltypes.Value{{sqltypes.NewVarBinary("")}},
    80  		},
    81  	)
    82  	qe = NewQueryEngine(env, se)
    83  	err := qe.Open()
    84  	wantErr := "require sql_mode to be STRICT_TRANS_TABLES or STRICT_ALL_TABLES: got ''"
    85  	if err == nil || err.Error() != wantErr {
    86  		t.Errorf("Open: %v, want %s", err, wantErr)
    87  	}
    88  	qe.Close()
    89  
    90  	// Test that we succeed if the enforcement flag is off.
    91  	config.EnforceStrictTransTables = false
    92  	qe = NewQueryEngine(env, se)
    93  	if err := qe.Open(); err != nil {
    94  		t.Fatal(err)
    95  	}
    96  	qe.Close()
    97  }
    98  
    99  func TestGetPlanPanicDuetoEmptyQuery(t *testing.T) {
   100  	db := fakesqldb.New(t)
   101  	defer db.Close()
   102  	schematest.AddDefaultQueries(db)
   103  	qe := newTestQueryEngine(10*time.Second, true, newDBConfigs(db))
   104  	qe.se.Open()
   105  	qe.Open()
   106  	defer qe.Close()
   107  
   108  	ctx := context.Background()
   109  	logStats := tabletenv.NewLogStats(ctx, "GetPlanStats")
   110  	_, err := qe.GetPlan(ctx, logStats, "", false)
   111  	require.EqualError(t, err, "Query was empty")
   112  }
   113  
   114  func addSchemaEngineQueries(db *fakesqldb.DB) {
   115  	db.AddQueryPattern(baseShowTablesPattern, &sqltypes.Result{
   116  		Fields: mysql.BaseShowTablesFields,
   117  		Rows: [][]sqltypes.Value{
   118  			mysql.BaseShowTablesRow("test_table_01", false, ""),
   119  			mysql.BaseShowTablesRow("test_table_02", false, ""),
   120  			mysql.BaseShowTablesRow("test_table_03", false, ""),
   121  			mysql.BaseShowTablesRow("seq", false, "vitess_sequence"),
   122  			mysql.BaseShowTablesRow("msg", false, "vitess_message,vt_ack_wait=30,vt_purge_after=120,vt_batch_size=1,vt_cache_size=10,vt_poller_interval=30"),
   123  		}})
   124  	db.AddQuery("show status like 'Innodb_rows_read'", sqltypes.MakeTestResult(sqltypes.MakeTestFields(
   125  		"Variable_name|Value",
   126  		"varchar|int64"),
   127  		"Innodb_rows_read|0",
   128  	))
   129  }
   130  
   131  func TestGetMessageStreamPlan(t *testing.T) {
   132  	db := fakesqldb.New(t)
   133  	defer db.Close()
   134  	schematest.AddDefaultQueries(db)
   135  
   136  	addSchemaEngineQueries(db)
   137  
   138  	qe := newTestQueryEngine(10*time.Second, true, newDBConfigs(db))
   139  	qe.se.Open()
   140  	qe.Open()
   141  	defer qe.Close()
   142  
   143  	plan, err := qe.GetMessageStreamPlan("msg")
   144  	if err != nil {
   145  		t.Fatal(err)
   146  	}
   147  	wantPlan := &planbuilder.Plan{
   148  		PlanID: planbuilder.PlanMessageStream,
   149  		Table:  qe.tables["msg"],
   150  		Permissions: []planbuilder.Permission{{
   151  			TableName: "msg",
   152  			Role:      tableacl.WRITER,
   153  		}},
   154  	}
   155  	if !reflect.DeepEqual(plan.Plan, wantPlan) {
   156  		t.Errorf("GetMessageStreamPlan(msg): %v, want %v", plan.Plan, wantPlan)
   157  	}
   158  	if plan.Rules == nil || plan.Authorized == nil {
   159  		t.Errorf("GetMessageStreamPlan(msg): Rules or ACLResult are nil. Rules: %v, Authorized: %v", plan.Rules, plan.Authorized)
   160  	}
   161  }
   162  
   163  func assertPlanCacheSize(t *testing.T, qe *QueryEngine, expected int) {
   164  	t.Helper()
   165  	var size int
   166  	qe.plans.Wait()
   167  	qe.plans.ForEach(func(_ any) bool {
   168  		size++
   169  		return true
   170  	})
   171  	require.Equal(t, expected, size, "expected query plan cache to contain %d entries, found %d", expected, size)
   172  }
   173  
   174  func TestQueryPlanCache(t *testing.T) {
   175  	db := fakesqldb.New(t)
   176  	defer db.Close()
   177  	schematest.AddDefaultQueries(db)
   178  
   179  	firstQuery := "select * from test_table_01"
   180  	secondQuery := "select * from test_table_02"
   181  	db.AddQuery("select * from test_table_01 where 1 != 1", &sqltypes.Result{})
   182  	db.AddQuery("select * from test_table_02 where 1 != 1", &sqltypes.Result{})
   183  
   184  	qe := newTestQueryEngine(10*time.Second, true, newDBConfigs(db))
   185  	qe.se.Open()
   186  	qe.Open()
   187  	defer qe.Close()
   188  
   189  	ctx := context.Background()
   190  	logStats := tabletenv.NewLogStats(ctx, "GetPlanStats")
   191  	if cache.DefaultConfig.LFU {
   192  		// this cache capacity is in bytes
   193  		qe.SetQueryPlanCacheCap(528)
   194  	} else {
   195  		// this cache capacity is in number of elements
   196  		qe.SetQueryPlanCacheCap(1)
   197  	}
   198  	firstPlan, err := qe.GetPlan(ctx, logStats, firstQuery, false)
   199  	require.NoError(t, err)
   200  	require.NotNil(t, firstPlan, "plan should not be nil")
   201  	secondPlan, err := qe.GetPlan(ctx, logStats, secondQuery, false)
   202  	fmt.Println(secondPlan.CachedSize(true))
   203  	require.NoError(t, err)
   204  	require.NotNil(t, secondPlan, "plan should not be nil")
   205  	expvar.Do(func(kv expvar.KeyValue) {
   206  		_ = kv.Value.String()
   207  	})
   208  	assertPlanCacheSize(t, qe, 1)
   209  	qe.ClearQueryPlanCache()
   210  }
   211  
   212  func TestNoQueryPlanCache(t *testing.T) {
   213  	db := fakesqldb.New(t)
   214  	defer db.Close()
   215  	schematest.AddDefaultQueries(db)
   216  
   217  	firstQuery := "select * from test_table_01"
   218  	db.AddQuery("select * from test_table_01 where 1 != 1", &sqltypes.Result{})
   219  	db.AddQuery("select * from test_table_02 where 1 != 1", &sqltypes.Result{})
   220  
   221  	qe := newTestQueryEngine(10*time.Second, true, newDBConfigs(db))
   222  	qe.se.Open()
   223  	qe.Open()
   224  	defer qe.Close()
   225  
   226  	ctx := context.Background()
   227  	logStats := tabletenv.NewLogStats(ctx, "GetPlanStats")
   228  	qe.SetQueryPlanCacheCap(1024)
   229  	firstPlan, err := qe.GetPlan(ctx, logStats, firstQuery, true)
   230  	if err != nil {
   231  		t.Fatal(err)
   232  	}
   233  	if firstPlan == nil {
   234  		t.Fatalf("plan should not be nil")
   235  	}
   236  	assertPlanCacheSize(t, qe, 0)
   237  	qe.ClearQueryPlanCache()
   238  }
   239  
   240  func TestNoQueryPlanCacheDirective(t *testing.T) {
   241  	db := fakesqldb.New(t)
   242  	defer db.Close()
   243  	schematest.AddDefaultQueries(db)
   244  
   245  	firstQuery := "select /*vt+ SKIP_QUERY_PLAN_CACHE=1 */ * from test_table_01"
   246  	db.AddQuery("select * from test_table_01 where 1 != 1", &sqltypes.Result{})
   247  	db.AddQuery("select /*vt+ SKIP_QUERY_PLAN_CACHE=1 */ * from test_table_01 where 1 != 1", &sqltypes.Result{})
   248  	db.AddQuery("select /*vt+ SKIP_QUERY_PLAN_CACHE=1 */ * from test_table_02 where 1 != 1", &sqltypes.Result{})
   249  
   250  	qe := newTestQueryEngine(10*time.Second, true, newDBConfigs(db))
   251  	qe.se.Open()
   252  	qe.Open()
   253  	defer qe.Close()
   254  
   255  	ctx := context.Background()
   256  	logStats := tabletenv.NewLogStats(ctx, "GetPlanStats")
   257  	qe.SetQueryPlanCacheCap(1024)
   258  	firstPlan, err := qe.GetPlan(ctx, logStats, firstQuery, false)
   259  	if err != nil {
   260  		t.Fatal(err)
   261  	}
   262  	if firstPlan == nil {
   263  		t.Fatalf("plan should not be nil")
   264  	}
   265  	assertPlanCacheSize(t, qe, 0)
   266  	qe.ClearQueryPlanCache()
   267  }
   268  
   269  func TestStatsURL(t *testing.T) {
   270  	db := fakesqldb.New(t)
   271  	defer db.Close()
   272  	schematest.AddDefaultQueries(db)
   273  	query := "select * from test_table_01"
   274  	db.AddQuery("select * from test_table_01 where 1 != 1", &sqltypes.Result{})
   275  	qe := newTestQueryEngine(1*time.Second, true, newDBConfigs(db))
   276  	qe.se.Open()
   277  	qe.Open()
   278  	defer qe.Close()
   279  	// warm up cache
   280  	ctx := context.Background()
   281  	logStats := tabletenv.NewLogStats(ctx, "GetPlanStats")
   282  	qe.GetPlan(ctx, logStats, query, false)
   283  
   284  	request, _ := http.NewRequest("GET", "/debug/tablet_plans", nil)
   285  	response := httptest.NewRecorder()
   286  	qe.handleHTTPQueryPlans(response, request)
   287  
   288  	request, _ = http.NewRequest("GET", "/debug/query_stats", nil)
   289  	response = httptest.NewRecorder()
   290  	qe.handleHTTPQueryStats(response, request)
   291  
   292  	request, _ = http.NewRequest("GET", "/debug/query_rules", nil)
   293  	response = httptest.NewRecorder()
   294  	qe.handleHTTPQueryRules(response, request)
   295  }
   296  
   297  func newTestQueryEngine(idleTimeout time.Duration, strict bool, dbcfgs *dbconfigs.DBConfigs) *QueryEngine {
   298  	config := tabletenv.NewDefaultConfig()
   299  	config.DB = dbcfgs
   300  	config.OltpReadPool.IdleTimeoutSeconds.Set(idleTimeout)
   301  	config.OlapReadPool.IdleTimeoutSeconds.Set(idleTimeout)
   302  	config.TxPool.IdleTimeoutSeconds.Set(idleTimeout)
   303  	env := tabletenv.NewEnv(config, "TabletServerTest")
   304  	se := schema.NewEngine(env)
   305  	qe := NewQueryEngine(env, se)
   306  	se.InitDBConfig(dbcfgs.DbaWithDB())
   307  	return qe
   308  }
   309  
   310  func runConsolidatedQuery(t *testing.T, sql string) *QueryEngine {
   311  	db := fakesqldb.New(t)
   312  	defer db.Close()
   313  
   314  	qe := newTestQueryEngine(1*time.Second, true, newDBConfigs(db))
   315  	qe.se.Open()
   316  	qe.Open()
   317  	defer qe.Close()
   318  
   319  	r1, ok := qe.consolidator.Create(sql)
   320  	if !ok {
   321  		t.Errorf("expected first consolidator ok")
   322  	}
   323  	r2, ok := qe.consolidator.Create(sql)
   324  	if ok {
   325  		t.Errorf("expected second consolidator not ok")
   326  	}
   327  
   328  	r1.Broadcast()
   329  	r2.Wait()
   330  
   331  	return qe
   332  }
   333  
   334  func TestConsolidationsUIRedaction(t *testing.T) {
   335  	// Reset to default redaction state.
   336  	defer func() {
   337  		streamlog.SetRedactDebugUIQueries(false)
   338  	}()
   339  
   340  	request, _ := http.NewRequest("GET", "/debug/consolidations", nil)
   341  
   342  	sql := "select * from test_db_01 where col = 'secret'"
   343  	redactedSQL := "select * from test_db_01 where col = :col"
   344  
   345  	// First with the redaction off
   346  	streamlog.SetRedactDebugUIQueries(false)
   347  	unRedactedResponse := httptest.NewRecorder()
   348  	qe := runConsolidatedQuery(t, sql)
   349  
   350  	qe.handleHTTPConsolidations(unRedactedResponse, request)
   351  	if !strings.Contains(unRedactedResponse.Body.String(), sql) {
   352  		t.Fatalf("Response is missing the consolidated query: %v %v", sql, unRedactedResponse.Body.String())
   353  	}
   354  
   355  	// Now with the redaction on
   356  	streamlog.SetRedactDebugUIQueries(true)
   357  	redactedResponse := httptest.NewRecorder()
   358  	qe.handleHTTPConsolidations(redactedResponse, request)
   359  
   360  	if strings.Contains(redactedResponse.Body.String(), "secret") {
   361  		t.Fatalf("Response contains unredacted consolidated query: %v %v", sql, redactedResponse.Body.String())
   362  	}
   363  
   364  	if !strings.Contains(redactedResponse.Body.String(), redactedSQL) {
   365  		t.Fatalf("Response missing redacted consolidated query: %v %v", redactedSQL, redactedResponse.Body.String())
   366  	}
   367  }
   368  
   369  func BenchmarkPlanCacheThroughput(b *testing.B) {
   370  	db := fakesqldb.New(b)
   371  	defer db.Close()
   372  
   373  	schematest.AddDefaultQueries(db)
   374  
   375  	db.AddQueryPattern(".*", &sqltypes.Result{})
   376  
   377  	qe := newTestQueryEngine(10*time.Second, true, newDBConfigs(db))
   378  	qe.se.Open()
   379  	qe.Open()
   380  	defer qe.Close()
   381  
   382  	ctx := context.Background()
   383  	logStats := tabletenv.NewLogStats(ctx, "GetPlanStats")
   384  
   385  	for i := 0; i < b.N; i++ {
   386  		query := fmt.Sprintf("SELECT (a, b, c) FROM test_table_%d", rand.Intn(500))
   387  		_, err := qe.GetPlan(ctx, logStats, query, false)
   388  		if err != nil {
   389  			b.Fatal(err)
   390  		}
   391  	}
   392  }
   393  
   394  func benchmarkPlanCache(b *testing.B, db *fakesqldb.DB, lfu bool, par int) {
   395  	b.Helper()
   396  
   397  	dbcfgs := newDBConfigs(db)
   398  	config := tabletenv.NewDefaultConfig()
   399  	config.DB = dbcfgs
   400  	config.QueryCacheLFU = lfu
   401  
   402  	env := tabletenv.NewEnv(config, "TabletServerTest")
   403  	se := schema.NewEngine(env)
   404  	qe := NewQueryEngine(env, se)
   405  
   406  	se.InitDBConfig(dbcfgs.DbaWithDB())
   407  	require.NoError(b, se.Open())
   408  	require.NoError(b, qe.Open())
   409  	defer qe.Close()
   410  
   411  	b.SetParallelism(par)
   412  	b.RunParallel(func(pb *testing.PB) {
   413  		ctx := context.Background()
   414  		logStats := tabletenv.NewLogStats(ctx, "GetPlanStats")
   415  
   416  		for pb.Next() {
   417  			query := fmt.Sprintf("SELECT (a, b, c) FROM test_table_%d", rand.Intn(500))
   418  			_, err := qe.GetPlan(ctx, logStats, query, false)
   419  			require.NoErrorf(b, err, "bad query: %s", query)
   420  		}
   421  	})
   422  }
   423  
   424  func BenchmarkPlanCacheContention(b *testing.B) {
   425  	db := fakesqldb.New(b)
   426  	defer db.Close()
   427  
   428  	schematest.AddDefaultQueries(db)
   429  
   430  	db.AddQueryPattern(".*", &sqltypes.Result{})
   431  
   432  	for par := 1; par <= 8; par *= 2 {
   433  		b.Run(fmt.Sprintf("ContentionLRU-%d", par), func(b *testing.B) {
   434  			benchmarkPlanCache(b, db, false, par)
   435  		})
   436  
   437  		b.Run(fmt.Sprintf("ContentionLFU-%d", par), func(b *testing.B) {
   438  			benchmarkPlanCache(b, db, true, par)
   439  		})
   440  	}
   441  }
   442  
   443  func TestPlanCachePollution(t *testing.T) {
   444  	plotPath := os.Getenv("CACHE_PLOT_PATH")
   445  	if plotPath == "" {
   446  		t.Skipf("CACHE_PLOT_PATH not set")
   447  	}
   448  
   449  	const NormalQueries = 500000
   450  	const PollutingQueries = NormalQueries / 2
   451  
   452  	db := fakesqldb.New(t)
   453  	defer db.Close()
   454  
   455  	schematest.AddDefaultQueries(db)
   456  
   457  	db.AddQueryPattern(".*", &sqltypes.Result{})
   458  
   459  	dbcfgs := newDBConfigs(db)
   460  	config := tabletenv.NewDefaultConfig()
   461  	config.DB = dbcfgs
   462  	// config.LFUQueryCacheSizeBytes = 3 * 1024 * 1024
   463  
   464  	env := tabletenv.NewEnv(config, "TabletServerTest")
   465  	se := schema.NewEngine(env)
   466  	qe := NewQueryEngine(env, se)
   467  
   468  	se.InitDBConfig(dbcfgs.DbaWithDB())
   469  	se.Open()
   470  
   471  	qe.Open()
   472  	defer qe.Close()
   473  
   474  	type Stats struct {
   475  		queries  uint64
   476  		cached   uint64
   477  		interval time.Duration
   478  	}
   479  
   480  	var stats1, stats2 Stats
   481  	var wg sync.WaitGroup
   482  
   483  	go func() {
   484  		cacheMode := "lru"
   485  		if config.QueryCacheLFU {
   486  			cacheMode = "lfu"
   487  		}
   488  
   489  		out, err := os.Create(path.Join(plotPath,
   490  			fmt.Sprintf("cache_plot_%d_%d_%s.dat",
   491  				config.QueryCacheSize, config.QueryCacheMemory, cacheMode,
   492  			)),
   493  		)
   494  		require.NoError(t, err)
   495  		defer out.Close()
   496  
   497  		var last1 uint64
   498  		var last2 uint64
   499  
   500  		for range time.Tick(100 * time.Millisecond) {
   501  			var avg1, avg2 time.Duration
   502  
   503  			if stats1.queries-last1 > 0 {
   504  				avg1 = stats1.interval / time.Duration(stats1.queries-last1)
   505  			}
   506  			if stats2.queries-last2 > 0 {
   507  				avg2 = stats2.interval / time.Duration(stats2.queries-last2)
   508  			}
   509  
   510  			stats1.interval = 0
   511  			last1 = stats1.queries
   512  			stats2.interval = 0
   513  			last2 = stats2.queries
   514  
   515  			cacheUsed, cacheCap := qe.plans.UsedCapacity(), qe.plans.MaxCapacity()
   516  
   517  			t.Logf("%d queries (%f hit rate), cache %d / %d (%f usage), %v %v",
   518  				stats1.queries+stats2.queries,
   519  				float64(stats1.cached)/float64(stats1.queries),
   520  				cacheUsed, cacheCap,
   521  				float64(cacheUsed)/float64(cacheCap), avg1, avg2)
   522  
   523  			if out != nil {
   524  				fmt.Fprintf(out, "%d %f %f %f %f %d %d\n",
   525  					stats1.queries+stats2.queries,
   526  					float64(stats1.queries)/float64(NormalQueries),
   527  					float64(stats2.queries)/float64(PollutingQueries),
   528  					float64(stats1.cached)/float64(stats1.queries),
   529  					float64(cacheUsed)/float64(cacheCap),
   530  					avg1.Microseconds(),
   531  					avg2.Microseconds(),
   532  				)
   533  			}
   534  		}
   535  	}()
   536  
   537  	runner := func(totalQueries uint64, stats *Stats, sample func() string) {
   538  		for i := uint64(0); i < totalQueries; i++ {
   539  			ctx := context.Background()
   540  			logStats := tabletenv.NewLogStats(ctx, "GetPlanStats")
   541  			query := sample()
   542  
   543  			start := time.Now()
   544  			_, err := qe.GetPlan(ctx, logStats, query, false)
   545  			require.NoErrorf(t, err, "bad query: %s", query)
   546  			stats.interval += time.Since(start)
   547  
   548  			atomic.AddUint64(&stats.queries, 1)
   549  			if logStats.CachedPlan {
   550  				atomic.AddUint64(&stats.cached, 1)
   551  			}
   552  		}
   553  	}
   554  
   555  	wg.Add(2)
   556  
   557  	go func() {
   558  		defer wg.Done()
   559  		runner(NormalQueries, &stats1, func() string {
   560  			return fmt.Sprintf("SELECT (a, b, c) FROM test_table_%d", rand.Intn(5000))
   561  		})
   562  	}()
   563  
   564  	go func() {
   565  		defer wg.Done()
   566  		time.Sleep(500 * time.Millisecond)
   567  		runner(PollutingQueries, &stats2, func() string {
   568  			return fmt.Sprintf("INSERT INTO test_table_00 VALUES (1, 2, 3, %d)", rand.Int())
   569  		})
   570  	}()
   571  
   572  	wg.Wait()
   573  }
   574  
   575  func TestAddQueryStats(t *testing.T) {
   576  	testcases := []struct {
   577  		name                      string
   578  		planType                  planbuilder.PlanType
   579  		tableName                 string
   580  		queryCount                int64
   581  		duration                  time.Duration
   582  		mysqlTime                 time.Duration
   583  		rowsAffected              int64
   584  		rowsReturned              int64
   585  		errorCount                int64
   586  		expectedQueryCounts       string
   587  		expectedQueryTimes        string
   588  		expectedQueryRowsAffected string
   589  		expectedQueryRowsReturned string
   590  		expectedQueryErrorCounts  string
   591  	}{
   592  		{
   593  			name:                      "select query",
   594  			planType:                  planbuilder.PlanSelect,
   595  			tableName:                 "A",
   596  			queryCount:                1,
   597  			duration:                  10,
   598  			rowsAffected:              0,
   599  			rowsReturned:              15,
   600  			errorCount:                0,
   601  			expectedQueryCounts:       `{"A.Select": 1}`,
   602  			expectedQueryTimes:        `{"A.Select": 10}`,
   603  			expectedQueryRowsAffected: `{}`,
   604  			expectedQueryRowsReturned: `{"A.Select": 15}`,
   605  			expectedQueryErrorCounts:  `{"A.Select": 0}`,
   606  		}, {
   607  			name:                      "select into query",
   608  			planType:                  planbuilder.PlanSelect,
   609  			tableName:                 "A",
   610  			queryCount:                1,
   611  			duration:                  10,
   612  			rowsAffected:              15,
   613  			rowsReturned:              0,
   614  			errorCount:                0,
   615  			expectedQueryCounts:       `{"A.Select": 1}`,
   616  			expectedQueryTimes:        `{"A.Select": 10}`,
   617  			expectedQueryRowsAffected: `{"A.Select": 15}`,
   618  			expectedQueryRowsReturned: `{"A.Select": 0}`,
   619  			expectedQueryErrorCounts:  `{"A.Select": 0}`,
   620  		}, {
   621  			name:                      "error",
   622  			planType:                  planbuilder.PlanSelect,
   623  			tableName:                 "A",
   624  			queryCount:                1,
   625  			duration:                  10,
   626  			rowsAffected:              0,
   627  			rowsReturned:              0,
   628  			errorCount:                1,
   629  			expectedQueryCounts:       `{"A.Select": 1}`,
   630  			expectedQueryTimes:        `{"A.Select": 10}`,
   631  			expectedQueryRowsAffected: `{}`,
   632  			expectedQueryRowsReturned: `{"A.Select": 0}`,
   633  			expectedQueryErrorCounts:  `{"A.Select": 1}`,
   634  		}, {
   635  			name:                      "insert query",
   636  			planType:                  planbuilder.PlanInsert,
   637  			tableName:                 "A",
   638  			queryCount:                1,
   639  			duration:                  10,
   640  			rowsAffected:              15,
   641  			rowsReturned:              0,
   642  			errorCount:                0,
   643  			expectedQueryCounts:       `{"A.Insert": 1}`,
   644  			expectedQueryTimes:        `{"A.Insert": 10}`,
   645  			expectedQueryRowsAffected: `{"A.Insert": 15}`,
   646  			expectedQueryRowsReturned: `{}`,
   647  			expectedQueryErrorCounts:  `{"A.Insert": 0}`,
   648  		},
   649  	}
   650  
   651  	t.Parallel()
   652  	for _, testcase := range testcases {
   653  		t.Run(testcase.name, func(t *testing.T) {
   654  			config := tabletenv.NewDefaultConfig()
   655  			config.DB = newDBConfigs(fakesqldb.New(t))
   656  			env := tabletenv.NewEnv(config, "TestAddQueryStats_"+testcase.name)
   657  			se := schema.NewEngine(env)
   658  			qe := NewQueryEngine(env, se)
   659  			qe.AddStats(testcase.planType, testcase.tableName, testcase.queryCount, testcase.duration, testcase.mysqlTime, testcase.rowsAffected, testcase.rowsReturned, testcase.errorCount)
   660  			assert.Equal(t, testcase.expectedQueryCounts, qe.queryCounts.String())
   661  			assert.Equal(t, testcase.expectedQueryTimes, qe.queryTimes.String())
   662  			assert.Equal(t, testcase.expectedQueryRowsAffected, qe.queryRowsAffected.String())
   663  			assert.Equal(t, testcase.expectedQueryRowsReturned, qe.queryRowsReturned.String())
   664  			assert.Equal(t, testcase.expectedQueryErrorCounts, qe.queryErrorCounts.String())
   665  		})
   666  	}
   667  }
   668  
   669  func TestPlanPoolUnsafe(t *testing.T) {
   670  	tcases := []struct {
   671  		name, query, err string
   672  	}{
   673  		{
   674  			"get_lock named locks are unsafe with server-side connection pooling",
   675  			"select get_lock('foo', 10) from dual",
   676  			"SelectLockFunc not allowed without reserved connection",
   677  		}, {
   678  			"setting system variables must happen inside reserved connections",
   679  			"set sql_safe_updates = false",
   680  			"Set not allowed without reserved connection",
   681  		}, {
   682  			"setting system variables must happen inside reserved connections",
   683  			"set @@sql_safe_updates = false",
   684  			"Set not allowed without reserved connection",
   685  		}, {
   686  			"setting system variables must happen inside reserved connections",
   687  			"set @udv = false",
   688  			"Set not allowed without reserved connection",
   689  		},
   690  	}
   691  	for _, tcase := range tcases {
   692  		t.Run(tcase.name, func(t *testing.T) {
   693  			statement, err := sqlparser.Parse(tcase.query)
   694  			require.NoError(t, err)
   695  			plan, err := planbuilder.Build(statement, map[string]*schema.Table{}, "dbName", false)
   696  			// Plan building will not fail, but it will mark that reserved connection is needed.
   697  			// checking plan is valid will fail.
   698  			require.NoError(t, err)
   699  			require.True(t, plan.NeedsReservedConn)
   700  			err = isValid(plan.PlanID, false, false)
   701  			require.EqualError(t, err, tcase.err)
   702  		})
   703  	}
   704  }