RBAC: Batch update on scope split migration (#72182)

* RBAC: Make the SplitScope migration concurrent

* Benchmark multiple alternatives: (updates in a loop, batch update, concurrent batch update)

* Only keep batching since mysql 5.7 does not seem to support concurrent batching

* Update pkg/services/accesscontrol/migrator/migrator.go

Co-authored-by: Ieva <ieva.vasiljeva@grafana.com>

---------

Co-authored-by: Ieva <ieva.vasiljeva@grafana.com>

pkg/services/accesscontrol/migrator/migrator.go

@@ -6,39 +6,106 @@ import (
 
 	"github.com/grafana/grafana/pkg/infra/db"
 	"github.com/grafana/grafana/pkg/infra/log"
-	"github.com/grafana/grafana/pkg/services/accesscontrol"
+	ac "github.com/grafana/grafana/pkg/services/accesscontrol"
 	"github.com/grafana/grafana/pkg/services/sqlstore"
+	"github.com/grafana/grafana/pkg/services/sqlstore/session"
+)
+
+var (
+	batchSize = 1000
 )
 
 func MigrateScopeSplit(db db.DB, log log.Logger) error {
 	t := time.Now()
-	var count = 0
+	ctx := context.Background()
-	err := db.WithTransactionalDbSession(context.Background(), func(sess *sqlstore.DBSession) error {
+	cnt := 0
-		var permissions []accesscontrol.Permission
 
-		err := sess.SQL("SELECT * FROM permission WHERE NOT scope = '' AND identifier = ''").Find(&permissions)
+	// Search for the permissions to update
-		if err != nil {
+	var permissions []ac.Permission
+	if errFind := db.WithTransactionalDbSession(ctx, func(sess *sqlstore.DBSession) error {
+		return sess.SQL("SELECT * FROM permission WHERE NOT scope = '' AND identifier = ''").Find(&permissions)
+	}); errFind != nil {
+		log.Error("could not search for permissions to update", "migration", "scopeSplit", "error", errFind)
+		return errFind
+	}
+
+	if len(permissions) == 0 {
+		log.Debug("no permission require a scope split", "migration", "scopeSplit")
+		return nil
+	}
+
+	errBatchUpdate := batch(len(permissions), batchSize, func(start, end int) error {
+		n := end - start
+
+		// IDs to remove
+		delQuery := "DELETE FROM permission WHERE id IN ("
+		delArgs := make([]interface{}, 0, n)
+
+		// Query to insert the updated permissions
+		insertQuery := "INSERT INTO permission (id, role_id, action, scope, kind, attribute, identifier, created, updated) VALUES "
+		insertArgs := make([]interface{}, 0, 9*n)
+
+		// Prepare batch of updated permissions
+		for i := start; i < end; i++ {
+			kind, attribute, identifier := permissions[i].SplitScope()
+
+			delQuery += "?,"
+			delArgs = append(delArgs, permissions[i].ID)
+
+			insertQuery += "(?, ?, ?, ?, ?, ?, ?, ?, ?),"
+			insertArgs = append(insertArgs, permissions[i].ID, permissions[i].RoleID,
+				permissions[i].Action, permissions[i].Scope,
+				kind, attribute, identifier,
+				permissions[i].Created, t,
+			)
+		}
+		// Remove trailing ','
+		insertQuery = insertQuery[:len(insertQuery)-1]
+
+		// Remove trailing ',' and close brackets
+		delQuery = delQuery[:len(delQuery)-1] + ")"
+
+		// Batch update the permissions
+		if errBatchUpdate := db.GetSqlxSession().WithTransaction(ctx, func(tx *session.SessionTx) error {
+			if _, errDel := tx.Exec(ctx, delQuery, delArgs...); errDel != nil {
+				log.Error("error deleting permissions", "migration", "scopeSplit", "error", errDel)
+				return errDel
+			}
+			if _, errInsert := tx.Exec(ctx, insertQuery, insertArgs...); errInsert != nil {
+				log.Error("error saving permissions", "migration", "scopeSplit", "error", errInsert)
+				return errInsert
+			}
+			return nil
+		}); errBatchUpdate != nil {
+			log.Error("error updating permission batch", "migration", "scopeSplit", "start", start, "end", end)
+			return errBatchUpdate
+		}
+
+		cnt += end - start
+		return nil
+	})
+	if errBatchUpdate != nil {
+		log.Error("could not migrate permissions", "migration", "scopeSplit", "total", len(permissions), "succeeded", cnt, "left", len(permissions)-cnt, "error", errBatchUpdate)
+		return errBatchUpdate
+	}
+
+	log.Debug("migrated permissions", "migration", "scopeSplit", "total", len(permissions), "succeeded", cnt, "in", time.Since(t))
+	return nil
+}
+
+func batch(count, batchSize int, eachFn func(start, end int) error) error {
+	for i := 0; i < count; {
+		end := i + batchSize
+		if end > count {
+			end = count
+		}
+
+		if err := eachFn(i, end); err != nil {
 			return err
 		}
 
-		for i, p := range permissions {
+		i = end
-			count++
+	}
-			kind, attribute, identifier := p.SplitScope()
 
-			permissions[i].Kind = kind
+	return nil
-			permissions[i].Attribute = attribute
-			permissions[i].Identifier = identifier
-
-			_, err := sess.Exec("UPDATE permission SET kind = ?, attribute = ?, identifier = ? WHERE id = ?", permissions[i].Kind, permissions[i].Attribute, permissions[i].Identifier, permissions[i].ID)
-			if err != nil {
-				return err
-			}
-		}
-
-		return nil
-	})
-
-	log.Debug("Migrated permissions ", "count", count, "in", time.Since(t))
-
-	return err
 }

pkg/services/accesscontrol/migrator/migrator_bench_test.go

@@ -0,0 +1,27 @@
+package migrator
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/grafana/grafana/pkg/infra/db"
+	"github.com/grafana/grafana/pkg/infra/log"
+)
+
+func benchScopeSplitConcurrent(b *testing.B, count int) {
+	store := db.InitTestDB(b)
+	// Populate permissions
+	require.NoError(b, batchInsertPermissions(count, store), "could not insert permissions")
+	logger := log.New("migrator.test")
+	b.ResetTimer()
+
+	for n := 0; n < b.N; n++ {
+		err := MigrateScopeSplit(store, logger)
+		require.NoError(b, err)
+	}
+}
+
+func BenchmarkMigrateScopeSplitConcurrent_50K(b *testing.B) { benchScopeSplitConcurrent(b, 50000) }
+
+func BenchmarkMigrateScopeSplitConcurrent_100K(b *testing.B) { benchScopeSplitConcurrent(b, 100000) }

pkg/services/accesscontrol/migrator/migrator_test.go

@@ -0,0 +1,61 @@
+package migrator
+
+import (
+	"context"
+	"fmt"
+	"testing"
+	"time"
+
+	"github.com/grafana/grafana/pkg/infra/db"
+	"github.com/grafana/grafana/pkg/infra/log"
+	ac "github.com/grafana/grafana/pkg/services/accesscontrol"
+	"github.com/grafana/grafana/pkg/services/sqlstore"
+	"github.com/stretchr/testify/require"
+)
+
+func batchInsertPermissions(cnt int, sqlStore db.DB) error {
+	now := time.Now()
+
+	return batch(cnt, batchSize, func(start, end int) error {
+		n := end - start
+		permissions := make([]ac.Permission, 0, n)
+		for i := start + 1; i < end+1; i++ {
+			permissions = append(permissions, ac.Permission{
+				RoleID:  1,
+				Action:  "action",
+				Scope:   fmt.Sprintf("resource:uid:%v", i),
+				Created: now,
+				Updated: now,
+			})
+		}
+		return sqlStore.WithDbSession(context.Background(), func(sess *db.Session) error {
+			_, err := sess.Insert(permissions)
+			return err
+		})
+	})
+}
+
+func TestMigrateScopeSplit(t *testing.T) {
+	sqlStore := db.InitTestDB(t)
+	logger := log.New("accesscontrol.migrator.test")
+
+	// Populate permissions
+	require.NoError(t, batchInsertPermissions(3*batchSize, sqlStore), "could not insert permissions")
+
+	// Migrate
+	require.NoError(t, MigrateScopeSplit(sqlStore, logger))
+
+	// Check migration result
+	permissions := make([]ac.Permission, 0, 3*batchSize)
+	errFind := sqlStore.WithDbSession(context.Background(), func(sess *sqlstore.DBSession) error {
+		return sess.Find(&permissions)
+	})
+	require.NoError(t, errFind, "could not find permissions in store")
+
+	for i := range permissions {
+		require.Equal(t, fmt.Sprintf("resource:uid:%v", i+1), permissions[i].Scope, "scope should have been preserved")
+		require.Equal(t, "resource", permissions[i].Kind)
+		require.Equal(t, "uid", permissions[i].Attribute)
+		require.Equal(t, fmt.Sprintf("%v", i+1), permissions[i].Identifier)
+	}
+}