diff --git a/go/vt/vtgate/engine/route.go b/go/vt/vtgate/engine/route.go index 5879252b6dc..77c60ea3c64 100644 --- a/go/vt/vtgate/engine/route.go +++ b/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], diff --git a/go/vt/vtgate/planbuilder/gen4_planner.go b/go/vt/vtgate/planbuilder/gen4_planner.go index a804f972b22..84c64d8d2d1 100644 --- a/go/vt/vtgate/planbuilder/gen4_planner.go +++ b/go/vt/vtgate/planbuilder/gen4_planner.go @@ -167,15 +167,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 @@ -270,3 +264,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 +} diff --git a/go/vt/vtgate/planbuilder/testdata/select_cases.txt b/go/vt/vtgate/planbuilder/testdata/select_cases.txt index d160744ca94..7ce44b59e37 100644 --- a/go/vt/vtgate/planbuilder/testdata/select_cases.txt +++ b/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" } }