Add a test to verify we respect the overall query timeout (#16800)

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

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

@@ -179,3 +179,30 @@ func TestQueryTimeoutWithoutVTGateDefault(t *testing.T) {
 	_, err = utils.ExecAllowError(t, mcmp.VtConn, "select sleep(5) from dual")
 	assert.Error(t, err)
 }
+
+// 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, "DeadlineExceeded desc = context deadline exceeded (errno 1317) (sqlstate 70100)")
+
+	// 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, "DeadlineExceeded desc = context deadline exceeded (errno 1317) (sqlstate 70100)")
+
+	// 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")
+}