context request type and trace id

Author: asmyasnikov

internal/meta/context.go

@@ -6,12 +6,17 @@ import (
 	"google.golang.org/grpc/metadata"
 )
 
-// WithTraceID returns a copy of parent deadline with traceID
+// WithTraceID returns a copy of parent context with traceID
 func WithTraceID(ctx context.Context, traceID string) context.Context {
 	return metadata.AppendToOutgoingContext(ctx, HeaderTraceID, traceID)
 }
 
-// WithUserAgent returns a copy of parent deadline with custom user-agent info
+// WithUserAgent returns a copy of parent context with custom user-agent info
 func WithUserAgent(ctx context.Context, userAgent string) context.Context {
 	return metadata.AppendToOutgoingContext(ctx, HeaderUserAgent, userAgent)
 }
+
+// WithRequestType returns a copy of parent context with custom request type
+func WithRequestType(ctx context.Context, requestType string) context.Context {
+	return metadata.AppendToOutgoingContext(ctx, HeaderRequestType, requestType)
+}

meta.go

@@ -0,0 +1,17 @@
+package ydb
+
+import (
+	"context"
+
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/meta"
+)
+
+// WithTraceID returns a copy of parent context with traceID
+func WithTraceID(ctx context.Context, traceID string) context.Context {
+	return meta.WithTraceID(ctx, traceID)
+}
+
+// WithRequestType returns a copy of parent context with custom request type
+func WithRequestType(ctx context.Context, requestType string) context.Context {
+	return meta.WithRequestType(ctx, requestType)
+}

options.go

@@ -61,6 +61,13 @@ func WithUserAgent(userAgent string) Option {
 	}
 }
 
+func WithRequestsType(requestsType string) Option {
+	return func(ctx context.Context, c *connection) error {
+		c.options = append(c.options, config.WithRequestsType(requestsType))
+		return nil
+	}
+}
+
 // WithConnectionString accept connection string like 'grpc[s]://{endpoint}/?database={database}'
 func WithConnectionString(connectionString string) Option {
 	return func(ctx context.Context, c *connection) error {