diff --git a/README.md b/README.md index f4a7206..2503d3a 100644 --- a/README.md +++ b/README.md @@ -80,8 +80,13 @@ func main() { //you can move the job to permanent storage if you can't handle it return bgjob.MoveToDlq(errors.New("move to dlq")) }) - - cli := bgjob.NewClient(bgjob.NewPgStore(db)) + + ctx := context.Background() + store, err := bgjob.NewPgStoreV2(ctx, db) + if err != nil { + panic(err) + } + cli := bgjob.NewClient(store) //if handler for job type wasn't provided, job will be moved to dlq handler := bgjob.NewMux(). @@ -98,7 +103,6 @@ func main() { bgjob.WithConcurrency(runtime.NumCPU()), //default 1 bgjob.WithPollInterval(500*time.Millisecond), //default 1s ) - ctx := context.Background() worker.Run(ctx) //call ones, non-blocking err = cli.Enqueue(ctx, bgjob.EnqueueRequest{ diff --git a/client_test.go b/client_test.go index c8b359b..06d3d95 100644 --- a/client_test.go +++ b/client_test.go @@ -367,7 +367,9 @@ func prepareTest(t *testing.T) (*require.Assertions, *db, *bgjob.Client) { err = applyMigration(db.DB) asserter.NoError(err) - cli := bgjob.NewClient(bgjob.NewPgStore(db.DB)) + store, err := bgjob.NewPgStoreV2(context.Background(), db.DB) + asserter.NoError(err) + cli := bgjob.NewClient(store) return asserter, db, cli } diff --git a/enqueue.go b/enqueue.go index 42b5773..32955f1 100644 --- a/enqueue.go +++ b/enqueue.go @@ -9,7 +9,7 @@ import ( ) type ExecerContext interface { - ExecContext(ctx context.Context, s string, args ...interface{}) (sql.Result, error) + ExecContext(ctx context.Context, s string, args ...any) (sql.Result, error) } func Enqueue(ctx context.Context, e ExecerContext, req EnqueueRequest) error { @@ -61,6 +61,7 @@ func requestsToJobs(list []EnqueueRequest) ([]Job, error) { Queue: req.Queue, Type: req.Type, Arg: req.Arg, + RequestId: req.RequestId, Attempt: 0, LastError: nil, NextRunAt: now.Add(req.Delay).Unix(), @@ -75,11 +76,11 @@ func requestsToJobs(list []EnqueueRequest) ([]Job, error) { func bulkInsert(ctx context.Context, e ExecerContext, jobs []Job) error { valueStrings := make([]string, 0, len(jobs)) - valueArgs := make([]interface{}, 0, len(jobs)*8) + valueArgs := make([]any, 0, len(jobs)*9) placeholderNum := 0 for _, job := range jobs { placeholders := make([]string, 0) - for i := 0; i < 8; i++ { + for i := 0; i < 9; i++ { placeholderNum++ placeholders = append(placeholders, fmt.Sprintf("$%d", placeholderNum)) } @@ -94,9 +95,10 @@ func bulkInsert(ctx context.Context, e ExecerContext, jobs []Job) error { job.NextRunAt, job.CreatedAt, job.UpdatedAt, + job.RequestId, ) } - query := fmt.Sprintf("INSERT INTO bgjob_job (id, queue, type, arg, attempt, next_run_at, created_at, updated_at) VALUES %s", + query := fmt.Sprintf("INSERT INTO bgjob_job (id, queue, type, arg, attempt, next_run_at, created_at, updated_at, request_id) VALUES %s", strings.Join(valueStrings, ",")) _, err := e.ExecContext(ctx, query, valueArgs...) if err != nil { diff --git a/enqueue_request.go b/enqueue_request.go index 4e95d49..6bca4bc 100644 --- a/enqueue_request.go +++ b/enqueue_request.go @@ -5,9 +5,10 @@ import ( ) type EnqueueRequest struct { - Id string //optional - Queue string //required - Type string //required - Arg []byte //optional - Delay time.Duration //optional + Id string //optional + Queue string //required + Type string //required + Arg []byte //optional + Delay time.Duration //optional + RequestId string //optional } diff --git a/example/main.go b/example/main.go index 25b989c..c043d45 100644 --- a/example/main.go +++ b/example/main.go @@ -45,7 +45,12 @@ func main() { return bgjob.MoveToDlq(errors.New("move to dlq")) }) - cli := bgjob.NewClient(bgjob.NewPgStore(db)) + ctx := context.Background() + store, err := bgjob.NewPgStoreV2(ctx, db) + if err != nil { + panic(err) + } + cli := bgjob.NewClient(store) //if handler for job type wasn't provided, job will be moved to dlq handler := bgjob.NewMux(). @@ -62,7 +67,6 @@ func main() { bgjob.WithConcurrency(runtime.NumCPU()), //default 1 bgjob.WithPollInterval(500*time.Millisecond), //default 1s ) - ctx := context.Background() worker.Run(ctx) //call ones, non-blocking err = cli.Enqueue(ctx, bgjob.EnqueueRequest{ diff --git a/job.go b/job.go index 385d26c..09dbf15 100644 --- a/job.go +++ b/job.go @@ -14,4 +14,5 @@ type Job struct { NextRunAt int64 CreatedAt time.Time UpdatedAt time.Time + RequestId string } diff --git a/migration/add_request_id.sql b/migration/add_request_id.sql new file mode 100644 index 0000000..7ee47cc --- /dev/null +++ b/migration/add_request_id.sql @@ -0,0 +1,21 @@ +alter table bgjob_job + add column request_id text; + +alter table bgjob_dead_job + add column request_id text; + +update bgjob_job + set request_id = '' + where request_id is null; + +update bgjob_dead_job + set request_id = '' + where request_id is null; + +alter table bgjob_job + alter column request_id set default '', + alter column request_id set not null; + +alter table bgjob_dead_job + alter column request_id set default '', + alter column request_id set not null; diff --git a/migration/init.sql b/migration/init.sql index 3d20013..d4b6197 100644 --- a/migration/init.sql +++ b/migration/init.sql @@ -8,7 +8,8 @@ create table bgjob_job last_error text, next_run_at int8 not null, created_at timestamp not null, - updated_at timestamp not null + updated_at timestamp not null, + request_id text not null default '' ); create index ix_bgjob_job__queue_next_run_at_created_at on bgjob_job (queue, next_run_at, created_at); @@ -25,5 +26,6 @@ create table bgjob_dead_job next_run_at int8 not null, job_created_at timestamp not null, job_updated_at timestamp not null, - moved_at timestamp not null + moved_at timestamp not null, + request_id text not null default '' ) diff --git a/pg_store.go b/pg_store.go index ac6c284..3c8a382 100644 --- a/pg_store.go +++ b/pg_store.go @@ -10,10 +10,21 @@ type pgStore struct { db *sql.DB } -func NewPgStore(db *sql.DB) *pgStore { - return &pgStore{ - db: db, +func NewPgStoreV2(ctx context.Context, db *sql.DB) (*pgStore, error) { + err := assertHasColumn(ctx, db, "bgjob_job", "request_id") + if err != nil { + return nil, fmt.Errorf("%w; please apply 'add_request_id.sql' migration adding column `request_id` to `bgjob_job`", err) + } + + err = assertHasColumn(ctx, db, "bgjob_dead_job", "request_id") + if err != nil { + return nil, fmt.Errorf("%w; please apply 'add_request_id.sql' migration adding column `request_id` to `bgjob_dead_job`", err) } + + return &pgStore{ + db: db, + }, + nil } func (p *pgStore) BulkInsert(ctx context.Context, jobs []Job) error { @@ -23,7 +34,7 @@ func (p *pgStore) BulkInsert(ctx context.Context, jobs []Job) error { func (p *pgStore) Acquire(ctx context.Context, queue string, handler func(tx Tx) error) error { return runTx(ctx, p.db, func(ctx context.Context, tx *sql.Tx) error { query := ` -SELECT id, queue, type, arg, attempt, last_error, next_run_at, created_at, updated_at +SELECT id, queue, type, arg, attempt, last_error, next_run_at, created_at, updated_at, request_id FROM bgjob_job WHERE queue = $1 AND next_run_at <= $2 ORDER BY next_run_at, created_at @@ -41,6 +52,7 @@ LIMIT 1 FOR UPDATE SKIP LOCKED &job.NextRunAt, &job.CreatedAt, &job.UpdatedAt, + &job.RequestId, ) if err == sql.ErrNoRows { return ErrEmptyQueue @@ -106,8 +118,8 @@ func (p *pgTx) Delete(ctx context.Context, id string) error { func (p *pgTx) SaveInDlq(ctx context.Context, job Job) error { query := `INSERT INTO bgjob_dead_job -(job_id, queue, type, arg, attempt, next_run_at, last_error, job_created_at, job_updated_at, moved_at) -VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10) +(job_id, queue, type, arg, attempt, next_run_at, last_error, job_created_at, job_updated_at, moved_at, request_id) +VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11) ` _, err := p.tx.ExecContext( ctx, @@ -122,6 +134,7 @@ VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10) job.CreatedAt, job.UpdatedAt, timeNow(), + job.RequestId, ) return err } @@ -131,3 +144,17 @@ func (p *pgTx) UpdateArg(ctx context.Context, id string, arg []byte) error { _, err := p.tx.ExecContext(ctx, query, arg, timeNow(), id) return err } + +func assertHasColumn(ctx context.Context, db *sql.DB, table, column string) error { + query := fmt.Sprintf("SELECT %s FROM %s LIMIT 1", column, table) + + err := db.QueryRowContext(ctx, query).Scan(new(any)) + if err != nil { + if err == sql.ErrNoRows { + return nil + } + return fmt.Errorf("bgjob: schema check error for %s.%s: %w", table, column, err) + } + + return nil +}