feat: Add public jobs.WithJobContext()

pconstantinou · Jun 26, 2024 · f5e7ee8 · f5e7ee8
1 parent e3d6237
commit f5e7ee8
Show file tree

Hide file tree

Showing 4 changed files with 8 additions and 18 deletions.
diff --git a/backends/memory/memory_backend.go b/backends/memory/memory_backend.go
@@ -331,7 +331,7 @@ func (m *MemBackend) scheduleFutureJobs(ctx context.Context) {
 }
 
 func (m *MemBackend) handleJob(ctx context.Context, job *jobs.Job, h handler.Handler) (err error) {
-	ctx = withJobContext(ctx, job)
+	ctx = jobs.WithJobContext(ctx, job)
 
 	m.logger.Debug(
 		"handling job",
@@ -391,8 +391,3 @@ func (m *MemBackend) removeFutureJob(jobID int64) {
 		m.futureJobs.Delete(job.ID)
 	}
 }
-
-// withJobContext creates a new context with the Job set
-func withJobContext(ctx context.Context, j *jobs.Job) context.Context {
-	return context.WithValue(ctx, internal.JobCtxVarKey, j)
-}
diff --git a/backends/postgres/postgres_backend.go b/backends/postgres/postgres_backend.go
@@ -939,7 +939,7 @@ func (p *PgBackend) handleJob(ctx context.Context, jobID string) (err error) {
 		return
 	}
 
-	ctx = withJobContext(ctx, job)
+	ctx = jobs.WithJobContext(ctx, job)
 	ctx = context.WithValue(ctx, txCtxVarKey, tx)
 
 	if job.RunAfter.After(time.Now()) {
@@ -1131,11 +1131,6 @@ func (p *PgBackend) acquire(ctx context.Context) (conn *pgxpool.Conn, err error)
 	}
 }
 
-// withJobContext creates a new context with the Job set
-func withJobContext(ctx context.Context, j *jobs.Job) context.Context {
-	return context.WithValue(ctx, internal.JobCtxVarKey, j)
-}
-
 func GetPQConnectionString(connectionString string) (string, error) {
 	pgxCfg, err := pgx.ParseConfig(connectionString)
 	if err != nil {

diff --git a/backends/redis/redis_backend.go b/backends/redis/redis_backend.go
@@ -253,7 +253,7 @@ func (b *RedisBackend) Start(_ context.Context, h handler.Handler) (err error) {
 			RunAfter:   ti.NextProcessAt,
 		}
 
-		ctx = withJobContext(ctx, job)
+		ctx = jobs.WithJobContext(ctx, job)
 		err = handler.Exec(ctx, h)
 		if err != nil {
 			b.logger.Error("error handling job", slog.Any("error", err))
@@ -360,8 +360,3 @@ func (b *RedisBackend) Shutdown(_ context.Context) {
 	b.client.Close()
 	b.server.Shutdown()
 }
-
-// withJobContext creates a new context with the Job set
-func withJobContext(ctx context.Context, j *jobs.Job) context.Context {
-	return context.WithValue(ctx, internal.JobCtxVarKey, j)
-}
diff --git a/jobs/jobs.go b/jobs/jobs.go
@@ -77,6 +77,11 @@ func FingerprintJob(j *Job) (err error) {
 	return
 }
 
+// WithJobContext adds a job to the provided context
+func WithJobContext(ctx context.Context, job *Job) context.Context {
+	return context.WithValue(ctx, internal.JobCtxVarKey, job)
+}
+
 // FromContext fetches the job from a context if the job context variable is set
 func FromContext(ctx context.Context) (j *Job, err error) {
 	var ok bool