Add AssertSchedule functionality (#46)

Now you can use this single function to upsert a schedule.

Author: rdner

pkg/queue/mock_test.go

@@ -60,6 +60,9 @@ type SchedulerMock struct {
 
 	EnsureError error
 	EnsureCount int
+
+	AssertError error
+	AssertCount int
 }
 
 func (s *SchedulerMock) Schedule(ctx context.Context, builder cdb.SQLBuilder, task TaskScheduleRequest) (err error) {
@@ -71,3 +74,8 @@ func (s *SchedulerMock) EnsureSchedule(ctx context.Context, builder cdb.SQLBuild
 	s.EnsureCount = s.EnsureCount + 1
 	return s.EnsureError
 }
+
+func (s *SchedulerMock) AssertSchedule(ctx context.Context, task TaskScheduleRequest) error {
+	s.AssertCount = s.AssertCount + 1
+	return s.AssertError
+}

pkg/queue/postgres/scheduler.go

@@ -3,20 +3,23 @@ package postgres
 import (
 	"context"
 	"database/sql"
+	"fmt"
 	"time"
 
 	"github.com/Masterminds/squirrel"
+	"github.com/contiamo/go-base/pkg/db"
 	cdb "github.com/contiamo/go-base/pkg/db"
+	"github.com/contiamo/go-base/pkg/queue"
 	"github.com/contiamo/go-base/pkg/tracing"
 	cvalidation "github.com/contiamo/go-base/pkg/validation"
-	"github.com/contiamo/go-base/pkg/queue"
 	uuid "github.com/satori/go.uuid"
 )
 
 // NewScheduler creates a new postgres task scheduler
 func NewScheduler(db *sql.DB) queue.Scheduler {
 	return &scheduler{
 		Tracer: tracing.NewTracer("queue", "PostgresScheduler"),
+		db:     db,
 	}
 }
 
@@ -28,6 +31,7 @@ func NewSchedulerWithMetrics(db *sql.DB) queue.Scheduler {
 // scheduler is a postgres backed implementation of the task scheduler
 type scheduler struct {
 	tracing.Tracer
+	db *sql.DB
 }
 
 func (q *scheduler) Schedule(ctx context.Context, builder cdb.SQLBuilder, task queue.TaskScheduleRequest) (err error) {
@@ -142,3 +146,51 @@ func (q *scheduler) EnsureSchedule(ctx context.Context, builder cdb.SQLBuilder,
 	// either value was 0 or err == sql.ErrorNoRows
 	return queue.ErrNotScheduled
 }
+
+func (q *scheduler) AssertSchedule(ctx context.Context, schedule queue.TaskScheduleRequest) (err error) {
+	span, ctx := q.StartSpan(ctx, "AssertSchedule")
+	defer func() {
+		q.FinishSpan(span, err)
+	}()
+
+	span.SetTag("queue", schedule.Queue)
+	span.SetTag("type", schedule.Type)
+	span.SetTag("cron", schedule.CronSchedule)
+
+	tx, err := q.db.BeginTx(ctx, nil)
+	if err != nil {
+		return err
+	}
+	defer func() {
+		if err != nil {
+			_ = tx.Rollback()
+			return
+		}
+		err = tx.Commit()
+	}()
+
+	builder := squirrel.StatementBuilder.
+		PlaceholderFormat(squirrel.Dollar).
+		RunWith(db.WrapWithTracing(tx))
+
+	_, err = tx.ExecContext(ctx, `LOCK TABLE schedules IN ACCESS EXCLUSIVE MODE;`)
+	if err != nil {
+		return fmt.Errorf("failed to lock `schedules`: %w", err)
+	}
+
+	err = q.EnsureSchedule(ctx, builder, schedule)
+	if err == nil {
+		span.SetTag("existed", true)
+		return nil
+	}
+	if err == queue.ErrNotScheduled {
+		span.SetTag("existed", false)
+		err = q.Schedule(ctx, builder, schedule)
+	}
+
+	if err != nil {
+		return err
+	}
+
+	return nil
+}

pkg/queue/postgres/scheduler_test.go

@@ -271,3 +271,55 @@ func TestEnsure(t *testing.T) {
 		})
 	}
 }
+
+func TestAssertSchedule(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
+	defer cancel()
+
+	logrus.SetOutput(ioutil.Discard)
+	defer logrus.SetOutput(os.Stdout)
+
+	schedule := queue.TaskScheduleRequest{
+		TaskBase: queue.TaskBase{
+			Queue: "test",
+			Type:  "type",
+			Spec:  queue.Spec(`{"some":"value"}`),
+		},
+		CronSchedule: "* * * * *", // every minute
+	}
+
+	t.Run("asserts and schedules a non existent schedule", func(t *testing.T) {
+		_, db := dbtest.GetDatabase(t)
+		defer db.Close()
+		require.NoError(t, SetupTables(ctx, db, nil))
+
+		q := NewScheduler(db)
+		err := q.AssertSchedule(ctx, schedule)
+		require.NoError(t, err)
+		dbtest.EqualCount(t, db, 1, "schedules", squirrel.Eq{
+			"task_queue":    schedule.Queue,
+			"task_type":     schedule.Type,
+			"cron_schedule": schedule.CronSchedule,
+			"task_spec":     []byte(schedule.Spec),
+		})
+	})
+
+	t.Run("asserts and does not create a schedule if exists", func(t *testing.T) {
+		_, db := dbtest.GetDatabase(t)
+		defer db.Close()
+		require.NoError(t, SetupTables(ctx, db, nil))
+
+		q := NewScheduler(db)
+		err := q.AssertSchedule(ctx, schedule)
+		require.NoError(t, err)
+		err = q.AssertSchedule(ctx, schedule)
+		require.NoError(t, err)
+		dbtest.EqualCount(t, db, 1, "schedules", squirrel.Eq{
+			"task_queue":    schedule.Queue,
+			"task_type":     schedule.Type,
+			"cron_schedule": schedule.CronSchedule,
+			"task_spec":     []byte(schedule.Spec),
+		})
+	})
+
+}

pkg/queue/scheduler.go

@@ -23,6 +23,9 @@ type Scheduler interface {
 	// currently exists. An error ErrNotScheduled is returned if a current task cannot be found.
 	// implementation and validation errors may also be returned and should be checked for.
 	EnsureSchedule(ctx context.Context, builder cdb.SQLBuilder, task TaskScheduleRequest) (err error)
+	// AssertSchedule makes sure a schedule with the given parameters exists, if it does not
+	// this function will create one.
+	AssertSchedule(ctx context.Context, task TaskScheduleRequest) (err error)
 }
 
 type schedulerWithMetrics struct {
@@ -48,3 +51,6 @@ func (s *schedulerWithMetrics) Schedule(ctx context.Context, builder cdb.SQLBuil
 func (s *schedulerWithMetrics) EnsureSchedule(ctx context.Context, builder cdb.SQLBuilder, task TaskScheduleRequest) error {
 	return s.s.EnsureSchedule(ctx, builder, task)
 }
+func (s *schedulerWithMetrics) AssertSchedule(ctx context.Context, task TaskScheduleRequest) error {
+	return s.s.AssertSchedule(ctx, task)
+}