test: add e2e testing for the changes

Signed-off-by: Manan Gupta <manan@planetscale.com>

go/test/endtoend/vtgate/queries/timeout/timeout_test.go

@@ -94,10 +94,41 @@ func TestQueryTimeoutWithTables(t *testing.T) {
 	utils.Exec(t, mcmp.VtConn, "select /*vt+ QUERY_TIMEOUT_MS=500 */ sleep(0.1) from t1 where id1 = 1")
 	_, err = utils.ExecAllowError(t, mcmp.VtConn, "select /*vt+ QUERY_TIMEOUT_MS=20 */ sleep(0.1) from t1 where id1 = 1")
 	require.Error(t, err)
-	assert.Contains(t, err.Error(), "context deadline exceeded")
+	if utils.BinaryIsAtLeastAtVersion(21, "vtgate") {
+		assert.ErrorContains(t, err, "Query execution was interrupted, maximum statement execution time exceeded")
+	} else {
+		assert.Contains(t, err.Error(), "context deadline exceeded")
+	}
 	assert.Contains(t, err.Error(), "(errno 1317) (sqlstate 70100)")
 }
 
+// TestOverallQueryTimeout tests that the query timeout is applied to the overall execution of a query
+// and not just individual routes.
+func TestOverallQueryTimeout(t *testing.T) {
+	utils.SkipIfBinaryIsBelowVersion(t, 21, "vtgate")
+	mcmp, closer := start(t)
+	defer closer()
+
+	mcmp.Exec("insert into t1(id1, id2) values (2,2),(3,3)")
+	// After inserting the rows above, if we run the following query, we will end up doing join on vtgate
+	// that issues one select query on the left side and 2 on the right side. The queries on the right side
+	// take 2 and 3 seconds each to run. If we have an overall timeout for 4 seconds, then it should fail.
+
+	_, err := utils.ExecAllowError(t, mcmp.VtConn, "select /*vt+ QUERY_TIMEOUT_MS=4000 */ sleep(u2.id2), u1.id2 from t1 u1 join t1 u2 where u1.id2 = u2.id1")
+	assert.Error(t, err)
+	assert.ErrorContains(t, err, "Query execution was interrupted, maximum statement execution time exceeded")
+
+	// Let's also check that setting the session variable also works.
+	utils.Exec(t, mcmp.VtConn, "set query_timeout=4000")
+	_, err = utils.ExecAllowError(t, mcmp.VtConn, "select sleep(u2.id2), u1.id2 from t1 u1 join t1 u2 where u1.id2 = u2.id1")
+	assert.Error(t, err)
+	assert.ErrorContains(t, err, "Query execution was interrupted, maximum statement execution time exceeded")
+
+	// Increasing the timeout should pass the query.
+	utils.Exec(t, mcmp.VtConn, "set query_timeout=10000")
+	_ = utils.Exec(t, mcmp.VtConn, "select sleep(u2.id2), u1.id2 from t1 u1 join t1 u2 where u1.id2 = u2.id1")
+}
+
 // TestQueryTimeoutWithShardTargeting tests the query timeout with shard targeting.
 func TestQueryTimeoutWithShardTargeting(t *testing.T) {
 	mcmp, closer := start(t)

go/vt/vtgate/engine/timeout_handler_test.go

@@ -13,35 +13,32 @@ import (
 // TestTimeoutHandler tests timeout handler primitive.
 func TestTimeoutHandler(t *testing.T) {
 	tests := []struct {
-		name    string
-		input   *TimeoutHandler
-		wantErr string
+		name      string
+		sleepTime time.Duration
+		timeout   int
+		wantErr   string
 	}{
 		{
-			name: "Timeout without failure",
-			input: NewTimeoutHandler(&fakePrimitive{
-				results:   nil,
-				sleepTime: 100 * time.Millisecond,
-			}, 1000),
-			wantErr: "",
+			name:      "Timeout without failure",
+			sleepTime: 100 * time.Millisecond,
+			timeout:   1000,
+			wantErr:   "",
 		}, {
-			name: "Timeout with failure",
-			input: NewTimeoutHandler(&fakePrimitive{
-				results:   nil,
-				sleepTime: 2 * time.Second,
-			}, 100),
-			wantErr: "VT15001: Query execution was interrupted, maximum statement execution time exceeded",
+			name:      "Timeout with failure",
+			sleepTime: 2 * time.Second,
+			timeout:   100,
+			wantErr:   "VT15001: Query execution was interrupted, maximum statement execution time exceeded",
 		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			_, err := tt.input.TryExecute(context.Background(), &noopVCursor{}, nil, false)
+			_, err := createTimeoutHandlerForTesting(tt.timeout, tt.sleepTime).TryExecute(context.Background(), &noopVCursor{}, nil, false)
 			if tt.wantErr != "" {
 				require.EqualError(t, err, tt.wantErr)
 			} else {
 				require.NoError(t, err)
 			}
-			err = tt.input.TryStreamExecute(context.Background(), &noopVCursor{}, nil, false, func(result *sqltypes.Result) error {
+			err = createTimeoutHandlerForTesting(tt.timeout, tt.sleepTime).TryStreamExecute(context.Background(), &noopVCursor{}, nil, false, func(result *sqltypes.Result) error {
 				return nil
 			})
 			if tt.wantErr != "" {
@@ -52,3 +49,11 @@ func TestTimeoutHandler(t *testing.T) {
 		})
 	}
 }
+
+// createTimeoutHandlerForTesting creates a TimeoutHandler for testing that has a fakePrimitive as an input.
+func createTimeoutHandlerForTesting(timeout int, sleepTime time.Duration) *TimeoutHandler {
+	return NewTimeoutHandler(&fakePrimitive{
+		results:   nil,
+		sleepTime: sleepTime,
+	}, timeout)
+}