Merge #61169

61169: multiregionccl: gate ADD REGION under ccl / enterprise license r=ajstorm a=otan Release note (enterprise change): ALTER DATABASE ... ADD REGION ... now requires an enterprise license. Release justification: low risk fix for new functionality Co-authored-by: Oliver Tan <[email protected]>
cockroachdb · Feb 28, 2021 · 9097271 · 9097271
2 parents 4bb49c3 + 6006e42
commit 9097271
Show file tree

Hide file tree

Showing 4 changed files with 79 additions and 29 deletions.
diff --git a/pkg/ccl/multiregionccl/BUILD.bazel b/pkg/ccl/multiregionccl/BUILD.bazel
@@ -11,6 +11,7 @@ go_library(
         "//pkg/sql",
         "//pkg/sql/catalog/catalogkv",
         "//pkg/sql/catalog/descpb",
+        "//pkg/sql/catalog/typedesc",
         "//pkg/sql/pgwire/pgcode",
         "//pkg/sql/pgwire/pgerror",
         "//pkg/sql/sem/tree",

diff --git a/pkg/ccl/multiregionccl/multiregion.go b/pkg/ccl/multiregionccl/multiregion.go
@@ -17,13 +17,15 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/catalogkv"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/descpb"
+	"github.com/cockroachdb/cockroach/pkg/sql/catalog/typedesc"
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgcode"
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgerror"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
 )
 
 func init() {
 	sql.CreateRegionConfigCCL = createRegionConfig
+	sql.GetMultiRegionEnumAddValuePlacementCCL = getMultiRegionEnumAddValuePlacement
 }
 
 func createRegionConfig(
@@ -128,3 +130,43 @@ func checkClusterSupportsMultiRegion(evalCtx *tree.EvalContext) error {
 	}
 	return nil
 }
+
+func getMultiRegionEnumAddValuePlacement(
+	execCfg *sql.ExecutorConfig, typeDesc *typedesc.Mutable, region tree.Name,
+) (tree.AlterTypeAddValue, error) {
+	if err := utilccl.CheckEnterpriseEnabled(
+		execCfg.Settings,
+		execCfg.ClusterID(),
+		execCfg.Organization(),
+		"ADD REGION",
+	); err != nil {
+		return tree.AlterTypeAddValue{}, err
+	}
+
+	// Find the location in the enum where we should insert the new value. We much search
+	// for the location (and not append to the end), as we want to keep the values in sorted
+	// order.
+	loc := sort.Search(
+		len(typeDesc.EnumMembers),
+		func(i int) bool {
+			return string(region) < typeDesc.EnumMembers[i].LogicalRepresentation
+		},
+	)
+
+	// If the above search couldn't find a value greater than the region being added, add the
+	// new region at the end of the enum.
+	before := true
+	if loc == len(typeDesc.EnumMembers) {
+		before = false
+		loc = len(typeDesc.EnumMembers) - 1
+	}
+
+	return tree.AlterTypeAddValue{
+		IfNotExists: false,
+		NewVal:      tree.EnumValue(region),
+		Placement: &tree.AlterTypeAddValuePlacement{
+			Before:      before,
+			ExistingVal: tree.EnumValue(typeDesc.EnumMembers[loc].LogicalRepresentation),
+		},
+	}, nil
+}
diff --git a/pkg/ccl/multiregionccl/multiregion_test.go b/pkg/ccl/multiregionccl/multiregion_test.go
@@ -73,13 +73,23 @@ CREATE TABLE t4 () LOCALITY REGIONAL BY ROW
 
 	// Test certain commands are no longer usable.
 	t.Run("no new multi-region items", func(t *testing.T) {
-		for _, errorStmt := range []string{
-			`CREATE DATABASE db WITH PRIMARY REGION "us-east1" REGIONS "us-east2"`,
+		for _, tc := range []struct {
+			stmt             string
+			expectedContains string
+		}{
+			{
+				stmt:             `CREATE DATABASE db WITH PRIMARY REGION "us-east1" REGIONS "us-east2"`,
+				expectedContains: multiRegionNoEnterpriseContains,
+			},
+			{
+				stmt:             `ALTER DATABASE test ADD REGION "us-east3"`,
+				expectedContains: "use of ADD REGION requires an enterprise license",
+			},
 		} {
-			t.Run(errorStmt, func(t *testing.T) {
-				_, err := sqlDB.Exec(errorStmt)
+			t.Run(tc.stmt, func(t *testing.T) {
+				_, err := sqlDB.Exec(tc.stmt)
 				require.Error(t, err)
-				require.Contains(t, err.Error(), multiRegionNoEnterpriseContains)
+				require.Contains(t, err.Error(), tc.expectedContains)
 			})
 		}
 	})

diff --git a/pkg/sql/alter_database.go b/pkg/sql/alter_database.go
@@ -13,7 +13,6 @@ package sql
 import (
 	"context"
 	"fmt"
-	"sort"
 
 	"github.com/cockroachdb/cockroach/pkg/keys"
 	"github.com/cockroachdb/cockroach/pkg/security"
@@ -26,6 +25,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgerror"
 	"github.com/cockroachdb/cockroach/pkg/sql/roleoption"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
+	"github.com/cockroachdb/cockroach/pkg/sql/sqlerrors"
 	"github.com/cockroachdb/cockroach/pkg/util/log/eventpb"
 	"github.com/cockroachdb/errors"
 )
@@ -138,6 +138,17 @@ func (p *planner) AlterDatabaseAddRegion(
 	return &alterDatabaseAddRegionNode{n: n, desc: dbDesc}, nil
 }
 
+// GetMultiRegionEnumAddValuePlacementCCL is the public hook point for the
+// CCL-licensed code to determine the placement for a new region inside
+// a region enum.
+var GetMultiRegionEnumAddValuePlacementCCL = func(
+	execCfg *ExecutorConfig, typeDesc *typedesc.Mutable, region tree.Name,
+) (tree.AlterTypeAddValue, error) {
+	return tree.AlterTypeAddValue{}, sqlerrors.NewCCLRequiredError(
+		errors.New("adding regions to a multi-region database requires a CCL binary"),
+	)
+}
+
 func (n *alterDatabaseAddRegionNode) startExec(params runParams) error {
 	// To add a region, the user has to have CREATEDB privileges, or be an admin user.
 	if err := params.p.CheckRoleOption(params.ctx, roleoption.CREATEDB); err != nil {
@@ -185,22 +196,13 @@ func (n *alterDatabaseAddRegionNode) startExec(params runParams) error {
 		return err
 	}
 
-	// Find the location in the enum where we should insert the new value. We much search
-	// for the location (and not append to the end), as we want to keep the values in sorted
-	// order.
-	loc := sort.Search(
-		len(typeDesc.EnumMembers),
-		func(i int) bool {
-			return string(n.n.Region) < typeDesc.EnumMembers[i].LogicalRepresentation
-		},
+	placement, err := GetMultiRegionEnumAddValuePlacementCCL(
+		params.p.ExecCfg(),
+		typeDesc,
+		n.n.Region,
 	)
-
-	// If the above search couldn't find a value greater than the region being added, add the
-	// new region at the end of the enum.
-	before := true
-	if loc == len(typeDesc.EnumMembers) {
-		before = false
-		loc = len(typeDesc.EnumMembers) - 1
+	if err != nil {
+		return err
 	}
 
 	// Add the new region value to the enum. This function adds the value to the enum and
@@ -209,14 +211,9 @@ func (n *alterDatabaseAddRegionNode) startExec(params runParams) error {
 	if err := params.p.addEnumValue(
 		params.ctx,
 		typeDesc,
-		&tree.AlterTypeAddValue{
-			IfNotExists: false,
-			NewVal:      tree.EnumValue(n.n.Region),
-			Placement: &tree.AlterTypeAddValuePlacement{
-				Before:      before,
-				ExistingVal: tree.EnumValue(typeDesc.EnumMembers[loc].LogicalRepresentation),
-			}},
-		jobDesc); err != nil {
+		&placement,
+		jobDesc,
+	); err != nil {
 		return err
 	}