Merge pull request #9008 from planetscale/gen4-qt-cd

Gen4: query timeout in comment directive

go/vt/vtgate/engine/route.go

@@ -820,6 +820,9 @@ func (route *Route) description() PrimitiveDescription {
 	if route.ScatterErrorsAsWarnings {
 		other["ScatterErrorsAsWarnings"] = true
 	}
+	if route.QueryTimeout > 0 {
+		other["QueryTimeout"] = route.QueryTimeout
+	}
 	return PrimitiveDescription{
 		OperatorType:      "Route",
 		Variant:           routeName[route.Opcode],

go/vt/vtgate/planbuilder/gen4_planner.go

@@ -177,15 +177,9 @@ func newBuildSelectPlan(selStmt sqlparser.SelectStatement, reservedVars *sqlpars
 		return nil, err
 	}
 
-	directives := sqlparser.ExtractCommentDirectives(sqlparser.GetFirstSelect(selStmt).Comments)
-	if directives.IsSet(sqlparser.DirectiveScatterErrorsAsWarnings) {
-		_, _ = visit(plan, func(logicalPlan logicalPlan) (bool, logicalPlan, error) {
-			switch plan := logicalPlan.(type) {
-			case *route:
-				plan.eroute.ScatterErrorsAsWarnings = true
-			}
-			return true, logicalPlan, nil
-		})
+	plan, err = pushCommentDirectivesOnPlan(plan, selStmt)
+	if err != nil {
+		return nil, err
 	}
 
 	return plan, nil
@@ -280,3 +274,24 @@ func planOrderByOnUnion(ctx *planningContext, plan logicalPlan, union *sqlparser
 	}
 	return plan, nil
 }
+
+func pushCommentDirectivesOnPlan(plan logicalPlan, stmt sqlparser.SelectStatement) (logicalPlan, error) {
+	directives := sqlparser.ExtractCommentDirectives(sqlparser.GetFirstSelect(stmt).Comments)
+	scatterAsWarns := false
+	if directives.IsSet(sqlparser.DirectiveScatterErrorsAsWarnings) {
+		scatterAsWarns = true
+	}
+	queryTimeout := queryTimeout(directives)
+	if scatterAsWarns || queryTimeout > 0 {
+		_, _ = visit(plan, func(logicalPlan logicalPlan) (bool, logicalPlan, error) {
+			switch plan := logicalPlan.(type) {
+			case *route:
+				plan.eroute.ScatterErrorsAsWarnings = scatterAsWarns
+				plan.eroute.QueryTimeout = queryTimeout
+			}
+			return true, logicalPlan, nil
+		})
+	}
+
+	return plan, nil
+}

go/vt/vtgate/planbuilder/testdata/select_cases.txt

@@ -69,6 +69,7 @@ Gen4 plan same as above
     },
     "FieldQuery": "select * from `user` where 1 != 1",
     "Query": "select /*vt+ QUERY_TIMEOUT_MS=1000 */ * from `user`",
+    "QueryTimeout": 1000,
     "Table": "`user`"
   }
 }
@@ -93,6 +94,7 @@ Gen4 plan same as above
         },
         "FieldQuery": "select count(*) from `user` where 1 != 1",
         "Query": "select /*vt+ QUERY_TIMEOUT_MS=1000 */ count(*) from `user`",
+        "QueryTimeout": 1000,
         "Table": "`user`"
       }
     ]
@@ -115,6 +117,7 @@ Gen4 plan same as above
         },
         "FieldQuery": "select count(*) from `user` where 1 != 1",
         "Query": "select /*vt+ QUERY_TIMEOUT_MS=1000 */ count(*) from `user`",
+        "QueryTimeout": 1000,
         "Table": "`user`"
       }
     ]
@@ -139,6 +142,7 @@ Gen4 plan same as above
         },
         "FieldQuery": "select * from `user` where 1 != 1",
         "Query": "select /*vt+ QUERY_TIMEOUT_MS=1000 */ * from `user` limit :__upper_limit",
+        "QueryTimeout": 1000,
         "Table": "`user`"
       }
     ]
@@ -1597,6 +1601,7 @@ Gen4 plan same as above
     },
     "FieldQuery": "select * from unsharded as route2 where 1 != 1",
     "Query": "select /*vt+ QUERY_TIMEOUT_MS=1000 */ * from unsharded as route2",
+    "QueryTimeout": 1000,
     "Table": "unsharded"
   }
 }