opt: add support for region column inference to catalog and test catalog

This commit adds support for the new storage param
`infer_rbr_region_col_using_constraint` to the optimizer and test catalogs.

Informs #148243

Release note: None

Author: DrewKimball

pkg/sql/opt/cat/table.go

@@ -171,6 +171,13 @@ type Table interface {
 	// different name in a `REGIONAL BY ROW AS` DDL clause.
 	HomeRegionColName() (colName string, ok bool)
 
+	// RegionalByRowUsingConstraint returns the foreign-key constraint that is
+	// used to look up the region for each row in a REGIONAL BY ROW table.
+	// This is only set if the infer_rbr_region_col_using_constraint storage param
+	// is set. If the storage param is not set, or if the table is not
+	// REGIONAL BY ROW, RegionalByRowUsingConstraint returns nil.
+	RegionalByRowUsingConstraint() ForeignKeyConstraint
+
 	// GetDatabaseID returns the owning database id of the table, or zero, if the
 	// owning database could not be determined.
 	GetDatabaseID() descpb.ID

pkg/sql/opt/exec/explain/plan_gist_factory.go

@@ -646,6 +646,11 @@ func (u *unknownTable) HomeRegionColName() (colName string, ok bool) {
 	return "", false
 }
 
+// RegionalByRowUsingConstraint is part of the cat.Table interface.
+func (ot *unknownTable) RegionalByRowUsingConstraint() cat.ForeignKeyConstraint {
+	return nil
+}
+
 // GetDatabaseID is part of the cat.Table interface.
 func (u *unknownTable) GetDatabaseID() descpb.ID {
 	return 0

pkg/sql/opt/testutils/testcat/create_table.go

@@ -421,13 +421,23 @@ OuterLoop:
 			cat.FamilyColumn{Column: col, Ordinal: colOrd})
 	}
 
+	// Allow specifying the name of a FK constraint to use for lookup of region
+	// column values for a REGIONAL BY ROW table.
+	var rbrUsingConstraintName string
+	if param := stmt.StorageParams.GetVal(catpb.RBRUsingConstraintTableSettingName); param != nil {
+		rbrUsingConstraintName = param.(*tree.StrVal).RawString()
+	}
+
 	// Search for foreign key constraints. We want to process them after first
 	// processing all the indexes (otherwise the foreign keys could add
 	// unnecessary indexes).
 	for _, def := range stmt.Defs {
 		switch def := def.(type) {
 		case *tree.ForeignKeyConstraintTableDef:
-			tc.resolveFK(tab, def)
+			fk := tc.resolveFK(tab, def)
+			if rbrUsingConstraintName != "" && string(def.Name) == rbrUsingConstraintName {
+				tab.regionalByRowUsingConstraint = fk
+			}
 		}
 	}
 
@@ -539,7 +549,9 @@ func (tc *Catalog) CreateTableAs(name tree.TableName, columns []cat.Column) *Tab
 }
 
 // resolveFK processes a foreign key constraint.
-func (tc *Catalog) resolveFK(tab *Table, d *tree.ForeignKeyConstraintTableDef) {
+func (tc *Catalog) resolveFK(
+	tab *Table, d *tree.ForeignKeyConstraintTableDef,
+) cat.ForeignKeyConstraint {
 	fromCols := make([]int, len(d.FromCols))
 	for i, c := range d.FromCols {
 		fromCols[i] = tab.FindOrdinal(string(c))
@@ -709,6 +721,7 @@ func (tc *Catalog) resolveFK(tab *Table, d *tree.ForeignKeyConstraintTableDef) {
 	}
 	tab.outboundFKs = append(tab.outboundFKs, fk)
 	targetTable.inboundFKs = append(targetTable.inboundFKs, fk)
+	return &fk
 }
 
 func (tt *Table) addCheckConstraint(check *tree.CheckConstraintTableDef) {

pkg/sql/opt/testutils/testcat/test_catalog.go

@@ -898,7 +898,8 @@ type Table struct {
 
 	multiRegion bool
 
-	implicitRBRIndexElem *tree.IndexElem
+	implicitRBRIndexElem         *tree.IndexElem
+	regionalByRowUsingConstraint cat.ForeignKeyConstraint
 
 	homeRegion string
 
@@ -1104,6 +1105,14 @@ func (tt *Table) HomeRegionColName() (colName string, ok bool) {
 	return string(tree.RegionalByRowRegionDefaultColName), true
 }
 
+// RegionalByRowUsingConstraint is part of the cat.Table interface.
+func (tt *Table) RegionalByRowUsingConstraint() cat.ForeignKeyConstraint {
+	if !tt.IsRegionalByRow() {
+		return nil
+	}
+	return tt.regionalByRowUsingConstraint
+}
+
 // GetDatabaseID is part of the cat.Table interface.
 func (tt *Table) GetDatabaseID() descpb.ID {
 	return tt.DatabaseID

pkg/sql/opt/testutils/testcat/types.go

@@ -97,7 +97,12 @@ func (tc *Catalog) ResolveType(
 
 // ResolveTypeByOID is part of the cat.Catalog interface.
 func (tc *Catalog) ResolveTypeByOID(ctx context.Context, typID oid.Oid) (*types.T, error) {
-	// First look for a matching user-defined enum type.
+	// Look for a builtin type first.
+	if typ, ok := types.OidToType[typID]; ok {
+		return typ, nil
+	}
+
+	// Look for a matching user-defined enum type.
 	for _, typ := range tc.enumTypes {
 		if typ.Oid() == typID {
 			return typ, nil

pkg/sql/opt_catalog.go

@@ -1179,6 +1179,7 @@ func newOptTable(
 			originColumns:     fk.ForeignKeyDesc().OriginColumnIDs,
 			referencedTable:   cat.StableID(fk.GetReferencedTableID()),
 			referencedColumns: fk.ForeignKeyDesc().ReferencedColumnIDs,
+			constraintID:      fk.GetConstraintID(),
 			validity:          fk.GetConstraintValidity(),
 			match:             tree.CompositeKeyMatchMethodType[fk.Match()],
 			deleteAction:      tree.ForeignKeyReferenceActionType[fk.OnDelete()],
@@ -1192,6 +1193,7 @@ func newOptTable(
 			originColumns:     fk.ForeignKeyDesc().OriginColumnIDs,
 			referencedTable:   ot.ID(),
 			referencedColumns: fk.ForeignKeyDesc().ReferencedColumnIDs,
+			constraintID:      fk.GetConstraintID(),
 			validity:          fk.GetConstraintValidity(),
 			match:             tree.CompositeKeyMatchMethodType[fk.Match()],
 			deleteAction:      tree.ForeignKeyReferenceActionType[fk.OnDelete()],
@@ -1570,6 +1572,21 @@ func (ot *optTable) HomeRegionColName() (colName string, ok bool) {
 	return *regionalByRowConfig.As, true
 }
 
+// RegionalByRowUsingConstraint is part of the cat.Table interface.
+func (ot *optTable) RegionalByRowUsingConstraint() cat.ForeignKeyConstraint {
+	if !ot.desc.IsLocalityRegionalByRow() {
+		return nil
+	}
+	if id := ot.desc.GetRegionalByRowUsingConstraint(); id != catid.ConstraintID(0) {
+		for i := range ot.outboundFKs {
+			if ot.outboundFKs[i].constraintID == id {
+				return &ot.outboundFKs[i]
+			}
+		}
+	}
+	return nil
+}
+
 // GetDatabaseID is part of the cat.Table interface.
 func (ot *optTable) GetDatabaseID() descpb.ID {
 	return ot.desc.GetParentID()
@@ -2283,6 +2300,7 @@ type optForeignKeyConstraint struct {
 	referencedTable   cat.StableID
 	referencedColumns []descpb.ColumnID
 
+	constraintID catid.ConstraintID
 	validity     descpb.ConstraintValidity
 	match        tree.CompositeKeyMatchMethod
 	deleteAction tree.ReferenceAction
@@ -2700,6 +2718,11 @@ func (ot *optVirtualTable) HomeRegionColName() (colName string, ok bool) {
 	return "", false
 }
 
+// RegionalByRowUsingConstraint is part of the cat.Table interface.
+func (ot *optVirtualTable) RegionalByRowUsingConstraint() cat.ForeignKeyConstraint {
+	return nil
+}
+
 // GetDatabaseID is part of the cat.Table interface.
 func (ot *optVirtualTable) GetDatabaseID() descpb.ID {
 	return 0