From 300545a8a3c4bb6f5701f6ea4b5c34a52842d0d0 Mon Sep 17 00:00:00 2001 From: glorv Date: Wed, 9 Aug 2023 13:03:15 +0800 Subject: [PATCH] add a function to set request source task type (#925) * add a function to set request source task type Signed-off-by: glorv --- util/request_source.go | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/util/request_source.go b/util/request_source.go index 66364877..6ebadfc7 100644 --- a/util/request_source.go +++ b/util/request_source.go @@ -48,10 +48,11 @@ const ( ExplicitTypeBR = "br" ExplicitTypeDumpling = "dumpling" ExplicitTypeBackground = "background" + ExplicitTypeDDL = "ddl" ) // ExplicitTypeList is the list of all explicit source types. -var ExplicitTypeList = []string{ExplicitTypeEmpty, ExplicitTypeLightning, ExplicitTypeBR, ExplicitTypeDumpling, ExplicitTypeBackground} +var ExplicitTypeList = []string{ExplicitTypeEmpty, ExplicitTypeLightning, ExplicitTypeBR, ExplicitTypeDumpling, ExplicitTypeBackground, ExplicitTypeDDL} const ( // InternalRequest is the scope of internal queries @@ -94,6 +95,15 @@ func WithInternalSourceType(ctx context.Context, source string) context.Context }) } +// WithInternalSourceAndTaskType create context with internal source and task name. +func WithInternalSourceAndTaskType(ctx context.Context, source, taskName string) context.Context { + return context.WithValue(ctx, RequestSourceKey, RequestSource{ + RequestSourceInternal: true, + RequestSourceType: source, + ExplicitRequestSourceType: taskName, + }) +} + // BuildRequestSource builds a request_source from internal, source and explicitSource. func BuildRequestSource(internal bool, source, explicitSource string) string { requestSource := RequestSource{