ServerLock: Fix missing return ID for postgres (#79878)

fix missing return ID for postgres

pkg/infra/serverlock/serverlock.go

@@ -3,7 +3,6 @@ package serverlock
 import (
 	"context"
 	"errors"
-	"fmt"
 	"math/rand"
 	"time"
 
@@ -12,6 +11,7 @@ import (
 	"github.com/grafana/grafana/pkg/infra/db"
 	"github.com/grafana/grafana/pkg/infra/log"
 	"github.com/grafana/grafana/pkg/infra/tracing"
+	"github.com/grafana/grafana/pkg/services/sqlstore"
 	"github.com/grafana/grafana/pkg/services/sqlstore/migrator"
 )
 
@@ -122,47 +122,8 @@ func (sl *ServerLockService) getOrCreate(ctx context.Context, actionName string)
 			OperationUID:  actionName,
 			LastExecution: 0,
 		}
-
-		affected := int64(1)
-		rawSQL := `INSERT INTO server_lock (operation_uid, last_execution, version) VALUES (?, ?, ?)`
-		if sl.SQLStore.GetDBType() == migrator.Postgres {
-			rawSQL += ` RETURNING id`
-			var id int64
-			_, err := dbSession.SQL(rawSQL, lockRow.OperationUID, lockRow.LastExecution, 0).Get(&id)
-			if err != nil {
-				return err
-			}
-			lockRow.Id = id
-		} else {
-			res, err := dbSession.Exec(
-				rawSQL,
-				lockRow.OperationUID, lockRow.LastExecution, 0)
-			if err != nil {
-				return err
-			}
-			lastID, err := res.LastInsertId()
-			if err != nil {
-				fmt.Println("Error getting last insert id", err)
-				sl.log.FromContext(ctx).Error("Error getting last insert id", "actionName", actionName, "error", err)
-			}
-			lockRow.Id = lastID
-
-			affected, err = res.RowsAffected()
-			if err != nil {
-				sl.log.FromContext(ctx).Error("Error getting rows affected", "actionName", actionName, "error", err)
-			}
-		}
-
-		if affected != 1 || lockRow.Id == 0 {
-			// this means that there was no error but there is something not working correctly
-			sl.log.FromContext(ctx).Error("Expected rows affected to be 1 if there was no error",
-				"actionName", actionName,
-				"rowsAffected", affected,
-				"lockRow ID", lockRow.Id)
-		}
-
-		result = lockRow
-		return nil
+		result, err = sl.createLock(ctx, lockRow, dbSession)
+		return err
 	})
 
 	return result, err
@@ -297,31 +258,11 @@ func (sl *ServerLockService) acquireForRelease(ctx context.Context, actionName s
 		if has {
 			if sl.isLockWithinInterval(result, maxInterval) {
 				return &ServerLockExistsError{actionName: actionName}
-			} else {
-				// lock has timed out, so we update the timestamp
-				result.LastExecution = time.Now().Unix()
-				res, err := dbSession.Exec("UPDATE server_lock SET last_execution = ? WHERE operation_uid = ?",
-					result.LastExecution, actionName)
-				if err != nil {
-					return err
-				}
-
-				affected, err := res.RowsAffected()
-				if err != nil {
-					ctxLogger.Error("Error getting rows affected", "actionName", actionName, "error", err)
-				}
-
-				if affected != 1 {
-					ctxLogger.Error("Expected rows affected to be 1 if there was no error", "actionName", actionName, "rowsAffected", affected)
-				}
-
-				return nil
 			}
-		} else {
-			// lock not found, creating it
-			res, err := dbSession.Exec(
-				"INSERT INTO server_lock (operation_uid, last_execution, version) VALUES (?, ?, ?)",
-				actionName, time.Now().Unix(), 0)
+			// lock has timed out, so we update the timestamp
+			result.LastExecution = time.Now().Unix()
+			res, err := dbSession.Exec("UPDATE server_lock SET last_execution = ? WHERE operation_uid = ?",
+				result.LastExecution, actionName)
 			if err != nil {
 				return err
 			}
@@ -331,20 +272,20 @@ func (sl *ServerLockService) acquireForRelease(ctx context.Context, actionName s
 				ctxLogger.Error("Error getting rows affected", "actionName", actionName, "error", err)
 			}
 
-			lastID, err := res.LastInsertId()
-			if err != nil {
-				ctxLogger.Error("Error getting last insert id", "actionName", actionName, "error", err)
+			if affected != 1 {
+				ctxLogger.Error("Expected rows affected to be 1 if there was no error", "actionName", actionName, "rowsAffected", affected)
 			}
 
-			if affected != 1 || lastID == 0 {
-				// this means that there was no error but there is something not working correctly
-				ctxLogger.Error("Expected rows affected to be 1 if there was no error",
-					"actionName", actionName,
-					"rowsAffected", affected,
-					"lastID", lastID)
-			}
+			return nil
 		}
-		return nil
+
+		// lock not found, creating it
+		lock := &serverLock{
+			OperationUID:  actionName,
+			LastExecution: time.Now().Unix(),
+		}
+		_, err = sl.createLock(ctx, lock, dbSession)
+		return err
 	})
 
 	return err
@@ -399,3 +340,44 @@ func (sl *ServerLockService) executeFunc(ctx context.Context, actionName string,
 
 	ctxLogger.Debug("Execution finished", "actionName", actionName, "duration", time.Since(start))
 }
+
+func (sl *ServerLockService) createLock(ctx context.Context,
+	lockRow *serverLock, dbSession *sqlstore.DBSession) (*serverLock, error) {
+	affected := int64(1)
+	rawSQL := `INSERT INTO server_lock (operation_uid, last_execution, version) VALUES (?, ?, ?)`
+	if sl.SQLStore.GetDBType() == migrator.Postgres {
+		rawSQL += ` RETURNING id`
+		var id int64
+		_, err := dbSession.SQL(rawSQL, lockRow.OperationUID, lockRow.LastExecution, 0).Get(&id)
+		if err != nil {
+			return nil, err
+		}
+		lockRow.Id = id
+	} else {
+		res, err := dbSession.Exec(
+			rawSQL,
+			lockRow.OperationUID, lockRow.LastExecution, 0)
+		if err != nil {
+			return nil, err
+		}
+		lastID, err := res.LastInsertId()
+		if err != nil {
+			sl.log.FromContext(ctx).Error("Error getting last insert id", "actionName", lockRow.OperationUID, "error", err)
+		}
+		lockRow.Id = lastID
+
+		affected, err = res.RowsAffected()
+		if err != nil {
+			sl.log.FromContext(ctx).Error("Error getting rows affected", "actionName", lockRow.OperationUID, "error", err)
+		}
+	}
+
+	if affected != 1 || lockRow.Id == 0 {
+		sl.log.FromContext(ctx).Error("Expected rows affected to be 1 if there was no error",
+			"actionName", lockRow.OperationUID,
+			"rowsAffected", affected,
+			"lockRow ID", lockRow.Id)
+	}
+
+	return lockRow, nil
+}