opt: add insertion checks for unique constraints

This commit adds checks for unique constraints when planning insertions
in the optimizer. This does not yet impact anything outside of the optimizer
tests, since UNIQUE WITHOUT INDEX is still not supported outside of the
optimizer test catalog.

Informs #41535

Release note: None

pkg/sql/opt/exec/execbuilder/mutation.go

@@ -45,8 +45,8 @@ func (b *Builder) buildMutationInput(
 	}
 
 	if p.WithID != 0 {
-		// The input might have extra columns that are used only by FK checks; make
-		// sure we don't project them away.
+		// The input might have extra columns that are used only by FK or unique
+		// checks; make sure we don't project them away.
 		cols := inputExpr.Relational().OutputCols.Copy()
 		for _, c := range colList {
 			cols.Remove(c)
@@ -101,7 +101,8 @@ func (b *Builder) buildInsert(ins *memo.InsertExpr) (execPlan, error) {
 		insertOrds,
 		returnOrds,
 		checkOrds,
-		b.allowAutoCommit && len(ins.Checks) == 0 && len(ins.FKCascades) == 0,
+		b.allowAutoCommit && len(ins.UniqueChecks) == 0 &&
+			len(ins.FKChecks) == 0 && len(ins.FKCascades) == 0,
 	)
 	if err != nil {
 		return execPlan{}, err
@@ -112,7 +113,9 @@ func (b *Builder) buildInsert(ins *memo.InsertExpr) (execPlan, error) {
 		ep.outputCols = mutationOutputColMap(ins)
 	}
 
-	if err := b.buildFKChecks(ins.Checks); err != nil {
+	// TODO(rytaft): build unique checks.
+
+	if err := b.buildFKChecks(ins.FKChecks); err != nil {
 		return execPlan{}, err
 	}
 
@@ -148,9 +151,9 @@ func (b *Builder) tryBuildFastPathInsert(ins *memo.InsertExpr) (_ execPlan, ok b
 
 	//  - there are no self-referencing foreign keys;
 	//  - all FK checks can be performed using direct lookups into unique indexes.
-	fkChecks := make([]exec.InsertFastPathFKCheck, len(ins.Checks))
-	for i := range ins.Checks {
-		c := &ins.Checks[i]
+	fkChecks := make([]exec.InsertFastPathFKCheck, len(ins.FKChecks))
+	for i := range ins.FKChecks {
+		c := &ins.FKChecks[i]
 		if md.Table(c.ReferencedTable).ID() == md.Table(ins.Table).ID() {
 			// Self-referencing FK.
 			return execPlan{}, false, nil
@@ -325,13 +328,16 @@ func (b *Builder) buildUpdate(upd *memo.UpdateExpr) (execPlan, error) {
 		returnColOrds,
 		checkOrds,
 		passthroughCols,
-		b.allowAutoCommit && len(upd.Checks) == 0 && len(upd.FKCascades) == 0,
+		b.allowAutoCommit && len(upd.UniqueChecks) == 0 &&
+			len(upd.FKChecks) == 0 && len(upd.FKCascades) == 0,
 	)
 	if err != nil {
 		return execPlan{}, err
 	}
 
-	if err := b.buildFKChecks(upd.Checks); err != nil {
+	// TODO(rytaft): build unique checks.
+
+	if err := b.buildFKChecks(upd.FKChecks); err != nil {
 		return execPlan{}, err
 	}
 
@@ -406,13 +412,16 @@ func (b *Builder) buildUpsert(ups *memo.UpsertExpr) (execPlan, error) {
 		updateColOrds,
 		returnColOrds,
 		checkOrds,
-		b.allowAutoCommit && len(ups.Checks) == 0 && len(ups.FKCascades) == 0,
+		b.allowAutoCommit && len(ups.UniqueChecks) == 0 &&
+			len(ups.FKChecks) == 0 && len(ups.FKCascades) == 0,
 	)
 	if err != nil {
 		return execPlan{}, err
 	}
 
-	if err := b.buildFKChecks(ups.Checks); err != nil {
+	// TODO(rytaft): build unique checks.
+
+	if err := b.buildFKChecks(ups.FKChecks); err != nil {
 		return execPlan{}, err
 	}
 
@@ -460,13 +469,13 @@ func (b *Builder) buildDelete(del *memo.DeleteExpr) (execPlan, error) {
 		tab,
 		fetchColOrds,
 		returnColOrds,
-		b.allowAutoCommit && len(del.Checks) == 0 && len(del.FKCascades) == 0,
+		b.allowAutoCommit && len(del.FKChecks) == 0 && len(del.FKCascades) == 0,
 	)
 	if err != nil {
 		return execPlan{}, err
 	}
 
-	if err := b.buildFKChecks(del.Checks); err != nil {
+	if err := b.buildFKChecks(del.FKChecks); err != nil {
 		return execPlan{}, err
 	}
 
@@ -528,7 +537,7 @@ func (b *Builder) tryBuildDeleteRange(del *memo.DeleteExpr) (_ execPlan, ok bool
 	if err != nil {
 		return execPlan{}, false, err
 	}
-	if err := b.buildFKChecks(del.Checks); err != nil {
+	if err := b.buildFKChecks(del.FKChecks); err != nil {
 		return execPlan{}, false, err
 	}
 	if err := b.buildFKCascades(del.WithID, del.FKCascades); err != nil {
@@ -677,7 +686,7 @@ func (b *Builder) buildDeleteRange(
 					autoCommit = true
 				}
 			}
-			if len(del.Checks) > 0 || len(del.FKCascades) > 0 {
+			if len(del.FKChecks) > 0 || len(del.FKCascades) > 0 {
 				// Do not allow autocommit if we have checks or cascades. This does not
 				// apply for the interleaved case, where we decided that the delete
 				// range takes care of all the FKs as well.

pkg/sql/opt/memo/expr_format.go

@@ -785,7 +785,7 @@ func (f *ExprFmtCtx) formatScalarWithLabel(
 		f.Buffer.WriteString(": ")
 	}
 	switch scalar.Op() {
-	case opt.ProjectionsOp, opt.AggregationsOp, opt.FKChecksOp, opt.KVOptionsOp:
+	case opt.ProjectionsOp, opt.AggregationsOp, opt.UniqueChecksOp, opt.FKChecksOp, opt.KVOptionsOp:
 		// Omit empty lists (except filters).
 		if scalar.ChildCount() == 0 {
 			return
@@ -912,7 +912,8 @@ func (f *ExprFmtCtx) formatScalarWithLabel(
 func (f *ExprFmtCtx) scalarPropsStrings(scalar opt.ScalarExpr) []string {
 	typ := scalar.DataType()
 	if typ == nil {
-		if scalar.Op() == opt.FKChecksItemOp || scalar.Op() == opt.KVOptionsItemOp {
+		if scalar.Op() == opt.UniqueChecksItemOp || scalar.Op() == opt.FKChecksItemOp ||
+			scalar.Op() == opt.KVOptionsItemOp {
 			// These are not true scalars and have no properties.
 			return nil
 		}
@@ -999,6 +1000,19 @@ func (f *ExprFmtCtx) formatScalarPrivate(scalar opt.ScalarExpr) {
 	case *KVOptionsItem:
 		fmt.Fprintf(f.Buffer, " %s", t.Key)
 
+	case *UniqueChecksItem:
+		tab := f.Memo.metadata.TableMeta(t.Table)
+		constraint := tab.Table.Unique(t.CheckOrdinal)
+		fmt.Fprintf(f.Buffer, ": %s(", tab.Alias.ObjectName)
+		for i := 0; i < constraint.ColumnCount(); i++ {
+			if i > 0 {
+				f.Buffer.WriteByte(',')
+			}
+			col := tab.Table.Column(constraint.ColumnOrdinal(tab.Table, i))
+			f.Buffer.WriteString(string(col.ColName()))
+		}
+		f.Buffer.WriteByte(')')
+
 	case *FKChecksItem:
 		origin := f.Memo.metadata.TableMeta(t.OriginTable)
 		referenced := f.Memo.metadata.TableMeta(t.ReferencedTable)

pkg/sql/opt/memo/interner.go

@@ -645,6 +645,12 @@ func (h *hasher) HashFKChecksExpr(val FKChecksExpr) {
 	}
 }
 
+func (h *hasher) HashUniqueChecksExpr(val UniqueChecksExpr) {
+	for i := range val {
+		h.HashRelExpr(val[i].Check)
+	}
+}
+
 func (h *hasher) HashKVOptionsExpr(val KVOptionsExpr) {
 	for i := range val {
 		h.HashString(val[i].Key)
@@ -1050,6 +1056,18 @@ func (h *hasher) IsFKChecksExprEqual(l, r FKChecksExpr) bool {
 	return true
 }
 
+func (h *hasher) IsUniqueChecksExprEqual(l, r UniqueChecksExpr) bool {
+	if len(l) != len(r) {
+		return false
+	}
+	for i := range l {
+		if l[i].Check != r[i].Check {
+			return false
+		}
+	}
+	return true
+}
+
 func (h *hasher) IsKVOptionsExprEqual(l, r KVOptionsExpr) bool {
 	if len(l) != len(r) {
 		return false

pkg/sql/opt/norm/prune_cols_funcs.go

@@ -43,7 +43,7 @@ func (c *CustomFuncs) NeededExplainCols(private *memo.ExplainPrivate) opt.ColSet
 // referenced by it. Other rules filter the FetchCols, CheckCols, etc. and can
 // in turn trigger the PruneMutationInputCols rule.
 func (c *CustomFuncs) NeededMutationCols(
-	private *memo.MutationPrivate, checks memo.FKChecksExpr,
+	private *memo.MutationPrivate, uniqueChecks memo.UniqueChecksExpr, fkChecks memo.FKChecksExpr,
 ) opt.ColSet {
 	var cols opt.ColSet
 
@@ -69,8 +69,12 @@ func (c *CustomFuncs) NeededMutationCols(
 	}
 
 	if private.WithID != 0 {
-		for i := range checks {
-			withUses := memo.WithUses(checks[i].Check)
+		for i := range uniqueChecks {
+			withUses := memo.WithUses(uniqueChecks[i].Check)
+			cols.UnionWith(withUses[private.WithID].UsedCols)
+		}
+		for i := range fkChecks {
+			withUses := memo.WithUses(fkChecks[i].Check)
 			cols.UnionWith(withUses[private.WithID].UsedCols)
 		}
 	}

pkg/sql/opt/norm/rules/prune_cols.opt

@@ -415,7 +415,8 @@
 [PruneMutationFetchCols, Normalize]
 (Update | Upsert | Delete
     $input:*
-    $checks:*
+    $uniqueChecks:*
+    $fkChecks:*
     $mutationPrivate:* &
         (CanPruneMutationFetchCols
             $mutationPrivate
@@ -428,7 +429,8 @@
 =>
 ((OpName)
     $input
-    $checks
+    $uniqueChecks
+    $fkChecks
     (PruneMutationFetchCols $mutationPrivate $needed)
 )
 
@@ -437,15 +439,25 @@
 [PruneMutationInputCols, Normalize]
 (Insert | Update | Upsert | Delete
     $input:*
-    $checks:*
+    $uniqueChecks:*
+    $fkChecks:*
     $mutationPrivate:* &
         (CanPruneCols
             $input
-            $needed:(NeededMutationCols $mutationPrivate $checks)
+            $needed:(NeededMutationCols
+                $mutationPrivate
+                $uniqueChecks
+                $fkChecks
+            )
         )
 )
 =>
-((OpName) (PruneCols $input $needed) $checks $mutationPrivate)
+((OpName)
+    (PruneCols $input $needed)
+    $uniqueChecks
+    $fkChecks
+    $mutationPrivate
+)
 
 # PruneReturningCols removes columns from the mutation operator's ReturnCols
 # set if they are not used in the RETURNING clause of the mutation.
@@ -457,7 +469,8 @@
 (Project
     $input:(Insert | Update | Upsert | Delete
         $innerInput:*
-        $checks:*
+        $uniqueChecks:*
+        $fkChecks:*
         $mutationPrivate:*
     )
     $projections:*
@@ -475,7 +488,8 @@
 (Project
     ((OpName $input)
         $innerInput
-        $checks
+        $uniqueChecks
+        $fkChecks
         (PruneMutationReturnCols $mutationPrivate $needed)
     )
     $projections

pkg/sql/opt/norm/testdata/rules/prune_cols

@@ -2585,6 +2585,9 @@ upsert checks
            ├── upsert_b:16 > 10 [as=check2:18, outer=(16)]
            └── column2:7 > upsert_b:16 [as=check3:19, outer=(7,16)]
 
+# TODO(rytaft): test that columns needed for unique checks are not pruned
+# from updates.
+
 # ------------------------------------------------------------------------------
 # PruneMutationReturnCols
 # ------------------------------------------------------------------------------

pkg/sql/opt/ops/mutation.opt

@@ -20,7 +20,8 @@
 [Relational, Mutation, WithBinding]
 define Insert {
     Input RelExpr
-    Checks FKChecksExpr
+    UniqueChecks UniqueChecksExpr
+    FKChecks FKChecksExpr
     _ MutationPrivate
 }
 
@@ -182,7 +183,8 @@ define MutationPrivate {
 [Relational, Mutation, WithBinding]
 define Update {
     Input RelExpr
-    Checks FKChecksExpr
+    UniqueChecks UniqueChecksExpr
+    FKChecks FKChecksExpr
     _ MutationPrivate
 }
 
@@ -205,7 +207,8 @@ define Update {
 [Relational, Mutation, WithBinding]
 define Upsert {
     Input RelExpr
-    Checks FKChecksExpr
+    UniqueChecks UniqueChecksExpr
+    FKChecks FKChecksExpr
     _ MutationPrivate
 }
 
@@ -217,7 +220,8 @@ define Upsert {
 [Relational, Mutation, WithBinding]
 define Delete {
     Input RelExpr
-    Checks FKChecksExpr
+    UniqueChecks UniqueChecksExpr
+    FKChecks FKChecksExpr
     _ MutationPrivate
 }
 
@@ -257,3 +261,32 @@ define FKChecksItemPrivate {
     # OpName is the name that should be used for this check in error messages.
     OpName string
 }
+
+# UniqueChecks is a list of uniqueness check queries, to be run after the main
+# query.
+[Scalar, List]
+define UniqueChecks {
+}
+
+# UniqueChecksItem is a unique check query, to be run after the main query.
+# An execution error will be generated if the query returns any results.
+[Scalar, ListItem]
+define UniqueChecksItem {
+    Check RelExpr
+    _ UniqueChecksItemPrivate
+}
+
+[Private]
+define UniqueChecksItemPrivate {
+    Table TableID
+
+    # This is the ordinal of the check in the table's unique constraints.
+    CheckOrdinal int
+
+    # KeyCols are the columns in the Check query that form the value tuple shown
+    # in the error message.
+    KeyCols ColList
+
+    # OpName is the name that should be used for this check in error messages.
+    OpName string
+}

pkg/sql/opt/optbuilder/BUILD.bazel

@@ -20,6 +20,7 @@ go_library(
         "misc_statements.go",
         "mutation_builder.go",
         "mutation_builder_fk.go",
+        "mutation_builder_unique.go",
         "opaque.go",
         "orderby.go",
         "partial_index.go",

pkg/sql/opt/optbuilder/delete.go

@@ -77,7 +77,9 @@ func (mb *mutationBuilder) buildDelete(returning tree.ReturningExprs) {
 	mb.buildFKChecksAndCascadesForDelete()
 
 	private := mb.makeMutationPrivate(returning != nil)
-	mb.outScope.expr = mb.b.factory.ConstructDelete(mb.outScope.expr, mb.checks, private)
+	mb.outScope.expr = mb.b.factory.ConstructDelete(
+		mb.outScope.expr, mb.uniqueChecks, mb.fkChecks, private,
+	)
 
 	mb.buildReturning(returning)
 }