diff --git a/telemetryutils/factory.go b/telemetryutils/factory.go index a6ba0cc..b9bc258 100644 --- a/telemetryutils/factory.go +++ b/telemetryutils/factory.go @@ -1,10 +1,13 @@ package telemetryutils import ( + "context" "os" + "github.com/flyteorg/flytestdlib/contextutils" "github.com/flyteorg/flytestdlib/version" + "go.opentelemetry.io/otel/attribute" "go.opentelemetry.io/otel/exporters/jaeger" "go.opentelemetry.io/otel/exporters/stdout/stdouttrace" "go.opentelemetry.io/otel/sdk/resource" @@ -93,3 +96,15 @@ func GetTracerProvider(serviceName string) rawtrace.TracerProvider { return noopTracerProvider } + +func NewSpan(ctx context.Context, serviceName string, spanName string) (context.Context, rawtrace.Span) { + var attributes []attribute.KeyValue + for key, value := range contextutils.GetLogFields(ctx) { + if value, ok := value.(string); ok { + attributes = append(attributes, attribute.String(key, value)) + } + } + + tracerProvider := GetTracerProvider(serviceName) + return tracerProvider.Tracer("default").Start(ctx, spanName, rawtrace.WithAttributes(attributes...)) +} diff --git a/telemetryutils/factory_test.go b/telemetryutils/factory_test.go new file mode 100644 index 0000000..e3647ab --- /dev/null +++ b/telemetryutils/factory_test.go @@ -0,0 +1,42 @@ +package telemetryutils + +import ( + "context" + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestRegisterTracerProvider(t *testing.T) { + serviceName := "foo" + + // register tracer provider with no exporters + emptyConfig := Config{} + err := RegisterTracerProvider(serviceName, &emptyConfig) + assert.Nil(t, err) + + // validate no tracerProviders are registered + assert.Len(t, tracerProviders, 0) + + // register tracer provider with all exporters + fullConfig := Config{ + FileConfig: FileConfig{ + Enabled: true, + Filename: "/dev/null", + }, + JaegerConfig: JaegerConfig{ + Enabled: true, + }, + } + err = RegisterTracerProvider(serviceName, &fullConfig) + assert.Nil(t, err) + + // validate tracerProvider is registered + assert.Len(t, tracerProviders, 1) +} + +func TestNewSpan(t *testing.T) { + ctx := context.TODO() + _, span := NewSpan(ctx, "bar", "baz/bat") + span.End() +} diff --git a/telemetryutils/utils.go b/telemetryutils/utils.go deleted file mode 100644 index bed0920..0000000 --- a/telemetryutils/utils.go +++ /dev/null @@ -1,22 +0,0 @@ -package telemetryutils - -import ( - "context" - - "github.com/flyteorg/flytestdlib/contextutils" - - "go.opentelemetry.io/otel/attribute" - "go.opentelemetry.io/otel/trace" -) - -func NewSpan(ctx context.Context, serviceName string, spanName string) (context.Context, trace.Span) { - var attributes []attribute.KeyValue - for key, value := range contextutils.GetLogFields(ctx) { - if value, ok := value.(string); ok { - attributes = append(attributes, attribute.String(key, value)) - } - } - - tracerProvider := GetTracerProvider(serviceName) - return tracerProvider.Tracer("default").Start(ctx, spanName, trace.WithAttributes(attributes...)) -}