[Backport 5.0] [fix] set unrestricted not working correctly (#49489)

## Description

Two problems fixed here:
1. When using unified perms table, the SQL call was removing all rows
where repo_id = %d. This is not correct, as we might have other rows in
the DB with the same repo_id
2. When sending more items than parameter limit, the query would fail.
Fixed by doing requests in chunks.

## Test plan

unit tests added/modified
 <br> Backport 1e088ff from #49421

Co-authored-by: Milan Freml <kopancek@users.noreply.github.com>

Author: github-actions[bot]

enterprise/internal/database/perms_store.go

@@ -888,9 +888,21 @@ DO UPDATE SET
    unrestricted = %s;
 `
 
-	q := sqlf.Sprintf(format, pq.Array(ids), unrestricted, unrestricted)
+	size := 65535 / 2 // 65535 is the max number of parameters in a query, 2 is the number of parameters in each row
+	chunks, err := collections.SplitIntoChunks(ids, size)
+	if err != nil {
+		return err
+	}
 
-	return errors.Wrap(s.Exec(ctx, q), "setting unrestricted flag")
+	for _, chunk := range chunks {
+		q := sqlf.Sprintf(format, pq.Array(chunk), unrestricted, unrestricted)
+		err := s.Exec(ctx, q)
+		if err != nil {
+			return errors.Wrap(err, "setting unrestricted flag")
+		}
+	}
+
+	return nil
 }
 
 func (s *permsStore) SetRepoPermissionsUnrestricted(ctx context.Context, ids []int32, unrestricted bool) error {
@@ -915,23 +927,56 @@ func (s *permsStore) SetRepoPermissionsUnrestricted(ctx context.Context, ids []i
 		return err
 	}
 
+	if !unrestricted {
+		return txs.unsetRepoPermissionsUnrestricted(ctx, ids)
+	}
+	return txs.setRepoPermissionsUnrestricted(ctx, ids)
+}
+
+func (s *permsStore) unsetRepoPermissionsUnrestricted(ctx context.Context, ids []int32) error {
+	format := `DELETE FROM user_repo_permissions WHERE repo_id = ANY(%s) AND user_id IS NULL;`
+	size := 65535 - 1 // for unsetting unrestricted, we have only 1 parameter per row
+	chunks, err := collections.SplitIntoChunks(ids, size)
+	if err != nil {
+		return err
+	}
+	for _, chunk := range chunks {
+		err := s.Exec(ctx, sqlf.Sprintf(format, pq.Array(chunk)))
+		if err != nil {
+			return errors.Wrap(err, "removing unrestricted flag")
+		}
+	}
+
+	return nil
+}
+
+func (s *permsStore) setRepoPermissionsUnrestricted(ctx context.Context, ids []int32) error {
 	currentTime := time.Now()
 	values := make([]*sqlf.Query, 0, len(ids))
 	for _, repoID := range ids {
 		values = append(values, sqlf.Sprintf("(NULL, %d, %s, %s, %s)", repoID, currentTime, currentTime, authz.SourceAPI))
 	}
 
-	q := sqlf.Sprintf(`
+	format := `
 INSERT INTO user_repo_permissions (user_id, repo_id, created_at, updated_at, source)
 VALUES %s
-ON CONFLICT DO NOTHING`,
-		sqlf.Join(values, ","),
-	)
-	if !unrestricted {
-		q = sqlf.Sprintf(`DELETE FROM user_repo_permissions WHERE repo_id = ANY(%s)`, pq.Array(ids))
+ON CONFLICT DO NOTHING;
+`
+
+	size := 65535 / 4 // 65535 is the max number of parameters in a query, 4 is the number of parameters in each row
+	chunks, err := collections.SplitIntoChunks(values, size)
+	if err != nil {
+		return err
+	}
+
+	for _, chunk := range chunks {
+		err = s.Exec(ctx, sqlf.Sprintf(format, sqlf.Join(chunk, ",")))
+		if err != nil {
+			errors.Wrapf(err, "setting repositories as unrestricted %v", chunk)
+		}
 	}
 
-	return errors.Wrapf(txs.Exec(ctx, q), "setting repositories as unrestricted %v %v", ids, unrestricted)
+	return nil
 }
 
 // upsertRepoPermissionsQuery upserts single row of repository permissions.

enterprise/internal/database/perms_store_test.go

@@ -26,6 +26,7 @@ import (
 	"github.com/sourcegraph/sourcegraph/internal/actor"
 	"github.com/sourcegraph/sourcegraph/internal/api"
 	"github.com/sourcegraph/sourcegraph/internal/authz"
+	"github.com/sourcegraph/sourcegraph/internal/collections"
 	"github.com/sourcegraph/sourcegraph/internal/conf"
 	"github.com/sourcegraph/sourcegraph/internal/database"
 	"github.com/sourcegraph/sourcegraph/internal/database/basestore"
@@ -786,17 +787,23 @@ func setupPermsRelatedEntities(t *testing.T, s *permsStore, permissions []authz.
 	}
 
 	defaultErrMessage := "setup test related entities before actual test"
-	usersQuery := sqlf.Sprintf(`INSERT INTO users(id, username) VALUES %s ON CONFLICT (id) DO NOTHING`, sqlf.Join(maps.Values(users), ","))
-	if err := s.execute(context.Background(), usersQuery); err != nil {
-		t.Fatal(defaultErrMessage, err)
+	if len(users) > 0 {
+		usersQuery := sqlf.Sprintf(`INSERT INTO users(id, username) VALUES %s ON CONFLICT (id) DO NOTHING`, sqlf.Join(maps.Values(users), ","))
+		if err := s.execute(context.Background(), usersQuery); err != nil {
+			t.Fatal(defaultErrMessage, err)
+		}
 	}
-	externalAccountsQuery := sqlf.Sprintf(`INSERT INTO user_external_accounts(id, user_id, service_type, service_id, account_id, client_id) VALUES %s ON CONFLICT(id) DO NOTHING`, sqlf.Join(maps.Values(externalAccounts), ","))
-	if err := s.execute(context.Background(), externalAccountsQuery); err != nil {
-		t.Fatal(defaultErrMessage, err)
+	if len(externalAccounts) > 0 {
+		externalAccountsQuery := sqlf.Sprintf(`INSERT INTO user_external_accounts(id, user_id, service_type, service_id, account_id, client_id) VALUES %s ON CONFLICT(id) DO NOTHING`, sqlf.Join(maps.Values(externalAccounts), ","))
+		if err := s.execute(context.Background(), externalAccountsQuery); err != nil {
+			t.Fatal(defaultErrMessage, err)
+		}
 	}
-	reposQuery := sqlf.Sprintf(`INSERT INTO repo(id, name) VALUES %s ON CONFLICT(id) DO NOTHING`, sqlf.Join(maps.Values(repos), ","))
-	if err := s.execute(context.Background(), reposQuery); err != nil {
-		t.Fatal(defaultErrMessage, err)
+	if len(repos) > 0 {
+		reposQuery := sqlf.Sprintf(`INSERT INTO repo(id, name) VALUES %s ON CONFLICT(id) DO NOTHING`, sqlf.Join(maps.Values(repos), ","))
+		if err := s.execute(context.Background(), reposQuery); err != nil {
+			t.Fatal(defaultErrMessage, err)
+		}
 	}
 }
 
@@ -1197,14 +1204,18 @@ func TestPermsStore_SetRepoPermissionsUnrestricted(t *testing.T) {
 
 		unrestricted := (len(p) == 1 && p[0].UserID == 0)
 
+		fmt.Printf("P: %v %v\n", p, unrestricted)
+
 		if unrestricted != want {
 			t.Fatalf("Want %v, got %v for %d", want, unrestricted, id)
 		}
 	}
 
 	assertUnrestricted := func(t *testing.T, id int32, want bool) {
 		t.Helper()
+		fmt.Printf("before legacyUnrestricted\n")
 		legacyUnrestricted(t, id, want)
+		fmt.Printf("after legacyUnrestricted\n")
 
 		type unrestrictedResult struct {
 			id     int32
@@ -1217,8 +1228,11 @@ func TestPermsStore_SetRepoPermissionsUnrestricted(t *testing.T) {
 			return r, err
 		})
 
+		fmt.Printf("before scanResults\n")
 		q := sqlf.Sprintf("SELECT repo_id, source FROM user_repo_permissions WHERE repo_id = %d AND user_id IS NULL", id)
 		results, err := scanResults(s.Handle().QueryContext(ctx, q.Query(sqlf.PostgresBindVar), q.Args()...))
+
+		fmt.Printf("after scanResults\n")
 		if err != nil {
 			t.Fatalf("loading user repo permissions for %d: %v", id, err)
 		}
@@ -1229,6 +1243,7 @@ func TestPermsStore_SetRepoPermissionsUnrestricted(t *testing.T) {
 			t.Fatalf("Want restricted, but found results for %d: %v", id, results)
 		}
 
+		fmt.Printf("Results: %v\n", results)
 		if want {
 			for _, r := range results {
 				require.Equal(t, authz.SourceAPI, r.source)
@@ -1243,30 +1258,46 @@ func TestPermsStore_SetRepoPermissionsUnrestricted(t *testing.T) {
 		VALUES (%d, %s, TRUE)`, id, fmt.Sprintf("repo-%d", id)))
 	}
 
-	// Add a couple of repos and a user
-	execQuery(t, ctx, s, sqlf.Sprintf(`INSERT INTO users (username) VALUES ('alice')`))
-	execQuery(t, ctx, s, sqlf.Sprintf(`INSERT INTO users (username) VALUES ('bob')`))
-	for i := 0; i < 2; i++ {
-		createRepo(t, i+1)
-		rp := &authz.RepoPermissions{
-			RepoID:  int32(i + 1),
-			Perm:    authz.Read,
-			UserIDs: toMapset(2),
-		}
-		if _, err := s.SetRepoPermissions(context.Background(), rp); err != nil {
-			t.Fatal(err)
-		}
-		if _, err := s.SetRepoPerms(context.Background(), int32(i+1), []authz.UserIDWithExternalAccountID{{UserID: 2}}, authz.SourceRepoSync); err != nil {
-			t.Fatal(err)
+	setupData := func() {
+		// Add a couple of repos and a user
+		execQuery(t, ctx, s, sqlf.Sprintf(`INSERT INTO users (username) VALUES ('alice')`))
+		execQuery(t, ctx, s, sqlf.Sprintf(`INSERT INTO users (username) VALUES ('bob')`))
+		for i := 0; i < 2; i++ {
+			createRepo(t, i+1)
+			rp := &authz.RepoPermissions{
+				RepoID:  int32(i + 1),
+				Perm:    authz.Read,
+				UserIDs: toMapset(2),
+			}
+			if _, err := s.SetRepoPermissions(context.Background(), rp); err != nil {
+				t.Fatal(err)
+			}
+			if _, err := s.SetRepoPerms(context.Background(), int32(i+1), []authz.UserIDWithExternalAccountID{{UserID: 2}}, authz.SourceRepoSync); err != nil {
+				t.Fatal(err)
+			}
 		}
 	}
 
+	cleanupTables := func() {
+		t.Helper()
+
+		cleanupPermsTables(t, s)
+		cleanupReposTable(t, s)
+		cleanupUsersTable(t, s)
+	}
+
 	t.Run("Both repos are restricted by default", func(t *testing.T) {
+		t.Cleanup(cleanupTables)
+		setupData()
+
 		assertUnrestricted(t, 1, false)
 		assertUnrestricted(t, 2, false)
 	})
 
 	t.Run("Set both repos to unrestricted", func(t *testing.T) {
+		t.Cleanup(cleanupTables)
+		setupData()
+
 		if err := s.SetRepoPermissionsUnrestricted(ctx, []int32{1, 2}, true); err != nil {
 			t.Fatal(err)
 		}
@@ -1275,39 +1306,63 @@ func TestPermsStore_SetRepoPermissionsUnrestricted(t *testing.T) {
 	})
 
 	t.Run("Set unrestricted on a repo not in permissions table", func(t *testing.T) {
+		t.Cleanup(cleanupTables)
+		setupData()
+
 		createRepo(t, 3)
-		if err := s.SetRepoPermissionsUnrestricted(ctx, []int32{3}, true); err != nil {
-			t.Fatal(err)
-		}
+		err := s.SetRepoPermissionsUnrestricted(ctx, []int32{1, 2, 3}, true)
+		require.NoError(t, err)
 
 		assertUnrestricted(t, 1, true)
 		assertUnrestricted(t, 2, true)
 		assertUnrestricted(t, 3, true)
 	})
 
 	t.Run("Unset restricted on a repo in and not in permissions table", func(t *testing.T) {
+		t.Cleanup(cleanupTables)
+		setupData()
+
+		createRepo(t, 3)
 		createRepo(t, 4)
-		if err := s.SetRepoPermissionsUnrestricted(ctx, []int32{2, 3, 4}, false); err != nil {
-			t.Fatal(err)
-		}
+
+		// set permissions on repo 4
+		_, err := s.SetRepoPerms(ctx, 4, []authz.UserIDWithExternalAccountID{{UserID: 2}}, authz.SourceRepoSync)
+		require.NoError(t, err)
+		err = s.SetRepoPermissionsUnrestricted(ctx, []int32{1, 2, 3, 4}, true)
+		require.NoError(t, err)
+		err = s.SetRepoPermissionsUnrestricted(ctx, []int32{2, 3, 4}, false)
+		require.NoError(t, err)
+
 		assertUnrestricted(t, 1, true)
 		assertUnrestricted(t, 2, false)
 		assertUnrestricted(t, 3, false)
 		assertUnrestricted(t, 4, false)
+		checkUserRepoPermissions(t, s, sqlf.Sprintf("repo_id = 4"), []authz.Permission{{UserID: 2, RepoID: 4, Source: authz.SourceRepoSync}})
 	})
 
-	t.Run("Set repos back to restricted again", func(t *testing.T) {
+	t.Run("Check parameter limit", func(t *testing.T) {
+		t.Cleanup(cleanupTables)
+
 		// Also checking that more than 65535 IDs can be processed without an error
 		var ids [66000]int32
+		p := make([]authz.Permission, len(ids))
 		for i := range ids {
 			ids[i] = int32(i + 1)
+			p[i] = authz.Permission{RepoID: ids[i], Source: authz.SourceAPI}
 		}
-		if err := s.SetRepoPermissionsUnrestricted(ctx, ids[:], false); err != nil {
+
+		chunks, err := collections.SplitIntoChunks(p, 15000)
+		require.NoError(t, err)
+
+		for _, chunk := range chunks {
+			setupPermsRelatedEntities(t, s, chunk)
+		}
+		if err := s.SetRepoPermissionsUnrestricted(ctx, ids[:], true); err != nil {
 			t.Fatal(err)
 		}
-		assertUnrestricted(t, 1, false)
-		assertUnrestricted(t, 500, false)
-		assertUnrestricted(t, 66000, false)
+		assertUnrestricted(t, 1, true)
+		assertUnrestricted(t, 500, true)
+		assertUnrestricted(t, 66000, true)
 	})
 }