Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

upgrades: fix backfilling of special roles in user ID migrations #99097

Merged
merged 2 commits into from
Mar 21, 2023
Merged

upgrades: fix backfilling of special roles in user ID migrations #99097

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
0440d30
upgrades: fix system.privileges user ID migration
andyyang890 Mar 21, 2023
de67a4d
upgrades: fix system.database_role_settings user ID migration
andyyang890 Mar 21, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 13 additions & 6 deletions pkg/upgrade/upgrades/database_role_settings_table_user_id_migration.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,9 +12,11 @@ package upgrades

import (
"context"
"fmt"

"github.com/cockroachdb/cockroach/pkg/clusterversion"
"github.com/cockroachdb/cockroach/pkg/keys"
"github.com/cockroachdb/cockroach/pkg/security/username"
"github.com/cockroachdb/cockroach/pkg/sql/catalog/systemschema"
"github.com/cockroachdb/cockroach/pkg/sql/sessiondata"
"github.com/cockroachdb/cockroach/pkg/upgrade"
Expand Down Expand Up @@ -57,13 +59,18 @@ func alterDatabaseRoleSettingsTableAddRoleIDColumn(
return nil
}

const backfillRoleIDColumnDatabaseRoleSettingsTableStmt = `
var backfillRoleIDColumnDatabaseRoleSettingsTableStmt = fmt.Sprintf(`
UPDATE system.database_role_settings
SET role_id = user_id
FROM system.users
WHERE role_id IS NULL AND role_name = username
SET role_id = (
SELECT CASE role_name
WHEN '%s' THEN %d
ELSE (SELECT user_id FROM system.users WHERE username = role_name)
END
)
WHERE role_id IS NULL
LIMIT 1000
`
`,
username.EmptyRole, username.EmptyRoleID)

const setRoleIDColumnToNotNullDatabaseRoleSettingsTableStmt = `
ALTER TABLE system.database_role_settings
Expand Down Expand Up @@ -92,7 +99,7 @@ func backfillDatabaseRoleSettingsTableRoleIDColumn(
// and any new rows inserted after the previous version (when the role_id column
// was added) will have had their role_id value populated at insertion time.
op := operation{
name: "set-user-id-not-null-database-role-settings-table",
name: "set-role-id-not-null-database-role-settings-table",
schemaList: []string{"role_id"},
query: setRoleIDColumnToNotNullDatabaseRoleSettingsTableStmt,
schemaExistsFn: columnExistsAndIsNotNull,
Expand Down
13 changes: 11 additions & 2 deletions pkg/upgrade/upgrades/database_role_settings_table_user_id_migration_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@ import (
"github.com/cockroachdb/cockroach/pkg/base"
"github.com/cockroachdb/cockroach/pkg/clusterversion"
"github.com/cockroachdb/cockroach/pkg/keys"
"github.com/cockroachdb/cockroach/pkg/security/username"
"github.com/cockroachdb/cockroach/pkg/server"
"github.com/cockroachdb/cockroach/pkg/settings/cluster"
"github.com/cockroachdb/cockroach/pkg/sql/catalog/catenumpb"
Expand Down Expand Up @@ -79,13 +80,17 @@ func runTestDatabaseRoleSettingsUserIDMigration(t *testing.T, numUsers int) {
// before the role_id column was added.
upgrades.InjectLegacyTable(ctx, t, s, systemschema.DatabaseRoleSettingsTable, getTableDescForDatabaseRoleSettingsTableBeforeRoleIDCol)

// Create test users.
// Create test users and add rows for each user to system.database_role_settings.
upgrades.ExecForCountInTxns(ctx, t, db, numUsers, 100 /* txCount */, func(txRunner *sqlutils.SQLRunner, i int) {
txRunner.Exec(t, fmt.Sprintf("CREATE USER testuser%d", i))
txRunner.Exec(t, fmt.Sprintf(`ALTER USER testuser%d SET application_name = 'roach sql'`, i))
})
tdb.CheckQueryResults(t, "SELECT count(*) FROM system.database_role_settings", [][]string{{strconv.Itoa(numUsers)}})

// Create a row in system.database_role_settings for the empty role.
tdb.Exec(t, "ALTER ROLE ALL SET timezone = 'America/New_York'")
tdb.CheckQueryResults(t, "SELECT count(*) FROM system.database_role_settings", [][]string{{strconv.Itoa(numUsers + 1)}})

// Run migrations.
_, err := tc.Conns[0].ExecContext(ctx, `SET CLUSTER SETTING version = $1`,
clusterversion.ByKey(clusterversion.V23_1DatabaseRoleSettingsHasRoleIDColumn).String())
Expand All @@ -110,8 +115,12 @@ func runTestDatabaseRoleSettingsUserIDMigration(t *testing.T, numUsers int) {

// Check that the backfill was successful and correct.
tdb.CheckQueryResults(t, "SELECT * FROM system.database_role_settings WHERE role_id IS NULL", [][]string{})
tdb.CheckQueryResults(t, "SELECT count(*) FROM system.database_role_settings", [][]string{{strconv.Itoa(numUsers)}})
tdb.CheckQueryResults(t, "SELECT count(*) FROM system.database_role_settings", [][]string{{strconv.Itoa(numUsers + 1)}})
tdb.CheckQueryResults(t, "SELECT count(*) FROM system.database_role_settings AS a JOIN system.users AS b ON a.role_name = b.username AND a.role_id <> b.user_id", [][]string{{"0"}})
tdb.CheckQueryResults(t,
fmt.Sprintf("SELECT count(*) FROM system.privileges WHERE username = '%s' AND user_id <> %d",
username.EmptyRole, username.EmptyRoleID),
[][]string{{"0"}})
}

func getTableDescForDatabaseRoleSettingsTableBeforeRoleIDCol() *descpb.TableDescriptor {
Expand Down
17 changes: 12 additions & 5 deletions pkg/upgrade/upgrades/system_privileges_user_id_migration.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,8 +12,10 @@ package upgrades

import (
"context"
"fmt"

"github.com/cockroachdb/cockroach/pkg/clusterversion"
"github.com/cockroachdb/cockroach/pkg/security/username"
"github.com/cockroachdb/cockroach/pkg/sql/catalog/descpb"
"github.com/cockroachdb/cockroach/pkg/sql/catalog/systemschema"
"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
Expand Down Expand Up @@ -69,13 +71,18 @@ func alterSystemPrivilegesAddUserIDColumn(
return nil
}

const backfillUserIDColumnSystemPrivilegesStmt = `
var backfillUserIDColumnSystemPrivilegesStmt = fmt.Sprintf(`
UPDATE system.privileges AS p
SET user_id = u.user_id
FROM system.users AS u
WHERE p.user_id is NULL AND p.username = u.username
SET user_id = (
SELECT CASE p.username
WHEN '%s' THEN %d
ELSE (SELECT user_id FROM system.users AS u WHERE u.username = p.username)
END
)
WHERE p.user_id IS NULL
LIMIT 1000
`
`,
username.PublicRole, username.PublicRoleID)

const setUserIDColumnToNotNullSystemPrivilegesStmt = `
ALTER TABLE system.privileges
Expand Down
12 changes: 10 additions & 2 deletions pkg/upgrade/upgrades/system_privileges_user_id_migration_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -83,13 +83,17 @@ func runTestSystemPrivilegesUserIDMigration(t *testing.T, numUsers int) {
// user_id column was added.
upgrades.InjectLegacyTable(ctx, t, s, systemschema.SystemPrivilegeTable, getTableDescForSystemPrivilegesTableBeforeUserIDCol)

// Create test users.
// Create test users and add rows for each user to system.privileges.
upgrades.ExecForCountInTxns(ctx, t, db, numUsers, 100 /* txCount */, func(txRunner *sqlutils.SQLRunner, i int) {
txRunner.Exec(t, fmt.Sprintf("CREATE USER testuser%d", i))
txRunner.Exec(t, fmt.Sprintf("GRANT SYSTEM MODIFYCLUSTERSETTING TO testuser%d", i))
})
tdb.CheckQueryResults(t, "SELECT count(*) FROM system.privileges", [][]string{{strconv.Itoa(numUsers)}})

// Create a row in system.privileges for the "public" role.
tdb.Exec(t, "REVOKE SELECT ON crdb_internal.tables FROM public")
tdb.CheckQueryResults(t, "SELECT count(*) FROM system.privileges", [][]string{{strconv.Itoa(numUsers + 1)}})

// Run migrations.
_, err := tc.Conns[0].ExecContext(ctx, `SET CLUSTER SETTING version = $1`,
clusterversion.ByKey(clusterversion.V23_1SystemPrivilegesTableHasUserIDColumn).String())
Expand All @@ -115,8 +119,12 @@ func runTestSystemPrivilegesUserIDMigration(t *testing.T, numUsers int) {

// Check that the backfill was successful and correct.
tdb.CheckQueryResults(t, "SELECT * FROM system.privileges WHERE user_id IS NULL", [][]string{})
tdb.CheckQueryResults(t, "SELECT count(*) FROM system.privileges", [][]string{{strconv.Itoa(numUsers)}})
tdb.CheckQueryResults(t, "SELECT count(*) FROM system.privileges", [][]string{{strconv.Itoa(numUsers + 1)}})
tdb.CheckQueryResults(t, "SELECT count(*) FROM system.privileges AS a JOIN system.users AS b ON a.username = b.username AND a.user_id <> b.user_id", [][]string{{"0"}})
tdb.CheckQueryResults(t,
fmt.Sprintf("SELECT count(*) FROM system.privileges WHERE username = '%s' AND user_id <> %d",
username.PublicRole, username.PublicRoleID),
[][]string{{"0"}})
}

func getTableDescForSystemPrivilegesTableBeforeUserIDCol() *descpb.TableDescriptor {
Expand Down