diff --git a/oncetask/context_keys.go b/oncetask/context_keys.go index c0ef42b..9b644b5 100644 --- a/oncetask/context_keys.go +++ b/oncetask/context_keys.go @@ -10,16 +10,18 @@ type contextKey string const ( taskIDContextKey contextKey = "oncetask.taskID" resourceKeyContextKey contextKey = "oncetask.resourceKey" + taskTypeContextKey contextKey = "oncetask.taskType" ) -// withTaskContext adds both task ID and resource key to the context for automatic logging -func withTaskContext(ctx context.Context, taskID, resourceKey string) context.Context { +// withTaskContext adds task ID, resource key, and task type to the context for automatic logging +func withTaskContext(ctx context.Context, taskID, resourceKey, taskType string) context.Context { if taskID != "" { ctx = context.WithValue(ctx, taskIDContextKey, taskID) } if resourceKey != "" { ctx = context.WithValue(ctx, resourceKeyContextKey, resourceKey) } + ctx = context.WithValue(ctx, taskTypeContextKey, taskType) return ctx } @@ -27,9 +29,7 @@ func withSingleTaskContext[TaskKind ~string](ctx context.Context, tasks []OnceTa if len(tasks) == 0 { return ctx } - taskID := tasks[0].Id - resourceKey := tasks[0].ResourceKey - return withTaskContext(ctx, taskID, resourceKey) + return withTaskContext(ctx, tasks[0].Id, tasks[0].ResourceKey, string(tasks[0].Type)) } // withResourceKeyTaskContext is used for resource key batched tasks and adds only the resource key to the context for automatic logging @@ -43,7 +43,7 @@ func withResourceKeyTaskContext[TaskKind ~string](ctx context.Context, tasks []O taskID = tasks[0].Id } - return withTaskContext(ctx, taskID, tasks[0].ResourceKey) + return withTaskContext(ctx, taskID, tasks[0].ResourceKey, string(tasks[0].Type)) } // GetCurrentTaskID returns the task ID stored in the context, or an empty string if not present. diff --git a/oncetask/context_keys_test.go b/oncetask/context_keys_test.go index 2e55ffe..f6a0020 100644 --- a/oncetask/context_keys_test.go +++ b/oncetask/context_keys_test.go @@ -14,14 +14,14 @@ func TestTaskIDFromContext(t *testing.T) { }) t.Run("returns task ID when present in context", func(t *testing.T) { - ctx := withTaskContext(context.Background(), "test-task-123", "") + ctx := withTaskContext(context.Background(), "test-task-123", "", "email") if got := GetCurrentTaskID(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") + ctx := withTaskContext(context.Background(), "task-456", "resource-789", "email") if got := GetCurrentTaskID(ctx); got != "task-456" { t.Errorf("Expected 'task-456', got: %q", got) } @@ -37,14 +37,14 @@ func TestResourceKeyFromContext(t *testing.T) { }) t.Run("returns resource key when present in context", func(t *testing.T) { - ctx := withTaskContext(context.Background(), "", "user-123") + ctx := withTaskContext(context.Background(), "", "user-123", "email") if got := GetCurrentTaskResourceKey(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") + ctx := withTaskContext(context.Background(), "task-456", "resource-789", "email") if got := GetCurrentTaskResourceKey(ctx); got != "resource-789" { t.Errorf("Expected 'resource-789', got: %q", got) } diff --git a/oncetask/context_logger.go b/oncetask/context_logger.go index 79d4588..dbc33d1 100644 --- a/oncetask/context_logger.go +++ b/oncetask/context_logger.go @@ -40,6 +40,10 @@ func (h *ContextHandler) Handle(ctx context.Context, r slog.Record) error { if resourceKey, ok := ctx.Value(resourceKeyContextKey).(string); ok { r.AddAttrs(slog.String("resourceKey", resourceKey)) } + // Extract task type from context and add it to the record + if taskType, ok := ctx.Value(taskTypeContextKey).(string); ok { + r.AddAttrs(slog.String("taskType", taskType)) + } return h.handler.Handle(ctx, r) } diff --git a/oncetask/context_logger_test.go b/oncetask/context_logger_test.go index ebc1551..e369e09 100644 --- a/oncetask/context_logger_test.go +++ b/oncetask/context_logger_test.go @@ -54,7 +54,7 @@ func TestContextHandler_WithTaskIDOnly(t *testing.T) { logger := slog.New(contextHandler) // Create a context with only task ID - ctx := withTaskContext(context.Background(), "test-task-123", "") + ctx := withTaskContext(context.Background(), "test-task-123", "", "email") // Log using the context logger.InfoContext(ctx, "Processing task", "operation", "sync") @@ -99,7 +99,7 @@ func TestContextHandler_WithResourceKeyOnly(t *testing.T) { logger := slog.New(contextHandler) // Create a context with only resource key (batch processing scenario) - ctx := withTaskContext(context.Background(), "", "user-123") + ctx := withTaskContext(context.Background(), "", "user-123", "email") // Log using the context logger.InfoContext(ctx, "Processing resource batch", "operation", "update") @@ -144,7 +144,7 @@ func TestContextHandler_WithBothTaskIDAndResourceKey(t *testing.T) { logger := slog.New(contextHandler) // Create a context with both task ID and resource key - ctx := withTaskContext(context.Background(), "task-999", "user-456") + ctx := withTaskContext(context.Background(), "task-999", "user-456", "email") // Log using the context logger.InfoContext(ctx, "Processing task with resource key") @@ -165,6 +165,27 @@ func TestContextHandler_WithBothTaskIDAndResourceKey(t *testing.T) { } } +func TestContextHandler_InjectsTaskType(t *testing.T) { + var buf bytes.Buffer + jsonHandler := slog.NewJSONHandler(&buf, &slog.HandlerOptions{ + Level: slog.LevelInfo, + }) + contextHandler := NewContextHandler(jsonHandler) + logger := slog.New(contextHandler) + + ctx := withTaskContext(context.Background(), "task-1", "user-1", "send_email") + logger.InfoContext(ctx, "Processing") + + var logEntry map[string]interface{} + if err := json.Unmarshal(buf.Bytes(), &logEntry); err != nil { + t.Fatalf("Failed to parse log output: %v", err) + } + + if taskType, ok := logEntry["taskType"].(string); !ok || taskType != "send_email" { + t.Errorf("Expected taskType to be 'send_email', got: %v", logEntry["taskType"]) + } +} + func TestContextHandler_WithAttrs(t *testing.T) { var buf bytes.Buffer jsonHandler := slog.NewJSONHandler(&buf, &slog.HandlerOptions{ @@ -175,7 +196,7 @@ func TestContextHandler_WithAttrs(t *testing.T) { // Create a logger with additional attributes logger := slog.New(contextHandler).With("service", "oncetask") - ctx := withTaskContext(context.Background(), "task-789", "") + ctx := withTaskContext(context.Background(), "task-789", "", "email") logger.InfoContext(ctx, "Test message") var logEntry map[string]interface{} diff --git a/oncetask/once_task_firestore.go b/oncetask/once_task_firestore.go index 5851b16..97b15d7 100644 --- a/oncetask/once_task_firestore.go +++ b/oncetask/once_task_firestore.go @@ -286,7 +286,7 @@ func (m *firestoreOnceTaskManager[TaskKind]) runLoop( if len(cancelledTasks) > 0 { cancellationHandler := getCancellationHandler[TaskKind](config) for _, task := range cancelledTasks { - ctx := withTaskContext(handlerCtx, task.Id, task.ResourceKey) + ctx := withTaskContext(handlerCtx, task.Id, task.ResourceKey, string(task.Type)) result, execErr := SafeExecute(ctx, cancellationHandler, &task) if err := m.completeBatch(m.ctx, []OnceTask[TaskKind]{task}, execErr, result, config); err != nil { slog.ErrorContext(m.ctx, "Failed to complete cancelled task", "error", err, "taskId", task.Id)