diff --git a/oncetask/context_logger.go b/oncetask/context_logger.go index 6b9d2b4..f005800 100644 --- a/oncetask/context_logger.go +++ b/oncetask/context_logger.go @@ -47,6 +47,24 @@ func withResourceKeyTaskContext[TaskKind ~string](ctx context.Context, tasks []O return withTaskContext(ctx, taskID, tasks[0].ResourceKey) } +// TaskIDFromContext returns the task ID stored in the context, or an empty string if not present. +// This is useful for debugging or when you need to access the current task ID within a handler. +func TaskIDFromContext(ctx context.Context) string { + if taskID, ok := ctx.Value(taskIDContextKey).(string); ok { + return taskID + } + return "" +} + +// ResourceKeyFromContext returns the resource key stored in the context, or an empty string if not present. +// This is useful for debugging or when you need to access the current resource key within a handler. +func ResourceKeyFromContext(ctx context.Context) string { + if resourceKey, ok := ctx.Value(resourceKeyContextKey).(string); ok { + return resourceKey + } + return "" +} + // ContextHandler is a slog.Handler that automatically extracts the task ID and resource key from context // and adds them as attributes to all log records. // diff --git a/oncetask/context_logger_test.go b/oncetask/context_logger_test.go index ebc1551..72e9ab0 100644 --- a/oncetask/context_logger_test.go +++ b/oncetask/context_logger_test.go @@ -192,3 +192,49 @@ func TestContextHandler_WithAttrs(t *testing.T) { t.Errorf("Expected service to be 'oncetask', got: %v", logEntry["service"]) } } + +func TestTaskIDFromContext(t *testing.T) { + t.Run("returns empty string when no task ID in context", func(t *testing.T) { + ctx := context.Background() + if got := TaskIDFromContext(ctx); got != "" { + t.Errorf("Expected empty string, got: %q", got) + } + }) + + t.Run("returns task ID when present in context", func(t *testing.T) { + ctx := withTaskContext(context.Background(), "test-task-123", "") + if got := TaskIDFromContext(ctx); got != "test-task-123" { + t.Errorf("Expected 'test-task-123', got: %q", got) + } + }) + + t.Run("returns task ID when both task ID and resource key are present", func(t *testing.T) { + ctx := withTaskContext(context.Background(), "task-456", "resource-789") + if got := TaskIDFromContext(ctx); got != "task-456" { + t.Errorf("Expected 'task-456', got: %q", got) + } + }) +} + +func TestResourceKeyFromContext(t *testing.T) { + t.Run("returns empty string when no resource key in context", func(t *testing.T) { + ctx := context.Background() + if got := ResourceKeyFromContext(ctx); got != "" { + t.Errorf("Expected empty string, got: %q", got) + } + }) + + t.Run("returns resource key when present in context", func(t *testing.T) { + ctx := withTaskContext(context.Background(), "", "user-123") + if got := ResourceKeyFromContext(ctx); got != "user-123" { + t.Errorf("Expected 'user-123', got: %q", got) + } + }) + + t.Run("returns resource key when both task ID and resource key are present", func(t *testing.T) { + ctx := withTaskContext(context.Background(), "task-456", "resource-789") + if got := ResourceKeyFromContext(ctx); got != "resource-789" { + t.Errorf("Expected 'resource-789', got: %q", got) + } + }) +}