Merge pull request #135833 from fqazi/backport24.3.0-rc-135737

release-24.3.0-rc: upgrade: adding is_draining to system.sql_instance can fail

pkg/ccl/multiregionccl/multiregion_system_table_test.go

@@ -19,6 +19,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/server"
 	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
+	"github.com/cockroachdb/cockroach/pkg/sql"
 	"github.com/cockroachdb/cockroach/pkg/sql/enum"
 	"github.com/cockroachdb/cockroach/pkg/sql/sqlinstance/instancestorage"
 	"github.com/cockroachdb/cockroach/pkg/sql/sqlliveness"
@@ -562,7 +563,7 @@ func TestMrSystemDatabaseUpgrade(t *testing.T) {
 
 	// Enable settings required for configuring a tenant's system database as multi-region.
 	makeSettings := func() *cluster.Settings {
-		cs := cluster.MakeTestingClusterSettingsWithVersions(clusterversion.V24_1.Version(),
+		cs := cluster.MakeTestingClusterSettingsWithVersions(clusterversion.Latest.Version(),
 			clusterversion.MinSupported.Version(),
 			false)
 		instancestorage.ReclaimLoopInterval.Override(ctx, &cs.SV, 150*time.Millisecond)
@@ -578,10 +579,14 @@ func TestMrSystemDatabaseUpgrade(t *testing.T) {
 		},
 		multiregionccltestutils.WithSettings(makeSettings()))
 	defer cleanup()
-
 	id, err := roachpb.MakeTenantID(11)
 	require.NoError(t, err)
 
+	// Disable license enforcement for this test.
+	for _, s := range cluster.Servers {
+		s.ExecutorConfig().(sql.ExecutorConfig).LicenseEnforcer.Disable(ctx)
+	}
+
 	tenantArgs := base.TestTenantArgs{
 		Settings: makeSettings(),
 		TestingKnobs: base.TestingKnobs{
@@ -593,7 +598,8 @@ func TestMrSystemDatabaseUpgrade(t *testing.T) {
 		TenantID: id,
 		Locality: cluster.Servers[0].Locality(),
 	}
-	_, tenantSQL := serverutils.StartTenant(t, cluster.Servers[0], tenantArgs)
+	appLayer, tenantSQL := serverutils.StartTenant(t, cluster.Servers[0], tenantArgs)
+	appLayer.ExecutorConfig().(sql.ExecutorConfig).LicenseEnforcer.Disable(ctx)
 
 	tDB := sqlutils.MakeSQLRunner(tenantSQL)
 

pkg/upgrade/upgrades/v24_3_sql_instances_add_draining.go

@@ -9,35 +9,22 @@ import (
 	"context"
 
 	"github.com/cockroachdb/cockroach/pkg/clusterversion"
-	"github.com/cockroachdb/cockroach/pkg/sql/catalog/descpb"
-	"github.com/cockroachdb/cockroach/pkg/sql/catalog/descs"
+	"github.com/cockroachdb/cockroach/pkg/keys"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/systemschema"
 	"github.com/cockroachdb/cockroach/pkg/upgrade"
-	"github.com/cockroachdb/cockroach/pkg/util/protoutil"
-	"github.com/cockroachdb/errors"
 )
 
 // sqlInstancesAddDrainingMigration adds a new column `is_draining` to the
 // system.sql_instances table.
 func sqlInstancesAddDrainingMigration(
-	ctx context.Context, cs clusterversion.ClusterVersion, deps upgrade.TenantDeps,
+	ctx context.Context, clusterVersion clusterversion.ClusterVersion, deps upgrade.TenantDeps,
 ) error {
-	finalDescriptor := systemschema.SQLInstancesTable()
-	// Replace the stored descriptor with the bootstrap descriptor.
-	err := deps.DB.DescsTxn(ctx, func(ctx context.Context, txn descs.Txn) error {
-		expectedDesc := finalDescriptor.TableDesc()
-		mutableDesc, err := txn.Descriptors().MutableByID(txn.KV()).Table(ctx, expectedDesc.GetID())
-		if err != nil {
-			return err
-		}
-		version := mutableDesc.Version
-		mutableDesc.TableDescriptor = *protoutil.Clone(expectedDesc).(*descpb.TableDescriptor)
-		mutableDesc.Version = version
-		return txn.Descriptors().WriteDesc(ctx, false, mutableDesc, txn.KV())
-	})
-	if err != nil {
-		return errors.Wrapf(err, "unable to replace system descriptor for system.%s (%+v)",
-			finalDescriptor.GetName(), finalDescriptor)
-	}
-	return err
+	return migrateTable(ctx, clusterVersion, deps, operation{
+		name:           "add-draining-column",
+		schemaList:     []string{"is_draining"},
+		schemaExistsFn: columnExists,
+		query:          `ALTER TABLE system.sql_instances ADD COLUMN IF NOT EXISTS is_draining BOOL NULL FAMILY "primary"`,
+	},
+		keys.SQLInstancesTableID,
+		systemschema.SQLInstancesTable())
 }