Use INSERT ON CONFLICT for retention task upsert (#60)

* Use INSERT ON CONFLICT for retention task upsert

**What**
- When deploying some of the apps with the retention policy, we saw the
  startup get stuck during the retention task assertion. Moving a an
  INSERT ON CONFLICT upsert method avoids the table lock and has the
  same effect.
- Add a partial unique index to the schedules table that is required for
  the upsert flow
- Any project using the retention policies will need to add this
  migration to their startup

```sql
CREATE UNIQUE INDEX IF NOT EXISTS unique_retention_idx
ON schedules (task_queue, task_type, (task_spec->>'queueName'), (task_spec->>'taskType'), (task_spec->>'status'))
WHERE task_type = 'retention';
```

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

Author: LucasRoesler

pkg/queue/postgres/retention.go

@@ -12,7 +12,9 @@ import (
 	cdb "github.com/contiamo/go-base/v2/pkg/db"
 	"github.com/contiamo/go-base/v2/pkg/queue"
 	"github.com/contiamo/go-base/v2/pkg/queue/handlers"
+	cvalidation "github.com/contiamo/go-base/v2/pkg/validation"
 	"github.com/opentracing/opentracing-go"
+	uuid "github.com/satori/go.uuid"
 )
 
 const (
@@ -59,67 +61,43 @@ func AssertRetentionSchedule(ctx context.Context, db *sql.DB, queueName string,
 		},
 		CronSchedule: fmt.Sprintf("%d * * * *", when), // every hour at minute "when"
 	}
-	tx, err := db.BeginTx(ctx, nil)
-	if err != nil {
-		return fmt.Errorf("can not start transaction for scheduling: %w", err)
-	}
-	defer func() {
-		if err != nil {
-			_ = tx.Rollback()
-			return
-		}
-		err = tx.Commit()
-	}()
-
-	_, err = tx.ExecContext(ctx, `LOCK TABLE schedules IN ACCESS EXCLUSIVE MODE;`)
+
+	err = cvalidation.CronTab(retentionSchedule.CronSchedule)
 	if err != nil {
-		return fmt.Errorf("failed to lock `schedules`: %w", err)
+		return err
 	}
 
 	builder := squirrel.StatementBuilder.
 		PlaceholderFormat(squirrel.Dollar).
-		RunWith(cdb.WrapWithTracing(tx))
-
-	var exists int
-	// use a unique error name here otherwise the sql.ErrNoRows might shadow
-	// us and things will break. This is also handled by the named error return
-	// variable, but this makes the code easier to copy and paste
-	existsErr := builder.Select("1").
-		From("schedules").
-		Where(squirrel.Eq{
-			"task_queue":              MaintenanceTaskQueue,
-			"task_type":               RetentionTask,
-			"task_spec->>'queueName'": queueName,
-			"task_spec->>'taskType'":  taskType,
-			"task_spec->>'status'":    status,
-		}).ScanContext(ctx, &exists)
-	if existsErr != nil && existsErr != sql.ErrNoRows {
-		return fmt.Errorf("can not verify existing schedule: %w", existsErr)
-	}
-
-	// will only non-zero if err is nil and task is not found
-	if exists == 0 {
-		span.SetTag("created", true)
-		// pass nil db because it doesn't need the raw db
-		return NewScheduler(nil).Schedule(ctx, builder, retentionSchedule)
-	}
-
-	span.SetTag("updated", true)
-	res, err := builder.Update("schedules").
-		Where(squirrel.Eq{
-			"task_queue":              MaintenanceTaskQueue,
-			"task_type":               RetentionTask,
-			"task_spec->>'queueName'": queueName,
-			"task_spec->>'taskType'":  taskType,
-			"task_spec->>'status'":    status,
-		}).
-		Set("updated_at", time.Now()).
-		Set("task_spec", retentionSchedule.Spec).
-		Set("cron_schedule", retentionSchedule.CronSchedule).
-		Set("next_execution_time", time.Now()).
-		ExecContext(ctx)
+		RunWith(cdb.WrapWithTracing(db))
+
+	scheduleID := uuid.NewV4().String()
+	res, err := builder.Insert("schedules").
+		Columns(
+			"schedule_id",
+			"task_queue",
+			"task_type",
+			"task_spec",
+			"cron_schedule",
+			"next_execution_time",
+		).
+		Values(
+			scheduleID,
+			retentionSchedule.Queue,
+			retentionSchedule.Type,
+			retentionSchedule.Spec,
+			retentionSchedule.CronSchedule,
+			time.Now(), // the schedule will enqueue the task immediately
+		).Suffix(`
+			ON CONFLICT (task_queue,task_type,(task_spec->>'queueName'),(task_spec->>'taskType'),(task_spec->>'status')) WHERE task_type='retention'
+			DO UPDATE SET
+				updated_at=now(),
+				next_execution_time=now(),
+				task_spec=EXCLUDED.task_spec,
+				cron_schedule=EXCLUDED.cron_schedule
+		`).ExecContext(ctx)
 	if err != nil {
-		return fmt.Errorf("can not update existing schdule: %w", err)
+		return fmt.Errorf("can not upsert retention schedule: %w", err)
 	}
 
 	updated, err := res.RowsAffected()

pkg/queue/postgres/scheduler.go

@@ -175,7 +175,7 @@ func (q *scheduler) AssertSchedule(ctx context.Context, schedule queue.TaskSched
 		PlaceholderFormat(squirrel.Dollar).
 		RunWith(cdb.WrapWithTracing(tx))
 
-	_, err = tx.ExecContext(ctx, `LOCK TABLE schedules IN ACCESS EXCLUSIVE MODE;`)
+	_, err = tx.ExecContext(ctx, `LOCK TABLE schedules IN SHARE MODE;`)
 	if err != nil {
 		return fmt.Errorf("failed to lock `schedules`: %w", err)
 	}

pkg/queue/postgres/setup.go

@@ -102,6 +102,13 @@ var (
 			Table:   SchedulesTable,
 			Columns: []string{"created_at DESC", "updated_at DESC"},
 		},
+		{
+			Table:     SchedulesTable,
+			Name:      "unique_retention_idx",
+			Columns:   []string{"task_queue", "task_type", "(task_spec->>'queueName')", "(task_spec->>'taskType')", "(task_spec->>'status')"},
+			Unique:    true,
+			Condition: fmt.Sprintf("task_type='%s'", RetentionTask),
+		},
 
 		// tasks
 		{
@@ -310,9 +317,10 @@ func syncTable(ctx context.Context, db db.SQLDB, tableName string, initColumns t
 }
 
 func listColumns(ctx context.Context, db db.SQLDB, tableName string) (columns map[string]nothing, err error) {
-	rows, err := db.QueryContext(ctx, `
-SELECT column_name FROM information_schema.columns WHERE table_name = $1;
-`, tableName)
+	rows, err := db.QueryContext(ctx,
+		`SELECT column_name FROM information_schema.columns WHERE table_name = $1;`,
+		tableName,
+	)
 	if err != nil {
 		return nil, err
 	}
@@ -349,9 +357,12 @@ func (s tableColumnSet) generateStatements() []string {
 
 // index describes all important properties of an index
 type index struct {
-	Table   string
-	Columns []string
-	Type    string
+	Name      string
+	Table     string
+	Columns   []string
+	Type      string
+	Unique    bool
+	Condition string
 }
 type indexList []index
 
@@ -361,22 +372,37 @@ func (l indexList) generateStatements() []string {
 	indexDefinitions := make([]string, 0, len(l))
 	for _, index := range l {
 		stmt := strings.Builder{}
-		stmt.WriteString("CREATE INDEX IF NOT EXISTS ")
+		stmt.WriteString("CREATE ")
+		if index.Unique {
+			stmt.WriteString("UNIQUE ")
+		}
+		stmt.WriteString("INDEX IF NOT EXISTS ")
 
 		columnList := strings.Join(index.Columns, ",")
 
-		stmt.WriteString(fmt.Sprintf(
-			"%s_%s_idx ON %s",
-			index.Table,
-			cstrings.ToUnderscoreCase(columnList),
-			index.Table,
-		))
+		// is a noop when string is empty
+		stmt.WriteString(index.Name)
+		if index.Name == "" {
+			stmt.WriteString(fmt.Sprintf(
+				"%s_%s_idx",
+				index.Table,
+				cstrings.ToUnderscoreCase(columnList),
+			))
+		}
+
+		stmt.WriteString(" ON ")
+		stmt.WriteString(index.Table)
 
 		if index.Type != "" {
 			stmt.WriteString(fmt.Sprintf(" USING %s ", index.Type))
 		}
 
-		stmt.WriteString("(" + columnList + ");")
+		stmt.WriteString("(" + columnList + ")")
+		if index.Condition != "" {
+			stmt.WriteString(" WHERE ")
+			stmt.WriteString(index.Condition)
+		}
+		stmt.WriteString(";")
 		indexDefinitions = append(indexDefinitions, stmt.String())
 	}
 

pkg/queue/postgres/setup_test.go

@@ -7,6 +7,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/Masterminds/squirrel"
 	dbtest "github.com/contiamo/go-base/v2/pkg/db/test"
 	"github.com/contiamo/go-base/v2/pkg/queue"
 	"github.com/sirupsen/logrus"
@@ -27,6 +28,20 @@ func TestSetupTables(t *testing.T) {
 		_, db := dbtest.GetDatabase(t)
 		defer db.Close()
 		require.NoError(t, SetupTables(ctx, db, nil))
+
+		dbtest.EqualCount(t, db, 1, "pg_indexes", squirrel.And{
+			squirrel.Eq{
+				"indexname": "unique_retention_idx",
+				"tablename": "schedules",
+			},
+			squirrel.Like{"indexdef": "CREATE UNIQUE INDEX%"},
+			squirrel.Like{"indexdef": "%task_queue%"},
+			squirrel.Like{"indexdef": "%task_spec%"},
+			squirrel.Like{"indexdef": "%queueName%"},
+			squirrel.Like{"indexdef": "%taskType%"},
+			squirrel.Like{"indexdef": "%status%"},
+			squirrel.Like{"indexdef": "%WHERE%"},
+		})
 	})
 
 	t.Run("bootstraps with references", func(t *testing.T) {