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 cockroachdb#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,16 @@
 [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 +460,8 @@
 (Project
     $input:(Insert | Update | Upsert | Delete
         $innerInput:*
-        $checks:*
+        $uniqueChecks:*
+        $fkChecks:*
         $mutationPrivate:*
     )
     $projections:*
@@ -475,7 +479,8 @@
 (Project
     ((OpName $input)
         $innerInput
-        $checks
+        $uniqueChecks
+        $fkChecks
         (PruneMutationReturnCols $mutationPrivate $needed)
     )
     $projections

pkg/sql/opt/ops/mutation.opt

@@ -19,8 +19,9 @@
 # mutation columns) by SQL users.
 [Relational, Mutation, WithBinding]
 define Insert {
-    Input RelExpr
-    Checks FKChecksExpr
+    Input        RelExpr
+    UniqueChecks UniqueChecksExpr
+    FKChecks     FKChecksExpr
     _ MutationPrivate
 }
 
@@ -181,8 +182,9 @@ define MutationPrivate {
 # columns that are computed.
 [Relational, Mutation, WithBinding]
 define Update {
-    Input RelExpr
-    Checks FKChecksExpr
+    Input        RelExpr
+    UniqueChecks UniqueChecksExpr
+    FKChecks     FKChecksExpr
     _ MutationPrivate
 }
 
@@ -204,8 +206,9 @@ define Update {
 # mutation columns that are computed.
 [Relational, Mutation, WithBinding]
 define Upsert {
-    Input RelExpr
-    Checks FKChecksExpr
+    Input        RelExpr
+    UniqueChecks UniqueChecksExpr
+    FKChecks     FKChecksExpr
     _ MutationPrivate
 }
 
@@ -216,8 +219,9 @@ define Upsert {
 #
 [Relational, Mutation, WithBinding]
 define Delete {
-    Input RelExpr
-    Checks FKChecksExpr
+    Input        RelExpr
+    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 unique 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 item checks that a new value in the table is unique.
+    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/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)
 }

pkg/sql/opt/optbuilder/insert.go

@@ -652,10 +652,14 @@ func (mb *mutationBuilder) buildInsert(returning tree.ReturningExprs) {
 	// Add any partial index put boolean columns to the input.
 	mb.projectPartialIndexPutCols(preCheckScope)
 
+	mb.buildUniqueChecksForInsert()
+
 	mb.buildFKChecksForInsert()
 
 	private := mb.makeMutationPrivate(returning != nil)
-	mb.outScope.expr = mb.b.factory.ConstructInsert(mb.outScope.expr, mb.checks, private)
+	mb.outScope.expr = mb.b.factory.ConstructInsert(
+		mb.outScope.expr, mb.uniqueChecks, mb.fkChecks, private,
+	)
 
 	mb.buildReturning(returning)
 }
@@ -1089,7 +1093,9 @@ func (mb *mutationBuilder) buildUpsert(returning tree.ReturningExprs) {
 	mb.buildFKChecksForUpsert()
 
 	private := mb.makeMutationPrivate(returning != nil)
-	mb.outScope.expr = mb.b.factory.ConstructUpsert(mb.outScope.expr, mb.checks, private)
+	mb.outScope.expr = mb.b.factory.ConstructUpsert(
+		mb.outScope.expr, mb.uniqueChecks, mb.fkChecks, private,
+	)
 
 	mb.buildReturning(returning)
 }