upgrades: fix idempotence issue in system table role id migrations

This patch fixes an idempotence issue in the role id migrations for the system.users and system.role_options tables. Release note: None
cockroachdb · Nov 15, 2022 · b8ca64f · b8ca64f
1 parent 03f6062
commit b8ca64f
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 32 deletions.
diff --git a/pkg/upgrade/upgrades/role_options_table_migration.go b/pkg/upgrade/upgrades/role_options_table_migration.go
@@ -13,12 +13,10 @@ package upgrades
 import (
 	"context"
 	"fmt"
-	"strings"
 
 	"github.com/cockroachdb/cockroach/pkg/clusterversion"
 	"github.com/cockroachdb/cockroach/pkg/jobs"
 	"github.com/cockroachdb/cockroach/pkg/keys"
-	"github.com/cockroachdb/cockroach/pkg/sql/catalog"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/systemschema"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
 	"github.com/cockroachdb/cockroach/pkg/sql/sessiondata"
@@ -46,7 +44,7 @@ func alterSystemRoleOptionsAddUserIDColumnWithIndex(
 			name:           "add-system-role-options-user-id-column",
 			schemaList:     []string{"user_id"},
 			query:          addUserIDColumnToRoleOptions,
-			schemaExistsFn: hasColumn,
+			schemaExistsFn: columnExists,
 		},
 		{
 			name:           "alter-system-role-options-add-index",
@@ -101,19 +99,10 @@ func setSystemRoleOptionsUserIDColumnNotNull(
 	ctx context.Context, cs clusterversion.ClusterVersion, d upgrade.TenantDeps, _ *jobs.Job,
 ) error {
 	op := operation{
-		name:       "alter-system-role-options-user-id-column-not-null",
-		schemaList: []string{"user_id"},
-		query:      updateUserIDColumnRoleOptionsSetNotNull,
-		schemaExistsFn: func(storedTable, _ catalog.TableDescriptor, colName string) (bool, error) {
-			storedCol, err := storedTable.FindColumnWithName(tree.Name(colName))
-			if err != nil {
-				if strings.Contains(err.Error(), "does not exist") {
-					return false, nil
-				}
-				return false, err
-			}
-			return !storedCol.IsNullable(), nil
-		},
+		name:           "alter-system-role-options-user-id-column-not-null",
+		schemaList:     []string{"user_id"},
+		query:          updateUserIDColumnRoleOptionsSetNotNull,
+		schemaExistsFn: columnExistsAndIsNotNull,
 	}
 	return migrateTable(ctx, cs, d, op, keys.RoleOptionsTableID, systemschema.RoleOptionsTable)
 }
diff --git a/pkg/upgrade/upgrades/system_users_role_id_migration.go b/pkg/upgrade/upgrades/system_users_role_id_migration.go
@@ -13,13 +13,11 @@ package upgrades
 import (
 	"context"
 	"fmt"
-	"strings"
 
 	"github.com/cockroachdb/cockroach/pkg/clusterversion"
 	"github.com/cockroachdb/cockroach/pkg/jobs"
 	"github.com/cockroachdb/cockroach/pkg/keys"
 	"github.com/cockroachdb/cockroach/pkg/security/username"
-	"github.com/cockroachdb/cockroach/pkg/sql/catalog"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/descidgen"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/systemschema"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
@@ -61,7 +59,7 @@ func alterSystemUsersAddUserIDColumnWithIndex(
 			name:           "add-system-users-user-id-column",
 			schemaList:     []string{"user_id"},
 			query:          addUserIDColumn,
-			schemaExistsFn: hasColumn,
+			schemaExistsFn: columnExists,
 		},
 		{
 			name:           "alter-system-users-add-index",
@@ -156,19 +154,10 @@ func setUserIDNotNull(
 	ctx context.Context, cs clusterversion.ClusterVersion, d upgrade.TenantDeps, _ *jobs.Job,
 ) error {
 	op := operation{
-		name:       "alter-system-users-user-id-column-not-null",
-		schemaList: []string{"user_id"},
-		query:      updateUserIDColumnSetNotNull,
-		schemaExistsFn: func(storedTable, _ catalog.TableDescriptor, colName string) (bool, error) {
-			storedCol, err := storedTable.FindColumnWithName(tree.Name(colName))
-			if err != nil {
-				if strings.Contains(err.Error(), "does not exist") {
-					return false, nil
-				}
-				return false, err
-			}
-			return !storedCol.IsNullable(), nil
-		},
+		name:           "alter-system-users-user-id-column-not-null",
+		schemaList:     []string{"user_id"},
+		query:          updateUserIDColumnSetNotNull,
+		schemaExistsFn: columnExistsAndIsNotNull,
 	}
 	return migrateTable(ctx, cs, d, op, keys.UsersTableID, systemschema.UsersTable)
 }