[MM-63314] Fix ClaimJob in HA environments (#30383)

* ClaimJob now returns newly claimed job

* internal code affected by change

* test changes required

* two branches: for mysql, use transaction; for postgres, use returning

* two branches: for mysql, use transaction; for postgres, use returning

* use same millis value for LastActivityAt and StartAt

* blank commit

---------

Co-authored-by: Mattermost Build <build@mattermost.com>

server/channels/store/sqlstore/job_store.go

@@ -172,34 +172,86 @@ func (jss SqlJobStore) UpdateStatus(id string, status string) (*model.Job, error
 	return job, nil
 }
 
-func (jss SqlJobStore) UpdateStatusOptimistically(id string, currentStatus string, newStatus string) (bool, error) {
+func (jss SqlJobStore) UpdateStatusOptimistically(id string, currentStatus string, newStatus string) (*model.Job, error) {
+	lastActivityAndStartTime := model.GetMillis()
+
+	if jss.DriverName() == model.DatabaseDriverMysql {
+		tx, err := jss.GetMaster().Beginx()
+		if err != nil {
+			return nil, errors.Wrap(err, "begin_transaction")
+		}
+		defer finalizeTransactionX(tx, &err)
+
+		builder := jss.getQueryBuilder().
+			Update("Jobs").
+			Set("LastActivityAt", lastActivityAndStartTime).
+			Set("Status", newStatus).
+			Where(sq.Eq{"Id": id, "Status": currentStatus})
+
+		if newStatus == model.JobStatusInProgress {
+			builder = builder.Set("StartAt", lastActivityAndStartTime)
+		}
+
+		sqlResult, err := tx.ExecBuilder(builder)
+		if err != nil {
+			return nil, errors.Wrapf(err, "failed to update Job with id=%s", id)
+		}
+		rows, err := sqlResult.RowsAffected()
+		if err != nil {
+			return nil, errors.Wrap(err, "unable to get rows affected")
+		}
+		if rows != 1 {
+			return nil, nil
+		}
+
+		getBuilder := jss.getQueryBuilder().
+			Select("*").
+			From("Jobs").
+			Where(sq.Eq{"Id": id, "Status": newStatus})
+
+		var job model.Job
+		if err = tx.GetBuilder(&job, getBuilder); err != nil {
+			if errors.Is(err, sql.ErrNoRows) {
+				return nil, store.NewErrNotFound("Job", id)
+			}
+			return nil, errors.Wrapf(err, "failed to get Job with id=%s", id)
+		}
+
+		err = tx.Commit()
+		if err != nil {
+			return nil, errors.Wrap(err, "commit_transaction")
+		}
+
+		return &job, nil
+	}
+
+	// For PostgreSQL, use RETURNING to get the updated job in a single query
 	builder := jss.getQueryBuilder().
 		Update("Jobs").
-		Set("LastActivityAt", model.GetMillis()).
+		Set("LastActivityAt", lastActivityAndStartTime).
 		Set("Status", newStatus).
-		Where(sq.Eq{"Id": id, "Status": currentStatus})
+		Where(sq.Eq{"Id": id, "Status": currentStatus}).
+		Suffix("RETURNING *")
 
 	if newStatus == model.JobStatusInProgress {
-		builder = builder.Set("StartAt", model.GetMillis())
-	}
-	query, args, err := builder.ToSql()
-	if err != nil {
-		return false, errors.Wrap(err, "job_tosql")
+		builder = builder.Set("StartAt", lastActivityAndStartTime)
 	}
 
-	sqlResult, err := jss.GetMaster().Exec(query, args...)
-	if err != nil {
-		return false, errors.Wrapf(err, "failed to update Job with id=%s", id)
-	}
-	rows, err := sqlResult.RowsAffected()
-	if err != nil {
-		return false, errors.Wrap(err, "unable to get rows affected")
-	}
-	if rows != 1 {
-		return false, nil
+	var job []*model.Job
+	if err := jss.GetMaster().SelectBuilder(&job, builder); err != nil {
+		if errors.Is(err, sql.ErrNoRows) {
+			return nil, store.NewErrNotFound("Job", id)
+		}
+		return nil, errors.Wrapf(err, "failed to update Job with id=%s", id)
 	}
 
-	return true, nil
+	// we are updating by id, so we should only ever update 1 job
+	if len(job) != 1 {
+		// no row was updated, but no error above, so to remain consistent we return nil, nil
+		return nil, nil
+	}
+
+	return job[0], nil
 }
 
 func (jss SqlJobStore) Get(c request.CTX, id string) (*model.Job, error) {
@@ -213,7 +265,7 @@ func (jss SqlJobStore) Get(c request.CTX, id string) (*model.Job, error) {
 
 	var status model.Job
 	if err = jss.GetReplica().Get(&status, query, args...); err != nil {
-		if err == sql.ErrNoRows {
+		if errors.Is(err, sql.ErrNoRows) {
 			return nil, store.NewErrNotFound("Job", id)
 		}
 		return nil, errors.Wrapf(err, "failed to get Job with id=%s", id)

server/channels/store/sqlstore/utils.go

@@ -5,6 +5,7 @@ package sqlstore
 
 import (
 	"database/sql"
+	"errors"
 	"fmt"
 	"io"
 	"net/url"
@@ -54,7 +55,7 @@ func MapStringsToQueryParams(list []string, paramPrefix string) (string, map[str
 // finalizeTransactionX ensures a transaction is closed after use, rolling back if not already committed.
 func finalizeTransactionX(transaction *sqlxTxWrapper, perr *error) {
 	// Rollback returns sql.ErrTxDone if the transaction was already closed.
-	if err := transaction.Rollback(); err != nil && err != sql.ErrTxDone {
+	if err := transaction.Rollback(); err != nil && !errors.Is(err, sql.ErrTxDone) {
 		*perr = merror.Append(*perr, err)
 	}
 }