sql/schemachanger: extend metadata updater to support for updating zone configs

pkg/ccl/backupccl/restore_job.go

@@ -896,6 +896,7 @@ func createImportingDescriptors(
 								regionConfig,
 								txn,
 								p.ExecCfg(),
+								descsCol,
 							); err != nil {
 								return err
 							}
@@ -1045,6 +1046,7 @@ func createImportingDescriptors(
 							ctx,
 							txn,
 							p.ExecCfg(),
+							descsCol,
 							regionConfig,
 							mutTable,
 							sql.ApplyZoneConfigForMultiRegionTableOptionTableAndIndexes,

pkg/sql/alter_database.go

@@ -530,6 +530,7 @@ func removeLocalityConfigFromAllTablesInDB(
 					ctx,
 					p.txn,
 					p.ExecCfg(),
+					p.Descriptors(),
 					multiregion.RegionConfig{}, // pass dummy config as it is not used.
 					tbDesc,
 					applyZoneConfigForMultiRegionTableOptionRemoveGlobalZoneConfig,
@@ -600,6 +601,7 @@ func (n *alterDatabaseDropRegionNode) startExec(params runParams) error {
 			n.desc.ID,
 			params.p.txn,
 			params.p.execCfg,
+			params.p.Descriptors(),
 		); err != nil {
 			return err
 		}
@@ -756,6 +758,7 @@ func (n *alterDatabasePrimaryRegionNode) switchPrimaryRegion(params runParams) e
 		updatedRegionConfig,
 		params.p.txn,
 		params.p.execCfg,
+		params.p.Descriptors(),
 	); err != nil {
 		return err
 	}
@@ -1129,6 +1132,7 @@ func (n *alterDatabaseSurvivalGoalNode) startExec(params runParams) error {
 		regionConfig,
 		params.p.txn,
 		params.p.execCfg,
+		params.p.Descriptors(),
 	); err != nil {
 		return err
 	}
@@ -1258,6 +1262,7 @@ func (n *alterDatabasePlacementNode) startExec(params runParams) error {
 		regionConfig,
 		params.p.txn,
 		params.p.execCfg,
+		params.p.Descriptors(),
 	); err != nil {
 		return err
 	}

pkg/sql/alter_index.go

@@ -127,6 +127,7 @@ func (n *alterIndexNode) startExec(params runParams) error {
 					params.ctx,
 					params.p.txn,
 					n.tableDesc,
+					params.p.Descriptors(),
 					n.index.GetID(),
 					oldPartitioning,
 					n.index.GetPartitioning(),

pkg/sql/alter_table.go

@@ -698,6 +698,7 @@ func (n *alterTableNode) startExec(params runParams) error {
 					params.ctx,
 					params.p.txn,
 					n.tableDesc,
+					params.p.Descriptors(),
 					n.tableDesc.GetPrimaryIndexID(),
 					oldPartitioning,
 					n.tableDesc.GetPrimaryIndex().GetPartitioning(),

pkg/sql/alter_table_locality.go

@@ -599,6 +599,7 @@ func (n *alterTableSetLocalityNode) writeNewTableLocalityAndZoneConfig(
 		params.ctx,
 		params.p.txn,
 		params.p.ExecCfg(),
+		params.p.Descriptors(),
 		regionConfig,
 		n.tableDesc,
 		ApplyZoneConfigForMultiRegionTableOptionTableAndIndexes,

pkg/sql/backfill.go

@@ -2326,7 +2326,7 @@ func runSchemaChangesInTxn(
 				}
 			} else if idx := m.AsIndex(); idx != nil {
 				if err := indexTruncateInTxn(
-					ctx, planner.Txn(), planner.ExecCfg(), planner.EvalContext(), immutDesc, idx, traceKV,
+					ctx, planner.Txn(), planner.ExecCfg(), planner.Descriptors(), planner.EvalContext(), immutDesc, idx, traceKV,
 				); err != nil {
 					return err
 				}
@@ -2698,6 +2698,7 @@ func indexTruncateInTxn(
 	ctx context.Context,
 	txn *kv.Txn,
 	execCfg *ExecutorConfig,
+	descriptors *descs.Collection,
 	evalCtx *eval.Context,
 	tableDesc catalog.TableDescriptor,
 	idx catalog.Index,
@@ -2724,7 +2725,7 @@ func indexTruncateInTxn(
 		}
 	}
 	// Remove index zone configs.
-	return RemoveIndexZoneConfigs(ctx, txn, execCfg, tableDesc, []uint32{uint32(idx.GetID())})
+	return RemoveIndexZoneConfigs(ctx, txn, execCfg, descriptors, tableDesc, []uint32{uint32(idx.GetID())})
 }
 
 // We note the time at the start of the merge in order to limit the set of

pkg/sql/create_index.go

@@ -913,6 +913,7 @@ func (p *planner) configureZoneConfigForNewIndexPartitioning(
 			ctx,
 			p.txn,
 			p.ExecCfg(),
+			p.Descriptors(),
 			regionConfig,
 			tableDesc,
 			applyZoneConfigForMultiRegionTableOptionNewIndexes(indexIDs...),

pkg/sql/create_table.go

@@ -477,6 +477,7 @@ func (n *createTableNode) startExec(params runParams) error {
 			params.ctx,
 			params.p.txn,
 			params.p.ExecCfg(),
+			params.p.Descriptors(),
 			regionConfig,
 			desc,
 			ApplyZoneConfigForMultiRegionTableOptionTableAndIndexes,

pkg/sql/create_view.go

@@ -333,6 +333,7 @@ func (n *createViewNode) startExec(params runParams) error {
 					params.ctx,
 					params.p.txn,
 					params.p.ExecCfg(),
+					params.p.Descriptors(),
 					regionConfig,
 					newDesc,
 					applyZoneConfigForMultiRegionTableOptionTableNewConfig(

pkg/sql/database_region_change_finalizer.go

@@ -134,7 +134,7 @@ func (r *databaseRegionChangeFinalizer) preDrop(ctx context.Context, txn *kv.Txn
 	}
 	for _, update := range zoneConfigUpdates {
 		if _, err := writeZoneConfigUpdate(
-			ctx, txn, r.localPlanner.ExecCfg(), update,
+			ctx, txn, r.localPlanner.ExecCfg(), r.localPlanner.Descriptors(), update,
 		); err != nil {
 			return err
 		}
@@ -210,6 +210,7 @@ func (r *databaseRegionChangeFinalizer) updateDatabaseZoneConfig(
 		regionConfig,
 		txn,
 		r.localPlanner.ExecCfg(),
+		r.localPlanner.Descriptors(),
 	)
 }
 

pkg/sql/descmetadata/BUILD.bazel

@@ -9,13 +9,16 @@ go_library(
     importpath = "github.com/cockroachdb/cockroach/pkg/sql/descmetadata",
     visibility = ["//visibility:public"],
     deps = [
+        "//pkg/config/zonepb",
         "//pkg/keys",
         "//pkg/kv",
         "//pkg/security/username",
         "//pkg/settings",
         "//pkg/sql/catalog",
         "//pkg/sql/catalog/descpb",
         "//pkg/sql/catalog/descs",
+        "//pkg/sql/pgwire/pgcode",
+        "//pkg/sql/pgwire/pgerror",
         "//pkg/sql/schemachanger/scbuild",
         "//pkg/sql/schemachanger/scexec",
         "//pkg/sql/sem/catid",
@@ -24,6 +27,7 @@ go_library(
         "//pkg/sql/sessiondatapb",
         "//pkg/sql/sessioninit",
         "//pkg/sql/sqlutil",
+        "//pkg/util/protoutil",
     ],
 )
 

pkg/sql/descmetadata/metadata_updater.go

@@ -15,19 +15,23 @@ import (
 	"fmt"
 	"strings"
 
+	"github.com/cockroachdb/cockroach/pkg/config/zonepb"
 	"github.com/cockroachdb/cockroach/pkg/keys"
 	"github.com/cockroachdb/cockroach/pkg/kv"
 	"github.com/cockroachdb/cockroach/pkg/security/username"
 	"github.com/cockroachdb/cockroach/pkg/settings"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/descpb"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/descs"
+	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgcode"
+	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgerror"
 	"github.com/cockroachdb/cockroach/pkg/sql/schemachanger/scexec"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
 	"github.com/cockroachdb/cockroach/pkg/sql/sessiondata"
 	"github.com/cockroachdb/cockroach/pkg/sql/sessiondatapb"
 	"github.com/cockroachdb/cockroach/pkg/sql/sessioninit"
 	"github.com/cockroachdb/cockroach/pkg/sql/sqlutil"
+	"github.com/cockroachdb/cockroach/pkg/util/protoutil"
 )
 
 // metadataUpdater which implements scexec.MetaDataUpdater that is used to update
@@ -217,3 +221,25 @@ func (mu metadataUpdater) DeleteSchedule(ctx context.Context, scheduleID int64)
 	)
 	return err
 }
+
+// DeleteZoneConfig implements scexec.DescriptorMetadataUpdater.
+func (mu metadataUpdater) DeleteZoneConfig(
+	ctx context.Context, id descpb.ID,
+) (numAffected int, err error) {
+	ie := mu.ieFactory(mu.ctx, mu.sessionData)
+	return ie.Exec(ctx, "delete-zone", mu.txn,
+		"DELETE FROM system.zones WHERE id = $1", id)
+}
+
+// UpsertZoneConfig implements scexec.DescriptorMetadataUpdater.
+func (mu metadataUpdater) UpsertZoneConfig(
+	ctx context.Context, id descpb.ID, zone *zonepb.ZoneConfig,
+) (numAffected int, err error) {
+	ie := mu.ieFactory(mu.ctx, mu.sessionData)
+	bytes, err := protoutil.Marshal(zone)
+	if err != nil {
+		return 0, pgerror.Wrap(err, pgcode.CheckViolation, "could not marshal zone config")
+	}
+	return ie.Exec(ctx, "upsert-zone", mu.txn,
+		"UPSERT INTO system.zones (id, config) VALUES ($1, $2)", id, bytes)
+}

pkg/sql/gcjob/index_garbage_collection.go

@@ -79,7 +79,7 @@ func gcIndexes(
 				return err
 			}
 			return sql.RemoveIndexZoneConfigs(
-				ctx, txn, execCfg, freshParentTableDesc, []uint32{uint32(index.IndexID)},
+				ctx, txn, execCfg, descriptors, freshParentTableDesc, []uint32{uint32(index.IndexID)},
 			)
 		}
 		if err := sql.DescsTxn(ctx, execCfg, removeIndexZoneConfigs); err != nil {

pkg/sql/region_util.go

@@ -825,6 +825,7 @@ func ApplyZoneConfigForMultiRegionTable(
 	ctx context.Context,
 	txn *kv.Txn,
 	execCfg *ExecutorConfig,
+	descriptors *descs.Collection,
 	regionConfig multiregion.RegionConfig,
 	table catalog.TableDescriptor,
 	opts ...applyZoneConfigForMultiRegionTableOption,
@@ -833,7 +834,7 @@ func ApplyZoneConfigForMultiRegionTable(
 	if update == nil || err != nil {
 		return err
 	}
-	_, err = writeZoneConfigUpdate(ctx, txn, execCfg, update)
+	_, err = writeZoneConfigUpdate(ctx, txn, execCfg, descriptors, update)
 	return err
 }
 
@@ -845,6 +846,7 @@ func ApplyZoneConfigFromDatabaseRegionConfig(
 	regionConfig multiregion.RegionConfig,
 	txn *kv.Txn,
 	execConfig *ExecutorConfig,
+	descriptors *descs.Collection,
 ) error {
 	// Build a zone config based on the RegionConfig information.
 	dbZoneConfig, err := zoneConfigForMultiRegionDatabase(regionConfig)
@@ -857,13 +859,18 @@ func ApplyZoneConfigFromDatabaseRegionConfig(
 		dbZoneConfig,
 		txn,
 		execConfig,
+		descriptors,
 	)
 }
 
 // discardMultiRegionFieldsForDatabaseZoneConfig resets the multi-region zone
 // config fields for a multi-region database.
 func discardMultiRegionFieldsForDatabaseZoneConfig(
-	ctx context.Context, dbID descpb.ID, txn *kv.Txn, execConfig *ExecutorConfig,
+	ctx context.Context,
+	dbID descpb.ID,
+	txn *kv.Txn,
+	execConfig *ExecutorConfig,
+	descriptors *descs.Collection,
 ) error {
 	// Merge with an empty zone config.
 	return applyZoneConfigForMultiRegionDatabase(
@@ -872,6 +879,7 @@ func discardMultiRegionFieldsForDatabaseZoneConfig(
 		*zonepb.NewZoneConfig(),
 		txn,
 		execConfig,
+		descriptors,
 	)
 }
 
@@ -881,6 +889,7 @@ func applyZoneConfigForMultiRegionDatabase(
 	mergeZoneConfig zonepb.ZoneConfig,
 	txn *kv.Txn,
 	execConfig *ExecutorConfig,
+	descriptors *descs.Collection,
 ) error {
 	currentZoneConfig, err := getZoneConfigRaw(ctx, txn, execConfig.Codec, execConfig.Settings, dbID)
 	if err != nil {
@@ -912,6 +921,7 @@ func applyZoneConfigForMultiRegionDatabase(
 		nil, /* table */
 		&newZoneConfig,
 		execConfig,
+		descriptors,
 		false, /* hasNewSubzones */
 	); err != nil {
 		return err
@@ -951,6 +961,7 @@ func (p *planner) refreshZoneConfigsForTables(
 					ctx,
 					p.txn,
 					p.ExecCfg(),
+					p.Descriptors(),
 					regionConfig,
 					tbDesc,
 					ApplyZoneConfigForMultiRegionTableOptionTableAndIndexes,
@@ -1014,7 +1025,8 @@ func (p *planner) maybeInitializeMultiRegionDatabase(
 		desc.ID,
 		*regionConfig,
 		p.txn,
-		p.execCfg); err != nil {
+		p.execCfg,
+		p.Descriptors()); err != nil {
 		return err
 	}
 
@@ -1099,6 +1111,7 @@ func (p *planner) ResetMultiRegionZoneConfigsForTable(ctx context.Context, id in
 		ctx,
 		p.txn,
 		p.ExecCfg(),
+		p.Descriptors(),
 		regionConfig,
 		desc,
 		ApplyZoneConfigForMultiRegionTableOptionTableAndIndexes,
@@ -1138,6 +1151,7 @@ func (p *planner) ResetMultiRegionZoneConfigsForDatabase(ctx context.Context, id
 		regionConfig,
 		p.txn,
 		p.execCfg,
+		p.Descriptors(),
 	); err != nil {
 		return err
 	}

pkg/sql/schema_changer.go

@@ -1892,6 +1892,7 @@ func maybeUpdateZoneConfigsForPKChange(
 	ctx context.Context,
 	txn *kv.Txn,
 	execCfg *ExecutorConfig,
+	descriptors *descs.Collection,
 	table *tabledesc.Mutable,
 	swapInfo *descpb.PrimaryKeySwap,
 ) error {
@@ -1934,7 +1935,7 @@ func maybeUpdateZoneConfigsForPKChange(
 
 	// Write the zone back. This call regenerates the index spans that apply
 	// to each partition in the index.
-	_, err = writeZoneConfig(ctx, txn, table.ID, table, zone, execCfg, false)
+	_, err = writeZoneConfig(ctx, txn, table.ID, table, zone, execCfg, descriptors, false)
 	if err != nil && !sqlerrors.IsCCLRequiredError(err) {
 		return err
 	}
@@ -3030,6 +3031,7 @@ func (sc *SchemaChanger) applyZoneConfigChangeForMutation(
 				ctx,
 				txn,
 				sc.execCfg,
+				descsCol,
 				regionConfig,
 				tableDesc,
 				opts...,
@@ -3042,7 +3044,7 @@ func (sc *SchemaChanger) applyZoneConfigChangeForMutation(
 		// Note this is done even for isDone = true, though not strictly
 		// necessary.
 		return maybeUpdateZoneConfigsForPKChange(
-			ctx, txn, sc.execCfg, tableDesc, pkSwap.PrimaryKeySwapDesc(),
+			ctx, txn, sc.execCfg, descsCol, tableDesc, pkSwap.PrimaryKeySwapDesc(),
 		)
 	}
 	return nil

pkg/sql/schemachanger/scdeps/sctestdeps/BUILD.bazel

@@ -13,6 +13,7 @@ go_library(
     visibility = ["//visibility:public"],
     deps = [
         "//pkg/clusterversion",
+        "//pkg/config/zonepb",
         "//pkg/jobs",
         "//pkg/jobs/jobspb",
         "//pkg/keys",

pkg/sql/schemachanger/scdeps/sctestdeps/config.go

@@ -13,12 +13,14 @@ package sctestdeps
 import (
 	"time"
 
+	"github.com/cockroachdb/cockroach/pkg/config/zonepb"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/descbuilder"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/nstree"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/tabledesc"
 	"github.com/cockroachdb/cockroach/pkg/sql/descmetadata"
 	"github.com/cockroachdb/cockroach/pkg/sql/schemachanger/scexec"
+	"github.com/cockroachdb/cockroach/pkg/sql/sem/catid"
 	"github.com/cockroachdb/cockroach/pkg/sql/sessiondata"
 	"github.com/cockroachdb/cockroach/pkg/util/hlc"
 )
@@ -149,5 +151,6 @@ var defaultOptions = []Option{
 		state.merger = &testBackfiller{s: state}
 		state.indexSpanSplitter = &indexSpanSplitter{}
 		state.approximateTimestamp = defaultCreatedAt
+		state.zoneConfigs = make(map[catid.DescID]*zonepb.ZoneConfig)
 	}),
 }