From 261045d1828c36e51f919e3131cbd75d98c85b78 Mon Sep 17 00:00:00 2001 From: Gabriel MABILLE Date: Fri, 28 Jul 2023 12:37:24 +0200 Subject: [PATCH] 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 --------- Co-authored-by: Ieva --- .../accesscontrol/migrator/migrator.go | 117 ++++++++++++++---- .../migrator/migrator_bench_test.go | 27 ++++ .../accesscontrol/migrator/migrator_test.go | 61 +++++++++ 3 files changed, 180 insertions(+), 25 deletions(-) create mode 100644 pkg/services/accesscontrol/migrator/migrator_bench_test.go create mode 100644 pkg/services/accesscontrol/migrator/migrator_test.go diff --git a/pkg/services/accesscontrol/migrator/migrator.go b/pkg/services/accesscontrol/migrator/migrator.go index b3f2c68fa7e..d4c0572cabc 100644 --- a/pkg/services/accesscontrol/migrator/migrator.go +++ b/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 - err := db.WithTransactionalDbSession(context.Background(), func(sess *sqlstore.DBSession) error { - var permissions []accesscontrol.Permission + ctx := context.Background() + cnt := 0 - err := sess.SQL("SELECT * FROM permission WHERE NOT scope = '' AND identifier = ''").Find(&permissions) - if err != nil { + // Search for the permissions to update + 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 { - count++ - kind, attribute, identifier := p.SplitScope() + i = end + } - permissions[i].Kind = kind - 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 + return nil } diff --git a/pkg/services/accesscontrol/migrator/migrator_bench_test.go b/pkg/services/accesscontrol/migrator/migrator_bench_test.go new file mode 100644 index 00000000000..5257a2fb2dd --- /dev/null +++ b/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) } diff --git a/pkg/services/accesscontrol/migrator/migrator_test.go b/pkg/services/accesscontrol/migrator/migrator_test.go new file mode 100644 index 00000000000..d09962e7709 --- /dev/null +++ b/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) + } +}