Add job context (#30)

* add no timescale option, add tests for manager

* add parent rid functionality, add tests, add context to queuer

* fix formatting

Author: siherrmann

core/runner.go

@@ -8,15 +8,21 @@ import (
 	"sync"
 	"time"
 
+	"github.com/google/uuid"
 	"github.com/siherrmann/queuer/helper"
 	"github.com/siherrmann/queuer/model"
 	vh "github.com/siherrmann/validator/helper"
 )
 
+type ContextKey string
+
+const JobRIDContextKey ContextKey = "JobRIDContextKey"
+
 type Runner struct {
 	cancel     context.CancelFunc
 	cancelMu   sync.RWMutex
 	Options    *model.Options
+	JobRID     *uuid.UUID
 	Task       interface{}
 	Parameters model.Parameters
 	// Result channel to return results
@@ -31,14 +37,22 @@ func NewRunner(options *model.Options, task interface{}, parameters ...interface
 	taskInputParameters, err := helper.GetInputParametersFromTask(task)
 	if err != nil {
 		return nil, helper.NewError("getting task input parameters", err)
-	} else if len(taskInputParameters) != len(parameters) {
-		return nil, fmt.Errorf("task expects %d parameters, got %d", len(taskInputParameters), len(parameters))
+	}
+
+	startIndex := 0
+	contextType := reflect.TypeOf((*context.Context)(nil)).Elem()
+	if len(taskInputParameters) > 0 && taskInputParameters[0] == contextType {
+		startIndex = 1
+	}
+
+	if len(taskInputParameters)-startIndex != len(parameters) {
+		return nil, fmt.Errorf("task expects %d parameters, got %d", len(taskInputParameters)-startIndex, len(parameters))
 	}
 
 	for i, param := range parameters {
-		paramConverted, err := vh.AnyToType(param, taskInputParameters[i])
+		paramConverted, err := vh.AnyToType(param, taskInputParameters[i+startIndex])
 		if err != nil {
-			return nil, fmt.Errorf("error converting parameter %d to type %s: %v", i, taskInputParameters[i].Kind(), err)
+			return nil, fmt.Errorf("error converting parameter %d to type %s: %v", i, taskInputParameters[i+startIndex].Kind(), err)
 		}
 		parameters[i] = paramConverted
 	}
@@ -73,6 +87,7 @@ func NewRunnerFromJob(task *model.Task, job *model.Job) (*Runner, error) {
 		return nil, fmt.Errorf("error creating runner from job: %v", err)
 	}
 
+	runner.JobRID = &job.RID
 	return runner, nil
 }
 
@@ -111,7 +126,21 @@ func (r *Runner) Run(ctx context.Context) {
 		}()
 
 		taskFunc := reflect.ValueOf(r.Task)
-		results := taskFunc.Call(r.Parameters.ToReflectValues())
+
+		var callParameters []reflect.Value
+		taskType := reflect.TypeOf(r.Task)
+		contextType := reflect.TypeOf((*context.Context)(nil)).Elem()
+
+		if taskType.NumIn() > 0 && taskType.In(0) == contextType {
+			jobCtx := ctx
+			if r.JobRID != nil {
+				jobCtx = context.WithValue(ctx, JobRIDContextKey, *r.JobRID)
+			}
+			callParameters = append(callParameters, reflect.ValueOf(jobCtx))
+		}
+
+		callParameters = append(callParameters, r.Parameters.ToReflectValues()...)
+		results := taskFunc.Call(callParameters)
 		resultValues := []interface{}{}
 		for _, result := range results {
 			resultValues = append(resultValues, result.Interface())
@@ -125,7 +154,8 @@ func (r *Runner) Run(ctx context.Context) {
 
 		var ok bool
 		if len(resultValues) > 0 {
-			if err, ok = resultValues[len(resultValues)-1].(error); ok || (outputParameters[1].String() == "error" && resultValues[len(resultValues)-1] == nil) {
+			lastOutIdx := len(outputParameters) - 1
+			if err, ok = resultValues[len(resultValues)-1].(error); ok || (lastOutIdx >= 0 && outputParameters[lastOutIdx].String() == "error" && resultValues[len(resultValues)-1] == nil) {
 				resultValues = resultValues[:len(resultValues)-1]
 			}
 		}

core/runner_test.go

@@ -692,3 +692,43 @@ func TestNewRunnerWithNestedStruct(t *testing.T) {
 		})
 	}
 }
+
+func TestRunnerContextInjection(t *testing.T) {
+	rid := uuid.New()
+	taskFn := func(ctx context.Context, p string) string {
+		ctxRid := ctx.Value(JobRIDContextKey)
+		if ctxRid != nil {
+			if id, ok := ctxRid.(uuid.UUID); ok {
+				return id.String() + "-" + p
+			}
+		}
+		return "no-ctx-" + p
+	}
+
+	task, err := model.NewTask(taskFn)
+	require.NoError(t, err)
+
+	job := &model.Job{
+		RID:        rid,
+		TaskName:   task.Name,
+		Parameters: []interface{}{"test"},
+	}
+
+	runner, err := NewRunnerFromJob(task, job)
+	require.NoError(t, err)
+
+	// Run the runner with a background context
+	go runner.Run(context.Background())
+
+	select {
+	case results := <-runner.ResultsChannel:
+		require.Len(t, results, 1)
+		resStr, ok := results[0].(string)
+		require.True(t, ok)
+		assert.Equal(t, rid.String()+"-test", resStr, "Context should contain the JobRID injected by Runner")
+	case err := <-runner.ErrorChannel:
+		t.Fatalf("Runner returned error: %v", err)
+	case <-time.After(1 * time.Second):
+		t.Fatal("Runner timed out")
+	}
+}

database/dbJob_test.go

@@ -819,3 +819,48 @@ func TestJobSelectAllJobsFromArchiveBySearch(t *testing.T) {
 	assert.NoError(t, err, "Expected SelectAllJobsFromArchiveBySearch to not return an error")
 	assert.Len(t, paginatedJobsBySearchFromArchive, pageLength, "Expected SelectAllJobsFromArchiveBySearch to return 3 archived jobs")
 }
+
+func TestAddRetentionArchive(t *testing.T) {
+	helper.SetTestDatabaseConfigEnvs(t, dbPort)
+	dbConfig, err := helper.NewDatabaseConfiguration()
+	if err != nil {
+		t.Fatalf("failed to create database configuration: %v", err)
+	}
+	database := helper.NewTestDatabase(dbConfig)
+	jobDbHandler, err := NewJobDBHandler(database, dbConfig)
+	require.NoError(t, err)
+
+	err = jobDbHandler.AddRetentionArchive(24 * time.Hour)
+	_ = err
+}
+
+func TestRemoveRetentionArchive(t *testing.T) {
+	helper.SetTestDatabaseConfigEnvs(t, dbPort)
+	dbConfig, err := helper.NewDatabaseConfiguration()
+	if err != nil {
+		t.Fatalf("failed to create database configuration: %v", err)
+	}
+	database := helper.NewTestDatabase(dbConfig)
+	jobDbHandler, err := NewJobDBHandler(database, dbConfig)
+	require.NoError(t, err)
+
+	err = jobDbHandler.RemoveRetentionArchive()
+	_ = err
+}
+
+func TestBatchInsertJobs(t *testing.T) {
+	helper.SetTestDatabaseConfigEnvs(t, dbPort)
+	dbConfig, err := helper.NewDatabaseConfiguration()
+	if err != nil {
+		t.Fatalf("failed to create database configuration: %v", err)
+	}
+	database := helper.NewTestDatabase(dbConfig)
+	jobDbHandler, err := NewJobDBHandler(database, dbConfig)
+	require.NoError(t, err)
+
+	job1, _ := model.NewJob("Task1", nil, nil)
+	job2, _ := model.NewJob("Task2", nil, nil)
+
+	err = jobDbHandler.BatchInsertJobs([]*model.Job{job1, job2})
+	assert.NoError(t, err)
+}

helper/database.go

@@ -354,6 +354,9 @@ func (d *Database) DropFunctionsFromPublicSchema(functionNames []string) error {
 			}
 			signatures = append(signatures, signature)
 		}
+		if err := rows.Err(); err != nil {
+			return NewError("rows iteration", err)
+		}
 
 		// Drop each overloaded function by its full signature
 		for _, signature := range signatures {

helper/database_test.go

@@ -2,6 +2,7 @@ package helper
 
 import (
 	"context"
+	"database/sql"
 	"log"
 	"testing"
 
@@ -302,3 +303,37 @@ func TestDropIndex(t *testing.T) {
 	err = database.DropIndex("test_drop_index", "name")
 	assert.NoError(t, err, "expected no error when dropping non-existing index")
 }
+
+func TestNewDatabaseWithDB(t *testing.T) {
+	db := NewDatabaseWithDB("test_with_db", nil, nil)
+	assert.NotNil(t, db)
+	assert.Equal(t, "test_with_db", db.Name)
+	assert.Nil(t, db.Instance)
+}
+
+func TestDatabaseClose(t *testing.T) {
+	dbConn, _ := sql.Open("postgres", "user=pqgotest dbname=pqgotest sslmode=verify-full")
+	db := NewDatabaseWithDB("test", dbConn, nil)
+	err := db.Close()
+	assert.NoError(t, err)
+}
+
+func TestDropFunctionsFromPublicSchema(t *testing.T) {
+	dbConn, _ := sql.Open("postgres", "user=pqgotest dbname=pqgotest sslmode=verify-full")
+	db := NewDatabaseWithDB("test", dbConn, nil)
+	// it will probably error with connection refused, but it tests the nil guard if any, or it tests the query logic
+	err := db.DropFunctionsFromPublicSchema([]string{"non_existent_func"})
+	assert.Error(t, err) // since it's a dummy connection string, it should error
+}
+
+func TestMustStartPostgresContainer(t *testing.T) {
+	teardown, port, err := MustStartPostgresContainer()
+	assert.NoError(t, err)
+	assert.NotEmpty(t, port)
+	assert.NotNil(t, teardown)
+
+	if teardown != nil {
+		err := teardown(context.Background())
+		assert.NoError(t, err)
+	}
+}

helper/prettyLog_test.go

@@ -0,0 +1,58 @@
+package helper
+
+import (
+	"bytes"
+	"context"
+	"log/slog"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestPrettyHandler(t *testing.T) {
+	var buf bytes.Buffer
+	handler := NewPrettyHandler(&buf, PrettyHandlerOptions{})
+
+	logger := slog.New(handler)
+	logger.Info("test message", "key", "value")
+
+	output := buf.String()
+	assert.Contains(t, output, "test message")
+	assert.Contains(t, output, "key")
+	assert.Contains(t, output, "value")
+	assert.Contains(t, output, "INFO")
+}
+
+func TestPrettyHandlerLevels(t *testing.T) {
+	var buf bytes.Buffer
+	handler := NewPrettyHandler(&buf, PrettyHandlerOptions{
+		SlogOpts: slog.HandlerOptions{
+			Level: slog.LevelDebug,
+		},
+	})
+
+	logger := slog.New(handler)
+	logger.Debug("debug msg")
+	logger.Warn("warn msg")
+	logger.Error("error msg")
+
+	output := buf.String()
+	assert.Contains(t, output, "debug msg")
+	assert.Contains(t, output, "warn msg")
+	assert.Contains(t, output, "error msg")
+	assert.Contains(t, output, "DEBUG")
+	assert.Contains(t, output, "WARN")
+	assert.Contains(t, output, "ERROR")
+}
+
+func TestPrettyHandlerHandleDirect(t *testing.T) {
+	var buf bytes.Buffer
+	handler := NewPrettyHandler(&buf, PrettyHandlerOptions{})
+
+	r := slog.NewRecord(time.Now(), slog.LevelInfo, "direct handle", 0)
+
+	err := handler.Handle(context.Background(), r)
+	assert.NoError(t, err)
+	assert.Contains(t, buf.String(), "direct handle")
+}

helper/task.go

@@ -1,6 +1,7 @@
 package helper
 
 import (
+	"context"
 	"fmt"
 	"reflect"
 	"runtime"
@@ -23,21 +24,29 @@ func CheckValidTask(task interface{}) error {
 }
 
 // CheckValidTaskWithParameters checks if the provided task and parameters are valid.
-// It checks if the task is a valid function and if the parameters match the task's input types.
+// It checks if the task is a valid function and if the parameters match the task's input types, ignoring an optional context.Context as first argument.
 func CheckValidTaskWithParameters(task interface{}, parameters ...interface{}) error {
 	err := CheckValidTask(task)
 	if err != nil {
 		return err
 	}
 
 	taskType := reflect.TypeOf(task)
-	if taskType.NumIn() != len(parameters) {
-		return fmt.Errorf("task expects %d parameters, got %d", taskType.NumIn(), len(parameters))
+	expectedParams := taskType.NumIn()
+	startIndex := 0
+
+	contextType := reflect.TypeOf((*context.Context)(nil)).Elem()
+	if expectedParams > 0 && taskType.In(0) == contextType {
+		startIndex = 1
+	}
+
+	if expectedParams-startIndex != len(parameters) {
+		return fmt.Errorf("task expects %d parameters, got %d", expectedParams-startIndex, len(parameters))
 	}
 
 	for i, param := range parameters {
-		if !reflect.TypeOf(param).AssignableTo(taskType.In(i)) {
-			return fmt.Errorf("parameter %d of task must be of type %s, got %s", i, taskType.In(i).Kind(), reflect.TypeOf(param).Kind())
+		if !reflect.TypeOf(param).AssignableTo(taskType.In(i + startIndex)) {
+			return fmt.Errorf("parameter %d of task must be of type %s, got %s", i, taskType.In(i+startIndex).Kind(), reflect.TypeOf(param).Kind())
 		}
 	}
 

helper/task_test.go

@@ -123,8 +123,8 @@ func TestCheckValidTaskWithParameters(t *testing.T) {
 		err := CheckValidTaskWithParameters(testFuncOneParam, "hello")
 		assert.NoError(t, err, "should not return error for matching params")
 
-		err = CheckValidTaskWithParameters(testFuncMultiParamsReturn, context.Background(), 10, 3.14)
-		assert.NoError(t, err, "should not return error for multiple matching params")
+		err = CheckValidTaskWithParameters(testFuncMultiParamsReturn, 10, 3.14)
+		assert.NoError(t, err, "should not return error for multiple matching params, ignoring context")
 	})
 
 	t.Run("Not a function", func(t *testing.T) {
@@ -152,15 +152,15 @@ func TestCheckValidTaskWithParameters(t *testing.T) {
 	})
 
 	t.Run("Function with wrong parameter type (multiple params - first wrong)", func(t *testing.T) {
-		err := CheckValidTaskWithParameters(testFuncMultiParamsReturn, "wrong", 10, 3.14)
+		err := CheckValidTaskWithParameters(testFuncMultiParamsReturn, "wrong", 3.14)
 		require.Error(t, err)
-		assert.Contains(t, err.Error(), "parameter 0 of task must be of type interface, got string", "should fail on first wrong type")
+		assert.Contains(t, err.Error(), "parameter 0 of task must be of type int, got string", "should fail on first wrong type")
 	})
 
 	t.Run("Function with wrong parameter type (multiple params - second wrong)", func(t *testing.T) {
-		err := CheckValidTaskWithParameters(testFuncMultiParamsReturn, context.Background(), "wrong", 3.14)
+		err := CheckValidTaskWithParameters(testFuncMultiParamsReturn, 10, "wrong")
 		require.Error(t, err)
-		assert.Contains(t, err.Error(), "parameter 1 of task must be of type int, got string", "should fail on second wrong type")
+		assert.Contains(t, err.Error(), "parameter 1 of task must be of type float64, got string", "should fail on second wrong type")
 	})
 
 	t.Run("Function with different parameter type (time.Duration vs int)", func(t *testing.T) {

model/options.go

@@ -5,12 +5,15 @@ import (
 	"encoding/json"
 	"errors"
 
+	"github.com/google/uuid"
 	"github.com/siherrmann/queuer/helper"
 )
 
 type Options struct {
-	OnError  *OnError  `json:"on_error,omitempty"`
-	Schedule *Schedule `json:"schedule,omitempty"`
+	OnError   *OnError               `json:"on_error,omitempty"`
+	Schedule  *Schedule              `json:"schedule,omitempty"`
+	Metadata  map[string]interface{} `json:"metadata,omitempty"`
+	ParentRID *uuid.UUID             `json:"parent_rid,omitempty"`
 }
 
 func (c *Options) IsValid() error {