schemachanger: add db zone config opgen rule

This patch adds an opgen rule for db zone configs in the ADD direction. Informs: cockroachdb#117574 Release note: None
annrpom · Jun 21, 2024 · 6bd82c1 · 6bd82c1
1 parent dcad64c
commit 6bd82c1
Show file tree

Hide file tree

Showing 18 changed files with 186 additions and 4 deletions.
diff --git a/pkg/sql/catalog/zone/zone_config.go b/pkg/sql/catalog/zone/zone_config.go
@@ -18,8 +18,14 @@ import (
 
 // ZoneConfigWithRawBytes wraps a zone config together with its expected
 // raw bytes. For cached modified zone configs, the raw bytes should be the
-// deserialized bytes from the zone config proto. Though, the raw bytes should
-// be the raw value read from kv when it's first read from storage.
+// deserialized bytes from the zone config proto. When performing a CPut,
+// raw bytes should represent the same bytes that we expect to read from kv
+// when it's first read from storage. This is to ensure that the existing
+// value hasn't changed since we last read it.
+//
+// N.B. if we don't expect any bytes to be read from storage (i.e. we are
+// inserting a new zone config -- not updating an existing one), then the value
+// of raw bytes should be nil.
 type ZoneConfigWithRawBytes struct {
 	zc *zonepb.ZoneConfig
 	rb []byte

diff --git a/pkg/sql/schemachanger/scdeps/BUILD.bazel b/pkg/sql/schemachanger/scdeps/BUILD.bazel
@@ -11,6 +11,7 @@ go_library(
     importpath = "github.com/cockroachdb/cockroach/pkg/sql/schemachanger/scdeps",
     visibility = ["//visibility:public"],
     deps = [
+        "//pkg/config/zonepb",
         "//pkg/jobs",
         "//pkg/jobs/jobspb",
         "//pkg/keys",
@@ -27,6 +28,7 @@ go_library(
         "//pkg/sql/catalog/descs",
         "//pkg/sql/catalog/nstree",
         "//pkg/sql/catalog/resolver",
+        "//pkg/sql/catalog/zone",
         "//pkg/sql/isql",
         "//pkg/sql/rowenc",
         "//pkg/sql/schemachanger/scbuild",

diff --git a/pkg/sql/schemachanger/scdeps/exec_deps.go b/pkg/sql/schemachanger/scdeps/exec_deps.go
@@ -15,6 +15,7 @@ import (
 	"fmt"
 	"time"
 
+	"github.com/cockroachdb/cockroach/pkg/config/zonepb"
 	"github.com/cockroachdb/cockroach/pkg/jobs"
 	"github.com/cockroachdb/cockroach/pkg/jobs/jobspb"
 	"github.com/cockroachdb/cockroach/pkg/keys"
@@ -27,6 +28,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/catalogkeys"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/descpb"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/descs"
+	"github.com/cockroachdb/cockroach/pkg/sql/catalog/zone"
 	"github.com/cockroachdb/cockroach/pkg/sql/isql"
 	"github.com/cockroachdb/cockroach/pkg/sql/schemachanger/scexec"
 	"github.com/cockroachdb/cockroach/pkg/sql/schemachanger/scexec/scmutationexec"
@@ -222,6 +224,25 @@ func (d *txnDeps) DeleteDescriptor(ctx context.Context, id descpb.ID) error {
 	return d.descsCollection.DeleteDescToBatch(ctx, d.kvTrace, id, d.getOrCreateBatch())
 }
 
+// UpdateZoneConfig implements the scexec.Catalog interface.
+func (d *txnDeps) UpdateZoneConfig(ctx context.Context, id descpb.ID, zc zonepb.ZoneConfig) error {
+	var newZc catalog.ZoneConfig
+	oldZc, err := d.descsCollection.GetZoneConfig(ctx, d.txn.KV(), id)
+	if err != nil {
+		return err
+	}
+
+	var rawBytes []byte
+	// If the zone config already exists, we need to preserve the raw bytes as the
+	// expected value that we will be updating. Otherwise, this will be a clean
+	// insert with no expected raw bytes.
+	if oldZc != nil {
+		rawBytes = oldZc.GetRawBytesInStorage()
+	}
+	newZc = zone.NewZoneConfigWithRawBytes(&zc, rawBytes)
+	return d.descsCollection.WriteZoneConfigToBatch(ctx, d.kvTrace, d.getOrCreateBatch(), id, newZc)
+}
+
 // DeleteZoneConfig implements the scexec.Catalog interface.
 func (d *txnDeps) DeleteZoneConfig(ctx context.Context, id descpb.ID) error {
 	return d.descsCollection.DeleteZoneConfigInBatch(ctx, d.kvTrace, d.getOrCreateBatch(), id)

diff --git a/pkg/sql/schemachanger/scdeps/sctestdeps/BUILD.bazel b/pkg/sql/schemachanger/scdeps/sctestdeps/BUILD.bazel
@@ -18,7 +18,9 @@ go_library(
         "//pkg/jobs",
         "//pkg/jobs/jobspb",
         "//pkg/keys",
+        "//pkg/roachpb",
         "//pkg/security/username",
+        "//pkg/server/serverpb",
         "//pkg/settings/cluster",
         "//pkg/sql/catalog",
         "//pkg/sql/catalog/catalogkeys",
@@ -37,6 +39,7 @@ go_library(
         "//pkg/sql/faketreeeval",
         "//pkg/sql/privilege",
         "//pkg/sql/schemachanger/scbuild",
+        "//pkg/sql/schemachanger/scdecomp",
         "//pkg/sql/schemachanger/scdeps",
         "//pkg/sql/schemachanger/scdeps/sctestutils",
         "//pkg/sql/schemachanger/scexec",

diff --git a/pkg/sql/schemachanger/scdeps/sctestdeps/test_deps.go b/pkg/sql/schemachanger/scdeps/sctestdeps/test_deps.go
@@ -18,10 +18,13 @@ import (
 	"time"
 
 	"github.com/cockroachdb/cockroach/pkg/clusterversion"
+	"github.com/cockroachdb/cockroach/pkg/config/zonepb"
 	"github.com/cockroachdb/cockroach/pkg/jobs"
 	"github.com/cockroachdb/cockroach/pkg/jobs/jobspb"
 	"github.com/cockroachdb/cockroach/pkg/keys"
+	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/security/username"
+	"github.com/cockroachdb/cockroach/pkg/server/serverpb"
 	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/catalogkeys"
@@ -35,6 +38,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/typedesc"
 	"github.com/cockroachdb/cockroach/pkg/sql/privilege"
 	"github.com/cockroachdb/cockroach/pkg/sql/schemachanger/scbuild"
+	"github.com/cockroachdb/cockroach/pkg/sql/schemachanger/scdecomp"
 	"github.com/cockroachdb/cockroach/pkg/sql/schemachanger/scdeps"
 	"github.com/cockroachdb/cockroach/pkg/sql/schemachanger/scdeps/sctestutils"
 	"github.com/cockroachdb/cockroach/pkg/sql/schemachanger/scexec"
@@ -797,6 +801,17 @@ func (s *TestState) DeleteDescriptor(ctx context.Context, id descpb.ID) error {
 	return nil
 }
 
+// UpdateZoneConfig implements the scexec.Catalog interface.
+func (s *TestState) UpdateZoneConfig(
+	ctx context.Context, id descpb.ID, zc zonepb.ZoneConfig,
+) error {
+	if s.catalogChanges.zoneConfigsToUpdate == nil {
+		s.catalogChanges.zoneConfigsToUpdate = make(map[descpb.ID]*zonepb.ZoneConfig)
+	}
+	s.catalogChanges.zoneConfigsToUpdate[id] = &zc
+	return nil
+}
+
 // DeleteZoneConfig implements the scexec.Catalog interface.
 func (s *TestState) DeleteZoneConfig(ctx context.Context, id descpb.ID) error {
 	s.catalogChanges.zoneConfigsToDelete.Add(id)
@@ -869,6 +884,18 @@ func (s *TestState) Validate(ctx context.Context) error {
 		s.LogSideEffectf("deleting zone config for #%d", deletedID)
 		s.uncommittedInMemory.DeleteZoneConfig(deletedID)
 	}
+	for upsertedID, zc := range s.catalogChanges.zoneConfigsToUpdate {
+		s.LogSideEffectf("upsert zone config for #%d", upsertedID)
+		var val roachpb.Value
+		if err := val.SetProto(zc); err != nil {
+			return err
+		}
+		valBytes, err := val.GetBytes()
+		if err != nil {
+			return err
+		}
+		s.uncommittedInMemory.UpsertZoneConfig(upsertedID, zc, valBytes)
+	}
 	commentKeys := make([]catalogkeys.CommentKey, 0, len(s.catalogChanges.commentsToUpdate))
 	for key := range s.catalogChanges.commentsToUpdate {
 		commentKeys = append(commentKeys, key)
@@ -1325,13 +1352,13 @@ func (s *TestState) ResolveFunctionByOID(
 }
 
 // ZoneConfigGetter implements scexec.Dependencies.
-func (s *TestState) ZoneConfigGetter() scbuild.ZoneConfigGetter {
+func (s *TestState) ZoneConfigGetter() scdecomp.ZoneConfigGetter {
 	return s
 }
 
 // GetZoneConfig implements scexec.Dependencies.
 func (s *TestState) GetZoneConfig(ctx context.Context, id descpb.ID) (catalog.ZoneConfig, error) {
-	return s.uncommittedInMemory.LookupZoneConfig(id), nil
+	return s.committed.LookupZoneConfig(id), nil
 }
 
 func (s *TestState) get(
@@ -1433,3 +1460,11 @@ func (s *TestState) InsertTemporarySchema(
 func (s *TestState) TemporarySchemaProvider() scbuild.TemporarySchemaProvider {
 	return s
 }
+
+func (s *TestState) NodesStatusServer() *serverpb.OptionalNodesStatusServer {
+	return &serverpb.OptionalNodesStatusServer{}
+}
+
+func (s *TestState) GetRegions() (*serverpb.RegionsResponse, error) {
+	return &serverpb.RegionsResponse{Regions: map[string]*serverpb.RegionsResponse_Region{}}, nil
+}
diff --git a/pkg/sql/schemachanger/scdeps/sctestdeps/test_state.go b/pkg/sql/schemachanger/scdeps/sctestdeps/test_state.go
@@ -104,6 +104,8 @@ func NewTestDependencies(options ...Option) *TestState {
 	for _, o := range options {
 		o.apply(&s)
 	}
+	zc := zonepb.DefaultSystemZoneConfigRef()
+	s.committed.UpsertZoneConfig(0, zc, nil)
 	s.uncommittedInMemory = catalogDeepCopy(s.committed.Catalog)
 	s.uncommittedInStorage = catalogDeepCopy(s.committed.Catalog)
 	return &s

diff --git a/pkg/sql/schemachanger/scexec/BUILD.bazel b/pkg/sql/schemachanger/scexec/BUILD.bazel
@@ -17,6 +17,7 @@ go_library(
     importpath = "github.com/cockroachdb/cockroach/pkg/sql/schemachanger/scexec",
     visibility = ["//visibility:public"],
     deps = [
+        "//pkg/config/zonepb",
         "//pkg/jobs",
         "//pkg/jobs/jobspb",
         "//pkg/jobs/jobsprotectedts",
@@ -60,6 +61,7 @@ go_test(
         ":scexec",
         "//pkg/base",
         "//pkg/clusterversion",
+        "//pkg/config/zonepb",
         "//pkg/jobs",
         "//pkg/jobs/jobspb",
         "//pkg/roachpb",

diff --git a/pkg/sql/schemachanger/scexec/dependencies.go b/pkg/sql/schemachanger/scexec/dependencies.go
@@ -13,6 +13,7 @@ package scexec
 import (
 	"context"
 
+	"github.com/cockroachdb/cockroach/pkg/config/zonepb"
 	"github.com/cockroachdb/cockroach/pkg/jobs"
 	"github.com/cockroachdb/cockroach/pkg/jobs/jobspb"
 	"github.com/cockroachdb/cockroach/pkg/jobs/jobsprotectedts"
@@ -71,6 +72,9 @@ type Catalog interface {
 	// DeleteDescriptor deletes a descriptor entry.
 	DeleteDescriptor(ctx context.Context, id descpb.ID) error
 
+	// UpdateZoneConfig upserts a zone config for a descriptor.
+	UpdateZoneConfig(ctx context.Context, id descpb.ID, zc zonepb.ZoneConfig) error
+
 	// DeleteZoneConfig deletes the zone config for a descriptor.
 	DeleteZoneConfig(ctx context.Context, id descpb.ID) error
 

diff --git a/pkg/sql/schemachanger/scexec/exec_immediate_mutation.go b/pkg/sql/schemachanger/scexec/exec_immediate_mutation.go
@@ -14,6 +14,7 @@ import (
 	"context"
 	"sort"
 
+	"github.com/cockroachdb/cockroach/pkg/config/zonepb"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/catalogkeys"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/descpb"
@@ -32,6 +33,7 @@ type immediateState struct {
 	withReset                  bool
 	sequencesToInit            []sequenceToInit
 	temporarySchemasToRegister map[descpb.ID]*temporarySchemaToRegister
+	zoneConfigsToUpdate        []zoneConfigToUpdate
 }
 
 type temporarySchemaToRegister struct {
@@ -51,6 +53,11 @@ type sequenceToInit struct {
 	startVal int64
 }
 
+type zoneConfigToUpdate struct {
+	id descpb.ID
+	zc zonepb.ZoneConfig
+}
+
 var _ scmutationexec.ImmediateMutationStateUpdater = (*immediateState)(nil)
 
 func (s *immediateState) AddToCheckedOutDescriptors(mut catalog.MutableDescriptor) {
@@ -129,6 +136,14 @@ func (s *immediateState) InitSequence(id descpb.ID, startVal int64) {
 		})
 }
 
+func (s *immediateState) UpdateZoneConfig(id descpb.ID, zc zonepb.ZoneConfig) {
+	s.zoneConfigsToUpdate = append(s.zoneConfigsToUpdate,
+		zoneConfigToUpdate{
+			id: id,
+			zc: zc,
+		})
+}
+
 func (s *immediateState) Reset() {
 	s.withReset = true
 }
@@ -194,6 +209,13 @@ func (s *immediateState) exec(ctx context.Context, c Catalog) error {
 	for tempIdxId, tempIdxToRegister := range s.temporarySchemasToRegister {
 		c.InsertTemporarySchema(tempIdxToRegister.schemaName, tempIdxToRegister.parentID, tempIdxId)
 	}
+
+	for _, zc := range s.zoneConfigsToUpdate {
+		if err := c.UpdateZoneConfig(ctx, zc.id, zc.zc); err != nil {
+			return err
+		}
+	}
+
 	return c.Validate(ctx)
 }
 

diff --git a/pkg/sql/schemachanger/scexec/mocks_generated_test.go b/pkg/sql/schemachanger/scexec/mocks_generated_test.go
diff --git a/pkg/sql/schemachanger/scexec/scmutationexec/BUILD.bazel b/pkg/sql/schemachanger/scexec/scmutationexec/BUILD.bazel
@@ -24,6 +24,7 @@ go_library(
     importpath = "github.com/cockroachdb/cockroach/pkg/sql/schemachanger/scexec/scmutationexec",
     visibility = ["//visibility:public"],
     deps = [
+        "//pkg/config/zonepb",
         "//pkg/jobs/jobspb",
         "//pkg/security/username",
         "//pkg/sql/catalog",

diff --git a/pkg/sql/schemachanger/scexec/scmutationexec/database.go b/pkg/sql/schemachanger/scexec/scmutationexec/database.go
@@ -37,3 +37,10 @@ func (i *immediateVisitor) CreateDatabaseDescriptor(
 	i.CreateDescriptor(mut)
 	return nil
 }
+
+func (i *immediateVisitor) AddDatabaseZoneConfig(
+	ctx context.Context, op scop.AddDatabaseZoneConfig,
+) error {
+	i.ImmediateMutationStateUpdater.UpdateZoneConfig(op.DatabaseID, *op.ZoneConfig)
+	return nil
+}
diff --git a/pkg/sql/schemachanger/scexec/scmutationexec/dependencies.go b/pkg/sql/schemachanger/scexec/scmutationexec/dependencies.go
@@ -14,6 +14,7 @@ import (
 	"context"
 	"time"
 
+	"github.com/cockroachdb/cockroach/pkg/config/zonepb"
 	"github.com/cockroachdb/cockroach/pkg/jobs/jobspb"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/catalogkeys"
@@ -82,6 +83,9 @@ type ImmediateMutationStateUpdater interface {
 	// InitSequence initializes a sequence.
 	InitSequence(id descpb.ID, startVal int64)
 
+	// UpdateZoneConfig updates a zone config.
+	UpdateZoneConfig(id descpb.ID, zc zonepb.ZoneConfig)
+
 	// Reset schedules a reset of the in-txn catalog state
 	// to undo the modifications from earlier stages.
 	Reset()

diff --git a/pkg/sql/schemachanger/scop/BUILD.bazel b/pkg/sql/schemachanger/scop/BUILD.bazel
@@ -21,6 +21,7 @@ go_library(
     importpath = "github.com/cockroachdb/cockroach/pkg/sql/schemachanger/scop",
     visibility = ["//visibility:public"],
     deps = [
+        "//pkg/config/zonepb",
         "//pkg/jobs/jobspb",
         "//pkg/sql/catalog/catenumpb",
         "//pkg/sql/catalog/catpb",

diff --git a/pkg/sql/schemachanger/scop/immediate_mutation.go b/pkg/sql/schemachanger/scop/immediate_mutation.go
@@ -11,6 +11,7 @@
 package scop
 
 import (
+	"github.com/cockroachdb/cockroach/pkg/config/zonepb"
 	"github.com/cockroachdb/cockroach/pkg/jobs/jobspb"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/catenumpb"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/catpb"
@@ -884,3 +885,11 @@ type CreateDatabaseDescriptor struct {
 	immediateMutationOp
 	DatabaseID descpb.ID
 }
+
+// AddDatabaseZoneConfig adds a zone config to a database.
+type AddDatabaseZoneConfig struct {
+	immediateMutationOp
+	DatabaseID descpb.ID
+	ZoneConfig *zonepb.ZoneConfig
+	SeqNum     uint32
+}
diff --git a/pkg/sql/schemachanger/scop/immediate_mutation_visitor_generated.go b/pkg/sql/schemachanger/scop/immediate_mutation_visitor_generated.go
diff --git a/pkg/sql/schemachanger/scplan/internal/opgen/BUILD.bazel b/pkg/sql/schemachanger/scplan/internal/opgen/BUILD.bazel
@@ -26,6 +26,7 @@ go_library(
         "opgen_database_data.go",
         "opgen_database_region_config.go",
         "opgen_database_role_setting.go",
+        "opgen_database_zone_config.go",
         "opgen_enum_type.go",
         "opgen_enum_type_value.go",
         "opgen_foreign_key_constraint.go",