opt: synthesize check constraints on enum columns

pkg/sql/opt/exec/execbuilder/testdata/enums

@@ -31,20 +31,20 @@ query T
 EXPLAIN (OPT) SELECT * FROM t WHERE x > 'hello'
 ----
 scan t
- └── constraint: /1: [/'howdy' - ]
+ └── constraint: /1: [/'howdy' - /'hi']
 
 # Test that we can perform constrained scans using secondary indexes too.
 query T
 EXPLAIN (OPT) SELECT * FROM t WHERE y = 'hello'
 ----
 scan t@i
- └── constraint: /2/1: [/'hello' - /'hello']
+ └── constraint: /2/1: [/'hello'/'hello' - /'hello'/'hi']
 
 query T
 EXPLAIN (OPT) SELECT * FROM t WHERE y > 'hello' AND y < 'hi'
 ----
 scan t@i
- └── constraint: /2/1: [/'howdy' - /'howdy']
+ └── constraint: /2/1: [/'howdy'/'hello' - /'howdy'/'hi']
 
 query T
 EXPLAIN (opt) SELECT * FROM t WHERE x IN ('hello', 'hi')
@@ -53,3 +53,16 @@ scan t
  └── constraint: /1
       ├── [/'hello' - /'hello']
       └── [/'hi' - /'hi']
+
+statement ok
+CREATE TABLE checks (x greeting NOT NULL, y int, INDEX (x, y))
+
+# Check that inferred check constraints from enum columns are used in plans.
+query T
+EXPLAIN (OPT) SELECT x, y FROM checks WHERE y = 2
+----
+scan checks@checks_x_y_idx
+ └── constraint: /1/2/3
+      ├── [/'hello'/2 - /'hello'/2]
+      ├── [/'howdy'/2 - /'howdy'/2]
+      └── [/'hi'/2 - /'hi'/2]

pkg/sql/opt_catalog.go

@@ -538,6 +538,11 @@ type optTable struct {
 	outboundFKs []optForeignKeyConstraint
 	inboundFKs  []optForeignKeyConstraint
 
+	// checkConstraints is the set of check constraints for this table. It
+	// can be different from desc's constraints because of synthesized
+	// constraints for user defined types.
+	checkConstraints []cat.CheckConstraint
+
 	// colMap is a mapping from unique ColumnID to column ordinal within the
 	// table. This is a common lookup that needs to be fast.
 	colMap map[sqlbase.ColumnID]int
@@ -626,6 +631,55 @@ func newOptTable(
 		ot.families[i].init(ot, &desc.Families[i+1])
 	}
 
+	// Synthesize any check constraints for user defined types.
+	var synthesizedChecks []cat.CheckConstraint
+	// TODO (rohany): We don't allow referencing columns in mutations in these
+	//  expressions. However, it seems like we will need to have these checks
+	//  operate on columns in mutations. Consider the following case:
+	//  * a user adds a column with an enum type.
+	//  * the column has a default expression of an enum that is not in the
+	//    writeable state.
+	//  * We will need a check constraint here to ensure that writes to the
+	//    column are not successful, but we wouldn't be able to add that now.
+	for i := 0; i < ot.ColumnCount(); i++ {
+		col := ot.Column(i)
+		colType := col.DatumType()
+		if colType.UserDefined() {
+			switch colType.Family() {
+			case types.EnumFamily:
+				// TODO (rohany): When we can alter types, this logic will change.
+				//  In particular, we will want to generate two check constraints if the
+				//  enum contains values that are read only. The first constraint will
+				//  be validated, and contain all of the members of the enum. The second
+				//  will be unvalidated, and will contain only the writeable members of
+				//  the enum. The unvalidated constraint ensures that only writeable
+				//  members of the enum are written. The validated constraint ensures
+				//  that all potentially written values of the enum are considered when
+				//  planning read operations.
+				// We synthesize an (x IN (v1, v2, v3...)) check for enum types.
+				expr := &tree.ComparisonExpr{
+					Operator: tree.In,
+					Left:     &tree.ColumnItem{ColumnName: col.ColName()},
+					Right:    tree.NewDTuple(colType, tree.MakeAllDEnumsInType(colType)...),
+				}
+				synthesizedChecks = append(synthesizedChecks, cat.CheckConstraint{
+					Constraint: tree.Serialize(expr),
+					Validated:  true,
+				})
+			}
+		}
+	}
+	// Move all existing and synthesized checks into the opt table.
+	activeChecks := desc.ActiveChecks()
+	ot.checkConstraints = make([]cat.CheckConstraint, 0, len(activeChecks)+len(synthesizedChecks))
+	for i := range activeChecks {
+		ot.checkConstraints = append(ot.checkConstraints, cat.CheckConstraint{
+			Constraint: activeChecks[i].Expr,
+			Validated:  activeChecks[i].Validity == sqlbase.ConstraintValidity_Validated,
+		})
+	}
+	ot.checkConstraints = append(ot.checkConstraints, synthesizedChecks...)
+
 	// Add stats last, now that other metadata is initialized.
 	if stats != nil {
 		ot.stats = make([]optTableStat, len(stats))
@@ -781,16 +835,12 @@ func (ot *optTable) Statistic(i int) cat.TableStatistic {
 
 // CheckCount is part of the cat.Table interface.
 func (ot *optTable) CheckCount() int {
-	return len(ot.desc.ActiveChecks())
+	return len(ot.checkConstraints)
 }
 
 // Check is part of the cat.Table interface.
 func (ot *optTable) Check(i int) cat.CheckConstraint {
-	check := ot.desc.ActiveChecks()[i]
-	return cat.CheckConstraint{
-		Constraint: check.Expr,
-		Validated:  check.Validity == sqlbase.ConstraintValidity_Validated,
-	}
+	return ot.checkConstraints[i]
 }
 
 // FamilyCount is part of the cat.Table interface.

pkg/sql/sem/tree/datum.go

@@ -3837,6 +3837,21 @@ func MakeDEnumFromLogicalRepresentation(typ *types.T, rep string) (*DEnum, error
 	}, nil
 }
 
+// MakeAllDEnumsInType generates a slice of all values in an enum.
+// TODO (rohany): In the future, take an option of whether to include
+//  non-writeable enum values or not.
+func MakeAllDEnumsInType(typ *types.T) []Datum {
+	result := make([]Datum, len(typ.TypeMeta.EnumData.LogicalRepresentations))
+	for i := 0; i < len(result); i++ {
+		result[i] = &DEnum{
+			EnumTyp:     typ,
+			PhysicalRep: typ.TypeMeta.EnumData.PhysicalRepresentations[i],
+			LogicalRep:  typ.TypeMeta.EnumData.LogicalRepresentations[i],
+		}
+	}
+	return result
+}
+
 // Format implements the NodeFormatter interface.
 func (d *DEnum) Format(ctx *FmtCtx) {
 	if ctx.HasFlags(fmtStaticallyFormatUserDefinedTypes) {