schemachanger: port over MR zone config helpers

Some precursor work to remove RBR fallbacks for ADD COLUMN.

Epic: CRDB-31282
Informs: #80545

Release note: None

pkg/sql/schemachanger/scbuild/internal/scbuildstmt/BUILD.bazel

@@ -84,6 +84,7 @@ go_library(
         "//pkg/sql/pgwire/pgerror",
         "//pkg/sql/pgwire/pgnotice",
         "//pkg/sql/privilege",
+        "//pkg/sql/regions",
         "//pkg/sql/rowenc",
         "//pkg/sql/rowenc/valueside",
         "//pkg/sql/schemachange",
@@ -117,6 +118,7 @@ go_library(
         "@com_github_cockroachdb_errors//:errors",
         "@com_github_cockroachdb_redact//:redact",
         "@com_github_lib_pq//oid",
+        "@org_golang_google_protobuf//proto",
     ],
 )
 

pkg/sql/schemachanger/scbuild/internal/scbuildstmt/zone_config_helpers.go

@@ -19,10 +19,12 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/descpb"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/fetchpb"
+	"github.com/cockroachdb/cockroach/pkg/sql/catalog/multiregion"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/zone"
 	"github.com/cockroachdb/cockroach/pkg/sql/covering"
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgcode"
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgerror"
+	"github.com/cockroachdb/cockroach/pkg/sql/regions"
 	"github.com/cockroachdb/cockroach/pkg/sql/rowenc"
 	"github.com/cockroachdb/cockroach/pkg/sql/rowenc/valueside"
 	"github.com/cockroachdb/cockroach/pkg/sql/schemachanger/scpb"
@@ -35,6 +37,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/util/encoding"
 	"github.com/cockroachdb/cockroach/pkg/util/protoutil"
 	"github.com/cockroachdb/errors"
+	"google.golang.org/protobuf/proto"
 )
 
 type zoneConfigAuthorizer interface {
@@ -1244,6 +1247,20 @@ func isCorrespondingTemporaryIndex(
 	return maybeCorresponding != nil
 }
 
+// findCorrespondingTemporaryIndexByID finds the temporary index that
+// corresponds to the currently mutated index identified by ID.
+//
+// Callers should take care that AllocateIDs() has been called before
+// using this function.
+func findCorrespondingTemporaryIndexByID(
+	b BuildCtx, tableID catid.DescID, indexID catid.IndexID,
+) *scpb.TemporaryIndex {
+	return b.QueryByID(tableID).FilterTemporaryIndex().
+		Filter(func(_ scpb.Status, _ scpb.TargetStatus, e *scpb.TemporaryIndex) bool {
+			return e.SourceIndexID == indexID
+		}).MustGetZeroOrOneElement()
+}
+
 // getSubzoneSpansWithIdx groups each subzone span by their subzoneIndexes
 // for a lookup of which subzone spans a particular subzone is referred to by.
 func getSubzoneSpansWithIdx(
@@ -1333,3 +1350,174 @@ func constructSideEffectPartitionElem(
 	}
 	return elem
 }
+
+// TODO(annie): This is unused for now.
+var _ = configureZoneConfigForNewIndexPartitioning
+
+// configureZoneConfigForNewIndexPartitioning configures the zone config for any
+// new index in a REGIONAL BY ROW table.
+// This *must* be done after the index ID has been allocated.
+func configureZoneConfigForNewIndexPartitioning(
+	b BuildCtx, tableID catid.DescID, indexDesc descpb.IndexDescriptor,
+) error {
+	if indexDesc.ID == 0 {
+		return errors.AssertionFailedf("index %s does not have id", indexDesc.Name)
+	}
+	// For REGIONAL BY ROW tables, correctly configure relevant zone configurations.
+	localityRBR := b.QueryByID(tableID).FilterTableLocalityRegionalByRow().MustGetZeroOrOneElement()
+	if localityRBR != nil {
+		dbID := b.QueryByID(tableID).FilterNamespace().MustGetOneElement().DatabaseID
+		regionConfig, err := b.SynthesizeRegionConfig(b, dbID)
+		if err != nil {
+			return err
+		}
+
+		indexIDs := []descpb.IndexID{indexDesc.ID}
+		if idx := findCorrespondingTemporaryIndexByID(b, tableID, indexDesc.ID); idx != nil {
+			indexIDs = append(indexIDs, idx.IndexID)
+		}
+
+		if err := ApplyZoneConfigForMultiRegionTable(
+			b,
+			regionConfig,
+			tableID,
+			applyZoneConfigForMultiRegionTableOptionNewIndexes(indexIDs...),
+		); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+// ApplyZoneConfigForMultiRegionTable applies zone config settings based
+// on the options provided and adds the scpb.TableZoneConfig to our builder.
+func ApplyZoneConfigForMultiRegionTable(
+	b BuildCtx,
+	regionConfig multiregion.RegionConfig,
+	tableID catid.DescID,
+	opts ...applyZoneConfigForMultiRegionTableOption,
+) error {
+	currentZoneConfigWithRaw, err := b.ZoneConfigGetter().GetZoneConfig(b, tableID)
+	if err != nil {
+		return err
+	}
+	if currentZoneConfigWithRaw == nil {
+		currentZoneConfigWithRaw = zone.NewZoneConfigWithRawBytes(zonepb.NewZoneConfig(), nil)
+	}
+	newZoneConfig := *currentZoneConfigWithRaw.ZoneConfigProto()
+
+	for _, opt := range opts {
+		modifiedNewZoneConfig, err := opt(
+			newZoneConfig,
+			regionConfig,
+			tableID,
+		)
+		if err != nil {
+			return err
+		}
+		newZoneConfig = modifiedNewZoneConfig
+	}
+
+	if regionConfig.HasSecondaryRegion() {
+		var newLeasePreferences []zonepb.LeasePreference
+		localityRBR := b.QueryByID(tableID).FilterTableLocalityRegionalByRow().MustGetZeroOrOneElement()
+		switch {
+		case localityRBR != nil:
+			region := b.QueryByID(tableID).FilterTableLocalitySecondaryRegion().MustGetZeroOrOneElement()
+			if region != nil {
+				newLeasePreferences = regions.SynthesizeLeasePreferences(region.RegionName, regionConfig.SecondaryRegion())
+			} else {
+				newLeasePreferences = regions.SynthesizeLeasePreferences(regionConfig.PrimaryRegion(), regionConfig.SecondaryRegion())
+			}
+		default:
+			newLeasePreferences = regions.SynthesizeLeasePreferences(regionConfig.PrimaryRegion(), regionConfig.SecondaryRegion())
+		}
+		newZoneConfig.LeasePreferences = newLeasePreferences
+	}
+
+	// Mark the NumReplicas as 0 if we have subzones but no other features
+	// in the zone config. This signifies a placeholder.
+	// Note we do not use hasNewSubzones here as there may be existing subzones
+	// on the zone config which may still be a placeholder.
+	if regions.IsPlaceholderZoneConfigForMultiRegion(newZoneConfig) {
+		newZoneConfig.NumReplicas = proto.Int32(0)
+	}
+
+	// Determine if we're rewriting or deleting the zone configuration.
+	newZoneConfigIsEmpty := newZoneConfig.Equal(zonepb.NewZoneConfig())
+	currentZoneConfigIsEmpty := currentZoneConfigWithRaw.ZoneConfigProto().Equal(zonepb.NewZoneConfig())
+	rewriteZoneConfig := !newZoneConfigIsEmpty
+	deleteZoneConfig := newZoneConfigIsEmpty && !currentZoneConfigIsEmpty
+
+	if deleteZoneConfig {
+		return nil
+	}
+	if !rewriteZoneConfig {
+		return nil
+	}
+
+	if err := newZoneConfig.Validate(); err != nil {
+		return pgerror.Wrap(
+			err,
+			pgcode.CheckViolation,
+			"could not validate zone config",
+		)
+	}
+	if err := newZoneConfig.ValidateTandemFields(); err != nil {
+		return pgerror.Wrap(
+			err,
+			pgcode.CheckViolation,
+			"could not validate zone config",
+		)
+	}
+	if len(newZoneConfig.Subzones) > 0 {
+		newZoneConfig.SubzoneSpans, err = generateSubzoneSpans(b, tableID, newZoneConfig.Subzones)
+		if err != nil {
+			return err
+		}
+	} else {
+		// To keep the Subzone and SubzoneSpan arrays consistent
+		newZoneConfig.SubzoneSpans = nil
+	}
+	if newZoneConfig.IsSubzonePlaceholder() && len(newZoneConfig.Subzones) == 0 {
+		return nil
+	}
+	tzc := &scpb.TableZoneConfig{TableID: tableID, ZoneConfig: &newZoneConfig, SeqNum: 1}
+	b.Add(tzc)
+	return nil
+}
+
+// applyZoneConfigForMultiRegionTableOptionNewIndexes applies table zone configs
+// for a newly added index which requires partitioning of individual indexes.
+func applyZoneConfigForMultiRegionTableOptionNewIndexes(
+	indexIDs ...descpb.IndexID,
+) applyZoneConfigForMultiRegionTableOption {
+	return func(
+		zoneConfig zonepb.ZoneConfig,
+		regionConfig multiregion.RegionConfig,
+		tableID catid.DescID,
+	) (newZoneConfig zonepb.ZoneConfig, err error) {
+		for _, indexID := range indexIDs {
+			for _, region := range regionConfig.Regions() {
+				zc, err := regions.ZoneConfigForMultiRegionPartition(region, regionConfig)
+				if err != nil {
+					return zoneConfig, err
+				}
+				zoneConfig.SetSubzone(zonepb.Subzone{
+					IndexID:       uint32(indexID),
+					PartitionName: string(region),
+					Config:        zc,
+				})
+			}
+		}
+		return zoneConfig, nil
+	}
+}
+
+// applyZoneConfigForMultiRegionTableOption is an option that can be passed into
+// applyZoneConfigForMultiRegionTable.
+type applyZoneConfigForMultiRegionTableOption func(
+	zoneConfig zonepb.ZoneConfig,
+	regionConfig multiregion.RegionConfig,
+	tableID catid.DescID,
+) (newZoneConfig zonepb.ZoneConfig, err error)