fix(permissions): chunked synchronization of role permission events (#9403)

# Which Problems Are Solved

Setup fails to push all role permission events when running Zitadel with
CockroachDB. `TransactionRetryError`s were visible in logs which finally
times out the setup job with `timeout: context deadline exceeded`

# How the Problems Are Solved

As suggested in the [Cockroach documentation](timeout: context deadline
exceeded), _"break down larger transactions"_. The commands to be pushed
for the role permissions are chunked in 50 events per push. This
chunking is only done with CockroachDB.

# Additional Changes

- gci run fixed some unrelated imports
- access to `command.Commands` for the setup job, so we can reuse the
sync logic.

# Additional Context

Closes #9293

---------

Co-authored-by: Silvan <27845747+adlerhurst@users.noreply.github.com>

internal/command/command.go

@@ -218,6 +218,33 @@ func (c *Commands) pushAppendAndReduce(ctx context.Context, object AppendReducer
 	return AppendAndReduce(object, events...)
 }
 
+// pushChunked pushes the commands in chunks of size to the eventstore.
+// This can be used to reduce the amount of events in a single transaction.
+// When an error occurs, the events that have been pushed so far will be returned.
+//
+// Warning: chunks are pushed in separate transactions.
+// Successful pushes will not be rolled back if a later chunk fails.
+// Only use this function when the caller is able to handle partial success
+// and is able to consolidate the state on errors.
+func (c *Commands) pushChunked(ctx context.Context, size uint16, cmds ...eventstore.Command) (_ []eventstore.Event, err error) {
+	ctx, span := tracing.NewSpan(ctx)
+	defer func() { span.EndWithError(err) }()
+
+	events := make([]eventstore.Event, 0, len(cmds))
+	for i := 0; i < len(cmds); i += int(size) {
+		end := i + int(size)
+		if end > len(cmds) {
+			end = len(cmds)
+		}
+		chunk, err := c.eventstore.Push(ctx, cmds[i:end]...)
+		if err != nil {
+			return events, err
+		}
+		events = append(events, chunk...)
+	}
+	return events, nil
+}
+
 type AppendReducerDetails interface {
 	AppendEvents(...eventstore.Event)
 	// TODO: Why is it allowed to return an error here?

internal/command/command_test.go

@@ -2,6 +2,7 @@ package command
 
 import (
 	"context"
+	"fmt"
 	"io"
 	"os"
 	"testing"
@@ -13,6 +14,7 @@ import (
 
 	"github.com/zitadel/zitadel/internal/eventstore"
 	"github.com/zitadel/zitadel/internal/i18n"
+	"github.com/zitadel/zitadel/internal/repository/permission"
 	"github.com/zitadel/zitadel/internal/repository/user"
 )
 
@@ -29,6 +31,93 @@ func TestMain(m *testing.M) {
 	os.Exit(m.Run())
 }
 
+func TestCommands_pushChunked(t *testing.T) {
+	aggregate := permission.NewAggregate("instanceID")
+	cmds := make([]eventstore.Command, 100)
+	for i := 0; i < 100; i++ {
+		cmds[i] = permission.NewAddedEvent(context.Background(), aggregate, "role", fmt.Sprintf("permission%d", i))
+	}
+	type args struct {
+		size uint16
+	}
+	tests := []struct {
+		name       string
+		args       args
+		eventstore func(*testing.T) *eventstore.Eventstore
+		wantEvents int
+		wantErr    error
+	}{
+		{
+			name: "push error",
+			args: args{
+				size: 100,
+			},
+			eventstore: expectEventstore(
+				expectPushFailed(io.ErrClosedPipe, cmds...),
+			),
+			wantEvents: 0,
+			wantErr:    io.ErrClosedPipe,
+		},
+		{
+			name: "single chunk",
+			args: args{
+				size: 100,
+			},
+			eventstore: expectEventstore(
+				expectPush(cmds...),
+			),
+			wantEvents: len(cmds),
+		},
+		{
+			name: "aligned chunks",
+			args: args{
+				size: 50,
+			},
+			eventstore: expectEventstore(
+				expectPush(cmds[0:50]...),
+				expectPush(cmds[50:100]...),
+			),
+			wantEvents: len(cmds),
+		},
+		{
+			name: "odd chunks",
+			args: args{
+				size: 30,
+			},
+			eventstore: expectEventstore(
+				expectPush(cmds[0:30]...),
+				expectPush(cmds[30:60]...),
+				expectPush(cmds[60:90]...),
+				expectPush(cmds[90:100]...),
+			),
+			wantEvents: len(cmds),
+		},
+		{
+			name: "partial error",
+			args: args{
+				size: 30,
+			},
+			eventstore: expectEventstore(
+				expectPush(cmds[0:30]...),
+				expectPush(cmds[30:60]...),
+				expectPushFailed(io.ErrClosedPipe, cmds[60:90]...),
+			),
+			wantEvents: len(cmds[0:60]),
+			wantErr:    io.ErrClosedPipe,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			c := &Commands{
+				eventstore: tt.eventstore(t),
+			}
+			gotEvents, err := c.pushChunked(context.Background(), tt.args.size, cmds...)
+			require.ErrorIs(t, err, tt.wantErr)
+			assert.Len(t, gotEvents, tt.wantEvents)
+		})
+	}
+}
+
 func TestCommands_asyncPush(t *testing.T) {
 	// make sure the test terminates on deadlock
 	background := context.Background()

internal/command/converter.go

@@ -15,6 +15,9 @@ func writeModelToObjectDetails(writeModel *eventstore.WriteModel) *domain.Object
 }
 
 func pushedEventsToObjectDetails(events []eventstore.Event) *domain.ObjectDetails {
+	if len(events) == 0 {
+		return &domain.ObjectDetails{}
+	}
 	return &domain.ObjectDetails{
 		Sequence:      events[len(events)-1].Sequence(),
 		EventDate:     events[len(events)-1].CreatedAt(),

internal/command/instance.go

@@ -233,21 +233,25 @@ func (c *Commands) SetUpInstance(ctx context.Context, setup *InstanceSetup) (str
 		return "", "", nil, nil, err
 	}
 
-	events, err := c.eventstore.Push(ctx, cmds...)
+	_, err = c.eventstore.Push(ctx, cmds...)
 	if err != nil {
 		return "", "", nil, nil, err
 	}
 
+	// RolePermissions need to be pushed in separate transaction.
+	// https://github.com/zitadel/zitadel/issues/9293
+	details, err := c.SynchronizeRolePermission(ctx, setup.zitadel.instanceID, setup.RolePermissionMappings)
+	if err != nil {
+		return "", "", nil, nil, err
+	}
+	details.ResourceOwner = setup.zitadel.orgID
+
 	var token string
 	if pat != nil {
 		token = pat.Token
 	}
 
-	return setup.zitadel.instanceID, token, machineKey, &domain.ObjectDetails{
-		Sequence:      events[len(events)-1].Sequence(),
-		EventDate:     events[len(events)-1].CreatedAt(),
-		ResourceOwner: setup.zitadel.orgID,
-	}, nil
+	return setup.zitadel.instanceID, token, machineKey, details, nil
 }
 
 func contextWithInstanceSetupInfo(ctx context.Context, instanceID, projectID, consoleAppID, externalDomain string) context.Context {
@@ -380,7 +384,6 @@ func setupInstanceElements(instanceAgg *instance.Aggregate, setup *InstanceSetup
 			setup.LabelPolicy.ThemeMode,
 		),
 		prepareAddDefaultEmailTemplate(instanceAgg, setup.EmailTemplate),
-		prepareAddRolePermissions(instanceAgg, setup.RolePermissionMappings),
 	}
 }
 

internal/command/instance_permissions.go

@@ -1,29 +0,0 @@
-package command
-
-import (
-	"context"
-	"strings"
-
-	"github.com/zitadel/zitadel/internal/api/authz"
-	"github.com/zitadel/zitadel/internal/command/preparation"
-	"github.com/zitadel/zitadel/internal/eventstore"
-	"github.com/zitadel/zitadel/internal/repository/instance"
-	"github.com/zitadel/zitadel/internal/repository/permission"
-)
-
-func prepareAddRolePermissions(a *instance.Aggregate, roles []authz.RoleMapping) preparation.Validation {
-	return func() (preparation.CreateCommands, error) {
-		return func(ctx context.Context, _ preparation.FilterToQueryReducer) (cmds []eventstore.Command, _ error) {
-			aggregate := permission.NewAggregate(a.InstanceID)
-			for _, r := range roles {
-				if strings.HasPrefix(r.Role, "SYSTEM") {
-					continue
-				}
-				for _, p := range r.Permissions {
-					cmds = append(cmds, permission.NewAddedEvent(ctx, aggregate, r.Role, p))
-				}
-			}
-			return cmds, nil
-		}, nil
-	}
-}

internal/command/instance_role_permissions.go

@@ -0,0 +1,101 @@
+package command
+
+import (
+	"context"
+	"database/sql"
+	_ "embed"
+	"strings"
+
+	"github.com/zitadel/logging"
+
+	"github.com/zitadel/zitadel/internal/api/authz"
+	"github.com/zitadel/zitadel/internal/database"
+	"github.com/zitadel/zitadel/internal/domain"
+	"github.com/zitadel/zitadel/internal/eventstore"
+	"github.com/zitadel/zitadel/internal/repository/permission"
+	"github.com/zitadel/zitadel/internal/telemetry/tracing"
+	"github.com/zitadel/zitadel/internal/zerrors"
+)
+
+const (
+	CockroachRollPermissionChunkSize uint16 = 50
+)
+
+// SynchronizeRolePermission checks the current state of role permissions in the eventstore for the aggregate.
+// It pushes the commands required to reach the desired state passed in target.
+// For system level permissions aggregateID must be set to `SYSTEM`, else it is the instance ID.
+//
+// In case cockroachDB is used, the commands are pushed in chunks of CockroachRollPermissionChunkSize.
+func (c *Commands) SynchronizeRolePermission(ctx context.Context, aggregateID string, target []authz.RoleMapping) (_ *domain.ObjectDetails, err error) {
+	ctx, span := tracing.NewSpan(ctx)
+	defer func() { span.EndWithError(err) }()
+
+	cmds, err := synchronizeRolePermissionCommands(ctx, c.eventstore.Client(), aggregateID,
+		rolePermissionMappingsToDatabaseMap(target, aggregateID == "SYSTEM"),
+	)
+	if err != nil {
+		return nil, zerrors.ThrowInternal(err, "COMMA-Iej2r", "Errors.Internal")
+	}
+	var events []eventstore.Event
+	if c.eventstore.Client().Database.Type() == "cockroach" {
+		events, err = c.pushChunked(ctx, CockroachRollPermissionChunkSize, cmds...)
+	} else {
+		events, err = c.eventstore.Push(ctx, cmds...)
+	}
+	if err != nil {
+		return nil, zerrors.ThrowInternal(err, "COMMA-AiV3u", "Errors.Internal")
+	}
+	return pushedEventsToObjectDetails(events), nil
+}
+
+func rolePermissionMappingsToDatabaseMap(mappings []authz.RoleMapping, system bool) database.Map[[]string] {
+	out := make(database.Map[[]string], len(mappings))
+	for _, m := range mappings {
+		if system == strings.HasPrefix(m.Role, "SYSTEM") {
+			out[m.Role] = m.Permissions
+		}
+	}
+	return out
+}
+
+var (
+	//go:embed instance_role_permissions_sync.sql
+	instanceRolePermissionsSyncQuery string
+)
+
+// synchronizeRolePermissionCommands checks the current state of role permissions in the eventstore for the aggregate.
+// It returns the commands required to reach the desired state passed in target.
+// For system level permissions aggregateID must be set to `SYSTEM`, else it is the instance ID.
+func synchronizeRolePermissionCommands(ctx context.Context, db *database.DB, aggregateID string, target database.Map[[]string]) (cmds []eventstore.Command, err error) {
+	ctx, span := tracing.NewSpan(ctx)
+	defer func() { span.EndWithError(err) }()
+
+	err = db.QueryContext(ctx,
+		rolePermissionScanner(ctx, permission.NewAggregate(aggregateID), &cmds),
+		instanceRolePermissionsSyncQuery,
+		aggregateID, target)
+	if err != nil {
+		return nil, err
+	}
+	return cmds, nil
+}
+
+func rolePermissionScanner(ctx context.Context, aggregate *eventstore.Aggregate, cmds *[]eventstore.Command) func(rows *sql.Rows) error {
+	return func(rows *sql.Rows) error {
+		for rows.Next() {
+			var operation, role, perm string
+			if err := rows.Scan(&operation, &role, &perm); err != nil {
+				return err
+			}
+			logging.WithFields("aggregate_id", aggregate.ID, "operation", operation, "role", role, "permission", perm).Debug("sync role permission")
+			switch operation {
+			case "add":
+				*cmds = append(*cmds, permission.NewAddedEvent(ctx, aggregate, role, perm))
+			case "remove":
+				*cmds = append(*cmds, permission.NewRemovedEvent(ctx, aggregate, role, perm))
+			}
+		}
+		return rows.Close()
+
+	}
+}

internal/command/instance_role_permissions_sync.sql

@@ -0,0 +1,52 @@
+/*
+This query creates a change set of permissions that need to be added or removed.
+It compares the current state in the fields table (thru the role_permissions view)
+against a passed role permission mapping as JSON, created from Zitadel's config:
+
+{
+	"IAM_ADMIN_IMPERSONATOR": ["admin.impersonation", "impersonation"],
+	"IAM_END_USER_IMPERSONATOR": ["impersonation"],
+	"FOO_BAR": ["foo.bar", "bar.foo"]
+ }
+
+It uses an aggregate_id as first argument which may be an instance_id or 'SYSTEM'
+for system level permissions.
+*/
+WITH target AS (
+	-- unmarshal JSON representation into flattened tabular data
+	SELECT
+		key AS role,
+		jsonb_array_elements_text(value) AS permission
+	FROM jsonb_each($2::jsonb)
+), add AS (
+    -- find all role permissions that exist in `target` and not in `role_permissions`
+	SELECT t.role, t.permission
+	FROM eventstore.role_permissions p
+	RIGHT JOIN target t
+		ON p.aggregate_id = $1::text
+		AND p.role = t.role
+		AND p.permission = t.permission
+	WHERE p.role IS NULL
+), remove AS (
+    -- find all role permissions that exist `role_permissions` and not in `target`
+	SELECT p.role, p.permission
+	FROM eventstore.role_permissions p
+	LEFT JOIN target t
+		ON p.role = t.role
+		AND p.permission = t.permission
+	WHERE p.aggregate_id = $1::text
+	AND t.role IS NULL
+)
+-- return the required operations
+SELECT
+	'add' AS operation,
+	role,
+	permission
+FROM add
+UNION ALL
+SELECT
+	'remove' AS operation,
+	role,
+	permission
+FROM remove
+;

internal/command/instance_role_permissions_test.go

@@ -0,0 +1,139 @@
+package command
+
+import (
+	"context"
+	"database/sql"
+	"database/sql/driver"
+	_ "embed"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/zitadel/zitadel/internal/api/authz"
+	"github.com/zitadel/zitadel/internal/database"
+	"github.com/zitadel/zitadel/internal/database/mock"
+	"github.com/zitadel/zitadel/internal/eventstore"
+	"github.com/zitadel/zitadel/internal/repository/permission"
+)
+
+func Test_rolePermissionMappingsToDatabaseMap(t *testing.T) {
+	type args struct {
+		mappings []authz.RoleMapping
+		system   bool
+	}
+	tests := []struct {
+		name string
+		args args
+		want database.Map[[]string]
+	}{
+		{
+			name: "instance",
+			args: args{
+				mappings: []authz.RoleMapping{
+					{Role: "role1", Permissions: []string{"permission1", "permission2"}},
+					{Role: "role2", Permissions: []string{"permission3", "permission4"}},
+					{Role: "SYSTEM_ROLE", Permissions: []string{"permission5", "permission6"}},
+				},
+				system: false,
+			},
+			want: database.Map[[]string]{
+				"role1": []string{"permission1", "permission2"},
+				"role2": []string{"permission3", "permission4"},
+			},
+		},
+		{
+			name: "system",
+			args: args{
+				mappings: []authz.RoleMapping{
+					{Role: "role1", Permissions: []string{"permission1", "permission2"}},
+					{Role: "role2", Permissions: []string{"permission3", "permission4"}},
+					{Role: "SYSTEM_ROLE", Permissions: []string{"permission5", "permission6"}},
+				},
+				system: true,
+			},
+			want: database.Map[[]string]{
+				"SYSTEM_ROLE": []string{"permission5", "permission6"},
+			},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got := rolePermissionMappingsToDatabaseMap(tt.args.mappings, tt.args.system)
+			assert.Equal(t, tt.want, got)
+		})
+	}
+}
+
+func Test_synchronizeRolePermissionCommands(t *testing.T) {
+	const aggregateID = "aggregateID"
+	aggregate := permission.NewAggregate(aggregateID)
+	target := database.Map[[]string]{
+		"role1": []string{"permission1", "permission2"},
+		"role2": []string{"permission3", "permission4"},
+	}
+	tests := []struct {
+		name     string
+		mock     func(*testing.T) *mock.SQLMock
+		wantCmds []eventstore.Command
+		wantErr  error
+	}{
+		{
+			name: "query error",
+			mock: func(t *testing.T) *mock.SQLMock {
+				return mock.NewSQLMock(t,
+					mock.ExpectQuery(instanceRolePermissionsSyncQuery,
+						mock.WithQueryArgs(aggregateID, target),
+						mock.WithQueryErr(sql.ErrConnDone),
+					),
+				)
+			},
+			wantErr: sql.ErrConnDone,
+		},
+		{
+			name: "no rows",
+			mock: func(t *testing.T) *mock.SQLMock {
+				return mock.NewSQLMock(t,
+					mock.ExpectQuery(instanceRolePermissionsSyncQuery,
+						mock.WithQueryArgs(aggregateID, target),
+						mock.WithQueryResult([]string{"operation", "role", "permission"}, [][]driver.Value{}),
+					),
+				)
+			},
+		},
+		{
+			name: "add and remove operations",
+			mock: func(t *testing.T) *mock.SQLMock {
+				return mock.NewSQLMock(t,
+					mock.ExpectQuery(instanceRolePermissionsSyncQuery,
+						mock.WithQueryArgs(aggregateID, target),
+						mock.WithQueryResult([]string{"operation", "role", "permission"}, [][]driver.Value{
+							{"add", "role1", "permission1"},
+							{"add", "role1", "permission2"},
+							{"remove", "role3", "permission5"},
+							{"remove", "role3", "permission6"},
+						}),
+					),
+				)
+			},
+			wantCmds: []eventstore.Command{
+				permission.NewAddedEvent(context.Background(), aggregate, "role1", "permission1"),
+				permission.NewAddedEvent(context.Background(), aggregate, "role1", "permission2"),
+				permission.NewRemovedEvent(context.Background(), aggregate, "role3", "permission5"),
+				permission.NewRemovedEvent(context.Background(), aggregate, "role3", "permission6"),
+			},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			mock := tt.mock(t)
+			defer mock.Assert(t)
+			db := &database.DB{
+				DB: mock.DB,
+			}
+			gotCmds, err := synchronizeRolePermissionCommands(context.Background(), db, aggregateID, target)
+			require.ErrorIs(t, err, tt.wantErr)
+			assert.Equal(t, tt.wantCmds, gotCmds)
+		})
+	}
+}