Fix more LastInsertId calls (#101892)

* Fix more places where LastInsertId is used without proper directive for Spanner.

* Fix tests.

pkg/infra/serverlock/serverlock.go

@@ -374,6 +374,10 @@ func (sl *ServerLockService) createLock(ctx context.Context,
 		}
 		lockRow.Id = id
 	} else {
+		if sl.SQLStore.GetDBType() == migrator.Spanner {
+			rawSQL += " THEN RETURN id" // Required for successful LastInsertId call.
+		}
+
 		res, err := dbSession.Exec(
 			rawSQL,
 			lockRow.OperationUID, lockRow.LastExecution, 0)

pkg/infra/serverlock/serverlock_test.go

@@ -36,6 +36,7 @@ func TestServerLock(t *testing.T) {
 
 	first, err := sl.getOrCreate(context.Background(), operationUID)
 	require.NoError(t, err)
+	require.NotZero(t, first.Id)
 
 	t.Run("trying to create three new row locks", func(t *testing.T) {
 		expectedLastExecution := first.LastExecution
@@ -44,8 +45,8 @@ func TestServerLock(t *testing.T) {
 		for i := 0; i < 3; i++ {
 			latest, err = sl.getOrCreate(context.Background(), operationUID)
 			require.NoError(t, err)
-			assert.Equal(t, operationUID, first.OperationUID)
-			assert.Equal(t, int64(1), first.Id)
+			assert.Equal(t, first.OperationUID, operationUID)
+			assert.Equal(t, first.Id, latest.Id)
 		}
 
 		assert.Equal(t,
@@ -114,7 +115,7 @@ func TestLockAndRelease(t *testing.T) {
 			affectedRows, err := sess.Insert(&lock)
 			require.NoError(t, err)
 			require.Equal(t, int64(1), affectedRows)
-			require.Equal(t, int64(1), lock.Id)
+			require.NotZero(t, lock.Id)
 			return nil
 		})
 		require.NoError(t, err)

pkg/services/sqlstore/session/session.go

@@ -101,19 +101,20 @@ func (gtx *SessionTx) ExecWithReturningId(ctx context.Context, query string, arg
 }
 
 func execWithReturningId(ctx context.Context, driverName string, query string, sess Session, args ...any) (int64, error) {
-	supported := false
 	var id int64
 	if driverName == "postgres" {
 		query = fmt.Sprintf("%s RETURNING id", query)
-		supported = true
-	}
-	if supported {
 		err := sess.Get(ctx, &id, query, args...)
 		if err != nil {
 			return id, err
 		}
 		return id, nil
 	} else {
+		if driverName == "spanner" {
+			// LastInsertId requires THEN RETURN directive.
+			query = fmt.Sprintf("%s THEN RETURN id", query)
+		}
+
 		res, err := sess.Exec(ctx, query, args...)
 		if err != nil {
 			return id, err