databricks-cli/libs/telemetry/context.go

Ignoring revisions in .git-blame-ignore-revs. Click here to bypass and see the normal blame view.

64 lines
1.7 KiB
Go
Raw Normal View History

2024-11-29 00:57:48 +00:00
package telemetry
import (
"context"
"fmt"
2024-11-29 00:57:48 +00:00
)
// Private type to store the telemetry logger in the context
type telemetryLogger int
// Key to store the telemetry logger in the context
var telemetryLoggerKey telemetryLogger
2025-01-03 06:17:45 +00:00
// TODO: Add tests for these methods.
func WithDefaultLogger(ctx context.Context) context.Context {
v := ctx.Value(telemetryLoggerKey)
2025-01-03 06:17:45 +00:00
// If no logger is set in the context, set the default logger.
if v == nil {
nctx := context.WithValue(ctx, telemetryLoggerKey, &defaultLogger{logs: []FrontendLog{}})
return nctx
}
switch v.(type) {
case *defaultLogger:
panic(fmt.Sprintf("default telemetry logger already set in the context: %v", v))
case *mockLogger:
// Do nothing. Unit and integration tests set the mock logger in the context
// to avoid making actual API calls. Thus WithDefaultLogger should silently
// ignore the mock logger.
default:
panic(fmt.Sprintf("unexpected telemetry logger type: %T", v))
2024-11-29 00:57:48 +00:00
}
2025-01-03 06:17:45 +00:00
return ctx
2024-11-29 00:57:48 +00:00
}
2025-01-03 06:17:45 +00:00
// WithMockLogger sets a mock telemetry logger in the context. It overrides the
// default logger if it is already set in the context.
func WithMockLogger(ctx context.Context) context.Context {
v := ctx.Value(telemetryLoggerKey)
if v != nil {
panic(fmt.Sprintf("telemetry logger already set in the context: %v", v))
}
return context.WithValue(ctx, telemetryLoggerKey, &mockLogger{})
}
func fromContext(ctx context.Context) Logger {
v := ctx.Value(telemetryLoggerKey)
if v == nil {
2024-11-29 00:57:48 +00:00
panic("telemetry logger not found in the context")
}
2025-01-03 06:17:45 +00:00
switch vv := v.(type) {
case *defaultLogger:
2025-01-03 06:17:45 +00:00
return vv
case *mockLogger:
2025-01-03 06:17:45 +00:00
return vv
default:
panic(fmt.Sprintf("unexpected telemetry logger type: %T", v))
}
2024-11-29 00:57:48 +00:00
}