diff --git a/util/request_source.go b/util/request_source.go index 663648771..6ebadfc7e 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{