Create task retention worker and scheduler (#53)

**What**
- Create new generic SQLTask spec and handler
- Create a new task handler that can be used to cleanup old tasks
  matching a specific filter. The caller can specify the queue, task
  type, status, age, and a custom sqlizer.
- Add utility to schedule this new cleanup task once per hour. The exact
  minute is chosen randomly as a basic attempt to spread our the
  workload.
- Improve the db test utility EqualCount to allow passing an error
  message and ars down to the require method inside of it. This allows
  the errors to be more expressive and hopefully easier to read
- Add new db setup utility that uses docker to ensure a db server is
  ready. This is intended as a dev mode only tool which is not left in
  the final committed code, but can be very useful to speedup test
  iteration and debugging.

Signed-off-by: Lucas Roesler <[email protected]>

Author: LucasRoesler

pkg/db/test/db.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"database/sql"
 	"fmt"
+	"os/exec"
 	"testing"
 	"time"
 
@@ -28,7 +29,7 @@ type DBInitializer func(context.Context, *sql.DB) error
 
 // EqualCount asserts that the count of rows matches the expected value given the table and WHERE query and args.EqualCount
 // Note that this is a simple  COUNT of rows in a single table. More complex queries should be constructed by hand.
-func EqualCount(t *testing.T, db *sql.DB, expected int, table string, filter squirrel.Sqlizer) int {
+func EqualCount(t *testing.T, db *sql.DB, expected int, table string, filter squirrel.Sqlizer, msgAndArgs ...interface{}) int {
 	var count int
 	err := squirrel.StatementBuilder.
 		PlaceholderFormat(squirrel.Dollar).
@@ -37,8 +38,8 @@ func EqualCount(t *testing.T, db *sql.DB, expected int, table string, filter squ
 		Where(filter).
 		RunWith(db).
 		Scan(&count)
-	require.NoError(t, err)
-	require.Equal(t, expected, count)
+	require.NoError(t, err, msgAndArgs...)
+	require.Equal(t, expected, count, msgAndArgs...)
 
 	return count
 }
@@ -112,3 +113,59 @@ func connectDB(name string) (db *sql.DB, err error) {
 
 	return sql.Open(cfg.DriverName, connStr+" password="+adminPassword)
 }
+
+// cleanupDB attempts to remove the test container created by EnsureDBReady
+func cleanupDB() {
+	cmd := exec.Command("docker", "rm", "-v", "-f", "go-base-postgres")
+	err := cmd.Run()
+	if err != nil {
+		fmt.Println("error during db cleanupdb: ", err)
+	}
+}
+
+// EnsureDBReady is a helper utility that can be inserted into a test during development to make it simpler
+// to run the test in isolation. Docker will be called to start a db that is read  to be used with GetDatabase.
+// Example usage:
+//
+//   err, cleanup := dbtest.EnsureDBReady(ctx)
+//   require.NoError(t, err)
+//   defer cleanup()
+//
+//   _, db := dbtest.GetDatabase(t)
+//   defer db.Close()
+//
+// While the code is somewhat robust to having existing dbs running, this is not intended to be left in the test code.
+// The goal is that you can run individual tests via your IDE integrations or using the CLI, e.g.
+//   go test -run ^TestRetentionHandler$`
+func EnsureDBReady(ctx context.Context) (error, func()) {
+	check := exec.CommandContext(ctx, "docker", "container", "inspect", "go-base-postgres", "-f", "{{.ID}}").Run()
+	if check == nil {
+		// container is already running
+		return nil, func() {}
+	}
+
+	if exitError, ok := check.(*exec.ExitError); ok && exitError.ExitCode() != 1 {
+		return fmt.Errorf("unexpected exit code when checking for running db: %w", check), func() {}
+	}
+
+	db := exec.CommandContext(ctx,
+		"docker",
+		"run",
+		"--rm",
+		"-d",
+		"--name", "go-base-postgres",
+		"-p", "0.0.0.0:5432:5432",
+		"-e", "POSTGRES_PASSWORD=localdev",
+		"-e", "POSTGRES_USER=contiamo_test",
+		"-e", "POSTGRES_DB=postgres",
+		"postgres:alpine",
+		"-c", "fsync=off",
+		"-c", "full_page_writes=off",
+		"-c", "synchronous_commit=off",
+	)
+
+	err := db.Run()
+
+	time.Sleep(3 * time.Second)
+	return err, cleanupDB
+}

pkg/queue/postgres/setup.go

@@ -11,6 +11,11 @@ import (
 )
 
 const (
+	// TasksTable is the name of the Postgres table that is used for tasks
+	TasksTable = "tasks"
+	// SchedulesTable is the name of the Postgres table used for schedules
+	SchedulesTable = "schedules"
+
 	createTableTmpl = `
 CREATE EXTENSION IF NOT EXISTS citext;
 
@@ -80,71 +85,71 @@ var (
 
 		// schedules
 		{
-			Table:   "schedules",
+			Table:   SchedulesTable,
 			Columns: []string{"next_execution_time DESC"},
 		},
 		{
-			Table:   "schedules",
+			Table:   SchedulesTable,
 			Columns: []string{"task_queue"},
 			Type:    "hash",
 		},
 		{
-			Table:   "schedules",
+			Table:   SchedulesTable,
 			Columns: []string{"task_type"},
 			Type:    "hash",
 		},
 		{
-			Table:   "schedules",
+			Table:   SchedulesTable,
 			Columns: []string{"created_at DESC", "updated_at DESC"},
 		},
 
 		// tasks
 		{
-			Table:   "tasks",
+			Table:   TasksTable,
 			Columns: []string{"queue"},
 			Type:    "hash",
 		},
 		{
-			Table:   "tasks",
+			Table:   TasksTable,
 			Columns: []string{"type"},
 			Type:    "hash",
 		},
 		{
-			Table:   "tasks",
+			Table:   TasksTable,
 			Columns: []string{"status"},
 			Type:    "hash",
 		},
 		{
-			Table:   "tasks",
+			Table:   TasksTable,
 			Columns: []string{"schedule_id"},
 			Type:    "hash",
 		},
 		{
-			Table:   "tasks",
+			Table:   TasksTable,
 			Columns: []string{"created_at", "last_heartbeat_at"},
 		},
 		{
-			Table:   "tasks",
+			Table:   TasksTable,
 			Columns: []string{"created_at DESC", "last_heartbeat_at DESC"},
 		},
 		{
-			Table:   "tasks",
+			Table:   TasksTable,
 			Columns: []string{"created_at DESC"},
 		},
 		{
-			Table:   "tasks",
+			Table:   TasksTable,
 			Columns: []string{"last_heartbeat_at DESC"},
 		},
 		{
-			Table:   "tasks",
+			Table:   TasksTable,
 			Columns: []string{"created_at DESC", "updated_at DESC"},
 		},
 		{
-			Table:   "tasks",
+			Table:   TasksTable,
 			Columns: []string{"started_at DESC"},
 		},
 		{
-			Table:   "tasks",
+			Table:   TasksTable,
 			Columns: []string{"finished_at DESC"},
 		},
 	}
@@ -170,14 +175,14 @@ func SetupTables(ctx context.Context, db db.SQLDB, references []ForeignReference
 	logrus.Debug("checking queue-related tables...")
 
 	logrus.Debug("checking `schedules` table...")
-	err = syncTable(ctx, db, "schedules", scheduleColumns, references)
+	err = syncTable(ctx, db, SchedulesTable, scheduleColumns, references)
 	if err != nil {
 		return err
 	}
 	logrus.Debug("`schedules` table is up to date")
 
 	logrus.Debug("checking `tasks` table...")
-	err = syncTable(ctx, db, "tasks", taskColumns, references)
+	err = syncTable(ctx, db, TasksTable, taskColumns, references)
 	if err != nil {
 		return err
 	}
@@ -194,12 +199,12 @@ func SetupTables(ctx context.Context, db db.SQLDB, references []ForeignReference
 	applyIndexes := make(indexList, 0, len(references)+len(indexes))
 	for _, ref := range references {
 		applyIndexes = append(applyIndexes, index{
-			Table:   "schedules",
+			Table:   SchedulesTable,
 			Columns: []string{ref.ColumnName},
 			Type:    "hash",
 		})
 		applyIndexes = append(applyIndexes, index{
-			Table:   "tasks",
+			Table:   TasksTable,
 			Columns: []string{ref.ColumnName},
 			Type:    "hash",
 		})

pkg/queue/workers/retention.go

@@ -0,0 +1,82 @@
+package workers
+
+import (
+	"context"
+	"database/sql"
+	"encoding/json"
+	"fmt"
+	"math/rand"
+	"time"
+
+	"github.com/Masterminds/squirrel"
+	"github.com/contiamo/go-base/pkg/queue"
+	"github.com/contiamo/go-base/pkg/queue/postgres"
+)
+
+const (
+	// MaintenanceTaskQueue task queue name used for all the periodic maintenance jobs.
+	// These are internal queue internal tasks
+	MaintenanceTaskQueue string = "queue-maintenance"
+
+	// RetentionTask is finished task cleanup type
+	RetentionTask queue.TaskType = "retention"
+)
+
+// NewRetentionHandler creates a task handler that will clean up old finished tasks
+func NewRetentionHandler(db *sql.DB) TaskHandler {
+	return NewSQLTaskHandler("RetentionHandler", db)
+}
+
+// AssertRetentionSchedule creates a new queue retention tasks for the supplied queue, finished tasks matching
+// the supplied parameters will be deleted
+func AssertRetentionSchedule(ctx context.Context, scheduler queue.Scheduler, queueName string, taskType queue.TaskType, status queue.TaskStatus, filter squirrel.Sqlizer, age time.Duration) error {
+	spec := createRetentionSpec(queueName, taskType, status, filter, age)
+	specBytes, err := json.Marshal(spec)
+	if err != nil {
+		return fmt.Errorf("can not build retention task spec: %w", err)
+	}
+	// randomly distribute the retention tasks throughout the hour
+	when := rand.Intn(60)
+	retentionSchedule := queue.TaskScheduleRequest{
+		TaskBase: queue.TaskBase{
+			Queue: MaintenanceTaskQueue,
+			Type:  RetentionTask,
+			Spec:  specBytes,
+		},
+		CronSchedule: fmt.Sprintf("%d * * * *", when), // every hour at minute "when"
+	}
+
+	return scheduler.AssertSchedule(ctx, retentionSchedule)
+}
+
+//createRetentionSpec builds the task retention job spec. It is split out to simplify test setup
+func createRetentionSpec(queueName string, taskType queue.TaskType, status queue.TaskStatus, filter squirrel.Sqlizer, age time.Duration) SQLExecTaskSpec {
+	spec := SQLExecTaskSpec{
+		SQL: "",
+	}
+
+	// use separate WHERE statements to make the order deterministic
+	deletionSQL := squirrel.Delete(postgres.TasksTable).
+		Where(squirrel.Eq{"status": status}).
+		Where(
+			// note that using this comparision allows us to use the index on
+			// finished_at, if yo use `age(now(), finished_at)`, this can not use the index
+			fmt.Sprintf("finished_at <= now() - interval '%f minutes'", age.Minutes()),
+		)
+
+	if queueName != "" {
+		deletionSQL = deletionSQL.Where(squirrel.Eq{"queue": queueName})
+	}
+
+	if taskType != "" {
+		deletionSQL = deletionSQL.Where(squirrel.Eq{"type": taskType})
+	}
+
+	if filter != nil {
+		deletionSQL = deletionSQL.Where(filter)
+	}
+
+	spec.SQL = squirrel.DebugSqlizer(deletionSQL)
+
+	return spec
+}