sql: ensure user has correct privileges when adding/removing regions

pkg/ccl/logictestccl/testdata/logic_test/multi_region_privileges

@@ -0,0 +1,102 @@
+# LogicTest: multiregion-9node-3region-3azs
+
+user root
+
+statement ok
+CREATE DATABASE db;
+CREATE TABLE db.t();
+GRANT CREATE ON DATABASE db TO testuser;
+CREATE TABLE db.t2();
+ALTER USER testuser CREATEDB;
+
+user testuser
+
+statement error user testuser must be owner of t or have CREATE privilege on t
+ALTER DATABASE db SET PRIMARY REGION "us-east-1"
+
+user root
+
+statement ok
+GRANT CREATE ON TABLE db.t TO testuser
+
+user testuser
+
+statement ok
+ALTER DATABASE db SET PRIMARY REGION "us-east-1"
+
+user root
+
+statement ok
+REVOKE CREATE ON TABLE db.t FROM testuser
+
+user testuser
+
+statement error user testuser must be owner of t or have CREATE privilege on t
+ALTER DATABASE db DROP REGION "us-east-1"
+
+user root
+
+statement ok
+GRANT CREATE ON TABLE db.t TO testuser
+
+user testuser
+
+statement ok
+ALTER DATABASE db DROP REGION "us-east-1"
+
+# Same thing, but this time testuser is the owner of the table (and doesn't have
+# CREATE privileges on it).
+user root
+
+statement ok
+REVOKE CREATE ON TABLE db.t FROM testuser;
+ALTER TABLE db.t OWNER TO testuser
+
+user testuser
+
+statement ok
+ALTER DATABASE db SET PRIMARY REGION "us-east-1"
+
+statement ok
+ALTER DATABASE db DROP REGION "us-east-1"
+
+subtest alter_table_locality_privs
+
+user root
+
+statement ok
+CREATE DATABASE alter_db PRIMARY REGION "us-east-1";
+CREATE TABLE alter_db.t();
+
+user testuser
+
+statement error pq: user testuser must be owner of t or have CREATE privilege on t
+ALTER TABLE alter_db.t SET LOCALITY GLOBAL
+
+user root
+
+statement ok
+GRANT CREATE ON TABLE alter_db.t TO testuser
+
+user testuser
+
+statement ok
+ALTER TABLE alter_db.t SET LOCALITY GLOBAL
+
+# Same thing, this time make testuser the owner.
+
+user root
+
+statement ok
+REVOKE CREATE ON TABLE alter_db.t FROM testuser
+
+# To be able to gain ownership of the table, testuser needs to have CREATE
+# privilege on the database.
+statement ok
+GRANT CREATE ON DATABASE alter_db to testuser;
+ALTER TABLE alter_db.t OWNER TO testuser
+
+user testuser
+
+statement ok
+ALTER TABLE alter_db.t SET LOCALITY REGIONAL

pkg/sql/alter_database.go

@@ -21,9 +21,11 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/dbdesc"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/descpb"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/multiregion"
+	"github.com/cockroachdb/cockroach/pkg/sql/catalog/tabledesc"
 	"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/privilege"
 	"github.com/cockroachdb/cockroach/pkg/sql/roleoption"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
 	"github.com/cockroachdb/cockroach/pkg/sql/sqlerrors"
@@ -358,6 +360,43 @@ func (p *planner) AlterDatabaseDropRegion(
 	}, nil
 }
 
+// checkPrivilegesForMultiRegionOp ensures the current user has the required
+// privileges to perform a multi-region operation of the given table descriptor.
+// A multi-region operation can be either altering the table's locality or
+// performing a region add/drop that implicitly repartitions the given table.
+// The user must:
+// - either be part of an admin role.
+// - or be an owner of the table.
+// - or have the CREATE privilege on the table.
+// privilege on the table descriptor.
+func (p *planner) checkPrivilegesForMultiRegionOp(
+	ctx context.Context, tableDesc catalog.TableDescriptor,
+) error {
+	hasAdminRole, err := p.HasAdminRole(ctx)
+	if err != nil {
+		return err
+	}
+	if !hasAdminRole {
+		// TODO(arul): It's worth noting CREATE isn't a thing on tables in postgres,
+		// so this will require some changes when (if) we move our privilege system
+		// to be more in line with postgres.
+		err := p.CheckPrivilege(ctx, tableDesc, privilege.CREATE)
+		// Wrap an insufficient privileges error a bit better to reflect the lack
+		// of ownership as well.
+		if pgerror.GetPGCode(err) == pgcode.InsufficientPrivilege {
+			return pgerror.Newf(pgcode.InsufficientPrivilege,
+				"user %s must be owner of %s or have %s privilege on %s",
+				p.SessionData().User(),
+				tableDesc.GetName(),
+				privilege.CREATE,
+				tableDesc.GetName(),
+			)
+		}
+		return err
+	}
+	return nil
+}
+
 // removeLocalityConfigFromAllTablesInDB removes the locality config from all
 // tables under the supplied database.
 func removeLocalityConfigFromAllTablesInDB(
@@ -370,51 +409,48 @@ func removeLocalityConfigFromAllTablesInDB(
 		)
 	}
 	b := p.Txn().NewBatch()
-	if err := forEachTableDesc(ctx, p, desc, hideVirtual,
-		func(immutable *dbdesc.Immutable, _ string, desc catalog.TableDescriptor) error {
-			mutDesc, err := p.Descriptors().GetMutableTableByID(ctx, p.txn, desc.GetID(), tree.ObjectLookupFlags{})
-			if err != nil {
+	if err := p.forEachTableInMultiRegionDatabase(ctx, desc, func(ctx context.Context, tbDesc *tabledesc.Mutable) error {
+		// The user must either be an admin or have the requisite privileges.
+		if err := p.checkPrivilegesForMultiRegionOp(ctx, tbDesc); err != nil {
+			return err
+		}
+
+		switch t := tbDesc.LocalityConfig.Locality.(type) {
+		case *descpb.TableDescriptor_LocalityConfig_Global_:
+			if err := ApplyZoneConfigForMultiRegionTable(
+				ctx,
+				p.txn,
+				p.ExecCfg(),
+				multiregion.RegionConfig{}, // pass dummy config as it is not used.
+				tbDesc,
+				applyZoneConfigForMultiRegionTableOptionRemoveGlobalZoneConfig,
+			); err != nil {
 				return err
 			}
-			switch t := mutDesc.LocalityConfig.Locality.(type) {
-			case *descpb.TableDescriptor_LocalityConfig_Global_:
-				if err := ApplyZoneConfigForMultiRegionTable(
-					ctx,
-					p.txn,
-					p.ExecCfg(),
-					multiregion.RegionConfig{}, // pass dummy config as it is not used.
-					mutDesc,
-					applyZoneConfigForMultiRegionTableOptionRemoveGlobalZoneConfig,
-				); err != nil {
-					return err
-				}
-			case *descpb.TableDescriptor_LocalityConfig_RegionalByTable_:
-				if t.RegionalByTable.Region != nil {
-					// This should error during the type descriptor changes.
-					return errors.AssertionFailedf(
-						"unexpected REGIONAL BY TABLE IN <region> on table %s during DROP REGION",
-						mutDesc.Name,
-					)
-				}
-			case *descpb.TableDescriptor_LocalityConfig_RegionalByRow_:
+		case *descpb.TableDescriptor_LocalityConfig_RegionalByTable_:
+			if t.RegionalByTable.Region != nil {
 				// This should error during the type descriptor changes.
 				return errors.AssertionFailedf(
-					"unexpected REGIONAL BY ROW on table %s during DROP REGION",
-					mutDesc.Name,
+					"unexpected REGIONAL BY TABLE IN <region> on table %s during DROP REGION",
+					tbDesc.Name,
 				)
-			default:
-				return errors.AssertionFailedf(
-					"unexpected locality %T on table %s during DROP REGION",
-					t,
-					mutDesc.Name,
-				)
-			}
-			mutDesc.LocalityConfig = nil
-			if err := p.writeSchemaChangeToBatch(ctx, mutDesc, b); err != nil {
-				return err
 			}
-			return nil
-		}); err != nil {
+		case *descpb.TableDescriptor_LocalityConfig_RegionalByRow_:
+			// This should error during the type descriptor changes.
+			return errors.AssertionFailedf(
+				"unexpected REGIONAL BY ROW on table %s during DROP REGION",
+				tbDesc.Name,
+			)
+		default:
+			return errors.AssertionFailedf(
+				"unexpected locality %T on table %s during DROP REGION",
+				t,
+				tbDesc.Name,
+			)
+		}
+		tbDesc.LocalityConfig = nil
+		return p.writeSchemaChangeToBatch(ctx, tbDesc, b)
+	}); err != nil {
 		return err
 	}
 	return p.Txn().Run(ctx, b)
@@ -608,38 +644,33 @@ func addDefaultLocalityConfigToAllTables(
 		)
 	}
 	b := p.Txn().NewBatch()
-	if err := forEachTableDesc(ctx, p, dbDesc, hideVirtual,
-		func(immutable *dbdesc.Immutable, _ string, desc catalog.TableDescriptor) error {
-			mutDesc, err := p.Descriptors().GetMutableTableByID(
-				ctx, p.txn, desc.GetID(), tree.ObjectLookupFlags{},
-			)
-			if err != nil {
-				return err
-			}
+	if err := p.forEachTableInMultiRegionDatabase(ctx, dbDesc, func(ctx context.Context, tbDesc *tabledesc.Mutable) error {
+		if err := p.checkPrivilegesForMultiRegionOp(ctx, tbDesc); err != nil {
+			return err
+		}
 
-			if err := checkCanConvertTableToMultiRegion(dbDesc, mutDesc); err != nil {
-				return err
-			}
+		if err := checkCanConvertTableToMultiRegion(dbDesc, tbDesc); err != nil {
+			return err
+		}
 
-			if mutDesc.MaterializedView() {
-				if err := p.alterTableDescLocalityToGlobal(
-					ctx, mutDesc, regionEnumID,
-				); err != nil {
-					return err
-				}
-			} else {
-				if err := p.alterTableDescLocalityToRegionalByTable(
-					ctx, tree.PrimaryRegionNotSpecifiedName, mutDesc, regionEnumID,
-				); err != nil {
-					return err
-				}
+		if tbDesc.MaterializedView() {
+			if err := p.alterTableDescLocalityToGlobal(
+				ctx, tbDesc, regionEnumID,
+			); err != nil {
+				return err
 			}
-
-			if err := p.writeSchemaChangeToBatch(ctx, mutDesc, b); err != nil {
+		} else {
+			if err := p.alterTableDescLocalityToRegionalByTable(
+				ctx, tree.PrimaryRegionNotSpecifiedName, tbDesc, regionEnumID,
+			); err != nil {
 				return err
 			}
-			return nil
-		}); err != nil {
+		}
+		if err := p.writeSchemaChangeToBatch(ctx, tbDesc, b); err != nil {
+			return err
+		}
+		return nil
+	}); err != nil {
 		return err
 	}
 	return p.Txn().Run(ctx, b)

pkg/sql/alter_table_locality.go

@@ -25,7 +25,6 @@ import (
 	"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/privilege"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
 	"github.com/cockroachdb/cockroach/pkg/sql/sqlerrors"
 	"github.com/cockroachdb/cockroach/pkg/sql/sqltelemetry"
@@ -61,11 +60,8 @@ func (p *planner) AlterTableLocality(
 		return newZeroNode(nil /* columns */), nil
 	}
 
-	// This check for CREATE privilege is kept for backwards compatibility.
-	if err := p.CheckPrivilege(ctx, tableDesc, privilege.CREATE); err != nil {
-		return nil, pgerror.Newf(pgcode.InsufficientPrivilege,
-			"must be owner of table %s or have CREATE privilege on table %s",
-			tree.Name(tableDesc.GetName()), tree.Name(tableDesc.GetName()))
+	if err := p.checkPrivilegesForMultiRegionOp(ctx, tableDesc); err != nil {
+		return nil, err
 	}
 
 	// Ensure that the database is multi-region enabled.

pkg/sql/authorization.go

@@ -135,6 +135,9 @@ func (p *planner) CheckPrivilegeForUser(
 
 // CheckPrivilege implements the AuthorizationAccessor interface.
 // Requires a valid transaction to be open.
+// TODO(arul): This CheckPrivileges method name is rather deceptive,
+// it should be probably be called CheckPrivilegesOrOwnership and return
+// a better error.
 func (p *planner) CheckPrivilege(
 	ctx context.Context, descriptor catalog.Descriptor, privilege privilege.Kind,
 ) error {