Merge pull request #21004 from benesch/configure-zone-sessiondb

sql: make zone configuration statements respect the session database

pkg/ccl/sqlccl/zone_test.go

@@ -222,6 +222,10 @@ func TestValidIndexPartitionSetShowZones(t *testing.T) {
 	// Ensure the shorthand index syntax works.
 	sqlutils.SetZoneConfig(t, sqlDB, `INDEX "primary"`, yamlOverride)
 	sqlutils.VerifyZoneConfigForTarget(t, sqlDB, `INDEX "primary"`, primaryRow)
+
+	// Ensure the session database is respected.
+	sqlutils.SetZoneConfig(t, sqlDB, "PARTITION p0 OF TABLE t", yamlOverride)
+	sqlutils.VerifyZoneConfigForTarget(t, sqlDB, "PARTITION p0 OF TABLE t", p0Row)
 }
 
 func TestInvalidIndexPartitionSetShowZones(t *testing.T) {

pkg/config/zone.go

@@ -151,13 +151,15 @@ func CLIZoneSpecifier(zs tree.ZoneSpecifier) string {
 		// The index is redundant when the partition is specified, so omit it.
 		ti.Index = ""
 	}
-	return ti.String()
+	return tree.AsStringWithFlags(&ti, tree.FmtSimpleQualified)
 }
 
 // ResolveZoneSpecifier converts a zone specifier to the ID of most specific
 // zone whose config applies.
 func ResolveZoneSpecifier(
-	zs *tree.ZoneSpecifier, resolveName func(parentID uint32, name string) (id uint32, err error),
+	zs *tree.ZoneSpecifier,
+	sessionDB string,
+	resolveName func(parentID uint32, name string) (id uint32, err error),
 ) (uint32, error) {
 	if zs.NamedZone != "" {
 		if zs.NamedZone == DefaultZoneName {
@@ -173,7 +175,7 @@ func ResolveZoneSpecifier(
 		return resolveName(keys.RootNamespaceID, string(zs.Database))
 	}
 
-	tn, err := zs.TableOrIndex.Table.Normalize()
+	tn, err := zs.TableOrIndex.Table.NormalizeWithDatabaseName(sessionDB)
 	if err != nil {
 		return 0, err
 	}

pkg/config/zone_test.go

@@ -279,7 +279,8 @@ func TestZoneSpecifiers(t *testing.T) {
 				if err != nil {
 					return err
 				}
-				id, err := config.ResolveZoneSpecifier(&zs, resolveName)
+				sessionDB := "" // the zone CLI never sets a session DB
+				id, err := config.ResolveZoneSpecifier(&zs, sessionDB, resolveName)
 				if err != nil {
 					return err
 				}

pkg/sql/set_zone_config.go

@@ -80,7 +80,7 @@ func (n *setZoneConfigNode) startExec(params runParams) error {
 		}
 	}
 
-	targetID, err := resolveZone(params.ctx, params.p.txn, &n.zoneSpecifier)
+	targetID, err := resolveZone(params.ctx, params.p.txn, &n.zoneSpecifier, params.p.session.Database)
 	if err != nil {
 		return err
 	}

pkg/sql/show_zone_config.go

@@ -81,7 +81,7 @@ func (n *showZoneConfigNode) startExec(params runParams) error {
 		}
 	}
 
-	targetID, err := resolveZone(params.ctx, params.p.txn, &n.zoneSpecifier)
+	targetID, err := resolveZone(params.ctx, params.p.txn, &n.zoneSpecifier, params.p.session.Database)
 	if err != nil {
 		return err
 	}
@@ -108,11 +108,8 @@ func (n *showZoneConfigNode) startExec(params runParams) error {
 
 	// Determine the CLI specifier for the zone config that actually applies
 	// without performing another KV lookup.
-	zs, err := ascendZoneSpecifier(n.zoneSpecifier, uint32(targetID), zoneID, subzone)
-	if err != nil {
-		return err
-	}
-	n.run.cliSpecifier = config.CLIZoneSpecifier(zs)
+	n.run.cliSpecifier = config.CLIZoneSpecifier(ascendZoneSpecifier(
+		n.zoneSpecifier, uint32(targetID), zoneID, subzone))
 
 	// Ensure subzone configs don't infect the output of config_bytes.
 	zone.Subzones = nil
@@ -157,19 +154,15 @@ func (*showZoneConfigNode) Close(context.Context) {}
 // finds without impacting performance.
 func ascendZoneSpecifier(
 	zs tree.ZoneSpecifier, resolvedID, actualID uint32, actualSubzone *config.Subzone,
-) (tree.ZoneSpecifier, error) {
+) tree.ZoneSpecifier {
 	if actualID == keys.RootNamespaceID {
 		// We had to traverse to the top of the hierarchy, so we're showing the
 		// default zone config.
 		zs.NamedZone = config.DefaultZoneName
 	} else if resolvedID != actualID {
 		// We traversed at least one level up, and we're not at the top of the
 		// hierarchy, so we're showing the database zone config.
-		tn, err := zs.TableOrIndex.Table.Normalize()
-		if err != nil {
-			return tree.ZoneSpecifier{}, err
-		}
-		zs.Database = tn.DatabaseName
+		zs.Database = zs.TableOrIndex.Table.TableName().DatabaseName
 	} else if actualSubzone == nil {
 		// We didn't find a subzone, so no index or partition zone config exists.
 		zs.TableOrIndex.Index = ""
@@ -178,5 +171,5 @@ func ascendZoneSpecifier(
 		// The resolved subzone did not name a partition, just an index.
 		zs.Partition = ""
 	}
-	return zs, nil
+	return zs
 }

pkg/sql/zone_config.go

@@ -173,9 +173,11 @@ func zoneSpecifierNotFoundError(zs tree.ZoneSpecifier) error {
 	}
 }
 
-func resolveZone(ctx context.Context, txn *client.Txn, zs *tree.ZoneSpecifier) (sqlbase.ID, error) {
+func resolveZone(
+	ctx context.Context, txn *client.Txn, zs *tree.ZoneSpecifier, sessionDB string,
+) (sqlbase.ID, error) {
 	errMissingKey := errors.New("missing key")
-	id, err := config.ResolveZoneSpecifier(zs,
+	id, err := config.ResolveZoneSpecifier(zs, sessionDB,
 		func(parentID uint32, name string) (uint32, error) {
 			kv, err := txn.Get(ctx, sqlbase.MakeNameMetadataKey(sqlbase.ID(parentID), name))
 			if err != nil {

pkg/sql/zone_test.go

@@ -199,6 +199,12 @@ func TestValidSetShowZones(t *testing.T) {
 		defaultRow, systemRow, jobsRow, livenessRow)
 	sqlutils.VerifyZoneConfigForTarget(t, sqlDB, "TABLE system.jobs", jobsRow)
 
+	// Verify that the session database is respected.
+	sqlutils.SetZoneConfig(t, sqlDB, "TABLE t", yamlOverride)
+	sqlutils.VerifyZoneConfigForTarget(t, sqlDB, "TABLE t", tableRow)
+	sqlutils.DeleteZoneConfig(t, sqlDB, "TABLE t")
+	sqlutils.VerifyZoneConfigForTarget(t, sqlDB, "TABLE t", defaultRow)
+
 	// Ensure zone configs are read transactionally instead of from the cached
 	// system config.
 	txn, err := db.Begin()
@@ -238,27 +244,35 @@ func TestInvalidSetShowZones(t *testing.T) {
 		},
 		{
 			"ALTER RANGE foo EXPERIMENTAL CONFIGURE ZONE ''",
-			"\"foo\" is not a built-in zone",
+			`"foo" is not a built-in zone`,
 		},
 		{
 			"ALTER DATABASE foo EXPERIMENTAL CONFIGURE ZONE ''",
-			"database \"foo\" does not exist",
+			`database "foo" does not exist`,
+		},
+		{
+			"ALTER TABLE system.foo EXPERIMENTAL CONFIGURE ZONE ''",
+			`relation "system.foo" does not exist`,
 		},
 		{
 			"ALTER TABLE foo EXPERIMENTAL CONFIGURE ZONE ''",
-			"relation \"foo\" does not exist",
+			`no database specified: "foo"`,
 		},
 		{
 			"EXPERIMENTAL SHOW ZONE CONFIGURATION FOR RANGE foo",
-			"\"foo\" is not a built-in zone",
+			`"foo" is not a built-in zone`,
 		},
 		{
 			"EXPERIMENTAL SHOW ZONE CONFIGURATION FOR DATABASE foo",
-			"database \"foo\" does not exist",
+			`database "foo" does not exist`,
 		},
 		{
 			"EXPERIMENTAL SHOW ZONE CONFIGURATION FOR TABLE foo",
-			"relation \"foo\" does not exist",
+			`no database specified: "foo"`,
+		},
+		{
+			"EXPERIMENTAL SHOW ZONE CONFIGURATION FOR TABLE system.foo",
+			`relation "system.foo" does not exist`,
 		},
 	} {
 		if _, err := db.Exec(tc.query); !testutils.IsError(err, tc.err) {