Automated cherry pick of #34247 (#34257)

Automatic Merge

server/channels/store/sqlstore/tokens_store.go

@@ -78,16 +78,41 @@ func (s SqlTokenStore) GetByToken(tokenString string) (*model.Token, error) {
 	return &token, nil
 }
 
-func (s SqlTokenStore) ConsumeOnce(tokenStr string) (*model.Token, error) {
+func (s SqlTokenStore) ConsumeOnce(tokenType, tokenStr string) (*model.Token, error) {
 	var token model.Token
 
-	query := `DELETE FROM Tokens WHERE Token = ? RETURNING *`
+	if s.DriverName() == model.DatabaseDriverPostgres {
+		query := `DELETE FROM Tokens WHERE Type = ? AND Token = ? RETURNING *`
+		if err := s.GetMaster().Get(&token, query, tokenType, tokenStr); err != nil {
+			if err == sql.ErrNoRows {
+				return nil, store.NewErrNotFound("Token", tokenStr)
+			}
+			return nil, errors.Wrapf(err, "failed to consume token with type %s", tokenType)
+		}
+		return &token, nil
+	}
 
-	if err := s.GetMaster().Get(&token, query, tokenStr); err != nil {
+	transaction, err := s.GetMaster().Beginx()
+	if err != nil {
+		return nil, errors.Wrap(err, "failed to begin transaction")
+	}
+	defer finalizeTransactionX(transaction, &err)
+
+	query := `SELECT * FROM Tokens WHERE Type = ? AND Token = ? FOR UPDATE`
+	if err = transaction.Get(&token, query, tokenType, tokenStr); err != nil {
 		if err == sql.ErrNoRows {
 			return nil, store.NewErrNotFound("Token", tokenStr)
 		}
-		return nil, errors.Wrapf(err, "failed to consume token")
+		return nil, errors.Wrapf(err, "failed to select token with type %s", tokenType)
+	}
+
+	deleteQuery := `DELETE FROM Tokens WHERE Type = ? AND Token = ?`
+	if _, err = transaction.Exec(deleteQuery, tokenType, tokenStr); err != nil {
+		return nil, errors.Wrapf(err, "failed to delete token with type %s", tokenType)
+	}
+
+	if err = transaction.Commit(); err != nil {
+		return nil, errors.Wrap(err, "failed to commit transaction")
 	}
 
 	return &token, nil